diff --git a/.changeset/calm-wires-persist.md b/.changeset/calm-wires-persist.md new file mode 100644 index 00000000000..3a8fe8f9a57 --- /dev/null +++ b/.changeset/calm-wires-persist.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Prevent recent ACP assistant responses from being lost when a session closes. diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index d2fae4f31d6..0628dda02ab 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -276,8 +276,12 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle compaction.abortController.abort(reason); } await Promise.all([loop.settled(), compactionSettled, prompt.drain(reason)]); - handle.dispose(); - this.onDidDisposeEmitter.fire(agentId); + try { + await handle.accessor.get(IEventDispatcher).flush(); + } finally { + handle.dispose(); + this.onDidDisposeEmitter.fire(agentId); + } } } diff --git a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts index 8445f83b75f..26bb92d0ed2 100644 --- a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts +++ b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts @@ -363,11 +363,14 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec if (handle === undefined) return; await this.announceWillClose({ sessionId, handle, reason: 'exit' }); this.sessions.delete(sessionId); - await this.drainAgents(handle); - await drainSessionMetadataWrites(); - await this.indexMirror.drain(); - handle.dispose(); - this._onDidCloseSession.fire({ sessionId }); + try { + await this.drainAgents(handle); + await drainSessionMetadataWrites(); + await this.indexMirror.drain(); + } finally { + handle.dispose(); + this._onDidCloseSession.fire({ sessionId }); + } } async archive(sessionId: string): Promise { diff --git a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts index d204da09a4f..b591b1160a1 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -492,6 +492,50 @@ describe('AgentLifecycleService', () => { expect(disposed).toEqual(['main']); }); + it('remove waits for pending agent state to flush before disposing the agent scope', async () => { + const svc = ix.get(IAgentLifecycleService); + const handle = await svc.create({ agentId: 'main' }); + let releaseFlush!: () => void; + let markFlushStarted!: () => void; + const flushStarted = new Promise((resolve) => { + markFlushStarted = resolve; + }); + const flush = vi.spyOn(handle.accessor.get(IEventDispatcher), 'flush').mockImplementation(() => { + markFlushStarted(); + return new Promise((resolve) => { + releaseFlush = resolve; + }); + }); + const disposed: string[] = []; + disposables.add(svc.onDidDispose((agentId) => disposed.push(agentId))); + + const removal = svc.remove('main'); + await flushStarted; + await Promise.resolve(); + + expect(disposed).toEqual([]); + + releaseFlush(); + await removal; + expect(flush).toHaveBeenCalledOnce(); + expect(disposed).toEqual(['main']); + }); + + it('remove still disposes the agent when the state flush rejects', async () => { + const svc = ix.get(IAgentLifecycleService); + const handle = await svc.create({ agentId: 'main' }); + const flushError = new Error('state flush failed'); + vi.spyOn(handle.accessor.get(IEventDispatcher), 'flush').mockRejectedValueOnce(flushError); + const disposed: string[] = []; + disposables.add(svc.onDidDispose((agentId) => disposed.push(agentId))); + + await expect(svc.remove('main')).rejects.toBe(flushError); + + expect(svc.get('main')).toBeUndefined(); + expect(disposed).toEqual(['main']); + expect(() => handle.accessor.get(IEventDispatcher)).toThrow(); + }); + it('remove cancels queued turns before waiting for the active turn to settle', async () => { loopActiveTurnId = 1; loopPendingTurnIds = [2, 3];