From 42baed8c8cbefddd365238a9120db267f36e1b05 Mon Sep 17 00:00:00 2001 From: Vanzeren <53075619+Vanzeren@users.noreply.github.com> Date: Wed, 22 Jul 2026 08:33:29 +0800 Subject: [PATCH] feat(checkpoint): dual-mode checkpoint storage with LangGraph DeltaChannel (#4292) * feat(checkpoint): dual-mode checkpoint storage with LangGraph DeltaChannel Add a restart-required database.checkpoint_channel_mode ("full" default, "delta") that stores the messages channel via LangGraph 1.2 DeltaChannel, cutting checkpoint storage from O(n^2) to O(n) for append-only history. Existing full checkpoints seed delta state transparently; no data migration. - config: mode schema + freeze-on-first-use with CheckpointModeReconfigurationError; mode marker persisted in checkpoint metadata; unsafe delta->full downgrade rejected fail-closed with CheckpointModeMismatchError (run-level error, failed state read) - state: delta message state schema; CheckpointStateAccessor centralizes materialized reads for all consumers (threads API, branches, regeneration, compaction, state updates, memory, goal workers) - runtime: raw writers (run durations, interrupted title, thread goal) parent their checkpoints to the checkpoint they derive from, preserving delta ancestry; rollback forks the pre-run lineage through a state mutation graph with Overwrite restores; InMemorySaver delta-history override delegates to the base walk (fixes dropped first write after migration, also present upstream) - tests: conformance suite over {memory, sqlite, postgres} covering migration replay, stable message IDs, storage shape and writer preservation; conftest fixture isolates the frozen mode between tests; stale config fakes refreshed - ci: backend unit tests gain a postgres service * fix(checkpoint): close materialization gaps in goal flow, guard public factory - Route goal-continuation message reads through CheckpointStateAccessor: raw channel_values reads see the delta sentinel in delta mode, which disabled goal continuation (stand_down=no_durable_end_of_turn) after durable assistant turns. Raw tuples remain for tuple-only metadata (checkpoint id, pending_writes). - Reject checkpoint_channel_mode='delta' + checkpointer in create_deerflow_agent at construction: factory-built persisted graphs bypass mode-marker injection and the fail-closed gate, reproducing silent mixed-mode state loss. Delta without persistence stays allowed. - Import the postgres saver lazily (pytest.importorskip in the fixture) so the documented default install collects the suite; add a CI job running pytest --collect-only on uv sync --group dev without extras. - Fix test_checkpointer fallback test to patch get_app_config at its use site (provider module), making it deterministic when a local config.yaml selects a persistent backend. * fix(gateway): preserve extension-owned channels in state mutations, bump config version - build_state_mutation_graph / build_checkpoint_state_mutation_accessor accept an explicit state_schema; branch and POST /state now compile the mutation graph from the thread's effective schema (graph_state_schema on the assistant graph). The base-ThreadState fallback silently discarded channels contributed by custom AgentMiddleware.state_schema on branch (data loss) and returned a false-success 200 on POST /state. - POST /state validates values keys against the mutation graph's channels and rejects unknown fields with 422 instead of ignoring them; reducer detection covers extension channels (BinaryOperatorAggregate or DeltaChannel) so Overwrite replace semantics work for middleware reducers in both modes. - Endpoint regression: custom AgentMiddleware.state_schema value survives branch, updates through POST /state, and an unknown field receives 422. - config_version 26 -> 27 for the new database.checkpoint_channel_mode (example, Helm chart values + README, support-bundle fixture), so existing installs get the outdated-config warning and make config-upgrade merges the field; covered by a test driving the real example file and the real config-upgrade script. * fix(gateway): resolve assistant schema via one boundary, copy branch reducer values with Overwrite GET /threads/{id}/state now resolves the thread's assistant_id through a single reusable boundary (thread metadata -> assistant_id -> effective graph), so channels contributed by a custom AgentMiddleware.state_schema are materialized instead of dropped by the default lead schema. POST /state uses the same boundary instead of resolving the schema ad hoc. Branch writes wrap every copied reducer channel in Overwrite (derived from the effective mutation graph: BinaryOperatorAggregate + DeltaChannel), not just messages, so already-aggregated values are never re-merged. Regression tests use a real AgentMiddleware.state_schema with a non-identity reducer in both full and delta modes: GET /state returns the extension value, POST /state replaces it, branch preserves it byte-for-byte; the unknown-field 422 is a separate assertion. * refactor(checkpoint): collapse read-path round-trips and ship dual-mode parity tests Address review round 4 on PR #4292: - Push ahistory/history limit through Pregel into checkpointer.alist (SQL LIMIT) instead of materializing all rows and breaking in Python - Fold the read-side mode-compat gate onto the returned snapshot's metadata; only writes keep the pre-write tuple fetch (fail-closed) - Cache factory-built accessor graphs per (assistant_id, mode) with factory-identity revalidation; state reads no longer build a lead agent per request - get_thread: one snapshot fetch + one raw pending_writes fetch on the resolved checkpoint (post-checkpoint __error__ writes never surface in snapshot.tasks; verified empirically) - DeerFlowClient.get_thread: single checkpointer.list walk collects pending_writes per checkpoint instead of N get_tuple calls - InMemorySaver delta-history patch: stand-down when the upstream override disappears, try/except guard, validated-version warning, guard tests - make_lead_agent mode precedence: first freeze is owned by app_config (client-supplied configurable key ignored); once frozen, injected key/app_config must match or fail closed - Rollback: lock in non-message channel restoration via fork inheritance with a dedicated reducer-channel test - Add tests/test_threads_checkpoint_mode.py and tests/test_gateway_checkpoint_mode.py referenced by AGENTS.md and the PR validation section: lifecycle parity (memory + sqlite), per-step blob-count storage guard, gateway endpoint parity Counted-saver tests pin checkpoint round-trips for aget/ahistory so these regressions cannot silently return. * fix(checkpoint): precise mode-mismatch HTTP mapping, gate E2E, and accessor resilience - threads router: map CheckpointModeMismatchError to 409 (with cause and thread id) and CheckpointModeReconfigurationError to 503 across all state endpoints instead of swallowing both into a generic 500 - gate coverage: seed a real delta checkpoint into AsyncSqliteSaver and assert aget/aupdate/ahistory fail closed; assert 409 at the HTTP boundary through the real route stack - rollback: compile the restore mutation graph with the thread's effective state schema per the build_state_mutation_graph contract - inheritance contract locks: rollback and manual compaction preserve middleware-contributed channels via checkpoint fork cloning - services: revalidate the accessor-graph cache against app_config identity so config.yaml hot-reloads never serve a stale compiled graph - services: degrade full-mode state reads to raw checkpointer reads when the agent factory is unavailable (delta gate still applies; delta mode has no fallback) - deps: override websockets==16.0 (langgraph-sdk 0.4.2's <16 pin silently downgraded 16.0 -> 15.0.1; pin is not grounded in any API incompatibility) and bump the langchain lower bound to what the lockfile actually resolves * fix(checkpoint): include anchor checkpoint in degraded history walk + cover get_thread - _RawCheckpointReadAccessor.ahistory: alist(before=...) is exclusive while pregel's get_state_history treats config.checkpoint_id as the inclusive start; fetch the anchor explicitly so both read paths paginate identically - extend the degraded-path gateway test: GET /thread returns raw values, and POST /history with before starts at the anchor checkpoint * fix(gateway): preserve degraded checkpoint timestamps * fix(gateway): harden degraded checkpoint access * fix(gateway): resolve assistants for checkpoint reads --------- Co-authored-by: Willem Jiang --- .github/workflows/backend-unit-tests.yml | 48 +- backend/AGENTS.md | 22 + backend/app/gateway/deps.py | 3 + backend/app/gateway/routers/runs.py | 20 +- backend/app/gateway/routers/thread_runs.py | 40 +- backend/app/gateway/routers/threads.py | 487 ++++++----- backend/app/gateway/services.py | 274 +++++++ .../harness/deerflow/agents/__init__.py | 11 +- .../harness/deerflow/agents/factory.py | 25 +- .../deerflow/agents/lead_agent/agent.py | 78 +- .../harness/deerflow/agents/thread_state.py | 75 +- .../harness/deerflow/checkpoint_patches.py | 99 +++ backend/packages/harness/deerflow/client.py | 112 ++- .../deerflow/config/database_config.py | 11 + .../harness/deerflow/runtime/__init__.py | 4 + .../deerflow/runtime/checkpoint_mode.py | 108 +++ .../deerflow/runtime/checkpoint_state.py | 187 +++++ .../deerflow/runtime/context_compaction.py | 71 +- .../packages/harness/deerflow/runtime/goal.py | 5 + .../harness/deerflow/runtime/runs/worker.py | 272 ++++--- backend/packages/harness/pyproject.toml | 10 +- backend/pyproject.toml | 13 + backend/tests/conftest.py | 16 + backend/tests/test_app_config_reload.py | 15 + backend/tests/test_checkpoint_mode.py | 286 +++++++ backend/tests/test_checkpoint_state.py | 306 +++++++ backend/tests/test_checkpointer.py | 4 +- backend/tests/test_client.py | 216 ++++- .../tests/test_client_langfuse_metadata.py | 1 + backend/tests/test_config_version.py | 48 ++ backend/tests/test_context_compaction.py | 264 +++--- backend/tests/test_create_deerflow_agent.py | 79 +- .../tests/test_delta_channel_checkpointers.py | 371 +++++++++ backend/tests/test_delta_channel_state.py | 184 +++++ backend/tests/test_gateway_checkpoint_mode.py | 276 +++++++ .../tests/test_gateway_config_freshness.py | 39 +- .../tests/test_gateway_run_drain_shutdown.py | 2 +- backend/tests/test_gateway_run_recovery.py | 4 +- backend/tests/test_gateway_services.py | 255 ++++++ backend/tests/test_goal_worker.py | 63 ++ .../tests/test_lead_agent_model_resolution.py | 91 ++- backend/tests/test_memory_tools.py | 2 + backend/tests/test_run_worker_rollback.py | 748 +++++++++++++---- backend/tests/test_runtime_lifecycle_e2e.py | 24 + backend/tests/test_support_bundle.py | 2 +- .../tests/test_thread_regenerate_prepare.py | 288 ++++++- backend/tests/test_thread_state_reducers.py | 14 + backend/tests/test_threads_checkpoint_mode.py | 147 ++++ backend/tests/test_threads_router.py | 756 +++++++++++++++++- backend/uv.lock | 114 ++- config.example.yaml | 3 + 51 files changed, 5818 insertions(+), 775 deletions(-) create mode 100644 backend/packages/harness/deerflow/checkpoint_patches.py create mode 100644 backend/packages/harness/deerflow/runtime/checkpoint_mode.py create mode 100644 backend/packages/harness/deerflow/runtime/checkpoint_state.py create mode 100644 backend/tests/test_checkpoint_mode.py create mode 100644 backend/tests/test_checkpoint_state.py create mode 100644 backend/tests/test_delta_channel_checkpointers.py create mode 100644 backend/tests/test_delta_channel_state.py create mode 100644 backend/tests/test_gateway_checkpoint_mode.py create mode 100644 backend/tests/test_threads_checkpoint_mode.py diff --git a/.github/workflows/backend-unit-tests.yml b/.github/workflows/backend-unit-tests.yml index 456ee5897..f6a3e3268 100644 --- a/.github/workflows/backend-unit-tests.yml +++ b/.github/workflows/backend-unit-tests.yml @@ -14,11 +14,57 @@ permissions: contents: read jobs: + default-install-collection: + # The main job installs --extra postgres; this job proves the documented + # contributor path (uv sync --group dev, no extras) still collects the + # whole suite, so optional-dependency imports must stay lazy. + if: github.event.pull_request.draft == false + runs-on: ubuntu-latest + timeout-minutes: 10 + + steps: + - name: Checkout + uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.12' + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Install backend dependencies (documented default) + working-directory: backend + run: uv sync --group dev + + - name: Collect backend tests + working-directory: backend + run: uv run pytest --collect-only -q + backend-unit-tests: if: github.event.pull_request.draft == false runs-on: ubuntu-latest timeout-minutes: 15 + services: + postgres: + image: postgres:17 + env: + POSTGRES_USER: deerflow + POSTGRES_PASSWORD: deerflow + POSTGRES_DB: deerflow_test + ports: + - 5432:5432 + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + + env: + TEST_POSTGRES_URI: postgresql://deerflow:deerflow@localhost:5432/deerflow_test?sslmode=disable + steps: - name: Checkout uses: actions/checkout@v6 @@ -33,7 +79,7 @@ jobs: - name: Install backend dependencies working-directory: backend - run: uv sync --group dev + run: uv sync --group dev --extra postgres - name: Run unit tests of backend working-directory: backend diff --git a/backend/AGENTS.md b/backend/AGENTS.md index d7abf0f79..4beacbb61 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -721,6 +721,28 @@ This invokes `alembic revision --autogenerate` against the live ORM models. Revi - `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch decision + locking - Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor) +### Checkpoint Channel Modes (`full` / `delta`) + +Checkpointer storage runs in one of two channel modes, selected by `checkpoint_channel_mode` in `config.yaml` (default `full`). `delta` mode adopts LangGraph 1.2's `DeltaChannel` for `messages`: checkpoints store a sentinel + per-step writes instead of the full message list, so storage/serde grows O(N) instead of O(N²) in turns. All checkpointer backends (memory/sqlite/postgres) serve both modes unchanged — the semantics live in the compiled graph's channel table, not in the saver. + +**Mode is process-frozen and restart-required.** `make_lead_agent` freezes the resolved mode (`runtime/checkpoint_mode.py::freeze_checkpoint_channel_mode`) before compiling the graph with the mode-matched schema (`agents/thread_state.py::get_thread_state_schema`, plus `adapt_state_schema_for_mode` / `normalize_middleware_state_schemas` for middleware state). A second, different mode in the same process raises `CheckpointModeReconfigurationError`. To switch: edit config, restart. + +**Compatibility is asymmetric and fail-closed.** Every checkpoint written in delta mode carries metadata marker `deerflow_checkpoint_channel_mode: "delta"` (injected via `inject_checkpoint_mode`; absence of marker = full, so pre-feature checkpoints need no migration). Before any state read/write, `ensure_checkpoint_mode_compatible` rejects a full-mode process opening a delta thread with `CheckpointModeMismatchError` (surfaced as HTTP 409 with the cause and thread id by the threads router; `CheckpointModeReconfigurationError` maps to 503) — a full-mode raw read of a delta blob would silently return empty/partial `messages`. The reverse direction is allowed: delta-mode processes read full checkpoints transparently (old full checkpoints seed the delta channel), so full → delta is the smooth migration path; delta → full requires materializing/converting the data first. Detection also honors upstream's `counters_since_delta_snapshot.messages` metadata, and an explicit config marker takes precedence over any ambient context value. + +**Never bypass `CheckpointStateAccessor` (`runtime/checkpoint_state.py`) for thread-state access.** It is the single choke point binding graph + checkpointer + mode: it injects the mode marker into configs, runs the compatibility check before every `get`/`update`/`history`, and returns materialized state (delta checkpoints lack `channel_values.messages` — raw `get_tuple` reads see a sentinel). Gateway `services.py` builds and passes the accessor; thread-owned reads (state/history/regeneration) must use `build_thread_checkpoint_state_accessor` so the recorded assistant's middleware schema materializes every channel. `history(limit)` semantics: `0` means zero items (explicit empty), `None` means unlimited — do not pass `limit=0` through to `graph.get_state_history`. Assistant metadata lookup is fail-closed for mutation accessors so a store outage cannot silently select the default schema and discard extension channels. In `full` mode the read path degrades to a raw checkpointer read (`_RawCheckpointReadAccessor`) when the agent factory cannot build the graph (bad model config, MCP outage) — full checkpoints carry complete `channel_values`, so reads don't need the graph; degraded snapshots take `created_at` from the standard checkpoint `ts` field, falling back to metadata only for compatibility. The delta gate still applies on the degraded path; `next`/`tasks` degrade to empty and thread status falls back to the stored status because task presence is not derivable, while delta mode has no fallback (materialization needs the channel table). + +**Wholesale message replacement uses a state-only mutation graph + `Overwrite`.** `update_state` values pass through channel reducers (`add_messages` merge in full, append in delta), so replacing `messages` wholesale (run rollback, context compaction) requires `{"messages": Overwrite([...])}`. The write goes through `build_state_mutation_graph(as_node, mode, state_schema)` — when the write carries materialized state, `state_schema` MUST be the thread's effective schema (`graph_state_schema(assistant_graph)`), because the base-ThreadState fallback silently discards written channels contributed by custom `AgentMiddleware.state_schema`. Channels absent from the write are unaffected: forked checkpoints clone the parent's channel blobs, so middleware channels survive rollback/compaction regardless of schema (locked by `test_rollback_preserves_middleware_contributed_channels` and `test_compact_thread_context_preserves_middleware_contributed_channels`) — a compiled graph with one no-op node (entry = finish) whose checkpoint machinery (channels/versions/metadata) is identical to the agent graph's but schedules no pending tasks, so the restored/compacted head stays idle instead of re-triggering the agent. Never hand-write checkpoints via `checkpointer.aput` for this; raw writers elsewhere must preserve checkpoint parentage — severed ancestry breaks delta replay (see `runtime/runs/worker.py` writer parenting and `checkpoint_patches.py`). + +**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes pre-run state (messages via accessor, raw `pending_writes` via `aget_tuple`) into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. Cancel-with-rollback then forks from the pre-run checkpoint via the mutation graph and replays the captured pending writes. + +**Where things live**: +- `runtime/checkpoint_mode.py` — mode freeze, marker injection, delta detection, compatibility gate, both error types +- `runtime/checkpoint_state.py` — `CheckpointStateAccessor`, `build_state_mutation_graph`, `RollbackPoint` +- `checkpoint_patches.py` (package root) — saver patches: delta-history folding for `InMemorySaver` (delegating to the base walk), stable message IDs across materialization, upstream first-write drop fix +- `agents/thread_state.py` — `ThreadState`/`DeltaThreadState`, `DELTA_MESSAGES_FIELD` (`DeltaChannel` with `snapshot_frequency=1000`), schema adaptation helpers +- `runtime/context_compaction.py` — compaction via accessor + mutation graph (reference consumer) +- Tests: `tests/test_checkpoint_mode.py` (freeze/detect/gate), `tests/test_checkpoint_state.py` (accessor/mutation graph), `tests/test_delta_channel_checkpointers.py` (saver parity), `tests/test_threads_checkpoint_mode.py`, `tests/test_gateway_checkpoint_mode.py` (dual-mode e2e parity), `tests/test_context_compaction.py` (mutation-graph write, no scheduling), `tests/test_run_worker_rollback.py` + ### Terminal Workbench / TUI (`packages/harness/deerflow/tui/`) A terminal-native UI over the embedded harness, exposed as the `deerflow` console script (`[project.scripts]` in `packages/harness/pyproject.toml`). It is a UI shell over `DeerFlowClient` and does **not** fork agent behavior. `textual` is an optional dependency (`deerflow-harness[tui]`; also in the backend dev group); the console script degrades to headless help when it is absent. Full guide: [docs/TUI.md](docs/TUI.md). diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index faf5634d8..decac902c 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -261,6 +261,7 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen """ from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config from deerflow.runtime import make_store, make_stream_bridge + from deerflow.runtime.checkpoint_mode import freeze_checkpoint_channel_mode from deerflow.runtime.checkpointer.async_provider import make_checkpointer from deerflow.runtime.events.store import make_run_event_store @@ -272,6 +273,7 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen async with AsyncExitStack() as stack: config = startup_config + app.state.checkpoint_channel_mode = freeze_checkpoint_channel_mode(config.database.checkpoint_channel_mode) app.state.stream_bridge = await stack.enter_async_context(make_stream_bridge(config)) @@ -433,6 +435,7 @@ def get_run_context(request: Request) -> RunContext: store=get_store(request), event_store=get_run_event_store(request), run_events_config=getattr(request.app.state, "run_events_config", None), + checkpoint_channel_mode=getattr(request.app.state, "checkpoint_channel_mode", "full"), thread_store=get_thread_store(request), app_config=get_config(), on_run_completed=getattr(request.app.state, "scheduled_task_service", None).handle_run_completion if getattr(request.app.state, "scheduled_task_service", None) is not None else None, diff --git a/backend/app/gateway/routers/runs.py b/backend/app/gateway/routers/runs.py index 9cf4e701b..23195d301 100644 --- a/backend/app/gateway/routers/runs.py +++ b/backend/app/gateway/routers/runs.py @@ -14,10 +14,10 @@ from fastapi import APIRouter, HTTPException, Query, Request from fastapi.responses import StreamingResponse from app.gateway.authz import require_permission -from app.gateway.deps import get_checkpointer, get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge +from app.gateway.deps import get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge from app.gateway.pagination import trim_run_message_page from app.gateway.routers.thread_runs import RunCreateRequest -from app.gateway.services import sse_consumer, start_run, wait_for_run_completion +from app.gateway.services import build_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion from deerflow.runtime import serialize_channel_values_for_api logger = logging.getLogger(__name__) @@ -75,14 +75,16 @@ async def stateless_wait(body: RunCreateRequest, request: Request) -> dict: completed = await wait_for_run_completion(bridge, record, request, run_mgr) if completed: - checkpointer = get_checkpointer(request) - config = {"configurable": {"thread_id": thread_id}} try: - checkpoint_tuple = await checkpointer.aget_tuple(config) - if checkpoint_tuple is not None: - checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} - channel_values = checkpoint.get("channel_values", {}) - return serialize_channel_values_for_api(channel_values) + accessor, config = build_checkpoint_state_accessor( + request, + thread_id=thread_id, + assistant_id=body.assistant_id, + ) + snapshot = await accessor.aget(config) + snapshot_config = snapshot.config or {} + if snapshot_config.get("configurable", {}).get("checkpoint_id"): + return serialize_channel_values_for_api(snapshot.values) except Exception: logger.exception("Failed to fetch final state for run %s", record.run_id) diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index e0a8c1096..e044362c6 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -23,9 +23,9 @@ from langchain_core.messages import BaseMessage from pydantic import BaseModel, Field from app.gateway.authz import require_permission -from app.gateway.deps import get_checkpointer, get_current_user, get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge +from app.gateway.deps import get_current_user, get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge from app.gateway.pagination import trim_run_message_page -from app.gateway.services import sse_consumer, start_run, wait_for_run_completion +from app.gateway.services import build_checkpoint_state_accessor, build_thread_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion from deerflow.runtime import CancelOutcome, RunRecord, RunStatus, serialize_channel_values_for_api from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, get_original_user_content_text, message_to_text from deerflow.workspace_changes import get_workspace_changes_response @@ -293,10 +293,9 @@ def _is_middleware_message_row(row: dict[str, Any]) -> bool: return str((row.get("metadata") or {}).get("caller", "")).startswith("middleware:") -def _checkpoint_messages(checkpoint_tuple: Any) -> list[Any]: - checkpoint = getattr(checkpoint_tuple, "checkpoint", None) or {} - channel_values = checkpoint.get("channel_values", {}) if isinstance(checkpoint, dict) else {} - messages = channel_values.get("messages", []) if isinstance(channel_values, dict) else [] +def _checkpoint_messages(snapshot: Any) -> list[Any]: + values = getattr(snapshot, "values", None) or {} + messages = values.get("messages", []) if isinstance(values, dict) else [] return messages if isinstance(messages, list) else [] @@ -387,10 +386,9 @@ async def _find_target_run_id(thread_id: str, message_id: str, target_message: A async def _find_base_checkpoint_before_human(thread_id: str, human_message_id: str, request: Request) -> Any: - checkpointer = get_checkpointer(request) - base_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + accessor, base_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id) try: - raw_checkpoints = [item async for item in checkpointer.alist(base_config, limit=REGENERATE_HISTORY_RAW_SCAN_LIMIT)] + raw_checkpoints = await accessor.ahistory(base_config, limit=REGENERATE_HISTORY_RAW_SCAN_LIMIT) checkpoints = [item for item in raw_checkpoints if not _is_duration_only_checkpoint(item)] except Exception as exc: logger.exception("Failed to list checkpoints for regenerate thread %s", thread_id) @@ -424,14 +422,14 @@ async def _find_base_checkpoint_before_human(thread_id: str, human_message_id: s async def _prepare_regenerate_payload(thread_id: str, message_id: str, request: Request) -> RegeneratePrepareResponse: - checkpointer = get_checkpointer(request) - latest_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + accessor, latest_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id) try: - latest_checkpoint = await checkpointer.aget_tuple(latest_config) + latest_checkpoint = await accessor.aget(latest_config) except Exception as exc: logger.exception("Failed to read latest checkpoint for regenerate thread %s", thread_id) raise HTTPException(status_code=500, detail="Failed to read latest checkpoint") from exc - if latest_checkpoint is None: + latest_checkpoint_id = _checkpoint_configurable(latest_checkpoint).get("checkpoint_id") + if not latest_checkpoint_id: raise HTTPException(status_code=404, detail=f"Thread {thread_id} has no checkpoint") messages = _checkpoint_messages(latest_checkpoint) @@ -534,14 +532,16 @@ async def wait_run(thread_id: str, body: RunCreateRequest, request: Request) -> completed = await wait_for_run_completion(bridge, record, request, run_mgr) if completed: - checkpointer = get_checkpointer(request) - config = {"configurable": {"thread_id": thread_id}} try: - checkpoint_tuple = await checkpointer.aget_tuple(config) - if checkpoint_tuple is not None: - checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} - channel_values = checkpoint.get("channel_values", {}) - return serialize_channel_values_for_api(channel_values) + accessor, config = build_checkpoint_state_accessor( + request, + thread_id=thread_id, + assistant_id=body.assistant_id, + ) + snapshot = await accessor.aget(config) + snapshot_config = snapshot.config or {} + if snapshot_config.get("configurable", {}).get("checkpoint_id"): + return serialize_channel_values_for_api(snapshot.values) except Exception: logger.exception("Failed to fetch final state for run %s", record.run_id) diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index 0ca523f3e..2567a526c 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -12,7 +12,6 @@ matching the LangGraph Platform wire format expected by the from __future__ import annotations -import copy import logging import shutil import uuid @@ -20,17 +19,27 @@ from pathlib import Path from typing import Any from fastapi import APIRouter, BackgroundTasks, HTTPException, Request -from langgraph.checkpoint.base import empty_checkpoint, uuid6 +from langgraph.checkpoint.base import empty_checkpoint +from langgraph.types import Overwrite from pydantic import BaseModel, Field, field_validator from sqlalchemy.exc import IntegrityError from app.gateway.authz import require_permission from app.gateway.deps import get_checkpointer, get_run_manager from app.gateway.internal_auth import get_trusted_internal_owner_user_id +from app.gateway.services import ( + build_checkpoint_state_accessor, + build_checkpoint_state_mutation_accessor, + build_thread_checkpoint_state_accessor, + build_thread_checkpoint_state_mutation_accessor, +) from app.gateway.utils import sanitize_log_param +from deerflow.agents.thread_state import THREAD_STATE_REDUCER_FIELDS from deerflow.config.paths import Paths, get_paths from deerflow.config.summarization_config import ContextSize from deerflow.runtime import serialize_channel_values_for_api +from deerflow.runtime.checkpoint_mode import CheckpointModeMismatchError, CheckpointModeReconfigurationError +from deerflow.runtime.checkpoint_state import graph_reducer_channels, graph_state_schema, graph_writable_channels from deerflow.runtime.context_compaction import ( ContextCompactionDisabled, ContextCompactionFailed, @@ -53,6 +62,22 @@ from deerflow.utils.time import coerce_iso, now_iso logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/threads", tags=["threads"]) +_CHECKPOINT_MODE_ERRORS = (CheckpointModeMismatchError, CheckpointModeReconfigurationError) + + +def _checkpoint_mode_http_error(exc: Exception, thread_id: str) -> HTTPException: + """Map checkpoint-mode guard failures to precise HTTP statuses. + + A mismatch means the thread's persisted checkpoints conflict with the + process's frozen mode (operator-actionable, 409); a reconfiguration means + the process itself is mid mode-flip (transient, 503). Both must surface + their message — a generic 500 would force operators to grep logs to + discover the root cause after a mode flip. + """ + if isinstance(exc, CheckpointModeMismatchError): + return HTTPException(status_code=409, detail=f"Thread {thread_id}: {exc}") + return HTTPException(status_code=503, detail=str(exc)) + # Metadata keys that the server controls; clients are not allowed to set # them. Pydantic ``@field_validator("metadata")`` strips them on every @@ -107,15 +132,14 @@ def _is_branch_assistant_message(message: Any) -> bool: return _message_type(message) == "ai" -def _checkpoint_messages(checkpoint_tuple: Any) -> list[Any]: - checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} - channel_values = checkpoint.get("channel_values", {}) or {} - messages = channel_values.get("messages") or [] +def _checkpoint_messages(snapshot: Any) -> list[Any]: + values = getattr(snapshot, "values", None) or {} + messages = values.get("messages") if isinstance(values, dict) else None return list(messages) if isinstance(messages, list) else [] -def _checkpoint_id(checkpoint_tuple: Any) -> str | None: - config = getattr(checkpoint_tuple, "config", {}) or {} +def _checkpoint_id(snapshot: Any) -> str | None: + config = getattr(snapshot, "config", {}) or {} raw = config.get("configurable", {}).get("checkpoint_id") return raw if isinstance(raw, str) and raw else None @@ -134,37 +158,38 @@ def _matches_branch_target(messages: list[Any], target_message_ids: set[str]) -> return not any(_is_branch_visible_message(message) for message in messages[target_end_index + 1 :]) -async def _find_branch_checkpoint(checkpointer: Any, thread_id: str, target_message_ids: set[str]) -> Any: - config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} +async def _find_branch_checkpoint( + accessor: Any, + config: dict[str, Any], + target_message_ids: set[str], +) -> Any: try: - async for checkpoint_tuple in checkpointer.alist(config, limit=_BRANCH_HISTORY_SCAN_LIMIT): - if _matches_branch_target(_checkpoint_messages(checkpoint_tuple), target_message_ids): - return checkpoint_tuple + for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_SCAN_LIMIT): + if _matches_branch_target(_checkpoint_messages(snapshot), target_message_ids): + return snapshot + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, config.get("configurable", {}).get("thread_id", "")) from exc except Exception: + thread_id = config.get("configurable", {}).get("thread_id", "") logger.exception("Failed to scan branch checkpoint history for thread %s", sanitize_log_param(thread_id)) raise HTTPException(status_code=500, detail="Failed to find branch checkpoint") raise HTTPException(status_code=409, detail="This turn can no longer be branched from.") -async def _branch_targets_latest_turn(checkpointer: Any, thread_id: str, target_message_ids: set[str]) -> bool: - """Return True when the target turn is the final visible turn in the current state. - - ``alist`` yields newest-first; we take the newest checkpoint that actually holds - messages (thread creation writes an empty checkpoint that must be skipped) and - reuse ``_matches_branch_target`` to check the target turn is its tail. Used to - decide whether cloning the (uncheckpointed) workspace onto a branch is safe: only - a branch from the latest turn shares the current workspace timeline. On any lookup - failure we fail closed (treat as historical) so a branch from an older turn never - inherits a later timeline's workspace files. - """ - config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} +async def _branch_targets_latest_turn( + accessor: Any, + config: dict[str, Any], + target_message_ids: set[str], +) -> bool: + """Return whether the target turn is the final visible turn.""" try: - async for checkpoint_tuple in checkpointer.alist(config, limit=_BRANCH_HISTORY_SCAN_LIMIT): - messages = _checkpoint_messages(checkpoint_tuple) + for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_SCAN_LIMIT): + messages = _checkpoint_messages(snapshot) if not messages: continue return _matches_branch_target(messages, target_message_ids) except Exception: + thread_id = config.get("configurable", {}).get("thread_id", "") logger.warning( "Failed to resolve latest turn for thread %s; treating branch as historical", sanitize_log_param(thread_id), @@ -416,19 +441,38 @@ def _delete_thread_data(thread_id: str, paths: Paths | None = None, *, user_id: return ThreadDeleteResponse(success=True, message=f"Deleted local thread data for {thread_id}") -def _derive_thread_status(checkpoint_tuple) -> str: - """Derive thread status from checkpoint metadata.""" - if checkpoint_tuple is None: - return "idle" - pending_writes = getattr(checkpoint_tuple, "pending_writes", None) or [] +async def _fetch_raw_pending_writes(checkpointer: Any, config: dict[str, Any]) -> list[Any]: + """Fetch pending writes attached to a specific checkpoint. - # Check for error in pending writes - for pw in pending_writes: - if len(pw) >= 2 and pw[1] == "__error__": + Snapshot ``tasks`` only reflect writes that were pending while a task was + still scheduled; writes attached to the latest checkpoint afterwards + (rollback reattachment, worker error fallback) never surface there, so the + status derivation needs one raw tuple fetch on the resolved checkpoint. + """ + raw_tuple = await checkpointer.aget_tuple(config) + if raw_tuple is None: + return [] + return list(getattr(raw_tuple, "pending_writes", ()) or ()) + + +def _derive_thread_status(snapshot: Any, pending_writes: list[Any], *, fallback_status: str = "idle") -> str: + """Derive thread status from the materialized snapshot plus the raw + pending writes attached to the resolved checkpoint.""" + if snapshot is None: + return "idle" + + for write in pending_writes: + if isinstance(write, (list, tuple)) and len(write) >= 2 and write[1] == "__error__": return "error" - # Check for pending next tasks (indicates interrupt) - tasks = getattr(checkpoint_tuple, "tasks", None) + tasks = getattr(snapshot, "tasks", None) or () + for task in tasks: + if getattr(task, "error", None) is not None: + return "error" + + if not getattr(snapshot, "tasks_known", True): + return fallback_status + if tasks: return "interrupted" @@ -632,7 +676,6 @@ async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Requ """Create a new main-thread branch from a completed assistant turn.""" from app.gateway.deps import get_thread_store - checkpointer = get_checkpointer(request) thread_store = get_thread_store(request) source_record = await thread_store.get(thread_id) @@ -642,10 +685,15 @@ async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Requ source_metadata = source_record.get("metadata") or {} if source_metadata.get(_SIDECAR_METADATA_KEY) is True: raise HTTPException(status_code=409, detail="Branching is only available in the main conversation.") + source_accessor, source_config = build_checkpoint_state_accessor( + request, + thread_id=thread_id, + assistant_id=source_record.get("assistant_id"), + ) target_message_ids = {body.message_id, *body.message_ids} - checkpoint_tuple = await _find_branch_checkpoint(checkpointer, thread_id, target_message_ids) - parent_checkpoint_id = _checkpoint_id(checkpoint_tuple) + snapshot = await _find_branch_checkpoint(source_accessor, source_config, target_message_ids) + parent_checkpoint_id = _checkpoint_id(snapshot) if not parent_checkpoint_id: raise HTTPException(status_code=409, detail="This turn can no longer be branched from.") @@ -654,7 +702,7 @@ async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Requ # after that turn (message history rolls back, workspace would not). Restrict the # best-effort clone to branches taken from the latest turn so history and workspace # stay consistent. - branch_from_latest_turn = await _branch_targets_latest_turn(checkpointer, thread_id, target_message_ids) + branch_from_latest_turn = await _branch_targets_latest_turn(source_accessor, source_config, target_message_ids) new_thread_id = str(uuid.uuid4()) now = now_iso() @@ -673,22 +721,37 @@ async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Requ thread_owner_user_id = get_trusted_internal_owner_user_id(request) thread_owner_kwargs = {"user_id": thread_owner_user_id} if thread_owner_user_id else {} - checkpoint = copy.deepcopy(getattr(checkpoint_tuple, "checkpoint", {}) or {}) - metadata = copy.deepcopy(getattr(checkpoint_tuple, "metadata", {}) or {}) - checkpoint["id"] = str(uuid6()) - metadata.update( + # Copy materialized values with replace semantics: reducer channels must + # not re-merge an already-aggregated value, so every copied reducer value + # is wrapped in Overwrite (not just messages). + branch_accessor, new_config = build_checkpoint_state_mutation_accessor( + request, + thread_id=new_thread_id, + as_node="branch", + # The branch write carries the full materialized snapshot; use the + # source assistant's effective schema so extension middleware channels + # survive instead of being silently discarded as unknown channels. + state_schema=graph_state_schema(getattr(source_accessor, "graph", None)), + ) + branch_reducer_fields = graph_reducer_channels(getattr(branch_accessor, "graph", None)) + if branch_reducer_fields is None: + branch_reducer_fields = THREAD_STATE_REDUCER_FIELDS + branch_values = {} + for key, value in dict(snapshot.values).items(): + if key in branch_reducer_fields: + branch_values[key] = Overwrite(list(value) if key == "messages" and isinstance(value, list) else value) + else: + branch_values[key] = value + new_config.setdefault("metadata", {}).update( { - "source": "branch", - "updated_at": now, - "created_at": now, **branch_metadata, + "source": "branch", } ) - - write_config = {"configurable": {"thread_id": new_thread_id, "checkpoint_ns": ""}} - new_versions = dict(checkpoint.get("channel_versions", {}) or {}) try: - await checkpointer.aput(write_config, checkpoint, metadata, new_versions) + await branch_accessor.aupdate(new_config, branch_values, as_node="branch") + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, new_thread_id) from exc except Exception: logger.exception("Failed to write branch checkpoint for thread %s", sanitize_log_param(new_thread_id)) raise HTTPException(status_code=500, detail="Failed to create branch") from None @@ -787,49 +850,45 @@ async def patch_thread(thread_id: str, body: ThreadPatchRequest, request: Reques @router.get("/{thread_id}", response_model=ThreadResponse) @require_permission("threads", "read", owner_check=True) async def get_thread(thread_id: str, request: Request) -> ThreadResponse: - """Get thread info. - - Reads metadata from the ThreadMetaStore and derives the accurate - execution status from the checkpointer. Falls back to the checkpointer - alone for threads that pre-date ThreadMetaStore adoption (backward compat). - """ + """Get thread info from metadata plus the graph's materialized state.""" from app.gateway.deps import get_thread_store thread_store = get_thread_store(request) checkpointer = get_checkpointer(request) - record: dict | None = await thread_store.get(thread_id) - - # Derive accurate status from the checkpointer - config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} try: - checkpoint_tuple = await checkpointer.aget_tuple(config) + accessor, config = build_checkpoint_state_accessor( + request, + thread_id=thread_id, + assistant_id=record.get("assistant_id") if record is not None else None, + ) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc + + try: + snapshot = await accessor.aget(config) + checkpoint_id = (snapshot.config or {}).get("configurable", {}).get("checkpoint_id") + pending_writes = await _fetch_raw_pending_writes(checkpointer, snapshot.config) if checkpoint_id else [] + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc except Exception: logger.exception("Failed to get checkpoint for thread %s", sanitize_log_param(thread_id)) raise HTTPException(status_code=500, detail="Failed to get thread") - if record is None and checkpoint_tuple is None: + if record is None and not checkpoint_id: raise HTTPException(status_code=404, detail=f"Thread {thread_id} not found") - # If the thread exists in the checkpointer but not in thread_meta (e.g. - # legacy data created before thread_meta adoption), synthesize a minimal - # record from the checkpoint metadata. - if record is None and checkpoint_tuple is not None: - ckpt_meta = getattr(checkpoint_tuple, "metadata", {}) or {} + metadata = snapshot.metadata or {} + if record is None: record = { "thread_id": thread_id, "status": "idle", - "created_at": coerce_iso(ckpt_meta.get("created_at", "")), - "updated_at": coerce_iso(ckpt_meta.get("updated_at", ckpt_meta.get("created_at", ""))), - "metadata": {k: v for k, v in ckpt_meta.items() if k not in ("created_at", "updated_at", "step", "source", "writes", "parents")}, + "created_at": coerce_iso(snapshot.created_at or metadata.get("created_at", "")), + "updated_at": coerce_iso(metadata.get("updated_at", snapshot.created_at or metadata.get("created_at", ""))), + "metadata": {key: value for key, value in metadata.items() if key not in ("created_at", "updated_at", "step", "source", "writes", "parents")}, } - - if record is None: - raise HTTPException(status_code=404, detail=f"Thread {thread_id} not found") - - status = _derive_thread_status(checkpoint_tuple) if checkpoint_tuple is not None else record.get("status", "idle") - checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} if checkpoint_tuple is not None else {} - channel_values = checkpoint.get("channel_values", {}) + stored_status = record.get("status", "idle") + status = _derive_thread_status(snapshot, pending_writes, fallback_status=stored_status) if checkpoint_id else stored_status return ThreadResponse( thread_id=thread_id, @@ -837,7 +896,7 @@ async def get_thread(thread_id: str, request: Request) -> ThreadResponse: created_at=coerce_iso(record.get("created_at", "")), updated_at=coerce_iso(record.get("updated_at", "")), metadata=record.get("metadata", {}), - values=serialize_channel_values_for_api(channel_values), + values=serialize_channel_values_for_api(snapshot.values), ) @@ -910,20 +969,33 @@ def _thread_compact_response(result: ThreadCompactionResult) -> ThreadCompactRes async def compact_thread(thread_id: str, body: ThreadCompactRequest, request: Request) -> ThreadCompactResponse: """Manually summarize old thread context while preserving the visible history.""" run_manager = get_run_manager(request) - checkpointer = get_checkpointer(request) + # Compaction writes only base-schema channels (messages + summary_text); + # every other channel — including middleware-contributed ones — is carried + # forward by checkpoint fork inheritance, so the base-schema mutation + # graph is sufficient (and avoids building the full lead graph per call). + try: + accessor, _ = build_checkpoint_state_mutation_accessor( + request, + thread_id=thread_id, + as_node="manual_compaction", + ) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc keep = body.keep.to_tuple() if body.keep is not None else None try: async with goal_thread_lock(thread_id): if await run_manager.has_inflight(thread_id): raise HTTPException(status_code=409, detail="Thread has a run in flight. Compact after the run finishes.") result = await compact_thread_context( - checkpointer, + accessor, thread_id, keep=keep, force=body.force, user_id=get_effective_user_id(), agent_name=body.agent_name, ) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc except ContextCompactionDisabled: raise HTTPException(status_code=409, detail="Context compaction is disabled.") from None except ContextCompactionFailed: @@ -942,51 +1014,41 @@ async def compact_thread(thread_id: str, body: ThreadCompactRequest, request: Re @router.get("/{thread_id}/state", response_model=ThreadStateResponse) @require_permission("threads", "read", owner_check=True) async def get_thread_state(thread_id: str, request: Request) -> ThreadStateResponse: - """Get the latest state snapshot for a thread. - - Channel values are serialized to ensure LangChain message objects - are converted to JSON-safe dicts. - """ - checkpointer = get_checkpointer(request) - - config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + """Get the latest materialized graph state for a thread.""" + # Resolve through the thread's assistant so custom middleware channels + # appear in the response instead of being dropped by the default schema. try: - checkpoint_tuple = await checkpointer.aget_tuple(config) + accessor, config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc + try: + snapshot = await accessor.aget(config) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc except Exception: logger.exception("Failed to get state for thread %s", sanitize_log_param(thread_id)) raise HTTPException(status_code=500, detail="Failed to get thread state") - if checkpoint_tuple is None: + snapshot_config = snapshot.config or {} + checkpoint_id = snapshot_config.get("configurable", {}).get("checkpoint_id") + if not checkpoint_id: raise HTTPException(status_code=404, detail=f"Thread {thread_id} not found") - checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} - metadata = getattr(checkpoint_tuple, "metadata", {}) or {} - checkpoint_id = None - ckpt_config = getattr(checkpoint_tuple, "config", {}) - if ckpt_config: - checkpoint_id = ckpt_config.get("configurable", {}).get("checkpoint_id") - - channel_values = checkpoint.get("channel_values", {}) - - parent_config = getattr(checkpoint_tuple, "parent_config", None) - parent_checkpoint_id = None - if parent_config: - parent_checkpoint_id = parent_config.get("configurable", {}).get("checkpoint_id") - - tasks_raw = getattr(checkpoint_tuple, "tasks", []) or [] - next_tasks = [t.name for t in tasks_raw if hasattr(t, "name")] - tasks = [{"id": getattr(t, "id", ""), "name": getattr(t, "name", "")} for t in tasks_raw] - - values = serialize_channel_values_for_api(channel_values) + parent_config = snapshot.parent_config or {} + parent_checkpoint_id = parent_config.get("configurable", {}).get("checkpoint_id") + metadata = snapshot.metadata or {} + created_at = snapshot.created_at or metadata.get("created_at", "") + tasks_raw = snapshot.tasks or () + tasks = [{"id": getattr(task, "id", ""), "name": getattr(task, "name", "")} for task in tasks_raw] return ThreadStateResponse( - values=values, - next=next_tasks, + values=serialize_channel_values_for_api(snapshot.values), + next=list(snapshot.next or ()), metadata=metadata, - checkpoint={"id": checkpoint_id, "ts": coerce_iso(metadata.get("created_at", ""))}, + checkpoint={"id": checkpoint_id, "ts": coerce_iso(created_at)}, checkpoint_id=checkpoint_id, parent_checkpoint_id=parent_checkpoint_id, - created_at=coerce_iso(metadata.get("created_at", "")), + created_at=coerce_iso(created_at), tasks=tasks, ) @@ -994,102 +1056,89 @@ async def get_thread_state(thread_id: str, request: Request) -> ThreadStateRespo @router.post("/{thread_id}/state", response_model=ThreadStateResponse) @require_permission("threads", "write", owner_check=True, require_existing=True) async def update_thread_state(thread_id: str, body: ThreadStateUpdateRequest, request: Request) -> ThreadStateResponse: - """Update thread state (e.g. for human-in-the-loop resume or title rename). - - Writes a new checkpoint that merges *body.values* into the latest - channel values, then syncs any updated ``title`` field through the - ThreadMetaStore abstraction so that ``/threads/search`` reflects the - change immediately in both sqlite and memory backends. - """ + """Replace selected thread-state fields through the materialized graph.""" from app.gateway.deps import get_thread_store - checkpointer = get_checkpointer(request) thread_store = get_thread_store(request) - - # checkpoint_ns must be present in the config for aput — default to "" - # (the root graph namespace). checkpoint_id is optional; omitting it - # fetches the latest checkpoint for the thread. - read_config: dict[str, Any] = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": "", + if body.checkpoint_id is not None: + if not body.checkpoint_id: + raise HTTPException(status_code=404, detail="Checkpoint not found") + selected_config = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + "checkpoint_id": body.checkpoint_id, + } } - } - if body.checkpoint_id: - read_config["configurable"]["checkpoint_id"] = body.checkpoint_id + try: + checkpoint_tuple = await get_checkpointer(request).aget_tuple(selected_config) + except Exception: + logger.exception("Failed to get state for thread %s", sanitize_log_param(thread_id)) + raise HTTPException(status_code=500, detail="Failed to get thread state") + if checkpoint_tuple is None: + raise HTTPException(status_code=404, detail=f"Checkpoint {body.checkpoint_id} not found") + mutation_node = body.as_node or "manual_state_update" + # Resolve through the shared boundary (thread metadata -> assistant_id -> + # effective schema) so extension middleware channels stay writable. + accessor, read_config = await build_thread_checkpoint_state_mutation_accessor( + request, + thread_id=thread_id, + as_node=mutation_node, + checkpoint_id=body.checkpoint_id, + ) + values = dict(body.values or {}) + writable_channels = graph_writable_channels(getattr(accessor, "graph", None)) + if writable_channels is not None: + unknown_fields = sorted(set(values) - writable_channels) + if unknown_fields: + raise HTTPException( + status_code=422, + detail=f"Unknown thread-state field(s): {', '.join(unknown_fields)}", + ) + reducer_fields = graph_reducer_channels(getattr(accessor, "graph", None)) + if reducer_fields is None: + reducer_fields = THREAD_STATE_REDUCER_FIELDS + updates = {key: Overwrite(value) if key in reducer_fields else value for key, value in values.items()} try: - checkpoint_tuple = await checkpointer.aget_tuple(read_config) - except Exception: - logger.exception("Failed to get state for thread %s", sanitize_log_param(thread_id)) - raise HTTPException(status_code=500, detail="Failed to get thread state") - - if checkpoint_tuple is None: - raise HTTPException(status_code=404, detail=f"Thread {thread_id} not found") - - # Work on mutable copies so we don't accidentally mutate cached objects. - checkpoint: dict[str, Any] = dict(getattr(checkpoint_tuple, "checkpoint", {}) or {}) - metadata: dict[str, Any] = dict(getattr(checkpoint_tuple, "metadata", {}) or {}) - channel_values: dict[str, Any] = dict(checkpoint.get("channel_values", {})) - - if body.values: - channel_values.update(body.values) - - checkpoint["channel_values"] = channel_values - metadata["updated_at"] = now_iso() - - if body.as_node: - metadata["source"] = "update" - metadata["step"] = metadata.get("step", 0) + 1 - metadata["writes"] = {body.as_node: body.values} - - # Assign a new checkpoint ID so aput performs an INSERT rather than an - # in-place REPLACE of the existing row. Use uuid6 (time-ordered) rather - # than uuid4 (random) so the new ID is always lexicographically greater - # than the previous one — LangGraph's checkpointers determine the "latest" - # checkpoint by max(checkpoint_ids) string order, matching the uuid6 epoch. - checkpoint["id"] = str(uuid6()) - - # aput requires checkpoint_ns in the config — use the same config used for the - # read (which always includes checkpoint_ns=""). The fresh checkpoint ID is - # assigned above via checkpoint["id"]; keep checkpoint_id out of the config so - # the write is keyed by the new checkpoint payload rather than the prior read. - # All supported savers (InMemorySaver, AsyncSqliteSaver, AsyncPostgresSaver) - # persist and echo back checkpoint["id"] verbatim — none mint their own — so - # the new_config below carries the uuid6 we assigned here. (Regression-locked - # by test_update_thread_state_inserts_new_checkpoint_each_call.) - write_config: dict[str, Any] = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": "", - } - } - try: - new_config = await checkpointer.aput(write_config, checkpoint, metadata, {}) + updated_config = await accessor.aupdate( + read_config, + updates, + as_node=mutation_node, + ) + snapshot = await accessor.aget(updated_config) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc except Exception: logger.exception("Failed to update state for thread %s", sanitize_log_param(thread_id)) raise HTTPException(status_code=500, detail="Failed to update thread state") - new_checkpoint_id: str | None = None - if isinstance(new_config, dict): - new_checkpoint_id = new_config.get("configurable", {}).get("checkpoint_id") - - # Sync title changes through the ThreadMetaStore abstraction so /threads/search - # reflects them immediately in both sqlite and memory backends. if thread_store and body.values and "title" in body.values: new_title = body.values["title"] - if new_title: # Skip empty strings and None + if new_title: try: await thread_store.update_display_name(thread_id, new_title) except Exception: logger.debug("Failed to sync title to thread_meta for %s (non-fatal)", sanitize_log_param(thread_id)) + snapshot_config = snapshot.config or {} + checkpoint_id = snapshot_config.get("configurable", {}).get("checkpoint_id") + parent_config = snapshot.parent_config or {} + parent_checkpoint_id = parent_config.get("configurable", {}).get("checkpoint_id") + metadata = snapshot.metadata or {} + created_at = snapshot.created_at or metadata.get("created_at", "") + tasks_raw = snapshot.tasks or () + tasks = [{"id": getattr(task, "id", ""), "name": getattr(task, "name", "")} for task in tasks_raw] + return ThreadStateResponse( - values=serialize_channel_values_for_api(channel_values), - next=[], + values=serialize_channel_values_for_api(snapshot.values), + next=list(snapshot.next or ()), metadata=metadata, - checkpoint_id=new_checkpoint_id, - created_at=coerce_iso(metadata.get("created_at", "")), + checkpoint={"id": checkpoint_id, "ts": coerce_iso(created_at)}, + checkpoint_id=checkpoint_id, + parent_checkpoint_id=parent_checkpoint_id, + created_at=coerce_iso(created_at), + tasks=tasks, ) @@ -1123,46 +1172,42 @@ async def get_thread_history( request: Request, background_tasks: BackgroundTasks, ) -> list[HistoryEntry]: - """Get checkpoint history for a thread. + """Get materialized graph state history for a thread. - Messages are read from the checkpointer's channel values (the - authoritative source) and serialized via - :func:`~deerflow.runtime.serialization.serialize_channel_values`. Only the latest (first) checkpoint carries the ``messages`` key to - avoid duplicating them across every entry. + avoid duplicating the complete conversation across every entry. """ checkpointer = get_checkpointer(request) - - config: dict[str, Any] = {"configurable": {"thread_id": thread_id}} - if body.before: - config["configurable"]["checkpoint_id"] = body.before + try: + accessor, config = await build_thread_checkpoint_state_accessor( + request, + thread_id=thread_id, + checkpoint_id=body.before, + ) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc entries: list[HistoryEntry] = [] is_latest_checkpoint = True try: - async for checkpoint_tuple in checkpointer.alist(config, limit=body.limit): - ckpt_config = getattr(checkpoint_tuple, "config", {}) - parent_config = getattr(checkpoint_tuple, "parent_config", None) - metadata = getattr(checkpoint_tuple, "metadata", {}) or {} - checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} + snapshots = await accessor.ahistory(config, limit=body.limit) + for snapshot in snapshots: + snapshot_config = snapshot.config or {} + parent_config = snapshot.parent_config or {} + metadata = snapshot.metadata or {} + materialized_values = snapshot.values if isinstance(snapshot.values, dict) else {} - checkpoint_id = ckpt_config.get("configurable", {}).get("checkpoint_id", "") - parent_id = None - if parent_config: - parent_id = parent_config.get("configurable", {}).get("checkpoint_id") + checkpoint_id = snapshot_config.get("configurable", {}).get("checkpoint_id", "") + parent_id = parent_config.get("configurable", {}).get("checkpoint_id") - channel_values = checkpoint.get("channel_values", {}) - - # Build values from checkpoint channel_values values: dict[str, Any] = {} - if title := channel_values.get("title"): + if title := materialized_values.get("title"): values["title"] = title - if thread_data := channel_values.get("thread_data"): + if thread_data := materialized_values.get("thread_data"): values["thread_data"] = thread_data - # Attach messages only to the latest checkpoint entry. if is_latest_checkpoint: - messages = channel_values.get("messages") + messages = materialized_values.get("messages") if messages: serialized_msgs = serialize_channel_values_for_api({"messages": messages}).get("messages", []) try: @@ -1244,9 +1289,7 @@ async def get_thread_history( is_latest_checkpoint = False - # Derive next tasks - tasks_raw = getattr(checkpoint_tuple, "tasks", []) or [] - next_tasks = [t.name for t in tasks_raw if hasattr(t, "name")] + next_tasks = list(snapshot.next or ()) # Strip LangGraph internal keys from metadata user_meta = {k: v for k, v in metadata.items() if k not in ("created_at", "updated_at", "step", "source", "writes", "parents", "run_durations")} @@ -1260,10 +1303,12 @@ async def get_thread_history( parent_checkpoint_id=parent_id, metadata=user_meta, values=values, - created_at=coerce_iso(metadata.get("created_at", "")), + created_at=coerce_iso(snapshot.created_at or metadata.get("created_at", "")), next=next_tasks, ) ) + except _CHECKPOINT_MODE_ERRORS as exc: + raise _checkpoint_mode_http_error(exc, thread_id) from exc except Exception: logger.exception("Failed to get history for thread %s", sanitize_log_param(thread_id)) raise HTTPException(status_code=500, detail="Failed to get thread history") diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index ef49a424c..453ef8a8f 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -34,6 +34,7 @@ from deerflow.config.app_config import get_app_config from deerflow.runtime import ( END_SENTINEL, HEARTBEAT_SENTINEL, + CheckpointStateAccessor, ConflictError, DisconnectMode, RunManager, @@ -41,8 +42,16 @@ from deerflow.runtime import ( RunStatus, StreamBridge, UnsupportedStrategyError, + build_state_mutation_graph, run_agent, ) +from deerflow.runtime.checkpoint_mode import ( + INTERNAL_CHECKPOINT_MODE_KEY, + CheckpointModeMismatchError, + checkpoint_tuple_uses_delta, + inject_checkpoint_mode, +) +from deerflow.runtime.checkpoint_state import graph_state_schema from deerflow.runtime.goal import goal_thread_lock from deerflow.runtime.runs.naming import resolve_root_run_name from deerflow.runtime.secret_context import redact_config_secrets @@ -494,6 +503,7 @@ def build_run_config( else: configurable = {"thread_id": thread_id} configurable.update(request_config.get("configurable") or {}) + configurable["thread_id"] = thread_id config["configurable"] = configurable for k, v in request_config.items(): if k not in ("configurable", "context"): @@ -536,11 +546,275 @@ def build_run_config( if isinstance(runtime_context, dict): runtime_context["agent_name"] = effective_agent_name config.setdefault("run_name", resolve_root_run_name(config, normalized)) + for section in ("configurable", "context"): + external_values = config.get(section) + if isinstance(external_values, dict): + external_values.pop(INTERNAL_CHECKPOINT_MODE_KEY, None) + if metadata: config.setdefault("metadata", {}).update(metadata) return config +def build_checkpoint_state_mutation_accessor( + request: Request, + *, + thread_id: str, + as_node: str, + checkpoint_id: str | None = None, + state_schema: Any | None = None, +) -> tuple[CheckpointStateAccessor, dict[str, Any]]: + """Build a state-only graph whose writer node finishes immediately. + + ``state_schema`` should be the thread's effective schema (from + :func:`graph_state_schema` on the assistant graph) whenever the write + carries materialized state; with the base-schema fallback, channels + contributed by custom middleware are silently discarded. + """ + mode = getattr(request.app.state, "checkpoint_channel_mode", "full") + config: dict[str, Any] = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + } + } + if checkpoint_id is not None: + config["configurable"]["checkpoint_id"] = checkpoint_id + inject_checkpoint_mode(config, mode) + + graph = build_state_mutation_graph(as_node, mode, state_schema) + accessor = CheckpointStateAccessor.bind( + graph, + get_checkpointer(request), + store=getattr(request.app.state, "store", None), + mode=mode, + ) + return accessor, config + + +# Cache of factory-built accessor graphs. Accessor operations (aget_state / +# aupdate_state) never execute graph nodes or middleware, so per-request +# variations (user, model, skills) cannot affect materialization semantics; +# the compiled graph is stable per (assistant_id, mode, app_config). The +# factory and app_config identities are re-validated on every call so patched +# factories take effect immediately and a config.yaml hot-reload (which +# rebuilds the AppConfig object) never serves a stale compiled graph — the +# cached reference keeps the old config alive, so id-reuse cannot produce a +# false hit. Bounded: cleared when too many distinct assistants appear. +_STATE_ACCESSOR_GRAPH_CACHE_MAX = 64 +_state_accessor_graph_cache: dict[tuple[str | None, str], tuple[Any, Any, Any]] = {} + + +def _state_accessor_graph(agent_factory: Any, assistant_id: str | None, mode: str, config: dict[str, Any]) -> Any: + app_config = (config.get("context") or {}).get("app_config") + key = (assistant_id, mode) + cached = _state_accessor_graph_cache.get(key) + if cached is not None and cached[0] is agent_factory and cached[1] is app_config: + return cached[2] + if len(_state_accessor_graph_cache) >= _STATE_ACCESSOR_GRAPH_CACHE_MAX: + _state_accessor_graph_cache.clear() + graph = agent_factory(config=config) + _state_accessor_graph_cache[key] = (agent_factory, app_config, graph) + return graph + + +class _RawCheckpointSnapshot: + """StateSnapshot-shaped view over a raw checkpoint tuple (full mode only). + + ``next``/``tasks`` are not derivable without the compiled graph and + degrade to empty; everything the read endpoints serialize (values, + metadata, config ancestry, created_at) comes straight from the tuple. + """ + + __slots__ = ("config", "values", "metadata", "parent_config", "created_at", "tasks", "tasks_known", "next") + + def __init__(self, config: dict[str, Any], tup: Any | None) -> None: + self.config = getattr(tup, "config", None) or config + checkpoint = getattr(tup, "checkpoint", None) or {} + self.values = dict(checkpoint.get("channel_values") or {}) + self.metadata = dict(getattr(tup, "metadata", None) or {}) + self.parent_config = getattr(tup, "parent_config", None) + self.created_at = checkpoint.get("ts") or self.metadata.get("created_at", "") + self.tasks: tuple = () + self.tasks_known = False + self.next: tuple = () + + +class _RawCheckpointReadAccessor: + """Degraded full-mode read accessor for when the agent factory is down. + + Full-mode checkpoints persist complete ``channel_values``, so reads do not + need the compiled graph. The fail-closed delta gate still applies: delta + checkpoints are rejected with :class:`CheckpointModeMismatchError` instead + of being served as partial state. Writes are unsupported — mutation paths + keep using the graph-backed accessor. + """ + + def __init__(self, checkpointer: Any, mode: str) -> None: + self.checkpointer = checkpointer + self.mode = mode + + @staticmethod + def _gate(tup: Any) -> None: + if checkpoint_tuple_uses_delta(tup): + raise CheckpointModeMismatchError("Thread requires delta mode; materialize and convert its checkpoints before using full mode.") + + async def aget(self, config: dict[str, Any]) -> _RawCheckpointSnapshot: + tup = await self.checkpointer.aget_tuple(config) + self._gate(tup) + return _RawCheckpointSnapshot(config, tup) + + async def ahistory(self, config: dict[str, Any], *, limit: int | None = None) -> list[_RawCheckpointSnapshot]: + if limit is not None and limit <= 0: + return [] + result: list[_RawCheckpointSnapshot] = [] + before = None + walk_config = config + if config.get("configurable", {}).get("checkpoint_id"): + # Pregel's get_state_history treats config.checkpoint_id as the + # inclusive start of the walk, while alist(before=...) is + # exclusive — fetch the anchor explicitly so the degraded path + # matches the graph path. + before = config + walk_config = { + **config, + "configurable": {k: v for k, v in config.get("configurable", {}).items() if k != "checkpoint_id"}, + } + anchor = await self.checkpointer.aget_tuple(before) + self._gate(anchor) + if anchor is not None: + result.append(_RawCheckpointSnapshot(config, anchor)) + if limit is None or len(result) < limit: + remaining = None if limit is None else limit - len(result) + async for tup in self.checkpointer.alist(walk_config, before=before, limit=remaining): + self._gate(tup) + result.append(_RawCheckpointSnapshot(config, tup)) + if limit is not None and len(result) >= limit: + break + return result + + +def build_checkpoint_state_accessor( + request: Request, + *, + thread_id: str, + assistant_id: str | None = None, + checkpoint_id: str | None = None, +) -> tuple[CheckpointStateAccessor, dict[str, Any]]: + """Build the mode-selected lead graph used for materialized checkpoint state.""" + ctx = get_run_context(request) + config = build_run_config(thread_id, None, None, assistant_id=assistant_id) + configurable = config.setdefault("configurable", {}) + configurable["checkpoint_ns"] = "" + if checkpoint_id is not None: + configurable["checkpoint_id"] = checkpoint_id + + if ctx.app_config is not None: + config.setdefault("context", {})["app_config"] = ctx.app_config + inject_checkpoint_mode(config, ctx.checkpoint_channel_mode) + + agent_factory = resolve_agent_factory(assistant_id) + try: + graph = _state_accessor_graph(agent_factory, assistant_id, ctx.checkpoint_channel_mode, config) + except Exception: + if ctx.checkpoint_channel_mode != "full": + # Delta materialization needs the graph's channel table; there is + # no degraded path. Surface the factory failure as-is. + raise + # Full-mode checkpoints carry complete channel_values: degrade to raw + # checkpointer reads so state endpoints survive a broken agent factory + # (bad model config, MCP server down, misconfigured skill). + logger.warning( + "Agent factory unavailable for thread %s; falling back to raw checkpointer reads", + thread_id, + exc_info=True, + ) + return _RawCheckpointReadAccessor(ctx.checkpointer, ctx.checkpoint_channel_mode), config + accessor = CheckpointStateAccessor.bind( + graph, + ctx.checkpointer, + store=ctx.store, + mode=ctx.checkpoint_channel_mode, + ) + return accessor, config + + +async def resolve_thread_assistant_id( + request: Request, + thread_id: str, + *, + fail_closed: bool = False, +) -> str | None: + """Return the assistant_id recorded in thread metadata, or ``None``. + + Missing records degrade to ``None`` (the default lead agent). Store + failures do the same for read callers, while mutation callers set + ``fail_closed`` so they cannot compile a write graph with the wrong schema. + """ + from app.gateway.deps import get_thread_store + + try: + thread_store = get_thread_store(request) + record = await thread_store.get(thread_id) + except Exception: + logger.warning("Failed to resolve assistant_id for thread %s", thread_id, exc_info=True) + if fail_closed: + raise + return None + return record.get("assistant_id") if isinstance(record, dict) else None + + +async def build_thread_checkpoint_state_accessor( + request: Request, + *, + thread_id: str, + checkpoint_id: str | None = None, + fail_closed: bool = False, +) -> tuple[CheckpointStateAccessor, dict[str, Any]]: + """Single resolution boundary for state endpoints. + + Thread metadata -> assistant_id -> effective assistant graph. Materializing + with the default lead schema would drop channels contributed by a custom + ``AgentMiddleware.state_schema`` from the response. + """ + assistant_id = await resolve_thread_assistant_id(request, thread_id, fail_closed=fail_closed) + return build_checkpoint_state_accessor( + request, + thread_id=thread_id, + assistant_id=assistant_id, + checkpoint_id=checkpoint_id, + ) + + +async def build_thread_checkpoint_state_mutation_accessor( + request: Request, + *, + thread_id: str, + as_node: str, + checkpoint_id: str | None = None, +) -> tuple[CheckpointStateAccessor, dict[str, Any]]: + """Mutation accessor compiled with the thread's effective state schema. + + Derives the schema through :func:`build_thread_checkpoint_state_accessor` + so writes carrying materialized state do not silently discard + extension-owned channels. + """ + read_accessor, _read_config = await build_thread_checkpoint_state_accessor( + request, + thread_id=thread_id, + checkpoint_id=checkpoint_id, + fail_closed=True, + ) + state_schema = graph_state_schema(getattr(read_accessor, "graph", None)) + return build_checkpoint_state_mutation_accessor( + request, + thread_id=thread_id, + as_node=as_node, + checkpoint_id=checkpoint_id, + state_schema=state_schema, + ) + + async def apply_checkpoint_to_run_config( config: dict[str, Any], *, diff --git a/backend/packages/harness/deerflow/agents/__init__.py b/backend/packages/harness/deerflow/agents/__init__.py index 4b56cbc06..5e75e90f3 100644 --- a/backend/packages/harness/deerflow/agents/__init__.py +++ b/backend/packages/harness/deerflow/agents/__init__.py @@ -7,6 +7,7 @@ __all__ = [ "Prev", "make_lead_agent", "SandboxState", + "DeltaThreadState", "ThreadState", ] @@ -28,10 +29,14 @@ def __getattr__(name: str): prime_enabled_skills_cache() globals()[name] = make_lead_agent return make_lead_agent - if name in {"SandboxState", "ThreadState"}: - from .thread_state import SandboxState, ThreadState + if name in {"DeltaThreadState", "SandboxState", "ThreadState"}: + from .thread_state import DeltaThreadState, SandboxState, ThreadState - exports = {"SandboxState": SandboxState, "ThreadState": ThreadState} + exports = { + "DeltaThreadState": DeltaThreadState, + "SandboxState": SandboxState, + "ThreadState": ThreadState, + } globals().update(exports) return exports[name] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/backend/packages/harness/deerflow/agents/factory.py b/backend/packages/harness/deerflow/agents/factory.py index b03490e66..e65772330 100644 --- a/backend/packages/harness/deerflow/agents/factory.py +++ b/backend/packages/harness/deerflow/agents/factory.py @@ -22,7 +22,8 @@ from deerflow.agents.features import RuntimeFeatures from deerflow.agents.middlewares.clarification_middleware import ClarificationMiddleware from deerflow.agents.middlewares.dangling_tool_call_middleware import DanglingToolCallMiddleware from deerflow.agents.middlewares.tool_error_handling_middleware import ToolErrorHandlingMiddleware -from deerflow.agents.thread_state import ThreadState +from deerflow.agents.thread_state import adapt_state_schema_for_mode, get_thread_state_schema, normalize_middleware_state_schemas +from deerflow.config.database_config import CheckpointChannelMode from deerflow.tools.builtins import ask_clarification_tool if TYPE_CHECKING: @@ -70,6 +71,7 @@ def create_deerflow_agent( extra_middleware: list[AgentMiddleware] | None = None, plan_mode: bool = False, state_schema: type | None = None, + checkpoint_channel_mode: CheckpointChannelMode = "full", checkpointer: BaseCheckpointSaver | None = None, name: str = "default", ) -> CompiledStateGraph: @@ -99,6 +101,12 @@ def create_deerflow_agent( Enable TodoMiddleware for task tracking. state_schema: LangGraph state type. Defaults to ``ThreadState``. + checkpoint_channel_mode: + Checkpoint representation for accumulating channels. Defaults to the + full-state compatibility schema. ``"delta"`` requires the guarded + persistence paths (mode markers + compatibility gate) and is therefore + rejected when combined with *checkpointer* in this factory; without a + checkpointer the graph is ephemeral and delta is allowed. checkpointer: Optional persistence backend. name: @@ -111,6 +119,14 @@ def create_deerflow_agent( """ if middleware is not None and features is not None: raise ValueError("Cannot specify both 'middleware' and 'features'. Use one or the other.") + if checkpoint_channel_mode == "delta" and checkpointer is not None: + raise ValueError( + "create_deerflow_agent does not support checkpoint_channel_mode='delta' with a checkpointer: " + "persisted graphs built here bypass checkpoint mode marker injection and the fail-closed " + "compatibility gate (see deerflow.runtime.checkpoint_mode), so a mixed-mode store would " + "silently corrupt thread state. Use the guarded application paths (make_lead_agent or " + "DeerFlowClient) for delta persistence; delta without a checkpointer is ephemeral and allowed." + ) if middleware is not None and extra_middleware: raise ValueError("Cannot use 'extra_middleware' with 'middleware' (full takeover).") if extra_middleware: @@ -119,7 +135,7 @@ def create_deerflow_agent( raise TypeError(f"extra_middleware items must be AgentMiddleware instances, got {type(mw).__name__}") effective_tools: list[BaseTool] = list(tools or []) - effective_state = state_schema or ThreadState + effective_state = get_thread_state_schema(checkpoint_channel_mode) if state_schema is None else adapt_state_schema_for_mode(state_schema, checkpoint_channel_mode) if middleware is not None: effective_middleware = list(middleware) @@ -138,6 +154,11 @@ def create_deerflow_agent( effective_tools.append(t) existing_names.add(t.name) + effective_middleware = normalize_middleware_state_schemas( + effective_middleware, + checkpoint_channel_mode, + ) + return create_agent( model=model, tools=effective_tools or None, diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index 29d5848ca..0ec8b4389 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -45,12 +45,18 @@ from deerflow.agents.middlewares.todo_middleware import TodoMiddleware from deerflow.agents.middlewares.token_usage_middleware import TokenUsageMiddleware from deerflow.agents.middlewares.tool_error_handling_middleware import build_lead_runtime_middlewares from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware -from deerflow.agents.thread_state import ThreadState +from deerflow.agents.thread_state import get_thread_state_schema, normalize_middleware_state_schemas from deerflow.config.agents_config import load_agent_config, validate_agent_name from deerflow.config.app_config import AppConfig, get_app_config from deerflow.config.memory_config import should_use_memory_tools from deerflow.config.subagents_config import DEFAULT_MAX_TOTAL_SUBAGENTS_PER_RUN from deerflow.models import create_chat_model +from deerflow.runtime.checkpoint_mode import ( + INTERNAL_CHECKPOINT_MODE_KEY, + freeze_checkpoint_channel_mode, + frozen_checkpoint_channel_mode, + inject_checkpoint_mode, +) from deerflow.skills.types import Skill from deerflow.tracing import build_tracing_callbacks @@ -453,7 +459,27 @@ def make_lead_agent(config: RunnableConfig): """LangGraph graph factory; keep the signature compatible with LangGraph Server.""" runtime_config = _get_runtime_config(config) runtime_app_config = runtime_config.get("app_config") - return _make_lead_agent(config, app_config=runtime_app_config or get_app_config()) + if not isinstance(runtime_app_config, AppConfig): + runtime_app_config = get_app_config() + # Mode selection precedence, pinned by test_checkpoint_mode.py: + # - First freeze: the app config owns the process mode; a client-supplied + # configurable key is ignored so a direct LangGraph request cannot + # reconfigure (or crash) a fresh process. + # - Once frozen: an internally injected key (run worker / gateway) or the + # app config must match the frozen mode; ``freeze_checkpoint_channel_mode`` + # fails closed on any mismatch, so neither a forged key nor a config.yaml + # change can silently reconfigure the process. + frozen_mode = frozen_checkpoint_channel_mode() + if frozen_mode is None: + requested_mode = runtime_app_config.database.checkpoint_channel_mode + else: + requested_mode = (config.get("configurable", {}) or {}).get( + INTERNAL_CHECKPOINT_MODE_KEY, + runtime_app_config.database.checkpoint_channel_mode, + ) + mode = freeze_checkpoint_channel_mode(requested_mode) + inject_checkpoint_mode(config, mode) + return _make_lead_agent(config, app_config=runtime_app_config) def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): @@ -464,6 +490,10 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): cfg = _get_runtime_config(config) resolved_app_config = app_config + mode = (config.get("configurable", {}) or {}).get( + INTERNAL_CHECKPOINT_MODE_KEY, + resolved_app_config.database.checkpoint_channel_mode, + ) # Extract user_id for user-scoped skill loading. # LangGraph gateway injects user_id into config["configurable"]; @@ -578,14 +608,17 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): return create_agent( model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled, app_config=resolved_app_config, attach_tracing=False), tools=final_tools, - middleware=build_middlewares( - config, - model_name=model_name, - available_skills=set(_BOOTSTRAP_SKILL_NAMES), - app_config=resolved_app_config, - deferred_setup=setup, - mcp_routing_middleware=mcp_routing_middleware, - user_id=resolved_user_id, + middleware=normalize_middleware_state_schemas( + build_middlewares( + config, + model_name=model_name, + available_skills=set(_BOOTSTRAP_SKILL_NAMES), + app_config=resolved_app_config, + deferred_setup=setup, + mcp_routing_middleware=mcp_routing_middleware, + user_id=resolved_user_id, + ), + mode, ), system_prompt=apply_prompt_template( subagent_enabled=subagent_enabled, @@ -597,7 +630,7 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): user_id=resolved_user_id, skill_names=skill_setup.skill_names or None, ), - state_schema=ThreadState, + state_schema=get_thread_state_schema(mode), ) # Custom agents can update their own SOUL.md / config via update_agent. @@ -645,15 +678,18 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): return create_agent( model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled, reasoning_effort=reasoning_effort, app_config=resolved_app_config, attach_tracing=False), tools=final_tools, - middleware=build_middlewares( - config, - model_name=model_name, - agent_name=agent_name, - available_skills=available_skills, - app_config=resolved_app_config, - deferred_setup=setup, - mcp_routing_middleware=mcp_routing_middleware, - user_id=resolved_user_id, + middleware=normalize_middleware_state_schemas( + build_middlewares( + config, + model_name=model_name, + agent_name=agent_name, + available_skills=available_skills, + app_config=resolved_app_config, + deferred_setup=setup, + mcp_routing_middleware=mcp_routing_middleware, + user_id=resolved_user_id, + ), + mode, ), system_prompt=apply_prompt_template( subagent_enabled=subagent_enabled, @@ -667,5 +703,5 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig): user_id=resolved_user_id, skill_names=skill_setup.skill_names or None, ), - state_schema=ThreadState, + state_schema=get_thread_state_schema(mode), ) diff --git a/backend/packages/harness/deerflow/agents/thread_state.py b/backend/packages/harness/deerflow/agents/thread_state.py index 0b10cd5d3..385794a67 100644 --- a/backend/packages/harness/deerflow/agents/thread_state.py +++ b/backend/packages/harness/deerflow/agents/thread_state.py @@ -1,9 +1,16 @@ -from collections.abc import Mapping -from typing import Annotated, NotRequired, TypedDict +import copy +from collections.abc import Mapping, Sequence +from functools import cache +from typing import Annotated, Any, NotRequired, TypedDict, get_type_hints from langchain.agents import AgentState +from langchain_core.messages import AnyMessage +from langgraph.channels import DeltaChannel +from langgraph.graph.message import add_messages +import deerflow.checkpoint_patches as _checkpoint_patches # noqa: F401 - import-time saver fixes from deerflow.agents.goal_state import GoalState +from deerflow.config.database_config import CheckpointChannelMode from deerflow.subagents.status_contract import SUBAGENT_STATUS_VALUES @@ -249,3 +256,67 @@ class ThreadState(AgentState): delegations: Annotated[list[DelegationEntry], merge_delegations] skill_context: Annotated[list[SkillEntry], merge_skill_context] summary_text: NotRequired[str | None] + + +def merge_message_writes(state: list[AnyMessage], writes: Sequence[Any]) -> list[AnyMessage]: + result = list(state) + for write in writes: + result = list(add_messages(result, write)) + return result + + +DELTA_MESSAGES_FIELD = Annotated[ + list[AnyMessage], + DeltaChannel(merge_message_writes, snapshot_frequency=1000), +] + + +class DeltaThreadState(ThreadState): + messages: DELTA_MESSAGES_FIELD + + +THREAD_STATE_REDUCER_FIELDS = frozenset( + { + "messages", + "sandbox", + "artifacts", + "todos", + "goal", + "viewed_images", + "promoted", + "delegations", + "skill_context", + } +) + + +def get_thread_state_schema(mode: CheckpointChannelMode) -> type: + return DeltaThreadState if mode == "delta" else ThreadState + + +@cache +def adapt_state_schema_for_mode(schema: type, mode: CheckpointChannelMode) -> type: + if mode == "full": + return schema + annotations = get_type_hints(schema, include_extras=True) + annotations["messages"] = DELTA_MESSAGES_FIELD + return TypedDict( + f"Delta{schema.__module__.replace('.', '_')}_{schema.__name__}", + annotations, + total=getattr(schema, "__total__", True), + ) + + +def normalize_middleware_state_schemas(middleware: Sequence[Any], mode: CheckpointChannelMode) -> list[Any]: + if mode == "full": + return list(middleware) + normalized = [] + for item in middleware: + schema = getattr(item, "state_schema", None) + if schema is None: + normalized.append(item) + continue + adapted = copy.copy(item) + adapted.state_schema = adapt_state_schema_for_mode(schema, mode) + normalized.append(adapted) + return normalized diff --git a/backend/packages/harness/deerflow/checkpoint_patches.py b/backend/packages/harness/deerflow/checkpoint_patches.py new file mode 100644 index 000000000..c82cc573c --- /dev/null +++ b/backend/packages/harness/deerflow/checkpoint_patches.py @@ -0,0 +1,99 @@ +"""Compatibility patches for third-party checkpoint savers. + +Lives at the top-level package (not ``deerflow.runtime``) so it can be +imported from ``deerflow.agents.thread_state`` without pulling in the heavy +``deerflow.runtime`` package __init__ (which eagerly imports the runs +machinery). Anchored from ``deerflow.agents.thread_state`` so every process +that builds a DeerFlow graph (gateway, workers, in-process LangGraph +runtime, tests) runs with the fixes in place. +""" + +from __future__ import annotations + +import importlib.metadata +import logging +from typing import Any + +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.memory import InMemorySaver +from packaging.version import Version + +logger = logging.getLogger(__name__) + +_PATCH_FLAG = "_deerflow_delta_history_patched" +# The patch was authored and verified against langgraph 1.2.9 +# (langgraph/checkpoint/memory/__init__.py::InMemorySaver.get_delta_channel_history). +# On any newer LangGraph the override must be re-inspected before keeping the +# patch: if upstream fixed or removed it, this module must stand down. +_PATCH_VALIDATED_LANGGRAPH_VERSION = Version("1.2.9") + + +def _get_delta_channel_history_via_base(self: Any, *, config: Any, channels: Any) -> Any: + return BaseCheckpointSaver.get_delta_channel_history(self, config=config, channels=channels) + + +async def _aget_delta_channel_history_via_base(self: Any, *, config: Any, channels: Any) -> Any: + return await BaseCheckpointSaver.aget_delta_channel_history(self, config=config, channels=channels) + + +def _upstream_override_present() -> bool: + """True while InMemorySaver still ships its own (buggy) override.""" + return ( + getattr(InMemorySaver, "get_delta_channel_history", None) is not None + and InMemorySaver.get_delta_channel_history is not BaseCheckpointSaver.get_delta_channel_history + and InMemorySaver.aget_delta_channel_history is not BaseCheckpointSaver.aget_delta_channel_history + ) + + +def ensure_inmemory_delta_history_patch() -> None: + """Fix InMemorySaver dropping writes on full -> delta migrated threads. + + ``InMemorySaver.get_delta_channel_history`` overrides the base walk with a + single-pass version that, upon reaching the first checkpoint carrying a + non-empty plain-value blob for a channel, skips that checkpoint's *own* + pending writes as "subsumed" by the blob. That is only true when the blob + was written by that same checkpoint. When the version was carried forward + from an older ancestor - exactly the first superstep after a full -> delta + migration, where the input write lands on a checkpoint still referencing + the pre-delta blob version - those pending writes postdate the blob and + are silently dropped: the first message appended after migration vanishes + from materialized state. + + Both the base implementation (used by the SQLite savers) and the Postgres + override collect the terminating checkpoint's writes *before* treating its + blob as the seed, which is the correct order. This patch delegates + InMemorySaver to the base implementation - one ``get_tuple`` per ancestor + instead of a single fused walk, which is fine for dict-backed storage. + + Idempotent. Guarded: stands down when the upstream override disappears or + the assignment fails, and warns once LangGraph moves past the validated + version so the patch is re-inspected instead of silently overriding an + upstream fix. Remove once LangGraph fixes the override upstream (no + upstream issue exists yet; re-check ``InMemorySaver.get_delta_channel_history`` + on every langgraph upgrade). + """ + if getattr(InMemorySaver, _PATCH_FLAG, False): + return + try: + langgraph_version = Version(importlib.metadata.version("langgraph")) + except Exception: + langgraph_version = _PATCH_VALIDATED_LANGGRAPH_VERSION + if langgraph_version > _PATCH_VALIDATED_LANGGRAPH_VERSION: + logger.warning( + "langgraph %s is newer than the version (%s) the InMemorySaver delta-history patch was validated against; re-inspect the upstream override before relying on the patch.", + langgraph_version, + _PATCH_VALIDATED_LANGGRAPH_VERSION, + ) + try: + if not _upstream_override_present(): + # Upstream removed its override (fixed or refactored): the base + # implementation is already in use, nothing to patch. + return + InMemorySaver.get_delta_channel_history = _get_delta_channel_history_via_base # type: ignore[method-assign] + InMemorySaver.aget_delta_channel_history = _aget_delta_channel_history_via_base # type: ignore[method-assign] + InMemorySaver._deerflow_delta_history_patched = True # type: ignore[attr-defined] + except (AttributeError, TypeError): + logger.warning("Failed to apply the InMemorySaver delta-history patch; leaving the upstream implementation untouched.", exc_info=True) + + +ensure_inmemory_delta_history_patch() diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index 0acfd05a9..966b59e3b 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -36,12 +36,18 @@ from langchain_core.runnables import RunnableConfig from deerflow.agents.lead_agent.agent import build_middlewares from deerflow.agents.lead_agent.prompt import apply_prompt_template, get_enabled_skills_for_config -from deerflow.agents.thread_state import ThreadState +from deerflow.agents.thread_state import get_thread_state_schema, normalize_middleware_state_schemas from deerflow.config.agents_config import AGENT_NAME_PATTERN from deerflow.config.app_config import get_app_config, is_trace_correlation_enabled, reload_app_config from deerflow.config.extensions_config import ExtensionsConfig, SkillStateConfig, get_extensions_config, reload_extensions_config from deerflow.config.paths import get_paths from deerflow.models import create_chat_model +from deerflow.runtime import CheckpointStateAccessor +from deerflow.runtime.checkpoint_mode import ( + ensure_checkpoint_mode_compatible, + freeze_checkpoint_channel_mode, + inject_checkpoint_mode, +) from deerflow.runtime.goal import DEFAULT_MAX_GOAL_CONTINUATIONS, build_goal_state, goal_thread_lock, read_thread_goal, write_thread_goal from deerflow.runtime.user_context import get_effective_user_id from deerflow.skills.describe import build_skill_search_setup @@ -171,6 +177,7 @@ class DeerFlowClient: if config_path is not None: reload_app_config(config_path) self._app_config = get_app_config() + self._checkpoint_channel_mode = freeze_checkpoint_channel_mode(self._app_config.database.checkpoint_channel_mode) if agent_name is not None and not AGENT_NAME_PATTERN.match(agent_name): raise ValueError(f"Invalid agent name '{agent_name}'. Must match pattern: {AGENT_NAME_PATTERN.pattern}") @@ -247,6 +254,7 @@ class DeerFlowClient: cfg.get("max_total_subagents"), self._agent_name, frozenset(self._available_skills) if self._available_skills is not None else None, + self._checkpoint_channel_mode, ) if self._agent is not None and self._agent_config_key == key: @@ -286,16 +294,19 @@ class DeerFlowClient: # Attaching them again on the model would emit duplicate spans. "model": create_chat_model(name=model_name, thinking_enabled=thinking_enabled, attach_tracing=False), "tools": final_tools, - "middleware": build_middlewares( - config, - model_name=model_name, - agent_name=self._agent_name, - available_skills=self._available_skills, - custom_middlewares=self._middlewares, - app_config=self._app_config, - deferred_setup=deferred_setup, - mcp_routing_middleware=mcp_routing_middleware, - user_id=get_effective_user_id(), + "middleware": normalize_middleware_state_schemas( + build_middlewares( + config, + model_name=model_name, + agent_name=self._agent_name, + available_skills=self._available_skills, + custom_middlewares=self._middlewares, + app_config=self._app_config, + deferred_setup=deferred_setup, + mcp_routing_middleware=mcp_routing_middleware, + user_id=get_effective_user_id(), + ), + self._checkpoint_channel_mode, ), "system_prompt": apply_prompt_template( subagent_enabled=subagent_enabled, @@ -309,7 +320,7 @@ class DeerFlowClient: user_id=get_effective_user_id(), skill_names=skill_setup.skill_names or None, ), - "state_schema": ThreadState, + "state_schema": get_thread_state_schema(self._checkpoint_channel_mode), } checkpointer = self._checkpointer if checkpointer is None: @@ -555,41 +566,50 @@ class DeerFlowClient: return {"thread_list": threads[:limit]} def get_thread(self, thread_id: str) -> dict: - """Get the complete thread record, including all node execution records. - - Args: - thread_id: Thread ID. - - Returns: - Dict containing the thread's full checkpoint history. - """ + """Get the complete materialized checkpoint history for a thread.""" checkpointer = self._get_thread_checkpointer() + config = self._get_runnable_config(thread_id) + self._ensure_agent(config) + if self._agent is None: + raise RuntimeError("Agent was not initialized") + + accessor = CheckpointStateAccessor.bind( + self._agent, + checkpointer, + mode=self._checkpoint_channel_mode, + ) + # One streaming walk collects pending_writes per checkpoint id; a + # per-snapshot get_tuple would cost one round-trip per checkpoint. + pending_writes_by_checkpoint: dict[str, list] = {} + for raw_tuple in checkpointer.list(config): + raw_checkpoint_id = raw_tuple.config.get("configurable", {}).get("checkpoint_id") + if raw_checkpoint_id: + pending_writes_by_checkpoint[raw_checkpoint_id] = list(getattr(raw_tuple, "pending_writes", ()) or ()) - config = {"configurable": {"thread_id": thread_id}} checkpoints = [] + for snapshot in accessor.history(config): + values = dict(snapshot.values or {}) + if "messages" in values: + values["messages"] = [self._serialize_message(message) if hasattr(message, "content") else message for message in values["messages"]] - for cp in checkpointer.list(config): - channel_values = dict(cp.checkpoint.get("channel_values", {})) - if "messages" in channel_values: - channel_values["messages"] = [self._serialize_message(m) if hasattr(m, "content") else m for m in channel_values["messages"]] - - cfg = cp.config.get("configurable", {}) - parent_cfg = cp.parent_config.get("configurable", {}) if cp.parent_config else {} + snapshot_config = snapshot.config or {} + configurable = snapshot_config.get("configurable", {}) + parent_config = snapshot.parent_config or {} + parent_configurable = parent_config.get("configurable", {}) + pending_writes = pending_writes_by_checkpoint.get(configurable.get("checkpoint_id"), []) checkpoints.append( { - "checkpoint_id": cfg.get("checkpoint_id"), - "parent_checkpoint_id": parent_cfg.get("checkpoint_id"), - "ts": cp.checkpoint.get("ts"), - "metadata": cp.metadata, - "values": channel_values, - "pending_writes": [{"task_id": w[0], "channel": w[1], "value": w[2]} for w in getattr(cp, "pending_writes", [])], + "checkpoint_id": configurable.get("checkpoint_id"), + "parent_checkpoint_id": parent_configurable.get("checkpoint_id"), + "ts": snapshot.created_at, + "metadata": snapshot.metadata, + "values": values, + "pending_writes": [{"task_id": write[0], "channel": write[1], "value": write[2]} for write in pending_writes], } ) - # Sort globally by timestamp to prevent partial ordering issues caused by different namespaces (e.g., subgraphs) - checkpoints.sort(key=lambda x: x["ts"] if x["ts"] else "") - + checkpoints.sort(key=lambda checkpoint: checkpoint["ts"] or "") return {"thread_id": thread_id, "checkpoints": checkpoints} # ------------------------------------------------------------------ @@ -736,6 +756,24 @@ class DeerFlowClient: thread_id = str(uuid.uuid4()) config = self._get_runnable_config(thread_id, **kwargs) + inject_checkpoint_mode(config, self._checkpoint_channel_mode) + checkpoint_config = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + } + } + checkpointer = self._checkpointer + if checkpointer is None: + from deerflow.runtime.checkpointer import get_checkpointer + + checkpointer = get_checkpointer() + if checkpointer is not None: + ensure_checkpoint_mode_compatible( + checkpointer, + checkpoint_config, + self._checkpoint_channel_mode, + ) # Inject tracing callbacks and Langfuse trace metadata at the graph # invocation root so the embedded client matches the gateway worker's diff --git a/backend/packages/harness/deerflow/config/database_config.py b/backend/packages/harness/deerflow/config/database_config.py index 000ebe482..4b970ac9e 100644 --- a/backend/packages/harness/deerflow/config/database_config.py +++ b/backend/packages/harness/deerflow/config/database_config.py @@ -36,12 +36,23 @@ from typing import Literal from pydantic import BaseModel, Field +CheckpointChannelMode = Literal["full", "delta"] + class DatabaseConfig(BaseModel): backend: Literal["memory", "sqlite", "postgres"] = Field( default="memory", description=("Storage backend for both checkpointer and application data. 'memory' for development (no persistence across restarts), 'sqlite' for single-node deployment, 'postgres' for production multi-node deployment."), ) + checkpoint_channel_mode: CheckpointChannelMode = Field( + default="full", + description=( + "Checkpoint representation for accumulating channels. " + "'full' preserves full-value message checkpoints; 'delta' uses " + "LangGraph DeltaChannel for messages. Restart is required, and all " + "processes sharing one checkpoint database must use the same value." + ), + ) sqlite_dir: str = Field( default=".deer-flow/data", description=("Directory for the SQLite database file. Both checkpointer and application data share {sqlite_dir}/deerflow.db."), diff --git a/backend/packages/harness/deerflow/runtime/__init__.py b/backend/packages/harness/deerflow/runtime/__init__.py index d87090748..2940766a5 100644 --- a/backend/packages/harness/deerflow/runtime/__init__.py +++ b/backend/packages/harness/deerflow/runtime/__init__.py @@ -5,6 +5,7 @@ Re-exports the public API of :mod:`~deerflow.runtime.runs` and directly from ``deerflow.runtime``. """ +from .checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph from .checkpointer import checkpointer_context, get_checkpointer, make_checkpointer, reset_checkpointer from .runs import CancelOutcome, ConflictError, DisconnectMode, RunContext, RunManager, RunRecord, RunStatus, UnsupportedStrategyError, run_agent from .serialization import serialize, serialize_channel_values, serialize_channel_values_for_api, serialize_lc_object, serialize_messages_tuple, strip_data_url_image_blocks @@ -16,6 +17,9 @@ from .store import get_store, make_store, reset_store, store_context from .stream_bridge import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamBridge, StreamEvent, make_stream_bridge __all__ = [ + # checkpoint state + "CheckpointStateAccessor", + "build_state_mutation_graph", # checkpointer "checkpointer_context", "get_checkpointer", diff --git a/backend/packages/harness/deerflow/runtime/checkpoint_mode.py b/backend/packages/harness/deerflow/runtime/checkpoint_mode.py new file mode 100644 index 000000000..ff1442ae6 --- /dev/null +++ b/backend/packages/harness/deerflow/runtime/checkpoint_mode.py @@ -0,0 +1,108 @@ +"""Dual-mode checkpoint channel safety: mode freeze, metadata markers, and the fail-closed gate. + +Checkpointer storage runs in ``full`` mode (whole-snapshot channel values) or +``delta`` mode (LangGraph ``DeltaChannel``: sentinel blobs + per-step writes). +The mode is process-frozen at agent-build time, stamped into each checkpoint's +metadata on write, and enforced before every state access: a full-mode process +opening a delta thread raises :class:`CheckpointModeMismatchError` instead of +silently materializing empty state. Delta-mode processes read legacy full +checkpoints transparently, so full -> delta is the supported migration path. +""" + +from __future__ import annotations + +from typing import Any + +from deerflow.config.database_config import CheckpointChannelMode + +INTERNAL_CHECKPOINT_MODE_KEY = "__deerflow_checkpoint_channel_mode" +CHECKPOINT_MODE_METADATA_KEY = "deerflow_checkpoint_channel_mode" + + +class CheckpointModeMismatchError(RuntimeError): + """Raised before a full-mode graph reads a Delta checkpoint.""" + + +class CheckpointModeReconfigurationError(RuntimeError): + """Raised when a process attempts to hot-switch its persistence mode.""" + + +_frozen_checkpoint_channel_mode: CheckpointChannelMode | None = None + + +def frozen_checkpoint_channel_mode() -> CheckpointChannelMode | None: + """Return the process-frozen checkpoint channel mode, if already frozen.""" + return _frozen_checkpoint_channel_mode + + +def freeze_checkpoint_channel_mode(mode: CheckpointChannelMode) -> CheckpointChannelMode: + global _frozen_checkpoint_channel_mode + if _frozen_checkpoint_channel_mode is None: + _frozen_checkpoint_channel_mode = mode + elif _frozen_checkpoint_channel_mode != mode: + raise CheckpointModeReconfigurationError("checkpoint_channel_mode is restart-required and cannot change in a running process") + return _frozen_checkpoint_channel_mode + + +def inject_checkpoint_mode(config: dict[str, Any], mode: CheckpointChannelMode) -> None: + configurable = config.setdefault("configurable", {}) + configurable[INTERNAL_CHECKPOINT_MODE_KEY] = mode + metadata = config.setdefault("metadata", {}) + if mode == "delta": + metadata[CHECKPOINT_MODE_METADATA_KEY] = "delta" + else: + metadata.pop(CHECKPOINT_MODE_METADATA_KEY, None) + + +def checkpoint_metadata_uses_delta(metadata: Any) -> bool: + """Whether checkpoint metadata carries the delta-mode marker.""" + if not metadata: + return False + if metadata.get(CHECKPOINT_MODE_METADATA_KEY) == "delta": + return True + counters = metadata.get("counters_since_delta_snapshot") + return isinstance(counters, dict) and "messages" in counters + + +def checkpoint_tuple_uses_delta(checkpoint_tuple: Any) -> bool: + if checkpoint_tuple is None: + return False + return checkpoint_metadata_uses_delta(getattr(checkpoint_tuple, "metadata", {}) or {}) + + +def state_snapshot_uses_delta(snapshot: Any) -> bool: + """Whether a materialized ``StateSnapshot`` originates from a delta checkpoint.""" + if snapshot is None: + return False + return checkpoint_metadata_uses_delta(getattr(snapshot, "metadata", {}) or {}) + + +def raise_if_snapshot_incompatible(snapshot: Any, mode: CheckpointChannelMode) -> None: + """Fail closed when a full-mode process materialized a delta checkpoint. + + Runs on the ``StateSnapshot`` returned by ``get_state``/``get_state_history``, + so reads cost a single checkpoint fetch: the marker lives in + ``snapshot.metadata``. Reading the blob is harmless; silently *using* the + empty/partial state is the danger, and the caller never receives it. + """ + if mode == "full" and state_snapshot_uses_delta(snapshot): + raise CheckpointModeMismatchError("Thread requires delta mode; materialize and convert its checkpoints before using full mode.") + + +def ensure_checkpoint_mode_compatible(checkpointer: Any, config: dict[str, Any], mode: CheckpointChannelMode) -> None: + """Pre-write gate: a write cannot be un-applied, so it checks ahead of time. + + Reads use :func:`raise_if_snapshot_incompatible` on the returned snapshot + instead, avoiding the extra fetch. + """ + if mode == "delta": + return + if checkpoint_tuple_uses_delta(checkpointer.get_tuple(config)): + raise CheckpointModeMismatchError("Thread requires delta mode; materialize and convert its checkpoints before using full mode.") + + +async def aensure_checkpoint_mode_compatible(checkpointer: Any, config: dict[str, Any], mode: CheckpointChannelMode) -> None: + if mode == "delta": + return + if checkpoint_tuple_uses_delta(await checkpointer.aget_tuple(config)): + raise CheckpointModeMismatchError("Thread requires delta mode; materialize and convert its checkpoints before using full mode.") diff --git a/backend/packages/harness/deerflow/runtime/checkpoint_state.py b/backend/packages/harness/deerflow/runtime/checkpoint_state.py new file mode 100644 index 000000000..9264dafd6 --- /dev/null +++ b/backend/packages/harness/deerflow/runtime/checkpoint_state.py @@ -0,0 +1,187 @@ +"""Materialized checkpoint-state access and state-only mutation graphs. + +:class:`CheckpointStateAccessor` is the single choke point for thread +checkpoint-state reads and writes. It binds a compiled graph (which carries +the mode-matched channel schema), a checkpointer, and the frozen channel mode: +every operation injects the mode marker into the config and passes the +compatibility gate before touching state. Delta checkpoints store no full +``channel_values`` — raw saver reads see sentinels — so consumers must go +through this accessor instead of calling the checkpointer directly. + +:func:`build_state_mutation_graph` compiles a state-only graph (one no-op +node, entry = finish) for wholesale state replacement such as rollback +restore and context compaction: it shares the agent graph's checkpoint +machinery but schedules no pending nodes, so the written head stays idle. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from deerflow.agents.thread_state import get_thread_state_schema +from deerflow.config.database_config import CheckpointChannelMode +from deerflow.runtime.checkpoint_mode import ( + aensure_checkpoint_mode_compatible, + ensure_checkpoint_mode_compatible, + inject_checkpoint_mode, + raise_if_snapshot_incompatible, +) + + +def _finish_state_mutation(_state: dict[str, Any]) -> dict[str, Any]: + return {} + + +def build_state_mutation_graph(as_node: str, mode: CheckpointChannelMode, state_schema: Any | None = None) -> Any: + """Compile a state-only graph whose single writer node finishes immediately. + + ``update_state(..., as_node=...)`` requires the node to be registered in + the graph; a dedicated single-node graph applies reducer writes and + finishes, so the mutation checkpoint schedules no agent nodes and has no + pending ``next`` nodes. + + ``state_schema`` must be the thread's *effective* schema (the class the + assistant graph was compiled with) whenever the write carries materialized + state: the base ThreadState fallback does not know channels contributed by + custom middleware, and writes to unknown channels are silently discarded. + """ + if not as_node: + raise ValueError("as_node is required for checkpoint state mutation") + from langgraph.graph import StateGraph + + builder = StateGraph(state_schema if state_schema is not None else get_thread_state_schema(mode)) + builder.add_node(as_node, _finish_state_mutation) + builder.set_entry_point(as_node) + builder.set_finish_point(as_node) + return builder.compile() + + +def graph_state_schema(graph: Any) -> Any | None: + """Return the state schema class a compiled graph was built with. + + The schema is the first entry of ``StateGraph.schemas`` (state schema is + registered before input/output schemas). Returns ``None`` for stub + accessors in tests that do not wrap a real compiled graph. + """ + schemas = getattr(getattr(graph, "builder", None), "schemas", None) + if not schemas: + return None + return next(iter(schemas)) + + +def graph_writable_channels(graph: Any) -> frozenset[str] | None: + """Return the user-visible state channel names of a compiled graph. + + Excludes Pregel-internal channels (``__*``) and branch fan-in channels + (``branch:*``). Returns ``None`` when the graph does not expose channels + (stub accessors), so callers can fall back to the base ThreadState set. + """ + channels = getattr(graph, "channels", None) + if not channels: + return None + return frozenset(name for name in channels if not name.startswith("__") and not name.startswith("branch:")) + + +def graph_reducer_channels(graph: Any) -> frozenset[str] | None: + """Return channel names whose writes merge through a reducer. + + Covers classic reducers (``BinaryOperatorAggregate``) and delta channels: + both require ``Overwrite`` wrapping for replace-style writes, in any mode. + Returns ``None`` when the graph does not expose channels (stub + accessors), so callers can fall back to the base ThreadState set. + """ + from langgraph.channels import BinaryOperatorAggregate, DeltaChannel + + channels = getattr(graph, "channels", None) + if channels is None: + return None + return frozenset(name for name, channel in channels.items() if isinstance(channel, (BinaryOperatorAggregate, DeltaChannel))) + + +@dataclass +class CheckpointStateAccessor: + graph: Any + checkpointer: Any + mode: CheckpointChannelMode + + @classmethod + def bind( + cls, + graph: Any, + checkpointer: Any, + *, + store: Any | None = None, + mode: CheckpointChannelMode = "full", + ) -> CheckpointStateAccessor: + graph.checkpointer = checkpointer + if store is not None: + graph.store = store + return cls(graph=graph, checkpointer=checkpointer, mode=mode) + + def _prepare_config(self, config: dict[str, Any]) -> dict[str, Any]: + prepared = { + **config, + "configurable": dict(config.get("configurable", {})), + "metadata": dict(config.get("metadata", {})), + } + inject_checkpoint_mode(prepared, self.mode) + return prepared + + def get(self, config: dict[str, Any]) -> Any: + prepared = self._prepare_config(config) + snapshot = self.graph.get_state(prepared) + raise_if_snapshot_incompatible(snapshot, self.mode) + return snapshot + + async def aget(self, config: dict[str, Any]) -> Any: + prepared = self._prepare_config(config) + snapshot = await self.graph.aget_state(prepared) + raise_if_snapshot_incompatible(snapshot, self.mode) + return snapshot + + def history(self, config: dict[str, Any], *, limit: int | None = None) -> list[Any]: + prepared = self._prepare_config(config) + if limit is not None and limit <= 0: + return [] + result = [] + for snapshot in self.graph.get_state_history(prepared, limit=limit): + raise_if_snapshot_incompatible(snapshot, self.mode) + result.append(snapshot) + if limit is not None and len(result) >= limit: + break + return result + + async def ahistory(self, config: dict[str, Any], *, limit: int | None = None) -> list[Any]: + prepared = self._prepare_config(config) + if limit is not None and limit <= 0: + return [] + result = [] + async for snapshot in self.graph.aget_state_history(prepared, limit=limit): + raise_if_snapshot_incompatible(snapshot, self.mode) + result.append(snapshot) + if limit is not None and len(result) >= limit: + break + return result + + def update( + self, + config: dict[str, Any], + values: dict[str, Any], + *, + as_node: str | None = None, + ) -> dict[str, Any]: + prepared = self._prepare_config(config) + ensure_checkpoint_mode_compatible(self.checkpointer, prepared, self.mode) + return self.graph.update_state(prepared, values, as_node=as_node) + + async def aupdate( + self, + config: dict[str, Any], + values: dict[str, Any], + *, + as_node: str | None = None, + ) -> dict[str, Any]: + prepared = self._prepare_config(config) + await aensure_checkpoint_mode_compatible(self.checkpointer, prepared, self.mode) + return await self.graph.aupdate_state(prepared, values, as_node=as_node) diff --git a/backend/packages/harness/deerflow/runtime/context_compaction.py b/backend/packages/harness/deerflow/runtime/context_compaction.py index 7bbab2fd1..b279c862c 100644 --- a/backend/packages/harness/deerflow/runtime/context_compaction.py +++ b/backend/packages/harness/deerflow/runtime/context_compaction.py @@ -2,17 +2,14 @@ from __future__ import annotations -import hashlib from dataclasses import dataclass from types import SimpleNamespace -from typing import Any -from langgraph.checkpoint.base import uuid6 +from langgraph.types import Overwrite from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, create_summarization_middleware from deerflow.config.app_config import AppConfig, get_app_config -from deerflow.runtime.goal import _call_checkpointer_method, _next_channel_version -from deerflow.utils.time import now_iso +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor class ContextCompactionDisabled(RuntimeError): @@ -48,15 +45,8 @@ def _create_compaction_middleware( return middleware -def _checkpoint_namespace(checkpoint_tuple: Any) -> str: - config = getattr(checkpoint_tuple, "config", {}) or {} - configurable = config.get("configurable", {}) if isinstance(config, dict) else {} - checkpoint_ns = configurable.get("checkpoint_ns", "") if isinstance(configurable, dict) else "" - return checkpoint_ns if isinstance(checkpoint_ns, str) else "" - - async def compact_thread_context( - checkpointer: Any, + accessor: CheckpointStateAccessor, thread_id: str, *, keep: tuple[str, int | float] | None = None, @@ -70,13 +60,13 @@ async def compact_thread_context( middleware = _create_compaction_middleware(app_config=resolved_app_config, keep=keep) read_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} - checkpoint_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", read_config) - if checkpoint_tuple is None: + snapshot = await accessor.aget(read_config) + snapshot_config = snapshot.config or {} + checkpoint_id = snapshot_config.get("configurable", {}).get("checkpoint_id") + if not checkpoint_id: raise LookupError(f"Thread {thread_id} checkpoint not found") - checkpoint: dict[str, Any] = dict(getattr(checkpoint_tuple, "checkpoint", {}) or {}) - metadata: dict[str, Any] = dict(getattr(checkpoint_tuple, "metadata", {}) or {}) - channel_values: dict[str, Any] = dict(checkpoint.get("channel_values", {}) or {}) + channel_values = snapshot.values or {} messages = channel_values.get("messages") if not isinstance(messages, list) or not messages: return ThreadCompactionResult(thread_id=thread_id, compacted=False, reason="not_enough_messages") @@ -94,42 +84,15 @@ async def compact_thread_context( if result is None: return ThreadCompactionResult(thread_id=thread_id, compacted=False, reason="not_enough_messages") - channel_values["messages"] = list(result.preserved_messages) - channel_values["summary_text"] = result.summary_text - checkpoint["channel_values"] = channel_values - - channel_versions = dict(checkpoint.get("channel_versions", {}) or {}) - new_versions: dict[str, Any] = {} - for channel in ("messages", "summary_text"): - next_version = _next_channel_version(checkpointer, channel_versions.get(channel)) - channel_versions[channel] = next_version - new_versions[channel] = next_version - checkpoint["channel_versions"] = channel_versions - checkpoint["id"] = str(uuid6()) - checkpoint["ts"] = now_iso() - - metadata["source"] = "update" - metadata["updated_at"] = now_iso() - prev_step = metadata.get("step") - metadata["step"] = (prev_step + 1) if isinstance(prev_step, int) else 1 - metadata["writes"] = { - "manual_compaction": { - "messages": { - "removed": len(result.messages_to_summarize), - "preserved": len(result.preserved_messages), - }, - "summary_text": { - "sha256": hashlib.sha256(result.summary_text.encode("utf-8")).hexdigest(), - "chars": len(result.summary_text), - }, - } - } - - write_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": _checkpoint_namespace(checkpoint_tuple)}} - new_config = await _call_checkpointer_method(checkpointer, "aput", "put", write_config, checkpoint, metadata, new_versions) - new_checkpoint_id = None - if isinstance(new_config, dict): - new_checkpoint_id = new_config.get("configurable", {}).get("checkpoint_id") + updated_config = await accessor.aupdate( + snapshot.config, + { + "messages": Overwrite(list(result.preserved_messages)), + "summary_text": result.summary_text, + }, + as_node="manual_compaction", + ) + new_checkpoint_id = updated_config.get("configurable", {}).get("checkpoint_id") return ThreadCompactionResult( thread_id=thread_id, diff --git a/backend/packages/harness/deerflow/runtime/goal.py b/backend/packages/harness/deerflow/runtime/goal.py index e4f4e87e6..ff0be9d4a 100644 --- a/backend/packages/harness/deerflow/runtime/goal.py +++ b/backend/packages/harness/deerflow/runtime/goal.py @@ -525,6 +525,11 @@ async def write_thread_goal( "configurable": { "thread_id": thread_id, "checkpoint_ns": "", + # Parent the new checkpoint to the one it was derived from. + # Without this the saver stores a parentless checkpoint, which + # severs Delta-channel replay ancestry (and truncates history + # walks in full mode too). + "checkpoint_id": _checkpoint_id_from_tuple(checkpoint_tuple), } } await _call_checkpointer_method(checkpointer, "aput", "put", write_config, checkpoint, metadata, {"goal": next_version}) diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 0427fe1f8..81b860fb5 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -30,9 +30,16 @@ from functools import lru_cache from typing import Any, Literal, cast from langgraph.checkpoint.base import empty_checkpoint +from langgraph.types import Overwrite from deerflow.agents.goal_state import GoalEvaluation, GoalState from deerflow.config.app_config import AppConfig +from deerflow.config.database_config import CheckpointChannelMode +from deerflow.runtime.checkpoint_mode import ( + aensure_checkpoint_mode_compatible, + inject_checkpoint_mode, +) +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph, graph_state_schema from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY from deerflow.runtime.goal import ( DEFAULT_MAX_GOAL_CONTINUATIONS, @@ -142,6 +149,7 @@ class RunContext: run_events_config: Any | None = field(default=None) thread_store: Any | None = field(default=None) app_config: AppConfig | None = field(default=None) + checkpoint_channel_mode: CheckpointChannelMode = "full" on_run_completed: Any | None = field(default=None) @@ -270,7 +278,6 @@ async def run_agent( thread_id = record.thread_id requested_modes: set[str] = set(stream_modes or ["values"]) pre_run_checkpoint_id: str | None = None - pre_run_snapshot: dict[str, Any] | None = None pre_run_workspace_snapshot: WorkspaceSnapshot | None = None workspace_changes_user_id: str | None = None snapshot_capture_failed = False @@ -281,6 +288,11 @@ async def run_agent( # history would mark every subsequent run on this thread as ``error``. pre_existing_message_ids: set[str] = set() + # Bound agent graph accessor + captured pre-run rollback point; assigned + # inside the try block so the finally rollback path can fork the pre-run + # checkpoint lineage (see below). + accessor: CheckpointStateAccessor | None = None + rollback_point: RollbackPoint | None = None journal = None # Buffers subagent step events for batched persistence (#3779); assigned once # streaming starts and flushed in the finally block. Pre-bound to None so the @@ -296,6 +308,37 @@ async def run_agent( try: await run_manager.wait_for_prior_finalizing(thread_id, run_id) + mode = ctx.checkpoint_channel_mode + inject_checkpoint_mode(config, mode) + checkpoint_config = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + } + } + if checkpointer is not None: + await aensure_checkpoint_mode_compatible( + checkpointer, + checkpoint_config, + mode, + ) + configurable = config["configurable"] + selected_configurable = { + "thread_id": thread_id, + "checkpoint_ns": configurable.get("checkpoint_ns", ""), + } + for selector_key in ("checkpoint_id", "checkpoint_map"): + if selector_key in configurable: + selected_configurable[selector_key] = configurable[selector_key] + selected_checkpoint_config = { + "configurable": selected_configurable, + } + if selected_checkpoint_config != checkpoint_config: + await aensure_checkpoint_mode_compatible( + checkpointer, + selected_checkpoint_config, + mode, + ) # Initialize RunJournal + write human_message event. # These are inside the try block so any exception (e.g. a DB @@ -327,25 +370,6 @@ async def run_agent( except Exception: logger.warning("Could not capture pre-run workspace snapshot for run %s", run_id, exc_info=True) - # Snapshot the latest pre-run checkpoint so rollback can restore it. - if checkpointer is not None: - try: - config_for_check = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} - ckpt_tuple = await checkpointer.aget_tuple(config_for_check) - if ckpt_tuple is not None: - ckpt_config = getattr(ckpt_tuple, "config", {}).get("configurable", {}) - pre_run_checkpoint_id = ckpt_config.get("checkpoint_id") - pre_run_snapshot = { - "checkpoint_ns": ckpt_config.get("checkpoint_ns", ""), - "checkpoint": copy.deepcopy(getattr(ckpt_tuple, "checkpoint", {})), - "metadata": copy.deepcopy(getattr(ckpt_tuple, "metadata", {})), - "pending_writes": copy.deepcopy(getattr(ckpt_tuple, "pending_writes", []) or []), - } - pre_existing_message_ids = _collect_pre_existing_message_ids(pre_run_snapshot) - except Exception: - snapshot_capture_failed = True - logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True) - # 2. Publish metadata — useStream needs both run_id AND thread_id await bridge.publish( run_id, @@ -365,7 +389,6 @@ async def run_agent( # manually here because we drive the graph through ``agent.astream(config=...)`` # without passing the official ``context=`` parameter. runtime_ctx = _build_runtime_context(thread_id, run_id, config.get("context"), ctx.app_config) - runtime_ctx[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = frozenset(pre_existing_message_ids) incoming_metadata = config.get("metadata") if isinstance(config.get("metadata"), dict) else {} deerflow_trace_id = resolve_deerflow_trace_id(incoming_metadata.get(DEERFLOW_TRACE_METADATA_KEY)) if deerflow_trace_id: @@ -422,6 +445,33 @@ async def run_agent( else: agent = agent_factory(config=initial_runnable_config) + accessor = CheckpointStateAccessor.bind( + agent, + checkpointer, + store=store, + mode=mode, + ) + + # Capture the pre-run rollback point (materialized state + raw pending + # writes) before this run mutates the thread. Raw checkpoint blobs + # cannot reconstruct Delta-channel messages (their checkpoints omit + # channel_values), so rollback forks the pre-run lineage through the + # graph and needs the materialized messages up front. Any capture + # failure disables rollback: restoring an empty or partial message + # history would silently truncate the thread. + if checkpointer is not None: + try: + rollback_point = await _capture_rollback_point(accessor, checkpointer, checkpoint_config) + except Exception: + snapshot_capture_failed = True + logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True) + if rollback_point is not None: + pre_run_checkpoint_id = rollback_point.config.get("configurable", {}).get("checkpoint_id") + pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(rollback_point.messages)}) + + runtime_ctx[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = frozenset(pre_existing_message_ids) + _install_runtime_context(config, runtime_ctx) + # Capture the effective (resolved) model name from the agent's metadata. # _resolve_model_name in agent.py may return the default model if the # requested name is not in the allowlist — this update ensures the @@ -535,6 +585,7 @@ async def run_agent( while not record.abort_event.is_set() and not llm_error_fallback_message and (journal is None or not journal.had_llm_error_fallback): continuation_input = await _prepare_goal_continuation_input( bridge=bridge, + accessor=accessor, checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, @@ -557,11 +608,11 @@ async def run_agent( await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user") try: await _rollback_to_pre_run_checkpoint( + accessor=accessor, checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, - pre_run_checkpoint_id=pre_run_checkpoint_id, - pre_run_snapshot=pre_run_snapshot, + rollback_point=rollback_point, snapshot_capture_failed=snapshot_capture_failed, ) logger.info("Run %s rolled back to pre-run checkpoint %s", run_id, pre_run_checkpoint_id) @@ -600,11 +651,11 @@ async def run_agent( await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user") try: await _rollback_to_pre_run_checkpoint( + accessor=accessor, checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, - pre_run_checkpoint_id=pre_run_checkpoint_id, - pre_run_snapshot=pre_run_snapshot, + rollback_point=rollback_point, snapshot_capture_failed=snapshot_capture_failed, ) logger.info("Run %s was cancelled and rolled back", run_id) @@ -745,11 +796,17 @@ def _goal_instance_matches(left: GoalState | None, right: GoalState | None) -> b return same_status and same_objective and same_created_at -def _read_checkpoint_messages(checkpoint_tuple: Any) -> list[Any]: - checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} - channel_values = checkpoint.get("channel_values", {}) if isinstance(checkpoint, dict) else {} - messages = channel_values.get("messages", []) if isinstance(channel_values, dict) else [] - return messages if isinstance(messages, list) else [] +async def _materialized_checkpoint_messages(accessor: CheckpointStateAccessor, thread_id: str) -> list[Any]: + """Read ``messages`` through the mode-matched accessor. + + Raw ``channel_values`` reads see a sentinel in delta mode; only a + materialized read reconstructs the list. Raw checkpoint tuples remain + valid for tuple-level metadata (checkpoint id, ``pending_writes``). + """ + snapshot = await accessor.aget({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}) + values = getattr(snapshot, "values", None) or {} + messages = values.get("messages") if isinstance(values, dict) else None + return list(messages) if isinstance(messages, list) else [] def _read_checkpoint_goal(checkpoint_tuple: Any) -> GoalState | None: @@ -866,6 +923,7 @@ async def _reread_goal_and_checkpoint(checkpointer: Any, thread_id: str) -> tupl async def _prepare_goal_continuation_input( *, bridge: StreamBridge, + accessor: CheckpointStateAccessor, checkpointer: Any, thread_id: str, run_id: str, @@ -928,7 +986,7 @@ async def _prepare_goal_continuation_input( if checkpoint_tuple is None: return None checkpoint_id_before = _checkpoint_id(checkpoint_tuple) - messages = _read_checkpoint_messages(checkpoint_tuple) + messages = await _materialized_checkpoint_messages(accessor, thread_id) conversation_signature_before = visible_conversation_signature(messages) evidence_signature = latest_visible_assistant_signature(messages) @@ -976,7 +1034,7 @@ async def _prepare_goal_continuation_input( return None checkpoint_changed = _checkpoint_id(current_checkpoint_tuple) != checkpoint_id_before - messages_changed = visible_conversation_signature(_read_checkpoint_messages(current_checkpoint_tuple)) != conversation_signature_before + messages_changed = visible_conversation_signature(await _materialized_checkpoint_messages(accessor, thread_id)) != conversation_signature_before if checkpoint_changed or messages_changed: await _persist(current_goal, evaluation, no_progress_count, stand_down_reason="thread_changed_after_evaluation") return None @@ -1028,7 +1086,7 @@ async def _prepare_goal_continuation_input( return None if not _goal_instance_matches(updated_goal, latest_goal) or latest_checkpoint_tuple is None: return None - if visible_conversation_signature(_read_checkpoint_messages(latest_checkpoint_tuple)) != conversation_signature_before: + if visible_conversation_signature(await _materialized_checkpoint_messages(accessor, thread_id)) != conversation_signature_before: # Do not pass continuation_count here: the persist above already # committed it (as next_count). Re-passing next_count would make # _persist_goal_evaluation's race guard (#4088) see that same write as @@ -1055,65 +1113,105 @@ async def _prepare_goal_continuation_input( return {"messages": [make_goal_continuation_message(updated_goal, evaluation)]} +@dataclass(frozen=True) +class RollbackPoint: + """Materialized pre-run state used to fork the pre-run checkpoint lineage. + + Raw checkpoint blobs cannot reconstruct Delta-channel messages (their + checkpoints omit ``channel_values``), so rollback restores messages by + applying an ``Overwrite`` through a state-mutation graph anchored at the + pre-run checkpoint instead of cloning the raw blob. + """ + + config: dict[str, Any] + messages: tuple[Any, ...] + metadata: dict[str, Any] + pending_writes: tuple[tuple[str, str, Any], ...] + + +async def _capture_rollback_point( + accessor: CheckpointStateAccessor, + checkpointer: Any, + read_config: dict[str, Any], +) -> RollbackPoint | None: + """Materialize the pre-run checkpoint state and its raw pending writes. + + Returns ``None`` when the thread has no checkpoint yet; the caller keeps + the existing delete/reset rollback contract for that case. + """ + snapshot = await accessor.aget(read_config) + snapshot_config = getattr(snapshot, "config", None) or {} + configurable = snapshot_config.get("configurable") or {} + if not configurable.get("checkpoint_id"): + return None + checkpoint_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", snapshot_config) + values = getattr(snapshot, "values", None) or {} + messages = values.get("messages") if isinstance(values, dict) else None + return RollbackPoint( + config={ + "configurable": { + "thread_id": configurable.get("thread_id"), + "checkpoint_ns": configurable.get("checkpoint_ns") or "", + "checkpoint_id": configurable.get("checkpoint_id"), + } + }, + messages=tuple(messages or ()), + metadata=dict(getattr(snapshot, "metadata", None) or {}), + pending_writes=tuple(getattr(checkpoint_tuple, "pending_writes", ()) or ()), + ) + + async def _rollback_to_pre_run_checkpoint( *, + accessor: CheckpointStateAccessor | None, checkpointer: Any, thread_id: str, run_id: str, - pre_run_checkpoint_id: str | None, - pre_run_snapshot: dict[str, Any] | None, + rollback_point: RollbackPoint | None, snapshot_capture_failed: bool, ) -> None: - """Restore thread state to the checkpoint snapshot captured before run start.""" + """Fork the pre-run checkpoint lineage with the pre-run messages restored. + + The fork is written through a state-only mutation graph (the synthetic + ``rollback_restore`` node must be registered for ``as_node`` and finishes + immediately so no agent nodes are scheduled). LangGraph owns the restored + checkpoint's source/step/channel versions/parent/timestamp; the parent + pointer back to the pre-run checkpoint is the audit trail. + """ if checkpointer is None: logger.info("Run %s rollback requested but no checkpointer is configured", run_id) return if snapshot_capture_failed: - logger.warning("Run %s rollback skipped: pre-run checkpoint snapshot capture failed", run_id) + logger.warning("Run %s rollback skipped: pre-run checkpoint capture failed", run_id) return - if pre_run_snapshot is None: + if rollback_point is None: await _call_checkpointer_method(checkpointer, "adelete_thread", "delete_thread", thread_id) logger.info("Run %s rollback reset thread %s to empty state", run_id, thread_id) return - checkpoint_to_restore = None - metadata_to_restore: dict[str, Any] = {} - checkpoint_ns = "" - checkpoint = pre_run_snapshot.get("checkpoint") - if not isinstance(checkpoint, dict): - logger.warning("Run %s rollback skipped: invalid pre-run checkpoint snapshot", run_id) - return - checkpoint_to_restore = checkpoint - if checkpoint_to_restore.get("id") is None and pre_run_checkpoint_id is not None: - checkpoint_to_restore = {**checkpoint_to_restore, "id": pre_run_checkpoint_id} - if checkpoint_to_restore.get("id") is None: + configurable = rollback_point.config.get("configurable", {}) + if not configurable.get("checkpoint_id"): logger.warning("Run %s rollback skipped: pre-run checkpoint has no checkpoint id", run_id) return - restore_marker = _new_checkpoint_marker() - checkpoint_to_restore = { - **checkpoint_to_restore, - "id": restore_marker["id"], - "ts": restore_marker["ts"], - } - metadata = pre_run_snapshot.get("metadata", {}) - metadata_to_restore = metadata if isinstance(metadata, dict) else {} - raw_checkpoint_ns = pre_run_snapshot.get("checkpoint_ns") - checkpoint_ns = raw_checkpoint_ns if isinstance(raw_checkpoint_ns, str) else "" - channel_versions = checkpoint_to_restore.get("channel_versions") - new_versions = dict(channel_versions) if isinstance(channel_versions, dict) else {} + if accessor is None: + # Unreachable in practice: a rollback point can only be captured + # through the bound accessor. Stay fail-closed. + logger.warning("Run %s rollback skipped: agent accessor unavailable", run_id) + return - restore_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}} - restored_config = await _call_checkpointer_method( - checkpointer, - "aput", - "put", - restore_config, - checkpoint_to_restore, - metadata_to_restore if isinstance(metadata_to_restore, dict) else {}, - new_versions, + # The restored checkpoint inherits every channel from the pre-run fork; + # compile the mutation graph with the thread's effective schema so + # middleware-contributed channels survive (the base ThreadState fallback + # would silently drop them). + mutation_graph = build_state_mutation_graph("rollback_restore", accessor.mode, graph_state_schema(getattr(accessor, "graph", None))) + mutation_accessor = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode=accessor.mode) + restored_config = await mutation_accessor.aupdate( + rollback_point.config, + {"messages": Overwrite(list(rollback_point.messages))}, + as_node="rollback_restore", ) if not isinstance(restored_config, dict): raise RuntimeError(f"Run {run_id} rollback restore returned invalid config: expected dict") @@ -1124,7 +1222,7 @@ async def _rollback_to_pre_run_checkpoint( if not restored_checkpoint_id: raise RuntimeError(f"Run {run_id} rollback restore did not return checkpoint_id") - pending_writes = pre_run_snapshot.get("pending_writes", []) + pending_writes = rollback_point.pending_writes if not pending_writes: return @@ -1386,7 +1484,10 @@ async def _ensure_interrupted_title(*, checkpointer: Any, thread_id: str, app_co metadata["writes"] = {"runtime_interrupt_title": {"title": title}} checkpoint_ns = _checkpoint_namespace(latest_tuple) - write_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}} + # Parent to the checkpoint this write was derived from - a parentless + # raw write would sever Delta-channel replay ancestry (and truncate + # full-mode history walks). + write_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": latest_identity}} await _call_checkpointer_method( checkpointer, "aput", @@ -1522,31 +1623,14 @@ def _extract_llm_error_fallback_message(value: Any, pre_existing_ids: set[str] | return walk(value) -def _collect_pre_existing_message_ids(snapshot: dict[str, Any] | None) -> set[str]: - """Pull stable message ids out of a pre-run checkpoint snapshot. - - Used by :func:`run_agent` to mask stale ``deerflow_error_fallback`` markers - on history messages so they don't trip the current run's failure path. A - missing or malformed snapshot yields an empty set (best-effort — we - intentionally never raise from this helper). - """ - if not isinstance(snapshot, dict): +def _collect_pre_existing_message_ids(values: Any) -> set[str]: + """Collect stable message IDs from graph-materialized channel values.""" + if not isinstance(values, dict): return set() - checkpoint = snapshot.get("checkpoint") - if not isinstance(checkpoint, dict): - return set() - channel_values = checkpoint.get("channel_values") - if not isinstance(channel_values, dict): - return set() - messages = channel_values.get("messages") + messages = values.get("messages") if not isinstance(messages, (list, tuple)): return set() - ids: set[str] = set() - for msg in messages: - msg_id = _message_id(msg) - if msg_id is not None: - ids.add(msg_id) - return ids + return {message_id for message in messages if (message_id := _message_id(message)) is not None} def _unpack_stream_item( diff --git a/backend/packages/harness/pyproject.toml b/backend/packages/harness/pyproject.toml index a8349a3bc..ee29c65fc 100644 --- a/backend/packages/harness/pyproject.toml +++ b/backend/packages/harness/pyproject.toml @@ -11,13 +11,15 @@ dependencies = [ "exa-py>=1.0.0", "httpx>=0.28.0", "kubernetes>=30.0.0", - "langchain>=1.2.15", + # Lower bound reflects what the lockfile resolves and tests run against + # (langgraph 1.2.9 pulls langchain >=1.3 transitively). + "langchain>=1.3", "langchain-anthropic>=1.4.1", "langchain-deepseek>=1.0.1", "langchain-mcp-adapters>=0.2.2", "langchain-openai>=1.2.1", "langfuse>=3.4.1", - "langgraph>=1.1.9", + "langgraph>=1.2.9,<1.3", "langgraph-api>=0.8.1", "langgraph-cli>=0.4.24", "langgraph-runtime-inmem>=0.28.0", @@ -32,7 +34,7 @@ dependencies = [ "ddgs>=9.10.0", "duckdb>=1.4.4", "langchain-google-genai>=4.2.1", - "langgraph-checkpoint-sqlite>=3.0.3", + "langgraph-checkpoint-sqlite>=3.1.0,<3.2", "langgraph-sdk>=0.1.51", "sqlalchemy[asyncio]>=2.0,<3.0", "aiosqlite>=0.19", @@ -55,7 +57,7 @@ groundroute = [] ollama = ["langchain-ollama>=0.3.0"] postgres = [ "asyncpg>=0.29", - "langgraph-checkpoint-postgres>=3.0.5", + "langgraph-checkpoint-postgres>=3.1.0,<3.2", "psycopg[binary]>=3.3.3", "psycopg-pool>=3.3.0", ] diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 3d0e8216b..66bb88a22 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -34,6 +34,7 @@ browser = ["deerflow-harness[browser]"] [dependency-groups] dev = [ "blockbuster>=1.5.26,<1.6", + "hypothesis>=6.100,<7", "jsonschema>=4.26.0", "prompt-toolkit>=3.0.0", "pytest>=9.0.3", @@ -60,6 +61,18 @@ markers = [ [tool.uv] index-url = "https://pypi.org/simple" +# langgraph-sdk 0.4.2 (pulled in by langgraph 1.2.9 for DeltaChannel) pins +# `websockets<16,>=14`, silently downgrading websockets 16.0 -> 15.0.1. The +# pin is not grounded in any API incompatibility: websockets 16's only +# breaking change is requiring Python >=3.10 (we require >=3.12), the sdk +# only imports `websockets.asyncio.client`/`websockets.exceptions` (both +# 16-compatible), and DeerFlow never uses the sdk's WebSocket transport +# (httpx/SSE only). Pin the exact pre-upgrade 16.0 for the IM channel +# integrations (dingtalk-stream, python-telegram-bot, etc.) that ran on it +# before. Remove once langgraph-sdk relaxes the pin upstream. Note: enabling +# the `openai[realtime]` or `slack-sdk[optional]` extras would conflict (they +# also cap websockets<16). +override-dependencies = ["websockets==16.0"] [tool.uv.workspace] members = ["packages/harness"] diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index a08a0da4a..286f03d9c 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -93,6 +93,22 @@ def _reset_skill_storage_singleton(): reset_skill_storage() +@pytest.fixture(autouse=True) +def _reset_frozen_checkpoint_channel_mode(monkeypatch): + """Reset the process-global frozen checkpoint channel mode between tests. + + Production treats ``checkpoint_channel_mode`` as restart-required: the + first client/app freezes it for the process. The test suite builds many + clients and apps with different modes in one process, so the freeze must + not leak across tests. Mirrors the per-test ``monkeypatch.setattr`` + resets already used in test_client.py / test_lead_agent_model_resolution.py. + """ + from deerflow.runtime import checkpoint_mode + + monkeypatch.setattr(checkpoint_mode, "_frozen_checkpoint_channel_mode", None) + yield + + @pytest.fixture(autouse=True) def _restore_title_config_singleton(): """Reset ``_title_config`` to its pristine default after every test. diff --git a/backend/tests/test_app_config_reload.py b/backend/tests/test_app_config_reload.py index 4ef30cced..3aa3a464a 100644 --- a/backend/tests/test_app_config_reload.py +++ b/backend/tests/test_app_config_reload.py @@ -13,6 +13,7 @@ from deerflow.config.acp_config import load_acp_config_from_dict from deerflow.config.agents_api_config import get_agents_api_config, load_agents_api_config_from_dict from deerflow.config.app_config import AppConfig, get_app_config, reset_app_config from deerflow.config.checkpointer_config import get_checkpointer_config, load_checkpointer_config_from_dict +from deerflow.config.database_config import DatabaseConfig from deerflow.config.guardrails_config import get_guardrails_config, load_guardrails_config_from_dict from deerflow.config.memory_config import get_memory_config, load_memory_config_from_dict from deerflow.config.stream_bridge_config import get_stream_bridge_config, load_stream_bridge_config_from_dict @@ -104,6 +105,20 @@ def _write_extensions_config(path: Path) -> None: path.write_text(json.dumps({"mcpServers": {}, "skills": {}}), encoding="utf-8") +def test_checkpoint_channel_mode_defaults_to_full() -> None: + assert DatabaseConfig().checkpoint_channel_mode == "full" + + +@pytest.mark.parametrize("mode", ["full", "delta"]) +def test_checkpoint_channel_mode_accepts_supported_values(mode: str) -> None: + assert DatabaseConfig(checkpoint_channel_mode=mode).checkpoint_channel_mode == mode + + +def test_checkpoint_channel_mode_rejects_unknown_value() -> None: + with pytest.raises(ValidationError): + DatabaseConfig(checkpoint_channel_mode="auto") + + def test_config_example_does_not_enable_empty_extensions_block_by_default(): config_example_path = Path(__file__).resolve().parents[2] / "config.example.yaml" diff --git a/backend/tests/test_checkpoint_mode.py b/backend/tests/test_checkpoint_mode.py new file mode 100644 index 000000000..66923263d --- /dev/null +++ b/backend/tests/test_checkpoint_mode.py @@ -0,0 +1,286 @@ +import os +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from deerflow.runtime.checkpoint_mode import ( + CHECKPOINT_MODE_METADATA_KEY, + CheckpointModeMismatchError, + aensure_checkpoint_mode_compatible, + checkpoint_metadata_uses_delta, + ensure_checkpoint_mode_compatible, + inject_checkpoint_mode, +) + + +def test_process_mode_change_requires_restart(monkeypatch: pytest.MonkeyPatch) -> None: + from deerflow.runtime import checkpoint_mode + + monkeypatch.setattr(checkpoint_mode, "_frozen_checkpoint_channel_mode", None) + assert checkpoint_mode.freeze_checkpoint_channel_mode("full") == "full" + with pytest.raises(checkpoint_mode.CheckpointModeReconfigurationError, match="restart"): + checkpoint_mode.freeze_checkpoint_channel_mode("delta") + + +def _config() -> dict: + return {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} + + +def test_inject_delta_mode_sets_internal_key_and_metadata_marker() -> None: + config = _config() + inject_checkpoint_mode(config, "delta") + assert config["configurable"]["__deerflow_checkpoint_channel_mode"] == "delta" + assert config["metadata"][CHECKPOINT_MODE_METADATA_KEY] == "delta" + + +def test_inject_full_mode_does_not_claim_delta_metadata() -> None: + config = _config() + inject_checkpoint_mode(config, "full") + assert config["configurable"]["__deerflow_checkpoint_channel_mode"] == "full" + assert CHECKPOINT_MODE_METADATA_KEY not in config.get("metadata", {}) + + +def test_sync_full_mode_rejects_delta_marker() -> None: + saver = MagicMock() + saver.get_tuple.return_value = SimpleNamespace( + metadata={CHECKPOINT_MODE_METADATA_KEY: "delta"}, + checkpoint={"channel_values": {}}, + ) + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + ensure_checkpoint_mode_compatible(saver, _config(), "full") + saver.put.assert_not_called() + + +@pytest.mark.anyio +async def test_async_full_mode_rejects_langgraph_delta_counters() -> None: + saver = AsyncMock() + saver.aget_tuple.return_value = SimpleNamespace( + metadata={"counters_since_delta_snapshot": {"messages": (1, 1)}}, + checkpoint={"channel_values": {}}, + ) + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + await aensure_checkpoint_mode_compatible(saver, _config(), "full") + saver.aput.assert_not_awaited() + + +@pytest.mark.anyio +async def test_delta_mode_accepts_plain_full_checkpoint() -> None: + saver = AsyncMock() + saver.aget_tuple.return_value = SimpleNamespace( + metadata={}, + checkpoint={"channel_values": {"messages": ["legacy"]}}, + ) + await aensure_checkpoint_mode_compatible(saver, _config(), "delta") + + +@pytest.mark.anyio +async def test_full_mode_accessor_rejects_real_delta_checkpoint_on_sqlite(tmp_path) -> None: + """Fail-closed gate against a real saver, not mocks. + + Seeds a delta checkpoint (marker + LangGraph delta counters) into a real + AsyncSqliteSaver, reopens the thread with a full-mode accessor, and + asserts every accessor surface raises before state is read or written. + This is the integration boundary the corruption-prevention design relies + on; the mock-based tests above cannot catch a wiring regression between + the accessor, the marker, and a real backend. + """ + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver + from langgraph.types import Overwrite + + from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph + + config = _config() + async with AsyncSqliteSaver.from_conn_string(str(tmp_path / "gate.db")) as saver: + await saver.setup() + delta_accessor = CheckpointStateAccessor.bind(build_state_mutation_graph("seed", "delta"), saver, mode="delta") + await delta_accessor.aupdate( + config, + {"messages": Overwrite([HumanMessage(content="hi", id="h1")])}, + as_node="seed", + ) + + # Sanity: the seeded head really is a delta checkpoint. + seeded = await saver.aget_tuple(config) + assert checkpoint_metadata_uses_delta(seeded.metadata) + + full_accessor = CheckpointStateAccessor.bind(build_state_mutation_graph("read", "full"), saver, mode="full") + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + await full_accessor.aget(config) + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + await full_accessor.aupdate(config, {"title": "x"}, as_node="read") + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + await full_accessor.ahistory(config) + + +def test_yaml_mode_change_is_rejected_when_graph_is_reconstructed(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + from deerflow.agents.lead_agent import agent as lead_agent + from deerflow.config.app_config import reset_app_config + from deerflow.runtime import checkpoint_mode + + config_path = tmp_path / "config.yaml" + + def write_config(mode: str) -> None: + config_path.write_text( + "\n".join( + ( + "sandbox:", + " use: deerflow.sandbox.local.provider:LocalSandboxProvider", + "database:", + f" checkpoint_channel_mode: {mode}", + ) + ) + + "\n", + encoding="utf-8", + ) + + write_config("full") + monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(config_path)) + monkeypatch.setattr(checkpoint_mode, "_frozen_checkpoint_channel_mode", None) + monkeypatch.setattr(lead_agent, "_make_lead_agent", lambda config, *, app_config: object()) + reset_app_config() + try: + lead_agent.make_lead_agent({"configurable": {}}) + + write_config("delta") + future_mtime = config_path.stat().st_mtime + 5 + os.utime(config_path, (future_mtime, future_mtime)) + + with pytest.raises(checkpoint_mode.CheckpointModeReconfigurationError, match="restart"): + lead_agent.make_lead_agent({"configurable": {}}) + finally: + reset_app_config() + + +@pytest.mark.asyncio +async def test_gateway_runtime_rejects_mode_different_from_frozen_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from contextlib import asynccontextmanager + + from fastapi import FastAPI + + from app.gateway.deps import langgraph_runtime + from deerflow.runtime import checkpoint_mode + + @asynccontextmanager + async def resource(_config): + yield object() + + async def noop(*_args, **_kwargs): + return None + + monkeypatch.setattr( + checkpoint_mode, + "_frozen_checkpoint_channel_mode", + "full", + ) + monkeypatch.setattr( + "deerflow.runtime.checkpointer.async_provider.make_checkpointer", + resource, + ) + monkeypatch.setattr("deerflow.runtime.make_stream_bridge", resource) + monkeypatch.setattr("deerflow.runtime.make_store", resource) + monkeypatch.setattr( + "deerflow.persistence.engine.init_engine_from_config", + noop, + ) + monkeypatch.setattr("deerflow.persistence.engine.close_engine", noop) + monkeypatch.setattr( + "deerflow.persistence.engine.get_session_factory", + lambda: None, + ) + monkeypatch.setattr( + "deerflow.runtime.events.store.make_run_event_store", + lambda _config: object(), + ) + monkeypatch.setattr( + "deerflow.persistence.thread_meta.make_thread_store", + lambda _session_factory, _store: object(), + ) + startup_config = SimpleNamespace( + database=SimpleNamespace( + backend="memory", + checkpoint_channel_mode="delta", + ), + run_events=None, + ) + + with pytest.raises( + checkpoint_mode.CheckpointModeReconfigurationError, + match="restart", + ): + async with langgraph_runtime(FastAPI(), startup_config): + pass + + +def test_direct_langgraph_request_cannot_select_delta_in_full_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from deerflow.agents.lead_agent import agent as lead_agent + from deerflow.config.app_config import AppConfig + from deerflow.runtime import checkpoint_mode + from deerflow.runtime.checkpoint_mode import INTERNAL_CHECKPOINT_MODE_KEY + + app_config = AppConfig.model_validate( + { + "sandbox": {"use": "deerflow.sandbox.local.provider:LocalSandboxProvider"}, + "database": {"checkpoint_channel_mode": "full"}, + } + ) + config = { + "configurable": { + INTERNAL_CHECKPOINT_MODE_KEY: "delta", + } + } + monkeypatch.setattr(checkpoint_mode, "_frozen_checkpoint_channel_mode", None) + monkeypatch.setattr(lead_agent, "get_app_config", lambda: app_config) + monkeypatch.setattr( + lead_agent, + "_make_lead_agent", + lambda config, *, app_config: object(), + ) + + lead_agent.make_lead_agent(config) + + assert checkpoint_mode._frozen_checkpoint_channel_mode == "full" + assert config["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "full" + assert CHECKPOINT_MODE_METADATA_KEY not in config["metadata"] + + +def test_gateway_runtime_app_config_can_supply_its_frozen_internal_mode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from deerflow.agents.lead_agent import agent as lead_agent + from deerflow.config.app_config import AppConfig + from deerflow.runtime import checkpoint_mode + from deerflow.runtime.checkpoint_mode import INTERNAL_CHECKPOINT_MODE_KEY + + reloaded_app_config = AppConfig.model_validate( + { + "sandbox": {"use": "deerflow.sandbox.local.provider:LocalSandboxProvider"}, + "database": {"checkpoint_channel_mode": "delta"}, + } + ) + config = { + "configurable": { + INTERNAL_CHECKPOINT_MODE_KEY: "full", + }, + "context": {"app_config": reloaded_app_config}, + } + monkeypatch.setattr( + checkpoint_mode, + "_frozen_checkpoint_channel_mode", + "full", + ) + monkeypatch.setattr( + lead_agent, + "_make_lead_agent", + lambda config, *, app_config: object(), + ) + + lead_agent.make_lead_agent(config) + + assert config["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "full" + assert CHECKPOINT_MODE_METADATA_KEY not in config["metadata"] diff --git a/backend/tests/test_checkpoint_state.py b/backend/tests/test_checkpoint_state.py new file mode 100644 index 000000000..377d3739d --- /dev/null +++ b/backend/tests/test_checkpoint_state.py @@ -0,0 +1,306 @@ +from __future__ import annotations + +from copy import deepcopy +from types import SimpleNamespace +from typing import Annotated, Any, TypedDict + +import pytest +from langchain_core.messages import AnyMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph.message import add_messages + +from deerflow.runtime import CheckpointStateAccessor +from deerflow.runtime.checkpoint_mode import CHECKPOINT_MODE_METADATA_KEY, INTERNAL_CHECKPOINT_MODE_KEY + + +class FakeCheckpointer: + def __init__(self) -> None: + self.sync_configs: list[dict[str, Any]] = [] + self.async_configs: list[dict[str, Any]] = [] + + def get_tuple(self, config: dict[str, Any]) -> None: + self.sync_configs.append(config) + return None + + async def aget_tuple(self, config: dict[str, Any]) -> None: + self.async_configs.append(config) + return None + + +class FakeGraph: + def __init__(self) -> None: + self.checkpointer: Any = None + self.store: Any = None + self.calls: list[tuple[Any, ...]] = [] + self.sync_history_yields = 0 + self.async_history_yields = 0 + + def get_state(self, config: dict[str, Any]) -> SimpleNamespace: + self.calls.append(("get", config)) + return SimpleNamespace(values={"messages": ["sync"]}) + + def get_state_history(self, config: dict[str, Any], *, limit: int | None = None): + self.calls.append(("history", config, limit)) + for index in range(4): + if limit is not None and self.sync_history_yields >= limit: + return + self.sync_history_yields += 1 + yield SimpleNamespace(values={"index": index}) + + def update_state(self, config: dict[str, Any], values: dict[str, Any], *, as_node: str | None = None) -> dict[str, Any]: + self.calls.append(("update", config, values, as_node)) + return {"updated": values, "as_node": as_node} + + async def aget_state(self, config: dict[str, Any]) -> SimpleNamespace: + self.calls.append(("aget", config)) + return SimpleNamespace(values={"messages": ["async"]}) + + async def aget_state_history(self, config: dict[str, Any], *, limit: int | None = None): + self.calls.append(("ahistory", config, limit)) + for index in range(4): + if limit is not None and self.async_history_yields >= limit: + return + self.async_history_yields += 1 + yield SimpleNamespace(values={"index": index}) + + async def aupdate_state( + self, + config: dict[str, Any], + values: dict[str, Any], + *, + as_node: str | None = None, + ) -> dict[str, Any]: + self.calls.append(("aupdate", config, values, as_node)) + return {"updated": values, "as_node": as_node} + + +def _assert_delta_config_is_copied(original: dict[str, Any], forwarded: dict[str, Any]) -> None: + assert forwarded is not original + assert forwarded["configurable"] is not original["configurable"] + assert forwarded["metadata"] is not original["metadata"] + assert forwarded["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "delta" + assert forwarded["metadata"][CHECKPOINT_MODE_METADATA_KEY] == "delta" + + +def test_sync_accessor_binds_persistence_guards_operations_and_preserves_input() -> None: + graph = FakeGraph() + saver = FakeCheckpointer() + store = object() + accessor = CheckpointStateAccessor.bind(graph, saver, store=store, mode="delta") + config = { + "configurable": {"thread_id": "thread-sync", "checkpoint_ns": ""}, + "metadata": {"caller": "test"}, + "tags": ["preserved"], + } + original = deepcopy(config) + + snapshot = accessor.get(config) + history = accessor.history(config, limit=2) + update = accessor.update(config, {"messages": ["changed"]}, as_node="tools") + + assert snapshot.values == {"messages": ["sync"]} + assert [item.values for item in history] == [{"index": 0}, {"index": 1}] + assert graph.sync_history_yields == 2 + assert update == {"updated": {"messages": ["changed"]}, "as_node": "tools"} + assert graph.checkpointer is saver + assert graph.store is store + assert config == original + for call in graph.calls: + _assert_delta_config_is_copied(config, call[1]) + assert graph.calls[-1][2:] == ({"messages": ["changed"]}, "tools") + + +@pytest.mark.anyio +async def test_async_accessor_binds_persistence_guards_operations_and_preserves_input() -> None: + graph = FakeGraph() + saver = FakeCheckpointer() + store = object() + accessor = CheckpointStateAccessor.bind(graph, saver, store=store, mode="delta") + config = { + "configurable": {"thread_id": "thread-async", "checkpoint_ns": ""}, + "metadata": {"caller": "test"}, + "tags": ["preserved"], + } + original = deepcopy(config) + + snapshot = await accessor.aget(config) + history = await accessor.ahistory(config, limit=2) + update = await accessor.aupdate(config, {"messages": ["changed"]}, as_node="agent") + + assert snapshot.values == {"messages": ["async"]} + assert [item.values for item in history] == [{"index": 0}, {"index": 1}] + assert graph.async_history_yields == 2 + assert update == {"updated": {"messages": ["changed"]}, "as_node": "agent"} + assert graph.checkpointer is saver + assert graph.store is store + assert config == original + for call in graph.calls: + _assert_delta_config_is_copied(config, call[1]) + assert graph.calls[-1][2:] == ({"messages": ["changed"]}, "agent") + + +def test_sync_history_zero_limit_guards_without_consuming_a_snapshot() -> None: + graph = FakeGraph() + saver = FakeCheckpointer() + accessor = CheckpointStateAccessor.bind(graph, saver) + config = {"configurable": {"thread_id": "thread-sync-zero"}} + + assert accessor.history(config, limit=0) == [] + # The read-side gate folds onto the returned snapshot: no standalone tuple fetch. + assert len(saver.sync_configs) == 0 + assert graph.sync_history_yields == 0 + + +@pytest.mark.anyio +async def test_async_history_zero_limit_guards_without_consuming_a_snapshot() -> None: + graph = FakeGraph() + saver = FakeCheckpointer() + accessor = CheckpointStateAccessor.bind(graph, saver) + config = {"configurable": {"thread_id": "thread-async-zero"}} + + assert await accessor.ahistory(config, limit=0) == [] + assert len(saver.async_configs) == 0 + assert graph.async_history_yields == 0 + + +@pytest.mark.anyio +async def test_full_accessor_gates_writes_and_checks_reads_on_the_returned_snapshot() -> None: + """Full mode: only writes pay the pre-write tuple fetch; reads check the + marker on the materialized snapshot's metadata instead.""" + graph = FakeGraph() + saver = FakeCheckpointer() + accessor = CheckpointStateAccessor.bind(graph, saver) + config = {"configurable": {"thread_id": "thread-full"}} + + accessor.get(config) + accessor.history(config, limit=1) + accessor.update(config, {}, as_node=None) + await accessor.aget(config) + await accessor.ahistory(config, limit=1) + await accessor.aupdate(config, {}, as_node=None) + + assert len(saver.sync_configs) == 1 + assert len(saver.async_configs) == 1 + for prepared in [*saver.sync_configs, *saver.async_configs]: + assert prepared["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "full" + assert CHECKPOINT_MODE_METADATA_KEY not in prepared["metadata"] + assert config == {"configurable": {"thread_id": "thread-full"}} + + +@pytest.mark.anyio +async def test_full_accessor_raises_when_the_returned_snapshot_is_delta_marked() -> None: + """A full-mode accessor must fail closed on a delta checkpoint, detected + via the returned snapshot metadata (no pre-read tuple fetch).""" + from deerflow.runtime.checkpoint_mode import CheckpointModeMismatchError + + graph = FakeGraph() + saver = FakeCheckpointer() + accessor = CheckpointStateAccessor.bind(graph, saver, mode="full") + config = {"configurable": {"thread_id": "thread-delta-marked"}} + graph.get_state = lambda _config: SimpleNamespace(values={}, metadata={CHECKPOINT_MODE_METADATA_KEY: "delta"}) + + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + accessor.get(config) + assert len(saver.sync_configs) == 0 + + async def _delta_history(_config, *, limit=None): + yield SimpleNamespace(values={"index": 0}, metadata={}) + yield SimpleNamespace(values={"index": 1}, metadata={CHECKPOINT_MODE_METADATA_KEY: "delta"}) + + graph.aget_state_history = _delta_history + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + await accessor.ahistory(config) + assert len(saver.async_configs) == 0 + + +@pytest.mark.anyio +async def test_full_accessor_writes_still_check_compatibility_before_writing() -> None: + """Writes cannot be un-applied, so the pre-write tuple fetch stays.""" + from deerflow.runtime.checkpoint_mode import CheckpointModeMismatchError + + graph = FakeGraph() + + class DeltaMarkedSaver(FakeCheckpointer): + async def aget_tuple(self, config): + self.async_configs.append(config) + return SimpleNamespace(metadata={CHECKPOINT_MODE_METADATA_KEY: "delta"}) + + saver = DeltaMarkedSaver() + accessor = CheckpointStateAccessor.bind(graph, saver, mode="full") + config = {"configurable": {"thread_id": "thread-delta-write"}} + + with pytest.raises(CheckpointModeMismatchError, match="requires delta mode"): + await accessor.aupdate(config, {"messages": []}, as_node=None) + assert len(saver.async_configs) == 1 + assert graph.calls == [] + + +class _CountingSaver(InMemorySaver): + """InMemorySaver that counts checkpoint round-trips.""" + + def __init__(self) -> None: + super().__init__() + self.aget_tuple_calls = 0 + self.alist_limits: list[int | None] = [] + + async def aget_tuple(self, config): + self.aget_tuple_calls += 1 + return await super().aget_tuple(config) + + async def alist(self, config, *, filter=None, before=None, limit=None): + self.alist_limits.append(limit) + async for item in super().alist(config, filter=filter, before=before, limit=limit): + yield item + + +def _build_counting_graph(saver): + from langchain_core.messages import HumanMessage + from langgraph.graph import StateGraph + + class _State(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + async def _append(state): + return {"messages": [HumanMessage(content=f"turn-{len(state.get('messages') or [])}")]} + + builder = StateGraph(_State) + builder.add_node("append", _append) + builder.set_entry_point("append") + builder.set_finish_point("append") + return builder.compile(checkpointer=saver) + + +@pytest.mark.anyio +async def test_ahistory_pushes_limit_into_alist_and_reads_each_snapshot_once() -> None: + """The history limit must reach ``checkpointer.alist`` (SQL LIMIT), and the + read-side compat gate must not add a standalone tuple fetch per call.""" + saver = _CountingSaver() + graph = _build_counting_graph(saver) + accessor = CheckpointStateAccessor.bind(graph, saver, mode="full") + config = {"configurable": {"thread_id": "thread-counted"}} + for _ in range(4): + await graph.ainvoke({}, config) + + saver.aget_tuple_calls = 0 + history = await accessor.ahistory(config, limit=2) + + assert len(history) == 2 + assert saver.alist_limits[-1] == 2 + # get_state_history walks via alist only; no aget_tuple in the read path. + assert saver.aget_tuple_calls == 0 + + +@pytest.mark.anyio +async def test_aget_fetches_the_checkpoint_exactly_once_in_full_mode() -> None: + """Full-mode reads: one fetch inside aget_state; the folded gate adds none.""" + saver = _CountingSaver() + graph = _build_counting_graph(saver) + accessor = CheckpointStateAccessor.bind(graph, saver, mode="full") + config = {"configurable": {"thread_id": "thread-counted-get"}} + await graph.ainvoke({}, config) + + saver.aget_tuple_calls = 0 + snapshot = await accessor.aget(config) + + assert [message.content for message in snapshot.values["messages"]] == ["turn-0"] + assert saver.aget_tuple_calls == 1 diff --git a/backend/tests/test_checkpointer.py b/backend/tests/test_checkpointer.py index 753f1053f..2befb2ad6 100644 --- a/backend/tests/test_checkpointer.py +++ b/backend/tests/test_checkpointer.py @@ -196,7 +196,7 @@ class TestHarnessPackaging: assert "postgres" in optional_dependencies assert optional_dependencies["postgres"] == [ "asyncpg>=0.29", - "langgraph-checkpoint-postgres>=3.0.5", + "langgraph-checkpoint-postgres>=3.1.0,<3.2", "psycopg[binary]>=3.3.3", "psycopg-pool>=3.3.0", ] @@ -225,7 +225,7 @@ class TestGetCheckpointer: """get_checkpointer should return InMemorySaver when not configured.""" from langgraph.checkpoint.memory import InMemorySaver - with patch("deerflow.config.app_config.get_app_config", side_effect=FileNotFoundError): + with patch("deerflow.runtime.checkpointer.provider.get_app_config", side_effect=FileNotFoundError): cp = get_checkpointer() assert cp is not None assert isinstance(cp, InMemorySaver) diff --git a/backend/tests/test_client.py b/backend/tests/test_client.py index ed08049eb..67c7bfd42 100644 --- a/backend/tests/test_client.py +++ b/backend/tests/test_client.py @@ -18,6 +18,8 @@ from app.gateway.routers.models import ModelResponse, ModelsListResponse from app.gateway.routers.skills import SkillInstallResponse, SkillResponse, SkillsListResponse from app.gateway.routers.threads import ThreadGoalResponse from app.gateway.routers.uploads import UploadResponse +from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware +from deerflow.agents.thread_state import DeltaThreadState, ThreadState from deerflow.client import DeerFlowClient from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig from deerflow.config.paths import Paths @@ -45,6 +47,7 @@ def mock_app_config(): config.skills.deferred_discovery = False config.skills.container_path = "/mnt/skills" config.tool_search.enabled = False + config.database.checkpoint_channel_mode = "full" return config @@ -119,6 +122,24 @@ class TestClientInit: c = DeerFlowClient(checkpointer=cp) assert c._checkpointer is cp + def test_process_mode_is_frozen_from_app_config(self, mock_app_config, monkeypatch: pytest.MonkeyPatch): + from deerflow.runtime import checkpoint_mode + + monkeypatch.setattr(checkpoint_mode, "_frozen_checkpoint_channel_mode", None) + with patch("deerflow.client.get_app_config", return_value=mock_app_config): + client = DeerFlowClient() + assert client._checkpoint_channel_mode == "full" + + mock_app_config.database.checkpoint_channel_mode = "delta" + with ( + patch("deerflow.client.get_app_config", return_value=mock_app_config), + pytest.raises( + checkpoint_mode.CheckpointModeReconfigurationError, + match="restart", + ), + ): + DeerFlowClient() + # --------------------------------------------------------------------------- # list_models / list_skills / get_memory @@ -281,6 +302,68 @@ class TestStream: assert call_kwargs["context"]["thread_id"] == "t1" assert call_kwargs["context"]["agent_name"] == "test-agent-1" + def test_full_mode_overwrites_internal_delta_before_agent_creation(self, client): + from deerflow.runtime.checkpoint_mode import ( + CHECKPOINT_MODE_METADATA_KEY, + INTERNAL_CHECKPOINT_MODE_KEY, + ) + + config = { + "configurable": { + "thread_id": "t-mode", + INTERNAL_CHECKPOINT_MODE_KEY: "delta", + }, + "metadata": {CHECKPOINT_MODE_METADATA_KEY: "delta"}, + } + checkpointer = MagicMock() + checkpointer.get_tuple.return_value = None + client._checkpointer = checkpointer + agent = _make_agent_mock([]) + + with ( + patch.object(client, "_get_runnable_config", return_value=config), + patch.object(client, "_ensure_agent"), + patch.object(client, "_agent", agent), + ): + list(client.stream("hi", thread_id="t-mode")) + + assert config["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "full" + assert CHECKPOINT_MODE_METADATA_KEY not in config["metadata"] + checkpointer.get_tuple.assert_called_once_with( + { + "configurable": { + "thread_id": "t-mode", + "checkpoint_ns": "", + } + } + ) + + def test_full_mode_rejects_delta_before_agent_creation(self, client): + from types import SimpleNamespace + + from deerflow.runtime.checkpoint_mode import ( + CHECKPOINT_MODE_METADATA_KEY, + CheckpointModeMismatchError, + ) + + checkpointer = MagicMock() + checkpointer.get_tuple.return_value = SimpleNamespace( + metadata={CHECKPOINT_MODE_METADATA_KEY: "delta"}, + checkpoint={"channel_values": {}}, + ) + client._checkpointer = checkpointer + agent = _make_agent_mock([]) + + with ( + patch.object(client, "_ensure_agent") as ensure_agent, + patch.object(client, "_agent", agent), + pytest.raises(CheckpointModeMismatchError, match="requires delta mode"), + ): + list(client.stream("hi", thread_id="t-delta")) + + ensure_agent.assert_not_called() + agent.stream.assert_not_called() + def test_stream_assigns_unique_run_id_per_call(self, client): """Each embedded client stream call has a run identity for per-run middleware.""" agent = MagicMock() @@ -947,7 +1030,7 @@ class TestEnsureAgent: with ( patch("deerflow.client.create_chat_model"), - patch("deerflow.client.create_agent", return_value=mock_agent), + patch("deerflow.client.create_agent", return_value=mock_agent) as mock_create_agent, patch("deerflow.client.build_middlewares", return_value=[]) as mock_build_middlewares, patch("deerflow.client.apply_prompt_template", return_value="prompt") as mock_apply_prompt, patch("deerflow.client.get_enabled_skills_for_config", return_value=[]), @@ -965,6 +1048,30 @@ class TestEnsureAgent: mock_apply_prompt.assert_called_once() assert mock_apply_prompt.call_args.kwargs.get("agent_name") == "custom-agent" assert mock_apply_prompt.call_args.kwargs.get("available_skills") == {"test_skill"} + assert mock_create_agent.call_args.kwargs["state_schema"] is ThreadState + + def test_delta_mode_selects_state_and_normalizes_middleware(self, client): + mock_agent = MagicMock() + middleware = ViewImageMiddleware() + original_schema = middleware.state_schema + client._checkpoint_channel_mode = "delta" + config = client._get_runnable_config("t-delta") + + with ( + patch("deerflow.client.create_chat_model"), + patch("deerflow.client.create_agent", return_value=mock_agent) as mock_create_agent, + patch("deerflow.client.build_middlewares", return_value=[middleware]), + patch("deerflow.client.apply_prompt_template", return_value="prompt"), + patch("deerflow.client.get_enabled_skills_for_config", return_value=[]), + patch.object(client, "_get_tools", return_value=[]), + patch("deerflow.runtime.checkpointer.get_checkpointer", return_value=None), + ): + client._ensure_agent(config) + + call_kwargs = mock_create_agent.call_args.kwargs + assert call_kwargs["state_schema"] is DeltaThreadState + assert call_kwargs["middleware"][0] is not middleware + assert middleware.state_schema is original_schema def test_uses_default_checkpointer_when_available(self, client): mock_agent = MagicMock() @@ -1034,7 +1141,7 @@ class TestEnsureAgent: """_ensure_agent does not recreate if config key unchanged.""" mock_agent = MagicMock() client._agent = mock_agent - client._agent_config_key = (None, True, False, False, None, None, None, None) + client._agent_config_key = (None, True, False, False, None, None, None, None, "full") config = client._get_runnable_config("t1") client._ensure_agent(config) @@ -1281,6 +1388,18 @@ class TestThreadQueries: cp.pending_writes = pending_writes or [] return cp + @staticmethod + def _make_mock_snapshot(checkpoint_tuple): + snapshot = MagicMock() + snapshot.values = dict(checkpoint_tuple.checkpoint["channel_values"]) + snapshot.config = checkpoint_tuple.config + snapshot.parent_config = checkpoint_tuple.parent_config + snapshot.metadata = checkpoint_tuple.metadata + snapshot.next = () + snapshot.tasks = () + snapshot.created_at = checkpoint_tuple.checkpoint["ts"] + return snapshot + def test_list_threads_empty(self, client): mock_checkpointer = MagicMock() mock_checkpointer.list.return_value = [] @@ -1334,56 +1453,113 @@ class TestThreadQueries: def test_get_thread(self, client): mock_checkpointer = MagicMock() client._checkpointer = mock_checkpointer + client._agent = MagicMock() msg1 = HumanMessage(content="Hello", id="m1") msg2 = AIMessage(content="Hi there", id="m2") cp1 = self._make_mock_checkpoint_tuple("t1", "c1", "2023-01-01T10:00:00Z", messages=[msg1]) - cp2 = self._make_mock_checkpoint_tuple("t1", "c2", "2023-01-01T10:01:00Z", parent_id="c1", messages=[msg1, msg2], pending_writes=[("task_1", "messages", {"text": "pending"})]) + cp2 = self._make_mock_checkpoint_tuple( + "t1", + "c2", + "2023-01-01T10:01:00Z", + parent_id="c1", + messages=[msg1, msg2], + pending_writes=[("task_1", "messages", {"text": "pending"})], + ) cp3_no_ts = self._make_mock_checkpoint_tuple("t1", "c3", None) + snapshots = [ + self._make_mock_snapshot(cp2), + self._make_mock_snapshot(cp1), + self._make_mock_snapshot(cp3_no_ts), + ] + # get_thread collects pending_writes via one checkpointer.list walk + # instead of a get_tuple round-trip per checkpoint. + mock_checkpointer.list.return_value = [cp1, cp2, cp3_no_ts] + accessor = MagicMock() + accessor.history.return_value = snapshots - # checkpointer.list yields in reverse time or random order, test sorting - mock_checkpointer.list.return_value = [cp2, cp1, cp3_no_ts] - - result = client.get_thread("t1") - - mock_checkpointer.list.assert_called_once_with({"configurable": {"thread_id": "t1"}}) + with ( + patch.object(client, "_ensure_agent"), + patch("deerflow.client.CheckpointStateAccessor.bind", return_value=accessor), + ): + result = client.get_thread("t1") assert result["thread_id"] == "t1" checkpoints = result["checkpoints"] assert len(checkpoints) == 3 - - # None timestamp remains None but is sorted first via a fallback key assert checkpoints[0]["checkpoint_id"] == "c3" assert checkpoints[0]["ts"] is None - - # Should be sorted by timestamp globally assert checkpoints[1]["checkpoint_id"] == "c1" assert checkpoints[1]["ts"] == "2023-01-01T10:00:00Z" assert len(checkpoints[1]["values"]["messages"]) == 1 - assert checkpoints[2]["checkpoint_id"] == "c2" assert checkpoints[2]["parent_checkpoint_id"] == "c1" assert checkpoints[2]["ts"] == "2023-01-01T10:01:00Z" assert len(checkpoints[2]["values"]["messages"]) == 2 - # Verify message serialization assert checkpoints[2]["values"]["messages"][1]["content"] == "Hi there" - - # Verify pending writes assert len(checkpoints[2]["pending_writes"]) == 1 assert checkpoints[2]["pending_writes"][0]["task_id"] == "task_1" assert checkpoints[2]["pending_writes"][0]["channel"] == "messages" + def test_get_thread_uses_materialized_snapshot_values(self, client): + mock_checkpointer = MagicMock() + raw_checkpoint = self._make_mock_checkpoint_tuple( + "thread-1", + "ckpt-2", + "2026-07-18T00:00:00Z", + parent_id="ckpt-1", + ) + mock_checkpointer.list.return_value = [raw_checkpoint] + mock_checkpointer.get_tuple.return_value = raw_checkpoint + client._checkpointer = mock_checkpointer + client._agent = MagicMock() + client._ensure_agent = MagicMock() + + snapshot = MagicMock() + snapshot.values = { + "messages": [ + HumanMessage(id="h1", content="question"), + AIMessage(id="a1", content="answer"), + ] + } + snapshot.config = { + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": "ckpt-2", + } + } + snapshot.parent_config = {"configurable": {"checkpoint_id": "ckpt-1"}} + snapshot.metadata = {"step": 2} + snapshot.next = () + snapshot.tasks = () + snapshot.created_at = "2026-07-18T00:00:00Z" + accessor = MagicMock() + accessor.history.return_value = [snapshot] + + with patch("deerflow.client.CheckpointStateAccessor", create=True) as accessor_type: + accessor_type.bind.return_value = accessor + result = client.get_thread("thread-1") + + messages = result["checkpoints"][0]["values"]["messages"] + assert [message["id"] for message in messages] == ["h1", "a1"] + def test_get_thread_fallback_checkpointer(self, client): mock_checkpointer = MagicMock() - mock_checkpointer.list.return_value = [] + accessor = MagicMock() + accessor.history.return_value = [] + client._agent = MagicMock() - with patch("deerflow.runtime.checkpointer.provider.get_checkpointer", return_value=mock_checkpointer): + with ( + patch("deerflow.runtime.checkpointer.provider.get_checkpointer", return_value=mock_checkpointer), + patch.object(client, "_ensure_agent"), + patch("deerflow.client.CheckpointStateAccessor.bind", return_value=accessor), + ): result = client.get_thread("t99") assert result["thread_id"] == "t99" assert result["checkpoints"] == [] - mock_checkpointer.list.assert_called_once_with({"configurable": {"thread_id": "t99"}}) # --------------------------------------------------------------------------- diff --git a/backend/tests/test_client_langfuse_metadata.py b/backend/tests/test_client_langfuse_metadata.py index 7f3a6ba10..595a2ab9d 100644 --- a/backend/tests/test_client_langfuse_metadata.py +++ b/backend/tests/test_client_langfuse_metadata.py @@ -72,6 +72,7 @@ def _make_client(_monkeypatch, *, enhance_enabled: bool = True) -> DeerFlowClien ) client = DeerFlowClient.__new__(DeerFlowClient) client._app_config = fake_app_config + client._checkpoint_channel_mode = "full" client._extensions_config = None client._model_name = "stub-model" client._thinking_enabled = False diff --git a/backend/tests/test_config_version.py b/backend/tests/test_config_version.py index cf1caafe7..657b7a830 100644 --- a/backend/tests/test_config_version.py +++ b/backend/tests/test_config_version.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import os import tempfile from pathlib import Path @@ -125,6 +126,53 @@ def test_newer_user_version_no_warning(caplog): assert "outdated" not in caplog.text +def test_version_26_config_upgrades_to_checkpoint_channel_mode(tmp_path, caplog): + """A v26 user config must be flagged outdated and merge the new persisted field. + + `database.checkpoint_channel_mode` shipped with config_version 27; the + upgrade path must add it with the safe default (``full``) without touching + the user's existing database backend settings. Uses the repository's real + config.example.yaml and the real config-upgrade script. + """ + import subprocess + + repo_root = Path(__file__).resolve().parents[2] + example_src = repo_root / "config.example.yaml" + example_data = yaml.safe_load(example_src.read_text(encoding="utf-8")) + expected_version = example_data["config_version"] + assert expected_version > 26, "config.example.yaml must be bumped past 26 for checkpoint_channel_mode" + + config_path = tmp_path / "config.yaml" + (tmp_path / "config.example.yaml").write_text(example_src.read_text(encoding="utf-8"), encoding="utf-8") + user_config = { + "config_version": 26, + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "database": {"backend": "sqlite", "sqlite_dir": "custom-data"}, + } + config_path.write_text(yaml.dump(user_config), encoding="utf-8") + + with caplog.at_level(logging.WARNING, logger="deerflow.config.app_config"): + AppConfig._check_config_version(dict(user_config), config_path) + assert "outdated" in caplog.text + assert "(version 26)" in caplog.text + + env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)} + result = subprocess.run( + ["bash", str(repo_root / "scripts" / "config-upgrade.sh")], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + + upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert upgraded["config_version"] == expected_version + assert upgraded["database"]["checkpoint_channel_mode"] == "full" + assert upgraded["database"]["backend"] == "sqlite" + assert upgraded["database"]["sqlite_dir"] == "custom-data" + + def _load_repo_example() -> dict: """Load the real repo config.example.yaml (first-run template).""" example_path = Path(__file__).resolve().parents[2] / "config.example.yaml" diff --git a/backend/tests/test_context_compaction.py b/backend/tests/test_context_compaction.py index d8bc378b6..29c47d8a5 100644 --- a/backend/tests/test_context_compaction.py +++ b/backend/tests/test_context_compaction.py @@ -1,35 +1,47 @@ from __future__ import annotations from types import SimpleNamespace +from typing import Annotated, NotRequired, TypedDict import pytest from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph.message import add_messages +from langgraph.types import Overwrite +from app.gateway import services as gateway_services from deerflow.runtime import context_compaction +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor from deerflow.runtime.context_compaction import compact_thread_context -class _FakeCheckpointer: - def __init__(self, checkpoint: dict, metadata: dict | None = None) -> None: - self.checkpoint = checkpoint - self.metadata = metadata or {"step": 4, "created_at": "2026-07-06T00:00:00+00:00"} - self.put_args = None - - async def aget_tuple(self, config): - return SimpleNamespace( - checkpoint=self.checkpoint, - metadata=self.metadata, - config={"configurable": {"thread_id": config["configurable"]["thread_id"], "checkpoint_id": "ckpt-old", "checkpoint_ns": ""}}, +class _FakeAccessor: + def __init__(self, values: dict) -> None: + self.snapshot = SimpleNamespace( + values=values, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_id": "ckpt-old", + "checkpoint_ns": "", + } + }, + metadata={"step": 4, "created_at": "2026-07-06T00:00:00+00:00"}, ) + self.update_args = None - def get_next_version(self, current_version, _channel): - if current_version is None: - return 1 - return current_version + 1 + async def aget(self, _config): + return self.snapshot - async def aput(self, config, checkpoint, metadata, new_versions): - self.put_args = (config, checkpoint, metadata, new_versions) - return {"configurable": {"checkpoint_id": checkpoint["id"]}} + async def aupdate(self, config, values, *, as_node=None): + self.update_args = (config, values, as_node) + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "checkpoint_ns": "", + "checkpoint_id": "ckpt-compacted", + } + } class _FakeCompactionMiddleware: @@ -58,46 +70,18 @@ class _FakeCompactionMiddleware: ) -class _SyncCheckpointer: - def __init__(self, checkpoint: dict, metadata: dict | None = None) -> None: - self.checkpoint = checkpoint - self.metadata = metadata or {"step": 4, "created_at": "2026-07-06T00:00:00+00:00"} - self.put_args = None - - def get_tuple(self, config): - return SimpleNamespace( - checkpoint=self.checkpoint, - metadata=self.metadata, - config={"configurable": {"thread_id": config["configurable"]["thread_id"], "checkpoint_id": "ckpt-old", "checkpoint_ns": ""}}, - ) - - def get_next_version(self, current_version, _channel): - if current_version is None: - return 1 - return current_version + 1 - - def put(self, config, checkpoint, metadata, new_versions): - self.put_args = (config, checkpoint, metadata, new_versions) - return {"configurable": {"checkpoint_id": checkpoint["id"]}} - - -class _RejectDeepcopy: - def __deepcopy__(self, _memo): - raise AssertionError("compact_thread_context must not deepcopy unrelated channel values") - - @pytest.mark.asyncio -async def test_compact_thread_context_writes_summary_and_bumps_changed_channels(monkeypatch): +async def test_compact_thread_context_reads_materialized_state_and_overwrites_messages(monkeypatch): messages = [ HumanMessage(content="old question"), AIMessage(content="old answer"), HumanMessage(content="latest question"), ] - checkpointer = _FakeCheckpointer( + accessor = _FakeAccessor( { - "id": "ckpt-old", - "channel_values": {"messages": messages, "summary_text": "OLD SUMMARY", "sandbox": _RejectDeepcopy()}, - "channel_versions": {"messages": 7, "summary_text": 3, "title": 2}, + "messages": messages, + "summary_text": "OLD SUMMARY", + "sandbox": object(), } ) middleware = _FakeCompactionMiddleware() @@ -108,7 +92,7 @@ async def test_compact_thread_context_writes_summary_and_bumps_changed_channels( ) result = await compact_thread_context( - checkpointer, + accessor, "thread-1", app_config=SimpleNamespace(), user_id="user-1", @@ -119,22 +103,16 @@ async def test_compact_thread_context_writes_summary_and_bumps_changed_channels( assert result.removed_message_count == 2 assert result.preserved_message_count == 1 assert result.summary_updated is True + assert result.checkpoint_id == "ckpt-compacted" assert result.total_tokens == 123 - assert checkpointer.put_args is not None - _config, written_checkpoint, written_metadata, new_versions = checkpointer.put_args - assert written_checkpoint["channel_values"]["messages"] == [messages[-1]] - assert written_checkpoint["channel_values"]["summary_text"] == "COMPRESSED SUMMARY" - assert isinstance(written_checkpoint["channel_values"]["sandbox"], _RejectDeepcopy) - assert written_checkpoint["channel_versions"]["messages"] == 8 - assert written_checkpoint["channel_versions"]["summary_text"] == 4 - assert written_checkpoint["channel_versions"]["title"] == 2 - assert new_versions == {"messages": 8, "summary_text": 4} - assert written_metadata["writes"]["manual_compaction"]["messages"] == { - "removed": 2, - "preserved": 1, - } - assert "COMPRESSED SUMMARY" not in str(written_metadata["writes"]) + assert accessor.update_args is not None + update_config, written_values, as_node = accessor.update_args + assert update_config == accessor.snapshot.config + assert isinstance(written_values["messages"], Overwrite) + assert written_values["messages"].value == [messages[-1]] + assert written_values["summary_text"] == "COMPRESSED SUMMARY" + assert as_node == "manual_compaction" assert middleware.prepare_calls == 1 assert middleware.runtime_contexts == [ {"thread_id": "thread-1", "user_id": "user-1", "agent_name": "research-agent"}, @@ -142,12 +120,126 @@ async def test_compact_thread_context_writes_summary_and_bumps_changed_channels( @pytest.mark.asyncio -async def test_compact_thread_context_returns_noop_without_writing(monkeypatch): - checkpointer = _FakeCheckpointer( +async def test_compact_thread_context_real_mutation_graph_finishes_without_scheduling(monkeypatch): + messages = [ + HumanMessage(id="h1", content="old question"), + AIMessage(id="a1", content="old answer"), + HumanMessage(id="h2", content="latest question"), + ] + request = SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + checkpointer=InMemorySaver(), + checkpoint_channel_mode="delta", + store=None, + ) + ) + ) + seed_accessor, seed_config = gateway_services.build_checkpoint_state_mutation_accessor( + request, + thread_id="thread-real-compaction", + as_node="seed", + ) + await seed_accessor.aupdate( + seed_config, { - "id": "ckpt-old", - "channel_values": {"messages": [HumanMessage(content="latest only")]}, - "channel_versions": {"messages": 1}, + "messages": Overwrite(messages), + "summary_text": "OLD SUMMARY", + }, + as_node="seed", + ) + accessor, config = gateway_services.build_checkpoint_state_mutation_accessor( + request, + thread_id="thread-real-compaction", + as_node="manual_compaction", + ) + monkeypatch.setattr( + context_compaction, + "_create_compaction_middleware", + lambda **_kwargs: _FakeCompactionMiddleware(), + ) + + result = await compact_thread_context( + accessor, + "thread-real-compaction", + app_config=SimpleNamespace(), + ) + snapshot = await accessor.aget(config) + + assert result.compacted is True + assert [message.id for message in snapshot.values["messages"]] == ["h2"] + assert snapshot.values["summary_text"] == "COMPRESSED SUMMARY" + assert snapshot.next == () + + +@pytest.mark.asyncio +async def test_compact_thread_context_preserves_middleware_contributed_channels(monkeypatch): + """Compaction must not drop channels the base ThreadState does not know. + + Contract lock for fork inheritance: the compaction write carries only + messages + summary_text (base channels), so a middleware-contributed + channel survives even though the mutation graph compiles with the base + schema. If LangGraph ever stops cloning unknown channels into forked + checkpoints, this test fails before production state is lost. + """ + from deerflow.runtime.checkpoint_state import build_state_mutation_graph + + class ExtensionState(TypedDict): + messages: Annotated[list, add_messages] + memory_notes: NotRequired[str] + + messages = [ + HumanMessage(id="h1", content="old question"), + AIMessage(id="a1", content="old answer"), + HumanMessage(id="h2", content="latest question"), + ] + request = SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + checkpointer=InMemorySaver(), + checkpoint_channel_mode="full", + store=None, + ) + ) + ) + seed_graph = build_state_mutation_graph("seed", "full", ExtensionState) + seed_accessor = CheckpointStateAccessor.bind(seed_graph, request.app.state.checkpointer, mode="full") + seed_config = {"configurable": {"thread_id": "thread-ext-compaction", "checkpoint_ns": ""}} + await seed_accessor.aupdate( + seed_config, + {"messages": messages, "memory_notes": "extension-value"}, + as_node="seed", + ) + + # Production path: compact uses the base-schema mutation accessor. + accessor, config = gateway_services.build_checkpoint_state_mutation_accessor( + request, + thread_id="thread-ext-compaction", + as_node="manual_compaction", + ) + monkeypatch.setattr( + context_compaction, + "_create_compaction_middleware", + lambda **_kwargs: _FakeCompactionMiddleware(), + ) + + result = await compact_thread_context( + accessor, + "thread-ext-compaction", + app_config=SimpleNamespace(), + ) + snapshot = await seed_accessor.aget(config) + + assert result.compacted is True + assert [message.id for message in snapshot.values["messages"]] == ["h2"] + assert snapshot.values["memory_notes"] == "extension-value" + + +@pytest.mark.asyncio +async def test_compact_thread_context_returns_noop_without_writing(monkeypatch): + accessor = _FakeAccessor( + { + "messages": [HumanMessage(content="latest only")], } ) middleware = _FakeCompactionMiddleware(should_compact=False) @@ -157,35 +249,9 @@ async def test_compact_thread_context_returns_noop_without_writing(monkeypatch): lambda **_kwargs: middleware, ) - result = await compact_thread_context(checkpointer, "thread-1", app_config=SimpleNamespace()) + result = await compact_thread_context(accessor, "thread-1", app_config=SimpleNamespace()) assert result.compacted is False assert result.reason == "not_enough_messages" - assert checkpointer.put_args is None + assert accessor.update_args is None assert middleware.prepare_calls == 1 - - -@pytest.mark.asyncio -async def test_compact_thread_context_supports_sync_checkpointer_methods(monkeypatch): - messages = [ - HumanMessage(content="old question"), - AIMessage(content="old answer"), - HumanMessage(content="latest question"), - ] - checkpointer = _SyncCheckpointer( - { - "id": "ckpt-old", - "channel_values": {"messages": messages, "summary_text": "OLD SUMMARY"}, - "channel_versions": {"messages": 7, "summary_text": 3}, - } - ) - monkeypatch.setattr( - context_compaction, - "_create_compaction_middleware", - lambda **_kwargs: _FakeCompactionMiddleware(), - ) - - result = await compact_thread_context(checkpointer, "thread-1", app_config=SimpleNamespace()) - - assert result.compacted is True - assert checkpointer.put_args is not None diff --git a/backend/tests/test_create_deerflow_agent.py b/backend/tests/test_create_deerflow_agent.py index dea50f1be..700de696c 100644 --- a/backend/tests/test_create_deerflow_agent.py +++ b/backend/tests/test_create_deerflow_agent.py @@ -4,17 +4,31 @@ from typing import get_type_hints from unittest.mock import MagicMock, patch import pytest +from langchain.agents import AgentState +from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel +from langchain_core.messages import AIMessage +from langgraph.channels import DeltaChannel +from langgraph.checkpoint.memory import InMemorySaver from deerflow.agents.factory import create_deerflow_agent from deerflow.agents.features import Next, Prev, RuntimeFeatures from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware -from deerflow.agents.thread_state import ThreadState +from deerflow.agents.thread_state import DeltaThreadState, ThreadState def _make_mock_model(): return MagicMock(name="mock_model") +class _FakeModel(FakeMessagesListChatModel): + def bind_tools(self, tools, **kwargs): # type: ignore[override] + return self + + +class _CustomState(AgentState): + custom_value: str + + def _make_mock_tool(name: str = "my_tool"): tool = MagicMock(name=name) tool.name = name @@ -36,6 +50,69 @@ def test_minimal_creation(mock_create_agent): call_kwargs = mock_create_agent.call_args[1] assert call_kwargs["model"] is model assert call_kwargs["system_prompt"] is None + assert call_kwargs["state_schema"] is ThreadState + + +@patch("deerflow.agents.factory.create_agent") +def test_delta_creation_selects_delta_state_and_copies_middleware(mock_create_agent): + mock_create_agent.return_value = MagicMock(name="compiled_graph") + middleware = ViewImageMiddleware() + original_schema = middleware.state_schema + + create_deerflow_agent( + _make_mock_model(), + middleware=[middleware], + checkpoint_channel_mode="delta", + ) + + call_kwargs = mock_create_agent.call_args.kwargs + assert call_kwargs["state_schema"] is DeltaThreadState + assert call_kwargs["middleware"][0] is not middleware + assert middleware.state_schema is original_schema + + +@patch("deerflow.agents.factory.create_agent") +def test_custom_state_schema_is_preserved_in_full_mode_and_adapted_in_delta_mode(mock_create_agent): + mock_create_agent.return_value = MagicMock(name="compiled_graph") + + create_deerflow_agent(_make_mock_model(), state_schema=_CustomState) + assert mock_create_agent.call_args.kwargs["state_schema"] is _CustomState + + create_deerflow_agent( + _make_mock_model(), + state_schema=_CustomState, + checkpoint_channel_mode="delta", + ) + adapted = mock_create_agent.call_args.kwargs["state_schema"] + hints = get_type_hints(adapted, include_extras=True) + assert "custom_value" in hints + assert any(isinstance(item, DeltaChannel) for item in hints["messages"].__metadata__) + + +def test_delta_checkpointer_combination_is_rejected_before_any_persistence(): + """Mixed-mode corruption guard: delta + checkpointer must fail loudly at + construction, before any state is read or written through the ungated graph.""" + saver = InMemorySaver() + with pytest.raises(ValueError, match="checkpoint_channel_mode='delta'"): + create_deerflow_agent( + _FakeModel(responses=[AIMessage(content="ok")]), + checkpoint_channel_mode="delta", + checkpointer=saver, + ) + assert list(saver.list(None)) == [] + + +def test_compiled_factory_graph_selects_full_and_delta_message_channels(): + full_graph = create_deerflow_agent( + _FakeModel(responses=[AIMessage(id="full-response", content="done")]), + ) + delta_graph = create_deerflow_agent( + _FakeModel(responses=[AIMessage(id="delta-response", content="done")]), + checkpoint_channel_mode="delta", + ) + + assert type(full_graph.channels["messages"]).__name__ == "BinaryOperatorAggregate" + assert isinstance(delta_graph.channels["messages"], DeltaChannel) # --------------------------------------------------------------------------- diff --git a/backend/tests/test_delta_channel_checkpointers.py b/backend/tests/test_delta_channel_checkpointers.py new file mode 100644 index 000000000..e79df0410 --- /dev/null +++ b/backend/tests/test_delta_channel_checkpointers.py @@ -0,0 +1,371 @@ +"""DeltaChannel migration, replay, and storage-shape contracts on real saver backends. + +These tests pin the LangGraph 1.2 DeltaChannel upgrade path that production +relies on: + +- full -> delta migration on the same thread (old plain-value blobs seed the + delta channel transparently), +- deterministic materialization with stable message ids (ids are stamped into + the persisted writes by ``put_writes``/``ensure_message_ids``), +- the on-disk storage shape (non-snapshot checkpoints omit ``messages`` from + ``channel_values``; snapshot checkpoints store a ``_DeltaSnapshot`` blob), +- non-Delta raw writers (goal / run-duration metadata / interrupted-title + helper) preserving Delta ancestry and the downgrade markers. + +Every contract runs against InMemorySaver, AsyncSqliteSaver, and - when +``TEST_POSTGRES_URI`` is set - AsyncPostgresSaver, because each backend +implements blob/version handling slightly differently. +""" + +from __future__ import annotations + +import os +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager +from typing import Annotated, Any, TypedDict +from uuid import uuid4 + +import pytest +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, RemoveMessage +from langgraph.channels import DeltaChannel +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.serde.types import _DeltaSnapshot +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages +from langgraph.types import Overwrite + +from deerflow.agents.goal_state import GoalState +from deerflow.agents.thread_state import DeltaThreadState, merge_message_writes +from deerflow.runtime.checkpoint_mode import CHECKPOINT_MODE_METADATA_KEY, checkpoint_tuple_uses_delta +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor +from deerflow.runtime.goal import write_thread_goal +from deerflow.runtime.runs.worker import _ensure_interrupted_title, persist_run_durations + + +class FullState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +class DeltaState(TypedDict): + messages: Annotated[ + list[AnyMessage], + DeltaChannel(merge_message_writes, snapshot_frequency=2), + ] + + +def _thread_id() -> str: + return f"delta-contract-{uuid4().hex}" + + +def _config(thread_id: str) -> dict[str, Any]: + return {"configurable": {"thread_id": thread_id}} + + +def _noop(state: dict[str, Any]) -> dict[str, Any]: + return {} + + +def _build_graph(schema: Any, checkpointer: Any) -> Any: + builder = StateGraph(schema) + builder.add_node("noop", _noop) + builder.set_entry_point("noop") + builder.set_finish_point("noop") + return builder.compile(checkpointer=checkpointer) + + +def _goal(objective: str) -> GoalState: + return { + "objective": objective, + "status": "active", + "created_at": "2026-07-18T00:00:00+00:00", + "updated_at": "2026-07-18T00:00:00+00:00", + "continuation_count": 0, + "max_continuations": 3, + "no_progress_count": 0, + "max_no_progress_continuations": 2, + } + + +class _SaverEnv: + """One saver instance with a reopen() that simulates a process restart. + + Reopening swaps in a brand-new saver over the same bytes (SQLite file or + Postgres schema) so replay-after-reopen contracts prove persistence, not + in-process caching. InMemorySaver keeps no external bytes, so reopen is a + no-op stand-in there. + """ + + def __init__(self, kind: str, open_saver: Callable[[], Any]) -> None: + self.kind = kind + self._open_saver = open_saver + self._cm: Any | None = None + self.saver: Any | None = None + + async def __aenter__(self) -> _SaverEnv: + self._cm = self._open_saver() + self.saver = await self._cm.__aenter__() + setup = getattr(self.saver, "setup", None) + if setup is not None: + await setup() + return self + + async def __aexit__(self, *exc: Any) -> None: + if self._cm is not None: + await self._cm.__aexit__(*exc) + self._cm = None + self.saver = None + + async def reopen(self) -> None: + if self.kind == "memory": + return + await self._cm.__aexit__(None, None, None) + self._cm = self._open_saver() + self.saver = await self._cm.__aenter__() + + +@asynccontextmanager +async def _open_sqlite(db_path: Any) -> AsyncIterator[Any]: + async with AsyncSqliteSaver.from_conn_string(str(db_path)) as saver: + await saver.setup() + yield saver + + +@asynccontextmanager +async def _open_postgres(uri: str) -> AsyncIterator[Any]: + aio = pytest.importorskip("langgraph.checkpoint.postgres.aio", reason="postgres extra not installed") + async with aio.AsyncPostgresSaver.from_conn_string(uri) as saver: + await saver.setup() + yield saver + + +@pytest.fixture(params=["memory", "sqlite", "postgres"]) +async def saver_env(request: pytest.FixtureRequest, tmp_path: Any) -> AsyncIterator[_SaverEnv]: + kind = request.param + if kind == "memory": + saver = InMemorySaver() + + @asynccontextmanager + async def open_memory() -> AsyncIterator[Any]: + yield saver + + open_saver = open_memory + elif kind == "sqlite": + db_path = tmp_path / "delta-contract.sqlite" + + def open_sqlite() -> Any: + return _open_sqlite(db_path) + + open_saver = open_sqlite + else: + uri = os.environ.get("TEST_POSTGRES_URI") + if not uri: + pytest.skip("TEST_POSTGRES_URI is not set") + + def open_postgres() -> Any: + return _open_postgres(uri) + + open_saver = open_postgres + + async with _SaverEnv(kind, open_saver) as env: + yield env + + +@pytest.mark.anyio +async def test_full_to_delta_migration_replays_on_same_thread(saver_env: _SaverEnv) -> None: + """A thread written by a full (pre-delta) graph must keep replaying after + the process swaps to a delta graph: the old plain-value ``messages`` blob + seeds the delta channel, and later delta writes append on top.""" + thread_id = _thread_id() + config = _config(thread_id) + + full_graph = _build_graph(FullState, saver_env.saver) + await full_graph.ainvoke({"messages": [HumanMessage(id="h1", content="seed from full mode")]}, config) + + delta_graph = _build_graph(DeltaState, saver_env.saver) + delta_accessor = CheckpointStateAccessor.bind(delta_graph, saver_env.saver, mode="delta") + + migrated = await delta_accessor.aget(config) + assert [m.id for m in migrated.values["messages"]] == ["h1"] + assert migrated.values["messages"][0].content == "seed from full mode" + + await delta_graph.ainvoke({"messages": [AIMessage(id="a1", content="delta reply")]}, config) + + await saver_env.reopen() + delta_graph = _build_graph(DeltaState, saver_env.saver) + delta_accessor = CheckpointStateAccessor.bind(delta_graph, saver_env.saver, mode="delta") + + replayed = await delta_accessor.aget(config) + assert [m.id for m in replayed.values["messages"]] == ["h1", "a1"] + assert [m.content for m in replayed.values["messages"]] == ["seed from full mode", "delta reply"] + + +@pytest.mark.anyio +async def test_materialization_is_deterministic_and_message_ids_are_stable(saver_env: _SaverEnv) -> None: + """Materializing the same thread twice - and once more after a saver + reopen - must produce identical values, and the id LangGraph stamps onto + an id-less message must survive persistence so it can drive RemoveMessage.""" + thread_id = _thread_id() + config = _config(thread_id) + graph = _build_graph(DeltaState, saver_env.saver) + accessor = CheckpointStateAccessor.bind(graph, saver_env.saver, mode="delta") + + await graph.ainvoke({"messages": [HumanMessage(content="no id on purpose")]}, config) + + first = await accessor.aget(config) + second = await accessor.aget(config) + assert first.values == second.values + + persisted_id = first.values["messages"][0].id + assert persisted_id + + await saver_env.reopen() + graph = _build_graph(DeltaState, saver_env.saver) + accessor = CheckpointStateAccessor.bind(graph, saver_env.saver, mode="delta") + + reopened = await accessor.aget(config) + assert reopened.values == first.values + assert reopened.values["messages"][0].id == persisted_id + + await graph.ainvoke({"messages": [RemoveMessage(id=persisted_id)]}, config) + after_remove = await accessor.aget(config) + assert after_remove.values["messages"] == [] + + +@pytest.mark.anyio +async def test_delta_storage_shape_and_snapshot_cadence(saver_env: _SaverEnv) -> None: + """Non-snapshot checkpoints must not carry ``messages`` in channel_values + (that is the whole storage win); every ``snapshot_frequency`` updates the + checkpoint carries a ``_DeltaSnapshot`` blob instead, and materialization + stays correct across both shapes.""" + thread_id = _thread_id() + config = _config(thread_id) + graph = _build_graph(DeltaState, saver_env.saver) + accessor = CheckpointStateAccessor.bind(graph, saver_env.saver, mode="delta") + + await graph.ainvoke({"messages": [HumanMessage(id="m1", content="one")]}, config) + first_tuple = await saver_env.saver.aget_tuple(config) + assert first_tuple is not None + assert "messages" not in first_tuple.checkpoint["channel_values"] + # The node's output persists as pending writes attached to the checkpoint + # saved *before* its superstep (an ancestor of the latest checkpoint). + chain_has_messages_write = False + async for chain_tuple in saver_env.saver.alist(config): + if any(channel == "messages" for _, channel, _ in chain_tuple.pending_writes or []): + chain_has_messages_write = True + break + assert chain_has_messages_write + + # Second update crosses snapshot_frequency=2 -> one checkpoint on the + # chain carries a _DeltaSnapshot blob for messages. (Which checkpoint + # reassembles it into channel_values differs per saver: InMemorySaver + # resolves the versioned blob on the latest checkpoint, AsyncSqliteSaver + # on its parent. The cadence contract is that the snapshot exists.) + await graph.ainvoke({"messages": [AIMessage(id="m2", content="two")]}, config) + snapshot_found = False + async for chain_tuple in saver_env.saver.alist(config): + if isinstance(chain_tuple.checkpoint["channel_values"].get("messages"), _DeltaSnapshot): + snapshot_found = True + break + assert snapshot_found + + snapshot = await accessor.aget(config) + assert [m.id for m in snapshot.values["messages"]] == ["m1", "m2"] + assert [m.content for m in snapshot.values["messages"]] == ["one", "two"] + + +@pytest.mark.anyio +async def test_non_delta_writers_preserve_delta_messages_and_markers(saver_env: _SaverEnv) -> None: + """Raw checkpoint writers that never touch the messages channel (thread + goal, run-duration metadata, interrupted-title helper) must not sever the + Delta parent lineage or drop the downgrade markers. Regression: a raw + ``aput`` whose write_config omits ``checkpoint_id`` stores a parentless + checkpoint; replay from it then walks an empty ancestor chain and the + whole message history silently disappears.""" + thread_id = _thread_id() + config = _config(thread_id) + write_config = {**config, "metadata": {CHECKPOINT_MODE_METADATA_KEY: "delta"}} + graph = _build_graph(DeltaThreadState, saver_env.saver) + accessor = CheckpointStateAccessor.bind(graph, saver_env.saver, mode="delta") + + await graph.ainvoke({"messages": Overwrite([HumanMessage(id="u1", content="hello goal writer")])}, write_config) + await graph.ainvoke({"messages": [AIMessage(id="a1", content="reply")]}, write_config) + + await write_thread_goal(saver_env.saver, thread_id, _goal("keep messages alive")) + durations_written = await persist_run_durations( + checkpointer=saver_env.saver, + thread_id=thread_id, + durations={"run-1": 7}, + ) + assert durations_written + # Delta checkpoints carry no messages in channel_values, so the title + # helper has nothing to derive from and must stay inert (no write). + title = await _ensure_interrupted_title(checkpointer=saver_env.saver, thread_id=thread_id, app_config=None) + assert title is None + + await saver_env.reopen() + graph = _build_graph(DeltaThreadState, saver_env.saver) + accessor = CheckpointStateAccessor.bind(graph, saver_env.saver, mode="delta") + + latest_tuple = await saver_env.saver.aget_tuple(config) + assert latest_tuple is not None + + # Production materialized read path (same one the Gateway uses). + snapshot = await accessor.aget(config) + assert [m.id for m in snapshot.values["messages"]] == ["u1", "a1"] + assert [m.content for m in snapshot.values["messages"]] == ["hello goal writer", "reply"] + + assert checkpoint_tuple_uses_delta(latest_tuple) + assert latest_tuple.metadata.get(CHECKPOINT_MODE_METADATA_KEY) == "delta" + assert latest_tuple.metadata.get("run_durations", {}).get("run-1") == 7 + + +# --------------------------------------------------------------------------- +# InMemorySaver delta-history patch guards (deerflow.checkpoint_patches) +# --------------------------------------------------------------------------- + + +def test_inmemory_delta_history_patch_is_active() -> None: + """The compatibility patch must be applied in every test/app process.""" + from langgraph.checkpoint.memory import InMemorySaver + + from deerflow import checkpoint_patches + + assert getattr(InMemorySaver, checkpoint_patches._PATCH_FLAG, False) is True + assert InMemorySaver.get_delta_channel_history is checkpoint_patches._get_delta_channel_history_via_base + assert InMemorySaver.aget_delta_channel_history is checkpoint_patches._aget_delta_channel_history_via_base + + +def test_inmemory_delta_history_patch_stands_down_without_upstream_override(monkeypatch) -> None: + """If upstream removes its (buggy) override, the patch must not reinstall.""" + from langgraph.checkpoint.memory import InMemorySaver + + from deerflow import checkpoint_patches + + monkeypatch.setattr(checkpoint_patches, "_upstream_override_present", lambda: False) + monkeypatch.delattr(InMemorySaver, checkpoint_patches._PATCH_FLAG, raising=False) + sentinel = object() + monkeypatch.setattr(InMemorySaver, "get_delta_channel_history", sentinel) + + checkpoint_patches.ensure_inmemory_delta_history_patch() + + assert getattr(InMemorySaver, checkpoint_patches._PATCH_FLAG, False) is False + assert InMemorySaver.get_delta_channel_history is sentinel + + +def test_inmemory_delta_history_patch_warns_on_unvalidated_langgraph(monkeypatch, caplog) -> None: + """A langgraph newer than the validated version must log a re-inspection warning.""" + import logging + + from langgraph.checkpoint.memory import InMemorySaver + + from deerflow import checkpoint_patches + + monkeypatch.setattr(checkpoint_patches.importlib.metadata, "version", lambda _name: "99.0.0") + monkeypatch.setattr(checkpoint_patches, "_upstream_override_present", lambda: False) + monkeypatch.delattr(InMemorySaver, checkpoint_patches._PATCH_FLAG, raising=False) + + with caplog.at_level(logging.WARNING, logger=checkpoint_patches.__name__): + checkpoint_patches.ensure_inmemory_delta_history_patch() + + assert any("newer than the version" in record.message for record in caplog.records) diff --git a/backend/tests/test_delta_channel_state.py b/backend/tests/test_delta_channel_state.py new file mode 100644 index 000000000..d9dcc7bc3 --- /dev/null +++ b/backend/tests/test_delta_channel_state.py @@ -0,0 +1,184 @@ +from typing import get_type_hints + +import pytest +from hypothesis import given +from hypothesis import strategies as st +from langchain.agents import AgentState, create_agent +from langchain.agents.middleware import AgentMiddleware +from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel +from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, RemoveMessage +from langgraph.channels import DeltaChannel +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph import StateGraph +from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages + +from deerflow.agents.thread_state import ( + DeltaThreadState, + ThreadState, + adapt_state_schema_for_mode, + get_thread_state_schema, + merge_message_writes, + normalize_middleware_state_schemas, +) + + +def _fold(state: list, writes: list) -> list: + result = list(state) + for write in writes: + result = list(add_messages(result, write)) + return result + + +@pytest.mark.parametrize( + "writes", + [ + [[HumanMessage(id="h1", content="one")], [AIMessage(id="a1", content="two")]], + [[AIMessage(id="same", content="old")], [AIMessage(id="same", content="new")]], + [[HumanMessage(id="h1", content="one")], [RemoveMessage(id="h1")]], + [ + [HumanMessage(id="h1", content="one"), AIMessage(id="a1", content="two")], + [RemoveMessage(id=REMOVE_ALL_MESSAGES), HumanMessage(id="h2", content="kept")], + ], + ], +) +def test_merge_message_writes_matches_sequential_add_messages(writes: list) -> None: + assert merge_message_writes([], writes) == _fold([], writes) + + +@given(split=st.integers(min_value=0, max_value=3)) +def test_merge_message_writes_is_batching_invariant(split: int) -> None: + state = [HumanMessage(id="h0", content="seed")] + writes = [ + [AIMessage(id="a1", content="first")], + [AIMessage(id="a1", content="replacement")], + [HumanMessage(id="h2", content="last")], + ] + xs = writes[:split] + ys = writes[split:] + assert merge_message_writes(merge_message_writes(state, xs), ys) == merge_message_writes(state, writes) + + +def test_merge_message_writes_matches_unknown_remove_error() -> None: + writes = [[RemoveMessage(id="missing")]] + + with pytest.raises(ValueError) as expected: + _fold([], writes) + with pytest.raises(type(expected.value)) as actual: + merge_message_writes([], writes) + + assert str(actual.value) == str(expected.value) + + +@pytest.mark.parametrize( + "write", + [ + [{"role": "user", "content": "from a dict", "id": "dict-1"}], + AIMessageChunk(id="chunk-1", content="from a chunk"), + ], +) +def test_merge_message_writes_matches_message_coercion(write: object) -> None: + assert merge_message_writes([], [write]) == _fold([], [write]) + + +def test_raw_tuple_coercion_matches_add_messages_reducer_parity(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("langgraph.graph.message.uuid.uuid4", lambda: "tuple-id") + writes = [[("human", "from a tuple")]] + + assert merge_message_writes([], writes) == _fold([], writes) + + +def test_mode_selects_expected_state_schema() -> None: + assert get_thread_state_schema("full") is ThreadState + assert get_thread_state_schema("delta") is DeltaThreadState + message_hint = get_type_hints(DeltaThreadState, include_extras=True)["messages"] + assert any(isinstance(item, DeltaChannel) for item in message_hint.__metadata__) + + +def test_delta_adaptation_replaces_agent_state_message_reducer() -> None: + adapted = adapt_state_schema_for_mode(AgentState, "delta") + hint = get_type_hints(adapted, include_extras=True)["messages"] + assert any(isinstance(item, DeltaChannel) for item in hint.__metadata__) + + +def test_agents_package_exports_delta_thread_state() -> None: + from deerflow.agents import DeltaThreadState as ExportedDeltaThreadState + + assert ExportedDeltaThreadState is DeltaThreadState + + +class _FirstState(AgentState): + first: str + + +class _SecondState(AgentState): + second: int + + +class _FirstMiddleware(AgentMiddleware): + state_schema = _FirstState + + +class _SecondMiddleware(AgentMiddleware): + state_schema = _SecondState + + +class _FakeModel(FakeMessagesListChatModel): + def bind_tools(self, tools, **kwargs): # type: ignore[override] + return self + + +def _compile_with_middleware(middleware: list[AgentMiddleware], mode: str): + return create_agent( + model=_FakeModel(responses=[AIMessage(id="response", content="done")]), + tools=None, + middleware=normalize_middleware_state_schemas(middleware, mode), + state_schema=get_thread_state_schema(mode), + ) + + +def test_delta_normalization_compiles_stable_channel_without_mutating_middleware() -> None: + first = _FirstMiddleware() + second = _SecondMiddleware() + middleware = [first, second] + + for _ in range(10): + graph = _compile_with_middleware(middleware, "delta") + assert isinstance(graph.channels["messages"], DeltaChannel) + + assert first.state_schema is _FirstState + assert second.state_schema is _SecondState + + full_graph = _compile_with_middleware(middleware, "full") + assert type(full_graph.channels["messages"]).__name__ == "BinaryOperatorAggregate" + assert first.state_schema is _FirstState + assert second.state_schema is _SecondState + + +@pytest.mark.parametrize( + "write", + [ + HumanMessage(content="root BaseMessage"), + {"role": "user", "content": "root message dict"}, + [HumanMessage(content="BaseMessage in list")], + [{"role": "user", "content": "message dict in list"}], + ], + ids=["base-message", "dict", "base-message-list", "dict-list"], +) +def test_production_message_forms_keep_assigned_ids_across_delta_replay(write: object) -> None: + builder = StateGraph(DeltaThreadState) + + def write_messages(_state): + return {"messages": write} + + builder.add_node("writer", write_messages) + builder.set_entry_point("writer") + builder.set_finish_point("writer") + graph = builder.compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "stable-message-replay"}} + + graph.invoke({}, config) + first = graph.get_state(config).values["messages"] + second = graph.get_state(config).values["messages"] + + assert first[0].id is not None + assert second[0].id == first[0].id diff --git a/backend/tests/test_gateway_checkpoint_mode.py b/backend/tests/test_gateway_checkpoint_mode.py new file mode 100644 index 000000000..7189432f9 --- /dev/null +++ b/backend/tests/test_gateway_checkpoint_mode.py @@ -0,0 +1,276 @@ +"""Dual-mode (full/delta) parity for the gateway thread-state endpoints. + +Drives ``GET /api/threads/{id}``, ``GET /api/threads/{id}/state`` and +``POST /api/threads/{id}/history`` through the real route stack +(``build_thread_checkpoint_state_accessor`` -> factory-built graph -> +``CheckpointStateAccessor``) against a real ``InMemorySaver``, once per +checkpoint channel mode, and asserts the wire responses are identical apart +from checkpoint ids/timestamps. The delta storage layout must be invisible +to API consumers. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from _router_auth_helpers import make_authed_test_app +from fastapi.testclient import TestClient +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph import StateGraph +from langgraph.store.memory import InMemoryStore + +from app.gateway import services as gateway_services +from app.gateway.routers import threads +from deerflow.agents.thread_state import get_thread_state_schema +from deerflow.config.app_config import AppConfig, reset_app_config, set_app_config +from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore +from deerflow.runtime.checkpoint_mode import checkpoint_metadata_uses_delta, inject_checkpoint_mode + +_THREAD_ID = "thread-gateway-parity" + + +@pytest.fixture +def _stub_app_config(): + set_app_config(AppConfig.model_validate({"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}})) + yield + reset_app_config() + + +def _build_reply_graph(mode: str, checkpointer: Any): + async def _reply(state: dict[str, Any]) -> dict[str, Any]: + n = len(state.get("messages") or []) + return {"messages": [AIMessage(content=f"answer-{n}", id=f"a{n}")]} + + builder = StateGraph(get_thread_state_schema(mode)) + builder.add_node("reply", _reply) + builder.set_entry_point("reply") + builder.set_finish_point("reply") + return builder.compile(checkpointer=checkpointer) + + +def _message_wire_shape(messages: list[dict[str, Any]]) -> list[tuple[str, str, str]]: + return [(message.get("type"), message.get("content"), message.get("id")) for message in messages] + + +def _run_gateway_flow(mode: str, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: + app = make_authed_test_app() + store = InMemoryStore() + checkpointer = InMemorySaver() + app.state.store = store + app.state.checkpointer = checkpointer + app.state.thread_store = MemoryThreadMetaStore(store) + app.state.checkpoint_channel_mode = mode + app.state.run_event_store = SimpleNamespace() + app.include_router(threads.router) + + graph = _build_reply_graph(mode, checkpointer) + monkeypatch.setattr( + gateway_services, + "resolve_agent_factory", + lambda assistant_id=None: lambda config: graph, + ) + + config: dict[str, Any] = {"configurable": {"thread_id": _THREAD_ID}} + inject_checkpoint_mode(config, mode) + for i in range(2): + asyncio.run(graph.ainvoke({"messages": [HumanMessage(content=f"question-{i}", id=f"h{i}")]}, config)) + + with TestClient(app) as client: + thread_response = client.get(f"/api/threads/{_THREAD_ID}") + state_response = client.get(f"/api/threads/{_THREAD_ID}/state") + history_response = client.post(f"/api/threads/{_THREAD_ID}/history", json={"limit": 10}) + + assert thread_response.status_code == 200, thread_response.text + assert state_response.status_code == 200, state_response.text + assert history_response.status_code == 200, history_response.text + + thread_payload = thread_response.json() + state_payload = state_response.json() + history_payload = history_response.json() + + return { + "thread_status": thread_payload["status"], + "thread_messages": _message_wire_shape(thread_payload["values"]["messages"]), + "state_messages": _message_wire_shape(state_payload["values"]["messages"]), + "history_messages": [_message_wire_shape(snapshot["values"].get("messages", [])) for snapshot in history_payload], + } + + +def test_thread_state_endpoints_are_mode_invariant(_stub_app_config, monkeypatch: pytest.MonkeyPatch) -> None: + full = _run_gateway_flow("full", monkeypatch) + monkeypatch.undo() + delta = _run_gateway_flow("delta", monkeypatch) + assert full == delta + # Guard against a vacuous pass: the flow must have observed real messages. + assert full["thread_messages"], "expected seeded messages in the thread response" + assert any(full["history_messages"]), "expected history snapshots with messages" + + +def test_full_mode_gateway_rejects_delta_thread_with_409(_stub_app_config, monkeypatch: pytest.MonkeyPatch) -> None: + """Fail-closed gate at the HTTP boundary, against a real checkpointer. + + A full-mode process opening a delta thread must get a precise 409 naming + the cause — not a generic 500 that forces operators to grep logs. Seeds a + real delta checkpoint through the delta graph (marker + LangGraph delta + counters land in checkpoint metadata), then exercises every state surface + of the threads router in full mode. + """ + app = make_authed_test_app() + store = InMemoryStore() + checkpointer = InMemorySaver() + app.state.store = store + app.state.checkpointer = checkpointer + app.state.thread_store.get = AsyncMock(return_value=None) + app.state.checkpoint_channel_mode = "full" + app.state.run_event_store = SimpleNamespace() + app.include_router(threads.router) + + full_graph = _build_reply_graph("full", checkpointer) + monkeypatch.setattr( + gateway_services, + "resolve_agent_factory", + lambda assistant_id=None: lambda config: full_graph, + ) + + # Seed through the delta graph so the checkpoint carries the delta marker. + delta_graph = _build_reply_graph("delta", checkpointer) + config: dict[str, Any] = {"configurable": {"thread_id": _THREAD_ID}} + inject_checkpoint_mode(config, "delta") + asyncio.run(delta_graph.ainvoke({"messages": [HumanMessage(content="question", id="h0")]}, config)) + latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": _THREAD_ID, "checkpoint_ns": ""}})) + assert checkpoint_metadata_uses_delta(latest.metadata), "seed did not produce a delta checkpoint" + + with TestClient(app) as client: + state_response = client.get(f"/api/threads/{_THREAD_ID}/state") + assert state_response.status_code == 409, state_response.text + assert "requires delta mode" in state_response.json()["detail"] + assert _THREAD_ID in state_response.json()["detail"] + + update_response = client.post(f"/api/threads/{_THREAD_ID}/state", json={"values": {"title": "x"}}) + assert update_response.status_code == 409, update_response.text + assert "requires delta mode" in update_response.json()["detail"] + + history_response = client.post(f"/api/threads/{_THREAD_ID}/history", json={"limit": 10}) + assert history_response.status_code == 409, history_response.text + assert "requires delta mode" in history_response.json()["detail"] + + thread_response = client.get(f"/api/threads/{_THREAD_ID}") + assert thread_response.status_code == 409, thread_response.text + assert "requires delta mode" in thread_response.json()["detail"] + + +def test_full_mode_state_reads_degrade_to_raw_checkpointer_when_factory_fails(_stub_app_config, monkeypatch: pytest.MonkeyPatch) -> None: + """Full-mode read endpoints survive a broken agent factory. + + Full-mode checkpoints persist complete channel_values, so when the agent + factory cannot build the graph (bad model config, MCP server down), state + reads degrade to raw checkpointer reads instead of 500ing. The fail-closed + delta gate must still apply on the degraded path. + """ + app = make_authed_test_app() + store = InMemoryStore() + checkpointer = InMemorySaver() + app.state.store = store + app.state.checkpointer = checkpointer + app.state.thread_store.get = AsyncMock(return_value=None) + app.state.checkpoint_channel_mode = "full" + app.state.run_event_store = SimpleNamespace() + app.include_router(threads.router) + + full_graph = _build_reply_graph("full", checkpointer) + config: dict[str, Any] = {"configurable": {"thread_id": _THREAD_ID}} + inject_checkpoint_mode(config, "full") + for i in range(2): + asyncio.run(full_graph.ainvoke({"messages": [HumanMessage(content=f"question-{i}", id=f"h{i}")]}, config)) + latest = asyncio.run(checkpointer.aget_tuple(config)) + assert latest is not None + latest_created_at = latest.checkpoint["ts"] + + def _broken_factory(assistant_id=None): + def _factory(config): + raise RuntimeError("model config broken") + + return _factory + + monkeypatch.setattr(gateway_services, "resolve_agent_factory", _broken_factory) + + with TestClient(app) as client: + state_response = client.get(f"/api/threads/{_THREAD_ID}/state") + assert state_response.status_code == 200, state_response.text + values = state_response.json()["values"] + assert state_response.json()["created_at"] == latest_created_at + assert state_response.json()["checkpoint"]["ts"] == latest_created_at + assert _message_wire_shape(values["messages"]) == [ + ("human", "question-0", "h0"), + ("ai", "answer-1", "a1"), + ("human", "question-1", "h1"), + ("ai", "answer-3", "a3"), + ] + # next/tasks are not derivable without the compiled graph. + assert state_response.json()["next"] == [] + + history_response = client.post(f"/api/threads/{_THREAD_ID}/history", json={"limit": 10}) + assert history_response.status_code == 200, history_response.text + entries = history_response.json() + assert len(entries) >= 2 + assert all(entry["created_at"] for entry in entries) + + # History pagination: config.checkpoint_id is the *inclusive* anchor + # (pregel semantics), so the degraded path must include it too. + anchor_id = entries[1]["checkpoint_id"] + paged_response = client.post(f"/api/threads/{_THREAD_ID}/history", json={"limit": 10, "before": anchor_id}) + assert paged_response.status_code == 200, paged_response.text + assert paged_response.json()[0]["checkpoint_id"] == anchor_id + + app.state.thread_store.get = AsyncMock( + return_value={ + "thread_id": _THREAD_ID, + "assistant_id": None, + "status": "interrupted", + "created_at": latest_created_at, + "updated_at": latest_created_at, + "metadata": {}, + } + ) + thread_response = client.get(f"/api/threads/{_THREAD_ID}") + assert thread_response.status_code == 200, thread_response.text + assert thread_response.json()["status"] == "interrupted" + assert _message_wire_shape(thread_response.json()["values"]["messages"]) == [ + ("human", "question-0", "h0"), + ("ai", "answer-1", "a1"), + ("human", "question-1", "h1"), + ("ai", "answer-3", "a3"), + ] + + # The fail-closed gate still applies on the degraded path: a delta + # checkpoint is a precise 409, never silently served as partial state. + delta_graph = _build_reply_graph("delta", checkpointer) + delta_config: dict[str, Any] = {"configurable": {"thread_id": "thread-degraded-delta"}} + inject_checkpoint_mode(delta_config, "delta") + asyncio.run(delta_graph.ainvoke({"messages": [HumanMessage(content="q", id="h0")]}, delta_config)) + delta_response = client.get("/api/threads/thread-degraded-delta/state") + assert delta_response.status_code == 409, delta_response.text + assert "requires delta mode" in delta_response.json()["detail"] + + +def test_mutation_accessor_fails_closed_when_thread_metadata_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None: + app = make_authed_test_app() + app.state.thread_store.get = AsyncMock(side_effect=RuntimeError("metadata store unavailable")) + resolve_factory = AsyncMock() + monkeypatch.setattr(gateway_services, "resolve_agent_factory", resolve_factory) + + with pytest.raises(RuntimeError, match="metadata store unavailable"): + asyncio.run( + gateway_services.build_thread_checkpoint_state_mutation_accessor( + SimpleNamespace(app=app), + thread_id="custom-assistant-thread", + as_node="manual_state_update", + ) + ) + + resolve_factory.assert_not_called() diff --git a/backend/tests/test_gateway_config_freshness.py b/backend/tests/test_gateway_config_freshness.py index 8f38ab6cc..919ab246b 100644 --- a/backend/tests/test_gateway_config_freshness.py +++ b/backend/tests/test_gateway_config_freshness.py @@ -40,13 +40,26 @@ def _isolate_app_config_singleton(): reset_app_config() -def _write_config_yaml(path: Path, *, log_level: str) -> None: +def _write_config_yaml( + path: Path, + *, + log_level: str, + checkpoint_channel_mode: str | None = None, +) -> None: + database = ( + "" + if checkpoint_channel_mode is None + else f""" +database: + checkpoint_channel_mode: {checkpoint_channel_mode} +""" + ) path.write_text( f""" sandbox: use: deerflow.sandbox.local.provider:LocalSandboxProvider log_level: {log_level} -""".strip() +{database}""".strip() + "\n", encoding="utf-8", ) @@ -157,6 +170,28 @@ def test_run_context_app_config_reflects_yaml_edit(tmp_path, monkeypatch): assert second == {"log_level": "debug", "run_events_config": {"frozen": "startup"}} +def test_run_context_freezes_checkpoint_channel_mode_at_startup(tmp_path, monkeypatch): + from unittest.mock import MagicMock + + from app.gateway.deps import get_run_context + + config_file = tmp_path / "config.yaml" + _write_config_yaml(config_file, log_level="info", checkpoint_channel_mode="delta") + monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(config_file)) + + request = MagicMock() + request.app.state.checkpointer = MagicMock() + request.app.state.store = MagicMock() + request.app.state.run_event_store = MagicMock() + request.app.state.run_events_config = {"frozen": "startup"} + request.app.state.thread_store = MagicMock() + request.app.state.checkpoint_channel_mode = "full" + + ctx = get_run_context(request) + assert ctx.app_config.database.checkpoint_channel_mode == "delta" + assert ctx.checkpoint_channel_mode == "full" + + @pytest.mark.parametrize( "exception", [ diff --git a/backend/tests/test_gateway_run_drain_shutdown.py b/backend/tests/test_gateway_run_drain_shutdown.py index 9d9ce8940..fd5aa975d 100644 --- a/backend/tests/test_gateway_run_drain_shutdown.py +++ b/backend/tests/test_gateway_run_drain_shutdown.py @@ -196,7 +196,7 @@ async def test_langgraph_runtime_drains_runs_before_closing_checkpointer(monkeyp monkeypatch.setattr(RunManager, "shutdown", spy_shutdown, raising=False) app = FastAPI() - startup_config = SimpleNamespace(database=SimpleNamespace(backend="memory"), run_events=None) + startup_config = SimpleNamespace(database=SimpleNamespace(backend="memory", checkpoint_channel_mode="full"), run_events=None) async with langgraph_runtime(app, startup_config): pass diff --git a/backend/tests/test_gateway_run_recovery.py b/backend/tests/test_gateway_run_recovery.py index df2eb2308..921162a48 100644 --- a/backend/tests/test_gateway_run_recovery.py +++ b/backend/tests/test_gateway_run_recovery.py @@ -100,7 +100,7 @@ async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch): """SQLite startup should recover stale active runs before serving requests.""" app = FastAPI() config = SimpleNamespace( - database=SimpleNamespace(backend="sqlite"), + database=SimpleNamespace(backend="sqlite", checkpoint_channel_mode="full"), run_events=SimpleNamespace(backend="memory"), stream_bridge=SimpleNamespace(recovered_stream_cleanup_delay_seconds=60.0), ) @@ -144,7 +144,7 @@ async def test_sqlite_runtime_does_not_mark_thread_error_when_newer_run_is_succe """Startup recovery should not let an old orphaned run overwrite a newer terminal thread state.""" app = FastAPI() config = SimpleNamespace( - database=SimpleNamespace(backend="sqlite"), + database=SimpleNamespace(backend="sqlite", checkpoint_channel_mode="full"), run_events=SimpleNamespace(backend="memory"), stream_bridge=SimpleNamespace(recovered_stream_cleanup_delay_seconds=60.0), ) diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index ba8d5999f..fa47641a5 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -322,6 +322,33 @@ def test_build_run_config_with_overrides(): assert config["metadata"]["user"] == "alice" +def test_build_run_config_route_thread_id_overrides_client_configurable(): + from app.gateway.services import build_run_config + + config = build_run_config( + "route-thread", + {"configurable": {"thread_id": "caller-thread"}}, + None, + ) + + assert config["configurable"]["thread_id"] == "route-thread" + + +@pytest.mark.parametrize("section", ["configurable", "context"]) +def test_build_run_config_strips_external_checkpoint_mode_override(section): + from app.gateway.services import build_run_config + from deerflow.runtime.checkpoint_mode import INTERNAL_CHECKPOINT_MODE_KEY + + config = build_run_config( + "thread-1", + {section: {INTERNAL_CHECKPOINT_MODE_KEY: "delta", "model_name": "gpt-4"}}, + None, + ) + + assert INTERNAL_CHECKPOINT_MODE_KEY not in config[section] + assert config[section]["model_name"] == "gpt-4" + + def test_build_run_config_context_path_still_sets_configurable_thread_id(_stub_app_config): """A caller-supplied context (e.g. request-scoped secrets, #3861) must not deprive the checkpointer of configurable.thread_id, which it always needs to @@ -480,6 +507,76 @@ def test_resolve_agent_factory_returns_make_lead_agent(): assert resolve_agent_factory("custom-agent-123") is make_lead_agent +@pytest.mark.parametrize( + ("checkpoint_id", "includes_checkpoint_id"), + [(None, False), ("checkpoint-1", True)], +) +def test_build_checkpoint_state_accessor_uses_frozen_mode_and_binds_runtime_persistence( + _stub_app_config, + checkpoint_id, + includes_checkpoint_id, +): + from types import SimpleNamespace + from unittest.mock import patch + + from app.gateway.services import build_checkpoint_state_accessor + from deerflow.config.app_config import get_app_config + from deerflow.runtime.checkpoint_mode import CHECKPOINT_MODE_METADATA_KEY, INTERNAL_CHECKPOINT_MODE_KEY + + class FakeGraph: + checkpointer = None + store = None + + graph = FakeGraph() + captured = {} + + def fake_factory(*, config): + captured["config"] = config + return graph + + checkpointer = object() + store = object() + ctx = SimpleNamespace( + checkpointer=checkpointer, + store=store, + checkpoint_channel_mode="delta", + app_config=get_app_config(), + ) + request = SimpleNamespace( + state=SimpleNamespace(checkpoint_channel_mode="full"), + ) + + with ( + patch("app.gateway.services.get_run_context", return_value=ctx), + patch("app.gateway.services.resolve_agent_factory", return_value=fake_factory) as resolve, + ): + accessor, config = build_checkpoint_state_accessor( + request, + thread_id="thread-1", + assistant_id="Research_Agent", + checkpoint_id=checkpoint_id, + ) + + resolve.assert_called_once_with("Research_Agent") + assert captured["config"] is config + assert accessor.graph is graph + assert accessor.checkpointer is checkpointer + assert accessor.mode == "delta" + assert graph.checkpointer is checkpointer + assert graph.store is store + assert config["configurable"]["thread_id"] == "thread-1" + assert config["configurable"]["checkpoint_ns"] == "" + assert config["configurable"]["agent_name"] == "research-agent" + assert config["context"]["agent_name"] == "research-agent" + assert config["context"]["app_config"] is ctx.app_config + assert config["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "delta" + assert config["metadata"][CHECKPOINT_MODE_METADATA_KEY] == "delta" + assert INTERNAL_CHECKPOINT_MODE_KEY not in config["context"] + assert ("checkpoint_id" in config["configurable"]) is includes_checkpoint_id + if checkpoint_id is not None: + assert config["configurable"]["checkpoint_id"] == checkpoint_id + + def test_build_run_config_configurable_custom_agent_dual_writes_agent_name(): """Regression for issue #3549: even when the caller uses the legacy ``configurable`` path, ``agent_name`` must also land in @@ -1633,3 +1730,161 @@ class TestInjectAuthenticatedUserContextAuthz: config = {"context": "not a dict"} with pytest.raises(TypeError, match="run context must be a mapping"): inject_authenticated_user_context(config, _make_request_with_auth_source("session")) + + +@pytest.mark.asyncio +async def test_run_agent_full_mode_rejects_delta_before_graph_invocation(): + import asyncio + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from deerflow.runtime.checkpoint_mode import ( + CHECKPOINT_MODE_METADATA_KEY, + INTERNAL_CHECKPOINT_MODE_KEY, + ) + from deerflow.runtime.runs.manager import RunRecord + from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus + from deerflow.runtime.runs.worker import RunContext, run_agent + + checkpointer = AsyncMock() + checkpointer.aget_tuple.return_value = SimpleNamespace( + metadata={CHECKPOINT_MODE_METADATA_KEY: "delta"}, + checkpoint={"channel_values": {}}, + ) + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + run_manager = SimpleNamespace( + wait_for_prior_finalizing=AsyncMock(), + set_status=AsyncMock(), + ) + record = RunRecord( + run_id="run-checkpoint-mode", + thread_id="thread-delta", + assistant_id="lead-agent", + status=RunStatus.pending, + on_disconnect=DisconnectMode.cancel, + ) + record.abort_event = asyncio.Event() + agent_factory = MagicMock() + config = { + "configurable": { + "thread_id": record.thread_id, + INTERNAL_CHECKPOINT_MODE_KEY: "delta", + }, + "metadata": {CHECKPOINT_MODE_METADATA_KEY: "delta"}, + } + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext( + checkpointer=checkpointer, + checkpoint_channel_mode="full", + ), + agent_factory=agent_factory, + graph_input={"messages": []}, + config=config, + ) + + assert config["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "full" + assert CHECKPOINT_MODE_METADATA_KEY not in config["metadata"] + agent_factory.assert_not_called() + run_manager.set_status.assert_any_await( + record.run_id, + RunStatus.error, + error="Thread requires delta mode; materialize and convert its checkpoints before using full mode.", + ) + + +@pytest.mark.asyncio +async def test_run_agent_full_mode_checks_selected_checkpoint_before_graph(): + import asyncio + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock, call + + from deerflow.runtime.checkpoint_mode import CHECKPOINT_MODE_METADATA_KEY + from deerflow.runtime.runs.manager import RunRecord + from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus + from deerflow.runtime.runs.worker import RunContext, run_agent + + checkpointer = AsyncMock() + checkpointer.aget_tuple.side_effect = [ + SimpleNamespace( + metadata={}, + checkpoint={"channel_values": {"messages": ["latest full"]}}, + ), + SimpleNamespace( + metadata={CHECKPOINT_MODE_METADATA_KEY: "delta"}, + checkpoint={"channel_values": {}}, + ), + ] + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + run_manager = SimpleNamespace( + wait_for_prior_finalizing=AsyncMock(), + set_status=AsyncMock(), + ) + record = RunRecord( + run_id="run-selected-checkpoint-mode", + thread_id="thread-selected-delta", + assistant_id="lead-agent", + status=RunStatus.pending, + on_disconnect=DisconnectMode.cancel, + ) + record.abort_event = asyncio.Event() + agent_factory = MagicMock() + selected_config = { + "configurable": { + "thread_id": record.thread_id, + "checkpoint_ns": "branch", + "checkpoint_id": "delta-checkpoint", + "checkpoint_map": {"": "delta-checkpoint"}, + } + } + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext( + checkpointer=checkpointer, + checkpoint_channel_mode="full", + ), + agent_factory=agent_factory, + graph_input={"messages": []}, + config=selected_config, + ) + + agent_factory.assert_not_called() + assert checkpointer.aget_tuple.await_args_list[:2] == [ + call( + { + "configurable": { + "thread_id": record.thread_id, + "checkpoint_ns": "", + } + } + ), + call( + { + "configurable": { + "thread_id": record.thread_id, + "checkpoint_ns": "branch", + "checkpoint_id": "delta-checkpoint", + "checkpoint_map": {"": "delta-checkpoint"}, + } + } + ), + ] + run_manager.set_status.assert_any_await( + record.run_id, + RunStatus.error, + error="Thread requires delta mode; materialize and convert its checkpoints before using full mode.", + ) diff --git a/backend/tests/test_goal_worker.py b/backend/tests/test_goal_worker.py index 6ade2962c..68f73a7c4 100644 --- a/backend/tests/test_goal_worker.py +++ b/backend/tests/test_goal_worker.py @@ -6,12 +6,19 @@ from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.base import empty_checkpoint, uuid6 from langgraph.checkpoint.memory import InMemorySaver +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph from deerflow.runtime.goal import GoalEvaluation, attach_goal_evaluation, build_goal_state, latest_visible_assistant_signature, read_thread_goal, write_thread_goal from deerflow.runtime.runs import worker from deerflow.runtime.runs.manager import RunRecord from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus +def _full_accessor(checkpointer) -> CheckpointStateAccessor: + """Bind a full-mode accessor over a state-only graph for materialized reads.""" + graph = build_state_mutation_graph("goal_evaluator", "full") + return CheckpointStateAccessor.bind(graph, checkpointer, mode="full") + + class _CollectingBridge: def __init__(self) -> None: self.events: list[tuple[str, object]] = [] @@ -157,6 +164,7 @@ async def test_goal_worker_returns_hidden_continuation_when_goal_is_unmet(monkey monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -197,6 +205,7 @@ async def test_goal_worker_clears_goal_when_evaluator_is_satisfied(monkeypatch): monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -210,6 +219,53 @@ async def test_goal_worker_clears_goal_when_evaluator_is_satisfied(monkeypatch): assert bridge.events[0][0] == "values" +@pytest.mark.asyncio +async def test_goal_worker_evaluates_materialized_messages_in_delta_mode(monkeypatch): + """Delta checkpoints store no ``channel_values.messages``; the goal flow must + read messages through the mode-matched accessor or it sees an empty list, + loses the durable-receipt check, and stands down every continuation. + """ + checkpointer = InMemorySaver() + thread_id = "delta-goal-thread" + accessor = CheckpointStateAccessor.bind(build_state_mutation_graph("goal_evaluator", "delta"), checkpointer, mode="delta") + await accessor.aupdate( + {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}, + { + "messages": [ + HumanMessage(content="Please finish this task."), + AIMessage(content="I made a start, but I am not done."), + ] + }, + as_node="goal_evaluator", + ) + await write_thread_goal(checkpointer, thread_id, build_goal_state("Finish all tests", max_continuations=2)) + bridge = _CollectingBridge() + seen: dict[str, list] = {} + + async def fake_evaluate_goal_completion(_goal, messages, **_kwargs): + seen["messages"] = list(messages) + return GoalEvaluation( + satisfied=False, + blocker="goal_not_met_yet", + reason="Tests have not passed yet.", + evidence_summary="Implementation is incomplete.", + ) + + monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) + + continuation = await worker._prepare_goal_continuation_input(accessor=accessor, bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, run_id="run-delta", model_name="test-model", app_config=None) + + assert continuation is not None + assert [message.content for message in seen["messages"]] == [ + "Please finish this task.", + "I made a start, but I am not done.", + ] + latest_goal = await read_thread_goal(checkpointer, thread_id) + assert latest_goal is not None + assert latest_goal["continuation_count"] == 1 + assert "stand_down_reason" not in latest_goal["last_evaluation"] + + @pytest.mark.asyncio async def test_goal_worker_stands_down_for_non_continuable_blocker(monkeypatch): checkpointer = InMemorySaver() @@ -228,6 +284,7 @@ async def test_goal_worker_stands_down_for_non_continuable_blocker(monkeypatch): monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -275,6 +332,7 @@ async def test_goal_worker_stands_down_when_no_progress_repeats(monkeypatch): monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -309,6 +367,7 @@ async def test_goal_worker_does_not_resurrect_goal_cleared_during_evaluation(mon monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -371,6 +430,7 @@ async def test_goal_worker_stops_when_abort_is_requested_during_evaluation(monke monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -410,6 +470,7 @@ async def test_goal_worker_stands_down_when_thread_changes_after_evaluation(monk monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -462,6 +523,7 @@ async def test_goal_worker_stands_down_when_thread_changes_before_continuation(m monkeypatch.setattr(worker, "evaluate_goal_completion", fake_evaluate_goal_completion) continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, @@ -492,6 +554,7 @@ async def test_goal_worker_stands_down_without_durable_assistant_receipt(): bridge = _CollectingBridge() continuation = await worker._prepare_goal_continuation_input( + accessor=_full_accessor(checkpointer), bridge=bridge, checkpointer=checkpointer, thread_id=thread_id, diff --git a/backend/tests/test_lead_agent_model_resolution.py b/backend/tests/test_lead_agent_model_resolution.py index e9a99a9a5..14c092e14 100644 --- a/backend/tests/test_lead_agent_model_resolution.py +++ b/backend/tests/test_lead_agent_model_resolution.py @@ -19,7 +19,8 @@ from deerflow.agents.lead_agent import agent as lead_agent_module from deerflow.agents.middlewares import summarization_middleware as summarization_middleware_module from deerflow.agents.middlewares.loop_detection_middleware import LoopDetectionMiddleware from deerflow.agents.middlewares.subagent_limit_middleware import SubagentLimitMiddleware -from deerflow.agents.thread_state import ThreadState +from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware +from deerflow.agents.thread_state import DeltaThreadState, ThreadState from deerflow.config.app_config import AppConfig from deerflow.config.extensions_config import ExtensionsConfig from deerflow.config.loop_detection_config import LoopDetectionConfig @@ -28,6 +29,7 @@ from deerflow.config.model_config import ModelConfig from deerflow.config.sandbox_config import SandboxConfig from deerflow.config.subagents_config import SubagentsAppConfig from deerflow.config.summarization_config import SummarizationConfig +from deerflow.runtime.checkpoint_mode import INTERNAL_CHECKPOINT_MODE_KEY from deerflow.runtime.secret_context import write_slash_skill_source_path from deerflow.skills.types import Skill, SkillCategory @@ -202,6 +204,93 @@ def test_internal_make_lead_agent_uses_explicit_app_config(monkeypatch): assert result["model"] is not None +@pytest.mark.parametrize("is_bootstrap", [False, True]) +def test_internal_make_lead_agent_selects_and_normalizes_delta_state(monkeypatch, is_bootstrap): + app_config = _make_app_config([_make_model("delta-model", supports_thinking=False)]) + middleware = ViewImageMiddleware() + original_schema = middleware.state_schema + + import deerflow.tools as tools_module + + monkeypatch.setattr(tools_module, "get_available_tools", lambda **kwargs: []) + monkeypatch.setattr( + lead_agent_module, + "build_middlewares", + lambda config, model_name, agent_name=None, **kwargs: [middleware], + ) + monkeypatch.setattr(lead_agent_module, "_load_enabled_available_skills", lambda *args, **kwargs: []) + monkeypatch.setattr(lead_agent_module, "create_chat_model", lambda **kwargs: object()) + monkeypatch.setattr(lead_agent_module, "create_agent", lambda **kwargs: kwargs) + + result = lead_agent_module._make_lead_agent( + { + "configurable": { + "model_name": "delta-model", + "is_bootstrap": is_bootstrap, + INTERNAL_CHECKPOINT_MODE_KEY: "delta", + } + }, + app_config=app_config, + ) + + assert result["state_schema"] is DeltaThreadState + assert result["middleware"][0] is not middleware + assert middleware.state_schema is original_schema + + +def test_internal_make_lead_agent_does_not_take_mode_from_runtime_context(monkeypatch): + app_config = _make_app_config([_make_model("full-model", supports_thinking=False)]) + + import deerflow.tools as tools_module + + monkeypatch.setattr(tools_module, "get_available_tools", lambda **kwargs: []) + monkeypatch.setattr(lead_agent_module, "build_middlewares", lambda *args, **kwargs: []) + monkeypatch.setattr(lead_agent_module, "_load_enabled_available_skills", lambda *args, **kwargs: []) + monkeypatch.setattr(lead_agent_module, "create_chat_model", lambda **kwargs: object()) + monkeypatch.setattr(lead_agent_module, "create_agent", lambda **kwargs: kwargs) + + result = lead_agent_module._make_lead_agent( + { + "configurable": { + "model_name": "full-model", + INTERNAL_CHECKPOINT_MODE_KEY: "full", + }, + "context": {INTERNAL_CHECKPOINT_MODE_KEY: "delta"}, + }, + app_config=app_config, + ) + + assert result["state_schema"] is ThreadState + + +def test_public_make_lead_agent_does_not_take_mode_from_runtime_context(monkeypatch): + from deerflow.runtime import checkpoint_mode + + app_config = _make_app_config([_make_model("full-model", supports_thinking=False)]) + captured: dict[str, object] = {} + + monkeypatch.setattr(checkpoint_mode, "_frozen_checkpoint_channel_mode", None) + + def _capture(config, *, app_config): + captured["config"] = config + captured["app_config"] = app_config + return object() + + monkeypatch.setattr(lead_agent_module, "_make_lead_agent", _capture) + config = { + "configurable": {"model_name": "full-model"}, + "context": { + "app_config": app_config, + INTERNAL_CHECKPOINT_MODE_KEY: "delta", + }, + } + + lead_agent_module.make_lead_agent(config) + + assert config["configurable"][INTERNAL_CHECKPOINT_MODE_KEY] == "full" + assert captured["app_config"] is app_config + + def test_make_lead_agent_uses_runtime_app_config_from_context_without_global_read(monkeypatch): app_config = _make_app_config([_make_model("context-model", supports_thinking=False)]) diff --git a/backend/tests/test_memory_tools.py b/backend/tests/test_memory_tools.py index f2c960b89..ff2be74df 100644 --- a/backend/tests/test_memory_tools.py +++ b/backend/tests/test_memory_tools.py @@ -508,6 +508,7 @@ class TestModeGating: memory=MemoryConfig(enabled=True, mode="tool"), skills=SimpleNamespace(deferred_discovery=False, container_path="/tmp/skills"), tool_search=SimpleNamespace(enabled=False, auto_promote_top_k=0), + database=SimpleNamespace(checkpoint_channel_mode="full"), ) agent_kwargs = lead_agent_module._make_lead_agent({"configurable": {"agent_name": "test-agent"}}, app_config=app_config) @@ -540,6 +541,7 @@ class TestModeGating: memory=MemoryConfig(enabled=True, mode="tool"), skills=SimpleNamespace(deferred_discovery=False, container_path="/tmp/skills"), tool_search=SimpleNamespace(enabled=False, auto_promote_top_k=0), + database=SimpleNamespace(checkpoint_channel_mode="full"), ) agent_kwargs = lead_agent_module._make_lead_agent({"configurable": {"agent_name": "test-agent"}}, app_config=app_config) diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 2cfbf32a2..9d498f84e 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -2,22 +2,31 @@ import asyncio import copy from contextlib import suppress from types import SimpleNamespace -from typing import Any -from unittest.mock import AsyncMock, call +from typing import Annotated, Any, NotRequired, TypedDict +from unittest.mock import AsyncMock, call, patch import pytest -from langchain_core.messages import AIMessage +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langgraph.channels.delta import DeltaChannel from langgraph.checkpoint.base import empty_checkpoint from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages +from langgraph.types import Overwrite +from deerflow.agents.thread_state import merge_artifacts, merge_message_writes +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY from deerflow.runtime.runs.manager import ConflictError, RunManager from deerflow.runtime.runs.schemas import RunStatus from deerflow.runtime.runs.worker import ( + RollbackPoint, RunContext, _agent_factory_supports_app_config, _build_runtime_context, _bump_channel_version, + _capture_rollback_point, _collect_pre_existing_message_ids, _ensure_interrupted_title, _extract_llm_error_fallback_message, @@ -29,18 +38,59 @@ from deerflow.runtime.runs.worker import ( class FakeCheckpointer: - def __init__(self, *, put_result): + def __init__(self): self.adelete_thread = AsyncMock() - self.aput = AsyncMock(return_value=put_result) + self.aget_tuple = AsyncMock(return_value=None) self.aput_writes = AsyncMock() -def _make_checkpoint(checkpoint_id: str, messages: list[str], version: int): - checkpoint = empty_checkpoint() - checkpoint["id"] = checkpoint_id - checkpoint["channel_values"] = {"messages": messages} - checkpoint["channel_versions"] = {"messages": version} - return checkpoint +def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pending_writes=()): + return RollbackPoint( + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": checkpoint_id, + } + }, + messages=tuple(messages), + metadata={"source": "input"}, + pending_writes=tuple(pending_writes), + ) + + +def _stub_mutation_graph(monkeypatch, *, restored_config): + """Replace the rollback mutation graph with a stub returning ``restored_config``.""" + mock_graph = SimpleNamespace() + mock_graph.aupdate_state = AsyncMock(return_value=restored_config) + monkeypatch.setattr( + "deerflow.runtime.runs.worker.build_state_mutation_graph", + lambda *args, **kwargs: mock_graph, + ) + return mock_graph + + +def _rollback_accessor(checkpointer): + return CheckpointStateAccessor(graph=SimpleNamespace(), checkpointer=checkpointer, mode="full") + + +class _DeltaChannelState(TypedDict): + messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)] + + +class _FullChannelState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +def _build_message_append_graph(state_schema: type, checkpointer: Any): + async def _append_message(state: dict[str, Any]) -> dict[str, Any]: + return {"messages": [HumanMessage(content=f"turn-{len(state.get('messages') or [])}")]} + + builder = StateGraph(state_schema) + builder.add_node("append_message", _append_message) + builder.set_entry_point("append_message") + builder.set_finish_point("append_message") + return builder.compile(checkpointer=checkpointer) def test_build_runtime_context_includes_app_config_when_present(): @@ -142,12 +192,34 @@ async def test_run_agent_threads_pre_existing_message_ids_into_runtime_context() async def aget_tuple(self, _config): return SimpleNamespace( config={"configurable": {"checkpoint_id": "checkpoint-1"}}, - checkpoint={"channel_values": {"messages": [AIMessage(id="old-ai", content="old")]}}, + checkpoint={"channel_values": {}}, metadata={}, pending_writes=[], ) class DummyAgent: + async def aget_state(self, _config): + return SimpleNamespace( + values={ + "messages": [ + HumanMessage(id="h1", content="question"), + AIMessage(id="a1", content="answer"), + ] + }, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": "checkpoint-1", + } + }, + parent_config=None, + metadata={}, + next=(), + tasks=(), + created_at=None, + ) + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): captured["context"] = config["context"] yield {"messages": []} @@ -166,7 +238,69 @@ async def test_run_agent_threads_pre_existing_message_ids_into_runtime_context() ) context = captured["context"] - assert context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"old-ai"}) + assert context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"h1", "a1"}) + + +@pytest.mark.anyio +async def test_run_agent_marks_rollback_unusable_when_capture_fails(): + """A failed pre-run capture must disable rollback entirely. + + Rollback now restores messages by forking the pre-run checkpoint through + the graph (raw checkpoint blobs cannot reconstruct Delta-channel history), + so there is no safe fallback when materialization fails: the worker must + report ``snapshot_capture_failed=True`` with no rollback point rather + than restoring an empty or partial message history. + """ + run_manager = RunManager() + record = await run_manager.create("thread-1") + record.abort_action = "rollback" + record.abort_event.set() + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + + class DummyCheckpointer: + async def aget_tuple(self, _config): + return SimpleNamespace( + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": "checkpoint-1", + } + }, + checkpoint={"id": "checkpoint-1", "channel_values": {}}, + metadata={"source": "loop"}, + pending_writes=[], + ) + + class DummyAgent: + async def aget_state(self, _config): + raise RuntimeError("materialization failed") + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + yield {"messages": []} + + with patch( + "deerflow.runtime.runs.worker._rollback_to_pre_run_checkpoint", + new_callable=AsyncMock, + ) as rollback: + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=DummyCheckpointer()), + agent_factory=lambda *, config: DummyAgent(), + graph_input={}, + config={}, + ) + + rollback.assert_awaited_once() + rollback_kwargs = rollback.await_args.kwargs + assert rollback_kwargs["snapshot_capture_failed"] is True + assert rollback_kwargs["rollback_point"] is None @pytest.mark.anyio @@ -351,42 +485,37 @@ async def test_run_agent_defaults_root_run_name_from_configurable_agent_name(): @pytest.mark.anyio -async def test_rollback_restores_snapshot_without_deleting_thread(): - checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}) +async def test_rollback_forks_pre_run_checkpoint_without_deleting_thread(monkeypatch): + checkpointer = FakeCheckpointer() + mock_graph = _stub_mutation_graph( + monkeypatch, + restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}, + ) + rollback_point = _make_rollback_point( + pending_writes=[ + ("task-a", "messages", {"content": "first"}), + ("task-a", "status", "done"), + ("task-b", "events", {"type": "tool"}), + ], + ) await _rollback_to_pre_run_checkpoint( + accessor=_rollback_accessor(checkpointer), checkpointer=checkpointer, thread_id="thread-1", run_id="run-1", - pre_run_checkpoint_id="ckpt-1", - pre_run_snapshot={ - "checkpoint_ns": "", - "checkpoint": { - "id": "ckpt-1", - "channel_versions": {"messages": 3}, - "channel_values": {"messages": ["before"]}, - }, - "metadata": {"source": "input"}, - "pending_writes": [ - ("task-a", "messages", {"content": "first"}), - ("task-a", "status", "done"), - ("task-b", "events", {"type": "tool"}), - ], - }, + rollback_point=rollback_point, snapshot_capture_failed=False, ) checkpointer.adelete_thread.assert_not_awaited() - checkpointer.aput.assert_awaited_once() - restore_config, restored_checkpoint, restored_metadata, new_versions = checkpointer.aput.await_args.args - assert restore_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} - assert restored_checkpoint["id"] != "ckpt-1" - assert "channel_versions" in restored_checkpoint - assert "channel_values" in restored_checkpoint - assert restored_checkpoint["channel_versions"] == {"messages": 3} - assert restored_checkpoint["channel_values"] == {"messages": ["before"]} - assert restored_metadata == {"source": "input"} - assert new_versions == {"messages": 3} + mock_graph.aupdate_state.assert_awaited_once() + update_config, update_values = mock_graph.aupdate_state.await_args.args[:2] + assert update_config["configurable"]["checkpoint_id"] == "ckpt-1" + overwrite = update_values["messages"] + assert isinstance(overwrite, Overwrite) + assert overwrite.value == ["before"] + assert mock_graph.aupdate_state.await_args.kwargs["as_node"] == "rollback_restore" assert checkpointer.aput_writes.await_args_list == [ call( {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}, @@ -402,191 +531,501 @@ async def test_rollback_restores_snapshot_without_deleting_thread(): @pytest.mark.anyio -async def test_rollback_restored_checkpoint_becomes_latest_with_real_checkpointer(): - checkpointer = InMemorySaver() - thread_config = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} - before_checkpoint = _make_checkpoint("0001", ["before"], 1) - before_config = checkpointer.put(thread_config, before_checkpoint, {"step": 1}, {"messages": 1}) - after_checkpoint = _make_checkpoint("0002", ["after"], 2) - after_config = checkpointer.put(before_config, after_checkpoint, {"step": 2}, {"messages": 2}) - checkpointer.put_writes(after_config, [("messages", "pending-after")], task_id="task-after") +async def test_rollback_skips_when_rollback_point_has_no_checkpoint_id(monkeypatch): + """A rollback point without a checkpoint id cannot anchor a fork - skip safely.""" + checkpointer = FakeCheckpointer() + mock_graph = _stub_mutation_graph( + monkeypatch, + restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}, + ) await _rollback_to_pre_run_checkpoint( + accessor=_rollback_accessor(checkpointer), checkpointer=checkpointer, thread_id="thread-1", run_id="run-1", - pre_run_checkpoint_id="0001", - pre_run_snapshot={ - "checkpoint_ns": "", - "checkpoint": before_checkpoint, - "metadata": {"step": 1}, - "pending_writes": [("task-before", "messages", "pending-before")], - }, + rollback_point=_make_rollback_point(checkpoint_id=None), snapshot_capture_failed=False, ) - latest = checkpointer.get_tuple(thread_config) + checkpointer.adelete_thread.assert_not_awaited() + mock_graph.aupdate_state.assert_not_awaited() + checkpointer.aput_writes.assert_not_awaited() - assert latest is not None - assert latest.config["configurable"]["checkpoint_id"] != "0001" - assert latest.config["configurable"]["checkpoint_id"] != "0002" - assert latest.checkpoint["channel_values"] == {"messages": ["before"]} - assert latest.pending_writes == [("task-before", "messages", "pending-before")] - assert ("task-after", "messages", "pending-after") not in latest.pending_writes + +class _ArtifactChannelState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + artifacts: Annotated[list[str], merge_artifacts] + + +def _build_message_and_artifact_graph(checkpointer: Any): + async def _append(state: dict[str, Any]) -> dict[str, Any]: + n = len(state.get("messages") or []) + return {"messages": [HumanMessage(content=f"turn-{n}")], "artifacts": [f"a{n}"]} + + builder = StateGraph(_ArtifactChannelState) + builder.add_node("append", _append) + builder.set_entry_point("append") + builder.set_finish_point("append") + return builder.compile(checkpointer=checkpointer) + + +class _ExtensionFullState(TypedDict): + """Full-mode schema with a channel base ThreadState does not know about. + + Mirrors a custom ``AgentMiddleware.state_schema`` contribution (memory, + todos, uploads, …): absent from the base schema the mutation graph falls + back to, so a base-schema restore silently drops it. + """ + + messages: Annotated[list[AnyMessage], add_messages] + memory_notes: NotRequired[str] + + +class _ExtensionDeltaState(TypedDict): + messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)] + memory_notes: NotRequired[str] + + +@pytest.mark.anyio +@pytest.mark.parametrize( + "mode,state_schema", + [("full", _ExtensionFullState), ("delta", _ExtensionDeltaState)], +) +async def test_rollback_preserves_middleware_contributed_channels(mode, state_schema): + """The rollback mutation graph must compile with the thread's effective schema. + + Regression: ``_rollback_to_pre_run_checkpoint`` used to build the mutation + graph from the base ThreadState fallback, silently dropping channels + contributed by custom middleware from the restored checkpoint. + """ + checkpointer = InMemorySaver() + + async def _step(state: dict[str, Any]) -> dict[str, Any]: + n = len(state.get("messages") or []) + return {"messages": [HumanMessage(content=f"turn-{n}")], "memory_notes": f"note-{n}"} + + builder = StateGraph(state_schema) + builder.add_node("step", _step) + builder.set_entry_point("step") + builder.set_finish_point("step") + graph = builder.compile(checkpointer=checkpointer) + + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode=mode) + thread_config = {"configurable": {"thread_id": "thread-1"}} + + await graph.ainvoke({}, thread_config) + rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config) + assert rollback_point is not None + + await graph.ainvoke({}, thread_config) + + await _rollback_to_pre_run_checkpoint( + accessor=accessor, + checkpointer=checkpointer, + thread_id="thread-1", + run_id="run-1", + rollback_point=rollback_point, + snapshot_capture_failed=False, + ) + + latest_snapshot = await accessor.aget(thread_config) + assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"] + assert latest_snapshot.values["memory_notes"] == "note-0" + + +@pytest.mark.anyio +async def test_rollback_restores_non_message_channels_via_fork_inheritance(): + """Rollback overwrites only messages; sibling reducer channels must be + restored to their pre-run values by forking the pre-run checkpoint + (non-updated channels inherit from the parent). Locks in the load-bearing + assumption of the rollback design.""" + checkpointer = InMemorySaver() + graph = _build_message_and_artifact_graph(checkpointer) + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="full") + thread_config = {"configurable": {"thread_id": "thread-1"}} + + await graph.ainvoke({}, thread_config) + rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config) + assert rollback_point is not None + + await graph.ainvoke({}, thread_config) + cancelled_snapshot = await accessor.aget(thread_config) + assert cancelled_snapshot.values["artifacts"] == ["a0", "a1"] + assert [message.content for message in cancelled_snapshot.values["messages"]] == ["turn-0", "turn-1"] + + await _rollback_to_pre_run_checkpoint( + accessor=accessor, + checkpointer=checkpointer, + thread_id="thread-1", + run_id="run-1", + rollback_point=rollback_point, + snapshot_capture_failed=False, + ) + + latest_snapshot = await accessor.aget(thread_config) + assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"] + assert latest_snapshot.values["artifacts"] == ["a0"] + + +@pytest.mark.anyio +async def test_rollback_restored_checkpoint_becomes_latest_with_real_checkpointer(): + """Full-mode: the restored fork becomes latest, keeps pre-run values and + pre-run pending writes, and drops writes attached to the cancelled run.""" + checkpointer = InMemorySaver() + graph = _build_message_append_graph(_FullChannelState, checkpointer) + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="full") + thread_config = {"configurable": {"thread_id": "thread-1"}} + + await graph.ainvoke({}, thread_config) + pre_run_snapshot = await accessor.aget(thread_config) + pre_run_checkpoint_id = pre_run_snapshot.config["configurable"]["checkpoint_id"] + pre_run_config = { + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": pre_run_checkpoint_id, + } + } + checkpointer.put_writes(pre_run_config, [("title", "pending-before")], task_id="task-before") + + rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config) + assert rollback_point is not None + assert rollback_point.pending_writes == (("task-before", "title", "pending-before"),) + + await graph.ainvoke({}, thread_config) + cancelled_snapshot = await accessor.aget(thread_config) + cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"] + checkpointer.put_writes( + {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": cancelled_checkpoint_id}}, + [("title", "pending-after")], + task_id="task-after", + ) + + await _rollback_to_pre_run_checkpoint( + accessor=accessor, + checkpointer=checkpointer, + thread_id="thread-1", + run_id="run-1", + rollback_point=rollback_point, + snapshot_capture_failed=False, + ) + + latest_snapshot = await accessor.aget(thread_config) + restored_checkpoint_id = latest_snapshot.config["configurable"]["checkpoint_id"] + assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id) + assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0"] + + latest_tuple = await checkpointer.aget_tuple(thread_config) + assert latest_tuple.pending_writes == [("task-before", "title", "pending-before")] + restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}}) + assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == pre_run_checkpoint_id @pytest.mark.anyio async def test_rollback_deletes_thread_when_no_snapshot_exists(): - checkpointer = FakeCheckpointer(put_result=None) + checkpointer = FakeCheckpointer() await _rollback_to_pre_run_checkpoint( + accessor=_rollback_accessor(checkpointer), checkpointer=checkpointer, thread_id="thread-1", run_id="run-1", - pre_run_checkpoint_id=None, - pre_run_snapshot=None, + rollback_point=None, snapshot_capture_failed=False, ) checkpointer.adelete_thread.assert_awaited_once_with("thread-1") - checkpointer.aput.assert_not_awaited() checkpointer.aput_writes.assert_not_awaited() @pytest.mark.anyio -async def test_rollback_raises_when_restore_config_has_no_checkpoint_id(): - checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}) +async def test_rollback_raises_when_restore_config_has_no_checkpoint_id(monkeypatch): + checkpointer = FakeCheckpointer() + _stub_mutation_graph( + monkeypatch, + restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}, + ) with pytest.raises(RuntimeError, match="did not return checkpoint_id"): await _rollback_to_pre_run_checkpoint( + accessor=_rollback_accessor(checkpointer), checkpointer=checkpointer, thread_id="thread-1", run_id="run-1", - pre_run_checkpoint_id="ckpt-1", - pre_run_snapshot={ - "checkpoint_ns": "", - "checkpoint": {"id": "ckpt-1", "channel_versions": {}}, - "metadata": {}, - "pending_writes": [("task-a", "messages", "value")], - }, + rollback_point=_make_rollback_point(pending_writes=[("task-a", "messages", "value")]), snapshot_capture_failed=False, ) checkpointer.adelete_thread.assert_not_awaited() - checkpointer.aput.assert_awaited_once() checkpointer.aput_writes.assert_not_awaited() @pytest.mark.anyio -async def test_rollback_normalizes_none_checkpoint_ns_to_root_namespace(): - checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}) +async def test_capture_rollback_point_normalizes_missing_checkpoint_ns(): + """Snapshots whose config omits ``checkpoint_ns`` must anchor the root namespace.""" - await _rollback_to_pre_run_checkpoint( - checkpointer=checkpointer, - thread_id="thread-1", - run_id="run-1", - pre_run_checkpoint_id="ckpt-1", - pre_run_snapshot={ - "checkpoint_ns": None, - "checkpoint": {"id": "ckpt-1", "channel_versions": {}}, - "metadata": {}, - "pending_writes": [], - }, - snapshot_capture_failed=False, - ) + class _TupleCheckpointer: + async def aget_tuple(self, _config): + return SimpleNamespace(pending_writes=[]) - checkpointer.aput.assert_awaited_once() - restore_config, restored_checkpoint, restored_metadata, new_versions = checkpointer.aput.await_args.args - assert restore_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} - assert restored_checkpoint["id"] != "ckpt-1" - assert restored_checkpoint["channel_versions"] == {} - assert restored_metadata == {} - assert new_versions == {} + class _SnapshotAgent: + async def aget_state(self, _config): + return SimpleNamespace( + values={"messages": [HumanMessage(id="h1", content="before")]}, + config={"configurable": {"thread_id": "thread-1", "checkpoint_id": "ckpt-1"}}, + metadata={"source": "loop"}, + ) + + checkpointer = _TupleCheckpointer() + accessor = CheckpointStateAccessor(graph=_SnapshotAgent(), checkpointer=checkpointer, mode="full") + + rollback_point = await _capture_rollback_point(accessor, checkpointer, {"configurable": {"thread_id": "thread-1"}}) + + assert rollback_point is not None + assert rollback_point.config["configurable"]["checkpoint_ns"] == "" + assert rollback_point.config["configurable"]["checkpoint_id"] == "ckpt-1" + assert [message.content for message in rollback_point.messages] == ["before"] + assert rollback_point.metadata == {"source": "loop"} @pytest.mark.anyio -async def test_rollback_raises_on_malformed_pending_write_not_a_tuple(): +async def test_capture_rollback_point_returns_none_without_checkpoint(): + """A thread with no checkpoints has no rollback point (delete/reset contract).""" + + class _EmptyAgent: + async def aget_state(self, _config): + return SimpleNamespace( + values={}, + config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}, + metadata={}, + ) + + checkpointer = SimpleNamespace(aget_tuple=AsyncMock(return_value=None)) + accessor = CheckpointStateAccessor(graph=_EmptyAgent(), checkpointer=checkpointer, mode="full") + + assert await _capture_rollback_point(accessor, checkpointer, {"configurable": {"thread_id": "thread-1"}}) is None + + +@pytest.mark.anyio +async def test_rollback_raises_on_malformed_pending_write_not_a_tuple(monkeypatch): """pending_writes containing a non-3-tuple item should raise RuntimeError.""" - checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}) + checkpointer = FakeCheckpointer() + _stub_mutation_graph( + monkeypatch, + restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}, + ) with pytest.raises(RuntimeError, match="rollback failed: pending_write is not a 3-tuple"): await _rollback_to_pre_run_checkpoint( + accessor=_rollback_accessor(checkpointer), checkpointer=checkpointer, thread_id="thread-1", run_id="run-1", - pre_run_checkpoint_id="ckpt-1", - pre_run_snapshot={ - "checkpoint_ns": "", - "checkpoint": {"id": "ckpt-1", "channel_versions": {}}, - "metadata": {}, - "pending_writes": [ + rollback_point=_make_rollback_point( + pending_writes=[ ("task-a", "messages", "valid"), # valid ["only", "two"], # malformed: only 2 elements ], - }, + ), snapshot_capture_failed=False, ) - # aput succeeded but aput_writes should not be called due to malformed data - checkpointer.aput.assert_awaited_once() checkpointer.aput_writes.assert_not_awaited() @pytest.mark.anyio -async def test_rollback_raises_on_malformed_pending_write_non_string_channel(): +async def test_rollback_raises_on_malformed_pending_write_non_string_channel(monkeypatch): """pending_writes containing a non-string channel should raise RuntimeError.""" - checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}) + checkpointer = FakeCheckpointer() + _stub_mutation_graph( + monkeypatch, + restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}, + ) with pytest.raises(RuntimeError, match="rollback failed: pending_write has non-string channel"): await _rollback_to_pre_run_checkpoint( + accessor=_rollback_accessor(checkpointer), checkpointer=checkpointer, thread_id="thread-1", run_id="run-1", - pre_run_checkpoint_id="ckpt-1", - pre_run_snapshot={ - "checkpoint_ns": "", - "checkpoint": {"id": "ckpt-1", "channel_versions": {}}, - "metadata": {}, - "pending_writes": [ + rollback_point=_make_rollback_point( + pending_writes=[ ("task-a", 123, "value"), # malformed: channel is not a string ], - }, + ), snapshot_capture_failed=False, ) - checkpointer.aput.assert_awaited_once() checkpointer.aput_writes.assert_not_awaited() @pytest.mark.anyio -async def test_rollback_propagates_aput_writes_failure(): +async def test_rollback_propagates_aput_writes_failure(monkeypatch): """If aput_writes fails, the exception should propagate (not be swallowed).""" - checkpointer = FakeCheckpointer(put_result={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}) - # Simulate aput_writes failure + checkpointer = FakeCheckpointer() checkpointer.aput_writes.side_effect = RuntimeError("Database connection lost") + _stub_mutation_graph( + monkeypatch, + restored_config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}, + ) with pytest.raises(RuntimeError, match="Database connection lost"): await _rollback_to_pre_run_checkpoint( + accessor=_rollback_accessor(checkpointer), checkpointer=checkpointer, thread_id="thread-1", run_id="run-1", - pre_run_checkpoint_id="ckpt-1", - pre_run_snapshot={ - "checkpoint_ns": "", - "checkpoint": {"id": "ckpt-1", "channel_versions": {}}, - "metadata": {}, - "pending_writes": [ + rollback_point=_make_rollback_point( + pending_writes=[ ("task-a", "messages", "value"), ], - }, + ), snapshot_capture_failed=False, ) - # aput succeeded, aput_writes was called but failed - checkpointer.aput.assert_awaited_once() checkpointer.aput_writes.assert_awaited_once() +@pytest.mark.anyio +async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints(): + """Delta channels omit messages from checkpoint blobs, so rollback must + fork the pre-run lineage through the graph: the restored checkpoint's + messages are reconstructed by replaying ancestor writes, and writes from + the cancelled run (not ancestors of the fork) must not leak back in.""" + checkpointer = InMemorySaver() + graph = _build_message_append_graph(_DeltaChannelState, checkpointer) + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta") + thread_config = {"configurable": {"thread_id": "thread-1"}} + + await graph.ainvoke({}, thread_config) + await graph.ainvoke({}, thread_config) + rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config) + assert rollback_point is not None + assert [message.content for message in rollback_point.messages] == ["turn-0", "turn-1"] + pre_run_checkpoint_id = rollback_point.config["configurable"]["checkpoint_id"] + + await graph.ainvoke({}, thread_config) # the run that gets cancelled + cancelled_snapshot = await accessor.aget(thread_config) + cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"] + assert len(cancelled_snapshot.values["messages"]) == 3 + + await _rollback_to_pre_run_checkpoint( + accessor=accessor, + checkpointer=checkpointer, + thread_id="thread-1", + run_id="run-1", + rollback_point=rollback_point, + snapshot_capture_failed=False, + ) + + latest_snapshot = await accessor.aget(thread_config) + restored_checkpoint_id = latest_snapshot.config["configurable"]["checkpoint_id"] + assert restored_checkpoint_id not in (pre_run_checkpoint_id, cancelled_checkpoint_id) + assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"] + + restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}}) + assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == pre_run_checkpoint_id + assert restored_tuple.metadata.get("source") == "update" + # A non-snapshot Delta checkpoint must not persist the full message list. + raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages") + assert not isinstance(raw_messages, list) + + +@pytest.mark.anyio +async def test_rollback_restores_pre_run_pending_writes_for_delta_checkpoints(): + """Pre-run pending writes are re-attached to the restored fork; writes + attached to the cancelled run are not.""" + checkpointer = InMemorySaver() + graph = _build_message_append_graph(_DeltaChannelState, checkpointer) + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta") + thread_config = {"configurable": {"thread_id": "thread-1"}} + + await graph.ainvoke({}, thread_config) + pre_run_snapshot = await accessor.aget(thread_config) + pre_run_checkpoint_id = pre_run_snapshot.config["configurable"]["checkpoint_id"] + checkpointer.put_writes( + {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": pre_run_checkpoint_id}}, + [("title", "in-flight-title")], + task_id="task-in-flight", + ) + + rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config) + assert rollback_point is not None + assert rollback_point.pending_writes == (("task-in-flight", "title", "in-flight-title"),) + + await graph.ainvoke({}, thread_config) # cancelled run + cancelled_snapshot = await accessor.aget(thread_config) + cancelled_checkpoint_id = cancelled_snapshot.config["configurable"]["checkpoint_id"] + checkpointer.put_writes( + {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": cancelled_checkpoint_id}}, + [("title", "cancelled-title")], + task_id="task-cancelled", + ) + + await _rollback_to_pre_run_checkpoint( + accessor=accessor, + checkpointer=checkpointer, + thread_id="thread-1", + run_id="run-1", + rollback_point=rollback_point, + snapshot_capture_failed=False, + ) + + latest_tuple = await checkpointer.aget_tuple(thread_config) + assert latest_tuple.pending_writes == [("task-in-flight", "title", "in-flight-title")] + assert [message.content for message in (await accessor.aget(thread_config)).values["messages"]] == ["turn-0"] + + +@pytest.mark.anyio +async def test_rollback_forks_pre_run_lineage_for_delta_checkpoints_sqlite_reopen(tmp_path): + """Same lineage contract against a disk-backed saver, verified after a + close/reopen so only persisted bytes can satisfy the assertions.""" + db_path = tmp_path / "rollback.sqlite3" + + async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer: + await checkpointer.setup() + graph = _build_message_append_graph(_DeltaChannelState, checkpointer) + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta") + thread_config = {"configurable": {"thread_id": "thread-1"}} + + await graph.ainvoke({}, thread_config) + await graph.ainvoke({}, thread_config) + rollback_point = await _capture_rollback_point(accessor, checkpointer, thread_config) + assert rollback_point is not None + pre_run_checkpoint_id = rollback_point.config["configurable"]["checkpoint_id"] + + await graph.ainvoke({}, thread_config) # cancelled run + + await _rollback_to_pre_run_checkpoint( + accessor=accessor, + checkpointer=checkpointer, + thread_id="thread-1", + run_id="run-1", + rollback_point=rollback_point, + snapshot_capture_failed=False, + ) + + restored_snapshot = await accessor.aget(thread_config) + restored_checkpoint_id = restored_snapshot.config["configurable"]["checkpoint_id"] + assert restored_checkpoint_id != pre_run_checkpoint_id + assert [message.content for message in restored_snapshot.values["messages"]] == ["turn-0", "turn-1"] + + async with AsyncSqliteSaver.from_conn_string(str(db_path)) as checkpointer: + graph = _build_message_append_graph(_DeltaChannelState, checkpointer) + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode="delta") + thread_config = {"configurable": {"thread_id": "thread-1"}} + + latest_snapshot = await accessor.aget(thread_config) + assert latest_snapshot.config["configurable"]["checkpoint_id"] == restored_checkpoint_id + assert [message.content for message in latest_snapshot.values["messages"]] == ["turn-0", "turn-1"] + + restored_tuple = await checkpointer.aget_tuple({"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": restored_checkpoint_id}}) + assert restored_tuple.parent_config["configurable"]["checkpoint_id"] == pre_run_checkpoint_id + raw_messages = restored_tuple.checkpoint.get("channel_values", {}).get("messages") + assert not isinstance(raw_messages, list) + + def test_agent_factory_supports_app_config_detects_supported_signature(): def factory(*, config, app_config=None): return (config, app_config) @@ -886,28 +1325,21 @@ def test_extract_llm_error_fallback_message_default_filter_is_empty(): assert _extract_llm_error_fallback_message(state) == "Connection error." -def test_collect_pre_existing_message_ids_pulls_ids_from_snapshot(): - snapshot = { - "checkpoint": { - "channel_values": { - "messages": [ - AIMessage(id="a", content="x"), - AIMessage(id="b", content="y"), - AIMessage(content="no-id-here"), # ignored - ] - } - } +def test_collect_pre_existing_message_ids_pulls_ids_from_values(): + values = { + "messages": [ + AIMessage(id="a", content="x"), + AIMessage(id="b", content="y"), + AIMessage(content="no-id-here"), # ignored + ] } - assert _collect_pre_existing_message_ids(snapshot) == {"a", "b"} + assert _collect_pre_existing_message_ids(values) == {"a", "b"} def test_collect_pre_existing_message_ids_handles_missing_pieces(): assert _collect_pre_existing_message_ids(None) == set() assert _collect_pre_existing_message_ids({}) == set() - assert _collect_pre_existing_message_ids({"checkpoint": None}) == set() - assert _collect_pre_existing_message_ids({"checkpoint": {}}) == set() - assert _collect_pre_existing_message_ids({"checkpoint": {"channel_values": None}}) == set() - assert _collect_pre_existing_message_ids({"checkpoint": {"channel_values": {"messages": None}}}) == set() + assert _collect_pre_existing_message_ids({"messages": None}) == set() @pytest.mark.anyio @@ -952,6 +1384,23 @@ async def test_run_agent_ignores_stale_llm_error_fallback_from_prior_run(): ) class DummyAgent: + async def aget_state(self, _config): + return SimpleNamespace( + values={"messages": [stale_fallback]}, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": "ckpt-stale", + } + }, + parent_config=None, + metadata={}, + next=(), + tasks=(), + created_at=None, + ) + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): # Replay the prior fallback message (as LangGraph would when using # stream_mode="values") and then yield a fresh successful AIMessage. @@ -1561,7 +2010,10 @@ async def test_ensure_interrupted_title_bumps_channel_version_and_declares_it_in assert written_metadata["source"] == "update" assert written_metadata["step"] == 8 assert written_metadata["writes"] == {"runtime_interrupt_title": {"title": "Generated Title"}} - assert write_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} + # The write is parented to the checkpoint it was derived from - without + # the parent pointer the saver stores a root checkpoint, severing + # Delta-channel replay ancestry and full-mode history walks. + assert write_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "ckpt-1"}} @pytest.mark.anyio diff --git a/backend/tests/test_runtime_lifecycle_e2e.py b/backend/tests/test_runtime_lifecycle_e2e.py index c093a2187..3521da9bb 100644 --- a/backend/tests/test_runtime_lifecycle_e2e.py +++ b/backend/tests/test_runtime_lifecycle_e2e.py @@ -24,6 +24,8 @@ import pytest from _agent_e2e_helpers import FakeToolCallingModel, build_single_tool_call_model from langchain_core.messages import AIMessage, HumanMessage +from deerflow.runtime.checkpoint_state import build_state_mutation_graph + pytestmark = pytest.mark.no_auto_user @@ -95,6 +97,28 @@ class _ScriptedAgent: self.interrupt_before_nodes = None self.interrupt_after_nodes = None self.model = FakeToolCallingModel(responses=[AIMessage(content=self.answer)]) + # Gateway read paths consume graph-materialized snapshots + # (``aget_state``/``get_state_history``), so the double delegates those + # reads to a real state-only graph bound to the same checkpointer. + self._state_graph = build_state_mutation_graph("scripted_state", "full") + + def _bound_state_graph(self): + self._state_graph.checkpointer = self.checkpointer + self._state_graph.store = self.store + return self._state_graph + + async def aget_state(self, config): + return await self._bound_state_graph().aget_state(config) + + def get_state(self, config): + return self._bound_state_graph().get_state(config) + + async def aget_state_history(self, config, **kwargs): + async for snapshot in self._bound_state_graph().aget_state_history(config, **kwargs): + yield snapshot + + def get_state_history(self, config, **kwargs): + yield from self._bound_state_graph().get_state_history(config, **kwargs) async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): del subgraphs diff --git a/backend/tests/test_support_bundle.py b/backend/tests/test_support_bundle.py index b0202ec47..d66433417 100644 --- a/backend/tests/test_support_bundle.py +++ b/backend/tests/test_support_bundle.py @@ -298,7 +298,7 @@ def test_create_support_bundle_masks_provider_config_secret_shaped_keys(tmp_path project_root = tmp_path / "project" project_root.mkdir() (project_root / "config.yaml").write_text( - "config_version: 26\n" + "config_version: 27\n" "models:\n - name: default\n" "guardrails:\n" " enabled: true\n" diff --git a/backend/tests/test_thread_regenerate_prepare.py b/backend/tests/test_thread_regenerate_prepare.py index e79e50940..992b9354a 100644 --- a/backend/tests/test_thread_regenerate_prepare.py +++ b/backend/tests/test_thread_regenerate_prepare.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio from types import SimpleNamespace +from unittest.mock import AsyncMock, patch import pytest from fastapi import HTTPException @@ -27,9 +28,11 @@ def _checkpoint(checkpoint_id: str, messages: list[object], *, metadata: dict | class FakeCheckpointer: - def __init__(self, history, *, latest=None): + def __init__(self, history, *, latest=None, materialized_history=None, materialized_latest=None): self.history = history self.latest = latest + self.materialized_history = materialized_history + self.materialized_latest = materialized_latest self.alist_limits = [] async def aget_tuple(self, config): @@ -44,6 +47,64 @@ class FakeCheckpointer: yield item +def _snapshot(checkpoint_id: str, messages: list[object], *, metadata: dict | None = None): + return SimpleNamespace( + values={"messages": messages}, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": checkpoint_id, + "checkpoint_map": None, + } + }, + metadata=metadata or {}, + ) + + +class FakeAccessor: + def __init__(self, checkpointer: FakeCheckpointer): + self.checkpointer = checkpointer + + @staticmethod + def _from_raw(checkpoint): + return SimpleNamespace( + values=dict(checkpoint.checkpoint.get("channel_values", {})), + config=checkpoint.config, + metadata=checkpoint.metadata, + ) + + async def aget(self, _config): + if self.checkpointer.materialized_latest is not None: + return self.checkpointer.materialized_latest + raw = self.checkpointer.latest or (self.checkpointer.history[0] if self.checkpointer.history else None) + return self._from_raw(raw) if raw is not None else SimpleNamespace(values={}, config={}, metadata={}) + + async def ahistory(self, _config, *, limit=None): + self.checkpointer.alist_limits.append(limit) + history = self.checkpointer.materialized_history + if history is None: + history = [self._from_raw(item) for item in self.checkpointer.history] + return history[:limit] + + +@pytest.fixture(autouse=True) +def _patch_checkpoint_accessor(monkeypatch): + from app.gateway.routers import thread_runs + + def build_accessor(request, *, thread_id, assistant_id=None, checkpoint_id=None): + config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + if checkpoint_id is not None: + config["configurable"]["checkpoint_id"] = checkpoint_id + return FakeAccessor(request.app.state.checkpointer), config + + async def build_thread_accessor(request, *, thread_id, checkpoint_id=None): + return build_accessor(request, thread_id=thread_id, checkpoint_id=checkpoint_id) + + monkeypatch.setattr(thread_runs, "build_checkpoint_state_accessor", build_accessor) + monkeypatch.setattr(thread_runs, "build_thread_checkpoint_state_accessor", build_thread_accessor) + + class FakeEventStore: def __init__(self, rows): self.rows = rows @@ -75,8 +136,174 @@ def _request(checkpointer, event_store, *, run_manager=None, user_id="user-1"): ) +def test_run_wait_readers_return_materialized_final_values() -> None: + from app.gateway.routers import runs, thread_runs + + snapshot = SimpleNamespace( + values={ + "messages": [ + HumanMessage(id="h1", content="question"), + AIMessage(id="a1", content="answer"), + ] + }, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": "ckpt-2", + } + }, + parent_config={"configurable": {"checkpoint_id": "ckpt-1"}}, + metadata={"step": 2}, + next=(), + tasks=(), + created_at=None, + ) + accessor = SimpleNamespace(aget=AsyncMock(return_value=snapshot)) + record = SimpleNamespace( + run_id="run-1", + thread_id="thread-1", + task=None, + status=RunStatus.success, + error=None, + ) + request = SimpleNamespace() + body = thread_runs.RunCreateRequest( + assistant_id="lead-agent", + config={"configurable": {"thread_id": "thread-1"}}, + ) + + async def _scenario() -> tuple[dict, dict]: + with ( + patch.object(thread_runs, "get_stream_bridge", return_value=object()), + patch.object(thread_runs, "get_run_manager", return_value=object()), + patch.object(thread_runs, "start_run", AsyncMock(return_value=record)), + patch.object( + thread_runs, + "build_checkpoint_state_accessor", + create=True, + return_value=(accessor, snapshot.config), + ), + patch.object(runs, "get_stream_bridge", return_value=object()), + patch.object(runs, "get_run_manager", return_value=object()), + patch.object(runs, "start_run", AsyncMock(return_value=record)), + patch.object( + runs, + "build_checkpoint_state_accessor", + create=True, + return_value=(accessor, snapshot.config), + ), + ): + thread_result = await thread_runs.wait_run.__wrapped__("thread-1", body, request) + stateless_result = await runs.stateless_wait(body, request) + return thread_result, stateless_result + + thread_result, stateless_result = asyncio.run(_scenario()) + + assert [message["id"] for message in thread_result["messages"]] == ["h1", "a1"] + assert [message["id"] for message in stateless_result["messages"]] == ["h1", "a1"] + + +def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None: + from app.gateway.routers import runs, thread_runs + + snapshot = SimpleNamespace( + values={}, + config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}}, + parent_config=None, + metadata={}, + next=(), + tasks=(), + created_at=None, + ) + accessor = SimpleNamespace(aget=AsyncMock(return_value=snapshot)) + record = SimpleNamespace( + run_id="run-1", + thread_id="thread-1", + task=None, + status=RunStatus.error, + error="run failed before checkpoint", + ) + request = SimpleNamespace() + body = thread_runs.RunCreateRequest(config={"configurable": {"thread_id": "thread-1"}}) + + async def _scenario() -> tuple[dict, dict]: + with ( + patch.object(thread_runs, "get_stream_bridge", return_value=object()), + patch.object(thread_runs, "get_run_manager", return_value=object()), + patch.object(thread_runs, "start_run", AsyncMock(return_value=record)), + patch.object( + thread_runs, + "build_checkpoint_state_accessor", + return_value=(accessor, snapshot.config), + ), + patch.object(runs, "get_stream_bridge", return_value=object()), + patch.object(runs, "get_run_manager", return_value=object()), + patch.object(runs, "start_run", AsyncMock(return_value=record)), + patch.object( + runs, + "build_checkpoint_state_accessor", + return_value=(accessor, snapshot.config), + ), + ): + thread_result = await thread_runs.wait_run.__wrapped__("thread-1", body, request) + stateless_result = await runs.stateless_wait(body, request) + return thread_result, stateless_result + + thread_result, stateless_result = asyncio.run(_scenario()) + + expected = {"status": "error", "error": "run failed before checkpoint"} + assert thread_result == expected + assert stateless_result == expected + + +@pytest.mark.parametrize("route_name", ["thread", "stateless"]) +def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(route_name: str) -> None: + from app.gateway.routers import runs, thread_runs + + record = SimpleNamespace( + run_id="run-1", + thread_id="thread-1", + task=None, + status=RunStatus.error, + error="run failed before checkpoint", + ) + request = SimpleNamespace() + body = thread_runs.RunCreateRequest(config={"configurable": {"thread_id": "thread-1"}}) + + async def _scenario() -> dict: + if route_name == "thread": + with ( + patch.object(thread_runs, "get_stream_bridge", return_value=object()), + patch.object(thread_runs, "get_run_manager", return_value=object()), + patch.object(thread_runs, "start_run", AsyncMock(return_value=record)), + patch.object( + thread_runs, + "build_checkpoint_state_accessor", + side_effect=RuntimeError("graph construction failed"), + ), + ): + return await thread_runs.wait_run.__wrapped__("thread-1", body, request) + + with ( + patch.object(runs, "get_stream_bridge", return_value=object()), + patch.object(runs, "get_run_manager", return_value=object()), + patch.object(runs, "start_run", AsyncMock(return_value=record)), + patch.object( + runs, + "build_checkpoint_state_accessor", + side_effect=RuntimeError("graph construction failed"), + ), + ): + return await runs.stateless_wait(body, request) + + result = asyncio.run(_scenario()) + + assert result == {"status": "error", "error": "run failed before checkpoint"} + + def test_prepare_regenerate_payload_returns_clean_input_and_base_checkpoint(): - from app.gateway.routers.thread_runs import _prepare_regenerate_payload + from app.gateway.routers import thread_runs human = HumanMessage( id="human-1", @@ -103,7 +330,12 @@ def test_prepare_regenerate_payload_returns_clean_input_and_base_checkpoint(): ] ) - response = asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store))) + original_builder = thread_runs.build_thread_checkpoint_state_accessor + thread_builder = AsyncMock(side_effect=original_builder) + with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", thread_builder): + response = asyncio.run(thread_runs._prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store))) + + assert [call.kwargs["thread_id"] for call in thread_builder.await_args_list] == ["thread-1", "thread-1"] assert response.checkpoint == { "checkpoint_ns": "", @@ -122,6 +354,54 @@ def test_prepare_regenerate_payload_returns_clean_input_and_base_checkpoint(): assert regenerated_human["additional_kwargs"] == {"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}]} +def test_prepare_regenerate_uses_materialized_history_when_raw_messages_are_omitted(): + from app.gateway.routers.thread_runs import _prepare_regenerate_payload + + earlier_human = HumanMessage(id="human-0", content="earlier question") + earlier_ai = AIMessage(id="ai-0", content="earlier answer") + target_human = HumanMessage(id="human-1", content="question") + target_ai = AIMessage(id="ai-1", content="answer") + + raw_latest = _checkpoint("ckpt-ai", []) + raw_after_human = _checkpoint("ckpt-human", []) + raw_base = _checkpoint("ckpt-base", []) + materialized_history = [ + _snapshot("ckpt-ai", [earlier_human, earlier_ai, target_human, target_ai]), + _snapshot("ckpt-human", [earlier_human, earlier_ai, target_human]), + _snapshot("ckpt-base", [earlier_human, earlier_ai]), + ] + checkpointer = FakeCheckpointer( + [raw_latest, raw_after_human, raw_base], + latest=raw_latest, + materialized_history=materialized_history, + materialized_latest=materialized_history[0], + ) + event_store = FakeEventStore( + [ + { + "run_id": "run-target", + "event_type": "ai_message", + "category": "message", + "content": {"id": "ai-1", "type": "ai", "content": "answer"}, + "metadata": {"caller": "lead_agent"}, + } + ] + ) + + response = asyncio.run( + _prepare_regenerate_payload( + "thread-1", + "ai-1", + _request(checkpointer, event_store), + ) + ) + + assert response.checkpoint["checkpoint_id"] == "ckpt-base" + assert response.metadata["regenerate_checkpoint_id"] == "ckpt-base" + assert response.input["messages"][0]["id"] == "human-1" + assert checkpointer.alist_limits == [400] + + def test_prepare_regenerate_payload_rejects_non_latest_assistant(): from app.gateway.routers.thread_runs import _prepare_regenerate_payload @@ -280,5 +560,5 @@ def test_find_base_checkpoint_ignores_duration_only_checkpoints() -> None: result = asyncio.run(_find_base_checkpoint_before_human("thread-1", "human-1", _request(checkpointer, FakeEventStore([])))) - assert result is base + assert result.config == base.config assert checkpointer.alist_limits == [400] diff --git a/backend/tests/test_thread_state_reducers.py b/backend/tests/test_thread_state_reducers.py index 864fb690a..77e6163be 100644 --- a/backend/tests/test_thread_state_reducers.py +++ b/backend/tests/test_thread_state_reducers.py @@ -13,6 +13,7 @@ from deerflow.agents import thread_state as thread_state_module from deerflow.agents.thread_state import ( _SKILL_CONTEXT_MAX_ENTRIES, TERMINAL_STATUSES, + THREAD_STATE_REDUCER_FIELDS, SkillEntry, ThreadState, merge_artifacts, @@ -352,3 +353,16 @@ class TestThreadStateAnnotations: def test_skill_context_field_is_wired_to_merge_skill_context(self): hints = get_type_hints(ThreadState, include_extras=True) assert merge_skill_context in hints["skill_context"].__metadata__ + + def test_public_reducer_field_names_match_thread_state_contract(self): + assert THREAD_STATE_REDUCER_FIELDS == { + "messages", + "sandbox", + "artifacts", + "todos", + "goal", + "viewed_images", + "promoted", + "delegations", + "skill_context", + } diff --git a/backend/tests/test_threads_checkpoint_mode.py b/backend/tests/test_threads_checkpoint_mode.py new file mode 100644 index 000000000..3f8a018ba --- /dev/null +++ b/backend/tests/test_threads_checkpoint_mode.py @@ -0,0 +1,147 @@ +"""Dual-mode (full/delta) end-to-end parity for thread checkpoint flows. + +Exercises the same thread lifecycle — multi-turn writes, latest-state read, +bounded history, branch from an earlier checkpoint, and a compaction-style +wholesale ``Overwrite`` write — against the production thread state schemas +(``get_thread_state_schema(mode)``) on both ``InMemorySaver`` and +``AsyncSqliteSaver``. Every assertion compares the materialized state +produced by the two modes: delta storage must be behaviorally invisible. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.graph import StateGraph +from langgraph.types import Overwrite + +from deerflow.agents.thread_state import get_thread_state_schema +from deerflow.runtime.checkpoint_mode import inject_checkpoint_mode +from deerflow.runtime.checkpoint_state import ( + CheckpointStateAccessor, + build_state_mutation_graph, +) + +pytestmark = pytest.mark.anyio + + +def _build_reply_graph(mode: str, checkpointer: Any): + """One-node graph on the production thread schema for ``mode``.""" + + async def _reply(state: dict[str, Any]) -> dict[str, Any]: + n = len(state.get("messages") or []) + return {"messages": [AIMessage(content=f"answer-{n}", id=f"a{n}")]} + + builder = StateGraph(get_thread_state_schema(mode)) + builder.add_node("reply", _reply) + builder.set_entry_point("reply") + builder.set_finish_point("reply") + return builder.compile(checkpointer=checkpointer) + + +def _normalize_messages(values: dict[str, Any]) -> list[tuple[str, str, str]]: + return [(message.type, message.content, message.id) for message in values.get("messages", [])] + + +async def _run_thread_lifecycle(mode: str, checkpointer: Any) -> dict[str, Any]: + """Drive the full thread lifecycle and return normalized observations.""" + graph = _build_reply_graph(mode, checkpointer) + accessor = CheckpointStateAccessor.bind(graph, checkpointer, mode=mode) + config: dict[str, Any] = {"configurable": {"thread_id": "thread-parity"}} + inject_checkpoint_mode(config, mode) + + # Multi-turn writes. + for i in range(3): + await graph.ainvoke({"messages": [HumanMessage(content=f"question-{i}", id=f"h{i}")]}, config) + + # Latest materialized state. + latest = await accessor.aget(config) + latest_messages = _normalize_messages(latest.values) + + # Bounded history: limit must hold and the walk must be newest-first. + history = await accessor.ahistory(config, limit=2) + history_shapes = [_normalize_messages(snapshot.values) for snapshot in history] + + # Branch from the first checkpoint (regenerate flow): the fork inherits + # the pre-branch state, then new writes append on top. + first_checkpoint_id = history[-1].config["configurable"]["checkpoint_id"] + branch_config: dict[str, Any] = {"configurable": {"thread_id": "thread-parity", "checkpoint_ns": "", "checkpoint_id": first_checkpoint_id}} + inject_checkpoint_mode(branch_config, mode) + await accessor.aupdate(branch_config, {"messages": [HumanMessage(content="retry", id="h-retry")]}, as_node="reply") + branched = await accessor.aget(branch_config) + branch_messages = _normalize_messages(branched.values) + + # Compaction-style wholesale write through the mutation graph: replace + # the message list with an Overwrite, no node scheduling. + mutation_graph = build_state_mutation_graph("compact", mode, get_thread_state_schema(mode)) + mutation_accessor = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode=mode) + compacted_messages = [HumanMessage(content="summary so far", id="h-summary")] + await mutation_accessor.aupdate(config, {"messages": Overwrite(compacted_messages)}) + compacted = await accessor.aget(config) + compacted_messages_seen = _normalize_messages(compacted.values) + + # The thread keeps working on top of the compacted state. + await graph.ainvoke({"messages": [HumanMessage(content="after-compact", id="h3b")]}, config) + resumed = await accessor.aget(config) + resumed_messages = _normalize_messages(resumed.values) + + return { + "latest": latest_messages, + "history": history_shapes, + "branch": branch_messages, + "compacted": compacted_messages_seen, + "resumed": resumed_messages, + } + + +async def test_thread_lifecycle_parity_memory_saver() -> None: + full = await _run_thread_lifecycle("full", InMemorySaver()) + delta = await _run_thread_lifecycle("delta", InMemorySaver()) + assert full == delta + + +async def test_thread_lifecycle_parity_sqlite_saver(tmp_path) -> None: + async with AsyncSqliteSaver.from_conn_string(str(tmp_path / "full.sqlite3")) as saver: + await saver.setup() + full = await _run_thread_lifecycle("full", saver) + async with AsyncSqliteSaver.from_conn_string(str(tmp_path / "delta.sqlite3")) as saver: + await saver.setup() + delta = await _run_thread_lifecycle("delta", saver) + assert full == delta + + +async def test_delta_thread_avoids_per_step_full_message_blobs(tmp_path) -> None: + """Storage-level guard: in delta mode the per-step checkpoints must not + carry a serialized ``messages`` blob (the channel persists as incremental + writes plus rare snapshots); full mode re-serializes the whole list into + every checkpoint, which is exactly the O(N^2) growth delta mode removes.""" + + async def _blob_checkpoint_counts(mode: str, db_name: str) -> tuple[int, int]: + async with AsyncSqliteSaver.from_conn_string(str(tmp_path / db_name)) as saver: + await saver.setup() + graph = _build_reply_graph(mode, saver) + config: dict[str, Any] = {"configurable": {"thread_id": f"thread-{mode}"}} + inject_checkpoint_mode(config, mode) + for i in range(4): + await graph.ainvoke({"messages": [HumanMessage(content=f"q{i}", id=f"h{i}")]}, config) + + total = 0 + with_messages_blob = 0 + async for checkpoint_tuple in saver.alist(config): + total += 1 + if "messages" in checkpoint_tuple.checkpoint.get("channel_values", {}): + with_messages_blob += 1 + return total, with_messages_blob + + delta_total, delta_blobs = await _blob_checkpoint_counts("delta", "delta.sqlite3") + full_total, full_blobs = await _blob_checkpoint_counts("full", "full.sqlite3") + + assert delta_total > 0 and full_total > 0 + # Delta: at most the periodic snapshot carries a blob; full: every + # message-writing checkpoint does. + assert delta_blobs <= 1, f"delta checkpoints re-serialized messages: {delta_blobs}/{delta_total}" + assert full_blobs >= 4, f"full mode should blob messages per step: {full_blobs}/{full_total}" diff --git a/backend/tests/test_threads_router.py b/backend/tests/test_threads_router.py index 63ddbd211..e41395edd 100644 --- a/backend/tests/test_threads_router.py +++ b/backend/tests/test_threads_router.py @@ -8,14 +8,17 @@ from _router_auth_helpers import make_authed_test_app from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from langchain_core.messages import AIMessage, HumanMessage, ToolMessage -from langgraph.checkpoint.base import empty_checkpoint +from langgraph.checkpoint.base import empty_checkpoint, uuid6 from langgraph.checkpoint.memory import InMemorySaver from langgraph.store.memory import InMemoryStore +from langgraph.types import Overwrite +from app.gateway import services as gateway_services from app.gateway.routers import threads from deerflow.config.paths import Paths from deerflow.persistence.thread_meta import InvalidMetadataFilterError from deerflow.persistence.thread_meta.memory import THREADS_NS, MemoryThreadMetaStore +from deerflow.runtime.checkpoint_state import CheckpointStateAccessor _ISO_TIMESTAMP_RE = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}") @@ -68,6 +71,136 @@ def _build_thread_app() -> tuple[FastAPI, InMemoryStore, InMemorySaver]: return app, store, checkpointer +class _RawStateAccessor: + def __init__(self, checkpointer: InMemorySaver): + self.checkpointer = checkpointer + + @staticmethod + def _snapshot(checkpoint_tuple, fallback_config): + if checkpoint_tuple is None: + return SimpleNamespace( + values={}, + config=fallback_config, + parent_config=None, + metadata={}, + next=(), + tasks=(), + created_at=None, + ) + checkpoint = checkpoint_tuple.checkpoint or {} + metadata = checkpoint_tuple.metadata or {} + return SimpleNamespace( + values=dict(checkpoint.get("channel_values", {})), + config=checkpoint_tuple.config, + parent_config=checkpoint_tuple.parent_config, + metadata=metadata, + next=(), + tasks=(), + created_at=checkpoint.get("ts") or metadata.get("created_at"), + ) + + async def aget(self, config): + checkpoint_tuple = await self.checkpointer.aget_tuple(config) + return self._snapshot(checkpoint_tuple, config) + + async def ahistory(self, config, *, limit=None): + snapshots = [] + async for checkpoint_tuple in self.checkpointer.alist(config, limit=limit): + snapshots.append(self._snapshot(checkpoint_tuple, config)) + return snapshots + + async def aupdate(self, config, values, *, as_node=None): + checkpoint_tuple = await self.checkpointer.aget_tuple(config) + checkpoint = dict(checkpoint_tuple.checkpoint if checkpoint_tuple is not None else empty_checkpoint()) + channel_values = dict(checkpoint.get("channel_values", {})) + channel_values.update({key: value.value if isinstance(value, Overwrite) else value for key, value in values.items()}) + checkpoint["channel_values"] = channel_values + channel_versions = dict(checkpoint.get("channel_versions", {})) + new_versions = {} + for key in values: + current_version = channel_versions.get(key) + next_version = current_version + 1 if isinstance(current_version, int) else 1 + channel_versions[key] = next_version + new_versions[key] = next_version + checkpoint["channel_versions"] = channel_versions + checkpoint["id"] = str(uuid6()) + metadata = dict(checkpoint_tuple.metadata if checkpoint_tuple is not None else {}) + metadata.update( + { + "source": "update", + "step": metadata.get("step", -1) + 1, + "writes": {as_node: values}, + } + ) + write_config = { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "checkpoint_ns": config["configurable"].get("checkpoint_ns", ""), + } + } + return await self.checkpointer.aput(write_config, checkpoint, metadata, new_versions) + + +@pytest.fixture(autouse=True) +def _patch_checkpoint_state_builder(monkeypatch): + def _builder(request, *, thread_id, assistant_id=None, checkpoint_id=None): + config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + if checkpoint_id is not None: + config["configurable"]["checkpoint_id"] = checkpoint_id + return _RawStateAccessor(request.app.state.checkpointer), config + + def _mutation_builder(request, *, thread_id, as_node, checkpoint_id=None, state_schema=None): + config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + if checkpoint_id is not None: + config["configurable"]["checkpoint_id"] = checkpoint_id + return _RawStateAccessor(request.app.state.checkpointer), config + + async def _read_boundary(request, *, thread_id, checkpoint_id=None): + return _builder(request, thread_id=thread_id, checkpoint_id=checkpoint_id) + + async def _mutation_boundary(request, *, thread_id, as_node, checkpoint_id=None): + return _mutation_builder(request, thread_id=thread_id, as_node=as_node, checkpoint_id=checkpoint_id) + + monkeypatch.setattr(threads, "build_checkpoint_state_accessor", _builder) + monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", _mutation_builder) + monkeypatch.setattr(threads, "build_thread_checkpoint_state_accessor", _read_boundary) + monkeypatch.setattr(threads, "build_thread_checkpoint_state_mutation_accessor", _mutation_boundary) + + +class _FakeStateAccessor: + def __init__(self, snapshot: SimpleNamespace): + self.snapshot = snapshot + + async def aget(self, config): + return self.snapshot + + async def ahistory(self, config, *, limit=None): + return [self.snapshot][:limit] + + +def _materialized_snapshot() -> SimpleNamespace: + return SimpleNamespace( + values={ + "messages": [ + HumanMessage(id="h1", content="question"), + AIMessage(id="a1", content="answer"), + ] + }, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": "ckpt-2", + } + }, + parent_config={"configurable": {"checkpoint_id": "ckpt-1"}}, + metadata={"step": 2}, + next=(), + tasks=(), + created_at=None, + ) + + async def _write_checkpoint( checkpointer: InMemorySaver, thread_id: str, @@ -562,6 +695,147 @@ def test_get_thread_returns_iso_for_legacy_unix_record() -> None: assert _ISO_TIMESTAMP_RE.match(body["updated_at"]), body["updated_at"] +def test_latest_thread_readers_use_materialized_snapshot_values() -> None: + app, store, checkpointer = _build_thread_app() + thread_id = "thread-1" + + async def _seed() -> None: + await store.aput( + THREADS_NS, + thread_id, + { + "thread_id": thread_id, + "status": "idle", + "created_at": "2026-07-18T00:00:00+00:00", + "updated_at": "2026-07-18T00:00:00+00:00", + "metadata": {}, + }, + ) + await checkpointer.aput( + {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}, + empty_checkpoint(), + {"step": 2, "source": "loop", "writes": {}, "parents": {}}, + {}, + ) + + asyncio.run(_seed()) + accessor = _FakeStateAccessor(_materialized_snapshot()) + thread_accessor = AsyncMock(return_value=(accessor, {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})) + + with ( + patch( + "app.gateway.routers.threads.build_checkpoint_state_accessor", + create=True, + return_value=(accessor, {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}), + ), + patch( + "app.gateway.routers.threads.build_thread_checkpoint_state_accessor", + new=thread_accessor, + ), + TestClient(app) as client, + ): + thread_response = client.get(f"/api/threads/{thread_id}") + state_response = client.get(f"/api/threads/{thread_id}/state") + history_response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10}) + + assert thread_response.status_code == 200, thread_response.text + assert state_response.status_code == 200, state_response.text + assert history_response.status_code == 200, history_response.text + assert [message["id"] for message in thread_response.json()["values"]["messages"]] == ["h1", "a1"] + assert [message["id"] for message in state_response.json()["values"]["messages"]] == ["h1", "a1"] + assert [message["id"] for message in history_response.json()[0]["values"]["messages"]] == ["h1", "a1"] + assert [call.kwargs["thread_id"] for call in thread_accessor.await_args_list] == [thread_id, thread_id] + + +def test_get_thread_status_uses_raw_pending_writes_for_materialized_checkpoint() -> None: + app, store, _checkpointer = _build_thread_app() + thread_id = "thread-1" + + async def _seed() -> None: + await store.aput( + THREADS_NS, + thread_id, + { + "thread_id": thread_id, + "status": "idle", + "created_at": "2026-07-18T00:00:00+00:00", + "updated_at": "2026-07-18T00:00:00+00:00", + "metadata": {}, + }, + ) + + asyncio.run(_seed()) + requested_configs = [] + + class _AdvancingCheckpointer: + async def aget_tuple(self, config): + requested_configs.append(config) + checkpoint_id = config.get("configurable", {}).get("checkpoint_id") + return SimpleNamespace( + pending_writes=[] if checkpoint_id == "ckpt-2" else [("task-old", "__error__", "stale")], + ) + + app.state.checkpointer = _AdvancingCheckpointer() + snapshot = _materialized_snapshot() + accessor = _FakeStateAccessor(snapshot) + + with ( + patch( + "app.gateway.routers.threads.build_checkpoint_state_accessor", + return_value=(accessor, {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}), + ), + TestClient(app) as client, + ): + response = client.get(f"/api/threads/{thread_id}") + + assert response.status_code == 200, response.text + assert response.json()["status"] == "idle" + assert requested_configs == [snapshot.config] + + +@pytest.mark.parametrize("stored_status", ["running", "error"]) +def test_get_thread_preserves_metadata_status_without_checkpoint(stored_status: str) -> None: + app, store, _checkpointer = _build_thread_app() + thread_id = "thread-without-checkpoint" + + async def _seed() -> None: + await store.aput( + THREADS_NS, + thread_id, + { + "thread_id": thread_id, + "status": stored_status, + "created_at": "2026-07-18T00:00:00+00:00", + "updated_at": "2026-07-18T00:00:00+00:00", + "metadata": {}, + }, + ) + + asyncio.run(_seed()) + snapshot = SimpleNamespace( + values={}, + config={"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}, + parent_config=None, + metadata={}, + next=(), + tasks=(), + created_at=None, + ) + accessor = _FakeStateAccessor(snapshot) + + with ( + patch( + "app.gateway.routers.threads.build_checkpoint_state_accessor", + return_value=(accessor, snapshot.config), + ), + TestClient(app) as client, + ): + response = client.get(f"/api/threads/{thread_id}") + + assert response.status_code == 200, response.text + assert response.json()["status"] == stored_status + + def test_patch_thread_returns_iso_and_advances_updated_at() -> None: app, store, _checkpointer = _build_thread_app() thread_id = "patch-target" @@ -888,6 +1162,333 @@ def test_branch_thread_from_older_assistant_turn_creates_truncated_thread() -> N assert branch_entry["values"]["title"] == "Original chat" +def test_branch_thread_uses_materialized_history_and_overwrites_fresh_seed(monkeypatch) -> None: + app, _store, _checkpointer = _build_thread_app() + source_thread_id = "source-materialized" + messages = [ + HumanMessage(id="h1", content="First question"), + AIMessage(id="a1", content="First answer"), + HumanMessage(id="h2", content="Second question"), + AIMessage(id="a2", content="Second answer"), + ] + + def snapshot(checkpoint_id: str, materialized_messages: list[object]) -> SimpleNamespace: + return SimpleNamespace( + values={"messages": materialized_messages, "title": "Materialized title"}, + config={ + "configurable": { + "thread_id": source_thread_id, + "checkpoint_ns": "", + "checkpoint_id": checkpoint_id, + } + }, + metadata={"step": int(checkpoint_id[-1])}, + ) + + source_accessor = SimpleNamespace() + source_history = [ + snapshot("ckpt-2", messages), + snapshot("ckpt-1", messages[:2]), + ] + branch_updates: list[tuple[dict, dict, str | None]] = [] + + async def source_ahistory(config, *, limit=None): + assert config["configurable"]["thread_id"] == source_thread_id + assert limit == 200 + return source_history + + async def branch_aupdate(config, values, *, as_node=None): + branch_updates.append((config, values, as_node)) + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "checkpoint_ns": "", + "checkpoint_id": "branch-seed", + } + } + + source_accessor.ahistory = source_ahistory + branch_accessor = SimpleNamespace(aupdate=branch_aupdate) + + def build_accessor(_request, *, thread_id, assistant_id=None, checkpoint_id=None): + assert thread_id == source_thread_id + return source_accessor, { + "configurable": { + "thread_id": source_thread_id, + "checkpoint_ns": "", + } + } + + def build_mutation_accessor(_request, *, thread_id, as_node, checkpoint_id=None, state_schema=None): + return branch_accessor, { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + } + } + + monkeypatch.setattr(threads, "build_checkpoint_state_accessor", build_accessor) + monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", build_mutation_accessor) + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": source_thread_id, "metadata": {}, "assistant_id": "agent"}, + ) + assert created.status_code == 200, created.text + response = client.post( + f"/api/threads/{source_thread_id}/branches", + json={"message_id": "a1", "message_ids": ["a1"]}, + ) + + assert response.status_code == 200, response.text + body = response.json() + assert body["parent_checkpoint_id"] == "ckpt-1" + assert len(branch_updates) == 1 + update_config, update_values, as_node = branch_updates[0] + assert isinstance(update_values["messages"], Overwrite) + assert [message.id for message in update_values["messages"].value] == ["h1", "a1"] + assert update_config["configurable"]["thread_id"] == body["thread_id"] + assert update_config["metadata"]["source"] == "branch" + assert as_node == "branch" + + +def test_branch_thread_real_mutation_graph_finishes_without_scheduling(monkeypatch) -> None: + app, _store, _checkpointer = _build_thread_app() + app.state.checkpoint_channel_mode = "delta" + source_thread_id = "source-real-branch" + messages = [ + HumanMessage(id="h1", content="First question"), + AIMessage(id="a1", content="First answer"), + HumanMessage(id="h2", content="Second question"), + AIMessage(id="a2", content="Second answer"), + ] + + def snapshot(checkpoint_id: str, materialized_messages: list[object]) -> SimpleNamespace: + return SimpleNamespace( + values={"messages": materialized_messages}, + config={ + "configurable": { + "thread_id": source_thread_id, + "checkpoint_ns": "", + "checkpoint_id": checkpoint_id, + } + }, + metadata={}, + ) + + source_accessor = SimpleNamespace( + ahistory=AsyncMock( + return_value=[ + snapshot("ckpt-2", messages), + snapshot("ckpt-1", messages[:2]), + ] + ) + ) + + def source_builder(_request, *, thread_id, assistant_id=None, checkpoint_id=None): + if thread_id != source_thread_id: + raise AssertionError("fresh branches must use the dedicated mutation graph") + return source_accessor, { + "configurable": { + "thread_id": source_thread_id, + "checkpoint_ns": "", + } + } + + real_mutation_builder = gateway_services.build_checkpoint_state_mutation_accessor + monkeypatch.setattr(threads, "build_checkpoint_state_accessor", source_builder) + monkeypatch.setattr( + threads, + "build_checkpoint_state_mutation_accessor", + real_mutation_builder, + raising=False, + ) + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": source_thread_id, "metadata": {}, "assistant_id": "agent"}, + ) + assert created.status_code == 200, created.text + response = client.post( + f"/api/threads/{source_thread_id}/branches", + json={"message_id": "a1", "message_ids": ["a1"]}, + ) + + assert response.status_code == 200, response.text + new_thread_id = response.json()["thread_id"] + accessor, config = real_mutation_builder( + SimpleNamespace(app=app), + thread_id=new_thread_id, + as_node="branch", + ) + branch_snapshot = asyncio.run(accessor.aget(config)) + assert [message.id for message in branch_snapshot.values["messages"]] == ["h1", "a1"] + assert branch_snapshot.next == () + assert branch_snapshot.metadata["deerflow_branch"] is True + assert branch_snapshot.metadata["branch_parent_checkpoint_id"] == "ckpt-1" + + +def _wire_extension_agent(monkeypatch, app, checkpointer, mode): + """Stub only the infra context + assistant factory; keep builders real. + + The production resolution path stays live: thread record -> assistant_id + -> resolve_agent_factory -> effective graph (base schema for ``mode`` plus + a non-identity reducer channel contributed by AgentMiddleware.state_schema). + """ + import operator + from typing import Annotated, NotRequired, TypedDict + + from langchain.agents import create_agent + from langchain.agents.middleware import AgentMiddleware + from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel + + from deerflow.agents.thread_state import get_thread_state_schema + + class ExtensionState(TypedDict): + ext_list: NotRequired[Annotated[list[str], operator.add]] + + class ExtensionMiddleware(AgentMiddleware): + state_schema = ExtensionState + + app.state.checkpoint_channel_mode = mode + model = FakeMessagesListChatModel(responses=[AIMessage(content="ok")]) + + def custom_factory(*, config=None): + return create_agent(model, middleware=[ExtensionMiddleware()], state_schema=get_thread_state_schema(mode)) + + def default_factory(*, config=None): + return create_agent(model, state_schema=get_thread_state_schema(mode)) + + def selective_factory(assistant_id): + # Only the thread's recorded assistant yields the extension graph; + # unresolved assistant_id must materialize with the default schema so + # the tests actually guard the resolution boundary. + return custom_factory if assistant_id == "extension-agent" else default_factory + + ctx = SimpleNamespace(checkpointer=checkpointer, store=None, checkpoint_channel_mode=mode, app_config=None) + monkeypatch.setattr(gateway_services, "get_run_context", lambda _request: ctx) + monkeypatch.setattr(gateway_services, "resolve_agent_factory", selective_factory) + monkeypatch.setattr(threads, "build_checkpoint_state_accessor", gateway_services.build_checkpoint_state_accessor) + monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", gateway_services.build_checkpoint_state_mutation_accessor) + monkeypatch.setattr(threads, "build_thread_checkpoint_state_accessor", gateway_services.build_thread_checkpoint_state_accessor) + monkeypatch.setattr(threads, "build_thread_checkpoint_state_mutation_accessor", gateway_services.build_thread_checkpoint_state_mutation_accessor) + return custom_factory + + +async def _seed_extension_source(checkpointer, custom_factory, mode, source_thread_id): + accessor = CheckpointStateAccessor.bind(custom_factory(), checkpointer, mode=mode) + await accessor.aupdate( + {"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}}, + {"messages": [HumanMessage(id="h1", content="question")], "ext_list": ["merged"]}, + as_node="model", + ) + await accessor.aupdate( + {"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}}, + {"messages": [AIMessage(id="a1", content="answer")], "ext_list": ["payload"]}, + as_node="model", + ) + + +@pytest.mark.parametrize("mode", ["full", "delta"]) +def test_state_endpoints_preserve_extension_reducer_channels(monkeypatch, mode) -> None: + """A non-identity middleware reducer channel survives state endpoints. + + GET /state must return the extension value (resolved via the thread's + assistant_id), POST /state must replace it, and branch must preserve it + byte-for-byte by copying reducer channels with Overwrite semantics. + """ + app, _store, checkpointer = _build_thread_app() + custom_factory = _wire_extension_agent(monkeypatch, app, checkpointer, mode) + + recorded_updates: list[dict] = [] + real_mutation_builder = gateway_services.build_checkpoint_state_mutation_accessor + + def recording_mutation_builder(request, *, thread_id, as_node, checkpoint_id=None, state_schema=None): + accessor, config = real_mutation_builder(request, thread_id=thread_id, as_node=as_node, checkpoint_id=checkpoint_id, state_schema=state_schema) + original_aupdate = accessor.aupdate + + async def recording_aupdate(config, values, *, as_node=None): + recorded_updates.append(dict(values)) + return await original_aupdate(config, values, as_node=as_node) + + accessor.aupdate = recording_aupdate + return accessor, config + + monkeypatch.setattr(gateway_services, "build_checkpoint_state_mutation_accessor", recording_mutation_builder) + monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", recording_mutation_builder) + + source_thread_id = "extension-source" + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": source_thread_id, "metadata": {}, "assistant_id": "extension-agent"}, + ) + assert created.status_code == 200, created.text + + # Seed after creation: create_thread writes an empty head checkpoint. + asyncio.run(_seed_extension_source(checkpointer, custom_factory, mode, source_thread_id)) + + read_response = client.get(f"/api/threads/{source_thread_id}/state") + assert read_response.status_code == 200, read_response.text + assert read_response.json()["values"]["ext_list"] == ["merged", "payload"] + + update_response = client.post( + f"/api/threads/{source_thread_id}/state", + json={"values": {"ext_list": ["replaced"]}}, + ) + assert update_response.status_code == 200, update_response.text + assert update_response.json()["values"]["ext_list"] == ["replaced"] + + branch_response = client.post( + f"/api/threads/{source_thread_id}/branches", + json={"message_id": "a1", "message_ids": ["a1"]}, + ) + assert branch_response.status_code == 200, branch_response.text + branch_thread_id = branch_response.json()["thread_id"] + + # The branch write must copy every reducer channel with replace semantics. + branch_update = recorded_updates[-1] + assert isinstance(branch_update["ext_list"], Overwrite) + assert branch_update["ext_list"].value == ["replaced"] + assert isinstance(branch_update["messages"], Overwrite) + + async def materialize(thread_id): + accessor = CheckpointStateAccessor.bind(custom_factory(), checkpointer, mode=mode) + snapshot = await accessor.aget({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}) + return snapshot.values + + branch_values = asyncio.run(materialize(branch_thread_id)) + assert branch_values["ext_list"] == ["replaced"] + assert [message.id for message in branch_values["messages"]] == ["h1", "a1"] + + +def test_update_thread_state_rejects_unknown_state_fields(monkeypatch) -> None: + """Unknown fields fail 422 instead of a false-success 200.""" + app, _store, checkpointer = _build_thread_app() + custom_factory = _wire_extension_agent(monkeypatch, app, checkpointer, "full") + source_thread_id = "extension-source-422" + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": source_thread_id, "metadata": {}, "assistant_id": "extension-agent"}, + ) + assert created.status_code == 200, created.text + + asyncio.run(_seed_extension_source(checkpointer, custom_factory, "full", source_thread_id)) + + response = client.post( + f"/api/threads/{source_thread_id}/state", + json={"values": {"not_a_state_field": 1}}, + ) + + assert response.status_code == 422, response.text + assert "not_a_state_field" in response.json()["detail"] + + def test_branch_display_name_strips_legacy_branch_prefix_only_for_branch_sources() -> None: assert threads._default_branch_display_name("Original chat") == "Original chat" assert threads._default_branch_display_name("Branch: Original chat") == "Branch: Original chat" @@ -1089,6 +1690,159 @@ def test_search_threads_succeeds_with_valid_metadata() -> None: # ── update_thread_state: each call inserts a new checkpoint (regression) ─────── +def test_update_thread_state_overwrites_reducer_fields_and_writes_last_values_directly(monkeypatch) -> None: + app, _store, _checkpointer = _build_thread_app() + update_calls: list[tuple[dict, dict, str | None]] = [] + updated_config = { + "configurable": { + "thread_id": "state-overwrite", + "checkpoint_ns": "", + "checkpoint_id": "ckpt-updated", + } + } + snapshot = SimpleNamespace( + values={ + "messages": [{"type": "human", "id": "h1", "content": "replacement"}], + "artifacts": ["artifact-1"], + "title": "Renamed", + }, + config=updated_config, + parent_config={"configurable": {"checkpoint_id": "ckpt-original"}}, + metadata={"source": "update", "step": 1}, + next=(), + tasks=(), + created_at="2026-07-18T00:00:00+00:00", + ) + + async def aupdate(config, values, *, as_node=None): + update_calls.append((config, values, as_node)) + return updated_config + + accessor = SimpleNamespace( + aupdate=aupdate, + aget=AsyncMock(return_value=snapshot), + ) + + async def build_accessor(_request, *, thread_id, as_node, checkpoint_id=None): + assert thread_id == "state-overwrite" + return accessor, { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + **({"checkpoint_id": checkpoint_id} if checkpoint_id else {}), + } + } + + monkeypatch.setattr(threads, "build_thread_checkpoint_state_mutation_accessor", build_accessor) + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": "state-overwrite", "metadata": {}}, + ) + assert created.status_code == 200, created.text + response = client.post( + "/api/threads/state-overwrite/state", + json={ + "values": { + "messages": [{"type": "human", "id": "h1", "content": "replacement"}], + "artifacts": ["artifact-1"], + "title": "Renamed", + } + }, + ) + + assert response.status_code == 200, response.text + assert len(update_calls) == 1 + read_config, updates, as_node = update_calls[0] + assert read_config["configurable"]["thread_id"] == "state-overwrite" + assert isinstance(updates["messages"], Overwrite) + assert updates["messages"].value[0]["id"] == "h1" + assert isinstance(updates["artifacts"], Overwrite) + assert updates["artifacts"].value == ["artifact-1"] + assert updates["title"] == "Renamed" + assert as_node == "manual_state_update" + accessor.aget.assert_awaited_once_with(updated_config) + assert response.json()["checkpoint_id"] == "ckpt-updated" + + +def test_update_thread_state_real_mutation_graph_finishes_without_scheduling(monkeypatch) -> None: + app, _store, _checkpointer = _build_thread_app() + app.state.checkpoint_channel_mode = "delta" + real_mutation_builder = gateway_services.build_checkpoint_state_mutation_accessor + + async def mutation_boundary(request, *, thread_id, as_node, checkpoint_id=None): + # No real assistant graph in this unit context: the boundary falls + # back to the base schema while writes still use the mutation graph. + return real_mutation_builder(request, thread_id=thread_id, as_node=as_node, checkpoint_id=checkpoint_id) + + monkeypatch.setattr(threads, "build_thread_checkpoint_state_mutation_accessor", mutation_boundary) + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": "state-real-mutation", "metadata": {}}, + ) + assert created.status_code == 200, created.text + response = client.post( + "/api/threads/state-real-mutation/state", + json={ + "values": { + "messages": [{"type": "human", "id": "h1", "content": "replacement"}], + "artifacts": ["artifact-1"], + "title": "Renamed", + } + }, + ) + + assert response.status_code == 200, response.text + body = response.json() + assert [message["id"] for message in body["values"]["messages"]] == ["h1"] + assert body["values"]["artifacts"] == ["artifact-1"] + assert body["values"]["title"] == "Renamed" + assert body["next"] == [] + + +@pytest.mark.parametrize("missing_checkpoint_id", ["does-not-exist", ""]) +def test_update_thread_state_rejects_missing_explicit_checkpoint_without_writing( + missing_checkpoint_id: str, +) -> None: + app, _store, checkpointer = _build_thread_app() + + with TestClient(app) as client: + created = client.post( + "/api/threads", + json={"thread_id": "state-missing-checkpoint", "metadata": {}}, + ) + assert created.status_code == 200, created.text + response = client.post( + "/api/threads/state-missing-checkpoint/state", + json={ + "checkpoint_id": missing_checkpoint_id, + "values": {"title": "Must not be written"}, + }, + ) + + assert response.status_code == 404, response.text + + async def collect_checkpoint_ids(): + return [ + item.config["configurable"]["checkpoint_id"] + async for item in checkpointer.alist( + { + "configurable": { + "thread_id": "state-missing-checkpoint", + "checkpoint_ns": "", + } + } + ) + ] + + checkpoint_ids = asyncio.run(collect_checkpoint_ids()) + assert missing_checkpoint_id not in checkpoint_ids + assert len(checkpoint_ids) == 1 + + def test_update_thread_state_inserts_new_checkpoint_each_call() -> None: """Each ``POST /state`` must INSERT a distinct, time-ordered checkpoint. diff --git a/backend/uv.lock b/backend/uv.lock index eee45d32e..086224d29 100644 --- a/backend/uv.lock +++ b/backend/uv.lock @@ -15,6 +15,7 @@ members = [ "deer-flow", "deerflow-harness", ] +overrides = [{ name = "websockets", specifier = "==16.0" }] [[package]] name = "agent-client-protocol" @@ -815,6 +816,7 @@ redis = [ [package.dev-dependencies] dev = [ { name = "blockbuster" }, + { name = "hypothesis" }, { name = "jsonschema" }, { name = "monocle-apptrace" }, { name = "prompt-toolkit" }, @@ -855,6 +857,7 @@ provides-extras = ["postgres", "redis", "discord", "monocle", "browser"] [package.metadata.requires-dev] dev = [ { name = "blockbuster", specifier = ">=1.5.26,<1.6" }, + { name = "hypothesis", specifier = ">=6.100,<7" }, { name = "jsonschema", specifier = ">=4.26.0" }, { name = "monocle-apptrace", specifier = ">=0.8.8" }, { name = "prompt-toolkit", specifier = ">=3.0.0" }, @@ -954,7 +957,7 @@ requires-dist = [ { name = "firecrawl-py", specifier = ">=1.15.0" }, { name = "httpx", specifier = ">=0.28.0" }, { name = "kubernetes", specifier = ">=30.0.0" }, - { name = "langchain", specifier = ">=1.2.15" }, + { name = "langchain", specifier = ">=1.3" }, { name = "langchain-anthropic", specifier = ">=1.4.1" }, { name = "langchain-deepseek", specifier = ">=1.0.1" }, { name = "langchain-google-genai", specifier = ">=4.2.1" }, @@ -962,10 +965,10 @@ requires-dist = [ { name = "langchain-ollama", marker = "extra == 'ollama'", specifier = ">=0.3.0" }, { name = "langchain-openai", specifier = ">=1.2.1" }, { name = "langfuse", specifier = ">=3.4.1" }, - { name = "langgraph", specifier = ">=1.1.9" }, + { name = "langgraph", specifier = ">=1.2.9,<1.3" }, { name = "langgraph-api", specifier = ">=0.8.1" }, - { name = "langgraph-checkpoint-postgres", marker = "extra == 'postgres'", specifier = ">=3.0.5" }, - { name = "langgraph-checkpoint-sqlite", specifier = ">=3.0.3" }, + { name = "langgraph-checkpoint-postgres", marker = "extra == 'postgres'", specifier = ">=3.1.0,<3.2" }, + { name = "langgraph-checkpoint-sqlite", specifier = ">=3.1.0,<3.2" }, { name = "langgraph-cli", specifier = ">=0.4.24" }, { name = "langgraph-runtime-inmem", specifier = ">=0.28.0" }, { name = "langgraph-sdk", specifier = ">=0.1.51" }, @@ -1660,6 +1663,55 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/30/47d0bf6072f7252e6521f3447ccfa40b421b6824517f82854703d0f5a98b/hyperframe-6.1.0-py3-none-any.whl", hash = "sha256:b03380493a519fce58ea5af42e4a42317bf9bd425596f7a0835ffce80f1a42e5", size = 13007, upload-time = "2025-01-22T21:41:47.295Z" }, ] +[[package]] +name = "hypothesis" +version = "6.156.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "sortedcontainers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/20/83/8dbe89bdb8c6f25a7a52e7898af6d82fe35dfef08e5c702f6e33231ce6c6/hypothesis-6.156.6.tar.gz", hash = "sha256:96de02faefa3ce079873541da96f42595583bb001e8e4219294ed7d4501cc4cc", size = 476304, upload-time = "2026-07-10T20:56:49.96Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1a/dc/0c2a851f06c91d5ac9ef0f3b9615efc1ed650411d2eee23b6334f491c85e/hypothesis-6.156.6-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:caf6a93d011c10972da111c38ceb34ced20feaa8581e2b350c0655b022e27875", size = 747998, upload-time = "2026-07-10T20:56:16.311Z" }, + { url = "https://files.pythonhosted.org/packages/8e/f8/59203ca978ab51595d12d6bc7e7a63300d7373431ab42ca3f1742e45db68/hypothesis-6.156.6-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:07f2bc9df1aeba80e12029c1618e2ee54abc440068c305d7075ffd6b85251843", size = 743073, upload-time = "2026-07-10T20:55:36.825Z" }, + { url = "https://files.pythonhosted.org/packages/68/d8/86a0023740434098d1b187a62bd5f99b198f098fb43e7fc58342283a8270/hypothesis-6.156.6-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7baca17f4803ad4aa151732326f3990baf54c3127df44aa872ac5bdf8a98a9a6", size = 1070169, upload-time = "2026-07-10T20:55:49.47Z" }, + { url = "https://files.pythonhosted.org/packages/9b/82/673453915fd0c67673f35a4876ba88f48c621335f293f3537d77b27d4286/hypothesis-6.156.6-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8083806645f84243aade727f4978185caaa0b7190af4318673999ee15fdbf424", size = 1121760, upload-time = "2026-07-10T20:55:53.502Z" }, + { url = "https://files.pythonhosted.org/packages/8a/c3/3a5557f52912f2fecc6ed59642dcf80dd8e89d0d9664502b68e23d66bf3d/hypothesis-6.156.6-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a922eedcd8618f9c2e17b79fa7b3f3f0b2df34e201958611cc3f0f46cca33c10", size = 1111440, upload-time = "2026-07-10T20:55:43.054Z" }, + { url = "https://files.pythonhosted.org/packages/38/a6/ae636d4ca7f996a1ccb4b3d5997d949f1718fba52b01559b3ab53b237b3f/hypothesis-6.156.6-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5291bd33c4704d274d7c214d5c200e77f372a06644f5cbbe96dcbe53cb2fbf10", size = 1244944, upload-time = "2026-07-10T20:55:56.109Z" }, + { url = "https://files.pythonhosted.org/packages/1e/79/c425d22d734be0268ca60d120c6296299e4220a1783cb1a4cc76232807bb/hypothesis-6.156.6-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:55f3ec50161b4a95bae63bff2b5166e45935b493013d3be30ede279bf6192318", size = 1288808, upload-time = "2026-07-10T20:56:06.249Z" }, + { url = "https://files.pythonhosted.org/packages/c5/3a/cc9f479d22cbdd36ddfc55a978378eddadd183b09339ebdb81be33bb18e7/hypothesis-6.156.6-cp310-abi3-win32.whl", hash = "sha256:e96570ca5cdd9a5f2ff9e80a6fb2fd5420ebf33b833d7de5b09b6ebb26a3eb6c", size = 634868, upload-time = "2026-07-10T20:55:37.959Z" }, + { url = "https://files.pythonhosted.org/packages/d6/89/2008d287289841a936456cb13443ca89d88da6e4527d611d482e9544164d/hypothesis-6.156.6-cp310-abi3-win_amd64.whl", hash = "sha256:32710718c22fe8c5571464e898bb87d282837b02617d6ad68130abf7cb4843cb", size = 640382, upload-time = "2026-07-10T20:55:30.634Z" }, + { url = "https://files.pythonhosted.org/packages/8c/45/9f009005b9c796f4a40424484ac7e70847bc088456fd940a937f96bb4b6d/hypothesis-6.156.6-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:a2a728b514fceb81e3f0464508911d5220fd74dadc3270f859427a686b60c4cf", size = 748844, upload-time = "2026-07-10T20:56:38.036Z" }, + { url = "https://files.pythonhosted.org/packages/02/2f/4d852bb8a9c73a68b18eca9b5b085285282122166e158f4d2a477639bfee/hypothesis-6.156.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7489b9a8f9df8227edd6c7cd8b9ccfab2483bab24da6a474c175973ca2294f58", size = 741936, upload-time = "2026-07-10T20:55:27.539Z" }, + { url = "https://files.pythonhosted.org/packages/74/89/b9968070ae042f9bf3149bb6ba6399d5f28f452e0fb7f638cafc69ff0b9a/hypothesis-6.156.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42760873d6db1069d6edbaa355a61b9078a9950259efcfc72fc695741d7db7cd", size = 1069749, upload-time = "2026-07-10T20:56:43.017Z" }, + { url = "https://files.pythonhosted.org/packages/00/a9/753806f5292b40aeab1d269e408e3a7e85be3c0d88828fb78ab4a34d6626/hypothesis-6.156.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b4e66aaa7385538a5d617174d47c198ee807f06de99e282a67c6cb724c69340d", size = 1120983, upload-time = "2026-07-10T20:56:25.424Z" }, + { url = "https://files.pythonhosted.org/packages/85/88/8386d064d680be27e936eba94f1448bc93ef6fa05473ee5034139f1c4284/hypothesis-6.156.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:08796b674c0b31a5dd4119b2173823390055921588d13eb77324e861b00fd7f8", size = 1243911, upload-time = "2026-07-10T20:55:54.799Z" }, + { url = "https://files.pythonhosted.org/packages/c1/8c/7524c1e5279e7728eb47c99f2357cbc5f08ae92e9bce49bf50118b53f9c9/hypothesis-6.156.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4ca8cc26ea2d31d22cf7710e92951cfaa921f0f8aa1b6db33a5176335f583a4f", size = 1287806, upload-time = "2026-07-10T20:56:02.176Z" }, + { url = "https://files.pythonhosted.org/packages/5a/b3/c347ad913e1c5f2988956fe17826c0400b4ce470b973e6c248e97b6a0acf/hypothesis-6.156.6-cp312-cp312-win_amd64.whl", hash = "sha256:c3363d3fb8015594636689572510bb6090602d8e8e838a5693c2d52d3b5b09d8", size = 637679, upload-time = "2026-07-10T20:55:39.056Z" }, + { url = "https://files.pythonhosted.org/packages/70/5d/9583fe153573523dac27226c89e041a86ad4aeeae08c868160cbb93d39d2/hypothesis-6.156.6-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:59a8def90d9a5a9b67e1ac529e903a2363ceb6cf873c209da6b4284c5daab671", size = 749264, upload-time = "2026-07-10T20:56:46.118Z" }, + { url = "https://files.pythonhosted.org/packages/86/35/e4113d06769b544f0fb77ffea9195b598b4c56a298905c21fd47c4eed388/hypothesis-6.156.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c574c3224563d730848bc5d1ef1683c4f83993400c0167899fe328f4bfcd4725", size = 742095, upload-time = "2026-07-10T20:56:41.412Z" }, + { url = "https://files.pythonhosted.org/packages/d8/5c/a47666ede10384e8978722cade7ab96a42df71d2ab577317092d0fed7c8a/hypothesis-6.156.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:01bb8270c46b3ef53b0c2d23ff613ea506d609d06f936d823ea57c58b66b05f7", size = 1069917, upload-time = "2026-07-10T20:56:04.845Z" }, + { url = "https://files.pythonhosted.org/packages/79/93/75f6057dadd9dc0134f37c08d5d14d04d3cd7374debbcb0cc4569c6712f1/hypothesis-6.156.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d4ea6559c13606e13b645927f2e0906e52b5ac5d99b40d3abaaeb2e8c7ceeb75", size = 1121204, upload-time = "2026-07-10T20:55:52.008Z" }, + { url = "https://files.pythonhosted.org/packages/62/87/308efef08bc60d1e673d035e8ca8e9663f4b6b3ba519c3cdebf6583c2b76/hypothesis-6.156.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2d47054d0230f0dd9b6868fc030126c7a6c25527144272ff376cc4e9c39f7540", size = 1244168, upload-time = "2026-07-10T20:55:40.288Z" }, + { url = "https://files.pythonhosted.org/packages/3b/66/de8fff5bd9a40a4056dafbe7f904887ef12632282bbbac90f1977c30dd3b/hypothesis-6.156.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:050c8c0815f88d47dd0875a92698d20d61639b7b721ee043a6d687c7f14ff7d8", size = 1288127, upload-time = "2026-07-10T20:56:00.541Z" }, + { url = "https://files.pythonhosted.org/packages/cf/8d/794fb26e1fd3ff004978f8f18b7aa7e1c2270ba72e1f977b987a812064f8/hypothesis-6.156.6-cp313-cp313-win_amd64.whl", hash = "sha256:f0d73edab7b8a0051b3634f2d04d62b7e7282f8f274963b11188ee4957d672ef", size = 637954, upload-time = "2026-07-10T20:56:33.35Z" }, + { url = "https://files.pythonhosted.org/packages/a5/be/5b4b27984cb43c60e95f570b069660335dad34cb67f7d226017c5d35d31e/hypothesis-6.156.6-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:34a70a7b8226e34d658072d8fb81d03f97f0a75ceb536329a321b94ce2232fd6", size = 749312, upload-time = "2026-07-10T20:55:46.902Z" }, + { url = "https://files.pythonhosted.org/packages/31/11/709cceffc28666c9d4cb75ffc6df5ce30db8c7dd5cc2c8b38a2fd837427f/hypothesis-6.156.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f1969646beead7d8cf6a2537d2765af89d73056e2cb218e7fae92b83802250a3", size = 742332, upload-time = "2026-07-10T20:56:30.254Z" }, + { url = "https://files.pythonhosted.org/packages/84/52/cfc79b13d8dd3cd6de6b9df921c557efe8528a9c90a3a7cd93b37188d57e/hypothesis-6.156.6-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cbc2ec7b7d905e6b6ec1635f6340bfa52aaab718101c59f052bc012a6b486cd8", size = 1070109, upload-time = "2026-07-10T20:55:48.244Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ac/1da4def1f006b5ad01187ff96379e24c37439d659ec10c3e944c03436c0f/hypothesis-6.156.6-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9367ae25dfa6dc1af37904785e43c4f8fe1c4118cafdc2f06514154fbdd90992", size = 1121528, upload-time = "2026-07-10T20:56:39.665Z" }, + { url = "https://files.pythonhosted.org/packages/68/47/744e4f5e3d635dea20dbedf3fa486e2a6fa5210e0a52a0d5c4da56babd84/hypothesis-6.156.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:455f09107ec07c78f2a83cb8fc19e23879c9d51cdc831de6f9cb6ec4059cb9af", size = 1244690, upload-time = "2026-07-10T20:56:31.854Z" }, + { url = "https://files.pythonhosted.org/packages/25/8a/42252fcd5e521d140dac532f29c2a13ca8f22908cb545ffdd64b5e225680/hypothesis-6.156.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c76634c45a3ceee4c4fdfed39aebd08b8b822ec8b0c556877ef82846fd777730", size = 1288519, upload-time = "2026-07-10T20:56:03.429Z" }, + { url = "https://files.pythonhosted.org/packages/44/e7/176df9e47cf583d2b8d234b78c0aac3a47075ad5d147e60b2c21a1338bb1/hypothesis-6.156.6-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:eb7e9f8343bc6b948937e6ec12e6879ed25a17b53ceccbd2b84adadd3d511698", size = 586452, upload-time = "2026-07-10T20:56:22.285Z" }, + { url = "https://files.pythonhosted.org/packages/8c/75/2c8a0411bbe76429f3ae738ef9a00107201bf6146d9534350014ce369d98/hypothesis-6.156.6-cp314-cp314-win_amd64.whl", hash = "sha256:f9631cd604ae6032c3edf99160dc1b9e33873f2e52762246b24f07fb758652ae", size = 637774, upload-time = "2026-07-10T20:56:34.97Z" }, + { url = "https://files.pythonhosted.org/packages/2a/22/8115005e9aa72c8d63d90e9db5e0b8425fd8950fbc5d6e332805d4d32c9e/hypothesis-6.156.6-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:1f81163d36d3763b09ffaef7c3a71e88174ca3e6816201fca9d1d159f448fdb5", size = 747428, upload-time = "2026-07-10T20:56:44.611Z" }, + { url = "https://files.pythonhosted.org/packages/e8/c2/66bfe9337f4a4b1f7754ee6d01d950280152a81d0d797e6c1d9eb0909750/hypothesis-6.156.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:556b905767e36147918634a64356aa05d8c956576f00aee01eb351678f193908", size = 740889, upload-time = "2026-07-10T20:55:57.656Z" }, + { url = "https://files.pythonhosted.org/packages/95/3b/69f45af2d4f0950b7d1af3cdbdd800b88a6c2370331481eda79d6171fbe3/hypothesis-6.156.6-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c3f2604b28d16d696aaaf4954d20f907b27e54034df98e64746a20c74c319f03", size = 1069270, upload-time = "2026-07-10T20:56:12.024Z" }, + { url = "https://files.pythonhosted.org/packages/f8/43/6b2549885da08f5e50ba34fb8e0d0a60b2f190ffd516fac220f8db5b5869/hypothesis-6.156.6-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ffe012ad66dbe7b8e8ddef6f6992ab1b36719ea64430c2bf1ff7135521052a15", size = 1120409, upload-time = "2026-07-10T20:55:34.551Z" }, + { url = "https://files.pythonhosted.org/packages/70/97/745c778c3eb29befa2367b1ded8437eecfbbe6932359d0f831275bc32170/hypothesis-6.156.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5bfa3c7b758f7278081c6bfec5f89b43c4eb075c0c9eb095323f7a9eb019b513", size = 1243111, upload-time = "2026-07-10T20:56:17.83Z" }, + { url = "https://files.pythonhosted.org/packages/ab/d7/c5ec6a442dc9b822f47064bda4b6d3e739dccdd1c5bf44c9a57fb6136830/hypothesis-6.156.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d0db1f4573800c618773622f03cb6533bb3377430ef938c9476ba10c39d22591", size = 1287262, upload-time = "2026-07-10T20:56:23.749Z" }, + { url = "https://files.pythonhosted.org/packages/11/0c/c134d61710e14b68b010215dcf6bd57d2ec05cd169dff8bfab8fefc2d410/hypothesis-6.156.6-cp314-cp314t-win_amd64.whl", hash = "sha256:38cd0c4a7b9f809f1e23a4d15adfa9c5d99869b9afc327350a5e563350b78e48", size = 637862, upload-time = "2026-07-10T20:56:13.347Z" }, +] + [[package]] name = "idna" version = "3.15" @@ -1872,16 +1924,16 @@ wheels = [ [[package]] name = "langchain" -version = "1.2.15" +version = "1.3.14" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "langgraph" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/98/3f/888a7099d2bd2917f8b0c3ffc7e347f1e664cf64267820b0b923c4f339fc/langchain-1.2.15.tar.gz", hash = "sha256:1717b6719daefae90b2728314a5e2a117ff916291e2862595b6c3d6fba33d652", size = 574732, upload-time = "2026-04-03T14:26:03.994Z" } +sdist = { url = "https://files.pythonhosted.org/packages/29/68/a6dbad9c22df4087a0f9e79ddd46226c442b30128bfeee538d5889492a73/langchain-1.3.14.tar.gz", hash = "sha256:1b6696c72ba3bbbce54d745e0180742c9f6ece8bbc59ed5a46c3e20b9a435929", size = 645181, upload-time = "2026-07-16T13:28:18.29Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/e8/a3b8cb0005553f6a876865073c81ef93bd7c5b18381bcb9ba4013af96ebc/langchain-1.2.15-py3-none-any.whl", hash = "sha256:e349db349cb3e9550c4044077cf90a1717691756cc236438404b23500e615874", size = 112714, upload-time = "2026-04-03T14:26:02.557Z" }, + { url = "https://files.pythonhosted.org/packages/a9/ec/0f942e78a621f8e3162ff1ed24284f469aaf51fb4607ee5831c626f2b2bc/langchain-1.3.14-py3-none-any.whl", hash = "sha256:4d10dbe91005952cddd56d0dc77aa108964da6bae90ab20063653957e901f782", size = 139560, upload-time = "2026-07-16T13:28:16.498Z" }, ] [[package]] @@ -1900,7 +1952,7 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.3.3" +version = "1.4.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonpatch" }, @@ -1913,9 +1965,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "uuid-utils" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d3/ae/8b74458fc3850ec3d150eb9f45e857db129dafa801fb5cf173dfc9f8bbf3/langchain_core-1.3.3.tar.gz", hash = "sha256:fa510a5db8efdc0c6ff41c0939fb5c00a0183c11f6b84233e892e3227ff69182", size = 915041, upload-time = "2026-05-05T19:02:36.612Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2a/b9/e937d0a90b26540bff07e7a7c64349f3b29c2dcc36257cd1cd3fdce17f2a/langchain_core-1.4.9.tar.gz", hash = "sha256:f8078901145bed0466755277500a5a22822a7b628808c4c0a28d4fc88895fcf2", size = 967294, upload-time = "2026-07-08T20:06:54.191Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/01/4771b7ab2af1d1aba5b710bd8f13d9225c609425214b357590a17b01be77/langchain_core-1.3.3-py3-none-any.whl", hash = "sha256:18aae8506f37da7f74398492279a7d6efcee4f8e23c4c41c7af080eeb7ef7bd1", size = 543857, upload-time = "2026-05-05T19:02:34.52Z" }, + { url = "https://files.pythonhosted.org/packages/84/70/ade2fada52772798ef815b6352b59e71b116aa0c32c3aef5be3dc2cbed12/langchain_core-1.4.9-py3-none-any.whl", hash = "sha256:28e3909e2a10cc81504952d795ac0a9e014c0018121ef89d48dd396fa09ec624", size = 558293, upload-time = "2026-07-08T20:06:52.382Z" }, ] [[package]] @@ -2020,7 +2072,7 @@ wheels = [ [[package]] name = "langgraph" -version = "1.1.9" +version = "1.2.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, @@ -2030,9 +2082,9 @@ dependencies = [ { name = "pydantic" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/8c/d5/9d9c65d5500a1ca7ea63d6d65aecfb248037018a74d7d4ef52e276bb4e4b/langgraph-1.1.9.tar.gz", hash = "sha256:bc5a49d5a5e71fda1f9c53c06c62f4caec9a95545b739d130a58b6ab3269e274", size = 560717, upload-time = "2026-04-21T13:43:06.809Z" } +sdist = { url = "https://files.pythonhosted.org/packages/41/4b/0d1130e26b41a99dcc88353bbe7162a1f255c4db746bd94024268e6af27b/langgraph-1.2.9.tar.gz", hash = "sha256:385f87bc1802c35af7e0aa479278ecba8582d103515eb48256cb2ddcd42d0bd4", size = 722869, upload-time = "2026-07-10T01:30:14.985Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/16/58/0380420e66619d12c992c1f8cfda0c7a04e8f0fe8a84752245b9e7b1cba7/langgraph-1.1.9-py3-none-any.whl", hash = "sha256:7db13ceecde4ea643df6c097dcc9e534895dcd9fcc6500eeff2f2cde0fab16b2", size = 173744, upload-time = "2026-04-21T13:43:05.513Z" }, + { url = "https://files.pythonhosted.org/packages/41/16/0b8dc48823f1326f3e0c8012a3c07a40da6f194299e2ec080df236287baf/langgraph-1.2.9-py3-none-any.whl", hash = "sha256:c2d98ad94333937922ba04148641c1da2bfe45b5b8e55d7b6dcb0bb2df809e76", size = 247473, upload-time = "2026-07-10T01:30:13.733Z" }, ] [[package]] @@ -2093,7 +2145,7 @@ wheels = [ [[package]] name = "langgraph-checkpoint-postgres" -version = "3.0.5" +version = "3.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langgraph-checkpoint" }, @@ -2101,23 +2153,23 @@ dependencies = [ { name = "psycopg" }, { name = "psycopg-pool" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/95/7a/8f439966643d32111248a225e6cb33a182d07c90de780c4dbfc1e0377832/langgraph_checkpoint_postgres-3.0.5.tar.gz", hash = "sha256:a8fd7278a63f4f849b5cbc7884a15ca8f41e7d5f7467d0a66b31e8c24492f7eb", size = 127856, upload-time = "2026-03-18T21:25:29.785Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6d/51/5a2dc42e8b5d5942b933b5b7237eae5a4dbc92508a04727c263dd383ad8a/langgraph_checkpoint_postgres-3.1.0.tar.gz", hash = "sha256:02bff4ab63d9dae8eab3a9640fce1d479da8965c9fba7b0dc04cb1f7c56f0a55", size = 148473, upload-time = "2026-05-12T03:40:10.599Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/87/b0f98b33a67204bca9d5619bcd9574222f6b025cf3c125eedcec9a50ecbc/langgraph_checkpoint_postgres-3.0.5-py3-none-any.whl", hash = "sha256:86d7040a88fd70087eaafb72251d796696a0a2d856168f5c11ef620771411552", size = 42907, upload-time = "2026-03-18T21:25:28.75Z" }, + { url = "https://files.pythonhosted.org/packages/2f/cd/eff9b82bc3b5f62d481b437099f44f3ef7b1d907f166fb4ee25e8f84a1e7/langgraph_checkpoint_postgres-3.1.0-py3-none-any.whl", hash = "sha256:814cce2ef35d792bf07b090a95eed004f1acac0724fe6605536b13f6d1e7032c", size = 48988, upload-time = "2026-05-12T03:40:08.925Z" }, ] [[package]] name = "langgraph-checkpoint-sqlite" -version = "3.0.3" +version = "3.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiosqlite" }, { name = "langgraph-checkpoint" }, { name = "sqlite-vec" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/04/61/40b7f8f29d6de92406e668c35265f409f57064907e31eae84ab3f2a3e3e1/langgraph_checkpoint_sqlite-3.0.3.tar.gz", hash = "sha256:438c234d37dabda979218954c9c6eb1db73bee6492c2f1d3a00552fe23fa34ed", size = 123876, upload-time = "2026-01-19T00:38:44.473Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/ea/83917c2369acf8a10a894d4247655fd063c07924ba5bc4e83c85d2eaeded/langgraph_checkpoint_sqlite-3.1.0.tar.gz", hash = "sha256:f926916ebc1b985d802cc9c820026036e84db9d910d62c97b57e4ba64f67d5ae", size = 147902, upload-time = "2026-05-12T03:34:52.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/d8/84ef22ee1cc485c4910df450108fd5e246497379522b3c6cfba896f71bf6/langgraph_checkpoint_sqlite-3.0.3-py3-none-any.whl", hash = "sha256:02eb683a79aa6fcda7cd4de43861062a5d160dbbb990ef8a9fd76c979998a952", size = 33593, upload-time = "2026-01-19T00:38:43.288Z" }, + { url = "https://files.pythonhosted.org/packages/97/07/b342811a16327900af2747c752ea19676172fcddf9b592cc384031076623/langgraph_checkpoint_sqlite-3.1.0-py3-none-any.whl", hash = "sha256:cc9b40df0076feae8a9ad42ae713621b148b00ac23adc09dc1dc66090a46e5ad", size = 38587, upload-time = "2026-05-12T03:34:51.231Z" }, ] [[package]] @@ -2138,15 +2190,15 @@ wheels = [ [[package]] name = "langgraph-prebuilt" -version = "1.0.11" +version = "1.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "langgraph-checkpoint" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/8d/bb/0e0b3eb33b1f2f32f8810a49aa24b7d11a5b0ed45f679386095946a59557/langgraph_prebuilt-1.0.11.tar.gz", hash = "sha256:0e71545f706a134b6a80a2a56916562797b499e3e4ab6eed5ce89396ac03d322", size = 171759, upload-time = "2026-04-24T18:18:34.528Z" } +sdist = { url = "https://files.pythonhosted.org/packages/29/66/ed9b93f56bc17ef22d551892f0ac2b225a97fe0fcf23a511b857f70d590b/langgraph_prebuilt-1.1.0.tar.gz", hash = "sha256:3c579cf6eed2d17f9c157c2d0fcaddcd8688524e7022d3b22b37a3bf4589d528", size = 178833, upload-time = "2026-05-12T03:37:49.332Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/8c/f4c574cb75ae9b8a474215d03a029ea723c919f65771ca1c82fe532d0297/langgraph_prebuilt-1.0.11-py3-none-any.whl", hash = "sha256:7afbaf5d64959e452976664c75bb8ec24098d3510cf9c205919baf443e7342ec", size = 36832, upload-time = "2026-04-24T18:18:33.586Z" }, + { url = "https://files.pythonhosted.org/packages/e9/43/3fe1a700b8490ed02679cdbbc8c915eb23a092faf496c9c1118abcd10be3/langgraph_prebuilt-1.1.0-py3-none-any.whl", hash = "sha256:51e311747d755b751d5c6b39b0c1446124d3a7643d2515017e6714b323508fc9", size = 41043, upload-time = "2026-05-12T03:37:48.007Z" }, ] [[package]] @@ -2169,15 +2221,18 @@ wheels = [ [[package]] name = "langgraph-sdk" -version = "0.3.13" +version = "0.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, + { name = "langchain-core" }, + { name = "langchain-protocol" }, { name = "orjson" }, + { name = "websockets" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0e/db/77a45127dddcfea5e4256ba916182903e4c31dc4cfca305b8c386f0a9e53/langgraph_sdk-0.3.13.tar.gz", hash = "sha256:419ca5663eec3cec192ad194ac0647c0c826866b446073eb40f384f950986cd5", size = 196360, upload-time = "2026-04-07T20:34:18.766Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b4/2b/bd8ac26d4e97f6df88ef05ce5b6a38945a3903e1025d926f4752aa88aa97/langgraph_sdk-0.4.2.tar.gz", hash = "sha256:b88f0f5f6328ac0680d6790614a905b2bcfa257f2276dba4e38f0e86db0aa738", size = 348327, upload-time = "2026-06-01T17:51:19.856Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fe/ef/64d64e9f8eea47ce7b939aa6da6863b674c8d418647813c20111645fcc62/langgraph_sdk-0.3.13-py3-none-any.whl", hash = "sha256:aee09e345c90775f6de9d6f4c7b847cfc652e49055c27a2aed0d981af2af3bd0", size = 96668, upload-time = "2026-04-07T20:34:17.866Z" }, + { url = "https://files.pythonhosted.org/packages/a0/05/aac507337cceae773c2cc9ab91eb6301963af7aeeb55b4217a00e15aff17/langgraph_sdk-0.4.2-py3-none-any.whl", hash = "sha256:75fa5096c1177ce39c847096a8fe3745ffd480ddb412995f836e9f5f884c43dd", size = 160521, upload-time = "2026-06-01T17:51:18.849Z" }, ] [[package]] @@ -4260,6 +4315,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/37/c3/6eeb6034408dac0fa653d126c9204ade96b819c936e136c5e8a6897eee9c/socksio-1.0.0-py3-none-any.whl", hash = "sha256:95dc1f15f9b34e8d7b16f06d74b8ccf48f609af32ab33c608d08761c5dcbb1f3", size = 12763, upload-time = "2020-04-17T15:50:31.878Z" }, ] +[[package]] +name = "sortedcontainers" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e8/c4/ba2f8066cceb6f23394729afe52f3bf7adec04bf9ed2c820b39e19299111/sortedcontainers-2.4.0.tar.gz", hash = "sha256:25caa5a06cc30b6b83d11423433f65d1f9d76c4c6a0c90e3379eaa43b9bfdb88", size = 30594, upload-time = "2021-05-16T22:03:42.897Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/46/9cb0e58b2deb7f82b84065f37f3bffeb12413f947f9388e4cac22c4621ce/sortedcontainers-2.4.0-py2.py3-none-any.whl", hash = "sha256:a163dcaede0f1c021485e957a39245190e74249897e2ae4b2aa38595db237ee0", size = 29575, upload-time = "2021-05-16T22:03:41.177Z" }, +] + [[package]] name = "soupsieve" version = "2.8.4" diff --git a/config.example.yaml b/config.example.yaml index 0b26bb31e..ac7c1e5a9 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1797,6 +1797,9 @@ database: # App ORM PostgreSQL command timeout. Set to null to disable it or raise it # for intentionally long commands. command_timeout: 30 + # Restart required. Use one value across every process sharing this database. + # full: current full-message checkpoints; delta: DeltaChannel for messages. + checkpoint_channel_mode: full # ============================================================================ # Run Events Configuration