fix(core): preserve latest user query when pruning chat history during tool loop (#13383)

core/llm/compileChatMessages.vitest.ts

@@ -0,0 +1,121 @@
+import { describe, expect, it } from "vitest";
+import { ChatMessage } from "../index.js";
+import { compileChatMessages } from "./countTokens.js";
+
+describe("compileChatMessages", () => {
+  it("should preserve the latest user message when pruning during a tool sequence loop", () => {
+    const toolCall = {
+      id: "call_123",
+      type: "function" as const,
+      function: {
+        name: "test_tool",
+        arguments: "{}",
+      },
+    };
+
+    const messages: ChatMessage[] = [
+      { role: "system", content: "You are a helpful assistant." },
+      {
+        role: "user",
+        content: "Initial user prompt that started the session.",
+      },
+      {
+        role: "assistant",
+        content: "Calling tool...",
+        toolCalls: [toolCall],
+      },
+      {
+        role: "tool",
+        content: "x".repeat(3000), // Large tool output causing context pressure
+        toolCallId: "call_123",
+      },
+    ];
+
+    const result = compileChatMessages({
+      modelName: "gpt-4",
+      msgs: messages,
+      knownContextLength: 4000,
+      maxTokens: 500,
+      supportsImages: false,
+    });
+
+    const userMessage = result.compiledChatMessages.find(
+      (m) => m.role === "user",
+    );
+    expect(userMessage).toBeDefined();
+    expect(userMessage?.content).toBe(
+      "Initial user prompt that started the session.",
+    );
+
+    // Ensure system message is preserved
+    const systemMessage = result.compiledChatMessages.find(
+      (m) => m.role === "system",
+    );
+    expect(systemMessage).toBeDefined();
+
+    // Ensure tool sequence is preserved
+    const toolMessage = result.compiledChatMessages.find(
+      (m) => m.role === "tool",
+    );
+    expect(toolMessage).toBeDefined();
+  });
+
+  it("should preserve latest user message and prune older intermediate turns", () => {
+    const toolCall1 = {
+      id: "call_1",
+      type: "function" as const,
+      function: { name: "search", arguments: "{}" },
+    };
+    const toolCall2 = {
+      id: "call_2",
+      type: "function" as const,
+      function: { name: "search", arguments: "{}" },
+    };
+
+    const messages: ChatMessage[] = [
+      { role: "system", content: "You are a helpful assistant." },
+      {
+        role: "user",
+        content:
+          "Turn 1 user request that is very old and long: " + "a".repeat(4000),
+      },
+      { role: "assistant", content: "Turn 1 response." },
+      { role: "user", content: "Turn 2 user request: do task B." },
+      {
+        role: "assistant",
+        content: "Calling search...",
+        toolCalls: [toolCall2],
+      },
+      {
+        role: "tool",
+        content: "Search results for task B: " + "b".repeat(500),
+        toolCallId: "call_2",
+      },
+    ];
+
+    const result = compileChatMessages({
+      modelName: "gpt-4",
+      msgs: messages,
+      knownContextLength: 800,
+      maxTokens: 100,
+      supportsImages: false,
+    });
+
+    // Old Turn 1 should be pruned under tight budget
+    const hasTurn1 = result.compiledChatMessages.some(
+      (m) =>
+        typeof m.content === "string" &&
+        m.content.includes("Turn 1 user request"),
+    );
+    expect(hasTurn1).toBe(false);
+
+    // Latest user message (Turn 2) MUST be preserved
+    const turn2User = result.compiledChatMessages.find(
+      (m) =>
+        typeof m.content === "string" &&
+        m.content.includes("Turn 2 user request"),
+    );
+    expect(turn2User).toBeDefined();
+    expect(result.didPrune).toBe(true);
+  });
+});

core/llm/countTokens.ts

@@ -493,6 +493,17 @@ function compileChatMessages({
   inputTokensAvailable -= systemMsgTokens;
   inputTokensAvailable -= lastMessagesTokens;
 
+  // Pin the latest user message if the trailing toolSequence does not already contain a user message
+  const hasUserInToolSequence = toolSequence.some((m) => m.role === "user");
+  const pinnedUserMsg = hasUserInToolSequence
+    ? undefined
+    : [...msgsCopy].reverse().find((m) => m.role === "user");
+  const pinnedUserTokens = pinnedUserMsg
+    ? countChatMessageTokens(modelName, pinnedUserMsg)
+    : 0;
+
+  inputTokensAvailable -= pinnedUserTokens;
+
   // Make sure there's enough context for the non-excludable items
   if (knownContextLength !== undefined && inputTokensAvailable < 0) {
     throw new Error(
@@ -508,38 +519,56 @@ function compileChatMessages({
   }
 
   // Now remove messages till we're under the limit
+  const pinnedIndex = pinnedUserMsg ? msgsCopy.indexOf(pinnedUserMsg) : -1;
+  const historyEntries: (ChatMessage & { tokens: number })[] = [];
   let currentTotal = 0;
-  const historyWithTokens = msgsCopy.map((message) => {
+  let pinnedEntry: (ChatMessage & { tokens: number }) | undefined;
+
+  for (let i = 0; i < msgsCopy.length; i++) {
+    const message = msgsCopy[i];
     const tokens = countChatMessageTokens(modelName, message);
-    currentTotal += tokens;
-    return {
-      ...message,
-      tokens,
-    };
-  });
-
-  while (historyWithTokens.length > 0 && currentTotal > inputTokensAvailable) {
-    const message = historyWithTokens.shift()!;
+    if (i === pinnedIndex) {
+      pinnedEntry = { ...message, tokens };
+    } else {
+      currentTotal += tokens;
+      historyEntries.push({ ...message, tokens });
+    }
+  }
+
+  let prunedCount = 0;
+  while (historyEntries.length > 0 && currentTotal > inputTokensAvailable) {
+    const message = historyEntries.shift()!;
     currentTotal -= message.tokens;
     didPrune = true;
+    prunedCount++;
 
     // At this point make sure no latent tool response without corresponding call
-    while (historyWithTokens[0]?.role === "tool") {
-      const message = historyWithTokens.shift()!;
+    while (historyEntries[0]?.role === "tool") {
+      const message = historyEntries.shift()!;
       currentTotal -= message.tokens;
+      prunedCount++;
     }
   }
 
+  if (pinnedEntry) {
+    const targetIndex = Math.max(0, pinnedIndex - prunedCount);
+    historyEntries.splice(targetIndex, 0, pinnedEntry);
+  }
+
   // Now reassemble
   const reassembled: ChatMessage[] = [];
   if (systemMsg) {
     reassembled.push(systemMsg);
   }
-  reassembled.push(...historyWithTokens.map(({ tokens, ...rest }) => rest));
+  reassembled.push(...historyEntries.map(({ tokens, ...rest }) => rest));
   reassembled.push(...toolSequence);
 
   const inputTokens =
-    currentTotal + systemMsgTokens + toolTokens + lastMessagesTokens;
+    currentTotal +
+    pinnedUserTokens +
+    systemMsgTokens +
+    toolTokens +
+    lastMessagesTokens;
   const availableTokens =
     contextLength - countingSafetyBuffer - minOutputTokens;
   const contextPercentage = inputTokens / availableTokens;