import { WsRpcGroup } from "@t3tools/contracts";import * as Duration from "effect/Duration";import * as Effect from "effect/Effect";import * as Layer from "effect/Layer";import * as Schedule from "effect/Schedule";import { RpcClient, RpcSerialization } from "effect/unstable/rpc";import * as Socket from "effect/unstable/socket/Socket"; import { acknowledgeRpcRequest, clearAllTrackedRpcRequests, trackRpcRequestSent,} from "./requestLatencyState";import { getWsReconnectDelayMsForRetry, recordWsConnectionAttempt, recordWsConnectionClosed, recordWsConnectionErrored, recordWsConnectionOpened, type WsConnectionMetadata, WS_RECONNECT_MAX_RETRIES,} from "./wsConnectionState"; export interface WsProtocolCloseContext { readonly intentional: boolean;} export interface WsProtocolLifecycleHandlers { readonly getConnectionLabel?: () => string | null; readonly getVersionMismatchHint?: () => string | null; readonly isCloseIntentional?: () => boolean; readonly isActive?: () => boolean; readonly onAttempt?: (socketUrl: string) => void; readonly onOpen?: () => void; readonly onHeartbeatPing?: () => void; readonly onHeartbeatPong?: () => void; readonly onHeartbeatTimeout?: () => void; readonly onRequestStart?: (info: { readonly id: string; readonly tag: string; readonly stream: boolean; }) => void; readonly onRequestChunk?: (info: { readonly id: string; readonly tag: string; readonly chunkCount: number; }) => void; readonly onRequestExit?: (info: { readonly id: string; readonly tag: string; readonly stream: boolean; }) => void; readonly onRequestInterrupt?: (info: { readonly id: string; readonly tag?: string }) => void; readonly onError?: (message: string) => void; readonly onClose?: ( details: { readonly code: number; readonly reason: string }, context: WsProtocolCloseContext, ) => void;} export const makeWsRpcProtocolClient = RpcClient.make(WsRpcGroup);type RpcClientFactory = typeof makeWsRpcProtocolClient;export type WsRpcProtocolClient = RpcClientFactory extends Effect.Effect<infer Client, any, any> ? Client : never;export type WsRpcProtocolSocketUrlProvider = string | (() => Promise<string>); function formatSocketErrorMessage(error: unknown): string { if (error instanceof Error && error.message.trim().length > 0) { return error.message; } return String(error);} function resolveWsRpcSocketUrl(rawUrl: string): string { const resolved = new URL(rawUrl); if (resolved.protocol !== "ws:" && resolved.protocol !== "wss:") { throw new Error(`Unsupported websocket transport URL protocol: ${resolved.protocol}`); } resolved.pathname = "/ws"; return resolved.toString();} function resolveConnectionMetadata(handlers?: WsProtocolLifecycleHandlers): WsConnectionMetadata { return { connectionLabel: handlers?.getConnectionLabel?.() ?? null, versionMismatchHint: handlers?.getVersionMismatchHint?.() ?? null, };} type ComposedWsProtocolLifecycleHandlers = Required< Pick<WsProtocolLifecycleHandlers, "isActive" | "onAttempt" | "onOpen" | "onError" | "onClose">>; function defaultLifecycleHandlers( handlers?: WsProtocolLifecycleHandlers,): ComposedWsProtocolLifecycleHandlers { return { isActive: () => true, onAttempt: (socketUrl) => { recordWsConnectionAttempt(socketUrl, resolveConnectionMetadata(handlers)); }, onOpen: () => { recordWsConnectionOpened(resolveConnectionMetadata(handlers)); }, onError: (message) => { clearAllTrackedRpcRequests(); recordWsConnectionErrored(message, resolveConnectionMetadata(handlers)); }, onClose: (details, context) => { clearAllTrackedRpcRequests(); if (context.intentional) { return; } recordWsConnectionClosed(details, resolveConnectionMetadata(handlers)); }, };} function composeLifecycleHandlers( handlers?: WsProtocolLifecycleHandlers,): ComposedWsProtocolLifecycleHandlers { const defaults = defaultLifecycleHandlers(handlers); const isActive = handlers?.isActive ?? defaults.isActive; return { isActive, onAttempt: (socketUrl) => { if (!isActive()) { return; } defaults.onAttempt(socketUrl); handlers?.onAttempt?.(socketUrl); }, onOpen: () => { if (!isActive()) { return; } defaults.onOpen(); handlers?.onOpen?.(); }, onError: (message) => { if (!isActive()) { return; } defaults.onError(message); handlers?.onError?.(message); }, onClose: (details, context) => { if (!isActive()) { return; } defaults.onClose(details, context); handlers?.onClose?.(details, context); }, };} export function createWsRpcProtocolLayer( url: WsRpcProtocolSocketUrlProvider, handlers?: WsProtocolLifecycleHandlers,) { const lifecycle = composeLifecycleHandlers(handlers); const resolvedUrl = typeof url === "function" ? Effect.promise(() => url()).pipe( Effect.map((rawUrl) => resolveWsRpcSocketUrl(rawUrl)), Effect.tapError((error) => Effect.sync(() => { lifecycle.onError(formatSocketErrorMessage(error)); }), ), Effect.orDie, ) : resolveWsRpcSocketUrl(url); const trackingWebSocketConstructorLayer = Layer.succeed( Socket.WebSocketConstructor, (socketUrl, protocols) => { lifecycle.onAttempt(socketUrl); const socket = new globalThis.WebSocket(socketUrl, protocols); socket.addEventListener( "open", () => { lifecycle.onOpen(); }, { once: true }, ); socket.addEventListener( "error", () => { lifecycle.onError("Unable to connect to the T3 server WebSocket."); }, { once: true }, ); socket.addEventListener( "close", (event) => { lifecycle.onClose( { code: event.code, reason: event.reason, }, { intentional: handlers?.isCloseIntentional?.() ?? false, }, ); }, { once: true }, ); return socket; }, ); const socketLayer = Socket.layerWebSocket(resolvedUrl).pipe( Layer.provide(trackingWebSocketConstructorLayer), ); const retryPolicy = Schedule.addDelay(Schedule.recurs(WS_RECONNECT_MAX_RETRIES), (retryCount) => Effect.succeed(Duration.millis(getWsReconnectDelayMsForRetry(retryCount) ?? 0)), ); const protocolLayer = Layer.effect( RpcClient.Protocol, Effect.map( RpcClient.makeProtocolSocket({ retryPolicy, retryTransientErrors: true, }), (protocol) => ({ ...protocol, run: (clientId, writeResponse) => protocol.run(clientId, (response) => { if (response._tag === "ClientProtocolError" || response._tag === "Defect") { clearAllTrackedRpcRequests(); } return writeResponse(response); }), }), ), ); const requestHooksLayer = Layer.succeed( RpcClient.RequestHooks, RpcClient.RequestHooks.of({ onRequestStart: (info) => Effect.sync(() => { if (!lifecycle.isActive()) { return; } handlers?.onRequestStart?.({ id: String(info.id), tag: info.tag, stream: info.stream, }); trackRpcRequestSent(String(info.id), info.tag); }), onRequestChunk: (info) => Effect.sync(() => { if (!lifecycle.isActive()) { return; } handlers?.onRequestChunk?.({ id: String(info.id), tag: info.tag, chunkCount: info.chunkCount, }); acknowledgeRpcRequest(String(info.id)); }), onRequestExit: (info) => Effect.sync(() => { if (!lifecycle.isActive()) { return; } handlers?.onRequestExit?.({ id: String(info.id), tag: info.tag, stream: info.stream, }); acknowledgeRpcRequest(String(info.id)); }), onRequestInterrupt: (info) => Effect.sync(() => { if (!lifecycle.isActive()) { return; } handlers?.onRequestInterrupt?.({ id: String(info.id), ...(info.tag === undefined ? {} : { tag: info.tag }), }); acknowledgeRpcRequest(String(info.id)); }), }), ); const connectionHooksLayer = Layer.succeed( RpcClient.ConnectionHooks, RpcClient.ConnectionHooks.of({ onConnect: Effect.void, onDisconnect: Effect.void, onPing: Effect.sync(() => { if (lifecycle.isActive()) { handlers?.onHeartbeatPing?.(); } }), onPong: Effect.sync(() => { if (lifecycle.isActive()) { handlers?.onHeartbeatPong?.(); } }), onPingTimeout: Effect.sync(() => { if (lifecycle.isActive()) { clearAllTrackedRpcRequests(); recordWsConnectionErrored( "WebSocket heartbeat timed out.", resolveConnectionMetadata(handlers), ); handlers?.onHeartbeatTimeout?.(); } }), }), ); return Layer.mergeAll( protocolLayer.pipe( Layer.provide(Layer.mergeAll(socketLayer, RpcSerialization.layerJson, connectionHooksLayer)), ), requestHooksLayer, connectionHooksLayer, );}
import { WsRpcGroup } from "@t3tools/contracts";
import * as Duration from "effect/Duration";
import * as Effect from "effect/Effect";
import * as Layer from "effect/Layer";
import * as Schedule from "effect/Schedule";
import { RpcClient, RpcSerialization } from "effect/unstable/rpc";
import * as Socket from "effect/unstable/socket/Socket";
import {
acknowledgeRpcRequest,
clearAllTrackedRpcRequests,
trackRpcRequestSent,
} from "./requestLatencyState";
import {
getWsReconnectDelayMsForRetry,
recordWsConnectionAttempt,
recordWsConnectionClosed,
recordWsConnectionErrored,
recordWsConnectionOpened,
type WsConnectionMetadata,
WS_RECONNECT_MAX_RETRIES,
} from "./wsConnectionState";
export interface WsProtocolCloseContext {
readonly intentional: boolean;
}
export interface WsProtocolLifecycleHandlers {
readonly getConnectionLabel?: () => string | null;
readonly getVersionMismatchHint?: () => string | null;
readonly isCloseIntentional?: () => boolean;
readonly isActive?: () => boolean;
readonly onAttempt?: (socketUrl: string) => void;
readonly onOpen?: () => void;
readonly onHeartbeatPing?: () => void;
readonly onHeartbeatPong?: () => void;
readonly onHeartbeatTimeout?: () => void;
readonly onRequestStart?: (info: {
readonly id: string;
readonly tag: string;
readonly stream: boolean;
}) => void;
readonly onRequestChunk?: (info: {
readonly id: string;
readonly tag: string;
readonly chunkCount: number;
}) => void;
readonly onRequestExit?: (info: {
readonly id: string;
readonly tag: string;
readonly stream: boolean;
}) => void;
readonly onRequestInterrupt?: (info: { readonly id: string; readonly tag?: string }) => void;
readonly onError?: (message: string) => void;
readonly onClose?: (
details: { readonly code: number; readonly reason: string },
context: WsProtocolCloseContext,
) => void;
}
export const makeWsRpcProtocolClient = RpcClient.make(WsRpcGroup);
type RpcClientFactory = typeof makeWsRpcProtocolClient;
export type WsRpcProtocolClient =
RpcClientFactory extends Effect.Effect<infer Client, any, any> ? Client : never;
export type WsRpcProtocolSocketUrlProvider = string | (() => Promise<string>);
function formatSocketErrorMessage(error: unknown): string {
if (error instanceof Error && error.message.trim().length > 0) {
return error.message;
}
return String(error);
}
function resolveWsRpcSocketUrl(rawUrl: string): string {
const resolved = new URL(rawUrl);
if (resolved.protocol !== "ws:" && resolved.protocol !== "wss:") {
throw new Error(`Unsupported websocket transport URL protocol: ${resolved.protocol}`);
}
resolved.pathname = "/ws";
return resolved.toString();
}
function resolveConnectionMetadata(handlers?: WsProtocolLifecycleHandlers): WsConnectionMetadata {
return {
connectionLabel: handlers?.getConnectionLabel?.() ?? null,
versionMismatchHint: handlers?.getVersionMismatchHint?.() ?? null,
};
}
type ComposedWsProtocolLifecycleHandlers = Required<
Pick<WsProtocolLifecycleHandlers, "isActive" | "onAttempt" | "onOpen" | "onError" | "onClose">
>;
function defaultLifecycleHandlers(
handlers?: WsProtocolLifecycleHandlers,
): ComposedWsProtocolLifecycleHandlers {
return {
isActive: () => true,
onAttempt: (socketUrl) => {
recordWsConnectionAttempt(socketUrl, resolveConnectionMetadata(handlers));
},
onOpen: () => {
recordWsConnectionOpened(resolveConnectionMetadata(handlers));
},
onError: (message) => {
clearAllTrackedRpcRequests();
recordWsConnectionErrored(message, resolveConnectionMetadata(handlers));
},
onClose: (details, context) => {
clearAllTrackedRpcRequests();
if (context.intentional) {
return;
}
recordWsConnectionClosed(details, resolveConnectionMetadata(handlers));
},
};
}
function composeLifecycleHandlers(
handlers?: WsProtocolLifecycleHandlers,
): ComposedWsProtocolLifecycleHandlers {
const defaults = defaultLifecycleHandlers(handlers);
const isActive = handlers?.isActive ?? defaults.isActive;
return {
isActive,
onAttempt: (socketUrl) => {
if (!isActive()) {
return;
}
defaults.onAttempt(socketUrl);
handlers?.onAttempt?.(socketUrl);
},
onOpen: () => {
if (!isActive()) {
return;
}
defaults.onOpen();
handlers?.onOpen?.();
},
onError: (message) => {
if (!isActive()) {
return;
}
defaults.onError(message);
handlers?.onError?.(message);
},
onClose: (details, context) => {
if (!isActive()) {
return;
}
defaults.onClose(details, context);
handlers?.onClose?.(details, context);
},
};
}
export function createWsRpcProtocolLayer(
url: WsRpcProtocolSocketUrlProvider,
handlers?: WsProtocolLifecycleHandlers,
) {
const lifecycle = composeLifecycleHandlers(handlers);
const resolvedUrl =
typeof url === "function"
? Effect.promise(() => url()).pipe(
Effect.map((rawUrl) => resolveWsRpcSocketUrl(rawUrl)),
Effect.tapError((error) =>
Effect.sync(() => {
lifecycle.onError(formatSocketErrorMessage(error));
}),
),
Effect.orDie,
)
: resolveWsRpcSocketUrl(url);
const trackingWebSocketConstructorLayer = Layer.succeed(
Socket.WebSocketConstructor,
(socketUrl, protocols) => {
lifecycle.onAttempt(socketUrl);
const socket = new globalThis.WebSocket(socketUrl, protocols);
socket.addEventListener(
"open",
() => {
lifecycle.onOpen();
},
{ once: true },
);
socket.addEventListener(
"error",
() => {
lifecycle.onError("Unable to connect to the T3 server WebSocket.");
},
{ once: true },
);
socket.addEventListener(
"close",
(event) => {
lifecycle.onClose(
{
code: event.code,
reason: event.reason,
},
{
intentional: handlers?.isCloseIntentional?.() ?? false,
},
);
},
{ once: true },
);
return socket;
},
);
const socketLayer = Socket.layerWebSocket(resolvedUrl).pipe(
Layer.provide(trackingWebSocketConstructorLayer),
);
const retryPolicy = Schedule.addDelay(Schedule.recurs(WS_RECONNECT_MAX_RETRIES), (retryCount) =>
Effect.succeed(Duration.millis(getWsReconnectDelayMsForRetry(retryCount) ?? 0)),
);
const protocolLayer = Layer.effect(
RpcClient.Protocol,
Effect.map(
RpcClient.makeProtocolSocket({
retryPolicy,
retryTransientErrors: true,
}),
(protocol) => ({
...protocol,
run: (clientId, writeResponse) =>
protocol.run(clientId, (response) => {
if (response._tag === "ClientProtocolError" || response._tag === "Defect") {
clearAllTrackedRpcRequests();
}
return writeResponse(response);
}),
}),
),
);
const requestHooksLayer = Layer.succeed(
RpcClient.RequestHooks,
RpcClient.RequestHooks.of({
onRequestStart: (info) =>
Effect.sync(() => {
if (!lifecycle.isActive()) {
return;
}
handlers?.onRequestStart?.({
id: String(info.id),
tag: info.tag,
stream: info.stream,
});
trackRpcRequestSent(String(info.id), info.tag);
}),
onRequestChunk: (info) =>
Effect.sync(() => {
if (!lifecycle.isActive()) {
return;
}
handlers?.onRequestChunk?.({
id: String(info.id),
tag: info.tag,
chunkCount: info.chunkCount,
});
acknowledgeRpcRequest(String(info.id));
}),
onRequestExit: (info) =>
Effect.sync(() => {
if (!lifecycle.isActive()) {
return;
}
handlers?.onRequestExit?.({
id: String(info.id),
tag: info.tag,
stream: info.stream,
});
acknowledgeRpcRequest(String(info.id));
}),
onRequestInterrupt: (info) =>
Effect.sync(() => {
if (!lifecycle.isActive()) {
return;
}
handlers?.onRequestInterrupt?.({
id: String(info.id),
...(info.tag === undefined ? {} : { tag: info.tag }),
});
acknowledgeRpcRequest(String(info.id));
}),
}),
);
const connectionHooksLayer = Layer.succeed(
RpcClient.ConnectionHooks,
RpcClient.ConnectionHooks.of({
onConnect: Effect.void,
onDisconnect: Effect.void,
onPing: Effect.sync(() => {
if (lifecycle.isActive()) {
handlers?.onHeartbeatPing?.();
}
}),
onPong: Effect.sync(() => {
if (lifecycle.isActive()) {
handlers?.onHeartbeatPong?.();
}
}),
onPingTimeout: Effect.sync(() => {
if (lifecycle.isActive()) {
clearAllTrackedRpcRequests();
recordWsConnectionErrored(
"WebSocket heartbeat timed out.",
resolveConnectionMetadata(handlers),
);
handlers?.onHeartbeatTimeout?.();
}
}),
}),
);
return Layer.mergeAll(
protocolLayer.pipe(
Layer.provide(Layer.mergeAll(socketLayer, RpcSerialization.layerJson, connectionHooksLayer)),
),
requestHooksLayer,
connectionHooksLayer,
);
}