diff --git a/core/autocomplete/templating/__tests__/renderPrompt.vitest.ts b/core/autocomplete/templating/__tests__/renderPrompt.vitest.ts index 039dde6cabe..33ef14e2204 100644 --- a/core/autocomplete/templating/__tests__/renderPrompt.vitest.ts +++ b/core/autocomplete/templating/__tests__/renderPrompt.vitest.ts @@ -26,12 +26,14 @@ vi.mock("../../../llm/countTokens", () => { const pruneLinesFromBottom = (str: string, allowed: number) => str.slice(0, allowed); const getTokenCountingBufferSafety = () => 0; + const MIN_RESPONSE_TOKENS = 1000; return { countTokens, pruneLinesFromTop, pruneLinesFromBottom, getTokenCountingBufferSafety, + MIN_RESPONSE_TOKENS, }; }); @@ -254,6 +256,31 @@ describe("renderPromptWithTokenLimit parity & pruning", () => { expect(compiledPrefix.length).toBeLessThan(120); }); + + it("does not wipe a small prompt when default maxTokens exceeds a small context length", () => { + // Regression test: a model whose contextLength (e.g. Ollama's num_ctx) is smaller + // than the default maxTokens reservation (4096) used to compute a negative prompt + // budget, which caused every prefix/suffix to be pruned to nothing regardless of + // how small the actual prompt was. + const shortPrefix = "A".repeat(50); + + const helper = makeHelper({ prunedPrefix: shortPrefix }); + + const llmStub = { + contextLength: 2048, + completionOptions: {}, + model: "test-model", + } as any; + + const { prefix: compiledPrefix } = renderPromptWithTokenLimit({ + snippetPayload: emptySnippetPayload, + workspaceDirs: ["file:///workspace"], + helper, + llm: llmStub, + }); + + expect(compiledPrefix.includes(shortPrefix)).toBe(true); + }); }); describe("stop-token merging", () => { diff --git a/core/autocomplete/templating/index.ts b/core/autocomplete/templating/index.ts index 3ebf29b71be..6de5b8df209 100644 --- a/core/autocomplete/templating/index.ts +++ b/core/autocomplete/templating/index.ts @@ -9,6 +9,7 @@ import { DEFAULT_MAX_TOKENS } from "../../llm/constants.js"; import { countTokens, getTokenCountingBufferSafety, + MIN_RESPONSE_TOKENS, pruneLinesFromBottom, pruneLinesFromTop, } from "../../llm/countTokens"; @@ -204,7 +205,12 @@ function pruneLength(llm: ILLM, prompt: string): number { const contextLength = llm.contextLength; const reservedTokens = llm.completionOptions.maxTokens ?? DEFAULT_MAX_TOKENS; const safetyBuffer = getTokenCountingBufferSafety(contextLength); - const maxAllowedPromptTokens = contextLength - reservedTokens - safetyBuffer; + // Reserve at most MIN_RESPONSE_TOKENS for the completion, same as compileChatMessages. + // Otherwise a large maxTokens (the default is 4096) can consume the entire context + // window on small-context models, leaving no room for the prompt and silently + // pruning it to nothing on every request. + const minOutputTokens = Math.min(MIN_RESPONSE_TOKENS, reservedTokens); + const maxAllowedPromptTokens = contextLength - minOutputTokens - safetyBuffer; const promptTokenCount = countTokens(prompt, llm.model); return promptTokenCount - maxAllowedPromptTokens; } diff --git a/core/llm/countTokens.ts b/core/llm/countTokens.ts index b742d70b0f0..a99a2a99ce4 100644 --- a/core/llm/countTokens.ts +++ b/core/llm/countTokens.ts @@ -562,6 +562,7 @@ export { countTokens, countTokensAsync, extractToolSequence, + MIN_RESPONSE_TOKENS, pruneLinesFromBottom, pruneLinesFromTop, pruneRawPromptFromTop,