diff --git a/proxy-server/src/conversation-manager.ts b/proxy-server/src/conversation-manager.ts index 81c8e23..e883fb1 100644 --- a/proxy-server/src/conversation-manager.ts +++ b/proxy-server/src/conversation-manager.ts @@ -72,10 +72,21 @@ export class ConversationManager implements ToolStateProvider { if (isPrimary) { this.primaryId = id; } else { - state.session.onSessionEnd(() => { + 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( @@ -103,7 +114,11 @@ export class ConversationManager implements ToolStateProvider { private evictStale(): void { for (const [id, conv] of this.conversations) { - if (!conv.isPrimary && !conv.state.session.sessionActive) { + if ( + !conv.isPrimary && + !conv.state.session.sessionActive && + !conv.state.toolRouter.hasPending + ) { conv.state.session.cleanup(); this.conversations.delete(id); this.logger.debug(`Evicted stale conversation ${id}`); diff --git a/proxy-server/src/tool-bridge/session-lifecycle.ts b/proxy-server/src/tool-bridge/session-lifecycle.ts index d63f765..e68108a 100644 --- a/proxy-server/src/tool-bridge/session-lifecycle.ts +++ b/proxy-server/src/tool-bridge/session-lifecycle.ts @@ -48,9 +48,10 @@ export class SessionLifecycle { } private fireSessionEnd(): void { - if (this._onSessionEnd) { - this._onSessionEnd(); - this._onSessionEnd = null; - } + // 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 46a5c75..a085083 100644 --- a/proxy-server/test/conversation-manager.test.ts +++ b/proxy-server/test/conversation-manager.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect } from "vitest"; +import { describe, it, expect, vi } from "vitest"; import { ConversationManager } from "../src/conversation-manager.js"; import { Logger } from "copilot-sdk-proxy"; @@ -189,6 +189,19 @@ describe("ConversationManager", () => { expect(manager.getState(conv.id)).toBe(conv.state); expect(manager.getPrimary()).toBe(conv); }); + + // 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", () => { + const manager = createManager(); + const conv = manager.create(); + conv.state.session.markSessionActive(); + conv.state.toolRouter.registerExpected("toolu_01abc", "XcodeRead"); + conv.state.session.markSessionInactive(); + + expect(manager.size).toBe(1); + expect(manager.getState(conv.id)).toBe(conv.state); + }); }); describe("tool_use continuation after session goes inactive", () => { @@ -400,6 +413,20 @@ describe("ConversationManager", () => { expect(manager.getPrimary()).toBe(primary); }); + it("calls cleanup on evicted conversations", () => { + const manager = createManager(); + const primary = manager.create({ isPrimary: true }); + primary.session = { on: () => () => {} } as never; + + const isolated = manager.create(); + const cleanup = vi.spyOn(isolated.state.session, "cleanup"); + + manager.findForNewRequest(); + + expect(cleanup).toHaveBeenCalledOnce(); + expect(manager.getState(isolated.id)).toBeUndefined(); + }); + it("does NOT evict active non-primary conversations", () => { const manager = createManager(); const primary = manager.create({ isPrimary: true }); @@ -414,7 +441,9 @@ describe("ConversationManager", () => { expect(manager.size).toBe(3); }); - it("calls cleanup on evicted conversations", async () => { + // 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 () => { const manager = createManager(); const primary = manager.create({ isPrimary: true }); primary.session = { on: () => () => {} } as never; @@ -426,7 +455,12 @@ describe("ConversationManager", () => { }); manager.findForNewRequest(); - await expect(resultPromise).rejects.toThrow(); + + expect(manager.getState(isolated.id)).toBe(isolated.state); + expect(isolated.state.toolRouter.resolveToolCall("call-1", "ok")).toBe( + true, + ); + await expect(resultPromise).resolves.toBe("ok"); }); }); });