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

package/src/messages/prune.ts

@@ -1,6 +1,7 @@
 import {
   AIMessage,
   BaseMessage,
+  ToolMessage,
   UsageMetadata,
 } from '@langchain/core/messages';
 import type {
@@ -9,7 +10,14 @@ import type {
   ReasoningContentText,
 } from '@/types/stream';
 import type { TokenCounter } from '@/types/run';
-import { ContentTypes, Providers } from '@/common';
+import { ContentTypes, Constants, Providers } from '@/common';
+import type { ContextPruningConfig } from '@/types/graph';
+import {
+  calculateMaxToolResultChars,
+  truncateToolResultContent,
+  truncateToolInput,
+} from '@/utils/truncation';
+import { applyContextPruning } from './contextPruning';
 
 export type PruneMessagesFactoryParams = {
   provider?: Providers;
@@ -18,6 +26,15 @@ export type PruneMessagesFactoryParams = {
   tokenCounter: TokenCounter;
   indexTokenCountMap: Record<string, number | undefined>;
   thinkingEnabled?: boolean;
+  /**
+   * Optional position-based content degradation config. When `enabled: true`,
+   * tool result messages outside the protected zone (system + last N assistant
+   * turns + image content) are soft-trimmed (head + tail) once their
+   * position-age exceeds `softTrimRatio`, then hard-cleared with a placeholder
+   * once age exceeds `hardClearRatio`. Defaults disabled to preserve existing
+   * behavior — opt in per deployment.
+   */
+  contextPruningConfig?: ContextPruningConfig;
 };
 export type PruneMessagesParams = {
   messages: BaseMessage[];
@@ -142,6 +159,7 @@ export function getMessagesWithinTokenLimit({
   tokenCounter,
   thinkingStartIndex: _thinkingStartIndex = -1,
   reasoningType = ContentTypes.THINKING,
+  instructionTokens: _instructionTokens = 0,
 }: {
   messages: BaseMessage[];
   maxContextTokens: number;
@@ -151,6 +169,13 @@ export function getMessagesWithinTokenLimit({
   tokenCounter: TokenCounter;
   thinkingStartIndex?: number;
   reasoningType?: ContentTypes.THINKING | ContentTypes.REASONING_CONTENT;
+  /**
+   * Token overhead for instructions (system message + tool schemas + summary)
+   * that are NOT included in `messages`. When messages[0] is already a
+   * SystemMessage the budget is deducted from its indexTokenCountMap entry
+   * as before; otherwise this value is subtracted from the available budget.
+   */
+  instructionTokens?: number;
 }): PruningResult {
   // Every reply is primed with <|start|>assistant<|message|>, so we
   // start with 3 tokens for the label after all messages have been counted.
@@ -158,7 +183,7 @@ export function getMessagesWithinTokenLimit({
   const instructions =
     _messages[0]?.getType() === 'system' ? _messages[0] : undefined;
   const instructionsTokenCount =
-    instructions != null ? (indexTokenCountMap[0] ?? 0) : 0;
+    instructions != null ? (indexTokenCountMap[0] ?? 0) : _instructionTokens;
   const initialContextTokens = maxContextTokens - instructionsTokenCount;
   let remainingContextTokens = initialContextTokens;
   let startType = _startType;
@@ -432,6 +457,559 @@ export function checkValidNumber(value: unknown): value is number {
   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: BaseMessage): Set<string> {
+  if (message.getType() !== 'ai') {
+    return new Set<string>();
+  }
+
+  const ids = new Set<string>();
+  const aiMessage = message as AIMessage;
+  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 as { type?: unknown; id?: unknown };
+      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: BaseMessage): string | null {
+  if (message.getType() !== 'tool') {
+    return null;
+  }
+  const toolMessage = message as ToolMessage & {
+    tool_call_id?: unknown;
+    toolCallId?: unknown;
+  };
+  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,
+}: {
+  message: BaseMessage;
+  messageIndexMap: Map<BaseMessage, number>;
+  tokenCounter: TokenCounter;
+  indexTokenCountMap: Record<string, number | undefined>;
+}): number {
+  const originalIndex = messageIndexMap.get(message) ?? -1;
+  if (originalIndex > -1 && indexTokenCountMap[originalIndex] != null) {
+    return indexTokenCountMap[originalIndex] as number;
+  }
+  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.
+ */
+export function repairOrphanedToolMessages({
+  context,
+  allMessages,
+  tokenCounter,
+  indexTokenCountMap,
+}: {
+  context: BaseMessage[];
+  allMessages: BaseMessage[];
+  tokenCounter: TokenCounter;
+  indexTokenCountMap: Record<string, number | undefined>;
+}): {
+  context: BaseMessage[];
+  reclaimedTokens: number;
+  droppedOrphanCount: number;
+  /** Messages removed from context during orphan repair.  These should be
+   *  appended to `messagesToRefine` so that summarization can still see them
+   *  (e.g. a ToolMessage whose parent AI was pruned). */
+  droppedMessages: BaseMessage[];
+} {
+  const messageIndexMap = new Map<BaseMessage, number>();
+  for (let i = 0; i < allMessages.length; i++) {
+    messageIndexMap.set(allMessages[i], i);
+  }
+
+  const validToolCallIds = new Set<string>();
+  const presentToolResultIds = new Set<string>();
+  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: BaseMessage[] = [];
+  const droppedMessages: BaseMessage[] = [];
+
+  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: AIMessage,
+  presentToolResultIds: Set<string>
+): AIMessage | null {
+  const keptToolCalls = (message.tool_calls ?? []).filter(
+    (tc) => typeof tc.id === 'string' && presentToolResultIds.has(tc.id)
+  );
+
+  let keptContent: MessageContentComplex[] | string;
+  if (Array.isArray(message.content)) {
+    const filtered = (message.content as MessageContentComplex[]).filter(
+      (block) => {
+        if (typeof block !== 'object') {
+          return true;
+        }
+        const record = block as { type?: unknown; id?: unknown };
+        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.
+ */
+export function sanitizeOrphanToolBlocks(
+  messages: BaseMessage[]
+): BaseMessage[] {
+  const allToolCallIds = new Set<string>();
+  const allToolResultIds = new Set<string>();
+
+  for (const msg of messages) {
+    const msgAny = msg as unknown as Record<string, unknown>;
+    const toolCalls = msgAny.tool_calls as Array<{ id?: string }> | undefined;
+    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 as Array<Record<string, unknown>>) {
+        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 as string | undefined;
+    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: BaseMessage[] = [];
+  const strippedAiIndices = new Set<number>();
+
+  for (const msg of messages) {
+    const msgAny = msg as unknown as Record<string, unknown>;
+    const msgType =
+      typeof (msg as { getType?: unknown }).getType === 'function'
+        ? msg.getType()
+        : ((msgAny.role as string | undefined) ??
+          (msgAny._type as string | undefined));
+
+    const toolCallId = msgAny.tool_call_id as string | undefined;
+    if (
+      (msgType === 'tool' || msg instanceof ToolMessage) &&
+      typeof toolCallId === 'string' &&
+      !allToolCallIds.has(toolCallId)
+    ) {
+      continue;
+    }
+
+    const toolCalls = msgAny.tool_calls as Array<{ id?: string }> | undefined;
+    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 as Array<Record<string, unknown>>).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 as BaseMessage);
+        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.
+ */
+export function preFlightTruncateToolResults(params: {
+  messages: BaseMessage[];
+  maxContextTokens: number;
+  indexTokenCountMap: Record<string, number | undefined>;
+  tokenCounter: TokenCounter;
+}): number {
+  const { messages, maxContextTokens, indexTokenCountMap, tokenCounter } =
+    params;
+  const baseMaxChars = calculateMaxToolResultChars(maxContextTokens);
+  let truncatedCount = 0;
+
+  const toolIndices: number[] = [];
+  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 as ToolMessage).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.
+ */
+export function preFlightTruncateToolCallInputs(params: {
+  messages: BaseMessage[];
+  maxContextTokens: number;
+  indexTokenCountMap: Record<string, number | undefined>;
+  tokenCounter: TokenCounter;
+}): number {
+  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 as MessageContentComplex[];
+    const state = { changed: false };
+    const newContent = originalContent.map((block) => {
+      if (typeof block !== 'object') {
+        return block;
+      }
+      const record = block as Record<string, unknown>;
+      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 as AIMessage;
+    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;
+}
+
 type ThinkingBlocks = {
   thinking_blocks?: Array<{
     type: 'thinking';
@@ -449,10 +1027,18 @@ export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
     0
   ) as number;
   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<number, string>();
   return function pruneMessages(params: PruneMessagesParams): {
     context: BaseMessage[];
     indexTokenCountMap: Record<string, number | undefined>;
     messagesToRefine: BaseMessage[];
+    originalToolContent?: Map<number, string>;
   } {
     if (
       factoryParams.provider === Providers.OPENAI &&
@@ -586,6 +1172,51 @@ export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
       };
     }
 
+    /*
+     * 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,
@@ -604,14 +1235,40 @@ export function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {
             : 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)),
+        (finalContext.length -
+          (finalContext[0]?.getType() === 'system' ? 1 : 0)),
       0
     );
 
-    return { context, indexTokenCountMap, messagesToRefine };
+    return {
+      context: finalContext,
+      indexTokenCountMap,
+      messagesToRefine: finalMessagesToRefine,
+      originalToolContent,
+    };
   };
 }