Skip to content

fix(runtime): preserve cross-step replay chronology #4750

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.

Sign up for GitHub

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

Merged
Astro-Han merged 1 commit into from
Sep 4, 2026
Merged

fix(runtime): preserve cross-step replay chronology #4750

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
186 changes: 186 additions & 0 deletions packages/runtime/src/__tests__/ai-sdk-backend.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down Expand Up @@ -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<Array<{ role: string; content: any[] }>> {
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,
Expand Down
40 changes: 30 additions & 10 deletions packages/runtime/src/ai-sdk-backend.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3899,6 +3899,10 @@ export class AiSdkBackend implements AgentBackend {
>();
const reasoningByStep = new Map<string, ThinkingItem[]>();
const textByStep = new Map<string, TextItem>();
const pendingStepOrder = new Set<string>();
const rememberPendingStep = (stepId: string) => {
pendingStepOrder.add(stepId);
};

const replaySupport = this.modelAdapter.runtimeEventReplaySupport();
const reasoningReplay = (item: ThinkingItem): ReplayReasoning | undefined => {
Expand Down Expand Up @@ -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 = [];
Expand All @@ -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(
Expand All @@ -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;
Expand Down Expand Up @@ -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 }] }, [
Expand All @@ -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();
Expand Down Expand Up @@ -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.
Expand Down