codecraft/web/src/store.ts

import { useSyncExternalStore } from "react";
import { get } from "./api";
import type {
  QueueItem,
  User,
  McpServerStatus,
  Message,
  Part,
  PermissionRequest,
  QuestionRequest,
  ServerEvent,
  TurnCheckInfo,
  Session,
} from "./types";

/** The assistant message currently being streamed. */
export interface Draft {
  id: string;
  parts: Part[];
}

export interface SessionError {
  id: number;
  message: string;
}

export interface State {
  me: User | null;
  /** Newer release (admins). */
  update: { version: string; url: string; notes: string } | null;
  connected: boolean;
  sessions: Session[];
  /** Includes subagent sessions loaded on demand. */
  sessionById: Record<string, Session>;
  messages: Record<string, Message[]>;
  /** Long chats are loaded in parts: are there older / newer messages than the loaded ones? */
  windows: Record<string, { older: boolean; newer: boolean } | undefined>;
  /** Sessions the current user may only read (shared with them). */
  readOnly: Record<string, boolean>;
  drafts: Record<string, Draft | undefined>;
  busy: Record<string, boolean>;
  progress: Record<string, string>;
  permissions: PermissionRequest[];
  questions: QuestionRequest[];
  errors: Record<string, SessionError[]>;
  mcp: Record<string, McpServerStatus>;
  /** The model request currently running per session (activity line). */
  steps: Record<string, { messageId: string; startedAt: number; promptTokens: number; estimated: boolean } | undefined>;
  /** Messages waiting for a busy session. */
  queue: Record<string, QueueItem[]>;
  /** Why a session waits: "model:<spec>" or "subagent:<name>". */
  waiting: Record<string, string | undefined>;
  /** Compaction running in a session (auto = context nearly full). */
  compacting: Record<string, { auto: boolean; tokens: number | null; limit: number | null; since: number } | undefined>;
  /** Check results per session, by the id of the turn's last answer. */
  checks: Record<string, Record<string, TurnCheckInfo> | undefined>;
}

let state: State = {
  me: null,
  update: null,
  connected: false,
  sessions: [],
  sessionById: {},
  messages: {},
  windows: {},
  readOnly: {},
  drafts: {},
  busy: {},
  progress: {},
  permissions: [],
  questions: [],
  errors: {},
  mcp: {},
  queue: {},
  steps: {},
  waiting: {},
  compacting: {},
  checks: {},
};

const listeners = new Set<() => void>();

export function setState(fn: (s: State) => Partial<State>) {
  state = { ...state, ...fn(state) };
  listeners.forEach((l) => l());
}

export function getState() {
  return state;
}

export function useStore<T>(selector: (s: State) => T): T {
  return useSyncExternalStore(
    (l) => {
      listeners.add(l);
      return () => listeners.delete(l);
    },
    () => selector(state),
  );
}

function upsertSession(s: State, session: Session): Partial<State> {
  const sessionById = { ...s.sessionById, [session.id]: session };
  if (session.parent_id) return { sessionById };
  const others = s.sessions.filter((x) => x.id !== session.id);
  const sessions = [session, ...others].sort((a, b) => b.updated_at - a.updated_at);
  return { sessions, sessionById };
}

let errorSeq = 0;

function draftFor(s: State, sid: string, mid: string): Draft {
  const d = s.drafts[sid];
  return d && d.id === mid ? { ...d, parts: [...d.parts] } : { id: mid, parts: [] };
}

