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

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

/**
 * @module ValidationDecorators
 * @description Method decorators for parameter and return value validation
 * 
 * This module provides TypeScript decorators for validating method parameters
 * and return values using Zod schemas. These decorators help ensure that data
 * flowing through the A2A protocol conforms to the expected schemas.
 */

import { z } from 'zod';
import { A2AError, ERROR_CODES } from '../errors';

/**
 * Method decorator for validating the first parameter of a method
 * 
 * This decorator validates the first parameter of a method against a Zod schema.
 * If validation fails, it throws an A2AError with details about the validation
 * failure. This is useful for ensuring that method parameters conform to the
 * expected schema before executing the method logic.
 * 
 * @param schema - Zod schema to validate the first parameter against
 * @returns Method decorator function
 * 
 * @example
 * ```typescript
 * // Define a schema for task parameters
 * const TaskParamsSchema = z.object({
 *   name: z.string(),
 *   description: z.string().optional(),
 *   priority: z.number().min(1).max(5).optional()
 * });
 * 
 * class TaskManager {
 *   // Apply the decorator to validate the first parameter
 *   @ValidateParams(TaskParamsSchema)
 *   createTask(params: any) {
 *     // This code only runs if params passes validation
 *     return { id: generateId(), ...params };
 *   }
 * }
 * ```
 */
export function ValidateParams(schema: z.ZodSchema) {
  return function(target: any, propertyKey: string, descriptor: PropertyDescriptor) {
    const originalMethod = descriptor.value;
    
    descriptor.value = function(...args: any[]) {
      try {
        schema.parse(args[0]); // Validate first argument
        return originalMethod.apply(this, args);
      } catch (err) {
        throw new A2AError('Invalid parameters', ERROR_CODES.INVALID_TASK_STATE, { validationError: err });
      }
    };
  };
}

/**
 * Method decorator for validating the return value of a method
 * 
 * This decorator validates the return value of a method against a Zod schema.
 * If validation fails, it throws an A2AError with details about the validation
 * failure. This is useful for ensuring that method return values conform to the
 * expected schema before they're passed to other components.
 * 
 * @param schema - Zod schema to validate the return value against
 * @returns Method decorator function
 * 
 * @example
 * ```typescript
 * // Define a schema for task objects
 * const TaskSchema = z.object({
 *   id: z.string(),
 *   name: z.string(),
 *   status: z.enum(['created', 'in_progress', 'completed', 'failed']),
 *   createdAt: z.string().datetime(),
 *   updatedAt: z.string().datetime()
 * });
 * 
 * class TaskManager {
 *   // Apply the decorator to validate the return value
 *   @ValidateReturn(TaskSchema)
 *   async getTask(id: string) {
 *     // Fetch task from database
 *     const task = await db.tasks.findOne({ id });
 *     return task; // Will be validated before being returned
 *   }
 * }
 * ```
 */
export function ValidateReturn(schema: z.ZodSchema) {
  return function(target: any, propertyKey: string, descriptor: PropertyDescriptor) {
    const originalMethod = descriptor.value;
    
    descriptor.value = async function(...args: any[]) {
      const result = await originalMethod.apply(this, args);
      try {
        return schema.parse(result);
      } catch (err) {
        throw new A2AError('Invalid return value', ERROR_CODES.TASK_NOT_FOUND, { validationError: err });
      }
    };
  };
}

/**
 * Method decorator for validating a task parameter
 * 
 * This decorator validates the first parameter of a method (expected to be a task)
 * against a Zod schema. If validation fails, it throws an A2AError with details
 * about the validation failure. This is specifically designed for methods that
 * operate on task objects and need to ensure the task is in a valid state.
 * 
 * @param schema - Zod schema to validate the task against
 * @returns Method decorator function
 * 
 * @example
 * ```typescript
 * // Define a schema for tasks in a specific state
 * const WorkingTaskSchema = TaskSchema.extend({
 *   status: z.literal('in_progress')
 * });
 * 
 * class TaskProcessor {
 *   // Apply the decorator to validate the task parameter
 *   @ValidateTask(WorkingTaskSchema)
 *   processTask(task: any, options: ProcessOptions) {
 *     // This code only runs if task passes validation
 *     // We know the task is in 'in_progress' status
 *     return this.performProcessing(task, options);
 *   }
 * }
 * ```
 */
export function ValidateTask(schema: z.ZodSchema) {
  return function(target: any, propertyKey: string, descriptor: PropertyDescriptor) {
    const originalMethod = descriptor.value;
    
    descriptor.value = function(task: any, ...args: any[]) {
      try {
        schema.parse(task);
        return originalMethod.apply(this, [task, ...args]);
      } catch (err) {
        throw new A2AError('Invalid task state', ERROR_CODES.TASK_ALREADY_COMPLETED, { validationError: err });
      }
    };
  };
}