From ed6300b35ef176c5c2ec84a0b84b5a50a1e1994d Mon Sep 17 00:00:00 2001 From: roboomp Date: Sat, 8 Aug 2026 09:42:35 +0000 Subject: [PATCH 1/4] fix(agent): release parked subagent session memory on dispose Keep-alive subagents are handed to AgentLifecycleManager.adopt, which stores their reviver closure in the process-global #adopted map. The closure is defined inside runSubagent's scope, which also captures the live AgentSession (extension-runner callbacks, buildSubagentSessionOptions), so its lexical environment pins the whole session graph. park() disposes and detaches the session but leaves the adoption record indefinitely, and #doDispose never dropped the in-memory transcript, session-manager entries, or the raw-SSE debug buffer (whose trimmed records retain slice() views of full wire frames), so every completed subagent's heavy state leaked for the process lifetime. dispose() is terminal and every revival path reopens from disk, so #doDispose now sheds retained conversation memory via agent.reset(), RawSseDebugBuffer.clear(), and SessionManager.releaseRetainedEntries(). The adoption record can still reference the session, but only as a husk. Fixes #8003 --- packages/coding-agent/CHANGELOG.md | 4 + .../coding-agent/src/debug/raw-sse-buffer.ts | 20 ++++ .../coding-agent/src/session/agent-session.ts | 10 ++ .../src/session/session-manager.ts | 14 +++ ...nt-session-dispose-releases-memory.test.ts | 93 +++++++++++++++++++ 5 files changed, 141 insertions(+) create mode 100644 packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 1cb7c55d0..8ccc452e1 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed long-running sessions leaking memory for every completed keep-alive `task`/scout subagent: a disposed (parked) subagent's `AgentSession` stayed pinned through the lifecycle adoption record's reviver closure, and `dispose()` never released the message array, session-manager entries, or the raw-SSE debug buffer, so heavy transcripts and captured provider wire frames accumulated for the process lifetime ([#8003](https://github.com/can1357/oh-my-pi/issues/8003)). + ## [17.2.11] - 2026-08-07 ### Added diff --git a/packages/coding-agent/src/debug/raw-sse-buffer.ts b/packages/coding-agent/src/debug/raw-sse-buffer.ts index 3490b4b4b..8bd6fad89 100644 --- a/packages/coding-agent/src/debug/raw-sse-buffer.ts +++ b/packages/coding-agent/src/debug/raw-sse-buffer.ts @@ -335,6 +335,26 @@ export class RawSseDebugBuffer { return body.length > 0 ? `${dropped}${body}` : dropped; } + /** + * Drop every retained record and reset accounting. Called from + * {@link AgentSession} teardown so a disposed (e.g. parked subagent) session + * stops pinning captured wire frames — each trimmed record holds a + * `slice()` of its parent SSE frame, which under JSC keeps the whole + * multi-MB frame alive. Notifies subscribers so a live debug viewer redraws + * empty. + */ + clear(): void { + this.#records = []; + this.#recordChars = []; + this.#head = 0; + this.#totalChars = 0; + this.#droppedRecords = 0; + this.#droppedChars = 0; + this.#totalEvents = 0; + this.#lastUpdatedAt = undefined; + this.#emit(); + } + #append(record: RawSseDebugRecord, chars: number): void { this.#records.push(record); this.#recordChars.push(chars); diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index a2ae42c32..fe52b58e2 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -3851,6 +3851,16 @@ export class AgentSession { } this.#eventListeners = []; this.#sessionChangeCallbacks.clear(); + + // Release retained conversation memory. dispose() is terminal, and every + // revival path reopens the transcript from disk (AgentLifecycleManager + // reviver / persisted-revive / `history://`), so the in-memory copy is + // dead weight from here on. Dropping it lets a parked subagent's session + // graph shed its heavy payloads even while the lifecycle adoption record's + // reviver closure still references the session object. Fixes #8003. + this.agent.reset(); + this.rawSseDebugBuffer.clear(); + this.sessionManager.releaseRetainedEntries(); } #closeAllProviderSessions(reason: string): void { diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index 5a311c508..e99844d8a 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/packages/coding-agent/src/session/session-manager.ts @@ -1707,6 +1707,20 @@ export class SessionManager { if (this.#diskFailure) throw this.#diskFailure; } + /** + * Drop the in-memory transcript after a terminal {@link close}. The entry + * journal and its index mirror the agent's message array (tool results, + * file contents, base64 frame images); on a disposed session — e.g. a + * parked subagent still referenced by the lifecycle adoption record — they + * would otherwise stay pinned for the process lifetime. Reads after this + * point reopen from disk (revival, `history://`), so releasing the + * in-memory copy is safe. Only call once, from session dispose. + */ + releaseRetainedEntries(): void { + this.#entries = []; + this.#index.clear(); + } + getCwd(): string { return this.#cwd; } diff --git a/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts new file mode 100644 index 000000000..b2d214fdb --- /dev/null +++ b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts @@ -0,0 +1,93 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; +import * as path from "node:path"; +import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; +import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { TempDir } from "@oh-my-pi/pi-utils"; + +// Regression: a keep-alive subagent's AgentSession is disposed at park() but +// stays reachable through the lifecycle adoption record's reviver closure +// (which shares the runSubagent lexical environment that captured the live +// session). Before the fix, dispose() left the message array, session-manager +// entries, and the raw-SSE debug buffer intact, so every completed subagent +// pinned its full transcript + captured wire frames for the process lifetime. +// dispose() must now shed that heavy state so the pinned graph is only a husk. +// See issue #8003. +describe("AgentSession dispose releases retained memory", () => { + let tempDir: TempDir; + let authStorage: AuthStorage; + let session: AgentSession | undefined; + + beforeEach(async () => { + tempDir = TempDir.createSync("@omp-dispose-release-"); + authStorage = await AuthStorage.create(path.join(tempDir.path(), "auth.db")); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + }); + + afterEach(async () => { + const current = session; + session = undefined; + if (current) await current.dispose(); + authStorage.close(); + AsyncJobManager.resetForTests(); + vi.restoreAllMocks(); + tempDir.removeSync(); + }); + + function createSession(): AgentSession { + const model = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!model) throw new Error("expected bundled model"); + const mock = createMockModel({ handler: () => ({ content: ["ok"] }) }); + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { model, systemPrompt: ["test"], tools: [] }, + streamFn: mock.stream, + }); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(tempDir.path()), + settings: Settings.isolated(), + modelRegistry: new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml")), + agentId: "Main", + }); + return session; + } + + it("clears messages, session entries, and the raw-SSE buffer on dispose", async () => { + const current = createSession(); + const bulk = "x".repeat(4096); + + const messages: AgentMessage[] = [ + { role: "user", content: [{ type: "text", text: bulk }], timestamp: Date.now() }, + ]; + current.agent.replaceMessages(messages); + current.sessionManager.appendMessage({ role: "user", content: bulk, timestamp: Date.now() }); + current.rawSseDebugBuffer.recordEvent( + { + event: "content_block_delta", + data: `data: ${bulk}`, + raw: ["event: content_block_delta", `data: ${bulk}`], + }, + current.agent.state.model, + ); + + // Precondition: the heavy state is actually present before dispose. + expect(current.agent.state.messages.length).toBeGreaterThan(0); + expect(current.sessionManager.getEntries().length).toBeGreaterThan(0); + expect(current.rawSseDebugBuffer.toRawText().length).toBeGreaterThan(0); + + await current.dispose(); + session = undefined; + + expect(current.agent.state.messages).toHaveLength(0); + expect(current.sessionManager.getEntries()).toHaveLength(0); + expect(current.rawSseDebugBuffer.toRawText()).toBe(""); + expect(current.rawSseDebugBuffer.snapshot().records).toHaveLength(0); + }); +}); From 4ca6c376b40be667ff8c8b4a8838a5d9608255d2 Mon Sep 17 00:00:00 2001 From: roboomp Date: Sat, 8 Aug 2026 09:52:15 +0000 Subject: [PATCH 2/4] fix(agent): detached append-only context on dispose Disposed sessions remained reachable through lifecycle reviver closures. Agent.reset() cleared the live message array but left AppendOnlyContextManager attached, retaining its normalized provider transcript and stable prompt/tool prefix. Detach the append-only manager during terminal disposal and extend the memory-release regression test to cover that second transcript copy. --- packages/coding-agent/CHANGELOG.md | 2 +- .../coding-agent/src/session/agent-session.ts | 1 + ...ent-session-dispose-releases-memory.test.ts | 18 ++++++++++++------ 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 8ccc452e1..ea14911db 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,7 +4,7 @@ ### Fixed -- Fixed long-running sessions leaking memory for every completed keep-alive `task`/scout subagent: a disposed (parked) subagent's `AgentSession` stayed pinned through the lifecycle adoption record's reviver closure, and `dispose()` never released the message array, session-manager entries, or the raw-SSE debug buffer, so heavy transcripts and captured provider wire frames accumulated for the process lifetime ([#8003](https://github.com/can1357/oh-my-pi/issues/8003)). +- Fixed long-running sessions leaking memory for every completed keep-alive `task`/scout subagent: a disposed (parked) subagent's `AgentSession` stayed pinned through the lifecycle adoption record's reviver closure, and `dispose()` never released the message array, append-only provider transcript, session-manager entries, or the raw-SSE debug buffer, so heavy transcripts and captured provider wire frames accumulated for the process lifetime ([#8003](https://github.com/can1357/oh-my-pi/issues/8003)). ## [17.2.11] - 2026-08-07 diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index fe52b58e2..87205e3c0 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -3859,6 +3859,7 @@ export class AgentSession { // graph shed its heavy payloads even while the lifecycle adoption record's // reviver closure still references the session object. Fixes #8003. this.agent.reset(); + this.agent.setAppendOnlyContext(undefined); this.rawSseDebugBuffer.clear(); this.sessionManager.releaseRetainedEntries(); } diff --git a/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts index b2d214fdb..5fd53d089 100644 --- a/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts +++ b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; -import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { Agent, type AgentMessage, AppendOnlyContextManager } from "@oh-my-pi/pi-agent-core"; import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async"; @@ -14,10 +14,10 @@ import { TempDir } from "@oh-my-pi/pi-utils"; // Regression: a keep-alive subagent's AgentSession is disposed at park() but // stays reachable through the lifecycle adoption record's reviver closure // (which shares the runSubagent lexical environment that captured the live -// session). Before the fix, dispose() left the message array, session-manager -// entries, and the raw-SSE debug buffer intact, so every completed subagent -// pinned its full transcript + captured wire frames for the process lifetime. -// dispose() must now shed that heavy state so the pinned graph is only a husk. +// session). Before the fix, dispose() left the message array, append-only +// provider transcript, session-manager entries, and the raw-SSE debug buffer +// intact, so every completed subagent pinned duplicate transcripts and captured +// wire frames. dispose() must shed that heavy state so the pinned graph is only a husk. // See issue #8003. describe("AgentSession dispose releases retained memory", () => { let tempDir: TempDir; @@ -59,7 +59,7 @@ describe("AgentSession dispose releases retained memory", () => { return session; } - it("clears messages, session entries, and the raw-SSE buffer on dispose", async () => { + it("releases all in-memory transcript copies and the raw-SSE buffer on dispose", async () => { const current = createSession(); const bulk = "x".repeat(4096); @@ -67,6 +67,9 @@ describe("AgentSession dispose releases retained memory", () => { { role: "user", content: [{ type: "text", text: bulk }], timestamp: Date.now() }, ]; current.agent.replaceMessages(messages); + const appendOnlyContext = new AppendOnlyContextManager(); + appendOnlyContext.syncMessages([{ role: "user", content: bulk }]); + current.agent.setAppendOnlyContext(appendOnlyContext); current.sessionManager.appendMessage({ role: "user", content: bulk, timestamp: Date.now() }); current.rawSseDebugBuffer.recordEvent( { @@ -79,6 +82,8 @@ describe("AgentSession dispose releases retained memory", () => { // Precondition: the heavy state is actually present before dispose. expect(current.agent.state.messages.length).toBeGreaterThan(0); + expect(current.agent.appendOnlyContext).toBe(appendOnlyContext); + expect(appendOnlyContext.log.length).toBeGreaterThan(0); expect(current.sessionManager.getEntries().length).toBeGreaterThan(0); expect(current.rawSseDebugBuffer.toRawText().length).toBeGreaterThan(0); @@ -88,6 +93,7 @@ describe("AgentSession dispose releases retained memory", () => { expect(current.agent.state.messages).toHaveLength(0); expect(current.sessionManager.getEntries()).toHaveLength(0); expect(current.rawSseDebugBuffer.toRawText()).toBe(""); + expect(current.agent.appendOnlyContext).toBeUndefined(); expect(current.rawSseDebugBuffer.snapshot().records).toHaveLength(0); }); }); From a7a32f35dd6a9b2a202b986f4fad06efa6955334 Mon Sep 17 00:00:00 2001 From: roboomp Date: Sat, 8 Aug 2026 10:06:55 +0000 Subject: [PATCH 3/4] fix(agent): settled active turn before clearing session memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dispose() only *signalled* the agent loop via abort(); it never awaited the run, so a mid-turn dispose (Ctrl-C/timeout/hard-killed subagent) could let the loop unwind after the release ran — its response/SSE interceptors re-recording wire frames into rawSseDebugBuffer and its terminal message re-appending to agent.state.messages, repopulating the disposed session with exactly the retained state the release drops. Detach the response/SSE interceptors and await a bounded agent.waitForIdle() before the reset/clear so it lands on a quiescent session. Added a deterministic regression test that gates the active turn and asserts dispose blocks on it before clearing. --- .../coding-agent/src/session/agent-session.ts | 19 +++++ ...nt-session-dispose-releases-memory.test.ts | 69 +++++++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 87205e3c0..fdcaf6de0 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -3852,6 +3852,25 @@ export class AgentSession { this.#eventListeners = []; this.#sessionChangeCallbacks.clear(); + // A dispose triggered mid-turn (Ctrl-C / timeout / hard-killed subagent) + // only *signals* the agent loop via the earlier abort(); the loop still + // unwinds asynchronously. Detach the response/SSE interceptors so a late + // frame cannot re-record into rawSseDebugBuffer, then wait (bounded) for + // the run to settle so its terminal message lands before — not after — the + // release below. Without this the clear races the unwind and a disposed + // session is repopulated with exactly the state we are trying to drop. + this.agent.setProviderResponseInterceptor(undefined); + this.agent.setRawSseEventInterceptor(undefined); + try { + await withTimeout( + this.agent.waitForIdle(), + POST_PROMPT_DRAIN_TIMEOUT_MS, + "Timed out waiting for the active agent run to settle during dispose", + ); + } catch (error) { + logger.warn("Active agent run still settling at dispose deadline", { error: String(error) }); + } + // Release retained conversation memory. dispose() is terminal, and every // revival path reopens the transcript from disk (AgentLifecycleManager // reviver / persisted-revive / `history://`), so the in-memory copy is diff --git a/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts index 5fd53d089..4c6257106 100644 --- a/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts +++ b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts @@ -96,4 +96,73 @@ describe("AgentSession dispose releases retained memory", () => { expect(current.agent.appendOnlyContext).toBeUndefined(); expect(current.rawSseDebugBuffer.snapshot().records).toHaveLength(0); }); + + it("waits for the active turn to settle before releasing memory", async () => { + const current = createSession(); + const bulk = "y".repeat(4096); + + // Seed a captured frame that dispose must ultimately drop. + current.rawSseDebugBuffer.recordEvent( + { + event: "content_block_delta", + data: `data: ${bulk}`, + raw: ["event: content_block_delta", `data: ${bulk}`], + }, + current.agent.state.model, + ); + + const order: string[] = []; + const reachedSettle = Promise.withResolvers(); + const settle = Promise.withResolvers(); + vi.spyOn(current.agent, "waitForIdle").mockImplementation(async () => { + order.push("waitForIdle:start"); + reachedSettle.resolve(); + await settle.promise; + // The aborted loop unwinds during the settle window: it appends its + // terminal message just before dispose clears the transcript. + current.agent.appendMessage({ + role: "assistant", + content: [{ type: "text", text: bulk }], + timestamp: Date.now(), + } as AgentMessage); + order.push("waitForIdle:end"); + }); + const detachResp = current.agent.setProviderResponseInterceptor.bind(current.agent); + vi.spyOn(current.agent, "setProviderResponseInterceptor").mockImplementation(fn => { + order.push(`detach:resp:${fn === undefined ? "off" : "on"}`); + detachResp(fn); + }); + const reset = current.agent.reset.bind(current.agent); + vi.spyOn(current.agent, "reset").mockImplementation(() => { + order.push("reset"); + reset(); + }); + + const disposeP = current.dispose(); + + // dispose must block on the still-running turn: reaching the settle await + // wins the race against dispose resolving. If dispose ever finished first + // it would have cleared the transcript mid-turn (the bug under test). + const winner = await Promise.race([ + reachedSettle.promise.then(() => "reached" as const), + disposeP.then(() => "disposed" as const), + ]); + expect(winner).toBe("reached"); + + // The response interceptor was detached before the wait, and nothing has + // been cleared yet. + expect(order).toContain("detach:resp:off"); + expect(order.indexOf("detach:resp:off")).toBeLessThan(order.indexOf("waitForIdle:start")); + expect(order).not.toContain("reset"); + + settle.resolve(); + await disposeP; + session = undefined; + + // reset ran only after the turn settled; the terminal message appended + // during the unwind and the seeded frame were both dropped. + expect(order.indexOf("reset")).toBeGreaterThan(order.indexOf("waitForIdle:end")); + expect(current.agent.state.messages).toHaveLength(0); + expect(current.rawSseDebugBuffer.snapshot().records).toHaveLength(0); + }); }); From 88021b90eaf8644512891809d59fd431beda9bce Mon Sep 17 00:00:00 2001 From: roboomp Date: Sat, 8 Aug 2026 10:26:35 +0000 Subject: [PATCH 4/4] fix(agent): drained in-flight session event handlers on dispose MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit agent-core dispatches the session's event subscriber fire-and-forget (agent.ts #emit), so a message_end/agent_end handler can still be awaiting extension/subscriber/maintenance work — and its sessionManager/agent.state append — after agent.waitForIdle() resolves. The earlier settle waited only on the core run, so a late handler could append the finished message/entries back into the disposed session and re-pin the transcript. Track every #handleAgentEvent dispatch in #inFlightEventHandlers and drain it (alongside agent.waitForIdle) inside the bounded settle before reset/clear/release. Added a regression test using a real extension whose message_end hook blocks before persistence, asserting dispose does not release memory until the in-flight handler settles. --- .../coding-agent/src/session/agent-session.ts | 56 +++++++++-- ...nt-session-dispose-releases-memory.test.ts | 99 +++++++++++++++++++ 2 files changed, 147 insertions(+), 8 deletions(-) diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index fdcaf6de0..c581bf0a1 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -2056,6 +2056,41 @@ export class AgentSession { */ #prunedTerminalRefusal: AssistantMessage | undefined = undefined; + /** + * In-flight {@link #dispatchAgentEvent} promises. agent-core invokes the + * event subscriber fire-and-forget, so a `message_end`/`agent_end` handler + * can still be awaiting extension/subscriber/maintenance work — and thus its + * `sessionManager`/`agent.state` append — after `agent.waitForIdle()` + * resolves. Dispose drains this set so the late append lands *before* the + * memory release, never after it. + */ + #inFlightEventHandlers = new Set>(); + + /** + * Subscriber entry point. Delegates to {@link #dispatchAgentEvent} and + * records the dispatch in {@link #inFlightEventHandlers} until it settles so + * {@link #drainInFlightEventHandlers} can await the session's async + * event/persistence pipeline during teardown. + */ + #handleAgentEvent = (event: AgentEvent): Promise => { + const processing = this.#dispatchAgentEvent(event); + this.#inFlightEventHandlers.add(processing); + void processing.finally(() => this.#inFlightEventHandlers.delete(processing)).catch(() => {}); + return processing; + }; + + /** + * Await every in-flight event handler (and any it chains into) so a late + * message/entry append cannot land after the caller clears session memory. + * The agent must already be idle — otherwise new events keep arriving and + * this never drains. + */ + async #drainInFlightEventHandlers(): Promise { + while (this.#inFlightEventHandlers.size > 0) { + await Promise.allSettled([...this.#inFlightEventHandlers]); + } + } + /** Internal handler for agent events - shared by subscribe and reconnect. * * `agent_end` handling schedules deferred post-prompt recovery work @@ -2070,7 +2105,7 @@ export class AgentSession { * `#postPromptTasksPromise` is set the moment `#emit` invokes this handler, so * the recovery wait always sees the in-flight handler and blocks until it — and * everything it schedules — settles. */ - #handleAgentEvent = async (event: AgentEvent): Promise => { + #dispatchAgentEvent = async (event: AgentEvent): Promise => { if (event.type === "tool_execution_end" && this.#isTerminalYieldToolResult(event)) { const alreadyTerminated = this.#synchronouslyTerminatedYieldToolCallIds.delete(event.toolCallId); if (!alreadyTerminated) { @@ -3853,17 +3888,22 @@ export class AgentSession { this.#sessionChangeCallbacks.clear(); // A dispose triggered mid-turn (Ctrl-C / timeout / hard-killed subagent) - // only *signals* the agent loop via the earlier abort(); the loop still - // unwinds asynchronously. Detach the response/SSE interceptors so a late - // frame cannot re-record into rawSseDebugBuffer, then wait (bounded) for - // the run to settle so its terminal message lands before — not after — the - // release below. Without this the clear races the unwind and a disposed - // session is repopulated with exactly the state we are trying to drop. + // only *signals* the agent loop via the earlier abort(); the loop and the + // session's fire-and-forget event handlers still unwind asynchronously. + // Detach the response/SSE interceptors so a late frame cannot re-record + // into rawSseDebugBuffer, then wait (bounded) for both the core run AND + // the in-flight event/persistence handlers to settle — the latter can + // still append the finished message/entries after agent.waitForIdle() + // alone. Without this the release races the unwind and a disposed session + // is repopulated with exactly the state we are trying to drop. this.agent.setProviderResponseInterceptor(undefined); this.agent.setRawSseEventInterceptor(undefined); try { await withTimeout( - this.agent.waitForIdle(), + (async () => { + await this.agent.waitForIdle(); + await this.#drainInFlightEventHandlers(); + })(), POST_PROMPT_DRAIN_TIMEOUT_MS, "Timed out waiting for the active agent run to settle during dispose", ); diff --git a/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts index 4c6257106..258305c57 100644 --- a/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts +++ b/packages/coding-agent/test/agent-session-dispose-releases-memory.test.ts @@ -1,14 +1,18 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; import { Agent, type AgentMessage, AppendOnlyContextManager } from "@oh-my-pi/pi-agent-core"; +import type { AssistantMessage } from "@oh-my-pi/pi-ai"; import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { ExtensionRuntime, loadExtensionFromFactory } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/loader"; +import { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/runner"; import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { EventBus } from "@oh-my-pi/pi-coding-agent/utils/event-bus"; import { TempDir } from "@oh-my-pi/pi-utils"; // Regression: a keep-alive subagent's AgentSession is disposed at park() but @@ -165,4 +169,99 @@ describe("AgentSession dispose releases retained memory", () => { expect(current.agent.state.messages).toHaveLength(0); expect(current.rawSseDebugBuffer.snapshot().records).toHaveLength(0); }); + + it("drains in-flight event handlers so a late persist cannot repopulate a disposed session", async () => { + const model = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!model) throw new Error("expected bundled model"); + const mock = createMockModel({ handler: () => ({ content: ["ok"] }) }); + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { model, systemPrompt: ["test"], tools: [] }, + streamFn: mock.stream, + }); + const sessionManager = SessionManager.inMemory(tempDir.path()); + const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml")); + + // A real extension whose message_end hook blocks. This models the exact + // gap the fix closes: agent-core dispatches the session's event handler + // fire-and-forget, and that handler awaits extension work BEFORE it + // persists the finished message — so agent.waitForIdle() alone is not + // enough to know the session is quiescent. + const reached = Promise.withResolvers(); + const release = Promise.withResolvers(); + const runtime = new ExtensionRuntime(); + const extension = await loadExtensionFromFactory( + pi => { + pi.on("message_end", async () => { + reached.resolve(); + await release.promise; + }); + }, + tempDir.path(), + new EventBus(), + runtime, + "block-message-end", + ); + const extensionRunner = new ExtensionRunner([extension], runtime, tempDir.path(), sessionManager, modelRegistry); + + const current = new AgentSession({ + agent, + sessionManager, + settings: Settings.isolated(), + modelRegistry, + agentId: "Main", + extensionRunner, + }); + session = current; + + const message: AssistantMessage = { + role: "assistant", + content: [{ type: "text", text: "x".repeat(4096) }], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; + + // Dispatch a real message_end: the session handler runs fire-and-forget + // and parks in the extension hook before it can persist the entry. + current.agent.emitExternalEvent({ type: "message_end", message }); + await reached.promise; + expect(current.sessionManager.getEntries()).toHaveLength(0); // persist not reached yet + + // Dispose must not release memory until that in-flight handler settles. + const reachedSettle = Promise.withResolvers(); + const detach = current.agent.setProviderResponseInterceptor.bind(current.agent); + vi.spyOn(current.agent, "setProviderResponseInterceptor").mockImplementation(fn => { + detach(fn); + reachedSettle.resolve(); + }); + const releaseSpy = vi.spyOn(current.sessionManager, "releaseRetainedEntries"); + + const disposeP = current.dispose(); + await reachedSettle.promise; + for (let i = 0; i < 10; i++) await Promise.resolve(); + + // Blocked draining the in-flight handler: memory release has not run. + expect(releaseSpy).not.toHaveBeenCalled(); + + release.resolve(); + await disposeP; + session = undefined; + + // The late persist landed during the drain and was then cleared, so the + // disposed session retains neither the entry nor the message. + expect(releaseSpy).toHaveBeenCalledTimes(1); + expect(current.sessionManager.getEntries()).toHaveLength(0); + expect(current.agent.state.messages).toHaveLength(0); + }); });