import { useState, useEffect, useRef, useCallback } from 'react'; import { GatewayClient, type JsonPayload } from '../lib/gateway'; import { genIdempotencyKey } from '../lib/utils'; import { getStoredCredentials, storeCredentials, clearCredentials } from '../lib/credentials'; import type { ChatMessage, MessageBlock, ConnectionStatus, Session } from '../types'; interface ChatPayloadMessage { content?: string | Array<{ type: string; text?: string }>; } function extractText(message: ChatPayloadMessage | undefined): string { if (!message) return ''; const content = message.content; if (typeof content === 'string') return content; if (Array.isArray(content)) { return content .filter((b) => b.type === 'text' && typeof b.text === 'string') .map((b) => b.text as string) .join('\n'); } return ''; } export function useGateway() { const clientRef = useRef(null); const [status, setStatus] = useState('disconnected'); const [messages, setMessages] = useState([]); const [sessions, setSessions] = useState([]); const [activeSession, setActiveSession] = useState('agent:main:main'); const [isGenerating, setIsGenerating] = useState(false); const [authenticated, setAuthenticated] = useState(null); // null = checking const [connectError, setConnectError] = useState(null); const [isConnecting, setIsConnecting] = useState(false); const isConnectingRef = useRef(false); const messagesRef = useRef(messages); const activeSessionRef = useRef(activeSession); // Sync refs in an effect to avoid ref writes during render useEffect(() => { messagesRef.current = messages; }, [messages]); useEffect(() => { activeSessionRef.current = activeSession; }, [activeSession]); const currentRunIdRef = useRef(null); const [activeSessions, setActiveSessions] = useState>(new Set()); const handleAgentEvent = useCallback((payload: JsonPayload) => { if (payload?.stream !== 'tool') return; const data = (payload.data ?? {}) as Record; const phase = data.phase as string | undefined; const toolCallId = data.toolCallId as string | undefined; const name = (data.name as string) || 'tool'; if (!toolCallId) return; setMessages(prev => { const last = prev[prev.length - 1]; if (!last || last.role !== 'assistant' || !last.isStreaming) return prev; const updated = { ...last, blocks: [...last.blocks] }; if (phase === 'start') { updated.blocks.push({ type: 'tool_use' as const, name, input: (data.args as Record) ?? {}, id: toolCallId, }); } else if (phase === 'result') { const rawResult = data.result; const result = typeof rawResult === 'string' ? rawResult : JSON.stringify(rawResult, null, 2); updated.blocks.push({ type: 'tool_result' as const, content: result?.slice(0, 500) || '', toolUseId: toolCallId, name, }); } return [...prev.slice(0, -1), updated]; }); }, []); const loadSessions = useCallback(async () => { try { const res = await clientRef.current?.send('sessions.list', {}); const sessionList = res?.sessions as Array> | undefined; if (sessionList) { setSessions(sessionList.map((s) => ({ key: (s.key || s.sessionKey) as string, label: (s.label || s.key || s.sessionKey) as string, messageCount: s.messageCount as number | undefined, totalTokens: s.totalTokens as number | undefined, contextTokens: s.contextTokens as number | undefined, inputTokens: s.inputTokens as number | undefined, outputTokens: s.outputTokens as number | undefined, }))); } } catch { // Silently ignore session list failures (e.g. disconnected) } }, []); const loadHistory = useCallback(async (sessionKey: string) => { try { const res = await clientRef.current?.send('chat.history', { sessionKey, limit: 100 }); const rawMsgs = res?.messages as Array> | undefined; if (rawMsgs) { /* eslint-disable @typescript-eslint/no-explicit-any -- raw gateway history messages have dynamic shape */ const msgs: ChatMessage[] = rawMsgs.map((m: Record, i: number) => { const blocks: MessageBlock[] = []; if (m.content) { if (Array.isArray(m.content)) { for (const block of m.content) { if (block.type === 'text') blocks.push({ type: 'text', text: block.text }); else if (block.type === 'thinking') blocks.push({ type: 'thinking', text: block.thinking || block.text || '' }); else if (block.type === 'image') { const src = block.source || {}; blocks.push({ type: 'image', mediaType: src.media_type || block.media_type || 'image/png', data: src.data || block.data, url: block.url || src.url }); } else if (block.type === 'image_url') { blocks.push({ type: 'image', mediaType: 'image/png', url: block.image_url?.url || block.url }); } else if (block.type === 'tool_use') blocks.push({ type: 'tool_use', name: block.name, input: block.input, id: block.id }); else if (block.type === 'tool_result') blocks.push({ type: 'tool_result', content: typeof block.content === 'string' ? block.content : JSON.stringify(block.content, null, 2), toolUseId: block.tool_use_id }); else if (block.type === 'toolCall') blocks.push({ type: 'tool_use', name: block.name, input: block.arguments || block.input, id: block.id }); else if (block.type === 'toolResult') blocks.push({ type: 'tool_result', content: typeof block.content === 'string' ? block.content : JSON.stringify(block.content, null, 2), toolUseId: block.toolCallId || block.tool_use_id, name: block.name }); } } else if (typeof m.content === 'string') { blocks.push({ type: 'text', text: m.content }); } } const role: 'user' | 'assistant' = m.role === 'user' ? 'user' : 'assistant'; if (m.role === 'toolResult') { const toolBlocks: MessageBlock[] = blocks.map(b => { if (b.type === 'text') { return { type: 'tool_result' as const, content: b.text, toolUseId: m.toolCallId }; } return b; }); return { id: m.id || `hist-${i}`, role: 'assistant' as const, content: '', timestamp: m.timestamp || Date.now(), blocks: toolBlocks, isToolResult: true, }; } return { id: m.id || `hist-${i}`, role, content: blocks.filter((b): b is Extract => b.type === 'text').map(b => b.text).join(''), timestamp: m.timestamp || Date.now(), blocks, }; }); const merged: ChatMessage[] = []; for (const msg of msgs) { const isToolResult = 'isToolResult' in msg && (msg as ChatMessage & { isToolResult?: boolean }).isToolResult; if (isToolResult && merged.length > 0 && merged[merged.length - 1].role === 'assistant') { merged[merged.length - 1] = { ...merged[merged.length - 1], blocks: [...merged[merged.length - 1].blocks, ...msg.blocks], }; } else if (isToolResult) { // skip orphan tool results } else { merged.push(msg); } } setMessages(merged); } } catch { // Silently ignore history load failures } }, []); const setupClient = useCallback((wsUrl: string, token: string) => { // Tear down existing client if (clientRef.current) { clientRef.current.disconnect(); } const client = new GatewayClient(wsUrl, token); clientRef.current = client; client.onStatus((s) => { setStatus(s); if (s === 'connected') { setAuthenticated(true); setConnectError(null); setIsConnecting(false); isConnectingRef.current = false; storeCredentials(wsUrl, token); loadSessions(); loadHistory(activeSessionRef.current); } else if (s === 'disconnected' && !client.isConnected) { // If we never connected successfully, this is an auth/connection error if (isConnectingRef.current) { setConnectError('Connection failed — check URL and token'); setIsConnecting(false); isConnectingRef.current = false; setAuthenticated(false); } } }); client.onEvent((event, payload) => { if (event === 'agent') { handleAgentEvent(payload); return; } if (event !== 'chat') return; const state = payload.state as string | undefined; const runId = payload.runId as string; const message = payload.message as ChatPayloadMessage | undefined; const errorMessage = payload.errorMessage as string | undefined; const evtSession = payload.sessionKey as string | undefined; if (evtSession) { if (state === 'delta') { setActiveSessions(prev => { if (prev.has(evtSession)) return prev; const next = new Set(prev); next.add(evtSession); return next; }); } else if (state === 'final' || state === 'error' || state === 'aborted') { setActiveSessions(prev => { if (!prev.has(evtSession)) return prev; const next = new Set(prev); next.delete(evtSession); return next; }); } } if (evtSession !== activeSessionRef.current) return; if (state === 'delta') { const text = extractText(message); currentRunIdRef.current = runId; setMessages(prev => { const last = prev[prev.length - 1]; if (last && last.role === 'assistant' && last.isStreaming && last.runId === runId) { const updated = { ...last }; updated.content = text; const nonTextBlocks = updated.blocks.filter(b => b.type !== 'text'); updated.blocks = [...nonTextBlocks, { type: 'text' as const, text }]; return [...prev.slice(0, -1), updated]; } const msg: ChatMessage = { id: runId + '-' + Date.now(), role: 'assistant', content: text, timestamp: Date.now(), blocks: [{ type: 'text', text }], isStreaming: true, runId, }; return [...prev, msg]; }); } else if (state === 'final') { currentRunIdRef.current = null; setIsGenerating(false); loadHistory(activeSessionRef.current); } else if (state === 'error') { currentRunIdRef.current = null; setIsGenerating(false); setMessages(prev => { const last = prev[prev.length - 1]; if (last && last.role === 'assistant' && last.isStreaming && last.runId === runId) { return [...prev.slice(0, -1), { ...last, isStreaming: false }]; } return [...prev, { id: 'error-' + Date.now(), role: 'assistant' as const, content: `Error: ${errorMessage || 'unknown error'}`, timestamp: Date.now(), blocks: [{ type: 'text' as const, text: `Error: ${errorMessage || 'unknown error'}` }], }]; }); } else if (state === 'aborted') { currentRunIdRef.current = null; setIsGenerating(false); setMessages(prev => { const last = prev[prev.length - 1]; if (last && last.role === 'assistant' && last.isStreaming) { return [...prev.slice(0, -1), { ...last, isStreaming: false }]; } return prev; }); } }); setIsConnecting(true); isConnectingRef.current = true; setConnectError(null); client.connect(); }, [handleAgentEvent, loadHistory, loadSessions]); // On mount: try stored credentials const initRef = useRef(false); useEffect(() => { if (initRef.current) return; initRef.current = true; const stored = getStoredCredentials(); if (stored) { // Init on mount — setupClient sets state as part of establishing the connection // eslint-disable-next-line react-hooks/set-state-in-effect setupClient(stored.url, stored.token); } else { setAuthenticated(false); } }, [setupClient]); const sendMessage = useCallback(async (text: string, attachments?: Array<{ mimeType: string; fileName: string; content: string }>) => { const userMsg: ChatMessage = { id: 'user-' + Date.now(), role: 'user', content: text, timestamp: Date.now(), blocks: [{ type: 'text', text }], }; setMessages(prev => [...prev, userMsg]); setIsGenerating(true); try { await clientRef.current?.send('chat.send', { sessionKey: activeSessionRef.current, message: text, deliver: false, idempotencyKey: genIdempotencyKey(), ...(attachments && attachments.length > 0 ? { attachments } : {}), }); } catch { // Failed to send — stop generating indicator setIsGenerating(false); } }, []); const abort = useCallback(async () => { try { await clientRef.current?.send('chat.abort', { sessionKey: activeSessionRef.current }); } catch { // Ignore abort failures } setIsGenerating(false); }, []); const switchSession = useCallback((key: string) => { setActiveSession(key); activeSessionRef.current = key; setMessages([]); loadHistory(key); }, [loadHistory]); const login = useCallback((url: string, token: string) => { setupClient(url, token); }, [setupClient]); const logout = useCallback(() => { if (clientRef.current) { clientRef.current.disconnect(); clientRef.current = null; } clearCredentials(); setAuthenticated(false); setMessages([]); setSessions([]); setStatus('disconnected'); setConnectError(null); }, []); // Periodic session refresh every 30s useEffect(() => { if (status !== 'connected') return; const interval = setInterval(loadSessions, 30000); return () => clearInterval(interval); }, [status, loadSessions]); const enrichedSessions = sessions.map(s => ({ ...s, isActive: activeSessions.has(s.key), })); return { status, messages, sessions: enrichedSessions, activeSession, isGenerating, sendMessage, abort, switchSession, loadSessions, authenticated, login, logout, connectError, isConnecting, }; }