mirror of
https://github.com/n8n-io/n8n.git
synced 2026-08-28 17:22:01 +08:00
feat(core): Add per-step checkpoints and crash resume to the agents SDK (no-changelog) (#33845)
This commit is contained in:
committed by
GitHub
parent
2ef112b883
commit
1fd337a734
@@ -0,0 +1,220 @@
|
||||
/**
|
||||
* Durable-log RFC (resilience phase), end to end through the public SDK
|
||||
* surface: run with `stepCheckpoints` on, crash mid-run, crash-resume from the
|
||||
* persisted step checkpoint with a fresh Agent instance.
|
||||
*
|
||||
* The model is scripted by injecting a `MockLanguageModelV3` through the
|
||||
* public `.model()` builder — the real `streamText`, loop, tool executor and
|
||||
* checkpoint machinery all run. (Module-level mocking is not available here:
|
||||
* the source lazy-loads the AI SDK via `require()`, which only the unit
|
||||
* config's require-rewrite plugin makes mockable.)
|
||||
*
|
||||
* Crash model: a graceful abort is NOT a crash — the stream session's shutdown
|
||||
* path deletes the run's checkpoint (cancel semantics). A crash is teardown
|
||||
* never running. So the test gates tool #2 on a deferred promise that is never
|
||||
* resolved: the loop freezes awaiting the tool batch, the abort signal is
|
||||
* never observed (abort checks sit at step boundaries), no shutdown code runs,
|
||||
* and the batch-1 checkpoint durably survives in the shared store — exactly
|
||||
* the state a dead process leaves behind. Everything is event-ordered; no
|
||||
* timers.
|
||||
*/
|
||||
import { convertArrayToReadableStream, MockLanguageModelV3 } from 'ai/test';
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { z } from 'zod';
|
||||
|
||||
import { chunksOfType, collectStreamChunks } from './helpers';
|
||||
import { Agent, Tool } from '../../index';
|
||||
import type { CheckpointStore, SerializableAgentState } from '../../types';
|
||||
|
||||
type MockStreamResult = Awaited<ReturnType<MockLanguageModelV3['doStream']>>;
|
||||
type MockStreamPart = MockStreamResult['stream'] extends ReadableStream<infer P> ? P : never;
|
||||
|
||||
const USAGE: Extract<MockStreamPart, { type: 'finish' }>['usage'] = {
|
||||
inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 },
|
||||
outputTokens: { total: 5, text: 5, reasoning: 0 },
|
||||
};
|
||||
|
||||
/** A scripted model turn that calls one tool. */
|
||||
function makeToolCallTurn(
|
||||
toolCallId: string,
|
||||
toolName: string,
|
||||
args: Record<string, unknown>,
|
||||
): MockStreamResult {
|
||||
const parts: MockStreamPart[] = [
|
||||
{ type: 'stream-start', warnings: [] },
|
||||
{ type: 'tool-call', toolCallId, toolName, input: JSON.stringify(args) },
|
||||
{ type: 'finish', finishReason: { unified: 'tool-calls', raw: 'tool_calls' }, usage: USAGE },
|
||||
];
|
||||
return { stream: convertArrayToReadableStream(parts) };
|
||||
}
|
||||
|
||||
/** A scripted model turn that finishes with plain text. */
|
||||
function makeTextTurn(text: string): MockStreamResult {
|
||||
const parts: MockStreamPart[] = [
|
||||
{ type: 'stream-start', warnings: [] },
|
||||
{ type: 'text-start', id: 'txt-1' },
|
||||
{ type: 'text-delta', id: 'txt-1', delta: text },
|
||||
{ type: 'text-end', id: 'txt-1' },
|
||||
{ type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage: USAGE },
|
||||
];
|
||||
return { stream: convertArrayToReadableStream(parts) };
|
||||
}
|
||||
|
||||
/**
|
||||
* A minimal CheckpointStore backed by a plain Map so it can be shared across
|
||||
* agent instances, standing in for durable external storage.
|
||||
*/
|
||||
class InMemoryCheckpointStore implements CheckpointStore {
|
||||
private store = new Map<string, SerializableAgentState>();
|
||||
|
||||
async save(key: string, state: SerializableAgentState): Promise<void> {
|
||||
this.store.set(key, structuredClone(state));
|
||||
}
|
||||
|
||||
async load(key: string): Promise<SerializableAgentState | undefined> {
|
||||
const state = this.store.get(key);
|
||||
return state ? structuredClone(state) : undefined;
|
||||
}
|
||||
|
||||
async delete(key: string): Promise<void> {
|
||||
this.store.delete(key);
|
||||
}
|
||||
|
||||
get size(): number {
|
||||
return this.store.size;
|
||||
}
|
||||
}
|
||||
|
||||
function createDeferred(): { promise: Promise<void>; resolve: () => void } {
|
||||
let resolve!: () => void;
|
||||
const promise = new Promise<void>((res) => {
|
||||
resolve = res;
|
||||
});
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
describe('step checkpoints + crash resume through the public Agent API', () => {
|
||||
it('resumes a crashed run from its step checkpoint under the original runId', async () => {
|
||||
const store = new InMemoryCheckpointStore();
|
||||
const counters: Record<string, { entered: number; completed: number }> = {
|
||||
tool_one: { entered: 0, completed: 0 },
|
||||
tool_two: { entered: 0, completed: 0 },
|
||||
tool_three: { entered: 0, completed: 0 },
|
||||
};
|
||||
|
||||
// One scripted model shared by both agent instances: turns 1-2 belong to
|
||||
// the original run, turns 3-5 to the crash-resumed one.
|
||||
const turns: MockStreamResult[] = [
|
||||
makeToolCallTurn('tc-1', 'tool_one', { value: 'first' }),
|
||||
makeToolCallTurn('tc-2', 'tool_two', { value: 'second' }),
|
||||
makeToolCallTurn('tc-2b', 'tool_two', { value: 'second-retry' }),
|
||||
makeToolCallTurn('tc-3', 'tool_three', { value: 'third' }),
|
||||
makeTextTurn('all three steps done'),
|
||||
];
|
||||
let nextTurn = 0;
|
||||
const model = new MockLanguageModelV3({
|
||||
provider: 'mock',
|
||||
modelId: 'scripted',
|
||||
doStream: async () => turns[nextTurn++],
|
||||
});
|
||||
|
||||
// Crash window controls: tool #2's first invocation signals entry and
|
||||
// then blocks on a deferred that is never resolved.
|
||||
const toolTwoEntered = createDeferred();
|
||||
const toolTwoGate = createDeferred();
|
||||
let gateArmed = true;
|
||||
|
||||
const makeRecordingTool = (name: string, step: number): Tool =>
|
||||
new Tool(name)
|
||||
.description(`Run step ${step} of the job`)
|
||||
.input(z.object({ value: z.string().optional() }))
|
||||
.handler(async () => {
|
||||
counters[name].entered++;
|
||||
if (name === 'tool_two' && gateArmed) {
|
||||
toolTwoEntered.resolve();
|
||||
await toolTwoGate.promise; // never resolves — the crash victim
|
||||
}
|
||||
counters[name].completed++;
|
||||
return { ok: true, step };
|
||||
});
|
||||
|
||||
const buildAgent = (): Agent =>
|
||||
new Agent('crash-resume-agent')
|
||||
.model(model)
|
||||
.instructions('Run the three steps in order using the tools. Be concise.')
|
||||
.tool(makeRecordingTool('tool_one', 1))
|
||||
.tool(makeRecordingTool('tool_two', 2))
|
||||
.tool(makeRecordingTool('tool_three', 3))
|
||||
.checkpoint(store);
|
||||
|
||||
// --- Half 1: run until the crash ---
|
||||
const agent1 = buildAgent();
|
||||
const controller = new AbortController();
|
||||
const first = await agent1.stream('Run steps one, two and three', {
|
||||
stepCheckpoints: true,
|
||||
abortSignal: controller.signal,
|
||||
});
|
||||
const originalRunId = first.runId;
|
||||
expect(originalRunId).toBeTruthy();
|
||||
|
||||
// Consume in the background so the loop isn't stalled by backpressure.
|
||||
// Never awaited: the run freezes inside tool #2 and this stream never ends.
|
||||
void collectStreamChunks(first.stream).catch(() => {});
|
||||
|
||||
await toolTwoEntered.promise;
|
||||
// Tool #2 is in flight, so the batch-1 step checkpoint is durably stored.
|
||||
// The abort is the kill signal; the frozen loop never observes it, so no
|
||||
// teardown (which would delete the checkpoint as a cancellation) runs.
|
||||
controller.abort();
|
||||
|
||||
const checkpointAtCrash = await store.load(originalRunId);
|
||||
expect(checkpointAtCrash?.status).toBe('running');
|
||||
expect(checkpointAtCrash?.pendingToolCalls).toEqual({});
|
||||
expect(JSON.stringify(checkpointAtCrash?.messageList)).toContain('tc-1');
|
||||
expect(counters.tool_one).toEqual({ entered: 1, completed: 1 });
|
||||
expect(model.doStreamCalls).toHaveLength(2);
|
||||
|
||||
// --- Half 2: fresh agent, same store, crash-resume ---
|
||||
gateArmed = false;
|
||||
const contextNote =
|
||||
'Your previous "tool_two" call was interrupted by a restart. Its effect is unverified — verify before retrying.';
|
||||
const agent2 = buildAgent();
|
||||
const resumed = await agent2.crashResume({
|
||||
runId: originalRunId,
|
||||
contextNotes: [contextNote],
|
||||
stepCheckpoints: true,
|
||||
});
|
||||
|
||||
// The run continues under the ORIGINAL runId.
|
||||
expect(resumed.runId).toBe(originalRunId);
|
||||
|
||||
const chunks = await collectStreamChunks(resumed.stream);
|
||||
expect(chunks.filter((c) => c.type === 'error')).toHaveLength(0);
|
||||
const finish = chunksOfType(chunks, 'finish').at(-1);
|
||||
expect(finish?.finishReason).toBe('stop');
|
||||
|
||||
// Every tool completed exactly once across both halves. The step-1 tool
|
||||
// was not re-run on resume (its result came from the checkpoint); tool #2's
|
||||
// first entry is the crash victim, re-issued by the scripted model as a
|
||||
// fresh call after reading the context note.
|
||||
expect(counters.tool_one).toEqual({ entered: 1, completed: 1 });
|
||||
expect(counters.tool_two).toEqual({ entered: 2, completed: 1 });
|
||||
expect(counters.tool_three).toEqual({ entered: 1, completed: 1 });
|
||||
|
||||
// The first resumed model call saw the checkpointed history AND the note.
|
||||
expect(model.doStreamCalls).toHaveLength(5);
|
||||
const resumedPrompt = JSON.stringify(model.doStreamCalls[2].prompt);
|
||||
expect(resumedPrompt).toContain('tc-1');
|
||||
expect(resumedPrompt).toContain('interrupted by a restart');
|
||||
|
||||
// Completion deleted the checkpoint...
|
||||
expect(await store.load(originalRunId)).toBeUndefined();
|
||||
expect(store.size).toBe(0);
|
||||
|
||||
// ...so a second crash-resume attempt finds nothing to resume.
|
||||
const agent3 = buildAgent();
|
||||
await expect(agent3.crashResume({ runId: originalRunId })).rejects.toThrow(
|
||||
/No checkpoint found/,
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,349 @@
|
||||
/**
|
||||
* Durable-log RFC (resilience phase) harness: per-step checkpoints and
|
||||
* crash-resume at the SDK level, with a mocked model. Proves:
|
||||
* - stepCheckpoints persists a `running` checkpoint at every step boundary,
|
||||
* - a fresh runtime resumes from that checkpoint via crashResume() and
|
||||
* continues to completion under the SAME runId,
|
||||
* - contextNotes (interrupted-tool and correction notes recovered by the
|
||||
* host from its durable log) reach the resumed model context,
|
||||
* - crashResume rejects suspended checkpoints (those belong to resume()).
|
||||
*/
|
||||
import * as aiModule from 'ai';
|
||||
import type { Mock } from 'vitest';
|
||||
import { z } from 'zod';
|
||||
|
||||
import type { CheckpointStore, SerializableAgentState } from '../../types';
|
||||
import type { StreamChunk } from '../../types/sdk/agent';
|
||||
import type { BuiltTool, InterruptibleToolContext } from '../../types/sdk/tool';
|
||||
import { AgentRuntime } from '../loop/agent-runtime';
|
||||
import { AgentEventBus } from '../state/event-bus';
|
||||
|
||||
// Mock provider packages so createModel() doesn't fail when no API key is set
|
||||
vi.mock('@ai-sdk/openai', () => ({
|
||||
createOpenAI: () =>
|
||||
Object.assign(() => ({ provider: 'openai', modelId: 'mock', specificationVersion: 'v3' }), {
|
||||
embeddingModel: () => ({ provider: 'openai', modelId: 'mock', specificationVersion: 'v2' }),
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock('@ai-sdk/anthropic', () => ({
|
||||
createAnthropic: () => () => ({
|
||||
provider: 'anthropic',
|
||||
modelId: 'mock',
|
||||
specificationVersion: 'v3',
|
||||
}),
|
||||
}));
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/consistent-type-imports
|
||||
type AiImport = typeof import('ai');
|
||||
|
||||
vi.mock('ai', async () => {
|
||||
const actual = await vi.importActual<AiImport>('ai');
|
||||
return {
|
||||
...actual,
|
||||
embed: vi.fn(),
|
||||
embedMany: vi.fn(),
|
||||
generateText: vi.fn(),
|
||||
streamText: vi.fn(),
|
||||
tool: vi.fn((config: unknown) => config),
|
||||
jsonSchema: vi.fn((schema: unknown) => ({ _type: 'jsonSchema', schema })),
|
||||
Output: {
|
||||
object: vi.fn(({ schema }: { schema: unknown }) => ({ _type: 'object', schema })),
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
const { streamText } = aiModule as unknown as { streamText: Mock };
|
||||
|
||||
function* makeChunkStream(
|
||||
chunks: Array<Record<string, unknown>>,
|
||||
): Generator<Record<string, unknown>> {
|
||||
for (const c of chunks) {
|
||||
yield c;
|
||||
}
|
||||
}
|
||||
|
||||
function makeStreamSuccess(text = 'Hello') {
|
||||
return {
|
||||
fullStream: makeChunkStream([{ type: 'text-delta', textDelta: text }]),
|
||||
finishReason: Promise.resolve('stop'),
|
||||
usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
|
||||
response: Promise.resolve({
|
||||
messages: [{ role: 'assistant', content: [{ type: 'text', text }] }],
|
||||
}),
|
||||
toolCalls: Promise.resolve([]),
|
||||
};
|
||||
}
|
||||
|
||||
function makeStreamWithToolCall(toolCallId: string, args: Record<string, unknown>) {
|
||||
return {
|
||||
fullStream: makeChunkStream([{ type: 'text-delta', textDelta: 'working...' }]),
|
||||
finishReason: Promise.resolve('tool-calls'),
|
||||
usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
|
||||
response: Promise.resolve({
|
||||
messages: [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [{ type: 'tool-call', toolCallId, toolName: 'lookup', args }],
|
||||
},
|
||||
],
|
||||
}),
|
||||
toolCalls: Promise.resolve([{ toolCallId, toolName: 'lookup', input: args }]),
|
||||
};
|
||||
}
|
||||
|
||||
async function collectChunks(stream: ReadableStream<unknown>): Promise<StreamChunk[]> {
|
||||
const chunks: StreamChunk[] = [];
|
||||
const reader = stream.getReader();
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
chunks.push(value as StreamChunk);
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
|
||||
/** In-memory CheckpointStore that also records every save for inspection. */
|
||||
class RecordingCheckpointStore implements CheckpointStore {
|
||||
map = new Map<string, SerializableAgentState>();
|
||||
|
||||
saves: Array<{ key: string; state: SerializableAgentState }> = [];
|
||||
|
||||
deletes: string[] = [];
|
||||
|
||||
async save(key: string, state: SerializableAgentState): Promise<void> {
|
||||
this.saves.push({ key, state: structuredClone(state) });
|
||||
this.map.set(key, state);
|
||||
await Promise.resolve();
|
||||
}
|
||||
|
||||
async load(key: string): Promise<SerializableAgentState | undefined> {
|
||||
return await Promise.resolve(this.map.get(key));
|
||||
}
|
||||
|
||||
async delete(key: string): Promise<void> {
|
||||
this.deletes.push(key);
|
||||
this.map.delete(key);
|
||||
await Promise.resolve();
|
||||
}
|
||||
}
|
||||
|
||||
/** Interruptible tool: suspends on first call (HITL), returns on resume. */
|
||||
const approveTool: BuiltTool = {
|
||||
name: 'approve',
|
||||
description: 'Requires approval',
|
||||
inputSchema: z.object({ question: z.string().optional() }),
|
||||
suspendSchema: z.object({ question: z.string() }),
|
||||
resumeSchema: z.object({ approved: z.boolean() }),
|
||||
handler: async (_input: unknown, ctx: unknown) => {
|
||||
const { suspend, resumeData } = ctx as InterruptibleToolContext;
|
||||
if (!resumeData) {
|
||||
return await suspend({ question: 'approve?' });
|
||||
}
|
||||
return { approved: true };
|
||||
},
|
||||
};
|
||||
|
||||
const lookupTool: BuiltTool = {
|
||||
name: 'lookup',
|
||||
description: 'Mock lookup tool',
|
||||
inputSchema: z.object({ value: z.string().optional() }),
|
||||
handler: async (input) =>
|
||||
await Promise.resolve({ found: (input as { value?: string }).value ?? 'nothing' }),
|
||||
};
|
||||
|
||||
function createRuntime(store: CheckpointStore, tools: BuiltTool[] = [lookupTool]) {
|
||||
const bus = new AgentEventBus();
|
||||
const runtime = new AgentRuntime({
|
||||
name: 'crash-test',
|
||||
model: 'openai/gpt-4o-mini',
|
||||
instructions: 'You are a test assistant.',
|
||||
tools,
|
||||
eventBus: bus,
|
||||
checkpointStorage: store,
|
||||
});
|
||||
return runtime;
|
||||
}
|
||||
|
||||
describe('step checkpoints + crash resume (durable-log RFC)', () => {
|
||||
beforeEach(() => {
|
||||
streamText.mockReset();
|
||||
});
|
||||
|
||||
it('persists a running checkpoint at every step boundary and resumes after a crash', async () => {
|
||||
const store = new RecordingCheckpointStore();
|
||||
const runtime = createRuntime(store);
|
||||
|
||||
streamText
|
||||
.mockReturnValueOnce(makeStreamWithToolCall('tc-1', { value: 'first' }))
|
||||
.mockReturnValueOnce(makeStreamWithToolCall('tc-2', { value: 'second' }))
|
||||
.mockReturnValueOnce(makeStreamSuccess('all done'));
|
||||
|
||||
const result = await runtime.stream('find things', { stepCheckpoints: true });
|
||||
const chunks = await collectChunks(result.stream);
|
||||
const finish = chunks.filter((c) => c.type === 'finish').at(-1) as
|
||||
| (StreamChunk & { type: 'finish'; finishReason: string })
|
||||
| undefined;
|
||||
expect(finish?.finishReason).toBe('stop');
|
||||
|
||||
// One checkpoint per completed tool step, all with status 'running' and
|
||||
// no pending tool calls (step boundary), under the run's id.
|
||||
const stepSaves = store.saves.filter((s) => s.state.status === 'running');
|
||||
expect(stepSaves).toHaveLength(2);
|
||||
for (const save of stepSaves) {
|
||||
expect(save.key).toBe(result.runId);
|
||||
expect(save.state.pendingToolCalls).toEqual({});
|
||||
expect(save.state.messageList).toBeDefined();
|
||||
}
|
||||
// The second checkpoint contains both settled tool calls.
|
||||
const secondJson = JSON.stringify(stepSaves[1].state.messageList);
|
||||
expect(secondJson).toContain('tc-1');
|
||||
expect(secondJson).toContain('tc-2');
|
||||
// The completed run deleted its checkpoint (no leak).
|
||||
expect(store.deletes).toContain(result.runId);
|
||||
|
||||
// CRASH SIMULATION: the process died after step 2's checkpoint and
|
||||
// before the run finished, so the store still holds that checkpoint.
|
||||
store.map.set(result.runId, stepSaves[1].state);
|
||||
|
||||
const runtime2 = createRuntime(store);
|
||||
streamText.mockReset();
|
||||
streamText.mockReturnValueOnce(makeStreamSuccess('resumed and done'));
|
||||
|
||||
const resumed = await runtime2.crashResume({
|
||||
runId: result.runId,
|
||||
contextNotes: [
|
||||
'Your previous "lookup" tool call was interrupted by a restart. Verify before retrying.',
|
||||
],
|
||||
});
|
||||
expect(resumed.runId).toBe(result.runId);
|
||||
const resumedChunks = await collectChunks(resumed.stream);
|
||||
const resumedFinish = resumedChunks.filter((c) => c.type === 'finish').at(-1) as
|
||||
| (StreamChunk & { type: 'finish'; finishReason: string })
|
||||
| undefined;
|
||||
expect(resumedFinish?.finishReason).toBe('stop');
|
||||
|
||||
// The resumed model call saw the checkpointed history AND the note.
|
||||
expect(streamText).toHaveBeenCalledTimes(1);
|
||||
const callArgs = streamText.mock.calls[0][0] as { messages: unknown };
|
||||
const contextJson = JSON.stringify(callArgs.messages);
|
||||
expect(contextJson).toContain('tc-1');
|
||||
expect(contextJson).toContain('tc-2');
|
||||
expect(contextJson).toContain('interrupted by a restart');
|
||||
});
|
||||
|
||||
it('a crash-resumed run that suspends at HITL persists a resumable suspended checkpoint', async () => {
|
||||
const store = new RecordingCheckpointStore();
|
||||
|
||||
// Original run: one completed tool step (writes the step checkpoint),
|
||||
// then the crash.
|
||||
const runtime = createRuntime(store, [lookupTool, approveTool]);
|
||||
streamText
|
||||
.mockReturnValueOnce(makeStreamWithToolCall('tc-1', { value: 'first' }))
|
||||
.mockReturnValueOnce(makeStreamSuccess('done'));
|
||||
const result = await runtime.stream('find things', { stepCheckpoints: true });
|
||||
await collectChunks(result.stream);
|
||||
const stepCheckpoint = store.saves.find((s) => s.state.status === 'running');
|
||||
expect(stepCheckpoint).toBeDefined();
|
||||
store.map.set(result.runId, stepCheckpoint!.state); // crash before completion
|
||||
|
||||
// Crash-resume: the next model turn calls the HITL tool, which suspends.
|
||||
const runtime2 = createRuntime(store, [lookupTool, approveTool]);
|
||||
streamText.mockReset();
|
||||
streamText.mockReturnValueOnce({
|
||||
fullStream: makeChunkStream([{ type: 'text-delta', textDelta: 'asking...' }]),
|
||||
finishReason: Promise.resolve('tool-calls'),
|
||||
usage: Promise.resolve({ inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
|
||||
response: Promise.resolve({
|
||||
messages: [
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{
|
||||
type: 'tool-call',
|
||||
toolCallId: 'tc-hitl',
|
||||
toolName: 'approve',
|
||||
args: { question: 'ok?' },
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
}),
|
||||
toolCalls: Promise.resolve([
|
||||
{ toolCallId: 'tc-hitl', toolName: 'approve', input: { question: 'ok?' } },
|
||||
]),
|
||||
});
|
||||
|
||||
const resumed = await runtime2.crashResume({ runId: result.runId, stepCheckpoints: true });
|
||||
const chunks = await collectChunks(resumed.stream);
|
||||
expect(chunks.some((c) => c.type === 'tool-call-suspended')).toBe(true);
|
||||
|
||||
// The suspension checkpoint is persisted under the same runId with the
|
||||
// pending HITL call, so the normal resume() path can pick it up.
|
||||
const suspended = store.map.get(result.runId);
|
||||
expect(suspended?.status).toBe('suspended');
|
||||
expect(Object.keys(suspended?.pendingToolCalls ?? {})).toEqual(['tc-hitl']);
|
||||
});
|
||||
|
||||
it('rejects crashResume on a suspended checkpoint (resume() owns those)', async () => {
|
||||
const store = new RecordingCheckpointStore();
|
||||
store.map.set('run_suspended', {
|
||||
status: 'suspended',
|
||||
messageList: { messages: [] } as never,
|
||||
pendingToolCalls: {},
|
||||
});
|
||||
const runtime = createRuntime(store);
|
||||
await expect(runtime.crashResume({ runId: 'run_suspended' })).rejects.toThrow(
|
||||
/crashResume only accepts step checkpoints/,
|
||||
);
|
||||
});
|
||||
|
||||
it('throws when no checkpoint exists for the runId', async () => {
|
||||
const runtime = createRuntime(new RecordingCheckpointStore());
|
||||
await expect(runtime.crashResume({ runId: 'run_missing' })).rejects.toThrow(
|
||||
/No checkpoint found/,
|
||||
);
|
||||
});
|
||||
|
||||
it('rejects a running checkpoint that still has pending tool calls (claimed HITL resume)', async () => {
|
||||
const store = new RecordingCheckpointStore();
|
||||
// The state claimResume() leaves behind if the process dies mid-resume:
|
||||
// status flipped to 'running', pending HITL calls not yet settled.
|
||||
store.map.set('run_claimed', {
|
||||
status: 'running',
|
||||
messageList: { messages: [] } as never,
|
||||
pendingToolCalls: {
|
||||
'tc-hitl': {
|
||||
suspended: true,
|
||||
toolCallId: 'tc-hitl',
|
||||
toolName: 'approve',
|
||||
input: {},
|
||||
suspendPayload: { question: 'approve?' },
|
||||
resumeSchema: {},
|
||||
runId: 'run_claimed',
|
||||
},
|
||||
},
|
||||
});
|
||||
const runtime = createRuntime(store);
|
||||
await expect(runtime.crashResume({ runId: 'run_claimed' })).rejects.toThrow(
|
||||
/pending tool calls/,
|
||||
);
|
||||
});
|
||||
|
||||
it('surfaces an error when maxIterations is decreased on crashResume (parity with resume)', async () => {
|
||||
const store = new RecordingCheckpointStore();
|
||||
store.map.set('run_max', {
|
||||
status: 'running',
|
||||
messageList: { messages: [], historyIds: [], inputIds: [], responseIds: [] },
|
||||
pendingToolCalls: {},
|
||||
executionOptions: { maxIterations: 5 },
|
||||
});
|
||||
const runtime = createRuntime(store);
|
||||
const result = await runtime.crashResume({ runId: 'run_max', maxIterations: 2 });
|
||||
const chunks = await collectChunks(result.stream);
|
||||
const errorChunk = chunks.find((c) => c.type === 'error') as
|
||||
| (StreamChunk & { type: 'error'; error: unknown })
|
||||
| undefined;
|
||||
expect(String(errorChunk?.error)).toContain('Cannot decrease maxIterations');
|
||||
});
|
||||
});
|
||||
@@ -453,6 +453,92 @@ export class AgentRuntime {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Durable-log RFC (resilience phase): re-drive a run from a `running`-status
|
||||
* step checkpoint after a process crash. Unlike resume(), there is no
|
||||
* pending tool call to settle — the checkpoint was written at a step
|
||||
* boundary — so the loop re-enters directly at the next model call.
|
||||
* `contextNotes` are appended as user messages before the model call: the
|
||||
* host uses them to surface interrupted tool calls ("effect unverified —
|
||||
* verify before retrying") and undrained steering corrections recovered
|
||||
* from its durable event log. Tool calls are never re-executed mechanically.
|
||||
*/
|
||||
async crashResume(
|
||||
options: { runId: string; contextNotes?: string[] } & ExecutionOptions,
|
||||
): Promise<StreamResult> {
|
||||
this.runId = options.runId;
|
||||
const state = await this.runState.loadForCrashResume(this.runId);
|
||||
if (!state) throw new Error(`No checkpoint found for runId: ${this.runId}`);
|
||||
if (state.status !== 'running') {
|
||||
throw new Error(
|
||||
`Checkpoint for runId ${this.runId} has status '${state.status}' — crashResume only accepts step checkpoints; use resume() for suspended runs`,
|
||||
);
|
||||
}
|
||||
// A claimed HITL resume also persists as 'running' but still carries its
|
||||
// pending tool calls; re-driving it would skip settling them. Step
|
||||
// checkpoints are always written with empty pendingToolCalls.
|
||||
if (Object.keys(state.pendingToolCalls).length > 0) {
|
||||
throw new Error(
|
||||
`Checkpoint for runId ${this.runId} has pending tool calls — crashResume only accepts step checkpoints`,
|
||||
);
|
||||
}
|
||||
|
||||
const list = AgentMessageList.deserialize(state.messageList);
|
||||
this.context.hydrateDeferredToolsFromList(list);
|
||||
|
||||
let abortScope: AgentAbortScope | undefined;
|
||||
try {
|
||||
const { runId: _rid, contextNotes, ...callerExecOptions } = options;
|
||||
const persisted = state.executionOptions ?? {};
|
||||
const persistedMaxIterations = persisted.maxIterations;
|
||||
const callerMaxIterations = callerExecOptions.maxIterations;
|
||||
if (
|
||||
callerMaxIterations !== undefined &&
|
||||
persistedMaxIterations !== undefined &&
|
||||
callerMaxIterations < persistedMaxIterations
|
||||
) {
|
||||
throw new Error(
|
||||
`Cannot decrease maxIterations when resuming a run. Expected >= ${persistedMaxIterations}, received ${callerMaxIterations}.`,
|
||||
);
|
||||
}
|
||||
const mergedMaxIterations = callerMaxIterations ?? persistedMaxIterations;
|
||||
const resumeOptions: RuntimeExecutionOptions = {
|
||||
persistence: state.persistence,
|
||||
...callerExecOptions,
|
||||
...(mergedMaxIterations !== undefined ? { maxIterations: mergedMaxIterations } : {}),
|
||||
...(state.iterationCount !== undefined ? { iterationCount: state.iterationCount } : {}),
|
||||
};
|
||||
|
||||
for (const note of contextNotes ?? []) {
|
||||
list.addInput([{ role: 'user', content: [{ type: 'text', text: note }] }]);
|
||||
}
|
||||
|
||||
abortScope = this.eventBus.createAbortScope(resumeOptions.abortSignal);
|
||||
const activeAbortScope = abortScope;
|
||||
|
||||
await this.ensureModelCost();
|
||||
await this.memory.setListObservationLogMemory(list, state.persistence);
|
||||
|
||||
return {
|
||||
runId: this.runId,
|
||||
stream: this.startStream({
|
||||
list,
|
||||
options: resumeOptions,
|
||||
abortScope: activeAbortScope,
|
||||
}),
|
||||
getState: () => this.getState(),
|
||||
};
|
||||
} catch (error) {
|
||||
const isAbort = abortScope?.isAborted ?? false;
|
||||
abortScope?.dispose();
|
||||
this.updateState({ status: isAbort ? 'cancelled' : 'failed' });
|
||||
if (!isAbort) {
|
||||
this.eventBus.emit({ type: AgentEvent.Error, message: String(error), error });
|
||||
}
|
||||
return { runId: this.runId, stream: makeErrorStream(error), getState: () => this.getState() };
|
||||
}
|
||||
}
|
||||
|
||||
// --- Private ---
|
||||
|
||||
/**
|
||||
@@ -757,6 +843,18 @@ export class AgentRuntime {
|
||||
|
||||
// Emit TurnEnd after all tool calls in this iteration are processed
|
||||
this.emitTurnEnd(turn.newMessages, extractSettledToolCalls(list.responseDelta()));
|
||||
|
||||
// Step boundary reached with nothing pending: durably checkpoint so a
|
||||
// crash before the next model call loses only the in-flight step.
|
||||
if (options?.stepCheckpoints) {
|
||||
await this.persistStepCheckpoint(
|
||||
list,
|
||||
totalUsage,
|
||||
options,
|
||||
maxIterations,
|
||||
iterationCount + 1,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if (!reachedStopCondition && iterationCount >= maxIterations) {
|
||||
@@ -843,6 +941,38 @@ export class AgentRuntime {
|
||||
return this.runId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Durable-log RFC (resilience phase): per-step checkpoint — the completion
|
||||
* of the "step boundary = durability boundary" rule. Called at the end of
|
||||
* each loop iteration (after tool results are appended to `list`, before
|
||||
* the next model call), gated on the `stepCheckpoints` opt-in so the write
|
||||
* cost is only paid where crash-resume matters. Reuses the suspension state
|
||||
* shape; pendingToolCalls is empty at a step boundary.
|
||||
*/
|
||||
private async persistStepCheckpoint(
|
||||
list: AgentMessageList,
|
||||
totalUsage: TokenUsage | undefined,
|
||||
options: RuntimeExecutionOptions | undefined,
|
||||
maxIterations?: number,
|
||||
iterationCount?: number,
|
||||
): Promise<void> {
|
||||
const resolvedMaxIterations = maxIterations ?? options?.maxIterations;
|
||||
const resolvedIterationCount = iterationCount ?? options?.iterationCount;
|
||||
const executionOptions: PersistedExecutionOptions | undefined =
|
||||
resolvedMaxIterations !== undefined ? { maxIterations: resolvedMaxIterations } : undefined;
|
||||
|
||||
const state: SerializableAgentState = {
|
||||
persistence: options?.persistence,
|
||||
status: 'running',
|
||||
messageList: list.serialize(),
|
||||
pendingToolCalls: {},
|
||||
usage: totalUsage,
|
||||
executionOptions,
|
||||
...(resolvedIterationCount !== undefined ? { iterationCount: resolvedIterationCount } : {}),
|
||||
};
|
||||
await this.runState.checkpointStep(this.runId, state);
|
||||
}
|
||||
|
||||
/** Clean up stored state for a run when it finishes without re-suspending. */
|
||||
private async cleanupRun(): Promise<void> {
|
||||
await this.runState.complete(this.runId);
|
||||
|
||||
@@ -48,6 +48,26 @@ export class RunStateManager {
|
||||
await this.store.save(runId, { ...state, status: 'suspended' });
|
||||
}
|
||||
|
||||
/**
|
||||
* Durable-log RFC (resilience phase): per-step checkpoint. Same
|
||||
* serialization and store as suspend(), but status stays 'running' — one
|
||||
* upserted row per run, overwritten at every step boundary, so a crash
|
||||
* loses only the in-flight step. Deleted by complete() like any checkpoint.
|
||||
*/
|
||||
async checkpointStep(runId: string, state: SerializableAgentState): Promise<void> {
|
||||
await this.store.save(runId, { ...state, status: 'running' });
|
||||
}
|
||||
|
||||
/**
|
||||
* Durable-log RFC (resilience phase): load a checkpoint after a crash,
|
||||
* where status is 'running' (mid-step death) rather than 'suspended'.
|
||||
* Callers pair this with the interrupted-run sweeper: in-flight tool calls
|
||||
* are resolved as `tool-interrupted` facts, never re-executed.
|
||||
*/
|
||||
async loadForCrashResume(runId: string): Promise<SerializableAgentState | undefined> {
|
||||
return await this.store.load(runId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Load a suspended run state for resumption. This is read-only so callers can
|
||||
* validate the resume request before claiming the checkpoint.
|
||||
|
||||
@@ -730,6 +730,26 @@ export class Agent implements BuiltAgent, AgentBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Durable-log RFC (resilience phase): re-drive a run from a `running`-status
|
||||
* step checkpoint after a process crash. There is no pending tool call to
|
||||
* settle — the loop re-enters at the next model call. See
|
||||
* AgentRuntime.crashResume for `contextNotes` semantics.
|
||||
*/
|
||||
async crashResume(
|
||||
options: { runId: string; contextNotes?: string[] } & ExecutionOptions,
|
||||
): Promise<StreamResult> {
|
||||
const config = await this.ensureBuilt();
|
||||
const active = this.createRuntime(config, options.runId);
|
||||
try {
|
||||
const result = await active.runtime.crashResume(options);
|
||||
return { ...result, stream: this.trackStreamRuntime(result.stream, active) };
|
||||
} catch (error) {
|
||||
await this.cleanupRuntime(active);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
approve(method: 'generate', options: ResumeOptions & ExecutionOptions): Promise<GenerateResult>;
|
||||
approve(method: 'stream', options: ResumeOptions & ExecutionOptions): Promise<StreamResult>;
|
||||
async approve(
|
||||
|
||||
@@ -176,6 +176,13 @@ export interface ExecutionOptions {
|
||||
executionCounter?: AgentExecutionCounter;
|
||||
onStepStart?: (event: OnStepStartEvent) => void | Promise<void>;
|
||||
onStepFinish?: (event: OnStepFinishEvent) => void | Promise<void>;
|
||||
/**
|
||||
* Durable-log RFC (resilience phase), opt-in: persist a `running`-status
|
||||
* checkpoint at every step boundary (after a tool batch settles, before the
|
||||
* next model call) so a crash loses only the in-flight step. Requires a
|
||||
* persistence-backed CheckpointStore; recover via `crashResume()`.
|
||||
*/
|
||||
stepCheckpoints?: boolean;
|
||||
}
|
||||
|
||||
export interface PersistedExecutionOptions {
|
||||
|
||||
Reference in New Issue
Block a user