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