);
}
// Context for sharing chat state with the message view
const ChatStateContext = createContext(null);
// Wildcard tool call renderer for all server-side tools
const toolCallRenderers = [
defineToolCallRenderer({
name: "*",
render: ({ name, args, result }) => (
}
/>
),
}),
];
// Custom message view that injects CitationBlocks after assistant responses.
// Uses CopilotChatMessageView's children render prop to post-process the
// rendered message elements and inject citations at the right positions.
function MessageViewWithCitations({
messages,
isRunning,
}: {
// biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union
messages: any[];
isRunning: boolean;
}) {
const chatState = useContext(ChatStateContext);
const cursor = isRunning ? (
) : null;
return (
{({ messageElements }) => {
if (!chatState?.qa_history?.length) {
return (
<>
{messageElements}
{cursor}
);
}
// CopilotChatMessageView renders one element per user/assistant/activity
// message (tool messages produce nothing). We correlate elements with
// messages to inject CitationBlocks after the right assistant responses.
//
// Tool call objects on messages only carry `id` (no `name`), so we
// can't identify which tool was called from the message alone. Instead
// we rely on the fact that qa_history only grows when the `ask` tool
// runs: after each assistant text response that followed tool calls,
// we inject any new qa_history citations.
const result: React.ReactNode[] = [];
let qaIdx = 0;
let seenToolCalls = false;
let elemIdx = 0;
for (const msg of messages) {
if (msg.role === "user") {
seenToolCalls = false;
}
if (
msg.role === "assistant" &&
Array.isArray(msg.toolCalls) &&
msg.toolCalls.length > 0
) {
seenToolCalls = true;
}
const isRendered =
msg.role === "user" ||
msg.role === "assistant" ||
msg.role === "activity";
if (!isRendered) continue;
if (elemIdx < messageElements.length) {
result.push(messageElements[elemIdx]);
elemIdx++;
}
// After an assistant text response that followed tool calls,
// inject the next qa_history entry's citations (one per turn)
if (msg.role === "assistant" && msg.content && seenToolCalls) {
if (qaIdx < chatState.qa_history.length) {
const qa = chatState.qa_history[qaIdx];
if (qa.citations?.length) {
result.push(
,
);
}
qaIdx++;
}
seenToolCalls = false;
}
}
while (elemIdx < messageElements.length) {
result.push(messageElements[elemIdx]);
elemIdx++;
}
return (
<>
{result}
{cursor}
);
}}
);
}
function ChatContentInner({
sessionId,
onSessionChange,
}: {
sessionId: string;
onSessionChange: (id: string) => void;
}) {
const [contextOpen, setContextOpen] = useState(false);
const [filterOpen, setFilterOpen] = useState(false);
const { agent } = useAgent({
agentId: "chat_agent",
updates: [
UseAgentUpdate.OnMessagesChanged,
UseAgentUpdate.OnStateChanged,
UseAgentUpdate.OnRunStatusChanged,
],
});
const { copilotkit: ck } = useCopilotKit();
// Set threadId (CopilotChat normally does this in its connect effect)
useEffect(() => {
agent.threadId = sessionId;
}, [agent, sessionId]);
const chatState = normalizeChatState(
(agent.state as AgentState)?.[AGUI_STATE_KEY],
);
const mergeChatState = (partial: Partial) => {
const current = normalizeChatState(
(agent.state as AgentState)?.[AGUI_STATE_KEY],
);
agent.setState({
...agent.state,
[AGUI_STATE_KEY]: {
...current,
...partial,
},
});
};
// Restore session from localStorage when agent reference changes.
// useAgent returns a provisional agent initially, then the real agent
// after runtime connects — re-run restore each time so messages stick.
useEffect(() => {
if (agent.messages.length > 0) return;
const session = getSession(sessionId);
if (!session) return;
if (session.chatState) {
agent.setState({
[AGUI_STATE_KEY]: normalizeChatState(session.chatState),
});
}
if (session.messages.length > 0) {
// biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union
agent.setMessages(session.messages as any[]);
}
}, [agent, sessionId]);
// Persist messages and state to localStorage.
// Read chatState from agent.state at effect time (not render time) so that
// restore and persist effects in the same commit see consistent state.
// biome-ignore lint/correctness/useExhaustiveDependencies: JSON.stringify tracks content changes
useEffect(() => {
if (sessionId && agent.messages.length > 0) {
const currentChatState = normalizeChatState(
(agent.state as AgentState)?.[AGUI_STATE_KEY],
);
updateSessionMessages(
sessionId,
serializeMessages(agent.messages),
currentChatState,
);
}
}, [JSON.stringify(agent.messages), chatState, sessionId]);
// biome-ignore lint/correctness/useExhaustiveDependencies: stable identity via agent ref
const messages = useMemo(
() => [...agent.messages],
[JSON.stringify(agent.messages)],
);
const onSubmitMessage = useCallback(
async (text: string) => {
agent.addMessage({
id: crypto.randomUUID(),
role: "user",
content: text,
});
try {
await ck.runAgent({ agent });
} catch (error) {
console.error("runAgent failed", error);
}
},
[agent, ck],
);
const onStop = useCallback(() => {
try {
ck.stopAgent({ agent });
} catch {
agent.abortRun();
}
}, [agent, ck]);
// Strip the scroll view's internal paddingBottom (inline style set by
// CopilotChatView to reserve space for its absolutely-positioned input
// container, which we've overridden to flow in a flex layout).
// Uses a ref callback so it runs when the element actually mounts.
const paddingObserver = useRef(null);
const scrollAreaCallbackRef = useCallback((node: HTMLDivElement | null) => {
if (paddingObserver.current) {
paddingObserver.current.disconnect();
paddingObserver.current = null;
}
if (!node) return;
const strip = () => {
for (const div of node.querySelectorAll("div")) {
if (div.style.paddingBottom && div.style.paddingBottom !== "1rem") {
div.style.paddingBottom = "1rem";
}
}
};
strip();
const observer = new MutationObserver(strip);
observer.observe(node, {
subtree: true,
attributes: true,
attributeFilter: ["style"],
});
paddingObserver.current = observer;
}, []);
const sessionContext = chatState.session_context;
const documentFilter = chatState.document_filter;
const initialContext = chatState.initial_context ?? "";
// Context is locked after first message (qa_history has entries)
const isContextLocked = (chatState.qa_history?.length ?? 0) > 0;
const handleFilterApply = (selected: string[]) => {
mergeChatState({ document_filter: selected });
};
const handleInitialContextChange = (value: string) => {
if (isContextLocked) return;
mergeChatState({ initial_context: value || null });
};
return (
{({ scrollView, feather, inputContainer }) => (
{scrollView}
{feather}
{inputContainer}
)}
setContextOpen(false)}
sessionContext={sessionContext}
initialContext={initialContext}
onInitialContextChange={handleInitialContextChange}
isLocked={isContextLocked}
/>
setFilterOpen(false)}
selected={documentFilter}
onApply={handleFilterApply}
/>
);
}
export default function Chat() {
const [activeSessionId, setActiveSessionId] = useState(null);
useEffect(() => {
let id = getActiveSessionId();
if (!id) {
id = createSession().id;
}
setActiveSessionId(id);
}, []);
if (!activeSessionId) return null;
return (
);
}