metaphorics
46f4993f65
Two defects dropped the first steering/follow-up message typed as auto-compaction began: - The compaction AbortController (which backs isCompacting) was installed AFTER auto_compaction_start was emitted. The emit awaits extension delivery and yields to the event loop, so a message typed as the loader appeared was read while isCompacting was false and mis-routed into the core agent queue (which the handoff reset then wiped). Install the controller before the emit, and move the emit to the first statement inside the existing try so the catch/finally cleanup still runs, in both #runAutoCompaction and #runAutoShake. The handoff branch now passes the run's local abort signal instead of the mutable controller field, so a superseded run bails at the handoff entry check rather than resetting the session. - handoff() calls agent.reset(), which clears the core steering/follow-up queues. Capture both queues immediately before reset and restore them immediately after (synchronous, no await gap), so queued steers and follow-ups -- including in-flight RPC/SDK steer()/followUp() and a hidden user companion such as an ultrathink notice -- survive the new-session reset instead of being silently dropped. Adds regression tests for both defects (controller-before-emit for the context-full and shake paths, queue preservation across the reset for both pre-enqueued and in-flight messages).
350 lines
13 KiB
TypeScript
350 lines
13 KiB
TypeScript
import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
|
||
import * as path from "node:path";
|
||
import { Agent } from "@oh-my-pi/pi-agent-core";
|
||
import type { AssistantMessage, ImageContent, ToolResultMessage } from "@oh-my-pi/pi-ai";
|
||
import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
|
||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||
import { AgentSession, type AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||
|
||
const usage = {
|
||
input: 16,
|
||
output: 8,
|
||
cacheRead: 0,
|
||
cacheWrite: 0,
|
||
totalTokens: 24,
|
||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||
};
|
||
|
||
describe("AgentSession shake", () => {
|
||
let tempDir: TempDir;
|
||
let session: AgentSession;
|
||
let sessionManager: SessionManager;
|
||
let authStorage: AuthStorage;
|
||
let modelRegistry: ModelRegistry;
|
||
let events: AgentSessionEvent[];
|
||
let apiInfo: { api: AssistantMessage["api"]; provider: AssistantMessage["provider"]; model: string };
|
||
|
||
beforeEach(async () => {
|
||
tempDir = TempDir.createSync("@pi-shake-");
|
||
authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db"));
|
||
authStorage.setRuntimeApiKey("anthropic", "test-key");
|
||
modelRegistry = new ModelRegistry(authStorage);
|
||
sessionManager = SessionManager.create(tempDir.path(), tempDir.path());
|
||
events = [];
|
||
|
||
const model = getBundledModel("anthropic", "claude-sonnet-4-5");
|
||
if (!model) throw new Error("Expected built-in anthropic model to exist");
|
||
apiInfo = { api: model.api, provider: model.provider, model: model.id };
|
||
|
||
const agent = new Agent({ initialState: { model, systemPrompt: ["Test"], tools: [], messages: [] } });
|
||
session = new AgentSession({
|
||
agent,
|
||
sessionManager,
|
||
settings: Settings.isolated({ "compaction.enabled": true, "compaction.autoContinue": false }),
|
||
modelRegistry,
|
||
});
|
||
session.subscribe(event => events.push(event));
|
||
});
|
||
|
||
afterEach(async () => {
|
||
if (session) await session.dispose();
|
||
authStorage.close();
|
||
try {
|
||
await tempDir.remove();
|
||
} catch {}
|
||
vi.restoreAllMocks();
|
||
});
|
||
|
||
/** Seed a user → assistant(toolCall) → toolResult turn carrying a heavy bash result. */
|
||
function seedHeavyToolResult(text: string, toolName = "bash"): void {
|
||
const toolCallId = `call_${toolName}_${Math.random().toString(36).slice(2)}`;
|
||
sessionManager.appendMessage({
|
||
role: "user",
|
||
content: [{ type: "text", text: "do it" }],
|
||
timestamp: Date.now() - 3,
|
||
});
|
||
sessionManager.appendMessage({
|
||
role: "assistant",
|
||
content: [
|
||
{ type: "text", text: "working" },
|
||
{ type: "toolCall", id: toolCallId, name: toolName, arguments: { command: "ls" } },
|
||
],
|
||
...apiInfo,
|
||
stopReason: "toolUse",
|
||
usage,
|
||
timestamp: Date.now() - 2,
|
||
});
|
||
sessionManager.appendMessage({
|
||
role: "toolResult",
|
||
toolCallId,
|
||
toolName,
|
||
content: [{ type: "text", text }],
|
||
isError: false,
|
||
timestamp: Date.now() - 1,
|
||
});
|
||
}
|
||
|
||
function branchToolResults(): ToolResultMessage[] {
|
||
return sessionManager
|
||
.getBranch()
|
||
.filter(e => e.type === "message" && (e.message as { role?: string }).role === "toolResult")
|
||
.map(e => (e as { message: ToolResultMessage }).message);
|
||
}
|
||
|
||
describe("elide", () => {
|
||
it("drops the tool result, offloads to an artifact, and embeds the recovery link", async () => {
|
||
seedHeavyToolResult("X".repeat(4000));
|
||
const replaceSpy = vi.spyOn(session.agent, "replaceMessages");
|
||
|
||
const result = await session.shake("elide");
|
||
|
||
expect(result.mode).toBe("elide");
|
||
expect(result.toolResultsDropped).toBe(1);
|
||
expect(result.tokensFreed).toBeGreaterThan(0);
|
||
expect(result.artifactId).toBeDefined();
|
||
expect(replaceSpy).toHaveBeenCalled();
|
||
|
||
const [tr] = branchToolResults();
|
||
expect(tr.prunedAt).toBeGreaterThan(0);
|
||
const text = tr.content.map(b => (b.type === "text" ? b.text : "")).join("");
|
||
expect(text).toContain(`artifact://${result.artifactId}`);
|
||
expect(text).toContain("shaken");
|
||
});
|
||
|
||
it("returns zero counts for an empty branch", async () => {
|
||
const result = await session.shake("elide");
|
||
expect(result.toolResultsDropped).toBe(0);
|
||
expect(result.blocksDropped).toBe(0);
|
||
expect(result.tokensFreed).toBe(0);
|
||
});
|
||
});
|
||
|
||
describe("images", () => {
|
||
it("mirrors dropImages and reports the removed image count", async () => {
|
||
const png: ImageContent = { type: "image", data: "iVBORw0KGgo", mimeType: "image/png" };
|
||
sessionManager.appendMessage({
|
||
role: "user",
|
||
content: [{ type: "text", text: "look" }, png],
|
||
timestamp: Date.now(),
|
||
});
|
||
|
||
const result = await session.shake("images");
|
||
|
||
expect(result.mode).toBe("images");
|
||
expect(result.imagesDropped).toBe(1);
|
||
const branch = sessionManager.getBranch();
|
||
const userMsg = branch.find(e => e.type === "message" && (e.message as { role?: string }).role === "user");
|
||
const content = (userMsg as { message: { content: unknown } }).message.content as Array<{ type: string }>;
|
||
expect(content.some(b => b.type === "image")).toBe(false);
|
||
});
|
||
});
|
||
|
||
describe("protected tools", () => {
|
||
it("never shakes skill results", async () => {
|
||
seedHeavyToolResult("S".repeat(4000), "skill");
|
||
const result = await session.shake("elide");
|
||
expect(result.toolResultsDropped).toBe(0);
|
||
});
|
||
});
|
||
|
||
describe("auto-shake strategy", () => {
|
||
it("dispatches the elide path and emits a shake action for threshold maintenance", async () => {
|
||
session.settings.set("compaction.strategy", "shake");
|
||
session.settings.set("compaction.thresholdPercent", 1);
|
||
session.settings.set("contextPromotion.enabled", false);
|
||
|
||
// Reclaim enough that the corrected (provider − tokensFreed) figure lands
|
||
// inside the 80% recovery band — otherwise the #2275 post-shake check would
|
||
// (correctly) declare pressure unresolved and fall back to context-full.
|
||
const shakeSpy = vi
|
||
.spyOn(session, "shake")
|
||
.mockResolvedValue({ mode: "elide", toolResultsDropped: 1, blocksDropped: 0, tokensFreed: 10_000 });
|
||
|
||
const assistantMessage: AssistantMessage = {
|
||
role: "assistant",
|
||
content: [{ type: "text", text: "trigger" }],
|
||
...apiInfo,
|
||
stopReason: "stop",
|
||
usage: {
|
||
input: 10_000,
|
||
output: 1_000,
|
||
cacheRead: 0,
|
||
cacheWrite: 0,
|
||
totalTokens: 11_000,
|
||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||
},
|
||
timestamp: Date.now(),
|
||
};
|
||
session.agent.emitExternalEvent({ type: "message_end", message: assistantMessage });
|
||
session.agent.emitExternalEvent({ type: "agent_end", messages: [assistantMessage] });
|
||
await Bun.sleep(20);
|
||
|
||
expect(shakeSpy).toHaveBeenCalledWith("elide", expect.anything());
|
||
const start = events.filter(e => e.type === "auto_compaction_start");
|
||
expect(start).toHaveLength(1);
|
||
expect(start[0]).toMatchObject({ type: "auto_compaction_start", reason: "threshold", action: "shake" });
|
||
const end = events.filter(e => e.type === "auto_compaction_end");
|
||
expect(end).toHaveLength(1);
|
||
expect(end[0]).toMatchObject({ type: "auto_compaction_end", action: "shake" });
|
||
});
|
||
|
||
it("has isCompacting true when the shake auto_compaction_start event fires", async () => {
|
||
// Defect 1 parity for the shake strategy: the controller backing isCompacting
|
||
// must be installed before auto_compaction_start is emitted, so a message
|
||
// typed as the loader appears is queued safely rather than mis-routed.
|
||
session.settings.set("compaction.strategy", "shake");
|
||
session.settings.set("compaction.thresholdPercent", 1);
|
||
session.settings.set("contextPromotion.enabled", false);
|
||
|
||
let capturedIsCompacting: boolean | undefined;
|
||
const { promise: shakeStarted, resolve: onShakeStarted } = Promise.withResolvers<void>();
|
||
session.subscribe(event => {
|
||
if (event.type === "auto_compaction_start" && event.action === "shake") {
|
||
capturedIsCompacting = session.isCompacting;
|
||
onShakeStarted();
|
||
}
|
||
});
|
||
|
||
vi.spyOn(session, "shake").mockResolvedValue({
|
||
mode: "elide",
|
||
toolResultsDropped: 1,
|
||
blocksDropped: 0,
|
||
tokensFreed: 10_000,
|
||
});
|
||
|
||
const assistantMessage: AssistantMessage = {
|
||
role: "assistant",
|
||
content: [{ type: "text", text: "trigger" }],
|
||
...apiInfo,
|
||
stopReason: "stop",
|
||
usage: {
|
||
input: 10_000,
|
||
output: 1_000,
|
||
cacheRead: 0,
|
||
cacheWrite: 0,
|
||
totalTokens: 11_000,
|
||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||
},
|
||
timestamp: Date.now(),
|
||
};
|
||
session.agent.emitExternalEvent({ type: "message_end", message: assistantMessage });
|
||
session.agent.emitExternalEvent({ type: "agent_end", messages: [assistantMessage] });
|
||
await shakeStarted;
|
||
|
||
expect(capturedIsCompacting).toBe(true);
|
||
});
|
||
|
||
it("falls back to context-full when shake cannot drop context below the threshold (regression #2119)", async () => {
|
||
session.settings.set("compaction.strategy", "shake");
|
||
session.settings.set("compaction.thresholdPercent", 1);
|
||
session.settings.set("contextPromotion.enabled", false);
|
||
|
||
// Seed agent state so the post-shake estimate is well above the 1% threshold
|
||
// (~2K tokens for a 200K window). The mocked shake returns reclaimed=true but
|
||
// does not modify state, mimicking the dead-loop scenario where shake removes
|
||
// nothing material yet the threshold check stays positive.
|
||
session.agent.replaceMessages([
|
||
{
|
||
role: "user",
|
||
content: [{ type: "text", text: "x".repeat(40000) }],
|
||
timestamp: Date.now(),
|
||
} as never,
|
||
]);
|
||
|
||
const shakeSpy = vi
|
||
.spyOn(session, "shake")
|
||
.mockResolvedValue({ mode: "elide", toolResultsDropped: 1, blocksDropped: 0, tokensFreed: 10 });
|
||
|
||
const assistantMessage: AssistantMessage = {
|
||
role: "assistant",
|
||
content: [{ type: "text", text: "trigger" }],
|
||
...apiInfo,
|
||
stopReason: "stop",
|
||
usage: {
|
||
input: 10_000,
|
||
output: 1_000,
|
||
cacheRead: 0,
|
||
cacheWrite: 0,
|
||
totalTokens: 11_000,
|
||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||
},
|
||
timestamp: Date.now(),
|
||
};
|
||
session.agent.emitExternalEvent({ type: "message_end", message: assistantMessage });
|
||
session.agent.emitExternalEvent({ type: "agent_end", messages: [assistantMessage] });
|
||
await Bun.sleep(50);
|
||
|
||
// Shake fires once. The pre-fix bug auto-continued, which would re-trigger shake
|
||
// on the next agent_end. The fix replaces that loop with a one-shot fallback.
|
||
expect(shakeSpy).toHaveBeenCalledTimes(1);
|
||
|
||
const shakeEnd = events.find(
|
||
e => e.type === "auto_compaction_end" && (e as { action?: string }).action === "shake",
|
||
) as { errorMessage?: string; skipped?: boolean } | undefined;
|
||
expect(shakeEnd).toBeDefined();
|
||
expect(shakeEnd?.errorMessage).toMatch(/falling back to context-full/i);
|
||
|
||
// Fallback enters the context-full path so the situation actually resolves.
|
||
const fullStart = events.find(
|
||
e => e.type === "auto_compaction_start" && (e as { action?: string }).action === "context-full",
|
||
);
|
||
expect(fullStart).toBeDefined();
|
||
});
|
||
|
||
it("falls back when provider-reported usage stays above the threshold even though the local estimate is below it (regression #2275)", async () => {
|
||
session.settings.set("compaction.strategy", "shake");
|
||
session.settings.set("compaction.thresholdTokens", 5_000);
|
||
session.settings.set("contextPromotion.enabled", false);
|
||
|
||
// Agent state holds almost no content, so #estimatePendingPromptTokens reads
|
||
// well below the 5K threshold. The pre-fix post-shake check trusted that
|
||
// estimate and treated the pressure as resolved, even though the assistant
|
||
// message's provider-reported usage (11K) was well above the threshold.
|
||
// This is the metric-divergence dead loop from #2275: thinking-heavy
|
||
// sessions hit it for real (thinkingSignature payloads aren't counted by
|
||
// the estimator), and an empty-state probe mimics it deterministically.
|
||
session.agent.replaceMessages([]);
|
||
|
||
const shakeSpy = vi
|
||
.spyOn(session, "shake")
|
||
.mockResolvedValue({ mode: "elide", toolResultsDropped: 1, blocksDropped: 0, tokensFreed: 10 });
|
||
|
||
const assistantMessage: AssistantMessage = {
|
||
role: "assistant",
|
||
content: [{ type: "text", text: "trigger" }],
|
||
...apiInfo,
|
||
stopReason: "stop",
|
||
usage: {
|
||
input: 10_000,
|
||
output: 1_000,
|
||
cacheRead: 0,
|
||
cacheWrite: 0,
|
||
totalTokens: 11_000,
|
||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||
},
|
||
timestamp: Date.now(),
|
||
};
|
||
session.agent.emitExternalEvent({ type: "message_end", message: assistantMessage });
|
||
session.agent.emitExternalEvent({ type: "agent_end", messages: [assistantMessage] });
|
||
await Bun.sleep(50);
|
||
|
||
expect(shakeSpy).toHaveBeenCalledTimes(1);
|
||
|
||
const shakeEnd = events.find(
|
||
e => e.type === "auto_compaction_end" && (e as { action?: string }).action === "shake",
|
||
) as { errorMessage?: string; skipped?: boolean } | undefined;
|
||
expect(shakeEnd).toBeDefined();
|
||
expect(shakeEnd?.errorMessage).toMatch(/falling back to context-full/i);
|
||
|
||
const fullStart = events.find(
|
||
e => e.type === "auto_compaction_start" && (e as { action?: string }).action === "context-full",
|
||
);
|
||
expect(fullStart).toBeDefined();
|
||
});
|
||
});
|
||
});
|