anthropic-client.ts

/**
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import type Anthropic from '@anthropic-ai/sdk';
import type {
  BetaMessage,
  BetaRawMessageStreamEvent,
} from '@anthropic-ai/sdk/resources/beta/messages.mjs';
import type {
  Message,
  MessageStreamEvent,
} from '@anthropic-ai/sdk/resources/messages.mjs';
import { mock } from 'node:test';

export interface MockAnthropicClientOptions {
  messageResponse?: Partial<Message>;
  sequentialResponses?: Partial<Message>[]; // For tool calling - multiple responses
  streamChunks?: MessageStreamEvent[];
  modelList?: Array<{ id: string; display_name?: string }>;
  shouldError?: Error;
  streamErrorAfterChunk?: number; // Throw error after this many chunks
  streamError?: Error; // Error to throw during streaming
  abortSignal?: AbortSignal; // Abort signal to check
}

/**
 * Creates a mock Anthropic client for testing
 */
export function createMockAnthropicClient(
  options: MockAnthropicClientOptions = {}
): Anthropic {
  const messageResponse = {
    ...mockDefaultMessage(),
    ...options.messageResponse,
  };
  const betaMessageResponse = toBetaMessage(messageResponse);

  // Support sequential responses for tool calling workflows
  let callCount = 0;
  const createStub = options.shouldError
    ? mock.fn(async () => {
        throw options.shouldError;
      })
    : options.sequentialResponses
      ? mock.fn(async () => {
          const response =
            options.sequentialResponses![callCount] || messageResponse;
          callCount++;
          return {
            ...mockDefaultMessage(),
            ...response,
          };
        })
      : mock.fn(async () => messageResponse);

  let betaCallCount = 0;
  const betaCreateStub = options.shouldError
    ? mock.fn(async () => {
        throw options.shouldError;
      })
    : options.sequentialResponses
      ? mock.fn(async () => {
          const response =
            options.sequentialResponses![betaCallCount] || messageResponse;
          betaCallCount++;
          return toBetaMessage({
            ...mockDefaultMessage(),
            ...response,
          });
        })
      : mock.fn(async () => betaMessageResponse);

  const streamStub = options.shouldError
    ? mock.fn(() => {
        throw options.shouldError;
      })
    : mock.fn((_body: any, opts?: { signal?: AbortSignal }) => {
        // Check abort signal before starting stream
        if (opts?.signal?.aborted) {
          throw new Error('AbortError');
        }
        return createMockStream(
          options.streamChunks || [],
          messageResponse as Message,
          options.streamErrorAfterChunk,
          options.streamError,
          opts?.signal
        );
      });

  const betaStreamStub = options.shouldError
    ? mock.fn(() => {
        throw options.shouldError;
      })
    : mock.fn((_body: any, opts?: { signal?: AbortSignal }) => {
        if (opts?.signal?.aborted) {
          throw new Error('AbortError');
        }
        const betaChunks = (options.streamChunks || []).map((chunk) =>
          toBetaStreamEvent(chunk)
        );
        return createMockStream(
          betaChunks,
          toBetaMessage(messageResponse),
          options.streamErrorAfterChunk,
          options.streamError,
          opts?.signal
        );
      });

  const listStub = options.shouldError
    ? mock.fn(async () => {
        throw options.shouldError;
      })
    : mock.fn(async () => ({
        data: options.modelList || mockDefaultModels(),
      }));

  return {
    messages: {
      create: createStub,
      stream: streamStub,
    },
    models: {
      list: listStub,
    },
    beta: {
      messages: {
        create: betaCreateStub,
        stream: betaStreamStub,
      },
    },
  } as unknown as Anthropic;
}

/**
 * Creates a mock async iterable stream for streaming responses
 */
function createMockStream<TMessageType, TEventType>(
  chunks: TEventType[],
  finalMsg: TMessageType,
  errorAfterChunk?: number,
  streamError?: Error,
  abortSignal?: AbortSignal
) {
  let index = 0;
  return {
    [Symbol.asyncIterator]() {
      return {
        async next() {
          // Check abort signal
          if (abortSignal?.aborted) {
            const error = new Error('AbortError');
            error.name = 'AbortError';
            throw error;
          }

          // Check if we should throw an error after this chunk
          if (
            errorAfterChunk !== undefined &&
            streamError &&
            index >= errorAfterChunk
          ) {
            throw streamError;
          }

          if (index < chunks.length) {
            return { value: chunks[index++] as TEventType, done: false };
          }
          return { value: undefined as unknown as TEventType, done: true };
        },
      };
    },
    async finalMessage() {
      // Check abort signal before returning final message
      if (abortSignal?.aborted) {
        const error = new Error('AbortError');
        error.name = 'AbortError';
        throw error;
      }
      return finalMsg as TMessageType;
    },
  };
}

