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.
12 kB · 369 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370import { createConnection, type Socket } from "node:net";import { serializeJsonLine } from "../rpc/jsonl.js";import { type PrivateFrame, PrivateFramedChannel } from "../session-worker/private-framing.js";import { type DaemonClientMessageListener, type DaemonClientRequestOptions, DaemonSocketClosedError,} from "./daemon-client.js";import type { DaemonClosingReason, DaemonCommand, DaemonOutbound, DaemonPeerTransportTicket, DaemonResponse, DaemonServerCapability,} from "./daemon-protocol.js";import { type DaemonPeerCommand, type DaemonPeerCommandBody, type DaemonWorkerCommand, type DaemonWorkerCommandBody, type DaemonWorkerFrameHeader, isDaemonWorkerFrameHeader,} from "./daemon-worker-protocol.js";
type DistributiveOmit<T, K extends keyof T> = T extends unknown ? Omit<T, K> : never;type DaemonCommandBody = DistributiveOmit<DaemonCommand, "id">;type DaemonWorkerWireCommandBody = DaemonCommandBody | DaemonWorkerCommandBody | DaemonPeerCommandBody;type DaemonWorkerWireCommand = DaemonCommand | DaemonWorkerCommand | DaemonPeerCommand;type DaemonWorkerAuthentication = Omit<Extract<DaemonWorkerCommand, { type: "worker_auth" }>, "id" | "type" | "token">;
export type DaemonWorkerFrameListener = (frame: PrivateFrame<DaemonWorkerFrameHeader>) => void;export type DaemonWorkerCloseListener = (error: Error) => void;type DaemonHello = Extract<DaemonOutbound, { type: "daemon_hello" }>;
export class DaemonWorkerAuthenticationError extends Error {}
/** A probe (hello/response/connect) timed out; recovery treats the worker as live-but-slow, never dead. */export class DaemonWorkerProbeTimeoutError extends Error {}
export class DaemonWorkerClient { private socket?: Socket; private channel?: PrivateFramedChannel<DaemonWorkerFrameHeader>; private readonly frameListeners = new Set<DaemonWorkerFrameListener>(); private readonly messageListeners = new Set<DaemonClientMessageListener>(); private readonly closeListeners = new Set<DaemonWorkerCloseListener>(); private readonly pending = new Map< string, { resolve: (response: DaemonResponse) => void; reject: (error: Error) => void; timeout: ReturnType<typeof setTimeout>; } >(); private requestId = 0; private helloMessage?: DaemonHello; private directPeer = false; private directClosingReason?: DaemonClosingReason; private readonly helloWaiters = new Set<{ resolve: (hello: DaemonHello) => void; reject: (error: Error) => void; timeout: ReturnType<typeof setTimeout>; }>();
constructor(private readonly socketPath: string) {}
get hello(): DaemonHello | undefined { return this.helloMessage; }
get isConnected(): boolean { return this.socket !== undefined && !this.socket.destroyed; }
supportsServerCapability(capability: DaemonServerCapability): boolean { return this.helloMessage?.serverCapabilities?.includes(capability) === true; }
async connect(timeoutMs = 3000): Promise<void> { if (this.socket) { throw new Error("Daemon worker client is already connected"); } const socket = createConnection(this.socketPath); this.socket = socket; this.channel = new PrivateFramedChannel(socket, isDaemonWorkerFrameHeader); this.channel.onFrame((frame) => this.handleFrame(frame));
await new Promise<void>((resolve, reject) => { const timeout = setTimeout(() => { cleanup(); socket.destroy(); reject(new Error(`Timed out connecting to daemon worker socket: ${this.socketPath}`)); }, timeoutMs); const cleanup = () => { clearTimeout(timeout); socket.off("connect", onConnect); socket.off("error", onError); }; const onConnect = () => { cleanup(); resolve(); }; const onError = (error: Error) => { cleanup(); reject(error); }; socket.once("connect", onConnect); socket.once("error", onError); });
socket.on("error", (error) => this.notifyClosed(socket, this.directCloseError(error))); socket.on("close", () => this.notifyClosed(socket, this.directCloseError(new Error("Daemon worker socket closed"))), ); }
waitForHello(timeoutMs = 3000): Promise<DaemonHello> { if (this.helloMessage) { return Promise.resolve(this.helloMessage); } if (!this.socket || this.socket.destroyed) { return Promise.reject(new Error("Daemon worker client is not connected")); } return new Promise((resolve, reject) => { const waiter = { resolve, reject, timeout: setTimeout(() => { this.helloWaiters.delete(waiter); reject(new DaemonWorkerProbeTimeoutError("Timed out waiting for daemon worker hello")); }, timeoutMs), }; this.helloWaiters.add(waiter); }); }
onFrame(listener: DaemonWorkerFrameListener): () => void { this.frameListeners.add(listener); return () => this.frameListeners.delete(listener); }
onMessage(listener: DaemonClientMessageListener): () => void { this.messageListeners.add(listener); return () => this.messageListeners.delete(listener); }
onClose(listener: DaemonWorkerCloseListener): () => void { this.closeListeners.add(listener); return () => this.closeListeners.delete(listener); }
request( command: DaemonCommandBody, timeoutMs = 30_000, // Progress/recovery options are supervisor-transport features; a direct request fails fast instead of replaying (no double execution). options: DaemonClientRequestOptions = {}, ): Promise<DaemonResponse> { return this.requestWire(command, timeoutMs, options.onResponse); }
requestWorker(command: DaemonWorkerCommandBody, timeoutMs = 30_000): Promise<DaemonResponse> { return this.requestWire(command, timeoutMs); }
async authenticateWorker( token: string, owner: DaemonWorkerAuthentication, timeoutMs = 3000, ): Promise<Extract<DaemonResponse, { success: true }>> { const response = await this.requestWorker({ type: "worker_auth", token, ...owner }, timeoutMs); if (!response.success) { throw new DaemonWorkerAuthenticationError(response.error); } return response; }
async authenticatePeer(ticket: DaemonPeerTransportTicket, timeoutMs = 3000): Promise<void> { const response = await this.requestWire( { type: "peer_auth", grantId: ticket.grantId, token: ticket.token, workerInstanceId: ticket.workerInstanceId, purpose: ticket.purpose, }, timeoutMs, ); if (!response.success) { throw new Error(response.error); } this.directPeer = true; }
close(): void { this.rejectAll(new Error("Daemon worker client closed")); this.channel?.close(); this.channel = undefined; this.socket?.destroy(); this.socket = undefined; this.directPeer = false; this.directClosingReason = undefined; }
private async requestWire( command: DaemonWorkerWireCommandBody, timeoutMs: number, onResponse?: DaemonClientRequestOptions["onResponse"], ): Promise<DaemonResponse> { if (!this.channel || !this.socket || this.socket.destroyed) { throw new Error("Daemon worker client is not connected"); } const id = `worker_${++this.requestId}`; const fullCommand = { ...command, id } as DaemonWorkerWireCommand; const response = new Promise<DaemonResponse>((resolve, reject) => { const timeout = setTimeout(() => { this.pending.delete(id); reject( new DaemonWorkerProbeTimeoutError(`Timed out waiting for daemon worker response to ${command.type}`), ); }, timeoutMs); this.pending.set(id, { resolve: (response) => { try { onResponse?.(response); resolve(response); } catch (error) { reject(error); } }, reject, timeout, }); }); try { await this.channel.send( { kind: "command", requestId: id, commandType: command.type }, Buffer.from(serializeJsonLine(fullCommand)), ); } catch (error) { const pending = this.pending.get(id); if (pending) { clearTimeout(pending.timeout); this.pending.delete(id); pending.reject(error instanceof Error ? error : new Error(String(error))); } } return response; }
private handleFrame(frame: PrivateFrame<DaemonWorkerFrameHeader>): void { if (frame.header.kind !== "outbound") { return; } if (frame.header.outboundType === "response" && frame.header.requestId) { const pending = this.pending.get(frame.header.requestId); if (pending) { let response: unknown; try { response = JSON.parse(frame.payload.toString("utf8")); } catch (error) { clearTimeout(pending.timeout); this.pending.delete(frame.header.requestId); pending.reject(new Error(`Invalid daemon worker response: ${String(error)}`)); return; } if (isDaemonResponse(response)) { clearTimeout(pending.timeout); this.pending.delete(frame.header.requestId); pending.resolve(response); return; } } } if (frame.header.outboundType === "daemon_hello") { try { const parsed = JSON.parse(frame.payload.toString("utf8")) as DaemonOutbound; if (parsed.type === "daemon_hello") { this.helloMessage = parsed; for (const waiter of [...this.helloWaiters]) { clearTimeout(waiter.timeout); this.helloWaiters.delete(waiter); waiter.resolve(parsed); } } } catch { // Invalid hello payloads are rejected by the timeout. } } for (const listener of this.frameListeners) { listener(frame); } if (this.directPeer && frame.header.outboundType !== "daemon_hello" && frame.header.outboundType !== "response") { this.emitDirectOutbound(frame); } }
/** A malformed outbound frame closes the direct link (the routed client falls back) instead of throwing into consumers. */ private emitDirectOutbound(frame: PrivateFrame<DaemonWorkerFrameHeader>): void { if (frame.header.kind !== "outbound") return; let message: DaemonOutbound; try { if (frame.header.payloadEncoding !== undefined && frame.header.payloadEncoding !== "jsonl") { throw new Error(`Direct worker sent an unsupported payload encoding: ${frame.header.payloadEncoding}`); } const parsed = JSON.parse(frame.payload.toString("utf8")) as unknown; if (!parsed || typeof parsed !== "object" || typeof (parsed as { type?: unknown }).type !== "string") { throw new Error("Direct worker sent an invalid outbound payload"); } message = parsed as DaemonOutbound; } catch (error) { const socket = this.socket; const directError = this.directCloseError(error instanceof Error ? error : new Error(String(error))); if (socket) { this.notifyClosed(socket, directError); socket.destroy(directError); } return; } if (message.type === "daemon_closing") { this.directClosingReason = message.reason; } for (const listener of [...this.messageListeners]) { listener(message); } }
private directCloseError(cause: Error): Error { return this.directPeer ? new DaemonSocketClosedError(this.socketPath, this.directClosingReason, cause.message) : cause; }
private rejectAll(error: Error): void { for (const [id, pending] of this.pending) { clearTimeout(pending.timeout); pending.reject(error); this.pending.delete(id); } for (const waiter of [...this.helloWaiters]) { clearTimeout(waiter.timeout); this.helloWaiters.delete(waiter); waiter.reject(error); } }
private notifyClosed(socket: Socket, error: Error): void { if (this.socket !== socket) { return; } this.socket = undefined; this.channel = undefined; this.directPeer = false; this.directClosingReason = undefined; this.rejectAll(error); for (const listener of [...this.closeListeners]) { listener(error); } }}
function isDaemonResponse(value: unknown): value is DaemonResponse { if (!value || typeof value !== "object") { return false; } const candidate = value as { type?: unknown; command?: unknown; success?: unknown }; return ( candidate.type === "response" && typeof candidate.command === "string" && typeof candidate.success === "boolean" );}