@strands-agents/sdk 0.6.0 → 1.0.0-rc.0

package/dist/src/agent/__tests__/agent.test.js

@@ -5,12 +5,13 @@ import { MockMessageModel } from '../../__fixtures__/mock-message-model.js';
 import { collectGenerator } from '../../__fixtures__/model-test-helpers.js';
 import { createMockTool, createRandomTool } from '../../__fixtures__/tool-helpers.js';
 import { ConcurrentInvocationError } from '../../errors.js';
-import { MaxTokensError, TextBlock, CachePointBlock, AgentResult, Message, ToolUseBlock, ToolResultBlock, ReasoningBlock, GuardContentBlock, ImageBlock, VideoBlock, DocumentBlock, } from '../../index.js';
+import { MaxTokensError, TextBlock, CachePointBlock, Message, ToolUseBlock, ToolResultBlock, ReasoningBlock, GuardContentBlock, ImageBlock, VideoBlock, DocumentBlock, } from '../../index.js';
 import { AgentPrinter } from '../printer.js';
-import { BeforeInvocationEvent, BeforeToolsEvent } from '../../hooks/events.js';
+import { AfterInvocationEvent, AfterToolCallEvent, AfterToolsEvent, BeforeInvocationEvent, BeforeToolsEvent, } from '../../hooks/events.js';
 import { BedrockModel } from '../../models/bedrock.js';
-import { StructuredOutputException } from '../../structured-output/exceptions.js';
+import { StructuredOutputError } from '../../errors.js';
 import { expectLoopMetrics } from '../../__fixtures__/metrics-helpers.js';
