-
Notifications
You must be signed in to change notification settings - Fork 3.2k
feat(ui): add output streaming to tool invocations #9960
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
feat(ui): add output streaming to tool invocations #9960
Conversation
…s in preliminary steps Signed-off-by: Nikunj <[email protected]>
Signed-off-by: Nikunj <[email protected]>
There was a problem hiding this 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:
- Create a tool call that enters
output-streamingstate with yields (preliminary results) - Then transition that same tool call to
output-errorstate - 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.
Signed-off-by: Nikunj <[email protected]>
Signed-off-by: Nikunj <[email protected]>
Background
Currently, when tools yield multiple intermediate results using async generators, all states are shown as
output-availablewithpreliminary: true. Theoutputproperty 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:Summary
This PR introduces enhanced tool progress tracking by:
Adding
output-streamingstate: A new explicit state that indicates a tool is actively streaming intermediate results, distinct fromoutput-available.Adding
yieldsarray: A new property that accumulates all yielded values during tool execution, allowing UI consumers to access the full history of intermediate results.Separating concerns: The
outputproperty now contains only the final returned value from the tool'sexecutefunction, while all intermediate yields are available in theyieldsarray.Fixing async generator return value: Corrected
executeToolto properly capture the final return value of async generators (whendone === true) rather than the last yielded value.The implementation uses a UI-only approach, interpreting existing
tool-output-availablechunks withpreliminary: trueasoutput-streamingstates. This ensures full backward compatibility with existingstreamTextfunctionality.Manual Verification
Checklist
pnpm changesetin the project root)Future Work
Related Issues
Fixes #9306