From 439d29844b60106f73e1691c9443a3659169df79 Mon Sep 17 00:00:00 2001 From: David Vance Date: Mon, 19 Oct 2026 22:27:45 +0200 Subject: [PATCH] fix(core): preserve latest user query when pruning chat history during tool loop (#13383) In long agent sessions with substantial tool execution outputs, compileChatMessages prunes message history from the top until it fits within the context budget. Because extractToolSequence separates only trailing assistant toolCalls and tool outputs, the opening and latest user queries were previously eligible for pruning from the history array. When all prior turns were pruned, the request reached models/providers (e.g. Ollama with Qwen 3.8 renderers) without any user turn, causing HTTP 500 ('no user query found in messages'). This patch pins the latest user turn when the trailing sequence is an assistant/tool loop, budgeting its tokens alongside the system message and tool sequence, and preserving it in place during history pruning. Adds vitest regression coverage. Closes #13383 Signed-off-by: David Vance --- core/llm/compileChatMessages.vitest.ts | 121 +++++++++++++++++++++++++ core/llm/countTokens.ts | 57 +++++++++--- 2 files changed, 164 insertions(+), 14 deletions(-) create mode 100644 core/llm/compileChatMessages.vitest.ts diff --git a/core/llm/compileChatMessages.vitest.ts b/core/llm/compileChatMessages.vitest.ts new file mode 100644 index 00000000000..ce628ef1ded --- /dev/null +++ b/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); + }); +}); diff --git a/core/llm/countTokens.ts b/core/llm/countTokens.ts index b742d70b0f0..469f83d9fce 100644 --- a/core/llm/countTokens.ts +++ b/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;