From ffee8fdbb2e0f559cc3b9db45302e0d768869d60 Mon Sep 17 00:00:00 2001 From: Louis Gu Date: Tue, 11 Aug 2026 17:21:21 +0800 Subject: [PATCH 1/3] fix(agent-core-v2): flush wire before agent disposal --- .changeset/calm-wires-persist.md | 5 +++ .../agentLifecycle/agentLifecycleService.ts | 6 ++-- .../agentLifecycle/agentLifecycle.test.ts | 32 +++++++++++++++++++ 3 files changed, 41 insertions(+), 2 deletions(-) create mode 100644 .changeset/calm-wires-persist.md diff --git a/.changeset/calm-wires-persist.md b/.changeset/calm-wires-persist.md new file mode 100644 index 0000000000..3a8fe8f9a5 --- /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 4ccbbdd6f5..e70d82bf2c 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -9,8 +9,9 @@ * Agent-scoped telemetry view. New logs receive a metadata * envelope while non-empty unversioned logs are rejected. Removal awaits the * agent task manager's graceful exit policy before draining turns and full - * compaction, then disposing the child scope. Fans session-level - * permission-mode switches out to every live agent. Bound at Session scope. + * compaction, flushing the replayable wire, then disposing the child scope. + * Fans session-level permission-mode switches out to every live agent. Bound + * at Session scope. * * No agent id is special here: the main agent is simply the agent created * with the conventional `MAIN_AGENT_ID`, and `fork` requires its source to @@ -268,6 +269,7 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle compaction.abortController.abort(reason); } await Promise.all([loop.settled(), compactionSettled]); + await handle.accessor.get(IWireService).flush(); handle.dispose(); this.onDidDisposeEmitter.fire(agentId); } 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 628b5af27d..204f052ea3 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -36,6 +36,7 @@ import { createMcpOAuthStore } from '#/app/mcpConfig/oauthStore'; import { ISessionSubagentService } from '#/session/subagent/subagent'; import { SessionSubagentService } from '#/session/subagent/subagentService'; import '#/agent/mcp/mcpService'; +import { IWireService } from '#/wire/wire'; import '#/wire/wireService'; import { IAgentTaskService } from '#/agent/task/task'; import { ISessionCronService } from '#/session/cron/sessionCronService'; @@ -455,6 +456,37 @@ describe('AgentLifecycleService', () => { expect(removed).toBe(true); }); + it('remove waits for the agent wire to flush before completing', async () => { + const svc = ix.get(IAgentLifecycleService); + const handle = await svc.create({ agentId: 'main' }); + let markFlushStarted!: () => void; + const flushStarted = new Promise((resolve) => { + markFlushStarted = resolve; + }); + let releaseFlush!: () => void; + const flushReleased = new Promise((resolve) => { + releaseFlush = resolve; + }); + const flush = vi + .spyOn(handle.accessor.get(IWireService), 'flush') + .mockImplementation(async () => { + markFlushStarted(); + await flushReleased; + }); + + let removed = false; + const removal = svc.remove('main').then(() => { + removed = true; + }); + await flushStarted; + expect(removed).toBe(false); + + releaseFlush(); + await removal; + expect(flush).toHaveBeenCalledOnce(); + expect(removed).toBe(true); + }); + it('ignites the self-wiring toolDedupe plugin so its listeners exist before the first turn', async () => { const svc = ix.get(IAgentLifecycleService); await svc.create({ agentId: 'main' }); From 1a0b87366c58f2bce45df7bff397413c45b167ce Mon Sep 17 00:00:00 2001 From: Louis Gu Date: Tue, 11 Aug 2026 17:27:45 +0800 Subject: [PATCH 2/3] fix(agent-core-v2): dispose agent after flush failures --- .../agentLifecycle/agentLifecycleService.ts | 9 ++++++--- .../session/agentLifecycle/agentLifecycle.test.ts | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index e70d82bf2c..5547db1e36 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -269,9 +269,12 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle compaction.abortController.abort(reason); } await Promise.all([loop.settled(), compactionSettled]); - await handle.accessor.get(IWireService).flush(); - handle.dispose(); - this.onDidDisposeEmitter.fire(agentId); + try { + await handle.accessor.get(IWireService).flush(); + } finally { + handle.dispose(); + this.onDidDisposeEmitter.fire(agentId); + } } } 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 204f052ea3..0642de13bf 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -487,6 +487,21 @@ describe('AgentLifecycleService', () => { expect(removed).toBe(true); }); + it('remove still disposes the agent when the wire flush rejects', async () => { + const svc = ix.get(IAgentLifecycleService); + const handle = await svc.create({ agentId: 'main' }); + const flushError = new Error('wire flush failed'); + vi.spyOn(handle.accessor.get(IWireService), 'flush').mockRejectedValueOnce(flushError); + const disposed: string[] = []; + disposables.add(svc.onDidDispose((id) => disposed.push(id))); + + await expect(svc.remove('main')).rejects.toBe(flushError); + + expect(svc.get('main')).toBeUndefined(); + expect(disposed).toEqual(['main']); + expect(() => handle.accessor.get(IWireService)).toThrow(); + }); + it('ignites the self-wiring toolDedupe plugin so its listeners exist before the first turn', async () => { const svc = ix.get(IAgentLifecycleService); await svc.create({ agentId: 'main' }); From 4f16c202e96119da6f870b175c0c8d0d2e939564 Mon Sep 17 00:00:00 2001 From: Louis Gu Date: Fri, 14 Aug 2026 01:08:57 +0800 Subject: [PATCH 3/3] fix(agent-core-v2): always dispose closing sessions --- .../sessionLifecycleService.ts | 13 +++++---- .../sessionLifecycle/sessionLifecycle.test.ts | 27 +++++++++++++++++++ 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts index 2eecd65f38..e486cd4d8c 100644 --- a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts +++ b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts @@ -399,11 +399,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/workspace/sessionLifecycle/sessionLifecycle.test.ts b/packages/agent-core-v2/test/workspace/sessionLifecycle/sessionLifecycle.test.ts index a848c9521b..3684bc729a 100644 --- a/packages/agent-core-v2/test/workspace/sessionLifecycle/sessionLifecycle.test.ts +++ b/packages/agent-core-v2/test/workspace/sessionLifecycle/sessionLifecycle.test.ts @@ -636,6 +636,33 @@ describe('SessionLifecycleService', () => { expect(svc.get('s1')).toBeUndefined(); }); + it('close disposes the session while preserving an agent removal failure', async () => { + const removeError = new Error('wire flush failed'); + const onTeardown = vi.fn(); + const agentHandle = { + id: 'main', + kind: LifecycleScope.Agent, + accessor: { get: () => ({}) }, + dispose: () => {}, + } as unknown as IAgentScopeHandle; + const svc = await build([ + stubPair(IAgentLifecycleService, { + ...agentLifecycleStub(), + list: () => [agentHandle], + remove: () => Promise.reject(removeError), + }), + ]); + svc.onWillCreateSession((event) => { + event.onSessionDispose(onTeardown); + }); + await svc.create({ sessionId: 's1', workDir: '/tmp/proj' }); + + await expect(svc.close('s1')).rejects.toBe(removeError); + + expect(svc.get('s1')).toBeUndefined(); + expect(onTeardown).toHaveBeenCalledOnce(); + }); + it('create seeds identity and materializes metadata', async () => { const svc = await build(); const h = await svc.create({ sessionId: 's1', workDir: '/tmp/proj' });