[Docs][Site][UI] - Search is now MCP based with UI updates. (#909)

<!--

Make sure you've read the CONTRIBUTING.md guidelines:
https://github.com/stack-auth/stack-auth/blob/dev/CONTRIBUTING.md

-->

Enhances the search functionality to now use the MCP server. 

Now handles API endpoints and webhooks. 

Now looks at what platform the user has selected, and searches based on
that. User can choose to filter differently if needed.

<img width="686" height="608" alt="image"
src="https://github.com/user-attachments/assets/641c9bd2-60d2-44b3-86ca-a4506257b430"
/>


<!-- RECURSEML_SUMMARY:START -->
## High-level PR Summary
This PR enhances the search functionality for Stack Auth documentation
by integrating it with their MCP (Model Control Plane) server. The
implementation replaces the previous client-side search algorithm with a
server-side approach that leverages the MCP's search capabilities. The
changes affect three main files: the MCP handler that now includes a new
`search_docs` tool, the search route handler that now forwards queries
to the MCP server instead of performing local searches, and the search
dialog UI that has been updated to better handle platform-specific
filtering and API documentation. The new implementation provides more
relevant search results and automatically filters based on the user's
current platform context, while allowing users to customize their search
filters as needed.

⏱️ Estimated Review Time: 30-90 minutes

<details>
<summary>💡 Review Order Suggestion</summary>

| Order | File Path |
|-------|-----------|
| 1 | `docs/src/app/api/internal/[transport]/route.ts` |
| 2 | `docs/src/app/api/search/route.ts` |
| 3 | `docs/src/components/layout/custom-search-dialog.tsx` |
</details>



<!-- RECURSEML_SUMMARY:END -->
<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Enhances search by integrating with MCP server, updating UI for
platform-specific filtering, and improving result presentation.
> 
>   - **Search Functionality**:
> - Integrates search with MCP server in `route.ts` files, replacing
client-side search.
> - Supports API and platform-specific results, filtering out admin API
endpoints.
>     - Sorts results by platform priority.
>   - **UI Updates**:
> - Updates `custom-search-dialog.tsx` to support platform-specific
filtering and API results.
>     - Adds platform badges and icons for different result types.
> - Auto-detects platform from URL and adjusts search filters
accordingly.
>   - **Miscellaneous**:
> - Adds new helper functions for platform extraction and MCP server
communication.
>     - Improves error handling and logging for search operations.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=stack-auth%2Fstack-auth&utm_source=github&utm_medium=referral)<sup>
for 9941d02bd5. You can
[customize](https://app.ellipsis.dev/stack-auth/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>

----


<!-- ELLIPSIS_HIDDEN -->

<!-- RECURSEML_ANALYSIS:START -->
## Review by RecurseML

_🔍 Review performed on
[7a0bf86..28264f7](7a0bf86cb3...28264f7f99)_

| &nbsp; Severity &nbsp; | &nbsp; Location &nbsp; | &nbsp; Issue &nbsp;
| &nbsp; Delete &nbsp; |
|:----------:|----------|-------|:--------:|
| ![Medium](https://img.shields.io/badge/Medium-yellow?style=plastic) |
[docs/src/app/api/internal/[transport]/route.ts:125](https://github.com/stack-auth/stack-auth/pull/909#discussion_r2368682927)
| API parameter 'query' uses camelCase instead of required snake_case |
[![](6854740bb0/?repo_owner=stack-auth&repo_name=stack-auth&pr_number=909)
|
| ![Medium](https://img.shields.io/badge/Medium-yellow?style=plastic) |
[docs/src/app/api/internal/[transport]/route.ts:126](https://github.com/stack-auth/stack-auth/pull/909#discussion_r2368683094)
| API parameter 'limit' uses camelCase instead of required snake_case |
[![](f4d772010d/?repo_owner=stack-auth&repo_name=stack-auth&pr_number=909)
|
| ![Medium](https://img.shields.io/badge/Medium-yellow?style=plastic) |
[docs/src/app/api/internal/[transport]/route.ts:174](https://github.com/stack-auth/stack-auth/pull/909#discussion_r2368683229)
| Async operation not wrapped with runAsynchronously |
[![](a573302c4b/?repo_owner=stack-auth&repo_name=stack-auth&pr_number=909)
|
| ![Medium](https://img.shields.io/badge/Medium-yellow?style=plastic) |
[docs/src/app/api/search/route.ts:23](https://github.com/stack-auth/stack-auth/pull/909#discussion_r2368683372)
| REST API parameters not using snake_case |
[![](d61261468a/?repo_owner=stack-auth&repo_name=stack-auth&pr_number=909)
|
| ![Medium](https://img.shields.io/badge/Medium-yellow?style=plastic) |
[docs/src/app/api/search/route.ts:130](https://github.com/stack-auth/stack-auth/pull/909#discussion_r2368683491)
| Async function call not wrapped in runAsynchronously |
[![](d315cd6c22/?repo_owner=stack-auth&repo_name=stack-auth&pr_number=909)
|
| ![Medium](https://img.shields.io/badge/Medium-yellow?style=plastic) |
[docs/src/app/api/search/route.ts:11](https://github.com/stack-auth/stack-auth/pull/909#discussion_r2368683588)
| Async function definition without proper runAsynchronously usage |
[![](37f4be1076/?repo_owner=stack-auth&repo_name=stack-auth&pr_number=909)
|

<details>
<summary>✅ Files analyzed, no issues (1)</summary>

  • `docs/src/components/layout/custom-search-dialog.tsx`
</details>

[![Need help? Join our
Discord](https://img.shields.io/badge/Need%20help%3F%20Join%20our%20Discord-5865F2?style=plastic&logo=discord&logoColor=white)](https://discord.gg/n3SsVDAW6U)
<!-- RECURSEML_ANALYSIS:END -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Search now includes API docs as an "API" result type with relevance
snippets, API-specific icons and group titles.
* Platform-aware filtering auto-detects platform on open and adds “API
only” and “Platform + API” views.

* **Improvements**
* Search is powered by a centralized streamed service for more
consistent results, better error handling, and fallbacks.
* Results are ordered with platform-priority, show clearer
counts/footers, and present more resilient, unified result formatting.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Konsti Wohlwend <n2d4xc@gmail.com>

docs/src/app/api/internal/[transport]/route.ts

@@ -134,6 +134,131 @@ const handler = createMcpHandler(
         };
       }
     );
+    server.tool(
+      "search_docs",
+      "Search through all Stack Auth documentation including API docs, guides, and examples. Returns ranked results with snippets and relevance scores.",
+      {
+        search_query: z.string().describe("The search query to find relevant documentation"),
+        result_limit: z.number().optional().describe("Maximum number of results to return (default: 50)")
+      },
+      async ({ search_query, result_limit = 50 }) => {
+        nodeClient?.capture({
+          event: "search_docs",
+          properties: { search_query, result_limit },
+          distinctId: "mcp-handler",
+        });
+
+        const results = [];
+        const queryLower = search_query.toLowerCase().trim();
+
+        // Search through all pages
+        for (const page of allPages) {
+          // Skip admin API endpoints
+          if (page.url.startsWith('/api/admin/')) {
+            continue;
+          }
+
+          let score = 0;
+          const title = page.data.title || '';
+          const description = page.data.description || '';
+
+          // Title matching (highest priority)
+          if (title.toLowerCase().includes(queryLower)) {
+            if (title.toLowerCase() === queryLower) {
+              score += 100; // Exact match
+            } else if (title.toLowerCase().startsWith(queryLower)) {
+              score += 80; // Starts with
+            } else {
+              score += 60; // Contains
+            }
+          }
+
+          // Description matching
+          if (description.toLowerCase().includes(queryLower)) {
+            score += 40;
+          }
+          // TOC/heading matching
+          for (const tocItem of page.data.toc) {
+            if (typeof tocItem.title === 'string' && tocItem.title.toLowerCase().includes(queryLower)) {
+              score += 30;
+            }
+          }
+
+          // Content matching (try to read the actual file)
+          try {
+            const filePath = `content/${page.file.path}`;
+            const content = await readFile(filePath, "utf-8");
+            const textContent = content
+              .replace(/^---[\s\S]*?---/, '') // Remove frontmatter
+              .replace(/<[^>]*>/g, ' ') // Remove JSX tags
+              .replace(/\{[^}]*\}/g, ' ') // Remove JSX expressions
+              .replace(/```[a-zA-Z]*\n/g, ' ') // Remove code block markers
+              .replace(/```/g, ' ')
+              .replace(/`([^`]*)`/g, '$1') // Remove inline code backticks
+              .replace(/\[([^\]]*)\]\([^)]*\)/g, '$1') // Extract link text
+              .replace(/[#*_~]/g, '') // Remove markdown formatting
+              .replace(/\s+/g, ' ')
+              .trim();
+
+            if (textContent.toLowerCase().includes(queryLower)) {
+              score += 20;
+
+              // Find snippet around the match
+              const matchIndex = textContent.toLowerCase().indexOf(queryLower);
+              const start = Math.max(0, matchIndex - 50);
+              const end = Math.min(textContent.length, matchIndex + 100);
+              const snippet = textContent.slice(start, end);
+
+              results.push({
+                title,
+                description,
+                url: page.url,
+                score,
+                snippet: `...${snippet}...`,
+                type: page.url.startsWith('/api/') ? 'api' : 'docs'
+              });
+            } else if (score > 0) {
+              // Add without snippet if title/description matched
+              results.push({
+                title,
+                description,
+                url: page.url,
+                score,
+                snippet: description || title,
+                type: page.url.startsWith('/api/') ? 'api' : 'docs'
+              });
+            }
+          } catch {
+            // If file reading fails but we have title/description matches
+            if (score > 0) {
+              results.push({
+                title,
+                description,
+                url: page.url,
+                score,
+                snippet: description || title,
+                type: page.url.startsWith('/api/') ? 'api' : 'docs'
+              });
+            }
+          }
+        }
+
+        // Sort by score (highest first) and limit results
+        const sortedResults = results
+          .sort((a, b) => b.score - a.score)
+          .slice(0, result_limit);
+
+        const searchResultText = sortedResults.length > 0
+          ? sortedResults.map(result =>
+              `Title: ${result.title}\nDescription: ${result.description}\nURL: ${result.url}\nType: ${result.type}\nScore: ${result.score}\nSnippet: ${result.snippet}\n`
+            ).join('\n---\n')
+          : `No results found for "${search_query}"`;
+
+        return {
+          content: [{ type: "text", text: searchResultText }],
+        };
+      }
+    );
     server.tool(
       "get_docs_by_id",
       "Use this tool to retrieve a specific Stack Auth Documentation page by its ID. It gives you the full content of the page so you can know exactly how to use specific Stack Auth APIs. Whenever using Stack Auth, you should always check the documentation first to have the most up-to-date information. When you write code using Stack Auth documentation you should reference the content you used in your comments.",

docs/src/app/api/search/route.ts

@@ -1,229 +1,158 @@
-import fs from 'fs';
-import { source } from 'lib/source';
 import type { NextRequest } from 'next/server';
-import path from 'path';
 
 type SearchResult = {
   id: string,
-  type: 'page' | 'heading' | 'text',
+  type: 'page' | 'heading' | 'text' | 'api',
   content: string,
   url: string,
-  score: number, // Add scoring for prioritization
+  title?: string,
 };
 
-// Helper function to calculate search relevance score
-function calculateScore(query: string, text: string, type: 'title' | 'description' | 'heading' | 'content'): number {
-  const queryLower = query.toLowerCase().trim();
-  const textLower = text.toLowerCase().trim();
-  // Base scores by type (higher = more important)
-  const baseScores = {
-    title: 100,
-    description: 70,
-    heading: 50,
-    content: 20
-  };
+// Helper function to call MCP server
+async function callMcpServer(search_query: string): Promise<SearchResult[]> {
+  try {
+    // Use localhost during development, production URL otherwise
+    const mcpUrl = process.env.NODE_ENV === 'development'
+      ? 'http://localhost:8104/api/internal/mcp'
+      : 'https://mcp.stack-auth.com/api/internal/mcp';
 
-  let score = 0;
-  let matchType = '';
+    const response = await fetch(mcpUrl, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'Accept': 'application/json, text/event-stream',
+      },
+      body: JSON.stringify({
+        jsonrpc: '2.0',
+        method: 'tools/call',
+        params: {
+          name: 'search_docs',
+          arguments: { search_query, result_limit: 20 },
+        },
+        id: Date.now(),
+      }),
+    });
 
-  // Exact match bonus (highest priority)
-  if (textLower === queryLower) {
-    score += baseScores[type] * 3; // Triple score for exact matches
-    matchType = 'exact';
-  }
-  // Starts with query bonus
-  else if (textLower.startsWith(queryLower)) {
-    score += baseScores[type] * 2; // Double score for starts with
-    matchType = 'starts-with';
-  }
-  // Contains as whole word bonus
-  else if (new RegExp(`\\b${queryLower.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}\\b`).test(textLower)) {
-    score += baseScores[type] * 1.5;
-    matchType = 'whole-word';
-  }
-  // Contains query bonus
-  else if (textLower.includes(queryLower)) {
-    score += baseScores[type];
-    matchType = 'contains';
-  }
-  else {
-    return 0; // No match
-  }
+    if (!response.ok) {
+      throw new Error(`MCP server error: ${response.status}`);
+    }
 
-  // Length penalty - shorter text with match is more relevant
-  const lengthPenalty = Math.min(text.length / 100, 0.3);
-  score -= lengthPenalty * 5;
+    // Parse Server-Sent Events format response
+    const text = await response.text();
+    const lines = text.split('\n');
+    let jsonData = null;
 
-  // Multiple occurrence bonus
-  const occurrences = (textLower.match(new RegExp(queryLower.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'g')) || []).length;
-  if (occurrences > 1) {
-    score += (occurrences - 1) * 15;
+    for (const line of lines) {
+      if (line.startsWith('data: ')) {
+        try {
+          jsonData = JSON.parse(line.substring(6));
+          break;
+        } catch (e) {
+          // Continue looking for valid JSON data
+        }
+      }
+    }
+
+    if (!jsonData) {
+      throw new Error('Invalid MCP response format');
+    }
+
+    if (jsonData.error) {
+      throw new Error(jsonData.error.message || 'MCP search failed');
+    }
+
+    // Parse the search results from the text response
+    const searchResultText = jsonData.result?.content?.[0]?.text || '';
+    if (searchResultText.includes('No results found')) {
+      return [];
+    }
+
+    const results: SearchResult[] = [];
+    const resultBlocks = searchResultText.split('\n---\n');
+
+    for (const block of resultBlocks) {
+      const lines = block.trim().split('\n');
+      let title = '';
+      let description = '';
+      let url = '';
+      let type = '';
+      let snippet = '';
+
+      for (const line of lines) {
+        if (line.startsWith('Title: ')) {
+          title = line.substring(7);
+        } else if (line.startsWith('Description: ')) {
+          description = line.substring(13);
+        } else if (line.startsWith('URL: ')) {
+          url = line.substring(5);
+        } else if (line.startsWith('Type: ')) {
+          type = line.substring(6);
+        } else if (line.startsWith('Snippet: ')) {
+          snippet = line.substring(9);
+        }
+      }
+
+      if (title && url) {
+        results.push({
+          id: `${url}-${type}`,
+          type: type === 'api' ? 'api' : 'page',
+          content: snippet || description || title,
+          url: url,
+          title: title,
+        });
+      }
+    }
+
+    return results;
+  } catch (error) {
+    console.error('MCP server call failed:', error);
+    // Fallback to empty results
+    return [];
   }
-
-  // Log scoring details for debugging
-  if (score > 0) {
-    console.log(`Score calculation: "${text}" (${type}) = ${score.toFixed(1)} [${matchType}, ${occurrences} occurrences]`);
-  }
-
-  return Math.max(score, 0);
 }
 
-// Helper function to extract text content from MDX
-function extractTextFromMDX(filePath: string): string {
-  try {
-    const content = fs.readFileSync(filePath, 'utf-8');
-    // Remove frontmatter
-    const withoutFrontmatter = content.replace(/^---[\s\S]*?---/, '');
-    // Remove JSX components and keep only text content
-    const textOnly = withoutFrontmatter
-      .replace(/<[^>]*>/g, ' ') // Remove JSX tags
-      .replace(/\{[^}]*\}/g, ' ') // Remove JSX expressions
-      .replace(/```[a-zA-Z]*\n/g, ' ') // Remove code block language markers
-      .replace(/```/g, ' ') // Remove code block delimiters but keep content
-      .replace(/`([^`]*)`/g, '$1') // Remove inline code backticks but keep content
-      .replace(/\[([^\]]*)\]\([^)]*\)/g, '$1') // Extract link text
-      .replace(/[#*_~]/g, '') // Remove markdown formatting (but keep backticks for now)
-      .replace(/\s+/g, ' ') // Normalize whitespace
-      .trim();
-    return textOnly;
-  } catch (error) {
-    console.error(`Error reading file ${filePath}:`, error);
-    return '';
-  }
+// Helper function to get platform priority for tie-breaking
+function getPlatformPriority(url: string): number {
+  // Higher number = higher priority
+  if (url.includes('/api/')) return 100; // API docs get highest priority
+  if (url.includes('/docs/next/')) return 90;
+  if (url.includes('/docs/react/')) return 80;
+  if (url.includes('/docs/js/')) return 70;
+  if (url.includes('/docs/python/')) return 60;
+  return 50; // Default priority
 }
 
 export async function GET(request: NextRequest) {
   const { searchParams } = new URL(request.url);
-  const query = searchParams.get('q');
+  const search_query = searchParams.get('q');
 
-  console.log('Search API called with query:', query);
+  console.log('Search API called with query:', search_query);
 
-  if (!query) {
+  if (!search_query) {
     return Response.json([]);
   }
 
   try {
-    // Get all pages from the source
-    const pages = source.getPages();
-    console.log(`Found ${pages.length} pages in source`);
+    // Call MCP server for search results
+    const results = await callMcpServer(search_query);
 
-    const results: SearchResult[] = [];
-    const queryLower = query.toLowerCase();
+    console.log(`Found ${results.length} search results from MCP server for "${search_query}"`);
 
-    // Search through all pages
-    pages.forEach((page, pageIndex) => {
-      const url = page.url;
-      const title = page.data.title || '';
-      const description = page.data.description || '';
+    // Filter out admin API endpoints as an additional safety measure
+    const filteredResults = results.filter(result => !result.url.startsWith('/api/admin'));
 
-      // Check if page title matches
-      const titleScore = calculateScore(query, title, 'title');
-      if (titleScore > 0) {
-        console.log(`Page title match: "${title}" at ${url} - Score: ${titleScore.toFixed(1)}`);
-        results.push({
-          id: `${url}-page`,
-          type: 'page',
-          content: title,
-          url: url,
-          score: titleScore
-        });
-      }
-
-      // Check if description matches
-      const descriptionScore = calculateScore(query, description, 'description');
-      if (descriptionScore > 0) {
-        results.push({
-          id: `${url}-description`,
-          type: 'text',
-          content: description,
-          url: url,
-          score: descriptionScore
-        });
-      }
-
-      // Search through TOC items (headings)
-      page.data.toc.forEach((tocItem, tocIndex) => {
-        const tocTitle = tocItem.title;
-        if (typeof tocTitle === 'string') {
-          const headingScore = calculateScore(query, tocTitle, 'heading');
-          if (headingScore > 0) {
-            results.push({
-              id: `${url}-${tocIndex}`,
-              type: 'heading',
-              content: tocTitle,
-              url: `${url}#${tocItem.url.slice(1)}`, // Remove the # from tocItem.url and add it back
-              score: headingScore
-            });
-          }
-        }
-      });
-
-      // Full content search by reading the actual MDX file
-      try {
-        // Construct file path from URL
-        const relativePath = url.replace('/docs/', './content/docs/') + '.mdx';
-        const fullPath = path.resolve(relativePath);
-
-        if (fs.existsSync(fullPath)) {
-          const textContent = extractTextFromMDX(fullPath);
-          const contentScore = calculateScore(query, textContent, 'content');
-
-          if (contentScore > 0) {
-            // Find a snippet around the match for better context
-            const matchIndex = textContent.toLowerCase().indexOf(queryLower);
-            const start = Math.max(0, matchIndex - 50);
-            const end = Math.min(textContent.length, matchIndex + 100);
-            const snippet = textContent.slice(start, end);
-
-            results.push({
-              id: `${url}-content-${pageIndex}`,
-              type: 'text',
-              content: `...${snippet}...`,
-              url: url,
-              score: contentScore
-            });
-          }
-        }
-      } catch (error) {
-        // Silently ignore file reading errors
-      }
+    // Sort by platform priority since MCP server already handles relevance
+    const sortedResults = filteredResults.sort((a, b) => {
+      return getPlatformPriority(b.url) - getPlatformPriority(a.url);
     });
 
-    // Sort results by score in descending order (highest score first)
-    results.sort((a, b) => b.score - a.score);
-
-    console.log(`\n=== RAW RESULTS FOR "${query}" ===`);
-    results.slice(0, 10).forEach((result, i) => {
-      console.log(`${i + 1}. "${result.content}" (${result.type}) - Score: ${result.score.toFixed(1)} - URL: ${result.url}`);
+    console.log(`\n=== MCP SEARCH RESULTS FOR "${search_query}" ===`);
+    sortedResults.slice(0, 10).forEach((result, i) => {
+      const priority = getPlatformPriority(result.url);
+      console.log(`${i + 1}. "${result.content}" (${result.type}) - Priority: ${priority} - URL: ${result.url}`);
     });
 
-    // Remove duplicate URLs and keep only the highest scoring result per URL
-    const seenUrls = new Set<string>();
-    const uniqueResults = results.filter(result => {
-      const baseUrl = result.url.split('#')[0]; // Remove fragment for deduplication
-      if (seenUrls.has(baseUrl)) {
-        console.log(`Duplicate URL filtered: ${result.content} (${result.score.toFixed(1)}) for ${baseUrl}`);
-        return false;
-      }
-      seenUrls.add(baseUrl);
-      return true;
-    });
-
-    // Re-sort after deduplication using the same logic
-    uniqueResults.sort((a, b) => b.score - a.score);
-
-    console.log(`\n=== FINAL RESULTS FOR "${query}" ===`);
-    uniqueResults.slice(0, 10).forEach((result, i) => {
-      console.log(`${i + 1}. "${result.content}" (${result.type}) - Score: ${result.score.toFixed(1)} - URL: ${result.url}`);
-    });
-
-    console.log(`\nFound ${uniqueResults.length} unique search results for "${query}"`);
-
-    // Remove score from response (internal use only)
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars -- score is used internally for sorting.
-    const responseResults = uniqueResults.map(({ score, ...result }) => result);
-
-    return Response.json(responseResults);
+    return Response.json(sortedResults);
 
   } catch (error) {
     console.error('Search error:', error);

docs/src/components/layout/custom-search-dialog.tsx

@@ -8,23 +8,106 @@ import { useSidebar } from '../layouts/sidebar-context';
 
 type SearchResult = {
   id: string,
-  type: 'page' | 'heading' | 'text',
+  type: 'page' | 'heading' | 'text' | 'api',
   content: string,
   url: string,
+  title?: string,
 };
 
+type DocumentCategory = 'api' | 'sdk' | 'component' | 'guide' | 'webhook';
+
 type GroupedResult = {
   basePath: string,
   title: string,
+  category: DocumentCategory,
+  categories: DocumentCategory[], // Support multiple categories (e.g., API + Webhook)
   results: SearchResult[],
 };
 
+function categorizeUrl(url: string): { primary: DocumentCategory, all: DocumentCategory[] } {
+  const categories: DocumentCategory[] = [];
+
+  // Check for API
+  if (url.startsWith('/api/')) {
+    categories.push('api');
+
+    // Check if it's also a webhook
+    if (url.includes('/webhook')) {
+      categories.push('webhook');
+      return { primary: 'webhook', all: categories };
+    }
+
+    return { primary: 'api', all: categories };
+  }
+
+  // Check for SDK
+  if (url.includes('/docs/sdk/') || url.includes('/sdk/')) {
+    categories.push('sdk');
+    return { primary: 'sdk', all: categories };
+  }
+
+  // Check for Component
+  if (url.includes('/docs/components/') || url.includes('/components/')) {
+    categories.push('component');
+    return { primary: 'component', all: categories };
+  }
+
+  // Default to guide
+  categories.push('guide');
+  return { primary: 'guide', all: categories };
+}
+
 function extractBasePathFromUrl(url: string): string {
-  // Extract everything after the platform but before any hash
+  // Handle API URLs
+  if (url.startsWith('/api/')) {
+    const match = url.match(/\/api\/([^#]+)/);
+    return match?.[1] || '';
+  }
+  // Handle docs URLs
   const match = url.match(/\/docs\/(.+?)(?:#|$)/);
   return match?.[1] || '';
 }
 
+function getCategoryLabel(category: DocumentCategory): string {
+  switch (category) {
+    case 'api': {
+      return 'API';
+    }
+    case 'sdk': {
+      return 'SDK';
+    }
+    case 'component': {
+      return 'COMP';
+    }
+    case 'guide': {
+      return 'GUIDE';
+    }
+    case 'webhook': {
+      return 'EVENT';
+    }
+  }
+}
+
+function getCategoryStyles(category: DocumentCategory): string {
+  switch (category) {
+    case 'api': {
+      return 'bg-red-500/10 text-red-700 dark:text-red-400 border border-red-500/20';
+    }
+    case 'sdk': {
+      return 'bg-blue-500/10 text-blue-700 dark:text-blue-400 border border-blue-500/20';
+    }
+    case 'component': {
+      return 'bg-cyan-500/10 text-cyan-700 dark:text-cyan-400 border border-cyan-500/20';
+    }
+    case 'guide': {
+      return 'bg-green-500/10 text-green-700 dark:text-green-400 border border-green-500/20';
+    }
+    case 'webhook': {
+      return 'bg-purple-500/10 text-purple-700 dark:text-purple-400 border border-purple-500/20';
+    }
+  }
+}
+
 function groupResultsByPage(results: SearchResult[]): GroupedResult[] {
   const grouped = new Map<string, GroupedResult>();
   const groupOrder: string[] = []; // Track the order groups are first encountered
@@ -32,15 +115,46 @@ function groupResultsByPage(results: SearchResult[]): GroupedResult[] {
   for (const result of results) {
     const basePath = extractBasePathFromUrl(result.url);
     const baseUrl = result.url.split('#')[0];
+    const { primary: category, all: categories } = categorizeUrl(result.url);
 
     if (!grouped.has(baseUrl)) {
-      // Find the page title from page-type results, fallback to path-based title
-      const pageResult = results.find(r => r.url === baseUrl && r.type === 'page');
-      const title = pageResult?.content || basePath.split('/').pop()?.replace(/-/g, ' ') || 'Unknown';
+      // Try to get title from the result itself first, then from other results with same base URL
+      let title = result.title;
+
+      if (!title) {
+        // Try to find a page-type result with this base URL that has a title
+        const pageResult = results.find(r => r.url.split('#')[0] === baseUrl && r.title);
+        title = pageResult?.title;
+      }
+
+      // Fallback to formatting the path
+      if (!title) {
+        // For API URLs, create readable titles
+        if (categories.includes('api')) {
+          const parts = basePath.split('/').filter(Boolean);
+          if (parts.length > 0) {
+            title = parts.map(part =>
+              part.split('-').map(word =>
+                word.charAt(0).toUpperCase() + word.slice(1)
+              ).join(' ')
+            ).join(' - ');
+          } else {
+            title = 'API Documentation';
+          }
+        } else {
+          // For docs URLs, format the last part of the path
+          const lastPart = basePath.split('/').pop() || basePath;
+          title = lastPart.split('-').map(word =>
+            word.charAt(0).toUpperCase() + word.slice(1)
+          ).join(' ');
+        }
+      }
 
       grouped.set(baseUrl, {
         basePath,
-        title,
+        title: title || 'Documentation',
+        category,
+        categories,
         results: []
       });
 
@@ -153,8 +267,18 @@ export function CustomSearchDialog({ open, onOpenChange }: CustomSearchDialogPro
 
   const groupedResults = groupResultsByPage(results);
 
-  // Use all results (no platform filtering)
-  const filteredResults = groupedResults;
+  // Sort results by category: guides first, then SDK, then API, then webhooks, then components
+  const categoryOrder: Record<DocumentCategory, number> = {
+    'guide': 1,
+    'sdk': 2,
+    'api': 3,
+    'webhook': 4,
+    'component': 5,
+  };
+
+  const filteredResults = groupedResults.sort((a, b) => {
+    return categoryOrder[a.category] - categoryOrder[b.category];
+  });
 
   // Flatten results for keyboard navigation
   const flatResults = filteredResults.flatMap(group =>
@@ -254,14 +378,26 @@ export function CustomSearchDialog({ open, onOpenChange }: CustomSearchDialogPro
           )}
 
           {!loading && filteredResults.map((group, groupIndex) => (
-            <div key={group.basePath || groupIndex} className="mb-6">
-              {/* Group Header */}
-              <div className="flex items-center gap-3 px-3 py-2 mb-3 bg-fd-muted/30 rounded-lg">
-                <h3 className="text-sm font-semibold text-fd-foreground">
+            <div key={group.basePath || groupIndex} className="mb-4">
+              {/* Group Header - grid layout for consistent badge alignment */}
+              <div className="grid grid-cols-[1fr_auto_auto] items-center gap-3 px-3 py-2.5 mb-2 bg-fd-muted/20 rounded-lg">
+                <h3 className="text-sm font-semibold text-fd-foreground truncate">
                   {group.title}
                 </h3>
-                <div className="flex-1" />
-                <span className="text-xs text-fd-muted-foreground">
+                <div className="flex items-center gap-1 flex-shrink-0">
+                  {group.categories.map((cat, idx) => (
+                    <span
+                      key={idx}
+                      className={cn(
+                        "inline-flex items-center justify-center px-2 py-0.5 rounded-md text-[10px] font-medium tracking-wide leading-none",
+                        getCategoryStyles(cat)
+                      )}
+                    >
+                      {getCategoryLabel(cat)}
+                    </span>
+                  ))}
+                </div>
+                <span className="text-xs text-fd-muted-foreground flex-shrink-0">
                   {group.results.length} result{group.results.length !== 1 ? 's' : ''}
                 </span>
               </div>
@@ -314,6 +450,11 @@ export function CustomSearchDialog({ open, onOpenChange }: CustomSearchDialogPro
                   );
                 })}
               </div>
+
+              {/* Separator between groups (except for last group) */}
+              {groupIndex < filteredResults.length - 1 && (
+                <div className="mt-4 mb-4 mx-3 border-t border-fd-border/30" />
+              )}
             </div>
           ))}