diff --git a/docs/adding-custom-metadata.md b/docs/adding-custom-metadata.md new file mode 100644 index 00000000..098c197f --- /dev/null +++ b/docs/adding-custom-metadata.md @@ -0,0 +1,116 @@ +# Adding Custom Metadata to Traces + +## Automatic Metadata from RequestContext + +Instead of manually adding metadata to each span, you can configure Mastra to automatically extract values from RequestContext and attach them as metadata to all spans in a trace. + +### Configuration-Level Extraction + +Define which RequestContext keys to extract in your tracing configuration: + +```typescript +export const mastra = new Mastra({ + observability: new Observability({ + configs: { + default: { + serviceName: "my-service", + requestContextKeys: ["userId", "environment", "tenantId"], + exporters: [new DefaultExporter()], + }, + }, + }), +}); +``` + +### Per-Request Additions + +You can add trace-specific keys using `tracingOptions.requestContextKeys`: + +```typescript +const requestContext = new RequestContext(); +requestContext.set("userId", "user-123"); +requestContext.set("environment", "production"); +requestContext.set("experimentId", "exp-789"); + +const result = await agent.generate({ + messages: [{ role: "user", content: "Hello" }], + requestContext, + tracingOptions: { + requestContextKeys: ["experimentId"], // Adds to configured keys + }, +}); +``` + +### Nested Value Extraction + +Use dot notation to extract nested values from RequestContext: + +```typescript +export const mastra = new Mastra({ + observability: new Observability({ + configs: { + default: { + requestContextKeys: ["user.id", "session.data.experimentId"], + exporters: [new DefaultExporter()], + }, + }, + }), +}); +``` + +## Adding Tags to Traces + +Tags are string labels that help you categorize and filter traces. Use `tracingOptions.tags`: + +```typescript +const result = await agent.generate({ + messages: [{ role: "user", content: "Hello" }], + tracingOptions: { + tags: ["production", "experiment-v2", "user-request"], + }, +}); +``` + +## Manual Metadata Addition + +You can add metadata to any span using the tracing context: + +```typescript +execute: async ({ inputData, tracingContext }) => { + const startTime = Date.now(); + const response = await fetch(inputData.endpoint); + + // Add custom metadata to the current span + tracingContext.currentSpan?.update({ + metadata: { + apiStatusCode: response.status, + endpoint: inputData.endpoint, + responseTimeMs: Date.now() - startTime, + userTier: inputData.userTier, + region: process.env.AWS_REGION, + }, + }); + + return await response.json(); +}; +``` + +## Child Spans and Metadata Extraction + +When creating child spans within tools or workflow steps, you can pass the `requestContext` parameter to enable metadata extraction: + +```typescript +execute: async ({ tracingContext, requestContext }) => { + // Create child span WITH requestContext - gets metadata extraction + const dbSpan = tracingContext.currentSpan?.createChildSpan({ + type: "generic", + name: "database-query", + requestContext, // Pass to enable metadata extraction + }); + + const results = await db.query("SELECT * FROM users"); + dbSpan?.end({ output: results }); + + return results; +}; +``` \ No newline at end of file diff --git a/docs/creating-child-spans.md b/docs/creating-child-spans.md new file mode 100644 index 00000000..898598a1 --- /dev/null +++ b/docs/creating-child-spans.md @@ -0,0 +1,75 @@ +# Creating Child Spans + +Child spans allow you to track fine-grained operations within your workflow steps or tools. They provide visibility into sub-operations like database queries, API calls, file operations, or complex calculations. This hierarchical structure helps you identify performance bottlenecks and understand the exact sequence of operations. + +## Creating Child Spans + +Create child spans inside a tool call or workflow step to track specific operations: + +```typescript +execute: async ({ inputData, tracingContext }) => { + // Create another child span for the main database operation + const querySpan = tracingContext.currentSpan?.createChildSpan({ + type: "generic", + name: "database-query", + input: { query: inputData.query }, + metadata: { database: "production" }, + }); + + try { + const results = await db.query(inputData.query); + querySpan?.end({ + output: results.data, + metadata: { + rowsReturned: results.length, + queryTimeMs: results.executionTime, + cacheHit: results.fromCache, + }, + }); + return results; + } catch (error) { + querySpan?.error({ + error, + metadata: { retryable: isRetryableError(error) }, + }); + throw error; + } +}; +``` + +Child spans automatically inherit the trace context from their parent, maintaining the relationship hierarchy in your observability platform. + +## Child Span Options + +When creating child spans, you can specify: + +- `type`: The span type ("generic", "llm", "tool", etc.) +- `name`: A descriptive name for the operation +- `input`: Input data for the operation +- `metadata`: Additional context information +- `requestContext`: For automatic metadata extraction (optional) + +## Error Handling in Child Spans + +Use the `error()` method to record errors in child spans: + +```typescript +querySpan?.error({ + error, + metadata: { retryable: isRetryableError(error) }, +}); +``` + +## Ending Child Spans + +Always end child spans with the `end()` method, providing output data and final metadata: + +```typescript +querySpan?.end({ + output: results, + metadata: { + rowsReturned: results.length, + executionTime: Date.now() - startTime, + }, +}); +``` \ No newline at end of file diff --git a/docs/opentelemetry-semantic-conventions.md b/docs/opentelemetry-semantic-conventions.md new file mode 100644 index 00000000..f5a8ed41 --- /dev/null +++ b/docs/opentelemetry-semantic-conventions.md @@ -0,0 +1,22 @@ +# OpenTelemetry Semantic Conventions + +## Span Naming + +• LLM Operations: `chat {model}` +• Tool Execution: `execute_tool {tool_name}` +• Agent Runs: `invoke_agent {agent_id}` +• Workflow Runs: `invoke_workflow {workflow_id}` + +## Key Attributes + +• `gen_ai.operation.name` - Operation type (chat, tool.execute, etc.) +• `gen_ai.provider.name` - AI provider (openai, anthropic, etc.) +• `gen_ai.request.model` - Model identifier +• `gen_ai.input.messages` - Chat history provided to the model +• `gen_ai.output.messages` - Messages returned by the model +• `gen_ai.usage.input_tokens` - Number of input tokens +• `gen_ai.usage.output_tokens` - Number of output tokens +• `gen_ai.request.temperature` - Sampling temperature +• `gen_ai.response.finish_reasons` - Completion reasons + +The exporter follows [OpenTelemetry Semantic Conventions for GenAI v1.38.0](https://github.com/open-telemetry/semantic-conventions/tree/v1.38.0/docs/gen-ai), ensuring compatibility with observability platforms. \ No newline at end of file diff --git a/docs/otel-semantic-conventions.html b/docs/otel-semantic-conventions.html new file mode 100644 index 00000000..4514893e --- /dev/null +++ b/docs/otel-semantic-conventions.html @@ -0,0 +1,122 @@ +OpenTelemetry Exporter | Tracing | Observability | Mastra Docs v1 Beta + + +
Skip to main content

OpenTelemetry Exporter

+

The OpenTelemetry (OTEL) exporter sends your traces to any OTEL-compatible observability platform using standardized OpenTelemetry Semantic Conventions for GenAI. This ensures broad compatibility with platforms like Datadog, New Relic, SigNoz, MLflow, Dash0, Traceloop, Laminar, and more.

+
Looking for bidirectional OTEL integration?

If you have existing OpenTelemetry instrumentation and want Mastra traces to inherit context from active OTEL spans, see the OpenTelemetry Bridge instead.

+

Installation
Direct link to Installation

+

Each provider requires specific protocol packages. Install the base exporter plus the protocol package for your provider:

+

For HTTP/Protobuf Providers (SigNoz, New Relic, Laminar, MLflow)
Direct link to For HTTP/Protobuf Providers (SigNoz, New Relic, Laminar, MLflow)

+
npm install @mastra/otel-exporter@beta @opentelemetry/exporter-trace-otlp-proto
+

For gRPC Providers (Dash0, Datadog)
Direct link to For gRPC Providers (Dash0, Datadog)

+
npm install @mastra/otel-exporter@beta @opentelemetry/exporter-trace-otlp-grpc @grpc/grpc-js
+

For HTTP/JSON Providers (Traceloop)
Direct link to For HTTP/JSON Providers (Traceloop)

+
npm install @mastra/otel-exporter@beta @opentelemetry/exporter-trace-otlp-http
+

Environment Variables
Direct link to Environment Variables

+

All providers support zero-config setup via environment variables. Set the appropriate variables and the exporter will automatically use them:

+
ProviderEnvironment Variables
Dash0DASH0_API_KEY (required), DASH0_ENDPOINT (required), DASH0_DATASET (optional)
SigNozSIGNOZ_API_KEY (required), SIGNOZ_REGION (optional), SIGNOZ_ENDPOINT (optional)
New RelicNEW_RELIC_LICENSE_KEY (required), NEW_RELIC_ENDPOINT (optional)
TraceloopTRACELOOP_API_KEY (required), TRACELOOP_DESTINATION_ID, TRACELOOP_ENDPOINT (optional)
LaminarLMNR_PROJECT_API_KEY (required), LAMINAR_ENDPOINT, LAMINAR_TEAM_ID (optional)
+

Provider Configurations
Direct link to Provider Configurations

+

MLflow
Direct link to MLflow

+

MLflow supports native Mastra tracing through its OTLP endpoint at /v1/traces. Use the custom provider with HTTP/Protobuf and include the experiment header so traces land in the correct MLflow experiment:

+
src/mastra/index.ts
new OtelExporter({
  provider: {
    custom: {
      endpoint: `${process.env.MLFLOW_TRACKING_URI}/v1/traces`,
      protocol: "http/protobuf",
      headers: {
        "x-mlflow-experiment-id": process.env.MLFLOW_EXPERIMENT_ID,
      },
    },
  },
})
+

Dash0
Direct link to Dash0

+

Dash0 provides real-time observability with automatic insights.

+

Zero-Config Setup
Direct link to Zero-Config Setup

+

Set environment variables and use the exporter with an empty config:

+
.env
# Required
DASH0_API_KEY=your-api-key
DASH0_ENDPOINT=ingress.us-west-2.aws.dash0.com:4317

# Optional
DASH0_DATASET=production
+
src/mastra/index.ts
import { Mastra } from "@mastra/core";
import { Observability } from "@mastra/observability";
import { OtelExporter } from "@mastra/otel-exporter";

export const mastra = new Mastra({
  observability: new Observability({
    configs: {
      otel: {
        serviceName: "my-service",
        exporters: [new OtelExporter({ provider: { dash0: {} } })],
      },
    },
  }),
});
+

Explicit Configuration
Direct link to Explicit Configuration

+
src/mastra/index.ts
import { Mastra } from "@mastra/core";
import { Observability } from "@mastra/observability";
import { OtelExporter } from "@mastra/otel-exporter";

export const mastra = new Mastra({
  observability: new Observability({
    configs: {
      otel: {
        serviceName: "my-service",
        exporters: [
          new OtelExporter({
            provider: {
              dash0: {
                apiKey: process.env.DASH0_API_KEY,
                endpoint: process.env.DASH0_ENDPOINT, // e.g., 'ingress.us-west-2.aws.dash0.com:4317'
                dataset: "production", // Optional dataset name
              },
            },
            resourceAttributes: {
              // Optional OpenTelemetry Resource Attributes for the trace
              ["deployment.environment"]: "dev",
            },
          }),
        ],
      },
    },
  }),
});
+
info

Get your Dash0 endpoint from your dashboard. It should be in the format ingress.{region}.aws.dash0.com:4317.

+

SigNoz
Direct link to SigNoz

+

SigNoz is an open-source APM alternative with built-in Tracing support.

+

Zero-Config Setup
Direct link to Zero-Config Setup

+
.env
# Required
SIGNOZ_API_KEY=your-api-key

# Optional
SIGNOZ_REGION=us  # 'us' | 'eu' | 'in'
SIGNOZ_ENDPOINT=https://my-signoz.example.com  # For self-hosted
+
src/mastra/index.ts
new OtelExporter({ provider: { signoz: {} } })
+

Explicit Configuration
Direct link to Explicit Configuration

+
src/mastra/index.ts
new OtelExporter({
  provider: {
    signoz: {
      apiKey: process.env.SIGNOZ_API_KEY,
      region: "us", // 'us' | 'eu' | 'in'
      // endpoint: 'https://my-signoz.example.com', // For self-hosted
    },
  },
});
+

New Relic
Direct link to New Relic

+

New Relic provides comprehensive observability with AI monitoring capabilities.

+

Zero-Config Setup
Direct link to Zero-Config Setup

+
.env
# Required
NEW_RELIC_LICENSE_KEY=your-license-key

# Optional
NEW_RELIC_ENDPOINT=https://otlp.eu01.nr-data.net  # For EU region
+
src/mastra/index.ts
new OtelExporter({ provider: { newrelic: {} } })
+

Explicit Configuration
Direct link to Explicit Configuration

+
src/mastra/index.ts
new OtelExporter({
  provider: {
    newrelic: {
      apiKey: process.env.NEW_RELIC_LICENSE_KEY,
      // endpoint: 'https://otlp.eu01.nr-data.net', // For EU region
    },
  },
});
+

Traceloop
Direct link to Traceloop

+

Traceloop specializes in LLM observability with automatic prompt tracking.

+

Zero-Config Setup
Direct link to Zero-Config Setup

+
.env
# Required
TRACELOOP_API_KEY=your-api-key

# Optional
TRACELOOP_DESTINATION_ID=my-destination
TRACELOOP_ENDPOINT=https://custom.traceloop.com
+
src/mastra/index.ts
new OtelExporter({ provider: { traceloop: {} } })
+

Explicit Configuration
Direct link to Explicit Configuration

+
src/mastra/index.ts
new OtelExporter({
  provider: {
    traceloop: {
      apiKey: process.env.TRACELOOP_API_KEY,
      destinationId: "my-destination", // Optional
    },
  },
});
+

Laminar
Direct link to Laminar

+

Laminar provides specialized LLM observability and analytics.

+

Zero-Config Setup
Direct link to Zero-Config Setup

+
.env
# Required
LMNR_PROJECT_API_KEY=your-api-key

# Optional
LAMINAR_ENDPOINT=https://api.lmnr.ai/v1/traces
LAMINAR_TEAM_ID=your-team-id  # For backwards compatibility
+
src/mastra/index.ts
new OtelExporter({ provider: { laminar: {} } })
+

Explicit Configuration
Direct link to Explicit Configuration

+
src/mastra/index.ts
new OtelExporter({
  provider: {
    laminar: {
      apiKey: process.env.LMNR_PROJECT_API_KEY,
      // teamId: process.env.LAMINAR_TEAM_ID, // Optional, for backwards compatibility
    },
  },
});
+

Datadog
Direct link to Datadog

+

Datadog APM provides application performance monitoring with distributed tracing. To send traces to Datadog via OTLP, you need the Datadog Agent running with OTLP ingestion enabled.

+

Datadog uses gRPC for OTLP ingestion, which requires explicit imports and bundler configuration to work correctly:

+
src/mastra/index.ts
// Explicitly import gRPC dependencies for the bundler
import "@grpc/grpc-js";
import "@opentelemetry/exporter-trace-otlp-grpc";
import { Mastra } from "@mastra/core";
import { Observability } from "@mastra/observability";
import { OtelExporter, type ExportProtocol } from "@mastra/otel-exporter";

export const mastra = new Mastra({
  // Add grpc-js to externals so it's handled at runtime
  bundler: {
    externals: ["@grpc/grpc-js"],
  },

  observability: new Observability({
    configs: {
      default: {
        serviceName: "my-service",
        exporters: [
          new OtelExporter({
            provider: {
              custom: {
                endpoint:
                  process.env.OTEL_EXPORTER_OTLP_ENDPOINT ||
                  "http://localhost:4317",
                protocol: (process.env.OTEL_EXPORTER_OTLP_PROTOCOL ||
                  "grpc") as ExportProtocol,
                headers: {},
              },
            },
          }),
        ],
      },
    },
  }),
});
+
info

The Datadog Agent must be configured with OTLP ingestion enabled. Add the following to your datadog.yaml:

otlp_config:
  receiver:
    protocols:
      grpc:
        endpoint: 0.0.0.0:4317

The default OTLP endpoint is http://localhost:4317 when running the Datadog Agent locally.

+
warning

The explicit imports of @grpc/grpc-js and @opentelemetry/exporter-trace-otlp-grpc at the top of the file, along with the bundler.externals configuration, are required for the gRPC transport to work correctly. Without these, you may encounter connection issues.

+

Custom/Generic OTEL Endpoints
Direct link to Custom/Generic OTEL Endpoints

+

For other OTEL-compatible platforms or custom collectors:

+
src/mastra/index.ts
new OtelExporter({
  provider: {
    custom: {
      endpoint: "https://your-collector.example.com/v1/traces",
      protocol: "http/protobuf", // 'http/json' | 'http/protobuf' | 'grpc'
      headers: {
        "x-api-key": process.env.API_KEY,
      },
    },
  },
});
+

Configuration Options
Direct link to Configuration Options

+

Complete Configuration
Direct link to Complete Configuration

+
new OtelExporter({
  // Provider configuration (required)
  provider: {
    // Use one of: dash0, signoz, newrelic, traceloop, laminar, custom
  },

  // Export configuration
  timeout: 30000, // Export timeout in milliseconds
  batchSize: 100, // Number of spans per batch

  // Debug options
  logLevel: "info", // 'debug' | 'info' | 'warn' | 'error'
});
+

OpenTelemetry Semantic Conventions
Direct link to OpenTelemetry Semantic Conventions

+

The exporter follows OpenTelemetry Semantic Conventions for GenAI v1.38.0, ensuring compatibility with observability platforms:

+

Span Naming
Direct link to Span Naming

+
    +
  • LLM Operations: chat {model}
    • +
  • Tool Execution: execute_tool {tool_name}
    • +
  • Agent Runs: invoke_agent {agent_id}
    • +
  • Workflow Runs: invoke_workflow {workflow_id}
    • +
+

Key Attributes
Direct link to Key Attributes

+
    +
  • gen_ai.operation.name - Operation type (chat, tool.execute, etc.)
    • +
  • gen_ai.provider.name - AI provider (openai, anthropic, etc.)
    • +
  • gen_ai.request.model - Model identifier
    • +
  • gen_ai.input.messages - Chat history provided to the model
    • +
  • gen_ai.output.messages - Messages returned by the model
    • +
  • gen_ai.usage.input_tokens - Number of input tokens
    • +
  • gen_ai.usage.output_tokens - Number of output tokens
    • +
  • gen_ai.request.temperature - Sampling temperature
    • +
  • gen_ai.response.finish_reasons - Completion reasons
    • +
+

Protocol Selection Guide
Direct link to Protocol Selection Guide

+

Choose the right protocol package based on your provider:

+
ProviderProtocolRequired Package
Dash0gRPC@opentelemetry/exporter-trace-otlp-grpc
DatadoggRPC@opentelemetry/exporter-trace-otlp-grpc
SigNozHTTP/Protobuf@opentelemetry/exporter-trace-otlp-proto
New RelicHTTP/Protobuf@opentelemetry/exporter-trace-otlp-proto
TraceloopHTTP/JSON@opentelemetry/exporter-trace-otlp-http
LaminarHTTP/Protobuf@opentelemetry/exporter-trace-otlp-proto
CustomVariesDepends on your collector
+
warning

Make sure to install the correct protocol package for your provider. The exporter will provide a helpful error message if the wrong package is installed.

+

Troubleshooting
Direct link to Troubleshooting

+

Missing Dependency Error
Direct link to Missing Dependency Error

+

If you see an error like:

+
HTTP/Protobuf exporter is not installed (required for signoz).
To use HTTP/Protobuf export, install the required package:
  npm install @opentelemetry/exporter-trace-otlp-proto
+

Install the suggested package for your provider.

+

Common Issues
Direct link to Common Issues

+
    +
  1. Wrong protocol package: Verify you installed the correct exporter for your provider
    2. +
  2. Invalid endpoint: Check endpoint format matches provider requirements
    3. +
  3. Authentication failures: Verify API keys and headers are correct
    4. +
+

Using Tags
Direct link to Using Tags

+

Tags help you categorize and filter traces in your observability platform. Add tags when executing agents or workflows:

+
const result = await agent.generate({
  messages: [{ role: "user", content: "Hello" }],
  tracingOptions: {
    tags: ["production", "experiment-v2", "user-request"],
  },
});
+

Tags are exported as a JSON string in the mastra.tags span attribute for broad backend compatibility. Common use cases include:

+
    +
  • Environment labels: "production", "staging"
    • +
  • Experiment tracking: "experiment-v1", "control-group"
    • +
  • Priority levels: "priority-high", "batch-job"
    • +
+ +
\ No newline at end of file diff --git a/docs/tracing-configuration-reference.md b/docs/tracing-configuration-reference.md new file mode 100644 index 00000000..272de19f --- /dev/null +++ b/docs/tracing-configuration-reference.md @@ -0,0 +1,128 @@ +# Mastra Tracing Configuration Reference + +## ObservabilityRegistryConfig + +```typescript +interface ObservabilityRegistryConfig { + default?: { enabled?: boolean }; + configs?: Record | ObservabilityInstance>; + configSelector?: ConfigSelector; +} +``` + +## ObservabilityInstanceConfig + +```typescript +interface ObservabilityInstanceConfig { + name: string; + serviceName: string; + sampling?: SamplingStrategy; + exporters?: ObservabilityExporter[]; + spanOutputProcessors?: SpanOutputProcessor[]; + includeInternalSpans?: boolean; + requestContextKeys?: string[]; + serializationOptions?: SerializationOptions; +} +``` + +## SerializationOptions + +Options for controlling how span data is serialized before export. + +```typescript +interface SerializationOptions { + maxStringLength?: number; + maxDepth?: number; + maxArrayLength?: number; + maxObjectKeys?: number; +} +``` + +## SamplingStrategy + +```typescript +type SamplingStrategy = + | { type: "always" } + | { type: "never" } + | { type: "ratio"; probability: number } + | { type: "custom"; sampler: (options?: TracingOptions) => boolean }; +``` + +## ConfigSelector + +```typescript +type ConfigSelector = ( + options: ConfigSelectorOptions, + availableConfigs: ReadonlyMap, +) => string | undefined; +``` + +## ConfigSelectorOptions + +```typescript +interface ConfigSelectorOptions { + requestContext?: RequestContext; +} +``` + +## Registry Methods + +### registerInstance +```typescript +registerInstance(name: string, instance: ObservabilityInstance, isDefault?: boolean): void; +``` +Registers an observability instance in the registry. + +### getInstance +```typescript +getInstance(name: string): ObservabilityInstance | undefined; +``` +Retrieves an observability instance by name. + +### getDefaultInstance +```typescript +getDefaultInstance(): ObservabilityInstance | undefined; +``` +Returns the default observability instance. + +### getSelectedInstance +```typescript +getSelectedInstance(options: ConfigSelectorOptions): ObservabilityInstance | undefined; +``` +Returns the observability instance selected by the config selector or default. + +### listInstances +```typescript +listInstances(): ReadonlyMap; +``` +Returns all registered observability instances. + +### hasInstance +```typescript +hasInstance(name: string): boolean; +``` +Checks if an observability instance exists. + +### setConfigSelector +```typescript +setConfigSelector(selector: ConfigSelector): void; +``` +Sets the config selector function. + +### unregisterInstance +```typescript +unregisterInstance(name: string): boolean; +``` +Removes an observability instance from the registry. + +### clear +```typescript +clear(): void; +``` +Clears all instances without shutdown. + +### shutdown +```typescript +async shutdown(): Promise; +``` +Shuts down all observability instances and clears the registry. \ No newline at end of file diff --git a/docs/tracing-instances-reference.md b/docs/tracing-instances-reference.md new file mode 100644 index 00000000..72c4c9fd --- /dev/null +++ b/docs/tracing-instances-reference.md @@ -0,0 +1,92 @@ +# Mastra Tracing Instances Reference + +## DefaultObservabilityInstance + +Default implementation of the ObservabilityInstance interface. + +### Constructor + +```typescript +new DefaultObservabilityInstance(config: ObservabilityInstanceConfig) +``` + +Creates a new DefaultObservabilityInstance with the specified configuration. + +### Properties + +Inherits all properties and methods from BaseObservabilityInstance. + +## BaseObservabilityInstance + +Base class for custom ObservabilityInstance implementations. + +### Methods + +#### getConfig + +```typescript +getConfig(): Readonly> +``` + +Returns the current observability configuration. + +#### getExporters + +```typescript +getExporters(): readonly ObservabilityExporter[] +``` + +Returns all configured exporters. + +#### getSpanOutputProcessors + +```typescript +getSpanOutputProcessors(): readonly SpanOutputProcessor[] +``` + +Returns all configured span output processors. + +#### getLogger + +```typescript +getLogger(): IMastraLogger +``` + +Returns the logger instance for exporters and other components. + +#### startSpan + +```typescript +startSpan(options: StartSpanOptions): Span +``` + +Start a new span of a specific SpanType. Creates the root span of a trace if no parent is provided. + +#### shutdown + +```typescript +async shutdown(): Promise +``` + +Shuts down all exporters and processors, cleaning up resources. + +## Custom Implementation + +To create a custom ObservabilityInstance implementation, extend BaseObservabilityInstance: + +```typescript +class CustomObservabilityInstance extends BaseObservabilityInstance { + constructor(config: ObservabilityInstanceConfig) { + super(config); + // Custom initialization + } + + // Override methods as needed + startSpan( + options: StartSpanOptions, + ): Span { + // Custom span creation logic + return super.startSpan(options); + } +} +``` \ No newline at end of file diff --git a/docs/tracing-interfaces-reference.md b/docs/tracing-interfaces-reference.md new file mode 100644 index 00000000..da56c8d4 --- /dev/null +++ b/docs/tracing-interfaces-reference.md @@ -0,0 +1,304 @@ +# Mastra Tracing Interfaces Reference + +## Core Interfaces + +### ObservabilityInstance + +Primary interface for observability. + +```typescript +interface ObservabilityInstance { + /** Get current configuration */ + getConfig(): Readonly>; + /** Get all exporters */ + getExporters(): readonly ObservabilityExporter[]; + /** Get all span output processors */ + getSpanOutputProcessors(): readonly SpanOutputProcessor[]; + /** Get the logger instance (for exporters and other components) */ + getLogger(): IMastraLogger; + /** Start a new span of a specific SpanType */ + startSpan( + options: StartSpanOptions, + ): Span; + /** Shutdown observability and clean up resources */ + shutdown(): Promise; +} +``` + +### Span + +Span interface, used internally for tracing. + +```typescript +interface Span { + readonly id: string; + readonly traceId: string; + readonly type: TType; + readonly name: string; + /** Is an internal span? (spans internal to the operation of mastra) */ + isInternal: boolean; + /** Parent span reference (undefined for root spans) */ + parent?: AnySpan; + /** Pointer to the ObservabilityInstance instance */ + observabilityInstance: ObservabilityInstance; + attributes?: SpanTypeMap[TType]; + metadata?: Record; + input?: any; + output?: any; + errorInfo?: any; + /** Tags for categorizing traces (only present on root spans) */ + tags?: string[]; + /** End the span */ + end(options?: EndSpanOptions): void; + /** Record an error for the span, optionally end the span as well */ + error(options: ErrorSpanOptions): void; + /** Update span attributes */ + update(options: UpdateSpanOptions): void; + /** Create child span - can be any span type independent of parent */ + createChildSpan( + options: ChildSpanOptions, + ): Span; + /** Create event span - can be any span type independent of parent */ + createEventSpan( + options: ChildEventOptions, + ): Span; + /** Returns TRUE if the span is the root span of a trace */ + get isRootSpan(): boolean; + /** Returns TRUE if the span is a valid span (not a NO-OP Span) */ + get isValid(): boolean; +} +``` + +### ObservabilityExporter + +Interface for observability exporters. + +```typescript +interface ObservabilityExporter { + /** Exporter name */ + name: string; + /** Initialize exporter with tracing configuration and/or access to Mastra */ + init?(options: InitExporterOptions): void; + /** Export tracing events */ + exportTracingEvent(event: TracingEvent): Promise; + /** Add score to a trace (optional) */ + addScoreToTrace?({ + traceId, + spanId, + score, + reason, + scorerName, + metadata, + }: { + traceId: string; + spanId?: string; + score: number; + reason?: string; + scorerName: string; + metadata?: Record; + }): Promise; + /** Shutdown exporter */ + shutdown(): Promise; +} +``` + +### SpanOutputProcessor + +Interface for span output processors. + +```typescript +interface SpanOutputProcessor { + /** Processor name */ + name: string; + /** Process span before export */ + process(span?: AnySpan): AnySpan | undefined; + /** Shutdown processor */ + shutdown(): Promise; +} +``` + +## Span Types + +### SpanType + +AI-specific span types with their associated metadata. + +```typescript +enum SpanType { + /** Agent run - root span for agent processes */ + AGENT_RUN = "agent_run", + /** Generic span for custom operations */ + GENERIC = "generic", + /** Model generation with model calls, token usage, prompts, completions */ + MODEL_GENERATION = "model_generation", + /** Single model execution step within a generation (one API call) */ + MODEL_STEP = "model_step", + /** Individual model streaming chunk/event */ + MODEL_CHUNK = "model_chunk", + /** MCP (Model Context Protocol) tool execution */ + MCP_TOOL_CALL = "mcp_tool_call", + /** Input or Output Processor execution */ + PROCESSOR_RUN = "processor_run", + /** Function/tool execution with inputs, outputs, errors */ + TOOL_CALL = "tool_call", + /** Workflow run - root span for workflow processes */ + WORKFLOW_RUN = "workflow_run", + /** Workflow step execution with step status, data flow */ + WORKFLOW_STEP = "workflow_step", + /** Workflow conditional execution with condition evaluation */ + WORKFLOW_CONDITIONAL = "workflow_conditional", + /** Individual condition evaluation within conditional */ + WORKFLOW_CONDITIONAL_EVAL = "workflow_conditional_eval", + /** Workflow parallel execution */ + WORKFLOW_PARALLEL = "workflow_parallel", + /** Workflow loop execution */ + WORKFLOW_LOOP = "workflow_loop", + /** Workflow sleep operation */ + WORKFLOW_SLEEP = "workflow_sleep", + /** Workflow wait for event operation */ + WORKFLOW_WAIT_EVENT = "workflow_wait_event", +} +``` + +### AnySpan + +Union type for cases that need to handle any span. + +```typescript +type AnySpan = Span; +``` + +## Key Span Attributes + +### AgentRunAttributes + +```typescript +interface AgentRunAttributes { + /** Agent identifier */ + agentId: string; + /** Agent Instructions */ + instructions?: string; + /** Agent Prompt */ + prompt?: string; + /** Available tools for this execution */ + availableTools?: string[]; + /** Maximum steps allowed */ + maxSteps?: number; +} +``` + +### ModelGenerationAttributes + +```typescript +interface ModelGenerationAttributes { + /** Model name (e.g., 'gpt-4', 'claude-3') */ + model?: string; + /** Model provider (e.g., 'openai', 'anthropic') */ + provider?: string; + /** Type of result/output this model call produced */ + resultType?: "tool_selection" | "response_generation" | "reasoning" | "planning"; + /** Token usage statistics */ + usage?: { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + promptCacheHitTokens?: number; + promptCacheMissTokens?: number; + }; + /** Model parameters */ + parameters?: { + maxOutputTokens?: number; + temperature?: number; + topP?: number; + topK?: number; + presencePenalty?: number; + frequencyPenalty?: number; + stopSequences?: string[]; + seed?: number; + maxRetries?: number; + }; + /** Whether this was a streaming response */ + streaming?: boolean; + /** Reason the generation finished */ + finishReason?: string; +} +``` + +### ToolCallAttributes + +```typescript +interface ToolCallAttributes { + toolId?: string; + toolType?: string; + toolDescription?: string; + success?: boolean; +} +``` + +## Options Types + +### StartSpanOptions + +```typescript +interface StartSpanOptions { + /** Span type */ + type: TType; + /** Span name */ + name: string; + /** Span attributes */ + attributes?: SpanTypeMap[TType]; + /** Span metadata */ + metadata?: Record; + /** Input data */ + input?: any; + /** Parent span */ + parent?: AnySpan; + /** Policy-level tracing configuration */ + tracingPolicy?: TracingPolicy; + /** Options passed when using a custom sampler strategy */ + customSamplerOptions?: CustomSamplerOptions; +} +``` + +### TracingOptions + +```typescript +interface TracingOptions { + /** Metadata to add to the root trace span */ + metadata?: Record; + /** + * Additional RequestContext keys to extract as metadata for this trace. + * These keys are added to the requestContextKeys config. + * Supports dot notation for nested values (e.g., 'user.id', 'session.data.experimentId'). + */ + requestContextKeys?: string[]; + /** + * Trace ID to use for this execution (1-32 hexadecimal characters). + * If provided, this trace will be part of the specified trace rather than starting a new one. + */ + traceId?: string; + /** + * Parent span ID to use for this execution (1-16 hexadecimal characters). + * If provided, the root span will be created as a child of this span. + */ + parentSpanId?: string; + /** + * Tags to apply to this trace. + * Tags are string labels that can be used to categorize and filter traces + * Note: Tags are only applied to the root span of a trace. + */ + tags?: string[]; +} +``` + +## Sampling Types + +### SamplingStrategy + +```typescript +type SamplingStrategy = + | { type: "always" } + | { type: "never" } + | { type: "ratio"; probability: number } + | { type: "custom"; sampler: (options?: CustomSamplerOptions) => boolean }; +``` \ No newline at end of file diff --git a/docs/tracing-spans-reference.md b/docs/tracing-spans-reference.md new file mode 100644 index 00000000..a7fbb5f9 --- /dev/null +++ b/docs/tracing-spans-reference.md @@ -0,0 +1,188 @@ +# Mastra Tracing Spans Reference + +## BaseSpan + +Base interface for all span types. + +```typescript +interface BaseSpan { + /** Unique span identifier */ + id: string; + /** OpenTelemetry-compatible trace ID (32 hex chars) */ + traceId: string; + /** Name of the span */ + name: string; + /** Type of the span */ + type: TType; + /** When span started */ + startTime: Date; + /** When span ended */ + endTime?: Date; + /** Type-specific attributes */ + attributes?: SpanTypeMap[TType]; + /** User-defined metadata */ + metadata?: Record; + /** Input passed at the start of the span */ + input?: any; + /** Output generated at the end of the span */ + output?: any; + /** Error information if span failed */ + errorInfo?: { + message: string; + id?: string; + domain?: string; + category?: string; + details?: Record; + }; + /** Is an event span? (occurs at startTime, has no endTime) */ + isEvent: boolean; +} +``` + +## Span + +Span interface, used internally for tracing. Extends BaseSpan with lifecycle methods and properties. + +```typescript +interface Span extends BaseSpan { + /** Is an internal span? (spans internal to the operation of mastra) */ + isInternal: boolean; + /** Parent span reference (undefined for root spans) */ + parent?: AnySpan; + /** Pointer to the ObservabilityInstance instance */ + observabilityInstance: ObservabilityInstance; +} +``` + +### Properties + +```typescript +/** Returns TRUE if the span is the root span of a trace */ +get isRootSpan(): boolean + +/** Returns TRUE if the span is a valid span (not a NO-OP Span) */ +get isValid(): boolean + +/** Get the closest parent spanId that isn't an internal span */ +getParentSpanId(includeInternalSpans?: boolean): string | undefined + +/** Returns a lightweight span ready for export */ +exportSpan(includeInternalSpans?: boolean): ExportedSpan | undefined +``` + +### Methods + +#### end + +```typescript +end(options?: EndSpanOptions): void +``` + +Ends the span and triggers export to configured exporters. Sets the `endTime` and optionally updates `output`, `metadata`, and `attributes`. + +#### error + +```typescript +error(options: ErrorSpanOptions): void +``` + +Records an error on the span. Sets the `errorInfo` field and can optionally end the span. + +#### update + +```typescript +update(options: UpdateSpanOptions): void +``` + +Updates span data while it's still active. Can modify `input`, `output`, `metadata`, and `attributes`. + +#### createChildSpan + +```typescript +createChildSpan( + options: ChildSpanOptions, +): Span +``` + +Creates a child span under this span. Child spans track sub-operations and inherit the trace context. + +#### createEventSpan + +```typescript +createEventSpan( + options: ChildEventOptions, +): Span +``` + +Creates an event span under this span. Event spans represent point-in-time occurrences with no duration. + +## ExportedSpan + +Exported Span interface, used for tracing exporters. A lightweight version of Span without methods or circular references. + +```typescript +interface ExportedSpan extends BaseSpan { + /** Parent span id reference (undefined for root spans) */ + parentSpanId?: string; + /** TRUE if the span is the root span of a trace */ + isRootSpan: boolean; +} +``` + +## Span Lifecycle Events + +### TracingEventType + +```typescript +enum TracingEventType { + /** Emitted when a span is created and started */ + SPAN_STARTED = "span_started", + /** Emitted when a span is updated via update() */ + SPAN_UPDATED = "span_updated", + /** Emitted when a span is ended via end() or error() */ + SPAN_ENDED = "span_ended", +} +``` + +### TracingEvent + +```typescript +type TracingEvent = + | { type: "span_started"; exportedSpan: AnyExportedSpan } + | { type: "span_updated"; exportedSpan: AnyExportedSpan } + | { type: "span_ended"; exportedSpan: AnyExportedSpan }; +``` + +Exporters receive these events to process and send trace data to observability platforms. + +## Union Types + +### AnySpan + +```typescript +type AnySpan = Span; +``` + +Union type for cases that need to handle any span type. + +### AnyExportedSpan + +```typescript +type AnyExportedSpan = ExportedSpan; +``` + +Union type for cases that need to handle any exported span type. + +## NO-OP Spans + +When tracing is disabled (sampling returns false), NO-OP spans are returned: + +### NoOpSpan + +A span that performs no operations. All methods are no-ops: + +- `id` returns `'no-op'` +- `traceId` returns `'no-op-trace'` +- `isValid` returns `false` +- `end()`, `error()`, `update()` do nothing +- `createChildSpan()` returns another NO-OP span \ No newline at end of file diff --git a/package-lock.json b/package-lock.json index 42805325..ae0cec20 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "agentstack", - "version": "1.1.1", + "version": "1.1.2", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "agentstack", - "version": "1.1.1", + "version": "1.1.2", "license": "ISC", "dependencies": { "@ai-sdk/google": "^2.0.52", diff --git a/src/mastra/agents/codingAgents.ts b/src/mastra/agents/codingAgents.ts index 69fc78a5..d6634ae3 100644 --- a/src/mastra/agents/codingAgents.ts +++ b/src/mastra/agents/codingAgents.ts @@ -11,7 +11,20 @@ import { upstashMemory } from '../config/upstash'; import { codeAnalysisTool } from '../tools/code-analysis.tool'; import { codeSearchTool } from '../tools/code-search.tool'; import { diffReviewTool } from '../tools/diff-review.tool'; -import * as e2bTools from '../tools/e2b'; +import { + checkFileExists, + createDirectory, + createSandbox, + deleteFile, + getFileInfo, + getFileSize, + listFiles, + runCode, + runCommand, + watchDirectory, + writeFile, + writeFiles +} from '../tools/e2b'; import { execaTool } from '../tools/execa-tool'; import { findReferencesTool } from '../tools/find-references.tool'; import { findSymbolTool } from '../tools/find-symbol.tool'; @@ -46,7 +59,7 @@ export const codeArchitectAgent = new Agent({ instructions: ({ requestContext }: { requestContext: RequestContext }) => { const userTier = requestContext.get('user-tier') ?? 'free' const language = requestContext.get('language') ?? 'en' - // const projectRoot = requestContext.get('projectRoot') ?? process.cwd() + // const projectRoot = requestContext.get('projectRoot') ?? process.cwd() return { role: 'system', @@ -334,7 +347,18 @@ Always use Vitest syntax: describe, it, expect, vi.mock, vi.fn.`, testGeneratorTool, codeSearchTool, execaTool, - ...e2bTools, + createSandbox, + writeFile, + writeFiles, + listFiles, + deleteFile, + createDirectory, + getFileInfo, + checkFileExists, + getFileSize, + watchDirectory, + runCommand, + runCode, // ...githubMCP.getTools(), }, memory: upstashMemory, @@ -456,7 +480,18 @@ For each refactoring: findReferencesTool, findSymbolTool, execaTool, - ...e2bTools, + createSandbox, + writeFile, + writeFiles, + listFiles, + deleteFile, + createDirectory, + getFileInfo, + checkFileExists, + getFileSize, + watchDirectory, + runCommand, + runCode, }, memory: upstashMemory, diff --git a/src/mastra/config/libsql.ts b/src/mastra/config/libsql.ts new file mode 100644 index 00000000..888b84b2 --- /dev/null +++ b/src/mastra/config/libsql.ts @@ -0,0 +1,119 @@ +import { google } from "@ai-sdk/google"; +import { LibSQLStore, LibSQLVector } from "@mastra/libsql"; +import { Memory } from "@mastra/memory"; +import { log } from "./logger"; + +const libsqlstorage = new LibSQLStore({ + id: 'libsql-storage', + url: process.env.TURSO_DATABASE_URL ?? "libsql://your-db-name.aws-ap-northeast-1.turso.io", + authToken: process.env.TURSO_AUTH_TOKEN, +}) + +// Create a new vector store instance +const libsqlvector = new LibSQLVector({ + id: 'libsql-vector', + connectionUrl: process.env.TURSO_DATABASE_URL ?? "file:./vectors.db", + // Optional: for Turso cloud databases + authToken: process.env.TURSO_AUTH_TOKEN, +}); + +// Create an index +await libsqlvector.createIndex({ + indexName: "memory_messages", + dimension: 3072, + metric: "cosine" +}); + +// Add vectors with metadata +const vectors = [[0.1, 0.2], [0.3, 0.4]]; +const metadata = [ + { text: "first document", category: "A" }, + { text: "second document", category: "B" } +]; +await libsqlvector.upsert({ + indexName: "memory_messages", + vectors, + metadata, +}); + +// Query similar vectors +const queryVector = [0.1, 0.2]; // Your query vector + +const results = await libsqlvector.query({ + indexName: "memory_messages", + queryVector, + topK: 10, // top K results + filter: { category: "A" } // optional metadata filter +}); + +export const LibsqlMemory = new Memory({ + storage: libsqlstorage, + vector: libsqlvector, + embedder: google.textEmbedding('gemini-embedding-001'), + options: { + // Message management + lastMessages: parseInt(process.env.MEMORY_LAST_MESSAGES ?? '500'), + generateTitle: true, + // Advanced semantic recall with HNSW index configuration + semanticRecall: { + topK: parseInt(process.env.SEMANTIC_TOP_K ?? '5'), + messageRange: { + before: parseInt(process.env.SEMANTIC_RANGE_BEFORE ?? '3'), + after: parseInt(process.env.SEMANTIC_RANGE_AFTER ?? '2'), + }, + scope: 'resource', // 'resource' | 'thread' + // HNSW index configuration to support high-dimensional embeddings (>2000 dimensions) + indexConfig: { + }, + threshold: 0.75, // Similarity threshold for semantic recall + indexName: 'memory_messages', // Index name for semantic recall + }, + // Enhanced working memory with supported template + workingMemory: { + enabled: true, + scope: 'resource', // 'resource' | 'thread' + version: 'vnext', + template: `# User Context + ## Profile + - Name/Role: + - Org/Loc: + - Style/Level: + ## Active + - Goals/Projects: + - Recent/Deadlines: + ## Insights + - Patterns/Habits: + - Session Focus: + - Action Items: + `, + }, + } +}) + + log.info('PG Store and Memory initialized with PgVector support', { + schema: process.env.DB_SCHEMA ?? 'mastra', + maxConnections: parseInt(process.env.DB_MAX_CONNECTIONS ?? '20'), + + memoryOptions: { + lastMessages: parseInt(process.env.MEMORY_LAST_MESSAGES ?? '500'), + semanticRecall: { + topK: parseInt(process.env.SEMANTIC_TOP_K ?? '5'), + messageRange: { + before: parseInt(process.env.SEMANTIC_RANGE_BEFORE ?? '3'), + after: parseInt(process.env.SEMANTIC_RANGE_AFTER ?? '2'), + }, + scope: 'resource', + indexConfig: { + type: 'ivfflat', // flat index type (supports dimensions > 4000, unlike HNSW limit of 2000) + metric: 'cosine', // Distance metric for normalized embeddings + } + }, + workingMemory: { + enabled: true, + scope: 'resource', + version: 'vnext', + } + + }, +}) + diff --git a/src/mastra/index.ts b/src/mastra/index.ts index 7dd65260..94968f5b 100644 --- a/src/mastra/index.ts +++ b/src/mastra/index.ts @@ -1,7 +1,6 @@ import { chatRoute, networkRoute, workflowRoute } from '@mastra/ai-sdk' import { Mastra } from '@mastra/core' -import { Observability, SamplingStrategyType, SensitiveDataFilter } from '@mastra/observability' -import { OtelExporter } from '@mastra/otel-exporter' +import { CloudExporter, Observability, SamplingStrategyType, SensitiveDataFilter } from '@mastra/observability' import { PostgresStore } from '@mastra/pg' // Config import { log } from './config/logger' @@ -235,7 +234,7 @@ export const mastra = new Mastra({ serviceName: "ai", sampling: { type: SamplingStrategyType.RATIO, probability: 0.5 }, // 50% sampling - requestContextKeys: ["userId", "environment", "tenantId"], + requestContextKeys: ["userId", "environment", "tenantId", "tool.id", "agent.id", "workflow.id", "memory.thread.id", "user.id"], spanOutputProcessors: [ new SensitiveDataFilter( { @@ -266,13 +265,14 @@ export const mastra = new Mastra({ { logger: log, logLevel: 'info', -// strategy: 'realtime', -// maxBatchSize: 500, -// maxBufferSize: 1000, -// maxBatchWaitMs: 6000, + strategy: 'auto', + maxBatchSize: 1000, + maxBufferSize: 10000, + maxBatchWaitMs: 5000, maxRetries: 5, } ), // Studio access0 + new CloudExporter(), ], includeInternalSpans: true, diff --git a/src/mastra/tools/AGENTS.md b/src/mastra/tools/AGENTS.md index 54ed6e60..d41f128b 100644 --- a/src/mastra/tools/AGENTS.md +++ b/src/mastra/tools/AGENTS.md @@ -67,6 +67,216 @@ The `createTool()` function returns a `Tool` object. - **`Tool`**: `object` - An object representing the defined tool, ready to be added to an agent. +## Tracing Implementation Patterns + +All tools should implement proper tracing using Mastra's observability system. The following pattern ensures consistent, structured tracing across all tools. + +### Required Imports + +> IMPORTANT: Tools MUST use Mastra's tracing helpers and the execution-provided `tracingContext`. Do NOT import or use OpenTelemetry APIs (for example, `@opentelemetry/api`) directly inside tools — instrumentation and exporter configuration belong at the platform / runtime startup layer, not inside individual tools. + +```typescript +import type { RequestContext } from '@mastra/core/request-context' +import type { InferUITool } from '@mastra/core/tools' +import { createTool } from '@mastra/core/tools' +import { SpanType } from '@mastra/core/observability' +import { z } from 'zod' +import { log } from '../config/logger' +``` + +### Context Interface Definition + +```typescript +export interface ToolNameContext extends RequestContext { + userId?: string + workspaceId?: string + // Add tool-specific context properties +} +``` + +### Tracing Context Usage (preferred) + +```typescript +execute: async (inputData, context) => { + const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ToolNameContext | undefined + + // Extract context values + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + + // Respect cancellation early + if (abortSignal?.aborted ?? false) { + throw new Error('Tool call cancelled') + } + + // Create child span for tool execution (Mastra tracingContext) + const toolSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'tool-operation-name', + input: inputData, + metadata: { + 'tool.id': 'tool-id', + 'tool.input.param': inputData.param, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) + + // Emit a structured progress event (required) + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `Input: param="${inputData.param}" - 📝 Operation started`, + stage: 'tool-id', + }, + id: 'tool-id', + }) + + try { + // Tool implementation + const result = await performToolLogic(inputData) + + // Update span with successful result and finish progress + toolSpan?.update({ + output: result, + metadata: { + 'tool.output.success': true, + 'tool.output.resultSize': Array.isArray(result) ? result.length : (typeof result === 'string' ? result.length : undefined), + }, + }) + toolSpan?.end() + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `Input: param="${inputData.param}" - ✅ Operation complete`, + stage: 'tool-id', + }, + id: 'tool-id', + }) + + return result + } catch (error) { + // Record error in span and mark progress done + toolSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `Input: param="${inputData.param}" - ❌ Operation error: ${error instanceof Error ? error.message : String(error)}`, + stage: 'tool-id', + }, + id: 'tool-id', + }) + + throw error + } +} +``` + +#### Migration: Replacing OpenTelemetry usage in tools + +If a tool currently imports and uses OpenTelemetry APIs directly (for example `trace.getTracer(...)`, `span.recordException(...)`, `span.setStatus(...)`), migrate it to the runtime `tracingContext` pattern above. Example migration: + +```diff +- // Old: do NOT use inside tools +- import { trace, SpanStatusCode } from '@opentelemetry/api' +- const tracer = trace.getTracer('pnpm-tool') +- const span = tracer.startSpan('pnpm-build', { attributes: { name, packagePath } }) +- try { /* ... */ } catch (e) { span.recordException(e); span.setStatus({ code: SpanStatusCode.ERROR, message: e.message }); span.end() } + ++ // New: preferred pattern inside tools ++ const tracingContext = context?.tracingContext ++ const span = tracingContext?.currentSpan?.createChildSpan({ ++ type: SpanType.TOOL_CALL, ++ name: 'pnpm-build', ++ input: { name, packagePath }, ++ metadata: { 'tool.id': 'pnpmBuild', 'operation': 'pnpm-build' }, ++ }) ++ try { /* ... */ } catch (e) { ++ span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: true }) ++ } +``` + +Checklist for safe migration: +- [ ] Replace `@opentelemetry/*` imports inside `src/mastra/tools` with the `tracingContext` pattern. +- [ ] Ensure each tool emits `data-tool-progress` events at start and completion (or stage updates). +- [ ] Add `abortSignal` checks before long-running operations. +- [ ] Use `SpanType.TOOL_CALL` for tool-level spans and include `tool.id` in span metadata. +- [ ] Add a small unit test that mocks `tracingContext.currentSpan.createChildSpan()` to verify spans are created, updated, and closed. + +### Span Lifecycle Management + +**Creating Child Spans:** + +```typescript +const childSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, // Use SpanType enum + name: 'operation-name', + input: inputData, + metadata: { + 'tool.id': 'tool-id', + 'operation.type': 'specific-operation', + }, +}) +``` + +**Updating Spans:** + +```typescript +span?.update({ + output: result, + metadata: { + 'operation.status': 'success', + 'result.size': result.length, + }, +}) +``` + +**Ending Spans:** + +```typescript +span?.end() +``` + +**Error Handling:** + +```typescript +span?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, +}) +``` + +### Progress Events + +Emit progress events for long-running operations: + +```typescript +await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', // "in-progress" | "done" + message: `Input: param="${inputData.param}" - 📝 Operation in progress`, + stage: 'tool-id', + }, + id: 'tool-id', +}) +``` + +### Complete Example + +See `weather-tool.ts` for a complete implementation example with proper tracing, error handling, and lifecycle hooks. + ## Categories ### 1. 💰 Financial Data APIs diff --git a/src/mastra/tools/alpha-vantage.tool.ts b/src/mastra/tools/alpha-vantage.tool.ts index 9cd64921..356e6379 100644 --- a/src/mastra/tools/alpha-vantage.tool.ts +++ b/src/mastra/tools/alpha-vantage.tool.ts @@ -1,11 +1,10 @@ -import { createTool } from '@mastra/core/tools' +import { SpanType } from '@mastra/core/observability' +import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { createTool } from '@mastra/core/tools' import { z } from 'zod' import { log } from '../config/logger' -import type { RequestContext } from '@mastra/core/request-context' - /** * Alpha Vantage Tools * @@ -15,12 +14,13 @@ import type { RequestContext } from '@mastra/core/request-context' * * Requires ALPHA_VANTAGE_API_KEY environment variable */ - +export type UserTier = 'free' | 'pro' | 'enterprise' /** * Base Request Context for Alpha Vantage */ export interface AlphaVantageRequestContext extends RequestContext { - apiKey?: string + apiKey?: string + 'user-tier': UserTier } // In-memory counter to track tool calls per request @@ -93,16 +93,7 @@ export const alphaVantageCryptoTool = createTool({ }), 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', { - attributes: { - 'tool.id': 'alpha-vantage-crypto', - 'tool.input.symbol': inputData.symbol, - 'tool.input.market': inputData.market, - 'tool.input.function': inputData.function, - }, - }) + const tracingContext = context?.tracingContext await context?.writer?.custom({ type: 'data-tool-progress', @@ -113,6 +104,20 @@ export const alphaVantageCryptoTool = createTool({ }, id: 'alpha-vantage-crypto', }) + + // Create child span for crypto data fetch + const cryptoSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'alpha-vantage-crypto-fetch', + input: inputData, + metadata: { + 'tool.id': 'alpha-vantage-crypto', + 'tool.input.symbol': inputData.symbol, + 'tool.input.market': inputData.market, + 'tool.input.function': inputData.function, + }, + }) + const apiKey = requestContext?.apiKey ?? process.env.ALPHA_VANTAGE_API_KEY if (typeof apiKey !== 'string' || apiKey.trim() === '') { @@ -270,7 +275,7 @@ export const alphaVantageCryptoTool = createTool({ time_zone: getMetadataValue('7. Time Zone') ?? undefined, }, } - span.end() + cryptoSpan?.end() return result } catch (error) { const errMsg = @@ -287,10 +292,11 @@ export const alphaVantageCryptoTool = createTool({ id: 'alpha-vantage-crypto', }) if (error instanceof Error) { - span.recordException(error) + cryptoSpan?.error({ + error, + endSpan: true, + }) } - span.setStatus({ code: 2, message: errMsg }) // ERROR status - span.end() throw error instanceof Error ? error : new Error(errMsg) } }, @@ -423,15 +429,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', { - attributes: { - 'tool.id': 'alpha-vantage-stock', - 'tool.input.symbol': inputData.symbol, - 'tool.input.function': inputData.function, - }, - }) + const tracingContext = context?.tracingContext await context?.writer?.custom({ type: 'data-tool-progress', @@ -442,6 +440,19 @@ export const alphaVantageStockTool = createTool({ }, id: 'alpha-vantage-stock', }) + + // Create child span for stock data fetch + const stockSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'alpha-vantage-stock-fetch', + input: inputData, + metadata: { + 'tool.id': 'alpha-vantage-stock', + 'tool.input.symbol': inputData.symbol, + 'tool.input.function': inputData.function, + }, + }) + const apiKey = requestContext?.apiKey ?? process.env.ALPHA_VANTAGE_API_KEY if (typeof apiKey !== 'string' || apiKey.trim() === '') { @@ -642,7 +653,7 @@ export const alphaVantageStockTool = createTool({ series_type: inputData.series_type, }, } - span.end() + stockSpan?.end() return result } catch (error) { const errorMessage = @@ -650,10 +661,11 @@ export const alphaVantageStockTool = createTool({ ? error.message : 'Unknown error occurred' if (error instanceof Error) { - span.recordException(error) + stockSpan?.error({ + error, + endSpan: true, + }) } - span.setStatus({ code: 2, message: errorMessage }) // ERROR status - span.end() throw error instanceof Error ? error : new Error(errorMessage) } }, @@ -816,15 +828,7 @@ export const alphaVantageTool = createTool({ }), execute: async (inputData, context) => { const requestContext = context?.requestContext as AlphaVantageRequestContext - const span = trace - .getTracer('alpha-vantage-tool', '1.0.0') - .startSpan('alpha-vantage', { - attributes: { - 'tool.id': 'alpha-vantage', - 'tool.input.function': inputData.function, - 'tool.input.symbol': inputData.symbol, - }, - }) + const tracingContext = context?.tracingContext await context?.writer?.custom({ type: 'data-tool-progress', @@ -835,6 +839,19 @@ export const alphaVantageTool = createTool({ }, id: 'alpha-vantage', }) + + // Create child span for general data fetch + const generalSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'alpha-vantage-general-fetch', + input: inputData, + metadata: { + 'tool.id': 'alpha-vantage', + 'tool.input.function': inputData.function, + 'tool.input.symbol': inputData.symbol, + }, + }) + const apiKey = requestContext?.apiKey ?? process.env.ALPHA_VANTAGE_API_KEY if (typeof apiKey !== 'string' || apiKey.trim() === '') { @@ -1050,7 +1067,7 @@ export const alphaVantageTool = createTool({ : undefined, }, } - span.end() + generalSpan?.end() return result } catch (error) { const errorMessage = @@ -1058,18 +1075,21 @@ export const alphaVantageTool = createTool({ ? error.message : 'Unknown error occurred' if (error instanceof Error) { - span.recordException(error) + generalSpan?.error({ + error, + endSpan: true, + }) } - span.setStatus({ code: 2, message: errorMessage }) // ERROR status - span.end() throw error instanceof Error ? error : new Error(errorMessage) } }, onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('alphaVantageTool tool input streaming started', { toolCallId, + log.info('alphaVantageTool tool input streaming started', { + toolCallId, messageCount: messages.length, abortSignal: abortSignal?.aborted, - hook: 'onInputStart' }); + hook: 'onInputStart' + }); }, onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { log.info('alphaVantageTool received input chunk', { diff --git a/src/mastra/tools/arxiv.tool.ts b/src/mastra/tools/arxiv.tool.ts index 59fbb7aa..1a205e3b 100644 --- a/src/mastra/tools/arxiv.tool.ts +++ b/src/mastra/tools/arxiv.tool.ts @@ -1,13 +1,16 @@ +import { SpanType } from '@mastra/core/observability' import type { InferUITool } from '@mastra/core/tools' import { createTool } 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' +export type UserTier = 'free' | 'pro' | 'enterprise' + export interface ArxivRequestContext extends RequestContext { - userId?: string + 'user-tier': UserTier + userId?: string } // In-memory counter to track tool calls per request @@ -35,437 +38,489 @@ import { extractText, getDocumentProxy } from 'unpdf' * Extract text content from PDF buffer using `unpdf` */ async function extractPdfText( - pdfBuffer: Buffer, - maxPages = 1000 + pdfBuffer: Buffer, + maxPages = 1000 ): Promise<{ - text: string - numpages: number - metadata?: Record + text: string + numpages: number + metadata?: Record }> { - try { - const pdf = await getDocumentProxy(new Uint8Array(pdfBuffer)) - const { totalPages, text } = await extractText(pdf, { - mergePages: true, - }) - - return { - text: text ?? '', - numpages: Math.min(totalPages ?? 1, maxPages), - metadata: undefined, - } - } catch (error) { - throw new Error( - `PDF parsing failed: ${error instanceof Error ? error.message : 'Unknown error'}` - ) + try { + const pdf = await getDocumentProxy(new Uint8Array(pdfBuffer)) + const { totalPages, text } = await extractText(pdf, { + mergePages: true, + }) + + return { + text: text ?? '', + numpages: Math.min(totalPages ?? 1, maxPages), + metadata: undefined, } + } catch (error) { + throw new Error( + `PDF parsing failed: ${error instanceof Error ? error.message : 'Unknown error'}` + ) + } } /** * Convert extracted PDF text to markdown format */ function convertPdfTextToMarkdown( - text: string, - metadata?: { title?: string; authors?: string[] } + text: string, + metadata?: { title?: string; authors?: string[] } ): string { - let markdown = text - - // Basic text normalization - markdown = markdown - .replace(/\f/g, '\n\n') // Form feeds to double newlines - .replace(/\r\n/g, '\n') // Normalize line endings - .replace(/\n\n\n+/g, '\n\n') // Multiple newlines to double - .replace(/\t/g, ' ') // Tabs to spaces - .replace(/(\w+)-\n(\w+)/g, '$1$2') // Fix hyphenated words - .replace(/\s+([.,!?;:])/g, '$1') // Clean punctuation spacing - .trim() - - // Add frontmatter if metadata available - const hasTitle = - typeof metadata?.title === 'string' && metadata.title.trim() !== '' - const hasAuthors = - Array.isArray(metadata?.authors) && metadata.authors.length > 0 - - if (hasTitle || hasAuthors) { - const frontmatter = [ - '---', - hasTitle - ? `title: "${String(metadata?.title).replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"` - : null, - hasAuthors ? `authors: ${JSON.stringify(metadata.authors)}` : null, - 'source: "arxiv-pdf"', - 'extracted_at: "' + new Date().toISOString() + '"', - '---', - '', - ] - .filter(Boolean) - .join('\n') - - markdown = frontmatter + markdown - } - - return markdown + let markdown = text + + // Basic text normalization + markdown = markdown + .replace(/\f/g, '\n\n') // Form feeds to double newlines + .replace(/\r\n/g, '\n') // Normalize line endings + .replace(/\n\n\n+/g, '\n\n') // Multiple newlines to double + .replace(/\t/g, ' ') // Tabs to spaces + .replace(/(\w+)-\n(\w+)/g, '$1$2') // Fix hyphenated words + .replace(/\s+([.,!?;:])/g, '$1') // Clean punctuation spacing + .trim() + + // Add frontmatter if metadata available + const hasTitle = + typeof metadata?.title === 'string' && metadata.title.trim() !== '' + const hasAuthors = + Array.isArray(metadata?.authors) && metadata.authors.length > 0 + + if (hasTitle || hasAuthors) { + const frontmatter = [ + '---', + hasTitle + ? `title: "${String(metadata?.title).replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"` + : null, + hasAuthors ? `authors: ${JSON.stringify(metadata.authors)}` : null, + 'source: "arxiv-pdf"', + 'extracted_at: "' + new Date().toISOString() + '"', + '---', + '', + ] + .filter(Boolean) + .join('\n') + + markdown = frontmatter + markdown + } + + return markdown } export const arxivTool = createTool({ - id: 'arxiv', - description: - 'Search and retrieve academic papers from arXiv.org including metadata, abstracts, and download links', - inputSchema: z.object({ - query: z - .string() - .optional() - .describe('Search query for papers (keywords, titles, authors)'), - id: z - .string() - .optional() - .describe( - "Specific arXiv paper ID (e.g., '2103.12345' or '2103.12345v1')" - ), - author: z.string().optional().describe('Search by author name'), - title: z.string().optional().describe('Search by paper title'), - category: z - .enum([ - 'cs.AI', - 'cs.CL', - 'cs.CV', - 'cs.LG', - 'cs.NE', - 'cs.RO', - 'math.AG', - 'math.AP', - 'math.CO', - 'math.CT', - 'math.DS', - 'physics.optics', - 'physics.comp-ph', - 'physics.gen-ph', - 'q-bio.BM', - 'q-bio.GN', - 'q-bio.MN', - 'q-bio.NC', - 'q-bio.OT', - 'q-bio.PE', - 'q-bio.QM', - 'stat.AP', - 'stat.CO', - 'stat.ME', - 'stat.ML', - 'stat.TH', - 'astro-ph.CO', - 'astro-ph.EP', - 'astro-ph.GA', - 'astro-ph.HE', - 'astro-ph.IM', - 'astro-ph.SR', - 'cond-mat.dis-nn', - 'cond-mat.mes-hall', - 'cond-mat.mtrl-sci', - 'cond-mat.other', - 'cond-mat.quant-gas', - 'cond-mat.soft', - 'cond-mat.stat-mech', - 'cond-mat.str-el', - 'cond-mat.supr-con', - 'hep-ex', - 'hep-lat', - 'hep-ph', - 'hep-th', - 'nucl-ex', - 'nucl-th', - 'quant-ph', - ]) - .optional() - .describe('ArXiv category/subject area'), - max_results: z - .number() - .min(1) - .max(100) - .optional() - .default(10) - .describe('Maximum number of results to return (1-100)'), - sort_by: z - .enum(['relevance', 'lastUpdatedDate', 'submittedDate']) - .optional() - .default('relevance') - .describe('Sort results by criteria'), - sort_order: z - .enum(['ascending', 'descending']) - .optional() - .default('descending') - .describe('Sort order'), - start: z - .number() - .min(0) - .optional() - .default(0) - .describe('Starting index for results (0-based)'), - include_abstract: z - .boolean() - .optional() - .default(true) - .describe('Whether to include paper abstracts in results'), - }), - outputSchema: z.object({ - papers: z.array( - z.object({ - id: z.string(), - title: z.string(), - authors: z.array(z.string()), - abstract: z.string().optional(), - categories: z.array(z.string()), - published: z.string(), - updated: z.string().optional(), - pdf_url: z.string(), - doi: z.string().optional(), - journal_ref: z.string().optional(), - comments: z.string().optional(), - }) - ), - total_results: z.number(), - start_index: z.number(), - max_results: z.number(), - }), - 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', { - attributes: { - 'tool.id': 'arxiv', - 'tool.input.query': inputData.query, - 'tool.input.id': inputData.id, - 'tool.input.category': inputData.category, - 'tool.input.maxResults': inputData.max_results, - }, - }) - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `📚 Searching arXiv for "${inputData.query ?? inputData.id ?? 'papers'}"`, - stage: 'arxiv', - }, - id: 'arxiv', - }) - toolCallCounters.set('arxiv', (toolCallCounters.get('arxiv') ?? 0) + 1) - try { - const params = new URLSearchParams() - - // Build search query - const searchTerms: string[] = [] - - if (inputData.query !== undefined && inputData.query !== null) { - searchTerms.push(inputData.query) - } - - if (inputData.author !== undefined && inputData.author !== null) { - searchTerms.push(`au:${inputData.author}`) - } + id: 'arxiv', + description: + 'Search and retrieve academic papers from arXiv.org including metadata, abstracts, and download links', + inputSchema: z.object({ + query: z + .string() + .optional() + .describe('Search query for papers (keywords, titles, authors)'), + id: z + .string() + .optional() + .describe( + "Specific arXiv paper ID (e.g., '2103.12345' or '2103.12345v1')" + ), + author: z.string().optional().describe('Search by author name'), + title: z.string().optional().describe('Search by paper title'), + category: z + .enum([ + 'cs.AI', + 'cs.CL', + 'cs.CV', + 'cs.LG', + 'cs.NE', + 'cs.RO', + 'math.AG', + 'math.AP', + 'math.CO', + 'math.CT', + 'math.DS', + 'physics.optics', + 'physics.comp-ph', + 'physics.gen-ph', + 'q-bio.BM', + 'q-bio.GN', + 'q-bio.MN', + 'q-bio.NC', + 'q-bio.OT', + 'q-bio.PE', + 'q-bio.QM', + 'stat.AP', + 'stat.CO', + 'stat.ME', + 'stat.ML', + 'stat.TH', + 'astro-ph.CO', + 'astro-ph.EP', + 'astro-ph.GA', + 'astro-ph.HE', + 'astro-ph.IM', + 'astro-ph.SR', + 'cond-mat.dis-nn', + 'cond-mat.mes-hall', + 'cond-mat.mtrl-sci', + 'cond-mat.other', + 'cond-mat.quant-gas', + 'cond-mat.soft', + 'cond-mat.stat-mech', + 'cond-mat.str-el', + 'cond-mat.supr-con', + 'hep-ex', + 'hep-lat', + 'hep-ph', + 'hep-th', + 'nucl-ex', + 'nucl-th', + 'quant-ph', + ]) + .optional() + .describe('ArXiv category/subject area'), + max_results: z + .number() + .min(1) + .max(100) + .optional() + .default(10) + .describe('Maximum number of results to return (1-100)'), + sort_by: z + .enum(['relevance', 'lastUpdatedDate', 'submittedDate']) + .optional() + .default('relevance') + .describe('Sort results by criteria'), + sort_order: z + .enum(['ascending', 'descending']) + .optional() + .default('descending') + .describe('Sort order'), + start: z + .number() + .min(0) + .optional() + .default(0) + .describe('Starting index for results (0-based)'), + include_abstract: z + .boolean() + .optional() + .default(true) + .describe('Whether to include paper abstracts in results'), + }), + outputSchema: z.object({ + papers: z.array( + z.object({ + id: z.string(), + title: z.string(), + authors: z.array(z.string()), + abstract: z.string().optional(), + categories: z.array(z.string()), + published: z.string(), + updated: z.string().optional(), + pdf_url: z.string(), + doi: z.string().optional(), + journal_ref: z.string().optional(), + comments: z.string().optional(), + }) + ), + total_results: z.number(), + start_index: z.number(), + max_results: z.number(), + }), + 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 (inputData.title !== undefined && inputData.title !== null) { - searchTerms.push(`ti:${inputData.title}`) - } + if (requestCtx) { + log.debug('ArXiv search request context available', { + tool: 'arxiv', + }) + } - if (inputData.category !== undefined) { - searchTerms.push(`cat:${inputData.category}`) - } + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'arxiv-search', + input: inputData, + metadata: { + 'tool.id': 'arxiv', + 'tool.input.query': inputData.query, + 'tool.input.id': inputData.id, + 'tool.input.category': inputData.category, + 'tool.input.maxResults': inputData.max_results, + }, + }) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `📚 Searching arXiv for "${inputData.query ?? inputData.id ?? 'papers'}"`, + stage: 'arxiv', + }, + id: 'arxiv', + }) + toolCallCounters.set('arxiv', (toolCallCounters.get('arxiv') ?? 0) + 1) + try { + const params = new URLSearchParams() - if (inputData.id !== undefined && inputData.id !== null) { - // If specific ID is provided, use it directly - params.append('id_list', inputData.id) - } else if (searchTerms.length > 0) { - params.append('search_query', searchTerms.join(' AND ')) - } else { - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: '❌ No search terms provided', - stage: 'arxiv', - }, - id: 'arxiv', - }) - throw new Error( - 'Either query, id, author, title, or category must be provided' - ) - } + // Build search query + const searchTerms: string[] = [] - // Add other parameters - params.append('max_results', String(inputData.max_results ?? 10)) - params.append('start', String(inputData.start ?? 0)) - - if (inputData.sort_by !== undefined) { - params.append( - 'sortBy', - inputData.sort_by === 'lastUpdatedDate' - ? 'lastUpdatedDate' - : inputData.sort_by === 'submittedDate' - ? 'submittedDate' - : 'relevance' - ) - } + if (inputData.query !== undefined && inputData.query !== null) { + searchTerms.push(inputData.query) + } - if (inputData.sort_order !== undefined) { - params.append('sortOrder', inputData.sort_order) - } + if (inputData.author !== undefined && inputData.author !== null) { + searchTerms.push(`au:${inputData.author}`) + } - const url = `http://export.arxiv.org/api/query?${params.toString()}` + if (inputData.title !== undefined && inputData.title !== null) { + searchTerms.push(`ti:${inputData.title}`) + } - // Check for cancellation before API call - if (abortSignal?.aborted) { - span.setStatus({ - code: 2, - message: 'Operation cancelled during API call', - }) - span.end() - throw new Error('ArXiv search cancelled during API call') - } + if (inputData.category !== undefined) { + searchTerms.push(`cat:${inputData.category}`) + } - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '📡 Fetching from arXiv API...', - stage: 'arxiv', - }, - id: 'arxiv', + if (inputData.id !== undefined && inputData.id !== null) { + // If specific ID is provided, use it directly + params.append('id_list', inputData.id) + } else if (searchTerms.length > 0) { + params.append('search_query', searchTerms.join(' AND ')) + } else { + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: '❌ No search terms provided', + stage: 'arxiv', + }, + id: 'arxiv', + }) + throw new Error( + 'Either query, id, author, title, or category must be provided' + ) + } + + // Add other parameters + params.append('max_results', String(inputData.max_results ?? 10)) + params.append('start', String(inputData.start ?? 0)) + + if (inputData.sort_by !== undefined) { + params.append( + 'sortBy', + inputData.sort_by === 'lastUpdatedDate' + ? 'lastUpdatedDate' + : inputData.sort_by === 'submittedDate' + ? 'submittedDate' + : 'relevance' + ) + } + + if (inputData.sort_order !== undefined) { + params.append('sortOrder', inputData.sort_order) + } + + const url = `http://export.arxiv.org/api/query?${params.toString()}` + + // Check for cancellation before API call + if (abortSignal?.aborted) { + if (span) { + const spanAny = span as any + if (typeof spanAny.setStatus === 'function') { + spanAny.setStatus({ + code: 2, + message: 'Operation cancelled during API call', }) - const response = await fetch(url, { signal: abortSignal }) - - if (!response.ok) { - throw new Error( - `ArXiv API error: ${response.status} ${response.statusText}` - ) - } + } else { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['error.code'] = 2 + spanAny.attributes['error.message'] = + 'Operation cancelled during API call' + } + if (typeof spanAny.end === 'function') { + spanAny.end() + } + } + throw new Error('ArXiv search cancelled during API call') + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '📡 Fetching from arXiv API...', + stage: 'arxiv', + }, + id: 'arxiv', + }) + const response = await fetch(url, { signal: abortSignal }) + + if (!response.ok) { + throw new Error( + `ArXiv API error: ${response.status} ${response.statusText}` + ) + } + + const xmlText = await response.text() + + // Parse XML response (simplified parsing - in production you might want a proper XML parser) + const papers = parseArxivXml( + xmlText, + Boolean(inputData.include_abstract) + ) + + // Extract total results from XML + const totalResultsRegex = + /(\d+)<\/opensearch:totalResults>/ + const totalResultsMatch = totalResultsRegex.exec(xmlText) + const totalResults = totalResultsMatch + ? parseInt(totalResultsMatch[1], 10) + : papers.length + + const startIndexRegex = + /(\d+)<\/opensearch:startIndex>/ + const startIndexMatch = startIndexRegex.exec(xmlText) + const startIndex = startIndexMatch + ? parseInt(startIndexMatch[1], 10) + : (inputData.start ?? 0) + + const result = { + papers, + total_results: totalResults, + start_index: startIndex, + max_results: inputData.max_results ?? 10, + } + if (span) { + const spanAny = span as any + if (typeof spanAny.setAttribute === 'function') { + spanAny.setAttribute('tool.output.totalResults', totalResults) + spanAny.setAttribute('tool.output.paperCount', papers.length) + } else { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['tool.output.totalResults'] = totalResults + spanAny.attributes['tool.output.paperCount'] = papers.length + } + if (typeof spanAny.end === 'function') { + spanAny.end() + } + } + return result + } catch (error) { + // Handle AbortError specifically + if (error instanceof Error && error.name === 'AbortError') { + const cancelMessage = `ArXiv search cancelled for "${inputData.query ?? inputData.id ?? 'papers'}"` + if (span) { + const spanAny = span as any + if (typeof spanAny.setStatus === 'function') { + spanAny.setStatus({ code: 2, message: cancelMessage }) + } else { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['error.code'] = 2 + spanAny.attributes['error.message'] = cancelMessage + } + if (typeof spanAny.end === 'function') { + spanAny.end() + } + } - const xmlText = await response.text() + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `🛑 ${cancelMessage}`, + stage: 'arxiv', + }, + id: 'arxiv', + }) - // Parse XML response (simplified parsing - in production you might want a proper XML parser) - const papers = parseArxivXml( - xmlText, - Boolean(inputData.include_abstract) - ) + log.warn(cancelMessage) + throw new Error(cancelMessage) + } + + const errorMessage = + error instanceof Error + ? error.message + : 'Unknown error occurred' + if (span) { + const spanAny = span as any + if (error instanceof Error && typeof spanAny.recordException === 'function') { + spanAny.recordException(error) + } else if (error instanceof Error) { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['error.exception'] = error.message + } - // Extract total results from XML - const totalResultsRegex = - /(\d+)<\/opensearch:totalResults>/ - const totalResultsMatch = totalResultsRegex.exec(xmlText) - const totalResults = totalResultsMatch - ? parseInt(totalResultsMatch[1], 10) - : papers.length - - const startIndexRegex = - /(\d+)<\/opensearch:startIndex>/ - const startIndexMatch = startIndexRegex.exec(xmlText) - const startIndex = startIndexMatch - ? parseInt(startIndexMatch[1], 10) - : (inputData.start ?? 0) - - const result = { - papers, - total_results: totalResults, - start_index: startIndex, - max_results: inputData.max_results ?? 10, - } - span.setAttribute('tool.output.totalResults', totalResults) - span.setAttribute('tool.output.paperCount', papers.length) - span.end() - return result - } catch (error) { - // Handle AbortError specifically - if (error instanceof Error && error.name === 'AbortError') { - const cancelMessage = `ArXiv search cancelled for "${inputData.query ?? inputData.id ?? 'papers'}"` - span.setStatus({ code: 2, message: cancelMessage }) - span.end() - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `🛑 ${cancelMessage}`, - stage: 'arxiv', - }, - id: 'arxiv', - }) - - log.warn(cancelMessage) - throw new Error(cancelMessage) - } + if (typeof spanAny.setStatus === 'function') { + spanAny.setStatus({ code: 2, message: errorMessage }) // ERROR status + } else { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['error.code'] = 2 + spanAny.attributes['error.message'] = errorMessage + } - const errorMessage = - error instanceof Error - ? error.message - : 'Unknown error occurred' - if (error instanceof Error) { - span.recordException(error) - } - span.setStatus({ code: 2, message: errorMessage }) // ERROR status - span.end() - throw error instanceof Error ? error : new Error(errorMessage) + if (typeof spanAny.end === 'function') { + spanAny.end() } - }, - 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', - }) - }, + } + 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 @@ -478,288 +533,362 @@ export type ArxivUITool = InferUITool */ export const arxivPdfParserTool = createTool({ - id: 'arxiv-pdf-parser', - description: - 'Download and parse arXiv PDF papers to markdown format using unpdf', - inputSchema: z.object({ - arxivId: z - .string() - .describe("ArXiv paper ID (e.g., '2103.12345' or '2103.12345v1')"), - maxPages: z - .number() - .min(1) - .max(1000) - .optional() - .default(100) - .describe('Maximum number of pages to parse (default: 100)'), - includeMetadata: z - .boolean() - .optional() - .default(true) - .describe('Include paper metadata in markdown frontmatter'), - normalizeText: z - .boolean() - .optional() - .default(true) - .describe('Apply text normalization and cleanup'), + id: 'arxiv-pdf-parser', + description: + 'Download and parse arXiv PDF papers to markdown format using unpdf', + inputSchema: z.object({ + arxivId: z + .string() + .describe("ArXiv paper ID (e.g., '2103.12345' or '2103.12345v1')"), + maxPages: z + .number() + .min(1) + .max(1000) + .optional() + .default(100) + .describe('Maximum number of pages to parse (default: 100)'), + includeMetadata: z + .boolean() + .optional() + .default(true) + .describe('Include paper metadata in markdown frontmatter'), + normalizeText: z + .boolean() + .optional() + .default(true) + .describe('Apply text normalization and cleanup'), + }), + outputSchema: z.object({ + success: z.boolean(), + arxivId: z.string(), + markdown: z.string(), + metadata: z + .object({ + title: z.string().optional(), + authors: z.array(z.string()).optional(), + pageCount: z.number(), + extractedAt: z.string(), + }) + .optional(), + statistics: z.object({ + pageCount: z.number(), + textLength: z.number(), + processingTimeMs: z.number(), }), - outputSchema: z.object({ - success: z.boolean(), - arxivId: z.string(), - markdown: z.string(), - metadata: z - .object({ - title: z.string().optional(), - authors: z.array(z.string()).optional(), - pageCount: z.number(), - extractedAt: z.string(), - }) - .optional(), - statistics: z.object({ - pageCount: z.number(), - textLength: z.number(), - processingTimeMs: z.number(), - }), - }), - 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', { - attributes: { - 'tool.id': 'arxiv-pdf-parser', - 'tool.input.arxivId': inputData.arxivId, - 'tool.input.maxPages': inputData.maxPages, - }, - }) + }), + execute: async (inputData, context) => { + const writer = context?.writer + const requestCtx = + context?.requestContext as ArxivRequestContext | undefined + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'arxiv-pdf-parser', + input: inputData, + metadata: { + 'tool.id': 'arxiv-pdf-parser', + 'tool.input.arxivId': inputData.arxivId, + 'tool.input.maxPages': inputData.maxPages, + 'user.id': requestCtx?.userId, + }, + }) + const abortSignal = context?.abortSignal + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: + '🚀 Starting arXiv PDF parser for ' + inputData.arxivId, + stage: 'arxiv-pdf-parser', + }, + id: 'arxiv-pdf-parser', + }) + toolCallCounters.set( + 'arxiv-pdf-parser', + (toolCallCounters.get('arxiv-pdf-parser') ?? 0) + 1 + ) + const startTime = Date.now() + try { + 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}` + + // Abort early if requested + if (abortSignal?.aborted) { + const cancelMessage = `ArXiv PDF parsing cancelled for ${inputData.arxivId}` + if (span) { + if (typeof (span as any).setStatus === 'function') { + ; (span as any).setStatus({ code: 2, message: cancelMessage }) + } + if (typeof (span as any).end === 'function') { + ; (span as any).end() + } + } + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `🛑 ${cancelMessage}`, + stage: 'arxiv-pdf-parser', + }, + id: 'arxiv-pdf-parser', + }) + log.warn(cancelMessage) + throw new Error(cancelMessage) + } + + // Download PDF + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🔄 Extracting text from PDF...', + stage: 'arxiv-pdf-parser', + }, + id: 'arxiv-pdf-parser', + }) + const response = await fetch(pdfUrl, { signal: abortSignal }) + + if (!response.ok) { await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: - '🚀 Starting arXiv PDF parser for ' + inputData.arxivId, - stage: 'arxiv-pdf-parser', - }, - id: 'arxiv-pdf-parser', + type: 'data-tool-progress', + data: { + status: 'done', + message: '❌ PDF download failed', + stage: 'arxiv-pdf-parser', + }, + id: 'arxiv-pdf-parser', }) - toolCallCounters.set( - 'arxiv-pdf-parser', - (toolCallCounters.get('arxiv-pdf-parser') ?? 0) + 1 + if (response.status === 404) { + throw new Error(`Paper not found: ${inputData.arxivId}`) + } + throw new Error( + `Failed to download PDF: ${response.status} ${response.statusText}` ) - const startTime = Date.now() + } - try { - 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 writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🔄 Extracting text from PDF...', - stage: 'arxiv-pdf-parser', - }, - id: 'arxiv-pdf-parser', - }) - const response = await fetch(pdfUrl) - - if (!response.ok) { - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: '❌ PDF download failed', - stage: 'arxiv-pdf-parser', - }, - id: 'arxiv-pdf-parser', - }) - if (response.status === 404) { - throw new Error(`Paper not found: ${inputData.arxivId}`) - } - throw new Error( - `Failed to download PDF: ${response.status} ${response.statusText}` - ) - } + const pdfBuffer = Buffer.from(await response.arrayBuffer()) - const pdfBuffer = Buffer.from(await response.arrayBuffer()) + // Extract text from PDF + const pdfContent = await extractPdfText( + pdfBuffer, + inputData.maxPages ?? 100 + ) - // Extract text from PDF - const pdfContent = await extractPdfText( - pdfBuffer, - inputData.maxPages ?? 100 - ) + // Try to get metadata from arXiv API + let paperMetadata: + | { title?: string; authors?: string[] } + | undefined + if (inputData.includeMetadata) { + try { + const apiUrl = `http://export.arxiv.org/api/query?id_list=${inputData.arxivId}&max_results=1` + const apiResponse = await fetch(apiUrl, { signal: abortSignal }) - // Try to get metadata from arXiv API - let paperMetadata: - | { title?: string; authors?: string[] } - | undefined - if (inputData.includeMetadata) { - try { - const apiUrl = `http://export.arxiv.org/api/query?id_list=${inputData.arxivId}&max_results=1` - const apiResponse = await fetch(apiUrl) - - if (apiResponse.ok) { - const xmlText = await apiResponse.text() - const titleRegex = /(.*?)<\/title>/ - const titleMatch = titleRegex.exec(xmlText) - const title = titleMatch - ? titleMatch[1] - .replace(/<!\[CDATA\[|\]\]>/g, '') - .trim() - : undefined - - const authors: string[] = [] - const authorRegex = - /<author><name>(.*?)<\/name><\/author>/g - let authorMatch - while ( - (authorMatch = authorRegex.exec(xmlText)) !== null - ) { - authors.push(authorMatch[1]) - } - - if ( - (title !== undefined && - title !== null && - title.length > 0) || - authors.length > 0 - ) { - paperMetadata = { - title, - authors: - authors.length > 0 ? authors : undefined, - } - } - } - } catch { - // Metadata fetch failed, continue without it - } + if (apiResponse.ok) { + const xmlText = await apiResponse.text() + const titleRegex = /<title>(.*?)<\/title>/ + const titleMatch = titleRegex.exec(xmlText) + const title = titleMatch + ? titleMatch[1] + .replace(/<!\[CDATA\[|\]\]>/g, '') + .trim() + : undefined + + const authors: string[] = [] + const authorRegex = + /<author><name>(.*?)<\/name><\/author>/g + let authorMatch + while ( + (authorMatch = authorRegex.exec(xmlText)) !== null + ) { + authors.push(authorMatch[1]) } - // Convert to markdown - let markdown = pdfContent.text - - if (inputData.normalizeText) { - markdown = convertPdfTextToMarkdown(markdown, paperMetadata) - } else if (paperMetadata && inputData.includeMetadata) { - // Add basic frontmatter even without normalization - const frontmatter = [ - '---', - typeof paperMetadata.title === 'string' && - paperMetadata.title.trim() !== '' - ? `title: "${paperMetadata.title.replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"` - : null, - paperMetadata.authors - ? `authors: ${JSON.stringify(paperMetadata.authors)}` - : null, - 'source: "arxiv-pdf"', - 'extracted_at: "' + new Date().toISOString() + '"', - '---', - '', - ] - .filter(Boolean) - .join('\n') - - markdown = frontmatter + markdown + if ( + (title !== undefined && + title !== null && + title.length > 0) || + authors.length > 0 + ) { + paperMetadata = { + title, + authors: + authors.length > 0 ? authors : undefined, + } } + } + } catch { + // Metadata fetch failed, continue without it + } + } - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: - '✅ PDF parsing complete: ' + - pdfContent.numpages + - ' pages', - stage: 'arxiv-pdf-parser', - }, - id: 'arxiv-pdf-parser', - }) - const processingTime = Date.now() - startTime - - const result = { - success: true, - arxivId: inputData.arxivId, - markdown, - metadata: inputData.includeMetadata - ? { - title: paperMetadata?.title, - authors: paperMetadata?.authors, - pageCount: pdfContent.numpages, - extractedAt: new Date().toISOString(), - } - : undefined, - statistics: { - pageCount: pdfContent.numpages, - textLength: markdown.length, - processingTimeMs: processingTime, - }, - } - span.setAttribute('tool.output.success', true) - span.setAttribute('tool.output.pageCount', pdfContent.numpages) - span.end() - return result - } catch (error) { - const errorMessage = - error instanceof Error - ? error.message - : 'Unknown error occurred' - if (error instanceof Error) { - span.recordException(error) - } - span.setStatus({ code: 2, message: errorMessage }) // ERROR status - span.end() + // Convert to markdown + let markdown = pdfContent.text + + if (inputData.normalizeText) { + markdown = convertPdfTextToMarkdown(markdown, paperMetadata) + } else if (paperMetadata && inputData.includeMetadata) { + // Add basic frontmatter even without normalization + const frontmatter = [ + '---', + typeof paperMetadata.title === 'string' && + paperMetadata.title.trim() !== '' + ? `title: "${paperMetadata.title.replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"` + : null, + paperMetadata.authors + ? `authors: ${JSON.stringify(paperMetadata.authors)}` + : null, + 'source: "arxiv-pdf"', + 'extracted_at: "' + new Date().toISOString() + '"', + '---', + '', + ] + .filter(Boolean) + .join('\n') - throw error instanceof Error ? error : new Error(errorMessage) + markdown = frontmatter + markdown + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: + '✅ PDF parsing complete: ' + + pdfContent.numpages + + ' pages', + stage: 'arxiv-pdf-parser', + }, + id: 'arxiv-pdf-parser', + }) + const processingTime = Date.now() - startTime + + const result = { + success: true, + arxivId: inputData.arxivId, + markdown, + metadata: inputData.includeMetadata + ? { + title: paperMetadata?.title, + authors: paperMetadata?.authors, + pageCount: pdfContent.numpages, + extractedAt: new Date().toISOString(), + } + : undefined, + statistics: { + pageCount: pdfContent.numpages, + textLength: markdown.length, + processingTimeMs: processingTime, + }, + } + + if (span) { + // Prefer using setAttribute if available; otherwise safely populate the attributes bag + const spanAny = span as any + if (typeof spanAny.setAttribute === 'function') { + spanAny.setAttribute('tool.output.success', true) + spanAny.setAttribute('tool.output.pageCount', pdfContent.numpages) + } else { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['tool.output.success'] = true + spanAny.attributes['tool.output.pageCount'] = pdfContent.numpages } - }, - 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', + if (typeof spanAny.end === 'function') { + spanAny.end() + } + } + return result + } catch (error) { + // Handle AbortError specifically + if (error instanceof Error && error.name === 'AbortError') { + const cancelMessage = `ArXiv PDF parsing cancelled for ${inputData.arxivId}` + if (span) { + if (typeof (span as any).setStatus === 'function') { + ; (span as any).setStatus({ code: 2, message: cancelMessage }) + } + if (typeof (span as any).end === 'function') { + ; (span as any).end() + } + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `🛑 ${cancelMessage}`, + stage: 'arxiv-pdf-parser', + }, + id: 'arxiv-pdf-parser', }) - }, + + log.warn(cancelMessage) + throw new Error(cancelMessage) + } + + const errorMessage = + error instanceof Error + ? error.message + : 'Unknown error occurred' + if (span) { + if (error instanceof Error && typeof (span as any).recordException === 'function') { + ; (span as any).recordException(error) + } + if (typeof (span as any).setStatus === 'function') { + ; (span as any).setStatus({ code: 2, message: errorMessage }) + } + if (typeof (span as any).end === 'function') { + ; (span as any).end() + } + } + + 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<typeof arxivPdfParserTool> @@ -771,269 +900,319 @@ export type ArxivPdfParserUITool = InferUITool<typeof arxivPdfParserTool> */ export const arxivPaperDownloaderTool = createTool({ - id: 'arxiv-paper-downloader', - description: - 'Download complete arXiv papers with metadata and optional PDF parsing to markdown', - inputSchema: z.object({ - arxivId: z - .string() - .describe("ArXiv paper ID (e.g., '2103.12345' or '2103.12345v1')"), - includePdfContent: z - .boolean() - .optional() - .default(false) - .describe('Whether to download and parse the PDF content'), - maxPages: z - .number() - .min(1) - .max(1000) - .optional() - .default(100) - .describe('Maximum pages to parse if includePdfContent is true'), - format: z - .enum(['metadata', 'markdown', 'both']) - .optional() - .default('both') - .describe('Output format: metadata only, markdown only, or both'), - }), - outputSchema: z.object({ - success: z.boolean(), - arxivId: z.string(), - metadata: z - .object({ - id: z.string(), - title: z.string(), - authors: z.array(z.string()), - abstract: z.string().optional(), - categories: z.array(z.string()), - published: z.string(), - updated: z.string().optional(), - pdf_url: z.string(), - doi: z.string().optional(), - journal_ref: z.string().optional(), - comments: z.string().optional(), - }) - .optional(), - pdfContent: z - .object({ - markdown: z.string(), - pageCount: z.number(), - textLength: z.number(), - }) - .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', { - attributes: { - 'tool.id': 'arxiv-paper-downloader', - 'tool.input.arxivId': inputData.arxivId, - 'tool.input.includePdfContent': inputData.includePdfContent, - 'tool.input.format': inputData.format, - }, - }) - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: - '🚀 Starting arXiv paper downloader for ' + - inputData.arxivId, - stage: 'arxiv-paper-downloader', - }, - id: 'arxiv-paper-downloader', + id: 'arxiv-paper-downloader', + description: + 'Download complete arXiv papers with metadata and optional PDF parsing to markdown', + inputSchema: z.object({ + arxivId: z + .string() + .describe("ArXiv paper ID (e.g., '2103.12345' or '2103.12345v1')"), + includePdfContent: z + .boolean() + .optional() + .default(false) + .describe('Whether to download and parse the PDF content'), + maxPages: z + .number() + .min(1) + .max(1000) + .optional() + .default(100) + .describe('Maximum pages to parse if includePdfContent is true'), + format: z + .enum(['metadata', 'markdown', 'both']) + .optional() + .default('both') + .describe('Output format: metadata only, markdown only, or both'), + }), + outputSchema: z.object({ + success: z.boolean(), + arxivId: z.string(), + metadata: z + .object({ + id: z.string(), + title: z.string(), + authors: z.array(z.string()), + abstract: z.string().optional(), + categories: z.array(z.string()), + published: z.string(), + updated: z.string().optional(), + pdf_url: z.string(), + doi: z.string().optional(), + journal_ref: z.string().optional(), + comments: z.string().optional(), + }) + .optional(), + pdfContent: z + .object({ + markdown: z.string(), + pageCount: z.number(), + textLength: z.number(), + }) + .optional(), + }), + execute: async (inputData, context) => { + const writer = context?.writer + const abortSignal = context?.abortSignal + const requestCtx = + context?.requestContext as ArxivRequestContext | undefined + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'arxiv-paper-downloader', + input: inputData, + metadata: { + 'tool.id': 'arxiv-paper-downloader', + 'tool.input.arxivId': inputData.arxivId, + 'tool.input.includePdfContent': inputData.includePdfContent, + 'tool.input.format': inputData.format, + 'user.id': requestCtx?.userId, + }, + }) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: + '🚀 Starting arXiv paper downloader for ' + + inputData.arxivId, + stage: 'arxiv-paper-downloader', + }, + id: 'arxiv-paper-downloader', + }) + toolCallCounters.set( + 'arxiv-paper-downloader', + (toolCallCounters.get('arxiv-paper-downloader') ?? 0) + 1 + ) + try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing arXiv paper downloader for user', { + userId: requestCtx.userId, }) - toolCallCounters.set( - 'arxiv-paper-downloader', - (toolCallCounters.get('arxiv-paper-downloader') ?? 0) + 1 + } + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '📡 Fetching paper metadata from arXiv API...', + stage: 'arxiv-paper-downloader', + }, + id: 'arxiv-paper-downloader', + }) + // 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, { signal: abortSignal }) + + if (!apiResponse.ok) { + throw new Error( + `Failed to fetch paper metadata: ${apiResponse.status} ${apiResponse.statusText}` ) - try { - 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', - message: '📡 Fetching paper metadata from arXiv API...', - stage: 'arxiv-paper-downloader', - }, - id: 'arxiv-paper-downloader', - }) - // 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, { signal: abortSignal }) - - if (!apiResponse.ok) { - throw new Error( - `Failed to fetch paper metadata: ${apiResponse.status} ${apiResponse.statusText}` - ) - } + } - const xmlText = await apiResponse.text() + const xmlText = await apiResponse.text() - const papers = parseArxivXml(xmlText, true) + const papers = parseArxivXml(xmlText, true) - if (papers?.length === 0) { - throw new Error(`Paper not found: ${inputData.arxivId}`) - } + if (papers?.length === 0) { + throw new Error(`Paper not found: ${inputData.arxivId}`) + } - const paperMetadata = papers[0] + const paperMetadata = papers[0] - let pdfContent: - | { markdown: string; pageCount: number; textLength: number } - | undefined + let pdfContent: + | { markdown: string; pageCount: number; textLength: number } + | undefined - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🔄 Processing paper data...', - stage: 'arxiv-paper-downloader', - }, - id: 'arxiv-paper-downloader', - }) - // Download and parse PDF if requested - if (inputData.includePdfContent) { - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '📥 Downloading PDF content...', - stage: 'arxiv-paper-downloader', - }, - id: 'arxiv-paper-downloader', - }) - try { - const pdfResponse = await fetch(paperMetadata.pdf_url) - - if (pdfResponse.ok) { - const pdfBuffer = Buffer.from( - await pdfResponse.arrayBuffer() - ) - const extracted = await extractPdfText( - pdfBuffer, - inputData.maxPages ?? 100 - ) - - const markdown = convertPdfTextToMarkdown( - extracted.text, - { - title: paperMetadata.title, - authors: paperMetadata.authors, - } - ) - - pdfContent = { - markdown, - pageCount: extracted.numpages, - textLength: markdown.length, - } - } else { - // PDF download failed, continue without it - } - } catch { - // PDF parsing failed, continue without it - } - } + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🔄 Processing paper data...', + stage: 'arxiv-paper-downloader', + }, + id: 'arxiv-paper-downloader', + }) + // Download and parse PDF if requested + if (inputData.includePdfContent) { + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '📥 Downloading PDF content...', + stage: 'arxiv-paper-downloader', + }, + id: 'arxiv-paper-downloader', + }) + try { + const pdfResponse = await fetch(paperMetadata.pdf_url, { + signal: abortSignal, + }) - // Return based on requested format - const result: { - success: boolean - arxivId: string - metadata?: typeof paperMetadata - pdfContent?: { - markdown: string - pageCount: number - textLength: number - } - } = { - success: true, - arxivId: inputData.arxivId, - } + if (pdfResponse.ok) { + const pdfBuffer = Buffer.from( + await pdfResponse.arrayBuffer() + ) + const extracted = await extractPdfText( + pdfBuffer, + inputData.maxPages ?? 100 + ) - if ( - inputData.format === 'metadata' || - inputData.format === 'both' - ) { - result.metadata = paperMetadata - } + const markdown = convertPdfTextToMarkdown( + extracted.text, + { + title: paperMetadata.title, + authors: paperMetadata.authors, + } + ) - if ( - (inputData.format === 'markdown' || - inputData.format === 'both') && - pdfContent - ) { - result.pdfContent = pdfContent + pdfContent = { + markdown, + pageCount: extracted.numpages, + textLength: markdown.length, } + } else { + // PDF download failed, continue without it + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: '❌ PDF download failed', + stage: 'arxiv-paper-downloader', + }, + id: 'arxiv-paper-downloader', + }) + log.warn( + `PDF download failed for ${inputData.arxivId}: ${pdfResponse.status} ${pdfResponse.statusText}` + ) + } + } catch (err) { + // PDF parsing or network failure, continue without it but log + const errMsg = + err instanceof Error ? err.message : 'Unknown error' + log.warn( + `PDF download/parse failed for ${inputData.arxivId}: ${errMsg}` + ) + } + } + + // Return based on requested format + const result: { + success: boolean + arxivId: string + metadata?: typeof paperMetadata + pdfContent?: { + markdown: string + pageCount: number + textLength: number + } + } = { + success: true, + arxivId: inputData.arxivId, + } + + if ( + inputData.format === 'metadata' || + inputData.format === 'both' + ) { + result.metadata = paperMetadata + } + + if ( + (inputData.format === 'markdown' || + inputData.format === 'both') && + pdfContent + ) { + result.pdfContent = pdfContent + } + + if (span) { + const spanAny = span as any + if (typeof spanAny.setAttribute === 'function') { + spanAny.setAttribute('tool.output.success', true) + spanAny.setAttribute('tool.output.hasPdfContent', !!pdfContent) + } else { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['tool.output.success'] = true + spanAny.attributes['tool.output.hasPdfContent'] = !!pdfContent + } + if (typeof spanAny.end === 'function') { + spanAny.end() + } + } + return result + } catch (error) { + const errorMessage = + error instanceof Error + ? error.message + : 'Unknown error occurred' + if (span) { + const spanAny = span as any + if (error instanceof Error && typeof spanAny.recordException === 'function') { + spanAny.recordException(error) + } else if (error instanceof Error) { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['error.exception'] = error.message + } - span.setAttribute('tool.output.success', true) - span.end() - return result - } catch (error) { - const errorMessage = - error instanceof Error - ? error.message - : 'Unknown error occurred' - if (error instanceof Error) { - span.recordException(error) - } - span.setStatus({ code: 2, message: errorMessage }) // ERROR status - span.end() - throw error instanceof Error ? error : new Error(errorMessage) + if (typeof spanAny.setStatus === 'function') { + spanAny.setStatus({ code: 2, message: errorMessage }) + } else { + spanAny.attributes = spanAny.attributes ?? {} + spanAny.attributes['error.code'] = 2 + spanAny.attributes['error.message'] = 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', - }) - }, + + if (typeof spanAny.end === 'function') { + spanAny.end() + } + } + 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< - typeof arxivPaperDownloaderTool + typeof arxivPaperDownloaderTool > /** @@ -1042,9 +1221,22 @@ export type ArxivPaperDownloaderUITool = InferUITool< */ // FIXME: Use a proper XML parser instead of regex @copilot fast-xml-parser function parseArxivXml( - xmlText: string, - includeAbstract: boolean + xmlText: string, + includeAbstract: boolean ): Array<{ + id: string + title: string + authors: string[] + abstract?: string + categories: string[] + published: string + updated?: string + pdf_url: string + doi?: string + journal_ref?: string + comments?: string +}> { + const papers: Array<{ id: string title: string authors: string[] @@ -1056,93 +1248,80 @@ function parseArxivXml( doi?: string journal_ref?: string comments?: string -}> { - const papers: Array<{ - id: string - title: string - authors: string[] - abstract?: string - categories: string[] - published: string - updated?: string - pdf_url: string - doi?: string - journal_ref?: string - comments?: string - }> = [] - - // Split by entry tags - const entryRegex = /<entry>([\s\S]*?)<\/entry>/g - let match - - while ((match = entryRegex.exec(xmlText)) !== null) { - const entryXml = match[1] - - // Extract fields using regex - const idRegex = /<id>(.*?)<\/id>/ - const idMatch = idRegex.exec(entryXml) - - const titleRegex = /<title>(.*?)<\/title>/ - const titleMatch = titleRegex.exec(entryXml) - - const publishedRegex = /<published>(.*?)<\/published>/ - const publishedMatch = publishedRegex.exec(entryXml) - - const updatedRegex = /<updated>(.*?)<\/updated>/ - const updatedMatch = updatedRegex.exec(entryXml) - - const summaryRegex = /<summary>(.*?)<\/summary>/ - const summaryMatch = summaryRegex.exec(entryXml) - - const doiRegex = /<arxiv:doi>(.*?)<\/arxiv:doi>/ - const doiMatch = doiRegex.exec(entryXml) - - const journalRefRegex = /<arxiv:journal_ref>(.*?)<\/arxiv:journal_ref>/ - const journalRefMatch = journalRefRegex.exec(entryXml) - - const commentsRegex = /<arxiv:comment>(.*?)<\/arxiv:comment>/ - const commentsMatch = commentsRegex.exec(entryXml) - - if (!idMatch || !titleMatch || !publishedMatch) { - continue // Skip malformed entries - } + }> = [] - // Extract authors - const authors: string[] = [] - const authorRegex = /<author><name>(.*?)<\/name><\/author>/g - let authorMatch - while ((authorMatch = authorRegex.exec(entryXml)) !== null) { - authors.push(authorMatch[1]) - } + // Split by entry tags + const entryRegex = /<entry>([\s\S]*?)<\/entry>/g + let match - // Extract categories - const categories: string[] = [] - const categoryRegex = /<category term="(.*?)"[^>]*\/>/g - let categoryMatch - while ((categoryMatch = categoryRegex.exec(entryXml)) !== null) { - categories.push(categoryMatch[1]) - } + while ((match = entryRegex.exec(xmlText)) !== null) { + const entryXml = match[1] - // Generate PDF URL from ID - const arxivId = idMatch[1].split('/').pop() ?? '' - - papers.push({ - id: arxivId, - title: titleMatch[1].replace(/<!\[CDATA\[|\]\]>/g, '').trim(), - authors, - abstract: - includeAbstract && summaryMatch - ? summaryMatch[1].replace(/<!\[CDATA\[|\]\]>/g, '').trim() - : undefined, - categories, - published: publishedMatch[1], - updated: updatedMatch ? updatedMatch[1] : undefined, - pdf_url: `https://arxiv.org/pdf/${arxivId}`, - doi: doiMatch ? doiMatch[1] : undefined, - journal_ref: journalRefMatch ? journalRefMatch[1] : undefined, - comments: commentsMatch ? commentsMatch[1] : undefined, - }) + // Extract fields using regex + const idRegex = /<id>(.*?)<\/id>/ + const idMatch = idRegex.exec(entryXml) + + const titleRegex = /<title>(.*?)<\/title>/ + const titleMatch = titleRegex.exec(entryXml) + + const publishedRegex = /<published>(.*?)<\/published>/ + const publishedMatch = publishedRegex.exec(entryXml) + + const updatedRegex = /<updated>(.*?)<\/updated>/ + const updatedMatch = updatedRegex.exec(entryXml) + + const summaryRegex = /<summary>(.*?)<\/summary>/ + const summaryMatch = summaryRegex.exec(entryXml) + + const doiRegex = /<arxiv:doi>(.*?)<\/arxiv:doi>/ + const doiMatch = doiRegex.exec(entryXml) + + const journalRefRegex = /<arxiv:journal_ref>(.*?)<\/arxiv:journal_ref>/ + const journalRefMatch = journalRefRegex.exec(entryXml) + + const commentsRegex = /<arxiv:comment>(.*?)<\/arxiv:comment>/ + const commentsMatch = commentsRegex.exec(entryXml) + + if (!idMatch || !titleMatch || !publishedMatch) { + continue // Skip malformed entries + } + + // Extract authors + const authors: string[] = [] + const authorRegex = /<author><name>(.*?)<\/name><\/author>/g + let authorMatch + while ((authorMatch = authorRegex.exec(entryXml)) !== null) { + authors.push(authorMatch[1]) + } + + // Extract categories + const categories: string[] = [] + const categoryRegex = /<category term="(.*?)"[^>]*\/>/g + let categoryMatch + while ((categoryMatch = categoryRegex.exec(entryXml)) !== null) { + categories.push(categoryMatch[1]) } - return papers + // Generate PDF URL from ID + const arxivId = idMatch[1].split('/').pop() ?? '' + + papers.push({ + id: arxivId, + title: titleMatch[1].replace(/<!\[CDATA\[|\]\]>/g, '').trim(), + authors, + abstract: + includeAbstract && summaryMatch + ? summaryMatch[1].replace(/<!\[CDATA\[|\]\]>/g, '').trim() + : undefined, + categories, + published: publishedMatch[1], + updated: updatedMatch ? updatedMatch[1] : undefined, + pdf_url: `https://arxiv.org/pdf/${arxivId}`, + doi: doiMatch ? doiMatch[1] : undefined, + journal_ref: journalRefMatch ? journalRefMatch[1] : undefined, + comments: commentsMatch ? commentsMatch[1] : undefined, + }) + } + + return papers } diff --git a/src/mastra/tools/browser-tool.ts b/src/mastra/tools/browser-tool.ts index 1c547bca..38d003f9 100644 --- a/src/mastra/tools/browser-tool.ts +++ b/src/mastra/tools/browser-tool.ts @@ -1,1180 +1,1243 @@ +import { SpanType } from '@mastra/core/observability' +import type { RequestContext } from '@mastra/core/request-context' 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 type UserTier = 'free' | 'pro' | 'enterprise' + export interface BrowserRequestContext extends RequestContext { - userId?: string; - sessionId?: string; + userId?: string; + sessionId?: string; + 'user-tier': UserTier; } // Browser instance cache for reuse let browserInstance: Browser | null = null async function getBrowser(): Promise<Browser> { - if (!(browserInstance?.isConnected() ?? false)) { - if (browserInstance) { - await browserInstance.close().catch(() => {}) // Close the old instance if it's not connected - } - browserInstance = await chromium.launch({ - headless: true, - chromiumSandbox: false, - }) + if (!(browserInstance?.isConnected() ?? false)) { + if (browserInstance) { + await browserInstance.close().catch(() => { }) // Close the old instance if it's not connected } - return browserInstance! + browserInstance = await chromium.launch({ + headless: true, + chromiumSandbox: false, + }) + } + return browserInstance! } export const browserTool = createTool({ - id: 'browserTool', - description: - 'Browser Tool, opens a browser and navigates to a url capturing the content', - inputSchema: z.object({ - url: z.string(), - }), - outputSchema: z.object({ - message: z.string(), - }), - 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) { - throw new Error('Browser operation cancelled') - } + id: 'browserTool', + description: + 'Browser Tool, opens a browser and navigates to a url capturing the content', + inputSchema: z.object({ + url: z.string(), + }), + outputSchema: z.object({ + message: z.string(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + + // Check if operation was already cancelled + if (abortSignal?.aborted === true) { + throw new Error('Browser operation cancelled') + } - const tracer = trace.getTracer('browser-tool', '1.0.0') - const span = tracer.startSpan('browser-scrape', { - attributes: { - 'tool.id': 'browser-scrape', - 'tool.input.url': inputData.url, - }, + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'browser-scrape', + input: inputData, + metadata: { + 'tool.id': 'browser-scrape', + 'tool.input.url': inputData.url, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `Input: url="${inputData.url}" - 🌐 Launching browser`, + stage: 'browserTool', + }, + id: 'browserTool', + }) + try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing browser scrape for user', { + userId: requestCtx.userId, }) + } - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `Input: url="${inputData.url}" - 🌐 Launching browser`, - stage: 'browserTool', - }, - id: 'browserTool', - }) - try { - if (typeof requestCtx?.userId === 'string') { - log.debug('Executing browser scrape for user', { - userId: requestCtx.userId, - }) - } + const browser = await getBrowser() + const page = await browser.newPage() - const browser = await getBrowser() - const page = await browser.newPage() - - await page.goto(inputData.url, { waitUntil: 'domcontentloaded' }) - - const docs = MDocument.fromHTML(await page.content()) - - await docs.chunk({ - strategy: 'html', - maxSize: 300, - sections: [ - ['h1', 'Header 1'], - ['h2', 'Header 2'], - ['h3', 'Header 3'], - ['h4', 'Header 4'], - ['h5', 'Header 5'], - ['h6', 'Header 6'], - ['p', 'Paragraph'], - ], - }) - - await page.close() - - if (!docs.getText().length) { - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '⚠️ No content found' }, - }) - span.setAttributes({ - 'tool.output.success': false, - 'tool.output.reason': 'No content', - }) - span.end() - return { message: 'No content' } - } + await page.goto(inputData.url, { waitUntil: 'domcontentloaded' }) - const result = docs.getText().join('\n') - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '✅ Content extracted successfully' }, - }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.contentLength': result.length, - }) - span.end() - return { message: result } - } catch (e) { - // Handle AbortError specifically - if (e instanceof Error && e.name === 'AbortError') { - const cancelMessage = `Browser operation cancelled for ${inputData.url}` - span.setStatus({ code: 2, message: cancelMessage }) - span.end() - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `🛑 ${cancelMessage}`, - stage: 'browserTool', - }, - id: 'browserTool', - }) - - log.warn(cancelMessage) - return { message: `Error: ${cancelMessage}` } - } + const docs = MDocument.fromHTML(await page.content()) - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`Browser scrape failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - 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', - }) - }, -}) + await docs.chunk({ + strategy: 'html', + maxSize: 300, + sections: [ + ['h1', 'Header 1'], + ['h2', 'Header 2'], + ['h3', 'Header 3'], + ['h4', 'Header 4'], + ['h5', 'Header 5'], + ['h6', 'Header 6'], + ['p', 'Paragraph'], + ], + }) -export const screenshotTool = createTool({ - id: 'screenshotTool', - description: - 'Capture a screenshot of a webpage. Returns base64-encoded PNG image.', - inputSchema: z.object({ - url: z.string().describe('URL to screenshot'), - fullPage: z - .boolean() - .optional() - .default(false) - .describe('Capture full page or viewport only'), - width: z.number().optional().default(1280).describe('Viewport width'), - height: z.number().optional().default(720).describe('Viewport height'), - }), - outputSchema: z.object({ - success: z.boolean(), - screenshot: z.string().optional(), - 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: { - 'tool.id': 'browser-screenshot', - 'tool.input.url': inputData.url, - 'tool.input.fullPage': inputData.fullPage, - }, - }) + await page.close() + if (!docs.getText().length) { await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: `📸 Taking screenshot of ${inputData.url}` }, + type: 'data-tool-progress', + data: { message: '⚠️ No content found' }, }) - - 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({ - width: inputData.width ?? 1280, - height: inputData.height ?? 720, - }) - await page.goto(inputData.url, { waitUntil: 'networkidle' }) - - const screenshot = await page.screenshot({ - fullPage: inputData.fullPage ?? false, - type: 'png', - }) - - await page.close() - - const base64 = screenshot.toString('base64') - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '✅ Screenshot captured' }, - }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.size': screenshot.length, - }) - span.end() - - return { success: true, screenshot: base64 } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`Screenshot failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - 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', + span?.update({ + output: { message: 'No content' }, + metadata: { + 'tool.output.success': false, + 'tool.output.reason': 'No content', + }, }) - }, - 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({ - id: 'pdfGeneratorTool', - description: 'Generate a PDF from a webpage. Returns base64-encoded PDF.', - inputSchema: z.object({ - url: z.string().describe('URL to convert to PDF'), - format: z.enum(['A4', 'Letter', 'Legal']).optional().default('A4'), - landscape: z.boolean().optional().default(false), - printBackground: z.boolean().optional().default(true), - }), - outputSchema: z.object({ - success: z.boolean(), - pdf: z.string().optional(), - 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: { - 'tool.id': 'browser-pdf-generate', - 'tool.input.url': inputData.url, - 'tool.input.format': inputData.format, - }, + span?.end() + return { message: 'No content' } + } + + const result = docs.getText().join('\n') + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '✅ Content extracted successfully' }, + }) + span?.update({ + output: { contentLength: result.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.contentLength': result.length, + }, + }) + span?.end() + return { message: result } + } catch (e) { + // Handle AbortError specifically + if (e instanceof Error && e.name === 'AbortError') { + const cancelMessage = `Browser operation cancelled for ${inputData.url}` + span?.update({ + metadata: { + 'tool.output.success': false, + 'tool.output.reason': cancelMessage, + } }) + span?.end() await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: `📄 Generating PDF from ${inputData.url}` }, + type: 'data-tool-progress', + data: { + status: 'done', + message: `🛑 ${cancelMessage}`, + stage: 'browserTool', + }, + id: 'browserTool', }) - 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' }) - - const pdf = await page.pdf({ - format: inputData.format ?? 'A4', - landscape: inputData.landscape ?? false, - printBackground: inputData.printBackground ?? true, - }) - - await page.close() - - const base64 = pdf.toString('base64') - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '✅ PDF generated' }, - }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.size': pdf.length, - }) - span.end() - - return { success: true, pdf: base64 } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`PDF generation failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - 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', + log.warn(cancelMessage) + return { message: `Error: ${cancelMessage}` } + } + + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`Browser scrape failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + 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({ + id: 'screenshotTool', + description: + 'Capture a screenshot of a webpage. Returns base64-encoded PNG image.', + inputSchema: z.object({ + url: z.string().describe('URL to screenshot'), + fullPage: z + .boolean() + .optional() + .default(false) + .describe('Capture full page or viewport only'), + width: z.number().optional().default(1280).describe('Viewport width'), + height: z.number().optional().default(720).describe('Viewport height'), + }), + outputSchema: z.object({ + success: z.boolean(), + screenshot: z.string().optional(), + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'browser-screenshot', + input: inputData, + metadata: { + 'tool.id': 'browser-screenshot', + 'tool.input.url': inputData.url, + 'tool.input.fullPage': inputData.fullPage, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: `📸 Taking screenshot of ${inputData.url}` }, + }) + + try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing browser screenshot for user', { + userId: requestCtx.userId, }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log[output.success ? 'info' : 'warn']('PDF generator tool completed', { - toolCallId, - toolName, - success: output.success, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', + } + if (abortSignal?.aborted === true) { + throw new Error('Screenshot operation cancelled') + } + const browser = await getBrowser() + const page = await browser.newPage() + await page.setViewportSize({ + width: inputData.width ?? 1280, + height: inputData.height ?? 720, + }) + await page.goto(inputData.url, { waitUntil: 'networkidle' }) + + const screenshot = await page.screenshot({ + fullPage: inputData.fullPage ?? false, + type: 'png', + }) + + await page.close() + + const base64 = screenshot.toString('base64') + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '✅ Screenshot captured' }, + }) + span?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true, + 'tool.output.size': screenshot.length, + }, + }) + span?.end() + + return { success: true, screenshot: base64 } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`Screenshot failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + 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({ + id: 'pdfGeneratorTool', + description: 'Generate a PDF from a webpage. Returns base64-encoded PDF.', + inputSchema: z.object({ + url: z.string().describe('URL to convert to PDF'), + format: z.enum(['A4', 'Letter', 'Legal']).optional().default('A4'), + landscape: z.boolean().optional().default(false), + printBackground: z.boolean().optional().default(true), + }), + outputSchema: z.object({ + success: z.boolean(), + pdf: z.string().optional(), + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'browser-pdf-generate', + input: inputData, + metadata: { + 'tool.id': 'browser-pdf-generate', + 'tool.input.url': inputData.url, + 'tool.input.format': inputData.format, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: `📄 Generating PDF from ${inputData.url}` }, + }) + + 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' }) + + const pdf = await page.pdf({ + format: inputData.format ?? 'A4', + landscape: inputData.landscape ?? false, + printBackground: inputData.printBackground ?? true, + }) + + await page.close() + + const base64 = pdf.toString('base64') + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '✅ PDF generated' }, + }) + span?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true, + 'tool.output.size': pdf.length, + }, + }) + span?.end() + + return { success: true, pdf: base64 } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`PDF generation failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + 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({ - id: 'clickAndExtractTool', - description: - 'Navigate to a URL, perform click actions, and extract content. Useful for SPAs and dynamic content.', - inputSchema: z.object({ - url: z.string().describe('URL to navigate to'), - clickSelector: z - .string() - .optional() - .describe('CSS selector of element to click before extraction'), - waitForSelector: z - .string() - .optional() - .describe('CSS selector to wait for before extraction'), - extractSelector: z - .string() - .optional() - .describe( - 'CSS selector to extract content from (defaults to body)' - ), - timeout: z - .number() - .optional() - .default(10000) - .describe('Timeout in milliseconds'), - }), - outputSchema: z.object({ - success: z.boolean(), - content: z.string().optional(), - 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: { - 'tool.id': 'browser-click-extract', - 'tool.input.url': inputData.url, - 'tool.input.clickSelector': inputData.clickSelector, - }, + id: 'clickAndExtractTool', + description: + 'Navigate to a URL, perform click actions, and extract content. Useful for SPAs and dynamic content.', + inputSchema: z.object({ + url: z.string().describe('URL to navigate to'), + clickSelector: z + .string() + .optional() + .describe('CSS selector of element to click before extraction'), + waitForSelector: z + .string() + .optional() + .describe('CSS selector to wait for before extraction'), + extractSelector: z + .string() + .optional() + .describe( + 'CSS selector to extract content from (defaults to body)' + ), + timeout: z + .number() + .optional() + .default(10000) + .describe('Timeout in milliseconds'), + }), + outputSchema: z.object({ + success: z.boolean(), + content: z.string().optional(), + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'browser-click-extract', + input: inputData, + metadata: { + 'tool.id': 'browser-click-extract', + 'tool.input.url': inputData.url, + 'tool.input.clickSelector': inputData.clickSelector, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: `🖱️ Navigating to ${inputData.url}` }, + }) + + 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' }) + + if ( + inputData.clickSelector !== undefined && + inputData.clickSelector !== null + ) { await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: `🖱️ Navigating to ${inputData.url}` }, + type: 'data-tool-progress', + data: { message: `🖱️ Clicking ${inputData.clickSelector}` }, }) - - 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' }) - - if ( - inputData.clickSelector !== undefined && - inputData.clickSelector !== null - ) { - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: `🖱️ Clicking ${inputData.clickSelector}` }, - }) - await page.click(inputData.clickSelector, { - timeout: inputData.timeout, - }) - } - - if ( - inputData.waitForSelector !== undefined && - inputData.waitForSelector !== null - ) { - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - message: `⏳ Waiting for ${inputData.waitForSelector}`, - }, - }) - await page.waitForSelector(inputData.waitForSelector, { - timeout: inputData.timeout, - }) - } - - const selector = inputData.extractSelector ?? 'body' - const content = await page.$eval( - selector, - (el) => el.textContent ?? '' - ) - - await page.close() - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '✅ Content extracted' }, - }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.contentLength': content.length, - }) - span.end() - - return { success: true, content: content.trim() } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`Click and extract failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - 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', + await page.click(inputData.clickSelector, { + timeout: inputData.timeout, }) - }, - 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', + } + + if ( + inputData.waitForSelector !== undefined && + inputData.waitForSelector !== null + ) { + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { + message: `⏳ Waiting for ${inputData.waitForSelector}`, + }, }) - }, - 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', + await page.waitForSelector(inputData.waitForSelector, { + timeout: inputData.timeout, }) - }, + } + + const selector = inputData.extractSelector ?? 'body' + const content = await page.$eval( + selector, + (el) => el.textContent ?? '' + ) + + await page.close() + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '✅ Content extracted' }, + }) + span?.update({ + output: { contentLength: content.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.contentLength': content.length, + }, + }) + span?.end() + + return { success: true, content: content.trim() } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`Click and extract failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + 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({ - id: 'fillFormTool', - description: 'Fill out a form on a webpage and optionally submit it.', - inputSchema: z.object({ - url: z.string().describe('URL containing the form'), - fields: z - .array( - z.object({ - selector: z - .string() - .describe('CSS selector for the input field'), - value: z.string().describe('Value to fill in'), - }) - ) - .describe('Form fields to fill'), - submitSelector: z + id: 'fillFormTool', + description: 'Fill out a form on a webpage and optionally submit it.', + inputSchema: z.object({ + url: z.string().describe('URL containing the form'), + fields: z + .array( + z.object({ + selector: z .string() - .optional() - .describe('CSS selector for submit button'), - waitForNavigation: z.boolean().optional().default(false), - }), - outputSchema: z.object({ - success: z.boolean(), - finalUrl: z.string().optional(), - 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: { - 'tool.id': 'browser-fill-form', - 'tool.input.url': inputData.url, - 'tool.input.fieldCount': inputData.fields.length, - }, + .describe('CSS selector for the input field'), + value: z.string().describe('Value to fill in'), }) - + ) + .describe('Form fields to fill'), + submitSelector: z + .string() + .optional() + .describe('CSS selector for submit button'), + waitForNavigation: z.boolean().optional().default(false), + }), + outputSchema: z.object({ + success: z.boolean(), + finalUrl: z.string().optional(), + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'browser-fill-form', + input: inputData, + metadata: { + 'tool.id': 'browser-fill-form', + 'tool.input.url': inputData.url, + 'tool.input.fieldCount': inputData.fields.length, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: `📝 Filling form on ${inputData.url}` }, + }) + + 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' }) + + for (const field of inputData.fields) { + await page.fill(field.selector, field.value) + } + + if ( + inputData.submitSelector !== undefined && + inputData.submitSelector !== null + ) { await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: `📝 Filling form on ${inputData.url}` }, + type: 'data-tool-progress', + data: { message: '📤 Submitting form...' }, }) - - 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' }) - - for (const field of inputData.fields) { - await page.fill(field.selector, field.value) - } - - if ( - inputData.submitSelector !== undefined && - inputData.submitSelector !== null - ) { - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '📤 Submitting form...' }, - }) - if (inputData.waitForNavigation) { - await Promise.all([ - page.waitForNavigation(), - page.click(inputData.submitSelector), - ]) - } else { - await page.click(inputData.submitSelector) - } - } - - const finalUrl = page.url() - await page.close() - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '✅ Form submitted' }, - }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.finalUrl': finalUrl, - }) - span.end() - - return { success: true, finalUrl } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`Form fill failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - return { success: false, message: errorMsg } + if (inputData.waitForNavigation) { + await Promise.all([ + page.waitForNavigation(), + page.click(inputData.submitSelector), + ]) + } else { + await page.click(inputData.submitSelector) } - }, - 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', - }) - }, + } + + const finalUrl = page.url() + await page.close() + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '✅ Form submitted' }, + }) + span?.update({ + output: { finalUrl }, + metadata: { + 'tool.output.success': true, + 'tool.output.finalUrl': finalUrl, + }, + }) + span?.end() + + return { success: true, finalUrl } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`Form fill failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + 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({ - id: 'googleSearch', - description: - 'Google Search. Passes the query to Google and returns the search results.', - inputSchema: z.object({ - query: z.string(), - }), - outputSchema: z.object({ - 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: { - 'tool.id': 'google-search', - 'tool.input.query': inputData.query, - }, + id: 'googleSearch', + description: + 'Google Search. Passes the query to Google and returns the search results.', + inputSchema: z.object({ + query: z.string(), + }), + outputSchema: z.object({ + message: z.string(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-search', + input: inputData, + metadata: { + 'tool.id': 'google-search', + 'tool.input.query': inputData.query, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { + message: + '🔍 Starting Google search for "' + inputData.query + '"', + }, + }) + + try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing Google search for user', { + userId: requestCtx.userId, }) - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - message: - '🔍 Starting Google search for "' + inputData.query + '"', - }, + } + if (abortSignal?.aborted === true) { + throw new Error('Google search operation cancelled') + } + const browser = await getBrowser() + const page = await browser.newPage() + await page.goto( + `https://www.google.com/search?q=${encodeURIComponent(inputData.query)}` + ) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '⏳ Waiting for search results...' }, + }) + + try { + await page.click('button:has-text("Accept all")', { + timeout: 5000, }) + } catch { + // Cookie dialog didn't appear, continue + } - 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( - `https://www.google.com/search?q=${encodeURIComponent(inputData.query)}` - ) + await page.waitForSelector('#search') - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '⏳ Waiting for search results...' }, - }) - - try { - await page.click('button:has-text("Accept all")', { - timeout: 5000, - }) - } catch { - // Cookie dialog didn't appear, continue - } + const text = await page.evaluate(() => { + const links: string[] = [] + const searchResults = document.querySelectorAll('div.g a') - await page.waitForSelector('#search') - - const text = await page.evaluate(() => { - const links: string[] = [] - const searchResults = document.querySelectorAll('div.g a') - - searchResults.forEach((link) => { - const href = link.getAttribute('href') - if (href?.startsWith('http') === true) { - links.push(href) - } - }) - - return links - }) - - await page.close() - - if (!text.length) { - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '⚠️ No results found' }, - }) - span.setAttributes({ - 'tool.output.success': false, - 'tool.output.reason': 'No results', - }) - span.end() - return { message: 'No results' } - } - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: '✅ Found ' + text.length + ' results' }, - }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.resultCount': text.length, - }) - span.end() - return { message: text.join('\n') } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`Google search failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - 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', + searchResults.forEach((link) => { + const href = link.getAttribute('href') + if (href?.startsWith('http') === true) { + links.push(href) + } }) - }, - 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', + + return links + }) + + await page.close() + + if (!text.length) { + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '⚠️ No results found' }, }) - }, - 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', + span?.update({ + output: { message: 'No results' }, + metadata: { + 'tool.output.success': false, + 'tool.output.reason': 'No results', + }, }) - }, + span?.end() + return { message: 'No results' } + } + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: '✅ Found ' + text.length + ' results' }, + }) + span?.update({ + output: { resultCount: text.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.resultCount': text.length, + }, + }) + span?.end() + return { message: text.join('\n') } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`Google search failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + 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({ - id: 'extractTablesTool', - description: - 'Extract all tables from a webpage and return as structured JSON data.', - inputSchema: z.object({ - url: z.string().describe('URL containing tables'), - tableIndex: z - .number() - .optional() - .describe( - 'Specific table index to extract (0-based), extracts all if not specified' - ), - }), - outputSchema: z.object({ - success: z.boolean(), - tables: z - .array( - z.object({ - headers: z.array(z.string()), - rows: z.array(z.array(z.string())), - }) - ) - .optional(), - 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: { - 'tool.id': 'browser-extract-tables', - 'tool.input.url': inputData.url, - }, + id: 'extractTablesTool', + description: + 'Extract all tables from a webpage and return as structured JSON data.', + inputSchema: z.object({ + url: z.string().describe('URL containing tables'), + tableIndex: z + .number() + .optional() + .describe( + 'Specific table index to extract (0-based), extracts all if not specified' + ), + }), + outputSchema: z.object({ + success: z.boolean(), + tables: z + .array( + z.object({ + headers: z.array(z.string()), + rows: z.array(z.array(z.string())), }) - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: `📊 Extracting tables from ${inputData.url}` }, + ) + .optional(), + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'browser-extract-tables', + input: inputData, + metadata: { + 'tool.id': 'browser-extract-tables', + 'tool.input.url': inputData.url, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: `📊 Extracting tables from ${inputData.url}` }, + }) + + try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing table extraction for user', { + userId: requestCtx.userId, }) - - 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' }) - - const tables = await page.evaluate((tableIndex) => { - const allTables = document.querySelectorAll('table') - const result: Array<{ headers: string[]; rows: string[][] }> = - [] - - const tablesToProcess = - tableIndex !== undefined - ? [allTables[tableIndex]].filter(Boolean) - : Array.from(allTables) - - for (const table of tablesToProcess) { - const headers: string[] = [] - const rows: string[][] = [] - - const headerCells = table.querySelectorAll('th') - headerCells.forEach((cell) => - headers.push(cell.textContent?.trim() ?? '') - ) - - const bodyRows = table.querySelectorAll('tbody tr') - bodyRows.forEach((row) => { - const rowData: string[] = [] - row.querySelectorAll('td').forEach((cell) => - rowData.push(cell.textContent?.trim() ?? '') - ) - if (rowData.length > 0) { - rows.push(rowData) - } - }) - - result.push({ headers, rows }) - } - - return result - }, inputData.tableIndex) - - await page.close() - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { message: `✅ Extracted ${tables.length} table(s)` }, - }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.tableCount': tables.length, - }) - span.end() - - return { success: true, tables } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`Table extraction failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) + } + 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' }) + + const tables = await page.evaluate((tableIndex) => { + const allTables = document.querySelectorAll('table') + const result: Array<{ headers: string[]; rows: string[][] }> = + [] + + const tablesToProcess = + tableIndex !== undefined + ? [allTables[tableIndex]].filter(Boolean) + : Array.from(allTables) + + for (const table of tablesToProcess) { + const headers: string[] = [] + const rows: string[][] = [] + + const headerCells = table.querySelectorAll('th') + headerCells.forEach((cell) => + headers.push(cell.textContent?.trim() ?? '') + ) + + const bodyRows = table.querySelectorAll('tbody tr') + bodyRows.forEach((row) => { + const rowData: string[] = [] + row.querySelectorAll('td').forEach((cell) => + rowData.push(cell.textContent?.trim() ?? '') + ) + if (rowData.length > 0) { + rows.push(rowData) } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - return { success: false, message: errorMsg } + }) + + result.push({ headers, rows }) } - }, - 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', - }) - }, + + return result + }, inputData.tableIndex) + + await page.close() + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: `✅ Extracted ${tables.length} table(s)` }, + }) + span?.update({ + output: { tableCount: tables.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.tableCount': tables.length, + }, + }) + span?.end() + + return { success: true, tables } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`Table extraction failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + 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({ - id: 'monitorPageTool', - description: - 'Monitor a webpage for changes by comparing content at intervals.', - inputSchema: z.object({ - url: z.string().describe('URL to monitor'), - selector: z - .string() - .optional() - .describe('CSS selector to monitor (defaults to body)'), - checkInterval: z - .number() - .optional() - .default(5000) - .describe('Interval in milliseconds between checks'), - maxChecks: z - .number() - .optional() - .default(10) - .describe('Maximum number of checks before stopping'), - }), - outputSchema: z.object({ - success: z.boolean(), - changed: z.boolean(), - previousContent: z.string().optional(), - currentContent: z.string().optional(), - checkCount: z.number(), - 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: { - 'tool.id': 'browser-monitor-page', - 'tool.input.url': inputData.url, - 'tool.input.selector': inputData.selector, - }, + id: 'monitorPageTool', + description: + 'Monitor a webpage for changes by comparing content at intervals.', + inputSchema: z.object({ + url: z.string().describe('URL to monitor'), + selector: z + .string() + .optional() + .describe('CSS selector to monitor (defaults to body)'), + checkInterval: z + .number() + .optional() + .default(5000) + .describe('Interval in milliseconds between checks'), + maxChecks: z + .number() + .optional() + .default(10) + .describe('Maximum number of checks before stopping'), + }), + outputSchema: z.object({ + success: z.boolean(), + changed: z.boolean(), + previousContent: z.string().optional(), + currentContent: z.string().optional(), + checkCount: z.number(), + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as BrowserRequestContext | undefined + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'browser-monitor-page', + input: inputData, + metadata: { + 'tool.id': 'browser-monitor-page', + 'tool.input.url': inputData.url, + 'tool.input.selector': inputData.selector, + 'user.id': requestCtx?.userId, + }, + }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { message: `👁️ Monitoring ${inputData.url}` }, + }) + + try { + if (typeof requestCtx?.userId === 'string') { + log.debug('Executing page monitor for user', { + userId: requestCtx.userId, }) - - await context?.writer?.custom({ + } + 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' + + await page.goto(inputData.url, { waitUntil: 'domcontentloaded' }) + const previousContent = await page.$eval( + selector, + (el) => el.textContent ?? '' + ) + + let checkCount = 0 + let changed = false + let currentContent = previousContent + + while (checkCount < (inputData.maxChecks ?? 10) && !changed) { + await new Promise((resolve) => + setTimeout(resolve, inputData.checkInterval ?? 5000) + ) + await page.reload({ waitUntil: 'domcontentloaded' }) + + currentContent = await page.$eval( + selector, + (el) => el.textContent ?? '' + ) + checkCount++ + + if (currentContent !== previousContent) { + changed = true + await context?.writer?.custom({ type: 'data-tool-progress', - data: { message: `👁️ Monitoring ${inputData.url}` }, - }) - - 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' - - await page.goto(inputData.url, { waitUntil: 'domcontentloaded' }) - const previousContent = await page.$eval( - selector, - (el) => el.textContent ?? '' - ) - - let checkCount = 0 - let changed = false - let currentContent = previousContent - - while (checkCount < (inputData.maxChecks ?? 10) && !changed) { - await new Promise((resolve) => - setTimeout(resolve, inputData.checkInterval ?? 5000) - ) - await page.reload({ waitUntil: 'domcontentloaded' }) - - currentContent = await page.$eval( - selector, - (el) => el.textContent ?? '' - ) - checkCount++ - - if (currentContent !== previousContent) { - changed = true - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - message: `🔔 Change detected after ${checkCount} checks`, - }, - }) - } - } - - await page.close() - - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.changed': changed, - 'tool.output.checkCount': checkCount, - }) - span.end() - return { - success: true, - changed, - previousContent, - currentContent, - checkCount, - } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : 'Unknown error' - log.error(`Page monitoring failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - return { - success: false, - changed: false, - checkCount: 0, - message: errorMsg, - } + data: { + message: `🔔 Change detected after ${checkCount} checks`, + }, + }) } - }, - 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', - }) - }, + } + + await page.close() + + span?.update({ + output: { success: true, changed, checkCount }, + metadata: { + 'tool.output.success': true, + 'tool.output.changed': changed, + 'tool.output.checkCount': checkCount, + }, + }) + span?.end() + return { + success: true, + changed, + previousContent, + currentContent, + checkCount, + } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : 'Unknown error' + log.error(`Page monitoring failed: ${errorMsg}`) + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + return { + success: false, + changed: false, + checkCount: 0, + message: errorMsg, + } + } + }, + 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<typeof browserTool> diff --git a/src/mastra/tools/calculator.tool.ts b/src/mastra/tools/calculator.tool.ts index 6c5939e9..8231e3d9 100644 --- a/src/mastra/tools/calculator.tool.ts +++ b/src/mastra/tools/calculator.tool.ts @@ -1,971 +1,1020 @@ +import { SpanType } from '@mastra/core/observability' +import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' import { createTool } 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' +export type UserTier = 'free' | 'pro' | 'enterprise' // RequestContext interface for Calculator tool export interface CalculatorToolContext extends RequestContext { - userId?: string - precision?: number - allowComplexExpressions?: boolean - maxExpressionLength?: number + 'user-tier': UserTier + userId?: string + precision?: number + allowComplexExpressions?: boolean + maxExpressionLength?: number } // Enhanced mathematical functions and constants const MATH_CONSTANTS = { - pi: Math.PI, - e: Math.E, - ln2: Math.LN2, - ln10: Math.LN10, - log2e: Math.LOG2E, - log10e: Math.LOG10E, - sqrt2: Math.SQRT2, - sqrt1_2: Math.SQRT1_2, - phi: (1 + Math.sqrt(5)) / 2, // Golden ratio - tau: 2 * Math.PI, // Tau (2π) + pi: Math.PI, + e: Math.E, + ln2: Math.LN2, + ln10: Math.LN10, + log2e: Math.LOG2E, + log10e: Math.LOG10E, + sqrt2: Math.SQRT2, + sqrt1_2: Math.SQRT1_2, + phi: (1 + Math.sqrt(5)) / 2, // Golden ratio + tau: 2 * Math.PI, // Tau (2π) } const MATH_FUNCTIONS = { - // Basic arithmetic - abs: Math.abs, - ceil: Math.ceil, - floor: Math.floor, - round: Math.round, - trunc: Math.trunc, - sign: Math.sign, - - // Trigonometric functions - sin: Math.sin, - cos: Math.cos, - tan: Math.tan, - asin: Math.asin, - acos: Math.acos, - atan: Math.atan, - atan2: Math.atan2, - - // Hyperbolic functions - sinh: Math.sinh, - cosh: Math.cosh, - tanh: Math.tanh, - asinh: Math.asinh, - acosh: Math.acosh, - atanh: Math.atanh, - - // Exponential and logarithmic - exp: Math.exp, - log: Math.log, - log10: Math.log10, - log2: Math.log2, - pow: Math.pow, - sqrt: Math.sqrt, - cbrt: Math.cbrt, - - // Min/Max - max: Math.max, - min: Math.min, - - // Random (with seed support) - random: Math.random, - - // Special functions - factorial, - fibonacci, - gcd, - lcm, - isPrime, - combinations, - permutations, + // Basic arithmetic + abs: Math.abs, + ceil: Math.ceil, + floor: Math.floor, + round: Math.round, + trunc: Math.trunc, + sign: Math.sign, + + // Trigonometric functions + sin: Math.sin, + cos: Math.cos, + tan: Math.tan, + asin: Math.asin, + acos: Math.acos, + atan: Math.atan, + atan2: Math.atan2, + + // Hyperbolic functions + sinh: Math.sinh, + cosh: Math.cosh, + tanh: Math.tanh, + asinh: Math.asinh, + acosh: Math.acosh, + atanh: Math.atanh, + + // Exponential and logarithmic + exp: Math.exp, + log: Math.log, + log10: Math.log10, + log2: Math.log2, + pow: Math.pow, + sqrt: Math.sqrt, + cbrt: Math.cbrt, + + // Min/Max + max: Math.max, + min: Math.min, + + // Random (with seed support) + random: Math.random, + + // Special functions + factorial, + fibonacci, + gcd, + lcm, + isPrime, + combinations, + permutations, } // Unit conversion factors (to base SI units) const UNIT_CONVERSIONS = { - // Length - mm: 0.001, - cm: 0.01, - dm: 0.1, - m: 1, - km: 1000, - in: 0.0254, - ft: 0.3048, - yd: 0.9144, - mi: 1609.344, - - // Mass - mg: 0.000001, - g: 0.001, - kg: 1, - t: 1000, - oz: 0.0283495, - lb: 0.453592, - st: 6.35029, - - // Temperature (special handling) - celsius: 'celsius', - fahrenheit: 'fahrenheit', - kelvin: 'kelvin', - - // Time - ms: 0.001, - s: 1, - min: 60, - h: 3600, - day: 86400, - - // Digital storage - b: 1, - B: 8, - KB: 8000, - MB: 8000000, - GB: 8000000000, - TB: 8000000000000, + // Length + mm: 0.001, + cm: 0.01, + dm: 0.1, + m: 1, + km: 1000, + in: 0.0254, + ft: 0.3048, + yd: 0.9144, + mi: 1609.344, + + // Mass + mg: 0.000001, + g: 0.001, + kg: 1, + t: 1000, + oz: 0.0283495, + lb: 0.453592, + st: 6.35029, + + // Temperature (special handling) + celsius: 'celsius', + fahrenheit: 'fahrenheit', + kelvin: 'kelvin', + + // Time + ms: 0.001, + s: 1, + min: 60, + h: 3600, + day: 86400, + + // Digital storage + b: 1, + B: 8, + KB: 8000, + MB: 8000000, + GB: 8000000000, + TB: 8000000000000, } // Create a safe evaluation context function createSafeContext() { - return { - ...MATH_CONSTANTS, - ...MATH_FUNCTIONS, - } + return { + ...MATH_CONSTANTS, + ...MATH_FUNCTIONS, + } } // Enhanced expression evaluator with variable support function evaluateExpression( - expression: string, - variables: Record<string, number> = {} + expression: string, + variables: Record<string, number> = {} ): number { - // Remove whitespace - const cleanExpr = expression.replace(/\s+/g, '') - - // Basic security checks - if (/[^0-9+\-*/().\s,a-zA-Z_]/.test(cleanExpr)) { - throw new Error('Invalid characters in expression') - } - - if ( - cleanExpr.includes('__proto__') || - cleanExpr.includes('prototype') || - cleanExpr.includes('constructor') - ) { - throw new Error('Potentially unsafe expression') - } - - // Create function with safe context and variables - const context = { ...createSafeContext(), ...variables } - const func = new Function(...Object.keys(context), `return ${cleanExpr}`) - - try { - const result = func(...Object.values(context)) - if (typeof result !== 'number' || !isFinite(result)) { - throw new Error('Expression did not evaluate to a valid number') - } - return result - } catch (error) { - throw new Error( - `Evaluation failed: ${error instanceof Error ? error.message : 'Unknown error'}` - ) + // Remove whitespace + const cleanExpr = expression.replace(/\s+/g, '') + + // Basic security checks + if (/[^0-9+\-*/().\s,a-zA-Z_]/.test(cleanExpr)) { + throw new Error('Invalid characters in expression') + } + + if ( + cleanExpr.includes('__proto__') || + cleanExpr.includes('prototype') || + cleanExpr.includes('constructor') + ) { + throw new Error('Potentially unsafe expression') + } + + // Create function with safe context and variables + const context = { ...createSafeContext(), ...variables } + const func = new Function(...Object.keys(context), `return ${cleanExpr}`) + + try { + const result = func(...Object.values(context)) + if (typeof result !== 'number' || !isFinite(result)) { + throw new Error('Expression did not evaluate to a valid number') } + return result + } catch (error) { + throw new Error( + `Evaluation failed: ${error instanceof Error ? error.message : 'Unknown error'}` + ) + } } // Unit conversion function function convertUnits(value: number, fromUnit: string, toUnit: string): number { - // Special handling for temperature - if ( - ['celsius', 'fahrenheit', 'kelvin'].includes(fromUnit) && - ['celsius', 'fahrenheit', 'kelvin'].includes(toUnit) - ) { - return convertTemperature(value, fromUnit as 'celsius' | 'fahrenheit' | 'kelvin', toUnit as 'celsius' | 'fahrenheit' | 'kelvin') - } - - // Regular unit conversion - const fromFactor = - UNIT_CONVERSIONS[fromUnit as keyof typeof UNIT_CONVERSIONS] - const toFactor = UNIT_CONVERSIONS[toUnit as keyof typeof UNIT_CONVERSIONS] - - if (typeof fromFactor !== 'number' || typeof toFactor !== 'number') { - throw new Error(`Unsupported unit conversion: ${fromUnit} to ${toUnit}`) - } - - // Convert to base unit, then to target unit - const baseValue = value * fromFactor - return baseValue / toFactor + // Special handling for temperature + if ( + ['celsius', 'fahrenheit', 'kelvin'].includes(fromUnit) && + ['celsius', 'fahrenheit', 'kelvin'].includes(toUnit) + ) { + return convertTemperature(value, fromUnit as 'celsius' | 'fahrenheit' | 'kelvin', toUnit as 'celsius' | 'fahrenheit' | 'kelvin') + } + + // Regular unit conversion + const fromFactor = + UNIT_CONVERSIONS[fromUnit as keyof typeof UNIT_CONVERSIONS] + const toFactor = UNIT_CONVERSIONS[toUnit as keyof typeof UNIT_CONVERSIONS] + + if (typeof fromFactor !== 'number' || typeof toFactor !== 'number') { + throw new Error(`Unsupported unit conversion: ${fromUnit} to ${toUnit}`) + } + + // Convert to base unit, then to target unit + const baseValue = value * fromFactor + return baseValue / toFactor } function convertTemperature( - value: number, - from: 'celsius' | 'fahrenheit' | 'kelvin', - to: 'celsius' | 'fahrenheit' | 'kelvin' + value: number, + from: 'celsius' | 'fahrenheit' | 'kelvin', + to: 'celsius' | 'fahrenheit' | 'kelvin' ): number { - let celsius: number - - // Convert to Celsius first - switch (from) { - case 'celsius': - celsius = value - break - case 'fahrenheit': - celsius = ((value - 32) * 5) / 9 - break - case 'kelvin': - celsius = value - 273.15 - break - default: - throw new Error(`Unsupported temperature unit: ${from}`) - } - - // Convert from Celsius to target - switch (to) { - case 'celsius': - return celsius - case 'fahrenheit': - return (celsius * 9) / 5 + 32 - case 'kelvin': - return celsius + 273.15 - default: - throw new Error(`Unsupported temperature unit: ${to}`) - } + let celsius: number + + // Convert to Celsius first + switch (from) { + case 'celsius': + celsius = value + break + case 'fahrenheit': + celsius = ((value - 32) * 5) / 9 + break + case 'kelvin': + celsius = value - 273.15 + break + default: + throw new Error(`Unsupported temperature unit: ${from}`) + } + + // Convert from Celsius to target + switch (to) { + case 'celsius': + return celsius + case 'fahrenheit': + return (celsius * 9) / 5 + 32 + case 'kelvin': + return celsius + 273.15 + default: + throw new Error(`Unsupported temperature unit: ${to}`) + } } // Matrix operations const MATRIX_OPERATIONS = { - add: (a: number[][], b: number[][]) => { - if (a.length !== b.length || a[0].length !== b[0].length) { - throw new Error('Matrix dimensions must match for addition') - } - return a.map((row, i) => row.map((val, j) => val + b[i][j])) - }, + add: (a: number[][], b: number[][]) => { + if (a.length !== b.length || a[0].length !== b[0].length) { + throw new Error('Matrix dimensions must match for addition') + } + return a.map((row, i) => row.map((val, j) => val + b[i][j])) + }, - subtract: (a: number[][], b: number[][]) => { - if (a.length !== b.length || a[0].length !== b[0].length) { - throw new Error('Matrix dimensions must match for subtraction') - } - return a.map((row, i) => row.map((val, j) => val - b[i][j])) - }, + subtract: (a: number[][], b: number[][]) => { + if (a.length !== b.length || a[0].length !== b[0].length) { + throw new Error('Matrix dimensions must match for subtraction') + } + return a.map((row, i) => row.map((val, j) => val - b[i][j])) + }, - multiply: (a: number[][], b: number[][]) => { - if (a[0].length !== b.length) { - throw new Error('Matrix dimensions incompatible for multiplication') - } - const result = Array(a.length) - .fill(0) - .map(() => Array(b[0].length).fill(0)) - for (let i = 0; i < a.length; i++) { - for (let j = 0; j < b[0].length; j++) { - for (let k = 0; k < b.length; k++) { - result[i][j] += a[i][k] * b[k][j] - } - } + multiply: (a: number[][], b: number[][]) => { + if (a[0].length !== b.length) { + throw new Error('Matrix dimensions incompatible for multiplication') + } + const result = Array(a.length) + .fill(0) + .map(() => Array(b[0].length).fill(0)) + for (let i = 0; i < a.length; i++) { + for (let j = 0; j < b[0].length; j++) { + for (let k = 0; k < b.length; k++) { + result[i][j] += a[i][k] * b[k][j] } - return result - }, + } + } + return result + }, - transpose: (matrix: number[][]) => { - return matrix[0].map((_, colIndex) => - matrix.map((row) => row[colIndex]) - ) - }, + transpose: (matrix: number[][]) => { + return matrix[0].map((_, colIndex) => + matrix.map((row) => row[colIndex]) + ) + }, - determinant: (matrix: number[][]) => { - if (matrix.length !== matrix[0].length) { - throw new Error('Matrix must be square for determinant calculation') - } - return calculateDeterminant(matrix) - }, + determinant: (matrix: number[][]) => { + if (matrix.length !== matrix[0].length) { + throw new Error('Matrix must be square for determinant calculation') + } + return calculateDeterminant(matrix) + }, - inverse: (matrix: number[][]) => { - const det = calculateDeterminant(matrix) - if (det === 0) { - throw new Error('Matrix is singular and cannot be inverted') - } - // Simplified inverse for 2x2 matrices - if (matrix.length === 2 && matrix[0].length === 2) { - const [[a, b], [c, d]] = matrix - return [ - [d / det, -b / det], - [-c / det, a / det], - ] - } - throw new Error('Inverse calculation only supported for 2x2 matrices') - }, + inverse: (matrix: number[][]) => { + const det = calculateDeterminant(matrix) + if (det === 0) { + throw new Error('Matrix is singular and cannot be inverted') + } + // Simplified inverse for 2x2 matrices + if (matrix.length === 2 && matrix[0].length === 2) { + const [[a, b], [c, d]] = matrix + return [ + [d / det, -b / det], + [-c / det, a / det], + ] + } + throw new Error('Inverse calculation only supported for 2x2 matrices') + }, } export const calculatorTool = createTool({ - id: 'calculator', - description: - 'Advanced mathematical calculator with expressions, unit conversions, and matrix operations', - inputSchema: z.object({ - expression: z - .string() - .describe( - 'Mathematical expression to evaluate (supports variables, functions, and complex expressions)' - ), - variables: z - .record(z.string(), z.number()) - .optional() - .describe('Variable definitions for the expression'), - precision: z - .number() - .optional() - .default(6) - .describe('Number of decimal places for the result'), - }), - outputSchema: z.object({ - success: z.boolean(), - result: z.number(), - formattedResult: z.string(), - expression: z.string(), - variables: z.record(z.string(), z.number()).optional(), - message: z.string().optional(), - }), + id: 'calculator', + description: + 'Advanced mathematical calculator with expressions, unit conversions, and matrix operations', + inputSchema: z.object({ + expression: z + .string() + .describe( + 'Mathematical expression to evaluate (supports variables, functions, and complex expressions)' + ), + variables: z + .record(z.string(), z.number()) + .optional() + .describe('Variable definitions for the expression'), + precision: z + .number() + .optional() + .default(6) + .describe('Number of decimal places for the result'), + }), + outputSchema: z.object({ + success: z.boolean(), + result: z.number(), + formattedResult: z.string(), + expression: z.string(), + variables: z.record(z.string(), z.number()).optional(), + message: z.string().optional(), + }), + + execute: async (inputData, context) => { + const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + + const requestCtx = context?.requestContext as CalculatorToolContext | undefined + const precision = requestCtx?.precision ?? 6 + const allowComplexExpressions = requestCtx?.allowComplexExpressions ?? true + const maxExpressionLength = requestCtx?.maxExpressionLength ?? 1000 + + // Create child span for calculator operation + const calculatorSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'calculator-operation', + input: inputData, + metadata: { + 'tool.id': 'calculator', + 'tool.input.expression': inputData.expression, + 'tool.input.precision': precision, + 'user.id': requestCtx?.userId, + }, + }) - execute: async (inputData, context) => { - const writer = context?.writer - - const requestCtx = context?.requestContext as CalculatorToolContext | undefined - const precision = requestCtx?.precision ?? 6 - const allowComplexExpressions = requestCtx?.allowComplexExpressions ?? true - const maxExpressionLength = requestCtx?.maxExpressionLength ?? 1000 - - const tracer = trace.getTracer('calculator-tool', '1.0.0') - const span = tracer.startSpan('calculate') - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🔢 Evaluating expression: ${inputData.expression}`, - stage: 'calculator', - }, - id: 'calculator', - }) + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🔢 Evaluating expression: ${inputData.expression}`, + stage: 'calculator', + }, + id: 'calculator', + }) - try { - // Validate expression length - if (inputData.expression.length > maxExpressionLength) { - throw new Error( - `Expression too long (max ${maxExpressionLength} characters)` - ) - } - - const vars = (inputData.variables ?? {}) - const result = evaluateExpression( - inputData.expression, - vars - ) - const actualPrecision = inputData.precision || precision - const formattedResult = result.toFixed(actualPrecision) - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Result: ${formattedResult}`, - stage: 'calculator', - }, - id: 'calculator', - }) - - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.result': result, - 'tool.input.expressionLength': inputData.expression.length, - }) - span.end() - - return { - success: true, - result, - formattedResult, - expression: inputData.expression, - variables: inputData.variables, - } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : String(e) - log.error(`Calculator evaluation failed: ${errorMsg}`, { - error: errorMsg, - }) - - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - - return { - success: false, - result: 0, - formattedResult: '0', - expression: inputData.expression, - variables: inputData.variables, - message: errorMsg, - } - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Calculator tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Calculator tool received input chunk', { - toolCallId, - inputTextDelta, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Calculator tool received input', { - toolCallId, - messageCount: messages.length, - inputData: { - expression: input.expression, - variablesCount: Object.keys(input.variables ?? {}).length, - }, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Calculator tool completed', { - toolCallId, - toolName, - outputData: { - success: output.success, - hasResult: output.result !== undefined, - }, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + try { + // Validate expression length + if (inputData.expression.length > maxExpressionLength) { + throw new Error( + `Expression too long (max ${maxExpressionLength} characters)` + ) + } + + const vars = (inputData.variables ?? {}) + const result = evaluateExpression( + inputData.expression, + vars + ) + const actualPrecision = inputData.precision || precision + const formattedResult = result.toFixed(actualPrecision) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Result: ${formattedResult}`, + stage: 'calculator', + }, + id: 'calculator', + }) + + // Update span with successful result + calculatorSpan?.update({ + output: { success: true, result, formattedResult }, + metadata: { + 'tool.output.success': true, + 'tool.output.result': result, + 'tool.input.expressionLength': inputData.expression.length, + }, + }) + calculatorSpan?.end() + + return { + success: true, + result, + formattedResult, + expression: inputData.expression, + variables: inputData.variables, + } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e) + log.error(`Calculator evaluation failed: ${errorMsg}`, { + error: errorMsg, + }) + + // Record error in span + calculatorSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + + return { + success: false, + result: 0, + formattedResult: '0', + expression: inputData.expression, + variables: inputData.variables, + message: errorMsg, + } + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Calculator tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Calculator tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Calculator tool received input', { + toolCallId, + messageCount: messages.length, + inputData: { + expression: input.expression, + variablesCount: Object.keys(input.variables ?? {}).length, + }, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Calculator tool completed', { + toolCallId, + toolName, + outputData: { + success: output.success, + hasResult: output.result !== undefined, + }, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const unitConverterTool = createTool({ - id: 'unit-converter', - description: - 'Convert between different units of measurement (length, mass, temperature, time, digital storage)', - inputSchema: z.object({ - value: z.number().describe('Value to convert'), - fromUnit: z - .string() - .describe('Source unit (e.g., "m", "kg", "celsius", "MB")'), - toUnit: z - .string() - .describe('Target unit (e.g., "ft", "lb", "fahrenheit", "GB")'), - precision: z - .number() - .optional() - .default(4) - .describe('Decimal precision for result'), - }), - outputSchema: z.object({ - success: z.boolean(), - result: z.number(), - formattedResult: z.string(), - value: z.number(), - fromUnit: z.string(), - toUnit: z.string(), - conversionFactor: z.number(), - message: z.string().optional(), - }), - 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') - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🔄 Converting ${inputData.value} ${inputData.fromUnit} to ${inputData.toUnit}`, - stage: 'unit-converter', - }, - id: 'unit-converter', - }) + id: 'unit-converter', + description: + 'Convert between different units of measurement (length, mass, temperature, time, digital storage)', + inputSchema: z.object({ + value: z.number().describe('Value to convert'), + fromUnit: z + .string() + .describe('Source unit (e.g., "m", "kg", "celsius", "MB")'), + toUnit: z + .string() + .describe('Target unit (e.g., "ft", "lb", "fahrenheit", "GB")'), + precision: z + .number() + .optional() + .default(4) + .describe('Decimal precision for result'), + }), + outputSchema: z.object({ + success: z.boolean(), + result: z.number(), + formattedResult: z.string(), + value: z.number(), + fromUnit: z.string(), + toUnit: z.string(), + conversionFactor: z.number(), + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as CalculatorToolContext | undefined + + // Create child span for unit conversion + const unitConverterSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'unit-conversion', + input: inputData, + metadata: { + 'tool.id': 'unit-converter', + 'tool.input.value': inputData.value, + 'tool.input.fromUnit': inputData.fromUnit, + 'tool.input.toUnit': inputData.toUnit, + 'user.id': requestCtx?.userId, + }, + }) - 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, - inputData.toUnit - ) - const formattedResult = result.toFixed(inputData.precision) - - // Calculate conversion factor - const fromFactor = - UNIT_CONVERSIONS[ - inputData.fromUnit as keyof typeof UNIT_CONVERSIONS - ] - const toFactor = - UNIT_CONVERSIONS[ - inputData.toUnit as keyof typeof UNIT_CONVERSIONS - ] - const conversionFactor = - typeof fromFactor === 'number' && typeof toFactor === 'number' - ? toFactor / fromFactor - : 1 - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Converted: ${formattedResult} ${inputData.toUnit}`, - stage: 'unit-converter', - }, - id: 'unit-converter', - }) - - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.result': result, - 'tool.input.value': inputData.value, - }) - span.end() - - return { - success: true, - result, - formattedResult, - value: inputData.value, - fromUnit: inputData.fromUnit, - toUnit: inputData.toUnit, - conversionFactor, - } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : String(e) - log.error(`Unit conversion failed: ${errorMsg}`, { - error: errorMsg, - }) - - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - - return { - success: false, - result: 0, - formattedResult: '0', - value: inputData.value, - fromUnit: inputData.fromUnit, - toUnit: inputData.toUnit, - conversionFactor: 1, - message: errorMsg, - } - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Unit Converter tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Unit Converter tool received input chunk', { - toolCallId, - inputTextDelta, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Unit Converter tool received input', { - toolCallId, - messageCount: messages.length, - inputData: { - value: input.value, - fromUnit: input.fromUnit, - toUnit: input.toUnit, - }, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Unit Converter tool completed', { - toolCallId, - toolName, - outputData: { - success: output.success, - result: output.formattedResult, - }, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🔄 Converting ${inputData.value} ${inputData.fromUnit} to ${inputData.toUnit}`, + stage: 'unit-converter', + }, + id: 'unit-converter', + }) + + 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, + inputData.toUnit + ) + const formattedResult = result.toFixed(inputData.precision) + + // Calculate conversion factor + const fromFactor = + UNIT_CONVERSIONS[ + inputData.fromUnit as keyof typeof UNIT_CONVERSIONS + ] + const toFactor = + UNIT_CONVERSIONS[ + inputData.toUnit as keyof typeof UNIT_CONVERSIONS + ] + const conversionFactor = + typeof fromFactor === 'number' && typeof toFactor === 'number' + ? toFactor / fromFactor + : 1 + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Converted: ${formattedResult} ${inputData.toUnit}`, + stage: 'unit-converter', + }, + id: 'unit-converter', + }) + + // Update span with successful result + unitConverterSpan?.update({ + output: { success: true, result, formattedResult }, + metadata: { + 'tool.output.success': true, + 'tool.output.result': result, + 'tool.input.value': inputData.value, + }, + }) + unitConverterSpan?.end() + + return { + success: true, + result, + formattedResult, + value: inputData.value, + fromUnit: inputData.fromUnit, + toUnit: inputData.toUnit, + conversionFactor, + } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e) + log.error(`Unit conversion failed: ${errorMsg}`, { + error: errorMsg, + }) + + // Record error in span + unitConverterSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + + return { + success: false, + result: 0, + formattedResult: '0', + value: inputData.value, + fromUnit: inputData.fromUnit, + toUnit: inputData.toUnit, + conversionFactor: 1, + message: errorMsg, + } + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Unit Converter tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Unit Converter tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Unit Converter tool received input', { + toolCallId, + messageCount: messages.length, + inputData: { + value: input.value, + fromUnit: input.fromUnit, + toUnit: input.toUnit, + }, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Unit Converter tool completed', { + toolCallId, + toolName, + outputData: { + success: output.success, + result: output.formattedResult, + }, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const matrixCalculatorTool = createTool({ - id: 'matrix-calculator', - description: - 'Perform matrix operations (add, subtract, multiply, transpose, determinant, inverse)', - inputSchema: z.object({ - operation: z - .enum([ - 'add', - 'subtract', - 'multiply', - 'transpose', - 'determinant', - 'inverse', - ]) - .describe('Matrix operation to perform'), - matrixA: z - .array(z.array(z.number())) - .describe('First matrix (or only matrix for unary operations)'), - matrixB: z - .array(z.array(z.number())) - .optional() - .describe('Second matrix (required for binary operations)'), - }), - outputSchema: z.object({ - success: z.boolean(), - result: z.union([z.array(z.array(z.number())), z.number()]), - operation: z.string(), - matrixA: z.array(z.array(z.number())), - matrixB: z.array(z.array(z.number())).optional(), - dimensions: z.object({ - a: z.array(z.number()), - b: z.array(z.number()).optional(), - result: z.array(z.number()), - }), - message: z.string().optional(), + id: 'matrix-calculator', + description: + 'Perform matrix operations (add, subtract, multiply, transpose, determinant, inverse)', + inputSchema: z.object({ + operation: z + .enum([ + 'add', + 'subtract', + 'multiply', + 'transpose', + 'determinant', + 'inverse', + ]) + .describe('Matrix operation to perform'), + matrixA: z + .array(z.array(z.number())) + .describe('First matrix (or only matrix for unary operations)'), + matrixB: z + .array(z.array(z.number())) + .optional() + .describe('Second matrix (required for binary operations)'), + }), + outputSchema: z.object({ + success: z.boolean(), + result: z.union([z.array(z.array(z.number())), z.number()]), + operation: z.string(), + matrixA: z.array(z.array(z.number())), + matrixB: z.array(z.array(z.number())).optional(), + dimensions: z.object({ + a: z.array(z.number()), + b: z.array(z.number()).optional(), + result: z.array(z.number()), }), - 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') - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `📐 Performing matrix ${inputData.operation}...`, - stage: 'matrix-calculator', - }, - id: 'matrix-calculator', - }) + message: z.string().optional(), + }), + execute: async (inputData, context) => { + const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as CalculatorToolContext | undefined + + // Create child span for matrix calculation + const matrixCalculatorSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'matrix-calculation', + input: inputData, + metadata: { + 'tool.id': 'matrix-calculator', + 'tool.input.operation': inputData.operation, + 'tool.input.matrixASize': inputData.matrixA.length, + 'user.id': requestCtx?.userId, + }, + }) - 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) { - case 'add': - if (!inputData.matrixB) { - throw new Error('Matrix B required for addition') - } - result = MATRIX_OPERATIONS.add( - inputData.matrixA, - inputData.matrixB - ) - break - case 'subtract': - if (!inputData.matrixB) { - throw new Error('Matrix B required for subtraction') - } - result = MATRIX_OPERATIONS.subtract( - inputData.matrixA, - inputData.matrixB - ) - break - case 'multiply': - if (!inputData.matrixB) { - throw new Error('Matrix B required for multiplication') - } - result = MATRIX_OPERATIONS.multiply( - inputData.matrixA, - inputData.matrixB - ) - break - case 'transpose': - result = MATRIX_OPERATIONS.transpose(inputData.matrixA) - break - case 'determinant': - result = MATRIX_OPERATIONS.determinant(inputData.matrixA) - break - case 'inverse': - result = MATRIX_OPERATIONS.inverse(inputData.matrixA) - break - default: - throw new Error( - `Unknown matrix operation: ${inputData.operation}` - ) - } - - const dimensions = { - a: [ - inputData.matrixA.length, - inputData.matrixA[0]?.length || 0, - ], - b: inputData.matrixB - ? [ - inputData.matrixB.length, - inputData.matrixB[0]?.length || 0, - ] - : undefined, - result: - Array.isArray(result) && Array.isArray(result[0]) - ? [result.length, result[0].length] - : [1, 1], - } - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Matrix ${inputData.operation} completed`, - stage: 'matrix-calculator', - }, - id: 'matrix-calculator', - }) - - span.setAttributes({ - 'tool.output.success': true, - 'tool.input.operation': inputData.operation, - }) - span.end() - - return { - success: true, - result, - operation: inputData.operation, - matrixA: inputData.matrixA, - matrixB: inputData.matrixB, - dimensions, - } - } catch (e) { - const errorMsg = e instanceof Error ? e.message : String(e) - log.error(`Matrix operation failed: ${errorMsg}`, { - error: errorMsg, - }) - - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() - - return { - success: false, - result: [], - operation: inputData.operation, - matrixA: inputData.matrixA, - matrixB: inputData.matrixB, - dimensions: { - a: [ - inputData.matrixA.length, - inputData.matrixA[0]?.length || 0, - ], - b: inputData.matrixB - ? [ - inputData.matrixB.length, - inputData.matrixB[0]?.length || 0, - ] - : undefined, - result: [0, 0], - }, - message: errorMsg, - } - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Matrix Calculator tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Matrix Calculator tool received input chunk', { - toolCallId, - inputTextDelta, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Matrix Calculator tool received input', { - toolCallId, - messageCount: messages.length, - inputData: { - operation: input.operation, - matrixADims: [input.matrixA.length, input.matrixA[0]?.length || 0], - }, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Matrix Calculator tool completed', { - toolCallId, - toolName, - outputData: { - success: output.success, - operation: output.operation, - }, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `📐 Performing matrix ${inputData.operation}...`, + stage: 'matrix-calculator', + }, + id: 'matrix-calculator', + }) + + 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) { + case 'add': + if (!inputData.matrixB) { + throw new Error('Matrix B required for addition') + } + result = MATRIX_OPERATIONS.add( + inputData.matrixA, + inputData.matrixB + ) + break + case 'subtract': + if (!inputData.matrixB) { + throw new Error('Matrix B required for subtraction') + } + result = MATRIX_OPERATIONS.subtract( + inputData.matrixA, + inputData.matrixB + ) + break + case 'multiply': + if (!inputData.matrixB) { + throw new Error('Matrix B required for multiplication') + } + result = MATRIX_OPERATIONS.multiply( + inputData.matrixA, + inputData.matrixB + ) + break + case 'transpose': + result = MATRIX_OPERATIONS.transpose(inputData.matrixA) + break + case 'determinant': + result = MATRIX_OPERATIONS.determinant(inputData.matrixA) + break + case 'inverse': + result = MATRIX_OPERATIONS.inverse(inputData.matrixA) + break + default: + throw new Error( + `Unknown matrix operation: ${inputData.operation}` + ) + } + + const dimensions = { + a: [ + inputData.matrixA.length, + inputData.matrixA[0]?.length || 0, + ], + b: inputData.matrixB + ? [ + inputData.matrixB.length, + inputData.matrixB[0]?.length || 0, + ] + : undefined, + result: + Array.isArray(result) && Array.isArray(result[0]) + ? [result.length, result[0].length] + : [1, 1], + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Matrix ${inputData.operation} completed`, + stage: 'matrix-calculator', + }, + id: 'matrix-calculator', + }) + + // Update span with successful result + matrixCalculatorSpan?.update({ + output: { success: true, operation: inputData.operation }, + metadata: { + 'tool.output.success': true, + 'tool.input.operation': inputData.operation, + }, + }) + matrixCalculatorSpan?.end() + + return { + success: true, + result, + operation: inputData.operation, + matrixA: inputData.matrixA, + matrixB: inputData.matrixB, + dimensions, + } + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e) + log.error(`Matrix operation failed: ${errorMsg}`, { + error: errorMsg, + }) + + // Record error in span + matrixCalculatorSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) + + return { + success: false, + result: [], + operation: inputData.operation, + matrixA: inputData.matrixA, + matrixB: inputData.matrixB, + dimensions: { + a: [ + inputData.matrixA.length, + inputData.matrixA[0]?.length || 0, + ], + b: inputData.matrixB + ? [ + inputData.matrixB.length, + inputData.matrixB[0]?.length || 0, + ] + : undefined, + result: [0, 0], + }, + message: errorMsg, + } + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Matrix Calculator tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Matrix Calculator tool received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Matrix Calculator tool received input', { + toolCallId, + messageCount: messages.length, + inputData: { + operation: input.operation, + matrixADims: [input.matrixA.length, input.matrixA[0]?.length || 0], + }, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Matrix Calculator tool completed', { + toolCallId, + toolName, + outputData: { + success: output.success, + operation: output.operation, + }, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) // Helper functions function factorial(n: number): number { - if (n < 0 || !Number.isInteger(n)) { - throw new Error('Factorial requires a non-negative integer') - } - if (n === 0 || n === 1) { - return 1 - } - let result = 1 - for (let i = 2; i <= n; i++) { - result *= i - } - return result + if (n < 0 || !Number.isInteger(n)) { + throw new Error('Factorial requires a non-negative integer') + } + if (n === 0 || n === 1) { + return 1 + } + let result = 1 + for (let i = 2; i <= n; i++) { + result *= i + } + return result } function fibonacci(n: number): number { - if (n < 0 || !Number.isInteger(n)) { - throw new Error('Fibonacci requires a non-negative integer') - } - if (n === 0) { - return 0 - } - if (n === 1) { - return 1 - } - let a = 0, - b = 1 - for (let i = 2; i <= n; i++) { - const temp = a + b - a = b - b = temp - } - return b + if (n < 0 || !Number.isInteger(n)) { + throw new Error('Fibonacci requires a non-negative integer') + } + if (n === 0) { + return 0 + } + if (n === 1) { + return 1 + } + let a = 0, + b = 1 + for (let i = 2; i <= n; i++) { + const temp = a + b + a = b + b = temp + } + return b } function gcd(...numbers: number[]): number { - if (numbers.length < 2) { - throw new Error('GCD requires at least two numbers') + if (numbers.length < 2) { + throw new Error('GCD requires at least two numbers') + } + return numbers.reduce((a, b) => { + a = Math.abs(a) + b = Math.abs(b) + while (b !== 0) { + const temp = b + b = a % b + a = temp } - return numbers.reduce((a, b) => { - a = Math.abs(a) - b = Math.abs(b) - while (b !== 0) { - const temp = b - b = a % b - a = temp - } - return a - }) + return a + }) } function lcm(...numbers: number[]): number { - if (numbers.length < 2) { - throw new Error('LCM requires at least two numbers') - } - return numbers.reduce((a, b) => Math.abs(a * b) / gcd(a, b)) + if (numbers.length < 2) { + throw new Error('LCM requires at least two numbers') + } + return numbers.reduce((a, b) => Math.abs(a * b) / gcd(a, b)) } function isPrime(n: number): boolean { - if (!Number.isInteger(n)) { - throw new Error('Prime check requires an integer') - } - if (n <= 1) { - return false - } - if (n <= 3) { - return true - } - if (n % 2 === 0 || n % 3 === 0) { - return false - } - for (let i = 5; i * i <= n; i += 6) { - if (n % i === 0 || n % (i + 2) === 0) { - return false - } - } + if (!Number.isInteger(n)) { + throw new Error('Prime check requires an integer') + } + if (n <= 1) { + return false + } + if (n <= 3) { return true + } + if (n % 2 === 0 || n % 3 === 0) { + return false + } + for (let i = 5; i * i <= n; i += 6) { + if (n % i === 0 || n % (i + 2) === 0) { + return false + } + } + return true } function combinations(n: number, k: number): number { - if ( - !Number.isInteger(n) || - !Number.isInteger(k) || - n < 0 || - k < 0 || - k > n - ) { - throw new Error('Invalid parameters for combinations') - } - return factorial(n) / (factorial(k) * factorial(n - k)) + if ( + !Number.isInteger(n) || + !Number.isInteger(k) || + n < 0 || + k < 0 || + k > n + ) { + throw new Error('Invalid parameters for combinations') + } + return factorial(n) / (factorial(k) * factorial(n - k)) } function permutations(n: number, k: number): number { - if ( - !Number.isInteger(n) || - !Number.isInteger(k) || - n < 0 || - k < 0 || - k > n - ) { - throw new Error('Invalid parameters for permutations') - } - return factorial(n) / factorial(n - k) + if ( + !Number.isInteger(n) || + !Number.isInteger(k) || + n < 0 || + k < 0 || + k > n + ) { + throw new Error('Invalid parameters for permutations') + } + return factorial(n) / factorial(n - k) } function calculateDeterminant(matrix: number[][]): number { - const n = matrix.length - if (n === 1) { - return matrix[0][0] - } - if (n === 2) { - return matrix[0][0] * matrix[1][1] - matrix[0][1] * matrix[1][0] - } - - let det = 0 - for (let i = 0; i < n; i++) { - const subMatrix = matrix - .slice(1) - .map((row) => row.filter((_, j) => j !== i)) - det += Math.pow(-1, i) * matrix[0][i] * calculateDeterminant(subMatrix) - } - return det + const n = matrix.length + if (n === 1) { + return matrix[0][0] + } + if (n === 2) { + return matrix[0][0] * matrix[1][1] - matrix[0][1] * matrix[1][0] + } + + let det = 0 + for (let i = 0; i < n; i++) { + const subMatrix = matrix + .slice(1) + .map((row) => row.filter((_, j) => j !== i)) + det += Math.pow(-1, i) * matrix[0][i] * calculateDeterminant(subMatrix) + } + return det } export type CalculatorUITool = InferUITool<typeof calculatorTool> diff --git a/src/mastra/tools/calendar-tool.ts b/src/mastra/tools/calendar-tool.ts index 51731c3e..4a98edb6 100644 --- a/src/mastra/tools/calendar-tool.ts +++ b/src/mastra/tools/calendar-tool.ts @@ -1,5 +1,5 @@ +import { SpanType } from '@mastra/core/observability' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' import { execSync } from 'child_process' import { z } from 'zod' import { log } from '../config/logger' @@ -173,18 +173,20 @@ export const listEvents = createTool({ }) }, execute: async (inputData, context) => { - const tracer = trace.getTracer('calendar-tool', '1.0.0') + const tracingContext = context?.tracingContext const requestCtx = context?.requestContext as CalendarRequestContext | undefined - const span = tracer.startSpan('list-calendar-events', { - attributes: { + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'list-calendar-events', + input: inputData, + metadata: { 'tool.id': 'list-calendar-events', 'tool.input.startDate': inputData.startDate, + 'user.id': requestCtx?.userId, }, }) - if (requestCtx?.userId) { - log.debug('Executing calendar list events for user', { userId: requestCtx.userId }) - } + log.debug('Executing calendar list events for user', { userId: requestCtx?.userId }) await context?.writer?.custom({ type: 'data-tool-progress', @@ -217,11 +219,14 @@ export const listEvents = createTool({ }, id: 'listEvents', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.eventCount': events.length, + span?.update({ + output: { eventCount: events.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.eventCount': events.length, + }, }) - span.end() + span?.end() return { content: JSON.stringify(formattedEvents, null, 2), @@ -231,11 +236,10 @@ export const listEvents = createTool({ } catch (e) { const errorMsg = e instanceof Error ? e.message : 'Unknown error' log.error(`Calendar read failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { content: `Error: ${errorMsg}` } } }, @@ -258,17 +262,19 @@ export const getTodayEvents = createTool({ count: z.number(), }), execute: async (inputData, context) => { - const tracer = trace.getTracer('calendar-tool', '1.0.0') + const tracingContext = context?.tracingContext const requestCtx = context?.requestContext as CalendarRequestContext | undefined - const span = tracer.startSpan('get-today-events', { - attributes: { + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'get-today-events', + input: inputData, + metadata: { 'tool.id': 'get-today-events', + 'user.id': requestCtx?.userId, }, }) - if (requestCtx?.userId) { - log.debug('Executing get today events for user', { userId: requestCtx.userId }) - } + log.debug('Executing get today events for user', { userId: requestCtx?.userId }) await context?.writer?.write({ type: 'progress', @@ -301,21 +307,23 @@ export const getTodayEvents = createTool({ message: `✅ Found ${todayEvents.length} events for today`, }, }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.eventCount': todayEvents.length, + span?.update({ + output: { eventCount: todayEvents.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.eventCount': todayEvents.length, + }, }) - span.end() + span?.end() return { events: formattedEvents, count: todayEvents.length } } catch (e) { const errorMsg = e instanceof Error ? e.message : 'Unknown error' log.error(`Today events read failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { events: [], count: 0 } } }, @@ -381,19 +389,21 @@ export const getUpcomingEvents = createTool({ count: z.number(), }), execute: async (inputData, context) => { - const tracer = trace.getTracer('calendar-tool', '1.0.0') + const tracingContext = context?.tracingContext const requestCtx = context?.requestContext as CalendarRequestContext | undefined - const span = tracer.startSpan('get-upcoming-events', { - attributes: { + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'get-upcoming-events', + input: inputData, + metadata: { 'tool.id': 'get-upcoming-events', 'tool.input.days': inputData.days, 'tool.input.limit': inputData.limit, + 'user.id': requestCtx?.userId, }, }) - if (requestCtx?.userId) { - log.debug('Executing get upcoming events for user', { userId: requestCtx.userId }) - } + log.debug('Executing get upcoming events for user', { userId: requestCtx?.userId }) await context?.writer?.write({ type: 'progress', @@ -436,21 +446,23 @@ export const getUpcomingEvents = createTool({ message: `✅ Found ${upcomingEvents.length} upcoming events`, }, }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.eventCount': upcomingEvents.length, + span?.update({ + output: { eventCount: upcomingEvents.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.eventCount': upcomingEvents.length, + }, }) - span.end() + span?.end() return { events: formattedEvents, count: upcomingEvents.length } } catch (e) { const errorMsg = e instanceof Error ? e.message : 'Unknown error' log.error(`Upcoming events read failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { events: [], count: 0 } } }, @@ -530,18 +542,20 @@ export const findFreeSlots = createTool({ ), }), execute: async (inputData, context) => { - const tracer = trace.getTracer('calendar-tool', '1.0.0') + const tracingContext = context?.tracingContext const requestCtx = context?.requestContext as CalendarRequestContext | undefined - const span = tracer.startSpan('find-free-slots', { - attributes: { + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'find-free-slots', + input: inputData, + metadata: { 'tool.id': 'find-free-slots', 'tool.input.date': inputData.date, + 'user.id': requestCtx?.userId, }, }) - if (requestCtx?.userId) { - log.debug('Executing find free slots for user', { userId: requestCtx.userId }) - } + log.debug('Executing find free slots for user', { userId: requestCtx?.userId }) await context?.writer?.write({ type: 'progress', @@ -626,21 +640,23 @@ export const findFreeSlots = createTool({ type: 'progress', data: { message: `✅ Found ${freeSlots.length} free slots` }, }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.freeSlotCount': freeSlots.length, + span?.update({ + output: { freeSlotCount: freeSlots.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.freeSlotCount': freeSlots.length, + }, }) - span.end() + span?.end() return { freeSlots, busyPeriods } } catch (e) { const errorMsg = e instanceof Error ? e.message : 'Unknown error' log.error(`Free slots search failed: ${errorMsg}`) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { freeSlots: [], busyPeriods: [] } } }, diff --git a/src/mastra/tools/code-analysis.tool.ts b/src/mastra/tools/code-analysis.tool.ts index 820e3d54..b831c4cf 100644 --- a/src/mastra/tools/code-analysis.tool.ts +++ b/src/mastra/tools/code-analysis.tool.ts @@ -1,146 +1,146 @@ +import { SpanType } from '@mastra/core/observability' +import type { RequestContext } from '@mastra/core/request-context' import { createTool } from '@mastra/core/tools' -import { z } from 'zod' +import type { Options as FastGlobOptions } from 'fast-glob' +import fg from 'fast-glob' import { promises as fs } from 'node:fs' import * as path from 'node:path' -import fg from 'fast-glob' -import type { Options as FastGlobOptions } from 'fast-glob' import stripComments from 'strip-comments' -import { trace, SpanStatusCode } from '@opentelemetry/api' -import { PythonParser } from './semantic-utils' +import { z } from 'zod' import { log } from '../config/logger' -import type { RequestContext } from '@mastra/core/request-context' +import { PythonParser } from './semantic-utils' export interface CodeAnalysisRequestContext extends RequestContext { - userId?: string; + userId?: string; } async function fileExists(filePath: string): Promise<boolean> { - try { - await fs.access(filePath) - return true - } catch { - return false - } + try { + await fs.access(filePath) + return true + } catch { + return false + } } const fastGlobDefaults: FastGlobOptions = { - onlyFiles: true, - absolute: true, - dot: true, - unique: true, + onlyFiles: true, + absolute: true, + dot: true, + unique: true, } async function globFiles( - pattern: string, - options?: FastGlobOptions + pattern: string, + options?: FastGlobOptions ): Promise<string[]> { - return fg(pattern, { ...fastGlobDefaults, ...options }) + return fg(pattern, { ...fastGlobDefaults, ...options }) } const DEFAULT_IGNORE = [ - '**/node_modules/**', - '**/.git/**', - '**/dist/**', - '**/build/**', + '**/node_modules/**', + '**/.git/**', + '**/dist/**', + '**/build/**', ] const codeAnalysisInputSchema = z.object({ - target: z - .union([ - z.string().describe('File path or glob pattern'), - z.array(z.string()).describe('Array of file paths'), - ]) - .describe('Files to analyze'), - options: z - .object({ - includeMetrics: z - .boolean() - .default(true) - .describe('Calculate LOC and complexity'), - detectPatterns: z - .boolean() - .default(true) - .describe('Detect code patterns and issues'), - maxFileSize: z - .number() - .default(100000) - .describe('Skip files larger than this (bytes)'), - }) - .optional(), + target: z + .union([ + z.string().describe('File path or glob pattern'), + z.array(z.string()).describe('Array of file paths'), + ]) + .describe('Files to analyze'), + options: z + .object({ + includeMetrics: z + .boolean() + .default(true) + .describe('Calculate LOC and complexity'), + detectPatterns: z + .boolean() + .default(true) + .describe('Detect code patterns and issues'), + maxFileSize: z + .number() + .default(100000) + .describe('Skip files larger than this (bytes)'), + }) + .optional(), }) const issueSchema = z.object({ - type: z.enum(['error', 'warning', 'info']), - line: z.number(), - column: z.number().optional(), - message: z.string(), - rule: z.string().optional(), + type: z.enum(['error', 'warning', 'info']), + line: z.number(), + column: z.number().optional(), + message: z.string(), + rule: z.string().optional(), }) const fileAnalysisSchema = z.object({ - path: z.string(), - language: z.string(), - loc: z.number().describe('Lines of code (non-empty, non-comment)'), - totalLines: z.number(), - complexity: z.number().describe('Estimated cyclomatic complexity'), - issues: z.array(issueSchema), + path: z.string(), + language: z.string(), + loc: z.number().describe('Lines of code (non-empty, non-comment)'), + totalLines: z.number(), + complexity: z.number().describe('Estimated cyclomatic complexity'), + issues: z.array(issueSchema), }) const codeAnalysisOutputSchema = z.object({ - files: z.array(fileAnalysisSchema), - summary: z.object({ - totalFiles: z.number(), - totalLoc: z.number(), - avgComplexity: z.number(), - issueCount: z.record(z.string(), z.number()), - }), + files: z.array(fileAnalysisSchema), + summary: z.object({ + totalFiles: z.number(), + totalLoc: z.number(), + avgComplexity: z.number(), + issueCount: z.record(z.string(), z.number()), + }), }) export type CodeAnalysisInput = z.infer<typeof codeAnalysisInputSchema> export type CodeAnalysisOutput = z.infer<typeof codeAnalysisOutputSchema> const LANGUAGE_MAP: Record<string, string> = { - '.ts': 'typescript', - '.tsx': 'typescript', - '.js': 'javascript', - '.jsx': 'javascript', - '.py': 'python', - '.rs': 'rust', - '.go': 'go', - '.java': 'java', - '.cpp': 'cpp', - '.c': 'c', - '.cs': 'csharp', - '.rb': 'ruby', - '.php': 'php', - '.swift': 'swift', - '.kt': 'kotlin', - '.md': 'markdown', - '.json': 'json', - '.yaml': 'yaml', - '.yml': 'yaml', + '.ts': 'typescript', + '.tsx': 'typescript', + '.js': 'javascript', + '.jsx': 'javascript', + '.py': 'python', + '.rs': 'rust', + '.go': 'go', + '.java': 'java', + '.cpp': 'cpp', + '.c': 'c', + '.cs': 'csharp', + '.rb': 'ruby', + '.php': 'php', + '.swift': 'swift', + '.kt': 'kotlin', + '.md': 'markdown', + '.json': 'json', + '.yaml': 'yaml', + '.yml': 'yaml', } function detectLanguage(filePath: string): string { - const ext = path.extname(filePath).toLowerCase() - return LANGUAGE_MAP[ext] || 'unknown' + const ext = path.extname(filePath).toLowerCase() + return LANGUAGE_MAP[ext] || 'unknown' } function calculateLoc(content: string): { loc: number; totalLines: number } { - const totalLines = content.split('\n').length - let sanitized = content + const totalLines = content.split('\n').length + let sanitized = content - try { - sanitized = stripComments(content, { preserveNewlines: true }) - } catch { - sanitized = content - } + try { + sanitized = stripComments(content, { preserveNewlines: true }) + } catch { + sanitized = content + } - const loc = sanitized - .split('\n') - .reduce((count, line) => (line.trim() ? count + 1 : count), 0) + const loc = sanitized + .split('\n') + .reduce((count, line) => (line.trim() ? count + 1 : count), 0) - return { loc, totalLines } + return { loc, totalLines } } /** @@ -152,480 +152,495 @@ function calculateLoc(content: string): { loc: number; totalLines: number } { * @returns The estimated complexity number. */ function estimateComplexity(content: string, language: string): number { - let complexity = 1 - - const languagePatterns: Record<string, RegExp[]> = { - typescript: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - javascript: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - python: [ - /\bif\b/g, - /\belif\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\band\b/g, - /\bor\b/g, - ], - java: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - cpp: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - c: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - csharp: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - rust: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bmatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - go: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bcase\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - ruby: [ - /\bif\b/g, - /\belsif\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\band\b/g, - /\bor\b/g, - ], - php: [ - /\bif\b/g, - /\belseif\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - swift: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bcase\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - kotlin: [ - /\bif\b/g, - /\belse\s+if\b/g, - /\bfor\b/g, - /\bwhile\b/g, - /\bwhen\b/g, - /\bcatch\b/g, - /\b\?\s*[^:]+\s*:/g, // ternary - /&&/g, - /\|\|/g, - ], - } - - const patterns = - language in languagePatterns - ? languagePatterns[language] - : languagePatterns['typescript'] // default to typescript patterns - - for (const pattern of patterns) { - const matches = content.match(pattern) - if (matches) { - complexity += matches.length - } + let complexity = 1 + + const languagePatterns: Record<string, RegExp[]> = { + typescript: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + javascript: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + python: [ + /\bif\b/g, + /\belif\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\band\b/g, + /\bor\b/g, + ], + java: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + cpp: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + c: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + csharp: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + rust: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bmatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + go: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bcase\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + ruby: [ + /\bif\b/g, + /\belsif\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\band\b/g, + /\bor\b/g, + ], + php: [ + /\bif\b/g, + /\belseif\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + swift: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bcase\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + kotlin: [ + /\bif\b/g, + /\belse\s+if\b/g, + /\bfor\b/g, + /\bwhile\b/g, + /\bwhen\b/g, + /\bcatch\b/g, + /\b\?\s*[^:]+\s*:/g, // ternary + /&&/g, + /\|\|/g, + ], + } + + const patterns = + language in languagePatterns + ? languagePatterns[language] + : languagePatterns['typescript'] // default to typescript patterns + + for (const pattern of patterns) { + const matches = content.match(pattern) + if (matches) { + complexity += matches.length } + } - return complexity + return complexity } function detectIssues( - content: string, - language: string + content: string, + language: string ): Array<z.infer<typeof issueSchema>> { - const issues: Array<z.infer<typeof issueSchema>> = [] - const lines = content.split('\n') - - const patterns = [ - { - regex: /console\.log\(/g, - message: 'Console statement found', - rule: 'no-console', - type: 'warning' as const, - }, - { - regex: /debugger/g, - message: 'Debugger statement found', - rule: 'no-debugger', - type: 'error' as const, - }, - { - regex: /TODO:/gi, - message: 'TODO comment found', - rule: 'no-todo', - type: 'info' as const, - }, - { - regex: /FIXME:/gi, - message: 'FIXME comment found', - rule: 'no-fixme', - type: 'warning' as const, - }, - { - regex: /HACK:/gi, - message: 'HACK comment found', - rule: 'no-hack', - type: 'warning' as const, - }, - { - regex: /any(?:\s|;|,|\))/g, - message: 'Explicit any type', - rule: 'no-explicit-any', - type: 'warning' as const, - }, - { - regex: /==(?!=)/g, - message: 'Loose equality used', - rule: 'eqeqeq', - type: 'warning' as const, - }, - { - regex: /\bvar\b/g, - message: 'var keyword used (prefer const/let)', - rule: 'no-var', - type: 'info' as const, - }, - ] - - for (let i = 0; i < lines.length; i++) { - const line = lines[i] - - for (const pattern of patterns) { - if ( - language === 'typescript' || - language === 'javascript' || - pattern.rule === 'no-todo' || - pattern.rule === 'no-fixme' || - pattern.rule === 'no-hack' - ) { - if (pattern.regex.test(line)) { - issues.push({ - type: pattern.type, - line: i + 1, - message: pattern.message, - rule: pattern.rule, - }) - } - pattern.regex.lastIndex = 0 - } - } + const issues: Array<z.infer<typeof issueSchema>> = [] + const lines = content.split('\n') + + const patterns = [ + { + regex: /console\.log\(/g, + message: 'Console statement found', + rule: 'no-console', + type: 'warning' as const, + }, + { + regex: /debugger/g, + message: 'Debugger statement found', + rule: 'no-debugger', + type: 'error' as const, + }, + { + regex: /TODO:/gi, + message: 'TODO comment found', + rule: 'no-todo', + type: 'info' as const, + }, + { + regex: /FIXME:/gi, + message: 'FIXME comment found', + rule: 'no-fixme', + type: 'warning' as const, + }, + { + regex: /HACK:/gi, + message: 'HACK comment found', + rule: 'no-hack', + type: 'warning' as const, + }, + { + regex: /any(?:\s|;|,|\))/g, + message: 'Explicit any type', + rule: 'no-explicit-any', + type: 'warning' as const, + }, + { + regex: /==(?!=)/g, + message: 'Loose equality used', + rule: 'eqeqeq', + type: 'warning' as const, + }, + { + regex: /\bvar\b/g, + message: 'var keyword used (prefer const/let)', + rule: 'no-var', + type: 'info' as const, + }, + ] + + for (let i = 0; i < lines.length; i++) { + const line = lines[i] - if (line.length > 120) { - issues.push({ - type: 'info', - line: i + 1, - message: `Line exceeds 120 characters (${line.length})`, - rule: 'max-line-length', - }) + for (const pattern of patterns) { + if ( + language === 'typescript' || + language === 'javascript' || + pattern.rule === 'no-todo' || + pattern.rule === 'no-fixme' || + pattern.rule === 'no-hack' + ) { + if (pattern.regex.test(line)) { + issues.push({ + type: pattern.type, + line: i + 1, + message: pattern.message, + rule: pattern.rule, + }) } + pattern.regex.lastIndex = 0 + } } - return issues + if (line.length > 120) { + issues.push({ + type: 'info', + line: i + 1, + message: `Line exceeds 120 characters (${line.length})`, + rule: 'max-line-length', + }) + } + } + + return issues } export const codeAnalysisTool = createTool({ - id: 'coding:codeAnalysis', - description: `Analyze source code files for metrics and issues. + id: 'coding:codeAnalysis', + description: `Analyze source code files for metrics and issues. Returns lines of code, complexity estimates, and detected patterns. Supports TypeScript, JavaScript, Python, and other languages. Use for code review preparation, quality assessment, and refactoring planning.`, - inputSchema: codeAnalysisInputSchema, - outputSchema: codeAnalysisOutputSchema, - execute: async (inputData, context): Promise<CodeAnalysisOutput> => { - const writer = context?.writer - const requestCtx = context?.requestContext as CodeAnalysisRequestContext | undefined + inputSchema: codeAnalysisInputSchema, + outputSchema: codeAnalysisOutputSchema, + execute: async (inputData, context): Promise<CodeAnalysisOutput> => { + const writer = context?.writer + const requestCtx = context?.requestContext as CodeAnalysisRequestContext | undefined + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🔬 Starting code analysis for target: ${Array.isArray(inputData.target) ? inputData.target.join(',') : inputData.target}`, + stage: 'coding:codeAnalysis', + }, + id: 'coding:codeAnalysis', + }) + + if (requestCtx?.userId) { + log.debug('Executing code analysis for user', { userId: requestCtx.userId }) + } - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🔬 Starting code analysis for target: ${Array.isArray(inputData.target) ? inputData.target.join(',') : inputData.target}`, - stage: 'coding:codeAnalysis', - }, - id: 'coding:codeAnalysis', - }) + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'code-analysis', + input: { target: inputData.target }, + metadata: { + 'tool.id': 'coding:codeAnalysis', + 'tool.input.target': Array.isArray(inputData.target) + ? inputData.target.join(',') + : inputData.target, + }, + }) - if (requestCtx?.userId) { - log.debug('Executing code analysis for user', { userId: requestCtx.userId }) + try { + const { target, options } = inputData + const includeMetrics = options?.includeMetrics ?? true + const detectPatterns = options?.detectPatterns ?? true + const maxFileSize = options?.maxFileSize ?? 100000 + + const targets = Array.isArray(target) ? target : [target] + let filePaths: string[] = [] + + for (const t of targets) { + if (t.includes('*')) { + const matches = await globFiles(t, { + ignore: DEFAULT_IGNORE, + }) + filePaths.push(...matches) + continue } - const tracer = trace.getTracer('code-analysis-tool', '1.0.0') - const span = tracer.startSpan('code-analysis', { - attributes: { - 'tool.id': 'code-analysis', - 'tool.input.target': Array.isArray(inputData.target) - ? inputData.target.join(',') - : inputData.target, - }, - }) + if (!(await fileExists(t))) { + continue + } - try { - const { target, options } = inputData - const includeMetrics = options?.includeMetrics ?? true - const detectPatterns = options?.detectPatterns ?? true - const maxFileSize = options?.maxFileSize ?? 100000 - - const targets = Array.isArray(target) ? target : [target] - let filePaths: string[] = [] - - for (const t of targets) { - if (t.includes('*')) { - const matches = await globFiles(t, { - ignore: DEFAULT_IGNORE, - }) - filePaths.push(...matches) - continue - } - - if (!(await fileExists(t))) { - continue - } - - const stat = await fs.stat(t) - if (stat.isDirectory()) { - const dirFiles = await globFiles(path.join(t, '**/*'), { - ignore: DEFAULT_IGNORE, - }) - filePaths.push(...dirFiles) - } else if (stat.isFile()) { - filePaths.push(t) - } - } + const stat = await fs.stat(t) + if (stat.isDirectory()) { + const dirFiles = await globFiles(path.join(t, '**/*'), { + ignore: DEFAULT_IGNORE, + }) + filePaths.push(...dirFiles) + } else if (stat.isFile()) { + filePaths.push(t) + } + } - filePaths = [...new Set(filePaths)] - - const fileAnalyses: Array<z.infer<typeof fileAnalysisSchema>> = [] - - for (const filePath of filePaths) { - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `📄 Analyzing file: ${filePath}`, - stage: 'coding:codeAnalysis', - }, - id: 'coding:codeAnalysis', - }) - - try { - const stats = await fs.stat(filePath) - if (!stats.isFile()) { - continue - } - if (stats.size > maxFileSize) { - continue - } - - const content = await fs.readFile(filePath, 'utf-8') - const language = detectLanguage(filePath) - - const { loc, totalLines } = includeMetrics - ? calculateLoc(content) - : { loc: 0, totalLines: 0 } - - let complexity = 0 - if (includeMetrics) { - if (language === 'python') { - try { - const complexityData = - await PythonParser.analyzeComplexity( - content - ) - complexity = complexityData.cyclomaticComplexity - } catch { - complexity = estimateComplexity( - content, - language - ) - } - } else { - complexity = estimateComplexity(content, language) - } - } - - const issues = detectPatterns - ? detectIssues(content, language) - : [] - - fileAnalyses.push({ - path: filePath, - language, - loc, - totalLines, - complexity, - issues, - }) - } catch { - // Skip unreadable files - } - } + filePaths = [...new Set(filePaths)] - const totalLoc = fileAnalyses.reduce((sum, f) => sum + f.loc, 0) - const avgComplexity = - fileAnalyses.length > 0 - ? fileAnalyses.reduce((sum, f) => sum + f.complexity, 0) / - fileAnalyses.length - : 0 - - const issueCount: Record<string, number> = {} - for (const file of fileAnalyses) { - for (const issue of file.issues) { - issueCount[issue.type] = (issueCount[issue.type] || 0) + 1 - } - } + const fileAnalyses: Array<z.infer<typeof fileAnalysisSchema>> = [] + + for (const filePath of filePaths) { + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `📄 Analyzing file: ${filePath}`, + stage: 'coding:codeAnalysis', + }, + id: 'coding:codeAnalysis', + }) - span.setAttributes({ - 'tool.output.totalFiles': fileAnalyses.length, - 'tool.output.totalLoc': totalLoc, - 'tool.output.avgComplexity': avgComplexity, - }) - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Code analysis complete: ${fileAnalyses.length} files analyzed`, - stage: 'coding:codeAnalysis', - }, - id: 'coding:codeAnalysis', - }) - span.end() - - return { - files: fileAnalyses, - summary: { - totalFiles: fileAnalyses.length, - totalLoc, - avgComplexity: Math.round(avgComplexity * 100) / 100, - issueCount, - }, + try { + const stats = await fs.stat(filePath) + if (!stats.isFile()) { + continue + } + if (stats.size > maxFileSize) { + continue + } + + const content = await fs.readFile(filePath, 'utf-8') + const language = detectLanguage(filePath) + + const { loc, totalLines } = includeMetrics + ? calculateLoc(content) + : { loc: 0, totalLines: 0 } + + let complexity = 0 + if (includeMetrics) { + if (language === 'python') { + try { + const complexityData = + await PythonParser.analyzeComplexity( + content + ) + complexity = complexityData.cyclomaticComplexity + } catch { + complexity = estimateComplexity( + content, + language + ) + } + } else { + complexity = estimateComplexity(content, language) } - } catch (error: unknown) { - const errorMessage = - error instanceof Error ? error.message : String(error) - span?.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - span?.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, - }) - span?.end() - throw error + } + + const issues = detectPatterns + ? detectIssues(content, language) + : [] + + fileAnalyses.push({ + path: filePath, + language, + loc, + totalLines, + complexity, + issues, + }) + } catch { + // Skip unreadable files } - }, - 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', - }) - }, + } + + const totalLoc = fileAnalyses.reduce((sum, f) => sum + f.loc, 0) + const avgComplexity = + fileAnalyses.length > 0 + ? fileAnalyses.reduce((sum, f) => sum + f.complexity, 0) / + fileAnalyses.length + : 0 + + const issueCount: Record<string, number> = {} + for (const file of fileAnalyses) { + for (const issue of file.issues) { + issueCount[issue.type] = (issueCount[issue.type] || 0) + 1 + } + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Code analysis complete: ${fileAnalyses.length} file(s), ${totalLoc} lines`, + stage: 'coding:codeAnalysis', + }, + id: 'coding:codeAnalysis', + }) + + span?.update({ + output: { files: fileAnalyses }, + metadata: { + 'tool.output.totalFiles': fileAnalyses.length, + 'tool.output.totalLoc': totalLoc, + 'tool.output.avgComplexity': avgComplexity, + 'tool.output.issueCount': issueCount, + }, + }) + span?.end() + + return { + files: fileAnalyses, + summary: { + totalFiles: fileAnalyses.length, + totalLoc, + avgComplexity: Math.round(avgComplexity * 100) / 100, + issueCount, + }, + } + } catch (error: unknown) { + const errorMessage = + error instanceof Error ? error.message : String(error) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `❌ Code analysis failed: ${errorMessage}`, + stage: 'coding:codeAnalysis', + }, + id: 'coding:codeAnalysis', + }) + + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) + 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 462d175b..1bcdcd1d 100644 --- a/src/mastra/tools/code-chunking.ts +++ b/src/mastra/tools/code-chunking.ts @@ -1,255 +1,311 @@ +import type { RequestContext } from '@mastra/core/request-context' 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 { z } from 'zod' -import { PythonParser } from './semantic-utils' +import { SpanType } from '@mastra/core/observability' import { MDocument } from '@mastra/rag' import { log } from '../config/logger' -import { trace, SpanStatusCode } from '@opentelemetry/api' +import { PythonParser } from './semantic-utils' export interface CodeChunkingRequestContext extends RequestContext { - chunkSize?: number - chunkOverlap?: number + chunkSize?: number + chunkOverlap?: number } const codeChunkingInputSchema = z.object({ - filePath: z.string(), - content: z.string(), - options: z - .object({ - chunkSize: z.number().default(512), - chunkOverlap: z.number().default(50), - }) - .optional(), + filePath: z.string(), + content: z.string(), + options: z + .object({ + chunkSize: z.number().default(512), + chunkOverlap: z.number().default(50), + }) + .optional(), }) const codeChunkingOutputSchema = z.object({ - chunks: z.array( - z.object({ - text: z.string(), - metadata: z.object({ - startLine: z.number(), - endLine: z.number(), - type: z.string(), - name: z.string().optional(), - }), - }) - ), + chunks: z.array( + z.object({ + text: z.string(), + metadata: z.object({ + startLine: z.number(), + endLine: z.number(), + type: z.string(), + name: z.string().optional(), + }), + }) + ), }) export const codeChunkerTool = createTool({ - id: 'code-chunker', - description: - '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, - execute: async (input: z.infer<typeof codeChunkingInputSchema>, context) => { - const { filePath, content, options } = input - const requestContext = context?.requestContext as CodeChunkingRequestContext | undefined - const ext = filePath.split('.').pop()?.toLowerCase() - interface Chunk { - text: string - metadata: { - startLine: number - endLine: number - type: string - name?: string - } - } - const chunks: Chunk[] = [] + id: 'code-chunker', + description: + '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, + execute: async (input: z.infer<typeof codeChunkingInputSchema>, context) => { + const { filePath, content, options } = input + const requestContext = context?.requestContext as CodeChunkingRequestContext | undefined + const ext = filePath.split('.').pop()?.toLowerCase() + interface Chunk { + text: string + metadata: { + startLine: number + endLine: number + type: string + name?: string + } + } + const chunks: Chunk[] = [] - const tracer = trace.getTracer('code-chunking') - const span = tracer.startSpan('code-chunker', { - attributes: { filePath, ext, operation: 'code-chunking' }, - }) + const writer = context?.writer + const abortSignal = context?.abortSignal - try { - if (['ts', 'tsx', 'js', 'jsx'].includes(ext ?? '')) { - // TypeScript/JavaScript Chunking - const project = new Project({ useInMemoryFileSystem: true }) - const sourceFile = project.createSourceFile(filePath, content) - - sourceFile.forEachChild((node) => { - // typedNode narrows down the commonly used shape we need for code chunking - const typedNode = node as unknown as { - getName?: () => string - getText: () => string - getStartLineNumber: () => number - getEndLineNumber: () => number - getKindName: () => string - } - let name: string | undefined - - if ( - node.isKind(SyntaxKind.FunctionDeclaration) || - node.isKind(SyntaxKind.ClassDeclaration) || - node.isKind(SyntaxKind.InterfaceDeclaration) || - node.isKind(SyntaxKind.TypeAliasDeclaration) || - node.isKind(SyntaxKind.EnumDeclaration) || - node.isKind(SyntaxKind.VariableStatement) - ) { - if (typeof typedNode.getName === 'function') { - name = typedNode.getName() - } - - chunks.push({ - text: typedNode.getText(), - metadata: { - startLine: typedNode.getStartLineNumber(), - endLine: typedNode.getEndLineNumber(), - type: typedNode.getKindName(), - name, - }, - }) - } - }) - - // If no structural chunks found, or file is mostly top-level code, might want to fallback? - // For now, if we got chunks, return them. If not, fallback to MDocument. - if (chunks.length === 0 && content.trim().length > 0) { - // Fallback to recursive - } else { - span.setAttribute('chunksCount', chunks.length) - span.end() - return { chunks } - } - } + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🔄 Starting code chunking for ${filePath}`, + stage: 'code-chunker', + }, + id: 'code-chunker', + }) + + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'code-chunker', + input: { filePath, ext }, + metadata: { + 'tool.id': 'code-chunker', + operation: 'code-chunking', + }, + }) - if (ext === 'py') { - // Python Chunking - try { - const symbols = await PythonParser.findSymbols(content) - const lines = content.split('\n') - - for (const symbol of symbols) { - if ( - (symbol.kind === 'function' || - symbol.kind === 'class') && - symbol.endLine !== null && - symbol.endLine !== undefined && - symbol.endLine > 0 - ) { - const chunkLines = lines.slice( - symbol.line - 1, - symbol.endLine - ) - chunks.push({ - text: chunkLines.join('\n'), - metadata: { - startLine: symbol.line, - endLine: symbol.endLine, - type: symbol.kind, - name: symbol.name, - }, - }) - } - } - - if (chunks.length > 0) { - span.setAttribute('chunksCount', chunks.length) - span.end() - return { chunks } - } - } catch (e) { - log.warn( - `Python parsing failed for ${filePath}, falling back to text chunking: ${e}` - ) - } + try { + if (['ts', 'tsx', 'js', 'jsx'].includes(ext ?? '')) { + // TypeScript/JavaScript Chunking + const project = new Project({ useInMemoryFileSystem: true }) + const sourceFile = project.createSourceFile(filePath, content) + + sourceFile.forEachChild((node) => { + // typedNode narrows down the commonly used shape we need for code chunking + const typedNode = node as unknown as { + getName?: () => string + getText: () => string + getStartLineNumber: () => number + getEndLineNumber: () => number + getKindName: () => string + } + let name: string | undefined + + if ( + node.isKind(SyntaxKind.FunctionDeclaration) || + node.isKind(SyntaxKind.ClassDeclaration) || + node.isKind(SyntaxKind.InterfaceDeclaration) || + node.isKind(SyntaxKind.TypeAliasDeclaration) || + node.isKind(SyntaxKind.EnumDeclaration) || + node.isKind(SyntaxKind.VariableStatement) + ) { + if (typeof typedNode.getName === 'function') { + name = typedNode.getName() } - // Fallback: Recursive Chunking via MDocument - const doc = new MDocument({ - docs: [{ text: content, metadata: { source: filePath } }], - type: 'document', + chunks.push({ + text: typedNode.getText(), + metadata: { + startLine: typedNode.getStartLineNumber(), + endLine: typedNode.getEndLineNumber(), + type: typedNode.getKindName(), + name, + }, }) + } + }) - const textChunks = await doc.chunk({ - strategy: 'recursive', - maxSize: options?.chunkSize ?? requestContext?.chunkSize ?? 512, - overlap: options?.chunkOverlap ?? requestContext?.chunkOverlap ?? 50, - separators: ['\n\n', '\n', ' '], - }) + // If no structural chunks found, or file is mostly top-level code, might want to fallback? + // For now, if we got chunks, return them. If not, fallback to MDocument. + if (chunks.length === 0 && content.trim().length > 0) { + // Fallback to recursive + } else { + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Chunking complete: ${chunks.length} chunk${chunks.length !== 1 ? 's' : ''}`, + stage: 'code-chunker', + }, + id: 'code-chunker', + }) + span?.update({ metadata: { 'tool.output.chunksCount': chunks.length } }) + span?.end() + return { chunks } + } + } + + if (ext === 'py') { + // Python Chunking + try { + const symbols = await PythonParser.findSymbols(content) + const lines = content.split('\n') - const mapped = textChunks.map((c, i) => ({ - text: c.text ?? '', + for (const symbol of symbols) { + if ( + (symbol.kind === 'function' || + symbol.kind === 'class') && + symbol.endLine !== null && + symbol.endLine !== undefined && + symbol.endLine > 0 + ) { + const chunkLines = lines.slice( + symbol.line - 1, + symbol.endLine + ) + chunks.push({ + text: chunkLines.join('\n'), metadata: { - startLine: 0, // MDocument doesn't give line numbers easily - endLine: 0, - type: 'text-chunk', - name: `chunk-${i}`, + startLine: symbol.line, + endLine: symbol.endLine, + type: symbol.kind, + name: symbol.name, }, - })) - - span.setAttribute('chunksCount', mapped.length) - span.end() - return { chunks: mapped } - } catch (error) { - const errorMeta = - error instanceof Error - ? { message: error.message, stack: error.stack } - : { value: String(error) } - - log.error(`Error chunking ${filePath}:`, errorMeta) - // Trace the exception and set status - span?.recordException( - error instanceof Error ? error : new Error(String(error)) - ) - span?.setStatus({ - code: SpanStatusCode.ERROR, - message: error instanceof Error ? error.message : String(error), + }) + } + } + + if (chunks.length > 0) { + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Chunking complete: ${chunks.length} chunk${chunks.length !== 1 ? 's' : ''}`, + stage: 'code-chunker', + }, + id: 'code-chunker', }) + span?.update({ metadata: { 'tool.output.chunksCount': chunks.length } }) span?.end() - - // Final fallback - return { - chunks: [ - { - text: content, - metadata: { - startLine: 1, - endLine: content.split('\n').length, - type: 'file', - name: 'full-file', - }, - }, - ], - } + return { chunks } + } + } catch (e) { + log.warn( + `Python parsing failed for ${filePath}, falling back to text chunking: ${e}` + ) } - }, - 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', + } + + // Fallback: Recursive Chunking via MDocument + const doc = new MDocument({ + docs: [{ text: content, metadata: { source: filePath } }], + type: 'document', + }) + + if (abortSignal?.aborted) { + span?.error({ error: new Error('Code chunking cancelled'), endSpan: true }) + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `🛑 Code chunking cancelled for ${filePath}`, + stage: 'code-chunker', + }, + id: 'code-chunker', }) - }, + throw new Error('Code chunking cancelled') + } + + const textChunks = await doc.chunk({ + strategy: 'recursive', + 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 ?? '', + metadata: { + startLine: 0, // MDocument doesn't give line numbers easily + endLine: 0, + type: 'text-chunk', + name: `chunk-${i}`, + }, + })) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Chunking complete: ${mapped.length} chunk${mapped.length !== 1 ? 's' : ''}`, + stage: 'code-chunker', + }, + id: 'code-chunker', + }) + span?.update({ metadata: { 'tool.output.chunksCount': mapped.length } }) + span?.end() + return { chunks: mapped } + } catch (error) { + const errorMeta = + error instanceof Error + ? { message: error.message, stack: error.stack } + : { value: String(error) } + + log.error(`Error chunking ${filePath}:`, errorMeta) + // Trace the exception and set status + span?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + + // Final fallback + return { + chunks: [ + { + text: content, + metadata: { + startLine: 1, + endLine: content.split('\n').length, + type: 'file', + name: 'full-file', + }, + }, + ], + } + } + }, + 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 b231f430..edeebfe5 100644 --- a/src/mastra/tools/code-search.tool.ts +++ b/src/mastra/tools/code-search.tool.ts @@ -1,377 +1,382 @@ +import { SpanType } from '@mastra/core/observability' +import type { RequestContext } from '@mastra/core/request-context' import { createTool } from '@mastra/core/tools' -import { SpanStatusCode, trace } from '@opentelemetry/api' import fg, { type Options as FastGlobOptions } from 'fast-glob' 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 + caseSensitive?: boolean + maxResults?: number + includeContext?: boolean + contextLines?: number + maxFileSize?: number } interface RegexLike { - lastIndex?: number - exec: RegExp['exec'] + lastIndex?: number + exec: RegExp['exec'] } async function fileExists(filePath: string): Promise<boolean> { - try { - await fs.access(filePath) - return true - } catch { - return false - } + try { + await fs.access(filePath) + return true + } catch { + return false + } } const defaultGlobOptions: FastGlobOptions = { - onlyFiles: true, - absolute: true, - dot: true, - unique: true, - ignore: ['**/node_modules/**', '**/.git/**', '**/dist/**', '**/build/**'], + onlyFiles: true, + absolute: true, + dot: true, + unique: true, + ignore: ['**/node_modules/**', '**/.git/**', '**/dist/**', '**/build/**'], } async function globFiles( - pattern: string, - options?: FastGlobOptions + pattern: string, + options?: FastGlobOptions ): Promise<string[]> { - return fg(pattern, { ...defaultGlobOptions, ...options }) + return fg(pattern, { ...defaultGlobOptions, ...options }) } const codeSearchInputSchema = z.object({ - pattern: z.string().describe('Search pattern (string or regex)'), - target: z - .union([ - z.string().describe('File path or glob pattern'), - z.array(z.string()).describe('Array of file paths or globs'), - ]) - .describe('Files to search'), - options: z - .object({ - isRegex: z.boolean().optional().describe('Treat pattern as regex'), - caseSensitive: z - .boolean() - .optional() - .describe('Case-sensitive search'), - maxResults: z - .number() - .optional() - .describe('Maximum results to return'), - includeContext: z - .boolean() - .optional() - .describe('Include surrounding lines'), - contextLines: z - .number() - .optional() - .describe('Number of context lines'), - maxFileSize: z - .number() - .optional() - .describe('Skip files larger than this (bytes)'), - }) - .optional(), + pattern: z.string().describe('Search pattern (string or regex)'), + target: z + .union([ + z.string().describe('File path or glob pattern'), + z.array(z.string()).describe('Array of file paths or globs'), + ]) + .describe('Files to search'), + options: z + .object({ + isRegex: z.boolean().optional().describe('Treat pattern as regex'), + caseSensitive: z + .boolean() + .optional() + .describe('Case-sensitive search'), + maxResults: z + .number() + .optional() + .describe('Maximum results to return'), + includeContext: z + .boolean() + .optional() + .describe('Include surrounding lines'), + contextLines: z + .number() + .optional() + .describe('Number of context lines'), + maxFileSize: z + .number() + .optional() + .describe('Skip files larger than this (bytes)'), + }) + .optional(), }) const matchSchema = z.object({ - file: z.string(), - line: z.number(), - column: z.number(), - content: z.string(), - context: z - .object({ - before: z.array(z.string()), - after: z.array(z.string()), - }) - .optional(), + file: z.string(), + line: z.number(), + column: z.number(), + content: z.string(), + context: z + .object({ + before: z.array(z.string()), + after: z.array(z.string()), + }) + .optional(), }) const codeSearchOutputSchema = z.object({ - matches: z.array(matchSchema), - stats: z.object({ - totalMatches: z.number(), - filesSearched: z.number(), - filesWithMatches: z.number(), - }), - truncated: z.boolean(), + matches: z.array(matchSchema), + stats: z.object({ + totalMatches: z.number(), + filesSearched: z.number(), + filesWithMatches: z.number(), + }), + truncated: z.boolean(), }) export type CodeSearchInput = z.infer<typeof codeSearchInputSchema> export type CodeSearchOutput = z.infer<typeof codeSearchOutputSchema> export const codeSearchTool = createTool({ - id: 'coding:codeSearch', - description: `Search for patterns across source files. + id: 'coding:codeSearch', + description: `Search for patterns across source files. Supports string and regex patterns with context lines. Use for finding usages, identifying patterns, and code exploration.`, - inputSchema: codeSearchInputSchema, - outputSchema: codeSearchOutputSchema, - execute: async (input: z.infer<typeof codeSearchInputSchema>, context): Promise<CodeSearchOutput> => { - 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) { - throw new Error('Code search cancelled') + inputSchema: codeSearchInputSchema, + outputSchema: codeSearchOutputSchema, + execute: async (input: z.infer<typeof codeSearchInputSchema>, context): Promise<CodeSearchOutput> => { + 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) { + throw new Error('Code search cancelled') + } + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🔎 Starting code search for pattern '${input.pattern}'`, + stage: 'coding:codeSearch', + }, + id: 'coding:codeSearch', + }) + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'code-search', + input: { + pattern: input.pattern, + targetCount: Array.isArray(input.target) ? input.target.length : 1, + }, + metadata: { + 'tool.id': 'coding:codeSearch', + 'tool.input.pattern': input.pattern, + operation: 'code-search', + }, + }) + + try { + const { pattern, target, options } = input + const isRegex = options?.isRegex ?? false + 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] + + for (const t of targets) { + if (t.includes('*')) { + const matches = await globFiles(t) + filePaths.push(...matches) + } else if (await fileExists(t)) { + const stat = await fs.stat(t) + if (stat.isFile()) { + filePaths.push(t) + } else if (stat.isDirectory()) { + const dirFiles = await globFiles(path.join(t, '**/*')) + filePaths.push(...dirFiles) + } } + } + + filePaths = [...new Set(filePaths)] + + // Check for cancellation before file processing + if (abortSignal?.aborted) { + const cancelMessage = 'Code search cancelled during file processing' + span?.error({ error: new Error(cancelMessage), endSpan: true }) await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🔎 Starting code search for pattern '${input.pattern}'`, - stage: 'coding:codeSearch', - }, - id: 'coding:codeSearch', - }) - const tracer = trace.getTracer('code-search') - const span = tracer.startSpan('code-search', { - attributes: { - pattern: input.pattern, - targetCount: Array.isArray(input.target) - ? input.target.length - : 1, - operation: 'code-search', - }, + type: 'data-tool-progress', + data: { + status: 'done', + message: `🛑 ${cancelMessage}`, + stage: 'coding:codeSearch', + }, + id: 'coding:codeSearch', }) + log.warn(cancelMessage) + throw new Error(cancelMessage) + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `📁 Files to search: ${filePaths.length}`, + stage: 'coding:codeSearch', + }, + id: 'coding:codeSearch', + }) + + const matches: Array<z.infer<typeof matchSchema>> = [] + const filesWithMatches = new Set<string>() + let truncated = false + + let searchRegex: RegexLike + try { + if (isRegex) { + searchRegex = new RE2( + pattern, + caseSensitive ? 'g' : 'gi' + ) as unknown as RegexLike + } else { + searchRegex = new RegExp( + pattern.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), + caseSensitive ? 'g' : 'gi' + ) + } + } catch (error) { + throw new Error( + `Invalid ${isRegex ? 'regex' : 'pattern'}: ${error instanceof Error ? error.message : String(error)}` + ) + } + + for (const filePath of filePaths) { + if (matches.length >= maxResults) { + truncated = true + break + } try { - const { pattern, target, options } = input - const isRegex = options?.isRegex ?? false - 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] - - for (const t of targets) { - if (t.includes('*')) { - const matches = await globFiles(t) - filePaths.push(...matches) - } else if (await fileExists(t)) { - const stat = await fs.stat(t) - if (stat.isFile()) { - filePaths.push(t) - } else if (stat.isDirectory()) { - const dirFiles = await globFiles(path.join(t, '**/*')) - filePaths.push(...dirFiles) - } - } + const stat = await fs.stat(filePath) + if (!stat.isFile()) { + continue + } + if (stat.size > maxFileSize) { + continue + } + const content = await fs.readFile(filePath, 'utf-8') + const lines = content.split('\n') + + for (let i = 0; i < lines.length; i++) { + if (matches.length >= maxResults) { + truncated = true + break } - filePaths = [...new Set(filePaths)] - - // Check for cancellation before file processing - if (abortSignal?.aborted) { - span.setStatus({ - code: 2, - message: 'Operation cancelled during file processing', - }) - span.end() - throw new Error('Code search cancelled during file processing') + const line = lines[i] + if (typeof searchRegex.lastIndex === 'number') { + searchRegex.lastIndex = 0 } - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `📁 Files to search: ${filePaths.length}`, - stage: 'coding:codeSearch', - }, - id: 'coding:codeSearch', - }) - - const matches: Array<z.infer<typeof matchSchema>> = [] - const filesWithMatches = new Set<string>() - let truncated = false - - let searchRegex: RegexLike - try { - if (isRegex) { - searchRegex = new RE2( - pattern, - caseSensitive ? 'g' : 'gi' - ) as unknown as RegexLike - } else { - searchRegex = new RegExp( - pattern.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), - caseSensitive ? 'g' : 'gi' - ) - } - } catch (error) { - throw new Error( - `Invalid ${isRegex ? 'regex' : 'pattern'}: ${error instanceof Error ? error.message : String(error)}` - ) - } - - for (const filePath of filePaths) { - if (matches.length >= maxResults) { - truncated = true - break - } - - try { - const stat = await fs.stat(filePath) - if (!stat.isFile()) { - continue - } - if (stat.size > maxFileSize) { - continue - } - const content = await fs.readFile(filePath, 'utf-8') - const lines = content.split('\n') - - for (let i = 0; i < lines.length; i++) { - if (matches.length >= maxResults) { - truncated = true - break - } - - const line = lines[i] - if (typeof searchRegex.lastIndex === 'number') { - searchRegex.lastIndex = 0 - } - let match - - while ((match = searchRegex.exec(line)) !== null) { - filesWithMatches.add(filePath) - - const contextObj = includeContext - ? { - before: lines.slice( - Math.max(0, i - contextLines), - i - ), - after: lines.slice( - i + 1, - i + 1 + contextLines - ), - } - : undefined - - matches.push({ - file: filePath, - line: i + 1, - column: match.index + 1, - content: line, - context: contextObj, - }) - - if (!isRegex) { - break - } - } - } - } catch { - // Skip binary or unreadable files + let match + + while ((match = searchRegex.exec(line)) !== null) { + filesWithMatches.add(filePath) + + const contextObj = includeContext + ? { + before: lines.slice( + Math.max(0, i - contextLines), + i + ), + after: lines.slice( + i + 1, + i + 1 + contextLines + ), } + : undefined + + matches.push({ + file: filePath, + line: i + 1, + column: match.index + 1, + content: line, + context: contextObj, + }) + + if (!isRegex) { + break + } } + } + } catch { + // Skip binary or unreadable files + } + } + + const result = { + matches, + stats: { + totalMatches: matches.length, + filesSearched: filePaths.length, + filesWithMatches: filesWithMatches.size, + }, + truncated, + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Code search complete: ${result.stats.totalMatches} matches across ${result.stats.filesWithMatches} files`, + stage: 'coding:codeSearch', + }, + id: 'coding:codeSearch', + }) + span?.end() + return result + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : String(error) + + // Handle AbortError specifically + if (error instanceof Error && error.name === 'AbortError') { + const cancelMessage = `Code search cancelled for pattern "${input.pattern}"` + span?.error({ error: new Error(cancelMessage), endSpan: true }) - const result = { - matches, - stats: { - totalMatches: matches.length, - filesSearched: filePaths.length, - filesWithMatches: filesWithMatches.size, - }, - truncated, - } + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `🛑 ${cancelMessage}`, + stage: 'coding:codeSearch', + }, + id: 'coding:codeSearch', + }) - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Code search complete: ${result.stats.totalMatches} matches across ${result.stats.filesWithMatches} files`, - stage: 'coding:codeSearch', - }, - id: 'coding:codeSearch', - }) - span.end() - return result - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : String(error) - - // Handle AbortError specifically - if (error instanceof Error && error.name === 'AbortError') { - const cancelMessage = `Code search cancelled for pattern "${input.pattern}"` - span.setStatus({ code: 2, message: cancelMessage }) - span.end() - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `🛑 ${cancelMessage}`, - stage: 'coding:codeSearch', - }, - id: 'coding:codeSearch', - }) - - log.warn(cancelMessage) - throw new Error(cancelMessage) - } + log.warn(cancelMessage) + throw new Error(cancelMessage) + } - span.recordException(new Error(errorMessage)) - span.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, - }) - span.end() - 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', - }) - }, + span?.error({ error: new Error(errorMessage), endSpan: true }) + 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 5f0f9ed4..0dd78b3c 100644 --- a/src/mastra/tools/color-change-tool.ts +++ b/src/mastra/tools/color-change-tool.ts @@ -1,85 +1,142 @@ +import { SpanType } from '@mastra/core/observability' +import type { RequestContext } from '@mastra/core/request-context' 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 + theme?: string + userId?: string + workspaceId?: string } export function changeBgColor(color: string) { - if (typeof window !== 'undefined') { - document.body.style.setProperty('--background', color) - } + if (typeof window !== 'undefined') { + document.body.style.setProperty('--background', color) + } } export const colorChangeTool = createTool({ - id: 'changeColor', - description: 'Changes the background color', - inputSchema: z.object({ - color: z.string(), - }), - execute: async (input, context) => { - const { color } = input; - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🎨 Changing background color to ${color}`, - stage: 'changeColor', - }, - id: 'changeColor', - }) + id: 'changeColor', + description: 'Changes the background color', + inputSchema: z.object({ + color: z.string(), + }), + execute: async (input, context) => { + const { color } = input; + const writer = context?.writer; + const abortSignal = context?.abortSignal; + const tracingContext = context?.tracingContext; + const requestCtx = context?.requestContext as ColorChangeRequestContext | undefined; + const userId = requestCtx?.userId; + const workspaceId = requestCtx?.workspaceId; + + // Respect cancellation early + if (abortSignal?.aborted ?? false) { + throw new Error('Change color operation cancelled') + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🎨 Changing background color to ${color}`, + stage: 'changeColor', + }, + id: 'changeColor', + }) - // On the server, we just return the color. - // The client-side UI will handle the actual style change when it receives the tool result or call. + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'change-color', + input: { color }, + metadata: { + 'tool.id': 'changeColor', + 'tool.input.color': color, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Background color changed to ${color}`, - stage: 'changeColor', - }, - id: 'changeColor', - }) + try { + // On the server, we just return the color. + // The client-side UI will handle the actual style change when it receives the tool result or call. - return { success: true, color } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Color change tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Color change received input chunk', { - toolCallId, - inputTextDelta, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Color change received complete input', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - color: input.color, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Color change completed', { - toolCallId, - toolName, - abortSignal: abortSignal?.aborted, - success: output.success, - color: output.color, - hook: 'onOutput', - }) - }, + span?.update({ + output: { success: true, color }, + metadata: { + 'tool.output.success': true, + }, + }) + span?.end() + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Background color changed to ${color}`, + stage: 'changeColor', + }, + id: 'changeColor', + }) + + return { success: true, color } + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error) + + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `Input: color="${color}" - ❌ Failed to change color: ${errorMessage}`, + stage: 'changeColor', + }, + id: 'changeColor', + }) + + log.error(`Failed to change background color to ${color}: ${errorMessage}`) + throw error + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Color change tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Color change received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Color change received complete input', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + color: input.color, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Color change completed', { + toolCallId, + toolName, + abortSignal: abortSignal?.aborted, + success: output.success, + color: output.color, + hook: 'onOutput', + }) + }, }) diff --git a/src/mastra/tools/copywriter-agent-tool.ts b/src/mastra/tools/copywriter-agent-tool.ts index 32b555b9..5e400545 100644 --- a/src/mastra/tools/copywriter-agent-tool.ts +++ b/src/mastra/tools/copywriter-agent-tool.ts @@ -1,7 +1,7 @@ +import { SpanType } from '@mastra/core/observability' import type { MastraModelOutput } from '@mastra/core/stream' 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' @@ -75,6 +75,7 @@ export const copywriterTool = createTool({ execute: async (input, context) => { const writer = context?.writer const mastra = context?.mastra + const tracingContext = context?.tracingContext const requestCtx = context?.requestContext as CopywriterRequestContext | undefined const userId = @@ -100,18 +101,21 @@ export const copywriterTool = createTool({ id: 'copywriter-agent', }) - const span = trace - .getTracer('copywriter-agent-tool', '1.0.0') - .startSpan('copywriter-generate', { - attributes: { - 'tool.id': 'copywriter-agent', - 'tool.input.topic': topic, - 'tool.input.contentType': contentType, - 'tool.input.targetAudience': targetAudience, - 'tool.input.tone': tone, - 'tool.input.length': length, - }, - }) + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'copywriter-generate', + input: { + topic, + contentType, + targetAudience, + tone, + length, + }, + metadata: { + 'tool.id': 'copywriter-agent', + 'user.id': requestCtx?.userId, + }, + }) try { const agent = @@ -283,18 +287,10 @@ export const copywriterTool = createTool({ } } catch (err) { const msg = err instanceof Error ? err.message : String(err) - span.recordException( - err instanceof Error ? err : new Error(msg) - ) - try { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: msg, - }) - } catch { - /* ignore */ - } - span.end() + span?.error({ + error: err instanceof Error ? err : new Error(msg), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', data: { @@ -354,10 +350,15 @@ export const copywriterTool = createTool({ ? firstParagraph.substring(0, 200) + '...' : firstParagraph - span.setAttribute('tool.output.success', true) - span.setAttribute('tool.output.wordCount', wordCount) - span.setAttribute('tool.output.contentLength', content.length) - span.end() + span?.update({ + output: { content, contentType, title, wordCount }, + metadata: { + 'tool.output.success': true, + 'tool.output.wordCount': wordCount, + 'tool.output.contentLength': content.length, + } + }) + span?.end() return { content, @@ -375,9 +376,10 @@ export const copywriterTool = createTool({ topic, contentType, }) - span?.setAttribute('tool.output.success', false) - span?.setAttribute('tool.output.error', errorMsg) - span?.end() + span?.error({ + error: error instanceof Error ? error : new Error(errorMsg), + endSpan: true, + }) throw new Error(`Failed to generate content: ${errorMsg}`) } }, diff --git a/src/mastra/tools/csv-to-json.tool.ts b/src/mastra/tools/csv-to-json.tool.ts index f892ed81..e871ab74 100644 --- a/src/mastra/tools/csv-to-json.tool.ts +++ b/src/mastra/tools/csv-to-json.tool.ts @@ -1,7 +1,7 @@ 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 { SpanType } from '@mastra/core/observability' import { parse } from 'csv-parse/sync' import * as fs from 'node:fs/promises' import { z } from 'zod' @@ -51,9 +51,10 @@ export const csvToJsonTool = createTool({ const writer = context?.writer const requestContext = context?.requestContext as CsvToJsonRequestContext | undefined const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext // Check if operation was already cancelled - if (abortSignal?.aborted === true) { + if (abortSignal?.aborted) { throw new Error('CSV to JSON conversion cancelled') } @@ -66,10 +67,16 @@ export const csvToJsonTool = createTool({ }, id: 'csv-to-json', }) - const tracer = trace.getTracer('csv-to-json', '1.0.0') - const rootSpan = tracer.startSpan('csv-to-json', { - attributes: { + + // Create child span for CSV conversion + const csvSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'csv-to-json-conversion', + input, + metadata: { 'tool.id': 'csv-to-json', + 'tool.input.hasFilePath': !!input.filePath, + 'tool.input.hasCsvData': !!input.csvData, }, }) @@ -78,12 +85,11 @@ export const csvToJsonTool = createTool({ const maxRows = config?.maxRows // Check for cancellation before file operations - if (abortSignal && abortSignal.aborted) { - rootSpan.setStatus({ - code: 2, - message: 'Operation cancelled during file operations', + if (abortSignal?.aborted) { + csvSpan?.error({ + error: new Error('Operation cancelled during file operations'), + endSpan: true, }) - rootSpan.end() throw new Error( 'CSV to JSON conversion cancelled during file operations' ) @@ -139,10 +145,14 @@ export const csvToJsonTool = createTool({ type: 'progress', data: { message: `✅ Converted ${records.length} records` }, }) - rootSpan.setAttributes({ - 'tool.output.recordCount': records.length, + csvSpan?.update({ + output: { data: records }, + metadata: { + 'tool.output.recordCount': records.length, + 'tool.output.success': true, + }, }) - rootSpan.end() + csvSpan?.end() return { data: records } } catch (error) { const errorMessage = @@ -153,8 +163,10 @@ export const csvToJsonTool = createTool({ // Handle AbortError specifically if (error instanceof Error && error.name === 'AbortError') { const cancelMessage = `CSV to JSON conversion cancelled` - rootSpan.setStatus({ code: 2, message: cancelMessage }) - rootSpan.end() + csvSpan?.error({ + error: new Error(cancelMessage), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', @@ -170,11 +182,10 @@ export const csvToJsonTool = createTool({ throw new Error(cancelMessage) } - rootSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - rootSpan.setStatus({ code: 2, message: errorMessage }) - rootSpan.end() + csvSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error instanceof Error ? error : new Error(errorMessage) } }, @@ -196,10 +207,8 @@ export const csvToJsonTool = createTool({ }) }, onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - const source = input.filePath - ? `file:${input.filePath}` - : 'raw CSV data' - const options = input.options || {} + 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, diff --git a/src/mastra/tools/data-file-manager.ts b/src/mastra/tools/data-file-manager.ts index ce3e9cbe..bfbaecaa 100644 --- a/src/mastra/tools/data-file-manager.ts +++ b/src/mastra/tools/data-file-manager.ts @@ -13,10 +13,10 @@ // approvedBy: sam // approvalDate: 9/22 //import type { RequestContext } from '@mastra/core/request-context'; +import { SpanType } from "@mastra/core/observability"; import type { RequestContext } from "@mastra/core/request-context"; import type { InferUITool } from "@mastra/core/tools"; import { createTool } from "@mastra/core/tools"; -import { trace } from "@opentelemetry/api"; import * as fs from 'node:fs/promises'; import * as path from 'node:path'; import { pipeline } from 'node:stream/promises'; @@ -78,14 +78,17 @@ export const readDataFileTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; await ensureDataDir(); 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: { + const readSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'read:file', + input, + metadata: { 'tool.id': 'read:file', 'tool.input.fileName': input.fileName, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -101,20 +104,22 @@ export const readDataFileTool = createTool({ } const content = await fs.readFile(realFullPath, 'utf-8') log.info(`Read file: ${fileName}`) - readSpan.setAttributes({ - 'tool.output.fileSize': content.length + readSpan?.update({ + output: { fileSize: content.length }, + metadata: { + 'tool.output.fileSize': content.length + } }); - readSpan.end(); + readSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ File read successfully', stage: 'read:file' }, id: 'read:file' }); return content } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - readSpan.recordException(error); - } - readSpan.setStatus({ code: 2, message: errorMessage }); - readSpan.end(); + readSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -122,6 +127,7 @@ export const readDataFileTool = createTool({ log.info('Read file tool input streaming started', { toolCallId, messageCount: messages.length, + abortSignal: abortSignal?.aborted, hook: 'onInputStart', }) }, @@ -130,6 +136,7 @@ export const readDataFileTool = createTool({ toolCallId, messageCount: messages.length, fileName: input.fileName, + abortSignal: abortSignal?.aborted, hook: 'onInputAvailable', }) }, @@ -138,6 +145,7 @@ export const readDataFileTool = createTool({ toolCallId, toolName, fileSize: output.length, + abortSignal: abortSignal?.aborted, hook: 'onOutput', }) }, @@ -163,14 +171,16 @@ export const writeDataFileTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const writeSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'write:file', + input, + metadata: { 'tool.id': 'write:file', 'tool.input.fileName': input.fileName, - 'tool.input.contentLength': input.content.length, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -194,20 +204,22 @@ export const writeDataFileTool = createTool({ await fs.mkdir(realDirPath, { recursive: true }) await fs.writeFile(realFullPath, content, 'utf-8') log.info(`Written to file: ${fileName}`) - writeSpan.setAttributes({ - 'tool.output.success': true + writeSpan?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true + } }); - writeSpan.end(); + writeSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ File written successfully', stage: 'write:file' }, id: 'write:file' }); return `File ${fileName} written successfully.` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - writeSpan.recordException(error); - } - writeSpan.setStatus({ code: 2, message: errorMessage }); - writeSpan.end(); + writeSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -255,13 +267,16 @@ export const deleteDataFileTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const deleteSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'delete:file', + input, + metadata: { 'tool.id': 'delete:file', 'tool.input.fileName': input.fileName, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -276,20 +291,22 @@ export const deleteDataFileTool = createTool({ } await fs.unlink(fullPath) log.info(`Deleted file: ${fileName}`) - deleteSpan.setAttributes({ - 'tool.output.success': true + deleteSpan?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true + } }); - deleteSpan.end(); + deleteSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ File deleted successfully', stage: 'delete:file' }, id: 'delete:file' }); return `File ${fileName} deleted successfully.` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - deleteSpan.recordException(error); - } - deleteSpan.setStatus({ code: 2, message: errorMessage }); - deleteSpan.end(); + deleteSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -338,13 +355,16 @@ export const listDataDirTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const listSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'list:directory', + input, + metadata: { 'tool.id': 'list:directory', 'tool.input.dirPath': input.dirPath ?? 'docs/data', - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -359,20 +379,22 @@ export const listDataDirTool = createTool({ } const contents = await fs.readdir(fullPath) log.info(`Listed directory: ${dirPath}`) - listSpan.setAttributes({ - 'tool.output.count': contents.length + listSpan?.update({ + output: { count: contents.length }, + metadata: { + 'tool.output.count': contents.length + } }); - listSpan.end(); + listSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Directory listed successfully', stage: 'list:directory' }, id: 'list:directory' }); return contents } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - listSpan.recordException(error); - } - listSpan.setStatus({ code: 2, message: errorMessage }); - listSpan.end(); + listSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -424,14 +446,17 @@ export const copyDataFileTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const copySpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'copy:file', + input, + metadata: { 'tool.id': 'copy:file', 'tool.input.sourceFile': input.sourceFile, 'tool.input.destFile': input.destFile, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -452,20 +477,22 @@ export const copyDataFileTool = createTool({ await fs.mkdir(destDir, { recursive: true }) await fs.copyFile(sourcePath, destPath) log.info(`Copied file: ${sourceFile} to ${destFile}`) - copySpan.setAttributes({ - 'tool.output.success': true + copySpan?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true + } }); - copySpan.end(); + copySpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ File copied successfully', stage: 'copy:file' }, id: 'copy:file' }); return `File ${sourceFile} copied to ${destFile} successfully.` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - copySpan.recordException(error); - } - copySpan.setStatus({ code: 2, message: errorMessage }); - copySpan.end(); + copySpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -518,14 +545,17 @@ export const moveDataFileTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const moveSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'move:file', + input, + metadata: { 'tool.id': 'move:file', 'tool.input.sourceFile': input.sourceFile, 'tool.input.destFile': input.destFile, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -546,20 +576,22 @@ export const moveDataFileTool = createTool({ await fs.mkdir(destDir, { recursive: true }) await fs.rename(sourcePath, destPath) log.info(`Moved file: ${sourceFile} to ${destFile}`) - moveSpan.setAttributes({ - 'tool.output.success': true + moveSpan?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true + } }); - moveSpan.end(); + moveSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ File moved successfully', stage: 'move:file' }, id: 'move:file' }); return `File ${sourceFile} moved to ${destFile} successfully.` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - moveSpan.recordException(error); - } - moveSpan.setStatus({ code: 2, message: errorMessage }); - moveSpan.end(); + moveSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -618,15 +650,18 @@ export const searchDataFilesTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const searchSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'search:files', + input, + metadata: { 'tool.id': 'search:files', 'tool.input.pattern': input.pattern, 'tool.input.searchContent': input.searchContent, 'tool.input.dirPath': input.dirPath, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -685,20 +720,22 @@ export const searchDataFilesTool = createTool({ await searchDir(searchPath) log.info(`Searched for pattern: ${pattern} in ${dirPath}`) - searchSpan.setAttributes({ - 'tool.output.resultCount': results.length + searchSpan?.update({ + output: { resultCount: results.length }, + metadata: { + 'tool.output.resultCount': results.length + } }); - searchSpan.end(); + searchSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `✅ Found ${results.length} matches`, stage: 'search:files' }, id: 'search:files' }); return results } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - searchSpan.recordException(error); - } - searchSpan.setStatus({ code: 2, message: errorMessage }); - searchSpan.end(); + searchSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -706,6 +743,7 @@ export const searchDataFilesTool = createTool({ log.info('Search files tool input streaming started', { toolCallId, messageCount: messages.length, + abortSignal: abortSignal?.aborted, hook: 'onInputStart', }) }, @@ -716,6 +754,7 @@ export const searchDataFilesTool = createTool({ pattern: input.pattern, searchContent: input.searchContent, dirPath: input.dirPath, + abortSignal: abortSignal?.aborted, hook: 'onInputAvailable', }) }, @@ -724,6 +763,7 @@ export const searchDataFilesTool = createTool({ toolCallId, toolName, resultCount: output.length, + abortSignal: abortSignal?.aborted, hook: 'onOutput', }) }, @@ -754,13 +794,16 @@ export const getDataFileInfoTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const infoSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'get:fileinfo', + input, + metadata: { 'tool.id': 'get:fileinfo', 'tool.input.fileName': input.fileName, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -783,21 +826,23 @@ export const getDataFileInfoTool = createTool({ isFile: stats.isFile(), isDirectory: stats.isDirectory(), } - infoSpan.setAttributes({ - 'tool.output.fileSize': stats.size, - 'tool.output.isFile': stats.isFile(), + infoSpan?.update({ + output: result, + metadata: { + 'tool.output.fileSize': stats.size, + 'tool.output.isFile': stats.isFile(), + } }); - infoSpan.end(); + infoSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ File info retrieved', stage: 'get:fileinfo' }, id: 'get:fileinfo' }); return result } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - infoSpan.recordException(error); - } - infoSpan.setStatus({ code: 2, message: errorMessage }); - infoSpan.end(); + infoSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -844,13 +889,16 @@ export const createDataDirTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const createDirSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'create:directory', + input, + metadata: { 'tool.id': 'create:directory', 'tool.input.dirPath': input.dirPath, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -864,20 +912,22 @@ export const createDataDirTool = createTool({ } await fs.mkdir(fullPath, { recursive: true }) log.info(`Created directory: ${dirPath}`) - createDirSpan.setAttributes({ - 'tool.output.success': true + createDirSpan?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true + } }); - createDirSpan.end(); + createDirSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Directory created successfully', stage: 'create:directory' }, id: 'create:directory' }); return `Directory ${dirPath} created successfully.` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - createDirSpan.recordException(error); - } - createDirSpan.setStatus({ code: 2, message: errorMessage }); - createDirSpan.end(); + createDirSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -924,13 +974,16 @@ export const removeDataDirTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const removeDirSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'remove:directory', + input, + metadata: { 'tool.id': 'remove:directory', 'tool.input.dirPath': input.dirPath, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -949,20 +1002,22 @@ export const removeDataDirTool = createTool({ } await fs.rmdir(fullPath) log.info(`Removed directory: ${dirPath}`) - removeDirSpan.setAttributes({ - 'tool.output.success': true + removeDirSpan?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true + } }); - removeDirSpan.end(); + removeDirSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Directory removed successfully', stage: 'remove:directory' }, id: 'remove:directory' }); return `Directory ${dirPath} removed successfully.` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - removeDirSpan.recordException(error); - } - removeDirSpan.setStatus({ code: 2, message: errorMessage }); - removeDirSpan.end(); + removeDirSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -1015,14 +1070,17 @@ export const archiveDataTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const archiveSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'archive:data', + input, + metadata: { 'tool.id': 'archive:data', 'tool.input.sourcePath': input.sourcePath, 'tool.input.archiveName': input.archiveName, - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -1048,20 +1106,22 @@ export const archiveDataTool = createTool({ await pipeline(sourceStream, gzip, archiveStream) log.info(`Archived: ${sourcePath} to ${archiveName}.gz`) - archiveSpan.setAttributes({ - 'tool.output.success': true + archiveSpan?.update({ + output: { success: true }, + metadata: { + 'tool.output.success': true + } }); - archiveSpan.end(); + archiveSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Archive created successfully', stage: 'archive:data' }, id: 'archive:data' }); return `File ${sourcePath} archived to ${archiveName}.gz successfully.` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - archiveSpan.recordException(error); - } - archiveSpan.setStatus({ code: 2, message: errorMessage }); - archiveSpan.end(); + archiveSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, @@ -1116,14 +1176,17 @@ export const backupDataTool = createTool({ execute: async (input, context) => { const writer = context?.writer; const requestCtx = context?.requestContext as DataFileManagerContext | undefined; + const tracingContext = context?.tracingContext; 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: { + const backupSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'backup:data', + input, + metadata: { 'tool.id': 'backup:data', 'tool.input.sourcePath': input.sourcePath, 'tool.input.backupDir': input.backupDir ?? 'backups', - 'tool.requestContext.userId': requestCtx?.userId, + 'user.id': requestCtx?.userId, } }); @@ -1151,20 +1214,22 @@ export const backupDataTool = createTool({ await fs.cp(sourceFullPath, backupFullPath, { recursive: true }) const relativeBackupPath = path.relative(DATA_DIR, backupFullPath) log.info(`Backed up: ${sourcePath} to ${relativeBackupPath}`) - backupSpan.setAttributes({ - 'tool.output.backupPath': relativeBackupPath + backupSpan?.update({ + output: { backupPath: relativeBackupPath }, + metadata: { + 'tool.output.backupPath': relativeBackupPath + } }); - backupSpan.end(); + backupSpan?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Backup created successfully', stage: 'backup:data' }, id: 'backup:data' }); return `Backup created: ${sourcePath} → ${relativeBackupPath}` } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - backupSpan.recordException(error); - } - backupSpan.setStatus({ code: 2, message: errorMessage }); - backupSpan.end(); + backupSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw error } }, diff --git a/src/mastra/tools/data-processing-tools.ts b/src/mastra/tools/data-processing-tools.ts index 9217554f..067b40f4 100644 --- a/src/mastra/tools/data-processing-tools.ts +++ b/src/mastra/tools/data-processing-tools.ts @@ -1,7 +1,7 @@ import type { RequestContext } from "@mastra/core/request-context"; import type { InferUITool } from "@mastra/core/tools"; import { createTool } from "@mastra/core/tools"; -import { trace } from "@opentelemetry/api"; +import { SpanType } from "@mastra/core/observability"; import excalidrawToSvg from 'excalidraw-to-svg'; import { XMLBuilder, XMLParser } from 'fast-xml-parser'; import * as fs from 'node:fs/promises'; @@ -146,9 +146,12 @@ export const readCSVDataTool = createTool({ await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '📖 Reading CSV file: ' + inputData.fileName, stage: 'read:CSVdata' }, id: 'read:CSVdata' }); - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('read_csv_data', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'read_csv_data', + input: inputData, + metadata: { 'tool.id': 'read:CSVdata', 'tool.input.fileName': inputData.fileName, } @@ -194,15 +197,19 @@ export const readCSVDataTool = createTool({ headers = rows.length > 0 ? Object.keys(rows[0]) : [] } - span.setAttributes({ 'tool.output.rowCount': rows.length, 'tool.output.headerCount': headers.length }); - span.end(); + span?.update({ + output: { rowCount: rows.length, headerCount: headers.length }, + metadata: { 'tool.output.rowCount': rows.length, 'tool.output.headerCount': headers.length } + }); + span?.end(); await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `✅ Read ${rows.length} rows from CSV`, stage: 'read:CSVdata' }, id: 'read:CSVdata' }); return { headers, rows, rawCSV: content } } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: 2, message: errorMessage }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw error } }, @@ -260,9 +267,12 @@ export const csvToExcalidrawTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('csv-to-excalidraw', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'csv-to-excalidraw', + input: inputData, + metadata: { 'tool.id': 'csv-to-excalidraw', 'tool.input.layoutType': inputData.layoutType, 'tool.input.hasHeaders': inputData.hasHeaders, @@ -491,15 +501,19 @@ export const csvToExcalidrawTool = createTool({ }, }; - span.setAttributes({ 'tool.output.elementCount': elements.length, 'tool.output.filename': filename }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.elementCount': elements.length, 'tool.output.filename': filename } + }); + span?.end(); return result; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: 2, message: errorMessage }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw error; } }, @@ -559,9 +573,12 @@ export const imageToCSVTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('image-to-csv', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'image-to-csv', + input: inputData, + metadata: { 'tool.id': 'image-to-csv', 'tool.input.elementCount': inputData.elements.length, } @@ -600,8 +617,11 @@ export const imageToCSVTool = createTool({ columns, }; - span.setAttributes({ 'tool.output.elementCount': elements.length, 'tool.output.filename': outputFilename }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.elementCount': elements.length, 'tool.output.filename': outputFilename } + }); + span?.end(); return result; }, onInputStart: ({ toolCallId, messages, abortSignal }) => { @@ -650,9 +670,12 @@ export const validateExcalidrawTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('validate-excalidraw', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'validate-excalidraw', + input: inputData, + metadata: { 'tool.id': 'validate-excalidraw', 'tool.input.autoFix': inputData.autoFix, } @@ -756,14 +779,18 @@ export const validateExcalidrawTool = createTool({ elementCount: Array.isArray(fixedData?.elements) ? fixedData.elements.length : 0, }; - span.setAttributes({ 'tool.output.isValid': result.isValid, 'tool.output.errorCount': errors.length }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.isValid': result.isValid, 'tool.output.errorCount': errors.length } + }); + span?.end(); return result; } catch (error) { const errorMsg = String(error); - span.recordException(error instanceof Error ? error : new Error(errorMsg)); - span.setAttributes({ 'tool.output.isValid': false }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMsg), + endSpan: true, + }); return { isValid: false, errors: [`Validation failed: ${errorMsg}`], @@ -825,9 +852,12 @@ export const processSVGTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('process-svg', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'process-svg', + input: inputData, + metadata: { 'tool.id': 'process-svg', 'tool.input.extractPaths': inputData.extractPaths, 'tool.input.extractText': inputData.extractText, @@ -911,14 +941,18 @@ export const processSVGTool = createTool({ result.elementCount = svgJson.children ? svgJson.children.length : 0 await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `✅ Processed SVG with ${result.elementCount} elements`, stage: 'process-svg' }, id: 'process-svg' }); - span.setAttributes({ 'tool.output.elementCount': result.elementCount }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.elementCount': result.elementCount } + }); + span?.end(); return result } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: 2, message: errorMessage }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw new Error(`SVG processing failed: ${errorMessage}`) } }, @@ -973,9 +1007,12 @@ export const processXMLTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('process-xml', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'process-xml', + input: inputData, + metadata: { 'tool.id': 'process-xml', 'tool.input.extractElements': inputData.extractElements?.join(','), } @@ -1101,14 +1138,18 @@ export const processXMLTool = createTool({ rootElement, }; - span.setAttributes({ 'tool.output.elementCount': elements.length, 'tool.output.rootElement': result.rootElement }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.elementCount': elements.length, 'tool.output.rootElement': result.rootElement } + }); + span?.end(); return result; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: 2, message: errorMessage }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw new Error(`XML processing failed: ${errorMessage}`) } }, @@ -1162,9 +1203,12 @@ export const convertDataFormatTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('convert-data-format', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'convert-data-format', + input: inputData, + metadata: { 'tool.id': 'convert-data-format', 'tool.input.inputFormat': inputData.inputFormat, 'tool.input.outputFormat': inputData.outputFormat, @@ -1276,14 +1320,18 @@ export const convertDataFormatTool = createTool({ metadata, }; - span.setAttributes({ 'tool.output.success': true, 'tool.output.format': outputFormat }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.success': true, 'tool.output.format': outputFormat } + }); + span?.end(); return result; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: 2, message: errorMessage }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw new Error(`Data conversion failed: ${errorMessage}`) } }, @@ -1333,9 +1381,12 @@ export const validateDataTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('validate-data', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'validate-data', + input: inputData, + metadata: { 'tool.id': 'validate-data', 'tool.input.schemaType': inputData.schemaType, 'tool.input.strict': inputData.strict, @@ -1467,14 +1518,18 @@ export const validateDataTool = createTool({ validatedData, }; - span.setAttributes({ 'tool.output.isValid': isValid, 'tool.output.errorCount': errors.length }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.isValid': isValid, 'tool.output.errorCount': errors.length } + }); + span?.end(); return result; } catch (error) { const errorMsg = String(error); - span.recordException(error instanceof Error ? error : new Error(errorMsg)); - span.setAttributes({ 'tool.output.isValid': false }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMsg), + endSpan: true, + }); return { isValid: false, errors: [`Validation failed: ${errorMsg}`], @@ -1591,9 +1646,12 @@ export const excalidrawToSVGTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('excalidraw-to-svg', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'excalidraw-to-svg', + input: inputData, + metadata: { 'tool.id': 'excalidraw-to-svg', 'tool.input.elementCount': inputData.excalidrawData.elements.length, 'tool.input.saveToFile': inputData.saveToFile, @@ -1659,18 +1717,22 @@ export const excalidrawToSVGTool = createTool({ elementCount: inputData.excalidrawData.elements.length, }; - span.setAttributes({ - 'tool.output.elementCount': result.elementCount, - 'tool.output.width': dimensions.width, - 'tool.output.height': dimensions.height, + span?.update({ + output: result, + metadata: { + 'tool.output.elementCount': result.elementCount, + 'tool.output.width': dimensions.width, + 'tool.output.height': dimensions.height, + } }); - span.end(); + span?.end(); return result; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: 2, message: errorMessage }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw new Error(`Excalidraw to SVG conversion failed: ${errorMessage}`) } }, @@ -1719,9 +1781,12 @@ export const svgToExcalidrawTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('data-processing'); - const span = tracer.startSpan('svg-to-excalidraw', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'svg-to-excalidraw', + input: inputData, + metadata: { 'tool.id': 'svg-to-excalidraw', } }); @@ -1937,14 +2002,18 @@ export const svgToExcalidrawTool = createTool({ warnings: warnings.length > 0 ? warnings : undefined, }; - span.setAttributes({ 'tool.output.elementCount': elements.length }); - span.end(); + span?.update({ + output: result, + metadata: { 'tool.output.elementCount': elements.length } + }); + span?.end(); return result; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: 2, message: errorMessage }); - span.end(); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw new Error(`SVG to Excalidraw conversion failed: ${errorMessage}`) } }, diff --git a/src/mastra/tools/data-validator.tool.ts b/src/mastra/tools/data-validator.tool.ts index e9e6dc26..474d7c3b 100644 --- a/src/mastra/tools/data-validator.tool.ts +++ b/src/mastra/tools/data-validator.tool.ts @@ -1,7 +1,7 @@ +import { SpanType } from '@mastra/core/observability' import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' import { z } from 'zod' import { log } from '../config/logger' @@ -10,6 +10,7 @@ const validatorContextSchema = z.object({ }) export interface DataValidatorContext extends RequestContext { + userId?: string validatorContext?: { maxErrors?: number } @@ -172,6 +173,7 @@ export const dataValidatorToolJSON = createTool({ execute: async (inputData, context) => { const writer = context?.writer const requestContext = context?.requestContext as DataValidatorContext | undefined + const tracingContext = context?.tracingContext await writer?.custom({ type: 'data-tool-progress', @@ -183,12 +185,13 @@ export const dataValidatorToolJSON = createTool({ id: 'data-validator-json', }) - const tracer = trace.getTracer('data-validator') - const rootSpan = tracer.startSpan('data-validator', { - attributes: { + const rootSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'data-validator', + input: inputData, + metadata: { 'tool.id': 'data-validator-json', - 'tool.input.hasData': String(inputData.data !== undefined), - 'tool.input.hasSchema': String(inputData.schema !== undefined), + 'user.id': requestContext?.userId, }, }) @@ -208,8 +211,11 @@ export const dataValidatorToolJSON = createTool({ }, id: 'data-validator-json', }) - rootSpan.setAttributes({ 'tool.output.valid': true }) - rootSpan.end() + rootSpan?.update({ + output: { valid: true }, + metadata: { 'tool.output.valid': true } + }) + rootSpan?.end() return { valid: true, cleanedData: result.data, @@ -232,11 +238,14 @@ export const dataValidatorToolJSON = createTool({ ) } - rootSpan.setAttributes({ - 'tool.output.valid': false, - 'tool.output.errorCount': errors.length, + rootSpan?.update({ + output: { valid: false, errorCount: errors.length }, + metadata: { + 'tool.output.valid': false, + 'tool.output.errorCount': errors.length, + }, }) - rootSpan.end() + rootSpan?.end() return { valid: false, errors, @@ -247,11 +256,10 @@ export const dataValidatorToolJSON = createTool({ error instanceof Error ? error.message : 'Unknown validation error' - rootSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - rootSpan.setStatus({ code: 2, message: errorMessage }) - rootSpan.end() + rootSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) return { valid: false, errors: [errorMessage], diff --git a/src/mastra/tools/datetime.tool.ts b/src/mastra/tools/datetime.tool.ts index 0e94997c..22c83374 100644 --- a/src/mastra/tools/datetime.tool.ts +++ b/src/mastra/tools/datetime.tool.ts @@ -1,6 +1,6 @@ import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import { z } from 'zod' import { log } from '../config/logger' @@ -91,6 +91,7 @@ export const dateTimeTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer const requestCtx = context?.requestContext as DateTimeToolContext | undefined + const tracingContext = context?.tracingContext const defaultTimezone = requestCtx?.defaultTimezone ?? 'UTC' const defaultLocale = requestCtx?.defaultLocale ?? 'en-US' const allowFutureDates = requestCtx?.allowFutureDates ?? true @@ -102,9 +103,11 @@ export const dateTimeTool = createTool({ allowFutureDates, }) - const tracer = trace.getTracer('datetime-tool', '1.0.0') - const span = tracer.startSpan('datetime-operation', { - attributes: { + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'datetime-operation', + input: inputData, + metadata: { 'tool.id': 'datetime', 'tool.input.operation': inputData.operation, 'tool.input.input': inputData.input, @@ -264,11 +267,14 @@ export const dateTimeTool = createTool({ id: 'datetime', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.operation': inputData.operation, + span?.update({ + output: { success: true, operation: inputData.operation, result }, + metadata: { + 'tool.output.success': true, + 'tool.output.operation': inputData.operation, + }, }) - span.end() + span?.end() return { success: true, @@ -282,11 +288,10 @@ export const dateTimeTool = createTool({ error: errorMsg, }) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { success: false, @@ -365,10 +370,13 @@ export const timeZoneTool = createTool({ }), execute: async (inputData, context) => { const writer = context?.writer + const tracingContext = context?.tracingContext - const tracer = trace.getTracer('timezone-tool', '1.0.0') - const span = tracer.startSpan('timezone-operation', { - attributes: { + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'timezone-operation', + input: inputData, + metadata: { 'tool.id': 'timezone', 'tool.input.operation': inputData.operation, }, @@ -442,11 +450,14 @@ export const timeZoneTool = createTool({ id: 'timezone', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.operation': inputData.operation, + span?.update({ + output: { success: true, operation: inputData.operation, result }, + metadata: { + 'tool.output.success': true, + 'tool.output.operation': inputData.operation, + }, }) - span.end() + span?.end() return { success: true, @@ -459,11 +470,10 @@ export const timeZoneTool = createTool({ error: errorMsg, }) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { success: false, diff --git a/src/mastra/tools/diff-review.tool.ts b/src/mastra/tools/diff-review.tool.ts index ff5e7ec8..efc71920 100644 --- a/src/mastra/tools/diff-review.tool.ts +++ b/src/mastra/tools/diff-review.tool.ts @@ -1,5 +1,5 @@ import { createTool } from '@mastra/core/tools' -import { SpanStatusCode, trace } from '@opentelemetry/api' +import { SpanType } from "@mastra/core/observability"; import { createPatch, structuredPatch } from 'diff' import { z } from 'zod' import { log } from '../config/logger' @@ -102,10 +102,17 @@ Use for code review, comparing versions, and analyzing modifications.`, } = inputData const writer = context?.writer const requestContext = context?.requestContext as DiffReviewContext | undefined + const tracingContext = context?.tracingContext - const tracer = trace.getTracer('diff-review') - const span = tracer.startSpan('diff-review', { - attributes: { filename, contextLines }, + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'diff-review', + input: inputData, + metadata: { + 'tool.id': 'coding:diffReview', + filename, + contextLines, + }, }) await writer?.custom({ type: 'data-tool-progress', @@ -230,9 +237,14 @@ Use for code review, comparing versions, and analyzing modifications.`, summary = `${totalChanges} change${totalChanges !== 1 ? 's' : ''}: ${parts.join(', ')} across ${hunks.length} hunk${hunks.length !== 1 ? 's' : ''}.` } - span?.setAttribute('additions', additions) - span?.setAttribute('deletions', deletions) - span?.setAttribute('totalChanges', totalChanges) + span?.update({ + output: { additions, deletions, totalChanges }, + metadata: { + additions, + deletions, + totalChanges, + } + }) await writer?.custom({ type: 'data-tool-progress', data: { @@ -258,14 +270,10 @@ Use for code review, comparing versions, and analyzing modifications.`, } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - span?.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - span?.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - span?.end() throw error } }, diff --git a/src/mastra/tools/document-chunking.tool.ts b/src/mastra/tools/document-chunking.tool.ts index e4e84b98..e947b1d1 100644 --- a/src/mastra/tools/document-chunking.tool.ts +++ b/src/mastra/tools/document-chunking.tool.ts @@ -1,10 +1,10 @@ +import { SpanType } from '@mastra/core/observability'; import type { InferUITool } from "@mastra/core/tools"; import { createTool } from "@mastra/core/tools"; import { MDocument, rerank, } from '@mastra/rag'; -import { SpanStatusCode, trace } from "@opentelemetry/api"; import { randomUUID } from 'crypto'; import { ModelRouterEmbeddingModel, ModelRouterLanguageModel } from "@mastra/core/llm"; @@ -221,11 +221,11 @@ Use this tool when you need advanced document processing with metadata extractio const startTime = Date.now() logToolExecution('mastra-chunker', { input: inputData }) - // Create a span for tracing - // Create a span for tracing - const tracer = trace.getTracer('document-chunking'); - const span = tracer.startSpan('mastra-chunker-tool', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'mastra-chunker', + input: { documentLength: inputData.documentContent.length, chunkingStrategy: inputData.chunkingStrategy, chunkSize: inputData.chunkSize, @@ -233,8 +233,11 @@ Use this tool when you need advanced document processing with metadata extractio extractSummary: inputData.extractSummary, extractKeywords: inputData.extractKeywords, extractQuestions: inputData.extractQuestions, - operation: 'mastra-chunker' - } + }, + metadata: { + 'tool.id': 'mastra-chunker', + operation: 'mastra-chunker', + }, }); try { @@ -416,12 +419,14 @@ Use this tool when you need advanced document processing with metadata extractio }) // Record error in tracing span - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw error instanceof Error ? error : new Error(errorMessage); } finally { - span.end(); + span?.end(); } }, }) @@ -495,18 +500,20 @@ content indexing, or semantic search capabilities. const startTime = Date.now() logToolExecution('mdocument-chunker', { input: inputData }) - // Create a span for tracing - // Create a span for tracing - const tracer = trace.getTracer('document-chunking'); - const span = tracer.startSpan('mdocument-chunker-tool', { - attributes: { + const tracingContext = context?.tracingContext; + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'mdocument-chunker', + input: { documentLength: inputData.documentContent.length, chunkingStrategy: inputData.chunkingStrategy, chunkSize: inputData.chunkSize, generateEmbeddings: inputData.generateEmbeddings, - indexName: 'memory_messages_3072', - operation: 'mdocument-chunker' - } + }, + metadata: { + 'tool.id': 'mdocument-chunker', + operation: 'mdocument-chunker', + }, }); try { @@ -769,12 +776,14 @@ content indexing, or semantic search capabilities. }) // Record error in tracing span - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw error instanceof Error ? error : new Error(errorMessage); } finally { - span.end(); + span?.end(); } }, }) @@ -865,15 +874,21 @@ Use this tool to improve retrieval quality by re-ranking initial search results. const startTime = Date.now() logToolExecution('document-reranker', { userQuery: inputData.userQuery }) - const tracer = trace.getTracer('document-chunking'); - const span = tracer.startSpan('document-reranker-tool', { - attributes: { + // Use the existing tracing context if available to create a child span. + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'document-reranker', + input: { userQuery: inputData.userQuery, indexName: inputData.indexName, topK: inputData.topK, initialTopK: inputData.initialTopK, - operation: 'document-reranker' - } + }, + metadata: { + 'tool.id': 'document:reranker', + operation: 'document-reranker', + }, }); try { @@ -1003,12 +1018,14 @@ Use this tool to improve retrieval quality by re-ranking initial search results. processingTimeMs: processingTime, }) - span.recordException(error instanceof Error ? error : new Error(errorMessage)); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMessage }); + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }); throw error instanceof Error ? error : new Error(errorMessage); } finally { - span.end(); + span?.end(); } }, }) diff --git a/src/mastra/tools/e2b.ts b/src/mastra/tools/e2b.ts index debb93ac..89308130 100644 --- a/src/mastra/tools/e2b.ts +++ b/src/mastra/tools/e2b.ts @@ -1,1651 +1,1548 @@ import { FilesystemEventType, FileType, Sandbox } from '@e2b/code-interpreter' import { createTool } from '@mastra/core/tools' -import { trace, SpanStatusCode, context as otelContext, propagation, metrics } from '@opentelemetry/api' -import type { BaggageEntry, BaggageEntryMetadata, Baggage } from '@opentelemetry/api' -import z from 'zod' -import { log, logError, logToolExecution } from '../config/logger' -import type { RequestContext } from '@mastra/core/request-context' -import type { InferUITool, InferUITools } from "@mastra/core/tools"; - - -export interface E2BRequestContext extends RequestContext { - userId?: string - sandboxLimits?: { - maxTimeout: number - maxFileSize: number - maxExecutionTime: number - allowedLanguages?: Array<'python' | 'javascript' | 'typescript'> - } -} -const tracer = trace.getTracer('e2b-tools', '1.0.0') +import { SpanType } from '@mastra/core/observability' -// Initialize metrics -const meter = metrics.getMeter('e2b-tools', '1.0.0') -const toolExecutionCounter = meter.createCounter('e2b_tool_executions_total', { - description: 'Total number of E2B tool executions' -}) -const toolExecutionDuration = meter.createHistogram('e2b_tool_execution_duration', { - description: 'Duration of E2B tool executions in milliseconds', - unit: 'ms' -}) -// Counter and histogram for sandbox-level operations (create/delete) -const sandboxOperationsCounter = meter.createCounter('e2b_sandbox_operations_total', { - description: 'Total number of sandbox operations (create/delete)' -}) -const sandboxOperationDuration = meter.createHistogram('e2b_sandbox_operation_duration', { - description: 'Duration of sandbox operations in milliseconds', - unit: 'ms' -}) +import type { RequestContext } from '@mastra/core/request-context' +import type { InferUITool, InferUITools } from "@mastra/core/tools" +import z from 'zod' +import { log, logError } from '../config/logger' -// Enhanced logging utilities with centralized logger -// Helper: return active span context info (traceId, spanId) -const getActiveSpanContextInfo = () => { - const span = trace.getActiveSpan() - if (!span) {return { traceId: undefined, spanId: undefined }} - const sc = span.spanContext() - return { traceId: sc.traceId, spanId: sc.spanId } -} -// Helper: convert Baggage to simple object for logging -const BaggageToObject = (b: Baggage | null | undefined) => { - if (b === null || b === undefined) { - return {} - } - const obj: Record<string, string> = {} - for (const [k, v] of b.getAllEntries()) { - obj[k] = String(v.value) - } - return obj +export interface E2BRequestContext extends RequestContext { + userId?: string + sandboxLimits?: { + maxTimeout: number + maxFileSize: number + maxExecutionTime: number + allowedLanguages?: Array<'python' | 'javascript' | 'typescript'> + } } -// Helper: merge provided entries into the current Baggage and set it on the active context -const mergeAndSetBaggage = (entries: Record<string, string>) => { - try { - const ctx = otelContext.active() - const currentBaggage = propagation.getBaggage(ctx) - const merged: Record<string, BaggageEntry> = {} - if (currentBaggage) { - for (const entry of currentBaggage.getAllEntries()) { - const [k, v] = entry - merged[k] = { value: String(v.value), metadata: v.metadata } - } - } - for (const k of Object.keys(entries)) { - merged[k] = { value: entries[k] } - } - const bag = propagation.createBaggage(merged) - const newCtx = propagation.setBaggage(ctx, bag) - // Make the new baggage active for the current synchronous call and its children - otelContext.with(newCtx, () => {}) - return bag - } catch (err) { - // Non-fatal - don't let baggage failures break execution - log.error('Failed to set baggage', { error: err }) - return undefined - } -} +/* tracing: use runtime tracingContext inside tools; tracer removed */ const logToolStart = (toolName: string, input: any, toolContext?: any) => { - const span = trace.getActiveSpan() - if (span) { - // Add span event for tool start - span.addEvent('tool.start', { - 'tool.name': toolName, - 'tool.input_size': JSON.stringify(input).length, - 'user.id': toolContext?.requestContext?.userId, - }) - - // Merge into baggage for downstream services and traces - mergeAndSetBaggage({ 'tool.name': toolName, 'tool.start_time': Date.now().toString() }) - - // Also set attributes on the span so it's available in OTEL exporters - try { - span.setAttribute('tool.name', toolName) - span.setAttribute('tool.input_size', JSON.stringify(input).length) - } catch (err) { - // ignore if span does not support setAttribute - } + const runtimeSpan = toolContext?.tracingContext?.currentSpan + if (runtimeSpan) { + const raw = runtimeSpan.raw + if (raw && typeof raw.addEvent === 'function') { + raw.addEvent('tool.start', { + 'tool.name': toolName, + 'tool.input_size': JSON.stringify(input).length, + 'user.id': toolContext?.requestContext?.userId, + }) + } else if (typeof runtimeSpan.update === 'function') { + runtimeSpan.update({ metadata: { 'tool.name': toolName, 'tool.input_size': JSON.stringify(input).length } }) } - - // Include trace info and baggage in the structured log - const traceInfo = getActiveSpanContextInfo() - const bag = propagation.getBaggage(otelContext.active()) - logToolExecution(toolName, { originalInput: input, _trace: traceInfo, _Baggage: BaggageToObject(bag), _userId: toolContext?.requestContext?.userId }) + } } - const logToolComplete = ( - toolName: string, - output: any, - duration: number, - toolContext?: any + toolName: string, + output: any, + duration: number, + toolContext?: any ) => { - const span = trace.getActiveSpan() - if (span) { - // Add span event for tool completion - span.addEvent('tool.complete', { - 'tool.name': toolName, - 'tool.duration_ms': duration, - 'tool.output_size': JSON.stringify(output).length, - 'user.id': toolContext?.requestContext?.userId, - }) - - // Update baggage - mergeAndSetBaggage({ 'tool.duration': duration.toString(), 'tool.completed_at': Date.now().toString() }) - - // Set attributes on span - try { - span.setAttribute('tool.duration_ms', duration) - span.setAttribute('tool.output_size', JSON.stringify(output).length) - } catch (err) { - // ignore - } - } - - // Record metrics - toolExecutionCounter.add(1, { + const runtimeSpan = toolContext?.tracingContext?.currentSpan + if (runtimeSpan) { + const raw = runtimeSpan.raw + if (raw && typeof raw.addEvent === 'function') { + raw.addEvent('tool.complete', { 'tool.name': toolName, - 'tool.status': 'success', + 'tool.duration_ms': duration, + 'tool.output_size': JSON.stringify(output).length, 'user.id': toolContext?.requestContext?.userId, - }) - toolExecutionDuration.record(duration, { - 'tool.name': toolName, - 'tool.status': 'success', - 'user.id': toolContext?.requestContext?.userId, - }) - - const traceInfo = getActiveSpanContextInfo() - const bag = propagation.getBaggage(otelContext.active()) - - logToolExecution( - toolName, - { _trace: traceInfo, _Baggage: BaggageToObject(bag), _userId: toolContext?.requestContext?.userId }, - output - ) + }) + } else if (typeof runtimeSpan.update === 'function') { + runtimeSpan.update({ metadata: { 'tool.duration_ms': duration, 'tool.output_size': JSON.stringify(output).length } }) + } + } } -const logToolError = (toolName: string, error: any, context?: any) => { - const span = trace.getActiveSpan() - if (span) { - // Add span event for tool error - span.addEvent('tool.error', { - 'tool.name': toolName, - 'error.message': error instanceof Error ? error.message : String(error), - 'error.stack': error instanceof Error ? error.stack : undefined, - 'user.id': context?.requestContext?.userId - }) - // Record error in span - span.recordException(error instanceof Error ? error : new Error(String(error))) +const logToolError = (toolName: string, error: any, context?: any) => { + const runtimeSpan = context?.tracingContext?.currentSpan + if (runtimeSpan) { + const raw = runtimeSpan.raw + if (raw && typeof raw.addEvent === 'function') { + raw.addEvent('tool.error', { + 'tool.name': toolName, + 'error.message': error instanceof Error ? error.message : String(error), + 'error.stack': error instanceof Error ? error.stack : undefined, + 'user.id': context?.requestContext?.userId + }) + } else if (typeof runtimeSpan.update === 'function') { + runtimeSpan.update({ metadata: { 'error.message': error instanceof Error ? error.message : String(error) } }) } - // Record error metrics - toolExecutionCounter.add(1, { - 'tool.name': toolName, - 'tool.status': 'error', - 'error.type': error instanceof Error ? error.constructor.name : 'Unknown' - }) + runtimeSpan?.error({ error: error instanceof Error ? error : new Error(String(error)), endSpan: false }) + } - logError(toolName, error, { - toolName, - userId: context?.requestContext?.userId - }) + logError(toolName, error, { + toolName, + userId: context?.requestContext?.userId + }) } export const createSandbox = createTool({ - id: 'createSandbox', - description: 'Create an e2b sandbox', - inputSchema: z.object({ - metadata: z - .record(z.string(), z.string()) - .optional() - .describe('Custom metadata for the sandbox'), - envs: z.record(z.string(), z.string()).optional().describe(` + id: 'createSandbox', + description: 'Create an e2b sandbox', + inputSchema: z.object({ + metadata: z + .record(z.string(), z.string()) + .optional() + .describe('Custom metadata for the sandbox'), + envs: z.record(z.string(), z.string()).optional().describe(` Custom environment variables for the sandbox. - Used when executing commands and code in the sandbox. + Used when executing commands and code in the sandbox, following the tracing implementation patterns described in AGENTS.md. Can be overridden with the \`envs\` argument when executing commands or code. `), - timeoutMS: z.number().optional().describe(` + timeoutMS: z.number().optional().describe(` Timeout for the sandbox in **milliseconds**. Maximum time a sandbox can be kept alive is 24 hours (86_400_000 milliseconds) for Pro users and 1 hour (3_600_000 milliseconds) for Hobby users. @default 300_000 // 5 minutes `), - }), - outputSchema: z.object({ - sandboxId: z.string(), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('createSandbox') - const startTime = Date.now() - const requestContext = context?.requestContext as E2BRequestContext - const userId = requestContext?.userId - if (userId && userId.length > 0) { - span.setAttribute('user.id', userId) - } + }), + outputSchema: z.object({ + sandboxId: z.string(), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'createSandbox', + input: inputData, + metadata: { 'tool.id': 'createSandbox' }, + }) + const startTime = Date.now() + const requestContext = context?.requestContext as E2BRequestContext + const userId = requestContext?.userId + if (userId && userId.length > 0) { + span?.update({ metadata: { 'user.id': userId } }) + } - // Add custom resource attributes - span.setAttribute('service.name', 'e2b-tools') - span.setAttribute('service.version', '1.0.0') - span.setAttribute('sandbox.operation', 'create') + // Add custom resource attributes + span?.update({ metadata: { 'service.name': 'e2b-tools', 'service.version': '1.0.0', 'sandbox.operation': 'create' } }) - // Set baggage for cross-service context propagation - const currentBaggage = propagation.getBaggage(otelContext.active()) - if (currentBaggage) { - const keys = Array.from(currentBaggage.getAllEntries()).map(([key]) => key) - span.setAttribute('baggage.keys', keys.join(',')) - } + // Set baggage for cross-service context propagation - // Add span event for input validation - span.addEvent('input.validation', { - 'input.has_metadata': !!inputData.metadata, - 'input.has_envs': !!inputData.envs, - 'input.timeout_ms': inputData.timeoutMS ?? 300000 - }) - try { - // Add start event and input size - const inputSize = JSON.stringify(inputData).length - span.addEvent('tool.start', { - ts: Date.now(), - input_size: inputSize, - }) - span.setAttribute('tool.input_size', inputSize) - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `Creating sandbox...`, - stage: 'createSandbox', - }, - id: 'createSandbox', - }) - - logToolStart('createSandbox', inputData, context) - - const sandbox = await Sandbox.create(inputData) - - // Record successful sandbox creation metrics - const creationDuration = Date.now() - startTime - sandboxOperationsCounter.add(1, { - 'operation.type': 'create', - 'operation.status': 'success', - 'sandbox.id': sandbox.sandboxId, - }) - sandboxOperationDuration.record(creationDuration, { - 'operation.type': 'create', - 'operation.status': 'success', - 'sandbox.id': sandbox.sandboxId, - }) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', sandbox.sandboxId) - span.setAttribute('sandbox.created', true) - - // Add span event for successful creation - span.addEvent('sandbox.created', { - 'sandbox.id': sandbox.sandboxId, - 'creation.time_ms': Date.now() - startTime - }) - - const result = { - sandboxId: sandbox.sandboxId, - } - - // Output size attribute & complete event - const outputSize = JSON.stringify(result).length - span.setAttribute('tool.output_size', outputSize) - span.addEvent('tool.complete', { - duration_ms: Date.now() - startTime, - }) - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `Sandbox created: ${sandbox.sandboxId}`, - stage: 'createSandbox', - }, - id: 'createSandbox', - }) - - logToolComplete( - 'createSandbox', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - // Record error metrics - sandboxOperationsCounter.add(1, { - 'operation.type': 'create', - 'operation.status': 'error', - 'error.type': e instanceof Error ? e.constructor.name : 'Unknown' - }) - - if (e instanceof Error) { - span.recordException(e) - span.addEvent('sandbox.creation.failed', { - 'error.message': e.message, - 'error.stack': e.stack?.slice(0, 500) - }) - } else { - span.recordException(new Error(String(e))) - } - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('createSandbox', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Create sandbox tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Create sandbox received complete input', { - toolCallId, - messageCount: messages.length, - hasMetadata: !!input.metadata, - hasEnvs: !!input.envs, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Create sandbox completed', { - toolCallId, - toolName, - sandboxId: output.sandboxId, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `Creating sandbox...`, + stage: 'createSandbox', + }, + id: 'createSandbox', + }) + + try { + logToolStart('createSandbox', inputData, context) + + const sandbox = await Sandbox.create(inputData) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': sandbox.sandboxId, 'sandbox.created': true } }) + + const result = { + sandboxId: sandbox.sandboxId, + } + + // Output size attribute & complete event + const outputSize = JSON.stringify(result).length + span?.update({ metadata: { 'tool.output_size': outputSize, 'tool.duration_ms': Date.now() - startTime } }) + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `Sandbox created: ${sandbox.sandboxId}`, + stage: 'createSandbox', + }, + id: 'createSandbox', + }) + + logToolComplete( + 'createSandbox', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: false }) + logToolError('createSandbox', e, context) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Create sandbox tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Create sandbox received complete input', { + toolCallId, + messageCount: messages.length, + hasMetadata: !!input.metadata, + hasEnvs: !!input.envs, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Create sandbox completed', { + toolCallId, + toolName, + sandboxId: output.sandboxId, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const writeFile = createTool({ - id: 'writeFile', - description: 'Write a file to the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe('The sandboxId for the sandbox to write the file to'), - path: z.string().describe('The path where the file should be written'), - content: z.string().describe('The content to write to the file'), - }), - outputSchema: z.object({ - success: z - .boolean() - .describe('Whether the file was written successfully'), - path: z.string().describe('The path where the file was written'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('writeFile') - const startTime = Date.now() - - try { - logToolStart('writeFile', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - await sandbox.files.write(inputData.path, inputData.content) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('file.path', inputData.path) - span.setAttribute('file.size', inputData.content.length) - - const result = { - success: true, - path: inputData.path, - } - - logToolComplete( - 'writeFile', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('writeFile', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Write file tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Write file received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - contentLength: input.content.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Write file completed', { - toolCallId, - toolName, - success: output.success, - path: output.path, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + id: 'writeFile', + description: 'Write a file to the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to write the file to'), + path: z.string().describe('The path where the file should be written'), + content: z.string().describe('The content to write to the file'), + }), + outputSchema: z.object({ + success: z + .boolean() + .describe('Whether the file was written successfully'), + path: z.string().describe('The path where the file was written'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'writeFile', + input: inputData, + metadata: { + 'tool.id': 'writeFile', + 'tool.input.sandboxId': inputData.sandboxId, + 'tool.input.path': inputData.path, + }, + }) + const startTime = Date.now() + + try { + logToolStart('writeFile', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + await sandbox.files.write(inputData.path, inputData.content) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'file.path': inputData.path, 'file.size': inputData.content.length } }) + span?.end() + + const result = { + success: true, + path: inputData.path, + } + + logToolComplete( + 'writeFile', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: true }) + const requestContext = context?.requestContext as E2BRequestContext + logError('writeFile', e, { + toolName: 'writeFile', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Write file tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Write file received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + contentLength: input.content.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Write file completed', { + toolCallId, + toolName, + success: output.success, + path: output.path, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const writeFiles = createTool({ - id: 'writeFiles', - description: 'Write multiple files to the e2b sandbox', - inputSchema: z.object({ - sandboxId: z + id: 'writeFiles', + description: 'Write multiple files to the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to write files to'), + files: z + .array( + z.object({ + path: z .string() - .describe('The sandboxId for the sandbox to write files to'), - files: z - .array( - z.object({ - path: z - .string() - .describe('The path where the file should be written'), - data: z - .string() - .describe('The content to write to the file'), - }) - ) - .describe('Array of files to write, each with path and data'), - }), - outputSchema: z.object({ - success: z - .boolean() - .describe('Whether all files were written successfully'), - filesWritten: z - .array(z.string()) - .describe('Array of file paths that were written'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('writeFiles') - const startTime = Date.now() - - try { - logToolStart('writeFiles', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - await sandbox.files.write(inputData.files) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('files.count', inputData.files.length) - span.setAttribute( - 'files.total_size', - inputData.files.reduce((sum, file) => sum + file.data.length, 0) - ) - - const result = { - success: true, - filesWritten: inputData.files.map((file) => file.path), - } - - logToolComplete( - 'writeFiles', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('writeFiles', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Write files tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Write files received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - fileCount: input.files.length, - totalContentSize: input.files.reduce( - (sum, file) => sum + file.data.length, - 0 - ), - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Write files completed', { - toolCallId, - toolName, - success: output.success, - filesWrittenCount: output.filesWritten.length, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', + .describe('The path where the file should be written'), + data: z + .string() + .describe('The content to write to the file'), }) - }, + ) + .describe('Array of files to write, each with path and data'), + }), + outputSchema: z.object({ + success: z + .boolean() + .describe('Whether all files were written successfully'), + filesWritten: z + .array(z.string()) + .describe('Array of file paths that were written'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'writeFiles', + input: inputData, + metadata: { + 'tool.id': 'writeFiles', + 'tool.input.sandboxId': inputData.sandboxId, + 'tool.input.fileCount': Array.isArray(inputData.files) ? inputData.files.length : undefined, + }, + }) + const startTime = Date.now() + + try { + logToolStart('writeFiles', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + await sandbox.files.write(inputData.files) + + const totalSize = inputData.files.reduce((sum, file) => sum + file.data.length, 0) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'files.count': inputData.files.length, 'files.total_size': totalSize } }) + span?.end() + + const result = { + success: true, + filesWritten: inputData.files.map((file) => file.path), + } + + span?.update({ metadata: { 'tool.output_size': JSON.stringify(result).length, 'tool.duration_ms': Date.now() - startTime } }) + + logToolComplete( + 'writeFiles', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: true }) + const requestContext = context?.requestContext as E2BRequestContext + logError('writeFiles', e, { + toolName: 'writeFiles', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Write files tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Write files received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + fileCount: input.files.length, + totalContentSize: input.files.reduce( + (sum, file) => sum + file.data.length, + 0 + ), + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Write files completed', { + toolCallId, + toolName, + success: output.success, + filesWrittenCount: output.filesWritten.length, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const listFiles = createTool({ - id: 'listFiles', - description: 'List files and directories in a path within the e2b sandbox', - inputSchema: z.object({ - sandboxId: z + id: 'listFiles', + description: 'List files and directories in a path within the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to list files from'), + path: z + .string() + .default('/') + .describe('The directory path to list files from'), + }), + outputSchema: z.object({ + files: z + .array( + z.object({ + name: z .string() - .describe('The sandboxId for the sandbox to list files from'), - path: z + .describe('The name of the file or directory'), + path: z .string() - .default('/') - .describe('The directory path to list files from'), - }), - outputSchema: z.object({ - files: z - .array( - z.object({ - name: z - .string() - .describe('The name of the file or directory'), - path: z - .string() - .describe('The full path of the file or directory'), - isDirectory: z - .boolean() - .describe('Whether this is a directory'), - }) - ) - .describe('Array of files and directories'), - path: z.string().describe('The path that was listed'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('listFiles') - const startTime = Date.now() - - try { - logToolStart('listFiles', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - const fileList = await sandbox.files.list(inputData.path) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('list.path', inputData.path) - span.setAttribute('list.file_count', fileList.length) - span.setAttribute( - 'list.directory_count', - fileList.filter((f) => f.type === FileType.DIR).length - ) - - const result = { - files: fileList.map((file) => ({ - name: file.name, - path: file.path, - isDirectory: file.type === FileType.DIR, - })), - path: inputData.path, - } - - logToolComplete( - 'listFiles', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('listFiles', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('List files tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('List files received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('List files completed', { - toolCallId, - toolName, - fileCount: output.files.length, - directoryCount: output.files.filter((f) => f.isDirectory).length, - path: output.path, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', + .describe('The full path of the file or directory'), + isDirectory: z + .boolean() + .describe('Whether this is a directory'), }) - }, + ) + .describe('Array of files and directories'), + path: z.string().describe('The path that was listed'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'listFiles', + input: inputData, + metadata: { + 'tool.id': 'listFiles', + 'tool.input.sandboxId': inputData.sandboxId, + 'tool.input.path': inputData.path, + }, + }) + const startTime = Date.now() + + try { + logToolStart('listFiles', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + const fileList = await sandbox.files.list(inputData.path) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'list.path': inputData.path, 'list.file_count': fileList.length, 'list.directory_count': fileList.filter((f) => f.type === FileType.DIR).length } }) + span?.end() + + const result = { + files: fileList.map((file) => ({ + name: file.name, + path: file.path, + isDirectory: file.type === FileType.DIR, + })), + path: inputData.path, + } + + logToolComplete( + 'listFiles', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: true }) + const requestContext = context?.requestContext as E2BRequestContext + logError('listFiles', e, { + toolName: 'listFiles', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('List files tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('List files received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('List files completed', { + toolCallId, + toolName, + fileCount: output.files.length, + directoryCount: output.files.filter((f) => f.isDirectory).length, + path: output.path, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const deleteFile = createTool({ - id: 'deleteFile', - description: 'Delete a file or directory from the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe('The sandboxId for the sandbox to delete the file from'), - path: z - .string() - .describe('The path to the file or directory to delete'), - }), - outputSchema: z.object({ - success: z - .boolean() - .describe('Whether the file was deleted successfully'), - path: z.string().describe('The path that was deleted'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('deleteFile') - const startTime = Date.now() - - try { - logToolStart('deleteFile', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - await sandbox.files.remove(inputData.path) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('file.path', inputData.path) - - const result = { - success: true, - path: inputData.path, - } - - logToolComplete( - 'deleteFile', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('deleteFile', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Delete file tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Delete file received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Delete file completed', { - toolCallId, - toolName, - success: output.success, - path: output.path, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + id: 'deleteFile', + description: 'Delete a file or directory from the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to delete the file from'), + path: z + .string() + .describe('The path to the file or directory to delete'), + }), + outputSchema: z.object({ + success: z + .boolean() + .describe('Whether the file was deleted successfully'), + path: z.string().describe('The path that was deleted'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'deleteFile', + input: inputData, + metadata: { 'tool.id': 'deleteFile', 'tool.input.sandboxId': inputData.sandboxId, 'tool.input.path': inputData.path }, + }) + const startTime = Date.now() + + try { + logToolStart('deleteFile', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + await sandbox.files.remove(inputData.path) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'file.path': inputData.path } }) + span?.end() + + const result = { + success: true, + path: inputData.path, + } + + logToolComplete( + 'deleteFile', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: true }) + const requestContext = context?.requestContext as E2BRequestContext + logError('deleteFile', e, { + toolName: 'deleteFile', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Delete file tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Delete file received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Delete file completed', { + toolCallId, + toolName, + success: output.success, + path: output.path, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const createDirectory = createTool({ - id: 'createDirectory', - description: 'Create a directory in the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe( - 'The sandboxId for the sandbox to create the directory in' - ), - path: z - .string() - .describe('The path where the directory should be created'), - }), - outputSchema: z.object({ - success: z - .boolean() - .describe('Whether the directory was created successfully'), - path: z.string().describe('The path where the directory was created'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('createDirectory') - const startTime = Date.now() - - try { - logToolStart('createDirectory', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - await sandbox.files.makeDir(inputData.path) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('directory.path', inputData.path) - - const result = { - success: true, - path: inputData.path, - } - - logToolComplete( - 'createDirectory', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('createDirectory', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Create directory tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Create directory received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Create directory completed', { - toolCallId, - toolName, - success: output.success, - path: output.path, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + id: 'createDirectory', + description: 'Create a directory in the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe( + 'The sandboxId for the sandbox to create the directory in' + ), + path: z + .string() + .describe('The path where the directory should be created'), + }), + outputSchema: z.object({ + success: z + .boolean() + .describe('Whether the directory was created successfully'), + path: z.string().describe('The path where the directory was created'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'createDirectory', + input: inputData, + metadata: { 'tool.id': 'createDirectory', 'tool.input.sandboxId': inputData.sandboxId, 'tool.input.path': inputData.path }, + }) + const startTime = Date.now() + + try { + logToolStart('createDirectory', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + await sandbox.files.makeDir(inputData.path) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'directory.path': inputData.path } }) + + const result = { + success: true, + path: inputData.path, + } + + logToolComplete( + 'createDirectory', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: true }) + const requestContext = context?.requestContext as E2BRequestContext + logError('createDirectory', e, { + toolName: 'createDirectory', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Create directory tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Create directory received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Create directory completed', { + toolCallId, + toolName, + success: output.success, + path: output.path, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const getFileInfo = createTool({ - id: 'getFileInfo', - description: - 'Get detailed information about a file or directory in the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe( - 'The sandboxId for the sandbox to get file information from' - ), - path: z - .string() - .describe( - 'The path to the file or directory to get information about' - ), - }), - outputSchema: z.object({ - name: z.string().describe('The name of the file or directory'), - type: z - .enum(FileType) - .optional() - .describe('Whether this is a file or directory'), - path: z.string().describe('The full path of the file or directory'), - size: z.number().describe('The size of the file or directory in bytes'), - mode: z - .number() - .describe('The file mode (permissions as octal number)'), - permissions: z.string().describe('Human-readable permissions string'), - owner: z.string().describe('The owner of the file or directory'), - group: z.string().describe('The group of the file or directory'), - modifiedTime: z - .date() - .optional() - .describe('The last modified time in ISO string format'), - symlinkTarget: z - .string() - .optional() - .describe('The target path if this is a symlink, null otherwise'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('getFileInfo') - const startTime = Date.now() - - try { - logToolStart('getFileInfo', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - const info = await sandbox.files.getInfo(inputData.path) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('file.path', inputData.path) - span.setAttribute('file.size', info.size) - if (info.type !== null && info.type !== undefined) { - span.setAttribute('file.type', info.type) - } - span.setAttribute('file.permissions', info.permissions) - - const result = { - name: info.name, - type: info.type, - path: info.path, - size: info.size, - mode: info.mode, - permissions: info.permissions, - owner: info.owner, - group: info.group, - modifiedTime: info.modifiedTime, - symlinkTarget: info.symlinkTarget, - } - - logToolComplete( - 'getFileInfo', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('getFileInfo', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Get file info tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Get file info received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Get file info completed', { - toolCallId, - toolName, - name: output.name, - size: output.size, - fileType: output.type, - permissions: output.permissions, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + id: 'getFileInfo', + description: + 'Get detailed information about a file or directory in the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe( + 'The sandboxId for the sandbox to get file information from' + ), + path: z + .string() + .describe( + 'The path to the file or directory to get information about' + ), + }), + outputSchema: z.object({ + name: z.string().describe('The name of the file or directory'), + type: z + .enum(FileType) + .optional() + .describe('Whether this is a file or directory'), + path: z.string().describe('The full path of the file or directory'), + size: z.number().describe('The size of the file or directory in bytes'), + mode: z + .number() + .describe('The file mode (permissions as octal number)'), + permissions: z.string().describe('Human-readable permissions string'), + owner: z.string().describe('The owner of the file or directory'), + group: z.string().describe('The group of the file or directory'), + modifiedTime: z + .date() + .optional() + .describe('The last modified time in ISO string format'), + symlinkTarget: z + .string() + .optional() + .describe('The target path if this is a symlink, null otherwise'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'getFileInfo', + input: inputData, + metadata: { 'tool.id': 'getFileInfo', 'tool.input.sandboxId': inputData.sandboxId, 'tool.input.path': inputData.path }, + }) + const startTime = Date.now() + + try { + logToolStart('getFileInfo', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + const info = await sandbox.files.getInfo(inputData.path) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'file.path': inputData.path, 'file.size': info.size } }) + if (info.type !== null && info.type !== undefined) { + span?.update({ metadata: { 'file.type': info.type } }) + } + span?.update({ metadata: { 'file.permissions': info.permissions } }) + + const result = { + name: info.name, + type: info.type, + path: info.path, + size: info.size, + mode: info.mode, + permissions: info.permissions, + owner: info.owner, + group: info.group, + modifiedTime: info.modifiedTime, + symlinkTarget: info.symlinkTarget, + } + + logToolComplete( + 'getFileInfo', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: true }) + const requestContext = context?.requestContext as E2BRequestContext + logError('getFileInfo', e, { + toolName: 'getFileInfo', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Get file info tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Get file info received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Get file info completed', { + toolCallId, + toolName, + name: output.name, + size: output.size, + fileType: output.type, + permissions: output.permissions, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const checkFileExists = createTool({ - id: 'checkFileExists', - description: 'Check if a file or directory exists in the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe( - 'The sandboxId for the sandbox to check file existence in' - ), - path: z.string().describe('The path to check for existence'), - }), - outputSchema: z - .object({ - exists: z - .boolean() - .describe('Whether the file or directory exists'), - path: z.string().describe('The path that was checked'), - type: z - .enum(FileType) - .optional() - .describe('The type if the path exists'), - }) - .or( - z.object({ - error: z - .string() - .describe('The error from a failed existence check'), - }) - ), - execute: async (inputData, context) => { - const span = tracer.startSpan('checkFileExists') - const startTime = Date.now() - - try { - logToolStart('checkFileExists', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - - try { - const info = await sandbox.files.getInfo(inputData.path) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('file.path', inputData.path) - span.setAttribute('file.exists', true) - if (info.type !== null && info.type !== undefined) { - span.setAttribute('file.type', info.type) - } - - const result = { - exists: true, - path: inputData.path, - type: info.type, - } - - logToolComplete( - 'checkFileExists', - result, - Date.now() - startTime, - context - ) - return result - } catch { - // If getInfo fails, the file doesn't exist - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('file.path', inputData.path) - span.setAttribute('file.exists', false) - - const result = { - exists: false, - path: inputData.path, - } - - logToolComplete( - 'checkFileExists', - result, - Date.now() - startTime, - context - ) - return result - } - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('checkFileExists', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() + id: 'checkFileExists', + description: 'Check if a file or directory exists in the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe( + 'The sandboxId for the sandbox to check file existence in' + ), + path: z.string().describe('The path to check for existence'), + }), + outputSchema: z + .object({ + exists: z + .boolean() + .describe('Whether the file or directory exists'), + path: z.string().describe('The path that was checked'), + type: z + .enum(FileType) + .optional() + .describe('The type if the path exists'), + }) + .or( + z.object({ + error: z + .string() + .describe('The error from a failed existence check'), + }) + ), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'checkFileExists', + input: inputData, + metadata: { 'tool.id': 'checkFileExists', 'tool.input.sandboxId': inputData.sandboxId, 'tool.input.path': inputData.path }, + }) + const startTime = Date.now() + + try { + logToolStart('checkFileExists', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + + try { + const info = await sandbox.files.getInfo(inputData.path) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'file.path': inputData.path, 'file.exists': true } }) + if (info.type !== null && info.type !== undefined) { + span?.update({ metadata: { 'file.type': info.type } }) } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Check file exists tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Check file exists received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Check file exists completed', { - toolCallId, - toolName, - exists: 'exists' in output ? output.exists : false, - path: 'path' in output ? output.path : '', - fileType: 'type' in output ? output.type : undefined, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + + const result = { + exists: true, + path: inputData.path, + type: info.type, + } + + logToolComplete( + 'checkFileExists', + result, + Date.now() - startTime, + context + ) + return result + } catch { + // If getInfo fails, the file doesn't exist + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'file.path': inputData.path, 'file.exists': false } }) + + const result = { + exists: false, + path: inputData.path, + } + + logToolComplete( + 'checkFileExists', + result, + Date.now() - startTime, + context + ) + return result + } + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: false }) + const requestContext = context?.requestContext as E2BRequestContext + logError('checkFileExists', e, { + toolName: 'checkFileExists', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Check file exists tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Check file exists received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Check file exists completed', { + toolCallId, + toolName, + exists: 'exists' in output ? output.exists : false, + path: 'path' in output ? output.path : '', + fileType: 'type' in output ? output.type : undefined, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const getFileSize = createTool({ - id: 'getFileSize', - description: 'Get the size of a file or directory in the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe('The sandboxId for the sandbox to get file size from'), - path: z.string().describe('The path to the file or directory'), - humanReadable: z - .boolean() - .default(false) - .describe( - "Whether to return size in human-readable format (e.g., '1.5 KB', '2.3 MB')" - ), - }), - outputSchema: z - .object({ - size: z.number().describe('The size in bytes'), - humanReadableSize: z - .string() - .optional() - .describe('Human-readable size string if requested'), - path: z.string().describe('The path that was checked'), - type: z - .enum(FileType) - .optional() - .describe('Whether this is a file or directory'), - }) - .or( - z.object({ - error: z - .string() - .describe('The error from a failed size check'), - }) - ), - execute: async (inputData, context) => { - const span = tracer.startSpan('getFileSize') - const startTime = Date.now() - - try { - logToolStart('getFileSize', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - const info = await sandbox.files.getInfo(inputData.path) - - let humanReadableSize: string | undefined - - if (inputData.humanReadable) { - const bytes = info.size - const sizes = ['B', 'KB', 'MB', 'GB', 'TB'] - if (bytes === 0) { - humanReadableSize = '0 B' - } else { - const i = Math.floor(Math.log(bytes) / Math.log(1024)) - const size = (bytes / Math.pow(1024, i)).toFixed(1) - humanReadableSize = `${size} ${sizes[i]}` - } - } - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('file.path', inputData.path) - span.setAttribute('file.size', info.size) - if (info.type !== null && info.type !== undefined) { - span.setAttribute('file.type', info.type) - } - span.setAttribute('humanReadable', inputData.humanReadable) - - const result = { - size: info.size, - humanReadableSize, - path: inputData.path, - type: info.type, - } - - logToolComplete( - 'getFileSize', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('getFileSize', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() + id: 'getFileSize', + description: 'Get the size of a file or directory in the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to get file size from'), + path: z.string().describe('The path to the file or directory'), + humanReadable: z + .boolean() + .default(false) + .describe( + "Whether to return size in human-readable format (e.g., '1.5 KB', '2.3 MB')" + ), + }), + outputSchema: z + .object({ + size: z.number().describe('The size in bytes'), + humanReadableSize: z + .string() + .optional() + .describe('Human-readable size string if requested'), + path: z.string().describe('The path that was checked'), + type: z + .enum(FileType) + .optional() + .describe('Whether this is a file or directory'), + }) + .or( + z.object({ + error: z + .string() + .describe('The error from a failed size check'), + }) + ), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'getFileSize', + input: inputData, + metadata: { 'tool.id': 'getFileSize', 'tool.input.sandboxId': inputData.sandboxId, 'tool.input.path': inputData.path }, + }) + const startTime = Date.now() + + try { + logToolStart('getFileSize', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + const info = await sandbox.files.getInfo(inputData.path) + + let humanReadableSize: string | undefined + + if (inputData.humanReadable) { + const bytes = info.size + const sizes = ['B', 'KB', 'MB', 'GB', 'TB'] + if (bytes === 0) { + humanReadableSize = '0 B' + } else { + const i = Math.floor(Math.log(bytes) / Math.log(1024)) + const size = (bytes / Math.pow(1024, i)).toFixed(1) + humanReadableSize = `${size} ${sizes[i]}` } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Get file size tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('Get file size received input chunk', { - toolCallId, - inputTextDelta, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Get file size received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - humanReadable: input.humanReadable, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Get file size completed', { - toolCallId, - toolName, - size: 'size' in output ? output.size : 0, - path: 'path' in output ? output.path : '', - humanReadable: - 'humanReadable' in output ? output.humanReadable : false, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + } + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'file.path': inputData.path, 'file.size': info.size, 'humanReadable': inputData.humanReadable } }) + if (info.type !== null && info.type !== undefined) { + span?.update({ metadata: { 'file.type': info.type } }) + } + + const result = { + size: info.size, + humanReadableSize, + path: inputData.path, + type: info.type, + } + + span?.update({ metadata: { 'tool.output_size': JSON.stringify(result).length, 'tool.duration_ms': Date.now() - startTime } }) + + logToolComplete( + 'getFileSize', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: false }) + const requestContext = context?.requestContext as E2BRequestContext + logError('getFileSize', e, { + toolName: 'getFileSize', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Get file size tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('Get file size received input chunk', { + toolCallId, + inputTextDelta, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Get file size received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + humanReadable: input.humanReadable, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Get file size completed', { + toolCallId, + toolName, + size: 'size' in output ? output.size : 0, + path: 'path' in output ? output.path : '', + humanReadable: + 'humanReadable' in output ? output.humanReadable : false, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const watchDirectory = createTool({ - id: 'watchDirectory', - description: 'Watch a directory for filesystem events in the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe('The sandboxId for the sandbox to watch directory in'), - path: z - .string() - .describe('The directory path to watch for events'), - recursive: z - .boolean() - .default(false) - .describe('Whether to watch subdirectories recursively'), - watchDuration: z - .number() - .default(5000) - .describe('How long to watch for events in milliseconds'), - }), - outputSchema: z - .object({ - watchStarted: z - .boolean() - .describe('Whether the watch was started successfully'), - path: z.string().describe('The path that was watched'), - events: z - .array( - z.object({ - type: z - .enum(FilesystemEventType) - .describe( - 'The type of filesystem event (WRITE, CREATE, DELETE, etc.)' - ), - name: z - .string() - .describe('The name of the file that changed'), - timestamp: z - .string() - .describe('When the event occurred'), - }) - ) - .describe( - 'Array of filesystem events that occurred during the watch period' - ), - }) - .or( - z.object({ - error: z - .string() - .describe('The error from a failed directory watch'), - }) + id: 'watchDirectory', + description: 'Watch a directory for filesystem events in the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to watch directory in'), + path: z + .string() + .describe('The directory path to watch for events'), + recursive: z + .boolean() + .default(false) + .describe('Whether to watch subdirectories recursively'), + watchDuration: z + .number() + .default(5000) + .describe('How long to watch for events in milliseconds'), + }), + outputSchema: z + .object({ + watchStarted: z + .boolean() + .describe('Whether the watch was started successfully'), + path: z.string().describe('The path that was watched'), + events: z + .array( + z.object({ + type: z + .enum(FilesystemEventType) + .describe( + 'The type of filesystem event (WRITE, CREATE, DELETE, etc.)' + ), + name: z + .string() + .describe('The name of the file that changed'), + timestamp: z + .string() + .describe('When the event occurred'), + }) + ) + .describe( + 'Array of filesystem events that occurred during the watch period' ), - execute: async (inputData, context) => { - const span = tracer.startSpan('watchDirectory') - const startTime = Date.now() - - try { - logToolStart('watchDirectory', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - const events: Array<{ - type: FilesystemEventType - name: string - timestamp: string - }> = [] - - // Start watching the directory - const handle = await sandbox.files.watchDir( - inputData.path, - async (event) => { - events.push({ - type: event.type, - name: event.name, - timestamp: new Date().toISOString(), - }) - }, - { - recursive: inputData.recursive, - } - ) - - // Watch for the specified duration - await new Promise((resolve) => - setTimeout(resolve, inputData.watchDuration) - ) - - // Stop watching - await handle.stop() - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('watch.path', inputData.path) - span.setAttribute('watch.recursive', inputData.recursive) - span.setAttribute('watch.duration', inputData.watchDuration) - span.setAttribute('events.count', events.length) - - const result = { - watchStarted: true, - path: inputData.path, - events, - } - - logToolComplete( - 'watchDirectory', - result, - Date.now() - startTime, - context - ) - return result - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('watchDirectory', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() + }) + .or( + z.object({ + error: z + .string() + .describe('The error from a failed directory watch'), + }) + ), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'watchDirectory', + input: inputData, + metadata: { + 'tool.id': 'watchDirectory', + 'tool.input.sandboxId': inputData.sandboxId, + 'tool.input.path': inputData.path, + 'tool.input.recursive': inputData.recursive, + 'tool.input.watchDuration': inputData.watchDuration, + }, + }) + const startTime = Date.now() + + try { + logToolStart('watchDirectory', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + const events: Array<{ + type: FilesystemEventType + name: string + timestamp: string + }> = [] + + // Start watching the directory + const handle = await sandbox.files.watchDir( + inputData.path, + async (event) => { + events.push({ + type: event.type, + name: event.name, + timestamp: new Date().toISOString(), + }) + }, + { + recursive: inputData.recursive, } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Watch directory tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Watch directory received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - path: input.path, - recursive: input.recursive, - watchDuration: input.watchDuration, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Watch directory completed', { - toolCallId, - toolName, - watchStarted: - 'watchStarted' in output ? output.watchStarted : false, - eventCount: 'events' in output ? output.events.length : 0, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + ) + + // Watch for the specified duration + await new Promise((resolve) => + setTimeout(resolve, inputData.watchDuration) + ) + + // Stop watching + await handle.stop() + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'watch.path': inputData.path, 'watch.recursive': inputData.recursive, 'watch.duration': inputData.watchDuration, 'events.count': events.length } }) + + const result = { + watchStarted: true, + path: inputData.path, + events, + } + + span?.update({ metadata: { 'tool.output_size': JSON.stringify(result).length, 'tool.duration_ms': Date.now() - startTime } }) + + logToolComplete( + 'watchDirectory', + result, + Date.now() - startTime, + context + ) + return result + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: false }) + const requestContext = context?.requestContext as E2BRequestContext + logError('watchDirectory', e, { + toolName: 'watchDirectory', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Watch directory tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Watch directory received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + path: input.path, + recursive: input.recursive, + watchDuration: input.watchDuration, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Watch directory completed', { + toolCallId, + toolName, + watchStarted: + 'watchStarted' in output ? output.watchStarted : false, + eventCount: 'events' in output ? output.events.length : 0, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const runCommand = createTool({ - id: 'runCommand', - description: 'Run a shell command in the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe('The sandboxId for the sandbox to run the command in'), - command: z.string().describe('The shell command to execute'), - workingDirectory: z - .string() - .optional() - .describe('The working directory to run the command in'), - timeoutMs: z - .number() - .default(30000) - .describe('Timeout for the command execution in milliseconds'), - captureOutput: z - .boolean() - .default(true) - .describe('Whether to capture stdout and stderr output'), - }), - outputSchema: z.object({ - success: z - .boolean() - .describe('Whether the command executed successfully'), - exitCode: z.number().describe('The exit code of the command'), - stdout: z.string().describe('The standard output from the command'), - stderr: z.string().describe('The standard error from the command'), - command: z.string().describe('The command that was executed'), - executionTime: z - .number() - .describe('How long the command took to execute in milliseconds'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('runCommand') - const startTime = Date.now() - - try { - logToolStart('runCommand', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - - const result = await sandbox.commands.run(inputData.command, { - cwd: inputData.workingDirectory, - timeoutMs: inputData.timeoutMs, - }) - - const executionTime = Date.now() - startTime - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('command.text', inputData.command) - span.setAttribute('command.cwd', inputData.workingDirectory ?? '/') - span.setAttribute('command.exit_code', result.exitCode) - span.setAttribute('command.execution_time', executionTime) - span.setAttribute( - 'command.stdout_length', - (result.stdout ?? '').length - ) - span.setAttribute( - 'command.stderr_length', - (result.stderr ?? '').length - ) - - const output = { - success: result.exitCode === 0, - exitCode: result.exitCode, - stdout: result.stdout ?? '', - stderr: result.stderr ?? '', - command: inputData.command, - executionTime, - } - - logToolComplete('runCommand', output, executionTime, context) - return output - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('runCommand', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Run command tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Run command received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - command: input.command, - workingDirectory: input.workingDirectory, - timeoutMs: input.timeoutMs, - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Run command completed', { - toolCallId, - toolName, - success: output.success, - exitCode: output.exitCode, - executionTime: output.executionTime, - stdoutLength: output.stdout.length, - stderrLength: output.stderr.length, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + id: 'runCommand', + description: 'Run a shell command in the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to run the command in'), + command: z.string().describe('The shell command to execute'), + workingDirectory: z + .string() + .optional() + .describe('The working directory to run the command in'), + timeoutMs: z + .number() + .default(30000) + .describe('Timeout for the command execution in milliseconds'), + captureOutput: z + .boolean() + .default(true) + .describe('Whether to capture stdout and stderr output'), + }), + outputSchema: z.object({ + success: z + .boolean() + .describe('Whether the command executed successfully'), + exitCode: z.number().describe('The exit code of the command'), + stdout: z.string().describe('The standard output from the command'), + stderr: z.string().describe('The standard error from the command'), + command: z.string().describe('The command that was executed'), + executionTime: z + .number() + .describe('How long the command took to execute in milliseconds'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'runCommand', + input: inputData, + metadata: { 'tool.id': 'runCommand' }, + }) + const startTime = Date.now() + + try { + logToolStart('runCommand', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + + const result = await sandbox.commands.run(inputData.command, { + cwd: inputData.workingDirectory, + timeoutMs: inputData.timeoutMs, + }) + + const executionTime = Date.now() - startTime + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'command.text': inputData.command, 'command.cwd': inputData.workingDirectory ?? '/', 'command.exit_code': result.exitCode, 'command.execution_time': executionTime, 'command.stdout_length': (result.stdout ?? '').length, 'command.stderr_length': (result.stderr ?? '').length } }) + + const output = { + success: result.exitCode === 0, + exitCode: result.exitCode, + stdout: result.stdout ?? '', + stderr: result.stderr ?? '', + command: inputData.command, + executionTime, + } + + // Output size attribute + const outputSize = JSON.stringify(output).length + span?.update({ metadata: { 'tool.output_size': outputSize, 'tool.duration_ms': Date.now() - startTime } }) + + logToolComplete('runCommand', output, executionTime, context) + return output + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: false }) + const requestContext = context?.requestContext as E2BRequestContext + logError('runCommand', e, { + toolName: 'runCommand', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Run command tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Run command received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + command: input.command, + workingDirectory: input.workingDirectory, + timeoutMs: input.timeoutMs, + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Run command completed', { + toolCallId, + toolName, + success: output.success, + exitCode: output.exitCode, + executionTime: output.executionTime, + stdoutLength: output.stdout.length, + stderrLength: output.stderr.length, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) export const runCode = createTool({ - id: 'runCode', - description: 'Execute code in the e2b sandbox', - inputSchema: z.object({ - sandboxId: z - .string() - .describe('The sandboxId for the sandbox to run code in'), - code: z.string().describe('The code to execute'), - runCodeOpts: z - .object({ - language: z - .enum(['python', 'javascript', 'typescript']) - .default('python') - .describe('The programming language of the code'), - }) - .optional() - .describe('Options for code execution'), - }), - outputSchema: z.object({ - execution: z.string().describe('JSON string containing execution results'), - }), - execute: async (inputData, context) => { - const span = tracer.startSpan('runCode') - const startTime = Date.now() - - try { - logToolStart('runCode', inputData, context) - - const sandbox = await Sandbox.connect(inputData.sandboxId) - - // Execute code using the appropriate interpreter based on language - const language = inputData.runCodeOpts?.language ?? 'python' - let command: string - - if (language === 'python') { - command = `python -c "${inputData.code.replace(/"/g, '\\"')}"` - } else if (language === 'javascript') { - command = `node -e "${inputData.code.replace(/"/g, '\\"')}"` - } else if (language === 'typescript') { - // For TypeScript, we need to compile and run - command = `npx tsx -e "${inputData.code.replace(/"/g, '\\"')}"` - } else { - throw new Error(`Unsupported language: ${language}`) - } - - const result = await sandbox.commands.run(command) - - span.setStatus({ code: SpanStatusCode.OK }) - span.setAttribute('sandbox.id', inputData.sandboxId) - span.setAttribute('code.language', language) - span.setAttribute('code.length', inputData.code.length) - - // Format result to match expected structure with logs.stdout and logs.stderr as arrays - const executionResult = { - logs: { - stdout: result.stdout ? result.stdout.split('\n') : [], - stderr: result.stderr ? result.stderr.split('\n') : [], - }, - exitCode: result.exitCode, - success: result.exitCode === 0, - } - - const output = { - execution: JSON.stringify(executionResult), - } - - logToolComplete('runCode', output, Date.now() - startTime, context) - return output - } catch (e) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: e instanceof Error ? e.message : String(e), - }) - logToolError('runCode', e, context) - const errObj = - e instanceof Error - ? { message: e.message, stack: e.stack } - : { message: String(e) } - throw new Error(JSON.stringify(errObj)) - } finally { - span.end() - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('Run code tool input streaming started', { - toolCallId, - messageCount: messages.length, - abortSignal: abortSignal?.aborted, - hook: 'onInputStart', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('Run code received complete input', { - toolCallId, - messageCount: messages.length, - sandboxId: input.sandboxId, - codeLength: input.code.length, - language: input.runCodeOpts?.language ?? 'python', - abortSignal: abortSignal?.aborted, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('Run code completed', { - toolCallId, - toolName, - executionLength: output.execution.length, - abortSignal: abortSignal?.aborted, - hook: 'onOutput', - }) - }, + id: 'runCode', + description: 'Execute code in the e2b sandbox', + inputSchema: z.object({ + sandboxId: z + .string() + .describe('The sandboxId for the sandbox to run code in'), + code: z.string().describe('The code to execute'), + runCodeOpts: z + .object({ + language: z + .enum(['python', 'javascript', 'typescript']) + .default('python') + .describe('The programming language of the code'), + }) + .optional() + .describe('Options for code execution'), + }), + outputSchema: z.object({ + execution: z.string().describe('JSON string containing execution results'), + }), + execute: async (inputData, context) => { + const tracingContext = context?.tracingContext + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'runCode', + input: inputData, + metadata: { + 'tool.id': 'runCode', + 'tool.input.sandboxId': inputData.sandboxId, + 'tool.input.codeLength': inputData.code.length, + 'tool.input.language': inputData.runCodeOpts?.language ?? 'python', + }, + }) + const startTime = Date.now() + + try { + logToolStart('runCode', inputData, context) + + const sandbox = await Sandbox.connect(inputData.sandboxId) + + // Execute code using the appropriate interpreter based on language + const language = inputData.runCodeOpts?.language ?? 'python' + let command: string + + if (language === 'python') { + command = `python -c "${inputData.code.replace(/"/g, '\\"')}"` + } else if (language === 'javascript') { + command = `node -e "${inputData.code.replace(/"/g, '\\"')}"` + } else if (language === 'typescript') { + // For TypeScript, we need to compile and run + command = `npx tsx -e "${inputData.code.replace(/"/g, '\\"')}"` + } else { + throw new Error(`Unsupported language: ${language}`) + } + + const result = await sandbox.commands.run(command) + + span?.update({ metadata: { 'operation.status': 'success', 'sandbox.id': inputData.sandboxId, 'code.language': language, 'code.length': inputData.code.length, 'command.exit_code': result.exitCode, 'command.stdout_length': (result.stdout ?? '').length, 'command.stderr_length': (result.stderr ?? '').length } }) + + // Format result to match expected structure with logs.stdout and logs.stderr as arrays + const executionResult = { + logs: { + stdout: result.stdout ? result.stdout.split('\n') : [], + stderr: result.stderr ? result.stderr.split('\n') : [], + }, + exitCode: result.exitCode, + success: result.exitCode === 0, + } + + const output = { + execution: JSON.stringify(executionResult), + } + + span?.update({ metadata: { 'tool.output_size': JSON.stringify(output).length, 'tool.duration_ms': Date.now() - startTime } }) + + logToolComplete('runCode', output, Date.now() - startTime, context) + return output + } catch (e) { + span?.error({ error: e instanceof Error ? e : new Error(String(e)), endSpan: false }) + const requestContext = context?.requestContext as E2BRequestContext + logError('runCode', e, { + toolName: 'runCode', + userId: requestContext?.userId + }) + const errObj = + e instanceof Error + ? { message: e.message, stack: e.stack } + : { message: String(e) } + throw new Error(JSON.stringify(errObj)) + } finally { + span?.end() + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('Run code tool input streaming started', { + toolCallId, + messageCount: messages.length, + abortSignal: abortSignal?.aborted, + hook: 'onInputStart', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('Run code received complete input', { + toolCallId, + messageCount: messages.length, + sandboxId: input.sandboxId, + codeLength: input.code.length, + language: input.runCodeOpts?.language ?? 'python', + abortSignal: abortSignal?.aborted, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('Run code completed', { + toolCallId, + toolName, + executionLength: output.execution.length, + abortSignal: abortSignal?.aborted, + hook: 'onOutput', + }) + }, }) // Type exports for UI components using InferUITool @@ -1664,33 +1561,32 @@ export type RunCodeUITool = InferUITool<typeof runCode> // Type union for all E2B UI tools export type E2BUITools = - | CreateSandboxUITool - | WriteFileUITool - | WriteFilesUITool - | ListFilesUITool - | DeleteFileUITool - | CreateDirectoryUITool - | GetFileInfoUITool - | CheckFileExistsUITool - | GetFileSizeUITool - | WatchDirectoryUITool - | RunCommandUITool - | RunCodeUITool + | CreateSandboxUITool + | WriteFileUITool + | WriteFilesUITool + | ListFilesUITool + | DeleteFileUITool + | CreateDirectoryUITool + | GetFileInfoUITool + | CheckFileExistsUITool + | GetFileSizeUITool + | WatchDirectoryUITool + | RunCommandUITool + | RunCodeUITool // Tool set for all E2B tools export const e2bToolSet = { - createSandbox, - writeFile, - writeFiles, - listFiles, - deleteFile, - createDirectory, - getFileInfo, - checkFileExists, - getFileSize, - watchDirectory, - runCommand, - runCode, + writeFile, + writeFiles, + listFiles, + deleteFile, + createDirectory, + getFileInfo, + checkFileExists, + getFileSize, + watchDirectory, + runCommand, + runCode, } as const // Type inference for the entire E2B tool set using InferUITools diff --git a/src/mastra/tools/editor-agent-tool.ts b/src/mastra/tools/editor-agent-tool.ts index 74f24846..b44f1540 100644 --- a/src/mastra/tools/editor-agent-tool.ts +++ b/src/mastra/tools/editor-agent-tool.ts @@ -1,7 +1,7 @@ import type { MastraModelOutput } from '@mastra/core/stream' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { SpanStatusCode, trace } from '@opentelemetry/api' +import { SpanType } from "@mastra/core/observability"; import { z } from 'zod' import { editorAgent } from '../agents/editorAgent' import { log } from '../config/logger' @@ -102,15 +102,15 @@ export const editorTool = createTool({ } = inputData const writer = context?.writer const requestContext = context?.requestContext as EditorAgentContext | undefined + const tracingContext = context?.tracingContext - const tracer = trace.getTracer('editor-agent-tool') - const span = tracer.startSpan('editor-agent', { - attributes: { + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'editor-agent', + input: inputData, + metadata: { + 'tool.id': 'editor-agent', contentType, - contentLength: content.length, - hasInstructions: - typeof instructions === 'string' && - instructions.trim().length > 0, tone: tone ?? 'not-specified', operation: 'editor-agent-run', }, @@ -201,10 +201,10 @@ export const editorTool = createTool({ } catch (err) { // Streaming/generation error — bubble up after logging const msg = err instanceof Error ? err.message : String(err) - span.recordException( - err instanceof Error ? err : new Error(msg) - ) - span.setStatus({ code: SpanStatusCode.ERROR, message: msg }) + span?.error({ + error: err instanceof Error ? err : new Error(msg), + endSpan: true, + }) throw err } @@ -222,7 +222,13 @@ export const editorTool = createTool({ } } - span.end() + span?.update({ + output: parsedResult, + metadata: { + 'tool.output.success': true, + } + }) + span?.end() await writer?.custom({ type: 'data-tool-progress', @@ -248,11 +254,10 @@ export const editorTool = createTool({ } catch (error) { const errorMsg = error instanceof Error ? error.message : 'Unknown error' - span.recordException( - error instanceof Error ? error : new Error(errorMsg) - ) - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }) - span.end() + span?.error({ + error: error instanceof Error ? error : new Error(errorMsg), + endSpan: true, + }) throw new Error(`Failed to edit content: ${errorMsg}`) } }, diff --git a/src/mastra/tools/evaluateResultTool.ts b/src/mastra/tools/evaluateResultTool.ts index 995ce564..5c920635 100644 --- a/src/mastra/tools/evaluateResultTool.ts +++ b/src/mastra/tools/evaluateResultTool.ts @@ -1,7 +1,7 @@ import type { MastraModelOutput } from '@mastra/core/stream' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { SpanStatusCode, trace } from '@opentelemetry/api' +import { SpanType } from "@mastra/core/observability"; import { z } from 'zod' import { evaluationAgent } from '../agents/evaluationAgent' import { log } from '../config/logger' @@ -47,7 +47,7 @@ export const evaluateResultTool = createTool({ toolCallId, query: input.query, url: input.result.url, - existingUrlsCount: input.existingUrls?.length || 0, + existingUrlsCount: input.existingUrls?.length ?? 0, hook: 'onInputAvailable', }) }, @@ -74,12 +74,15 @@ export const evaluateResultTool = createTool({ id: 'evaluate-result', }) - const tracer = trace.getTracer('evaluate-result') - const evalSpan = tracer.startSpan('evaluate_result', { - attributes: { + const tracingContext = context?.tracingContext + const evalSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'evaluate_result', + input: inputData, + metadata: { + 'tool.id': 'evaluate-result', query: inputData.query, url: inputData.result.url, - existingUrlsCount: inputData.existingUrls?.length ?? 0, operation: 'evaluate_result', }, }) @@ -90,7 +93,11 @@ export const evaluateResultTool = createTool({ // Check if URL already exists (only if existingUrls was provided) if (existingUrls?.includes(result.url)) { - evalSpan.end() + evalSpan?.update({ + output: { isRelevant: false, reason: 'URL already processed' }, + metadata: { 'tool.output.status': 'already-processed' } + }) + evalSpan?.end() await context?.writer?.custom({ type: 'data-tool-progress', data: { @@ -167,12 +174,10 @@ export const evaluateResultTool = createTool({ } catch (err) { const errorMessage = err instanceof Error ? err.message : String(err) - evalSpan.recordException(new Error(errorMessage)) - evalSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + evalSpan?.error({ + error: err instanceof Error ? err : new Error(errorMessage), + endSpan: true, }) - evalSpan.end() await context?.writer?.custom({ type: 'data-tool-progress', @@ -215,13 +220,11 @@ export const evaluateResultTool = createTool({ response: responseObject, }) const error = 'Invalid response format from evaluation agent' - evalSpan.recordException(new Error(error)) - evalSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: error, + evalSpan?.error({ + error: new Error(error), + endSpan: true, }) - evalSpan.end() await context?.writer?.custom({ type: 'data-tool-progress', data: { @@ -237,7 +240,7 @@ export const evaluateResultTool = createTool({ } } - evalSpan.end() + evalSpan?.end() await context?.writer?.custom({ type: 'data-tool-progress', @@ -259,12 +262,10 @@ export const evaluateResultTool = createTool({ const errorMessage = error instanceof Error ? error.message : String(error) - evalSpan.recordException(new Error(errorMessage)) - evalSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + evalSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - evalSpan.end() // Emit final progress event with error message await context?.writer?.custom({ diff --git a/src/mastra/tools/execa-tool.ts b/src/mastra/tools/execa-tool.ts index b909bb12..57ca5b08 100644 --- a/src/mastra/tools/execa-tool.ts +++ b/src/mastra/tools/execa-tool.ts @@ -1,6 +1,6 @@ import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { SpanType } from "@mastra/core/observability"; import chalk from 'chalk' import type { ExecaError as ExecaErrorType } from 'execa' import execa from 'execa' @@ -43,20 +43,19 @@ export const execaTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer const requestContext = context?.requestContext as ExecaContext | undefined + const tracingContext = context?.tracingContext - const tracer = trace.getTracer('execa-tool', '1.0.0') - const span = tracer.startSpan('execa-tool', { - attributes: { - 'tool.id': 'execa-tool', + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'execa-tool', + input: inputData, + metadata: { + 'tool.id': 'execaTool', 'tool.input.command': inputData.command, - 'tool.input.args': inputData.args.join(' '), - 'tool.requestContext.userId': requestContext?.userId, + 'user.id': requestContext?.userId, }, }) - // Ensure colorTransform is "used" to satisfy strict linting without removal - span.setAttribute('internal.colorTransform', !!colorTransform) - const { command, args, cwd, timeout, env } = inputData await writer?.custom({ type: 'data-tool-progress', @@ -92,18 +91,22 @@ export const execaTool = createTool({ }, id: 'execaTool', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.outputLength': output.length, + span?.update({ + output: { success: true, outputLength: output.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.outputLength': output.length, + } }) - span.end() + span?.end() return { message: chalk.green(output) } } catch (e) { const errorMsg = e instanceof Error ? e.message : String(e) log.error(errorMsg) - span.recordException(e instanceof Error ? e : new Error(errorMsg)) - span.setStatus({ code: 2, message: errorMsg }) - span.end() + span?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) const execaErr = e as ExecaErrorType if (e instanceof Error && 'all' in e) { return { diff --git a/src/mastra/tools/fs.ts b/src/mastra/tools/fs.ts index 5acd1755..2eee889b 100644 --- a/src/mastra/tools/fs.ts +++ b/src/mastra/tools/fs.ts @@ -1,5 +1,5 @@ import { createTool } from '@mastra/core/tools' -import { SpanStatusCode, trace } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import { promises as fsPromises } from 'node:fs' import { z } from 'zod' import { log } from '../config/logger' @@ -24,13 +24,22 @@ export const fsTool = createTool({ const writer = context?.writer const abortSignal = context?.abortSignal const requestContext = context?.requestContext as FsToolContext | undefined + const tracingContext = context?.tracingContext + if (abortSignal?.aborted === true) { throw new Error('FS operation cancelled') } - const tracer = trace.getTracer('tools/fs-tool') - const span = tracer.startSpan('fs-tool', { - attributes: { action: inputData.action, file: inputData.file }, + // Create child span for FS operation + const fsSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'fs-operation', + input: inputData, + metadata: { + 'tool.id': 'fsTool', + 'tool.input.action': inputData.action, + 'tool.input.file': inputData.file, + }, }) const { action, file, data } = inputData @@ -46,11 +55,10 @@ export const fsTool = createTool({ try { // Check for cancellation before file operations if (abortSignal?.aborted) { - span?.setStatus({ - code: 2, - message: 'Operation cancelled during file operations', + fsSpan?.error({ + error: new Error('Operation cancelled during file operations'), + endSpan: true, }) - span?.end() throw new Error('FS operation cancelled during file operations') } @@ -77,8 +85,13 @@ export const fsTool = createTool({ }, id: 'fsTool', }) - span?.setAttribute('success', true) - span?.end() + fsSpan?.update({ + output: { message: 'Success' }, + metadata: { + 'tool.output.success': true, + }, + }) + fsSpan?.end() return { message: 'Success' } } catch (e) { const errorMsg = e instanceof Error ? e.message : String(e) @@ -86,8 +99,10 @@ export const fsTool = createTool({ // Handle AbortError specifically if (e instanceof Error && e.name === 'AbortError') { const cancelMessage = `FS operation cancelled` - span?.setStatus({ code: 2, message: cancelMessage }) - span?.end() + fsSpan?.error({ + error: new Error(cancelMessage), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', @@ -115,9 +130,10 @@ export const fsTool = createTool({ id: 'fsTool', }) log.error(`FS operation failed: ${errorMsg}`) - span?.recordException(e instanceof Error ? e : new Error(errorMsg)) - span?.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }) - span?.end() + fsSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { message: `Error: ${errorMsg}`, } diff --git a/src/mastra/tools/json-to-csv.tool.ts b/src/mastra/tools/json-to-csv.tool.ts index f8acf313..01288e66 100644 --- a/src/mastra/tools/json-to-csv.tool.ts +++ b/src/mastra/tools/json-to-csv.tool.ts @@ -1,7 +1,7 @@ import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import { z } from 'zod' import { log } from '../config/logger' @@ -39,6 +39,7 @@ export const jsonToCsvTool = createTool({ const writer = context?.writer const requestContext = context?.requestContext as JsonToCsvRequestContext | undefined const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext // Check if operation was already cancelled if (abortSignal?.aborted === true) { @@ -55,9 +56,12 @@ export const jsonToCsvTool = createTool({ id: 'json-to-csv', }) - const tracer = trace.getTracer('json-to-csv', '1.0.0') - const rootSpan = tracer.startSpan('json-to-csv', { - attributes: { + // Create child span for JSON to CSV conversion + const jsonCsvSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'json-to-csv-conversion', + input, + metadata: { 'tool.id': 'json-to-csv', 'tool.input.recordCount': input.data.length, }, @@ -66,7 +70,7 @@ export const jsonToCsvTool = createTool({ try { const { data, options } = input if (data === undefined || data === null || data.length === 0) { - rootSpan.end() + jsonCsvSpan?.end() return { csv: '' } } @@ -124,10 +128,14 @@ export const jsonToCsvTool = createTool({ const csvOutput = rows.join('\n') - rootSpan.setAttributes({ - 'tool.output.csvLength': csvOutput.length, + jsonCsvSpan?.update({ + output: { csv: csvOutput }, + metadata: { + 'tool.output.csvLength': csvOutput.length, + 'tool.output.success': true, + }, }) - rootSpan.end() + jsonCsvSpan?.end() return { csv: csvOutput } } catch (error) { const errorMessage = @@ -138,8 +146,10 @@ export const jsonToCsvTool = createTool({ // Handle AbortError specifically if (error instanceof Error && error.name === 'AbortError') { const cancelMessage = `JSON to CSV conversion cancelled` - rootSpan.setStatus({ code: 2, message: cancelMessage }) - rootSpan.end() + jsonCsvSpan?.error({ + error: new Error(cancelMessage), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', @@ -155,11 +165,10 @@ export const jsonToCsvTool = createTool({ throw new Error(cancelMessage) } - rootSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - rootSpan.setStatus({ code: 2, message: errorMessage }) - rootSpan.end() + jsonCsvSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) throw new Error(errorMessage) } }, @@ -181,7 +190,7 @@ export const jsonToCsvTool = createTool({ }) }, onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - const options = input.options || {} + const options = input.options ?? {} log.info('JSON to CSV received complete input', { toolCallId, messageCount: messages.length, diff --git a/src/mastra/tools/pnpm-tool.ts b/src/mastra/tools/pnpm-tool.ts index b46d11be..b59e238a 100644 --- a/src/mastra/tools/pnpm-tool.ts +++ b/src/mastra/tools/pnpm-tool.ts @@ -1,7 +1,7 @@ +import { SpanType } from '@mastra/core/observability'; import type { RequestContext } from '@mastra/core/request-context'; import type { InferUITool } from "@mastra/core/tools"; import { createTool } from "@mastra/core/tools"; -import { SpanStatusCode, trace } from "@opentelemetry/api"; import chalk from 'chalk'; import type { ExecaError as ExecaErrorType } from 'execa'; import execa from 'execa'; @@ -17,7 +17,7 @@ const pnpmContextSchema = z.object({ export type PnpmContext = z.infer<typeof pnpmContextSchema>; export interface PnpmRequestContext extends RequestContext { - pnpmToolContext?: PnpmContext + pnpmToolContext?: PnpmContext } export const pnpmBuild = createTool({ @@ -55,17 +55,23 @@ export const pnpmBuild = createTool({ }, execute: async (inputData, context) => { const writer = context?.writer; + const abortSignal = context?.abortSignal; + const tracingContext = context?.tracingContext; const requestContext = context?.requestContext as PnpmRequestContext | undefined; const pnpmContext = requestContext?.pnpmToolContext; const { verbose } = pnpmContextSchema.parse(pnpmContext ?? {}); - const tracer = trace.getTracer('pnpm-tool'); - const span = tracer.startSpan('pnpm-build', { - attributes: { - name: inputData.name, - packagePath: inputData.packagePath, - operation: 'pnpm-build' - } + // Create child span for pnpm build using Mastra tracing context + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'pnpm-build', + input: { name: inputData.name, packagePath: inputData.packagePath }, + metadata: { + 'tool.id': 'pnpmBuild', + 'operation': 'pnpm-build', + 'package.name': inputData.name, + 'package.path': inputData.packagePath, + }, }); const { name, packagePath } = inputData @@ -81,15 +87,25 @@ export const pnpmBuild = createTool({ await p if (verbose) { await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `✅ Build complete for ${name}`, stage: 'pnpmBuild' }, id: 'pnpmBuild' }); } - span.end(); + span?.end(); return { message: 'Done' } } catch (e) { - const error = e as ExecaErrorType - const errorMsg = error.message || String(e) - log.error(errorMsg) - span.recordException(new Error(errorMsg)); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }); - span.end(); + const error = e as ExecaErrorType; + const errorMsg = error.message || String(e); + log.error(errorMsg); + span?.error({ + error: error instanceof Error ? error : new Error(errorMsg), + endSpan: true, + }); + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `Input: name="${inputData.name}" - ❌ Build failed: ${errorMsg}`, + stage: 'pnpmBuild', + }, + id: 'pnpmBuild', + }); return { message: errorMsg || 'Error' } } }, @@ -126,11 +142,17 @@ export const pnpmChangesetStatus = createTool({ }, execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('pnpm-tool'); - const span = tracer.startSpan('pnpm-changeset-status', { - attributes: { - operation: 'pnpm-changeset-status' - } + const abortSignal = context?.abortSignal; + const tracingContext = context?.tracingContext; + + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'pnpm-changeset-status', + input: {}, + metadata: { + 'tool.id': 'pnpmChangesetStatus', + 'operation': 'pnpm-changeset-status', + }, }); await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🔍 Checking changeset status...', stage: 'pnpmChangesetStatus' }, id: 'pnpmChangesetStatus' }); @@ -157,16 +179,15 @@ export const pnpmChangesetStatus = createTool({ await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `✅ Found ${packages.length} packages to publish`, stage: 'pnpmChangesetStatus' }, id: 'pnpmChangesetStatus' }); - span.setAttribute('packageCount', packages.length); - span.end(); + span?.update({ output: { message: packages }, metadata: { packageCount: packages.length, 'tool.output.success': true } }); + span?.end(); return { message: packages } } catch (e) { - const error = e as ExecaErrorType + const error = e as ExecaErrorType; const errorMsg = error.message || String(e); - log.error(errorMsg) - span.recordException(error); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }); - span.end(); + log.error(errorMsg); + span?.error({ error: error instanceof Error ? error : new Error(errorMsg), endSpan: true }); + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `🔍 Changeset check failed: ${errorMsg}`, stage: 'pnpmChangesetStatus' }, id: 'pnpmChangesetStatus' }); return { message: [] } } }, @@ -203,11 +224,17 @@ export const pnpmChangesetPublish = createTool({ }, execute: async (inputData, context) => { const writer = context?.writer; - const tracer = trace.getTracer('pnpm-tool'); - const span = tracer.startSpan('pnpm-changeset-publish', { - attributes: { - operation: 'pnpm-changeset-publish' - } + const abortSignal = context?.abortSignal; + const tracingContext = context?.tracingContext; + + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'pnpm-changeset-publish', + input: {}, + metadata: { + 'tool.id': 'pnpmChangesetPublish', + 'operation': 'pnpm-changeset-publish', + }, }); // const {value} = input // unused @@ -222,15 +249,15 @@ export const pnpmChangesetPublish = createTool({ await p await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Publish complete', stage: 'pnpmChangesetPublish' }, id: 'pnpmChangesetPublish' }); - span.end(); + span?.update({ output: { message: 'Done' }, metadata: { 'tool.output.success': true } }); + span?.end(); return { message: 'Done' } } catch (e) { const error = e as ExecaErrorType const errorMsg = error.message || String(e) log.error(errorMsg) - span.recordException(new Error(errorMsg)); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }); - span.end(); + span?.error({ error: error instanceof Error ? error : new Error(errorMsg), endSpan: true }); + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `🚨 Publish failed: ${errorMsg}`, stage: 'pnpmChangesetPublish' }, id: 'pnpmChangesetPublish' }); if (e instanceof Error && 'message' in e) { return { message: errorMsg } } @@ -274,12 +301,18 @@ export const activeDistTag = createTool({ }, execute: async (input, context) => { const writer = context?.writer; - const tracer = trace.getTracer('pnpm-tool'); - const span = tracer.startSpan('active-dist-tag', { - attributes: { - packagePath: input.packagePath, - operation: 'active-dist-tag' - } + const abortSignal = context?.abortSignal; + const tracingContext = context?.tracingContext; + + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'active-dist-tag', + input: { packagePath: input.packagePath }, + metadata: { + 'tool.id': 'activeDistTag', + 'operation': 'active-dist-tag', + 'package.path': input.packagePath, + }, }); const { packagePath } = input @@ -305,15 +338,16 @@ export const activeDistTag = createTool({ log.info(`\n`) await p await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `✅ Dist tag set to latest for ${pkgJson.name}@${version}`, stage: 'activeDistTag' }, id: 'activeDistTag' }); - span.end(); + + span?.update({ output: { message: 'Done' }, metadata: { 'tool.output.success': true } }); + span?.end(); return { message: 'Done' } } catch (e) { const error = e as ExecaErrorType const errorMsg = error.message || String(e) log.error(errorMsg) - span.recordException(new Error(errorMsg)); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }); - span.end(); + span?.error({ error: error instanceof Error ? error : new Error(errorMsg), endSpan: true }); + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `🚨 Dist tag failed: ${errorMsg}`, stage: 'activeDistTag' }, id: 'activeDistTag' }); if (e instanceof Error && 'message' in e) { return { message: errorMsg } } @@ -361,14 +395,20 @@ export const pnpmRun = createTool({ }, execute: async (input, context) => { const writer = context?.writer; - const tracer = trace.getTracer('pnpm-tool'); - const span = tracer.startSpan('pnpm-run', { - attributes: { + const abortSignal = context?.abortSignal; + const tracingContext = context?.tracingContext; + + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'pnpm-run', + input: { script: input.script, args: input.args, packagePath: input.packagePath }, + metadata: { + 'tool.id': 'pnpmRun', + 'operation': 'pnpm-run', script: input.script, args: JSON.stringify(input.args), - packagePath: input.packagePath, - operation: 'pnpm-run' - } + 'package.path': input.packagePath, + }, }); const { script, args = [], packagePath } = input @@ -383,15 +423,16 @@ export const pnpmRun = createTool({ log.info(`\n`) await p await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `✅ Script ${script} complete`, stage: 'pnpmRun' }, id: 'pnpmRun' }); - span.end(); + + span?.update({ output: { message: 'Done' }, metadata: { 'tool.output.success': true } }); + span?.end(); return { message: 'Done' } } catch (e) { const error = e as ExecaErrorType const errorMsg = error.message || String(e) log.error(errorMsg) - span.recordException(new Error(errorMsg)); - span.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }); - span.end(); + span?.error({ error: error instanceof Error ? error : new Error(errorMsg), endSpan: true }); + await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: `❌ Script failed: ${errorMsg}`, stage: 'pnpmRun' }, id: 'pnpmRun' }); return { message: errorMsg || 'Error' } } }, diff --git a/src/mastra/tools/serpapi-academic-local.tool.ts b/src/mastra/tools/serpapi-academic-local.tool.ts index 362fdeb5..6f15570d 100644 --- a/src/mastra/tools/serpapi-academic-local.tool.ts +++ b/src/mastra/tools/serpapi-academic-local.tool.ts @@ -6,7 +6,7 @@ * @module serpapi-academic-local-tool */ import type { RequestContext } from '@mastra/core/request-context' -import { trace, SpanStatusCode } from '@opentelemetry/api' +import { SpanType } from "@mastra/core/observability"; import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' import { getJson } from 'serpapi' @@ -82,20 +82,24 @@ export const googleScholarTool = createTool({ const writer = context?.writer const abortSignal = context?.abortSignal const requestContext = context?.requestContext as SerpApiContext | undefined + const tracingContext = context?.tracingContext // Check if operation was already cancelled if (abortSignal?.aborted === true) { throw new Error('Google Scholar search cancelled') } - const tracer = trace.getTracer('serpapi-academic-local-tool') - const scholarSpan = tracer.startSpan('google-scholar-tool', { - attributes: { + const scholarSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-scholar-tool', + input, + metadata: { + 'tool.id': 'googlescholar', query: input.query, yearRange: `${input.yearStart}-${input.yearEnd}`, operation: 'google-scholar', - }, - }) + } + }); log.info('Executing Google Scholar search', { query: input.query }) await writer?.custom({ type: 'data-tool-progress', @@ -121,11 +125,7 @@ export const googleScholarTool = createTool({ } // Check for cancellation before API call if (abortSignal?.aborted) { - scholarSpan.setStatus({ - code: 2, - message: 'Operation cancelled during API call', - }) - scholarSpan.end() + scholarSpan?.end() throw new Error( 'Google Scholar search cancelled during API call' ) @@ -171,7 +171,13 @@ export const googleScholarTool = createTool({ }) ) ?? [] const result = { papers } - scholarSpan.end() + scholarSpan?.update({ + output: result, + metadata: { + 'tool.output.paperCount': papers.length, + } + }) + scholarSpan?.end() log.info('Google Scholar search completed', { query: input.query, paperCount: papers.length, @@ -181,8 +187,10 @@ export const googleScholarTool = createTool({ // Handle AbortError specifically if (error instanceof Error && error.name === 'AbortError') { const cancelMessage = `Google Scholar search cancelled for "${input.query}"` - scholarSpan.setStatus({ code: 2, message: cancelMessage }) - scholarSpan.end() + scholarSpan?.error({ + error: new Error(cancelMessage), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', @@ -200,14 +208,10 @@ export const googleScholarTool = createTool({ const errorMessage = error instanceof Error ? error.message : String(error) - scholarSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - scholarSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + scholarSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - scholarSpan.end() log.error('Google Scholar search failed', { query: input.query, error: errorMessage, @@ -303,14 +307,18 @@ export const googleFinanceTool = createTool({ validateSerpApiKey() const writer = context?.writer const requestContext = context?.requestContext as SerpApiContext | undefined + const tracingContext = context?.tracingContext - const tracer = trace.getTracer('serpapi-academic-local-tool') - const financeSpan = tracer.startSpan('google-finance-tool', { - attributes: { + const financeSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-finance-tool', + input, + metadata: { + 'tool.id': 'google-finance', query: input.query, operation: 'google-finance', - }, - }) + } + }); log.info('Executing Google Finance search', { query: input.query }) await writer?.custom({ type: 'data-tool-progress', @@ -357,7 +365,13 @@ export const googleFinanceTool = createTool({ previousClose: summary?.previous_close, news, } - financeSpan.end() + financeSpan?.update({ + output: result, + metadata: { + 'tool.output.symbol': result.symbol, + } + }) + financeSpan?.end() log.info('Google Finance search completed', { query: input.query, symbol: result.symbol, @@ -366,14 +380,10 @@ export const googleFinanceTool = createTool({ } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - financeSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - financeSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + financeSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - financeSpan.end() log.error('Google Finance search failed', { query: input.query, error: errorMessage, @@ -442,15 +452,19 @@ export const yelpSearchTool = createTool({ validateSerpApiKey() const writer = context?.writer const requestContext = context?.requestContext as SerpApiContext | undefined + const tracingContext = context?.tracingContext - const tracer = trace.getTracer('serpapi-academic-local-tool') - const yelpSpan = tracer.startSpan('yelp-search-tool', { - attributes: { + const yelpSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'yelp-search-tool', + input, + metadata: { + 'tool.id': 'yelp-search', query: input.query, location: input.location, operation: 'yelp-search', - }, - }) + } + }); log.info('Executing Yelp search', { query: input.query, location: input.location, @@ -511,7 +525,13 @@ export const yelpSearchTool = createTool({ }) ) ?? [] const result = { businesses } - yelpSpan.end() + yelpSpan?.update({ + output: result, + metadata: { + 'tool.output.businessCount': businesses.length, + } + }) + yelpSpan?.end() log.info('Yelp search completed', { query: input.query, location: input.location, @@ -521,14 +541,10 @@ export const yelpSearchTool = createTool({ } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - yelpSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - yelpSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + yelpSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - yelpSpan.end() log.error('Yelp search failed', { query: input.query, location: input.location, diff --git a/src/mastra/tools/serpapi-news-trends.tool.ts b/src/mastra/tools/serpapi-news-trends.tool.ts index 631afec8..8a5f04f4 100644 --- a/src/mastra/tools/serpapi-news-trends.tool.ts +++ b/src/mastra/tools/serpapi-news-trends.tool.ts @@ -6,7 +6,7 @@ * * @module serpapi-news-trends-tool */ -import { trace, SpanStatusCode } from '@opentelemetry/api' +import { SpanType } from "@mastra/core/observability"; import { createTool } from '@mastra/core/tools' import { z } from 'zod' import { getJson } from 'serpapi' @@ -126,6 +126,7 @@ export const googleNewsTool = createTool({ execute: async (input, context) => { const writer = context?.writer const requestContext = context?.requestContext as SerpApiContext | undefined + const tracingContext = context?.tracingContext await writer?.custom({ type: 'data-tool-progress', data: { @@ -137,15 +138,18 @@ export const googleNewsTool = createTool({ }) validateSerpApiKey() - const tracer = trace.getTracer('serpapi-news-trends-tool') - const newsSpan = tracer.startSpan('google-news-tool', { - attributes: { + const newsSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-news-tool', + input, + metadata: { + 'tool.id': 'google-news', query: input.query, timeRange: input.timeRange, topic: input.topic, operation: 'google-news', - }, - }) + } + }); await writer?.custom({ type: 'data-tool-progress', data: { @@ -212,7 +216,13 @@ export const googleNewsTool = createTool({ }, id: 'google-news', }) - newsSpan.end() + newsSpan?.update({ + output: result, + metadata: { + 'tool.output.articlesFound': newsArticles.length, + } + }) + newsSpan?.end() log.info('Google News search completed', { query: input.query, resultCount: newsArticles.length, @@ -221,14 +231,10 @@ export const googleNewsTool = createTool({ } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - newsSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - newsSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + newsSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - newsSpan.end() log.error('Google News search failed', { query: input.query, error: errorMessage, @@ -266,6 +272,7 @@ export const googleNewsLiteTool = createTool({ }), execute: async (input, context) => { const writer = context?.writer + const tracingContext = context?.tracingContext await writer?.custom({ type: 'data-tool-progress', data: { @@ -280,13 +287,16 @@ export const googleNewsLiteTool = createTool({ }) validateSerpApiKey() - const tracer = trace.getTracer('serpapi-news-trends-tool') - const newsLiteSpan = tracer.startSpan('google-news-lite-tool', { - attributes: { + const newsLiteSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-news-lite-tool', + input, + metadata: { + 'tool.id': 'google-news-lite', query: input.query, operation: 'google-news-lite', - }, - }) + } + }); await writer?.custom({ type: 'data-tool-progress', data: { @@ -329,7 +339,13 @@ export const googleNewsLiteTool = createTool({ }, id: 'google-news-lite', }) - newsLiteSpan.end() + newsLiteSpan?.update({ + output: result, + metadata: { + 'tool.output.articlesFound': newsArticles.length, + } + }) + newsLiteSpan?.end() log.info('Google News Lite search completed', { query: input.query, resultCount: newsArticles.length, @@ -338,14 +354,10 @@ export const googleNewsLiteTool = createTool({ } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - newsLiteSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - newsLiteSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + newsLiteSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - newsLiteSpan.end() log.error('Google News Lite search failed', { query: input.query, error: errorMessage, @@ -421,6 +433,7 @@ export const googleTrendsTool = createTool({ outputSchema: googleTrendsOutputSchema, execute: async (input, context) => { const writer = context?.writer + const tracingContext = context?.tracingContext await writer?.custom({ type: 'data-tool-progress', data: { @@ -435,14 +448,17 @@ export const googleTrendsTool = createTool({ }) validateSerpApiKey() - const tracer = trace.getTracer('serpapi-news-trends-tool') - const trendsSpan = tracer.startSpan('google-trends-tool', { - attributes: { + const trendsSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-trends-tool', + input, + metadata: { + 'tool.id': 'google-trends', query: input.query, timeRange: input.timeRange, operation: 'google-trends', - }, - }) + } + }); await writer?.custom({ type: 'data-tool-progress', @@ -514,7 +530,13 @@ export const googleTrendsTool = createTool({ }, id: 'google-trends', }) - trendsSpan.end() + trendsSpan?.update({ + output: result, + metadata: { + 'tool.output.dataPoints': interestOverTime.length, + } + }) + trendsSpan?.end() log.info('Google Trends search completed', { query: input.query, dataPoints: interestOverTime.length, @@ -524,14 +546,10 @@ export const googleTrendsTool = createTool({ } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - trendsSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - trendsSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + trendsSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - trendsSpan.end() log.error('Google Trends search failed', { query: input.query, error: errorMessage, @@ -578,6 +596,7 @@ export const googleAutocompleteTool = createTool({ outputSchema: googleAutocompleteOutputSchema, execute: async (input, context) => { const writer = context?.writer + const tracingContext = context?.tracingContext await writer?.custom({ type: 'data-tool-progress', data: { @@ -592,13 +611,16 @@ export const googleAutocompleteTool = createTool({ }) validateSerpApiKey() - const tracer = trace.getTracer('serpapi-news-trends-tool') - const autocompleteSpan = tracer.startSpan('google-autocomplete-tool', { - attributes: { + const autocompleteSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-autocomplete-tool', + input, + metadata: { + 'tool.id': 'google-autocomplete', query: input.query, operation: 'google-autocomplete', - }, - }) + } + }); await writer?.custom({ type: 'data-tool-progress', @@ -644,7 +666,13 @@ export const googleAutocompleteTool = createTool({ }, id: 'google-autocomplete', }) - autocompleteSpan.end() + autocompleteSpan?.update({ + output: result, + metadata: { + 'tool.output.suggestionCount': suggestions.length, + } + }) + autocompleteSpan?.end() log.info('Google Autocomplete completed', { query: input.query, suggestionCount: suggestions.length, @@ -654,14 +682,10 @@ export const googleAutocompleteTool = createTool({ } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - autocompleteSpan.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - autocompleteSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + autocompleteSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - autocompleteSpan.end() log.error('Google Autocomplete failed', { query: input.query, error: errorMessage, diff --git a/src/mastra/tools/serpapi-search.tool.ts b/src/mastra/tools/serpapi-search.tool.ts index 45009b64..d4da5517 100644 --- a/src/mastra/tools/serpapi-search.tool.ts +++ b/src/mastra/tools/serpapi-search.tool.ts @@ -7,7 +7,7 @@ * @module serpapi-search-tool */ import type { RequestContext } from '@mastra/core/request-context' -import { trace, SpanStatusCode } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import { createTool } from '@mastra/core/tools' import { getJson } from 'serpapi' import { z } from 'zod' @@ -103,6 +103,7 @@ export const googleSearchTool = createTool({ const writer = context?.writer const abortSignal = context?.abortSignal const requestContext = context?.requestContext as SerpApiContext | undefined + const tracingContext = context?.tracingContext // Check if operation was already cancelled if (abortSignal?.aborted === true) { @@ -118,15 +119,21 @@ export const googleSearchTool = createTool({ id: 'serpapi-search', }) - const tracer = trace.getTracer('serpapi-search') - const searchSpan = tracer.startSpan('google-search-tool', { - attributes: { - query: input.query, - numResults: input.numResults, - location: input.location, - operation: 'google-search', + // Create child span for Google search + const searchSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-search', + input, + metadata: { + 'tool.id': 'google-search', + 'tool.input.query': input.query, + 'tool.input.numResults': input.numResults, + 'tool.input.location': input.location, + 'tool.input.language': input.language, + 'tool.input.device': input.device, }, }) + log.info('Executing Google search', { query: input.query, numResults: input.numResults, @@ -155,11 +162,10 @@ export const googleSearchTool = createTool({ } // Check for cancellation before API call if (abortSignal?.aborted) { - searchSpan.setStatus({ - code: 2, - message: 'Operation cancelled during API call', + searchSpan?.error({ + error: new Error('Operation cancelled during API call'), + endSpan: true, }) - searchSpan.end() throw new Error('Google search cancelled during API call') } @@ -223,7 +229,16 @@ export const googleSearchTool = createTool({ searchInfo, } - searchSpan.end() + searchSpan?.update({ + output: result, + metadata: { + 'tool.output.organicResults': organicResults.length, + 'tool.output.hasKnowledgeGraph': !!knowledgeGraph, + 'tool.output.relatedSearches': relatedSearches?.length ?? 0, + 'tool.output.success': true, + }, + }) + searchSpan?.end() log.info('Google search completed successfully', { query: input.query, resultCount: organicResults.length, @@ -245,8 +260,10 @@ export const googleSearchTool = createTool({ // Handle AbortError specifically if (error instanceof Error && error.name === 'AbortError') { const cancelMessage = `Google search cancelled for "${input.query}"` - searchSpan.setStatus({ code: 2, message: cancelMessage }) - searchSpan.end() + searchSpan?.error({ + error: new Error(cancelMessage), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', @@ -271,12 +288,10 @@ export const googleSearchTool = createTool({ }, id: 'serpapi-search', }) - searchSpan.recordException(new Error(errorMessage)) - searchSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + searchSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - searchSpan.end() log.error('Google search failed', { query: input.query, error: errorMessage, @@ -284,7 +299,7 @@ export const googleSearchTool = createTool({ throw error } }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { + onInputStart: ({ toolCallId, abortSignal }) => { log.info('Google search tool input streaming started', { toolCallId, abortSignal: abortSignal?.aborted, @@ -300,7 +315,7 @@ export const googleSearchTool = createTool({ hook: 'onInputDelta', }) }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + onInputAvailable: ({ input, toolCallId, abortSignal }) => { log.info('Google search received complete input', { toolCallId, abortSignal: abortSignal?.aborted, @@ -319,7 +334,7 @@ export const googleSearchTool = createTool({ abortSignal: abortSignal?.aborted, organicResults: output.organicResults.length, hasKnowledgeGraph: !!output.knowledgeGraph, - relatedSearches: output.relatedSearches?.length || 0, + relatedSearches: output.relatedSearches?.length ?? 0, hook: 'onOutput', }) }, @@ -378,7 +393,8 @@ export const googleAiOverviewTool = createTool({ outputSchema: googleAiOverviewOutputSchema, execute: async (input, context) => { const writer = context?.writer - const requestContext = context?.requestContext as SerpApiContext | undefined + const tracingContext = context?.tracingContext + await writer?.custom({ type: 'data-tool-progress', data: { @@ -391,15 +407,19 @@ export const googleAiOverviewTool = createTool({ // Validate API key validateSerpApiKey() - const tracer = trace.getTracer('serpapi-search') - const overviewSpan = tracer.startSpan('google-ai-overview-tool', { - attributes: { - query: input.query, - location: input.location, - includeScrapedContent: input.includeScrapedContent, - operation: 'google-ai-overview', + // Create child span for AI overview + const overviewSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'google-ai-overview', + input, + metadata: { + 'tool.id': 'google-ai-overview', + 'tool.input.query': input.query, + 'tool.input.location': input.location, + 'tool.input.includeScrapedContent': input.includeScrapedContent, }, }) + await writer?.custom({ type: 'data-tool-progress', data: { @@ -448,7 +468,16 @@ export const googleAiOverviewTool = createTool({ : undefined, available, } - overviewSpan.end() + overviewSpan?.update({ + output: result, + metadata: { + 'tool.output.available': available, + 'tool.output.sourcesCount': sources.length, + 'tool.output.hasAiOverview': Boolean(aiOverviewData?.text), + 'tool.output.success': true, + }, + }) + overviewSpan?.end() await writer?.custom({ type: 'data-tool-progress', data: { @@ -476,12 +505,10 @@ export const googleAiOverviewTool = createTool({ }, id: 'google-ai-overview', }) - overviewSpan.recordException(new Error(errorMessage)) - overviewSpan.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, + overviewSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, }) - overviewSpan.end() log.error('Google AI Overview failed', { query: input.query, error: errorMessage, diff --git a/src/mastra/tools/serpapi-shopping.tool.ts b/src/mastra/tools/serpapi-shopping.tool.ts index 175c4ebc..e9c3c276 100644 --- a/src/mastra/tools/serpapi-shopping.tool.ts +++ b/src/mastra/tools/serpapi-shopping.tool.ts @@ -6,7 +6,7 @@ * @module serpapi-shopping-tool */ import { createTool } from '@mastra/core/tools'; -import { trace } from "@opentelemetry/api"; +import { SpanType } from "@mastra/core/observability"; import { getJson } from 'serpapi'; import { z } from 'zod'; import { log } from '../config/logger'; @@ -72,14 +72,16 @@ export const amazonSearchTool = createTool({ }, execute: async (inputData, context) => { const writer = context?.writer; + const tracingContext = context?.tracingContext; await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🛒 Starting Amazon search for "' + inputData.query + '"', stage: 'amazon-search' }, id: 'amazon-search' }); validateSerpApiKey() - // Get tracer from OpenTelemetry API - const tracer = trace.getTracer('amazon-search-tool', '1.0.0'); - const amazonSpan = tracer.startSpan('amazon-search', { - attributes: { + const amazonSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'amazon-search', + input: inputData, + metadata: { 'tool.id': 'amazon-search', 'tool.input.query': inputData.query, 'tool.input.sortBy': inputData.sortBy, @@ -147,18 +149,22 @@ export const amazonSearchTool = createTool({ const result = { products } await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Amazon search complete: ' + products.length + ' products', stage: 'amazon-search' }, id: 'amazon-search' }); - amazonSpan.setAttribute('tool.output.productCount', products.length); - amazonSpan.end(); + amazonSpan?.update({ + output: result, + metadata: { + 'tool.output.productCount': products.length + } + }); + amazonSpan?.end(); log.info('Amazon search completed', { query: inputData.query, productCount: products.length }) return result } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - amazonSpan.recordException(error); - } - amazonSpan.setStatus({ code: 2, message: errorMessage }); // ERROR status - amazonSpan.end(); + amazonSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) log.error('Amazon search failed', { query: inputData.query, error: errorMessage }) throw new Error(`Amazon search failed: ${errorMessage}`) } @@ -221,14 +227,16 @@ export const walmartSearchTool = createTool({ }, execute: async (inputData, context) => { const writer = context?.writer; + const tracingContext = context?.tracingContext; await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🛒 Starting Walmart search for "' + inputData.query + '"', stage: 'walmart-search' }, id: 'walmart-search' }); validateSerpApiKey() - // Get tracer from OpenTelemetry API - const tracer = trace.getTracer('walmart-search-tool', '1.0.0'); - const walmartSpan = tracer.startSpan('walmart-search', { - attributes: { + const walmartSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'walmart-search', + input: inputData, + metadata: { 'tool.id': 'walmart-search', 'tool.input.query': inputData.query, 'tool.input.sortBy': inputData.sortBy, @@ -280,18 +288,22 @@ export const walmartSearchTool = createTool({ const result = { products } await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Walmart search complete: ' + products.length + ' products', stage: 'walmart-search' }, id: 'walmart-search' }); - walmartSpan.setAttribute('tool.output.productCount', products.length); - walmartSpan.end(); + walmartSpan?.update({ + output: result, + metadata: { + 'tool.output.productCount': products.length + } + }); + walmartSpan?.end(); log.info('Walmart search completed', { query: inputData.query, productCount: products.length }) return result } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - walmartSpan.recordException(error); - } - walmartSpan.setStatus({ code: 2, message: errorMessage }); // ERROR status - walmartSpan.end(); + walmartSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) log.error('Walmart search failed', { query: inputData.query, error: errorMessage }) throw new Error(`Walmart search failed: ${errorMessage}`) } @@ -356,14 +368,16 @@ export const ebaySearchTool = createTool({ }, execute: async (inputData, context) => { const writer = context?.writer; + const tracingContext = context?.tracingContext; await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🛒 Starting eBay search for "' + inputData.query + '"', stage: 'ebay-search' }, id: 'ebay-search' }); validateSerpApiKey() - // Get tracer from OpenTelemetry API - const tracer = trace.getTracer('ebay-search-tool', '1.0.0'); - const ebaySpan = tracer.startSpan('ebay-search', { - attributes: { + const ebaySpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'ebay-search', + input: inputData, + metadata: { 'tool.id': 'ebay-search', 'tool.input.query': inputData.query, 'tool.input.condition': inputData.condition, @@ -421,17 +435,21 @@ export const ebaySearchTool = createTool({ ) ?? [] const result = { products } await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ eBay search complete: ' + products.length + ' products', stage: 'ebay-search' }, id: 'ebay-search' }); - ebaySpan.setAttribute('tool.output.productCount', products.length); - ebaySpan.end(); + ebaySpan?.update({ + output: result, + metadata: { + 'tool.output.productCount': products.length + } + }); + ebaySpan?.end(); log.info('eBay search completed', { query: inputData.query, productCount: products.length }) return result } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - ebaySpan.recordException(error); - } - ebaySpan.setStatus({ code: 2, message: errorMessage }); // ERROR status - ebaySpan.end(); + ebaySpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) log.error('eBay search failed', { query: inputData.query, error: errorMessage }) throw new Error(`eBay search failed: ${errorMessage}`) } @@ -493,14 +511,16 @@ export const homeDepotSearchTool = createTool({ }, execute: async (inputData, context) => { const writer = context?.writer; + const tracingContext = context?.tracingContext; await writer?.custom({ type: 'data-tool-progress', data: { status: 'in-progress', message: '🛒 Starting Home Depot search for "' + inputData.query + '"', stage: 'home-depot-search' }, id: 'home-depot-search' }); validateSerpApiKey() - // Get tracer from OpenTelemetry API - const tracer = trace.getTracer('home-depot-search-tool', '1.0.0'); - const homeDepotSpan = tracer.startSpan('home-depot-search', { - attributes: { + const homeDepotSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'home-depot-search', + input: inputData, + metadata: { 'tool.id': 'home-depot-search', 'tool.input.query': inputData.query, 'tool.input.sortBy': inputData.sortBy, @@ -546,17 +566,21 @@ export const homeDepotSearchTool = createTool({ ) ?? [] const result = { products } await writer?.custom({ type: 'data-tool-progress', data: { status: 'done', message: '✅ Home Depot search complete: ' + products.length + ' products', stage: 'home-depot-search' }, id: 'home-depot-search' }); - homeDepotSpan.setAttribute('tool.output.productCount', products.length); - homeDepotSpan.end(); + homeDepotSpan?.update({ + output: result, + metadata: { + 'tool.output.productCount': products.length + } + }); + homeDepotSpan?.end(); log.info('Home Depot search completed', { query: inputData.query, productCount: products.length }) return result } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - homeDepotSpan.recordException(error); - } - homeDepotSpan.setStatus({ code: 2, message: errorMessage }); // ERROR status - homeDepotSpan.end(); + homeDepotSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) log.error('Home Depot search failed', { query: inputData.query, error: errorMessage }) throw new Error(`Home Depot search failed: ${errorMessage}`) } diff --git a/src/mastra/tools/test-generator.tool.ts b/src/mastra/tools/test-generator.tool.ts index 43d6d5e5..c58842ab 100644 --- a/src/mastra/tools/test-generator.tool.ts +++ b/src/mastra/tools/test-generator.tool.ts @@ -1,178 +1,178 @@ +import { SpanType } from '@mastra/core/observability' import { createTool } from '@mastra/core/tools' -import { SpanStatusCode, trace } from '@opentelemetry/api' import { promises as fs } from 'node:fs' import * as path from 'node:path' import { z } from 'zod' import { log } from '../config/logger' async function fileExists(filePath: string): Promise<boolean> { - try { - await fs.access(filePath) - return true - } catch { - return false - } + try { + await fs.access(filePath) + return true + } catch { + return false + } } const testGeneratorInputSchema = z.object({ - sourceFile: z - .string() - .describe('Path to the source file to generate tests for'), - outputPath: z - .string() - .optional() - .describe('Output path for test file (defaults to __tests__ folder)'), - options: z - .object({ - includeEdgeCases: z - .boolean() - .default(true) - .describe('Generate edge case tests'), - mockExternals: z - .boolean() - .default(true) - .describe('Add mock setup for external dependencies'), - testStyle: z - .enum(['describe-it', 'test-only']) - .default('describe-it') - .describe('Test structure style'), - }) - .optional(), + sourceFile: z + .string() + .describe('Path to the source file to generate tests for'), + outputPath: z + .string() + .optional() + .describe('Output path for test file (defaults to __tests__ folder)'), + options: z + .object({ + includeEdgeCases: z + .boolean() + .default(true) + .describe('Generate edge case tests'), + mockExternals: z + .boolean() + .default(true) + .describe('Add mock setup for external dependencies'), + testStyle: z + .enum(['describe-it', 'test-only']) + .default('describe-it') + .describe('Test structure style'), + }) + .optional(), }) const testCaseSchema = z.object({ - name: z.string(), - type: z.enum(['unit', 'integration', 'edge-case']), - code: z.string(), + name: z.string(), + type: z.enum(['unit', 'integration', 'edge-case']), + code: z.string(), }) const testGeneratorOutputSchema = z - .object({ - testFile: z.string(), - sourceFile: z.string(), - framework: z.literal('vitest'), - content: z.string(), - tests: z.array(testCaseSchema), - coverage: z.object({ - functions: z.array(z.string()), - exports: z.array(z.string()), - }), - runCommand: z.string(), - }) - .describe('Test generator output') + .object({ + testFile: z.string(), + sourceFile: z.string(), + framework: z.literal('vitest'), + content: z.string(), + tests: z.array(testCaseSchema), + coverage: z.object({ + functions: z.array(z.string()), + exports: z.array(z.string()), + }), + runCommand: z.string(), + }) + .describe('Test generator output') export type TestGeneratorInput = z.infer<typeof testGeneratorInputSchema> export type TestGeneratorOutput = z.infer<typeof testGeneratorOutputSchema> interface ParsedFunction { - name: string - isAsync: boolean - isExported: boolean - params: string[] - returnType?: string + name: string + isAsync: boolean + isExported: boolean + params: string[] + returnType?: string } function parseTypeScriptFunctions(content: string): ParsedFunction[] { - const functions: ParsedFunction[] = [] - - // Match exported functions: export function name(params): type - const exportFuncRegex = - /export\s+(async\s+)?function\s+(\w+)\s*\(([^)]*)\)(?:\s*:\s*([^{]+))?/g - let match: RegExpExecArray | null - while ((match = exportFuncRegex.exec(content)) !== null) { - functions.push({ - name: match[2], - isAsync: !!match[1], - isExported: true, - params: match[3] - ? match[3] - .split(',') - .map((p) => p.trim().split(':')[0].trim()) - .filter(Boolean) - : [], - returnType: match[4]?.trim(), - }) - } - - // Match exported const arrow functions: export const name = (params) => - const exportConstRegex = - /export\s+const\s+(\w+)\s*=\s*(async\s+)?\(([^)]*)\)(?:\s*:\s*([^=]+))?\s*=>/g - while ((match = exportConstRegex.exec(content)) !== null) { - functions.push({ - name: match[1], - isAsync: !!match[2], - isExported: true, - params: match[3] - ? match[3] - .split(',') - .map((p) => p.trim().split(':')[0].trim()) - .filter(Boolean) - : [], - returnType: match[4]?.trim(), - }) - } - - // Match regular functions (not exported) - const funcRegex = - /(?<!export\s+)(async\s+)?function\s+(\w+)\s*\(([^)]*)\)(?:\s*:\s*([^{]+))?/g - while ((match = funcRegex.exec(content)) !== null) { - const matchedName = match[2] - if (!functions.some((f) => f.name === matchedName)) { - functions.push({ - name: matchedName, - isAsync: !!match[1], - isExported: false, - params: match[3] - ? match[3] - .split(',') - .map((p) => p.trim().split(':')[0].trim()) - .filter(Boolean) - : [], - returnType: match[4]?.trim(), - }) - } + const functions: ParsedFunction[] = [] + + // Match exported functions: export function name(params): type + const exportFuncRegex = + /export\s+(async\s+)?function\s+(\w+)\s*\(([^)]*)\)(?:\s*:\s*([^{]+))?/g + let match: RegExpExecArray | null + while ((match = exportFuncRegex.exec(content)) !== null) { + functions.push({ + name: match[2], + isAsync: !!match[1], + isExported: true, + params: match[3] + ? match[3] + .split(',') + .map((p) => p.trim().split(':')[0].trim()) + .filter(Boolean) + : [], + returnType: match[4]?.trim(), + }) + } + + // Match exported const arrow functions: export const name = (params) => + const exportConstRegex = + /export\s+const\s+(\w+)\s*=\s*(async\s+)?\(([^)]*)\)(?:\s*:\s*([^=]+))?\s*=>/g + while ((match = exportConstRegex.exec(content)) !== null) { + functions.push({ + name: match[1], + isAsync: !!match[2], + isExported: true, + params: match[3] + ? match[3] + .split(',') + .map((p) => p.trim().split(':')[0].trim()) + .filter(Boolean) + : [], + returnType: match[4]?.trim(), + }) + } + + // Match regular functions (not exported) + const funcRegex = + /(?<!export\s+)(async\s+)?function\s+(\w+)\s*\(([^)]*)\)(?:\s*:\s*([^{]+))?/g + while ((match = funcRegex.exec(content)) !== null) { + const matchedName = match[2] + if (!functions.some((f) => f.name === matchedName)) { + functions.push({ + name: matchedName, + isAsync: !!match[1], + isExported: false, + params: match[3] + ? match[3] + .split(',') + .map((p) => p.trim().split(':')[0].trim()) + .filter(Boolean) + : [], + returnType: match[4]?.trim(), + }) } + } - return functions + return functions } function parseExports(content: string): string[] { - const exports: string[] = [] - - const namedExportRegex = - /export\s+(?:const|let|var|function|class|type|interface)\s+(\w+)/g - let exportMatch: RegExpExecArray | null - while ((exportMatch = namedExportRegex.exec(content)) !== null) { - exports.push(exportMatch[1]) - } - - const defaultExportRegex = /export\s+default\s+(?:function\s+)?(\w+)?/g - while ((exportMatch = defaultExportRegex.exec(content)) !== null) { - if (exportMatch[1]) { - exports.push(`default (${exportMatch[1]})`) - } else { - exports.push('default') - } + const exports: string[] = [] + + const namedExportRegex = + /export\s+(?:const|let|var|function|class|type|interface)\s+(\w+)/g + let exportMatch: RegExpExecArray | null + while ((exportMatch = namedExportRegex.exec(content)) !== null) { + exports.push(exportMatch[1]) + } + + const defaultExportRegex = /export\s+default\s+(?:function\s+)?(\w+)?/g + while ((exportMatch = defaultExportRegex.exec(content)) !== null) { + if (exportMatch[1]) { + exports.push(`default (${exportMatch[1]})`) + } else { + exports.push('default') } + } - return [...new Set(exports)] + return [...new Set(exports)] } function generateTestCase( - func: ParsedFunction, - sourceFile: string, - options: { includeEdgeCases?: boolean; testStyle?: string } + func: ParsedFunction, + sourceFile: string, + options: { includeEdgeCases?: boolean; testStyle?: string } ): Array<z.infer<typeof testCaseSchema>> { - const tests: Array<z.infer<typeof testCaseSchema>> = [] + const tests: Array<z.infer<typeof testCaseSchema>> = [] - // Basic unit test - const asyncPrefix = func.isAsync ? 'async ' : '' - const awaitPrefix = func.isAsync ? 'await ' : '' + // Basic unit test + const asyncPrefix = func.isAsync ? 'async ' : '' + const awaitPrefix = func.isAsync ? 'await ' : '' - tests.push({ - name: `${func.name} - should work correctly`, - type: 'unit', - code: ` it('should work correctly', ${asyncPrefix}() => { + tests.push({ + name: `${func.name} - should work correctly`, + type: 'unit', + code: ` it('should work correctly', ${asyncPrefix}() => { // Arrange ${func.params.length > 0 ? func.params.map((p) => `const ${p} = undefined // TODO: Add test value`).join('\n ') : '// No parameters'} @@ -182,60 +182,60 @@ function generateTestCase( // Assert expect(result).toBeDefined() })`, - }) - - // Edge case tests - if (options.includeEdgeCases !== false) { - if (func.params.length > 0) { - tests.push({ - name: `${func.name} - should handle empty/null inputs`, - type: 'edge-case', - code: ` it('should handle edge cases', ${asyncPrefix}() => { + }) + + // Edge case tests + if (options.includeEdgeCases !== false) { + if (func.params.length > 0) { + tests.push({ + name: `${func.name} - should handle empty/null inputs`, + type: 'edge-case', + code: ` it('should handle edge cases', ${asyncPrefix}() => { // Test with undefined/null inputs ${func.isAsync ? 'await ' : ''}expect(() => ${func.name}(${func.params.map(() => 'undefined').join(', ')})).not.toThrow() })`, - }) - } - - if (func.isAsync) { - tests.push({ - name: `${func.name} - should handle async errors`, - type: 'edge-case', - code: ` it('should handle async errors gracefully', async () => { + }) + } + + if (func.isAsync) { + tests.push({ + name: `${func.name} - should handle async errors`, + type: 'edge-case', + code: ` it('should handle async errors gracefully', async () => { // Test error handling // TODO: Mock dependencies to throw and verify error handling expect(true).toBe(true) })`, - }) - } + }) } + } - return tests + return tests } function generateTestFile( - sourceFile: string, - functions: ParsedFunction[], - exports: string[], - options: { - testStyle?: string - mockExternals?: boolean - includeEdgeCases?: boolean - } + sourceFile: string, + functions: ParsedFunction[], + exports: string[], + options: { + testStyle?: string + mockExternals?: boolean + includeEdgeCases?: boolean + } ): { content: string; tests: Array<z.infer<typeof testCaseSchema>> } { - const moduleName = path.basename(sourceFile, path.extname(sourceFile)) - const allTests: Array<z.infer<typeof testCaseSchema>> = [] + const moduleName = path.basename(sourceFile, path.extname(sourceFile)) + const allTests: Array<z.infer<typeof testCaseSchema>> = [] - const exportedFuncs = functions.filter((f) => f.isExported) - const importNames = exportedFuncs.map((f) => f.name).join(', ') + const exportedFuncs = functions.filter((f) => f.isExported) + const importNames = exportedFuncs.map((f) => f.name).join(', ') - let content = `import { describe, it, expect, vi, beforeEach } from 'vitest' + let content = `import { describe, it, expect, vi, beforeEach } from 'vitest' import { ${importNames || moduleName} } from './${moduleName}' ` - if (options.mockExternals !== false) { - content += `// Mock external dependencies + if (options.mockExternals !== false) { + content += `// Mock external dependencies // vi.mock('dependency-name', () => ({ ... })) beforeEach(() => { @@ -243,155 +243,172 @@ beforeEach(() => { }) ` - } + } - for (const func of exportedFuncs) { - const testCases = generateTestCase(func, sourceFile, options) - allTests.push(...testCases) + for (const func of exportedFuncs) { + const testCases = generateTestCase(func, sourceFile, options) + allTests.push(...testCases) - if (options.testStyle === 'describe-it') { - content += `describe('${func.name}', () => { + if (options.testStyle === 'describe-it') { + content += `describe('${func.name}', () => { ${testCases.map((t) => t.code).join('\n\n')} }) ` - } else { - for (const test of testCases) { - content += `test('${func.name} - ${test.name}', ${func.isAsync ? 'async ' : ''}() => { + } else { + for (const test of testCases) { + content += `test('${func.name} - ${test.name}', ${func.isAsync ? 'async ' : ''}() => { ${test.code.replace(/^ {2}/gm, '')} }) ` - } - } + } } + } - return { content: content.trim() + '\n', tests: allTests } + return { content: content.trim() + '\n', tests: allTests } } export const testGeneratorTool = createTool({ - id: 'coding:testGenerator', - description: `Generate Vitest test scaffolds for TypeScript/JavaScript source files. + id: 'coding:testGenerator', + description: `Generate Vitest test scaffolds for TypeScript/JavaScript source files. Parses exported functions and creates test cases with describe/it blocks. Supports edge case generation and mock setup for external dependencies. Use for increasing test coverage and establishing testing patterns.`, - inputSchema: testGeneratorInputSchema, - outputSchema: testGeneratorOutputSchema, - - execute: async (inputData, context?) => { - const { sourceFile, outputPath, options } = inputData - const { includeEdgeCases, mockExternals, testStyle } = options ?? {} - - const tracer = trace.getTracer('test-generator-tool') - const span = tracer.startSpan('test-generator', { - attributes: { - sourceFile, - testStyle, - operation: 'generate-tests', - }, - }) + inputSchema: testGeneratorInputSchema, + outputSchema: testGeneratorOutputSchema, + + execute: async (inputData, context?) => { + const { sourceFile, outputPath, options } = inputData + const { includeEdgeCases, mockExternals, testStyle } = options ?? {} + + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as any | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const span = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'test_generator', + input: { sourceFile, testStyle, operation: 'generate-tests' }, + metadata: { + 'tool.id': 'coding:testGenerator', + 'tool.input.sourceFile': sourceFile, + 'tool.input.testStyle': testStyle, + 'tool.input.operation': 'generate-tests', + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - try { - if (!(await fileExists(sourceFile))) { - throw new Error(`Source file not found: ${sourceFile}`) - } - - await context?.writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🧪 Generating tests for: ${sourceFile}`, - stage: 'coding:testGenerator', - }, - id: 'coding:testGenerator', - }) - - const content = await fs.readFile(sourceFile, 'utf-8') - const functions = parseTypeScriptFunctions(content) - const exports = parseExports(content) - - const { content: testContent, tests } = generateTestFile( - sourceFile, - functions, - exports, - { testStyle, mockExternals, includeEdgeCases } - ) - - const sourceDir = path.dirname(sourceFile) - const sourceName = path.basename( - sourceFile, - path.extname(sourceFile) - ) - const testFile = - outputPath ?? - path.join(sourceDir, '__tests__', `${sourceName}.test.ts`) - - span.end() - return { - testFile, - sourceFile, - framework: 'vitest' as const, - content: testContent, - tests, - coverage: { - functions: functions - .filter((f) => f.isExported) - .map((f) => f.name), - exports, - }, - runCommand: `npx vitest ${testFile}`, - } - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : String(error) - span.recordException( - error instanceof Error ? error : new Error(errorMessage) - ) - span.setStatus({ - code: SpanStatusCode.ERROR, - message: errorMessage, - }) - span.end() - throw error - } - }, - onInputStart: ({ toolCallId, messages, abortSignal }) => { - log.info('testGeneratorTool tool input streaming started', { - toolCallId, - hook: 'onInputStart', - }) - }, - onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { - log.info('testGeneratorTool received input chunk', { - toolCallId, - inputTextDelta, - hook: 'onInputDelta', - }) - }, - onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { - log.info('testGeneratorTool received input', { - toolCallId, - inputData: { - sourceFile: input.sourceFile, - outputPath: input.outputPath, - options: input.options, - }, - hook: 'onInputAvailable', - }) - }, - onOutput: ({ output, toolCallId, toolName, abortSignal }) => { - log.info('testGeneratorTool completed', { - toolCallId, - toolName, - outputData: { - testFile: output.testFile, - sourceFile: output.sourceFile, - framework: output.framework, - tests: output.tests, - coverage: output.coverage, - runCommand: output.runCommand, - }, - hook: 'onOutput', + try { + if (!(await fileExists(sourceFile))) { + throw new Error(`Source file not found: ${sourceFile}`) + } + + await context?.writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🧪 Generating tests for: ${sourceFile}`, + stage: 'coding:testGenerator', + }, + id: 'coding:testGenerator', + }) + + const content = await fs.readFile(sourceFile, 'utf-8') + const functions = parseTypeScriptFunctions(content) + const exports = parseExports(content) + + const { content: testContent, tests } = generateTestFile( + sourceFile, + functions, + exports, + { testStyle, mockExternals, includeEdgeCases } + ) + + const sourceDir = path.dirname(sourceFile) + const sourceName = path.basename( + sourceFile, + path.extname(sourceFile) + ) + const testFile = + outputPath ?? + path.join(sourceDir, '__tests__', `${sourceName}.test.ts`) + + const result = { + testFile, + sourceFile, + framework: 'vitest' as const, + content: testContent, + tests, + coverage: { + functions: functions + .filter((f) => f.isExported) + .map((f) => f.name), + exports, + }, + runCommand: `npx vitest ${testFile}`, + } + try { + span?.update({ + output: result, + metadata: { + 'tool.output.testCount': tests.length, + 'tool.output.testFile': testFile + }, }) - }, + } catch { + // ignore tracing update errors + } + span?.end() + return result + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : String(error) + span?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) + throw error + } + }, + onInputStart: ({ toolCallId, messages, abortSignal }) => { + log.info('testGeneratorTool tool input streaming started', { + toolCallId, + hook: 'onInputStart', + }) + }, + onInputDelta: ({ inputTextDelta, toolCallId, messages, abortSignal }) => { + log.info('testGeneratorTool received input chunk', { + toolCallId, + inputTextDelta, + hook: 'onInputDelta', + }) + }, + onInputAvailable: ({ input, toolCallId, messages, abortSignal }) => { + log.info('testGeneratorTool received input', { + toolCallId, + inputData: { + sourceFile: input.sourceFile, + outputPath: input.outputPath, + options: input.options, + }, + hook: 'onInputAvailable', + }) + }, + onOutput: ({ output, toolCallId, toolName, abortSignal }) => { + log.info('testGeneratorTool completed', { + toolCallId, + toolName, + outputData: { + testFile: output.testFile, + sourceFile: output.sourceFile, + framework: output.framework, + tests: output.tests, + coverage: output.coverage, + runCommand: output.runCommand, + }, + hook: 'onOutput', + }) + }, }) diff --git a/src/mastra/tools/text-analysis.tool.ts b/src/mastra/tools/text-analysis.tool.ts index f989cc75..75bf5107 100644 --- a/src/mastra/tools/text-analysis.tool.ts +++ b/src/mastra/tools/text-analysis.tool.ts @@ -1,14 +1,14 @@ +import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import { z } from 'zod' import { log } from '../config/logger' -import type { RequestContext } from '@mastra/core/request-context' - export interface TextAnalysisToolContext extends RequestContext { defaultLanguage?: string includeAdvancedMetrics?: boolean + userId?: string } export const textAnalysisTool = createTool({ @@ -62,16 +62,22 @@ export const textAnalysisTool = createTool({ }), execute: async (inputData, context) => { const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext const requestCtx = context?.requestContext as TextAnalysisToolContext | undefined const defaultLanguage = requestCtx?.defaultLanguage ?? 'en' const includeAdvancedMetrics = requestCtx?.includeAdvancedMetrics ?? false - const tracer = trace.getTracer('text-analysis-tool', '1.0.0') - const span = tracer.startSpan('text-analysis', { - attributes: { + // Create child span for text analysis + const textAnalysisSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'text-analysis', + input: inputData, + metadata: { 'tool.id': 'text-analysis', 'tool.input.textLength': inputData.text.length, - 'tool.input.operations': inputData.operations.join(','), + 'tool.input.operationsCount': inputData.operations.length, + 'user.id': requestCtx?.userId, }, }) @@ -143,11 +149,15 @@ export const textAnalysisTool = createTool({ id: 'text-analysis', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.operationsCount': inputData.operations.length, + // Update span with successful result + textAnalysisSpan?.update({ + output: { success: true, operationsCount: inputData.operations.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.operationsCount': inputData.operations.length, + }, }) - span.end() + textAnalysisSpan?.end() return { success: true, @@ -159,11 +169,11 @@ export const textAnalysisTool = createTool({ const errorMsg = e instanceof Error ? e.message : String(e) log.error(`Text analysis failed: ${errorMsg}`, { error: errorMsg }) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + // Record error in span + textAnalysisSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { success: false, @@ -259,13 +269,20 @@ export const textProcessingTool = createTool({ }), execute: async (inputData, context) => { const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as TextAnalysisToolContext | undefined - const tracer = trace.getTracer('text-processing-tool', '1.0.0') - const span = tracer.startSpan('text-processing', { - attributes: { + // Create child span for text processing + const textProcessingSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'text-processing', + input: inputData, + metadata: { 'tool.id': 'text-processing', 'tool.input.textLength': inputData.text.length, - 'tool.input.operations': inputData.operations.join(','), + 'tool.input.operationsCount': inputData.operations.length, + 'user.id': requestCtx?.userId, }, }) @@ -346,12 +363,16 @@ export const textProcessingTool = createTool({ id: 'text-processing', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.operationsCount': inputData.operations.length, - 'tool.output.lengthChange': statistics.changes, + // Update span with successful result + textProcessingSpan?.update({ + output: { success: true, operationsCount: inputData.operations.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.operationsCount': inputData.operations.length, + 'tool.output.lengthChange': statistics.changes, + }, }) - span.end() + textProcessingSpan?.end() return { success: true, @@ -367,11 +388,11 @@ export const textProcessingTool = createTool({ error: errorMsg, }) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + // Record error in span + textProcessingSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { success: false, diff --git a/src/mastra/tools/url-tool.ts b/src/mastra/tools/url-tool.ts index 1c25085f..aecc1b55 100644 --- a/src/mastra/tools/url-tool.ts +++ b/src/mastra/tools/url-tool.ts @@ -1,15 +1,15 @@ +import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import { z } from 'zod' import { log } from '../config/logger' -import type { RequestContext } from '@mastra/core/request-context' - export interface UrlToolContext extends RequestContext { defaultProtocol?: string allowLocalhost?: boolean timeout?: number + userId?: string } export const urlValidationTool = createTool({ @@ -71,6 +71,8 @@ export const urlValidationTool = createTool({ }), execute: async (inputData, context) => { const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext const requestCtx = context?.requestContext as UrlToolContext | undefined @@ -78,8 +80,18 @@ export const urlValidationTool = createTool({ const allowLocalhost = requestCtx?.allowLocalhost ?? false const timeout = requestCtx?.timeout ?? 5000 - const tracer = trace.getTracer('url-validation-tool', '1.0.0') - const span = tracer.startSpan('url-validation') + // Create child span for URL validation + const urlValidationSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'url-validation', + input: inputData, + metadata: { + 'tool.id': 'url-validation', + 'tool.input.url': inputData.url, + 'tool.input.operationsCount': inputData.operations.length, + 'user.id': requestCtx?.userId, + }, + }) await writer?.custom({ type: 'data-tool-progress', @@ -172,11 +184,15 @@ export const urlValidationTool = createTool({ id: 'url-validation', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.operationsCount': inputData.operations.length, + // Update span with successful result + urlValidationSpan?.update({ + output: { success: true, operationsCount: inputData.operations.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.operationsCount': inputData.operations.length, + }, }) - span.end() + urlValidationSpan?.end() return { success: true, @@ -188,11 +204,11 @@ export const urlValidationTool = createTool({ const errorMsg = e instanceof Error ? e.message : String(e) log.error(`URL validation failed: ${errorMsg}`, { error: errorMsg }) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + // Record error in span + urlValidationSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { success: false, @@ -292,9 +308,22 @@ export const urlManipulationTool = createTool({ }), execute: async (inputData, context) => { const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as UrlToolContext | undefined - const tracer = trace.getTracer('url-manipulation-tool', '1.0.0') - const span = tracer.startSpan('url-manipulation') + // Create child span for URL manipulation + const urlManipulationSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'url-manipulation', + input: inputData, + metadata: { + 'tool.id': 'url-manipulation', + 'tool.input.baseUrl': inputData.baseUrl, + 'tool.input.operationsCount': inputData.operations.length, + 'user.id': requestCtx?.userId, + }, + }) await writer?.custom({ type: 'data-tool-progress', @@ -328,7 +357,7 @@ export const urlManipulationTool = createTool({ currentUrl, inputData.parameters.query as Record< string, - any + string | number | boolean > ) } @@ -353,7 +382,7 @@ export const urlManipulationTool = createTool({ currentUrl, inputData.parameters.query as Record< string, - any + string | number | boolean | null | undefined > ) } @@ -455,11 +484,15 @@ export const urlManipulationTool = createTool({ id: 'url-manipulation', }) - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.operationsCount': inputData.operations.length, + // Update span with successful result + urlManipulationSpan?.update({ + output: { success: true, operationsCount: inputData.operations.length }, + metadata: { + 'tool.output.success': true, + 'tool.output.operationsCount': inputData.operations.length, + }, }) - span.end() + urlManipulationSpan?.end() return { success: true, @@ -474,11 +507,11 @@ export const urlManipulationTool = createTool({ error: errorMsg, }) - if (e instanceof Error) { - span.recordException(e) - } - span.setStatus({ code: 2, message: errorMsg }) - span.end() + // Record error in span + urlManipulationSpan?.error({ + error: e instanceof Error ? e : new Error(errorMsg), + endSpan: true, + }) return { success: false, @@ -624,7 +657,7 @@ async function checkUrlReachability( const controller = new AbortController() const timeoutId = setTimeout(() => controller.abort(), timeout) - const headers = userAgent ? { 'User-Agent': userAgent } : undefined + const headers = userAgent && userAgent.length > 0 ? { 'User-Agent': userAgent } : undefined const response = await fetch(url, { method: 'HEAD', signal: controller.signal, @@ -646,7 +679,7 @@ async function getUrlMetadata(url: string, timeout = 5000, userAgent?: string) { const response = await fetch(url, { method: 'HEAD', signal: controller.signal, - headers: userAgent ? { 'User-Agent': userAgent } : undefined, + headers: userAgent && userAgent.length > 0 ? { 'User-Agent': userAgent } : undefined, }) clearTimeout(timeoutId) diff --git a/src/mastra/tools/weather-tool.ts b/src/mastra/tools/weather-tool.ts index c74d3f6e..5f1a2e8a 100644 --- a/src/mastra/tools/weather-tool.ts +++ b/src/mastra/tools/weather-tool.ts @@ -1,24 +1,28 @@ import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import { z } from 'zod' import { log } from '../config/logger' export interface WeatherToolContext extends RequestContext { temperatureUnit?: 'celsius' | 'fahrenheit' + userId?: string + workspaceId?: string + maxRows?: number } export type WeatherToolContextType = WeatherToolContext -interface GeocodingResponse extends RequestContext { +interface GeocodingResponse { results?: Array<{ latitude: number longitude: number name: string }> } -interface WeatherResponse extends RequestContext { + +interface WeatherResponse { current: { time: string temperature_2m: number @@ -50,9 +54,10 @@ export const weatherTool = createTool({ execute: async (inputData, context) => { const writer = context?.writer const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext // Check if operation was already cancelled - if (abortSignal?.aborted === true) { + if (abortSignal?.aborted ?? false) { throw new Error('Weather lookup cancelled') } @@ -68,18 +73,25 @@ export const weatherTool = createTool({ const requestCtx = context?.requestContext as WeatherToolContext | undefined const temperatureUnit = requestCtx?.temperatureUnit ?? 'celsius' + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId log.info( - `Fetching weather for location: ${inputData.location} in ${temperatureUnit}` + `Fetching weather for location: ${inputData.location} in ${temperatureUnit}`, + { userId, workspaceId } ) - // Get tracer from OpenTelemetry API - const tracer = trace.getTracer('weather-tool', '1.0.0') - const weatherSpan = tracer.startSpan('get-weather', { - attributes: { + // Create child span for weather lookup using Mastra tracing context + const weatherSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'weather-lookup', + input: { location: inputData.location, temperatureUnit }, + metadata: { 'tool.id': 'get-weather', 'tool.input.location': inputData.location, 'tool.input.temperatureUnit': temperatureUnit, + 'user.id': userId, + 'workspace.id': workspaceId, }, }) @@ -95,12 +107,11 @@ export const weatherTool = createTool({ }) // Check for cancellation before geocoding - if (abortSignal?.aborted) { - weatherSpan.setStatus({ - code: 2, - message: 'Operation cancelled during geocoding', + if (abortSignal?.aborted ?? false) { + weatherSpan?.error({ + error: new Error('Operation cancelled during geocoding'), + endSpan: true, }) - weatherSpan.end() throw new Error('Weather lookup cancelled during geocoding') } @@ -120,12 +131,17 @@ export const weatherTool = createTool({ id: 'get-weather', }) - weatherSpan.setAttributes({ - 'tool.output.location': result.location, - 'tool.output.temperature': result.temperature, - 'tool.output.conditions': result.conditions, + // Update span with successful result + weatherSpan?.update({ + output: result, + metadata: { + 'tool.output.location': result.location, + 'tool.output.temperature': result.temperature, + 'tool.output.conditions': result.conditions, + 'tool.output.success': true, + }, }) - weatherSpan.end() + weatherSpan?.end() log.info(`Weather fetched successfully for ${inputData.location}`) const finalResult = { @@ -151,8 +167,10 @@ export const weatherTool = createTool({ // Handle AbortError specifically if (error instanceof Error && error.name === 'AbortError') { const cancelMessage = `Weather lookup cancelled for ${inputData.location}` - weatherSpan.setStatus({ code: 2, message: cancelMessage }) - weatherSpan.end() + weatherSpan?.error({ + error: new Error(cancelMessage), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', @@ -168,11 +186,11 @@ export const weatherTool = createTool({ throw new Error(cancelMessage) } - if (error instanceof Error) { - weatherSpan.recordException(error) - } - weatherSpan.setStatus({ code: 2, message: errorMessage }) // ERROR status - weatherSpan.end() + // Record error in span + weatherSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) await writer?.custom({ type: 'data-tool-progress', diff --git a/src/mastra/tools/web-scraper-tool.ts b/src/mastra/tools/web-scraper-tool.ts index 22745c33..5798a762 100644 --- a/src/mastra/tools/web-scraper-tool.ts +++ b/src/mastra/tools/web-scraper-tool.ts @@ -14,9 +14,9 @@ // approvedBy: sam // approvalDate: 9/22 +import { SpanType } from '@mastra/core/observability' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' import * as cheerio from 'cheerio' import { CheerioCrawler, Request } from 'crawlee' import { JSDOM } from 'jsdom' @@ -30,9 +30,11 @@ import { log } from '../config/logger' import type { RequestContext } from '@mastra/core/request-context' export interface ScraperToolContext extends RequestContext { - allowedDomains?: string[] - userAgent?: string - timeout?: number + allowedDomains?: string[] + userAgent?: string + timeout?: number + userId?: string + workspaceId?: string } // Centralized data directory constant for consistency @@ -42,607 +44,607 @@ const DATA_DIR = path.resolve(process.cwd(), './data') * Ensures the data directory exists and is accessible */ async function ensureDataDir(): Promise<void> { - try { - await fs.mkdir(DATA_DIR, { recursive: true }) - } catch (error) { - log.error('Failed to create data directory', { - error: error instanceof Error ? error.message : String(error), - path: DATA_DIR, - }) - throw new ScrapingError( - 'Failed to create data directory', - 'DATA_DIR_ERROR' - ) - } + try { + await fs.mkdir(DATA_DIR, { recursive: true }) + } catch (error) { + log.error('Failed to create data directory', { + error: error instanceof Error ? error.message : String(error), + path: DATA_DIR, + }) + throw new ScrapingError( + 'Failed to create data directory', + 'DATA_DIR_ERROR' + ) + } } // Enhanced HTML processing with JSDOM const DANGEROUS_TAGS = new Set([ - 'script', - 'style', - 'iframe', - 'embed', - 'object', - 'noscript', - 'meta', - 'link', - 'form', - 'input', - 'button', - 'select', - 'textarea', - 'frame', - 'frameset', + 'script', + 'style', + 'iframe', + 'embed', + 'object', + 'noscript', + 'meta', + 'link', + 'form', + 'input', + 'button', + 'select', + 'textarea', + 'frame', + 'frameset', ]) const DANGEROUS_ATTRS = new Set([ - 'onload', - 'onerror', - 'onclick', - 'onmouseover', - 'onmouseout', - 'onkeydown', - 'onkeyup', - 'onkeypress', - 'onfocus', - 'onblur', - 'formaction', + 'onload', + 'onerror', + 'onclick', + 'onmouseover', + 'onmouseout', + 'onkeydown', + 'onkeyup', + 'onkeypress', + 'onfocus', + 'onblur', + 'formaction', ]) function sanitizeHtml(html: string): string { - try { - // Ensure input is a string and explicitly treat as HTML to avoid ambiguous parsing - const safeHtml = String(html) - const jsdom = new JSDOM(safeHtml, { - contentType: 'text/html', - includeNodeLocations: false, - }) - const { document } = jsdom.window + try { + // Ensure input is a string and explicitly treat as HTML to avoid ambiguous parsing + const safeHtml = String(html) + const jsdom = new JSDOM(safeHtml, { + contentType: 'text/html', + includeNodeLocations: false, + }) + const { document } = jsdom.window - // Remove dangerous elements - DANGEROUS_TAGS.forEach((tagName) => { - const elements = document.querySelectorAll(tagName) - elements.forEach((element) => element.remove()) - }) + // Remove dangerous elements + DANGEROUS_TAGS.forEach((tagName) => { + const elements = document.querySelectorAll(tagName) + elements.forEach((element) => element.remove()) + }) - // Remove event handler attributes - const allElements = document.querySelectorAll('*') - allElements.forEach((element) => { - Array.from(element.attributes).forEach((attr) => { - if ( - attr.name.startsWith('on') || - DANGEROUS_ATTRS.has(attr.name.toLowerCase()) - ) { - element.removeAttribute(attr.name) - } - }) - }) + // Remove event handler attributes + const allElements = document.querySelectorAll('*') + allElements.forEach((element) => { + Array.from(element.attributes).forEach((attr) => { + if ( + attr.name.startsWith('on') || + DANGEROUS_ATTRS.has(attr.name.toLowerCase()) + ) { + element.removeAttribute(attr.name) + } + }) + }) - return document.body.innerHTML - } catch (error) { - log.warn('JSDOM sanitization failed, falling back to cheerio', { - error, - }) - // Pre-sanitize the raw HTML to avoid passing dangerous content directly into cheerio. - // Remove dangerous elements and inline event handlers, and neutralize javascript: hrefs. - const preSanitizedHtml = String(html) - // Remove dangerous elements with their content: script, style, iframe, embed, object, noscript, meta, link, form, input, button, select, textarea, frame, frameset - .replace( - /<\s*(script|style|iframe|embed|object|noscript|meta|link|form|input|button|select|textarea|frame|frameset)\b[^>]*>[\s\S]*?<\/\s*\1\s*>/gi, - '' - ) - // Remove self-closing dangerous tags (meta, link, input), but keep img tags intact - .replace(/<\s*(meta|link|input|br|hr)[^>]*\/?>/gi, '') - // Remove inline event handler attributes like onclick="..." and onmouseover='...' - .replace(/ on\w+\s*=\s*(?:"[^"]*"|'[^']*'|[^\s>]+)/gi, ' ') - // Neutralize javascript: URIs in href attributes - .replace(/href\s*=\s*("|\\')\s*javascript:[^"\\']*\1/gi, 'href="#"') - - const $ = cheerio.load(preSanitizedHtml, { xmlMode: false }) - DANGEROUS_TAGS.forEach((tag) => $(tag).remove()) - $('*').each((_i, element) => { - const el = $(element) - const attrs = el.attr() - if (attrs) { - Object.keys(attrs).forEach((attr) => { - if ( - DANGEROUS_ATTRS.has(attr.toLowerCase()) || - attr.toLowerCase().startsWith('on') - ) { - el.removeAttr(attr) - } - }) - } - }) - return $.html() - } + return document.body.innerHTML + } catch (error) { + log.warn('JSDOM sanitization failed, falling back to cheerio', { + error, + }) + // Pre-sanitize the raw HTML to avoid passing dangerous content directly into cheerio. + // Remove dangerous elements and inline event handlers, and neutralize javascript: hrefs. + const preSanitizedHtml = String(html) + // Remove dangerous elements with their content: script, style, iframe, embed, object, noscript, meta, link, form, input, button, select, textarea, frame, frameset + .replace( + /<\s*(script|style|iframe|embed|object|noscript|meta|link|form|input|button|select|textarea|frame|frameset)\b[^>]*>[\s\S]*?<\/\s*\1\s*>/gi, + '' + ) + // Remove self-closing dangerous tags (meta, link, input), but keep img tags intact + .replace(/<\s*(meta|link|input|br|hr)[^>]*\/?>/gi, '') + // Remove inline event handler attributes like onclick="..." and onmouseover='...' + .replace(/ on\w+\s*=\s*(?:"[^"]*"|'[^']*'|[^\s>]+)/gi, ' ') + // Neutralize javascript: URIs in href attributes + .replace(/href\s*=\s*("|\\')\s*javascript:[^"\\']*\1/gi, 'href="#"') + + const $ = cheerio.load(preSanitizedHtml, { xmlMode: false }) + DANGEROUS_TAGS.forEach((tag) => $(tag).remove()) + $('*').each((_i, element) => { + const el = $(element) + const attrs = el.attr() + if (attrs) { + Object.keys(attrs).forEach((attr) => { + if ( + DANGEROUS_ATTRS.has(attr.toLowerCase()) || + attr.toLowerCase().startsWith('on') + ) { + el.removeAttr(attr) + } + }) + } + }) + return $.html() + } } function extractTextContent(html: string): string { - try { - const dom = new JSDOM(html, { includeNodeLocations: false }) - return dom.window.document.body.textContent?.trim() ?? '' - } catch (error) { - log.warn('JSDOM text extraction failed, falling back to cheerio', { - error, - }) - const $ = cheerio.load(html) - return $.text().trim() - } + try { + const dom = new JSDOM(html, { includeNodeLocations: false }) + return dom.window.document.body.textContent?.trim() ?? '' + } catch (error) { + log.warn('JSDOM text extraction failed, falling back to cheerio', { + error, + }) + const $ = cheerio.load(html) + return $.text().trim() + } } function htmlToMarkdown(html: string): string { - try { - const sanitizedHtml = sanitizeHtml(html) - const dom = new JSDOM(sanitizedHtml, { includeNodeLocations: true }) - const { document } = dom.window - - const convertNode = (node: Node): string => { - if (node.nodeType === dom.window.Node.TEXT_NODE) { - return node.textContent?.trim() ?? '' - } - - if (node.nodeType === dom.window.Node.ELEMENT_NODE) { - const element = node as Element - const tagName = element.tagName.toLowerCase() - switch (tagName) { - case 'h1': - return `# ${getTextContent(element)}\n\n` - case 'h2': - return `## ${getTextContent(element)}\n\n` - case 'h3': - return `### ${getTextContent(element)}\n\n` - case 'h4': - return `#### ${getTextContent(element)}\n\n` - case 'h5': - return `##### ${getTextContent(element)}\n\n` - case 'h6': - return `###### ${getTextContent(element)}\n\n` - case 'p': - return `${getTextContent(element)}\n\n` - case 'br': - return '\n' - case 'strong': - case 'b': - return `**${getTextContent(element)}**` - case 'em': - case 'i': - return `*${getTextContent(element)}*` - case 'code': - return `\`${getTextContent(element)}\`` - case 'pre': - return `\`\`\`\n${getTextContent(element).trim()}\n\`\`\`\n\n` - case 'a': { - const href = element.getAttribute('href') - const text = getTextContent(element) - return href !== null ? `[${text}](${href})` : text - } - case 'ul': - return `${convertChildren(element)}\n` - case 'ol': - return `${convertChildren(element)}\n` - case 'li': - return `- ${getTextContent(element)}\n` - case 'blockquote': - return `> ${getTextContent(element)}\n\n` - case 'hr': - return '---\n\n' - case 'img': { - // Normalize src/alt and skip images without a usable src - const srcAttr = element.getAttribute('src') - const altAttr = element.getAttribute('alt') ?? '' - const trimmedSrc = (srcAttr ?? '').trim() - if (trimmedSrc === '') { - return '' - } - const alt = altAttr.trim() - return `![${alt}](${trimmedSrc})` - } - case 'table': - return convertTable(element) - default: - return convertChildren(element) - } + try { + const sanitizedHtml = sanitizeHtml(html) + const dom = new JSDOM(sanitizedHtml, { includeNodeLocations: true }) + const { document } = dom.window + + const convertNode = (node: Node): string => { + if (node.nodeType === dom.window.Node.TEXT_NODE) { + return node.textContent?.trim() ?? '' + } + + if (node.nodeType === dom.window.Node.ELEMENT_NODE) { + const element = node as Element + const tagName = element.tagName.toLowerCase() + switch (tagName) { + case 'h1': + return `# ${getTextContent(element)}\n\n` + case 'h2': + return `## ${getTextContent(element)}\n\n` + case 'h3': + return `### ${getTextContent(element)}\n\n` + case 'h4': + return `#### ${getTextContent(element)}\n\n` + case 'h5': + return `##### ${getTextContent(element)}\n\n` + case 'h6': + return `###### ${getTextContent(element)}\n\n` + case 'p': + return `${getTextContent(element)}\n\n` + case 'br': + return '\n' + case 'strong': + case 'b': + return `**${getTextContent(element)}**` + case 'em': + case 'i': + return `*${getTextContent(element)}*` + case 'code': + return `\`${getTextContent(element)}\`` + case 'pre': + return `\`\`\`\n${getTextContent(element).trim()}\n\`\`\`\n\n` + case 'a': { + const href = element.getAttribute('href') + const text = getTextContent(element) + return href !== null ? `[${text}](${href})` : text + } + case 'ul': + return `${convertChildren(element)}\n` + case 'ol': + return `${convertChildren(element)}\n` + case 'li': + return `- ${getTextContent(element)}\n` + case 'blockquote': + return `> ${getTextContent(element)}\n\n` + case 'hr': + return '---\n\n' + case 'img': { + // Normalize src/alt and skip images without a usable src + const srcAttr = element.getAttribute('src') + const altAttr = element.getAttribute('alt') ?? '' + const trimmedSrc = (srcAttr ?? '').trim() + if (trimmedSrc === '') { + return '' } - return '' + const alt = altAttr.trim() + return `![${alt}](${trimmedSrc})` + } + case 'table': + return convertTable(element) + default: + return convertChildren(element) } + } + return '' + } - const convertChildren = (element: Element): string => { - let result = '' - element.childNodes.forEach((child) => { - result += convertNode(child) - }) - return result - } + const convertChildren = (element: Element): string => { + let result = '' + element.childNodes.forEach((child) => { + result += convertNode(child) + }) + return result + } - const getTextContent = (element: Element): string => { - let result = '' - element.childNodes.forEach((child) => { - result += convertNode(child) - }) - return result - } + const getTextContent = (element: Element): string => { + let result = '' + element.childNodes.forEach((child) => { + result += convertNode(child) + }) + return result + } - const convertTable = (table: Element): string => { - const rows = Array.from(table.querySelectorAll('tr')) - if (rows.length === 0) { - return '' - } + const convertTable = (table: Element): string => { + const rows = Array.from(table.querySelectorAll('tr')) + if (rows.length === 0) { + return '' + } - let markdown = '' - rows.forEach((row, index) => { - const cells = Array.from(row.querySelectorAll('td, th')) - const cellTexts = cells.map((cell) => getTextContent(cell)) - markdown += '| ' + cellTexts.join(' | ') + ' |\n' + let markdown = '' + rows.forEach((row, index) => { + const cells = Array.from(row.querySelectorAll('td, th')) + const cellTexts = cells.map((cell) => getTextContent(cell)) + markdown += '| ' + cellTexts.join(' | ') + ' |\n' - if (index === 0) { - markdown += - '| ' + cellTexts.map(() => '---').join(' | ') + ' |\n' - } - }) - return markdown + '\n' + if (index === 0) { + markdown += + '| ' + cellTexts.map(() => '---').join(' | ') + ' |\n' } - - return convertNode(document.body).trim() - } catch (error) { - log.warn('JSDOM conversion failed, using marked fallback', { error }) - return extractTextContent(html) + }) + return markdown + '\n' } + + return convertNode(document.body).trim() + } catch (error) { + log.warn('JSDOM conversion failed, using marked fallback', { error }) + return extractTextContent(html) + } } function sanitizeMarkdown(markdown: string): string { - try { - // Comprehensive HTML entity escaping to prevent XSS and parsing issues - // Escape dangerous characters that could interfere with JSON or HTML rendering - return String(markdown) - .replace(/&/g, '&') // Must be first to avoid double-escaping - .replace(/</g, '<') - .replace(/>/g, '>') - .replace(/'/g, ''') - .replace(/"/g, '"') - .replace(/\\/g, '\') // Escape backslashes - .replace(/\n/g, '\\n') // Escape newlines for JSON safety - .replace(/\r/g, '\\r') // Escape carriage returns - .replace(/\t/g, '\\t') // Escape tabs - .replace(/\f/g, '\\f') // Escape form feeds - .replace(/\0/g, '\\0') // Escape null characters - } catch { - return '' - } + try { + // Comprehensive HTML entity escaping to prevent XSS and parsing issues + // Escape dangerous characters that could interfere with JSON or HTML rendering + return String(markdown) + .replace(/&/g, '&') // Must be first to avoid double-escaping + .replace(/</g, '<') + .replace(/>/g, '>') + .replace(/'/g, ''') + .replace(/"/g, '"') + .replace(/\\/g, '\') // Escape backslashes + .replace(/\n/g, '\\n') // Escape newlines for JSON safety + .replace(/\r/g, '\\r') // Escape carriage returns + .replace(/\t/g, '\\t') // Escape tabs + .replace(/\f/g, '\\f') // Escape form feeds + .replace(/\0/g, '\\0') // Escape null characters + } catch { + return '' + } } export const HtmlProcessor = { - DANGEROUS_TAGS, - DANGEROUS_ATTRS, - sanitizeHtml, - extractTextContent, - htmlToMarkdown, - sanitizeMarkdown, + DANGEROUS_TAGS, + DANGEROUS_ATTRS, + sanitizeHtml, + extractTextContent, + htmlToMarkdown, + sanitizeMarkdown, } // Enhanced error handling utilities export class ScrapingError extends Error { - public readonly code: string - public readonly statusCode?: number - public readonly url?: string - - constructor( - message: string, - code: string, - statusCode?: number, - url?: string - ) { - super(message) - this.code = code - this.statusCode = statusCode - this.url = url - // Include the error code in the message so the property is referenced and available for logs - this.message = `${message}${this.code ? ` (code=${this.code})` : ''}` - this.name = 'ScrapingError' - } + public readonly code: string + public readonly statusCode?: number + public readonly url?: string + + constructor( + message: string, + code: string, + statusCode?: number, + url?: string + ) { + super(message) + this.code = code + this.statusCode = statusCode + this.url = url + // Include the error code in the message so the property is referenced and available for logs + this.message = `${message}${this.code ? ` (code=${this.code})` : ''}` + this.name = 'ScrapingError' + } } export class ValidationUtils { - static getAllowedDomains(): string[] { - const raw = process.env.WEB_SCRAPER_ALLOWED_DOMAINS ?? '' - return raw - .split(',') - .map((d) => d.trim().toLowerCase()) - .filter((d) => d.length > 0) + static getAllowedDomains(): string[] { + const raw = process.env.WEB_SCRAPER_ALLOWED_DOMAINS ?? '' + return raw + .split(',') + .map((d) => d.trim().toLowerCase()) + .filter((d) => d.length > 0) + } + + static validateUrl(url: string): boolean { + try { + const parsedUrl = new URL(url) + return ['http:', 'https:'].includes(parsedUrl.protocol) + } catch { + return false } + } - static validateUrl(url: string): boolean { - try { - const parsedUrl = new URL(url) - return ['http:', 'https:'].includes(parsedUrl.protocol) - } catch { - return false - } + static isUrlAllowed(url: string, allowedDomains: string[]): boolean { + if (allowedDomains.length === 0) { + return true } - - static isUrlAllowed(url: string, allowedDomains: string[]): boolean { - if (allowedDomains.length === 0) { - return true - } - try { - const parsedUrl = new URL(url) - const hostname = parsedUrl.hostname.toLowerCase() - return allowedDomains.some( - (d) => hostname === d || hostname.endsWith(`.${d}`) - ) - } catch { - return false - } + try { + const parsedUrl = new URL(url) + const hostname = parsedUrl.hostname.toLowerCase() + return allowedDomains.some( + (d) => hostname === d || hostname.endsWith(`.${d}`) + ) + } catch { + return false } + } - static sanitizeFileName(fileName: string): string { - return fileName.replace(/[^a-zA-Z0-9\-_.]/g, '_') - } + static sanitizeFileName(fileName: string): string { + return fileName.replace(/[^a-zA-Z0-9\-_.]/g, '_') + } - static validateFilePath(filePath: string, intendedDir: string): boolean { - const resolvedPath = path.resolve(filePath) - const resolvedIntendedDir = path.resolve(intendedDir) - return resolvedPath.startsWith(resolvedIntendedDir) - } + static validateFilePath(filePath: string, intendedDir: string): boolean { + const resolvedPath = path.resolve(filePath) + const resolvedIntendedDir = path.resolve(intendedDir) + return resolvedPath.startsWith(resolvedIntendedDir) + } } // Input Schema const crawlingOptionsSchema = z - .object({ - depth: z - .number() - .min(1) - .max(5) - .optional() - .describe( - 'Maximum crawling depth for following links (default: 1, max: 5).' - ), - maxPages: z - .number() - .min(1) - .max(100) - .optional() - .describe( - 'Maximum number of pages to crawl (default: 1, max: 100).' - ), - followLinks: z - .boolean() - .optional() - .describe( - 'Whether to follow and crawl internal links (default: false).' - ), - respectRobotsTxt: z - .boolean() - .optional() - .describe( - 'Whether to check and respect robots.txt (default: true).' - ), - delayBetweenRequests: z - .number() - .min(0) - .max(10000) - .optional() - .describe( - 'Delay between requests in milliseconds for rate limiting (default: 1000).' - ), - }) - .optional() + .object({ + depth: z + .number() + .min(1) + .max(5) + .optional() + .describe( + 'Maximum crawling depth for following links (default: 1, max: 5).' + ), + maxPages: z + .number() + .min(1) + .max(100) + .optional() + .describe( + 'Maximum number of pages to crawl (default: 1, max: 100).' + ), + followLinks: z + .boolean() + .optional() + .describe( + 'Whether to follow and crawl internal links (default: false).' + ), + respectRobotsTxt: z + .boolean() + .optional() + .describe( + 'Whether to check and respect robots.txt (default: true).' + ), + delayBetweenRequests: z + .number() + .min(0) + .max(10000) + .optional() + .describe( + 'Delay between requests in milliseconds for rate limiting (default: 1000).' + ), + }) + .optional() const extractionOptionsSchema = z - .object({ - includeImages: z - .boolean() - .optional() - .describe( - 'Whether to extract image URLs and metadata (default: false).' - ), - extractMetadata: z - .boolean() - .optional() - .describe( - 'Whether to extract page metadata (title, description, keywords, etc.) (default: true).' - ), - contentType: z - .enum(['text', 'links', 'images', 'metadata', 'structured', 'all']) - .optional() - .describe( - 'Type of content to prioritize extracting (default: "all").' - ), - extractStructuredData: z - .boolean() - .optional() - .describe( - 'Whether to extract structured data (JSON-LD, microdata, RDFa) (default: false).' - ), - languageDetection: z - .boolean() - .optional() - .describe( - 'Whether to detect and return content language (default: false).' - ), - contentFiltering: z - .object({ - minLength: z.number().optional(), - maxLength: z.number().optional(), - keywords: z.array(z.string()).optional(), - excludePatterns: z.array(z.string()).optional(), - }) - .optional() - .describe( - 'Content filtering options (min/max length, keywords, exclude patterns).' - ), - }) - .optional() + .object({ + includeImages: z + .boolean() + .optional() + .describe( + 'Whether to extract image URLs and metadata (default: false).' + ), + extractMetadata: z + .boolean() + .optional() + .describe( + 'Whether to extract page metadata (title, description, keywords, etc.) (default: true).' + ), + contentType: z + .enum(['text', 'links', 'images', 'metadata', 'structured', 'all']) + .optional() + .describe( + 'Type of content to prioritize extracting (default: "all").' + ), + extractStructuredData: z + .boolean() + .optional() + .describe( + 'Whether to extract structured data (JSON-LD, microdata, RDFa) (default: false).' + ), + languageDetection: z + .boolean() + .optional() + .describe( + 'Whether to detect and return content language (default: false).' + ), + contentFiltering: z + .object({ + minLength: z.number().optional(), + maxLength: z.number().optional(), + keywords: z.array(z.string()).optional(), + excludePatterns: z.array(z.string()).optional(), + }) + .optional() + .describe( + 'Content filtering options (min/max length, keywords, exclude patterns).' + ), + }) + .optional() const requestOptionsSchema = z - .object({ - timeout: z - .number() - .min(1000) - .max(60000) - .optional() - .describe( - 'Request timeout in milliseconds (default: 30000, max: 60000).' - ), - userAgent: z - .string() - .optional() - .describe('Custom User-Agent string for requests.'), - headers: z - .record(z.string(), z.string()) - .optional() - .describe('Custom HTTP headers to send with requests.'), - retryAttempts: z - .number() - .min(0) - .max(5) - .optional() - .describe( - 'Number of retry attempts for failed requests (default: 2, max: 5).' - ), - compression: z - .boolean() - .optional() - .describe( - 'Whether to handle compressed responses (gzip, deflate) (default: true).' - ), - cookies: z - .record(z.string(), z.string()) - .optional() - .describe('Custom cookies to send with requests.'), - }) - .optional() + .object({ + timeout: z + .number() + .min(1000) + .max(60000) + .optional() + .describe( + 'Request timeout in milliseconds (default: 30000, max: 60000).' + ), + userAgent: z + .string() + .optional() + .describe('Custom User-Agent string for requests.'), + headers: z + .record(z.string(), z.string()) + .optional() + .describe('Custom HTTP headers to send with requests.'), + retryAttempts: z + .number() + .min(0) + .max(5) + .optional() + .describe( + 'Number of retry attempts for failed requests (default: 2, max: 5).' + ), + compression: z + .boolean() + .optional() + .describe( + 'Whether to handle compressed responses (gzip, deflate) (default: true).' + ), + cookies: z + .record(z.string(), z.string()) + .optional() + .describe('Custom cookies to send with requests.'), + }) + .optional() const storageOptionsSchema = z - .object({ - saveMarkdown: z - .boolean() - .optional() - .describe( - 'Whether to save the scraped content as markdown to the data directory.' - ), - markdownFileName: z - .string() - .optional() - .describe( - 'Optional filename for the markdown file (relative to data/ directory).' - ), - outputFormat: z - .enum(['json', 'markdown', 'html', 'text']) - .optional() - .describe('Output format for extracted content (default: "json").'), - }) - .optional() + .object({ + saveMarkdown: z + .boolean() + .optional() + .describe( + 'Whether to save the scraped content as markdown to the data directory.' + ), + markdownFileName: z + .string() + .optional() + .describe( + 'Optional filename for the markdown file (relative to data/ directory).' + ), + outputFormat: z + .enum(['json', 'markdown', 'html', 'text']) + .optional() + .describe('Output format for extracted content (default: "json").'), + }) + .optional() const webScraperInputSchema = z - .object({ - url: z - .url() - .refine( - (v) => ValidationUtils.validateUrl(v), - 'Invalid URL format' - ), - selector: z - .string() - .optional() - .describe( - "CSS selector for the elements to extract (e.g., 'h1', '.product-title'). If not provided, the entire page content will be extracted." - ), - extractAttributes: z - .array(z.string()) - .optional() - .describe( - "Array of HTML attributes to extract from selected elements (e.g., 'href', 'src', 'alt')." - ), - crawling: crawlingOptionsSchema, - extraction: extractionOptionsSchema, - request: requestOptionsSchema, - storage: storageOptionsSchema, - }) - .strict() + .object({ + url: z + .url() + .refine( + (v) => ValidationUtils.validateUrl(v), + 'Invalid URL format' + ), + selector: z + .string() + .optional() + .describe( + "CSS selector for the elements to extract (e.g., 'h1', '.product-title'). If not provided, the entire page content will be extracted." + ), + extractAttributes: z + .array(z.string()) + .optional() + .describe( + "Array of HTML attributes to extract from selected elements (e.g., 'href', 'src', 'alt')." + ), + crawling: crawlingOptionsSchema, + extraction: extractionOptionsSchema, + request: requestOptionsSchema, + storage: storageOptionsSchema, + }) + .strict() // Output Schema const webScraperOutputSchema = z - .object({ - url: z.url().describe('The URL that was scraped.'), - status: z - .string() - .describe( - "Status of the scraping operation (e.g., 'success', 'failed')." - ), - errorMessage: z - .string() - .optional() - .describe('Error message if the operation failed.'), - content: z - .object({ - extractedData: z - .array(z.record(z.string(), z.string())) - .describe( - 'Array of extracted data, where each object represents an element and its extracted attributes/text.' - ), - rawContent: z - .string() - .optional() - .describe( - 'The full raw HTML content of the page (if no selector is provided).' - ), - markdownContent: z - .string() - .optional() - .describe( - 'The scraped content converted to markdown format.' - ), - }) - .describe('Scraped content results.'), - storage: z - .object({ - savedFilePath: z - .string() - .optional() - .describe( - 'Path to the saved markdown file (if saveMarkdown was true).' - ), - }) - .optional() - .describe('Storage information.'), - analysis: z - .object({ - metadata: z - .record(z.string(), z.string().optional()) - .optional() - .describe( - 'Page metadata (title, description, keywords, etc.) extracted from meta tags.' - ), - images: z - .array( - z.object({ - src: z.string(), - alt: z.string().optional(), - title: z.string().optional(), - }) - ) - .optional() - .describe( - 'Array of image URLs and metadata found on the page.' - ), - structuredData: z - .array(z.unknown()) - .optional() - .describe( - 'Structured data extracted from JSON-LD and microdata.' - ), - detectedLanguage: z - .string() - .optional() - .describe( - 'Detected content language from HTML lang attribute or meta tags.' - ), + .object({ + url: z.url().describe('The URL that was scraped.'), + status: z + .string() + .describe( + "Status of the scraping operation (e.g., 'success', 'failed')." + ), + errorMessage: z + .string() + .optional() + .describe('Error message if the operation failed.'), + content: z + .object({ + extractedData: z + .array(z.record(z.string(), z.string())) + .describe( + 'Array of extracted data, where each object represents an element and its extracted attributes/text.' + ), + rawContent: z + .string() + .optional() + .describe( + 'The full raw HTML content of the page (if no selector is provided).' + ), + markdownContent: z + .string() + .optional() + .describe( + 'The scraped content converted to markdown format.' + ), + }) + .describe('Scraped content results.'), + storage: z + .object({ + savedFilePath: z + .string() + .optional() + .describe( + 'Path to the saved markdown file (if saveMarkdown was true).' + ), + }) + .optional() + .describe('Storage information.'), + analysis: z + .object({ + metadata: z + .record(z.string(), z.string().optional()) + .optional() + .describe( + 'Page metadata (title, description, keywords, etc.) extracted from meta tags.' + ), + images: z + .array( + z.object({ + src: z.string(), + alt: z.string().optional(), + title: z.string().optional(), }) - .optional() - .describe('Analysis results.'), - }) - .strict() + ) + .optional() + .describe( + 'Array of image URLs and metadata found on the page.' + ), + structuredData: z + .array(z.unknown()) + .optional() + .describe( + 'Structured data extracted from JSON-LD and microdata.' + ), + detectedLanguage: z + .string() + .optional() + .describe( + 'Detected content language from HTML lang attribute or meta tags.' + ), + }) + .optional() + .describe('Analysis results.'), + }) + .strict() // In-memory counter to track tool calls per request // Add this line at the beginning of each tool's execute function to track usage: @@ -650,952 +652,1168 @@ const webScraperOutputSchema = z const toolCallCounters = new Map<string, number>() export const webScraperTool = createTool({ - id: 'web:scraper', - description: - 'Extracts structured data from web pages using JSDOM and Cheerio with enhanced security and error handling.', - inputSchema: webScraperInputSchema, - outputSchema: webScraperOutputSchema, - execute: async (inputData, context) => { - const requestContext = context?.requestContext as ScraperToolContext - const writer = context?.writer - const abortSignal = context?.abortSignal - - // Check if operation was already cancelled - if (abortSignal?.aborted === true) { - throw new Error('Web scraping cancelled') - } + id: 'web:scraper', + description: + 'Extracts structured data from web pages using JSDOM and Cheerio with enhanced security and error handling.', + inputSchema: webScraperInputSchema, + outputSchema: webScraperOutputSchema, + execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext + const writer = context?.writer + const abortSignal = context?.abortSignal + + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + + const toolSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'web_scrape', + input: inputData, + metadata: { + 'tool.id': 'web-scraper', + 'tool.input.url': inputData.url, + 'tool.input.selector': inputData.selector ?? undefined, + 'tool.input.saveMarkdown': inputData.storage?.saveMarkdown ?? false, + 'tool.input.followLinks': inputData.crawling?.followLinks ?? false, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - const allowedDomains = - requestContext?.allowedDomains ?? ValidationUtils.getAllowedDomains() - if (!ValidationUtils.isUrlAllowed(inputData.url, allowedDomains)) { - throw new ScrapingError( - `Domain is not allowlisted for scraping: ${inputData.url}`, - 'DOMAIN_NOT_ALLOWED', - undefined, - inputData.url - ) - } + // Check if operation was already cancelled + if (abortSignal?.aborted === true) { + throw new Error('Web scraping cancelled') + } - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🌐 Starting web scrape for ${inputData.url}`, - stage: 'web:scraper', - }, - id: 'web:scraper', - }) - toolCallCounters.set( - 'web:scraper', - (toolCallCounters.get('web:scraper') ?? 0) + 1 - ) + const allowedDomains = + requestContext?.allowedDomains ?? ValidationUtils.getAllowedDomains() + if (!ValidationUtils.isUrlAllowed(inputData.url, allowedDomains)) { + throw new ScrapingError( + `Domain is not allowlisted for scraping: ${inputData.url}`, + 'DOMAIN_NOT_ALLOWED', + undefined, + inputData.url + ) + } - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🐛 Initializing crawler...', - stage: 'web:scraper', - }, - id: 'web:scraper', - }) - log.info('Starting enhanced web scraping with JSDOM', { - url: inputData.url, - selector: inputData.selector, - saveMarkdown: inputData.storage?.saveMarkdown, - }) + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🌐 Starting web scrape for ${inputData.url}`, + stage: 'web:scraper', + }, + id: 'web:scraper', + }) + toolCallCounters.set( + 'web:scraper', + (toolCallCounters.get('web:scraper') ?? 0) + 1 + ) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🐛 Initializing crawler...', + stage: 'web:scraper', + }, + id: 'web:scraper', + }) + log.info('Starting enhanced web scraping with JSDOM', { + url: inputData.url, + selector: inputData.selector, + saveMarkdown: inputData.storage?.saveMarkdown, + }) - let rawContent: string | undefined - let markdownContent: string | undefined - let savedFilePath: string | undefined - const extractedData: Array<Record<string, string>> = [] - let status = 'failed' - let errorMessage: string | undefined - let scrapedUrl: string = inputData.url + let rawContent: string | undefined + let markdownContent: string | undefined + let savedFilePath: string | undefined + const extractedData: Array<Record<string, string>> = [] + let status = 'failed' + let errorMessage: string | undefined + let scrapedUrl: string = inputData.url - // Extract additional data based on options - const metadata: Record<string, string | undefined> = {} - const images: Array<{ src: string; alt?: string; title?: string }> = [] - const structuredData: unknown[] = [] - let detectedLanguage: string | undefined + // Extract additional data based on options + const metadata: Record<string, string | undefined> = {} + const images: Array<{ src: string; alt?: string; title?: string }> = [] + const structuredData: unknown[] = [] + let detectedLanguage: string | undefined + + try { + const headers: Record<string, string> = { + ...(inputData.request?.headers ?? {}), + } + if ( + typeof inputData.request?.userAgent === 'string' && + inputData.request.userAgent.trim() !== '' + ) { + headers['user-agent'] = inputData.request.userAgent + } + + const maxDepth = inputData.crawling?.depth ?? 1 + const maxPages = inputData.crawling?.maxPages ?? 1 + const followLinks = inputData.crawling?.followLinks ?? false + const retryAttempts = inputData.request?.retryAttempts ?? 2 + const delayBetweenRequests = + inputData.crawling?.delayBetweenRequests ?? 1000 + + const crawler = new CheerioCrawler({ + maxRequestsPerCrawl: followLinks ? Math.min(maxPages, 50) : 1, + maxConcurrency: 10, + maxRequestRetries: retryAttempts, + sameDomainDelaySecs: delayBetweenRequests / 1000, + requestHandlerTimeoutSecs: + (inputData.request?.timeout ?? 30000) / 1000, + async requestHandler({ + request, + body, + response, + enqueueLinks, + }) { + try { + scrapedUrl = request.url - try { - const headers: Record<string, string> = { - ...(inputData.request?.headers ?? {}), - } if ( - typeof inputData.request?.userAgent === 'string' && - inputData.request.userAgent.trim() !== '' + !ValidationUtils.isUrlAllowed( + request.url, + allowedDomains + ) ) { - headers['user-agent'] = inputData.request.userAgent + throw new ScrapingError( + `Domain is not allowlisted for scraping: ${request.url}`, + 'DOMAIN_NOT_ALLOWED', + undefined, + request.url + ) } - const maxDepth = inputData.crawling?.depth ?? 1 - const maxPages = inputData.crawling?.maxPages ?? 1 - const followLinks = inputData.crawling?.followLinks ?? false - const retryAttempts = inputData.request?.retryAttempts ?? 2 - const delayBetweenRequests = - inputData.crawling?.delayBetweenRequests ?? 1000 - - const crawler = new CheerioCrawler({ - maxRequestsPerCrawl: followLinks ? Math.min(maxPages, 50) : 1, - maxConcurrency: 10, - maxRequestRetries: retryAttempts, - sameDomainDelaySecs: delayBetweenRequests / 1000, - requestHandlerTimeoutSecs: - (inputData.request?.timeout ?? 30000) / 1000, - async requestHandler({ - request, - body, - response, - enqueueLinks, - }) { - try { - scrapedUrl = request.url - - if ( - !ValidationUtils.isUrlAllowed( - request.url, - allowedDomains - ) - ) { - throw new ScrapingError( - `Domain is not allowlisted for scraping: ${request.url}`, - 'DOMAIN_NOT_ALLOWED', - undefined, - request.url - ) - } - - if ( - typeof response?.statusCode === 'number' && - Number.isFinite(response.statusCode) && - response.statusCode >= 400 - ) { - throw new ScrapingError( - `HTTP ${response.statusCode}: ${response.statusMessage}`, - 'HTTP_ERROR', - response.statusCode, - request.url - ) - } - - rawContent = body.toString() - - // Sanitize HTML using JSDOM - rawContent = HtmlProcessor.sanitizeHtml(rawContent) - - if (followLinks) { - const currentDepth = - typeof request.userData?.depth === 'number' - ? (request.userData.depth as number) - : 0 - - if (currentDepth < maxDepth) { - try { - await enqueueLinks({ - selector: 'a[href]', - baseUrl: request.url, - strategy: 'same-hostname', - limit: Math.max( - 1, - Math.min(maxPages, 100) - ), - transformRequestFunction: (req) => { - const nextDepth = currentDepth + 1 - const nextUrl = req.url - if ( - !ValidationUtils.isUrlAllowed( - nextUrl, - allowedDomains - ) - ) { - return null - } - return { - ...req, - headers, - userData: { - ...(req.userData ?? {}), - depth: nextDepth, - }, - } - }, - }) - } catch (error) { - log.warn('Failed to enqueue links', { - error: - error instanceof Error - ? error.message - : String(error), - url: request.url, - }) - } - } - } - - // Extract metadata if requested - if (inputData.extraction?.extractMetadata !== false) { - const dom = new JSDOM(rawContent, { - includeNodeLocations: false, - }) - const { document } = dom.window - - metadata.title = document - .querySelector('title') - ?.textContent?.trim() - metadata.description = - document - .querySelector('meta[name="description"]') - ?.getAttribute('content') ?? undefined - metadata.keywords = - document - .querySelector('meta[name="keywords"]') - ?.getAttribute('content') ?? undefined - metadata.author = - document - .querySelector('meta[name="author"]') - ?.getAttribute('content') ?? undefined - metadata.robots = - document - .querySelector('meta[name="robots"]') - ?.getAttribute('content') ?? undefined - metadata.viewport = - document - .querySelector('meta[name="viewport"]') - ?.getAttribute('content') ?? undefined - metadata.charset = - document - .querySelector('meta[charset]') - ?.getAttribute('charset') ?? undefined - metadata.canonical = - document - .querySelector('link[rel="canonical"]') - ?.getAttribute('href') ?? undefined - } - - // Extract images if requested - if (inputData.extraction?.includeImages ?? false) { - const dom = new JSDOM(rawContent, { - includeNodeLocations: false, - }) - const { document } = dom.window - const imgElements = document.querySelectorAll('img') - imgElements.forEach((img) => { - const src = img.getAttribute('src') - const trimmedSrc = (src ?? '').trim() - if (trimmedSrc !== '') { - images.push({ - src: new URL(trimmedSrc, request.url) - .href, - alt: - img.getAttribute('alt') ?? - undefined, - title: - img.getAttribute('title') ?? - undefined, - }) - } - }) - } + if ( + typeof response?.statusCode === 'number' && + Number.isFinite(response.statusCode) && + response.statusCode >= 400 + ) { + throw new ScrapingError( + `HTTP ${response.statusCode}: ${response.statusMessage}`, + 'HTTP_ERROR', + response.statusCode, + request.url + ) + } - // Extract structured data if requested - if ( - inputData.extraction?.extractStructuredData ?? - false - ) { - const dom = new JSDOM(rawContent, { - includeNodeLocations: false, - }) - const { document } = dom.window - - // JSON-LD - const jsonLdScripts = document.querySelectorAll( - 'script[type="application/ld+json"]' - ) - jsonLdScripts.forEach((script) => { - try { - const data = JSON.parse( - script.textContent || '' - ) - structuredData.push(data) - } catch { - // Skip invalid JSON-LD - } - }) - - // Microdata - const microdataElements = - document.querySelectorAll('[itemscope]') - microdataElements.forEach((element) => { - const itemType = - element.getAttribute('itemtype') - const trimmedItemType = (itemType ?? '').trim() - if (trimmedItemType !== '') { - const microdata: Record<string, unknown> = { - '@type': trimmedItemType, - } - const props = - element.querySelectorAll('[itemprop]') - props.forEach((prop) => { - const propName = - prop.getAttribute('itemprop') - const trimmedPropName = ( - propName ?? '' - ).trim() - const propValueRaw = - prop.textContent?.trim() ?? - prop.getAttribute('content') - const trimmedPropValue = ( - propValueRaw ?? '' - ).trim() - if ( - trimmedPropName !== '' && - trimmedPropValue !== '' - ) { - microdata[trimmedPropName] = - trimmedPropValue - } - }) - if (Object.keys(microdata).length > 1) { - structuredData.push(microdata) - } - } - }) - } + rawContent = body.toString() + + // Sanitize HTML using JSDOM + rawContent = HtmlProcessor.sanitizeHtml(rawContent) + + if (followLinks) { + const currentDepth = + typeof request.userData?.depth === 'number' + ? (request.userData.depth as number) + : 0 + + if (currentDepth < maxDepth) { + const enqueueSpan = toolSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'enqueue_links', + input: { url: request.url, depth: currentDepth + 1 }, + metadata: { + 'tool.id': 'web-scraper', + 'operation.type': 'enqueue_links', + }, + }) + try { + await enqueueLinks({ + selector: 'a[href]', + baseUrl: request.url, + strategy: 'same-hostname', + limit: Math.max( + 1, + Math.min(maxPages, 100) + ), + transformRequestFunction: (req) => { + const nextDepth = currentDepth + 1 + const nextUrl = req.url + if ( + !ValidationUtils.isUrlAllowed( + nextUrl, + allowedDomains + ) + ) { + return null + } + return { + ...req, + headers, + userData: { + ...(req.userData ?? {}), + depth: nextDepth, + }, + } + }, + }) + try { + enqueueSpan?.update({ + metadata: { 'operation.status': 'success' }, + }) + } catch { + // ignore tracing update errors + } + } catch (error) { + enqueueSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + log.warn('Failed to enqueue links', { + error: + error instanceof Error + ? error.message + : String(error), + url: request.url, + }) + } finally { + enqueueSpan?.end() + } + } + } - // Language detection (basic) - if (inputData.extraction?.languageDetection ?? false) { - const dom = new JSDOM(rawContent, { - includeNodeLocations: false, - }) - const { document } = dom.window - detectedLanguage = - document.documentElement.getAttribute('lang') ?? - document - .querySelector( - 'meta[http-equiv="content-language"]' - ) - ?.getAttribute('content') ?? - document - .querySelector('meta[name="language"]') - ?.getAttribute('content') ?? - undefined - } + // Content extraction with tracing span + const contentSpan = toolSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'content_extraction', + input: { url: request.url }, + metadata: { + 'tool.id': 'web-scraper', + 'operation.type': 'content_extraction', + }, + }) + try { + // Extract metadata if requested + if (inputData.extraction?.extractMetadata !== false) { + const dom = new JSDOM(rawContent, { + includeNodeLocations: false, + }) + const { document } = dom.window + + metadata.title = document + .querySelector('title') + ?.textContent?.trim() + metadata.description = + document + .querySelector('meta[name="description"]') + ?.getAttribute('content') ?? undefined + metadata.keywords = + document + .querySelector('meta[name="keywords"]') + ?.getAttribute('content') ?? undefined + metadata.author = + document + .querySelector('meta[name="author"]') + ?.getAttribute('content') ?? undefined + metadata.robots = + document + .querySelector('meta[name="robots"]') + ?.getAttribute('content') ?? undefined + metadata.viewport = + document + .querySelector('meta[name="viewport"]') + ?.getAttribute('content') ?? undefined + metadata.charset = + document + .querySelector('meta[charset]') + ?.getAttribute('charset') ?? undefined + metadata.canonical = + document + .querySelector('link[rel="canonical"]') + ?.getAttribute('href') ?? undefined + } + + // Extract images if requested + if (inputData.extraction?.includeImages ?? false) { + const dom = new JSDOM(rawContent, { + includeNodeLocations: false, + }) + const { document } = dom.window + const imgElements = document.querySelectorAll('img') + imgElements.forEach((img) => { + const src = img.getAttribute('src') + const trimmedSrc = (src ?? '').trim() + if (trimmedSrc !== '') { + images.push({ + src: new URL(trimmedSrc, request.url) + .href, + alt: + img.getAttribute('alt') ?? + undefined, + title: + img.getAttribute('title') ?? + undefined, + }) + } + }) + } + + // Extract structured data if requested + if ( + inputData.extraction?.extractStructuredData ?? + false + ) { + const dom = new JSDOM(rawContent, { + includeNodeLocations: false, + }) + const { document } = dom.window - if ( - typeof inputData.selector === 'string' && - inputData.selector.trim() !== '' - ) { - // Use JSDOM for better DOM manipulation - const dom = new JSDOM(rawContent, { - includeNodeLocations: false, - }) - const { document } = dom.window - - const selector = inputData.selector.trim() - const elements = document.querySelectorAll(selector) - elements.forEach((element) => { - const data = new Map<string, string>() - data.set( - 'text', - element.textContent?.trim() ?? '' - ) - - if (inputData.extractAttributes) { - inputData.extractAttributes.forEach( - (attr) => { - if ( - typeof attr === 'string' && - !Object.prototype.hasOwnProperty.call( - Object.prototype, - attr - ) && - attr !== '__proto__' && - attr !== 'constructor' && - attr !== 'prototype' && - !attr.includes('<') && - !attr.includes('>') - ) { - const attrValue = - element.getAttribute(attr) - if ( - attrValue !== null && - attrValue !== undefined - ) { - data.set( - `attr_${attr}`, - attrValue - ) - } - } - } - ) - } - extractedData.push(Object.fromEntries(data)) - }) - } + // JSON-LD + const jsonLdScripts = document.querySelectorAll( + 'script[type="application/ld+json"]' + ) + jsonLdScripts.forEach((script) => { + try { + const data = JSON.parse( + script.textContent || '' + ) + structuredData.push(data) + } catch { + // Skip invalid JSON-LD + } + }) - status = 'success' - } catch (error) { - if (error instanceof ScrapingError) { - throw error - } - throw new ScrapingError( - `Request handler failed: ${error instanceof Error ? error.message : String(error)}`, - 'REQUEST_HANDLER_ERROR', - undefined, - request.url - ) + // Microdata + const microdataElements = + document.querySelectorAll('[itemscope]') + microdataElements.forEach((element) => { + const itemType = + element.getAttribute('itemtype') + const trimmedItemType = (itemType ?? '').trim() + if (trimmedItemType !== '') { + const microdata: Record<string, unknown> = { + '@type': trimmedItemType, } - }, - failedRequestHandler({ request, error }) { - const scrapingError = new ScrapingError( - `Failed to scrape ${request.url}: ${error instanceof Error ? error.message : String(error)}`, - 'REQUEST_FAILED', - undefined, - request.url - ) - errorMessage = scrapingError.message - log.error(scrapingError.message) - }, - }) + const props = + element.querySelectorAll('[itemprop]') + props.forEach((prop) => { + const propName = + prop.getAttribute('itemprop') + const trimmedPropName = ( + propName ?? '' + ).trim() + const propValueRaw = + prop.textContent?.trim() ?? + prop.getAttribute('content') + const trimmedPropValue = ( + propValueRaw ?? '' + ).trim() + if ( + trimmedPropName !== '' && + trimmedPropValue !== '' + ) { + microdata[trimmedPropName] = + trimmedPropValue + } + }) + if (Object.keys(microdata).length > 1) { + structuredData.push(microdata) + } + } + }) + } - // Check for cancellation before crawling - if (abortSignal?.aborted) { - throw new Error('Web scraping cancelled during crawling') + // Language detection (basic) + if (inputData.extraction?.languageDetection ?? false) { + const dom = new JSDOM(rawContent, { + includeNodeLocations: false, + }) + const { document } = dom.window + detectedLanguage = + document.documentElement.getAttribute('lang') ?? + document + .querySelector( + 'meta[http-equiv="content-language"]' + ) + ?.getAttribute('content') ?? + document + .querySelector('meta[name="language"]') + ?.getAttribute('content') ?? + undefined + } + + // annotate contentSpan and close + try { + contentSpan?.update({ + metadata: { + 'operation.status': 'success', + 'operation.metadataCount': Object.keys(metadata).length, + 'operation.imagesCount': images.length, + 'operation.structuredDataCount': structuredData.length, + 'operation.detectedLanguage': detectedLanguage ?? undefined, + }, + }) + } catch { + // ignore tracing update errors + } finally { + contentSpan?.end() + } + } catch (err) { + contentSpan?.error({ + error: err instanceof Error ? err : new Error(String(err)), + endSpan: true, + }) + throw err } - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '📥 Fetching and parsing page...', - stage: 'web:scraper', - }, - id: 'web:scraper', - }) - await crawler.run([ - new Request({ - url: inputData.url, - headers, - userData: { depth: 0 }, - }), - ]) - - // Enhanced HTML to markdown conversion using JSDOM if ( - typeof rawContent === 'string' && - rawContent.trim().length > 0 + typeof inputData.selector === 'string' && + inputData.selector.trim() !== '' ) { - try { - markdownContent = HtmlProcessor.htmlToMarkdown(rawContent) - markdownContent = - HtmlProcessor.sanitizeMarkdown(markdownContent) - } catch (error) { - log.warn( - 'Enhanced HTML to markdown conversion failed, using fallback', - { - error: - error instanceof Error - ? error.message - : String(error), - } - ) - try { - // Sanitize raw HTML before passing to marked to prevent XSS via raw HTML content. - const sanitizedForMarked = - HtmlProcessor.sanitizeHtml(rawContent) - markdownContent = await marked.parse(sanitizedForMarked) - } catch (fallbackError) { - log.warn('Fallback conversion also failed', { - error: - fallbackError instanceof Error - ? fallbackError.message - : String(fallbackError), - }) - markdownContent = - HtmlProcessor.extractTextContent(rawContent) - } - } - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '✂️ Converting to markdown...', - stage: 'web:scraper', - }, - id: 'web:scraper', - }) - if ( - inputData.storage?.saveMarkdown === true && - typeof markdownContent === 'string' && - markdownContent.trim() !== '' - ) { - try { - const fileName = - typeof inputData.storage?.markdownFileName === - 'string' && - inputData.storage.markdownFileName.trim() !== '' - ? ValidationUtils.sanitizeFileName( - inputData.storage.markdownFileName - ) - : `scraped_${new Date().toISOString().replace(/[:.]/g, '-')}.md` - - await ensureDataDir() - const fullPath = path.join(DATA_DIR, fileName) + // Use JSDOM for better DOM manipulation + const dom = new JSDOM(rawContent, { + includeNodeLocations: false, + }) + const { document } = dom.window + + const selector = inputData.selector.trim() + const elements = document.querySelectorAll(selector) + elements.forEach((element) => { + const data = new Map<string, string>() + data.set( + 'text', + element.textContent?.trim() ?? '' + ) + if (inputData.extractAttributes) { + inputData.extractAttributes.forEach( + (attr) => { + if ( + typeof attr === 'string' && + !Object.prototype.hasOwnProperty.call( + Object.prototype, + attr + ) && + attr !== '__proto__' && + attr !== 'constructor' && + attr !== 'prototype' && + !attr.includes('<') && + !attr.includes('>') + ) { + const attrValue = + element.getAttribute(attr) if ( - !ValidationUtils.validateFilePath( - fullPath, - DATA_DIR - ) + attrValue !== null && + attrValue !== undefined ) { - throw new ScrapingError( - 'Invalid file path', - 'INVALID_FILE_PATH' - ) - } - // Open the file handle and write via the handle to avoid using fs.writeFile with a non-literal argument. - // Path was already validated by ValidationUtils.validateFilePath above. - const fileHandle = await fs.open(fullPath, 'w') - try { - await fileHandle.writeFile(markdownContent, 'utf-8') - } finally { - await fileHandle.close() + data.set( + `attr_${attr}`, + attrValue + ) } - savedFilePath = fileName - log.info('Markdown content saved securely', { - fileName, - }) - } catch (error) { - if (error instanceof ScrapingError) { - throw error - } - log.error('Failed to save markdown file', { - error: - error instanceof Error - ? error.message - : String(error), - }) + } } + ) } + extractedData.push(Object.fromEntries(data)) + }) } - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Scraping complete: ${extractedData.length} elements${typeof savedFilePath === 'string' ? ', saved to ' + savedFilePath : ''}`, - stage: 'web:scraper', - }, - id: 'web:scraper', - }) + status = 'success' + } catch (error) { + if (error instanceof ScrapingError) { + throw error + } + throw new ScrapingError( + `Request handler failed: ${error instanceof Error ? error.message : String(error)}`, + 'REQUEST_HANDLER_ERROR', + undefined, + request.url + ) + } + }, + failedRequestHandler({ request, error }) { + const scrapingError = new ScrapingError( + `Failed to scrape ${request.url}: ${error instanceof Error ? error.message : String(error)}`, + 'REQUEST_FAILED', + undefined, + request.url + ) + errorMessage = scrapingError.message + log.error(scrapingError.message) + }, + }) + + // Check for cancellation before crawling + if (abortSignal?.aborted) { + throw new Error('Web scraping cancelled during crawling') + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '📥 Fetching and parsing page...', + stage: 'web:scraper', + }, + id: 'web:scraper', + }) + const crawlSpan = toolSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'crawler_run', + input: { url: inputData.url, maxDepth, maxPages, followLinks }, + metadata: { + 'tool.id': 'web-scraper', + 'operation.type': 'crawler.run', + }, + }) + try { + await crawler.run([ + new Request({ + url: inputData.url, + headers, + userData: { depth: 0 }, + }), + ]) + crawlSpan?.update({ + metadata: { + 'operation.status': 'success', + 'result.extractedElements': extractedData.length, + }, + }) + } catch (err) { + crawlSpan?.error({ + error: err instanceof Error ? err : new Error(String(err)), + endSpan: true, + }) + throw err + } finally { + crawlSpan?.end() + } + + // Enhanced HTML to markdown conversion using JSDOM + if ( + typeof rawContent === 'string' && + rawContent.trim().length > 0 + ) { + const convertSpan = toolSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'html_to_markdown', + input: { htmlLength: rawContent.length }, + metadata: { 'tool.id': 'web-scraper', 'operation.type': 'convert' }, + }) - return webScraperOutputSchema.parse({ - url: scrapedUrl, - status, - errorMessage, - content: { - extractedData, - rawContent: - typeof inputData.selector === 'string' && - inputData.selector.trim() !== '' - ? undefined - : typeof rawContent === 'string' && - rawContent.trim().length > 0 - ? rawContent - : undefined, - markdownContent, - }, - storage: - typeof savedFilePath === 'string' - ? { savedFilePath } - : undefined, - analysis: { - metadata: - Object.keys(metadata).length > 0 ? metadata : undefined, - images: images.length > 0 ? images : undefined, - structuredData: - structuredData.length > 0 ? structuredData : undefined, - detectedLanguage, - }, + try { + // Primary conversion attempt + markdownContent = HtmlProcessor.htmlToMarkdown(rawContent) + markdownContent = + HtmlProcessor.sanitizeMarkdown(markdownContent) + + try { + convertSpan?.update({ + metadata: { + 'operation.status': 'primary_success', + 'result.markdownLength': markdownContent.length, + }, }) + } catch { + // ignore tracing update errors + } } catch (error) { - // Handle AbortError specifically - if (error instanceof Error && error.name === 'AbortError') { - const cancelMessage = `Web scraping cancelled for ${inputData.url}` - log.warn(cancelMessage) - throw new Error(cancelMessage) + try { + convertSpan?.update({ metadata: { 'operation.status': 'primary_failed' } }) + } catch { + // ignore tracing update errors + } + + log.warn( + 'Enhanced HTML to markdown conversion failed, using fallback', + { + error: + error instanceof Error + ? error.message + : String(error), } + ) + try { + // Sanitize raw HTML before passing to marked to prevent XSS via raw HTML content. + const sanitizedForMarked = + HtmlProcessor.sanitizeHtml(rawContent) + markdownContent = await marked.parse(sanitizedForMarked) - errorMessage = `Web scraping failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - 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', - }) - }, -}) + try { + convertSpan?.update({ + metadata: { + 'operation.status': 'fallback_success', + 'result.markdownLength': markdownContent.length, + }, + }) + } catch { + // ignore tracing update errors + } + } catch (fallbackError) { + log.warn('Fallback conversion also failed', { + error: + fallbackError instanceof Error + ? fallbackError.message + : String(fallbackError), + }) + try { + convertSpan?.error({ + error: fallbackError instanceof Error ? fallbackError : new Error(String(fallbackError)), + endSpan: true, + }) + } catch { + // ignore tracing errors + } + markdownContent = + HtmlProcessor.extractTextContent(rawContent) + } + } finally { + convertSpan?.end() + } + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '✂️ Converting to markdown...', + stage: 'web:scraper', + }, + id: 'web:scraper', + }) + if ( + inputData.storage?.saveMarkdown === true && + typeof markdownContent === 'string' && + markdownContent.trim() !== '' + ) { + const saveSpan = toolSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'save_markdown', + input: { markdownLength: markdownContent.length }, + metadata: { 'tool.id': 'web-scraper', 'operation.type': 'save' }, + }) + try { + const fileName = + typeof inputData.storage?.markdownFileName === + 'string' && + inputData.storage.markdownFileName.trim() !== '' + ? ValidationUtils.sanitizeFileName( + inputData.storage.markdownFileName + ) + : `scraped_${new Date().toISOString().replace(/[:.]/g, '-')}.md` + + await ensureDataDir() + const fullPath = path.join(DATA_DIR, fileName) + + if ( + !ValidationUtils.validateFilePath( + fullPath, + DATA_DIR + ) + ) { + throw new ScrapingError( + 'Invalid file path', + 'INVALID_FILE_PATH' + ) + } + // Open the file handle and write via the handle to avoid using fs.writeFile with a non-literal argument. + // Path was already validated by ValidationUtils.validateFilePath above. + const fileHandle = await fs.open(fullPath, 'w') + try { + await fileHandle.writeFile(markdownContent, 'utf-8') + try { + saveSpan?.update({ + metadata: { + 'operation.status': 'success', + 'savedFileName': fileName, + 'result.savedFileSize': typeof markdownContent === 'string' ? markdownContent.length : 0, + }, + }) + } catch { + // ignore tracing update errors + } + } finally { + await fileHandle.close() + } + savedFilePath = fileName + log.info('Markdown content saved securely', { + fileName, + }) + } catch (error) { + try { + saveSpan?.error({ error: error instanceof Error ? error : new Error(String(error)), endSpan: true }) + } catch { + // ignore tracing errors + } + if (error instanceof ScrapingError) { + throw error + } + log.error('Failed to save markdown file', { + error: + error instanceof Error + ? error.message + : String(error), + }) + } finally { + saveSpan?.end() + } + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Scraping complete: ${extractedData.length} elements${typeof savedFilePath === 'string' ? ', saved to ' + savedFilePath : ''}`, + stage: 'web:scraper', + }, + id: 'web:scraper', + }) + + const result = webScraperOutputSchema.parse({ + url: scrapedUrl, + status, + errorMessage, + content: { + extractedData, + rawContent: + typeof inputData.selector === 'string' && + inputData.selector.trim() !== '' + ? undefined + : typeof rawContent === 'string' && + rawContent.trim().length > 0 + ? rawContent + : undefined, + markdownContent, + }, + storage: + typeof savedFilePath === 'string' + ? { savedFilePath } + : undefined, + analysis: { + metadata: + Object.keys(metadata).length > 0 ? metadata : undefined, + images: images.length > 0 ? images : undefined, + structuredData: + structuredData.length > 0 ? structuredData : undefined, + detectedLanguage, + }, + }) + + toolSpan?.update({ + output: result, + metadata: { + 'tool.output.success': status === 'success', + 'tool.output.extractedElements': extractedData.length, + 'tool.output.markdownLength': typeof markdownContent === 'string' ? markdownContent.length : 0, + 'tool.output.savedFile': savedFilePath ?? undefined, + }, + }) + toolSpan?.end() + + return result + } catch (error) { + // Handle AbortError specifically + if (error instanceof Error && error.name === 'AbortError') { + const cancelMessage = `Web scraping cancelled for ${inputData.url}` + log.warn(cancelMessage) + try { + toolSpan?.update({ + metadata: { + 'operation.status': 'cancelled', + }, + }) + } catch { + // ignore tracing update errors + } + try { + toolSpan?.end() + } catch { + // ignore tracing end errors + } + throw new Error(cancelMessage) + } + + errorMessage = `Web scraping failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + try { + toolSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + } catch { + // ignore tracing errors + } + 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<typeof webScraperTool> // ===== BATCH WEB SCRAPER TOOL ===== const batchWebScraperInputSchema = z.object({ - urls: z - .array(z.url()) - .describe('Array of URLs to scrape.') - .max(10, 'Maximum 10 URLs allowed for batch scraping'), - selector: z - .string() - .optional() - .describe('CSS selector for elements to extract from each page.'), - maxConcurrent: z - .number() - .min(1) - .max(15) - .optional() - .describe( - 'Maximum number of concurrent requests (default: 3, max: 15).' - ), - saveResults: z - .boolean() - .optional() - .describe('Whether to save results to data directory.'), - baseFileName: z - .string() - .optional() - .describe("Base filename for saved results (default: 'batch_scrape')."), + urls: z + .array(z.url()) + .describe('Array of URLs to scrape.') + .max(10, 'Maximum 10 URLs allowed for batch scraping'), + selector: z + .string() + .optional() + .describe('CSS selector for elements to extract from each page.'), + maxConcurrent: z + .number() + .min(1) + .max(15) + .optional() + .describe( + 'Maximum number of concurrent requests (default: 3, max: 15).' + ), + saveResults: z + .boolean() + .optional() + .describe('Whether to save results to data directory.'), + baseFileName: z + .string() + .optional() + .describe("Base filename for saved results (default: 'batch_scrape')."), }) const batchWebScraperOutputSchema = z.object({ - results: z - .array( - z.object({ - url: z.string(), - success: z.boolean(), - extractedData: z - .array(z.record(z.string(), z.string())) - .optional(), - markdownContent: z.string().optional(), - errorMessage: z.string().optional(), - }) - ) - .describe('Results for each scraped URL.'), - savedFilePath: z - .string() - .optional() - .describe('Path to the saved batch results file.'), - totalProcessed: z.number().describe('Total number of URLs processed.'), - successful: z.number().describe('Number of successful scrapes.'), - failed: z.number().describe('Number of failed scrapes.'), + results: z + .array( + z.object({ + url: z.string(), + success: z.boolean(), + extractedData: z + .array(z.record(z.string(), z.string())) + .optional(), + markdownContent: z.string().optional(), + errorMessage: z.string().optional(), + }) + ) + .describe('Results for each scraped URL.'), + savedFilePath: z + .string() + .optional() + .describe('Path to the saved batch results file.'), + totalProcessed: z.number().describe('Total number of URLs processed.'), + successful: z.number().describe('Number of successful scrapes.'), + failed: z.number().describe('Number of failed scrapes.'), }) export const batchWebScraperTool = createTool({ - id: 'batch-web-scraper', - description: - 'Scrape multiple web pages concurrently with enhanced JSDOM processing and rate limiting.', - inputSchema: batchWebScraperInputSchema, - outputSchema: batchWebScraperOutputSchema, - execute: async (inputData, context) => { - const requestContext = context?.requestContext as ScraperToolContext - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🌐 Batch scraping ${inputData.urls.length} URLs`, - stage: 'batch-web-scraper', - }, - id: 'batch-web-scraper', - }) - toolCallCounters.set( - 'batch-web-scraper', - (toolCallCounters.get('batch-web-scraper') ?? 0) + 1 - ) - - const tracer = trace.getTracer('batch-web-scraper', '1.0.0') - const batchSpan = tracer.startSpan('batch_scrape', { - attributes: { - 'tool.id': 'batch-web-scraper', - 'tool.input.urlCount': inputData.urls.length, - 'tool.input.maxConcurrent': inputData.maxConcurrent ?? 3, - }, - }) - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🐛 Initializing batch crawlers...', - stage: 'batch-web-scraper', - }, - id: 'batch-web-scraper', - }) - log.info('Starting enhanced batch web scraping with JSDOM', { - urlCount: inputData.urls.length, - maxConcurrent: inputData.maxConcurrent ?? 3, - saveResults: inputData.saveResults, - }) - - const results: Array<{ - url: string - success: boolean - extractedData?: Array<Record<string, string>> - markdownContent?: string - errorMessage?: string - }> = [] + id: 'batch-web-scraper', + description: + 'Scrape multiple web pages concurrently with enhanced JSDOM processing and rate limiting.', + inputSchema: batchWebScraperInputSchema, + outputSchema: batchWebScraperOutputSchema, + execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🌐 Batch scraping ${inputData.urls.length} URLs`, + stage: 'batch-web-scraper', + }, + id: 'batch-web-scraper', + }) + toolCallCounters.set( + 'batch-web-scraper', + (toolCallCounters.get('batch-web-scraper') ?? 0) + 1 + ) + + const tracingContext = context?.tracingContext + const requestCtx = requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const batchSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'batch_scrape', + input: inputData, + metadata: { + 'tool.id': 'batch-web-scraper', + 'tool.input.urlCount': inputData.urls.length, + 'tool.input.maxConcurrent': inputData.maxConcurrent ?? 3, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - let savedFilePath: string | undefined - const maxConcurrent = Math.min(inputData.maxConcurrent ?? 3, 5) + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🐛 Initializing batch crawlers...', + stage: 'batch-web-scraper', + }, + id: 'batch-web-scraper', + }) + log.info('Starting enhanced batch web scraping with JSDOM', { + urlCount: inputData.urls.length, + maxConcurrent: inputData.maxConcurrent ?? 3, + saveResults: inputData.saveResults, + }) - try { - for (let i = 0; i < inputData.urls.length; i += maxConcurrent) { - const batch = inputData.urls.slice(i, i + maxConcurrent) - const batchPromises = batch.map(async (url: string) => { - try { - const crawler = new CheerioCrawler({ - maxRequestsPerCrawl: 1, - requestHandlerTimeoutSecs: 20, - async requestHandler({ body }) { - const rawContent = body.toString() - const sanitizedHtml = - HtmlProcessor.sanitizeHtml(rawContent) - const extractedData: Array< - Record<string, string> - > = [] - - if ( - typeof inputData.selector === 'string' && - inputData.selector.trim() !== '' - ) { - // Use JSDOM for better element selection - // Ensure we pass a string and explicitly set contentType to avoid ambiguous parsing and suppress XSS/static analysis warnings. - const jsdom = new JSDOM( - String(sanitizedHtml), - { - contentType: 'text/html', - includeNodeLocations: false, - } - ) - const { document } = jsdom.window - - const selector = inputData.selector.trim() - const elements = - document.querySelectorAll(selector) - elements.forEach((element) => { - const data = new Map<string, string>() - data.set( - 'text', - element.textContent?.trim() ?? '' - ) - extractedData.push( - Object.fromEntries(data) - ) - }) - } - - // Convert to markdown and sanitize the markdown output so we do not return unencoded HTML. - let markdownContent = - HtmlProcessor.htmlToMarkdown(sanitizedHtml) - markdownContent = - HtmlProcessor.sanitizeMarkdown( - markdownContent - ) - - results.push({ - url, - success: true, - extractedData, - markdownContent, - }) - }, - failedRequestHandler({ error }) { - results.push({ - url, - success: false, - errorMessage: - error instanceof Error - ? error.message - : String(error), - }) - }, - }) - - await crawler.run([new Request({ url })]) - } catch (error) { - results.push({ - url, - success: false, - errorMessage: - error instanceof Error - ? error.message - : String(error), - }) - } - }) + const results: Array<{ + url: string + success: boolean + extractedData?: Array<Record<string, string>> + markdownContent?: string + errorMessage?: string + }> = [] - await Promise.all(batchPromises) + let savedFilePath: string | undefined + const maxConcurrent = Math.min(inputData.maxConcurrent ?? 3, 5) - if (i + maxConcurrent < inputData.urls.length) { - await new Promise((resolve) => setTimeout(resolve, 1000)) - } - } + try { + for (let i = 0; i < inputData.urls.length; i += maxConcurrent) { + const batch = inputData.urls.slice(i, i + maxConcurrent) + const batchPromises = batch.map(async (url: string) => { + try { + const crawler = new CheerioCrawler({ + maxRequestsPerCrawl: 1, + requestHandlerTimeoutSecs: 20, + async requestHandler({ body }) { + const rawContent = body.toString() + const sanitizedHtml = + HtmlProcessor.sanitizeHtml(rawContent) + const extractedData: Array< + Record<string, string> + > = [] - if (inputData.saveResults ?? false) { - try { - const timestamp = new Date() - .toISOString() - .replace(/[:.]/g, '-') - const fileName = `${ValidationUtils.sanitizeFileName(inputData.baseFileName ?? 'batch_scrape')}_${timestamp}.json` - await ensureDataDir() - const fullPath = path.join(DATA_DIR, fileName) - - if (!ValidationUtils.validateFilePath(fullPath, DATA_DIR)) { - throw new ScrapingError( - 'Invalid file path', - 'INVALID_FILE_PATH' - ) - } - // Open the file handle and write via the handle to avoid using fs.writeFile with a non-literal argument. - // Path was already validated by ValidationUtils.validateFilePath above. - const fileHandle = await fs.open(fullPath, 'w') - try { - await fileHandle.writeFile( - JSON.stringify(results, null, 2), - 'utf-8' - ) - } finally { - await fileHandle.close() + if ( + typeof inputData.selector === 'string' && + inputData.selector.trim() !== '' + ) { + // Use JSDOM for better element selection + // Ensure we pass a string and explicitly set contentType to avoid ambiguous parsing and suppress XSS/static analysis warnings. + const jsdom = new JSDOM( + String(sanitizedHtml), + { + contentType: 'text/html', + includeNodeLocations: false, } - savedFilePath = path.relative(DATA_DIR, fullPath) - log.info('Batch results saved securely', { - fileName: savedFilePath, - }) - } catch (error) { - log.error('Failed to save batch results', { - error: - error instanceof Error - ? error.message - : String(error), - }) + ) + const { document } = jsdom.window + + const selector = inputData.selector.trim() + const elements = + document.querySelectorAll(selector) + elements.forEach((element) => { + const data = new Map<string, string>() + data.set( + 'text', + element.textContent?.trim() ?? '' + ) + extractedData.push( + Object.fromEntries(data) + ) + }) } - } - const successful = results.filter((r) => r.success).length - const failed = results.length - successful - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Batch complete: ${successful}/${results.length} successful`, - stage: 'batch-web-scraper', - }, - id: 'batch-web-scraper', + // Convert to markdown and sanitize the markdown output so we do not return unencoded HTML. + let markdownContent = + HtmlProcessor.htmlToMarkdown(sanitizedHtml) + markdownContent = + HtmlProcessor.sanitizeMarkdown( + markdownContent + ) + + results.push({ + url, + success: true, + extractedData, + markdownContent, + }) + }, + failedRequestHandler({ error }) { + results.push({ + url, + success: false, + errorMessage: + error instanceof Error + ? error.message + : String(error), + }) + }, }) - batchSpan.setAttributes({ - 'tool.output.totalProcessed': results.length, - 'tool.output.successful': successful, - 'tool.output.failed': failed, - }) - batchSpan.end() - - return batchWebScraperOutputSchema.parse({ - results, - savedFilePath, - totalProcessed: results.length, - successful, - failed, + await crawler.run([new Request({ url })]) + } catch (error) { + results.push({ + url, + success: false, + errorMessage: + error instanceof Error + ? error.message + : String(error), }) + } + }) + + await Promise.all(batchPromises) + + if (i + maxConcurrent < inputData.urls.length) { + await new Promise((resolve) => setTimeout(resolve, 1000)) + } + } + + if (inputData.saveResults ?? false) { + try { + const timestamp = new Date() + .toISOString() + .replace(/[:.]/g, '-') + const fileName = `${ValidationUtils.sanitizeFileName(inputData.baseFileName ?? 'batch_scrape')}_${timestamp}.json` + await ensureDataDir() + const fullPath = path.join(DATA_DIR, fileName) + + if (!ValidationUtils.validateFilePath(fullPath, DATA_DIR)) { + throw new ScrapingError( + 'Invalid file path', + 'INVALID_FILE_PATH' + ) + } + // Open the file handle and write via the handle to avoid using fs.writeFile with a non-literal argument. + // Path was already validated by ValidationUtils.validateFilePath above. + const fileHandle = await fs.open(fullPath, 'w') + try { + await fileHandle.writeFile( + JSON.stringify(results, null, 2), + 'utf-8' + ) + } finally { + await fileHandle.close() + } + savedFilePath = path.relative(DATA_DIR, fullPath) + log.info('Batch results saved securely', { + fileName: savedFilePath, + }) } catch (error) { - const errorMessage = `Batch scraping failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - batchSpan.recordException(error) - } - batchSpan.setStatus({ code: 2, message: errorMessage }) - batchSpan.end() - throw error + log.error('Failed to save batch results', { + error: + error instanceof Error + ? error.message + : String(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', - }) - }, + } + + const successful = results.filter((r) => r.success).length + const failed = results.length - successful + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Batch complete: ${successful}/${results.length} successful`, + stage: 'batch-web-scraper', + }, + id: 'batch-web-scraper', + }) + + batchSpan?.update({ + metadata: { + 'tool.output.totalProcessed': results.length, + 'tool.output.successful': successful, + 'tool.output.failed': failed, + }, + }) + batchSpan?.end() + + return batchWebScraperOutputSchema.parse({ + results, + savedFilePath, + totalProcessed: results.length, + successful, + failed, + }) + } catch (error) { + const errorMessage = `Batch scraping failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + batchSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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<typeof batchWebScraperTool> @@ -1603,821 +1821,843 @@ export type BatchWebScraperUITool = InferUITool<typeof batchWebScraperTool> // ===== SITE MAP EXTRACTOR TOOL ===== const siteMapExtractorInputSchema = z.object({ - url: z - .url() - .describe('The base URL of the website to extract sitemap from.') - .refine((v) => ValidationUtils.validateUrl(v), 'Invalid URL format'), - maxDepth: z - .number() - .min(1) - .max(5) - .optional() - .describe('Maximum depth to crawl (default: 2, max: 5).'), - maxPages: z - .number() - .min(1) - .max(200) - .optional() - .describe('Maximum number of pages to crawl (default: 50, max: 200).'), - includeExternal: z - .boolean() - .optional() - .describe('Whether to include external links (default: false).'), - saveMap: z - .boolean() - .optional() - .describe('Whether to save the site map to data directory.'), + url: z + .url() + .describe('The base URL of the website to extract sitemap from.') + .refine((v) => ValidationUtils.validateUrl(v), 'Invalid URL format'), + maxDepth: z + .number() + .min(1) + .max(5) + .optional() + .describe('Maximum depth to crawl (default: 2, max: 5).'), + maxPages: z + .number() + .min(1) + .max(200) + .optional() + .describe('Maximum number of pages to crawl (default: 50, max: 200).'), + includeExternal: z + .boolean() + .optional() + .describe('Whether to include external links (default: false).'), + saveMap: z + .boolean() + .optional() + .describe('Whether to save the site map to data directory.'), }) const siteMapExtractorOutputSchema = z.object({ - baseUrl: z.string().describe('The base URL that was crawled.'), - pages: z - .array( - z.object({ - url: z.string(), - title: z.string().optional(), - depth: z.number(), - internalLinks: z.array(z.string()), - externalLinks: z.array(z.string()), - }) - ) - .describe('Array of discovered pages with their metadata.'), - totalPages: z.number().describe('Total number of pages discovered.'), - savedFilePath: z - .string() - .optional() - .describe('Path to the saved site map file.'), + baseUrl: z.string().describe('The base URL that was crawled.'), + pages: z + .array( + z.object({ + url: z.string(), + title: z.string().optional(), + depth: z.number(), + internalLinks: z.array(z.string()), + externalLinks: z.array(z.string()), + }) + ) + .describe('Array of discovered pages with their metadata.'), + totalPages: z.number().describe('Total number of pages discovered.'), + savedFilePath: z + .string() + .optional() + .describe('Path to the saved site map file.'), }) export const siteMapExtractorTool = createTool({ - id: 'site-map-extractor', - description: - 'Extract a comprehensive site map by crawling internal links with enhanced JSDOM processing and rate limiting.', - inputSchema: siteMapExtractorInputSchema, - outputSchema: siteMapExtractorOutputSchema, - execute: async (inputData, context) => { - const requestContext = context?.requestContext as ScraperToolContext - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🗺️ Starting site map extraction for ${inputData.url}`, - stage: 'site-map-extractor', - }, - id: 'site-map-extractor', - }) - toolCallCounters.set( - 'site-map-extractor', - (toolCallCounters.get('site-map-extractor') ?? 0) + 1 - ) - const tracer = trace.getTracer('site-map-extractor', '1.0.0') - const mapSpan = tracer.startSpan('site_map_extraction', { - attributes: { - 'tool.id': 'site-map-extractor', - 'tool.input.url': inputData.url, - 'tool.input.maxDepth': inputData.maxDepth ?? 2, - 'tool.input.maxPages': inputData.maxPages ?? 50, - 'tool.input.includeExternal': - inputData.includeExternal ?? false, - }, - }) - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🔍 Crawling internal links...', - stage: 'site-map-extractor', - }, - id: 'site-map-extractor', - }) - log.info('Starting enhanced site map extraction with JSDOM', { - url: inputData.url, - maxDepth: inputData.maxDepth ?? 2, - maxPages: inputData.maxPages ?? 50, - }) - - const baseUrl = new URL(inputData.url) - const visited = new Set<string>() - const pages: Array<{ - url: string - title?: string - depth: number - internalLinks: string[] - externalLinks: string[] - }> = [] - - let savedFilePath: string | undefined - - // Define crawlPage at the function body root to avoid nested function declarations. - const crawlPage = async (url: string, depth: number): Promise<void> => { - if ( - visited.has(url) || - depth > (inputData.maxDepth ?? 2) || - pages.length >= (inputData.maxPages ?? 50) - ) { - return - } - - visited.add(url) - - try { - const crawler = new CheerioCrawler({ - maxRequestsPerCrawl: 5, - maxConcurrency: 10, - requestHandlerTimeoutSecs: 15, - async requestHandler({ body }) { - const rawContent = body.toString() - const sanitizedHtml = - HtmlProcessor.sanitizeHtml(rawContent) - - // Use JSDOM for better link extraction - const dom = new JSDOM(sanitizedHtml, { - includeNodeLocations: false, - }) - const { document } = dom.window - - const title = - document - .querySelector('title') - ?.textContent?.trim() ?? - document.querySelector('h1')?.textContent?.trim() - const internalLinks: string[] = [] - const externalLinks: string[] = [] - - const links = document.querySelectorAll('a[href]') - links.forEach((link) => { - const href = link.getAttribute('href') - if (href !== null && href !== undefined) { - try { - const absoluteUrl = new URL(href, url).href - const linkUrl = new URL(absoluteUrl) - - if (linkUrl.hostname === baseUrl.hostname) { - if (!visited.has(absoluteUrl)) { - internalLinks.push(absoluteUrl) - } - } else if ( - inputData.includeExternal === true - ) { - // Align with schema default (false) - externalLinks.push(absoluteUrl) - } - } catch { - // Invalid URL, skip silently - } - } - }) - - pages.push({ - url, - title, - depth, - internalLinks, - externalLinks, - }) - - for (const link of internalLinks) { - if ( - !visited.has(link) && - pages.length < (inputData.maxPages ?? 50) - ) { - await crawlPage(link, depth + 1) - await new Promise((resolve) => - setTimeout(resolve, 200) - ) - } - } - }, - failedRequestHandler({ error }) { - log.warn(`Failed to crawl ${url}`, { - error: - error instanceof Error - ? error.message - : String(error), - }) - }, - }) + id: 'site-map-extractor', + description: + 'Extract a comprehensive site map by crawling internal links with enhanced JSDOM processing and rate limiting.', + inputSchema: siteMapExtractorInputSchema, + outputSchema: siteMapExtractorOutputSchema, + execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🗺️ Starting site map extraction for ${inputData.url}`, + stage: 'site-map-extractor', + }, + id: 'site-map-extractor', + }) + toolCallCounters.set( + 'site-map-extractor', + (toolCallCounters.get('site-map-extractor') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const mapSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'site_map_extraction', + input: inputData, + metadata: { + 'tool.id': 'site-map-extractor', + 'tool.input.url': inputData.url, + 'tool.input.maxDepth': inputData.maxDepth ?? 2, + 'tool.input.maxPages': inputData.maxPages ?? 50, + 'tool.input.includeExternal': inputData.includeExternal ?? false, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - await crawler.run([new Request({ url })]) - } catch (error) { - log.warn(`Error crawling ${url}`, { - error: - error instanceof Error ? error.message : String(error), - }) - } - } + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🔍 Crawling internal links...', + stage: 'site-map-extractor', + }, + id: 'site-map-extractor', + }) + log.info('Starting enhanced site map extraction with JSDOM', { + url: inputData.url, + maxDepth: inputData.maxDepth ?? 2, + maxPages: inputData.maxPages ?? 50, + }) - try { - await crawlPage(inputData.url, 0) + const baseUrl = new URL(inputData.url) + const visited = new Set<string>() + const pages: Array<{ + url: string + title?: string + depth: number + internalLinks: string[] + externalLinks: string[] + }> = [] + + let savedFilePath: string | undefined + + // Define crawlPage at the function body root to avoid nested function declarations. + const crawlPage = async (url: string, depth: number): Promise<void> => { + if ( + visited.has(url) || + depth > (inputData.maxDepth ?? 2) || + pages.length >= (inputData.maxPages ?? 50) + ) { + return + } + + visited.add(url) + + try { + const crawler = new CheerioCrawler({ + maxRequestsPerCrawl: 5, + maxConcurrency: 10, + requestHandlerTimeoutSecs: 15, + async requestHandler({ body }) { + const rawContent = body.toString() + const sanitizedHtml = + HtmlProcessor.sanitizeHtml(rawContent) + + // Use JSDOM for better link extraction + const dom = new JSDOM(sanitizedHtml, { + includeNodeLocations: false, + }) + const { document } = dom.window - if (inputData.saveMap ?? false) { + const title = + document + .querySelector('title') + ?.textContent?.trim() ?? + document.querySelector('h1')?.textContent?.trim() + const internalLinks: string[] = [] + const externalLinks: string[] = [] + + const links = document.querySelectorAll('a[href]') + links.forEach((link) => { + const href = link.getAttribute('href') + if (href !== null && href !== undefined) { try { - const timestamp = new Date() - .toISOString() - .replace(/[:.]/g, '-') - const fileName = `sitemap_${baseUrl.hostname}_${timestamp}.json` - await ensureDataDir() - const fullPath = path.join(DATA_DIR, fileName) - - if (!ValidationUtils.validateFilePath(fullPath, DATA_DIR)) { - throw new ScrapingError( - 'Invalid file path', - 'INVALID_FILE_PATH' - ) - } - await fs.writeFile( - fullPath, - JSON.stringify( - { - baseUrl: inputData.url, - crawledAt: new Date().toISOString(), - pages, - }, - null, - 2 - ), - 'utf-8' - ) + const absoluteUrl = new URL(href, url).href + const linkUrl = new URL(absoluteUrl) - savedFilePath = path.relative(DATA_DIR, fullPath) // Consistent with other tools - log.info('Site map saved securely', { - fileName: savedFilePath, - }) - } catch (error) { - log.error('Failed to save site map', { - error: - error instanceof Error - ? error.message - : String(error), - }) + if (linkUrl.hostname === baseUrl.hostname) { + if (!visited.has(absoluteUrl)) { + internalLinks.push(absoluteUrl) + } + } else if ( + inputData.includeExternal === true + ) { + // Align with schema default (false) + externalLinks.push(absoluteUrl) + } + } catch { + // Invalid URL, skip silently } - } + } + }) - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'done', - message: `✅ Site map complete: ${pages.length} pages discovered`, - stage: 'site-map-extractor', - }, - id: 'site-map-extractor', + pages.push({ + url, + title, + depth, + internalLinks, + externalLinks, }) - mapSpan.setAttributes({ - 'tool.output.totalPages': pages.length, - 'tool.output.savedFile': - typeof savedFilePath === 'string' && - savedFilePath.trim().length > 0, + + for (const link of internalLinks) { + if ( + !visited.has(link) && + pages.length < (inputData.maxPages ?? 50) + ) { + await crawlPage(link, depth + 1) + await new Promise((resolve) => + setTimeout(resolve, 200) + ) + } + } + }, + failedRequestHandler({ error }) { + log.warn(`Failed to crawl ${url}`, { + error: + error instanceof Error + ? error.message + : String(error), }) - mapSpan.end() + }, + }) + + await crawler.run([new Request({ url })]) + } catch (error) { + log.warn(`Error crawling ${url}`, { + error: + error instanceof Error ? error.message : String(error), + }) + } + } - return siteMapExtractorOutputSchema.parse({ + try { + await crawlPage(inputData.url, 0) + + if (inputData.saveMap ?? false) { + try { + const timestamp = new Date() + .toISOString() + .replace(/[:.]/g, '-') + const fileName = `sitemap_${baseUrl.hostname}_${timestamp}.json` + await ensureDataDir() + const fullPath = path.join(DATA_DIR, fileName) + + if (!ValidationUtils.validateFilePath(fullPath, DATA_DIR)) { + throw new ScrapingError( + 'Invalid file path', + 'INVALID_FILE_PATH' + ) + } + await fs.writeFile( + fullPath, + JSON.stringify( + { baseUrl: inputData.url, + crawledAt: new Date().toISOString(), pages, - totalPages: pages.length, - savedFilePath, - }) + }, + null, + 2 + ), + 'utf-8' + ) + + savedFilePath = path.relative(DATA_DIR, fullPath) // Consistent with other tools + log.info('Site map saved securely', { + fileName: savedFilePath, + }) } catch (error) { - const errorMessage = `Site map extraction failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - mapSpan.recordException(error) - } - mapSpan.setStatus({ code: 2, message: errorMessage }) - mapSpan.end() - throw error + log.error('Failed to save site map', { + error: + error instanceof Error + ? error.message + : String(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', - }) - }, + } + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `✅ Site map complete: ${pages.length} pages discovered`, + stage: 'site-map-extractor', + }, + id: 'site-map-extractor', + }) + mapSpan?.update({ + metadata: { + 'tool.output.totalPages': pages.length, + 'tool.output.savedFile': typeof savedFilePath === 'string' && savedFilePath.trim().length > 0, + }, + }) + mapSpan?.end() + + return siteMapExtractorOutputSchema.parse({ + baseUrl: inputData.url, + pages, + totalPages: pages.length, + savedFilePath, + }) + } catch (error) { + const errorMessage = `Site map extraction failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + mapSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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<typeof siteMapExtractorTool> // ===== LINK EXTRACTOR TOOL ===== const linkExtractorInputSchema = z.object({ - url: z - .url() - .describe('The URL of the web page to extract links from.') - .refine((v) => ValidationUtils.validateUrl(v), 'Invalid URL format'), - linkTypes: z - .array(z.enum(['internal', 'external', 'all'])) - .optional() - .describe("Types of links to extract (default: ['all'])."), - includeAnchors: z - .boolean() - .optional() - .describe('Whether to include anchor text with links (default: true).'), - filterPatterns: z - .array(z.string()) - .optional() - .describe('Regex patterns to filter links by href.'), + url: z + .url() + .describe('The URL of the web page to extract links from.') + .refine((v) => ValidationUtils.validateUrl(v), 'Invalid URL format'), + linkTypes: z + .array(z.enum(['internal', 'external', 'all'])) + .optional() + .describe("Types of links to extract (default: ['all'])."), + includeAnchors: z + .boolean() + .optional() + .describe('Whether to include anchor text with links (default: true).'), + filterPatterns: z + .array(z.string()) + .optional() + .describe('Regex patterns to filter links by href.'), }) const linkExtractorOutputSchema = z.object({ - url: z.string().describe('The URL that was analyzed.'), - links: z - .array( - z.object({ - href: z.string(), - text: z.string(), - type: z.enum(['internal', 'external']), - isValid: z.boolean(), - }) - ) - .describe('Array of extracted links with metadata.'), - summary: z - .object({ - total: z.number(), - internal: z.number(), - external: z.number(), - invalid: z.number(), - }) - .describe('Summary statistics of extracted links.'), + url: z.string().describe('The URL that was analyzed.'), + links: z + .array( + z.object({ + href: z.string(), + text: z.string(), + type: z.enum(['internal', 'external']), + isValid: z.boolean(), + }) + ) + .describe('Array of extracted links with metadata.'), + summary: z + .object({ + total: z.number(), + internal: z.number(), + external: z.number(), + invalid: z.number(), + }) + .describe('Summary statistics of extracted links.'), }) export const linkExtractorTool = createTool({ - id: 'link-extractor', - description: - 'Extract and analyze all links from a web page with enhanced JSDOM processing and filtering.', - inputSchema: linkExtractorInputSchema, - outputSchema: linkExtractorOutputSchema, - execute: async (inputData, context) => { - const requestContext = context?.requestContext as ScraperToolContext - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🔗 Extracting links from ${inputData.url}`, - stage: 'link-extractor', - }, - id: 'link-extractor', - }) - toolCallCounters.set( - 'link-extractor', - (toolCallCounters.get('link-extractor') ?? 0) + 1 - ) - const tracer = trace.getTracer('link-extractor', '1.0.0') - const linkSpan = tracer.startSpan('link_extraction', { - attributes: { - 'tool.id': 'link-extractor', - 'tool.input.url': inputData.url, - 'tool.input.linkTypes': inputData.linkTypes ?? ['all'], - 'tool.input.includeAnchors': inputData.includeAnchors ?? true, - 'tool.input.filterCount': inputData.filterPatterns?.length ?? 0, - }, - }) - - log.info('Starting enhanced link extraction with JSDOM', { - url: inputData.url, - linkTypes: inputData.linkTypes ?? ['all'], - }) - - try { - let rawContent: string | undefined - let scrapedUrl: string = inputData.url - - const crawler = new CheerioCrawler({ - maxRequestsPerCrawl: 10, - maxConcurrency: 10, - requestHandlerTimeoutSecs: 20, - async requestHandler({ request, body }) { - scrapedUrl = request.url - rawContent = body.toString() - }, - failedRequestHandler({ error }) { - throw new ScrapingError( - `Failed to fetch ${inputData.url}: ${error instanceof Error ? error.message : String(error)}`, - 'FETCH_FAILED', - undefined, - inputData.url - ) - }, - }) + id: 'link-extractor', + description: + 'Extract and analyze all links from a web page with enhanced JSDOM processing and filtering.', + inputSchema: linkExtractorInputSchema, + outputSchema: linkExtractorOutputSchema, + execute: async (inputData, context) => { + const requestContext = context?.requestContext as ScraperToolContext + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🔗 Extracting links from ${inputData.url}`, + stage: 'link-extractor', + }, + id: 'link-extractor', + }) + toolCallCounters.set( + 'link-extractor', + (toolCallCounters.get('link-extractor') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const linkSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'link_extraction', + input: inputData, + metadata: { + 'tool.id': 'link-extractor', + 'tool.input.url': inputData.url, + 'tool.input.linkTypes': inputData.linkTypes ?? ['all'], + 'tool.input.includeAnchors': inputData.includeAnchors ?? true, + 'tool.input.filterCount': inputData.filterPatterns?.length ?? 0, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - await crawler.run([new Request({ url: inputData.url })]) + log.info('Starting enhanced link extraction with JSDOM', { + url: inputData.url, + linkTypes: inputData.linkTypes ?? ['all'], + }) - if ( - typeof rawContent !== 'string' || - rawContent.trim().length === 0 - ) { - throw new ScrapingError( - 'Failed to retrieve page content', - 'NO_CONTENT' - ) - } + try { + let rawContent: string | undefined + let scrapedUrl: string = inputData.url + + const crawler = new CheerioCrawler({ + maxRequestsPerCrawl: 10, + maxConcurrency: 10, + requestHandlerTimeoutSecs: 20, + async requestHandler({ request, body }) { + scrapedUrl = request.url + rawContent = body.toString() + }, + failedRequestHandler({ error }) { + throw new ScrapingError( + `Failed to fetch ${inputData.url}: ${error instanceof Error ? error.message : String(error)}`, + 'FETCH_FAILED', + undefined, + inputData.url + ) + }, + }) + + await crawler.run([new Request({ url: inputData.url })]) + + if ( + typeof rawContent !== 'string' || + rawContent.trim().length === 0 + ) { + throw new ScrapingError( + 'Failed to retrieve page content', + 'NO_CONTENT' + ) + } - const sanitizedHtml = HtmlProcessor.sanitizeHtml(rawContent) + const sanitizedHtml = HtmlProcessor.sanitizeHtml(rawContent) - // Use JSDOM for comprehensive link extraction - const dom = new JSDOM(sanitizedHtml, { - includeNodeLocations: false, - }) - const { document } = dom.window - const baseUrl = new URL(scrapedUrl) + // Use JSDOM for comprehensive link extraction + const dom = new JSDOM(sanitizedHtml, { + includeNodeLocations: false, + }) + const { document } = dom.window + const baseUrl = new URL(scrapedUrl) - const wildcardMatch = (pattern: string, input: string): boolean => { - const trimmed = pattern.trim() - if (trimmed === '') { - return false - } - const parts = trimmed - .split('*') - .map((p) => p.trim()) - .filter((p) => p.length > 0) - if (parts.length === 0) { - return true - } - let position = 0 - for (const part of parts) { - const idx = input.indexOf(part, position) - if (idx === -1) { - return false - } - position = idx + part.length + const wildcardMatch = (pattern: string, input: string): boolean => { + const trimmed = pattern.trim() + if (trimmed === '') { + return false + } + const parts = trimmed + .split('*') + .map((p) => p.trim()) + .filter((p) => p.length > 0) + if (parts.length === 0) { + return true + } + let position = 0 + for (const part of parts) { + const idx = input.indexOf(part, position) + if (idx === -1) { + return false + } + position = idx + part.length + } + return true + } + + const compiledFilters = (inputData.filterPatterns ?? []) + .filter((p) => typeof p === 'string' && p.trim() !== '') + .map((pattern) => { + try { + return { pattern, re2: new RE2(pattern) } + } catch { + return { pattern, re2: null as RE2 | null } + } + }) + + const links: Array<{ + href: string + text: string + type: 'internal' | 'external' + isValid: boolean + }> = [] + + const linkElements = document.querySelectorAll('a[href]') + linkElements.forEach((link) => { + const href = link.getAttribute('href') + const text = link.textContent?.trim() ?? href ?? '' + + if (href !== null) { + try { + const absoluteUrl = new URL(href, scrapedUrl).href + const linkUrl = new URL(absoluteUrl) + + if (compiledFilters.length > 0) { + const matchesFilter = compiledFilters.some( + ({ pattern, re2 }) => { + if (re2) { + return re2.test(absoluteUrl) + } + return wildcardMatch(pattern, absoluteUrl) } - return true - } + ) - const compiledFilters = (inputData.filterPatterns ?? []) - .filter((p) => typeof p === 'string' && p.trim() !== '') - .map((pattern) => { - try { - return { pattern, re2: new RE2(pattern) } - } catch { - return { pattern, re2: null as RE2 | null } - } - }) - - const links: Array<{ - href: string - text: string - type: 'internal' | 'external' - isValid: boolean - }> = [] - - const linkElements = document.querySelectorAll('a[href]') - linkElements.forEach((link) => { - const href = link.getAttribute('href') - const text = link.textContent?.trim() ?? href ?? '' - - if (href !== null) { - try { - const absoluteUrl = new URL(href, scrapedUrl).href - const linkUrl = new URL(absoluteUrl) - - if (compiledFilters.length > 0) { - const matchesFilter = compiledFilters.some( - ({ pattern, re2 }) => { - if (re2) { - return re2.test(absoluteUrl) - } - return wildcardMatch(pattern, absoluteUrl) - } - ) - - if (!matchesFilter) { - return - } - } - - const isInternal = linkUrl.hostname === baseUrl.hostname - const linkType = isInternal ? 'internal' : 'external' - - const requestedTypes = inputData.linkTypes ?? ['all'] - if ( - !requestedTypes.includes('all') && - !requestedTypes.includes(linkType) - ) { - return - } + if (!matchesFilter) { + return + } + } - // Validate and sanitize href before returning it to avoid unencoded values reaching HTML contexts. - const isValidAbsolute = - ValidationUtils.validateUrl(absoluteUrl) - const safeHref = isValidAbsolute - ? absoluteUrl - : HtmlProcessor.sanitizeMarkdown( - String(absoluteUrl) - ) - - links.push({ - href: safeHref, - text: - inputData.includeAnchors !== false ? text : '', - type: linkType, - isValid: isValidAbsolute, - }) - } catch { - // In case of any parsing error, return a sanitized href and mark it invalid. - const safeHref = HtmlProcessor.sanitizeMarkdown( - String(href) - ) - links.push({ - href: safeHref, - text: - inputData.includeAnchors !== false ? text : '', - type: 'external', - isValid: false, - }) - } - } - }) + const isInternal = linkUrl.hostname === baseUrl.hostname + const linkType = isInternal ? 'internal' : 'external' - const summary = { - total: links.length, - internal: links.filter((l) => l.type === 'internal').length, - external: links.filter((l) => l.type === 'external').length, - invalid: links.filter((l) => !l.isValid).length, + const requestedTypes = inputData.linkTypes ?? ['all'] + if ( + !requestedTypes.includes('all') && + !requestedTypes.includes(linkType) + ) { + return } - linkSpan.setAttributes({ - 'tool.output.linkCount': links.length, - 'tool.output.summary': JSON.stringify(summary), - }) - linkSpan.end() - return linkExtractorOutputSchema.parse({ - url: scrapedUrl, - links, - summary, + // Validate and sanitize href before returning it to avoid unencoded values reaching HTML contexts. + const isValidAbsolute = + ValidationUtils.validateUrl(absoluteUrl) + const safeHref = isValidAbsolute + ? absoluteUrl + : HtmlProcessor.sanitizeMarkdown( + String(absoluteUrl) + ) + + links.push({ + href: safeHref, + text: + inputData.includeAnchors !== false ? text : '', + type: linkType, + isValid: isValidAbsolute, }) - } catch (error) { - const errorMessage = `Link extraction failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - linkSpan.recordException(error) - } - linkSpan.setStatus({ code: 2, message: errorMessage }) - linkSpan.end() - throw error + } catch { + // In case of any parsing error, return a sanitized href and mark it invalid. + const safeHref = HtmlProcessor.sanitizeMarkdown( + String(href) + ) + links.push({ + href: safeHref, + text: + inputData.includeAnchors !== false ? text : '', + type: 'external', + isValid: false, + }) + } } - }, - 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', - }) - }, + }) + + const summary = { + total: links.length, + internal: links.filter((l) => l.type === 'internal').length, + external: links.filter((l) => l.type === 'external').length, + invalid: links.filter((l) => !l.isValid).length, + } + linkSpan?.update({ + metadata: { + 'tool.output.linkCount': links.length, + 'tool.output.summary': JSON.stringify(summary), + }, + }) + linkSpan?.end() + + return linkExtractorOutputSchema.parse({ + url: scrapedUrl, + links, + summary, + }) + } catch (error) { + const errorMessage = `Link extraction failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + linkSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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<typeof linkExtractorTool> // ===== HTML TO MARKDOWN CONVERTER TOOL ===== const htmlToMarkdownOutputSchema = z - .object({ - // Ensure we always pass a real string into the sanitizer to avoid any unencoded/ambiguous values. - markdown: z - .string() - .transform((s) => HtmlProcessor.sanitizeMarkdown(String(s))) - .describe( - 'The converted markdown content (HTML-encoded to prevent XSS).' - ), - savedFilePath: z - .string() - .optional() - .describe('Path to the saved file if saveToFile was true.'), - }) - .strict() + .object({ + // Ensure we always pass a real string into the sanitizer to avoid any unencoded/ambiguous values. + markdown: z + .string() + .transform((s) => HtmlProcessor.sanitizeMarkdown(String(s))) + .describe( + 'The converted markdown content (HTML-encoded to prevent XSS).' + ), + savedFilePath: z + .string() + .optional() + .describe('Path to the saved file if saveToFile was true.'), + }) + .strict() const htmlToMarkdownInputSchema = z - .object({ - html: z - .string() - .transform((s) => HtmlProcessor.sanitizeHtml(String(s))) - .describe( - 'The HTML content to convert to markdown (will be sanitized by the schema to prevent raw HTML storage).' - ), - saveToFile: z - .boolean() - .optional() - .describe('Whether to save the markdown to a file.'), - fileName: z - .string() - .optional() - .describe( - 'Filename for the saved markdown (relative to data directory).' - ), - }) - .strict() - -export const htmlToMarkdownTool = createTool({ - id: 'html-to-markdown', - description: - 'Convert HTML content to well-formatted markdown with enhanced JSDOM parsing and security.', - inputSchema: htmlToMarkdownInputSchema, - outputSchema: htmlToMarkdownOutputSchema, - execute: async (inputData, context) => { - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🔄 Converting HTML to markdown...', - stage: 'html-to-markdown', - }, - id: 'html-to-markdown', - }) - toolCallCounters.set( - 'html-to-markdown', - (toolCallCounters.get('html-to-markdown') ?? 0) + 1 - ) - const tracer = trace.getTracer('html-to-markdown', '1.0.0') - const convertSpan = tracer.startSpan('html_to_markdown', { - attributes: { - 'tool.id': 'html-to-markdown', - 'tool.input.htmlLength': inputData.html.length, - 'tool.input.saveToFile': inputData.saveToFile, - 'tool.input.fileName': inputData.fileName, - }, - }) - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🧹 Sanitizing HTML...', - stage: 'html-to-markdown', - }, - id: 'html-to-markdown', - }) - log.info('Converting HTML to markdown with enhanced JSDOM processing', { - htmlLength: inputData.html.length, - saveToFile: inputData.saveToFile, - }) - - let savedFilePath: string | undefined - - try { - const sanitizedHtml = HtmlProcessor.sanitizeHtml(inputData.html) - const markdown = HtmlProcessor.htmlToMarkdown(sanitizedHtml) - - if (inputData.saveToFile === true) { - try { - const providedName = inputData.fileName - const fileName = - typeof providedName === 'string' && - providedName.trim().length > 0 - ? ValidationUtils.sanitizeFileName(providedName) - : `converted_${new Date().toISOString().replace(/[:.]/g, '-')}.md` - await ensureDataDir() - const fullPath = path.join(DATA_DIR, fileName) - - if (!ValidationUtils.validateFilePath(fullPath, DATA_DIR)) { - throw new ScrapingError( - 'Invalid file path', - 'INVALID_FILE_PATH' - ) - } - // Open the file handle and write via the handle to avoid using fs.writeFile with a non-literal argument. - // Path was already validated by ValidationUtils.validateFilePath above. - const fileHandle = await fs.open(fullPath, 'w') - try { - await fileHandle.writeFile(markdown, 'utf-8') - } finally { - await fileHandle.close() - } - savedFilePath = path.relative(DATA_DIR, fullPath) - log.info('Markdown saved securely', { - fileName: savedFilePath, - }) - } catch (error) { - if (error instanceof ScrapingError) { - throw error - } - log.error('Failed to save markdown file', { - error: - error instanceof Error - ? error.message - : String(error), - }) - } - } + .object({ + html: z + .string() + .transform((s) => HtmlProcessor.sanitizeHtml(String(s))) + .describe( + 'The HTML content to convert to markdown (will be sanitized by the schema to prevent raw HTML storage).' + ), + saveToFile: z + .boolean() + .optional() + .describe('Whether to save the markdown to a file.'), + fileName: z + .string() + .optional() + .describe( + 'Filename for the saved markdown (relative to data directory).' + ), + }) + .strict() - convertSpan.setAttributes({ - 'tool.output.markdownLength': markdown.length, - 'tool.output.savedFile': - typeof savedFilePath === 'string' && - savedFilePath.trim().length > 0, - }) - convertSpan.end() +export const htmlToMarkdownTool = createTool({ + id: 'html-to-markdown', + description: + 'Convert HTML content to well-formatted markdown with enhanced JSDOM parsing and security.', + inputSchema: htmlToMarkdownInputSchema, + outputSchema: htmlToMarkdownOutputSchema, + execute: async (inputData, context) => { + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🔄 Converting HTML to markdown...', + stage: 'html-to-markdown', + }, + id: 'html-to-markdown', + }) + toolCallCounters.set( + 'html-to-markdown', + (toolCallCounters.get('html-to-markdown') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const convertSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'html_to_markdown', + input: inputData, + metadata: { + 'tool.id': 'html-to-markdown', + 'tool.input.htmlLength': inputData.html.length, + 'tool.input.saveToFile': inputData.saveToFile, + 'tool.input.fileName': inputData.fileName, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - return htmlToMarkdownOutputSchema.parse({ - markdown, - savedFilePath, - }) + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🧹 Sanitizing HTML...', + stage: 'html-to-markdown', + }, + id: 'html-to-markdown', + }) + log.info('Converting HTML to markdown with enhanced JSDOM processing', { + htmlLength: inputData.html.length, + saveToFile: inputData.saveToFile, + }) + + let savedFilePath: string | undefined + + try { + const sanitizedHtml = HtmlProcessor.sanitizeHtml(inputData.html) + const markdown = HtmlProcessor.htmlToMarkdown(sanitizedHtml) + + if (inputData.saveToFile === true) { + try { + const providedName = inputData.fileName + const fileName = + typeof providedName === 'string' && + providedName.trim().length > 0 + ? ValidationUtils.sanitizeFileName(providedName) + : `converted_${new Date().toISOString().replace(/[:.]/g, '-')}.md` + await ensureDataDir() + const fullPath = path.join(DATA_DIR, fileName) + + if (!ValidationUtils.validateFilePath(fullPath, DATA_DIR)) { + throw new ScrapingError( + 'Invalid file path', + 'INVALID_FILE_PATH' + ) + } + // Open the file handle and write via the handle to avoid using fs.writeFile with a non-literal argument. + // Path was already validated by ValidationUtils.validateFilePath above. + const fileHandle = await fs.open(fullPath, 'w') + try { + await fileHandle.writeFile(markdown, 'utf-8') + } finally { + await fileHandle.close() + } + savedFilePath = path.relative(DATA_DIR, fullPath) + log.info('Markdown saved securely', { + fileName: savedFilePath, + }) } catch (error) { - const errorMessage = `HTML to markdown conversion failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - convertSpan.recordException(error) - } - convertSpan.setStatus({ code: 2, message: errorMessage }) - convertSpan.end() + if (error instanceof ScrapingError) { throw error + } + log.error('Failed to save markdown file', { + error: + error instanceof Error + ? error.message + : String(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', - }) - }, + } + + convertSpan?.update({ + metadata: { + 'tool.output.markdownLength': markdown.length, + 'tool.output.savedFile': typeof savedFilePath === 'string' && savedFilePath.trim().length > 0, + }, + }) + convertSpan?.end() + + return htmlToMarkdownOutputSchema.parse({ + markdown, + savedFilePath, + }) + } catch (error) { + const errorMessage = `HTML to markdown conversion failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + convertSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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<typeof htmlToMarkdownTool> @@ -2425,962 +2665,1007 @@ export type HtmlToMarkdownUITool = InferUITool<typeof htmlToMarkdownTool> // ===== SCRAPED CONTENT MANAGER TOOL ===== const listScrapedContentInputSchema = z.object({ - pattern: z - .string() - .optional() - .describe("Pattern to filter files (e.g., '*.md', 'scraped_*')."), - includeMetadata: z - .boolean() - .optional() - .describe('Whether to include file metadata (default: true).'), + pattern: z + .string() + .optional() + .describe("Pattern to filter files (e.g., '*.md', 'scraped_*')."), + includeMetadata: z + .boolean() + .optional() + .describe('Whether to include file metadata (default: true).'), }) const listScrapedContentOutputSchema = z.object({ - files: z - .array( - z.object({ - name: z.string(), - path: z.string(), - size: z.number().optional(), - modified: z.string().optional(), - created: z.string().optional(), - }) - ) - .describe('Array of scraped content files.'), - totalFiles: z.number().describe('Total number of files found.'), - totalSize: z.number().describe('Total size of all files in bytes.'), + files: z + .array( + z.object({ + name: z.string(), + path: z.string(), + size: z.number().optional(), + modified: z.string().optional(), + created: z.string().optional(), + }) + ) + .describe('Array of scraped content files.'), + totalFiles: z.number().describe('Total number of files found.'), + totalSize: z.number().describe('Total size of all files in bytes.'), }) export const listScrapedContentTool = createTool({ - id: 'list-scraped-content', - description: - 'List all scraped content files stored in the data directory with enhanced security.', - inputSchema: listScrapedContentInputSchema, - outputSchema: listScrapedContentOutputSchema, - execute: async (inputData, context) => { - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '📂 Listing scraped content files...', - stage: 'list-scraped-content', - }, - id: 'list-scraped-content', - }) - toolCallCounters.set( - 'list-scraped-content', - (toolCallCounters.get('list-scraped-content') ?? 0) + 1 - ) - const tracer = trace.getTracer('list-scraped-content', '1.0.0') - const listSpan = tracer.startSpan('list_scraped_content', { - attributes: { - 'tool.id': 'list-scraped-content', - 'tool.input.pattern': inputData.pattern, - 'tool.input.includeMetadata': inputData.includeMetadata ?? true, - }, - }) - - log.info('Listing scraped content with security validation', { - pattern: inputData.pattern, - }) + id: 'list-scraped-content', + description: + 'List all scraped content files stored in the data directory with enhanced security.', + inputSchema: listScrapedContentInputSchema, + outputSchema: listScrapedContentOutputSchema, + execute: async (inputData, context) => { + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '📂 Listing scraped content files...', + stage: 'list-scraped-content', + }, + id: 'list-scraped-content', + }) + toolCallCounters.set( + 'list-scraped-content', + (toolCallCounters.get('list-scraped-content') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const listSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'list_scraped_content', + input: inputData, + metadata: { + 'tool.id': 'list-scraped-content', + 'tool.input.pattern': inputData.pattern, + 'tool.input.includeMetadata': inputData.includeMetadata ?? true, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - try { - await ensureDataDir() + log.info('Listing scraped content with security validation', { + pattern: inputData.pattern, + }) + try { + await ensureDataDir() + + try { + await fs.access(DATA_DIR) + } catch { + return listScrapedContentOutputSchema.parse({ + files: [], + totalFiles: 0, + totalSize: 0, + }) + } + + const items = await fs.readdir(DATA_DIR, { withFileTypes: true }) + const files: Array<{ + name: string + path: string + size?: number + modified?: string + created?: string + }> = [] + + let totalSize = 0 + + for (const item of items) { + if (item.isFile()) { + if ( + typeof inputData.pattern === 'string' && + inputData.pattern.trim() !== '' + ) { try { - await fs.access(DATA_DIR) - } catch { - return listScrapedContentOutputSchema.parse({ - files: [], - totalFiles: 0, - totalSize: 0, - }) - } - - const items = await fs.readdir(DATA_DIR, { withFileTypes: true }) - const files: Array<{ - name: string - path: string - size?: number - modified?: string - created?: string - }> = [] - - let totalSize = 0 - - for (const item of items) { - if (item.isFile()) { + const patternStr = inputData.pattern.trim() + + // Only allow a safe subset of characters in the pattern: + // letters, numbers, underscore, hyphen, dot, space and the wildcard '*'. + // Reject patterns containing other special characters to avoid any regex-like parsing. + if (!/^[a-zA-Z0-9_\-.\\*\s]+$/.test(patternStr)) { + // Unsafe pattern detected; skip filtering to avoid using untrusted input. + } else { + // Simple and safe wildcard matcher: supports '*' only, avoids RegExp construction. + const matchWithWildcard = ( + name: string, + pattern: string + ): boolean => { + // Match-all + if (pattern === '*') { + return true + } + const parts = pattern.split('*') + let pos = 0 + for (let i = 0; i < parts.length; i++) { + const part = parts[i] + if (part === '') { + continue + } + const idx = name.indexOf(part, pos) + if (idx === -1) { + return false + } + // If pattern does not start with '*', first part must be at the start if ( - typeof inputData.pattern === 'string' && - inputData.pattern.trim() !== '' + i === 0 && + !pattern.startsWith('*') && + idx !== 0 ) { - try { - const patternStr = inputData.pattern.trim() - - // Only allow a safe subset of characters in the pattern: - // letters, numbers, underscore, hyphen, dot, space and the wildcard '*'. - // Reject patterns containing other special characters to avoid any regex-like parsing. - if (!/^[a-zA-Z0-9_\-.\\*\s]+$/.test(patternStr)) { - // Unsafe pattern detected; skip filtering to avoid using untrusted input. - } else { - // Simple and safe wildcard matcher: supports '*' only, avoids RegExp construction. - const matchWithWildcard = ( - name: string, - pattern: string - ): boolean => { - // Match-all - if (pattern === '*') { - return true - } - const parts = pattern.split('*') - let pos = 0 - for (let i = 0; i < parts.length; i++) { - const part = parts[i] - if (part === '') { - continue - } - const idx = name.indexOf(part, pos) - if (idx === -1) { - return false - } - // If pattern does not start with '*', first part must be at the start - if ( - i === 0 && - !pattern.startsWith('*') && - idx !== 0 - ) { - return false - } - pos = idx + part.length - } - // If pattern does not end with '*', last part must align with the end - if ( - !pattern.endsWith('*') && - parts[parts.length - 1] !== '' - ) { - const last = parts[parts.length - 1] - if (!name.endsWith(last)) { - return false - } - } - return true - } - if (!matchWithWildcard(item.name, patternStr)) { - continue - } - } - } catch { - // Invalid pattern, skip filtering - } + return false } - - const filePath = path.join(DATA_DIR, item.name) - const relativePath = path.relative(DATA_DIR, filePath) - - let metadata - if (inputData.includeMetadata !== false) { - try { - const stats = await fs.stat(filePath) - metadata = { - size: stats.size, - modified: stats.mtime.toISOString(), - created: stats.birthtime.toISOString(), - } - totalSize += stats.size - } catch { - // Can't get metadata, continue without it - } + pos = idx + part.length + } + // If pattern does not end with '*', last part must align with the end + if ( + !pattern.endsWith('*') && + parts[parts.length - 1] !== '' + ) { + const last = parts[parts.length - 1] + if (!name.endsWith(last)) { + return false } - - files.push({ - name: item.name, - path: relativePath, - ...metadata, - }) + } + return true + } + if (!matchWithWildcard(item.name, patternStr)) { + continue } + } + } catch { + // Invalid pattern, skip filtering } + } - listSpan.setAttributes({ - 'tool.output.totalFiles': files.length, - 'tool.output.totalSize': totalSize, - }) - listSpan.end() + const filePath = path.join(DATA_DIR, item.name) + const relativePath = path.relative(DATA_DIR, filePath) - return listScrapedContentOutputSchema.parse({ - files, - totalFiles: files.length, - totalSize, - }) - } catch (error) { - const errorMessage = `Failed to list scraped content: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - listSpan.recordException(error) + let metadata + if (inputData.includeMetadata !== false) { + try { + const stats = await fs.stat(filePath) + metadata = { + size: stats.size, + modified: stats.mtime.toISOString(), + created: stats.birthtime.toISOString(), + } + totalSize += stats.size + } catch { + // Can't get metadata, continue without it } - listSpan.setStatus({ code: 2, message: errorMessage }) - listSpan.end() - throw error + } + + files.push({ + name: item.name, + path: relativePath, + ...metadata, + }) } - }, - 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', - }) - }, + } + + listSpan?.update({ + metadata: { + 'tool.output.totalFiles': files.length, + 'tool.output.totalSize': totalSize, + }, + }) + listSpan?.end() + + return listScrapedContentOutputSchema.parse({ + files, + totalFiles: files.length, + totalSize, + }) + } catch (error) { + const errorMessage = `Failed to list scraped content: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + listSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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< - typeof listScrapedContentTool + typeof listScrapedContentTool > // ===== CONTENT CLEANER TOOL ===== const contentCleanerInputSchema = z.object({ - html: z.string().describe('The HTML content to clean.'), - removeScripts: z - .boolean() - .optional() - .describe('Remove script tags (default: true).'), - removeStyles: z - .boolean() - .optional() - .describe('Remove style tags (default: true).'), - removeComments: z - .boolean() - .optional() - .describe('Remove HTML comments (default: true).'), - preserveStructure: z - .boolean() - .optional() - .describe('Preserve document structure (default: true).'), + html: z.string().describe('The HTML content to clean.'), + removeScripts: z + .boolean() + .optional() + .describe('Remove script tags (default: true).'), + removeStyles: z + .boolean() + .optional() + .describe('Remove style tags (default: true).'), + removeComments: z + .boolean() + .optional() + .describe('Remove HTML comments (default: true).'), + preserveStructure: z + .boolean() + .optional() + .describe('Preserve document structure (default: true).'), }) const contentCleanerOutputSchema = z.object({ - cleanedHtml: z.string().describe('The cleaned HTML content.'), - originalSize: z.number().describe('Original HTML size in characters.'), - cleanedSize: z.number().describe('Cleaned HTML size in characters.'), - reductionPercent: z.number().describe('Percentage of content removed.'), + cleanedHtml: z.string().describe('The cleaned HTML content.'), + originalSize: z.number().describe('Original HTML size in characters.'), + cleanedSize: z.number().describe('Cleaned HTML size in characters.'), + reductionPercent: z.number().describe('Percentage of content removed.'), }) export const contentCleanerTool = createTool({ - id: 'content-cleaner', - description: - 'Clean HTML content by removing unwanted elements with enhanced JSDOM processing and security.', - inputSchema: contentCleanerInputSchema, - outputSchema: contentCleanerOutputSchema, - execute: async (inputData, context) => { - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: '🧹 Starting content cleaning...', - stage: 'content-cleaner', - }, - id: 'content-cleaner', - }) - toolCallCounters.set( - 'content-cleaner', - (toolCallCounters.get('content-cleaner') ?? 0) + 1 - ) - const tracer = trace.getTracer('content-cleaner', '1.0.0') - const cleanSpan = tracer.startSpan('content_cleaning', { - attributes: { - 'tool.id': 'content-cleaner', - 'tool.input.htmlLength': inputData.html.length, - 'tool.input.removeScripts': inputData.removeScripts ?? true, - 'tool.input.removeStyles': inputData.removeStyles ?? true, - 'tool.input.removeComments': inputData.removeComments ?? true, - }, - }) - - log.info('Cleaning HTML content with enhanced JSDOM security', { - originalLength: inputData.html.length, - removeScripts: inputData.removeScripts ?? true, - removeStyles: inputData.removeStyles ?? true, - }) - - try { - const dom = new JSDOM(inputData.html, { - // Removed runScripts: 'dangerously' for safety - includeNodeLocations: false, - }) + id: 'content-cleaner', + description: + 'Clean HTML content by removing unwanted elements with enhanced JSDOM processing and security.', + inputSchema: contentCleanerInputSchema, + outputSchema: contentCleanerOutputSchema, + execute: async (inputData, context) => { + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: '🧹 Starting content cleaning...', + stage: 'content-cleaner', + }, + id: 'content-cleaner', + }) + toolCallCounters.set( + 'content-cleaner', + (toolCallCounters.get('content-cleaner') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const cleanSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'content_cleaning', + input: inputData, + metadata: { + 'tool.id': 'content-cleaner', + 'tool.input.htmlLength': inputData.html.length, + 'tool.input.removeScripts': inputData.removeScripts ?? true, + 'tool.input.removeStyles': inputData.removeStyles ?? true, + 'tool.input.removeComments': inputData.removeComments ?? true, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - const { document } = dom.window - const originalSize = inputData.html.length + log.info('Cleaning HTML content with enhanced JSDOM security', { + originalLength: inputData.html.length, + removeScripts: inputData.removeScripts ?? true, + removeStyles: inputData.removeStyles ?? true, + }) - // Remove dangerous elements using JSDOM - const dangerousElements = document.querySelectorAll( - 'script, style, iframe, embed, object, noscript, meta, link[rel="stylesheet"], form, input, button, select, textarea' + try { + const dom = new JSDOM(inputData.html, { + // Removed runScripts: 'dangerously' for safety + includeNodeLocations: false, + }) + + const { document } = dom.window + const originalSize = inputData.html.length + + // Remove dangerous elements using JSDOM + const dangerousElements = document.querySelectorAll( + 'script, style, iframe, embed, object, noscript, meta, link[rel="stylesheet"], form, input, button, select, textarea' + ) + dangerousElements.forEach((element) => element.remove()) + + // Remove event handler attributes + const allElements = document.querySelectorAll('*') + allElements.forEach((element) => { + Array.from(element.attributes).forEach((attr) => { + if ( + attr.name.startsWith('on') || + HtmlProcessor.DANGEROUS_ATTRS.has( + attr.name.toLowerCase() ) - dangerousElements.forEach((element) => element.remove()) - - // Remove event handler attributes - const allElements = document.querySelectorAll('*') - allElements.forEach((element) => { - Array.from(element.attributes).forEach((attr) => { - if ( - attr.name.startsWith('on') || - HtmlProcessor.DANGEROUS_ATTRS.has( - attr.name.toLowerCase() - ) - ) { - element.removeAttribute(attr.name) - } - }) - }) - - // Remove HTML comments - if (inputData.removeComments !== false) { - const removeComments = (node: Node) => { - const childNodes = Array.from(node.childNodes) - childNodes.forEach((child) => { - if (child.nodeType === dom.window.Node.COMMENT_NODE) { - child.remove() - } else if ( - child.nodeType === dom.window.Node.ELEMENT_NODE - ) { - removeComments(child) - } - }) - } - removeComments(document.body) - } - - const cleanedHtml = - inputData.preserveStructure !== false - ? document.body.innerHTML - : (document.body.textContent ?? '') - const cleanedSize = cleanedHtml.length - const reductionPercent = - originalSize > 0 - ? ((originalSize - cleanedSize) / originalSize) * 100 - : 0 - - cleanSpan.setAttributes({ - 'tool.output.originalSize': originalSize, - 'tool.output.cleanedSize': cleanedSize, - 'tool.output.reductionPercent': - Math.round(reductionPercent * 100) / 100, - }) - cleanSpan.end() - - return contentCleanerOutputSchema.parse({ - cleanedHtml, - originalSize, - cleanedSize, - reductionPercent: Math.round(reductionPercent * 100) / 100, - }) - } catch (error) { - const errorMessage = `Content cleaning failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - cleanSpan.recordException(error) + ) { + element.removeAttribute(attr.name) + } + }) + }) + + // Remove HTML comments + if (inputData.removeComments !== false) { + const removeComments = (node: Node) => { + const childNodes = Array.from(node.childNodes) + childNodes.forEach((child) => { + if (child.nodeType === dom.window.Node.COMMENT_NODE) { + child.remove() + } else if ( + child.nodeType === dom.window.Node.ELEMENT_NODE + ) { + removeComments(child) } - cleanSpan.setStatus({ code: 2, message: errorMessage }) - cleanSpan.end() - 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', - }) - }, + removeComments(document.body) + } + + const cleanedHtml = + inputData.preserveStructure !== false + ? document.body.innerHTML + : (document.body.textContent ?? '') + const cleanedSize = cleanedHtml.length + const reductionPercent = + originalSize > 0 + ? ((originalSize - cleanedSize) / originalSize) * 100 + : 0 + + cleanSpan?.update({ + metadata: { + 'tool.output.originalSize': originalSize, + 'tool.output.cleanedSize': cleanedSize, + 'tool.output.reductionPercent': + Math.round(reductionPercent * 100) / 100, + }, + }) + cleanSpan?.end() + + return contentCleanerOutputSchema.parse({ + cleanedHtml, + originalSize, + cleanedSize, + reductionPercent: Math.round(reductionPercent * 100) / 100, + }) + } catch (error) { + const errorMessage = `Content cleaning failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + cleanSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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<typeof contentCleanerTool> // ===== API DATA FETCHER TOOL ===== const apiDataFetcherInputSchema = z.object({ - url: z.url().describe('API endpoint URL.'), - method: z - .enum(['GET', 'POST', 'PUT', 'DELETE']) - .optional() - .describe('HTTP method (default: GET).'), - headers: z - .record(z.string(), z.string()) - .optional() - .describe('HTTP headers.'), - body: z.any().optional().describe('Request body for POST/PUT.'), - auth: z - .object({ - type: z.enum(['bearer', 'basic', 'api-key']), - token: z.string(), - headerName: z.string().optional(), - }) - .optional() - .describe('Authentication configuration.'), - timeout: z - .number() - .min(1000) - .max(60000) - .optional() - .describe('Request timeout in ms (default: 30000).'), + url: z.url().describe('API endpoint URL.'), + method: z + .enum(['GET', 'POST', 'PUT', 'DELETE']) + .optional() + .describe('HTTP method (default: GET).'), + headers: z + .record(z.string(), z.string()) + .optional() + .describe('HTTP headers.'), + body: z.any().optional().describe('Request body for POST/PUT.'), + auth: z + .object({ + type: z.enum(['bearer', 'basic', 'api-key']), + token: z.string(), + headerName: z.string().optional(), + }) + .optional() + .describe('Authentication configuration.'), + timeout: z + .number() + .min(1000) + .max(60000) + .optional() + .describe('Request timeout in ms (default: 30000).'), }) const apiDataFetcherOutputSchema = z.object({ - data: z.any().describe('Response data.'), - status: z.number().describe('HTTP status code.'), - headers: z.record(z.string(), z.string()).describe('Response headers.'), - responseTime: z.number().describe('Response time in milliseconds.'), + data: z.any().describe('Response data.'), + status: z.number().describe('HTTP status code.'), + headers: z.record(z.string(), z.string()).describe('Response headers.'), + responseTime: z.number().describe('Response time in milliseconds.'), }) export const apiDataFetcherTool = createTool({ - id: 'api-data-fetcher', - description: - 'Fetch data from JSON APIs with authentication and error handling.', - inputSchema: apiDataFetcherInputSchema, - outputSchema: apiDataFetcherOutputSchema, - execute: async (inputData, context) => { - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `🌐 Fetching data from ${inputData.url}`, - stage: 'api-data-fetcher', - }, - id: 'api-data-fetcher', - }) - toolCallCounters.set( - 'api-data-fetcher', - (toolCallCounters.get('api-data-fetcher') ?? 0) + 1 - ) - const tracer = trace.getTracer('api-data-fetcher', '1.0.0') - const fetchSpan = tracer.startSpan('api_data_fetch', { - attributes: { - 'tool.id': 'api-data-fetcher', - 'tool.input.url': inputData.url, - 'tool.input.method': inputData.method ?? 'GET', - 'tool.input.hasAuth': !!inputData.auth, - }, - }) - - log.info('Fetching data from API', { - url: inputData.url, - method: inputData.method, - }) - - try { - const startTime = Date.now() - - // Prepare headers - const headers: Record<string, string> = { ...inputData.headers } - - // Add authentication - if (inputData.auth) { - switch (inputData.auth.type) { - case 'bearer': - headers['Authorization'] = - `Bearer ${inputData.auth.token}` - break - case 'basic': - headers['Authorization'] = - `Basic ${Buffer.from(inputData.auth.token).toString('base64')}` - break - case 'api-key': { - const headerName = - inputData.auth.headerName ?? 'X-API-Key' - headers[headerName] = inputData.auth.token - break - } - } - } - - // Make the request - const response = await fetch(inputData.url, { - method: inputData.method ?? 'GET', - headers, - body: - inputData.body !== null && inputData.body !== undefined - ? JSON.stringify(inputData.body) - : undefined, - signal: AbortSignal.timeout(inputData.timeout ?? 30000), - }) - - const responseTime = Date.now() - startTime - - if (!response.ok) { - throw new ScrapingError( - `API request failed: ${response.status} ${response.statusText}`, - 'API_ERROR', - response.status, - inputData.url - ) - } - - const contentType = response.headers.get('content-type') - let data: unknown - - if (contentType?.includes('application/json') ?? false) { - data = await response.json() - } else { - data = await response.text() - } - - const responseHeaders: Record<string, string> = {} - response.headers.forEach((value, key) => { - responseHeaders[key] = value - }) + id: 'api-data-fetcher', + description: + 'Fetch data from JSON APIs with authentication and error handling.', + inputSchema: apiDataFetcherInputSchema, + outputSchema: apiDataFetcherOutputSchema, + execute: async (inputData, context) => { + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `🌐 Fetching data from ${inputData.url}`, + stage: 'api-data-fetcher', + }, + id: 'api-data-fetcher', + }) + toolCallCounters.set( + 'api-data-fetcher', + (toolCallCounters.get('api-data-fetcher') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const fetchSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'api_data_fetch', + input: inputData, + metadata: { + 'tool.id': 'api-data-fetcher', + 'tool.input.url': inputData.url, + 'tool.input.method': inputData.method ?? 'GET', + 'tool.input.hasAuth': !!inputData.auth, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - fetchSpan.setAttributes({ - 'tool.output.status': response.status, - 'tool.output.responseTime': responseTime, - 'tool.output.dataType': typeof data, - }) - fetchSpan.end() + log.info('Fetching data from API', { + url: inputData.url, + method: inputData.method, + }) - return apiDataFetcherOutputSchema.parse({ - data, - status: response.status, - headers: responseHeaders, - responseTime, - }) - } catch (error) { - const errorMessage = `API data fetch failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - fetchSpan.recordException(error) - } - fetchSpan.setStatus({ code: 2, message: errorMessage }) - fetchSpan.end() - throw error + try { + const startTime = Date.now() + + // Prepare headers + const headers: Record<string, string> = { ...inputData.headers } + + // Add authentication + if (inputData.auth) { + switch (inputData.auth.type) { + case 'bearer': + headers['Authorization'] = + `Bearer ${inputData.auth.token}` + break + case 'basic': + headers['Authorization'] = + `Basic ${Buffer.from(inputData.auth.token).toString('base64')}` + break + case 'api-key': { + const headerName = + inputData.auth.headerName ?? 'X-API-Key' + headers[headerName] = inputData.auth.token + break + } } - }, - 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', - }) - }, + } + + // Make the request + const response = await fetch(inputData.url, { + method: inputData.method ?? 'GET', + headers, + body: + inputData.body !== null && inputData.body !== undefined + ? JSON.stringify(inputData.body) + : undefined, + signal: AbortSignal.timeout(inputData.timeout ?? 30000), + }) + + const responseTime = Date.now() - startTime + + if (!response.ok) { + throw new ScrapingError( + `API request failed: ${response.status} ${response.statusText}`, + 'API_ERROR', + response.status, + inputData.url + ) + } + + const contentType = response.headers.get('content-type') + let data: unknown + + if (contentType?.includes('application/json') ?? false) { + data = await response.json() + } else { + data = await response.text() + } + + const responseHeaders: Record<string, string> = {} + response.headers.forEach((value, key) => { + responseHeaders[key] = value + }) + + fetchSpan?.update({ + metadata: { + 'tool.output.status': response.status, + 'tool.output.responseTime': responseTime, + 'tool.output.dataType': typeof data, + }, + }) + fetchSpan?.end() + + return apiDataFetcherOutputSchema.parse({ + data, + status: response.status, + headers: responseHeaders, + responseTime, + }) + } catch (error) { + const errorMessage = `API data fetch failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + fetchSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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 ===== const scrapingSchedulerInputSchema = z.object({ - urls: z.array(z.string().url()).describe('URLs to scrape periodically.'), - schedule: z - .string() - .describe( - 'Cron expression for scheduling (e.g., "0 */6 * * *" for every 6 hours).' - ), - config: z - .record(z.string(), z.any()) - .optional() - .describe('Scraping configuration to apply.'), - maxRuns: z - .number() - .min(1) - .max(100) - .optional() - .describe('Maximum number of scheduled runs (default: 10).'), + urls: z.array(z.string().url()).describe('URLs to scrape periodically.'), + schedule: z + .string() + .describe( + 'Cron expression for scheduling (e.g., "0 */6 * * *" for every 6 hours).' + ), + config: z + .record(z.string(), z.any()) + .optional() + .describe('Scraping configuration to apply.'), + maxRuns: z + .number() + .min(1) + .max(100) + .optional() + .describe('Maximum number of scheduled runs (default: 10).'), }) const scrapingSchedulerOutputSchema = z.object({ - jobId: z.string().describe('Unique ID for the scheduled job.'), - nextRun: z.string().describe('Next scheduled run time.'), - totalRuns: z.number().describe('Total number of runs scheduled.'), - status: z.string().describe('Scheduler status.'), + jobId: z.string().describe('Unique ID for the scheduled job.'), + nextRun: z.string().describe('Next scheduled run time.'), + totalRuns: z.number().describe('Total number of runs scheduled.'), + status: z.string().describe('Scheduler status.'), }) export const scrapingSchedulerTool = createTool({ - id: 'scraping-scheduler', - description: 'Schedule periodic scraping jobs with cron-like syntax.', - inputSchema: scrapingSchedulerInputSchema, - outputSchema: scrapingSchedulerOutputSchema, - execute: async (inputData, context) => { - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `⏰ Scheduling scraping for ${inputData.urls.length} URLs`, - stage: 'scraping-scheduler', - }, - id: 'scraping-scheduler', - }) - toolCallCounters.set( - 'scraping-scheduler', - (toolCallCounters.get('scraping-scheduler') ?? 0) + 1 - ) - const tracer = trace.getTracer('scraping-scheduler', '1.0.0') - const scheduleSpan = tracer.startSpan('scraping_scheduler', { - attributes: { - 'tool.id': 'scraping-scheduler', - 'tool.input.urlCount': inputData.urls.length, - 'tool.input.schedule': inputData.schedule, - 'tool.input.maxRuns': inputData.maxRuns ?? 10, - }, - }) - - log.info('Scheduling periodic scraping', { - urlCount: inputData.urls.length, - schedule: inputData.schedule, - }) - - try { - // Simple in-memory scheduler (in production, use a proper job scheduler) - const jobId = `scrape_${Date.now()}_${Math.random().toString(36).substr(2, 9)}` - const maxRuns = inputData.maxRuns ?? 10 - - // Calculate next run time (simplified - doesn't parse full cron) - const nextRun = new Date(Date.now() + 3600000).toISOString() // 1 hour from now - - // In a real implementation, this would integrate with a cron job system - log.info('Scraping job scheduled', { jobId, nextRun }) + id: 'scraping-scheduler', + description: 'Schedule periodic scraping jobs with cron-like syntax.', + inputSchema: scrapingSchedulerInputSchema, + outputSchema: scrapingSchedulerOutputSchema, + execute: async (inputData, context) => { + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `⏰ Scheduling scraping for ${inputData.urls.length} URLs`, + stage: 'scraping-scheduler', + }, + id: 'scraping-scheduler', + }) + toolCallCounters.set( + 'scraping-scheduler', + (toolCallCounters.get('scraping-scheduler') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const scheduleSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'scraping_scheduler', + input: inputData, + metadata: { + 'tool.id': 'scraping-scheduler', + 'tool.input.urlCount': inputData.urls.length, + 'tool.input.schedule': inputData.schedule, + 'tool.input.maxRuns': inputData.maxRuns ?? 10, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - scheduleSpan.setAttributes({ - 'tool.output.jobId': jobId, - 'tool.output.nextRun': nextRun, - 'tool.output.totalRuns': maxRuns, - }) - scheduleSpan.end() + log.info('Scheduling periodic scraping', { + urlCount: inputData.urls.length, + schedule: inputData.schedule, + }) - return scrapingSchedulerOutputSchema.parse({ - jobId, - nextRun, - totalRuns: maxRuns, - status: 'scheduled', - }) - } catch (error) { - const errorMessage = `Scheduling failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - scheduleSpan.recordException(error) - } - scheduleSpan.setStatus({ code: 2, message: errorMessage }) - scheduleSpan.end() - 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', - }) - }, + try { + // Simple in-memory scheduler (in production, use a proper job scheduler) + const jobId = `scrape_${Date.now()}_${Math.random().toString(36).substr(2, 9)}` + const maxRuns = inputData.maxRuns ?? 10 + + // Calculate next run time (simplified - doesn't parse full cron) + const nextRun = new Date(Date.now() + 3600000).toISOString() // 1 hour from now + + // In a real implementation, this would integrate with a cron job system + log.info('Scraping job scheduled', { jobId, nextRun }) + + scheduleSpan?.update({ + metadata: { + 'tool.output.jobId': jobId, + 'tool.output.nextRun': nextRun, + 'tool.output.totalRuns': maxRuns, + }, + }) + scheduleSpan?.end() + + return scrapingSchedulerOutputSchema.parse({ + jobId, + nextRun, + totalRuns: maxRuns, + status: 'scheduled', + }) + } catch (error) { + const errorMessage = `Scheduling failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + scheduleSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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 ===== const dataExporterInputSchema = z.object({ - data: z - .array(z.record(z.string(), z.unknown())) - .describe('Data to export.'), - format: z - .enum(['json', 'csv', 'xml', 'database']) - .describe('Export format.'), - destination: z - .string() - .describe('Export destination (file path or DB connection string).'), - options: z - .record(z.string(), z.unknown()) - .optional() - .describe('Format-specific options.'), + data: z + .array(z.record(z.string(), z.unknown())) + .describe('Data to export.'), + format: z + .enum(['json', 'csv', 'xml', 'database']) + .describe('Export format.'), + destination: z + .string() + .describe('Export destination (file path or DB connection string).'), + options: z + .record(z.string(), z.unknown()) + .optional() + .describe('Format-specific options.'), }) const dataExporterOutputSchema = z.object({ - success: z.boolean().describe('Whether the export was successful.'), - exportedCount: z.number().describe('Number of records exported.'), - filePath: z.string().optional().describe('File path if exported to file.'), - message: z.string().describe('Export result message.'), + success: z.boolean().describe('Whether the export was successful.'), + exportedCount: z.number().describe('Number of records exported.'), + filePath: z.string().optional().describe('File path if exported to file.'), + message: z.string().describe('Export result message.'), }) export const dataExporterTool = createTool({ - id: 'data-exporter', - description: 'Export scraped data to various formats and destinations.', - inputSchema: dataExporterInputSchema, - outputSchema: dataExporterOutputSchema, - execute: async (inputData, context) => { - const writer = context?.writer - - await writer?.custom({ - type: 'data-tool-progress', - data: { - status: 'in-progress', - message: `📤 Exporting ${inputData.data.length} records to ${inputData.format}...`, - stage: 'data-exporter', - }, - id: 'data-exporter', - }) - toolCallCounters.set( - 'data-exporter', - (toolCallCounters.get('data-exporter') ?? 0) + 1 - ) - const tracer = trace.getTracer('data-exporter', '1.0.0') - const exportSpan = tracer.startSpan('data_export', { - attributes: { - 'tool.id': 'data-exporter', - 'tool.input.dataCount': inputData.data.length, - 'tool.input.format': inputData.format, - 'tool.input.destination': inputData.destination, - }, - }) - - log.info('Exporting data', { - dataCount: inputData.data.length, - format: inputData.format, - destination: inputData.destination, - }) - - try { - let filePath: string | undefined - let message = 'Export completed successfully' - - switch (inputData.format) { - case 'json': { - const jsonContent = JSON.stringify(inputData.data, null, 2) - await ensureDataDir() - filePath = path.isAbsolute(inputData.destination) - ? inputData.destination - : path.join(DATA_DIR, inputData.destination) - await fs.mkdir(path.dirname(filePath), { recursive: true }) - await fs.writeFile(filePath, jsonContent, 'utf-8') - break - } - - case 'csv': { - // Simple CSV conversion (in production, use a proper CSV library) - if (inputData.data.length === 0) { - throw new ScrapingError('No data to export', 'NO_DATA') - } - const headers = Object.keys(inputData.data[0]) - const csvRows = [ - headers.join(','), - ...inputData.data.map((row) => - headers - .map((header) => - JSON.stringify(row[header] ?? '') - ) - .join(',') - ), - ] - const csvContent = csvRows.join('\n') - await ensureDataDir() - filePath = path.isAbsolute(inputData.destination) - ? inputData.destination - : path.join(DATA_DIR, inputData.destination) - await fs.mkdir(path.dirname(filePath), { recursive: true }) - await fs.writeFile(filePath, csvContent, 'utf-8') - break - } - - case 'xml': { - // Simple XML conversion - let xmlContent = - '<?xml version="1.0" encoding="UTF-8"?>\n<data>\n' - for (const item of inputData.data) { - xmlContent += ' <item>\n' - for (const [key, value] of Object.entries(item)) { - xmlContent += ` <${key}>${String(value ?? '').replace(/[<>&"]/g, (c) => ({ '<': '<', '>': '>', '&': '&', '"': '"' })[c] ?? c)}</${key}>\n` - } - xmlContent += ' </item>\n' - } - xmlContent += '</data>' - await ensureDataDir() - filePath = path.isAbsolute(inputData.destination) - ? inputData.destination - : path.join(DATA_DIR, inputData.destination) - await fs.mkdir(path.dirname(filePath), { recursive: true }) - await fs.writeFile(filePath, xmlContent, 'utf-8') - break - } + id: 'data-exporter', + description: 'Export scraped data to various formats and destinations.', + inputSchema: dataExporterInputSchema, + outputSchema: dataExporterOutputSchema, + execute: async (inputData, context) => { + const writer = context?.writer + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `📤 Exporting ${inputData.data.length} records to ${inputData.format}...`, + stage: 'data-exporter', + }, + id: 'data-exporter', + }) + toolCallCounters.set( + 'data-exporter', + (toolCallCounters.get('data-exporter') ?? 0) + 1 + ) + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as ScraperToolContext | undefined + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + const exportSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'data_export', + input: inputData, + metadata: { + 'tool.id': 'data-exporter', + 'tool.input.dataCount': inputData.data.length, + 'tool.input.format': inputData.format, + 'tool.input.destination': inputData.destination, + 'user.id': userId, + 'workspace.id': workspaceId, + }, + }) - case 'database': - // Placeholder for database export (would need DB connection) - message = `Database export to ${inputData.destination} not implemented (placeholder)` - break + log.info('Exporting data', { + dataCount: inputData.data.length, + format: inputData.format, + destination: inputData.destination, + }) - default: - throw new ScrapingError( - `Unsupported format: ${inputData.format}`, - 'UNSUPPORTED_FORMAT' - ) - } + try { + let filePath: string | undefined + let message = 'Export completed successfully' + + switch (inputData.format) { + case 'json': { + const jsonContent = JSON.stringify(inputData.data, null, 2) + await ensureDataDir() + filePath = path.isAbsolute(inputData.destination) + ? inputData.destination + : path.join(DATA_DIR, inputData.destination) + await fs.mkdir(path.dirname(filePath), { recursive: true }) + await fs.writeFile(filePath, jsonContent, 'utf-8') + break + } - exportSpan.setAttributes({ - 'tool.output.success': true, - 'tool.output.exportedCount': inputData.data.length, - 'tool.output.filePath': filePath, - }) - exportSpan.end() + case 'csv': { + // Simple CSV conversion (in production, use a proper CSV library) + if (inputData.data.length === 0) { + throw new ScrapingError('No data to export', 'NO_DATA') + } + const headers = Object.keys(inputData.data[0]) + const csvRows = [ + headers.join(','), + ...inputData.data.map((row) => + headers + .map((header) => + JSON.stringify(row[header] ?? '') + ) + .join(',') + ), + ] + const csvContent = csvRows.join('\n') + await ensureDataDir() + filePath = path.isAbsolute(inputData.destination) + ? inputData.destination + : path.join(DATA_DIR, inputData.destination) + await fs.mkdir(path.dirname(filePath), { recursive: true }) + await fs.writeFile(filePath, csvContent, 'utf-8') + break + } - return dataExporterOutputSchema.parse({ - success: true, - exportedCount: inputData.data.length, - filePath, - message, - }) - } catch (error) { - const errorMessage = `Data export failed: ${error instanceof Error ? error.message : String(error)}` - log.error(errorMessage) - if (error instanceof Error) { - exportSpan.recordException(error) + case 'xml': { + // Simple XML conversion + let xmlContent = + '<?xml version="1.0" encoding="UTF-8"?>\n<data>\n' + for (const item of inputData.data) { + xmlContent += ' <item>\n' + for (const [key, value] of Object.entries(item)) { + xmlContent += ` <${key}>${String(value ?? '').replace(/[<>&"]/g, (c) => ({ '<': '<', '>': '>', '&': '&', '"': '"' })[c] ?? c)}</${key}>\n` } - exportSpan.setStatus({ code: 2, message: errorMessage }) - exportSpan.end() - throw error + xmlContent += ' </item>\n' + } + xmlContent += '</data>' + await ensureDataDir() + filePath = path.isAbsolute(inputData.destination) + ? inputData.destination + : path.join(DATA_DIR, inputData.destination) + await fs.mkdir(path.dirname(filePath), { recursive: true }) + await fs.writeFile(filePath, xmlContent, 'utf-8') + break } - }, - 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', - }) - }, + + case 'database': + // Placeholder for database export (would need DB connection) + message = `Database export to ${inputData.destination} not implemented (placeholder)` + break + + default: + throw new ScrapingError( + `Unsupported format: ${inputData.format}`, + 'UNSUPPORTED_FORMAT' + ) + } + + exportSpan?.update({ + metadata: { + 'tool.output.success': true, + 'tool.output.exportedCount': inputData.data.length, + 'tool.output.filePath': filePath, + }, + }) + exportSpan?.end() + + return dataExporterOutputSchema.parse({ + success: true, + exportedCount: inputData.data.length, + filePath, + message, + }) + } catch (error) { + const errorMessage = `Data export failed: ${error instanceof Error ? error.message : String(error)}` + log.error(errorMessage) + exportSpan?.error({ + error: error instanceof Error ? error : new Error(String(error)), + endSpan: true, + }) + 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', + }) + }, }) diff --git a/src/mastra/tools/write-note.ts b/src/mastra/tools/write-note.ts index d07a133a..05eb3f13 100644 --- a/src/mastra/tools/write-note.ts +++ b/src/mastra/tools/write-note.ts @@ -1,11 +1,17 @@ +import type { RequestContext } from '@mastra/core/request-context' import type { InferUITool } from '@mastra/core/tools' import { createTool } from '@mastra/core/tools' -import { trace } from '@opentelemetry/api' +import { SpanType } from '@mastra/core/observability' import * as fs from 'node:fs/promises' import * as path from 'node:path' import { z } from 'zod' import { log } from '../config/logger' +export interface WriteNoteContext extends RequestContext { + userId?: string + workspaceId?: string +} + const NOTES_DIR = path.join(process.cwd(), 'notes') export const writeNoteTool = createTool({ @@ -22,21 +28,47 @@ export const writeNoteTool = createTool({ .describe('The markdown content of the note.'), }), outputSchema: z.string().nonempty(), - execute: async (inputData) => { + execute: async (inputData, context) => { + const writer = context?.writer + const abortSignal = context?.abortSignal + const tracingContext = context?.tracingContext + const requestCtx = context?.requestContext as WriteNoteContext | undefined - const startTime = Date.now() + // Check if operation was already cancelled + if (abortSignal?.aborted ?? false) { + throw new Error('Write note operation cancelled') + } - // Get tracer from OpenTelemetry API - const tracer = trace.getTracer('write-note-tool', '1.0.0') - const span = tracer.startSpan('write-note', { - attributes: { - 'tool.id': 'write-note', + const userId = requestCtx?.userId + const workspaceId = requestCtx?.workspaceId + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'in-progress', + message: `Input: title="${inputData.title}" - 📝 Starting note write operation`, + stage: 'write', + }, + id: 'write', + }) + + // Create child span for note writing operation + const noteSpan = tracingContext?.currentSpan?.createChildSpan({ + type: SpanType.TOOL_CALL, + name: 'write-note', + input: { title: inputData.title, contentLength: inputData.content.length }, + metadata: { + 'tool.id': 'write', 'tool.input.title': inputData.title, 'tool.input.contentLength': inputData.content.length, 'notes.dir': NOTES_DIR, + 'user.id': userId, + 'workspace.id': workspaceId, }, }) + const startTime = Date.now() + try { const { title, content } = inputData const filePath = path.join(NOTES_DIR, `${title}.md`) @@ -44,26 +76,50 @@ export const writeNoteTool = createTool({ await fs.writeFile(filePath, content, 'utf-8') const result = `Successfully wrote to note "${title}".` - span.setAttributes({ - 'tool.output.success': true, - 'tool.output.filePath': filePath, - 'tool.output.processingTimeMs': Date.now() - startTime, + + // Update span with successful result + noteSpan?.update({ + output: { success: true, filePath }, + metadata: { + 'tool.output.success': true, + 'tool.output.filePath': filePath, + 'tool.output.processingTimeMs': Date.now() - startTime, + }, + }) + noteSpan?.end() + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `Input: title="${inputData.title}" - ✅ Note written successfully`, + stage: 'write', + }, + id: 'write', }) - span.end() + return result - } catch (error: unknown) { - const errorMessage = - error instanceof Error ? error.message : String(error) - if (error instanceof Error) { - span.recordException(error) - } - span.setStatus({ code: 2, message: errorMessage }) // ERROR status - span.setAttribute( - 'tool.output.processingTimeMs', - Date.now() - startTime - ) - span.end() - return `Error writing note: ${errorMessage}` + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error) + + // Record error in span + noteSpan?.error({ + error: error instanceof Error ? error : new Error(errorMessage), + endSpan: true, + }) + + await writer?.custom({ + type: 'data-tool-progress', + data: { + status: 'done', + message: `Input: title="${inputData.title}" - ❌ Note write failed: ${errorMessage}`, + stage: 'write', + }, + id: 'write', + }) + + log.error(`Failed to write note "${inputData.title}": ${errorMessage}`) + throw error } }, onInputStart: ({ toolCallId, messages, abortSignal }) => {