@dexwox-labs/a2a-core/src/telemetry/decorators.ts

Core types, validation and telemetry for Google's Agent-to-Agent (A2A) protocol - shared foundation for client and server implementations

/**
 * @module Telemetry
 * @description Decorators for instrumenting code with OpenTelemetry
 */

import { context, trace, type Context, type Span } from '@opentelemetry/api';
import { TelemetryService } from './service';

/** Type definition for any function */
type AnyFunction = (...args: any[]) => any;

/**
 * Configuration options for the Trace decorator
 */
type TraceOptions = {
  /** Custom name for the span (defaults to ClassName.methodName) */
  name?: string;
  /** Parent context for the span (for manual context propagation) */
  parentContext?: Context;
};

/**
 * Method decorator that automatically creates a span for the decorated method
 * 
 * This decorator instruments methods with OpenTelemetry tracing, creating spans
 * that track method execution, timing, and errors. It supports both synchronous
 * and asynchronous methods.
 * 
 * @param nameOrOptions - Optional span name or configuration object
 * 
 * @example
 * ```typescript
 * class UserService {
 *   // Basic usage
 *   @Trace()
 *   async getUser(id: string) {
 *     // Method implementation
 *   }
 * 
 *   // With custom span name
 *   @Trace('FetchUserDetails')
 *   async getUserDetails(id: string) {
 *     // Method implementation
 *   }
 * 
 *   // With options object
 *   @Trace({ name: 'UserAuthentication' })
 *   async authenticateUser(username: string, password: string) {
 *     // Method implementation
 *   }
 * }
 * ```
 */
export function Trace(nameOrOptions?: string | TraceOptions): any {
  // This function can be used as @Trace or @Trace(options)
  return function(
    target: any, 
    propertyKey: string | symbol, 
    descriptor: PropertyDescriptor
  ): PropertyDescriptor {
    if (!descriptor || typeof descriptor.value !== 'function') {
      throw new Error('@Trace can only be applied to method declarations');
    }
    
    const originalMethod = descriptor.value;
    const className = target.constructor.name;
    const methodName = String(propertyKey);
    
    // Determine the span name and parent context
    let spanName: string;
    let parentContext: Context | undefined;
    
    if (typeof nameOrOptions === 'string') {
      spanName = nameOrOptions;
    } else if (nameOrOptions?.name) {
      spanName = nameOrOptions.name;
      parentContext = nameOrOptions.parentContext;
    } else {
      spanName = `${className}.${methodName}`;
    }
    
    // Handle both synchronous and asynchronous methods
    const isAsync = originalMethod.constructor.name === 'AsyncFunction';
    
    if (isAsync) {
      descriptor.value = async function(...args: any[]) {
        const telemetry = TelemetryService.initialize();
        if (!telemetry.isEnabled()) {
          return originalMethod.apply(this, args);
        }

        const { span } = telemetry.startSpan(spanName, parentContext);
        
        try {
          // Record method arguments if detailed telemetry is enabled
          if (telemetry['config'].collectionLevel === 'detailed') {
            span.setAttribute('method', methodName);
            span.setAttribute('class', className);
            
            if (args && args.length > 0) {
              span.setAttribute('args', JSON.stringify(args, (_, v) => 
                typeof v === 'bigint' ? v.toString() : v
              ));
            }
          }
          
          const result = await originalMethod.apply(this, args);
          span.setStatus({ code: 1 }); // OK
          return result;
        } catch (error) {
          span.recordException(error as Error);
          span.setStatus({ 
            code: 2, 
            message: error instanceof Error ? error.message : String(error) 
          });
          throw error;
        } finally {
          telemetry.endSpan(span);
        }
      };
    } else {
      descriptor.value = function(...args: any[]) {
        const telemetry = TelemetryService.initialize();
        if (!telemetry.isEnabled()) {
          return originalMethod.apply(this, args);
        }

        const { span } = telemetry.startSpan(spanName, parentContext);
        
        try {
          // Record method arguments if detailed telemetry is enabled
          if (telemetry['config'].collectionLevel === 'detailed') {
            span.setAttribute('method', methodName);
            span.setAttribute('class', className);
            
            if (args && args.length > 0) {
              span.setAttribute('args', JSON.stringify(args, (_, v) => 
                typeof v === 'bigint' ? v.toString() : v
              ));
            }
          }
          
          const result = originalMethod.apply(this, args);
          span.setStatus({ code: 1 }); // OK
          return result;
        } catch (error) {
          span.recordException(error as Error);
          span.setStatus({ 
            code: 2, 
            message: error instanceof Error ? error.message : String(error) 
          });
          throw error;
        } finally {
          telemetry.endSpan(span);
        }
      };
    }

    return descriptor;
  };
}

