@strands-agents/sdk 1.0.0-rc.4 → 1.0.0

package/dist/src/agent/__tests__/agent.invocation-state.test.js

@@ -0,0 +1,219 @@
+import { describe, expect, it } from 'vitest';
+import { Agent } from '../agent.js';
+import { AfterInvocationEvent, AfterModelCallEvent, AfterToolCallEvent, BeforeInvocationEvent, BeforeModelCallEvent, MessageAddedEvent, } from '../../hooks/events.js';
+import { MockMessageModel } from '../../__fixtures__/mock-message-model.js';
+import { createMockTool } from '../../__fixtures__/tool-helpers.js';
+import { ToolResultBlock, TextBlock } from '../../types/messages.js';
+describe('invocationState', () => {
+    describe('round-trip', () => {
+        it('returns an empty object on AgentResult when no invocationState is passed', async () => {
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+            const agent = new Agent({ model });
+            const result = await agent.invoke('Hi');
+            expect(result.invocationState).toEqual({});
+        });
+        it('returns the passed invocationState on AgentResult', async () => {
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+            const agent = new Agent({ model });
+            const result = await agent.invoke('Hi', { invocationState: { userId: 'u-1', traceId: 't-1' } });
+            expect(result.invocationState).toEqual({ userId: 'u-1', traceId: 't-1' });
+        });
+        it('preserves reference identity: caller keeps the same object they passed in', async () => {
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+            const agent = new Agent({ model });
+            const state = { userId: 'u-1' };
+            const result = await agent.invoke('Hi', { invocationState: state });
+            expect(result.invocationState).toBe(state);
+        });
+    });
+    describe('hook mutation', () => {
+        it('propagates mutations from BeforeModelCallEvent to AfterModelCallEvent and AgentResult', async () => {
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+            const agent = new Agent({ model });
+            let seenInAfter;
+            agent.addHook(BeforeModelCallEvent, (event) => {
+                event.invocationState.counter = event.invocationState.counter ?? 0;
+                event.invocationState.counter = event.invocationState.counter + 1;
+            });
+            agent.addHook(AfterModelCallEvent, (event) => {
+                seenInAfter = event.invocationState;
+            });
+            const result = await agent.invoke('Hi');
+            expect(seenInAfter).toEqual({ counter: 1 });
+            expect(result.invocationState).toEqual({ counter: 1 });
+        });
+        it('shares the same invocationState object across all lifecycle events in one invocation', async () => {
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+            const agent = new Agent({ model });
+            const seen = [];
+            const collect = (event) => {
+                seen.push(event.invocationState);
+            };
+            agent.addHook(BeforeInvocationEvent, collect);
+            agent.addHook(BeforeModelCallEvent, collect);
+            agent.addHook(AfterModelCallEvent, collect);
+            agent.addHook(MessageAddedEvent, collect);
+            agent.addHook(AfterInvocationEvent, collect);
+            const result = await agent.invoke('Hi');
+            // Every hook, plus the result, sees the same reference.
+            expect(seen.length).toBeGreaterThan(0);
+            for (const observed of seen) {
+                expect(observed).toBe(result.invocationState);
+            }
+        });
+    });
+    describe('multi-cycle persistence', () => {
+        it('persists mutations across recursive agent loop cycles (tool-use scenario)', async () => {
+            const tool = createMockTool('ping', () => new ToolResultBlock({
+                toolUseId: 'tool-1',
+                status: 'success',
+                content: [new TextBlock('pong')],
+            }));
+            const model = new MockMessageModel()
+                .addTurn([{ type: 'toolUseBlock', name: 'ping', toolUseId: 'tool-1', input: {} }])
+                .addTurn({ type: 'textBlock', text: 'Done' });
+            const agent = new Agent({ model, tools: [tool] });
+            // Write in AfterToolCallEvent during cycle 1; read in BeforeModelCallEvent during cycle 2.
+            let cycle2State;
+            let modelCalls = 0;
+            agent.addHook(AfterToolCallEvent, (event) => {
+                event.invocationState.toolCompleted = true;
+            });
+            agent.addHook(BeforeModelCallEvent, (event) => {
+                modelCalls++;
+                if (modelCalls === 2) {
+                    cycle2State = event.invocationState;
+                }
+            });
+            const result = await agent.invoke('Run ping');
+            expect(modelCalls).toBe(2);
+            expect(cycle2State).toEqual({ toolCompleted: true });
+            expect(result.invocationState).toEqual({ toolCompleted: true });
+        });
+    });
+    describe('tool access', () => {
+        it('passes invocationState to tools via ToolContext and surfaces mutations on the result', async () => {
+            const tool = createMockTool('writer', () => {
+                throw new Error('unused');
+            });
+            // Override stream to read/write invocationState.
+            // eslint-disable-next-line require-yield
+            tool.stream = async function* (context) {
+                const prev = context.invocationState.callCount ?? 0;
+                context.invocationState.callCount = prev + 1;
+                context.invocationState.lastToolSeenUserId = context.invocationState.userId;
+                return new ToolResultBlock({
+                    toolUseId: context.toolUse.toolUseId,
+                    status: 'success',
+                    content: [new TextBlock('ok')],
+                });
+            };
+            const model = new MockMessageModel()
+                .addTurn([{ type: 'toolUseBlock', name: 'writer', toolUseId: 'tu-1', input: {} }])
+                .addTurn({ type: 'textBlock', text: 'Done' });
+            const agent = new Agent({ model, tools: [tool] });
+            const result = await agent.invoke('Run writer', { invocationState: { userId: 'u-42' } });
+            expect(result.invocationState).toEqual({
+                userId: 'u-42',
+                callCount: 1,
+                lastToolSeenUserId: 'u-42',
+            });
+        });
+    });
+    describe('isolation from appState', () => {
+        it('does not touch agent.appState when invocationState is mutated', async () => {
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+            const agent = new Agent({ model, appState: { persistent: 'yes' } });
+            agent.addHook(BeforeModelCallEvent, (event) => {
+                event.invocationState.ephemeral = 'only-this-run';
+            });
+            const result = await agent.invoke('Hi', { invocationState: { requestId: 'r-1' } });
+            expect(result.invocationState).toEqual({ requestId: 'r-1', ephemeral: 'only-this-run' });
+            expect(agent.appState.get('persistent')).toBe('yes');
+            expect(agent.appState.get('ephemeral')).toBeUndefined();
+            expect(agent.appState.get('requestId')).toBeUndefined();
+        });
+    });
+    describe('across invocations', () => {
+        it('does not leak state between invocations on the same agent (default bag)', async () => {
+            const model = new MockMessageModel()
+                .addTurn({ type: 'textBlock', text: 'A' })
+                .addTurn({ type: 'textBlock', text: 'B' });
+            const agent = new Agent({ model });
+            agent.addHook(BeforeModelCallEvent, (event) => {
+                event.invocationState.seen = true;
+            });
+            const first = await agent.invoke('1');
+            const second = await agent.invoke('2');
+            expect(first.invocationState).toEqual({ seen: true });
+            expect(second.invocationState).toEqual({ seen: true });
+            expect(first.invocationState).not.toBe(second.invocationState);
+        });
+    });
+    describe('retry paths', () => {
+        it('preserves same invocationState reference across AfterModelCallEvent retry', async () => {
+            const model = new MockMessageModel()
+                .addTurn(new Error('transient failure'))
+                .addTurn({ type: 'textBlock', text: 'Success after retry' });
+            const agent = new Agent({ model, printer: false });
+            let retried = false;
+            const seen = [];
+            agent.addHook(BeforeModelCallEvent, (event) => {
+                seen.push(event.invocationState);
+                event.invocationState.modelCalls = event.invocationState.modelCalls ?? 0;
+                event.invocationState.modelCalls = event.invocationState.modelCalls + 1;
+            });
+            agent.addHook(AfterModelCallEvent, (event) => {
+                seen.push(event.invocationState);
+                if (!retried && event.error) {
+                    retried = true;
+                    event.retry = true;
+                }
+            });
+            const result = await agent.invoke('Test', { invocationState: { userId: 'u-1' } });
+            // Retry path was exercised: two Before + two After observations.
+            expect(seen.length).toBe(4);
+            // Every observation is the same object the caller passed in.
+            for (const observed of seen) {
+                expect(observed).toBe(result.invocationState);
+            }
+            // Mutations from the first attempt survive into the retry.
+            expect(result.invocationState).toEqual({ userId: 'u-1', modelCalls: 2 });
+        });
+        it('preserves same invocationState reference across AfterToolCallEvent retry', async () => {
+            let toolCalls = 0;
+            const tool = createMockTool('flaky', () => {
+                toolCalls++;
+                return new ToolResultBlock({
+                    toolUseId: 'tu-1',
+                    status: toolCalls === 1 ? 'error' : 'success',
+                    content: [new TextBlock(toolCalls === 1 ? 'fail' : 'ok')],
+                });
+            });
+            const model = new MockMessageModel()
+                .addTurn([{ type: 'toolUseBlock', name: 'flaky', toolUseId: 'tu-1', input: {} }])
+                .addTurn({ type: 'textBlock', text: 'Done' });
+            const agent = new Agent({ model, tools: [tool], printer: false });
+            let retried = false;
+            const seen = [];
+            agent.addHook(AfterToolCallEvent, (event) => {
+                seen.push(event.invocationState);
+                event.invocationState.toolAttempts = event.invocationState.toolAttempts ?? 0;
+                event.invocationState.toolAttempts = event.invocationState.toolAttempts + 1;
+                if (!retried && event.result.status === 'error') {
+                    retried = true;
+                    event.retry = true;
+                }
+            });
+            const result = await agent.invoke('Run flaky', { invocationState: { requestId: 'r-1' } });
+            // Retry fired twice: failed attempt + successful attempt.
+            expect(toolCalls).toBe(2);
+            expect(seen.length).toBe(2);
+            for (const observed of seen) {
+                expect(observed).toBe(result.invocationState);
+            }
+            expect(result.invocationState).toEqual({ requestId: 'r-1', toolAttempts: 2 });
+        });
+    });
+});
+//# sourceMappingURL=agent.invocation-state.test.js.map

