learning_ai_common_plat/packages/ai-ui/src/useToolCalls.ts

import { useCallback, useMemo, useState } from 'react';
import type { ToolCallStatus, ToolInvocation } from './types.js';

export interface UseToolCallsHelpers {
  /** Map of id → ToolInvocation, stable across renders. */
  invocations: Record<string, ToolInvocation>;
  /** Ordered list of invocations (insertion order). */
  ordered: ToolInvocation[];
  /** Begin a new tool call. Idempotent — repeated calls update in place. */
  begin: (
    init: Pick<ToolInvocation, 'id' | 'name'> & Partial<Omit<ToolInvocation, 'id' | 'name'>>,
  ) => void;
  /** Merge a partial update into an existing invocation. */
  update: (id: string, patch: Partial<ToolInvocation>) => void;
  /** Transition an invocation to its terminal state (success / error). */
  settle: (
    id: string,
    result:
      | { status: 'success'; output?: unknown }
      | { status: 'error'; error: string },
  ) => void;
  /** Remove a specific invocation (e.g. on reload). */
  clear: (id: string) => void;
  /** Remove all invocations. */
  clearAll: () => void;
}

/**
 * `useToolCalls` — per-turn tool-invocation state machine.
 *
 * Stores tool calls in two parallel shapes — a Record for O(1) lookup
 * by id and an ordered array for rendering. Insertion order is preserved
 * across `update` calls so vertical scroll position stays stable while
 * a stream is in flight.
 *
 * Typical wiring (Wave 2 0.2 — usage with raw streams):
 *
 * ```ts
 * const tools = useToolCalls();
 * for await (const event of parseToolStream(response.body)) {
 *   if (event.type === 'tool-call')   tools.begin(event);
 *   if (event.type === 'tool-delta')  tools.update(event.id, event);
 *   if (event.type === 'tool-result') tools.settle(event.id, event.result);
 * }
 * ```
 *
 * 0.4 (`<ToolPalette>` with MCP discovery) will layer auto-wiring on top
 * of this hook.
 */
export function useToolCalls(initial: ToolInvocation[] = []): UseToolCallsHelpers {
  const [invocations, setInvocations] = useState<Record<string, ToolInvocation>>(
    () => Object.fromEntries(initial.map(t => [t.id, t])),
  );
  const [order, setOrder] = useState<string[]>(() => initial.map(t => t.id));

  const begin = useCallback<UseToolCallsHelpers['begin']>(init => {
    const id = init.id;
    const status: ToolCallStatus = init.status ?? 'pending';
    const startedAt = performance.now();
    setInvocations(prev => {
      const existing = prev[id];
      const next: ToolInvocation = {
        ...existing,
        status,
        ...init,
        // Carry the start timestamp into a private field for later
        // settle() to compute durationMs. We tuck it under a symbol-
        // free convention so consumers can still read everything.
      };
      // Stash startedAt on a hidden field. We can't use Symbol because
      // we want this hook to survive JSON serialization in tests.
      (next as { __startedAt?: number }).__startedAt = startedAt;
      return { ...prev, [id]: next };
    });
    setOrder(prev => (prev.includes(id) ? prev : [...prev, id]));
  }, []);

  const update = useCallback<UseToolCallsHelpers['update']>((id, patch) => {
    setInvocations(prev => {
      const existing = prev[id];
      if (!existing) return prev;
      return { ...prev, [id]: { ...existing, ...patch } };
    });
  }, []);

  const settle = useCallback<UseToolCallsHelpers['settle']>((id, result) => {
    setInvocations(prev => {
      const existing = prev[id];
      if (!existing) return prev;
      const startedAt = (existing as { __startedAt?: number }).__startedAt;
      const durationMs =
        typeof startedAt === 'number'
          ? Math.max(0, Math.round(performance.now() - startedAt))
          : existing.durationMs;
      const next: ToolInvocation =
        result.status === 'success'
          ? {
              ...existing,
              status: 'success',
              output: result.output ?? existing.output,
              error: undefined,
              durationMs,
            }
          : {
              ...existing,
              status: 'error',
              error: result.error,
              durationMs,
            };
      return { ...prev, [id]: next };
    });
  }, []);

  const clear = useCallback<UseToolCallsHelpers['clear']>(id => {
    setInvocations(prev => {
      if (!(id in prev)) return prev;
      const { [id]: _drop, ...rest } = prev;
      return rest;
    });
    setOrder(prev => prev.filter(x => x !== id));
  }, []);

  const clearAll = useCallback(() => {
    setInvocations({});
    setOrder([]);
  }, []);

  const ordered = useMemo(
    () =>
      order
        .map(id => invocations[id])
        .filter((x): x is ToolInvocation => x !== undefined),
    [order, invocations],
  );

  return { invocations, ordered, begin, update, settle, clear, clearAll };
}