@librechat/agents 2.2.9 → 2.3.0

package/dist/cjs/messages/prune.cjs

@@ -64,8 +64,15 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
                 context = context.slice(requiredTypeIndex);
             }
         }
+        // Add system message if it exists
+        if (instructions && _messages.length > 0) {
+            context.push(_messages[0]);
+            messages$1.shift();
+        }
         // Handle thinking mode requirement for Anthropic
         if (thinkingEnabled && context.length > 0 && tokenCounter) {
+            // Check if the latest message is an assistant message
+            const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';
             // Process only if we have an assistant message in the context
             const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');
             if (firstAssistantIndex >= 0) {
@@ -75,10 +82,11 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
                     firstAssistantMsg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');
                 // Only proceed if we need to add thinking blocks
                 if (!hasThinkingBlock) {
-                    // Collect thinking blocks from pruned assistant messages
+                    // Collect thinking blocks from pruned assistant messages, starting from the most recent
                     const thinkingBlocks = [];
-                    // Look through pruned messages for thinking blocks
-                    for (const msg of messages$1) {
+                    // Look through pruned messages for thinking blocks, starting from the end (most recent)
+                    for (let i = messages$1.length - 1; i >= 0; i--) {
+                        const msg = messages$1[i];
                         if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
                             for (const item of msg.content) {
                                 if (item && typeof item === 'object' && item.type === 'thinking') {
@@ -125,28 +133,100 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
                         context[firstAssistantIndex] = newMessage;
                         // If we've exceeded the token limit, we need to prune more messages
                         if (currentTokenCount > remainingContextTokens) {
-                            // Remove messages from the end of the context until we're under the token limit
-                            // But make sure to keep the first assistant message with thinking block
+                            // Build a map of tool call IDs to track AI <--> tool message correspondences
+                            const toolCallIdMap = new Map();
+                            // Identify tool call IDs in the context
+                            for (let i = 0; i < context.length; i++) {
+                                const msg = context[i];
+                                // Check for tool calls in AI messages
+                                if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
+                                    for (const item of msg.content) {
+                                        if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                                            toolCallIdMap.set(item.id, i);
+                                        }
+                                    }
+                                }
+                                // Check for tool messages
+                                if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {
+                                    toolCallIdMap.set(msg.tool_call_id, i);
+                                }
+                            }
+                            // Track which messages to remove
+                            const indicesToRemove = new Set();
+                            // Start removing messages from the end, but preserve AI <--> tool message correspondences
                             let i = context.length - 1;
                             while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {
                                 const msgToRemove = context[i];
-                                const msgTokenCount = tokenCounter(msgToRemove);
-                                context.splice(i, 1);
-                                currentTokenCount -= msgTokenCount;
+                                // Check if this is a tool message or has tool calls
+                                let canRemove = true;
+                                if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {
+                                    // If this is a tool message, check if we need to keep its corresponding AI message
+                                    const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);
+                                    if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {
+                                        // We need to remove both the tool message and its corresponding AI message
+                                        indicesToRemove.add(i);
+                                        indicesToRemove.add(aiIndex);
+                                        currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));
+                                        canRemove = false;
+                                    }
+                                }
+                                else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {
+                                    // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages
+                                    for (const item of msgToRemove.content) {
+                                        if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                                            const toolIndex = toolCallIdMap.get(item.id);
+                                            if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {
+                                                // We need to remove both the AI message and its corresponding tool message
+                                                indicesToRemove.add(i);
+                                                indicesToRemove.add(toolIndex);
+                                                currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));
+                                                canRemove = false;
+                                                break;
+                                            }
+                                        }
+                                    }
+                                }
+                                // If we can remove this message individually
+                                if (canRemove && !indicesToRemove.has(i)) {
+                                    indicesToRemove.add(i);
+                                    currentTokenCount -= tokenCounter(msgToRemove);
+                                }
                                 i--;
                             }
+                            // Remove messages in reverse order to avoid index shifting
+                            const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);
+                            for (const index of sortedIndices) {
+                                context.splice(index, 1);
+                            }
                             // Update remainingContextTokens to reflect the new token count
                             remainingContextTokens = maxContextTokens - currentTokenCount;
                         }
                     }
                 }
             }
+            // If the latest message is an assistant message, ensure an assistant message appears early in the context
+            // but maintain system message precedence
+            if (latestMessageIsAssistant && context.length > 0) {
+                // Find the first assistant message in the context
+                const assistantIndex = context.findIndex(msg => msg.getType() === 'ai');
+                if (assistantIndex > 0) {
+                    // Check if there's a system message at the beginning
+                    const hasSystemFirst = context[0].getType() === 'system';
+                    // Move the assistant message to the appropriate position
+                    const assistantMsg = context[assistantIndex];
+                    context.splice(assistantIndex, 1);
+                    if (hasSystemFirst) {
+                        // Insert after the system message
+                        context.splice(1, 0, assistantMsg);
+                    }
+                    else {
+                        // Insert at the beginning if no system message
+                        context.unshift(assistantMsg);
+                    }
+                }
+            }
         }
     }
-    if (instructions && _messages.length > 0) {
-        context.push(_messages[0]);
-        messages$1.shift();
-    }
     const prunedMemory = messages$1;
     summaryIndex = prunedMemory.length - 1;
     remainingContextTokens -= currentTokenCount;

package/dist/cjs/messages/prune.cjs.map

