mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-25 07:28:07 +00:00
fix(gateway): preserve regenerate state in branched threads (#4358)
* fix(gateway): preserve regenerate state in branched threads * test(gateway): isolate branch regenerate regression config * fix(gateway): preserve branching for legacy histories * fix(gateway): harden branch regenerate lineage * docs(gateway): clarify branch checkpoint behavior --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
Huixin615
GitHub
parent
959f052750
commit
fbc1463809
@ -693,7 +693,7 @@ Interrupted first-turn runs still persist a fallback conversation title, so stop
|
||||
|
||||
Streaming Markdown responses animate only newly arrived words; text that is already visible is not faded out and replayed when the next chunk extends the same block.
|
||||
|
||||
In the Web UI, completed assistant turns can be branched into a new main conversation. The new thread starts from that turn's checkpoint. Because workspace files are not checkpointed, the branch only receives a best-effort copy of the current workspace when you branch from the latest turn; branching from an older turn keeps just the restored message history so the branch never inherits files that were created in a later part of the conversation.
|
||||
In the Web UI, completed assistant turns can be branched into a new main conversation. The new thread starts from that turn's checkpoint and keeps the preceding replay checkpoint, so the branched response can be regenerated immediately. Legacy or imported histories without checkpoint parent links use a bounded chronological fallback; if no earlier replay checkpoint exists, branching still succeeds with the legacy single-checkpoint shape, while regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are left unchanged rather than attempting an unsafe checkpoint copy. Because workspace files are not checkpointed, the branch only receives a best-effort copy of the current workspace when you branch from the latest turn; branching from an older turn keeps just the restored message history so the branch never inherits files that were created in a later part of the conversation.
|
||||
|
||||
The Web UI reports completed task time once per run. This is total wall-clock time—including model reasoning, tool calls, and waiting—not a per-step or model-only thinking duration. Reasoning content remains available through its own separate disclosure.
|
||||
|
||||
|
||||
@ -548,7 +548,7 @@ Tools 也是同样的思路。DeerFlow 自带一组核心工具：网页搜索
|
||||
|
||||
Gateway 生成后续建议时，现在会先把普通字符串输出和 block/list 风格的富文本内容统一归一化，再去解析 JSON 数组响应，因此不同 provider 的内容包装方式不会再悄悄把建议吞掉。
|
||||
|
||||
Web UI 支持从已完成的 assistant 回复分叉出一个新的主对话。新 thread 会从该回复对应的 checkpoint 开始，并尽力复制当前 thread 的工作区文件。
|
||||
Web UI 支持从已完成的 assistant 回复分叉出一个新的主对话。新 thread 会保留该轮回复的 checkpoint 以及用户消息之前的重放 checkpoint，因此分叉后可以立即重新生成该回复。对于缺少 checkpoint 父链接的旧历史或导入历史，Gateway 会进行有界的时间顺序查找；如果不存在更早的重放 checkpoint，分叉仍会按旧版单-checkpoint 形态成功创建，但无法重新生成继承的回复。已有的单-checkpoint 分叉会保持不变，不会通过不安全的 checkpoint 复制尝试修复。只有从最新回合分叉时才会尽力复制当前 thread 的工作区文件；从历史回合分叉不会带入后续时间线创建的文件。
|
||||
|
||||
```text
|
||||
# sandbox 容器内的路径
|
||||
|
||||
@ -417,7 +417,7 @@ Localhost persistence deliberately reads the direct request `Host` and ignores `
|
||||
| **Skills** (`/api/skills`) | `GET /` - list skills; `GET /{name}` - details; `PUT /{name}` - update enabled; `POST /install` - install from .skill archive (accepts standard optional frontmatter like `version`, `author`, `compatibility`); `POST /reload` - admin-only process-local prompt-cache invalidation after trusted external filesystem changes |
|
||||
| **Memory** (`/api/memory`) | `GET /` - memory data; `POST /reload` - force reload; `GET /config` - config; `GET /status` - config + data |
|
||||
| **Uploads** (`/api/threads/{id}/uploads`) | `POST /` - upload files (auto-converts PDF/PPT/Excel/Word); `GET /list` - list; `DELETE /{filename}` - delete |
|
||||
| **Threads** (`/api/threads/{id}`) | `DELETE /` - remove DeerFlow-managed local thread data after LangGraph thread deletion; `POST /branches` - create a new main-thread branch from a completed assistant turn checkpoint. Workspace files are not checkpointed, so the branch only best-effort copies the current workspace when branching from the **latest** turn (`workspace_clone_mode="current_thread_best_effort"`); branching from an older/historical turn skips the copy (`workspace_clone_mode="skipped_historical_turn"`) so the branch never inherits files that only exist in a later timeline. Branch creation also seeds the new thread's run-event feed from the branch checkpoint's visible messages (`history_seed_mode` in the response): the thread feed reads run_events, not checkpoints, so without the seed the inherited history disappears from the UI after the branch's first run (#4380); `GET /goal`, `PUT /goal`, `DELETE /goal` - read, set, and clear the active thread goal; `POST /compact` - manually summarize older active context into `summary_text` and retain the recent message window, blocked while a run is in flight; unexpected failures are logged server-side and return a generic 500 detail |
|
||||
| **Threads** (`/api/threads/{id}`) | `DELETE /` - remove DeerFlow-managed local thread data after LangGraph thread deletion; `POST /branches` - create a new main-thread branch from a completed assistant turn checkpoint and, when an addressable pre-user replay checkpoint exists, materialize it into the branch namespace so the inherited response remains regeneratable. Workspace files are not checkpointed, so the branch only best-effort copies the current workspace when branching from the **latest** turn (`workspace_clone_mode="current_thread_best_effort"`); branching from an older/historical turn skips the copy (`workspace_clone_mode="skipped_historical_turn"`) so the branch never inherits files that only exist in a later timeline. Branch creation also seeds the new thread's run-event feed from the branch checkpoint's visible messages (`history_seed_mode` in the response): the thread feed reads run_events, not checkpoints, so without the seed the inherited history disappears from the UI after the branch's first run (#4380); `GET /goal`, `PUT /goal`, `DELETE /goal` - read, set, and clear the active thread goal; `POST /compact` - manually summarize older active context into `summary_text` and retain the recent message window, blocked while a run is in flight; unexpected failures are logged server-side and return a generic 500 detail |
|
||||
| **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - serve artifacts; active content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) are always forced as download attachments to reduce XSS risk; `?download=true` still forces download for other file types |
|
||||
| **Suggestions** (`/api/suggestions`) | `GET /config` - returns global suggestions config boolean; `POST /threads/{id}/suggestions` - generate follow-up questions; rich list/block model content is normalized and inline reasoning (`<think>...</think>`, including unclosed/truncated blocks from reasoning models like MiniMax-M3) is stripped before JSON parsing |
|
||||
| **Input Polish** (`/api/input-polish`) | `POST /` - rewrite a composer draft before it is sent. This is a short authenticated `runs:create` LLM request using `input_polish` config; it does not create a LangGraph run, persist a message, or modify thread state. Shares the non-graph one-shot LLM path (`deerflow.utils.oneshot_llm.run_oneshot_llm`) with the suggestions route so model build + Langfuse metadata + invoke stay in one place; validates the same stripped view of the draft it sends to the model, and preserves literal `<think>` substrings in the rewrite (`strip_think_blocks(truncate_unclosed=False)`) |
|
||||
@ -449,6 +449,16 @@ metadata only.
|
||||
|
||||
Proxied through nginx: `/api/langgraph/*` → Gateway LangGraph-compatible runtime, all other `/api/*` → Gateway REST APIs.
|
||||
|
||||
**Branch/regenerate checkpoint invariant**: `app/gateway/checkpoint_lineage.py`
|
||||
walks `parent_config` rather than globally ordered checkpoint history so replay
|
||||
anchors stay on the selected lineage after regenerations create sibling branches.
|
||||
New conversation branches persist the pre-user replay anchor before their visible
|
||||
head through the state mutation graph, which preserves materialized state in both
|
||||
full and delta checkpoint modes. Only an explicitly absent legacy parent link may
|
||||
use chronological compatibility lookup; cycles, dangling links, and depth-limit
|
||||
exhaustion fail closed. Existing single-checkpoint branches are never repaired by
|
||||
copying a raw checkpoint because delta state is not self-contained in one tuple.
|
||||
|
||||
### Sandbox System (`packages/harness/deerflow/sandbox/`)
|
||||
|
||||
**Interface**: Abstract `Sandbox` with `execute_command(command, env=None)`, `read_file`, `write_file`, `list_dir`. The optional `env` injects per-call environment variables (request-scoped secrets — see Request-Scoped Secrets below); `LocalSandbox` merges it via `subprocess.run(env=...)` and `AioSandbox` routes env-bearing commands through the `bash.exec(env=...)` API on a fresh session.
|
||||
@ -811,6 +821,8 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
|
||||
|
||||
**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).
|
||||
|
||||
**Replay checkpoint lookup prefers lineage and degrades only for an explicitly missing legacy parent link.** Branch and regenerate paths first walk `parent_config`, which prevents a global chronological scan from selecting a sibling created by regeneration. `CheckpointParentMissingError` alone enables the bounded newest-first history fallback in `app/gateway/checkpoint_lineage.py`; cycles, dangling/non-addressable parents, target mismatches, and depth exhaustion raise `CheckpointLineageIntegrityError` and fail closed instead of selecting a sibling. The compatibility scans request 400 raw checkpoints so up to 200 duration-only entries do not consume the effective branch-history budget; the fallback scans oldest-to-newest internally, skips duration-only checkpoints, and accepts only checkpoints with an addressable id as the replay base. A source history with no discoverable pre-user checkpoint preserves the historical single-checkpoint branch behavior instead of rejecting the branch; regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are not mutated by regenerate preparation, and no raw checkpoint tuple is copied across threads because delta state depends on ancestry and pending writes. Regenerate source-run lookup uses the current thread's exact event, then the server-stamped `run_id` on the copied human message, then verified RunManager content matching; it does not read parent-thread events. Storage or checkpoint-mode failures are not treated as a missing base and still fail closed.
|
||||
|
||||
**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.
|
||||
|
||||
164
backend/app/gateway/checkpoint_lineage.py
Normal file
164
backend/app/gateway/checkpoint_lineage.py
Normal file
@ -0,0 +1,164 @@
|
||||
"""Shared helpers for resolving replay checkpoints on one checkpoint lineage."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
|
||||
class CheckpointLineageError(RuntimeError):
|
||||
"""Raised when a requested checkpoint ancestor cannot be resolved safely."""
|
||||
|
||||
|
||||
class CheckpointParentMissingError(CheckpointLineageError):
|
||||
"""Raised when a legacy checkpoint does not record its parent link."""
|
||||
|
||||
|
||||
class CheckpointLineageIntegrityError(CheckpointLineageError):
|
||||
"""Raised when recorded checkpoint lineage is present but unsafe to use."""
|
||||
|
||||
|
||||
def checkpoint_messages(checkpoint_tuple: Any) -> list[Any]:
|
||||
values = getattr(checkpoint_tuple, "values", None)
|
||||
if isinstance(values, dict):
|
||||
messages = values.get("messages", [])
|
||||
return list(messages) if isinstance(messages, list) else []
|
||||
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 []
|
||||
return list(messages) if isinstance(messages, list) else []
|
||||
|
||||
|
||||
def checkpoint_configurable(checkpoint_tuple: Any) -> dict[str, Any]:
|
||||
config = getattr(checkpoint_tuple, "config", None) or {}
|
||||
configurable = config.get("configurable", {}) if isinstance(config, dict) else {}
|
||||
return dict(configurable) if isinstance(configurable, dict) else {}
|
||||
|
||||
|
||||
def checkpoint_metadata(checkpoint_tuple: Any) -> dict[str, Any]:
|
||||
metadata = getattr(checkpoint_tuple, "metadata", None) or {}
|
||||
return dict(metadata) if isinstance(metadata, dict) else {}
|
||||
|
||||
|
||||
def is_duration_only_checkpoint(checkpoint_tuple: Any) -> bool:
|
||||
writes = checkpoint_metadata(checkpoint_tuple).get("writes")
|
||||
return isinstance(writes, dict) and "runtime_run_duration" in writes
|
||||
|
||||
|
||||
def _message_id(message: Any) -> str | None:
|
||||
value = getattr(message, "id", None)
|
||||
if value is None and isinstance(message, dict):
|
||||
value = message.get("id")
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
def _config_identity(config: dict[str, Any]) -> tuple[str, str, str] | None:
|
||||
configurable = config.get("configurable", {})
|
||||
thread_id = configurable.get("thread_id")
|
||||
checkpoint_ns = configurable.get("checkpoint_ns", "")
|
||||
checkpoint_id = configurable.get("checkpoint_id")
|
||||
if not isinstance(thread_id, str) or not thread_id or not isinstance(checkpoint_id, str) or not checkpoint_id:
|
||||
return None
|
||||
return thread_id, str(checkpoint_ns or ""), checkpoint_id
|
||||
|
||||
|
||||
def _checkpoint_identity(checkpoint_tuple: Any) -> tuple[str, str, str] | None:
|
||||
return _config_identity(getattr(checkpoint_tuple, "config", {}) or {})
|
||||
|
||||
|
||||
def _checkpoint_exists(checkpoint_tuple: Any) -> bool:
|
||||
"""Distinguish a persisted empty checkpoint from an accessor miss.
|
||||
|
||||
LangGraph represents a missing explicit ``checkpoint_id`` as an empty
|
||||
snapshot that echoes the requested config. Persisted snapshots always
|
||||
carry metadata, a creation timestamp, or a raw checkpoint payload.
|
||||
"""
|
||||
|
||||
explicit = getattr(checkpoint_tuple, "checkpoint_exists", None)
|
||||
if isinstance(explicit, bool):
|
||||
return explicit
|
||||
if getattr(checkpoint_tuple, "metadata", None) is not None:
|
||||
return True
|
||||
if getattr(checkpoint_tuple, "created_at", None) is not None:
|
||||
return True
|
||||
return isinstance(getattr(checkpoint_tuple, "checkpoint", None), dict)
|
||||
|
||||
|
||||
async def find_checkpoint_before_message(
|
||||
accessor: Any,
|
||||
head_checkpoint: Any,
|
||||
message_id: str,
|
||||
*,
|
||||
max_depth: int,
|
||||
) -> Any:
|
||||
"""Walk one parent lineage and return the first checkpoint before ``message_id``.
|
||||
|
||||
Following ``parent_config`` is important after a regenerate: a thread can contain
|
||||
sibling checkpoint branches, and a global time-ordered scan can otherwise select
|
||||
a checkpoint from the wrong branch. Duration-only metadata checkpoints do not
|
||||
represent an addressable conversation state and are skipped.
|
||||
"""
|
||||
|
||||
if message_id not in {_message_id(message) for message in checkpoint_messages(head_checkpoint)}:
|
||||
raise CheckpointLineageIntegrityError("Target message is not present in the checkpoint head")
|
||||
|
||||
current = head_checkpoint
|
||||
visited: set[tuple[str, str, str]] = set()
|
||||
current_identity = _checkpoint_identity(current)
|
||||
if current_identity is not None:
|
||||
visited.add(current_identity)
|
||||
|
||||
# Each step performs one ancestor read, but normal branch/regenerate
|
||||
# histories cross the target boundary within 1–3 reads. Keep max_depth as
|
||||
# a conservative safety cap for valid histories with many intermediate or
|
||||
# duration-only checkpoints.
|
||||
for _ in range(max_depth):
|
||||
parent_config = getattr(current, "parent_config", None)
|
||||
if not isinstance(parent_config, dict):
|
||||
raise CheckpointParentMissingError("Checkpoint lineage ended before the target message")
|
||||
|
||||
parent = await accessor.aget(parent_config)
|
||||
parent_identity = _checkpoint_identity(parent)
|
||||
requested_parent_identity = _config_identity(parent_config)
|
||||
if parent_identity is None or not _checkpoint_exists(parent) or (requested_parent_identity is not None and parent_identity != requested_parent_identity):
|
||||
raise CheckpointLineageIntegrityError("Checkpoint parent link is not addressable")
|
||||
if parent_identity is not None:
|
||||
if parent_identity in visited:
|
||||
raise CheckpointLineageIntegrityError("Checkpoint lineage contains a cycle")
|
||||
visited.add(parent_identity)
|
||||
|
||||
if is_duration_only_checkpoint(parent):
|
||||
current = parent
|
||||
continue
|
||||
|
||||
parent_message_ids = {_message_id(message) for message in checkpoint_messages(parent)}
|
||||
if message_id not in parent_message_ids:
|
||||
return parent
|
||||
current = parent
|
||||
|
||||
raise CheckpointLineageIntegrityError(f"Checkpoint lineage exceeded the scan limit ({max_depth})")
|
||||
|
||||
|
||||
def find_checkpoint_before_message_chronologically(
|
||||
checkpoints: Sequence[Any],
|
||||
message_id: str,
|
||||
) -> tuple[Any | None, bool]:
|
||||
"""Return ``(replay_base, target_found)`` from newest-first history.
|
||||
|
||||
This is a compatibility fallback for imported or legacy checkpoints that do
|
||||
not carry ``parent_config`` links. Callers must prefer the lineage walk when
|
||||
links are available because a chronological scan cannot distinguish sibling
|
||||
checkpoint branches. Duration-only checkpoints are ignored, and only
|
||||
checkpoints with an addressable id can become the replay base.
|
||||
"""
|
||||
|
||||
previous_checkpoint = None
|
||||
for checkpoint_tuple in reversed(checkpoints):
|
||||
if is_duration_only_checkpoint(checkpoint_tuple):
|
||||
continue
|
||||
message_ids = {_message_id(message) for message in checkpoint_messages(checkpoint_tuple)}
|
||||
if message_id in message_ids:
|
||||
return previous_checkpoint, True
|
||||
if checkpoint_configurable(checkpoint_tuple).get("checkpoint_id"):
|
||||
previous_checkpoint = checkpoint_tuple
|
||||
return None, False
|
||||
@ -23,9 +23,19 @@ from langchain_core.messages import BaseMessage
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.authz import require_permission
|
||||
from app.gateway.checkpoint_lineage import (
|
||||
CheckpointLineageError,
|
||||
CheckpointParentMissingError,
|
||||
checkpoint_configurable,
|
||||
checkpoint_messages,
|
||||
find_checkpoint_before_message,
|
||||
find_checkpoint_before_message_chronologically,
|
||||
is_duration_only_checkpoint,
|
||||
)
|
||||
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 build_checkpoint_state_accessor, build_thread_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion
|
||||
from app.gateway.utils import sanitize_log_param
|
||||
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
|
||||
@ -37,12 +47,12 @@ REGENERATE_HISTORY_SCAN_LIMIT = 200
|
||||
# (one per successful run in steady state) consume roughly half of history.
|
||||
REGENERATE_HISTORY_RAW_SCAN_LIMIT = REGENERATE_HISTORY_SCAN_LIMIT * 2
|
||||
THREAD_MESSAGE_PAGE_SCAN_BATCH = 201
|
||||
_MISSING_REGENERATE_BASE_DETAIL = "Could not find an addressable checkpoint before the target user message"
|
||||
_UNSAFE_REGENERATE_LINEAGE_DETAIL = "Could not safely resolve the checkpoint before the target user message"
|
||||
|
||||
|
||||
def _is_duration_only_checkpoint(checkpoint_tuple: Any) -> bool:
|
||||
metadata = getattr(checkpoint_tuple, "metadata", None)
|
||||
writes = metadata.get("writes") if isinstance(metadata, dict) else None
|
||||
return isinstance(writes, dict) and "runtime_run_duration" in writes
|
||||
return is_duration_only_checkpoint(checkpoint_tuple)
|
||||
|
||||
|
||||
def compute_run_durations(runs) -> dict[str, int]:
|
||||
@ -294,15 +304,11 @@ def _is_middleware_message_row(row: dict[str, Any]) -> bool:
|
||||
|
||||
|
||||
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 []
|
||||
return checkpoint_messages(snapshot)
|
||||
|
||||
|
||||
def _checkpoint_configurable(checkpoint_tuple: Any) -> dict[str, Any]:
|
||||
config = getattr(checkpoint_tuple, "config", None) or {}
|
||||
configurable = config.get("configurable", {}) if isinstance(config, dict) else {}
|
||||
return dict(configurable) if isinstance(configurable, dict) else {}
|
||||
return checkpoint_configurable(checkpoint_tuple)
|
||||
|
||||
|
||||
def _checkpoint_response(checkpoint_tuple: Any) -> dict[str, Any]:
|
||||
@ -356,7 +362,13 @@ def _run_last_ai_matches_message(record: RunRecord, message: Any) -> bool:
|
||||
return last_ai_message == target_text[: len(last_ai_message)]
|
||||
|
||||
|
||||
async def _find_target_run_id(thread_id: str, message_id: str, target_message: Any, request: Request) -> str:
|
||||
async def _find_target_run_id(
|
||||
thread_id: str,
|
||||
message_id: str,
|
||||
target_message: Any,
|
||||
source_human: Any,
|
||||
request: Request,
|
||||
) -> str:
|
||||
event_store = get_run_event_store(request)
|
||||
rows = await event_store.list_messages(thread_id, limit=REGENERATE_HISTORY_SCAN_LIMIT)
|
||||
for row in reversed(rows):
|
||||
@ -366,6 +378,11 @@ async def _find_target_run_id(thread_id: str, message_id: str, target_message: A
|
||||
run_id = row.get("run_id")
|
||||
if isinstance(run_id, str) and run_id:
|
||||
return run_id
|
||||
|
||||
source_run_id = _message_additional_kwargs(source_human).get("run_id")
|
||||
if isinstance(source_run_id, str) and source_run_id:
|
||||
return source_run_id
|
||||
|
||||
run_mgr = get_run_manager(request)
|
||||
user_id = await get_current_user(request)
|
||||
records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=10)
|
||||
@ -385,8 +402,37 @@ async def _find_target_run_id(thread_id: str, message_id: str, target_message: A
|
||||
raise HTTPException(status_code=409, detail="Could not find source run for assistant message")
|
||||
|
||||
|
||||
async def _find_base_checkpoint_before_human(thread_id: str, human_message_id: str, request: Request) -> Any:
|
||||
async def _find_base_checkpoint_before_human(
|
||||
thread_id: str,
|
||||
human_message_id: str,
|
||||
request: Request,
|
||||
*,
|
||||
head_checkpoint: Any | None = None,
|
||||
) -> Any:
|
||||
accessor, base_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id)
|
||||
if head_checkpoint is not None:
|
||||
try:
|
||||
return await find_checkpoint_before_message(
|
||||
accessor,
|
||||
head_checkpoint,
|
||||
human_message_id,
|
||||
max_depth=REGENERATE_HISTORY_RAW_SCAN_LIMIT,
|
||||
)
|
||||
except CheckpointParentMissingError:
|
||||
# Old checkpoints and imported histories may not have parent links.
|
||||
# Preserve the bounded chronological fallback for those records.
|
||||
logger.debug(
|
||||
"Could not resolve parent lineage for regenerate thread %s; falling back to history scan",
|
||||
sanitize_log_param(thread_id),
|
||||
exc_info=True,
|
||||
)
|
||||
except CheckpointLineageError as exc:
|
||||
logger.warning(
|
||||
"Rejected unsafe checkpoint lineage for regenerate thread %s",
|
||||
sanitize_log_param(thread_id),
|
||||
exc_info=True,
|
||||
)
|
||||
raise HTTPException(status_code=409, detail=_UNSAFE_REGENERATE_LINEAGE_DETAIL) from exc
|
||||
try:
|
||||
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)]
|
||||
@ -394,19 +440,14 @@ async def _find_base_checkpoint_before_human(thread_id: str, human_message_id: s
|
||||
logger.exception("Failed to list checkpoints for regenerate thread %s", thread_id)
|
||||
raise HTTPException(status_code=500, detail="Failed to inspect checkpoint history") from exc
|
||||
|
||||
previous_checkpoint = None
|
||||
for checkpoint_tuple in reversed(checkpoints):
|
||||
messages = _checkpoint_messages(checkpoint_tuple)
|
||||
message_ids = {_message_id(message) for message in messages}
|
||||
if human_message_id in message_ids:
|
||||
if previous_checkpoint is None:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Could not find an addressable checkpoint before the target user message",
|
||||
)
|
||||
return previous_checkpoint
|
||||
if _checkpoint_configurable(checkpoint_tuple).get("checkpoint_id"):
|
||||
previous_checkpoint = checkpoint_tuple
|
||||
previous_checkpoint, target_found = find_checkpoint_before_message_chronologically(raw_checkpoints, human_message_id)
|
||||
if target_found:
|
||||
if previous_checkpoint is None:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=_MISSING_REGENERATE_BASE_DETAIL,
|
||||
)
|
||||
return previous_checkpoint
|
||||
|
||||
if len(checkpoints) >= REGENERATE_HISTORY_SCAN_LIMIT:
|
||||
logger.warning(
|
||||
@ -451,8 +492,19 @@ async def _prepare_regenerate_payload(thread_id: str, message_id: str, request:
|
||||
if not previous_human_id:
|
||||
raise HTTPException(status_code=409, detail="The source user message is missing an id")
|
||||
|
||||
base_checkpoint_tuple = await _find_base_checkpoint_before_human(thread_id, previous_human_id, request)
|
||||
target_run_id = await _find_target_run_id(thread_id, message_id, target_message, request)
|
||||
base_checkpoint_tuple = await _find_base_checkpoint_before_human(
|
||||
thread_id,
|
||||
previous_human_id,
|
||||
request,
|
||||
head_checkpoint=latest_checkpoint,
|
||||
)
|
||||
target_run_id = await _find_target_run_id(
|
||||
thread_id,
|
||||
message_id,
|
||||
target_message,
|
||||
previous_human,
|
||||
request,
|
||||
)
|
||||
checkpoint = _checkpoint_response(base_checkpoint_tuple)
|
||||
metadata = {
|
||||
"regenerate_from_message_id": message_id,
|
||||
|
||||
@ -25,6 +25,13 @@ from pydantic import BaseModel, Field, field_validator
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from app.gateway.authz import require_permission
|
||||
from app.gateway.checkpoint_lineage import (
|
||||
CheckpointLineageError,
|
||||
CheckpointParentMissingError,
|
||||
find_checkpoint_before_message,
|
||||
find_checkpoint_before_message_chronologically,
|
||||
is_duration_only_checkpoint,
|
||||
)
|
||||
from app.gateway.deps import get_checkpointer, get_run_event_store, get_run_manager
|
||||
from app.gateway.internal_auth import get_trusted_internal_owner_user_id
|
||||
from app.gateway.services import (
|
||||
@ -90,6 +97,7 @@ _SERVER_RESERVED_METADATA_KEYS: frozenset[str] = frozenset({"owner_id", "user_id
|
||||
_SIDECAR_METADATA_KEY = "deerflow_sidecar"
|
||||
_BRANCH_METADATA_KEY = "deerflow_branch"
|
||||
_BRANCH_HISTORY_SCAN_LIMIT = 200
|
||||
_BRANCH_HISTORY_RAW_SCAN_LIMIT = _BRANCH_HISTORY_SCAN_LIMIT * 2
|
||||
|
||||
|
||||
def _strip_reserved_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]:
|
||||
@ -159,13 +167,26 @@ 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 :])
|
||||
|
||||
|
||||
def _branch_target_human_message(messages: list[Any], target_message_ids: set[str]) -> Any | None:
|
||||
index_by_id = {_message_id(message): index for index, message in enumerate(messages) if _message_id(message)}
|
||||
if not target_message_ids.issubset(index_by_id.keys()):
|
||||
return None
|
||||
target_start_index = min(index_by_id[message_id] for message_id in target_message_ids)
|
||||
return next(
|
||||
(message for message in reversed(messages[:target_start_index]) if _message_type(message) == "human" and _is_branch_visible_message(message)),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
async def _find_branch_checkpoint(
|
||||
accessor: Any,
|
||||
config: dict[str, Any],
|
||||
target_message_ids: set[str],
|
||||
) -> Any:
|
||||
try:
|
||||
for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_SCAN_LIMIT):
|
||||
for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_RAW_SCAN_LIMIT):
|
||||
if is_duration_only_checkpoint(snapshot):
|
||||
continue
|
||||
if _matches_branch_target(_checkpoint_messages(snapshot), target_message_ids):
|
||||
return snapshot
|
||||
except _CHECKPOINT_MODE_ERRORS as exc:
|
||||
@ -184,7 +205,9 @@ async def _branch_targets_latest_turn(
|
||||
) -> bool:
|
||||
"""Return whether the target turn is the final visible turn."""
|
||||
try:
|
||||
for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_SCAN_LIMIT):
|
||||
for snapshot in await accessor.ahistory(config, limit=_BRANCH_HISTORY_RAW_SCAN_LIMIT):
|
||||
if is_duration_only_checkpoint(snapshot):
|
||||
continue
|
||||
messages = _checkpoint_messages(snapshot)
|
||||
if not messages:
|
||||
continue
|
||||
@ -199,6 +222,57 @@ async def _branch_targets_latest_turn(
|
||||
return False
|
||||
|
||||
|
||||
async def _find_branch_replay_base(
|
||||
accessor: Any,
|
||||
config: dict[str, Any],
|
||||
snapshot: Any,
|
||||
target_human_id: str,
|
||||
) -> Any | None:
|
||||
"""Resolve a replay base while preserving unlinked legacy histories."""
|
||||
|
||||
try:
|
||||
return await find_checkpoint_before_message(
|
||||
accessor,
|
||||
snapshot,
|
||||
target_human_id,
|
||||
max_depth=_BRANCH_HISTORY_RAW_SCAN_LIMIT,
|
||||
)
|
||||
except CheckpointParentMissingError:
|
||||
thread_id = config.get("configurable", {}).get("thread_id", "")
|
||||
logger.debug(
|
||||
"Could not resolve parent lineage for branch thread %s; falling back to history scan",
|
||||
sanitize_log_param(thread_id),
|
||||
exc_info=True,
|
||||
)
|
||||
except CheckpointLineageError as exc:
|
||||
thread_id = config.get("configurable", {}).get("thread_id", "")
|
||||
logger.warning(
|
||||
"Rejected unsafe checkpoint lineage for branch thread %s",
|
||||
sanitize_log_param(thread_id),
|
||||
exc_info=True,
|
||||
)
|
||||
raise HTTPException(status_code=409, detail="This turn can no longer be branched from.") from exc
|
||||
|
||||
try:
|
||||
history = await accessor.ahistory(config, limit=_BRANCH_HISTORY_RAW_SCAN_LIMIT)
|
||||
except _CHECKPOINT_MODE_ERRORS as exc:
|
||||
thread_id = config.get("configurable", {}).get("thread_id", "")
|
||||
raise _checkpoint_mode_http_error(exc, thread_id) from exc
|
||||
except Exception as exc:
|
||||
thread_id = config.get("configurable", {}).get("thread_id", "")
|
||||
logger.exception("Failed to scan replay checkpoint history for thread %s", sanitize_log_param(thread_id))
|
||||
raise HTTPException(status_code=500, detail="Failed to inspect checkpoint history") from exc
|
||||
|
||||
replay_base, target_found = find_checkpoint_before_message_chronologically(history, target_human_id)
|
||||
if not target_found:
|
||||
logger.warning(
|
||||
"Could not locate branch user message %s in chronological history for thread %s",
|
||||
sanitize_log_param(target_human_id),
|
||||
sanitize_log_param(config.get("configurable", {}).get("thread_id", "")),
|
||||
)
|
||||
return replay_base
|
||||
|
||||
|
||||
def _ignore_branch_user_data(directory: str, names: list[str]) -> set[str]:
|
||||
ignored: set[str] = set()
|
||||
base = Path(directory)
|
||||
@ -700,6 +774,16 @@ async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Requ
|
||||
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.")
|
||||
target_human = _branch_target_human_message(_checkpoint_messages(snapshot), target_message_ids)
|
||||
target_human_id = _message_id(target_human)
|
||||
if not target_human_id:
|
||||
raise HTTPException(status_code=409, detail="This turn can no longer be branched from.")
|
||||
replay_base_tuple = await _find_branch_replay_base(
|
||||
source_accessor,
|
||||
source_config,
|
||||
snapshot,
|
||||
target_human_id,
|
||||
)
|
||||
|
||||
# Workspace files are not checkpointed, so they only reflect the *current* thread
|
||||
# state. Cloning them onto a branch from an older turn would leak files created
|
||||
@ -740,20 +824,39 @@ async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Requ
|
||||
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(
|
||||
{
|
||||
**branch_metadata,
|
||||
"source": "branch",
|
||||
}
|
||||
)
|
||||
|
||||
def branch_values(source_snapshot: Any) -> dict[str, Any]:
|
||||
values: dict[str, Any] = {}
|
||||
for key, value in dict(source_snapshot.values).items():
|
||||
if key in branch_reducer_fields:
|
||||
values[key] = Overwrite(list(value) if key == "messages" and isinstance(value, list) else value)
|
||||
else:
|
||||
values[key] = value
|
||||
return values
|
||||
|
||||
# Stamp both synthetic checkpoints with the branch-creation time because
|
||||
# serializers fall back to metadata when snapshot.created_at is absent.
|
||||
checkpoint_metadata_updates = {
|
||||
**branch_metadata,
|
||||
"source": "branch",
|
||||
"updated_at": now,
|
||||
"created_at": now,
|
||||
}
|
||||
new_config.setdefault("metadata", {}).update(checkpoint_metadata_updates)
|
||||
try:
|
||||
await branch_accessor.aupdate(new_config, branch_values, as_node="branch")
|
||||
head_config = new_config
|
||||
if replay_base_tuple is not None:
|
||||
head_config = await branch_accessor.aupdate(
|
||||
new_config,
|
||||
branch_values(replay_base_tuple),
|
||||
as_node="branch",
|
||||
)
|
||||
head_config.setdefault("metadata", {}).update(checkpoint_metadata_updates)
|
||||
await branch_accessor.aupdate(
|
||||
head_config,
|
||||
branch_values(snapshot),
|
||||
as_node="branch",
|
||||
)
|
||||
except _CHECKPOINT_MODE_ERRORS as exc:
|
||||
raise _checkpoint_mode_http_error(exc, new_thread_id) from exc
|
||||
except Exception:
|
||||
|
||||
@ -629,9 +629,10 @@ class _RawCheckpointSnapshot:
|
||||
metadata, config ancestry, created_at) comes straight from the tuple.
|
||||
"""
|
||||
|
||||
__slots__ = ("config", "values", "metadata", "parent_config", "created_at", "tasks", "tasks_known", "next")
|
||||
__slots__ = ("checkpoint_exists", "config", "values", "metadata", "parent_config", "created_at", "tasks", "tasks_known", "next")
|
||||
|
||||
def __init__(self, config: dict[str, Any], tup: Any | None) -> None:
|
||||
self.checkpoint_exists = tup is not None
|
||||
self.config = getattr(tup, "config", None) or config
|
||||
checkpoint = getattr(tup, "checkpoint", None) or {}
|
||||
self.values = dict(checkpoint.get("channel_values") or {})
|
||||
|
||||
@ -1,12 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import empty_checkpoint, uuid6
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from deerflow.runtime import RunStatus
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
@ -27,6 +30,38 @@ def _checkpoint(checkpoint_id: str, messages: list[object], *, metadata: dict |
|
||||
)
|
||||
|
||||
|
||||
async def _put_memory_checkpoint(
|
||||
checkpointer: InMemorySaver,
|
||||
thread_id: str,
|
||||
messages: list[object],
|
||||
*,
|
||||
step: int,
|
||||
parent_config: dict | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> dict:
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint["id"] = str(uuid6())
|
||||
checkpoint["channel_values"] = {"messages": messages}
|
||||
checkpoint["channel_versions"] = {"messages": step}
|
||||
checkpoint_metadata = {
|
||||
"step": step,
|
||||
"source": "loop",
|
||||
"writes": {"test": {"messages": messages}},
|
||||
"parents": {},
|
||||
}
|
||||
checkpoint_metadata.update(metadata or {})
|
||||
return await checkpointer.aput(
|
||||
parent_config or {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
||||
checkpoint,
|
||||
checkpoint_metadata,
|
||||
{"messages": step},
|
||||
)
|
||||
|
||||
|
||||
async def _collect_checkpoints(checkpointer: InMemorySaver, config: dict) -> list:
|
||||
return [checkpoint async for checkpoint in checkpointer.alist(config)]
|
||||
|
||||
|
||||
class FakeCheckpointer:
|
||||
def __init__(self, history, *, latest=None, materialized_history=None, materialized_latest=None):
|
||||
self.history = history
|
||||
@ -72,19 +107,26 @@ class FakeAccessor:
|
||||
values=dict(checkpoint.checkpoint.get("channel_values", {})),
|
||||
config=checkpoint.config,
|
||||
metadata=checkpoint.metadata,
|
||||
parent_config=getattr(checkpoint, "parent_config", None),
|
||||
)
|
||||
|
||||
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)
|
||||
async def aget(self, config):
|
||||
materialized_latest = getattr(self.checkpointer, "materialized_latest", None)
|
||||
if materialized_latest is not None and not config.get("configurable", {}).get("checkpoint_id"):
|
||||
return materialized_latest
|
||||
raw = await self.checkpointer.aget_tuple(config)
|
||||
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
|
||||
async def ahistory(self, config, *, limit=None):
|
||||
alist_limits = getattr(self.checkpointer, "alist_limits", None)
|
||||
if alist_limits is not None:
|
||||
alist_limits.append(limit)
|
||||
history = getattr(self.checkpointer, "materialized_history", None)
|
||||
if history is None:
|
||||
history = [self._from_raw(item) for item in self.checkpointer.history]
|
||||
if hasattr(self.checkpointer, "history"):
|
||||
history = [self._from_raw(item) for item in self.checkpointer.history]
|
||||
else:
|
||||
history = [self._from_raw(item) async for item in self.checkpointer.alist(config, limit=limit)]
|
||||
return history[:limit]
|
||||
|
||||
|
||||
@ -354,6 +396,110 @@ 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_payload_does_not_mutate_legacy_single_checkpoint_branch():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
source_thread_id = "source-thread"
|
||||
branch_thread_id = "legacy-branch"
|
||||
source_run_id = "source-run"
|
||||
human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": source_run_id})
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
|
||||
async def _seed() -> str:
|
||||
source_base_config = await _put_memory_checkpoint(checkpointer, source_thread_id, [], step=0)
|
||||
after_human = await _put_memory_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
[human],
|
||||
step=1,
|
||||
parent_config=source_base_config,
|
||||
)
|
||||
source_head_config = await _put_memory_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
[human, ai],
|
||||
step=2,
|
||||
parent_config=after_human,
|
||||
)
|
||||
source_head = await checkpointer.aget_tuple(source_head_config)
|
||||
assert source_head is not None
|
||||
|
||||
legacy_head = copy.deepcopy(source_head.checkpoint)
|
||||
legacy_head_id = str(uuid6())
|
||||
legacy_head["id"] = legacy_head_id
|
||||
legacy_metadata = copy.deepcopy(source_head.metadata)
|
||||
legacy_metadata.update(
|
||||
{
|
||||
"source": "branch",
|
||||
"deerflow_branch": True,
|
||||
"branch_parent_thread_id": source_thread_id,
|
||||
"branch_parent_checkpoint_id": source_head_config["configurable"]["checkpoint_id"],
|
||||
"branch_parent_message_id": "ai-1",
|
||||
}
|
||||
)
|
||||
await checkpointer.aput(
|
||||
{"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}},
|
||||
legacy_head,
|
||||
legacy_metadata,
|
||||
dict(legacy_head["channel_versions"]),
|
||||
)
|
||||
return legacy_head_id
|
||||
|
||||
legacy_head_id = asyncio.run(_seed())
|
||||
request = _request(checkpointer, FakeEventStore([]))
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(_prepare_regenerate_payload(branch_thread_id, "ai-1", request))
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
|
||||
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
|
||||
assert latest is not None
|
||||
assert latest.config["configurable"]["checkpoint_id"] == legacy_head_id
|
||||
branch_history = asyncio.run(_collect_checkpoints(checkpointer, {"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
|
||||
assert [item.config["configurable"]["checkpoint_id"] for item in branch_history] == [legacy_head_id]
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_rejects_legacy_branch_when_source_checkpoint_is_missing():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
branch_thread_id = "legacy-orphan"
|
||||
human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": "source-run"})
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
|
||||
async def _seed() -> None:
|
||||
await _put_memory_checkpoint(
|
||||
checkpointer,
|
||||
branch_thread_id,
|
||||
[human, ai],
|
||||
step=1,
|
||||
metadata={
|
||||
"source": "branch",
|
||||
"deerflow_branch": True,
|
||||
"branch_parent_thread_id": "deleted-source",
|
||||
"branch_parent_checkpoint_id": "missing-checkpoint",
|
||||
"branch_parent_message_id": "ai-1",
|
||||
},
|
||||
)
|
||||
|
||||
asyncio.run(_seed())
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(
|
||||
_prepare_regenerate_payload(
|
||||
branch_thread_id,
|
||||
"ai-1",
|
||||
_request(checkpointer, FakeEventStore([])),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
|
||||
branch_history = asyncio.run(_collect_checkpoints(checkpointer, {"configurable": {"thread_id": branch_thread_id, "checkpoint_ns": ""}}))
|
||||
assert len(branch_history) == 1
|
||||
|
||||
|
||||
def test_prepare_regenerate_uses_materialized_history_when_raw_messages_are_omitted():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
@ -402,6 +548,124 @@ def test_prepare_regenerate_uses_materialized_history_when_raw_messages_are_omit
|
||||
assert checkpointer.alist_limits == [400]
|
||||
|
||||
|
||||
def test_prepare_regenerate_rejects_cyclic_lineage_without_chronological_fallback():
|
||||
from app.gateway.routers import thread_runs
|
||||
|
||||
human = HumanMessage(id="human-1", content="question")
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
|
||||
def linked_snapshot(checkpoint_id: str, messages: list[object], parent_id: str) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
values={"messages": messages},
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
metadata={},
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": parent_id,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
head = linked_snapshot("head", [human, ai], "cycle")
|
||||
cycle = linked_snapshot("cycle", [human], "head")
|
||||
wrong_sibling_base = linked_snapshot("wrong-sibling", [], "root")
|
||||
by_id = {"head": head, "cycle": cycle}
|
||||
|
||||
async def aget(config):
|
||||
return by_id[config["configurable"]["checkpoint_id"]]
|
||||
|
||||
accessor = SimpleNamespace(
|
||||
aget=aget,
|
||||
ahistory=AsyncMock(return_value=[head, wrong_sibling_base]),
|
||||
)
|
||||
builder = AsyncMock(
|
||||
return_value=(
|
||||
accessor,
|
||||
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
|
||||
)
|
||||
)
|
||||
|
||||
with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", builder):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(
|
||||
thread_runs._find_base_checkpoint_before_human(
|
||||
"thread-1",
|
||||
"human-1",
|
||||
_request(FakeCheckpointer([]), FakeEventStore([])),
|
||||
head_checkpoint=head,
|
||||
)
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Could not safely resolve the checkpoint before the target user message"
|
||||
accessor.ahistory.assert_not_awaited()
|
||||
|
||||
|
||||
def test_prepare_regenerate_rejects_dangling_parent_without_chronological_fallback():
|
||||
from app.gateway.routers import thread_runs
|
||||
|
||||
human = HumanMessage(id="human-1", content="question")
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
head = SimpleNamespace(
|
||||
values={"messages": [human, ai]},
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "head",
|
||||
}
|
||||
},
|
||||
metadata={},
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "missing",
|
||||
}
|
||||
},
|
||||
)
|
||||
missing = SimpleNamespace(
|
||||
values={},
|
||||
config=head.parent_config,
|
||||
metadata=None,
|
||||
created_at=None,
|
||||
parent_config=None,
|
||||
)
|
||||
accessor = SimpleNamespace(
|
||||
aget=AsyncMock(return_value=missing),
|
||||
ahistory=AsyncMock(return_value=[head, _snapshot("wrong-sibling", [])]),
|
||||
)
|
||||
builder = AsyncMock(
|
||||
return_value=(
|
||||
accessor,
|
||||
{"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
|
||||
)
|
||||
)
|
||||
|
||||
with patch.object(thread_runs, "build_thread_checkpoint_state_accessor", builder):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(
|
||||
thread_runs._find_base_checkpoint_before_human(
|
||||
"thread-1",
|
||||
"human-1",
|
||||
_request(FakeCheckpointer([]), FakeEventStore([])),
|
||||
head_checkpoint=head,
|
||||
)
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Could not safely resolve the checkpoint before the target user message"
|
||||
accessor.ahistory.assert_not_awaited()
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_rejects_non_latest_assistant():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
@ -459,6 +723,30 @@ def test_prepare_regenerate_payload_falls_back_to_matching_run_when_events_are_m
|
||||
assert response.metadata["regenerate_from_run_id"] == "run-latest"
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_uses_server_stamped_human_run_id_without_parent_events():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
human = HumanMessage(id="human-1", content="question", additional_kwargs={"run_id": "parent-run"})
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
base = _checkpoint("ckpt-base", [])
|
||||
after_human = _checkpoint("ckpt-human", [human])
|
||||
latest = _checkpoint(
|
||||
"ckpt-ai",
|
||||
[human, ai],
|
||||
metadata={
|
||||
"deerflow_branch": True,
|
||||
"branch_parent_thread_id": "parent-thread",
|
||||
"branch_parent_checkpoint_id": "parent-checkpoint",
|
||||
},
|
||||
)
|
||||
checkpointer = FakeCheckpointer([latest, after_human, base])
|
||||
event_store = FakeEventStore([])
|
||||
|
||||
response = asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
||||
|
||||
assert response.target_run_id == "parent-run"
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_rejects_unverified_run_fallback_when_events_are_missing():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
|
||||
@ -14,7 +14,7 @@ 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 app.gateway.routers import thread_runs, threads
|
||||
from deerflow.config.paths import Paths
|
||||
from deerflow.persistence.thread_meta import InvalidMetadataFilterError
|
||||
from deerflow.persistence.thread_meta.memory import THREADS_NS, MemoryThreadMetaStore
|
||||
@ -165,6 +165,7 @@ def _patch_checkpoint_state_builder(monkeypatch):
|
||||
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)
|
||||
monkeypatch.setattr(thread_runs, "build_thread_checkpoint_state_accessor", _read_boundary)
|
||||
|
||||
|
||||
class _FakeStateAccessor:
|
||||
@ -209,6 +210,7 @@ async def _write_checkpoint(
|
||||
*,
|
||||
step: int,
|
||||
metadata: dict | None = None,
|
||||
parent_config: dict | None = None,
|
||||
) -> dict:
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint["id"] = checkpoint_id
|
||||
@ -223,7 +225,7 @@ async def _write_checkpoint(
|
||||
}
|
||||
checkpoint_metadata.update(metadata or {})
|
||||
return await checkpointer.aput(
|
||||
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
||||
parent_config or {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
||||
checkpoint,
|
||||
checkpoint_metadata,
|
||||
{"messages": step},
|
||||
@ -1100,6 +1102,86 @@ def test_ai_message_lacks_duration_only_for_unannotated_ai_messages() -> None:
|
||||
# ── branch threads from completed assistant turns ─────────────────────────────
|
||||
|
||||
|
||||
def test_branch_thread_can_prepare_regenerate_without_branch_run_events() -> None:
|
||||
app, _store, checkpointer = _build_thread_app()
|
||||
app.include_router(thread_runs.router)
|
||||
source_thread_id = "source-regenerate"
|
||||
source_run_id = "source-run"
|
||||
|
||||
async def list_messages(_thread_id: str, *, limit: int, **_kwargs) -> list[dict]:
|
||||
assert limit == thread_runs.REGENERATE_HISTORY_SCAN_LIMIT
|
||||
return []
|
||||
|
||||
async def list_by_thread(_thread_id: str, *, user_id=None, limit: int = 100) -> list:
|
||||
return []
|
||||
|
||||
app.state.run_event_store = SimpleNamespace(list_messages=list_messages)
|
||||
app.state.run_manager = SimpleNamespace(list_by_thread=list_by_thread)
|
||||
|
||||
human = HumanMessage(id="human-1", content="Question", additional_kwargs={"run_id": source_run_id})
|
||||
ai = AIMessage(id="ai-1", content="Answer")
|
||||
|
||||
with TestClient(app) as client:
|
||||
created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}})
|
||||
assert created.status_code == 200, created.text
|
||||
|
||||
initial = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}}))
|
||||
assert initial is not None
|
||||
after_human = asyncio.run(
|
||||
_write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human],
|
||||
step=1,
|
||||
parent_config=initial.config,
|
||||
)
|
||||
)
|
||||
asyncio.run(
|
||||
_write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human, ai],
|
||||
step=2,
|
||||
parent_config=after_human,
|
||||
)
|
||||
)
|
||||
|
||||
branch_response = client.post(
|
||||
f"/api/threads/{source_thread_id}/branches",
|
||||
json={"message_id": "ai-1", "message_ids": ["ai-1"]},
|
||||
)
|
||||
assert branch_response.status_code == 200, branch_response.text
|
||||
branch_thread_id = branch_response.json()["thread_id"]
|
||||
|
||||
prepare_response = client.post(
|
||||
f"/api/threads/{branch_thread_id}/runs/regenerate/prepare",
|
||||
json={"message_id": "ai-1"},
|
||||
)
|
||||
|
||||
assert prepare_response.status_code == 200, prepare_response.text
|
||||
prepared = prepare_response.json()
|
||||
assert prepared["target_run_id"] == source_run_id
|
||||
branch_base_id = prepared["checkpoint"]["checkpoint_id"]
|
||||
assert prepared["input"]["messages"][0]["id"] == "human-1"
|
||||
assert prepared["input"]["messages"][0]["content"] == [{"type": "text", "text": "Question"}]
|
||||
|
||||
branch_base = asyncio.run(
|
||||
checkpointer.aget_tuple(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": branch_thread_id,
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": branch_base_id,
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
assert branch_base is not None
|
||||
assert branch_base.checkpoint.get("channel_values", {}).get("messages", []) == []
|
||||
|
||||
|
||||
def test_branch_thread_from_older_assistant_turn_creates_truncated_thread() -> None:
|
||||
app, store, checkpointer = _build_thread_app()
|
||||
source_thread_id = "source-thread"
|
||||
@ -1111,16 +1193,49 @@ def test_branch_thread_from_older_assistant_turn_creates_truncated_thread() -> N
|
||||
human_3 = HumanMessage(id="human-3", content="Third question")
|
||||
ai_3 = AIMessage(id="ai-3", content="Third answer")
|
||||
|
||||
async def _seed() -> None:
|
||||
await _write_checkpoint(checkpointer, source_thread_id, "0001", [human_1, ai_1], step=1)
|
||||
await _write_checkpoint(checkpointer, source_thread_id, "0002", [human_1, ai_1, human_2, ai_2], step=2)
|
||||
await _write_checkpoint(checkpointer, source_thread_id, "0003", [human_1, ai_1, human_2, ai_2, human_3, ai_3], step=3)
|
||||
|
||||
asyncio.run(_seed())
|
||||
async def _seed(parent_config: dict) -> dict:
|
||||
after_human_1 = await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human_1], step=1, parent_config=parent_config)
|
||||
after_ai_1 = await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human_1, ai_1], step=2, parent_config=after_human_1)
|
||||
after_human_2 = await _write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human_1, ai_1, human_2],
|
||||
step=3,
|
||||
parent_config=after_ai_1,
|
||||
)
|
||||
after_ai_2 = await _write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human_1, ai_1, human_2, ai_2],
|
||||
step=4,
|
||||
parent_config=after_human_2,
|
||||
)
|
||||
after_human_3 = await _write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human_1, ai_1, human_2, ai_2, human_3],
|
||||
step=5,
|
||||
parent_config=after_ai_2,
|
||||
)
|
||||
await _write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human_1, ai_1, human_2, ai_2, human_3, ai_3],
|
||||
step=6,
|
||||
parent_config=after_human_3,
|
||||
)
|
||||
return after_ai_2
|
||||
|
||||
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
|
||||
initial = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}}))
|
||||
assert initial is not None
|
||||
target_checkpoint_config = asyncio.run(_seed(initial.config))
|
||||
asyncio.run(
|
||||
store.aput(
|
||||
THREADS_NS,
|
||||
@ -1149,7 +1264,7 @@ def test_branch_thread_from_older_assistant_turn_creates_truncated_thread() -> N
|
||||
search_response = client.post("/api/threads/search", json={"limit": 10})
|
||||
|
||||
assert body["parent_thread_id"] == source_thread_id
|
||||
assert body["parent_checkpoint_id"] == "0002"
|
||||
assert body["parent_checkpoint_id"] == target_checkpoint_config["configurable"]["checkpoint_id"]
|
||||
assert body["branched_from_message_id"] == "ai-2"
|
||||
assert body["workspace_clone_mode"] == "skipped_historical_turn"
|
||||
|
||||
@ -1172,7 +1287,7 @@ def test_branch_thread_uses_materialized_history_and_overwrites_fresh_seed(monke
|
||||
AIMessage(id="a2", content="Second answer"),
|
||||
]
|
||||
|
||||
def snapshot(checkpoint_id: str, materialized_messages: list[object]) -> SimpleNamespace:
|
||||
def snapshot(checkpoint_id: str, materialized_messages: list[object], *, parent_id: str | None = None) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
values={"messages": materialized_messages, "title": "Materialized title"},
|
||||
config={
|
||||
@ -1183,31 +1298,48 @@ def test_branch_thread_uses_materialized_history_and_overwrites_fresh_seed(monke
|
||||
}
|
||||
},
|
||||
metadata={"step": int(checkpoint_id[-1])},
|
||||
parent_config=(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": source_thread_id,
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": parent_id,
|
||||
}
|
||||
}
|
||||
if parent_id is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
source_accessor = SimpleNamespace()
|
||||
source_history = [
|
||||
snapshot("ckpt-2", messages),
|
||||
snapshot("ckpt-1", messages[:2]),
|
||||
snapshot("ckpt-2", messages, parent_id="ckpt-1"),
|
||||
snapshot("ckpt-1", messages[:2], parent_id="ckpt-0"),
|
||||
snapshot("ckpt-0", []),
|
||||
]
|
||||
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
|
||||
assert limit == threads._BRANCH_HISTORY_RAW_SCAN_LIMIT
|
||||
return source_history
|
||||
|
||||
async def source_aget(config):
|
||||
checkpoint_id = config["configurable"]["checkpoint_id"]
|
||||
return next(item for item in source_history if item.config["configurable"]["checkpoint_id"] == checkpoint_id)
|
||||
|
||||
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",
|
||||
"checkpoint_id": f"branch-{len(branch_updates)}",
|
||||
}
|
||||
}
|
||||
|
||||
source_accessor.ahistory = source_ahistory
|
||||
source_accessor.aget = source_aget
|
||||
branch_accessor = SimpleNamespace(aupdate=branch_aupdate)
|
||||
|
||||
def build_accessor(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
|
||||
@ -1244,13 +1376,179 @@ def test_branch_thread_uses_materialized_history_and_overwrites_fresh_seed(monke
|
||||
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"
|
||||
assert len(branch_updates) == 2
|
||||
replay_config, replay_values, replay_node = branch_updates[0]
|
||||
assert isinstance(replay_values["messages"], Overwrite)
|
||||
assert replay_values["messages"].value == []
|
||||
assert replay_config["configurable"]["thread_id"] == body["thread_id"]
|
||||
assert replay_config["metadata"]["source"] == "branch"
|
||||
assert replay_node == "branch"
|
||||
|
||||
head_config, head_values, head_node = branch_updates[1]
|
||||
assert isinstance(head_values["messages"], Overwrite)
|
||||
assert [message.id for message in head_values["messages"].value] == ["h1", "a1"]
|
||||
assert head_config["configurable"]["checkpoint_id"] == "branch-1"
|
||||
assert head_config["metadata"]["source"] == "branch"
|
||||
assert head_node == "branch"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("include_replay_base", "expected_message_ids"),
|
||||
[
|
||||
(True, [[], ["h1", "a1"]]),
|
||||
(False, [["h1", "a1"]]),
|
||||
],
|
||||
ids=["chronological-replay-base", "legacy-single-checkpoint"],
|
||||
)
|
||||
def test_branch_thread_preserves_unlinked_legacy_histories(
|
||||
monkeypatch,
|
||||
include_replay_base: bool,
|
||||
expected_message_ids: list[list[str]],
|
||||
) -> None:
|
||||
app, _store, _checkpointer = _build_thread_app()
|
||||
source_thread_id = "source-unlinked"
|
||||
messages = [
|
||||
HumanMessage(id="h1", content="Question"),
|
||||
AIMessage(id="a1", content="Answer"),
|
||||
]
|
||||
|
||||
def snapshot(checkpoint_id: str, snapshot_messages: list[object], *, duration_only: bool = False) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
values={"messages": snapshot_messages},
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": source_thread_id,
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": checkpoint_id,
|
||||
}
|
||||
},
|
||||
metadata={"writes": {"runtime_run_duration": 1}} if duration_only else {},
|
||||
parent_config=None,
|
||||
)
|
||||
|
||||
source_history = [snapshot("ckpt-1", messages)]
|
||||
if include_replay_base:
|
||||
source_history.extend(
|
||||
[
|
||||
snapshot("ckpt-duration", [], duration_only=True),
|
||||
snapshot("ckpt-0", []),
|
||||
]
|
||||
)
|
||||
|
||||
history_limits: list[int | None] = []
|
||||
|
||||
async def source_ahistory(config, *, limit=None):
|
||||
assert config["configurable"]["thread_id"] == source_thread_id
|
||||
history_limits.append(limit)
|
||||
return source_history
|
||||
|
||||
async def unexpected_lineage_read(_config):
|
||||
raise AssertionError("unlinked checkpoints must use chronological history")
|
||||
|
||||
branch_updates: list[dict] = []
|
||||
|
||||
async def branch_aupdate(config, values, *, as_node=None):
|
||||
assert as_node == "branch"
|
||||
branch_updates.append(values)
|
||||
return {
|
||||
"configurable": {
|
||||
"thread_id": config["configurable"]["thread_id"],
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": f"branch-{len(branch_updates)}",
|
||||
}
|
||||
}
|
||||
|
||||
source_accessor = SimpleNamespace(ahistory=source_ahistory, aget=unexpected_lineage_read)
|
||||
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
|
||||
assert response.json()["parent_checkpoint_id"] == "ckpt-1"
|
||||
assert [[message.id for message in update["messages"].value] for update in branch_updates] == expected_message_ids
|
||||
assert history_limits == [
|
||||
threads._BRANCH_HISTORY_RAW_SCAN_LIMIT,
|
||||
threads._BRANCH_HISTORY_RAW_SCAN_LIMIT,
|
||||
threads._BRANCH_HISTORY_RAW_SCAN_LIMIT,
|
||||
]
|
||||
|
||||
|
||||
def test_branch_history_scans_budget_for_duration_only_checkpoints() -> None:
|
||||
target = SimpleNamespace(
|
||||
values={
|
||||
"messages": [
|
||||
HumanMessage(id="h1", content="Question"),
|
||||
AIMessage(id="a1", content="Answer"),
|
||||
]
|
||||
},
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "source-duration-budget",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "target",
|
||||
}
|
||||
},
|
||||
metadata={},
|
||||
)
|
||||
duration_only = [
|
||||
SimpleNamespace(
|
||||
values={"messages": []},
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "source-duration-budget",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": f"duration-{index}",
|
||||
}
|
||||
},
|
||||
metadata={"writes": {"runtime_run_duration": index}},
|
||||
)
|
||||
for index in range(threads._BRANCH_HISTORY_SCAN_LIMIT)
|
||||
]
|
||||
history = [*duration_only, target]
|
||||
limits: list[int | None] = []
|
||||
|
||||
async def ahistory(_config, *, limit=None):
|
||||
limits.append(limit)
|
||||
return history[:limit]
|
||||
|
||||
accessor = SimpleNamespace(ahistory=ahistory)
|
||||
config = {"configurable": {"thread_id": "source-duration-budget", "checkpoint_ns": ""}}
|
||||
|
||||
found = asyncio.run(threads._find_branch_checkpoint(accessor, config, {"a1"}))
|
||||
targets_latest = asyncio.run(threads._branch_targets_latest_turn(accessor, config, {"a1"}))
|
||||
|
||||
assert found is target
|
||||
assert targets_latest is True
|
||||
assert limits == [threads._BRANCH_HISTORY_RAW_SCAN_LIMIT, threads._BRANCH_HISTORY_RAW_SCAN_LIMIT]
|
||||
|
||||
|
||||
def test_branch_thread_real_mutation_graph_finishes_without_scheduling(monkeypatch) -> None:
|
||||
@ -1264,7 +1562,7 @@ def test_branch_thread_real_mutation_graph_finishes_without_scheduling(monkeypat
|
||||
AIMessage(id="a2", content="Second answer"),
|
||||
]
|
||||
|
||||
def snapshot(checkpoint_id: str, materialized_messages: list[object]) -> SimpleNamespace:
|
||||
def snapshot(checkpoint_id: str, materialized_messages: list[object], *, parent_id: str | None = None) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
values={"messages": materialized_messages},
|
||||
config={
|
||||
@ -1275,15 +1573,27 @@ def test_branch_thread_real_mutation_graph_finishes_without_scheduling(monkeypat
|
||||
}
|
||||
},
|
||||
metadata={},
|
||||
parent_config=(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": source_thread_id,
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": parent_id,
|
||||
}
|
||||
}
|
||||
if parent_id is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
source_history = [
|
||||
snapshot("ckpt-2", messages, parent_id="ckpt-1"),
|
||||
snapshot("ckpt-1", messages[:2], parent_id="ckpt-0"),
|
||||
snapshot("ckpt-0", []),
|
||||
]
|
||||
source_accessor = SimpleNamespace(
|
||||
ahistory=AsyncMock(
|
||||
return_value=[
|
||||
snapshot("ckpt-2", messages),
|
||||
snapshot("ckpt-1", messages[:2]),
|
||||
]
|
||||
)
|
||||
ahistory=AsyncMock(return_value=source_history),
|
||||
aget=AsyncMock(side_effect=lambda config: next(item for item in source_history if item.config["configurable"]["checkpoint_id"] == config["configurable"]["checkpoint_id"])),
|
||||
)
|
||||
|
||||
def source_builder(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
|
||||
@ -1374,6 +1684,7 @@ def _wire_extension_agent(monkeypatch, app, checkpointer, mode):
|
||||
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)
|
||||
monkeypatch.setattr(thread_runs, "build_thread_checkpoint_state_accessor", gateway_services.build_thread_checkpoint_state_accessor)
|
||||
return custom_factory
|
||||
|
||||
|
||||
@ -1389,6 +1700,24 @@ async def _seed_extension_source(checkpointer, custom_factory, mode, source_thre
|
||||
{"messages": [AIMessage(id="a1", content="answer")], "ext_list": ["payload"]},
|
||||
as_node="model",
|
||||
)
|
||||
await accessor.aupdate(
|
||||
{"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}},
|
||||
{
|
||||
"messages": [
|
||||
HumanMessage(
|
||||
id="h2",
|
||||
content="follow-up",
|
||||
additional_kwargs={"run_id": "source-run"},
|
||||
)
|
||||
]
|
||||
},
|
||||
as_node="model",
|
||||
)
|
||||
await accessor.aupdate(
|
||||
{"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}},
|
||||
{"messages": [AIMessage(id="a2", content="follow-up answer")]},
|
||||
as_node="model",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["full", "delta"])
|
||||
@ -1397,11 +1726,23 @@ def test_state_endpoints_preserve_extension_reducer_channels(monkeypatch, mode)
|
||||
|
||||
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.
|
||||
byte-for-byte by copying reducer channels with Overwrite semantics. The
|
||||
copied pre-user checkpoint must also remain materializable for regenerate.
|
||||
"""
|
||||
app, _store, checkpointer = _build_thread_app()
|
||||
app.include_router(thread_runs.router)
|
||||
custom_factory = _wire_extension_agent(monkeypatch, app, checkpointer, mode)
|
||||
|
||||
async def list_messages(_thread_id: str, *, limit: int, **_kwargs) -> list[dict]:
|
||||
assert limit == thread_runs.REGENERATE_HISTORY_SCAN_LIMIT
|
||||
return []
|
||||
|
||||
async def list_by_thread(_thread_id: str, *, user_id=None, limit: int = 100) -> list:
|
||||
return []
|
||||
|
||||
app.state.run_event_store = SimpleNamespace(list_messages=list_messages)
|
||||
app.state.run_manager = SimpleNamespace(list_by_thread=list_by_thread)
|
||||
|
||||
recorded_updates: list[dict] = []
|
||||
real_mutation_builder = gateway_services.build_checkpoint_state_mutation_accessor
|
||||
|
||||
@ -1444,11 +1785,17 @@ def test_state_endpoints_preserve_extension_reducer_channels(monkeypatch, mode)
|
||||
|
||||
branch_response = client.post(
|
||||
f"/api/threads/{source_thread_id}/branches",
|
||||
json={"message_id": "a1", "message_ids": ["a1"]},
|
||||
json={"message_id": "a2", "message_ids": ["a2"]},
|
||||
)
|
||||
assert branch_response.status_code == 200, branch_response.text
|
||||
branch_thread_id = branch_response.json()["thread_id"]
|
||||
|
||||
prepare_response = client.post(
|
||||
f"/api/threads/{branch_thread_id}/runs/regenerate/prepare",
|
||||
json={"message_id": "a2"},
|
||||
)
|
||||
assert prepare_response.status_code == 200, prepare_response.text
|
||||
|
||||
# The branch write must copy every reducer channel with replace semantics.
|
||||
branch_update = recorded_updates[-1]
|
||||
assert isinstance(branch_update["ext_list"], Overwrite)
|
||||
@ -1462,7 +1809,24 @@ def test_state_endpoints_preserve_extension_reducer_channels(monkeypatch, mode)
|
||||
|
||||
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"]
|
||||
assert [message.id for message in branch_values["messages"]] == ["h1", "a1", "h2", "a2"]
|
||||
|
||||
prepared = prepare_response.json()
|
||||
assert prepared["target_run_id"] == "source-run"
|
||||
assert prepared["input"]["messages"][0]["id"] == "h2"
|
||||
base_accessor = CheckpointStateAccessor.bind(custom_factory(), checkpointer, mode=mode)
|
||||
base_values = asyncio.run(
|
||||
base_accessor.aget(
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": branch_thread_id,
|
||||
"checkpoint_ns": prepared["checkpoint"]["checkpoint_ns"],
|
||||
"checkpoint_id": prepared["checkpoint"]["checkpoint_id"],
|
||||
}
|
||||
}
|
||||
)
|
||||
).values
|
||||
assert [message.id for message in base_values["messages"]] == ["h1", "a1"]
|
||||
|
||||
|
||||
async def _seed_branch_history_source(checkpointer, custom_factory, mode, source_thread_id):
|
||||
@ -1625,14 +1989,16 @@ def test_branch_thread_rejects_non_assistant_targets() -> None:
|
||||
human = HumanMessage(id="human-1", content="Question")
|
||||
ai = AIMessage(id="ai-1", content="Answer")
|
||||
|
||||
async def _seed() -> None:
|
||||
await _write_checkpoint(checkpointer, source_thread_id, "0001", [human, ai], step=1)
|
||||
|
||||
asyncio.run(_seed())
|
||||
async def _seed(parent_config: dict) -> None:
|
||||
after_human = await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human], step=1, parent_config=parent_config)
|
||||
await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human, ai], step=2, parent_config=after_human)
|
||||
|
||||
with TestClient(app) as client:
|
||||
created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}})
|
||||
assert created.status_code == 200, created.text
|
||||
initial = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}}))
|
||||
assert initial is not None
|
||||
asyncio.run(_seed(initial.config))
|
||||
|
||||
response = client.post(
|
||||
f"/api/threads/{source_thread_id}/branches",
|
||||
@ -1660,10 +2026,9 @@ def test_branch_thread_best_effort_copies_current_workspace(tmp_path) -> None:
|
||||
human = HumanMessage(id="human-file", content="Make a file")
|
||||
ai = AIMessage(id="ai-file", content="Done")
|
||||
|
||||
async def _seed() -> None:
|
||||
await _write_checkpoint(checkpointer, source_thread_id, "0001", [human, ai], step=1)
|
||||
|
||||
asyncio.run(_seed())
|
||||
async def _seed(parent_config: dict) -> None:
|
||||
after_human = await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human], step=1, parent_config=parent_config)
|
||||
await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human, ai], step=2, parent_config=after_human)
|
||||
|
||||
with (
|
||||
patch("app.gateway.routers.threads.get_paths", return_value=paths),
|
||||
@ -1672,6 +2037,9 @@ def test_branch_thread_best_effort_copies_current_workspace(tmp_path) -> None:
|
||||
):
|
||||
created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}})
|
||||
assert created.status_code == 200, created.text
|
||||
initial = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}}))
|
||||
assert initial is not None
|
||||
asyncio.run(_seed(initial.config))
|
||||
|
||||
response = client.post(
|
||||
f"/api/threads/{source_thread_id}/branches",
|
||||
@ -1711,11 +2079,26 @@ def test_branch_thread_from_historical_turn_skips_workspace_clone(tmp_path) -> N
|
||||
human_2 = HumanMessage(id="human-2", content="Second question")
|
||||
ai_2 = AIMessage(id="ai-2", content="Second answer")
|
||||
|
||||
async def _seed() -> None:
|
||||
await _write_checkpoint(checkpointer, source_thread_id, "0001", [human_1, ai_1], step=1)
|
||||
await _write_checkpoint(checkpointer, source_thread_id, "0002", [human_1, ai_1, human_2, ai_2], step=2)
|
||||
|
||||
asyncio.run(_seed())
|
||||
async def _seed(parent_config: dict) -> dict:
|
||||
after_human_1 = await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human_1], step=1, parent_config=parent_config)
|
||||
after_ai_1 = await _write_checkpoint(checkpointer, source_thread_id, str(uuid6()), [human_1, ai_1], step=2, parent_config=after_human_1)
|
||||
after_human_2 = await _write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human_1, ai_1, human_2],
|
||||
step=3,
|
||||
parent_config=after_ai_1,
|
||||
)
|
||||
await _write_checkpoint(
|
||||
checkpointer,
|
||||
source_thread_id,
|
||||
str(uuid6()),
|
||||
[human_1, ai_1, human_2, ai_2],
|
||||
step=4,
|
||||
parent_config=after_human_2,
|
||||
)
|
||||
return after_ai_1
|
||||
|
||||
with (
|
||||
patch("app.gateway.routers.threads.get_paths", return_value=paths),
|
||||
@ -1724,6 +2107,9 @@ def test_branch_thread_from_historical_turn_skips_workspace_clone(tmp_path) -> N
|
||||
):
|
||||
created = client.post("/api/threads", json={"thread_id": source_thread_id, "metadata": {}})
|
||||
assert created.status_code == 200, created.text
|
||||
initial = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": source_thread_id, "checkpoint_ns": ""}}))
|
||||
assert initial is not None
|
||||
target_checkpoint_config = asyncio.run(_seed(initial.config))
|
||||
|
||||
response = client.post(
|
||||
f"/api/threads/{source_thread_id}/branches",
|
||||
@ -1732,7 +2118,7 @@ def test_branch_thread_from_historical_turn_skips_workspace_clone(tmp_path) -> N
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
body = response.json()
|
||||
assert body["parent_checkpoint_id"] == "0001"
|
||||
assert body["parent_checkpoint_id"] == target_checkpoint_config["configurable"]["checkpoint_id"]
|
||||
assert body["workspace_clone_mode"] == "skipped_historical_turn"
|
||||
|
||||
target_user_data = paths.sandbox_user_data_dir(body["thread_id"], user_id=user_id)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user