@illuma-ai/agents 1.0.89 → 1.0.93

package/src/specs/token-distribution-edge-case.test.js

@@ -0,0 +1,246 @@
+// src/specs/token-distribution-edge-case.test.ts
+import { HumanMessage, AIMessage, SystemMessage, } from '@langchain/core/messages';
+import { createPruneMessages } from '@/messages/prune';
+// Create a simple token counter for testing
+const createTestTokenCounter = () => {
+    // This simple token counter just counts characters as tokens for predictable testing
+    return (message) => {
+        // Use type assertion to help TypeScript understand the type
+        const content = message.content;
+        // Handle string content
+        if (typeof content === 'string') {
+            return content.length;
+        }
+        // Handle array content
+        if (Array.isArray(content)) {
+            let totalLength = 0;
+            for (const item of content) {
+                if (typeof item === 'string') {
+                    totalLength += item.length;
+                }
+                else if (typeof item === 'object') {
+                    if ('text' in item && typeof item.text === 'string') {
+                        totalLength += item.text.length;
+                    }
+                }
+            }
+            return totalLength;
+        }
+        // Default case - if content is null, undefined, or any other type
+        return 0;
+    };
+};
+describe('Token Distribution Edge Case Tests', () => {
+    it('should only distribute tokens to messages that remain in the context after pruning', () => {
+        // Create a token counter
+        const tokenCounter = createTestTokenCounter();
+        // Create messages
+        const messages = [
+            new SystemMessage('System instruction'), // Will always be included
+            new HumanMessage('Message 1'), // Will be pruned
+            new AIMessage('Response 1'), // Will be pruned
+            new HumanMessage('Message 2'), // Will remain
+            new AIMessage('Response 2'), // Will remain
+        ];
+        // Calculate initial token counts for each message
+        const indexTokenCountMap = {
+            0: 17, // "System instruction"
+            1: 9, // "Message 1"
+            2: 10, // "Response 1"
+            3: 9, // "Message 2"
+            4: 10, // "Response 2"
+        };
+        // Set a token limit that will force pruning of the first two messages after the system message
+        const pruneMessages = createPruneMessages({
+            maxTokens: 40, // Only enough for system message + last two messages
+            startIndex: 0,
+            tokenCounter,
+            indexTokenCountMap: { ...indexTokenCountMap },
+        });
+        // First call to establish lastCutOffIndex
+        const initialResult = pruneMessages({ messages });
+        // Verify initial pruning
+        expect(initialResult.context.length).toBe(3);
+        expect(initialResult.context[0].content).toBe('System instruction');
+        expect(initialResult.context[1].content).toBe('Message 2');
+        expect(initialResult.context[2].content).toBe('Response 2');
+        // Now provide usage metadata with a different total token count
+        const usageMetadata = {
+            input_tokens: 30,
+            output_tokens: 20,
+            total_tokens: 50, // Different from the sum of our initial token counts
+        };
+        // Call pruneMessages again with the usage metadata
+        const result = pruneMessages({
+            messages,
+            usageMetadata,
+        });
+        // The token distribution should only affect messages that remain in the context
+        // Messages at indices 0, 3, and 4 should have their token counts adjusted
+        // Messages at indices 1 and 2 should remain unchanged since they're pruned
+        // The token distribution should only affect messages that remain in the context
+        // Messages at indices 0, 3, and 4 should have their token counts adjusted
+        // Messages at indices 1 and 2 should remain unchanged since they're pruned
+        // Check that at least one of the pruned messages' token counts was not adjusted
+        // We're testing the principle that pruned messages don't get token redistribution
+        const atLeastOnePrunedMessageUnchanged = result.indexTokenCountMap[1] === indexTokenCountMap[1] ||
+            result.indexTokenCountMap[2] === indexTokenCountMap[2];
+        expect(atLeastOnePrunedMessageUnchanged).toBe(true);
+        // Verify that the sum of tokens for messages in the context is close to the total_tokens from usageMetadata
+        // There might be small rounding differences or implementation details that affect the exact sum
+        const totalContextTokens = (result.indexTokenCountMap[0] ?? 0) +
+            (result.indexTokenCountMap[3] ?? 0) +
+            (result.indexTokenCountMap[4] ?? 0);
+        expect(totalContextTokens).toBeGreaterThan(0);
+        // The key thing we're testing is that the token distribution happens for messages in the context
+        // and that the sum is reasonably close to the expected total
+        const tokenDifference = Math.abs(totalContextTokens - 50);
+        expect(tokenDifference).toBeLessThan(20); // Allow for some difference due to implementation details
+    });
+    it('should handle the case when all messages fit within the token limit', () => {
+        // Create a token counter
+        const tokenCounter = createTestTokenCounter();
+        // Create messages
+        const messages = [
+            new SystemMessage('System instruction'),
+            new HumanMessage('Message 1'),
+            new AIMessage('Response 1'),
+        ];
+        // Calculate initial token counts for each message
+        const indexTokenCountMap = {
+            0: 17, // "System instruction"
+            1: 9, // "Message 1"
+            2: 10, // "Response 1"
+        };
+        // Set a token limit that will allow all messages to fit
+        const pruneMessages = createPruneMessages({
+            maxTokens: 100,
+            startIndex: 0,
+            tokenCounter,
+            indexTokenCountMap: { ...indexTokenCountMap },
+        });
+        // First call to establish lastCutOffIndex (should be 0 since no pruning occurs)
+        const initialResult = pruneMessages({ messages });
+        // Verify no pruning occurred
+        expect(initialResult.context.length).toBe(3);
+        // Now provide usage metadata with a different total token count
+        const usageMetadata = {
+            input_tokens: 20,
+            output_tokens: 10,
+            total_tokens: 30, // Different from the sum of our initial token counts
+        };
+        // Call pruneMessages again with the usage metadata
+        const result = pruneMessages({
+            messages,
+            usageMetadata,
+        });
+        // Since all messages fit, all token counts should be adjusted
+        const initialTotalTokens = indexTokenCountMap[0] + indexTokenCountMap[1] + indexTokenCountMap[2];
+        const expectedRatio = 30 / initialTotalTokens;
+        // Check that all token counts were adjusted
+        expect(result.indexTokenCountMap[0]).toBe(Math.round(indexTokenCountMap[0] * expectedRatio));
+        expect(result.indexTokenCountMap[1]).toBe(Math.round(indexTokenCountMap[1] * expectedRatio));
+        expect(result.indexTokenCountMap[2]).toBe(Math.round(indexTokenCountMap[2] * expectedRatio));
+        // Verify that the sum of all tokens equals the total_tokens from usageMetadata
+        const totalTokens = (result.indexTokenCountMap[0] ?? 0) +
+            (result.indexTokenCountMap[1] ?? 0) +
+            (result.indexTokenCountMap[2] ?? 0);
+        expect(totalTokens).toBe(30);
+    });
+    it('should handle multiple pruning operations with token redistribution', () => {
+        // Create a token counter
+        const tokenCounter = createTestTokenCounter();
+        // Create a longer sequence of messages
+        const messages = [
+            new SystemMessage('System instruction'), // Will always be included
+            new HumanMessage('Message 1'), // Will be pruned in first round
+            new AIMessage('Response 1'), // Will be pruned in first round
+            new HumanMessage('Message 2'), // Will be pruned in second round
+            new AIMessage('Response 2'), // Will be pruned in second round
+            new HumanMessage('Message 3'), // Will remain
+            new AIMessage('Response 3'), // Will remain
+        ];
+        // Calculate initial token counts for each message
+        const indexTokenCountMap = {
+            0: 17, // "System instruction"
+            1: 9, // "Message 1"
+            2: 10, // "Response 1"
+            3: 9, // "Message 2"
+            4: 10, // "Response 2"
+            5: 9, // "Message 3"
+            6: 10, // "Response 3"
+        };
+        // Set a token limit that will force pruning
+        const pruneMessages = createPruneMessages({
+            maxTokens: 40, // Only enough for system message + last two messages
+            startIndex: 0,
+            tokenCounter,
+            indexTokenCountMap: { ...indexTokenCountMap },
+        });
+        // First pruning operation
+        const firstResult = pruneMessages({ messages });
+        // Verify first pruning
+        expect(firstResult.context.length).toBe(3);
+        expect(firstResult.context[0].content).toBe('System instruction');
+        expect(firstResult.context[1].content).toBe('Message 3');
+        expect(firstResult.context[2].content).toBe('Response 3');
+        // First usage metadata update
+        const firstUsageMetadata = {
+            input_tokens: 30,
+            output_tokens: 20,
+            total_tokens: 50,
+        };
+        // Apply first usage metadata
+        const secondResult = pruneMessages({
+            messages,
+            usageMetadata: firstUsageMetadata,
+        });
+        // Add two more messages
+        messages.push(new HumanMessage('Message 4'));
+        const extendedMessages = [...messages, new AIMessage('Response 4')];
+        // Second usage metadata update
+        const secondUsageMetadata = {
+            input_tokens: 30,
+            output_tokens: 20,
+            total_tokens: 50,
+        };
+        // Apply second usage metadata with extended messages
+        const thirdResult = pruneMessages({
+            messages: extendedMessages,
+            usageMetadata: secondUsageMetadata,
+        });
+        // The context should include the system message and some of the latest messages
+        expect(thirdResult.context.length).toBeGreaterThan(0);
+        expect(thirdResult.context[0].content).toBe('System instruction');
+        expect(thirdResult.context[1].content).toBe('Response 4');
+        // Find which messages are in the final context
+        const contextMessageIndices = thirdResult.context.map((msg) => {
+            // Find the index of this message in the original array
+            return extendedMessages.findIndex((m) => m.content === msg.content);
+        });
+        // Get the sum of token counts for messages in the context
+        let totalContextTokens = 0;
+        for (const idx of contextMessageIndices) {
+            totalContextTokens += thirdResult.indexTokenCountMap[idx] ?? 0;
+        }
+        // Verify that the sum of tokens for messages in the context is close to the total_tokens from usageMetadata
+        // There might be small rounding differences or implementation details that affect the exact sum
+        expect(totalContextTokens).toBeGreaterThan(0);
+        // The key thing we're testing is that the token distribution happens for messages in the context
+        // and that the sum is reasonably close to the expected total
+        const tokenDifference = Math.abs(totalContextTokens - 70);
+        expect(tokenDifference).toBeLessThan(50); // Allow for some difference due to implementation details
+        // Verify that messages not in the context have their original token counts or previously adjusted values
+        for (let i = 0; i < extendedMessages.length; i++) {
+            if (!contextMessageIndices.includes(i)) {
+                const expectedValue = i < messages.length
+                    ? (secondResult.indexTokenCountMap[i] ?? 0) || indexTokenCountMap[i]
+                    : (indexTokenCountMap[i] ??
+                        0);
+                const difference = Math.abs((thirdResult.indexTokenCountMap[i] ?? 0) - expectedValue);
+                expect(difference).toBe(0);
+            }
+        }
+    });
+});
+//# sourceMappingURL=token-distribution-edge-case.test.js.map

