fix(coding-agent): simplify agent session settlement

This commit is contained in:
Mario Zechner
2026-05-19 09:52:58 +02:00
parent 6a21a52d91
commit 32bcdc9739
5 changed files with 187 additions and 173 deletions

View File

@@ -85,8 +85,8 @@ describe("AgentSession bash and persistence characterization", () => {
releaseToolExecution?.();
await firstPrompt;
expect(harness.session.hasPendingBashMessages).toBe(true);
expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(false);
expect(harness.session.hasPendingBashMessages).toBe(false);
expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(true);
await harness.session.prompt("next turn");

View File

@@ -8,8 +8,8 @@ import { afterEach, describe, expect, it, vi } from "vitest";
import { createHarness, type Harness } from "./harness.js";
type SessionWithCompactionInternals = {
_checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise<void>;
_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<void>;
_checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise<boolean>;
_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<boolean>;
};
function createUsage(totalTokens: number) {
@@ -217,13 +217,9 @@ describe("AgentSession compaction characterization", () => {
timestamp: Date.now(),
});
const continueSpy = vi.spyOn(harness.session.agent, "continue").mockResolvedValue();
const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
await sessionInternals._runAutoCompaction("threshold", false);
await vi.advanceTimersByTimeAsync(100);
expect(continueSpy).toHaveBeenCalledTimes(1);
await expect(sessionInternals._runAutoCompaction("threshold", false)).resolves.toBe(true);
});
it("does not retry overflow recovery more than once", async () => {
@@ -235,7 +231,7 @@ describe("AgentSession compaction characterization", () => {
errorMessage: "prompt is too long",
timestamp: Date.now(),
});
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false);
const compactionErrors: string[] = [];
harness.session.subscribe((event) => {
if (event.type === "compaction_end" && event.errorMessage) {
@@ -283,7 +279,7 @@ describe("AgentSession compaction characterization", () => {
timestamp: Date.now(),
});
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false);
await sessionInternals._checkCompaction(staleAssistant, false);
@@ -311,7 +307,7 @@ describe("AgentSession compaction characterization", () => {
errorAssistant,
];
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false);
await sessionInternals._checkCompaction(errorAssistant);
@@ -332,7 +328,7 @@ describe("AgentSession compaction characterization", () => {
errorAssistant,
];
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false);
await sessionInternals._checkCompaction(errorAssistant);
@@ -377,7 +373,7 @@ describe("AgentSession compaction characterization", () => {
errorAssistant,
];
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(false);
await sessionInternals._checkCompaction(errorAssistant);
@@ -395,8 +391,8 @@ describe("AgentSession compaction characterization", () => {
const belowThresholdInternals = belowThresholdHarness.session as unknown as SessionWithCompactionInternals;
const disabledInternals = disabledHarness.session as unknown as SessionWithCompactionInternals;
const belowThresholdSpy = vi.spyOn(belowThresholdInternals, "_runAutoCompaction").mockResolvedValue();
const disabledSpy = vi.spyOn(disabledInternals, "_runAutoCompaction").mockResolvedValue();
const belowThresholdSpy = vi.spyOn(belowThresholdInternals, "_runAutoCompaction").mockResolvedValue(false);
const disabledSpy = vi.spyOn(disabledInternals, "_runAutoCompaction").mockResolvedValue(false);
await belowThresholdInternals._checkCompaction(
createAssistant(belowThresholdHarness, { stopReason: "stop", totalTokens: 1_000, timestamp: Date.now() }),

View File

@@ -0,0 +1,95 @@
import type { AgentTool } from "@earendil-works/pi-agent-core";
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import { Type } from "typebox";
import { afterEach, describe, expect, it } from "vitest";
import { createHarness, type Harness } from "../harness.js";
function createEchoTool(): AgentTool {
return {
name: "echo",
label: "Echo",
description: "Echo text back",
parameters: Type.Object({ text: Type.String() }),
execute: async (_toolCallId, params) => {
const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : "";
return { content: [{ type: "text", text }], details: { text } };
},
};
}
describe("regressions #1717/#2113: agent session event settlement", () => {
const harnesses: Harness[] = [];
afterEach(() => {
while (harnesses.length > 0) {
harnesses.pop()?.cleanup();
}
});
it("keeps persisted assistant/toolResult message order when extension message_end handlers yield", async () => {
const harness = await createHarness({
tools: [createEchoTool()],
extensionFactories: [
(pi) => {
pi.on("message_end", async (event) => {
if (event.message.role === "assistant") {
await new Promise((resolve) => setTimeout(resolve, 20));
}
});
},
],
});
harnesses.push(harness);
harness.setResponses([
fauxAssistantMessage([fauxToolCall("echo", { text: "one" }), fauxToolCall("echo", { text: "two" })], {
stopReason: "toolUse",
}),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("run tools");
const branchMessages = harness.sessionManager
.getBranch()
.filter((entry) => entry.type === "message")
.map((entry) => entry.message);
expect(branchMessages.map((message) => message.role)).toEqual([
"user",
"assistant",
"toolResult",
"toolResult",
"assistant",
]);
const firstToolResultIndex = branchMessages.findIndex((message) => message.role === "toolResult");
expect(firstToolResultIndex).toBeGreaterThan(0);
expect(branchMessages[firstToolResultIndex - 1]?.role).toBe("assistant");
});
it("runs tool_call handlers after the assistant tool-use message is settled in the session", async () => {
let harness: Harness;
const branchRolesAtToolCall: string[][] = [];
harness = await createHarness({
tools: [createEchoTool()],
extensionFactories: [
(pi) => {
pi.on("tool_call", () => {
branchRolesAtToolCall.push(
harness.sessionManager
.getBranch()
.filter((entry) => entry.type === "message")
.map((entry) => entry.message.role),
);
});
},
],
});
harnesses.push(harness);
harness.setResponses([
fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("run tool");
expect(branchRolesAtToolCall).toEqual([["user", "assistant"]]);
});
});