From fcbf0609b05c8adb9517259e8b715d98a3323212 Mon Sep 17 00:00:00 2001 From: Ryker_Feng <90562015+18062706139fcz@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:46:51 +0800 Subject: [PATCH] feat(chat): edit and rerun latest user turn (#4377) Co-authored-by: Willem Jiang --- README.md | 2 + backend/AGENTS.md | 7 +- backend/app/gateway/routers/thread_runs.py | 387 +++++++++++++++--- .../harness/deerflow/persistence/run/sql.py | 22 + .../harness/deerflow/runtime/runs/manager.py | 62 +++ .../deerflow/runtime/runs/store/base.py | 16 + .../deerflow/runtime/runs/store/memory.py | 16 + .../harness/deerflow/runtime/runs/worker.py | 76 +++- backend/tests/test_history_batch_queries.py | 119 ++++++ backend/tests/test_run_worker_rollback.py | 144 +++++++ .../tests/test_thread_messages_feedback.py | 3 + backend/tests/test_thread_messages_page.py | 84 +++- .../tests/test_thread_regenerate_prepare.py | 350 +++++++++++++++- .../tests/test_thread_run_query_validation.py | 10 +- ...st_tool_args_schema_no_pydantic_warning.py | 10 + frontend/AGENTS.md | 3 + .../[agent_name]/chats/[thread_id]/page.tsx | 16 + .../app/workspace/chats/[thread_id]/page.tsx | 17 + .../workspace/messages/message-list-item.tsx | 136 +++++- .../workspace/messages/message-list.tsx | 59 ++- frontend/src/core/i18n/locales/en-US.ts | 4 + frontend/src/core/i18n/locales/types.ts | 3 + frontend/src/core/i18n/locales/zh-CN.ts | 4 + frontend/src/core/messages/utils.ts | 83 ++++ frontend/src/core/threads/hooks.ts | 226 ++++++++-- frontend/tests/e2e/agent-chat.spec.ts | 157 +++++++ .../tests/unit/core/messages/utils.test.ts | 83 ++++ .../unit/core/threads/message-merge.test.ts | 32 ++ 28 files changed, 2024 insertions(+), 107 deletions(-) diff --git a/README.md b/README.md index 3e0c3fb9b..8f7906862 100644 --- a/README.md +++ b/README.md @@ -788,6 +788,8 @@ In the Web UI, completed assistant turns can be branched into a new main convers 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. +In the Web UI, the latest completed user turn can also be edited and rerun from the message toolbar. DeerFlow restores the conversation checkpoint before that user message, submits the edited text as a new user message, and hides the superseded turn once the replay is in progress or succeeds. This is a conversation-state replay only: files, memory updates, and external tool side effects are not undone. + Web UI chat links percent-encode custom thread identifiers before placing them in route segments, so reserved URL characters such as `#` and `?` do not change which conversation is opened. ``` diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 9c89850bf..1b996f014 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -441,7 +441,7 @@ Localhost persistence deliberately reads the direct request `Host` and ignores ` | **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 (`...`, 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 `` substrings in the rewrite (`strip_think_blocks(truncate_unclosed=False)`) | -| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block; `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest assistant answer, carrying the latest non-empty thread title in graph input so resuming an older checkpoint cannot roll back a later manual rename (#4457); `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated per-run messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET /../messages` - legacy thread message array; `GET /../messages/page` - backward thread-global `seq` history page with middleware/subagent-AI/successful-regenerate filtering and page-run-scoped feedback enrichment; subagent AI callbacks remain available through run events while parent `task` ToolMessages stay visible for card restoration; `GET /../token-usage` - aggregate tokens | +| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block; `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest assistant answer; carrying the latest non-empty thread title in graph input so resuming an older checkpoint cannot roll back a later manual rename (#4457); `POST /edit-regenerate/prepare` - prepare a checkpoint replay from the latest editable human turn with a replacement user message and edit replay metadata; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated per-run messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET /../messages` - legacy thread message array; `GET /../messages/page` - backward thread-global `seq` history page with middleware/subagent-AI/successful-regenerate/edit-replay filtering and page-run-scoped feedback enrichment; subagent AI callbacks remain available through run events while parent `task` ToolMessages stay visible for card restoration; `GET /../token-usage` - aggregate tokens | | **Feedback** (`/api/threads/{id}/runs/{rid}/feedback`) | `PUT /` - upsert feedback; `DELETE /` - delete user feedback; `POST /` - create feedback; `GET /` - list feedback; `GET /stats` - aggregate stats; `DELETE /{fid}` - delete specific | | **Runs** (`/api/runs`) | `POST /stream` - stateless run + SSE; `POST /wait` - stateless run + block; `GET /{rid}/messages` - paginated messages by run_id `{data, has_more}` (cursor: `after_seq`/`before_seq`); `GET /{rid}/feedback` - list feedback by run_id | | **GitHub Webhooks** (`/api/webhooks/github`) | `POST /` - receive GitHub App / repo webhook deliveries. Verifies `X-Hub-Signature-256` against `GITHUB_WEBHOOK_SECRET`; exempt from auth + CSRF because authenticity is enforced by HMAC. The route is fail-closed: mounted only when `GITHUB_WEBHOOK_SECRET` is set, or when explicit dev opt-in `DEER_FLOW_ALLOW_UNVERIFIED_GITHUB_WEBHOOKS=1` is set. Recognized events include `ping`, `issues`, `issue_comment`, `pull_request`, `pull_request_review`, and `pull_request_review_comment`; unknown events return 200 with `handled=false`. Fan-out runtime failures return 503, keeping the delivery recorded as failed for manual/API/scripted redelivery (GitHub does not automatically retry any failed delivery, 5xx included); permanent/non-retryable conditions such as `channels.github.enabled: false`, unknown events, malformed payloads, or unavailable channel service return 200 with a skipped/handled response. | @@ -504,7 +504,8 @@ JSONL event stores when `GATEWAY_WORKERS > 1`. **RunManager / RunStore contract**: - LangGraph-compatible run requests validate their supported subset before creating a run. `runtime/stream_modes.py` is the shared backend contract for public stream modes and the worker's `graph.astream` mapping; the public `messages-tuple` mode maps to LangGraph's internal `messages` mode, while public `messages`, `events`, and other unsupported modes are rejected instead of being dropped or replaced with `values`. `app/gateway/run_models.py::RunCreateRequest` is shared by HTTP and internal scheduled launch paths, retains only truthful compatibility defaults for unimplemented options (`if_not_exists="create"` plus `None` placeholders), returns 422 for unsupported values including `on_completion="complete"`, `on_completion="continue"`, and `multitask_strategy="enqueue"`, and forbids undeclared SDK options so fields such as `checkpoint_during` and `durability` cannot be silently discarded. A placeholder must still accept the stock SDK's own default: `langgraph_sdk` drops only `None` from its run payload, so `stream_resumable=False` reaches every request and means "non-resumable", which is what DeerFlow serves — rejecting it 422'd every IM channel run (#4466). `tests/test_run_request_validation.py::test_gateway_accepts_langgraph_sdk_default_payload` pins the real SDK payload against this boundary; channel tests mock the SDK client and cannot catch this class of drift. - `RunManager.get()` is async; direct callers must `await` it. -- The history batch helpers `list_successful_regenerate_sources()` and `get_many_by_thread()` default to `user_id=AUTO`: they resolve the request user and fail closed when no user context exists. Migration/admin callers that intentionally need an unscoped read must pass `user_id=None` explicitly. +- The history batch helpers `list_successful_regenerate_sources()`, `list_edit_regenerate_runs()`, and `get_many_by_thread()` default to `user_id=AUTO`: they resolve the request user and fail closed when no user context exists. Migration/admin callers that intentionally need an unscoped read must pass `user_id=None` explicitly. +- Edit-and-rerun visibility is derived from edit replay runs (`metadata.replay_kind="edit"` plus `regenerate_from_run_id`) by `RunManager.list_edit_replay_visibility()`: the newest attempt for each source run is authoritative. Pending/running/success attempts hide the original source run; failed, timed-out, or interrupted attempts hide only the failed attempt so the original conversation reappears. - When a persistent `RunStore` is configured, `get()` and `list_by_thread()` hydrate historical runs from the store. In-memory records win for the same `run_id` so task, abort, and stream-control state stays attached to active local runs. - Thread metadata status switches to `running` only after `RunManager.try_start()` succeeds. Pending-cancelled runs therefore skip the old `running` projection, while clients may observe the prior thread status during the short worker-startup window. - `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (owner's lease is still alive — caller should return 409 + `Retry-After`), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions. @@ -933,7 +934,7 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c **Wholesale state 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 reducer values requires `Overwrite` rather than an ordinary update. Full-mode rollback and context compaction replace `messages`; delta resume and delta rollback replace every materialized channel and reset current-head-only channels to their schema default (or `None`). These writes go through `build_state_mutation_graph(as_node, mode, state_schema)`, and `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 a full-mode fork write inherit the parent's channel blobs, so middleware channels survive rollback/compaction (locked by `test_rollback_preserves_middleware_contributed_channels` and `test_compact_thread_context_preserves_middleware_contributed_channels`). The compiled mutation graph has 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 the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Both modes reattach only the captured pre-run pending writes to the restored checkpoint. +**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Both modes reattach only the captured pre-run pending writes to the restored checkpoint. Edit replay runs (`metadata.replay_kind="edit"`) also restore the pre-run checkpoint on failed, timed-out, or interrupted completion and publish the restored `values` snapshot to the stream before `end`, so clients do not remain on a transient edited branch when the replay did not produce a successful replacement. **Where things live**: - `runtime/checkpoint_mode.py` — mode freeze, marker injection, delta detection, compatibility gate, both error types diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index 77905389e..5d1069697 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -13,6 +13,7 @@ from __future__ import annotations import asyncio import logging +import uuid from copy import deepcopy from datetime import UTC, datetime from typing import Any, Literal @@ -50,6 +51,7 @@ 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" +THREAD_MESSAGE_LEGACY_SCAN_BATCH = 201 def _is_duration_only_checkpoint(checkpoint_tuple: Any) -> bool: @@ -90,6 +92,16 @@ class RegeneratePrepareResponse(BaseModel): target_run_id: str +class EditRegeneratePrepareRequest(BaseModel): + human_message_id: str = Field(..., min_length=1, description="Source human message id to edit and rerun") + replacement_text: str = Field(..., min_length=1, description="Replacement user-visible text") + + +class EditRegeneratePrepareResponse(RegeneratePrepareResponse): + replacement_human_message_id: str + source_message_ids: list[str] + + class ThreadMessagesPageResponse(BaseModel): data: list[dict[str, Any]] has_more: bool @@ -263,6 +275,15 @@ def _message_additional_kwargs(message: Any) -> dict[str, Any]: return dict(value or {}) if isinstance(value, dict) else {} +def _message_tool_calls(message: Any) -> list[Any]: + value = getattr(message, "tool_calls", None) + if value is None and isinstance(message, dict): + value = message.get("tool_calls") + if value is None: + value = _message_additional_kwargs(message).get("tool_calls") + return list(value) if isinstance(value, list) else [] + + def _is_hidden_or_control_message(message: Any) -> bool: message_type = _message_type(message) additional_kwargs = _message_additional_kwargs(message) @@ -286,6 +307,11 @@ def _checkpoint_messages(snapshot: Any) -> list[Any]: return checkpoint_messages(snapshot) +def _checkpoint_values(snapshot: Any) -> dict[str, Any]: + values = getattr(snapshot, "values", None) + return dict(values) if isinstance(values, dict) else {} + + def _checkpoint_configurable(checkpoint_tuple: Any) -> dict[str, Any]: return checkpoint_configurable(checkpoint_tuple) @@ -322,6 +348,54 @@ def _clean_human_message_for_regenerate(message: Any) -> dict[str, Any]: return clean_message +def _clean_human_message_for_edit(message: Any, *, replacement_id: str, replacement_text: str) -> dict[str, Any]: + source_kwargs = _message_additional_kwargs(message) + additional_kwargs: dict[str, Any] = {} + for key in ("files", "referenced_message_contexts"): + if key in source_kwargs: + additional_kwargs[key] = deepcopy(source_kwargs[key]) + + clean_message: dict[str, Any] = { + "type": "human", + "id": replacement_id, + "content": [{"type": "text", "text": replacement_text}], + "additional_kwargs": additional_kwargs, + } + name = _message_name(message) + if name: + clean_message["name"] = name + return clean_message + + +def _is_terminal_assistant_text_message(message: Any) -> bool: + return _is_visible_ai_message(message) and bool(_message_text(message).strip()) and not _message_tool_calls(message) + + +def _has_active_goal(snapshot: Any) -> bool: + goal = _checkpoint_values(snapshot).get("goal") + return isinstance(goal, dict) and goal.get("status") == "active" + + +def _latest_editable_turn(messages: list[Any], human_message_id: str) -> tuple[int, Any, int, Any, list[str]]: + latest_human_index = next((index for index in range(len(messages) - 1, -1, -1) if _is_visible_human_message(messages[index])), None) + if latest_human_index is None or _message_id(messages[latest_human_index]) != human_message_id: + raise HTTPException(status_code=409, detail="Only the latest completed user turn can be edited") + + source_human = messages[latest_human_index] + last_ai_index: int | None = None + for index, message in enumerate(messages[latest_human_index + 1 :], start=latest_human_index + 1): + if _is_visible_human_message(message): + break + if _is_visible_ai_message(message): + last_ai_index = index + + if last_ai_index is None or not _is_terminal_assistant_text_message(messages[last_ai_index]): + raise HTTPException(status_code=409, detail="Only completed assistant text turns can be edited") + + source_message_ids = [message_id for message in messages[latest_human_index : last_ai_index + 1] if (message_id := _message_id(message))] + return latest_human_index, source_human, last_ai_index, messages[last_ai_index], source_message_ids + + def _event_message_id(row: dict[str, Any]) -> str | None: content = row.get("content") if isinstance(content, BaseMessage): @@ -441,6 +515,32 @@ async def _find_base_checkpoint_before_human( ) +def _run_status_value(record: Any) -> str | None: + status = getattr(record, "status", None) + if isinstance(status, RunStatus): + return status.value + return str(status) if status is not None else None + + +async def _require_successful_source_run(thread_id: str, run_id: str, request: Request) -> RunRecord: + run_mgr = get_run_manager(request) + user_id = await get_current_user(request) + record = await run_mgr.get(run_id, user_id=user_id) + if record is None: + # The run-event journal is the authoritative lookup above. This fallback + # only covers recent in-memory/store hydration gaps for the latest turn. + records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=20) + record = next((candidate for candidate in records if getattr(candidate, "run_id", None) == run_id), None) + if record is None: + raise HTTPException(status_code=409, detail="Could not find source run for assistant message") + record_thread_id = getattr(record, "thread_id", None) + if isinstance(record_thread_id, str) and record_thread_id and record_thread_id != thread_id: + raise HTTPException(status_code=409, detail="Could not find source run for assistant message") + if _run_status_value(record) != RunStatus.success.value: + raise HTTPException(status_code=409, detail="Only successful assistant runs can be edited and rerun") + return record + + async def _prepare_regenerate_payload(thread_id: str, message_id: str, request: Request) -> RegeneratePrepareResponse: accessor, latest_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id) try: @@ -507,6 +607,85 @@ async def _prepare_regenerate_payload(thread_id: str, message_id: str, request: ) +async def _prepare_edit_regenerate_payload( + thread_id: str, + human_message_id: str, + replacement_text: str, + request: Request, +) -> EditRegeneratePrepareResponse: + normalized_text = replacement_text.strip() + if not normalized_text: + raise HTTPException(status_code=409, detail="Edited message cannot be empty") + + accessor, latest_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id) + try: + latest_checkpoint = await accessor.aget(latest_config) + except Exception as exc: + logger.exception("Failed to read latest checkpoint for edit replay thread %s", thread_id) + raise HTTPException(status_code=500, detail="Failed to read latest checkpoint") from exc + 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) + if _has_active_goal(latest_checkpoint): + raise HTTPException(status_code=409, detail="Cannot edit while a goal is active") + + _, source_human, _, source_ai, source_message_ids = _latest_editable_turn(messages, human_message_id) + source_text = get_original_user_content_text(_message_content(source_human), _message_additional_kwargs(source_human)).strip() + if normalized_text == source_text: + raise HTTPException(status_code=409, detail="Edited message is unchanged") + + source_human_id = _message_id(source_human) + source_ai_id = _message_id(source_ai) + if not source_human_id: + raise HTTPException(status_code=409, detail="The source user message is missing an id") + if not source_ai_id: + raise HTTPException(status_code=409, detail="The source assistant message is missing an id") + + base_checkpoint_tuple = await _find_base_checkpoint_before_human(thread_id, source_human_id, request) + target_run_id = await _find_target_run_id(thread_id, source_ai_id, source_ai, source_human, request) + source_record = await _require_successful_source_run(thread_id, target_run_id, request) + checkpoint = _checkpoint_response(base_checkpoint_tuple) + replacement_human_message_id = str(uuid.uuid4()) + source_metadata = getattr(source_record, "metadata", None) or {} + existing_group_id = source_metadata.get("edit_version_group_id") if isinstance(source_metadata, dict) else None + # Reserved for future edit-chain grouping across repeated edits of the same + # original prompt; current visibility still keys off regenerate_from_run_id. + edit_version_group_id = existing_group_id if isinstance(existing_group_id, str) and existing_group_id else source_human_id + metadata = { + "replay_kind": "edit", + "regenerate_from_message_id": source_ai_id, + "regenerate_from_run_id": target_run_id, + "regenerate_checkpoint_id": checkpoint["checkpoint_id"], + "edit_from_message_id": source_human_id, + "edit_message_id": replacement_human_message_id, + "edit_version_group_id": edit_version_group_id, + } + return EditRegeneratePrepareResponse( + input={ + "messages": [ + _clean_human_message_for_edit( + source_human, + replacement_id=replacement_human_message_id, + replacement_text=normalized_text, + ) + ] + }, + checkpoint=checkpoint, + metadata=metadata, + target_run_id=target_run_id, + replacement_human_message_id=replacement_human_message_id, + source_message_ids=source_message_ids, + ) + + +async def _default_history_hidden_run_ids(run_mgr: Any, thread_id: str, *, user_id: str | None) -> set[str]: + superseded_run_ids = await run_mgr.list_successful_regenerate_sources(thread_id, user_id=user_id) + edit_visibility = await run_mgr.list_edit_replay_visibility(thread_id, user_id=user_id) + return set(superseded_run_ids) | set(edit_visibility.hidden_source_run_ids) | set(edit_visibility.hidden_attempt_run_ids) + + # --------------------------------------------------------------------------- # Endpoints # --------------------------------------------------------------------------- @@ -523,6 +702,17 @@ async def prepare_regenerate_run( return await _prepare_regenerate_payload(thread_id, body.message_id, request) +@router.post("/{thread_id}/runs/edit-regenerate/prepare", response_model=EditRegeneratePrepareResponse) +@require_permission("runs", "create", owner_check=True, require_existing=True) +async def prepare_edit_regenerate_run( + thread_id: str, + body: EditRegeneratePrepareRequest, + request: Request, +) -> EditRegeneratePrepareResponse: + """Prepare input and checkpoint for editing then rerunning the latest user turn.""" + return await _prepare_edit_regenerate_payload(thread_id, body.human_message_id, body.replacement_text, request) + + @router.post("/{thread_id}/runs", response_model=RunResponse) @require_permission("runs", "create", owner_check=True, require_existing=True) async def create_run(thread_id: str, body: RunCreateRequest, request: Request) -> RunResponse: @@ -753,12 +943,23 @@ async def list_thread_messages( after_seq: int | None = Query(default=None, ge=1), ) -> list[dict]: """Return displayable messages for a thread (across all runs), with feedback attached.""" - event_store = get_run_event_store(request) - messages = await event_store.list_messages(thread_id, limit=limit, before_seq=before_seq, after_seq=after_seq) - # Resolve the caller once; it is needed both to scope the feedback query # below and to list the thread's runs for turn-duration injection. user_id = await get_current_user(request) + run_mgr = get_run_manager(request) + hidden_run_ids = await _default_history_hidden_run_ids(run_mgr, thread_id, user_id=user_id) + messages, _ = await _scan_visible_thread_messages( + thread_id, + limit=limit, + before_seq=before_seq, + after_seq=after_seq, + request=request, + user_id=user_id, + hidden_run_ids=hidden_run_ids, + include_middleware=True, + include_extra=False, + batch_size=THREAD_MESSAGE_LEGACY_SCAN_BATCH, + ) # Find the last AI message per run_id. AI messages are persisted by # RunJournal with event_type "llm.ai.response" (see runtime/journal.py); @@ -793,7 +994,6 @@ async def list_thread_messages( else: msg["feedback"] = None - run_mgr = get_run_manager(request) runs = await run_mgr.list_by_thread(thread_id, user_id=user_id) run_durations = compute_run_durations(runs) @@ -810,6 +1010,122 @@ async def list_thread_messages( return messages +async def _scan_visible_thread_messages( + thread_id: str, + *, + limit: int, + before_seq: int | None, + after_seq: int | None, + request: Request, + user_id: str | None, + hidden_run_ids: set[str], + include_middleware: bool, + include_extra: bool, + batch_size: int, +) -> tuple[list[dict[str, Any]], bool]: + """Scan raw message rows until ``limit`` visible rows survive filtering.""" + event_store = get_run_event_store(request) + needed = limit + 1 if include_extra else limit + + if after_seq is not None: + visible: list[dict[str, Any]] = [] + scan_after = after_seq + while len(visible) < needed: + raw = await event_store.list_messages( + thread_id, + limit=batch_size, + after_seq=scan_after, + user_id=user_id, + ) + if not raw: + break + _validate_message_scan_rows(raw, thread_id=thread_id, scan_before=None, scan_after=scan_after) + reached_before_bound = False + for row in raw: + if before_seq is not None and row["seq"] >= before_seq: + reached_before_bound = True + break + if (not include_middleware and _is_thread_history_hidden_message_row(row)) or row.get("run_id") in hidden_run_ids: + continue + visible.append(row) + if len(visible) == needed: + break + next_scan_after = max(row["seq"] for row in raw) + if next_scan_after <= scan_after: + _raise_non_advancing_message_scan(thread_id=thread_id, scan_before=None, scan_after=scan_after, next_cursor=next_scan_after, row_count=len(raw)) + scan_after = next_scan_after + if reached_before_bound or len(raw) < batch_size: + break + has_more = len(visible) > limit + return visible[:limit], has_more + + visible_desc: list[dict[str, Any]] = [] + scan_before = before_seq + while len(visible_desc) < needed: + raw = await event_store.list_messages( + thread_id, + limit=batch_size, + before_seq=scan_before, + user_id=user_id, + ) + if not raw: + break + _validate_message_scan_rows(raw, thread_id=thread_id, scan_before=scan_before, scan_after=None) + for row in reversed(raw): + if (not include_middleware and _is_thread_history_hidden_message_row(row)) or row.get("run_id") in hidden_run_ids: + continue + visible_desc.append(row) + if len(visible_desc) == needed: + break + next_scan_before = min(row["seq"] for row in raw) + if scan_before is not None and next_scan_before >= scan_before: + _raise_non_advancing_message_scan(thread_id=thread_id, scan_before=scan_before, scan_after=None, next_cursor=next_scan_before, row_count=len(raw)) + scan_before = next_scan_before + if len(raw) < batch_size: + break + has_more = len(visible_desc) > limit + return list(reversed(visible_desc[:limit])), has_more + + +def _validate_message_scan_rows( + rows: list[dict[str, Any]], + *, + thread_id: str, + scan_before: int | None, + scan_after: int | None, +) -> None: + invalid_seq_rows = [row for row in rows if not isinstance(row.get("seq"), int)] + if invalid_seq_rows: + logger.error( + "Thread message scan found rows without sequence values: thread_id=%s scan_before=%s scan_after=%s row_count=%d invalid_count=%d", + thread_id, + scan_before, + scan_after, + len(rows), + len(invalid_seq_rows), + ) + raise RuntimeError("Run event message rows are missing sequence values") + + +def _raise_non_advancing_message_scan( + *, + thread_id: str, + scan_before: int | None, + scan_after: int | None, + next_cursor: int, + row_count: int, +) -> None: + logger.error( + "Thread message scan cursor did not advance: thread_id=%s scan_before=%s scan_after=%s next_cursor=%s row_count=%d", + thread_id, + scan_before, + scan_after, + next_cursor, + row_count, + ) + raise RuntimeError("Run event message scan did not advance its cursor") + + async def _scan_thread_message_page( thread_id: str, *, @@ -819,57 +1135,20 @@ async def _scan_thread_message_page( user_id: str | None, ) -> tuple[list[dict[str, Any]], bool]: """Select the newest ``limit + 1`` page-eligible rows before a cursor.""" - event_store = get_run_event_store(request) run_mgr = get_run_manager(request) - superseded_run_ids = await run_mgr.list_successful_regenerate_sources(thread_id, user_id=user_id) - visible_desc: list[dict[str, Any]] = [] - scan_before = before_seq - - while len(visible_desc) < limit + 1: - raw = await event_store.list_messages( - thread_id, - limit=THREAD_MESSAGE_PAGE_SCAN_BATCH, - before_seq=scan_before, - user_id=user_id, - ) - if not raw: - break - - invalid_seq_rows = [row for row in raw if not isinstance(row.get("seq"), int)] - if invalid_seq_rows: - logger.error( - "Thread message scan found rows without sequence values: thread_id=%s scan_before=%s row_count=%d invalid_count=%d", - thread_id, - scan_before, - len(raw), - len(invalid_seq_rows), - ) - raise RuntimeError("Run event message rows are missing sequence values") - - for row in reversed(raw): - if _is_thread_history_hidden_message_row(row) or row.get("run_id") in superseded_run_ids: - continue - visible_desc.append(row) - if len(visible_desc) == limit + 1: - break - - raw_seqs = [row["seq"] for row in raw] - next_scan_before = min(raw_seqs) - if scan_before is not None and next_scan_before >= scan_before: - logger.error( - "Thread message scan cursor did not advance: thread_id=%s scan_before=%s next_scan_before=%s row_count=%d", - thread_id, - scan_before, - next_scan_before, - len(raw), - ) - raise RuntimeError("Run event message scan did not advance its cursor") - scan_before = next_scan_before - if len(raw) < THREAD_MESSAGE_PAGE_SCAN_BATCH: - break - - has_more = len(visible_desc) > limit - return list(reversed(visible_desc[:limit])), has_more + hidden_run_ids = await _default_history_hidden_run_ids(run_mgr, thread_id, user_id=user_id) + return await _scan_visible_thread_messages( + thread_id, + limit=limit, + before_seq=before_seq, + after_seq=None, + request=request, + user_id=user_id, + hidden_run_ids=hidden_run_ids, + include_middleware=False, + include_extra=True, + batch_size=THREAD_MESSAGE_PAGE_SCAN_BATCH, + ) async def _enrich_thread_message_page( diff --git a/backend/packages/harness/deerflow/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index a0ea18afb..d4e475de7 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -191,6 +191,28 @@ class RunRepository(RunStore): result = await session.execute(stmt) return {value for value in result.scalars() if isinstance(value, str) and value} + async def list_edit_regenerate_runs( + self, + thread_id, + *, + user_id: str | None | _AutoSentinel = AUTO, + ): + resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_edit_regenerate_runs") + replay_kind = RunRow.metadata_json["replay_kind"].as_string() + source = RunRow.metadata_json["regenerate_from_run_id"].as_string() + stmt = select(RunRow).where( + RunRow.thread_id == thread_id, + replay_kind == "edit", + source.is_not(None), + source != "", + ) + if resolved_user_id is not None: + stmt = stmt.where(RunRow.user_id == resolved_user_id) + stmt = stmt.order_by(RunRow.created_at.asc()) + async with self._sf() as session: + result = await session.execute(stmt) + return [self._row_to_dict(row) for row in result.scalars()] + async def get_many_by_thread( self, thread_id, diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index eeb965c8b..f658437ef 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -21,6 +21,7 @@ from deerflow.utils.time import is_lease_expired from deerflow.utils.time import now_iso as _now_iso from .schemas import DisconnectMode, RunStatus, ThreadOperationKind +from .store.base import EditReplayVisibility if TYPE_CHECKING: from deerflow.config.run_ownership_config import RunOwnershipConfig @@ -706,6 +707,67 @@ class RunManager: sources.add(source) return sources + @staticmethod + def _record_status_value(record: RunRecord) -> str: + status = record.status + return status.value if isinstance(status, RunStatus) else str(status) + + @staticmethod + def _compute_edit_replay_visibility(records: list[RunRecord]) -> EditReplayVisibility: + latest_attempt_by_source: dict[str, tuple[str, str]] = {} + failed_attempts: set[str] = set() + for record in sorted(records, key=lambda item: item.created_at): + metadata = record.metadata or {} + if metadata.get("replay_kind") != "edit": + continue + source = metadata.get("regenerate_from_run_id") + if not isinstance(source, str) or not source: + continue + status = RunManager._record_status_value(record) + latest_attempt_by_source[source] = (record.run_id, status) + if status in {RunStatus.error.value, RunStatus.timeout.value, RunStatus.interrupted.value}: + failed_attempts.add(record.run_id) + + hidden_sources: set[str] = set() + for source, (_, status) in latest_attempt_by_source.items(): + if status in {RunStatus.pending.value, RunStatus.running.value, RunStatus.success.value}: + hidden_sources.add(source) + return EditReplayVisibility( + hidden_source_run_ids=hidden_sources, + hidden_attempt_run_ids=failed_attempts, + ) + + async def list_edit_replay_visibility( + self, + thread_id: str, + *, + user_id: str | None | _AutoSentinel = AUTO, + ) -> EditReplayVisibility: + """Return run-id visibility rules for edit-and-rerun attempts. + + Store rows cover reload/multi-worker history. Current-process records + override the same run ids because they may have newer terminal status + than the persisted snapshot visible when this query started. + """ + resolved_user_id = resolve_user_id(user_id, method_name="RunManager.list_edit_replay_visibility") + records_by_id: dict[str, RunRecord] = {} + if self._store is not None: + rows = await self._store.list_edit_regenerate_runs(thread_id, user_id=resolved_user_id) + for row in rows: + try: + record = self._record_from_store(row) + except Exception: + logger.warning("Failed to map edit replay run row for %s", row.get("run_id"), exc_info=True) + continue + records_by_id[record.run_id] = record + + async with self._lock: + memory_records = [record for record in self._thread_records_locked(thread_id) if resolved_user_id is None or record.user_id == resolved_user_id] + for record in memory_records: + records_by_id[record.run_id] = record + + return self._compute_edit_replay_visibility(list(records_by_id.values())) + async def try_start(self, run_id: str) -> RunStartOutcome: """Transition an uncancelled pending run to running before building the agent.""" async with self._lock: diff --git a/backend/packages/harness/deerflow/runtime/runs/store/base.py b/backend/packages/harness/deerflow/runtime/runs/store/base.py index 385d9714e..bf705eb27 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/base.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/base.py @@ -11,9 +11,16 @@ When user_id is None, no user filtering is applied (single-user mode). from __future__ import annotations import abc +from dataclasses import dataclass, field from typing import Any +@dataclass(frozen=True) +class EditReplayVisibility: + hidden_source_run_ids: set[str] = field(default_factory=set) + hidden_attempt_run_ids: set[str] = field(default_factory=set) + + class RunStore(abc.ABC): @abc.abstractmethod async def put( @@ -69,6 +76,15 @@ class RunStore(abc.ABC): """ raise NotImplementedError + async def list_edit_regenerate_runs( + self, + thread_id: str, + *, + user_id: str | None = None, + ) -> list[dict[str, Any]]: + """Return all edit-regenerate attempt runs for one thread, oldest first.""" + raise NotImplementedError + async def get_many_by_thread( self, thread_id: str, diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index e92356085..ba8eb4fd5 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -105,6 +105,22 @@ class MemoryRunStore(RunStore): sources.add(source) return sources + async def list_edit_regenerate_runs(self, thread_id, *, user_id=None): + run_ids = self._runs_by_thread.get(thread_id) or () + results = [] + for run_id in run_ids: + run = self._runs.get(run_id) + if run is None: + continue + if user_id is not None and run.get("user_id") != user_id: + continue + metadata = run.get("metadata") or {} + source = metadata.get("regenerate_from_run_id") + if metadata.get("replay_kind") == "edit" and isinstance(source, str) and source: + results.append(run) + results.sort(key=lambda r: r["created_at"]) + return results + async def get_many_by_thread(self, thread_id, run_ids, *, user_id=None): thread_run_ids = self._runs_by_thread.get(thread_id) or () return {run_id: run for run_id in thread_run_ids if run_id in run_ids and (run := self._runs.get(run_id)) is not None and run.get("operation_kind", "run") == "run" and (user_id is None or run.get("user_id") == user_id)} diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 146b6037a..b2cc97612 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -521,6 +521,7 @@ async def run_agent( workspace_changes_user_id: str | None = None snapshot_capture_failed = False llm_error_fallback_message: str | None = None + checkpoint_rollback_completed = False # Message ids checkpointed *before* this run started. The stream loop uses # this set to mask out ``deerflow_error_fallback`` markers that belong to # earlier runs on the same thread — without it, one stale fallback in @@ -890,7 +891,7 @@ async def run_agent( **terminal_status_kwargs, ) try: - await _rollback_to_pre_run_checkpoint( + checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint( accessor=accessor, checkpointer=checkpointer, thread_id=thread_id, @@ -912,6 +913,7 @@ async def run_agent( if error_msg is None and journal is not None: error_msg = journal.llm_error_fallback_message error_msg = error_msg or "LLM provider failed after retries" + await _ensure_finalizing_before_edit_failure(run_manager, record) await run_manager.set_status( run_id, RunStatus.error, @@ -964,7 +966,7 @@ async def run_agent( **terminal_status_kwargs, ) try: - await _rollback_to_pre_run_checkpoint( + checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint( accessor=accessor, checkpointer=checkpointer, thread_id=thread_id, @@ -976,6 +978,7 @@ async def run_agent( except Exception: logger.warning("Run %s cancellation rollback failed", run_id, exc_info=True) else: + await _ensure_finalizing_before_edit_failure(run_manager, record) await run_manager.set_status( run_id, RunStatus.interrupted, @@ -986,6 +989,7 @@ async def run_agent( except Exception as exc: error_msg = f"{exc}" logger.exception("Run %s failed: %s", run_id, error_msg) + await _ensure_finalizing_before_edit_failure(run_manager, record) await run_manager.set_status( run_id, RunStatus.error, @@ -1008,6 +1012,30 @@ async def run_agent( run_id, ) + if not record.ownership_lost and _is_edit_replay_run(record) and record.status != RunStatus.success: + if not record.finalizing: + await run_manager.set_finalizing(run_id, True) + try: + if not checkpoint_rollback_completed: + checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint( + accessor=accessor, + checkpointer=checkpointer, + thread_id=thread_id, + run_id=run_id, + rollback_point=rollback_point, + snapshot_capture_failed=snapshot_capture_failed, + ) + if checkpoint_rollback_completed: + await _publish_restored_checkpoint_values( + bridge=bridge, + run_id=run_id, + accessor=accessor, + thread_id=thread_id, + ) + logger.info("Run %s edit replay restored pre-run checkpoint %s", run_id, pre_run_checkpoint_id) + except Exception: + logger.warning("Run %s edit replay rollback failed", run_id, exc_info=True) + # Persist any subagent step events still buffered (#3779) — including on # abort/exception paths, where the stream loop broke before its own flush. if not record.ownership_lost and subagent_events is not None: @@ -1076,7 +1104,7 @@ async def run_agent( except Exception: logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True) - if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted: + if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted and not _is_edit_replay_run(record): try: await run_manager.wait_for_prior_finalizing(thread_id, run_id) if not await run_manager.has_later_started_run(thread_id, run_id): @@ -1479,6 +1507,31 @@ async def _prepare_goal_continuation_input( return {"messages": [make_goal_continuation_message(updated_goal, evaluation)]} +def _is_edit_replay_run(record: RunRecord) -> bool: + metadata = record.metadata or {} + return metadata.get("replay_kind") == "edit" + + +async def _ensure_finalizing_before_edit_failure(run_manager: RunManager, record: RunRecord) -> None: + if _is_edit_replay_run(record) and not record.finalizing: + await run_manager.set_finalizing(record.run_id, True) + + +async def _publish_restored_checkpoint_values( + *, + bridge: StreamBridge, + run_id: str, + accessor: CheckpointStateAccessor | None, + thread_id: str, +) -> None: + if accessor is None: + return + snapshot = await accessor.aget({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}) + values = getattr(snapshot, "values", None) + if isinstance(values, dict): + await bridge.publish(run_id, "values", serialize(values, mode="values")) + + @dataclass(frozen=True) class RollbackPoint: """Materialized pre-run state used to restore the thread after cancellation. @@ -1657,8 +1710,8 @@ async def _rollback_to_pre_run_checkpoint( run_id: str, rollback_point: RollbackPoint | None, snapshot_capture_failed: bool, -) -> None: - """Restore the complete pre-run state after a cancelled run. +) -> bool: + """Restore the complete pre-run state and report whether it completed. Full mode forks the captured pre-run checkpoint and overwrites messages; all other channels inherit from that parent. Delta mode cannot safely fork @@ -1669,27 +1722,27 @@ async def _rollback_to_pre_run_checkpoint( """ if checkpointer is None: logger.info("Run %s rollback requested but no checkpointer is configured", run_id) - return + return False if snapshot_capture_failed: logger.warning("Run %s rollback skipped: pre-run checkpoint capture failed", run_id) - return + return False 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 + return True 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 + return False 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 + return False # Compile with the thread's effective schema so middleware-contributed # channels survive (the base ThreadState fallback would silently drop @@ -1733,7 +1786,7 @@ async def _rollback_to_pre_run_checkpoint( pending_writes = rollback_point.pending_writes if not pending_writes: - return + return True writes_by_task: dict[str, list[tuple[str, Any]]] = {} for item in pending_writes: @@ -1753,6 +1806,7 @@ async def _rollback_to_pre_run_checkpoint( writes, task_id=task_id, ) + return True def _new_checkpoint_marker() -> dict[str, str]: diff --git a/backend/tests/test_history_batch_queries.py b/backend/tests/test_history_batch_queries.py index 0ea6696e7..6f8f70d7a 100644 --- a/backend/tests/test_history_batch_queries.py +++ b/backend/tests/test_history_batch_queries.py @@ -6,6 +6,7 @@ import pytest from deerflow.runtime import RunManager, RunStatus from deerflow.runtime.events.store.memory import MemoryRunEventStore +from deerflow.runtime.runs.manager import EditReplayVisibility from deerflow.runtime.runs.store.memory import MemoryRunStore @@ -288,3 +289,121 @@ async def test_run_manager_batch_history_methods_allow_explicit_unscoped_access( assert sources == {"source-alice", "source-bob"} assert set(records) == {"regen-alice", "regen-bob"} + + +@pytest.mark.anyio +async def test_memory_run_store_lists_edit_replay_runs_unbounded_and_owner_scoped(): + store = MemoryRunStore() + for index in range(105): + await store.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success") + await store.put( + "edit-success", + thread_id="t1", + user_id="alice", + status="success", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"}, + ) + await store.put( + "edit-error", + thread_id="t1", + user_id="alice", + status="error", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"}, + ) + await store.put( + "edit-bob", + thread_id="t1", + user_id="bob", + status="success", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"}, + ) + + rows = await store.list_edit_regenerate_runs("t1", user_id="alice") + + assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"] + + +@pytest.mark.anyio +async def test_run_repository_lists_edit_replay_runs_unbounded_and_owner_scoped(tmp_path): + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.persistence.run import RunRepository + + await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}", sqlite_dir=str(tmp_path)) + try: + repo = RunRepository(get_session_factory()) + for index in range(105): + await repo.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success") + await repo.put( + "edit-success", + thread_id="t1", + user_id="alice", + status="success", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"}, + ) + await repo.put( + "edit-error", + thread_id="t1", + user_id="alice", + status="error", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"}, + ) + await repo.put( + "edit-bob", + thread_id="t1", + user_id="bob", + status="success", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"}, + ) + + rows = await repo.list_edit_regenerate_runs("t1", user_id="alice") + + assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"] + finally: + await close_engine() + + +@pytest.mark.anyio +async def test_run_manager_computes_edit_replay_visibility_by_latest_attempt(): + manager = RunManager() + running = await manager.create( + "t1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-running"}, + ) + running.status = RunStatus.running + success = await manager.create( + "t1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"}, + ) + success.status = RunStatus.success + failed = await manager.create( + "t1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-failed"}, + ) + failed.status = RunStatus.error + older_success = await manager.create( + "t1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"}, + ) + older_success.status = RunStatus.success + newer_failed = await manager.create( + "t1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"}, + ) + newer_failed.status = RunStatus.interrupted + older_failed_then_success = await manager.create( + "t1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"}, + ) + older_failed_then_success.status = RunStatus.error + newer_success_after_failure = await manager.create( + "t1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"}, + ) + newer_success_after_failure.status = RunStatus.success + + visibility = await manager.list_edit_replay_visibility("t1", user_id=None) + + assert visibility == EditReplayVisibility( + hidden_source_run_ids={"source-running", "source-success", "source-retry-success"}, + hidden_attempt_run_ids={failed.run_id, newer_failed.run_id, older_failed_then_success.run_id}, + ) diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index abe3f92d1..0fa6cf525 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -814,6 +814,150 @@ async def test_run_agent_marks_llm_error_fallback_as_error_status(): bridge.publish_end.assert_awaited_once_with(record.run_id) +@pytest.mark.anyio +async def test_run_agent_rolls_back_failed_edit_replay_and_publishes_restored_values(): + run_manager = RunManager() + record = await run_manager.create( + "thread-1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-run"}, + ) + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + before_messages = [HumanMessage(id="h1", content="original"), AIMessage(id="a1", content="answer")] + + 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): + return SimpleNamespace( + values={"messages": before_messages}, + 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): + raise RuntimeError("edit replay failed") + if False: + yield # pragma: no cover - keep this an async generator + + error_status_finalizing_states: list[bool] = [] + original_set_status = run_manager.set_status + + async def _set_status(run_id, status, **kwargs): + if run_id == record.run_id and status == RunStatus.error: + error_status_finalizing_states.append(record.finalizing) + return await original_set_status(run_id, status, **kwargs) + + run_manager.set_status = _set_status # type: ignore[method-assign] + with patch( + "deerflow.runtime.runs.worker._rollback_to_pre_run_checkpoint", + new_callable=AsyncMock, + ) as rollback: + rollback.return_value = True + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=DummyCheckpointer()), + agent_factory=lambda *, config: DummyAgent(), + graph_input={}, + config={}, + ) + + fetched = await run_manager.get(record.run_id) + assert fetched is not None + assert fetched.status == RunStatus.error + assert error_status_finalizing_states == [True] + rollback.assert_awaited_once() + publish_events = [call_args.args[1] for call_args in bridge.publish.await_args_list] + assert "error" in publish_events + assert "values" in publish_events + assert publish_events.index("values") > publish_events.index("error") + bridge.publish_end.assert_awaited_once_with(record.run_id) + + +@pytest.mark.anyio +async def test_failed_edit_replay_does_not_publish_restored_values_when_snapshot_capture_failed(): + run_manager = RunManager() + record = await run_manager.create( + "thread-1", + metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-run"}, + ) + 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("snapshot capture failed") + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + raise RuntimeError("edit replay failed") + if False: + yield # pragma: no cover - keep this an async generator + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=DummyCheckpointer()), + agent_factory=lambda *, config: DummyAgent(), + graph_input={}, + config={}, + ) + + fetched = await run_manager.get(record.run_id) + assert fetched is not None + assert fetched.status == RunStatus.error + publish_events = [call_args.args[1] for call_args in bridge.publish.await_args_list] + assert "error" in publish_events + assert "values" not in publish_events + bridge.publish_end.assert_awaited_once_with(record.run_id) + + @pytest.mark.anyio async def test_run_agent_defaults_root_run_name_from_assistant_id(): run_manager = RunManager() diff --git a/backend/tests/test_thread_messages_feedback.py b/backend/tests/test_thread_messages_feedback.py index 027fac0f3..935d8b884 100644 --- a/backend/tests/test_thread_messages_feedback.py +++ b/backend/tests/test_thread_messages_feedback.py @@ -15,6 +15,7 @@ from _router_auth_helpers import make_authed_test_app from fastapi.testclient import TestClient from app.gateway.routers import thread_runs +from deerflow.runtime.runs.manager import EditReplayVisibility def _make_app(messages, feedback_grouped): @@ -32,6 +33,8 @@ def _make_app(messages, feedback_grouped): # list_thread_messages also calls run_manager.list_by_thread to inject # turn durations; stub it to return no runs so that path stays inert. run_manager = MagicMock() + run_manager.list_successful_regenerate_sources = AsyncMock(return_value=set()) + run_manager.list_edit_replay_visibility = AsyncMock(return_value=EditReplayVisibility()) run_manager.list_by_thread = AsyncMock(return_value=[]) app.state.run_manager = run_manager diff --git a/backend/tests/test_thread_messages_page.py b/backend/tests/test_thread_messages_page.py index 07bfb6129..6b6360db2 100644 --- a/backend/tests/test_thread_messages_page.py +++ b/backend/tests/test_thread_messages_page.py @@ -15,18 +15,28 @@ from app.gateway.routers import thread_runs from deerflow.runtime import RunRecord from deerflow.runtime.events.store.memory import MemoryRunEventStore from deerflow.runtime.journal import build_branch_history_seed_events +from deerflow.runtime.runs.manager import EditReplayVisibility -def _make_app(event_store: MemoryRunEventStore, *, superseded: set[str] | None = None, records=None, feedback=None): +def _make_app( + event_store: MemoryRunEventStore, + *, + superseded: set[str] | None = None, + edit_visibility: EditReplayVisibility | None = None, + records=None, + feedback=None, +): app = make_authed_test_app() app.include_router(thread_runs.router) app.state.run_event_store = event_store run_manager = AsyncMock() run_manager.list_successful_regenerate_sources.return_value = superseded or set() + run_manager.list_edit_replay_visibility.return_value = edit_visibility or EditReplayVisibility() run_manager.get_many_by_thread.return_value = records or {} app.state.run_manager = run_manager feedback_repo = AsyncMock() feedback_repo.list_by_run_ids.return_value = feedback or {} + feedback_repo.list_by_thread_grouped.return_value = feedback or {} app.state.feedback_repo = feedback_repo return app @@ -187,6 +197,76 @@ def test_thread_page_filters_all_successfully_superseded_runs_before_filling(): assert body["next_before_seq"] is None +def test_thread_page_applies_edit_replay_visibility_before_filling(): + store = MemoryRunEventStore() + + async def seed(): + await _put_message(store, "source-success", "ai", "source-success-answer") + await _put_message(store, "edit-success", "ai", "edit-success-answer") + await _put_message(store, "source-failed", "ai", "source-failed-answer") + await _put_message(store, "edit-failed", "ai", "edit-failed-answer") + + asyncio.run(seed()) + app = _make_app( + store, + edit_visibility=EditReplayVisibility( + hidden_source_run_ids={"source-success"}, + hidden_attempt_run_ids={"edit-failed"}, + ), + ) + with TestClient(app) as client: + response = client.get("/api/threads/thread-1/messages/page?limit=4") + + body = response.json() + assert [row["run_id"] for row in body["data"]] == ["edit-success", "source-failed"] + assert body["has_more"] is False + assert body["next_before_seq"] is None + + +def test_legacy_thread_messages_apply_regenerate_and_edit_replay_visibility(): + store = MemoryRunEventStore() + + async def seed(): + await _put_message(store, "regen-source", "ai", "regen-source-answer") + await _put_message(store, "source-success", "ai", "source-success-answer") + await _put_message(store, "edit-success", "ai", "edit-success-answer") + await _put_message(store, "source-failed", "ai", "source-failed-answer") + await _put_message(store, "edit-failed", "ai", "edit-failed-answer") + + asyncio.run(seed()) + app = _make_app( + store, + superseded={"regen-source"}, + edit_visibility=EditReplayVisibility( + hidden_source_run_ids={"source-success"}, + hidden_attempt_run_ids={"edit-failed"}, + ), + ) + with TestClient(app) as client: + response = client.get("/api/threads/thread-1/messages?limit=5") + + assert [row["run_id"] for row in response.json()] == ["edit-success", "source-failed"] + + +def test_legacy_thread_messages_applies_limit_after_visibility_filtering(monkeypatch): + monkeypatch.setattr(thread_runs, "THREAD_MESSAGE_LEGACY_SCAN_BATCH", 2) + store = MemoryRunEventStore() + + async def seed(): + await _put_message(store, "visible-old-1", "ai", "visible-old-1-answer") + await _put_message(store, "visible-old-2", "ai", "visible-old-2-answer") + await _put_message(store, "hidden-new-1", "ai", "hidden-new-1-answer") + await _put_message(store, "hidden-new-2", "ai", "hidden-new-2-answer") + + asyncio.run(seed()) + app = _make_app(store, superseded={"hidden-new-1", "hidden-new-2"}) + with TestClient(app) as client: + response = client.get("/api/threads/thread-1/messages?limit=2") + + assert response.status_code == 200 + assert [row["run_id"] for row in response.json()] == ["visible-old-1", "visible-old-2"] + + def test_thread_page_keeps_earlier_branch_turns_when_the_last_one_is_regenerated(): """Regenerating a branch's inherited answer must not delete the turns before it (#4458). @@ -253,7 +333,7 @@ def test_thread_page_logs_when_scan_cursor_does_not_advance(caplog): assert "Thread message scan cursor did not advance" in caplog.text assert "thread_id=thread-1" in caplog.text assert "scan_before=10" in caplog.text - assert "next_scan_before=10" in caplog.text + assert "next_cursor=10" in caplog.text assert "row_count=1" in caplog.text diff --git a/backend/tests/test_thread_regenerate_prepare.py b/backend/tests/test_thread_regenerate_prepare.py index 21958343b..7672fa1ec 100644 --- a/backend/tests/test_thread_regenerate_prepare.py +++ b/backend/tests/test_thread_regenerate_prepare.py @@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, patch import pytest from fastapi import HTTPException -from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from langgraph.checkpoint.base import empty_checkpoint, uuid6 from langgraph.checkpoint.memory import InMemorySaver @@ -15,7 +15,16 @@ from deerflow.runtime import RunStatus from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY -def _checkpoint(checkpoint_id: str, messages: list[object], *, metadata: dict | None = None): +def _checkpoint( + checkpoint_id: str, + messages: list[object], + *, + metadata: dict | None = None, + goal: dict | None = None, +): + channel_values = {"messages": messages} + if goal is not None: + channel_values["goal"] = goal return SimpleNamespace( config={ "configurable": { @@ -25,7 +34,7 @@ def _checkpoint(checkpoint_id: str, messages: list[object], *, metadata: dict | "checkpoint_map": None, } }, - checkpoint={"channel_values": {"messages": messages}}, + checkpoint={"channel_values": channel_values}, metadata=metadata or {}, ) @@ -162,6 +171,9 @@ class FakeRunManager: async def list_by_thread(self, thread_id, *, user_id=None, limit=100): return self.records[:limit] + async def get(self, run_id, *, user_id=None): + return next((record for record in self.records if record.run_id == run_id), None) + def _request(checkpointer, event_store, *, run_manager=None, user_id="user-1"): from app.gateway.auth_disabled import AUTH_SOURCE_SESSION @@ -534,6 +546,338 @@ def test_prepare_regenerate_payload_rejects_legacy_branch_when_source_checkpoint assert len(branch_history) == 1 +def test_prepare_edit_regenerate_payload_returns_new_human_and_edit_metadata(): + from app.gateway.routers import thread_runs + + human = HumanMessage( + id="human-1", + content="injected\n\noriginal question", + name="researcher", + additional_kwargs={ + ORIGINAL_USER_CONTENT_KEY: "original question", + "files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}], + "referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}], + "hide_from_ui": False, + "run_id": "old-run", + "timestamp": "2026-07-22T00:00:00Z", + "middleware_private": "do-not-copy", + }, + ) + ai = AIMessage(id="ai-1", content="answer v1") + base = _checkpoint("ckpt-base", []) + after_human = _checkpoint("ckpt-human", [human]) + latest = _checkpoint("ckpt-ai", [human, ai]) + checkpointer = FakeCheckpointer([latest, after_human, base]) + event_store = FakeEventStore( + [ + { + "run_id": "run-old", + "event_type": "llm.ai.response", + "category": "message", + "content": {"id": "ai-1", "type": "ai", "content": "answer v1"}, + "metadata": {"caller": "lead_agent"}, + } + ] + ) + run_manager = FakeRunManager( + [ + SimpleNamespace( + run_id="run-old", + status=RunStatus.success, + metadata={}, + last_ai_message="answer v1", + ) + ] + ) + + response = asyncio.run( + thread_runs._prepare_edit_regenerate_payload( + "thread-1", + "human-1", + " updated question\nwith details ", + _request(checkpointer, event_store, run_manager=run_manager), + ) + ) + + assert response.checkpoint == { + "checkpoint_ns": "", + "checkpoint_id": "ckpt-base", + "checkpoint_map": None, + } + assert response.target_run_id == "run-old" + assert response.replacement_human_message_id != "human-1" + assert response.source_message_ids == ["human-1", "ai-1"] + assert response.metadata == { + "replay_kind": "edit", + "regenerate_from_message_id": "ai-1", + "regenerate_from_run_id": "run-old", + "regenerate_checkpoint_id": "ckpt-base", + "edit_from_message_id": "human-1", + "edit_message_id": response.replacement_human_message_id, + "edit_version_group_id": "human-1", + } + replacement = response.input["messages"][0] + assert replacement == { + "type": "human", + "id": response.replacement_human_message_id, + "name": "researcher", + "content": [{"type": "text", "text": "updated question\nwith details"}], + "additional_kwargs": { + "files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}], + "referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}], + }, + } + + +@pytest.mark.parametrize( + ("replacement_text", "detail"), + [ + (" \n\t", "Edited message cannot be empty"), + (" original question ", "Edited message is unchanged"), + ], +) +def test_prepare_edit_regenerate_payload_rejects_empty_or_unchanged_text(replacement_text: str, detail: str): + from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload + + human = HumanMessage(id="human-1", content="original question") + ai = AIMessage(id="ai-1", content="answer") + checkpointer = FakeCheckpointer( + [ + _checkpoint("ckpt-ai", [human, ai]), + _checkpoint("ckpt-human", [human]), + _checkpoint("ckpt-base", []), + ] + ) + event_store = FakeEventStore( + [ + { + "run_id": "run-old", + "event_type": "llm.ai.response", + "category": "message", + "content": {"id": "ai-1", "type": "ai", "content": "answer"}, + "metadata": {"caller": "lead_agent"}, + } + ] + ) + run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.success, metadata={}, last_ai_message="answer")]) + + with pytest.raises(HTTPException) as exc: + asyncio.run( + _prepare_edit_regenerate_payload( + "thread-1", + "human-1", + replacement_text, + _request(checkpointer, event_store, run_manager=run_manager), + ) + ) + + assert exc.value.status_code == 409 + assert exc.value.detail == detail + + +def test_prepare_edit_regenerate_payload_requires_latest_human_turn(): + from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload + + old_human = HumanMessage(id="human-old", content="old question") + old_ai = AIMessage(id="ai-old", content="old answer") + latest_human = HumanMessage(id="human-latest", content="latest question") + latest_ai = AIMessage(id="ai-latest", content="latest answer") + checkpointer = FakeCheckpointer( + [ + _checkpoint("ckpt-latest", [old_human, old_ai, latest_human, latest_ai]), + _checkpoint("ckpt-base", []), + ] + ) + + with pytest.raises(HTTPException) as exc: + asyncio.run( + _prepare_edit_regenerate_payload( + "thread-1", + "human-old", + "edited old question", + _request(checkpointer, FakeEventStore([])), + ) + ) + + assert exc.value.status_code == 409 + assert exc.value.detail == "Only the latest completed user turn can be edited" + + +def test_prepare_edit_regenerate_payload_allows_answered_historical_clarification(): + from app.gateway.routers import thread_runs + + old_human = HumanMessage(id="human-old", content="ambiguous request") + clarification = ToolMessage( + id="tool-clarify", + tool_call_id="call-clarify", + content="need input", + artifact={ + "human_input": { + "version": 1, + "kind": "human_input_request", + "request_id": "clarify-1", + "prompt": "Which format?", + } + }, + ) + clarification_answer = HumanMessage( + id="human-clarify-answer", + content="Use Markdown", + additional_kwargs={ + "hide_from_ui": True, + "human_input_response": { + "version": 1, + "kind": "human_input_response", + "request_id": "clarify-1", + "source": "user", + "value": "Use Markdown", + }, + }, + ) + latest_human = HumanMessage(id="human-latest", content="latest question") + latest_ai = AIMessage(id="ai-latest", content="latest answer") + historical_messages = [old_human, clarification, clarification_answer] + checkpointer = FakeCheckpointer( + [ + _checkpoint("ckpt-ai", [*historical_messages, latest_human, latest_ai]), + _checkpoint("ckpt-human", [*historical_messages, latest_human]), + _checkpoint("ckpt-base", historical_messages), + ] + ) + event_store = FakeEventStore( + [ + { + "run_id": "run-latest", + "event_type": "llm.ai.response", + "category": "message", + "content": {"id": "ai-latest", "type": "ai", "content": "latest answer"}, + "metadata": {"caller": "lead_agent"}, + } + ] + ) + run_manager = FakeRunManager( + [ + SimpleNamespace( + run_id="run-latest", + status=RunStatus.success, + metadata={}, + last_ai_message="latest answer", + ) + ] + ) + + response = asyncio.run( + thread_runs._prepare_edit_regenerate_payload( + "thread-1", + "human-latest", + "updated latest question", + _request(checkpointer, event_store, run_manager=run_manager), + ) + ) + + assert response.checkpoint["checkpoint_id"] == "ckpt-base" + assert response.target_run_id == "run-latest" + assert response.source_message_ids == ["human-latest", "ai-latest"] + + +def test_prepare_edit_regenerate_payload_rejects_open_clarification_turn(): + from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload + + human = HumanMessage(id="human-1", content="question") + tool = ToolMessage( + id="tool-1", + tool_call_id="call-1", + content="need input", + artifact={"human_input": {"request_id": "clarify-1", "status": "pending"}}, + ) + checkpointer = FakeCheckpointer( + [ + _checkpoint("ckpt-tool", [human, tool]), + _checkpoint("ckpt-human", [human]), + _checkpoint("ckpt-base", []), + ] + ) + + with pytest.raises(HTTPException) as exc: + asyncio.run( + _prepare_edit_regenerate_payload( + "thread-1", + "human-1", + "updated question", + _request(checkpointer, FakeEventStore([])), + ) + ) + + assert exc.value.status_code == 409 + assert exc.value.detail == "Only completed assistant text turns can be edited" + + +def test_prepare_edit_regenerate_payload_rejects_active_goal(): + from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload + + human = HumanMessage(id="human-1", content="question") + ai = AIMessage(id="ai-1", content="answer") + checkpointer = FakeCheckpointer( + [ + _checkpoint("ckpt-ai", [human, ai], goal={"status": "active", "objective": "finish"}), + _checkpoint("ckpt-human", [human]), + _checkpoint("ckpt-base", []), + ] + ) + + with pytest.raises(HTTPException) as exc: + asyncio.run( + _prepare_edit_regenerate_payload( + "thread-1", + "human-1", + "updated question", + _request(checkpointer, FakeEventStore([])), + ) + ) + + assert exc.value.status_code == 409 + assert exc.value.detail == "Cannot edit while a goal is active" + + +def test_prepare_edit_regenerate_payload_requires_successful_source_run(): + from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload + + human = HumanMessage(id="human-1", content="question") + ai = AIMessage(id="ai-1", content="answer") + checkpointer = FakeCheckpointer( + [ + _checkpoint("ckpt-ai", [human, ai]), + _checkpoint("ckpt-human", [human]), + _checkpoint("ckpt-base", []), + ] + ) + event_store = FakeEventStore( + [ + { + "run_id": "run-old", + "event_type": "llm.ai.response", + "category": "message", + "content": {"id": "ai-1", "type": "ai", "content": "answer"}, + "metadata": {"caller": "lead_agent"}, + } + ] + ) + run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.error, metadata={}, last_ai_message="answer")]) + + with pytest.raises(HTTPException) as exc: + asyncio.run( + _prepare_edit_regenerate_payload( + "thread-1", + "human-1", + "updated question", + _request(checkpointer, event_store, run_manager=run_manager), + ) + ) + + assert exc.value.status_code == 409 + assert exc.value.detail == "Only successful assistant runs can be edited and rerun" + + def test_prepare_regenerate_uses_materialized_history_when_raw_messages_are_omitted(): from app.gateway.routers.thread_runs import _prepare_regenerate_payload diff --git a/backend/tests/test_thread_run_query_validation.py b/backend/tests/test_thread_run_query_validation.py index 0f4596610..be3490ffd 100644 --- a/backend/tests/test_thread_run_query_validation.py +++ b/backend/tests/test_thread_run_query_validation.py @@ -7,6 +7,7 @@ from _router_auth_helpers import make_authed_test_app from fastapi.testclient import TestClient from app.gateway.routers import thread_runs +from deerflow.runtime.runs.manager import EditReplayVisibility def _make_app(): @@ -20,6 +21,8 @@ def _make_app(): app.state.run_event_store = event_store run_manager = MagicMock() + run_manager.list_successful_regenerate_sources = AsyncMock(return_value=set()) + run_manager.list_edit_replay_visibility = AsyncMock(return_value=EditReplayVisibility()) run_manager.list_by_thread = AsyncMock(return_value=[]) app.state.run_manager = run_manager return app @@ -61,6 +64,10 @@ def test_read_endpoints_reject_non_positive_seq_cursors(path: str, cursor: str, def test_read_endpoints_accept_positive_limits_and_hit_store(): app = _make_app() + app.state.run_event_store.list_messages.return_value = [ + {"seq": 1, "run_id": "run-1", "event_type": "llm.human.input", "content": {"type": "human", "id": "h1"}}, + {"seq": 2, "run_id": "run-1", "event_type": "llm.human.input", "content": {"type": "human", "id": "h2"}}, + ] with TestClient(app) as client: thread_messages = client.get("/api/threads/thread-1/messages", params={"limit": 1}) run_messages = client.get("/api/threads/thread-1/runs/run-1/messages", params={"limit": 1}) @@ -69,7 +76,8 @@ def test_read_endpoints_accept_positive_limits_and_hit_store(): assert thread_messages.status_code == 200 assert run_messages.status_code == 200 assert run_events.status_code == 200 - app.state.run_event_store.list_messages.assert_awaited_once_with("thread-1", limit=1, before_seq=None, after_seq=None) + assert len(thread_messages.json()) == 1 + app.state.run_event_store.list_messages.assert_awaited_once_with("thread-1", limit=thread_runs.THREAD_MESSAGE_LEGACY_SCAN_BATCH, before_seq=None, user_id=None) app.state.run_event_store.list_messages_by_run.assert_awaited_once_with( "thread-1", "run-1", diff --git a/backend/tests/test_tool_args_schema_no_pydantic_warning.py b/backend/tests/test_tool_args_schema_no_pydantic_warning.py index 6da56347f..63cd572ce 100644 --- a/backend/tests/test_tool_args_schema_no_pydantic_warning.py +++ b/backend/tests/test_tool_args_schema_no_pydantic_warning.py @@ -17,6 +17,7 @@ import warnings import pytest from langchain.tools import ToolRuntime +from langchain_core.utils.function_calling import convert_to_openai_tool from deerflow.sandbox.tools import ( bash_tool, @@ -27,6 +28,7 @@ from deerflow.sandbox.tools import ( str_replace_tool, write_file_tool, ) +from deerflow.tools.builtins.list_uploaded_files_tool import list_uploaded_files from deerflow.tools.builtins.present_file_tool import present_file_tool from deerflow.tools.builtins.setup_agent_tool import setup_agent from deerflow.tools.builtins.task_tool import task_tool @@ -55,6 +57,7 @@ _TOOL_CASES = [ (read_file_tool, {"description": "read", "path": "/tmp/x"}), (write_file_tool, {"description": "write", "path": "/tmp/x", "content": "hi"}), (str_replace_tool, {"description": "replace", "path": "/tmp/x", "old_str": "a", "new_str": "b"}), + (list_uploaded_files, {"include_outline": False, "max_results": 20}), (present_file_tool, {"filepaths": ["/tmp/x"], "tool_call_id": "call-1"}), (view_image_tool, {"image_path": "/tmp/img.png", "tool_call_id": "call-1"}), (task_tool, {"description": "do", "prompt": "go", "subagent_type": "general-purpose", "tool_call_id": "call-1"}), @@ -101,6 +104,13 @@ def test_write_file_append_is_discoverable_in_tool_schema() -> None: assert "append" in append_field.description +def test_list_uploaded_files_model_schema_excludes_injected_runtime() -> None: + """The model-facing schema must not expose ToolRuntime internals.""" + parameters = convert_to_openai_tool(list_uploaded_files)["function"]["parameters"] + + assert set(parameters["properties"]) == {"include_outline", "max_results"} + + @pytest.mark.parametrize("tool_obj", [case[0] for case in _TOOL_CASES], ids=[case[0].name for case in _TOOL_CASES]) def test_model_facing_tool_parameters_have_descriptions(tool_obj) -> None: """Every model-facing tool parameter should explain when and how to use it.""" diff --git a/frontend/AGENTS.md b/frontend/AGENTS.md index 0b28fcf94..a463b5df1 100644 --- a/frontend/AGENTS.md +++ b/frontend/AGENTS.md @@ -85,6 +85,8 @@ Human input requests are a structured message protocol layered on normal chat hi Tool-calling AI messages can contain user-visible text as well as `tool_calls`. `core/messages/utils.ts` keeps these turns in an `assistant:processing` group, and `components/workspace/messages/message-group.tsx` must render the visible text as a processing step instead of treating the message as only tool metadata. This preserves provider text such as error explanations or "trying another approach" notes during tool-heavy runs. +Edit-and-rerun is deliberately latest-turn-only. `core/messages/utils.ts::getLatestEditableTurn()` exposes a human turn only when the transcript is idle and the most recent visible turn ends in a terminal assistant message. `core/threads/hooks.ts::editAndRegenerateMessage()` calls `POST /api/threads/{id}/runs/edit-regenerate/prepare`, submits the returned replacement message/checkpoint/metadata through the same LangGraph stream path as regenerate, optimistically hides the superseded message ids, and clears the optimistic replacement once the persisted replacement arrives. + `MessageGroup` builds its tool-result and browser-preview lookups once per processing group before converting messages to steps. The lookup preserves the first non-empty result and first screenshot-bearing browser view for each tool-call ID, matching the streamed-message display semantics without repeatedly scanning the full group for every tool call. ### Key Patterns @@ -103,6 +105,7 @@ Tool-calling AI messages can contain user-visible text as well as `tool_calls`. - `src/app/workspace/chats/[thread_id]/page.tsx` owns composer busy-state wiring. - `src/app/workspace/chats/[thread_id]/page.tsx` owns branch-from-turn submission and navigation; sidecar `MessageList` instances do not receive the branch action. +- `src/app/workspace/chats/[thread_id]/page.tsx` and `src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx` own edit-and-rerun submission wiring because the page must preserve normal/custom-agent run context; `MessageList` only detects the latest editable user turn and renders the inline editor. - `src/app/workspace/chats/[thread_id]/page.tsx` gates the Workspace Browser trigger and browser right panel on `/api/features -> browser_control.enabled`; default/failed feature discovery hides the browser control so optional backend installs do not show a dead Live socket. - `src/app/workspace/chats/[thread_id]/page.tsx` and `src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx` own active-goal display state for their composer overlays. - `src/components/workspace/messages/message-list.tsx` owns human-input card answered/latest/pending gating; entry pages only translate a submitted card response into `sendMessage` calls. diff --git a/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx b/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx index a92d07ef4..78d76de19 100644 --- a/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx +++ b/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx @@ -92,6 +92,7 @@ export default function AgentChatPage() { pendingUsageMessages, sendMessage, regenerateMessage, + editAndRegenerateMessage, isUploading, isHistoryLoading, hasMoreHistory, @@ -209,6 +210,11 @@ export default function AgentChatPage() { regenerateMessage(threadId, messageId, supersededMessageIds), [regenerateMessage, threadId], ); + const handleEditAndRegenerate = useCallback( + (messageId: string, replacementText: string) => + editAndRegenerateMessage(threadId, messageId, replacementText), + [editAndRegenerateMessage, threadId], + ); const tokenUsageInlineMode = tokenUsageEnabled ? localSettings.tokenUsage.inlineMode @@ -307,6 +313,16 @@ export default function AgentChatPage() { !thread.isLoading } onRegenerateMessage={handleRegenerate} + canEdit={ + !isNewThread && + !isMock && + env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY !== "true" && + !isUploading && + !thread.isLoading && + !hasGoal && + !hasOpenHumanInputCard + } + onEditAndRegenerateMessage={handleEditAndRegenerate} onSubmitHumanInput={ isMock || env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true" ? undefined diff --git a/frontend/src/app/workspace/chats/[thread_id]/page.tsx b/frontend/src/app/workspace/chats/[thread_id]/page.tsx index e78218a20..858702ba3 100644 --- a/frontend/src/app/workspace/chats/[thread_id]/page.tsx +++ b/frontend/src/app/workspace/chats/[thread_id]/page.tsx @@ -104,6 +104,7 @@ export default function ChatPage() { pendingUsageMessages, sendMessage, regenerateMessage, + editAndRegenerateMessage, isUploading, isHistoryLoading, hasMoreHistory, @@ -212,6 +213,11 @@ export default function ChatPage() { regenerateMessage(threadId, messageId, supersededMessageIds), [regenerateMessage, threadId], ); + const handleEditAndRegenerate = useCallback( + (messageId: string, replacementText: string) => + editAndRegenerateMessage(threadId, messageId, replacementText), + [editAndRegenerateMessage, threadId], + ); const handleBranchTurn = useCallback( async (messageId: string, messageIds: string[]) => { if ( @@ -319,6 +325,17 @@ export default function ChatPage() { !thread.isLoading } onRegenerateMessage={handleRegenerate} + canEdit={ + !isNewThread && + !isMock && + env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY !== "true" && + !isUploading && + !thread.isLoading && + !branchThread.isPending && + !hasGoal && + !hasOpenHumanInputCard + } + onEditAndRegenerateMessage={handleEditAndRegenerate} onSubmitHumanInput={ isMock || env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true" ? undefined diff --git a/frontend/src/components/workspace/messages/message-list-item.tsx b/frontend/src/components/workspace/messages/message-list-item.tsx index ccdbe5f80..d69cab389 100644 --- a/frontend/src/components/workspace/messages/message-list-item.tsx +++ b/frontend/src/components/workspace/messages/message-list-item.tsx @@ -1,9 +1,12 @@ import type { Message } from "@langchain/langgraph-sdk"; import { + CheckIcon, FileIcon, Loader2Icon, + PencilIcon, ThumbsDownIcon, ThumbsUpIcon, + XIcon, } from "lucide-react"; import { memo, @@ -26,6 +29,8 @@ import { import { Shimmer } from "@/components/ai-elements/shimmer"; import { Task, TaskTrigger } from "@/components/ai-elements/task"; import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Textarea } from "@/components/ui/textarea"; import { deleteFeedback, upsertFeedback, @@ -59,6 +64,7 @@ import { CitationSourcesPanel } from "../citations/citation-sources-panel"; import { CopyButton } from "../copy-button"; import { ReferenceAttachmentSummary } from "../sidecar/reference-attachments"; import { SlashSkillChip } from "../slash-skill-chip"; +import { Tooltip } from "../tooltip"; import { MarkdownContent } from "./markdown-content"; import { createMarkdownLinkComponent } from "./markdown-link"; @@ -139,6 +145,9 @@ export function MessageListItem({ threadId, artifactPaths = [], showCopyButton = true, + canEdit = false, + isEditPending = false, + onEditAndRegenerate, }: { className?: string; message: Message; @@ -148,8 +157,50 @@ export function MessageListItem({ feedback?: FeedbackData | null; runId?: string; showCopyButton?: boolean; + canEdit?: boolean; + isEditPending?: boolean; + onEditAndRegenerate?: (replacementText: string) => void | Promise; }) { + const { t } = useI18n(); const isHuman = message.type === "human"; + const editableText = useMemo( + () => (isHuman ? (getMessageCopyData(message) ?? "") : ""), + [isHuman, message], + ); + const [isEditing, setIsEditing] = useState(false); + const [draft, setDraft] = useState(""); + const [isSubmittingEdit, setIsSubmittingEdit] = useState(false); + const trimmedDraft = draft.trim(); + const editSubmitDisabled = + isEditPending || + isSubmittingEdit || + trimmedDraft.length === 0 || + trimmedDraft === editableText.trim(); + + const startEditing = useCallback(() => { + setDraft(editableText); + setIsEditing(true); + }, [editableText]); + const cancelEditing = useCallback(() => { + setIsEditing(false); + setDraft(""); + }, []); + const submitEdit = useCallback(async () => { + if (editSubmitDisabled || !onEditAndRegenerate) { + return; + } + setIsSubmittingEdit(true); + try { + const result = await onEditAndRegenerate(trimmedDraft); + if (result !== false) { + setIsEditing(false); + setDraft(""); + } + } finally { + setIsSubmittingEdit(false); + } + }, [editSubmitDisabled, onEditAndRegenerate, trimmedDraft]); + return ( {!isLoading && showCopyButton && (
+ {canEdit && isHuman && onEditAndRegenerate && !isEditing && ( + + + + )} {feedback !== undefined && runId && threadId && ( void; + onDraftChange: (value: string) => void; + onSubmit: () => void | Promise; + }; }) { const { t } = useI18n(); const isHuman = message.type === "human"; @@ -411,11 +497,57 @@ function MessageContent_({ /> )} {filesList} - {contentToDisplay && ( + {editState ? ( +
+