From 45722170380d585761221e64d74efe7cd875516b Mon Sep 17 00:00:00 2001 From: Chetan Sharma Date: Sun, 21 Jun 2026 07:05:45 +0530 Subject: [PATCH] feat: persist AI turn duration in backend and UI (#3663) * feat: persist AI turn duration in backend and UI * chore: restore uv.lock to match main * refactor(frontend): use Math.floor instead of Math.ceil for reasoning timer --- backend/app/gateway/routers/thread_runs.py | 49 +++++++++++ backend/app/gateway/routers/threads.py | 42 ++++++++- .../test_thread_run_messages_pagination.py | 88 ++++++++++++++++++- .../src/components/ai-elements/reasoning.tsx | 9 +- .../workspace/messages/message-list-item.tsx | 70 +++++++++++++-- frontend/src/core/threads/hooks.ts | 79 ++++++++++++++--- 6 files changed, 310 insertions(+), 27 deletions(-) diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index 381252db2..42e8ead23 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -32,6 +32,24 @@ router = APIRouter(prefix="/api/threads", tags=["runs"]) REGENERATE_HISTORY_SCAN_LIMIT = 200 +def compute_run_durations(runs) -> dict[str, int]: + """Map run_id -> duration in seconds from run timestamps.""" + from datetime import datetime + + durations: dict[str, int] = {} + for r in runs: + if r.created_at and r.updated_at: + try: + created = datetime.fromisoformat(r.created_at.replace("Z", "+00:00")) + updated = datetime.fromisoformat(r.updated_at.replace("Z", "+00:00")) + # Note: updated_at - created_at represents the row's total lifetime, + # which can slightly overshoot the actual AI turn end if the row is mutated later. + durations[r.run_id] = int((updated - created).total_seconds()) + except Exception: + logger.warning("Failed to parse timestamps for run %s", r.run_id, exc_info=True) + return durations + + # --------------------------------------------------------------------------- # Request / response models # --------------------------------------------------------------------------- @@ -657,6 +675,20 @@ async def list_thread_messages( else: msg["feedback"] = None + run_mgr = get_run_manager(request) + runs = await run_mgr.list_by_thread(thread_id, user_id=user_id) + run_durations = compute_run_durations(runs) + + 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] + return messages @@ -683,6 +715,23 @@ async def list_run_messages( after_seq=after_seq, ) data, has_more = trim_run_message_page(rows, limit=limit, after_seq=after_seq) + + if data: + run_mgr = get_run_manager(request) + 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 + return {"data": data, "has_more": has_more} diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index 0570f0519..a2549847b 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -640,7 +640,47 @@ async def get_thread_history(thread_id: str, body: ThreadHistoryRequest, request if is_latest_checkpoint: messages = channel_values.get("messages") if messages: - values["messages"] = serialize_channel_values_for_api({"messages": messages}).get("messages", []) + serialized_msgs = serialize_channel_values_for_api({"messages": messages}).get("messages", []) + try: + from app.gateway.deps import get_run_event_store, get_run_manager + from app.gateway.routers.thread_runs import compute_run_durations + + run_mgr = get_run_manager(request) + event_store = get_run_event_store(request) + + runs = await run_mgr.list_by_thread(thread_id) + + # FIXME: Fetching limit=1000 silently drops durations for messages older than the cap on long threads. + # We do this full fetch because raw LangGraph messages lack a native run_id link. + + events = await event_store.list_messages(thread_id, limit=1000) + + if runs and serialized_msgs: + # 1. Map each run_id to its actual duration + run_durations = compute_run_durations(runs) + + # 2. Map every message id directly to its parent run_id + msg_to_run = {} + for e in events: + content = e.get("content", {}) + if isinstance(content, dict) and content.get("type") == "ai" and "id" in content: + msg_to_run[content["id"]] = e["run_id"] + + # 3. Inject the exact correct duration into each AI message + for msg in serialized_msgs: + if msg.get("type") == "ai": + msg_id = msg.get("id") + run_id = msg_to_run.get(msg_id) + if run_id and run_id in run_durations: + if "additional_kwargs" not in msg: + msg["additional_kwargs"] = {} + msg["additional_kwargs"]["turn_duration"] = run_durations[run_id] + + except Exception: + logger.warning("Failed to inject turn_duration for thread %s", thread_id, exc_info=True) + + values["messages"] = serialized_msgs + is_latest_checkpoint = False # Derive next tasks diff --git a/backend/tests/test_thread_run_messages_pagination.py b/backend/tests/test_thread_run_messages_pagination.py index 6d36001ae..feae377a2 100644 --- a/backend/tests/test_thread_run_messages_pagination.py +++ b/backend/tests/test_thread_run_messages_pagination.py @@ -25,8 +25,10 @@ def _make_app(event_store=None, run_manager=None): if event_store is not None: app.state.run_event_store = event_store - if run_manager is not None: - app.state.run_manager = run_manager + if run_manager is None: + run_manager = AsyncMock() + run_manager.get.return_value = None + app.state.run_manager = run_manager return app @@ -231,3 +233,85 @@ def test_stream_store_only_run_returns_409(): assert response.status_code == 409 assert "not active on this worker" in response.json()["detail"] + + +def test_list_run_messages_injects_turn_duration(): + """Verify that list_run_messages injects turn_duration into ALL AI messages for the run.""" + from unittest.mock import AsyncMock + + from deerflow.runtime import RunRecord + + # Mock a run record that took exactly 5 seconds + 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": "human", "text": "Hello"}}, + {"seq": 2, "run_id": "run-1", "content": {"type": "ai", "text": "Thinking..."}}, + {"seq": 3, "run_id": "run-1", "content": {"type": "ai", "text": "Response"}}, + ] + + 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 "turn_duration" not in data[0]["content"].get("additional_kwargs", {}) + + assert data[1]["content"]["additional_kwargs"]["turn_duration"] == 5 + 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.""" + 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": "human", "text": "Hello"}}, + {"seq": 2, "run_id": "run-1", "content": {"type": "ai", "text": "Response"}}, + ] + + event_store = MagicMock() + event_store.list_messages = AsyncMock(return_value=rows) + + run_manager = AsyncMock() + run_manager.list_by_thread = AsyncMock(return_value=[mock_run]) + + 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 data[1]["content"]["additional_kwargs"]["turn_duration"] == 5 diff --git a/frontend/src/components/ai-elements/reasoning.tsx b/frontend/src/components/ai-elements/reasoning.tsx index d23b213e8..aea85e59c 100644 --- a/frontend/src/components/ai-elements/reasoning.tsx +++ b/frontend/src/components/ai-elements/reasoning.tsx @@ -39,6 +39,7 @@ export type ReasoningProps = ComponentProps & { onOpenChange?: (open: boolean) => void; duration?: number; startTimeProp?: number | null; + onTurnDurationChange?: (duration: number | undefined) => void; }; const AUTO_CLOSE_DELAY = 1000; @@ -53,6 +54,7 @@ export const Reasoning = memo( onOpenChange, duration: durationProp, startTimeProp, + onTurnDurationChange, children, ...props }: ReasoningProps) => { @@ -61,9 +63,10 @@ export const Reasoning = memo( defaultProp: defaultOpen, onChange: onOpenChange, }); - const [duration, setDuration] = useControllableState({ + const [duration, setDuration] = useControllableState({ prop: durationProp, defaultProp: undefined, + onChange: onTurnDurationChange, }); const [hasAutoClosed, setHasAutoClosed] = useState(false); @@ -81,7 +84,7 @@ export const Reasoning = memo( setStartTime(Date.now()); } } else if (startTime !== null) { - setDuration(Math.ceil((Date.now() - startTime) / MS_IN_S)); + setDuration(Math.floor((Date.now() - startTime) / MS_IN_S)); setStartTime(null); } }, [isStreaming, startTimeProp, startTime, setDuration]); @@ -135,7 +138,7 @@ const LiveTimer = ({ startTime }: { startTime: number }) => { const [elapsed, setElapsed] = useState(0); useEffect(() => { - const calculateElapsed = () => Math.ceil((Date.now() - startTime) / 1000); + const calculateElapsed = () => Math.floor((Date.now() - startTime) / 1000); setElapsed(calculateElapsed()); const interval = setInterval(() => { diff --git a/frontend/src/components/workspace/messages/message-list-item.tsx b/frontend/src/components/workspace/messages/message-list-item.tsx index 425998af8..695163580 100644 --- a/frontend/src/components/workspace/messages/message-list-item.tsx +++ b/frontend/src/components/workspace/messages/message-list-item.tsx @@ -210,6 +210,8 @@ function MessageImage({ ); } +const clientTurnDurations = new Map(); + function MessageContent_({ className, message, @@ -225,6 +227,45 @@ function MessageContent_({ }) { const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading); const isHuman = message.type === "human"; + const rawTurnDuration = message.additional_kwargs?.turn_duration as + | number + | undefined; + + const [cachedDuration, setCachedDuration] = useState( + () => + message.id + ? clientTurnDurations.get(`${threadId}:${message.id}`) + : undefined, + ); + const turnDuration = rawTurnDuration ?? cachedDuration; + + useEffect(() => { + if (rawTurnDuration !== undefined && message.id) { + clientTurnDurations.set(`${threadId}:${message.id}`, rawTurnDuration); + setCachedDuration(rawTurnDuration); + } + }, [rawTurnDuration, message.id]); + + const handleDurationChange = useCallback( + (d: number | undefined) => { + if (d !== undefined && message.id) { + clientTurnDurations.set(`${threadId}:${message.id}`, d); + setCachedDuration(d); + } + }, + [message.id], + ); + + useEffect(() => { + return () => { + for (const key of clientTurnDurations.keys()) { + if (key.startsWith(`${threadId}:`)) { + clientTurnDurations.delete(key); + } + } + }; + }, [threadId]); + const [wasLoading, setWasLoading] = useState(isLoading); useEffect(() => { if (isLoading) setWasLoading(true); @@ -299,7 +340,12 @@ function MessageContent_({ if (!isHuman && reasoningContent && !rawContent) { return ( - + {reasoningContent} @@ -334,14 +380,20 @@ function MessageContent_({ return ( {filesList} - {!isHuman && (!!reasoningContent || wasLoading) && ( - - - {reasoningContent && ( - {reasoningContent} - )} - - )} + {!isHuman && + (!!reasoningContent || wasLoading || turnDuration !== undefined) && ( + + + {reasoningContent && ( + {reasoningContent} + )} + + )}