@librechat/agents 3.1.56 → 3.1.60

package/src/specs/prune.test.ts

@@ -2,18 +2,26 @@
 import { config } from 'dotenv';
 config();
 import {
-  HumanMessage,
   AIMessage,
-  SystemMessage,
   BaseMessage,
   ToolMessage,
+  HumanMessage,
+  isBaseMessage,
+  SystemMessage,
+  AIMessageChunk,
 } from '@langchain/core/messages';
 import type { RunnableConfig } from '@langchain/core/runnables';
 import type { UsageMetadata } from '@langchain/core/messages';
 import type * as t from '@/types';
-import { createPruneMessages } from '@/messages/prune';
+import {
+  getMessagesWithinTokenLimit as realGetMessagesWithinTokenLimit,
+  preFlightTruncateToolCallInputs,
+  repairOrphanedToolMessages,
+  sanitizeOrphanToolBlocks,
+  createPruneMessages,
+} from '@/messages/prune';
 import { getLLMConfig } from '@/utils/llmConfig';
-import { Providers } from '@/common';
+import { Providers, ContentTypes } from '@/common';
 import { Run } from '@/run';
 
 // Create a simple token counter for testing
@@ -42,6 +50,14 @@ const createTestTokenCounter = (): t.TokenCounter => {
           if ('text' in item && typeof item.text === 'string') {
             totalLength += item.text.length;
           }
+          // Count tool_use input fields (serialized args contribute to token count)
+          if ('input' in item && item.input != null) {
+            const input = item.input;
+            totalLength +=
+              typeof input === 'string'
+                ? input.length
+                : JSON.stringify(input).length;
+          }
         }
       }
 
