Refactor session runtime persistence

2026-06-12 16:17:58 +08:00
parent 77d360c1e8
commit 953126e1d3
15 changed files with 717 additions and 253 deletions
@@ -21,6 +21,8 @@ export const maxDuration = 3600;
 type ClawChatResponse = {
 	final_output?: string;
 	session_id?: string;
+	run_id?: string;
+	status?: string;
 	tool_calls?: number;
 	stop_reason?: string;
 	elapsed_ms?: number;
@@ -130,16 +132,14 @@ export async function POST(req: Request) {
 				writer.write({
 					type: "reasoning-delta",
 					id: "reasoning-1",
-					delta: "请求已发送，等待后端开始处理。",
+					delta: "任务已提交到后端运行队列。",
 				});
-				const payload = await callClawBackendStream(
+				const payload = await startClawBackendRun(
 					userPrompt,
 					runtimeContext,
 					account.id,
 					sessionId,
 					resumeId,
-					writer,
-					streamState,
 				);
 				const elapsedMs = payload.elapsed_ms ?? Date.now() - startedAt;
 				const text = formatClawResponse(payload);
@@ -148,13 +148,13 @@ export async function POST(req: Request) {
 					writer.write({
 						type: "reasoning-delta",
 						id: "reasoning-1",
-						delta: `\n思考并执行完成，用时 ${formatDuration(elapsedMs)}。`,
+						delta: `\n后端已接管执行，用时 ${formatDuration(elapsedMs)}。`,
 					});
 					writer.write({ type: "reasoning-end", id: "reasoning-1" });
 					streamState.reasoningEnded = true;
 				}
 				writeToolTrace(writer, payload.transcript, streamState);
-				if (!streamState.textStreamed) {
+				if (!streamState.textStreamed && text) {
 					writeTextDelta(writer, text, streamState);
 				}
 				endTextPart(writer, streamState);
@@ -164,6 +164,7 @@ export async function POST(req: Request) {
 					finishReason: payload.error ? "error" : "stop",
 					messageMetadata: {
 						sessionId: payload.session_id,
+						runId: payload.run_id,
 						toolCalls: payload.tool_calls,
 						stopReason: payload.stop_reason,
 						elapsedMs,
@@ -411,17 +412,15 @@ function safeFilename(value: string) {
 	return value.replace(/[^a-zA-Z0-9._-]+/g, "_").slice(0, 120) || "attachment";
 }

-async function callClawBackendStream(
+async function startClawBackendRun(
 	prompt: string,
 	runtimeContext: string,
 	accountId: string,
 	sessionId: string,
 	resumeSessionId?: string,
-	writer?: UIMessageStreamWriter<UIMessage>,
-	streamState?: StreamState,
 ): Promise<ClawChatResponse> {
 	try {
-		const response = await fetch(`${CLAW_API_URL}/api/chat/stream`, {
+		const response = await fetch(`${CLAW_API_URL}/api/chat/start`, {
 			method: "POST",
 			headers: { "content-type": "application/json" },
 			body: JSON.stringify({
@@ -443,9 +442,7 @@ async function callClawBackendStream(
 					`Claw backend returned ${response.status}`,
 			};
 		}
-		if (!response.body)
-			return { error: "Claw backend returned an empty stream" };
-		return await consumeClawStream(response.body, writer, streamState);
+		return (await response.json()) as ClawChatResponse;
 	} catch (err) {
 		return {
 			error:
@@ -71,9 +71,9 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 							: null;
 					const selectedResumeSessionId = selectedSessionId;
 					const outgoingSessionId =
-						lastSessionId ??
 						selectedResumeSessionId ??
 						pendingWorkspaceSessionId ??
+						lastSessionId ??
 						options.id;
 					const lastUserMessage = [...messages]
 						.reverse()
@@ -87,7 +87,7 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 							id: outgoingSessionId,
 							messages: lastUserMessage ? [lastUserMessage] : [],
 							resumeSessionId:
-								lastSessionId ?? selectedResumeSessionId ?? undefined,
+								selectedResumeSessionId ?? lastSessionId ?? undefined,
 						},
 					};
 				},
@@ -102,15 +102,13 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 			console.log("[replay-session] called", { sessionId });
 			writeActiveSessionId(sessionId);
 			pushSessionUrl(sessionId);
-			// 后端已经落盘/结束后，用回放结果接管当前会话。
-			// 先取消本地仍挂起的 assistant-ui run，避免 UI 残留“运行中”且无法停止。
-			if (runtime.thread.getState().isRunning) {
-				runtime.thread.cancelRun();
-			}
 			runtime.thread.import(repository);
 		},
 		[runtime],
 	);
+	const clearSession = useCallback(() => {
+		runtime.thread.import({ headId: null, messages: [] });
+	}, [runtime]);

 	useEffect(() => {
 		const sessionId = initialSessionId?.trim();
@@ -162,7 +160,7 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {

 	return (
 		<AssistantRuntimeProvider runtime={runtime}>
-			<ClawSessionReplayProvider value={{ replaySession }}>
+			<ClawSessionReplayProvider value={{ replaySession, clearSession }}>
 				<ActivityProvider>
 					<SidebarProvider>
 						<div className="flex h-dvh w-full">
@@ -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,
@@ -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
@@ -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>
@@ -8,6 +8,7 @@ type ClawSessionReplayContextValue = {
 		sessionId: string,
 		repository: ExportedMessageRepository,
 	) => void;
+	clearSession: () => void;
 };

 const ClawSessionReplayContext =