Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 83 additions & 72 deletions extensions/cli/src/compaction.infiniteLoop.test.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { ModelConfig } from "@continuedev/config-yaml";
import { BaseLlmApi } from "@continuedev/openai-adapters";
import { convertToUnifiedHistory } from "core/util/messageConversion.js";
import { describe, expect, it, vi } from "vitest";
import { beforeEach, describe, expect, it, vi } from "vitest";

import { compactChatHistory } from "./compaction.js";
import { streamChatResponse } from "./stream/streamChatResponse.js";
Expand All @@ -12,7 +12,10 @@ vi.mock("./stream/streamChatResponse.js", () => ({
}));

vi.mock("./util/tokenizer.js", () => ({
countChatHistoryItemTokens: vi.fn(),
countChatHistoryTokens: vi.fn(),
countToolDefinitionTokens: vi.fn(),
countTotalInputTokens: vi.fn(),
getModelContextLimit: vi.fn(),
getModelMaxTokens: vi.fn(),
}));
Expand All @@ -30,17 +33,51 @@ describe("compaction infinite loop prevention", () => {

const mockLlmApi = {} as BaseLlmApi;

it("should not loop infinitely when pruning doesn't reduce history size", async () => {
const { countChatHistoryTokens, getModelContextLimit } = await import(
"./util/tokenizer.js"
);
// Keep tokenizer/stream spy call histories isolated per test so exact
// call-count assertions (e.g. "pruned until it fits") are per-test.
beforeEach(() => {
vi.clearAllMocks();
});

/**
* A history large enough that the pinned prefix + recent tail do not cover
* it all, so compactChatHistory reaches the pruning loop instead of taking
* the no-op (everything-fits) early return.
*/
const buildLargeHistory = () => {
const history = convertToUnifiedHistory([
{ role: "system", content: "System message" },
{ role: "user", content: "Initial task" },
]);
for (let i = 0; i < 20; i++) {
history.push({
message: { role: "assistant", content: `Assistant response ${i}` },
contextItems: [],
});
history.push({
message: { role: "user", content: `User followup ${i}` },
contextItems: [],
});
}
return history;
};

const setupDefaultMocks = async () => {
const {
countChatHistoryItemTokens,
countChatHistoryTokens,
getModelContextLimit,
getModelMaxTokens,
} = await import("./util/tokenizer.js");
const mockStreamResponse = vi.mocked(streamChatResponse);
const mockCountTokens = vi.mocked(countChatHistoryTokens);
const mockGetContextLimit = vi.mocked(getModelContextLimit);

// Setup mocks
mockGetContextLimit.mockReturnValue(4000);
mockCountTokens.mockReturnValue(5000); // Always too big
vi.mocked(getModelContextLimit).mockReturnValue(4000);
vi.mocked(getModelMaxTokens).mockReturnValue(1000);
// Per-message token estimate: keep the tail budget (~12K) from covering
// the whole 40+ message history, so the fold region is non-empty.
vi.mocked(countChatHistoryItemTokens).mockReturnValue(1000);
vi.mocked(countChatHistoryTokens).mockReturnValue(5000);

mockStreamResponse.mockImplementation(
async (history, model, api, controller, callbacks) => {
callbacks?.onContent?.("Summary");
Expand All @@ -49,91 +86,65 @@ describe("compaction infinite loop prevention", () => {
},
);

// History that can't be pruned further (only system message)
const history = convertToUnifiedHistory([
{ role: "system", content: "System message" },
]);
return {
mockStreamResponse,
mockCountHistoryTokens: vi.mocked(countChatHistoryTokens),
mockGetContextLimit: vi.mocked(getModelContextLimit),
};
};

it("should not loop infinitely when pruning doesn't reduce history size", async () => {
const { mockCountHistoryTokens } = await setupDefaultMocks();

// Token count is always over the available-for-input budget, so the
// pruning loop must keep pruning and eventually exit (when the history is
// exhausted) instead of hanging.
const history = buildLargeHistory();

// This should not hang - it should break out of the loop
const result = await compactChatHistory(history, mockModel, mockLlmApi);

// Should complete successfully even though token count is still too high
expect(result.compactedHistory).toBeDefined();
expect(mockCountTokens).toHaveBeenCalled();
expect(mockCountHistoryTokens).toHaveBeenCalled();
});

it("should not loop infinitely with history ending in assistant message", async () => {
const { countChatHistoryTokens, getModelContextLimit } = await import(
"./util/tokenizer.js"
);
const mockStreamResponse = vi.mocked(streamChatResponse);
const mockCountTokens = vi.mocked(countChatHistoryTokens);
const mockGetContextLimit = vi.mocked(getModelContextLimit);

// Setup mocks
mockGetContextLimit.mockReturnValue(4000);
mockCountTokens.mockReturnValue(5000); // Always too big
mockStreamResponse.mockImplementation(
async (history, model, api, controller, callbacks) => {
callbacks?.onContent?.("Summary");
callbacks?.onContentComplete?.("Summary");
return "Summary";
},
);

// History that ends with assistant - pruning won't change it
const history = convertToUnifiedHistory([
{ role: "system", content: "System message" },
{ role: "user", content: "Hello" },
{ role: "assistant", content: "Hi there" },
]);
await setupDefaultMocks();

// Ends with an assistant message (after a user turn); pruneLastMessage
// must remove the pair and the loop must terminate.
const history = buildLargeHistory();
history.push({
message: { role: "user", content: "Last user turn" },
contextItems: [],
});
history.push({
message: { role: "assistant", content: "Last assistant turn" },
contextItems: [],
});

// This should not hang
const result = await compactChatHistory(history, mockModel, mockLlmApi);

expect(result.compactedHistory).toBeDefined();
});

it("should successfully prune when pruning actually reduces size", async () => {
const { countChatHistoryTokens, getModelContextLimit } = await import(
"./util/tokenizer.js"
);
const mockStreamResponse = vi.mocked(streamChatResponse);
const mockCountTokens = vi.mocked(countChatHistoryTokens);
const mockGetContextLimit = vi.mocked(getModelContextLimit);

// Setup mocks
mockGetContextLimit.mockReturnValue(4000);
const { mockCountHistoryTokens } = await setupDefaultMocks();

// Mock token counting to show reduction after pruning
// Available for input = 4000 - 1000 - 700 (prompt) = 2300. First check is
// over budget, second is still over, third fits -> loop exits.
let callCount = 0;
mockCountTokens.mockImplementation(() => {
mockCountHistoryTokens.mockImplementation(() => {
callCount++;
if (callCount === 1) return 5000; // Initial too big
if (callCount === 2) return 3000; // After pruning, fits
return 2000; // Subsequent calls
if (callCount === 2) return 3000; // After first prune, still too big
return 2000; // Subsequent calls fit
});

mockStreamResponse.mockImplementation(
async (history, model, api, controller, callbacks) => {
callbacks?.onContent?.("Summary");
callbacks?.onContentComplete?.("Summary");
return "Summary";
},
);

// History that can be successfully pruned
const history = convertToUnifiedHistory([
{ role: "system", content: "System message" },
{ role: "user", content: "Hello" },
{ role: "assistant", content: "Hi there" },
{ role: "user", content: "Another question" },
]);
const history = buildLargeHistory();

const result = await compactChatHistory(history, mockModel, mockLlmApi);

expect(result.compactedHistory).toBeDefined();
// The function will call countTokens multiple times during the process
expect(mockCountTokens).toHaveBeenCalled();
expect(mockCountHistoryTokens).toHaveBeenCalledTimes(3);
});
});
Loading
Loading