@backmatter/studio 0.0.1

package/src/chat/ChatProvider.tsx

@@ -0,0 +1,661 @@
+// packages/studio/src/chat/ChatProvider.tsx
+//
+// Global chat state: panel mode, thread selection, model, context chips, and the
+// active chat session (AI SDK useChat). Mounted once in AppShell for workspace routes.
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useReducer,
+  useRef,
+  useState,
+} from "react";
+import { useChat } from "@ai-sdk/react";
+import {
+  DefaultChatTransport,
+  getToolName,
+  isToolUIPart,
+  type UIMessage,
+} from "ai";
+import {
+  aiMessagesUrl,
+  createThread,
+  fetchAiModels,
+  fetchThread,
+  keepAllChanges,
+  keepChange,
+  notifyHeadChanged,
+  notifyTreeChanged,
+  undoAllChanges,
+  undoChange,
+  type AiModelInfo,
+  type AiTier,
+  type ApiError,
+  type ChatChangeSummary,
+  type ChatThread,
+} from "../api";
+import { apiCredentials, apiHeaders } from "../studio-config";
+import {
+  docKey,
+  notifyDocChanged,
+  useEditorFocus,
+  type DocRef,
+  type EditorFocus,
+} from "../chrome/EditorFocusContext";
+import {
+  addMention as addMentionTo,
+  buildChatContext,
+  removeMention as removeMentionFrom,
+  type ChatRequestContext,
+} from "./chat-context";
+import {
+  changeIdsInMessage,
+  isWriteTool,
+  treeStatusPatches,
+  type TreeStatusPatch,
+  type WriteToolOutput,
+} from "./chat-helpers";
+import {
+  initialPanelState,
+  panelReducer,
+  type PanelAction,
+  type PanelState,
+} from "./chat-state";
+
+// ---------------------------------------------------------------------------
+// UI-level context (panel, thread selection, model, chips)
+// ---------------------------------------------------------------------------
+
+export type ChatUIValue = {
+  workspaceId: string;
+  panel: PanelState;
+  dispatch: (action: PanelAction) => void;
+  /** null while loading; false when the server has no AI provider configured. */
+  aiAvailable: boolean | null;
+  models: AiModelInfo[];
+  tier: AiTier;
+  model: string | null;
+  setModel: (id: string) => void;
+  threadId: string | null;
+  openThread: (threadId: string) => void;
+  newThread: () => void;
+  mentions: DocRef[];
+  addMention: (ref: DocRef) => void;
+  removeMention: (ref: DocRef) => void;
+  includeFocus: boolean;
+  setIncludeFocus: (v: boolean) => void;
+  focus: EditorFocus | null;
+  /** Bumped whenever a thread is created/renamed/deleted so history can refetch. */
+  threadsVersion: number;
+  bumpThreads: () => void;
+};
+
+const ChatUIContext = createContext<ChatUIValue | null>(null);
+
+export function useChatUI(): ChatUIValue {
+  const ctx = useContext(ChatUIContext);
+  if (!ctx) throw new Error("useChatUI must be used within ChatProvider");
+  return ctx;
+}
+
+/** For chrome that renders on routes with and without a workspace. */
+export function useChatUIOptional(): ChatUIValue | null {
+  return useContext(ChatUIContext);
+}
+
+// ---------------------------------------------------------------------------
+// Session-level context (messages, status, changes)
+// ---------------------------------------------------------------------------
+
+export type ChatSessionValue = {
+  threadId: string | null;
+  thread: ChatThread | null;
+  loading: boolean;
+  messages: UIMessage[];
+  status: "submitted" | "streaming" | "ready" | "error";
+  error: Error | undefined;
+  clearError: () => void;
+  send: (text: string) => Promise<void>;
+  stop: () => void;
+  changes: Record<string, ChatChangeSummary>;
+  keep: (changeId: string) => Promise<void>;
+  undo: (changeId: string) => Promise<void>;
+  keepAll: (messageId: string) => Promise<void>;
+  undoAll: (messageId: string) => Promise<void>;
+  /** Change ids attached to a message id (from tool outputs seen in this session). */
+  changeIdsForMessage: (message: UIMessage) => string[];
+};
+
+const ChatSessionContext = createContext<ChatSessionValue | null>(null);
+
+export function useChatSession(): ChatSessionValue {
+  const ctx = useContext(ChatSessionContext);
+  if (!ctx) throw new Error("useChatSession must be used within ChatProvider");
+  return ctx;
+}
+
+// ---------------------------------------------------------------------------
+
+const modelStorageKey = (ws: string) => `backmatter-chat-model:${ws}`;
+
+/** Apply tool statuses locally so the top bar moves before the next git poll. */
+function publishTreePatches(workspaceId: string, patches: TreeStatusPatch[]) {
+  let published = false;
+  for (const patch of patches) {
+    notifyTreeChanged({
+      workspaceId,
+      locale: patch.locale,
+      type: patch.type,
+      slug: patch.slug,
+      ...(patch.absent ? { remove: true } : { status: patch.status }),
+    });
+    // A published file stays on screen; head-changed reloads it once. A file
+    // that disappeared still needs the open editor to drop it.
+    if (patch.absent || patch.status !== "clean") {
+      notifyDocChanged({ locale: patch.locale, type: patch.type, slug: patch.slug });
+    }
+    if (patch.status === "clean") published = true;
+  }
+  if (published) notifyHeadChanged();
+}
+
+function refFromInput(input: unknown): DocRef | null {
+  const ref = (input as { docRef?: Partial<DocRef> } | undefined)?.docRef;
+  if (!ref || !ref.locale || !ref.type || !ref.slug) return null;
+  return { locale: ref.locale, type: ref.type, slug: ref.slug };
+}
+
+export function ChatProvider({
+  workspaceId,
+  startOpen = false,
+  children,
+}: {
+  workspaceId: string;
+  startOpen?: boolean;
+  children: React.ReactNode;
+}) {
+  const [panel, dispatch] = useReducer(
+    panelReducer,
+    startOpen ? { ...initialPanelState, mode: "open" as const } : initialPanelState,
+  );
+  const [aiAvailable, setAiAvailable] = useState<boolean | null>(null);
+  const [models, setModels] = useState<AiModelInfo[]>([]);
+  const [tier, setTier] = useState<AiTier>("ask");
+  const [model, setModelState] = useState<string | null>(null);
+  const [threadId, setThreadId] = useState<string | null>(null);
+  const [freshThread, setFreshThread] = useState(false);
+  const [pendingText, setPendingText] = useState<string | null>(null);
+  const [mentions, setMentions] = useState<DocRef[]>([]);
+  const [includeFocus, setIncludeFocus] = useState(true);
+  const [threadsVersion, setThreadsVersion] = useState(0);
+  const { focus } = useEditorFocus();
+
+  // Models + availability.
+  useEffect(() => {
+    let cancelled = false;
+    setAiAvailable(null);
+    fetchAiModels(workspaceId)
+      .then((res) => {
+        if (cancelled) return;
+        setModels(res.models);
+        setTier(res.tier);
+        setAiAvailable(true);
+        const stored = window.localStorage.getItem(modelStorageKey(workspaceId));
+        const valid = res.models.some((m) => m.id === stored);
+        setModelState(valid ? stored : res.defaultModel);
+      })
+      .catch((err: ApiError) => {
+        if (cancelled) return;
+        setAiAvailable(false);
+        if (err.status !== 503) console.error("ai models failed", err);
+      });
+    return () => {
+      cancelled = true;
+    };
+  }, [workspaceId]);
+
+  // Reset per-workspace session state.
+  useEffect(() => {
+    setThreadId(null);
+    setFreshThread(false);
+    setPendingText(null);
+    setMentions([]);
+    setIncludeFocus(true);
+  }, [workspaceId]);
+
+  const setModel = useCallback(
+    (id: string) => {
+      setModelState(id);
+      window.localStorage.setItem(modelStorageKey(workspaceId), id);
+    },
+    [workspaceId],
+  );
+
+  const openThread = useCallback((id: string) => {
+    setThreadId(id);
+    setFreshThread(false);
+    setPendingText(null);
+    dispatch({ type: "showChat" });
+  }, []);
+
+  const newThread = useCallback(() => {
+    setThreadId(null);
+    setFreshThread(false);
+    setPendingText(null);
+    setMentions([]);
+    dispatch({ type: "showChat" });
+  }, []);
+
+  const addMention = useCallback((ref: DocRef) => {
+    setMentions((prev) => addMentionTo(prev, ref));
+  }, []);
+  const removeMention = useCallback((ref: DocRef) => {
+    setMentions((prev) => removeMentionFrom(prev, ref));
+  }, []);
+  const bumpThreads = useCallback(() => setThreadsVersion((v) => v + 1), []);
+
+  // Context getter kept in a ref so the session never re-renders because of it.
+  const contextRef = useRef<() => ChatRequestContext>(() => ({
+    focusDoc: null,
+    selection: null,
+    mentions: [],
+  }));
+  contextRef.current = () => buildChatContext({ focus, includeFocus, mentions });
+  const modelRef = useRef<string | null>(model);
+  modelRef.current = model;
+
+  // Creating a thread lazily on first send.
+  const startThread = useCallback(
+    async (text: string) => {
+      const res = await createThread(workspaceId);
+      setPendingText(text);
+      setFreshThread(true);
+      setThreadId(res.thread.id);
+      bumpThreads();
+    },
+    [workspaceId, bumpThreads],
+  );
+
+  const ui = useMemo<ChatUIValue>(
+    () => ({
+      workspaceId,
+      panel,
+      dispatch,
+      aiAvailable,
+      models,
+      tier,
+      model,
+      setModel,
+      threadId,
+      openThread,
+      newThread,
+      mentions,
+      addMention,
+      removeMention,
+      includeFocus,
+      setIncludeFocus,
+      focus,
+      threadsVersion,
+      bumpThreads,
+    }),
+    [
+      workspaceId,
+      panel,
+      aiAvailable,
+      models,
+      tier,
+      model,
+      setModel,
+      threadId,
+      openThread,
+      newThread,
+      mentions,
+      addMention,
+      removeMention,
+      includeFocus,
+      focus,
+      threadsVersion,
+      bumpThreads,
+    ],
+  );
+
+  // The session subtree is keyed by thread and remounts when a chat starts or
+  // switches. Only the chat UI (see ChatSessionScope) lives inside it, so the
+  // editor and the rest of the app never remount because of chat state.
+  const renderSession = useCallback(
+    (scoped: React.ReactNode) =>
+      threadId ? (
+        <ChatSession
+          key={threadId}
+          workspaceId={workspaceId}
+          threadId={threadId}
+          fresh={freshThread}
+          pendingText={pendingText}
+          onPendingConsumed={() => setPendingText(null)}
+          getContext={() => contextRef.current()}
+          getModel={() => modelRef.current}
+          onThreadTouched={bumpThreads}
+        >
+          {scoped}
+        </ChatSession>
+      ) : (
+        <EmptySession onFirstSend={startThread}>{scoped}</EmptySession>
+      ),
+    [threadId, workspaceId, freshThread, pendingText, bumpThreads, startThread],
+  );
+
+  return (
+    <ChatUIContext.Provider value={ui}>
+      <SessionHostContext.Provider value={renderSession}>{children}</SessionHostContext.Provider>
+    </ChatUIContext.Provider>
+  );
+}
+
+const SessionHostContext = createContext<((scoped: React.ReactNode) => React.ReactNode) | null>(
+  null,
+);
+
+/**
+ * Mounts the per-thread chat session around its children. Wrap only the chat UI
+ * (bar + panel) in this, never page content: the session remounts per thread.
+ */
+export function ChatSessionScope({ children }: { children: React.ReactNode }) {
+  const render = useContext(SessionHostContext);
+  if (!render) throw new Error("ChatSessionScope must be inside ChatProvider");
+  return <>{render(children)}</>;
+}
+
+// ---------------------------------------------------------------------------
+
+function EmptySession({
+  onFirstSend,
+  children,
+}: {
+  onFirstSend: (text: string) => Promise<void>;
+  children: React.ReactNode;
+}) {
+  const [error, setError] = useState<Error | undefined>();
+  const [status, setStatus] = useState<ChatSessionValue["status"]>("ready");
+  const value = useMemo<ChatSessionValue>(
+    () => ({
+      threadId: null,
+      thread: null,
+      loading: false,
+      messages: [],
+      status,
+      error,
+      clearError: () => setError(undefined),
+      send: async (text) => {
+        setStatus("submitted");
+        try {
+          await onFirstSend(text);
+        } catch (err) {
+          setError(err instanceof Error ? err : new Error(String(err)));
+          setStatus("error");
+        }
+      },
+      stop: () => {},
+      changes: {},
+      keep: async () => {},
+      undo: async () => {},
+      keepAll: async () => {},
+      undoAll: async () => {},
+      changeIdsForMessage: () => [],
+    }),
+    [onFirstSend, error, status],
+  );
+  return <ChatSessionContext.Provider value={value}>{children}</ChatSessionContext.Provider>;
+}
+
+type SessionProps = {
+  workspaceId: string;
+  threadId: string;
+  fresh: boolean;
+  pendingText: string | null;
+  onPendingConsumed: () => void;
+  getContext: () => ChatRequestContext;
+  getModel: () => string | null;
+  onThreadTouched: () => void;
+  children: React.ReactNode;
+};
+
+/** Loads thread history (unless the thread was just created), then mounts useChat. */
+function ChatSession(props: SessionProps) {
+  const { workspaceId, threadId, fresh } = props;
+  const [loaded, setLoaded] = useState(fresh);
+  const [thread, setThread] = useState<ChatThread | null>(null);
+  const [initialMessages, setInitialMessages] = useState<UIMessage[]>([]);
+  const [initialChanges, setInitialChanges] = useState<ChatChangeSummary[]>([]);
+  const [loadError, setLoadError] = useState<Error | undefined>();
+
+  useEffect(() => {
+    if (fresh) return;
+    let cancelled = false;
+    fetchThread(workspaceId, threadId)
+      .then((res) => {
+        if (cancelled) return;
+        setThread(res.thread);
+        setInitialMessages(res.messages as UIMessage[]);
+        setInitialChanges(res.changes);
+        setLoaded(true);
+      })
+      .catch((err: Error) => {
+        if (!cancelled) setLoadError(err);
+      });
+    return () => {
+      cancelled = true;
+    };
+  }, [workspaceId, threadId, fresh]);
+
+  if (!loaded) {
+    return (
+      <ChatSessionContext.Provider
+        value={{
+          threadId,
+          thread: null,
+          loading: !loadError,
+          messages: [],
+          status: loadError ? "error" : "ready",
+          error: loadError,
+          clearError: () => setLoadError(undefined),
+          send: async () => {},
+          stop: () => {},
+          changes: {},
+          keep: async () => {},
+          undo: async () => {},
+          keepAll: async () => {},
+          undoAll: async () => {},
+          changeIdsForMessage: () => [],
+        }}
+      >
+        {props.children}
+      </ChatSessionContext.Provider>
+    );
+  }
+
+  return (
+    <LiveSession
+      {...props}
+      thread={thread}
+      initialMessages={initialMessages}
+      initialChanges={initialChanges}
+    />
+  );
+}
+
+function LiveSession({
+  workspaceId,
+  threadId,
+  thread,
+  initialMessages,
+  initialChanges,
+  pendingText,
+  onPendingConsumed,
+  getContext,
+  getModel,
+  onThreadTouched,
+  children,
+}: SessionProps & {
+  thread: ChatThread | null;
+  initialMessages: UIMessage[];
+  initialChanges: ChatChangeSummary[];
+}) {
+  const { setAiEditingDocKey } = useEditorFocus();
+  const [changes, setChanges] = useState<Record<string, ChatChangeSummary>>(() =>
+    Object.fromEntries(initialChanges.map((c) => [c.id, c])),
+  );
+
+  const transport = useMemo(
+    () =>
+      new DefaultChatTransport<UIMessage>({
+        api: aiMessagesUrl(workspaceId, threadId),
+        credentials: apiCredentials(),
+        headers: apiHeaders(),
+        prepareSendMessagesRequest: ({ messages, body }) => ({
+          body: {
+            message: messages[messages.length - 1],
+            ...(body ?? {}),
+          },
+        }),
+      }),
+    [workspaceId, threadId],
+  );
+
+  const chat = useChat<UIMessage>({
+    id: threadId,
+    messages: initialMessages,
+    transport,
+    onFinish: () => onThreadTouched(),
+  });
+
+  const { messages, status, sendMessage } = chat;
+
+  const send = useCallback(
+    async (text: string) => {
+      await sendMessage(
+        { text },
+        { body: { model: getModel(), context: getContext() } },
+      );
+    },
+    [sendMessage, getModel, getContext],
+  );
+
+  // Send the message that created this thread.
+  const sentPending = useRef(false);
+  useEffect(() => {
+    if (pendingText === null || sentPending.current) return;
+    sentPending.current = true;
+    onPendingConsumed();
+    void send(pendingText);
+  }, [pendingText, onPendingConsumed, send]);
+
+  // Watch tool parts: mark the focused doc as "AI editing" while a write tool
+  // targets it, and notify the editor/preview when a change lands.
+  const seenOutputs = useRef<Set<string> | null>(null);
+  useEffect(() => {
+    // History is already loaded when this session mounts. Ignore those tool
+    // results so an old commit does not clear the current working tree.
+    if (seenOutputs.current === null) {
+      const seeded = new Set<string>();
+      for (const message of messages) {
+        for (const part of message.parts ?? []) {
+          if (isToolUIPart(part) && part.state === "output-available") {
+            seeded.add(part.toolCallId);
+          }
+        }
+      }
+      seenOutputs.current = seeded;
+    }
+    const seen = seenOutputs.current;
+    const streaming = status === "streaming" || status === "submitted";
+    const last = messages[messages.length - 1];
+    let editingKey: string | null = null;
+    if (last?.role === "assistant") {
+      for (const part of last.parts) {
+        if (!isToolUIPart(part)) continue;
+        const name = getToolName(part);
+        const write = isWriteTool(name);
+        const ref = write ? refFromInput(part.input) : null;
+        if (streaming && write && ref) editingKey = docKey(ref);
+        if (part.state === "output-available" && !seen.has(part.toolCallId)) {
+          seen.add(part.toolCallId);
+          const patches = treeStatusPatches(part.output);
+          if (patches.length > 0) publishTreePatches(workspaceId, patches);
+          const out = part.output as WriteToolOutput | undefined;
+          if (write && out?.ok) {
+            setChanges((prev) => ({ ...prev, [out.change.id]: out.change }));
+            if (patches.length === 0) {
+              if (ref) notifyDocChanged(ref);
+              notifyTreeChanged();
+            }
+          }
+        }
+      }
+    }
+    setAiEditingDocKey(streaming ? editingKey : null);
+  }, [messages, status, setAiEditingDocKey, workspaceId]);
+
+  useEffect(() => () => setAiEditingDocKey(null), [setAiEditingDocKey]);
+
+  const applyChangeResult = useCallback((change: ChatChangeSummary, notify: boolean) => {
+    setChanges((prev) => ({ ...prev, [change.id]: { ...prev[change.id], ...change } }));
+    if (notify) {
+      notifyDocChanged({ locale: change.locale, type: change.type, slug: change.slug });
+      notifyTreeChanged();
+    }
+  }, []);
+
+  const keep = useCallback(
+    async (changeId: string) => {
+      const res = await keepChange(workspaceId, changeId);
+      applyChangeResult(res.change, false);
+    },
+    [workspaceId, applyChangeResult],
+  );
+  const undo = useCallback(
+    async (changeId: string) => {
+      const res = await undoChange(workspaceId, changeId);
+      applyChangeResult(res.change, true);
+    },
+    [workspaceId, applyChangeResult],
+  );
+  const keepAll = useCallback(
+    async (messageId: string) => {
+      const res = await keepAllChanges(workspaceId, threadId, messageId);
+      for (const c of res.changes) applyChangeResult(c, false);
+    },
+    [workspaceId, threadId, applyChangeResult],
+  );
+  const undoAll = useCallback(
+    async (messageId: string) => {
+      const res = await undoAllChanges(workspaceId, threadId, messageId);
+      for (const c of res.changes) applyChangeResult(c, true);
+    },
+    [workspaceId, threadId, applyChangeResult],
+  );
+
+  const value = useMemo<ChatSessionValue>(
+    () => ({
+      threadId,
+      thread,
+      loading: false,
+      messages,
+      status,
+      error: chat.error,
+      clearError: chat.clearError,
+      send,
+      stop: chat.stop,
+      changes,
+      keep,
+      undo,
+      keepAll,
+      undoAll,
+      changeIdsForMessage: changeIdsInMessage,
+    }),
+    [threadId, thread, messages, status, chat.error, chat.clearError, send, chat.stop, changes, keep, undo, keepAll, undoAll],
+  );
+
+  return <ChatSessionContext.Provider value={value}>{children}</ChatSessionContext.Provider>;
+}