Preserve interrupted post-tool context

packages/core/agent-loop/src/agent.ts

@@ -8,7 +8,7 @@
 
 import type { Context } from 'cordis'
 import { agentEvents } from '@deepseek-ai/dsh-agent'
-import type { AgentId, AgentOptions, AgentStatus, SendOptions } from '@deepseek-ai/dsh-agent'
+import type { AgentId, AgentOptions, AgentStatus, HookContext, SendOptions } from '@deepseek-ai/dsh-agent'
 import type { Agent } from '@deepseek-ai/dsh-agent'
 import { deepFreeze } from '@deepseek-ai/dsh-llm'
 import type { ContentBlock, MessageSource } from '@deepseek-ai/dsh-llm'
@@ -288,11 +288,21 @@ export class ReactLoopAgent implements Agent {
     }
   }
 
-  /** Run one tool-call batch and drain its deferred context before resolving or rejecting. */
-  private async withToolBatch<T>(run: () => Promise<T>): Promise<T> {
+  /**
+   * Run one tool-call batch and drain its deferred context before settlement.
+   * The loop-owned acceptor remains valid after public disposal begins because
+   * the interrupted turn stays open until this batch settles.
+   */
+  private async withToolBatch<T>(
+    run: (acceptContext: (context: HookContext) => void) => Promise<T>,
+  ): Promise<T> {
     this.toolBatchActive = true
+    const acceptContext = (context: HookContext): void => {
+      const accepted = this.acceptMessage(context.content, { source: context.source })
+      this.deferredInjections.push(accepted)
+    }
     try {
-      return await run()
+      return await run(acceptContext)
     } finally {
       this.toolBatchActive = false
       this.drainDeferredInjections()

packages/core/agent-loop/src/loop.ts

@@ -85,8 +85,8 @@ export interface LoopHandle {
   clearCancel(): void
   /** Settle idle waiters when pre-running cancellation skips a turn, without emitting `agent/status`. */
   settleIdle(): void
-  /** Run an active tool-call batch and drain deferred context before resolving or rejecting. */
-  readonly withToolBatch: <T>(run: () => Promise<T>) => Promise<T>
+  /** Run an active tool-call batch, accepting post-tool context into the FIFO drained before settlement. */
+  readonly withToolBatch: <T>(run: (acceptContext: (context: HookContext) => void) => Promise<T>) => Promise<T>
 }
 
 /**
@@ -559,9 +559,7 @@ async function runStep(
   // Tool execution stays sequential; recheck abort around each normalized result.
   const toolCalls = message.content.filter(block => block.type === 'tool-call')
   if (toolCalls.length === 0) return { hadToolCalls: false, finish: assembler.finish }
-  return handle.withToolBatch(async () => {
-    // Buffer context until all results are appended to preserve call/result adjacency.
-    const pendingContext: HookContext[] = []
+  return handle.withToolBatch(async (acceptContext) => {
     for (const call of toolCalls) {
       /* v8 ignore next -- signal.reason always set: cancel()/disposal provide a default */
       if (signal.aborted) throw new Error(String(signal.reason ?? 'aborted'))
@@ -591,7 +589,9 @@ async function runStep(
         // Persist tool-owned presentation data for replay.
         ...result.meta !== undefined ? { meta: result.meta } : {},
       }, { surfaceOp: 'append', sourceEventSeqs: [callEvent.seq] })
-      if (result.additionalContext) pendingContext.push(result.additionalContext)
+      // Accept into the batch FIFO immediately; it remains deferred until every
+      // result settles and survives abort, cancellation, or disposal afterward.
+      if (result.additionalContext) acceptContext(result.additionalContext)
       // The signal may flip while the tool is awaited.
       /* v8 ignore start -- signal.reason default unreachable: cancel()/disposal always set it */
       // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
@@ -599,11 +599,6 @@ async function runStep(
       /* v8 ignore stop */
     }
 
-    // Append buffered context after the complete result batch.
-    for (const context of pendingContext) {
-      agent.inject(context.content, { source: context.source })
-    }
-
     return { hadToolCalls: true, finish: assembler.finish }
   })
 }

packages/core/agent-loop/tests/contract-regressions.spec.ts

@@ -3,7 +3,7 @@ import { Context } from 'cordis'
 import LlmService, { CallId, ContentBlock, MessageSource, StreamChunk } from '@deepseek-ai/dsh-llm'
 import SessionStore, { Session, SessionEvent, SessionId, TurnEndReason } from '@deepseek-ai/dsh-session'
 import SystemPrompt from '@deepseek-ai/dsh-system-prompt'
-import ToolRegistry, { defineTool } from '@deepseek-ai/dsh-tools'
+import ToolRegistry, { defineTool, type PostToolDecision } from '@deepseek-ai/dsh-tools'
 import AgentRegistry, { AgentId, type ContinuationDecision } from '@deepseek-ai/dsh-agent'
 import AgentLoop, { ReactLoopAgent } from '@deepseek-ai/dsh-agent-loop'
 import { prepareReactLoopAgent } from '../src/agent.ts'
