fix(runs): stamp elapsed duration once per run across message APIs (#4163)

* fix(runs): show run duration once per run and label it as elapsed work

turn_duration is the run's wall-clock lifetime (updated_at - created_at),
but both message endpoints stamped it onto every AI message of the run
and the frontend rendered each stamp as 'Thought for X seconds' — the
same number repeated per message, and tool-wait time presented as model
thinking latency (#4152).

- share one stamping helper between list_messages and list_run_messages:
  only the run's final non-middleware AI message carries turn_duration
  (list_run_messages already tried to do this but iterated reversed()
  without stopping at the first match)
- completed-state trigger copy becomes 'Worked for X seconds' since the
  number includes tool execution, not just thinking

Fixes #4152

* fix(threads): stamp turn_duration once per run in /history replay too

get_thread_history stamped turn_duration on every AI message of a run
instead of only the last one, so a run with several AI messages (e.g.
a tool call followed by a final answer) showed the same duration badge
more than once on reload.

Rebasing onto main picked up #4118, which replaced the /history
duration path with a checkpoint-metadata fast path plus an
event-store/run-manager fallback for legacy checkpoints - both of
which stamped every AI message per run_id, reintroducing this same
bug on both paths. This commit folds the once-per-run stamping
(stamp_turn_duration_on_last_ai) into both paths instead of only the
legacy one, and narrows the fast-path/fallback boundary to a per-run
completeness check (turn_run_ids - checkpoint_run_durations) instead
of a per-message one, so a run whose duration is already in checkpoint
metadata never re-triggers the event-store correlation fallback just
because its non-final AI messages correctly have no turn_duration.

* fix(frontend): stop caching a client-measured turn_duration per message

A run can produce several standalone content-only AI messages (e.g. a
subagent handoff), each briefly becoming the newest message and each
mounting its own Reasoning timer via the shared turnStartTime. Wiring
onTurnDurationChange let that per-message timer cache and display a
"Worked for X seconds" badge the moment a later message superseded it
- with a different, premature number, since the timer only measured
that message's own streaming window, not the run's total elapsed
time. That reintroduced the #4152 duplicate-badge bug through a path
the backend fix (once-per-run stamping) doesn't cover.

cachedDuration now only ever mirrors a backend-confirmed
turn_duration (already handled by the other effect here), so a
superseded message can no longer display a stale or wrong duration.

* fix(frontend): use 'Worked for' framing even without a persisted duration

The completed-state copy diverged depending on whether turn_duration
had landed yet: "Thought for a few seconds" (no duration) vs. "Worked
for X seconds" (duration present). Both describe the same completed
run, just with or without a persisted number, so the framing
shouldn't disagree on what happened.

* docs(runs): note why the middleware skip is inert on /history, add test

stamp_turn_duration_on_last_ai's middleware-caller skip only has an
effect on the event-store message shape (metadata.caller); checkpoint
messages replayed on /history never carry that field. Document why
that's not a gap - middleware writes go to thread metadata, not the
messages channel, so no middleware message ever reaches a
checkpoint's messages list.

Also lock the middleware-skip contract on list_thread_messages with
its own regression test, mirroring the existing one for
list_run_messages - the thread feed only gained this skip through the
same shared helper, and previously stamped every AI message,
middleware included.

* chore(frontend): retain current run duration UI

The frontend behavior was superseded by #4348; keep the latest main implementation while retaining this PR's backend fix.

* test(frontend): retain current reasoning coverage

Align the old PR test with the frontend implementation already merged in #4348.

---------

Co-authored-by: fancyboi999 <fancyboi999@users.noreply.github.com>
This commit is contained in:
Xinmin Zeng 2026-07-29 12:47:30 +08:00 committed by GitHub
parent 4e66acbbb4
commit e56481d9e3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 270 additions and 54 deletions

View File

@ -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}

View File

@ -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

View File

@ -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", {})

View File

@ -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 ─────────────────────────────