mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-09 16:08:31 +00:00
feat: support regenerating latest answer (#3637)
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
Huixin615
GitHub
parent
cce794a124
commit
ca9428d0cd
13 changed files with 1114 additions and 23 deletions
|
|
@ -273,7 +273,7 @@ CORS is same-origin by default when requests enter through nginx on port 2026. S
|
|||
| **Threads** (`/api/threads/{id}`) | `DELETE /` - remove DeerFlow-managed local thread data after LangGraph thread deletion; unexpected failures are logged server-side and return a generic 500 detail |
|
||||
| **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - serve artifacts; active content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) are always forced as download attachments to reduce XSS risk; `?download=true` still forces download for other file types |
|
||||
| **Suggestions** (`/api/suggestions`) | `GET /config` - returns global suggestions config boolean; `POST /threads/{id}/suggestions` - generate follow-up questions; rich list/block model content is normalized and inline reasoning (`<think>...</think>`, including unclosed/truncated blocks from reasoning models like MiniMax-M3) is stripped before JSON parsing |
|
||||
| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /../messages` - thread messages with feedback; `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; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /../messages` - thread messages with feedback; `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 |
|
||||
|
||||
|
|
@ -283,6 +283,7 @@ CORS is same-origin by default when requests enter through nginx on port 2026. S
|
|||
- `cancel()` and `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persist interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions.
|
||||
- Store-only hydrated runs are readable history. If the current worker has no in-memory task/control state for that run, cancellation APIs can return 409 because this worker cannot stop the task.
|
||||
- `POST /wait` (both thread-scoped and `/api/runs/wait`) drains the stream bridge via `wait_for_run_completion()` instead of bare `await record.task`, so it honours the run's `on_disconnect` setting and cancels the background run on real client disconnect rather than returning a stale checkpoint (issue #3265).
|
||||
- Thread-scoped run creation accepts `checkpoint` / `checkpoint_id`; Gateway validates the checkpoint belongs to the request thread before writing `checkpoint_id` / `checkpoint_ns` into `config.configurable` for LangGraph branching.
|
||||
|
||||
Proxied through nginx: `/api/langgraph/*` → Gateway LangGraph-compatible runtime, all other `/api/*` → Gateway REST APIs.
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from typing import Any, Literal
|
|||
|
||||
from fastapi import APIRouter, HTTPException, Query, Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
from langchain_core.messages import BaseMessage
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.authz import require_permission
|
||||
|
|
@ -24,9 +25,11 @@ from app.gateway.deps import get_checkpointer, get_current_user, get_feedback_re
|
|||
from app.gateway.pagination import trim_run_message_page
|
||||
from app.gateway.services import sse_consumer, start_run, wait_for_run_completion
|
||||
from deerflow.runtime import RunRecord, RunStatus, serialize_channel_values_for_api
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, get_original_user_content_text
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/api/threads", tags=["runs"])
|
||||
REGENERATE_HISTORY_SCAN_LIMIT = 200
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -57,6 +60,17 @@ class RunCreateRequest(BaseModel):
|
|||
feedback_keys: list[str] | None = Field(default=None, description="LangSmith feedback keys")
|
||||
|
||||
|
||||
class RegeneratePrepareRequest(BaseModel):
|
||||
message_id: str = Field(..., min_length=1, description="Assistant message id to regenerate")
|
||||
|
||||
|
||||
class RegeneratePrepareResponse(BaseModel):
|
||||
input: dict[str, Any]
|
||||
checkpoint: dict[str, Any]
|
||||
metadata: dict[str, Any]
|
||||
target_run_id: str
|
||||
|
||||
|
||||
class RunResponse(BaseModel):
|
||||
run_id: str
|
||||
thread_id: str
|
||||
|
|
@ -134,11 +148,273 @@ def _record_to_response(record: RunRecord) -> RunResponse:
|
|||
)
|
||||
|
||||
|
||||
def _message_id(message: Any) -> str | None:
|
||||
value = getattr(message, "id", None)
|
||||
if value is None and isinstance(message, dict):
|
||||
value = message.get("id")
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
def _message_type(message: Any) -> str | None:
|
||||
value = getattr(message, "type", None)
|
||||
if value is None and isinstance(message, dict):
|
||||
value = message.get("type") or message.get("role")
|
||||
if value == "assistant":
|
||||
return "ai"
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
def _message_name(message: Any) -> str | None:
|
||||
value = getattr(message, "name", None)
|
||||
if value is None and isinstance(message, dict):
|
||||
value = message.get("name")
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
def _message_content(message: Any) -> Any:
|
||||
if isinstance(message, dict):
|
||||
return message.get("content")
|
||||
return getattr(message, "content", None)
|
||||
|
||||
|
||||
def _message_text(message: Any) -> str:
|
||||
content = _message_content(message)
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for block in content:
|
||||
if isinstance(block, str):
|
||||
parts.append(block)
|
||||
elif isinstance(block, dict):
|
||||
text = block.get("text")
|
||||
if isinstance(text, str):
|
||||
parts.append(text)
|
||||
else:
|
||||
nested = block.get("content")
|
||||
if isinstance(nested, str):
|
||||
parts.append(nested)
|
||||
return "".join(parts)
|
||||
if isinstance(content, dict):
|
||||
for key in ("text", "content"):
|
||||
value = content.get(key)
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return ""
|
||||
|
||||
|
||||
def _message_additional_kwargs(message: Any) -> dict[str, Any]:
|
||||
value = getattr(message, "additional_kwargs", None)
|
||||
if value is None and isinstance(message, dict):
|
||||
value = message.get("additional_kwargs")
|
||||
return dict(value or {}) if isinstance(value, dict) else {}
|
||||
|
||||
|
||||
def _is_hidden_or_control_message(message: Any) -> bool:
|
||||
message_type = _message_type(message)
|
||||
additional_kwargs = _message_additional_kwargs(message)
|
||||
return message_type == "remove" or _message_name(message) == "summary" or additional_kwargs.get("hide_from_ui") is True
|
||||
|
||||
|
||||
def _is_visible_human_message(message: Any) -> bool:
|
||||
return _message_type(message) == "human" and not _is_hidden_or_control_message(message)
|
||||
|
||||
|
||||
def _is_visible_ai_message(message: Any) -> bool:
|
||||
return _message_type(message) == "ai" and not _is_hidden_or_control_message(message)
|
||||
|
||||
|
||||
def _checkpoint_messages(checkpoint_tuple: Any) -> list[Any]:
|
||||
checkpoint = getattr(checkpoint_tuple, "checkpoint", None) or {}
|
||||
channel_values = checkpoint.get("channel_values", {}) if isinstance(checkpoint, dict) else {}
|
||||
messages = channel_values.get("messages", []) if isinstance(channel_values, dict) else []
|
||||
return messages if isinstance(messages, list) else []
|
||||
|
||||
|
||||
def _checkpoint_configurable(checkpoint_tuple: Any) -> dict[str, Any]:
|
||||
config = getattr(checkpoint_tuple, "config", None) or {}
|
||||
configurable = config.get("configurable", {}) if isinstance(config, dict) else {}
|
||||
return dict(configurable) if isinstance(configurable, dict) else {}
|
||||
|
||||
|
||||
def _checkpoint_response(checkpoint_tuple: Any) -> dict[str, Any]:
|
||||
configurable = _checkpoint_configurable(checkpoint_tuple)
|
||||
checkpoint_id = configurable.get("checkpoint_id")
|
||||
if not checkpoint_id:
|
||||
raise HTTPException(status_code=409, detail="Checkpoint is missing checkpoint_id")
|
||||
return {
|
||||
"checkpoint_ns": str(configurable.get("checkpoint_ns") or ""),
|
||||
"checkpoint_id": str(checkpoint_id),
|
||||
"checkpoint_map": configurable.get("checkpoint_map"),
|
||||
}
|
||||
|
||||
|
||||
def _clean_human_message_for_regenerate(message: Any) -> dict[str, Any]:
|
||||
additional_kwargs = _message_additional_kwargs(message)
|
||||
content = get_original_user_content_text(_message_content(message), additional_kwargs)
|
||||
additional_kwargs.pop(ORIGINAL_USER_CONTENT_KEY, None)
|
||||
additional_kwargs.pop("hide_from_ui", None)
|
||||
|
||||
clean_message: dict[str, Any] = {
|
||||
"type": "human",
|
||||
"content": [{"type": "text", "text": content}],
|
||||
"additional_kwargs": additional_kwargs,
|
||||
}
|
||||
message_id = _message_id(message)
|
||||
if message_id:
|
||||
clean_message["id"] = message_id
|
||||
name = _message_name(message)
|
||||
if name:
|
||||
clean_message["name"] = name
|
||||
return clean_message
|
||||
|
||||
|
||||
def _event_message_id(row: dict[str, Any]) -> str | None:
|
||||
content = row.get("content")
|
||||
if isinstance(content, BaseMessage):
|
||||
return _message_id(content)
|
||||
if isinstance(content, dict):
|
||||
return _message_id(content)
|
||||
return None
|
||||
|
||||
|
||||
def _run_last_ai_matches_message(record: RunRecord, message: Any) -> bool:
|
||||
last_ai_message = (record.last_ai_message or "").strip()
|
||||
if not last_ai_message:
|
||||
return False
|
||||
target_text = _message_text(message).strip()
|
||||
if not target_text:
|
||||
return False
|
||||
return last_ai_message == target_text[: len(last_ai_message)]
|
||||
|
||||
|
||||
async def _find_target_run_id(thread_id: str, message_id: str, target_message: Any, request: Request) -> str:
|
||||
event_store = get_run_event_store(request)
|
||||
rows = await event_store.list_messages(thread_id, limit=REGENERATE_HISTORY_SCAN_LIMIT)
|
||||
for row in reversed(rows):
|
||||
if row.get("event_type") not in {"ai_message", "llm.ai.response"}:
|
||||
continue
|
||||
if _event_message_id(row) == message_id:
|
||||
run_id = row.get("run_id")
|
||||
if isinstance(run_id, str) and run_id:
|
||||
return run_id
|
||||
run_mgr = get_run_manager(request)
|
||||
user_id = await get_current_user(request)
|
||||
records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=10)
|
||||
fallback_record = next(
|
||||
(record for record in records if record.status == RunStatus.success and _run_last_ai_matches_message(record, target_message)),
|
||||
None,
|
||||
)
|
||||
if fallback_record is not None:
|
||||
return fallback_record.run_id
|
||||
if len(rows) >= REGENERATE_HISTORY_SCAN_LIMIT:
|
||||
logger.warning(
|
||||
"Could not find source run for regenerate message %s in recent run events for thread %s (limit=%s)",
|
||||
message_id,
|
||||
thread_id,
|
||||
REGENERATE_HISTORY_SCAN_LIMIT,
|
||||
)
|
||||
raise HTTPException(status_code=409, detail="Could not find source run for assistant message")
|
||||
|
||||
|
||||
async def _find_base_checkpoint_before_human(thread_id: str, human_message_id: str, request: Request) -> Any:
|
||||
checkpointer = get_checkpointer(request)
|
||||
base_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
try:
|
||||
checkpoints = [item async for item in checkpointer.alist(base_config, limit=REGENERATE_HISTORY_SCAN_LIMIT)]
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to list checkpoints for regenerate thread %s", thread_id)
|
||||
raise HTTPException(status_code=500, detail="Failed to inspect checkpoint history") from exc
|
||||
|
||||
previous_checkpoint = None
|
||||
for checkpoint_tuple in reversed(checkpoints):
|
||||
messages = _checkpoint_messages(checkpoint_tuple)
|
||||
message_ids = {_message_id(message) for message in messages}
|
||||
if human_message_id in message_ids:
|
||||
if previous_checkpoint is None:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Could not find an addressable checkpoint before the target user message",
|
||||
)
|
||||
return previous_checkpoint
|
||||
if _checkpoint_configurable(checkpoint_tuple).get("checkpoint_id"):
|
||||
previous_checkpoint = checkpoint_tuple
|
||||
|
||||
if len(checkpoints) >= REGENERATE_HISTORY_SCAN_LIMIT:
|
||||
logger.warning(
|
||||
"Could not locate target user message %s in recent checkpoint history for thread %s (limit=%s)",
|
||||
human_message_id,
|
||||
thread_id,
|
||||
REGENERATE_HISTORY_SCAN_LIMIT,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=(f"Could not locate target user message in recent checkpoint history (limit={REGENERATE_HISTORY_SCAN_LIMIT})"),
|
||||
)
|
||||
|
||||
|
||||
async def _prepare_regenerate_payload(thread_id: str, message_id: str, request: Request) -> RegeneratePrepareResponse:
|
||||
checkpointer = get_checkpointer(request)
|
||||
latest_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
try:
|
||||
latest_checkpoint = await checkpointer.aget_tuple(latest_config)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to read latest checkpoint for regenerate thread %s", thread_id)
|
||||
raise HTTPException(status_code=500, detail="Failed to read latest checkpoint") from exc
|
||||
if latest_checkpoint is None:
|
||||
raise HTTPException(status_code=404, detail=f"Thread {thread_id} has no checkpoint")
|
||||
|
||||
messages = _checkpoint_messages(latest_checkpoint)
|
||||
target_index = next((i for i, message in enumerate(messages) if _message_id(message) == message_id), None)
|
||||
if target_index is None:
|
||||
raise HTTPException(status_code=404, detail=f"Message {message_id} not found")
|
||||
target_message = messages[target_index]
|
||||
if not _is_visible_ai_message(target_message):
|
||||
raise HTTPException(status_code=409, detail="Only visible assistant messages can be regenerated")
|
||||
|
||||
latest_visible_ai = next((message for message in reversed(messages) if _is_visible_ai_message(message)), None)
|
||||
if _message_id(latest_visible_ai) != message_id:
|
||||
raise HTTPException(status_code=409, detail="Only the latest assistant message can be regenerated")
|
||||
|
||||
previous_human = next((message for message in reversed(messages[:target_index]) if _is_visible_human_message(message)), None)
|
||||
if previous_human is None:
|
||||
raise HTTPException(status_code=409, detail="Could not find the user message for this assistant response")
|
||||
previous_human_id = _message_id(previous_human)
|
||||
if not previous_human_id:
|
||||
raise HTTPException(status_code=409, detail="The source user message is missing an id")
|
||||
|
||||
base_checkpoint_tuple = await _find_base_checkpoint_before_human(thread_id, previous_human_id, request)
|
||||
target_run_id = await _find_target_run_id(thread_id, message_id, target_message, request)
|
||||
checkpoint = _checkpoint_response(base_checkpoint_tuple)
|
||||
metadata = {
|
||||
"regenerate_from_message_id": message_id,
|
||||
"regenerate_from_run_id": target_run_id,
|
||||
"regenerate_checkpoint_id": checkpoint["checkpoint_id"],
|
||||
}
|
||||
return RegeneratePrepareResponse(
|
||||
input={"messages": [_clean_human_message_for_regenerate(previous_human)]},
|
||||
checkpoint=checkpoint,
|
||||
metadata=metadata,
|
||||
target_run_id=target_run_id,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Endpoints
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("/{thread_id}/runs/regenerate/prepare", response_model=RegeneratePrepareResponse)
|
||||
@require_permission("runs", "create", owner_check=True, require_existing=True)
|
||||
async def prepare_regenerate_run(
|
||||
thread_id: str,
|
||||
body: RegeneratePrepareRequest,
|
||||
request: Request,
|
||||
) -> RegeneratePrepareResponse:
|
||||
"""Prepare input and checkpoint for regenerating the latest assistant turn."""
|
||||
return await _prepare_regenerate_payload(thread_id, body.message_id, 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:
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from fastapi import HTTPException, Request
|
|||
from langchain_core.messages import BaseMessage
|
||||
from langchain_core.messages.utils import convert_to_messages
|
||||
|
||||
from app.gateway.deps import get_run_context, get_run_manager, get_stream_bridge
|
||||
from app.gateway.deps import get_checkpointer, get_run_context, get_run_manager, get_stream_bridge
|
||||
from app.gateway.internal_auth import INTERNAL_SYSTEM_ROLE, get_trusted_internal_owner_user_id
|
||||
from app.gateway.utils import sanitize_log_param
|
||||
from deerflow.config.app_config import get_app_config
|
||||
|
|
@ -284,6 +284,65 @@ def build_run_config(
|
|||
return config
|
||||
|
||||
|
||||
async def apply_checkpoint_to_run_config(
|
||||
config: dict[str, Any],
|
||||
*,
|
||||
body: Any,
|
||||
thread_id: str,
|
||||
request: Request,
|
||||
) -> None:
|
||||
"""Validate an optional run checkpoint and attach it to RunnableConfig."""
|
||||
checkpoint = getattr(body, "checkpoint", None)
|
||||
checkpoint_id = getattr(body, "checkpoint_id", None)
|
||||
checkpoint_ns = ""
|
||||
checkpoint_map = None
|
||||
|
||||
if checkpoint:
|
||||
if not isinstance(checkpoint, Mapping):
|
||||
raise HTTPException(status_code=400, detail="checkpoint must be an object")
|
||||
checkpoint_thread_id = checkpoint.get("thread_id")
|
||||
if checkpoint_thread_id is not None and str(checkpoint_thread_id) != thread_id:
|
||||
raise HTTPException(status_code=400, detail="checkpoint thread_id does not match request thread_id")
|
||||
raw_checkpoint_id = checkpoint.get("checkpoint_id")
|
||||
if raw_checkpoint_id:
|
||||
checkpoint_id = str(raw_checkpoint_id)
|
||||
raw_checkpoint_ns = checkpoint.get("checkpoint_ns")
|
||||
if raw_checkpoint_ns is not None:
|
||||
checkpoint_ns = str(raw_checkpoint_ns)
|
||||
checkpoint_map = checkpoint.get("checkpoint_map")
|
||||
|
||||
if not checkpoint_id:
|
||||
return
|
||||
|
||||
read_config: dict[str, Any] = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": checkpoint_ns,
|
||||
"checkpoint_id": str(checkpoint_id),
|
||||
}
|
||||
}
|
||||
if checkpoint_map is not None:
|
||||
read_config["configurable"]["checkpoint_map"] = checkpoint_map
|
||||
|
||||
checkpointer = get_checkpointer(request)
|
||||
try:
|
||||
checkpoint_tuple = await checkpointer.aget_tuple(read_config)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to validate checkpoint %s for thread %s", checkpoint_id, sanitize_log_param(thread_id))
|
||||
raise HTTPException(status_code=500, detail="Failed to validate checkpoint") from exc
|
||||
if checkpoint_tuple is None:
|
||||
raise HTTPException(status_code=404, detail=f"Checkpoint {checkpoint_id} not found")
|
||||
|
||||
configurable = config.setdefault("configurable", {})
|
||||
if not isinstance(configurable, dict):
|
||||
raise HTTPException(status_code=400, detail="request config configurable must be an object")
|
||||
configurable["thread_id"] = thread_id
|
||||
configurable["checkpoint_ns"] = checkpoint_ns
|
||||
configurable["checkpoint_id"] = str(checkpoint_id)
|
||||
if checkpoint_map is not None:
|
||||
configurable["checkpoint_map"] = checkpoint_map
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Run lifecycle
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -396,6 +455,7 @@ async def start_run(
|
|||
agent_factory = resolve_agent_factory(body.assistant_id)
|
||||
graph_input = normalize_input(body.input)
|
||||
config = build_run_config(thread_id, body.config, body.metadata, assistant_id=body.assistant_id)
|
||||
await apply_checkpoint_to_run_config(config, body=body, thread_id=thread_id, request=request)
|
||||
|
||||
# Merge DeerFlow-specific context overrides into both ``configurable`` and ``context``.
|
||||
# The ``context`` field is a custom extension for the langgraph-compat layer
|
||||
|
|
|
|||
|
|
@ -373,6 +373,70 @@ def test_run_create_request_context_defaults_to_none():
|
|||
assert body.context is None
|
||||
|
||||
|
||||
def test_apply_checkpoint_to_run_config_writes_checkpoint_fields():
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.services import apply_checkpoint_to_run_config
|
||||
|
||||
class FakeCheckpointer:
|
||||
def __init__(self):
|
||||
self.seen_config = None
|
||||
|
||||
async def aget_tuple(self, config):
|
||||
self.seen_config = config
|
||||
return SimpleNamespace(config=config, checkpoint={"channel_values": {}})
|
||||
|
||||
checkpointer = FakeCheckpointer()
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(checkpointer=checkpointer)))
|
||||
body = SimpleNamespace(
|
||||
checkpoint={
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "ckpt-1",
|
||||
"checkpoint_map": {"": "ckpt-1"},
|
||||
},
|
||||
checkpoint_id=None,
|
||||
)
|
||||
config = {"configurable": {"thread_id": "thread-1"}}
|
||||
|
||||
asyncio.run(apply_checkpoint_to_run_config(config, body=body, thread_id="thread-1", request=request))
|
||||
|
||||
assert checkpointer.seen_config == {
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "ckpt-1",
|
||||
"checkpoint_map": {"": "ckpt-1"},
|
||||
}
|
||||
}
|
||||
assert config["configurable"]["checkpoint_id"] == "ckpt-1"
|
||||
assert config["configurable"]["checkpoint_ns"] == ""
|
||||
assert config["configurable"]["checkpoint_map"] == {"": "ckpt-1"}
|
||||
|
||||
|
||||
def test_apply_checkpoint_to_run_config_rejects_missing_checkpoint():
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from app.gateway.services import apply_checkpoint_to_run_config
|
||||
|
||||
class FakeCheckpointer:
|
||||
async def aget_tuple(self, config):
|
||||
return None
|
||||
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(checkpointer=FakeCheckpointer())))
|
||||
body = SimpleNamespace(checkpoint=None, checkpoint_id="missing")
|
||||
config = {"configurable": {"thread_id": "thread-1"}}
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(apply_checkpoint_to_run_config(config, body=body, thread_id="thread-1", request=request))
|
||||
|
||||
assert exc.value.status_code == 404
|
||||
assert "missing" in exc.value.detail
|
||||
|
||||
|
||||
def test_context_merges_into_configurable():
|
||||
"""Context values must be merged into config['configurable'] by start_run.
|
||||
|
||||
|
|
|
|||
261
backend/tests/test_thread_regenerate_prepare.py
Normal file
261
backend/tests/test_thread_regenerate_prepare.py
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from deerflow.runtime import RunStatus
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
|
||||
def _checkpoint(checkpoint_id: str, messages: list[object]):
|
||||
return SimpleNamespace(
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "thread-1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": checkpoint_id,
|
||||
"checkpoint_map": None,
|
||||
}
|
||||
},
|
||||
checkpoint={"channel_values": {"messages": messages}},
|
||||
)
|
||||
|
||||
|
||||
class FakeCheckpointer:
|
||||
def __init__(self, history, *, latest=None):
|
||||
self.history = history
|
||||
self.latest = latest
|
||||
self.alist_limits = []
|
||||
|
||||
async def aget_tuple(self, config):
|
||||
checkpoint_id = config.get("configurable", {}).get("checkpoint_id")
|
||||
if checkpoint_id:
|
||||
return next((item for item in self.history if item.config["configurable"]["checkpoint_id"] == checkpoint_id), None)
|
||||
return self.latest or (self.history[0] if self.history else None)
|
||||
|
||||
async def alist(self, config, limit=200):
|
||||
self.alist_limits.append(limit)
|
||||
for item in self.history[:limit]:
|
||||
yield item
|
||||
|
||||
|
||||
class FakeEventStore:
|
||||
def __init__(self, rows):
|
||||
self.rows = rows
|
||||
|
||||
async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq=None):
|
||||
return self.rows[-limit:]
|
||||
|
||||
|
||||
class FakeRunManager:
|
||||
def __init__(self, records):
|
||||
self.records = records
|
||||
|
||||
async def list_by_thread(self, thread_id, *, user_id=None, limit=100):
|
||||
return self.records[:limit]
|
||||
|
||||
|
||||
def _request(checkpointer, event_store, *, run_manager=None, user_id="user-1"):
|
||||
from app.gateway.auth_disabled import AUTH_SOURCE_SESSION
|
||||
|
||||
return SimpleNamespace(
|
||||
app=SimpleNamespace(
|
||||
state=SimpleNamespace(
|
||||
checkpointer=checkpointer,
|
||||
run_event_store=event_store,
|
||||
run_manager=run_manager or FakeRunManager([]),
|
||||
)
|
||||
),
|
||||
state=SimpleNamespace(user=SimpleNamespace(id=user_id), auth_source=AUTH_SOURCE_SESSION),
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_returns_clean_input_and_base_checkpoint():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
human = HumanMessage(
|
||||
id="human-1",
|
||||
content="<uploaded_files>injected</uploaded_files>\n\n/data-analysis analyze data.csv",
|
||||
additional_kwargs={
|
||||
ORIGINAL_USER_CONTENT_KEY: "/data-analysis analyze data.csv",
|
||||
"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
|
||||
},
|
||||
)
|
||||
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"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
response = asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
||||
|
||||
assert response.checkpoint == {
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": "ckpt-base",
|
||||
"checkpoint_map": None,
|
||||
}
|
||||
assert response.target_run_id == "run-old"
|
||||
assert response.metadata == {
|
||||
"regenerate_from_message_id": "ai-1",
|
||||
"regenerate_from_run_id": "run-old",
|
||||
"regenerate_checkpoint_id": "ckpt-base",
|
||||
}
|
||||
regenerated_human = response.input["messages"][0]
|
||||
assert regenerated_human["id"] == "human-1"
|
||||
assert regenerated_human["content"] == [{"type": "text", "text": "/data-analysis analyze data.csv"}]
|
||||
assert regenerated_human["additional_kwargs"] == {"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}]}
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_rejects_non_latest_assistant():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
human = HumanMessage(id="human-1", content="question")
|
||||
old_ai = AIMessage(id="ai-old", content="old")
|
||||
latest_ai = AIMessage(id="ai-latest", content="latest")
|
||||
base = _checkpoint("ckpt-base", [])
|
||||
after_human = _checkpoint("ckpt-human", [human])
|
||||
latest = _checkpoint("ckpt-latest", [human, old_ai, latest_ai])
|
||||
checkpointer = FakeCheckpointer([latest, after_human, base])
|
||||
event_store = FakeEventStore(
|
||||
[
|
||||
{
|
||||
"run_id": "run-old",
|
||||
"event_type": "ai_message",
|
||||
"category": "message",
|
||||
"content": {"id": "ai-old", "type": "ai", "content": "old"},
|
||||
"metadata": {"caller": "lead_agent"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(_prepare_regenerate_payload("thread-1", "ai-old", _request(checkpointer, event_store)))
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Only the latest assistant message can be regenerated"
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_falls_back_to_matching_run_when_events_are_missing():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
human = HumanMessage(id="human-1", content="question")
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
base = _checkpoint("ckpt-base", [])
|
||||
after_human = _checkpoint("ckpt-human", [human])
|
||||
latest = _checkpoint("ckpt-ai", [human, ai])
|
||||
checkpointer = FakeCheckpointer([latest, after_human, base])
|
||||
run_manager = FakeRunManager(
|
||||
[
|
||||
SimpleNamespace(run_id="run-latest", status=RunStatus.success, last_ai_message="answer"),
|
||||
SimpleNamespace(run_id="run-older", status=RunStatus.error, last_ai_message="answer"),
|
||||
]
|
||||
)
|
||||
|
||||
response = asyncio.run(
|
||||
_prepare_regenerate_payload(
|
||||
"thread-1",
|
||||
"ai-1",
|
||||
_request(checkpointer, FakeEventStore([]), run_manager=run_manager),
|
||||
)
|
||||
)
|
||||
|
||||
assert response.target_run_id == "run-latest"
|
||||
assert response.metadata["regenerate_from_run_id"] == "run-latest"
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_rejects_unverified_run_fallback_when_events_are_missing():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
human = HumanMessage(id="human-1", content="question")
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
base = _checkpoint("ckpt-base", [])
|
||||
after_human = _checkpoint("ckpt-human", [human])
|
||||
latest = _checkpoint("ckpt-ai", [human, ai])
|
||||
checkpointer = FakeCheckpointer([latest, after_human, base])
|
||||
run_manager = FakeRunManager(
|
||||
[
|
||||
SimpleNamespace(run_id="run-latest", status=RunStatus.success, last_ai_message="different"),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(
|
||||
_prepare_regenerate_payload(
|
||||
"thread-1",
|
||||
"ai-1",
|
||||
_request(checkpointer, FakeEventStore([]), run_manager=run_manager),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Could not find source run for assistant message"
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_requires_addressable_checkpoint_before_human():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
human = HumanMessage(id="human-1", content="question")
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
latest = _checkpoint("ckpt-ai", [human, ai])
|
||||
checkpointer = FakeCheckpointer([latest])
|
||||
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"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Could not find an addressable checkpoint before the target user message"
|
||||
assert checkpointer.alist_limits == [200]
|
||||
|
||||
|
||||
def test_prepare_regenerate_payload_reports_recent_checkpoint_scan_limit():
|
||||
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
|
||||
|
||||
human = HumanMessage(id="human-1", content="question")
|
||||
ai = AIMessage(id="ai-1", content="answer")
|
||||
latest = _checkpoint("ckpt-latest", [human, ai])
|
||||
history_without_human = [_checkpoint(f"ckpt-{index}", []) for index in range(201)]
|
||||
checkpointer = FakeCheckpointer(history_without_human, latest=latest)
|
||||
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"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(_prepare_regenerate_payload("thread-1", "ai-1", _request(checkpointer, event_store)))
|
||||
|
||||
assert exc.value.status_code == 409
|
||||
assert exc.value.detail == "Could not locate target user message in recent checkpoint history (limit=200)"
|
||||
assert checkpointer.alist_limits == [200]
|
||||
|
|
@ -80,6 +80,7 @@ export default function ChatPage() {
|
|||
thread,
|
||||
pendingUsageMessages,
|
||||
sendMessage,
|
||||
regenerateMessage,
|
||||
isUploading,
|
||||
isHistoryLoading,
|
||||
hasMoreHistory,
|
||||
|
|
@ -159,6 +160,11 @@ export default function ChatPage() {
|
|||
const handleStop = useCallback(async () => {
|
||||
await thread.stop();
|
||||
}, [thread]);
|
||||
const handleRegenerate = useCallback(
|
||||
(messageId: string, supersededMessageIds: string[]) =>
|
||||
regenerateMessage(threadId, messageId, supersededMessageIds),
|
||||
[regenerateMessage, threadId],
|
||||
);
|
||||
|
||||
const tokenUsageInlineMode = tokenUsageEnabled
|
||||
? localSettings.tokenUsage.inlineMode
|
||||
|
|
@ -208,6 +214,14 @@ export default function ChatPage() {
|
|||
loadMoreHistory={loadMoreHistory}
|
||||
isHistoryLoading={isHistoryLoading}
|
||||
tokenUsageInlineMode={tokenUsageInlineMode}
|
||||
canRegenerate={
|
||||
!isNewThread &&
|
||||
!isMock &&
|
||||
env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY !== "true" &&
|
||||
!isUploading &&
|
||||
!thread.isLoading
|
||||
}
|
||||
onRegenerateMessage={handleRegenerate}
|
||||
/>
|
||||
</div>
|
||||
<div
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import type { Message } from "@langchain/langgraph-sdk";
|
||||
import type { BaseStream } from "@langchain/langgraph-sdk/react";
|
||||
import { ChevronUpIcon, Loader2Icon } from "lucide-react";
|
||||
import { ChevronUpIcon, Loader2Icon, RefreshCcwIcon } from "lucide-react";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
|
||||
import {
|
||||
|
|
@ -43,6 +43,7 @@ import { cn } from "@/lib/utils";
|
|||
import { ArtifactFileList } from "../artifacts/artifact-file-list";
|
||||
import { CopyButton } from "../copy-button";
|
||||
import { StreamingIndicator } from "../streaming-indicator";
|
||||
import { Tooltip } from "../tooltip";
|
||||
|
||||
import { MarkdownContent } from "./markdown-content";
|
||||
import { MessageGroup } from "./message-group";
|
||||
|
|
@ -172,6 +173,8 @@ export function MessageList({
|
|||
hasMoreHistory,
|
||||
loadMoreHistory,
|
||||
isHistoryLoading,
|
||||
onRegenerateMessage,
|
||||
canRegenerate = false,
|
||||
}: {
|
||||
className?: string;
|
||||
threadId: string;
|
||||
|
|
@ -181,6 +184,11 @@ export function MessageList({
|
|||
hasMoreHistory?: boolean;
|
||||
loadMoreHistory?: () => void;
|
||||
isHistoryLoading?: boolean;
|
||||
onRegenerateMessage?: (
|
||||
messageId: string,
|
||||
supersededMessageIds: string[],
|
||||
) => void | Promise<void>;
|
||||
canRegenerate?: boolean;
|
||||
}) {
|
||||
const { t } = useI18n();
|
||||
const [turnStartTime, setTurnStartTime] = useState<number | null>(null);
|
||||
|
|
@ -194,6 +202,9 @@ export function MessageList({
|
|||
}, [thread.isLoading]);
|
||||
const messages = thread.messages;
|
||||
const groupedMessages = getMessageGroups(messages);
|
||||
const [regeneratingMessageId, setRegeneratingMessageId] = useState<
|
||||
string | null
|
||||
>(null);
|
||||
const hasActiveAssistantText = useMemo(() => {
|
||||
let lastHumanIndex = -1;
|
||||
for (let i = groupedMessages.length - 1; i >= 0; i--) {
|
||||
|
|
@ -226,21 +237,86 @@ export function MessageList({
|
|||
[messages, thread.getMessagesMetadata, thread.isLoading],
|
||||
);
|
||||
|
||||
const renderAssistantCopyButton = useCallback(
|
||||
(messages: Message[], isStreaming: boolean) => {
|
||||
const clipboardData = getAssistantTurnCopyData(messages, { isStreaming });
|
||||
const latestAssistantGroupId = useMemo(() => {
|
||||
if (thread.isLoading) {
|
||||
return null;
|
||||
}
|
||||
for (let i = groupedMessages.length - 1; i >= 0; i -= 1) {
|
||||
const group = groupedMessages[i];
|
||||
if (group?.type === "assistant") {
|
||||
return group.id;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}, [groupedMessages, thread.isLoading]);
|
||||
|
||||
if (!clipboardData) {
|
||||
const renderAssistantActions = useCallback(
|
||||
(
|
||||
messages: Message[],
|
||||
isStreaming: boolean,
|
||||
enableRegenerateForTurn: boolean,
|
||||
) => {
|
||||
const clipboardData = getAssistantTurnCopyData(messages, { isStreaming });
|
||||
const regenerateTarget = [...messages]
|
||||
.reverse()
|
||||
.find((message) => message.type === "ai" && message.id);
|
||||
const supersededMessageIds = messages
|
||||
.filter((message) => message.type === "ai" && message.id)
|
||||
.map((message) => message.id)
|
||||
.filter((id): id is string => typeof id === "string");
|
||||
|
||||
if (!clipboardData && !regenerateTarget) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="mt-2 flex justify-start opacity-0 transition-opacity delay-200 duration-300 group-hover/assistant-turn:opacity-100">
|
||||
<CopyButton clipboardData={clipboardData} />
|
||||
<div className="mt-2 flex justify-start gap-1 opacity-0 transition-opacity delay-200 duration-300 group-hover/assistant-turn:opacity-100">
|
||||
{clipboardData && <CopyButton clipboardData={clipboardData} />}
|
||||
{enableRegenerateForTurn &&
|
||||
regenerateTarget?.id &&
|
||||
onRegenerateMessage && (
|
||||
<Tooltip content={t.common.regenerate}>
|
||||
<Button
|
||||
aria-label={t.common.regenerate}
|
||||
size="icon-sm"
|
||||
type="button"
|
||||
variant="ghost"
|
||||
disabled={
|
||||
!canRegenerate ||
|
||||
regeneratingMessageId === regenerateTarget.id
|
||||
}
|
||||
onClick={() => {
|
||||
const targetId = regenerateTarget.id;
|
||||
if (!targetId) {
|
||||
return;
|
||||
}
|
||||
setRegeneratingMessageId(targetId);
|
||||
void Promise.resolve(
|
||||
onRegenerateMessage?.(targetId, supersededMessageIds),
|
||||
).finally(() => {
|
||||
setRegeneratingMessageId(null);
|
||||
});
|
||||
}}
|
||||
>
|
||||
<RefreshCcwIcon
|
||||
className={cn(
|
||||
"size-3",
|
||||
regeneratingMessageId === regenerateTarget.id &&
|
||||
"animate-spin",
|
||||
)}
|
||||
/>
|
||||
</Button>
|
||||
</Tooltip>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
},
|
||||
[],
|
||||
[
|
||||
canRegenerate,
|
||||
onRegenerateMessage,
|
||||
regeneratingMessageId,
|
||||
t.common.regenerate,
|
||||
],
|
||||
);
|
||||
|
||||
const renderTokenUsage = useCallback(
|
||||
|
|
@ -341,12 +417,13 @@ export function MessageList({
|
|||
turnUsageMessages,
|
||||
})}
|
||||
{group.type === "assistant" &&
|
||||
renderAssistantCopyButton(
|
||||
renderAssistantActions(
|
||||
group.messages,
|
||||
isAssistantMessageGroupStreaming(
|
||||
group.messages,
|
||||
streamingMessages,
|
||||
),
|
||||
group.id === latestAssistantGroupId,
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ export const enUS: Translations = {
|
|||
exportAsMarkdown: "Export as Markdown",
|
||||
exportAsJSON: "Export as JSON",
|
||||
exportSuccess: "Conversation exported",
|
||||
regenerate: "Regenerate",
|
||||
},
|
||||
|
||||
// Home
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ export interface Translations {
|
|||
exportAsMarkdown: string;
|
||||
exportAsJSON: string;
|
||||
exportSuccess: string;
|
||||
regenerate: string;
|
||||
};
|
||||
|
||||
home: {
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ export const zhCN: Translations = {
|
|||
exportAsMarkdown: "导出为 Markdown",
|
||||
exportAsJSON: "导出为 JSON",
|
||||
exportSuccess: "对话已导出",
|
||||
regenerate: "重新生成",
|
||||
},
|
||||
|
||||
// Home
|
||||
|
|
|
|||
|
|
@ -59,6 +59,17 @@ type SendMessageOptions = {
|
|||
additionalKwargs?: Record<string, unknown>;
|
||||
};
|
||||
|
||||
type RegeneratePrepareResponse = {
|
||||
input: Partial<AgentThreadState>;
|
||||
checkpoint: {
|
||||
checkpoint_ns: string;
|
||||
checkpoint_id: string;
|
||||
checkpoint_map: Record<string, unknown> | null;
|
||||
};
|
||||
metadata: Record<string, unknown>;
|
||||
target_run_id: string;
|
||||
};
|
||||
|
||||
const EMPTY_THREAD_VALUES: AgentThreadState = {
|
||||
title: "",
|
||||
messages: [],
|
||||
|
|
@ -121,6 +132,69 @@ function dedupeMessagesByIdentity(messages: Message[]): Message[] {
|
|||
});
|
||||
}
|
||||
|
||||
function dedupeRunMessagesByIdentity(messages: RunMessage[]): RunMessage[] {
|
||||
const lastIndexByIdentity = new Map<string, number>();
|
||||
messages.forEach((message, index) => {
|
||||
const identity = messageIdentity(message.content);
|
||||
if (identity) {
|
||||
lastIndexByIdentity.set(`${message.run_id}:${identity}`, index);
|
||||
}
|
||||
});
|
||||
|
||||
return messages.filter((message, index) => {
|
||||
const identity = messageIdentity(message.content);
|
||||
if (!identity) {
|
||||
return true;
|
||||
}
|
||||
return lastIndexByIdentity.get(`${message.run_id}:${identity}`) === index;
|
||||
});
|
||||
}
|
||||
|
||||
export function getSupersededRunIds(
|
||||
runs: Run[] | undefined,
|
||||
pendingSupersededRunIds?: ReadonlySet<string>,
|
||||
) {
|
||||
const ids = new Set(pendingSupersededRunIds ?? []);
|
||||
for (const run of runs ?? []) {
|
||||
if (run.status !== "success") {
|
||||
continue;
|
||||
}
|
||||
const metadata = run.metadata;
|
||||
if (metadata && typeof metadata === "object") {
|
||||
const fromRunId = Reflect.get(metadata, "regenerate_from_run_id");
|
||||
if (typeof fromRunId === "string" && fromRunId) {
|
||||
ids.add(fromRunId);
|
||||
}
|
||||
}
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
export function removeSetItems<T>(
|
||||
values: ReadonlySet<T>,
|
||||
itemsToRemove: Iterable<T>,
|
||||
) {
|
||||
const next = new Set(values);
|
||||
for (const item of itemsToRemove) {
|
||||
next.delete(item);
|
||||
}
|
||||
return next;
|
||||
}
|
||||
|
||||
export function buildVisibleHistoryMessages(
|
||||
messageRows: RunMessage[],
|
||||
supersededRunIds: ReadonlySet<string>,
|
||||
appendedMessages: Message[],
|
||||
) {
|
||||
const visibleRows = messageRows.filter(
|
||||
(message) => !supersededRunIds.has(message.run_id),
|
||||
);
|
||||
return dedupeMessagesByIdentity([
|
||||
...visibleRows.map((message) => message.content),
|
||||
...appendedMessages,
|
||||
]);
|
||||
}
|
||||
|
||||
export function findLatestUnloadedRunIndex(
|
||||
runs: Run[],
|
||||
loadedRunIds: ReadonlySet<string>,
|
||||
|
|
@ -399,6 +473,21 @@ function getStreamErrorMessage(error: unknown): string {
|
|||
return "Request failed.";
|
||||
}
|
||||
|
||||
async function readResponseErrorMessage(
|
||||
response: Response,
|
||||
fallback = "Request failed.",
|
||||
) {
|
||||
try {
|
||||
const data = await response.json();
|
||||
if (typeof data?.detail === "string" && data.detail.trim()) {
|
||||
return data.detail;
|
||||
}
|
||||
} catch {
|
||||
// Use the fallback below when the response body is not JSON.
|
||||
}
|
||||
return response.statusText || fallback;
|
||||
}
|
||||
|
||||
function getHttpStatus(error: unknown): number | undefined {
|
||||
if (typeof error !== "object" || error === null) {
|
||||
return undefined;
|
||||
|
|
@ -449,6 +538,11 @@ export function useThreadStream({
|
|||
const [liveMessagesThreadId, setLiveMessagesThreadId] = useState<
|
||||
string | null
|
||||
>(null);
|
||||
const [pendingSupersededRunIds, setPendingSupersededRunIds] = useState<
|
||||
ReadonlySet<string>
|
||||
>(() => new Set());
|
||||
const [pendingSupersededMessageIds, setPendingSupersededMessageIds] =
|
||||
useState<ReadonlySet<string>>(() => new Set());
|
||||
const [isUploading, setIsUploading] = useState(false);
|
||||
// Track the thread ID that is currently streaming to handle thread changes during streaming
|
||||
const [onStreamThreadId, setOnStreamThreadId] = useState(() => threadId);
|
||||
|
|
@ -470,7 +564,10 @@ export function useThreadStream({
|
|||
loadMore: loadMoreHistory,
|
||||
loading: isHistoryLoading,
|
||||
appendMessages,
|
||||
} = useThreadHistory(onStreamThreadId ?? "", { enabled: !isMock });
|
||||
} = useThreadHistory(onStreamThreadId ?? "", {
|
||||
enabled: !isMock,
|
||||
pendingSupersededRunIds,
|
||||
});
|
||||
|
||||
// Keep listeners ref updated with latest callbacks
|
||||
useEffect(() => {
|
||||
|
|
@ -686,6 +783,8 @@ export function useThreadStream({
|
|||
setOptimisticMessages([]);
|
||||
setOptimisticThreadId(null);
|
||||
setLiveMessagesThreadId(null);
|
||||
setPendingSupersededRunIds(new Set());
|
||||
setPendingSupersededMessageIds(new Set());
|
||||
toast.error(getStreamErrorMessage(error));
|
||||
pendingUsageBaselineMessageIdsRef.current = new Set(
|
||||
messagesRef.current
|
||||
|
|
@ -710,6 +809,9 @@ export function useThreadStream({
|
|||
queryKey: INFINITE_THREADS_QUERY_KEY_PREFIX,
|
||||
});
|
||||
if (threadIdRef.current && !isMock) {
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: ["thread", threadIdRef.current],
|
||||
});
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: threadTokenUsageQueryKey(threadIdRef.current),
|
||||
});
|
||||
|
|
@ -720,8 +822,14 @@ export function useThreadStream({
|
|||
const hasVisibleStreamState =
|
||||
Boolean(threadId) || liveMessagesThreadId === currentViewThreadId;
|
||||
const persistedMessages = useMemo(
|
||||
() => (hasVisibleStreamState ? thread.messages : []),
|
||||
[hasVisibleStreamState, thread.messages],
|
||||
() =>
|
||||
hasVisibleStreamState
|
||||
? thread.messages.filter(
|
||||
(message) =>
|
||||
!message.id || !pendingSupersededMessageIds.has(message.id),
|
||||
)
|
||||
: [],
|
||||
[hasVisibleStreamState, pendingSupersededMessageIds, thread.messages],
|
||||
);
|
||||
const visibleHistory = useMemo(
|
||||
() => (threadId ? history : []),
|
||||
|
|
@ -751,6 +859,8 @@ export function useThreadStream({
|
|||
messagesRef.current = [];
|
||||
summarizedRef.current = new Set<string>();
|
||||
pendingUsageBaselineMessageIdsRef.current = new Set();
|
||||
setPendingSupersededRunIds(new Set());
|
||||
setPendingSupersededMessageIds(new Set());
|
||||
prevHumanMsgCountRef.current =
|
||||
latestMessageCountsRef.current.humanMessageCount;
|
||||
}, [threadId]);
|
||||
|
|
@ -1011,6 +1121,113 @@ export function useThreadStream({
|
|||
],
|
||||
);
|
||||
|
||||
const regenerateMessage = useCallback(
|
||||
async (
|
||||
threadId: string,
|
||||
messageId: string,
|
||||
supersededMessageIds: string[] = [messageId],
|
||||
) => {
|
||||
if (sendInFlightRef.current || !threadId || !messageId) {
|
||||
return;
|
||||
}
|
||||
sendInFlightRef.current = true;
|
||||
prevHumanMsgCountRef.current = humanMessageCount;
|
||||
pendingUsageBaselineMessageIdsRef.current = new Set(
|
||||
persistedMessages
|
||||
.map(messageIdentity)
|
||||
.filter((id): id is string => Boolean(id)),
|
||||
);
|
||||
setLiveMessagesThreadId(threadId);
|
||||
listeners.current.onSend?.(threadId);
|
||||
let preparedSupersededRunId: string | null = null;
|
||||
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;
|
||||
preparedSupersededRunId = prepared.target_run_id;
|
||||
preparedSupersededMessageIds = supersededMessageIds;
|
||||
setPendingSupersededRunIds((current) => {
|
||||
const next = new Set(current);
|
||||
next.add(prepared.target_run_id);
|
||||
return next;
|
||||
});
|
||||
setPendingSupersededMessageIds((current) => {
|
||||
const next = new Set(current);
|
||||
for (const id of supersededMessageIds) {
|
||||
next.add(id);
|
||||
}
|
||||
return next;
|
||||
});
|
||||
|
||||
await thread.submit(prepared.input, {
|
||||
threadId,
|
||||
checkpoint: prepared.checkpoint,
|
||||
metadata: prepared.metadata,
|
||||
streamSubgraphs: true,
|
||||
streamResumable: true,
|
||||
config: {
|
||||
recursion_limit: 1000,
|
||||
},
|
||||
context: {
|
||||
...context,
|
||||
thinking_enabled: context.mode !== "flash",
|
||||
is_plan_mode: context.mode === "pro" || context.mode === "ultra",
|
||||
subagent_enabled: context.mode === "ultra",
|
||||
reasoning_effort:
|
||||
context.reasoning_effort ??
|
||||
(context.mode === "ultra"
|
||||
? "high"
|
||||
: context.mode === "pro"
|
||||
? "medium"
|
||||
: context.mode === "thinking"
|
||||
? "low"
|
||||
: undefined),
|
||||
thread_id: threadId,
|
||||
},
|
||||
});
|
||||
void queryClient.invalidateQueries({ queryKey: ["thread", threadId] });
|
||||
void queryClient.invalidateQueries({ queryKey: ["threads", "search"] });
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: INFINITE_THREADS_QUERY_KEY_PREFIX,
|
||||
});
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: threadTokenUsageQueryKey(threadId),
|
||||
});
|
||||
} catch (error) {
|
||||
setLiveMessagesThreadId(null);
|
||||
if (preparedSupersededRunId) {
|
||||
const supersededRunId = preparedSupersededRunId;
|
||||
setPendingSupersededRunIds((current) =>
|
||||
removeSetItems(current, [supersededRunId]),
|
||||
);
|
||||
setPendingSupersededMessageIds((current) =>
|
||||
removeSetItems(current, preparedSupersededMessageIds),
|
||||
);
|
||||
}
|
||||
toast.error(getStreamErrorMessage(error));
|
||||
} finally {
|
||||
sendInFlightRef.current = false;
|
||||
}
|
||||
},
|
||||
[context, humanMessageCount, persistedMessages, queryClient, thread],
|
||||
);
|
||||
|
||||
// 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) {
|
||||
|
|
@ -1047,6 +1264,7 @@ export function useThreadStream({
|
|||
thread: mergedThread,
|
||||
pendingUsageMessages,
|
||||
sendMessage,
|
||||
regenerateMessage,
|
||||
isUploading,
|
||||
isHistoryLoading,
|
||||
hasMoreHistory,
|
||||
|
|
@ -1056,11 +1274,12 @@ export function useThreadStream({
|
|||
|
||||
type ThreadHistoryOptions = {
|
||||
enabled?: boolean;
|
||||
pendingSupersededRunIds?: ReadonlySet<string>;
|
||||
};
|
||||
|
||||
export function useThreadHistory(
|
||||
threadId: string,
|
||||
{ enabled = true }: ThreadHistoryOptions = {},
|
||||
{ enabled = true, pendingSupersededRunIds }: ThreadHistoryOptions = {},
|
||||
) {
|
||||
const runs = useThreadRuns(threadId, { enabled });
|
||||
const threadIdRef = useRef(threadId);
|
||||
|
|
@ -1073,7 +1292,20 @@ export function useThreadHistory(
|
|||
const runBeforeSeqRef = useRef<Map<string, number>>(new Map());
|
||||
const loadGenerationRef = useRef(0);
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [messages, setMessages] = useState<Message[]>([]);
|
||||
const [messageRows, setMessageRows] = useState<RunMessage[]>([]);
|
||||
const [appendedMessages, setAppendedMessages] = useState<Message[]>([]);
|
||||
|
||||
const supersededRunIds = useMemo(() => {
|
||||
return getSupersededRunIds(runs.data, pendingSupersededRunIds);
|
||||
}, [pendingSupersededRunIds, runs.data]);
|
||||
|
||||
const messages = useMemo(() => {
|
||||
return buildVisibleHistoryMessages(
|
||||
messageRows,
|
||||
supersededRunIds,
|
||||
appendedMessages,
|
||||
);
|
||||
}, [appendedMessages, messageRows, supersededRunIds]);
|
||||
|
||||
const loadMessages = useCallback(async () => {
|
||||
if (!enabled) {
|
||||
|
|
@ -1139,11 +1371,11 @@ export function useThreadHistory(
|
|||
) {
|
||||
return;
|
||||
}
|
||||
const _messages = result.data
|
||||
.filter((m) => !m.metadata.caller?.startsWith("middleware:"))
|
||||
.map((m) => m.content);
|
||||
setMessages((prev) =>
|
||||
dedupeMessagesByIdentity([..._messages, ...prev]),
|
||||
const _messages = result.data.filter(
|
||||
(m) => !m.metadata.caller?.startsWith("middleware:"),
|
||||
);
|
||||
setMessageRows((prev) =>
|
||||
dedupeRunMessagesByIdentity([..._messages, ...prev]),
|
||||
);
|
||||
const nextBeforeSeq = getNextRunMessagesBeforeSeq(result);
|
||||
if (typeof nextBeforeSeq === "number") {
|
||||
|
|
@ -1197,7 +1429,8 @@ export function useThreadHistory(
|
|||
runBeforeSeqRef.current = new Map();
|
||||
loadingRef.current = false;
|
||||
setLoading(false);
|
||||
setMessages([]);
|
||||
setMessageRows([]);
|
||||
setAppendedMessages([]);
|
||||
}
|
||||
|
||||
if (!enabled) {
|
||||
|
|
@ -1217,7 +1450,7 @@ export function useThreadHistory(
|
|||
}, [enabled, threadId, runs.data, loadMessages]);
|
||||
|
||||
const appendMessages = useCallback((_messages: Message[]) => {
|
||||
setMessages((prev) => {
|
||||
setAppendedMessages((prev) => {
|
||||
return dedupeMessagesByIdentity([...prev, ..._messages]);
|
||||
});
|
||||
}, []);
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ export interface RunMessage {
|
|||
content: Message;
|
||||
metadata: {
|
||||
caller: string;
|
||||
[key: string]: unknown;
|
||||
};
|
||||
created_at: string;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,13 +3,16 @@ import { expect, test } from "vitest";
|
|||
|
||||
import {
|
||||
buildRunMessagesUrl,
|
||||
buildVisibleHistoryMessages,
|
||||
findLatestUnloadedRunIndex,
|
||||
getNextRunMessagesBeforeSeq,
|
||||
getOldestRunMessageSeq,
|
||||
getSupersededRunIds,
|
||||
getSummarizationMiddlewareMessages,
|
||||
getVisibleOptimisticMessages,
|
||||
MAX_CONSECUTIVE_EMPTY_RUN_LOADS,
|
||||
mergeMessages,
|
||||
removeSetItems,
|
||||
runMessagesPageHasMore,
|
||||
shouldAutoContinueOnEmptyRun,
|
||||
} from "@/core/threads/hooks";
|
||||
|
|
@ -355,6 +358,104 @@ test("findLatestUnloadedRunIndex returns -1 when every run is already loaded", (
|
|||
expect(findLatestUnloadedRunIndex(runs, new Set(["R1", "R2"]))).toBe(-1);
|
||||
});
|
||||
|
||||
test("getSupersededRunIds combines completed regenerate metadata with pending ids", () => {
|
||||
const runs = [
|
||||
{
|
||||
run_id: "run-new",
|
||||
status: "success",
|
||||
metadata: { regenerate_from_run_id: "run-old" },
|
||||
},
|
||||
{
|
||||
run_id: "run-normal",
|
||||
status: "success",
|
||||
metadata: {},
|
||||
},
|
||||
] as unknown as Run[];
|
||||
|
||||
expect(getSupersededRunIds(runs, new Set(["run-pending"]))).toEqual(
|
||||
new Set(["run-old", "run-pending"]),
|
||||
);
|
||||
});
|
||||
|
||||
test("getSupersededRunIds ignores failed regenerate runs but keeps pending ids", () => {
|
||||
const runs = [
|
||||
{
|
||||
run_id: "run-error",
|
||||
status: "error",
|
||||
metadata: { regenerate_from_run_id: "run-old" },
|
||||
},
|
||||
{
|
||||
run_id: "run-interrupted",
|
||||
status: "interrupted",
|
||||
metadata: { regenerate_from_run_id: "run-older" },
|
||||
},
|
||||
] as unknown as Run[];
|
||||
|
||||
expect(getSupersededRunIds(runs, new Set(["run-pending"]))).toEqual(
|
||||
new Set(["run-pending"]),
|
||||
);
|
||||
});
|
||||
|
||||
test("removeSetItems removes pending superseded ids after submit failure", () => {
|
||||
expect(
|
||||
removeSetItems(new Set(["run-old", "run-other"]), ["run-old"]),
|
||||
).toEqual(new Set(["run-other"]));
|
||||
});
|
||||
|
||||
test("buildVisibleHistoryMessages filters superseded runs but keeps regenerated run", () => {
|
||||
const oldHuman = {
|
||||
id: "human-1",
|
||||
type: "human",
|
||||
content: "question",
|
||||
} as Message;
|
||||
const oldAi = {
|
||||
id: "ai-old",
|
||||
type: "ai",
|
||||
content: "old answer",
|
||||
} as Message;
|
||||
const newHuman = {
|
||||
id: "human-1",
|
||||
type: "human",
|
||||
content: "question",
|
||||
} as Message;
|
||||
const newAi = {
|
||||
id: "ai-new",
|
||||
type: "ai",
|
||||
content: "new answer",
|
||||
} as Message;
|
||||
const rows: RunMessage[] = [
|
||||
{
|
||||
run_id: "run-old",
|
||||
content: oldHuman,
|
||||
metadata: { caller: "lead_agent" },
|
||||
created_at: "2026-06-18T00:00:00Z",
|
||||
},
|
||||
{
|
||||
run_id: "run-old",
|
||||
content: oldAi,
|
||||
metadata: { caller: "lead_agent" },
|
||||
created_at: "2026-06-18T00:00:01Z",
|
||||
},
|
||||
{
|
||||
run_id: "run-new",
|
||||
content: newHuman,
|
||||
metadata: { caller: "lead_agent" },
|
||||
created_at: "2026-06-18T00:00:02Z",
|
||||
},
|
||||
{
|
||||
run_id: "run-new",
|
||||
content: newAi,
|
||||
metadata: { caller: "lead_agent" },
|
||||
created_at: "2026-06-18T00:00:03Z",
|
||||
},
|
||||
];
|
||||
|
||||
expect(buildVisibleHistoryMessages(rows, new Set(["run-old"]), [])).toEqual([
|
||||
newHuman,
|
||||
newAi,
|
||||
]);
|
||||
});
|
||||
|
||||
test("loading runs in newest-first order and prepending pages yields chronological messages (regression for #3352)", () => {
|
||||
// Simulate backend list_by_thread returning newest first.
|
||||
const runs = [
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue