zk-data-agent/frontend/app/app/assistant.tsx

"use client";

import type { ExportedMessageRepository } from "@assistant-ui/core";
import { AssistantRuntimeProvider } from "@assistant-ui/react";
import {
	AssistantChatTransport,
	useChatRuntime,
} from "@assistant-ui/react-ai-sdk";
import type { UIMessage } from "ai";
import { useCallback, useEffect, useMemo, useRef } from "react";
import {
	ActivityPanel,
	ActivityProvider,
} from "@/components/assistant-ui/activity-panel";
import {
	fetchLatestRunStatus,
	toReplayRepository,
} from "@/components/assistant-ui/thread-list";
import { Thread } from "@/components/assistant-ui/thread";
import { ThreadListSidebar } from "@/components/assistant-ui/threadlist-sidebar";
import {
	SidebarInset,
	SidebarProvider,
	SidebarTrigger,
} from "@/components/ui/sidebar";
import {
	readActiveSessionId,
	readPendingWorkspaceSessionId,
	writeActiveSessionId,
} from "@/lib/claw-active-session";
import { pushSessionUrl } from "@/lib/claw-session-url";
import { ClawSessionReplayProvider } from "@/lib/claw-session-replay";
import { ClawAccountGate, useClawAccount } from "./claw-account-gate";

type AssistantProps = {
	initialSessionId?: string;
};

export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
	const { account, isLoading, setAccount } = useClawAccount();
	const loadedInitialSessionRef = useRef<string | null>(null);
	const transport = useMemo(
		() =>
			new AssistantChatTransport({
				api: "/api/chat",
				prepareSendMessagesRequest: async (options) => {
					const body = options.body as Record<string, unknown>;
					const messages = options.messages as UIMessage[];
					const selectedSessionId = readActiveSessionId();
					const lastSessionId = getLastSessionId(messages);
					const pendingWorkspaceSessionId =
						messages.length <= 1 ? readPendingWorkspaceSessionId() : null;
					const selectedResumeSessionId =
						messages.length > 1 ? selectedSessionId : null;
					const outgoingSessionId =
						lastSessionId ??
						selectedResumeSessionId ??
						pendingWorkspaceSessionId ??
						options.id;
					const lastUserMessage = [...messages]
						.reverse()
						.find((message) => message.role === "user");
					writeActiveSessionId(outgoingSessionId);

					return {
						body: {
							...body,
							id: outgoingSessionId,
							messages: lastUserMessage ? [lastUserMessage] : [],
							resumeSessionId:
								lastSessionId ?? selectedResumeSessionId ?? undefined,
						},
					};
				},
			}),
		[],
	);
	const runtime = useChatRuntime({
		transport,
	});
	const replaySession = useCallback(
		(sessionId: string, repository: ExportedMessageRepository) => {
			writeActiveSessionId(sessionId);
			// 后端已经落盘/结束后，用回放结果接管当前会话。
			// 先取消本地仍挂起的 assistant-ui run，避免 UI 残留“运行中”且无法停止。
			if (runtime.thread.getState().isRunning) {
				runtime.thread.cancelRun();
			}
			runtime.thread.import(repository);
		},
		[runtime],
	);

	useEffect(() => {
		const sessionId = initialSessionId?.trim();
		if (!account || !sessionId) return;
		if (loadedInitialSessionRef.current === sessionId) return;
		loadedInitialSessionRef.current = sessionId;
		const targetSessionId = sessionId;
		let cancelled = false;

		async function loadSession() {
			try {
				writeActiveSessionId(targetSessionId);
				const response = await fetch(`/api/claw/sessions/${targetSessionId}`, {
					cache: "no-store",
				});
				if (!response.ok || cancelled) return;
				const payload = await response.json();
				const runStatus = await fetchLatestRunStatus(targetSessionId);
				if (cancelled) return;
				replaySession(
					targetSessionId,
					toReplayRepository(payload, targetSessionId, runStatus),
				);
				pushSessionUrl(targetSessionId);
			} catch {
				if (!cancelled) {
					loadedInitialSessionRef.current = null;
				}
			}
		}

		void loadSession();
		return () => {
			cancelled = true;
		};
	}, [account, initialSessionId, replaySession]);

	if (isLoading) {
		return (
			<div className="flex min-h-dvh items-center justify-center bg-background text-muted-foreground text-sm">
				正在加载账号...
			</div>
		);
	}

	if (!account) {
		return <ClawAccountGate onAccount={setAccount} />;
	}

	return (
		<AssistantRuntimeProvider runtime={runtime}>
			<ClawSessionReplayProvider value={{ replaySession }}>
				<ActivityProvider>
					<SidebarProvider>
						<div className="flex h-dvh w-full">
							<SidebarTrigger className="fixed top-3 left-3 z-40 size-9 rounded-lg bg-background/90 shadow-sm ring-1 ring-border backdrop-blur sm:hidden" />
							<ThreadListSidebar
								account={account}
								onLogout={() => setAccount(null)}
							/>
							<SidebarInset>
								<div className="flex min-h-0 flex-1 overflow-hidden">
									<div className="min-w-0 flex-1 overflow-hidden">
										<Thread />
									</div>
									<ActivityPanel />
								</div>
							</SidebarInset>
						</div>
					</SidebarProvider>
				</ActivityProvider>
			</ClawSessionReplayProvider>
		</AssistantRuntimeProvider>
	);
};

function getLastSessionId(messages: UIMessage[]) {
	for (const message of [...messages].reverse()) {
		if (message.role !== "assistant") continue;
		const metadata = message.metadata;
		if (metadata && typeof metadata === "object" && "sessionId" in metadata) {
			const sessionId = (metadata as { sessionId?: unknown }).sessionId;
			if (typeof sessionId === "string" && sessionId) return sessionId;
		}
	}
	return undefined;
}