Fix deferred injection lifecycle

packages/core/agent-loop/README.md

@@ -44,7 +44,7 @@ Configured agents start automatically. `cwd` applies only to fresh sessions; `re
 
 - `ReactLoopAgent` — the concrete `Agent` implementation. Its inbox is a JavaScript native-private field, and one prepared session can be claimed by only one concrete driver. Everything observable happens through session events and the `agent/*` event taxonomy.
 
-`Inbox`, `runLoop`, and the instance-bound publication/start controls are package-internal. The package root does not export them, and the package exports map exposes no `./src/*` escape hatch; lifecycle owners create agents through `ctx.agents` rather than constructing or starting the driver internals. `ReactLoopAgent.send()`, running `steer()`, and open-turn `inject()` materialize content plus resolved source once as detached, deeply frozen lossless JSON; malformed data throws before enqueue or append. An injection that arrives while assistant tool calls await results stays in a FIFO until the complete result batch is logged, keeping provider tool messages contiguous.
+`Inbox`, `runLoop`, and the instance-bound publication/start controls are package-internal. The package root does not export them, and the package exports map exposes no `./src/*` escape hatch; lifecycle owners create agents through `ctx.agents` rather than constructing or starting the driver internals. `ReactLoopAgent.send()`, running `steer()`, and open-turn `inject()` materialize content plus resolved source once as detached, deeply frozen lossless JSON; malformed data throws before enqueue or append. An injection that arrives while the current step executes assistant tool calls stays in a FIFO until the batch settles; successful batches place it after the complete result batch, and interrupted batches drain it before the turn closes.
 
 ### Loop lifecycle (`loop.ts`)
 

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

@@ -12,7 +12,7 @@ import type { AgentId, AgentOptions, AgentStatus, SendOptions } from '@deepseek-
 import type { Agent } from '@deepseek-ai/dsh-agent'
 import { deepFreeze } from '@deepseek-ai/dsh-llm'
 import type { ContentBlock, MessageSource } from '@deepseek-ai/dsh-llm'
-import { isToolPairingBalanced, snapshotJsonValue, type Session } from '@deepseek-ai/dsh-session'
+import { snapshotJsonValue, type Session } from '@deepseek-ai/dsh-session'
 import { Inbox, type InboxMessage } from './inbox.ts'
 import { isTurnOpen, lastTurnNumber, runLoop } from './loop.ts'
 
@@ -149,6 +149,8 @@ export class ReactLoopAgent implements Agent {
    * this set before the lifecycle unregisters the agent or detaches its session.
    */
   private pendingIdleFlushes = new Set<Promise<void>>()
+  /** Whether the current step is executing an assistant tool-call batch. */
+  private toolBatchActive = false
   /** Open-turn injections waiting for the active assistant tool-call batch to close. */
   private deferredInjections: InboxMessage[] = []
 
