@strands-agents/sdk 1.0.0-rc.5 → 1.1.0

package/dist/src/models/__tests__/bedrock.test.js

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { BedrockRuntimeClient, ConverseStreamCommand, ValidationException } from '@aws-sdk/client-bedrock-runtime';
+import { BedrockRuntimeClient, ConverseStreamCommand, CountTokensCommand, ValidationException, } from '@aws-sdk/client-bedrock-runtime';
 import { isNode } from '../../__fixtures__/environment.js';
 import { BedrockModel } from '../bedrock.js';
 import { ContextWindowOverflowError, ModelThrottledError } from '../../errors.js';
@@ -9,6 +9,7 @@ import { ImageBlock, VideoBlock, DocumentBlock } from '../../types/media.js';
 import { CitationsBlock } from '../../types/citations.js';
 import { collectIterator } from '../../__fixtures__/model-test-helpers.js';
 import { NOOP_TOOL_SPEC } from '../../tools/noop-tool.js';
+import { warnOnce } from '../../logging/warn-once.js';
 /**
  * Helper function to mock BedrockRuntimeClient implementation with customizable config.
  * @param options - Optional configuration for mock region, useFipsEndpoint, and send functions
@@ -92,6 +93,8 @@ vi.mock('@aws-sdk/client-bedrock-runtime', async (importOriginal) => {
         }
         throw new Error('Unhandled command type in mock');
     });
+    // Create a mock CountTokensCommand class
+    const CountTokensCommand = vi.fn();
     // Create a mock ValidationException class
     class MockValidationException extends Error {
         constructor(opts) {
@@ -113,9 +116,13 @@ vi.mock('@aws-sdk/client-bedrock-runtime', async (importOriginal) => {
         }),
         ConverseStreamCommand,
         ConverseCommand,
+        CountTokensCommand,
         ValidationException: MockValidationException,
     };
 });
+vi.mock('../../logging/warn-once.js', () => ({
+    warnOnce: vi.fn(),
+}));
 describe('BedrockModel', () => {
     const BEDROCK_NOOP_TOOL_CONFIG = {
         tools: [{ toolSpec: { ...NOOP_TOOL_SPEC, inputSchema: { json: NOOP_TOOL_SPEC.inputSchema } } }],
@@ -145,11 +152,20 @@ describe('BedrockModel', () => {
             const config = provider.getConfig();
             expect(config.modelId).toBeDefined();
         });
+        it('warns when modelId is not explicitly set', () => {
+            new BedrockModel();
+            expect(warnOnce).toHaveBeenCalledWith(expect.objectContaining({ warn: expect.any(Function) }), expect.stringContaining('using default modelId'));
+        });
+        it('does not warn when modelId is explicitly set', () => {
+            new BedrockModel({ modelId: 'us.anthropic.claude-3-5-sonnet-20241022-v2:0' });
+            expect(warnOnce).not.toHaveBeenCalledWith(expect.objectContaining({ warn: expect.any(Function) }), expect.stringContaining('using default modelId'));
+        });
         it('uses provided model ID ', () => {
             const customModelId = 'us.anthropic.claude-3-5-sonnet-20241022-v2:0';
             const provider = new BedrockModel({ modelId: customModelId });
             expect(provider.getConfig()).toStrictEqual({
                 modelId: customModelId,
+                contextWindowLimit: 200_000,
             });
         });
         it('uses provided region', () => {
@@ -158,6 +174,7 @@ describe('BedrockModel', () => {
             expect(BedrockRuntimeClient).toHaveBeenCalledWith({
                 region: customRegion,
                 customUserAgent: 'strands-agents-ts-sdk',
+                requestHandler: { requestTimeout: 120_000 },
             });
         });
         it('extends custom user agent if provided', () => {
@@ -166,6 +183,7 @@ describe('BedrockModel', () => {
             expect(BedrockRuntimeClient).toHaveBeenCalledWith({
                 region: 'us-west-2',
                 customUserAgent: 'my-app/1.0 strands-agents-ts-sdk',
+                requestHandler: { requestTimeout: 120_000 },
             });
         });
         it('passes custom endpoint to client', () => {
@@ -176,6 +194,7 @@ describe('BedrockModel', () => {
                 region,
                 endpoint,
                 customUserAgent: 'strands-agents-ts-sdk',
+                requestHandler: { requestTimeout: 120_000 },
             });
         });
         it('passes custom credentials to client', () => {
@@ -189,8 +208,26 @@ describe('BedrockModel', () => {
                 region,
                 credentials,
                 customUserAgent: 'strands-agents-ts-sdk',
+                requestHandler: { requestTimeout: 120_000 },
             });
         });
+        it('applies a default 120s request timeout', () => {
+            new BedrockModel({ region: 'us-west-2' });
+            expect(BedrockRuntimeClient).toHaveBeenCalledWith(expect.objectContaining({ requestHandler: { requestTimeout: 120_000 } }));
+        });
+        it('lets the caller override requestTimeout', () => {
+            new BedrockModel({ region: 'us-west-2', clientConfig: { requestHandler: { requestTimeout: 5_000 } } });
+            expect(BedrockRuntimeClient).toHaveBeenCalledWith(expect.objectContaining({ requestHandler: { requestTimeout: 5_000 } }));
+        });
+        it('merges the default timeout with other requestHandler options', () => {
+            new BedrockModel({ region: 'us-west-2', clientConfig: { requestHandler: { connectionTimeout: 1_000 } } });
+            expect(BedrockRuntimeClient).toHaveBeenCalledWith(expect.objectContaining({ requestHandler: { requestTimeout: 120_000, connectionTimeout: 1_000 } }));
+        });
+        it('passes a user-provided handler instance through untouched', () => {
+            const handler = { handle: vi.fn(), updateHttpClientConfig: vi.fn(), httpHandlerConfigs: vi.fn() };
+            new BedrockModel({ region: 'us-west-2', clientConfig: { requestHandler: handler } });
+            expect(BedrockRuntimeClient).toHaveBeenCalledWith(expect.objectContaining({ requestHandler: handler }));
+        });
         it('adds api key middleware when apiKey is provided', () => {
             const provider = new BedrockModel({ region: 'us-east-1', apiKey: 'br-test-key' });
             const mockAdd = provider['_client'].middlewareStack.add;
@@ -222,6 +259,7 @@ describe('BedrockModel', () => {
             expect(config).toStrictEqual({
                 modelId: 'global.anthropic.claude-sonnet-4-6',
                 temperature: 0.5,
+                contextWindowLimit: 1_000_000,
             });
         });
         it('includes contextWindowLimit in config when provided', () => {
@@ -234,6 +272,43 @@ describe('BedrockModel', () => {
                 contextWindowLimit: 200_000,
             });
         });
+        it('auto-populates contextWindowLimit from model ID lookup', () => {
+            const provider = new BedrockModel({ modelId: 'anthropic.claude-sonnet-4-20250514-v1:0' });
+            expect(provider.getConfig()).toStrictEqual({
+                modelId: 'anthropic.claude-sonnet-4-20250514-v1:0',
+                contextWindowLimit: 1_000_000,
+            });
+        });
+        it('auto-populates contextWindowLimit for cross-region model IDs', () => {
+            const provider = new BedrockModel({ modelId: 'us.anthropic.claude-sonnet-4-6' });
+            expect(provider.getConfig()).toStrictEqual({
+                modelId: 'us.anthropic.claude-sonnet-4-6',
+                contextWindowLimit: 1_000_000,
+            });
+        });
+        it('auto-populates contextWindowLimit for default model ID', () => {
+            const provider = new BedrockModel();
+            expect(provider.getConfig()).toStrictEqual({
+                modelId: 'global.anthropic.claude-sonnet-4-6',
+                contextWindowLimit: 1_000_000,
+            });
+        });
+        it('does not override explicit contextWindowLimit', () => {
+            const provider = new BedrockModel({
+                modelId: 'anthropic.claude-sonnet-4-20250514-v1:0',
+                contextWindowLimit: 100_000,
+            });
+            expect(provider.getConfig()).toStrictEqual({
+                modelId: 'anthropic.claude-sonnet-4-20250514-v1:0',
+                contextWindowLimit: 100_000,
+            });
+        });
+        it('leaves contextWindowLimit undefined for unknown model IDs', () => {
+            const provider = new BedrockModel({ modelId: 'unknown.model-v1:0' });
+            expect(provider.getConfig()).toStrictEqual({
+                modelId: 'unknown.model-v1:0',
+            });
+        });
     });
     describe('updateConfig', () => {
         it('merges new config with existing config', () => {
@@ -243,6 +318,7 @@ describe('BedrockModel', () => {
                 modelId: 'global.anthropic.claude-sonnet-4-6',
                 temperature: 0.8,
                 maxTokens: 2048,
+                contextWindowLimit: 1_000_000,
             });
         });
         it('preserves fields not included in the update', () => {
@@ -259,6 +335,24 @@ describe('BedrockModel', () => {
                 maxTokens: 1024,
             });
         });
+        it('re-resolves contextWindowLimit when modelId changes and it was auto-resolved', () => {
+            const provider = new BedrockModel({ region: 'us-west-2' });
+            expect(provider.getConfig().contextWindowLimit).toBe(1_000_000);
+            provider.updateConfig({ modelId: 'anthropic.claude-haiku-4-5-20251001-v1:0' });
+            expect(provider.getConfig().contextWindowLimit).toBe(200_000);
+        });
+        it('clears contextWindowLimit when modelId changes to unknown model', () => {
+            const provider = new BedrockModel({ region: 'us-west-2' });
+            expect(provider.getConfig().contextWindowLimit).toBe(1_000_000);
+            provider.updateConfig({ modelId: 'my-custom-finetuned-model' });
+            expect(provider.getConfig().contextWindowLimit).toBeUndefined();
+        });
+        it('preserves explicit contextWindowLimit when modelId changes', () => {
+            const provider = new BedrockModel({ region: 'us-west-2', contextWindowLimit: 50_000 });
+            expect(provider.getConfig().contextWindowLimit).toBe(50_000);
+            provider.updateConfig({ modelId: 'anthropic.claude-haiku-4-5-20251001-v1:0' });
+            expect(provider.getConfig().contextWindowLimit).toBe(50_000);
+        });
     });
     describe('getConfig', () => {
         it('returns the current configuration', () => {
@@ -1174,8 +1268,8 @@ describe('BedrockModel', () => {
             });
         });
         it('does not warn when array system prompt is provided without cacheConfig', async () => {
-            const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => { });
             const provider = new BedrockModel();
+            const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => { });
             const messages = [new Message({ role: 'user', content: [new TextBlock('Hello')] })];
             const options = {
                 systemPrompt: [
@@ -3530,5 +3624,137 @@ describe('BedrockModel', () => {
             }));
         });
     });
+    describe('countTokens', () => {
+        const messages = [new Message({ role: 'user', content: [new TextBlock('hello')] })];
+        const toolSpecs = [
+            { name: 'test_tool', description: 'A test tool', inputSchema: { type: 'object', properties: {} } },
+        ];
+        beforeEach(() => {
+            vi.clearAllMocks();
+            BedrockModel.clearCountTokensCache();
+        });
+        it('should return native token count on success', async () => {
+            const mockSend = vi.fn(async () => ({ inputTokens: 42 }));
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel();
+            const result = await model.countTokens(messages);
+            expect(result).toBe(42);
+            expect(mockSend).toHaveBeenCalledOnce();
+        });
+        it('should include system prompt in request', async () => {
+            const mockSend = vi.fn(async () => ({ inputTokens: 55 }));
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel();
+            const result = await model.countTokens(messages, { systemPrompt: 'Be helpful.' });
+            expect(result).toBe(55);
+            const commandInput = vi.mocked(CountTokensCommand).mock.calls[0][0];
+            expect(commandInput).toStrictEqual({
+                modelId: expect.any(String),
+                input: {
+                    converse: {
+                        messages: [{ role: 'user', content: [{ text: 'hello' }] }],
+                        system: [{ text: 'Be helpful.' }],
+                    },
+                },
+            });
+        });
+        it('should include tool specs in request', async () => {
+            const mockSend = vi.fn(async () => ({ inputTokens: 100 }));
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel();
+            const result = await model.countTokens(messages, { toolSpecs });
+            expect(result).toBe(100);
+            const commandInput = vi.mocked(CountTokensCommand).mock.calls[0][0];
+            expect(commandInput).toStrictEqual({
+                modelId: expect.any(String),
+                input: {
+                    converse: {
+                        messages: [{ role: 'user', content: [{ text: 'hello' }] }],
+                        toolConfig: {
+                            tools: [
+                                {
+                                    toolSpec: {
+                                        name: 'test_tool',
+                                        description: 'A test tool',
+                                        inputSchema: { json: { type: 'object', properties: {} } },
+                                    },
+                                },
+                            ],
+                        },
+                    },
+                },
+            });
+        });
+        it('should strip inferenceConfig from request', async () => {
+            const mockSend = vi.fn(async () => ({ inputTokens: 10 }));
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel({ maxTokens: 100 });
+            await model.countTokens(messages);
+            const commandInput = vi.mocked(CountTokensCommand).mock.calls[0][0];
+            expect(commandInput).toStrictEqual({
+                modelId: expect.any(String),
+                input: {
+                    converse: {
+                        messages: [{ role: 'user', content: [{ text: 'hello' }] }],
+                    },
+                },
+            });
+        });
+        it('should fall back to estimation on API error', async () => {
+            const mockSend = vi.fn(async () => {
+                throw new Error('API error');
+            });
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel();
+            const result = await model.countTokens(messages);
+            expect(typeof result).toBe('number');
+            expect(result).toBeGreaterThanOrEqual(0);
+        });
+        it('should fall back to estimation on generic exception', async () => {
+            const mockSend = vi.fn(async () => {
+                throw new Error('Connection failed');
+            });
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel();
+            const result = await model.countTokens(messages);
+            expect(typeof result).toBe('number');
+            expect(result).toBeGreaterThanOrEqual(0);
+        });
+        it('should cache model ID and skip API call when model does not support counting tokens', async () => {
+            const unsupportedError = new Error("The provided model doesn't support counting tokens");
+            unsupportedError.name = 'ValidationException';
+            const mockSend = vi.fn(async () => {
+                throw unsupportedError;
+            });
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel();
+            // First call: hits API, gets error, caches
+            await model.countTokens(messages);
+            expect(mockSend).toHaveBeenCalledOnce();
+            // Second call: skips API entirely
+            await model.countTokens(messages);
+            expect(mockSend).toHaveBeenCalledOnce();
+        });
+        it('should not cache model ID for other errors', async () => {
+            const mockSend = vi.fn(async () => {
+                throw new Error('Transient network error');
+            });
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel();
+            await model.countTokens(messages);
+            expect(mockSend).toHaveBeenCalledTimes(1);
+            // Second call should still attempt the API
+            await model.countTokens(messages);
+            expect(mockSend).toHaveBeenCalledTimes(2);
+        });
+        it('should skip native API and use heuristic when useNativeTokenCount is false', async () => {
+            const mockSend = vi.fn();
+            mockBedrockClientImplementation({ send: mockSend });
+            const model = new BedrockModel({ useNativeTokenCount: false });
+            const result = await model.countTokens(messages);
+            expect(mockSend).not.toHaveBeenCalled();
+            expect(result).toBe(2); // heuristic: Math.ceil('hello'.length / 4)
+        });
+    });
 });
 //# sourceMappingURL=bedrock.test.js.map