flynn/src/models/bedrock.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { BedrockClient } from './bedrock.js';
import type { ChatStreamEvent } from './types.js';

const mockSend = vi.fn().mockResolvedValue({
  output: {
    message: {
      content: [{ text: 'Hello from Bedrock!' }],
    },
  },
  stopReason: 'end_turn',
  usage: { inputTokens: 10, outputTokens: 5 },
});

vi.mock('@aws-sdk/client-bedrock-runtime', () => ({
  BedrockRuntimeClient: vi.fn().mockImplementation(() => ({
    send: mockSend,
  })),
  ConverseCommand: vi.fn().mockImplementation((params) => params),
  ConverseStreamCommand: vi.fn().mockImplementation((params) => params),
}));

describe('BedrockClient', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockSend.mockResolvedValue({
      output: {
        message: {
          content: [{ text: 'Hello from Bedrock!' }],
        },
      },
      stopReason: 'end_turn',
      usage: { inputTokens: 10, outputTokens: 5 },
    });
  });

  it('sends messages and returns response', async () => {
    const client = new BedrockClient({
      model: 'anthropic.claude-3-sonnet',
      region: 'us-east-1',
    });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    expect(response.content).toBe('Hello from Bedrock!');
    expect(response.stopReason).toBe('end_turn');
    expect(response.usage.inputTokens).toBe(10);
    expect(response.usage.outputTokens).toBe(5);
  });

  it('parses tool use response', async () => {
    mockSend.mockResolvedValueOnce({
      output: {
        message: {
          content: [{
            toolUse: {
              toolUseId: 'tool_01',
              name: 'shell.exec',
              input: { command: 'ls' },
            },
          }],
        },
      },
      stopReason: 'tool_use',
      usage: { inputTokens: 20, outputTokens: 15 },
    });

    const client = new BedrockClient({
      model: 'anthropic.claude-3-sonnet',
      region: 'us-east-1',
    });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'list files' }],
      tools: [{
        name: 'shell.exec',
        description: 'Run shell command',
        input_schema: { type: 'object', properties: { command: { type: 'string' } }, required: ['command'] },
      }],
    });

    expect(response.stopReason).toBe('tool_use');
    expect(response.toolCalls).toHaveLength(1);
    expect(response.toolCalls![0].name).toBe('shell.exec');
    expect(response.toolCalls![0].args).toEqual({ command: 'ls' });
  });

  it('uses default region when none provided', async () => {
    const client = new BedrockClient({
      model: 'anthropic.claude-3-sonnet',
    });

    const response = await client.chat({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    expect(response.content).toBe('Hello from Bedrock!');
  });

  it('passes system prompt to API', async () => {
    const client = new BedrockClient({
      model: 'anthropic.claude-3-sonnet',
      region: 'us-east-1',
    });

    await client.chat({
      messages: [{ role: 'user', content: 'Hello' }],
      system: 'You are a helpful assistant.',
    });

    expect(mockSend).toHaveBeenCalledTimes(1);
    // ConverseCommand is called with params that include system
    const { ConverseCommand } = await import('@aws-sdk/client-bedrock-runtime');
    expect(ConverseCommand).toHaveBeenCalledWith(
      expect.objectContaining({
        system: [{ text: 'You are a helpful assistant.' }],
      }),
    );
  });
});

describe('BedrockClient streaming', () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  it('streams content events', async () => {
    mockSend.mockResolvedValueOnce({
      stream: (async function* () {
        yield { contentBlockDelta: { delta: { text: 'Hello ' } } };
        yield { contentBlockDelta: { delta: { text: 'from Bedrock!' } } };
        yield { metadata: { usage: { inputTokens: 10, outputTokens: 5 } } };
      })(),
    });

    const client = new BedrockClient({
      model: 'anthropic.claude-3-sonnet',
      region: 'us-east-1',
    });

    const chunks: string[] = [];
    let finalUsage: { inputTokens: number; outputTokens: number } | undefined;

    for await (const event of client.chatStream({
      messages: [{ role: 'user', content: 'Hello' }],
    })) {
      if (event.type === 'content' && event.content) {
        chunks.push(event.content);
      }
      if (event.type === 'done' && event.usage) {
        finalUsage = event.usage;
      }
    }

    expect(chunks.join('')).toBe('Hello from Bedrock!');
    expect(finalUsage).toEqual({ inputTokens: 10, outputTokens: 5 });
  });

  it('yields error event on failure', async () => {
    mockSend.mockRejectedValueOnce(new Error('Service unavailable'));

    const client = new BedrockClient({
      model: 'anthropic.claude-3-sonnet',
      region: 'us-east-1',
    });

    const events: ChatStreamEvent[] = [];
    for await (const event of client.chatStream({
      messages: [{ role: 'user', content: 'Hello' }],
    })) {
      events.push(event);
    }

    expect(events).toHaveLength(1);
    expect(events[0].type).toBe('error');
    expect(events[0].error?.message).toBe('Service unavailable');
  });
});