From d3edf87ce99b83082ec580006c0b51d32ad7fc5b Mon Sep 17 00:00:00 2001 From: zz Date: Fri, 1 May 2026 23:52:24 -0700 Subject: [PATCH] Add option to suppress tool input deltas Signed-off-by: zz --- .changeset/quiet-tools-stream.md | 5 + packages/ai/src/agent/do-stream-step.test.ts | 106 ++++++++++++++++++ packages/ai/src/agent/do-stream-step.ts | 12 +- ...e-agent-suppress-tool-input-deltas.test.ts | 51 +++++++++ packages/ai/src/agent/durable-agent.ts | 10 ++ packages/ai/src/agent/stream-text-iterator.ts | 7 ++ 6 files changed, 190 insertions(+), 1 deletion(-) create mode 100644 .changeset/quiet-tools-stream.md create mode 100644 packages/ai/src/agent/durable-agent-suppress-tool-input-deltas.test.ts diff --git a/.changeset/quiet-tools-stream.md b/.changeset/quiet-tools-stream.md new file mode 100644 index 0000000000..4c8f7952c5 --- /dev/null +++ b/.changeset/quiet-tools-stream.md @@ -0,0 +1,5 @@ +--- +"@workflow/ai": minor +--- + +Add `suppressToolInputDeltas` to `DurableAgent.stream()` to avoid writing incremental tool input deltas to UI streams while preserving final tool inputs and execution state. diff --git a/packages/ai/src/agent/do-stream-step.test.ts b/packages/ai/src/agent/do-stream-step.test.ts index 28f6122881..0d0041cd6e 100644 --- a/packages/ai/src/agent/do-stream-step.test.ts +++ b/packages/ai/src/agent/do-stream-step.test.ts @@ -3,6 +3,73 @@ import { describe, expect, it } from 'vitest'; import { doStreamStep, normalizeFinishReason } from './do-stream-step.js'; import { safeParseToolCallInput } from './safe-parse-tool-call-input.js'; +function createToolInputDeltaModel() { + return async () => + ({ + provider: 'mock', + modelId: 'mock-tool-input-deltas', + doStream: async () => ({ + stream: new ReadableStream({ + start(controller) { + for (const part of [ + { type: 'stream-start', warnings: [] }, + { + type: 'response-metadata', + id: 'response-1', + modelId: 'mock-tool-input-deltas', + timestamp: new Date(), + }, + { + type: 'tool-input-start', + id: 'call-1', + toolName: 'writeFile', + }, + { + type: 'tool-input-delta', + id: 'call-1', + delta: '{"path":"src/App.tsx"', + }, + { + type: 'tool-input-delta', + id: 'call-1', + delta: ',"content":"x"}', + }, + { type: 'tool-input-end', id: 'call-1' }, + { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'writeFile', + input: '{"path":"src/App.tsx","content":"x"}', + }, + { + type: 'finish', + finishReason: { type: 'tool-calls' }, + usage: { + inputTokens: { total: 1 }, + outputTokens: { total: 1 }, + }, + }, + ]) { + controller.enqueue(part); + } + controller.close(); + }, + }), + }), + }) as any; +} + +function createChunkCollector() { + const chunks: any[] = []; + const writable = new WritableStream({ + write(chunk) { + chunks.push(chunk); + }, + }); + + return { chunks, writable }; +} + describe('normalizeFinishReason', () => { describe('string finish reasons', () => { it('should pass through "stop"', () => { @@ -198,4 +265,43 @@ describe('doStreamStep', () => { }) ); }); + + it('can suppress tool input deltas without dropping final tool input', async () => { + const { chunks, writable } = createChunkCollector(); + + const result = await doStreamStep( + [{ role: 'user', content: [{ type: 'text', text: 'write a file' }] }], + createToolInputDeltaModel(), + writable, + [], + { suppressToolInputDeltas: true } + ); + + expect(chunks.map((chunk) => chunk.type)).toEqual([ + 'start-step', + 'tool-input-start', + 'tool-input-available', + 'finish-step', + ]); + expect(result.toolCalls[0]?.input).toBe( + '{"path":"src/App.tsx","content":"x"}' + ); + expect(result.step.toolCalls[0]?.input).toEqual({ + path: 'src/App.tsx', + content: 'x', + }); + }); + + it('streams tool input deltas by default', async () => { + const { chunks, writable } = createChunkCollector(); + + await doStreamStep( + [{ role: 'user', content: [{ type: 'text', text: 'write a file' }] }], + createToolInputDeltaModel(), + writable, + [] + ); + + expect(chunks.map((chunk) => chunk.type)).toContain('tool-input-delta'); + }); }); diff --git a/packages/ai/src/agent/do-stream-step.ts b/packages/ai/src/agent/do-stream-step.ts index 15ba779dad..8bfc29baee 100644 --- a/packages/ai/src/agent/do-stream-step.ts +++ b/packages/ai/src/agent/do-stream-step.ts @@ -16,12 +16,12 @@ import { type ToolSet, type UIMessageChunk, } from 'ai'; +import { getErrorMessage } from '../get-error-message.js'; import type { ProviderOptions, StreamTextTransform, TelemetrySettings, } from './durable-agent.js'; -import { getErrorMessage } from '../get-error-message.js'; import { safeParseToolCallInput } from './safe-parse-tool-call-input.js'; import { recordSpan } from './telemetry.js'; import type { CompatibleLanguageModel } from './types.js'; @@ -71,6 +71,11 @@ export interface DoStreamStepOptions { providerOptions?: ProviderOptions; toolChoice?: ToolChoice; includeRawChunks?: boolean; + /** + * If true, suppresses incremental tool-input-delta UIMessageChunks while + * preserving raw stream processing and final tool input availability. + */ + suppressToolInputDeltas?: boolean; experimental_telemetry?: TelemetrySettings; transforms?: Array>; responseFormat?: LanguageModelV3CallOptions['responseFormat']; @@ -221,6 +226,7 @@ export async function doStreamStep( const chunks: LanguageModelV3StreamPart[] = []; const includeRawChunks = options?.includeRawChunks ?? false; const collectUIChunks = options?.collectUIChunks ?? false; + const suppressToolInputDeltas = options?.suppressToolInputDeltas ?? false; const uiChunks: UIMessageChunk[] = []; let msToFirstChunk: number | undefined; @@ -437,6 +443,10 @@ export async function doStreamStep( } case 'tool-input-delta': { + if (suppressToolInputDeltas) { + break; + } + controller.enqueue({ type: 'tool-input-delta', toolCallId: part.id, diff --git a/packages/ai/src/agent/durable-agent-suppress-tool-input-deltas.test.ts b/packages/ai/src/agent/durable-agent-suppress-tool-input-deltas.test.ts new file mode 100644 index 0000000000..af3c09f2be --- /dev/null +++ b/packages/ai/src/agent/durable-agent-suppress-tool-input-deltas.test.ts @@ -0,0 +1,51 @@ +import type { LanguageModelV3 } from '@ai-sdk/provider'; +import { describe, expect, it, vi } from 'vitest'; + +vi.mock('./stream-text-iterator.js', () => ({ + streamTextIterator: vi.fn(), +})); + +const { DurableAgent } = await import('./durable-agent.js'); + +function createMockModel(): LanguageModelV3 { + return { + specificationVersion: 'v3' as const, + provider: 'test', + modelId: 'test-model', + doGenerate: vi.fn(), + doStream: vi.fn(), + supportedUrls: {}, + }; +} + +describe('DurableAgent suppressToolInputDeltas', () => { + it('passes suppressToolInputDeltas to streamTextIterator', async () => { + const mockModel = createMockModel(); + const agent = new DurableAgent({ + model: async () => mockModel, + tools: {}, + }); + const mockWritable = new WritableStream({ + write: vi.fn(), + close: vi.fn(), + }); + + const { streamTextIterator } = await import('./stream-text-iterator.js'); + const mockIterator = { + next: vi.fn().mockResolvedValueOnce({ done: true, value: [] }), + }; + vi.mocked(streamTextIterator).mockReturnValue(mockIterator as any); + + await agent.stream({ + messages: [{ role: 'user', content: 'test' }], + writable: mockWritable, + suppressToolInputDeltas: true, + }); + + expect(streamTextIterator).toHaveBeenCalledWith( + expect.objectContaining({ + suppressToolInputDeltas: true, + }) + ); + }); +}); diff --git a/packages/ai/src/agent/durable-agent.ts b/packages/ai/src/agent/durable-agent.ts index a84d7214dc..7a027a0c81 100644 --- a/packages/ai/src/agent/durable-agent.ts +++ b/packages/ai/src/agent/durable-agent.ts @@ -580,6 +580,15 @@ export interface DurableAgentStreamOptions< */ includeRawChunks?: boolean; + /** + * If true, suppresses incremental `tool-input-delta` UI stream chunks. + * + * This only affects UI stream persistence/transport. The raw model stream, + * final parsed tool input, tool execution, and conversation state are + * preserved. + */ + suppressToolInputDeltas?: boolean; + /** * A function that attempts to repair a tool call that failed to parse. */ @@ -1007,6 +1016,7 @@ export class DurableAgent { experimental_context: experimentalContext, experimental_telemetry: effectiveTelemetry, includeRawChunks: options.includeRawChunks ?? false, + suppressToolInputDeltas: options.suppressToolInputDeltas ?? false, experimental_transform: options.experimental_transform as | StreamTextTransform | Array>, diff --git a/packages/ai/src/agent/stream-text-iterator.ts b/packages/ai/src/agent/stream-text-iterator.ts index 07b7e87bd5..6e341ecc9d 100644 --- a/packages/ai/src/agent/stream-text-iterator.ts +++ b/packages/ai/src/agent/stream-text-iterator.ts @@ -80,6 +80,7 @@ export async function* streamTextIterator({ experimental_context, experimental_telemetry, includeRawChunks = false, + suppressToolInputDeltas = false, experimental_transform, responseFormat, collectUIChunks = false, @@ -99,6 +100,11 @@ export async function* streamTextIterator({ experimental_context?: unknown; experimental_telemetry?: TelemetrySettings; includeRawChunks?: boolean; + /** + * If true, suppresses incremental tool-input-delta UIMessageChunks while + * preserving raw stream processing and final tool input availability. + */ + suppressToolInputDeltas?: boolean; experimental_transform?: | StreamTextTransform | Array>; @@ -305,6 +311,7 @@ export async function* streamTextIterator({ ...currentGenerationSettings, toolChoice: currentToolChoice, includeRawChunks, + suppressToolInputDeltas, experimental_telemetry, transforms, responseFormat,