diff --git a/sdk/src/__tests__/byok-run.test.ts b/sdk/src/__tests__/byok-run.test.ts index 8dbadc1264..3ba668e3b5 100644 --- a/sdk/src/__tests__/byok-run.test.ts +++ b/sdk/src/__tests__/byok-run.test.ts @@ -2,7 +2,7 @@ import { mkdtemp, readFile, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import path from 'node:path' -import { afterEach, describe, expect, test } from 'bun:test' +import { afterEach, describe, expect, spyOn, test } from 'bun:test' import { COMPACTION_SUMMARY_PROMPT_OPENING } from '@codebuff/common/constants/compaction-summary' import { CodebuffClient } from '../client' @@ -385,6 +385,47 @@ describe('direct BYOK SDK runs', () => { expect(requests.every((url) => url === 'http://127.0.0.1:9876/v1/chat/completions')).toBe(true) }, 30_000) + test('schedules checkpoints at the host-scaled interval', async () => { + globalThis.fetch = (async (_input, init) => { + const body = JSON.parse(String(init?.body)) + return sse({ + id: 'work', object: 'chat.completion.chunk', created: 1, model: body.model, + choices: [{ index: 0, delta: { content: 'Done.' }, finish_reason: 'stop' }], + }) + }) as typeof fetch + + const delays: number[] = [] + const realSetInterval = globalThis.setInterval + const spy = spyOn(globalThis, 'setInterval').mockImplementation((( + fn: unknown, + delay?: number, + ...rest: unknown[] + ) => { + delays.push(Number(delay)) + return realSetInterval(fn as TimerHandler, delay, ...(rest as never[])) + }) as typeof setInterval) + try { + // A large persisted chat asks for a longer interval than the 5s default. + const scaled = await new CodebuffClient({ agentDefinitions: [agent], byok: connection() }) + .run({ + agent: agent.id, + prompt: 'Scaled.', + stateSnapshotIntervalMs: 12_345, + onStateSnapshot: () => {}, + }) + expect(scaled.output.type).not.toBe('error') + + const defaults = await new CodebuffClient({ agentDefinitions: [agent], byok: connection() }) + .run({ agent: agent.id, prompt: 'Default.', onStateSnapshot: () => {} }) + expect(defaults.output.type).not.toBe('error') + } finally { + spy.mockRestore() + } + + expect(delays).toContain(12_345) + expect(delays).toContain(5_000) + }) + test('returns a local cancellation result without provider or hosted fallback requests', async () => { const controller = new AbortController() controller.abort(new Error('Cancelled by user')) diff --git a/sdk/src/__tests__/cancelled-session-state.test.ts b/sdk/src/__tests__/cancelled-session-state.test.ts index 9874cab23b..b618628aeb 100644 --- a/sdk/src/__tests__/cancelled-session-state.test.ts +++ b/sdk/src/__tests__/cancelled-session-state.test.ts @@ -110,6 +110,30 @@ describe('buildCancelledSessionState', () => { ) }) + test('a checkpoint snapshot shares its history entries with the live run', () => { + // What a host persisting deltas needs: the snapshot's prefix IS the live + // history's prefix, by identity, so "what changed since my last save" is a + // pointer walk rather than a serialization of the whole transcript. + const history = [text('user', 'hello'), text('assistant', 'hi')] + const live = sessionStateWith(history) + + const snapshot = buildCancelledSessionState({ + sessionState: live, + runtimeMadeProgress: true, + message: 'The session ended before this response completed.', + shareMessageEntries: true, + }) + + expect(snapshot.mainAgentState.messageHistory[0]).toBe(history[0]) + expect(snapshot.mainAgentState.messageHistory[1]).toBe(history[1]) + // ...while the appended note exists only in the snapshot. + expect(snapshot.mainAgentState.messageHistory).toHaveLength(3) + expect(live.mainAgentState.messageHistory).toHaveLength(2) + expect(live.mainAgentState.contextTokenCount).toBe( + countTokensMessages(live.mainAgentState.messageHistory) + SYSTEM_AND_TOOLS, + ) + }) + test('does not disturb the live session', () => { // Same contract cloneSessionState exists for: the run is still unwinding // and the caller keeps using the original. diff --git a/sdk/src/__tests__/snapshot-session-state.test.ts b/sdk/src/__tests__/snapshot-session-state.test.ts new file mode 100644 index 0000000000..56a3fc4c5b --- /dev/null +++ b/sdk/src/__tests__/snapshot-session-state.test.ts @@ -0,0 +1,120 @@ +import { getInitialSessionState } from '@codebuff/common/types/session-state' +import { getStubProjectFileContext } from '@codebuff/common/util/file' +import { describe, expect, it } from 'bun:test' + +import { snapshotSessionState } from '../run' + +import type { SessionState } from '@codebuff/common/types/session-state' + +const text = (role: 'user' | 'assistant', body: string) => + ({ role, content: [{ type: 'text', text: body }] }) as any + +function makeSession(): SessionState { + const state = getInitialSessionState(getStubProjectFileContext()) + state.mainAgentState.messageHistory = [ + text('user', 'hello'), + text('assistant', 'world'), + ] + state.mainAgentState.agentContext = { + goal: { objective: 'ship it', logs: ['started'] }, + } + return state +} + +describe('snapshotSessionState', () => { + it('copies the containers a run mutates, so a snapshot stays put', () => { + const live = makeSession() + const snapshot = snapshotSessionState(live) + + // The mutations the run actually makes: push a message, replace the array, + // append a subgoal log, record a spawned child. + snapshot.mainAgentState.messageHistory.push(text('user', 'interrupted')) + live.mainAgentState.messageHistory.push(text('assistant', 'more')) + ;(snapshot.mainAgentState.agentContext.goal.logs as string[]).push('note') + live.mainAgentState.childRunIds.push('child-run') + + expect(snapshot.mainAgentState.messageHistory).toHaveLength(3) + expect(live.mainAgentState.messageHistory).toHaveLength(3) + expect(live.mainAgentState.agentContext.goal.logs).toEqual(['started']) + expect(snapshot.mainAgentState).not.toBe(live.mainAgentState) + expect(snapshot.mainAgentState.messageHistory).not.toBe( + live.mainAgentState.messageHistory, + ) + // A spawn after the snapshot must not appear in the snapshot the host is + // still about to serialize. + expect(snapshot.mainAgentState.childRunIds).toEqual([]) + }) + + it('shares message objects, which is what makes the delta cheap to find', () => { + const live = makeSession() + const first = snapshotSessionState(live) + // The runtime appends a new message and replaces the array, as it does at a + // step boundary. + live.mainAgentState.messageHistory = [ + ...live.mainAgentState.messageHistory, + text('assistant', 'step'), + ] + const second = snapshotSessionState(live) + + // A host can walk the two histories and find the append by identity, + // instead of serializing both to compare them. + expect(second.mainAgentState.messageHistory[0]).toBe( + first.mainAgentState.messageHistory[0], + ) + expect(second.mainAgentState.messageHistory[1]).toBe( + first.mainAgentState.messageHistory[1], + ) + expect( + second.mainAgentState.messageHistory[first.mainAgentState.messageHistory.length], + ).not.toBe(undefined) + }) + + it('copies subagent containers at every level', () => { + const live = makeSession() + const child = { + ...getInitialSessionState(getStubProjectFileContext()).mainAgentState, + agentId: 'child', + messageHistory: [text('assistant', 'sub work')], + } + live.mainAgentState.subagents = [child] + + const snapshot = snapshotSessionState(live) + + expect(snapshot.mainAgentState.subagents[0]).not.toBe(child) + expect(snapshot.mainAgentState.subagents[0].messageHistory).not.toBe( + child.messageHistory, + ) + // Entries are still shared; only the containers are copied. + expect(snapshot.mainAgentState.subagents[0].messageHistory[0]).toBe( + child.messageHistory[0], + ) + snapshot.mainAgentState.subagents[0].messageHistory.push(text('user', 'x')) + expect(child.messageHistory).toHaveLength(1) + }) + + it('shares fileContext, and does not fail on values JSON cannot clone', () => { + const live = makeSession() + ;(live.fileContext as any).customToolDefinitions = { + mcpTool: { inputSchema: { parse: () => ({}), _def: {} } }, + } + + const snapshot = snapshotSessionState(live) + + expect(snapshot.fileContext).toBe(live.fileContext) + expect( + (snapshot.fileContext as any).customToolDefinitions.mcpTool, + ).toBe((live.fileContext as any).customToolDefinitions.mcpTool) + }) + + it('is a structural copy of the live state, not a stale one', () => { + const live = makeSession() + live.mainAgentState.contextTokenCount = 123 + + const snapshot = snapshotSessionState(live) + + expect(snapshot.mainAgentState.contextTokenCount).toBe(123) + expect(snapshot.mainAgentState.stepsRemaining).toBe( + live.mainAgentState.stepsRemaining, + ) + }) +}) diff --git a/sdk/src/__tests__/state-snapshot-interval.test.ts b/sdk/src/__tests__/state-snapshot-interval.test.ts new file mode 100644 index 0000000000..1252c255e2 --- /dev/null +++ b/sdk/src/__tests__/state-snapshot-interval.test.ts @@ -0,0 +1,24 @@ +import { describe, expect, test } from 'bun:test' + +import { resolveStateSnapshotIntervalMs } from '../run' + +describe('resolveStateSnapshotIntervalMs', () => { + test('defaults to 5s when the host does not ask for anything', () => { + expect(resolveStateSnapshotIntervalMs()).toBe(5_000) + expect(resolveStateSnapshotIntervalMs(undefined)).toBe(5_000) + }) + + test('keeps an interval a host scaled to its transcript size', () => { + expect(resolveStateSnapshotIntervalMs(10_000)).toBe(10_000) + expect(resolveStateSnapshotIntervalMs(20_000)).toBe(20_000) + expect(resolveStateSnapshotIntervalMs(1_500.4)).toBe(1_500) + }) + + test('clamps, so checkpointing can never be disabled by the option', () => { + expect(resolveStateSnapshotIntervalMs(0)).toBe(1_000) + expect(resolveStateSnapshotIntervalMs(-5)).toBe(1_000) + expect(resolveStateSnapshotIntervalMs(Number.POSITIVE_INFINITY)).toBe(5_000) + expect(resolveStateSnapshotIntervalMs(Number.NaN)).toBe(5_000) + expect(resolveStateSnapshotIntervalMs(10 * 60_000)).toBe(60_000) + }) +}) diff --git a/sdk/src/run.ts b/sdk/src/run.ts index 3bbcc5d2ea..126db8b8e3 100644 --- a/sdk/src/run.ts +++ b/sdk/src/run.ts @@ -82,6 +82,7 @@ import type { import type { ToolResultOutput } from '@codebuff/common/types/messages/content-part' import type { PrintModeEvent } from '@codebuff/common/types/print-mode' import type { + AgentState, ConversationSummary, SessionState, } from '@codebuff/common/types/session-state' @@ -304,6 +305,15 @@ export type RunOptions = { * lose the in-flight turn. The final resolved RunState supersedes any * snapshot; no snapshots are emitted after the run settles. */ onStateSnapshot?: (runState: RunState) => void + /** How often `onStateSnapshot` fires while a run is in flight. Defaults to + * 5s. A host that knows how large the state it persists has grown can + * stretch this: cloning + serializing + fsyncing a multi-MB transcript every + * 5s costs seconds of CPU per minute of wall time, so the checkpoint interval + * should grow with the transcript, not stay fixed. Clamped to + * [1s, 60s] — the interval bounds how much a killed process loses, so a host + * cannot disable checkpointing by passing 0 or Infinity. The immediate + * checkpoint at run start is unaffected. */ + stateSnapshotIntervalMs?: number /** Provider-reported usage for each root-agent model request. */ onUsage?: (usage: AgentUsageData) => void /** A model request ended before an exact provider usage receipt arrived. */ @@ -381,6 +391,21 @@ export function applyTerminalShell( /** How often onStateSnapshot fires while a run is in flight. */ const STATE_SNAPSHOT_INTERVAL_MS = 5_000 +/** Bounds for `RunOptions.stateSnapshotIntervalMs`. The floor keeps the + * checkpoint live; the ceiling bounds the work a killed process can lose. */ +const MIN_STATE_SNAPSHOT_INTERVAL_MS = 1_000 +const MAX_STATE_SNAPSHOT_INTERVAL_MS = 60_000 + +export function resolveStateSnapshotIntervalMs(requested?: number): number { + if (typeof requested !== 'number' || !Number.isFinite(requested)) { + return STATE_SNAPSHOT_INTERVAL_MS + } + return Math.min( + MAX_STATE_SNAPSHOT_INTERVAL_MS, + Math.max(MIN_STATE_SNAPSHOT_INTERVAL_MS, Math.round(requested)), + ) +} + export const STATE_SNAPSHOT_INTERRUPTION_MESSAGE = 'The session ended before this response completed. Partial progress has been preserved.' @@ -428,6 +453,58 @@ export function cloneSessionState( return { fileContext: state.fileContext, mainAgentState } } +/** + * Copy an agent state's containers while sharing its message objects. + * + * Everything the runtime mutates in place is copied: the state object, the + * `messageHistory` array (pushed to), each subgoal's `logs` array (appended + * to), `childRunIds`/`ancestorRunIds` (pushed to when a subagent spawns), and + * the same recursively for subagent states. The message objects themselves are + * shared, which is what makes a snapshot cheap and — the reason this exists — + * lets a host recognise that a snapshot extends the previous one by pointer + * identity instead of re-serializing the whole history to find out. + * + * Message objects are immutable once appended: the runtime appends new ones and + * replaces the history array at step boundaries, and the one place that rewrites + * an existing message's content (`fitToolResults`, terminal-output shortening) + * copies it first, precisely so the UI's and trace writer's copies stay intact. + * The step loop itself already relies on this — it detects that history is an + * extension of what it sent with `inputHistory.every((m, i) => history[i] === m)`. + */ +function copyAgentStateSharingMessages(state: AgentState): AgentState { + const agentContext: AgentState['agentContext'] = {} + for (const [key, subgoal] of Object.entries(state.agentContext ?? {})) { + agentContext[key] = { ...subgoal, logs: [...(subgoal.logs ?? [])] } + } + return { + ...state, + messageHistory: [...state.messageHistory], + agentContext, + childRunIds: [...(state.childRunIds ?? [])], + ancestorRunIds: [...(state.ancestorRunIds ?? [])], + subagents: (state.subagents ?? []).map(copyAgentStateSharingMessages), + } +} + +/** + * The state a periodic checkpoint is built from: independent containers, shared + * message objects (see copyAgentStateSharingMessages). + * + * A host may keep a snapshot, serialize it, or resume from it, but must not + * mutate a message inside one: those objects belong to the live run. + * + * This is what deep-cloning on every snapshot costs: a JSON round-trip of the + * whole agent state, ~26 ms on a 32 MB history, on the host's thread every + * checkpoint interval — and it made the snapshots mutually unrelated, so a host + * persisting them could only ever write the entire transcript again. + */ +export function snapshotSessionState(state: SessionState): SessionState { + return { + fileContext: state.fileContext, + mainAgentState: copyAgentStateSharingMessages(state.mainAgentState), + } +} + /** * The session state a cancelled or errored turn persists, and the state a * follow-up prompt resumes from. @@ -453,11 +530,24 @@ export function buildCancelledSessionState(params: { /** Why the turn ended. Appended as a system-tagged user message. */ message: string logger?: Logger + /** Share message objects instead of deep-copying them (see + * snapshotSessionState). For periodic checkpoints, which are serialized + * immediately; every other caller keeps the deep copy, because it may hold + * the state while the run is still unwinding. */ + shareMessageEntries?: boolean }): SessionState { - const { sessionState, runtimeMadeProgress, promptMessage, message, logger } = - params + const { + sessionState, + runtimeMadeProgress, + promptMessage, + message, + logger, + shareMessageEntries, + } = params - const state = cloneSessionState(sessionState, logger) + const state = shareMessageEntries + ? snapshotSessionState(sessionState) + : cloneSessionState(sessionState, logger) const previousHistory = state.mainAgentState.messageHistory // A checkpoint can land after an assistant tool call is recorded but before // its result arrives. Drop that half-step at the persistence boundary so a @@ -619,6 +709,7 @@ async function runOnce({ costMode, extraCodebuffMetadata, onStateSnapshot, + stateSnapshotIntervalMs, onUsage, onUsageIncomplete, onCompaction, @@ -756,13 +847,17 @@ async function runOnce({ * any progress (replaced the shared messageHistory), those messages are * preserved. Otherwise the user's message is added so it isn't lost. */ - function getCancelledSessionState(message: string): SessionState { + function getCancelledSessionState( + message: string, + shareMessageEntries?: boolean, + ): SessionState { const runtimeMadeProgress = sessionState.mainAgentState.messageHistory !== initialMessageHistory return buildCancelledSessionState({ sessionState, runtimeMadeProgress, + shareMessageEntries, // Only add the user's message if the runtime didn't get a chance to. promptMessage: prompt || content @@ -776,10 +871,13 @@ async function runOnce({ logger, }) } - function getCancelledRunState(message?: string): RunState { + function getCancelledRunState( + message?: string, + shareMessageEntries?: boolean, + ): RunState { message = message ?? 'Run cancelled by user.' return { - sessionState: getCancelledSessionState(message), + sessionState: getCancelledSessionState(message, shareMessageEntries), traceSessionId, output: { type: 'error', @@ -1057,9 +1155,12 @@ async function runOnce({ if (onStateSnapshot) { // The runtime replaces mainAgentState.messageHistory with a new array at // each step boundary, so reference identity is a cheap "has anything - // durable changed" check. Skipping unchanged ticks avoids deep-cloning a - // potentially multi-MB sessionState every interval while the run is just - // waiting on a slow LLM call. + // durable changed" check. Skipping unchanged ticks avoids building a + // snapshot every interval while the run is just waiting on a slow LLM call. + // + // Snapshots share their message objects (snapshotSessionState) rather than + // deep-copying them: a host persisting them can then write only what is new + // since its own last save instead of re-serializing the whole transcript. let lastSnapshotHistory: unknown = null const emitStateSnapshot = () => { if (settled || signal?.aborted) { @@ -1072,7 +1173,7 @@ async function runOnce({ lastSnapshotHistory = history try { onStateSnapshot( - getCancelledRunState(STATE_SNAPSHOT_INTERRUPTION_MESSAGE), + getCancelledRunState(STATE_SNAPSHOT_INTERRUPTION_MESSAGE, true), ) } catch (error) { logger?.debug?.( @@ -1084,7 +1185,10 @@ async function runOnce({ // Emit immediately so the user's prompt is checkpointed as soon as the // run starts, then keep checkpointing progress while it is in flight. emitStateSnapshot() - snapshotTimer = setInterval(emitStateSnapshot, STATE_SNAPSHOT_INTERVAL_MS) + snapshotTimer = setInterval( + emitStateSnapshot, + resolveStateSnapshotIntervalMs(stateSnapshotIntervalMs), + ) // Don't let the checkpoint timer keep the host process alive. if (typeof (snapshotTimer as any).unref === 'function') { ;(snapshotTimer as any).unref()