fix(agent-core-v2): flush wire before agent disposal

.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.

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);
+    }
   }
 }
 

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<void> {

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<void>((resolve) => {
+      markFlushStarted = resolve;
+    });
+    const flush = vi.spyOn(handle.accessor.get(IEventDispatcher), 'flush').mockImplementation(() => {
+      markFlushStarted();
+      return new Promise<void>((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];