@illuma-ai/agents 1.4.0-alpha.5 → 1.5.0

package/dist/esm/messages/prune.mjs

@@ -1,6 +1,8 @@
-import { AIMessage } from '@langchain/core/messages';
-import { Providers, ContentTypes } from '../common/enum.mjs';
+import { AIMessage, ToolMessage } from '@langchain/core/messages';
+import { Providers, ContentTypes, Constants } from '../common/enum.mjs';
 import '../tools/approval/constants.mjs';
+import { truncateToolInput, truncateToolResultContent, calculateMaxToolResultChars } from '../utils/truncation.mjs';
+import { applyContextPruning } from './contextPruning.mjs';
 
 function isIndexInContext(arrayA, arrayB, targetIndex) {
     const startingIndexInA = arrayA.length - arrayB.length;
@@ -82,12 +84,12 @@ function buildEmergencyContext(messages, endIndex) {
     }
     return result;
 }
-function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, indexTokenCountMap, startType: _startType, thinkingEnabled, tokenCounter, thinkingStartIndex: _thinkingStartIndex = -1, reasoningType = ContentTypes.THINKING, }) {
+function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, indexTokenCountMap, startType: _startType, thinkingEnabled, tokenCounter, thinkingStartIndex: _thinkingStartIndex = -1, reasoningType = ContentTypes.THINKING, instructionTokens: _instructionTokens = 0, }) {
     // Every reply is primed with <|start|>assistant<|message|>, so we
     // start with 3 tokens for the label after all messages have been counted.
     let currentTokenCount = 3;
     const instructions = _messages[0]?.getType() === 'system' ? _messages[0] : undefined;
-    const instructionsTokenCount = instructions != null ? (indexTokenCountMap[0] ?? 0) : 0;
+    const instructionsTokenCount = instructions != null ? (indexTokenCountMap[0] ?? 0) : _instructionTokens;
     const initialContextTokens = maxContextTokens - instructionsTokenCount;
     let remainingContextTokens = initialContextTokens;
     let startType = _startType;
@@ -306,12 +308,443 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
 function checkValidNumber(value) {
     return typeof value === 'number' && !isNaN(value) && value > 0;
 }
+/**
+ * Returns the set of tool_call ids referenced by an AIMessage — both via the
+ * canonical `tool_calls` array and via inline `tool_use` / `tool_call` content
+ * blocks (Anthropic-style content arrays).
+ */
+function getToolCallIds(message) {
+    if (message.getType() !== 'ai') {
+        return new Set();
+    }
+    const ids = new Set();
+    const aiMessage = message;
+    for (const toolCall of aiMessage.tool_calls ?? []) {
+        if (typeof toolCall.id === 'string' && toolCall.id.length > 0) {
+            ids.add(toolCall.id);
+        }
+    }
+    if (Array.isArray(aiMessage.content)) {
+        for (const part of aiMessage.content) {
+            if (typeof part !== 'object') {
+                continue;
+            }
+            const record = part;
+            if ((record.type === 'tool_use' || record.type === 'tool_call') &&
+                typeof record.id === 'string' &&
+                record.id.length > 0) {
+                ids.add(record.id);
+            }
+        }
+    }
+    return ids;
+}
+function getToolResultId(message) {
+    if (message.getType() !== 'tool') {
+        return null;
+    }
+    const toolMessage = message;
+    if (typeof toolMessage.tool_call_id === 'string' &&
+        toolMessage.tool_call_id.length > 0) {
+        return toolMessage.tool_call_id;
+    }
+    if (typeof toolMessage.toolCallId === 'string' &&
+        toolMessage.toolCallId.length > 0) {
+        return toolMessage.toolCallId;
+    }
+    return null;
+}
+function resolveTokenCountForMessage({ message, messageIndexMap, tokenCounter, indexTokenCountMap, }) {
+    const originalIndex = messageIndexMap.get(message) ?? -1;
+    if (originalIndex > -1 && indexTokenCountMap[originalIndex] != null) {
+        return indexTokenCountMap[originalIndex];
+    }
+    return tokenCounter(message);
+}
+/**
+ * Repairs orphan tool messages in a pruned context. Drops `ToolMessage`s whose
+ * matching `AIMessage` was pruned, and strips orphan `tool_use` content blocks
+ * (and `tool_calls` entries) from `AIMessage`s whose corresponding tool results
+ * are missing. Without this, Anthropic / Bedrock / OpenAI reject the request
+ * with a structural validation 400.
+ *
+ * Intended to run on the post-prune `context` array (not the pre-prune
+ * `allMessages`). The dropped messages are returned so the caller can append
+ * them to `messagesToRefine` for summarization fidelity.
+ */
+function repairOrphanedToolMessages({ context, allMessages, tokenCounter, indexTokenCountMap, }) {
+    const messageIndexMap = new Map();
+    for (let i = 0; i < allMessages.length; i++) {
+        messageIndexMap.set(allMessages[i], i);
+    }
+    const validToolCallIds = new Set();
+    const presentToolResultIds = new Set();
+    for (const message of context) {
+        for (const id of getToolCallIds(message)) {
+            validToolCallIds.add(id);
+        }
+        const resultId = getToolResultId(message);
+        if (resultId != null) {
+            presentToolResultIds.add(resultId);
+        }
+    }
+    let reclaimedTokens = 0;
+    let droppedOrphanCount = 0;
+    const repairedContext = [];
+    const droppedMessages = [];
+    for (const message of context) {
+        if (message.getType() === 'tool') {
+            const toolResultId = getToolResultId(message);
+            if (toolResultId == null || !validToolCallIds.has(toolResultId)) {
+                droppedOrphanCount += 1;
+                reclaimedTokens += resolveTokenCountForMessage({
+                    message,
+                    tokenCounter,
+                    messageIndexMap,
+                    indexTokenCountMap,
+                });
+                droppedMessages.push(message);
+                continue;
+            }
+            repairedContext.push(message);
+            continue;
+        }
+        if (message.getType() === 'ai' && message instanceof AIMessage) {
+            const toolCallIds = getToolCallIds(message);
+            if (toolCallIds.size > 0) {
+                let hasOrphanToolCalls = false;
+                for (const id of toolCallIds) {
+                    if (!presentToolResultIds.has(id)) {
+                        hasOrphanToolCalls = true;
+                        break;
+                    }
+                }
+                if (hasOrphanToolCalls) {
+                    const originalTokens = resolveTokenCountForMessage({
+                        message,
+                        messageIndexMap,
+                        tokenCounter,
+                        indexTokenCountMap,
+                    });
+                    const stripped = stripOrphanToolUseBlocks(message, presentToolResultIds);
+                    if (stripped != null) {
+                        const strippedTokens = tokenCounter(stripped);
+                        reclaimedTokens += originalTokens - strippedTokens;
+                        repairedContext.push(stripped);
+                    }
+                    else {
+                        droppedOrphanCount += 1;
+                        reclaimedTokens += originalTokens;
+                        droppedMessages.push(message);
+                    }
+                    continue;
+                }
+            }
+        }
+        repairedContext.push(message);
+    }
+    return {
+        context: repairedContext,
+        reclaimedTokens,
+        droppedOrphanCount,
+        droppedMessages,
+    };
+}
+/**
+ * Strips tool_use content blocks and tool_calls entries from an AI message
+ * when their corresponding ToolMessages are not in the context.
+ * Returns null if the message has no content left after stripping.
+ */
+function stripOrphanToolUseBlocks(message, presentToolResultIds) {
+    const keptToolCalls = (message.tool_calls ?? []).filter((tc) => typeof tc.id === 'string' && presentToolResultIds.has(tc.id));
+    let keptContent;
+    if (Array.isArray(message.content)) {
+        const filtered = message.content.filter((block) => {
+            if (typeof block !== 'object') {
+                return true;
+            }
+            const record = block;
+            if ((record.type === 'tool_use' || record.type === 'tool_call') &&
+                typeof record.id === 'string') {
+                return presentToolResultIds.has(record.id);
+            }
+            return true;
+        });
+        if (filtered.length === 0) {
+            return null;
+        }
+        keptContent = filtered;
+    }
+    else {
+        keptContent = message.content;
+    }
+    return new AIMessage({
+        ...message,
+        content: keptContent,
+        tool_calls: keptToolCalls.length > 0 ? keptToolCalls : undefined,
+    });
+}
+/**
+ * Lightweight structural cleanup: strips orphan tool_use blocks from AI messages
+ * and drops orphan ToolMessages whose AI counterpart is missing.
+ *
+ * Unlike `repairOrphanedToolMessages`, this does NOT track tokens — it is
+ * intended as a final safety net right before model invocation to prevent
+ * Anthropic / Bedrock structural validation errors.
+ *
+ * Uses duck-typing instead of `getType()` because messages at this stage
+ * may be plain objects (from LangGraph state serialization) rather than
+ * proper BaseMessage class instances.
+ *
+ * Includes a fast-path: if every tool_call has a matching tool_result and
+ * vice-versa, the original array is returned immediately with zero allocation.
+ */
+function sanitizeOrphanToolBlocks(messages) {
+    const allToolCallIds = new Set();
+    const allToolResultIds = new Set();
+    for (const msg of messages) {
+        const msgAny = msg;
+        const toolCalls = msgAny.tool_calls;
+        if (Array.isArray(toolCalls)) {
+            for (const tc of toolCalls) {
+                if (typeof tc.id === 'string' &&
+                    tc.id.length > 0 &&
+                    !tc.id.startsWith(Constants.ANTHROPIC_SERVER_TOOL_PREFIX)) {
+                    allToolCallIds.add(tc.id);
+                }
+            }
+        }
+        if (Array.isArray(msgAny.content)) {
+            for (const block of msgAny.content) {
+                if (typeof block === 'object' &&
+                    (block.type === 'tool_use' || block.type === 'tool_call') &&
+                    typeof block.id === 'string' &&
+                    !block.id.startsWith(Constants.ANTHROPIC_SERVER_TOOL_PREFIX)) {
+                    allToolCallIds.add(block.id);
+                }
+            }
+        }
+        const toolCallId = msgAny.tool_call_id;
+        if (typeof toolCallId === 'string' && toolCallId.length > 0) {
+            allToolResultIds.add(toolCallId);
+        }
+    }
+    let hasOrphans = false;
+    for (const id of allToolCallIds) {
+        if (!allToolResultIds.has(id)) {
+            hasOrphans = true;
+            break;
+        }
+    }
+    if (!hasOrphans) {
+        for (const id of allToolResultIds) {
+            if (!allToolCallIds.has(id)) {
+                hasOrphans = true;
+                break;
+            }
+        }
+    }
+    if (!hasOrphans) {
+        return messages;
+    }
+    const result = [];
+    const strippedAiIndices = new Set();
+    for (const msg of messages) {
+        const msgAny = msg;
+        const msgType = typeof msg.getType === 'function'
+            ? msg.getType()
+            : (msgAny.role ??
+                msgAny._type);
+        const toolCallId = msgAny.tool_call_id;
+        if ((msgType === 'tool' || msg instanceof ToolMessage) &&
+            typeof toolCallId === 'string' &&
+            !allToolCallIds.has(toolCallId)) {
+            continue;
+        }
+        const toolCalls = msgAny.tool_calls;
+        if ((msgType === 'ai' ||
+            msgType === 'assistant' ||
+            msg instanceof AIMessage) &&
+            Array.isArray(toolCalls) &&
+            toolCalls.length > 0) {
+            const hasOrphanCalls = toolCalls.some((tc) => typeof tc.id === 'string' && !allToolResultIds.has(tc.id));
+            if (hasOrphanCalls) {
+                if (msg instanceof AIMessage) {
+                    const stripped = stripOrphanToolUseBlocks(msg, allToolResultIds);
+                    if (stripped != null) {
+                        strippedAiIndices.add(result.length);
+                        result.push(stripped);
+                    }
+                    continue;
+                }
+                const keptToolCalls = toolCalls.filter((tc) => typeof tc.id === 'string' && allToolResultIds.has(tc.id));
+                const keptContent = Array.isArray(msgAny.content)
+                    ? msgAny.content.filter((block) => {
+                        if (typeof block !== 'object')
+                            return true;
+                        if ((block.type === 'tool_use' || block.type === 'tool_call') &&
+                            typeof block.id === 'string') {
+                            return allToolResultIds.has(block.id);
+                        }
+                        return true;
+                    })
+                    : msgAny.content;
+                if (keptToolCalls.length === 0 &&
+                    Array.isArray(keptContent) &&
+                    keptContent.length === 0) {
+                    continue;
+                }
+                strippedAiIndices.add(result.length);
+                const patched = Object.create(Object.getPrototypeOf(msg), Object.getOwnPropertyDescriptors(msg));
+                patched.tool_calls = keptToolCalls.length > 0 ? keptToolCalls : [];
+                patched.content = keptContent;
+                result.push(patched);
+                continue;
+            }
+        }
+        result.push(msg);
+    }
+    // Bedrock/Anthropic require the conversation to end with a user message;
+    // a stripped AI message (tool_use removed) represents a dead-end exchange.
+    while (result.length > 0 && strippedAiIndices.has(result.length - 1)) {
+        result.pop();
+    }
+    return result;
+}
+/**
+ * Pre-flight truncation: truncates oversized tool result string content before
+ * the budget walker runs. Older results get smaller budgets (recency factor),
+ * with a 200-char floor so even ancient entries keep diagnostic signal. Returns
+ * the number of ToolMessages truncated.
+ */
+function preFlightTruncateToolResults(params) {
+    const { messages, maxContextTokens, indexTokenCountMap, tokenCounter } = params;
+    const baseMaxChars = calculateMaxToolResultChars(maxContextTokens);
+    let truncatedCount = 0;
+    const toolIndices = [];
+    for (let i = 0; i < messages.length; i++) {
+        if (messages[i].getType() === 'tool') {
+            toolIndices.push(i);
+        }
+    }
+    for (let t = 0; t < toolIndices.length; t++) {
+        const i = toolIndices[t];
+        const message = messages[i];
+        const content = message.content;
+        if (typeof content !== 'string') {
+            continue;
+        }
+        const position = toolIndices.length > 1 ? t / (toolIndices.length - 1) : 1;
+        const recencyFactor = 0.2 + 0.8 * position;
+        const maxChars = Math.max(200, Math.floor(baseMaxChars * recencyFactor));
+        if (content.length <= maxChars) {
+            continue;
+        }
+        const truncated = truncateToolResultContent(content, maxChars);
+        const cloned = new ToolMessage({
+            content: truncated,
+            tool_call_id: message.tool_call_id,
+            name: message.name,
+            id: message.id,
+            additional_kwargs: message.additional_kwargs,
+            response_metadata: message.response_metadata,
+        });
+        messages[i] = cloned;
+        indexTokenCountMap[i] = tokenCounter(cloned);
+        truncatedCount++;
+    }
+    return truncatedCount;
+}
+/**
+ * Pre-flight truncation: truncates oversized `tool_use` input fields in AI messages.
+ *
+ * Tool call inputs (arguments) can be very large — e.g., code evaluation payloads from
+ * MCP tools like chrome-devtools. Since these tool calls have already been executed,
+ * the model only needs a summary of what was called, not the full arguments. Truncating
+ * them before pruning can prevent entire messages from being dropped.
+ *
+ * Uses 15% of the context window (in estimated characters, ~4 chars/token) as the
+ * per-input cap, capped at 200K chars.
+ *
+ * @returns The number of AI messages that had tool_use inputs truncated.
+ */
+function preFlightTruncateToolCallInputs(params) {
+    const { messages, maxContextTokens, indexTokenCountMap, tokenCounter } = params;
+    const maxInputChars = Math.min(Math.floor(maxContextTokens * 0.15) * 4, 200_000);
+    let truncatedCount = 0;
+    for (let i = 0; i < messages.length; i++) {
+        const message = messages[i];
+        if (message.getType() !== 'ai') {
+            continue;
+        }
+        if (!Array.isArray(message.content)) {
+            continue;
+        }
+        const originalContent = message.content;
+        const state = { changed: false };
+        const newContent = originalContent.map((block) => {
+            if (typeof block !== 'object') {
+                return block;
+            }
+            const record = block;
+            if (record.type !== 'tool_use' && record.type !== 'tool_call') {
+                return block;
+            }
+            const input = record.input;
+            if (input == null) {
+                return block;
+            }
+            const serialized = typeof input === 'string' ? input : JSON.stringify(input);
+            if (serialized.length <= maxInputChars) {
+                return block;
+            }
+            state.changed = true;
+            // Replaces original input with { _truncated, _originalChars } —
+            // safe because the tool call already executed in a prior turn.
+            return {
+                ...record,
+                input: truncateToolInput(serialized, maxInputChars),
+            };
+        });
+        if (!state.changed) {
+            continue;
+        }
+        const aiMsg = message;
+        const newToolCalls = (aiMsg.tool_calls ?? []).map((tc) => {
+            const serializedArgs = JSON.stringify(tc.args);
+            if (serializedArgs.length <= maxInputChars) {
+                return tc;
+            }
+            // Replaces original args with { _truncated, _originalChars } —
+            // safe because the tool call already executed in a prior turn.
+            return {
+                ...tc,
+                args: truncateToolInput(serializedArgs, maxInputChars),
+            };
+        });
+        messages[i] = new AIMessage({
+            ...aiMsg,
+            content: newContent,
+            tool_calls: newToolCalls.length > 0 ? newToolCalls : undefined,
+        });
+        indexTokenCountMap[i] = tokenCounter(messages[i]);
+        truncatedCount++;
+    }
+    return truncatedCount;
+}
 function createPruneMessages(factoryParams) {
     const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };
     let lastTurnStartIndex = factoryParams.startIndex;
     let lastCutOffIndex = 0;
     let totalTokens = Object.values(indexTokenCountMap).reduce((a = 0, b = 0) => a + b, 0);
     let runThinkingStartIndex = -1;
+    /**
+     * Original (pre-truncation) tool result content keyed by message index.
+     * Lets the summarizer see full tool outputs even after pre-flight truncation
+     * has shortened them in the live messages array. Cleared whenever the
+     * factory is recreated (e.g. after a summarization checkpoint).
+     */
+    const originalToolContent = new Map();
     return function pruneMessages(params) {
         if (factoryParams.provider === Providers.OPENAI &&
             factoryParams.thinkingEnabled === true) {
@@ -410,6 +843,47 @@ function createPruneMessages(factoryParams) {
                 messagesToRefine: [],
             };
         }
+        /*
+         * Pre-flight truncation: chop oversized tool result content and oversized
+         * tool_use input args BEFORE the budget walker runs. Older results get
+         * smaller per-message budgets (recencyFactor) so they degrade first while
+         * recent turns stay intact. Saves the originals in `originalToolContent`
+         * keyed by message index so the summarizer can later reconstruct full
+         * fidelity even after the live array was shortened.
+         */
+        for (let i = 0; i < params.messages.length; i++) {
+            const m = params.messages[i];
+            if (m.getType() === 'tool' &&
+                typeof m.content === 'string' &&
+                !originalToolContent.has(i)) {
+                originalToolContent.set(i, m.content);
+            }
+        }
+        preFlightTruncateToolResults({
+            messages: params.messages,
+            maxContextTokens: factoryParams.maxTokens,
+            indexTokenCountMap,
+            tokenCounter: factoryParams.tokenCounter,
+        });
+        preFlightTruncateToolCallInputs({
+            messages: params.messages,
+            maxContextTokens: factoryParams.maxTokens,
+            indexTokenCountMap,
+            tokenCounter: factoryParams.tokenCounter,
+        });
+        /*
+         * Position-based content degradation: soft-trim then hard-clear older tool
+         * results before the budget walker runs. Off by default — caller opts in
+         * via `contextPruningConfig.enabled`.
+         */
+        if (factoryParams.contextPruningConfig?.enabled) {
+            applyContextPruning({
+                messages: params.messages,
+                indexTokenCountMap,
+                tokenCounter: factoryParams.tokenCounter,
+                config: factoryParams.contextPruningConfig,
+            });
+        }
         const { context, thinkingStartIndex, messagesToRefine } = getMessagesWithinTokenLimit({
             maxContextTokens: factoryParams.maxTokens,
             messages: params.messages,
@@ -425,10 +899,33 @@ function createPruneMessages(factoryParams) {
                 : undefined,
         });
         runThinkingStartIndex = thinkingStartIndex ?? -1;
+        /*
+         * Orphan repair pass — drops tool messages whose AIMessage was pruned and
+         * strips orphan tool_use blocks from AIMessages whose tool results are
+         * missing. Without this, the budget walker can split tool_use ↔ tool_result
+         * pairs and the provider rejects the request with a structural 400.
+         * Dropped messages flow into messagesToRefine for summarization.
+         */
+        const repaired = repairOrphanedToolMessages({
+            context,
+            allMessages: params.messages,
+            tokenCounter: factoryParams.tokenCounter,
+            indexTokenCountMap,
+        });
+        const finalContext = repaired.context;
+        const finalMessagesToRefine = repaired.droppedMessages.length > 0
+            ? [...messagesToRefine, ...repaired.droppedMessages]
+            : messagesToRefine;
         /** The index is the first value of `context`, index relative to `params.messages` */
         lastCutOffIndex = Math.max(params.messages.length -
-            (context.length - (context[0]?.getType() === 'system' ? 1 : 0)), 0);
-        return { context, indexTokenCountMap, messagesToRefine };
+            (finalContext.length -
+                (finalContext[0]?.getType() === 'system' ? 1 : 0)), 0);
+        return {
+            context: finalContext,
+            indexTokenCountMap,
+            messagesToRefine: finalMessagesToRefine,
+            originalToolContent,
+        };
     };
 }
 /** Calculates context utilization as a percentage (0-100) */
@@ -442,5 +939,5 @@ function getContextUtilization(indexTokenCountMap, instructionTokens, maxContext
     return (totalTokens / maxContextTokens) * 100;
 }
 
-export { calculateTotalTokens, checkValidNumber, createPruneMessages, getContextUtilization, getMessagesWithinTokenLimit };
+export { calculateTotalTokens, checkValidNumber, createPruneMessages, getContextUtilization, getMessagesWithinTokenLimit, preFlightTruncateToolCallInputs, preFlightTruncateToolResults, repairOrphanedToolMessages, sanitizeOrphanToolBlocks };
 //# sourceMappingURL=prune.mjs.map