From 1a445089dd1e958fb9aa32937d59e370b09a3670 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 3 Apr 2026 16:14:22 -0400 Subject: [PATCH 1/4] refactor: simplify URL fetching around Firecrawl Move chat and autogen link retrieval onto a single Firecrawl v2 client, use batch scrape for multi-URL fetches, and remove the old hybrid scraping mode and URL hinting path. Made-with: Cursor --- .env.example | 8 - src/app/api/chat/route.ts | 69 +--- src/app/api/workspaces/autogen/route.ts | 5 +- .../assistant-ui/URLContextToolUI.tsx | 68 ++-- .../ai/__tests__/process-urls-shared.test.ts | 5 +- src/lib/ai/process-urls-shared.ts | 39 +- src/lib/ai/tools/process-urls.ts | 124 ++++--- src/lib/ai/utils/firecrawl.ts | 344 +++++++++++++++--- src/lib/ai/utils/url-processor.ts | 326 ----------------- 9 files changed, 441 insertions(+), 547 deletions(-) delete mode 100644 src/lib/ai/utils/url-processor.ts diff --git a/.env.example b/.env.example index 15e4a773..51dc4aee 100644 --- a/.env.example +++ b/.env.example @@ -38,14 +38,6 @@ GOOGLE_GENERATIVE_AI_API_KEY=AIza... # Get your API key from firecrawl.dev FIRECRAWL_API_KEY=fc_... -# Scraping Mode -# Options: -# - hybrid (Default: Try Google Context -> Fallback to Firecrawl if key exists -> Fallback to Direct) -# - firecrawl-only (Force Firecrawl for all scraping) -# - google-only (Force Google Context only) -# - direct-only (Force Direct Fetch only) -SCRAPING_MODE=hybrid - # FastAPI Service (optional - file conversion, doc-to-markdown, audio/video analysis) # FASTAPI_BASE_URL=https://your-fastapi-service.com # FASTAPI_API_KEY=your-service-api-key diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts index 947ef6f3..05bec851 100644 --- a/src/app/api/chat/route.ts +++ b/src/app/api/chat/route.ts @@ -19,10 +19,6 @@ import { getPostHogServerClient } from "@/lib/posthog-server"; import { withServerObservability } from "@/lib/with-server-observability"; import { normalizeLegacyToolMessages } from "@/lib/ai/legacy-tool-message-compat"; -// Regex patterns as constants (compiled once, reused for all requests) -const URL_CONTEXT_REGEX = /\[URL_CONTEXT:(.+?)\]/g; -const DIRECT_URL_REGEX = /https?:\/\/[^\s]+/g; - /** * Extract workspaceId from system context or request body */ @@ -41,50 +37,20 @@ function extractWorkspaceId(body: any): string | null { } /** - * Process messages in a single pass: extract URL context URLs and clean markers + * Process messages in a single pass: clean URL context markers * File attachments are handled natively as file parts via the SupabaseAttachmentAdapter. */ function processMessages(messages: any[]): { - urlContextUrls: string[]; cleanedMessages: any[]; } { - const urlContextUrlsSet = new Set(); - const cleanedMessages = messages.map((message) => { if (message.content && Array.isArray(message.content)) { const updatedContent = message.content.map((part: any) => { if (part.type === "text" && typeof part.text === "string") { const text = part.text; - // Extract URL context URLs (use Set for O(1) lookups) - const urlContextRegexLocal = new RegExp( - URL_CONTEXT_REGEX.source, - URL_CONTEXT_REGEX.flags, - ); - const urlContextMatches = text.matchAll(urlContextRegexLocal); - for (const urlMatch of urlContextMatches) { - const url = urlMatch[1]; - if (url) urlContextUrlsSet.add(url); - } - - // Extract direct URLs - const directUrlRegexLocal = new RegExp( - DIRECT_URL_REGEX.source, - DIRECT_URL_REGEX.flags, - ); - const directUrlMatches = text.matchAll(directUrlRegexLocal); - for (const directMatch of directUrlMatches) { - const url = directMatch[0]; - if (url) urlContextUrlsSet.add(url); - } - - // Clean URL_CONTEXT markers (create new instance to avoid global regex state issues) - const urlContextReplaceRegex = new RegExp( - URL_CONTEXT_REGEX.source, - URL_CONTEXT_REGEX.flags, - ); const updatedText = text.replace( - urlContextReplaceRegex, + /\[URL_CONTEXT:(.+?)\]/g, (_match: string, url: string) => { return url; }, @@ -100,7 +66,6 @@ function processMessages(messages: any[]): { }); return { - urlContextUrls: Array.from(urlContextUrlsSet), cleanedMessages, }; } @@ -166,26 +131,6 @@ function injectSelectionContext( } } -/** - * Build the enhanced system prompt with guidelines and detection hints - * Uses array join for better performance than string concatenation - */ -function buildSystemPrompt( - baseSystem: string, - urlContextUrls: string[], -): string { - const parts: string[] = [baseSystem]; - - // Add URL detection hint if URLs are present - if (urlContextUrls.length > 0) { - parts.push( - `\n\nURL DETECTION: The user's message contains ${urlContextUrls.length} URL(s): ${urlContextUrls.join(", ")}. Call the processUrls tool with a structured object like { "urls": ["https://example.com"] } to analyze them. Do not wrap the URLs in a JSON string.`, - ); - } - - return parts.join(""); -} - async function handlePOST(req: Request) { let workspaceId: string | null = null; let activeFolderId: string | undefined; @@ -261,9 +206,8 @@ async function handlePOST(req: Request) { emptyMessages: "remove", }); - // Process messages in single pass: extract URLs and clean markers - const { urlContextUrls, cleanedMessages } = - processMessages(convertedMessages); + // Process messages in single pass: clean URL context markers + const { cleanedMessages } = processMessages(convertedMessages); // Get pre-formatted selected cards context from client (no DB fetch needed) const selectedCardsContext = getSelectedCardsContext(body); @@ -282,9 +226,6 @@ async function handlePOST(req: Request) { modelId = `anthropic/${modelId}`; } - // Build system prompt (identity, guidelines, URL hints — no selected cards) - const finalSystemPrompt = buildSystemPrompt(system, urlContextUrls); - // Inject selected cards + reply selections into the last user message injectSelectionContext( cleanedMessages, @@ -355,7 +296,7 @@ async function handlePOST(req: Request) { const result = streamText({ model: model, temperature: 1.0, - system: finalSystemPrompt, + system, messages: cleanedMessages, stopWhen: stepCountIs(25), tools, diff --git a/src/app/api/workspaces/autogen/route.ts b/src/app/api/workspaces/autogen/route.ts index 49ce5b72..ed0bb205 100644 --- a/src/app/api/workspaces/autogen/route.ts +++ b/src/app/api/workspaces/autogen/route.ts @@ -10,7 +10,7 @@ import { db, workspaces } from "@/lib/db/client"; import { generateSlug } from "@/lib/workspace/slug"; import { workspaceWorker, type CreateItemParams } from "@/lib/ai/workers"; import { searchVideos } from "@/lib/youtube"; -import { UrlProcessor } from "@/lib/ai/utils/url-processor"; +import { FirecrawlClient } from "@/lib/ai/utils/firecrawl"; import { findNextAvailablePosition } from "@/lib/workspace-state/grid-layout-helpers"; import { generateItemId } from "@/lib/workspace-state/item-helpers"; import type { Item, QuizQuestion } from "@/lib/workspace-state/types"; @@ -173,7 +173,8 @@ async function runDistillationAgent( let linkContext = ""; if (nonYtLinks.length > 0) { send({ type: "toolCall", data: { toolName: "urlFetch", status: "fetching" } }); - const results = await UrlProcessor.processUrls(nonYtLinks); + const client = new FirecrawlClient(); + const results = await client.scrapeUrls(nonYtLinks); const successful = results.filter((r) => r.success && r.content); if (successful.length > 0) { linkContext = diff --git a/src/components/assistant-ui/URLContextToolUI.tsx b/src/components/assistant-ui/URLContextToolUI.tsx index e0f2ce60..6a73fccf 100644 --- a/src/components/assistant-ui/URLContextToolUI.tsx +++ b/src/components/assistant-ui/URLContextToolUI.tsx @@ -1,6 +1,6 @@ "use client"; -import { LinkIcon, ChevronDownIcon, CheckIcon, ExternalLinkIcon, AlertCircleIcon } from "lucide-react"; +import { LinkIcon, ChevronDownIcon, CheckIcon, ExternalLinkIcon } from "lucide-react"; import { useCallback, useRef, @@ -222,7 +222,6 @@ type ProcessUrlsResult = export const URLContextToolUI = makeAssistantToolUI<{ urls?: string[]; - instruction?: string; jsonInput?: string; }, ProcessUrlsResult>({ toolName: "processUrls", @@ -231,17 +230,21 @@ export const URLContextToolUI = makeAssistantToolUI<{ const isComplete = status.type === "complete"; const parsedResult = result != null ? parseURLContextResult(result) : null; - type Meta = { urlMetadata?: URLMetadata[]; groundingChunks?: unknown[]; sources?: SourceMetadata[] }; + type Meta = { + urlMetadata?: URLMetadata[]; + groundingChunks?: unknown[]; + sources?: SourceMetadata[]; + }; const metadata = (typeof parsedResult === "object" && parsedResult !== null && "metadata" in parsedResult ? (parsedResult as { metadata?: Meta }).metadata : null) as Meta | null; const urlMetadata = metadata?.urlMetadata ?? null; - const groundingChunks = metadata?.groundingChunks ?? null; - const sources = metadata?.sources ?? null; const normalizedArgs = normalizeProcessUrlsArgs(args); const urls = normalizedArgs?.urls ?? []; - const instruction = normalizedArgs?.instruction; const urlCount = urls.length; + const successfulCount = + urlMetadata?.filter((m) => m.urlRetrievalStatus === "URL_RETRIEVAL_STATUS_SUCCESS").length ?? 0; + const failedCount = (urlMetadata?.length ?? 0) - successfulCount; // Helper to get status badge color const getStatusColor = (status: string) => { @@ -265,7 +268,7 @@ export const URLContextToolUI = makeAssistantToolUI<{ } /> @@ -274,23 +277,12 @@ export const URLContextToolUI = makeAssistantToolUI<{
{urlCount > 0 && (
- URLs: - {instruction && ( -
-
- -
- Custom instruction: -

{instruction}

-
-
-
- )} + Links:
- {urls.map((url, index) => { + {urls.map((url) => { const urlMeta = urlMetadata?.find((m) => m.retrievedUrl === url); return ( -
+
- Analyzing {urlCount} URL{urlCount !== 1 ? "s" : ""}... + Processing {urlCount} link{urlCount !== 1 ? "s" : ""}...
)} @@ -333,34 +325,16 @@ export const URLContextToolUI = makeAssistantToolUI<{
- Successfully processed {urlCount} URL{urlCount !== 1 ? "s" : ""} + Processed {urlCount} link{urlCount !== 1 ? "s" : ""}
- {metadata && ( -
- {groundingChunks && groundingChunks.length > 0 && ( -
- Grounding chunks: - {groundingChunks.length} -
- )} - - {sources && sources.length > 0 && ( -
- Sources: - {sources.length} -
- )} - - {urlMetadata && urlMetadata.length > 0 && ( -
- Retrieved: - - {urlMetadata.filter((m) => m.urlRetrievalStatus === "URL_RETRIEVAL_STATUS_SUCCESS").length} / {urlMetadata.length} - -
- )} + {metadata && urlMetadata && urlMetadata.length > 0 && ( +
+ + {successfulCount} loaded + {failedCount > 0 ? `, ${failedCount} unavailable` : ""} +
)}
diff --git a/src/lib/ai/__tests__/process-urls-shared.test.ts b/src/lib/ai/__tests__/process-urls-shared.test.ts index f402290b..c3c91fcb 100644 --- a/src/lib/ai/__tests__/process-urls-shared.test.ts +++ b/src/lib/ai/__tests__/process-urls-shared.test.ts @@ -6,15 +6,13 @@ describe("normalizeProcessUrlsArgs", () => { expect( normalizeProcessUrlsArgs({ urls: ["https://example.com", "https://vercel.com"], - instruction: "Summarize the key points.", }), ).toEqual({ urls: ["https://example.com", "https://vercel.com"], - instruction: "Summarize the key points.", }); }); - it("normalizes legacy jsonInput payloads", () => { + it("normalizes legacy jsonInput payloads and drops instruction", () => { expect( normalizeProcessUrlsArgs({ jsonInput: JSON.stringify({ @@ -24,7 +22,6 @@ describe("normalizeProcessUrlsArgs", () => { }), ).toEqual({ urls: ["https://example.com"], - instruction: "Extract dates.", }); }); diff --git a/src/lib/ai/process-urls-shared.ts b/src/lib/ai/process-urls-shared.ts index 41b58dbf..1ce8c7b5 100644 --- a/src/lib/ai/process-urls-shared.ts +++ b/src/lib/ai/process-urls-shared.ts @@ -7,18 +7,41 @@ export const ProcessUrlsInputSchema = z.object({ .array(z.string().min(1)) .min(1) .max(MAX_PROCESS_URLS) - .describe("Web URLs to analyze. Pass them directly as an array of strings."), - instruction: z - .string() - .trim() - .min(1) - .max(1000) - .optional() - .describe("Optional extra instruction describing what to extract from the URLs."), + .describe("Web URLs to fetch. Pass them directly as an array of strings."), }); export type ProcessUrlsInput = z.infer; +export const ProcessUrlsOutputSchema = z.object({ + text: z.string(), + metadata: z + .object({ + provider: z.string().optional(), + urlMetadata: z + .array( + z.object({ + retrievedUrl: z.string(), + urlRetrievalStatus: z.string(), + }), + ) + .nullable() + .optional(), + groundingChunks: z.array(z.unknown()).nullable().optional(), + sources: z + .array( + z.object({ + uri: z.string(), + title: z.string(), + }), + ) + .nullable() + .optional(), + }) + .optional(), +}); + +export type ProcessUrlsOutput = z.infer; + export function normalizeProcessUrlsArgs( input: unknown, ): ProcessUrlsInput | null { diff --git a/src/lib/ai/tools/process-urls.ts b/src/lib/ai/tools/process-urls.ts index b3c7652b..610b8b2e 100644 --- a/src/lib/ai/tools/process-urls.ts +++ b/src/lib/ai/tools/process-urls.ts @@ -1,7 +1,21 @@ -import { google, type GoogleGenerativeAIProviderMetadata } from "@ai-sdk/google"; -import { generateText, tool, zodSchema } from "ai"; +import { tool, zodSchema } from "ai"; import { logger } from "@/lib/utils/logger"; -import { ProcessUrlsInputSchema } from "@/lib/ai/process-urls-shared"; +import { + ProcessUrlsInputSchema, + ProcessUrlsOutputSchema, + type ProcessUrlsOutput, +} from "@/lib/ai/process-urls-shared"; +import { FirecrawlClient } from "@/lib/ai/utils/firecrawl"; + +const MAX_CONTENT_CHARS_PER_URL = 12000; + +function truncateContent(content: string): string { + if (content.length <= MAX_CONTENT_CHARS_PER_URL) { + return content; + } + + return `${content.slice(0, MAX_CONTENT_CHARS_PER_URL).trim()}\n\n[Content truncated for length]`; +} /** * Create the processUrls tool for analyzing web pages @@ -13,10 +27,11 @@ import { ProcessUrlsInputSchema } from "@/lib/ai/process-urls-shared"; */ export function createProcessUrlsTool() { return tool({ - description: "Analyze web pages using Google's URL Context API. Extracts content, key information, and metadata from regular web URLs (http/https). Use this for web pages, articles, documentation, and other web content. This tool does not handle uploaded files or videos.", + description: "Fetch the content of web pages from the provided URLs. Use this for web pages, articles, and documentation when the model needs the actual page text before answering. This tool does not handle uploaded files or video URLs.", inputSchema: zodSchema(ProcessUrlsInputSchema), + outputSchema: zodSchema(ProcessUrlsOutputSchema), strict: true, - execute: async ({ urls: urlList, instruction }) => { + execute: async ({ urls: urlList }): Promise => { logger.debug("🔗 [URL_TOOL] Processing web URLs:", urlList); const fileUrls = urlList.filter((url: string) => @@ -26,57 +41,73 @@ export function createProcessUrlsTool() { if (fileUrls.length > 0) { logger.warn("🔗 [URL_TOOL] File/video URLs detected for web URL tool:", fileUrls); - return `Error: This tool only handles regular web URLs, not uploaded files or video URLs (${fileUrls.join(', ')})`; + return { + text: `Error: This tool only handles regular web URLs, not uploaded files or video URLs (${fileUrls.join(", ")})`, + metadata: { + provider: "firecrawl", + urlMetadata: urlList.map((url) => ({ + retrievedUrl: url, + urlRetrievalStatus: fileUrls.includes(url) + ? "URL_RETRIEVAL_STATUS_UNSUPPORTED" + : "URL_RETRIEVAL_STATUS_SKIPPED", + })), + groundingChunks: null, + sources: null, + }, + }; } try { - const urlsBlock = urlList.map((url, index) => `${index + 1}. ${url}`).join("\n"); - const instructionBlock = instruction?.trim() - ? `Focus instruction: ${instruction.trim()}\n\n` - : ""; + const client = new FirecrawlClient(); + const scrapedResults = (await client.scrapeUrls(urlList)).map((result) => ({ + ...result, + content: truncateContent(result.content), + success: result.success && result.content.length > 0, + })); - const { text, providerMetadata } = await generateText({ - model: google("gemini-3-flash-preview"), - tools: { - url_context: google.tools.urlContext({}), - }, - prompt: `${instructionBlock}Analyze these URLs directly using URL context: + const successfulResults = scrapedResults.filter((result) => result.success); + const failedResults = scrapedResults.filter((result) => !result.success); -${urlsBlock} + if (failedResults.length > 0) { + logger.warn( + `🔗 [URL_TOOL] ${failedResults.length} URL(s) failed to process:`, + failedResults.map((result) => result.url), + ); + } -Provide a clear, accurate answer in this format: -Summary: [1-2 sentences] -Key information: -- [Point 1] -- [Point 2] -- [Additional points as needed] -Details: [Important details, specs, dates, or relevant factual context]`, - }); - - const googleMetadata = providerMetadata?.google as GoogleGenerativeAIProviderMetadata | undefined; - const groundingMetadata = googleMetadata?.groundingMetadata ?? null; - const urlContextMetadata = googleMetadata?.urlContextMetadata ?? null; - const groundingChunks = groundingMetadata?.groundingChunks ?? null; - const sources = groundingChunks - ?.flatMap((chunk) => { - const uri = chunk.web?.uri; - if (!uri) return []; + if (successfulResults.length === 0) { + return { + text: `Failed to process any of the provided URLs. Errors: ${failedResults.map((result) => `${result.url}: ${result.error || "Unknown error"}`).join("; ")}`, + metadata: { + provider: "firecrawl", + urlMetadata: scrapedResults.map((result) => ({ + retrievedUrl: result.url, + urlRetrievalStatus: result.success + ? "URL_RETRIEVAL_STATUS_SUCCESS" + : "URL_RETRIEVAL_STATUS_FAILED", + })), + groundingChunks: null, + sources: null, + }, + }; + } - return [{ - uri, - title: chunk.web?.title || uri, - }]; - }) - .filter((source, index, array) => - array.findIndex((candidate) => candidate.uri === source.uri) === index - ) ?? null; + const combinedText = successfulResults + .map((result) => `# ${result.title}\nURL: ${result.url}\n\n${result.content}`) + .join("\n\n---\n\n"); return { - text, + text: combinedText, metadata: { - urlMetadata: urlContextMetadata?.urlMetadata ?? null, - groundingChunks, - sources, + provider: "firecrawl", + urlMetadata: scrapedResults.map((result) => ({ + retrievedUrl: result.url, + urlRetrievalStatus: result.success + ? "URL_RETRIEVAL_STATUS_SUCCESS" + : "URL_RETRIEVAL_STATUS_FAILED", + })), + groundingChunks: null, + sources: successfulResults.map(({ url, title }) => ({ uri: url, title })), }, }; @@ -88,6 +119,7 @@ Details: [Important details, specs, dates, or relevant factual context]`, return { text: `Error processing web URLs: ${error instanceof Error ? error.message : String(error)}`, metadata: { + provider: "firecrawl", urlMetadata: null, groundingChunks: null, sources: null, diff --git a/src/lib/ai/utils/firecrawl.ts b/src/lib/ai/utils/firecrawl.ts index de7a7903..238d5f1e 100644 --- a/src/lib/ai/utils/firecrawl.ts +++ b/src/lib/ai/utils/firecrawl.ts @@ -1,24 +1,43 @@ import { logger } from "@/lib/utils/logger"; -export interface FirecrawlScrapeResponse { +const DEFAULT_SCRAPE_OPTIONS = { + formats: ["markdown"], + onlyMainContent: true, + waitFor: 1000, +}; + +const BATCH_POLL_INTERVAL_MS = 1500; +const BATCH_POLL_TIMEOUT_MS = 45000; + +export interface FirecrawlMetadata { + title?: string | string[]; + description?: string | string[]; + language?: string | string[] | null; + sourceURL?: string; + url?: string; + statusCode?: number; + error?: string | null; + [key: string]: unknown; +} + +export interface FirecrawlDocument { + content?: string; + markdown?: string; + metadata?: FirecrawlMetadata; +} + +export interface FirecrawlPageResult { + url: string; + title: string; + content: string; success: boolean; - data?: { - content?: string; - markdown?: string; - metadata?: { - title?: string; - description?: string; - language?: string; - sourceURL?: string; - [key: string]: any; - }; - }; error?: string; + metadata?: FirecrawlMetadata; } export class FirecrawlClient { private apiKey: string; - private baseUrl = "https://api.firecrawl.dev/v1"; + private baseUrl = "https://api.firecrawl.dev/v2"; constructor(apiKey?: string) { this.apiKey = apiKey || process.env.FIRECRAWL_API_KEY || ""; @@ -27,12 +46,188 @@ export class FirecrawlClient { } } - /** - * Scrape a single URL using Firecrawl - */ - async scrapeUrl(url: string): Promise { + private getHeaders() { + return { + "Content-Type": "application/json", + "Authorization": `Bearer ${this.apiKey}`, + }; + } + + private async parseErrorResponse(response: Response): Promise { + if (response.status === 401) { + return "Invalid Firecrawl API key"; + } + + if (response.status === 429) { + return "Firecrawl rate limit exceeded"; + } + + try { + const errorText = await response.text(); + if (!errorText.trim()) { + return `Firecrawl API error: ${response.status}`; + } + + const errorJson = JSON.parse(errorText); + if (typeof errorJson?.error === "string") { + return errorJson.error; + } + + return `Firecrawl API error: ${response.status} ${errorText}`; + } catch { + // Fall through to generic status message if the body is unreadable or not JSON. + } + + return `Firecrawl API error: ${response.status}`; + } + + private firstString(value: unknown): string | undefined { + if (typeof value === "string") { + return value; + } + + if (Array.isArray(value)) { + const first = value.find((item) => typeof item === "string"); + return typeof first === "string" ? first : undefined; + } + + return undefined; + } + + private normalizePageResult(url: string, document: unknown): FirecrawlPageResult { + if (!document || typeof document !== "object" || Array.isArray(document)) { + return { + url, + title: url, + content: "", + success: false, + error: "Invalid Firecrawl response", + }; + } + + const record = document as { + markdown?: unknown; + content?: unknown; + metadata?: unknown; + }; + + const metadata = + record.metadata && typeof record.metadata === "object" && !Array.isArray(record.metadata) + ? (record.metadata as FirecrawlMetadata) + : undefined; + + const content = + typeof record.markdown === "string" + ? record.markdown + : typeof record.content === "string" + ? record.content + : ""; + + const metadataError = + typeof metadata?.error === "string" && metadata.error.trim().length > 0 + ? metadata.error + : undefined; + + return { + url, + title: + this.firstString(metadata?.title) ?? + this.firstString(metadata?.description) ?? + metadata?.sourceURL ?? + url, + content, + success: !metadataError && content.trim().length > 0, + error: metadataError ?? (content.trim().length === 0 ? "No content returned" : undefined), + metadata, + }; + } + + private async pollBatchScrape(jobId: string, urls: string[]): Promise { + const startedAt = Date.now(); + + while (Date.now() - startedAt < BATCH_POLL_TIMEOUT_MS) { + const response = await fetch(`${this.baseUrl}/batch/scrape/${jobId}`, { + method: "GET", + headers: this.getHeaders(), + }); + + if (!response.ok) { + throw new Error(await this.parseErrorResponse(response)); + } + + const result = await response.json(); + if (result?.status === "completed" && Array.isArray(result?.data)) { + return this.normalizeBatchResults(urls, result.data); + } + + if (result?.status === "failed") { + throw new Error( + typeof result?.error === "string" + ? result.error + : "Firecrawl batch scrape failed", + ); + } + + await new Promise((resolve) => setTimeout(resolve, BATCH_POLL_INTERVAL_MS)); + } + + throw new Error("Firecrawl batch scrape timed out"); + } + + private normalizeBatchResults(urls: string[], documents: unknown[]): FirecrawlPageResult[] { + const unusedDocuments = [...documents]; + const bySourceUrl = new Map(); + + for (const document of documents) { + if (!document || typeof document !== "object" || Array.isArray(document)) { + continue; + } + + const metadata = + "metadata" in document && document.metadata && typeof document.metadata === "object" && !Array.isArray(document.metadata) + ? (document.metadata as FirecrawlMetadata) + : undefined; + + const sourceUrl = + typeof metadata?.sourceURL === "string" + ? metadata.sourceURL + : typeof metadata?.url === "string" + ? metadata.url + : undefined; + + if (sourceUrl && !bySourceUrl.has(sourceUrl)) { + bySourceUrl.set(sourceUrl, document); + } + } + + return urls.map((url, index) => { + const matched = bySourceUrl.get(url); + if (matched) { + return this.normalizePageResult(url, matched); + } + + const fallback = unusedDocuments[index]; + return fallback + ? this.normalizePageResult(url, fallback) + : { + url, + title: url, + content: "", + success: false, + error: "No result returned for URL", + }; + }); + } + + async scrapeUrl(url: string): Promise { if (!this.apiKey) { - return { success: false, error: "Firecrawl API key not configured" }; + return { + url, + title: url, + content: "", + success: false, + error: "Firecrawl API key not configured", + }; } try { @@ -40,45 +235,110 @@ export class FirecrawlClient { const response = await fetch(`${this.baseUrl}/scrape`, { method: "POST", - headers: { - "Content-Type": "application/json", - "Authorization": `Bearer ${this.apiKey}`, - }, + headers: this.getHeaders(), body: JSON.stringify({ url, - formats: ["markdown"], - onlyMainContent: true, - waitFor: 1000, // Wait for dynamic content + ...DEFAULT_SCRAPE_OPTIONS, }), }); if (!response.ok) { - if (response.status === 401) { - return { success: false, error: "Invalid Firecrawl API key" }; - } - if (response.status === 429) { - return { success: false, error: "Firecrawl rate limit exceeded" }; - } - const errorText = await response.text(); - throw new Error(`Firecrawl API error: ${response.status} ${errorText}`); + return { + url, + title: url, + content: "", + success: false, + error: await this.parseErrorResponse(response), + }; } const result = await response.json(); - - if (!result.success) { - return { success: false, error: result.error || "Unknown Firecrawl error" }; + if (result?.success === false) { + return { + url, + title: url, + content: "", + success: false, + error: typeof result?.error === "string" ? result.error : "Unknown Firecrawl error", + }; } - return { - success: true, - data: result.data, - }; - } catch (error: any) { + return this.normalizePageResult(url, result?.data); + } catch (error) { logger.error(`❌ [Firecrawl] Error scraping ${url}:`, error); return { + url, + title: url, + content: "", success: false, - error: error.message || String(error), + error: error instanceof Error ? error.message : String(error), }; } } + + async scrapeUrls(urls: string[]): Promise { + if (urls.length === 0) { + return []; + } + + if (!this.apiKey) { + return urls.map((url) => ({ + url, + title: url, + content: "", + success: false, + error: "Firecrawl API key not configured", + })); + } + + if (urls.length === 1) { + return [await this.scrapeUrl(urls[0])]; + } + + try { + logger.debug(`🔥 [Firecrawl] Batch scraping ${urls.length} URLs`); + + const response = await fetch(`${this.baseUrl}/batch/scrape`, { + method: "POST", + headers: this.getHeaders(), + body: JSON.stringify({ + urls, + ...DEFAULT_SCRAPE_OPTIONS, + }), + }); + + if (!response.ok) { + const error = await this.parseErrorResponse(response); + return urls.map((url) => ({ + url, + title: url, + content: "", + success: false, + error, + })); + } + + const result = await response.json(); + + if (Array.isArray(result?.data)) { + return this.normalizeBatchResults(urls, result.data); + } + + if (typeof result?.id === "string") { + return await this.pollBatchScrape(result.id, urls); + } + + throw new Error("Unexpected Firecrawl batch scrape response"); + } catch (error) { + logger.error("❌ [Firecrawl] Error batch scraping URLs:", error); + const message = error instanceof Error ? error.message : String(error); + return urls.map((url) => ({ + url, + title: url, + content: "", + success: false, + error: message, + })); + } + } } diff --git a/src/lib/ai/utils/url-processor.ts b/src/lib/ai/utils/url-processor.ts deleted file mode 100644 index 3b96cb13..00000000 --- a/src/lib/ai/utils/url-processor.ts +++ /dev/null @@ -1,326 +0,0 @@ -import { google } from "@ai-sdk/google"; -import { generateText } from "ai"; - -import { logger } from "@/lib/utils/logger"; -import { FirecrawlClient } from "@/lib/ai/utils/firecrawl"; - -export interface UrlContent { - url: string; - title: string; - content: string; - success: boolean; - error?: string; - source?: "google-context" | "firecrawl" | "google-grounding"; -} - -// Scraping Mode Configuration -type ScrapingMode = "hybrid" | "firecrawl-only" | "google-only" | "direct-only"; - -function getScrapingMode(): ScrapingMode { - const mode = process.env.SCRAPING_MODE?.toLowerCase(); - - // If explicitly set, respect it - if (mode === "firecrawl-only") return "firecrawl-only"; - if (mode === "google-only") return "google-only"; - if (mode === "direct-only") return "direct-only"; - - // Default to hybrid - return "hybrid"; -} - -/** - * STRATEGY INTERFACE - */ -interface ScrapingStrategy { - name: string; - process(url: string): Promise; -} - -/** - * STRATEGY 1: Google URL Context (Official/Smart) - */ -class GoogleContextStrategy implements ScrapingStrategy { - name = "google-context"; - - async process(url: string): Promise { - try { - logger.debug(`🌐 [Scraper:Google] Trying URL Context for: ${url}`); - - // Add timeout to prevent hanging - const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), 15000); // 15 second timeout - - try { - const result = await generateText({ - model: google("gemini-2.5-flash-lite"), - tools: { - urlContext: google.tools.urlContext({}) as any, - }, - // @ts-ignore - maxToolRoundtrips is valid but missing from current type defs - maxToolRoundtrips: 2, - prompt: `Please read the content of the following URL: ${url} - - If you can assume the content, extract the main title and the full text content of the page. - Do not summarize wildly, just provide the content as accurately as possible. - - Return the result in this format: - Title: [Title] - Content: [Content]`, - abortSignal: controller.signal, - }); - - clearTimeout(timeoutId); - - const { text } = result; - - // --- DEBUG LOGGING --- - logger.debug(`🔍 [Scraper:Google] Response received:`, { - textLength: text?.length || 0, - textPreview: text?.substring(0, 200) || '(empty)', - hasText: !!text, - }); - - // --- SIMPLIFED VALIDATION --- - - // --- SIMPLIFED VALIDATION --- - - // 1. Check if we actually got text back - if (!text || text.trim().length === 0) { - logger.warn(`⚠️ [Scraper:Google] No text returned.`); - return null; // Fallback to Firecrawl - } - - // 2. Check for common "Access Denied" or "Unable to read" phrases from LLM - const lowerText = text.toLowerCase(); - if ( - lowerText.includes("i cannot access") || - lowerText.includes("i am unable to read") || - lowerText.includes("access is denied") || - lowerText.includes("403 forbidden") - ) { - logger.warn(`⚠️ [Scraper:Google] LLM reported access failure.`); - return null; // Fallback to Firecrawl - } - - // If we passed checks, we assume success. - // Try to parse title/content if possible, or just use whole text. - const titleMatch = text.match(/Title:\s*(.*?)(\n|$)/); - const contentMatch = text.match(/Content:\s*([\s\S]*)/); - - let title = url; - let content = text; - - if (titleMatch && contentMatch) { - title = titleMatch[1].trim(); - content = contentMatch[1].trim(); - } - - return { - url, - title, - content, - success: true, - source: "google-context", - }; - } catch (error: any) { - clearTimeout(timeoutId); - - if (error.name === 'AbortError') { - logger.warn(`⏱️ [Scraper:Google] Timeout after 15s for ${url}`); - return null; - } - throw error; // Re-throw non-timeout errors - } - - } catch (error) { - logger.error(`❌ [Scraper:Google] Exception caught for ${url}:`, { - error: error instanceof Error ? error.message : String(error), - stack: error instanceof Error ? error.stack : undefined, - }); - return null; - } - } -} - -/** - * STRATEGY 2: Firecrawl (Robust/Stealth/Dynamic) - */ -class FirecrawlStrategy implements ScrapingStrategy { - name = "firecrawl"; - private client: FirecrawlClient; - - constructor() { - this.client = new FirecrawlClient(); - } - - async process(url: string): Promise { - try { - // Check if configured - if (!process.env.FIRECRAWL_API_KEY) { - logger.debug("⚠️ [Scraper:Firecrawl] Skipped: No API key provided"); - return null; - } - - logger.debug(`🔥 [Scraper:Firecrawl] Scraping URL: ${url}`); - - const result = await this.client.scrapeUrl(url); - - if (!result.success || !result.data) { - logger.warn(`⚠️ [Scraper:Firecrawl] Failed: ${result.error}`); - return null; - } - - // Prefer robust metadata, fall back to simple props - const title = result.data.metadata?.title || result.data.metadata?.ogTitle || url; - const content = result.data.markdown || result.data.content || ""; - - if (!content) return null; - - return { - url, - title, - content, - success: true, - source: "firecrawl", - }; - - } catch (error) { - logger.error(`❌ [Scraper:Firecrawl] Error:`, error); - return null; - } - } -} - -/** - * STRATEGY 3: Google Grounding (Search Fallback) - * If we can't read the page directly, just ask Google what it's about. - */ -class GoogleGroundingStrategy implements ScrapingStrategy { - name = "google-grounding"; - - async process(url: string): Promise { - try { - logger.debug(`🔍 [Scraper:Grounding] Falling back to Search for: ${url}`); - - const { text } = await generateText({ - model: google("gemini-2.5-flash"), - tools: { - googleSearch: google.tools.googleSearch({}), - }, - // @ts-ignore - maxToolRoundtrips is valid but missing from current type defs - maxToolRoundtrips: 2, - prompt: `Find information about this specific URL: ${url} - - Provide a detailed summary of the content found on this page. - Focus on the main topics, key points, and purpose of the page. - - Format the output as: - Title: [Page Title] - Content: [Detailed Summary]`, - }); - - // Flexible parsing - const titleMatch = text.match(/Title:\s*(.*?)(\n|$)/); - const contentMatch = text.match(/Content:\s*([\s\S]*)/); - - const title = titleMatch ? titleMatch[1].trim() : url; - let content = contentMatch ? contentMatch[1].trim() : text; - - if (content.length < 50) return null; - - return { - url, - title, - content, - success: true, - source: "google-grounding", - }; - - } catch (error) { - logger.error(`❌ [Scraper:Grounding] Error:`, error); - return null; - } - } -} - -/** - * Shared utility to process URLs uses Strategy Pattern - */ -export class UrlProcessor { - - static async processUrl(url: string): Promise { - const mode = getScrapingMode(); - logger.debug(`🌐 [UrlProcessor] Processing ${url} (Mode: ${mode})`); - - // Define strategies - const strategies: ScrapingStrategy[] = []; - - // --- STRATEGY SELECTION LOGIC --- - - if (mode === "firecrawl-only") { - strategies.push(new FirecrawlStrategy()); - } else if (mode === "google-only") { - strategies.push(new GoogleContextStrategy()); - strategies.push(new GoogleGroundingStrategy()); // Add grounding as backup even in google-only - } else { - // HYBRID (Default) - - // 1. Google Context (Best for "Reading" specific pages) - strategies.push(new GoogleContextStrategy()); - - // 2. Firecrawl (Best for "Accessing" blocked/JS pages) - if (process.env.FIRECRAWL_API_KEY) { - strategies.push(new FirecrawlStrategy()); - } - - // 3. Google Grounding (Absolute "Last Resort" - Search for the page info) - strategies.push(new GoogleGroundingStrategy()); - } - - // --- EXECUTION LOOP --- - - for (const strategy of strategies) { - const result = await strategy.process(url); - if (result && result.success) { - logger.debug(`✅ [UrlProcessor] Success via ${strategy.name}`); - return result; - } - } - - // --- FINAL FAILURE --- - - logger.error(`❌ [UrlProcessor] All strategies failed for ${url}`); - return { - url, - title: url, - content: "", - success: false, - error: `Failed to fetch content using strategies: ${strategies.map(s => s.name).join(', ')}`, - }; - } - - /** - * Process multiple URLs in parallel with graceful partial failure handling - */ - static async processUrls(urls: string[]): Promise { - const results = await Promise.allSettled( - urls.map(url => this.processUrl(url)) - ); - - return results.map((result, index) => { - if (result.status === 'fulfilled') { - return result.value; - } else { - // Log the error but return a failed UrlContent object - logger.error(`❌ [UrlProcessor] Failed to process URL ${urls[index]}:`, result.reason); - return { - url: urls[index], - title: urls[index], - content: '', - success: false, - error: result.reason instanceof Error ? result.reason.message : String(result.reason), - }; - } - }); - } -} From 4902bd9e014338c3a30b2994773ce9c8298cb775 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 3 Apr 2026 16:14:34 -0400 Subject: [PATCH 2/4] fix: align tool outputs with typed UI contracts Add output schemas for workspace tools, return typed web search results instead of stringified JSON, and normalize legacy tool payloads so persisted threads stay compatible. Made-with: Cursor --- .../assistant-ui/WebSearchToolUI.tsx | 69 ++++--------------- src/hooks/ai/use-create-card-from-message.ts | 46 ++++++++----- .../legacy-tool-message-compat.test.ts | 40 ++++++++++- src/lib/ai/legacy-tool-message-compat.ts | 15 ++++ src/lib/ai/tool-result-schemas.ts | 26 ++++--- src/lib/ai/tools/read-workspace.ts | 28 ++++++++ src/lib/ai/tools/search-workspace.ts | 25 ++++++- src/lib/ai/tools/web-search.ts | 14 ++-- src/lib/ai/web-search-shared.ts | 58 ++++++++++++++++ 9 files changed, 228 insertions(+), 93 deletions(-) create mode 100644 src/lib/ai/web-search-shared.ts diff --git a/src/components/assistant-ui/WebSearchToolUI.tsx b/src/components/assistant-ui/WebSearchToolUI.tsx index ee1fa643..c75f7d87 100644 --- a/src/components/assistant-ui/WebSearchToolUI.tsx +++ b/src/components/assistant-ui/WebSearchToolUI.tsx @@ -15,6 +15,8 @@ import { } from "@assistant-ui/react"; import { ToolUIErrorBoundary } from "@/components/tool-ui/shared"; +import { parseWebSearchResult } from "@/lib/ai/tool-result-schemas"; +import type { WebSearchResult } from "@/lib/ai/web-search-shared"; import { Collapsible, CollapsibleContent, @@ -204,39 +206,15 @@ const getDomain = (url: string) => { } }; -type GroundingChunk = { - web?: { - uri?: string; - title?: string; - }; -}; - -type WebSearchResultPayload = { - text?: string; - groundingMetadata?: { - webSearchQueries?: string[]; - groundingChunks?: GroundingChunk[]; - }; -}; - const WebSearchContent: FC<{ status: { type: string }; - result: string | null; + result: WebSearchResult | null; }> = ({ status, result }) => { const isRunning = status.type === "running"; - - let parsed: WebSearchResultPayload | null = null; - try { - if (result) { - parsed = JSON.parse(result) as WebSearchResultPayload; - } - } catch { - parsed = { text: result }; - } - + const parsed = result ? parseWebSearchResult(result) : null; const metadata = parsed?.groundingMetadata; - const queries = metadata?.webSearchQueries as string[] | undefined; - const chunks = metadata?.groundingChunks; + const queries = metadata?.webSearchQueries; + const sources = parsed?.sources ?? []; return ( @@ -254,8 +232,8 @@ const WebSearchContent: FC<{
Search Queries:
- {queries.map((q, i) => ( -
+ {queries.map((q) => ( +
{q}
))} @@ -263,39 +241,22 @@ const WebSearchContent: FC<{
)} - {/* 2. Show the sources found (Grounding Chunks) */} - {chunks && chunks.length > 0 ? ( + {sources.length > 0 ? (
Sources:
- {chunks.map((chunk, i) => { - const uri = chunk.web?.uri || ""; - const title = chunk.web?.title || "Untitled Source"; - - // Use title for favicon domain as requested, fallback to uri domain - const faviconDomain = getDomain(title).includes('.') ? getDomain(title) : getDomain(uri); - const faviconUrl = `https://www.google.com/s2/favicons?domain=${faviconDomain}&sz=32`; - + {sources.map((source) => { + const uri = source.url; + const title = source.title || getDomain(uri); return ( - {/* Favicon */} -
- {/* eslint-disable-next-line @next/next/no-img-element */} - { - e.currentTarget.style.display = 'none'; - }} - /> -
+
{title} @@ -332,7 +293,7 @@ WebSearchContent.displayName = "WebSearchContent"; */ export const WebSearchToolUI = makeAssistantToolUI<{ query: string; -}, string>({ +}, WebSearchResult>({ toolName: "webSearch", render: function WebSearchToolUI({ status, result }) { return ( diff --git a/src/hooks/ai/use-create-card-from-message.ts b/src/hooks/ai/use-create-card-from-message.ts index 4ef713ba..8c051e3d 100644 --- a/src/hooks/ai/use-create-card-from-message.ts +++ b/src/hooks/ai/use-create-card-from-message.ts @@ -7,6 +7,7 @@ import { useWorkspaceStore } from "@/lib/stores/workspace-store"; import { useUIStore } from "@/lib/stores/ui-store"; import { useQueryClient } from "@tanstack/react-query"; import { logger } from "@/lib/utils/logger"; +import { normalizeWebSearchResult } from "@/lib/ai/web-search-shared"; interface CreateCardOptions { debounceMs?: number; @@ -27,24 +28,31 @@ export function useCreateCardFromMessage(options: CreateCardOptions = {}) { const currentWorkspaceId = useWorkspaceStore((state) => state.currentWorkspaceId); const queryClient = useQueryClient(); - // Helper to extract sources from a tool result JSON string - const extractSourcesFromToolResult = (resultJson: string) => { - try { - const parsed = JSON.parse(resultJson); - const chunks = parsed?.groundingMetadata?.groundingChunks || []; - const extractedSources: Array<{ title: string; url: string; favicon?: string }> = []; - - for (const chunk of chunks) { - const uri = chunk?.web?.uri; - const title = chunk?.web?.title; - if (uri && title) { - extractedSources.push({ title, url: uri }); - } - } - return extractedSources; - } catch (e) { + const extractSourcesFromToolResult = (result: unknown) => { + const parsed = normalizeWebSearchResult(result); + if (!parsed) { return []; } + + if (parsed.sources.length > 0) { + return parsed.sources.map((source) => ({ + title: source.title, + url: source.url, + })); + } + + const chunks = parsed.groundingMetadata?.groundingChunks || []; + const extractedSources: Array<{ title: string; url: string; favicon?: string }> = []; + + for (const chunk of chunks) { + const uri = chunk?.web?.uri; + const title = chunk?.web?.title; + if (uri && title) { + extractedSources.push({ title, url: uri }); + } + } + + return extractedSources; }; const createCard = useCallback(async () => { @@ -120,14 +128,14 @@ export function useCreateCardFromMessage(options: CreateCardOptions = {}) { if (part.type === 'tool-result' && (part as any).toolName === 'webSearch') { const result = (part as any).result; if (result) { - allSources.push(...extractSourcesFromToolResult(JSON.stringify(result))); + allSources.push(...extractSourcesFromToolResult(result)); } } } } // Check simplified Vercel AI SDK structure if ((msg as any).toolName === 'webSearch' && (msg as any).content) { - allSources.push(...extractSourcesFromToolResult(JSON.stringify(msg.content))); + allSources.push(...extractSourcesFromToolResult((msg as any).content)); } } @@ -153,7 +161,7 @@ export function useCreateCardFromMessage(options: CreateCardOptions = {}) { if (Array.isArray(invocations)) { for (const tool of invocations) { if (tool.toolName === 'webSearch' && tool.state === 'result') { - allSources.push(...extractSourcesFromToolResult(JSON.stringify(tool.result))); + allSources.push(...extractSourcesFromToolResult(tool.result)); } } } diff --git a/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts b/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts index 823c4891..be38c7ed 100644 --- a/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts +++ b/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts @@ -3,7 +3,7 @@ import type { UIMessage } from "ai"; import { normalizeLegacyToolMessages } from "../legacy-tool-message-compat"; describe("normalizeLegacyToolMessages", () => { - it("normalizes legacy processUrls jsonInput tool args", () => { + it("normalizes legacy processUrls jsonInput tool args and drops instruction", () => { const messages = [ { id: "1", @@ -31,7 +31,6 @@ describe("normalizeLegacyToolMessages", () => { type: "tool-processUrls", input: { urls: ["https://example.com"], - instruction: "Extract dates.", }, }); }); @@ -61,4 +60,41 @@ describe("normalizeLegacyToolMessages", () => { output: "The answer is 6765.", }); }); + + it("normalizes legacy webSearch string outputs", () => { + const messages = [ + { + id: "1", + role: "assistant", + parts: [ + { + type: "tool-webSearch", + toolCallId: "call_3", + state: "output-available", + input: { query: "latest AI news" }, + output: JSON.stringify({ + text: "Summary text", + sources: [{ title: "Example", url: "https://example.com" }], + groundingMetadata: { + groundingChunks: [ + { web: { uri: "https://example.com", title: "Example" } }, + ], + }, + }), + }, + ], + }, + ] as UIMessage[]; + + const normalized = normalizeLegacyToolMessages(messages); + const part = normalized[0]?.parts[0]; + + expect(part).toMatchObject({ + type: "tool-webSearch", + output: { + text: "Summary text", + sources: [{ title: "Example", url: "https://example.com" }], + }, + }); + }); }); diff --git a/src/lib/ai/legacy-tool-message-compat.ts b/src/lib/ai/legacy-tool-message-compat.ts index d64ff765..1c447942 100644 --- a/src/lib/ai/legacy-tool-message-compat.ts +++ b/src/lib/ai/legacy-tool-message-compat.ts @@ -1,5 +1,6 @@ import type { UIMessage } from "ai"; import { normalizeProcessUrlsArgs } from "./process-urls-shared"; +import { normalizeWebSearchResult } from "./web-search-shared"; function normalizeExecuteCodeOutput(output: unknown): unknown { if (typeof output === "string") { @@ -16,6 +17,11 @@ function normalizeExecuteCodeOutput(output: unknown): unknown { return output; } +function normalizeWebSearchOutput(output: unknown): unknown { + const normalized = normalizeWebSearchResult(output); + return normalized ?? output; +} + export function normalizeLegacyToolMessages(messages: UIMessage[]): UIMessage[] { return messages.map((message) => { if (!Array.isArray(message.parts)) { @@ -41,6 +47,15 @@ export function normalizeLegacyToolMessages(messages: UIMessage[]): UIMessage[] return { ...part, output: normalizeExecuteCodeOutput(part.output) }; } + if ( + part.type === "tool-webSearch" && + "state" in part && + part.state === "output-available" && + "output" in part + ) { + return { ...part, output: normalizeWebSearchOutput(part.output) }; + } + return part; }); diff --git a/src/lib/ai/tool-result-schemas.ts b/src/lib/ai/tool-result-schemas.ts index b326619f..19f5d050 100644 --- a/src/lib/ai/tool-result-schemas.ts +++ b/src/lib/ai/tool-result-schemas.ts @@ -1,5 +1,10 @@ import { z } from "zod"; import { parseWithSchema } from "@/components/tool-ui/shared"; +import { ProcessUrlsOutputSchema } from "@/lib/ai/process-urls-shared"; +import { + WebSearchResultSchema, + normalizeWebSearchResult, +} from "@/lib/ai/web-search-shared"; /** * Shared schemas and parsers for tool results. Used by assistant-ui Tool UIs @@ -145,18 +150,21 @@ export function parseStringResult(input: unknown): string { } /** processUrls – result is string or { text, metadata } */ -export const URLContextResultSchema = z.union([ - z.string(), - z - .object({ - text: z.string().optional(), - metadata: z.any().optional(), - }) - .passthrough(), -]); +export const URLContextResultSchema = z.union([z.string(), ProcessUrlsOutputSchema]); export type URLContextResult = z.infer; export function parseURLContextResult(input: unknown): URLContextResult { return parseWithSchema(URLContextResultSchema, input, "URLContextResult"); } + +export type WebSearchResult = z.infer; + +export function parseWebSearchResult(input: unknown): WebSearchResult { + const normalized = normalizeWebSearchResult(input); + if (normalized) { + return normalized; + } + + return parseWithSchema(WebSearchResultSchema, input, "WebSearchResult"); +} diff --git a/src/lib/ai/tools/read-workspace.ts b/src/lib/ai/tools/read-workspace.ts index 73c4f455..93957b5d 100644 --- a/src/lib/ai/tools/read-workspace.ts +++ b/src/lib/ai/tools/read-workspace.ts @@ -12,6 +12,33 @@ const DEFAULT_LIMIT = 500; const MAX_LIMIT = 2000; const MAX_LINE_LENGTH = 2000; +const ReadWorkspaceResultSchema = z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + message: z.string(), + }), + z.object({ + success: z.literal(true), + itemName: z.string(), + type: z.string(), + path: z.string(), + content: z.string(), + totalLines: z.number().int().nonnegative(), + lineStart: z.number().int().min(1), + lineEnd: z.number().int().nonnegative(), + hasMore: z.boolean(), + rangeNote: z.string(), + nextLineStart: z.number().int().min(1).optional(), + totalPages: z.number().int().min(1).optional(), + pageRange: z + .object({ + start: z.number().int().min(1).optional(), + end: z.number().int().min(1).optional(), + }) + .optional(), + }), +]); + export function createReadWorkspaceTool(ctx: WorkspaceToolContext) { return tool({ description: @@ -57,6 +84,7 @@ export function createReadWorkspaceTool(ctx: WorkspaceToolContext) { .describe("For PDFs only: 1-indexed end page inclusive (e.g. 10 for pages 5–10). Use with pageStart."), }) ), + outputSchema: zodSchema(ReadWorkspaceResultSchema), strict: true, execute: async ({ path, itemName, lineStart = 1, limit = DEFAULT_LIMIT, pageStart, pageEnd }) => { if (!path?.trim() && !itemName?.trim()) { diff --git a/src/lib/ai/tools/search-workspace.ts b/src/lib/ai/tools/search-workspace.ts index dd664fc0..f54c7f00 100644 --- a/src/lib/ai/tools/search-workspace.ts +++ b/src/lib/ai/tools/search-workspace.ts @@ -9,6 +9,21 @@ import type { Item } from "@/lib/workspace-state/types"; const MAX_LINE_LENGTH = 2000; const MAX_MATCHES = 100; +const SearchWorkspaceResultSchema = z.discriminatedUnion("success", [ + z.object({ + success: z.literal(false), + message: z.string(), + matches: z.number().int().nonnegative(), + output: z.string(), + }), + z.object({ + success: z.literal(true), + matches: z.number().int().nonnegative(), + truncated: z.boolean(), + output: z.string(), + }), +]); + function buildRegex(pattern: string): RegExp { const hasRegexChars = /[.*+?^${}()|[\]\\]/.test(pattern); if (hasRegexChars) { @@ -33,6 +48,7 @@ export function createSearchWorkspaceTool(ctx: WorkspaceToolContext) { path: z.string().optional().describe("Folder prefix (Physics/) or exact item path (Physics/documents/File.md)"), }) ), + outputSchema: zodSchema(SearchWorkspaceResultSchema), strict: true, execute: async ({ pattern, include, path: pathPrefix }) => { if (!pattern?.trim()) { @@ -40,7 +56,14 @@ export function createSearchWorkspaceTool(ctx: WorkspaceToolContext) { } const accessResult = await loadStateForTool(ctx); - if (!accessResult.success) return accessResult; + if (!accessResult.success) { + return { + success: false, + message: accessResult.message, + matches: 0, + output: "", + }; + } const { state } = accessResult; let items: Item[] = state.items.filter((i) => i.type !== "folder"); diff --git a/src/lib/ai/tools/web-search.ts b/src/lib/ai/tools/web-search.ts index 9a776e2e..55ed761a 100644 --- a/src/lib/ai/tools/web-search.ts +++ b/src/lib/ai/tools/web-search.ts @@ -1,6 +1,10 @@ import { z } from "zod"; import { tool, generateText, stepCountIs, zodSchema } from "ai"; import { google } from "@ai-sdk/google"; +import { + WebSearchResultSchema, + type WebSearchResult, +} from "@/lib/ai/web-search-shared"; const VERTEX_REDIRECT_HOST = "vertexaisearch.cloud.google.com"; const VERTEX_REDIRECT_PATH = "/grounding-api-redirect/"; @@ -90,12 +94,6 @@ export async function resolveGroundingChunksToSources( return resolved.filter((s): s is { title: string; url: string } => s !== null); } -export type WebSearchResult = { - text: string; - sources: Array<{ title: string; url: string }>; - groundingMetadata?: { groundingChunks?: unknown[] }; -}; - /** * Execute a web search and return text + sources. Used by both chat webSearch tool and autogen. */ @@ -152,9 +150,9 @@ export function createWebSearchTool() { }) ), strict: true, + outputSchema: zodSchema(WebSearchResultSchema), execute: async ({ query }) => { - const { text, groundingMetadata } = await executeWebSearch(query); - return JSON.stringify({ text, groundingMetadata }); + return await executeWebSearch(query); }, }); } diff --git a/src/lib/ai/web-search-shared.ts b/src/lib/ai/web-search-shared.ts new file mode 100644 index 00000000..aca02643 --- /dev/null +++ b/src/lib/ai/web-search-shared.ts @@ -0,0 +1,58 @@ +import { z } from "zod"; + +export const WebSearchSourceSchema = z.object({ + title: z.string(), + url: z.string(), +}); + +export const GroundingChunkSchema = z + .object({ + web: z + .object({ + uri: z.string().optional(), + title: z.string().optional(), + }) + .optional(), + }) + .passthrough(); + +export const WebSearchGroundingMetadataSchema = z + .object({ + webSearchQueries: z.array(z.string()).optional(), + groundingChunks: z.array(GroundingChunkSchema).optional(), + }) + .passthrough(); + +export const WebSearchResultSchema = z.object({ + text: z.string(), + sources: z.array(WebSearchSourceSchema), + groundingMetadata: WebSearchGroundingMetadataSchema.optional(), +}); + +export type WebSearchSource = z.infer; +export type GroundingChunk = z.infer; +export type WebSearchGroundingMetadata = z.infer< + typeof WebSearchGroundingMetadataSchema +>; +export type WebSearchResult = z.infer; + +export function normalizeWebSearchResult( + input: unknown, +): WebSearchResult | null { + const direct = WebSearchResultSchema.safeParse(input); + if (direct.success) { + return direct.data; + } + + if (typeof input !== "string") { + return null; + } + + try { + const parsed = JSON.parse(input); + const normalized = WebSearchResultSchema.safeParse(parsed); + return normalized.success ? normalized.data : null; + } catch { + return null; + } +} From 0afdbc380b8ffc3acad2cf0df6e4e934139d1de3 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 3 Apr 2026 16:19:09 -0400 Subject: [PATCH 3/4] fix: cleanup --- .../assistant-ui/URLContextToolUI.tsx | 2 - src/lib/ai/process-urls-shared.ts | 2 - src/lib/ai/tools/process-urls.ts | 8 -- src/lib/ai/utils/firecrawl.ts | 133 +----------------- 4 files changed, 2 insertions(+), 143 deletions(-) diff --git a/src/components/assistant-ui/URLContextToolUI.tsx b/src/components/assistant-ui/URLContextToolUI.tsx index 6a73fccf..05b527d8 100644 --- a/src/components/assistant-ui/URLContextToolUI.tsx +++ b/src/components/assistant-ui/URLContextToolUI.tsx @@ -215,7 +215,6 @@ type ProcessUrlsResult = text: string; metadata?: { urlMetadata?: URLMetadata[] | null; - groundingChunks?: unknown[] | null; sources?: SourceMetadata[] | null; }; }; @@ -232,7 +231,6 @@ export const URLContextToolUI = makeAssistantToolUI<{ const parsedResult = result != null ? parseURLContextResult(result) : null; type Meta = { urlMetadata?: URLMetadata[]; - groundingChunks?: unknown[]; sources?: SourceMetadata[]; }; const metadata = (typeof parsedResult === "object" && parsedResult !== null && "metadata" in parsedResult ? (parsedResult as { metadata?: Meta }).metadata : null) as Meta | null; diff --git a/src/lib/ai/process-urls-shared.ts b/src/lib/ai/process-urls-shared.ts index 1ce8c7b5..eda96d5b 100644 --- a/src/lib/ai/process-urls-shared.ts +++ b/src/lib/ai/process-urls-shared.ts @@ -16,7 +16,6 @@ export const ProcessUrlsOutputSchema = z.object({ text: z.string(), metadata: z .object({ - provider: z.string().optional(), urlMetadata: z .array( z.object({ @@ -26,7 +25,6 @@ export const ProcessUrlsOutputSchema = z.object({ ) .nullable() .optional(), - groundingChunks: z.array(z.unknown()).nullable().optional(), sources: z .array( z.object({ diff --git a/src/lib/ai/tools/process-urls.ts b/src/lib/ai/tools/process-urls.ts index 610b8b2e..892a778f 100644 --- a/src/lib/ai/tools/process-urls.ts +++ b/src/lib/ai/tools/process-urls.ts @@ -44,14 +44,12 @@ export function createProcessUrlsTool() { return { text: `Error: This tool only handles regular web URLs, not uploaded files or video URLs (${fileUrls.join(", ")})`, metadata: { - provider: "firecrawl", urlMetadata: urlList.map((url) => ({ retrievedUrl: url, urlRetrievalStatus: fileUrls.includes(url) ? "URL_RETRIEVAL_STATUS_UNSUPPORTED" : "URL_RETRIEVAL_STATUS_SKIPPED", })), - groundingChunks: null, sources: null, }, }; @@ -79,14 +77,12 @@ export function createProcessUrlsTool() { return { text: `Failed to process any of the provided URLs. Errors: ${failedResults.map((result) => `${result.url}: ${result.error || "Unknown error"}`).join("; ")}`, metadata: { - provider: "firecrawl", urlMetadata: scrapedResults.map((result) => ({ retrievedUrl: result.url, urlRetrievalStatus: result.success ? "URL_RETRIEVAL_STATUS_SUCCESS" : "URL_RETRIEVAL_STATUS_FAILED", })), - groundingChunks: null, sources: null, }, }; @@ -99,14 +95,12 @@ export function createProcessUrlsTool() { return { text: combinedText, metadata: { - provider: "firecrawl", urlMetadata: scrapedResults.map((result) => ({ retrievedUrl: result.url, urlRetrievalStatus: result.success ? "URL_RETRIEVAL_STATUS_SUCCESS" : "URL_RETRIEVAL_STATUS_FAILED", })), - groundingChunks: null, sources: successfulResults.map(({ url, title }) => ({ uri: url, title })), }, }; @@ -119,9 +113,7 @@ export function createProcessUrlsTool() { return { text: `Error processing web URLs: ${error instanceof Error ? error.message : String(error)}`, metadata: { - provider: "firecrawl", urlMetadata: null, - groundingChunks: null, sources: null, }, }; diff --git a/src/lib/ai/utils/firecrawl.ts b/src/lib/ai/utils/firecrawl.ts index 238d5f1e..5290b81e 100644 --- a/src/lib/ai/utils/firecrawl.ts +++ b/src/lib/ai/utils/firecrawl.ts @@ -6,9 +6,6 @@ const DEFAULT_SCRAPE_OPTIONS = { waitFor: 1000, }; -const BATCH_POLL_INTERVAL_MS = 1500; -const BATCH_POLL_TIMEOUT_MS = 45000; - export interface FirecrawlMetadata { title?: string | string[]; description?: string | string[]; @@ -20,12 +17,6 @@ export interface FirecrawlMetadata { [key: string]: unknown; } -export interface FirecrawlDocument { - content?: string; - markdown?: string; - metadata?: FirecrawlMetadata; -} - export interface FirecrawlPageResult { url: string; title: string; @@ -142,83 +133,6 @@ export class FirecrawlClient { }; } - private async pollBatchScrape(jobId: string, urls: string[]): Promise { - const startedAt = Date.now(); - - while (Date.now() - startedAt < BATCH_POLL_TIMEOUT_MS) { - const response = await fetch(`${this.baseUrl}/batch/scrape/${jobId}`, { - method: "GET", - headers: this.getHeaders(), - }); - - if (!response.ok) { - throw new Error(await this.parseErrorResponse(response)); - } - - const result = await response.json(); - if (result?.status === "completed" && Array.isArray(result?.data)) { - return this.normalizeBatchResults(urls, result.data); - } - - if (result?.status === "failed") { - throw new Error( - typeof result?.error === "string" - ? result.error - : "Firecrawl batch scrape failed", - ); - } - - await new Promise((resolve) => setTimeout(resolve, BATCH_POLL_INTERVAL_MS)); - } - - throw new Error("Firecrawl batch scrape timed out"); - } - - private normalizeBatchResults(urls: string[], documents: unknown[]): FirecrawlPageResult[] { - const unusedDocuments = [...documents]; - const bySourceUrl = new Map(); - - for (const document of documents) { - if (!document || typeof document !== "object" || Array.isArray(document)) { - continue; - } - - const metadata = - "metadata" in document && document.metadata && typeof document.metadata === "object" && !Array.isArray(document.metadata) - ? (document.metadata as FirecrawlMetadata) - : undefined; - - const sourceUrl = - typeof metadata?.sourceURL === "string" - ? metadata.sourceURL - : typeof metadata?.url === "string" - ? metadata.url - : undefined; - - if (sourceUrl && !bySourceUrl.has(sourceUrl)) { - bySourceUrl.set(sourceUrl, document); - } - } - - return urls.map((url, index) => { - const matched = bySourceUrl.get(url); - if (matched) { - return this.normalizePageResult(url, matched); - } - - const fallback = unusedDocuments[index]; - return fallback - ? this.normalizePageResult(url, fallback) - : { - url, - title: url, - content: "", - success: false, - error: "No result returned for URL", - }; - }); - } - async scrapeUrl(url: string): Promise { if (!this.apiKey) { return { @@ -295,50 +209,7 @@ export class FirecrawlClient { return [await this.scrapeUrl(urls[0])]; } - try { - logger.debug(`🔥 [Firecrawl] Batch scraping ${urls.length} URLs`); - - const response = await fetch(`${this.baseUrl}/batch/scrape`, { - method: "POST", - headers: this.getHeaders(), - body: JSON.stringify({ - urls, - ...DEFAULT_SCRAPE_OPTIONS, - }), - }); - - if (!response.ok) { - const error = await this.parseErrorResponse(response); - return urls.map((url) => ({ - url, - title: url, - content: "", - success: false, - error, - })); - } - - const result = await response.json(); - - if (Array.isArray(result?.data)) { - return this.normalizeBatchResults(urls, result.data); - } - - if (typeof result?.id === "string") { - return await this.pollBatchScrape(result.id, urls); - } - - throw new Error("Unexpected Firecrawl batch scrape response"); - } catch (error) { - logger.error("❌ [Firecrawl] Error batch scraping URLs:", error); - const message = error instanceof Error ? error.message : String(error); - return urls.map((url) => ({ - url, - title: url, - content: "", - success: false, - error: message, - })); - } + logger.debug(`🔥 [Firecrawl] Scraping ${urls.length} URLs in parallel`); + return await Promise.all(urls.map((url) => this.scrapeUrl(url))); } } From 45eead0de3a5b3a6e2186987e837e2988a1849af Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 3 Apr 2026 16:34:37 -0400 Subject: [PATCH 4/4] fix: cleanup non-critical compatibility paths Simplify source extraction for create-card-from-message, trim unused URL tool UI types, and keep only the legacy webSearch compatibility needed so persisted threads continue rendering. Made-with: Cursor --- .../assistant-ui/URLContextToolUI.tsx | 7 -- src/hooks/ai/use-create-card-from-message.ts | 100 +++++------------- .../legacy-tool-message-compat.test.ts | 36 +++++++ src/lib/ai/tools/process-urls.ts | 5 - src/lib/ai/web-search-shared.ts | 2 +- 5 files changed, 64 insertions(+), 86 deletions(-) diff --git a/src/components/assistant-ui/URLContextToolUI.tsx b/src/components/assistant-ui/URLContextToolUI.tsx index 05b527d8..988ebcb2 100644 --- a/src/components/assistant-ui/URLContextToolUI.tsx +++ b/src/components/assistant-ui/URLContextToolUI.tsx @@ -204,18 +204,12 @@ type URLMetadata = { urlRetrievalStatus?: string; }; -type SourceMetadata = { - uri?: string; - title?: string; -}; - type ProcessUrlsResult = | string | { text: string; metadata?: { urlMetadata?: URLMetadata[] | null; - sources?: SourceMetadata[] | null; }; }; @@ -231,7 +225,6 @@ export const URLContextToolUI = makeAssistantToolUI<{ const parsedResult = result != null ? parseURLContextResult(result) : null; type Meta = { urlMetadata?: URLMetadata[]; - sources?: SourceMetadata[]; }; const metadata = (typeof parsedResult === "object" && parsedResult !== null && "metadata" in parsedResult ? (parsedResult as { metadata?: Meta }).metadata : null) as Meta | null; const urlMetadata = metadata?.urlMetadata ?? null; diff --git a/src/hooks/ai/use-create-card-from-message.ts b/src/hooks/ai/use-create-card-from-message.ts index 8c051e3d..05ca50d2 100644 --- a/src/hooks/ai/use-create-card-from-message.ts +++ b/src/hooks/ai/use-create-card-from-message.ts @@ -28,31 +28,39 @@ export function useCreateCardFromMessage(options: CreateCardOptions = {}) { const currentWorkspaceId = useWorkspaceStore((state) => state.currentWorkspaceId); const queryClient = useQueryClient(); - const extractSourcesFromToolResult = (result: unknown) => { - const parsed = normalizeWebSearchResult(result); - if (!parsed) { + const extractSourcesFromParts = (parts: unknown) => { + if (!Array.isArray(parts)) { return []; } - if (parsed.sources.length > 0) { - return parsed.sources.map((source) => ({ - title: source.title, - url: source.url, - })); - } + return parts.flatMap((part) => { + if (!part || typeof part !== "object" || Array.isArray(part)) { + return []; + } - const chunks = parsed.groundingMetadata?.groundingChunks || []; - const extractedSources: Array<{ title: string; url: string; favicon?: string }> = []; + const toolPart = part as { + type?: string; + state?: string; + output?: unknown; + }; + + if ( + toolPart.type !== "tool-webSearch" || + toolPart.state !== "output-available" + ) { + return []; + } - for (const chunk of chunks) { - const uri = chunk?.web?.uri; - const title = chunk?.web?.title; - if (uri && title) { - extractedSources.push({ title, url: uri }); + const parsed = normalizeWebSearchResult(toolPart.output); + if (!parsed) { + return []; } - } - return extractedSources; + return parsed.sources.map((source) => ({ + title: source.title, + url: source.url, + })); + }); }; const createCard = useCallback(async () => { @@ -104,68 +112,14 @@ export function useCreateCardFromMessage(options: CreateCardOptions = {}) { if (currentIndex !== -1) { const allSources: Array<{ title: string; url: string; favicon?: string }> = []; - // Look backwards from current message up to the last user message - // or a reasonable limit to find the associated tool result for (let i = currentIndex; i >= 0; i--) { const msg = messages[i]; - // Stop if we hit a user message (start of the turn) - // But carefully: sometimes the user message triggers the tool immediately if (msg.role === 'user' && i !== currentIndex) { - // If we found sources, great. If not, maybe check this user message too if it has attachments? - // For now, break here as tool results usually come after user input. break; } - // Check 1: Tool message (role='tool') with webSearch content - if (msg.role === 'tool') { - // Determine if this is a webSearch tool - // Often inferred from content or toolName if available - const content = msg.content; - // Assistant UI uses parts; check for tool-result part - if (Array.isArray(content)) { - for (const part of content) { - if (part.type === 'tool-result' && (part as any).toolName === 'webSearch') { - const result = (part as any).result; - if (result) { - allSources.push(...extractSourcesFromToolResult(result)); - } - } - } - } - // Check simplified Vercel AI SDK structure - if ((msg as any).toolName === 'webSearch' && (msg as any).content) { - allSources.push(...extractSourcesFromToolResult((msg as any).content)); - } - } - - // Check 2: Assistant message with toolInvocations (AI SDK 3.x+ style) - if (msg.role === 'assistant') { - // Check extracted tool calls if stored in helper fields - // Or specialized parts - if (Array.isArray(msg.content)) { - for (const part of msg.content) { - if (part.type === 'tool-call' && (part as any).toolName === 'webSearch') { - // Sometimes the result is attached to the call in the UI state - const result = (part as any).result || (part as any).args; // Result is usually separate - // Actually, in Assistant UI, result might be embedded if completed - if ((part as any).result) { - allSources.push(...extractSourcesFromToolResult((part as any).result)); - } - } - } - } - - // Check 'toolInvocations' property if exposed directly - const invocations = (msg as any).toolInvocations; - if (Array.isArray(invocations)) { - for (const tool of invocations) { - if (tool.toolName === 'webSearch' && tool.state === 'result') { - allSources.push(...extractSourcesFromToolResult(tool.result)); - } - } - } - } + allSources.push(...extractSourcesFromParts((msg as any).parts)); } if (allSources.length > 0) { diff --git a/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts b/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts index be38c7ed..468471c6 100644 --- a/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts +++ b/src/lib/ai/__tests__/legacy-tool-message-compat.test.ts @@ -97,4 +97,40 @@ describe("normalizeLegacyToolMessages", () => { }, }); }); + + it("normalizes legacy webSearch string outputs without sources", () => { + const messages = [ + { + id: "1", + role: "assistant", + parts: [ + { + type: "tool-webSearch", + toolCallId: "call_4", + state: "output-available", + input: { query: "latest AI news" }, + output: JSON.stringify({ + text: "Summary text", + groundingMetadata: { + groundingChunks: [ + { web: { uri: "https://example.com", title: "Example" } }, + ], + }, + }), + }, + ], + }, + ] as UIMessage[]; + + const normalized = normalizeLegacyToolMessages(messages); + const part = normalized[0]?.parts[0]; + + expect(part).toMatchObject({ + type: "tool-webSearch", + output: { + text: "Summary text", + sources: [], + }, + }); + }); }); diff --git a/src/lib/ai/tools/process-urls.ts b/src/lib/ai/tools/process-urls.ts index 892a778f..a3f40c79 100644 --- a/src/lib/ai/tools/process-urls.ts +++ b/src/lib/ai/tools/process-urls.ts @@ -17,11 +17,6 @@ function truncateContent(content: string): string { return `${content.slice(0, MAX_CONTENT_CHARS_PER_URL).trim()}\n\n[Content truncated for length]`; } -/** - * Create the processUrls tool for analyzing web pages - */ - - /** * Create the processUrls tool for analyzing web pages */ diff --git a/src/lib/ai/web-search-shared.ts b/src/lib/ai/web-search-shared.ts index aca02643..36ade4e2 100644 --- a/src/lib/ai/web-search-shared.ts +++ b/src/lib/ai/web-search-shared.ts @@ -25,7 +25,7 @@ export const WebSearchGroundingMetadataSchema = z export const WebSearchResultSchema = z.object({ text: z.string(), - sources: z.array(WebSearchSourceSchema), + sources: z.array(WebSearchSourceSchema).optional().default([]), groundingMetadata: WebSearchGroundingMetadataSchema.optional(), });