diff --git a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts index f9693fd817..052516d52d 100644 --- a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts +++ b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts @@ -3335,6 +3335,133 @@ describe('AiSdkBackend model history', () => { assert.match(JSON.stringify(assistant), /Maka shipped the feature/); }); + test('preserves pending assistant steps before a following client tool step', async () => { + const prompt = await replayPrompt([ + runtimeTextEvent({ + id: 'rt-user', + turnId: 'turn-prev', + role: 'user', + author: 'user', + text: 'inspect the workspace', + }), + runtimeEvent({ + id: 'rt-progress-a', + turnId: 'turn-prev', + role: 'model', + author: 'agent', + refs: { providerEventId: 'progress-step-a' }, + content: { kind: 'text', text: 'I found the relevant package.' }, + }), + runtimeEvent({ + id: 'rt-progress-b', + turnId: 'turn-prev', + role: 'model', + author: 'agent', + refs: { providerEventId: 'progress-step-b' }, + content: { kind: 'text', text: 'I will inspect its configuration.' }, + }), + clientToolCallEvent('rt-read-call', 'tool-step'), + clientToolResultEvent('rt-read-result'), + ]); + + assert.deepEqual( + prompt.slice(0, 5).map((message) => ({ + role: message.role, + types: message.content.map((part) => part.type), + text: message.content.find((part) => part.type === 'text')?.text, + })), + [ + { role: 'user', types: ['text'], text: 'inspect the workspace' }, + { + role: 'assistant', + types: ['text'], + text: 'I found the relevant package.', + }, + { + role: 'assistant', + types: ['text'], + text: 'I will inspect its configuration.', + }, + { role: 'assistant', types: ['tool-call'], text: undefined }, + { role: 'tool', types: ['tool-result'], text: undefined }, + ], + ); + }); + + test('groups a client tool only with the immediately pending matching step', async () => { + const contiguousPrompt = await replayPrompt([ + runtimeTextEvent({ + id: 'rt-user', + turnId: 'turn-prev', + role: 'user', + author: 'user', + text: 'read the file', + }), + runtimeEvent({ + id: 'rt-progress', + turnId: 'turn-prev', + role: 'model', + author: 'agent', + refs: { providerEventId: 'shared-step' }, + content: { kind: 'text', text: 'I will read the file now.' }, + }), + clientToolCallEvent('rt-read-call', 'shared-step'), + clientToolResultEvent('rt-read-result'), + ]); + assert.deepEqual( + contiguousPrompt.slice(0, 3).map((message) => ({ + role: message.role, + types: message.content.map((part) => part.type), + })), + [ + { role: 'user', types: ['text'] }, + { role: 'assistant', types: ['text', 'tool-call'] }, + { role: 'tool', types: ['tool-result'] }, + ], + ); + + const interruptedPrompt = await replayPrompt([ + runtimeTextEvent({ + id: 'rt-user', + turnId: 'turn-prev', + role: 'user', + author: 'user', + text: 'read the file', + }), + runtimeEvent({ + id: 'rt-progress-a', + turnId: 'turn-prev', + role: 'model', + author: 'agent', + refs: { providerEventId: 'shared-step' }, + content: { kind: 'text', text: 'I will read the file now.' }, + }), + runtimeEvent({ + id: 'rt-progress-b', + turnId: 'turn-prev', + role: 'model', + author: 'agent', + refs: { providerEventId: 'intervening-step' }, + content: { kind: 'text', text: 'Another step was persisted.' }, + }), + clientToolCallEvent('rt-read-call', 'shared-step'), + clientToolResultEvent('rt-read-result'), + ]); + assert.deepEqual( + interruptedPrompt.slice(0, 5).map((message) => ({ + role: message.role, + types: message.content.map((part) => part.type), + })), + [ + { role: 'user', types: ['text'] }, + { role: 'assistant', types: ['text'] }, + { role: 'assistant', types: ['text'] }, + { role: 'assistant', types: ['tool-call'] }, + { role: 'tool', types: ['tool-result'] }, + ], + ); + }); + test('falls back to grounded text when Open Responses cannot replay a hosted tool pair', async () => { const model = completionModel(); const backend = createTestAiSdkBackend({ @@ -16040,6 +16167,65 @@ function runtimeEvent(input: { }; } +function clientToolCallEvent(id: string, stepId: string): RuntimeEvent { + return runtimeEvent({ + id, + turnId: 'turn-prev', + role: 'model', + author: 'agent', + refs: { stepId }, + content: { + kind: 'function_call', + id: 'read-1', + name: 'Read', + args: { path: 'notes.md' }, + }, + }); +} + +function clientToolResultEvent(id: string): RuntimeEvent { + return runtimeEvent({ + id, + turnId: 'turn-prev', + role: 'tool', + author: 'tool', + content: { + kind: 'function_response', + id: 'read-1', + name: 'Read', + result: { kind: 'text', text: 'file contents' }, + isError: false, + }, + }); +} + +async function replayPrompt( + runtimeContext: RuntimeEvent[], +): Promise> { + const model = completionModel(); + const backend = createTestAiSdkBackend({ + sessionId: 'session-1', + header: header(), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + modelFactory: () => model, + tools: [], + newId: idGenerator(), + now: monotonicClock(), + }); + await drain( + backend.send({ + turnId: 'turn-current', + text: 'continue', + context: [], + runtimeContext, + }), + ); + return compactPrompt(model) as Array<{ role: string; content: any[] }>; +} + function compactPrompt(model: MockLanguageModelV4): unknown { return model.doStreamCalls[0]?.prompt.map((message) => ({ role: message.role, diff --git a/packages/runtime/src/ai-sdk-backend.ts b/packages/runtime/src/ai-sdk-backend.ts index 8f641c56dd..8fc7daf013 100644 --- a/packages/runtime/src/ai-sdk-backend.ts +++ b/packages/runtime/src/ai-sdk-backend.ts @@ -3899,6 +3899,10 @@ export class AiSdkBackend implements AgentBackend { >(); const reasoningByStep = new Map(); const textByStep = new Map(); + const pendingStepOrder = new Set(); + const rememberPendingStep = (stepId: string) => { + pendingStepOrder.add(stepId); + }; const replaySupport = this.modelAdapter.runtimeEventReplaySupport(); const reasoningReplay = (item: ThinkingItem): ReplayReasoning | undefined => { @@ -4157,6 +4161,7 @@ export class AiSdkBackend implements AgentBackend { if (stepId !== undefined) { reasoningByStep.delete(stepId); textByStep.delete(stepId); + pendingStepOrder.delete(stepId); } await emitStep(reasoning, text, group); group = []; @@ -4173,18 +4178,28 @@ export class AiSdkBackend implements AgentBackend { bufferedCalls = []; await emitGroupedCalls(calls); }; + const flushPendingStep = async (stepId: string) => { + const text = textByStep.get(stepId); + const reasoning = reasoningByStep.get(stepId); + textByStep.delete(stepId); + reasoningByStep.delete(stepId); + pendingStepOrder.delete(stepId); + await emitStep(reasoning, text, []); + }; + const flushPendingStepsBefore = async (stepId: string | undefined) => { + const pendingStepIds = [...pendingStepOrder]; + const lastPendingStepId = pendingStepIds.at(-1); + const earlierStepIds = + stepId !== undefined && lastPendingStepId === stepId + ? pendingStepIds.slice(0, -1) + : pendingStepIds; + if (earlierStepIds.length === 0) return; + await flushLooseCalls(); + for (const pendingStepId of earlierStepIds) await flushPendingStep(pendingStepId); + }; const flushPendingSteps = async () => { await flushLooseCalls(); - for (const [stepId, text] of textByStep) { - textByStep.delete(stepId); - const reasoning = reasoningByStep.get(stepId); - reasoningByStep.delete(stepId); - await emitStep(reasoning, text, []); - } - for (const [stepId, reasoning] of reasoningByStep) { - reasoningByStep.delete(stepId); - await emitStep(reasoning, undefined, []); - } + for (const stepId of [...pendingStepOrder]) await flushPendingStep(stepId); }; for (const item of admitProviderReasoningReplayItems( @@ -4193,6 +4208,7 @@ export class AiSdkBackend implements AgentBackend { )) { switch (item.kind) { case 'tool_call': + await flushPendingStepsBefore(item.stepId); if (item.toolName !== 'apply_patch') { bufferedCalls.push(item); break; @@ -4248,6 +4264,7 @@ export class AiSdkBackend implements AgentBackend { ts: downgradedCall.ts, }); } + rememberPendingStep(downgradedCall.stepId); } else { await flushPendingSteps(); push({ role: 'assistant', content: [{ type: 'text', text: replayFact }] }, [ @@ -4262,6 +4279,7 @@ export class AiSdkBackend implements AgentBackend { const stepReasoning = reasoningByStep.get(item.stepId) ?? []; stepReasoning.push(item); reasoningByStep.set(item.stepId, stepReasoning); + rememberPendingStep(item.stepId); } else { // Legacy standalone reasoning (pure-reasoning turn): emit on its own. await flushPendingSteps(); @@ -4297,11 +4315,13 @@ export class AiSdkBackend implements AgentBackend { if (thisCalls.length > 0) { await emitStep(reasoningByStep.get(stepId), item, thisCalls); reasoningByStep.delete(stepId); + pendingStepOrder.delete(stepId); } else { // Runtime-owned settlement persists assistant facts before the // matching tool calls. Hold the step closer until those calls // arrive; a terminal text-only step flushes below. textByStep.set(stepId, item); + rememberPendingStep(stepId); } } else { // Legacy per-turn assistant text: standalone after any tool block.