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 environmentLabel: string;
35 readonly method: string;
36}
37
38export class EnvironmentRpcRequestObserver extends Context.Reference<{
39 readonly observe: (
40 request: EnvironmentRpcRequestObservation,
41 ) => Effect.Effect<Effect.Effect<void>>;
42}>("@t3tools/client-runtime/rpc/EnvironmentRpcRequestObserver", {
43 defaultValue: () => ({
44 observe: () => Effect.succeed(Effect.void),
45 }),
46}) {}
47
48export type EnvironmentRpcTag = keyof WsRpcProtocolClient & string;
49type RpcMethod<TTag extends EnvironmentRpcTag> = WsRpcProtocolClient[TTag];
50
51export type EnvironmentSubscriptionRpcTag =
52 | typeof WS_METHODS.providerAuthSubscribe
53 | typeof WS_METHODS.providerInstallSubscribe
54 | typeof ORCHESTRATION_V2_WS_METHODS.subscribeShell
55 | typeof ORCHESTRATION_V2_WS_METHODS.subscribeThread
56 | typeof ORCHESTRATION_V2_WS_METHODS.subscribeCommandOutput
57 | typeof ORCHESTRATION_V2_WS_METHODS.subscribePromptSuggestion
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.serverGetDatabaseStorage
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.subscribeVcsStatus
72 | typeof WS_METHODS.subscribeVcsBranch
73 | typeof WS_METHODS.subscribeWorktreeSetup
74 | typeof WS_METHODS.subscribeProjectClones
75 | typeof WS_METHODS.terminalAttach
76 | typeof WS_METHODS.terminalObserve;
77
78export type EnvironmentStreamCommandRpcTag =
79 | typeof WS_METHODS.chatGptHandoffSubscribe
80 | typeof WS_METHODS.cloudInstallRelayClient
81 | typeof WS_METHODS.serverUpdateServerWithProgress
82 | typeof WS_METHODS.gitRunStackedAction;
83
84export type EnvironmentStreamRpcTag =
85 | EnvironmentSubscriptionRpcTag
86 | EnvironmentStreamCommandRpcTag;
87
88export type EnvironmentUnaryRpcTag = Exclude<EnvironmentRpcTag, EnvironmentStreamRpcTag>;
89
90export interface EnvironmentRpcSubscriptionObservation {
91 readonly environmentId: string;
92 readonly method: EnvironmentSubscriptionRpcTag;
93 readonly input: unknown;
94}
95
96export class EnvironmentRpcSubscriptionObserver extends Context.Reference<{
97 readonly observe: (
98 subscription: EnvironmentRpcSubscriptionObservation,
99 ) => Effect.Effect<Effect.Effect<void>>;
100}>("@t3tools/client-runtime/rpc/EnvironmentRpcSubscriptionObserver", {
101 defaultValue: () => ({
102 observe: () => Effect.succeed(Effect.void),
103 }),
104}) {}
105
106export const isRpcClientError = Schema.is(RpcClientError.RpcClientError);
107const isEnvironmentAuthorizationError = Schema.is(EnvironmentAuthorizationError);
108
109/** A subscription never retries an authorization failure on its own. */
110export const isAuthorizationFailure = (cause: Cause.Cause<unknown>) =>
111 cause.reasons.some(
112 (reason) => reason._tag === "Fail" && isEnvironmentAuthorizationError(reason.error),
113 );
114
115const MAX_EXPECTED_FAILURE_RETRY_DELAY_MS = 30_000;
116
117export const subscriptionRetryDelay = (base: Duration.Input, retries: number) =>
118 Duration.millis(
119 Math.min(Duration.toMillis(base) * 2 ** retries, MAX_EXPECTED_FAILURE_RETRY_DELAY_MS),
120 );
121
122export type EnvironmentRpcInput<TTag extends EnvironmentRpcTag> = Parameters<RpcMethod<TTag>>[0];
123
124export type EnvironmentRpcSuccess<TTag extends EnvironmentUnaryRpcTag> =
125 RpcMethod<TTag> extends (input: any, options?: any) => Effect.Effect<infer A, any, any>
126 ? A
127 : never;
128
129export type EnvironmentRpcFailure<TTag extends EnvironmentUnaryRpcTag> =
130 RpcMethod<TTag> extends (input: any, options?: any) => Effect.Effect<any, infer E, any>
131 ? E
132 : never;
133
134export type EnvironmentRpcStreamValue<TTag extends EnvironmentStreamRpcTag> =
135 RpcMethod<TTag> extends (input: any, options?: any) => Stream.Stream<infer A, any, any>
136 ? A
137 : never;
138
139export type EnvironmentRpcStreamFailure<TTag extends EnvironmentStreamRpcTag> =
140 RpcMethod<TTag> extends (input: any, options?: any) => Stream.Stream<any, infer E, any>
141 ? E
142 : never;
143
144export const currentSession = Effect.fn("EnvironmentRpc.currentSession")(function* () {
145 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
146 return yield* SubscriptionRef.get(supervisor.session).pipe(
147 Effect.flatMap(
148 Option.match({
149 onNone: () =>
150 Effect.fail(
151 new EnvironmentRpcUnavailableError({
152 environmentId: supervisor.target.environmentId,
153 message: `${supervisor.target.label} is not connected.`,
154 }),
155 ),
156 onSome: Effect.succeed,
157 }),
158 ),
159 );
160});
161
162export const getInitialServerConfig = Effect.fn("EnvironmentRpc.getInitialServerConfig")(
163 function* () {
164 const session = yield* currentSession();
165 return yield* session.initialConfig;
166 },
167);
168
169/** Installed by the shared command boundary. A missing boundary denies protected writes. */
170export class RpcPermissionGuard extends Context.Reference<{
171 readonly authorize: (
172 environmentId: EnvironmentId,
173 method: string,
174 input: unknown,
175 ) => Effect.Effect<void, EnvironmentAuthorizationError>;
176}>("@t3tools/client-runtime/rpc/RpcPermissionGuard", {
177 defaultValue: () => ({
178 authorize: (_environmentId, method, input) => {
179 const scope = clientRpcRequiredScopes(method, input)[0];
180 return scope === undefined
181 ? Effect.void
182 : Effect.fail(
183 new EnvironmentAuthorizationError({
184 ...authScopeRequiredResponse(scope),
185 message: `This connection requires ${scope}.`,
186 }),
187 );
188 },
189 }),
190}) {}
191
192const authorizeRequest = Effect.fn("EnvironmentRpc.authorize")(function* (
193 method: string,
194 input: unknown,
195) {
196 if (clientRpcRequiredScopes(method, input).length === 0) return;
197 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
198 const guard = yield* RpcPermissionGuard;
199 yield* guard.authorize(supervisor.target.environmentId, method, input);
200});
201
202export const requestInSession = Effect.fn("EnvironmentRpc.request")(function* <
203 TTag extends EnvironmentUnaryRpcTag,
204>(session: RpcSession, tag: TTag, input: EnvironmentRpcInput<TTag>) {
205 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
206 yield* Effect.annotateCurrentSpan({
207 "environment.id": supervisor.target.environmentId,
208 "rpc.method": tag,
209 });
210 yield* authorizeRequest(tag, input);
211 const observer = yield* EnvironmentRpcRequestObserver;
212 const method = session.client[tag] as (
213 input: EnvironmentRpcInput<TTag>,
214 ) => Effect.Effect<EnvironmentRpcSuccess<TTag>, EnvironmentRpcFailure<TTag>>;
215 const completeObservation = yield* observer.observe({
216 environmentId: supervisor.target.environmentId,
217 environmentLabel: supervisor.target.label,
218 method: tag,
219 });
220 return yield* method(input).pipe(Effect.ensuring(completeObservation));
221});
222
223// The session-bound request owns the existing EnvironmentRpc.request span.
224export const requestGuarded = Effect.fnUntraced(function* <TTag extends EnvironmentUnaryRpcTag>(
225 tag: TTag,
226 input: EnvironmentRpcInput<TTag>,
227) {
228 return yield* requestInSession(yield* currentSession(), tag, input);
229});
230
231export function runStreamGuarded<TTag extends EnvironmentStreamCommandRpcTag>(
232 tag: TTag,
233 input: EnvironmentRpcInput<TTag>,
234): Stream.Stream<
235 EnvironmentRpcStreamValue<TTag>,
236 | EnvironmentRpcStreamFailure<TTag>
237 | EnvironmentRpcUnavailableError
238 | EnvironmentAuthorizationError,
239 EnvironmentSupervisor.EnvironmentSupervisor
240> {
241 return Stream.unwrap(
242 currentSession().pipe(Effect.map((session) => runStreamInSession(session, tag, input))),
243 );
244}
245
246export function runStreamInSession<TTag extends EnvironmentStreamCommandRpcTag>(
247 session: RpcSession,
248 tag: TTag,
249 input: EnvironmentRpcInput<TTag>,
250): Stream.Stream<
251 EnvironmentRpcStreamValue<TTag>,
252 EnvironmentRpcStreamFailure<TTag> | EnvironmentAuthorizationError,
253 EnvironmentSupervisor.EnvironmentSupervisor
254> {
255 const method = session.client[tag] as (
256 input: EnvironmentRpcInput<TTag>,
257 ) => Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>;
258 return Stream.unwrap(authorizeRequest(tag, input).pipe(Effect.map(() => method(input)))).pipe(
259 Stream.withSpan("EnvironmentRpc.runStream", {
260 attributes: { "rpc.method": tag },
261 }),
262 );
263}
264
265interface SubscriptionOptions<TTag extends EnvironmentSubscriptionRpcTag> {
266 /** Reports defects; the subscription then waits for the next session or resubscribe. */
267 readonly onDefect?: (
268 cause: Cause.Cause<EnvironmentRpcStreamFailure<TTag>>,
269 ) => Effect.Effect<void, never, never>;
270 readonly onExpectedFailure?: (
271 cause: Cause.Cause<EnvironmentRpcStreamFailure<TTag>>,
272 ) => Effect.Effect<void, never, never>;
273 readonly retryExpectedFailureAfter?: Duration.Input;
274 readonly resubscribe?: Stream.Stream<unknown, never, never>;
275}
276
277function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
278 tag: TTag,
279 makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
280 mapStream: (
281 session: RpcSession,
282 stream: Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>,
283 ) => Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>>,
284 options?: SubscriptionOptions<TTag>,
285): Stream.Stream<
286 A,
287 EnvironmentRpcStreamFailure<TTag>,
288 EnvironmentSupervisor.EnvironmentSupervisor
289> {
290 return Stream.unwrap(
291 Effect.gen(function* () {
292 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
293 const observer = yield* EnvironmentRpcSubscriptionObserver;
294 const sessionChanges = SubscriptionRef.changes(supervisor.session);
295 const sessions =
296 options?.resubscribe === undefined
297 ? sessionChanges
298 : Stream.merge(
299 sessionChanges,
300 options.resubscribe.pipe(
301 Stream.mapEffect(() => SubscriptionRef.get(supervisor.session)),
302 ),
303 );
304 return sessions.pipe(
305 Stream.switchMap(
306 Option.match({
307 onNone: () => Stream.empty,
308 onSome: (session) => {
309 const method = (
310 tag === WS_METHODS.subscribeServerConfig
311 ? session.subscribeServerConfig
312 : session.client[tag]
313 ) as (
314 input: EnvironmentRpcInput<TTag>,
315 ) => Stream.Stream<
316 EnvironmentRpcStreamValue<TTag>,
317 EnvironmentRpcStreamFailure<TTag>
318 >;
319 let expectedFailureRetries = 0;
320 const subscribeToSession = (): Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>> =>
321 Stream.suspend(() =>
322 Stream.unwrap(
323 Effect.gen(function* () {
324 const input = yield* makeInput(session);
325 const completeObservation = yield* observer.observe({
326 environmentId: supervisor.target.environmentId,
327 method: tag,
328 input,
329 });
330 const stream = mapStream(
331 session,
332 Stream.suspend(() => method(input)),
333 ).pipe(
334 Stream.onFirst(() =>
335 Effect.sync(() => {
336 expectedFailureRetries = 0;
337 }),
338 ),
339 );
340 return stream.pipe(Stream.ensuring(completeObservation));
341 }),
342 ).pipe(
343 Stream.tapCause((cause) =>
344 options?.onDefect !== undefined &&
345 cause.reasons.some(
346 (reason) =>
347 reason._tag === "Die" ||
348 (reason._tag === "Fail" &&
349 isRpcClientError(reason.error) &&
350 reason.error.reason._tag === "RpcClientDefect"),
351 )
352 ? options.onDefect(cause)
353 : Effect.void,
354 ),
355 Stream.catchCause((cause) => {
356 const hasOnlyExpectedFailures =
357 cause.reasons.length > 0 &&
358 cause.reasons.every((reason) => reason._tag === "Fail");
359 const isTransportFailure =
360 hasOnlyExpectedFailures &&
361 cause.reasons.every(
362 (reason) => reason._tag === "Fail" && isRpcClientError(reason.error),
363 );
364 if (isTransportFailure) {
365 return Stream.fromEffect(
366 Effect.logWarning(
367 "Durable RPC subscription lost its transport; waiting for the next session.",
368 {
369 cause: Cause.pretty(cause),
370 method: tag,
371 environmentId: supervisor.target.environmentId,
372 },
373 ),
374 ).pipe(Stream.drain);
375 }
376 if (hasOnlyExpectedFailures && options?.onExpectedFailure !== undefined) {
377 const handled = Stream.fromEffect(options.onExpectedFailure(cause)).pipe(
378 Stream.drain,
379 );
380 if (
381 options.retryExpectedFailureAfter === undefined ||
382 isAuthorizationFailure(cause)
383 ) {
384 return handled;
385 }
386 const retryDelay = subscriptionRetryDelay(
387 options.retryExpectedFailureAfter,
388 expectedFailureRetries,
389 );
390 expectedFailureRetries += 1;
391 return handled.pipe(
392 Stream.concat(
393 Stream.fromEffect(Effect.sleep(retryDelay)).pipe(Stream.drain),
394 ),
395 Stream.concat(subscribeToSession()),
396 );
397 }
398 // A reported defect parks until the next session or resubscribe. Failing would
399 // end the whole subscription, so a mounted screen never syncs again.
400 if (options?.onDefect !== undefined && Cause.hasDies(cause)) {
401 return Stream.empty;
402 }
403 return Stream.failCause(cause);
404 }),
405 ),
406 );
407 return subscribeToSession();
408 },
409 }),
410 ),
411 );
412 }),
413 ).pipe(
414 Stream.withSpan("EnvironmentRpc.subscribe", {
415 attributes: { "rpc.method": tag },
416 }),
417 );
418}
419
420export function subscribeDynamic<TTag extends EnvironmentSubscriptionRpcTag>(
421 tag: TTag,
422 makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
423 options?: SubscriptionOptions<TTag>,
424): Stream.Stream<
425 EnvironmentRpcStreamValue<TTag>,
426 EnvironmentRpcStreamFailure<TTag>,
427 EnvironmentSupervisor.EnvironmentSupervisor
428> {
429 return subscribeDynamicMapped(tag, makeInput, (_session, stream) => stream, options);
430}
431
432/** Tags each value before `switchMap` can buffer it across a session change. */
433export function subscribeDynamicWithSession<TTag extends EnvironmentSubscriptionRpcTag>(
434 tag: TTag,
435 makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
436 options?: SubscriptionOptions<TTag>,
437): Stream.Stream<
438 readonly [session: RpcSession, value: EnvironmentRpcStreamValue<TTag>],
439 EnvironmentRpcStreamFailure<TTag>,
440 EnvironmentSupervisor.EnvironmentSupervisor
441> {
442 return subscribeDynamicMapped(
443 tag,
444 makeInput,
445 (session, stream) => stream.pipe(Stream.map((value) => [session, value] as const)),
446 options,
447 );
448}
449
450export function subscribe<TTag extends EnvironmentSubscriptionRpcTag>(
451 tag: TTag,
452 input: EnvironmentRpcInput<TTag>,
453 options?: SubscriptionOptions<TTag>,
454): Stream.Stream<
455 EnvironmentRpcStreamValue<TTag>,
456 EnvironmentRpcStreamFailure<TTag>,
457 EnvironmentSupervisor.EnvironmentSupervisor
458> {
459 return subscribeDynamic(tag, () => Effect.succeed(input), options);
460}
461
462export const subscribeVcsBranch = (
463 input: EnvironmentRpcInput<typeof WS_METHODS.subscribeVcsBranch>,
464) =>
465 subscribeDynamicMapped(
466 WS_METHODS.subscribeVcsBranch,
467 () => Effect.succeed(input),
468 (session, stream) =>
469 Stream.unwrap(
470 session.initialConfig.pipe(
471 Effect.mapError(
472 (error) =>
473 new RpcClientError.RpcClientError({
474 reason: new RpcClientError.RpcClientDefect({
475 message: error.message,
476 cause: error,
477 }),
478 }),
479 ),
480 Effect.map((config) =>
481 config.environment.capabilities.vcsBranchStatus === true
482 ? stream
483 : session.client[WS_METHODS.subscribeVcsStatus]({
484 ...input,
485 includeRemote: false,
486 }).pipe(
487 Stream.filter((event) => event._tag !== "remoteUpdated"),
488 Stream.map(({ local }) => ({ refName: local.refName, isRepo: local.isRepo })),
489 ),
490 ),
491 ),
492 ),
493 );
494
495export const refreshVcsBranch = Effect.fn("EnvironmentRpc.refreshVcsBranch")(function* (
496 input: EnvironmentRpcInput<typeof WS_METHODS.vcsRefreshBranch>,
497) {
498 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
499 const session = yield* currentSession();
500 const config = yield* session.initialConfig.pipe(
501 Effect.mapError(
502 (error) =>
503 new EnvironmentRpcUnavailableError({
504 environmentId: supervisor.target.environmentId,
505 message: error.message,
506 }),
507 ),
508 );
509 if (config.environment.capabilities.vcsBranchStatus === true) {
510 return yield* requestInSession(session, WS_METHODS.vcsRefreshBranch, input);
511 }
512 // Older hosts retain their existing status stream without adding per-row wake scans.
513});
514
515export const request = <TTag extends Exclude<EnvironmentUnaryRpcTag, ClientGuardedRpcTag>>(
516 tag: TTag,
517 input: EnvironmentRpcInput<TTag>,
518) => requestGuarded(tag, input);
519
520export const runStream = <
521 TTag extends Exclude<EnvironmentStreamCommandRpcTag, ClientGuardedRpcTag>,
522>(
523 tag: TTag,
524 input: EnvironmentRpcInput<TTag>,
525) => runStreamGuarded(tag, input);