buster/packages/ai/tests/utils/retry/healing-proof-unit.test.ts

import type { CoreMessage } from 'ai';
import { NoSuchToolError } from 'ai';
import { describe, expect, it, vi } from 'vitest';
import { retryableAgentStreamWithHealing } from '../../../src/utils/retry';
import type { MastraAgent } from '../../../src/utils/retry/types';

/**
 * Unit test to prove healing mechanism works without running full agents
 */
describe('Healing Mechanism Unit Test - Definitive Proof', () => {
  it('PROOF: onError callback receives tool errors and returns healing response', async () => {
    // Create a mock agent that will capture the onError callback
    let capturedOnError: ((error: unknown) => unknown) | undefined;
    let simulatedToolError: NoSuchToolError | undefined = undefined;

    const mockAgent: MastraAgent = {
      name: 'test-agent',
      description: 'Test agent',
      model: {} as any,
      tools: {
        'valid-tool': {} as any,
        'another-tool': {} as any,
      },
      stream: vi.fn().mockImplementation(async (messages, options) => {
        // Capture the onError callback from options
        capturedOnError = options?.onError;

        // Create a mock stream that will trigger the error
        const mockStream = {
          fullStream: {
            async *[Symbol.asyncIterator]() {
              // Yield some initial chunks
              yield { type: 'text-delta', text: 'Starting...' };

              // Now simulate a tool error occurring during streaming
              if (capturedOnError && simulatedToolError) {
                const healingResponse = capturedOnError(simulatedToolError);

                // The healing response should be returned to the LLM
                expect(healingResponse).toBeDefined();
                expect(healingResponse).toHaveProperty('error');

                // Yield the healing response as if it were a tool result
                yield {
                  type: 'tool-result',
                  toolCallId: simulatedToolError.toolCallId,
                  toolName: simulatedToolError.toolName,
                  result: healingResponse,
                };
              }

              // Continue with more chunks after healing
              yield { type: 'text-delta', text: 'Continuing after healing...' };
            },
          },
        };

        return mockStream as any;
      }),
    } as any;

    const messages: CoreMessage[] = [{ role: 'user', content: 'Test message' }];

    // Create the tool error we'll simulate
    simulatedToolError = new NoSuchToolError({
      toolName: 'non-existent-tool',
      availableTools: ['valid-tool', 'another-tool'],
    });
    (simulatedToolError as any).toolCallId = 'test-call-123';

    let healingOccurred = false;
    let healingMessage = '';

    const result = await retryableAgentStreamWithHealing({
      agent: mockAgent,
      messages,
      options: {
        toolCallStreaming: true,
        runtimeContext: {} as any, // Mock context
      },
      retryConfig: {
        maxRetries: 3,
        onRetry: (error, attempt) => {
          healingOccurred = true;
          if (error.healingMessage.role === 'tool' && Array.isArray(error.healingMessage.content)) {
            const toolResult = error.healingMessage.content[0];
            if ('result' in toolResult && toolResult.result && 'error' in toolResult.result) {
              healingMessage = toolResult.result.error as string;
            }
          }
        },
      },
    });

    // Verify the stream was created
    expect(result.stream).toBeDefined();
    expect(mockAgent.stream).toHaveBeenCalledTimes(1);

    // Verify onError callback was passed to the agent
    expect(capturedOnError).toBeDefined();

    // Test the onError callback directly
    const errorResponse = capturedOnError!(simulatedToolError);

    // THIS IS THE PROOF: The error response is returned to heal the agent
    expect(errorResponse).toEqual({
      error:
        'Tool "non-existent-tool" is not available. Available tools: valid-tool, another-tool. Please use one of the available tools instead.',
    });

    // Verify healing was tracked
    expect(healingOccurred).toBe(true);
    expect(healingMessage).toContain('Tool "non-existent-tool" is not available');
    expect(healingMessage).toContain('valid-tool');
    expect(healingMessage).toContain('another-tool');

    console.log('✅ HEALING MECHANISM PROVEN TO WORK!');
    console.log('   - onError callback successfully captures tool errors');
    console.log('   - Healing response is returned to the LLM');
    console.log('   - Agent can continue execution after healing');
  });

  it('PROOF: Multiple error types are handled correctly', async () => {
    const mockAgent: MastraAgent = {
      stream: vi.fn().mockResolvedValue({ fullStream: [] }),
    } as any;

    let capturedOnError: ((error: unknown) => unknown) | undefined;

    await retryableAgentStreamWithHealing({
      agent: mockAgent,
      messages: [{ role: 'user', content: 'Test' }],
      options: {
        runtimeContext: {} as any,
        onError: (cb) => {
          capturedOnError = cb;
        },
      },
    });

    // Get the onError from the mock call
    const mockCall = vi.mocked(mockAgent.stream).mock.calls[0];
    capturedOnError = mockCall[1]?.onError;

    expect(capturedOnError).toBeDefined();

    // Test 1: NoSuchToolError
    const noSuchToolError = new NoSuchToolError({
      toolName: 'bad-tool',
      availableTools: ['good-tool'],
    });
    (noSuchToolError as any).toolCallId = 'call1';

    const response1 = capturedOnError!(noSuchToolError);
    expect(response1).toHaveProperty('error');
    expect((response1 as any).error).toContain('Tool "bad-tool" is not available');

    // Test 2: InvalidToolArgumentsError
    const invalidArgsError = new Error('Invalid args');
    invalidArgsError.name = 'AI_InvalidToolArgumentsError';
    (invalidArgsError as any).toolCallId = 'call2';
    (invalidArgsError as any).toolName = 'some-tool';
    (invalidArgsError as any).cause = {
      errors: [
        { path: ['field1'], message: 'Required' },
        { path: ['field2'], message: 'Must be number' },
      ],
    };

    const response2 = capturedOnError!(invalidArgsError);
    expect(response2).toHaveProperty('error');
    expect((response2 as any).error).toContain('Invalid tool arguments');
    expect((response2 as any).error).toContain('field1: Required');
    expect((response2 as any).error).toContain('field2: Must be number');

    // Test 3: Non-healable error returns undefined
    const genericError = new Error('Some other error');
    const response3 = capturedOnError!(genericError);
    expect(response3).toBeUndefined();

    console.log('✅ ALL ERROR TYPES HANDLED CORRECTLY!');
  });

  it('PROOF: Healing attempts are limited by maxRetries', async () => {
    const mockAgent: MastraAgent = {
      stream: vi.fn().mockResolvedValue({ fullStream: [] }),
    } as any;

    let healingAttempts = 0;

    await retryableAgentStreamWithHealing({
      agent: mockAgent,
      messages: [{ role: 'user', content: 'Test' }],
      options: { runtimeContext: {} as any },
      retryConfig: {
        maxRetries: 2,
        onRetry: () => {
          healingAttempts++;
        },
      },
    });

    const capturedOnError = vi.mocked(mockAgent.stream).mock.calls[0][1]?.onError;
    expect(capturedOnError).toBeDefined();

    const error = new NoSuchToolError({
      toolName: 'test',
      availableTools: [],
    });

    // First attempt - should heal
    let response = capturedOnError!(error);
    expect(response).toBeDefined();
    expect(healingAttempts).toBe(1);

    // Second attempt - should heal
    response = capturedOnError!(error);
    expect(response).toBeDefined();
    expect(healingAttempts).toBe(2);

    // Third attempt - should NOT heal (exceeds maxRetries)
    response = capturedOnError!(error);
    expect(response).toBeUndefined();
    expect(healingAttempts).toBe(2); // No additional retry

    console.log('✅ MAX RETRIES LIMIT PROVEN TO WORK!');
  });
});