Skip to content

Conversation

@Nikunj0601
Copy link

Background

Currently, when tools yield multiple intermediate results using async generators, all states are shown as output-available with preliminary: true. The output property only contains the last yielded value, and the final return value from the async generator is lost. This makes it difficult for UI consumers to:

  • Distinguish between streaming progress and final results
  • Access the complete history of intermediate yields
  • Retrieve the actual final return value from the tool

Summary

This PR introduces enhanced tool progress tracking by:

  1. Adding output-streaming state: A new explicit state that indicates a tool is actively streaming intermediate results, distinct from output-available.

  2. Adding yields array: A new property that accumulates all yielded values during tool execution, allowing UI consumers to access the full history of intermediate results.

  3. Separating concerns: The output property now contains only the final returned value from the tool's execute function, while all intermediate yields are available in the yields array.

  4. Fixing async generator return value: Corrected executeTool to properly capture the final return value of async generators (when done === true) rather than the last yielded value.

The implementation uses a UI-only approach, interpreting existing tool-output-available chunks with preliminary: true as output-streaming states. This ensures full backward compatibility with existing streamText functionality.

Manual Verification

Checklist

  • Tests have been added / updated (for bug fixes / features)
  • Documentation has been added / updated (for bug fixes / features)
  • A patch changeset for relevant packages has been added (for bug fixes / features - run pnpm changeset in the project root)
  • I have reviewed this pull request (self-review)

Future Work

Related Issues

Fixes #9306

Copy link
Contributor

@vercel vercel bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Additional Comments:

packages/ai/src/ui/process-ui-message-stream.ts (lines 622-646):
When a tool transitions to the output-error state after yielding preliminary results, the yields are not cleared and persist on the tool part, violating the type contract that specifies yields should never exist in error states.

View Details
📝 Patch Details
diff --git a/packages/ai/src/ui/process-ui-message-stream.ts b/packages/ai/src/ui/process-ui-message-stream.ts
index dbff7597d..e9a3447ce 100644
--- a/packages/ai/src/ui/process-ui-message-stream.ts
+++ b/packages/ai/src/ui/process-ui-message-stream.ts
@@ -632,6 +632,16 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   input: (toolInvocation as any).input,
                   errorText: chunk.errorText,
                 });
+
+                // Clear yields for error state to comply with type contract
+                const part = state.message.parts.find(
+                  part =>
+                    part.type === 'dynamic-tool' &&
+                    part.toolCallId === chunk.toolCallId,
+                ) as any;
+                if (part) {
+                  delete part.yields;
+                }
               } else {
                 const toolInvocation = getToolInvocation(chunk.toolCallId);
 
@@ -643,6 +653,15 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   rawInput: (toolInvocation as any).rawInput,
                   errorText: chunk.errorText,
                 });
+
+                // Clear yields for error state to comply with type contract
+                const part = state.message.parts.find(
+                  part =>
+                    isToolUIPart(part) && part.toolCallId === chunk.toolCallId,
+                ) as any;
+                if (part) {
+                  delete part.yields;
+                }
               }
 
               write();

Analysis

Tool yields not cleared when transitioning to output-error state violates type contract

What fails: When a tool part transitions from output-streaming state (which has yields array) to output-error state, the yields field is not cleared from the part object, violating the type contract where DynamicToolUIPart specifies yields?: never; in the error case, and UIToolInvocation should forbid yields in error state.

How to reproduce:

  1. Create a tool call that enters output-streaming state with yields (preliminary results)
  2. Then transition that same tool call to output-error state
  3. Examine the tool part object - it will still contain the yields array from step 1

What happens: The tool-output-error case handler (lines 622-646) calls updateToolPart() or updateDynamicToolPart() without providing yields. The update functions only set yields if provided (if (anyOptions.yields != null)) and never clear existing yields. This leaves stale yields on the part after transitioning to error state.

Expected behavior: The yields field should be deleted/cleared when transitioning to output-error state to comply with the type contract that forbids yields in error states.

Root cause: Lines 193-194 in updateToolPart() and similar logic at lines 278-279 in updateDynamicToolPart() only set yields if provided in options, never clearing them:

if (anyOptions.yields != null) {
  anyPart.yields = anyOptions.yields;
}

When transitioning to error state, yields is not passed, so existing yields persist.

Type contract reference: ui-messages.ts lines 210-248 shows UIToolInvocation with output-error state having no yields field, and lines 264-312 show DynamicToolUIPart with yields?: never; explicitly forbidding yields in error state.


packages/ai/src/ui/ui-messages.ts (lines 247-255):
The UIToolInvocation type definition for output-error state is missing yields?: never, which is inconsistent with the DynamicToolUIPart type and would allow yields in error states contrary to the design intent.

View Details
📝 Patch Details
diff --git a/packages/ai/src/ui/ui-messages.ts b/packages/ai/src/ui/ui-messages.ts
index 107a452cd..e8ea9bedc 100644
--- a/packages/ai/src/ui/ui-messages.ts
+++ b/packages/ai/src/ui/ui-messages.ts
@@ -250,6 +250,7 @@ export type UIToolInvocation<TOOL extends UITool | Tool> = {
       rawInput?: unknown; // TODO AI SDK 6: remove this field, input should be unknown
       output?: never;
       errorText: string;
+      yields?: never;
       providerExecuted?: boolean;
       callProviderMetadata?: ProviderMetadata;
     }

Analysis

Missing yields?: never constraint in UIToolInvocation error state

What fails: The UIToolInvocation type definition for output-error state (lines 247-255 in packages/ai/src/ui/ui-messages.ts) is missing the yields?: never constraint, creating an inconsistency with DynamicToolUIPart and allowing yields properties in error states contrary to the design intent.

How to reproduce:

import { UIToolInvocation, UITool } from './packages/ai/src/ui/ui-messages';

type TestTool = UITool & {
  input: { value: string };
  output: string;
};

const toolInvocationWithYields: UIToolInvocation<TestTool> = {
  toolCallId: 'test-1',
  state: 'output-error',
  input: undefined,
  errorText: 'test error',
  yields: [], // Before fix: TypeScript allows this (incorrect)
};

Result: Before fix - TypeScript compiles successfully, allowing yields in error states.

Expected: After fix - TypeScript should error with "Type 'undefined[]' is not assignable to type 'never'" to match the DynamicToolUIPart constraint at line 304.

Implementation: Added yields?: never; to the output-error case of UIToolInvocation type definition to enforce consistency with DynamicToolUIPart and prevent yields in error states.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Feature request: add output-streaming to tool invocations

1 participant