@@ -1 +1 @@
-{"version":3,"file":"prune.cjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n  maxTokens: number;\n  startIndex: number;\n  tokenCounter: TokenCounter;\n  indexTokenCountMap: Record<string, number>;\n  thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n  messages: BaseMessage[];\n  usageMetadata?: Partial<UsageMetadata>;\n  startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n  const baseInputTokens = Number(usage.input_tokens) || 0;\n  const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n  const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n  \n  const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n  const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n  return {\n    input_tokens: totalInputTokens,\n    output_tokens: totalOutputTokens,\n    total_tokens: totalInputTokens + totalOutputTokens\n  };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n  messages: _messages,\n  maxContextTokens,\n  indexTokenCountMap,\n  startOnMessageType,\n  thinkingEnabled,\n  tokenCounter,\n}: {\n  messages: BaseMessage[];\n  maxContextTokens: number;\n  indexTokenCountMap: Record<string, number>;\n  startOnMessageType?: string;\n  thinkingEnabled?: boolean;\n  tokenCounter?: TokenCounter;\n}): {\n  context: BaseMessage[];\n  remainingContextTokens: number;\n  messagesToRefine: BaseMessage[];\n  summaryIndex: number;\n} {\n  // Every reply is primed with <|start|>assistant<|message|>, so we\n  // start with 3 tokens for the label after all messages have been counted.\n  let summaryIndex = -1;\n  let currentTokenCount = 3;\n  const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n  const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n  let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n  const messages = [..._messages];\n  let context: BaseMessage[] = [];\n\n  if (currentTokenCount < remainingContextTokens) {\n    let currentIndex = messages.length;\n    while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n      currentIndex--;\n      if (messages.length === 1 && instructions) {\n        break;\n      }\n      const poppedMessage = messages.pop();\n      if (!poppedMessage) continue;\n      \n      const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n      if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n        context.push(poppedMessage);\n        currentTokenCount += tokenCount;\n      } else {\n        messages.push(poppedMessage);\n        break;\n      }\n    }\n    \n  // Handle startOnMessageType requirement\n  if (startOnMessageType && context.length > 0) {\n    const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n    \n    if (requiredTypeIndex > 0) {\n      context = context.slice(requiredTypeIndex);\n    }\n  }\n  \n  // Handle thinking mode requirement for Anthropic\n  if (thinkingEnabled && context.length > 0 && tokenCounter) {\n    // Process only if we have an assistant message in the context\n    const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n    if (firstAssistantIndex >= 0) {\n      const firstAssistantMsg = context[firstAssistantIndex];\n      \n      // Check if the first assistant message already has a thinking block\n      const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n        firstAssistantMsg.content.some(item => \n          item && typeof item === 'object' && item.type === 'thinking');\n      \n      // Only proceed if we need to add thinking blocks\n      if (!hasThinkingBlock) {\n        // Collect thinking blocks from pruned assistant messages\n        const thinkingBlocks: any[] = [];\n        \n        // Look through pruned messages for thinking blocks\n        for (const msg of messages) {\n          if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n            for (const item of msg.content) {\n              if (item && typeof item === 'object' && item.type === 'thinking') {\n                thinkingBlocks.push(item);\n                // We only need one thinking block\n                break;\n              }\n            }\n            if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n          }\n        }\n        \n        // If we found thinking blocks, add them to the first assistant message\n        if (thinkingBlocks.length > 0) {\n          // Calculate token count of original message\n          const originalTokenCount = tokenCounter(firstAssistantMsg);\n          \n          // Create a new content array with thinking blocks at the beginning\n          let newContent: any[];\n          \n          if (Array.isArray(firstAssistantMsg.content)) {\n            // Keep the original content (excluding any existing thinking blocks)\n            const originalContent = firstAssistantMsg.content.filter(item => \n              !(item && typeof item === 'object' && item.type === 'thinking'));\n            \n            newContent = [...thinkingBlocks, ...originalContent];\n          } else if (typeof firstAssistantMsg.content === 'string') {\n            newContent = [\n              ...thinkingBlocks,\n              { type: 'text', text: firstAssistantMsg.content }\n            ];\n          } else {\n            newContent = thinkingBlocks;\n          }\n          \n          // Create a new message with the updated content\n          const newMessage = new AIMessage({\n            content: newContent,\n            additional_kwargs: firstAssistantMsg.additional_kwargs,\n            response_metadata: firstAssistantMsg.response_metadata,\n          });\n          \n          // Calculate token count of new message\n          const newTokenCount = tokenCounter(newMessage);\n          \n          // Adjust current token count\n          currentTokenCount += (newTokenCount - originalTokenCount);\n          \n          // Replace the first assistant message\n          context[firstAssistantIndex] = newMessage;\n          \n          // If we've exceeded the token limit, we need to prune more messages\n          if (currentTokenCount > remainingContextTokens) {\n            // Remove messages from the end of the context until we're under the token limit\n            // But make sure to keep the first assistant message with thinking block\n            let i = context.length - 1;\n            while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n              const msgToRemove = context[i];\n              const msgTokenCount = tokenCounter(msgToRemove);\n              context.splice(i, 1);\n              currentTokenCount -= msgTokenCount;\n              i--;\n            }\n            \n            // Update remainingContextTokens to reflect the new token count\n            remainingContextTokens = maxContextTokens - currentTokenCount;\n          }\n        }\n      }\n    }\n  }\n  }\n\n  if (instructions && _messages.length > 0) {\n    context.push(_messages[0] as BaseMessage);\n    messages.shift();\n  }\n\n  const prunedMemory = messages;\n  summaryIndex = prunedMemory.length - 1;\n  remainingContextTokens -= currentTokenCount;\n\n  return {\n    summaryIndex,\n    remainingContextTokens,\n    context: context.reverse(),\n    messagesToRefine: prunedMemory,\n  };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n  return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n  const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n  let lastTurnStartIndex = factoryParams.startIndex;\n  let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n  \n  return function pruneMessages(params: PruneMessagesParams): {\n    context: BaseMessage[];\n    indexTokenCountMap: Record<string, number>;\n  } {\n    let currentUsage: UsageMetadata | undefined;\n    if (params.usageMetadata && (\n      checkValidNumber(params.usageMetadata.input_tokens)\n      || (\n        checkValidNumber(params.usageMetadata.input_token_details)\n        && (\n          checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n          || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n        )\n      )\n    ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n      currentUsage = calculateTotalTokens(params.usageMetadata);\n      totalTokens = currentUsage.total_tokens;\n    }\n\n    for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n      const message = params.messages[i];\n      if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n        indexTokenCountMap[i] = currentUsage.output_tokens;\n      } else if (indexTokenCountMap[i] === undefined) {\n        indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n        totalTokens += indexTokenCountMap[i];\n      }\n    }\n\n    // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n    // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n    // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n    // relative the manually counted tokens in `indexTokenCountMap`.\n    if (currentUsage) {\n      const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n      const ratio = currentUsage.total_tokens / totalIndexTokens;\n      for (const key in indexTokenCountMap) {\n        indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n      }\n    }\n\n    lastTurnStartIndex = params.messages.length;\n    if (totalTokens <= factoryParams.maxTokens) {\n      return { context: params.messages, indexTokenCountMap };\n    }\n\n    // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n    const { context } = getMessagesWithinTokenLimit({\n      maxContextTokens: factoryParams.maxTokens,\n      messages: params.messages,\n      indexTokenCountMap,\n      startOnMessageType: params.startOnMessageType,\n      thinkingEnabled: factoryParams.thinkingEnabled,\n      tokenCounter: factoryParams.tokenCounter,\n    });\n\n    return { context, indexTokenCountMap };\n  }\n}\n"],"names":["messages","AIMessage"],"mappings":";;;;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAMA,UAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAGA,UAAQ,CAAC,MAAM;AAClC,QAAA,OAAOA,UAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAIA,UAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAGA,UAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAAA,UAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;AAEzD,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAC5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,MAAM,GAAG,IAAIA,UAAQ,EAAE;AAC1B,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAIC,kBAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;;AAG9C,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;AAC9B,gCAAA,MAAM,aAAa,GAAG,YAAY,CAAC,WAAW,CAAC;AAC/C,gCAAA,OAAO,CAAC,MAAM,CAAC,CAAC,EAAE,CAAC,CAAC;gCACpB,iBAAiB,IAAI,aAAa;AAClC,gCAAA,CAAC,EAAE;;;AAIL,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;IAQvE,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;QACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;QACzCD,UAAQ,CAAC,KAAK,EAAE;;IAGlB,MAAM,YAAY,GAAGA,UAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}
+{"version":3,"file":"prune.cjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n  maxTokens: number;\n  startIndex: number;\n  tokenCounter: TokenCounter;\n  indexTokenCountMap: Record<string, number>;\n  thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n  messages: BaseMessage[];\n  usageMetadata?: Partial<UsageMetadata>;\n  startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n  const baseInputTokens = Number(usage.input_tokens) || 0;\n  const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n  const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n  \n  const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n  const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n  return {\n    input_tokens: totalInputTokens,\n    output_tokens: totalOutputTokens,\n    total_tokens: totalInputTokens + totalOutputTokens\n  };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n  messages: _messages,\n  maxContextTokens,\n  indexTokenCountMap,\n  startOnMessageType,\n  thinkingEnabled,\n  tokenCounter,\n}: {\n  messages: BaseMessage[];\n  maxContextTokens: number;\n  indexTokenCountMap: Record<string, number>;\n  startOnMessageType?: string;\n  thinkingEnabled?: boolean;\n  tokenCounter?: TokenCounter;\n}): {\n  context: BaseMessage[];\n  remainingContextTokens: number;\n  messagesToRefine: BaseMessage[];\n  summaryIndex: number;\n} {\n  // Every reply is primed with <|start|>assistant<|message|>, so we\n  // start with 3 tokens for the label after all messages have been counted.\n  let summaryIndex = -1;\n  let currentTokenCount = 3;\n  const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n  const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n  let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n  const messages = [..._messages];\n  let context: BaseMessage[] = [];\n\n  if (currentTokenCount < remainingContextTokens) {\n    let currentIndex = messages.length;\n    while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n      currentIndex--;\n      if (messages.length === 1 && instructions) {\n        break;\n      }\n      const poppedMessage = messages.pop();\n      if (!poppedMessage) continue;\n      \n      const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n      if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n        context.push(poppedMessage);\n        currentTokenCount += tokenCount;\n      } else {\n        messages.push(poppedMessage);\n        break;\n      }\n    }\n    \n  // Handle startOnMessageType requirement\n  if (startOnMessageType && context.length > 0) {\n    const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n    \n    if (requiredTypeIndex > 0) {\n      context = context.slice(requiredTypeIndex);\n    }\n  }\n  \n  // Add system message if it exists\n  if (instructions && _messages.length > 0) {\n    context.push(_messages[0] as BaseMessage);\n    messages.shift();\n  }\n  \n  // Handle thinking mode requirement for Anthropic\n  if (thinkingEnabled && context.length > 0 && tokenCounter) {\n    // Check if the latest message is an assistant message\n    const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';\n    \n    // Process only if we have an assistant message in the context\n    const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n    \n    if (firstAssistantIndex >= 0) {\n      const firstAssistantMsg = context[firstAssistantIndex];\n      \n      // Check if the first assistant message already has a thinking block\n      const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n        firstAssistantMsg.content.some(item => \n          item && typeof item === 'object' && item.type === 'thinking');\n      \n      // Only proceed if we need to add thinking blocks\n      if (!hasThinkingBlock) {\n        // Collect thinking blocks from pruned assistant messages, starting from the most recent\n        const thinkingBlocks: any[] = [];\n        \n        // Look through pruned messages for thinking blocks, starting from the end (most recent)\n        for (let i = messages.length - 1; i >= 0; i--) {\n          const msg = messages[i];\n          if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n            for (const item of msg.content) {\n              if (item && typeof item === 'object' && item.type === 'thinking') {\n                thinkingBlocks.push(item);\n                // We only need one thinking block\n                break;\n              }\n            }\n            if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n          }\n        }\n        \n        // If we found thinking blocks, add them to the first assistant message\n        if (thinkingBlocks.length > 0) {\n          // Calculate token count of original message\n          const originalTokenCount = tokenCounter(firstAssistantMsg);\n          \n          // Create a new content array with thinking blocks at the beginning\n          let newContent: any[];\n          \n          if (Array.isArray(firstAssistantMsg.content)) {\n            // Keep the original content (excluding any existing thinking blocks)\n            const originalContent = firstAssistantMsg.content.filter(item => \n              !(item && typeof item === 'object' && item.type === 'thinking'));\n            \n            newContent = [...thinkingBlocks, ...originalContent];\n          } else if (typeof firstAssistantMsg.content === 'string') {\n            newContent = [\n              ...thinkingBlocks,\n              { type: 'text', text: firstAssistantMsg.content }\n            ];\n          } else {\n            newContent = thinkingBlocks;\n          }\n          \n          // Create a new message with the updated content\n          const newMessage = new AIMessage({\n            content: newContent,\n            additional_kwargs: firstAssistantMsg.additional_kwargs,\n            response_metadata: firstAssistantMsg.response_metadata,\n          });\n          \n          // Calculate token count of new message\n          const newTokenCount = tokenCounter(newMessage);\n          \n          // Adjust current token count\n          currentTokenCount += (newTokenCount - originalTokenCount);\n          \n          // Replace the first assistant message\n          context[firstAssistantIndex] = newMessage;\n          \n          // If we've exceeded the token limit, we need to prune more messages\n          if (currentTokenCount > remainingContextTokens) {\n            // Build a map of tool call IDs to track AI <--> tool message correspondences\n            const toolCallIdMap = new Map<string, number>();\n            \n            // Identify tool call IDs in the context\n            for (let i = 0; i < context.length; i++) {\n              const msg = context[i];\n              \n              // Check for tool calls in AI messages\n              if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n                for (const item of msg.content) {\n                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n                    toolCallIdMap.set(item.id, i);\n                  }\n                }\n              }\n              \n              // Check for tool messages\n              if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {\n                toolCallIdMap.set(msg.tool_call_id, i);\n              }\n            }\n            \n            // Track which messages to remove\n            const indicesToRemove = new Set<number>();\n            \n            // Start removing messages from the end, but preserve AI <--> tool message correspondences\n            let i = context.length - 1;\n            while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n              const msgToRemove = context[i];\n              \n              // Check if this is a tool message or has tool calls\n              let canRemove = true;\n              \n              if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {\n                // If this is a tool message, check if we need to keep its corresponding AI message\n                const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);\n                if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {\n                  // We need to remove both the tool message and its corresponding AI message\n                  indicesToRemove.add(i);\n                  indicesToRemove.add(aiIndex);\n                  currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));\n                  canRemove = false;\n                }\n              } else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {\n                // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages\n                for (const item of msgToRemove.content) {\n                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n                    const toolIndex = toolCallIdMap.get(item.id as string);\n                    if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {\n                      // We need to remove both the AI message and its corresponding tool message\n                      indicesToRemove.add(i);\n                      indicesToRemove.add(toolIndex);\n                      currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));\n                      canRemove = false;\n                      break;\n                    }\n                  }\n                }\n              }\n              \n              // If we can remove this message individually\n              if (canRemove && !indicesToRemove.has(i)) {\n                indicesToRemove.add(i);\n                currentTokenCount -= tokenCounter(msgToRemove);\n              }\n              \n              i--;\n            }\n            \n            // Remove messages in reverse order to avoid index shifting\n            const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);\n            for (const index of sortedIndices) {\n              context.splice(index, 1);\n            }\n            \n            // Update remainingContextTokens to reflect the new token count\n            remainingContextTokens = maxContextTokens - currentTokenCount;\n          }\n        }\n      }\n    }\n    \n    // If the latest message is an assistant message, ensure an assistant message appears early in the context\n    // but maintain system message precedence\n    if (latestMessageIsAssistant && context.length > 0) {\n      // Find the first assistant message in the context\n      const assistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n      if (assistantIndex > 0) {\n        // Check if there's a system message at the beginning\n        const hasSystemFirst = context[0].getType() === 'system';\n        \n        // Move the assistant message to the appropriate position\n        const assistantMsg = context[assistantIndex];\n        context.splice(assistantIndex, 1);\n        \n        if (hasSystemFirst) {\n          // Insert after the system message\n          context.splice(1, 0, assistantMsg);\n        } else {\n          // Insert at the beginning if no system message\n          context.unshift(assistantMsg);\n        }\n      }\n    }\n  }\n  }\n\n  const prunedMemory = messages;\n  summaryIndex = prunedMemory.length - 1;\n  remainingContextTokens -= currentTokenCount;\n\n  return {\n    summaryIndex,\n    remainingContextTokens,\n    context: context.reverse(),\n    messagesToRefine: prunedMemory,\n  };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n  return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n  const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n  let lastTurnStartIndex = factoryParams.startIndex;\n  let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n  \n  return function pruneMessages(params: PruneMessagesParams): {\n    context: BaseMessage[];\n    indexTokenCountMap: Record<string, number>;\n  } {\n    let currentUsage: UsageMetadata | undefined;\n    if (params.usageMetadata && (\n      checkValidNumber(params.usageMetadata.input_tokens)\n      || (\n        checkValidNumber(params.usageMetadata.input_token_details)\n        && (\n          checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n          || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n        )\n      )\n    ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n      currentUsage = calculateTotalTokens(params.usageMetadata);\n      totalTokens = currentUsage.total_tokens;\n    }\n\n    for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n      const message = params.messages[i];\n      if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n        indexTokenCountMap[i] = currentUsage.output_tokens;\n      } else if (indexTokenCountMap[i] === undefined) {\n        indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n        totalTokens += indexTokenCountMap[i];\n      }\n    }\n\n    // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n    // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n    // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n    // relative the manually counted tokens in `indexTokenCountMap`.\n    if (currentUsage) {\n      const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n      const ratio = currentUsage.total_tokens / totalIndexTokens;\n      for (const key in indexTokenCountMap) {\n        indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n      }\n    }\n\n    lastTurnStartIndex = params.messages.length;\n    if (totalTokens <= factoryParams.maxTokens) {\n      return { context: params.messages, indexTokenCountMap };\n    }\n\n    // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n    const { context } = getMessagesWithinTokenLimit({\n      maxContextTokens: factoryParams.maxTokens,\n      messages: params.messages,\n      indexTokenCountMap,\n      startOnMessageType: params.startOnMessageType,\n      thinkingEnabled: factoryParams.thinkingEnabled,\n      tokenCounter: factoryParams.tokenCounter,\n    });\n\n    return { context, indexTokenCountMap };\n  }\n}\n"],"names":["messages","AIMessage"],"mappings":";;;;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAMA,UAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAGA,UAAQ,CAAC,MAAM;AAClC,QAAA,OAAOA,UAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAIA,UAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAGA,UAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAAA,UAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;YACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;YACzCA,UAAQ,CAAC,KAAK,EAAE;;;QAIlB,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;YAEzD,MAAM,wBAAwB,GAAG,SAAS,CAAC,MAAM,GAAG,CAAC,IAAI,SAAS,CAAC,SAAS,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,IAAI;;AAG3G,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAE5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,IAAI,CAAC,GAAGA,UAAQ,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,EAAE;AAC7C,wBAAA,MAAM,GAAG,GAAGA,UAAQ,CAAC,CAAC,CAAC;AACvB,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAIC,kBAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;AAE9C,4BAAA,MAAM,aAAa,GAAG,IAAI,GAAG,EAAkB;;AAG/C,4BAAA,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;AACvC,gCAAA,MAAM,GAAG,GAAG,OAAO,CAAC,CAAC,CAAC;;AAGtB,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,oCAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,CAAC,CAAC;;;;;AAMnC,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,GAAG,IAAI,OAAO,GAAG,CAAC,YAAY,KAAK,QAAQ,EAAE;oCAC7F,aAAa,CAAC,GAAG,CAAC,GAAG,CAAC,YAAY,EAAE,CAAC,CAAC;;;;AAK1C,4BAAA,MAAM,eAAe,GAAG,IAAI,GAAG,EAAU;;AAGzC,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;;gCAG9B,IAAI,SAAS,GAAG,IAAI;AAEpB,gCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,WAAW,IAAI,OAAO,WAAW,CAAC,YAAY,KAAK,QAAQ,EAAE;;oCAErH,MAAM,OAAO,GAAG,aAAa,CAAC,GAAG,CAAC,WAAW,CAAC,YAAY,CAAC;AAC3D,oCAAA,IAAI,OAAO,KAAK,SAAS,IAAI,OAAO,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;;AAE3E,wCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,wCAAA,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC;AAC5B,wCAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC;wCACjF,SAAS,GAAG,KAAK;;;AAEd,qCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,CAAC,EAAE;;AAE/E,oCAAA,KAAK,MAAM,IAAI,IAAI,WAAW,CAAC,OAAO,EAAE;AACtC,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,MAAM,SAAS,GAAG,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAY,CAAC;AACtD,4CAAA,IAAI,SAAS,KAAK,SAAS,IAAI,SAAS,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE;;AAEjF,gDAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,gDAAA,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC;AAC9B,gDAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC,CAAC;gDACnF,SAAS,GAAG,KAAK;gDACjB;;;;;;gCAOR,IAAI,SAAS,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE;AACxC,oCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,oCAAA,iBAAiB,IAAI,YAAY,CAAC,WAAW,CAAC;;AAGhD,gCAAA,CAAC,EAAE;;;4BAIL,MAAM,aAAa,GAAG,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;AACvE,4BAAA,KAAK,MAAM,KAAK,IAAI,aAAa,EAAE;AACjC,gCAAA,OAAO,CAAC,MAAM,CAAC,KAAK,EAAE,CAAC,CAAC;;;AAI1B,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;YAQrE,IAAI,wBAAwB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;;AAElD,gBAAA,MAAM,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AACvE,gBAAA,IAAI,cAAc,GAAG,CAAC,EAAE;;oBAEtB,MAAM,cAAc,GAAG,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,QAAQ;;AAGxD,oBAAA,MAAM,YAAY,GAAG,OAAO,CAAC,cAAc,CAAC;AAC5C,oBAAA,OAAO,CAAC,MAAM,CAAC,cAAc,EAAE,CAAC,CAAC;oBAEjC,IAAI,cAAc,EAAE;;wBAElB,OAAO,CAAC,MAAM,CAAC,CAAC,EAAE,CAAC,EAAE,YAAY,CAAC;;yBAC7B;;AAEL,wBAAA,OAAO,CAAC,OAAO,CAAC,YAAY,CAAC;;;;;;IAOrC,MAAM,YAAY,GAAGD,UAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}

