import { useCallback, useEffect, useRef, useState } from "react"; import { runAgent } from "../agent"; import { buildMessages, composeSystemPrompt } from "../history"; import { OllamaError } from "../ollama"; import type { ResolvedPreferences } from "../preferences"; import { createSession, createTurn, saveSession, withTitleFromFirstTurn } from "../sessions"; import { toolsFor } from "../tools"; import type { Session } from "../types"; import { errorMessage, isAbortError } from "../utils"; /** * The chat state machine. * * Kept out of the component so `chat.tsx` stays a layout file, and so the * streaming behaviour has one obvious home. * * Two details are worth calling out: * * 1. Streaming deltas are accumulated in refs and flushed to React state on a * timer. A long answer arrives as thousands of tokens, and re-rendering the * whole transcript per token makes the launcher visibly janky. * 2. The live session is mirrored in a ref. `ask` needs the session as it is at * call time, and reading it from state inside a stable callback would close * over a stale value. */ /** Coalescing window for streamed deltas. Short enough to feel live, long enough to stay smooth. */ const FLUSH_INTERVAL_MS = 60; export interface UseChatOptions { preferences: ResolvedPreferences; /** Restores an existing conversation, e.g. when reopening from the sessions command. */ initialSession?: Session; } export interface UseChat { session: Session; isStreaming: boolean; streamingResponse: string; streamingThinking: string; /** Short activity line such as "Searching the web", shown while a turn is in flight. */ status?: string; error?: string; ask: (prompt: string) => void; cancel: () => void; replaceSession: (session: Session) => void; newSession: () => void; setModel: (model: string) => void; } export function useChat(options: UseChatOptions): UseChat { const [session, setSession] = useState( () => options.initialSession ?? createSession(options.preferences.model), ); const [isStreaming, setIsStreaming] = useState(false); const [streamingResponse, setStreamingResponse] = useState(""); const [streamingThinking, setStreamingThinking] = useState(""); const [status, setStatus] = useState(undefined); const [error, setError] = useState(undefined); const sessionRef = useRef(session); const preferencesRef = useRef(options.preferences); const abortRef = useRef(null); const liveRef = useRef(false); const responseRef = useRef(""); const thinkingRef = useRef(""); const statusRef = useRef(undefined); const flushTimerRef = useRef | null>(null); preferencesRef.current = options.preferences; /** Updates state and the ref together, so `ask` never reads a stale session. */ const commit = useCallback((next: Session) => { sessionRef.current = next; setSession(next); }, []); const clearFlushTimer = useCallback(() => { if (flushTimerRef.current !== null) { clearTimeout(flushTimerRef.current); flushTimerRef.current = null; } }, []); const scheduleFlush = useCallback(() => { if (flushTimerRef.current !== null) return; flushTimerRef.current = setTimeout(() => { flushTimerRef.current = null; if (!liveRef.current) return; setStreamingResponse(responseRef.current); setStreamingThinking(thinkingRef.current); }, FLUSH_INTERVAL_MS); }, []); const flushNow = useCallback(() => { clearFlushTimer(); setStreamingResponse(responseRef.current); setStreamingThinking(thinkingRef.current); }, [clearFlushTimer]); const setStatusValue = useCallback((value: string | undefined) => { if (statusRef.current === value) return; statusRef.current = value; setStatus(value); }, []); const cancel = useCallback(() => { abortRef.current?.abort(); }, []); /** Aborts any in-flight request when the command unmounts, so no stream survives the view. */ useEffect(() => { return () => { liveRef.current = false; abortRef.current?.abort(); if (flushTimerRef.current !== null) clearTimeout(flushTimerRef.current); }; }, []); const ask = useCallback( (rawPrompt: string) => { const prompt = rawPrompt.trim(); if (!prompt || liveRef.current) return; const preferences = preferencesRef.current; const base = sessionRef.current; // The turn is inserted immediately with an empty response so the UI shows // the question while the answer streams in beneath it. const turn = createTurn(prompt, preferences.model); const started = withTitleFromFirstTurn({ ...base, model: preferences.model, turns: [...base.turns, turn], }); commit(started); responseRef.current = ""; thinkingRef.current = ""; setStreamingResponse(""); setStreamingThinking(""); setError(undefined); setStatusValue("Thinking"); liveRef.current = true; setIsStreaming(true); const controller = new AbortController(); abortRef.current = controller; const messages = buildMessages( composeSystemPrompt(preferences.systemPrompt, preferences.webEnabled), base.turns, prompt, ); void (async () => { let failure: string | undefined; let toolCalls = turn.toolCalls; try { const result = await runAgent({ host: preferences.host, apiKey: preferences.apiKey, model: preferences.model, think: preferences.think, messages, tools: toolsFor({ webEnabled: preferences.webEnabled }), toolContext: { host: preferences.host, apiKey: preferences.apiKey }, signal: controller.signal, events: { onContentDelta: (delta) => { responseRef.current += delta; setStatusValue(undefined); scheduleFlush(); }, onThinkingDelta: (delta) => { thinkingRef.current += delta; scheduleFlush(); }, onStatus: (value) => setStatusValue(value), onToolCall: (record) => { toolCalls = [...toolCalls, record]; }, }, }); responseRef.current = result.content; thinkingRef.current = result.thinking; toolCalls = result.toolCalls; } catch (thrown) { if (isAbortError(thrown)) { failure = "Cancelled."; } else { const hint = thrown instanceof OllamaError ? thrown.hint : undefined; failure = hint ? `${errorMessage(thrown)} ${hint}` : errorMessage(thrown); } } const finished = { ...turn, response: responseRef.current, thinking: thinkingRef.current, toolCalls, error: failure, }; const updated = withTitleFromFirstTurn({ ...started, turns: [...base.turns, finished], updatedAt: Date.now(), }); liveRef.current = false; abortRef.current = null; commit(updated); setIsStreaming(false); setStatusValue(undefined); setError(failure); flushNow(); // Persistence is best-effort: a storage failure should not lose the // answer that is already on screen. void saveSession(updated).catch(() => undefined); })(); }, [commit, flushNow, scheduleFlush, setStatusValue], ); const replaceSession = useCallback( (next: Session) => { abortRef.current?.abort(); liveRef.current = false; setIsStreaming(false); setStatusValue(undefined); setError(undefined); setStreamingResponse(""); setStreamingThinking(""); commit(next); }, [commit, setStatusValue], ); const newSession = useCallback(() => { replaceSession(createSession(preferencesRef.current.model)); }, [replaceSession]); const setModel = useCallback( (model: string) => { commit({ ...sessionRef.current, model, updatedAt: Date.now() }); }, [commit], ); return { session, isStreaming, streamingResponse, streamingThinking, status, error, ask, cancel, replaceSession, newSession, setModel, }; }