+import { expectAgentResult } from '../../__fixtures__/agent-helpers.js';
 describe('Agent', () => {
     describe('stream', () => {
         describe('basic streaming', () => {
@@ -41,14 +42,14 @@ describe('Agent', () => {
                 const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
                 const agent = new Agent({ model });
                 const { result } = await collectGenerator(agent.stream('Test prompt'));
-                expect(result).toEqual(new AgentResult({
+                expect(result).toEqual(expectAgentResult({
                     stopReason: 'endTurn',
-                    lastMessage: expect.objectContaining({
-                        role: 'assistant',
-                        content: expect.arrayContaining([expect.objectContaining({ type: 'textBlock', text: 'Hello' })]),
-                    }),
-                    metrics: expectLoopMetrics({ cycleCount: 1 }),
+                    messageText: 'Hello',
+                    cycleCount: 1,
+                    traceCount: 1,
                 }));
+                // Verify trace structure
+                expect(result.traces?.[0]?.children).toEqual(expect.arrayContaining([expect.objectContaining({ name: 'stream_messages' })]));
             });
         });
         describe('with tool use', () => {
@@ -108,13 +109,34 @@ describe('Agent', () => {
         });
         describe('error handling', () => {
             it('throws MaxTokensError when model hits token limit', async () => {
-                const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial...' }, 'maxTokens');
+                const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial...' }, { stopReason: 'maxTokens' });
                 const agent = new Agent({ model });
                 await expect(async () => {
                     await collectGenerator(agent.stream('Test'));
                 }).rejects.toThrow(MaxTokensError);
             });
         });
+        describe('hook error cleanup', () => {
+            it('fires AfterInvocationEvent when consumer breaks from stream', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-1', input: {} })
+                    .addTurn({ type: 'textBlock', text: 'Done' });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('ok')],
+                }));
+                const agent = new Agent({ model, tools: [tool], printer: false });
+                const afterInvocationCallback = vi.fn();
+                agent.addHook(AfterInvocationEvent, afterInvocationCallback);
+                for await (const event of agent.stream('Test')) {
+                    if (event.type === 'beforeToolsEvent') {
+                        break;
+                    }
+                }
+                expect(afterInvocationCallback).toHaveBeenCalledOnce();
+            });
+        });
     });
     describe('invoke', () => {
         describe('basic invocation', () => {
@@ -131,21 +153,18 @@ describe('Agent', () => {
                 const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Response text' });
                 const agent = new Agent({ model });
                 const result = await agent.invoke('Test prompt');
-                expect(result).toEqual(new AgentResult({
+                expect(result).toEqual(expectAgentResult({
                     stopReason: 'endTurn',
-                    lastMessage: expect.objectContaining({
-                        type: 'message',
-                        role: 'assistant',
-                        content: expect.arrayContaining([expect.objectContaining({ type: 'textBlock', text: 'Response text' })]),
-                    }),
-                    metrics: expectLoopMetrics({ cycleCount: 1 }),
+                    messageText: 'Response text',
+                    cycleCount: 1,
+                    traceCount: 1,
                 }));
             });
             it('consumes stream events internally', async () => {
                 const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
                 const agent = new Agent({ model });
                 const result = await agent.invoke('Test');
-                expect(result).toEqual(new AgentResult({
+                expect(result).toEqual(expect.objectContaining({
                     stopReason: 'endTurn',
                     lastMessage: expect.objectContaining({
                         type: 'message',
@@ -159,8 +178,12 @@ describe('Agent', () => {
         describe('with tool use', () => {
             it('executes tools and returns final result', async () => {
                 const model = new MockMessageModel()
-                    .addTurn({ type: 'toolUseBlock', name: 'calc', toolUseId: 'tool-1', input: { a: 1, b: 2 } })
-                    .addTurn({ type: 'textBlock', text: 'The answer is 3' });
+                    .addTurn({ type: 'toolUseBlock', name: 'calc', toolUseId: 'tool-1', input: { a: 1, b: 2 } }, {
+                    usage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
+                })
+                    .addTurn({ type: 'textBlock', text: 'The answer is 3' }, {
+                    usage: { inputTokens: 200, outputTokens: 30, totalTokens: 230 },
+                });
                 const tool = createMockTool('calc', () => new ToolResultBlock({
                     toolUseId: 'tool-1',
                     status: 'success',
@@ -168,26 +191,245 @@ describe('Agent', () => {
                 }));
                 const agent = new Agent({ model, tools: [tool] });
                 const result = await agent.invoke('What is 1 + 2?');
-                expect(result).toEqual(new AgentResult({
+                expect(result).toEqual(expectAgentResult({
                     stopReason: 'endTurn',
-                    lastMessage: expect.objectContaining({
-                        type: 'message',
-                        role: 'assistant',
-                        content: expect.arrayContaining([
-                            expect.objectContaining({ type: 'textBlock', text: 'The answer is 3' }),
-                        ]),
+                    messageText: 'The answer is 3',
+                    cycleCount: 2,
+                    toolNames: ['calc'],
+                    traceCount: 2,
+                    usage: { inputTokens: 300, outputTokens: 80, totalTokens: 380 },
+                }));
+                // Verify detailed trace children structure
+                expect(result.traces?.[0]?.children).toEqual(expect.arrayContaining([
+                    expect.objectContaining({ name: 'stream_messages' }),
+                    expect.objectContaining({ name: 'Tool: calc' }),
+                ]));
+                expect(result.traces?.[1]?.children).toEqual(expect.arrayContaining([expect.objectContaining({ name: 'stream_messages' })]));
+            });
+            it('stores cycleId in trace metadata', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'calc', toolUseId: 'tool-1', input: {} })
+                    .addTurn({ type: 'textBlock', text: 'Done' });
+                const tool = createMockTool('calc', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('result')],
+                }));
+                const agent = new Agent({ model, tools: [tool] });
+                const result = await agent.invoke('Test');
+                expect(result.traces).toEqual([
+                    expect.objectContaining({
+                        name: 'Cycle 1',
+                        metadata: expect.objectContaining({ cycleId: 'cycle-1' }),
+                    }),
+                    expect.objectContaining({
+                        name: 'Cycle 2',
+                        metadata: expect.objectContaining({ cycleId: 'cycle-2' }),
                     }),
-                    metrics: expectLoopMetrics({ cycleCount: 2, toolNames: ['calc'] }),
+                ]);
+            });
+            it('stores tool metadata in trace children', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-abc123', input: {} })
+                    .addTurn({ type: 'textBlock', text: 'Done' });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-abc123',
+                    status: 'success',
+                    content: [new TextBlock('result')],
                 }));
+                const agent = new Agent({ model, tools: [tool] });
+                const result = await agent.invoke('Test');
+                expect(result.traces).toEqual(expect.arrayContaining([
+                    expect.objectContaining({
+                        name: 'Cycle 1',
+                        children: expect.arrayContaining([
+                            expect.objectContaining({
+                                name: 'Tool: testTool',
+                                metadata: expect.objectContaining({
+                                    toolUseId: 'tool-abc123',
+                                    toolName: 'testTool',
+                                }),
+                            }),
+                        ]),
+                    }),
+                ]));
             });
         });
         describe('error handling', () => {
             it('propagates maxTokens error', async () => {
-                const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial' }, 'maxTokens');
+                const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial' }, { stopReason: 'maxTokens' });
                 const agent = new Agent({ model });
                 await expect(agent.invoke('Test')).rejects.toThrow(MaxTokensError);
             });
         });
+        describe('metrics on errors', () => {
+            it('tracks cycle count when maxTokens error occurs', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-1', input: {} }, {
+                    usage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
+                })
+                    .addTurn({ type: 'textBlock', text: 'Partial' }, {
+                    stopReason: 'maxTokens',
+                    usage: { inputTokens: 80, outputTokens: 20, totalTokens: 100 },
+                });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('Done')],
+                }));
+                const agent = new Agent({ model, tools: [tool] });
+                const meter = agent._meter;
+                await expect(agent.invoke('Test')).rejects.toThrow(MaxTokensError);
+                expect(meter.metrics.cycleCount).toBe(2);
+                // Only the first turn's usage is accumulated; the second turn throws
+                // MaxTokensError inside streamAggregated before metadata reaches updateCycle
+                expect(meter.metrics.accumulatedUsage).toStrictEqual({
+                    inputTokens: 100,
+                    outputTokens: 50,
+                    totalTokens: 150,
+                });
+                expect(meter.metrics.accumulatedMetrics).toStrictEqual({
+                    latencyMs: expect.any(Number),
+                });
+                expect(meter.metrics.toolMetrics).toStrictEqual({
+                    testTool: {
+                        callCount: 1,
+                        successCount: 1,
+                        errorCount: 0,
+                        totalTime: expect.any(Number),
+                    },
+                });
+            });
+            it('collects local traces for completed cycles when error occurs mid-run', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-1', input: {} }, {
+                    usage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
+                })
+                    .addTurn({ type: 'textBlock', text: 'Partial' }, {
+                    stopReason: 'maxTokens',
+                    usage: { inputTokens: 80, outputTokens: 20, totalTokens: 100 },
+                });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('Done')],
+                }));
+                const agent = new Agent({ model, tools: [tool] });
+                const tracer = agent._tracer;
+                await expect(agent.invoke('Test')).rejects.toThrow(MaxTokensError);
+                // Cycle 1 completed (tool use), cycle 2 errored (maxTokens)
+                expect(tracer.localTraces).toEqual([
+                    expect.objectContaining({
+                        name: 'Cycle 1',
+                        children: [
+                            expect.objectContaining({ name: 'stream_messages' }),
+                            expect.objectContaining({ name: 'Tool: testTool' }),
+                        ],
+                    }),
+                    expect.objectContaining({
+                        name: 'Cycle 2',
+                        children: [expect.objectContaining({ name: 'stream_messages' })],
+                    }),
+                ]);
+            });
+            it('tracks metrics when a hook throws an error', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-1', input: {} }, {
+                    usage: { inputTokens: 60, outputTokens: 25, totalTokens: 85 },
+                })
+                    .addTurn({ type: 'textBlock', text: 'Done' });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('Result')],
+                }));
+                const agent = new Agent({ model, tools: [tool] });
+                agent.addHook(BeforeToolsEvent, () => {
+                    throw new Error('Hook failure');
+                });
+                const meter = agent._meter;
+                await expect(agent.invoke('Test')).rejects.toThrow('Hook failure');
+                // The hook throws after the model returns but before tools execute,
+                // so the first cycle's model usage is recorded but no tool metrics exist
+                expect(meter.metrics.cycleCount).toBe(1);
+                expect(meter.metrics.accumulatedUsage).toStrictEqual({
+                    inputTokens: 60,
+                    outputTokens: 25,
+                    totalTokens: 85,
+                });
+                expect(meter.metrics.accumulatedMetrics).toStrictEqual({
+                    latencyMs: expect.any(Number),
+                });
+                expect(meter.metrics.toolMetrics).toStrictEqual({});
+            });
+        });
+        describe('hook error cleanup', () => {
+            it('fires AfterInvocationEvent when a mid-stream hook throws', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-1', input: {} })
+                    .addTurn({ type: 'textBlock', text: 'Done' });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('ok')],
+                }));
+                const agent = new Agent({ model, tools: [tool], printer: false });
+                agent.addHook(AfterToolCallEvent, () => {
+                    throw new Error('hook error');
+                });
+                const afterInvocationCallback = vi.fn();
+                agent.addHook(AfterInvocationEvent, afterInvocationCallback);
+                await expect(agent.invoke('Test')).rejects.toThrow('hook error');
+                expect(afterInvocationCallback).toHaveBeenCalledOnce();
+            });
+            it('fires AfterToolsEvent when a mid-stream hook throws', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-1', input: {} })
+                    .addTurn({ type: 'textBlock', text: 'Done' });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('ok')],
+                }));
+                const agent = new Agent({ model, tools: [tool], printer: false });
+                agent.addHook(AfterToolCallEvent, () => {
+                    throw new Error('hook error');
+                });
+                const afterToolsCallback = vi.fn();
+                agent.addHook(AfterToolsEvent, afterToolsCallback);
+                await expect(agent.invoke('Test')).rejects.toThrow('hook error');
+                expect(afterToolsCallback).toHaveBeenCalledOnce();
+            });
+            it('does not fire AfterInvocationEvent when BeforeInvocationEvent hook throws', async () => {
+                const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+                const agent = new Agent({ model, printer: false });
+                agent.addHook(BeforeInvocationEvent, () => {
+                    throw new Error('before hook error');
+                });
+                const afterInvocationCallback = vi.fn();
+                agent.addHook(AfterInvocationEvent, afterInvocationCallback);
+                await expect(agent.invoke('Test')).rejects.toThrow('before hook error');
+                expect(afterInvocationCallback).not.toHaveBeenCalled();
+            });
+            it('does not fire AfterToolsEvent when BeforeToolsEvent hook throws', async () => {
+                const model = new MockMessageModel()
+                    .addTurn({ type: 'toolUseBlock', name: 'testTool', toolUseId: 'tool-1', input: {} })
+                    .addTurn({ type: 'textBlock', text: 'Done' });
+                const tool = createMockTool('testTool', () => new ToolResultBlock({
+                    toolUseId: 'tool-1',
+                    status: 'success',
+                    content: [new TextBlock('ok')],
+                }));
+                const agent = new Agent({ model, tools: [tool], printer: false });
+                agent.addHook(BeforeToolsEvent, () => {
+                    throw new Error('before tools hook error');
+                });
+                const afterToolsCallback = vi.fn();
+                agent.addHook(AfterToolsEvent, afterToolsCallback);
+                await expect(agent.invoke('Test')).rejects.toThrow('before tools hook error');
+                expect(afterToolsCallback).not.toHaveBeenCalled();
+            });
+        });
     });
     describe('API consistency', () => {
         it('invoke() and stream() produce same final result', async () => {
@@ -218,29 +460,36 @@ describe('Agent', () => {
             const agent2 = new Agent({ model: model2, tools: [tool2] });
             const invokeResult = await agent1.invoke('Use tool');
             const { result: streamResult } = await collectGenerator(agent2.stream('Use tool'));
-            expect(invokeResult.stopReason).toBe(streamResult.stopReason);
-            expect(invokeResult.lastMessage).toEqual(streamResult.lastMessage);
+            expect(invokeResult).toEqual(expect.objectContaining({
+                stopReason: streamResult.stopReason,
+                lastMessage: streamResult.lastMessage,
+                traces: streamResult.traces?.map((t) => expect.objectContaining({
+                    name: t.name,
+                    children: expect.arrayContaining(Array(t.children.length).fill(expect.objectContaining({ name: expect.any(String) }))),
+                })),
+            }));
         });
     });
     describe('messages', () => {
         it('returns array of messages', () => {
             const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
             const agent = new Agent({ model });
-            const messages = agent.messages;
-            expect(messages).toBeDefined();
-            expect(Array.isArray(messages)).toBe(true);
+            expect(agent.messages).toEqual([]);
         });
         it('reflects conversation history after invoke', async () => {
             const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Response' });
             const agent = new Agent({ model });
             await agent.invoke('Hello');
-            const messages = agent.messages;
-            expect(messages.length).toBeGreaterThan(0);
-            expect(messages.length).toBe(2);
-            expect(messages[0]?.role).toBe('user');
-            expect(messages[0]?.content).toEqual([{ type: 'textBlock', text: 'Hello' }]);
-            expect(messages[1]?.role).toBe('assistant');
-            expect(messages[1]?.content).toEqual([{ type: 'textBlock', text: 'Response' }]);
+            expect(agent.messages).toEqual([
+                expect.objectContaining({
+                    role: 'user',
+                    content: [{ type: 'textBlock', text: 'Hello' }],
+                }),
+                expect.objectContaining({
+                    role: 'assistant',
+                    content: [{ type: 'textBlock', text: 'Response' }],
+                }),
+            ]);
         });
     });
     describe('printer configuration', () => {
@@ -300,7 +549,7 @@ describe('Agent', () => {
         it('releases lock after errors and abandoned streams', async () => {
             // Test error case
             const model = new MockMessageModel()
-                .addTurn({ type: 'textBlock', text: 'Partial' }, 'maxTokens')
+                .addTurn({ type: 'textBlock', text: 'Partial' }, { stopReason: 'maxTokens' })
                 .addTurn({ type: 'textBlock', text: 'Success' });
             const agent = new Agent({ model });
             await expect(agent.invoke('First')).rejects.toThrow(MaxTokensError);
@@ -499,7 +748,7 @@ describe('Agent', () => {
                     }),
                     new VideoBlock({
                         format: 'mp4',
-                        source: { s3Location: { uri: 's3://bucket/video.mp4' } },
+                        source: { location: { type: 's3', uri: 's3://bucket/video.mp4' } },
                     }),
                     new DocumentBlock({
                         format: 'pdf',
@@ -551,7 +800,7 @@ describe('Agent', () => {
                     {
                         video: {
                             format: 'mp4',
-                            source: { s3Location: { uri: 's3://bucket/video.mp4' } },
+                            source: { location: { type: 's3', uri: 's3://bucket/video.mp4' } },
                         },
                     },
                     {
@@ -694,6 +943,7 @@ describe('Agent', () => {
             const agent = new Agent({ model, structuredOutputSchema: schema });
             const result = await agent.invoke('Test');
             expect(result.structuredOutput).toEqual({ name: 'John', age: 30 });
+            expect(model.callCount).toBe(1);
         });
         it('forces structured output tool when model does not use it', async () => {
             const schema = z.object({ value: z.number() });
@@ -705,16 +955,16 @@ describe('Agent', () => {
             const result = await agent.invoke('Test');
             expect(result.structuredOutput).toEqual({ value: 42 });
         });
-        it('throws StructuredOutputException when model refuses to use tool after forcing', async () => {
+        it('throws StructuredOutputError when model refuses to use tool after forcing', async () => {
             const schema = z.object({ value: z.number() });
             // Model returns text twice - once normally, once when forced
             const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Response' });
             const agent = new Agent({ model, structuredOutputSchema: schema });
-            await expect(agent.invoke('Test')).rejects.toThrow(StructuredOutputException);
+            await expect(agent.invoke('Test')).rejects.toThrow(StructuredOutputError);
         });
         it('throws MaxTokensError when maxTokens reached before structured output', async () => {
             const schema = z.object({ value: z.number() });
-            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial...' }, 'maxTokens');
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial...' }, { stopReason: 'maxTokens' });
             const agent = new Agent({ model, structuredOutputSchema: schema });
             await expect(agent.invoke('Test')).rejects.toThrow(MaxTokensError);
         });
@@ -756,7 +1006,7 @@ describe('Agent', () => {
         });
         it('cleans up structured output tool even when error occurs', async () => {
             const schema = z.object({ value: z.number() });
-            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial...' }, 'maxTokens');
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Partial...' }, { stopReason: 'maxTokens' });
             const agent = new Agent({ model, structuredOutputSchema: schema });
             await expect(agent.invoke('Test')).rejects.toThrow();
             const toolNames = agent.tools.map((t) => t.name);