package/dist/esm/messages/prune.mjs

@@ -62,8 +62,15 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
                 context = context.slice(requiredTypeIndex);
             }
         }
+        // Add system message if it exists
+        if (instructions && _messages.length > 0) {
+            context.push(_messages[0]);
+            messages.shift();
+        }
         // Handle thinking mode requirement for Anthropic
         if (thinkingEnabled && context.length > 0 && tokenCounter) {
+            // Check if the latest message is an assistant message
+            const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';
             // Process only if we have an assistant message in the context
             const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');
             if (firstAssistantIndex >= 0) {
@@ -73,10 +80,11 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
                     firstAssistantMsg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');
                 // Only proceed if we need to add thinking blocks
                 if (!hasThinkingBlock) {
-                    // Collect thinking blocks from pruned assistant messages
+                    // Collect thinking blocks from pruned assistant messages, starting from the most recent
                     const thinkingBlocks = [];
-                    // Look through pruned messages for thinking blocks
-                    for (const msg of messages) {
+                    // Look through pruned messages for thinking blocks, starting from the end (most recent)
+                    for (let i = messages.length - 1; i >= 0; i--) {
+                        const msg = messages[i];
                         if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
                             for (const item of msg.content) {
                                 if (item && typeof item === 'object' && item.type === 'thinking') {
@@ -123,28 +131,100 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
                         context[firstAssistantIndex] = newMessage;
                         // If we've exceeded the token limit, we need to prune more messages
                         if (currentTokenCount > remainingContextTokens) {
-                            // Remove messages from the end of the context until we're under the token limit
-                            // But make sure to keep the first assistant message with thinking block
+                            // Build a map of tool call IDs to track AI <--> tool message correspondences
+                            const toolCallIdMap = new Map();
+                            // Identify tool call IDs in the context
+                            for (let i = 0; i < context.length; i++) {
+                                const msg = context[i];
+                                // Check for tool calls in AI messages
+                                if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
+                                    for (const item of msg.content) {
+                                        if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                                            toolCallIdMap.set(item.id, i);
+                                        }
+                                    }
+                                }
+                                // Check for tool messages
+                                if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {
+                                    toolCallIdMap.set(msg.tool_call_id, i);
+                                }
+                            }
+                            // Track which messages to remove
+                            const indicesToRemove = new Set();
+                            // Start removing messages from the end, but preserve AI <--> tool message correspondences
                             let i = context.length - 1;
                             while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {
                                 const msgToRemove = context[i];
-                                const msgTokenCount = tokenCounter(msgToRemove);
-                                context.splice(i, 1);
-                                currentTokenCount -= msgTokenCount;
+                                // Check if this is a tool message or has tool calls
+                                let canRemove = true;
+                                if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {
+                                    // If this is a tool message, check if we need to keep its corresponding AI message
+                                    const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);
+                                    if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {
+                                        // We need to remove both the tool message and its corresponding AI message
+                                        indicesToRemove.add(i);
+                                        indicesToRemove.add(aiIndex);
+                                        currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));
+                                        canRemove = false;
+                                    }
+                                }
+                                else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {
+                                    // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages
+                                    for (const item of msgToRemove.content) {
+                                        if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                                            const toolIndex = toolCallIdMap.get(item.id);
+                                            if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {
+                                                // We need to remove both the AI message and its corresponding tool message
+                                                indicesToRemove.add(i);
+                                                indicesToRemove.add(toolIndex);
+                                                currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));
+                                                canRemove = false;
+                                                break;
+                                            }
+                                        }
+                                    }
+                                }
+                                // If we can remove this message individually
+                                if (canRemove && !indicesToRemove.has(i)) {
+                                    indicesToRemove.add(i);
+                                    currentTokenCount -= tokenCounter(msgToRemove);
+                                }
                                 i--;
                             }
