From a1986a340ac3ff31e228a9e44230b36f3533b15b Mon Sep 17 00:00:00 2001 From: Jack Chen Date: Fri, 19 Jun 2026 15:23:11 +0800 Subject: [PATCH] feat(checkpoint): cut snapshot volume and re-run synthesis on runTeam restore (#314) Cuts the per-task O(N^2) shared-memory snapshot volume (reused-store skip) and re-runs coordinator synthesis on runTeam restore, plus docs/checkpoint.md and README coverage. Closes #20. Refs #294 #312 #313. --- CLAUDE.md | 1 + README.md | 1 + README_zh.md | 1 + docs/checkpoint.md | 134 +++++++++ packages/core/README.md | 3 + packages/core/README_zh.md | 3 + packages/core/src/memory/shared.ts | 10 + .../core/src/orchestrator/orchestrator.ts | 262 +++++++++++++----- packages/core/src/types.ts | 15 + packages/core/tests/checkpoint.test.ts | 149 +++++++++- 10 files changed, 505 insertions(+), 74 deletions(-) create mode 100644 docs/checkpoint.md diff --git a/CLAUDE.md b/CLAUDE.md index 2292486..a86907e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -79,6 +79,7 @@ Detailed behavior is documented in `docs/` — the single source of truth, so up | Tool presets, custom tools, sandbox, delegation, MCP | `tool/` | [tool-configuration.md](docs/tool-configuration.md) | | Providers, env vars, local servers, AI SDK bridge | `llm/` | [providers.md](docs/providers.md) | | Shared memory + custom backends | `memory/` | [shared-memory.md](docs/shared-memory.md) | +| Checkpoint/resume over `MemoryStore` | `memory/checkpoint.ts`, `orchestrator/orchestrator.ts` | [checkpoint.md](docs/checkpoint.md) | | Tracing, progress events, dashboard | `utils/trace.ts`, `dashboard/` | [observability.md](docs/observability.md) | | CLI usage + JSON schemas | `cli/oma.ts` | [cli.md](docs/cli.md) | diff --git a/README.md b/README.md index 895aea4..05527bf 100644 --- a/README.md +++ b/README.md @@ -146,6 +146,7 @@ npm test # run the test suite - [Tool configuration](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/tool-configuration.md) — tool presets, custom tools, the filesystem sandbox, and MCP. - [Observability](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/observability.md) — `onProgress` events, `onTrace` spans, and the post-run dashboard. - [Shared memory](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/shared-memory.md) — the default store and custom `MemoryStore` backends. +- [Checkpoint & resume](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/checkpoint.md) — opt-in per-run snapshot/resume over any `MemoryStore`; survive crashes and restarts. - [Context management](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/context-management.md) — sliding window, summarization, compaction, and custom compressors. - [CLI](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/cli.md) — the JSON-first `oma` binary for shell and CI. - [Consensus](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/consensus.md) — the `runConsensus` proposer→judge primitive, the per-task `verify` hook, and the budget invariant. diff --git a/README_zh.md b/README_zh.md index 4e1f40f..5571758 100644 --- a/README_zh.md +++ b/README_zh.md @@ -146,6 +146,7 @@ npm test # 跑测试套件 - [工具配置](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/tool-configuration.md) — 工具预设、自定义工具、文件系统沙箱、MCP。 - [可观测性](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/observability.md) — `onProgress` 事件、`onTrace` span、运行后 dashboard。 - [共享记忆](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/shared-memory.md) — 默认存储与自定义 `MemoryStore` 后端。 +- [Checkpoint & resume](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/checkpoint.md) — 可选的按运行快照/恢复，跑在任意 `MemoryStore` 上；崩溃、重启后可续跑。 - [上下文管理](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/context-management.md) — 滑动窗口、摘要、压缩、自定义压缩器。 - [CLI](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/cli.md) — 面向 shell 和 CI 的 JSON-first `oma` 命令行。 - [模型路由](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/model-routing.md) — 可选的 `modelRouting` 策略：按 phase / agent / role / priority / leaf 匹配，first match wins。 diff --git a/docs/checkpoint.md b/docs/checkpoint.md new file mode 100644 index 0000000..3c8d122 --- /dev/null +++ b/docs/checkpoint.md @@ -0,0 +1,134 @@ +# Checkpoint & Resume + +Long-running task workflows can persist their progress and resume after a crash, an abort, or a process restart. Checkpointing is **opt-in** and runs entirely over the existing [`MemoryStore`](shared-memory.md) interface, so the same in-memory, Redis, Postgres, or custom backend that holds shared memory also holds checkpoints — no extra storage layer. + +It covers the orchestration paths (`runTeam`, `runTasks`, `runFromPlan`, and `restore`). A single `runAgent` call has nothing to resume and is not checkpointed. + +## Enable it + +Pass `checkpoint` per call, or set a default for every run via `OrchestratorConfig.checkpoint`. Per-call options override the config default. + +```typescript +import { OpenMultiAgent, Team, InMemoryStore } from '@open-multi-agent/core' + +const store = new InMemoryStore() // or a durable MemoryStore (Redis, SQLite, ...) + +const team = new Team({ + name: 'research', + agents: [researcher, writer], + sharedMemoryStore: store, +}) + +const orchestrator = new OpenMultiAgent() + +// Snapshots are written after each completed task. +await orchestrator.runTasks(team, tasks, { checkpoint: { store } }) +``` + +`checkpoint: true` is shorthand: it reuses the team's shared-memory store when the team has one, otherwise a private in-memory store scoped to the orchestrator instance. + +```typescript +const orchestrator = new OpenMultiAgent({ checkpoint: true }) // default for all runs +``` + +### `CheckpointOptions` + +| Field | Type | Default | Purpose | +|-------|------|---------|---------| +| `enabled` | `boolean` | `true` | Set `false` to disable for a single run when a config default is on. | +| `store` | `MemoryStore` | team's shared-memory store | Durable backend for checkpoint records. | +| `runId` | `string` | — | Logical run id; derives a per-run checkpoint key. | +| `key` | `string` | — | Exact store key. Takes precedence over `runId`. | + +> **A `runId`, `key`, or explicit `store` is required when the team has no shared-memory store.** The instance-level fallback store is shared across every run on the orchestrator, so without a distinct key two concurrent runs would overwrite each other at the default checkpoint key. The call throws rather than risk a silent stomp. + +## Resume + +`restore()` loads the latest checkpoint, rebuilds the task queue and shared memory, skips completed tasks, and runs the remainder. + +```typescript +// After a crash/restart: same team wiring, same store. +const resumedTeam = new Team({ + name: 'research', + agents: [researcher, writer], + sharedMemoryStore: store, +}) + +const result = await orchestrator.restore(resumedTeam, { checkpoint: { store } }) +``` + +A restored `runTeam` run re-runs the coordinator synthesis, so you get the same synthesized final answer (under `result.agentResults.get('coordinator')`) as a fresh `runTeam`, not just the raw per-task outputs. Re-supply the coordinator config you used originally — the checkpoint can't persist a live adapter: + +```typescript +const result = await orchestrator.restore(resumedTeam, { + checkpoint: { store }, + coordinator: { provider: 'anthropic', model: 'claude-sonnet-4-6' }, // same as the original runTeam +}) +``` + +If synthesis can't run (no usable coordinator config or credentials) or the synthesis call fails, restore is best-effort: it returns the raw per-task outputs without a `'coordinator'` entry and emits an `onProgress` `synthesis_failed` event. `runTasks` / `runFromPlan` runs never synthesize. + +If no checkpoint is found, `restore()` falls back to a normal run of the tasks or plan you pass — so the same call works for both first run and resume: + +```typescript +// Fresh store → runs all tasks. Existing checkpoint → resumes, skipping done tasks. +await orchestrator.restore(team, tasks, { checkpoint: { store } }) +await orchestrator.restore(team, plan, { checkpoint: { store } }) // PlanArtifact +await orchestrator.restore(team, { checkpoint: { store } }) // resume-only, no-op on empty store +``` + +## What gets saved + +On each successfully completed task, the orchestrator writes a `CheckpointSnapshot`: + +- **Task queue state** — every task and its status partition (pending / in-progress / completed / failed / blocked / skipped). +- **Shared memory** — the turn counter is always recorded. The full entry snapshot is embedded **only when the checkpoint store differs from the team's shared-memory store**. When they are the same store (the default for `checkpoint: true`), the entries are already durable there, so re-embedding them every task would be wasted ~O(N²) write volume across a long run; resume reads them straight from the store instead. Either way, resume rehydrates shared memory correctly. +- **Completed task results** — `taskId`, `assignee`, and `result` for each finished task, so resumed agents see prior outputs. + +Snapshots are stored as JSON under a reserved namespace: `__oma_checkpoint__//latest` (or `__oma_checkpoint__/latest` when no `runId` is set). Keys under `__oma_checkpoint__/` are reserved — shared-memory snapshot/restore deliberately skips them so one store can hold both agent memory and checkpoints. + +### Saves are best-effort + +A checkpoint write must never take down the run it protects. If the store rejects (a transient Redis/SQLite error), the failure is surfaced via `onProgress` and the run continues; the next completed task retries the write. + +```typescript +const orchestrator = new OpenMultiAgent({ + onProgress(event) { + if (event.type === 'error' && event.data?.kind === 'checkpoint_save_failed') { + console.warn('checkpoint write failed, run continues:', event.data.error) + } + }, +}) +``` + +## Advanced: the `Checkpoint` class + +For inspecting or managing checkpoints directly, the manager and key helpers are exported: + +```typescript +import { + Checkpoint, + checkpointKey, + isCheckpointKey, + CHECKPOINT_KEY_PREFIX, + DEFAULT_CHECKPOINT_KEY, +} from '@open-multi-agent/core' + +const cp = new Checkpoint(store, { runId: 'nightly-2026-06-18' }) +const snapshot = await cp.loadLatest() // CheckpointSnapshot | null +await cp.delete() // drop the persisted checkpoint +``` + +## Limitations + +Per-run snapshot/restore over `MemoryStore`. What it does *not* yet do: + +- **Resume is task-grained, not mid-task.** A task interrupted while running re-runs from the start on resume — the in-flight agent's conversation history inside a running task is not persisted. Recovery happens at task boundaries. +- **Snapshot-based, not event-sourced.** Each checkpoint overwrites the previous one; there is no transition log to replay. + +Two notes on the shared-memory optimization described above: + +- A *separate* durable checkpoint store (shared memory in store X, `checkpoint: { store: Y }`) still embeds the full memory snapshot on each save — necessary, since Y holds no other copy of the entries. +- The reused-store path does not point-in-time roll back shared memory. The default framework writes results only on task completion (so a crashed in-progress task has written nothing), but a custom tool that writes to shared memory mid-task would not have those partial writes rolled back on resume. + +These are tracked as follow-ups: [#312](https://github.com/open-multi-agent/open-multi-agent/issues/312) (mid-task recovery) and [#313](https://github.com/open-multi-agent/open-multi-agent/issues/313) (event-sourced replay). diff --git a/packages/core/README.md b/packages/core/README.md index 984adc6..53909f8 100644 --- a/packages/core/README.md +++ b/packages/core/README.md @@ -169,6 +169,7 @@ Route orchestration phases to different models with an opt-in `modelRouting` pol | **Configurable coordinator** | Override the coordinator's model, provider, adapter, system prompt, or tools via `runTeam(team, goal, { coordinator })`. | | **Observability** | `onProgress` events, `onTrace` spans, post-run HTML dashboard rendering the executed task DAG. API keys and tokens are redacted from traces, bash output, and the dashboard. ([observability guide](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/observability.md)) | | **Pluggable shared memory** | Default in-process KV; swap in Redis / Postgres / your own backend by implementing `MemoryStore`. ([shared memory](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/shared-memory.md)) | +| **Checkpoint & resume** | Opt-in per-run checkpointing over any `MemoryStore`: snapshot on each completed task, then `restore()` skips finished tasks to continue after a crash or restart. Best-effort saves never take the run down. ([checkpoint & resume](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/checkpoint.md)) | | **Sandboxed filesystem workspace** | Built-in filesystem tools are sandboxed to `/.agent-workspace` by default; agents sharing the default configuration share this root. For per-agent isolation, set `AgentConfig.cwd`; for a different shared root, set `OrchestratorConfig.defaultCwd`; pass `null` to disable. ([sandbox config](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/tool-configuration.md)) | Production controls ([context strategies](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/context-management.md), task retry with backoff, loop detection, tool output truncation/compression) are covered in the [Production Checklist](#production-checklist). @@ -395,6 +396,7 @@ Before going live, wire up the controls that protect token spend, recover from f | Bound wall-clock time | `timeoutMs` per agent (aborts a run that hangs, common with local models) | `AgentConfig` | | Cap tool output | `maxToolOutputChars` (or per-tool `maxOutputChars`) + `compressToolResults: true` | `AgentConfig` and `defineTool()` | | Recover from failure | Per-task `maxRetries`, `retryDelayMs`, `retryBackoff` (exponential multiplier) | Task config used via `runTasks()` | +| Survive a crash or restart | `checkpoint` (pass a `runId` or a durable `MemoryStore`) + `restore()` to resume, skipping completed tasks | `OrchestratorConfig` / run options | | Hard-cap spend | `maxTokenBudget` on the orchestrator | `OrchestratorConfig` | | Catch stuck agents | `loopDetection` with `onLoopDetected: 'terminate'` (or a custom handler) | `AgentConfig` | | Trace and audit | `onTrace` to your tracing backend; persist `renderTeamRunDashboard(result)` | `OrchestratorConfig` | @@ -408,6 +410,7 @@ Before going live, wire up the controls that protect token spend, recover from f - [Tool configuration](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/tool-configuration.md) — tool presets, custom tools, the filesystem sandbox, and MCP. - [Observability](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/observability.md) — `onProgress` events, `onTrace` spans, and the post-run dashboard. - [Shared memory](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/shared-memory.md) — the default store and custom `MemoryStore` backends. +- [Checkpoint & resume](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/checkpoint.md) — opt-in per-run snapshot/resume over any `MemoryStore`; survive crashes and restarts. - [Context management](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/context-management.md) — sliding window, summarization, compaction, and custom compressors. - [CLI](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/cli.md) — the JSON-first `oma` binary for shell and CI. - [Consensus](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/consensus.md) — the `runConsensus` proposer→judge primitive, the per-task `verify` hook, and the budget invariant. diff --git a/packages/core/README_zh.md b/packages/core/README_zh.md index f6013cb..49f68e9 100644 --- a/packages/core/README_zh.md +++ b/packages/core/README_zh.md @@ -167,6 +167,7 @@ const result = await orchestrator.runFromPlan(team, plan) | **可配置协调者** | 通过 `runTeam(team, goal, { coordinator })` 覆盖协调者的 model、provider、adapter、system prompt 或工具。 | | **可观测性** | `onProgress` 事件、`onTrace` span，运行结束后渲染任务 DAG 的 HTML dashboard。API key 和 token 会从 trace、bash 输出和 dashboard 中自动脱敏。([可观测性指南](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/observability.md)) | | **可插拔共享记忆** | 默认进程内 KV；实现 `MemoryStore` 接口即可换 Redis / Postgres / 自家后端。([共享记忆](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/shared-memory.md)) | +| **Checkpoint & resume** | 可选的按运行 checkpoint，跑在任意 `MemoryStore` 上：每个任务完成时快照，`restore()` 跳过已完成任务，崩溃或重启后续跑。存盘 best-effort，不会拖垮运行。([checkpoint & resume](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/checkpoint.md)) | | **沙箱化文件系统工作目录** | 内置文件系统工具默认沙箱化在 `/.agent-workspace`；继承默认配置的 agent 共享同一根目录。需要 per-agent 隔离时显式设置 `AgentConfig.cwd`；改换共享根目录用 `OrchestratorConfig.defaultCwd`；传 `null` 关闭沙箱。([沙箱配置](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/tool-configuration.md)) | 生产级控制（[上下文策略](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/context-management.md)、任务重试退避、循环检测、工具输出截断/压缩）见 [生产级检查清单](#生产级检查清单)。 @@ -393,6 +394,7 @@ await oma.runAgent( | 控制运行时长 | `timeoutMs`（每个 agent，运行挂起时中止；本地模型常见） | `AgentConfig` | | 限制工具输出 | `maxToolOutputChars`（或单工具 `maxOutputChars`）+ `compressToolResults: true` | `AgentConfig` 和 `defineTool()` | | 失败重试 | 任务级 `maxRetries`、`retryDelayMs`、`retryBackoff`（指数退避倍率） | 通过 `runTasks()` 用的任务配置 | +| 崩溃/重启后续跑 | `checkpoint`（给 `runId` 或持久化 `MemoryStore`）+ `restore()` 续跑，跳过已完成任务 | `OrchestratorConfig` / 运行选项 | | 总额封顶 | orchestrator 上设 `maxTokenBudget` | `OrchestratorConfig` | | 卡死检测 | `loopDetection` + `onLoopDetected: 'terminate'`（或自定义 handler） | `AgentConfig` | | 追踪与审计 | `onTrace` 接你的 tracing 后端；落盘 `renderTeamRunDashboard(result)` | `OrchestratorConfig` | @@ -406,6 +408,7 @@ await oma.runAgent( - [工具配置](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/tool-configuration.md) — 工具预设、自定义工具、文件系统沙箱、MCP。 - [可观测性](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/observability.md) — `onProgress` 事件、`onTrace` span、运行后 dashboard。 - [共享记忆](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/shared-memory.md) — 默认存储与自定义 `MemoryStore` 后端。 +- [Checkpoint & resume](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/checkpoint.md) — 可选的按运行快照/恢复，跑在任意 `MemoryStore` 上；崩溃、重启后可续跑。 - [上下文管理](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/context-management.md) — 滑动窗口、摘要、压缩、自定义压缩器。 - [CLI](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/cli.md) — 面向 shell 和 CI 的 JSON-first `oma` 命令行。 - [模型路由](https://github.com/open-multi-agent/open-multi-agent/blob/main/docs/model-routing.md) — 可选的 `modelRouting` 策略：按 phase / agent / role / priority / leaf 匹配，first match wins。 diff --git a/packages/core/src/memory/shared.ts b/packages/core/src/memory/shared.ts index 36dd7bb..88b6e22 100644 --- a/packages/core/src/memory/shared.ts +++ b/packages/core/src/memory/shared.ts @@ -118,6 +118,16 @@ export class SharedMemory { return this.turnCount } + /** + * Restore the turn counter without rehydrating entries. Used on checkpoint + * resume when the shared-memory store is reused: the entries are already + * durable in the store, so only the counter needs restoring to keep TTL + * expiry correct. + */ + setTurnCount(turnCount: number): void { + this.turnCount = turnCount + } + // --------------------------------------------------------------------------- // Snapshot / restore // --------------------------------------------------------------------------- diff --git a/packages/core/src/orchestrator/orchestrator.ts b/packages/core/src/orchestrator/orchestrator.ts index a316d30..165dce0 100644 --- a/packages/core/src/orchestrator/orchestrator.ts +++ b/packages/core/src/orchestrator/orchestrator.ts @@ -626,6 +626,13 @@ interface ActiveCheckpoint { readonly mode: CheckpointSnapshot['mode'] readonly goal?: string readonly runId?: string + /** + * True when the checkpoint store is the same object as the team's + * shared-memory store. In that case the memory entries are already durable + * in the store, so the checkpoint omits the full shared-memory snapshot + * (avoids ~O(N^2) write volume) and persists only the turn counter. + */ + readonly reusesSharedMemoryStore: boolean saveChain: Promise } @@ -748,7 +755,14 @@ async function saveRunCheckpoint(queue: TaskQueue, ctx: RunContext): Promise 0), - maxTurns: coordinatorOverrides?.maxTurns ?? 3, - maxTokens: coordinatorOverrides?.maxTokens, - temperature: coordinatorOverrides?.temperature, - topP: coordinatorOverrides?.topP, - topK: coordinatorOverrides?.topK, - minP: coordinatorOverrides?.minP, - parallelToolCalls: coordinatorOverrides?.parallelToolCalls, - frequencyPenalty: coordinatorOverrides?.frequencyPenalty, - presencePenalty: coordinatorOverrides?.presencePenalty, - extraBody: coordinatorOverrides?.extraBody, - toolPreset: coordinatorOverrides?.toolPreset, - tools: coordinatorOverrides?.tools, - disallowedTools: coordinatorOverrides?.disallowedTools, - cwd: coordinatorOverrides?.cwd === undefined - ? this.config.defaultCwd - : coordinatorOverrides.cwd, - loopDetection: coordinatorOverrides?.loopDetection, - timeoutMs: coordinatorOverrides?.timeoutMs, - } + const coordinatorBaseConfig = this.buildCoordinatorBaseConfig( + coordinatorOverrides, + agentConfigs, + (options?.verifyJudges?.length ?? 0) > 0, + ) const coordinatorConfig = withModelRoute( coordinatorBaseConfig, routeMatches(options?.modelRouting, { phase: 'coordinator', agent: 'coordinator' }), @@ -1966,46 +1958,19 @@ export class OpenMultiAgent { // ------------------------------------------------------------------ // Step 5: Coordinator synthesises final result // ------------------------------------------------------------------ - if (options?.abortSignal?.aborted) { - return this.buildTeamRunResult(agentResults, goal, taskRecords) - } - if ( - maxTokenBudget !== undefined - && cumulativeUsage.input_tokens + cumulativeUsage.output_tokens > maxTokenBudget - ) { - return this.buildTeamRunResult(agentResults, goal, taskRecords) - } - const synthesisPrompt = await this.buildSynthesisPrompt(goal, queue.list(), team) - const synthesisAgent = buildAgent(withModelRoute( - coordinatorBaseConfig, - routeMatches(options?.modelRouting, { phase: 'synthesis', agent: 'coordinator' }), - )) - const synthTraceOptions: Partial | undefined = this.config.onTrace - ? { onTrace: this.config.onTrace, runId: runId ?? '', traceAgent: 'coordinator' } - : undefined - const synthesisResult = await synthesisAgent.run(synthesisPrompt, synthTraceOptions) - agentResults.set('coordinator', synthesisResult) - cumulativeUsage = addUsage(cumulativeUsage, synthesisResult.tokenUsage) - if ( - maxTokenBudget !== undefined - && cumulativeUsage.input_tokens + cumulativeUsage.output_tokens > maxTokenBudget - ) { - this.config.onProgress?.({ - type: 'budget_exceeded', - agent: 'coordinator', - data: new TokenBudgetExceededError( - 'coordinator', - cumulativeUsage.input_tokens + cumulativeUsage.output_tokens, - maxTokenBudget, - ), - }) - } - - this.config.onProgress?.({ - type: 'agent_complete', - agent: 'coordinator', - data: synthesisResult, + const synthesis = await this.runCoordinatorSynthesis(team, queue, goal, coordinatorBaseConfig, { + modelRouting: options?.modelRouting, + runId, + abortSignal: options?.abortSignal, + cumulativeUsage, + maxTokenBudget, }) + if (synthesis === null) { + // Aborted or already over budget — return raw task outputs, no synthesis. + return this.buildTeamRunResult(agentResults, goal, taskRecords) + } + agentResults.set('coordinator', synthesis.result) + cumulativeUsage = synthesis.cumulativeUsage // Note: coordinator decompose and synthesis are internal meta-steps. // Only actual user tasks (non-coordinator keys) are counted in @@ -2088,10 +2053,14 @@ export class OpenMultiAgent { * remainder. When no checkpoint is found the call falls back to a normal run of * the provided tasks/plan (or a no-op when neither is given). * - * Execution always flows through the explicit-task-queue path, so — like - * {@link OpenMultiAgent.runFromPlan} — a restored `runTeam` run returns the raw - * per-task outputs in `tasks`, not a coordinator-synthesized final answer. Pass - * the original tasks/plan to resume a `runTasks`/`runFromPlan` run unchanged. + * A resumed `runTeam` run re-runs the coordinator synthesis so the result + * matches a fresh `runTeam` (a synthesized final answer under the + * `'coordinator'` key in `agentResults`, not just raw per-task outputs). + * Re-supply the coordinator via `options.coordinator` — the checkpoint cannot + * persist a live adapter. If no usable coordinator config is available or the + * synthesis call fails, restore falls back to raw outputs and emits an + * `onProgress` `synthesis_failed` event. A restored `runTasks`/`runFromPlan` + * run never synthesizes; pass the original tasks/plan to resume it unchanged. */ async restore( team: Team, @@ -2183,6 +2152,10 @@ export class OpenMultiAgent { const sharedMem = team.getSharedMemoryInstance() if (sharedMem && snapshot.sharedMemory) { await sharedMem.restore(snapshot.sharedMemory) + } else if (sharedMem && snapshot.turnCount !== undefined) { + // Reused-store checkpoint: entries are already in the store; only the + // turn counter needs restoring so TTL expiry resumes correctly. + sharedMem.setTurnCount(snapshot.turnCount) } const queue = TaskQueue.fromSnapshot(snapshot.queue, { resetInProgress: true }) @@ -2203,6 +2176,7 @@ export class OpenMultiAgent { snapshot.goal ?? options?.goal, agentResults, checkpointForResume, + options?.coordinator, ) } @@ -2500,6 +2474,109 @@ export class OpenMultiAgent { ].join('\n') } + /** + * Build the base coordinator {@link AgentConfig} shared by the decomposition + * and synthesis passes. Falls back to orchestrator defaults for any field the + * caller's {@link CoordinatorConfig} leaves unset. + */ + private buildCoordinatorBaseConfig( + coordinatorOverrides: CoordinatorConfig | undefined, + agentConfigs: AgentConfig[], + hasVerifyJudges: boolean, + ): AgentConfig { + return { + name: 'coordinator', + model: coordinatorOverrides?.model ?? this.config.defaultModel, + ...(coordinatorOverrides?.adapter !== undefined ? { adapter: coordinatorOverrides.adapter } : {}), + provider: coordinatorOverrides?.provider ?? this.config.defaultProvider, + baseURL: coordinatorOverrides?.baseURL ?? this.config.defaultBaseURL, + apiKey: coordinatorOverrides?.apiKey ?? this.config.defaultApiKey, + systemPrompt: this.buildCoordinatorPrompt(agentConfigs, coordinatorOverrides, hasVerifyJudges), + maxTurns: coordinatorOverrides?.maxTurns ?? 3, + maxTokens: coordinatorOverrides?.maxTokens, + temperature: coordinatorOverrides?.temperature, + topP: coordinatorOverrides?.topP, + topK: coordinatorOverrides?.topK, + minP: coordinatorOverrides?.minP, + parallelToolCalls: coordinatorOverrides?.parallelToolCalls, + frequencyPenalty: coordinatorOverrides?.frequencyPenalty, + presencePenalty: coordinatorOverrides?.presencePenalty, + extraBody: coordinatorOverrides?.extraBody, + toolPreset: coordinatorOverrides?.toolPreset, + tools: coordinatorOverrides?.tools, + disallowedTools: coordinatorOverrides?.disallowedTools, + cwd: coordinatorOverrides?.cwd === undefined + ? this.config.defaultCwd + : coordinatorOverrides.cwd, + loopDetection: coordinatorOverrides?.loopDetection, + timeoutMs: coordinatorOverrides?.timeoutMs, + } + } + + /** + * Run the coordinator synthesis pass over completed task results. Returns the + * synthesis result plus updated cumulative usage, or `null` when synthesis is + * skipped (run aborted, or the token budget was already exhausted before the + * pass). Emits `budget_exceeded` (when synthesis tips over budget) and + * `agent_complete`, mirroring the inline `runTeam` path. Does not mutate + * `agentResults` — the caller records the `'coordinator'` entry. + */ + private async runCoordinatorSynthesis( + team: Team, + queue: TaskQueue, + goal: string, + coordinatorBaseConfig: AgentConfig, + opts: { + readonly modelRouting?: ModelRoutingPolicy + readonly runId?: string + readonly abortSignal?: AbortSignal + readonly cumulativeUsage: TokenUsage + readonly maxTokenBudget?: number + }, + ): Promise<{ readonly result: AgentRunResult; readonly cumulativeUsage: TokenUsage } | null> { + if (opts.abortSignal?.aborted) return null + if ( + opts.maxTokenBudget !== undefined + && opts.cumulativeUsage.input_tokens + opts.cumulativeUsage.output_tokens > opts.maxTokenBudget + ) { + return null + } + + const synthesisPrompt = await this.buildSynthesisPrompt(goal, queue.list(), team) + const synthesisAgent = buildAgent(withModelRoute( + coordinatorBaseConfig, + routeMatches(opts.modelRouting, { phase: 'synthesis', agent: 'coordinator' }), + )) + const synthTraceOptions: Partial | undefined = this.config.onTrace + ? { onTrace: this.config.onTrace, runId: opts.runId ?? '', traceAgent: 'coordinator' } + : undefined + const result = await synthesisAgent.run(synthesisPrompt, synthTraceOptions) + const cumulativeUsage = addUsage(opts.cumulativeUsage, result.tokenUsage) + + if ( + opts.maxTokenBudget !== undefined + && cumulativeUsage.input_tokens + cumulativeUsage.output_tokens > opts.maxTokenBudget + ) { + this.config.onProgress?.({ + type: 'budget_exceeded', + agent: 'coordinator', + data: new TokenBudgetExceededError( + 'coordinator', + cumulativeUsage.input_tokens + cumulativeUsage.output_tokens, + opts.maxTokenBudget, + ), + }) + } + + this.config.onProgress?.({ + type: 'agent_complete', + agent: 'coordinator', + data: result, + }) + + return { result, cumulativeUsage } + } + /** Build the synthesis prompt shown to the coordinator after all tasks complete. */ private async buildSynthesisPrompt( goal: string, @@ -2572,6 +2649,7 @@ export class OpenMultiAgent { goal?: string, initialAgentResults?: Map, activeCheckpoint?: ActiveCheckpoint, + coordinatorForSynthesis?: CoordinatorConfig, ): Promise { const agentConfigs = team.getAgents() const scheduler = new Scheduler('dependency-first') @@ -2606,6 +2684,44 @@ export class OpenMultiAgent { await executeQueue(queue, ctx) + // A resumed `runTeam` re-runs the coordinator synthesis so the restored + // result matches a fresh `runTeam` (a synthesized final answer, not raw + // per-task outputs). Best-effort: a missing/unusable coordinator config or + // a failing synthesis call must not discard the recovered work — on failure + // we surface `synthesis_failed` and fall back to raw outputs. + if (checkpoint?.mode === 'runTeam' && goal !== undefined) { + try { + const coordinatorBaseConfig = this.buildCoordinatorBaseConfig(coordinatorForSynthesis, agentConfigs, false) + const synthesis = await this.runCoordinatorSynthesis(team, queue, goal, coordinatorBaseConfig, { + modelRouting: options?.modelRouting, + runId: ctx.runId, + abortSignal: options?.abortSignal, + cumulativeUsage: ctx.cumulativeUsage, + maxTokenBudget: ctx.maxTokenBudget, + }) + if (synthesis !== null && synthesis.result.success) { + agentResults.set('coordinator', synthesis.result) + ctx.cumulativeUsage = synthesis.cumulativeUsage + } else if (synthesis !== null) { + // Synthesis ran but the coordinator agent failed (e.g. the LLM call + // errored). Keep the recovered task outputs and surface the failure + // rather than attaching a failed answer under `'coordinator'`. + this.config.onProgress?.({ + type: 'error', + data: { + kind: 'synthesis_failed', + error: new Error(synthesis.result.output || 'coordinator synthesis failed'), + }, + }) + } + } catch (error) { + this.config.onProgress?.({ + type: 'error', + data: { kind: 'synthesis_failed', error }, + }) + } + } + const taskRecords: readonly TaskExecutionRecord[] = queue.list().map((task) => ({ id: task.id, title: task.title, @@ -2638,7 +2754,8 @@ export class OpenMultiAgent { // orchestrator, so concurrent runs would overwrite each other at the // default checkpoint key. Require a `runId` (or an explicit `key`/`store`) // before falling back, so each run resolves to a distinct, resumable key. - const explicitStore = options.store ?? team.getSharedMemory() + const sharedStore = team.getSharedMemory() + const explicitStore = options.store ?? sharedStore if (!explicitStore && options.runId === undefined && options.key === undefined) { throw new Error( 'Checkpoint requires a `runId` (or an explicit `store`/`key`) when the team has no ' + @@ -2652,6 +2769,7 @@ export class OpenMultiAgent { mode, ...(goal !== undefined ? { goal } : {}), ...(options.runId !== undefined ? { runId: options.runId } : {}), + reusesSharedMemoryStore: sharedStore !== undefined && store === sharedStore, saveChain: Promise.resolve(), } } diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index b9e271c..acc7c1d 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -1118,6 +1118,15 @@ export interface CheckpointOptions { export interface RestoreOptions extends RunTasksOptions { /** Optional goal attached to the restored result when no checkpoint goal exists. */ readonly goal?: string + /** + * Coordinator config used to re-run synthesis when resuming a `runTeam` run. + * Re-supply the same config passed to the original `runTeam` call — the + * checkpoint cannot persist a live {@link CoordinatorConfig.adapter}. When + * synthesis cannot run (no usable config/credentials) or the LLM call fails, + * restore falls back to raw per-task outputs and emits an `onProgress` + * `synthesis_failed` event. + */ + readonly coordinator?: CoordinatorConfig } /** Serializable form of a {@link MemoryEntry}. */ @@ -1183,6 +1192,12 @@ export interface CheckpointSnapshot { readonly goal?: string readonly queue: TaskQueueSnapshot readonly sharedMemory?: SharedMemorySnapshot + /** + * Shared-memory turn counter at checkpoint time. Persisted separately from + * {@link sharedMemory} so TTL/expiry stays correct on resume even when the + * full entry snapshot is omitted (checkpoint store === shared-memory store). + */ + readonly turnCount?: number readonly completedTaskResults: readonly CompletedTaskCheckpoint[] } diff --git a/packages/core/tests/checkpoint.test.ts b/packages/core/tests/checkpoint.test.ts index 796d3c9..4e036d4 100644 --- a/packages/core/tests/checkpoint.test.ts +++ b/packages/core/tests/checkpoint.test.ts @@ -211,7 +211,11 @@ describe('OpenMultiAgent checkpoint/restore', () => { }) it('restores after an aborted run, skips completed tasks, and rehydrates shared memory', async () => { + // Separate checkpoint store: the embedded shared-memory snapshot is what + // rehydrates `store` after it is wiped (simulating a non-durable shared + // store across a restart). The reused-store path is covered separately. const store = new InMemoryStore() + const checkpointStore = new InMemoryStore() const scripted = scriptedAdapter(['first output', 'second output']) const abort = new AbortController() const orchestrator = new OpenMultiAgent({ @@ -230,7 +234,7 @@ describe('OpenMultiAgent checkpoint/restore', () => { await orchestrator.runTasks(team, tasks, { abortSignal: abort.signal, - checkpoint: { store }, + checkpoint: { store: checkpointStore }, }) expect(scripted.calls()).toBe(1) @@ -241,7 +245,7 @@ describe('OpenMultiAgent checkpoint/restore', () => { agents: [worker('worker', scripted.adapter)], sharedMemoryStore: store, }) - const restored = await orchestrator.restore(resumedTeam, { checkpoint: { store } }) + const restored = await orchestrator.restore(resumedTeam, { checkpoint: { store: checkpointStore } }) expect(scripted.calls()).toBe(2) expect(scripted.prompts[1]).toContain('first output') @@ -325,6 +329,69 @@ describe('OpenMultiAgent checkpoint/restore', () => { expect(scripted.calls()).toBe(2) expect(result.tasks?.map((record) => record.status)).toEqual(['completed', 'completed']) }) + + it('reused store omits the shared-memory snapshot but persists the turn counter', async () => { + const store = new InMemoryStore() + const scripted = scriptedAdapter(['first output', 'second output']) + const abort = new AbortController() + const orchestrator = new OpenMultiAgent({ + onProgress(event) { + if (event.type === 'task_complete') abort.abort() + }, + }) + const team = new Team({ + name: 'team', + agents: [worker('worker', scripted.adapter)], + sharedMemoryStore: store, + }) + await team.getSharedMemoryInstance()!.writeExpiring('seed', 'ttl', 'short', 5) + + await orchestrator.runTasks(team, tasks, { abortSignal: abort.signal, checkpoint: { store } }) + + // Checkpoint store === shared-memory store: the entries are already durable + // in the store, so the snapshot omits them and records only the turn count. + const persisted = await new Checkpoint(store, {}).loadLatest() + expect(persisted?.sharedMemory).toBeUndefined() + expect(persisted?.turnCount).toBe(1) + + // Resume restores the turn counter so TTL expiry continues correctly. + const resumedTeam = new Team({ + name: 'team', + agents: [worker('worker', scripted.adapter)], + sharedMemoryStore: store, + }) + await orchestrator.restore(resumedTeam, { checkpoint: { store } }) + expect(resumedTeam.getSharedMemoryInstance()!.getTurnCount()).toBe(2) + expect((await resumedTeam.getSharedMemoryInstance()!.read('seed/ttl'))?.value).toBe('short') + }) + + it('separate checkpoint store still embeds the shared-memory snapshot', async () => { + const sharedStore = new InMemoryStore() + const checkpointStore = new InMemoryStore() + const scripted = scriptedAdapter(['first output', 'second output']) + const abort = new AbortController() + const orchestrator = new OpenMultiAgent({ + onProgress(event) { + if (event.type === 'task_complete') abort.abort() + }, + }) + const team = new Team({ + name: 'team', + agents: [worker('worker', scripted.adapter)], + sharedMemoryStore: sharedStore, + }) + await team.getSharedMemoryInstance()!.write('seed', 'note', { keep: true }) + + await orchestrator.runTasks(team, tasks, { + abortSignal: abort.signal, + checkpoint: { store: checkpointStore }, + }) + + // Checkpoint store differs from the shared-memory store, so the snapshot must + // embed the entries — the checkpoint store holds no other copy. + const persisted = await new Checkpoint(checkpointStore, {}).loadLatest() + expect(persisted?.sharedMemory?.entries.some((entry) => entry.key === 'seed/note')).toBe(true) + }) }) /** A store whose writes always reject, to exercise best-effort checkpointing. */ @@ -430,3 +497,81 @@ describe('checkpoint resilience and key safety', () => { expect(result.tasks?.map((record) => record.status)).toEqual(['completed', 'completed']) }) }) + +describe('runTeam restore synthesis', () => { + /** Persist a runTeam-mode checkpoint with `first` completed and `second` pending. */ + function saveRunTeamCheckpoint(store: MemoryStore, goal = 'achieve the goal') { + const queue = new TaskQueue() + queue.add(task('first', { assignee: 'worker' })) + queue.add(task('second', { assignee: 'worker', dependsOn: ['first'] })) + queue.complete('first', 'first output') + return new Checkpoint(store, {}).save({ + version: 1, + mode: 'runTeam', + createdAt: new Date().toISOString(), + goal, + queue: queue.snapshot(), + completedTaskResults: [{ taskId: 'first', assignee: 'worker', result: 'first output' }], + }) + } + + it('re-runs the coordinator synthesis and returns the synthesized answer', async () => { + const store = new InMemoryStore() + const workerAdapter = scriptedAdapter(['second output']) + const coordinator = scriptedAdapter(['SYNTHESIZED ANSWER']) + const orchestrator = new OpenMultiAgent() + const team = new Team({ + name: 'team', + agents: [worker('worker', workerAdapter.adapter)], + sharedMemoryStore: store, + }) + await saveRunTeamCheckpoint(store) + + const restored = await orchestrator.restore(team, { + checkpoint: { store }, + coordinator: { model: 'mock-model', adapter: coordinator.adapter }, + }) + + expect(workerAdapter.calls()).toBe(1) // only the pending 'second' task ran + expect(coordinator.calls()).toBe(1) // synthesis ran; restore does not re-decompose + expect(restored.agentResults.get('coordinator')?.output).toBe('SYNTHESIZED ANSWER') + expect(restored.tasks?.every((record) => record.status === 'completed')).toBe(true) + }) + + it('is best-effort when synthesis fails: raw outputs plus a synthesis_failed event', async () => { + const store = new InMemoryStore() + const workerAdapter = scriptedAdapter(['second output']) + const throwingCoordinator: LLMAdapter = { + name: 'throwing-coordinator', + async chat() { + throw new Error('synthesis boom') + }, + async *stream() { + yield { type: 'done' as const, data: textResponse('unused', 'mock-model') } + }, + } + const events: OrchestratorEvent[] = [] + const orchestrator = new OpenMultiAgent({ onProgress: (event) => events.push(event) }) + const team = new Team({ + name: 'team', + agents: [worker('worker', workerAdapter.adapter)], + sharedMemoryStore: store, + }) + await saveRunTeamCheckpoint(store) + + const restored = await orchestrator.restore(team, { + checkpoint: { store }, + coordinator: { model: 'mock-model', adapter: throwingCoordinator }, + }) + + expect(restored.agentResults.has('coordinator')).toBe(false) // synthesis skipped + expect(restored.tasks?.every((record) => record.status === 'completed')).toBe(true) // work preserved + expect( + events.some( + (event) => + event.type === 'error' && + (event.data as { kind?: string } | undefined)?.kind === 'synthesis_failed', + ), + ).toBe(true) + }) +})