diff --git a/backend/src/abort-registry.ts b/backend/src/abort-registry.ts index 3911ea3..2f3e520 100644 --- a/backend/src/abort-registry.ts +++ b/backend/src/abort-registry.ts @@ -3,8 +3,8 @@ * * Allows the /stop HTTP route to cancel an in-flight populate or update * workflow. The AbortSignal is retrieved inside Mastra workflow steps - * (which receive no signal parameter from the framework) via `getSignal()` - * and passed explicitly to each `agent.generate()` call. + * and dataset-scoped web tools via `getSignal()`, then passed explicitly + * to each `agent.generate()` and TinyFish request. * * Keyed by datasetId because: * - The /stop route knows the datasetId (from the request body). diff --git a/backend/src/mastra/agents/investigate.ts b/backend/src/mastra/agents/investigate.ts index 63a7b8c..6ab8a77 100644 --- a/backend/src/mastra/agents/investigate.ts +++ b/backend/src/mastra/agents/investigate.ts @@ -1,7 +1,7 @@ import { Agent } from "@mastra/core/agent"; import { createOpenRouter } from "@openrouter/ai-sdk-provider"; import { buildPopulateTools } from "../tools/dataset-tools.js"; -import { searchWebTool, fetchPageTool } from "../tools/web-tools.js"; +import { buildWebTools } from "../tools/web-tools.js"; import type { AuthContext } from "../workflows/populate.js"; import type { PopulateColumn } from "../../pipeline/populate.js"; @@ -67,6 +67,7 @@ export function buildInvestigateAgent( authorizedDatasetId, authContext, ); + const { searchWebTool, fetchPageTool } = buildWebTools(authorizedDatasetId); return new Agent({ id: "investigate-agent", name: "Dataset Investigate Agent", diff --git a/backend/src/mastra/agents/populate.ts b/backend/src/mastra/agents/populate.ts index eb9b26f..70a6806 100644 --- a/backend/src/mastra/agents/populate.ts +++ b/backend/src/mastra/agents/populate.ts @@ -1,7 +1,7 @@ import { Agent } from "@mastra/core/agent"; import { createOpenRouter } from "@openrouter/ai-sdk-provider"; import { buildSubagentTool } from "../tools/investigate-tool.js"; -import { searchWebTool, fetchPageTool } from "../tools/web-tools.js"; +import { buildWebTools } from "../tools/web-tools.js"; import type { AuthContext } from "../workflows/populate.js"; import type { PopulateColumn } from "../../pipeline/populate.js"; import type { RunMetrics } from "../run-metrics.js"; @@ -49,6 +49,7 @@ export function buildPopulateAgent( apiKey: openRouterApiKey, baseURL: process.env.OPENROUTER_BASE_URL, }); + const { searchWebTool, fetchPageTool } = buildWebTools(authorizedDatasetId); return new Agent({ id: "populate-agent", diff --git a/backend/src/mastra/agents/refresh.ts b/backend/src/mastra/agents/refresh.ts index 144065a..add9331 100644 --- a/backend/src/mastra/agents/refresh.ts +++ b/backend/src/mastra/agents/refresh.ts @@ -1,7 +1,7 @@ import { Agent } from "@mastra/core/agent"; import { createOpenRouter } from "@openrouter/ai-sdk-provider"; import { buildPopulateTools } from "../tools/dataset-tools.js"; -import { searchWebTool, fetchPageTool } from "../tools/web-tools.js"; +import { buildWebTools } from "../tools/web-tools.js"; import type { AuthContext } from "../workflows/populate.js"; import type { PopulateColumn } from "../../pipeline/populate.js"; @@ -62,6 +62,7 @@ export function buildRefreshAgent( authorizedDatasetId, authContext, ); + const { searchWebTool, fetchPageTool } = buildWebTools(authorizedDatasetId); return new Agent({ id: "refresh-agent", name: "Dataset Refresh Agent", diff --git a/backend/src/mastra/tools/web-tools.ts b/backend/src/mastra/tools/web-tools.ts index 245ee43..e6ad5e3 100644 --- a/backend/src/mastra/tools/web-tools.ts +++ b/backend/src/mastra/tools/web-tools.ts @@ -1,15 +1,36 @@ import { createTool } from "@mastra/core/tools"; import { z } from "zod"; +import { getSignal } from "../../abort-registry.js"; import { FETCH_TIMEOUT_MS } from "../../fetch-timeout.js"; import { getTinyFishApiKey, tinyFishHeaders } from "../../local-credentials.js"; +function requestSignal(datasetId: string) { + const workflowSignal = getSignal(datasetId); + rethrowIfStopped(workflowSignal); + const timeoutSignal = AbortSignal.timeout(FETCH_TIMEOUT_MS); + return { + workflowSignal, + timeoutSignal, + signal: workflowSignal + ? AbortSignal.any([workflowSignal, timeoutSignal]) + : timeoutSignal, + }; +} + +function rethrowIfStopped(signal: AbortSignal | undefined): void { + if (!signal?.aborted) return; + throw signal.reason instanceof Error + ? signal.reason + : new DOMException("Dataset run stopped.", "AbortError"); +} + const searchResultSchema = z.object({ title: z.string(), snippet: z.string(), url: z.string(), }); -export const searchWebTool = createTool({ +const buildSearchWebTool = (datasetId: string) => createTool({ id: "search_web", description: 'Search the web for information. Returns a list of results with titles, snippets, and URLs. Call with: {"query": "your search terms"}', @@ -31,14 +52,12 @@ export const searchWebTool = createTool({ const url = `https://api.search.tinyfish.ai?query=${encodeURIComponent(query)}`; console.log(`[search_web] Searching: "${query}"`); - const controller = new AbortController(); - const timeout = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS); + const signals = requestSignal(datasetId); try { const res = await fetch(url, { headers: tinyFishHeaders(apiKey), - signal: controller.signal, + signal: signals.signal, }); - clearTimeout(timeout); if (!res.ok) { const body = await res.text(); @@ -62,8 +81,8 @@ export const searchWebTool = createTool({ return { results: [], error: "No results found for this query. Try a broader search or use synthetic data." }; return { results }; } catch (err) { - clearTimeout(timeout); - if (err instanceof Error && err.name === "AbortError") + rethrowIfStopped(signals.workflowSignal); + if (signals.timeoutSignal.aborted) return { error: "Search timed out. Skip web search and use synthetic data." }; const msg = err instanceof Error ? err.message : String(err); console.error(`[search_web] Failed:`, msg); @@ -72,7 +91,7 @@ export const searchWebTool = createTool({ }, }); -export const fetchPageTool = createTool({ +const buildFetchPageTool = (datasetId: string) => createTool({ id: "fetch_page", description: 'Fetch a web page and extract its content as clean markdown text. Call with: {"url": "https://example.com/page"}', @@ -96,8 +115,7 @@ export const fetchPageTool = createTool({ console.log(`[fetch_page] Fetching: ${targetUrl}`); - const controller = new AbortController(); - const timeout = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS); + const signals = requestSignal(datasetId); try { const res = await fetch("https://api.fetch.tinyfish.ai", { method: "POST", @@ -106,9 +124,8 @@ export const fetchPageTool = createTool({ ...tinyFishHeaders(apiKey), }, body: JSON.stringify({ urls: [targetUrl], format: "markdown" }), - signal: controller.signal, + signal: signals.signal, }); - clearTimeout(timeout); if (!res.ok) { const body = await res.text(); @@ -151,8 +168,8 @@ export const fetchPageTool = createTool({ text, }; } catch (err) { - clearTimeout(timeout); - if (err instanceof Error && err.name === "AbortError") + rethrowIfStopped(signals.workflowSignal); + if (signals.timeoutSignal.aborted) return { error: "Page fetch timed out. Try a different URL or use search snippet data." }; const msg = err instanceof Error ? err.message : String(err); console.error(`[fetch_page] Failed:`, msg); @@ -160,3 +177,10 @@ export const fetchPageTool = createTool({ } }, }); + +export function buildWebTools(datasetId: string) { + return { + searchWebTool: buildSearchWebTool(datasetId), + fetchPageTool: buildFetchPageTool(datasetId), + }; +}