+                            // Remove messages in reverse order to avoid index shifting
+                            const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);
+                            for (const index of sortedIndices) {
+                                context.splice(index, 1);
+                            }
                             // Update remainingContextTokens to reflect the new token count
                             remainingContextTokens = maxContextTokens - currentTokenCount;
                         }
                     }
                 }
             }
+            // If the latest message is an assistant message, ensure an assistant message appears early in the context
+            // but maintain system message precedence
+            if (latestMessageIsAssistant && context.length > 0) {
+                // Find the first assistant message in the context
+                const assistantIndex = context.findIndex(msg => msg.getType() === 'ai');
+                if (assistantIndex > 0) {
+                    // Check if there's a system message at the beginning
+                    const hasSystemFirst = context[0].getType() === 'system';
+                    // Move the assistant message to the appropriate position
+                    const assistantMsg = context[assistantIndex];
+                    context.splice(assistantIndex, 1);
+                    if (hasSystemFirst) {
+                        // Insert after the system message
+                        context.splice(1, 0, assistantMsg);
+                    }
+                    else {
+                        // Insert at the beginning if no system message
+                        context.unshift(assistantMsg);
+                    }
+                }
+            }
         }
     }
-    if (instructions && _messages.length > 0) {
-        context.push(_messages[0]);
-        messages.shift();
-    }
     const prunedMemory = messages;
     summaryIndex = prunedMemory.length - 1;
     remainingContextTokens -= currentTokenCount;

package/dist/esm/messages/prune.mjs.map

