Showing Code Mode in the UI

import { useState } from "react";
import { parsePartialJSON } from "@tanstack/ai";
import { useChat, fetchServerSentEvents } from "@tanstack/ai-react";

interface CodeModeEvent {
  eventType: string;
  data: unknown;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null;
}

export function CodeModeChat() {
  const [input, setInput] = useState("");
  const [events, setEvents] = useState<Record<string, Array<CodeModeEvent>>>(
    {},
  );

  const { messages, sendMessage, isLoading } = useChat({
    connection: fetchServerSentEvents("/api/chat"),
    threadId: "support-chat",
    // Keep each event under the execute_typescript call that sent it.
    onCustomEvent: (eventType, data, { toolCallId }) => {
      if (!toolCallId) return;
      setEvents((prev) => ({
        ...prev,
        [toolCallId]: [...(prev[toolCallId] ?? []), { eventType, data }],
      }));
    },
  });

  return (
    <>
      {messages.map((message) => (
        <div key={message.id}>
          {message.parts.map((part, index) => {
            if (part.type === "text") {
              return <p key={index}>{part.content}</p>;
            }
            if (
              part.type === "tool-call" &&
              part.name === "execute_typescript"
            ) {
              return (
                <CodeExecutionPanel
                  key={part.id}
                  args={part.arguments}
                  output={part.output}
                  events={events[part.id] ?? []}
                />
              );
            }
            return null;
          })}
        </div>
      ))}
      <form
        onSubmit={(event) => {
          event.preventDefault();
          if (input.trim() === "") return;
          sendMessage(input);
          setInput("");
        }}
      >
        <input value={input} onChange={(event) => setInput(event.target.value)} />
        <button type="submit" disabled={isLoading}>
          Send
        </button>
      </form>
    </>
  );
}

function CodeExecutionPanel({
  args,
  output,
  events,
}: {
  args: string;
  output: unknown;
  events: Array<CodeModeEvent>;
}) {
  // The arguments stream in as partial JSON, so the code shows while the model writes it.
  const parsed: unknown = parsePartialJSON(args);
  const code =
    isRecord(parsed) && typeof parsed.typescriptCode === "string"
      ? parsed.typescriptCode
      : "";

  return (
    <div>
      <details open>
        <summary>TypeScript code</summary>
        <pre>
          <code>{code}</code>
        </pre>
      </details>
      <ul>
        {events.map((event, index) => (
          <li key={index}>{describeEvent(event)}</li>
        ))}
      </ul>
      <pre>{describeResult(output)}</pre>
    </div>
  );
}

function describeEvent({ eventType, data }: CodeModeEvent): string {
  if (!isRecord(data)) return eventType;
  switch (eventType) {
    case "code_mode:execution_started":
      return "Started";
    case "code_mode:console":
      return `[${String(data.level)}] ${String(data.message)}`;
    case "code_mode:external_call":
      return `Call ${String(data.function)}(${JSON.stringify(data.args)})`;
    case "code_mode:external_result":
      return `Done ${String(data.function)} in ${String(data.duration)} ms`;
    case "code_mode:external_error":
      return `Failed ${String(data.function)}: ${String(data.error)}`;
    case "code_mode:execution_finished":
      return `Finished in ${String(data.durationMs)} ms`;
    default:
      return eventType;
  }
}

function describeResult(output: unknown): string {
  if (output === undefined) return "Running...";
  if (!isRecord(output)) return JSON.stringify(output, null, 2);
  if (output.success === false) {
    const message = isRecord(output.error) ? output.error.message : undefined;
    return `Error: ${String(message ?? "unknown error")}`;
  }
  return JSON.stringify(output.result, null, 2);
}

(eventType: string, data: unknown, context: { toolCallId?: string }) => void