package/dist/src/agent/__tests__/agent.invocation-state.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"agent.invocation-state.test.js","sourceRoot":"","sources":["../../../../src/agent/__tests__/agent.invocation-state.test.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,QAAQ,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,QAAQ,CAAA;AAC7C,OAAO,EAAE,KAAK,EAAE,MAAM,aAAa,CAAA;AACnC,OAAO,EACL,oBAAoB,EACpB,mBAAmB,EACnB,kBAAkB,EAClB,qBAAqB,EACrB,oBAAoB,EACpB,iBAAiB,GAClB,MAAM,uBAAuB,CAAA;AAC9B,OAAO,EAAE,gBAAgB,EAAE,MAAM,0CAA0C,CAAA;AAC3E,OAAO,EAAE,cAAc,EAAE,MAAM,oCAAoC,CAAA;AACnE,OAAO,EAAE,eAAe,EAAE,SAAS,EAAE,MAAM,yBAAyB,CAAA;AAGpE,QAAQ,CAAC,iBAAiB,EAAE,GAAG,EAAE;IAC/B,QAAQ,CAAC,YAAY,EAAE,GAAG,EAAE;QAC1B,EAAE,CAAC,0EAA0E,EAAE,KAAK,IAAI,EAAE;YACxF,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YAClF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,CAAC,CAAA;YAElC,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,CAAC,CAAA;YAEvC,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,CAAC,CAAA;QAC5C,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,mDAAmD,EAAE,KAAK,IAAI,EAAE;YACjE,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YAClF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,CAAC,CAAA;YAElC,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,EAAE,EAAE,eAAe,EAAE,EAAE,MAAM,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,EAAE,CAAC,CAAA;YAE/F,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;QAC3E,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,2EAA2E,EAAE,KAAK,IAAI,EAAE;YACzF,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YAClF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,CAAC,CAAA;YAElC,MAAM,KAAK,GAAoB,EAAE,MAAM,EAAE,KAAK,EAAE,CAAA;YAChD,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,EAAE,EAAE,eAAe,EAAE,KAAK,EAAE,CAAC,CAAA;YAEnE,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAA;QAC5C,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,eAAe,EAAE,GAAG,EAAE;QAC7B,EAAE,CAAC,uFAAuF,EAAE,KAAK,IAAI,EAAE;YACrG,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YAClF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,CAAC,CAAA;YAElC,IAAI,WAAwC,CAAA;YAC5C,KAAK,CAAC,OAAO,CAAC,oBAAoB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC5C,KAAK,CAAC,eAAe,CAAC,OAAO,GAAI,KAAK,CAAC,eAAe,CAAC,OAA8B,IAAI,CAAC,CAAA;gBAC1F,KAAK,CAAC,eAAe,CAAC,OAAO,GAAI,KAAK,CAAC,eAAe,CAAC,OAAkB,GAAG,CAAC,CAAA;YAC/E,CAAC,CAAC,CAAA;YACF,KAAK,CAAC,OAAO,CAAC,mBAAmB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC3C,WAAW,GAAG,KAAK,CAAC,eAAe,CAAA;YACrC,CAAC,CAAC,CAAA;YAEF,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,CAAC,CAAA;YAEvC,MAAM,CAAC,WAAW,CAAC,CAAC,OAAO,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,CAAC,CAAA;YAC3C,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,CAAC,CAAA;QACxD,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,sFAAsF,EAAE,KAAK,IAAI,EAAE;YACpG,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YAClF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,CAAC,CAAA;YAElC,MAAM,IAAI,GAAsB,EAAE,CAAA;YAClC,MAAM,OAAO,GAAG,CAAC,KAA2C,EAAQ,EAAE;gBACpE,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,eAAe,CAAC,CAAA;YAClC,CAAC,CAAA;YAED,KAAK,CAAC,OAAO,CAAC,qBAAqB,EAAE,OAAO,CAAC,CAAA;YAC7C,KAAK,CAAC,OAAO,CAAC,oBAAoB,EAAE,OAAO,CAAC,CAAA;YAC5C,KAAK,CAAC,OAAO,CAAC,mBAAmB,EAAE,OAAO,CAAC,CAAA;YAC3C,KAAK,CAAC,OAAO,CAAC,iBAAiB,EAAE,OAAO,CAAC,CAAA;YACzC,KAAK,CAAC,OAAO,CAAC,oBAAoB,EAAE,OAAO,CAAC,CAAA;YAE5C,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,CAAC,CAAA;YAEvC,wDAAwD;YACxD,MAAM,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAA;YACtC,KAAK,MAAM,QAAQ,IAAI,IAAI,EAAE,CAAC;gBAC5B,MAAM,CAAC,QAAQ,CAAC,CAAC,IAAI,CAAC,MAAM,CAAC,eAAe,CAAC,CAAA;YAC/C,CAAC;QACH,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,yBAAyB,EAAE,GAAG,EAAE;QACvC,EAAE,CAAC,2EAA2E,EAAE,KAAK,IAAI,EAAE;YACzF,MAAM,IAAI,GAAG,cAAc,CACzB,MAAM,EACN,GAAG,EAAE,CACH,IAAI,eAAe,CAAC;gBAClB,SAAS,EAAE,QAAQ;gBACnB,MAAM,EAAE,SAAS;gBACjB,OAAO,EAAE,CAAC,IAAI,SAAS,CAAC,MAAM,CAAC,CAAC;aACjC,CAAC,CACL,CAAA;YAED,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE;iBACjC,OAAO,CAAC,CAAC,EAAE,IAAI,EAAE,cAAc,EAAE,IAAI,EAAE,MAAM,EAAE,SAAS,EAAE,QAAQ,EAAE,KAAK,EAAE,EAAE,EAAE,CAAC,CAAC;iBACjF,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,CAAA;YAC/C,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,CAAA;YAEjD,2FAA2F;YAC3F,IAAI,WAAwC,CAAA;YAC5C,IAAI,UAAU,GAAG,CAAC,CAAA;YAClB,KAAK,CAAC,OAAO,CAAC,kBAAkB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC1C,KAAK,CAAC,eAAe,CAAC,aAAa,GAAG,IAAI,CAAA;YAC5C,CAAC,CAAC,CAAA;YACF,KAAK,CAAC,OAAO,CAAC,oBAAoB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC5C,UAAU,EAAE,CAAA;gBACZ,IAAI,UAAU,KAAK,CAAC,EAAE,CAAC;oBACrB,WAAW,GAAG,KAAK,CAAC,eAAe,CAAA;gBACrC,CAAC;YACH,CAAC,CAAC,CAAA;YAEF,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,UAAU,CAAC,CAAA;YAE7C,MAAM,CAAC,UAAU,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;YAC1B,MAAM,CAAC,WAAW,CAAC,CAAC,OAAO,CAAC,EAAE,aAAa,EAAE,IAAI,EAAE,CAAC,CAAA;YACpD,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,aAAa,EAAE,IAAI,EAAE,CAAC,CAAA;QACjE,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,aAAa,EAAE,GAAG,EAAE;QAC3B,EAAE,CAAC,sFAAsF,EAAE,KAAK,IAAI,EAAE;YACpG,MAAM,IAAI,GAAG,cAAc,CAAC,QAAQ,EAAE,GAAG,EAAE;gBACzC,MAAM,IAAI,KAAK,CAAC,QAAQ,CAAC,CAAA;YAC3B,CAAC,CAAC,CAAA;YACF,iDAAiD;YACjD,yCAAyC;YACzC,IAAI,CAAC,MAAM,GAAG,KAAK,SAAS,CAAC,EAAE,OAAO;gBACpC,MAAM,IAAI,GAAI,OAAO,CAAC,eAAe,CAAC,SAAgC,IAAI,CAAC,CAAA;gBAC3E,OAAO,CAAC,eAAe,CAAC,SAAS,GAAG,IAAI,GAAG,CAAC,CAAA;gBAC5C,OAAO,CAAC,eAAe,CAAC,kBAAkB,GAAG,OAAO,CAAC,eAAe,CAAC,MAAM,CAAA;gBAC3E,OAAO,IAAI,eAAe,CAAC;oBACzB,SAAS,EAAE,OAAO,CAAC,OAAO,CAAC,SAAS;oBACpC,MAAM,EAAE,SAAS;oBACjB,OAAO,EAAE,CAAC,IAAI,SAAS,CAAC,IAAI,CAAC,CAAC;iBAC/B,CAAC,CAAA;YACJ,CAAC,CAAA;YAED,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE;iBACjC,OAAO,CAAC,CAAC,EAAE,IAAI,EAAE,cAAc,EAAE,IAAI,EAAE,QAAQ,EAAE,SAAS,EAAE,MAAM,EAAE,KAAK,EAAE,EAAE,EAAE,CAAC,CAAC;iBACjF,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,CAAA;YAC/C,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,CAAA;YAEjD,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,YAAY,EAAE,EAAE,eAAe,EAAE,EAAE,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,CAAA;YAExF,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC;gBACrC,MAAM,EAAE,MAAM;gBACd,SAAS,EAAE,CAAC;gBACZ,kBAAkB,EAAE,MAAM;aAC3B,CAAC,CAAA;QACJ,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,yBAAyB,EAAE,GAAG,EAAE;QACvC,EAAE,CAAC,+DAA+D,EAAE,KAAK,IAAI,EAAE;YAC7E,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YAClF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,QAAQ,EAAE,EAAE,UAAU,EAAE,KAAK,EAAE,EAAE,CAAC,CAAA;YAEnE,KAAK,CAAC,OAAO,CAAC,oBAAoB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC5C,KAAK,CAAC,eAAe,CAAC,SAAS,GAAG,eAAe,CAAA;YACnD,CAAC,CAAC,CAAA;YAEF,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,EAAE,EAAE,eAAe,EAAE,EAAE,SAAS,EAAE,KAAK,EAAE,EAAE,CAAC,CAAA;YAElF,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,SAAS,EAAE,KAAK,EAAE,SAAS,EAAE,eAAe,EAAE,CAAC,CAAA;YACxF,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,GAAG,CAAC,YAAY,CAAC,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAA;YACpD,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,GAAG,CAAC,WAAW,CAAC,CAAC,CAAC,aAAa,EAAE,CAAA;YACvD,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,GAAG,CAAC,WAAW,CAAC,CAAC,CAAC,aAAa,EAAE,CAAA;QACzD,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,oBAAoB,EAAE,GAAG,EAAE;QAClC,EAAE,CAAC,yEAAyE,EAAE,KAAK,IAAI,EAAE;YACvF,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE;iBACjC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,GAAG,EAAE,CAAC;iBACzC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,GAAG,EAAE,CAAC,CAAA;YAC5C,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,CAAC,CAAA;YAElC,KAAK,CAAC,OAAO,CAAC,oBAAoB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC5C,KAAK,CAAC,eAAe,CAAC,IAAI,GAAG,IAAI,CAAA;YACnC,CAAC,CAAC,CAAA;YAEF,MAAM,KAAK,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,GAAG,CAAC,CAAA;YACrC,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,GAAG,CAAC,CAAA;YAEtC,MAAM,CAAC,KAAK,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,IAAI,EAAE,CAAC,CAAA;YACrD,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,IAAI,EAAE,CAAC,CAAA;YACtD,MAAM,CAAC,KAAK,CAAC,eAAe,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,eAAe,CAAC,CAAA;QAChE,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,aAAa,EAAE,GAAG,EAAE;QAC3B,EAAE,CAAC,2EAA2E,EAAE,KAAK,IAAI,EAAE;YACzF,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE;iBACjC,OAAO,CAAC,IAAI,KAAK,CAAC,mBAAmB,CAAC,CAAC;iBACvC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,qBAAqB,EAAE,CAAC,CAAA;YAC9D,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAElD,IAAI,OAAO,GAAG,KAAK,CAAA;YACnB,MAAM,IAAI,GAAsB,EAAE,CAAA;YAElC,KAAK,CAAC,OAAO,CAAC,oBAAoB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC5C,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,eAAe,CAAC,CAAA;gBAChC,KAAK,CAAC,eAAe,CAAC,UAAU,GAAI,KAAK,CAAC,eAAe,CAAC,UAAiC,IAAI,CAAC,CAAA;gBAChG,KAAK,CAAC,eAAe,CAAC,UAAU,GAAI,KAAK,CAAC,eAAe,CAAC,UAAqB,GAAG,CAAC,CAAA;YACrF,CAAC,CAAC,CAAA;YACF,KAAK,CAAC,OAAO,CAAC,mBAAmB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC3C,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,eAAe,CAAC,CAAA;gBAChC,IAAI,CAAC,OAAO,IAAI,KAAK,CAAC,KAAK,EAAE,CAAC;oBAC5B,OAAO,GAAG,IAAI,CAAA;oBACd,KAAK,CAAC,KAAK,GAAG,IAAI,CAAA;gBACpB,CAAC;YACH,CAAC,CAAC,CAAA;YAEF,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,MAAM,EAAE,EAAE,eAAe,EAAE,EAAE,MAAM,EAAE,KAAK,EAAE,EAAE,CAAC,CAAA;YAEjF,iEAAiE;YACjE,MAAM,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;YAC3B,6DAA6D;YAC7D,KAAK,MAAM,QAAQ,IAAI,IAAI,EAAE,CAAC;gBAC5B,MAAM,CAAC,QAAQ,CAAC,CAAC,IAAI,CAAC,MAAM,CAAC,eAAe,CAAC,CAAA;YAC/C,CAAC;YACD,2DAA2D;YAC3D,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,UAAU,EAAE,CAAC,EAAE,CAAC,CAAA;QAC1E,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,0EAA0E,EAAE,KAAK,IAAI,EAAE;YACxF,IAAI,SAAS,GAAG,CAAC,CAAA;YACjB,MAAM,IAAI,GAAG,cAAc,CAAC,OAAO,EAAE,GAAG,EAAE;gBACxC,SAAS,EAAE,CAAA;gBACX,OAAO,IAAI,eAAe,CAAC;oBACzB,SAAS,EAAE,MAAM;oBACjB,MAAM,EAAE,SAAS,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,SAAS;oBAC7C,OAAO,EAAE,CAAC,IAAI,SAAS,CAAC,SAAS,KAAK,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC;iBAC1D,CAAC,CAAA;YACJ,CAAC,CAAC,CAAA;YAEF,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE;iBACjC,OAAO,CAAC,CAAC,EAAE,IAAI,EAAE,cAAc,EAAE,IAAI,EAAE,OAAO,EAAE,SAAS,EAAE,MAAM,EAAE,KAAK,EAAE,EAAE,EAAE,CAAC,CAAC;iBAChF,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,CAAA;YAC/C,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,IAAI,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAEjE,IAAI,OAAO,GAAG,KAAK,CAAA;YACnB,MAAM,IAAI,GAAsB,EAAE,CAAA;YAElC,KAAK,CAAC,OAAO,CAAC,kBAAkB,EAAE,CAAC,KAAK,EAAE,EAAE;gBAC1C,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,eAAe,CAAC,CAAA;gBAChC,KAAK,CAAC,eAAe,CAAC,YAAY,GAAI,KAAK,CAAC,eAAe,CAAC,YAAmC,IAAI,CAAC,CAAA;gBACpG,KAAK,CAAC,eAAe,CAAC,YAAY,GAAI,KAAK,CAAC,eAAe,CAAC,YAAuB,GAAG,CAAC,CAAA;gBACvF,IAAI,CAAC,OAAO,IAAI,KAAK,CAAC,MAAM,CAAC,MAAM,KAAK,OAAO,EAAE,CAAC;oBAChD,OAAO,GAAG,IAAI,CAAA;oBACd,KAAK,CAAC,KAAK,GAAG,IAAI,CAAA;gBACpB,CAAC;YACH,CAAC,CAAC,CAAA;YAEF,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,MAAM,CAAC,WAAW,EAAE,EAAE,eAAe,EAAE,EAAE,SAAS,EAAE,KAAK,EAAE,EAAE,CAAC,CAAA;YAEzF,0DAA0D;YAC1D,MAAM,CAAC,SAAS,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;YACzB,MAAM,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;YAC3B,KAAK,MAAM,QAAQ,IAAI,IAAI,EAAE,CAAC;gBAC5B,MAAM,CAAC,QAAQ,CAAC,CAAC,IAAI,CAAC,MAAM,CAAC,eAAe,CAAC,CAAA;YAC/C,CAAC;YACD,MAAM,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,EAAE,SAAS,EAAE,KAAK,EAAE,YAAY,EAAE,CAAC,EAAE,CAAC,CAAA;QAC/E,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;AACJ,CAAC,CAAC,CAAA"}

