"use client"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { useChat } from "@ai-sdk/react"; import { DefaultChatTransport } from "ai"; import useSWR from "swr"; import useSWRImmutable from "swr/immutable"; import type { PersonaState } from "@/shared/ai-elements/persona"; import { showMapNotice } from "@/features/map/core/components/notice-actions"; import type { AgentApprovalMode, AgentChatMessage, AgentModelOption, AgentPermissionReply, AgentPermissionRequest, AgentQuestionRequest, AgentStreamRenderState } from "@/features/agent"; import { createAgentApiClient, type AgentSessionStreamEvent } from "@/features/agent/api/client"; import { agentBootstrapKey, agentSessionsKey, fetchAgentBootstrap, fetchAgentSessions } from "@/features/agent/api/swr"; import { isUiEnvelopeAllowed, parseUiEnvelopePayload, type UIEnvelopePayload, type UIRegistry } from "@/features/agent/ui-envelope"; import type { FrontendActionRequest } from "@/features/agent/frontend-action"; import { FrontendActionExecutor } from "@/features/agent/frontend-action/executor"; import { applySessionStreamEvent, createCompletedStreamRenderState, getDataString, markLastAssistantStreamDone, markStreamRenderPending, readBodyString, toAgentChatMessages, toAgentUiMessages, toPermissionStatus, type AgentDataPart, type AgentUiMessage, type PermissionOverride, type QuestionOverride } from "./agent-session-message-state"; const AGENT_PANEL_COLLAPSE_MS = 180; type AgentRuntimeAvailability = "connecting" | "connected" | "unavailable" | "failed"; type UseWorkbenchAgentOptions = { onUiEnvelope: (payload: UIEnvelopePayload, sessionId: string) => Promise | void; onFrontendAction: (request: FrontendActionRequest, signal: AbortSignal) => Promise; }; export function useWorkbenchAgent({ onUiEnvelope, onFrontendAction }: UseWorkbenchAgentOptions) { const collapseTimerRef = useRef(null); const mobileCollapseTimerRef = useRef(null); const sessionStreamAbortRef = useRef(null); const clientRef = useRef(createAgentApiClient()); const sessionIdRef = useRef(null); const approvalModeRef = useRef("request"); const registryRef = useRef(null); const frontendActionEnabledRef = useRef(false); const onFrontendActionRef = useRef(onFrontendAction); const processedEnvelopeIdsRef = useRef(new Set()); const [panelOpen, setPanelOpen] = useState(true); const [panelCollapsing, setPanelCollapsing] = useState(false); const [mobileOpen, setMobileOpen] = useState(false); const [mobilePanelCollapsing, setMobilePanelCollapsing] = useState(false); const [sessionId, setSessionId] = useState(null); const [registry, setRegistry] = useState(null); const [sessionTitle, setSessionTitle] = useState("新对话"); const [statusLabel, setStatusLabel] = useState("Agent 后端连接中"); const [runtimeAvailability, setRuntimeAvailability] = useState("connecting"); const [resumedStreaming, setResumedStreaming] = useState(false); const [modelOptions, setModelOptions] = useState([]); const [selectedModel, setSelectedModel] = useState(""); const [approvalMode, setApprovalModeState] = useState("request"); const [streamRenderState, setStreamRenderState] = useState({}); const [permissionOverrides, setPermissionOverrides] = useState>({}); const [questionOverrides, setQuestionOverrides] = useState>({}); useEffect(() => { sessionIdRef.current = sessionId; }, [sessionId]); const setApprovalMode = useCallback((mode: AgentApprovalMode) => { approvalModeRef.current = mode; setApprovalModeState(mode); }, []); useEffect(() => { registryRef.current = registry; }, [registry]); useEffect(() => { onFrontendActionRef.current = onFrontendAction; }, [onFrontendAction]); const applySessionId = useCallback((nextSessionId: string | undefined) => { if (!nextSessionId || sessionIdRef.current === nextSessionId) { return; } sessionIdRef.current = nextSessionId; setSessionId(nextSessionId); }, []); const clearSessionId = useCallback(() => { sessionIdRef.current = null; setSessionId(null); }, []); const resolveRenderRef = useCallback((renderRef: string, nextSessionId: string) => { return clientRef.current.resolveRenderRef(renderRef, nextSessionId); }, []); const frontendActionExecutor = useMemo( () => new FrontendActionExecutor( (request, signal) => onFrontendActionRef.current(request, signal), (nextSessionId, actionId, result) => clientRef.current.submitFrontendActionResult(nextSessionId, actionId, result) ), [] ); const handleFrontendAction = useCallback( (value: unknown) => frontendActionExecutor.handle(value, sessionIdRef.current), [frontendActionExecutor] ); const claimEnvelope = useCallback((payload: UIEnvelopePayload, activeSessionId: string) => { if (payload.session_id !== activeSessionId) return false; const key = `${activeSessionId}:${payload.envelope_id}`; if (processedEnvelopeIdsRef.current.has(key)) return false; processedEnvelopeIdsRef.current.add(key); if (processedEnvelopeIdsRef.current.size > 256) { const oldest = processedEnvelopeIdsRef.current.values().next().value; if (oldest) processedEnvelopeIdsRef.current.delete(oldest); } return true; }, []); const handleDataPart = useCallback( (part: AgentDataPart) => { const eventSessionId = getDataString(part.data, "session_id"); if (eventSessionId && sessionIdRef.current && eventSessionId !== sessionIdRef.current) { return; } applySessionId(eventSessionId); if (part.type === "data-stream_token") { markStreamRenderPending( part.data, chatRef.current.messages, setStreamRenderState ); return; } if (part.type === "data-progress") { const title = getDataString(part.data, "title"); if (title) { setStatusLabel(title); } return; } if (part.type === "data-frontend_action") { void handleFrontendAction(part.data); return; } if (part.type === "data-session_title") { const title = getDataString(part.data, "title"); if (title) { setSessionTitle(title); } return; } if (part.type !== "data-ui_envelope") { return; } const payload = parseUiEnvelopePayload(part.data); const activeSessionId = getDataString(part.data, "session_id") ?? sessionIdRef.current; if (!payload || !activeSessionId || !isUiEnvelopeAllowed(payload.envelope, registryRef.current)) { showMapNotice({ tone: "warning", title: "Agent 展示已降级", message: "收到的 UIEnvelope 未通过前端白名单校验。" }); return; } if (!claimEnvelope(payload, activeSessionId)) return; void Promise.resolve(onUiEnvelope(payload, activeSessionId)).catch((error) => { showMapNotice({ tone: "error", title: "Agent UIEnvelope 处理失败", message: error instanceof Error ? error.message : "工作台处理 Agent 展示结果失败。" }); }); }, [applySessionId, claimEnvelope, handleFrontendAction, onUiEnvelope] ); const transport = useMemo( () => new DefaultChatTransport({ api: "/api/v1/agent/chat/stream", prepareSendMessagesRequest({ id, messages, body, trigger, messageId }) { return { body: { ...body, id, messages, trigger, messageId, session_id: readBodyString(body, "session_id") ?? sessionIdRef.current ?? undefined, approval_mode: readBodyString(body, "approval_mode") ?? approvalModeRef.current, capabilities: frontendActionEnabledRef.current ? ["frontend-action@1"] : [] } }; } }), [] ); const chat = useChat({ transport, onData: (part) => handleDataPart(part as AgentDataPart), onError: (error) => { markLastAssistantStreamDone(chatRef.current.messages, setStreamRenderState); setRuntimeAvailability("failed"); setStatusLabel("Agent 后端请求失败"); showMapNotice({ id: "agent-stream-status", tone: "error", title: "Agent 请求失败", message: error.message || "请确认后端服务已启动。" }); }, onFinish: () => { markLastAssistantStreamDone(chatRef.current.messages, setStreamRenderState); setStatusLabel("Agent 分析完成"); } }); const chatRef = useRef(chat); useEffect(() => { chatRef.current = chat; }, [chat]); const { data: backendConnection, error: backendConnectionError } = useSWRImmutable(agentBootstrapKey, () => fetchAgentBootstrap(clientRef.current), { shouldRetryOnError: false }); const { data: sessionHistory = [], isLoading: sessionHistoryLoading, isValidating: sessionHistoryValidating, mutate: mutateSessionHistory } = useSWR(agentSessionsKey, () => fetchAgentSessions(clientRef.current), { revalidateOnFocus: false, shouldRetryOnError: false }); const clearCollapseTimer = useCallback(() => { if (collapseTimerRef.current !== null) { window.clearTimeout(collapseTimerRef.current); collapseTimerRef.current = null; } }, []); const clearMobileCollapseTimer = useCallback(() => { if (mobileCollapseTimerRef.current !== null) { window.clearTimeout(mobileCollapseTimerRef.current); mobileCollapseTimerRef.current = null; } }, []); const stopSessionStreamSubscription = useCallback(() => { sessionStreamAbortRef.current?.abort(); sessionStreamAbortRef.current = null; setResumedStreaming(false); }, []); useEffect(() => { return () => { clearCollapseTimer(); clearMobileCollapseTimer(); stopSessionStreamSubscription(); frontendActionExecutor.cancelAll(); chatRef.current.stop(); }; }, [clearCollapseTimer, clearMobileCollapseTimer, frontendActionExecutor, stopSessionStreamSubscription]); useEffect(() => { if (!backendConnection) { return; } registryRef.current = backendConnection.registry; frontendActionEnabledRef.current = Boolean(backendConnection.frontendActionRegistry); setRegistry(backendConnection.registry); setModelOptions(backendConnection.models); setSelectedModel((current) => current || backendConnection.defaultModel || backendConnection.models[0]?.id || ""); if (chatRef.current.status === "ready") { setRuntimeAvailability("connected"); setStatusLabel("准备就绪"); } }, [backendConnection]); useEffect(() => { if (!backendConnectionError) { return; } setRuntimeAvailability("unavailable"); setStatusLabel("Agent 后端不可用"); }, [backendConnectionError]); const collapsePanel = useCallback(() => { clearCollapseTimer(); if (prefersReducedMotion()) { setPanelOpen(false); setPanelCollapsing(false); return; } setPanelCollapsing(true); collapseTimerRef.current = window.setTimeout(() => { setPanelOpen(false); setPanelCollapsing(false); collapseTimerRef.current = null; }, AGENT_PANEL_COLLAPSE_MS); }, [clearCollapseTimer]); const expandPanel = useCallback(() => { clearCollapseTimer(); setPanelCollapsing(false); setPanelOpen(true); }, [clearCollapseTimer]); const openMobilePanel = useCallback(() => { clearMobileCollapseTimer(); setMobilePanelCollapsing(false); setMobileOpen(true); }, [clearMobileCollapseTimer]); const closeMobilePanel = useCallback(() => { clearMobileCollapseTimer(); if (prefersReducedMotion()) { setMobileOpen(false); setMobilePanelCollapsing(false); return; } setMobilePanelCollapsing(true); mobileCollapseTimerRef.current = window.setTimeout(() => { setMobileOpen(false); setMobilePanelCollapsing(false); mobileCollapseTimerRef.current = null; }, AGENT_PANEL_COLLAPSE_MS); }, [clearMobileCollapseTimer]); const streaming = chat.status === "submitted" || chat.status === "streaming" || resumedStreaming; const messages = useMemo( () => toAgentChatMessages(chat.messages, permissionOverrides, questionOverrides), [chat.messages, permissionOverrides, questionOverrides] ); const personaState = useMemo( () => deriveAgentPersonaState(messages, { chatStatus: chat.status, runtimeAvailability, statusLabel }), [chat.status, messages, runtimeAvailability, statusLabel] ); const refreshSessionHistory = useCallback(async () => { try { await mutateSessionHistory(); } catch (error) { showMapNotice({ id: "agent-history-status", tone: "error", title: "Agent 历史记录加载失败", message: error instanceof Error ? error.message : "无法获取 Agent 会话历史。" }); } }, [mutateSessionHistory]); const handleSessionStreamEvent = useCallback( (event: AgentSessionStreamEvent, expectedSessionId: string) => { if (event.type === "state") { if (sessionIdRef.current !== expectedSessionId && sessionIdRef.current !== event.sessionId) { return; } const nextMessages = toAgentUiMessages(event.messages); chatRef.current.setMessages(nextMessages); setStreamRenderState(createCompletedStreamRenderState(nextMessages)); setResumedStreaming(event.isStreaming); setStatusLabel(event.isStreaming ? "Agent 会话流已恢复" : "Agent 历史会话已打开"); return; } const eventSessionId = event.sessionId; if (eventSessionId && eventSessionId !== sessionIdRef.current) { return; } if (event.type === "session_title") { const title = getDataString(event.data, "title"); if (title) { setSessionTitle(title); } } else if (event.type === "token") { markStreamRenderPending( event.data, chatRef.current.messages, setStreamRenderState ); } else if (event.type === "ui_envelope") { const payload = parseUiEnvelopePayload(event.data); const activeSessionId = eventSessionId ?? sessionIdRef.current; if (!payload || !activeSessionId || !isUiEnvelopeAllowed(payload.envelope, registryRef.current)) { showMapNotice({ tone: "warning", title: "Agent 展示已降级", message: "收到的 UIEnvelope 未通过前端白名单校验。" }); return; } if (!claimEnvelope(payload, activeSessionId)) return; void Promise.resolve(onUiEnvelope(payload, activeSessionId)).catch((error) => { showMapNotice({ tone: "error", title: "Agent UIEnvelope 处理失败", message: error instanceof Error ? error.message : "工作台处理 Agent 展示结果失败。" }); }); } else if (event.type === "frontend_action") { void handleFrontendAction(event.data); } else if (event.type === "progress") { const title = getDataString(event.data, "title"); if (title) { setStatusLabel(title); } } else if (event.type === "done") { setResumedStreaming(false); markLastAssistantStreamDone(chatRef.current.messages, setStreamRenderState); setStatusLabel("Agent 分析完成"); void refreshSessionHistory(); } else if (event.type === "error") { setResumedStreaming(false); markLastAssistantStreamDone(chatRef.current.messages, setStreamRenderState); setStatusLabel("Agent 后端请求失败"); void refreshSessionHistory(); } chatRef.current.setMessages((current) => applySessionStreamEvent(current, event)); }, [claimEnvelope, handleFrontendAction, onUiEnvelope, refreshSessionHistory] ); const startSessionStreamSubscription = useCallback( (nextSessionId: string) => { stopSessionStreamSubscription(); const controller = new AbortController(); sessionStreamAbortRef.current = controller; setResumedStreaming(true); void clientRef.current .streamSession(nextSessionId, { signal: controller.signal, onEvent: (event) => handleSessionStreamEvent(event, nextSessionId) }) .catch((error) => { if (!controller.signal.aborted) { setRuntimeAvailability("failed"); setStatusLabel("Agent 会话流恢复失败"); showMapNotice({ id: "agent-stream-status", tone: "error", title: "Agent 会话流恢复失败", message: error instanceof Error ? error.message : "无法恢复这个运行中的会话。" }); } }) .finally(() => { if (sessionStreamAbortRef.current === controller) { sessionStreamAbortRef.current = null; setResumedStreaming(false); void refreshSessionHistory(); } }); }, [handleSessionStreamEvent, refreshSessionHistory, stopSessionStreamSubscription] ); const loadHistorySession = useCallback( async (nextSessionId: string) => { if (streaming) { chat.stop(); stopSessionStreamSubscription(); } try { const loadedSession = await clientRef.current.loadSession(nextSessionId); if (!loadedSession) { showMapNotice({ id: "agent-history-status", tone: "warning", title: "Agent 历史记录不可用", message: "这个历史会话不存在或已被删除。" }); await refreshSessionHistory(); return; } applySessionId(loadedSession.sessionId); setSessionTitle(loadedSession.title); setStatusLabel("Agent 历史会话已打开"); setPermissionOverrides({}); setQuestionOverrides({}); const nextMessages = toAgentUiMessages(loadedSession.messages); chat.setMessages(nextMessages); setStreamRenderState(createCompletedStreamRenderState(nextMessages)); if (loadedSession.runStatus === "running" || loadedSession.isStreaming) { startSessionStreamSubscription(loadedSession.sessionId); } await refreshSessionHistory(); } catch (error) { showMapNotice({ id: "agent-history-status", tone: "error", title: "Agent 历史记录打开失败", message: error instanceof Error ? error.message : "无法打开这个历史会话。" }); } }, [applySessionId, chat, refreshSessionHistory, startSessionStreamSubscription, stopSessionStreamSubscription, streaming] ); const renameHistorySession = useCallback( async (nextSessionId: string, title: string) => { const normalizedTitle = title.trim(); if (!normalizedTitle) { showMapNotice({ tone: "warning", message: "对话主题不能为空。" }); return; } try { await clientRef.current.updateSessionTitle(nextSessionId, normalizedTitle, { isTitleManuallyEdited: true }); await mutateSessionHistory( (current = []) => current.map((session) => session.id === nextSessionId ? { ...session, title: normalizedTitle, updatedAt: Date.now() } : session ), { revalidate: false } ); if (sessionIdRef.current === nextSessionId) { setSessionTitle(normalizedTitle); } await refreshSessionHistory(); } catch (error) { showMapNotice({ id: "agent-history-status", tone: "error", title: "Agent 对话重命名失败", message: error instanceof Error ? error.message : "无法更新这个对话主题。" }); } }, [mutateSessionHistory, refreshSessionHistory] ); const resetDraftSession = useCallback( (nextStatusLabel = "Agent 新对话已准备") => { clearSessionId(); setSessionTitle("新对话"); setStatusLabel(nextStatusLabel); setRuntimeAvailability("connected"); setPermissionOverrides({}); setQuestionOverrides({}); setStreamRenderState({}); chat.setMessages([]); }, [chat, clearSessionId] ); const deleteHistorySession = useCallback( async (nextSessionId: string) => { if (streaming) { showMapNotice({ tone: "warning", message: "Agent 正在处理上一条指令,完成后再删除会话。" }); return false; } try { await clientRef.current.deleteSession(nextSessionId); await mutateSessionHistory( (current = []) => current.filter((session) => session.id !== nextSessionId), { revalidate: false } ); if (sessionIdRef.current === nextSessionId) { resetDraftSession(); } await refreshSessionHistory(); return true; } catch (error) { showMapNotice({ id: "agent-history-status", tone: "error", title: "Agent 历史记录删除失败", message: error instanceof Error ? error.message : "无法删除这个历史会话。" }); return false; } }, [mutateSessionHistory, refreshSessionHistory, resetDraftSession, streaming] ); const startNewSession = useCallback(async () => { if (chat.status === "submitted" || chat.status === "streaming") { chat.stop(); } stopSessionStreamSubscription(); resetDraftSession(); }, [chat, resetDraftSession, stopSessionStreamSubscription]); const submitPrompt = useCallback( async (prompt: string) => { const normalizedPrompt = prompt.trim(); if (!normalizedPrompt) { return; } if (streaming) { showMapNotice({ tone: "warning", message: "Agent 正在处理上一条指令。" }); return; } setStatusLabel("Agent 正在分析"); await chat.sendMessage( { text: normalizedPrompt }, { body: { session_id: sessionIdRef.current ?? undefined, model: selectedModel || undefined, approval_mode: approvalMode } } ); }, [approvalMode, chat, selectedModel, streaming] ); const stopPrompt = useCallback(() => { chat.stop(); stopSessionStreamSubscription(); const activeSessionId = sessionIdRef.current; if (activeSessionId) { void clientRef.current.abort(activeSessionId).catch(() => undefined); } markLastAssistantStreamDone(chatRef.current.messages, setStreamRenderState); }, [chat, stopSessionStreamSubscription]); const replyPermission = useCallback(async (permission: AgentPermissionRequest, reply: AgentPermissionReply) => { setPermissionOverrides((current) => ({ ...current, [permission.requestId]: { status: "submitting" } })); try { await clientRef.current.replyPermission(permission.requestId, { sessionId: permission.sessionId, reply }); setPermissionOverrides((current) => ({ ...current, [permission.requestId]: { status: toPermissionStatus(reply) } })); } catch (error) { const message = error instanceof Error ? error.message : "权限请求提交失败。"; setPermissionOverrides((current) => ({ ...current, [permission.requestId]: { status: "error", error: message } })); showMapNotice({ id: "agent-permission-status", tone: "error", title: "Agent 权限提交失败", message }); } }, []); const replyQuestion = useCallback(async (question: AgentQuestionRequest, answers: string[][]) => { setQuestionOverrides((current) => ({ ...current, [question.requestId]: { status: "submitting" } })); try { await clientRef.current.replyQuestion(question.requestId, { sessionId: question.sessionId, answers }); setQuestionOverrides((current) => ({ ...current, [question.requestId]: { status: "answered", answers } })); } catch (error) { const message = error instanceof Error ? error.message : "问题回复提交失败。"; setQuestionOverrides((current) => ({ ...current, [question.requestId]: { status: "error", error: message } })); showMapNotice({ id: "agent-question-status", tone: "error", title: "Agent 问题提交失败", message }); } }, []); const rejectQuestion = useCallback(async (question: AgentQuestionRequest) => { setQuestionOverrides((current) => ({ ...current, [question.requestId]: { status: "submitting" } })); try { await clientRef.current.rejectQuestion(question.requestId, { sessionId: question.sessionId }); setQuestionOverrides((current) => ({ ...current, [question.requestId]: { status: "rejected" } })); } catch (error) { const message = error instanceof Error ? error.message : "问题跳过提交失败。"; setQuestionOverrides((current) => ({ ...current, [question.requestId]: { status: "error", error: message } })); showMapNotice({ id: "agent-question-status", tone: "error", title: "Agent 问题提交失败", message }); } }, []); return { approvalMode, collapsePanel, expandPanel, mobileOpen, mobilePanelCollapsing, messages, modelOptions, openMobilePanel, closeMobilePanel, panelCollapsing, panelOpen, personaState, refreshSessionHistory, startNewSession, loadHistorySession, renameHistorySession, deleteHistorySession, sessionHistory, sessionHistoryLoading: sessionHistoryLoading || sessionHistoryValidating, sessionId, sessionTitle, statusLabel, streamRenderState, selectedModel, setApprovalMode, setSelectedModel, stopPrompt, streaming, rejectQuestion, replyPermission, replyQuestion, resolveRenderRef, submitPrompt }; } function deriveAgentPersonaState( messages: AgentChatMessage[], { chatStatus, runtimeAvailability, statusLabel }: { chatStatus: string; runtimeAvailability: AgentRuntimeAvailability; statusLabel: string; } ): PersonaState { if ( runtimeAvailability === "unavailable" || runtimeAvailability === "failed" || /失败|不可用|错误|降级/.test(statusLabel) ) { return "asleep"; } if (hasPendingOperatorInput(messages)) { return "listening"; } if ( chatStatus === "submitted" || chatStatus === "streaming" || hasRunningSessionWork(messages) || runtimeAvailability === "connecting" ) { return "thinking"; } return "idle"; } function hasPendingOperatorInput(messages: AgentChatMessage[]) { return messages.some( (message) => message.permissions?.some((permission) => permission.status === "pending" || permission.status === "error") || message.questions?.some((question) => question.status === "pending" || question.status === "error") ); } function hasRunningSessionWork(messages: AgentChatMessage[]) { const latestProgress = messages.flatMap((message) => message.progress ?? []).at(-1); const latestTodo = messages.flatMap((message) => message.todos?.todos ?? []).at(-1); return latestProgress?.status === "running" || latestTodo?.status === "in_progress"; } function prefersReducedMotion() { return window.matchMedia("(prefers-reduced-motion: reduce)").matches; }