diff --git a/packages/coding-agent/src/experimental/server.ts b/packages/coding-agent/src/experimental/server.ts index e31a383f2..fd34285f0 100644 --- a/packages/coding-agent/src/experimental/server.ts +++ b/packages/coding-agent/src/experimental/server.ts @@ -331,23 +331,23 @@ async function startServerBackend( const repo = new JsonlSessionRepo({ fileSystem: executionEnv, sessionsRoot: sessionDir }); const host: PiServerHost = { sessions: { - list: async () => { + list: async (context) => { const sessions = new Map( - (await repo.list(undefined, TODO_CONTEXT)).map((metadata) => [metadata.path, metadata]), + (await repo.list(undefined, context)).map((metadata) => [metadata.path, metadata]), ); for (const metadata of workers.trackedSessions) sessions.set(metadata.path, metadata); return [...sessions.values()]; }, - create: async (createOptions) => { - const session = await repo.create(createOptions, TODO_CONTEXT); + create: async (createOptions, context) => { + const session = await repo.create(createOptions, context); try { return session.metadata; } finally { - await session.close(TODO_CONTEXT); + await session.close(context); } }, }, - createHarness: (metadata) => workers.createHarness(metadata), + createHarness: (metadata, context) => workers.createHarness(metadata, context), }; const socketPath = options.path; const closeCatalog = async (): Promise => { diff --git a/packages/coding-agent/src/experimental/session-worker-manager.ts b/packages/coding-agent/src/experimental/session-worker-manager.ts index 6ad55ee5c..2dbc22ec6 100644 --- a/packages/coding-agent/src/experimental/session-worker-manager.ts +++ b/packages/coding-agent/src/experimental/session-worker-manager.ts @@ -1,7 +1,12 @@ import type { ChildProcess } from "node:child_process"; import { randomUUID } from "node:crypto"; import { isAbsolute } from "node:path"; -import type { JsonlSessionMetadata } from "@earendil-works/pi-agent-core"; +import { + BACKGROUND_CONTEXT, + type Context, + type JsonlSessionMetadata, + TODO_CONTEXT, +} from "@earendil-works/pi-agent-core"; import { createRpcClient, type LaneEvent } from "@earendil-works/pi-protocol"; import type { HostedHarnessAttachment, HostedHarnessHandle, HostedHarnessWatch } from "@earendil-works/pi-server"; import { Check } from "typebox/value"; @@ -51,7 +56,7 @@ interface PendingDemand { interface WorkerLaneWatch { readonly worker: WorkerRecord; readonly scope: WorkerOperationScope; - listener?: (event: LaneEvent) => void | Promise; + listener?: (event: LaneEvent, context: Context) => void | Promise; deliveryTail: Promise; } @@ -139,27 +144,33 @@ export class SessionWorkerManager { this.#resolveDiscovery = undefined; } - async createHarness(metadata: JsonlSessionMetadata): Promise { + async createHarness(metadata: JsonlSessionMetadata, context: Context): Promise { if (this.#detached || this.#shuttingDown) throw new Error("Experimental server is shutting down"); const existing = this.#workersBySession.get(metadata.path); if (existing) return this.#hostedHandle(existing); const pending = this.#pending.get(metadata.path); if (pending) return this.#hostedHandle(await pending.promise); - return this.#hostedHandle(await this.#launch(metadata)); + return this.#hostedHandle(await this.#launch(metadata, context)); } #hostedHandle(worker: WorkerRecord): HostedHarnessHandle { - const operations = createRpcClient(SessionWorkerOperations, (call) => this.#invoke(worker, call)); return { terminated: worker.terminated, - attachClient: () => this.#attachClient(worker), - prompt: (prompt) => operations.prompt(prompt), - watch: () => this.#createLaneWatch(worker, operations), - close: () => this.#stopWorker(worker), + attachClient: (context) => this.#attachClient(worker, context), + prompt: (prompt, context) => this.#operations(worker, context).prompt(prompt), + watch: (context) => this.#createLaneWatch(worker, context), + close: (context) => this.#stopWorker(worker, context), }; } - async #attachClient(worker: WorkerRecord): Promise { + #operations( + worker: WorkerRecord, + context: Context, + ): ReturnType> { + return createRpcClient(SessionWorkerOperations, (call) => this.#invoke(worker, call, context)); + } + + async #attachClient(worker: WorkerRecord, context: Context): Promise { if (this.#detached || this.#shuttingDown || worker.stopping) { throw new Error("Experimental Session worker is stopping"); } @@ -170,19 +181,19 @@ export class SessionWorkerManager { const attachmentId = randomUUID(); worker.attachmentId = attachmentId; try { - await this.#applyDemand(worker, attachmentId); + await this.#applyDemand(worker, attachmentId, true, context); } catch (error) { if (worker.attachmentId === attachmentId) worker.attachmentId = undefined; throw error; } let released = false; return { - release: async () => { + release: async (releaseContext) => { if (released) return; released = true; if (this.#detached || worker.attachmentId !== attachmentId) return; try { - await this.#applyDemand(worker, null); + await this.#applyDemand(worker, null, true, releaseContext); } catch (error) { if (!this.#detached && !this.#coordinator.wasReplaced) throw error; } finally { @@ -192,14 +203,13 @@ export class SessionWorkerManager { }; } - async #createLaneWatch( - worker: WorkerRecord, - operations: ReturnType>, - ): Promise { + async #createLaneWatch(worker: WorkerRecord, context: Context): Promise { const scope = this.#operationScope(worker); - const { watchId, snapshot } = await operations.watch(); + const { watchId, snapshot } = await this.#operations(worker, context).watch(); if (this.#laneWatches.has(watchId)) { - await operations.stopWatch(watchId).catch(() => {}); + await this.#operations(worker, context) + .stopWatch(watchId) + .catch(() => {}); throw new Error("Session worker reused an active lane watch ID"); } const watch: WorkerLaneWatch = { worker, scope, deliveryTail: Promise.resolve() }; @@ -207,33 +217,40 @@ export class SessionWorkerManager { let state: "ready" | "started" | "unsubscribed" = "ready"; return { snapshot, - start: async (listener) => { + start: async (listener, startContext) => { if (state !== "ready" || this.#laneWatches.get(watchId) !== watch) { throw new Error("Session worker lane watch may be started only once"); } state = "started"; watch.listener = listener; try { - await operations.startWatch(watchId); + await this.#operations(worker, startContext).startWatch(watchId); } catch (error) { this.#laneWatches.delete(watchId); delete watch.listener; state = "unsubscribed"; - await operations.stopWatch(watchId).catch(() => {}); + await this.#operations(worker, startContext) + .stopWatch(watchId) + .catch(() => {}); throw error; } }, - unsubscribe: async () => { + unsubscribe: async (unsubscribeContext) => { if (state === "unsubscribed") return; state = "unsubscribed"; if (this.#laneWatches.get(watchId) === watch) this.#laneWatches.delete(watchId); delete watch.listener; - await operations.stopWatch(watchId); + await this.#operations(worker, unsubscribeContext).stopWatch(watchId); }, }; } - async #applyDemand(worker: WorkerRecord, attachmentId: string | null, compensateOnTimeout = true): Promise { + async #applyDemand( + worker: WorkerRecord, + attachmentId: string | null, + compensateOnTimeout: boolean, + _context: Context, + ): Promise { if (worker.stopping || this.#workersByPeer.get(worker.peerId) !== worker) { throw new Error("Experimental Session worker is stopping"); } @@ -265,7 +282,7 @@ export class SessionWorkerManager { } /** Worker operations have no wall-clock timeout: completion, disconnect, replacement, or shutdown settles them. */ - #invoke(worker: WorkerRecord, call: SessionWorkerOperationCall): Promise { + #invoke(worker: WorkerRecord, call: SessionWorkerOperationCall, _context: Context): Promise { let scope: WorkerOperationScope; try { scope = this.#operationScope(worker); @@ -301,7 +318,7 @@ export class SessionWorkerManager { }; } - #stopWorker(worker: WorkerRecord): Promise { + #stopWorker(worker: WorkerRecord, _context: Context): Promise { if (this.#detached || this.#workersByPeer.get(worker.peerId) !== worker) return Promise.resolve(); worker.stopPromise ??= this.#stopWorkerInternal(worker); return worker.stopPromise; @@ -359,7 +376,7 @@ export class SessionWorkerManager { })(); await Promise.all([ stopPending, - ...[...this.#workersBySession.values()].map((worker) => this.#stopWorker(worker)), + ...[...this.#workersBySession.values()].map((worker) => this.#stopWorker(worker, BACKGROUND_CONTEXT)), ]); this.#detachState(); } @@ -378,7 +395,7 @@ export class SessionWorkerManager { this.#detachState(); } - #launch(metadata: JsonlSessionMetadata): Promise { + #launch(metadata: JsonlSessionMetadata, _context: Context): Promise { const sessionKey = metadata.path; const peerId = `worker-${randomUUID()}`; const token = randomUUID(); @@ -543,7 +560,7 @@ export class SessionWorkerManager { ) { return; } - watch.deliveryTail = watch.deliveryTail.then(() => listener(event)).catch(() => {}); + watch.deliveryTail = watch.deliveryTail.then(() => listener(event, TODO_CONTEXT)).catch(() => {}); } #recordReadyWorker(peerId: string, message: Extract): void { @@ -651,12 +668,12 @@ export class SessionWorkerManager { clearTimeout(pending.timer); const timeoutError = new Error("Session worker demand update timed out"); try { - await this.#applyDemand(pending.worker, null, false); + await this.#applyDemand(pending.worker, null, false, BACKGROUND_CONTEXT); if (pending.attachmentId === null) pending.resolve(); else pending.reject(timeoutError); } catch (cleanupError) { try { - await this.#stopWorker(pending.worker); + await this.#stopWorker(pending.worker, BACKGROUND_CONTEXT); } catch (stopError) { pending.reject( new AggregateError( diff --git a/packages/coding-agent/test/experimental-session-worker-manager.test.ts b/packages/coding-agent/test/experimental-session-worker-manager.test.ts index e8ef772d3..cc807cdaa 100644 --- a/packages/coding-agent/test/experimental-session-worker-manager.test.ts +++ b/packages/coding-agent/test/experimental-session-worker-manager.test.ts @@ -1,4 +1,4 @@ -import type { JsonlSessionMetadata } from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT, type JsonlSessionMetadata } from "@earendil-works/pi-agent-core"; import { afterEach, describe, expect, test, vi } from "vitest"; import type { CoordinatorConnectionEvent } from "../src/experimental/coordinator.ts"; import { SessionWorkerManager } from "../src/experimental/session-worker-manager.ts"; @@ -65,7 +65,7 @@ async function createAttachedWorker(): Promise<{ const coordinator = new FakeCoordinator(); const workers = new SessionWorkerManager(coordinator, "/tmp"); await workers.discover(new Set(["worker-1"])); - const handle = await workers.createHarness(metadata); + const handle = await workers.createHarness(metadata, BACKGROUND_CONTEXT); coordinator.onSend = (peerId, payload) => { if (payload.type !== "session_demand") return; queueMicrotask(() => @@ -82,8 +82,13 @@ async function createAttachedWorker(): Promise<{ }), ); }; - const attachment = await handle.attachClient!(); - return { coordinator, workers, handle, release: () => Promise.resolve(attachment.release()) }; + const attachment = await handle.attachClient!(BACKGROUND_CONTEXT); + return { + coordinator, + workers, + handle, + release: () => Promise.resolve(attachment.release(BACKGROUND_CONTEXT)), + }; } describe("Session worker lifecycle failures", () => { @@ -92,7 +97,7 @@ describe("Session worker lifecycle failures", () => { const coordinator = new FakeCoordinator(); const workers = new SessionWorkerManager(coordinator, "/tmp"); await workers.discover(new Set(["worker-1"])); - const handle = await workers.createHarness(metadata); + const handle = await workers.createHarness(metadata, BACKGROUND_CONTEXT); const demands: (string | null)[] = []; coordinator.onSend = (peerId, payload) => { if (payload.type !== "session_demand") return; @@ -113,7 +118,7 @@ describe("Session worker lifecycle failures", () => { ); }; - const attaching = expect(handle.attachClient!()).rejects.toThrow("timed out"); + const attaching = expect(handle.attachClient!(BACKGROUND_CONTEXT)).rejects.toThrow("timed out"); await vi.advanceTimersByTimeAsync(5_000); await attaching; expect(demands).toHaveLength(2); @@ -128,10 +133,10 @@ describe("Session worker lifecycle failures", () => { const coordinator = new FakeCoordinator(); const workers = new SessionWorkerManager(coordinator, "/tmp"); await workers.discover(new Set(["worker-1"])); - const handle = await workers.createHarness(metadata); + const handle = await workers.createHarness(metadata, BACKGROUND_CONTEXT); coordinator.onSend = () => {}; - const attaching = expect(handle.attachClient!()).rejects.toThrow("worker was terminated"); + const attaching = expect(handle.attachClient!(BACKGROUND_CONTEXT)).rejects.toThrow("worker was terminated"); await vi.advanceTimersByTimeAsync(5_000); await vi.advanceTimersByTimeAsync(5_000); await vi.advanceTimersByTimeAsync(10_000); @@ -148,7 +153,7 @@ describe("Session worker lifecycle failures", () => { await release(); coordinator.onSend = () => {}; - const closing = handle.close(); + const closing = handle.close(BACKGROUND_CONTEXT); await vi.advanceTimersByTimeAsync(10_000); await closing; expect(kill).toHaveBeenCalledWith(123, "SIGKILL"); @@ -185,7 +190,7 @@ describe("Session worker operations", () => { }); }; - await expect(handle.prompt(["Hello"])).resolves.toEqual({ + await expect(handle.prompt(["Hello"], BACKGROUND_CONTEXT)).resolves.toEqual({ ok: true, value: { kind: "completed", runId: "run-1", leafId: "leaf-1" }, }); @@ -226,7 +231,7 @@ describe("Session worker operations", () => { ); }; - await expect(handle.prompt(["Hello"])).rejects.toThrow(/mismatched operation response/); + await expect(handle.prompt(["Hello"], BACKGROUND_CONTEXT)).rejects.toThrow(/mismatched operation response/); workers.detach(); await release(); }); @@ -257,7 +262,7 @@ describe("Session worker operations", () => { ); }; - await expect(handle.prompt(["Hello"])).rejects.toThrow(/invalid operation response/); + await expect(handle.prompt(["Hello"], BACKGROUND_CONTEXT)).rejects.toThrow(/invalid operation response/); workers.detach(); await release(); }); @@ -265,7 +270,7 @@ describe("Session worker operations", () => { test("rejects pending prompts on replacement without stopping the worker", async () => { const { coordinator, workers, handle } = await createAttachedWorker(); coordinator.onSend = () => {}; - const prompting = handle.prompt(["Hello"]); + const prompting = handle.prompt(["Hello"], BACKGROUND_CONTEXT); workers.detach(); await expect(prompting).rejects.toThrow(/replaced during a worker operation/); diff --git a/packages/server/src/hosted-harness-manager.ts b/packages/server/src/hosted-harness-manager.ts index 99ceca092..9739df5f0 100644 --- a/packages/server/src/hosted-harness-manager.ts +++ b/packages/server/src/hosted-harness-manager.ts @@ -1,5 +1,5 @@ import { randomUUID } from "node:crypto"; -import type { SessionMetadata } from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT, type Context, type SessionMetadata } from "@earendil-works/pi-agent-core"; import { createRpcDispatcher, type EventEnvelope, @@ -56,7 +56,8 @@ interface ClientAttachment { interface RpcContext { readonly client: object; - publish(message: EventEnvelope): Promise; + readonly context: Context; + publish(message: EventEnvelope, context: Context): Promise; } interface HostedSession { @@ -86,25 +87,29 @@ export class HostedHarnessManager (await this.options.host.sessions.list()).map(toProtocolSessionMetadata), - create: async ({ client }, createOptions) => + list: async ({ context }, ..._args: never[]) => + (await this.options.host.sessions.list(context)).map(toProtocolSessionMetadata), + create: async ({ client, context }, createOptions) => this.runForClient(client, async () => { if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError(); - return toProtocolSessionMetadata(await this.options.host.sessions.create(createOptions)); + return toProtocolSessionMetadata(await this.options.host.sessions.create(createOptions, context)); }), - attach: async ({ client }, sessionId) => { + attach: async ({ client, context }, sessionId) => { if (this.options.isClosing()) throw new ServerDrainingError(); - return this.runForClient(client, () => this.attachClient(client, sessionId)); + return this.runForClient(client, () => this.attachClient(client, sessionId, context)); }, - prompt: async ({ client }, sessionId, prompt) => { - const admitted = await this.runForClient(client, () => this.startPrompt(client, sessionId, prompt)); + prompt: async ({ client, context }, sessionId, prompt) => { + const admitted = await this.runForClient(client, () => + this.startPrompt(client, sessionId, prompt, context), + ); return admitted.result; }, - watch: ({ client }, sessionId) => this.runForClient(client, () => this.createWatch(client, sessionId)), - startWatch: ({ client, publish }, sessionId, watchId) => - this.runForClient(client, () => this.startWatch(client, sessionId, watchId, publish)), - stopWatch: ({ client }, sessionId, watchId) => - this.runForClient(client, () => this.stopWatch(client, sessionId, watchId)), + watch: ({ client, context }, sessionId) => + this.runForClient(client, () => this.createWatch(client, sessionId, context)), + startWatch: ({ client, publish, context }, sessionId, watchId) => + this.runForClient(client, () => this.startWatch(client, sessionId, watchId, publish, context)), + stopWatch: ({ client, context }, sessionId, watchId) => + this.runForClient(client, () => this.stopWatch(client, sessionId, watchId, context)), }, (message) => new ProtocolValidationError(message), ); @@ -113,29 +118,30 @@ export class HostedHarnessManager Promise, + publish: (message: EventEnvelope, context: Context) => Promise, + context: Context, ): Promise { - return this.dispatchRpc(call, { client, publish }); + return this.dispatchRpc(call, { client, publish, context }); } - async disconnect(client: object): Promise { + async disconnect(client: object, context: Context): Promise { this.disconnectedClients.add(client); try { await this.runForClient(client, async () => { const attachment = this.attachmentsByClient.get(client); - if (attachment) await this.releaseAttachment(attachment); + if (attachment) await this.releaseAttachment(attachment, context); }); } finally { this.disconnectedClients.delete(client); } } - close(): Promise { - this.closePromise ??= this.closeInternal(); + close(context: Context): Promise { + this.closePromise ??= this.closeInternal(context); return this.closePromise; } - private async closeInternal(): Promise { + private async closeInternal(context: Context): Promise { const operationPromises = [...this.clientOperations.values()]; const openingPromises = [...this.openingSessions.values()]; const [operationResults, openingResults] = await Promise.all([ @@ -150,14 +156,14 @@ export class HostedHarnessManager - session.attachment ? [this.releaseAttachment(session.attachment)] : [], + session.attachment ? [this.releaseAttachment(session.attachment, context)] : [], ), ); for (const result of attachmentResults) { if (result.status === "rejected") closeErrors.push(result.reason); } const hosted = [...this.hostedSessions.values()]; - const closeResults = await Promise.allSettled(hosted.map(({ harness }) => harness.close())); + const closeResults = await Promise.allSettled(hosted.map(({ harness }) => harness.close(context))); for (let index = 0; index < closeResults.length; index++) { const result = closeResults[index]!; const session = hosted[index]!; @@ -187,13 +193,13 @@ export class HostedHarnessManager { + private async attachClient(client: object, sessionId: string, context: Context): Promise<{ sessionId: string }> { if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError(); const current = this.attachmentsByClient.get(client); if (current?.session.id === sessionId) return { sessionId }; - if (current) await this.releaseAttachment(current); + if (current) await this.releaseAttachment(current, context); - const hosted = await this.acquire(sessionId); + const hosted = await this.acquire(sessionId, context); if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError(); const occupied = hosted.attachment; if (occupied) { @@ -205,7 +211,7 @@ export class HostedHarnessManager {} }, + hosted.harness.attachClient ? hosted.harness.attachClient(context) : { release: (_context: Context) => {} }, ); attachment.acquiring = acquiring; attachment.lease = await acquiring; @@ -219,7 +225,7 @@ export class HostedHarnessManager }> { if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError(); const attachment = this.attachmentsByClient.get(client); if (!attachment || attachment.session.id !== sessionId) throw new SessionNotAttachedError(); - const result = attachment.session.harness.prompt(prompt); + const result = attachment.session.harness.prompt(prompt, context); attachment.prompts.add(result); const remove = (): void => { attachment.prompts.delete(result); @@ -246,12 +253,13 @@ export class HostedHarnessManager { const attachment = this.requireAttachment(client, sessionId); if (attachment.watch !== undefined) throw new WatchInUseError(); const create = attachment.session.harness.watch; if (create === undefined) throw new NotSupportedError("Hosted Harness does not support lane watches"); - const handle = await create.call(attachment.session.harness); + const handle = await create.call(attachment.session.harness, context); const watch = { id: randomUUID(), handle, started: false } satisfies ActiveWatch; attachment.watch = watch; return { watchId: watch.id, snapshot: handle.snapshot }; @@ -261,7 +269,8 @@ export class HostedHarnessManager Promise, + publish: (message: EventEnvelope, context: Context) => Promise, + context: Context, ): Promise<{ watchId: string }> { const attachment = this.requireAttachment(client, sessionId); const watch = attachment.watch; @@ -269,11 +278,14 @@ export class HostedHarnessManager publish({ type: "event", watchId, event })); + await watch.handle.start( + (event, eventContext) => publish({ type: "event", watchId, event }, eventContext), + context, + ); } catch (error) { delete attachment.watch; try { - await watch.handle.unsubscribe(); + await watch.handle.unsubscribe(context); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "Failed to start and clean up lane watch"); } @@ -283,12 +295,17 @@ export class HostedHarnessManager { + private async stopWatch( + client: object, + sessionId: string, + watchId: string, + context: Context, + ): Promise<{ watchId: string }> { const attachment = this.requireAttachment(client, sessionId); const watch = attachment.watch; if (watch?.id !== watchId) throw new WatchNotFoundError(); delete attachment.watch; - await watch.handle.unsubscribe(); + await watch.handle.unsubscribe(context); return { watchId }; } @@ -299,21 +316,21 @@ export class HostedHarnessManager { + private releaseAttachment(attachment: ClientAttachment, context: Context): Promise { attachment.releasing ??= (async () => { const errors: unknown[] = []; try { const watch = attachment.watch; delete attachment.watch; try { - await watch?.handle.unsubscribe(); + await watch?.handle.unsubscribe(context); } catch (error) { errors.push(error); } await Promise.allSettled(attachment.prompts); try { const lease = attachment.lease ?? (await attachment.acquiring); - if (lease) await lease.release(); + if (lease) await lease.release(context); } catch (error) { errors.push(error); } @@ -333,12 +350,12 @@ export class HostedHarnessManager { + private async acquire(sessionId: string, context: Context): Promise { const existing = this.hostedSessions.get(sessionId); if (existing) return existing; const opening = this.openingSessions.get(sessionId); if (opening) return opening; - const pending = this.open(sessionId); + const pending = this.open(sessionId, context); this.openingSessions.set(sessionId, pending); try { return await pending; @@ -347,15 +364,17 @@ export class HostedHarnessManager { - const matches = (await this.options.host.sessions.list()).filter((candidate) => candidate.id === sessionId); + private async open(sessionId: string, context: Context): Promise { + const matches = (await this.options.host.sessions.list(context)).filter( + (candidate) => candidate.id === sessionId, + ); if (matches.length === 0) throw new SessionNotFoundError(`Unknown session: ${sessionId}`); if (matches.length > 1) throw new SessionAmbiguousError(); const metadata = matches[0]!; - const harness = await this.options.host.createHarness(metadata); + const harness = await this.options.host.createHarness(metadata, context); if (this.options.isClosing()) { try { - await harness.close(); + await harness.close(context); } catch (error) { this.options.reportError(error); throw new HarnessCleanupError( @@ -381,7 +400,7 @@ export class HostedHarnessManager + void this.releaseAttachment(attachment, BACKGROUND_CONTEXT).catch((releaseError: unknown) => this.options.reportError(releaseError), ); } diff --git a/packages/server/src/server.ts b/packages/server/src/server.ts index fa20a826d..62d3a5621 100644 --- a/packages/server/src/server.ts +++ b/packages/server/src/server.ts @@ -1,4 +1,4 @@ -import type { SessionMetadata } from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT, type SessionMetadata, TODO_CONTEXT } from "@earendil-works/pi-agent-core"; import { type ClientHello, type ClientMessage, @@ -258,9 +258,14 @@ export class PiServer { private async handleRequest(state: ConnectionState, envelope: RequestEnvelope): Promise { try { if (envelope.serverId !== this.serverId) throw new WrongServerError(); - const result: ProtocolRpcResult = await this.sessions.executeCall(envelope.call, state, async (message) => { - await this.sendMessage(state, message); - }); + const result: ProtocolRpcResult = await this.sessions.executeCall( + envelope.call, + state, + async (message, _context) => { + await this.sendMessage(state, message); + }, + TODO_CONTEXT, + ); await this.sendMessage(state, { type: "response", id: envelope.id, @@ -294,7 +299,7 @@ export class PiServer { connection.stage = "closed"; clearTimeout(connection.handshakeTimeout); if (this.connections.delete(connection)) this.notifyConnectionCountChanged(); - void this.sessions.disconnect(connection).catch((error: unknown) => this.reportError(error)); + void this.sessions.disconnect(connection, TODO_CONTEXT).catch((error: unknown) => this.reportError(error)); } private async sendMessage(connection: ConnectionState, message: ServerMessage): Promise { @@ -342,7 +347,7 @@ export class PiServer { } await Promise.all(connections.map((connection) => this.closeConnection(connection.connection))); for (const connection of connections) this.disconnect(connection); - await this.sessions.close(); + await this.sessions.close(BACKGROUND_CONTEXT); this.connections.clear(); } diff --git a/packages/server/src/testing/host.ts b/packages/server/src/testing/host.ts index a40804238..b97db7539 100644 --- a/packages/server/src/testing/host.ts +++ b/packages/server/src/testing/host.ts @@ -1,4 +1,4 @@ -import type { Session, SessionMetadata } from "@earendil-works/pi-agent-core"; +import type { Context, Session, SessionMetadata } from "@earendil-works/pi-agent-core"; import { BACKGROUND_CONTEXT, MemorySessionRepo } from "@earendil-works/pi-agent-core"; import type { LaneEvent, LaneSnapshot, PromptArguments, RunResult } from "@earendil-works/pi-protocol"; import type { HostedHarnessHandle, HostedHarnessWatch, PiServerHost } from "../types.ts"; @@ -35,8 +35,8 @@ const emptyLaneSnapshot: LaneSnapshot = { class TestHarnessWatch implements HostedHarnessWatch { readonly snapshot: LaneSnapshot; - private readonly buffered: LaneEvent[] = []; - private listener: ((event: LaneEvent) => void | Promise) | undefined; + private readonly buffered: Array<{ event: LaneEvent; context: Context }> = []; + private listener: ((event: LaneEvent, context: Context) => void | Promise) | undefined; private tail: Promise = Promise.resolve(); private state: "buffering" | "started" | "unsubscribed" = "buffering"; @@ -44,32 +44,32 @@ class TestHarnessWatch implements HostedHarnessWatch { this.snapshot = structuredClone(snapshot); } - start(listener: (event: LaneEvent) => void | Promise): void { + start(listener: (event: LaneEvent, context: Context) => void | Promise, _context: Context): void { if (this.state !== "buffering") throw new Error("Test Harness watch may be started only once"); this.state = "started"; this.listener = listener; - for (const event of this.buffered.splice(0)) this.enqueue(event); + for (const buffered of this.buffered.splice(0)) this.enqueue(buffered.event, buffered.context); } - unsubscribe(): void { + unsubscribe(_context: Context): void { this.state = "unsubscribed"; this.buffered.splice(0); this.listener = undefined; } - push(event: LaneEvent): Promise { + push(event: LaneEvent, context: Context): Promise { if (this.state === "unsubscribed") return Promise.resolve(); if (this.state === "buffering") { - this.buffered.push(structuredClone(event)); + this.buffered.push({ event: structuredClone(event), context }); return Promise.resolve(); } - return this.enqueue(event); + return this.enqueue(event, context); } - private enqueue(event: LaneEvent): Promise { + private enqueue(event: LaneEvent, context: Context): Promise { const listener = this.listener; if (listener === undefined) return Promise.resolve(); - const delivery = this.tail.then(() => listener(structuredClone(event))); + const delivery = this.tail.then(() => listener(structuredClone(event), context)); this.tail = delivery.catch(() => {}); return delivery; } @@ -97,11 +97,11 @@ export class TestHarness { this.session = session; } - attachClient(): { release(): void } { + attachClient(_context: Context): { release(context: Context): void } { this.attachedClients += 1; let released = false; return { - release: () => { + release: (_context) => { if (released) return; this.attachmentReleaseCount += 1; if (this.failAttachmentRelease) throw this.failAttachmentRelease; @@ -111,24 +111,24 @@ export class TestHarness { }; } - async watch(): Promise { + async watch(_context: Context): Promise { const watch = new TestHarnessWatch(this.watchSnapshot); this.watches.add(watch); return { snapshot: watch.snapshot, - start: (listener) => watch.start(listener), - unsubscribe: () => { - watch.unsubscribe(); + start: (listener, context) => watch.start(listener, context), + unsubscribe: (context) => { + watch.unsubscribe(context); this.watches.delete(watch); }, }; } - async emitEvent(event: LaneEvent): Promise { - await Promise.all([...this.watches].map((watch) => watch.push(event))); + async emitEvent(event: LaneEvent, context: Context = BACKGROUND_CONTEXT): Promise { + await Promise.all([...this.watches].map((watch) => watch.push(event, context))); } - async prompt(prompt: PromptArguments): Promise { + async prompt(prompt: PromptArguments, _context: Context): Promise { this.promptCalls.push(prompt); if (this.nextPromptError) { const error = this.nextPromptError; @@ -149,7 +149,7 @@ export class TestHarness { return result; } - async close(): Promise { + async close(context: Context): Promise { this.closeCount += 1; const gate = this.nextCloseGate; if (gate) { @@ -162,7 +162,7 @@ export class TestHarness { this.failClose = undefined; throw error; } - await this.session.close(BACKGROUND_CONTEXT); + await this.session.close(context); this.closed.resolve(undefined); this.#termination.resolve(undefined); } @@ -197,28 +197,28 @@ export class TestServerHost implements PiServerHost { nextCreateHarnessError?: Error; nextHarnessCloseError?: Error; readonly sessions: PiServerHost["sessions"] = { - list: async () => { + list: async (context) => { const delay = this.nextListDelay; if (delay) { this.nextListDelay = undefined; delay.entered.resolve(undefined); await delay.release.promise; } - return this.repo.list(undefined, BACKGROUND_CONTEXT); + return this.repo.list(undefined, context); }, - create: async ({ id }) => { - const session = await this.repo.create({ id }, BACKGROUND_CONTEXT); + create: async ({ id }, context) => { + const session = await this.repo.create({ id }, context); try { return session.metadata; } finally { - await session.close(BACKGROUND_CONTEXT); + await session.close(context); } }, }; private nextListDelay?: ListDelay; private nextCreateHarnessGate?: OpenGate; - async createHarness(metadata: SessionMetadata): Promise { + async createHarness(metadata: SessionMetadata, context: Context): Promise { this.createHarnessCount += 1; const gate = this.nextCreateHarnessGate; if (gate) { @@ -226,7 +226,7 @@ export class TestServerHost implements PiServerHost { gate.entered.resolve(undefined); await gate.release.promise; } - const session = await this.repo.open(metadata, BACKGROUND_CONTEXT); + const session = await this.repo.open(metadata, context); try { if (this.nextCreateHarnessError) { const error = this.nextCreateHarnessError; @@ -243,7 +243,7 @@ export class TestServerHost implements PiServerHost { this.harnesses.set(metadata.id, harnesses); return harness; } catch (error) { - await session.close(BACKGROUND_CONTEXT); + await session.close(context); throw error; } } diff --git a/packages/server/src/types.ts b/packages/server/src/types.ts index 18b243eb0..45e66e1d2 100644 --- a/packages/server/src/types.ts +++ b/packages/server/src/types.ts @@ -1,4 +1,4 @@ -import type { SessionMetadata } from "@earendil-works/pi-agent-core"; +import type { Context, SessionMetadata } from "@earendil-works/pi-agent-core"; import type { LaneEvent, LaneSnapshot, @@ -22,34 +22,34 @@ export type MaybePromise = T | Promise; /** One client connection's exclusive attachment to a hosted Session. */ export interface HostedHarnessAttachment { - release(): MaybePromise; + release(context: Context): MaybePromise; } /** Snapshot-first observation handle supplied by a hosted Harness adapter. */ export interface HostedHarnessWatch { readonly snapshot: LaneSnapshot; - start(listener: (event: LaneEvent) => MaybePromise): MaybePromise; - unsubscribe(): MaybePromise; + start(listener: (event: LaneEvent, context: Context) => MaybePromise, context: Context): MaybePromise; + unsubscribe(context: Context): MaybePromise; } /** A process-safe handle for the hosted Harness operations exposed by this server. */ export interface HostedHarnessHandle { /** Acquire the Session for one client connection. Sessions permit only one attachment at a time. */ - attachClient?(): MaybePromise; + attachClient?(context: Context): MaybePromise; /** Execute one serializable prompt against the hosted Harness. */ - prompt(prompt: PromptArguments): Promise; + prompt(prompt: PromptArguments, context: Context): Promise; /** Observe the attached main lane when supported by this host. */ - watch?(): Promise; + watch?(context: Context): Promise; /** Resolves with an error for unexpected termination, or undefined after an expected close. */ readonly terminated?: Promise; - close(): Promise; + close(context: Context): Promise; } /** Host capabilities used directly by Session RPC operations. */ export interface PiServerHost { readonly sessions: { - list(): Promise; - create(options: SessionCreateOptions): Promise; + list(context: Context): Promise; + create(options: SessionCreateOptions, context: Context): Promise; }; - createHarness(metadata: TMetadata): Promise; + createHarness(metadata: TMetadata, context: Context): Promise; } diff --git a/packages/server/test/conformance.test.ts b/packages/server/test/conformance.test.ts index b9e5ec5ba..3c67118c5 100644 --- a/packages/server/test/conformance.test.ts +++ b/packages/server/test/conformance.test.ts @@ -1,4 +1,4 @@ -import type { SessionMetadata } from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT, type SessionMetadata } from "@earendil-works/pi-agent-core"; import { afterEach, describe, expect, test } from "vitest"; import type { ByteConnection, ByteConnectionHandler } from "../src/connection.ts"; import { PiServer } from "../src/server.ts"; @@ -652,6 +652,6 @@ describe("hosted Harness acquisition failures", () => { expect(host.latestHarness("session-1").closeCount).toBe(1); servers.delete(server); - await host.latestHarness("session-1").close(); + await host.latestHarness("session-1").close(BACKGROUND_CONTEXT); }); });