@@ -410,6 +426,8 @@ describe('Prune Messages Tests', () => {
 
       expect(result.context.length).toBe(3);
       expect(result.context).toEqual(messages);
+      expect(result.messagesToRefine).toEqual([]);
+      expect(result.remainingContextTokens).toBeGreaterThan(0);
     });
 
     it('should prune messages when over token limit', () => {
@@ -436,6 +454,7 @@ describe('Prune Messages Tests', () => {
         startIndex: 0,
         tokenCounter,
         indexTokenCountMap,
+        reserveRatio: 0,
       });
 
       const result = pruneMessages({ messages });
@@ -445,6 +464,9 @@ describe('Prune Messages Tests', () => {
       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
+      expect(Array.isArray(result.messagesToRefine)).toBe(true);
+      expect(result.messagesToRefine?.length).toBe(2);
+      expect(typeof result.remainingContextTokens).toBe('number');
     });
 
     it('should respect startType parameter', () => {
@@ -520,26 +542,229 @@ describe('Prune Messages Tests', () => {
         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(
+      // Map stays in raw tiktoken space — calibrationRatio captures the multiplier.
+      // rawSum * calibrationRatio should approximate input_tokens (50).
+      const rawSum = Object.values(result.indexTokenCountMap).reduce(
         (a = 0, b = 0) => a + b,
         0
+      ) as number;
+      const calibratedEstimate = Math.round(
+        rawSum * (result.calibrationRatio ?? 1)
       );
-      expect(totalTokens).toBe(75);
+      expect(Math.abs(calibratedEstimate - 50)).toBeLessThanOrEqual(3);
     });
   });
 
   describe('Tool Message Handling', () => {
+    it('should drop orphan tool messages that no longer have matching AI tool calls', () => {
+      const tokenCounter = createTestTokenCounter();
+      const context = [
+        new SystemMessage('System instruction'),
+        new ToolMessage({
+          content: 'Orphan result',
+          tool_call_id: 'tool-orphan',
+        }),
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'I will call a tool now' },
+            {
+              type: 'tool_use',
+              id: 'tool-valid',
+              name: 'read_file',
+              input: '{"path":"README.md"}',
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Valid result',
+          tool_call_id: 'tool-valid',
+        }),
+      ];
+
+      const indexTokenCountMap = {
+        0: tokenCounter(context[0]),
+        1: tokenCounter(context[1]),
+        2: tokenCounter(context[2]),
+        3: tokenCounter(context[3]),
+      };
+
+      const repaired = repairOrphanedToolMessages({
+        context,
+        allMessages: context,
+        tokenCounter,
+        indexTokenCountMap,
+      });
+
+      expect(repaired.context).toHaveLength(3);
+      expect(repaired.context[0]).toBe(context[0]);
+      expect(repaired.context[1]).toBe(context[2]);
+      expect(repaired.context[2]).toBe(context[3]);
+      expect(repaired.droppedOrphanCount).toBe(1);
+      expect(repaired.reclaimedTokens).toBe(indexTokenCountMap[1]);
+    });
+
+    it('should strip orphan tool_use blocks from AI messages when ToolMessages are not in context', () => {
+      const tokenCounter = createTestTokenCounter();
+      const context = [
+        new HumanMessage('Show me something cool'),
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'Let me create an animation.' },
+            {
+              type: 'tool_use',
+              id: 'tool-navigate',
+              name: 'navigate_page',
+              input: '{"url":"about:blank"}',
+            },
+            {
+              type: 'tool_use',
+              id: 'tool-script',
+              name: 'evaluate_script',
+              input: '{"function":"' + 'x'.repeat(3000) + '"}',
+            },
+          ],
+          tool_calls: [
+            {
+              id: 'tool-navigate',
+              name: 'navigate_page',
+              args: { url: 'about:blank' },
+            },
+            {
+              id: 'tool-script',
+              name: 'evaluate_script',
+              args: { fn: 'x'.repeat(3000) },
+            },
+          ],
+        }),
+        // ToolMessages for both tool calls are NOT in context (pruned)
+      ];
+
+      const indexTokenCountMap = {
+        0: tokenCounter(context[0]),
+        1: tokenCounter(context[1]),
+      };
+
+      const repaired = repairOrphanedToolMessages({
+        context,
+        allMessages: context,
+        tokenCounter,
+        indexTokenCountMap,
+      });
+
+      // AI message should survive but with tool_use blocks stripped
+      expect(repaired.context).toHaveLength(2);
+      const repairedAI = repaired.context[1] as AIMessage;
+      expect(repairedAI.getType()).toBe('ai');
+
+      // Should only have the text block, no tool_use blocks
+      const content = repairedAI.content as Array<{ type: string }>;
+      expect(content).toHaveLength(1);
+      expect(content[0].type).toBe('text');
+
+      // tool_calls should be empty
+      expect(repairedAI.tool_calls ?? []).toHaveLength(0);
+
+      // Token savings from stripping the large tool_use blocks
+      expect(repaired.reclaimedTokens).toBeGreaterThan(0);
+    });
+
+    it('should drop AI message entirely when it has only tool_use blocks with no text', () => {
+      const tokenCounter = createTestTokenCounter();
+      const context = [
+        new HumanMessage('Do something'),
+        new AIMessage({
+          content: [
+            {
+              type: 'tool_use',
+              id: 'tool-only',
+              name: 'some_tool',
+              input: '{"query":"test"}',
+            },
+          ],
+          tool_calls: [
+            { id: 'tool-only', name: 'some_tool', args: { query: 'test' } },
+          ],
+        }),
+      ];
+
+      const indexTokenCountMap = {
+        0: tokenCounter(context[0]),
+        1: tokenCounter(context[1]),
+      };
+
+      const repaired = repairOrphanedToolMessages({
+        context,
+        allMessages: context,
+        tokenCounter,
+        indexTokenCountMap,
+      });
+
+      // AI message should be dropped since it had only tool_use blocks
+      expect(repaired.context).toHaveLength(1);
+      expect(repaired.context[0].getType()).toBe('human');
+      expect(repaired.droppedOrphanCount).toBe(1);
+    });
+
+    it('should keep tool_use blocks when their ToolMessages ARE in context', () => {
+      const tokenCounter = createTestTokenCounter();
+      const context = [
+        new HumanMessage('Do something'),
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'Calling tool' },
+            {
+              type: 'tool_use',
+              id: 'tool-present',
+              name: 'read_file',
+              input: '{"path":"test.txt"}',
+            },
+          ],
+          tool_calls: [
+            {
+              id: 'tool-present',
+              name: 'read_file',
+              args: { path: 'test.txt' },
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'File contents here',
+          tool_call_id: 'tool-present',
+        }),
+      ];
+
+      const indexTokenCountMap = {
+        0: tokenCounter(context[0]),
+        1: tokenCounter(context[1]),
+        2: tokenCounter(context[2]),
+      };
+
+      const repaired = repairOrphanedToolMessages({
+        context,
+        allMessages: context,
+        tokenCounter,
+        indexTokenCountMap,
+      });
+
+      // Nothing should change — all tool_use blocks have matching ToolMessages
+      expect(repaired.context).toHaveLength(3);
+      expect(repaired.reclaimedTokens).toBe(0);
+      expect(repaired.droppedOrphanCount).toBe(0);
+    });
+
     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 AIMessage({
+          content: [{ type: 'text', text: 'AI msg 1' }],
+          tool_calls: [{ id: 'tool1', name: 'read_file', args: { p: '1' } }],
+        }),
         new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
-        new AIMessage('AI message 2'),
+        new AIMessage({
+          content: [{ type: 'text', text: 'AI msg 2' }],
+          tool_calls: [{ id: 'tool2', name: 'read_file', args: { p: '2' } }],
+        }),
         new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' }),
       ];
 
@@ -562,10 +787,10 @@ describe('Prune Messages Tests', () => {
       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
+      // AI message 1 + Tool result 1 are pruned. Tool result 1 is orphaned (AI 1 pruned).
       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[1].getType()).toBe('ai'); // AI message 2
       expect(result.context[2]).toBe(messages[4]); // Tool result 2
     });
 
@@ -577,6 +802,7 @@ describe('Prune Messages Tests', () => {
         new AIMessage('AI message 1'),
         new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
         new HumanMessage('Human message 2'),
+        // Tool result 2 has no parent AI tool_call — this is an orphan
         new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' }),
       ];
 
@@ -595,16 +821,16 @@ describe('Prune Messages Tests', () => {
         startIndex: 0,
         tokenCounter,
         indexTokenCountMap: { ...indexTokenCountMap },
+        reserveRatio: 0,
       });
 
       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);
