From 9426b01d3930414c9ac333b33c67b56a11d3c2ee Mon Sep 17 00:00:00 2001 From: Nathan Rajlich Date: Sat, 14 Mar 2026 10:31:55 -0700 Subject: [PATCH] Fix prepareStep system message lost when messages is also returned Apply messages override before system override so that the system message is prepended to the new prompt instead of being discarded. Previously, system was applied first then messages replaced the entire conversation prompt, losing the system message. Add tests for prepareStep system/messages ordering: - system only: prepended to prompt - system + messages: system prepended to replaced messages - system + messages with existing system: replaces existing - system updates across multi-step tool call rounds --- .../fix-prepare-step-system-ordering.md | 5 + .../ai/src/agent/stream-text-iterator.test.ts | 217 +++++++++++++++++- packages/ai/src/agent/stream-text-iterator.ts | 10 +- 3 files changed, 227 insertions(+), 5 deletions(-) create mode 100644 .changeset/fix-prepare-step-system-ordering.md diff --git a/.changeset/fix-prepare-step-system-ordering.md b/.changeset/fix-prepare-step-system-ordering.md new file mode 100644 index 0000000000..dca364b916 --- /dev/null +++ b/.changeset/fix-prepare-step-system-ordering.md @@ -0,0 +1,5 @@ +--- +'@workflow/ai': patch +--- + +Fix `prepareStep` system message being discarded when `messages` is also returned diff --git a/packages/ai/src/agent/stream-text-iterator.test.ts b/packages/ai/src/agent/stream-text-iterator.test.ts index 6284bc7bf4..a5b6e4236b 100644 --- a/packages/ai/src/agent/stream-text-iterator.test.ts +++ b/packages/ai/src/agent/stream-text-iterator.test.ts @@ -10,9 +10,10 @@ import type { LanguageModelV3Prompt, LanguageModelV3ToolCall, LanguageModelV3ToolResult, + LanguageModelV3ToolResultPart, } from '@ai-sdk/provider'; import type { StepResult, ToolSet, UIMessageChunk } from 'ai'; -import { describe, expect, it, vi, beforeEach } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; // Mock doStreamStep vi.mock('./do-stream-step.js', () => ({ @@ -651,4 +652,218 @@ describe('streamTextIterator', () => { }); }); }); + + describe('prepareStep system and messages ordering', () => { + it('should apply system message when prepareStep returns only system', async () => { + const mockWritable = createMockWritable(); + const mockModel = vi.fn(); + + let capturedPrompt: LanguageModelV3Prompt | undefined; + + vi.mocked(doStreamStep).mockImplementationOnce(async (prompt) => { + capturedPrompt = prompt; + return { + toolCalls: [], + finish: { finishReason: 'stop' }, + step: createMockStepResult({ finishReason: 'stop' }), + }; + }); + + const iterator = streamTextIterator({ + prompt: [{ role: 'user', content: [{ type: 'text', text: 'hello' }] }], + tools: {} as ToolSet, + writable: mockWritable, + model: mockModel as any, + prepareStep: () => ({ + system: 'You are a helpful assistant.', + }), + }); + + await iterator.next(); + + expect(capturedPrompt).toBeDefined(); + expect(capturedPrompt![0]).toEqual({ + role: 'system', + content: 'You are a helpful assistant.', + }); + expect(capturedPrompt![1]).toEqual({ + role: 'user', + content: [{ type: 'text', text: 'hello' }], + }); + }); + + it('should preserve system message when prepareStep returns both system and messages', async () => { + const mockWritable = createMockWritable(); + const mockModel = vi.fn(); + + let capturedPrompt: LanguageModelV3Prompt | undefined; + + vi.mocked(doStreamStep).mockImplementationOnce(async (prompt) => { + capturedPrompt = prompt; + return { + toolCalls: [], + finish: { finishReason: 'stop' }, + step: createMockStepResult({ finishReason: 'stop' }), + }; + }); + + // prepareStep returns both system and messages — system should NOT be lost + const customMessages: LanguageModelV3Prompt = [ + { + role: 'user', + content: [{ type: 'text', text: 'modified message' }], + }, + ]; + + const iterator = streamTextIterator({ + prompt: [ + { role: 'user', content: [{ type: 'text', text: 'original' }] }, + ], + tools: {} as ToolSet, + writable: mockWritable, + model: mockModel as any, + prepareStep: () => ({ + system: 'Dynamic system prompt.', + messages: customMessages, + }), + }); + + await iterator.next(); + + expect(capturedPrompt).toBeDefined(); + // System message should be prepended to the replaced messages + expect(capturedPrompt!).toHaveLength(2); + expect(capturedPrompt![0]).toEqual({ + role: 'system', + content: 'Dynamic system prompt.', + }); + expect(capturedPrompt![1]).toEqual({ + role: 'user', + content: [{ type: 'text', text: 'modified message' }], + }); + }); + + it('should replace existing system message when messages already contains one', async () => { + const mockWritable = createMockWritable(); + const mockModel = vi.fn(); + + let capturedPrompt: LanguageModelV3Prompt | undefined; + + vi.mocked(doStreamStep).mockImplementationOnce(async (prompt) => { + capturedPrompt = prompt; + return { + toolCalls: [], + finish: { finishReason: 'stop' }, + step: createMockStepResult({ finishReason: 'stop' }), + }; + }); + + // Messages already include a system message — prepareStep's system should replace it + const customMessages: LanguageModelV3Prompt = [ + { role: 'system', content: 'Old system prompt.' }, + { role: 'user', content: [{ type: 'text', text: 'hello' }] }, + ]; + + const iterator = streamTextIterator({ + prompt: [ + { role: 'user', content: [{ type: 'text', text: 'original' }] }, + ], + tools: {} as ToolSet, + writable: mockWritable, + model: mockModel as any, + prepareStep: () => ({ + system: 'New system prompt.', + messages: customMessages, + }), + }); + + await iterator.next(); + + expect(capturedPrompt).toBeDefined(); + expect(capturedPrompt!).toHaveLength(2); + expect(capturedPrompt![0]).toEqual({ + role: 'system', + content: 'New system prompt.', + }); + expect(capturedPrompt![1]).toEqual({ + role: 'user', + content: [{ type: 'text', text: 'hello' }], + }); + }); + + it('should update system message on subsequent steps', async () => { + const mockWritable = createMockWritable(); + const mockModel = vi.fn(); + + const capturedPrompts: LanguageModelV3Prompt[] = []; + + const toolCall: LanguageModelV3ToolCall = { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'testTool', + input: '{}', + }; + + vi.mocked(doStreamStep) + .mockImplementationOnce(async (prompt) => { + capturedPrompts.push([...prompt]); + return { + toolCalls: [toolCall], + finish: { finishReason: 'tool-calls' }, + step: createMockStepResult({ finishReason: 'tool-calls' }), + }; + }) + .mockImplementationOnce(async (prompt) => { + capturedPrompts.push([...prompt]); + return { + toolCalls: [], + finish: { finishReason: 'stop' }, + step: createMockStepResult({ finishReason: 'stop' }), + }; + }); + + const iterator = streamTextIterator({ + prompt: [{ role: 'user', content: [{ type: 'text', text: 'hello' }] }], + tools: { + testTool: { + description: 'Test', + execute: async () => ({ ok: true }), + }, + } as ToolSet, + writable: mockWritable, + model: mockModel as any, + prepareStep: ({ stepNumber: sn }) => ({ + system: `System prompt v${sn}`, + }), + }); + + // First step + await iterator.next(); + + // Provide tool results + const toolResults: LanguageModelV3ToolResultPart[] = [ + { + type: 'tool-result', + toolCallId: 'call-1', + toolName: 'testTool', + output: { type: 'text', value: '{"ok":true}' }, + }, + ]; + + // Second step + await iterator.next(toolResults); + + expect(capturedPrompts).toHaveLength(2); + // First step should have system v0 + expect(capturedPrompts[0][0]).toEqual({ + role: 'system', + content: 'System prompt v0', + }); + // Second step should have system v1 + expect(capturedPrompts[1][0]).toEqual({ + role: 'system', + content: 'System prompt v1', + }); + }); + }); }); diff --git a/packages/ai/src/agent/stream-text-iterator.ts b/packages/ai/src/agent/stream-text-iterator.ts index 6ff9b2fee5..754ac8ebd6 100644 --- a/packages/ai/src/agent/stream-text-iterator.ts +++ b/packages/ai/src/agent/stream-text-iterator.ts @@ -148,8 +148,13 @@ export async function* streamTextIterator({ if (prepareResult.model !== undefined) { currentModel = prepareResult.model; } + if (prepareResult.messages !== undefined) { + conversationPrompt = [...prepareResult.messages]; + } if (prepareResult.system !== undefined) { - // Update or prepend system message in the conversation prompt + // Update or prepend system message in the conversation prompt. + // Applied AFTER messages override so the system message isn't + // lost when messages replaces the prompt. if ( conversationPrompt.length > 0 && conversationPrompt[0].role === 'system' @@ -167,9 +172,6 @@ export async function* streamTextIterator({ }); } } - if (prepareResult.messages !== undefined) { - conversationPrompt = [...prepareResult.messages]; - } if (prepareResult.experimental_context !== undefined) { currentContext = prepareResult.experimental_context; }