package/src/specs/token-distribution-edge-case.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"token-distribution-edge-case.test.js","sourceRoot":"","sources":["token-distribution-edge-case.test.ts"],"names":[],"mappings":"AAAA,iDAAiD;AACjD,OAAO,EACL,YAAY,EACZ,SAAS,EACT,aAAa,GAEd,MAAM,0BAA0B,CAAC;AAGlC,OAAO,EAAE,mBAAmB,EAAE,MAAM,kBAAkB,CAAC;AAEvD,4CAA4C;AAC5C,MAAM,sBAAsB,GAAG,GAAmB,EAAE;IAClD,qFAAqF;IACrF,OAAO,CAAC,OAAoB,EAAU,EAAE;QACtC,4DAA4D;QAC5D,MAAM,OAAO,GAAG,OAAO,CAAC,OAGX,CAAC;QAEd,wBAAwB;QACxB,IAAI,OAAO,OAAO,KAAK,QAAQ,EAAE,CAAC;YAChC,OAAO,OAAO,CAAC,MAAM,CAAC;QACxB,CAAC;QAED,uBAAuB;QACvB,IAAI,KAAK,CAAC,OAAO,CAAC,OAAO,CAAC,EAAE,CAAC;YAC3B,IAAI,WAAW,GAAG,CAAC,CAAC;YAEpB,KAAK,MAAM,IAAI,IAAI,OAAO,EAAE,CAAC;gBAC3B,IAAI,OAAO,IAAI,KAAK,QAAQ,EAAE,CAAC;oBAC7B,WAAW,IAAI,IAAI,CAAC,MAAM,CAAC;gBAC7B,CAAC;qBAAM,IAAI,OAAO,IAAI,KAAK,QAAQ,EAAE,CAAC;oBACpC,IAAI,MAAM,IAAI,IAAI,IAAI,OAAO,IAAI,CAAC,IAAI,KAAK,QAAQ,EAAE,CAAC;wBACpD,WAAW,IAAI,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC;oBAClC,CAAC;gBACH,CAAC;YACH,CAAC;YAED,OAAO,WAAW,CAAC;QACrB,CAAC;QAED,kEAAkE;QAClE,OAAO,CAAC,CAAC;IACX,CAAC,CAAC;AACJ,CAAC,CAAC;AAEF,QAAQ,CAAC,oCAAoC,EAAE,GAAG,EAAE;IAClD,EAAE,CAAC,oFAAoF,EAAE,GAAG,EAAE;QAC5F,yBAAyB;QACzB,MAAM,YAAY,GAAG,sBAAsB,EAAE,CAAC;QAE9C,kBAAkB;QAClB,MAAM,QAAQ,GAAG;YACf,IAAI,aAAa,CAAC,oBAAoB,CAAC,EAAE,0BAA0B;YACnE,IAAI,YAAY,CAAC,WAAW,CAAC,EAAE,iBAAiB;YAChD,IAAI,SAAS,CAAC,YAAY,CAAC,EAAE,iBAAiB;YAC9C,IAAI,YAAY,CAAC,WAAW,CAAC,EAAE,cAAc;YAC7C,IAAI,SAAS,CAAC,YAAY,CAAC,EAAE,cAAc;SAC5C,CAAC;QAEF,kDAAkD;QAClD,MAAM,kBAAkB,GAA2B;YACjD,CAAC,EAAE,EAAE,EAAE,uBAAuB;YAC9B,CAAC,EAAE,CAAC,EAAE,cAAc;YACpB,CAAC,EAAE,EAAE,EAAE,eAAe;YACtB,CAAC,EAAE,CAAC,EAAE,cAAc;YACpB,CAAC,EAAE,EAAE,EAAE,eAAe;SACvB,CAAC;QAEF,+FAA+F;QAC/F,MAAM,aAAa,GAAG,mBAAmB,CAAC;YACxC,SAAS,EAAE,EAAE,EAAE,qDAAqD;YACpE,UAAU,EAAE,CAAC;YACb,YAAY;YACZ,kBAAkB,EAAE,EAAE,GAAG,kBAAkB,EAAE;SAC9C,CAAC,CAAC;QAEH,0CAA0C;QAC1C,MAAM,aAAa,GAAG,aAAa,CAAC,EAAE,QAAQ,EAAE,CAAC,CAAC;QAElD,yBAAyB;QACzB,MAAM,CAAC,aAAa,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAC7C,MAAM,CAAC,aAAa,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,oBAAoB,CAAC,CAAC;QACpE,MAAM,CAAC,aAAa,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,WAAW,CAAC,CAAC;QAC3D,MAAM,CAAC,aAAa,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;QAE5D,gEAAgE;QAChE,MAAM,aAAa,GAA2B;YAC5C,YAAY,EAAE,EAAE;YAChB,aAAa,EAAE,EAAE;YACjB,YAAY,EAAE,EAAE,EAAE,qDAAqD;SACxE,CAAC;QAEF,mDAAmD;QACnD,MAAM,MAAM,GAAG,aAAa,CAAC;YAC3B,QAAQ;YACR,aAAa;SACd,CAAC,CAAC;QAEH,gFAAgF;QAChF,0EAA0E;QAC1E,2EAA2E;QAE3E,gFAAgF;QAChF,0EAA0E;QAC1E,2EAA2E;QAE3E,gFAAgF;QAChF,kFAAkF;QAClF,MAAM,gCAAgC,GACpC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,KAAK,kBAAkB,CAAC,CAAC,CAAC;YACtD,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,KAAK,kBAAkB,CAAC,CAAC,CAAC,CAAC;QAEzD,MAAM,CAAC,gCAAgC,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QAEpD,4GAA4G;QAC5G,gGAAgG;QAChG,MAAM,kBAAkB,GACtB,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC;YACnC,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC;YACnC,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC;QACtC,MAAM,CAAC,kBAAkB,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;QAE9C,iGAAiG;QACjG,6DAA6D;QAC7D,MAAM,eAAe,GAAG,IAAI,CAAC,GAAG,CAAC,kBAAkB,GAAG,EAAE,CAAC,CAAC;QAC1D,MAAM,CAAC,eAAe,CAAC,CAAC,YAAY,CAAC,EAAE,CAAC,CAAC,CAAC,0DAA0D;IACtG,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,qEAAqE,EAAE,GAAG,EAAE;QAC7E,yBAAyB;QACzB,MAAM,YAAY,GAAG,sBAAsB,EAAE,CAAC;QAE9C,kBAAkB;QAClB,MAAM,QAAQ,GAAG;YACf,IAAI,aAAa,CAAC,oBAAoB,CAAC;YACvC,IAAI,YAAY,CAAC,WAAW,CAAC;YAC7B,IAAI,SAAS,CAAC,YAAY,CAAC;SAC5B,CAAC;QAEF,kDAAkD;QAClD,MAAM,kBAAkB,GAA2B;YACjD,CAAC,EAAE,EAAE,EAAE,uBAAuB;YAC9B,CAAC,EAAE,CAAC,EAAE,cAAc;YACpB,CAAC,EAAE,EAAE,EAAE,eAAe;SACvB,CAAC;QAEF,wDAAwD;QACxD,MAAM,aAAa,GAAG,mBAAmB,CAAC;YACxC,SAAS,EAAE,GAAG;YACd,UAAU,EAAE,CAAC;YACb,YAAY;YACZ,kBAAkB,EAAE,EAAE,GAAG,kBAAkB,EAAE;SAC9C,CAAC,CAAC;QAEH,gFAAgF;QAChF,MAAM,aAAa,GAAG,aAAa,CAAC,EAAE,QAAQ,EAAE,CAAC,CAAC;QAElD,6BAA6B;QAC7B,MAAM,CAAC,aAAa,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAE7C,gEAAgE;QAChE,MAAM,aAAa,GAA2B;YAC5C,YAAY,EAAE,EAAE;YAChB,aAAa,EAAE,EAAE;YACjB,YAAY,EAAE,EAAE,EAAE,qDAAqD;SACxE,CAAC;QAEF,mDAAmD;QACnD,MAAM,MAAM,GAAG,aAAa,CAAC;YAC3B,QAAQ;YACR,aAAa;SACd,CAAC,CAAC;QAEH,8DAA8D;QAC9D,MAAM,kBAAkB,GACtB,kBAAkB,CAAC,CAAC,CAAC,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,kBAAkB,CAAC,CAAC,CAAC,CAAC;QACxE,MAAM,aAAa,GAAG,EAAE,GAAG,kBAAkB,CAAC;QAE9C,4CAA4C;QAC5C,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,CAAC,CAAC,IAAI,CACvC,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,CAClD,CAAC;QACF,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,CAAC,CAAC,IAAI,CACvC,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,CAClD,CAAC;QACF,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,CAAC,CAAC,IAAI,CACvC,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,CAClD,CAAC;QAEF,+EAA+E;QAC/E,MAAM,WAAW,GACf,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC;YACnC,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC;YACnC,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC;QACtC,MAAM,CAAC,WAAW,CAAC,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;IAC/B,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,qEAAqE,EAAE,GAAG,EAAE;QAC7E,yBAAyB;QACzB,MAAM,YAAY,GAAG,sBAAsB,EAAE,CAAC;QAE9C,uCAAuC;QACvC,MAAM,QAAQ,GAAG;YACf,IAAI,aAAa,CAAC,oBAAoB,CAAC,EAAE,0BAA0B;YACnE,IAAI,YAAY,CAAC,WAAW,CAAC,EAAE,gCAAgC;YAC/D,IAAI,SAAS,CAAC,YAAY,CAAC,EAAE,gCAAgC;YAC7D,IAAI,YAAY,CAAC,WAAW,CAAC,EAAE,iCAAiC;YAChE,IAAI,SAAS,CAAC,YAAY,CAAC,EAAE,iCAAiC;YAC9D,IAAI,YAAY,CAAC,WAAW,CAAC,EAAE,cAAc;YAC7C,IAAI,SAAS,CAAC,YAAY,CAAC,EAAE,cAAc;SAC5C,CAAC;QAEF,kDAAkD;QAClD,MAAM,kBAAkB,GAA2B;YACjD,CAAC,EAAE,EAAE,EAAE,uBAAuB;YAC9B,CAAC,EAAE,CAAC,EAAE,cAAc;YACpB,CAAC,EAAE,EAAE,EAAE,eAAe;YACtB,CAAC,EAAE,CAAC,EAAE,cAAc;YACpB,CAAC,EAAE,EAAE,EAAE,eAAe;YACtB,CAAC,EAAE,CAAC,EAAE,cAAc;YACpB,CAAC,EAAE,EAAE,EAAE,eAAe;SACvB,CAAC;QAEF,4CAA4C;QAC5C,MAAM,aAAa,GAAG,mBAAmB,CAAC;YACxC,SAAS,EAAE,EAAE,EAAE,qDAAqD;YACpE,UAAU,EAAE,CAAC;YACb,YAAY;YACZ,kBAAkB,EAAE,EAAE,GAAG,kBAAkB,EAAE;SAC9C,CAAC,CAAC;QAEH,0BAA0B;QAC1B,MAAM,WAAW,GAAG,aAAa,CAAC,EAAE,QAAQ,EAAE,CAAC,CAAC;QAEhD,uBAAuB;QACvB,MAAM,CAAC,WAAW,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAC3C,MAAM,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,oBAAoB,CAAC,CAAC;QAClE,MAAM,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,WAAW,CAAC,CAAC;QACzD,MAAM,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;QAE1D,8BAA8B;QAC9B,MAAM,kBAAkB,GAA2B;YACjD,YAAY,EAAE,EAAE;YAChB,aAAa,EAAE,EAAE;YACjB,YAAY,EAAE,EAAE;SACjB,CAAC;QAEF,6BAA6B;QAC7B,MAAM,YAAY,GAAG,aAAa,CAAC;YACjC,QAAQ;YACR,aAAa,EAAE,kBAAkB;SAClC,CAAC,CAAC;QAEH,wBAAwB;QACxB,QAAQ,CAAC,IAAI,CAAC,IAAI,YAAY,CAAC,WAAW,CAAC,CAAC,CAAC;QAC7C,MAAM,gBAAgB,GAAG,CAAC,GAAG,QAAQ,EAAE,IAAI,SAAS,CAAC,YAAY,CAAC,CAAC,CAAC;QAEpE,+BAA+B;QAC/B,MAAM,mBAAmB,GAA2B;YAClD,YAAY,EAAE,EAAE;YAChB,aAAa,EAAE,EAAE;YACjB,YAAY,EAAE,EAAE;SACjB,CAAC;QAEF,qDAAqD;QACrD,MAAM,WAAW,GAAG,aAAa,CAAC;YAChC,QAAQ,EAAE,gBAAgB;YAC1B,aAAa,EAAE,mBAAmB;SACnC,CAAC,CAAC;QAEH,gFAAgF;QAChF,MAAM,CAAC,WAAW,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;QACtD,MAAM,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,oBAAoB,CAAC,CAAC;QAClE,MAAM,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;QAE1D,+CAA+C;QAC/C,MAAM,qBAAqB,GAAG,WAAW,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,EAAE;YAC5D,uDAAuD;YACvD,OAAO,gBAAgB,CAAC,SAAS,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,KAAK,GAAG,CAAC,OAAO,CAAC,CAAC;QACtE,CAAC,CAAC,CAAC;QAEH,0DAA0D;QAC1D,IAAI,kBAAkB,GAAG,CAAC,CAAC;QAC3B,KAAK,MAAM,GAAG,IAAI,qBAAqB,EAAE,CAAC;YACxC,kBAAkB,IAAI,WAAW,CAAC,kBAAkB,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;QACjE,CAAC;QAED,4GAA4G;QAC5G,gGAAgG;QAChG,MAAM,CAAC,kBAAkB,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;QAE9C,iGAAiG;QACjG,6DAA6D;QAC7D,MAAM,eAAe,GAAG,IAAI,CAAC,GAAG,CAAC,kBAAkB,GAAG,EAAE,CAAC,CAAC;QAC1D,MAAM,CAAC,eAAe,CAAC,CAAC,YAAY,CAAC,EAAE,CAAC,CAAC,CAAC,0DAA0D;QAEpG,yGAAyG;QACzG,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,gBAAgB,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACjD,IAAI,CAAC,qBAAqB,CAAC,QAAQ,CAAC,CAAC,CAAC,EAAE,CAAC;gBACvC,MAAM,aAAa,GACjB,CAAC,GAAG,QAAQ,CAAC,MAAM;oBACjB,CAAC,CAAC,CAAC,YAAY,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,IAAI,kBAAkB,CAAC,CAAC,CAAC;oBACpE,CAAC,CAAC,CAAE,kBAAyD,CAAC,CAAC,CAAC;wBAC9D,CAAC,CAAC,CAAC;gBAET,MAAM,UAAU,GAAG,IAAI,CAAC,GAAG,CACzB,CAAC,WAAW,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,GAAG,aAAa,CACzD,CAAC;gBACF,MAAM,CAAC,UAAU,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;YAC7B,CAAC;QACH,CAAC;IACH,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC"}