+      // Tool result 2 is an orphan (no AI message with tool_call_id 'tool2' in context)
+      // so it gets dropped. Context is system + human message 2.
+      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
-      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', () => {
@@ -676,11 +902,20 @@ describe('Prune Messages Tests', () => {
       const messages = [
         new SystemMessage('System instruction'),
         new HumanMessage('Human message 1'),
-        new AIMessage('AI message 1 with tool use'),
+        new AIMessage({
+          content: [{ type: 'text', text: 'AI message 1' }],
+          tool_calls: [{ id: 'tool1', name: 'read_file', args: { path: 'a' } }],
+        }),
         new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
-        new AIMessage('AI message 2 with tool use'),
+        new AIMessage({
+          content: [{ type: 'text', text: 'AI message 2' }],
+          tool_calls: [{ id: 'tool2', name: 'read_file', args: { path: 'b' } }],
+        }),
         new ToolMessage({ content: 'Tool result 2', tool_call_id: 'tool2' }),
-        new AIMessage('AI message 3 with tool use'),
+        new AIMessage({
+          content: [{ type: 'text', text: 'AI message 3' }],
+          tool_calls: [{ id: 'tool3', name: 'read_file', args: { path: 'c' } }],
+        }),
         new ToolMessage({ content: 'Tool result 3', tool_call_id: 'tool3' }),
       ];
 
@@ -704,15 +939,557 @@ describe('Prune Messages Tests', () => {
 
       const result = pruneMessages({ messages });
 
+      // AI messages 2 & 3 with their ToolMessages fit; AI1+Tool1 pruned; Tool1 is orphan (AI1 pruned)
       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[1].getType()).toBe('ai'); // AI message 2
       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[3].getType()).toBe('ai'); // AI message 3
       expect(result.context[4]).toBe(messages[7]); // Tool result 3
     });
   });
 
