@probelabs/probe 0.6.0-rc100

package/src/agent/simpleTelemetry.js

@@ -0,0 +1,267 @@
+import { existsSync, mkdirSync, createWriteStream } from 'fs';
+import { dirname } from 'path';
+
+/**
+ * Simple telemetry implementation for probe-agent
+ * This provides basic tracing functionality without complex OpenTelemetry dependencies
+ */
+export class SimpleTelemetry {
+  constructor(options = {}) {
+    this.serviceName = options.serviceName || 'probe-agent';
+    this.enableFile = options.enableFile || false;
+    this.enableConsole = options.enableConsole || false;
+    this.filePath = options.filePath || './traces.jsonl';
+    this.stream = null;
+    
+    if (this.enableFile) {
+      this.initializeFileExporter();
+    }
+  }
+
+  initializeFileExporter() {
+    try {
+      const dir = dirname(this.filePath);
+      if (!existsSync(dir)) {
+        mkdirSync(dir, { recursive: true });
+      }
+      
+      this.stream = createWriteStream(this.filePath, { flags: 'a' });
+      this.stream.on('error', (error) => {
+        console.error(`[SimpleTelemetry] Stream error: ${error.message}`);
+      });
+      
+      console.log(`[SimpleTelemetry] File exporter initialized: ${this.filePath}`);
+    } catch (error) {
+      console.error(`[SimpleTelemetry] Failed to initialize file exporter: ${error.message}`);
+    }
+  }
+
+  createSpan(name, attributes = {}) {
+    const span = {
+      traceId: this.generateTraceId(),
+      spanId: this.generateSpanId(),
+      name,
+      startTime: Date.now(),
+      attributes: { ...attributes, service: this.serviceName },
+      events: [],
+      status: 'OK'
+    };
+    
+    return {
+      ...span,
+      addEvent: (eventName, eventAttributes = {}) => {
+        span.events.push({
+          name: eventName,
+          time: Date.now(),
+          attributes: eventAttributes
+        });
+      },
+      setAttributes: (attrs) => {
+        Object.assign(span.attributes, attrs);
+      },
+      setStatus: (status) => {
+        span.status = status;
+      },
+      end: () => {
+        span.endTime = Date.now();
+        span.duration = span.endTime - span.startTime;
+        this.exportSpan(span);
+      }
+    };
+  }
+
+  exportSpan(span) {
+    const spanData = {
+      ...span,
+      timestamp: new Date().toISOString()
+    };
+
+    if (this.enableConsole) {
+      console.log('[Trace]', JSON.stringify(spanData, null, 2));
+    }
+
+    if (this.enableFile && this.stream) {
+      this.stream.write(JSON.stringify(spanData) + '\n');
+    }
+  }
+
+  generateTraceId() {
+    return Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
+  }
+
+  generateSpanId() {
+    return Math.random().toString(36).substring(2, 10);
+  }
+
+  async flush() {
+    if (this.stream) {
+      return new Promise((resolve) => {
+        this.stream.once('drain', resolve);
+        if (!this.stream.writableNeedDrain) {
+          resolve();
+        }
+      });
+    }
+  }
+
+  async shutdown() {
+    if (this.stream) {
+      return new Promise((resolve) => {
+        this.stream.end(() => {
+          console.log(`[SimpleTelemetry] File stream closed: ${this.filePath}`);
+          resolve();
+        });
+      });
+    }
+  }
+}
+
+/**
+ * Simple tracer for application-level tracing
+ */
+export class SimpleAppTracer {
+  constructor(telemetry, sessionId = null) {
+    this.telemetry = telemetry;
+    this.sessionId = sessionId || this.generateSessionId();
+  }
+
+  generateSessionId() {
+    return Math.random().toString(36).substring(2, 15);
+  }
+
+  isEnabled() {
+    return this.telemetry !== null;
+  }
+
+  createSessionSpan(attributes = {}) {
+    if (!this.isEnabled()) return null;
+
+    return this.telemetry.createSpan('agent.session', {
+      'session.id': this.sessionId,
+      ...attributes
+    });
+  }
+
+  createAISpan(modelName, provider, attributes = {}) {
+    if (!this.isEnabled()) return null;
+
+    return this.telemetry.createSpan('ai.request', {
+      'ai.model': modelName,
+      'ai.provider': provider,
+      'session.id': this.sessionId,
+      ...attributes
+    });
+  }
+
+  createToolSpan(toolName, attributes = {}) {
+    if (!this.isEnabled()) return null;
+
+    return this.telemetry.createSpan('tool.call', {
+      'tool.name': toolName,
+      'session.id': this.sessionId,
+      ...attributes
+    });
+  }
+
+  addEvent(name, attributes = {}) {
+    // For simplicity, just log events when no active span
+    if (this.telemetry && this.telemetry.enableConsole) {
+      console.log('[Event]', name, attributes);
+    }
+  }
+
+  /**
+   * Record delegation events
+   */
+  recordDelegationEvent(eventType, data = {}) {
+    if (!this.isEnabled()) return;
+
+    this.addEvent(`delegation.${eventType}`, {
+      'session.id': this.sessionId,
+      ...data
+    });
+  }
+
+  /**
+   * Record JSON validation events
+   */
+  recordJsonValidationEvent(eventType, data = {}) {
+    if (!this.isEnabled()) return;
+
+    this.addEvent(`json_validation.${eventType}`, {
+      'session.id': this.sessionId,
+      ...data
+    });
+  }
+
+  /**
+   * Record Mermaid validation events
+   */
+  recordMermaidValidationEvent(eventType, data = {}) {
+    if (!this.isEnabled()) return;
+
+    this.addEvent(`mermaid_validation.${eventType}`, {
+      'session.id': this.sessionId,
+      ...data
+    });
+  }
+
+  setAttributes(attributes) {
+    // For simplicity, just log attributes when no active span
+    if (this.telemetry && this.telemetry.enableConsole) {
+      console.log('[Attributes]', attributes);
+    }
+  }
+
+  async withSpan(spanName, fn, attributes = {}) {
+    if (!this.isEnabled()) {
+      return fn();
+    }
+
+    const span = this.telemetry.createSpan(spanName, {
+      'session.id': this.sessionId,
+      ...attributes
+    });
+
+    try {
+      const result = await fn();
+      span.setStatus('OK');
+      return result;
+    } catch (error) {
+      span.setStatus('ERROR');
+      span.addEvent('exception', { 
+        'exception.message': error.message,
+        'exception.stack': error.stack 
+      });
+      throw error;
+    } finally {
+      span.end();
+    }
+  }
+
+  async flush() {
+    if (this.telemetry) {
+      await this.telemetry.flush();
+    }
+  }
+
+  async shutdown() {
+    if (this.telemetry) {
+      await this.telemetry.shutdown();
+    }
+  }
+}
+
+/**
+ * Initialize simple telemetry from CLI options
+ */
+export function initializeSimpleTelemetryFromOptions(options) {
+  const telemetry = new SimpleTelemetry({
+    serviceName: 'probe-agent',
+    enableFile: options.traceFile !== undefined,
+    enableConsole: options.traceConsole,
+    filePath: options.traceFile || './traces.jsonl'
+  });
+
+  return telemetry;
+}