export interface CreateMockAnthropicMessageOptions {
  id?: string;
  text?: string;
  toolUse?: {
    id?: string;
    name: string;
    input: any;
  };
  stopReason?: Message['stop_reason'];
  usage?: Partial<Message['usage']>;
}

/**
 * Creates a customizable mock Anthropic Message response
 *
 * @example
 * // Simple text response
 * createMockAnthropicMessage({ text: 'Hi there!' })
 *
 * // Tool use response
 * createMockAnthropicMessage({
 *   toolUse: { name: 'get_weather', input: { city: 'NYC' } }
 * })
 *
 * // Custom usage
 * createMockAnthropicMessage({ usage: { input_tokens: 5, output_tokens: 15 } })
 */
export function createMockAnthropicMessage(
  options: CreateMockAnthropicMessageOptions = {}
): Message {
  const content: Message['content'] = [];

  if (options.toolUse) {
    content.push({
      type: 'tool_use',
      id: options.toolUse.id || 'toolu_test123',
      name: options.toolUse.name,
      input: options.toolUse.input,
    });
  } else {
    content.push({
      type: 'text',
      text: options.text || 'Hello! How can I help you today?',
      citations: null,
    });
  }

  const usage: Message['usage'] = {
    cache_creation: null,
    cache_creation_input_tokens: 0,
    cache_read_input_tokens: 0,
    input_tokens: 10,
    output_tokens: 20,
    server_tool_use: null,
    service_tier: null,
    ...(options.usage ?? {}),
  };

  return {
    id: options.id || 'msg_test123',
    type: 'message',
    role: 'assistant',
    model: 'claude-3-5-sonnet-20241022',
    content,
    stop_reason:
      options.stopReason || (options.toolUse ? 'tool_use' : 'end_turn'),
    stop_sequence: null,
    usage,
  };
}

/**
 * Creates a default mock Message response
 */
export function mockDefaultMessage(): Message {
  return createMockAnthropicMessage();
}

/**
 * Creates a mock text content block chunk event
 */
export function mockTextChunk(text: string): MessageStreamEvent {
  return {
    type: 'content_block_delta',
    index: 0,
    delta: {
      type: 'text_delta',
      text,
    },
  } as MessageStreamEvent;
}

/**
 * Creates a mock content block start event with text
 */
export function mockContentBlockStart(text: string): MessageStreamEvent {
  return {
    type: 'content_block_start',
    index: 0,
    content_block: {
      type: 'text',
      text,
    },
  } as MessageStreamEvent;
}

/**
 * Creates a mock tool use content block
 */
export function mockToolUseChunk(
  id: string,
  name: string,
  input: any
): MessageStreamEvent {
  return {
    type: 'content_block_start',
    index: 0,
    content_block: {
      type: 'tool_use',
      id,
      name,
      input,
    },
  } as MessageStreamEvent;
}

/**
 * Creates a default list of mock models
 */
export function mockDefaultModels() {
  return [
    { id: 'claude-3-5-sonnet-20241022', display_name: 'Claude 3.5 Sonnet' },
    { id: 'claude-3-5-haiku-20241022', display_name: 'Claude 3.5 Haiku' },
    { id: 'claude-3-opus-20240229', display_name: 'Claude 3 Opus' },
  ];
}

/**
 * Creates a mock Message with tool use
 */
export function mockMessageWithToolUse(
  toolName: string,
  toolInput: any
): Partial<Message> {
  return {
    content: [
      {
        type: 'tool_use',
        id: 'toolu_test123',
        name: toolName,
        input: toolInput,
      },
    ],
    stop_reason: 'tool_use',
  };
}

/**
 * Creates a mock Message with custom content
 */
export function mockMessageWithContent(
  content: Message['content']
): Partial<Message> {
  return {
    content,
    stop_reason: 'end_turn',
  };
}

function toBetaMessage(message: Message): BetaMessage {
  // @ts-ignore
  return {
    ...message,
    container: null,
    context_management: null,
    usage: {
      cache_creation: message.usage.cache_creation,
      cache_creation_input_tokens: message.usage.cache_creation_input_tokens,
      cache_read_input_tokens: message.usage.cache_read_input_tokens,
      input_tokens: message.usage.input_tokens,
      output_tokens: message.usage.output_tokens,
      server_tool_use: message.usage.server_tool_use as any,
      service_tier: message.usage.service_tier,
    },
  };
}

function toBetaStreamEvent(
  event: MessageStreamEvent
): BetaRawMessageStreamEvent {
  return event as unknown as BetaRawMessageStreamEvent;
}