package/dist/src/agent/__tests__/agent.stateful-model.test.d.ts

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=agent.stateful-model.test.d.ts.map

package/dist/src/agent/__tests__/agent.stateful-model.test.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"agent.stateful-model.test.d.ts","sourceRoot":"","sources":["../../../../src/agent/__tests__/agent.stateful-model.test.ts"],"names":[],"mappings":""}

package/dist/src/agent/__tests__/agent.stateful-model.test.js

@@ -0,0 +1,169 @@
+import { describe, expect, it } from 'vitest';
+import { Agent } from '../agent.js';
+import { MockMessageModel } from '../../__fixtures__/mock-message-model.js';
+import { MockSnapshotStorage } from '../../__fixtures__/mock-storage-provider.js';
+import { SlidingWindowConversationManager } from '../../conversation-manager/sliding-window-conversation-manager.js';
+import { NullConversationManager } from '../../conversation-manager/null-conversation-manager.js';
+import { SessionManager } from '../../session/session-manager.js';
+import { SNAPSHOT_SCHEMA_VERSION } from '../../types/snapshot.js';
+import { Message } from '../../types/messages.js';
+/**
+ * Mock model that advertises itself as stateful and records the modelState
+ * object it receives, so tests can verify the agent's modelState flows through.
+ */
+class StatefulMockModel extends MockMessageModel {
+    receivedOptions = [];
+    _responseIds;
+    constructor(responseIds = ['resp_1', 'resp_2', 'resp_3']) {
+        super();
+        this._responseIds = responseIds;
+    }
+    get stateful() {
+        return true;
+    }
+    async *stream(messages, options) {
+        this.receivedOptions.push(options ?? {});
+        // Simulate that the provider captured a fresh response id on the wire.
+        if (options?.modelState) {
+            const next = this._responseIds[this.receivedOptions.length - 1];
+            if (next !== undefined) {
+                options.modelState.set('responseId', next);
+            }
+        }
+        yield* super.stream(messages, options);
+    }
+}
+describe('Agent with stateful model', () => {
+    describe('constructor', () => {
+        it('throws when a conversationManager is supplied alongside a stateful model', () => {
+            const model = new StatefulMockModel();
+            expect(() => new Agent({ model, conversationManager: new SlidingWindowConversationManager({ windowSize: 5 }) })).toThrow(/stateful model/);
+        });
+        it('assigns NullConversationManager when the model is stateful', () => {
+            const model = new StatefulMockModel();
+            const agent = new Agent({ model, printer: false });
+            // Private field; access through bracket notation to avoid making it public.
+            expect(agent._conversationManager).toBeInstanceOf(NullConversationManager);
+        });
+        it('initializes modelState as an empty store', () => {
+            const model = new StatefulMockModel();
+            const agent = new Agent({ model, printer: false });
+            expect(agent.modelState.getAll()).toEqual({});
+        });
+        it('hydrates modelState from AgentConfig.modelState', () => {
+            const model = new StatefulMockModel();
+            const agent = new Agent({ model, printer: false, modelState: { responseId: 'resp_restored' } });
+            expect(agent.modelState.getAll()).toEqual({ responseId: 'resp_restored' });
+        });
+    });
+    describe('invocation', () => {
+        it('passes agent.modelState to the model via streamOptions.modelState', async () => {
+            const model = new StatefulMockModel(['resp_first']).addTurn({ type: 'textBlock', text: 'Hi' });
+            const agent = new Agent({ model, printer: false });
+            await agent.invoke('Hello');
+            expect(model.receivedOptions[0]?.modelState).toBe(agent.modelState);
+            expect(agent.modelState.getAll()).toEqual({ responseId: 'resp_first' });
+        });
+        it('clears messages after invocation since the server holds history', async () => {
+            const model = new StatefulMockModel().addTurn({ type: 'textBlock', text: 'Hi there' });
+            const agent = new Agent({ model, printer: false });
+            await agent.invoke('First turn');
+            expect(agent.messages).toEqual([]);
+        });
+        it('clears messages before SessionManager snapshots on AfterInvocationEvent', async () => {
+            // Guards the ordering of ModelPlugin vs SessionManager hooks on
+            // AfterInvocationEvent: ModelPlugin must clear messages *before*
+            // SessionManager persists the snapshot, otherwise the stored snapshot
+            // would duplicate history that the server already owns.
+            const storage = new MockSnapshotStorage();
+            const sessionManager = new SessionManager({
+                sessionId: 'test-session',
+                storage: { snapshot: storage },
+            });
+            const model = new StatefulMockModel().addTurn({ type: 'textBlock', text: 'reply' });
+            const agent = new Agent({ id: 'agent-1', model, sessionManager, printer: false });
+            await agent.invoke('hi');
+            const snapshot = await storage.loadSnapshot({
+                location: { sessionId: 'test-session', scope: 'agent', scopeId: 'agent-1' },
+            });
+            expect(snapshot).not.toBeNull();
+            expect(snapshot.data.messages).toEqual([]);
+        });
+        it('preserves modelState across invocations so previous_response_id chains', async () => {
+            const model = new StatefulMockModel(['resp_1', 'resp_2'])
+                .addTurn({ type: 'textBlock', text: 'one' })
+                .addTurn({ type: 'textBlock', text: 'two' });
+            const agent = new Agent({ model, printer: false });
+            await agent.invoke('turn 1');
+            expect(agent.modelState.getAll()).toEqual({ responseId: 'resp_1' });
+            await agent.invoke('turn 2');
+            expect(agent.modelState.getAll()).toEqual({ responseId: 'resp_2' });
+            // Both turns should have seen the state at invocation time.
+            expect(model.receivedOptions).toHaveLength(2);
+        });
+    });
+    describe('stateless model (default)', () => {
+        it('does not clear messages after invocation', async () => {
+            const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
+            const agent = new Agent({ model, printer: false });
+            await agent.invoke('Hi');
+            // user message + assistant reply
+            expect(agent.messages.length).toBe(2);
+        });
+        it('uses the caller-provided conversationManager', () => {
+            const model = new MockMessageModel();
+            const convo = new SlidingWindowConversationManager({ windowSize: 7 });
+            const agent = new Agent({ model, conversationManager: convo });
+            expect(agent._conversationManager).toBe(convo);
+        });
+    });
+    describe('SessionManager restore guard', () => {
+        // Pre-seeds a session snapshot with messages, then verifies that SessionManager
+        // discards those messages on restore when the model is stateful.
+        async function setupStorageWithMessages(agentId, sessionId) {
+            const storage = new MockSnapshotStorage();
+            await storage.saveSnapshot({
+                location: { sessionId, scope: 'agent', scopeId: agentId },
+                snapshotId: 'latest',
+                isLatest: true,
+                snapshot: {
+                    scope: 'agent',
+                    schemaVersion: SNAPSHOT_SCHEMA_VERSION,
+                    createdAt: new Date().toISOString(),
+                    data: {
+                        messages: [{ role: 'user', content: [{ text: 'old turn' }] }],
+                        state: {},
+                        systemPrompt: null,
+                        modelState: {},
+                    },
+                    appData: {},
+                },
+            });
+            return storage;
+        }
+        it('discards restored messages when the model is stateful', async () => {
+            const storage = await setupStorageWithMessages('agent-1', 'session-stateful');
+            const sessionManager = new SessionManager({
+                sessionId: 'session-stateful',
+                storage: { snapshot: storage },
+            });
+            const model = new StatefulMockModel();
+            const agent = new Agent({ id: 'agent-1', model, sessionManager, printer: false });
+            await agent.initialize();
+            expect(agent.messages).toEqual([]);
+        });
+        it('restores messages when the model is stateless', async () => {
+            const storage = await setupStorageWithMessages('agent-2', 'session-stateless');
+            const sessionManager = new SessionManager({
+                sessionId: 'session-stateless',
+                storage: { snapshot: storage },
+            });
+            const model = new MockMessageModel();
+            const agent = new Agent({ id: 'agent-2', model, sessionManager, printer: false });
+            await agent.initialize();
+            expect(agent.messages).toHaveLength(1);
+            expect(agent.messages[0].role).toBe('user');
+        });
+    });
+});
+//# sourceMappingURL=agent.stateful-model.test.js.map

