mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-28 00:48:07 +00:00
feat(chat): edit and rerun latest user turn (#4377)
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
Ryker_Feng
GitHub
parent
6f53fd5e99
commit
fcbf0609b0
@ -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.
|
||||
|
||||
```
|
||||
|
||||
@ -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 (`<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)`) |
|
||||
| **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
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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)}
|
||||
|
||||
@ -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]:
|
||||
|
||||
@ -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},
|
||||
)
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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="<uploaded_files>injected</uploaded_files>\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
|
||||
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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."""
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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<boolean>;
|
||||
}) {
|
||||
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 (
|
||||
<AIElementMessage
|
||||
className={cn("group/conversation-message relative w-full", className)}
|
||||
@ -162,6 +213,18 @@ export function MessageListItem({
|
||||
threadId={threadId}
|
||||
artifactPaths={artifactPaths}
|
||||
runId={runId}
|
||||
editState={
|
||||
isHuman && isEditing
|
||||
? {
|
||||
draft,
|
||||
disabled: isEditPending || isSubmittingEdit,
|
||||
submitDisabled: editSubmitDisabled,
|
||||
onCancel: cancelEditing,
|
||||
onDraftChange: setDraft,
|
||||
onSubmit: submitEdit,
|
||||
}
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
{!isLoading && showCopyButton && (
|
||||
<MessageToolbar
|
||||
@ -174,6 +237,20 @@ export function MessageListItem({
|
||||
>
|
||||
<div className="pointer-events-auto flex gap-1">
|
||||
<CopyButton clipboardData={getMessageCopyData(message)} />
|
||||
{canEdit && isHuman && onEditAndRegenerate && !isEditing && (
|
||||
<Tooltip content={t.common.editAndRerun}>
|
||||
<Button
|
||||
aria-label={t.common.editAndRerun}
|
||||
size="icon-sm"
|
||||
type="button"
|
||||
variant="ghost"
|
||||
disabled={isEditPending || isSubmittingEdit}
|
||||
onClick={startEditing}
|
||||
>
|
||||
<PencilIcon className="size-3" />
|
||||
</Button>
|
||||
</Tooltip>
|
||||
)}
|
||||
{feedback !== undefined && runId && threadId && (
|
||||
<FeedbackButtons
|
||||
threadId={threadId}
|
||||
@ -285,6 +362,7 @@ function MessageContent_({
|
||||
threadId,
|
||||
artifactPaths,
|
||||
runId,
|
||||
editState,
|
||||
}: {
|
||||
className?: string;
|
||||
message: Message;
|
||||
@ -292,6 +370,14 @@ function MessageContent_({
|
||||
threadId: string;
|
||||
artifactPaths: readonly string[];
|
||||
runId?: string;
|
||||
editState?: {
|
||||
draft: string;
|
||||
disabled: boolean;
|
||||
submitDisabled: boolean;
|
||||
onCancel: () => void;
|
||||
onDraftChange: (value: string) => void;
|
||||
onSubmit: () => void | Promise<void>;
|
||||
};
|
||||
}) {
|
||||
const { t } = useI18n();
|
||||
const isHuman = message.type === "human";
|
||||
@ -411,11 +497,57 @@ function MessageContent_({
|
||||
/>
|
||||
)}
|
||||
{filesList}
|
||||
{contentToDisplay && (
|
||||
{editState ? (
|
||||
<div className="bg-background border-border flex w-full min-w-0 flex-col gap-2 rounded-lg border p-2 shadow-sm">
|
||||
<Textarea
|
||||
autoFocus
|
||||
className="min-h-24 resize-y"
|
||||
disabled={editState.disabled}
|
||||
value={editState.draft}
|
||||
onChange={(event) =>
|
||||
editState.onDraftChange(event.currentTarget.value)
|
||||
}
|
||||
onKeyDown={(event) => {
|
||||
if (event.key === "Escape") {
|
||||
event.preventDefault();
|
||||
editState.onCancel();
|
||||
}
|
||||
if (event.key === "Enter" && (event.metaKey || event.ctrlKey)) {
|
||||
event.preventDefault();
|
||||
void editState.onSubmit();
|
||||
}
|
||||
}}
|
||||
/>
|
||||
<div className="text-muted-foreground text-xs">
|
||||
{t.common.editRerunWarning}
|
||||
</div>
|
||||
<div className="flex justify-end gap-1">
|
||||
<Button
|
||||
size="sm"
|
||||
type="button"
|
||||
variant="ghost"
|
||||
disabled={editState.disabled}
|
||||
onClick={editState.onCancel}
|
||||
>
|
||||
<XIcon className="size-3" />
|
||||
{t.common.cancel}
|
||||
</Button>
|
||||
<Button
|
||||
size="sm"
|
||||
type="button"
|
||||
disabled={editState.submitDisabled}
|
||||
onClick={() => void editState.onSubmit()}
|
||||
>
|
||||
<CheckIcon className="size-3" />
|
||||
{t.common.updateAndRerun}
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
) : contentToDisplay ? (
|
||||
<AIElementMessageContent className="w-full max-w-full">
|
||||
<HumanMessageText content={contentToDisplay} />
|
||||
</AIElementMessageContent>
|
||||
)}
|
||||
) : null}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
@ -50,6 +50,7 @@ import {
|
||||
getAssistantTurnCopyData,
|
||||
getAssistantTurnUsageMessages,
|
||||
getBranchableAssistantGroupIds,
|
||||
getLatestEditableTurn,
|
||||
getMessageGroups,
|
||||
getStreamingMessageLookup,
|
||||
hasContent,
|
||||
@ -332,9 +333,11 @@ export function MessageList({
|
||||
loadMoreHistory,
|
||||
isHistoryLoading,
|
||||
onRegenerateMessage,
|
||||
onEditAndRegenerateMessage,
|
||||
onSubmitHumanInput,
|
||||
onBranchTurn,
|
||||
canRegenerate = false,
|
||||
canEdit = false,
|
||||
canBranch = false,
|
||||
enableSidecarActions = true,
|
||||
sidecarSurface = false,
|
||||
@ -353,7 +356,11 @@ export function MessageList({
|
||||
onRegenerateMessage?: (
|
||||
messageId: string,
|
||||
supersededMessageIds: string[],
|
||||
) => void | Promise<void>;
|
||||
) => boolean | void | Promise<boolean | void>;
|
||||
onEditAndRegenerateMessage?: (
|
||||
messageId: string,
|
||||
replacementText: string,
|
||||
) => boolean | Promise<boolean>;
|
||||
onSubmitHumanInput?: (
|
||||
request: HumanInputRequest,
|
||||
response: HumanInputResponse,
|
||||
@ -363,6 +370,7 @@ export function MessageList({
|
||||
messageIds: string[],
|
||||
) => void | Promise<void>;
|
||||
canRegenerate?: boolean;
|
||||
canEdit?: boolean;
|
||||
canBranch?: boolean;
|
||||
enableSidecarActions?: boolean;
|
||||
sidecarSurface?: boolean;
|
||||
@ -463,6 +471,7 @@ export function MessageList({
|
||||
const [branchingMessageId, setBranchingMessageId] = useState<string | null>(
|
||||
null,
|
||||
);
|
||||
const [editingMessageId, setEditingMessageId] = useState<string | null>(null);
|
||||
const hasActiveAssistantText = useMemo(() => {
|
||||
let lastHumanIndex = -1;
|
||||
for (let i = groupedMessages.length - 1; i >= 0; i--) {
|
||||
@ -649,6 +658,15 @@ export function MessageList({
|
||||
() => getBranchableAssistantGroupIds(groupedMessages, thread.isLoading),
|
||||
[groupedMessages, thread.isLoading],
|
||||
);
|
||||
const latestEditableTurn = useMemo(
|
||||
() => getLatestEditableTurn(groupedMessages, thread.isLoading),
|
||||
[groupedMessages, thread.isLoading],
|
||||
);
|
||||
const latestEditableHumanMessageId = latestEditableTurn?.humanMessage.id;
|
||||
const replayActionBusy =
|
||||
regeneratingMessageId != null ||
|
||||
branchingMessageId != null ||
|
||||
editingMessageId != null;
|
||||
|
||||
const clearSelectionToolbar = useCallback(() => {
|
||||
setSelectionToolbar(null);
|
||||
@ -818,7 +836,9 @@ export function MessageList({
|
||||
type="button"
|
||||
variant="ghost"
|
||||
disabled={
|
||||
!canBranch || branchingMessageId === actionTarget.id
|
||||
!canBranch ||
|
||||
replayActionBusy ||
|
||||
branchingMessageId === actionTarget.id
|
||||
}
|
||||
onClick={() => {
|
||||
const targetId = actionTarget.id;
|
||||
@ -852,7 +872,9 @@ export function MessageList({
|
||||
type="button"
|
||||
variant="ghost"
|
||||
disabled={
|
||||
!canRegenerate || regeneratingMessageId === actionTarget.id
|
||||
!canRegenerate ||
|
||||
replayActionBusy ||
|
||||
regeneratingMessageId === actionTarget.id
|
||||
}
|
||||
onClick={() => {
|
||||
const targetId = actionTarget.id;
|
||||
@ -887,6 +909,7 @@ export function MessageList({
|
||||
onBranchTurn,
|
||||
onRegenerateMessage,
|
||||
regeneratingMessageId,
|
||||
replayActionBusy,
|
||||
t.common.branch,
|
||||
t.common.regenerate,
|
||||
],
|
||||
@ -1040,6 +1063,36 @@ export function MessageList({
|
||||
: undefined
|
||||
}
|
||||
showCopyButton={group.type !== "assistant"}
|
||||
canEdit={
|
||||
group.type === "human" &&
|
||||
Boolean(msg.id) &&
|
||||
msg.id === latestEditableHumanMessageId &&
|
||||
canEdit &&
|
||||
!replayActionBusy &&
|
||||
Boolean(onEditAndRegenerateMessage)
|
||||
}
|
||||
isEditPending={editingMessageId === msg.id}
|
||||
onEditAndRegenerate={
|
||||
group.type === "human" &&
|
||||
msg.id &&
|
||||
onEditAndRegenerateMessage
|
||||
? async (replacementText) => {
|
||||
const targetId = msg.id;
|
||||
if (!targetId) {
|
||||
return false;
|
||||
}
|
||||
setEditingMessageId(targetId);
|
||||
try {
|
||||
return await onEditAndRegenerateMessage(
|
||||
targetId,
|
||||
replacementText,
|
||||
);
|
||||
} finally {
|
||||
setEditingMessageId(null);
|
||||
}
|
||||
}
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
);
|
||||
|
||||
|
||||
@ -52,6 +52,10 @@ export const enUS: Translations = {
|
||||
exportAsJSON: "Export as JSON",
|
||||
exportSuccess: "Conversation exported",
|
||||
regenerate: "Regenerate",
|
||||
editAndRerun: "Edit and rerun",
|
||||
updateAndRerun: "Update and rerun",
|
||||
editRerunWarning:
|
||||
"Rerunning restores conversation state only. Files, memory, and external actions are not undone.",
|
||||
branch: "Branch conversation",
|
||||
showArtifacts: "Show artifacts of this conversation",
|
||||
browser: "Browser",
|
||||
|
||||
@ -41,6 +41,9 @@ export interface Translations {
|
||||
exportAsJSON: string;
|
||||
exportSuccess: string;
|
||||
regenerate: string;
|
||||
editAndRerun: string;
|
||||
updateAndRerun: string;
|
||||
editRerunWarning: string;
|
||||
branch: string;
|
||||
showArtifacts: string;
|
||||
browser: string;
|
||||
|
||||
@ -52,6 +52,10 @@ export const zhCN: Translations = {
|
||||
exportAsJSON: "导出为 JSON",
|
||||
exportSuccess: "对话已导出",
|
||||
regenerate: "重新生成",
|
||||
editAndRerun: "编辑并重新运行",
|
||||
updateAndRerun: "更新并重新运行",
|
||||
editRerunWarning:
|
||||
"重新运行只恢复对话状态，不会撤销文件、Memory 或外部操作。",
|
||||
branch: "分叉",
|
||||
showArtifacts: "查看此对话的文件",
|
||||
browser: "浏览器",
|
||||
|
||||
@ -205,6 +205,89 @@ export function getBranchableAssistantGroupIds(
|
||||
return branchableGroupIds;
|
||||
}
|
||||
|
||||
export type EditableTurn = {
|
||||
humanMessage: Message;
|
||||
};
|
||||
|
||||
function isTerminalAssistantTextMessage(message: Message | undefined): boolean {
|
||||
return (
|
||||
message?.type === "ai" &&
|
||||
Boolean(extractTextFromMessage(message).trim()) &&
|
||||
!hasToolCalls(message)
|
||||
);
|
||||
}
|
||||
|
||||
export function getLatestEditableTurn(
|
||||
groups: MessageGroup[],
|
||||
isCurrentTurnLoading: boolean,
|
||||
): EditableTurn | null {
|
||||
if (isCurrentTurnLoading) {
|
||||
return null;
|
||||
}
|
||||
|
||||
let candidate: EditableTurn | null = null;
|
||||
let currentHumanGroup: MessageGroup | null = null;
|
||||
let currentTurnGroups: MessageGroup[] = [];
|
||||
let lastAIGroup: MessageGroup | null = null;
|
||||
|
||||
const completeTurn = () => {
|
||||
if (!currentHumanGroup) {
|
||||
currentTurnGroups = [];
|
||||
lastAIGroup = null;
|
||||
return;
|
||||
}
|
||||
|
||||
const humanMessage = currentHumanGroup?.messages.find(
|
||||
(message) => message.type === "human" && message.id,
|
||||
);
|
||||
let assistantMessage: Message | undefined;
|
||||
for (let i = (lastAIGroup?.messages.length ?? 0) - 1; i >= 0; i -= 1) {
|
||||
const message = lastAIGroup?.messages[i];
|
||||
if (message?.type === "ai" && message.id) {
|
||||
assistantMessage = message;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (
|
||||
currentHumanGroup &&
|
||||
lastAIGroup?.type === "assistant" &&
|
||||
humanMessage &&
|
||||
isTerminalAssistantTextMessage(assistantMessage)
|
||||
) {
|
||||
candidate = {
|
||||
humanMessage,
|
||||
};
|
||||
} else {
|
||||
candidate = null;
|
||||
}
|
||||
|
||||
currentHumanGroup = null;
|
||||
currentTurnGroups = [];
|
||||
lastAIGroup = null;
|
||||
};
|
||||
|
||||
for (const group of groups) {
|
||||
if (group.type === "human") {
|
||||
completeTurn();
|
||||
currentHumanGroup = group;
|
||||
currentTurnGroups = [group];
|
||||
continue;
|
||||
}
|
||||
|
||||
if (currentHumanGroup) {
|
||||
currentTurnGroups.push(group);
|
||||
}
|
||||
|
||||
if (group.messages.some((message) => message.type === "ai")) {
|
||||
lastAIGroup = group;
|
||||
}
|
||||
}
|
||||
|
||||
completeTurn();
|
||||
return candidate;
|
||||
}
|
||||
|
||||
export function groupMessages<T>(
|
||||
messages: Message[],
|
||||
mapper: (group: MessageGroup) => T,
|
||||
|
||||
@ -96,6 +96,18 @@ type RegeneratePrepareResponse = {
|
||||
target_run_id: string;
|
||||
};
|
||||
|
||||
type EditRegeneratePrepareResponse = RegeneratePrepareResponse & {
|
||||
replacement_human_message_id: string;
|
||||
source_message_ids: string[];
|
||||
};
|
||||
|
||||
export type PendingPreparedReplayMask = {
|
||||
kind: "regenerate" | "edit";
|
||||
targetRunId: string;
|
||||
supersededMessageIds: string[];
|
||||
replacementHumanMessageId?: string;
|
||||
};
|
||||
|
||||
export function hasToolResult(messages: Message[], toolName: string): boolean {
|
||||
const matchingToolCallIds = new Set<string>();
|
||||
for (const message of messages) {
|
||||
@ -742,6 +754,24 @@ export function getVisibleOptimisticMessages(
|
||||
return optimisticMessages;
|
||||
}
|
||||
|
||||
export function areOptimisticMessagesConfirmed(
|
||||
optimisticMessages: Message[],
|
||||
persistedMessages: Message[],
|
||||
): boolean {
|
||||
const optimisticIdentities = optimisticMessages
|
||||
.map(messageIdentity)
|
||||
.filter(isNonEmptyString);
|
||||
if (optimisticIdentities.length === 0) {
|
||||
return false;
|
||||
}
|
||||
const persistedIdentities = new Set(
|
||||
persistedMessages.map(messageIdentity).filter(isNonEmptyString),
|
||||
);
|
||||
return optimisticIdentities.every((identity) =>
|
||||
persistedIdentities.has(identity),
|
||||
);
|
||||
}
|
||||
|
||||
export function getSummarizationMiddlewareMessages(
|
||||
data: unknown,
|
||||
): Message[] | undefined {
|
||||
@ -1143,6 +1173,9 @@ export function useThreadStream({
|
||||
const threadIdRef = useRef<string | null>(threadId ?? null);
|
||||
const startedRef = useRef(false);
|
||||
const pendingUsageBaselineMessageIdsRef = useRef<Set<string>>(new Set());
|
||||
const pendingPreparedReplayRef = useRef<PendingPreparedReplayMask | null>(
|
||||
null,
|
||||
);
|
||||
const listeners = useRef({
|
||||
onSend,
|
||||
onStart,
|
||||
@ -1211,6 +1244,24 @@ export function useThreadStream({
|
||||
const { tasksRef, setTasks } = useSubtaskContext();
|
||||
const updateSubtask = useUpdateSubtask();
|
||||
|
||||
const clearPreparedReplayMasks = useCallback(
|
||||
(replay: PendingPreparedReplayMask | null) => {
|
||||
if (!replay) {
|
||||
return;
|
||||
}
|
||||
setPendingSupersededRunIds((current) =>
|
||||
removeSetItems(current, [replay.targetRunId]),
|
||||
);
|
||||
setPendingSupersededMessageIds((current) =>
|
||||
removeSetItems(current, replay.supersededMessageIds),
|
||||
);
|
||||
if (pendingPreparedReplayRef.current === replay) {
|
||||
pendingPreparedReplayRef.current = null;
|
||||
}
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
const thread = useStream<AgentThreadState>({
|
||||
client: getAPIClient(isMock),
|
||||
assistantId: "lead_agent",
|
||||
@ -1397,6 +1448,7 @@ export function useThreadStream({
|
||||
setOptimisticMessages([]);
|
||||
setOptimisticThreadId(null);
|
||||
setLiveMessagesThreadId(null);
|
||||
pendingPreparedReplayRef.current = null;
|
||||
setPendingSupersededRunIds(new Set());
|
||||
setPendingSupersededMessageIds(new Set());
|
||||
toast.error(getStreamErrorMessage(error));
|
||||
@ -1416,6 +1468,7 @@ export function useThreadStream({
|
||||
},
|
||||
onFinish(state) {
|
||||
listeners.current.onFinish?.(state.values);
|
||||
pendingPreparedReplayRef.current = null;
|
||||
pendingUsageBaselineMessageIdsRef.current = new Set(
|
||||
messagesRef.current
|
||||
.map(messageIdentity)
|
||||
@ -1428,13 +1481,26 @@ export function useThreadStream({
|
||||
const stopThread = useCallback(async () => {
|
||||
const stoppedThreadId =
|
||||
threadIdRef.current ?? displayThreadId ?? threadId ?? null;
|
||||
const pendingReplay = pendingPreparedReplayRef.current;
|
||||
await stopThreadAndInvalidateCaches(
|
||||
queryClient,
|
||||
() => thread.stop(),
|
||||
stoppedThreadId,
|
||||
isMock,
|
||||
);
|
||||
}, [displayThreadId, isMock, queryClient, thread, threadId]);
|
||||
if (pendingReplay) {
|
||||
setOptimisticMessages([]);
|
||||
setOptimisticThreadId(null);
|
||||
clearPreparedReplayMasks(pendingReplay);
|
||||
}
|
||||
}, [
|
||||
clearPreparedReplayMasks,
|
||||
displayThreadId,
|
||||
isMock,
|
||||
queryClient,
|
||||
thread,
|
||||
threadId,
|
||||
]);
|
||||
|
||||
const hasVisibleStreamState =
|
||||
Boolean(threadId) || liveMessagesThreadId === currentViewThreadId;
|
||||
@ -1490,6 +1556,7 @@ export function useThreadStream({
|
||||
transientHistoryThreadIdRef.current = null;
|
||||
summarizedRef.current = new Set<string>();
|
||||
pendingUsageBaselineMessageIdsRef.current = new Set();
|
||||
pendingPreparedReplayRef.current = null;
|
||||
setPendingSupersededRunIds(new Set());
|
||||
setPendingSupersededMessageIds(new Set());
|
||||
prevHumanMsgCountRef.current =
|
||||
@ -1554,6 +1621,16 @@ export function useThreadStream({
|
||||
}
|
||||
}, [hasHumanOptimistic, humanMessageCount, optimisticMessageCount]);
|
||||
|
||||
useEffect(() => {
|
||||
if (
|
||||
optimisticMessageCount > 0 &&
|
||||
areOptimisticMessagesConfirmed(optimisticMessages, persistedMessages)
|
||||
) {
|
||||
setOptimisticMessages([]);
|
||||
setOptimisticThreadId(null);
|
||||
}
|
||||
}, [optimisticMessageCount, optimisticMessages, persistedMessages]);
|
||||
|
||||
const sendMessage = useCallback(
|
||||
async (
|
||||
threadId: string,
|
||||
@ -1761,14 +1838,20 @@ export function useThreadStream({
|
||||
],
|
||||
);
|
||||
|
||||
const regenerateMessage = useCallback(
|
||||
async (
|
||||
threadId: string,
|
||||
messageId: string,
|
||||
supersededMessageIds: string[] = [messageId],
|
||||
) => {
|
||||
if (sendInFlightRef.current || !threadId || !messageId) {
|
||||
return;
|
||||
const submitPreparedReplay = useCallback(
|
||||
async <TPrepared extends RegeneratePrepareResponse>({
|
||||
threadId,
|
||||
prepare,
|
||||
getSupersededMessageIds,
|
||||
getOptimisticMessages,
|
||||
}: {
|
||||
threadId: string;
|
||||
prepare: () => Promise<TPrepared>;
|
||||
getSupersededMessageIds: (prepared: TPrepared) => string[];
|
||||
getOptimisticMessages?: (prepared: TPrepared) => Message[];
|
||||
}) => {
|
||||
if (sendInFlightRef.current || !threadId) {
|
||||
return false;
|
||||
}
|
||||
sendInFlightRef.current = true;
|
||||
prevHumanMsgCountRef.current = humanMessageCount;
|
||||
@ -1783,25 +1866,21 @@ export function useThreadStream({
|
||||
let preparedSupersededMessageIds: string[] = [];
|
||||
|
||||
try {
|
||||
const response = await fetch(
|
||||
`${getBackendBaseURL()}/api/threads/${encodeURIComponent(
|
||||
threadId,
|
||||
)}/runs/regenerate/prepare`,
|
||||
{
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
credentials: "include",
|
||||
body: JSON.stringify({ message_id: messageId }),
|
||||
},
|
||||
);
|
||||
if (!response.ok) {
|
||||
throw new Error(await readResponseErrorMessage(response));
|
||||
}
|
||||
const prepared = (await response.json()) as RegeneratePrepareResponse;
|
||||
const prepared = await prepare();
|
||||
preparedSupersededRunId = prepared.target_run_id;
|
||||
preparedSupersededMessageIds = supersededMessageIds;
|
||||
preparedSupersededMessageIds = getSupersededMessageIds(prepared);
|
||||
const replacementHumanMessageId =
|
||||
"replacement_human_message_id" in prepared &&
|
||||
typeof prepared.replacement_human_message_id === "string"
|
||||
? prepared.replacement_human_message_id
|
||||
: undefined;
|
||||
const pendingReplay: PendingPreparedReplayMask = {
|
||||
kind: replacementHumanMessageId ? "edit" : "regenerate",
|
||||
targetRunId: prepared.target_run_id,
|
||||
supersededMessageIds: preparedSupersededMessageIds,
|
||||
replacementHumanMessageId,
|
||||
};
|
||||
pendingPreparedReplayRef.current = pendingReplay;
|
||||
setPendingSupersededRunIds((current) => {
|
||||
const next = new Set(current);
|
||||
next.add(prepared.target_run_id);
|
||||
@ -1809,12 +1888,18 @@ export function useThreadStream({
|
||||
});
|
||||
setPendingSupersededMessageIds((current) => {
|
||||
const next = new Set(current);
|
||||
for (const id of supersededMessageIds) {
|
||||
for (const id of preparedSupersededMessageIds) {
|
||||
next.add(id);
|
||||
}
|
||||
return next;
|
||||
});
|
||||
|
||||
const nextOptimisticMessages = getOptimisticMessages?.(prepared) ?? [];
|
||||
if (nextOptimisticMessages.length > 0) {
|
||||
setOptimisticThreadId(threadId);
|
||||
setOptimisticMessages(nextOptimisticMessages);
|
||||
}
|
||||
|
||||
await thread.submit(prepared.input, {
|
||||
threadId,
|
||||
checkpoint: prepared.checkpoint,
|
||||
@ -1846,13 +1931,20 @@ export function useThreadStream({
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: INFINITE_THREADS_QUERY_KEY_PREFIX,
|
||||
});
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: threadHistoryQueryKey(threadId),
|
||||
});
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: threadTokenUsageQueryKey(threadId),
|
||||
});
|
||||
return true;
|
||||
} catch (error) {
|
||||
setOptimisticMessages([]);
|
||||
setOptimisticThreadId(null);
|
||||
setLiveMessagesThreadId(null);
|
||||
if (preparedSupersededRunId) {
|
||||
const supersededRunId = preparedSupersededRunId;
|
||||
pendingPreparedReplayRef.current = null;
|
||||
setPendingSupersededRunIds((current) =>
|
||||
removeSetItems(current, [supersededRunId]),
|
||||
);
|
||||
@ -1861,6 +1953,7 @@ export function useThreadStream({
|
||||
);
|
||||
}
|
||||
toast.error(getStreamErrorMessage(error));
|
||||
return false;
|
||||
} finally {
|
||||
sendInFlightRef.current = false;
|
||||
}
|
||||
@ -1868,6 +1961,82 @@ export function useThreadStream({
|
||||
[context, humanMessageCount, persistedMessages, queryClient, thread],
|
||||
);
|
||||
|
||||
const regenerateMessage = useCallback(
|
||||
async (
|
||||
threadId: string,
|
||||
messageId: string,
|
||||
supersededMessageIds: string[] = [messageId],
|
||||
) => {
|
||||
if (!messageId) {
|
||||
return false;
|
||||
}
|
||||
return submitPreparedReplay({
|
||||
threadId,
|
||||
prepare: async () => {
|
||||
const response = await fetch(
|
||||
`${getBackendBaseURL()}/api/threads/${encodeURIComponent(
|
||||
threadId,
|
||||
)}/runs/regenerate/prepare`,
|
||||
{
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
credentials: "include",
|
||||
body: JSON.stringify({ message_id: messageId }),
|
||||
},
|
||||
);
|
||||
if (!response.ok) {
|
||||
throw new Error(await readResponseErrorMessage(response));
|
||||
}
|
||||
return (await response.json()) as RegeneratePrepareResponse;
|
||||
},
|
||||
getSupersededMessageIds: () => supersededMessageIds,
|
||||
});
|
||||
},
|
||||
[submitPreparedReplay],
|
||||
);
|
||||
|
||||
const editAndRegenerateMessage = useCallback(
|
||||
async (
|
||||
threadId: string,
|
||||
humanMessageId: string,
|
||||
replacementText: string,
|
||||
) => {
|
||||
if (!humanMessageId) {
|
||||
return false;
|
||||
}
|
||||
return submitPreparedReplay<EditRegeneratePrepareResponse>({
|
||||
threadId,
|
||||
prepare: async () => {
|
||||
const response = await fetch(
|
||||
`${getBackendBaseURL()}/api/threads/${encodeURIComponent(
|
||||
threadId,
|
||||
)}/runs/edit-regenerate/prepare`,
|
||||
{
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
credentials: "include",
|
||||
body: JSON.stringify({
|
||||
human_message_id: humanMessageId,
|
||||
replacement_text: replacementText,
|
||||
}),
|
||||
},
|
||||
);
|
||||
if (!response.ok) {
|
||||
throw new Error(await readResponseErrorMessage(response));
|
||||
}
|
||||
return (await response.json()) as EditRegeneratePrepareResponse;
|
||||
},
|
||||
getSupersededMessageIds: (prepared) => prepared.source_message_ids,
|
||||
getOptimisticMessages: (prepared) => prepared.input.messages ?? [],
|
||||
});
|
||||
},
|
||||
[submitPreparedReplay],
|
||||
);
|
||||
|
||||
// Cache the latest thread messages in a ref to compare against incoming history messages for deduplication,
|
||||
// and to allow access to the full message list in onUpdateEvent without causing re-renders.
|
||||
if (persistedMessages.length >= messagesRef.current.length) {
|
||||
@ -1961,6 +2130,7 @@ export function useThreadStream({
|
||||
pendingUsageMessages,
|
||||
sendMessage,
|
||||
regenerateMessage,
|
||||
editAndRegenerateMessage,
|
||||
isUploading,
|
||||
isHistoryLoading,
|
||||
hasMoreHistory,
|
||||
|
||||
@ -3,6 +3,7 @@ import { expect, test } from "@playwright/test";
|
||||
import {
|
||||
handleRunStream,
|
||||
mockLangGraphAPI,
|
||||
MOCK_RUN_ID,
|
||||
MOCK_THREAD_ID,
|
||||
} from "./utils/mock-api";
|
||||
|
||||
@ -184,4 +185,160 @@ test.describe("Agent chat", () => {
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
test("agent chat can edit and rerun its latest user message", async ({
|
||||
page,
|
||||
}) => {
|
||||
const humanMessage = {
|
||||
type: "human",
|
||||
id: "msg-human-agent",
|
||||
content: [{ type: "text", text: "Original agent question" }],
|
||||
};
|
||||
const replacementHumanMessage = {
|
||||
type: "human",
|
||||
id: "msg-human-agent-edited",
|
||||
content: [{ type: "text", text: "Edited agent question" }],
|
||||
};
|
||||
const aiMessage = {
|
||||
type: "ai",
|
||||
id: "msg-ai-agent",
|
||||
content: "Custom agent response",
|
||||
};
|
||||
mockLangGraphAPI(page, {
|
||||
agents: MOCK_AGENTS,
|
||||
threads: [
|
||||
{
|
||||
thread_id: MOCK_THREAD_ID,
|
||||
title: "Agent conversation",
|
||||
agent_name: "test-agent",
|
||||
messages: [humanMessage, aiMessage],
|
||||
},
|
||||
],
|
||||
});
|
||||
let historyRows = [
|
||||
{ run_id: `run-${MOCK_THREAD_ID}`, content: humanMessage },
|
||||
{ run_id: `run-${MOCK_THREAD_ID}`, content: aiMessage },
|
||||
];
|
||||
await page.route(
|
||||
`**/api/threads/${MOCK_THREAD_ID}/messages/page`,
|
||||
(route) =>
|
||||
route.fulfill({
|
||||
status: 200,
|
||||
contentType: "application/json",
|
||||
body: JSON.stringify({
|
||||
data: historyRows.map((row, index) => ({
|
||||
run_id: row.run_id,
|
||||
seq: index + 1,
|
||||
content: row.content,
|
||||
metadata: { caller: "lead_agent" },
|
||||
created_at: `2025-01-01T00:00:${String(index).padStart(2, "0")}Z`,
|
||||
})),
|
||||
has_more: false,
|
||||
next_before_seq: null,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
let prepareBody:
|
||||
| { human_message_id?: string; replacement_text?: string }
|
||||
| undefined;
|
||||
let streamBody: Record<string, unknown> | undefined;
|
||||
await page.route(
|
||||
`**/api/threads/${MOCK_THREAD_ID}/runs/edit-regenerate/prepare`,
|
||||
(route) => {
|
||||
prepareBody = route.request().postDataJSON() as {
|
||||
human_message_id?: string;
|
||||
replacement_text?: string;
|
||||
};
|
||||
return route.fulfill({
|
||||
status: 200,
|
||||
contentType: "application/json",
|
||||
body: JSON.stringify({
|
||||
input: { messages: [replacementHumanMessage] },
|
||||
checkpoint: {
|
||||
checkpoint_id: "checkpoint-before-human",
|
||||
checkpoint_ns: "",
|
||||
checkpoint_map: null,
|
||||
},
|
||||
metadata: {
|
||||
replay_kind: "edit",
|
||||
regenerate_from_message_id: aiMessage.id,
|
||||
regenerate_from_run_id: `run-${MOCK_THREAD_ID}`,
|
||||
regenerate_checkpoint_id: "checkpoint-before-human",
|
||||
edit_from_message_id: humanMessage.id,
|
||||
edit_message_id: replacementHumanMessage.id,
|
||||
edit_version_group_id: humanMessage.id,
|
||||
},
|
||||
target_run_id: `run-${MOCK_THREAD_ID}`,
|
||||
replacement_human_message_id: replacementHumanMessage.id,
|
||||
source_message_ids: [humanMessage.id, aiMessage.id],
|
||||
}),
|
||||
});
|
||||
},
|
||||
);
|
||||
await page.route(
|
||||
`**/api/langgraph/threads/${MOCK_THREAD_ID}/runs/stream`,
|
||||
(route) => {
|
||||
streamBody = route.request().postDataJSON() as Record<string, unknown>;
|
||||
historyRows = [
|
||||
{ run_id: MOCK_RUN_ID, content: replacementHumanMessage },
|
||||
{
|
||||
run_id: MOCK_RUN_ID,
|
||||
content: {
|
||||
type: "ai",
|
||||
id: "msg-ai-1",
|
||||
content: "Hello from DeerFlow!",
|
||||
},
|
||||
},
|
||||
];
|
||||
return handleRunStream(route);
|
||||
},
|
||||
);
|
||||
|
||||
await page.goto(`/workspace/agents/test-agent/chats/${MOCK_THREAD_ID}`);
|
||||
await expect(page.getByText("Original agent question")).toBeVisible({
|
||||
timeout: 15_000,
|
||||
});
|
||||
|
||||
const humanTurn = page.getByText("Original agent question");
|
||||
await humanTurn.hover();
|
||||
await page.getByRole("button", { name: "Edit and rerun" }).click();
|
||||
|
||||
const editor = page.locator("textarea").first();
|
||||
await expect(editor).toHaveValue("Original agent question");
|
||||
await editor.fill("Edited agent question");
|
||||
await page.getByRole("button", { name: "Update and rerun" }).click();
|
||||
|
||||
await expect
|
||||
.poll(() => prepareBody)
|
||||
.toEqual({
|
||||
human_message_id: humanMessage.id,
|
||||
replacement_text: "Edited agent question",
|
||||
});
|
||||
await expect.poll(() => streamBody).toBeDefined();
|
||||
expect(streamBody).toMatchObject({
|
||||
input: { messages: [replacementHumanMessage] },
|
||||
checkpoint: {
|
||||
checkpoint_id: "checkpoint-before-human",
|
||||
checkpoint_ns: "",
|
||||
checkpoint_map: null,
|
||||
},
|
||||
metadata: {
|
||||
replay_kind: "edit",
|
||||
regenerate_from_message_id: aiMessage.id,
|
||||
regenerate_from_run_id: `run-${MOCK_THREAD_ID}`,
|
||||
regenerate_checkpoint_id: "checkpoint-before-human",
|
||||
edit_from_message_id: humanMessage.id,
|
||||
edit_message_id: replacementHumanMessage.id,
|
||||
edit_version_group_id: humanMessage.id,
|
||||
},
|
||||
context: {
|
||||
agent_name: "test-agent",
|
||||
thread_id: MOCK_THREAD_ID,
|
||||
},
|
||||
});
|
||||
await expect(page.getByText("Edited agent question")).toBeVisible();
|
||||
await expect(page.getByText("Original agent question")).not.toBeVisible();
|
||||
await expect(page.getByText("Hello from DeerFlow!")).toBeVisible();
|
||||
});
|
||||
});
|
||||
|
||||
@ -6,6 +6,7 @@ import {
|
||||
extractTextFromMessage,
|
||||
extractReasoningContentFromMessage,
|
||||
getBranchableAssistantGroupIds,
|
||||
getLatestEditableTurn,
|
||||
getMessageCopyData,
|
||||
getAssistantTurnCopyData,
|
||||
getAssistantTurnUsageMessages,
|
||||
@ -134,6 +135,88 @@ describe("branchable assistant groups", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("latest editable user turn", () => {
|
||||
test("returns the latest completed human turn", () => {
|
||||
const messages = [
|
||||
{ id: "human-1", type: "human", content: "First question" },
|
||||
{ id: "ai-history", type: "ai", content: "Historical answer" },
|
||||
{ id: "human-2", type: "human", content: "Use tools" },
|
||||
{
|
||||
id: "ai-tool",
|
||||
type: "ai",
|
||||
content: "",
|
||||
tool_calls: [{ id: "tool-1", name: "web_search", args: {} }],
|
||||
},
|
||||
{
|
||||
id: "tool-result",
|
||||
type: "tool",
|
||||
name: "web_search",
|
||||
tool_call_id: "tool-1",
|
||||
content: "result",
|
||||
},
|
||||
{ id: "ai-final", type: "ai", content: "Final answer" },
|
||||
] as Message[];
|
||||
|
||||
const turn = getLatestEditableTurn(getMessageGroups(messages), false);
|
||||
|
||||
expect(turn?.humanMessage.id).toBe("human-2");
|
||||
});
|
||||
|
||||
test("returns null while the current turn is loading", () => {
|
||||
const messages = [
|
||||
{ id: "human-1", type: "human", content: "Question" },
|
||||
{ id: "ai-1", type: "ai", content: "Answer" },
|
||||
] as Message[];
|
||||
|
||||
expect(getLatestEditableTurn(getMessageGroups(messages), true)).toBeNull();
|
||||
});
|
||||
|
||||
test("returns null when the latest turn has no final assistant text", () => {
|
||||
const messages = [
|
||||
{ id: "human-1", type: "human", content: "Question" },
|
||||
{
|
||||
id: "ai-tool",
|
||||
type: "ai",
|
||||
content: "",
|
||||
tool_calls: [{ id: "tool-1", name: "web_search", args: {} }],
|
||||
},
|
||||
{
|
||||
id: "tool-result",
|
||||
type: "tool",
|
||||
name: "web_search",
|
||||
tool_call_id: "tool-1",
|
||||
content: "result",
|
||||
},
|
||||
] as Message[];
|
||||
|
||||
expect(getLatestEditableTurn(getMessageGroups(messages), false)).toBeNull();
|
||||
});
|
||||
|
||||
test("returns null when the latest assistant text still has tool calls", () => {
|
||||
const messages = [
|
||||
{ id: "human-1", type: "human", content: "Question" },
|
||||
{
|
||||
id: "ai-tool",
|
||||
type: "ai",
|
||||
content: "Let me search first.",
|
||||
tool_calls: [{ id: "tool-1", name: "web_search", args: {} }],
|
||||
},
|
||||
] as Message[];
|
||||
|
||||
expect(getLatestEditableTurn(getMessageGroups(messages), false)).toBeNull();
|
||||
});
|
||||
|
||||
test("returns null when a later human turn is incomplete", () => {
|
||||
const messages = [
|
||||
{ id: "human-1", type: "human", content: "First question" },
|
||||
{ id: "ai-1", type: "ai", content: "First answer" },
|
||||
{ id: "human-2", type: "human", content: "Follow up" },
|
||||
] as Message[];
|
||||
|
||||
expect(getLatestEditableTurn(getMessageGroups(messages), false)).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
test("reasoning + content (no tool calls) yields a single assistant bubble, not a duplicate processing group", () => {
|
||||
// Regression for #3868: in thinking/pro/ultra modes the final assistant
|
||||
// message carries both reasoning_content and answer text. It must surface its
|
||||
|
||||
@ -5,6 +5,7 @@ import { InfiniteQueryObserver, QueryClient } from "@tanstack/react-query";
|
||||
import {
|
||||
buildThreadMessagesPageUrl,
|
||||
buildVisibleHistoryMessages,
|
||||
areOptimisticMessagesConfirmed,
|
||||
computeSummarizationTransientMessages,
|
||||
flattenThreadHistoryPages,
|
||||
getSummarizationMiddlewareMessages,
|
||||
@ -440,6 +441,37 @@ test("getVisibleOptimisticMessages hides optimistic user input after later serve
|
||||
]);
|
||||
});
|
||||
|
||||
test("areOptimisticMessagesConfirmed returns true when server messages contain every optimistic id", () => {
|
||||
const optimisticHuman = {
|
||||
id: "replacement-human-1",
|
||||
type: "human",
|
||||
content: "edited question",
|
||||
} as Message;
|
||||
const serverHuman = {
|
||||
id: "replacement-human-1",
|
||||
type: "human",
|
||||
content: "edited question",
|
||||
} as Message;
|
||||
const serverAi = {
|
||||
id: "replacement-ai-1",
|
||||
type: "ai",
|
||||
content: "new answer",
|
||||
} as Message;
|
||||
|
||||
expect(
|
||||
areOptimisticMessagesConfirmed([optimisticHuman], [serverHuman, serverAi]),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
test("areOptimisticMessagesConfirmed ignores optimistic messages without stable ids", () => {
|
||||
const optimisticHuman = {
|
||||
type: "human",
|
||||
content: "edited question",
|
||||
} as Message;
|
||||
|
||||
expect(areOptimisticMessagesConfirmed([optimisticHuman], [])).toBe(false);
|
||||
});
|
||||
|
||||
test("buildThreadMessagesPageUrl encodes the thread and backward cursor", () => {
|
||||
expect(
|
||||
buildThreadMessagesPageUrl(
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user