@illuma-ai/agents 1.0.89 → 1.0.93

package/src/specs/prune.test.js

@@ -0,0 +1,733 @@
+// src/specs/prune.test.ts
+import { config } from 'dotenv';
+config();
+import { HumanMessage, AIMessage, SystemMessage, BaseMessage, ToolMessage, } from '@langchain/core/messages';
+import { createPruneMessages } from '@/messages/prune';
+import { getLLMConfig } from '@/utils/llmConfig';
+import { Providers } from '@/common';
+import { Run } from '@/run';
+// 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;
+    };
+};
+// Since the internal functions in prune.ts are not exported, we'll reimplement them here for testing
+// This is based on the implementation in src/messages/prune.ts
+function calculateTotalTokens(usage) {
+    const baseInputTokens = Number(usage.input_tokens) || 0;
+    const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;
+    const cacheRead = Number(usage.input_token_details?.cache_read) || 0;
+    const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;
+    const totalOutputTokens = Number(usage.output_tokens) || 0;
+    return {
+        input_tokens: totalInputTokens,
+        output_tokens: totalOutputTokens,
+        total_tokens: totalInputTokens + totalOutputTokens,
+    };
+}
+function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, indexTokenCountMap, startType, }) {
+    // Every reply is primed with <|start|>assistant<|message|>, so we
+    // start with 3 tokens for the label after all messages have been counted.
+    let summaryIndex = -1;
+    let currentTokenCount = 3;
+    const instructions = _messages[0]?.getType() === 'system' ? _messages[0] : undefined;
+    const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;
+    let remainingContextTokens = maxContextTokens - instructionsTokenCount;
+    const messages = [..._messages];
+    const context = [];
+    if (currentTokenCount < remainingContextTokens) {
+        let currentIndex = messages.length;
+        while (messages.length > 0 &&
+            currentTokenCount < remainingContextTokens &&
+            currentIndex > 1) {
+            currentIndex--;
+            if (messages.length === 1 && instructions) {
+                break;
+            }
+            const poppedMessage = messages.pop();
+            if (!poppedMessage)
+                continue;
+            const tokenCount = indexTokenCountMap[currentIndex] || 0;
+            if (currentTokenCount + tokenCount <= remainingContextTokens) {
+                context.push(poppedMessage);
+                currentTokenCount += tokenCount;
+            }
+            else {
+                messages.push(poppedMessage);
+                break;
+            }
+        }
+        // If startType is specified, discard messages until we find one of the required type
+        if (startType != null && startType && context.length > 0) {
+            const requiredTypeIndex = context.findIndex((msg) => msg.getType() === startType);
+            if (requiredTypeIndex > 0) {
+                // If we found a message of the required type, discard all messages before it
+                const remainingMessages = context.slice(requiredTypeIndex);
+                context.length = 0; // Clear the array
+                context.push(...remainingMessages);
+            }
+        }
+    }
+    if (instructions && _messages.length > 0) {
+        context.push(_messages[0]);
+        messages.shift();
+    }
+    const prunedMemory = messages;
+    summaryIndex = prunedMemory.length - 1;
+    remainingContextTokens -= currentTokenCount;
+    return {
+        summaryIndex,
+        remainingContextTokens,
+        context: context.reverse(),
+        messagesToRefine: prunedMemory,
+    };
+}
+function checkValidNumber(value) {
+    return typeof value === 'number' && !isNaN(value) && value > 0;
+}
+describe('Prune Messages Tests', () => {
+    jest.setTimeout(30000);
+    describe('calculateTotalTokens', () => {
+        it('should calculate total tokens correctly with all fields present', () => {
+            const usage = {
+                input_tokens: 100,
+                output_tokens: 50,
+                input_token_details: {
+                    cache_creation: 10,
+                    cache_read: 5,
+                },
+            };
+            const result = calculateTotalTokens(usage);
+            expect(result.input_tokens).toBe(115); // 100 + 10 + 5
+            expect(result.output_tokens).toBe(50);
+            expect(result.total_tokens).toBe(165); // 115 + 50
+        });
+        it('should handle missing fields gracefully', () => {
+            const usage = {
+                input_tokens: 100,
+                output_tokens: 50,
+            };
+            const result = calculateTotalTokens(usage);
+            expect(result.input_tokens).toBe(100);
+            expect(result.output_tokens).toBe(50);
+            expect(result.total_tokens).toBe(150);
+        });
+        it('should handle empty usage object', () => {
+            const usage = {};
+            const result = calculateTotalTokens(usage);
+            expect(result.input_tokens).toBe(0);
+            expect(result.output_tokens).toBe(0);
+            expect(result.total_tokens).toBe(0);
+        });
+    });
+    describe('getMessagesWithinTokenLimit', () => {
+        it('should include all messages when under token limit', () => {
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Hello'),
+                new AIMessage('Hi there'),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // "System instruction"
+                1: 5, // "Hello"
+                2: 8, // "Hi there"
+            };
+            const result = getMessagesWithinTokenLimit({
+                messages,
+                maxContextTokens: 100,
+                indexTokenCountMap,
+            });
+            expect(result.context.length).toBe(3);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[0].getType()).toBe('system'); // System message
+            expect(result.remainingContextTokens).toBe(100 - 17 - 5 - 8 - 3); // -3 for the assistant label tokens
+            expect(result.messagesToRefine.length).toBe(0);
+        });
+        it('should prune oldest messages when over token limit', () => {
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Message 1'),
+                new AIMessage('Response 1'),
+                new HumanMessage('Message 2'),
+                new AIMessage('Response 2'),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // "System instruction"
+                1: 9, // "Message 1"
+                2: 10, // "Response 1"
+                3: 9, // "Message 2"
+                4: 10, // "Response 2"
+            };
+            // Set a limit that can only fit the system message and the last two messages
+            const result = getMessagesWithinTokenLimit({
+                messages,
+                maxContextTokens: 40,
+                indexTokenCountMap,
+            });
+            // Should include system message and the last two messages
+            expect(result.context.length).toBe(3);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[0].getType()).toBe('system'); // System message
+            expect(result.context[1]).toBe(messages[3]); // Message 2
+            expect(result.context[2]).toBe(messages[4]); // Response 2
+            // Should have the first two messages in messagesToRefine
+            expect(result.messagesToRefine.length).toBe(2);
+            expect(result.messagesToRefine[0]).toBe(messages[1]); // Message 1
+            expect(result.messagesToRefine[1]).toBe(messages[2]); // Response 1
+        });
+        it('should always include system message even when at token limit', () => {
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Hello'),
+                new AIMessage('Hi there'),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // "System instruction"
+                1: 5, // "Hello"
+                2: 8, // "Hi there"
+            };
+            // Set a limit that can only fit the system message
+            const result = getMessagesWithinTokenLimit({
+                messages,
+                maxContextTokens: 20,
+                indexTokenCountMap,
+            });
+            expect(result.context.length).toBe(1);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.messagesToRefine.length).toBe(2);
+        });
+        it('should start context with a specific message type when startType is specified', () => {
+            const messages = [
+                new SystemMessage('System instruction'),
+                new AIMessage('AI message 1'),
+                new HumanMessage('Human message 1'),
+                new AIMessage('AI message 2'),
+                new HumanMessage('Human message 2'),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // "System instruction"
+                1: 12, // "AI message 1"
+                2: 15, // "Human message 1"
+                3: 12, // "AI message 2"
+                4: 15, // "Human message 2"
+            };
+            // Set a limit that can fit all messages
+            const result = getMessagesWithinTokenLimit({
+                messages,
+                maxContextTokens: 100,
+                indexTokenCountMap,
+                startType: 'human',
+            });
+            // All messages should be included since we're under the token limit
+            expect(result.context.length).toBe(5);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[1]); // AI message 1
+            expect(result.context[2]).toBe(messages[2]); // Human message 1
+            expect(result.context[3]).toBe(messages[3]); // AI message 2
+            expect(result.context[4]).toBe(messages[4]); // Human message 2
+            // All messages should be included since we're under the token limit
+            expect(result.messagesToRefine.length).toBe(0);
+        });
+        it('should keep all messages if no message of required type is found', () => {
+            const messages = [
+                new SystemMessage('System instruction'),
+                new AIMessage('AI message 1'),
+                new AIMessage('AI message 2'),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // "System instruction"
+                1: 12, // "AI message 1"
+                2: 12, // "AI message 2"
+            };
+            // Set a limit that can fit all messages
+            const result = getMessagesWithinTokenLimit({
+                messages,
+                maxContextTokens: 100,
+                indexTokenCountMap,
+                startType: 'human',
+            });
+            // Should include all messages since no human messages exist to start from
+            expect(result.context.length).toBe(3);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[1]); // AI message 1
+            expect(result.context[2]).toBe(messages[2]); // AI message 2
+            expect(result.messagesToRefine.length).toBe(0);
+        });
+    });
+    describe('checkValidNumber', () => {
+        it('should return true for valid positive numbers', () => {
+            expect(checkValidNumber(5)).toBe(true);
+            expect(checkValidNumber(1.5)).toBe(true);
+            expect(checkValidNumber(Number.MAX_SAFE_INTEGER)).toBe(true);
+        });
+        it('should return false for zero, negative numbers, and NaN', () => {
+            expect(checkValidNumber(0)).toBe(false);
+            expect(checkValidNumber(-5)).toBe(false);
+            expect(checkValidNumber(NaN)).toBe(false);
+        });
+        it('should return false for non-number types', () => {
+            expect(checkValidNumber('5')).toBe(false);
+            expect(checkValidNumber(null)).toBe(false);
+            expect(checkValidNumber(undefined)).toBe(false);
+            expect(checkValidNumber({})).toBe(false);
+            expect(checkValidNumber([])).toBe(false);
+        });
+    });
+    describe('createPruneMessages', () => {
+        it('should return all messages when under token limit', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Hello'),
+                new AIMessage('Hi there'),
+            ];
+            const indexTokenCountMap = {
+                0: tokenCounter(messages[0]),
+                1: tokenCounter(messages[1]),
+                2: tokenCounter(messages[2]),
+            };
+            const pruneMessages = createPruneMessages({
+                maxTokens: 100,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap,
+            });
+            const result = pruneMessages({ messages });
+            expect(result.context.length).toBe(3);
+            expect(result.context).toEqual(messages);
+        });
+        it('should prune messages when over token limit', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Message 1'),
+                new AIMessage('Response 1'),
+                new HumanMessage('Message 2'),
+                new AIMessage('Response 2'),
+            ];
+            const indexTokenCountMap = {
+                0: tokenCounter(messages[0]),
+                1: tokenCounter(messages[1]),
+                2: tokenCounter(messages[2]),
+                3: tokenCounter(messages[3]),
+                4: tokenCounter(messages[4]),
+            };
+            // Set a limit that can only fit the system message and the last two messages
+            const pruneMessages = createPruneMessages({
+                maxTokens: 40,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap,
+            });
+            const result = pruneMessages({ messages });
+            // Should include system message and the last two messages
+            expect(result.context.length).toBe(3);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[3]); // Message 2
+            expect(result.context[2]).toBe(messages[4]); // Response 2
+        });
+        it('should respect startType parameter', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new AIMessage('AI message 1'),
+                new HumanMessage('Human message 1'),
+                new AIMessage('AI message 2'),
+                new HumanMessage('Human message 2'),
+            ];
+            const indexTokenCountMap = {
+                0: tokenCounter(messages[0]),
+                1: tokenCounter(messages[1]),
+                2: tokenCounter(messages[2]),
+                3: tokenCounter(messages[3]),
+                4: tokenCounter(messages[4]),
+            };
+            // Set a limit that can fit all messages
+            const pruneMessages = createPruneMessages({
+                maxTokens: 100,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({
+                messages,
+                startType: 'human',
+            });
+            // All messages should be included since we're under the token limit
+            expect(result.context.length).toBe(5);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[1]); // AI message 1
+            expect(result.context[2]).toBe(messages[2]); // Human message 1
+            expect(result.context[3]).toBe(messages[3]); // AI message 2
+            expect(result.context[4]).toBe(messages[4]); // Human message 2
+        });
+        it('should update token counts when usage metadata is provided', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Hello'),
+                new AIMessage('Hi there'),
+            ];
+            const indexTokenCountMap = {
+                0: tokenCounter(messages[0]),
+                1: tokenCounter(messages[1]),
+                2: tokenCounter(messages[2]),
+            };
+            const pruneMessages = createPruneMessages({
+                maxTokens: 100,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            // Provide usage metadata that indicates different token counts
+            const usageMetadata = {
+                input_tokens: 50,
+                output_tokens: 25,
+                total_tokens: 75,
+            };
+            const result = pruneMessages({
+                messages,
+                usageMetadata,
+            });
+            // The function should have updated the indexTokenCountMap based on the usage metadata
+            expect(result.indexTokenCountMap).not.toEqual(indexTokenCountMap);
+            // The total of all values in indexTokenCountMap should equal the total_tokens from usageMetadata
+            const totalTokens = Object.values(result.indexTokenCountMap).reduce((a = 0, b = 0) => a + b, 0);
+            expect(totalTokens).toBe(75);
+        });
+    });
+    describe('Tool Message Handling', () => {
+        it('should ensure context does not start with a tool message by finding an AI message', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new AIMessage('AI message 1'),
+                new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
+                new AIMessage('AI message 2'),
+                new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' }),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // System instruction
+                1: 12, // AI message 1
+                2: 13, // Tool result 1
+                3: 12, // AI message 2
+                4: 13, // Tool result 2
+            };
+            // Create a pruneMessages function with a token limit that will only include the last few messages
+            const pruneMessages = createPruneMessages({
+                maxTokens: 58, // Only enough for system + last 3 messages + 3, but should not include a parent-less tool message
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            // The context should include the system message, AI message 2, and Tool result 2
+            // It should NOT start with Tool result 2 alone
+            expect(result.context.length).toBe(3);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[3]); // AI message 2
+            expect(result.context[2]).toBe(messages[4]); // Tool result 2
+        });
+        it('should ensure context does not start with a tool message by finding a human message', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Human message 1'),
+                new AIMessage('AI message 1'),
+                new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
+                new HumanMessage('Human message 2'),
+                new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' }),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // System instruction
+                1: 15, // Human message 1
+                2: 12, // AI message 1
+                3: 13, // Tool result 1
+                4: 15, // Human message 2
+                5: 13, // Tool result 2
+            };
+            // Create a pruneMessages function with a token limit that will only include the last few messages
+            const pruneMessages = createPruneMessages({
+                maxTokens: 48, // Only enough for system + last 2 messages
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            // The context should include the system message, Human message 2, and Tool result 2
+            // It should NOT start with Tool result 2 alone
+            expect(result.context.length).toBe(3);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[4]); // Human message 2
+            expect(result.context[2]).toBe(messages[5]); // Tool result 2
+        });
+        it('should handle the case where a tool message is followed by an AI message', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Human message'),
+                new AIMessage('AI message with tool use'),
+                new ToolMessage({ content: 'Tool result', tool_call_id: 'tool1' }),
+                new AIMessage('AI message after tool'),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // System instruction
+                1: 13, // Human message
+                2: 22, // AI message with tool use
+                3: 11, // Tool result
+                4: 19, // AI message after tool
+            };
+            const pruneMessages = createPruneMessages({
+                maxTokens: 50,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            expect(result.context.length).toBe(2);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[4]); // AI message after tool
+        });
+        it('should handle the case where a tool message is followed by a human message', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Human message 1'),
+                new AIMessage('AI message with tool use'),
+                new ToolMessage({ content: 'Tool result', tool_call_id: 'tool1' }),
+                new HumanMessage('Human message 2'),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // System instruction
+                1: 15, // Human message 1
+                2: 22, // AI message with tool use
+                3: 11, // Tool result
+                4: 15, // Human message 2
+            };
+            const pruneMessages = createPruneMessages({
+                maxTokens: 46,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            expect(result.context.length).toBe(2);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[4]); // Human message 2
+        });
+        it('should handle complex sequence with multiple tool messages', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Human message 1'),
+                new AIMessage('AI message 1 with tool use'),
+                new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
+                new AIMessage('AI message 2 with tool use'),
+                new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' }),
+                new AIMessage('AI message 3 with tool use'),
+                new ToolMessage({ content: 'Tool result 3', tool_call_id: 'tool3' }),
+            ];
+            const indexTokenCountMap = {
+                0: 17, // System instruction
+                1: 15, // Human message 1
+                2: 26, // AI message 1 with tool use
+                3: 13, // Tool result 1
+                4: 26, // AI message 2 with tool use
+                5: 13, // Tool result 2
+                6: 26, // AI message 3 with tool use
+                7: 13, // Tool result 3
+            };
+            const pruneMessages = createPruneMessages({
+                maxTokens: 111,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            expect(result.context.length).toBe(5);
+            expect(result.context[0]).toBe(messages[0]); // System message
+            expect(result.context[1]).toBe(messages[4]); // AI message 2 with tool use
+            expect(result.context[2]).toBe(messages[5]); // Tool result 2
+            expect(result.context[3]).toBe(messages[6]); // AI message 3 with tool use
+            expect(result.context[4]).toBe(messages[7]); // Tool result 3
+        });
+    });
+    describe('messagesToRefine return value', () => {
+        it('should return empty messagesToRefine when all messages fit', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'),
+                new HumanMessage('Hello'),
+                new AIMessage('Hi there'),
+            ];
+            const indexTokenCountMap = {
+                0: tokenCounter(messages[0]),
+                1: tokenCounter(messages[1]),
+                2: tokenCounter(messages[2]),
+            };
+            const pruneMessages = createPruneMessages({
+                maxTokens: 500,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            expect(result.messagesToRefine).toEqual([]);
+            expect(Array.isArray(result.messagesToRefine)).toBe(true);
+            expect(result.messagesToRefine.length).toBe(0);
+            // All messages should be in context
+            expect(result.context.length).toBe(3);
+            expect(result.context).toEqual(messages);
+        });
+        it('should return pruned messages in messagesToRefine when context is exceeded', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('System instruction'), // 18 chars
+                new HumanMessage('First message from user'), // 23 chars
+                new AIMessage('First response from assistant'), // 29 chars
+                new HumanMessage('Second message from user'), // 24 chars
+                new AIMessage('Second response'), // 15 chars
+                new HumanMessage('Third message'), // 13 chars
+                new AIMessage('Third response'), // 14 chars
+            ];
+            const indexTokenCountMap = {};
+            for (let i = 0; i < messages.length; i++) {
+                indexTokenCountMap[i] = tokenCounter(messages[i]);
+            }
+            // Set a low maxTokens so that only system + last couple of messages fit
+            // System instruction = 18, plus 3 base tokens = 21 reserved
+            // Remaining budget = 55 - 18 = 37, minus 3 base = 34
+            // Third response (14) + Third message (13) = 27, fits
+            // Second response (15) would make 42, exceeds 34
+            const pruneMessages = createPruneMessages({
+                maxTokens: 55,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            // The context should contain the system message and the most recent messages that fit
+            expect(result.context.length).toBeGreaterThanOrEqual(2);
+            expect(result.context[0]).toBe(messages[0]); // System message always included
+            // messagesToRefine should contain the oldest messages that were discarded
+            expect(result.messagesToRefine.length).toBeGreaterThan(0);
+            // The pruned messages should be the ones NOT in context (excluding system message)
+            // Verify no overlap: messagesToRefine should not contain any message from context
+            for (const prunedMsg of result.messagesToRefine) {
+                expect(result.context).not.toContain(prunedMsg);
+            }
+            // messagesToRefine contains the boundary messages that triggered the overflow
+            // (the pruning loop pops from newest to oldest, and stops when a message doesn't fit)
+            // So messagesToRefine will have at least the first message that exceeded the budget
+            expect(result.messagesToRefine.length).toBeGreaterThanOrEqual(1);
+            // Each message in messagesToRefine should be a valid BaseMessage
+            for (const msg of result.messagesToRefine) {
+                expect(msg).toBeDefined();
+                expect(typeof msg.getType()).toBe('string');
+            }
+            ;
+        });
+        it('should return messagesToRefine compatible with summarization callbacks', () => {
+            const tokenCounter = createTestTokenCounter();
+            const messages = [
+                new SystemMessage('You are a helpful assistant'),
+                new HumanMessage('Tell me about TypeScript'),
+                new AIMessage('TypeScript is a typed superset of JavaScript'),
+                new HumanMessage('What about generics?'),
+                new AIMessage('Generics allow creating reusable components'),
+                new HumanMessage('Show me an example'),
+                new AIMessage('Here is a generic function example'),
+            ];
+            const indexTokenCountMap = {};
+            for (let i = 0; i < messages.length; i++) {
+                indexTokenCountMap[i] = tokenCounter(messages[i]);
+            }
+            // Use a low token limit to force pruning
+            const pruneMessages = createPruneMessages({
+                maxTokens: 80,
+                startIndex: 0,
+                tokenCounter,
+                indexTokenCountMap: { ...indexTokenCountMap },
+            });
+            const result = pruneMessages({ messages });
+            // Verify messagesToRefine contains valid BaseMessage instances
+            expect(result.messagesToRefine.length).toBeGreaterThan(0);
+            for (const msg of result.messagesToRefine) {
+                // Each message should be an instance of BaseMessage
+                expect(msg).toBeInstanceOf(BaseMessage);
+                // Each message should have a valid type
+                const msgType = msg.getType();
+                expect(['human', 'ai', 'system', 'tool', 'generic', 'function']).toContain(msgType);
+                // Each message should have accessible content (string or array)
+                expect(msg.content).toBeDefined();
+                expect(typeof msg.content === 'string' || Array.isArray(msg.content)).toBe(true);
+            }
+            // Verify the messages could be serialized (important for passing to callbacks)
+            const serialized = result.messagesToRefine.map((msg) => ({
+                type: msg.getType(),
+                content: msg.content,
+            }));
+            expect(serialized.length).toBe(result.messagesToRefine.length);
+            expect(serialized.every((s) => s.type && s.content !== undefined)).toBe(true);
+        });
+    });
+    describe('Integration with Run', () => {
+        it('should initialize Run with custom token counter and process messages', async () => {
+            const provider = Providers.OPENAI;
+            const llmConfig = getLLMConfig(provider);
+            const tokenCounter = createTestTokenCounter();
+            const run = await Run.create({
+                runId: 'test-prune-run',
+                graphConfig: {
+                    type: 'standard',
+                    llmConfig,
+                    instructions: 'You are a helpful assistant.',
+                    maxContextTokens: 1000,
+                },
+                returnContent: true,
+                skipCleanup: true,
+                tokenCounter,
+                indexTokenCountMap: {},
+            });
+            // Override the model to use a fake LLM
+            run.Graph?.overrideTestModel(['This is a test response'], 1);
+            const messages = [new HumanMessage('Hello, how are you?')];
+            const config = {
+                configurable: {
+                    thread_id: 'test-thread',
+                },
+                streamMode: 'values',
+                version: 'v2',
+            };
+            await run.processStream({ messages }, config);
+            const finalMessages = run.getRunMessages();
+            expect(finalMessages).toBeDefined();
+            expect(finalMessages?.length).toBeGreaterThan(0);
+        });
+    });
+});
+//# sourceMappingURL=prune.test.js.map