@@ -1 +1 @@
-{"version":3,"file":"prune.mjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n  maxTokens: number;\n  startIndex: number;\n  tokenCounter: TokenCounter;\n  indexTokenCountMap: Record<string, number>;\n  thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n  messages: BaseMessage[];\n  usageMetadata?: Partial<UsageMetadata>;\n  startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n  const baseInputTokens = Number(usage.input_tokens) || 0;\n  const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n  const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n  \n  const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n  const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n  return {\n    input_tokens: totalInputTokens,\n    output_tokens: totalOutputTokens,\n    total_tokens: totalInputTokens + totalOutputTokens\n  };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n  messages: _messages,\n  maxContextTokens,\n  indexTokenCountMap,\n  startOnMessageType,\n  thinkingEnabled,\n  tokenCounter,\n}: {\n  messages: BaseMessage[];\n  maxContextTokens: number;\n  indexTokenCountMap: Record<string, number>;\n  startOnMessageType?: string;\n  thinkingEnabled?: boolean;\n  tokenCounter?: TokenCounter;\n}): {\n  context: BaseMessage[];\n  remainingContextTokens: number;\n  messagesToRefine: BaseMessage[];\n  summaryIndex: number;\n} {\n  // Every reply is primed with <|start|>assistant<|message|>, so we\n  // start with 3 tokens for the label after all messages have been counted.\n  let summaryIndex = -1;\n  let currentTokenCount = 3;\n  const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n  const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n  let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n  const messages = [..._messages];\n  let context: BaseMessage[] = [];\n\n  if (currentTokenCount < remainingContextTokens) {\n    let currentIndex = messages.length;\n    while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n      currentIndex--;\n      if (messages.length === 1 && instructions) {\n        break;\n      }\n      const poppedMessage = messages.pop();\n      if (!poppedMessage) continue;\n      \n      const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n      if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n        context.push(poppedMessage);\n        currentTokenCount += tokenCount;\n      } else {\n        messages.push(poppedMessage);\n        break;\n      }\n    }\n    \n  // Handle startOnMessageType requirement\n  if (startOnMessageType && context.length > 0) {\n    const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n    \n    if (requiredTypeIndex > 0) {\n      context = context.slice(requiredTypeIndex);\n    }\n  }\n  \n  // Handle thinking mode requirement for Anthropic\n  if (thinkingEnabled && context.length > 0 && tokenCounter) {\n    // Process only if we have an assistant message in the context\n    const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n    if (firstAssistantIndex >= 0) {\n      const firstAssistantMsg = context[firstAssistantIndex];\n      \n      // Check if the first assistant message already has a thinking block\n      const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n        firstAssistantMsg.content.some(item => \n          item && typeof item === 'object' && item.type === 'thinking');\n      \n      // Only proceed if we need to add thinking blocks\n      if (!hasThinkingBlock) {\n        // Collect thinking blocks from pruned assistant messages\n        const thinkingBlocks: any[] = [];\n        \n        // Look through pruned messages for thinking blocks\n        for (const msg of messages) {\n          if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n            for (const item of msg.content) {\n              if (item && typeof item === 'object' && item.type === 'thinking') {\n                thinkingBlocks.push(item);\n                // We only need one thinking block\n                break;\n              }\n            }\n            if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n          }\n        }\n        \n        // If we found thinking blocks, add them to the first assistant message\n        if (thinkingBlocks.length > 0) {\n          // Calculate token count of original message\n          const originalTokenCount = tokenCounter(firstAssistantMsg);\n          \n          // Create a new content array with thinking blocks at the beginning\n          let newContent: any[];\n          \n          if (Array.isArray(firstAssistantMsg.content)) {\n            // Keep the original content (excluding any existing thinking blocks)\n            const originalContent = firstAssistantMsg.content.filter(item => \n              !(item && typeof item === 'object' && item.type === 'thinking'));\n            \n            newContent = [...thinkingBlocks, ...originalContent];\n          } else if (typeof firstAssistantMsg.content === 'string') {\n            newContent = [\n              ...thinkingBlocks,\n              { type: 'text', text: firstAssistantMsg.content }\n            ];\n          } else {\n            newContent = thinkingBlocks;\n          }\n          \n          // Create a new message with the updated content\n          const newMessage = new AIMessage({\n            content: newContent,\n            additional_kwargs: firstAssistantMsg.additional_kwargs,\n            response_metadata: firstAssistantMsg.response_metadata,\n          });\n          \n          // Calculate token count of new message\n          const newTokenCount = tokenCounter(newMessage);\n          \n          // Adjust current token count\n          currentTokenCount += (newTokenCount - originalTokenCount);\n          \n          // Replace the first assistant message\n          context[firstAssistantIndex] = newMessage;\n          \n          // If we've exceeded the token limit, we need to prune more messages\n          if (currentTokenCount > remainingContextTokens) {\n            // Remove messages from the end of the context until we're under the token limit\n            // But make sure to keep the first assistant message with thinking block\n            let i = context.length - 1;\n            while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n              const msgToRemove = context[i];\n              const msgTokenCount = tokenCounter(msgToRemove);\n              context.splice(i, 1);\n              currentTokenCount -= msgTokenCount;\n              i--;\n            }\n            \n            // Update remainingContextTokens to reflect the new token count\n            remainingContextTokens = maxContextTokens - currentTokenCount;\n          }\n        }\n      }\n    }\n  }\n  }\n\n  if (instructions && _messages.length > 0) {\n    context.push(_messages[0] as BaseMessage);\n    messages.shift();\n  }\n\n  const prunedMemory = messages;\n  summaryIndex = prunedMemory.length - 1;\n  remainingContextTokens -= currentTokenCount;\n\n  return {\n    summaryIndex,\n    remainingContextTokens,\n    context: context.reverse(),\n    messagesToRefine: prunedMemory,\n  };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n  return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n  const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n  let lastTurnStartIndex = factoryParams.startIndex;\n  let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n  \n  return function pruneMessages(params: PruneMessagesParams): {\n    context: BaseMessage[];\n    indexTokenCountMap: Record<string, number>;\n  } {\n    let currentUsage: UsageMetadata | undefined;\n    if (params.usageMetadata && (\n      checkValidNumber(params.usageMetadata.input_tokens)\n      || (\n        checkValidNumber(params.usageMetadata.input_token_details)\n        && (\n          checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n          || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n        )\n      )\n    ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n      currentUsage = calculateTotalTokens(params.usageMetadata);\n      totalTokens = currentUsage.total_tokens;\n    }\n\n    for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n      const message = params.messages[i];\n      if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n        indexTokenCountMap[i] = currentUsage.output_tokens;\n      } else if (indexTokenCountMap[i] === undefined) {\n        indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n        totalTokens += indexTokenCountMap[i];\n      }\n    }\n\n    // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n    // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n    // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n    // relative the manually counted tokens in `indexTokenCountMap`.\n    if (currentUsage) {\n      const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n      const ratio = currentUsage.total_tokens / totalIndexTokens;\n      for (const key in indexTokenCountMap) {\n        indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n      }\n    }\n\n    lastTurnStartIndex = params.messages.length;\n    if (totalTokens <= factoryParams.maxTokens) {\n      return { context: params.messages, indexTokenCountMap };\n    }\n\n    // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n    const { context } = getMessagesWithinTokenLimit({\n      maxContextTokens: factoryParams.maxTokens,\n      messages: params.messages,\n      indexTokenCountMap,\n      startOnMessageType: params.startOnMessageType,\n      thinkingEnabled: factoryParams.thinkingEnabled,\n      tokenCounter: factoryParams.tokenCounter,\n    });\n\n    return { context, indexTokenCountMap };\n  }\n}\n"],"names":[],"mappings":";;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAM,QAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAG,QAAQ,CAAC,MAAM;AAClC,QAAA,OAAO,QAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAI,QAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAG,QAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAA,QAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;AAEzD,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAC5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,MAAM,GAAG,IAAI,QAAQ,EAAE;AAC1B,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAI,SAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;;AAG9C,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;AAC9B,gCAAA,MAAM,aAAa,GAAG,YAAY,CAAC,WAAW,CAAC;AAC/C,gCAAA,OAAO,CAAC,MAAM,CAAC,CAAC,EAAE,CAAC,CAAC;gCACpB,iBAAiB,IAAI,aAAa;AAClC,gCAAA,CAAC,EAAE;;;AAIL,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;IAQvE,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;QACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;QACzC,QAAQ,CAAC,KAAK,EAAE;;IAGlB,MAAM,YAAY,GAAG,QAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}
+{"version":3,"file":"prune.mjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n  maxTokens: number;\n  startIndex: number;\n  tokenCounter: TokenCounter;\n  indexTokenCountMap: Record<string, number>;\n  thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n  messages: BaseMessage[];\n  usageMetadata?: Partial<UsageMetadata>;\n  startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n  const baseInputTokens = Number(usage.input_tokens) || 0;\n  const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n  const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n  \n  const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n  const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n  return {\n    input_tokens: totalInputTokens,\n    output_tokens: totalOutputTokens,\n    total_tokens: totalInputTokens + totalOutputTokens\n  };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n  messages: _messages,\n  maxContextTokens,\n  indexTokenCountMap,\n  startOnMessageType,\n  thinkingEnabled,\n  tokenCounter,\n}: {\n  messages: BaseMessage[];\n  maxContextTokens: number;\n  indexTokenCountMap: Record<string, number>;\n  startOnMessageType?: string;\n  thinkingEnabled?: boolean;\n  tokenCounter?: TokenCounter;\n}): {\n  context: BaseMessage[];\n  remainingContextTokens: number;\n  messagesToRefine: BaseMessage[];\n  summaryIndex: number;\n} {\n  // Every reply is primed with <|start|>assistant<|message|>, so we\n  // start with 3 tokens for the label after all messages have been counted.\n  let summaryIndex = -1;\n  let currentTokenCount = 3;\n  const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n  const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n  let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n  const messages = [..._messages];\n  let context: BaseMessage[] = [];\n\n  if (currentTokenCount < remainingContextTokens) {\n    let currentIndex = messages.length;\n    while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n      currentIndex--;\n      if (messages.length === 1 && instructions) {\n        break;\n      }\n      const poppedMessage = messages.pop();\n      if (!poppedMessage) continue;\n      \n      const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n      if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n        context.push(poppedMessage);\n        currentTokenCount += tokenCount;\n      } else {\n        messages.push(poppedMessage);\n        break;\n      }\n    }\n    \n  // Handle startOnMessageType requirement\n  if (startOnMessageType && context.length > 0) {\n    const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n    \n    if (requiredTypeIndex > 0) {\n      context = context.slice(requiredTypeIndex);\n    }\n  }\n  \n  // Add system message if it exists\n  if (instructions && _messages.length > 0) {\n    context.push(_messages[0] as BaseMessage);\n    messages.shift();\n  }\n  \n  // Handle thinking mode requirement for Anthropic\n  if (thinkingEnabled && context.length > 0 && tokenCounter) {\n    // Check if the latest message is an assistant message\n    const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';\n    \n    // Process only if we have an assistant message in the context\n    const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n    \n    if (firstAssistantIndex >= 0) {\n      const firstAssistantMsg = context[firstAssistantIndex];\n      \n      // Check if the first assistant message already has a thinking block\n      const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n        firstAssistantMsg.content.some(item => \n          item && typeof item === 'object' && item.type === 'thinking');\n      \n      // Only proceed if we need to add thinking blocks\n      if (!hasThinkingBlock) {\n        // Collect thinking blocks from pruned assistant messages, starting from the most recent\n        const thinkingBlocks: any[] = [];\n        \n        // Look through pruned messages for thinking blocks, starting from the end (most recent)\n        for (let i = messages.length - 1; i >= 0; i--) {\n          const msg = messages[i];\n          if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n            for (const item of msg.content) {\n              if (item && typeof item === 'object' && item.type === 'thinking') {\n                thinkingBlocks.push(item);\n                // We only need one thinking block\n                break;\n              }\n            }\n            if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n          }\n        }\n        \n        // If we found thinking blocks, add them to the first assistant message\n        if (thinkingBlocks.length > 0) {\n          // Calculate token count of original message\n          const originalTokenCount = tokenCounter(firstAssistantMsg);\n          \n          // Create a new content array with thinking blocks at the beginning\n          let newContent: any[];\n          \n          if (Array.isArray(firstAssistantMsg.content)) {\n            // Keep the original content (excluding any existing thinking blocks)\n            const originalContent = firstAssistantMsg.content.filter(item => \n              !(item && typeof item === 'object' && item.type === 'thinking'));\n            \n            newContent = [...thinkingBlocks, ...originalContent];\n          } else if (typeof firstAssistantMsg.content === 'string') {\n            newContent = [\n              ...thinkingBlocks,\n              { type: 'text', text: firstAssistantMsg.content }\n            ];\n          } else {\n            newContent = thinkingBlocks;\n          }\n          \n          // Create a new message with the updated content\n          const newMessage = new AIMessage({\n            content: newContent,\n            additional_kwargs: firstAssistantMsg.additional_kwargs,\n            response_metadata: firstAssistantMsg.response_metadata,\n          });\n          \n          // Calculate token count of new message\n          const newTokenCount = tokenCounter(newMessage);\n          \n          // Adjust current token count\n          currentTokenCount += (newTokenCount - originalTokenCount);\n          \n          // Replace the first assistant message\n          context[firstAssistantIndex] = newMessage;\n          \n          // If we've exceeded the token limit, we need to prune more messages\n          if (currentTokenCount > remainingContextTokens) {\n            // Build a map of tool call IDs to track AI <--> tool message correspondences\n            const toolCallIdMap = new Map<string, number>();\n            \n            // Identify tool call IDs in the context\n            for (let i = 0; i < context.length; i++) {\n              const msg = context[i];\n              \n              // Check for tool calls in AI messages\n              if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n                for (const item of msg.content) {\n                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n                    toolCallIdMap.set(item.id, i);\n                  }\n                }\n              }\n              \n              // Check for tool messages\n              if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {\n                toolCallIdMap.set(msg.tool_call_id, i);\n              }\n            }\n            \n            // Track which messages to remove\n            const indicesToRemove = new Set<number>();\n            \n            // Start removing messages from the end, but preserve AI <--> tool message correspondences\n            let i = context.length - 1;\n            while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n              const msgToRemove = context[i];\n              \n              // Check if this is a tool message or has tool calls\n              let canRemove = true;\n              \n              if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {\n                // If this is a tool message, check if we need to keep its corresponding AI message\n                const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);\n                if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {\n                  // We need to remove both the tool message and its corresponding AI message\n                  indicesToRemove.add(i);\n                  indicesToRemove.add(aiIndex);\n                  currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));\n                  canRemove = false;\n                }\n              } else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {\n                // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages\n                for (const item of msgToRemove.content) {\n                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n                    const toolIndex = toolCallIdMap.get(item.id as string);\n                    if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {\n                      // We need to remove both the AI message and its corresponding tool message\n                      indicesToRemove.add(i);\n                      indicesToRemove.add(toolIndex);\n                      currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));\n                      canRemove = false;\n                      break;\n                    }\n                  }\n                }\n              }\n              \n              // If we can remove this message individually\n              if (canRemove && !indicesToRemove.has(i)) {\n                indicesToRemove.add(i);\n                currentTokenCount -= tokenCounter(msgToRemove);\n              }\n              \n              i--;\n            }\n            \n            // Remove messages in reverse order to avoid index shifting\n            const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);\n            for (const index of sortedIndices) {\n              context.splice(index, 1);\n            }\n            \n            // Update remainingContextTokens to reflect the new token count\n            remainingContextTokens = maxContextTokens - currentTokenCount;\n          }\n        }\n      }\n    }\n    \n    // If the latest message is an assistant message, ensure an assistant message appears early in the context\n    // but maintain system message precedence\n    if (latestMessageIsAssistant && context.length > 0) {\n      // Find the first assistant message in the context\n      const assistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n      if (assistantIndex > 0) {\n        // Check if there's a system message at the beginning\n        const hasSystemFirst = context[0].getType() === 'system';\n        \n        // Move the assistant message to the appropriate position\n        const assistantMsg = context[assistantIndex];\n        context.splice(assistantIndex, 1);\n        \n        if (hasSystemFirst) {\n          // Insert after the system message\n          context.splice(1, 0, assistantMsg);\n        } else {\n          // Insert at the beginning if no system message\n          context.unshift(assistantMsg);\n        }\n      }\n    }\n  }\n  }\n\n  const prunedMemory = messages;\n  summaryIndex = prunedMemory.length - 1;\n  remainingContextTokens -= currentTokenCount;\n\n  return {\n    summaryIndex,\n    remainingContextTokens,\n    context: context.reverse(),\n    messagesToRefine: prunedMemory,\n  };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n  return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n  const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n  let lastTurnStartIndex = factoryParams.startIndex;\n  let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n  \n  return function pruneMessages(params: PruneMessagesParams): {\n    context: BaseMessage[];\n    indexTokenCountMap: Record<string, number>;\n  } {\n    let currentUsage: UsageMetadata | undefined;\n    if (params.usageMetadata && (\n      checkValidNumber(params.usageMetadata.input_tokens)\n      || (\n        checkValidNumber(params.usageMetadata.input_token_details)\n        && (\n          checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n          || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n        )\n      )\n    ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n      currentUsage = calculateTotalTokens(params.usageMetadata);\n      totalTokens = currentUsage.total_tokens;\n    }\n\n    for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n      const message = params.messages[i];\n      if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n        indexTokenCountMap[i] = currentUsage.output_tokens;\n      } else if (indexTokenCountMap[i] === undefined) {\n        indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n        totalTokens += indexTokenCountMap[i];\n      }\n    }\n\n    // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n    // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n    // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n    // relative the manually counted tokens in `indexTokenCountMap`.\n    if (currentUsage) {\n      const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n      const ratio = currentUsage.total_tokens / totalIndexTokens;\n      for (const key in indexTokenCountMap) {\n        indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n      }\n    }\n\n    lastTurnStartIndex = params.messages.length;\n    if (totalTokens <= factoryParams.maxTokens) {\n      return { context: params.messages, indexTokenCountMap };\n    }\n\n    // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n    const { context } = getMessagesWithinTokenLimit({\n      maxContextTokens: factoryParams.maxTokens,\n      messages: params.messages,\n      indexTokenCountMap,\n      startOnMessageType: params.startOnMessageType,\n      thinkingEnabled: factoryParams.thinkingEnabled,\n      tokenCounter: factoryParams.tokenCounter,\n    });\n\n    return { context, indexTokenCountMap };\n  }\n}\n"],"names":[],"mappings":";;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAM,QAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAG,QAAQ,CAAC,MAAM;AAClC,QAAA,OAAO,QAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAI,QAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAG,QAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAA,QAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;YACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;YACzC,QAAQ,CAAC,KAAK,EAAE;;;QAIlB,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;YAEzD,MAAM,wBAAwB,GAAG,SAAS,CAAC,MAAM,GAAG,CAAC,IAAI,SAAS,CAAC,SAAS,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,IAAI;;AAG3G,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAE5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,IAAI,CAAC,GAAG,QAAQ,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,EAAE;AAC7C,wBAAA,MAAM,GAAG,GAAG,QAAQ,CAAC,CAAC,CAAC;AACvB,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAI,SAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;AAE9C,4BAAA,MAAM,aAAa,GAAG,IAAI,GAAG,EAAkB;;AAG/C,4BAAA,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;AACvC,gCAAA,MAAM,GAAG,GAAG,OAAO,CAAC,CAAC,CAAC;;AAGtB,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,oCAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,CAAC,CAAC;;;;;AAMnC,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,GAAG,IAAI,OAAO,GAAG,CAAC,YAAY,KAAK,QAAQ,EAAE;oCAC7F,aAAa,CAAC,GAAG,CAAC,GAAG,CAAC,YAAY,EAAE,CAAC,CAAC;;;;AAK1C,4BAAA,MAAM,eAAe,GAAG,IAAI,GAAG,EAAU;;AAGzC,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;;gCAG9B,IAAI,SAAS,GAAG,IAAI;AAEpB,gCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,WAAW,IAAI,OAAO,WAAW,CAAC,YAAY,KAAK,QAAQ,EAAE;;oCAErH,MAAM,OAAO,GAAG,aAAa,CAAC,GAAG,CAAC,WAAW,CAAC,YAAY,CAAC;AAC3D,oCAAA,IAAI,OAAO,KAAK,SAAS,IAAI,OAAO,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;;AAE3E,wCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,wCAAA,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC;AAC5B,wCAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC;wCACjF,SAAS,GAAG,KAAK;;;AAEd,qCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,CAAC,EAAE;;AAE/E,oCAAA,KAAK,MAAM,IAAI,IAAI,WAAW,CAAC,OAAO,EAAE;AACtC,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,MAAM,SAAS,GAAG,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAY,CAAC;AACtD,4CAAA,IAAI,SAAS,KAAK,SAAS,IAAI,SAAS,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE;;AAEjF,gDAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,gDAAA,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC;AAC9B,gDAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC,CAAC;gDACnF,SAAS,GAAG,KAAK;gDACjB;;;;;;gCAOR,IAAI,SAAS,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE;AACxC,oCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,oCAAA,iBAAiB,IAAI,YAAY,CAAC,WAAW,CAAC;;AAGhD,gCAAA,CAAC,EAAE;;;4BAIL,MAAM,aAAa,GAAG,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;AACvE,4BAAA,KAAK,MAAM,KAAK,IAAI,aAAa,EAAE;AACjC,gCAAA,OAAO,CAAC,MAAM,CAAC,KAAK,EAAE,CAAC,CAAC;;;AAI1B,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;YAQrE,IAAI,wBAAwB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;;AAElD,gBAAA,MAAM,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AACvE,gBAAA,IAAI,cAAc,GAAG,CAAC,EAAE;;oBAEtB,MAAM,cAAc,GAAG,OAAO,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,QAAQ;;AAGxD,oBAAA,MAAM,YAAY,GAAG,OAAO,CAAC,cAAc,CAAC;AAC5C,oBAAA,OAAO,CAAC,MAAM,CAAC,cAAc,EAAE,CAAC,CAAC;oBAEjC,IAAI,cAAc,EAAE;;wBAElB,OAAO,CAAC,MAAM,CAAC,CAAC,EAAE,CAAC,EAAE,YAAY,CAAC;;yBAC7B;;AAEL,wBAAA,OAAO,CAAC,OAAO,CAAC,YAAY,CAAC;;;;;;IAOrC,MAAM,YAAY,GAAG,QAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@librechat/agents",
-  "version": "2.2.9",
+  "version": "2.3.0",
   "main": "./dist/cjs/main.cjs",
   "module": "./dist/esm/main.mjs",
   "types": "./dist/types/index.d.ts",

