diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index 79a9663f2..475491042 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -78,6 +78,41 @@ def compute_run_durations(runs) -> dict[str, int]: return durations +def stamp_turn_duration_on_last_ai(messages, run_durations: dict[str, int]) -> None: + """Attach each run's elapsed seconds to that run's final visible AI message only. + + ``turn_duration`` is the run's wall-clock lifetime (``compute_run_durations``), + not model thinking time — it belongs to the run, not to individual messages. + Stamping every AI message made the UI repeat the same number once per + message and let tool-wait time read as thinking latency (#4152). + Middleware-caller messages (e.g. title generation) are skipped so the badge + lands on the assistant's actual final answer. + + Accepts both message shapes that carry ``run_id``: event-store rows, which + wrap the message payload in a ``content`` dict, and flat serialized + checkpoint messages (``/history``), where the payload is the row itself. + + The middleware skip is only effective on the event-store shape: checkpoint + messages replayed on ``/history`` never carry ``metadata.caller`` (they are + plain serialized LangChain messages), so this skip is inert there. That is + not a gap in practice — middleware writes (e.g. title generation) go to + thread metadata, not the ``messages`` channel, so no middleware message + reaches a checkpoint's ``messages`` list to begin with. + """ + stamped: set[str] = set() + for msg in reversed(messages): + rid = msg.get("run_id") + if not rid or rid in stamped or rid not in run_durations: + continue + content = msg.get("content") + payload = content if isinstance(content, dict) else msg + metadata = msg.get("metadata") or {} + is_middleware = str(metadata.get("caller", "")).startswith("middleware:") + if payload.get("type") == "ai" and not is_middleware: + payload.setdefault("additional_kwargs", {})["turn_duration"] = run_durations[rid] + stamped.add(rid) + + # --------------------------------------------------------------------------- # Request / response models # --------------------------------------------------------------------------- @@ -1102,14 +1137,7 @@ async def list_thread_messages( run_durations = compute_run_durations(runs) if run_durations: - for msg in messages: - content = msg.get("content", {}) - if isinstance(content, dict) and content.get("type") == "ai": - rid = msg.get("run_id") - if rid and rid in run_durations: - if "additional_kwargs" not in content: - content["additional_kwargs"] = {} - content["additional_kwargs"]["turn_duration"] = run_durations[rid] + stamp_turn_duration_on_last_ai(messages, run_durations) return messages @@ -1355,16 +1383,8 @@ async def list_run_messages( record = await run_mgr.get(run_id) if record: durations = compute_run_durations([record]) - duration = durations.get(run_id) - if duration is not None: - for msg in reversed(data): - content = msg.get("content") - metadata = msg.get("metadata", {}) - is_middleware = str(metadata.get("caller", "")).startswith("middleware:") - if isinstance(content, dict) and content.get("type") == "ai" and not is_middleware: - if "additional_kwargs" not in content: - content["additional_kwargs"] = {} - content["additional_kwargs"]["turn_duration"] = duration + if durations: + stamp_turn_duration_on_last_ai(data, durations) return {"data": data, "has_more": has_more} diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index acc67c214..df45306a9 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -1312,11 +1312,6 @@ async def update_thread_state(thread_id: str, body: ThreadStateUpdateRequest, re ) -def _ai_message_lacks_duration(message: dict[str, Any]) -> bool: - additional_kwargs = message.get("additional_kwargs") - return message.get("type") == "ai" and (not isinstance(additional_kwargs, dict) or "turn_duration" not in additional_kwargs) - - def _checkpoint_run_durations(metadata: Any) -> dict[str, int]: raw_durations = metadata.get("run_durations") if isinstance(metadata, dict) else None if not isinstance(raw_durations, dict): @@ -1324,16 +1319,6 @@ def _checkpoint_run_durations(metadata: Any) -> dict[str, int]: return {run_id: duration_seconds for run_id, duration_seconds in raw_durations.items() if valid_duration_entry(run_id, duration_seconds)} -def _set_message_turn_duration(message: dict[str, Any], run_id: str, run_durations: dict[str, int]) -> None: - if message.get("type") != "ai" or run_id not in run_durations: - return - additional_kwargs = message.get("additional_kwargs") - if not isinstance(additional_kwargs, dict): - additional_kwargs = {} - message["additional_kwargs"] = additional_kwargs - additional_kwargs.setdefault("turn_duration", run_durations[run_id]) - - @router.post("/{thread_id}/history", response_model=list[HistoryEntry]) @require_permission("threads", "read", owner_check=True) async def get_thread_history( @@ -1381,11 +1366,14 @@ async def get_thread_history( if messages: serialized_msgs = serialize_channel_values_for_api({"messages": messages}).get("messages", []) try: + from app.gateway.routers.thread_runs import stamp_turn_duration_on_last_ai + # Human messages define turn boundaries. New checkpoints # carry the completed turns' durations in metadata, so the # messages channel stays unchanged. checkpoint_run_durations = _checkpoint_run_durations(metadata) current_turn_run_id = None + turn_run_ids: set[str] = set() for msg in serialized_msgs: if msg.get("type") == "human": additional_kwargs = msg.get("additional_kwargs") @@ -1399,12 +1387,21 @@ async def get_thread_history( continue msg.setdefault("run_id", current_turn_run_id) - _set_message_turn_duration(msg, current_turn_run_id, checkpoint_run_durations) + if msg.get("type") == "ai": + turn_run_ids.add(current_turn_run_id) - # Legacy checkpoints without duration metadata are - # correlated once via event-store + run-manager, then - # upgraded by a metadata-only checkpoint write. - if any(_ai_message_lacks_duration(msg) for msg in serialized_msgs): + # Stamp each run's duration on its last AI message only, + # same as the live message endpoints — never every AI + # message in a multi-message turn (#4152). + stamp_turn_duration_on_last_ai(serialized_msgs, checkpoint_run_durations) + + # Runs referenced by this checkpoint's AI messages but + # absent from checkpoint metadata are either legacy + # (never migrated) or just completed. Correlate once via + # event-store + run-manager, then upgrade by a + # metadata-only checkpoint write. + missing_run_ids = turn_run_ids - set(checkpoint_run_durations) + if missing_run_ids: from app.gateway.deps import get_run_event_store, get_run_manager from app.gateway.routers.thread_runs import compute_run_durations from deerflow.runtime.runs.worker import persist_run_durations @@ -1439,7 +1436,8 @@ async def get_thread_history( run_id = msg_to_run.get(msg.get("id")) or current_turn_run_id if run_id: msg["run_id"] = run_id - _set_message_turn_duration(msg, run_id, run_durations) + + stamp_turn_duration_on_last_ai(serialized_msgs, run_durations) # Intentional, best-effort write-on-read migration: # persist legacy metadata after the response so the diff --git a/backend/tests/test_thread_run_messages_pagination.py b/backend/tests/test_thread_run_messages_pagination.py index b1773e957..f6669bbb5 100644 --- a/backend/tests/test_thread_run_messages_pagination.py +++ b/backend/tests/test_thread_run_messages_pagination.py @@ -265,8 +265,9 @@ def test_join_store_only_run_allowed_with_cross_process_bridge(): assert "event: end" in response.text -def test_list_run_messages_injects_turn_duration(): - """Verify that list_run_messages injects turn_duration into ALL AI messages for the run.""" +def test_list_run_messages_injects_turn_duration_only_on_last_ai(): + """#4152: turn_duration is the run's wall-clock elapsed and belongs to the run, + so only the run's final AI message carries it — not every AI message.""" from unittest.mock import AsyncMock from deerflow.runtime import RunRecord @@ -300,13 +301,13 @@ def test_list_run_messages_injects_turn_duration(): data = response.json()["data"] assert "turn_duration" not in data[0]["content"].get("additional_kwargs", {}) - - assert data[1]["content"]["additional_kwargs"]["turn_duration"] == 5 + assert "turn_duration" not in data[1]["content"].get("additional_kwargs", {}) assert data[2]["content"]["additional_kwargs"]["turn_duration"] == 5 -def test_list_thread_messages_injects_turn_duration(): - """Verify that list_thread_messages injects turn_duration into the inner content.""" +def test_list_run_messages_turn_duration_skips_middleware_tail(): + """The duration badge lands on the assistant's final answer, not on a + trailing middleware-caller message (e.g. title generation).""" from unittest.mock import AsyncMock from deerflow.runtime import RunRecord @@ -320,9 +321,101 @@ def test_list_thread_messages_injects_turn_duration(): created_at="2026-06-20T10:00:00Z", updated_at="2026-06-20T10:00:05Z", ) + + rows = [ + {"seq": 1, "run_id": "run-1", "content": {"type": "ai", "text": "Response"}}, + {"seq": 2, "run_id": "run-1", "content": {"type": "ai", "text": "Title"}, "metadata": {"caller": "middleware:title"}}, + ] + + event_store = _make_event_store(rows) + run_manager = AsyncMock() + run_manager.get.return_value = mock_run + app = _make_app(event_store=event_store, run_manager=run_manager) + + with TestClient(app) as client: + response = client.get("/api/threads/thread-1/runs/run-1/messages") + + assert response.status_code == 200 + data = response.json()["data"] + + assert data[0]["content"]["additional_kwargs"]["turn_duration"] == 5 + assert "turn_duration" not in data[1]["content"].get("additional_kwargs", {}) + + +def test_list_thread_messages_injects_turn_duration_once_per_run(): + """#4152: in the thread feed each run contributes exactly one duration badge — + on its last AI message — even when the run produced several AI messages and + other runs interleave in the same thread.""" + from unittest.mock import AsyncMock + + from deerflow.runtime import RunRecord + + def _run(run_id: str, seconds: int) -> RunRecord: + return RunRecord( + run_id=run_id, + thread_id="thread-1", + assistant_id=None, + status="success", + on_disconnect="cancel", + created_at="2026-06-20T10:00:00Z", + updated_at=f"2026-06-20T10:00:{seconds:02d}Z", + ) + rows = [ {"seq": 1, "run_id": "run-1", "content": {"type": "human", "text": "Hello"}}, - {"seq": 2, "run_id": "run-1", "content": {"type": "ai", "text": "Response"}}, + {"seq": 2, "run_id": "run-1", "content": {"type": "ai", "text": "Before tools"}}, + {"seq": 3, "run_id": "run-1", "content": {"type": "ai", "text": "Final answer"}}, + {"seq": 4, "run_id": "run-2", "content": {"type": "human", "text": "Again"}}, + {"seq": 5, "run_id": "run-2", "content": {"type": "ai", "text": "Second answer"}}, + ] + + event_store = MagicMock() + event_store.list_messages = AsyncMock(return_value=rows) + + run_manager = AsyncMock() + run_manager.list_by_thread = AsyncMock(return_value=[_run("run-1", 5), _run("run-2", 9)]) + + feedback_repo = MagicMock() + feedback_repo.list_by_thread_grouped = AsyncMock(return_value={}) + + app = _make_app(event_store=event_store, run_manager=run_manager) + app.state.feedback_repo = feedback_repo + + with TestClient(app) as client: + response = client.get("/api/threads/thread-1/messages") + + assert response.status_code == 200 + data = response.json() + + assert "turn_duration" not in data[0].get("content", {}).get("additional_kwargs", {}) + assert "turn_duration" not in data[1]["content"].get("additional_kwargs", {}) + assert data[2]["content"]["additional_kwargs"]["turn_duration"] == 5 + assert "turn_duration" not in data[3].get("content", {}).get("additional_kwargs", {}) + assert data[4]["content"]["additional_kwargs"]["turn_duration"] == 9 + + +def test_list_thread_messages_turn_duration_skips_middleware_tail(): + """The thread feed gained the middleware skip through the same + ``stamp_turn_duration_on_last_ai`` helper as the per-run endpoint — before + that it stamped every AI message, middleware included. Lock the contract + on this path too, not just ``list_run_messages``.""" + from unittest.mock import AsyncMock + + from deerflow.runtime import RunRecord + + mock_run = RunRecord( + run_id="run-1", + thread_id="thread-1", + assistant_id=None, + status="success", + on_disconnect="cancel", + created_at="2026-06-20T10:00:00Z", + updated_at="2026-06-20T10:00:05Z", + ) + + rows = [ + {"seq": 1, "run_id": "run-1", "content": {"type": "ai", "text": "Response"}}, + {"seq": 2, "run_id": "run-1", "content": {"type": "ai", "text": "Title"}, "metadata": {"caller": "middleware:title"}}, ] event_store = MagicMock() @@ -343,5 +436,5 @@ def test_list_thread_messages_injects_turn_duration(): assert response.status_code == 200 data = response.json() - assert "turn_duration" not in data[0].get("content", {}).get("additional_kwargs", {}) - assert data[1]["content"]["additional_kwargs"]["turn_duration"] == 5 + assert data[0]["content"]["additional_kwargs"]["turn_duration"] == 5 + assert "turn_duration" not in data[1]["content"].get("additional_kwargs", {}) diff --git a/backend/tests/test_threads_router.py b/backend/tests/test_threads_router.py index e2c0f5f2f..9c2269925 100644 --- a/backend/tests/test_threads_router.py +++ b/backend/tests/test_threads_router.py @@ -1241,7 +1241,67 @@ def test_get_thread_history_associates_tool_messages_from_checkpoint_turn() -> N history_messages = response.json()[0]["values"]["messages"] assert [message.get("run_id") for message in history_messages[1:]] == ["run-1", "run-1", "run-1"] - assert [message["additional_kwargs"]["turn_duration"] for message in history_messages if message["type"] == "ai"] == [4, 4] + # #4152: turn_duration belongs to the run, not to every AI message in it — + # only the run's last AI message ("Done") gets stamped, not the + # tool-calling one that precedes it. + ai_messages = [message for message in history_messages if message["type"] == "ai"] + assert len(ai_messages) == 2 + assert "turn_duration" not in (ai_messages[0].get("additional_kwargs") or {}) + assert ai_messages[1]["additional_kwargs"]["turn_duration"] == 4 + + +def test_get_thread_history_fast_path_skips_runs_already_in_checkpoint_metadata() -> None: + """A checkpoint can carry metadata for some runs but not others (e.g. the + latest run just completed and hasn't been persisted yet). Only the + missing run should trigger the event-store/run-manager correlation + fallback; the already-migrated run must not re-query it.""" + app, _store, checkpointer = _build_thread_app() + thread_id = "history-partial-migration" + messages = [ + HumanMessage(id="human-1", content="First", additional_kwargs={"run_id": "run-migrated"}), + AIMessage(id="ai-1", content="First answer"), + HumanMessage(id="human-2", content="Second", additional_kwargs={"run_id": "run-pending"}), + AIMessage(id="ai-2", content="Second answer"), + ] + asyncio.run( + _write_checkpoint( + checkpointer, + thread_id, + "checkpoint-partial", + messages, + step=1, + metadata={"run_durations": {"run-migrated": 4}}, + ) + ) + + async def list_by_thread(_: str) -> list[SimpleNamespace]: + return [ + SimpleNamespace( + run_id="run-pending", + created_at="2026-07-05T00:00:00+00:00", + updated_at="2026-07-05T00:00:06+00:00", + ), + ] + + list_messages_calls: list[str] = [] + + async def list_messages(thread: str, *, limit: int) -> list[dict]: + list_messages_calls.append(thread) + return [] + + app.state.run_manager = SimpleNamespace(list_by_thread=list_by_thread) + app.state.run_event_store = SimpleNamespace(list_messages=list_messages) + + with TestClient(app) as client: + response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10}) + + assert response.status_code == 200, response.text + history_messages = response.json()[0]["values"]["messages"] + assert history_messages[1]["additional_kwargs"]["turn_duration"] == 4 + assert history_messages[3]["additional_kwargs"]["turn_duration"] == 6 + # The fallback still runs (run-pending was missing), but it is the only + # reason it ran — proven by it firing exactly once, not skipped entirely. + assert list_messages_calls == [thread_id] def test_get_thread_history_backfills_legacy_durations_with_exact_event_run_id() -> None: @@ -1291,11 +1351,56 @@ def test_get_thread_history_backfills_legacy_durations_with_exact_event_run_id() assert latest.metadata["run_durations"] == {"boundary-run": 3, "exact-run": 7} -def test_ai_message_lacks_duration_only_for_unannotated_ai_messages() -> None: - assert threads._ai_message_lacks_duration({"type": "ai"}) - assert threads._ai_message_lacks_duration({"type": "ai", "additional_kwargs": []}) - assert not threads._ai_message_lacks_duration({"type": "tool"}) - assert not threads._ai_message_lacks_duration({"type": "ai", "additional_kwargs": {"turn_duration": 0}}) +def test_get_thread_history_injects_turn_duration_once_per_run() -> None: + """#4152: ``/history`` replays checkpoint messages on reload, so it must + stamp ``turn_duration`` the same way the message endpoints do — once per + run, on that run's last AI message — even when a run produced several AI + messages.""" + from unittest.mock import AsyncMock, MagicMock + + from deerflow.runtime import RunRecord + + def _run(run_id: str, seconds: int) -> RunRecord: + return RunRecord( + run_id=run_id, + thread_id="history-durations", + assistant_id=None, + status="success", + on_disconnect="cancel", + created_at="2026-06-20T10:00:00Z", + updated_at=f"2026-06-20T10:00:{seconds:02d}Z", + ) + + app, _store, checkpointer = _build_thread_app() + thread_id = "history-durations" + + messages = [ + HumanMessage(id="human-1", content="Hello", additional_kwargs={"run_id": "run-1"}), + AIMessage(id="ai-1a", content="Before tools"), + AIMessage(id="ai-1b", content="Final answer"), + HumanMessage(id="human-2", content="Again", additional_kwargs={"run_id": "run-2"}), + AIMessage(id="ai-2", content="Second answer"), + ] + asyncio.run(_write_checkpoint(checkpointer, thread_id, "0001", messages, step=1)) + + run_manager = AsyncMock() + run_manager.list_by_thread = AsyncMock(return_value=[_run("run-1", 5), _run("run-2", 9)]) + event_store = MagicMock() + event_store.list_messages = AsyncMock(return_value=[]) + app.state.run_manager = run_manager + app.state.run_event_store = event_store + + with TestClient(app) as client: + response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10}) + + assert response.status_code == 200, response.text + replayed = response.json()[0]["values"]["messages"] + + assert "turn_duration" not in (replayed[0].get("additional_kwargs") or {}) + assert "turn_duration" not in (replayed[1].get("additional_kwargs") or {}) + assert replayed[2]["additional_kwargs"]["turn_duration"] == 5 + assert "turn_duration" not in (replayed[3].get("additional_kwargs") or {}) + assert replayed[4]["additional_kwargs"]["turn_duration"] == 9 # ── branch threads from completed assistant turns ─────────────────────────────