Skip to content

Stop active TinyFish requests when dataset runs are cancelled #151

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dunkeln wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Stop active TinyFish requests when dataset runs are cancelled #151

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions backend/src/abort-registry.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
3 changes: 2 additions & 1 deletion backend/src/mastra/agents/investigate.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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";

Expand Down Expand Up @@ -67,6 +67,7 @@ export function buildInvestigateAgent(
authorizedDatasetId,
authContext,
);
const { searchWebTool, fetchPageTool } = buildWebTools(authorizedDatasetId);
return new Agent({
id: "investigate-agent",
name: "Dataset Investigate Agent",
Expand Down
3 changes: 2 additions & 1 deletion backend/src/mastra/agents/populate.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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",
Expand Down
3 changes: 2 additions & 1 deletion backend/src/mastra/agents/refresh.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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";

Expand Down Expand Up @@ -62,6 +62,7 @@ export function buildRefreshAgent(
authorizedDatasetId,
authContext,
);
const { searchWebTool, fetchPageTool } = buildWebTools(authorizedDatasetId);
return new Agent({
id: "refresh-agent",
name: "Dataset Refresh Agent",
Expand Down
52 changes: 38 additions & 14 deletions backend/src/mastra/tools/web-tools.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"}',
Expand All @@ -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();
Expand All @@ -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);
Expand All @@ -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"}',
Expand All @@ -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",
Expand All @@ -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();
Expand Down Expand Up @@ -151,12 +168,19 @@ 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);
return { error: `Fetch failed: ${msg}. Use data from search snippets instead.` };
}
},
});

export function buildWebTools(datasetId: string) {
return {
searchWebTool: buildSearchWebTool(datasetId),
fetchPageTool: buildFetchPageTool(datasetId),
};
Comment on lines +181 to +185

@coderabbitai coderabbitai Bot Jul 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟠 Major | ⚡ Quick win

Capture authContext in the web-tool closures.

buildWebTools() accepts only a dataset ID, so these tools cannot perform the required user/run audit logging or CAPABILITY_VIOLATION tracking.

  • backend/src/mastra/tools/web-tools.ts#L181-L185: accept authContext and pass it into both tool factories so each execution closes over it.
  • backend/src/mastra/agents/investigate.ts#L70-L70: pass the existing authContext to buildWebTools.
  • backend/src/mastra/agents/populate.ts#L52-L52: pass the existing authContext to buildWebTools.
  • backend/src/mastra/agents/refresh.ts#L65-L65: pass the existing authContext to buildWebTools.

As per coding guidelines, “Capture authContext (Clerk userId + workflow run id) in tools for security audit logging and CAPABILITY_VIOLATION PostHog event tracking.”

📍 Affects 4 files
  • backend/src/mastra/tools/web-tools.ts#L181-L185 (this comment)
  • backend/src/mastra/agents/investigate.ts#L70-L70
  • backend/src/mastra/agents/populate.ts#L52-L52
  • backend/src/mastra/agents/refresh.ts#L65-L65
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/src/mastra/tools/web-tools.ts` around lines 181 - 185, Update
buildWebTools in backend/src/mastra/tools/web-tools.ts to accept authContext and
pass it to both buildSearchWebTool and buildFetchPageTool so each closure can
perform audit logging and CAPABILITY_VIOLATION tracking. Update the existing
buildWebTools calls in backend/src/mastra/agents/investigate.ts (line 70),
backend/src/mastra/agents/populate.ts (line 52), and
backend/src/mastra/agents/refresh.ts (line 65) to pass their existing
authContext.

Source: Coding guidelines

}