NodeAgentThread.tsx

/**
 * NodeAgentThread — a custom-styled assistant-ui Thread.
 *
 * Built from assistant-ui's headless primitives (`ThreadPrimitive`,
 * `MessagePrimitive`, `ComposerPrimitive`) and themed with the NodeBench design
 * DNA — no Tailwind, no shadcn. The assistant's message renders text plus the
 * four tool UIs inline (registered in toolUIs.tsx).
 */

import { useRef, type MouseEvent } from "react";
import {
  ActionBarPrimitive,
  AuiIf,
  ComposerPrimitive,
  MessagePrimitive,
  ThreadPrimitive,
  useAui,
  useAuiState,
} from "@assistant-ui/react";
import { DEMO_QUESTION } from "../demoScenario";

function focusComposer(event: MouseEvent<HTMLButtonElement>) {
  event.currentTarget.closest(".na-thread")?.querySelector<HTMLTextAreaElement>("textarea[name=input]")?.focus();
}

function UserMessage() {
  return (
    <MessagePrimitive.Root className="na-msg na-msg-user">
      <div className="na-bubble na-bubble-user">
        <MessagePrimitive.Parts />
      </div>
    </MessagePrimitive.Root>
  );
}

function AssistantMessage() {
  const aui = useAui();
  const messageId = useAuiState((s) => s.message.id);
  const submittedResponse = useRef<string | null>(null);
  const retryResponse = (event: MouseEvent<HTMLButtonElement>) => {
    if (aui.thread().getState().isRunning || submittedResponse.current === messageId) {
      event.preventDefault();
      return;
    }
    submittedResponse.current = messageId;
    focusComposer(event);
  };
  return (
    <MessagePrimitive.Root className="na-msg na-msg-assistant">
      <span className="na-av" aria-hidden>
        ◆
      </span>
      <div className="na-bubble na-bubble-assistant">
        <MessagePrimitive.Parts />
        <AuiIf condition={(s) => s.message.status?.type === "incomplete" && (s.message.status.reason === "cancelled" || s.message.status.reason === "error")}>
          <div className="na-recovery">
            <AuiIf condition={(s) => s.message.status?.type === "incomplete" && s.message.status.reason === "cancelled"}>
              <p role="status">Stopped this response. Completed work is kept.</p>
            </AuiIf>
            <MessagePrimitive.Error>
              <p role="alert">This response failed. Earlier results are kept.</p>
            </MessagePrimitive.Error>
            <ActionBarPrimitive.Reload onClick={retryResponse} onFocus={(event) => {
              if (event.currentTarget.matches(":focus-visible")) {
                event.currentTarget.parentElement?.scrollIntoView({ block: "start" });
              }
            }}>Retry response</ActionBarPrimitive.Reload>
          </div>
        </AuiIf>
      </div>
    </MessagePrimitive.Root>
  );
}

export function NodeAgentThread() {
  return (
    <ThreadPrimitive.Root className="na-thread">
      <ThreadPrimitive.Viewport className="na-viewport">
        <AuiIf condition={(s) => s.thread.isEmpty}>
          <div className="na-empty">
            <span className="na-av na-av-lg" aria-hidden>
              ◆
            </span>
            <h2>Ask the room</h2>
            <p>
              NodeAgent gathers the relevant context, finds the right document for the right
              answer, updates the model as a versioned delta, and writes the cited memo — and
              you watch each step render inline.
            </p>
            <div className="na-suggestions">
              <ThreadPrimitive.Suggestion
                className="na-chip"
                prompt={DEMO_QUESTION}
                method="replace"
                autoSend
              >
                Does our wedge hold + survive 18 months? →
              </ThreadPrimitive.Suggestion>
              <ThreadPrimitive.Suggestion
                className="na-chip"
                prompt="What changed since the last investor update, and does the model still hold?"
                method="replace"
                autoSend
              >
                What changed since last update? →
              </ThreadPrimitive.Suggestion>
            </div>
          </div>
        </AuiIf>

        <ThreadPrimitive.Messages>
          {({ message }) =>
            message.role === "user" ? <UserMessage /> : <AssistantMessage />
          }
        </ThreadPrimitive.Messages>

        <ThreadPrimitive.ViewportFooter className="na-footer">
          <ComposerPrimitive.Root className="na-composer">
            <ComposerPrimitive.Input
              className="na-composer-input"
              placeholder="Ask the room anything…"
              rows={1}
              autoFocus
            />
            <AuiIf condition={(s) => !s.thread.isRunning}>
              <ComposerPrimitive.Send className="na-send" aria-label="Send">
                ↑
              </ComposerPrimitive.Send>
            </AuiIf>
            <AuiIf condition={(s) => s.thread.isRunning}>
              <ComposerPrimitive.Cancel className="na-send na-stop" aria-label="Stop response" onClick={focusComposer}>
                Stop
              </ComposerPrimitive.Cancel>
            </AuiIf>
          </ComposerPrimitive.Root>
          <p className="na-disclaimer">
            Deterministic demo — no keys. The agent runs the real ported modules over a fixed
            scenario.
          </p>
        </ThreadPrimitive.ViewportFooter>
      </ThreadPrimitive.Viewport>
    </ThreadPrimitive.Root>
  );
}