Merge pull request #929 from AnishSarkar22/impr/thinking-steps

Impr: Updated to assistant-ui v0.12.x & many UI changes

surfsense_backend/app/agents/new_chat/__init__.py

@@ -5,7 +5,7 @@ This module provides the SurfSense deep agent with configurable tools
 for knowledge base search, podcast generation, and more.
 
 Directory Structure:
-- tools/: All agent tools (knowledge_base, podcast, link_preview, etc.)
+- tools/: All agent tools (knowledge_base, podcast, generate_image, etc.)
 - chat_deepagent.py: Main agent factory
 - system_prompt.py: System prompts and instructions
 - context.py: Context schema for the agent
@@ -37,9 +37,7 @@ from .tools import (
     BUILTIN_TOOLS,
     ToolDefinition,
     build_tools,
-    create_display_image_tool,
     create_generate_podcast_tool,
-    create_link_preview_tool,
     create_scrape_webpage_tool,
     create_search_knowledge_base_tool,
     format_documents_for_context,
@@ -63,9 +61,7 @@ __all__ = [
     # LLM config
     "create_chat_litellm_from_config",
     # Tool factories
-    "create_display_image_tool",
     "create_generate_podcast_tool",
-    "create_link_preview_tool",
     "create_scrape_webpage_tool",
     "create_search_knowledge_base_tool",
     # Agent factory

surfsense_backend/app/agents/new_chat/chat_deepagent.py

@@ -150,8 +150,6 @@ async def create_surfsense_deep_agent(
     - search_knowledge_base: Search the user's personal knowledge base
     - generate_podcast: Generate audio podcasts from content
     - generate_image: Generate images from text descriptions using AI models
-    - link_preview: Fetch rich previews for URLs
-    - display_image: Display images in chat
     - scrape_webpage: Extract content from webpages
     - save_memory: Store facts/preferences about the user
     - recall_memory: Retrieve relevant user memories
@@ -207,7 +205,7 @@ async def create_surfsense_deep_agent(
         # Create agent with only specific tools
         agent = create_surfsense_deep_agent(
             llm, search_space_id, db_session, ...,
-            enabled_tools=["search_knowledge_base", "link_preview"]
+            enabled_tools=["search_knowledge_base", "scrape_webpage"]
         )
 
         # Create agent without podcast generation

surfsense_backend/app/agents/new_chat/system_prompt.py

@@ -184,48 +184,6 @@ _TOOL_INSTRUCTIONS["generate_report"] = """
   - AFTER CALLING THIS TOOL: Do NOT repeat, summarize, or reproduce the report content in the chat. The report is already displayed as an interactive card that the user can open, read, copy, and export. Simply confirm that the report was generated (e.g., "I've generated your report on [topic]. You can view the Markdown report now, and export it in various formats from the card."). NEVER write out the report text in the chat.
 """
 
-_TOOL_INSTRUCTIONS["link_preview"] = """
-- link_preview: Fetch metadata for a URL to display a rich preview card.
-  - IMPORTANT: Use this tool WHENEVER the user shares or mentions a URL/link in their message.
-  - This fetches the page's Open Graph metadata (title, description, thumbnail) to show a preview card.
-  - NOTE: This tool only fetches metadata, NOT the full page content. It cannot read the article text.
-  - Trigger scenarios:
-    * User shares a URL (e.g., "Check out https://example.com")
-    * User pastes a link in their message
-    * User asks about a URL or link
-  - Args:
-    - url: The URL to fetch metadata for (must be a valid HTTP/HTTPS URL)
-  - Returns: A rich preview card with title, description, thumbnail, and domain
-  - The preview card will automatically be displayed in the chat.
-"""
-
-_TOOL_INSTRUCTIONS["display_image"] = """
-- display_image: Display an image in the chat with metadata.
-  - Use this tool ONLY when you have a valid public HTTP/HTTPS image URL to show.
-  - This displays the image with an optional title, description, and source attribution.
-  - Valid use cases:
-    * Showing an image from a URL the user explicitly mentioned in their message
-    * Displaying images found in scraped webpage content (from scrape_webpage tool)
-    * Showing a publicly accessible diagram or chart from a known URL
-    * Displaying an AI-generated image after calling the generate_image tool (ALWAYS required)
-  
-  CRITICAL - NEVER USE THIS TOOL FOR USER-UPLOADED ATTACHMENTS:
-  When a user uploads/attaches an image file to their message:
-    * The image is ALREADY VISIBLE in the chat UI as a thumbnail on their message
-    * You do NOT have a URL for their uploaded image - only extracted text/description
-    * Calling display_image will FAIL and show "Image not available" error
-    * Simply analyze the image content and respond with your analysis - DO NOT try to display it
-    * The user can already see their own uploaded image - they don't need you to show it again
-  
-  - Args:
-    - src: The URL of the image (MUST be a valid public HTTP/HTTPS URL that you know exists)
-    - alt: Alternative text describing the image (for accessibility)
-    - title: Optional title to display below the image
-    - description: Optional description providing context about the image
-  - Returns: An image card with the image, title, and description
-  - The image will automatically be displayed in the chat.
-"""
-
 _TOOL_INSTRUCTIONS["generate_image"] = """
 - generate_image: Generate images from text descriptions using AI image models.
   - Use this when the user asks you to create, generate, draw, design, or make an image.
@@ -233,10 +191,7 @@ _TOOL_INSTRUCTIONS["generate_image"] = """
   - Args:
     - prompt: A detailed text description of the image to generate. Be specific about subject, style, colors, composition, and mood.
     - n: Number of images to generate (1-4, default: 1)
-  - Returns: A dictionary with the generated image URL in the "src" field, along with metadata.
-  - CRITICAL: After calling generate_image, you MUST call `display_image` with the returned "src" URL
-    to actually show the image in the chat. The generate_image tool only generates the image and returns
-    the URL — it does NOT display anything. You must always follow up with display_image.
+  - Returns: A dictionary with the generated image metadata. The image will automatically be displayed in the chat.
   - IMPORTANT: Write a detailed, descriptive prompt for best results. Don't just pass the user's words verbatim -
     expand and improve the prompt with specific details about style, lighting, composition, and mood.
   - If the user's request is vague (e.g., "make me an image of a cat"), enhance the prompt with artistic details.
@@ -245,14 +200,11 @@ _TOOL_INSTRUCTIONS["generate_image"] = """
 _TOOL_INSTRUCTIONS["scrape_webpage"] = """
 - scrape_webpage: Scrape and extract the main content from a webpage.
   - Use this when the user wants you to READ and UNDERSTAND the actual content of a webpage.
-  - IMPORTANT: This is different from link_preview:
-    * link_preview: Only fetches metadata (title, description, thumbnail) for display
-    * scrape_webpage: Actually reads the FULL page content so you can analyze/summarize it
   - CRITICAL — WHEN TO USE (always attempt scraping, never refuse before trying):
     * When a user asks to "get", "fetch", "pull", "grab", "scrape", or "read" content from a URL
     * When the user wants live/dynamic data from a specific webpage (e.g., tables, scores, stats, prices)
     * When a URL was mentioned earlier in the conversation and the user asks for its actual content
-    * When link_preview or search_knowledge_base returned insufficient data and the user wants more
+    * When search_knowledge_base returned insufficient data and the user wants more
   - Trigger scenarios:
     * "Read this article and summarize it"
     * "What does this page say about X?"
@@ -268,9 +220,10 @@ _TOOL_INSTRUCTIONS["scrape_webpage"] = """
     - url: The URL of the webpage to scrape (must be HTTP/HTTPS)
     - max_length: Maximum content length to return (default: 50000 chars)
   - Returns: The page title, description, full content (in markdown), word count, and metadata
-  - After scraping, you will have the full article text and can analyze, summarize, or answer questions about it.
+  - After scraping, provide a comprehensive, well-structured summary with key takeaways using headings or bullet points.
+  - Reference the source using markdown links [descriptive text](url) — never bare URLs.
   - IMAGES: The scraped content may contain image URLs in markdown format like `![alt text](image_url)`.
-    * When you find relevant/important images in the scraped content, use the `display_image` tool to show them to the user.
+    * When you find relevant/important images in the scraped content, include them in your response using standard markdown image syntax: `![alt text](image_url)`.
     * This makes your response more visual and engaging.
     * Prioritize showing: diagrams, charts, infographics, key illustrations, or images that help explain the content.
     * Don't show every image - just the most relevant 1-3 images that enhance understanding.
@@ -292,6 +245,8 @@ _TOOL_INSTRUCTIONS["web_search"] = """
   - Args:
     - query: The search query - use specific, descriptive terms
     - top_k: Number of results to retrieve (default: 10, max: 50)
+  - If search snippets are insufficient for the user's question, use `scrape_webpage` on the most relevant result URL for full content.
+  - When presenting results, reference sources as markdown links [descriptive text](url) — never bare URLs.
 """
 
 # Memory tool instructions have private and shared variants.
@@ -476,32 +431,31 @@ _TOOL_EXAMPLES["generate_report"] = """
 
 _TOOL_EXAMPLES["scrape_webpage"] = """
 - User: "Check out https://dev.to/some-article"
-  - Call: `link_preview(url="https://dev.to/some-article")`
   - Call: `scrape_webpage(url="https://dev.to/some-article")`
-  - Then provide your analysis of the content.
+  - Respond with a structured analysis — key points, takeaways.
 - User: "Read this article and summarize it for me: https://example.com/blog/ai-trends"
   - Call: `scrape_webpage(url="https://example.com/blog/ai-trends")`
-  - Then provide a summary based on the scraped text.
+  - Respond with a thorough summary using headings and bullet points.
 - User: (after discussing https://example.com/stats) "Can you get the live data from that page?"
   - Call: `scrape_webpage(url="https://example.com/stats")`
   - IMPORTANT: Always attempt scraping first. Never refuse before trying the tool.
-"""
-
-_TOOL_EXAMPLES["display_image"] = """
-- User: "Show me this image: https://example.com/image.png"
-  - Call: `display_image(src="https://example.com/image.png", alt="User shared image")`
-- User uploads an image file and asks: "What is this image about?"
-  - DO NOT call display_image! The user's uploaded image is already visible in the chat.
-  - Simply analyze the image content and respond directly.
+- User: "https://example.com/blog/weekend-recipes"
+  - Call: `scrape_webpage(url="https://example.com/blog/weekend-recipes")`
+  - When a user sends just a URL with no instructions, scrape it and provide a concise summary of the content.
 """
 
 _TOOL_EXAMPLES["generate_image"] = """
 - User: "Generate an image of a cat"
-  - Step 1: `generate_image(prompt="A fluffy orange tabby cat sitting on a windowsill, bathed in warm golden sunlight, soft bokeh background with green houseplants, photorealistic style, cozy atmosphere")`
-  - Step 2: Use the returned "src" URL to display it: `display_image(src="<returned_url>", alt="A fluffy orange tabby cat on a windowsill", title="Generated Image")`
+  - Call: `generate_image(prompt="A fluffy orange tabby cat sitting on a windowsill, bathed in warm golden sunlight, soft bokeh background with green houseplants, photorealistic style, cozy atmosphere")`
+  - The generated image will automatically be displayed in the chat.
 - User: "Draw me a logo for a coffee shop called Bean Dream"
-  - Step 1: `generate_image(prompt="Minimalist modern logo design for a coffee shop called 'Bean Dream', featuring a stylized coffee bean with dream-like swirls of steam, clean vector style, warm brown and cream color palette, white background, professional branding")`
-  - Step 2: `display_image(src="<returned_url>", alt="Bean Dream coffee shop logo", title="Generated Image")`
+  - Call: `generate_image(prompt="Minimalist modern logo design for a coffee shop called 'Bean Dream', featuring a stylized coffee bean with dream-like swirls of steam, clean vector style, warm brown and cream color palette, white background, professional branding")`
+  - The generated image will automatically be displayed in the chat.
+- User: "Show me this image: https://example.com/image.png"
+  - Simply include it in your response using markdown: `![Image](https://example.com/image.png)`
+- User uploads an image file and asks: "What is this image about?"
+  - The user's uploaded image is already visible in the chat.
+  - Simply analyze the image content and respond directly.
 """
 
 _TOOL_EXAMPLES["web_search"] = """
@@ -522,8 +476,6 @@ _ALL_TOOL_NAMES_ORDERED = [
     "generate_podcast",
     "generate_video_presentation",
     "generate_report",
-    "link_preview",
-    "display_image",
     "generate_image",
     "scrape_webpage",
     "save_memory",
@@ -764,7 +716,7 @@ Do not use the sandbox for:
 
 When your code creates output files (images, CSVs, PDFs, etc.) in the sandbox:
 - **Print the absolute path** at the end of your script so the user can download the file. Example: `print("SANDBOX_FILE: /tmp/chart.png")`
-- **DO NOT call `display_image`** for files created inside the sandbox. Sandbox files are not accessible via public URLs, so `display_image` will always show "Image not available". The frontend automatically renders a download button from the `SANDBOX_FILE:` marker.
+- **DO NOT use markdown image syntax** for files created inside the sandbox. Sandbox files are not accessible via public URLs and will show "Image not available". The frontend automatically renders a download button from the `SANDBOX_FILE:` marker.
 - You can output multiple files, one per line: `print("SANDBOX_FILE: /tmp/report.csv")`, `print("SANDBOX_FILE: /tmp/chart.png")`
 - Always describe what the file contains in your response text so the user knows what they are downloading.
 - IMPORTANT: Every `execute` call that saves a file MUST print the `SANDBOX_FILE: <path>` marker. Without it the user cannot download the file.

surfsense_backend/app/agents/new_chat/tools/__init__.py

@@ -10,8 +10,6 @@ Available tools:
 - generate_podcast: Generate audio podcasts from content
 - generate_video_presentation: Generate video presentations with slides and narration
 - generate_image: Generate images from text descriptions using AI models
-- link_preview: Fetch rich previews for URLs
-- display_image: Display images in chat
 - scrape_webpage: Extract content from webpages
 - save_memory: Store facts/preferences about the user
 - recall_memory: Retrieve relevant user memories
@@ -19,7 +17,6 @@ Available tools:
 
 # Registry exports
 # Tool factory exports (for direct use)
-from .display_image import create_display_image_tool
 from .generate_image import create_generate_image_tool
 from .knowledge_base import (
     CONNECTOR_DESCRIPTIONS,
@@ -27,7 +24,6 @@ from .knowledge_base import (
     format_documents_for_context,
     search_knowledge_base_async,
 )
-from .link_preview import create_link_preview_tool
 from .podcast import create_generate_podcast_tool
 from .registry import (
     BUILTIN_TOOLS,
@@ -50,11 +46,9 @@ __all__ = [
     "ToolDefinition",
     "build_tools",
     # Tool factories
-    "create_display_image_tool",
     "create_generate_image_tool",
     "create_generate_podcast_tool",
     "create_generate_video_presentation_tool",
-    "create_link_preview_tool",
     "create_recall_memory_tool",
     "create_save_memory_tool",
     "create_scrape_webpage_tool",

surfsense_backend/app/agents/new_chat/tools/display_image.py

@@ -1,111 +0,0 @@
-"""
-Display image tool for the SurfSense agent.
-
-This module provides a tool for displaying images in the chat UI
-with metadata like title, description, and source attribution.
-"""
-
-import hashlib
-from typing import Any
-from urllib.parse import urlparse
-
-from langchain_core.tools import tool
-
-
-def extract_domain(url: str) -> str:
-    """Extract the domain from a URL."""
-    try:
-        parsed = urlparse(url)
-        domain = parsed.netloc
-        # Remove 'www.' prefix if present
-        if domain.startswith("www."):
-            domain = domain[4:]
-        return domain
-    except Exception:
-        return ""
-
-
-def generate_image_id(src: str) -> str:
-    """Generate a unique ID for an image."""
-    hash_val = hashlib.md5(src.encode()).hexdigest()[:12]
-    return f"image-{hash_val}"
-
-
-def create_display_image_tool():
-    """
-    Factory function to create the display_image tool.
-
-    Returns:
-        A configured tool function for displaying images.
-    """
-
-    @tool
-    async def display_image(
-        src: str,
-        alt: str = "Image",
-        title: str | None = None,
-        description: str | None = None,
-    ) -> dict[str, Any]:
-        """
-        Display an image in the chat with metadata.
-
-        Use this tool when you want to show an image to the user.
-        This displays the image with an optional title, description,
-        and source attribution.
-
-        Common use cases:
-        - Showing an image from a URL the user mentioned
-        - Displaying a diagram or chart you're referencing
-        - Showing example images when explaining concepts
-
-        Args:
-            src: The URL of the image to display (must be a valid HTTP/HTTPS URL)
-            alt: Alternative text describing the image (for accessibility)
-            title: Optional title to display below the image
-            description: Optional description providing context about the image
-
-        Returns:
-            A dictionary containing image metadata for the UI to render:
-            - id: Unique identifier for this image
-            - assetId: The image URL (for deduplication)
-            - src: The image URL
-            - alt: Alt text for accessibility
-            - title: Image title (if provided)
-            - description: Image description (if provided)
-            - domain: Source domain
-        """
-        image_id = generate_image_id(src)
-
-        # Ensure URL has protocol
-        if not src.startswith(("http://", "https://")):
-            src = f"https://{src}"
-
-        domain = extract_domain(src)
-
-        # Determine aspect ratio based on image source
-        # AI-generated images should use "auto" to preserve their native ratio
-        is_generated = "/image-generations/" in src
-        if is_generated:
-            ratio = "auto"
-            domain = "ai-generated"
-        elif "unsplash.com" in src or "pexels.com" in src:
-            ratio = "16:9"
-        elif (
-            "imgur.com" in src or "github.com" in src or "githubusercontent.com" in src
-        ):
-            ratio = "auto"
-        else:
-            ratio = "auto"
-
-        return {
-            "id": image_id,
-            "assetId": src,
-            "src": src,
-            "alt": alt,
-            "title": title,
-            "description": description,
-            "domain": domain,
-            "ratio": ratio,
-        }
-
-    return display_image

surfsense_backend/app/agents/new_chat/tools/generate_image.py

@@ -2,8 +2,7 @@
 Image generation tool for the SurfSense agent.
 
 This module provides a tool that generates images using litellm.aimage_generation()
-and returns the result via the existing display_image tool format so the frontend
-renders the generated image inline in the chat.
+and returns the result directly in a format the frontend Image component can render.
 
 Config resolution:
 1. Uses the search space's image_generation_config_id preference
@@ -11,6 +10,7 @@ Config resolution:
 3. Supports global YAML configs (negative IDs) and user DB configs (positive IDs)
 """
 
+import hashlib
 import logging
 from typing import Any
 
@@ -222,11 +222,17 @@ def create_generate_image_tool(
             else:
                 return {"error": "No displayable image data in the response"}
 
+            image_id = f"image-{hashlib.md5(image_url.encode()).hexdigest()[:12]}"
+
             return {
+                "id": image_id,
+                "assetId": image_url,
                 "src": image_url,
                 "alt": revised_prompt or prompt,
                 "title": "Generated Image",
                 "description": revised_prompt if revised_prompt != prompt else None,
+                "domain": "ai-generated",
+                "ratio": "auto",
                 "generated": True,
                 "prompt": prompt,
                 "image_count": len(images),

surfsense_backend/app/agents/new_chat/tools/link_preview.py

@@ -1,465 +0,0 @@
-"""
-Link preview tool for the SurfSense agent.
-
-This module provides a tool for fetching URL metadata (title, description,
-Open Graph image, etc.) to display rich link previews in the chat UI.
-"""
-
-import asyncio
-import hashlib
-import logging
-import re
-from typing import Any
-from urllib.parse import urlparse
-
-import httpx
-import trafilatura
-from fake_useragent import UserAgent
-from langchain_core.tools import tool
-from playwright.sync_api import sync_playwright
-
-from app.utils.proxy_config import get_playwright_proxy, get_residential_proxy_url
-
-logger = logging.getLogger(__name__)
-
-
-def extract_domain(url: str) -> str:
-    """Extract the domain from a URL."""
-    try:
-        parsed = urlparse(url)
-        domain = parsed.netloc
-        # Remove 'www.' prefix if present
-        if domain.startswith("www."):
-            domain = domain[4:]
-        return domain
-    except Exception:
-        return ""
-
-
-def extract_og_content(html: str, property_name: str) -> str | None:
-    """Extract Open Graph meta content from HTML."""
-    # Try og:property first
-    pattern = rf'<meta[^>]+property=["\']og:{property_name}["\'][^>]+content=["\']([^"\']+)["\']'
-    match = re.search(pattern, html, re.IGNORECASE)
-    if match:
-        return match.group(1)
-
-    # Try content before property
-    pattern = rf'<meta[^>]+content=["\']([^"\']+)["\'][^>]+property=["\']og:{property_name}["\']'
-    match = re.search(pattern, html, re.IGNORECASE)
-    if match:
-        return match.group(1)
-
-    return None
-
-
-def extract_twitter_content(html: str, name: str) -> str | None:
-    """Extract Twitter Card meta content from HTML."""
-    pattern = (
-        rf'<meta[^>]+name=["\']twitter:{name}["\'][^>]+content=["\']([^"\']+)["\']'
-    )
-    match = re.search(pattern, html, re.IGNORECASE)
-    if match:
-        return match.group(1)
-
-    # Try content before name
-    pattern = (
-        rf'<meta[^>]+content=["\']([^"\']+)["\'][^>]+name=["\']twitter:{name}["\']'
-    )
-    match = re.search(pattern, html, re.IGNORECASE)
-    if match:
-        return match.group(1)
-
-    return None
-
-
-def extract_meta_description(html: str) -> str | None:
-    """Extract meta description from HTML."""
-    pattern = r'<meta[^>]+name=["\']description["\'][^>]+content=["\']([^"\']+)["\']'
-    match = re.search(pattern, html, re.IGNORECASE)
-    if match:
-        return match.group(1)
-
-    # Try content before name
-    pattern = r'<meta[^>]+content=["\']([^"\']+)["\'][^>]+name=["\']description["\']'
-    match = re.search(pattern, html, re.IGNORECASE)
-    if match:
-        return match.group(1)
-
-    return None
-
-
-def extract_title(html: str) -> str | None:
-    """Extract title from HTML."""
-    # Try og:title first
-    og_title = extract_og_content(html, "title")
-    if og_title:
-        return og_title
-
-    # Try twitter:title
-    twitter_title = extract_twitter_content(html, "title")
-    if twitter_title:
-        return twitter_title
-
-    # Fall back to <title> tag
-    pattern = r"<title[^>]*>([^<]+)</title>"
-    match = re.search(pattern, html, re.IGNORECASE)
-    if match:
-        return match.group(1).strip()
-
-    return None
-
-
-def extract_description(html: str) -> str | None:
-    """Extract description from HTML."""
-    # Try og:description first
-    og_desc = extract_og_content(html, "description")
-    if og_desc:
-        return og_desc
-
-    # Try twitter:description
-    twitter_desc = extract_twitter_content(html, "description")
-    if twitter_desc:
-        return twitter_desc
-
-    # Fall back to meta description
-    return extract_meta_description(html)
-
-
-def extract_image(html: str) -> str | None:
-    """Extract image URL from HTML."""
-    # Try og:image first
-    og_image = extract_og_content(html, "image")
-    if og_image:
-        return og_image
-
-    # Try twitter:image
-    twitter_image = extract_twitter_content(html, "image")
-    if twitter_image:
-        return twitter_image
-
-    return None
-
-
-def generate_preview_id(url: str) -> str:
-    """Generate a unique ID for a link preview."""
-    hash_val = hashlib.md5(url.encode()).hexdigest()[:12]
-    return f"link-preview-{hash_val}"
-
-
-def _unescape_html(text: str) -> str:
-    """Unescape common HTML entities."""
-    return (
-        text.replace("&amp;", "&")
-        .replace("&lt;", "<")
-        .replace("&gt;", ">")
-        .replace("&quot;", '"')
-        .replace("&#39;", "'")
-        .replace("&apos;", "'")
-    )
-
-
-def _make_absolute_url(image_url: str, base_url: str) -> str:
-    """Convert a relative image URL to an absolute URL."""
-    if image_url.startswith(("http://", "https://")):
-        return image_url
-    if image_url.startswith("//"):
-        return f"https:{image_url}"
-    if image_url.startswith("/"):
-        parsed = urlparse(base_url)
-        return f"{parsed.scheme}://{parsed.netloc}{image_url}"
-    return image_url
-
-
-async def fetch_with_chromium(url: str) -> dict[str, Any] | None:
-    """
-    Fetch page content using headless Chromium browser via Playwright.
-    Used as a fallback when simple HTTP requests are blocked (403, etc.).
-
-    Runs the sync Playwright API in a thread so it works on any event
-    loop, including Windows ``SelectorEventLoop``.
-
-    Args:
-        url: URL to fetch
-
-    Returns:
-        Dict with title, description, image, and raw_html, or None if failed
-    """
-    try:
-        return await asyncio.to_thread(_fetch_with_chromium_sync, url)
-    except Exception as e:
-        logger.error(f"[link_preview] Chromium fallback failed for {url}: {e}")
-        return None
-
-
-def _fetch_with_chromium_sync(url: str) -> dict[str, Any] | None:
-    """Synchronous Playwright fetch executed in a worker thread."""
-    logger.info(f"[link_preview] Falling back to Chromium for {url}")
-
-    ua = UserAgent()
-    user_agent = ua.random
-
-    playwright_proxy = get_playwright_proxy()
-
-    with sync_playwright() as p:
-        launch_kwargs: dict = {"headless": True}
-        if playwright_proxy:
-            launch_kwargs["proxy"] = playwright_proxy
-        browser = p.chromium.launch(**launch_kwargs)
-        context = browser.new_context(user_agent=user_agent)
-        page = context.new_page()
-
-        try:
-            page.goto(url, wait_until="domcontentloaded", timeout=30000)
-            raw_html = page.content()
-        finally:
-            browser.close()
-
-    if not raw_html or len(raw_html.strip()) == 0:
-        logger.warning(f"[link_preview] Chromium returned empty content for {url}")
-        return None
-
-    trafilatura_metadata = trafilatura.extract_metadata(raw_html)
-
-    image = extract_image(raw_html)
-
-    result: dict[str, Any] = {
-        "title": None,
-        "description": None,
-        "image": image,
-        "raw_html": raw_html,
-    }
-
-    if trafilatura_metadata:
-        result["title"] = trafilatura_metadata.title
-        result["description"] = trafilatura_metadata.description
-
-    if not result["title"]:
-        result["title"] = extract_title(raw_html)
-    if not result["description"]:
-        result["description"] = extract_description(raw_html)
-
-    logger.info(f"[link_preview] Successfully fetched {url} via Chromium")
-    return result
-
-
-def create_link_preview_tool():
-    """
-    Factory function to create the link_preview tool.
-
-    Returns:
-        A configured tool function for fetching link previews.
-    """
-
-    @tool
-    async def link_preview(url: str) -> dict[str, Any]:
-        """
-        Fetch metadata for a URL to display a rich link preview.
-
-        Use this tool when the user shares a URL or asks about a specific webpage.
-        This tool fetches the page's Open Graph metadata (title, description, image)
-        to display a nice preview card in the chat.
-
-        Common triggers include:
-        - User shares a URL in the chat
-        - User asks "What's this link about?" or similar
-        - User says "Show me a preview of this page"
-        - User wants to preview an article or webpage
-
-        Args:
-            url: The URL to fetch metadata for. Must be a valid HTTP/HTTPS URL.
-
-        Returns:
-            A dictionary containing:
-            - id: Unique identifier for this preview
-            - assetId: The URL itself (for deduplication)
-            - kind: "link" (type of media card)
-            - href: The URL to open when clicked
-            - title: Page title
-            - description: Page description (if available)
-            - thumb: Thumbnail/preview image URL (if available)
-            - domain: The domain name
-            - error: Error message (if fetch failed)
-        """
-        preview_id = generate_preview_id(url)
-        domain = extract_domain(url)
-
-        # Validate URL
-        if not url.startswith(("http://", "https://")):
-            url = f"https://{url}"
-
-        try:
-            # Generate a random User-Agent to avoid bot detection
-            ua = UserAgent()
-            user_agent = ua.random
-
-            # Use residential proxy if configured
-            proxy_url = get_residential_proxy_url()
-
-            # Use a browser-like User-Agent to fetch Open Graph metadata.
-            # We're only fetching publicly available metadata (title, description, thumbnail)
-            # that websites intentionally expose via OG tags for link preview purposes.
-            async with httpx.AsyncClient(
-                timeout=10.0,
-                follow_redirects=True,
-                proxy=proxy_url,
-                headers={
-                    "User-Agent": user_agent,
-                    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8",
-                    "Accept-Language": "en-US,en;q=0.9",
-                    "Accept-Encoding": "gzip, deflate, br",
-                    "Cache-Control": "no-cache",
-                    "Pragma": "no-cache",
-                },
-            ) as client:
-                response = await client.get(url)
-                response.raise_for_status()
-
-                # Get content type to ensure it's HTML
-                content_type = response.headers.get("content-type", "")
-                if "text/html" not in content_type.lower():
-                    # Not an HTML page, return basic info
-                    return {
-                        "id": preview_id,
-                        "assetId": url,
-                        "kind": "link",
-                        "href": url,
-                        "title": url.split("/")[-1] or domain,
-                        "description": f"File from {domain}",
-                        "domain": domain,
-                    }
-
-                html = response.text
-
-                # Extract metadata
-                title = extract_title(html) or domain
-                description = extract_description(html)
-                image = extract_image(html)
-
-                # Make sure image URL is absolute
-                if image:
-                    image = _make_absolute_url(image, url)
-
-                # Clean up title and description (unescape HTML entities)
-                if title:
-                    title = _unescape_html(title)
-                if description:
-                    description = _unescape_html(description)
-                    # Truncate long descriptions
-                    if len(description) > 200:
-                        description = description[:197] + "..."
-
-                return {
-                    "id": preview_id,
-                    "assetId": url,
-                    "kind": "link",
-                    "href": url,
-                    "title": title,
-                    "description": description,
-                    "thumb": image,
-                    "domain": domain,
-                }
-
-        except httpx.TimeoutException:
-            # Timeout - try Chromium fallback
-            logger.warning(
-                f"[link_preview] Timeout for {url}, trying Chromium fallback"
-            )
-            chromium_result = await fetch_with_chromium(url)
-            if chromium_result:
-                title = chromium_result.get("title") or domain
-                description = chromium_result.get("description")
-                image = chromium_result.get("image")
-
-                # Clean up and truncate
-                if title:
-                    title = _unescape_html(title)
-                if description:
-                    description = _unescape_html(description)
-                    if len(description) > 200:
-                        description = description[:197] + "..."
-
-                # Make sure image URL is absolute
-                if image:
-                    image = _make_absolute_url(image, url)
-
-                return {
-                    "id": preview_id,
-                    "assetId": url,
-                    "kind": "link",
-                    "href": url,
-                    "title": title,
-                    "description": description,
-                    "thumb": image,
-                    "domain": domain,
-                }
-
-            return {
-                "id": preview_id,
-                "assetId": url,
-                "kind": "link",
-                "href": url,
-                "title": domain or "Link",
-                "domain": domain,
-                "error": "Request timed out",
-            }
-        except httpx.HTTPStatusError as e:
-            status_code = e.response.status_code
-
-            # For 403 (Forbidden) and similar bot-detection errors, try Chromium fallback
-            if status_code in (403, 401, 406, 429):
-                logger.warning(
-                    f"[link_preview] HTTP {status_code} for {url}, trying Chromium fallback"
-                )
-                chromium_result = await fetch_with_chromium(url)
-                if chromium_result:
-                    title = chromium_result.get("title") or domain
-                    description = chromium_result.get("description")
-                    image = chromium_result.get("image")
-
-                    # Clean up and truncate
-                    if title:
-                        title = _unescape_html(title)
-                    if description:
-                        description = _unescape_html(description)
-                        if len(description) > 200:
-                            description = description[:197] + "..."
-
-                    # Make sure image URL is absolute
-                    if image:
-                        image = _make_absolute_url(image, url)
-
-                    return {
-                        "id": preview_id,
-                        "assetId": url,
-                        "kind": "link",
-                        "href": url,
-                        "title": title,
-                        "description": description,
-                        "thumb": image,
-                        "domain": domain,
-                    }
-
-            return {
-                "id": preview_id,
-                "assetId": url,
-                "kind": "link",
-                "href": url,
-                "title": domain or "Link",
-                "domain": domain,
-                "error": f"HTTP {status_code}",
-            }
-        except Exception as e:
-            error_message = str(e)
-            logger.error(f"[link_preview] Error fetching {url}: {error_message}")
-            return {
-                "id": preview_id,
-                "assetId": url,
-                "kind": "link",
-                "href": url,
-                "title": domain or "Link",
-                "domain": domain,
-                "error": f"Failed to fetch: {error_message[:50]}",
-            }
-
-    return link_preview

surfsense_backend/app/agents/new_chat/tools/registry.py

@@ -50,7 +50,6 @@ from .confluence import (
     create_delete_confluence_page_tool,
     create_update_confluence_page_tool,
 )
-from .display_image import create_display_image_tool
 from .generate_image import create_generate_image_tool
 from .gmail import (
     create_create_gmail_draft_tool,
@@ -78,7 +77,6 @@ from .linear import (
     create_delete_linear_issue_tool,
     create_update_linear_issue_tool,
 )
-from .link_preview import create_link_preview_tool
 from .mcp_tool import load_mcp_tools
 from .notion import (
     create_create_notion_page_tool,
@@ -187,20 +185,6 @@ BUILTIN_TOOLS: list[ToolDefinition] = [
         # are optional — when missing, source_strategy="kb_search" degrades
         # gracefully to "provided"
     ),
-    # Link preview tool - fetches Open Graph metadata for URLs
-    ToolDefinition(
-        name="link_preview",
-        description="Fetch metadata for a URL to display a rich preview card",
-        factory=lambda deps: create_link_preview_tool(),
-        requires=[],
-    ),
-    # Display image tool - shows images in the chat
-    ToolDefinition(
-        name="display_image",
-        description="Display an image in the chat with metadata",
-        factory=lambda deps: create_display_image_tool(),
-        requires=[],
-    ),
     # Generate image tool - creates images using AI models (DALL-E, GPT Image, etc.)
     ToolDefinition(
         name="generate_image",
@@ -567,7 +551,7 @@ def build_tools(
         tools = build_tools(deps)
 
         # Use only specific tools
-        tools = build_tools(deps, enabled_tools=["search_knowledge_base", "link_preview"])
+        tools = build_tools(deps, enabled_tools=["search_knowledge_base"])
 
         # Use defaults but disable podcast
         tools = build_tools(deps, disabled_tools=["generate_podcast"])

surfsense_backend/app/services/public_chat_service.py

@@ -38,13 +38,10 @@ from app.db import (
 from app.utils.rbac import check_permission
 
 UI_TOOLS = {
-    "display_image",
-    "link_preview",
+    "generate_image",
     "generate_podcast",
     "generate_report",
     "generate_video_presentation",
-    "scrape_webpage",
-    "multi_link_preview",
 }
 
 

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -335,38 +335,19 @@ async def _stream_agent_events(
                     status="in_progress",
                     items=last_active_step_items,
                 )
-            elif tool_name == "link_preview":
-                url = (
-                    tool_input.get("url", "")
+            elif tool_name == "generate_image":
+                prompt = (
+                    tool_input.get("prompt", "")
                     if isinstance(tool_input, dict)
                     else str(tool_input)
                 )
-                last_active_step_title = "Fetching link preview"
+                last_active_step_title = "Generating image"
                 last_active_step_items = [
-                    f"URL: {url[:80]}{'...' if len(url) > 80 else ''}"
+                    f"Prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}"
                 ]
                 yield streaming_service.format_thinking_step(
                     step_id=tool_step_id,
-                    title="Fetching link preview",
-                    status="in_progress",
-                    items=last_active_step_items,
-                )
-            elif tool_name == "display_image":
-                src = (
-                    tool_input.get("src", "")
-                    if isinstance(tool_input, dict)
-                    else str(tool_input)
-                )
-                title = (
-                    tool_input.get("title", "") if isinstance(tool_input, dict) else ""
-                )
-                last_active_step_title = "Analyzing the image"
-                last_active_step_items = [
-                    f"Analyzing: {title[:50] if title else src[:50]}{'...' if len(title or src) > 50 else ''}"
-                ]
-                yield streaming_service.format_thinking_step(
-                    step_id=tool_step_id,
-                    title="Analyzing the image",
+                    title="Generating image",
                     status="in_progress",
                     items=last_active_step_items,
                 )
@@ -507,44 +488,22 @@ async def _stream_agent_events(
                     status="completed",
                     items=completed_items,
                 )
-            elif tool_name == "link_preview":
-                if isinstance(tool_output, dict):
-                    title = tool_output.get("title", "Link")
-                    domain = tool_output.get("domain", "")
-                    has_error = "error" in tool_output
-                    if has_error:
-                        completed_items = [
-                            *last_active_step_items,
-                            f"Error: {tool_output.get('error', 'Failed to fetch')}",
-                        ]
-                    else:
-                        completed_items = [
-                            *last_active_step_items,
-                            f"Title: {title[:60]}{'...' if len(title) > 60 else ''}",
-                            f"Domain: {domain}" if domain else "Preview loaded",
-                        ]
-                else:
-                    completed_items = [*last_active_step_items, "Preview loaded"]
-                yield streaming_service.format_thinking_step(
-                    step_id=original_step_id,
-                    title="Fetching link preview",
-                    status="completed",
-                    items=completed_items,
-                )
-            elif tool_name == "display_image":
-                if isinstance(tool_output, dict):
-                    title = tool_output.get("title", "")
-                    alt = tool_output.get("alt", "Image")
-                    display_name = title or alt
+            elif tool_name == "generate_image":
+                if isinstance(tool_output, dict) and not tool_output.get("error"):
                     completed_items = [
                         *last_active_step_items,
-                        f"Analyzed: {display_name[:50]}{'...' if len(display_name) > 50 else ''}",
+                        "Image generated successfully",
                     ]
                 else:
-                    completed_items = [*last_active_step_items, "Image analyzed"]
+                    error_msg = (
+                        tool_output.get("error", "Generation failed")
+                        if isinstance(tool_output, dict)
+                        else "Generation failed"
+                    )
+                    completed_items = [*last_active_step_items, f"Error: {error_msg}"]
                 yield streaming_service.format_thinking_step(
                     step_id=original_step_id,
-                    title="Analyzing the image",
+                    title="Generating image",
                     status="completed",
                     items=completed_items,
                 )
@@ -819,30 +778,7 @@ async def _stream_agent_events(
                         f"Presentation generation failed: {error_msg}",
                         "error",
                     )
-            elif tool_name == "link_preview":
-                yield streaming_service.format_tool_output_available(
-                    tool_call_id,
-                    tool_output
-                    if isinstance(tool_output, dict)
-                    else {"result": tool_output},
-                )
-                if isinstance(tool_output, dict) and "error" not in tool_output:
-                    title = tool_output.get("title", "Link")
-                    yield streaming_service.format_terminal_info(
-                        f"Link preview loaded: {title[:50]}{'...' if len(title) > 50 else ''}",
-                        "success",
-                    )
-                else:
-                    error_msg = (
-                        tool_output.get("error", "Failed to fetch")
-                        if isinstance(tool_output, dict)
-                        else "Failed to fetch"
-                    )
-                    yield streaming_service.format_terminal_info(
-                        f"Link preview failed: {error_msg}",
-                        "error",
-                    )
-            elif tool_name == "display_image":
+            elif tool_name == "generate_image":
                 yield streaming_service.format_tool_output_available(
                     tool_call_id,
                     tool_output
@@ -850,11 +786,16 @@ async def _stream_agent_events(
                     else {"result": tool_output},
                 )
                 if isinstance(tool_output, dict):
-                    title = tool_output.get("title") or tool_output.get("alt", "Image")
-                    yield streaming_service.format_terminal_info(
-                        f"Image analyzed: {title[:40]}{'...' if len(title) > 40 else ''}",
-                        "success",
-                    )
+                    if tool_output.get("error"):
+                        yield streaming_service.format_terminal_info(
+                            f"Image generation failed: {tool_output['error'][:60]}",
+                            "error",
+                        )
+                    else:
+                        yield streaming_service.format_terminal_info(
+                            "Image generated successfully",
+                            "success",
+                        )
             elif tool_name == "scrape_webpage":
                 if isinstance(tool_output, dict):
                     display_output = {

surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx

@@ -33,59 +33,15 @@ import { closeReportPanelAtom } from "@/atoms/chat/report-panel.atom";
 import { closeEditorPanelAtom } from "@/atoms/editor/editor-panel.atom";
 import { membersAtom } from "@/atoms/members/members-query.atoms";
 import { currentUserAtom } from "@/atoms/user/user-query.atoms";
+import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps";
 import { Thread } from "@/components/assistant-ui/thread";
 import { MobileEditorPanel } from "@/components/editor-panel/editor-panel";
 import { MobileHitlEditPanel } from "@/components/hitl-edit-panel/hitl-edit-panel";
 import { MobileReportPanel } from "@/components/report-panel/report-panel";
-import {
-	CreateConfluencePageToolUI,
-	DeleteConfluencePageToolUI,
-	UpdateConfluencePageToolUI,
-} from "@/components/tool-ui/confluence";
-import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
-import { DisplayImageToolUI } from "@/components/tool-ui/display-image";
-import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
-import { GenerateReportToolUI } from "@/components/tool-ui/generate-report";
-import {
-	CreateGmailDraftToolUI,
-	SendGmailEmailToolUI,
-	TrashGmailEmailToolUI,
-	UpdateGmailDraftToolUI,
-} from "@/components/tool-ui/gmail";
-import {
-	CreateCalendarEventToolUI,
-	DeleteCalendarEventToolUI,
-	UpdateCalendarEventToolUI,
-} from "@/components/tool-ui/google-calendar";
-import {
-	CreateGoogleDriveFileToolUI,
-	DeleteGoogleDriveFileToolUI,
-} from "@/components/tool-ui/google-drive";
-import {
-	CreateJiraIssueToolUI,
-	DeleteJiraIssueToolUI,
-	UpdateJiraIssueToolUI,
-} from "@/components/tool-ui/jira";
-import {
-	CreateLinearIssueToolUI,
-	DeleteLinearIssueToolUI,
-	UpdateLinearIssueToolUI,
-} from "@/components/tool-ui/linear";
-import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview";
-import {
-	CreateNotionPageToolUI,
-	DeleteNotionPageToolUI,
-	UpdateNotionPageToolUI,
-} from "@/components/tool-ui/notion";
-import { SandboxExecuteToolUI } from "@/components/tool-ui/sandbox-execute";
-import { ScrapeWebpageToolUI } from "@/components/tool-ui/scrape-webpage";
-import { RecallMemoryToolUI, SaveMemoryToolUI } from "@/components/tool-ui/user-memory";
-import { GenerateVideoPresentationToolUI } from "@/components/tool-ui/video-presentation";
 import { Skeleton } from "@/components/ui/skeleton";
 import { useChatSessionStateSync } from "@/hooks/use-chat-session-state";
 import { useMessagesSync } from "@/hooks/use-messages-sync";
 import { documentsApiService } from "@/lib/apis/documents-api.service";
-// import { WriteTodosToolUI } from "@/components/tool-ui/write-todos";
 import { getBearerToken } from "@/lib/auth-utils";
 import { convertToThreadMessage } from "@/lib/chat/message-utils";
 import {
@@ -101,6 +57,7 @@ import {
 	type ContentPartsState,
 	readSSEStream,
 	type ThinkingStepData,
+	updateThinkingSteps,
 	updateToolCall,
 } from "@/lib/chat/streaming-state";
 import {
@@ -137,23 +94,6 @@ function markInterruptsCompleted(contentParts: Array<{ type: string; result?: un
 	}
 }
 
-/**
- * Extract thinking steps from message content
- */
-function extractThinkingSteps(content: unknown): ThinkingStep[] {
-	if (!Array.isArray(content)) return [];
-
-	const thinkingPart = content.find(
-		(part: unknown) =>
-			typeof part === "object" &&
-			part !== null &&
-			"type" in part &&
-			(part as { type: string }).type === "thinking-steps"
-	) as { type: "thinking-steps"; steps: ThinkingStep[] } | undefined;
-
-	return thinkingPart?.steps || [];
-}
-
 /**
  * Zod schema for mentioned document info (for type-safe parsing)
  */
@@ -191,10 +131,9 @@ const TOOLS_WITH_UI = new Set([
 	"generate_podcast",
 	"generate_report",
 	"generate_video_presentation",
-	"link_preview",
 	"display_image",
+	"generate_image",
 	"delete_notion_page",
-	"scrape_webpage",
 	"create_notion_page",
 	"update_notion_page",
 	"create_linear_issue",
@@ -227,11 +166,6 @@ export default function NewChatPage() {
 	const [currentThread, setCurrentThread] = useState<ThreadRecord | null>(null);
 	const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
 	const [isRunning, setIsRunning] = useState(false);
-	// Store thinking steps per message ID - kept separate from content to avoid
-	// "unsupported part type" errors from assistant-ui
-	const [messageThinkingSteps, setMessageThinkingSteps] = useState<Map<string, ThinkingStep[]>>(
-		new Map()
-	);
 	const abortControllerRef = useRef<AbortController | null>(null);
 	const [pendingInterrupt, setPendingInterrupt] = useState<{
 		threadId: number;
@@ -332,6 +266,7 @@ export default function NewChatPage() {
 
 	// Initialize thread and load messages
 	// For new chats (no urlChatId), we use lazy creation - thread is created on first message
+	// biome-ignore lint/correctness/useExhaustiveDependencies: searchSpaceId triggers re-init when switching spaces with the same urlChatId
 	const initializeThread = useCallback(async () => {
 		setIsInitializing(true);
 
@@ -339,7 +274,6 @@ export default function NewChatPage() {
 		setMessages([]);
 		setThreadId(null);
 		setCurrentThread(null);
-		setMessageThinkingSteps(new Map());
 		setMentionedDocuments([]);
 		setSidebarDocuments([]);
 		setMessageDocumentsMap({});
@@ -364,18 +298,8 @@ export default function NewChatPage() {
 					const loadedMessages = messagesResponse.messages.map(convertToThreadMessage);
 					setMessages(loadedMessages);
 
-					// Extract and restore thinking steps from persisted messages
-					const restoredThinkingSteps = new Map<string, ThinkingStep[]>();
-					// Extract and restore mentioned documents from persisted messages
 					const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
-
 					for (const msg of messagesResponse.messages) {
-						if (msg.role === "assistant") {
-							const steps = extractThinkingSteps(msg.content);
-							if (steps.length > 0) {
-								restoredThinkingSteps.set(`msg-${msg.id}`, steps);
-							}
-						}
 						if (msg.role === "user") {
 							const docs = extractMentionedDocuments(msg.content);
 							if (docs.length > 0) {
@@ -383,9 +307,6 @@ export default function NewChatPage() {
 							}
 						}
 					}
-					if (restoredThinkingSteps.size > 0) {
-						setMessageThinkingSteps(restoredThinkingSteps);
-					}
 					if (Object.keys(restoredDocsMap).length > 0) {
 						setMessageDocumentsMap(restoredDocsMap);
 					}
@@ -789,18 +710,17 @@ export default function NewChatPage() {
 						}
 
 						case "data-thinking-step": {
-							// Handle thinking step events for chain-of-thought display
 							const stepData = parsed.data as ThinkingStepData;
 							if (stepData?.id) {
 								currentThinkingSteps.set(stepData.id, stepData);
-								// Update thinking steps state for rendering
-								// The ThinkingStepsScrollHandler in Thread component
-								// will handle auto-scrolling when this state changes
-								setMessageThinkingSteps((prev) => {
-									const newMap = new Map(prev);
-									newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values()));
-									return newMap;
-								});
+								updateThinkingSteps(contentPartsState, currentThinkingSteps);
+								setMessages((prev) =>
+									prev.map((m) =>
+										m.id === assistantMsgId
+											? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
+											: m
+									)
+								);
 							}
 							break;
 						}
@@ -865,13 +785,8 @@ export default function NewChatPage() {
 					}
 				}
 
-				// Persist assistant message (with thinking steps for restoration on refresh)
 				// Skip persistence for interrupted messages -- handleResume will persist the final version
-				const finalContent = buildContentForPersistence(
-					contentPartsState,
-					TOOLS_WITH_UI,
-					currentThinkingSteps
-				);
+				const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
 				if (contentParts.length > 0 && !wasInterrupted) {
 					try {
 						const savedMessage = await appendMessage(currentThreadId, {
@@ -891,18 +806,6 @@ export default function NewChatPage() {
 								? { ...prev, assistantMsgId: newMsgId }
 								: prev
 						);
-
-						// Also update thinking steps map with new ID
-						setMessageThinkingSteps((prev) => {
-							const steps = prev.get(assistantMsgId);
-							if (steps) {
-								const newMap = new Map(prev);
-								newMap.delete(assistantMsgId);
-								newMap.set(newMsgId, steps);
-								return newMap;
-							}
-							return prev;
-						});
 					} catch (err) {
 						console.error("Failed to persist assistant message:", err);
 					}
@@ -919,11 +822,7 @@ export default function NewChatPage() {
 							(part.type === "tool-call" && TOOLS_WITH_UI.has(part.toolName))
 					);
 					if (hasContent && currentThreadId) {
-						const partialContent = buildContentForPersistence(
-							contentPartsState,
-							TOOLS_WITH_UI,
-							currentThinkingSteps
-						);
+						const partialContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
 						try {
 							const savedMessage = await appendMessage(currentThreadId, {
 								role: "assistant",
@@ -970,7 +869,6 @@ export default function NewChatPage() {
 			} finally {
 				setIsRunning(false);
 				abortControllerRef.current = null;
-				// Note: We no longer clear thinking steps - they persist with the message
 			}
 		},
 		[
@@ -1013,9 +911,7 @@ export default function NewChatPage() {
 			const controller = new AbortController();
 			abortControllerRef.current = controller;
 
-			const currentThinkingSteps = new Map<string, ThinkingStepData>(
-				(messageThinkingSteps.get(assistantMsgId) ?? []).map((s) => [s.id, s])
-			);
+			const currentThinkingSteps = new Map<string, ThinkingStepData>();
 
 			const contentPartsState: ContentPartsState = {
 				contentParts: [],
@@ -1042,6 +938,15 @@ export default function NewChatPage() {
 								result: p.result as unknown,
 							});
 							contentPartsState.currentTextPartIndex = -1;
+						} else if (p.type === "data-thinking-steps") {
+							const stepsData = p.data as { steps: ThinkingStepData[] } | undefined;
+							contentParts.push({
+								type: "data-thinking-steps",
+								data: { steps: stepsData?.steps ?? [] },
+							});
+							for (const step of stepsData?.steps ?? []) {
+								currentThinkingSteps.set(step.id, step);
+							}
 						}
 					}
 				}
@@ -1159,11 +1064,14 @@ export default function NewChatPage() {
 							const stepData = parsed.data as ThinkingStepData;
 							if (stepData?.id) {
 								currentThinkingSteps.set(stepData.id, stepData);
-								setMessageThinkingSteps((prev) => {
-									const newMap = new Map(prev);
-									newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values()));
-									return newMap;
-								});
+								updateThinkingSteps(contentPartsState, currentThinkingSteps);
+								setMessages((prev) =>
+									prev.map((m) =>
+										m.id === assistantMsgId
+											? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
+											: m
+									)
+								);
 							}
 							break;
 						}
@@ -1217,11 +1125,7 @@ export default function NewChatPage() {
 					}
 				}
 
-				const finalContent = buildContentForPersistence(
-					contentPartsState,
-					TOOLS_WITH_UI,
-					currentThinkingSteps
-				);
+				const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
 				if (contentParts.length > 0) {
 					try {
 						const savedMessage = await appendMessage(resumeThreadId, {
@@ -1232,16 +1136,6 @@ export default function NewChatPage() {
 						setMessages((prev) =>
 							prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m))
 						);
-						setMessageThinkingSteps((prev) => {
-							const steps = prev.get(assistantMsgId);
-							if (steps) {
-								const newMap = new Map(prev);
-								newMap.delete(assistantMsgId);
-								newMap.set(newMsgId, steps);
-								return newMap;
-							}
-							return prev;
-						});
 					} catch (err) {
 						console.error("Failed to persist resumed assistant message:", err);
 					}
@@ -1257,7 +1151,7 @@ export default function NewChatPage() {
 				abortControllerRef.current = null;
 			}
 		},
-		[pendingInterrupt, messages, searchSpaceId, messageThinkingSteps]
+		[pendingInterrupt, messages, searchSpaceId]
 	);
 
 	useEffect(() => {
@@ -1376,20 +1270,6 @@ export default function NewChatPage() {
 				return prev;
 			});
 
-			// Clear thinking steps for the removed messages
-			setMessageThinkingSteps((prev) => {
-				const newMap = new Map(prev);
-				// Remove thinking steps for the last two messages
-				const lastTwoIds = messages
-					.slice(-2)
-					.map((m) => m.id)
-					.filter((id): id is string => !!id);
-				for (const id of lastTwoIds) {
-					newMap.delete(id);
-				}
-				return newMap;
-			});
-
 			// Start streaming
 			setIsRunning(true);
 			const controller = new AbortController();
@@ -1520,11 +1400,14 @@ export default function NewChatPage() {
 							const stepData = parsed.data as ThinkingStepData;
 							if (stepData?.id) {
 								currentThinkingSteps.set(stepData.id, stepData);
-								setMessageThinkingSteps((prev) => {
-									const newMap = new Map(prev);
-									newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values()));
-									return newMap;
-								});
+								updateThinkingSteps(contentPartsState, currentThinkingSteps);
+								setMessages((prev) =>
+									prev.map((m) =>
+										m.id === assistantMsgId
+											? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
+											: m
+									)
+								);
 							}
 							break;
 						}
@@ -1535,11 +1418,7 @@ export default function NewChatPage() {
 				}
 
 				// Persist messages after streaming completes
-				const finalContent = buildContentForPersistence(
-					contentPartsState,
-					TOOLS_WITH_UI,
-					currentThinkingSteps
-				);
+				const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
 				if (contentParts.length > 0) {
 					try {
 						// Persist user message (for both edit and reload modes, since backend deleted it)
@@ -1570,18 +1449,6 @@ export default function NewChatPage() {
 							prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m))
 						);
 
-						setMessageThinkingSteps((prev) => {
-							const steps = prev.get(assistantMsgId);
-							if (steps) {
-								const newMap = new Map(prev);
-								newMap.delete(assistantMsgId);
-								newMap.set(newMsgId, steps);
-								return newMap;
-							}
-							return prev;
-						});
-
-						// Track successful response
 						trackChatResponseReceived(searchSpaceId, threadId);
 					} catch (err) {
 						console.error("Failed to persist regenerated message:", err);
@@ -1614,7 +1481,7 @@ export default function NewChatPage() {
 				abortControllerRef.current = null;
 			}
 		},
-		[threadId, searchSpaceId, messages, setMessageThinkingSteps, disabledTools]
+		[threadId, searchSpaceId, messages, disabledTools]
 	);
 
 	// Handle editing a message - truncates history and regenerates with new query
@@ -1719,40 +1586,10 @@ export default function NewChatPage() {
 
 	return (
 		<AssistantRuntimeProvider runtime={runtime}>
-			<GeneratePodcastToolUI />
-			<GenerateReportToolUI />
-			<GenerateVideoPresentationToolUI />
-			<LinkPreviewToolUI />
-			<DisplayImageToolUI />
-			<ScrapeWebpageToolUI />
-			<SaveMemoryToolUI />
-			<RecallMemoryToolUI />
-			<CreateNotionPageToolUI />
-			<UpdateNotionPageToolUI />
-			<DeleteNotionPageToolUI />
-			<CreateLinearIssueToolUI />
-			<UpdateLinearIssueToolUI />
-			<DeleteLinearIssueToolUI />
-			<CreateGoogleDriveFileToolUI />
-			<DeleteGoogleDriveFileToolUI />
-			<CreateCalendarEventToolUI />
-			<UpdateCalendarEventToolUI />
-			<DeleteCalendarEventToolUI />
-			<CreateGmailDraftToolUI />
-			<UpdateGmailDraftToolUI />
-			<SendGmailEmailToolUI />
-			<TrashGmailEmailToolUI />
-			<CreateJiraIssueToolUI />
-			<UpdateJiraIssueToolUI />
-			<DeleteJiraIssueToolUI />
-			<CreateConfluencePageToolUI />
-			<UpdateConfluencePageToolUI />
-			<DeleteConfluencePageToolUI />
-			<SandboxExecuteToolUI />
-			{/* <WriteTodosToolUI /> Disabled for now */}
+			<ThinkingStepsDataUI />
 			<div key={searchSpaceId} className="flex h-[calc(100dvh-64px)] overflow-hidden">
 				<div className="flex-1 flex flex-col min-w-0 overflow-hidden">
-					<Thread messageThinkingSteps={messageThinkingSteps} />
+					<Thread />
 				</div>
 				<MobileReportPanel />
 				<MobileEditorPanel />

surfsense_web/components/assistant-ui/assistant-message.tsx

@@ -1,26 +1,62 @@
 import {
 	ActionBarPrimitive,
-	AssistantIf,
+	AuiIf,
 	ErrorPrimitive,
 	MessagePrimitive,
-	useAssistantState,
-	useMessage,
+	useAuiState,
 } from "@assistant-ui/react";
 import { useAtomValue } from "jotai";
 import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react";
 import type { FC } from "react";
-import { useContext, useEffect, useMemo, useRef, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import { commentsEnabledAtom, targetCommentIdAtom } from "@/atoms/chat/current-thread.atom";
 import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
-import {
-	ThinkingStepsContext,
-	ThinkingStepsDisplay,
-} from "@/components/assistant-ui/thinking-steps";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container";
 import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet";
+import {
+	CreateConfluencePageToolUI,
+	DeleteConfluencePageToolUI,
+	UpdateConfluencePageToolUI,
+} from "@/components/tool-ui/confluence";
+import { GenerateImageToolUI } from "@/components/tool-ui/generate-image";
+import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
+import { GenerateReportToolUI } from "@/components/tool-ui/generate-report";
+import {
+	CreateGmailDraftToolUI,
+	SendGmailEmailToolUI,
+	TrashGmailEmailToolUI,
+	UpdateGmailDraftToolUI,
+} from "@/components/tool-ui/gmail";
+import {
+	CreateCalendarEventToolUI,
+	DeleteCalendarEventToolUI,
+	UpdateCalendarEventToolUI,
+} from "@/components/tool-ui/google-calendar";
+import {
+	CreateGoogleDriveFileToolUI,
+	DeleteGoogleDriveFileToolUI,
+} from "@/components/tool-ui/google-drive";
+import {
+	CreateJiraIssueToolUI,
+	DeleteJiraIssueToolUI,
+	UpdateJiraIssueToolUI,
+} from "@/components/tool-ui/jira";
+import {
+	CreateLinearIssueToolUI,
+	DeleteLinearIssueToolUI,
+	UpdateLinearIssueToolUI,
+} from "@/components/tool-ui/linear";
+import {
+	CreateNotionPageToolUI,
+	DeleteNotionPageToolUI,
+	UpdateNotionPageToolUI,
+} from "@/components/tool-ui/notion";
+import { SandboxExecuteToolUI } from "@/components/tool-ui/sandbox-execute";
+import { RecallMemoryToolUI, SaveMemoryToolUI } from "@/components/tool-ui/user-memory";
+import { GenerateVideoPresentationToolUI } from "@/components/tool-ui/video-presentation";
 import { useComments } from "@/hooks/use-comments";
 import { useMediaQuery } from "@/hooks/use-media-query";
 import { cn } from "@/lib/utils";
@@ -35,42 +71,50 @@ export const MessageError: FC = () => {
 	);
 };
 
-/**
- * Custom component to render thinking steps from Context
- */
-const ThinkingStepsPart: FC = () => {
-	const thinkingStepsMap = useContext(ThinkingStepsContext);
-
-	// Get the current message ID to look up thinking steps
-	const messageId = useAssistantState(({ message }) => message?.id);
-	const thinkingSteps = thinkingStepsMap.get(messageId) || [];
-
-	// Check if this specific message is currently streaming
-	// A message is streaming if: thread is running AND this is the last assistant message
-	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
-	const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false);
-	const isMessageStreaming = isThreadRunning && isLastMessage;
-
-	if (thinkingSteps.length === 0) return null;
-
-	return (
-		<div className="mb-3">
-			<ThinkingStepsDisplay steps={thinkingSteps} isThreadRunning={isMessageStreaming} />
-		</div>
-	);
-};
-
 const AssistantMessageInner: FC = () => {
 	return (
 		<>
-			{/* Render thinking steps from message content - this ensures proper scroll tracking */}
-			<ThinkingStepsPart />
-
 			<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
 				<MessagePrimitive.Parts
 					components={{
 						Text: MarkdownText,
-						tools: { Fallback: ToolFallback },
+						tools: {
+							by_name: {
+								generate_report: GenerateReportToolUI,
+								generate_podcast: GeneratePodcastToolUI,
+								generate_video_presentation: GenerateVideoPresentationToolUI,
+								display_image: GenerateImageToolUI,
+								generate_image: GenerateImageToolUI,
+								save_memory: SaveMemoryToolUI,
+								recall_memory: RecallMemoryToolUI,
+								execute: SandboxExecuteToolUI,
+								create_notion_page: CreateNotionPageToolUI,
+								update_notion_page: UpdateNotionPageToolUI,
+								delete_notion_page: DeleteNotionPageToolUI,
+								create_linear_issue: CreateLinearIssueToolUI,
+								update_linear_issue: UpdateLinearIssueToolUI,
+								delete_linear_issue: DeleteLinearIssueToolUI,
+								create_google_drive_file: CreateGoogleDriveFileToolUI,
+								delete_google_drive_file: DeleteGoogleDriveFileToolUI,
+								create_calendar_event: CreateCalendarEventToolUI,
+								update_calendar_event: UpdateCalendarEventToolUI,
+								delete_calendar_event: DeleteCalendarEventToolUI,
+								create_gmail_draft: CreateGmailDraftToolUI,
+								update_gmail_draft: UpdateGmailDraftToolUI,
+								send_gmail_email: SendGmailEmailToolUI,
+								trash_gmail_email: TrashGmailEmailToolUI,
+								create_jira_issue: CreateJiraIssueToolUI,
+								update_jira_issue: UpdateJiraIssueToolUI,
+								delete_jira_issue: DeleteJiraIssueToolUI,
+								create_confluence_page: CreateConfluencePageToolUI,
+								update_confluence_page: UpdateConfluencePageToolUI,
+								delete_confluence_page: DeleteConfluencePageToolUI,
+								link_preview: () => null,
+								multi_link_preview: () => null,
+								scrape_webpage: () => null,
+							},
+							Fallback: ToolFallback,
+						},
 					}}
 				/>
 				<MessageError />
@@ -95,7 +139,7 @@ export const AssistantMessage: FC = () => {
 	const messageRef = useRef<HTMLDivElement>(null);
 	const commentPanelRef = useRef<HTMLDivElement>(null);
 	const commentTriggerRef = useRef<HTMLButtonElement>(null);
-	const messageId = useAssistantState(({ message }) => message?.id);
+	const messageId = useAuiState(({ message }) => message?.id);
 	const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom);
 	const dbMessageId = parseMessageId(messageId);
 	const commentsEnabled = useAtomValue(commentsEnabledAtom);
@@ -104,8 +148,8 @@ export const AssistantMessage: FC = () => {
 	const isMediumScreen = useMediaQuery("(min-width: 768px) and (max-width: 1023px)");
 	const isDesktop = useMediaQuery("(min-width: 1024px)");
 
-	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
-	const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false);
+	const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
+	const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
 	const isMessageStreaming = isThreadRunning && isLastMessage;
 
 	const { data: commentsData, isSuccess: commentsLoaded } = useComments({
@@ -227,7 +271,7 @@ export const AssistantMessage: FC = () => {
 };
 
 const AssistantActionBar: FC = () => {
-	const { isLast } = useMessage();
+	const isLast = useAuiState((s) => s.message.isLast);
 
 	return (
 		<ActionBarPrimitive.Root
@@ -238,12 +282,12 @@ const AssistantActionBar: FC = () => {
 		>
 			<ActionBarPrimitive.Copy asChild>
 				<TooltipIconButton tooltip="Copy">
-					<AssistantIf condition={({ message }) => message.isCopied}>
+					<AuiIf condition={({ message }) => message.isCopied}>
 						<CheckIcon />
-					</AssistantIf>
-					<AssistantIf condition={({ message }) => !message.isCopied}>
+					</AuiIf>
+					<AuiIf condition={({ message }) => !message.isCopied}>
 						<CopyIcon />
-					</AssistantIf>
+					</AuiIf>
 				</TooltipIconButton>
 			</ActionBarPrimitive.Copy>
 			<ActionBarPrimitive.ExportMarkdown asChild>

surfsense_web/components/assistant-ui/image.tsx

@@ -0,0 +1,227 @@
+"use client";
+
+import type { ImageMessagePartComponent } from "@assistant-ui/react";
+import { cva, type VariantProps } from "class-variance-authority";
+import { ImageIcon, ImageOffIcon } from "lucide-react";
+import { memo, type PropsWithChildren, useEffect, useRef, useState } from "react";
+import { createPortal } from "react-dom";
+import { cn } from "@/lib/utils";
+
+const imageVariants = cva("aui-image-root relative overflow-hidden rounded-lg", {
+	variants: {
+		variant: {
+			outline: "border border-border",
+			ghost: "",
+			muted: "bg-muted/50",
+		},
+		size: {
+			sm: "max-w-64",
+			default: "max-w-96",
+			lg: "max-w-[512px]",
+			full: "w-full",
+		},
+	},
+	defaultVariants: {
+		variant: "outline",
+		size: "default",
+	},
+});
+
+export type ImageRootProps = React.ComponentProps<"div"> & VariantProps<typeof imageVariants>;
+
+function ImageRoot({ className, variant, size, children, ...props }: ImageRootProps) {
+	return (
+		<div
+			data-slot="image-root"
+			data-variant={variant}
+			data-size={size}
+			className={cn(imageVariants({ variant, size, className }))}
+			{...props}
+		>
+			{children}
+		</div>
+	);
+}
+
+type ImagePreviewProps = Omit<React.ComponentProps<"img">, "children"> & {
+	containerClassName?: string;
+};
+
+function ImagePreview({
+	className,
+	containerClassName,
+	onLoad,
+	onError,
+	alt = "Image content",
+	src,
+	...props
+}: ImagePreviewProps) {
+	const imgRef = useRef<HTMLImageElement>(null);
+	const [loadedSrc, setLoadedSrc] = useState<string | undefined>(undefined);
+	const [errorSrc, setErrorSrc] = useState<string | undefined>(undefined);
+
+	const loaded = loadedSrc === src;
+	const error = errorSrc === src;
+
+	useEffect(() => {
+		if (typeof src === "string" && imgRef.current?.complete && imgRef.current.naturalWidth > 0) {
+			setLoadedSrc(src);
+		}
+	}, [src]);
+
+	return (
+		<div data-slot="image-preview" className={cn("relative min-h-32", containerClassName)}>
+			{!loaded && !error && (
+				<div
+					data-slot="image-preview-loading"
+					className="absolute inset-0 flex items-center justify-center bg-muted/50"
+				>
+					<ImageIcon className="size-8 animate-pulse text-muted-foreground" />
+				</div>
+			)}
+			{error ? (
+				<div
+					data-slot="image-preview-error"
+					className="flex min-h-32 items-center justify-center bg-muted/50 p-4"
+				>
+					<ImageOffIcon className="size-8 text-muted-foreground" />
+				</div>
+			) : (
+				// biome-ignore lint/performance/noImgElement: intentional for dynamic external URLs
+				<img
+					ref={imgRef}
+					src={src}
+					alt={alt}
+					className={cn("block h-auto w-full object-contain", !loaded && "invisible", className)}
+					onLoad={(e) => {
+						if (typeof src === "string") setLoadedSrc(src);
+						onLoad?.(e);
+					}}
+					onError={(e) => {
+						if (typeof src === "string") setErrorSrc(src);
+						onError?.(e);
+					}}
+					{...props}
+				/>
+			)}
+		</div>
+	);
+}
+
+function ImageFilename({ className, children, ...props }: React.ComponentProps<"span">) {
+	if (!children) return null;
+
+	return (
+		<span
+			data-slot="image-filename"
+			className={cn("block truncate px-2 py-1.5 text-muted-foreground text-xs", className)}
+			{...props}
+		>
+			{children}
+		</span>
+	);
+}
+
+type ImageZoomProps = PropsWithChildren<{
+	src: string;
+	alt?: string;
+}>;
+
+function ImageZoom({ src, alt = "Image preview", children }: ImageZoomProps) {
+	const [isMounted, setIsMounted] = useState(false);
+	const [isOpen, setIsOpen] = useState(false);
+
+	useEffect(() => {
+		setIsMounted(true);
+	}, []);
+
+	const handleOpen = () => setIsOpen(true);
+	const handleClose = () => setIsOpen(false);
+
+	useEffect(() => {
+		if (!isOpen) return;
+		const handleKeyDown = (e: KeyboardEvent) => {
+			if (e.key === "Escape") setIsOpen(false);
+		};
+		document.addEventListener("keydown", handleKeyDown);
+		return () => document.removeEventListener("keydown", handleKeyDown);
+	}, [isOpen]);
+
+	useEffect(() => {
+		if (!isOpen) return;
+		const originalOverflow = document.body.style.overflow;
+		document.body.style.overflow = "hidden";
+		return () => {
+			document.body.style.overflow = originalOverflow;
+		};
+	}, [isOpen]);
+
+	return (
+		<>
+			<button
+				type="button"
+				onClick={handleOpen}
+				className="aui-image-zoom-trigger cursor-zoom-in border-0 bg-transparent p-0 text-left"
+				aria-label="Click to zoom image"
+			>
+				{children}
+			</button>
+			{isMounted &&
+				isOpen &&
+				createPortal(
+					<button
+						type="button"
+						data-slot="image-zoom-overlay"
+						className="aui-image-zoom-overlay fade-in fixed inset-0 z-50 flex animate-in cursor-zoom-out items-center justify-center border-0 bg-black/80 p-0 duration-200"
+						onClick={handleClose}
+						aria-label="Close zoomed image"
+					>
+						{/** biome-ignore lint/performance/noImgElement: <explanation> */}
+						<img
+							data-slot="image-zoom-content"
+							src={src}
+							alt={alt}
+							className="aui-image-zoom-content fade-in zoom-in-95 max-h-[90vh] max-w-[90vw] animate-in object-contain duration-200"
+							onClick={(e) => {
+								e.stopPropagation();
+								handleClose();
+							}}
+							onKeyDown={(e) => {
+								if (e.key === "Enter") {
+									e.stopPropagation();
+									handleClose();
+								}
+							}}
+						/>
+					</button>,
+					document.body
+				)}
+		</>
+	);
+}
+
+const ImageImpl: ImageMessagePartComponent = ({ image, filename }) => {
+	return (
+		<ImageRoot>
+			<ImageZoom src={image} alt={filename || "Image content"}>
+				<ImagePreview src={image} alt={filename || "Image content"} />
+			</ImageZoom>
+			<ImageFilename>{filename}</ImageFilename>
+		</ImageRoot>
+	);
+};
+
+const Image = memo(ImageImpl) as unknown as ImageMessagePartComponent & {
+	Root: typeof ImageRoot;
+	Preview: typeof ImagePreview;
+	Filename: typeof ImageFilename;
+	Zoom: typeof ImageZoom;
+};
+
+Image.displayName = "Image";
+Image.Root = ImageRoot;
+Image.Preview = ImagePreview;
+Image.Filename = ImageFilename;
+Image.Zoom = ImageZoom;
+
+export { Image, ImageRoot, ImagePreview, ImageFilename, ImageZoom, imageVariants };

surfsense_web/components/assistant-ui/markdown-text.tsx

@@ -3,21 +3,50 @@
 import "@assistant-ui/react-markdown/styles/dot.css";
 
 import {
-	type CodeHeaderProps,
 	MarkdownTextPrimitive,
 	unstable_memoizeMarkdownComponents as memoizeMarkdownComponents,
 	useIsMarkdownCodeBlock,
 } from "@assistant-ui/react-markdown";
-import { CheckIcon, CopyIcon } from "lucide-react";
+import { CheckIcon, CopyIcon, ExternalLinkIcon } from "lucide-react";
+import { useTheme } from "next-themes";
+import type { CSSProperties } from "react";
 import { type FC, memo, type ReactNode, useState } from "react";
+import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
+import { materialDark, materialLight } from "react-syntax-highlighter/dist/esm/styles/prism";
 import rehypeKatex from "rehype-katex";
 import remarkGfm from "remark-gfm";
 import remarkMath from "remark-math";
+import { ImagePreview, ImageRoot, ImageZoom } from "@/components/assistant-ui/image";
 import "katex/dist/katex.min.css";
 import { InlineCitation, UrlCitation } from "@/components/assistant-ui/inline-citation";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import {
+	Table,
+	TableBody,
+	TableCell,
+	TableHead,
+	TableHeader,
+	TableRow,
+} from "@/components/ui/table";
 import { cn } from "@/lib/utils";
 
+function stripThemeBackgrounds(
+	theme: Record<string, CSSProperties>
+): Record<string, CSSProperties> {
+	const cleaned: Record<string, CSSProperties> = {};
+	for (const key of Object.keys(theme)) {
+		const { background, backgroundColor, ...rest } = theme[key] as CSSProperties & {
+			background?: string;
+			backgroundColor?: string;
+		};
+		cleaned[key] = rest;
+	}
+	return cleaned;
+}
+
+const cleanMaterialDark = stripThemeBackgrounds(materialDark);
+const cleanMaterialLight = stripThemeBackgrounds(materialLight);
+
 // Storage for URL citations replaced during preprocess to avoid GFM autolink interference.
 // Populated in preprocessMarkdown, consumed in parseTextWithCitations.
 let _pendingUrlCitations = new Map<string, string>();
@@ -149,7 +178,7 @@ const MarkdownTextImpl = () => {
 
 export const MarkdownText = memo(MarkdownTextImpl);
 
-const CodeHeader: FC<CodeHeaderProps> = ({ language, code }) => {
+const InlineCodeHeader: FC<{ language: string; code: string }> = ({ language, code }) => {
 	const { isCopied, copyToClipboard } = useCopyToClipboard();
 	const onCopy = () => {
 		if (!code || isCopied) return;
@@ -157,8 +186,8 @@ const CodeHeader: FC<CodeHeaderProps> = ({ language, code }) => {
 	};
 
 	return (
-		<div className="aui-code-header-root mt-4 flex items-center justify-between gap-4 rounded-t-lg bg-muted-foreground/15 px-4 py-2 font-semibold text-foreground text-sm dark:bg-muted-foreground/20">
-			<span className="aui-code-header-language lowercase [&>span]:text-xs">{language}</span>
+		<div className="flex items-center justify-between gap-4 px-4 py-2 font-semibold text-muted-foreground text-sm">
+			<span className="lowercase text-xs">{language}</span>
 			<TooltipIconButton tooltip="Copy" onClick={onCopy}>
 				{!isCopied && <CopyIcon />}
 				{isCopied && <CheckIcon />}
@@ -188,17 +217,17 @@ const useCopyToClipboard = ({ copiedDuration = 3000 }: { copiedDuration?: number
 function processChildrenWithCitations(children: ReactNode): ReactNode {
 	if (typeof children === "string") {
 		const parsed = parseTextWithCitations(children);
-		return parsed.length === 1 && typeof parsed[0] === "string" ? children : <>{parsed}</>;
+		return parsed.length === 1 && typeof parsed[0] === "string" ? children : parsed;
 	}
 
 	if (Array.isArray(children)) {
-		return children.map((child, index) => {
+		return children.map((child) => {
 			if (typeof child === "string") {
 				const parsed = parseTextWithCitations(child);
 				return parsed.length === 1 && typeof parsed[0] === "string" ? (
 					child
 				) : (
-					<span key={index}>{parsed}</span>
+					<span key={child}>{parsed}</span>
 				);
 			}
 			return child;
@@ -208,6 +237,54 @@ function processChildrenWithCitations(children: ReactNode): ReactNode {
 	return children;
 }
 
+function extractDomain(url: string): string {
+	try {
+		const parsed = new URL(url);
+		return parsed.hostname.replace(/^www\./, "");
+	} catch {
+		return "";
+	}
+}
+
+function MarkdownImage({ src, alt }: { src?: string; alt?: string }) {
+	if (!src) return null;
+
+	const domain = extractDomain(src);
+
+	return (
+		<div className="my-4 w-fit max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<ImageRoot variant="ghost" size="full">
+				<ImageZoom src={src} alt={alt || "Image"}>
+					<ImagePreview
+						src={src}
+						alt={alt || "Image"}
+						className="max-h-[20rem] w-auto max-w-full object-contain"
+					/>
+				</ImageZoom>
+			</ImageRoot>
+
+			<div className="flex items-center justify-between px-5 py-3">
+				<div className="min-w-0 flex-1">
+					{alt && alt !== "Image" && (
+						<p className="text-sm font-semibold text-foreground line-clamp-2">{alt}</p>
+					)}
+					{domain && <p className="text-xs text-muted-foreground mt-0.5 truncate">{domain}</p>}
+				</div>
+				<a
+					href={src}
+					target="_blank"
+					rel="noopener noreferrer"
+					className="ml-3 shrink-0 inline-flex items-center gap-1.5 rounded-lg border px-3 py-1.5 text-xs font-medium text-foreground transition-colors hover:bg-muted"
+					onClick={(e) => e.stopPropagation()}
+				>
+					Open
+					<ExternalLinkIcon className="size-3" />
+				</a>
+			</div>
+		</div>
+	);
+}
+
 const defaultComponents = memoizeMarkdownComponents({
 	h1: ({ className, children, ...props }) => (
 		<h1
@@ -299,66 +376,69 @@ const defaultComponents = memoizeMarkdownComponents({
 		<hr className={cn("aui-md-hr my-5 border-b", className)} {...props} />
 	),
 	table: ({ className, ...props }) => (
-		<div className="aui-md-table-wrapper my-5 w-full overflow-x-auto">
-			<table
-				className={cn("aui-md-table w-full min-w-max border-separate border-spacing-0", className)}
-				{...props}
-			/>
+		<div className="aui-md-table-wrapper my-5 overflow-hidden rounded-2xl border">
+			<Table className={cn("aui-md-table", className)} {...props} />
 		</div>
 	),
+	thead: ({ className, ...props }) => (
+		<TableHeader className={cn("aui-md-thead", className)} {...props} />
+	),
+	tbody: ({ className, ...props }) => (
+		<TableBody className={cn("aui-md-tbody", className)} {...props} />
+	),
 	th: ({ className, children, ...props }) => (
-		<th
+		<TableHead
 			className={cn(
-				"aui-md-th bg-muted px-4 py-2 text-left font-bold first:rounded-tl-lg last:rounded-tr-lg [[align=center]]:text-center [[align=right]]:text-right",
+				"aui-md-th bg-muted/50 whitespace-normal [[align=center]]:text-center [[align=right]]:text-right",
 				className
 			)}
 			{...props}
 		>
 			{processChildrenWithCitations(children)}
-		</th>
+		</TableHead>
 	),
 	td: ({ className, children, ...props }) => (
-		<td
+		<TableCell
 			className={cn(
-				"aui-md-td border-b border-l px-4 py-2 text-left last:border-r [[align=center]]:text-center [[align=right]]:text-right",
+				"aui-md-td whitespace-normal [[align=center]]:text-center [[align=right]]:text-right",
 				className
 			)}
 			{...props}
 		>
 			{processChildrenWithCitations(children)}
-		</td>
-	),
-	tr: ({ className, ...props }) => (
-		<tr
-			className={cn(
-				"aui-md-tr m-0 border-b p-0 first:border-t [&:last-child>td:first-child]:rounded-bl-lg [&:last-child>td:last-child]:rounded-br-lg",
-				className
-			)}
-			{...props}
-		/>
+		</TableCell>
 	),
+	tr: ({ className, ...props }) => <TableRow className={cn("aui-md-tr", className)} {...props} />,
 	sup: ({ className, ...props }) => (
 		<sup className={cn("aui-md-sup [&>a]:text-xs [&>a]:no-underline", className)} {...props} />
 	),
-	pre: ({ className, ...props }) => (
-		<pre
-			className={cn(
-				"aui-md-pre overflow-x-auto rounded-t-none! rounded-b-lg bg-black p-4 text-white",
-				className
-			)}
-			{...props}
-		/>
-	),
-	code: function Code({ className, ...props }) {
+	pre: ({ children }) => <>{children}</>,
+	code: function Code({ className, children, ...props }) {
 		const isCodeBlock = useIsMarkdownCodeBlock();
+		const { resolvedTheme } = useTheme();
+		if (!isCodeBlock) {
+			return (
+				<code
+					className={cn("aui-md-inline-code rounded border bg-muted font-semibold", className)}
+					{...props}
+				/>
+			);
+		}
+		const language = /language-(\w+)/.exec(className || "")?.[1] ?? "text";
+		const codeString = String(children).replace(/\n$/, "");
+		const syntaxStyle = resolvedTheme === "dark" ? cleanMaterialDark : cleanMaterialLight;
 		return (
-			<code
-				className={cn(
-					!isCodeBlock && "aui-md-inline-code rounded border bg-muted font-semibold",
-					className
-				)}
-				{...props}
-			/>
+			<div className="mt-4 overflow-hidden rounded-2xl" style={{ background: "var(--syntax-bg)" }}>
+				<InlineCodeHeader language={language} code={codeString} />
+				<SyntaxHighlighter
+					style={syntaxStyle}
+					language={language}
+					PreTag="div"
+					customStyle={{ margin: 0, background: "transparent" }}
+				>
+					{codeString}
+				</SyntaxHighlighter>
+			</div>
 		);
 	},
 	strong: ({ className, children, ...props }) => (
@@ -371,5 +451,8 @@ const defaultComponents = memoizeMarkdownComponents({
 			{processChildrenWithCitations(children)}
 		</em>
 	),
-	CodeHeader,
+	img: ({ src, alt }) => (
+		<MarkdownImage src={typeof src === "string" ? src : undefined} alt={alt} />
+	),
+	CodeHeader: () => null,
 });

surfsense_web/components/assistant-ui/thinking-steps.tsx

@@ -1,14 +1,17 @@
-import { useAssistantState, useThreadViewport } from "@assistant-ui/react";
+import { makeAssistantDataUI, useAuiState } from "@assistant-ui/react";
 import { ChevronRightIcon } from "lucide-react";
 import type { FC } from "react";
-import { createContext, useCallback, useContext, useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
-import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
 import { cn } from "@/lib/utils";
 
-// Context to pass thinking steps to AssistantMessage
-export const ThinkingStepsContext = createContext<Map<string, ThinkingStep[]>>(new Map());
+export interface ThinkingStep {
+	id: string;
+	title: string;
+	items: string[];
+	status: "pending" | "in_progress" | "completed";
+}
 
 /**
  * Chain of thought display component - single collapsible dropdown design
@@ -19,7 +22,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 }) => {
 	const [isOpen, setIsOpen] = useState(true);
 
-	// Derive effective status for each step
 	const getEffectiveStatus = useCallback(
 		(step: ThinkingStep): "pending" | "in_progress" | "completed" => {
 			if (step.status === "in_progress" && !isThreadRunning) {
@@ -37,7 +39,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 		steps.every((s) => getEffectiveStatus(s) === "completed");
 	const isProcessing = isThreadRunning && !allCompleted;
 
-	// Auto-collapse when all tasks are completed
 	useEffect(() => {
 		if (allCompleted) {
 			setIsOpen(false);
@@ -62,7 +63,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 	return (
 		<div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2">
 			<div className="rounded-lg">
-				{/* Main collapsible header */}
 				<button
 					type="button"
 					onClick={() => setIsOpen(!isOpen)}
@@ -71,20 +71,17 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 						"text-muted-foreground hover:text-foreground"
 					)}
 				>
-					{/* Header text with shimmer if processing (streaming) */}
 					{isProcessing ? (
 						<TextShimmerLoader text={getHeaderText()} size="sm" />
 					) : (
 						<span>{getHeaderText()}</span>
 					)}
 
-					{/* Chevron */}
 					<ChevronRightIcon
 						className={cn("size-4 transition-transform duration-200", isOpen && "rotate-90")}
 					/>
 				</button>
 
-				{/* Collapsible content with CSS grid animation */}
 				<div
 					className={cn(
 						"grid transition-[grid-template-rows] duration-300 ease-out",
@@ -99,13 +96,10 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 
 								return (
 									<div key={step.id} className="relative flex gap-3">
-										{/* Dot and line column */}
 										<div className="relative flex flex-col items-center w-2">
-											{/* Vertical connection line - extends to next dot */}
 											{!isLast && (
 												<div className="absolute left-1/2 top-[15px] -bottom-[7px] w-px -translate-x-1/2 bg-muted-foreground/30" />
 											)}
-											{/* Step dot - on top of line */}
 											<div className="relative z-10 mt-[7px] flex shrink-0 items-center justify-center">
 												{effectiveStatus === "in_progress" ? (
 													<span className="relative flex size-2">
@@ -118,9 +112,7 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 											</div>
 										</div>
 
-										{/* Step content */}
 										<div className="flex-1 min-w-0 pb-4">
-											{/* Step title */}
 											<div
 												className={cn(
 													"text-sm leading-5",
@@ -132,11 +124,10 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 												{step.title}
 											</div>
 
-											{/* Step items (sub-content) */}
 											{step.items && step.items.length > 0 && (
 												<div className="mt-1 space-y-0.5">
-													{step.items.map((item, idx) => (
-														<ChainOfThoughtItem key={`${step.id}-item-${idx}`} className="text-xs">
+													{step.items.map((item) => (
+														<ChainOfThoughtItem key={`${step.id}-${item}`} className="text-xs">
 															{item}
 														</ChainOfThoughtItem>
 													))}
@@ -155,51 +146,26 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 };
 
 /**
- * Component that handles auto-scroll when thinking steps update.
- * Uses useThreadViewport to scroll to bottom when thinking steps change,
- * ensuring the user always sees the latest content during streaming.
+ * assistant-ui data UI component that renders thinking steps from message content.
+ * Registered globally via makeAssistantDataUI — renders inside MessagePrimitive.Parts
+ * at the position of the data part in the content array.
  */
-export const ThinkingStepsScrollHandler: FC = () => {
-	const thinkingStepsMap = useContext(ThinkingStepsContext);
-	const viewport = useThreadViewport();
-	const isRunning = useAssistantState(({ thread }) => thread.isRunning);
-	// Track the serialized state to detect any changes
-	const prevStateRef = useRef<string>("");
+function ThinkingStepsDataRenderer({ data }: { name: string; data: unknown }) {
+	const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
+	const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
+	const isMessageStreaming = isThreadRunning && isLastMessage;
 
-	useEffect(() => {
-		// Only act during streaming
-		if (!isRunning) {
-			prevStateRef.current = "";
-			return;
-		}
+	const steps = (data as { steps: ThinkingStep[] } | null)?.steps ?? [];
+	if (steps.length === 0) return null;
 
-		// Serialize the thinking steps state to detect any changes
-		// This catches new steps, status changes, and item additions
-		let stateString = "";
-		thinkingStepsMap.forEach((steps, msgId) => {
-			steps.forEach((step) => {
-				stateString += `${msgId}:${step.id}:${step.status}:${step.items?.length || 0};`;
-			});
-		});
+	return (
+		<div className="mb-3 -mx-2 leading-normal">
+			<ThinkingStepsDisplay steps={steps} isThreadRunning={isMessageStreaming} />
+		</div>
+	);
+}
 
-		// If state changed at all during streaming, scroll
-		if (stateString !== prevStateRef.current && stateString !== "") {
-			prevStateRef.current = stateString;
-
-			// Multiple attempts to ensure scroll happens after DOM updates
-			const scrollAttempt = () => {
-				try {
-					viewport.scrollToBottom();
-				} catch {
-					// Ignore errors - viewport might not be ready
-				}
-			};
-
-			// Delayed attempts to handle async DOM updates
-			requestAnimationFrame(scrollAttempt);
-			setTimeout(scrollAttempt, 100);
-		}
-	}, [thinkingStepsMap, viewport, isRunning]);
-
-	return null; // This component doesn't render anything
-};
+export const ThinkingStepsDataUI = makeAssistantDataUI({
+	name: "thinking-steps",
+	render: ThinkingStepsDataRenderer,
+});

surfsense_web/components/assistant-ui/thread.tsx

@@ -1,27 +1,18 @@
 import {
-	ActionBarPrimitive,
-	AssistantIf,
-	BranchPickerPrimitive,
+	AuiIf,
 	ComposerPrimitive,
-	ErrorPrimitive,
 	MessagePrimitive,
 	ThreadPrimitive,
-	useAssistantState,
-	useComposerRuntime,
+	useAui,
+	useAuiState,
 } from "@assistant-ui/react";
 import { useAtom, useAtomValue, useSetAtom } from "jotai";
 import {
 	AlertCircle,
 	ArrowDownIcon,
 	ArrowUpIcon,
-	CheckIcon,
-	ChevronLeftIcon,
-	ChevronRightIcon,
-	CopyIcon,
-	DownloadIcon,
 	Globe,
 	Plus,
-	RefreshCwIcon,
 	Settings2,
 	SquareIcon,
 	Unplug,
@@ -32,7 +23,7 @@ import {
 import { AnimatePresence, motion } from "motion/react";
 import Image from "next/image";
 import { useParams } from "next/navigation";
-import { type FC, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react";
+import { type FC, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { createPortal } from "react-dom";
 import {
 	agentToolsAtom,
@@ -63,12 +54,6 @@ import {
 	InlineMentionEditor,
 	type InlineMentionEditorRef,
 } from "@/components/assistant-ui/inline-mention-editor";
-import { MarkdownText } from "@/components/assistant-ui/markdown-text";
-import {
-	ThinkingStepsContext,
-	ThinkingStepsDisplay,
-} from "@/components/assistant-ui/thinking-steps";
-import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { UserMessage } from "@/components/assistant-ui/user-message";
 import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/components/layout/ui/sidebar/SidebarSlideOutPanel";
@@ -76,7 +61,6 @@ import {
 	DocumentMentionPicker,
 	type DocumentMentionPickerRef,
 } from "@/components/new-chat/document-mention-picker";
-import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
 import { Avatar, AvatarFallback, AvatarGroup } from "@/components/ui/avatar";
 import { Button } from "@/components/ui/button";
 import { Drawer, DrawerContent, DrawerHandle, DrawerTitle } from "@/components/ui/drawer";
@@ -111,16 +95,8 @@ const CYCLING_PLACEHOLDERS = [
 	"Check if this week's Slack messages reference any GitHub issues",
 ];
 
-interface ThreadProps {
-	messageThinkingSteps?: Map<string, ThinkingStep[]>;
-}
-
-export const Thread: FC<ThreadProps> = ({ messageThinkingSteps = new Map() }) => {
-	return (
-		<ThinkingStepsContext.Provider value={messageThinkingSteps}>
-			<ThreadContent />
-		</ThinkingStepsContext.Provider>
-	);
+export const Thread: FC = () => {
+	return <ThreadContent />;
 };
 
 const ThreadContent: FC = () => {
@@ -135,9 +111,9 @@ const ThreadContent: FC = () => {
 				turnAnchor="top"
 				className="aui-thread-viewport relative flex flex-1 min-h-0 flex-col overflow-y-auto px-4 pt-4"
 			>
-				<AssistantIf condition={({ thread }) => thread.isEmpty}>
+				<AuiIf condition={({ thread }) => thread.isEmpty}>
 					<ThreadWelcome />
-				</AssistantIf>
+				</AuiIf>
 
 				<ThreadPrimitive.Messages
 					components={{
@@ -152,11 +128,11 @@ const ThreadContent: FC = () => {
 					style={{ paddingBottom: "max(1rem, env(safe-area-inset-bottom))" }}
 				>
 					<ThreadScrollToBottom />
-					<AssistantIf condition={({ thread }) => !thread.isEmpty}>
+					<AuiIf condition={({ thread }) => !thread.isEmpty}>
 						<div className="fade-in slide-in-from-bottom-4 animate-in duration-500 ease-out fill-mode-both">
 							<Composer />
 						</div>
-					</AssistantIf>
+					</AuiIf>
 				</ThreadPrimitive.ViewportFooter>
 			</ThreadPrimitive.Viewport>
 		</ThreadPrimitive.Root>
@@ -327,11 +303,11 @@ const Composer: FC = () => {
 	const editorContainerRef = useRef<HTMLDivElement>(null);
 	const documentPickerRef = useRef<DocumentMentionPickerRef>(null);
 	const { search_space_id, chat_id } = useParams();
-	const composerRuntime = useComposerRuntime();
+	const aui = useAui();
 	const hasAutoFocusedRef = useRef(false);
 
-	const isThreadEmpty = useAssistantState(({ thread }) => thread.isEmpty);
-	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
+	const isThreadEmpty = useAuiState(({ thread }) => thread.isEmpty);
+	const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
 
 	// Cycling placeholder state - only cycles in new chats
 	const [placeholderIndex, setPlaceholderIndex] = useState(0);
@@ -378,7 +354,7 @@ const Composer: FC = () => {
 	// hooks never fire their own network requests (eliminates N+1 API calls).
 	// Return a primitive string from the selector so useSyncExternalStore can
 	// compare snapshots by value and avoid infinite re-render loops.
-	const assistantIdsKey = useAssistantState(({ thread }) =>
+	const assistantIdsKey = useAuiState(({ thread }) =>
 		thread.messages
 			.filter((m) => m.role === "assistant" && m.id?.startsWith("msg-"))
 			.map((m) => m.id?.replace("msg-", ""))
@@ -414,9 +390,9 @@ const Composer: FC = () => {
 	// Sync editor text with assistant-ui composer runtime
 	const handleEditorChange = useCallback(
 		(text: string) => {
-			composerRuntime.setText(text);
+			aui.composer().setText(text);
 		},
-		[composerRuntime]
+		[aui]
 	);
 
 	// Open document picker when @ mention is triggered
@@ -469,7 +445,7 @@ const Composer: FC = () => {
 			return;
 		}
 		if (!showDocumentPopover) {
-			composerRuntime.send();
+			aui.composer().send();
 			editorRef.current?.clear();
 			setMentionedDocuments([]);
 			setSidebarDocs([]);
@@ -478,7 +454,7 @@ const Composer: FC = () => {
 		showDocumentPopover,
 		isThreadRunning,
 		isBlockedByOtherUser,
-		composerRuntime,
+		aui,
 		setMentionedDocuments,
 		setSidebarDocs,
 	]);
@@ -591,7 +567,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 		const atBottom = el.scrollHeight - el.scrollTop - el.clientHeight <= 2;
 		setToolsScrollPos(atTop ? "top" : atBottom ? "bottom" : "middle");
 	}, []);
-	const isComposerTextEmpty = useAssistantState(({ composer }) => {
+	const isComposerTextEmpty = useAuiState(({ composer }) => {
 		const text = composer.text?.trim() || "";
 		return text.length === 0;
 	});
@@ -1007,16 +983,14 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 					</button>
 				)}
 			</div>
-
 			{!hasModelConfigured && (
 				<div className="flex items-center gap-1.5 text-amber-600 dark:text-amber-400 text-xs">
 					<AlertCircle className="size-3" />
 					<span>Select a model</span>
 				</div>
 			)}
-
 			<div className="flex items-center gap-2">
-				<AssistantIf condition={({ thread }) => !thread.isRunning}>
+				<AuiIf condition={({ thread }) => !thread.isRunning}>
 					<ComposerPrimitive.Send asChild disabled={isSendDisabled}>
 						<TooltipIconButton
 							tooltip={
@@ -1042,9 +1016,9 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 							<ArrowUpIcon className="aui-composer-send-icon size-4" />
 						</TooltipIconButton>
 					</ComposerPrimitive.Send>
-				</AssistantIf>
+				</AuiIf>
 
-				<AssistantIf condition={({ thread }) => thread.isRunning}>
+				<AuiIf condition={({ thread }) => thread.isRunning}>
 					<ComposerPrimitive.Cancel asChild>
 						<Button
 							type="button"
@@ -1056,7 +1030,7 @@ const ComposerAction: FC<ComposerActionProps> = ({ isBlockedByOtherUser = false
 							<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
 						</Button>
 					</ComposerPrimitive.Cancel>
-				</AssistantIf>
+				</AuiIf>
 			</div>
 		</div>
 	);
@@ -1080,17 +1054,11 @@ interface ToolGroup {
 const TOOL_GROUPS: ToolGroup[] = [
 	{
 		label: "Research",
-		tools: ["search_knowledge_base", "search_surfsense_docs", "scrape_webpage", "link_preview"],
+		tools: ["search_knowledge_base", "search_surfsense_docs", "scrape_webpage"],
 	},
 	{
 		label: "Generate",
-		tools: [
-			"generate_podcast",
-			"generate_video_presentation",
-			"generate_report",
-			"generate_image",
-			"display_image",
-		],
+		tools: ["generate_podcast", "generate_video_presentation", "generate_report", "generate_image"],
 	},
 	{
 		label: "Memory",
@@ -1140,97 +1108,6 @@ const TOOL_GROUPS: ToolGroup[] = [
 	},
 ];
 
-const MessageError: FC = () => {
-	return (
-		<MessagePrimitive.Error>
-			<ErrorPrimitive.Root className="aui-message-error-root mt-2 rounded-md border border-destructive bg-destructive/10 p-3 text-destructive text-sm dark:bg-destructive/5 dark:text-red-200">
-				<ErrorPrimitive.Message className="aui-message-error-message line-clamp-2" />
-			</ErrorPrimitive.Root>
-		</MessagePrimitive.Error>
-	);
-};
-
-/**
- * Custom component to render thinking steps from Context
- */
-const ThinkingStepsPart: FC = () => {
-	const thinkingStepsMap = useContext(ThinkingStepsContext);
-
-	// Get the current message ID to look up thinking steps
-	const messageId = useAssistantState(({ message }) => message?.id);
-	const thinkingSteps = thinkingStepsMap.get(messageId) || [];
-
-	// Check if this specific message is currently streaming
-	// A message is streaming if: thread is running AND this is the last assistant message
-	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
-	const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false);
-	const isMessageStreaming = isThreadRunning && isLastMessage;
-
-	if (thinkingSteps.length === 0) return null;
-
-	return (
-		<div className="mb-3">
-			<ThinkingStepsDisplay steps={thinkingSteps} isThreadRunning={isMessageStreaming} />
-		</div>
-	);
-};
-
-const AssistantMessageInner: FC = () => {
-	return (
-		<>
-			{/* Render thinking steps from message content - this ensures proper scroll tracking */}
-			<ThinkingStepsPart />
-
-			<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
-				<MessagePrimitive.Parts
-					components={{
-						Text: MarkdownText,
-						tools: { Fallback: ToolFallback },
-					}}
-				/>
-				<MessageError />
-			</div>
-
-			<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
-				<BranchPicker />
-				<AssistantActionBar />
-			</div>
-		</>
-	);
-};
-
-const AssistantActionBar: FC = () => {
-	return (
-		<ActionBarPrimitive.Root
-			hideWhenRunning
-			autohide="not-last"
-			autohideFloat="single-branch"
-			className="aui-assistant-action-bar-root -ml-1 col-start-3 row-start-2 flex gap-1 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm"
-		>
-			<ActionBarPrimitive.Copy asChild>
-				<TooltipIconButton tooltip="Copy">
-					<AssistantIf condition={({ message }) => message.isCopied}>
-						<CheckIcon />
-					</AssistantIf>
-					<AssistantIf condition={({ message }) => !message.isCopied}>
-						<CopyIcon />
-					</AssistantIf>
-				</TooltipIconButton>
-			</ActionBarPrimitive.Copy>
-			<ActionBarPrimitive.ExportMarkdown asChild>
-				<TooltipIconButton tooltip="Export as Markdown">
-					<DownloadIcon />
-				</TooltipIconButton>
-			</ActionBarPrimitive.ExportMarkdown>
-			<ActionBarPrimitive.Reload asChild>
-				<TooltipIconButton tooltip="Refresh">
-					<RefreshCwIcon />
-				</TooltipIconButton>
-			</ActionBarPrimitive.Reload>
-		</ActionBarPrimitive.Root>
-	);
-};
-
 const EditComposer: FC = () => {
 	return (
 		<MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
@@ -1253,30 +1130,3 @@ const EditComposer: FC = () => {
 		</MessagePrimitive.Root>
 	);
 };
-
-const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({ className, ...rest }) => {
-	return (
-		<BranchPickerPrimitive.Root
-			hideWhenSingleBranch
-			className={cn(
-				"aui-branch-picker-root -ml-2 mr-2 inline-flex items-center text-muted-foreground text-xs",
-				className
-			)}
-			{...rest}
-		>
-			<BranchPickerPrimitive.Previous asChild>
-				<TooltipIconButton tooltip="Previous">
-					<ChevronLeftIcon />
-				</TooltipIconButton>
-			</BranchPickerPrimitive.Previous>
-			<span className="aui-branch-picker-state font-medium">
-				<BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
-			</span>
-			<BranchPickerPrimitive.Next asChild>
-				<TooltipIconButton tooltip="Next">
-					<ChevronRightIcon />
-				</TooltipIconButton>
-			</BranchPickerPrimitive.Next>
-		</BranchPickerPrimitive.Root>
-	);
-};

surfsense_web/components/assistant-ui/tool-fallback.tsx

@@ -1,75 +1,134 @@
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { CheckIcon, ChevronDownIcon, ChevronUpIcon, XCircleIcon } from "lucide-react";
 import { useState } from "react";
-import { Button } from "@/components/ui/button";
+import { getToolIcon } from "@/contracts/enums/toolIcons";
 import { cn } from "@/lib/utils";
 
+function formatToolName(name: string): string {
+	return name.replace(/_/g, " ").replace(/\b\w/g, (c) => c.toUpperCase());
+}
+
 export const ToolFallback: ToolCallMessagePartComponent = ({
 	toolName,
 	argsText,
 	result,
 	status,
 }) => {
-	const [isCollapsed, setIsCollapsed] = useState(true);
+	const [isExpanded, setIsExpanded] = useState(false);
 
 	const isCancelled = status?.type === "incomplete" && status.reason === "cancelled";
+	const isError = status?.type === "incomplete" && status.reason === "error";
+	const isRunning = status?.type === "running" || status?.type === "requires-action";
 	const cancelledReason =
 		isCancelled && status.error
 			? typeof status.error === "string"
 				? status.error
 				: JSON.stringify(status.error)
 			: null;
+	const errorReason =
+		isError && status.error
+			? typeof status.error === "string"
+				? status.error
+				: JSON.stringify(status.error)
+			: null;
+
+	const Icon = getToolIcon(toolName);
+	const displayName = formatToolName(toolName);
 
 	return (
 		<div
 			className={cn(
-				"aui-tool-fallback-root mb-4 flex w-full flex-col gap-3 rounded-lg border py-3",
-				isCancelled && "border-muted-foreground/30 bg-muted/30"
+				"my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none",
+				isCancelled && "opacity-60",
+				isError && "border-destructive/20 bg-destructive/5"
 			)}
 		>
-			<div className="aui-tool-fallback-header flex items-center gap-2 px-4">
-				{isCancelled ? (
-					<XCircleIcon className="aui-tool-fallback-icon size-4 text-muted-foreground" />
-				) : (
-					<CheckIcon className="aui-tool-fallback-icon size-4" />
-				)}
-				<p
+			<button
+				type="button"
+				onClick={() => setIsExpanded(!isExpanded)}
+				className="flex w-full items-center gap-3 px-5 py-4 text-left transition-colors hover:bg-muted/50 focus:outline-none focus-visible:outline-none"
+			>
+				<div
 					className={cn(
-						"aui-tool-fallback-title grow",
-						isCancelled && "text-muted-foreground line-through"
+						"flex size-8 shrink-0 items-center justify-center rounded-lg",
+						isError ? "bg-destructive/10" : isCancelled ? "bg-muted" : "bg-primary/10"
 					)}
 				>
-					{isCancelled ? "Cancelled tool: " : "Used tool: "}
-					<b>{toolName}</b>
-				</p>
-				<Button onClick={() => setIsCollapsed(!isCollapsed)}>
-					{isCollapsed ? <ChevronUpIcon /> : <ChevronDownIcon />}
-				</Button>
-			</div>
-			{!isCollapsed && (
-				<div className="aui-tool-fallback-content flex flex-col gap-2 border-t pt-2">
-					{cancelledReason && (
-						<div className="aui-tool-fallback-cancelled-root px-4">
-							<p className="aui-tool-fallback-cancelled-header font-semibold text-muted-foreground">
-								Cancelled reason:
-							</p>
-							<p className="aui-tool-fallback-cancelled-reason text-muted-foreground">
-								{cancelledReason}
-							</p>
-						</div>
-					)}
-					<div className={cn("aui-tool-fallback-args-root px-4", isCancelled && "opacity-60")}>
-						<pre className="aui-tool-fallback-args-value whitespace-pre-wrap">{argsText}</pre>
-					</div>
-					{!isCancelled && result !== undefined && (
-						<div className="aui-tool-fallback-result-root border-t border-dashed px-4 pt-2">
-							<p className="aui-tool-fallback-result-header font-semibold">Result:</p>
-							<pre className="aui-tool-fallback-result-content whitespace-pre-wrap">
-								{typeof result === "string" ? result : JSON.stringify(result, null, 2)}
-							</pre>
-						</div>
+					{isError ? (
+						<XCircleIcon className="size-4 text-destructive" />
+					) : isCancelled ? (
+						<XCircleIcon className="size-4 text-muted-foreground" />
+					) : isRunning ? (
+						<Icon className="size-4 text-primary animate-pulse" />
+					) : (
+						<CheckIcon className="size-4 text-primary" />
 					)}
 				</div>
+
+				<div className="flex-1 min-w-0">
+					<p
+						className={cn(
+							"text-sm font-semibold",
+							isError
+								? "text-destructive"
+								: isCancelled
+									? "text-muted-foreground line-through"
+									: "text-foreground"
+						)}
+					>
+						{isRunning
+							? displayName
+							: isCancelled
+								? `Cancelled: ${displayName}`
+								: isError
+									? `Failed: ${displayName}`
+									: displayName}
+					</p>
+					{isRunning && <p className="text-xs text-muted-foreground mt-0.5">Running...</p>}
+					{cancelledReason && (
+						<p className="text-xs text-muted-foreground mt-0.5 truncate">{cancelledReason}</p>
+					)}
+					{errorReason && (
+						<p className="text-xs text-destructive/80 mt-0.5 truncate">{errorReason}</p>
+					)}
+				</div>
+
+				{!isRunning && (
+					<div className="shrink-0 text-muted-foreground">
+						{isExpanded ? (
+							<ChevronDownIcon className="size-4" />
+						) : (
+							<ChevronUpIcon className="size-4" />
+						)}
+					</div>
+				)}
+			</button>
+
+			{isExpanded && !isRunning && (
+				<>
+					<div className="mx-5 h-px bg-border/50" />
+					<div className="px-5 py-3 space-y-3">
+						{argsText && (
+							<div>
+								<p className="text-xs font-medium text-muted-foreground mb-1">Arguments</p>
+								<pre className="text-xs text-foreground/80 whitespace-pre-wrap break-all">
+									{argsText}
+								</pre>
+							</div>
+						)}
+						{!isCancelled && result !== undefined && (
+							<>
+								<div className="h-px bg-border/30" />
+								<div>
+									<p className="text-xs font-medium text-muted-foreground mb-1">Result</p>
+									<pre className="text-xs text-foreground/80 whitespace-pre-wrap break-all">
+										{typeof result === "string" ? result : JSON.stringify(result, null, 2)}
+									</pre>
+								</div>
+							</>
+						)}
+					</div>
+				</>
 			)}
 		</div>
 	);

surfsense_web/components/assistant-ui/user-message.tsx

@@ -1,6 +1,7 @@
-import { ActionBarPrimitive, MessagePrimitive, useAssistantState } from "@assistant-ui/react";
+import { ActionBarPrimitive, AuiIf, MessagePrimitive, useAuiState } from "@assistant-ui/react";
 import { useAtomValue } from "jotai";
-import { FileText, Pen } from "lucide-react";
+import { CheckIcon, CopyIcon, FileText, Pen } from "lucide-react";
+import Image from "next/image";
 import { type FC, useState } from "react";
 import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
@@ -24,82 +25,81 @@ const UserAvatar: FC<AuthorMetadata> = ({ displayName, avatarUrl }) => {
 
 	if (avatarUrl && !hasError) {
 		return (
-			<img
+			<Image
 				src={avatarUrl}
 				alt={displayName || "User"}
-				className="size-8 rounded-full object-cover"
+				width={32}
+				height={32}
+				className="size-8 rounded-full object-cover select-none"
 				referrerPolicy="no-referrer"
 				onError={() => setHasError(true)}
+				unoptimized
 			/>
 		);
 	}
 
 	return (
-		<div className="flex size-8 items-center justify-center rounded-full bg-primary/10 text-xs font-medium text-primary">
+		<div className="flex size-8 items-center justify-center rounded-full bg-primary/10 text-xs font-medium text-primary select-none">
 			{initials}
 		</div>
 	);
 };
 
 export const UserMessage: FC = () => {
-	const messageId = useAssistantState(({ message }) => message?.id);
+	const messageId = useAuiState(({ message }) => message?.id);
 	const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
 	const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined;
-	const metadata = useAssistantState(({ message }) => message?.metadata);
+	const metadata = useAuiState(({ message }) => message?.metadata);
 	const author = metadata?.custom?.author as AuthorMetadata | undefined;
 
 	return (
 		<MessagePrimitive.Root
-			className="aui-user-message-root fade-in slide-in-from-bottom-1 mx-auto grid w-full max-w-(--thread-max-width) animate-in auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] content-start gap-y-2 px-2 py-3 duration-150 [&:where(>*)]:col-start-2"
+			className="group/user-msg aui-user-message-root fade-in slide-in-from-bottom-1 mx-auto grid w-full max-w-(--thread-max-width) animate-in auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] content-start gap-y-2 px-2 pt-3 pb-8 duration-150 [&:where(>*)]:col-start-2"
 			data-role="user"
 		>
-			<div className="aui-user-message-content-wrapper col-start-2 min-w-0 flex items-end gap-2">
-				<div className="flex-1 min-w-0">
-					{/* Display mentioned documents */}
-					{mentionedDocs && mentionedDocs.length > 0 && (
-						<div className="flex flex-wrap items-end gap-2 mb-2 justify-end">
-							{/* Mentioned documents as chips */}
-							{mentionedDocs?.map((doc) => (
-								<span
-									key={`${doc.document_type}:${doc.id}`}
-									className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20"
-									title={doc.title}
-								>
-									<FileText className="size-3" />
-									<span className="max-w-[150px] truncate">{doc.title}</span>
-								</span>
-							))}
-						</div>
-					)}
-					{/* Message bubble with action bar positioned relative to it */}
-					<div className="relative">
+			<div className="col-start-2 min-w-0">
+				<div className="aui-user-message-content-wrapper flex items-end gap-2">
+					<div className="relative flex-1 min-w-0">
+						{mentionedDocs && mentionedDocs.length > 0 && (
+							<div className="flex flex-wrap items-end gap-2 mb-2 justify-end">
+								{mentionedDocs?.map((doc) => (
+									<span
+										key={`${doc.document_type}:${doc.id}`}
+										className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20"
+										title={doc.title}
+									>
+										<FileText className="size-3" />
+										<span className="max-w-[150px] truncate">{doc.title}</span>
+									</span>
+								))}
+							</div>
+						)}
 						<div className="aui-user-message-content wrap-break-word rounded-2xl bg-muted px-4 py-2.5 text-foreground">
 							<MessagePrimitive.Parts />
 						</div>
-						<div className="aui-user-action-bar-wrapper absolute top-1/2 right-full -translate-y-1/2 pr-1">
+						<div className="absolute right-0 top-full mt-1 z-10 opacity-100 pointer-events-auto md:opacity-0 md:pointer-events-none md:transition-opacity md:duration-200 md:delay-300 md:group-hover/user-msg:opacity-100 md:group-hover/user-msg:delay-0 md:group-hover/user-msg:pointer-events-auto">
 							<UserActionBar />
 						</div>
 					</div>
+					{author && (
+						<div className="shrink-0 mb-1.5">
+							<UserAvatar displayName={author.displayName} avatarUrl={author.avatarUrl} />
+						</div>
+					)}
 				</div>
-				{/* User avatar - only shown in shared chats */}
-				{author && (
-					<div className="shrink-0 mb-1.5">
-						<UserAvatar displayName={author.displayName} avatarUrl={author.avatarUrl} />
-					</div>
-				)}
 			</div>
 		</MessagePrimitive.Root>
 	);
 };
 
 const UserActionBar: FC = () => {
-	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
+	const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
 
 	// Get current message ID
-	const currentMessageId = useAssistantState(({ message }) => message?.id);
+	const currentMessageId = useAuiState(({ message }) => message?.id);
 
 	// Find the last user message ID in the thread (computed once, memoized by selector)
-	const lastUserMessageId = useAssistantState(({ thread }) => {
+	const lastUserMessageId = useAuiState(({ thread }) => {
 		const messages = thread.messages;
 		for (let i = messages.length - 1; i >= 0; i--) {
 			if (messages[i].role === "user") {
@@ -118,13 +118,21 @@ const UserActionBar: FC = () => {
 	return (
 		<ActionBarPrimitive.Root
 			hideWhenRunning
-			autohide="not-last"
-			className="aui-user-action-bar-root flex flex-col items-end"
+			className="aui-user-action-bar-root flex items-center justify-end gap-1 text-muted-foreground"
 		>
-			{/* Only allow editing the last user message */}
+			<ActionBarPrimitive.Copy asChild>
+				<TooltipIconButton tooltip="Copy">
+					<AuiIf condition={({ message }) => message.isCopied}>
+						<CheckIcon />
+					</AuiIf>
+					<AuiIf condition={({ message }) => !message.isCopied}>
+						<CopyIcon />
+					</AuiIf>
+				</TooltipIconButton>
+			</ActionBarPrimitive.Copy>
 			{canEdit && (
 				<ActionBarPrimitive.Edit asChild>
-					<TooltipIconButton tooltip="Edit" className="aui-user-action-edit p-4">
+					<TooltipIconButton tooltip="Edit" className="aui-user-action-edit">
 						<Pen />
 					</TooltipIconButton>
 				</ActionBarPrimitive.Edit>

surfsense_web/components/public-chat/public-chat-view.tsx

@@ -1,14 +1,9 @@
 "use client";
 
 import { AssistantRuntimeProvider } from "@assistant-ui/react";
+import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps";
 import { Navbar } from "@/components/homepage/navbar";
 import { ReportPanel } from "@/components/report-panel/report-panel";
-import { DisplayImageToolUI } from "@/components/tool-ui/display-image";
-import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
-import { GenerateReportToolUI } from "@/components/tool-ui/generate-report";
-import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview";
-import { ScrapeWebpageToolUI } from "@/components/tool-ui/scrape-webpage";
-import { GenerateVideoPresentationToolUI } from "@/components/tool-ui/video-presentation";
 import { Spinner } from "@/components/ui/spinner";
 import { usePublicChat } from "@/hooks/use-public-chat";
 import { usePublicChatRuntime } from "@/hooks/use-public-chat-runtime";
@@ -45,14 +40,7 @@ export function PublicChatView({ shareToken }: PublicChatViewProps) {
 		<main className="min-h-screen bg-main-panel text-foreground overflow-x-hidden">
 			<Navbar scrolledBgClassName={navbarScrolledBg} />
 			<AssistantRuntimeProvider runtime={runtime}>
-				{/* Tool UIs for rendering tool results */}
-				<GeneratePodcastToolUI />
-				<GenerateReportToolUI />
-				<GenerateVideoPresentationToolUI />
-				<LinkPreviewToolUI />
-				<DisplayImageToolUI />
-				<ScrapeWebpageToolUI />
-
+				<ThinkingStepsDataUI />
 				<div className="flex h-screen pt-16 overflow-hidden">
 					<div className="flex-1 flex flex-col min-w-0 overflow-hidden">
 						<PublicThread footer={<PublicChatFooter shareToken={shareToken} />} />

surfsense_web/components/public-chat/public-thread.tsx

@@ -2,16 +2,20 @@
 
 import {
 	ActionBarPrimitive,
-	AssistantIf,
+	AuiIf,
 	MessagePrimitive,
 	ThreadPrimitive,
-	useAssistantState,
+	useAuiState,
 } from "@assistant-ui/react";
 import { CheckIcon, CopyIcon } from "lucide-react";
 import { type FC, type ReactNode, useState } from "react";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { GenerateImageToolUI } from "@/components/tool-ui/generate-image";
+import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
+import { GenerateReportToolUI } from "@/components/tool-ui/generate-report";
+import { GenerateVideoPresentationToolUI } from "@/components/tool-ui/video-presentation";
 
 interface PublicThreadProps {
 	footer?: ReactNode;
@@ -75,6 +79,7 @@ const UserAvatar: FC<AuthorMetadata & { hasError: boolean; onError: () => void }
 
 	if (avatarUrl && !hasError) {
 		return (
+			// biome-ignore lint/performance/noImgElement: external OAuth/profile avatar URL
 			<img
 				src={avatarUrl}
 				alt={displayName || "User"}
@@ -93,7 +98,7 @@ const UserAvatar: FC<AuthorMetadata & { hasError: boolean; onError: () => void }
 };
 
 const PublicUserMessage: FC = () => {
-	const metadata = useAssistantState(({ message }) => message?.metadata);
+	const metadata = useAuiState(({ message }) => message?.metadata);
 	const author = metadata?.custom?.author as AuthorMetadata | undefined;
 
 	return (
@@ -139,7 +144,19 @@ const PublicAssistantMessage: FC = () => {
 				<MessagePrimitive.Parts
 					components={{
 						Text: MarkdownText,
-						tools: { Fallback: ToolFallback },
+						tools: {
+							by_name: {
+								generate_podcast: GeneratePodcastToolUI,
+								generate_report: GenerateReportToolUI,
+								generate_video_presentation: GenerateVideoPresentationToolUI,
+								display_image: GenerateImageToolUI,
+								generate_image: GenerateImageToolUI,
+								link_preview: () => null,
+								multi_link_preview: () => null,
+								scrape_webpage: () => null,
+							},
+							Fallback: ToolFallback,
+						},
 					}}
 				/>
 			</div>
@@ -160,12 +177,12 @@ const PublicAssistantActionBar: FC = () => {
 		>
 			<ActionBarPrimitive.Copy asChild>
 				<TooltipIconButton tooltip="Copy">
-					<AssistantIf condition={({ message }) => message.isCopied}>
+					<AuiIf condition={({ message }) => message.isCopied}>
 						<CheckIcon />
-					</AssistantIf>
-					<AssistantIf condition={({ message }) => !message.isCopied}>
+					</AuiIf>
+					<AuiIf condition={({ message }) => !message.isCopied}>
 						<CopyIcon />
-					</AssistantIf>
+					</AuiIf>
 				</TooltipIconButton>
 			</ActionBarPrimitive.Copy>
 		</ActionBarPrimitive.Root>

surfsense_web/components/tool-ui/article/index.tsx

@@ -1,425 +0,0 @@
-"use client";
-
-import {
-	AlertCircleIcon,
-	BookOpenIcon,
-	CalendarIcon,
-	ExternalLinkIcon,
-	FileTextIcon,
-	UserIcon,
-} from "lucide-react";
-import Image from "next/image";
-import { Component, type ReactNode, useCallback, useState } from "react";
-import { z } from "zod";
-import { Card, CardContent } from "@/components/ui/card";
-import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
-import { cn } from "@/lib/utils";
-
-/**
- * Zod schema for serializable article data (from backend)
- */
-const SerializableArticleSchema = z.object({
-	id: z.string().default("article-unknown"),
-	assetId: z.string().nullish(),
-	kind: z.literal("article").nullish(),
-	title: z.string().default("Untitled Article"),
-	description: z.string().nullish(),
-	content: z.string().nullish(),
-	href: z.string().url().nullish(),
-	domain: z.string().nullish(),
-	author: z.string().nullish(),
-	date: z.string().nullish(),
-	word_count: z.number().nullish(),
-	wordCount: z.number().nullish(),
-	was_truncated: z.boolean().nullish(),
-	wasTruncated: z.boolean().nullish(),
-	error: z.string().nullish(),
-});
-
-/**
- * Serializable article data type (from backend)
- */
-export type SerializableArticle = z.infer<typeof SerializableArticleSchema>;
-
-/**
- * Article component props
- */
-export interface ArticleProps {
-	/** Unique identifier for the article */
-	id: string;
-	/** Asset identifier (usually the URL) */
-	assetId?: string;
-	/** Article title */
-	title: string;
-	/** Brief description or excerpt */
-	description?: string;
-	/** Full content of the article (markdown) */
-	content?: string;
-	/** URL to the original article */
-	href?: string;
-	/** Domain of the article source */
-	domain?: string;
-	/** Author name */
-	author?: string;
-	/** Publication date */
-	date?: string;
-	/** Word count */
-	wordCount?: number;
-	/** Whether content was truncated */
-	wasTruncated?: boolean;
-	/** Optional max width */
-	maxWidth?: string;
-	/** Optional error message */
-	error?: string;
-	/** Optional className */
-	className?: string;
-	/** Response actions */
-	responseActions?: Array<{
-		id: string;
-		label: string;
-		variant?: "default" | "outline";
-	}>;
-	/** Response action handler */
-	onResponseAction?: (actionId: string) => void;
-}
-
-/**
- * Parse and validate serializable article data to ArticleProps
- */
-export function parseSerializableArticle(data: unknown): ArticleProps {
-	const result = SerializableArticleSchema.safeParse(data);
-
-	if (!result.success) {
-		console.warn("Invalid article data:", result.error.issues);
-		// Return fallback with basic info
-		const obj = (data && typeof data === "object" ? data : {}) as Record<string, unknown>;
-		return {
-			id: String(obj.id || "article-unknown"),
-			title: String(obj.title || "Untitled Article"),
-			error: "Failed to parse article data",
-		};
-	}
-
-	const parsed = result.data;
-	return {
-		id: parsed.id,
-		assetId: parsed.assetId,
-		title: parsed.title,
-		description: parsed.description,
-		content: parsed.content,
-		href: parsed.href,
-		domain: parsed.domain,
-		author: parsed.author,
-		date: parsed.date,
-		wordCount: parsed.word_count ?? parsed.wordCount,
-		wasTruncated: parsed.was_truncated ?? parsed.wasTruncated,
-		error: parsed.error,
-	};
-}
-
-/**
- * Format word count for display
- */
-function formatWordCount(count: number): string {
-	if (count >= 1000) {
-		return `${(count / 1000).toFixed(1)}k words`;
-	}
-	return `${count} words`;
-}
-
-/**
- * Favicon component that fetches the site icon via Google's favicon service,
- * falling back to BookOpenIcon on error.
- */
-function SiteFavicon({ domain }: { domain: string }) {
-	const [failed, setFailed] = useState(false);
-
-	if (failed) {
-		return <BookOpenIcon className="size-5 text-primary" />;
-	}
-
-	return (
-		<Image
-			src={`https://www.google.com/s2/favicons?domain=${encodeURIComponent(domain)}&sz=64`}
-			alt={`${domain} favicon`}
-			width={28}
-			height={28}
-			className="size-5 sm:size-7 rounded-sm"
-			onError={() => setFailed(true)}
-			unoptimized
-		/>
-	);
-}
-
-/**
- * Article card component for displaying scraped webpage content
- */
-export function Article({
-	id,
-	title,
-	description,
-	content,
-	href,
-	domain,
-	author,
-	date,
-	wordCount,
-	wasTruncated,
-	maxWidth = "100%",
-	error,
-	className,
-	responseActions,
-	onResponseAction,
-}: ArticleProps) {
-	const handleCardClick = useCallback(() => {
-		if (href) {
-			window.open(href, "_blank", "noopener,noreferrer");
-		}
-	}, [href]);
-
-	// Error state
-	if (error) {
-		return (
-			<Card
-				id={id}
-				className={cn("overflow-hidden border-destructive/20 bg-destructive/5", className)}
-				style={{ maxWidth }}
-			>
-				<CardContent className="p-4">
-					<div className="flex items-center gap-3">
-						<div className="flex size-10 shrink-0 items-center justify-center rounded-lg bg-destructive/10">
-							<AlertCircleIcon className="size-5 text-destructive" />
-						</div>
-						<div className="flex-1 min-w-0">
-							<p className="font-medium text-destructive text-sm">Failed to scrape webpage</p>
-							{href && <p className="text-muted-foreground text-xs mt-0.5 truncate">{href}</p>}
-							<p className="text-muted-foreground text-xs mt-1">{error}</p>
-						</div>
-					</div>
-				</CardContent>
-			</Card>
-		);
-	}
-
-	return (
-		<TooltipProvider>
-			<Card
-				id={id}
-				className={cn(
-					"group relative overflow-hidden transition-all duration-200",
-					"hover:shadow-lg hover:border-primary/20",
-					href && "cursor-pointer",
-					className
-				)}
-				style={{ maxWidth }}
-				onClick={href ? handleCardClick : undefined}
-				role={href ? "link" : undefined}
-				tabIndex={href ? 0 : undefined}
-				onKeyDown={(e) => {
-					if (href && (e.key === "Enter" || e.key === " ")) {
-						e.preventDefault();
-						handleCardClick();
-					}
-				}}
-			>
-				{/* Header */}
-				<CardContent className="p-3 sm:p-4">
-					<div className="flex items-start gap-2.5 sm:gap-3">
-						{/* Favicon / Icon */}
-						{domain ? (
-							<div className="flex size-8 sm:size-10 shrink-0 items-center justify-center">
-								<SiteFavicon domain={domain} />
-							</div>
-						) : (
-							<div className="flex size-8 sm:size-10 shrink-0 items-center justify-center rounded-lg bg-primary/10">
-								<BookOpenIcon className="size-4 sm:size-5 text-primary" />
-							</div>
-						)}
-
-						{/* Content */}
-						<div className="flex-1 min-w-0">
-							{/* Title */}
-							<h3 className="font-semibold text-xs sm:text-sm line-clamp-2 group-hover:text-primary transition-colors">
-								{title}
-							</h3>
-
-							{/* Description */}
-							{description && (
-								<p className="text-muted-foreground text-[10px] sm:text-xs mt-1 line-clamp-2">
-									{description}
-								</p>
-							)}
-
-							{/* Metadata row */}
-							<div className="flex flex-wrap items-center gap-x-2 sm:gap-x-3 gap-y-1 mt-1.5 sm:mt-2 text-[10px] sm:text-xs text-muted-foreground">
-								{domain && (
-									<Tooltip>
-										<TooltipTrigger asChild>
-											<span className="flex items-center gap-1">
-												<ExternalLinkIcon className="size-3" />
-												<span className="truncate max-w-[120px]">{domain}</span>
-											</span>
-										</TooltipTrigger>
-										<TooltipContent>
-											<p>Source: {domain}</p>
-										</TooltipContent>
-									</Tooltip>
-								)}
-
-								{author && (
-									<Tooltip>
-										<TooltipTrigger asChild>
-											<span className="flex items-center gap-1">
-												<UserIcon className="size-3" />
-												<span className="truncate max-w-[100px]">{author}</span>
-											</span>
-										</TooltipTrigger>
-										<TooltipContent>
-											<p>Author: {author}</p>
-										</TooltipContent>
-									</Tooltip>
-								)}
-
-								{date && (
-									<span className="flex items-center gap-1">
-										<CalendarIcon className="size-3" />
-										<span>{date}</span>
-									</span>
-								)}
-
-								{wordCount && (
-									<Tooltip>
-										<TooltipTrigger asChild>
-											<span className="flex items-center gap-1">
-												<FileTextIcon className="size-3" />
-												<span>{formatWordCount(wordCount)}</span>
-												{wasTruncated && <span className="text-warning">(truncated)</span>}
-											</span>
-										</TooltipTrigger>
-										<TooltipContent>
-											<p>
-												{wasTruncated
-													? "Content was truncated due to length"
-													: "Full article content available"}
-											</p>
-										</TooltipContent>
-									</Tooltip>
-								)}
-							</div>
-						</div>
-					</div>
-
-					{/* Response actions */}
-					{responseActions && responseActions.length > 0 && (
-						<div className="flex gap-2 mt-3 pt-3 border-t">
-							{responseActions.map((action) => (
-								<button
-									key={action.id}
-									type="button"
-									onClick={(e) => {
-										e.stopPropagation();
-										onResponseAction?.(action.id);
-									}}
-									className={cn(
-										"px-3 py-1.5 text-xs font-medium rounded-md transition-colors",
-										action.variant === "outline"
-											? "border border-input bg-background hover:bg-accent hover:text-accent-foreground"
-											: "bg-primary text-primary-foreground hover:bg-primary/90"
-									)}
-								>
-									{action.label}
-								</button>
-							))}
-						</div>
-					)}
-				</CardContent>
-			</Card>
-		</TooltipProvider>
-	);
-}
-
-/**
- * Loading state for article component
- */
-export function ArticleLoading({ title = "Loading article..." }: { title?: string }) {
-	return (
-		<Card className="overflow-hidden animate-pulse">
-			<CardContent className="p-4">
-				<div className="flex items-start gap-3">
-					<div className="size-10 rounded-lg bg-muted" />
-					<div className="flex-1 space-y-2">
-						<div className="h-4 bg-muted rounded w-3/4" />
-						<div className="h-3 bg-muted rounded w-full" />
-						<div className="h-3 bg-muted rounded w-1/2" />
-					</div>
-				</div>
-				<p className="text-xs text-muted-foreground mt-3">{title}</p>
-			</CardContent>
-		</Card>
-	);
-}
-
-/**
- * Skeleton for article component
- */
-export function ArticleSkeleton() {
-	return (
-		<Card className="overflow-hidden">
-			<CardContent className="p-4">
-				<div className="flex items-start gap-3 animate-pulse">
-					<div className="size-10 rounded-lg bg-muted" />
-					<div className="flex-1 space-y-2">
-						<div className="h-4 bg-muted rounded w-3/4" />
-						<div className="h-3 bg-muted rounded w-full" />
-						<div className="h-3 bg-muted rounded w-2/3" />
-					</div>
-				</div>
-			</CardContent>
-		</Card>
-	);
-}
-
-/**
- * Error boundary props
- */
-interface ErrorBoundaryProps {
-	children: ReactNode;
-	fallback?: ReactNode;
-}
-
-interface ErrorBoundaryState {
-	hasError: boolean;
-}
-
-/**
- * Error boundary for article component
- */
-export class ArticleErrorBoundary extends Component<ErrorBoundaryProps, ErrorBoundaryState> {
-	constructor(props: ErrorBoundaryProps) {
-		super(props);
-		this.state = { hasError: false };
-	}
-
-	static getDerivedStateFromError(): ErrorBoundaryState {
-		return { hasError: true };
-	}
-
-	render() {
-		if (this.state.hasError) {
-			return (
-				this.props.fallback || (
-					<Card className="overflow-hidden border-destructive/20 bg-destructive/5">
-						<CardContent className="p-4">
-							<div className="flex items-center gap-3">
-								<AlertCircleIcon className="size-5 text-destructive" />
-								<p className="text-sm text-destructive">Failed to render article</p>
-							</div>
-						</CardContent>
-					</Card>
-				)
-			);
-		}
-
-		return this.props.children;
-	}
-}

surfsense_web/components/tool-ui/audio.tsx

@@ -1,12 +1,6 @@
 "use client";
 
-import {
-	DownloadIcon,
-	PauseIcon,
-	PlayIcon,
-	Volume2Icon,
-	VolumeXIcon,
-} from "lucide-react";
+import { DownloadIcon, PauseIcon, PlayIcon, Volume2Icon, VolumeXIcon } from "lucide-react";
 import { useCallback, useEffect, useRef, useState } from "react";
 import { Button } from "@/components/ui/button";
 import { Slider } from "@/components/ui/slider";

surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -457,42 +457,42 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const CreateConfluencePageToolUI = makeAssistantToolUI<
+export const CreateConfluencePageToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{ title: string; content?: string; space_id?: string },
 	CreateConfluencePageResult
->({
-	toolName: "create_confluence_page",
-	render: function CreateConfluencePageUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { CornerDownLeftIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
@@ -396,44 +396,43 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const DeleteConfluencePageToolUI = makeAssistantToolUI<
+export const DeleteConfluencePageToolUI = ({
+	result,
+}: ToolCallMessagePartProps<
 	{ page_title_or_id: string; delete_from_kb?: boolean },
 	DeleteConfluencePageResult
->({
-	toolName: "delete_confluence_page",
-	render: function DeleteConfluencePageUI({ result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isWarningResult(result)) return <WarningCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isWarningResult(result)) return <WarningCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
@@ -493,47 +493,47 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const UpdateConfluencePageToolUI = makeAssistantToolUI<
+export const UpdateConfluencePageToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{
 		page_title_or_id: string;
 		new_title?: string;
 		new_content?: string;
 	},
 	UpdateConfluencePageResult
->({
-	toolName: "update_confluence_page",
-	render: function UpdateConfluencePageUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/deepagent-thinking.tsx

@@ -1,406 +0,0 @@
-"use client";
-
-import { makeAssistantToolUI } from "@assistant-ui/react";
-import { Brain, CheckCircle2, Loader2, Search, Sparkles } from "lucide-react";
-import type { FC, ReactNode } from "react";
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
-import { z } from "zod";
-import {
-	ChainOfThought,
-	ChainOfThoughtContent,
-	ChainOfThoughtItem,
-	ChainOfThoughtStep,
-	ChainOfThoughtTrigger,
-} from "@/components/prompt-kit/chain-of-thought";
-import { cn } from "@/lib/utils";
-
-// ============================================================================
-// Constants
-// ============================================================================
-
-/** Step status values */
-const STEP_STATUS = {
-	PENDING: "pending",
-	IN_PROGRESS: "in_progress",
-	COMPLETED: "completed",
-} as const;
-
-/** Agent thinking status values */
-const THINKING_STATUS = {
-	THINKING: "thinking",
-	SEARCHING: "searching",
-	SYNTHESIZING: "synthesizing",
-	COMPLETED: "completed",
-} as const;
-
-/** Keywords for icon detection */
-const STEP_KEYWORDS = {
-	SEARCH: ["search", "knowledge"] as const,
-	ANALYSIS: ["analy", "understand"] as const,
-} as const;
-
-/** Icon size class */
-const ICON_SIZE_CLASS = "size-4" as const;
-
-/** Status text mapping */
-const STATUS_TEXT_MAP: Record<string, string> = {
-	[THINKING_STATUS.SEARCHING]: "Searching knowledge base...",
-	[THINKING_STATUS.SYNTHESIZING]: "Synthesizing response...",
-	[THINKING_STATUS.THINKING]: "Thinking...",
-} as const;
-
-// ============================================================================
-// Type Definitions
-// ============================================================================
-
-type StepStatus = (typeof STEP_STATUS)[keyof typeof STEP_STATUS];
-type ThinkingStatus = (typeof THINKING_STATUS)[keyof typeof THINKING_STATUS];
-
-// ============================================================================
-// Zod Schemas
-// ============================================================================
-
-const ThinkingStepSchema = z.object({
-	id: z.string(),
-	title: z.string(),
-	items: z.array(z.string()).default([]),
-	status: z
-		.enum([STEP_STATUS.PENDING, STEP_STATUS.IN_PROGRESS, STEP_STATUS.COMPLETED])
-		.default(STEP_STATUS.PENDING),
-});
-
-const DeepAgentThinkingArgsSchema = z.object({
-	query: z.string().nullish(),
-	context: z.string().nullish(),
-});
-
-const DeepAgentThinkingResultSchema = z.object({
-	steps: z.array(ThinkingStepSchema).nullish(),
-	status: z
-		.enum([
-			THINKING_STATUS.THINKING,
-			THINKING_STATUS.SEARCHING,
-			THINKING_STATUS.SYNTHESIZING,
-			THINKING_STATUS.COMPLETED,
-		])
-		.nullish(),
-	summary: z.string().nullish(),
-});
-
-/** Types derived from Zod schemas */
-type ThinkingStep = z.infer<typeof ThinkingStepSchema>;
-type DeepAgentThinkingArgs = z.infer<typeof DeepAgentThinkingArgsSchema>;
-type DeepAgentThinkingResult = z.infer<typeof DeepAgentThinkingResultSchema>;
-
-// ============================================================================
-// Parser Functions
-// ============================================================================
-
-/** Default fallback step when parsing fails */
-const DEFAULT_FALLBACK_STEP: ThinkingStep = {
-	id: "unknown",
-	title: "Processing...",
-	items: [],
-	status: STEP_STATUS.PENDING,
-} as const;
-
-/**
- * Parse and validate a single thinking step
- */
-export function parseThinkingStep(data: unknown): ThinkingStep {
-	const result = ThinkingStepSchema.safeParse(data);
-	if (!result.success) {
-		console.warn("Invalid thinking step data:", result.error.issues);
-		return DEFAULT_FALLBACK_STEP;
-	}
-	return result.data;
-}
-
-/**
- * Parse and validate thinking result
- */
-export function parseThinkingResult(data: unknown): DeepAgentThinkingResult {
-	const result = DeepAgentThinkingResultSchema.safeParse(data);
-	if (!result.success) {
-		console.warn("Invalid thinking result data:", result.error.issues);
-		return {};
-	}
-	return result.data;
-}
-
-// ============================================================================
-// Icon Utilities
-// ============================================================================
-
-/**
- * Check if title contains any of the keywords
- */
-function titleContainsKeywords(title: string, keywords: readonly string[]): boolean {
-	const titleLower = title.toLowerCase();
-	return keywords.some((keyword) => titleLower.includes(keyword));
-}
-
-/**
- * Get icon based on step status and title
- */
-function getStepIcon(status: StepStatus, title: string): ReactNode {
-	if (status === STEP_STATUS.IN_PROGRESS) {
-		return <Loader2 className={cn(ICON_SIZE_CLASS, "animate-spin text-primary")} />;
-	}
-
-	if (status === STEP_STATUS.COMPLETED) {
-		return <CheckCircle2 className={cn(ICON_SIZE_CLASS, "text-emerald-500")} />;
-	}
-
-	// Default icons based on step type keywords
-	if (titleContainsKeywords(title, STEP_KEYWORDS.SEARCH)) {
-		return <Search className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
-	}
-
-	if (titleContainsKeywords(title, STEP_KEYWORDS.ANALYSIS)) {
-		return <Brain className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
-	}
-
-	return <Sparkles className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
-}
-
-// ============================================================================
-// Sub-Components
-// ============================================================================
-
-interface ThinkingStepDisplayProps {
-	step: ThinkingStep;
-	isOpen: boolean;
-	onToggle: () => void;
-}
-
-/**
- * Component to display a single thinking step with controlled open state
- */
-const ThinkingStepDisplay: FC<ThinkingStepDisplayProps> = ({ step, isOpen, onToggle }) => {
-	const icon = useMemo(() => getStepIcon(step.status, step.title), [step.status, step.title]);
-
-	const isInProgress = step.status === STEP_STATUS.IN_PROGRESS;
-	const isCompleted = step.status === STEP_STATUS.COMPLETED;
-
-	return (
-		<ChainOfThoughtStep open={isOpen} onOpenChange={onToggle}>
-			<ChainOfThoughtTrigger
-				leftIcon={icon}
-				swapIconOnHover={!isInProgress}
-				className={cn(
-					isInProgress && "text-foreground font-medium",
-					isCompleted && "text-muted-foreground"
-				)}
-			>
-				{step.title}
-			</ChainOfThoughtTrigger>
-			<ChainOfThoughtContent>
-				{step.items.map((item, index) => (
-					<ChainOfThoughtItem key={`${step.id}-item-${index}`}>{item}</ChainOfThoughtItem>
-				))}
-			</ChainOfThoughtContent>
-		</ChainOfThoughtStep>
-	);
-};
-
-interface ThinkingLoadingStateProps {
-	status?: ThinkingStatus | string;
-}
-
-/**
- * Loading state with animated thinking indicator
- */
-const ThinkingLoadingState: FC<ThinkingLoadingStateProps> = ({ status }) => {
-	const statusText = useMemo(() => {
-		if (status && status in STATUS_TEXT_MAP) {
-			return STATUS_TEXT_MAP[status];
-		}
-		return STATUS_TEXT_MAP[THINKING_STATUS.THINKING];
-	}, [status]);
-
-	return (
-		<div className="my-3 flex items-center gap-2 rounded-lg border border-border/50 bg-muted/30 px-4 py-3">
-			<div className="relative">
-				<Brain className="size-5 text-primary" />
-				<span className="absolute -right-0.5 -top-0.5 flex size-2">
-					<span className="absolute inline-flex size-full animate-ping rounded-full bg-primary/60" />
-					<span className="relative inline-flex size-2 rounded-full bg-primary" />
-				</span>
-			</div>
-			<span className="text-sm text-muted-foreground">{statusText}</span>
-		</div>
-	);
-};
-
-interface SmartChainOfThoughtProps {
-	steps: ThinkingStep[];
-}
-
-/** Type for tracking step override states */
-type StepOverrides = Record<string, boolean>;
-
-/** Type for tracking step status history */
-type StepStatusHistory = Record<string, StepStatus>;
-
-/**
- * Smart chain of thought renderer with state management
- */
-const SmartChainOfThought: FC<SmartChainOfThoughtProps> = ({ steps }) => {
-	// Track which steps the user has manually toggled
-	const [manualOverrides, setManualOverrides] = useState<StepOverrides>({});
-	// Track previous step statuses to detect changes
-	const prevStatusesRef = useRef<StepStatusHistory>({});
-
-	// Clear manual overrides when a step's status changes
-	useEffect(() => {
-		const currentStatuses: StepStatusHistory = {};
-		steps.forEach((step) => {
-			currentStatuses[step.id] = step.status;
-			// If status changed, clear any manual override for this step
-			const prevStatus = prevStatusesRef.current[step.id];
-			if (prevStatus && prevStatus !== step.status) {
-				setManualOverrides((prev) => {
-					const next = { ...prev };
-					delete next[step.id];
-					return next;
-				});
-			}
-		});
-		prevStatusesRef.current = currentStatuses;
-	}, [steps]);
-
-	const getStepOpenState = useCallback(
-		(step: ThinkingStep): boolean => {
-			// If user has manually toggled, respect that
-			if (manualOverrides[step.id] !== undefined) {
-				return manualOverrides[step.id];
-			}
-			// Auto behavior: open if in progress
-			if (step.status === STEP_STATUS.IN_PROGRESS) {
-				return true;
-			}
-			// Default: collapsed (all steps collapse when processing is done)
-			return false;
-		},
-		[manualOverrides]
-	);
-
-	const handleToggle = useCallback((stepId: string, currentOpen: boolean) => {
-		setManualOverrides((prev) => ({
-			...prev,
-			[stepId]: !currentOpen,
-		}));
-	}, []);
-
-	return (
-		<ChainOfThought>
-			{steps.map((step) => {
-				const isOpen = getStepOpenState(step);
-				return (
-					<ThinkingStepDisplay
-						key={step.id}
-						step={step}
-						isOpen={isOpen}
-						onToggle={() => handleToggle(step.id, isOpen)}
-					/>
-				);
-			})}
-		</ChainOfThought>
-	);
-};
-
-/**
- * DeepAgent Thinking Tool UI Component
- *
- * This component displays the agent's chain-of-thought reasoning
- * when the deepagent is processing a query. It shows thinking steps
- * in a collapsible, hierarchical format.
- */
-export const DeepAgentThinkingToolUI = makeAssistantToolUI<
-	DeepAgentThinkingArgs,
-	DeepAgentThinkingResult
->({
-	toolName: "deepagent_thinking",
-	render: function DeepAgentThinkingUI({ result, status }) {
-		// Loading state - tool is still running
-		if (status.type === "running" || status.type === "requires-action") {
-			return <ThinkingLoadingState status={result?.status ?? undefined} />;
-		}
-
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return null; // Don't show anything if cancelled
-			}
-			if (status.reason === "error") {
-				return null; // Don't show error for thinking - it's not critical
-			}
-		}
-
-		// No result or no steps - don't render anything
-		if (!result?.steps || result.steps.length === 0) {
-			return null;
-		}
-
-		// Render the chain of thought
-		return (
-			<div className="my-3 w-full">
-				<SmartChainOfThought steps={result.steps} />
-			</div>
-		);
-	},
-});
-
-// ============================================================================
-// Public Components
-// ============================================================================
-
-export interface InlineThinkingDisplayProps {
-	/** The thinking steps to display */
-	steps: ThinkingStep[];
-	/** Whether content is currently streaming */
-	isStreaming?: boolean;
-	/** Additional CSS class names */
-	className?: string;
-}
-
-/**
- * Inline Thinking Display Component
- *
- * A simpler version that can be used inline with the message content
- * for displaying reasoning without the full tool UI infrastructure.
- */
-export const InlineThinkingDisplay: FC<InlineThinkingDisplayProps> = ({
-	steps,
-	isStreaming = false,
-	className,
-}) => {
-	if (steps.length === 0 && !isStreaming) {
-		return null;
-	}
-
-	return (
-		<div className={cn("my-3 w-full", className)}>
-			{isStreaming && steps.length === 0 ? (
-				<ThinkingLoadingState />
-			) : (
-				<SmartChainOfThought steps={steps} />
-			)}
-		</div>
-	);
-};
-
-// ============================================================================
-// Exports
-// ============================================================================
-
-export type {
-	ThinkingStep,
-	DeepAgentThinkingArgs,
-	DeepAgentThinkingResult,
-	StepStatus,
-	ThinkingStatus,
-};
-
-export { STEP_STATUS, THINKING_STATUS };

surfsense_web/components/tool-ui/display-image.tsx

@@ -1,165 +0,0 @@
-"use client";
-
-import { makeAssistantToolUI } from "@assistant-ui/react";
-import { AlertCircleIcon, ImageIcon } from "lucide-react";
-import { z } from "zod";
-import {
-	Image,
-	ImageErrorBoundary,
-	ImageLoading,
-	parseSerializableImage,
-} from "@/components/tool-ui/image";
-
-// ============================================================================
-// Zod Schemas
-// ============================================================================
-
-/**
- * Schema for display_image tool arguments
- */
-const DisplayImageArgsSchema = z.object({
-	src: z.string(),
-	alt: z.string().nullish(),
-	title: z.string().nullish(),
-	description: z.string().nullish(),
-});
-
-/**
- * Schema for display_image tool result
- */
-const DisplayImageResultSchema = z.object({
-	id: z.string(),
-	assetId: z.string(),
-	src: z.string(),
-	alt: z.string().nullish(),
-	title: z.string().nullish(),
-	description: z.string().nullish(),
-	domain: z.string().nullish(),
-	ratio: z.string().nullish(),
-	error: z.string().nullish(),
-});
-
-// ============================================================================
-// Types
-// ============================================================================
-
-type DisplayImageArgs = z.infer<typeof DisplayImageArgsSchema>;
-type DisplayImageResult = z.infer<typeof DisplayImageResultSchema>;
-
-/**
- * Error state component shown when image display fails
- */
-function ImageErrorState({ src, error }: { src: string; error: string }) {
-	return (
-		<div className="my-4 overflow-hidden rounded-xl border border-destructive/20 bg-destructive/5 p-4 max-w-md">
-			<div className="flex items-center gap-4">
-				<div className="flex size-12 shrink-0 items-center justify-center rounded-lg bg-destructive/10">
-					<AlertCircleIcon className="size-6 text-destructive" />
-				</div>
-				<div className="flex-1 min-w-0">
-					<p className="font-medium text-destructive text-sm">Failed to display image</p>
-					<p className="text-muted-foreground text-xs mt-0.5 truncate">{src}</p>
-					<p className="text-muted-foreground text-xs mt-1">{error}</p>
-				</div>
-			</div>
-		</div>
-	);
-}
-
-/**
- * Cancelled state component
- */
-function ImageCancelledState({ src }: { src: string }) {
-	return (
-		<div className="my-4 rounded-xl border border-muted p-4 text-muted-foreground max-w-md">
-			<p className="flex items-center gap-2">
-				<ImageIcon className="size-4" />
-				<span className="line-through truncate">Image: {src}</span>
-			</p>
-		</div>
-	);
-}
-
-/**
- * Parsed Image component with error handling
- * Note: Image component has built-in click handling via href/src,
- * so no additional responseActions needed.
- */
-function ParsedImage({ result }: { result: unknown }) {
-	const image = parseSerializableImage(result);
-
-	return <Image {...image} maxWidth="512px" />;
-}
-
-/**
- * Display Image Tool UI Component
- *
- * This component is registered with assistant-ui to render an image
- * when the display_image tool is called by the agent.
- *
- * It displays images with:
- * - Title and description
- * - Source attribution
- * - Hover overlay effects
- * - Click to open full size
- */
-export const DisplayImageToolUI = makeAssistantToolUI<DisplayImageArgs, DisplayImageResult>({
-	toolName: "display_image",
-	render: function DisplayImageUI({ args, result, status }) {
-		const src = args.src || "Unknown";
-
-		// Loading state - tool is still running
-		if (status.type === "running" || status.type === "requires-action") {
-			return (
-				<div className="my-4">
-					<ImageLoading title={`Loading image...`} />
-				</div>
-			);
-		}
-
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return <ImageCancelledState src={src} />;
-			}
-			if (status.reason === "error") {
-				return (
-					<ImageErrorState
-						src={src}
-						error={typeof status.error === "string" ? status.error : "An error occurred"}
-					/>
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return (
-				<div className="my-4">
-					<ImageLoading title="Preparing image..." />
-				</div>
-			);
-		}
-
-		// Error result from the tool
-		if (result.error) {
-			return <ImageErrorState src={src} error={result.error} />;
-		}
-
-		// Success - render the image
-		return (
-			<div className="my-4">
-				<ImageErrorBoundary>
-					<ParsedImage result={result} />
-				</ImageErrorBoundary>
-			</div>
-		);
-	},
-});
-
-export {
-	DisplayImageArgsSchema,
-	DisplayImageResultSchema,
-	type DisplayImageArgs,
-	type DisplayImageResult,
-};

surfsense_web/components/tool-ui/generate-image.tsx

@@ -0,0 +1,142 @@
+"use client";
+
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
+import { AlertCircleIcon, ImageIcon } from "lucide-react";
+import { z } from "zod";
+import {
+	Image,
+	ImageErrorBoundary,
+	ImageLoading,
+	parseSerializableImage,
+} from "@/components/tool-ui/image";
+
+const GenerateImageArgsSchema = z.object({
+	prompt: z.string(),
+	n: z.number().nullish(),
+});
+
+const GenerateImageResultSchema = z.object({
+	id: z.string(),
+	assetId: z.string(),
+	src: z.string(),
+	alt: z.string().nullish(),
+	title: z.string().nullish(),
+	description: z.string().nullish(),
+	domain: z.string().nullish(),
+	ratio: z.string().nullish(),
+	generated: z.boolean().nullish(),
+	prompt: z.string().nullish(),
+	image_count: z.number().nullish(),
+	error: z.string().nullish(),
+});
+
+type GenerateImageArgs = z.infer<typeof GenerateImageArgsSchema>;
+type GenerateImageResult = z.infer<typeof GenerateImageResultSchema>;
+
+function ImageErrorState({ prompt, error }: { prompt: string; error: string }) {
+	return (
+		<div className="my-4 overflow-hidden rounded-xl border border-destructive/20 bg-destructive/5 p-4 max-w-md">
+			<div className="flex items-center gap-4">
+				<div className="flex size-12 shrink-0 items-center justify-center rounded-lg bg-destructive/10">
+					<AlertCircleIcon className="size-6 text-destructive" />
+				</div>
+				<div className="flex-1 min-w-0">
+					<p className="font-medium text-destructive text-sm">Image generation failed</p>
+					<p className="text-muted-foreground text-xs mt-0.5 truncate">{prompt}</p>
+					<p className="text-muted-foreground text-xs mt-1">{error}</p>
+				</div>
+			</div>
+		</div>
+	);
+}
+
+function ImageCancelledState({ prompt }: { prompt: string }) {
+	return (
+		<div className="my-4 rounded-xl border border-muted p-4 text-muted-foreground max-w-md">
+			<p className="flex items-center gap-2">
+				<ImageIcon className="size-4" />
+				<span className="line-through truncate">Generate: {prompt}</span>
+			</p>
+		</div>
+	);
+}
+
+function ParsedImage({ result }: { result: unknown }) {
+	const image = parseSerializableImage(result);
+	return (
+		<Image
+			id={image.id}
+			assetId={image.assetId}
+			src={image.src}
+			alt={image.alt}
+			title={image.title ?? undefined}
+			description={image.description ?? undefined}
+			href={image.href ?? undefined}
+			domain={image.domain ?? undefined}
+			ratio={image.ratio ?? undefined}
+			source={image.source ?? undefined}
+			maxWidth="512px"
+		/>
+	);
+}
+
+/**
+ * Tool UI for generate_image — renders the generated image directly
+ * from the tool result directly.
+ */
+export const GenerateImageToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps<GenerateImageArgs, GenerateImageResult>) => {
+	const prompt = args.prompt || "Generating image...";
+
+	if (status.type === "running" || status.type === "requires-action") {
+		return (
+			<div className="my-4">
+				<ImageLoading title="Generating image" />
+			</div>
+		);
+	}
+
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
+			return <ImageCancelledState prompt={prompt} />;
+		}
+		if (status.reason === "error") {
+			return (
+				<ImageErrorState
+					prompt={prompt}
+					error={typeof status.error === "string" ? status.error : "An error occurred"}
+				/>
+			);
+		}
+	}
+
+	if (!result) {
+		return (
+			<div className="my-4">
+				<ImageLoading title="Loading" />
+			</div>
+		);
+	}
+
+	if (result.error) {
+		return <ImageErrorState prompt={prompt} error={result.error} />;
+	}
+
+	return (
+		<div className="my-4">
+			<ImageErrorBoundary>
+				<ParsedImage result={result} />
+			</ImageErrorBoundary>
+		</div>
+	);
+};
+
+export {
+	GenerateImageArgsSchema,
+	GenerateImageResultSchema,
+	type GenerateImageArgs,
+	type GenerateImageResult,
+};

surfsense_web/components/tool-ui/generate-podcast.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useParams, usePathname } from "next/navigation";
 import { useCallback, useEffect, useRef, useState } from "react";
 import { z } from "zod";
@@ -372,95 +372,91 @@ function PodcastStatusPoller({ podcastId, title }: { podcastId: number; title: s
  *
  * It polls for task completion and auto-updates when the podcast is ready.
  */
-export const GeneratePodcastToolUI = makeAssistantToolUI<
-	GeneratePodcastArgs,
-	GeneratePodcastResult
->({
-	toolName: "generate_podcast",
-	render: function GeneratePodcastUI({ args, result, status }) {
-		const title = args.podcast_title || "SurfSense Podcast";
+export const GeneratePodcastToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps<GeneratePodcastArgs, GeneratePodcastResult>) => {
+	const title = args.podcast_title || "SurfSense Podcast";
 
-		// Loading state - tool is still running (agent processing)
-		if (status.type === "running" || status.type === "requires-action") {
-			return <PodcastGeneratingState title={title} />;
-		}
+	// Loading state - tool is still running (agent processing)
+	if (status.type === "running" || status.type === "requires-action") {
+		return <PodcastGeneratingState title={title} />;
+	}
 
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return (
-					<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
-						<div className="px-5 pt-5 pb-4">
-							<p className="text-sm font-semibold text-muted-foreground">Podcast Cancelled</p>
-							<p className="text-xs text-muted-foreground mt-0.5">
-								Podcast generation was cancelled
-							</p>
-						</div>
-					</div>
-				);
-			}
-			if (status.reason === "error") {
-				return (
-					<PodcastErrorState
-						title={title}
-						error={typeof status.error === "string" ? status.error : "An error occurred"}
-					/>
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return <PodcastGeneratingState title={title} />;
-		}
-
-		// Failed result (new: "failed", legacy: "error")
-		if (result.status === "failed" || result.status === "error") {
-			return <PodcastErrorState title={title} error={result.error || "Generation failed"} />;
-		}
-
-		// Already generating - show simple warning, don't create another poller
-		// The FIRST tool call will display the podcast when ready
-		// (new: "generating", legacy: "already_generating")
-		if (result.status === "generating" || result.status === "already_generating") {
+	// Incomplete/cancelled state
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
 			return (
 				<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
 					<div className="px-5 pt-5 pb-4">
-						<p className="text-sm font-semibold text-foreground">Podcast already in progress</p>
-						<p className="text-xs text-muted-foreground mt-0.5">
-							Please wait for the current podcast to complete.
-						</p>
+						<p className="text-sm font-semibold text-muted-foreground">Podcast Cancelled</p>
+						<p className="text-xs text-muted-foreground mt-0.5">Podcast generation was cancelled</p>
 					</div>
 				</div>
 			);
 		}
-
-		// Pending - poll for completion (new: "pending" with podcast_id)
-		if (result.status === "pending" && result.podcast_id) {
-			return <PodcastStatusPoller podcastId={result.podcast_id} title={result.title || title} />;
-		}
-
-		// Ready with podcast_id (new: "ready", legacy: "success")
-		if ((result.status === "ready" || result.status === "success") && result.podcast_id) {
-			return <PodcastPlayer podcastId={result.podcast_id} title={result.title || title} />;
-		}
-
-		// Legacy: old chats with Celery task_id (status: "processing" or "success" without podcast_id)
-		// These can't be recovered since the old task polling endpoint no longer exists
-		if (result.task_id && !result.podcast_id) {
+		if (status.reason === "error") {
 			return (
-				<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
-					<div className="px-5 pt-5 pb-4">
-						<p className="text-sm font-semibold text-muted-foreground">Podcast Unavailable</p>
-						<p className="text-xs text-muted-foreground mt-0.5">
-							This podcast was generated with an older version. Please generate a new one.
-						</p>
-					</div>
-				</div>
+				<PodcastErrorState
+					title={title}
+					error={typeof status.error === "string" ? status.error : "An error occurred"}
+				/>
 			);
 		}
+	}
 
-		// Fallback - missing required data
-		return <PodcastErrorState title={title} error="Missing podcast ID" />;
-	},
-});
+	// No result yet
+	if (!result) {
+		return <PodcastGeneratingState title={title} />;
+	}
+
+	// Failed result (new: "failed", legacy: "error")
+	if (result.status === "failed" || result.status === "error") {
+		return <PodcastErrorState title={title} error={result.error || "Generation failed"} />;
+	}
+
+	// Already generating - show simple warning, don't create another poller
+	// The FIRST tool call will display the podcast when ready
+	// (new: "generating", legacy: "already_generating")
+	if (result.status === "generating" || result.status === "already_generating") {
+		return (
+			<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+				<div className="px-5 pt-5 pb-4">
+					<p className="text-sm font-semibold text-foreground">Podcast already in progress</p>
+					<p className="text-xs text-muted-foreground mt-0.5">
+						Please wait for the current podcast to complete.
+					</p>
+				</div>
+			</div>
+		);
+	}
+
+	// Pending - poll for completion (new: "pending" with podcast_id)
+	if (result.status === "pending" && result.podcast_id) {
+		return <PodcastStatusPoller podcastId={result.podcast_id} title={result.title || title} />;
+	}
+
+	// Ready with podcast_id (new: "ready", legacy: "success")
+	if ((result.status === "ready" || result.status === "success") && result.podcast_id) {
+		return <PodcastPlayer podcastId={result.podcast_id} title={result.title || title} />;
+	}
+
+	// Legacy: old chats with Celery task_id (status: "processing" or "success" without podcast_id)
+	// These can't be recovered since the old task polling endpoint no longer exists
+	if (result.task_id && !result.podcast_id) {
+		return (
+			<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+				<div className="px-5 pt-5 pb-4">
+					<p className="text-sm font-semibold text-muted-foreground">Podcast Unavailable</p>
+					<p className="text-xs text-muted-foreground mt-0.5">
+						This podcast was generated with an older version. Please generate a new one.
+					</p>
+				</div>
+			</div>
+		);
+	}
+
+	// Fallback - missing required data
+	return <PodcastErrorState title={title} error="Missing podcast ID" />;
+};

surfsense_web/components/tool-ui/generate-report.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useAtomValue, useSetAtom } from "jotai";
 import { Dot } from "lucide-react";
 import { useParams, usePathname } from "next/navigation";
@@ -273,64 +273,62 @@ function ReportCard({
  * Generate Report Tool UI — renders custom UI inline in chat
  * when the generate_report tool is called by the agent.
  */
-export const GenerateReportToolUI = makeAssistantToolUI<GenerateReportArgs, GenerateReportResult>({
-	toolName: "generate_report",
-	render: function GenerateReportUI({ args, result, status }) {
-		const params = useParams();
-		const pathname = usePathname();
-		const isPublicRoute = pathname?.startsWith("/public/");
-		const shareToken = isPublicRoute && typeof params?.token === "string" ? params.token : null;
+export const GenerateReportToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps<GenerateReportArgs, GenerateReportResult>) => {
+	const params = useParams();
+	const pathname = usePathname();
+	const isPublicRoute = pathname?.startsWith("/public/");
+	const shareToken = isPublicRoute && typeof params?.token === "string" ? params.token : null;
 
-		const topic = args.topic || "Report";
+	const topic = args.topic || "Report";
 
-		const sawRunningRef = useRef(false);
-		if (status.type === "running" || status.type === "requires-action") {
-			sawRunningRef.current = true;
+	const sawRunningRef = useRef(false);
+	if (status.type === "running" || status.type === "requires-action") {
+		sawRunningRef.current = true;
+	}
+
+	if (status.type === "running" || status.type === "requires-action") {
+		return <ReportGeneratingState topic={topic} />;
+	}
+
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
+			return <ReportCancelledState />;
 		}
-
-		if (status.type === "running" || status.type === "requires-action") {
-			return <ReportGeneratingState topic={topic} />;
-		}
-
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return <ReportCancelledState />;
-			}
-			if (status.reason === "error") {
-				return (
-					<ReportErrorState
-						title={topic}
-						error={typeof status.error === "string" ? status.error : "An error occurred"}
-					/>
-				);
-			}
-		}
-
-		if (!result) {
-			return <ReportGeneratingState topic={topic} />;
-		}
-
-		if (result.status === "failed") {
+		if (status.reason === "error") {
 			return (
 				<ReportErrorState
-					title={result.title || topic}
-					error={result.error || "Generation failed"}
+					title={topic}
+					error={typeof status.error === "string" ? status.error : "An error occurred"}
 				/>
 			);
 		}
+	}
 
-		if (result.status === "ready" && result.report_id) {
-			return (
-				<ReportCard
-					reportId={result.report_id}
-					title={result.title || topic}
-					wordCount={result.word_count ?? undefined}
-					shareToken={shareToken}
-					autoOpen={sawRunningRef.current}
-				/>
-			);
-		}
+	if (!result) {
+		return <ReportGeneratingState topic={topic} />;
+	}
 
-		return <ReportErrorState title={topic} error="Missing report ID" />;
-	},
-});
+	if (result.status === "failed") {
+		return (
+			<ReportErrorState title={result.title || topic} error={result.error || "Generation failed"} />
+		);
+	}
+
+	if (result.status === "ready" && result.report_id) {
+		return (
+			<ReportCard
+				reportId={result.report_id}
+				title={result.title || topic}
+				wordCount={result.word_count ?? undefined}
+				shareToken={shareToken}
+				autoOpen={sawRunningRef.current}
+			/>
+		);
+	}
+
+	return <ReportErrorState title={topic} error="Missing report ID" />;
+};

surfsense_web/components/tool-ui/gmail/create-draft.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen, UserIcon, UsersIcon } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -466,42 +466,42 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const CreateGmailDraftToolUI = makeAssistantToolUI<
+export const CreateGmailDraftToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{ to: string; subject: string; body: string; cc?: string; bcc?: string },
 	CreateGmailDraftResult
->({
-	toolName: "create_gmail_draft",
-	render: function CreateGmailDraftUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/gmail/send-email.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, MailIcon, Pen, UserIcon, UsersIcon } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -464,42 +464,42 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const SendGmailEmailToolUI = makeAssistantToolUI<
+export const SendGmailEmailToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{ to: string; subject: string; body: string; cc?: string; bcc?: string },
 	SendGmailEmailResult
->({
-	toolName: "send_gmail_email",
-	render: function SendGmailEmailUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/gmail/trash-email.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { CalendarIcon, CornerDownLeftIcon, MailIcon, UserIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
@@ -379,43 +379,42 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const TrashGmailEmailToolUI = makeAssistantToolUI<
+export const TrashGmailEmailToolUI = ({
+	result,
+}: ToolCallMessagePartProps<
 	{ email_subject_or_id: string; delete_from_kb?: boolean },
 	TrashGmailEmailResult
->({
-	toolName: "trash_gmail_email",
-	render: function TrashGmailEmailUI({ result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/gmail/update-draft.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, MailIcon, Pen, UserIcon, UsersIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
@@ -508,7 +508,10 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const UpdateGmailDraftToolUI = makeAssistantToolUI<
+export const UpdateGmailDraftToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{
 		draft_subject_or_id: string;
 		body: string;
@@ -518,42 +521,39 @@ export const UpdateGmailDraftToolUI = makeAssistantToolUI<
 		bcc?: string;
 	},
 	UpdateGmailDraftResult
->({
-	toolName: "update_gmail_draft",
-	render: function UpdateGmailDraftUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", {
-								detail: { decisions: [decision] },
-							})
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", {
+							detail: { decisions: [decision] },
+						})
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/google-calendar/create-event.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { ClockIcon, CornerDownLeftIcon, GlobeIcon, MapPinIcon, Pen, UsersIcon } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -606,7 +606,10 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const CreateCalendarEventToolUI = makeAssistantToolUI<
+export const CreateCalendarEventToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{
 		summary: string;
 		start_datetime: string;
@@ -616,39 +619,36 @@ export const CreateCalendarEventToolUI = makeAssistantToolUI<
 		attendees?: string[];
 	},
 	CreateCalendarEventResult
->({
-	toolName: "create_calendar_event",
-	render: function CreateCalendarEventUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/google-calendar/delete-event.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { CalendarIcon, ClockIcon, CornerDownLeftIcon, MapPinIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
@@ -431,44 +431,43 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const DeleteCalendarEventToolUI = makeAssistantToolUI<
+export const DeleteCalendarEventToolUI = ({
+	result,
+}: ToolCallMessagePartProps<
 	{ event_title_or_id: string; delete_from_kb?: boolean },
 	DeleteCalendarEventResult
->({
-	toolName: "delete_calendar_event",
-	render: function DeleteCalendarEventUI({ result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isWarningResult(result)) return <WarningCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isWarningResult(result)) return <WarningCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/google-calendar/update-event.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import {
 	ArrowRightIcon,
@@ -653,7 +653,10 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const UpdateCalendarEventToolUI = makeAssistantToolUI<
+export const UpdateCalendarEventToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{
 		event_ref: string;
 		new_summary?: string;
@@ -664,40 +667,37 @@ export const UpdateCalendarEventToolUI = makeAssistantToolUI<
 		new_attendees?: string[];
 	},
 	UpdateCalendarEventResult
->({
-	toolName: "update_calendar_event",
-	render: function UpdateCalendarEventUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/google-drive/create-file.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, FileIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -492,44 +492,44 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const CreateGoogleDriveFileToolUI = makeAssistantToolUI<
+export const CreateGoogleDriveFileToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{ name: string; file_type: string; content?: string },
 	CreateGoogleDriveFileResult
->({
-	toolName: "create_google_drive_file",
-	render: function CreateGoogleDriveFileUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
 
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
 
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/google-drive/trash-file.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { CornerDownLeftIcon, InfoIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
@@ -410,46 +410,45 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const DeleteGoogleDriveFileToolUI = makeAssistantToolUI<
+export const DeleteGoogleDriveFileToolUI = ({
+	result,
+}: ToolCallMessagePartProps<
 	{ file_name: string; delete_from_kb?: boolean },
 	DeleteGoogleDriveFileResult
->({
-	toolName: "delete_google_drive_file",
-	render: function DeleteGoogleDriveFileUI({ result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
 
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isWarningResult(result)) return <WarningCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isWarningResult(result)) return <WarningCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/image/index.tsx

@@ -4,9 +4,9 @@ import { ExternalLinkIcon, ImageIcon, SparklesIcon } from "lucide-react";
 import NextImage from "next/image";
 import { Component, type ReactNode, useState } from "react";
 import { z } from "zod";
+import { TextShimmerLoader } from "@/components/prompt-kit/loader";
 import { Badge } from "@/components/ui/badge";
 import { Card } from "@/components/ui/card";
-import { Spinner } from "@/components/ui/spinner";
 import { cn } from "@/lib/utils";
 
 /**
@@ -145,7 +145,7 @@ export class ImageErrorBoundary extends Component<
 	render() {
 		if (this.state.hasError) {
 			return (
-				<Card className="w-full max-w-md overflow-hidden">
+				<Card className="w-full max-w-md overflow-hidden rounded-2xl border-0 shadow-none select-none">
 					<div className="aspect-square bg-muted flex items-center justify-center">
 						<div className="flex flex-col items-center gap-2 text-muted-foreground">
 							<ImageIcon className="size-8" />
@@ -165,7 +165,10 @@ export class ImageErrorBoundary extends Component<
  */
 export function ImageSkeleton({ maxWidth = "512px" }: { maxWidth?: string }) {
 	return (
-		<Card className="w-full overflow-hidden animate-pulse" style={{ maxWidth }}>
+		<Card
+			className="w-full overflow-hidden rounded-2xl border-0 shadow-none select-none animate-pulse"
+			style={{ maxWidth }}
+		>
 			<div className="aspect-square bg-muted flex items-center justify-center">
 				<ImageIcon className="size-12 text-muted-foreground/30" />
 			</div>
@@ -176,14 +179,20 @@ export function ImageSkeleton({ maxWidth = "512px" }: { maxWidth?: string }) {
 /**
  * Image Loading State
  */
-export function ImageLoading({ title = "Loading image..." }: { title?: string }) {
+export function ImageLoading({
+	title = "Loading",
+	maxWidth = "512px",
+}: {
+	title?: string;
+	maxWidth?: string;
+}) {
 	return (
-		<Card className="w-full max-w-md overflow-hidden">
+		<Card
+			className="w-full overflow-hidden rounded-2xl border-0 shadow-none select-none"
+			style={{ maxWidth }}
+		>
 			<div className="aspect-square bg-muted flex items-center justify-center">
-				<div className="flex flex-col items-center gap-3">
-					<Spinner size="lg" className="text-muted-foreground" />
-					<p className="text-muted-foreground text-sm">{title}</p>
-				</div>
+				<TextShimmerLoader text={title} size="md" />
 			</div>
 		</Card>
 	);
@@ -214,8 +223,8 @@ export function Image({
 	const [isHovered, setIsHovered] = useState(false);
 	const [imageError, setImageError] = useState(false);
 	const [imageLoaded, setImageLoaded] = useState(false);
-	const displayDomain = domain || source?.label;
 	const isGenerated = domain === "ai-generated";
+	const displayDomain = isGenerated ? "AI Generated" : domain || source?.label;
 	const isAutoRatio = !ratio || ratio === "auto";
 
 	const handleClick = () => {
@@ -227,7 +236,14 @@ export function Image({
 
 	if (imageError) {
 		return (
-			<Card id={id} className={cn("w-full overflow-hidden", className)} style={{ maxWidth }}>
+			<Card
+				id={id}
+				className={cn(
+					"w-full overflow-hidden rounded-2xl border-0 shadow-none select-none",
+					className
+				)}
+				style={{ maxWidth }}
+			>
 				<div className="aspect-square bg-muted flex items-center justify-center">
 					<div className="flex flex-col items-center gap-2 text-muted-foreground">
 						<ImageIcon className="size-8" />
@@ -242,8 +258,7 @@ export function Image({
 		<Card
 			id={id}
 			className={cn(
-				"group w-full overflow-hidden cursor-pointer transition-shadow duration-200 hover:shadow-lg",
-				isGenerated && "ring-1 ring-primary/10",
+				"group w-full overflow-hidden rounded-2xl border-0 shadow-none select-none cursor-pointer transition-shadow duration-200 hover:shadow-lg",
 				className
 			)}
 			style={{ maxWidth }}
@@ -265,18 +280,22 @@ export function Image({
 					<>
 						{!imageLoaded && (
 							<div className="aspect-square flex items-center justify-center">
-								<Spinner size="lg" className="text-muted-foreground" />
+								<TextShimmerLoader text="Loading" size="md" />
 							</div>
 						)}
-						{/* eslint-disable-next-line @next/next/no-img-element */}
-						<img
+						<NextImage
 							src={src}
 							alt={alt}
+							width={0}
+							height={0}
+							sizes="100vw"
+							loading="eager"
 							className={cn(
 								"w-full h-auto transition-transform duration-300",
 								isHovered && "scale-[1.02]",
 								!imageLoaded && "hidden"
 							)}
+							unoptimized
 							onLoad={() => setImageLoaded(true)}
 							onError={() => setImageError(true)}
 						/>
@@ -316,11 +335,9 @@ export function Image({
 						{description && (
 							<p className="text-white/80 text-xs line-clamp-2 mb-1.5">{description}</p>
 						)}
-						{displayDomain && (
+						{displayDomain && !isGenerated && (
 							<div className="flex items-center gap-1.5">
-								{isGenerated ? (
-									<SparklesIcon className="size-3.5 text-white/70" />
-								) : source?.iconUrl ? (
+								{source?.iconUrl ? (
 									<NextImage
 										src={source.iconUrl}
 										alt={source.label}

surfsense_web/components/tool-ui/index.ts

@@ -6,30 +6,14 @@
  * rich UI when specific tools are called by the agent.
  */
 
-export {
-	Article,
-	ArticleErrorBoundary,
-	ArticleLoading,
-	type ArticleProps,
-	ArticleSkeleton,
-	parseSerializableArticle,
-	type SerializableArticle,
-} from "./article";
 export { Audio } from "./audio";
 export {
-	type DeepAgentThinkingArgs,
-	type DeepAgentThinkingResult,
-	DeepAgentThinkingToolUI,
-	InlineThinkingDisplay,
-	type ThinkingStep,
-} from "./deepagent-thinking";
-export {
-	type DisplayImageArgs,
-	DisplayImageArgsSchema,
-	type DisplayImageResult,
-	DisplayImageResultSchema,
-	DisplayImageToolUI,
-} from "./display-image";
+	type GenerateImageArgs,
+	GenerateImageArgsSchema,
+	type GenerateImageResult,
+	GenerateImageResultSchema,
+	GenerateImageToolUI,
+} from "./generate-image";
 export { GeneratePodcastToolUI } from "./generate-podcast";
 export { GenerateReportToolUI } from "./generate-report";
 export { CreateGoogleDriveFileToolUI, DeleteGoogleDriveFileToolUI } from "./google-drive";
@@ -47,27 +31,6 @@ export {
 	DeleteLinearIssueToolUI,
 	UpdateLinearIssueToolUI,
 } from "./linear";
-export {
-	type LinkPreviewArgs,
-	LinkPreviewArgsSchema,
-	type LinkPreviewResult,
-	LinkPreviewResultSchema,
-	LinkPreviewToolUI,
-	type MultiLinkPreviewArgs,
-	MultiLinkPreviewArgsSchema,
-	type MultiLinkPreviewResult,
-	MultiLinkPreviewResultSchema,
-	MultiLinkPreviewToolUI,
-} from "./link-preview";
-export {
-	MediaCard,
-	MediaCardErrorBoundary,
-	MediaCardLoading,
-	type MediaCardProps,
-	MediaCardSkeleton,
-	parseSerializableMediaCard,
-	type SerializableMediaCard,
-} from "./media-card";
 export { CreateNotionPageToolUI, DeleteNotionPageToolUI, UpdateNotionPageToolUI } from "./notion";
 export {
 	Plan,
@@ -85,13 +48,6 @@ export {
 	ExecuteResultSchema,
 	SandboxExecuteToolUI,
 } from "./sandbox-execute";
-export {
-	type ScrapeWebpageArgs,
-	ScrapeWebpageArgsSchema,
-	type ScrapeWebpageResult,
-	ScrapeWebpageResultSchema,
-	ScrapeWebpageToolUI,
-} from "./scrape-webpage";
 export {
 	type MemoryItem,
 	type RecallMemoryArgs,

surfsense_web/components/tool-ui/jira/create-jira-issue.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -536,7 +536,10 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const CreateJiraIssueToolUI = makeAssistantToolUI<
+export const CreateJiraIssueToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{
 		project_key: string;
 		summary: string;
@@ -545,39 +548,36 @@ export const CreateJiraIssueToolUI = makeAssistantToolUI<
 		priority?: string;
 	},
 	CreateJiraIssueResult
->({
-	toolName: "create_jira_issue",
-	render: function CreateJiraIssueUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/jira/delete-jira-issue.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { CornerDownLeftIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
@@ -393,44 +393,43 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const DeleteJiraIssueToolUI = makeAssistantToolUI<
+export const DeleteJiraIssueToolUI = ({
+	result,
+}: ToolCallMessagePartProps<
 	{ issue_title_or_key: string; delete_from_kb?: boolean },
 	DeleteJiraIssueResult
->({
-	toolName: "delete_jira_issue",
-	render: function DeleteJiraIssueUI({ result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isWarningResult(result)) return <WarningCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isWarningResult(result)) return <WarningCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/jira/update-jira-issue.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
@@ -553,7 +553,10 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const UpdateJiraIssueToolUI = makeAssistantToolUI<
+export const UpdateJiraIssueToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{
 		issue_title_or_key: string;
 		new_summary?: string;
@@ -561,40 +564,37 @@ export const UpdateJiraIssueToolUI = makeAssistantToolUI<
 		new_priority?: string;
 	},
 	UpdateJiraIssueResult
->({
-	toolName: "update_jira_issue",
-	render: function UpdateJiraIssueUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isInsufficientPermissionsResult(result))
-			return <InsufficientPermissionsCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isInsufficientPermissionsResult(result))
+		return <InsufficientPermissionsCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/linear/create-linear-issue.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -605,40 +605,37 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const CreateLinearIssueToolUI = makeAssistantToolUI<
-	{ title: string; description?: string },
-	CreateLinearIssueResult
->({
-	toolName: "create_linear_issue",
-	render: function CreateLinearIssueUI({ args, result }) {
-		if (!result) return null;
+export const CreateLinearIssueToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<{ title: string; description?: string }, CreateLinearIssueResult>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/linear/delete-linear-issue.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { CornerDownLeftIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
@@ -360,42 +360,41 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const DeleteLinearIssueToolUI = makeAssistantToolUI<
+export const DeleteLinearIssueToolUI = ({
+	result,
+}: ToolCallMessagePartProps<
 	{ issue_ref: string; delete_from_kb?: boolean },
 	DeleteLinearIssueResult
->({
-	toolName: "delete_linear_issue",
-	render: function DeleteLinearIssueUI({ result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isWarningResult(result)) return <WarningCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isWarningResult(result)) return <WarningCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/linear/update-linear-issue.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
@@ -739,7 +739,10 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const UpdateLinearIssueToolUI = makeAssistantToolUI<
+export const UpdateLinearIssueToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<
 	{
 		issue_ref: string;
 		new_title?: string;
@@ -750,38 +753,35 @@ export const UpdateLinearIssueToolUI = makeAssistantToolUI<
 		new_label_names?: string[];
 	},
 	UpdateLinearIssueResult
->({
-	toolName: "update_linear_issue",
-	render: function UpdateLinearIssueUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					window.dispatchEvent(
+						new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
+					);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
-		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
-		if (isErrorResult(result)) return <ErrorCard result={result} />;
+	if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+	if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+	if (isErrorResult(result)) return <ErrorCard result={result} />;
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/link-preview.tsx

@@ -1,259 +0,0 @@
-"use client";
-
-import { makeAssistantToolUI } from "@assistant-ui/react";
-import { AlertCircleIcon, ExternalLinkIcon, LinkIcon } from "lucide-react";
-import { z } from "zod";
-import {
-	MediaCard,
-	MediaCardErrorBoundary,
-	MediaCardLoading,
-	parseSerializableMediaCard,
-	type SerializableMediaCard,
-} from "@/components/tool-ui/media-card";
-
-// ============================================================================
-// Zod Schemas
-// ============================================================================
-
-/**
- * Schema for link_preview tool arguments
- */
-const LinkPreviewArgsSchema = z.object({
-	url: z.string(),
-	title: z.string().nullish(),
-});
-
-/**
- * Schema for link_preview tool result
- */
-const LinkPreviewResultSchema = z.object({
-	id: z.string(),
-	assetId: z.string(),
-	kind: z.literal("link"),
-	href: z.string(),
-	title: z.string(),
-	description: z.string().nullish(),
-	thumb: z.string().nullish(),
-	domain: z.string().nullish(),
-	error: z.string().nullish(),
-});
-
-// ============================================================================
-// Types
-// ============================================================================
-
-type LinkPreviewArgs = z.infer<typeof LinkPreviewArgsSchema>;
-type LinkPreviewResult = z.infer<typeof LinkPreviewResultSchema>;
-
-/**
- * Error state component shown when link preview fails
- */
-function LinkPreviewErrorState({ url, error }: { url: string; error: string }) {
-	return (
-		<div className="my-4 overflow-hidden rounded-xl border border-destructive/20 bg-destructive/5 p-4 max-w-md">
-			<div className="flex items-center gap-4">
-				<div className="flex size-12 shrink-0 items-center justify-center rounded-lg bg-destructive/10">
-					<AlertCircleIcon className="size-6 text-destructive" />
-				</div>
-				<div className="flex-1 min-w-0">
-					<p className="font-medium text-destructive text-sm">Failed to load preview</p>
-					<p className="text-muted-foreground text-xs mt-0.5 truncate">{url}</p>
-					<p className="text-muted-foreground text-xs mt-1">{error}</p>
-				</div>
-			</div>
-		</div>
-	);
-}
-
-/**
- * Cancelled state component
- */
-function LinkPreviewCancelledState({ url }: { url: string }) {
-	return (
-		<div className="my-4 rounded-xl border border-muted p-4 text-muted-foreground max-w-md">
-			<p className="flex items-center gap-2">
-				<LinkIcon className="size-4" />
-				<span className="line-through truncate">Preview: {url}</span>
-			</p>
-		</div>
-	);
-}
-
-/**
- * Parsed MediaCard component with error handling
- */
-function ParsedMediaCard({ result }: { result: unknown }) {
-	const card = parseSerializableMediaCard(result);
-
-	return (
-		<MediaCard
-			{...card}
-			maxWidth="420px"
-			responseActions={[{ id: "open", label: "Open", variant: "default" }]}
-			onResponseAction={(id) => {
-				if (id === "open" && card.href) {
-					window.open(card.href, "_blank", "noopener,noreferrer");
-				}
-			}}
-		/>
-	);
-}
-
-/**
- * Link Preview Tool UI Component
- *
- * This component is registered with assistant-ui to render a rich
- * link preview card when the link_preview tool is called by the agent.
- *
- * It displays website metadata including:
- * - Title and description
- * - Thumbnail/Open Graph image
- * - Domain name
- * - Clickable link to open in new tab
- */
-export const LinkPreviewToolUI = makeAssistantToolUI<LinkPreviewArgs, LinkPreviewResult>({
-	toolName: "link_preview",
-	render: function LinkPreviewUI({ args, result, status }) {
-		const url = args.url || "Unknown URL";
-
-		// Loading state - tool is still running
-		if (status.type === "running" || status.type === "requires-action") {
-			return (
-				<div className="my-4">
-					<MediaCardLoading title={`Loading preview for ${url}...`} />
-				</div>
-			);
-		}
-
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return <LinkPreviewCancelledState url={url} />;
-			}
-			if (status.reason === "error") {
-				return (
-					<LinkPreviewErrorState
-						url={url}
-						error={typeof status.error === "string" ? status.error : "An error occurred"}
-					/>
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return (
-				<div className="my-4">
-					<MediaCardLoading title={`Fetching metadata for ${url}...`} />
-				</div>
-			);
-		}
-
-		// Error result from the tool
-		if (result.error) {
-			return <LinkPreviewErrorState url={url} error={result.error} />;
-		}
-
-		// Success - render the media card
-		return (
-			<div className="my-4">
-				<MediaCardErrorBoundary>
-					<ParsedMediaCard result={result} />
-				</MediaCardErrorBoundary>
-			</div>
-		);
-	},
-});
-
-// ============================================================================
-// Multi Link Preview Schemas
-// ============================================================================
-
-/**
- * Schema for multi_link_preview tool arguments
- */
-const MultiLinkPreviewArgsSchema = z.object({
-	urls: z.array(z.string()),
-});
-
-/**
- * Schema for error items in multi_link_preview result
- */
-const MultiLinkPreviewErrorSchema = z.object({
-	url: z.string(),
-	error: z.string(),
-});
-
-/**
- * Schema for multi_link_preview tool result
- */
-const MultiLinkPreviewResultSchema = z.object({
-	previews: z.array(LinkPreviewResultSchema),
-	errors: z.array(MultiLinkPreviewErrorSchema).nullish(),
-});
-
-type MultiLinkPreviewArgs = z.infer<typeof MultiLinkPreviewArgsSchema>;
-type MultiLinkPreviewResult = z.infer<typeof MultiLinkPreviewResultSchema>;
-
-export const MultiLinkPreviewToolUI = makeAssistantToolUI<
-	MultiLinkPreviewArgs,
-	MultiLinkPreviewResult
->({
-	toolName: "multi_link_preview",
-	render: function MultiLinkPreviewUI({ args, result, status }) {
-		const urls = args.urls || [];
-
-		// Loading state
-		if (status.type === "running" || status.type === "requires-action") {
-			return (
-				<div className="my-4 grid gap-4 sm:grid-cols-2">
-					{urls.slice(0, 4).map((url, index) => (
-						<MediaCardLoading key={`loading-${url}-${index}`} title="Loading..." />
-					))}
-				</div>
-			);
-		}
-
-		// Incomplete state
-		if (status.type === "incomplete") {
-			return (
-				<div className="my-4 text-muted-foreground text-sm">
-					<p className="flex items-center gap-2">
-						<LinkIcon className="size-4" />
-						<span>Link previews cancelled</span>
-					</p>
-				</div>
-			);
-		}
-
-		// No result
-		if (!result || !result.previews) {
-			return null;
-		}
-
-		// Render grid of previews
-		return (
-			<div className="my-4 grid gap-4 sm:grid-cols-2">
-				{result.previews.map((preview) => (
-					<MediaCardErrorBoundary key={preview.id}>
-						<ParsedMediaCard result={preview} />
-					</MediaCardErrorBoundary>
-				))}
-				{result.errors?.map((err) => (
-					<LinkPreviewErrorState key={err.url} url={err.url} error={err.error} />
-				))}
-			</div>
-		);
-	},
-});
-
-export {
-	LinkPreviewArgsSchema,
-	LinkPreviewResultSchema,
-	MultiLinkPreviewArgsSchema,
-	MultiLinkPreviewResultSchema,
-	type LinkPreviewArgs,
-	type LinkPreviewResult,
-	type MultiLinkPreviewArgs,
-	type MultiLinkPreviewResult,
-};

surfsense_web/components/tool-ui/media-card/index.tsx

@@ -1,354 +0,0 @@
-"use client";
-
-import { ExternalLinkIcon, Globe, ImageIcon, LinkIcon } from "lucide-react";
-import Image from "next/image";
-import { Component, type ReactNode } from "react";
-import { z } from "zod";
-import { Badge } from "@/components/ui/badge";
-import { Button } from "@/components/ui/button";
-import { Card, CardContent } from "@/components/ui/card";
-import { Spinner } from "@/components/ui/spinner";
-import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
-import { cn } from "@/lib/utils";
-
-/**
- * Zod schemas for runtime validation
- */
-const AspectRatioSchema = z.enum(["1:1", "4:3", "16:9", "9:16", "21:9", "auto"]);
-const MediaCardKindSchema = z.enum(["link", "image", "video", "audio"]);
-
-const ResponseActionSchema = z.object({
-	id: z.string(),
-	label: z.string(),
-	variant: z.enum(["default", "secondary", "outline", "destructive", "ghost"]).nullish(),
-	confirmLabel: z.string().nullish(),
-});
-
-const SerializableMediaCardSchema = z.object({
-	id: z.string(),
-	assetId: z.string(),
-	kind: MediaCardKindSchema,
-	href: z.string().nullish(),
-	src: z.string().nullish(),
-	title: z.string(),
-	description: z.string().nullish(),
-	thumb: z.string().nullish(),
-	ratio: AspectRatioSchema.nullish(),
-	domain: z.string().nullish(),
-});
-
-/**
- * Types derived from Zod schemas
- */
-type AspectRatio = z.infer<typeof AspectRatioSchema>;
-type MediaCardKind = z.infer<typeof MediaCardKindSchema>;
-type ResponseAction = z.infer<typeof ResponseActionSchema>;
-export type SerializableMediaCard = z.infer<typeof SerializableMediaCardSchema>;
-
-/**
- * Props for the MediaCard component
- */
-export interface MediaCardProps {
-	id: string;
-	assetId: string;
-	kind: MediaCardKind;
-	href?: string;
-	src?: string;
-	title: string;
-	description?: string;
-	thumb?: string;
-	ratio?: AspectRatio;
-	domain?: string;
-	maxWidth?: string;
-	alt?: string;
-	className?: string;
-	responseActions?: ResponseAction[];
-	onResponseAction?: (id: string) => void;
-}
-
-/**
- * Parse and validate serializable media card from tool result
- */
-export function parseSerializableMediaCard(result: unknown): SerializableMediaCard {
-	const parsed = SerializableMediaCardSchema.safeParse(result);
-
-	if (!parsed.success) {
-		console.warn("Invalid media card data:", parsed.error.issues);
-		throw new Error(`Invalid media card: ${parsed.error.issues.map((i) => i.message).join(", ")}`);
-	}
-
-	return parsed.data;
-}
-
-/**
- * Get aspect ratio class based on ratio prop
- */
-function getAspectRatioClass(ratio?: AspectRatio): string {
-	switch (ratio) {
-		case "1:1":
-			return "aspect-square";
-		case "4:3":
-			return "aspect-[4/3]";
-		case "16:9":
-			return "aspect-video";
-		case "9:16":
-			return "aspect-[9/16]";
-		case "21:9":
-			return "aspect-[21/9]";
-		case "auto":
-		default:
-			return "aspect-[2/1]";
-	}
-}
-
-/**
- * Get icon based on media card kind
- */
-function getKindIcon(kind: MediaCardKind) {
-	switch (kind) {
-		case "link":
-			return <LinkIcon className="size-5" />;
-		case "image":
-			return <ImageIcon className="size-5" />;
-		case "video":
-		case "audio":
-			return <Globe className="size-5" />;
-		default:
-			return <LinkIcon className="size-5" />;
-	}
-}
-
-/**
- * Error boundary for MediaCard
- */
-interface MediaCardErrorBoundaryState {
-	hasError: boolean;
-	error?: Error;
-}
-
-export class MediaCardErrorBoundary extends Component<
-	{ children: ReactNode },
-	MediaCardErrorBoundaryState
-> {
-	constructor(props: { children: ReactNode }) {
-		super(props);
-		this.state = { hasError: false };
-	}
-
-	static getDerivedStateFromError(error: Error): MediaCardErrorBoundaryState {
-		return { hasError: true, error };
-	}
-
-	render() {
-		if (this.state.hasError) {
-			return (
-				<Card className="w-full max-w-md border-destructive/20 bg-destructive/5">
-					<CardContent className="flex items-center gap-3 p-4">
-						<div className="flex size-10 shrink-0 items-center justify-center rounded-lg bg-destructive/10">
-							<LinkIcon className="size-5 text-destructive" />
-						</div>
-						<div className="min-w-0 flex-1">
-							<p className="font-medium text-destructive text-sm">Failed to load preview</p>
-							<p className="text-muted-foreground text-xs truncate">
-								{this.state.error?.message || "An error occurred"}
-							</p>
-						</div>
-					</CardContent>
-				</Card>
-			);
-		}
-
-		return this.props.children;
-	}
-}
-
-/**
- * Loading skeleton for MediaCard
- */
-export function MediaCardSkeleton({ maxWidth = "420px" }: { maxWidth?: string }) {
-	return (
-		<Card className="w-full overflow-hidden animate-pulse" style={{ maxWidth }}>
-			<div className="aspect-[2/1] bg-muted" />
-			<CardContent className="p-4">
-				<div className="h-4 w-3/4 rounded bg-muted" />
-				<div className="mt-2 h-3 w-full rounded bg-muted" />
-				<div className="mt-1 h-3 w-2/3 rounded bg-muted" />
-			</CardContent>
-		</Card>
-	);
-}
-
-/**
- * MediaCard Component
- *
- * A rich media card for displaying link previews, images, and other media
- * in AI chat applications. Supports thumbnails, descriptions, and actions.
- */
-export function MediaCard({
-	id,
-	kind,
-	href,
-	title,
-	description,
-	thumb,
-	ratio = "auto",
-	domain,
-	maxWidth = "420px",
-	alt,
-	className,
-	responseActions,
-	onResponseAction,
-}: MediaCardProps) {
-	const aspectRatioClass = getAspectRatioClass(ratio);
-	const displayDomain = domain || (href ? new URL(href).hostname.replace("www.", "") : undefined);
-
-	const handleCardClick = () => {
-		if (href) {
-			window.open(href, "_blank", "noopener,noreferrer");
-		}
-	};
-
-	return (
-		<TooltipProvider>
-			<Card
-				id={id}
-				className={cn(
-					"group relative w-full overflow-hidden transition-all duration-200",
-					"hover:shadow-lg hover:border-primary/20",
-					href && "cursor-pointer",
-					className
-				)}
-				style={{ maxWidth }}
-				onClick={href ? handleCardClick : undefined}
-				role={href ? "link" : undefined}
-				tabIndex={href ? 0 : undefined}
-				onKeyDown={(e) => {
-					if (href && (e.key === "Enter" || e.key === " ")) {
-						e.preventDefault();
-						handleCardClick();
-					}
-				}}
-			>
-				{/* Thumbnail */}
-				{thumb && (
-					<div className={cn("relative w-full overflow-hidden bg-muted", aspectRatioClass)}>
-						<Image
-							src={thumb}
-							alt={alt || title}
-							fill
-							className="object-cover transition-transform duration-300 group-hover:scale-105"
-							unoptimized
-							onError={(e) => {
-								// Hide broken images
-								e.currentTarget.style.display = "none";
-							}}
-						/>
-						{/* Gradient overlay */}
-						<div className="absolute inset-0 bg-gradient-to-t from-black/20 to-transparent opacity-0 transition-opacity group-hover:opacity-100" />
-					</div>
-				)}
-
-				{/* Fallback when no thumbnail */}
-				{!thumb && (
-					<div
-						className={cn(
-							"relative flex w-full items-center justify-center bg-gradient-to-br from-muted to-muted/50",
-							aspectRatioClass
-						)}
-					>
-						<div className="flex flex-col items-center gap-2 text-muted-foreground">
-							{getKindIcon(kind)}
-							<span className="text-xs">{kind === "link" ? "Link Preview" : kind}</span>
-						</div>
-					</div>
-				)}
-
-				{/* Content */}
-				<CardContent className="p-4">
-					<div className="flex items-start gap-3">
-						{/* Domain favicon placeholder */}
-						<div className="flex size-10 shrink-0 items-center justify-center rounded-lg bg-muted">
-							<Globe className="size-5 text-muted-foreground" />
-						</div>
-
-						<div className="min-w-0 flex-1">
-							{/* Domain badge */}
-							{displayDomain && (
-								<div className="mb-1.5 flex items-center gap-1.5">
-									<Badge variant="secondary" className="text-xs font-normal">
-										{displayDomain}
-									</Badge>
-									{href && (
-										<ExternalLinkIcon className="size-3 text-muted-foreground opacity-0 transition-opacity group-hover:opacity-100" />
-									)}
-								</div>
-							)}
-
-							{/* Title */}
-							<h3 className="font-semibold text-foreground text-sm leading-tight line-clamp-2 group-hover:text-primary transition-colors">
-								{title}
-							</h3>
-
-							{/* Description */}
-							{description && (
-								<p className="mt-1.5 text-muted-foreground text-xs leading-relaxed line-clamp-2">
-									{description}
-								</p>
-							)}
-						</div>
-					</div>
-
-					{/* Response Actions */}
-					{responseActions && responseActions.length > 0 && (
-						<div className="mt-4 flex items-center justify-end gap-2 border-t pt-3">
-							{responseActions.map((action) => (
-								<Tooltip key={action.id}>
-									<TooltipTrigger asChild>
-										<Button
-											variant={action.variant || "secondary"}
-											size="sm"
-											onClick={(e) => {
-												e.stopPropagation();
-												onResponseAction?.(action.id);
-											}}
-										>
-											{action.label}
-										</Button>
-									</TooltipTrigger>
-									{action.confirmLabel && (
-										<TooltipContent>
-											<p>{action.confirmLabel}</p>
-										</TooltipContent>
-									)}
-								</Tooltip>
-							))}
-						</div>
-					)}
-				</CardContent>
-			</Card>
-		</TooltipProvider>
-	);
-}
-
-/**
- * MediaCard Loading State
- */
-export function MediaCardLoading({ title = "Loading preview..." }: { title?: string }) {
-	return (
-		<Card className="w-full max-w-md overflow-hidden">
-			<div className="aspect-[2/1] bg-muted animate-pulse flex items-center justify-center">
-				<Spinner size="lg" className="text-muted-foreground" />
-			</div>
-			<CardContent className="p-4">
-				<div className="flex items-center gap-3">
-					<div className="size-10 rounded-lg bg-muted animate-pulse" />
-					<div className="flex-1 space-y-2">
-						<div className="h-4 w-3/4 rounded bg-muted animate-pulse" />
-						<div className="h-3 w-1/2 rounded bg-muted animate-pulse" />
-					</div>
-				</div>
-				<p className="mt-3 text-center text-muted-foreground text-sm">{title}</p>
-			</CardContent>
-		</Card>
-	);
-}

surfsense_web/components/tool-ui/notion/create-notion-page.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
@@ -445,46 +445,43 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const CreateNotionPageToolUI = makeAssistantToolUI<
-	{ title: string; content: string },
-	CreateNotionPageResult
->({
-	toolName: "create_notion_page",
-	render: function CreateNotionPageUI({ args, result }) {
-		if (!result) return null;
+export const CreateNotionPageToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<{ title: string; content: string }, CreateNotionPageResult>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (isAuthErrorResult(result)) {
-			return <AuthErrorCard result={result} />;
-		}
+	if (isAuthErrorResult(result)) {
+		return <AuthErrorCard result={result} />;
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isErrorResult(result)) {
-			return <ErrorCard result={result} />;
-		}
+	if (isErrorResult(result)) {
+		return <ErrorCard result={result} />;
+	}
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/notion/delete-notion-page.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { CornerDownLeftIcon } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
@@ -372,53 +372,52 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const DeleteNotionPageToolUI = makeAssistantToolUI<
+export const DeleteNotionPageToolUI = ({
+	result,
+}: ToolCallMessagePartProps<
 	{ page_title: string; delete_from_kb?: boolean },
 	DeleteNotionPageResult
->({
-	toolName: "delete_notion_page",
-	render: function DeleteNotionPageUI({ result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isInfoResult(result)) {
-			return <InfoCard result={result} />;
-		}
+	if (isInfoResult(result)) {
+		return <InfoCard result={result} />;
+	}
 
-		if (isWarningResult(result)) {
-			return <WarningCard result={result} />;
-		}
+	if (isWarningResult(result)) {
+		return <WarningCard result={result} />;
+	}
 
-		if (isAuthErrorResult(result)) {
-			return <AuthErrorCard result={result} />;
-		}
+	if (isAuthErrorResult(result)) {
+		return <AuthErrorCard result={result} />;
+	}
 
-		if (isErrorResult(result)) {
-			return <ErrorCard result={result} />;
-		}
+	if (isErrorResult(result)) {
+		return <ErrorCard result={result} />;
+	}
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/notion/update-notion-page.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { useSetAtom } from "jotai";
 import { CornerDownLeftIcon, Pen } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
@@ -395,50 +395,47 @@ function SuccessCard({ result }: { result: SuccessResult }) {
 	);
 }
 
-export const UpdateNotionPageToolUI = makeAssistantToolUI<
-	{ page_title: string; content: string },
-	UpdateNotionPageResult
->({
-	toolName: "update_notion_page",
-	render: function UpdateNotionPageUI({ args, result }) {
-		if (!result) return null;
+export const UpdateNotionPageToolUI = ({
+	args,
+	result,
+}: ToolCallMessagePartProps<{ page_title: string; content: string }, UpdateNotionPageResult>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				<ApprovalCard
-					args={args}
-					interruptData={result}
-					onDecision={(decision) => {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			<ApprovalCard
+				args={args}
+				interruptData={result}
+				onDecision={(decision) => {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (
-			typeof result === "object" &&
-			result !== null &&
-			"status" in result &&
-			(result as { status: string }).status === "rejected"
-		) {
-			return null;
-		}
+	if (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as { status: string }).status === "rejected"
+	) {
+		return null;
+	}
 
-		if (isInfoResult(result)) {
-			return <InfoCard result={result} />;
-		}
+	if (isInfoResult(result)) {
+		return <InfoCard result={result} />;
+	}
 
-		if (isAuthErrorResult(result)) {
-			return <AuthErrorCard result={result} />;
-		}
+	if (isAuthErrorResult(result)) {
+		return <AuthErrorCard result={result} />;
+	}
 
-		if (isErrorResult(result)) {
-			return <ErrorCard result={result} />;
-		}
+	if (isErrorResult(result)) {
+		return <ErrorCard result={result} />;
+	}
 
-		return <SuccessCard result={result as SuccessResult} />;
-	},
-});
+	return <SuccessCard result={result as SuccessResult} />;
+};

surfsense_web/components/tool-ui/sandbox-execute.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import {
 	AlertCircleIcon,
 	CheckCircle2Icon,
@@ -380,41 +380,42 @@ function ExecuteCompleted({
 // Tool UI
 // ============================================================================
 
-export const SandboxExecuteToolUI = makeAssistantToolUI<ExecuteArgs, ExecuteResult>({
-	toolName: "execute",
-	render: function SandboxExecuteUI({ args, result, status }) {
-		const command = args.command || "…";
+export const SandboxExecuteToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps<ExecuteArgs, ExecuteResult>) => {
+	const command = args.command || "…";
 
-		if (status.type === "running" || status.type === "requires-action") {
-			return <ExecuteLoading command={command} />;
+	if (status.type === "running" || status.type === "requires-action") {
+		return <ExecuteLoading command={command} />;
+	}
+
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
+			return <ExecuteCancelledState command={command} />;
 		}
-
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return <ExecuteCancelledState command={command} />;
-			}
-			if (status.reason === "error") {
-				return (
-					<ExecuteErrorState
-						command={command}
-						error={typeof status.error === "string" ? status.error : "An error occurred"}
-					/>
-				);
-			}
+		if (status.reason === "error") {
+			return (
+				<ExecuteErrorState
+					command={command}
+					error={typeof status.error === "string" ? status.error : "An error occurred"}
+				/>
+			);
 		}
+	}
 
-		if (!result) {
-			return <ExecuteLoading command={command} />;
-		}
+	if (!result) {
+		return <ExecuteLoading command={command} />;
+	}
 
-		if (result.error && !result.result && !result.output) {
-			return <ExecuteErrorState command={command} error={result.error} />;
-		}
+	if (result.error && !result.result && !result.output) {
+		return <ExecuteErrorState command={command} error={result.error} />;
+	}
 
-		const parsed = parseExecuteResult(result);
-		const threadId = result.thread_id || null;
-		return <ExecuteCompleted command={command} parsed={parsed} threadId={threadId} />;
-	},
-});
+	const parsed = parseExecuteResult(result);
+	const threadId = result.thread_id || null;
+	return <ExecuteCompleted command={command} parsed={parsed} threadId={threadId} />;
+};
 
 export { ExecuteArgsSchema, ExecuteResultSchema, type ExecuteArgs, type ExecuteResult };

surfsense_web/components/tool-ui/scrape-webpage.tsx

@@ -1,166 +0,0 @@
-"use client";
-
-import { makeAssistantToolUI } from "@assistant-ui/react";
-import { AlertCircleIcon, FileTextIcon } from "lucide-react";
-import { z } from "zod";
-import {
-	Article,
-	ArticleErrorBoundary,
-	ArticleLoading,
-	parseSerializableArticle,
-} from "@/components/tool-ui/article";
-
-// ============================================================================
-// Zod Schemas
-// ============================================================================
-
-/**
- * Schema for scrape_webpage tool arguments
- */
-const ScrapeWebpageArgsSchema = z.object({
-	url: z.string(),
-	max_length: z.number().nullish(),
-});
-
-/**
- * Schema for scrape_webpage tool result
- */
-const ScrapeWebpageResultSchema = z.object({
-	id: z.string(),
-	assetId: z.string(),
-	kind: z.literal("article"),
-	href: z.string(),
-	title: z.string(),
-	description: z.string().nullish(),
-	content: z.string().nullish(),
-	domain: z.string().nullish(),
-	author: z.string().nullish(),
-	date: z.string().nullish(),
-	word_count: z.number().nullish(),
-	was_truncated: z.boolean().nullish(),
-	crawler_type: z.string().nullish(),
-	error: z.string().nullish(),
-});
-
-// ============================================================================
-// Types
-// ============================================================================
-
-type ScrapeWebpageArgs = z.infer<typeof ScrapeWebpageArgsSchema>;
-type ScrapeWebpageResult = z.infer<typeof ScrapeWebpageResultSchema>;
-
-/**
- * Error state component shown when webpage scraping fails
- */
-function ScrapeErrorState({ url, error }: { url: string; error: string }) {
-	return (
-		<div className="my-4 overflow-hidden rounded-xl border border-destructive/20 bg-destructive/5 p-4 max-w-md">
-			<div className="flex items-center gap-4">
-				<div className="flex size-12 shrink-0 items-center justify-center rounded-lg bg-destructive/10">
-					<AlertCircleIcon className="size-6 text-destructive" />
-				</div>
-				<div className="flex-1 min-w-0">
-					<p className="font-medium text-destructive text-sm">Failed to scrape webpage</p>
-					<p className="text-muted-foreground text-xs mt-0.5 truncate">{url}</p>
-					<p className="text-muted-foreground text-xs mt-1">{error}</p>
-				</div>
-			</div>
-		</div>
-	);
-}
-
-/**
- * Cancelled state component
- */
-function ScrapeCancelledState({ url }: { url: string }) {
-	return (
-		<div className="my-4 rounded-xl border border-muted p-4 text-muted-foreground max-w-md">
-			<p className="flex items-center gap-2">
-				<FileTextIcon className="size-4" />
-				<span className="line-through truncate">Scraping: {url}</span>
-			</p>
-		</div>
-	);
-}
-
-/**
- * Parsed Article component with error handling
- */
-function ParsedArticle({ result }: { result: unknown }) {
-	const { description, ...article } = parseSerializableArticle(result);
-
-	return <Article {...article} maxWidth="480px" />;
-}
-
-/**
- * Scrape Webpage Tool UI Component
- *
- * This component is registered with assistant-ui to render an article card
- * when the scrape_webpage tool is called by the agent.
- *
- * It displays scraped webpage content including:
- * - Title and description
- * - Author and date (if available)
- * - Word count
- * - Link to original source
- */
-export const ScrapeWebpageToolUI = makeAssistantToolUI<ScrapeWebpageArgs, ScrapeWebpageResult>({
-	toolName: "scrape_webpage",
-	render: function ScrapeWebpageUI({ args, result, status }) {
-		const url = args.url || "Unknown URL";
-
-		// Loading state - tool is still running
-		if (status.type === "running" || status.type === "requires-action") {
-			return (
-				<div className="my-4">
-					<ArticleLoading title={`Scraping ${url}...`} />
-				</div>
-			);
-		}
-
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return <ScrapeCancelledState url={url} />;
-			}
-			if (status.reason === "error") {
-				return (
-					<ScrapeErrorState
-						url={url}
-						error={typeof status.error === "string" ? status.error : "An error occurred"}
-					/>
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return (
-				<div className="my-4">
-					<ArticleLoading title={`Extracting content from ${url}...`} />
-				</div>
-			);
-		}
-
-		// Error result from the tool
-		if (result.error) {
-			return <ScrapeErrorState url={url} error={result.error} />;
-		}
-
-		// Success - render the article card
-		return (
-			<div className="my-4">
-				<ArticleErrorBoundary>
-					<ParsedArticle result={result} />
-				</ArticleErrorBoundary>
-			</div>
-		);
-	},
-});
-
-export {
-	ScrapeWebpageArgsSchema,
-	ScrapeWebpageResultSchema,
-	type ScrapeWebpageArgs,
-	type ScrapeWebpageResult,
-};

surfsense_web/components/tool-ui/user-memory.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
 import { BrainIcon, CheckIcon, Loader2Icon, SearchIcon, XIcon } from "lucide-react";
 import { z } from "zod";
 
@@ -80,191 +80,193 @@ function CategoryBadge({ category }: { category: string }) {
 // Save Memory Tool UI
 // ============================================================================
 
-export const SaveMemoryToolUI = makeAssistantToolUI<SaveMemoryArgs, SaveMemoryResult>({
-	toolName: "save_memory",
-	render: function SaveMemoryUI({ args, result, status }) {
-		const isRunning = status.type === "running" || status.type === "requires-action";
-		const isComplete = status.type === "complete";
-		const isError = result?.status === "error";
+export const SaveMemoryToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps<SaveMemoryArgs, SaveMemoryResult>) => {
+	const isRunning = status.type === "running" || status.type === "requires-action";
+	const isComplete = status.type === "complete";
+	const isError = result?.status === "error";
 
-		// Parse args safely
-		const parsedArgs = SaveMemoryArgsSchema.safeParse(args);
-		const content = parsedArgs.success ? parsedArgs.data.content : "";
-		const category = parsedArgs.success ? parsedArgs.data.category : "fact";
+	// Parse args safely
+	const parsedArgs = SaveMemoryArgsSchema.safeParse(args);
+	const content = parsedArgs.success ? parsedArgs.data.content : "";
+	const category = parsedArgs.success ? parsedArgs.data.category : "fact";
 
-		// Loading state
-		if (isRunning) {
-			return (
-				<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
-					<div className="flex size-8 items-center justify-center rounded-full bg-primary/10">
-						<Loader2Icon className="size-4 animate-spin text-primary" />
-					</div>
-					<div className="flex-1">
-						<span className="text-sm text-muted-foreground">Saving to memory...</span>
-					</div>
+	// Loading state
+	if (isRunning) {
+		return (
+			<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
+				<div className="flex size-8 items-center justify-center rounded-full bg-primary/10">
+					<Loader2Icon className="size-4 animate-spin text-primary" />
 				</div>
-			);
-		}
-
-		// Error state
-		if (isError) {
-			return (
-				<div className="my-3 flex items-center gap-3 rounded-lg border border-destructive/20 bg-destructive/5 px-4 py-3">
-					<div className="flex size-8 items-center justify-center rounded-full bg-destructive/10">
-						<XIcon className="size-4 text-destructive" />
-					</div>
-					<div className="flex-1">
-						<span className="text-sm text-destructive">Failed to save memory</span>
-						{result?.error && <p className="mt-1 text-xs text-destructive/70">{result.error}</p>}
-					</div>
+				<div className="flex-1">
+					<span className="text-sm text-muted-foreground">Saving to memory...</span>
 				</div>
-			);
-		}
+			</div>
+		);
+	}
 
-		// Success state
-		if (isComplete && result?.status === "saved") {
-			return (
-				<div className="my-3 flex items-center gap-3 rounded-lg border border-primary/20 bg-primary/5 px-4 py-3">
-					<div className="flex size-8 items-center justify-center rounded-full bg-primary/10">
-						<BrainIcon className="size-4 text-primary" />
-					</div>
-					<div className="flex-1 min-w-0">
-						<div className="flex items-center gap-2">
-							<CheckIcon className="size-3 text-green-500 shrink-0" />
-							<span className="text-sm font-medium text-foreground">Memory saved</span>
-							<CategoryBadge category={category} />
-						</div>
-						<p className="mt-1 truncate text-sm text-muted-foreground">{content}</p>
-					</div>
+	// Error state
+	if (isError) {
+		return (
+			<div className="my-3 flex items-center gap-3 rounded-lg border border-destructive/20 bg-destructive/5 px-4 py-3">
+				<div className="flex size-8 items-center justify-center rounded-full bg-destructive/10">
+					<XIcon className="size-4 text-destructive" />
 				</div>
-			);
-		}
-
-		// Default/incomplete state - show what's being saved
-		if (content) {
-			return (
-				<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
-					<div className="flex size-8 items-center justify-center rounded-full bg-muted">
-						<BrainIcon className="size-4 text-muted-foreground" />
-					</div>
-					<div className="flex-1 min-w-0">
-						<div className="flex items-center gap-2">
-							<span className="text-sm text-muted-foreground">Saving memory</span>
-							<CategoryBadge category={category} />
-						</div>
-						<p className="mt-1 truncate text-sm text-muted-foreground">{content}</p>
-					</div>
+				<div className="flex-1">
+					<span className="text-sm text-destructive">Failed to save memory</span>
+					{result?.error && <p className="mt-1 text-xs text-destructive/70">{result.error}</p>}
 				</div>
-			);
-		}
+			</div>
+		);
+	}
 
-		return null;
-	},
-});
+	// Success state
+	if (isComplete && result?.status === "saved") {
+		return (
+			<div className="my-3 flex items-center gap-3 rounded-lg border border-primary/20 bg-primary/5 px-4 py-3">
+				<div className="flex size-8 items-center justify-center rounded-full bg-primary/10">
+					<BrainIcon className="size-4 text-primary" />
+				</div>
+				<div className="flex-1 min-w-0">
+					<div className="flex items-center gap-2">
+						<CheckIcon className="size-3 text-green-500 shrink-0" />
+						<span className="text-sm font-medium text-foreground">Memory saved</span>
+						<CategoryBadge category={category} />
+					</div>
+					<p className="mt-1 truncate text-sm text-muted-foreground">{content}</p>
+				</div>
+			</div>
+		);
+	}
+
+	// Default/incomplete state - show what's being saved
+	if (content) {
+		return (
+			<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
+				<div className="flex size-8 items-center justify-center rounded-full bg-muted">
+					<BrainIcon className="size-4 text-muted-foreground" />
+				</div>
+				<div className="flex-1 min-w-0">
+					<div className="flex items-center gap-2">
+						<span className="text-sm text-muted-foreground">Saving memory</span>
+						<CategoryBadge category={category} />
+					</div>
+					<p className="mt-1 truncate text-sm text-muted-foreground">{content}</p>
+				</div>
+			</div>
+		);
+	}
+
+	return null;
+};
 
 // ============================================================================
 // Recall Memory Tool UI
 // ============================================================================
 
-export const RecallMemoryToolUI = makeAssistantToolUI<RecallMemoryArgs, RecallMemoryResult>({
-	toolName: "recall_memory",
-	render: function RecallMemoryUI({ args, result, status }) {
-		const isRunning = status.type === "running" || status.type === "requires-action";
-		const isComplete = status.type === "complete";
-		const isError = result?.status === "error";
+export const RecallMemoryToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps<RecallMemoryArgs, RecallMemoryResult>) => {
+	const isRunning = status.type === "running" || status.type === "requires-action";
+	const isComplete = status.type === "complete";
+	const isError = result?.status === "error";
 
-		// Parse args safely
-		const parsedArgs = RecallMemoryArgsSchema.safeParse(args);
-		const query = parsedArgs.success ? parsedArgs.data.query : null;
+	// Parse args safely
+	const parsedArgs = RecallMemoryArgsSchema.safeParse(args);
+	const query = parsedArgs.success ? parsedArgs.data.query : null;
 
-		// Loading state
-		if (isRunning) {
-			return (
-				<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
-					<div className="flex size-8 items-center justify-center rounded-full bg-primary/10">
-						<Loader2Icon className="size-4 animate-spin text-primary" />
-					</div>
-					<div className="flex-1">
-						<span className="text-sm text-muted-foreground">
-							{query ? `Searching memories for "${query}"...` : "Recalling memories..."}
-						</span>
-					</div>
+	// Loading state
+	if (isRunning) {
+		return (
+			<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
+				<div className="flex size-8 items-center justify-center rounded-full bg-primary/10">
+					<Loader2Icon className="size-4 animate-spin text-primary" />
 				</div>
-			);
-		}
-
-		// Error state
-		if (isError) {
-			return (
-				<div className="my-3 flex items-center gap-3 rounded-lg border border-destructive/20 bg-destructive/5 px-4 py-3">
-					<div className="flex size-8 items-center justify-center rounded-full bg-destructive/10">
-						<XIcon className="size-4 text-destructive" />
-					</div>
-					<div className="flex-1">
-						<span className="text-sm text-destructive">Failed to recall memories</span>
-						{result?.error && <p className="mt-1 text-xs text-destructive/70">{result.error}</p>}
-					</div>
+				<div className="flex-1">
+					<span className="text-sm text-muted-foreground">
+						{query ? `Searching memories for "${query}"...` : "Recalling memories..."}
+					</span>
 				</div>
-			);
-		}
+			</div>
+		);
+	}
 
-		// Success state with memories
-		if (isComplete && result?.status === "success") {
-			const memories = result.memories || [];
-			const count = result.count || 0;
-
-			if (count === 0) {
-				return (
-					<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
-						<div className="flex size-8 items-center justify-center rounded-full bg-muted">
-							<SearchIcon className="size-4 text-muted-foreground" />
-						</div>
-						<span className="text-sm text-muted-foreground">No memories found</span>
-					</div>
-				);
-			}
-
-			return (
-				<div className="my-3 rounded-lg border bg-card/60 px-4 py-3">
-					<div className="flex items-center gap-2 mb-2">
-						<BrainIcon className="size-4 text-primary" />
-						<span className="text-sm font-medium text-foreground">
-							Recalled {count} {count === 1 ? "memory" : "memories"}
-						</span>
-					</div>
-					<div className="space-y-2">
-						{memories.slice(0, 5).map((memory: MemoryItem) => (
-							<div
-								key={memory.id}
-								className="flex items-start gap-2 rounded-md bg-muted/50 px-3 py-2"
-							>
-								<CategoryBadge category={memory.category} />
-								<span className="text-sm text-muted-foreground flex-1">{memory.memory_text}</span>
-							</div>
-						))}
-						{memories.length > 5 && (
-							<p className="text-xs text-muted-foreground">...and {memories.length - 5} more</p>
-						)}
-					</div>
+	// Error state
+	if (isError) {
+		return (
+			<div className="my-3 flex items-center gap-3 rounded-lg border border-destructive/20 bg-destructive/5 px-4 py-3">
+				<div className="flex size-8 items-center justify-center rounded-full bg-destructive/10">
+					<XIcon className="size-4 text-destructive" />
 				</div>
-			);
-		}
+				<div className="flex-1">
+					<span className="text-sm text-destructive">Failed to recall memories</span>
+					{result?.error && <p className="mt-1 text-xs text-destructive/70">{result.error}</p>}
+				</div>
+			</div>
+		);
+	}
 
-		// Default/incomplete state
-		if (query) {
+	// Success state with memories
+	if (isComplete && result?.status === "success") {
+		const memories = result.memories || [];
+		const count = result.count || 0;
+
+		if (count === 0) {
 			return (
 				<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
 					<div className="flex size-8 items-center justify-center rounded-full bg-muted">
 						<SearchIcon className="size-4 text-muted-foreground" />
 					</div>
-					<span className="text-sm text-muted-foreground">Searching memories for "{query}"</span>
+					<span className="text-sm text-muted-foreground">No memories found</span>
 				</div>
 			);
 		}
 
-		return null;
-	},
-});
+		return (
+			<div className="my-3 rounded-lg border bg-card/60 px-4 py-3">
+				<div className="flex items-center gap-2 mb-2">
+					<BrainIcon className="size-4 text-primary" />
+					<span className="text-sm font-medium text-foreground">
+						Recalled {count} {count === 1 ? "memory" : "memories"}
+					</span>
+				</div>
+				<div className="space-y-2">
+					{memories.slice(0, 5).map((memory: MemoryItem) => (
+						<div
+							key={memory.id}
+							className="flex items-start gap-2 rounded-md bg-muted/50 px-3 py-2"
+						>
+							<CategoryBadge category={memory.category} />
+							<span className="text-sm text-muted-foreground flex-1">{memory.memory_text}</span>
+						</div>
+					))}
+					{memories.length > 5 && (
+						<p className="text-xs text-muted-foreground">...and {memories.length - 5} more</p>
+					)}
+				</div>
+			</div>
+		);
+	}
+
+	// Default/incomplete state
+	if (query) {
+		return (
+			<div className="my-3 flex items-center gap-3 rounded-lg border bg-card/60 px-4 py-3">
+				<div className="flex size-8 items-center justify-center rounded-full bg-muted">
+					<SearchIcon className="size-4 text-muted-foreground" />
+				</div>
+				<span className="text-sm text-muted-foreground">Searching memories for "{query}"</span>
+			</div>
+		);
+	}
+
+	return null;
+};
 
 // ============================================================================
 // Exports

surfsense_web/components/tool-ui/video-presentation/combined-player.tsx

@@ -118,7 +118,7 @@ export function CombinedPlayer({ slides }: CombinedPlayerProps) {
 	);
 
 	return (
-		<div className="overflow-hidden rounded-xl border shadow-2xl shadow-purple-500/5">
+		<div className="overflow-hidden rounded-xl">
 			<Player
 				component={CompositionWithScenes}
 				durationInFrames={totalFrames}

surfsense_web/components/tool-ui/video-presentation/errors.ts

@@ -0,0 +1,46 @@
+export function getVideoDownloadErrorToast(err: unknown): { title: string; description: string } {
+	const msg = err instanceof Error ? err.message.toLowerCase() : "";
+
+	if (
+		msg.includes("webcodecs") ||
+		msg.includes("canrendermediaonweb") ||
+		msg.includes("not support")
+	) {
+		return {
+			title: "Browser Not Supported",
+			description: "Video rendering requires Chrome, Edge, or Firefox 130+.",
+		};
+	}
+
+	if (msg.includes("memory") || msg.includes("oom") || msg.includes("allocation")) {
+		return {
+			title: "Out of Memory",
+			description: "The presentation is too large to render. Try closing other tabs.",
+		};
+	}
+
+	return {
+		title: "Download Failed",
+		description: "Something went wrong while rendering. Please try again.",
+	};
+}
+
+export function getPptxExportErrorToast(err: unknown): { title: string; description: string } {
+	const msg = err instanceof Error ? err.message.toLowerCase() : "";
+
+	if (
+		msg.includes("dynamically imported") ||
+		msg.includes("failed to fetch") ||
+		msg.includes("network")
+	) {
+		return {
+			title: "Export Unavailable",
+			description: "Could not load the export module. Check your network and try again.",
+		};
+	}
+
+	return {
+		title: "PPTX Export Failed",
+		description: "Something went wrong while exporting. Please try again.",
+	};
+}

surfsense_web/components/tool-ui/video-presentation/generate-video-presentation.tsx

@@ -1,11 +1,13 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
-import { AlertCircleIcon, Download, Film, Loader2, Presentation, X } from "lucide-react";
+import type { ToolCallMessagePartProps } from "@assistant-ui/react";
+import { Dot, Download, Loader2, Presentation, X } from "lucide-react";
 import { useParams, usePathname } from "next/navigation";
 import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { toast } from "sonner";
 import { z } from "zod";
-import { Spinner } from "@/components/ui/spinner";
+import { TextShimmerLoader } from "@/components/prompt-kit/loader";
+import { Button } from "@/components/ui/button";
 import { baseApiService } from "@/lib/apis/base-api.service";
 import { authenticatedFetch } from "@/lib/auth-utils";
 import { compileCheck, compileToComponent } from "@/lib/remotion/compile-check";
@@ -16,6 +18,7 @@ import {
 	CombinedPlayer,
 	type CompiledSlide,
 } from "./combined-player";
+import { getPptxExportErrorToast, getVideoDownloadErrorToast } from "./errors";
 
 const GenerateVideoPresentationArgsSchema = z.object({
 	source_content: z.string(),
@@ -77,30 +80,10 @@ function parseStatusResponse(data: unknown): VideoPresentationStatusResponse | n
 
 function GeneratingState({ title }: { title: string }) {
 	return (
-		<div className="my-4 overflow-hidden rounded-xl border border-primary/20 bg-linear-to-br from-primary/5 to-primary/10 p-4 sm:p-6">
-			<div className="flex items-center gap-3 sm:gap-4">
-				<div className="relative shrink-0">
-					<div className="flex size-12 sm:size-16 items-center justify-center rounded-full bg-primary/20">
-						<Film className="size-6 sm:size-8 text-primary" />
-					</div>
-					<div className="absolute inset-1 animate-ping rounded-full bg-primary/20" />
-				</div>
-				<div className="flex-1 min-w-0">
-					<h3 className="font-semibold text-foreground text-sm sm:text-lg leading-tight">
-						{title}
-					</h3>
-					<div className="mt-1.5 sm:mt-2 flex items-center gap-1.5 sm:gap-2 text-muted-foreground">
-						<Spinner size="sm" className="size-3 sm:size-4" />
-						<span className="text-xs sm:text-sm">
-							Generating video presentation. This may take a few minutes.
-						</span>
-					</div>
-					<div className="mt-2 sm:mt-3">
-						<div className="h-1 sm:h-1.5 w-full overflow-hidden rounded-full bg-primary/10">
-							<div className="h-full w-1/3 animate-pulse rounded-full bg-primary" />
-						</div>
-					</div>
-				</div>
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-foreground line-clamp-2">{title}</p>
+				<TextShimmerLoader text="Generating video presentation" size="sm" />
 			</div>
 		</div>
 	);
@@ -108,20 +91,14 @@ function GeneratingState({ title }: { title: string }) {
 
 function ErrorState({ title, error }: { title: string; error: string }) {
 	return (
-		<div className="my-4 overflow-hidden rounded-xl border border-destructive/20 bg-destructive/5 p-4 sm:p-6">
-			<div className="flex items-center gap-3 sm:gap-4">
-				<div className="flex size-12 sm:size-16 shrink-0 items-center justify-center rounded-full bg-destructive/10">
-					<AlertCircleIcon className="size-6 sm:size-8 text-destructive" />
-				</div>
-				<div className="flex-1 min-w-0">
-					<h3 className="font-semibold text-foreground text-sm sm:text-base leading-tight">
-						{title}
-					</h3>
-					<p className="mt-1 text-destructive text-xs sm:text-sm">
-						Failed to generate video presentation
-					</p>
-					<p className="mt-1.5 sm:mt-2 text-muted-foreground text-xs sm:text-sm">{error}</p>
-				</div>
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-destructive">Video Generation Failed</p>
+			</div>
+			<div className="mx-5 h-px bg-border/50" />
+			<div className="px-5 py-4">
+				<p className="text-sm font-medium text-foreground line-clamp-2">{title}</p>
+				<p className="text-sm text-muted-foreground mt-1">{error}</p>
 			</div>
 		</div>
 	);
@@ -129,20 +106,10 @@ function ErrorState({ title, error }: { title: string; error: string }) {
 
 function CompilationLoadingState({ title }: { title: string }) {
 	return (
-		<div className="my-4 overflow-hidden rounded-xl border bg-muted/30 p-4 sm:p-6">
-			<div className="flex items-center gap-3 sm:gap-4">
-				<div className="flex size-12 sm:size-16 shrink-0 items-center justify-center rounded-full bg-primary/10">
-					<Film className="size-6 sm:size-8 text-primary/50" />
-				</div>
-				<div className="flex-1 min-w-0">
-					<h3 className="font-semibold text-foreground text-sm sm:text-base leading-tight">
-						{title}
-					</h3>
-					<div className="mt-1.5 sm:mt-2 flex items-center gap-1.5 sm:gap-2 text-muted-foreground">
-						<Spinner size="sm" className="size-3 sm:size-4" />
-						<span className="text-xs sm:text-sm">Compiling scenes...</span>
-					</div>
-				</div>
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-foreground line-clamp-2">{title}</p>
+				<TextShimmerLoader text="Compiling scenes" size="sm" />
 			</div>
 		</div>
 	);
@@ -163,7 +130,6 @@ function VideoPresentationPlayer({
 
 	const [isRendering, setIsRendering] = useState(false);
 	const [renderProgress, setRenderProgress] = useState<number | null>(null);
-	const [renderError, setRenderError] = useState<string | null>(null);
 	const [renderFormat, setRenderFormat] = useState<string | null>(null);
 	const abortControllerRef = useRef<AbortController | null>(null);
 
@@ -277,7 +243,6 @@ function VideoPresentationPlayer({
 
 		setIsRendering(true);
 		setRenderProgress(0);
-		setRenderError(null);
 		setRenderFormat(null);
 
 		const controller = new AbortController();
@@ -346,10 +311,9 @@ function VideoPresentationPlayer({
 			document.body.removeChild(a);
 			URL.revokeObjectURL(url);
 		} catch (err) {
-			if ((err as Error).name === "AbortError") {
-				// User cancelled
-			} else {
-				setRenderError(err instanceof Error ? err.message : "Failed to render video");
+			if ((err as Error).name !== "AbortError") {
+				const { title, description } = getVideoDownloadErrorToast(err);
+				toast.error(title, { description });
 			}
 		} finally {
 			setIsRendering(false);
@@ -367,7 +331,6 @@ function VideoPresentationPlayer({
 
 		setIsPptxExporting(true);
 		setPptxProgress("Preparing...");
-		setRenderError(null);
 
 		try {
 			const { exportToPptx } = await import("dom-to-pptx");
@@ -420,10 +383,11 @@ function VideoPresentationPlayer({
 				fileName: "presentation.pptx",
 			});
 
-			roots.forEach((r) => r.unmount());
+			for (const r of roots) r.unmount();
 			document.body.removeChild(offscreen);
 		} catch (err) {
-			setRenderError(err instanceof Error ? err.message : "Failed to export PPTX");
+			const { title, description } = getPptxExportErrorToast(err);
+			toast.error(title, { description });
 		} finally {
 			setIsPptxExporting(false);
 			setPptxProgress(null);
@@ -439,92 +403,84 @@ function VideoPresentationPlayer({
 	}
 
 	return (
-		<div className="my-4 space-y-3">
-			{/* Title bar with actions */}
-			<div className="flex items-center justify-between flex-wrap gap-2">
-				<div className="flex items-center gap-3 min-w-0">
-					<div className="flex size-8 shrink-0 items-center justify-center rounded-lg bg-primary/10">
-						<Film className="size-4 text-primary" />
-					</div>
-					<div className="min-w-0">
-						<h3 className="text-sm font-semibold text-foreground truncate">{title}</h3>
-						<p className="text-xs text-muted-foreground">
-							{compiledSlides.length} slides &middot; {totalDuration.toFixed(1)}s &middot; {FPS}fps
-						</p>
-					</div>
-				</div>
-
-				<div className="flex items-center gap-2">
-					{isRendering ? (
-						<>
-							<div className="flex items-center gap-2 rounded-lg border bg-card px-3 py-1.5">
-								<Loader2 className="size-3.5 animate-spin text-primary" />
-								<span className="text-xs font-medium">
-									Rendering {renderFormat ?? ""}{" "}
-									{renderProgress !== null ? `${Math.round(renderProgress * 100)}%` : "..."}
-								</span>
-								<div className="h-1.5 w-20 overflow-hidden rounded-full bg-secondary">
-									<div
-										className="h-full rounded-full bg-primary transition-all duration-300"
-										style={{ width: `${(renderProgress ?? 0) * 100}%` }}
-									/>
-								</div>
-							</div>
-							<button
-								onClick={handleCancelRender}
-								className="rounded-lg border p-1.5 text-muted-foreground transition-colors hover:bg-secondary hover:text-foreground"
-								title="Cancel render"
-								type="button"
-							>
-								<X className="size-3.5" />
-							</button>
-						</>
-					) : (
-						<>
-							<button
-								onClick={handleDownload}
-								className="inline-flex items-center gap-1.5 rounded-lg border bg-card px-3 py-1.5 text-xs font-medium text-foreground transition-colors hover:bg-secondary"
-								type="button"
-							>
-								<Download className="size-3.5" />
-								Download MP4
-							</button>
-							<button
-								onClick={handleDownloadPPTX}
-								disabled={isPptxExporting}
-								className="inline-flex items-center gap-1.5 rounded-lg border bg-card px-3 py-1.5 text-xs font-medium text-foreground transition-colors hover:bg-secondary disabled:opacity-50 disabled:cursor-not-allowed"
-								type="button"
-							>
-								{isPptxExporting ? (
-									<>
-										<Loader2 className="size-3.5 animate-spin" />
-										{pptxProgress ?? "Exporting..."}
-									</>
-								) : (
-									<>
-										<Presentation className="size-3.5" />
-										Download PPTX
-									</>
-								)}
-							</button>
-						</>
-					)}
-				</div>
+		<div className="my-4 max-w-2xl overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			{/* Header */}
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-foreground line-clamp-2">{title}</p>
+				<p className="text-xs text-muted-foreground mt-0.5 flex items-center">
+					{compiledSlides.length} slides <Dot className="size-4" /> {totalDuration.toFixed(1)}s{" "}
+					<Dot className="size-4" /> {FPS}fps
+				</p>
 			</div>
 
-			{/* Render error */}
-			{renderError && (
-				<div className="flex items-start gap-3 rounded-xl border border-destructive/20 bg-destructive/5 p-3">
-					<AlertCircleIcon className="mt-0.5 size-4 shrink-0 text-destructive" />
-					<div>
-						<p className="text-sm font-medium text-destructive">Download Failed</p>
-						<p className="mt-1 text-xs text-destructive/70 whitespace-pre-wrap">{renderError}</p>
-					</div>
-				</div>
-			)}
+			<div className="mx-5 h-px bg-border/50" />
 
-			{/* Combined Remotion Player */}
-			<CombinedPlayer slides={compiledSlides} />
+			{/* Remotion Player */}
+			<div className="px-5 pt-3">
+				<CombinedPlayer slides={compiledSlides} />
+			</div>
+
+			<div className="mx-5 mt-3 h-px bg-border/50" />
+
+			{/* Action buttons */}
+			<div className="px-5 py-3 flex items-center gap-2 flex-wrap">
+				{isRendering ? (
+					<>
+						<div className="flex items-center gap-2">
+							<Loader2 className="size-3.5 animate-spin text-muted-foreground" />
+							<span className="text-xs font-medium text-muted-foreground">
+								Rendering {renderFormat ?? ""}{" "}
+								{renderProgress !== null ? `${Math.round(renderProgress * 100)}%` : "..."}
+							</span>
+							<div className="h-1.5 w-20 overflow-hidden rounded-full bg-muted">
+								<div
+									className="h-full rounded-full bg-muted-foreground/60 transition-all duration-300"
+									style={{ width: `${(renderProgress ?? 0) * 100}%` }}
+								/>
+							</div>
+						</div>
+						<Button
+							variant="ghost"
+							size="icon"
+							onClick={handleCancelRender}
+							className="size-7 text-muted-foreground"
+						>
+							<X className="size-3.5" />
+						</Button>
+					</>
+				) : (
+					<>
+						<Button
+							variant="ghost"
+							size="sm"
+							onClick={handleDownload}
+							className="gap-1.5 h-7 px-2.5 text-xs text-muted-foreground"
+						>
+							<Download className="size-3.5" />
+							Download MP4
+						</Button>
+						<Button
+							variant="ghost"
+							size="sm"
+							onClick={handleDownloadPPTX}
+							disabled={isPptxExporting}
+							className="gap-1.5 h-7 px-2.5 text-xs text-muted-foreground"
+						>
+							{isPptxExporting ? (
+								<>
+									<Loader2 className="size-3.5 animate-spin" />
+									{pptxProgress ?? "Exporting..."}
+								</>
+							) : (
+								<>
+									<Presentation className="size-3.5" />
+									Download PPTX
+								</>
+							)}
+						</Button>
+					</>
+				)}
+			</div>
 		</div>
 	);
 }
@@ -595,92 +551,85 @@ function StatusPoller({
 	return <ErrorState title={title} error="Unexpected state" />;
 }
 
-export const GenerateVideoPresentationToolUI = makeAssistantToolUI<
-	GenerateVideoPresentationArgs,
-	GenerateVideoPresentationResult
->({
-	toolName: "generate_video_presentation",
-	render: function GenerateVideoPresentationUI({ args, result, status }) {
-		const params = useParams();
-		const pathname = usePathname();
-		const isPublicRoute = pathname?.startsWith("/public/");
-		const shareToken = isPublicRoute && typeof params?.token === "string" ? params.token : null;
+export const GenerateVideoPresentationToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps<GenerateVideoPresentationArgs, GenerateVideoPresentationResult>) => {
+	const params = useParams();
+	const pathname = usePathname();
+	const isPublicRoute = pathname?.startsWith("/public/");
+	const shareToken = isPublicRoute && typeof params?.token === "string" ? params.token : null;
 
-		const title = args.video_title || "SurfSense Presentation";
+	const title = args.video_title || "SurfSense Presentation";
 
-		if (status.type === "running" || status.type === "requires-action") {
-			return <GeneratingState title={title} />;
-		}
+	if (status.type === "running" || status.type === "requires-action") {
+		return <GeneratingState title={title} />;
+	}
 
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return (
-					<div className="my-4 rounded-xl border border-muted p-3 sm:p-4 text-muted-foreground">
-						<p className="flex items-center gap-1.5 sm:gap-2 text-xs sm:text-sm">
-							<Film className="size-3.5 sm:size-4" />
-							<span className="line-through">Presentation generation cancelled</span>
-						</p>
-					</div>
-				);
-			}
-			if (status.reason === "error") {
-				return (
-					<ErrorState
-						title={title}
-						error={typeof status.error === "string" ? status.error : "An error occurred"}
-					/>
-				);
-			}
-		}
-
-		if (!result) {
-			return <GeneratingState title={title} />;
-		}
-
-		if (result.status === "failed") {
-			return <ErrorState title={title} error={result.error || "Generation failed"} />;
-		}
-
-		if (result.status === "generating") {
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
 			return (
-				<div className="my-4 overflow-hidden rounded-xl border border-amber-500/20 bg-amber-500/5 p-3 sm:p-4">
-					<div className="flex items-center gap-2.5 sm:gap-3">
-						<div className="flex size-8 sm:size-10 shrink-0 items-center justify-center rounded-full bg-amber-500/20">
-							<Film className="size-4 sm:size-5 text-amber-500" />
-						</div>
-						<div className="min-w-0">
-							<p className="text-amber-600 dark:text-amber-400 text-xs sm:text-sm font-medium">
-								Presentation already in progress
-							</p>
-							<p className="text-muted-foreground text-[10px] sm:text-xs mt-0.5">
-								Please wait for the current presentation to complete.
-							</p>
-						</div>
+				<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+					<div className="px-5 pt-5 pb-4">
+						<p className="text-sm font-semibold text-muted-foreground">Presentation Cancelled</p>
+						<p className="text-xs text-muted-foreground mt-0.5">
+							Presentation generation was cancelled
+						</p>
 					</div>
 				</div>
 			);
 		}
-
-		if (result.status === "pending" && result.video_presentation_id) {
+		if (status.reason === "error") {
 			return (
-				<StatusPoller
-					presentationId={result.video_presentation_id}
-					title={result.title || title}
-					shareToken={shareToken}
+				<ErrorState
+					title={title}
+					error={typeof status.error === "string" ? status.error : "An error occurred"}
 				/>
 			);
 		}
+	}
 
-		if (result.status === "ready" && result.video_presentation_id) {
-			return (
-				<VideoPresentationPlayer
-					presentationId={result.video_presentation_id}
-					title={result.title || title}
-					shareToken={shareToken}
-				/>
-			);
-		}
+	if (!result) {
+		return <GeneratingState title={title} />;
+	}
 
-		return <ErrorState title={title} error="Missing presentation ID" />;
-	},
-});
+	if (result.status === "failed") {
+		return <ErrorState title={title} error={result.error || "Generation failed"} />;
+	}
+
+	if (result.status === "generating") {
+		return (
+			<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+				<div className="px-5 pt-5 pb-4">
+					<p className="text-sm font-semibold text-foreground">Presentation already in progress</p>
+					<p className="text-xs text-muted-foreground mt-0.5">
+						Please wait for the current presentation to complete.
+					</p>
+				</div>
+			</div>
+		);
+	}
+
+	if (result.status === "pending" && result.video_presentation_id) {
+		return (
+			<StatusPoller
+				presentationId={result.video_presentation_id}
+				title={result.title || title}
+				shareToken={shareToken}
+			/>
+		);
+	}
+
+	if (result.status === "ready" && result.video_presentation_id) {
+		return (
+			<VideoPresentationPlayer
+				presentationId={result.video_presentation_id}
+				title={result.title || title}
+				shareToken={shareToken}
+			/>
+		);
+	}
+
+	return <ErrorState title={title} error="Missing presentation ID" />;
+};

surfsense_web/components/tool-ui/write-todos.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI, useAssistantState } from "@assistant-ui/react";
+import { type ToolCallMessagePartProps, useAuiState } from "@assistant-ui/react";
 import { useAtomValue, useSetAtom } from "jotai";
 import { useEffect, useMemo } from "react";
 import { z } from "zod";
@@ -63,96 +63,98 @@ function WriteTodosLoading() {
  * only the FIRST component renders. Subsequent updates just update the
  * shared state, and the first component reads from it.
  */
-export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosData, WriteTodosData>({
-	toolName: "write_todos",
-	render: function WriteTodosUI({ args, result, status, toolCallId }) {
-		const updatePlanState = useSetAtom(updatePlanStateAtom);
-		const planStates = useAtomValue(planStatesAtom);
+export const WriteTodosToolUI = ({
+	args,
+	result,
+	status,
+	toolCallId,
+}: ToolCallMessagePartProps<WriteTodosData, WriteTodosData>) => {
+	const updatePlanState = useSetAtom(updatePlanStateAtom);
+	const planStates = useAtomValue(planStatesAtom);
 
-		// Check if the THREAD is running
-		const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
+	// Check if the THREAD is running
+	const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
 
-		// Use result if available, otherwise args (for streaming)
-		const data = result || args;
-		const hasTodos = data?.todos && data.todos.length > 0;
+	// Use result if available, otherwise args (for streaming)
+	const data = result || args;
+	const hasTodos = data?.todos && data.todos.length > 0;
 
-		// Fixed title for all plans in conversation
-		const planTitle = "Plan";
+	// Fixed title for all plans in conversation
+	const planTitle = "Plan";
 
-		// SYNCHRONOUS ownership check
-		const isOwner = useMemo(() => {
-			return registerPlanOwner(planTitle, toolCallId);
-		}, [planTitle, toolCallId]);
+	// SYNCHRONOUS ownership check
+	const isOwner = useMemo(() => {
+		return registerPlanOwner(planTitle, toolCallId);
+	}, [planTitle, toolCallId]);
 
-		// Get canonical title
-		const canonicalTitle = useMemo(() => getCanonicalPlanTitle(planTitle), [planTitle]);
+	// Get canonical title
+	const canonicalTitle = useMemo(() => getCanonicalPlanTitle(planTitle), [planTitle]);
 
-		// Register/update the plan state
-		useEffect(() => {
-			if (hasTodos) {
-				const normalizedPlan = parseSerializablePlan({ todos: data.todos });
-				updatePlanState({
-					id: normalizedPlan.id,
-					title: canonicalTitle,
-					todos: normalizedPlan.todos,
-					toolCallId,
-				});
-			}
-		}, [data, hasTodos, canonicalTitle, updatePlanState, toolCallId]);
-
-		// Get the current plan state
-		const currentPlanState = planStates.get(canonicalTitle);
-
-		// If we're NOT the owner, render nothing
-		if (!isOwner) {
-			return null;
+	// Register/update the plan state
+	useEffect(() => {
+		if (hasTodos) {
+			const normalizedPlan = parseSerializablePlan({ todos: data.todos });
+			updatePlanState({
+				id: normalizedPlan.id,
+				title: canonicalTitle,
+				todos: normalizedPlan.todos,
+				toolCallId,
+			});
 		}
+	}, [data, hasTodos, canonicalTitle, updatePlanState, toolCallId]);
 
-		// Loading state
-		if (status.type === "running" || status.type === "requires-action") {
-			if (hasTodos) {
-				const plan = parseSerializablePlan({ todos: data.todos });
-				return (
-					<div className="my-4">
-						<PlanErrorBoundary>
-							<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
-						</PlanErrorBoundary>
-					</div>
-				);
-			}
-			return <WriteTodosLoading />;
+	// Get the current plan state
+	const currentPlanState = planStates.get(canonicalTitle);
+
+	// If we're NOT the owner, render nothing
+	if (!isOwner) {
+		return null;
+	}
+
+	// Loading state
+	if (status.type === "running" || status.type === "requires-action") {
+		if (hasTodos) {
+			const plan = parseSerializablePlan({ todos: data.todos });
+			return (
+				<div className="my-4">
+					<PlanErrorBoundary>
+						<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
+					</PlanErrorBoundary>
+				</div>
+			);
 		}
+		return <WriteTodosLoading />;
+	}
 
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (currentPlanState || hasTodos) {
-				const plan = currentPlanState || parseSerializablePlan({ todos: data?.todos || [] });
-				return (
-					<div className="my-4">
-						<PlanErrorBoundary>
-							<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
-						</PlanErrorBoundary>
-					</div>
-				);
-			}
-			return null;
+	// Incomplete/cancelled state
+	if (status.type === "incomplete") {
+		if (currentPlanState || hasTodos) {
+			const plan = currentPlanState || parseSerializablePlan({ todos: data?.todos || [] });
+			return (
+				<div className="my-4">
+					<PlanErrorBoundary>
+						<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
+					</PlanErrorBoundary>
+				</div>
+			);
 		}
+		return null;
+	}
 
-		// Success - render the plan
-		const planToRender =
-			currentPlanState || (hasTodos ? parseSerializablePlan({ todos: data.todos }) : null);
-		if (!planToRender) {
-			return <WriteTodosLoading />;
-		}
+	// Success - render the plan
+	const planToRender =
+		currentPlanState || (hasTodos ? parseSerializablePlan({ todos: data.todos }) : null);
+	if (!planToRender) {
+		return <WriteTodosLoading />;
+	}
 
-		return (
-			<div className="my-4">
-				<PlanErrorBoundary>
-					<Plan {...planToRender} showProgress={true} isStreaming={isThreadRunning} />
-				</PlanErrorBoundary>
-			</div>
-		);
-	},
-});
+	return (
+		<div className="my-4">
+			<PlanErrorBoundary>
+				<Plan {...planToRender} showProgress={true} isStreaming={isThreadRunning} />
+			</PlanErrorBoundary>
+		</div>
+	);
+};
 
 export { WriteTodosSchema, type WriteTodosData };

surfsense_web/contracts/enums/toolIcons.tsx

@@ -5,8 +5,6 @@ import {
 	FileText,
 	Film,
 	Globe,
-	ImageIcon,
-	Link2,
 	type LucideIcon,
 	Podcast,
 	ScanLine,
@@ -19,8 +17,6 @@ const TOOL_ICONS: Record<string, LucideIcon> = {
 	generate_podcast: Podcast,
 	generate_video_presentation: Film,
 	generate_report: FileText,
-	link_preview: Link2,
-	display_image: ImageIcon,
 	generate_image: Sparkles,
 	scrape_webpage: ScanLine,
 	web_search: Globe,

surfsense_web/lib/chat/message-utils.ts

@@ -2,8 +2,8 @@ import type { ThreadMessageLike } from "@assistant-ui/react";
 import type { MessageRecord } from "./thread-persistence";
 
 /**
- * Convert backend message to assistant-ui ThreadMessageLike format
- * Filters out 'thinking-steps' part as it's handled separately via messageThinkingSteps
+ * Convert backend message to assistant-ui ThreadMessageLike format.
+ * Migrates legacy `thinking-steps` parts to `data-thinking-steps` (assistant-ui data parts).
  */
 export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
 	let content: ThreadMessageLike["content"];
@@ -11,26 +11,34 @@ export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
 	if (typeof msg.content === "string") {
 		content = [{ type: "text", text: msg.content }];
 	} else if (Array.isArray(msg.content)) {
-		// Filter out custom metadata parts - they're handled separately
-		const filteredContent = msg.content.filter((part: unknown) => {
-			if (typeof part !== "object" || part === null || !("type" in part)) return true;
-			const partType = (part as { type: string }).type;
-			// Filter out metadata parts not directly renderable by assistant-ui
-			return (
-				partType !== "thinking-steps" &&
-				partType !== "mentioned-documents" &&
-				partType !== "attachments"
-			);
-		});
+		const convertedContent = msg.content
+			.filter((part: unknown) => {
+				if (typeof part !== "object" || part === null || !("type" in part)) return true;
+				const partType = (part as { type: string }).type;
+				return partType !== "mentioned-documents" && partType !== "attachments";
+			})
+			.map((part: unknown) => {
+				if (
+					typeof part === "object" &&
+					part !== null &&
+					"type" in part &&
+					(part as { type: string }).type === "thinking-steps"
+				) {
+					return {
+						type: "data-thinking-steps",
+						data: { steps: (part as { steps: unknown[] }).steps ?? [] },
+					};
+				}
+				return part;
+			});
 		content =
-			filteredContent.length > 0
-				? (filteredContent as ThreadMessageLike["content"])
+			convertedContent.length > 0
+				? (convertedContent as ThreadMessageLike["content"])
 				: [{ type: "text", text: "" }];
 	} else {
 		content = [{ type: "text", text: String(msg.content) }];
 	}
 
-	// Build metadata.custom for author display in shared chats
 	const metadata = msg.author_id
 		? {
 				custom: {

surfsense_web/lib/chat/streaming-state.ts

@@ -15,6 +15,10 @@ export type ContentPart =
 			toolName: string;
 			args: Record<string, unknown>;
 			result?: unknown;
+	  }
+	| {
+			type: "data-thinking-steps";
+			data: { steps: ThinkingStepData[] };
 	  };
 
 export interface ContentPartsState {
@@ -23,6 +27,32 @@ export interface ContentPartsState {
 	toolCallIndices: Map<string, number>;
 }
 
+export function updateThinkingSteps(
+	state: ContentPartsState,
+	steps: Map<string, ThinkingStepData>
+): void {
+	const stepsArray = Array.from(steps.values());
+	const existingIdx = state.contentParts.findIndex((p) => p.type === "data-thinking-steps");
+
+	if (existingIdx >= 0) {
+		state.contentParts[existingIdx] = {
+			type: "data-thinking-steps",
+			data: { steps: stepsArray },
+		};
+	} else {
+		state.contentParts.unshift({
+			type: "data-thinking-steps",
+			data: { steps: stepsArray },
+		});
+		if (state.currentTextPartIndex >= 0) {
+			state.currentTextPartIndex += 1;
+		}
+		for (const [id, idx] of state.toolCallIndices) {
+			state.toolCallIndices.set(id, idx + 1);
+		}
+	}
+}
+
 export function appendText(state: ContentPartsState, delta: string): void {
 	if (
 		state.currentTextPartIndex >= 0 &&
@@ -75,6 +105,7 @@ export function buildContentForUI(
 	const filtered = state.contentParts.filter((part) => {
 		if (part.type === "text") return part.text.length > 0;
 		if (part.type === "tool-call") return toolsWithUI.has(part.toolName);
+		if (part.type === "data-thinking-steps") return true;
 		return false;
 	});
 	return filtered.length > 0
@@ -84,23 +115,17 @@ export function buildContentForUI(
 
 export function buildContentForPersistence(
 	state: ContentPartsState,
-	toolsWithUI: Set<string>,
-	currentThinkingSteps: Map<string, ThinkingStepData>
+	toolsWithUI: Set<string>
 ): unknown[] {
 	const parts: unknown[] = [];
 
-	if (currentThinkingSteps.size > 0) {
-		parts.push({
-			type: "thinking-steps",
-			steps: Array.from(currentThinkingSteps.values()),
-		});
-	}
-
 	for (const part of state.contentParts) {
 		if (part.type === "text" && part.text.length > 0) {
 			parts.push(part);
 		} else if (part.type === "tool-call" && toolsWithUI.has(part.toolName)) {
 			parts.push(part);
+		} else if (part.type === "data-thinking-steps") {
+			parts.push(part);
 		}
 	}
 

surfsense_web/package.json

@@ -23,9 +23,8 @@
 	"dependencies": {
 		"@ai-sdk/react": "^1.2.12",
 		"@ariakit/react": "^0.4.21",
-		"@assistant-ui/react": "^0.11.53",
-		"@assistant-ui/react-ai-sdk": "^1.1.20",
-		"@assistant-ui/react-markdown": "^0.11.9",
+		"@assistant-ui/react": "^0.12.19",
+		"@assistant-ui/react-markdown": "^0.12.6",
 		"@babel/standalone": "^7.29.2",
 		"@hookform/resolvers": "^5.2.2",
 		"@number-flow/react": "^0.5.10",

surfsense_web/pnpm-lock.yaml

@@ -15,14 +15,11 @@ importers:
         specifier: ^0.4.21
         version: 0.4.21(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
       '@assistant-ui/react':
-        specifier: ^0.11.53
-        version: 0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))
-      '@assistant-ui/react-ai-sdk':
-        specifier: ^1.1.20
-        version: 1.3.8(@assistant-ui/react@0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))(@types/react@19.2.14)(assistant-cloud@0.1.18)(react@19.2.4)
+        specifier: ^0.12.19
+        version: 0.12.19(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))
       '@assistant-ui/react-markdown':
-        specifier: ^0.11.9
-        version: 0.11.10(@assistant-ui/react@0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
+        specifier: ^0.12.6
+        version: 0.12.6(@assistant-ui/react@0.12.19(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
       '@babel/standalone':
         specifier: ^7.29.2
         version: 7.29.2
@@ -432,32 +429,16 @@ importers:
 
 packages:
 
-  '@ai-sdk/gateway@3.0.53':
-    resolution: {integrity: sha512-QT3FEoNARMRlk8JJVR7L98exiK9C8AGfrEJVbRxBT1yIXKs/N19o/+PsjTRVsARgDJNcy9JbJp1FspKucEat0Q==}
-    engines: {node: '>=18'}
-    peerDependencies:
-      zod: ^3.25.76 || ^4.1.8
-
   '@ai-sdk/provider-utils@2.2.8':
     resolution: {integrity: sha512-fqhG+4sCVv8x7nFzYnFo19ryhAa3w096Kmc3hWxMQfW/TubPOmt3A6tYZhl4mUfQWWQMsuSkLrtjlWuXBVSGQA==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.23.8
 
-  '@ai-sdk/provider-utils@4.0.15':
-    resolution: {integrity: sha512-8XiKWbemmCbvNN0CLR9u3PQiet4gtEVIrX4zzLxnCj06AwsEDJwJVBbKrEI4t6qE8XRSIvU2irka0dcpziKW6w==}
-    engines: {node: '>=18'}
-    peerDependencies:
-      zod: ^3.25.76 || ^4.1.8
-
   '@ai-sdk/provider@1.1.3':
     resolution: {integrity: sha512-qZMxYJ0qqX/RfnuIaab+zp8UAeJn/ygXXAffR5I4N0n1IrvA6qBsjc8hXLmBiMV2zoXlifkacF7sEFnYnjBcqg==}
     engines: {node: '>=18'}
 
-  '@ai-sdk/provider@3.0.8':
-    resolution: {integrity: sha512-oGMAgGoQdBXbZqNG0Ze56CHjDZ1IDYOwGYxYjO5KLSlz5HiNQ9udIXsPZ61VWaHGZ5XW/jyjmr6t2xz2jGVwbQ==}
-    engines: {node: '>=18'}
-
   '@ai-sdk/react@1.2.12':
     resolution: {integrity: sha512-jK1IZZ22evPZoQW3vlkZ7wvjYGYF+tRBKXtrcolduIkQ/m/sOAVcVeVDUDvh1T91xCnWCdUGCPZg2avZ90mv3g==}
     engines: {node: '>=18'}
@@ -468,12 +449,6 @@ packages:
       zod:
         optional: true
 
-  '@ai-sdk/react@3.0.99':
-    resolution: {integrity: sha512-xMsp5br4Dpr/3BYq/jrE8q4YLgViU1KHVq8VB0+dzdLJFU3jKA83uoxpbWqzV/edQOBPgGBSb2CgmV5v77rvzA==}
-    engines: {node: '>=18'}
-    peerDependencies:
-      react: ^18 || ~19.0.1 || ~19.1.2 || ^19.2.1
-
   '@ai-sdk/ui-utils@1.2.11':
     resolution: {integrity: sha512-3zcwCc8ezzFlwp3ZD15wAPjf2Au4s3vAbKsXQVyhxODHcmu0iyPO2Eua6D/vicq/AUm/BAo60r97O6HU+EI0+w==}
     engines: {node: '>=18'}
@@ -499,31 +474,37 @@ packages:
       react: ^17.0.0 || ^18.0.0 || ^19.0.0
       react-dom: ^17.0.0 || ^18.0.0 || ^19.0.0
 
-  '@assistant-ui/react-ai-sdk@1.3.8':
-    resolution: {integrity: sha512-LJ2k2r4SYDfH2gmd5xIsu7XBNGucN7ipLgzHmZ4nd8MX8/S/lBmfiNIUko7MPbwbauq6G4KPmRVsiJ5QrqIx6A==}
+  '@assistant-ui/core@0.1.7':
+    resolution: {integrity: sha512-219T42ihVOicbJXZLWgD2CW5Bylg9Nk7geC331X4RfJxTDYlm2zIjViGlGaqfj6URXBp6kMulO2BTUrHGmAvdw==}
     peerDependencies:
-      '@assistant-ui/react': ^0.12.11
+      '@assistant-ui/store': ^0.2.3
+      '@assistant-ui/tap': ^0.5.3
       '@types/react': '*'
-      assistant-cloud: '*'
+      assistant-cloud: ^0.1.22
       react: ^18 || ^19
+      zustand: ^5.0.11
     peerDependenciesMeta:
       '@types/react':
         optional: true
       assistant-cloud:
         optional: true
+      react:
+        optional: true
+      zustand:
+        optional: true
 
-  '@assistant-ui/react-markdown@0.11.10':
-    resolution: {integrity: sha512-7JFd9/s/ZzOtUAHfrxvij4Ti+4V42FVyjF9veWRUsGKKcw7bBZvBxyb2cBMr93sUf0R1eQHsIV39hZjil8J7lw==}
+  '@assistant-ui/react-markdown@0.12.6':
+    resolution: {integrity: sha512-utJqsdDXB3UVZfOa3ErLpaTHraeXkDshR0D34shWdTHrmLyx9e/HypTu4+BgiSsxS+ME6t9WO9M3VeGDprfUcQ==}
     peerDependencies:
-      '@assistant-ui/react': ^0.11.58
+      '@assistant-ui/react': ^0.12.19
       '@types/react': '*'
       react: ^18 || ^19
     peerDependenciesMeta:
       '@types/react':
         optional: true
 
-  '@assistant-ui/react@0.11.58':
-    resolution: {integrity: sha512-5VbparS71X36Q7g+mHwXZvo4eaJohKkQzMP8jBZD9V/Bl26I8s/s3q9WjRqYWMRWaiyYaoEgnQhESM9yyBtW2g==}
+  '@assistant-ui/react@0.12.19':
+    resolution: {integrity: sha512-scAf0o8cwjuHT9Y44EFGXcE2y6BSmpeMvt0NxOn8+Y/HBlNttQMLNvrM0p2AjacXCUufagiafAnWybzBV3nKEQ==}
     peerDependencies:
       '@types/react': '*'
       '@types/react-dom': '*'
@@ -535,8 +516,18 @@ packages:
       '@types/react-dom':
         optional: true
 
-  '@assistant-ui/tap@0.3.6':
-    resolution: {integrity: sha512-4IAN32J9820qbwdc7DeR5HxJVTj+cRPVSMwa9Fv2oP2eMFPAV1eZ8+/co6mgtuM9jSc38vYtZntPsGSHwL7rTg==}
+  '@assistant-ui/store@0.2.3':
+    resolution: {integrity: sha512-daStbgSQiX7+csqK6Cvo7A8p8UZkTCSMxBHxbhJvwrlVbp7BRJWTxq3U3rpTkSGIar23SXIyVRRfXU8VW7pswA==}
+    peerDependencies:
+      '@assistant-ui/tap': ^0.5.3
+      '@types/react': '*'
+      react: ^18 || ^19
+    peerDependenciesMeta:
+      '@types/react':
+        optional: true
+
+  '@assistant-ui/tap@0.5.3':
+    resolution: {integrity: sha512-wy06ksqF2LfFxe4JXy31Ns89N/be1Dy3c+mG363cFHFp3CbLkRu8CrCN2SQSgCkXt628E+D8QyzqdBcl9kD4NQ==}
     peerDependencies:
       '@types/react': '*'
       react: ^18 || ^19
@@ -1089,6 +1080,10 @@ packages:
     resolution: {integrity: sha512-05WQkdpL9COIMz4LjTxGpPNCdlpyimKppYNoJ5Di5EUObifl8t4tuLuUBBZEpoLYOmfvIWrsp9fCl0HoPRVTdA==}
     engines: {node: '>=6.9.0'}
 
+  '@babel/runtime@7.29.2':
+    resolution: {integrity: sha512-JiDShH45zKHWyGe4ZNVRrCjBz8Nh9TMmZG1kh4QTK8hCBTWBi8Da+i7s1fJw7/lYpM4ccepSNfqzZ/QvABBi5g==}
+    engines: {node: '>=6.9.0'}
+
   '@babel/standalone@7.29.2':
     resolution: {integrity: sha512-VSuvywmVRS8efooKrvJzs6BlVSxRvAdLeGrAKUrWoBx1fFBSeE/oBpUZCQ5BcprLyXy04W8skzz7JT8GqlNRJg==}
     engines: {node: '>=6.9.0'}
@@ -4728,10 +4723,6 @@ packages:
     cpu: [x64]
     os: [win32]
 
-  '@vercel/oidc@3.1.0':
-    resolution: {integrity: sha512-Fw28YZpRnA3cAHHDlkt7xQHiJ0fcL+NRcIqsocZQUSmbzeIKRpwttJjik5ZGanXP+vlA4SbTg+AbA3bP363l+w==}
-    engines: {node: '>= 20'}
-
   '@xmldom/xmldom@0.8.11':
     resolution: {integrity: sha512-cQzWCtO6C8TQiYl1ruKNn2U6Ao4o4WBBcbL61yJl84x+j5sOWWFU9X7DpND8XZG3daDppSsigMdfAIl2upQBRw==}
     engines: {node: '>=10.0.0'}
@@ -4768,12 +4759,6 @@ packages:
       react:
         optional: true
 
-  ai@6.0.97:
-    resolution: {integrity: sha512-eZIAcBymwGhBwncRH/v9pillZNMeRCDkc4BwcvwXerXd7sxjVxRis3ZNCNCpP02pVH4NLs81ljm4cElC4vbNcQ==}
-    engines: {node: '>=18'}
-    peerDependencies:
-      zod: ^3.25.76 || ^4.1.8
-
   ajv-formats@2.1.1:
     resolution: {integrity: sha512-Wx0Kx52hxE7C18hkMEggYlEifqWZtYaRgouJor+WMdPnQyEK13vgEWyVNup7SoeeoLMsr4kf5h6dOW11I15MUA==}
     peerDependencies:
@@ -4851,14 +4836,11 @@ packages:
     resolution: {integrity: sha512-BNoCY6SXXPQ7gF2opIP4GBE+Xw7U+pHMYKuzjgCN3GwiaIR09UUeKfheyIry77QtrCBlC0KK0q5/TER/tYh3PQ==}
     engines: {node: '>= 0.4'}
 
-  assistant-cloud@0.1.18:
-    resolution: {integrity: sha512-6tq2jPGIBjkjsLQ/Fd4r6PGj4hf05oM2jBl4hBs7YIkaJ3qBVUWiHary2+faNpsPOoY71brsVukl/qz5B1rQkA==}
+  assistant-cloud@0.1.22:
+    resolution: {integrity: sha512-AEE9shV+oFrGDv/MRTRERctNKpIYS0n34UpAQXXICiOkSWD6QZnS1ljLqruFko7fJoT5CIWq8dNeJWdzQLTBLg==}
 
-  assistant-stream@0.2.47:
-    resolution: {integrity: sha512-0f+yVwoh7GVwYqaWh6vT+P/zflvEyqysJJzGhjqOPxUYjbNOjcifBw+fVwQPtxysyzye2TZCQtmOWjP0ggvnqw==}
-
-  assistant-stream@0.3.3:
-    resolution: {integrity: sha512-Ne/uTseMIiZx740dTbr/SWxONM8nYj4Z5BRmUfqQN+TNgtOCgWOlC/oTUQ+A7LIUHtmGbcoyZwDf8yd2RASnDA==}
+  assistant-stream@0.3.6:
+    resolution: {integrity: sha512-NdtSRrQfWCDA/aqQ1xhobf/xnhuMZkhFAw9xzAt5iAoL3ouxVXOowSRN87OL4MYBQEvqtcjw9/CE6YcsXoBtuw==}
 
   ast-types-flow@0.0.8:
     resolution: {integrity: sha512-OH/2E5Fg20h2aPrbe+QL8JZQFko0YZaF+j4mnQ7BGhfavO7OpSLa8a0y9sBwomHdSbkhTS8TQNayBfnW5DwbvQ==}
@@ -5686,10 +5668,6 @@ packages:
   eventemitter3@5.0.4:
     resolution: {integrity: sha512-mlsTRyGaPBjPedk6Bvw+aqbsXDtoAyAzm5MO7JgU+yVRyMQ5O8bD4Kcci7BS85f93veegeCPkL8R4GLClnjLFw==}
 
-  eventsource-parser@3.0.6:
-    resolution: {integrity: sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg==}
-    engines: {node: '>=18.0.0'}
-
   execa@5.1.1:
     resolution: {integrity: sha512-8uSpZZocAZRBAPIEINJj3Lo9HyGitllczc27Eh5YYojjMFMn8yHMDMaUHE2Jqfq05D/wucwI4JGURyXt1vchyg==}
     engines: {node: '>=10'}
@@ -6944,6 +6922,11 @@ packages:
     engines: {node: ^18 || >=20}
     hasBin: true
 
+  nanoid@5.1.7:
+    resolution: {integrity: sha512-ua3NDgISf6jdwezAheMOk4mbE1LXjm1DfMUDMuJf4AqxLFK3ccGpgWizwa5YV7Yz9EpXwEaWoRXSb/BnV0t5dQ==}
+    engines: {node: ^18 || >=20}
+    hasBin: true
+
   napi-build-utils@2.0.0:
     resolution: {integrity: sha512-GEbrYkbfF7MoNaoh2iGG84Mnf/WZfB0GdGEsM8wz7Expx/LlWf5U8t9nvJKXSp3qr5IsEbK04cBGhol/KwOsWA==}
 
@@ -8238,6 +8221,11 @@ packages:
     peerDependencies:
       react: '>=16.8.0'
 
+  use-effect-event@2.0.3:
+    resolution: {integrity: sha512-fz1en+z3fYXCXx3nMB8hXDMuygBltifNKZq29zDx+xNJ+1vEs6oJlYd9sK31vxJ0YI534VUsHEBY0k2BATsmBQ==}
+    peerDependencies:
+      react: ^18.3 || ^19.0.0-0
+
   use-intl@4.8.3:
     resolution: {integrity: sha512-nLxlC/RH+le6g3amA508Itnn/00mE+J22ui21QhOWo5V9hCEC43+WtnRAITbJW0ztVZphev5X9gvOf2/Dk9PLA==}
     peerDependencies:
@@ -8477,13 +8465,6 @@ packages:
 
 snapshots:
 
-  '@ai-sdk/gateway@3.0.53(zod@4.3.6)':
-    dependencies:
-      '@ai-sdk/provider': 3.0.8
-      '@ai-sdk/provider-utils': 4.0.15(zod@4.3.6)
-      '@vercel/oidc': 3.1.0
-      zod: 4.3.6
-
   '@ai-sdk/provider-utils@2.2.8(zod@4.3.6)':
     dependencies:
       '@ai-sdk/provider': 1.1.3
@@ -8491,21 +8472,10 @@ snapshots:
       secure-json-parse: 2.7.0
       zod: 4.3.6
 
-  '@ai-sdk/provider-utils@4.0.15(zod@4.3.6)':
-    dependencies:
-      '@ai-sdk/provider': 3.0.8
-      '@standard-schema/spec': 1.1.0
-      eventsource-parser: 3.0.6
-      zod: 4.3.6
-
   '@ai-sdk/provider@1.1.3':
     dependencies:
       json-schema: 0.4.0
 
-  '@ai-sdk/provider@3.0.8':
-    dependencies:
-      json-schema: 0.4.0
-
   '@ai-sdk/react@1.2.12(react@19.2.4)(zod@4.3.6)':
     dependencies:
       '@ai-sdk/provider-utils': 2.2.8(zod@4.3.6)
@@ -8516,16 +8486,6 @@ snapshots:
     optionalDependencies:
       zod: 4.3.6
 
-  '@ai-sdk/react@3.0.99(react@19.2.4)(zod@4.3.6)':
-    dependencies:
-      '@ai-sdk/provider-utils': 4.0.15(zod@4.3.6)
-      ai: 6.0.97(zod@4.3.6)
-      react: 19.2.4
-      swr: 2.4.0(react@19.2.4)
-      throttleit: 2.1.0
-    transitivePeerDependencies:
-      - zod
-
   '@ai-sdk/ui-utils@1.2.11(zod@4.3.6)':
     dependencies:
       '@ai-sdk/provider': 1.1.3
@@ -8551,20 +8511,21 @@ snapshots:
       react: 19.2.4
       react-dom: 19.2.4(react@19.2.4)
 
-  '@assistant-ui/react-ai-sdk@1.3.8(@assistant-ui/react@0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))(@types/react@19.2.14)(assistant-cloud@0.1.18)(react@19.2.4)':
+  '@assistant-ui/core@0.1.7(@assistant-ui/store@0.2.3(@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4))(@types/react@19.2.14)(react@19.2.4))(@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4))(@types/react@19.2.14)(assistant-cloud@0.1.22)(react@19.2.4)(zustand@5.0.11(@types/react@19.2.14)(immer@10.2.0)(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))':
     dependencies:
-      '@ai-sdk/react': 3.0.99(react@19.2.4)(zod@4.3.6)
-      '@assistant-ui/react': 0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))
-      ai: 6.0.97(zod@4.3.6)
-      react: 19.2.4
-      zod: 4.3.6
+      '@assistant-ui/store': 0.2.3(@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4))(@types/react@19.2.14)(react@19.2.4)
+      '@assistant-ui/tap': 0.5.3(@types/react@19.2.14)(react@19.2.4)
+      assistant-stream: 0.3.6
+      nanoid: 5.1.7
     optionalDependencies:
       '@types/react': 19.2.14
-      assistant-cloud: 0.1.18
+      assistant-cloud: 0.1.22
+      react: 19.2.4
+      zustand: 5.0.11(@types/react@19.2.14)(immer@10.2.0)(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))
 
-  '@assistant-ui/react-markdown@0.11.10(@assistant-ui/react@0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)':
+  '@assistant-ui/react-markdown@0.12.6(@assistant-ui/react@0.12.19(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)':
     dependencies:
-      '@assistant-ui/react': 0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))
+      '@assistant-ui/react': 0.12.19(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))
       '@radix-ui/react-primitive': 2.1.4(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
       '@radix-ui/react-use-callback-ref': 1.1.1(@types/react@19.2.14)(react@19.2.4)
       classnames: 2.5.1
@@ -8577,21 +8538,21 @@ snapshots:
       - react-dom
       - supports-color
 
-  '@assistant-ui/react@0.11.58(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))':
+  '@assistant-ui/react@0.12.19(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(immer@10.2.0)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4))':
     dependencies:
-      '@assistant-ui/tap': 0.3.6(@types/react@19.2.14)(react@19.2.4)
+      '@assistant-ui/core': 0.1.7(@assistant-ui/store@0.2.3(@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4))(@types/react@19.2.14)(react@19.2.4))(@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4))(@types/react@19.2.14)(assistant-cloud@0.1.22)(react@19.2.4)(zustand@5.0.11(@types/react@19.2.14)(immer@10.2.0)(react@19.2.4)(use-sync-external-store@1.6.0(react@19.2.4)))
+      '@assistant-ui/store': 0.2.3(@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4))(@types/react@19.2.14)(react@19.2.4)
+      '@assistant-ui/tap': 0.5.3(@types/react@19.2.14)(react@19.2.4)
       '@radix-ui/primitive': 1.1.3
       '@radix-ui/react-compose-refs': 1.1.2(@types/react@19.2.14)(react@19.2.4)
       '@radix-ui/react-context': 1.1.3(@types/react@19.2.14)(react@19.2.4)
-      '@radix-ui/react-dropdown-menu': 2.1.16(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
-      '@radix-ui/react-popover': 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
       '@radix-ui/react-primitive': 2.1.4(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
-      '@radix-ui/react-slot': 1.2.4(@types/react@19.2.14)(react@19.2.4)
       '@radix-ui/react-use-callback-ref': 1.1.1(@types/react@19.2.14)(react@19.2.4)
       '@radix-ui/react-use-escape-keydown': 1.1.1(@types/react@19.2.14)(react@19.2.4)
-      assistant-cloud: 0.1.18
-      assistant-stream: 0.2.47
-      nanoid: 5.1.6
+      assistant-cloud: 0.1.22
+      assistant-stream: 0.3.6
+      nanoid: 5.1.7
+      radix-ui: 1.4.3(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4)
       react: 19.2.4
       react-dom: 19.2.4(react@19.2.4)
       react-textarea-autosize: 8.5.9(@types/react@19.2.14)(react@19.2.4)
@@ -8604,7 +8565,15 @@ snapshots:
       - immer
       - use-sync-external-store
 
-  '@assistant-ui/tap@0.3.6(@types/react@19.2.14)(react@19.2.4)':
+  '@assistant-ui/store@0.2.3(@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4))(@types/react@19.2.14)(react@19.2.4)':
+    dependencies:
+      '@assistant-ui/tap': 0.5.3(@types/react@19.2.14)(react@19.2.4)
+      react: 19.2.4
+      use-effect-event: 2.0.3(react@19.2.4)
+    optionalDependencies:
+      '@types/react': 19.2.14
+
+  '@assistant-ui/tap@0.5.3(@types/react@19.2.14)(react@19.2.4)':
     optionalDependencies:
       '@types/react': 19.2.14
       react: 19.2.4
@@ -9318,6 +9287,8 @@ snapshots:
 
   '@babel/runtime@7.28.6': {}
 
+  '@babel/runtime@7.29.2': {}
+
   '@babel/standalone@7.29.2': {}
 
   '@babel/template@7.28.6':
@@ -13074,8 +13045,6 @@ snapshots:
   '@unrs/resolver-binding-win32-x64-msvc@1.11.1':
     optional: true
 
-  '@vercel/oidc@3.1.0': {}
-
   '@xmldom/xmldom@0.8.11': {}
 
   abstract-logging@2.0.1: {}
@@ -13104,14 +13073,6 @@ snapshots:
     optionalDependencies:
       react: 19.2.4
 
-  ai@6.0.97(zod@4.3.6):
-    dependencies:
-      '@ai-sdk/gateway': 3.0.53(zod@4.3.6)
-      '@ai-sdk/provider': 3.0.8
-      '@ai-sdk/provider-utils': 4.0.15(zod@4.3.6)
-      '@opentelemetry/api': 1.9.0
-      zod: 4.3.6
-
   ajv-formats@2.1.1(ajv@8.18.0):
     optionalDependencies:
       ajv: 8.18.0
@@ -13217,20 +13178,14 @@ snapshots:
       get-intrinsic: 1.3.0
       is-array-buffer: 3.0.5
 
-  assistant-cloud@0.1.18:
+  assistant-cloud@0.1.22:
     dependencies:
-      assistant-stream: 0.3.3
+      assistant-stream: 0.3.6
 
-  assistant-stream@0.2.47:
+  assistant-stream@0.3.6:
     dependencies:
       '@standard-schema/spec': 1.1.0
-      nanoid: 5.1.6
-      secure-json-parse: 4.1.0
-
-  assistant-stream@0.3.3:
-    dependencies:
-      '@standard-schema/spec': 1.1.0
-      nanoid: 5.1.6
+      nanoid: 5.1.7
       secure-json-parse: 4.1.0
 
   ast-types-flow@0.0.8: {}
@@ -14214,8 +14169,6 @@ snapshots:
 
   eventemitter3@5.0.4: {}
 
-  eventsource-parser@3.0.6: {}
-
   execa@5.1.1:
     dependencies:
       cross-spawn: 7.0.6
@@ -15765,6 +15718,8 @@ snapshots:
 
   nanoid@5.1.6: {}
 
+  nanoid@5.1.7: {}
+
   napi-build-utils@2.0.0: {}
 
   napi-postinstall@0.3.4: {}
@@ -16448,7 +16403,7 @@ snapshots:
 
   react-textarea-autosize@8.5.9(@types/react@19.2.14)(react@19.2.4):
     dependencies:
-      '@babel/runtime': 7.28.6
+      '@babel/runtime': 7.29.2
       react: 19.2.4
       use-composed-ref: 1.4.0(@types/react@19.2.14)(react@19.2.4)
       use-latest: 1.3.0(@types/react@19.2.14)(react@19.2.4)
@@ -17409,6 +17364,10 @@ snapshots:
       dequal: 2.0.3
       react: 19.2.4
 
+  use-effect-event@2.0.3(react@19.2.4):
+    dependencies:
+      react: 19.2.4
+
   use-intl@4.8.3(react@19.2.4):
     dependencies:
       '@formatjs/fast-memoize': 3.1.0