sdk: host-configurable state snapshot interval, and snapshots that share messages

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'))

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.

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,
+    )
+  })
+})

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)
+  })
+})