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

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');
@@ -1074,6 +1075,25 @@ describe('Agent', () => {
             const second = await agent.invoke('Second');
             expect(second.structuredOutput).toEqual({ name: 'Bob' });
         });
+        it('skips structured output extraction when AfterToolsEvent.endTurn halts the loop', async () => {
+            const schema = z.object({ name: z.string() });
+            const model = new MockMessageModel()
+                .addTurn({
+                type: 'toolUseBlock',
+                name: 'strands_structured_output',
+                toolUseId: 'tool-1',
+                input: { name: 'John' },
+            })
+                .addTurn({ type: 'textBlock', text: 'Done' });
+            const agent = new Agent({ model, structuredOutputSchema: schema });
+            agent.addHook(AfterToolsEvent, (event) => {
+                event.endTurn = true;
+            });
+            const result = await agent.invoke('Test');
+            expect(result.stopReason).toBe('endTurn');
+            expect(result.structuredOutput).toBeUndefined();
+            expect(model.callCount).toBe(1);
+        });
     });
 });
 describe('Agent._redactLastMessage', () => {
@@ -1155,4 +1175,199 @@ 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);
+    });
+});
+describe('normalizeToolUseNames', () => {
+    it('replaces invalid tool-use names with INVALID_TOOL_NAME before calling model', async () => {
+        const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'ok' });
+        const streamSpy = vi.spyOn(model, 'stream');
+        const agent = new Agent({
+            model,
+            printer: false,
+            messages: [
+                new Message({ role: 'user', content: [new TextBlock('do thing')] }),
+                new Message({
+                    role: 'assistant',
+                    content: [new ToolUseBlock({ name: 'bad name!', toolUseId: 'tu-1', input: {} })],
+                }),
+                new Message({
+                    role: 'user',
+                    content: [
+                        new ToolResultBlock({
+                            toolUseId: 'tu-1',
+                            status: 'success',
+                            content: [new TextBlock('result')],
+                        }),
+                    ],
+                }),
+            ],
+        });
+        await agent.invoke('continue');
+        const sentMessages = streamSpy.mock.calls[0]?.[0];
+        const sentToolUse = sentMessages
+            .find((m) => m.role === 'assistant')
+            .content.find((b) => b.type === 'toolUseBlock');
+        expect(sentToolUse).toStrictEqual(new ToolUseBlock({ name: 'INVALID_TOOL_NAME', toolUseId: 'tu-1', input: {} }));
+        // Agent's stored history is not mutated.
+        const storedToolUse = agent.messages
+            .find((m) => m.role === 'assistant')
+            .content.find((b) => b.type === 'toolUseBlock');
+        expect(storedToolUse).toStrictEqual(new ToolUseBlock({ name: 'bad name!', toolUseId: 'tu-1', input: {} }));
+    });
+    it('preserves reasoningSignature on replaced tool-use blocks', async () => {
+        const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'ok' });
+        const streamSpy = vi.spyOn(model, 'stream');
+        const agent = new Agent({
+            model,
+            printer: false,
+            messages: [
+                new Message({ role: 'user', content: [new TextBlock('do thing')] }),
+                new Message({
+                    role: 'assistant',
+                    content: [new ToolUseBlock({ name: 'bad!', toolUseId: 'tu-1', input: {}, reasoningSignature: 'sig-abc' })],
+                }),
+                new Message({
+                    role: 'user',
+                    content: [new ToolResultBlock({ toolUseId: 'tu-1', status: 'success', content: [new TextBlock('ok')] })],
+                }),
+            ],
+        });
+        await agent.invoke('continue');
+        const sentMessages = streamSpy.mock.calls[0]?.[0];
+        const sentToolUse = sentMessages
+            .find((m) => m.role === 'assistant')
+            .content.find((b) => b.type === 'toolUseBlock');
+        expect(sentToolUse).toStrictEqual(new ToolUseBlock({
+            name: 'INVALID_TOOL_NAME',
+            toolUseId: 'tu-1',
+            input: {},
+            reasoningSignature: 'sig-abc',
+        }));
+    });
+    it('leaves valid names untouched', async () => {
+        const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'ok' });
+        const streamSpy = vi.spyOn(model, 'stream');
+        const agent = new Agent({
+            model,
+            printer: false,
+            messages: [
+                new Message({ role: 'user', content: [new TextBlock('do thing')] }),
+                new Message({
+                    role: 'assistant',
+                    content: [new ToolUseBlock({ name: 'good_tool-1', toolUseId: 'tu-1', input: {} })],
+                }),
+                new Message({
+                    role: 'user',
+                    content: [
+                        new ToolResultBlock({
+                            toolUseId: 'tu-1',
+                            status: 'success',
+                            content: [new TextBlock('result')],
+                        }),
+                    ],
+                }),
+            ],
+        });
+        await agent.invoke('continue');
+        const sentMessages = streamSpy.mock.calls[0]?.[0];
+        const sentToolUse = sentMessages
+            .find((m) => m.role === 'assistant')
+            .content.find((b) => b.type === 'toolUseBlock');
+        expect(sentToolUse).toStrictEqual(new ToolUseBlock({ name: 'good_tool-1', toolUseId: 'tu-1', input: {} }));
+    });
+});
 //# sourceMappingURL=agent.test.js.map