Refactor session runtime persistence

frontend/app/components/assistant-ui/thread-list.tsx

@@ -124,24 +124,22 @@ export const ThreadList: FC = () => {
 };
 
 const ThreadListNew: FC = () => {
+	const { clearSession } = useClawSessionReplay();
 	return (
-		<div
-			onClickCapture={() => {
+		<Button
+			type="button"
+			variant="outline"
+			className="aui-thread-list-new h-9 w-full justify-start gap-2 rounded-lg px-3 text-sm hover:bg-muted data-active:bg-muted"
+			onClick={() => {
 				clearActiveSessionId();
 				pushHomeUrl();
+				clearSession();
 				window.dispatchEvent(new Event("claw-active-session-cleared"));
 			}}
 		>
-			<ThreadListPrimitive.New asChild>
-				<Button
-					variant="outline"
-					className="aui-thread-list-new h-9 w-full justify-start gap-2 rounded-lg px-3 text-sm hover:bg-muted data-active:bg-muted"
-				>
-					<PlusIcon className="size-4" />
-					New Task
-				</Button>
-			</ThreadListPrimitive.New>
-		</div>
+			<PlusIcon className="size-4" />
+			New Task
+		</Button>
 	);
 };
 
@@ -754,6 +752,7 @@ export function toReplayRepository(
 		}
 		const id = `${fallbackSessionId}-run-${runStatus.run_id ?? "active"}`;
 		const runStartedAtMs = resolveRunStartedAtMs(runStatus);
+		const runCreatedAt = new Date(runStartedAtMs ?? Date.now());
 		const parts = buildActiveRunParts(runStatus) as UIMessage["parts"];
 		const uiMessage: UIMessage = {
 			id,
@@ -763,12 +762,17 @@ export function toReplayRepository(
 				sessionId: session.session_id ?? fallbackSessionId,
 				runId: runStatus.run_id,
 				runStartedAtMs,
+				custom: {
+					sessionId: session.session_id ?? fallbackSessionId,
+					runId: runStatus.run_id,
+					runStartedAtMs,
+				},
 			},
 		} as UIMessage;
 		const threadMessage = {
 			id,
 			role: "assistant",
-			createdAt: new Date(),
+			createdAt: runCreatedAt,
 			content: toThreadMessageContent(parts),
 			status: { type: "running" },
 			metadata: {
@@ -776,6 +780,9 @@ export function toReplayRepository(
 				unstable_annotations: [],
 				unstable_data: [],
 				steps: [],
+				sessionId: session.session_id ?? fallbackSessionId,
+				runId: runStatus.run_id,
+				runStartedAtMs,
 				custom: {
 					sessionId: session.session_id ?? fallbackSessionId,
 					runId: runStatus.run_id,

frontend/app/components/assistant-ui/thread.tsx

@@ -434,11 +434,15 @@ function useCurrentThreadSessionId({
 			clearPendingWorkspaceSessionId();
 		}
 	}, [currentRun.sessionId, pendingSessionId]);
-	const replayRunSessionId = currentRun.active ? currentRun.sessionId : null;
+	const replayRunSessionId =
+		currentRun.active &&
+		(!activeSessionId || currentRun.sessionId === activeSessionId)
+			? currentRun.sessionId
+			: null;
 	return (
-		replayRunSessionId ??
-		(includePending ? pendingSessionId : null) ??
 		(includeActive ? activeSessionId : null) ??
+		(includePending ? pendingSessionId : null) ??
+		replayRunSessionId ??
 		currentRun.sessionId
 	);
 }
@@ -894,6 +898,10 @@ const Composer: FC = () => {
 		includePending: true,
 		includeActive: true,
 	});
+	const replayedRunIsCurrent =
+		replayedRun.active &&
+		Boolean(workspaceSessionId) &&
+		replayedRun.sessionId === workspaceSessionId;
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
 			<ComposerPrimitive.AttachmentDropzone asChild>
@@ -909,7 +917,7 @@ const Composer: FC = () => {
 						rows={1}
 						autoFocus
 						aria-label="Message input"
-						disabled={replayedRun.active}
+						disabled={replayedRunIsCurrent}
 					/>
 					<ComposerAction />
 				</div>
@@ -960,9 +968,18 @@ const ComposerAction: FC = () => {
 		includeActive: true,
 	});
 	const [runtimeCancelling, setRuntimeCancelling] = useState(false);
-	const cancelSessionId = replayedRun.sessionId ?? currentSessionId;
+	const replayedRunIsCurrent =
+		replayedRun.active &&
+		Boolean(currentSessionId) &&
+		replayedRun.sessionId === currentSessionId;
+	const cancelSessionId = replayedRunIsCurrent
+		? replayedRun.sessionId
+		: currentSessionId;
 	const backendRunStatus = useBackendActiveRunStatus(cancelSessionId);
-	const cancelRunId = replayedRun.runId ?? backendRunStatus?.run_id ?? null;
+	const cancelRunId =
+		(replayedRunIsCurrent ? replayedRun.runId : null) ??
+		backendRunStatus?.run_id ??
+		null;
 	useEffect(() => {
 		if (!runtimeRunning) setRuntimeCancelling(false);
 	}, [runtimeRunning]);
@@ -973,7 +990,7 @@ const ComposerAction: FC = () => {
 				<ComposerAssistButtons />
 			</div>
 			<ComposerContextStatus />
-			{!runtimeRunning && !replayedRun.active && !backendRunStatus ? (
+			{!runtimeRunning && !replayedRunIsCurrent && !backendRunStatus ? (
 				<ComposerPrimitive.Send asChild>
 					<TooltipIconButton
 						tooltip="Send message"
@@ -1012,7 +1029,7 @@ const ComposerAction: FC = () => {
 					<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
 				</Button>
 			) : null}
-			{!runtimeRunning && (replayedRun.active || backendRunStatus) ? (
+			{!runtimeRunning && (replayedRunIsCurrent || backendRunStatus) ? (
 				<ReplayedRunCancelButton
 					sessionId={cancelSessionId}
 					runId={cancelRunId}
@@ -1323,16 +1340,6 @@ function useComposerContextStatus() {
 		setModel: async () => {},
 	});
 
-	useEffect(() => {
-		const cleanLatestSessionId = normalizeSessionId(latestSessionId);
-		if (cleanLatestSessionId && !isRunning) {
-			writeActiveSessionId(cleanLatestSessionId);
-			window.setTimeout(() => {
-				window.dispatchEvent(new Event("claw-sessions-changed"));
-			}, 0);
-		}
-	}, [isRunning, latestSessionId]);
-
 	useEffect(() => {
 		if (wasRunningRef.current && !isRunning) {
 			window.dispatchEvent(new Event("claw-sessions-changed"));
@@ -1359,11 +1366,8 @@ function useComposerContextStatus() {
 				);
 				const latestMessageSessionId = normalizeSessionId(latestSessionId);
 				let sessionId = normalizeSessionId(
-					isRunning
-						? activeSessionId || latestMessageSessionId
-						: latestMessageSessionId || activeSessionId,
+					activeSessionId || latestMessageSessionId,
 				);
-				if (sessionId) writeActiveSessionId(sessionId);
 				let sessionPayload: ClawStoredSession | null = null;
 				let contextBudget: ContextBudget | undefined;
 				if (sessionId) {
@@ -2063,10 +2067,10 @@ const ActivityChainSummary: FC<{
 		metadataRunStartedAtMs ?? contentRunStartedAtMs ?? messageCreatedAtMs;
 	const liveStartedAtRef = useRef<number | null>(initialStartedAtMs);
 	const [elapsedMs, setElapsedMs] = useState<number | null>(() => {
-		if (storedElapsedMs !== null) return storedElapsedMs;
 		if (initialStartedAtMs !== null) {
 			return Math.max(0, Date.now() - initialStartedAtMs);
 		}
+		if (storedElapsedMs !== null) return storedElapsedMs;
 		return null;
 	});
 	const label = useAuiState((s) => {
@@ -2095,6 +2099,14 @@ const ActivityChainSummary: FC<{
 			liveStartedAtRef.current = authoritativeStartedAt;
 		} else if (liveStartedAtRef.current === null && storedElapsedMs !== null) {
 			liveStartedAtRef.current = Date.now() - storedElapsedMs;
+		} else if (
+			liveStartedAtRef.current !== null &&
+			storedElapsedMs !== null
+		) {
+			liveStartedAtRef.current = Math.min(
+				liveStartedAtRef.current,
+				Date.now() - storedElapsedMs,
+			);
 		} else if (
 			liveStartedAtRef.current === null &&
 			messageCreatedAtMs !== null

frontend/app/components/assistant-ui/threadlist-sidebar.tsx

@@ -1,5 +1,4 @@
 import type { ExportedMessageRepository } from "@assistant-ui/core";
-import { ThreadListPrimitive } from "@assistant-ui/react";
 import {
 	BarChart3Icon,
 	BotIcon,
@@ -185,6 +184,7 @@ function CollapsedSidebarRail({
 	onLogout: () => void;
 }) {
 	const { setOpen } = useSidebar();
+	const { clearSession } = useClawSessionReplay();
 
 	return (
 		<div className="flex h-full w-full flex-col items-center border-r bg-sidebar py-3 text-sidebar-foreground">
@@ -200,19 +200,17 @@ function CollapsedSidebarRail({
 			</button>
 
 			<div className="flex flex-col items-center gap-3">
-				<div
-					onClickCapture={() => {
+				<CollapsedIconButton
+					label="新任务"
+					onClick={() => {
 						clearActiveSessionId();
 						pushHomeUrl();
+						clearSession();
 						window.dispatchEvent(new Event("claw-active-session-cleared"));
 					}}
 				>
-					<ThreadListPrimitive.New asChild>
-						<CollapsedIconButton label="新任务">
-							<SquarePenIcon className="size-4" />
-						</CollapsedIconButton>
-					</ThreadListPrimitive.New>
-				</div>
+					<SquarePenIcon className="size-4" />
+				</CollapsedIconButton>
 				<CollapsedSessionSearch />
 				<CollapsedRecentSessions />
 			</div>