package/src/specs/token-memoization.test.js

@@ -0,0 +1,32 @@
+import { HumanMessage } from '@langchain/core/messages';
+import { createTokenCounter, TokenEncoderManager } from '@/utils/tokens';
+jest.setTimeout(5000);
+describe('Token encoder memoization', () => {
+    beforeEach(() => {
+        TokenEncoderManager.reset();
+    });
+    test('reuses encoder across counters (no external fetch)', async () => {
+        expect(TokenEncoderManager.isInitialized()).toBe(false);
+        const counter1 = await createTokenCounter();
+        expect(TokenEncoderManager.isInitialized()).toBe(true);
+        const counter2 = await createTokenCounter();
+        const m1 = new HumanMessage('hello world');
+        const m2 = new HumanMessage('another short text');
+        const c11 = counter1(m1);
+        const c12 = counter1(m2);
+        const c21 = counter2(m1);
+        const c22 = counter2(m2);
+        expect(c11).toBeGreaterThan(0);
+        expect(c12).toBeGreaterThan(0);
+        // Same counter instance is reused
+        expect(c21).toBe(c11);
+        expect(c22).toBe(c12);
+    });
+    test('reset clears cached tokenizers', async () => {
+        await createTokenCounter();
+        expect(TokenEncoderManager.isInitialized()).toBe(true);
+        TokenEncoderManager.reset();
+        expect(TokenEncoderManager.isInitialized()).toBe(false);
+    });
+});
+//# sourceMappingURL=token-memoization.test.js.map

