diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 46d2bd4f8..3e4f1fe81 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -15,7 +15,7 @@ import { AsyncLocalStorage } from "node:async_hooks"; import { randomUUID } from "node:crypto"; -import { existsSync, mkdirSync, mkdtempSync, readFileSync, writeFileSync } from "node:fs"; +import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { basename, dirname, join, resolve } from "node:path"; import { @@ -8760,8 +8760,8 @@ export class AgentSession { /** Typed handlers for host requests arriving from the IPython kernel comm bridge. */ private _createKernelHostHandlers(): HostRequestHandlers { const handlers: HostRequestHandlers = { - "rlm.run": createRlmRunHostHandler(async ({ prompt, kwargs, cellSourceCode }) => ({ - ...(await this.runRlmChild(prompt, kwargs, cellSourceCode)), + "rlm.run": createRlmRunHostHandler(async ({ prompt, kwargs, cellSourceCode }, signal) => ({ + ...(await this.runRlmChild(prompt, kwargs, cellSourceCode, signal)), })), "rlm.find_models": createRlmFindModelsHostHandler((query, limit) => this.findRlmModels(query, limit)), "rlm.list_subagents": createRlmListSubagentsHostHandler(() => this.listRlmSubagents()), @@ -9685,7 +9685,9 @@ export class AgentSession { prompt: string, kwargs: Record = {}, spawnCode?: string, + signal?: AbortSignal, ): Promise { + signal?.throwIfAborted(); const { name: rawName, model: rawModel, ...unsupported } = kwargs; const unsupportedKwargs = Object.keys(unsupported); if (unsupportedKwargs.length > 0) { @@ -9712,12 +9714,19 @@ export class AgentSession { } finally { if (requestedSessionName) this._pendingRlmSubagentSessionNames.delete(requestedSessionName); } + signal?.throwIfAborted(); if (this._disposed || this._disposing) throw new Error("Cannot spawn a subagent after its parent was disposed"); const childSessionDir = this._createChildRlmSessionDir(); const childNodeId = basename(childSessionDir); const sessionName = requestedSessionName ?? createDefaultRlmSubagentSessionName(prompt, childNodeId); - if (!requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(sessionName); + try { + if (!requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(sessionName); + signal?.throwIfAborted(); + } catch (error) { + rmSync(childSessionDir, { recursive: true, force: true }); + throw error; + } const startedAt = Date.now(); const parentAssistantForUsage = this._findLastAssistantMessage(); const label = rlmChildLabel(prompt); @@ -9741,6 +9750,15 @@ export class AgentSession { if (run.status === "cancelled") throw new Error(run.error ?? "RLM child cancelled"); }; this._activeRlmChildRuns.set(run.id, run); + const abortFromHost = () => { + const reason = signal?.reason; + this._cancelRlmChildRun(run, reason instanceof Error ? reason.message : "IPython kernel host request aborted"); + }; + if (signal?.aborted) { + abortFromHost(); + } else { + signal?.addEventListener("abort", abortFromHost, { once: true }); + } const emitChildUpdate = () => { const childModel = childSession?.model ?? modelSelection.model; this._emit({ @@ -9995,6 +10013,7 @@ export class AgentSession { } } } finally { + signal?.removeEventListener("abort", abortFromHost); if (run.detachedDeletion && childRuntime) { try { await this._deleteRlmSubagentSession(run.id, childRuntime.session); @@ -10036,8 +10055,9 @@ export class AgentSession { prompt: string, kwargs: Record = {}, spawnCode?: string, + signal?: AbortSignal, ): Promise { - return this._startRlmChildRun(prompt, kwargs, spawnCode); + return this._startRlmChildRun(prompt, kwargs, spawnCode, signal); } // ========================================================================= diff --git a/packages/coding-agent/src/core/kernel/index.ts b/packages/coding-agent/src/core/kernel/index.ts index f7916eab8..757eeea65 100644 --- a/packages/coding-agent/src/core/kernel/index.ts +++ b/packages/coding-agent/src/core/kernel/index.ts @@ -62,7 +62,10 @@ export const HOST_COMM_TARGET = "host.request"; * This legacy unary compatibility alias remains the dispatcher and registration * contract while context-aware handlers are staged separately below. */ -export type HostRequestHandler = (payload: Record) => Promise>; +export type HostRequestHandler = ( + payload: Record, + signal?: AbortSignal, +) => Promise>; /** * Per-call authority supplied by the host-request dispatcher. @@ -616,6 +619,7 @@ export class KernelManager { // attribute their spawning program. private lastCellCode?: string; private readonly inFlightHostRequests = new Set>(); + private hostRequestController = new AbortController(); private state: "idle" | "starting" | "running" | "shutdown" = "idle"; /** Memoized so concurrent callers all await the same in-flight startup. */ private startPromise?: Promise; @@ -658,6 +662,9 @@ export class KernelManager { private async doStart(startOptions: KernelStartOptions): Promise { if (this.state !== "idle") return; + if (this.hostRequestController.signal.aborted) { + this.hostRequestController = new AbortController(); + } this.state = "starting"; installSignalHandlersOnce(); // Tracked from the moment startup begins so session cleanup and signal @@ -1300,9 +1307,10 @@ export class KernelManager { } this.handledHostRequestCommIds.add(commId); + const signal = this.hostRequestController.signal; const task = (async () => { try { - const result = await this.handleHostRequest(data); + const result = await this.handleHostRequest(data, signal); try { await this.sendCommMessage(commId, { status: "ok", ...result }); } catch (replyError) { @@ -1327,7 +1335,7 @@ export class KernelManager { }); } - private async handleHostRequest(data: unknown): Promise> { + private async handleHostRequest(data: unknown, signal: AbortSignal): Promise> { if (!isRecord(data)) { throw new Error("host request payload must be an object"); } @@ -1343,7 +1351,7 @@ export class KernelManager { // the in-flight execution; detached spawns (asyncio.create_task) fire after // the scheduling cell goes idle, so fall back to that last cell's source. const cellSourceCode = this.activeExecution?.code ?? this.lastCellCode; - return handler({ ...data, cellSourceCode }); + return handler({ ...data, cellSourceCode }, signal); } private async sendCommMessage(commId: string, data: Record): Promise { @@ -1362,6 +1370,7 @@ export class KernelManager { } private cleanupResources(killSignal: NodeJS.Signals = "SIGTERM"): void { + this.abortHostRequests("IPython kernel stopped"); this.clearSnapshotTimer(); this.lateSentAgentMessageHandlers.clear(); if (this.forkedLivenessTimer) { @@ -1401,6 +1410,12 @@ export class KernelManager { this.startPromise = undefined; } + private abortHostRequests(message: string): void { + if (!this.hostRequestController.signal.aborted) { + this.hostRequestController.abort(new Error(message)); + } + } + private async waitForHostRequestsToSettle(tasks: Promise[], timeoutMs: number): Promise { let timeout: ReturnType | undefined; const timeoutPromise = new Promise<"timeout">((resolve) => { @@ -1429,6 +1444,7 @@ export class KernelManager { } // Best-effort final flush (bounded) before teardown — used by signal handlers // so a SIGINT/SIGTERM exit doesn't lose work the debounced snapshot hasn't saved. + this.abortHostRequests("IPython kernel shut down"); if (opts.snapshot) { await this.flushSnapshotForDispose(); } @@ -1469,6 +1485,7 @@ export class KernelManager { } async kill(): Promise { + this.abortHostRequests("IPython kernel killed"); this.state = "shutdown"; liveKernels.delete(this); this.cleanupResources("SIGKILL"); @@ -1575,12 +1592,12 @@ export class KernelManager { /** Graceful cleanup. Waits briefly for in-flight host request handlers before closing sockets. */ dispose(): Promise { return (async () => { + this.abortHostRequests("IPython kernel disposed"); // Final namespace flush while the kernel is still live (session end / reload). await this.flushSnapshotForDispose(); this.state = "shutdown"; liveKernels.delete(this); const inFlightHostRequests = [...this.inFlightHostRequests]; - // TODO: plumb AbortSignal through AgentSession.prompt so disposal can cancel long-running child loops. try { if (inFlightHostRequests.length > 0) { await this.waitForHostRequestsToSettle(inFlightHostRequests, HOST_REQUEST_DISPOSE_TIMEOUT_MS); @@ -1593,6 +1610,7 @@ export class KernelManager { /** Synchronous best-effort cleanup. Safe to call from `process.on('exit')`. */ disposeSync(): void { + this.abortHostRequests("IPython kernel disposed"); this.state = "shutdown"; liveKernels.delete(this); // TODO: replace this best-effort hard-exit path if Node exposes an awaitable process-exit cleanup hook. diff --git a/packages/coding-agent/src/core/rlm-runtime.ts b/packages/coding-agent/src/core/rlm-runtime.ts index e89472fce..9a5c66652 100644 --- a/packages/coding-agent/src/core/rlm-runtime.ts +++ b/packages/coding-agent/src/core/rlm-runtime.ts @@ -49,7 +49,7 @@ export interface RlmFindModelsResult { models: RlmModelMatch[]; } -export type RlmRunHandler = (request: RlmRunRequest) => Promise>; +export type RlmRunHandler = (request: RlmRunRequest, signal?: AbortSignal) => Promise>; export type RlmListSubagentsHandler = () => RlmListSubagentsResult | Promise; export type RlmDeleteSubagentHandler = (target: string) => Promise; export type RlmFindModelsHandler = (query: string, limit: number) => RlmFindModelsResult | Promise; @@ -150,17 +150,20 @@ export function findRlmModelMatches(query: string, models: Model[], limit: /** Adapt an RlmRunHandler into the typed "rlm.run" handler for the kernel host bridge. */ export function createRlmRunHostHandler(handler: RlmRunHandler): HostRequestHandler { - return async (payload) => { + return async (payload, signal) => { if (typeof payload.prompt !== "string") { throw new Error("rlm.run prompt must be a string"); } const kwargs = isRecord(payload.kwargs) ? payload.kwargs : {}; const cellSourceCode = typeof payload.cellSourceCode === "string" ? payload.cellSourceCode : undefined; - const result = await handler({ - prompt: payload.prompt, - kwargs, - cellSourceCode, - }); + const result = await handler( + { + prompt: payload.prompt, + kwargs, + cellSourceCode, + }, + signal, + ); return result as unknown as Record; }; } diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index f066f02b1..48e66e002 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -144,6 +144,7 @@ interface KernelExecuteTestApi { start: () => Promise; state: "idle" | "starting" | "running" | "shutdown"; activeExecution?: unknown; + snapshotState?: () => Promise; shell?: { send(frames: Buffer[]): Promise; close(): void; @@ -2168,6 +2169,43 @@ describe("AgentSession rlm recursion", () => { ); }); + it("removes an unadmitted child directory when host cancellation wins name validation", async () => { + let releaseNameCheck: () => void = () => {}; + const nameCheckGate = new Promise((resolve) => { + releaseNameCheck = resolve; + }); + let nameCheckStarted = false; + const controller = new AbortController(); + const root = createSession({ + agentMessageController: { + assertSessionNameAvailable: async () => { + nameCheckStarted = true; + await nameCheckGate; + }, + listAgents: async () => ({ + current: { activeSessionId: "parent-active", sessionId: "parent-session" }, + agents: [], + }), + sendAgentMessage: async () => { + throw new Error("unexpected send"); + }, + }, + }); + + const run = root.runRlmChild("cancel during default name validation", {}, undefined, controller.signal); + await waitFor(() => nameCheckStarted); + const childDirs = readdirSync(tempDir, { recursive: true, encoding: "utf8" }).filter((path) => + basename(path).startsWith("sub-"), + ); + expect(childDirs).toHaveLength(1); + + controller.abort(new Error("host disposed")); + releaseNameCheck(); + + await expect(run).rejects.toThrow("host disposed"); + expect(existsSync(join(tempDir, childDirs[0] ?? ""))).toBe(false); + }); + it("cancels active rlm children when the parent session is disposed", async () => { let releaseChild: () => void = () => {}; const release = new Promise((resolve) => { @@ -2893,6 +2931,57 @@ describe("AgentSession rlm recursion", () => { await waitFor(() => rootRun.status === "done"); }); + it("cancels an admitted child in promptAndWait when its kernel host is disposed", async () => { + let releaseChild: () => void = () => {}; + const release = new Promise((resolve) => { + releaseChild = resolve; + }); + let childStarted = false; + const root = createSession({ + streamFn: (_model, context) => { + const text = userText(context); + const stream = createAssistantMessageEventStream(); + if (text === "kernel-owned shard") { + childStarted = true; + void release.then(() => { + stream.push({ type: "done", reason: "stop", message: assistantMessage(`child answer: ${text}`) }); + }); + } + return stream; + }, + }); + const manager = new KernelManager({ + python: process.execPath, + hostHandlers: (root as unknown as InspectableRlmSession)._createKernelHostHandlers(), + }); + const replies: CapturedCommReply[] = []; + const kernel = manager as unknown as KernelCommTestApi; + kernel.sendCommMessage = async (commId, data) => { + replies.push({ commId, data }); + }; + + try { + kernel.handleCommMessage(rlmCommOpen("comm-real-child", "kernel-owned shard")); + await waitFor(() => replies.some((reply) => reply.commId === "comm-real-child")); + await waitFor(() => childStarted); + const runs = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; + const run = [...runs.values()][0]; + if (!run?.session) throw new Error("Missing admitted child session"); + const childAbort = vi.spyOn(run.session, "abort"); + + await manager.dispose(); + + expect(run.status).toBe("cancelled"); + expect(run.error).toBe("IPython kernel disposed"); + expect(childAbort).toHaveBeenCalledTimes(1); + releaseChild(); + await waitFor(() => !runs.has(run.id)); + } finally { + releaseChild(); + await manager.dispose(); + } + }); + it("runs parallel rlm comm requests independently", async () => { let active = 0; let maxActive = 0; @@ -3122,26 +3211,59 @@ print(_result.name) } }); - it("waits for in-flight rlm comm work during dispose and buffers failures", async () => { + it.each(["dispose", "shutdown"] as const)( + "aborts host work before the final snapshot during %s", + async (teardown) => { + let hostSignal: AbortSignal | undefined; + const manager = new KernelManager({ + python: process.execPath, + snapshot: { path: "/tmp/unused.dill", manifestPath: "/tmp/unused.json" }, + hostHandlers: { + "rlm.run": createRlmRunHostHandler(async (_request, signal) => { + hostSignal = signal; + await new Promise((resolve) => { + if (signal?.aborted) resolve(); + else signal?.addEventListener("abort", () => resolve(), { once: true }); + }); + return {}; + }), + }, + }); + const kernel = manager as unknown as KernelCommTestApi & KernelExecuteTestApi; + kernel.state = "running"; + kernel.sendCommMessage = async () => {}; + kernel.snapshotState = async () => { + expect(hostSignal?.aborted).toBe(true); + return null; + }; + + kernel.handleCommMessage(rlmCommOpen(`comm-${teardown}`, "slow child")); + await waitFor(() => hostSignal !== undefined); + + if (teardown === "dispose") await manager.dispose(); + else await manager.shutdown({ snapshot: true }); + + expect(hostSignal?.aborted).toBe(true); + }, + ); + + it("aborts in-flight rlm comm work during dispose and buffers failures", async () => { let started = false; let handlerSettled = false; - let released = false; - let releaseChild: () => void = () => {}; - const release = new Promise((resolve) => { - releaseChild = () => { - if (released) return; - released = true; - resolve(); - }; - }); + let receivedSignal: AbortSignal | undefined; const manager = new KernelManager({ python: process.execPath, hostHandlers: { - "rlm.run": createRlmRunHostHandler(async () => { + "rlm.run": createRlmRunHostHandler(async (_request, signal) => { started = true; + receivedSignal = signal; try { - await release; - throw new Error("child failed after dispose"); + await new Promise((_resolve, reject) => { + const onAbort = () => reject(signal?.reason ?? new Error("aborted")); + if (signal?.aborted) onAbort(); + else signal?.addEventListener("abort", onAbort, { once: true }); + }); + return {}; } finally { handlerSettled = true; } @@ -3152,21 +3274,11 @@ print(_result.name) try { const kernel = manager as unknown as KernelCommTestApi; - kernel.handleCommMessage(rlmCommOpen("comm-dispose", "slow child")); await waitFor(() => started); - const disposePromise = manager.dispose(); - let disposeSettled = false; - const trackedDispose = disposePromise.then(() => { - disposeSettled = true; - }); - - await sleep(25); - expect(disposeSettled).toBe(false); - - releaseChild(); - await expectSettlesWithin(trackedDispose, 1000); + await expectSettlesWithin(manager.dispose(), 1000); + expect(receivedSignal?.aborted).toBe(true); expect(handlerSettled).toBe(true); const kernelStderr = (manager as unknown as { kernelStderr: string }).kernelStderr; @@ -3174,7 +3286,6 @@ print(_result.name) expect(kernelStderr).toContain("[kernel] failed to send host request error reply for comm comm-dispose"); expect(stderrSpy).not.toHaveBeenCalled(); } finally { - releaseChild(); await manager.dispose(); stderrSpy.mockRestore(); }