mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-30 01:46:01 +00:00
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:
parent
4e66acbbb4
commit
e56481d9e3
@ -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}
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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", {})
|
||||
|
||||
@ -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 ─────────────────────────────
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user