chore: import upstream snapshot with attribution

2026-07-13 13:24:08 +08:00
commit 0d3cb498a3
5438 changed files with 1316560 additions and 0 deletions
@@ -0,0 +1,782 @@
+import { NodeHttpHandler } from '@smithy/node-http-handler';
+import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import logger from '../../../src/logger';
+import { AwsBedrockKnowledgeBaseProvider } from '../../../src/providers/bedrock/knowledgeBase';
+import { sha256 } from '../../../src/util/createHash';
+import { createEmptyTokenUsage } from '../../../src/util/tokenUsageUtils';
+import { mockProcessEnv } from '../../util/utils';
+
+const mockSend = vi.fn();
+const mockBedrockClient = {
+  send: mockSend,
+};
+
+vi.mock('@aws-sdk/client-bedrock-agent-runtime', async (importOriginal) => {
+  return {
+    ...(await importOriginal()),
+
+    BedrockAgentRuntimeClient: vi.fn().mockImplementation(function () {
+      return mockBedrockClient;
+    }),
+
+    RetrieveAndGenerateCommand: vi.fn().mockImplementation(function (params) {
+      return params;
+    }),
+  };
+});
+
+// Module imports - loaded in beforeAll
+let BedrockAgentRuntimeClient: typeof import('@aws-sdk/client-bedrock-agent-runtime').BedrockAgentRuntimeClient;
+let RetrieveAndGenerateCommand: typeof import('@aws-sdk/client-bedrock-agent-runtime').RetrieveAndGenerateCommand;
+const NodeHttpHandlerMock = vi.mocked(NodeHttpHandler);
+
+// Mock @smithy/node-http-handler with ESM-compatible exports
+vi.mock('@smithy/node-http-handler', () => ({
+  __esModule: true,
+  NodeHttpHandler: vi.fn().mockImplementation(function () {
+    return {
+      handle: vi.fn(),
+    };
+  }),
+  default: vi.fn().mockImplementation(function () {
+    return {
+      handle: vi.fn(),
+    };
+  }),
+}));
+
+// Mock proxy-agent with ESM-compatible exports
+vi.mock('proxy-agent', () => ({
+  __esModule: true,
+  ProxyAgent: vi.fn(function ProxyAgentMock() {}),
+  default: vi.fn(function ProxyAgentMock() {}),
+}));
+
+vi.mock('../../../src/logger', () => ({
+  default: {
+    debug: vi.fn(),
+    error: vi.fn(),
+    warn: vi.fn(),
+  },
+}));
+
+const mockGet = vi.hoisted(() => vi.fn());
+
+const mockSet = vi.hoisted(() => vi.fn());
+
+const mockIsCacheEnabled = vi.fn().mockReturnValue(false);
+
+function buildKnowledgeBaseCacheKey({
+  knowledgeBaseId,
+  modelArn,
+  modelName,
+  prompt,
+  region,
+  kbConfig,
+}: {
+  knowledgeBaseId: string;
+  modelArn?: string;
+  modelName: string;
+  prompt: string;
+  region: string;
+  kbConfig: Record<string, unknown>;
+}) {
+  const cacheConfig = {
+    region,
+    modelName,
+    ...Object.fromEntries(
+      Object.entries(kbConfig).filter(
+        ([key]) => !['accessKeyId', 'secretAccessKey', 'sessionToken'].includes(key),
+      ),
+    ),
+  };
+  const configStr = JSON.stringify(cacheConfig, Object.keys(cacheConfig).sort());
+
+  return `bedrock-kb:${knowledgeBaseId}:${modelArn}:${region}:${sha256(
+    JSON.stringify({
+      configStr,
+      prompt,
+    }),
+  )}`;
+}
+
+vi.mock('../../../src/cache', async (importOriginal) => {
+  return {
+    ...(await importOriginal()),
+
+    getCache: vi.fn().mockImplementation(function () {
+      return {
+        get: mockGet,
+        set: mockSet,
+      };
+    }),
+
+    isCacheEnabled: () => mockIsCacheEnabled(),
+  };
+});
+
+describe('AwsBedrockKnowledgeBaseProvider', () => {
+  beforeAll(async () => {
+    const bedrockModule = await import('@aws-sdk/client-bedrock-agent-runtime');
+    BedrockAgentRuntimeClient = bedrockModule.BedrockAgentRuntimeClient;
+    RetrieveAndGenerateCommand = bedrockModule.RetrieveAndGenerateCommand;
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockGet.mockReset();
+    mockSet.mockReset();
+    mockIsCacheEnabled.mockReset().mockReturnValue(false);
+    mockProcessEnv({ AWS_BEDROCK_MAX_RETRIES: undefined });
+    mockProcessEnv({ AWS_BEARER_TOKEN_BEDROCK: undefined });
+    mockProcessEnv({ HTTPS_PROXY: undefined });
+    mockProcessEnv({ https_proxy: undefined });
+    mockProcessEnv({ HTTP_PROXY: undefined });
+    mockProcessEnv({ http_proxy: undefined });
+    mockProcessEnv({ npm_config_https_proxy: undefined });
+    mockProcessEnv({ npm_config_http_proxy: undefined });
+    mockProcessEnv({ npm_config_proxy: undefined });
+    mockProcessEnv({ all_proxy: undefined });
+  });
+
+  afterEach(() => {
+    vi.clearAllMocks();
+    mockProcessEnv({ AWS_BEARER_TOKEN_BEDROCK: undefined });
+    mockProcessEnv({ HTTPS_PROXY: undefined });
+    mockProcessEnv({ https_proxy: undefined });
+    mockProcessEnv({ HTTP_PROXY: undefined });
+    mockProcessEnv({ http_proxy: undefined });
+    mockProcessEnv({ npm_config_https_proxy: undefined });
+    mockProcessEnv({ npm_config_http_proxy: undefined });
+    mockProcessEnv({ npm_config_proxy: undefined });
+    mockProcessEnv({ all_proxy: undefined });
+  });
+
+  it('should throw an error if knowledgeBaseId is not provided', () => {
+    expect(() => {
+      new AwsBedrockKnowledgeBaseProvider('us.anthropic.claude-3-7-sonnet-20241022-v2:0', {
+        config: {} as any,
+      });
+    }).toThrow('Knowledge Base ID is required');
+  });
+
+  it('should create provider with required options', () => {
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    expect(provider).toBeDefined();
+    expect(provider.kbConfig.knowledgeBaseId).toBe('kb-123');
+    expect(provider.getRegion()).toBe('us-east-1');
+  });
+
+  it('should create knowledge base client without proxy settings', async () => {
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    await provider.getKnowledgeBaseClient();
+
+    // client-bedrock-agent-runtime already defaults to HTTP/1.1,
+    // so no custom handler is needed without proxy or apiKey
+    expect(NodeHttpHandlerMock).not.toHaveBeenCalled();
+    expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({
+      region: 'us-east-1',
+      retryMode: 'adaptive',
+      maxAttempts: 10,
+    });
+  });
+
+  it('should create knowledge base client with credentials', async () => {
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          accessKeyId: 'test-access-key',
+          secretAccessKey: 'test-secret-key',
+        },
+      },
+    );
+
+    await provider.getKnowledgeBaseClient();
+
+    expect(NodeHttpHandlerMock).not.toHaveBeenCalled();
+    expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({
+      region: 'us-east-1',
+      retryMode: 'adaptive',
+      maxAttempts: 10,
+      credentials: {
+        accessKeyId: 'test-access-key',
+        secretAccessKey: 'test-secret-key',
+      },
+    });
+  });
+
+  it('should respect AWS_BEDROCK_MAX_RETRIES environment variable', async () => {
+    mockProcessEnv({ AWS_BEDROCK_MAX_RETRIES: '5' });
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    await provider.getKnowledgeBaseClient();
+
+    expect(NodeHttpHandlerMock).not.toHaveBeenCalled();
+    expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({
+      region: 'us-east-1',
+      retryMode: 'adaptive',
+      maxAttempts: 5,
+    });
+  });
+
+  it('should call the knowledge base API with correct parameters', async () => {
+    const mockResponse = {
+      output: {
+        text: 'This is the response from the knowledge base',
+      },
+      citations: [
+        {
+          retrievedReferences: [
+            {
+              content: {
+                text: 'This is a citation',
+              },
+              location: {
+                type: 's3',
+                s3Location: {
+                  uri: 's3://bucket/key',
+                },
+              },
+            },
+          ],
+          generatedResponsePart: {
+            textResponsePart: {
+              text: 'part of the response',
+              span: {
+                start: 0,
+                end: 10,
+              },
+            },
+          },
+        },
+      ],
+    };
+
+    mockSend.mockResolvedValueOnce(mockResponse);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    const result = await provider.callApi('What is the capital of France?');
+
+    const expectedCommand = {
+      input: { text: 'What is the capital of France?' },
+      retrieveAndGenerateConfiguration: {
+        type: 'KNOWLEDGE_BASE',
+        knowledgeBaseConfiguration: {
+          knowledgeBaseId: 'kb-123',
+          modelArn: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+        },
+      },
+    };
+
+    expect(RetrieveAndGenerateCommand).toHaveBeenCalledWith(expectedCommand);
+    expect(mockSend).toHaveBeenCalledWith(expectedCommand);
+    expect(result).toEqual({
+      output: 'This is the response from the knowledge base',
+      metadata: { citations: mockResponse.citations },
+      tokenUsage: createEmptyTokenUsage(),
+    });
+  });
+
+  it('should handle API errors gracefully', async () => {
+    mockSend.mockRejectedValueOnce(new Error('API error'));
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    const result = await provider.callApi('What is the capital of France?');
+
+    expect(result).toEqual({
+      error: 'Bedrock Knowledge Base API error: Error: API error',
+    });
+  });
+
+  it('should use custom modelArn if provided', async () => {
+    const mockResponse = {
+      output: {
+        text: 'This is the response from the knowledge base',
+      },
+      citations: [],
+    };
+
+    mockSend.mockResolvedValueOnce(mockResponse);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider('amazon.nova-lite-v1:0', {
+      config: {
+        knowledgeBaseId: 'kb-123',
+        region: 'us-east-1',
+        modelArn: 'custom:model:arn',
+      },
+    });
+
+    await provider.callApi('What is the capital of France?');
+
+    const expectedCommand = {
+      input: { text: 'What is the capital of France?' },
+      retrieveAndGenerateConfiguration: {
+        type: 'KNOWLEDGE_BASE',
+        knowledgeBaseConfiguration: {
+          knowledgeBaseId: 'kb-123',
+          modelArn: 'custom:model:arn',
+        },
+      },
+    };
+
+    expect(RetrieveAndGenerateCommand).toHaveBeenCalledWith(expectedCommand);
+  });
+
+  it('should pass along config parameters but not create generationConfiguration', async () => {
+    const mockResponse = {
+      output: {
+        text: 'This is the response from the knowledge base',
+      },
+      citations: [],
+    };
+
+    mockSend.mockResolvedValueOnce(mockResponse);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider('amazon.nova-lite-v1:0', {
+      config: {
+        knowledgeBaseId: 'kb-123',
+        region: 'us-east-1',
+      } as any,
+    });
+
+    await provider.callApi('What is the capital of France?');
+
+    const expectedCommand = {
+      input: { text: 'What is the capital of France?' },
+      retrieveAndGenerateConfiguration: {
+        type: 'KNOWLEDGE_BASE',
+        knowledgeBaseConfiguration: {
+          knowledgeBaseId: 'kb-123',
+          modelArn: 'arn:aws:bedrock:us-east-1::foundation-model/amazon.nova-lite-v1:0',
+        },
+      },
+    };
+
+    expect(RetrieveAndGenerateCommand).toHaveBeenCalledWith(expectedCommand);
+  });
+
+  it('should not include retrievalConfiguration when numberOfResults is not provided', async () => {
+    const mockResponse = {
+      output: {
+        text: 'This is the response from the knowledge base',
+      },
+      citations: [],
+    };
+
+    mockSend.mockResolvedValueOnce(mockResponse);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    await provider.callApi('What is the capital of France?');
+
+    const expectedCommand = {
+      input: { text: 'What is the capital of France?' },
+      retrieveAndGenerateConfiguration: {
+        type: 'KNOWLEDGE_BASE',
+        knowledgeBaseConfiguration: {
+          knowledgeBaseId: 'kb-123',
+          modelArn: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+        },
+      },
+    };
+
+    expect(RetrieveAndGenerateCommand).toHaveBeenCalledWith(expectedCommand);
+  });
+
+  it('should use custom numberOfResults when provided', async () => {
+    const mockResponse = {
+      output: {
+        text: 'This is the response from the knowledge base',
+      },
+      citations: [],
+    };
+
+    mockSend.mockResolvedValueOnce(mockResponse);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          numberOfResults: 10,
+        },
+      },
+    );
+
+    await provider.callApi('What is the capital of France?');
+
+    const expectedCommand = {
+      input: { text: 'What is the capital of France?' },
+      retrieveAndGenerateConfiguration: {
+        type: 'KNOWLEDGE_BASE',
+        knowledgeBaseConfiguration: {
+          knowledgeBaseId: 'kb-123',
+          modelArn: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+          retrievalConfiguration: {
+            vectorSearchConfiguration: {
+              numberOfResults: 10,
+            },
+          },
+        },
+      },
+    };
+
+    expect(RetrieveAndGenerateCommand).toHaveBeenCalledWith(expectedCommand);
+  });
+
+  it('should retrieve citations from cache when available', async () => {
+    mockIsCacheEnabled.mockReturnValue(true);
+
+    const cachedResponse = JSON.stringify({
+      output: 'Cached response from knowledge base',
+      citations: [
+        {
+          retrievedReferences: [
+            {
+              content: { text: 'Citation from cache' },
+              location: { s3Location: { uri: 'https://example.com/cached' } },
+            },
+          ],
+        },
+      ],
+    });
+
+    mockGet.mockResolvedValueOnce(cachedResponse);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    const result = await provider.callApi('What is the capital of France?');
+
+    const cacheKey = mockGet.mock.calls[0][0];
+
+    expect(cacheKey).toBe(
+      buildKnowledgeBaseCacheKey({
+        knowledgeBaseId: 'kb-123',
+        modelArn: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+        modelName: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+        prompt: 'What is the capital of France?',
+        region: 'us-east-1',
+        kbConfig: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      }),
+    );
+    expect(cacheKey).not.toContain('What is the capital of France?');
+    const cacheHitLog = vi
+      .mocked(logger.debug)
+      .mock.calls.find(
+        ([message]) => message === 'Returning cached Bedrock Knowledge Base response',
+      );
+    expect(cacheHitLog).toBeDefined();
+    expect(JSON.stringify(cacheHitLog)).not.toContain('What is the capital of France?');
+
+    expect(result).toEqual({
+      output: 'Cached response from knowledge base',
+      metadata: {
+        citations: [
+          {
+            retrievedReferences: [
+              {
+                content: { text: 'Citation from cache' },
+                location: { s3Location: { uri: 'https://example.com/cached' } },
+              },
+            ],
+          },
+        ],
+      },
+      tokenUsage: createEmptyTokenUsage(),
+      cached: true,
+    });
+
+    mockIsCacheEnabled.mockReturnValue(false);
+  });
+
+  it('should use custom modelArn in cache key when provided', async () => {
+    mockIsCacheEnabled.mockReturnValue(true);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider('amazon.nova-lite-v1:0', {
+      config: {
+        knowledgeBaseId: 'kb-123',
+        region: 'us-east-1',
+        modelArn: 'custom:model:arn',
+      },
+    });
+
+    mockGet.mockResolvedValueOnce(null);
+
+    const mockResponse = {
+      output: {
+        text: 'Response with custom model ARN',
+      },
+      citations: [],
+    };
+    mockSend.mockResolvedValueOnce(mockResponse);
+
+    await provider.callApi('What is the capital of France?');
+
+    const cacheKey = mockGet.mock.calls[0][0];
+
+    expect(cacheKey).toBe(
+      buildKnowledgeBaseCacheKey({
+        knowledgeBaseId: 'kb-123',
+        modelArn: 'custom:model:arn',
+        modelName: 'amazon.nova-lite-v1:0',
+        prompt: 'What is the capital of France?',
+        region: 'us-east-1',
+        kbConfig: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          modelArn: 'custom:model:arn',
+        },
+      }),
+    );
+    expect(cacheKey).not.toContain('What is the capital of France?');
+
+    expect(mockSet).toHaveBeenCalledWith(cacheKey, expect.any(String));
+
+    mockIsCacheEnabled.mockReturnValue(false);
+  });
+
+  it('should include numberOfResults in cache key when provided', async () => {
+    mockIsCacheEnabled.mockReturnValue(true);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          numberOfResults: 10,
+        },
+      },
+    );
+
+    mockGet.mockResolvedValueOnce(null);
+
+    const mockResponse = {
+      output: {
+        text: 'Response with custom numberOfResults',
+      },
+      citations: [],
+    };
+    mockSend.mockResolvedValueOnce(mockResponse);
+
+    await provider.callApi('What is the capital of France?');
+
+    const cacheKey = mockGet.mock.calls[0][0];
+
+    expect(cacheKey).toBe(
+      buildKnowledgeBaseCacheKey({
+        knowledgeBaseId: 'kb-123',
+        modelArn: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+        modelName: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+        prompt: 'What is the capital of France?',
+        region: 'us-east-1',
+        kbConfig: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          numberOfResults: 10,
+        },
+      }),
+    );
+    expect(cacheKey).not.toContain('What is the capital of France?');
+
+    expect(mockSet).toHaveBeenCalledWith(cacheKey, expect.any(String));
+
+    mockIsCacheEnabled.mockReturnValue(false);
+  });
+
+  it('should hash prompt and config values in the cache key', async () => {
+    mockIsCacheEnabled.mockReturnValue(true);
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          apiKey: 'SECRET_API_KEY',
+          modelArn: 'custom:model:arn',
+        },
+      },
+    );
+
+    mockGet.mockResolvedValueOnce(null);
+    mockSend.mockResolvedValueOnce({
+      output: {
+        text: 'SECRET_RESPONSE_VALUE',
+      },
+      citations: [{ retrievedReferences: [{ content: { text: 'SECRET_CITATION_VALUE' } }] }],
+    });
+
+    await provider.callApi('SECRET_PROMPT_VALUE');
+
+    const cacheKey = mockGet.mock.calls[0][0];
+    const debugLogs = JSON.stringify(vi.mocked(logger.debug).mock.calls);
+
+    expect(cacheKey).not.toContain('SECRET_PROMPT_VALUE');
+    expect(cacheKey).not.toContain('SECRET_API_KEY');
+    expect(debugLogs).not.toContain('SECRET_PROMPT_VALUE');
+    expect(debugLogs).not.toContain('SECRET_RESPONSE_VALUE');
+    expect(debugLogs).not.toContain('SECRET_CITATION_VALUE');
+    expect(cacheKey).toBe(
+      buildKnowledgeBaseCacheKey({
+        knowledgeBaseId: 'kb-123',
+        modelArn: 'custom:model:arn',
+        modelName: 'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+        prompt: 'SECRET_PROMPT_VALUE',
+        region: 'us-east-1',
+        kbConfig: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          apiKey: 'SECRET_API_KEY',
+          modelArn: 'custom:model:arn',
+        },
+      }),
+    );
+
+    mockIsCacheEnabled.mockReturnValue(false);
+  });
+
+  it('should create knowledge base client with API key authentication from config', async () => {
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          apiKey: 'test-api-key',
+        },
+      },
+    );
+
+    await provider.getKnowledgeBaseClient();
+
+    expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({
+      region: 'us-east-1',
+      retryMode: 'adaptive',
+      maxAttempts: 10,
+      requestHandler: expect.any(Object),
+    });
+  });
+
+  it('should create knowledge base client with API key authentication from environment', async () => {
+    mockProcessEnv({ AWS_BEARER_TOKEN_BEDROCK: 'test-env-api-key' });
+
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+        },
+      },
+    );
+
+    await provider.getKnowledgeBaseClient();
+
+    expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({
+      region: 'us-east-1',
+      retryMode: 'adaptive',
+      maxAttempts: 10,
+      requestHandler: expect.any(Object),
+    });
+
+    mockProcessEnv({ AWS_BEARER_TOKEN_BEDROCK: undefined });
+  });
+
+  it('should prioritize explicit credentials over API key for knowledge base', async () => {
+    const provider = new AwsBedrockKnowledgeBaseProvider(
+      'us.anthropic.claude-3-7-sonnet-20241022-v2:0',
+      {
+        config: {
+          knowledgeBaseId: 'kb-123',
+          region: 'us-east-1',
+          apiKey: 'test-api-key',
+          accessKeyId: 'test-access-key',
+          secretAccessKey: 'test-secret-key',
+        },
+      },
+    );
+
+    await provider.getKnowledgeBaseClient();
+
+    // Should use explicit credentials (highest priority) instead of API key
+    expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({
+      region: 'us-east-1',
+      retryMode: 'adaptive',
+      maxAttempts: 10,
+      credentials: {
+        accessKeyId: 'test-access-key',
+        secretAccessKey: 'test-secret-key',
+        sessionToken: undefined,
+      },
+      requestHandler: expect.any(Object), // Still has handler for API key scenario
+    });
+  });
+});