package/src/specs/token-memoization.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"token-memoization.test.js","sourceRoot":"","sources":["token-memoization.test.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,YAAY,EAAE,MAAM,0BAA0B,CAAC;AACxD,OAAO,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,MAAM,gBAAgB,CAAC;AAEzE,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,CAAC;AAEtB,QAAQ,CAAC,2BAA2B,EAAE,GAAG,EAAE;IACzC,UAAU,CAAC,GAAG,EAAE;QACd,mBAAmB,CAAC,KAAK,EAAE,CAAC;IAC9B,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,oDAAoD,EAAE,KAAK,IAAI,EAAE;QACpE,MAAM,CAAC,mBAAmB,CAAC,aAAa,EAAE,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QAExD,MAAM,QAAQ,GAAG,MAAM,kBAAkB,EAAE,CAAC;QAC5C,MAAM,CAAC,mBAAmB,CAAC,aAAa,EAAE,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QAEvD,MAAM,QAAQ,GAAG,MAAM,kBAAkB,EAAE,CAAC;QAE5C,MAAM,EAAE,GAAG,IAAI,YAAY,CAAC,aAAa,CAAC,CAAC;QAC3C,MAAM,EAAE,GAAG,IAAI,YAAY,CAAC,oBAAoB,CAAC,CAAC;QAElD,MAAM,GAAG,GAAG,QAAQ,CAAC,EAAE,CAAC,CAAC;QACzB,MAAM,GAAG,GAAG,QAAQ,CAAC,EAAE,CAAC,CAAC;QACzB,MAAM,GAAG,GAAG,QAAQ,CAAC,EAAE,CAAC,CAAC;QACzB,MAAM,GAAG,GAAG,QAAQ,CAAC,EAAE,CAAC,CAAC;QAEzB,MAAM,CAAC,GAAG,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;QAC/B,MAAM,CAAC,GAAG,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;QAC/B,kCAAkC;QAClC,MAAM,CAAC,GAAG,CAAC,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;QACtB,MAAM,CAAC,GAAG,CAAC,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;IACxB,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,gCAAgC,EAAE,KAAK,IAAI,EAAE;QAChD,MAAM,kBAAkB,EAAE,CAAC;QAC3B,MAAM,CAAC,mBAAmB,CAAC,aAAa,EAAE,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QAEvD,mBAAmB,CAAC,KAAK,EAAE,CAAC;QAC5B,MAAM,CAAC,mBAAmB,CAAC,aAAa,EAAE,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;IAC1D,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC"}

