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

import { useCallback, useEffect, useRef, useState } from "react";
import { environment } from "@vicinae/api";
import { runAgent } from "../agent";
import { confirmWithAlert, createApprovalLedger } from "../approval";
import { toStoredAttachment } from "../attachments";
import { createPolicy, displayPath } from "../files";
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 { EMPTY_TRACE, afterToolCall, withFinalTrace, withThinking, type ThinkingTrace } from "../thinking";
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;
  /**
   * The reasoning to show right now: only the part since the last tool call.
   *
   * Deliberately not the whole trace. The full one is carried on the turn for
   * storage; this is the live view, and it is cleared each time a tool starts so
   * a long research turn does not stack every round's thinking on screen.
   */
  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<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 [activity, setActivity] = useState<ToolActivity | undefined>(undefined);
  const [error, setError] = useState<string | undefined>(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<VisionSupport>(options.vision ?? "unknown");
  const abortRef = useRef<AbortController | null>(null);
  const liveRef = useRef(false);
  const responseRef = useRef("");
  const thinkingRef = useRef<ThinkingTrace>(EMPTY_TRACE);
  const statusRef = useRef<string | undefined>(undefined);
  const flushTimerRef = useRef<ReturnType<typeof setTimeout> | 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.current);
    }, FLUSH_INTERVAL_MS);
  }, []);

  const flushNow = useCallback(() => {
    clearFlushTimer();
    setStreamingResponse(responseRef.current);
    setStreamingThinking(thinkingRef.current.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 = EMPTY_TRACE;
      setStreamingResponse("");
      setStreamingThinking("");
      setError(undefined);
      setActivity(undefined);
      setStatusValue("Thinking");
      liveRef.current = true;
      setIsStreaming(true);

      const controller = new AbortController();
      abortRef.current = controller;

      // Built per turn rather than per session: the ledger inside it is the record of
      // what the user approved while answering this one question, and the next question
      // starts by asking again. The policy is cheap to rebuild and holds no state.
      const policy = createPolicy({
        supportPath: environment.supportPath,
        readRoots: preferences.extraReadRoots,
        writeRoots: preferences.extraWriteRoots,
      });
      const fileAccess = preferences.fileEnabled
        ? { policy, approve: createApprovalLedger(confirmWithAlert).approve }
        : undefined;

      const messages = buildMessages(
        composeSystemPrompt(preferences.systemPrompt, {
          model: preferences.model,
          host: preferences.host,
          selfHosted: isSelfHosted(preferences.host),
          webTools: preferences.webEnabled,
          vision: visionRef.current,
          fileTools: preferences.fileEnabled,
          fileRoots: {
            read: policy.readRoots.map((root) => displayPath(root, policy.home)),
            write: policy.writeRoots.map((root) => displayPath(root, policy.home)),
          },
        }),
        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,
              fileEnabled: preferences.fileEnabled,
            }),
            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,
              files: fileAccess,
            },
            signal: controller.signal,
            events: {
              onContentDelta: (delta) => {
                responseRef.current += delta;
                setStatusValue(undefined);
                scheduleFlush();
              },
              onThinkingDelta: (delta) => {
                thinkingRef.current = withThinking(thinkingRef.current, delta);
                scheduleFlush();
              },
              onStatus: (value) => setStatusValue(value),
              onToolStart: (_name, label, count) => {
                // The reasoning that produced this call is spent, so the live
                // segment is cleared and the tool line takes its place. The full
                // trace is untouched, because the stored turn keeps all of it.
                thinkingRef.current = afterToolCall(thinkingRef.current);
                setActivity({ label, count });
                flushNow();
              },
              onToolCall: (record) => {
                toolCalls = [...toolCalls, record];
              },
            },
          });

          responseRef.current = result.content;
          thinkingRef.current = withFinalTrace(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.full,
          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("");
      thinkingRef.current = EMPTY_TRACE;
      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,
  };
}