Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
9.1 kB · 275 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276import { DaemonCapabilityUnavailableError, type DaemonClientCloseListener, type DaemonClientMessageListener, type DaemonClientRequestOptions, type DaemonCommandBody, type DaemonHello, DaemonSocketClosedError, type DaemonTransportClient, getDaemonSocketCloseReason,} from "./daemon-client.js";import { type DaemonClosingReason, type DaemonCommand, type DaemonPeerTransportTicket, type DaemonResponse, type DaemonServerCapability, getDaemonCommandCompatibilities, isSessionPlaneDaemonCommand, meetsDaemonCommandCompatibility,} from "./daemon-protocol.js";import { getDaemonSocketIdentity } from "./daemon-socket.js";import { DaemonWorkerClient } from "./daemon-worker-client.js";
export class DaemonDirectTransportClosedError extends DaemonSocketClosedError { constructor(cause: Error) { super("direct-worker", getDaemonSocketCloseReason(cause), cause.message); this.name = "DaemonDirectTransportClosedError"; }}
export class DaemonControlPlaneTransportError extends Error { constructor(cause: Error) { super(`Daemon control-plane transport failed: ${cause.message}`, { cause }); this.name = "DaemonControlPlaneTransportError"; }}
/** * One logical daemon connection over two sockets: session-plane commands go direct to the worker, * the rest to the supervisor. Any direct-path loss degrades silently to supervisor routing; an * established direct link keeps serving the session while the supervisor reconnects. */export class DaemonRoutedClient implements DaemonTransportClient { private direct?: DaemonWorkerClient; private readonly messageListeners = new Set<DaemonClientMessageListener>(); private readonly closeListeners = new Set<DaemonClientCloseListener>(); private readonly unsubscribeSupervisorMessage: () => void; private readonly unsubscribeSupervisorClose: () => void; private unsubscribeDirectMessage?: () => void; private unsubscribeDirectClose?: () => void; private closed = false;
constructor( private readonly supervisor: DaemonTransportClient, direct: DaemonWorkerClient, ) { this.direct = direct; this.unsubscribeSupervisorMessage = supervisor.onMessage((message) => this.emitMessage(message)); this.unsubscribeSupervisorClose = supervisor.onClose((error) => { if (!this.closed) this.emitClose(error); }); this.bindDirect(direct); }
get hello(): DaemonHello | undefined { return this.supervisor.hello ?? this.direct?.hello; }
get isConnected(): boolean { return this.supervisor.isConnected || this.direct?.isConnected === true; }
get hasDirectTransport(): boolean { return this.direct?.isConnected === true; }
get controlPlaneTransport(): DaemonTransportClient { return this.supervisor; }
get isControlPlaneReady(): boolean { return this.supervisor.isConnected && this.supervisor.hello !== undefined; }
supportsServerCapability(capability: DaemonServerCapability): boolean { return ( this.supervisor.supportsServerCapability(capability) || this.direct?.supportsServerCapability(capability) === true ); }
waitForHello(timeoutMs = 3000): Promise<DaemonHello> { return this.supervisor.waitForHello(timeoutMs); }
async connect(timeoutMs = 3000): Promise<void> { if (!this.supervisor.isConnected) await this.supervisor.connect(timeoutMs); }
async reconnect(timeoutMs = 3000): Promise<void> { if (!this.supervisor.isConnected) await this.supervisor.reconnect(timeoutMs); }
disconnectForReconnect(reason: DaemonClosingReason): void { this.fallbackToSupervisor(); this.supervisor.disconnectForReconnect(reason); }
resetTransportForReconnect(): void { this.supervisor.resetTransportForReconnect(); }
onMessage(listener: DaemonClientMessageListener): () => void { this.messageListeners.add(listener); return () => this.messageListeners.delete(listener); }
onClose(listener: DaemonClientCloseListener): () => void { this.closeListeners.add(listener); return () => this.closeListeners.delete(listener); }
enableRequestRecovery(): void { this.supervisor.enableRequestRecovery(); }
request( command: DaemonCommandBody, timeoutMs = 30_000, options: DaemonClientRequestOptions = {}, ): Promise<DaemonResponse> { if (command.type === "reattach") { // Reattach may land on a different worker; the direct link to the old one is stale on success. return this.requestControlPlane(command, timeoutMs, options).then((response) => { if (response.success) this.fallbackToSupervisor(); return response; }); } const direct = this.direct; if (direct?.isConnected && this.servesDirect(direct, command)) { return direct.request(command, timeoutMs, options); } return this.requestControlPlane(command, timeoutMs, options); }
private servesDirect(direct: DaemonWorkerClient, command: DaemonCommandBody): boolean { if (!isSessionPlaneDaemonCommand(command.type)) return false; const hello = direct.hello; if (!hello) return false; return getDaemonCommandCompatibilities(command as DaemonCommand).every((compatibility) => meetsDaemonCommandCompatibility(hello, compatibility), ); }
private requestControlPlane( command: DaemonCommandBody, timeoutMs: number, options: DaemonClientRequestOptions, ): Promise<DaemonResponse> { return this.supervisor.request(command, timeoutMs, options).catch((error: unknown) => { if (error instanceof DaemonCapabilityUnavailableError) throw error; throw new DaemonControlPlaneTransportError(error instanceof Error ? error : new Error(String(error))); }); }
fallbackToSupervisor(): void { const direct = this.direct; if (!direct) return; this.direct = undefined; this.unsubscribeDirectMessage?.(); this.unsubscribeDirectMessage = undefined; this.unsubscribeDirectClose?.(); this.unsubscribeDirectClose = undefined; direct.close(); }
close(): void { if (this.closed) return; this.closed = true; this.fallbackToSupervisor(); this.unsubscribeSupervisorMessage(); this.unsubscribeSupervisorClose(); this.supervisor.close(); this.messageListeners.clear(); this.closeListeners.clear(); }
private bindDirect(direct: DaemonWorkerClient): void { this.unsubscribeDirectMessage = direct.onMessage((message) => this.emitMessage(message)); this.unsubscribeDirectClose = direct.onClose((error) => { if (this.closed || this.direct !== direct) return; this.direct = undefined; this.unsubscribeDirectMessage?.(); this.unsubscribeDirectMessage = undefined; this.unsubscribeDirectClose = undefined; this.emitClose(new DaemonDirectTransportClosedError(error)); }); }
private emitMessage(message: Parameters<DaemonClientMessageListener>[0]): void { for (const listener of [...this.messageListeners]) listener(message); }
private emitClose(error: Error): void { for (const listener of [...this.closeListeners]) listener(error); }}
/** Upgrade a supervisor connection with a direct worker link; every failure returns the supervisor unchanged. */export async function createDaemonSessionTransport( supervisor: DaemonTransportClient, activeSessionId: string, directDisabled: boolean,): Promise<DaemonTransportClient> { if ( directDisabled || supervisor instanceof DaemonRoutedClient || !supervisor.supportsServerCapability("direct_peer_transport") ) { return supervisor; } let direct: DaemonWorkerClient | undefined; try { // recoverable:false — this caller owns the fallback; a parked ticket request would pend the attach forever. const response = await supervisor.request({ type: "get_direct_worker_transport", activeSessionId }, 5000, { recoverable: false, }); if (!response.success) return supervisor; const ticket = readSessionTransportTicket(response.data); if (!ticket || ticket.activeSessionId !== activeSessionId || Date.parse(ticket.expiresAt) <= Date.now()) { return supervisor; } const currentIdentity = getDaemonSocketIdentity(ticket.socketPath); if ( !currentIdentity || currentIdentity.dev !== ticket.socketIdentity.dev || currentIdentity.ino !== ticket.socketIdentity.ino ) { return supervisor; } direct = new DaemonWorkerClient(ticket.socketPath); await direct.connect(1000); await direct.waitForHello(); await direct.authenticatePeer(ticket); return new DaemonRoutedClient(supervisor, direct); } catch { direct?.close(); return supervisor; }}
function readSessionTransportTicket(value: unknown): DaemonPeerTransportTicket | undefined { if (!value || typeof value !== "object") return undefined; const candidate = value as Partial<DaemonPeerTransportTicket>; if ( candidate.purpose !== "session_client" || typeof candidate.socketPath !== "string" || typeof candidate.workerInstanceId !== "string" || typeof candidate.activeSessionId !== "string" || typeof candidate.grantId !== "string" || typeof candidate.token !== "string" || typeof candidate.expiresAt !== "string" ) { return undefined; } if ( !candidate.socketIdentity || typeof candidate.socketIdentity.dev !== "number" || typeof candidate.socketIdentity.ino !== "number" ) { return undefined; } return candidate as DaemonPeerTransportTicket;}