package/dist/src/agent/__tests__/agent.stateful-model.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"agent.stateful-model.test.js","sourceRoot":"","sources":["../../../../src/agent/__tests__/agent.stateful-model.test.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,QAAQ,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,QAAQ,CAAA;AAC7C,OAAO,EAAE,KAAK,EAAE,MAAM,aAAa,CAAA;AACnC,OAAO,EAAE,gBAAgB,EAAE,MAAM,0CAA0C,CAAA;AAC3E,OAAO,EAAE,mBAAmB,EAAE,MAAM,6CAA6C,CAAA;AACjF,OAAO,EAAE,gCAAgC,EAAE,MAAM,mEAAmE,CAAA;AACpH,OAAO,EAAE,uBAAuB,EAAE,MAAM,yDAAyD,CAAA;AACjG,OAAO,EAAE,cAAc,EAAE,MAAM,kCAAkC,CAAA;AACjE,OAAO,EAAE,uBAAuB,EAAE,MAAM,yBAAyB,CAAA;AACjE,OAAO,EAAE,OAAO,EAAE,MAAM,yBAAyB,CAAA;AAKjD;;;GAGG;AACH,MAAM,iBAAkB,SAAQ,gBAAgB;IACrC,eAAe,GAAoB,EAAE,CAAA;IAC7B,YAAY,CAAU;IAEvC,YAAY,cAAwB,CAAC,QAAQ,EAAE,QAAQ,EAAE,QAAQ,CAAC;QAChE,KAAK,EAAE,CAAA;QACP,IAAI,CAAC,YAAY,GAAG,WAAW,CAAA;IACjC,CAAC;IAED,IAAa,QAAQ;QACnB,OAAO,IAAI,CAAA;IACb,CAAC;IAEQ,KAAK,CAAC,CAAC,MAAM,CAAC,QAAmB,EAAE,OAAuB;QACjE,IAAI,CAAC,eAAe,CAAC,IAAI,CAAC,OAAO,IAAI,EAAE,CAAC,CAAA;QACxC,uEAAuE;QACvE,IAAI,OAAO,EAAE,UAAU,EAAE,CAAC;YACxB,MAAM,IAAI,GAAG,IAAI,CAAC,YAAY,CAAC,IAAI,CAAC,eAAe,CAAC,MAAM,GAAG,CAAC,CAAC,CAAA;YAC/D,IAAI,IAAI,KAAK,SAAS,EAAE,CAAC;gBACvB,OAAO,CAAC,UAAU,CAAC,GAAG,CAAC,YAAY,EAAE,IAAI,CAAC,CAAA;YAC5C,CAAC;QACH,CAAC;QACD,KAAK,CAAC,CAAC,KAAK,CAAC,MAAM,CAAC,QAAQ,EAAE,OAAO,CAAC,CAAA;IACxC,CAAC;CACF;AAED,QAAQ,CAAC,2BAA2B,EAAE,GAAG,EAAE;IACzC,QAAQ,CAAC,aAAa,EAAE,GAAG,EAAE;QAC3B,EAAE,CAAC,0EAA0E,EAAE,GAAG,EAAE;YAClF,MAAM,KAAK,GAAG,IAAI,iBAAiB,EAAE,CAAA;YACrC,MAAM,CACJ,GAAG,EAAE,CAAC,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,mBAAmB,EAAE,IAAI,gCAAgC,CAAC,EAAE,UAAU,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CACzG,CAAC,OAAO,CAAC,gBAAgB,CAAC,CAAA;QAC7B,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,4DAA4D,EAAE,GAAG,EAAE;YACpE,MAAM,KAAK,GAAG,IAAI,iBAAiB,EAAE,CAAA;YACrC,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAClD,4EAA4E;YAC5E,MAAM,CAAE,KAAsD,CAAC,oBAAoB,CAAC,CAAC,cAAc,CACjG,uBAAuB,CACxB,CAAA;QACH,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,0CAA0C,EAAE,GAAG,EAAE;YAClD,MAAM,KAAK,GAAG,IAAI,iBAAiB,EAAE,CAAA;YACrC,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAClD,MAAM,CAAC,KAAK,CAAC,UAAU,CAAC,MAAM,EAAE,CAAC,CAAC,OAAO,CAAC,EAAE,CAAC,CAAA;QAC/C,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,iDAAiD,EAAE,GAAG,EAAE;YACzD,MAAM,KAAK,GAAG,IAAI,iBAAiB,EAAE,CAAA;YACrC,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,UAAU,EAAE,EAAE,UAAU,EAAE,eAAe,EAAE,EAAE,CAAC,CAAA;YAC/F,MAAM,CAAC,KAAK,CAAC,UAAU,CAAC,MAAM,EAAE,CAAC,CAAC,OAAO,CAAC,EAAE,UAAU,EAAE,eAAe,EAAE,CAAC,CAAA;QAC5E,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,YAAY,EAAE,GAAG,EAAE;QAC1B,EAAE,CAAC,mEAAmE,EAAE,KAAK,IAAI,EAAE;YACjF,MAAM,KAAK,GAAG,IAAI,iBAAiB,CAAC,CAAC,YAAY,CAAC,CAAC,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,IAAI,EAAE,CAAC,CAAA;YAC9F,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAClD,MAAM,KAAK,CAAC,MAAM,CAAC,OAAO,CAAC,CAAA;YAC3B,MAAM,CAAC,KAAK,CAAC,eAAe,CAAC,CAAC,CAAC,EAAE,UAAU,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,UAAU,CAAC,CAAA;YACnE,MAAM,CAAC,KAAK,CAAC,UAAU,CAAC,MAAM,EAAE,CAAC,CAAC,OAAO,CAAC,EAAE,UAAU,EAAE,YAAY,EAAE,CAAC,CAAA;QACzE,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,iEAAiE,EAAE,KAAK,IAAI,EAAE;YAC/E,MAAM,KAAK,GAAG,IAAI,iBAAiB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,UAAU,EAAE,CAAC,CAAA;YACtF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAClD,MAAM,KAAK,CAAC,MAAM,CAAC,YAAY,CAAC,CAAA;YAChC,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,CAAC,OAAO,CAAC,EAAE,CAAC,CAAA;QACpC,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,yEAAyE,EAAE,KAAK,IAAI,EAAE;YACvF,gEAAgE;YAChE,iEAAiE;YACjE,sEAAsE;YACtE,wDAAwD;YACxD,MAAM,OAAO,GAAG,IAAI,mBAAmB,EAAE,CAAA;YACzC,MAAM,cAAc,GAAG,IAAI,cAAc,CAAC;gBACxC,SAAS,EAAE,cAAc;gBACzB,OAAO,EAAE,EAAE,QAAQ,EAAE,OAAO,EAAE;aAC/B,CAAC,CAAA;YACF,MAAM,KAAK,GAAG,IAAI,iBAAiB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YACnF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,EAAE,EAAE,SAAS,EAAE,KAAK,EAAE,cAAc,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAEjF,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,CAAC,CAAA;YAExB,MAAM,QAAQ,GAAG,MAAM,OAAO,CAAC,YAAY,CAAC;gBAC1C,QAAQ,EAAE,EAAE,SAAS,EAAE,cAAc,EAAE,KAAK,EAAE,OAAO,EAAE,OAAO,EAAE,SAAS,EAAE;aAC5E,CAAC,CAAA;YACF,MAAM,CAAC,QAAQ,CAAC,CAAC,GAAG,CAAC,QAAQ,EAAE,CAAA;YAC/B,MAAM,CAAE,QAAS,CAAC,IAAgC,CAAC,QAAQ,CAAC,CAAC,OAAO,CAAC,EAAE,CAAC,CAAA;QAC1E,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,wEAAwE,EAAE,KAAK,IAAI,EAAE;YACtF,MAAM,KAAK,GAAG,IAAI,iBAAiB,CAAC,CAAC,QAAQ,EAAE,QAAQ,CAAC,CAAC;iBACtD,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,KAAK,EAAE,CAAC;iBAC3C,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,KAAK,EAAE,CAAC,CAAA;YAC9C,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAElD,MAAM,KAAK,CAAC,MAAM,CAAC,QAAQ,CAAC,CAAA;YAC5B,MAAM,CAAC,KAAK,CAAC,UAAU,CAAC,MAAM,EAAE,CAAC,CAAC,OAAO,CAAC,EAAE,UAAU,EAAE,QAAQ,EAAE,CAAC,CAAA;YAEnE,MAAM,KAAK,CAAC,MAAM,CAAC,QAAQ,CAAC,CAAA;YAC5B,MAAM,CAAC,KAAK,CAAC,UAAU,CAAC,MAAM,EAAE,CAAC,CAAC,OAAO,CAAC,EAAE,UAAU,EAAE,QAAQ,EAAE,CAAC,CAAA;YAEnE,4DAA4D;YAC5D,MAAM,CAAC,KAAK,CAAC,eAAe,CAAC,CAAC,YAAY,CAAC,CAAC,CAAC,CAAA;QAC/C,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,2BAA2B,EAAE,GAAG,EAAE;QACzC,EAAE,CAAC,0CAA0C,EAAE,KAAK,IAAI,EAAE;YACxD,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,WAAW,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;YAClF,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YAClD,MAAM,KAAK,CAAC,MAAM,CAAC,IAAI,CAAC,CAAA;YACxB,iCAAiC;YACjC,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;QACvC,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,8CAA8C,EAAE,GAAG,EAAE;YACtD,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAA;YACpC,MAAM,KAAK,GAAG,IAAI,gCAAgC,CAAC,EAAE,UAAU,EAAE,CAAC,EAAE,CAAC,CAAA;YACrE,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,KAAK,EAAE,mBAAmB,EAAE,KAAK,EAAE,CAAC,CAAA;YAC9D,MAAM,CAAE,KAAsD,CAAC,oBAAoB,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAA;QAClG,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,QAAQ,CAAC,8BAA8B,EAAE,GAAG,EAAE;QAC5C,gFAAgF;QAChF,iEAAiE;QACjE,KAAK,UAAU,wBAAwB,CAAC,OAAe,EAAE,SAAiB;YACxE,MAAM,OAAO,GAAG,IAAI,mBAAmB,EAAE,CAAA;YACzC,MAAM,OAAO,CAAC,YAAY,CAAC;gBACzB,QAAQ,EAAE,EAAE,SAAS,EAAE,KAAK,EAAE,OAAO,EAAE,OAAO,EAAE,OAAO,EAAE;gBACzD,UAAU,EAAE,QAAQ;gBACpB,QAAQ,EAAE,IAAI;gBACd,QAAQ,EAAE;oBACR,KAAK,EAAE,OAAO;oBACd,aAAa,EAAE,uBAAuB;oBACtC,SAAS,EAAE,IAAI,IAAI,EAAE,CAAC,WAAW,EAAE;oBACnC,IAAI,EAAE;wBACJ,QAAQ,EAAE,CAAC,EAAE,IAAI,EAAE,MAAM,EAAE,OAAO,EAAE,CAAC,EAAE,IAAI,EAAE,UAAU,EAAE,CAAC,EAAE,CAAyB;wBACrF,KAAK,EAAE,EAAE;wBACT,YAAY,EAAE,IAAI;wBAClB,UAAU,EAAE,EAAE;qBACf;oBACD,OAAO,EAAE,EAAE;iBACZ;aACF,CAAC,CAAA;YACF,OAAO,OAAO,CAAA;QAChB,CAAC;QAED,EAAE,CAAC,uDAAuD,EAAE,KAAK,IAAI,EAAE;YACrE,MAAM,OAAO,GAAG,MAAM,wBAAwB,CAAC,SAAS,EAAE,kBAAkB,CAAC,CAAA;YAC7E,MAAM,cAAc,GAAG,IAAI,cAAc,CAAC;gBACxC,SAAS,EAAE,kBAAkB;gBAC7B,OAAO,EAAE,EAAE,QAAQ,EAAE,OAAO,EAAE;aAC/B,CAAC,CAAA;YACF,MAAM,KAAK,GAAG,IAAI,iBAAiB,EAAE,CAAA;YACrC,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,EAAE,EAAE,SAAS,EAAE,KAAK,EAAE,cAAc,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YACjF,MAAM,KAAK,CAAC,UAAU,EAAE,CAAA;YACxB,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,CAAC,OAAO,CAAC,EAAE,CAAC,CAAA;QACpC,CAAC,CAAC,CAAA;QAEF,EAAE,CAAC,+CAA+C,EAAE,KAAK,IAAI,EAAE;YAC7D,MAAM,OAAO,GAAG,MAAM,wBAAwB,CAAC,SAAS,EAAE,mBAAmB,CAAC,CAAA;YAC9E,MAAM,cAAc,GAAG,IAAI,cAAc,CAAC;gBACxC,SAAS,EAAE,mBAAmB;gBAC9B,OAAO,EAAE,EAAE,QAAQ,EAAE,OAAO,EAAE;aAC/B,CAAC,CAAA;YACF,MAAM,KAAK,GAAG,IAAI,gBAAgB,EAAE,CAAA;YACpC,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC,EAAE,EAAE,EAAE,SAAS,EAAE,KAAK,EAAE,cAAc,EAAE,OAAO,EAAE,KAAK,EAAE,CAAC,CAAA;YACjF,MAAM,KAAK,CAAC,UAAU,EAAE,CAAA;YACxB,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,CAAC,YAAY,CAAC,CAAC,CAAC,CAAA;YACtC,MAAM,CAAC,KAAK,CAAC,QAAQ,CAAC,CAAC,CAAE,CAAC,IAAI,CAAC,CAAC,IAAI,CAAC,MAAM,CAAC,CAAA;QAC9C,CAAC,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;AACJ,CAAC,CAAC,CAAA"}

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

