feat(coding-agent): coordinate replaceable servers and workers

Keep the stable socket in an opaque detached coordinator while server processes own detached worker discovery and lifecycle. Remove the public drain RPC in favor of server replacement through coordinator routing.
This commit is contained in:
Armin Ronacher 2026-08-15 11:49:53 +02:00
parent 33eef6d1c6
commit 8dce5da769
23 changed files with 1725 additions and 649 deletions

View file

@ -0,0 +1,353 @@
import { type ChildProcess, spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import { isAbsolute } from "node:path";
import type { JsonlSessionMetadata } from "@earendil-works/pi-agent-core";
import type { HostedHarnessHandle } from "@earendil-works/pi-server";
import Type, { type Static } from "typebox";
import { Check } from "typebox/value";
import type { CoordinatorServer, CoordinatorServerEvent } from "./coordinator-client.ts";
import {
createExperimentalSessionWorkerLaunchSpec,
SESSION_WORKER_CONTROL_ADDRESS_ENV,
SESSION_WORKER_CONTROL_TOKEN_ENV,
SESSION_WORKER_COORDINATED_ENV,
SESSION_WORKER_PEER_ID_ENV,
SESSION_WORKER_SESSION_KEY_ENV,
} from "./session-worker.ts";
const WORKER_STARTUP_TIMEOUT_MS = 15_000;
const WORKER_SHUTDOWN_TIMEOUT_MS = 10_000;
const WORKER_DISCOVERY_TIMEOUT_MS = 5_000;
const JsonlSessionMetadataSchema = Type.Object({
id: Type.String({ minLength: 1 }),
createdAt: Type.Integer(),
storageVersion: Type.Integer(),
cwd: Type.String(),
path: Type.String(),
modifiedAt: Type.Number(),
parentSessionId: Type.Optional(Type.String()),
legacyParentSessionPath: Type.Optional(Type.String()),
});
const WorkerMessageSchema = Type.Union([
Type.Object({
type: Type.Literal("worker_ready"),
token: Type.String(),
sessionKey: Type.String(),
sessionId: Type.String(),
pid: Type.Integer({ minimum: 1 }),
metadata: JsonlSessionMetadataSchema,
}),
Type.Object({
type: Type.Literal("worker_failed"),
token: Type.String(),
sessionKey: Type.String(),
message: Type.String(),
}),
]);
type WorkerMessage = Static<typeof WorkerMessageSchema>;
interface WorkerRecord {
readonly peerId: string;
readonly metadata: JsonlSessionMetadata;
readonly pid: number;
readonly terminated: Promise<Error | undefined>;
resolveTerminated(error: Error | undefined): void;
expectedStop: boolean;
}
interface PendingLaunch {
readonly sessionKey: string;
readonly peerId: string;
readonly token: string;
readonly child: ChildProcess;
readonly timer: NodeJS.Timeout;
readonly promise: Promise<WorkerRecord>;
resolve(worker: WorkerRecord): void;
reject(error: Error): void;
}
/** Session and process bookkeeping owned by one replaceable Pi server process. */
export class CoordinatedSessionWorkers {
readonly workerPids = new Map<string, number>();
readonly #coordinator: CoordinatorServer;
readonly #sessionDir: string;
readonly #workersBySession = new Map<string, WorkerRecord>();
readonly #workersByPeer = new Map<string, WorkerRecord>();
readonly #pending = new Map<string, PendingLaunch>();
readonly #removeListener: () => void;
#discoveryPeers?: Set<string>;
#resolveDiscovery?: () => void;
#detached = false;
#shuttingDown = false;
constructor(coordinator: CoordinatorServer, sessionDir: string) {
this.#coordinator = coordinator;
this.#sessionDir = sessionDir;
this.#removeListener = coordinator.onEvent((event) => this.#handleCoordinatorEvent(event));
}
get trackedSessions(): readonly JsonlSessionMetadata[] {
return [...this.#workersBySession.values()].map((worker) => worker.metadata);
}
async discover(peerIds: ReadonlySet<string>): Promise<void> {
if (this.#detached) return;
const undiscovered = new Set(
[...peerIds].filter((peerId) => !this.#workersByPeer.has(peerId) && !this.#pendingPeer(peerId)),
);
if (undiscovered.size === 0) return;
this.#discoveryPeers = undiscovered;
const discovered = new Promise<void>((resolve) => {
this.#resolveDiscovery = resolve;
});
await this.#coordinator.broadcast({ type: "discover_workers" });
let timer: NodeJS.Timeout | undefined;
await Promise.race([
discovered,
new Promise<void>((resolve) => {
timer = setTimeout(resolve, WORKER_DISCOVERY_TIMEOUT_MS);
timer.unref();
}),
]);
if (timer) clearTimeout(timer);
this.#discoveryPeers = undefined;
this.#resolveDiscovery = undefined;
}
async createHarness(metadata: JsonlSessionMetadata): Promise<HostedHarnessHandle> {
if (this.#detached || this.#shuttingDown) throw new Error("Experimental server is shutting down");
const existing = this.#workersBySession.get(metadata.path);
if (existing) return { terminated: existing.terminated, close: async () => {} };
const pending = this.#pending.get(metadata.path);
if (pending) {
const worker = await pending.promise;
return { terminated: worker.terminated, close: async () => {} };
}
const worker = await this.#launch(metadata);
return { terminated: worker.terminated, close: async () => {} };
}
async shutdown(): Promise<void> {
if (this.#detached || this.#shuttingDown) return;
this.#shuttingDown = true;
for (const pending of this.#pending.values()) {
void this.#coordinator.send(pending.peerId, { type: "shutdown" }).catch(() => {});
}
const workers = [...this.#workersBySession.values()];
for (const worker of workers) {
worker.expectedStop = true;
void this.#coordinator.send(worker.peerId, { type: "shutdown" }).catch(() => {});
}
const finished = Promise.all([
...workers.map((worker) => worker.terminated),
...[...this.#pending.values()].map((pending) => {
if (pending.child.exitCode !== null || pending.child.signalCode !== null) return Promise.resolve();
return new Promise<void>((resolve) => pending.child.once("exit", () => resolve()));
}),
]).then(() => undefined);
let timer: NodeJS.Timeout | undefined;
const timedOut = new Promise<boolean>((resolve) => {
timer = setTimeout(() => resolve(true), WORKER_SHUTDOWN_TIMEOUT_MS);
timer.unref();
});
if ((await Promise.race([finished.then(() => false), timedOut])) === true) {
for (const worker of workers) {
try {
process.kill(worker.pid, "SIGKILL");
} catch (error) {
if (!(error instanceof Error && "code" in error && error.code === "ESRCH")) throw error;
}
}
for (const pending of this.#pending.values()) pending.child.kill("SIGKILL");
await finished;
}
if (timer) clearTimeout(timer);
this.#detachState();
}
/** Forget workers without stopping them when this server is replaced. */
detach(): void {
if (this.#detached) return;
this.#detached = true;
for (const pending of this.#pending.values()) {
clearTimeout(pending.timer);
pending.reject(new Error("Experimental server was replaced"));
}
this.#detachState();
}
#launch(metadata: JsonlSessionMetadata): Promise<WorkerRecord> {
const sessionKey = metadata.path;
const peerId = `worker-${randomUUID()}`;
const token = randomUUID();
const launch = createExperimentalSessionWorkerLaunchSpec(metadata, { sessionDir: this.#sessionDir });
let child: ChildProcess;
try {
child = spawn(launch.command, launch.args, {
cwd: launch.cwd,
detached: true,
env: {
...launch.env,
[SESSION_WORKER_CONTROL_ADDRESS_ENV]: this.#coordinator.controlPath,
[SESSION_WORKER_CONTROL_TOKEN_ENV]: token,
[SESSION_WORKER_SESSION_KEY_ENV]: Buffer.from(sessionKey).toString("base64url"),
[SESSION_WORKER_COORDINATED_ENV]: "1",
[SESSION_WORKER_PEER_ID_ENV]: peerId,
},
stdio: "ignore",
windowsHide: true,
});
child.unref();
} catch (error) {
return Promise.reject(error);
}
let resolve!: (worker: WorkerRecord) => void;
let reject!: (error: Error) => void;
const promise = new Promise<WorkerRecord>((resolvePromise, rejectPromise) => {
resolve = resolvePromise;
reject = rejectPromise;
});
const timer = setTimeout(
() => this.#failPending(sessionKey, new Error("Session worker startup timed out")),
WORKER_STARTUP_TIMEOUT_MS,
);
timer.unref();
const pending = { sessionKey, peerId, token, child, timer, promise, resolve, reject };
this.#pending.set(sessionKey, pending);
child.once("error", (error) => this.#failPending(sessionKey, error));
child.once("exit", (code, signal) => this.#childExited(pending, code, signal));
return promise;
}
#handleCoordinatorEvent(event: CoordinatorServerEvent): void {
if (this.#detached) return;
if (event.type === "peer_disconnected") {
this.#markDiscovered(event.peerId);
const worker = this.#workersByPeer.get(event.peerId);
if (worker) {
this.#removeWorker(
worker,
worker.expectedStop
? undefined
: new Error(`Session worker ${worker.metadata.id} disconnected unexpectedly`),
);
}
const pending = this.#pendingPeer(event.peerId);
if (pending) this.#failPending(pending.sessionKey, new Error("Session worker disconnected during startup"));
return;
}
if (event.type !== "message" || !Check(WorkerMessageSchema, event.payload)) return;
const message: WorkerMessage = event.payload;
if (message.type === "worker_failed") {
const pending = this.#pending.get(message.sessionKey);
if (pending?.peerId === event.from && pending.token === message.token) {
this.#failPending(message.sessionKey, new Error(`Session worker failed: ${message.message}`));
}
return;
}
this.#recordReadyWorker(event.from, message);
}
#recordReadyWorker(peerId: string, message: Extract<WorkerMessage, { type: "worker_ready" }>): void {
if (
message.sessionKey !== message.metadata.path ||
message.sessionId !== message.metadata.id ||
!isAbsolute(message.metadata.cwd) ||
!isAbsolute(message.metadata.path)
) {
return;
}
this.#markDiscovered(peerId);
const existing = this.#workersBySession.get(message.sessionKey);
if (existing && existing.peerId !== peerId) {
void this.#coordinator.send(peerId, { type: "shutdown" }).catch(() => {});
return;
}
const pending = this.#pending.get(message.sessionKey);
if (
pending &&
(pending.peerId !== peerId || pending.token !== message.token || pending.child.pid !== message.pid)
) {
return;
}
if (existing) return;
let resolveTerminated!: (error: Error | undefined) => void;
const terminated = new Promise<Error | undefined>((resolve) => {
resolveTerminated = resolve;
});
const worker: WorkerRecord = {
peerId,
metadata: message.metadata,
pid: message.pid,
terminated,
resolveTerminated,
expectedStop: false,
};
this.#workersBySession.set(message.sessionKey, worker);
this.#workersByPeer.set(peerId, worker);
this.workerPids.set(message.sessionId, message.pid);
if (pending) {
this.#pending.delete(message.sessionKey);
clearTimeout(pending.timer);
pending.resolve(worker);
}
}
#childExited(pending: PendingLaunch, code: number | null, signal: NodeJS.Signals | null): void {
if (this.#pending.get(pending.sessionKey) === pending) {
this.#failPending(
pending.sessionKey,
new Error(`Session worker exited before readiness (${signal ?? code ?? "unknown"})`),
);
return;
}
const worker = this.#workersByPeer.get(pending.peerId);
if (worker) {
this.#removeWorker(
worker,
worker.expectedStop
? undefined
: new Error(`Session worker ${worker.metadata.id} exited unexpectedly (${signal ?? code ?? "unknown"})`),
);
}
}
#failPending(sessionKey: string, error: Error): void {
const pending = this.#pending.get(sessionKey);
if (!pending) return;
this.#pending.delete(sessionKey);
clearTimeout(pending.timer);
if (pending.child.exitCode === null && pending.child.signalCode === null) pending.child.kill("SIGKILL");
pending.reject(error);
}
#removeWorker(worker: WorkerRecord, error: Error | undefined): void {
if (this.#workersByPeer.get(worker.peerId) !== worker) return;
this.#workersByPeer.delete(worker.peerId);
this.#workersBySession.delete(worker.metadata.path);
if (this.workerPids.get(worker.metadata.id) === worker.pid) this.workerPids.delete(worker.metadata.id);
worker.resolveTerminated(error);
}
#pendingPeer(peerId: string): PendingLaunch | undefined {
return [...this.#pending.values()].find((pending) => pending.peerId === peerId);
}
#markDiscovered(peerId: string): void {
if (!this.#discoveryPeers?.delete(peerId) || this.#discoveryPeers.size !== 0) return;
this.#resolveDiscovery?.();
}
#detachState(): void {
this.#removeListener();
this.#pending.clear();
this.#workersByPeer.clear();
this.#workersBySession.clear();
this.workerPids.clear();
this.#discoveryPeers = undefined;
this.#resolveDiscovery?.();
this.#resolveDiscovery = undefined;
}
}

View file

@ -0,0 +1,282 @@
import { spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import { createConnection, type Socket } from "node:net";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import Type, { type Static } from "typebox";
import { Check } from "typebox/value";
const COORDINATOR_PROTOCOL_VERSION = 1;
const COORDINATOR_START_TIMEOUT_MS = 10_000;
const COORDINATOR_RETRY_MS = 10;
const MAX_CONTROL_LINE_BYTES = 1024 * 1024;
const CoordinatorMessageSchema = Type.Union([
Type.Object({
type: Type.Literal("server_registered"),
serverConnectionId: Type.String(),
peers: Type.Array(Type.String()),
}),
Type.Object({ type: Type.Literal("server_replaced") }),
Type.Object({ type: Type.Literal("peer_connected"), peerId: Type.String() }),
Type.Object({ type: Type.Literal("peer_disconnected"), peerId: Type.String() }),
Type.Object({ type: Type.Literal("message"), from: Type.String(), payload: Type.Unknown() }),
]);
type CoordinatorMessage = Static<typeof CoordinatorMessageSchema>;
export type CoordinatorServerEvent =
| { readonly type: "peer_connected"; readonly peerId: string }
| { readonly type: "peer_disconnected"; readonly peerId: string }
| { readonly type: "message"; readonly from: string; readonly payload: unknown };
export interface CoordinatorServerOptions {
readonly controlPath: string;
readonly endpoint: string;
readonly serverConnectionId?: string;
}
/** The server-side endpoint of the coordinator's intentionally opaque message router. */
export class CoordinatorServer {
readonly serverConnectionId: string;
readonly replaced: Promise<void>;
readonly peerIds = new Set<string>();
readonly #controlPath: string;
readonly #endpoint: string;
readonly #listeners = new Set<(event: CoordinatorServerEvent) => void>();
#socket?: Socket;
#registered = false;
#closed = false;
#replacedValue = false;
#resolveRegistered?: () => void;
#rejectRegistered?: (error: Error) => void;
#resolveReplaced!: () => void;
constructor(options: CoordinatorServerOptions) {
this.serverConnectionId = options.serverConnectionId ?? randomUUID();
this.#controlPath = options.controlPath;
this.#endpoint = options.endpoint;
this.replaced = new Promise((resolve) => {
this.#resolveReplaced = resolve;
});
}
get controlPath(): string {
return this.#controlPath;
}
get wasReplaced(): boolean {
return this.#replacedValue;
}
onEvent(listener: (event: CoordinatorServerEvent) => void): () => void {
this.#listeners.add(listener);
return () => this.#listeners.delete(listener);
}
async connect(): Promise<void> {
if (this.#socket) throw new Error("Coordinator server is already connected");
const socket = await connectSocket(this.#controlPath);
this.#socket = socket;
attachJsonLineReader(socket, (message) => this.#handleMessage(message));
const registered = new Promise<void>((resolve, reject) => {
this.#resolveRegistered = resolve;
this.#rejectRegistered = reject;
});
socket.once("close", () => this.#disconnected(new Error("Coordinator connection closed")));
socket.once("error", (error) => this.#disconnected(error));
await writeJsonLine(socket, {
type: "register_server",
protocol: COORDINATOR_PROTOCOL_VERSION,
serverConnectionId: this.serverConnectionId,
endpoint: this.#endpoint,
});
await registered;
}
send(peerId: string, payload: unknown): Promise<void> {
return this.#write({ type: "send", to: peerId, payload });
}
broadcast(payload: unknown): Promise<void> {
return this.#write({ type: "broadcast", payload });
}
close(): void {
if (this.#closed) return;
this.#closed = true;
this.#socket?.destroy();
this.#socket = undefined;
this.peerIds.clear();
this.#rejectRegistered?.(new Error("Coordinator server closed"));
this.#resolveRegistered = undefined;
this.#rejectRegistered = undefined;
}
async #write(message: unknown): Promise<void> {
if (!this.#registered || !this.#socket || this.#closed) throw new Error("Coordinator server is not connected");
await writeJsonLine(this.#socket, message);
}
#handleMessage(value: unknown): void {
if (!Check(CoordinatorMessageSchema, value)) {
this.#socket?.destroy(new Error("Coordinator sent an invalid message"));
return;
}
const message: CoordinatorMessage = value;
if (message.type === "server_registered") {
if (message.serverConnectionId !== this.serverConnectionId) {
this.#socket?.destroy(new Error("Coordinator returned an invalid server registration"));
return;
}
for (const peerId of message.peers) this.peerIds.add(peerId);
this.#registered = true;
this.#resolveRegistered?.();
this.#resolveRegistered = undefined;
this.#rejectRegistered = undefined;
return;
}
if (message.type === "server_replaced") {
this.#markReplaced();
return;
}
if (message.type === "peer_connected") {
this.peerIds.add(message.peerId);
this.#emit({ type: "peer_connected", peerId: message.peerId });
return;
}
if (message.type === "peer_disconnected") {
this.peerIds.delete(message.peerId);
this.#emit({ type: "peer_disconnected", peerId: message.peerId });
return;
}
if (message.type === "message") {
this.#emit({ type: "message", from: message.from, payload: message.payload });
return;
}
this.#socket?.destroy(new Error("Coordinator sent an unsupported message"));
}
#emit(event: CoordinatorServerEvent): void {
for (const listener of this.#listeners) listener(event);
}
#disconnected(error: Error): void {
this.#socket = undefined;
if (this.#closed) return;
this.#rejectRegistered?.(error);
this.#resolveRegistered = undefined;
this.#rejectRegistered = undefined;
this.#markReplaced();
}
#markReplaced(): void {
if (this.#replacedValue) return;
this.#replacedValue = true;
this.#resolveReplaced();
}
}
export function getExperimentalCoordinatorControlPath(serverId: string, socketDirectory: string): string {
return join(socketDirectory, `.c-${serverId.slice(0, 12)}.sock`);
}
export function getExperimentalServerPath(socketDirectory: string): string {
return join(socketDirectory, `.s-${randomUUID().slice(0, 12)}.sock`);
}
export interface CoordinatorStartupLease {
close(): void;
}
export async function ensureExperimentalCoordinator(
publicPath: string,
controlPath: string,
coordinatorUrl = new URL(
import.meta.url.endsWith(".js") ? "coordinator-process.js" : "coordinator-process.ts",
import.meta.url,
),
): Promise<CoordinatorStartupLease> {
const existing = await tryConnect(controlPath);
if (existing) return { close: () => existing.destroy() };
const args = [
...(coordinatorUrl.pathname.endsWith(".ts") ? ["--import", "tsx"] : []),
fileURLToPath(coordinatorUrl),
publicPath,
controlPath,
];
const child = spawn(process.execPath, args, {
cwd: process.cwd(),
detached: true,
env: process.env,
stdio: "ignore",
windowsHide: true,
});
child.unref();
const deadline = Date.now() + COORDINATOR_START_TIMEOUT_MS;
while (true) {
const socket = await tryConnect(controlPath);
if (socket) return { close: () => socket.destroy() };
if (child.exitCode !== null || child.signalCode !== null) throw new Error("Coordinator exited during startup");
if (Date.now() >= deadline) throw new Error("Timed out waiting for coordinator startup");
await new Promise<void>((resolve) => setTimeout(resolve, COORDINATOR_RETRY_MS));
}
}
function connectSocket(path: string): Promise<Socket> {
return new Promise((resolve, reject) => {
const socket = createConnection(path);
const onConnect = (): void => {
socket.off("error", onError);
resolve(socket);
};
const onError = (error: Error): void => {
socket.off("connect", onConnect);
reject(error);
};
socket.once("connect", onConnect);
socket.once("error", onError);
});
}
async function tryConnect(path: string): Promise<Socket | undefined> {
try {
return await connectSocket(path);
} catch (error) {
const code = error instanceof Error && "code" in error ? error.code : undefined;
if (code === "ENOENT" || code === "ECONNREFUSED") return undefined;
throw error;
}
}
function attachJsonLineReader(socket: Socket, onMessage: (message: unknown) => void): void {
let buffered = "";
socket.setEncoding("utf8");
socket.on("data", (chunk: string) => {
buffered += chunk;
if (Buffer.byteLength(buffered) > MAX_CONTROL_LINE_BYTES) {
socket.destroy(new Error("Coordinator message is too large"));
return;
}
while (true) {
const newline = buffered.indexOf("\n");
if (newline === -1) return;
const line = buffered.slice(0, newline);
buffered = buffered.slice(newline + 1);
try {
onMessage(JSON.parse(line));
} catch {
socket.destroy(new Error("Coordinator sent invalid JSON"));
return;
}
}
});
}
function writeJsonLine(socket: Socket, message: unknown): Promise<void> {
return new Promise((resolve, reject) => {
socket.write(`${JSON.stringify(message)}\n`, (error) => {
if (error) reject(error);
else resolve();
});
});
}

View file

@ -0,0 +1,329 @@
import { chmod, lstat, unlink } from "node:fs/promises";
import { createConnection, createServer, type Server, type Socket } from "node:net";
// This process is intentionally a transport shim. It depends only on Node
// built-ins and never interprets Pi, session, worker, or lifecycle payloads.
const COORDINATOR_PROTOCOL_VERSION = 1;
const MAX_CONTROL_LINE_BYTES = 1024 * 1024;
const EMPTY_STARTUP_GRACE_MS = 30_000;
const EMPTY_SHUTDOWN_GRACE_MS = 250;
interface ServerPeer {
readonly serverConnectionId: string;
readonly endpoint: string;
readonly socket: Socket;
}
interface RoutedPeer {
readonly peerId: string;
readonly socket: Socket;
}
interface ControlMessage {
readonly type: string;
readonly protocol?: unknown;
readonly serverConnectionId?: unknown;
readonly endpoint?: unknown;
readonly peerId?: unknown;
readonly to?: unknown;
readonly payload?: unknown;
}
const publicPath = process.argv[2];
const controlPath = process.argv[3];
if (!publicPath || !controlPath) throw new Error("Coordinator requires public and control socket paths");
const peers = new Map<string, RoutedPeer>();
const controlConnections = new Set<Socket>();
const publicConnections = new Map<Socket, Socket>();
let currentServer: ServerPeer | undefined;
let shuttingDown = false;
let emptyTimer: NodeJS.Timeout | undefined;
const publicServer = createServer((socket) => acceptPublicConnection(socket));
const controlServer = createServer((socket) => acceptControlConnection(socket));
async function main(): Promise<void> {
await removeStaleSocket(controlPath);
await removeStaleSocket(publicPath);
try {
await listen(controlServer, controlPath);
await restrictSocket(controlPath);
await listen(publicServer, publicPath);
await restrictSocket(publicPath);
} catch (error) {
await Promise.all([cleanupSocket(controlPath), cleanupSocket(publicPath)]);
throw error;
}
scheduleEmptyShutdown(EMPTY_STARTUP_GRACE_MS);
}
function acceptControlConnection(socket: Socket): void {
if (shuttingDown) {
socket.destroy();
return;
}
controlConnections.add(socket);
let server: ServerPeer | undefined;
let peer: RoutedPeer | undefined;
attachJsonLineReader(socket, (message) => {
if (!server && !peer) {
if (message.type === "register_server") {
server = registerServer(socket, message);
return;
}
if (message.type === "register_peer") {
peer = registerPeer(socket, message);
return;
}
throw new Error("Coordinator connection did not register a role");
}
if (server) {
if (server === currentServer) handleRoutedMessage("server", message);
return;
}
handleRoutedMessage(peer!.peerId, message);
});
const disconnect = (): void => {
controlConnections.delete(socket);
if (server && currentServer === server) currentServer = undefined;
if (peer && peers.get(peer.peerId) === peer) {
peers.delete(peer.peerId);
if (currentServer) writeJsonLine(currentServer.socket, { type: "peer_disconnected", peerId: peer.peerId });
}
checkEmpty();
};
socket.once("close", disconnect);
socket.once("error", () => socket.destroy());
}
function registerServer(socket: Socket, message: ControlMessage): ServerPeer {
if (message.protocol !== COORDINATOR_PROTOCOL_VERSION) throw new Error("Unsupported coordinator protocol");
if (typeof message.serverConnectionId !== "string" || message.serverConnectionId.length === 0) {
throw new Error("Coordinator serverConnectionId must be a string");
}
if (typeof message.endpoint !== "string" || message.endpoint.length === 0) {
throw new Error("Coordinator endpoint must be a string");
}
const { serverConnectionId, endpoint } = message;
const server = { serverConnectionId, endpoint, socket };
const previous = currentServer;
currentServer = server;
cancelEmptyShutdown();
writeJsonLine(socket, {
type: "server_registered",
serverConnectionId,
peers: [...peers.keys()],
});
if (previous && previous !== server) {
closePublicConnections();
writeJsonLine(previous.socket, { type: "server_replaced" });
}
return server;
}
function registerPeer(socket: Socket, message: ControlMessage): RoutedPeer {
if (message.protocol !== COORDINATOR_PROTOCOL_VERSION) throw new Error("Unsupported coordinator protocol");
if (typeof message.peerId !== "string" || message.peerId.length === 0) {
throw new Error("Coordinator peerId must be a string");
}
const { peerId } = message;
if (peerId === "server" || peers.has(peerId)) throw new Error(`Coordinator peer is already connected: ${peerId}`);
const peer = { peerId, socket };
peers.set(peerId, peer);
cancelEmptyShutdown();
writeJsonLine(socket, { type: "peer_registered", peerId });
if (currentServer) {
writeJsonLine(currentServer.socket, { type: "peer_connected", peerId });
}
return peer;
}
function handleRoutedMessage(from: string, message: ControlMessage): void {
if (message.type === "send") {
if (typeof message.to !== "string" || message.to.length === 0) {
throw new Error("Coordinator message target must be a string");
}
const { to } = message;
const target = to === "server" ? currentServer?.socket : peers.get(to)?.socket;
if (target) writeJsonLine(target, { type: "message", from, payload: message.payload });
return;
}
if (message.type === "broadcast") {
if (from !== "server") throw new Error("Only the current server may broadcast");
for (const peer of peers.values()) {
writeJsonLine(peer.socket, { type: "message", from, payload: message.payload });
}
return;
}
throw new Error(`Unknown coordinator routing message: ${message.type}`);
}
function acceptPublicConnection(client: Socket): void {
if (shuttingDown || !currentServer) {
client.destroy();
return;
}
cancelEmptyShutdown();
const upstream = createConnection(currentServer.endpoint);
publicConnections.set(client, upstream);
let finalized = false;
const finalize = (): void => {
if (finalized) return;
finalized = true;
publicConnections.delete(client);
checkEmpty();
};
upstream.once("connect", () => {
client.pipe(upstream);
upstream.pipe(client);
});
upstream.once("error", () => client.destroy());
upstream.once("close", () => {
client.destroy();
finalize();
});
client.once("error", () => upstream.destroy());
client.once("close", () => {
upstream.destroy();
finalize();
});
}
function closePublicConnections(): void {
for (const [client, upstream] of publicConnections) {
client.destroy();
upstream.destroy();
}
publicConnections.clear();
}
function checkEmpty(): void {
if (shuttingDown || currentServer || peers.size > 0 || publicConnections.size > 0 || controlConnections.size > 0) {
cancelEmptyShutdown();
return;
}
scheduleEmptyShutdown(EMPTY_SHUTDOWN_GRACE_MS);
}
function scheduleEmptyShutdown(delayMs: number): void {
if (emptyTimer || shuttingDown) return;
emptyTimer = setTimeout(() => {
emptyTimer = undefined;
if (!currentServer && peers.size === 0 && publicConnections.size === 0 && controlConnections.size === 0) {
void shutdownCoordinator();
}
}, delayMs);
emptyTimer.unref();
}
function cancelEmptyShutdown(): void {
if (!emptyTimer) return;
clearTimeout(emptyTimer);
emptyTimer = undefined;
}
async function shutdownCoordinator(): Promise<void> {
if (shuttingDown) return;
shuttingDown = true;
cancelEmptyShutdown();
closePublicConnections();
for (const connection of controlConnections) connection.destroy();
await Promise.all([closeServer(publicServer), closeServer(controlServer)]);
await Promise.all([cleanupSocket(publicPath), cleanupSocket(controlPath)]);
process.exit(0);
}
function writeJsonLine(socket: Socket, message: unknown): void {
if (!socket.destroyed) socket.write(`${JSON.stringify(message)}\n`);
}
function attachJsonLineReader(socket: Socket, onMessage: (message: ControlMessage) => void): void {
let buffered = "";
socket.setEncoding("utf8");
socket.on("data", (chunk: string) => {
buffered += chunk;
if (Buffer.byteLength(buffered) > MAX_CONTROL_LINE_BYTES) {
socket.destroy(new Error("Coordinator message is too large"));
return;
}
while (true) {
const newline = buffered.indexOf("\n");
if (newline === -1) return;
const line = buffered.slice(0, newline);
buffered = buffered.slice(newline + 1);
try {
const message = JSON.parse(line) as ControlMessage | null;
if (typeof message?.type !== "string") throw new Error("Coordinator message must have a type");
onMessage(message);
} catch (error) {
socket.destroy(error instanceof Error ? error : new Error(String(error)));
return;
}
}
});
}
function listen(server: Server, path: string): Promise<void> {
return new Promise((resolve, reject) => {
const onError = (error: Error): void => {
server.off("listening", onListening);
reject(error);
};
const onListening = (): void => {
server.off("error", onError);
resolve();
};
server.once("error", onError);
server.once("listening", onListening);
server.listen(path);
});
}
function closeServer(server: Server): Promise<void> {
if (!server.listening) return Promise.resolve();
return new Promise((resolve) => server.close(() => resolve()));
}
async function restrictSocket(path: string): Promise<void> {
if (process.platform !== "win32") await chmod(path, 0o600);
}
async function removeStaleSocket(path: string): Promise<void> {
try {
const stats = await lstat(path);
if (!stats.isSocket()) throw new Error(`Coordinator path is not a socket: ${path}`);
} catch (error) {
if (error instanceof Error && "code" in error && error.code === "ENOENT") return;
throw error;
}
const live = await new Promise<boolean>((resolve, reject) => {
const socket = createConnection(path);
socket.once("connect", () => {
socket.destroy();
resolve(true);
});
socket.once("error", (error: NodeJS.ErrnoException) => {
socket.destroy();
if (error.code === "ECONNREFUSED" || error.code === "ENOENT") resolve(false);
else reject(error);
});
});
if (live) throw new Error(`Coordinator socket is already active: ${path}`);
await cleanupSocket(path);
}
async function cleanupSocket(path: string): Promise<void> {
if (process.platform === "win32") return;
await unlink(path).catch((error: unknown) => {
if (!(error instanceof Error && "code" in error && error.code === "ENOENT")) throw error;
});
}
process.once("SIGINT", () => void shutdownCoordinator());
process.once("SIGTERM", () => void shutdownCoordinator());
void main().catch((error: unknown) => {
console.error(error);
process.exit(1);
});

View file

@ -2,23 +2,27 @@ import { randomUUID } from "node:crypto";
import { chmod, lstat, mkdir } from "node:fs/promises";
import { homedir } from "node:os";
import { basename, join } from "node:path";
import { setTimeout as delay } from "node:timers/promises";
import { type JsonlSessionMetadata, JsonlSessionRepo } from "@earendil-works/pi-agent-core";
import { NodeExecutionEnv } from "@earendil-works/pi-agent-core/node";
import { PiClient } from "@earendil-works/pi-client";
import { requestServerDrain } from "@earendil-works/pi-client/control";
import { createUnixTransportFactory, discoverUnixServers, type UnixServerRoute } from "@earendil-works/pi-client/unix";
import { isServerId } from "@earendil-works/pi-protocol";
import type { PiServer, PiServerHost } from "@earendil-works/pi-server";
import type { HostedHarnessHandle, PiServer, PiServerHost } from "@earendil-works/pi-server";
import { createUnixServer, getUnixSocketPath } from "@earendil-works/pi-server/unix";
import { getAgentDir } from "../../config.ts";
import { resolvePath } from "../../utils/paths.ts";
import type { ClientCommand } from "./commands/client.ts";
import { CoordinatedSessionWorkers } from "./coordinated-session-workers.ts";
import {
CoordinatorServer,
type CoordinatorStartupLease,
ensureExperimentalCoordinator,
getExperimentalCoordinatorControlPath,
getExperimentalServerPath,
} from "./coordinator-client.ts";
import { acquireExperimentalServerProfile } from "./server-profile.ts";
import { startExperimentalSessionWorker } from "./session-worker.ts";
const SOCKET_RELEASE_TIMEOUT_MS = 10_000;
const SOCKET_RELEASE_POLL_MS = 10;
const EXPERIMENTAL_SOCKET_ROOT = "/tmp";
export interface ExperimentalServer {
@ -47,7 +51,7 @@ export interface StartExperimentalServerOptions {
readonly sessionDir?: string;
}
export interface StartExperimentalServerGenerationOptions {
export interface StartExperimentalCoordinatedServerOptions {
/** Persistent profile directory. Defaults to ~/.pi/server. */
readonly directory?: string;
/** Physical socket directory. Defaults to a short, private per-user runtime directory. */
@ -65,18 +69,21 @@ export function resolveExperimentalSessionDirectory(sessionDir?: string): string
return resolvePath(sessionDir ?? join(getAgentDir(), "experimental", "sessions"));
}
interface ExperimentalWorkerController {
readonly workerPids: ReadonlyMap<string, number>;
readonly trackedSessions: readonly JsonlSessionMetadata[];
createHarness(metadata: JsonlSessionMetadata): Promise<HostedHarnessHandle>;
}
/** Start the experimental durable list-and-attach server composition. */
export async function startExperimentalServer(
options: StartExperimentalServerOptions = {},
): Promise<ExperimentalServer> {
const serverId = options.serverId ?? randomUUID();
const sessionDir = resolveExperimentalSessionDirectory(options.sessionDir);
const executionEnv = new NodeExecutionEnv({ cwd: process.cwd() });
const repo = new JsonlSessionRepo({ fileSystem: executionEnv, sessionsRoot: sessionDir });
const workerPids = new Map<string, number>();
const host: PiServerHost<JsonlSessionMetadata> = {
sessions: repo,
const controller: ExperimentalWorkerController = {
workerPids,
trackedSessions: [],
createHarness: async (metadata) => {
const sessionId = metadata.id;
const worker = await startExperimentalSessionWorker(metadata, { sessionDir });
@ -96,6 +103,27 @@ export async function startExperimentalServer(
};
},
};
return startExperimentalServerBackend(options, controller);
}
async function startExperimentalServerBackend(
options: StartExperimentalServerOptions,
workers: ExperimentalWorkerController,
): Promise<ExperimentalServer> {
const serverId = options.serverId ?? randomUUID();
const sessionDir = resolveExperimentalSessionDirectory(options.sessionDir);
const executionEnv = new NodeExecutionEnv({ cwd: process.cwd() });
const repo = new JsonlSessionRepo({ fileSystem: executionEnv, sessionsRoot: sessionDir });
const host: PiServerHost<JsonlSessionMetadata> = {
sessions: {
list: async () => {
const sessions = new Map((await repo.list()).map((metadata) => [metadata.path, metadata]));
for (const metadata of workers.trackedSessions) sessions.set(metadata.path, metadata);
return [...sessions.values()];
},
},
createHarness: (metadata) => workers.createHarness(metadata),
};
let socketPath = options.path;
if (socketPath === undefined) {
const socketDirectory = options.directory ?? getExperimentalSocketDirectory();
@ -137,72 +165,100 @@ export async function startExperimentalServer(
sessionDir,
socketPath,
server,
workerPids,
workerPids: workers.workerPids,
closed,
close() {
if (closePromise === undefined)
closePromise = server.close().then(
() => closed,
() => closed,
);
closePromise ??= server.close().then(
() => closed,
() => closed,
);
return closePromise;
},
};
}
/** Start the experimental local server, replacing an existing generation for the same profile. */
export async function startExperimentalServerGeneration(
options: StartExperimentalServerGenerationOptions = {},
/** Start a replaceable experimental server behind the stable coordinator endpoint. */
export async function startExperimentalCoordinatedServer(
options: StartExperimentalCoordinatedServerOptions = {},
): Promise<ExperimentalServer> {
const directory = options.directory ?? join(homedir(), ".pi", "server");
const socketDirectory = options.socketDirectory ?? getExperimentalSocketDirectory();
const { serverId, release } = await acquireExperimentalServerProfile(directory);
let runtime: ExperimentalServer;
let backend: ExperimentalServer | undefined;
let coordinator: CoordinatorServer | undefined;
let startupLease: CoordinatorStartupLease | undefined;
let workers: CoordinatedSessionWorkers | undefined;
let released = false;
try {
await ensurePrivateSocketDirectory(socketDirectory);
const socketPath = getUnixSocketPath(serverId, socketDirectory);
await drainExistingGeneration(serverId, socketPath);
runtime = await startExperimentalServer({
directory,
path: socketPath,
serverId,
sessionDir: options.sessionDir,
});
} catch (error) {
try {
await release();
} catch (releaseError) {
throw new AggregateError([error, releaseError], "Server generation failed and launcher lock release failed");
}
throw error;
}
const controlPath = getExperimentalCoordinatorControlPath(serverId, socketDirectory);
const serverPath = getExperimentalServerPath(socketDirectory);
startupLease = await ensureExperimentalCoordinator(socketPath, controlPath);
coordinator = new CoordinatorServer({ controlPath, endpoint: serverPath });
const sessionDir = resolveExperimentalSessionDirectory(options.sessionDir);
workers = new CoordinatedSessionWorkers(coordinator, sessionDir);
backend = await startExperimentalServerBackend(
{ path: serverPath, serverId, sessionDir: options.sessionDir },
workers,
);
await coordinator.connect();
startupLease.close();
startupLease = undefined;
await workers.discover(coordinator.peerIds);
try {
const activeBackend = backend;
const activeCoordinator = coordinator;
const activeWorkers = workers;
void activeCoordinator.replaced
.then(async () => {
activeWorkers.detach();
await activeBackend.close();
})
.finally(() => activeCoordinator.close())
.catch(() => {});
let closePromise: Promise<void> | undefined;
const runtime: ExperimentalServer = {
serverId,
sessionDir: activeBackend.sessionDir,
socketPath,
server: activeBackend.server,
workerPids: activeWorkers.workerPids,
closed: activeBackend.closed,
close() {
closePromise ??= (async () => {
try {
await activeBackend.close();
} finally {
try {
if (activeCoordinator.wasReplaced) activeWorkers.detach();
else await activeWorkers.shutdown();
} finally {
activeCoordinator.close();
}
}
})();
return closePromise;
},
};
released = true;
await release();
return runtime;
} catch (error) {
try {
await runtime.close();
} catch (closeError) {
throw new AggregateError([error, closeError], "Launcher lock release and server cleanup failed");
startupLease?.close();
if (coordinator?.wasReplaced) workers?.detach();
const cleanup = await Promise.allSettled([
backend?.close(),
coordinator?.wasReplaced ? undefined : workers?.shutdown(),
Promise.resolve(coordinator?.close()),
released ? undefined : release(),
]);
const cleanupErrors = cleanup.flatMap((result) => (result.status === "rejected" ? [result.reason] : []));
if (cleanupErrors.length > 0) {
throw new AggregateError([error, ...cleanupErrors], "Coordinated server startup and cleanup failed");
}
throw error;
}
return runtime;
}
async function drainExistingGeneration(serverId: string, socketPath: string): Promise<void> {
try {
await requestServerDrain({
serverId,
transportFactory: createUnixTransportFactory({ path: socketPath }),
});
await waitForSocketRelease(socketPath);
} catch (error) {
// At this launcher boundary, a missing or refused stable socket proves that
// no generation accepted the drain request. All ambiguous failures propagate.
if (hasErrorCode(error, "ENOENT") || hasErrorCode(error, "ECONNREFUSED")) return;
throw error;
}
}
/** Discover local servers, then list sessions or attach to one selected session. */
@ -269,20 +325,6 @@ function routeFromExplicitPath(path: string): UnixServerRoute {
return { serverId, path };
}
async function waitForSocketRelease(path: string): Promise<void> {
const deadline = Date.now() + SOCKET_RELEASE_TIMEOUT_MS;
while (true) {
try {
await lstat(path);
} catch (error) {
if (hasErrorCode(error, "ENOENT")) return;
throw error;
}
if (Date.now() >= deadline) throw new Error(`Timed out waiting for server socket to close: ${path}`);
await delay(SOCKET_RELEASE_POLL_MS);
}
}
function getExperimentalSocketDirectory(): string {
if (process.platform === "win32" || typeof process.getuid !== "function") {
throw new Error("Experimental Unix server transport requires a POSIX user ID");
@ -299,14 +341,3 @@ async function ensurePrivateSocketDirectory(directory: string): Promise<void> {
throw new Error(`Unix socket directory is not owned by the current user: ${directory}`);
await chmod(directory, 0o700);
}
function hasErrorCode(error: unknown, code: string): boolean {
let current = error;
const seen = new Set<unknown>();
while (current instanceof Error && !seen.has(current)) {
seen.add(current);
if ("code" in current && current.code === code) return true;
current = current.cause;
}
return false;
}

View file

@ -1,3 +1,4 @@
import { createConnection, type Socket } from "node:net";
import { isAbsolute } from "node:path";
import {
AgentHarness,
@ -7,22 +8,163 @@ import {
type Session,
} from "@earendil-works/pi-agent-core";
import { NodeExecutionEnv } from "@earendil-works/pi-agent-core/node";
import lockfile from "proper-lockfile";
import Type, { type Static } from "typebox";
import { Check } from "typebox/value";
import { ModelRuntime } from "../../core/model-runtime.ts";
import {
SESSION_WORKER_CONTROL_ADDRESS_ENV,
SESSION_WORKER_CONTROL_TOKEN_ENV,
SESSION_WORKER_COORDINATED_ENV,
SESSION_WORKER_PEER_ID_ENV,
SESSION_WORKER_SESSION_KEY_ENV,
} from "./session-worker.ts";
// Prototype-only process control. Agent operations will use the shared
// transport-independent protocol rather than additional ad hoc IPC messages.
// Process lifecycle control only. Agent operations use the shared
// transport-independent protocol rather than this channel.
export type SessionWorkerCommand = { type: "shutdown" };
export type SessionWorkerEvent =
| { type: "ready"; sessionId: string; pid: number }
| { type: "failed"; message: string };
function send(event: SessionWorkerEvent): Promise<void> {
return new Promise((resolve, reject) => {
if (!process.send) {
reject(new Error("Session worker requires an IPC channel"));
const WorkerMetadataSchema = Type.Object({
id: Type.String({ minLength: 1 }),
createdAt: Type.Integer(),
storageVersion: Type.Integer(),
cwd: Type.String(),
path: Type.String(),
modifiedAt: Type.Number(),
parentSessionId: Type.Optional(Type.String()),
legacyParentSessionPath: Type.Optional(Type.String()),
});
const CoordinatorInputSchema = Type.Union([
Type.Object({ type: Type.Literal("peer_registered"), peerId: Type.String() }),
Type.Object({ type: Type.Literal("message"), from: Type.Literal("server"), payload: Type.Unknown() }),
]);
const WorkerCommandSchema = Type.Union([
Type.Object({ type: Type.Literal("shutdown") }),
Type.Object({ type: Type.Literal("discover_workers") }),
]);
type CoordinatorInput = Static<typeof CoordinatorInputSchema>;
type WorkerCommand = Static<typeof WorkerCommandSchema>;
type SessionWorkerControlEvent =
| {
type: "register_worker";
protocol: 1;
token: string;
sessionKey: string;
sessionId: string;
pid: number;
}
| {
type: "worker_ready";
token: string;
sessionKey: string;
sessionId: string;
pid: number;
metadata: JsonlSessionMetadata;
}
| { type: "worker_failed"; token: string; sessionKey: string; message: string };
interface WorkerControl {
readonly coordinated: boolean;
readonly messages: AsyncIterable<unknown>;
readonly socket: Socket;
send(event: SessionWorkerControlEvent): Promise<void>;
}
let failureControl: WorkerControl | undefined;
async function connectControl(): Promise<WorkerControl> {
const address = process.env[SESSION_WORKER_CONTROL_ADDRESS_ENV];
const token = process.env[SESSION_WORKER_CONTROL_TOKEN_ENV];
const encodedSessionKey = process.env[SESSION_WORKER_SESSION_KEY_ENV];
if (!address || !token || !encodedSessionKey) throw new Error("Session worker requires a control address");
const socket = createConnection(address);
await new Promise<void>((resolve, reject) => {
socket.once("connect", resolve);
socket.once("error", reject);
});
const messages = createJsonLineMessages(socket);
const coordinated = process.env[SESSION_WORKER_COORDINATED_ENV] === "1";
if (coordinated) {
const peerId = process.env[SESSION_WORKER_PEER_ID_ENV];
if (!peerId) throw new Error("Coordinated session worker requires a peer ID");
await writeJsonLine(socket, { type: "register_peer", protocol: 1, peerId });
const registered = await messages[Symbol.asyncIterator]().next();
if (
registered.done ||
!Check(CoordinatorInputSchema, registered.value) ||
registered.value.type !== "peer_registered"
) {
throw new Error("Coordinator rejected the session worker registration");
}
}
return {
coordinated,
messages,
socket,
send: (event) => writeJsonLine(socket, coordinated ? { type: "send", to: "server", payload: event } : event),
};
}
async function readCommands(control: WorkerControl, onShutdown: () => void, onDiscovery: () => void): Promise<void> {
for await (const value of control.messages) {
if (!control.coordinated) {
if (Check(WorkerCommandSchema, value) && value.type === "shutdown") onShutdown();
continue;
}
if (!Check(CoordinatorInputSchema, value)) {
control.socket.destroy(new Error("Coordinator sent an invalid worker message"));
return;
}
process.send(event, (error) => {
const message: CoordinatorInput = value;
if (message.type !== "message" || !Check(WorkerCommandSchema, message.payload)) continue;
const command: WorkerCommand = message.payload;
if (command.type === "shutdown") onShutdown();
else onDiscovery();
}
}
function createJsonLineMessages(socket: Socket): AsyncIterable<unknown> {
const queued: unknown[] = [];
const waiters: ((value: unknown) => void)[] = [];
let buffered = "";
socket.setEncoding("utf8");
socket.on("data", (chunk: string) => {
buffered += chunk;
while (true) {
const newline = buffered.indexOf("\n");
if (newline === -1) return;
const line = buffered.slice(0, newline);
buffered = buffered.slice(newline + 1);
try {
const value: unknown = JSON.parse(line);
const waiter = waiters.shift();
if (waiter) waiter(value);
else queued.push(value);
} catch {
socket.destroy(new Error("Session worker received invalid control JSON"));
return;
}
}
});
return {
[Symbol.asyncIterator]() {
return {
next: async () => {
const value = queued.shift() ?? (await new Promise<unknown>((resolve) => waiters.push(resolve)));
return { done: false as const, value };
},
};
},
};
}
function writeJsonLine(socket: Socket, message: unknown): Promise<void> {
return new Promise((resolve, reject) => {
socket.write(`${JSON.stringify(message)}\n`, (error) => {
if (error) reject(error);
else resolve();
});
@ -37,27 +179,10 @@ function parseMetadata(value: string | undefined): JsonlSessionMetadata {
} catch (error) {
throw new Error("Session worker received invalid session metadata", { cause: error });
}
if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
if (!Check(WorkerMetadataSchema, parsed) || !isAbsolute(parsed.cwd) || !isAbsolute(parsed.path)) {
throw new Error("Session worker received invalid session metadata");
}
const candidate = parsed as Record<string, unknown>;
if (
typeof candidate.id !== "string" ||
candidate.id.length === 0 ||
!Number.isSafeInteger(candidate.createdAt) ||
!Number.isSafeInteger(candidate.storageVersion) ||
typeof candidate.cwd !== "string" ||
!isAbsolute(candidate.cwd) ||
typeof candidate.path !== "string" ||
!isAbsolute(candidate.path) ||
typeof candidate.modifiedAt !== "number" ||
!Number.isFinite(candidate.modifiedAt) ||
(candidate.parentSessionId !== undefined && typeof candidate.parentSessionId !== "string") ||
(candidate.legacyParentSessionPath !== undefined && typeof candidate.legacyParentSessionPath !== "string")
) {
throw new Error("Session worker received invalid session metadata");
}
return candidate as unknown as JsonlSessionMetadata;
return parsed;
}
async function closeResources(resources: {
@ -65,6 +190,7 @@ async function closeResources(resources: {
session?: Session<JsonlSessionMetadata>;
repo: JsonlSessionRepo;
executionEnv: NodeExecutionEnv;
releaseOwnership: () => Promise<void>;
}): Promise<void> {
const errors: unknown[] = [];
try {
@ -83,6 +209,11 @@ async function closeResources(resources: {
} catch (error) {
errors.push(error);
}
try {
await resources.releaseOwnership();
} catch (error) {
errors.push(error);
}
if (errors.length === 1) throw errors[0];
if (errors.length > 1) throw new AggregateError(errors, "Session worker cleanup failed");
}
@ -92,7 +223,27 @@ async function run(): Promise<void> {
if (!sessionDir || !isAbsolute(sessionDir)) throw new Error("Session worker requires an absolute session directory");
const metadata = parseMetadata(process.argv[3]);
const sessionId = metadata.id;
const control = await connectControl();
const token = process.env[SESSION_WORKER_CONTROL_TOKEN_ENV]!;
const sessionKey = Buffer.from(process.env[SESSION_WORKER_SESSION_KEY_ENV]!, "base64url").toString();
failureControl = control;
if (!control.coordinated) {
await control.send({
type: "register_worker",
protocol: 1,
token,
sessionKey,
sessionId,
pid: process.pid,
});
}
const releaseOwnership = await lockfile.lock(metadata.path, {
realpath: true,
stale: 2_000,
update: 1_000,
retries: { retries: 320, factor: 1, minTimeout: 25, maxTimeout: 25, maxRetryTime: 8_000 },
});
const executionEnv = new NodeExecutionEnv({ cwd: metadata.cwd });
const repo = new JsonlSessionRepo({ fileSystem: executionEnv, sessionsRoot: sessionDir });
let session: Session<JsonlSessionMetadata> | undefined;
@ -105,7 +256,7 @@ async function run(): Promise<void> {
({ harness } = await AgentHarness.create({ session, models: modelRuntime, model, tools: [], resources: {} }));
} catch (error) {
try {
await closeResources({ harness, session, repo, executionEnv });
await closeResources({ harness, session, repo, executionEnv, releaseOwnership });
} catch (cleanupError) {
throw new AggregateError([error, cleanupError], "Session worker startup and cleanup failed");
}
@ -114,7 +265,7 @@ async function run(): Promise<void> {
let closing: Promise<void> | undefined;
const close = (): Promise<void> => {
closing ??= closeResources({ harness, repo, executionEnv });
closing ??= closeResources({ harness, repo, executionEnv, releaseOwnership });
return closing;
};
const closeAndExit = (): void => {
@ -127,15 +278,29 @@ async function run(): Promise<void> {
);
};
process.on("message", (message: SessionWorkerCommand) => {
if (message?.type === "shutdown") closeAndExit();
});
process.once("disconnect", closeAndExit);
let ready = false;
const announce = (): void => {
if (!ready) return;
void control
.send({
type: "worker_ready",
token,
sessionKey,
sessionId,
pid: process.pid,
metadata,
})
.catch(() => closeAndExit());
};
void readCommands(control, closeAndExit, announce);
control.socket.once("close", closeAndExit);
control.socket.once("error", () => closeAndExit());
process.once("SIGTERM", closeAndExit);
process.once("SIGINT", closeAndExit);
try {
await send({ type: "ready", sessionId, pid: process.pid });
ready = true;
await control.send({ type: "worker_ready", token, sessionKey, sessionId, pid: process.pid, metadata });
} catch (error) {
try {
await close();
@ -148,6 +313,11 @@ async function run(): Promise<void> {
void run().catch(async (error: unknown) => {
const message = error instanceof Error ? error.message : String(error);
await send({ type: "failed", message }).catch(() => {});
const token = process.env[SESSION_WORKER_CONTROL_TOKEN_ENV];
const encodedSessionKey = process.env[SESSION_WORKER_SESSION_KEY_ENV];
if (token && encodedSessionKey) {
const sessionKey = Buffer.from(encodedSessionKey, "base64url").toString();
await failureControl?.send({ type: "worker_failed", token, sessionKey, message }).catch(() => {});
}
process.exit(1);
});

View file

@ -1,11 +1,35 @@
import { type ChildProcess, fork } from "node:child_process";
import { isAbsolute } from "node:path";
import { type ChildProcess, spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import { chmod, unlink } from "node:fs/promises";
import { createServer, type Server, type Socket } from "node:net";
import { tmpdir } from "node:os";
import { isAbsolute, join } from "node:path";
import { fileURLToPath } from "node:url";
import type { JsonlSessionMetadata } from "@earendil-works/pi-agent-core";
import Type, { type Static } from "typebox";
import { Check } from "typebox/value";
import type { SessionWorkerCommand, SessionWorkerEvent } from "./session-worker-process.ts";
const DEFAULT_STARTUP_TIMEOUT_MS = 5_000;
const DEFAULT_STARTUP_TIMEOUT_MS = 10_000;
const DEFAULT_SHUTDOWN_TIMEOUT_MS = 5_000;
const MAX_TIMER_DELAY_MS = 2_147_483_647;
const WorkerControlInputSchema = Type.Union([
Type.Object({ type: Type.Literal("worker_failed"), message: Type.String() }),
Type.Object({ type: Type.Literal("worker_ready"), sessionId: Type.String(), pid: Type.Integer({ minimum: 1 }) }),
Type.Object({
type: Type.Literal("register_worker"),
protocol: Type.Literal(1),
token: Type.String(),
sessionId: Type.String(),
pid: Type.Integer({ minimum: 1 }),
}),
]);
type WorkerControlInput = Static<typeof WorkerControlInputSchema>;
export const SESSION_WORKER_CONTROL_ADDRESS_ENV = "PI_SESSION_WORKER_CONTROL_ADDRESS";
export const SESSION_WORKER_CONTROL_TOKEN_ENV = "PI_SESSION_WORKER_CONTROL_TOKEN";
export const SESSION_WORKER_SESSION_KEY_ENV = "PI_SESSION_WORKER_SESSION_KEY_BASE64";
export const SESSION_WORKER_COORDINATED_ENV = "PI_SESSION_WORKER_COORDINATED";
export const SESSION_WORKER_PEER_ID_ENV = "PI_SESSION_WORKER_PEER_ID";
export interface ExperimentalSessionWorker {
readonly sessionId: string;
@ -14,6 +38,13 @@ export interface ExperimentalSessionWorker {
close(): Promise<void>;
}
export interface SessionWorkerLaunchSpec {
readonly command: string;
readonly args: readonly string[];
readonly cwd: string;
readonly env: Readonly<Record<string, string>>;
}
export interface StartExperimentalSessionWorkerOptions {
readonly sessionDir: string;
readonly startupTimeoutMs?: number;
@ -21,29 +52,74 @@ export interface StartExperimentalSessionWorkerOptions {
readonly workerUrl?: URL;
}
export function createExperimentalSessionWorkerLaunchSpec(
metadata: JsonlSessionMetadata,
options: Pick<StartExperimentalSessionWorkerOptions, "sessionDir" | "workerUrl">,
): SessionWorkerLaunchSpec {
if (!isAbsolute(options.sessionDir)) throw new TypeError("Session worker sessionDir must be absolute");
const workerUrl =
options.workerUrl ??
new URL(
import.meta.url.endsWith(".js") ? "session-worker-process.js" : "session-worker-process.ts",
import.meta.url,
);
return {
command: process.execPath,
args: [
...(workerUrl.pathname.endsWith(".ts") ? ["--import", "tsx"] : []),
fileURLToPath(workerUrl),
options.sessionDir,
JSON.stringify(metadata),
],
cwd: process.cwd(),
env: Object.fromEntries(
Object.entries(process.env).flatMap(([name, value]) => (value === undefined ? [] : [[name, value]])),
),
};
}
export async function startExperimentalSessionWorker(
metadata: JsonlSessionMetadata,
options: StartExperimentalSessionWorkerOptions,
): Promise<ExperimentalSessionWorker> {
if (!isAbsolute(options.sessionDir)) throw new TypeError("Session worker sessionDir must be absolute");
const sessionId = metadata.id;
const startupTimeoutMs = validateTimeout(options.startupTimeoutMs ?? DEFAULT_STARTUP_TIMEOUT_MS, "startupTimeoutMs");
const shutdownTimeoutMs = validateTimeout(
options.shutdownTimeoutMs ?? DEFAULT_SHUTDOWN_TIMEOUT_MS,
"shutdownTimeoutMs",
);
const workerUrl = options.workerUrl ?? defaultWorkerUrl();
const child = fork(workerUrl, [options.sessionDir, JSON.stringify(metadata)], {
execArgv: workerExecArgv(workerUrl),
stdio: ["ignore", "ignore", "inherit", "ipc"],
});
const launch = createExperimentalSessionWorkerLaunchSpec(metadata, options);
const token = randomUUID();
const controlId = randomUUID();
const controlAddress =
process.platform === "win32" ? `\\\\.\\pipe\\pi-worker-${controlId}` : join(tmpdir(), `pi-w-${controlId}.sock`);
const control = await createWorkerControlServer(controlAddress, token, metadata.id);
let child: ChildProcess;
try {
child = spawn(launch.command, launch.args, {
cwd: launch.cwd,
detached: true,
env: {
...launch.env,
[SESSION_WORKER_CONTROL_ADDRESS_ENV]: controlAddress,
[SESSION_WORKER_CONTROL_TOKEN_ENV]: token,
[SESSION_WORKER_SESSION_KEY_ENV]: Buffer.from(metadata.path).toString("base64url"),
},
stdio: "ignore",
windowsHide: true,
});
child.unref();
} catch (error) {
await control.close();
throw error;
}
try {
const ready = await waitForReady(child, sessionId, startupTimeoutMs);
return new ChildSessionWorker(child, ready.sessionId, ready.pid, shutdownTimeoutMs);
const ready = await waitForReady(child, control, metadata.id, startupTimeoutMs);
return new SpawnedSessionWorker(child, control, ready.sessionId, ready.pid, shutdownTimeoutMs);
} catch (error) {
if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL");
await waitForExit(child);
await control.close();
throw error;
}
}
@ -55,38 +131,163 @@ function validateTimeout(value: number, name: string): number {
return value;
}
function workerExecArgv(workerUrl: URL): string[] {
return workerUrl.pathname.endsWith(".ts") ? ["--import", "tsx"] : process.execArgv;
type WorkerControllerEvent = SessionWorkerEvent | { type: "hello"; token: string; sessionId: string; pid: number };
interface WorkerControlServer {
readonly connected: Promise<Socket>;
readonly events: AsyncIterable<WorkerControllerEvent>;
send(message: SessionWorkerCommand): Promise<void>;
close(): Promise<void>;
}
function defaultWorkerUrl(): URL {
return new URL(
import.meta.url.endsWith(".js") ? "session-worker-process.js" : "session-worker-process.ts",
import.meta.url,
);
async function createWorkerControlServer(
address: string,
token: string,
expectedSessionId: string,
): Promise<WorkerControlServer> {
let accepted: Socket | undefined;
let resolveConnected!: (socket: Socket) => void;
const connected = new Promise<Socket>((resolve) => {
resolveConnected = resolve;
});
const messages: WorkerControllerEvent[] = [];
const waiters: ((message: WorkerControllerEvent) => void)[] = [];
const server = createServer((socket) => {
if (accepted) {
socket.destroy();
return;
}
accepted = socket;
resolveConnected(socket);
attachJsonLineReader(socket, (value) => {
if (!Check(WorkerControlInputSchema, value)) {
socket.destroy(new Error("Session worker sent an invalid control message"));
return;
}
const message: WorkerControlInput = value;
let event: WorkerControllerEvent;
if (message.type === "worker_failed") {
event = { type: "failed", message: message.message };
} else if (message.type === "worker_ready") {
event = { type: "ready", sessionId: message.sessionId, pid: message.pid };
} else {
event = {
type: "hello",
token: message.token,
sessionId: message.sessionId,
pid: message.pid,
};
}
if (event.type === "hello" && (event.token !== token || event.sessionId !== expectedSessionId)) {
socket.destroy(new Error("Session worker sent invalid control credentials"));
return;
}
const waiter = waiters.shift();
if (waiter) waiter(event);
else messages.push(event);
});
});
await listen(server, address);
if (process.platform !== "win32") await chmod(address, 0o600);
let closePromise: Promise<void> | undefined;
return {
connected,
events: {
[Symbol.asyncIterator]() {
return {
next: async () => {
const message =
messages.shift() ??
(await new Promise<(typeof messages)[number]>((resolve) => waiters.push(resolve)));
return { done: false as const, value: message };
},
};
},
},
async send(message) {
const socket = await connected;
await new Promise<void>((resolve, reject) => {
socket.write(`${JSON.stringify(message)}\n`, (error) => {
if (error) reject(error);
else resolve();
});
});
},
close() {
closePromise ??= closeWorkerControl(server, accepted, address);
return closePromise;
},
};
}
class ChildSessionWorker implements ExperimentalSessionWorker {
async function waitForReady(
child: ChildProcess,
control: WorkerControlServer,
expectedSessionId: string,
startupTimeoutMs: number,
): Promise<Extract<SessionWorkerEvent, { type: "ready" }>> {
const events = control.events[Symbol.asyncIterator]();
const startup = (async () => {
const hello = (await events.next()).value;
if (hello.type !== "hello" || hello.sessionId !== expectedSessionId || hello.pid !== child.pid) {
throw new Error("Session worker reported an invalid identity");
}
const event = (await events.next()).value;
if (event.type === "failed") throw new Error(`Session worker failed: ${event.message}`);
if (event.type !== "ready" || event.sessionId !== expectedSessionId || event.pid !== child.pid) {
throw new Error("Session worker reported an invalid identity");
}
return event;
})();
const exited = waitForExit(child).then(({ code, signal }) => {
throw new Error(`Session worker exited before readiness (${signal ?? code ?? "unknown"})`);
});
let timeout: NodeJS.Timeout | undefined;
const timedOut = new Promise<never>((_, reject) => {
timeout = setTimeout(
() => reject(new Error(`Session worker startup timed out after ${startupTimeoutMs}ms`)),
startupTimeoutMs,
);
timeout.unref();
});
try {
return await Promise.race([startup, exited, timedOut]);
} finally {
if (timeout) clearTimeout(timeout);
}
}
class SpawnedSessionWorker implements ExperimentalSessionWorker {
readonly terminated: Promise<Error | undefined>;
readonly sessionId: string;
readonly pid: number;
readonly #child: ChildProcess;
readonly #control: WorkerControlServer;
readonly #exit: Promise<{ code: number | null; signal: NodeJS.Signals | null }>;
readonly #shutdownTimeoutMs: number;
#closePromise?: Promise<void>;
#closeRequested = false;
constructor(child: ChildProcess, sessionId: string, pid: number, shutdownTimeoutMs: number) {
constructor(
child: ChildProcess,
control: WorkerControlServer,
sessionId: string,
pid: number,
shutdownTimeoutMs: number,
) {
this.#child = child;
this.#control = control;
this.sessionId = sessionId;
this.pid = pid;
this.#shutdownTimeoutMs = shutdownTimeoutMs;
this.#exit = waitForExit(child);
this.terminated = this.#exit.then(({ code, signal }) =>
this.#closeRequested
this.terminated = this.#exit.then(async ({ code, signal }) => {
await this.#control.close();
return this.#closeRequested
? undefined
: new Error(`Session worker ${sessionId} exited unexpectedly (${signal ?? code ?? "unknown"})`),
);
: new Error(`Session worker ${sessionId} exited unexpectedly (${signal ?? code ?? "unknown"})`);
});
}
close(): Promise<void> {
@ -96,69 +297,62 @@ class ChildSessionWorker implements ExperimentalSessionWorker {
}
async #close(): Promise<void> {
if (this.#child.exitCode !== null || this.#child.signalCode !== null) return;
if (!this.#child.connected) {
this.#child.kill("SIGKILL");
} else {
this.#child.send({ type: "shutdown" } satisfies SessionWorkerCommand, (error) => {
if (error) this.#child.kill("SIGKILL");
});
}
const timeout = setTimeout(() => this.#child.kill("SIGKILL"), this.#shutdownTimeoutMs);
timeout.unref();
try {
await this.#exit;
} finally {
clearTimeout(timeout);
if (this.#child.exitCode === null && this.#child.signalCode === null) {
await this.#control.send({ type: "shutdown" }).catch(() => this.#child.kill("SIGKILL"));
const timeout = setTimeout(() => this.#child.kill("SIGKILL"), this.#shutdownTimeoutMs);
timeout.unref();
try {
await this.#exit;
} finally {
clearTimeout(timeout);
}
}
await this.#control.close();
}
}
function waitForReady(
child: ChildProcess,
expectedSessionId: string,
startupTimeoutMs: number,
): Promise<Extract<SessionWorkerEvent, { type: "ready" }>> {
function attachJsonLineReader(socket: Socket, onMessage: (message: unknown) => void): void {
let buffered = "";
socket.setEncoding("utf8");
socket.on("data", (chunk: string) => {
buffered += chunk;
while (true) {
const newline = buffered.indexOf("\n");
if (newline === -1) return;
const line = buffered.slice(0, newline);
buffered = buffered.slice(newline + 1);
try {
onMessage(JSON.parse(line));
} catch {
socket.destroy(new Error("Session worker sent invalid JSON"));
return;
}
}
});
}
function listen(server: Server, address: string): Promise<void> {
return new Promise((resolve, reject) => {
let timeout: NodeJS.Timeout | undefined;
const cleanup = (): void => {
if (timeout) clearTimeout(timeout);
child.off("message", onMessage);
child.off("error", onError);
child.off("exit", onExit);
};
const fail = (error: Error): void => {
cleanup();
const onError = (error: Error): void => {
server.off("listening", onListening);
reject(error);
};
const onMessage = (message: SessionWorkerEvent): void => {
if (message?.type === "failed") {
fail(new Error(`Session worker failed: ${message.message}`));
return;
}
if (message?.type !== "ready") return;
if (message.sessionId !== expectedSessionId || message.pid !== child.pid) {
fail(new Error("Session worker reported an invalid identity"));
return;
}
cleanup();
resolve(message);
const onListening = (): void => {
server.off("error", onError);
resolve();
};
const onError = (error: Error): void => fail(error);
const onExit = (code: number | null, signal: NodeJS.Signals | null): void => {
fail(new Error(`Session worker exited before readiness (${signal ?? code ?? "unknown"})`));
};
timeout = setTimeout(
() => fail(new Error(`Session worker startup timed out after ${startupTimeoutMs}ms`)),
startupTimeoutMs,
);
timeout.unref();
child.on("message", onMessage);
child.once("error", onError);
child.once("exit", onExit);
server.once("error", onError);
server.once("listening", onListening);
server.listen(address);
});
}
async function closeWorkerControl(server: Server, socket: Socket | undefined, address: string): Promise<void> {
socket?.destroy();
if (server.listening) await new Promise<void>((resolve) => server.close(() => resolve()));
if (process.platform !== "win32") await unlink(address).catch(() => undefined);
}
function waitForExit(child: ChildProcess): Promise<{ code: number | null; signal: NodeJS.Signals | null }> {
if (child.exitCode !== null || child.signalCode !== null) {
return Promise.resolve({ code: child.exitCode, signal: child.signalCode });