fix(acp): stream engine-triggered follow-up turns

.changeset/fresh-pumas-stream.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix ACP clients missing follow-up messages after background tasks complete.

packages/acp-server/src/session.ts

@@ -167,9 +167,9 @@ interface TurnDriver {
   cancelRequested?: boolean;
   /**
    * Turn-scoped events that arrived while `turnId` was still unknown. Once
-   * the launch resolves, the entries matching the driver's turn are replayed
-   * in arrival order; the rest (a still-draining prior turn) are dropped —
-   * the same verdict the live path would have given.
+   * the launch resolves, all entries are replayed in arrival order. Only the
+   * matching turn may settle the driver; every event still belongs to this
+   * session's main-agent stream and must reach the ACP client.
    */
   early: Array<{ readonly turnId: number; readonly dispatch: () => void }>;
 }
@@ -639,9 +639,9 @@ export class AcpSession {
 
   /**
    * Submit the prompt and drive the turn to completion: `agent.prompt()`
-   * returns the launched turn id, which the session-level event handlers use
-   * to attribute events to this driver. Settles on `turn.ended`; a no-launch
-   * result (hook-blocked / not runnable) settles with `end_turn`.
+   * returns the launched turn id, which identifies the `turn.ended` event that
+   * settles this driver. A no-launch result (hook-blocked / not runnable)
+   * settles with `end_turn`.
    */
   private driveTurn(input: readonly ContentPart[]): Promise<PromptResponse> {
     this.assertNoActiveTurn();
@@ -650,8 +650,8 @@ export class AcpSession {
 
   /**
    * Shared turn settlement for every launch path (`agent.prompt`,
-   * `agent.activateSkill`): the returned turn id attributes subsequent events
-   * to this driver; `undefined` means no turn launched (hook-blocked / not
+   * `agent.activateSkill`): the returned turn id identifies which turn may
+   * settle this driver; `undefined` means no turn launched (hook-blocked / not
    * runnable — the busy case never gets this far, see
    * {@link assertNoActiveTurn}), so the prompt settles gracefully with
    * `end_turn`.
@@ -664,6 +664,7 @@ export class AcpSession {
         (launched) => {
           if (driver.settled) return;
           if (launched === undefined) {
+            this.replayEarlyEvents(driver);
             // No turn will emit `turn.ended`, so settle gracefully. The engine
             // publishes a `prompt.completed` with reason 'blocked' for the
             // hook-blocked case; the wire carries no blocking message to
@@ -688,11 +689,10 @@ export class AcpSession {
           // Replay the events the turn emitted before its id arrived (a fast
           // turn can outrun the launch round-trip — `activateSkill` returns
           // only after the prompt-metadata update).
-          for (const early of driver.early.splice(0)) {
-            if (early.turnId === driver.turnId) early.dispatch();
-          }
+          this.replayEarlyEvents(driver);
         },
         (error) => {
+          this.replayEarlyEvents(driver);
           this.settleDriver(driver, () => {
             reject(mapPromptLaunchError(error, this.sessionId));
           });
@@ -714,9 +714,13 @@ export class AcpSession {
     dispatch();
   }
 
+  private replayEarlyEvents(driver: TurnDriver): void {
+    for (const early of driver.early.splice(0)) early.dispatch();
+  }
+
   /**
    * Settle the driver exactly once and detach it from the session so later
-   * events of its turn are ignored.
+   * events cannot affect the client prompt response.
    */
   private settleDriver(driver: TurnDriver, action: () => void): void {
     if (driver.settled) return;
@@ -725,7 +729,7 @@ export class AcpSession {
     action();
   }
 
-  /** The active driver, but only for events of ITS turn. */
+  /** The active driver, but only for settling its client-initiated turn. */
   private driverFor(turnId: number): TurnDriver | undefined {
     const driver = this.driver;
     if (driver === undefined || driver.turnId === undefined || driver.turnId !== turnId) {
@@ -735,17 +739,14 @@ export class AcpSession {
   }
 
   private onAssistantDelta(event: AgentEventPayloads['assistant.delta']): void {
-    if (this.driverFor(event.turnId) === undefined) return;
     this.emit(assistantDeltaToSessionUpdate(this.sessionId, event));
   }
 
   private onThinkingDelta(event: AgentEventPayloads['thinking.delta']): void {
-    if (this.driverFor(event.turnId) === undefined) return;
     this.emit(thinkingDeltaToSessionUpdate(this.sessionId, event));
   }
 
   private onToolCallStarted(event: AgentEventPayloads['tool.call.started']): void {
-    if (this.driverFor(event.turnId) === undefined) return;
     // The klient payload mirrors `ToolCallStartedEvent` (`args` / `display`
     // arrive as `unknown` — cast at this seam).
     const mapped = event as unknown as ToolCallStartedEvent;
@@ -787,7 +788,6 @@ export class AcpSession {
   }
 
   private onToolCallDelta(event: AgentEventPayloads['tool.call.delta']): void {
-    if (this.driverFor(event.turnId) === undefined) return;
     // The klient payload mirrors `ToolCallDeltaEvent` field-for-field.
     const mapped = event as unknown as ToolCallDeltaEvent;
     const key = acpToolCallId(event.turnId, event.toolCallId);
@@ -809,15 +809,13 @@ export class AcpSession {
   }
 
   private onToolProgress(event: AgentEventPayloads['tool.progress']): void {
-    if (this.driverFor(event.turnId) === undefined) return;
     // The klient payload mirrors `ToolProgressEvent` field-for-field; the
     // helper forwards only `status` updates with text (as a title refresh)
     // and returns null for everything else, which `emit` drops.
     this.emit(toolProgressToSessionUpdate(this.sessionId, event as unknown as ToolProgressEvent));
   }
 
   private onToolResult(event: AgentEventPayloads['tool.result']): void {
-    if (this.driverFor(event.turnId) === undefined) return;
     const key = acpToolCallId(event.turnId, event.toolCallId);
     const locations = this.toolLocations.get(key);
     this.toolLocations.delete(key);
@@ -906,17 +904,20 @@ export class AcpSession {
 
   private onTurnEnded(event: AgentEventPayloads['turn.ended']): void {
     const driver = this.driverFor(event.turnId);
-    if (driver === undefined) return;
-    const error = event.error as { readonly code: string; readonly message?: string } | undefined;
-    this.settleDriver(driver, () => {
-      // Auth failures must surface as a JSON-RPC `auth_required` error
-      // so the client triggers its re-auth flow, not a silent `end_turn`.
-      if (event.reason === 'failed' && isAuthError(error)) {
-        driver.reject(RequestError.authRequired(undefined, error?.message));
-        return;
-      }
-      driver.resolve({ stopReason: turnEndReasonToStopReason(event.reason, error) });
-    });
+    if (driver !== undefined) {
+      const error = event.error as
+        | { readonly code: string; readonly message?: string }
+        | undefined;
+      this.settleDriver(driver, () => {
+        // Auth failures must surface as a JSON-RPC `auth_required` error
+        // so the client triggers its re-auth flow, not a silent `end_turn`.
+        if (event.reason === 'failed' && isAuthError(error)) {
+          driver.reject(RequestError.authRequired(undefined, error?.message));
+          return;
+        }
+        driver.resolve({ stopReason: turnEndReasonToStopReason(event.reason, error) });
+      });
+    }
     void this.emitUsageUpdate();
   }
 

packages/acp-server/test/e2e-turn.test.ts

@@ -15,9 +15,9 @@ import { join } from 'node:path';
 import { fileURLToPath } from 'node:url';
 
 import { getLiveSessionById, IAgentLifecycleService, IEventBus } from '@moonshot-ai/agent-core-v2';
-import { afterEach, describe, expect, it } from 'vitest';
+import { afterEach, describe, expect, it, vi } from 'vitest';
 
-import { mapPromptLaunchError } from '../src/session';
+import { AcpSession, mapPromptLaunchError } from '../src/session';
 import { createTestClient, type TestClient } from './_helpers/acpClient';
 import { writeFakeModelConfig } from './_helpers/fakeModelConfig';
 import { solidPngBase64 } from './_helpers/png';
@@ -28,6 +28,72 @@ const STDIO_MCP_FIXTURE = fileURLToPath(
   new URL('../../agent-core-v2/test/mcpCore/fixtures/mock-stdio-server.mjs', import.meta.url),
 );
 
+interface TurnDriverHarness {
+  driveLaunch(
+    launch: Promise<{ readonly turn_id: number } | undefined>,
+  ): Promise<{ readonly stopReason: string }>;
+  dispatchTurnEvent(turnId: number, dispatch: () => void): void;
+  onTurnEnded(event: { readonly turnId: number; readonly reason: 'completed' }): void;
+}
+
+function createTurnDriverHarness(): TurnDriverHarness {
+  const session = Object.create(AcpSession.prototype) as AcpSession;
+  Object.defineProperty(session, 'emitUsageUpdate', {
+    value: async (): Promise<void> => {},
+  });
+  return session as unknown as TurnDriverHarness;
+}
+
+describe('AcpSession turn driver buffering', () => {
+  it('replays unrelated events without letting them settle the client prompt', async () => {
+    const session = createTurnDriverHarness();
+    let resolveLaunch!: (result: { readonly turn_id: number }) => void;
+    const prompt = session.driveLaunch(
+      new Promise((resolve) => {
+        resolveLaunch = resolve;
+      }),
+    );
+    let promptSettled = false;
+    void prompt.then(() => {
+      promptSettled = true;
+    });
+    const delivered: string[] = [];
+
+    session.dispatchTurnEvent(99, () => delivered.push('background message'));
+    session.dispatchTurnEvent(99, () => {
+      session.onTurnEnded({ turnId: 99, reason: 'completed' });
+    });
+    resolveLaunch({ turn_id: 100 });
+
+    await vi.waitFor(() => {
+      expect(delivered).toEqual(['background message']);
+    });
+    expect(promptSettled).toBe(false);
+
+    session.dispatchTurnEvent(100, () => {
+      session.onTurnEnded({ turnId: 100, reason: 'completed' });
+    });
+    await expect(prompt).resolves.toEqual({ stopReason: 'end_turn' });
+  });
+
+  it('replays buffered events when the client prompt launches no turn', async () => {
+    const session = createTurnDriverHarness();
+    let resolveLaunch!: (result: undefined) => void;
+    const prompt = session.driveLaunch(
+      new Promise((resolve) => {
+        resolveLaunch = resolve;
+      }),
+    );
+    const delivered: string[] = [];
+
+    session.dispatchTurnEvent(99, () => delivered.push('background message'));
+    resolveLaunch(undefined);
+
+    await expect(prompt).resolves.toEqual({ stopReason: 'end_turn' });
+    expect(delivered).toEqual(['background message']);
+  });
+});
+
 describe('acp-server real prompt turn (scripted LLM)', () => {
   let homeDir: string | undefined;
   let client: TestClient | undefined;
@@ -91,6 +157,78 @@ describe('acp-server real prompt turn (scripted LLM)', () => {
     expect(usageUpdate?.cost).toBeUndefined();
   }, 30_000);
 
+  it('streams an engine-triggered follow-up turn after the client prompt settles', async () => {
+    const c = await boot();
+    scripted!.mockNextText('background task started');
+
+    const created = (await c.send('session/new', { cwd: homeDir, mcpServers: [] })) as {
+      sessionId: string;
+    };
+    await c.waitForSessionUpdate('available_commands_update', 10_000);
+
+    const result = (await c.send('session/prompt', {
+      sessionId: created.sessionId,
+      prompt: [{ type: 'text', text: 'start a background task' }],
+    })) as { stopReason: string };
+    expect(result.stopReason).toBe('end_turn');
+
+    const updatesBeforeFollowUp = c.sessionUpdates().length;
+    const session = getLiveSessionById(c.server.core.accessor, created.sessionId);
+    const agentHandle = session?.accessor.get(IAgentLifecycleService).get('main');
+    const bus = agentHandle?.accessor.get(IEventBus);
+    expect(bus).toBeDefined();
+
+    bus!.publish({
+      type: 'assistant.delta',
+      turnId: 99,
+      delta: 'background task finished',
+    });
+    bus!.publish({
+      type: 'tool.call.started',
+      turnId: 99,
+      toolCallId: 'background-read',
+      name: 'Read',
+      args: { path: '/tmp/background-output.log' },
+    });
+    bus!.publish({
+      type: 'tool.result',
+      turnId: 99,
+      toolCallId: 'background-read',
+      output: 'done',
+    });
+    bus!.publish({ type: 'turn.ended', turnId: 99, reason: 'completed' });
+
+    await vi.waitFor(() => {
+      const updates = c
+        .sessionUpdates()
+        .slice(updatesBeforeFollowUp)
+        .map((message) => message.params as { update?: Record<string, unknown> });
+      expect(updates).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({
+            update: expect.objectContaining({
+              sessionUpdate: 'agent_message_chunk',
+              content: { type: 'text', text: 'background task finished' },
+            }),
+          }),
+          expect.objectContaining({
+            update: expect.objectContaining({
+              sessionUpdate: 'tool_call',
+              toolCallId: '99:background-read',
+            }),
+          }),
+          expect.objectContaining({
+            update: expect.objectContaining({
+              sessionUpdate: 'tool_call_update',
+              toolCallId: '99:background-read',
+              status: 'completed',
+            }),
+          }),
+        ]),
+      );
+    });
+  }, 30_000);
+
   it('runs a tool call and bridges the approval request to the client', async () => {
     const c = await boot();
     // First model response: a Bash tool call. Second: a short text wrap-up