@@ -154,6 +154,13 @@ describe('abort during tool execution ends the turn', () => {
         return [{ type: 'text', text: 'done' }]
       },
     }))
+    ctx.on('tools/post-execute', async (): Promise<PostToolDecision> => ({
+      kind: 'accept',
+      additionalContext: {
+        content: [{ type: 'text', text: 'accepted result context after abort' }],
+        source: { kind: 'plugin', plugin: 'test' },
+      },
+    }))
 
     send(agent, 'go')
     await waitForIdle(ctx, agent)
@@ -163,9 +170,65 @@ describe('abort during tool execution ends the turn', () => {
       .filter(event => event.type === 'tool/result' || event.type === 'context/message'
         || event.type === 'step/end' || event.type === 'turn/end')
       .map(event => event.type))
-      .toEqual(['tool/result', 'context/message', 'step/end', 'turn/end'])
+      .toEqual(['tool/result', 'context/message', 'context/message', 'step/end', 'turn/end'])
+    expect(events
+      .filter(event => event.type === 'context/message')
+      .map(event => event.data.content))
+      .toEqual([
+        [{ type: 'text', text: 'accepted before abort' }],
+        [{ type: 'text', text: 'accepted result context after abort' }],
+      ])
+  })
+
+  it('records post-tool context when a later call aborts the batch', async () => {
+    const adapter = new MockAdapter([[
+      { type: 'block-start', index: 0, blockType: 'tool-call' },
+      { type: 'block-end', index: 0, block: { type: 'tool-call', id: CallId('c1'), name: 'first', arguments: '{}' } },
+      { type: 'block-start', index: 1, blockType: 'tool-call' },
+      { type: 'block-end', index: 1, block: { type: 'tool-call', id: CallId('c2'), name: 'aborter', arguments: '{}' } },
+      { type: 'finish', reason: { kind: 'tool-calls' } },
+    ] satisfies StreamChunk[]])
+    const ctx = await harness(adapter)
+    const agent = ctx.agentLoop.create(AgentId('a-later-abort-context'), { model: 'mock' })
+    ctx.tools.register(defineTool({
+      name: 'first',
+      description: '',
+      parameters: {},
+      async execute() {
+        return [{ type: 'text', text: 'first done' }]
+      },
+    }))
+    ctx.tools.register(defineTool({
+      name: 'aborter',
+      description: '',
+      parameters: {},
+      async execute() {
+        ;(agent as unknown as { currentAbort?: AbortController }).currentAbort?.abort('user interrupt')
+        return [{ type: 'text', text: 'aborted' }]
+      },
+    }))
+    ctx.on('tools/post-execute', async (exec, _result, next): Promise<PostToolDecision> => {
+      if (exec.callId !== CallId('c1')) return next()
+      return {
+        kind: 'accept',
+        additionalContext: {
+          content: [{ type: 'text', text: 'accepted after first result' }],
+          source: { kind: 'plugin', plugin: 'test' },
+        },
+      }
+    })
+
+    send(agent, 'go')
+    await waitForIdle(ctx, agent)
+
+    const events = [...agent.session.events]
+    expect(events
+      .filter(event => event.type === 'tool/result' || event.type === 'context/message'
+        || event.type === 'step/end' || event.type === 'turn/end')
+      .map(event => event.type))
+      .toEqual(['tool/result', 'tool/result', 'context/message', 'step/end', 'turn/end'])
     expect(events.find(event => event.type === 'context/message')?.data.content)
-      .toEqual([{ type: 'text', text: 'accepted before abort' }])
+      .toEqual([{ type: 'text', text: 'accepted after first result' }])
   })
 
   it('drains deferred context before disposal reaches quiescence', async () => {
@@ -192,13 +255,25 @@ describe('abort during tool execution ends the turn', () => {
         return [{ type: 'text', text: 'done' }]
       },
     }))
+    ctx.on('tools/post-execute', async (): Promise<PostToolDecision> => ({
+      kind: 'accept',
+      additionalContext: {
+        content: [{ type: 'text', text: 'accepted result context during disposal' }],
+        source: { kind: 'plugin', plugin: 'test' },
+      },
+    }))
 
     send(agent, 'go')
     await started.promise
     await fiber.dispose()
 
-    expect(agent.session.events.find(event => event.type === 'context/message')?.data.content)
-      .toEqual([{ type: 'text', text: 'accepted before disposal' }])
+    expect(agent.session.events
+      .filter(event => event.type === 'context/message')
+      .map(event => event.data.content))
+      .toEqual([
+        [{ type: 'text', text: 'accepted before disposal' }],
+        [{ type: 'text', text: 'accepted result context during disposal' }],
+      ])
     expect(agent.session.events.find(event => event.type === 'turn/end')?.data.reason)
       .toEqual({ kind: 'disposed' })
   })