package/src/messages/prune.ts

@@ -102,10 +102,20 @@ function getMessagesWithinTokenLimit({
     }
   }
   
+  // Add system message if it exists
+  if (instructions && _messages.length > 0) {
+    context.push(_messages[0] as BaseMessage);
+    messages.shift();
+  }
+  
   // Handle thinking mode requirement for Anthropic
   if (thinkingEnabled && context.length > 0 && tokenCounter) {
+    // Check if the latest message is an assistant message
+    const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';
+    
     // Process only if we have an assistant message in the context
     const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');
+    
     if (firstAssistantIndex >= 0) {
       const firstAssistantMsg = context[firstAssistantIndex];
       
@@ -116,11 +126,12 @@ function getMessagesWithinTokenLimit({
       
       // Only proceed if we need to add thinking blocks
       if (!hasThinkingBlock) {
-        // Collect thinking blocks from pruned assistant messages
+        // Collect thinking blocks from pruned assistant messages, starting from the most recent
         const thinkingBlocks: any[] = [];
         
-        // Look through pruned messages for thinking blocks
-        for (const msg of messages) {
+        // Look through pruned messages for thinking blocks, starting from the end (most recent)
+        for (let i = messages.length - 1; i >= 0; i--) {
+          const msg = messages[i];
           if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
             for (const item of msg.content) {
               if (item && typeof item === 'object' && item.type === 'thinking') {
@@ -174,31 +185,113 @@ function getMessagesWithinTokenLimit({
           
           // If we've exceeded the token limit, we need to prune more messages
           if (currentTokenCount > remainingContextTokens) {
-            // Remove messages from the end of the context until we're under the token limit
-            // But make sure to keep the first assistant message with thinking block
+            // Build a map of tool call IDs to track AI <--> tool message correspondences
+            const toolCallIdMap = new Map<string, number>();
+            
+            // Identify tool call IDs in the context
+            for (let i = 0; i < context.length; i++) {
+              const msg = context[i];
+              
+              // Check for tool calls in AI messages
+              if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
+                for (const item of msg.content) {
+                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                    toolCallIdMap.set(item.id, i);
+                  }
+                }
+              }
+              
+              // Check for tool messages
+              if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {
+                toolCallIdMap.set(msg.tool_call_id, i);
+              }
+            }
+            
+            // Track which messages to remove
+            const indicesToRemove = new Set<number>();
+            
+            // Start removing messages from the end, but preserve AI <--> tool message correspondences
             let i = context.length - 1;
             while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {
               const msgToRemove = context[i];
-              const msgTokenCount = tokenCounter(msgToRemove);
-              context.splice(i, 1);
-              currentTokenCount -= msgTokenCount;
+              
+              // Check if this is a tool message or has tool calls
+              let canRemove = true;
+              
+              if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {
+                // If this is a tool message, check if we need to keep its corresponding AI message
+                const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);
+                if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {
+                  // We need to remove both the tool message and its corresponding AI message
+                  indicesToRemove.add(i);
+                  indicesToRemove.add(aiIndex);
+                  currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));
+                  canRemove = false;
+                }
+              } else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {
+                // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages
+                for (const item of msgToRemove.content) {
+                  if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
+                    const toolIndex = toolCallIdMap.get(item.id as string);
+                    if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {
+                      // We need to remove both the AI message and its corresponding tool message
+                      indicesToRemove.add(i);
+                      indicesToRemove.add(toolIndex);
+                      currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));
+                      canRemove = false;
+                      break;
+                    }
+                  }
+                }
+              }
+              
+              // If we can remove this message individually
+              if (canRemove && !indicesToRemove.has(i)) {
+                indicesToRemove.add(i);
+                currentTokenCount -= tokenCounter(msgToRemove);
+              }
+              
               i--;
             }
             
+            // Remove messages in reverse order to avoid index shifting
+            const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);
+            for (const index of sortedIndices) {
+              context.splice(index, 1);
+            }
+            
             // Update remainingContextTokens to reflect the new token count
             remainingContextTokens = maxContextTokens - currentTokenCount;
           }
         }
       }
     }
