import { useCallback, useEffect, useRef, useState } from "react"; import { runAgent } from "../agent"; import { toStoredAttachment } from "../attachments"; import { ipPlace } from "../geo"; import { buildMessages, composeSystemPrompt } from "../history"; import { isSelfHosted } from "../ollama"; import { OllamaError } from "../ollama"; import type { ResolvedPreferences } from "../preferences"; import { createSession, createTurn, saveSession, withTitleFromFirstTurn } from "../sessions"; import { toolsFor } from "../tools"; import type { PendingAttachment, Session, VisionSupport } 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; /** * Whether the model can read images, for the system prompt. * * Passed in rather than checked here because the answer is already known to the * view: the snip action is gated on the same value, and asking twice would mean * two capability requests per model change. */ vision?: VisionSupport; } /** Which tool is running right now, and how stale it is getting. */ export interface ToolActivity { /** Human-readable label, e.g. "Searching the web". */ label: string; /** * How many times this turn has now started this *identical* call. * * Counted by the agent over name and arguments together, so two different * searches are two tools doing their job rather than a repeat. */ count: number; } 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; /** * The tool running right now. * * Separate from {@link UseChat.status} because `status` goes quiet as soon as * the answer starts arriving, and the whole point of this one is to stay on * screen: a model that has started circling looks identical to a model that is * working, unless the tool and its repeat count stay visible. */ activity?: ToolActivity; error?: string; ask: (prompt: string, attachment?: PendingAttachment) => 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 [activity, setActivity] = useState(undefined); const [error, setError] = useState(undefined); const sessionRef = useRef(session); const preferencesRef = useRef(options.preferences); /** * Read through a ref for the same reason as the preferences: `ask` is a stable * callback, and the capability check lands after the first render. */ const visionRef = useRef(options.vision ?? "unknown"); 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; // Only ever forwards. The check starts as "unknown", and a value arriving is // news; going back to "unknown" would be the view telling us less than before. if (options.vision) visionRef.current = options.vision; /** 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, attachment?: PendingAttachment) => { const prompt = rawPrompt.trim(); if (!prompt || liveRef.current) return; const preferences = preferencesRef.current; const base = sessionRef.current; // Only the path, mime type and dimensions are kept on the turn. The base64 // payload is used for this one request and then dropped, so it never // reaches storage. The dimensions have to survive because the transcript // lays the preview out long after the payload is gone. const attachments = attachment ? [toStoredAttachment(attachment)] : []; // 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, Date.now(), attachments); const started = withTitleFromFirstTurn({ ...base, model: preferences.model, turns: [...base.turns, turn], }); commit(started); responseRef.current = ""; thinkingRef.current = ""; setStreamingResponse(""); setStreamingThinking(""); setError(undefined); setActivity(undefined); setStatusValue("Thinking"); liveRef.current = true; setIsStreaming(true); const controller = new AbortController(); abortRef.current = controller; const messages = buildMessages( composeSystemPrompt(preferences.systemPrompt, { model: preferences.model, host: preferences.host, selfHosted: isSelfHosted(preferences.host), webTools: preferences.webEnabled, vision: visionRef.current, }), base.turns, prompt, attachment ? [attachment.base64] : undefined, ); 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, // `undefined` is "unlimited" in the preferences, and `null` is what the // agent reads as "no cap": an omitted option there means the default // cap, so the two conventions meet here. maxToolRounds: preferences.maxToolRounds ?? null, messages, tools: toolsFor({ webEnabled: preferences.webEnabled }), toolContext: { host: preferences.host, apiKey: preferences.apiKey, place: preferences.place, // Only wired up when the user switched it on, so an off setting // cannot reach the network for a location at all. ipPlace: preferences.ipGeolocation ? ipPlace : undefined, }, signal: controller.signal, events: { onContentDelta: (delta) => { responseRef.current += delta; setStatusValue(undefined); scheduleFlush(); }, onThinkingDelta: (delta) => { thinkingRef.current += delta; scheduleFlush(); }, onStatus: (value) => setStatusValue(value), onToolStart: (_name, label, count) => setActivity({ label, count }), 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); setActivity(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); setActivity(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, activity, error, ask, cancel, replaceSession, newSession, setModel, }; }