/** * Chat Hook * * 管理 WebSocket 连接和消息状态 */ import { useState, useEffect, useCallback, useRef } from 'react'; import { createWebSocket, getMessages, type Message } from '../api/client.js'; import type { PermissionRequest } from '../components/PermissionDialog.js'; import type { ConfigErrorPayload, ToolStartPayload, ToolEndPayload, MessagePart, ToolMessagePart, AgentModeType, SubagentStartPayload, SubagentEndPayload, SubagentStreamPayload, SubagentToolStartPayload, SubagentToolEndPayload, SubagentState, SubagentToolInfo, } from '../api/types.js'; interface UseChatOptions { sessionId: string; onError?: (error: Error) => void; onSessionNotFound?: () => void; onSessionUpdated?: (sessionId: string, name: string) => void; /** 配置错误回调(如 API Key 未配置) */ onConfigError?: (error: ConfigErrorPayload) => void; } interface ChatState { messages: Message[]; isConnected: boolean; isLoading: boolean; /** 流式消息对象,复用 Message 结构 */ streamingMessage: Message | null; permissionRequest: PermissionRequest | null; /** Agent 模式 (会话级别) */ agentMode: AgentModeType; /** 是否自动授权文件写入/编辑 (会话级别) */ autoApprove: boolean; /** 当前正在执行的 Agent 名称 */ currentAgent: string; /** 当前正在执行的子 Agent 状态 */ currentSubagent: SubagentState | null; } export function useChat({ sessionId, onError, onSessionNotFound, onSessionUpdated, onConfigError }: UseChatOptions) { const [state, setState] = useState({ messages: [], isConnected: false, isLoading: false, streamingMessage: null, permissionRequest: null, agentMode: 'build', autoApprove: false, currentAgent: 'build', currentSubagent: null, }); const wsRef = useRef(null); const reconnectTimeoutRef = useRef>(); const reconnectAttemptsRef = useRef(0); const maxReconnectAttempts = 5; // 标记是否正在主动关闭连接(切换 session 时) const isClosingRef = useRef(false); // 用 ref 存储回调,避免依赖变化导致无限循环 const onErrorRef = useRef(onError); const onSessionNotFoundRef = useRef(onSessionNotFound); const onSessionUpdatedRef = useRef(onSessionUpdated); const onConfigErrorRef = useRef(onConfigError); onErrorRef.current = onError; onSessionNotFoundRef.current = onSessionNotFound; onSessionUpdatedRef.current = onSessionUpdated; onConfigErrorRef.current = onConfigError; // 加载历史消息 const loadMessages = useCallback(async () => { try { const { data } = await getMessages(sessionId); setState((prev) => ({ ...prev, messages: data })); } catch (error) { // 会话不存在(404 或 "Session not found"),通知上层重新创建 const msg = error instanceof Error ? error.message : ''; if (msg.includes('404') || msg.toLowerCase().includes('not found')) { onSessionNotFoundRef.current?.(); return; } onErrorRef.current?.(error instanceof Error ? error : new Error('Failed to load messages')); } }, [sessionId]); // 连接 WebSocket const connect = useCallback(() => { // 如果正在关闭,不要连接 if (isClosingRef.current) return; // 如果已经连接,不要重复连接 if (wsRef.current?.readyState === WebSocket.OPEN) return; // 如果正在连接中,不要重复连接 if (wsRef.current?.readyState === WebSocket.CONNECTING) return; const ws = createWebSocket(sessionId); ws.onopen = () => { // 如果在连接过程中组件已卸载,立即关闭 if (isClosingRef.current) { ws.close(); return; } reconnectAttemptsRef.current = 0; // 连接成功,重置重连次数 setState((prev) => ({ ...prev, isConnected: true })); }; ws.onclose = () => { setState((prev) => ({ ...prev, isConnected: false })); // 主动关闭时不重连 if (isClosingRef.current) { isClosingRef.current = false; return; } // 限制重连次数 if (reconnectAttemptsRef.current < maxReconnectAttempts) { reconnectAttemptsRef.current++; reconnectTimeoutRef.current = setTimeout(connect, 3000); } }; ws.onerror = () => { // 主动关闭时不报错 if (isClosingRef.current) return; onErrorRef.current?.(new Error('WebSocket connection error')); }; ws.onmessage = (event) => { try { const message = JSON.parse(event.data); switch (message.type) { case 'chunk': { const chunkContent = message.payload?.content || ''; setState((prev) => { // 初始化或获取当前流式消息 const streaming = prev.streamingMessage || { id: `streaming-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`, role: 'assistant' as const, timestamp: new Date().toISOString(), parts: [] as MessagePart[], content: '', metadata: { agentName: prev.currentAgent }, }; // 复制 parts 数组以进行修改 const parts = [...streaming.parts]; const lastPart = parts[parts.length - 1]; // 如果最后一个 part 是 text,追加内容;否则创建新 text part if (lastPart?.type === 'text') { parts[parts.length - 1] = { ...lastPart, text: lastPart.text + chunkContent, }; } else if (chunkContent) { parts.push({ type: 'text', id: `text-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`, text: chunkContent, }); } return { ...prev, streamingMessage: { ...streaming, parts, content: (streaming.content || '') + chunkContent, }, }; }); break; } case 'tool_start': { const payload = message.payload as ToolStartPayload; setState((prev) => { // 初始化或获取当前流式消息 const streaming = prev.streamingMessage || { id: `streaming-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`, role: 'assistant' as const, timestamp: new Date().toISOString(), parts: [] as MessagePart[], content: '', metadata: { agentName: prev.currentAgent }, }; // 检查是否已存在相同 ID 的工具调用(去重) const existingTool = streaming.parts.find( (part) => part.type === 'tool' && part.id === payload.id ); if (existingTool) { // 已存在相同 ID 的工具调用,跳过 return prev; } // 添加工具调用 part const toolPart: ToolMessagePart = { type: 'tool', id: payload.id, toolCallId: payload.id, toolName: payload.toolName, status: 'running', arguments: payload.arguments, }; // 如果是 task 工具,切换到子 agent const newAgent = payload.toolName === 'task' && payload.arguments?.subagent_type ? (payload.arguments.subagent_type as string) : prev.currentAgent; return { ...prev, currentAgent: newAgent, streamingMessage: { ...streaming, parts: [...streaming.parts, toolPart], metadata: { ...streaming.metadata, agentName: newAgent }, }, }; }); break; } case 'tool_end': { const payload = message.payload as ToolEndPayload; setState((prev) => { if (!prev.streamingMessage) return prev; // 查找并更新对应的工具 part const parts = prev.streamingMessage.parts.map((part) => { if (part.type === 'tool' && part.id === payload.id) { return { ...part, status: payload.status, result: payload.result, error: payload.error, duration: payload.duration, } as ToolMessagePart; } return part; }); // 查找完成的工具是否为 task,如果是则恢复主 agent const completedTool = prev.streamingMessage.parts.find( (part) => part.type === 'tool' && part.id === payload.id ); const isTaskTool = completedTool?.type === 'tool' && completedTool.toolName === 'task'; const newAgent = isTaskTool ? prev.agentMode : prev.currentAgent; return { ...prev, currentAgent: newAgent, streamingMessage: { ...prev.streamingMessage, parts, metadata: { ...prev.streamingMessage.metadata, agentName: newAgent }, }, }; }); break; } case 'done': setState((prev) => { // 使用流式消息或创建新消息 const streaming = prev.streamingMessage; const content = message.payload?.content || streaming?.content || ''; // 从服务器 payload 获取 agentName,或使用当前 agentMode const agentName = message.payload?.agentName || prev.agentMode; const newMessage: Message = streaming ? { ...streaming, id: message.payload?.id || streaming.id, timestamp: message.payload?.timestamp || streaming.timestamp, content, metadata: { ...streaming.metadata, agentName }, } : { id: message.payload?.id || `assistant-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`, role: 'assistant', timestamp: message.payload?.timestamp || new Date().toISOString(), parts: [{ type: 'text', id: `text-${Date.now()}`, text: content }], content, metadata: { agentName }, }; return { ...prev, messages: [...prev.messages, newMessage], streamingMessage: null, isLoading: false, currentAgent: prev.agentMode, // 恢复为主 agent }; }); break; case 'message_received': // 用户消息已确认 - 构建完整的消息对象 setState((prev) => { const content = message.payload?.content || ''; const userMessage: Message = { id: `user-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`, role: 'user', timestamp: new Date().toISOString(), parts: [{ type: 'text', id: `text-${Date.now()}`, text: content }], content, }; return { ...prev, messages: [...prev.messages, userMessage], }; }); break; case 'error': // 检查是否为配置错误 if (message.payload?.type === 'config_error') { onConfigErrorRef.current?.(message.payload as ConfigErrorPayload); } else { onErrorRef.current?.(new Error(message.payload?.message || 'Unknown error')); } setState((prev) => ({ ...prev, isLoading: false, streamingMessage: null })); break; case 'session_updated': // 会话信息更新(如标题) if (message.payload?.id && message.payload?.name) { onSessionUpdatedRef.current?.(message.payload.id, message.payload.name); } break; case 'permission_request': // 权限请求 if (message.payload) { setState((prev) => ({ ...prev, permissionRequest: message.payload as PermissionRequest, })); } break; // ============ 子 Agent 事件处理 ============ case 'subagent:start': { const payload = message.payload as SubagentStartPayload; setState((prev) => ({ ...prev, currentAgent: payload.agentName, currentSubagent: { id: payload.agentId, name: payload.agentName, description: payload.description, status: 'running', tools: [], streamContent: '', }, })); break; } case 'subagent:tool_start': { const payload = message.payload as SubagentToolStartPayload; setState((prev) => { if (!prev.currentSubagent || prev.currentSubagent.id !== payload.agentId) { return prev; } // 检查是否已存在相同 ID 的工具调用(去重) const existingTool = prev.currentSubagent.tools.find( (tool) => tool.id === payload.toolCallId ); if (existingTool) { // 已存在相同 ID 的工具调用,跳过 return prev; } const newTool: SubagentToolInfo = { id: payload.toolCallId, toolName: payload.toolName, status: 'running', args: payload.args, }; return { ...prev, currentSubagent: { ...prev.currentSubagent, tools: [...prev.currentSubagent.tools, newTool], }, }; }); break; } case 'subagent:stream': { const payload = message.payload as SubagentStreamPayload; setState((prev) => { if (!prev.currentSubagent || prev.currentSubagent.id !== payload.agentId) { return prev; } return { ...prev, currentSubagent: { ...prev.currentSubagent, streamContent: prev.currentSubagent.streamContent + payload.content, }, }; }); break; } case 'subagent:tool_end': { const payload = message.payload as SubagentToolEndPayload; setState((prev) => { if (!prev.currentSubagent || prev.currentSubagent.id !== payload.agentId) { return prev; } const updatedTools = prev.currentSubagent.tools.map((tool) => { if (tool.id === payload.toolCallId) { return { ...tool, status: payload.status === 'completed' ? 'completed' : 'error', result: payload.result, error: payload.error, duration: payload.duration, } as SubagentToolInfo; } return tool; }); return { ...prev, currentSubagent: { ...prev.currentSubagent, tools: updatedTools, }, }; }); break; } case 'subagent:end': { const payload = message.payload as SubagentEndPayload; setState((prev) => { // 只有当 agentId 匹配时才处理 if (!prev.currentSubagent || prev.currentSubagent.id !== payload.agentId) { return prev; } return { ...prev, currentAgent: prev.agentMode, // 恢复为主 Agent currentSubagent: { ...prev.currentSubagent, status: payload.success ? 'completed' : 'error', duration: payload.duration, error: payload.error, }, }; }); // 完成后短暂延迟再清除,让 UI 能显示最终状态 setTimeout(() => { setState((prev) => ({ ...prev, currentSubagent: null, })); }, 1000); break; } } } catch { // 忽略解析错误 } }; wsRef.current = ws; }, [sessionId]); // 发送消息 const sendMessage = useCallback( (content: string) => { if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) { onErrorRef.current?.(new Error('WebSocket not connected')); return; } setState((prev) => ({ ...prev, isLoading: true })); wsRef.current.send( JSON.stringify({ type: 'message', sessionId, payload: { content, agentMode: state.agentMode, autoApprove: state.autoApprove, }, }) ); }, [sessionId, state.agentMode, state.autoApprove] ); // 取消处理 const cancelProcessing = useCallback(() => { if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return; wsRef.current.send( JSON.stringify({ type: 'cancel', sessionId, }) ); // 保留已流式输出的内容作为消息 setState((prev) => { const streaming = prev.streamingMessage; // 如果有流式消息且有内容,保存到 messages 中 if (streaming && (streaming.content || streaming.parts.length > 0)) { const cancelledMessage: Message = { ...streaming, id: streaming.id || `cancelled-${Date.now()}`, // 标记为已取消(可选:在 content 末尾添加提示) content: streaming.content || '', }; return { ...prev, messages: [...prev.messages, cancelledMessage], isLoading: false, streamingMessage: null, }; } return { ...prev, isLoading: false, streamingMessage: null }; }); }, [sessionId]); // 发送权限响应 const respondToPermission = useCallback( (requestId: string, allow: boolean, remember?: boolean) => { if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) { onErrorRef.current?.(new Error('WebSocket not connected')); return; } wsRef.current.send( JSON.stringify({ type: 'permission_response', sessionId, payload: { requestId, allow, remember }, }) ); // 清除权限请求状态 setState((prev) => ({ ...prev, permissionRequest: null })); }, [sessionId] ); // 允许权限请求 const allowPermission = useCallback( (requestId: string, remember?: boolean) => { respondToPermission(requestId, true, remember); }, [respondToPermission] ); // 拒绝权限请求 const denyPermission = useCallback( (requestId: string, remember?: boolean) => { respondToPermission(requestId, false, remember); }, [respondToPermission] ); // 设置 Agent 模式 (会话级别) const setAgentMode = useCallback((mode: AgentModeType) => { setState((prev) => ({ ...prev, agentMode: mode })); }, []); // 设置自动授权 (会话级别,立即生效) const setAutoApprove = useCallback( (enabled: boolean) => { setState((prev) => ({ ...prev, autoApprove: enabled })); // 立即通过 WebSocket 发送配置更新,使其对当前执行生效 if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) { wsRef.current.send( JSON.stringify({ type: 'config_update', sessionId, payload: { autoApprove: enabled }, }) ); } }, [sessionId] ); // 初始化 useEffect(() => { // 重置状态 isClosingRef.current = false; setState({ messages: [], isConnected: false, isLoading: false, streamingMessage: null, permissionRequest: null, agentMode: 'build', autoApprove: false, currentAgent: 'build', currentSubagent: null, }); reconnectAttemptsRef.current = 0; loadMessages(); connect(); return () => { clearTimeout(reconnectTimeoutRef.current); // 标记为主动关闭,避免触发错误回调和重连 isClosingRef.current = true; // 只关闭已建立的连接 if (wsRef.current) { const ws = wsRef.current; // 清除引用,防止后续操作 wsRef.current = null; // 清除事件处理器,避免关闭时触发错误 ws.onclose = null; ws.onerror = null; ws.onmessage = null; ws.onopen = null; // 只关闭已经建立的连接,避免 "closed before established" 警告 if (ws.readyState === WebSocket.OPEN) { ws.close(); } // 对于 CONNECTING 状态,等待连接建立后再关闭 // 这样可以避免浏览器警告 } }; }, [loadMessages, connect]); return { ...state, sendMessage, cancelProcessing, reload: loadMessages, allowPermission, denyPermission, setAgentMode, setAutoApprove, }; }