@@ -7,7 +7,7 @@ import { createMockTool, createRandomTool } from '../../__fixtures__/tool-helper
 import { ConcurrentInvocationError } from '../../errors.js';
 import { MaxTokensError, TextBlock, CachePointBlock, Message, ToolUseBlock, ToolResultBlock, ReasoningBlock, GuardContentBlock, ImageBlock, VideoBlock, DocumentBlock, } from '../../index.js';
 import { AgentPrinter } from '../printer.js';
-import { AfterInvocationEvent, AfterToolCallEvent, AfterToolsEvent, BeforeInvocationEvent, BeforeToolsEvent, } from '../../hooks/events.js';
+import { AfterInvocationEvent, AfterToolCallEvent, AfterToolsEvent, BeforeInvocationEvent, BeforeModelCallEvent, BeforeToolsEvent, } from '../../hooks/events.js';
 import { BedrockModel } from '../../models/bedrock.js';
 import { StructuredOutputError } from '../../errors.js';
 import { expectLoopMetrics } from '../../__fixtures__/metrics-helpers.js';
@@ -36,7 +36,7 @@ describe('Agent', () => {
                 const { items } = await collectGenerator(agent.stream('Test prompt'));
                 expect(items.length).toBeGreaterThan(0);
                 const firstItem = items[0];
-                expect(firstItem).toEqual(new BeforeInvocationEvent({ agent: agent }));
+                expect(firstItem).toEqual(new BeforeInvocationEvent({ agent: agent, invocationState: {} }));
             });
             it('returns AgentResult as generator return value', async () => {
                 const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' });
@@ -89,6 +89,7 @@ describe('Agent', () => {
                         role: 'assistant',
                         content: [new ToolUseBlock({ name: 'testTool', toolUseId: 'tool-1', input: {} })],
                     }),
+                    invocationState: {},
                 }));
                 expect(afterTools).toBeDefined();
                 expect(afterTools?.type).toBe('afterToolsEvent');
