canopy-ui 0.3.0 → 0.6.0

package/src/chat/sessionReducer.ts

@@ -0,0 +1,368 @@
+import type { Draft, Message, SessionState, WsEvent } from "./protocol";
+
+// Pure reducer for SessionState — extracted from useSessionSocket so it
+// can be unit-tested without WebSocket plumbing. Side-effect events
+// (session.title_updated → optional injected callback; session.error →
+// setLastError + clear draft debounce) stay in the hook itself.
+//
+// Keep this file dependency-free (no React) so a vitest run doesn't pull
+// jsdom or RTL.
+//
+// canopy adaptation vs ace: message/draft ids are STRINGS, and
+// `chat.stream_start` UPSERTS the assistant message — canopy's
+// `draft.committed` carries only `user_message_id` (no assistant id to
+// pre-insert), so the assistant row is created lazily when its first stream
+// frame arrives.
+export function sessionReducer(prev: SessionState, frame: WsEvent): SessionState {
+  switch (frame.event) {
+    case "session.state":
+      return frame.data;
+
+    case "chat.stream_start": {
+      // Upsert: if the assistant message already exists (rare — a runner that
+      // pre-inserts it), flip it to streaming; otherwise create it. canopy's
+      // draft.committed cannot pre-send the assistant id, so this is the
+      // normal path for making the streamed reply visible.
+      const exists = prev.messages.some((m) => m.id === frame.data.message_id);
+      if (exists) {
+        return {
+          ...prev,
+          messages: prev.messages.map((m) =>
+            m.id === frame.data.message_id
+              ? { ...m, status: "streaming" as const }
+              : m,
+          ),
+        };
+      }
+      const nowIso = new Date().toISOString();
+      const assistant: Message = {
+        id: frame.data.message_id,
+        turn_index: frame.data.turn_index,
+        role: "assistant",
+        content: {},
+        plaintext: "",
+        status: "streaming",
+        error_detail: null,
+        started_at: nowIso,
+        completed_at: null,
+        created_at: nowIso,
+      };
+      return { ...prev, messages: [...prev.messages, assistant] };
+    }
+
+    case "session.activity":
+      return { ...prev, activity: frame.data.state };
+
+    case "chat.user_message": {
+      // Someone typed into emdash, OR into this page. Both reach here, and that
+      // is why matching on turn_index alone is not enough: a web send writes its
+      // row at a DENSE index (services._next_index), then the agent reads the
+      // message and the transcript re-ships the same text at a COMPOSITE ordinal
+      // (record * BLOCK_STRIDE + block). Same words, two different indices, two
+      // different ids — so the upsert missed and the message rendered twice
+      // live, while a reload showed it once (get_or_create dedupes server-side).
+      //
+      // Falling back to matching identical text on a recent user row closes it.
+      // Deliberately narrow: same role, same text, and only against the tail, so
+      // a genuine repeat of a short message ("yes") sent much later still lands
+      // as its own row.
+      const RECENT_USER_ROWS = 6;
+      const sameText = (m: Message) =>
+        m.role === "user" &&
+        m.plaintext.trim() !== "" &&
+        m.plaintext.trim() === frame.data.plaintext.trim();
+      const recentUsers = prev.messages.filter((m) => m.role === "user").slice(-RECENT_USER_ROWS);
+      const existing =
+        prev.messages.find(
+          (m) => m.id === frame.data.message_id ||
+            (m.role === "user" && m.turn_index === frame.data.turn_index),
+        ) ?? recentUsers.find(sameText);
+      if (existing) {
+        return {
+          ...prev,
+          messages: prev.messages.map((m) =>
+            m === existing
+              ? {
+                  ...m,
+                  id: frame.data.message_id,
+                  // Take the incoming ordinal: the transcript's composite index
+                  // is the durable one, so the row sorts where a reload puts it.
+                  turn_index: frame.data.turn_index,
+                  plaintext: frame.data.plaintext,
+                }
+              : m,
+          ),
+        };
+      }
+      const nowIso = new Date().toISOString();
+      const user: Message = {
+        id: frame.data.message_id,
+        turn_index: frame.data.turn_index,
+        role: "user",
+        content: { text: frame.data.plaintext },
+        plaintext: frame.data.plaintext,
+        status: "complete",
+        error_detail: null,
+        started_at: null,
+        completed_at: null,
+        created_at: nowIso,
+      };
+      return { ...prev, messages: [...prev.messages, user] };
+    }
+
+    case "chat.delta":
+      return {
+        ...prev,
+        messages: prev.messages.map((m) =>
+          m.id === frame.data.message_id
+            ? { ...m, plaintext: m.plaintext + frame.data.text }
+            : m,
+        ),
+      };
+
+    case "chat.stream_complete":
+      return {
+        ...prev,
+        messages: prev.messages.map((m) =>
+          m.id === frame.data.message_id
+            ? {
+                ...m,
+                plaintext: frame.data.plaintext,
+                status: "complete" as const,
+              }
+            : m,
+        ),
+      };
+
+    case "chat.stream_error":
+      // NOTE: backend emits chat.stream_error with detail="cancelled"
+      // for stop-driven cancellation; there's no separate
+      // chat.stream_cancelled event in practice. Distinguished by detail.
+      return {
+        ...prev,
+        messages: prev.messages.map((m) =>
+          m.id === frame.data.message_id
+            ? {
+                ...m,
+                status: "error" as const,
+                error_detail: frame.data.detail,
+              }
+            : m,
+        ),
+      };
+
+    case "chat.stream_cancelled":
+      return {
+        ...prev,
+        messages: prev.messages.map((m) =>
+          m.id === frame.data.message_id
+            ? {
+                ...m,
+                status: "error" as const,
+                error_detail: `cancelled (partial: ${frame.data.partial_len} chars)`,
+              }
+            : m,
+        ),
+      };
+
+    case "chat.tool_use":
+    case "chat.tool_result": {
+      // Tool rows are real Message rows on the server; these frames are the
+      // same row arriving live. Upsert by id so a re-delivery (reconnect
+      // catch-up, a retried post) updates in place instead of doubling the
+      // row — and so a tool_result that lands twice can't orphan its pair.
+      const role = frame.event === "chat.tool_use" ? "tool_use" : "tool_result";
+      const block = frame.data.block ?? {};
+      const plaintext = typeof block.text === "string" ? block.text : "";
+      const id = frame.data.tool_message_id;
+      // Reconcile on the CORRELATION key, not the message id. The same tool call
+      // arrives twice by design: once live from a hook (no ordinal, id `seq:-1`)
+      // and once durably from the transcript (ordinal-keyed, a real id). They
+      // share only `tool_use_id` / `id` in content, so matching on that is what
+      // makes the live row a placeholder the durable row REPLACES rather than a
+      // duplicate that sits beside it forever.
+      const correlation =
+        role === "tool_use"
+          ? (block.id as string | undefined)
+          : (block.tool_use_id as string | undefined);
+      const existing = prev.messages.find(
+        (m) =>
+          m.id === id ||
+          (m.role === role &&
+            correlation !== undefined &&
+            correlation !== "" &&
+            (m.content as Record<string, unknown>)?.[
+              role === "tool_use" ? "id" : "tool_use_id"
+            ] === correlation),
+      );
+      if (existing) {
+        // Adopt the incoming id and ordinal: a durable row superseding a live
+        // placeholder must take over its identity, or the next update keys on a
+        // `seq:-1` that no longer means anything.
+        return {
+          ...prev,
+          messages: prev.messages.map((m) =>
+            m === existing
+              ? {
+                  ...m,
+                  id,
+                  turn_index: frame.data.turn_index ?? m.turn_index,
+                  content: block,
+                  plaintext,
+                }
+              : m,
+          ),
+        };
+      }
+      const nowIso = new Date().toISOString();
+      const message: Message = {
+        id,
+        // Fall back to appending after the newest row when the server didn't
+        // send an ordinal — order still holds, since frames arrive in order.
+        turn_index:
+          frame.data.turn_index ??
+          (prev.messages[prev.messages.length - 1]?.turn_index ?? 0) + 1,
+        role,
+        content: block,
+        plaintext,
+        status: block.is_error === true ? "error" : "complete",
+        error_detail: null,
+        started_at: nowIso,
+        completed_at: nowIso,
+        created_at: nowIso,
+      };
+      return { ...prev, messages: [...prev.messages, message] };
+    }
+
+    case "draft.updated": {
+      const incoming = frame.data as Draft;
+      // If we're the current editor, keep our local body — the server
+      // echo is stale relative to keystrokes that happened since the
+      // debounced send. Only accept metadata (version, last_editor, etc).
+      if (
+        prev.active_draft &&
+        incoming.last_editor === prev.current_user_id
+      ) {
+        return {
+          ...prev,
+          active_draft: {
+            ...prev.active_draft,
+            version: incoming.version,
+            last_editor: incoming.last_editor,
+            last_edit_at: incoming.last_edit_at,
+          },
+        };
+      }
+      return { ...prev, active_draft: incoming };
+    }
+
+    case "draft.lock_changed":
+      if (prev.active_draft && prev.active_draft.id === frame.data.draft_id) {
+        return {
+          ...prev,
+          active_draft: {
+            ...prev.active_draft,
+            last_editor: frame.data.holder_user_id ?? prev.active_draft.last_editor,
+          },
+        };
+      }
+      return prev;
+
+    case "draft.committed": {
+      // Insert the optimistic USER message from the draft body that's about
+      // to be cleared. The assistant reply is NOT inserted here — canopy's
+      // draft.committed carries no assistant id; `chat.stream_start` upserts
+      // that row when the reply begins.
+      //
+      // Also clear active_draft.body here. The server creates a new empty
+      // draft with last_editor=sender, so the follow-up draft.updated hits
+      // the "keep local body" branch above and would otherwise leave the
+      // just-sent text in the textarea — which lets Enter re-send the same
+      // turn repeatedly.
+      const prevDraftBody = prev.active_draft?.body ?? "";
+      const maxTurnIndex = prev.messages.reduce(
+        (acc, msg) => Math.max(acc, msg.turn_index),
+        0,
+      );
+      const nowIso = new Date().toISOString();
+      const userMessage: Message = {
+        id: frame.data.user_message_id,
+        turn_index: maxTurnIndex + 1,
+        role: "user",
+        content: { text: prevDraftBody },
+        plaintext: prevDraftBody,
+        status: "complete",
+        error_detail: null,
+        started_at: null,
+        completed_at: nowIso,
+        created_at: nowIso,
+      };
+      return {
+        ...prev,
+        active_draft: prev.active_draft
+          ? { ...prev.active_draft, body: "" }
+          : prev.active_draft,
+        messages: [...prev.messages, userMessage],
+      };
+    }
+
+    case "draft.discarded":
+      if (prev.active_draft && prev.active_draft.id === frame.data.draft_id) {
+        return {
+          ...prev,
+          active_draft: { ...prev.active_draft, body: "" },
+        };
+      }
+      return prev;
+
+    case "presence.joined": {
+      const ids = new Set(prev.presence_user_ids);
+      ids.add(frame.data.user_id);
+      return { ...prev, presence_user_ids: [...ids] };
+    }
+
+    case "presence.left":
+      return {
+        ...prev,
+        presence_user_ids: prev.presence_user_ids.filter(
+          (id) => id !== frame.data.user_id,
+        ),
+      };
+
+    case "session.error": {
+      // Side effects (setLastError, clear draft debounce) are handled
+      // by the hook; the reducer only knows about the version-mismatch
+      // recovery, which mutates active_draft.
+      if (
+        frame.data.code === "draft_version_mismatch" &&
+        frame.data.detail &&
+        typeof frame.data.detail === "object"
+      ) {
+        const detail = frame.data.detail as {
+          current_version: number;
+          current_body: string;
+        };
+        return prev.active_draft
+          ? {
+              ...prev,
+              active_draft: {
+                ...prev.active_draft,
+                version: detail.current_version,
+                body: detail.current_body,
+              },
+            }
+          : prev;
+      }
+      return prev;
+    }
+
+    case "session.title_updated":
+      // Pure reducer leaves this alone — the hook calls its optional
+      // onTitleUpdated callback on receipt and short-circuits. Included
+      // here so an exhaustive switch type-checks.
+      return prev;
+
+    default:
+      return prev;
+  }
+}

