1import {
2 EnvironmentAuthorizationError,
3 clientRpcRequiredScopes,
4 authScopeRequiredResponse,
5 type EnvironmentId,
6 type ClientGuardedRpcTag,
7 ORCHESTRATION_V2_WS_METHODS,
8 WS_METHODS,
9} from "@t3tools/contracts";
10import * as Cause from "effect/Cause";
11import * as Context from "effect/Context";
12import * as Duration from "effect/Duration";
13import * as Effect from "effect/Effect";
14import * as Option from "effect/Option";
15import * as Schema from "effect/Schema";
16import * as Stream from "effect/Stream";
17import * as SubscriptionRef from "effect/SubscriptionRef";
18import { RpcClientError } from "effect/rpc";
19
20import * as EnvironmentSupervisor from "../connection/supervisor.ts";
21import type { WsRpcProtocolClient } from "../rpc/protocol.ts";
22import type { RpcSession } from "../rpc/session.ts";
23
24export class EnvironmentRpcUnavailableError extends Schema.TaggedError<EnvironmentRpcUnavailableError>()(
25 "EnvironmentRpcUnavailableError",
26 {
27 environmentId: Schema.String,
28 message: Schema.String,
29 },
30) {}
31
32export interface EnvironmentRpcRequestObservation {
33 readonly environmentId: string;
34 readonly method: string;
35}
36
37export class EnvironmentRpcRequestObserver extends Context.Reference<{
38 readonly observe: (
39 request: EnvironmentRpcRequestObservation,
40 ) => Effect.Effect<Effect.Effect<void>>;
41}>("@t3tools/client-runtime/rpc/EnvironmentRpcRequestObserver", {
42 defaultValue: () => ({
43 observe: () => Effect.succeed(Effect.void),
44 }),
45}) {}
46
47export type EnvironmentRpcTag = keyof WsRpcProtocolClient & string;
48type RpcMethod<TTag extends EnvironmentRpcTag> = WsRpcProtocolClient[TTag];
49
50export type EnvironmentSubscriptionRpcTag =
51 | typeof WS_METHODS.uiSettingsSubscribe
52 | typeof WS_METHODS.subscribeScriptProgress
53 | typeof WS_METHODS.codexAuthCallbackSubscribe
54 | typeof WS_METHODS.providerAuthSubscribe
55 | typeof WS_METHODS.providerInstallSubscribe
56 | typeof ORCHESTRATION_V2_WS_METHODS.subscribeShell
57 | typeof ORCHESTRATION_V2_WS_METHODS.subscribeThread
58 | typeof ORCHESTRATION_V2_WS_METHODS.searchThreadStream
59 | typeof WS_METHODS.subscribeAuthAccess
60 | typeof WS_METHODS.subscribeServerConfig
61 | typeof WS_METHODS.subscribeServerLifecycle
62 | typeof WS_METHODS.scheduledTasksSubscribe
63 | typeof WS_METHODS.serverGetStorageCleanupReport
64 | typeof WS_METHODS.subscribeTerminalEvents
65 | typeof WS_METHODS.subscribeTerminalMetadata
66 | typeof WS_METHODS.subscribePreviewEvents
67 | typeof WS_METHODS.subscribeDiscoveredLocalServers
68 | typeof WS_METHODS.subscribeDeviceState
69 | typeof WS_METHODS.subscribeResourceTelemetry
70 | typeof WS_METHODS.pullRequestsSubscribeRefreshes
71 | typeof WS_METHODS.mergebotSubscribe
72 | typeof WS_METHODS.mergebotSubscribePullRequestStatus
73 | typeof WS_METHODS.mergebotSubscribeDeployments
74 | typeof WS_METHODS.jjSubscribeStack
75 | typeof WS_METHODS.subscribeVcsStatus
76 | typeof WS_METHODS.subscribeWorktreeSetup
77 | typeof WS_METHODS.subscribeProjectClones
78 | typeof WS_METHODS.projectsWatch
79 | typeof WS_METHODS.terminalAttach
80 | typeof WS_METHODS.terminalObserve;
81
82export type EnvironmentStreamCommandRpcTag =
83 | typeof WS_METHODS.chatGptHandoffSubscribe
84 | typeof WS_METHODS.cloudInstallRelayClient
85 | typeof WS_METHODS.serverUpdateServerWithProgress
86 | typeof WS_METHODS.gitRunStackedAction;
87
88export type EnvironmentStreamRpcTag =
89 | EnvironmentSubscriptionRpcTag
90 | EnvironmentStreamCommandRpcTag;
91
92export type EnvironmentUnaryRpcTag = Exclude<EnvironmentRpcTag, EnvironmentStreamRpcTag>;
93
94export interface EnvironmentRpcSubscriptionObservation {
95 readonly environmentId: string;
96 readonly method: EnvironmentSubscriptionRpcTag;
97 readonly input: unknown;
98}
99
100export class EnvironmentRpcSubscriptionObserver extends Context.Reference<{
101 readonly observe: (
102 subscription: EnvironmentRpcSubscriptionObservation,
103 ) => Effect.Effect<Effect.Effect<void>>;
104}>("@t3tools/client-runtime/rpc/EnvironmentRpcSubscriptionObserver", {
105 defaultValue: () => ({
106 observe: () => Effect.succeed(Effect.void),
107 }),
108}) {}
109
110export const isRpcClientError = Schema.is(RpcClientError.RpcClientError);
111const isEnvironmentAuthorizationError = Schema.is(EnvironmentAuthorizationError);
112
113/** Ceiling for the doubling delay between same-session expected-failure retries. */
114const MAX_EXPECTED_FAILURE_RETRY_DELAY_MS = 30_000;
115
116export type EnvironmentRpcInput<TTag extends EnvironmentRpcTag> = Parameters<RpcMethod<TTag>>[0];
117
118export type EnvironmentRpcSuccess<TTag extends EnvironmentUnaryRpcTag> =
119 RpcMethod<TTag> extends (input: any, options?: any) => Effect.Effect<infer A, any, any>
120 ? A
121 : never;
122
123export type EnvironmentRpcFailure<TTag extends EnvironmentUnaryRpcTag> =
124 RpcMethod<TTag> extends (input: any, options?: any) => Effect.Effect<any, infer E, any>
125 ? E
126 : never;
127
128export type EnvironmentRpcStreamValue<TTag extends EnvironmentStreamRpcTag> =
129 RpcMethod<TTag> extends (input: any, options?: any) => Stream.Stream<infer A, any, any>
130 ? A
131 : never;
132
133export type EnvironmentRpcStreamFailure<TTag extends EnvironmentStreamRpcTag> =
134 RpcMethod<TTag> extends (input: any, options?: any) => Stream.Stream<any, infer E, any>
135 ? E
136 : never;
137
138const currentSession = Effect.fn("EnvironmentRpc.currentSession")(function* () {
139 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
140 return yield* SubscriptionRef.get(supervisor.session).pipe(
141 Effect.flatMap(
142 Option.match({
143 onNone: () =>
144 Effect.fail(
145 new EnvironmentRpcUnavailableError({
146 environmentId: supervisor.target.environmentId,
147 message: `${supervisor.target.label} is not connected.`,
148 }),
149 ),
150 onSome: Effect.succeed,
151 }),
152 ),
153 );
154});
155
156export const getInitialServerConfig = Effect.fn("EnvironmentRpc.getInitialServerConfig")(
157 function* () {
158 const session = yield* currentSession();
159 return yield* session.initialConfig;
160 },
161);
162
163/** Installed by the shared command boundary. A missing boundary denies protected writes. */
164export class RpcPermissionGuard extends Context.Reference<{
165 readonly authorize: (
166 environmentId: EnvironmentId,
167 method: string,
168 input: unknown,
169 ) => Effect.Effect<void, EnvironmentAuthorizationError>;
170}>("@t3tools/client-runtime/rpc/RpcPermissionGuard", {
171 defaultValue: () => ({
172 authorize: (_environmentId, method, input) => {
173 const scope = clientRpcRequiredScopes(method, input)[0];
174 return scope === undefined
175 ? Effect.void
176 : Effect.fail(
177 new EnvironmentAuthorizationError({
178 ...authScopeRequiredResponse(scope),
179 message: `This connection requires ${scope}.`,
180 }),
181 );
182 },
183 }),
184}) {}
185
186const authorizeRequest = Effect.fn("EnvironmentRpc.authorize")(function* (
187 method: string,
188 input: unknown,
189) {
190 if (clientRpcRequiredScopes(method, input).length === 0) return;
191 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
192 const guard = yield* RpcPermissionGuard;
193 yield* guard.authorize(supervisor.target.environmentId, method, input);
194});
195
196export const requestGuarded = Effect.fn("EnvironmentRpc.request")(function* <
197 TTag extends EnvironmentUnaryRpcTag,
198>(tag: TTag, input: EnvironmentRpcInput<TTag>) {
199 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
200 yield* Effect.annotateCurrentSpan({
201 "environment.id": supervisor.target.environmentId,
202 "rpc.method": tag,
203 });
204 const session = yield* currentSession();
205 yield* authorizeRequest(tag, input);
206 const observer = yield* EnvironmentRpcRequestObserver;
207 const method = session.client[tag] as (
208 input: EnvironmentRpcInput<TTag>,
209 ) => Effect.Effect<EnvironmentRpcSuccess<TTag>, EnvironmentRpcFailure<TTag>>;
210 const completeObservation = yield* observer.observe({
211 environmentId: supervisor.target.environmentId,
212 method: tag,
213 });
214 return yield* method(input).pipe(Effect.ensuring(completeObservation));
215});
216
217export function runStreamGuarded<TTag extends EnvironmentStreamCommandRpcTag>(
218 tag: TTag,
219 input: EnvironmentRpcInput<TTag>,
220): Stream.Stream<
221 EnvironmentRpcStreamValue<TTag>,
222 | EnvironmentRpcStreamFailure<TTag>
223 | EnvironmentRpcUnavailableError
224 | EnvironmentAuthorizationError,
225 EnvironmentSupervisor.EnvironmentSupervisor
226> {
227 return Stream.unwrap(
228 authorizeRequest(tag, input).pipe(
229 Effect.andThen(currentSession()),
230 Effect.map((session) => {
231 const method = session.client[tag] as (
232 input: EnvironmentRpcInput<TTag>,
233 ) => Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>;
234 return method(input);
235 }),
236 ),
237 ).pipe(
238 Stream.withSpan("EnvironmentRpc.runStream", {
239 attributes: { "rpc.method": tag },
240 }),
241 );
242}
243
244interface SubscriptionOptions<TTag extends EnvironmentSubscriptionRpcTag> {
245 /** Reports protocol or programming defects without changing their recovery policy. */
246 readonly onDefect?: (
247 cause: Cause.Cause<EnvironmentRpcStreamFailure<TTag>>,
248 ) => Effect.Effect<void, never, never>;
249 readonly onExpectedFailure?: (
250 cause: Cause.Cause<EnvironmentRpcStreamFailure<TTag>>,
251 ) => Effect.Effect<void, never, never>;
252 /**
253 * First delay before resubscribing on the same session after an expected
254 * failure. Each consecutive failure doubles it up to 30 seconds, and the
255 * first value from a healthy stream resets it. Authorization failures are
256 * not retried; they wait for the next session or `resubscribe` signal.
257 */
258 readonly retryExpectedFailureAfter?: Duration.Input;
259 readonly resubscribe?: Stream.Stream<unknown, never, never>;
260}
261
262function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
263 tag: TTag,
264 makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
265 mapStream: (
266 session: RpcSession,
267 stream: Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>,
268 ) => Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>>,
269 options?: SubscriptionOptions<TTag>,
270): Stream.Stream<
271 A,
272 EnvironmentRpcStreamFailure<TTag>,
273 EnvironmentSupervisor.EnvironmentSupervisor
274> {
275 return Stream.unwrap(
276 Effect.gen(function* () {
277 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
278 const observer = yield* EnvironmentRpcSubscriptionObserver;
279 const sessionChanges = SubscriptionRef.changes(supervisor.session);
280 const sessions =
281 options?.resubscribe === undefined
282 ? sessionChanges
283 : Stream.merge(
284 sessionChanges,
285 options.resubscribe.pipe(
286 Stream.mapEffect(() => SubscriptionRef.get(supervisor.session)),
287 ),
288 );
289 return sessions.pipe(
290 Stream.switchMap(
291 Option.match({
292 onNone: () => Stream.empty,
293 onSome: (session) => {
294 const method = (
295 tag === WS_METHODS.subscribeServerConfig
296 ? session.subscribeServerConfig
297 : session.client[tag]
298 ) as (
299 input: EnvironmentRpcInput<TTag>,
300 ) => Stream.Stream<
301 EnvironmentRpcStreamValue<TTag>,
302 EnvironmentRpcStreamFailure<TTag>
303 >;
304 // Consecutive expected-failure retries on this session. Reset by
305 // the first value a resubscribed stream delivers.
306 let expectedFailureRetries = 0;
307 const subscribeToSession = (): Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>> =>
308 Stream.suspend(() =>
309 Stream.unwrap(
310 Effect.gen(function* () {
311 const input = yield* makeInput(session);
312 const completeObservation = yield* observer.observe({
313 environmentId: supervisor.target.environmentId,
314 method: tag,
315 input,
316 });
317 const stream = mapStream(session, method(input)).pipe(
318 Stream.onFirst(() =>
319 Effect.sync(() => {
320 expectedFailureRetries = 0;
321 }),
322 ),
323 );
324 return stream.pipe(Stream.ensuring(completeObservation));
325 }),
326 ).pipe(
327 Stream.tapCause((cause) =>
328 options?.onDefect !== undefined &&
329 cause.reasons.some(
330 (reason) =>
331 reason._tag === "Die" ||
332 (reason._tag === "Fail" &&
333 isRpcClientError(reason.error) &&
334 reason.error.reason._tag === "RpcClientDefect"),
335 )
336 ? options.onDefect(cause)
337 : Effect.void,
338 ),
339 Stream.catchCause((cause) => {
340 const hasOnlyExpectedFailures =
341 cause.reasons.length > 0 &&
342 cause.reasons.every((reason) => reason._tag === "Fail");
343 const isTransportFailure =
344 hasOnlyExpectedFailures &&
345 cause.reasons.every(
346 (reason) => reason._tag === "Fail" && isRpcClientError(reason.error),
347 );
348 if (isTransportFailure) {
349 return Stream.fromEffect(
350 Effect.logWarning(
351 "Durable RPC subscription lost its transport; waiting for the next session.",
352 {
353 cause: Cause.pretty(cause),
354 method: tag,
355 environmentId: supervisor.target.environmentId,
356 },
357 ),
358 ).pipe(Stream.drain);
359 }
360 if (hasOnlyExpectedFailures && options?.onExpectedFailure !== undefined) {
361 const handled = Stream.fromEffect(options.onExpectedFailure(cause)).pipe(
362 Stream.drain,
363 );
364 const isAuthorizationFailure = cause.reasons.some(
365 (reason) =>
366 reason._tag === "Fail" && isEnvironmentAuthorizationError(reason.error),
367 );
368 if (
369 options.retryExpectedFailureAfter === undefined ||
370 isAuthorizationFailure
371 ) {
372 return handled;
373 }
374 const retryDelay = Duration.millis(
375 Math.min(
376 Duration.toMillis(options.retryExpectedFailureAfter) *
377 2 ** expectedFailureRetries,
378 MAX_EXPECTED_FAILURE_RETRY_DELAY_MS,
379 ),
380 );
381 expectedFailureRetries += 1;
382 return handled.pipe(
383 Stream.concat(
384 Stream.fromEffect(Effect.sleep(retryDelay)).pipe(Stream.drain),
385 ),
386 Stream.concat(subscribeToSession()),
387 );
388 }
389 return Stream.failCause(cause);
390 }),
391 ),
392 );
393 return subscribeToSession();
394 },
395 }),
396 ),
397 );
398 }),
399 ).pipe(
400 Stream.withSpan("EnvironmentRpc.subscribe", {
401 attributes: { "rpc.method": tag },
402 }),
403 );
404}
405
406export function subscribeDynamic<TTag extends EnvironmentSubscriptionRpcTag>(
407 tag: TTag,
408 makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
409 options?: SubscriptionOptions<TTag>,
410): Stream.Stream<
411 EnvironmentRpcStreamValue<TTag>,
412 EnvironmentRpcStreamFailure<TTag>,
413 EnvironmentSupervisor.EnvironmentSupervisor
414> {
415 return subscribeDynamicMapped(tag, makeInput, (_session, stream) => stream, options);
416}
417
418/** Tags each value before `switchMap` can buffer it across a session change. */
419export function subscribeDynamicWithSession<TTag extends EnvironmentSubscriptionRpcTag>(
420 tag: TTag,
421 makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
422 options?: SubscriptionOptions<TTag>,
423): Stream.Stream<
424 readonly [session: RpcSession, value: EnvironmentRpcStreamValue<TTag>],
425 EnvironmentRpcStreamFailure<TTag>,
426 EnvironmentSupervisor.EnvironmentSupervisor
427> {
428 return subscribeDynamicMapped(
429 tag,
430 makeInput,
431 (session, stream) => stream.pipe(Stream.map((value) => [session, value] as const)),
432 options,
433 );
434}
435
436export function subscribe<TTag extends EnvironmentSubscriptionRpcTag>(
437 tag: TTag,
438 input: EnvironmentRpcInput<TTag>,
439 options?: SubscriptionOptions<TTag>,
440): Stream.Stream<
441 EnvironmentRpcStreamValue<TTag>,
442 EnvironmentRpcStreamFailure<TTag>,
443 EnvironmentSupervisor.EnvironmentSupervisor
444> {
445 return subscribeDynamic(tag, () => Effect.succeed(input), options);
446}
447
448/** Protected writes must go through the permission-aware command layer. */
449export const request = <TTag extends Exclude<EnvironmentUnaryRpcTag, ClientGuardedRpcTag>>(
450 tag: TTag,
451 input: EnvironmentRpcInput<TTag>,
452) => requestGuarded(tag, input);
453
454export const runStream = <
455 TTag extends Exclude<EnvironmentStreamCommandRpcTag, ClientGuardedRpcTag>,
456>(
457 tag: TTag,
458 input: EnvironmentRpcInput<TTag>,
459) => runStreamGuarded(tag, input);