ai-agent/frontend/src/lib/spawn.tsx

import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from "react";
import type { RichTag } from "@gabvdl/ui";
import {
  fetchConversations,
  spawnConversation,
  type UploadedFile,
} from "@/api";

// Misclick guard: after submit we show the optimistic card but hold the real
// API call for this long, during which tapping the card fully un-sends it.
const QUEUE_DELAY_MS = 3_000;
// How long to wait for the new transcript to sync before giving up the redirect.
const SYNC_TIMEOUT_MS = 90_000;
const SYNC_POLL_MS = 1_500;

const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));

/**
 * Poll the conversation list until a transcript with this session id syncs
 * (~2s after a spawn/fork), and return its conversation id — or null on
 * timeout/abort. Shared by the spawn lifecycle below and the fork flow.
 */
export async function waitForConversation(
  sessionId: string,
  opts?: { timeoutMs?: number; aborted?: () => boolean },
): Promise<string | null> {
  const deadline = Date.now() + (opts?.timeoutMs ?? SYNC_TIMEOUT_MS);
  while (!opts?.aborted?.() && Date.now() < deadline) {
    try {
      const { conversations } = await fetchConversations();
      const hit = conversations.find((c) => c.sessionId === sessionId);
      if (hit) return hit.id;
    } catch {
      /* transient — keep polling */
    }
    await sleep(SYNC_POLL_MS);
  }
  return null;
}

type SpawnStatus = "idle" | "queued" | "spawning" | "waiting" | "ready" | "error";

interface EnqueueArgs {
  /** The raw text the user typed (restored to the composer on cancel). */
  base: string;
  /** The fully composed prompt (with tags/guidelines/attachments) to send. */
  prompt: string;
  /** Files echoed in the optimistic card + restored on cancel. */
  files: UploadedFile[];
  /** Tags active at submit (guidelines + locations), echoed in the sent card. */
  tags?: RichTag[];
  /** `--model` argument from the composer's model tag (undefined = default). */
  model?: string;
  /** Agent harness the model routes to: "claude" (default) or "pi". */
  harness?: "claude" | "pi";
  /** false ⇒ spawn with thinking disabled (the composer's thinking toggle). */
  thinking?: boolean;
  /** claude only: `--effort` level from the composer's effort select. */
  effort?: string;
}

interface SpawnContextValue {
  status: SpawnStatus;
  /** The optimistic echo of the just-submitted prompt (null when idle). */
  sentText: string | null;
  sentFiles: UploadedFile[];
  /** The guideline/location tags that went out with the send (chips in the card). */
  sentTags: RichTag[];
  error: string | null;
  /** Seconds left in the un-send window (0 once the API call has fired). */
  countdown: number;
  /**
   * The synced conversation id once the spawn has landed (status `"ready"`) —
   * the sent card offers "tap to open" instead of auto-navigating there.
   */
  readyId: string | null;
  /** Queue a spawn; the real API call fires after QUEUE_DELAY_MS. */
  enqueue: (args: EnqueueArgs) => void;
  /**
   * Un-send / detach the current spawn and return the raw text so the composer
   * can restore it. Only works during the queued window (before the API call).
   */
  cancel: () => { base: string; files: UploadedFile[] } | null;
  /**
   * Hide the sent card / error notification without cancelling anything: a
   * spawn already fired keeps running on the host, this just stops tracking it
   * (swipe-away, or collapsing the spawn box).
   */
  dismiss: () => void;
}

const SpawnContext = createContext<SpawnContextValue | null>(null);

/**
 * App-level owner of the "start a new session" lifecycle. Lives above the
 * router so the queued API call still fires — and the transcript sync keeps
 * being watched — even if the composer (SpawnBox) unmounts because the user
 * navigated away before the 3s un-send window elapsed.
 */