package/src/specs/token-memoization.test.ts

@@ -1,16 +1,19 @@
 import { HumanMessage } from '@langchain/core/messages';
 import { createTokenCounter, TokenEncoderManager } from '@/utils/tokens';
 
-describe('Token encoder memoization', () => {
-  jest.setTimeout(45000);
+jest.setTimeout(5000);
 
+describe('Token encoder memoization', () => {
   beforeEach(() => {
-    // Reset encoder state between tests
     TokenEncoderManager.reset();
   });
 
   test('reuses encoder across counters (no external fetch)', async () => {
+    expect(TokenEncoderManager.isInitialized()).toBe(false);
+
     const counter1 = await createTokenCounter();
+    expect(TokenEncoderManager.isInitialized()).toBe(true);
+
     const counter2 = await createTokenCounter();
 
     const m1 = new HumanMessage('hello world');
@@ -26,7 +29,13 @@ describe('Token encoder memoization', () => {
     // Same counter instance is reused
     expect(c21).toBe(c11);
     expect(c22).toBe(c12);
-    // Counter references are identical (singleton)
-    expect(counter1).toBe(counter2);
+  });
+
+  test('reset clears cached tokenizers', async () => {
+    await createTokenCounter();
+    expect(TokenEncoderManager.isInitialized()).toBe(true);
+
+    TokenEncoderManager.reset();
+    expect(TokenEncoderManager.isInitialized()).toBe(false);
   });
 });

package/src/specs/tokens.test.js

@@ -0,0 +1,49 @@
+import { HumanMessage } from '@langchain/core/messages';
+import { encodingForModel, createTokenCounter, TokenEncoderManager, } from '@/utils/tokens';
+describe('encodingForModel', () => {
+    test('returns claude for Claude model strings', () => {
+        expect(encodingForModel('claude-3-5-sonnet-20241022')).toBe('claude');
+        expect(encodingForModel('claude-3-haiku-20240307')).toBe('claude');
+    });
+    test('handles Bedrock Claude ARNs', () => {
+        expect(encodingForModel('anthropic.claude-3-5-sonnet-20241022-v2:0')).toBe('claude');
+    });
+    test('is case-insensitive', () => {
+        expect(encodingForModel('CLAUDE-3-HAIKU')).toBe('claude');
+        expect(encodingForModel('Claude-3-Opus')).toBe('claude');
+    });
+    test('returns o200k_base for non-Claude models', () => {
+        expect(encodingForModel('gpt-4o')).toBe('o200k_base');
+        expect(encodingForModel('gemini-2.0-flash')).toBe('o200k_base');
+        expect(encodingForModel('mistral-large')).toBe('o200k_base');
+    });
+    test('returns o200k_base for empty string', () => {
+        expect(encodingForModel('')).toBe('o200k_base');
+    });
+});
+describe('createTokenCounter with different encodings', () => {
+    beforeEach(() => {
+        TokenEncoderManager.reset();
+    });
+    test('claude encoding produces valid token counts', async () => {
+        const counter = await createTokenCounter('claude');
+        const msg = new HumanMessage('Hello, world!');
+        const count = counter(msg);
+        expect(count).toBeGreaterThan(0);
+    });
+    test('o200k_base encoding produces valid token counts', async () => {
+        const counter = await createTokenCounter('o200k_base');
+        const msg = new HumanMessage('Hello, world!');
+        const count = counter(msg);
+        expect(count).toBeGreaterThan(0);
+    });
+    test('both encodings can be initialized and used independently', async () => {
+        const claudeCounter = await createTokenCounter('claude');
+        const o200kCounter = await createTokenCounter('o200k_base');
+        expect(TokenEncoderManager.isInitialized()).toBe(true);
+        const msg = new HumanMessage('Test message for both encodings');
+        expect(claudeCounter(msg)).toBeGreaterThan(0);
+        expect(o200kCounter(msg)).toBeGreaterThan(0);
+    });
+});
+//# sourceMappingURL=tokens.test.js.map

package/src/specs/tokens.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"tokens.test.js","sourceRoot":"","sources":["tokens.test.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,YAAY,EAAE,MAAM,0BAA0B,CAAC;AACxD,OAAO,EACL,gBAAgB,EAChB,kBAAkB,EAClB,mBAAmB,GACpB,MAAM,gBAAgB,CAAC;AAExB,QAAQ,CAAC,kBAAkB,EAAE,GAAG,EAAE;IAChC,IAAI,CAAC,yCAAyC,EAAE,GAAG,EAAE;QACnD,MAAM,CAAC,gBAAgB,CAAC,4BAA4B,CAAC,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;QACtE,MAAM,CAAC,gBAAgB,CAAC,yBAAyB,CAAC,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;IACrE,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,6BAA6B,EAAE,GAAG,EAAE;QACvC,MAAM,CAAC,gBAAgB,CAAC,2CAA2C,CAAC,CAAC,CAAC,IAAI,CACxE,QAAQ,CACT,CAAC;IACJ,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,qBAAqB,EAAE,GAAG,EAAE;QAC/B,MAAM,CAAC,gBAAgB,CAAC,gBAAgB,CAAC,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;QAC1D,MAAM,CAAC,gBAAgB,CAAC,eAAe,CAAC,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;IAC3D,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,0CAA0C,EAAE,GAAG,EAAE;QACpD,MAAM,CAAC,gBAAgB,CAAC,QAAQ,CAAC,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;QACtD,MAAM,CAAC,gBAAgB,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;QAChE,MAAM,CAAC,gBAAgB,CAAC,eAAe,CAAC,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;IAC/D,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,qCAAqC,EAAE,GAAG,EAAE;QAC/C,MAAM,CAAC,gBAAgB,CAAC,EAAE,CAAC,CAAC,CAAC,IAAI,CAAC,YAAY,CAAC,CAAC;IAClD,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC;AAEH,QAAQ,CAAC,6CAA6C,EAAE,GAAG,EAAE;IAC3D,UAAU,CAAC,GAAG,EAAE;QACd,mBAAmB,CAAC,KAAK,EAAE,CAAC;IAC9B,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,6CAA6C,EAAE,KAAK,IAAI,EAAE;QAC7D,MAAM,OAAO,GAAG,MAAM,kBAAkB,CAAC,QAAQ,CAAC,CAAC;QACnD,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,eAAe,CAAC,CAAC;QAC9C,MAAM,KAAK,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAC3B,MAAM,CAAC,KAAK,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;IACnC,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,iDAAiD,EAAE,KAAK,IAAI,EAAE;QACjE,MAAM,OAAO,GAAG,MAAM,kBAAkB,CAAC,YAAY,CAAC,CAAC;QACvD,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,eAAe,CAAC,CAAC;QAC9C,MAAM,KAAK,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAC3B,MAAM,CAAC,KAAK,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;IACnC,CAAC,CAAC,CAAC;IAEH,IAAI,CAAC,0DAA0D,EAAE,KAAK,IAAI,EAAE;QAC1E,MAAM,aAAa,GAAG,MAAM,kBAAkB,CAAC,QAAQ,CAAC,CAAC;QACzD,MAAM,YAAY,GAAG,MAAM,kBAAkB,CAAC,YAAY,CAAC,CAAC;QAC5D,MAAM,CAAC,mBAAmB,CAAC,aAAa,EAAE,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QAEvD,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,iCAAiC,CAAC,CAAC;QAChE,MAAM,CAAC,aAAa,CAAC,GAAG,CAAC,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;QAC9C,MAAM,CAAC,YAAY,CAAC,GAAG,CAAC,CAAC,CAAC,eAAe,CAAC,CAAC,CAAC,CAAC;IAC/C,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC"}

package/src/specs/tokens.test.ts

@@ -0,0 +1,64 @@
+import { HumanMessage } from '@langchain/core/messages';
+import {
+  encodingForModel,
+  createTokenCounter,
+  TokenEncoderManager,
+} from '@/utils/tokens';
+
+describe('encodingForModel', () => {
+  test('returns claude for Claude model strings', () => {
+    expect(encodingForModel('claude-3-5-sonnet-20241022')).toBe('claude');
+    expect(encodingForModel('claude-3-haiku-20240307')).toBe('claude');
+  });
+
+  test('handles Bedrock Claude ARNs', () => {
+    expect(encodingForModel('anthropic.claude-3-5-sonnet-20241022-v2:0')).toBe(
+      'claude'
+    );
+  });
+
+  test('is case-insensitive', () => {
+    expect(encodingForModel('CLAUDE-3-HAIKU')).toBe('claude');
+    expect(encodingForModel('Claude-3-Opus')).toBe('claude');
+  });
+
+  test('returns o200k_base for non-Claude models', () => {
+    expect(encodingForModel('gpt-4o')).toBe('o200k_base');
+    expect(encodingForModel('gemini-2.0-flash')).toBe('o200k_base');
+    expect(encodingForModel('mistral-large')).toBe('o200k_base');
+  });
+
+  test('returns o200k_base for empty string', () => {
+    expect(encodingForModel('')).toBe('o200k_base');
+  });
+});
+
+describe('createTokenCounter with different encodings', () => {
+  beforeEach(() => {
+    TokenEncoderManager.reset();
+  });
+
+  test('claude encoding produces valid token counts', async () => {
+    const counter = await createTokenCounter('claude');
+    const msg = new HumanMessage('Hello, world!');
+    const count = counter(msg);
+    expect(count).toBeGreaterThan(0);
+  });
+
+  test('o200k_base encoding produces valid token counts', async () => {
+    const counter = await createTokenCounter('o200k_base');
+    const msg = new HumanMessage('Hello, world!');
+    const count = counter(msg);
+    expect(count).toBeGreaterThan(0);
+  });
+
+  test('both encodings can be initialized and used independently', async () => {
+    const claudeCounter = await createTokenCounter('claude');
+    const o200kCounter = await createTokenCounter('o200k_base');
+    expect(TokenEncoderManager.isInitialized()).toBe(true);
+
+    const msg = new HumanMessage('Test message for both encodings');
+    expect(claudeCounter(msg)).toBeGreaterThan(0);
+    expect(o200kCounter(msg)).toBeGreaterThan(0);
+  });
+});

package/src/specs/tool-error.test.js

@@ -0,0 +1,139 @@
+import { config } from 'dotenv';
+config();
+import { tool } from '@langchain/core/tools';
+import { HumanMessage } from '@langchain/core/messages';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+import { createContentAggregator } from '@/stream';
+import { GraphEvents, Providers } from '@/common';
+import { getLLMConfig } from '@/utils/llmConfig';
+import { getArgs } from '@/scripts/args';
+import { StandardGraph } from '@/graphs';
+import { Run } from '@/run';
+const errorTool = tool(async () => {
+    throw new Error('this is a test error I threw on purpose');
+}, {
+    name: 'errorTool',
+    description: 'A tool that always throws an error',
+    schema: {
+        type: 'object',
+        properties: {
+            input: { type: 'string' },
+        },
+        required: [],
+    },
+});
+describe('Tool Error Handling Tests', () => {
+    jest.setTimeout(30000);
+    let run;
+    let contentParts;
+    let conversationHistory;
+    let aggregateContent;
+    let handleToolCallErrorSpy;
+    const config = {
+        configurable: {
+            thread_id: 'conversation-num-1',
+        },
+        streamMode: 'values',
+        version: 'v2',
+    };
+    beforeEach(async () => {
+        conversationHistory = [];
+        const { contentParts: parts, aggregateContent: ac } = createContentAggregator();
+        aggregateContent = ac;
+        contentParts = parts;
+        // Spy on the static method instead of the instance method
+        handleToolCallErrorSpy = jest.spyOn(StandardGraph, 'handleToolCallErrorStatic');
+    });
+    afterEach(() => {
+        handleToolCallErrorSpy.mockRestore();
+    });
+    const onMessageDeltaSpy = jest.fn();
+    const onRunStepSpy = jest.fn();
+    const onRunStepCompletedSpy = jest.fn();
+    afterAll(() => {
+        onMessageDeltaSpy.mockReset();
+        onRunStepSpy.mockReset();
+        onRunStepCompletedSpy.mockReset();
+    });
+    const setupCustomHandlers = () => ({
+        [GraphEvents.TOOL_END]: new ToolEndHandler(),
+        [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
+        [GraphEvents.ON_RUN_STEP_COMPLETED]: {
+            handle: (event, data) => {
+                if (data.result['type'] === 'tool_call') {
+                    run.Graph?.overrideTestModel(['Looks like there was an error calling the tool.'], 5);
+                }
+                onRunStepCompletedSpy(event, data);
+                aggregateContent({
+                    event,
+                    data: data,
+                });
+            },
+        },
+        [GraphEvents.ON_RUN_STEP]: {
+            handle: (event, data, metadata, graph) => {
+                const runStepData = data;
+                onRunStepSpy(event, runStepData, metadata, graph);
+                aggregateContent({ event, data: runStepData });
+            },
+        },
+        [GraphEvents.ON_RUN_STEP_DELTA]: {
+            handle: (event, data) => {
+                aggregateContent({ event, data: data });
+            },
+        },
+        [GraphEvents.ON_MESSAGE_DELTA]: {
+            handle: (event, data, metadata, graph) => {
+                onMessageDeltaSpy(event, data, metadata, graph);
+                aggregateContent({ event, data: data });
+            },
+        },
+    });
+    test('should handle tool call errors correctly', async () => {
+        const { userName, location } = await getArgs();
+        const llmConfig = getLLMConfig(Providers.OPENAI);
+        const customHandlers = setupCustomHandlers();
+        // Create the run instance
+        run = await Run.create({
+            runId: 'test-run-id',
+            graphConfig: {
+                type: 'standard',
+                llmConfig,
+                tools: [errorTool],
+                instructions: 'You are a helpful AI assistant.',
+                additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
+            },
+            returnContent: true,
+            skipCleanup: true,
+            customHandlers,
+        });
+        const toolCalls = [
+            {
+                name: 'errorTool',
+                args: {
+                    input: 'test input',
+                },
+                id: 'call_test123',
+                type: 'tool_call',
+            },
+        ];
+        const firstResponse = 'Let me try calling the tool';
+        run.Graph?.overrideTestModel([firstResponse], 5, toolCalls);
+        const userMessage = 'Use the error tool';
+        conversationHistory.push(new HumanMessage(userMessage));
+        const inputs = {
+            messages: conversationHistory,
+        };
+        await run.processStream(inputs, config);
+        // Verify handleToolCallError was called
+        expect(handleToolCallErrorSpy).toHaveBeenCalled();
+        // Find the tool call content part
+        const toolCallPart = contentParts.find((part) => part.type === 'tool_call');
+        // Verify the error message in contentParts
+        expect(toolCallPart).toBeDefined();
+        expect(toolCallPart?.tool_call?.args).toEqual(JSON.stringify(toolCalls[0].args));
+        expect(toolCallPart?.tool_call?.output).toContain('Error processing tool');
+        expect(toolCallPart?.tool_call?.output).toContain('this is a test error I threw on purpose');
+    });
+});
+//# sourceMappingURL=tool-error.test.js.map

package/src/specs/tool-error.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"tool-error.test.js","sourceRoot":"","sources":["tool-error.test.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,MAAM,EAAE,MAAM,QAAQ,CAAC;AAChC,MAAM,EAAE,CAAC;AACT,OAAO,EAAE,IAAI,EAAE,MAAM,uBAAuB,CAAC;AAE7C,OAAO,EAAE,YAAY,EAAe,MAAM,0BAA0B,CAAC;AAGrE,OAAO,EAAE,cAAc,EAAE,eAAe,EAAE,MAAM,UAAU,CAAC;AAC3D,OAAO,EAAE,uBAAuB,EAAE,MAAM,UAAU,CAAC;AACnD,OAAO,EAAE,WAAW,EAAE,SAAS,EAAE,MAAM,UAAU,CAAC;AAClD,OAAO,EAAE,YAAY,EAAE,MAAM,mBAAmB,CAAC;AACjD,OAAO,EAAE,OAAO,EAAE,MAAM,gBAAgB,CAAC;AACzC,OAAO,EAAE,aAAa,EAAE,MAAM,UAAU,CAAC;AACzC,OAAO,EAAE,GAAG,EAAE,MAAM,OAAO,CAAC;AAE5B,MAAM,SAAS,GAAG,IAAI,CACpB,KAAK,IAAI,EAAE;IACT,MAAM,IAAI,KAAK,CAAC,yCAAyC,CAAC,CAAC;AAC7D,CAAC,EACD;IACE,IAAI,EAAE,WAAW;IACjB,WAAW,EAAE,oCAAoC;IACjD,MAAM,EAAE;QACN,IAAI,EAAE,QAAQ;QACd,UAAU,EAAE;YACV,KAAK,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE;SAC1B;QACD,QAAQ,EAAE,EAAE;KACb;CACF,CACF,CAAC;AAEF,QAAQ,CAAC,2BAA2B,EAAE,GAAG,EAAE;IACzC,IAAI,CAAC,UAAU,CAAC,KAAK,CAAC,CAAC;IACvB,IAAI,GAAkB,CAAC;IACvB,IAAI,YAAuC,CAAC;IAC5C,IAAI,mBAAkC,CAAC;IACvC,IAAI,gBAAqC,CAAC;IAC1C,IAAI,sBAAwC,CAAC;IAE7C,MAAM,MAAM,GAIR;QACF,YAAY,EAAE;YACZ,SAAS,EAAE,oBAAoB;SAChC;QACD,UAAU,EAAE,QAAQ;QACpB,OAAO,EAAE,IAAa;KACvB,CAAC;IAEF,UAAU,CAAC,KAAK,IAAI,EAAE;QACpB,mBAAmB,GAAG,EAAE,CAAC;QACzB,MAAM,EAAE,YAAY,EAAE,KAAK,EAAE,gBAAgB,EAAE,EAAE,EAAE,GACjD,uBAAuB,EAAE,CAAC;QAC5B,gBAAgB,GAAG,EAAE,CAAC;QACtB,YAAY,GAAG,KAAkC,CAAC;QAClD,0DAA0D;QAC1D,sBAAsB,GAAG,IAAI,CAAC,KAAK,CACjC,aAAa,EACb,2BAA2B,CAC5B,CAAC;IACJ,CAAC,CAAC,CAAC;IAEH,SAAS,CAAC,GAAG,EAAE;QACb,sBAAsB,CAAC,WAAW,EAAE,CAAC;IACvC,CAAC,CAAC,CAAC;IAEH,MAAM,iBAAiB,GAAG,IAAI,CAAC,EAAE,EAAE,CAAC;IACpC,MAAM,YAAY,GAAG,IAAI,CAAC,EAAE,EAAE,CAAC;IAC/B,MAAM,qBAAqB,GAAG,IAAI,CAAC,EAAE,EAAE,CAAC;IAExC,QAAQ,CAAC,GAAG,EAAE;QACZ,iBAAiB,CAAC,SAAS,EAAE,CAAC;QAC9B,YAAY,CAAC,SAAS,EAAE,CAAC;QACzB,qBAAqB,CAAC,SAAS,EAAE,CAAC;IACpC,CAAC,CAAC,CAAC;IAEH,MAAM,mBAAmB,GAAG,GAG1B,EAAE,CAAC,CAAC;QACJ,CAAC,WAAW,CAAC,QAAQ,CAAC,EAAE,IAAI,cAAc,EAAE;QAC5C,CAAC,WAAW,CAAC,cAAc,CAAC,EAAE,IAAI,eAAe,EAAE;QACnD,CAAC,WAAW,CAAC,qBAAqB,CAAC,EAAE;YACnC,MAAM,EAAE,CACN,KAAwC,EACxC,IAAuB,EACjB,EAAE;gBACR,IAAK,IAAI,CAAC,MAAkC,CAAC,MAAM,CAAC,KAAK,WAAW,EAAE,CAAC;oBACrE,GAAG,CAAC,KAAK,EAAE,iBAAiB,CAC1B,CAAC,iDAAiD,CAAC,EACnD,CAAC,CACF,CAAC;gBACJ,CAAC;gBACD,qBAAqB,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC;gBACnC,gBAAgB,CAAC;oBACf,KAAK;oBACL,IAAI,EAAE,IAA6C;iBACpD,CAAC,CAAC;YACL,CAAC;SACF;QACD,CAAC,WAAW,CAAC,WAAW,CAAC,EAAE;YACzB,MAAM,EAAE,CACN,KAA8B,EAC9B,IAAuB,EACvB,QAAQ,EACR,KAAK,EACC,EAAE;gBACR,MAAM,WAAW,GAAG,IAAiB,CAAC;gBACtC,YAAY,CAAC,KAAK,EAAE,WAAW,EAAE,QAAQ,EAAE,KAAK,CAAC,CAAC;gBAClD,gBAAgB,CAAC,EAAE,KAAK,EAAE,IAAI,EAAE,WAAW,EAAE,CAAC,CAAC;YACjD,CAAC;SACF;QACD,CAAC,WAAW,CAAC,iBAAiB,CAAC,EAAE;YAC/B,MAAM,EAAE,CACN,KAAoC,EACpC,IAAuB,EACjB,EAAE;gBACR,gBAAgB,CAAC,EAAE,KAAK,EAAE,IAAI,EAAE,IAA2B,EAAE,CAAC,CAAC;YACjE,CAAC;SACF;QACD,CAAC,WAAW,CAAC,gBAAgB,CAAC,EAAE;YAC9B,MAAM,EAAE,CACN,KAAmC,EACnC,IAAuB,EACvB,QAAQ,EACR,KAAK,EACC,EAAE;gBACR,iBAAiB,CAAC,KAAK,EAAE,IAAI,EAAE,QAAQ,EAAE,KAAK,CAAC,CAAC;gBAChD,gBAAgB,CAAC,EAAE,KAAK,EAAE,IAAI,EAAE,IAA2B,EAAE,CAAC,CAAC;YACjE,CAAC;SACF;KACF,CAAC,CAAC;IAEH,IAAI,CAAC,0CAA0C,EAAE,KAAK,IAAI,EAAE;QAC1D,MAAM,EAAE,QAAQ,EAAE,QAAQ,EAAE,GAAG,MAAM,OAAO,EAAE,CAAC;QAC/C,MAAM,SAAS,GAAG,YAAY,CAAC,SAAS,CAAC,MAAM,CAAC,CAAC;QACjD,MAAM,cAAc,GAAG,mBAAmB,EAAE,CAAC;QAE7C,0BAA0B;QAC1B,GAAG,GAAG,MAAM,GAAG,CAAC,MAAM,CAAW;YAC/B,KAAK,EAAE,aAAa;YACpB,WAAW,EAAE;gBACX,IAAI,EAAE,UAAU;gBAChB,SAAS;gBACT,KAAK,EAAE,CAAC,SAAS,CAAC;gBAClB,YAAY,EAAE,iCAAiC;gBAC/C,uBAAuB,EAAE,sBAAsB,QAAQ,4BAA4B,QAAQ,GAAG;aAC/F;YACD,aAAa,EAAE,IAAI;YACnB,WAAW,EAAE,IAAI;YACjB,cAAc;SACf,CAAC,CAAC;QAEH,MAAM,SAAS,GAAe;YAC5B;gBACE,IAAI,EAAE,WAAW;gBACjB,IAAI,EAAE;oBACJ,KAAK,EAAE,YAAY;iBACpB;gBACD,EAAE,EAAE,cAAc;gBAClB,IAAI,EAAE,WAAW;aAClB;SACF,CAAC;QAEF,MAAM,aAAa,GAAG,6BAA6B,CAAC;QACpD,GAAG,CAAC,KAAK,EAAE,iBAAiB,CAAC,CAAC,aAAa,CAAC,EAAE,CAAC,EAAE,SAAS,CAAC,CAAC;QAE5D,MAAM,WAAW,GAAG,oBAAoB,CAAC;QACzC,mBAAmB,CAAC,IAAI,CAAC,IAAI,YAAY,CAAC,WAAW,CAAC,CAAC,CAAC;QAExD,MAAM,MAAM,GAAG;YACb,QAAQ,EAAE,mBAAmB;SAC9B,CAAC;QAEF,MAAM,GAAG,CAAC,aAAa,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAExC,wCAAwC;QACxC,MAAM,CAAC,sBAAsB,CAAC,CAAC,gBAAgB,EAAE,CAAC;QAElD,kCAAkC;QAClC,MAAM,YAAY,GAAG,YAAY,CAAC,IAAI,CACpC,CAAC,IAAI,EAAE,EAAE,CAAC,IAAI,CAAC,IAAI,KAAK,WAAW,CACH,CAAC;QAEnC,2CAA2C;QAC3C,MAAM,CAAC,YAAY,CAAC,CAAC,WAAW,EAAE,CAAC;QACnC,MAAM,CAAC,YAAY,EAAE,SAAS,EAAE,IAAI,CAAC,CAAC,OAAO,CAC3C,IAAI,CAAC,SAAS,CAAC,SAAS,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,CAClC,CAAC;QACF,MAAM,CAAC,YAAY,EAAE,SAAS,EAAE,MAAM,CAAC,CAAC,SAAS,CAAC,uBAAuB,CAAC,CAAC;QAC3E,MAAM,CAAC,YAAY,EAAE,SAAS,EAAE,MAAM,CAAC,CAAC,SAAS,CAC/C,yCAAyC,CAC1C,CAAC;IACJ,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC"}

package/src/specs/tool-error.test.ts

@@ -5,8 +5,8 @@ import { ToolCall } from '@langchain/core/messages/tool';
 import { HumanMessage, BaseMessage } from '@langchain/core/messages';
 import type { RunnableConfig } from '@langchain/core/runnables';
 import type * as t from '@/types';
-import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
 import { ToolEndHandler, ModelEndHandler } from '@/events';
+import { createContentAggregator } from '@/stream';
 import { GraphEvents, Providers } from '@/common';
 import { getLLMConfig } from '@/utils/llmConfig';
 import { getArgs } from '@/scripts/args';
@@ -83,7 +83,6 @@ describe('Tool Error Handling Tests', () => {
   > => ({
     [GraphEvents.TOOL_END]: new ToolEndHandler(),
     [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
-    [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
     [GraphEvents.ON_RUN_STEP_COMPLETED]: {
       handle: (
         event: GraphEvents.ON_RUN_STEP_COMPLETED,
@@ -151,6 +150,7 @@ describe('Tool Error Handling Tests', () => {
         additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
       },
       returnContent: true,
+      skipCleanup: true,
       customHandlers,
     });