UIHeadless

Building Custom UI

This guide walks through building a complete chat interface from scratch using only @aomi-labs/react hooks. You control every pixel.

Setup

Install the headless library:

npm install @aomi-labs/react @assistant-ui/react

Wrap your app with the runtime provider:

// app/providers.tsx
import { AomiRuntimeProvider } from "@aomi-labs/react";

export function Providers({ children }: { children: React.ReactNode }) {
  return (
    <AomiRuntimeProvider backendUrl={process.env.NEXT_PUBLIC_BACKEND_URL!}>
      {children}
    </AomiRuntimeProvider>
  );
}

Message List Component

Display messages for the current thread:

// components/message-list.tsx
"use client";

import { useAomiRuntime } from "@aomi-labs/react";

export function MessageList() {
  const { currentThreadId, getMessages, isRunning } = useAomiRuntime();
  const messages = getMessages(currentThreadId);

  return (
    <div className="flex-1 overflow-y-auto p-4 space-y-4">
      {messages.map((msg, i) => (
        <div
          key={i}
          className={`flex ${msg.role === "user" ? "justify-end" : "justify-start"}`}
        >
          <div
            className={`max-w-[80%] rounded-lg px-4 py-2 ${
              msg.role === "user"
                ? "bg-blue-600 text-white"
                : "bg-gray-100 text-gray-900"
            }`}
          >
            {msg.content.map((block, j) => (
              <span key={j}>
                {block.type === "text" ? block.text : null}
              </span>
            ))}
          </div>
        </div>
      ))}

      {isRunning && (
        <div className="flex justify-start">
          <div className="bg-gray-100 rounded-lg px-4 py-2 text-gray-500">
            Thinking...
          </div>
        </div>
      )}
    </div>
  );
}

Chat Input Component

A composer that sends messages:

// components/chat-input.tsx
"use client";

import { useState } from "react";
import { useAomiRuntime } from "@aomi-labs/react";

export function ChatInput() {
  const { sendMessage, isRunning, cancelGeneration } = useAomiRuntime();
  const [input, setInput] = useState("");

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!input.trim() || isRunning) return;

    const text = input;
    setInput("");
    await sendMessage(text);
  };

  return (
    <form onSubmit={handleSubmit} className="border-t p-4 flex gap-2">
      <input
        type="text"
        value={input}
        onChange={(e) => setInput(e.target.value)}
        placeholder="Type a message..."
        className="flex-1 rounded-lg border px-4 py-2"
        disabled={isRunning}
      />
      {isRunning ? (
        <button
          type="button"
          onClick={cancelGeneration}
          className="rounded-lg bg-red-500 px-4 py-2 text-white"
        >
          Stop
        </button>
      ) : (
        <button
          type="submit"
          disabled={!input.trim()}
          className="rounded-lg bg-blue-600 px-4 py-2 text-white disabled:opacity-50"
        >
          Send
        </button>
      )}
    </form>
  );
}

Thread Switcher Component

Navigate between conversation threads:

// components/thread-switcher.tsx
"use client";

import { useAomiRuntime } from "@aomi-labs/react";

export function ThreadSwitcher() {
  const {
    currentThreadId,
    threadMetadata,
    selectThread,
    createThread,
    deleteThread,
  } = useAomiRuntime();

  const threads = Array.from(threadMetadata.entries());

  return (
    <div className="w-64 border-r h-full flex flex-col">
      <div className="p-3 border-b">
        <button
          onClick={() => createThread()}
          className="w-full rounded-lg bg-gray-100 px-3 py-2 text-sm hover:bg-gray-200"
        >
          + New Thread
        </button>
      </div>

      <div className="flex-1 overflow-y-auto">
        {threads.map(([id, meta]) => (
          <div
            key={id}
            onClick={() => selectThread(id)}
            className={`flex items-center justify-between px-3 py-2 cursor-pointer text-sm ${
              id === currentThreadId
                ? "bg-blue-50 text-blue-700"
                : "hover:bg-gray-50"
            }`}
          >
            <span className="truncate">{meta.title || "Untitled"}</span>
            <button
              onClick={(e) => {
                e.stopPropagation();
                deleteThread(id);
              }}
              className="text-gray-400 hover:text-red-500 text-xs"
            >
              x
            </button>
          </div>
        ))}
      </div>
    </div>
  );
}

Composing Everything

Put it all together on a page:

// app/chat/page.tsx
"use client";

import { MessageList } from "@/components/message-list";
import { ChatInput } from "@/components/chat-input";
import { ThreadSwitcher } from "@/components/thread-switcher";

export default function ChatPage() {
  return (
    <div className="flex h-screen">
      <ThreadSwitcher />
      <div className="flex flex-1 flex-col">
        <MessageList />
        <ChatInput />
      </div>
    </div>
  );
}

Handling Streaming Responses

The isRunning flag from useAomiRuntime is true while the assistant is generating. Messages update in real time as tokens stream in, so your MessageList component will re-render with updated content automatically.

For more granular streaming control, subscribe to SSE events:

import { useEffect } from "react";
import { useEventContext } from "@aomi-labs/react";

function StreamMonitor() {
  const { subscribe, sseStatus } = useEventContext();

  useEffect(() => {
    const unsub = subscribe("streaming_text", (event) => {
      console.log("Streaming chunk:", event.payload);
    });
    return unsub;
  }, [subscribe]);

  return <div>SSE: {sseStatus}</div>;
}

Key Principle

The headless library provides all the state and actions. Your components are purely presentational. There are no layout constraints, no CSS to override, and no component hierarchy to follow. You decide how messages look, where threads are listed, and how controls are arranged.

On this page