export function SpawnProvider({ children }: { children: React.ReactNode }) {
  const [status, setStatus] = useState<SpawnStatus>("idle");
  const [sentText, setSentText] = useState<string | null>(null);
  const [sentFiles, setSentFiles] = useState<UploadedFile[]>([]);
  const [sentTags, setSentTags] = useState<RichTag[]>([]);
  const [error, setError] = useState<string | null>(null);
  const [countdown, setCountdown] = useState(0);
  const [readyId, setReadyId] = useState<string | null>(null);

  // The queued (not-yet-sent) request + its timers, and an abort token the
  // in-flight sync loop checks so cancel() can detach it.
  const pending = useRef<EnqueueArgs | null>(null);
  const queueTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const tickTimer = useRef<ReturnType<typeof setInterval> | null>(null);
  const runId = useRef(0);
  // The run id whose API call has already been fired, so run() can never spawn a
  // second conversation for the same submission (idempotency guard).
  const firedFor = useRef(-1);

  const clearTimers = useCallback(() => {
    if (queueTimer.current) clearTimeout(queueTimer.current);
    if (tickTimer.current) clearInterval(tickTimer.current);
    queueTimer.current = null;
    tickTimer.current = null;
  }, []);

  const reset = useCallback(() => {
    clearTimers();
    pending.current = null;
    setStatus("idle");
    setSentText(null);
    setSentFiles([]);
    setSentTags([]);
    setCountdown(0);
    setReadyId(null);
  }, [clearTimers]);

  // Fire the actual API call + poll for the transcript. Deliberately NO
  // auto-redirect once it lands: the sent card flips to "ready — tap to open"
  // and the user decides when (or whether) to jump in.
  const run = useCallback(
    async (args: EnqueueArgs, id: number) => {
      // Fire exactly once per submission — guards against the timer callback (or
      // a stray re-entry) ever kicking off a second `claude -p` for one prompt.
      if (firedFor.current === id) return;
      firedFor.current = id;
      const aborted = () => runId.current !== id;
      setStatus("spawning");
      setCountdown(0);
      try {
        const { sessionId } = await spawnConversation(
          args.prompt,
          args.model,
          args.harness,
          args.thinking,
          args.effort,
        );
        if (aborted()) return;
        setStatus("waiting");
        const hit = await waitForConversation(sessionId, { aborted });
        if (aborted()) return;
        if (hit) {
          setReadyId(hit);
          setStatus("ready");
          return;
        }
        // The session started but its transcript never synced in time — say so
        // rather than silently dropping the card.
        setError(
          "Session started, but its transcript hasn't synced yet — it will show up in the conversation list shortly.",
        );
        setStatus("error");
      } catch (e) {
        if (aborted()) return;
        clearTimers();
        pending.current = args;
        setError(e instanceof Error ? e.message : String(e));
        setStatus("error");
      }
    },
    [clearTimers],
  );

  const enqueue = useCallback(
    (args: EnqueueArgs) => {
      clearTimers();
      const id = ++runId.current;
      pending.current = args;
      setError(null);
      setSentText(args.base);
      setSentFiles(args.files);
      setSentTags(args.tags ?? []);
      setStatus("queued");
      setCountdown(Math.ceil(QUEUE_DELAY_MS / 1000));
      tickTimer.current = setInterval(
        () => setCountdown((c) => Math.max(0, c - 1)),
        1000,
      );
      queueTimer.current = setTimeout(() => {
        clearTimers();
        void run(args, id);
      }, QUEUE_DELAY_MS);
    },
    [clearTimers, run],
  );

  const cancel = useCallback(() => {
    // Only the pre-API queued window can be un-sent. `queueTimer` is set exactly
    // while we're still counting down and the API call has NOT yet fired; once
    // it fires the timer is cleared. Refusing to "cancel" after that is what
    // guarantees a single conversation — otherwise a tap racing the timer would
    // silently leave a spawned session behind while restoring the composer, and
    // the user would re-submit and get a duplicate.
    if (!queueTimer.current) return null;
    const args = pending.current;
    setError(null);
    reset();
    return args ? { base: args.base, files: args.files } : null;
  }, [reset]);

  const dismiss = useCallback(() => {
    // Detach any in-flight sync loop (its `aborted()` check trips on the bump),
    // then clear the card/error. Anything already sent keeps running unwatched.
    runId.current += 1;
    setError(null);
    reset();
  }, [reset]);

  useEffect(() => () => clearTimers(), [clearTimers]);

  const value = useMemo<SpawnContextValue>(
    () => ({ status, sentText, sentFiles, sentTags, error, countdown, readyId, enqueue, cancel, dismiss }),
    [status, sentText, sentFiles, sentTags, error, countdown, readyId, enqueue, cancel, dismiss],
  );

  return <SpawnContext.Provider value={value}>{children}</SpawnContext.Provider>;
}

export function useSpawn(): SpawnContextValue {
  const ctx = useContext(SpawnContext);
  if (!ctx) throw new Error("useSpawn must be used within a SpawnProvider");
  return ctx;
}