smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(session): re-arm mid-turn compaction after new cut points

Browse Source 
The dead-end mark parked the live tool-loop context permanently, but the array keeps growing: a later smaller tool result can move the oversized turn to the summarizable side. Recheck prepareCompaction each boundary and re-attempt once a cut point appears, instead of suppressing until provider overflow.

Added a regression proving maintenance re-attempts once a cut point becomes available.

Fixes #7151
This commit is contained in:
roboomp
2026-07-31 08:37:56 +00:00
parent ae7038f9d3
commit f6878d3739
2 changed files with 166 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/session/session-maintenance.ts
+24 -2
View File
@@ -268,7 +268,12 @@ export interface SessionMaintenanceHost {
export class SessionMaintenance {
#compactionAbortController: AbortController | undefined;
#autoCompactionAbortController: AbortController | undefined;
/** Live tool-loop contexts whose mid-turn maintenance reached a no-progress dead end. */
/**
* Live tool-loop contexts parked after mid-turn maintenance hit a no-progress
* dead end. Membership suppresses the repeated rescue + warning while no cut
* point exists; {@link maintainContextMidRun} re-arms the entry once a later
* tool result makes `prepareCompaction` viable again.
*/
readonly #midTurnCompactionDeadEnds = new WeakSet<AgentMessage[]>();
#skipPostTurnMaintenanceAssistantTimestamp: number | undefined;
readonly #host: SessionMaintenanceHost;
@@ -1034,7 +1039,24 @@ export class SessionMaintenance {
) {
return;
}
if (this.#midTurnCompactionDeadEnds.has(activeMessages)) return;
if (this.#midTurnCompactionDeadEnds.has(activeMessages)) {
// A prior boundary already ran the dead-end rescue and could not reduce
// this turn. Re-running the rescue and re-emitting its warning on every
// following tool boundary is wasted work while nothing summarizable
// exists. But the tool loop keeps appending turns: once a later
// (smaller) tool result gives prepareCompaction a cut point before the
// now-older oversized turn, compaction can finally make progress and
// MUST run rather than stay suppressed until provider overflow (#7153
// review). Stay parked only while no cut point is available; re-arm as
// soon as one appears.
if (
!model ||
prepareCompaction(this.#host.sessionManager.getBranch(), compactionSettings, model) === undefined
) {
return;
}
this.#midTurnCompactionDeadEnds.delete(activeMessages);
}
const lastAssistant = [...activeMessages]
.reverse()
packages/coding-agent/test/agent-session-mid-turn-compaction-dead-end.test.ts
+142 -42
View File
@@ -1,16 +1,19 @@
import { afterEach, describe, expect, it, vi } from "bun:test";
import * as fs from "node:fs";
import * as path from "node:path";
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
import * as compactionModule from "@oh-my-pi/pi-agent-core/compaction";
import { z } from "@oh-my-pi/pi-ai";
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
import { createMockModel, type MockResponse } from "@oh-my-pi/pi-ai/providers/mock";
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
import { loadExtensions } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/loader";
import { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/runner";
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages";
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
import { TempDir } from "@oh-my-pi/pi-utils";
import { getProjectAgentDir, TempDir } from "@oh-my-pi/pi-utils";
const noopSchema = z.object({});
const noopTool: AgentTool<typeof noopSchema, undefined> = {
@@ -25,6 +28,11 @@ const noopTool: AgentTool<typeof noopSchema, undefined> = {
const DEAD_END_WARNING = "Compaction freed too little context to make progress";
/** One threshold-tripping tool-call turn. */
function toolTurn(id: string): MockResponse {
return { content: [{ type: "toolCall", id, name: "noop", arguments: {} }], usage: { input: 190_000 } };
}
describe("AgentSession mid-turn compaction dead-end", () => {
let tempDir: TempDir;
let authStorage: AuthStorage;
@@ -37,43 +45,53 @@ describe("AgentSession mid-turn compaction dead-end", () => {
vi.restoreAllMocks();
});
it("attempts and warns once per oversized tool-loop turn", async () => {
vi.spyOn(compactionModule, "prepareCompaction").mockReturnValue(undefined);
async function createSession(options: {
responses: MockResponse[];
/** Optional `session_before_compact` short-circuit so a viable compaction makes no LLM call. */
shortCircuitCompaction?: boolean;
}): Promise<{ notices: string[]; compactionStarts: number[]; compactionResults: number }> {
tempDir = TempDir.createSync("@pi-mid-turn-compaction-dead-end-");
authStorage = await AuthStorage.create(path.join(tempDir.path(), "auth.db"));
authStorage.setRuntimeApiKey("mock", "test-key");
const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml"));
const mock = createMockModel({
responses: [
{
content: [{ type: "toolCall", id: "noop-1", name: "noop", arguments: {} }],
usage: { input: 190_000 },
},
{
content: [{ type: "toolCall", id: "noop-2", name: "noop", arguments: {} }],
usage: { input: 190_000 },
},
{ content: ["done"] },
{
content: [{ type: "toolCall", id: "noop-3", name: "noop", arguments: {} }],
usage: { input: 190_000 },
},
{
content: [{ type: "toolCall", id: "noop-4", name: "noop", arguments: {} }],
usage: { input: 190_000 },
},
{ content: ["done again"] },
],
});
const mock = createMockModel({ responses: options.responses });
vi.spyOn(modelRegistry, "getAvailable").mockReturnValue([mock]);
let extensionRunner: ExtensionRunner | undefined;
const sessionManager = SessionManager.inMemory(tempDir.path());
if (options.shortCircuitCompaction) {
const extensionsDir = path.join(getProjectAgentDir(tempDir.path()), "extensions");
fs.mkdirSync(extensionsDir, { recursive: true });
const extensionPath = path.join(extensionsDir, "compaction-short-circuit.ts");
fs.writeFileSync(
extensionPath,
[
"export default function(pi) {",
'\tpi.on("session_before_compact", async (event) => ({',
"\t\tcompaction: {",
'\t\t\tsummary: "compacted",',
"\t\t\tshortSummary: undefined,",
"\t\t\tfirstKeptEntryId: event.preparation.firstKeptEntryId,",
"\t\t\ttokensBefore: event.preparation.tokensBefore,",
"\t\t\tdetails: {},",
"\t\t},",
"\t}));",
"}",
].join("\n"),
);
const loaded = await loadExtensions([extensionPath], tempDir.path());
extensionRunner = new ExtensionRunner(
loaded.extensions,
loaded.runtime,
tempDir.path(),
sessionManager,
modelRegistry,
);
}
const agent = new Agent({
getApiKey: () => "test-key",
initialState: {
model: mock,
systemPrompt: ["Test"],
tools: [noopTool],
messages: [],
},
initialState: { model: mock, systemPrompt: ["Test"], tools: [noopTool], messages: [] },
convertToLlm,
streamFn: mock.stream,
});
@@ -87,28 +105,110 @@ describe("AgentSession mid-turn compaction dead-end", () => {
});
session = new AgentSession({
agent,
sessionManager: SessionManager.inMemory(tempDir.path()),
sessionManager,
settings,
modelRegistry,
toolRegistry: new Map([[noopTool.name, noopTool]]),
extensionRunner,
});
const notices: string[] = [];
let compactionStarts = 0;
const compactionStarts: number[] = [];
const compactionResults = 0;
const state = { notices, compactionStarts, compactionResults };
session.subscribe(event => {
if (event.type === "notice") notices.push(event.message);
if (event.type === "auto_compaction_start") compactionStarts++;
else if (event.type === "auto_compaction_start") compactionStarts.push(1);
else if (event.type === "auto_compaction_end" && event.result) state.compactionResults++;
});
return state;
}
it("attempts and warns once per oversized tool-loop turn when no cut point ever appears", async () => {
// The genuinely-unrecoverable shape: prepareCompaction can never find a cut
// point (the whole oversized region is the most recent turn). Re-running the
// rescue and re-emitting the warning on every following tool boundary is
// wasted work, so it must fire exactly once per turn — and re-arm cleanly
// for the next user turn.
vi.spyOn(compactionModule, "prepareCompaction").mockReturnValue(undefined);
const state = await createSession({
responses: [
toolTurn("noop-1"),
toolTurn("noop-2"),
{ content: ["done"] },
toolTurn("noop-3"),
toolTurn("noop-4"),
{ content: ["done again"] },
],
});
await session.prompt("Run both tools before answering");
expect(mock.calls).toHaveLength(3);
expect(notices).toEqual([expect.stringContaining(DEAD_END_WARNING)]);
expect(compactionStarts).toBe(1);
expect(state.notices).toEqual([expect.stringContaining(DEAD_END_WARNING)]);
expect(state.compactionStarts).toHaveLength(1);
await session.prompt("Run two more tools before answering");
expect(state.notices).toEqual([
expect.stringContaining(DEAD_END_WARNING),
expect.stringContaining(DEAD_END_WARNING),
]);
expect(state.compactionStarts).toHaveLength(2);
});
expect(mock.calls).toHaveLength(6);
expect(notices).toEqual([expect.stringContaining(DEAD_END_WARNING), expect.stringContaining(DEAD_END_WARNING)]);
expect(compactionStarts).toBe(2);
it("re-attempts compaction once a later tool result creates a cut point", async () => {
// Reviewer regression (#7153): the live context keeps growing, so a dead end
// at the first boundary MUST NOT permanently suppress maintenance. Once a
// later, smaller turn lets prepareCompaction cut before the now-older
// oversized turn, compaction has to run instead of parking until overflow.
let preparable = false;
const branch = () => session.sessionManager.getBranch();
const fakePreparation = (): compactionModule.CompactionPreparation => {
const entries = branch();
const firstKeptEntryId = entries[entries.length - 1]?.id;
if (!firstKeptEntryId) throw new Error("branch has no entry to keep");
return {
firstKeptEntryId,
messagesToSummarize: [{ role: "user", content: "old", timestamp: Date.now() }],
turnPrefixMessages: [],
recentMessages: [],
isSplitTurn: false,
tokensBefore: 190_000,
fileOps: { read: new Set(), written: new Set(), edited: new Set() },
settings: session.settings.getGroup("compaction"),
};
};
vi.spyOn(compactionModule, "prepareCompaction").mockImplementation(() =>
preparable ? fakePreparation() : undefined,
);
const state = await createSession({
responses: [toolTurn("noop-1"), toolTurn("noop-2"), toolTurn("noop-3"), { content: ["done"] }],
shortCircuitCompaction: true,
});
// The first boundary hits the no-preparation dead end (nothing summarizable
// and the elide rescue frees nothing). Once its warning lands, mark the
// context "preparable" to model the later cut point appearing.
vi.spyOn(session, "shake").mockResolvedValue({
mode: "elide",
toolResultsDropped: 0,
blocksDropped: 0,
tokensFreed: 0,
});
vi.spyOn(session, "getContextUsage").mockImplementation(() =>
preparable
? { tokens: 1_000, contextWindow: 200_000, percent: 0.5 }
: { tokens: 190_000, contextWindow: 200_000, percent: 95 },
);
session.subscribe(event => {
if (event.type === "notice" && event.message.includes(DEAD_END_WARNING)) preparable = true;
});
await session.prompt("Run three tools before answering");
// First dead end warned once; a later boundary re-attempted (start #2) and,
// with a cut point now available, produced a real compaction result instead
// of staying parked.
expect(state.notices.filter(m => m.includes(DEAD_END_WARNING))).toHaveLength(1);
expect(state.compactionStarts.length).toBeGreaterThanOrEqual(2);
expect(state.compactionResults).toBeGreaterThanOrEqual(1);
});
});