export function applyEvent(ev: ServerEvent) {
  switch (ev.type) {
    case "session_updated":
      setState((s) => upsertSession(s, ev.session));
      break;
    case "session_deleted":
      setState((s) => {
        const { [ev.session_id]: _, ...rest } = s.sessionById;
        return { sessions: s.sessions.filter((x) => x.id !== ev.session_id), sessionById: rest };
      });
      break;
    case "session_status":
      setState((s) => ({
        busy: { ...s.busy, [ev.session_id]: ev.busy },
        drafts: ev.busy ? s.drafts : { ...s.drafts, [ev.session_id]: undefined },
        steps: ev.busy ? s.steps : { ...s.steps, [ev.session_id]: undefined },
      }));
      break;
    case "message_added":
      setState((s) => {
        const list = s.messages[ev.session_id];
        const drafts = s.drafts[ev.session_id]?.id === ev.message.id ? { ...s.drafts, [ev.session_id]: undefined } : s.drafts;
        // Viewing an older part of the chat: the new message comes with "back to the end".
        if (!list || s.windows[ev.session_id]?.newer) return { drafts };
        if (list.some((m) => m.id === ev.message.id)) return { drafts };
        return { messages: { ...s.messages, [ev.session_id]: [...list, ev.message] }, drafts };
      });
      break;
    case "text_delta":
    case "reasoning_delta":
      setState((s) => {
        const d = draftFor(s, ev.session_id, ev.message_id);
        const kind = ev.type === "text_delta" ? "text" : "reasoning";
        const last = d.parts[d.parts.length - 1];
        if (last && last.type === kind) {
          d.parts[d.parts.length - 1] = { ...last, text: last.text + ev.text } as Part;
        } else {
          d.parts.push({ type: kind, text: ev.text } as Part);
        }
        return { drafts: { ...s.drafts, [ev.session_id]: d } };
      });
      break;
    case "tool_started":
      setState((s) => {
        const d = draftFor(s, ev.session_id, ev.message_id);
        d.parts.push({ type: "tool_call", id: ev.call_id, name: ev.name, input: ev.input });
        return { drafts: { ...s.drafts, [ev.session_id]: d } };
      });
      break;
    case "tool_progress":
      setState((s) => ({ progress: { ...s.progress, [ev.call_id]: ev.output } }));
      break;
    case "permission_request":
      setState((s) => ({ permissions: [...s.permissions.filter((p) => p.id !== ev.id), ev] }));
      break;
    case "permission_resolved":
      setState((s) => ({ permissions: s.permissions.filter((p) => p.id !== ev.id) }));
      break;
    case "question_request":
      setState((s) => ({ questions: [...s.questions.filter((q) => q.id !== ev.id), ev] }));
      notify(ev.questions[0]?.question ?? "");
      break;
    case "question_resolved":
      setState((s) => ({ questions: s.questions.filter((q) => q.id !== ev.id) }));
      break;
    case "error":
      setState((s) => ({
        errors: {
          ...s.errors,
          [ev.session_id]: [...(s.errors[ev.session_id] ?? []), { id: ++errorSeq, message: ev.message }].slice(-5),
        },
      }));
      break;
    case "step_started":
      setState((s) => ({
        steps: { ...s.steps, [ev.session_id]: { messageId: ev.message_id, startedAt: ev.started_at, promptTokens: ev.prompt_tokens, estimated: ev.estimated } },
      }));
      break;
    case "queue_updated":
      setState((s) => ({ queue: { ...s.queue, [ev.session_id]: ev.items } }));
      break;
    case "waiting":
      setState((s) => ({ waiting: { ...s.waiting, [ev.session_id]: ev.reason ?? undefined } }));
      break;
    case "compaction":
      setState((s) => ({
        compacting: {
          ...s.compacting,
          [ev.session_id]: ev.state === "started" ? { auto: ev.auto, tokens: ev.tokens, limit: ev.limit, since: Date.now() } : undefined,
        },
      }));
      break;
    case "check_result": {
      const info: TurnCheckInfo = { message_id: ev.message_id, status: ev.status, exit_code: ev.exit_code, duration_ms: ev.duration_ms, output: ev.output };
      setState((s) => ({ checks: { ...s.checks, [ev.session_id]: { ...s.checks[ev.session_id], [ev.message_id]: info } } }));
      break;
    }
    case "mcp_status":
      setState(() => ({ mcp: ev.servers }));
      break;
    case "resync":
      void resync();
      break;
  }
}

/** Notify when the app is in the background and an agent needs input. */
function notify(body: string) {
  if (document.visibilityState === "visible" || !("Notification" in window)) return;
  // With Web Push on this device the server notifies (avoid duplicates).
  if (localStorage.getItem("codecraft.push") === "1") return;
  if (Notification.permission === "granted") {
    new Notification("CodeCraft", { body, icon: "/icon-192.png" });
  }
}

