fix(agent): render deferred tool UI results

Author: LukasParke

packages/agent/src/inner-loop/resume-tool-results.ts

@@ -221,6 +221,12 @@ export async function resumeToolResults<TTools extends readonly Tool[]>(
   const settledIds = new Set(state.settledAsyncCallIds ?? []);
 
   const envelopes: models.BaseInputsUnion[] = [];
+  const uiToolResults: Array<{
+    callId: string;
+    name: string;
+    input: Record<string, unknown>;
+    output: unknown;
+  }> = [];
   /** callId → the terminal lifecycle status persisted for that entry. */
   const settledNow = new Map<string, ToolTaskStatus>();
 
@@ -249,6 +255,7 @@ export async function resumeToolResults<TTools extends readonly Tool[]>(
     const envelope = buildResumeEnvelope(entry, task, request.tools);
 
     envelopes.push(buildTaskResultMessage(envelope));
+    collectUiToolResult(uiToolResults, envelope, task);
     // Persist the entry's real terminal status. 'expired' / 'timed_out'
     // have no ToolTaskStatus member — they persist as 'failed'.
     settledNow.set(
@@ -313,7 +320,7 @@ export async function resumeToolResults<TTools extends readonly Tool[]>(
 
   // Continue the conversation: the envelopes are already in persisted
   // history, so no fresh input is supplied.
-  return callModel(
+  const result = callModel(
     client,
     {
       ...request.run,
@@ -325,6 +332,28 @@ export async function resumeToolResults<TTools extends readonly Tool[]>(
     },
     options,
   );
+  result.queueUiToolResults(uiToolResults);
+  return result;
+}
+
+function collectUiToolResult(
+  results: Array<{
+    callId: string;
+    name: string;
+    input: Record<string, unknown>;
+    output: unknown;
+  }>,
+  envelope: ToolTaskResultEnvelope,
+  task: PendingAsyncTool,
+): void {
+  if (envelope.status === 'completed' && task.input !== undefined) {
+    results.push({
+      callId: task.callId,
+      name: task.name,
+      input: task.input,
+      output: envelope.result,
+    });
+  }
 }
 
 /**

packages/agent/src/lib/async-tool-registry.ts

@@ -147,6 +147,7 @@ export class AsyncToolRegistry {
     callId: string;
     taskId: string;
     name: string;
+    input: Record<string, unknown>;
     expiresAt?: number;
     pollAfterMs?: number;
   }): ToolTask {
@@ -155,6 +156,7 @@ export class AsyncToolRegistry {
       callId: entry.callId,
       toolName: entry.name,
       mode: 'defer',
+      input: entry.input,
       ...(entry.expiresAt !== undefined && {
         expiresAt: entry.expiresAt,
       }),
@@ -301,6 +303,9 @@ export class AsyncToolRegistry {
       mode: t.mode,
       status: t.status,
       startedAt: t.startedAt,
+      ...(t.input !== undefined && {
+        input: t.input,
+      }),
       ...(t.expiresAt !== undefined && {
         expiresAt: t.expiresAt,
       }),

packages/agent/src/lib/model-result.ts

@@ -707,6 +707,12 @@ export class ModelResult<
   > | null = null;
   private uiBroadcaster: ToolEventBroadcaster<ToolUiFragmentEvent> | null = null;
   private pendingUiFragments = new Set<Promise<void>>();
+  private queuedUiToolResults: Array<{
+    callId: string;
+    name: string;
+    input: Record<string, unknown>;
+    output: unknown;
+  }> = [];
   private uiBroadcasterCompletionPromise: Promise<void> | null = null;
   private turnBroadcasterCompletionPromise: Promise<void> | null = null;
   private initialStreamPipeStarted = false;
@@ -2775,10 +2781,12 @@ export class ModelResult<
     }
 
     const registry = this.ensureAsyncToolRegistry();
+    const input = (tc.arguments ?? {}) as Record<string, unknown>;
     const liveTask = registry.trackDeferred({
       callId: tc.id,
       taskId: invocation.taskId,
       name: String(tc.name),
+      input,
       ...(invocation.pollAfterMs !== undefined && {
         pollAfterMs: invocation.pollAfterMs,
       }),
@@ -2793,6 +2801,7 @@ export class ModelResult<
       mode: 'defer',
       status: 'working',
       startedAt: liveTask.startedAt,
+      input,
       ...(invocation.pollAfterMs !== undefined && {
         pollAfterMs: invocation.pollAfterMs,
       }),
@@ -3797,10 +3806,12 @@ export class ModelResult<
       if (collision) {
         return collision;
       }
+      const input = (toolCall.arguments ?? {}) as Record<string, unknown>;
       const liveTask = registry.trackDeferred({
         callId: toolCall.id,
         taskId: invocation.taskId,
         name: String(toolCall.name),
+        input,
         ...(invocation.expiresAt !== undefined && {
           expiresAt: invocation.expiresAt,
         }),
@@ -3816,6 +3827,7 @@ export class ModelResult<
         mode: 'defer',
         status: 'working',
         startedAt: liveTask.startedAt,
+        input,
         ...(invocation.pollAfterMs !== undefined && {
           pollAfterMs: invocation.pollAfterMs,
         }),
@@ -4422,6 +4434,41 @@ export class ModelResult<
     };
   }
 
+  /** @internal Queue externally resumed tool results for the UI lifecycle. */
+  queueUiToolResults(
+    results: Array<{
+      callId: string;
+      name: string;
+      input: Record<string, unknown>;
+      output: unknown;
+    }>,
+  ): void {
+    this.queuedUiToolResults.push(...results);
+  }
+
+  private dispatchQueuedUiToolResults(): void {
+    const queued = this.queuedUiToolResults;
+    this.queuedUiToolResults = [];
+    for (const result of queued) {
+      const tool = this.options.tools?.find(
+        (candidate) => isClientTool(candidate) && candidate.function.name === result.name,
+      );
+      if (tool) {
+        this.dispatchUiFragment({
+          toolCall: {
+            id: result.callId,
+            name: result.name,
+            arguments: result.input,
+          } as ParsedToolCall<Tool>,
+          tool,
+          result: {
+            result: result.output,
+          },
+        });
+      }
+    }
+  }
+
   private dispatchUiFragment(value: {
     toolCall: ParsedToolCall<Tool>;
     tool: Tool;
@@ -6945,9 +6992,8 @@ export class ModelResult<
    */
   getUiStream(): AsyncIterableIterator<UiStreamEvent> {
     return async function* (this: ModelResult<TTools, TShared>) {
-      await this.initStreamGuarded();
-
       if (!this.options.tools?.length) {
+        await this.initStreamGuarded();
         let streamFailed = false;
         try {
           if (this.reusableStream) {
@@ -6973,6 +7019,8 @@ export class ModelResult<
       const uiBroadcaster = this.uiBroadcaster;
       const uiConsumer = uiBroadcaster.createConsumer();
       try {
+        this.dispatchQueuedUiToolResults();
+        await this.initStreamGuarded();
         const { consumer, executionPromise } = this.startTurnBroadcasterExecution();
         if (!this.uiBroadcasterCompletionPromise) {
           this.uiBroadcasterCompletionPromise = executionPromise.finally(async () => {

packages/agent/src/lib/tool-types.ts

@@ -1622,6 +1622,8 @@ export interface PendingAsyncTool {
   status: ToolTaskStatus;
   /** Unix ms when the task started. */
   startedAt: number;
+  /** The original call arguments, retained for deferred UI rendering. */
+  input?: Record<string, unknown>;
   /** Unix ms after which the task is considered expired. */
   expiresAt?: number;
   /** Poll-interval hint surfaced to the model and external pollers. */

packages/agent/tests/unit/async-tool-deferred.test.ts

@@ -152,6 +152,9 @@ describe('tool.deferred — pause & placeholder', () => {
       name: 'request_legal_review',
       mode: 'defer',
       status: 'working',
+      input: {
+        contractId: 'c-9',
+      },
     });
 
     // No follow-up request — the loop paused after the placeholder.
@@ -724,6 +727,9 @@ describe('tool.deferred — cross-process resume', () => {
       callId: 'call_d1',
       taskId: 'ticket_c-9',
       mode: 'defer',
+      input: {
+        contractId: 'c-9',
+      },
     });
   });
 

packages/agent/tests/unit/async-tool-registry.test.ts

@@ -84,6 +84,35 @@ describe('AsyncToolRegistry — timeout settlement', () => {
   });
 });
 
+describe('AsyncToolRegistry — deferred input', () => {
+  it('retains input in settlement and persistence snapshots', () => {
+    const registry = new AsyncToolRegistry();
+    const task = registry.trackDeferred({
+      callId: 'call_d1',
+      taskId: 'task_d1',
+      name: 'weather',
+      input: {
+        city: 'Lisbon',
+      },
+    });
+
+    expect(task.input).toEqual({
+      city: 'Lisbon',
+    });
+    expect(registry.snapshot()[0]).toMatchObject({
+      input: {
+        city: 'Lisbon',
+      },
+    });
+    registry.cancelTask('task_d1');
+    expect(registry.takeSettled()[0]).toMatchObject({
+      input: {
+        city: 'Lisbon',
+      },
+    });
+  });
+});
+
 describe('AsyncToolRegistry — grace-window visibility (register/untrack)', () => {
   it('a registered (not yet tracked) task is reachable by steer and cancel', () => {
     const registry = new AsyncToolRegistry();