diff --git a/proxy-server/src/conversation-manager.ts b/proxy-server/src/conversation-manager.ts index e883fb1..5a39a93 100644 --- a/proxy-server/src/conversation-manager.ts +++ b/proxy-server/src/conversation-manager.ts @@ -71,22 +71,6 @@ export class ConversationManager implements ToolStateProvider { if (isPrimary) { this.primaryId = id; - } else { - const onSessionEnd = (): void => { - // The bridge sends tools/call right after the session goes idle. - // Removing the conversation while those calls are still pending would - // make them fail with "Conversation not found", so defer until they drain. - if (state.toolRouter.hasPending) { - this.logger.debug( - `Conversation ${id} session ended but tool calls pending, deferring removal`, - ); - state.session.onSessionEnd(onSessionEnd); - return; - } - this.logger.debug(`Conversation ${id} session ended, removing`); - this.conversations.delete(id); - }; - state.session.onSessionEnd(onSessionEnd); } this.logger.debug( @@ -112,22 +96,30 @@ export class ConversationManager implements ToolStateProvider { } } - private evictStale(): void { + // A conversation can be removed once no future request can route to it. That + // means it is not the primary, its session is idle, and no tool calls are + // pending. The bridge sends tools/call just after the session goes idle, so + // the pending check keeps an isolated conversation alive to answer them. + private isDisposable(conv: Conversation): boolean { + return ( + !conv.isPrimary && + !conv.state.session.sessionActive && + !conv.state.toolRouter.hasPending + ); + } + + private evictDisposable(): void { for (const [id, conv] of this.conversations) { - if ( - !conv.isPrimary && - !conv.state.session.sessionActive && - !conv.state.toolRouter.hasPending - ) { + if (this.isDisposable(conv)) { conv.state.session.cleanup(); this.conversations.delete(id); - this.logger.debug(`Evicted stale conversation ${id}`); + this.logger.debug(`Evicted idle conversation ${id}`); } } } findForNewRequest(): { conversation: Conversation; isReuse: boolean } { - this.evictStale(); + this.evictDisposable(); const primary = this.getPrimary(); if (primary) { diff --git a/proxy-server/src/tool-bridge/session-lifecycle.ts b/proxy-server/src/tool-bridge/session-lifecycle.ts index e68108a..2a5bf1c 100644 --- a/proxy-server/src/tool-bridge/session-lifecycle.ts +++ b/proxy-server/src/tool-bridge/session-lifecycle.ts @@ -4,7 +4,6 @@ export class SessionLifecycle { private readonly toolRouter: ToolRouter; private _sessionActive = false; private _hadError = false; - private _onSessionEnd: (() => void) | null = null; constructor(toolRouter: ToolRouter) { this.toolRouter = toolRouter; @@ -29,29 +28,14 @@ export class SessionLifecycle { this._hadError = true; } - onSessionEnd(callback: () => void): void { - this._onSessionEnd = callback; - } - markSessionInactive(): void { + // Don't reject expected entries here. The session goes idle before the + // bridge's tools/call requests arrive, so they are still needed. this._sessionActive = false; - // Don't reject expected entries here. handler-core's finally block - // calls this before the MCP tools/call request arrives, so rejecting - // would wipe entries the bridge still needs. cleanup() handles that. - this.fireSessionEnd(); } cleanup(): void { this._sessionActive = false; this.toolRouter.rejectAll("Session cleanup"); - this.fireSessionEnd(); - } - - private fireSessionEnd(): void { - // Clear before calling so the callback can register a new one without - // this method overwriting it. - const callback = this._onSessionEnd; - this._onSessionEnd = null; - if (callback) callback(); } } diff --git a/proxy-server/test/conversation-manager.test.ts b/proxy-server/test/conversation-manager.test.ts index a085083..6224cf2 100644 --- a/proxy-server/test/conversation-manager.test.ts +++ b/proxy-server/test/conversation-manager.test.ts @@ -157,51 +157,91 @@ describe("ConversationManager", () => { }); }); - describe("auto-removal via onSessionEnd", () => { - it("removes non-primary conversation when session becomes inactive", () => { + describe("removing conversations", () => { + it("keeps an idle conversation until the next request", () => { const manager = createManager(); + const primary = manager.create({ isPrimary: true }); + primary.session = { on: () => () => {} } as never; + const conv = manager.create(); conv.state.session.markSessionActive(); - - expect(manager.size).toBe(1); conv.state.session.markSessionInactive(); - expect(manager.size).toBe(0); - expect(manager.getState(conv.id)).toBeUndefined(); - }); - it("removes non-primary conversation on cleanup", () => { - const manager = createManager(); - const conv = manager.create(); + // Going idle does not remove it on its own. + expect(manager.getState(conv.id)).toBe(conv.state); - expect(manager.size).toBe(1); - conv.state.session.cleanup(); - expect(manager.size).toBe(0); + // The next request sweeps it. + manager.findForNewRequest(); + expect(manager.getState(conv.id)).toBeUndefined(); }); - it("does NOT auto-remove primary conversation on session idle", () => { + it("keeps the primary conversation when idle", () => { const manager = createManager(); - const conv = manager.create({ isPrimary: true }); - conv.state.session.markSessionActive(); + const primary = manager.create({ isPrimary: true }); + primary.session = { on: () => () => {} } as never; + primary.state.session.markSessionActive(); + primary.state.session.markSessionInactive(); - expect(manager.size).toBe(1); - conv.state.session.markSessionInactive(); - expect(manager.size).toBe(1); - expect(manager.getState(conv.id)).toBe(conv.state); - expect(manager.getPrimary()).toBe(conv); + manager.findForNewRequest(); + expect(manager.getPrimary()).toBe(primary); }); // An isolated conversation that goes idle right after emitting tool // requests must stay alive because the bridge still has tools/call to deliver. - it("does NOT auto-remove non-primary conversation that has outstanding tool calls", () => { + it("keeps a conversation that asked for a tool", () => { const manager = createManager(); + const primary = manager.create({ isPrimary: true }); + primary.session = { on: () => () => {} } as never; + const conv = manager.create(); conv.state.session.markSessionActive(); conv.state.toolRouter.registerExpected("toolu_01abc", "XcodeRead"); conv.state.session.markSessionInactive(); - expect(manager.size).toBe(1); + manager.findForNewRequest(); expect(manager.getState(conv.id)).toBe(conv.state); }); + + // While a finished conversation is still in the map, nothing should route + // to it, just as if it had already been removed. + it("does not route to a finished conversation", () => { + const manager = createManager(); + const primary = manager.create({ isPrimary: true }); + primary.session = { on: () => () => {} } as never; + + // An isolated conversation runs a tool call to completion, then goes idle. + const finished = manager.create(); + finished.state.session.markSessionActive(); + finished.state.toolRouter.registerExpected("toolu_done", "Read"); + finished.state.toolRouter.registerMCPRequest( + "Read", + () => {}, + () => {}, + ); + finished.state.toolRouter.resolveToolCall("toolu_done", "result"); + finished.state.session.markSessionInactive(); + + // It is still in the map, but nothing can route to it. + expect(manager.getState(finished.id)).toBe(finished.state); + expect(manager.findByContinuationIds(["toolu_done"])).toBeUndefined(); + expect(manager.findByExpectedTool("Read")).toBeUndefined(); + }); + + it("removes a finished conversation on the next request", () => { + const manager = createManager(); + const primary = manager.create({ isPrimary: true }); + primary.session = { on: () => () => {} } as never; + + const finished = manager.create(); + finished.state.session.markSessionActive(); + finished.state.session.markSessionInactive(); + + const { conversation, isReuse } = manager.findForNewRequest(); + + expect(isReuse).toBe(true); + expect(conversation).toBe(primary); + expect(manager.getState(finished.id)).toBeUndefined(); + }); }); describe("tool_use continuation after session goes inactive", () => { @@ -443,7 +483,7 @@ describe("ConversationManager", () => { // A new request must not evict an isolated conversation that still has a // tool call in flight, or the bridge loses it before the result comes back. - it("does NOT evict conversations with in-flight tool calls", async () => { + it("does not evict a conversation with a pending tool call", async () => { const manager = createManager(); const primary = manager.create({ isPrimary: true }); primary.session = { on: () => () => {} } as never; diff --git a/proxy-server/test/tool-bridge/session-lifecycle.test.ts b/proxy-server/test/tool-bridge/session-lifecycle.test.ts index e5b2318..552fb10 100644 --- a/proxy-server/test/tool-bridge/session-lifecycle.test.ts +++ b/proxy-server/test/tool-bridge/session-lifecycle.test.ts @@ -82,33 +82,6 @@ describe("SessionLifecycle", () => { }); }); - describe("onSessionEnd callback", () => { - it("fires when markSessionInactive is called", () => { - const { session } = create(); - const callback = vi.fn(); - session.onSessionEnd(callback); - session.markSessionInactive(); - expect(callback).toHaveBeenCalledOnce(); - }); - - it("fires when cleanup is called", () => { - const { session } = create(); - const callback = vi.fn(); - session.onSessionEnd(callback); - session.cleanup(); - expect(callback).toHaveBeenCalledOnce(); - }); - - it("is cleared after firing (not called twice)", () => { - const { session } = create(); - const callback = vi.fn(); - session.onSessionEnd(callback); - session.markSessionInactive(); - session.markSessionInactive(); - expect(callback).toHaveBeenCalledOnce(); - }); - }); - describe("cleanup", () => { it("rejects all pending with 'Session cleanup' error", async () => { const { router, session } = create(); diff --git a/proxy-server/test/tool-bridge/state.test.ts b/proxy-server/test/tool-bridge/state.test.ts index 93f67b7..d2dd0e2 100644 --- a/proxy-server/test/tool-bridge/state.test.ts +++ b/proxy-server/test/tool-bridge/state.test.ts @@ -75,12 +75,4 @@ describe("ToolBridgeState", () => { await expect(promise).rejects.toThrow("Session cleanup"); expect(state.toolRouter.hasPending).toBe(false); }); - - it("cleanup fires onSessionEnd callback", () => { - const state = new ToolBridgeState(); - const callback = vi.fn(); - state.session.onSessionEnd(callback); - state.session.cleanup(); - expect(callback).toHaveBeenCalledOnce(); - }); });