diff --git a/src/mastra/tools/alpha-vantage.tool.ts b/src/mastra/tools/alpha-vantage.tool.ts index dfaea24a..9cd64921 100644 --- a/src/mastra/tools/alpha-vantage.tool.ts +++ b/src/mastra/tools/alpha-vantage.tool.ts @@ -1,9 +1,11 @@ -import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' +import type { InferUITool } from '@mastra/core/tools' import { trace } from '@opentelemetry/api' import { z } from 'zod' import { log } from '../config/logger' +import type { RequestContext } from '@mastra/core/request-context' + /** * Alpha Vantage Tools * @@ -14,6 +16,13 @@ import { log } from '../config/logger' * Requires ALPHA_VANTAGE_API_KEY environment variable */ +/** + * Base Request Context for Alpha Vantage + */ +export interface AlphaVantageRequestContext extends RequestContext { + apiKey?: string +} + // In-memory counter to track tool calls per request // Add this line at the beginning of each tool's execute function to track usage: // toolCallCounters.set('tool-id', (toolCallCounters.get('tool-id') ?? 0) + 1) @@ -83,6 +92,7 @@ export const alphaVantageCryptoTool = createTool({ .optional(), }), execute: async (inputData, context) => { + const requestContext = context?.requestContext as AlphaVantageRequestContext const span = trace .getTracer('alpha-vantage-crypto-tool', '1.0.0') .startSpan('alpha-vantage-crypto', { @@ -103,7 +113,7 @@ export const alphaVantageCryptoTool = createTool({ }, id: 'alpha-vantage-crypto', }) - const apiKey = process.env.ALPHA_VANTAGE_API_KEY + const apiKey = requestContext?.apiKey ?? process.env.ALPHA_VANTAGE_API_KEY if (typeof apiKey !== 'string' || apiKey.trim() === '') { await context?.writer?.custom({ @@ -313,7 +323,7 @@ export const alphaVantageCryptoTool = createTool({ }) }, onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - const dataKeys = output.data ? Object.keys(output.data).length : 0 + const dataKeys = (output.data !== null && typeof output.data === 'object') ? Object.keys(output.data as object).length : 0 log.info('Alpha Vantage crypto completed', { toolCallId, toolName, @@ -412,6 +422,7 @@ export const alphaVantageStockTool = createTool({ }), execute: async (inputData, context) => { + const requestContext = context?.requestContext as AlphaVantageRequestContext const span = trace .getTracer('alpha-vantage-stock-tool', '1.0.0') .startSpan('alpha-vantage-stock', { @@ -431,7 +442,7 @@ export const alphaVantageStockTool = createTool({ }, id: 'alpha-vantage-stock', }) - const apiKey = process.env.ALPHA_VANTAGE_API_KEY + const apiKey = requestContext?.apiKey ?? process.env.ALPHA_VANTAGE_API_KEY if (typeof apiKey !== 'string' || apiKey.trim() === '') { await context?.writer?.custom({ @@ -674,7 +685,7 @@ export const alphaVantageStockTool = createTool({ }) }, onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - const dataKeys = output.data ? Object.keys(output.data).length : 0 + const dataKeys = (output.data !== null && typeof output.data === 'object') ? Object.keys(output.data as object).length : 0 log.info('Alpha Vantage stock completed', { toolCallId, toolName, @@ -804,6 +815,7 @@ export const alphaVantageTool = createTool({ .optional(), }), execute: async (inputData, context) => { + const requestContext = context?.requestContext as AlphaVantageRequestContext const span = trace .getTracer('alpha-vantage-tool', '1.0.0') .startSpan('alpha-vantage', { @@ -823,7 +835,7 @@ export const alphaVantageTool = createTool({ }, id: 'alpha-vantage', }) - const apiKey = process.env.ALPHA_VANTAGE_API_KEY + const apiKey = requestContext?.apiKey ?? process.env.ALPHA_VANTAGE_API_KEY if (typeof apiKey !== 'string' || apiKey.trim() === '') { await context?.writer?.custom({ diff --git a/src/mastra/tools/arxiv.tool.ts b/src/mastra/tools/arxiv.tool.ts index fca57228..59fbb7aa 100644 --- a/src/mastra/tools/arxiv.tool.ts +++ b/src/mastra/tools/arxiv.tool.ts @@ -4,6 +4,12 @@ import { trace } from '@opentelemetry/api' import { z } from 'zod' import { log } from '../config/logger' +import type { RequestContext } from '@mastra/core/request-context' + +export interface ArxivRequestContext extends RequestContext { + userId?: string +} + // In-memory counter to track tool calls per request // Add this line at the beginning of each tool's execute function to track usage: // toolCallCounters.set('tool-id', (toolCallCounters.get('tool-id') ?? 0) + 1) @@ -219,60 +225,23 @@ export const arxivTool = createTool({ start_index: z.number(), max_results: z.number(), }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('ArXiv tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('ArXiv tool received input chunk', { - toolCallId, - inputTextDelta, - abortSignal: abortSignal?.aborted, - messageCount: messages.length, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - const searchType = input.id - ? 'by-id' - : input.query - ? 'by-query' - : 'general' - log.info('ArXiv tool received input', { - toolCallId, - messageCount: messages.length, - searchType, - query: input.query, - id: input.id, - author: input.author, - category: input.category, - maxResults: input.max_results, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('ArXiv search completed', { - toolCallId, - toolName, - abortSignal: abortSignal?.aborted, - papersFound: output.papers.length, - totalResults: output.total_results, - hook: 'onOutput', - }) - }, execute: async (inputData, context) => { const abortSignal = context?.abortSignal + const writer = context?.writer + const requestCtx = + context?.requestContext as ArxivRequestContext | undefined // Check if operation was already cancelled if (abortSignal?.aborted === true) { throw new Error('ArXiv search cancelled') } + if (requestCtx) { + log.debug('ArXiv search request context available', { + tool: 'arxiv', + }) + } + const span = trace .getTracer('arxiv-tool', '1.0.0') .startSpan('arxiv-search', { @@ -285,7 +254,7 @@ export const arxivTool = createTool({ }, }) - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -323,7 +292,7 @@ export const arxivTool = createTool({ } else if (searchTerms.length > 0) { params.append('search_query', searchTerms.join(' AND ')) } else { - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', @@ -359,7 +328,7 @@ export const arxivTool = createTool({ const url = `http://export.arxiv.org/api/query?${params.toString()}` // Check for cancellation before API call - if (abortSignal && abortSignal.aborted) { + if (abortSignal?.aborted) { span.setStatus({ code: 2, message: 'Operation cancelled during API call', @@ -368,7 +337,7 @@ export const arxivTool = createTool({ throw new Error('ArXiv search cancelled during API call') } - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -377,7 +346,7 @@ export const arxivTool = createTool({ }, id: 'arxiv', }) - const response = await fetch(url) + const response = await fetch(url, { signal: abortSignal }) if (!response.ok) { throw new Error( @@ -425,7 +394,7 @@ export const arxivTool = createTool({ span.setStatus({ code: 2, message: cancelMessage }) span.end() - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', @@ -451,6 +420,52 @@ export const arxivTool = createTool({ throw error instanceof Error ? error : new Error(errorMessage) } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('ArXiv tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('ArXiv tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + const searchType = input.id + ? 'by-id' + : input.query + ? 'by-query' + : 'general' + log.info('ArXiv tool received input', { + toolCallId, + messageCount: messages.length, + searchType, + query: input.query, + id: input.id, + author: input.author, + category: input.category, + maxResults: input.max_results, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('ArXiv search completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + papersFound: output.papers.length, + totalResults: output.total_results, + hook: 'onOutput', + }) + }, }) export type ArxivUITool = InferUITool @@ -507,6 +522,9 @@ export const arxivPdfParserTool = createTool({ }), }), execute: async (inputData, context) => { + const writer = context?.writer + const requestCtx = + context?.requestContext as ArxivRequestContext | undefined const span = trace .getTracer('arxiv-pdf-parser-tool', '1.0.0') .startSpan('arxiv-pdf-parser', { @@ -517,7 +535,7 @@ export const arxivPdfParserTool = createTool({ }, }) - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -534,20 +552,16 @@ export const arxivPdfParserTool = createTool({ const startTime = Date.now() try { - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '📥 Downloading PDF from arXiv...', - stage: 'arxiv-pdf-parser', - }, - id: 'arxiv-pdf-parser', - }) + if (requestCtx) { + log.debug('Arxiv PDF parser request context available', { + tool: 'arxiv-pdf-parser', + }) + } // Construct PDF URL const pdfUrl = `https://arxiv.org/pdf/${inputData.arxivId}` // Download PDF - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -559,7 +573,7 @@ export const arxivPdfParserTool = createTool({ const response = await fetch(pdfUrl) if (!response.ok) { - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', @@ -658,7 +672,7 @@ export const arxivPdfParserTool = createTool({ markdown = frontmatter + markdown } - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', @@ -695,7 +709,6 @@ export const arxivPdfParserTool = createTool({ span.end() return result } catch (error) { - const processingTime = Date.now() - startTime const errorMessage = error instanceof Error ? error.message @@ -709,6 +722,44 @@ export const arxivPdfParserTool = createTool({ throw error instanceof Error ? error : new Error(errorMessage) } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('ArXiv PDF parser input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('ArXiv PDF parser received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('ArXiv PDF parser received input', { + toolCallId, + messageCount: messages.length, + arxivId: input.arxivId, + maxPages: input.maxPages, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('ArXiv PDF parsing completed', { + toolCallId, + toolName, + arxivId: output.arxivId, + success: output.success, + pageCount: output.statistics.pageCount, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export type ArxivPdfParserUITool = InferUITool @@ -772,6 +823,10 @@ export const arxivPaperDownloaderTool = createTool({ .optional(), }), execute: async (inputData, context) => { + const writer = context?.writer + const abortSignal = context?.abortSignal + const requestCtx = + context?.requestContext as ArxivRequestContext | undefined const span = trace .getTracer('arxiv-paper-downloader-tool', '1.0.0') .startSpan('arxiv-paper-downloader', { @@ -783,7 +838,7 @@ export const arxivPaperDownloaderTool = createTool({ }, }) - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -799,7 +854,12 @@ export const arxivPaperDownloaderTool = createTool({ (toolCallCounters.get('arxiv-paper-downloader') ?? 0) + 1 ) try { - await context?.writer?.custom({ + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing arXiv paper downloader for user', { + userId: requestCtx.userId, + }) + } + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -810,7 +870,7 @@ export const arxivPaperDownloaderTool = createTool({ }) // Get metadata from arXiv API const apiUrl = `http://export.arxiv.org/api/query?id_list=${inputData.arxivId}&max_results=1` - const apiResponse = await fetch(apiUrl) + const apiResponse = await fetch(apiUrl, { signal: abortSignal }) if (!apiResponse.ok) { throw new Error( @@ -832,7 +892,7 @@ export const arxivPaperDownloaderTool = createTool({ | { markdown: string; pageCount: number; textLength: number } | undefined - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -843,7 +903,7 @@ export const arxivPaperDownloaderTool = createTool({ }) // Download and parse PDF if requested if (inputData.includePdfContent) { - await context?.writer?.custom({ + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', @@ -931,6 +991,45 @@ export const arxivPaperDownloaderTool = createTool({ throw error instanceof Error ? error : new Error(errorMessage) } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('ArXiv paper downloader input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('ArXiv paper downloader received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('ArXiv paper downloader received input', { + toolCallId, + messageCount: messages.length, + arxivId: input.arxivId, + includePdf: input.includePdfContent, + format: input.format, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('ArXiv paper download completed', { + toolCallId, + toolName, + arxivId: output.arxivId, + success: output.success, + hasPdf: !!output.pdfContent, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export type ArxivPaperDownloaderUITool = InferUITool< diff --git a/src/mastra/tools/browser-tool.ts b/src/mastra/tools/browser-tool.ts index fc000b25..1c547bca 100644 --- a/src/mastra/tools/browser-tool.ts +++ b/src/mastra/tools/browser-tool.ts @@ -1,11 +1,15 @@ -import type { InferUITool } from '@mastra/core/tools' -import { createTool } from '@mastra/core/tools' +import { createTool, type InferUITool } from '@mastra/core/tools' import { MDocument } from '@mastra/rag' import { trace } from '@opentelemetry/api' import type { Browser } from 'playwright-core' import { chromium } from 'playwright-core' import { z } from 'zod' import { log } from '../config/logger' +import type { RequestContext } from '@mastra/core/request-context' +export interface BrowserRequestContext extends RequestContext { + userId?: string; + sessionId?: string; +} // Browser instance cache for reuse let browserInstance: Browser | null = null @@ -33,47 +37,9 @@ export const browserTool = createTool({ outputSchema: z.object({ message: z.string(), }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Browser tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Browser tool received input chunk', { - toolCallId, - inputTextDelta, - abortSignal: abortSignal?.aborted, - messageCount: messages.length, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Browser tool received input', { - toolCallId, - messageCount: messages.length, - url: input.url, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - const isSuccess = - !output.message.includes('Error:') && - !output.message.includes('Failed') - log[isSuccess ? 'info' : 'warn']('Browser tool completed', { - toolCallId, - toolName, - success: isSuccess, - message: output.message.substring(0, 100), - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, execute: async (inputData, context) => { const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined // Check if operation was already cancelled if (abortSignal?.aborted === true) { @@ -98,14 +64,10 @@ export const browserTool = createTool({ id: 'browserTool', }) try { - // Check for cancellation before browser operations - if (abortSignal && abortSignal.aborted) { - span.setStatus({ - code: 2, - message: 'Operation cancelled during browser operations', + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing browser scrape for user', { + userId: requestCtx.userId, }) - span.end() - throw new Error('Browser operation cancelled during setup') } const browser = await getBrowser() @@ -186,6 +148,45 @@ export const browserTool = createTool({ return { message: `Error: ${errorMsg}` } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Browser tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Browser tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Browser tool received input', { + toolCallId, + messageCount: messages.length, + url: input.url, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + const isSuccess = + !output.message.includes('Error:') && + !output.message.includes('Failed') + log[isSuccess ? 'info' : 'warn']('Browser tool completed', { + toolCallId, + toolName, + success: isSuccess, + message: output.message.substring(0, 100), + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const screenshotTool = createTool({ @@ -208,6 +209,8 @@ export const screenshotTool = createTool({ message: z.string().optional(), }), execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined const tracer = trace.getTracer('screenshot-tool', '1.0.0') const span = tracer.startSpan('browser-screenshot', { attributes: { @@ -223,6 +226,14 @@ export const screenshotTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing browser screenshot for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Screenshot operation cancelled') + } const browser = await getBrowser() const page = await browser.newPage() await page.setViewportSize({ @@ -261,6 +272,41 @@ export const screenshotTool = createTool({ return { success: false, message: errorMsg } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Screenshot tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Screenshot tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Screenshot tool received input', { + toolCallId, + messageCount: messages.length, + url: input.url, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log[output.success ? 'info' : 'warn']('Screenshot tool completed', { + toolCallId, + toolName, + success: output.success, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const pdfGeneratorTool = createTool({ @@ -278,6 +324,8 @@ export const pdfGeneratorTool = createTool({ message: z.string().optional(), }), execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined const tracer = trace.getTracer('pdf-generator-tool', '1.0.0') const span = tracer.startSpan('browser-pdf-generate', { attributes: { @@ -293,6 +341,14 @@ export const pdfGeneratorTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing PDF generation for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('PDF generation operation cancelled') + } const browser = await getBrowser() const page = await browser.newPage() await page.goto(inputData.url, { waitUntil: 'networkidle' }) @@ -328,6 +384,41 @@ export const pdfGeneratorTool = createTool({ return { success: false, message: errorMsg } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('PDF generator tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('PDF generator tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('PDF generator tool received input', { + toolCallId, + messageCount: messages.length, + url: input.url, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log[output.success ? 'info' : 'warn']('PDF generator tool completed', { + toolCallId, + toolName, + success: output.success, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const clickAndExtractTool = createTool({ @@ -362,6 +453,8 @@ export const clickAndExtractTool = createTool({ message: z.string().optional(), }), execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined const tracer = trace.getTracer('click-extract-tool', '1.0.0') const span = tracer.startSpan('browser-click-extract', { attributes: { @@ -377,6 +470,14 @@ export const clickAndExtractTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing click and extract for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Click and extract operation cancelled') + } const browser = await getBrowser() const page = await browser.newPage() await page.goto(inputData.url, { waitUntil: 'domcontentloaded' }) @@ -439,6 +540,41 @@ export const clickAndExtractTool = createTool({ return { success: false, message: errorMsg } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Click and extract tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Click and extract tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Click and extract tool received input', { + toolCallId, + messageCount: messages.length, + url: input.url, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log[output.success ? 'info' : 'warn']('Click and extract tool completed', { + toolCallId, + toolName, + success: output.success, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const fillFormTool = createTool({ @@ -468,6 +604,8 @@ export const fillFormTool = createTool({ message: z.string().optional(), }), execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined const tracer = trace.getTracer('fill-form-tool', '1.0.0') const span = tracer.startSpan('browser-fill-form', { attributes: { @@ -483,6 +621,14 @@ export const fillFormTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing fill form for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Fill form operation cancelled') + } const browser = await getBrowser() const page = await browser.newPage() await page.goto(inputData.url, { waitUntil: 'domcontentloaded' }) @@ -534,6 +680,41 @@ export const fillFormTool = createTool({ return { success: false, message: errorMsg } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Fill form tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Fill form tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Fill form tool received input', { + toolCallId, + messageCount: messages.length, + url: input.url, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log[output.success ? 'info' : 'warn']('Fill form tool completed', { + toolCallId, + toolName, + success: output.success, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const googleSearch = createTool({ @@ -547,6 +728,8 @@ export const googleSearch = createTool({ message: z.string(), }), execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined const tracer = trace.getTracer('google-search-tool', '1.0.0') const span = tracer.startSpan('google-search', { attributes: { @@ -564,6 +747,14 @@ export const googleSearch = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing Google search for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Google search operation cancelled') + } const browser = await getBrowser() const page = await browser.newPage() await page.goto( @@ -635,6 +826,44 @@ export const googleSearch = createTool({ return { message: `Error: ${errorMsg}` } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Google search tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Google search tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Google search tool received input', { + toolCallId, + messageCount: messages.length, + query: input.query, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + const isSuccess = + !output.message.includes('Error:') && + !output.message.includes('Failed') + log[isSuccess ? 'info' : 'warn']('Google search tool completed', { + toolCallId, + toolName, + success: isSuccess, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const extractTablesTool = createTool({ @@ -663,6 +892,8 @@ export const extractTablesTool = createTool({ message: z.string().optional(), }), execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined const tracer = trace.getTracer('extract-tables-tool', '1.0.0') const span = tracer.startSpan('browser-extract-tables', { attributes: { @@ -677,6 +908,14 @@ export const extractTablesTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing table extraction for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Table extraction operation cancelled') + } const browser = await getBrowser() const page = await browser.newPage() await page.goto(inputData.url, { waitUntil: 'domcontentloaded' }) @@ -741,6 +980,41 @@ export const extractTablesTool = createTool({ return { success: false, message: errorMsg } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Extract tables tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Extract tables tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Extract tables tool received input', { + toolCallId, + messageCount: messages.length, + url: input.url, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log[output.success ? 'info' : 'warn']('Extract tables tool completed', { + toolCallId, + toolName, + success: output.success, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const monitorPageTool = createTool({ @@ -773,6 +1047,8 @@ export const monitorPageTool = createTool({ message: z.string().optional(), }), execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as BrowserRequestContext | undefined const tracer = trace.getTracer('monitor-page-tool', '1.0.0') const span = tracer.startSpan('browser-monitor-page', { attributes: { @@ -788,6 +1064,14 @@ export const monitorPageTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing page monitor for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Page monitor operation cancelled') + } const browser = await getBrowser() const page = await browser.newPage() const selector = inputData.selector ?? 'body' @@ -856,6 +1140,41 @@ export const monitorPageTool = createTool({ } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Monitor page tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Monitor page tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Monitor page tool received input', { + toolCallId, + messageCount: messages.length, + url: input.url, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log[output.success ? 'info' : 'warn']('Monitor page tool completed', { + toolCallId, + toolName, + success: output.success, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export type BrowserToolUITool = InferUITool diff --git a/src/mastra/tools/calculator.tool.ts b/src/mastra/tools/calculator.tool.ts index bcd07069..6c5939e9 100644 --- a/src/mastra/tools/calculator.tool.ts +++ b/src/mastra/tools/calculator.tool.ts @@ -8,6 +8,7 @@ import type { RequestContext } from '@mastra/core/request-context' // RequestContext interface for Calculator tool export interface CalculatorToolContext extends RequestContext { + userId?: string precision?: number allowComplexExpressions?: boolean maxExpressionLength?: number @@ -476,6 +477,8 @@ export const unitConverterTool = createTool({ }), execute: async (inputData, context) => { const writer = context?.writer + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as CalculatorToolContext | undefined const tracer = trace.getTracer('unit-converter-tool', '1.0.0') const span = tracer.startSpan('convert-units') @@ -491,6 +494,14 @@ export const unitConverterTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing unit conversion for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Unit conversion operation cancelled') + } const result = convertUnits( inputData.value, inputData.fromUnit, @@ -644,6 +655,8 @@ export const matrixCalculatorTool = createTool({ }), execute: async (inputData, context) => { const writer = context?.writer + const abortSignal = context?.abortSignal + const requestCtx = context?.requestContext as CalculatorToolContext | undefined const tracer = trace.getTracer('matrix-calculator-tool', '1.0.0') const span = tracer.startSpan('matrix-operation') @@ -659,6 +672,14 @@ export const matrixCalculatorTool = createTool({ }) try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing matrix calculation for user', { + userId: requestCtx.userId, + }) + } + if (abortSignal?.aborted === true) { + throw new Error('Matrix calculation operation cancelled') + } let result: number[][] | number switch (inputData.operation) { diff --git a/src/mastra/tools/calendar-tool.ts b/src/mastra/tools/calendar-tool.ts index c5c536bc..51731c3e 100644 --- a/src/mastra/tools/calendar-tool.ts +++ b/src/mastra/tools/calendar-tool.ts @@ -4,6 +4,12 @@ import { execSync } from 'child_process' import { z } from 'zod' import { log } from '../config/logger' +import type { RequestContext } from '@mastra/core/request-context' + +export interface CalendarRequestContext extends RequestContext { + userId?: string; +} + interface CalendarEvent { title: string startDate: Date @@ -168,6 +174,7 @@ export const listEvents = createTool({ }, execute: async (inputData, context) => { const tracer = trace.getTracer('calendar-tool', '1.0.0') + const requestCtx = context?.requestContext as CalendarRequestContext | undefined const span = tracer.startSpan('list-calendar-events', { attributes: { 'tool.id': 'list-calendar-events', @@ -175,6 +182,10 @@ export const listEvents = createTool({ }, }) + if (requestCtx?.userId) { + log.debug('Executing calendar list events for user', { userId: requestCtx.userId }) + } + await context?.writer?.custom({ type: 'data-tool-progress', data: { @@ -248,12 +259,17 @@ export const getTodayEvents = createTool({ }), execute: async (inputData, context) => { const tracer = trace.getTracer('calendar-tool', '1.0.0') + const requestCtx = context?.requestContext as CalendarRequestContext | undefined const span = tracer.startSpan('get-today-events', { attributes: { 'tool.id': 'get-today-events', }, }) + if (requestCtx?.userId) { + log.debug('Executing get today events for user', { userId: requestCtx.userId }) + } + await context?.writer?.write({ type: 'progress', data: { message: "📅 Getting today's events..." }, @@ -303,6 +319,40 @@ export const getTodayEvents = createTool({ return { events: [], count: 0 } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Today events tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Today events tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ toolCallId, messages, abortSignal }) => { + log.info('Today events tool received input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Today events tool completed', { + toolCallId, + toolName, + eventsFound: output.count, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const getUpcomingEvents = createTool({ @@ -332,6 +382,7 @@ export const getUpcomingEvents = createTool({ }), execute: async (inputData, context) => { const tracer = trace.getTracer('calendar-tool', '1.0.0') + const requestCtx = context?.requestContext as CalendarRequestContext | undefined const span = tracer.startSpan('get-upcoming-events', { attributes: { 'tool.id': 'get-upcoming-events', @@ -340,6 +391,10 @@ export const getUpcomingEvents = createTool({ }, }) + if (requestCtx?.userId) { + log.debug('Executing get upcoming events for user', { userId: requestCtx.userId }) + } + await context?.writer?.write({ type: 'progress', data: { @@ -399,6 +454,42 @@ export const getUpcomingEvents = createTool({ return { events: [], count: 0 } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Upcoming events tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Upcoming events tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Upcoming events tool received input', { + toolCallId, + messageCount: messages.length, + days: input.days, + limit: input.limit, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Upcoming events tool completed', { + toolCallId, + toolName, + eventsFound: output.count, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const findFreeSlots = createTool({ @@ -440,6 +531,7 @@ export const findFreeSlots = createTool({ }), execute: async (inputData, context) => { const tracer = trace.getTracer('calendar-tool', '1.0.0') + const requestCtx = context?.requestContext as CalendarRequestContext | undefined const span = tracer.startSpan('find-free-slots', { attributes: { 'tool.id': 'find-free-slots', @@ -447,6 +539,10 @@ export const findFreeSlots = createTool({ }, }) + if (requestCtx?.userId) { + log.debug('Executing find free slots for user', { userId: requestCtx.userId }) + } + await context?.writer?.write({ type: 'progress', data: { message: '📅 Finding free time slots...' }, @@ -548,4 +644,41 @@ export const findFreeSlots = createTool({ return { freeSlots: [], busyPeriods: [] } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Find free slots tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Find free slots tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Find free slots tool received input', { + toolCallId, + messageCount: messages.length, + date: input.date, + workdayStart: input.workdayStart, + workdayEnd: input.workdayEnd, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Find free slots tool completed', { + toolCallId, + toolName, + slotsFound: output.freeSlots.length, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) diff --git a/src/mastra/tools/code-analysis.tool.ts b/src/mastra/tools/code-analysis.tool.ts index 67dbff5d..820e3d54 100644 --- a/src/mastra/tools/code-analysis.tool.ts +++ b/src/mastra/tools/code-analysis.tool.ts @@ -8,7 +8,11 @@ import stripComments from 'strip-comments' import { trace, SpanStatusCode } from '@opentelemetry/api' import { PythonParser } from './semantic-utils' import { log } from '../config/logger' -// RequestContext typing not used to avoid incompatibility with ToolExecutionContext +import type { RequestContext } from '@mastra/core/request-context' + +export interface CodeAnalysisRequestContext extends RequestContext { + userId?: string; +} async function fileExists(filePath: string): Promise { try { @@ -404,45 +408,10 @@ Supports TypeScript, JavaScript, Python, and other languages. Use for code review preparation, quality assessment, and refactoring planning.`, inputSchema: codeAnalysisInputSchema, outputSchema: codeAnalysisOutputSchema, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Code analysis tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Code analysis tool received input chunk', { - toolCallId, - inputTextDelta, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Code analysis tool received input', { - toolCallId, - messageCount: messages.length, - inputData: { target: input.target, options: input.options }, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Code analysis tool completed', { - toolCallId, - toolName, - outputData: { - totalFiles: output.summary.totalFiles, - totalLoc: output.summary.totalLoc, - }, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, execute: async (inputData, context): Promise => { const writer = context?.writer + const requestCtx = context?.requestContext as CodeAnalysisRequestContext | undefined + await writer?.custom({ type: 'data-tool-progress', data: { @@ -452,6 +421,11 @@ Use for code review preparation, quality assessment, and refactoring planning.`, }, id: 'coding:codeAnalysis', }) + + if (requestCtx?.userId) { + log.debug('Executing code analysis for user', { userId: requestCtx.userId }) + } + const tracer = trace.getTracer('code-analysis-tool', '1.0.0') const span = tracer.startSpan('code-analysis', { attributes: { @@ -616,4 +590,42 @@ Use for code review preparation, quality assessment, and refactoring planning.`, throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Code analysis tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Code analysis tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Code analysis tool received input', { + toolCallId, + messageCount: messages.length, + inputData: { target: input.target, options: input.options }, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Code analysis tool completed', { + toolCallId, + toolName, + outputData: { + totalFiles: output.summary.totalFiles, + totalLoc: output.summary.totalLoc, + }, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) diff --git a/src/mastra/tools/code-chunking.ts b/src/mastra/tools/code-chunking.ts index 1d970f19..462d175b 100644 --- a/src/mastra/tools/code-chunking.ts +++ b/src/mastra/tools/code-chunking.ts @@ -1,11 +1,18 @@ import { createTool } from '@mastra/core/tools' import { z } from 'zod' import { Project, SyntaxKind } from 'ts-morph' +import type { RequestContext } from '@mastra/core/request-context' + import { PythonParser } from './semantic-utils' import { MDocument } from '@mastra/rag' import { log } from '../config/logger' import { trace, SpanStatusCode } from '@opentelemetry/api' +export interface CodeChunkingRequestContext extends RequestContext { + chunkSize?: number + chunkOverlap?: number +} + const codeChunkingInputSchema = z.object({ filePath: z.string(), content: z.string(), @@ -37,43 +44,9 @@ export const codeChunkerTool = createTool({ 'Intelligently chunks code files based on syntax (functions, classes) for TS/JS and Python, falling back to recursive chunking for others.', inputSchema: codeChunkingInputSchema, outputSchema: codeChunkingOutputSchema, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Code chunker tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Code chunker received input chunk', { - toolCallId, - inputTextDelta, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Code chunker received input', { - toolCallId, - messageCount: messages.length, - inputData: { filePath: input.filePath }, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Code chunker completed', { - toolCallId, - toolName, - chunksCreated: output.chunks.length, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, - execute: async (inputData) => { - const { filePath, content, options } = inputData + execute: async (input: z.infer, context) => { + const { filePath, content, options } = input + const requestContext = context?.requestContext as CodeChunkingRequestContext | undefined const ext = filePath.split('.').pop()?.toLowerCase() interface Chunk { text: string @@ -193,13 +166,13 @@ export const codeChunkerTool = createTool({ const textChunks = await doc.chunk({ strategy: 'recursive', - maxSize: options?.chunkSize ?? 512, - overlap: options?.chunkOverlap ?? 50, + maxSize: options?.chunkSize ?? requestContext?.chunkSize ?? 512, + overlap: options?.chunkOverlap ?? requestContext?.chunkOverlap ?? 50, separators: ['\n\n', '\n', ' '], }) const mapped = textChunks.map((c, i) => ({ - text: c.text || '', + text: c.text ?? '', metadata: { startLine: 0, // MDocument doesn't give line numbers easily endLine: 0, @@ -244,4 +217,39 @@ export const codeChunkerTool = createTool({ } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Code chunker tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Code chunker received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Code chunker received input', { + toolCallId, + messageCount: messages.length, + inputData: { filePath: input.filePath }, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Code chunker completed', { + toolCallId, + toolName, + chunksCreated: output.chunks.length, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) diff --git a/src/mastra/tools/code-search.tool.ts b/src/mastra/tools/code-search.tool.ts index bfca12a1..b231f430 100644 --- a/src/mastra/tools/code-search.tool.ts +++ b/src/mastra/tools/code-search.tool.ts @@ -5,8 +5,17 @@ import { promises as fs } from 'node:fs' import * as path from 'node:path' import RE2 from 're2' import { z } from 'zod' +import type { RequestContext } from '@mastra/core/request-context' import { log } from '../config/logger' +export interface CodeSearchRequestContext extends RequestContext { + caseSensitive?: boolean + maxResults?: number + includeContext?: boolean + contextLines?: number + maxFileSize?: number +} + interface RegexLike { lastIndex?: number exec: RegExp['exec'] @@ -104,52 +113,10 @@ Supports string and regex patterns with context lines. Use for finding usages, identifying patterns, and code exploration.`, inputSchema: codeSearchInputSchema, outputSchema: codeSearchOutputSchema, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Code search tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Code search tool received input chunk', { - toolCallId, - inputTextDelta, - abortSignal: abortSignal?.aborted, - messageCount: messages.length, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - const targetDesc = Array.isArray(input.target) - ? `${input.target.length} targets` - : input.target - log.info('Code search received complete input', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - pattern: input.pattern, - target: targetDesc, - isRegex: input.options?.isRegex ?? false, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Code search completed', { - toolCallId, - toolName, - abortSignal: abortSignal?.aborted, - totalMatches: output.stats.totalMatches, - filesWithMatches: output.stats.filesWithMatches, - filesSearched: output.stats.filesSearched, - truncated: output.truncated, - hook: 'onOutput', - }) - }, - execute: async (inputData, context): Promise => { + execute: async (input: z.infer, context): Promise => { const writer = context?.writer const abortSignal = context?.abortSignal + const requestContext = context?.requestContext as CodeSearchRequestContext | undefined // Check if operation was already cancelled if (abortSignal?.aborted === true) { @@ -159,7 +126,7 @@ Use for finding usages, identifying patterns, and code exploration.`, type: 'data-tool-progress', data: { status: 'in-progress', - message: `🔎 Starting code search for pattern '${inputData.pattern}'`, + message: `🔎 Starting code search for pattern '${input.pattern}'`, stage: 'coding:codeSearch', }, id: 'coding:codeSearch', @@ -167,22 +134,22 @@ Use for finding usages, identifying patterns, and code exploration.`, const tracer = trace.getTracer('code-search') const span = tracer.startSpan('code-search', { attributes: { - pattern: inputData.pattern, - targetCount: Array.isArray(inputData.target) - ? inputData.target.length + pattern: input.pattern, + targetCount: Array.isArray(input.target) + ? input.target.length : 1, operation: 'code-search', }, }) try { - const { pattern, target, options } = inputData + const { pattern, target, options } = input const isRegex = options?.isRegex ?? false - const caseSensitive = options?.caseSensitive ?? false - const maxResults = options?.maxResults ?? 100 - const includeContext = options?.includeContext ?? true - const contextLines = options?.contextLines ?? 2 - const maxFileSize = options?.maxFileSize ?? 1_000_000 + const caseSensitive = options?.caseSensitive ?? requestContext?.caseSensitive ?? false + const maxResults = options?.maxResults ?? requestContext?.maxResults ?? 100 + const includeContext = options?.includeContext ?? requestContext?.includeContext ?? true + const contextLines = options?.contextLines ?? requestContext?.contextLines ?? 2 + const maxFileSize = options?.maxFileSize ?? requestContext?.maxFileSize ?? 1_000_000 let filePaths: string[] = [] const targets = Array.isArray(target) ? target : [target] @@ -205,7 +172,7 @@ Use for finding usages, identifying patterns, and code exploration.`, filePaths = [...new Set(filePaths)] // Check for cancellation before file processing - if (abortSignal && abortSignal.aborted) { + if (abortSignal?.aborted) { span.setStatus({ code: 2, message: 'Operation cancelled during file processing', @@ -337,7 +304,7 @@ Use for finding usages, identifying patterns, and code exploration.`, // Handle AbortError specifically if (error instanceof Error && error.name === 'AbortError') { - const cancelMessage = `Code search cancelled for pattern "${inputData.pattern}"` + const cancelMessage = `Code search cancelled for pattern "${input.pattern}"` span.setStatus({ code: 2, message: cancelMessage }) span.end() @@ -364,4 +331,47 @@ Use for finding usages, identifying patterns, and code exploration.`, throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Code search tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Code search tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + const targetDesc = Array.isArray(input.target) + ? `${input.target.length} targets` + : input.target + log.info('Code search received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + pattern: input.pattern, + target: targetDesc, + isRegex: input.options?.isRegex ?? false, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Code search completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + totalMatches: output.stats.totalMatches, + filesWithMatches: output.stats.filesWithMatches, + filesSearched: output.stats.filesSearched, + truncated: output.truncated, + hook: 'onOutput', + }) + }, }) diff --git a/src/mastra/tools/color-change-tool.ts b/src/mastra/tools/color-change-tool.ts index a157912e..5f0f9ed4 100644 --- a/src/mastra/tools/color-change-tool.ts +++ b/src/mastra/tools/color-change-tool.ts @@ -1,6 +1,11 @@ import { createTool } from '@mastra/core/tools' import z from 'zod' import { log } from '../config/logger' +import type { RequestContext } from '@mastra/core/request-context' + +export interface ColorChangeRequestContext extends RequestContext { + theme?: string +} export function changeBgColor(color: string) { if (typeof window !== 'undefined') { @@ -14,7 +19,8 @@ export const colorChangeTool = createTool({ inputSchema: z.object({ color: z.string(), }), - execute: async ({ color }, context) => { + execute: async (input, context) => { + const { color } = input; await context?.writer?.custom({ type: 'data-tool-progress', data: { diff --git a/src/mastra/tools/copywriter-agent-tool.ts b/src/mastra/tools/copywriter-agent-tool.ts index 0e63328c..32b555b9 100644 --- a/src/mastra/tools/copywriter-agent-tool.ts +++ b/src/mastra/tools/copywriter-agent-tool.ts @@ -3,9 +3,14 @@ import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' import { SpanStatusCode, trace } from '@opentelemetry/api' import { z } from 'zod' +import type { RequestContext } from '@mastra/core/request-context' import { copywriterAgent } from '../agents/copywriterAgent' import { log } from '../config/logger' +export interface CopywriterRequestContext extends RequestContext { + userId?: string +} + log.info('Initializing Enhanced Copywriter Agent Tool...') export const copywriterTool = createTool({ @@ -67,14 +72,14 @@ export const copywriterTool = createTool({ .optional() .describe('Approximate word count of the content'), }), - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer const mastra = context?.mastra + const requestCtx = context?.requestContext as CopywriterRequestContext | undefined - const userIdVal = context?.requestContext?.get('userId') const userId = - typeof userIdVal === 'string' && userIdVal.trim().length > 0 - ? userIdVal + typeof requestCtx?.userId === 'string' && requestCtx.userId.trim().length > 0 + ? requestCtx.userId : 'anonymous' const { topic, @@ -83,7 +88,7 @@ export const copywriterTool = createTool({ tone, length = 'medium', specificRequirements, - } = inputData + } = input await writer?.custom({ type: 'data-tool-progress', diff --git a/src/mastra/tools/csv-to-json.tool.ts b/src/mastra/tools/csv-to-json.tool.ts index dc28f1c8..f892ed81 100644 --- a/src/mastra/tools/csv-to-json.tool.ts +++ b/src/mastra/tools/csv-to-json.tool.ts @@ -1,14 +1,17 @@ import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' +import type { RequestContext } from '@mastra/core/request-context' import { trace } from '@opentelemetry/api' import { parse } from 'csv-parse/sync' import * as fs from 'node:fs/promises' import { z } from 'zod' import { log } from '../config/logger' -const csvToolContextSchema = z.object({ - maxRows: z.number().optional(), -}) +export interface CsvToJsonRequestContext extends RequestContext { + csvToolContext?: { + maxRows?: number + } +} export const csvToJsonTool = createTool({ id: 'csv-to-json', @@ -44,53 +47,9 @@ export const csvToJsonTool = createTool({ skip_empty_lines: true, }), }), - outputSchema: z.object({ - data: z.array(z.any()).describe('Parsed JSON data'), - }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('CSV to JSON tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('CSV to JSON tool received input chunk', { - toolCallId, - inputTextDelta, - abortSignal: abortSignal?.aborted, - messageCount: messages.length, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - const source = input.filePath - ? `file:${input.filePath}` - : 'raw CSV data' - const options = input.options || {} - log.info('CSV to JSON received complete input', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - source, - delimiter: options.delimiter, - columns: options.columns, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('CSV to JSON conversion completed', { - toolCallId, - toolName, - abortSignal: abortSignal?.aborted, - recordsProcessed: output.data.length, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer - const requestContext = context?.requestContext + const requestContext = context?.requestContext as CsvToJsonRequestContext | undefined const abortSignal = context?.abortSignal // Check if operation was already cancelled @@ -115,11 +74,8 @@ export const csvToJsonTool = createTool({ }) try { - const config = requestContext?.get('csvToolContext') - const { maxRows } = - config !== undefined - ? csvToolContextSchema.parse(config) - : { maxRows: undefined } + const config = requestContext?.csvToolContext + const maxRows = config?.maxRows // Check for cancellation before file operations if (abortSignal && abortSignal.aborted) { @@ -133,20 +89,20 @@ export const csvToJsonTool = createTool({ ) } - let contentToParse = inputData.csvData + let contentToParse = input.csvData if ( - inputData.filePath !== undefined && - inputData.filePath !== null + input.filePath !== undefined && + input.filePath !== null ) { try { contentToParse = await fs.readFile( - inputData.filePath, + input.filePath, 'utf-8' ) } catch (err) { throw new Error( - `Failed to read file at ${inputData.filePath}: ${(err as Error).message}` + `Failed to read file at ${input.filePath}: ${(err as Error).message}` ) } } @@ -159,7 +115,7 @@ export const csvToJsonTool = createTool({ throw new Error('Either csvData or filePath must be provided') } - const options = inputData.options ?? { + const options = input.options ?? { delimiter: ',', columns: true, trim: true, @@ -222,6 +178,47 @@ export const csvToJsonTool = createTool({ throw error instanceof Error ? error : new Error(errorMessage) } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('CSV to JSON tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('CSV to JSON tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + const source = input.filePath + ? `file:${input.filePath}` + : 'raw CSV data' + const options = input.options || {} + log.info('CSV to JSON received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + source, + delimiter: options.delimiter, + columns: options.columns, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('CSV to JSON conversion completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + recordsProcessed: output.data.length, + hook: 'onOutput', + }) + }, }) export type CsvToJsonUITool = InferUITool diff --git a/src/mastra/tools/data-file-manager.ts b/src/mastra/tools/data-file-manager.ts index 92481803..608f7769 100644 --- a/src/mastra/tools/data-file-manager.ts +++ b/src/mastra/tools/data-file-manager.ts @@ -75,43 +75,22 @@ export const readDataFileTool = createTool({ ), }), outputSchema: z.string().describe('The content of the file as a string.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Read file tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Read file received complete input', { - toolCallId, - messageCount: messages.length, - fileName: input.fileName, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Read file completed', { - toolCallId, - toolName, - fileSize: output.length, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer; + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; await ensureDataDir(); - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '📖 Reading file: ' + inputData.fileName, stage: 'read:file' }, id: 'read:file' }); + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '📖 Reading file: ' + input.fileName, stage: 'read:file' }, id: 'read:file' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const readSpan = tracer.startSpan('read:file', { attributes: { 'tool.id': 'read:file', - 'tool.input.fileName': inputData.fileName, + 'tool.input.fileName': input.fileName, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { fileName } = inputData + const { fileName } = input const fullPath = validateDataPath(fileName) // Resolve the real path to protect against symlink/relative attacks, then validate it is inside DATA_DIR. const realFullPath = await fs.realpath(fullPath) @@ -139,64 +118,64 @@ export const readDataFileTool = createTool({ throw error } }, -}) - -export type ReadDataFileUITool = InferUITool; - -export const writeDataFileTool = createTool({ - id: 'write:file', - description: - 'Writes content to a file within the data directory. If the file does not exist, it will be created. If it exists, its content will be overwritten.', - inputSchema: z.object({ - fileName: z - .string() - .describe( - 'The name of the file (relative to the data/ directory).' - ), - content: z.string().describe('The content to write to the file.'), - }), - outputSchema: z - .string() - .describe('A confirmation string indicating success.'), onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Write file tool input streaming started', { + log.info('Read file tool input streaming started', { toolCallId, messageCount: messages.length, hook: 'onInputStart', }) }, onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Write file received complete input', { + log.info('Read file received complete input', { toolCallId, messageCount: messages.length, fileName: input.fileName, - contentLength: input.content.length, hook: 'onInputAvailable', }) }, onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Write file completed', { + log.info('Read file completed', { toolCallId, toolName, - success: output.includes('successfully'), + fileSize: output.length, hook: 'onOutput', }) }, - execute: async (inputData, context) => { - const writer = context?.writer; +}) - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '💾 Writing to file: ' + inputData.fileName, stage: 'write:file' }, id: 'write:file' }); +export type ReadDataFileUITool = InferUITool; + +export const writeDataFileTool = createTool({ + id: 'write:file', + description: + 'Writes content to a file within the data directory. If the file does not exist, it will be created. If it exists, its content will be overwritten.', + inputSchema: z.object({ + fileName: z + .string() + .describe( + 'The name of the file (relative to the data/ directory).' + ), + content: z.string().describe('The content to write to the file.'), + }), + outputSchema: z + .string() + .describe('A confirmation string indicating success.'), + execute: async (input, context) => { + const writer = context?.writer; + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '💾 Writing to file: ' + input.fileName, stage: 'write:file' }, id: 'write:file' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const writeSpan = tracer.startSpan('write:file', { attributes: { 'tool.id': 'write:file', - 'tool.input.fileName': inputData.fileName, - 'tool.input.contentLength': inputData.content.length, + 'tool.input.fileName': input.fileName, + 'tool.input.contentLength': input.content.length, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { fileName, content } = inputData + const { fileName, content } = input const fullPath = validateDataPath(fileName) // Resolve the real path to protect against symlink/relative attacks, then validate it is inside DATA_DIR. const realFullPath = await fs.realpath(fullPath) @@ -232,59 +211,62 @@ export const writeDataFileTool = createTool({ throw error } }, -}) - -export type WriteDataFileUITool = InferUITool; - -export const deleteDataFileTool = createTool({ - id: 'delete:file', - description: 'Deletes a file within the data directory.', - inputSchema: z.object({ - fileName: z - .string() - .describe( - 'The name of the file (relative to the data/ directory).' - ), - }), - outputSchema: z - .string() - .describe('A confirmation string indicating success.'), onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Delete file tool input streaming started', { + log.info('Write file tool input streaming started', { toolCallId, messageCount: messages.length, hook: 'onInputStart', }) }, onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Delete file received complete input', { + log.info('Write file received complete input', { toolCallId, messageCount: messages.length, fileName: input.fileName, + contentLength: input.content.length, hook: 'onInputAvailable', }) }, onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Delete file completed', { + log.info('Write file completed', { toolCallId, toolName, success: output.includes('successfully'), hook: 'onOutput', }) }, - execute: async (inputData, context) => { +}) + +export type WriteDataFileUITool = InferUITool; + +export const deleteDataFileTool = createTool({ + id: 'delete:file', + description: 'Deletes a file within the data directory.', + inputSchema: z.object({ + fileName: z + .string() + .describe( + 'The name of the file (relative to the data/ directory).' + ), + }), + outputSchema: z + .string() + .describe('A confirmation string indicating success.'), + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🗑️ Deleting file: ' + inputData.fileName, stage: 'delete:file' }, id: 'delete:file' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🗑️ Deleting file: ' + input.fileName, stage: 'delete:file' }, id: 'delete:file' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const deleteSpan = tracer.startSpan('delete:file', { attributes: { 'tool.id': 'delete:file', - 'tool.input.fileName': inputData.fileName, + 'tool.input.fileName': input.fileName, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { fileName } = inputData + const { fileName } = input const fullPath = validateDataPath(fileName) // Defensive: Ensure fullPath is within DATA_DIR before deleting if (!fullPath.startsWith(DATA_DIR)) { @@ -311,6 +293,29 @@ export const deleteDataFileTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Delete file tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Delete file received complete input', { + toolCallId, + messageCount: messages.length, + fileName: input.fileName, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Delete file completed', { + toolCallId, + toolName, + success: output.includes('successfully'), + hook: 'onOutput', + }) + }, }) export type DeleteDataFileUITool = InferUITool; @@ -330,42 +335,21 @@ export const listDataDirTool = createTool({ outputSchema: z .array(z.string()) .describe('An array of file and directory names.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('List directory tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('List directory received complete input', { - toolCallId, - messageCount: messages.length, - dirPath: input.dirPath, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('List directory completed', { - toolCallId, - toolName, - itemCount: output.length, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '📂 Listing directory: ' + (inputData.dirPath ?? 'docs/data'), stage: 'list:directory' }, id: 'list:directory' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '📂 Listing directory: ' + (input.dirPath ?? 'docs/data'), stage: 'list:directory' }, id: 'list:directory' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const listSpan = tracer.startSpan('list:directory', { attributes: { 'tool.id': 'list:directory', - 'tool.input.dirPath': inputData.dirPath ?? 'docs/data', + 'tool.input.dirPath': input.dirPath ?? 'docs/data', + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { dirPath = 'docs/data' } = inputData + const { dirPath = 'docs/data' } = input const fullPath = validateDataPath(dirPath) // Defensive: Ensure fullPath is within DATA_DIR before reading directory if (!fullPath.startsWith(DATA_DIR)) { @@ -392,6 +376,29 @@ export const listDataDirTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('List directory tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('List directory received complete input', { + toolCallId, + messageCount: messages.length, + dirPath: input.dirPath, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('List directory completed', { + toolCallId, + toolName, + itemCount: output.length, + hook: 'onOutput', + }) + }, }) export type ListDataDirUITool = InferUITool; @@ -414,44 +421,22 @@ export const copyDataFileTool = createTool({ outputSchema: z .string() .describe('A confirmation string indicating success.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Copy file tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Copy file received complete input', { - toolCallId, - messageCount: messages.length, - sourceFile: input.sourceFile, - destFile: input.destFile, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Copy file completed', { - toolCallId, - toolName, - success: output.includes('successfully'), - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `📋 Copying file: ${inputData.sourceFile} to ${inputData.destFile}`, stage: 'copy:file' }, id: 'copy:file' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `📋 Copying file: ${input.sourceFile} to ${input.destFile}`, stage: 'copy:file' }, id: 'copy:file' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const copySpan = tracer.startSpan('copy:file', { attributes: { 'tool.id': 'copy:file', - 'tool.input.sourceFile': inputData.sourceFile, - 'tool.input.destFile': inputData.destFile, + 'tool.input.sourceFile': input.sourceFile, + 'tool.input.destFile': input.destFile, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { sourceFile, destFile } = inputData + const { sourceFile, destFile } = input const sourcePath = validateDataPath(sourceFile) const destPath = validateDataPath(destFile) // Defensive: Ensure both paths are within DATA_DIR @@ -484,6 +469,30 @@ export const copyDataFileTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Copy file tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Copy file received complete input', { + toolCallId, + messageCount: messages.length, + sourceFile: input.sourceFile, + destFile: input.destFile, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Copy file completed', { + toolCallId, + toolName, + success: output.includes('successfully'), + hook: 'onOutput', + }) + }, }) export type CopyDataFileUITool = InferUITool; @@ -506,44 +515,22 @@ export const moveDataFileTool = createTool({ outputSchema: z .string() .describe('A confirmation string indicating success.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Move file tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Move file received complete input', { - toolCallId, - messageCount: messages.length, - sourceFile: input.sourceFile, - destFile: input.destFile, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Move file completed', { - toolCallId, - toolName, - success: output.includes('successfully'), - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `🚚 Moving file: ${inputData.sourceFile} to ${inputData.destFile}`, stage: 'move:file' }, id: 'move:file' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `🚚 Moving file: ${input.sourceFile} to ${input.destFile}`, stage: 'move:file' }, id: 'move:file' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const moveSpan = tracer.startSpan('move:file', { attributes: { 'tool.id': 'move:file', - 'tool.input.sourceFile': inputData.sourceFile, - 'tool.input.destFile': inputData.destFile, + 'tool.input.sourceFile': input.sourceFile, + 'tool.input.destFile': input.destFile, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { sourceFile, destFile } = inputData + const { sourceFile, destFile } = input const sourcePath = validateDataPath(sourceFile) const destPath = validateDataPath(destFile) // Defensive: Ensure both paths are within DATA_DIR @@ -576,6 +563,30 @@ export const moveDataFileTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Move file tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Move file received complete input', { + toolCallId, + messageCount: messages.length, + sourceFile: input.sourceFile, + destFile: input.destFile, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Move file completed', { + toolCallId, + toolName, + success: output.includes('successfully'), + hook: 'onOutput', + }) + }, }) export type MoveDataFileUITool = InferUITool; @@ -599,46 +610,23 @@ export const searchDataFilesTool = createTool({ .optional() .describe( "The directory to search in (relative to data/, default: '')." - ), - }), - outputSchema: z - .array(z.string()) - .describe('An array of matching file paths.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Search files tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Search files received complete input', { - toolCallId, - messageCount: messages.length, - pattern: input.pattern, - searchContent: input.searchContent, - dirPath: input.dirPath, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Search files completed', { - toolCallId, - toolName, - resultCount: output.length, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + ), + }), + outputSchema: z + .array(z.string()) + .describe('An array of matching file paths.'), + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `🔍 Searching for pattern: "${inputData.pattern}"`, stage: 'search:files' }, id: 'search:files' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `🔍 Searching for pattern: "${input.pattern}"`, stage: 'search:files' }, id: 'search:files' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const searchSpan = tracer.startSpan('search:files', { attributes: { 'tool.id': 'search:files', - 'tool.input.pattern': inputData.pattern, - 'tool.input.searchContent': inputData.searchContent, - 'tool.input.dirPath': inputData.dirPath, + 'tool.input.pattern': input.pattern, + 'tool.input.searchContent': input.searchContent, + 'tool.input.dirPath': input.dirPath, + 'tool.requestContext.userId': requestCtx?.userId, } }); @@ -647,7 +635,7 @@ export const searchDataFilesTool = createTool({ pattern, searchContent = false, dirPath = 'docs/data', - } = inputData + } = input const searchPath = validateDataPath(dirPath) if (!searchPath.startsWith(DATA_DIR)) { throw new Error( @@ -714,6 +702,31 @@ export const searchDataFilesTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Search files tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Search files received complete input', { + toolCallId, + messageCount: messages.length, + pattern: input.pattern, + searchContent: input.searchContent, + dirPath: input.dirPath, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Search files completed', { + toolCallId, + toolName, + resultCount: output.length, + hook: 'onOutput', + }) + }, }) export type SearchFilesUITool = InferUITool; @@ -738,43 +751,21 @@ export const getDataFileInfoTool = createTool({ isDirectory: z.boolean(), }) .describe('File metadata information.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Get file info tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Get file info received complete input', { - toolCallId, - messageCount: messages.length, - fileName: input.fileName, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Get file info completed', { - toolCallId, - toolName, - fileSize: output.size, - isFile: output.isFile, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: 'ℹ️ Getting info for file: ' + inputData.fileName, stage: 'get:fileinfo' }, id: 'get:fileinfo' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: 'ℹ️ Getting info for file: ' + input.fileName, stage: 'get:fileinfo' }, id: 'get:fileinfo' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const infoSpan = tracer.startSpan('get:fileinfo', { attributes: { 'tool.id': 'get:fileinfo', - 'tool.input.fileName': inputData.fileName, + 'tool.input.fileName': input.fileName, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { fileName } = inputData + const { fileName } = input const fullPath = validateDataPath(fileName) // Resolve the real path to protect against symlink/relative attacks, then validate it is inside DATA_DIR. const realFullPath = await fs.realpath(fullPath) @@ -810,58 +801,61 @@ export const getDataFileInfoTool = createTool({ throw error } }, -}) - -export type GetDataFileInfoUITool = InferUITool; -export const createDataDirTool = createTool({ - id: 'create:directory', - description: 'Creates a new directory within the data directory.', - inputSchema: z.object({ - dirPath: z - .string() - .describe( - 'The path of the directory to create (relative to the data/ directory).' - ), - }), - outputSchema: z - .string() - .describe('A confirmation string indicating success.'), onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Create directory tool input streaming started', { + log.info('Get file info tool input streaming started', { toolCallId, messageCount: messages.length, hook: 'onInputStart', }) }, onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Create directory received complete input', { + log.info('Get file info received complete input', { toolCallId, messageCount: messages.length, - dirPath: input.dirPath, + fileName: input.fileName, hook: 'onInputAvailable', }) }, onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Create directory completed', { + log.info('Get file info completed', { toolCallId, toolName, - success: output.includes('successfully'), + fileSize: output.size, + isFile: output.isFile, hook: 'onOutput', }) }, - execute: async (inputData, context) => { +}) + +export type GetDataFileInfoUITool = InferUITool; +export const createDataDirTool = createTool({ + id: 'create:directory', + description: 'Creates a new directory within the data directory.', + inputSchema: z.object({ + dirPath: z + .string() + .describe( + 'The path of the directory to create (relative to the data/ directory).' + ), + }), + outputSchema: z + .string() + .describe('A confirmation string indicating success.'), + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '📁 Creating directory: ' + inputData.dirPath, stage: 'create:directory' }, id: 'create:directory' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '📁 Creating directory: ' + input.dirPath, stage: 'create:directory' }, id: 'create:directory' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const createDirSpan = tracer.startSpan('create:directory', { attributes: { 'tool.id': 'create:directory', - 'tool.input.dirPath': inputData.dirPath, + 'tool.input.dirPath': input.dirPath, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { dirPath } = inputData + const { dirPath } = input const fullPath = validateDataPath(dirPath) if (!fullPath.startsWith(DATA_DIR)) { throw new Error( @@ -887,32 +881,15 @@ export const createDataDirTool = createTool({ throw error } }, -}) - -export type CreateDataDirUITool = InferUITool; - -export const removeDataDirTool = createTool({ - id: 'remove:directory', - description: 'Removes an empty directory within the data directory.', - inputSchema: z.object({ - dirPath: z - .string() - .describe( - 'The path of the directory to remove (relative to the data/ directory).' - ), - }), - outputSchema: z - .string() - .describe('A confirmation string indicating success.'), onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Remove directory tool input streaming started', { + log.info('Create directory tool input streaming started', { toolCallId, messageCount: messages.length, hook: 'onInputStart', }) }, onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Remove directory received complete input', { + log.info('Create directory received complete input', { toolCallId, messageCount: messages.length, dirPath: input.dirPath, @@ -920,26 +897,45 @@ export const removeDataDirTool = createTool({ }) }, onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Remove directory completed', { + log.info('Create directory completed', { toolCallId, toolName, success: output.includes('successfully'), hook: 'onOutput', }) }, - execute: async (inputData, context) => { +}) + +export type CreateDataDirUITool = InferUITool; + +export const removeDataDirTool = createTool({ + id: 'remove:directory', + description: 'Removes an empty directory within the data directory.', + inputSchema: z.object({ + dirPath: z + .string() + .describe( + 'The path of the directory to remove (relative to the data/ directory).' + ), + }), + outputSchema: z + .string() + .describe('A confirmation string indicating success.'), + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🗑️ Removing directory: ' + inputData.dirPath, stage: 'remove:directory' }, id: 'remove:directory' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🗑️ Removing directory: ' + input.dirPath, stage: 'remove:directory' }, id: 'remove:directory' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const removeDirSpan = tracer.startSpan('remove:directory', { attributes: { 'tool.id': 'remove:directory', - 'tool.input.dirPath': inputData.dirPath, + 'tool.input.dirPath': input.dirPath, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { dirPath } = inputData + const { dirPath } = input const fullPath = validateDataPath(dirPath) if (!fullPath.startsWith(DATA_DIR)) { throw new Error( @@ -970,6 +966,29 @@ export const removeDataDirTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Remove directory tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Remove directory received complete input', { + toolCallId, + messageCount: messages.length, + dirPath: input.dirPath, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Remove directory completed', { + toolCallId, + toolName, + success: output.includes('successfully'), + hook: 'onOutput', + }) + }, }) export type RemoveDataDirUITool = InferUITool; @@ -993,44 +1012,22 @@ export const archiveDataTool = createTool({ outputSchema: z .string() .describe('A confirmation string indicating success.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Archive data tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Archive data received complete input', { - toolCallId, - messageCount: messages.length, - sourcePath: input.sourcePath, - archiveName: input.archiveName, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Archive data completed', { - toolCallId, - toolName, - success: output.includes('successfully'), - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `📦 Archiving: ${inputData.sourcePath} to ${inputData.archiveName}.gz`, stage: 'archive:data' }, id: 'archive:data' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `📦 Archiving: ${input.sourcePath} to ${input.archiveName}.gz`, stage: 'archive:data' }, id: 'archive:data' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const archiveSpan = tracer.startSpan('archive:data', { attributes: { 'tool.id': 'archive:data', - 'tool.input.sourcePath': inputData.sourcePath, - 'tool.input.archiveName': inputData.archiveName, + 'tool.input.sourcePath': input.sourcePath, + 'tool.input.archiveName': input.archiveName, + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { sourcePath, archiveName } = inputData + const { sourcePath, archiveName } = input const sourceFullPath = validateDataPath(sourcePath) const archiveFullPath = validateDataPath(archiveName + '.gz') if ( @@ -1068,6 +1065,30 @@ export const archiveDataTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Archive data tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Archive data received complete input', { + toolCallId, + messageCount: messages.length, + sourcePath: input.sourcePath, + archiveName: input.archiveName, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Archive data completed', { + toolCallId, + toolName, + success: output.includes('successfully'), + hook: 'onOutput', + }) + }, }) export type ArchiveDataUITool = InferUITool; @@ -1092,44 +1113,22 @@ export const backupDataTool = createTool({ outputSchema: z .string() .describe('A confirmation string indicating success with backup path.'), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Backup data tool input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Backup data received complete input', { - toolCallId, - messageCount: messages.length, - sourcePath: input.sourcePath, - backupDir: input.backupDir, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Backup data completed', { - toolCallId, - toolName, - success: output.includes('Backup created'), - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer; - await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `💾 Creating backup for: ${inputData.sourcePath}`, stage: 'backup:data' }, id: 'backup:data' }); + const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: `💾 Creating backup for: ${input.sourcePath}`, stage: 'backup:data' }, id: 'backup:data' }); const tracer = trace.getTracer('data-file-manager', '1.0.0'); const backupSpan = tracer.startSpan('backup:data', { attributes: { 'tool.id': 'backup:data', - 'tool.input.sourcePath': inputData.sourcePath, - 'tool.input.backupDir': inputData.backupDir ?? 'backups', + 'tool.input.sourcePath': input.sourcePath, + 'tool.input.backupDir': input.backupDir ?? 'backups', + 'tool.requestContext.userId': requestCtx?.userId, } }); try { - const { sourcePath, backupDir = 'backups' } = inputData + const { sourcePath, backupDir = 'backups' } = input const sourceFullPath = validateDataPath(sourcePath) if (!sourceFullPath.startsWith(DATA_DIR)) { throw new Error( @@ -1169,6 +1168,30 @@ export const backupDataTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Backup data tool input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Backup data received complete input', { + toolCallId, + messageCount: messages.length, + sourcePath: input.sourcePath, + backupDir: input.backupDir, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Backup data completed', { + toolCallId, + toolName, + success: output.includes('Backup created'), + hook: 'onOutput', + }) + }, }) export type BackupDataUITool = InferUITool; diff --git a/src/mastra/tools/finnhub-tools.ts b/src/mastra/tools/finnhub-tools.ts index af46ef8f..41d40b1d 100644 --- a/src/mastra/tools/finnhub-tools.ts +++ b/src/mastra/tools/finnhub-tools.ts @@ -1,5 +1,6 @@ import type { InferUITool } from "@mastra/core/tools"; import { createTool } from "@mastra/core/tools"; +import type { RequestContext } from "@mastra/core/request-context"; import { SpanStatusCode, trace } from "@opentelemetry/api"; import { z } from "zod"; import { @@ -10,6 +11,10 @@ import { logToolExecution, } from "../config/logger"; +export interface FinnhubRequestContext extends RequestContext { + userId?: string; +} + /** * Finnhub Quotes Tool * @@ -31,35 +36,14 @@ export const finnhubQuotesTool = createTool({ }).optional(), message: z.string().optional() }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Finnhub quotes input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Finnhub quotes received complete input', { - toolCallId, - messageCount: messages.length, - symbol: input.symbol, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Finnhub quotes completed', { - toolCallId, - toolName, - hasData: !!output.data, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { + const requestCtx = context?.requestContext as FinnhubRequestContext | undefined; const tracer = trace.getTracer('finnhub-tools'); const rootSpan = tracer.startSpan('finnhub-quotes', { attributes: { - symbol: inputData.symbol, - operation: 'finnhub-quotes' + symbol: input.symbol, + operation: 'finnhub-quotes', + 'tool.requestContext.userId': requestCtx?.userId, } }); @@ -67,13 +51,13 @@ export const finnhubQuotesTool = createTool({ type: 'data-tool-progress', data: { status: 'in-progress', - message: `Fetching quote for ${inputData.symbol}`, + message: `Fetching quote for ${input.symbol}`, stage: 'finnhub-quotes' }, id: 'finnhub-quotes' }); - logToolExecution('finnhubQuotesTool', { input: inputData }); + logToolExecution('finnhubQuotesTool', { input: input }); const apiKey = process.env.FINNHUB_API_KEY; @@ -82,7 +66,7 @@ export const finnhubQuotesTool = createTool({ rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubQuotesTool', new Error(message), { symbol: inputData.symbol }); + logError('finnhubQuotesTool', new Error(message), { symbol: input.symbol }); return { data: null, message @@ -92,7 +76,7 @@ export const finnhubQuotesTool = createTool({ try { const params = new URLSearchParams(); params.append('token', apiKey); - params.append('symbol', inputData.symbol); + params.append('symbol', input.symbol); const url = `https://finnhub.io/api/v1/quote?${params.toString()}`; @@ -105,7 +89,7 @@ export const finnhubQuotesTool = createTool({ }); logStepStart('finnhub-api-call', { - symbol: inputData.symbol, + symbol: input.symbol, url: url.replace(apiKey, '[REDACTED]') }); @@ -114,19 +98,7 @@ export const finnhubQuotesTool = createTool({ const data = await response.json(); const apiDuration = Date.now() - apiStartTime; - /*apiSpan.end({ - output: { ... } - }); -> OTEL doesn't accept object in end(). Only timestamp.*/ - - // Set attributes on apiSpan - // apiSpan.setAttribute('http.status_code', response.status); - // ... - // But I can't replicate `end({output})` directly. - // I will set attributes and then end. - apiSpan?.end(); - // Wait, the previous code used `apiSpan?.end({ output: ... })`. This was a custom wrapper. - // I will just end it. logStepEnd('finnhub-api-call', { status: response.status, @@ -140,7 +112,7 @@ export const finnhubQuotesTool = createTool({ const message = String(errorValue); rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); - logError('finnhubQuotesTool', new Error(message), { symbol: inputData.symbol, apiError: message }); + logError('finnhubQuotesTool', new Error(message), { symbol: input.symbol, apiError: message }); return { data: null, message @@ -151,14 +123,14 @@ export const finnhubQuotesTool = createTool({ const result = { data, metadata: { - symbol: inputData.symbol + symbol: input.symbol }, message: undefined }; rootSpan.end(); - logToolExecution('finnhubQuotesTool', { output: { symbol: inputData.symbol } }); + logToolExecution('finnhubQuotesTool', { output: { symbol: input.symbol } }); return result; @@ -167,14 +139,37 @@ export const finnhubQuotesTool = createTool({ rootSpan.recordException(error instanceof Error ? error : new Error(errorMessage)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); rootSpan.end(); - logError('finnhubQuotesTool', error instanceof Error ? error : new Error(errorMessage), { symbol: inputData.symbol }); + logError('finnhubQuotesTool', error instanceof Error ? error : new Error(errorMessage), { symbol: input.symbol }); return { data: null, message: errorMessage }; } - } -}); + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Finnhub quotes input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Finnhub quotes received complete input', { + toolCallId, + messageCount: messages.length, + symbol: input.symbol, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Finnhub quotes completed', { + toolCallId, + toolName, + hasData: !!output.data, + hook: 'onOutput', + }) + }, +}) /** * Finnhub Company Tool @@ -207,37 +202,15 @@ export const finnhubCompanyTool = createTool({ }).optional(), message: z.string().optional() }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Finnhub company input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Finnhub company received complete input', { - toolCallId, - messageCount: messages.length, - function: input.function, - symbol: input.symbol, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Finnhub company completed', { - toolCallId, - toolName, - hasData: !!output.data, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { + const requestCtx = context?.requestContext as FinnhubRequestContext | undefined; const tracer = trace.getTracer('finnhub-tools'); const rootSpan = tracer.startSpan('finnhub-company', { attributes: { - function: inputData.function, - symbol: inputData.symbol, - operation: 'finnhub-company' + function: input.function, + symbol: input.symbol, + operation: 'finnhub-company', + 'tool.requestContext.userId': requestCtx?.userId, } }); @@ -245,13 +218,13 @@ export const finnhubCompanyTool = createTool({ type: 'data-tool-progress', data: { status: 'in-progress', - message: `Fetching company data (${inputData.function}) for ${inputData.symbol}`, + message: `Fetching company data (${input.function}) for ${input.symbol}`, stage: 'finnhub-company' }, id: 'finnhub-company' }); - logToolExecution('finnhubCompanyTool', { input: inputData }); + logToolExecution('finnhubCompanyTool', { input: input }); const apiKey = process.env.FINNHUB_API_KEY; @@ -260,7 +233,7 @@ export const finnhubCompanyTool = createTool({ rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubCompanyTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubCompanyTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -272,28 +245,28 @@ export const finnhubCompanyTool = createTool({ const params = new URLSearchParams(); params.append('token', apiKey); - switch (inputData.function) { + switch (input.function) { case "COMPANY_PROFILE": - params.append('symbol', inputData.symbol); + params.append('symbol', input.symbol); url = `https://finnhub.io/api/v1/stock/profile2?${params.toString()}`; break; case "COMPANY_NEWS": { - params.append('symbol', inputData.symbol); - if (inputData.from !== undefined && inputData.from !== null && inputData.from.trim() !== '') { - params.append('from', inputData.from); + params.append('symbol', input.symbol); + if (input.from !== undefined && input.from !== null && input.from.trim() !== '') { + params.append('from', input.from); } - if (inputData.to !== undefined && inputData.to !== null && inputData.to.trim() !== '') { - params.append('to', inputData.to); + if (input.to !== undefined && input.to !== null && input.to.trim() !== '') { + params.append('to', input.to); } url = `https://finnhub.io/api/v1/company-news?${params.toString()}`; break; } default: { - const message = `Unsupported function: ${inputData.function}`; + const message = `Unsupported function: ${input.function}`; rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubCompanyTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubCompanyTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -315,8 +288,8 @@ export const finnhubCompanyTool = createTool({ }); logStepStart('finnhub-api-call', { - function: inputData.function, - symbol: inputData.symbol, + function: input.function, + symbol: input.symbol, url: url.replace(apiKey, '[REDACTED]') }); @@ -339,7 +312,7 @@ export const finnhubCompanyTool = createTool({ const message = String(errorValue); rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); - logError('finnhubCompanyTool', new Error(message), { function: inputData.function, symbol: inputData.symbol, apiError: message }); + logError('finnhubCompanyTool', new Error(message), { function: input.function, symbol: input.symbol, apiError: message }); return { data: null, message @@ -350,17 +323,17 @@ export const finnhubCompanyTool = createTool({ const result = { data, metadata: { - function: inputData.function, - symbol: inputData.symbol, - from: inputData.from, - to: inputData.to + function: input.function, + symbol: input.symbol, + from: input.from, + to: input.to }, message: undefined }; rootSpan.end(); - logToolExecution('finnhubCompanyTool', { output: { function: inputData.function, symbol: inputData.symbol } }); + logToolExecution('finnhubCompanyTool', { output: { function: input.function, symbol: input.symbol } }); return result; @@ -369,14 +342,38 @@ export const finnhubCompanyTool = createTool({ rootSpan.recordException(error instanceof Error ? error : new Error(errorMessage)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); rootSpan.end(); - logError('finnhubCompanyTool', error instanceof Error ? error : new Error(errorMessage), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubCompanyTool', error instanceof Error ? error : new Error(errorMessage), { function: input.function, symbol: input.symbol }); return { data: null, message: errorMessage }; } - } -}); + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Finnhub company input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Finnhub company received complete input', { + toolCallId, + messageCount: messages.length, + function: input.function, + symbol: input.symbol, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Finnhub company completed', { + toolCallId, + toolName, + hasData: !!output.data, + hook: 'onOutput', + }) + }, +}) /** * Finnhub Financials Tool @@ -409,37 +406,15 @@ export const finnhubFinancialsTool = createTool({ }).optional(), message: z.string().optional() }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Finnhub financials input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Finnhub financials received complete input', { - toolCallId, - messageCount: messages.length, - function: input.function, - symbol: input.symbol, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Finnhub financials completed', { - toolCallId, - toolName, - hasData: !!output.data, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { + const requestCtx = context?.requestContext as FinnhubRequestContext | undefined; const tracer = trace.getTracer('finnhub-tools'); const rootSpan = tracer.startSpan('finnhub-financials', { attributes: { - function: inputData.function, - symbol: inputData.symbol, - operation: 'finnhub-financials' + function: input.function, + symbol: input.symbol, + operation: 'finnhub-financials', + 'tool.requestContext.userId': requestCtx?.userId, } }); @@ -447,13 +422,13 @@ export const finnhubFinancialsTool = createTool({ type: 'data-tool-progress', data: { status: 'in-progress', - message: `Fetching financials (${inputData.function}) for ${inputData.symbol}`, + message: `Fetching financials (${input.function}) for ${input.symbol}`, stage: 'finnhub-financials' }, id: 'finnhub-financials' }); - logToolExecution('finnhubFinancialsTool', { input: inputData }); + logToolExecution('finnhubFinancialsTool', { input: input }); const apiKey = process.env.FINNHUB_API_KEY; @@ -462,7 +437,7 @@ export const finnhubFinancialsTool = createTool({ rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubFinancialsTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubFinancialsTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -473,9 +448,9 @@ export const finnhubFinancialsTool = createTool({ let url: string; const params = new URLSearchParams(); params.append('token', apiKey); - params.append('symbol', inputData.symbol); + params.append('symbol', input.symbol); - switch (inputData.function) { + switch (input.function) { case "FINANCIAL_STATEMENTS": { url = `https://finnhub.io/api/v1/stock/financials?${params.toString()}`; break; @@ -493,11 +468,11 @@ export const finnhubFinancialsTool = createTool({ break; } default: { - const message = `Unsupported function: ${inputData.function}`; + const message = `Unsupported function: ${input.function}`; rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubFinancialsTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubFinancialsTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -514,8 +489,8 @@ export const finnhubFinancialsTool = createTool({ }); logStepStart('finnhub-api-call', { - function: inputData.function, - symbol: inputData.symbol, + function: input.function, + symbol: input.symbol, url: url.replace(apiKey, '[REDACTED]') }); @@ -538,7 +513,7 @@ export const finnhubFinancialsTool = createTool({ const message = String(errorValue); rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); - logError('finnhubFinancialsTool', new Error(message), { function: inputData.function, symbol: inputData.symbol, apiError: message }); + logError('finnhubFinancialsTool', new Error(message), { function: input.function, symbol: input.symbol, apiError: message }); return { data: null, message @@ -549,15 +524,15 @@ export const finnhubFinancialsTool = createTool({ const result = { data, metadata: { - function: inputData.function, - symbol: inputData.symbol + function: input.function, + symbol: input.symbol }, message: undefined }; rootSpan.end(); - logToolExecution('finnhubFinancialsTool', { output: { function: inputData.function, symbol: inputData.symbol } }); + logToolExecution('finnhubFinancialsTool', { output: { function: input.function, symbol: input.symbol } }); return result; @@ -566,14 +541,38 @@ export const finnhubFinancialsTool = createTool({ rootSpan.recordException(error instanceof Error ? error : new Error(errorMessage)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); rootSpan.end(); - logError('finnhubFinancialsTool', error instanceof Error ? error : new Error(errorMessage), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubFinancialsTool', error instanceof Error ? error : new Error(errorMessage), { function: input.function, symbol: input.symbol }); return { data: null, message: errorMessage }; } - } -}); + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Finnhub financials input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Finnhub financials received complete input', { + toolCallId, + messageCount: messages.length, + function: input.function, + symbol: input.symbol, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Finnhub financials completed', { + toolCallId, + toolName, + hasData: !!output.data, + hook: 'onOutput', + }) + }, +}) /** * Finnhub Analysis Tool @@ -602,37 +601,15 @@ export const finnhubAnalysisTool = createTool({ }).optional(), message: z.string().optional() }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Finnhub analysis input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Finnhub analysis received complete input', { - toolCallId, - messageCount: messages.length, - function: input.function, - symbol: input.symbol, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Finnhub analysis completed', { - toolCallId, - toolName, - hasData: !!output.data, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { + const requestCtx = context?.requestContext as FinnhubRequestContext | undefined; const tracer = trace.getTracer('finnhub-tools'); const rootSpan = tracer.startSpan('finnhub-analysis', { attributes: { - function: inputData.function, - symbol: inputData.symbol, - operation: 'finnhub-analysis' + function: input.function, + symbol: input.symbol, + operation: 'finnhub-analysis', + 'tool.requestContext.userId': requestCtx?.userId, } }); @@ -640,13 +617,13 @@ export const finnhubAnalysisTool = createTool({ type: 'data-tool-progress', data: { status: 'in-progress', - message: `Fetching analysis (${inputData.function}) for ${inputData.symbol}`, + message: `Fetching analysis (${input.function}) for ${input.symbol}`, stage: 'finnhub-analysis' }, id: 'finnhub-analysis' }); - logToolExecution('finnhubAnalysisTool', { input: inputData }); + logToolExecution('finnhubAnalysisTool', { input: input }); const apiKey = process.env.FINNHUB_API_KEY; @@ -655,7 +632,7 @@ export const finnhubAnalysisTool = createTool({ rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubAnalysisTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubAnalysisTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -666,9 +643,9 @@ export const finnhubAnalysisTool = createTool({ let url: string; const params = new URLSearchParams(); params.append('token', apiKey); - params.append('symbol', inputData.symbol); + params.append('symbol', input.symbol); - switch (inputData.function) { + switch (input.function) { case "RECOMMENDATION_TRENDS": { url = `https://finnhub.io/api/v1/stock/recommendation?${params.toString()}`; break; @@ -678,11 +655,11 @@ export const finnhubAnalysisTool = createTool({ break; } default: { - const message = `Unsupported function: ${inputData.function}`; + const message = `Unsupported function: ${input.function}`; rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubAnalysisTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubAnalysisTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -699,8 +676,8 @@ export const finnhubAnalysisTool = createTool({ }); logStepStart('finnhub-api-call', { - function: inputData.function, - symbol: inputData.symbol, + function: input.function, + symbol: input.symbol, url: url.replace(apiKey, '[REDACTED]') }); @@ -723,7 +700,7 @@ export const finnhubAnalysisTool = createTool({ const message = String(errorValue); rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); - logError('finnhubAnalysisTool', new Error(message), { function: inputData.function, symbol: inputData.symbol, apiError: message }); + logError('finnhubAnalysisTool', new Error(message), { function: input.function, symbol: input.symbol, apiError: message }); return { data: null, message @@ -734,15 +711,15 @@ export const finnhubAnalysisTool = createTool({ const result = { data, metadata: { - function: inputData.function, - symbol: inputData.symbol + function: input.function, + symbol: input.symbol }, message: undefined }; rootSpan.end(); - logToolExecution('finnhubAnalysisTool', { output: { function: inputData.function, symbol: inputData.symbol } }); + logToolExecution('finnhubAnalysisTool', { output: { function: input.function, symbol: input.symbol } }); return result; @@ -751,14 +728,38 @@ export const finnhubAnalysisTool = createTool({ rootSpan.recordException(error instanceof Error ? error : new Error(errorMessage)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); rootSpan.end(); - logError('finnhubAnalysisTool', error instanceof Error ? error : new Error(errorMessage), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubAnalysisTool', error instanceof Error ? error : new Error(errorMessage), { function: input.function, symbol: input.symbol }); return { data: null, message: errorMessage }; } - } -}); + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Finnhub analysis input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Finnhub analysis received complete input', { + toolCallId, + messageCount: messages.length, + function: input.function, + symbol: input.symbol, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Finnhub analysis completed', { + toolCallId, + toolName, + hasData: !!output.data, + hook: 'onOutput', + }) + }, +}) /** * Finnhub Technical Tool @@ -799,38 +800,15 @@ export const finnhubTechnicalTool = createTool({ }).optional(), message: z.string().optional() }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Finnhub technical input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Finnhub technical received complete input', { - toolCallId, - messageCount: messages.length, - function: input.function, - symbol: input.symbol, - resolution: input.resolution, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Finnhub technical completed', { - toolCallId, - toolName, - hasData: !!output.data, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { + const requestCtx = context?.requestContext as FinnhubRequestContext | undefined; const tracer = trace.getTracer('finnhub-tools'); const rootSpan = tracer.startSpan('finnhub-technical', { attributes: { - function: inputData.function, - symbol: inputData.symbol, - operation: 'finnhub-technical' + function: input.function, + symbol: input.symbol, + operation: 'finnhub-technical', + 'tool.requestContext.userId': requestCtx?.userId, } }); @@ -838,13 +816,13 @@ export const finnhubTechnicalTool = createTool({ type: 'data-tool-progress', data: { status: 'in-progress', - message: `Fetching technical data (${inputData.function}) for ${inputData.symbol}`, + message: `Fetching technical data (${input.function}) for ${input.symbol}`, stage: 'finnhub-technical' }, id: 'finnhub-technical' }); - logToolExecution('finnhubTechnicalTool', { input: inputData }); + logToolExecution('finnhubTechnicalTool', { input: input }); const apiKey = process.env.FINNHUB_API_KEY; @@ -853,7 +831,7 @@ export const finnhubTechnicalTool = createTool({ rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubTechnicalTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubTechnicalTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -864,27 +842,27 @@ export const finnhubTechnicalTool = createTool({ let url: string; const params = new URLSearchParams(); params.append('token', apiKey); - params.append('symbol', inputData.symbol); - params.append('resolution', inputData.resolution); + params.append('symbol', input.symbol); + params.append('resolution', input.resolution); - switch (inputData.function) { + switch (input.function) { case "TECHNICAL_INDICATOR": { - if (inputData.indicator === undefined || inputData.indicator === null || inputData.indicator.trim() === '' || - inputData.timeperiod === undefined || inputData.timeperiod === null || - inputData.series_type === undefined || inputData.series_type === null || inputData.series_type.trim() === '') { + if (input.indicator === undefined || input.indicator === null || input.indicator.trim() === '' || + input.timeperiod === undefined || input.timeperiod === null || + input.series_type === undefined || input.series_type === null || input.series_type.trim() === '') { const message = "TECHNICAL_INDICATOR function requires indicator, timeperiod, and series_type parameters"; rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubTechnicalTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubTechnicalTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message }; } - params.append('indicator', inputData.indicator); - params.append('timeperiod', inputData.timeperiod.toString()); - params.append('series_type', inputData.series_type); + params.append('indicator', input.indicator); + params.append('timeperiod', input.timeperiod.toString()); + params.append('series_type', input.series_type); url = `https://finnhub.io/api/v1/indicator?${params.toString()}`; break; } @@ -901,11 +879,11 @@ export const finnhubTechnicalTool = createTool({ break; } default: { - const message = `Unsupported function: ${inputData.function}`; + const message = `Unsupported function: ${input.function}`; rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubTechnicalTool', new Error(message), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubTechnicalTool', new Error(message), { function: input.function, symbol: input.symbol }); return { data: null, message @@ -922,8 +900,8 @@ export const finnhubTechnicalTool = createTool({ }); logStepStart('finnhub-api-call', { - function: inputData.function, - symbol: inputData.symbol, + function: input.function, + symbol: input.symbol, url: url.replace(apiKey, '[REDACTED]') }); @@ -946,7 +924,7 @@ export const finnhubTechnicalTool = createTool({ const message = String(errorValue); rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); - logError('finnhubTechnicalTool', new Error(message), { function: inputData.function, symbol: inputData.symbol, apiError: message }); + logError('finnhubTechnicalTool', new Error(message), { function: input.function, symbol: input.symbol, apiError: message }); return { data: null, message @@ -957,19 +935,19 @@ export const finnhubTechnicalTool = createTool({ const result = { data, metadata: { - function: inputData.function, - symbol: inputData.symbol, - resolution: inputData.resolution, - indicator: inputData.indicator, - timeperiod: inputData.timeperiod, - series_type: inputData.series_type + function: input.function, + symbol: input.symbol, + resolution: input.resolution, + indicator: input.indicator, + timeperiod: input.timeperiod, + series_type: input.series_type }, message: undefined }; rootSpan.end(); - logToolExecution('finnhubTechnicalTool', { output: { function: inputData.function, symbol: inputData.symbol } }); + logToolExecution('finnhubTechnicalTool', { output: { function: input.function, symbol: input.symbol } }); return result; @@ -978,14 +956,39 @@ export const finnhubTechnicalTool = createTool({ rootSpan.recordException(error instanceof Error ? error : new Error(errorMessage)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); rootSpan.end(); - logError('finnhubTechnicalTool', error instanceof Error ? error : new Error(errorMessage), { function: inputData.function, symbol: inputData.symbol }); + logError('finnhubTechnicalTool', error instanceof Error ? error : new Error(errorMessage), { function: input.function, symbol: input.symbol }); return { data: null, message: errorMessage }; } - } -}); + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Finnhub technical input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Finnhub technical received complete input', { + toolCallId, + messageCount: messages.length, + function: input.function, + symbol: input.symbol, + resolution: input.resolution, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Finnhub technical completed', { + toolCallId, + toolName, + hasData: !!output.data, + hook: 'onOutput', + }) + }, +}) /** * Finnhub Economic Tool @@ -1008,35 +1011,14 @@ export const finnhubEconomicTool = createTool({ }).optional(), message: z.string().optional() }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Finnhub economic input streaming started', { - toolCallId, - messageCount: messages.length, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Finnhub economic received complete input', { - toolCallId, - messageCount: messages.length, - economic_code: input.economic_code, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Finnhub economic completed', { - toolCallId, - toolName, - hasData: !!output.data, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { + const requestCtx = context?.requestContext as FinnhubRequestContext | undefined; const tracer = trace.getTracer('finnhub-tools'); const rootSpan = tracer.startSpan('finnhub-economic', { attributes: { - economic_code: inputData.economic_code, - operation: 'finnhub-economic' + economic_code: input.economic_code, + operation: 'finnhub-economic', + 'tool.requestContext.userId': requestCtx?.userId, } }); @@ -1044,13 +1026,13 @@ export const finnhubEconomicTool = createTool({ type: 'data-tool-progress', data: { status: 'in-progress', - message: `Fetching economic data for ${inputData.economic_code}`, + message: `Fetching economic data for ${input.economic_code}`, stage: 'finnhub-economic' }, id: 'finnhub-economic' }); - logToolExecution('finnhubEconomicTool', { input: inputData }); + logToolExecution('finnhubEconomicTool', { input: input }); const apiKey = process.env.FINNHUB_API_KEY; @@ -1059,7 +1041,7 @@ export const finnhubEconomicTool = createTool({ rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); rootSpan.end(); - logError('finnhubEconomicTool', new Error(message), { economic_code: inputData.economic_code }); + logError('finnhubEconomicTool', new Error(message), { economic_code: input.economic_code }); return { data: null, message @@ -1069,7 +1051,7 @@ export const finnhubEconomicTool = createTool({ try { const params = new URLSearchParams(); params.append('token', apiKey); - params.append('code', inputData.economic_code); + params.append('code', input.economic_code); const url = `https://finnhub.io/api/v1/economic?${params.toString()}`; @@ -1082,7 +1064,7 @@ export const finnhubEconomicTool = createTool({ }); logStepStart('finnhub-api-call', { - economic_code: inputData.economic_code, + economic_code: input.economic_code, url: url.replace(apiKey, '[REDACTED]') }); @@ -1105,7 +1087,7 @@ export const finnhubEconomicTool = createTool({ const message = String(errorValue); rootSpan.recordException(new Error(message)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message }); - logError('finnhubEconomicTool', new Error(message), { economic_code: inputData.economic_code, apiError: message }); + logError('finnhubEconomicTool', new Error(message), { economic_code: input.economic_code, apiError: message }); return { data: null, message @@ -1116,14 +1098,14 @@ export const finnhubEconomicTool = createTool({ const result = { data, metadata: { - economic_code: inputData.economic_code + economic_code: input.economic_code }, message: undefined }; rootSpan.end(); - logToolExecution('finnhubEconomicTool', { output: { economic_code: inputData.economic_code } }); + logToolExecution('finnhubEconomicTool', { output: { economic_code: input.economic_code } }); return result; @@ -1132,14 +1114,37 @@ export const finnhubEconomicTool = createTool({ rootSpan.recordException(error instanceof Error ? error : new Error(errorMessage)); rootSpan.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); rootSpan.end(); - logError('finnhubEconomicTool', error instanceof Error ? error : new Error(errorMessage), { economic_code: inputData.economic_code }); + logError('finnhubEconomicTool', error instanceof Error ? error : new Error(errorMessage), { economic_code: input.economic_code }); return { data: null, message: errorMessage }; } - } -}); + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Finnhub economic input streaming started', { + toolCallId, + messageCount: messages.length, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Finnhub economic received complete input', { + toolCallId, + messageCount: messages.length, + economic_code: input.economic_code, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Finnhub economic completed', { + toolCallId, + toolName, + hasData: !!output.data, + hook: 'onOutput', + }) + }, +}) export type FinnhubEconomicUITool = InferUITool; export type FinnhubTechnicalUITool = InferUITool; diff --git a/src/mastra/tools/git-local.tool.ts b/src/mastra/tools/git-local.tool.ts index 616edd01..166cfcc5 100644 --- a/src/mastra/tools/git-local.tool.ts +++ b/src/mastra/tools/git-local.tool.ts @@ -16,6 +16,14 @@ export interface GitToolContext extends RequestContext { timeout?: number } +const gitToolContextSchema = z.object({ + defaultBranch: z.string().optional(), + allowForce: z.boolean().optional(), + allowAmend: z.boolean().optional(), + maxCommits: z.number().optional(), + timeout: z.number().optional(), +}) + // Enhanced Git Status Tool with more detailed information export const gitStatusTool = createTool({ id: 'git:status', @@ -82,7 +90,7 @@ export const gitStatusTool = createTool({ }) try { - const cwd = inputData.repoPath || process.cwd() + const cwd = inputData.repoPath ?? process.cwd() const args = ['status', '--porcelain'] if (!inputData.porcelain) { @@ -355,7 +363,7 @@ export const gitDiffTool = createTool({ }) try { - const cwd = inputData.repoPath || process.cwd() + const cwd = inputData.repoPath ?? process.cwd() const args = ['diff'] if (inputData.target) { @@ -559,7 +567,7 @@ export const gitCommitTool = createTool({ }) try { - const cwd = inputData.repoPath || process.cwd() + const cwd = inputData.repoPath ?? process.cwd() // Validate amend permission if (inputData.amend && !allowAmend) { @@ -1094,7 +1102,7 @@ export const gitBranchTool = createTool({ }) try { - const cwd = inputData.repoPath || process.cwd() + const cwd = inputData.repoPath ?? process.cwd() let result: any = {} switch (inputData.operation) { @@ -1296,7 +1304,7 @@ export const gitBranchTool = createTool({ const mergeResult = await execa('git', args, { cwd, stdio: 'pipe', - timeout: timeout * 2, // Merges can take longer + timeout: (timeout ?? 30000) * 2, // Merges can take longer }) // Check for conflicts @@ -1335,7 +1343,7 @@ export const gitBranchTool = createTool({ await execa('git', ['rebase', inputData.targetBranch], { cwd, stdio: 'pipe', - timeout: timeout * 2, + timeout: (timeout ?? 30000) * 2, // Merges can take longer }) result = { diff --git a/src/mastra/tools/json-to-csv.tool.ts b/src/mastra/tools/json-to-csv.tool.ts index 2192b946..f8acf313 100644 --- a/src/mastra/tools/json-to-csv.tool.ts +++ b/src/mastra/tools/json-to-csv.tool.ts @@ -5,9 +5,11 @@ import { trace } from '@opentelemetry/api' import { z } from 'zod' import { log } from '../config/logger' -const csvToolContextSchema = z.object({ - maxRows: z.number().optional(), -}) +export interface JsonToCsvRequestContext extends RequestContext { + csvToolContext?: { + maxRows?: number + } +} export const jsonToCsvTool = createTool({ id: 'json-to-csv', @@ -33,54 +35,9 @@ export const jsonToCsvTool = createTool({ includeHeaders: true, }), }), - outputSchema: z.object({ - csv: z.string().describe('Generated CSV string'), - }), - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('JSON to CSV tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('JSON to CSV tool received input chunk', { - toolCallId, - inputTextDelta, - abortSignal: abortSignal?.aborted, - messageCount: messages.length, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - const options = input.options || {} - log.info('JSON to CSV received complete input', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - recordCount: input.data.length, - delimiter: options.delimiter, - includeHeaders: options.includeHeaders, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - const csvLines = output.csv.split('\n').length - log.info('JSON to CSV conversion completed', { - toolCallId, - toolName, - abortSignal: abortSignal?.aborted, - csvLines, - csvLength: output.csv.length, - hook: 'onOutput', - }) - }, - execute: async (inputData, context) => { + execute: async (input, context) => { const writer = context?.writer - const requestContext = context?.requestContext as RequestContext<{ - csvToolContext: unknown - }> + const requestContext = context?.requestContext as JsonToCsvRequestContext | undefined const abortSignal = context?.abortSignal // Check if operation was already cancelled @@ -92,7 +49,7 @@ export const jsonToCsvTool = createTool({ type: 'data-tool-progress', data: { status: 'in-progress', - message: `📊 Converting ${inputData.data.length} JSON records to CSV`, + message: `📊 Converting ${input.data.length} JSON records to CSV`, stage: 'json-to-csv', }, id: 'json-to-csv', @@ -102,22 +59,19 @@ export const jsonToCsvTool = createTool({ const rootSpan = tracer.startSpan('json-to-csv', { attributes: { 'tool.id': 'json-to-csv', - 'tool.input.recordCount': inputData.data.length, + 'tool.input.recordCount': input.data.length, }, }) try { - const { data, options } = inputData + const { data, options } = input if (data === undefined || data === null || data.length === 0) { rootSpan.end() return { csv: '' } } - const config = requestContext?.get('csvToolContext') - const { maxRows } = - config !== undefined - ? csvToolContextSchema.parse(config) - : { maxRows: undefined } + const config = requestContext?.csvToolContext + const maxRows = config?.maxRows if (maxRows !== undefined && data.length > maxRows) { throw new Error( @@ -209,6 +163,46 @@ export const jsonToCsvTool = createTool({ throw new Error(errorMessage) } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('JSON to CSV tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('JSON to CSV tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + const options = input.options || {} + log.info('JSON to CSV received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + recordCount: input.data.length, + delimiter: options.delimiter, + includeHeaders: options.includeHeaders, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + const csvLines = output.csv.split('\n').length + log.info('JSON to CSV conversion completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + csvLines, + csvLength: output.csv.length, + hook: 'onOutput', + }) + }, }) export type JsonToCsvUITool = InferUITool diff --git a/src/mastra/tools/random-generator.tool.ts b/src/mastra/tools/random-generator.tool.ts index 97e5a4fb..50d94bfc 100644 --- a/src/mastra/tools/random-generator.tool.ts +++ b/src/mastra/tools/random-generator.tool.ts @@ -2,16 +2,15 @@ import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' import { trace } from '@opentelemetry/api' import { z } from 'zod' -import { log, logError } from '../config/logger' +import { log } from '../config/logger' -// Define the Zod schema for the runtime context -const randomToolContextSchema = z.object({ - seed: z.number().optional(), - locale: z.string().default('en'), -}) +import type { RequestContext } from '@mastra/core/request-context' -// Infer the TypeScript type from the Zod schema -export type RandomToolContext = z.infer +// Define the interface for the runtime context +export interface RandomToolContext extends RequestContext { + seed?: number + locale?: string +} export const randomGeneratorTool = createTool({ id: 'random-generator', @@ -75,15 +74,14 @@ export const randomGeneratorTool = createTool({ message: z.string().optional(), }), execute: async (inputData, context) => { - const requestContext = context?.requestContext as RandomGeneratorRequestContext - const localeFromContext = requestContext?.locale ?? 'en' + const requestContext = context?.requestContext as RandomToolContext + const locale = + inputData.options?.locale ?? requestContext?.locale ?? 'en' const tracer = trace.getTracer('random-generator-tool', '1.0.0') const span = tracer.startSpan('random-generate') const writer = context?.writer - - await writer?.custom({ type: 'data-tool-progress', data: { @@ -95,17 +93,30 @@ export const randomGeneratorTool = createTool({ }) try { - let result: any + let result: + | string + | number + | boolean + | unknown[] + | Record + | null = null if (inputData.count === 1) { - result = generateRandomItem(inputData.type, inputData.options) + result = generateRandomItem(inputData.type, { + ...inputData.options, + locale, + }) } else { - result = [] + const items: unknown[] = [] for (let i = 0; i < inputData.count; i++) { - result.push( - generateRandomItem(inputData.type, inputData.options) + items.push( + generateRandomItem(inputData.type, { + ...inputData.options, + locale, + }) ) } + result = items } await writer?.custom({ @@ -150,10 +161,63 @@ export const randomGeneratorTool = createTool({ } } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Random generator tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Random generator tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Random generator tool received input', { + toolCallId, + messageCount: messages.length, + inputData: { type: input.type, count: input.count }, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Random generator tool completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + outputData: { + success: output.success, + type: output.type, + count: output.count, + }, + hook: 'onOutput', + }) + }, }) +interface RandomOptions { + min?: number + max?: number + length?: number + format?: string + includeSpecial?: boolean + locale?: string + itemType?: string + properties?: number +} + // Helper functions -function generateRandomItem(type: string, options?: any): any { +function generateRandomItem( + type: string, + options?: RandomOptions +): string | number | boolean | unknown[] | Record { switch (type) { case 'string': return generateRandomString( @@ -177,7 +241,7 @@ function generateRandomItem(type: string, options?: any): any { case 'array': return generateRandomArray( options?.length ?? 5, - options?.itemType ?? 'string' + (options?.itemType as string) ?? 'string' ) case 'object': return generateRandomObject(options?.properties ?? 3) @@ -352,7 +416,7 @@ function generateRandomAddress(locale = 'en'): string { return `${number} ${street}, ${city}` } -function generateRandomArray(length: number, itemType: string): any[] { +function generateRandomArray(length: number, itemType: string): unknown[] { const result = [] for (let i = 0; i < length; i++) { result.push(generateRandomItem(itemType)) @@ -360,8 +424,10 @@ function generateRandomArray(length: number, itemType: string): any[] { return result } -function generateRandomObject(properties: number): Record { - const result: Record = {} +function generateRandomObject( + properties: number +): Record { + const result: Record = {} for (let i = 0; i < properties; i++) { const key = `prop${i + 1}` const type = ['string', 'number', 'boolean'][ diff --git a/src/mastra/tools/web-scraper-tool.ts b/src/mastra/tools/web-scraper-tool.ts index 8a30d2c4..4c524b2f 100644 --- a/src/mastra/tools/web-scraper-tool.ts +++ b/src/mastra/tools/web-scraper-tool.ts @@ -27,6 +27,14 @@ import RE2 from 're2' import { z } from 'zod' import { log } from '../config/logger' +import type { RequestContext } from '@mastra/core/request-context' + +export interface ScraperToolContext extends RequestContext { + allowedDomains?: string[] + userAgent?: string + timeout?: number +} + // Centralized data directory constant for consistency const DATA_DIR = path.resolve(process.cwd(), './data') @@ -647,46 +655,8 @@ export const webScraperTool = createTool({ 'Extracts structured data from web pages using JSDOM and Cheerio with enhanced security and error handling.', inputSchema: webScraperInputSchema, outputSchema: webScraperOutputSchema, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Web scraper tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Web scraper tool received input chunk', { - toolCallId, - inputTextDelta, - abortSignal: abortSignal?.aborted, - messageCount: messages.length, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Web scraper received complete input', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - url: input.url, - selector: input.selector || 'none', - saveMarkdown: input.storage?.saveMarkdown || false, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Web scraper completed', { - toolCallId, - toolName, - abortSignal: abortSignal?.aborted, - status: output.status, - elementsExtracted: output.content.extractedData.length, - savedFilePath: output.storage?.savedFilePath, - hook: 'onOutput', - }) - }, execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext const writer = context?.writer const abortSignal = context?.abortSignal @@ -695,7 +665,8 @@ export const webScraperTool = createTool({ throw new Error('Web scraping cancelled') } - const allowedDomains = ValidationUtils.getAllowedDomains() + const allowedDomains = + requestContext?.allowedDomains ?? ValidationUtils.getAllowedDomains() if (!ValidationUtils.isUrlAllowed(inputData.url, allowedDomains)) { throw new ScrapingError( `Domain is not allowlisted for scraping: ${inputData.url}`, @@ -1091,7 +1062,7 @@ export const webScraperTool = createTool({ }) // Check for cancellation before crawling - if (abortSignal?.aborted) { + if (abortSignal?.aborted === true) { throw new Error('Web scraping cancelled during crawling') } @@ -1264,6 +1235,45 @@ export const webScraperTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Web scraper tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Web scraper tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Web scraper received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + url: input.url, + selector: input.selector || 'none', + saveMarkdown: input.storage?.saveMarkdown || false, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Web scraper completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + status: output.status, + elementsExtracted: output.content.extractedData.length, + savedFilePath: output.storage?.savedFilePath, + hook: 'onOutput', + }) + }, }) export type WebScraperUITool = InferUITool @@ -1326,6 +1336,7 @@ export const batchWebScraperTool = createTool({ inputSchema: batchWebScraperInputSchema, outputSchema: batchWebScraperOutputSchema, execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext const writer = context?.writer await writer?.custom({ @@ -1547,6 +1558,44 @@ export const batchWebScraperTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Batch web scraper tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Batch web scraper tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Batch web scraper received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + urlCount: input.urls.length, + selector: input.selector ?? 'none', + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Batch web scraper completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + totalProcessed: output.totalProcessed, + successful: output.successful, + failed: output.failed, + hook: 'onOutput', + }) + }, }) export type BatchWebScraperUITool = InferUITool @@ -1607,6 +1656,7 @@ export const siteMapExtractorTool = createTool({ inputSchema: siteMapExtractorInputSchema, outputSchema: siteMapExtractorOutputSchema, execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext const writer = context?.writer await writer?.custom({ @@ -1840,6 +1890,43 @@ export const siteMapExtractorTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Site map extractor tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Site map extractor tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Site map extractor received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + url: input.url, + maxDepth: input.maxDepth ?? 2, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Site map extractor completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + totalPages: output.totalPages, + savedFilePath: output.savedFilePath, + hook: 'onOutput', + }) + }, }) export type SiteMapExtractorUITool = InferUITool @@ -1893,6 +1980,7 @@ export const linkExtractorTool = createTool({ inputSchema: linkExtractorInputSchema, outputSchema: linkExtractorOutputSchema, execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext const writer = context?.writer await writer?.custom({ @@ -2103,6 +2191,41 @@ export const linkExtractorTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Link extractor tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Link extractor tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Link extractor received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + url: input.url, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Link extractor completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + linkCount: output.links.length, + hook: 'onOutput', + }) + }, }) export type LinkExtractorUITool = InferUITool // ===== HTML TO MARKDOWN CONVERTER TOOL ===== @@ -2260,6 +2383,41 @@ export const htmlToMarkdownTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('HTML to Markdown tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('HTML to Markdown tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('HTML to Markdown received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + htmlLength: input.html.length, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('HTML to Markdown completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + markdownLength: output.markdown.length, + hook: 'onOutput', + }) + }, }) export type HtmlToMarkdownUITool = InferUITool @@ -2466,6 +2624,41 @@ export const listScrapedContentTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('List scraped content tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('List scraped content tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('List scraped content received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + pattern: input.pattern, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('List scraped content completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + fileCount: output.files.length, + hook: 'onOutput', + }) + }, }) export type ListScrapedContentUITool = InferUITool< @@ -2622,6 +2815,41 @@ export const contentCleanerTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Content cleaner tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Content cleaner tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Content cleaner received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + htmlLength: input.html.length, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Content cleaner completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + cleanedLength: output.cleanedHtml.length, + hook: 'onOutput', + }) + }, }) export type ContentCleanerUITool = InferUITool @@ -2784,6 +3012,41 @@ export const apiDataFetcherTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('API data fetcher tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('API data fetcher tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('API data fetcher received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + url: input.url, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('API data fetcher completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + status: output.status, + hook: 'onOutput', + }) + }, }) // ===== SCRAPING SCHEDULER TOOL ===== @@ -2885,6 +3148,41 @@ export const scrapingSchedulerTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Scraping scheduler tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Scraping scheduler tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Scraping scheduler received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + urlCount: input.urls.length, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Scraping scheduler completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + jobId: output.jobId, + hook: 'onOutput', + }) + }, }) // ===== EXPORT AND INTEGRATION TOOL ===== @@ -3048,4 +3346,41 @@ export const dataExporterTool = createTool({ throw error } }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Data exporter tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Data exporter tool received input chunk', { + toolCallId, + inputTextDelta, + abortSignal: abortSignal?.aborted, + messageCount: messages.length, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Data exporter received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + dataCount: input.data.length, + format: input.format, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Data exporter completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + success: output.success, + exportedCount: output.exportedCount, + hook: 'onOutput', + }) + }, })