+    
+    // If the latest message is an assistant message, ensure an assistant message appears early in the context
+    // but maintain system message precedence
+    if (latestMessageIsAssistant && context.length > 0) {
+      // Find the first assistant message in the context
+      const assistantIndex = context.findIndex(msg => msg.getType() === 'ai');
+      if (assistantIndex > 0) {
+        // Check if there's a system message at the beginning
+        const hasSystemFirst = context[0].getType() === 'system';
+        
+        // Move the assistant message to the appropriate position
+        const assistantMsg = context[assistantIndex];
+        context.splice(assistantIndex, 1);
+        
+        if (hasSystemFirst) {
+          // Insert after the system message
+          context.splice(1, 0, assistantMsg);
+        } else {
+          // Insert at the beginning if no system message
+          context.unshift(assistantMsg);
+        }
+      }
+    }
   }
   }
 
-  if (instructions && _messages.length > 0) {
-    context.push(_messages[0] as BaseMessage);
-    messages.shift();
-  }
-
   const prunedMemory = messages;
   summaryIndex = prunedMemory.length - 1;
   remainingContextTokens -= currentTokenCount;

package/src/specs/thinking-prune.test.ts

@@ -390,4 +390,266 @@ describe('Prune Messages with Thinking Mode Tests', () => {
     // The function should not throw an error even though no thinking blocks are found
     expect(() => pruneMessages({ messages })).not.toThrow();
   });