@@ -1155,4 +1156,100 @@ describe('Agent._redactLastMessage', () => {
         expect(agent['messages']).toHaveLength(0);
     });
 });
+describe('_estimateInputTokens', () => {
+    function captureProjectedTokens(agent) {
+        return new Promise((resolve) => {
+            agent.addHook(BeforeModelCallEvent, (event) => {
+                resolve(event.projectedInputTokens);
+            });
+        });
+    }
+    it('uses full estimation on cold start (no prior usage metadata)', async () => {
+        const model = new MockMessageModel();
+        model.addTurn({ type: 'textBlock', text: 'Hello' });
+        const countTokensSpy = vi.spyOn(model, 'countTokens');
+        countTokensSpy.mockResolvedValue(42);
+        const agent = new Agent({ model, printer: false });
+        const tokenPromise = captureProjectedTokens(agent);
+        await agent.invoke('Hi');
+        expect(await tokenPromise).toBe(42);
+        expect(countTokensSpy).toHaveBeenCalledWith(expect.any(Array), expect.any(Object));
+    });
+    it('uses known baseline when no new messages after last assistant', async () => {
+        const model = new MockMessageModel();
+        model.addTurn({ type: 'textBlock', text: 'Hello' });
+        const agent = new Agent({
+            model,
+            printer: false,
+            messages: [
+                new Message({ role: 'user', content: [new TextBlock('Hi')] }),
+                new Message({
+                    role: 'assistant',
+                    content: [new TextBlock('Hello')],
+                    metadata: { usage: { inputTokens: 100, outputTokens: 20, totalTokens: 120 } },
+                }),
+            ],
+        });
+        // Invoke with no args — no new user message appended, so the last assistant
+        // message is still the final message and newMessages.length === 0
+        const tokenPromise = captureProjectedTokens(agent);
+        await agent.invoke([]);
+        // baseline = inputTokens(100) + outputTokens(20) = 120
+        expect(await tokenPromise).toBe(120);
+    });
+    it('returns undefined projectedInputTokens when estimation fails', async () => {
+        const model = new MockMessageModel();
+        model.addTurn({ type: 'textBlock', text: 'Hello' });
+        vi.spyOn(model, 'countTokens').mockRejectedValue(new Error('API unavailable'));
+        const agent = new Agent({ model, printer: false });
+        const tokenPromise = captureProjectedTokens(agent);
+        await agent.invoke('Hi');
+        expect(await tokenPromise).toBeUndefined();
+    });
+    it('estimates delta for new messages after last assistant', async () => {
+        const model = new MockMessageModel();
+        model
+            .addTurn([{ type: 'toolUseBlock', name: 'test', toolUseId: 'id-1', input: {} }], {
+            usage: { inputTokens: 100, outputTokens: 30, totalTokens: 130 },
+        })
+            .addTurn({ type: 'textBlock', text: 'Done' });
+        const countTokensSpy = vi.spyOn(model, 'countTokens');
+        countTokensSpy.mockResolvedValue(50);
+        const tool = createMockTool('test', () => new ToolResultBlock({
+            toolUseId: 'id-1',
+            status: 'success',
+            content: [new TextBlock('result')],
+        }));
+        const agent = new Agent({ model, tools: [tool], printer: false });
+        // Capture the second BeforeModelCallEvent (after tool execution)
+        let callCount = 0;
+        const tokenPromise = new Promise((resolve) => {
+            agent.addHook(BeforeModelCallEvent, (event) => {
+                callCount++;
+                if (callCount === 2)
+                    resolve(event.projectedInputTokens);
+            });
+        });
+        await agent.invoke('Use the tool');
+        // baseline (100+30) + estimated delta (50) = 180
+        expect(await tokenPromise).toBe(180);
+        expect(countTokensSpy).toHaveBeenCalled();
+    });
+    it('uses baseline from prior invocation on second invoke', async () => {
+        const model = new MockMessageModel();
+        model
+            .addTurn({ type: 'textBlock', text: 'First response' }, { usage: { inputTokens: 200, outputTokens: 50, totalTokens: 250 } })
+            .addTurn({ type: 'textBlock', text: 'Second response' });
+        const countTokensSpy = vi.spyOn(model, 'countTokens');
+        countTokensSpy.mockResolvedValue(15);
+        const agent = new Agent({ model, printer: false });
+        await agent.invoke('First question');
+        // Second invocation — the user message "Second question" is appended after
+        // the assistant message with usage metadata, so it hits the baseline + delta path
+        const tokenPromise = captureProjectedTokens(agent);
+        await agent.invoke('Second question');
+        // baseline (200+50) + estimated delta for new user message (15) = 265
+        expect(await tokenPromise).toBe(265);
+    });
+});
 //# sourceMappingURL=agent.test.js.map