export function dismissError(sessionId: string, id: number) {
  setState((s) => ({ errors: { ...s.errors, [sessionId]: (s.errors[sessionId] ?? []).filter((e) => e.id !== id) } }));
}

export async function loadSessions() {
  const list = await get<Session[]>("/sessions");
  setState((s) => ({
    sessions: list,
    sessionById: { ...s.sessionById, ...Object.fromEntries(list.map((x) => [x.id, x])) },
    busy: { ...s.busy, ...Object.fromEntries(list.map((x) => [x.id, !!x.busy])) },
  }));
}

/** Messages loaded when a chat opens (from the last summary on, at most this many). */
export const PAGE = 200;

type Page = { messages: Message[]; has_older: boolean; has_newer: boolean };

export async function loadSession(id: string) {
  const data = await get<{ session: Session; messages: Message[]; busy: boolean; can_write: boolean; queue: QueueItem[]; step?: any; has_older?: boolean; has_newer?: boolean }>(
    `/sessions/${id}?limit=${PAGE}`,
  );
  setState((s) => ({
    windows: { ...s.windows, [id]: { older: !!data.has_older, newer: !!data.has_newer } },
    ...upsertSession(s, data.session),
    sessionById: { ...s.sessionById, [id]: data.session },
    messages: { ...s.messages, [id]: data.messages },
    busy: { ...s.busy, [id]: data.busy },
    readOnly: { ...s.readOnly, [id]: !data.can_write },
    queue: { ...s.queue, [id]: data.queue ?? [] },
    steps: {
      ...s.steps,
      [id]: data.step ? { messageId: data.step.message_id, startedAt: data.step.started_at, promptTokens: data.step.prompt_tokens, estimated: data.step.estimated } : s.steps[id],
    },
  }));
}

/** Check results of a chat (badges on its turns). */
export async function loadChecks(id: string) {
  const list = await get<TurnCheckInfo[]>(`/sessions/${id}/checks`).catch(() => [] as TurnCheckInfo[]);
  setState((s) => ({ checks: { ...s.checks, [id]: { ...Object.fromEntries(list.map((c) => [c.message_id, c])), ...s.checks[id] } } }));
}

/** Older messages before the loaded ones (scrolling up). Returns how many came. */
export async function loadOlder(id: string, stopAtSummary = true): Promise<number> {
  const first = getState().messages[id]?.[0];
  if (!first) return 0;
  const page = await get<Page>(`/sessions/${id}/messages?before=${encodeURIComponent(first.id)}&limit=${PAGE}${stopAtSummary ? "&stop_at_summary=true" : ""}`);
  setState((s) => {
    const cur = s.messages[id] ?? [];
    if (cur[0]?.id !== first.id) return {}; // changed meanwhile
    return {
      messages: { ...s.messages, [id]: [...page.messages, ...cur] },
      windows: { ...s.windows, [id]: { older: page.has_older, newer: s.windows[id]?.newer ?? false } },
    };
  });
  return page.messages.length;
}

/** Show the part of the chat around a message (search hit, media). */
export async function loadAround(id: string, messageId: string) {
  const page = await get<Page>(`/sessions/${id}/messages?around=${encodeURIComponent(messageId)}&limit=${PAGE}`);
  setState((s) => ({
    messages: { ...s.messages, [id]: page.messages },
    windows: { ...s.windows, [id]: { older: page.has_older, newer: page.has_newer } },
  }));
}

/** After reconnects: reload everything that could have changed. */
export async function resync() {
  const isAdmin = getState().me?.role === "admin";
  const [pending, mcp] = await Promise.all([get<any[]>("/pending"), isAdmin ? get("/mcp") : Promise.resolve({})]);
  setState(() => ({
    permissions: pending.filter((e) => e.type === "permission_request"),
    questions: pending.filter((e) => e.type === "question_request"),
    mcp,
  }));
  await loadSessions();
  const loaded = Object.keys(getState().messages);
  await Promise.all(loaded.map((id) => loadSession(id).catch(() => undefined)));
}