+  
+  it('should preserve AI <--> tool message correspondences when pruning', () => {
+    // Create a token counter
+    const tokenCounter = createTestTokenCounter();
+    
+    // Create messages with tool calls
+    const assistantMessageWithToolCall = new AIMessage({
+      content: [
+        {
+          type: "text",
+          text: "Let me check that file:",
+        },
+        {
+          type: "tool_use",
+          id: "tool123",
+          name: "text_editor_mcp_textEditor",
+          input: "{\"command\": \"view\", \"path\": \"/path/to/file.txt\"}",
+        },
+      ],
+    });
+    
+    const toolResponseMessage = new ToolMessage({
+      content: [
+        {
+          type: "text",
+          text: "{\"success\":true,\"message\":\"File content\"}",
+        },
+      ],
+      tool_call_id: "tool123",
+      name: "text_editor_mcp_textEditor",
+    });
+    
+    const assistantMessageWithThinking = new AIMessage({
+      content: [
+        {
+          type: "thinking",
+          thinking: "This is a thinking block",
+        },
+        {
+          type: "text",
+          text: "Response with thinking",
+        },
+      ],
+    });
+    
+    const messages = [
+      new SystemMessage("System instruction"),
+      new HumanMessage("Hello"),
+      assistantMessageWithToolCall,
+      toolResponseMessage,
+      new HumanMessage("Next message"),
+      assistantMessageWithThinking,
+    ];
+    
+    // Calculate token counts for each message
+    const indexTokenCountMap: Record<string, number> = {};
+    for (let i = 0; i < messages.length; i++) {
+      indexTokenCountMap[i] = tokenCounter(messages[i]);
+    }
+    
+    // Create pruneMessages function with thinking mode enabled and a low token limit
+    const pruneMessages = createPruneMessages({
+      maxTokens: 100, // Set a low token limit to force pruning
+      startIndex: 0,
+      tokenCounter,
+      indexTokenCountMap: { ...indexTokenCountMap },
+      thinkingEnabled: true,
+    });
+    
+    // Prune messages
+    const result = pruneMessages({ messages });
+    
+    // Find assistant message with tool call and its corresponding tool message in the pruned context
+    const assistantIndex = result.context.findIndex(msg => 
+      msg.getType() === 'ai' && 
+      Array.isArray(msg.content) && 
+      msg.content.some(item => item && typeof item === 'object' && item.type === 'tool_use' && item.id === 'tool123')
+    );
+    
+    // If the assistant message with tool call is in the context, its corresponding tool message should also be there
+    if (assistantIndex !== -1) {
+      const toolIndex = result.context.findIndex(msg => 
+        msg.getType() === 'tool' && 
+        'tool_call_id' in msg && 
+        msg.tool_call_id === 'tool123'
+      );
+      
+      expect(toolIndex).not.toBe(-1);
+    }
+    
+    // If the tool message is in the context, its corresponding assistant message should also be there
+    const toolIndex = result.context.findIndex(msg => 
+      msg.getType() === 'tool' && 
+      'tool_call_id' in msg && 
+      msg.tool_call_id === 'tool123'
+    );
+    
+    if (toolIndex !== -1) {
+      const assistantWithToolIndex = result.context.findIndex(msg => 
+        msg.getType() === 'ai' && 
+        Array.isArray(msg.content) && 
+        msg.content.some(item => item && typeof item === 'object' && item.type === 'tool_use' && item.id === 'tool123')
+      );
+      
+      expect(assistantWithToolIndex).not.toBe(-1);
+    }
+  });
+  
+  it('should ensure an assistant message appears early in the context when the latest message is an assistant message', () => {
+    // Create a token counter
+    const tokenCounter = createTestTokenCounter();
+    
+    // Create messages with the latest message being an assistant message
+    const assistantMessageWithThinking = new AIMessage({
+      content: [
+        {
+          type: "thinking",
+          thinking: "This is a thinking block",
+        },
+        {
+          type: "text",
+          text: "Response with thinking",
+        },
+      ],
+    });
+    
+    // Test case without system message
+    const messagesWithoutSystem = [
+      new HumanMessage("Hello"),
+      new AIMessage("First assistant response"),
+      new HumanMessage("Next message"),
+      assistantMessageWithThinking, // Latest message is an assistant message
+    ];
+    
+    // Calculate token counts for each message
+    const indexTokenCountMapWithoutSystem: Record<string, number> = {};
+    for (let i = 0; i < messagesWithoutSystem.length; i++) {
+      indexTokenCountMapWithoutSystem[i] = tokenCounter(messagesWithoutSystem[i]);
+    }
+    
+    // Create pruneMessages function with thinking mode enabled
+    const pruneMessagesWithoutSystem = createPruneMessages({
+      maxTokens: 60, // Set a lower token limit to force more pruning
+      startIndex: 0,
+      tokenCounter,
+      indexTokenCountMap: { ...indexTokenCountMapWithoutSystem },
+      thinkingEnabled: true,
+    });
+    
+    // Prune messages
+    const resultWithoutSystem = pruneMessagesWithoutSystem({ messages: messagesWithoutSystem });
+    
+    // Verify that the first message in the pruned context is an assistant message when no system message
+    expect(resultWithoutSystem.context.length).toBeGreaterThan(0);
+    expect(resultWithoutSystem.context[0].getType()).toBe('ai');
+    
+    // Test case with system message
+    const messagesWithSystem = [
+      new SystemMessage("System instruction"),
+      new HumanMessage("Hello"),
+      new AIMessage("First assistant response"),
+      new HumanMessage("Next message"),
+      assistantMessageWithThinking, // Latest message is an assistant message
+    ];
+    
+    // Calculate token counts for each message
+    const indexTokenCountMapWithSystem: Record<string, number> = {};
+    for (let i = 0; i < messagesWithSystem.length; i++) {
+      indexTokenCountMapWithSystem[i] = tokenCounter(messagesWithSystem[i]);
+    }
+    
+    // Create pruneMessages function with thinking mode enabled
+    const pruneMessagesWithSystem = createPruneMessages({
+      maxTokens: 70, // Set a token limit to force some pruning but keep system message
+      startIndex: 0,
+      tokenCounter,
+      indexTokenCountMap: { ...indexTokenCountMapWithSystem },
+      thinkingEnabled: true,
+    });
+    
+    // Prune messages
+    const resultWithSystem = pruneMessagesWithSystem({ messages: messagesWithSystem });
+    
+    // Verify that the system message remains first, followed by an assistant message
+    expect(resultWithSystem.context.length).toBeGreaterThan(1);
+    expect(resultWithSystem.context[0].getType()).toBe('system');
+    expect(resultWithSystem.context[1].getType()).toBe('ai');
+  });
+  
+  it('should look for thinking blocks starting from the most recent messages', () => {
+    // Create a token counter
+    const tokenCounter = createTestTokenCounter();
+    
+    // Create messages with multiple thinking blocks
+    const olderAssistantMessageWithThinking = new AIMessage({
+      content: [
+        {
+          type: "thinking",
+          thinking: "This is an older thinking block",
+        },
+        {
+          type: "text",
+          text: "Older response with thinking",
+        },
+      ],
+    });
+    
+    const newerAssistantMessageWithThinking = new AIMessage({
+      content: [
+        {
+          type: "thinking",
+          thinking: "This is a newer thinking block",
+        },
+        {
+          type: "text",
+          text: "Newer response with thinking",
+        },
+      ],
+    });
+    
+    const messages = [
+      new SystemMessage("System instruction"),
+      new HumanMessage("Hello"),
+      olderAssistantMessageWithThinking,
+      new HumanMessage("Next message"),
+      newerAssistantMessageWithThinking,
+    ];
+    
+    // Calculate token counts for each message
+    const indexTokenCountMap: Record<string, number> = {};
+    for (let i = 0; i < messages.length; i++) {
+      indexTokenCountMap[i] = tokenCounter(messages[i]);
+    }
+    
+    // Create pruneMessages function with thinking mode enabled
+    // Set a token limit that will force pruning of the older assistant message
+    const pruneMessages = createPruneMessages({
+      maxTokens: 80,
+      startIndex: 0,
+      tokenCounter,
+      indexTokenCountMap: { ...indexTokenCountMap },
+      thinkingEnabled: true,
+    });
+    
+    // Prune messages
+    const result = pruneMessages({ messages });
+    
+    // Find the first assistant message in the pruned context
+    const firstAssistantIndex = result.context.findIndex(msg => msg.getType() === 'ai');
+    expect(firstAssistantIndex).not.toBe(-1);
+    
+    const firstAssistantMsg = result.context[firstAssistantIndex];
+    expect(Array.isArray(firstAssistantMsg.content)).toBe(true);
+    
+    // Verify that the first assistant message has a thinking block
+    const thinkingBlock = (firstAssistantMsg.content as any[]).find(item => 
+      item && typeof item === 'object' && item.type === 'thinking');
+    expect(thinkingBlock).toBeDefined();
+    
+    // Verify that it's the newer thinking block
+    expect(thinkingBlock.thinking).toContain("newer thinking block");
+  });
 });