lam-agent-vicinae/src/lib/hooks/useChat.ts

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<Session>(
    () => options.initialSession ?? createSession(options.preferences.model),
  );
  const [isStreaming, setIsStreaming] = useState(false);
  const [streamingResponse, setStreamingResponse] = useState("");
  const [streamingThinking, setStreamingThinking] = useState("");
  const [status, setStatus] = useState<string | undefined>(undefined);
  const [error, setError] = useState<string | undefined>(undefined);

  const sessionRef = useRef(session);
  const preferencesRef = useRef(options.preferences);
  const abortRef = useRef<AbortController | null>(null);
  const liveRef = useRef(false);
  const responseRef = useRef("");
  const thinkingRef = useRef("");
  const statusRef = useRef<string | undefined>(undefined);
  const flushTimerRef = useRef<ReturnType<typeof setTimeout> | 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,
  };
}