"use client" import { useEffect, useMemo, useRef, useState } from "react" import type { AiChatContext } from "@/lib/ai/chat-context" import { buildEmbedSearchParams, systemPromptFromContext } from "@/lib/ai/chat-context" import { buildAiEmbedUrl, resolveAiEmbedOrigin } from "@/lib/ai/embed-url" import { useAiIframeExternalLinks } from "@/lib/ai/use-ai-iframe-navigation" import { useAiConfig, useCreateAiSession } from "@/lib/api/hooks/use-ai-queries" import { useTheme } from "next-themes" type AiChatIframeProps = { publicPath?: string context: AiChatContext className?: string } export function AiChatIframe({ publicPath = "/ai", context, className }: AiChatIframeProps) { const iframeRef = useRef(null) const { resolvedTheme } = useTheme() const { data: config, isSuccess } = useAiConfig() const createSession = useCreateAiSession() const [sessionToken, setSessionToken] = useState() const [sessionId, setSessionId] = useState() const sessionContextKey = useMemo( () => [ context.app, context.temporary, context.messageId, context.accountId, context.drivePath, context.fileId, context.contactId, ].join("|"), [ context.app, context.temporary, context.messageId, context.accountId, context.drivePath, context.fileId, context.contactId, ] ) const embedOrigin = useMemo(() => resolveAiEmbedOrigin(publicPath), [publicPath]) const enabledTools = config?.enabled_tools ?? [] const mcpUrl = config?.mcp_url ?? "/api/v1/ai/mcp" const iframeSrc = useMemo(() => { if (!isSuccess || !config?.enabled) return null const qs = buildEmbedSearchParams(context, config.default_model) return buildAiEmbedUrl(publicPath, qs) }, [isSuccess, config, publicPath, context]) useAiIframeExternalLinks(embedOrigin) useEffect(() => { if (!config?.enabled) return let cancelled = false createSession .mutateAsync(context) .then((session) => { if (cancelled) return setSessionToken(session.token_secret) setSessionId(session.session_id) }) .catch(() => { if (!cancelled) { setSessionToken(undefined) setSessionId(undefined) } }) return () => { cancelled = true } }, [config?.enabled, sessionContextKey, context, createSession.mutateAsync]) useEffect(() => { const iframe = iframeRef.current if (!iframe?.contentWindow || !embedOrigin) return iframe.contentWindow.postMessage( { type: "ULTI_THEME", theme: resolvedTheme === "dark" ? "dark" : "light" }, embedOrigin ) }, [resolvedTheme, embedOrigin]) useEffect(() => { const iframe = iframeRef.current if (!iframe?.contentWindow || !embedOrigin) return const systemPrompt = [ systemPromptFromContext(context), context.systemPromptExtra, ] .filter(Boolean) .join("\n\n") iframe.contentWindow.postMessage( { type: "ULTI_CONTEXT_UPDATE", context, systemPrompt: systemPrompt || undefined, }, embedOrigin ) }, [context, embedOrigin]) useEffect(() => { const iframe = iframeRef.current if (!iframe?.contentWindow || !embedOrigin || !sessionToken) return iframe.contentWindow.postMessage( { type: "ULTI_SESSION", token_secret: sessionToken, session_id: sessionId, mcp_url: mcpUrl, enabled_tools: enabledTools, default_model: config?.default_model, }, embedOrigin ) }, [sessionToken, sessionId, mcpUrl, enabledTools, config?.default_model, embedOrigin]) if (!iframeSrc) { return (
) } return (