+  describe('preFlightTruncateToolCallInputs', () => {
+    it('should truncate oversized tool_use input fields in AI messages', () => {
+      const tokenCounter = createTestTokenCounter();
+      const largeInput = '{"function":"' + 'x'.repeat(5000) + '"}';
+      const messages: BaseMessage[] = [
+        new HumanMessage('Run this script'),
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'I will execute the script.' },
+            {
+              type: 'tool_use',
+              id: 'tool-exec',
+              name: 'evaluate_script',
+              input: largeInput,
+            },
+          ],
+          tool_calls: [
+            {
+              id: 'tool-exec',
+              name: 'evaluate_script',
+              args: { function: 'x'.repeat(5000) },
+            },
+          ],
+        }),
+        new ToolMessage({ content: 'Result: OK', tool_call_id: 'tool-exec' }),
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {
+        0: tokenCounter(messages[0]),
+        1: tokenCounter(messages[1]),
+        2: tokenCounter(messages[2]),
+      };
+
+      const originalTokens = indexTokenCountMap[1] as number;
+      expect(originalTokens).toBeGreaterThan(5000); // Large input counted
+
+      // maxContextTokens: 1000 → maxInputChars = floor(1000 * 0.15) * 4 = 600
+      const truncated = preFlightTruncateToolCallInputs({
+        messages,
+        maxContextTokens: 1000,
+        indexTokenCountMap,
+        tokenCounter,
+      });
+
+      expect(truncated).toBe(1);
+      const newTokens = indexTokenCountMap[1] as number;
+      expect(newTokens).toBeLessThan(originalTokens);
+
+      // Verify the content block was truncated
+      const aiMsg = messages[1] as AIMessage;
+      const toolUseBlock = (
+        aiMsg.content as Array<Record<string, unknown>>
+      ).find((b) => b.type === 'tool_use');
+      expect(toolUseBlock).toBeDefined();
+      const truncatedInput = toolUseBlock!.input as {
+        _truncated: string;
+        _originalChars: number;
+      };
+      expect(truncatedInput._truncated).toContain('truncated');
+      expect(truncatedInput._originalChars).toBeGreaterThan(600);
+
+      // Verify tool_calls args were also truncated
+      expect(aiMsg.tool_calls).toBeDefined();
+      const tc = aiMsg.tool_calls![0];
+      expect(tc.args).toHaveProperty('_truncated');
+    });
+
+    it('should not truncate inputs that fit within the budget', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages: BaseMessage[] = [
+        new HumanMessage('Read a file'),
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'Reading file.' },
+            {
+              type: 'tool_use',
+              id: 'tool-read',
+              name: 'read_file',
+              input: '{"path":"test.txt"}',
+            },
+          ],
+          tool_calls: [
+            { id: 'tool-read', name: 'read_file', args: { path: 'test.txt' } },
+          ],
+        }),
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {
+        0: tokenCounter(messages[0]),
+        1: tokenCounter(messages[1]),
+      };
+
+      const originalTokens = indexTokenCountMap[1];
+
+      const truncated = preFlightTruncateToolCallInputs({
+        messages,
+        maxContextTokens: 1000,
+        indexTokenCountMap,
+        tokenCounter,
+      });
+
+      expect(truncated).toBe(0);
+      expect(indexTokenCountMap[1]).toBe(originalTokens);
+    });
+
+    it('should skip non-AI messages', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages: BaseMessage[] = [
+        new HumanMessage('Hello'),
+        new ToolMessage({ content: 'x'.repeat(5000), tool_call_id: 'tool-1' }),
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {
+        0: tokenCounter(messages[0]),
+        1: tokenCounter(messages[1]),
+      };
+
+      const truncated = preFlightTruncateToolCallInputs({
+        messages,
+        maxContextTokens: 1000,
+        indexTokenCountMap,
+        tokenCounter,
+      });
+
+      // Should not touch ToolMessages (that's preFlightTruncateToolResults' job)
+      expect(truncated).toBe(0);
+    });
+  });
+
+  describe('Instruction token budget reservation (getInstructionTokens)', () => {
+    it('should reserve budget for instruction tokens when no system message is present', () => {
+      const tokenCounter = createTestTokenCounter();
+      // Agent flow: messages do NOT include a system message.
+      // The system message is prepended later by buildSystemRunnable.
+      const messages = [
+        new HumanMessage('Hello there'), // 11 chars
+        new AIMessage('Hi'), // 2 chars
+        new HumanMessage('How are you?'), // 12 chars
+        new AIMessage('Good'), // 4 chars
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      // Total message tokens: 11 + 2 + 12 + 4 = 29
+      // Instruction tokens: 20 (simulating system prompt overhead)
+      // Effective budget for messages: 50 - 20 = 30 → fits all 29 tokens
+      const pruneMessages = createPruneMessages({
+        maxTokens: 50,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        reserveRatio: 0,
+        getInstructionTokens: () => 20,
+      });
+
+      const result = pruneMessages({ messages });
+
+      // All messages should fit: 29 message tokens + 20 instruction = 49 ≤ 50
+      expect(result.context.length).toBe(4);
+      expect(result.context).toEqual(messages);
+      expect(result.messagesToRefine).toEqual([]);
+    });
+
+    it('should prune when messages + instruction tokens exceed budget', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new HumanMessage('Hello there'), // 11 chars
+        new AIMessage('Hi'), // 2 chars
+        new HumanMessage('How are you?'), // 12 chars
+        new AIMessage('Good'), // 4 chars
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      // Total message tokens: 29
+      // Instruction tokens: 25 (simulating large tool schema overhead)
+      // Effective budget: 40 - 25 = 15 → must prune older messages
+      const pruneMessages = createPruneMessages({
+        maxTokens: 40,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        getInstructionTokens: () => 25,
+      });
+
+      const result = pruneMessages({ messages });
+
+      // Should prune older messages to fit within 15 available tokens.
+      // Working backwards: "Good" (4) + "How are you?" (12) = 16 > 15
+      // So only "Good" (4) fits, context starts on that AI message.
+      // But startType may require a human message...
+      // Actually with no startType and 3 tokens of overhead,
+      // available = 15 - 3 = 12: "Good" (4) fits, "How are you?" (12) → 4+12=16 > 12
+      // So only "Good" (4) fits.
+      expect(result.context.length).toBeLessThan(4);
+      expect(Array.isArray(result.messagesToRefine)).toBe(true);
+      expect(result.messagesToRefine!.length).toBeGreaterThan(0);
+    });
+
+    it('should correctly account for instruction tokens in early-return path', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new HumanMessage('Hi'), // 2 chars
+        new AIMessage('Hello'), // 5 chars
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      // Message tokens: 7
+      // Instruction tokens: 100 (simulating 26 MCP tools ~5000 chars)
+      // Budget: 50 → 7 + 100 = 107 > 50, so early-return should NOT fire
+      const pruneMessages = createPruneMessages({
+        maxTokens: 50,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        getInstructionTokens: () => 100,
+      });
+
+      const result = pruneMessages({ messages });
+
+      // Even though messages alone (7) fit in 50, the instruction overhead (100)
+      // means pruning must occur.  With only 50 - 100 = -50 effective budget,
+      // nothing fits → all messages pruned.
+      expect(result.messagesToRefine!.length).toBeGreaterThan(0);
+    });
+
+    it('should not double-subtract when messages include a system message', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new SystemMessage('System'), // 6 chars
+        new HumanMessage('Hello there'), // 11 chars
+        new AIMessage('Hi'), // 2 chars
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      // When messages[0] IS a system message, getMessagesWithinTokenLimit uses
+      // indexTokenCountMap[0] (6) to subtract from budget, ignoring instructionTokens.
+      // getInstructionTokens is only used when no system message is at index 0.
+      const pruneMessages = createPruneMessages({
+        maxTokens: 30,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        getInstructionTokens: () => 999, // Should be ignored for system message path
+      });
+
+      const result = pruneMessages({ messages });
+
+      // Budget: 30 - 6 (system) = 24 available.
+      // "Hi" (2) + "Hello there" (11) + 3 overhead = 16, fits in 24.
+      // All messages should be kept.
+      expect(result.context.length).toBe(3);
+      expect(result.context[0]).toBe(messages[0]); // System message preserved
+    });
+
+    it('index 0 should NOT be inflated when getInstructionTokens is provided', () => {
+      const tokenCounter = createTestTokenCounter();
+      const messages = [
+        new HumanMessage('Hello there'), // 11 chars
+        new AIMessage('Hi'), // 2 chars
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      const pruneMessages = createPruneMessages({
+        maxTokens: 50,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        getInstructionTokens: () => 10,
+      });
+
+      // Before and after pruning, index 0 should remain 11 (the real token count)
+      expect(indexTokenCountMap[0]).toBe(11);
+      pruneMessages({ messages });
+      // The returned indexTokenCountMap should still have the real count at index 0
+      expect(indexTokenCountMap[0]).toBe(11);
+    });
+
+    it('pre-flight truncation uses effective budget after instruction overhead', () => {
+      const tokenCounter = createTestTokenCounter();
+      // Simulate the real scenario: AI message has a massive tool_call input
+      // (like the chrome-devtools evaluate_script with a 7000-char JS payload)
+      const hugeInput = 'x'.repeat(7000);
+      const messages = [
+        new HumanMessage('show me something'), // 17 chars
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'Creating animation' },
+            {
+              type: 'tool_use',
+              id: 'tool_1',
+              name: 'evaluate_script',
+              input: { function: hugeInput },
+            },
+          ],
+          tool_calls: [
+            {
+              id: 'tool_1',
+              name: 'evaluate_script',
+              args: { function: hugeInput },
+              type: 'tool_call' as const,
+            },
+          ],
+        }),
+        new ToolMessage({ content: 'Script executed', tool_call_id: 'tool_1' }),
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      // Pre-flight truncation uses maxTokens for the truncation threshold:
+      // Math.floor(8000*0.15)*4 = 4800 chars.  The AI message's tool_use
+      // input (~7015 chars) shrinks to ~4800 chars, giving an AI token
+      // count of ~4850.
+      //
+      // The effective pruning budget subtracts instruction overhead:
+      // effectiveMax = 8000 - 2000 = 6000, which is enough for all three
+      // messages (~4850 + 17 + 15 ≈ 4882).
+      const instructionTokens = 2000;
+      const pruneMessages = createPruneMessages({
+        maxTokens: 8000,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        getInstructionTokens: () => instructionTokens,
+        reserveRatio: 0,
+      });
+
+      const result = pruneMessages({ messages });
+
+      // The AI message should survive pruning (not be in messagesToRefine)
+      // because pre-flight truncation used the effective budget
+      const aiMessagesInContext = result.context.filter(
+        (m) => m.getType() === 'ai'
+      );
+      expect(aiMessagesInContext.length).toBe(1);
+      expect(result.context.length).toBe(3); // All 3 messages fit after truncation
+    });
+
+    it('emergency truncation recovers when initial prune produces empty context', () => {
+      const tokenCounter = createTestTokenCounter();
+      // Simulate post-summarization state: only 4 messages remain, but one
+      // has a huge tool_call input that exceeds available budget alone.
+      // With char-based counter, the AI message with 4000-char input is ~4000 tokens.
+      // Available budget: 5000 - 4500 = 500. Nothing fits on first pass.
+      const hugeInput = 'x'.repeat(4000);
+      const messages = [
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'Running script' },
+            {
+              type: 'tool_use',
+              id: 'tool_1',
+              name: 'evaluate_script',
+              input: { function: hugeInput },
+            },
+          ],
+          tool_calls: [
+            {
+              id: 'tool_1',
+              name: 'evaluate_script',
+              args: { function: hugeInput },
+              type: 'tool_call' as const,
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Script ran successfully',
+          tool_call_id: 'tool_1',
+        }),
+        new HumanMessage('that looks great'),
+        new AIMessage('Thanks! Want more?'),
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      // Available budget is extremely tight: 500 tokens for messages.
+      // The AI message alone is ~4000+ tokens. Initial prune: nothing fits.
+      // Emergency truncation should reduce tool inputs to 150 chars,
+      // making the AI message fit.
+      const pruneMessages = createPruneMessages({
+        maxTokens: 5000,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        getInstructionTokens: () => 4500,
+      });
+
+      const result = pruneMessages({ messages });
+
+      // Emergency truncation should have recovered — context is NOT empty
+      expect(result.context.length).toBeGreaterThan(0);
+      // At minimum, the newest messages should be present
+      const types = result.context.map((m) => m.getType());
+      expect(types).toContain('human');
+    });
+  });
+
+  describe('Empty messages guard', () => {
+    it('returns empty context without crashing when messages array is empty', () => {
+      const tokenCounter = createTestTokenCounter();
+      const pruneMessages = createPruneMessages({
+        maxTokens: 8000,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap: {},
+        getInstructionTokens: () => 4000,
+      });
+
+      // Simulate post-summarization state where REMOVE_ALL left an empty messages array
+      const result = pruneMessages({
+        messages: [],
+        usageMetadata: {
+          input_tokens: 100,
+          output_tokens: 50,
+          total_tokens: 150,
+        } as UsageMetadata,
+      });
+
+      expect(result.context).toEqual([]);
+      expect(result.messagesToRefine).toEqual([]);
+      expect(result.prePruneContextTokens).toBe(0);
+      expect(result.remainingContextTokens).toBe(8000);
+    });
+  });
+
+  describe('Dropped orphan ToolMessages appear in messagesToRefine', () => {
+    it('appends orphan ToolMessage (whose parent AI was pruned) to messagesToRefine for summarization', () => {
+      const tokenCounter = createTestTokenCounter();
+
+      // Build messages where the large AI(evaluate) won't fit in a tight budget,
+      // but its smaller ToolMessage(evaluate) does.  After backward iteration,
+      // the ToolMessage lands in context while its parent AI is in prunedMemory.
+      // repairOrphanedToolMessages then drops the orphan ToolMessage from context.
+      // The fix: that dropped ToolMessage must appear in messagesToRefine so
+      // summarization sees the tool result (otherwise summary says "in progress").
+      const messages: BaseMessage[] = [
+        new HumanMessage('Build me a solar system simulation'),
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'I will write the code now.' },
+            {
+              type: 'tool_use',
+              id: 'tc_eval',
+              name: 'evaluate_script',
+              // Large input that consumes most of the budget
+              input: { code: 'x'.repeat(3000) },
+            },
+          ],
+          tool_calls: [
+            {
+              id: 'tc_eval',
+              name: 'evaluate_script',
+              args: { code: 'x'.repeat(3000) },
+              type: 'tool_call' as const,
+            },
+          ],
+        }),
+        new ToolMessage({
+          // Small result — fits in budget individually
+          content: 'Solar system simulation launched successfully!',
+          tool_call_id: 'tc_eval',
+          name: 'evaluate_script',
+        }),
+      ];
+
+      const indexTokenCountMap: Record<string, number | undefined> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+      }
+
+      // Budget is tight enough that the large AI message won't fit
+      // even after emergency truncation, but HumanMessage and ToolMessage
+      // individually can.  Budget must be low enough that proportional
+      // emergency truncation (budget / messages * 4 chars) still leaves
+      // the AI message too large to fit.
+      const pruneMessages = createPruneMessages({
+        maxTokens: 100,
+        startIndex: 0,
+        tokenCounter,
+        indexTokenCountMap,
+        getInstructionTokens: () => 0,
+      });
+
+      const result = pruneMessages({ messages });
+
+      // The orphan ToolMessage(evaluate) should NOT be in context
+      // (its parent AI was pruned away)
+      const contextToolMsgs = result.context.filter(
+        (m) => m.getType() === 'tool'
+      );
+      const orphanInContext = contextToolMsgs.some(
+        (m) => (m as ToolMessage).tool_call_id === 'tc_eval'
+      );
+      expect(orphanInContext).toBe(false);
+
+      // The key assertion: the dropped ToolMessage MUST appear in messagesToRefine
+      // so that summarization can see "Solar system simulation launched successfully!"
+      expect(result.messagesToRefine).toBeDefined();
+      const refineToolMsgs = result.messagesToRefine!.filter(
+        (m) => m.getType() === 'tool'
+      );
+      const toolInRefine = refineToolMsgs.some(
+        (m) => (m as ToolMessage).tool_call_id === 'tc_eval'
+      );
+      expect(toolInRefine).toBe(true);
+
+      // The parent AI message should also be in messagesToRefine (from prunedMemory)
+      const refineAiMsgs = result.messagesToRefine!.filter(
+        (m) => m.getType() === 'ai'
+      );
+      const aiInRefine = refineAiMsgs.some((m) =>
+        ((m as AIMessage).tool_calls ?? []).some((tc) => tc.id === 'tc_eval')
+      );
+      expect(aiInRefine).toBe(true);
+    });
+  });
+
   describe('Integration with Run', () => {
     it('should initialize Run with custom token counter and process messages', async () => {
       const provider = Providers.OPENAI;
@@ -757,3 +1534,399 @@ describe('Prune Messages Tests', () => {
     });
   });
 });
