From 0dc900126c21608d3aa197c06b0dc915e8e0f09c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=AD=A6=E9=98=B3?= <wuy398@chinaunicom.cn>
Date: Fri, 8 May 2026 21:15:39 +0800
Subject: [PATCH] Carry elapsed time through replay tool parts

---
 .../assistant-ui/activity-panel.tsx           | 13 ++++---
 .../components/assistant-ui/thread-list.tsx   | 38 +++++++++++++++++--
 .../app/components/assistant-ui/thread.tsx    | 20 ++++++++--
 3 files changed, 58 insertions(+), 13 deletions(-)

diff --git a/frontend/app/components/assistant-ui/activity-panel.tsx b/frontend/app/components/assistant-ui/activity-panel.tsx
index 8e77680..e8d0454 100644
--- a/frontend/app/components/assistant-ui/activity-panel.tsx
+++ b/frontend/app/components/assistant-ui/activity-panel.tsx
@@ -585,7 +585,7 @@ function collectActivityItems(messages: readonly MessageState[]) {
 				});
 			}
 			if (part.type === "tool-call") {
-				const input = stripStageNote(part.args);
+				const input = stripClawToolMetadata(part.args);
 				const stageNote = getStageNote(part.args);
 				const argsText = formatToolArgs(input);
 				if (stageNote) {
@@ -632,12 +632,13 @@ function getStageNote(value: unknown) {
 	return typeof note === "string" ? note.trim() : "";
 }
 
-function stripStageNote(value: unknown): unknown {
+function stripClawToolMetadata(value: unknown): unknown {
 	if (!value || typeof value !== "object" || Array.isArray(value)) return value;
-	const { __claw_stage_note: _stageNote, ...rest } = value as Record<
-		string,
-		unknown
-	>;
+	const {
+		__claw_stage_note: _stageNote,
+		__claw_elapsed_ms: _elapsedMs,
+		...rest
+	} = value as Record<string, unknown>;
 	return rest;
 }
 
diff --git a/frontend/app/components/assistant-ui/thread-list.tsx b/frontend/app/components/assistant-ui/thread-list.tsx
index 9ba785a..ed7217d 100644
--- a/frontend/app/components/assistant-ui/thread-list.tsx
+++ b/frontend/app/components/assistant-ui/thread-list.tsx
@@ -425,10 +425,11 @@ export function toReplayRepository(
 					: toolParts
 				: [{ type: "text", text: content }]
 		) as UIMessage["parts"];
+		const replayParts = withReplayElapsedMs(parts, elapsedMs);
 		const uiMessage: UIMessage = {
 			id,
 			role: message.role,
-			parts,
+			parts: replayParts,
 			...(message.role === "assistant"
 				? {
 						metadata: {
@@ -442,7 +443,7 @@ export function toReplayRepository(
 			id,
 			role: message.role,
 			createdAt: new Date(),
-			content: toThreadMessageContent(parts, elapsedMs),
+			content: toThreadMessageContent(replayParts, elapsedMs),
 			...(message.role === "assistant"
 				? {
 						status: { type: "complete", reason: "stop" },
@@ -797,6 +798,28 @@ function attachStageNote(input: unknown, stageNote?: string) {
 	return { value: input, __claw_stage_note: note };
 }
 
+function withReplayElapsedMs(parts: UIMessage["parts"], elapsedMs?: number) {
+	if (elapsedMs === undefined) return parts;
+	return parts.map((part) => {
+		if (part.type === "reasoning") return { ...part, elapsedMs };
+		if (part.type === "dynamic-tool") {
+			return {
+				...part,
+				elapsedMs,
+				input: attachReplayElapsed(part.input ?? {}, elapsedMs),
+			};
+		}
+		return part;
+	}) as UIMessage["parts"];
+}
+
+function attachReplayElapsed(input: unknown, elapsedMs: number) {
+	if (input && typeof input === "object" && !Array.isArray(input)) {
+		return { ...input, __claw_elapsed_ms: elapsedMs };
+	}
+	return { value: input, __claw_elapsed_ms: elapsedMs };
+}
+
 function toThreadMessageContent(parts: UIMessage["parts"], elapsedMs?: number) {
 	return parts.map((part) => {
 		const elapsedPayload = elapsedMs === undefined ? {} : { elapsedMs };
@@ -808,7 +831,7 @@ function toThreadMessageContent(parts: UIMessage["parts"], elapsedMs?: number) {
 				toolName: part.toolName,
 				toolCallId: part.toolCallId,
 				args: part.input ?? {},
-				argsText: JSON.stringify(part.input ?? {}),
+				argsText: JSON.stringify(stripReplayElapsed(part.input ?? {})),
 				...elapsedPayload,
 				...(part.state === "output-available" ? { result: part.output } : {}),
 			};
@@ -817,6 +840,15 @@ function toThreadMessageContent(parts: UIMessage["parts"], elapsedMs?: number) {
 	});
 }
 
+function stripReplayElapsed(value: unknown): unknown {
+	if (!value || typeof value !== "object" || Array.isArray(value)) return value;
+	const { __claw_elapsed_ms: _elapsedMs, ...rest } = value as Record<
+		string,
+		unknown
+	>;
+	return rest;
+}
+
 function parseToolInput(value: unknown) {
 	if (typeof value !== "string") return value ?? {};
 	try {
diff --git a/frontend/app/components/assistant-ui/thread.tsx b/frontend/app/components/assistant-ui/thread.tsx
index 57df736..592f9ca 100644
--- a/frontend/app/components/assistant-ui/thread.tsx
+++ b/frontend/app/components/assistant-ui/thread.tsx
@@ -1316,10 +1316,7 @@ const ActivityChainSummary: FC<{
 	);
 	const storedElapsedMs = useAuiState((s) => {
 		const value = readMetadataValue(s.message.metadata, "elapsedMs");
-		const contentValue = findNumberInActivityParts(
-			s.message.content,
-			"elapsedMs",
-		);
+		const contentValue = findElapsedMsInActivityParts(s.message.content);
 		if (typeof value !== "number" || !Number.isFinite(value)) {
 			return contentValue;
 		}
@@ -1397,6 +1394,21 @@ function findNumberInActivityParts(
 	return null;
 }
 
+function findElapsedMsInActivityParts(content: readonly ActivitySummaryPart[]) {
+	for (const part of content) {
+		if (typeof part.elapsedMs === "number" && Number.isFinite(part.elapsedMs)) {
+			return Math.max(0, Math.round(part.elapsedMs));
+		}
+		const args = part.args;
+		if (!args || typeof args !== "object" || Array.isArray(args)) continue;
+		const value = (args as Record<string, unknown>).__claw_elapsed_ms;
+		if (typeof value === "number" && Number.isFinite(value)) {
+			return Math.max(0, Math.round(value));
+		}
+	}
+	return null;
+}
+
 function summarizeActivityChainLabel(
 	content: readonly ActivitySummaryPart[],
 	status: string | undefined,