Fix deferred injection lifecycle

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 {