package/src/agent/telemetry.js

@@ -0,0 +1,225 @@
+import nodeSDKPkg from '@opentelemetry/sdk-node';
+import resourcesPkg from '@opentelemetry/resources';
+import semanticConventionsPkg from '@opentelemetry/semantic-conventions';
+import { trace, context, SpanStatusCode } from '@opentelemetry/api';
+import otlpPkg from '@opentelemetry/exporter-trace-otlp-http';
+import spanPkg from '@opentelemetry/sdk-trace-base';
+
+import { existsSync, mkdirSync } from 'fs';
+import { dirname } from 'path';
+import { FileSpanExporter } from './fileSpanExporter.js';
+
+const { NodeSDK } = nodeSDKPkg;
+const { resourceFromAttributes } = resourcesPkg;
+const { ATTR_SERVICE_NAME, ATTR_SERVICE_VERSION } = semanticConventionsPkg;
+const { OTLPTraceExporter } = otlpPkg;
+const { BatchSpanProcessor, ConsoleSpanExporter } = spanPkg;
+
+/**
+ * Custom OpenTelemetry configuration for probe-agent
+ */
+export class TelemetryConfig {
+  constructor(options = {}) {
+    this.serviceName = options.serviceName || 'probe-agent';
+    this.serviceVersion = options.serviceVersion || '1.0.0';
+    this.enableFile = options.enableFile || false;
+    this.enableRemote = options.enableRemote || false;
+    this.enableConsole = options.enableConsole || false;
+    this.filePath = options.filePath || './traces.jsonl';
+    this.remoteEndpoint = options.remoteEndpoint || 'http://localhost:4318/v1/traces';
+    this.sdk = null;
+    this.tracer = null;
+  }
+
+  /**
+   * Initialize OpenTelemetry SDK
+   */
+  initialize() {
+    if (this.sdk) {
+      console.warn('Telemetry already initialized');
+      return;
+    }
+
+    const resource = resourceFromAttributes({
+      [ATTR_SERVICE_NAME]: this.serviceName,
+      [ATTR_SERVICE_VERSION]: this.serviceVersion,
+    });
+
+    const spanProcessors = [];
+
+    // Add file exporter if enabled
+    if (this.enableFile) {
+      try {
+        // Ensure the directory exists
+        const dir = dirname(this.filePath);
+        if (!existsSync(dir)) {
+          mkdirSync(dir, { recursive: true });
+        }
+        
+        const fileExporter = new FileSpanExporter(this.filePath);
+        spanProcessors.push(new BatchSpanProcessor(fileExporter, {
+          maxQueueSize: 2048,
+          maxExportBatchSize: 512,
+          scheduledDelayMillis: 500,
+          exportTimeoutMillis: 30000,
+        }));
+        console.log(`[Telemetry] File exporter enabled, writing to: ${this.filePath}`);
+      } catch (error) {
+        console.error(`[Telemetry] Failed to initialize file exporter: ${error.message}`);
+      }
+    }
+
+    // Add remote exporter if enabled
+    if (this.enableRemote) {
+      try {
+        const remoteExporter = new OTLPTraceExporter({
+          url: this.remoteEndpoint,
+        });
+        spanProcessors.push(new BatchSpanProcessor(remoteExporter, {
+          maxQueueSize: 2048,
+          maxExportBatchSize: 512,
+          scheduledDelayMillis: 500,
+          exportTimeoutMillis: 30000,
+        }));
+        console.log(`[Telemetry] Remote exporter enabled, endpoint: ${this.remoteEndpoint}`);
+      } catch (error) {
+        console.error(`[Telemetry] Failed to initialize remote exporter: ${error.message}`);
+      }
+    }
+
+    // Add console exporter if enabled (useful for debugging)
+    if (this.enableConsole) {
+      const consoleExporter = new ConsoleSpanExporter();
+      spanProcessors.push(new BatchSpanProcessor(consoleExporter, {
+        maxQueueSize: 2048,
+        maxExportBatchSize: 512,
+        scheduledDelayMillis: 500,
+        exportTimeoutMillis: 30000,
+      }));
+      console.log(`[Telemetry] Console exporter enabled`);
+    }
+
+    if (spanProcessors.length === 0) {
+      console.log('[Telemetry] No exporters configured, telemetry will not be collected');
+      return;
+    }
+
+    this.sdk = new NodeSDK({
+      resource,
+      spanProcessors,
+    });
+
+    try {
+      this.sdk.start();
+      this.tracer = trace.getTracer(this.serviceName, this.serviceVersion);
+      console.log(`[Telemetry] OpenTelemetry SDK initialized successfully`);
+    } catch (error) {
+      console.error(`[Telemetry] Failed to start OpenTelemetry SDK: ${error.message}`);
+    }
+  }
+
+  /**
+   * Get the tracer instance
+   */
+  getTracer() {
+    return this.tracer;
+  }
+
+  /**
+   * Create a span with the given name and attributes
+   */
+  createSpan(name, attributes = {}) {
+    if (!this.tracer) {
+      return null;
+    }
+
+    return this.tracer.startSpan(name, {
+      attributes,
+    });
+  }
+
+  /**
+   * Wrap a function to automatically create spans
+   */
+  wrapFunction(name, fn, attributes = {}) {
+    if (!this.tracer) {
+      return fn;
+    }
+
+    return async (...args) => {
+      const span = this.createSpan(name, attributes);
+      if (!span) {
+        return fn(...args);
+      }
+
+      try {
+        const result = await context.with(trace.setSpan(context.active(), span), () => fn(...args));
+        span.setStatus({ code: SpanStatusCode.OK });
+        return result;
+      } catch (error) {
+        span.setStatus({
+          code: SpanStatusCode.ERROR,
+          message: error.message,
+        });
+        span.recordException(error);
+        throw error;
+      } finally {
+        span.end();
+      }
+    };
+  }
+
+  /**
+   * Force flush all pending spans
+   */
+  async forceFlush() {
+    if (this.sdk) {
+      try {
+        const tracerProvider = trace.getTracerProvider();
+        
+        if (tracerProvider && typeof tracerProvider.forceFlush === 'function') {
+          await tracerProvider.forceFlush();
+        }
+        
+        // Add a small delay to ensure file writes complete
+        await new Promise(resolve => setTimeout(resolve, 100));
+        
+        console.log('[Telemetry] OpenTelemetry spans flushed successfully');
+      } catch (error) {
+        console.error(`[Telemetry] Failed to flush OpenTelemetry spans: ${error.message}`);
+      }
+    }
+  }
+
+  /**
+   * Shutdown telemetry
+   */
+  async shutdown() {
+    if (this.sdk) {
+      try {
+        await this.sdk.shutdown();
+        console.log('[Telemetry] OpenTelemetry SDK shutdown successfully');
+      } catch (error) {
+        console.error(`[Telemetry] Failed to shutdown OpenTelemetry SDK: ${error.message}`);
+      }
+    }
+  }
+}
+
+/**
+ * Initialize telemetry from CLI options
+ */
+export function initializeTelemetryFromOptions(options) {
+  const config = new TelemetryConfig({
+    serviceName: 'probe-agent',
+    serviceVersion: '1.0.0',
+    enableFile: options.traceFile !== undefined,
+    enableRemote: options.traceRemote !== undefined,
+    enableConsole: options.traceConsole,
+    filePath: options.traceFile || './traces.jsonl',
+    remoteEndpoint: options.traceRemote || 'http://localhost:4318/v1/traces',
+  });
+
+  config.initialize();
+  return config;
+}