+
+describe('sanitizeOrphanToolBlocks', () => {
+  it('strips orphan tool_use blocks from AI messages with no matching ToolMessage', () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('Hello'),
+      new AIMessage({
+        content: [
+          { type: 'text', text: 'Let me check.' },
+          { type: 'tool_use', id: 'tool_1', name: 'calc', input: { x: 1 } },
+        ],
+        tool_calls: [
+          { id: 'tool_1', name: 'calc', args: { x: 1 }, type: 'tool_call' },
+        ],
+      }),
+      // No ToolMessage for tool_1 — orphan
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    // The stripped AI message was the last message → dropped (incomplete tool call)
+    expect(result).toHaveLength(1);
+    expect(result[0].getType()).toBe('human');
+  });
+
+  it('drops orphan ToolMessages whose AI message is missing', () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('Hello'),
+      new ToolMessage({
+        content: 'result',
+        tool_call_id: 'tool_orphan',
+        name: 'calc',
+      }),
+      new AIMessage('Some response'),
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    expect(result).toHaveLength(2); // HumanMessage + AIMessage, orphan ToolMessage dropped
+    expect(result[0].getType()).toBe('human');
+    expect(result[1].getType()).toBe('ai');
+  });
+
+  it('preserves correctly paired tool_use and ToolMessages', () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('Compute 1+1'),
+      new AIMessage({
+        content: [
+          { type: 'text', text: 'Let me calculate.' },
+          { type: 'tool_use', id: 'tool_a', name: 'calc', input: { x: 1 } },
+        ],
+        tool_calls: [
+          { id: 'tool_a', name: 'calc', args: { x: 1 }, type: 'tool_call' },
+        ],
+      }),
+      new ToolMessage({
+        content: '2',
+        tool_call_id: 'tool_a',
+        name: 'calc',
+      }),
+      new AIMessage('The answer is 2.'),
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    expect(result).toHaveLength(4); // All messages preserved
+    expect(result.map((m) => m.getType())).toEqual([
+      'human',
+      'ai',
+      'tool',
+      'ai',
+    ]);
+  });
+
+  it('drops AI message entirely when it only contained orphan tool_use blocks', () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('Do something'),
+      new AIMessage({
+        content: [{ type: 'tool_use', id: 'tool_x', name: 'run', input: {} }],
+        tool_calls: [
+          { id: 'tool_x', name: 'run', args: {}, type: 'tool_call' },
+        ],
+      }),
+      // No ToolMessage for tool_x
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    // The AI message had only tool_use blocks, stripping them leaves nothing → dropped
+    expect(result).toHaveLength(1);
+    expect(result[0].getType()).toBe('human');
+  });
+
+  it('keeps stripped AI message in the middle but drops stripped trailing AI', () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('First question'),
+      new AIMessage({
+        content: [
+          { type: 'text', text: 'Let me use two tools.' },
+          { type: 'tool_use', id: 'tool_a', name: 'calc', input: { x: 1 } },
+          {
+            type: 'tool_use',
+            id: 'tool_orphan',
+            name: 'search',
+            input: { q: 'test' },
+          },
+        ],
+        tool_calls: [
+          { id: 'tool_a', name: 'calc', args: { x: 1 }, type: 'tool_call' },
+          {
+            id: 'tool_orphan',
+            name: 'search',
+            args: { q: 'test' },
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: '42',
+        tool_call_id: 'tool_a',
+        name: 'calc',
+      }),
+      // No ToolMessage for tool_orphan, but conversation continues:
+      new AIMessage({
+        content: [{ type: 'text', text: 'Got the calc result.' }],
+        tool_calls: [
+          { id: 'tool_b', name: 'run', args: {}, type: 'tool_call' },
+        ],
+      }),
+      // tool_b is also orphan → stripped, and this AI is last → dropped
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    // message[1]: AI has tool_orphan stripped but tool_a kept → stays (middle, not trailing)
+    // message[3]: AI has tool_b stripped, is trailing → dropped
+    expect(result).toHaveLength(3); // HumanMessage, stripped AI (kept tool_a), ToolMessage
+    const ai = result[1] as AIMessage;
+    expect(ai.tool_calls).toHaveLength(1);
+    expect(ai.tool_calls![0].id).toBe('tool_a');
+    expect(result[2].getType()).toBe('tool');
+  });
+
+  it('keeps unmodified trailing AI message (no orphan tool_use)', () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('Hello'),
+      new ToolMessage({
+        content: 'result',
+        tool_call_id: 'tool_orphan',
+        name: 'calc',
+      }),
+      new AIMessage('Final response without tool calls.'),
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    // orphan ToolMessage dropped, trailing AI kept (was not stripped)
+    expect(result).toHaveLength(2);
+    expect(result[0].getType()).toBe('human');
+    expect(result[1].getType()).toBe('ai');
+  });
+
+  it('preserves BaseMessage prototype on stripped AIMessage instances', () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('Hello'),
+      new AIMessage({
+        content: [
+          { type: 'text', text: 'Let me search and calculate.' },
+          {
+            type: 'tool_use',
+            id: 'tool_a',
+            name: 'search',
+            input: { q: 'test' },
+          },
+          { type: 'tool_use', id: 'tool_b', name: 'calc', input: { x: 1 } },
+        ],
+        tool_calls: [
+          {
+            id: 'tool_a',
+            name: 'search',
+            args: { q: 'test' },
+            type: 'tool_call' as const,
+          },
+          {
+            id: 'tool_b',
+            name: 'calc',
+            args: { x: 1 },
+            type: 'tool_call' as const,
+          },
+        ],
+      }),
+      new ToolMessage({ content: 'result', tool_call_id: 'tool_b' }),
+      // No ToolMessage for tool_a — orphan
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    // AI message should survive (tool_a stripped, tool_b kept)
+    expect(result).toHaveLength(3);
+
+    // Every output message must pass isBaseMessage and have getType()
+    for (const msg of result) {
+      expect(isBaseMessage(msg)).toBe(true);
+      expect(typeof msg.getType()).toBe('string');
+    }
+    expect(result[1].getType()).toBe('ai');
+    expect(result[1]).toBeInstanceOf(AIMessage);
+  });
+
+  it('preserves AIMessageChunk prototype on stripped messages', () => {
+    // Simulate what happens in real graph execution: model returns AIMessageChunk,
+    // state passes through LangGraph, sanitizeOrphanToolBlocks strips orphan server tools.
+    const chunk = new AIMessageChunk({
+      content: [
+        { type: 'text', text: 'Searching...' },
+        { type: 'tool_use', id: 'srvtoolu_1', name: 'web_search', input: '' },
+        { type: 'tool_use', id: 'toolu_2', name: 'calculator', input: '2+2' },
+      ],
+      tool_call_chunks: [
+        { id: 'srvtoolu_1', index: 0, name: 'web_search', args: '' },
+        { id: 'toolu_2', index: 2, name: 'calculator', args: '2+2' },
+      ],
+    });
+
+    const messages: BaseMessage[] = [
+      new HumanMessage('Search and calculate'),
+      chunk,
+      new ToolMessage({ content: '4', tool_call_id: 'toolu_2' }),
+      // No ToolMessage for srvtoolu_1 — server tool, orphan
+    ];
+
+    const result = sanitizeOrphanToolBlocks(messages);
+    expect(result).toHaveLength(3);
+
+    // The AIMessageChunk must retain its prototype so LangChain's
+    // coerceMessageLikeToMessage recognizes it as a BaseMessage.
+    const aiMsg = result[1];
+    expect(isBaseMessage(aiMsg)).toBe(true);
+    expect(typeof aiMsg.getType()).toBe('string');
+    expect(aiMsg.getType()).toBe('ai');
+  });
+
+  it('preserves prototype on plain-object messages with duck-typed patching', () => {
+    // Simulate deserialized messages that still have a prototype (e.g. from
+    // LangGraph subgraph state transfer) but aren't class instances.
+    const proto = { _getType: (): string => 'ai', getType: (): string => 'ai' };
+    const plainAi = Object.create(proto);
+    Object.assign(plainAi, {
+      role: 'assistant',
+      content: [
+        { type: 'text', text: 'checking' },
+        { type: 'tool_use', id: 'orphan_1', name: 'tool', input: {} },
+      ],
+      tool_calls: [
+        { id: 'orphan_1', name: 'tool', args: {}, type: 'tool_call' },
+      ],
+    });
+
+    const messages = [plainAi] as BaseMessage[];
+    sanitizeOrphanToolBlocks(messages);
+
+    // Stripped AI was trailing → dropped. But if we add a human after:
+    const messages2 = [
+      new HumanMessage('hi'),
+      plainAi,
+      new HumanMessage('follow up'),
+    ] as BaseMessage[];
+    const result2 = sanitizeOrphanToolBlocks(messages2);
+
+    // The patched message in the middle must still have _getType from proto
+    const middleMsg = result2[1];
+    expect(typeof middleMsg._getType).toBe('function');
+    expect(middleMsg._getType()).toBe('ai');
+  });
+
+  it('handles plain objects (non-BaseMessage instances) via duck typing', () => {
+    // Simulate messages that have lost their class instances (LangGraph state serialization)
+    const plainMessages = [
+      { role: 'user', content: 'Hello', _type: 'human' },
+      {
+        role: 'assistant',
+        _type: 'ai',
+        content: [
+          { type: 'text', text: 'Let me check.' },
+          { type: 'tool_use', id: 'tool_1', name: 'calc', input: { x: 1 } },
+        ],
+        tool_calls: [
+          { id: 'tool_1', name: 'calc', args: { x: 1 }, type: 'tool_call' },
+        ],
+      },
+      // No ToolMessage for tool_1 — orphan
+    ] as unknown as BaseMessage[];
+
+    // Should not throw "getType is not a function"
+    const result = sanitizeOrphanToolBlocks(plainMessages);
+    // The stripped AI message was the last message → dropped (incomplete tool call)
+    expect(result).toHaveLength(1);
+  });
+});
+
+describe('prunedMemory ordering with thinking enabled', () => {
+  it('messagesToRefine preserves chronological order when thinking search pops multiple messages', () => {
+    const tokenCounter = createTestTokenCounter();
+    const messages: BaseMessage[] = [
+      new HumanMessage('Hello'),
+      new AIMessage({
+        content: [
+          {
+            type: ContentTypes.REASONING_CONTENT,
+            reasoningText: {
+              text: 'Thinking about navigation...',
+              signature: 'sig1',
+            },
+          },
+          { type: 'text', text: 'Navigating now.' },
+        ],
+        tool_calls: [
+          {
+            id: 'tc_nav',
+            name: 'navigate',
+            args: { url: 'about:blank' },
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: 'Navigated to about:blank.',
+        tool_call_id: 'tc_nav',
+        name: 'navigate',
+      }),
+      new AIMessage({
+        content: [
+          {
+            type: ContentTypes.REASONING_CONTENT,
+            reasoningText: {
+              text: 'Now I will write code...',
+              signature: 'sig2',
+            },
+          },
+          { type: 'text', text: 'Running script.' },
+        ],
+        tool_calls: [
+          {
+            id: 'tc_eval',
+            name: 'evaluate',
+            args: { code: 'x'.repeat(5000) },
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: 'y'.repeat(5000), // large tool result
+        tool_call_id: 'tc_eval',
+        name: 'evaluate',
+      }),
+    ];
+
+    const indexTokenCountMap: Record<string, number | undefined> = {};
+    for (let i = 0; i < messages.length; i++) {
+      indexTokenCountMap[i] = tokenCounter(messages[i]);
+    }
+
+    // Use a very tight budget so the backward iteration must prune messages
+    // The thinking search will cause the loop to `continue` past the large ToolMessage
+    const result = realGetMessagesWithinTokenLimit({
+      messages,
+      maxContextTokens: 200, // very tight
+      indexTokenCountMap,
+      thinkingEnabled: true,
+      tokenCounter,
+      reasoningType: ContentTypes.REASONING_CONTENT,
+    });
+
+    // The key assertion: messagesToRefine must be in chronological order.
+    // AI(evaluate) at index 3 must come BEFORE ToolMessage(evaluate) at index 4.
+    for (let i = 0; i < result.messagesToRefine.length - 1; i++) {
+      const current = result.messagesToRefine[i];
+      const next = result.messagesToRefine[i + 1];
+      // A ToolMessage should never come before its AI message
+      if (next.getType() === 'ai' && current.getType() === 'tool') {
+        const toolId = (current as ToolMessage).tool_call_id;
+        const aiToolIds = ((next as AIMessage).tool_calls ?? []).map(
+          (tc) => tc.id
+        );
+        expect(aiToolIds).not.toContain(toolId);
+      }
+    }
+
+    // Verify the specific ordering: if both AI(evaluate) and Tool(evaluate) are in
+    // messagesToRefine, AI must come first.
+    const evalAiIdx = result.messagesToRefine.findIndex(
+      (m) =>
+        m.getType() === 'ai' &&
+        ((m as AIMessage).tool_calls ?? []).some((tc) => tc.id === 'tc_eval')
+    );
+    const evalToolIdx = result.messagesToRefine.findIndex(
+      (m) =>
+        m.getType() === 'tool' && (m as ToolMessage).tool_call_id === 'tc_eval'
+    );
+    if (evalAiIdx >= 0 && evalToolIdx >= 0) {
+      expect(evalAiIdx).toBeLessThan(evalToolIdx);
+    }
+  });
+});