/**
 * Method decorator for recording metrics when a method is called
 * 
 * This decorator records metrics each time the decorated method is called,
 * allowing for monitoring method usage, performance, and other custom metrics.
 * 
 * @param name - The name of the metric to record
 * @param value - Optional value to record (default: 1)
 * @param attributes - Optional attributes to include with the metric
 * @returns A method decorator
 * 
 * @example
 * ```typescript
 * class MessageService {
 *   // Count message sends
 *   @Metric('messages.sent')
 *   async sendMessage(message: Message) {
 *     // Implementation
 *   }
 * 
 *   // Record message size with custom value
 *   @Metric('message.size', message.content.length)
 *   async processMessage(message: Message) {
 *     // Implementation
 *   }
 * 
 *   // Include additional attributes
 *   @Metric('api.call', 1, { endpoint: '/messages' })
 *   async fetchMessages() {
 *     // Implementation
 *   }
 * }
 * ```
 */
export function Metric(name: string, value: number = 1, attributes?: Record<string, unknown>): any {
  // This is a decorator factory, it returns the actual decorator function
  return function(
    target: any, 
    propertyKey: string | symbol, 
    descriptor: PropertyDescriptor
  ): PropertyDescriptor {
    // Ensure we're decorating a method
    if (!descriptor || typeof descriptor.value !== 'function') {
      throw new Error('@Metric can only be applied to method declarations');
    }
    
    // Store the original method
    const originalMethod = descriptor.value;
    const className = target.constructor.name;
    const methodName = String(propertyKey);
    
    // Replace the original method with our instrumented version
    descriptor.value = function(...args: any[]) {
      const telemetry = TelemetryService.initialize();
      if (telemetry.isEnabled()) {
        telemetry.recordMetric(name, value, {
          ...attributes,
          class: className,
          method: methodName
        });
      }
      
      return originalMethod.apply(this, args);
    };
    
    return descriptor;
  };
}

/**
 * Helper function to apply the metric decorator logic
 * 
 * This is an internal utility function used by the Metric decorator to
 * apply consistent metric recording logic.
 * 
 * @internal
 */
function applyMetricDecorator(
  target: any, 
  propertyKey: string | symbol, 
  descriptor: PropertyDescriptor,
  name: string,
  value: number = 1,
  attributes?: Record<string, unknown>
): PropertyDescriptor {
  if (!descriptor || typeof descriptor.value !== 'function') {
    throw new Error('@Metric can only be applied to method declarations');
  }
  
  const originalMethod = descriptor.value;
  const className = target.constructor.name;
  const methodName = String(propertyKey);
  
  descriptor.value = function(...args: any[]) {
    const telemetry = TelemetryService.initialize();
    if (telemetry.isEnabled()) {
      telemetry.recordMetric(name, value, {
        ...attributes,
        class: className,
        method: methodName
      });
    }
    
    return originalMethod.apply(this, args);
  };
  
  return descriptor;
}

/**
 * Class decorator to automatically trace all methods in a class
 * 
 * This decorator applies the Trace decorator to all methods in a class,
 * making it easy to instrument an entire class without decorating each
 * method individually.
 * 
 * @param name - Optional custom name prefix for the spans (defaults to class name)
 * @returns A class decorator
 * 
 * @example
 * ```typescript
 * // Basic usage - traces all methods with ClassName.methodName
 * @TraceClass()
 * class UserService {
 *   async getUser(id: string) { 
 *     // Method implementation
 *   }
 *   async updateUser(id: string, data: any) { 
 *     // Method implementation
 *   }
 * }
 * 
 * // With custom name prefix
 * @TraceClass('Users')
 * class UserService {
 *   // Will be traced as 'Users.getUser'
 *   async getUser(id: string) { 
 *     // Method implementation
 *   }
 * }
 * ```
 */
export function TraceClass(name?: string): ClassDecorator {
  return function(constructor: Function) {
    const className = name || constructor.name;
    
    for (const propertyName of Object.getOwnPropertyNames(constructor.prototype)) {
      if (propertyName === 'constructor') continue;
      
      const descriptor = Object.getOwnPropertyDescriptor(
        constructor.prototype,
        propertyName
      );
      
      if (descriptor && typeof descriptor.value === 'function') {
        const traceDescriptor = Trace(`${className}.${propertyName}`)(
          constructor.prototype,
          propertyName,
          descriptor
        );
        
        if (traceDescriptor) {
          Object.defineProperty(
            constructor.prototype,
            propertyName,
            traceDescriptor
          );
        }
      }
    }
  };
}