package/src/chat/useSessionSocket.ts

@@ -0,0 +1,332 @@
+import { useCallback, useEffect, useRef, useState } from "react";
+
+import type { Message, SessionState, WsEvent } from "./protocol";
+import { shouldSyncDraftLive } from "./drafts";
+import { prependHistory } from "./history";
+import { sessionReducer } from "./sessionReducer";
+
+const HEARTBEAT_INTERVAL_MS = 20_000;
+const RECONNECT_DELAYS_MS = [1_000, 2_000, 5_000, 10_000];
+const DRAFT_UPDATE_DEBOUNCE_MS = 150;
+
+const INITIAL_STATE: SessionState = {
+  messages: [],
+  active_draft: null,
+  participants: [],
+  presence_user_ids: [],
+  current_user_id: 0,
+};
+
+export interface UseSessionSocketOptions {
+  /** The chat session id (UUID string). */
+  sessionId: string;
+  /**
+   * App-injected WebSocket URL builder. The kit never imports app routing/base
+   * helpers; the container passes one (e.g. canopy's `wsUrl`). Called with the
+   * relative path `ws/canopy-sessions/${sessionId}/`.
+   */
+  wsUrl: (path: string) => string;
+  /**
+   * Optional side-effect callback fired when the server broadcasts a
+   * `session.title_updated` (replaces ace's `notifySessionsUpdated`). The kit
+   * has no opinion on what to do with it.
+   */
+  onTitleUpdated?: () => void;
+}
+
+export interface UseSessionSocketResult {
+  state: SessionState;
+  connected: boolean;
+  /** A send is outstanding with no reply yet — the turn is QUEUED, waiting for
+   *  a runner. Nothing in `state` can express this (there is no assistant
+   *  message until the first token), and it is what keeps Stop reachable while
+   *  a turn is stuck. */
+  awaitingReply: boolean;
+  sendChat: () => void;
+  stopChat: (messageId: string | null) => void;
+  updateDraft: (body: string) => void;
+  takeOverDraft: () => void;
+  discardDraft: () => void;
+  prependMessages: (older: Message[]) => void;
+  lastError: string | null;
+}
+
+export function useSessionSocket({
+  sessionId,
+  wsUrl,
+  onTitleUpdated,
+}: UseSessionSocketOptions): UseSessionSocketResult {
+  const [state, setState] = useState<SessionState>(INITIAL_STATE);
+  const [connected, setConnected] = useState(false);
+  const [lastError, setLastError] = useState<string | null>(null);
+  // A send has gone out but no reply has begun — i.e. the turn is QUEUED,
+  // waiting for a runner to claim it. There is no assistant message during
+  // this window, so nothing else in the state can express it, and without it
+  // the Stop control is unreachable exactly when the turn is stuck.
+  const [awaitingReply, setAwaitingReply] = useState(false);
+
+  const socketRef = useRef<WebSocket | null>(null);
+  const stateRef = useRef<SessionState>(INITIAL_STATE);
+  const reconnectAttemptRef = useRef(0);
+  const heartbeatTimerRef = useRef<number | null>(null);
+  const draftDebounceRef = useRef<number | null>(null);
+  const pendingDraftBodyRef = useRef<string | null>(null);
+  const closedByUserRef = useRef(false);
+  const onTitleUpdatedRef = useRef(onTitleUpdated);
+  // Control frames that must not be lost across a reconnect (currently
+  // only chat.stop). The WS-world analogue of an abortable chat transport.
+  const pendingFramesRef = useRef<{ action: string; data: unknown }[]>([]);
+
+  useEffect(() => {
+    stateRef.current = state;
+  }, [state]);
+
+  useEffect(() => {
+    onTitleUpdatedRef.current = onTitleUpdated;
+  }, [onTitleUpdated]);
+
+  const send = useCallback((frame: { action: string; data: unknown }) => {
+    const ws = socketRef.current;
+    if (ws && ws.readyState === WebSocket.OPEN) {
+      ws.send(JSON.stringify(frame));
+      return;
+    }
+    // Queue chat.stop so a stop clicked while the socket is reconnecting
+    // is delivered on next OPEN instead of silently dropped. Draft updates
+    // are intentionally NOT queued — they have a version guard and the
+    // user's next keystroke will refresh the body anyway.
+    if (frame.action === "chat.stop") {
+      pendingFramesRef.current.push(frame);
+    }
+  }, []);
+
+  const applyEvent = useCallback((frame: WsEvent) => {
+    // Any of these means the queued window is over: the reply began, ended,
+    // was cancelled, or the send failed outright.
+    if (
+      frame.event === "chat.stream_start" ||
+      frame.event === "chat.stream_complete" ||
+      frame.event === "chat.stream_error" ||
+      frame.event === "chat.stream_cancelled" ||
+      frame.event === "session.error"
+    ) {
+      setAwaitingReply(false);
+    }
+    // Side-effect events: handle BEFORE setState so React strict-mode's
+    // double-invocation of the updater doesn't double-fire the effect.
+    if (frame.event === "session.title_updated") {
+      onTitleUpdatedRef.current?.();
+      return;
+    }
+    if (frame.event === "session.error") {
+      setLastError(frame.data.message);
+      if (
+        frame.data.code === "draft_version_mismatch" &&
+        frame.data.detail &&
+        typeof frame.data.detail === "object"
+      ) {
+        // Clear any pending optimistic body so the user's stale local
+        // text doesn't auto-re-send with the new version.
+        pendingDraftBodyRef.current = null;
+        if (draftDebounceRef.current != null) {
+          window.clearTimeout(draftDebounceRef.current);
+          draftDebounceRef.current = null;
+        }
+      }
+    }
+    setState((prev) => sessionReducer(prev, frame));
+  }, []);
+
+  const connect = useCallback(() => {
+    if (closedByUserRef.current) return;
+    const ws = new WebSocket(wsUrl(`ws/canopy-sessions/${sessionId}/`));
+    socketRef.current = ws;
+
+    ws.onopen = () => {
+      setConnected(true);
+      reconnectAttemptRef.current = 0;
+      // Flush any control frames that were queued while the socket was
+      // closed. See `send` above.
+      const queued = pendingFramesRef.current;
+      pendingFramesRef.current = [];
+      for (const frame of queued) {
+        ws.send(JSON.stringify(frame));
+      }
+      if (heartbeatTimerRef.current != null) {
+        window.clearInterval(heartbeatTimerRef.current);
+      }
+      heartbeatTimerRef.current = window.setInterval(() => {
+        send({ action: "presence.heartbeat", data: {} });
+      }, HEARTBEAT_INTERVAL_MS);
+    };
+
+    ws.onmessage = (e) => {
+      try {
+        const frame = JSON.parse(e.data) as WsEvent;
+        applyEvent(frame);
+      } catch {
+        // ignore malformed frames
+      }
+    };
+
+    ws.onclose = () => {
+      setConnected(false);
+      if (heartbeatTimerRef.current != null) {
+        window.clearInterval(heartbeatTimerRef.current);
+        heartbeatTimerRef.current = null;
+      }
+      if (closedByUserRef.current) return;
+      const attempt = reconnectAttemptRef.current;
+      const delay =
+        RECONNECT_DELAYS_MS[Math.min(attempt, RECONNECT_DELAYS_MS.length - 1)];
+      reconnectAttemptRef.current = attempt + 1;
+      window.setTimeout(connect, delay);
+    };
+
+    ws.onerror = () => {
+      // onclose will fire next; nothing to do here.
+    };
+  }, [applyEvent, send, sessionId, wsUrl]);
+
+  useEffect(() => {
+    closedByUserRef.current = false;
+    reconnectAttemptRef.current = 0;
+    connect();
+    return () => {
+      closedByUserRef.current = true;
+      if (heartbeatTimerRef.current != null) {
+        window.clearInterval(heartbeatTimerRef.current);
+      }
+      if (socketRef.current) {
+        socketRef.current.close();
+        socketRef.current = null;
+      }
+    };
+  }, [connect]);
+
+  const sendChat = useCallback(() => {
+    // Flush the local body BEFORE committing. `chat.send` commits the SERVER's
+    // draft, so this is the moment the body has to exist there — and when
+    // live sync is off (single-player) it is the ONLY time it is sent.
+    //
+    // Unconditional on purpose: keying this off a pending debounce timer meant
+    // nothing was flushed when there was no timer, which is now the normal case.
+    if (draftDebounceRef.current != null) {
+      window.clearTimeout(draftDebounceRef.current);
+      draftDebounceRef.current = null;
+    }
+    if (pendingDraftBodyRef.current != null && stateRef.current.active_draft) {
+      send({
+        action: "draft.update",
+        data: {
+          version: stateRef.current.active_draft.version,
+          body: pendingDraftBodyRef.current,
+        },
+      });
+    }
+    pendingDraftBodyRef.current = null;
+    setAwaitingReply(true);
+    send({ action: "chat.send", data: {} });
+  }, [send]);
+
+  const stopChat = useCallback(
+    (messageId: string | null) => {
+      // messageId is null when the turn is still queued. The server's
+      // chat.stop cancels every non-terminal turn on the session and only
+      // echoes the id back, so a null one cancels just as effectively.
+      setAwaitingReply(false);
+      send({ action: "chat.stop", data: { message_id: messageId } });
+    },
+    [send],
+  );
+
+  const updateDraft = useCallback(
+    (body: string) => {
+      // Optimistic local update so the textarea feels snappy.
+      setState((prev) =>
+        prev.active_draft
+          ? { ...prev, active_draft: { ...prev.active_draft, body } }
+          : prev,
+      );
+      pendingDraftBodyRef.current = body;
+      // Alone in the session? Don't mirror keystrokes at all. The body is
+      // flushed once by sendChat, which is the only moment the server actually
+      // needs it. This is what makes single-player typing purely local — no
+      // round trip, no echo, no version to disagree about.
+      if (!shouldSyncDraftLive(stateRef.current.presence_user_ids)) {
+        if (draftDebounceRef.current != null) {
+          window.clearTimeout(draftDebounceRef.current);
+          draftDebounceRef.current = null;
+        }
+        return;
+      }
+      if (draftDebounceRef.current != null) {
+        window.clearTimeout(draftDebounceRef.current);
+      }
+      draftDebounceRef.current = window.setTimeout(() => {
+        draftDebounceRef.current = null;
+        const current = stateRef.current.active_draft;
+        const pending = pendingDraftBodyRef.current;
+        if (current != null && pending != null) {
+          // Only consumed once it has actually gone out. Clearing it
+          // unconditionally dropped anything typed before session.state
+          // arrived (no draft yet ⇒ nothing sent, body forgotten).
+          pendingDraftBodyRef.current = null;
+          send({
+            action: "draft.update",
+            data: { version: current.version, body: pending },
+          });
+        }
+      }, DRAFT_UPDATE_DEBOUNCE_MS);
+    },
+    [send],
+  );
+
+  const takeOverDraft = useCallback(() => {
+    send({ action: "draft.take_over", data: {} });
+  }, [send]);
+
+  const discardDraft = useCallback(() => {
+    send({ action: "draft.discard", data: {} });
+  }, [send]);
+
+  // Someone joining mid-compose must see what is ALREADY typed. Nothing was
+  // mirrored while we were alone, so without this one catch-up flush their view
+  // would sit empty until the next keystroke. Closes the only gap that skipping
+  // live sync opens up.
+  const liveSync = shouldSyncDraftLive(state.presence_user_ids);
+  useEffect(() => {
+    if (!liveSync) return;
+    const pending = pendingDraftBodyRef.current;
+    const current = stateRef.current.active_draft;
+    if (pending == null || current == null) return;
+    pendingDraftBodyRef.current = null;
+    send({
+      action: "draft.update",
+      data: { version: current.version, body: pending },
+    });
+  }, [liveSync, send]);
+
+  const prependMessages = useCallback((older: Message[]) => {
+    // Apply a REST "Load earlier" page into the live socket state. A later
+    // session.state snapshot (e.g. reconnect) resets to the tail — acceptable;
+    // the user re-loads earlier if needed.
+    setState((prev) => {
+      const merged = prependHistory(prev.messages, older);
+      return merged === prev.messages ? prev : { ...prev, messages: merged };
+    });
+  }, []);
+
+  return {
+    state,
+    connected,
+    awaitingReply,
+    sendChat,
+    stopChat,
+    updateDraft,
+    takeOverDraft,
+    discardDraft,
+    prependMessages,
+    lastError,
+  };
+}