Inside your chat UI

'use client';

import type { UIMessage } from 'ai';
import {
  DiffReview,
  ToolApprovalCard,
  getToolPartName,
  isToolPart,
  type DiffReviewResult,
  type FileChange,
  type ToolApprovalResponse,
} from 'signoff-ui';

export function MessageParts({
  message,
  onApproval,
  onReview,
}: {
  message: UIMessage;
  /** `useChat().addToolApprovalResponse` */
  onApproval: (response: ToolApprovalResponse) => void;
  /** Calls `useChat().addToolOutput` with `reviewToolOutput(review)` as the tool's output. */
  onReview: (toolCallId: string, review: DiffReviewResult) => void;
}) {
  return message.parts.map((part, index) => {
    if (part.type === 'text') return <p key={index}>{part.text}</p>;
    if (!isToolPart(part)) return null;
    if (getToolPartName(part) === 'review_changes' && part.state === 'input-available') {
      const { files } = part.input as { files: FileChange[] };
      return (
        <DiffReview key={part.toolCallId} files={files} onSubmit={(review) => onReview(part.toolCallId, review)} />
      );
    }
    // Renders nothing for calls that need no approval.
    return <ToolApprovalCard key={part.toolCallId} part={part} onRespond={onApproval} />;
  });
}

if (isToolPart(part) && getToolPartName(part) === 'review_changes' && part.state === 'input-available') {
  return <DiffReview key={part.toolCallId} files={(part.input as { files: FileChange[] }).files} onSubmit={…} />;
}
if (isToolPart(part) && part.approval) {
  return <ToolApprovalCard key={part.toolCallId} part={part} onRespond={addToolApprovalResponse} />;
}
// …and AI Elements' Tool for the other calls

'use client';

import { AssistantRuntimeProvider, ComposerPrimitive, MessagePrimitive, ThreadPrimitive } from '@assistant-ui/react';
import { useChatRuntime } from '@assistant-ui/react-ai-sdk';
import { lastAssistantMessageIsCompleteWithApprovalResponses, lastAssistantMessageIsCompleteWithToolCalls } from 'ai';
import { useApprovalPolicy } from 'signoff-ui';
import { SignoffToolsProvider, signoffTools } from 'signoff-ui/assistant-ui';

// review_changes gets a DiffReview; any call with an approval gate gets the approval card.
const tools = signoffTools({ review: ['review_changes'] });

function AssistantMessage() {
  return (
    <MessagePrimitive.Root>
      <MessagePrimitive.Parts components={{ tools }} />
    </MessagePrimitive.Root>
  );
}

function UserMessage() {
  return (
    <MessagePrimitive.Root>
      <MessagePrimitive.Parts />
    </MessagePrimitive.Root>
  );
}

export function Thread() {
  const runtime = useChatRuntime({
    // Continue the run once the review is in and every approval has an answer.
    sendAutomaticallyWhen: (chat) =>
      lastAssistantMessageIsCompleteWithToolCalls(chat) || lastAssistantMessageIsCompleteWithApprovalResponses(chat),
  });
  const policy = useApprovalPolicy();
  return (
    <AssistantRuntimeProvider runtime={runtime}>
      <SignoffToolsProvider policy={policy} tools={{ run_command: { risk: 'high' } }}>
        <ThreadPrimitive.Root>
          <ThreadPrimitive.Viewport>
            <ThreadPrimitive.Messages>
              {({ message }) => (message.role === 'assistant' ? <AssistantMessage /> : <UserMessage />)}
            </ThreadPrimitive.Messages>
          </ThreadPrimitive.Viewport>
          <ComposerPrimitive.Root>
            <ComposerPrimitive.Input aria-label="Message the agent" />
            <ComposerPrimitive.Send>Send</ComposerPrimitive.Send>
          </ComposerPrimitive.Root>
        </ThreadPrimitive.Root>
      </SignoffToolsProvider>
    </AssistantRuntimeProvider>
  );
}