feat: initial Lam Agent extension (#1)

src/lib/hooks/useChat.ts

@@ -0,0 +1,268 @@
+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,
+  };
+}