diff --git a/surfsense_backend/app/agents/new_chat/__init__.py b/surfsense_backend/app/agents/new_chat/__init__.py index eccb7a5c3..4238d17de 100644 --- a/surfsense_backend/app/agents/new_chat/__init__.py +++ b/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 diff --git a/surfsense_backend/app/agents/new_chat/chat_deepagent.py b/surfsense_backend/app/agents/new_chat/chat_deepagent.py index c69ba1063..2857be4a7 100644 --- a/surfsense_backend/app/agents/new_chat/chat_deepagent.py +++ b/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 diff --git a/surfsense_backend/app/agents/new_chat/system_prompt.py b/surfsense_backend/app/agents/new_chat/system_prompt.py index f8ac62787..b53251a1d 100644 --- a/surfsense_backend/app/agents/new_chat/system_prompt.py +++ b/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="", 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="", 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: ` marker. Without it the user cannot download the file. diff --git a/surfsense_backend/app/agents/new_chat/tools/__init__.py b/surfsense_backend/app/agents/new_chat/tools/__init__.py index 5002e69bb..404926d19 100644 --- a/surfsense_backend/app/agents/new_chat/tools/__init__.py +++ b/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", diff --git a/surfsense_backend/app/agents/new_chat/tools/display_image.py b/surfsense_backend/app/agents/new_chat/tools/display_image.py deleted file mode 100644 index 4424cc0d3..000000000 --- a/surfsense_backend/app/agents/new_chat/tools/display_image.py +++ /dev/null @@ -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 diff --git a/surfsense_backend/app/agents/new_chat/tools/generate_image.py b/surfsense_backend/app/agents/new_chat/tools/generate_image.py index 8ffa4ecde..d94d55b1a 100644 --- a/surfsense_backend/app/agents/new_chat/tools/generate_image.py +++ b/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), diff --git a/surfsense_backend/app/agents/new_chat/tools/link_preview.py b/surfsense_backend/app/agents/new_chat/tools/link_preview.py deleted file mode 100644 index 81d91d54c..000000000 --- a/surfsense_backend/app/agents/new_chat/tools/link_preview.py +++ /dev/null @@ -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']+property=["\']og:{property_name}["\'][^>]+content=["\']([^"\']+)["\']' - match = re.search(pattern, html, re.IGNORECASE) - if match: - return match.group(1) - - # Try content before property - pattern = rf']+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']+name=["\']twitter:{name}["\'][^>]+content=["\']([^"\']+)["\']' - ) - match = re.search(pattern, html, re.IGNORECASE) - if match: - return match.group(1) - - # Try content before name - pattern = ( - rf']+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']+name=["\']description["\'][^>]+content=["\']([^"\']+)["\']' - match = re.search(pattern, html, re.IGNORECASE) - if match: - return match.group(1) - - # Try content before name - pattern = r']+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 tag - pattern = r"<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("&", "&") - .replace("<", "<") - .replace(">", ">") - .replace(""", '"') - .replace("'", "'") - .replace("'", "'") - ) - - -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 diff --git a/surfsense_backend/app/agents/new_chat/tools/registry.py b/surfsense_backend/app/agents/new_chat/tools/registry.py index 4ee8023d2..7700d47d3 100644 --- a/surfsense_backend/app/agents/new_chat/tools/registry.py +++ b/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"]) diff --git a/surfsense_backend/app/services/public_chat_service.py b/surfsense_backend/app/services/public_chat_service.py index a75eb73f8..376db974f 100644 --- a/surfsense_backend/app/services/public_chat_service.py +++ b/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", } diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 82d98da18..1f3eaa179 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/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 = { diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index e7da4393f..8578d2dcb 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/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(null); const [messages, setMessages] = useState([]); 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>( - new Map() - ); const abortControllerRef = useRef(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(); - // Extract and restore mentioned documents from persisted messages const restoredDocsMap: Record = {}; - 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( - (messageThinkingSteps.get(assistantMsgId) ?? []).map((s) => [s.id, s]) - ); + const currentThinkingSteps = new Map(); 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 ( - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - {/* Disabled for now */} +
- +
diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index 4af5b07ee..9fefecb1c 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/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 ( -
- -
- ); -}; - const AssistantMessageInner: FC = () => { return ( <> - {/* Render thinking steps from message content - this ensures proper scroll tracking */} - -
null, + multi_link_preview: () => null, + scrape_webpage: () => null, + }, + Fallback: ToolFallback, + }, }} /> @@ -95,7 +139,7 @@ export const AssistantMessage: FC = () => { const messageRef = useRef(null); const commentPanelRef = useRef(null); const commentTriggerRef = useRef(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 ( { > - message.isCopied}> + message.isCopied}> - - !message.isCopied}> + + !message.isCopied}> - + diff --git a/surfsense_web/components/assistant-ui/image.tsx b/surfsense_web/components/assistant-ui/image.tsx new file mode 100644 index 000000000..65059bcdc --- /dev/null +++ b/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; + +function ImageRoot({ className, variant, size, children, ...props }: ImageRootProps) { + return ( +
+ {children} +
+ ); +} + +type ImagePreviewProps = Omit, "children"> & { + containerClassName?: string; +}; + +function ImagePreview({ + className, + containerClassName, + onLoad, + onError, + alt = "Image content", + src, + ...props +}: ImagePreviewProps) { + const imgRef = useRef(null); + const [loadedSrc, setLoadedSrc] = useState(undefined); + const [errorSrc, setErrorSrc] = useState(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 ( +
+ {!loaded && !error && ( +
+ +
+ )} + {error ? ( +
+ +
+ ) : ( + // biome-ignore lint/performance/noImgElement: intentional for dynamic external URLs + {alt} { + if (typeof src === "string") setLoadedSrc(src); + onLoad?.(e); + }} + onError={(e) => { + if (typeof src === "string") setErrorSrc(src); + onError?.(e); + }} + {...props} + /> + )} +
+ ); +} + +function ImageFilename({ className, children, ...props }: React.ComponentProps<"span">) { + if (!children) return null; + + return ( + + {children} + + ); +} + +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 ( + <> + + {isMounted && + isOpen && + createPortal( + , + document.body + )} + + ); +} + +const ImageImpl: ImageMessagePartComponent = ({ image, filename }) => { + return ( + + + + + {filename} + + ); +}; + +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 }; diff --git a/surfsense_web/components/assistant-ui/markdown-text.tsx b/surfsense_web/components/assistant-ui/markdown-text.tsx index 10ce85c5a..3d33463b2 100644 --- a/surfsense_web/components/assistant-ui/markdown-text.tsx +++ b/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 +): Record { + const cleaned: Record = {}; + 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(); @@ -149,7 +178,7 @@ const MarkdownTextImpl = () => { export const MarkdownText = memo(MarkdownTextImpl); -const CodeHeader: FC = ({ 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 = ({ language, code }) => { }; return ( -
- {language} +
+ {language} {!isCopied && } {isCopied && } @@ -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 ) : ( - {parsed} + {parsed} ); } 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 ( +
+ + + + + + +
+
+ {alt && alt !== "Image" && ( +

{alt}

+ )} + {domain &&

{domain}

} +
+ e.stopPropagation()} + > + Open + + +
+
+ ); +} + const defaultComponents = memoizeMarkdownComponents({ h1: ({ className, children, ...props }) => (

), table: ({ className, ...props }) => ( -
- +
+
), + thead: ({ className, ...props }) => ( + + ), + tbody: ({ className, ...props }) => ( + + ), th: ({ className, children, ...props }) => ( - + ), td: ({ className, children, ...props }) => ( - - ), - tr: ({ className, ...props }) => ( - td:first-child]:rounded-bl-lg [&:last-child>td:last-child]:rounded-br-lg", - className - )} - {...props} - /> + ), + tr: ({ className, ...props }) => , sup: ({ className, ...props }) => ( a]:text-xs [&>a]:no-underline", className)} {...props} /> ), - pre: ({ 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 (
+				
+			);
+		}
+		const language = /language-(\w+)/.exec(className || "")?.[1] ?? "text";
+		const codeString = String(children).replace(/\n$/, "");
+		const syntaxStyle = resolvedTheme === "dark" ? cleanMaterialDark : cleanMaterialLight;
 		return (
-			
+			

+				
+				
+					{codeString}
+				
+			

 		);
 	},
 	strong: ({ className, children, ...props }) => (
@@ -371,5 +451,8 @@ const defaultComponents = memoizeMarkdownComponents({
 			{processChildrenWithCitations(children)}
 		
 	),
-	CodeHeader,
+	img: ({ src, alt }) => (
+		
+	),
+	CodeHeader: () => null,
 });
diff --git a/surfsense_web/components/assistant-ui/thinking-steps.tsx b/surfsense_web/components/assistant-ui/thinking-steps.tsx
index c773824f8..900fc7b09 100644
--- a/surfsense_web/components/assistant-ui/thinking-steps.tsx
+++ b/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>(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 (
 		

 			

-				{/* Main collapsible header */}
 				
 
-				{/* Collapsible content with CSS grid animation */}
 				

-										{/* Dot and line column */}
 										

-											{/* Vertical connection line - extends to next dot */}
 											{!isLast && (
 												

 											)}
-											{/* Step dot - on top of line */}
 											

 												{effectiveStatus === "in_progress" ? (
 													
@@ -118,9 +112,7 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
 											

 										

 
-										{/* Step content */}
 										

-											{/* Step title */}
 											

 
-											{/* Step items (sub-content) */}
 											{step.items && step.items.length > 0 && (
 												

-													{step.items.map((item, idx) => (
-														
+													{step.items.map((item) => (
+														
 															{item}
 														
 													))}
@@ -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("");
+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 (
+		

+			
+		

+	);
+}
 
-		// 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,
+});
diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx
index 081e234a8..195afc090 100644
--- a/surfsense_web/components/assistant-ui/thread.tsx
+++ b/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;
-}
-
-export const Thread: FC = ({ messageThinkingSteps = new Map() }) => {
-	return (
-		
-			
-		
-	);
+export const Thread: FC = () => {
+	return ;
 };
 
 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"
 			>
-				 thread.isEmpty}>
+				 thread.isEmpty}>
 					
-				
+				
 
 				 {
 					style={{ paddingBottom: "max(1rem, env(safe-area-inset-bottom))" }}
 				>
 					
-					 !thread.isEmpty}>
+					 !thread.isEmpty}>
 						

 							
 						

-					
+					
 				
 			
 		
@@ -327,11 +303,11 @@ const Composer: FC = () => {
 	const editorContainerRef = useRef(null);
 	const documentPickerRef = useRef(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 = ({ 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 = ({ isBlockedByOtherUser = false
 					
 				)}
 			

-
 			{!hasModelConfigured && (
 				

 					
 					Select a model
 				

 			)}
-
 			

-				 !thread.isRunning}>
+				 !thread.isRunning}>
 					
 						 = ({ isBlockedByOtherUser = false
 							
 						
 					
-				
+				
 
-				 thread.isRunning}>
+				 thread.isRunning}>
 					
 						
 					
-				
+				
 			

 		

 	);
@@ -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 (
-		
-			
-				
-			
-		
-	);
-};
-
-/**
- * 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 (
-		

-			
-		

-	);
-};
-
-const AssistantMessageInner: FC = () => {
-	return (
-		<>
-			{/* Render thinking steps from message content - this ensures proper scroll tracking */}
-			
-
-			

-				
-				
-			

-
-			

-				
-				
-			

-		
-	);
-};
-
-const AssistantActionBar: FC = () => {
-	return (
-		
-			
-				
-					 message.isCopied}>
-						
-					
-					 !message.isCopied}>
-						
-					
-				
-			
-			
-				
-					
-				
-			
-			
-				
-					
-				
-			
-		
-	);
-};
-
 const EditComposer: FC = () => {
 	return (
 		
@@ -1253,30 +1130,3 @@ const EditComposer: FC = () => {
 		
 	);
 };
-
-const BranchPicker: FC = ({ className, ...rest }) => {
-	return (
-		
-			
-				
-					
-				
-			
-			
-				 / 
-			
-			
-				
-					
-				
-			
-		
-	);
-};
diff --git a/surfsense_web/components/assistant-ui/tool-fallback.tsx b/surfsense_web/components/assistant-ui/tool-fallback.tsx
index 636b43c36..89498fbca 100644
--- a/surfsense_web/components/assistant-ui/tool-fallback.tsx
+++ b/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 (
 		

-			

-				{isCancelled ? (
-					
-				) : (
-					
-				)}
-				
 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"
+			>
+				

-					{isCancelled ? "Cancelled tool: " : "Used tool: "}
-					{toolName}
-				

-				
-			

-			{!isCollapsed && (
-				

-					{cancelledReason && (
-						

-							

-								Cancelled reason:
-							

-							

-								{cancelledReason}
-							

-						

-					)}
-					

-						
{argsText}

-					

-					{!isCancelled && result !== undefined && (
-						

-							
Result:

-							
-								{typeof result === "string" ? result : JSON.stringify(result, null, 2)}
-							

-						

+					{isError ? (
+						
+					) : isCancelled ? (
+						
+					) : isRunning ? (
+						
+					) : (
+						
 					)}
 				

+
+				

+					

+						{isRunning
+							? displayName
+							: isCancelled
+								? `Cancelled: ${displayName}`
+								: isError
+									? `Failed: ${displayName}`
+									: displayName}
+					

+					{isRunning && 
Running...
}
+					{cancelledReason && (
+						
{cancelledReason}

+					)}
+					{errorReason && (
+						
{errorReason}

+					)}
+				

+
+				{!isRunning && (
+					

+						{isExpanded ? (
+							
+						) : (
+							
+						)}
+					

+				)}
+			
+
+			{isExpanded && !isRunning && (
+				<>
+					

+					

+						{argsText && (
+							

+								
Arguments

+								
+									{argsText}
+								

+							

+						)}
+						{!isCancelled && result !== undefined && (
+							<>
+								

+								

+									
Result

+									
+										{typeof result === "string" ? result : JSON.stringify(result, null, 2)}
+									

+								

+							
+						)}
+					

+				
 			)}
 		

 	);
diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx
index 1c0525277..c01e8e486 100644
--- a/surfsense_web/components/assistant-ui/user-message.tsx
+++ b/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 = ({ displayName, avatarUrl }) => {
 
 	if (avatarUrl && !hasError) {
 		return (
-			{displayName setHasError(true)}
+				unoptimized
 			/>
 		);
 	}
 
 	return (
-		

+		

 			{initials}
 		

 	);
 };
 
 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 (
 		
-			

-				

-					{/* Display mentioned documents */}
-					{mentionedDocs && mentionedDocs.length > 0 && (
-						

-							{/* Mentioned documents as chips */}
-							{mentionedDocs?.map((doc) => (
-								
-									
-									{doc.title}
-								
-							))}
-						

-					)}
-					{/* Message bubble with action bar positioned relative to it */}
-					

+			

+				

+					

+						{mentionedDocs && mentionedDocs.length > 0 && (
+							

+								{mentionedDocs?.map((doc) => (
+									
+										
+										{doc.title}
+									
+								))}
+							

+						)}
 						

 							
 						

-						

+						

 							
 						

 					

+					{author && (
+						

+							
+						

+					)}
 				

-				{/* User avatar - only shown in shared chats */}
-				{author && (
-					

-						
-					

-				)}
 			

 		
 	);
 };
 
 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 (
 		
-			{/* Only allow editing the last user message */}
+			
+				
+					 message.isCopied}>
+						
+					
+					 !message.isCopied}>
+						
+					
+				
+			
 			{canEdit && (
 				
-					
+					
 						
 					
 				
diff --git a/surfsense_web/components/public-chat/public-chat-view.tsx b/surfsense_web/components/public-chat/public-chat-view.tsx
index 706d47ca6..f8dd6db5a 100644
--- a/surfsense_web/components/public-chat/public-chat-view.tsx
+++ b/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) {
 		

 			
 			
-				{/* Tool UIs for rendering tool results */}
-				
-				
-				
-				
-				
-				
-
+				
 				

 					

 						} />
diff --git a/surfsense_web/components/public-chat/public-thread.tsx b/surfsense_web/components/public-chat/public-thread.tsx
index 9b31a1a02..3a224374e 100644
--- a/surfsense_web/components/public-chat/public-thread.tsx
+++ b/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 void }
 
 	if (avatarUrl && !hasError) {
 		return (
+			// biome-ignore lint/performance/noImgElement: external OAuth/profile avatar URL
 			{displayName 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 = () => {
 				 null,
+								multi_link_preview: () => null,
+								scrape_webpage: () => null,
+							},
+							Fallback: ToolFallback,
+						},
 					}}
 				/>
 			

@@ -160,12 +177,12 @@ const PublicAssistantActionBar: FC = () => {
 		>
 			
 				
-					 message.isCopied}>
+					 message.isCopied}>
 						
-					
-					 !message.isCopied}>
+					
+					 !message.isCopied}>
 						
-					
+					
 				
 			
 		
diff --git a/surfsense_web/components/tool-ui/article/index.tsx b/surfsense_web/components/tool-ui/article/index.tsx
deleted file mode 100644
index 43ea7c4c9..000000000
--- a/surfsense_web/components/tool-ui/article/index.tsx
+++ /dev/null
@@ -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;
-
-/**
- * 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;
-		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 ;
-	}
-
-	return (
-		 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 (
-			
-				
-					

-						

-							
-						

-						

-							
Failed to scrape webpage

-							{href && 
{href}
}
-							
{error}

-						

-					

-				
-			
-		);
-	}
-
-	return (
-		
-			 {
-					if (href && (e.key === "Enter" || e.key === " ")) {
-						e.preventDefault();
-						handleCardClick();
-					}
-				}}
-			>
-				{/* Header */}
-				
-					

-						{/* Favicon / Icon */}
-						{domain ? (
-							

-								
-							

-						) : (
-							

-								
-							

-						)}
-
-						{/* Content */}
-						

-							{/* Title */}
-							

-								{title}
-							

-
-							{/* Description */}
-							{description && (
-								

-									{description}
-								

-							)}
-
-							{/* Metadata row */}
-							

-								{domain && (
-									
-										
-											
-												
-												{domain}
-											
-										
-										
-											
Source: {domain}

-										
-									
-								)}
-
-								{author && (
-									
-										
-											
-												
-												{author}
-											
-										
-										
-											
Author: {author}

-										
-									
-								)}
-
-								{date && (
-									
-										
-										{date}
-									
-								)}
-
-								{wordCount && (
-									
-										
-											
-												
-												{formatWordCount(wordCount)}
-												{wasTruncated && (truncated)}
-											
-										
-										
-											

-												{wasTruncated
-													? "Content was truncated due to length"
-													: "Full article content available"}
-											

-										
-									
-								)}
-							

-						

-					

-
-					{/* Response actions */}
-					{responseActions && responseActions.length > 0 && (
-						

-							{responseActions.map((action) => (
-								
-							))}
-						

-					)}
-				
-			
-		
-	);
-}
-
-/**
- * Loading state for article component
- */
-export function ArticleLoading({ title = "Loading article..." }: { title?: string }) {
-	return (
-		
-			
-				

-					

-					

-						

-						

-						

-					

-				

-				
{title}

-			
-		
-	);
-}
-
-/**
- * Skeleton for article component
- */
-export function ArticleSkeleton() {
-	return (
-		
-			
-				

-					

-					

-						

-						

-						

-					

-				

-			
-		
-	);
-}
-
-/**
- * Error boundary props
- */
-interface ErrorBoundaryProps {
-	children: ReactNode;
-	fallback?: ReactNode;
-}
-
-interface ErrorBoundaryState {
-	hasError: boolean;
-}
-
-/**
- * Error boundary for article component
- */
-export class ArticleErrorBoundary extends Component {
-	constructor(props: ErrorBoundaryProps) {
-		super(props);
-		this.state = { hasError: false };
-	}
-
-	static getDerivedStateFromError(): ErrorBoundaryState {
-		return { hasError: true };
-	}
-
-	render() {
-		if (this.state.hasError) {
-			return (
-				this.props.fallback || (
-					
-						
-							

-								
-								
Failed to render article

-							

-						
-					
-				)
-			);
-		}
-
-		return this.props.children;
-	}
-}
diff --git a/surfsense_web/components/tool-ui/audio.tsx b/surfsense_web/components/tool-ui/audio.tsx
index f9634f34b..dae752034 100644
--- a/surfsense_web/components/tool-ui/audio.tsx
+++ b/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";
diff --git a/surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx b/surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx
index ea4434852..d45c6879e 100644
--- a/surfsense_web/components/tool-ui/confluence/create-confluence-page.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isErrorResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx b/surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx
index 211ee3388..538b84bfe 100644
--- a/surfsense_web/components/tool-ui/confluence/delete-confluence-page.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isAuthErrorResult(result)) return ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isWarningResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isWarningResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx b/surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx
index 286981d51..b42e26f69 100644
--- a/surfsense_web/components/tool-ui/confluence/update-confluence-page.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isAuthErrorResult(result)) return ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/deepagent-thinking.tsx b/surfsense_web/components/tool-ui/deepagent-thinking.tsx
deleted file mode 100644
index 3e6f668a8..000000000
--- a/surfsense_web/components/tool-ui/deepagent-thinking.tsx
+++ /dev/null
@@ -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 = {
-	[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;
-type DeepAgentThinkingArgs = z.infer;
-type DeepAgentThinkingResult = z.infer;
-
-// ============================================================================
-// 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 ;
-	}
-
-	if (status === STEP_STATUS.COMPLETED) {
-		return ;
-	}
-
-	// Default icons based on step type keywords
-	if (titleContainsKeywords(title, STEP_KEYWORDS.SEARCH)) {
-		return ;
-	}
-
-	if (titleContainsKeywords(title, STEP_KEYWORDS.ANALYSIS)) {
-		return ;
-	}
-
-	return ;
-}
-
-// ============================================================================
-// Sub-Components
-// ============================================================================
-
-interface ThinkingStepDisplayProps {
-	step: ThinkingStep;
-	isOpen: boolean;
-	onToggle: () => void;
-}
-
-/**
- * Component to display a single thinking step with controlled open state
- */
-const ThinkingStepDisplay: FC = ({ 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 (
-		
-			
-				{step.title}
-			
-			
-				{step.items.map((item, index) => (
-					{item}
-				))}
-			
-		
-	);
-};
-
-interface ThinkingLoadingStateProps {
-	status?: ThinkingStatus | string;
-}
-
-/**
- * Loading state with animated thinking indicator
- */
-const ThinkingLoadingState: FC = ({ 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 (
-		

-			

-				
-				
-					
-					
-				
-			

-			{statusText}
-		

-	);
-};
-
-interface SmartChainOfThoughtProps {
-	steps: ThinkingStep[];
-}
-
-/** Type for tracking step override states */
-type StepOverrides = Record;
-
-/** Type for tracking step status history */
-type StepStatusHistory = Record;
-
-/**
- * Smart chain of thought renderer with state management
- */
-const SmartChainOfThought: FC = ({ steps }) => {
-	// Track which steps the user has manually toggled
-	const [manualOverrides, setManualOverrides] = useState({});
-	// Track previous step statuses to detect changes
-	const prevStatusesRef = useRef({});
-
-	// 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 (
-		
-			{steps.map((step) => {
-				const isOpen = getStepOpenState(step);
-				return (
-					 handleToggle(step.id, isOpen)}
-					/>
-				);
-			})}
-		
-	);
-};
-
-/**
- * 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 ;
-		}
-
-		// 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 (
-			

-				
-			

-		);
-	},
-});
-
-// ============================================================================
-// 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 = ({
-	steps,
-	isStreaming = false,
-	className,
-}) => {
-	if (steps.length === 0 && !isStreaming) {
-		return null;
-	}
-
-	return (
-		

-			{isStreaming && steps.length === 0 ? (
-				
-			) : (
-				
-			)}
-		

-	);
-};
-
-// ============================================================================
-// Exports
-// ============================================================================
-
-export type {
-	ThinkingStep,
-	DeepAgentThinkingArgs,
-	DeepAgentThinkingResult,
-	StepStatus,
-	ThinkingStatus,
-};
-
-export { STEP_STATUS, THINKING_STATUS };
diff --git a/surfsense_web/components/tool-ui/display-image.tsx b/surfsense_web/components/tool-ui/display-image.tsx
deleted file mode 100644
index b5fccbc78..000000000
--- a/surfsense_web/components/tool-ui/display-image.tsx
+++ /dev/null
@@ -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;
-type DisplayImageResult = z.infer;
-
-/**
- * Error state component shown when image display fails
- */
-function ImageErrorState({ src, error }: { src: string; error: string }) {
-	return (
-		

-			

-				

-					
-				

-				

-					
Failed to display image

-					
{src}

-					
{error}

-				

-			

-		

-	);
-}
-
-/**
- * Cancelled state component
- */
-function ImageCancelledState({ src }: { src: string }) {
-	return (
-		

-			

-				
-				Image: {src}
-			

-		

-	);
-}
-
-/**
- * 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 ;
-}
-
-/**
- * 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({
-	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 (
-				

-					
-				

-			);
-		}
-
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return ;
-			}
-			if (status.reason === "error") {
-				return (
-					
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return (
-				

-					
-				

-			);
-		}
-
-		// Error result from the tool
-		if (result.error) {
-			return ;
-		}
-
-		// Success - render the image
-		return (
-			

-				
-					
-				
-			

-		);
-	},
-});
-
-export {
-	DisplayImageArgsSchema,
-	DisplayImageResultSchema,
-	type DisplayImageArgs,
-	type DisplayImageResult,
-};
diff --git a/surfsense_web/components/tool-ui/generate-image.tsx b/surfsense_web/components/tool-ui/generate-image.tsx
new file mode 100644
index 000000000..f077dcdad
--- /dev/null
+++ b/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;
+type GenerateImageResult = z.infer;
+
+function ImageErrorState({ prompt, error }: { prompt: string; error: string }) {
+	return (
+		

+			

+				

+					
+				

+				

+					
Image generation failed

+					
{prompt}

+					
{error}

+				

+			

+		

+	);
+}
+
+function ImageCancelledState({ prompt }: { prompt: string }) {
+	return (
+		

+			

+				
+				Generate: {prompt}
+			

+		

+	);
+}
+
+function ParsedImage({ result }: { result: unknown }) {
+	const image = parseSerializableImage(result);
+	return (
+		
+	);
+}
+
+/**
+ * Tool UI for generate_image — renders the generated image directly
+ * from the tool result directly.
+ */
+export const GenerateImageToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps) => {
+	const prompt = args.prompt || "Generating image...";
+
+	if (status.type === "running" || status.type === "requires-action") {
+		return (
+			

+				
+			

+		);
+	}
+
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
+			return ;
+		}
+		if (status.reason === "error") {
+			return (
+				
+			);
+		}
+	}
+
+	if (!result) {
+		return (
+			

+				
+			

+		);
+	}
+
+	if (result.error) {
+		return ;
+	}
+
+	return (
+		

+			
+				
+			
+		

+	);
+};
+
+export {
+	GenerateImageArgsSchema,
+	GenerateImageResultSchema,
+	type GenerateImageArgs,
+	type GenerateImageResult,
+};
diff --git a/surfsense_web/components/tool-ui/generate-podcast.tsx b/surfsense_web/components/tool-ui/generate-podcast.tsx
index bac5b3d5c..02f53efad 100644
--- a/surfsense_web/components/tool-ui/generate-podcast.tsx
+++ b/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) => {
+	const title = args.podcast_title || "SurfSense Podcast";
 
-		// Loading state - tool is still running (agent processing)
-		if (status.type === "running" || status.type === "requires-action") {
-			return ;
-		}
+	// Loading state - tool is still running (agent processing)
+	if (status.type === "running" || status.type === "requires-action") {
+		return ;
+	}
 
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return (
-					

-						

-							
Podcast Cancelled

-							

-								Podcast generation was cancelled
-							

-						

-					

-				);
-			}
-			if (status.reason === "error") {
-				return (
-					
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return ;
-		}
-
-		// Failed result (new: "failed", legacy: "error")
-		if (result.status === "failed" || result.status === "error") {
-			return ;
-		}
-
-		// 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 (
 				

 					

-						
Podcast already in progress

-						

-							Please wait for the current podcast to complete.
-						

+						
Podcast Cancelled

+						
Podcast generation was cancelled

 					

 				

 			);
 		}
-
-		// Pending - poll for completion (new: "pending" with podcast_id)
-		if (result.status === "pending" && result.podcast_id) {
-			return ;
-		}
-
-		// Ready with podcast_id (new: "ready", legacy: "success")
-		if ((result.status === "ready" || result.status === "success") && result.podcast_id) {
-			return ;
-		}
-
-		// 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 (
-				

-					

-						
Podcast Unavailable

-						

-							This podcast was generated with an older version. Please generate a new one.
-						

-					

-				

+				
 			);
 		}
+	}
 
-		// Fallback - missing required data
-		return ;
-	},
-});
+	// No result yet
+	if (!result) {
+		return ;
+	}
+
+	// Failed result (new: "failed", legacy: "error")
+	if (result.status === "failed" || result.status === "error") {
+		return ;
+	}
+
+	// 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 (
+			

+				

+					
Podcast already in progress

+					

+						Please wait for the current podcast to complete.
+					

+				

+			

+		);
+	}
+
+	// Pending - poll for completion (new: "pending" with podcast_id)
+	if (result.status === "pending" && result.podcast_id) {
+		return ;
+	}
+
+	// Ready with podcast_id (new: "ready", legacy: "success")
+	if ((result.status === "ready" || result.status === "success") && result.podcast_id) {
+		return ;
+	}
+
+	// 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 (
+			

+				

+					
Podcast Unavailable

+					

+						This podcast was generated with an older version. Please generate a new one.
+					

+				

+			

+		);
+	}
+
+	// Fallback - missing required data
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/generate-report.tsx b/surfsense_web/components/tool-ui/generate-report.tsx
index dd81d3403..4f22bbb8f 100644
--- a/surfsense_web/components/tool-ui/generate-report.tsx
+++ b/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({
-	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) => {
+	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 ;
+	}
+
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
+			return ;
 		}
-
-		if (status.type === "running" || status.type === "requires-action") {
-			return ;
-		}
-
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return ;
-			}
-			if (status.reason === "error") {
-				return (
-					
-				);
-			}
-		}
-
-		if (!result) {
-			return ;
-		}
-
-		if (result.status === "failed") {
+		if (status.reason === "error") {
 			return (
 				
 			);
 		}
+	}
 
-		if (result.status === "ready" && result.report_id) {
-			return (
-				
-			);
-		}
+	if (!result) {
+		return ;
+	}
 
-		return ;
-	},
-});
+	if (result.status === "failed") {
+		return (
+			
+		);
+	}
+
+	if (result.status === "ready" && result.report_id) {
+		return (
+			
+		);
+	}
+
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/gmail/create-draft.tsx b/surfsense_web/components/tool-ui/gmail/create-draft.tsx
index bca1bba80..aa8f58c72 100644
--- a/surfsense_web/components/tool-ui/gmail/create-draft.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isErrorResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/gmail/send-email.tsx b/surfsense_web/components/tool-ui/gmail/send-email.tsx
index d3cf9d639..fda375e51 100644
--- a/surfsense_web/components/tool-ui/gmail/send-email.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isErrorResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/gmail/trash-email.tsx b/surfsense_web/components/tool-ui/gmail/trash-email.tsx
index d68c4b03f..f79d093f0 100644
--- a/surfsense_web/components/tool-ui/gmail/trash-email.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isNotFoundResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isNotFoundResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/gmail/update-draft.tsx b/surfsense_web/components/tool-ui/gmail/update-draft.tsx
index 8c9fac109..7789368b2 100644
--- a/surfsense_web/components/tool-ui/gmail/update-draft.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", {
-								detail: { decisions: [decision] },
-							})
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isNotFoundResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isNotFoundResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/google-calendar/create-event.tsx b/surfsense_web/components/tool-ui/google-calendar/create-event.tsx
index 8d337c8f7..a2e23dd36 100644
--- a/surfsense_web/components/tool-ui/google-calendar/create-event.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isErrorResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/google-calendar/delete-event.tsx b/surfsense_web/components/tool-ui/google-calendar/delete-event.tsx
index 8df47c3d8..404a6ced2 100644
--- a/surfsense_web/components/tool-ui/google-calendar/delete-event.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isWarningResult(result)) return ;
-		if (isAuthErrorResult(result)) return ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isWarningResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/google-calendar/update-event.tsx b/surfsense_web/components/tool-ui/google-calendar/update-event.tsx
index c5e261c11..cc941bab8 100644
--- a/surfsense_web/components/tool-ui/google-calendar/update-event.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isAuthErrorResult(result)) return ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/google-drive/create-file.tsx b/surfsense_web/components/tool-ui/google-drive/create-file.tsx
index c0b38db8e..efbd59453 100644
--- a/surfsense_web/components/tool-ui/google-drive/create-file.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
+	if (isAuthErrorResult(result)) return ;
 
-		if (isInsufficientPermissionsResult(result))
-			return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
 
-		if (isErrorResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/google-drive/trash-file.tsx b/surfsense_web/components/tool-ui/google-drive/trash-file.tsx
index 35559bb30..9d6a2fafb 100644
--- a/surfsense_web/components/tool-ui/google-drive/trash-file.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
+	if (isAuthErrorResult(result)) return ;
 
-		if (isInsufficientPermissionsResult(result))
-			return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
 
-		if (isNotFoundResult(result)) return ;
-		if (isWarningResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isWarningResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/image/index.tsx b/surfsense_web/components/tool-ui/image/index.tsx
index ec04b779f..81c55d10a 100644
--- a/surfsense_web/components/tool-ui/image/index.tsx
+++ b/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 (
-				
+				
 					

 						

 							
@@ -165,7 +165,10 @@ export class ImageErrorBoundary extends Component<
  */
 export function ImageSkeleton({ maxWidth = "512px" }: { maxWidth?: string }) {
 	return (
-		
+		
 			

 				
 			

@@ -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 (
-		
+		
 			

-				

-					
-					
{title}

-				

+				
 			

 		
 	);
@@ -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 (
-			
+			
 				

 					

 						
@@ -242,8 +258,7 @@ export function Image({
 		
 						{!imageLoaded && (
 							

-								
+								
 							

 						)}
-						{/* eslint-disable-next-line @next/next/no-img-element */}
-						{alt} setImageLoaded(true)}
 							onError={() => setImageError(true)}
 						/>
@@ -316,11 +335,9 @@ export function Image({
 						{description && (
 							
{description}

 						)}
-						{displayDomain && (
+						{displayDomain && !isGenerated && (
 							

-								{isGenerated ? (
-									
-								) : source?.iconUrl ? (
+								{source?.iconUrl ? (
 									({
-	toolName: "create_jira_issue",
-	render: function CreateJiraIssueUI({ args, result }) {
-		if (!result) return null;
+>) => {
+	if (!result) return null;
 
-		if (isInterruptResult(result)) {
-			return (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isAuthErrorResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/jira/delete-jira-issue.tsx b/surfsense_web/components/tool-ui/jira/delete-jira-issue.tsx
index 0ad5be5bd..01e682f33 100644
--- a/surfsense_web/components/tool-ui/jira/delete-jira-issue.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isAuthErrorResult(result)) return ;
-		if (isWarningResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isWarningResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/jira/update-jira-issue.tsx b/surfsense_web/components/tool-ui/jira/update-jira-issue.tsx
index f085a04d6..d5ff8fc87 100644
--- a/surfsense_web/components/tool-ui/jira/update-jira-issue.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isInsufficientPermissionsResult(result))
-			return ;
-		if (isAuthErrorResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isInsufficientPermissionsResult(result))
+		return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/linear/create-linear-issue.tsx b/surfsense_web/components/tool-ui/linear/create-linear-issue.tsx
index 39e689a46..b03eda86c 100644
--- a/surfsense_web/components/tool-ui/linear/create-linear-issue.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isErrorResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/linear/delete-linear-issue.tsx b/surfsense_web/components/tool-ui/linear/delete-linear-issue.tsx
index 592f01555..59a66999d 100644
--- a/surfsense_web/components/tool-ui/linear/delete-linear-issue.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isAuthErrorResult(result)) return ;
-		if (isWarningResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isWarningResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/linear/update-linear-issue.tsx b/surfsense_web/components/tool-ui/linear/update-linear-issue.tsx
index 0b0aa4623..dc5bcd0e8 100644
--- a/surfsense_web/components/tool-ui/linear/update-linear-issue.tsx
+++ b/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 (
-				 {
-						window.dispatchEvent(
-							new CustomEvent("hitl-decision", { detail: { decisions: [decision] } })
-						);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		if (isAuthErrorResult(result)) return ;
-		if (isErrorResult(result)) return ;
+	if (isNotFoundResult(result)) return ;
+	if (isAuthErrorResult(result)) return ;
+	if (isErrorResult(result)) return ;
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/link-preview.tsx b/surfsense_web/components/tool-ui/link-preview.tsx
deleted file mode 100644
index 5c1a952b2..000000000
--- a/surfsense_web/components/tool-ui/link-preview.tsx
+++ /dev/null
@@ -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;
-type LinkPreviewResult = z.infer;
-
-/**
- * Error state component shown when link preview fails
- */
-function LinkPreviewErrorState({ url, error }: { url: string; error: string }) {
-	return (
-		

-			

-				

-					
-				

-				

-					
Failed to load preview

-					
{url}

-					
{error}

-				

-			

-		

-	);
-}
-
-/**
- * Cancelled state component
- */
-function LinkPreviewCancelledState({ url }: { url: string }) {
-	return (
-		

-			

-				
-				Preview: {url}
-			

-		

-	);
-}
-
-/**
- * Parsed MediaCard component with error handling
- */
-function ParsedMediaCard({ result }: { result: unknown }) {
-	const card = parseSerializableMediaCard(result);
-
-	return (
-		 {
-				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({
-	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 (
-				

-					
-				

-			);
-		}
-
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return ;
-			}
-			if (status.reason === "error") {
-				return (
-					
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return (
-				

-					
-				

-			);
-		}
-
-		// Error result from the tool
-		if (result.error) {
-			return ;
-		}
-
-		// Success - render the media card
-		return (
-			

-				
-					
-				
-			

-		);
-	},
-});
-
-// ============================================================================
-// 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;
-type MultiLinkPreviewResult = z.infer;
-
-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 (
-				

-					{urls.slice(0, 4).map((url, index) => (
-						
-					))}
-				

-			);
-		}
-
-		// Incomplete state
-		if (status.type === "incomplete") {
-			return (
-				

-					

-						
-						Link previews cancelled
-					

-				

-			);
-		}
-
-		// No result
-		if (!result || !result.previews) {
-			return null;
-		}
-
-		// Render grid of previews
-		return (
-			

-				{result.previews.map((preview) => (
-					
-						
-					
-				))}
-				{result.errors?.map((err) => (
-					
-				))}
-			

-		);
-	},
-});
-
-export {
-	LinkPreviewArgsSchema,
-	LinkPreviewResultSchema,
-	MultiLinkPreviewArgsSchema,
-	MultiLinkPreviewResultSchema,
-	type LinkPreviewArgs,
-	type LinkPreviewResult,
-	type MultiLinkPreviewArgs,
-	type MultiLinkPreviewResult,
-};
diff --git a/surfsense_web/components/tool-ui/media-card/index.tsx b/surfsense_web/components/tool-ui/media-card/index.tsx
deleted file mode 100644
index c7c8cfdf2..000000000
--- a/surfsense_web/components/tool-ui/media-card/index.tsx
+++ /dev/null
@@ -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;
-type MediaCardKind = z.infer;
-type ResponseAction = z.infer;
-export type SerializableMediaCard = z.infer;
-
-/**
- * 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 ;
-		case "image":
-			return ;
-		case "video":
-		case "audio":
-			return ;
-		default:
-			return ;
-	}
-}
-
-/**
- * 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 (
-				
-					
-						

-							
-						

-						

-							
Failed to load preview

-							

-								{this.state.error?.message || "An error occurred"}
-							

-						

-					
-				
-			);
-		}
-
-		return this.props.children;
-	}
-}
-
-/**
- * Loading skeleton for MediaCard
- */
-export function MediaCardSkeleton({ maxWidth = "420px" }: { maxWidth?: string }) {
-	return (
-		
-			

-			
-				

-				

-				

-			
-		
-	);
-}
-
-/**
- * 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 (
-		
-			 {
-					if (href && (e.key === "Enter" || e.key === " ")) {
-						e.preventDefault();
-						handleCardClick();
-					}
-				}}
-			>
-				{/* Thumbnail */}
-				{thumb && (
-					

-						 {
-								// Hide broken images
-								e.currentTarget.style.display = "none";
-							}}
-						/>
-						{/* Gradient overlay */}
-						

-					

-				)}
-
-				{/* Fallback when no thumbnail */}
-				{!thumb && (
-					

-						

-							{getKindIcon(kind)}
-							{kind === "link" ? "Link Preview" : kind}
-						

-					

-				)}
-
-				{/* Content */}
-				
-					

-						{/* Domain favicon placeholder */}
-						

-							
-						

-
-						

-							{/* Domain badge */}
-							{displayDomain && (
-								

-									
-										{displayDomain}
-									
-									{href && (
-										
-									)}
-								

-							)}
-
-							{/* Title */}
-							

-								{title}
-							

-
-							{/* Description */}
-							{description && (
-								

-									{description}
-								

-							)}
-						

-					

-
-					{/* Response Actions */}
-					{responseActions && responseActions.length > 0 && (
-						

-							{responseActions.map((action) => (
-								
-									
-										
-									
-									{action.confirmLabel && (
-										
-											
{action.confirmLabel}

-										
-									)}
-								
-							))}
-						

-					)}
-				
-			
-		
-	);
-}
-
-/**
- * MediaCard Loading State
- */
-export function MediaCardLoading({ title = "Loading preview..." }: { title?: string }) {
-	return (
-		
-			

-				
-			

-			
-				

-					

-					

-						

-						

-					

-				

-				
{title}

-			
-		
-	);
-}
diff --git a/surfsense_web/components/tool-ui/notion/create-notion-page.tsx b/surfsense_web/components/tool-ui/notion/create-notion-page.tsx
index 3b89b3c4a..9394fd9e0 100644
--- a/surfsense_web/components/tool-ui/notion/create-notion-page.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					const event = new CustomEvent("hitl-decision", {
+						detail: { decisions: [decision] },
+					});
+					window.dispatchEvent(event);
+				}}
+			/>
+		);
+	}
 
-		if (isAuthErrorResult(result)) {
-			return ;
-		}
+	if (isAuthErrorResult(result)) {
+		return ;
+	}
 
-		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 ;
-		}
+	if (isErrorResult(result)) {
+		return ;
+	}
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/notion/delete-notion-page.tsx b/surfsense_web/components/tool-ui/notion/delete-notion-page.tsx
index c3f78209d..a2dd70752 100644
--- a/surfsense_web/components/tool-ui/notion/delete-notion-page.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		}
+	if (isInfoResult(result)) {
+		return ;
+	}
 
-		if (isWarningResult(result)) {
-			return ;
-		}
+	if (isWarningResult(result)) {
+		return ;
+	}
 
-		if (isAuthErrorResult(result)) {
-			return ;
-		}
+	if (isAuthErrorResult(result)) {
+		return ;
+	}
 
-		if (isErrorResult(result)) {
-			return ;
-		}
+	if (isErrorResult(result)) {
+		return ;
+	}
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/notion/update-notion-page.tsx b/surfsense_web/components/tool-ui/notion/update-notion-page.tsx
index b3bb05117..56a557e47 100644
--- a/surfsense_web/components/tool-ui/notion/update-notion-page.tsx
+++ b/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 (
-				 {
-						const event = new CustomEvent("hitl-decision", {
-							detail: { decisions: [decision] },
-						});
-						window.dispatchEvent(event);
-					}}
-				/>
-			);
-		}
+	if (isInterruptResult(result)) {
+		return (
+			 {
+					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 ;
-		}
+	if (isInfoResult(result)) {
+		return ;
+	}
 
-		if (isAuthErrorResult(result)) {
-			return ;
-		}
+	if (isAuthErrorResult(result)) {
+		return ;
+	}
 
-		if (isErrorResult(result)) {
-			return ;
-		}
+	if (isErrorResult(result)) {
+		return ;
+	}
 
-		return ;
-	},
-});
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/sandbox-execute.tsx b/surfsense_web/components/tool-ui/sandbox-execute.tsx
index d4e60de9e..f8483b042 100644
--- a/surfsense_web/components/tool-ui/sandbox-execute.tsx
+++ b/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({
-	toolName: "execute",
-	render: function SandboxExecuteUI({ args, result, status }) {
-		const command = args.command || "…";
+export const SandboxExecuteToolUI = ({
+	args,
+	result,
+	status,
+}: ToolCallMessagePartProps) => {
+	const command = args.command || "…";
 
-		if (status.type === "running" || status.type === "requires-action") {
-			return ;
+	if (status.type === "running" || status.type === "requires-action") {
+		return ;
+	}
+
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
+			return ;
 		}
-
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return ;
-			}
-			if (status.reason === "error") {
-				return (
-					
-				);
-			}
+		if (status.reason === "error") {
+			return (
+				
+			);
 		}
+	}
 
-		if (!result) {
-			return ;
-		}
+	if (!result) {
+		return ;
+	}
 
-		if (result.error && !result.result && !result.output) {
-			return ;
-		}
+	if (result.error && !result.result && !result.output) {
+		return ;
+	}
 
-		const parsed = parseExecuteResult(result);
-		const threadId = result.thread_id || null;
-		return ;
-	},
-});
+	const parsed = parseExecuteResult(result);
+	const threadId = result.thread_id || null;
+	return ;
+};
 
 export { ExecuteArgsSchema, ExecuteResultSchema, type ExecuteArgs, type ExecuteResult };
diff --git a/surfsense_web/components/tool-ui/scrape-webpage.tsx b/surfsense_web/components/tool-ui/scrape-webpage.tsx
deleted file mode 100644
index 87fae8868..000000000
--- a/surfsense_web/components/tool-ui/scrape-webpage.tsx
+++ /dev/null
@@ -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;
-type ScrapeWebpageResult = z.infer;
-
-/**
- * Error state component shown when webpage scraping fails
- */
-function ScrapeErrorState({ url, error }: { url: string; error: string }) {
-	return (
-		

-			

-				

-					
-				

-				

-					
Failed to scrape webpage

-					
{url}

-					
{error}

-				

-			

-		

-	);
-}
-
-/**
- * Cancelled state component
- */
-function ScrapeCancelledState({ url }: { url: string }) {
-	return (
-		

-			

-				
-				Scraping: {url}
-			

-		

-	);
-}
-
-/**
- * Parsed Article component with error handling
- */
-function ParsedArticle({ result }: { result: unknown }) {
-	const { description, ...article } = parseSerializableArticle(result);
-
-	return 
;
-}
-
-/**
- * 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({
-	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 (
-				

-					
-				

-			);
-		}
-
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return ;
-			}
-			if (status.reason === "error") {
-				return (
-					
-				);
-			}
-		}
-
-		// No result yet
-		if (!result) {
-			return (
-				

-					
-				

-			);
-		}
-
-		// Error result from the tool
-		if (result.error) {
-			return ;
-		}
-
-		// Success - render the article card
-		return (
-			

-				
-					
-				
-			

-		);
-	},
-});
-
-export {
-	ScrapeWebpageArgsSchema,
-	ScrapeWebpageResultSchema,
-	type ScrapeWebpageArgs,
-	type ScrapeWebpageResult,
-};
diff --git a/surfsense_web/components/tool-ui/user-memory.tsx b/surfsense_web/components/tool-ui/user-memory.tsx
index f7c002887..e232bdcc7 100644
--- a/surfsense_web/components/tool-ui/user-memory.tsx
+++ b/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({
-	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) => {
+	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 (
-				

-					

-						
-					

-					

-						Saving to memory...
-					

+	// Loading state
+	if (isRunning) {
+		return (
+			

+				

+					
 				

-			);
-		}
-
-		// Error state
-		if (isError) {
-			return (
-				

-					

-						
-					

-					

-						Failed to save memory
-						{result?.error && 
{result.error}
}
-					

+				

+					Saving to memory...
 				

-			);
-		}
+			

+		);
+	}
 
-		// Success state
-		if (isComplete && result?.status === "saved") {
-			return (
-				

-					

-						
-					

-					

-						

-							
-							Memory saved
-							
-						

-						
{content}

-					

+	// Error state
+	if (isError) {
+		return (
+			

+				

+					
 				

-			);
-		}
-
-		// Default/incomplete state - show what's being saved
-		if (content) {
-			return (
-				

-					

-						
-					

-					

-						

-							Saving memory
-							
-						

-						
{content}

-					

+				

+					Failed to save memory
+					{result?.error && 
{result.error}
}
 				

-			);
-		}
+			

+		);
+	}
 
-		return null;
-	},
-});
+	// Success state
+	if (isComplete && result?.status === "saved") {
+		return (
+			

+				

+					
+				

+				

+					

+						
+						Memory saved
+						
+					

+					
{content}

+				

+			

+		);
+	}
+
+	// Default/incomplete state - show what's being saved
+	if (content) {
+		return (
+			

+				

+					
+				

+				

+					

+						Saving memory
+						
+					

+					
{content}

+				

+			

+		);
+	}
+
+	return null;
+};
 
 // ============================================================================
 // Recall Memory Tool UI
 // ============================================================================
 
-export const RecallMemoryToolUI = makeAssistantToolUI({
-	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) => {
+	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 (
-				

-					

-						
-					

-					

-						
-							{query ? `Searching memories for "${query}"...` : "Recalling memories..."}
-						
-					

+	// Loading state
+	if (isRunning) {
+		return (
+			

+				

+					
 				

-			);
-		}
-
-		// Error state
-		if (isError) {
-			return (
-				

-					

-						
-					

-					

-						Failed to recall memories
-						{result?.error && 
{result.error}
}
-					

+				

+					
+						{query ? `Searching memories for "${query}"...` : "Recalling memories..."}
+					
 				

-			);
-		}
+			

+		);
+	}
 
-		// Success state with memories
-		if (isComplete && result?.status === "success") {
-			const memories = result.memories || [];
-			const count = result.count || 0;
-
-			if (count === 0) {
-				return (
-					

-						

-							
-						

-						No memories found
-					

-				);
-			}
-
-			return (
-				

-					

-						
-						
-							Recalled {count} {count === 1 ? "memory" : "memories"}
-						
-					

-					

-						{memories.slice(0, 5).map((memory: MemoryItem) => (
-							

-								
-								{memory.memory_text}
-							

-						))}
-						{memories.length > 5 && (
-							
...and {memories.length - 5} more

-						)}
-					

+	// Error state
+	if (isError) {
+		return (
+			

+				

+					
 				

-			);
-		}
+				

+					Failed to recall memories
+					{result?.error && 
{result.error}
}
+				

+			

+		);
+	}
 
-		// 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 (
 				

 					

 						
 					

-					Searching memories for "{query}"
+					No memories found
 				

 			);
 		}
 
-		return null;
-	},
-});
+		return (
+			

+				

+					
+					
+						Recalled {count} {count === 1 ? "memory" : "memories"}
+					
+				

+				

+					{memories.slice(0, 5).map((memory: MemoryItem) => (
+						

+							
+							{memory.memory_text}
+						

+					))}
+					{memories.length > 5 && (
+						
...and {memories.length - 5} more

+					)}
+				

+			

+		);
+	}
+
+	// Default/incomplete state
+	if (query) {
+		return (
+			

+				

+					
+				

+				Searching memories for "{query}"
+			

+		);
+	}
+
+	return null;
+};
 
 // ============================================================================
 // Exports
diff --git a/surfsense_web/components/tool-ui/video-presentation/combined-player.tsx b/surfsense_web/components/tool-ui/video-presentation/combined-player.tsx
index 9a87c48d2..c630008db 100644
--- a/surfsense_web/components/tool-ui/video-presentation/combined-player.tsx
+++ b/surfsense_web/components/tool-ui/video-presentation/combined-player.tsx
@@ -118,7 +118,7 @@ export function CombinedPlayer({ slides }: CombinedPlayerProps) {
 	);
 
 	return (
-		

+		

 			
-			

-				

-					

-						
-					

-					

-				

-				

-					

-						{title}
-					

-					

-						
-						
-							Generating video presentation. This may take a few minutes.
-						
-					

-					

-						

-							

-						

-					

-				

+		

+			

+				
{title}

+				
 			

 		

 	);
@@ -108,20 +91,14 @@ function GeneratingState({ title }: { title: string }) {
 
 function ErrorState({ title, error }: { title: string; error: string }) {
 	return (
-		

-			

-				

-					
-				

-				

-					

-						{title}
-					

-					

-						Failed to generate video presentation
-					

-					
{error}

-				

+		

+			

+				
Video Generation Failed

+			

+			

+			

+				
{title}

+				
{error}

 			

 		

 	);
@@ -129,20 +106,10 @@ function ErrorState({ title, error }: { title: string; error: string }) {
 
 function CompilationLoadingState({ title }: { title: string }) {
 	return (
-		

-			

-				

-					
-				

-				

-					

-						{title}
-					

-					

-						
-						Compiling scenes...
-					

-				

+		

+			

+				
{title}

+				
 			

 		

 	);
@@ -163,7 +130,6 @@ function VideoPresentationPlayer({
 
 	const [isRendering, setIsRendering] = useState(false);
 	const [renderProgress, setRenderProgress] = useState(null);
-	const [renderError, setRenderError] = useState(null);
 	const [renderFormat, setRenderFormat] = useState(null);
 	const abortControllerRef = useRef(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 (
-		

-			{/* Title bar with actions */}
-			

-				

-					

-						
-					

-					

-						
{title}

-						

-							{compiledSlides.length} slides · {totalDuration.toFixed(1)}s · {FPS}fps
-						

-					

-				

-
-				

-					{isRendering ? (
-						<>
-							

-								
-								
-									Rendering {renderFormat ?? ""}{" "}
-									{renderProgress !== null ? `${Math.round(renderProgress * 100)}%` : "..."}
-								
-								

-									

-								

-							

-							
-						
-					) : (
-						<>
-							
-							
-						
-					)}
-				

+		

+			{/* Header */}
+			

+				
{title}

+				

+					{compiledSlides.length} slides  {totalDuration.toFixed(1)}s{" "}
+					 {FPS}fps
+				

 			

 
-			{/* Render error */}
-			{renderError && (
-				

-					
-					

-						
Download Failed

-						
{renderError}

-					

-				

-			)}
+			

 
-			{/* Combined Remotion Player */}
-			
+			{/* Remotion Player */}
+			

+				
+			

+
+			

+
+			{/* Action buttons */}
+			

+				{isRendering ? (
+					<>
+						

+							
+							
+								Rendering {renderFormat ?? ""}{" "}
+								{renderProgress !== null ? `${Math.round(renderProgress * 100)}%` : "..."}
+							
+							

+								

+							

+						

+						
+					
+				) : (
+					<>
+						
+						
+					
+				)}
+			

 		

 	);
 }
@@ -595,92 +551,85 @@ function StatusPoller({
 	return ;
 }
 
-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) => {
+	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 ;
-		}
+	if (status.type === "running" || status.type === "requires-action") {
+		return ;
+	}
 
-		if (status.type === "incomplete") {
-			if (status.reason === "cancelled") {
-				return (
-					

-						

-							
-							Presentation generation cancelled
-						

-					

-				);
-			}
-			if (status.reason === "error") {
-				return (
-					
-				);
-			}
-		}
-
-		if (!result) {
-			return ;
-		}
-
-		if (result.status === "failed") {
-			return ;
-		}
-
-		if (result.status === "generating") {
+	if (status.type === "incomplete") {
+		if (status.reason === "cancelled") {
 			return (
-				

-					

-						

-							
-						

-						

-							

-								Presentation already in progress
-							

-							

-								Please wait for the current presentation to complete.
-							

-						

+				

+					

+						
Presentation Cancelled

+						

+							Presentation generation was cancelled
+						

 					

 				

 			);
 		}
-
-		if (result.status === "pending" && result.video_presentation_id) {
+		if (status.reason === "error") {
 			return (
-				
 			);
 		}
+	}
 
-		if (result.status === "ready" && result.video_presentation_id) {
-			return (
-				
-			);
-		}
+	if (!result) {
+		return ;
+	}
 
-		return ;
-	},
-});
+	if (result.status === "failed") {
+		return ;
+	}
+
+	if (result.status === "generating") {
+		return (
+			

+				

+					
Presentation already in progress

+					

+						Please wait for the current presentation to complete.
+					

+				

+			

+		);
+	}
+
+	if (result.status === "pending" && result.video_presentation_id) {
+		return (
+			
+		);
+	}
+
+	if (result.status === "ready" && result.video_presentation_id) {
+		return (
+			
+		);
+	}
+
+	return ;
+};
diff --git a/surfsense_web/components/tool-ui/write-todos.tsx b/surfsense_web/components/tool-ui/write-todos.tsx
index 9b959bd33..104cbcf44 100644
--- a/surfsense_web/components/tool-ui/write-todos.tsx
+++ b/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({
-	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) => {
+	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 (
-					

-						
-							
-						
-					

-				);
-			}
-			return ;
+	// 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 (
+				

+					
+						
+					
+				

+			);
 		}
+		return ;
+	}
 
-		// Incomplete/cancelled state
-		if (status.type === "incomplete") {
-			if (currentPlanState || hasTodos) {
-				const plan = currentPlanState || parseSerializablePlan({ todos: data?.todos || [] });
-				return (
-					

-						
-							
-						
-					

-				);
-			}
-			return null;
+	// Incomplete/cancelled state
+	if (status.type === "incomplete") {
+		if (currentPlanState || hasTodos) {
+			const plan = currentPlanState || parseSerializablePlan({ todos: data?.todos || [] });
+			return (
+				

+					
+						
+					
+				

+			);
 		}
+		return null;
+	}
 
-		// Success - render the plan
-		const planToRender =
-			currentPlanState || (hasTodos ? parseSerializablePlan({ todos: data.todos }) : null);
-		if (!planToRender) {
-			return ;
-		}
+	// Success - render the plan
+	const planToRender =
+		currentPlanState || (hasTodos ? parseSerializablePlan({ todos: data.todos }) : null);
+	if (!planToRender) {
+		return ;
+	}
 
-		return (
-			

-				
-					
-				
-			

-		);
-	},
-});
+	return (
+		

+			
+				
+			
+		

+	);
+};
 
 export { WriteTodosSchema, type WriteTodosData };
diff --git a/surfsense_web/contracts/enums/toolIcons.tsx b/surfsense_web/contracts/enums/toolIcons.tsx
index 4d92713f8..90ec7a544 100644
--- a/surfsense_web/contracts/enums/toolIcons.tsx
+++ b/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 = {
 	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,
diff --git a/surfsense_web/lib/chat/message-utils.ts b/surfsense_web/lib/chat/message-utils.ts
index 81538731b..7c0da03c4 100644
--- a/surfsense_web/lib/chat/message-utils.ts
+++ b/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: {
diff --git a/surfsense_web/lib/chat/streaming-state.ts b/surfsense_web/lib/chat/streaming-state.ts
index 4364fd515..3f1c498b6 100644
--- a/surfsense_web/lib/chat/streaming-state.ts
+++ b/surfsense_web/lib/chat/streaming-state.ts
@@ -15,6 +15,10 @@ export type ContentPart =
 			toolName: string;
 			args: Record;
 			result?: unknown;
+	  }
+	| {
+			type: "data-thinking-steps";
+			data: { steps: ThinkingStepData[] };
 	  };
 
 export interface ContentPartsState {
@@ -23,6 +27,32 @@ export interface ContentPartsState {
 	toolCallIndices: Map;
 }
 
+export function updateThinkingSteps(
+	state: ContentPartsState,
+	steps: Map
+): 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,
-	currentThinkingSteps: Map
+	toolsWithUI: Set
 ): 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);
 		}
 	}
 
diff --git a/surfsense_web/package.json b/surfsense_web/package.json
index 4868ac864..7b7ef1e56 100644
--- a/surfsense_web/package.json
+++ b/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",
diff --git a/surfsense_web/pnpm-lock.yaml b/surfsense_web/pnpm-lock.yaml
index 0254eb9eb..03a9f9241 100644
--- a/surfsense_web/pnpm-lock.yaml
+++ b/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
{processChildrenWithCitations(children)} - {processChildrenWithCitations(children)} -