1import {
2 type ServerConfig,
3 type ServerConfigStreamEvent,
4 WsSubscribeServerConfigRpc,
5 WS_METHODS,
6} from "@t3tools/contracts";
7import {
8 joinSocketFrameParts,
9 SOCKET_FRAME_PARTS_QUERY_PARAM,
10} from "@t3tools/shared/socketFrameParts";
11import * as Cause from "effect/Cause";
12import * as Context from "effect/Context";
13import * as Deferred from "effect/Deferred";
14import * as Duration from "effect/Duration";
15import * as Effect from "effect/Effect";
16import * as Equal from "effect/Equal";
17import * as Exit from "effect/Exit";
18import * as Layer from "effect/Layer";
19import * as Option from "effect/Option";
20import * as PubSub from "effect/PubSub";
21import * as Ref from "effect/Ref";
22import * as Schedule from "effect/Schedule";
23import * as Schema from "effect/Schema";
24import type * as Scope from "effect/Scope";
25import * as Stream from "effect/Stream";
26import type * as Rpc from "effect/rpc/Rpc";
27import * as RpcClient from "effect/rpc/RpcClient";
28import * as RpcClientError from "effect/rpc/RpcClientError";
29import * as RpcSerialization from "effect/rpc/RpcSerialization";
30import * as Socket from "effect/socket/Socket";
31
32import { makeWsRpcProtocolClient, PING_TIMEOUT, type WsRpcProtocolClient } from "./protocol.ts";
33import { NETWORK_BLOCKING_HINT, UNRESPONSIVE_HOST_HINT } from "../errors/network.ts";
34import type {
35 ConnectionAttemptError,
36 ConnectionTransientError,
37 PreparedConnection,
38} from "../connection/model.ts";
39import {
40 ConnectionBlockedError,
41 ConnectionTransientError as ConnectionTransientErrorClass,
42} from "../connection/model.ts";
43import {
44 applyServerConfigProjection,
45 type ServerConfigProjection,
46 withoutEnvironmentThemes,
47} from "../state/serverConfigProjection.ts";
48import { environmentMismatchError } from "../connection/errors.ts";
49
50const SOCKET_OPEN_TIMEOUT = "15 seconds";
51
52export interface RpcSession {
53 readonly client: WsRpcProtocolClient;
54 readonly initialConfig: Effect.Effect<ServerConfig, ConnectionAttemptError>;
55 readonly subscribeServerConfig: (
56 input?: ServerConfigSubscriptionInput,
57 ) => ServerConfigSubscription;
58 readonly ready: Effect.Effect<void, ConnectionAttemptError>;
59 readonly probe: Effect.Effect<void, ConnectionAttemptError>;
60 /** Frames received so far: proof the socket is alive that needs no reply. */
61 readonly inboundFrames: Effect.Effect<number>;
62 readonly closed: Effect.Effect<never, ConnectionAttemptError>;
63}
64
65export interface RpcSessionOptions {
66 readonly environmentThemes?: boolean;
67 readonly usageLimitSources?: boolean;
68 readonly usageLimitsCommand?: boolean;
69 readonly providerCatalogsOnDemand?: boolean;
70 readonly providerCatalogDigests?: boolean;
71 readonly providerStatusDeltas?: boolean;
72}
73
74export class RpcSessionFactory extends Context.Service<
75 RpcSessionFactory,
76 {
77 readonly connect: (
78 connection: PreparedConnection,
79 ) => Effect.Effect<RpcSession, ConnectionAttemptError, Scope.Scope>;
80 }
81>()("@t3tools/client-runtime/rpc/session/RpcSessionFactory") {}
82
83type InitialConfigError = Effect.Error<
84 ReturnType<WsRpcProtocolClient[typeof WS_METHODS.serverGetConfig]>
85>;
86type ProbeError = Effect.Error<ReturnType<WsRpcProtocolClient[typeof WS_METHODS.serverProbe]>>;
87type ServerConfigSubscriptionError =
88 | Rpc.ErrorExit<typeof WsSubscribeServerConfigRpc>
89 | RpcClientError.RpcClientError;
90type ServerConfigSubscription = Stream.Stream<
91 ServerConfigStreamEvent,
92 ServerConfigSubscriptionError
93>;
94type ServerConfigSubscriptionInput = Parameters<
95 WsRpcProtocolClient[typeof WS_METHODS.subscribeServerConfig]
96>[0];
97type EnvironmentThemesUpdatedEvent = Extract<
98 ServerConfigStreamEvent,
99 { readonly type: "environmentThemesUpdated" }
100>;
101type UsageLimitSourcesUpdatedEvent = Extract<
102 ServerConfigStreamEvent,
103 { readonly type: "usageLimitSourcesUpdated" }
104>;
105
106interface ServerConfigReplayState {
107 readonly projection: ServerConfigProjection;
108 readonly revision: number;
109 readonly themesEvent: EnvironmentThemesUpdatedEvent | undefined;
110 readonly sourcesEvent: UsageLimitSourcesUpdatedEvent | undefined;
111}
112
113interface BufferedServerConfigEvent {
114 readonly event: ServerConfigStreamEvent;
115 readonly replay: ServerConfigReplayState;
116 readonly revision: number;
117}
118
119function serverConfigReplayEvents(
120 state: ServerConfigReplayState,
121): ReadonlyArray<ServerConfigStreamEvent> {
122 const snapshot = {
123 version: 1 as const,
124 type: "snapshot" as const,
125 config: withoutEnvironmentThemes(state.projection.config),
126 };
127 return [
128 snapshot,
129 ...(state.themesEvent === undefined ? [] : [state.themesEvent]),
130 ...(state.sourcesEvent === undefined ? [] : [state.sourcesEvent]),
131 ];
132}
133
134const isSocketErrorReason = Schema.is(Socket.SocketErrorReason);
135const PING_TIMEOUT_SECONDS = Duration.toSeconds(PING_TIMEOUT);
136
137function mapSessionRpcError(
138 error: InitialConfigError | ProbeError | ServerConfigSubscriptionError,
139 networkHint: string,
140 via: "relay" | "direct",
141): ConnectionAttemptError {
142 switch (error._tag) {
143 case "EnvironmentAuthorizationError":
144 return new ConnectionBlockedError({
145 reason: "permission",
146 detail: error.message,
147 hop: "server",
148 via,
149 });
150 case "KeybindingsConfigParseError":
151 case "ServerSettingsError":
152 return new ConnectionTransientErrorClass({
153 reason: "remote-unavailable",
154 detail: error.message,
155 hop: "server",
156 via,
157 });
158 case "RpcClientError":
159 return new ConnectionTransientErrorClass({
160 reason: "transport",
161 detail: `${error.message}${isSocketErrorReason(error.reason) ? networkHint : ""}`,
162 ...socketFailureLocation(via),
163 });
164 }
165}
166
167/** Over T3 Connect the relay already reached the host, so a socket that drops failed between this device and the tunnel. */
168function socketFailureLocation(via: "relay" | "direct") {
169 return { hop: via === "relay" ? "tunnel" : "server", via } as const;
170}
171
172/** @public */
173export const make = Effect.fn("RpcSessionFactory.make")(function* (
174 options: RpcSessionOptions = {},
175) {
176 const webSocketConstructor = yield* Socket.WebSocketConstructor;
177 const serverConfigInput: ServerConfigSubscriptionInput = {
178 ...(options.environmentThemes === true ? { environmentThemes: true } : {}),
179 ...(options.usageLimitSources === true ? { usageLimitSources: true } : {}),
180 ...(options.usageLimitsCommand === true ? { usageLimitsCommand: true } : {}),
181 ...(options.providerCatalogsOnDemand === true ? { providerCatalogsOnDemand: true } : {}),
182 ...(options.providerCatalogDigests === true ? { providerCatalogDigests: true } : {}),
183 ...(options.providerStatusDeltas === true ? { providerStatusDeltas: true } : {}),
184 };
185 // Config state merges the snapshots it fetches back in, so no subscriber needs them broadcast.
186 const onDemandSnapshots = (input: ServerConfigSubscriptionInput) => ({
187 ...input,
188 workspaceSnapshotsOnDemand: true,
189 });
190
191 const connect = Effect.fnUntraced(function* (connection: PreparedConnection) {
192 const via = connection.target._tag === "RelayConnectionTarget" ? "relay" : "direct";
193 const networkHint = via === "relay" ? ` ${NETWORK_BLOCKING_HINT}` : "";
194 const mapRpcError = (error: Parameters<typeof mapSessionRpcError>[0]) =>
195 mapSessionRpcError(error, networkHint, via);
196 yield* Effect.annotateCurrentSpan({
197 "connection.environment.id": connection.environmentId,
198 });
199
200 const connected = yield* Deferred.make<void>();
201 const disconnected = yield* Deferred.make<never, ConnectionTransientError>();
202 const pingTimedOut = yield* Ref.make(false);
203 const hooks = RpcClient.ConnectionHooks.of({
204 onConnect: Deferred.succeed(connected, undefined).pipe(Effect.asVoid),
205 onDisconnect: Effect.all([Deferred.isDone(connected), Ref.get(pingTimedOut)]).pipe(
206 Effect.flatMap(([wasConnected, timedOut]) =>
207 Deferred.fail(
208 disconnected,
209 new ConnectionTransientErrorClass({
210 reason: "transport",
211 // The network hint only fits a socket that never opened. Once it
212 // opened, the path worked, and a host that goes quiet is more
213 // often overloaded or asleep than blocked.
214 detail: !wasConnected
215 ? `${connection.label} could not establish a WebSocket connection.${networkHint}`
216 : timedOut
217 ? `${connection.label} stopped responding for ${PING_TIMEOUT_SECONDS}s. ${UNRESPONSIVE_HOST_HINT}`
218 : `${connection.label} disconnected.`,
219 ...socketFailureLocation(via),
220 }),
221 ),
222 ),
223 Effect.asVoid,
224 ),
225 });
226 // The protocol never reopens its socket, so every frame counted here came
227 // from this session's one socket.
228 let inboundFrames = 0;
229 const countingWebSocketConstructor: typeof webSocketConstructor = (url, options) => {
230 const socket = webSocketConstructor(url, options);
231 socket.addEventListener("message", () => {
232 inboundFrames += 1;
233 });
234 return socket;
235 };
236 const socketUrl = new URL(connection.socketUrl);
237 socketUrl.searchParams.set(SOCKET_FRAME_PARTS_QUERY_PARAM, "1");
238 const layerSocket = Layer.effect(
239 Socket.Socket,
240 Effect.map(Effect.service(Socket.Socket), joinSocketFrameParts),
241 ).pipe(
242 Layer.provide(
243 Socket.layerWebSocket(socketUrl.toString(), {
244 openTimeout: SOCKET_OPEN_TIMEOUT,
245 }),
246 ),
247 Layer.provide(Layer.succeed(Socket.WebSocketConstructor, countingWebSocketConstructor)),
248 );
249 const layerProtocol = Layer.effect(
250 RpcClient.Protocol,
251 RpcClient.makeProtocolSocket({
252 pingTimeout: PING_TIMEOUT,
253 retryTransientErrors: false,
254 retryPolicy: Schedule.recurs(0),
255 onPingTimeout: Ref.set(pingTimedOut, true),
256 }).pipe(
257 // Teardown sends an interrupt per open stream, and a write to a closed
258 // socket waits out the client's 1s interrupt timeout before reconnecting.
259 // Requests still go through, so they fail with the protocol error.
260 Effect.map((protocol) => ({
261 ...protocol,
262 send: (clientId, request, transferables) =>
263 request._tag !== "Request" && Deferred.isDoneUnsafe(disconnected)
264 ? Effect.void
265 : protocol.send(clientId, request, transferables),
266 })),
267 ),
268 ).pipe(
269 Layer.provide(
270 Layer.mergeAll(
271 layerSocket,
272 RpcSerialization.layerJson,
273 Layer.succeed(RpcClient.ConnectionHooks, hooks),
274 ),
275 ),
276 );
277 const protocolContext = yield* Layer.build(layerProtocol).pipe(
278 Effect.withSpan("environment.websocket.connect"),
279 );
280 // Building the protocol only starts the socket; this span times the
281 // handshake up to the first open, without making `connect` wait for it.
282 yield* Deferred.await(connected).pipe(
283 Effect.raceFirst(Deferred.await(disconnected)),
284 Effect.withSpan("environment.websocket.open"),
285 Effect.ignore,
286 Effect.forkScoped,
287 );
288 const protocolClient = yield* makeWsRpcProtocolClient.pipe(Effect.provide(protocolContext));
289 const initialConfigDeferred = yield* Deferred.make<ServerConfig>();
290 const serverConfigExit = yield* Deferred.make<void, ServerConfigSubscriptionError>();
291 const configSubscriptionClosed = yield* Deferred.make<never, ConnectionAttemptError>();
292 const serverConfigState = yield* Ref.make(Option.none<ServerConfigReplayState>());
293 const serverConfigUpdates = yield* PubSub.sliding<BufferedServerConfigEvent>(64);
294 const configSubscriptionEndedError = new ConnectionTransientErrorClass({
295 reason: "remote-unavailable",
296 detail: `${connection.label} config subscription ended.`,
297 });
298 const serverConfigSource = protocolClient[WS_METHODS.subscribeServerConfig](
299 onDemandSnapshots(serverConfigInput),
300 ).pipe(
301 Stream.runForEach((event) =>
302 Effect.gen(function* () {
303 const buffered = yield* Ref.modify(serverConfigState, (current) => {
304 const projection = applyServerConfigProjection(
305 Option.map(current, (state) => state.projection),
306 event,
307 );
308 if (Option.isNone(projection)) {
309 return [Option.none<BufferedServerConfigEvent>(), current] as const;
310 }
311 const next = {
312 projection: projection.value,
313 revision: Option.match(current, {
314 onNone: () => 1,
315 onSome: (state) => state.revision + 1,
316 }),
317 themesEvent:
318 event.type === "environmentThemesUpdated"
319 ? event
320 : Option.getOrUndefined(current)?.themesEvent,
321 sourcesEvent:
322 event.type === "usageLimitSourcesUpdated"
323 ? event
324 : Option.getOrUndefined(current)?.sourcesEvent,
325 } satisfies ServerConfigReplayState;
326 return [
327 Option.some({ event, replay: next, revision: next.revision }),
328 Option.some(next),
329 ] as const;
330 });
331 if (Option.isSome(buffered)) {
332 yield* PubSub.publish(serverConfigUpdates, buffered.value);
333 }
334 if (event.type === "snapshot") {
335 yield* Deferred.succeed(initialConfigDeferred, event.config);
336 }
337 }),
338 ),
339 Effect.onExit((exit) => {
340 if (Exit.isSuccess(exit)) {
341 return Effect.all([
342 Deferred.succeed(serverConfigExit, undefined),
343 Deferred.fail(configSubscriptionClosed, configSubscriptionEndedError),
344 ]).pipe(Effect.asVoid);
345 }
346 if (Cause.hasInterruptsOnly(exit.cause)) {
347 return Effect.void;
348 }
349 return Effect.all([
350 Deferred.failCause(serverConfigExit, exit.cause),
351 Deferred.failCause(configSubscriptionClosed, Cause.map(exit.cause, mapRpcError)),
352 ]).pipe(Effect.asVoid);
353 }),
354 );
355 yield* serverConfigSource.pipe(Effect.forkScoped);
356 const initialConfig = Effect.raceFirst(
357 Deferred.await(initialConfigDeferred),
358 Deferred.await(serverConfigExit).pipe(
359 Effect.mapError(mapRpcError),
360 Effect.flatMap(() => Effect.fail(configSubscriptionEndedError)),
361 ),
362 ).pipe(
363 Effect.filterOrElse(
364 (config) => config.environment.environmentId === connection.environmentId,
365 (config) =>
366 environmentMismatchError({
367 expected: connection.environmentId,
368 actual: config.environment.environmentId,
369 }),
370 ),
371 Effect.withSpan("environment.initialSync"),
372 );
373 const serverConfigEvents = Stream.unwrap(
374 Effect.gen(function* () {
375 const subscription = yield* PubSub.subscribe(serverConfigUpdates);
376 const snapshot = yield* Ref.get(serverConfigState);
377 if (Option.isNone(snapshot)) {
378 return Stream.empty;
379 }
380 const updates = Stream.fromSubscription(subscription).pipe(
381 Stream.filter((buffered) => buffered.revision > snapshot.value.revision),
382 Stream.mapAccum(
383 () => snapshot.value.revision,
384 (revision, buffered) => [
385 buffered.revision,
386 buffered.revision === revision + 1
387 ? [buffered.event]
388 : serverConfigReplayEvents(buffered.replay),
389 ],
390 ),
391 );
392 const terminal = Stream.fromEffect(Deferred.await(serverConfigExit)).pipe(Stream.drain);
393 return Stream.concat(
394 Stream.fromIterable(serverConfigReplayEvents(snapshot.value)),
395 Stream.merge(updates, terminal, { haltStrategy: "either" }),
396 );
397 }),
398 ).pipe(
399 Stream.catchCause((cause) => {
400 if (Cause.hasInterruptsOnly(cause)) {
401 return Stream.failCause(cause);
402 }
403 // The supervisor keeps the original cause. Shared durable consumers
404 // need a transport-shaped failure so they wait for its replacement.
405 return Stream.fail(
406 new RpcClientError.RpcClientError({
407 reason: new RpcClientError.RpcClientDefect({
408 message: `${connection.label} config subscription failed.`,
409 cause,
410 }),
411 }),
412 );
413 }),
414 );
415 const validatedInitialConfig = initialConfig.pipe(
416 Effect.mapError(
417 (cause) =>
418 new RpcClientError.RpcClientError({
419 reason: new RpcClientError.RpcClientDefect({
420 message: `${connection.label} config subscription failed.`,
421 cause,
422 }),
423 }),
424 ),
425 );
426 const subscribeServerConfig = (input: ServerConfigSubscriptionInput = serverConfigInput) =>
427 Stream.unwrap(
428 validatedInitialConfig.pipe(
429 Effect.as(
430 Equal.equals(input, serverConfigInput)
431 ? serverConfigEvents
432 : protocolClient[WS_METHODS.subscribeServerConfig](onDemandSnapshots(input)),
433 ),
434 ),
435 );
436 const probe = initialConfig.pipe(
437 Effect.flatMap(() =>
438 protocolClient[WS_METHODS.serverProbe]({}).pipe(Effect.mapError(mapRpcError)),
439 ),
440 Effect.asVoid,
441 Effect.withSpan("clientRuntime.connection.rpcSession.probe"),
442 );
443
444 return {
445 client: protocolClient,
446 initialConfig,
447 subscribeServerConfig,
448 ready: Deferred.await(connected).pipe(
449 Effect.andThen(initialConfig),
450 Effect.asVoid,
451 Effect.raceFirst(Deferred.await(disconnected)),
452 ),
453 probe,
454 inboundFrames: Effect.sync(() => inboundFrames),
455 closed: Effect.raceFirst(
456 Deferred.await(disconnected),
457 Deferred.await(configSubscriptionClosed),
458 ),
459 } satisfies RpcSession;
460 });
461
462 return RpcSessionFactory.of({ connect });
463});
464
465export const layer = (options: RpcSessionOptions) => Layer.effect(RpcSessionFactory, make(options));