@@ -231,11 +233,9 @@ export class ReactLoopAgent implements Agent {
     if (isTurnOpen(this.session)) {
       const accepted = this.acceptMessage(content, options)
       // Provider protocols require every assistant tool-call batch to be
-      // followed only by its tool results. Queue arbitrary asynchronous context
-      // until the tail cut is balanced; an existing queue preserves FIFO in the
-      // narrow window after the last result and before the loop drains it.
-      if (this.deferredInjections.length > 0
-        || !isToolPairingBalanced(this.session.surface.nodes, this.session.events, null)) {
+      // followed only by its tool results. Historical interrupted batches do
+      // not own new context; only the currently executing batch may defer it.
+      if (this.toolBatchActive) {
         this.deferredInjections.push(accepted)
         return
       }
@@ -288,6 +288,17 @@ 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> {
+    this.toolBatchActive = true
+    try {
+      return await run()
+    } finally {
+      this.toolBatchActive = false
+      this.drainDeferredInjections()
+    }
+  }
+
   cancel(reason?: string): void {
     // Arm only for current work; an idle marker would cancel the next prompt.
     if (this._status === 'running' || this.currentAbort !== undefined || this.#inbox.hasQueued || this.#inbox.hasSteering) {
@@ -352,7 +363,7 @@ export class ReactLoopAgent implements Agent {
       isCancelled: () => this.cancelRequested,
       cancelReason: () => this.cancelReason,
       clearCancel: () => { this.cancelRequested = false },
-      drainDeferredInjections: () => { this.drainDeferredInjections() },
+      withToolBatch: run => this.withToolBatch(run),
       // Pre-step cancellation re-parks without emitting a status transition.
       settleIdle: () => { this.settleIdleWaiters() },
     })

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
-  /** Append context that arrived while an assistant tool-call batch was awaiting its results. */
-  readonly drainDeferredInjections: () => void
+  /** Run an active tool-call batch and drain deferred context before resolving or rejecting. */
+  readonly withToolBatch: <T>(run: () => Promise<T>) => Promise<T>
 }
 
 /**
@@ -327,7 +327,7 @@ async function runTurn(
       let stepOutcome: { hadToolCalls: boolean; finish: FinishReason } | { error: Error }
       try {
         stepOutcome = await runStep(
-          ctx, events, agent, turn, step, assembly, fullSystemPrompt, boundaryMessages, transmission, abort.signal)
+          ctx, events, agent, handle, turn, step, assembly, fullSystemPrompt, boundaryMessages, transmission, abort.signal)
       } catch (error: unknown) {
         stepOutcome = { error: toError(error) }
       } finally {
@@ -351,10 +351,6 @@ async function runTurn(
         break
       }
 
-      // A successful tool step has committed its complete result batch. Context
-      // accepted while that batch was pending can now join the open turn.
-      if (stepOutcome.hadToolCalls) handle.drainDeferredInjections()
-
       // Preserve max-token completion unless a later disposal, abort, or error wins.
       const stepReason = stepFinishReason(stepOutcome.finish)
       if (stepReason) reason = stepReason
@@ -468,6 +464,7 @@ async function runStep(
   ctx: Context,
   events: AgentEventDispatch,
   agent: ReactLoopAgent,
+  handle: LoopHandle,
   turn: number,
   step: number,
   assembly: PromptAssembly,
@@ -561,51 +558,54 @@ async function runStep(
 
   // Tool execution stays sequential; recheck abort around each normalized result.
   const toolCalls = message.content.filter(block => block.type === 'tool-call')
-  // Buffer context until all results are appended to preserve call/result adjacency.
-  const pendingContext: HookContext[] = []
-  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'))
-    const callEvent = session.append('tool/call', { turn, step, callId: call.id, name: call.name, arguments: call.arguments })
-    let parsedArguments: unknown
-    try {
-      parsedArguments = call.arguments ? JSON.parse(call.arguments) : {}
-    } catch {
-      parsedArguments = call.arguments
+  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[] = []
+    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'))
+      const callEvent = session.append('tool/call', { turn, step, callId: call.id, name: call.name, arguments: call.arguments })
+      let parsedArguments: unknown
+      try {
+        parsedArguments = call.arguments ? JSON.parse(call.arguments) : {}
+      } catch {
+        parsedArguments = call.arguments
+      }
+      // TODO(pre-tool-input-rewrite): Keep logged history and live presentation aligned;
+      // see docs/rfc/proposed/feature/2026-06-30-pre-tool-input-rewrite.md.
+      const result = await ctx.tools.execute({
+        callId: call.id,
+        name: call.name,
+        arguments: parsedArguments,
+        agent,
+        signal,
+      })
+      session.append('tool/result', {
+        turn, step,
+        // Preserve transcript pairing even if a post-execute listener returns another id.
+        callId: call.id,
+        content: result.content,
+        isError: result.isError,
+        ...result.error ? { error: result.error } : {},
+        // 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)
+      // 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
+      if (signal.aborted) throw new Error(String(signal.reason ?? 'aborted'))
+      /* v8 ignore stop */
     }
-    // TODO(pre-tool-input-rewrite): Keep logged history and live presentation aligned;
-    // see docs/rfc/proposed/feature/2026-06-30-pre-tool-input-rewrite.md.
-    const result = await ctx.tools.execute({
-      callId: call.id,
-      name: call.name,
-      arguments: parsedArguments,
-      agent,
-      signal,
-    })
-    session.append('tool/result', {
-      turn, step,
-      // Preserve transcript pairing even if a post-execute listener returns another id.
-      callId: call.id,
-      content: result.content,
-      isError: result.isError,
-      ...result.error ? { error: result.error } : {},
-      // 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)
-    // 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
-    if (signal.aborted) throw new Error(String(signal.reason ?? 'aborted'))
-    /* v8 ignore stop */
-  }
 
-  // Append buffered context after the complete result batch.
-  for (const context of pendingContext) {
-    agent.inject(context.content, { source: context.source })
-  }
+    // Append buffered context after the complete result batch.
+    for (const context of pendingContext) {
+      agent.inject(context.content, { source: context.source })
+    }
 
-  return { hadToolCalls: toolCalls.length > 0, finish: assembler.finish }
+    return { hadToolCalls: true, finish: assembler.finish }
+  })
 }
 
 function withoutToolCalls(message: Message): Message {

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

@@ -139,6 +139,115 @@ describe('abort during tool execution ends the turn', () => {
     expect(adapter.requests).toHaveLength(1)        // no follow-up model call
     expect(reasons).toEqual([{ kind: 'aborted', reason: 'user interrupt' }])
   })
+
+  it('records context accepted before a tool-step abort in the same turn', async () => {
+    const adapter = new MockAdapter([toolCallResponse('c1', 'aborter', {})])
+    const ctx = await harness(adapter)
+    const agent = ctx.agentLoop.create(AgentId('a-abort-injection'), { model: 'mock' })
+    ctx.tools.register(defineTool({
+      name: 'aborter',
+      description: '',
+      parameters: {},
+      async execute() {
+        agent.inject([{ type: 'text', text: 'accepted before abort' }], { source: { kind: 'plugin', plugin: 'test' } })
+        ;(agent as unknown as { currentAbort?: AbortController }).currentAbort?.abort('user interrupt')
+        return [{ type: 'text', text: 'done' }]
+      },
+    }))
+
+    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', 'context/message', 'step/end', 'turn/end'])
+    expect(events.find(event => event.type === 'context/message')?.data.content)
+      .toEqual([{ type: 'text', text: 'accepted before abort' }])
+  })
+
+  it('drains deferred context before disposal reaches quiescence', async () => {
+    const adapter = new MockAdapter([toolCallResponse('c1', 'waiter', {})])
+    const ctx = await harness(adapter)
+    const started = Promise.withResolvers<undefined>()
+    let agent!: ReactLoopAgent
+    const fiber = await ctx.plugin(Object.assign((inner: Context) => {
+      agent = inner.agentLoop.create(AgentId('a-dispose-injection'), { model: 'mock' })
+    }, { inject: ['agentLoop'] }))
+    ctx.tools.register(defineTool({
+      name: 'waiter',
+      description: '',
+      parameters: {},
+      async execute(_args, exec) {
+        agent.inject([{ type: 'text', text: 'accepted before disposal' }], { source: { kind: 'plugin', plugin: 'test' } })
+        started.resolve(undefined)
+        const signal = exec.signal
+        if (!signal) throw new Error('tool execution signal is missing')
+        await new Promise<void>((resolve) => {
+          if (signal.aborted) resolve()
+          else signal.addEventListener('abort', () => { resolve() }, { once: true })
+        })
+        return [{ type: 'text', text: 'done' }]
+      },
+    }))
+
+    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.find(event => event.type === 'turn/end')?.data.reason)
+      .toEqual({ kind: 'disposed' })
+  })
+
+  it('limits injection deferral to the current tool 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: 'aborter', arguments: '{}' } },
+        { type: 'block-start', index: 1, blockType: 'tool-call' },
+        { type: 'block-end', index: 1, block: { type: 'tool-call', id: CallId('c2'), name: 'second', arguments: '{}' } },
+        { type: 'finish', reason: { kind: 'tool-calls' } },
+      ] satisfies StreamChunk[],
+      textResponse('later turn'),
+    ])
+    const ctx = await harness(adapter)
+    const agent = ctx.agentLoop.create(AgentId('a-historical-tool-pair'), { model: 'mock' })
+    ctx.tools.register(defineTool({
+      name: 'aborter',
+      description: '',
+      parameters: {},
+      async execute() {
+        ;(agent as unknown as { currentAbort?: AbortController }).currentAbort?.abort('user interrupt')
+        return [{ type: 'text', text: 'done' }]
+      },
+    }))
+    ctx.tools.register(defineTool({
+      name: 'second',
+      description: '',
+      parameters: {},
+      async execute() {
+        return [{ type: 'text', text: 'must not run' }]
+      },
+    }))
+
+    send(agent, 'leave an unmatched historical call')
+    await waitForIdle(ctx, agent)
+    ctx.on('agent/pre-step', (subject, turn) => {
+      if (subject === agent && turn === 2) {
+        agent.inject([{ type: 'text', text: 'new turn context' }], { source: { kind: 'plugin', plugin: 'test' } })
+      }
+    })
+    send(agent, 'start a text-only turn')
+    await waitForIdle(ctx, agent)
+
+    expect(agent.session.events.find(event => event.type === 'context/message')?.data.content)
+      .toEqual([{ type: 'text', text: 'new turn context' }])
+    expect(JSON.stringify(adapter.requests[1]?.messages)).toContain('new turn context')
+  })
 })
 
 describe('steering from late extension points is never stranded', () => {