fix: trim custom message exclusion tests

packages/agent/test/agent.test.ts

@@ -1,19 +1,7 @@
-import {
-	type AssistantMessage,
-	type AssistantMessageEvent,
-	EventStream,
-	getModel,
-	type Message,
-} from "@earendil-works/pi-ai";
+import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@earendil-works/pi-ai";
 import { Type } from "typebox";
 import { describe, expect, it } from "vitest";
-import {
-	Agent,
-	type AgentEvent,
-	type AgentMessage,
-	type AgentTool,
-	type AgentToolUpdateCallback,
-} from "../src/index.ts";
+import { Agent, type AgentEvent, type AgentTool, type AgentToolUpdateCallback } from "../src/index.ts";
 
 // Mock stream that mimics AssistantMessageEventStream
 class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
@@ -560,67 +548,6 @@ describe("Agent", () => {
 		await firstPrompt.catch(() => {});
 	});
 
-	it("continue() should reserve the active run before async LLM conversion and use the filtered context tail", async () => {
-		const convertStarted = createDeferred();
-		const releaseConvert = createDeferred();
-		let convertCallCount = 0;
-		let providerMessages: Message[] = [];
-		const displayOnlyMessage = {
-			role: "displayOnly",
-			content: "status",
-			timestamp: Date.now(),
-		} as unknown as AgentMessage;
-		const agent = new Agent({
-			convertToLlm: async (messages) => {
-				convertCallCount++;
-				if (convertCallCount === 1) {
-					convertStarted.resolve();
-					await releaseConvert.promise;
-				}
-				return messages
-					.filter((message) => (message as { role: string }).role !== "displayOnly")
-					.filter(
-						(message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult",
-					) as Message[];
-			},
-			streamFn: (_model, context) => {
-				providerMessages = context.messages;
-				const stream = new MockAssistantStream();
-				queueMicrotask(() => {
-					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") });
-				});
-				return stream;
-			},
-		});
-		agent.state.messages = [
-			{
-				role: "user",
-				content: [{ type: "text", text: "Initial" }],
-				timestamp: Date.now() - 10,
-			},
-			createAssistantMessage("Initial response"),
-			displayOnlyMessage,
-		];
-		agent.followUp({
-			role: "user",
-			content: [{ type: "text", text: "Queued follow-up" }],
-			timestamp: Date.now(),
-		});
-
-		const continuePromise = agent.continue();
-		await convertStarted.promise;
-
-		expect(agent.state.isStreaming).toBe(true);
-		await expect(agent.prompt("Second message")).rejects.toThrow("Agent is already processing a prompt");
-
-		releaseConvert.resolve();
-		await continuePromise;
-
-		expect(convertCallCount).toBe(2);
-		expect(providerMessages[providerMessages.length - 1]?.role).toBe("user");
-		expect(agent.state.messages).toContain(displayOnlyMessage);
-	});
-
 	it("continue() should process queued follow-up messages after an assistant turn", async () => {
 		const agent = new Agent({
 			streamFn: () => {
@@ -632,6 +559,14 @@ describe("Agent", () => {
 			},
 		});
 
+		const excludedCustom = {
+			role: "custom" as const,
+			customType: "status",
+			content: "display only",
+			display: true,
+			excludeFromContext: true,
+			timestamp: Date.now() - 5,
+		};
 		agent.state.messages = [
 			{
 				role: "user",
@@ -639,6 +574,7 @@ describe("Agent", () => {
 				timestamp: Date.now() - 10,
 			},
 			createAssistantMessage("Initial response"),
+			excludedCustom,
 		];
 
 		agent.followUp({
@@ -656,6 +592,7 @@ describe("Agent", () => {
 		});
 
 		expect(hasQueuedFollowUp).toBe(true);
+		expect(agent.state.messages).toContain(excludedCustom);
 		expect(agent.state.messages[agent.state.messages.length - 1].role).toBe("assistant");
 	});
 

packages/agent/test/harness/compaction.test.ts

@@ -8,7 +8,6 @@ import {
 	type Usage,
 } from "@earendil-works/pi-ai";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
-import { prepareBranchEntries } from "../../src/harness/compaction/branch-summarization.ts";
 import {
 	type CompactionPreparation,
 	calculateContextTokens,
@@ -373,9 +372,19 @@ describe("harness compaction", () => {
 			details: { readFiles: ["old-read.ts"], modifiedFiles: ["old-edit.ts"] },
 		};
 		const u2 = createMessageEntry(createUserMessage("large turn"), compaction1.id);
-		const a2 = createMessageEntry(createAssistantMessage("large assistant message"), u2.id);
+		const excludedCustom: CustomMessageEntry = {
+			type: "custom_message",
+			id: createId(),
+			parentId: u2.id,
+			timestamp: new Date().toISOString(),
+			customType: "status",
+			content: "tool is running",
+			display: true,
+			excludeFromContext: true,
+		};
+		const a2 = createMessageEntry(createAssistantMessage("large assistant message"), excludedCustom.id);
 		const preparation = getOrThrow(
-			prepareCompaction([u1, a1, compaction1, u2, a2], {
+			prepareCompaction([u1, a1, compaction1, u2, excludedCustom, a2], {
 				enabled: true,
 				reserveTokens: 100,
 				keepRecentTokens: 1,
@@ -389,76 +398,6 @@ describe("harness compaction", () => {
 		expect([...preparation!.fileOps.written]).toContain("written.ts");
 	});
 
-	it("ignores excluded custom messages during compaction and branch-summary preparation", () => {
-		const excludedCustom: CustomMessageEntry = {
-			type: "custom_message",
-			id: createId(),
-			parentId: null,
-			timestamp: new Date().toISOString(),
-			customType: "status",
-			content: "x".repeat(1000),
-			display: true,
-			excludeFromContext: true,
-		};
-		const user = createMessageEntry(createUserMessage("keep"), excludedCustom.id);
-
-		const simplePreparation = getOrThrow(
-			prepareCompaction([excludedCustom, user], { enabled: true, reserveTokens: 0, keepRecentTokens: 1 }),
-		);
-		const branchPreparation = prepareBranchEntries([excludedCustom, user], 10);
-
-		expect(simplePreparation?.tokensBefore).toBe(1);
-		expect(simplePreparation?.messagesToSummarize).toEqual([]);
-		expect(branchPreparation.messages.map((message) => message.role)).toEqual(["user"]);
-		expect(branchPreparation.totalTokens).toBe(1);
-
-		const splitUser = createMessageEntry(createUserMessage("inspect file"));
-		const assistantWithToolCall = createMessageEntry(
-			{
-				...createAssistantMessage("calling tool"),
-				content: [{ type: "toolCall", id: "call-1", name: "read", arguments: { path: "file.ts" } }],
-			},
-			splitUser.id,
-		);
-		const splitExcludedCustom: CustomMessageEntry = {
-			type: "custom_message",
-			id: createId(),
-			parentId: assistantWithToolCall.id,
-			timestamp: new Date().toISOString(),
-			customType: "status",
-			content: "tool is running",
-			display: true,
-			excludeFromContext: true,
-		};
-		const toolResult = createMessageEntry(
-			{
-				role: "toolResult",
-				toolCallId: "call-1",
-				toolName: "read",
-				content: [{ type: "text", text: "x".repeat(1000) }],
-				isError: false,
-				timestamp: Date.now(),
-			},
-			splitExcludedCustom.id,
-		);
-		const assistantFinal = createMessageEntry(createAssistantMessage("done"), toolResult.id);
-		const splitPreparation = getOrThrow(
-			prepareCompaction([splitUser, assistantWithToolCall, splitExcludedCustom, toolResult, assistantFinal], {
-				enabled: true,
-				reserveTokens: 0,
-				keepRecentTokens: 1,
-			}),
-		);
-
-		expect(splitPreparation?.isSplitTurn).toBe(true);
-		expect(splitPreparation?.firstKeptEntryId).toBe(assistantFinal.id);
-		expect(splitPreparation?.turnPrefixMessages.map((message) => message.role)).toEqual([
-			"user",
-			"assistant",
-			"toolResult",
-		]);
-	});
-
 	it("prepares custom and branch summary entries for summarization", () => {
 		const branchSummary: BranchSummaryEntry = {
 			type: "branch_summary",

packages/coding-agent/test/compaction.test.ts

@@ -4,7 +4,6 @@ import { getModel } from "@earendil-works/pi-ai";
 import { readFileSync } from "fs";
 import { join } from "path";
 import { beforeEach, describe, expect, it } from "vitest";
-import { prepareBranchEntries } from "../src/core/compaction/branch-summarization.ts";
 import {
 	type CompactionSettings,
 	calculateContextTokens,
@@ -19,7 +18,6 @@ import {
 import {
 	buildSessionContext,
 	type CompactionEntry,
-	type CustomMessageEntry,
 	type ModelChangeEntry,
 	migrateSessionEntries,
 	parseSessionEntries,
@@ -94,22 +92,6 @@ function createMessageEntry(message: AgentMessage): SessionMessageEntry {
 	return entry;
 }
 
-function createCustomMessageEntry(content: string, excludeFromContext?: boolean): CustomMessageEntry {
-	const id = `test-id-${entryCounter++}`;
-	const entry: CustomMessageEntry = {
-		type: "custom_message",
-		id,
-		parentId: lastId,
-		timestamp: new Date().toISOString(),
-		customType: "status",
-		content,
-		display: true,
-		excludeFromContext,
-	};
-	lastId = id;
-	return entry;
-}
-
 function createCompactionEntry(summary: string, firstKeptEntryId: string): CompactionEntry {
 	const id = `test-id-${entryCounter++}`;
 	const entry: CompactionEntry = {
@@ -413,59 +395,6 @@ describe("buildSessionContext", () => {
 	});
 });
 
-describe("prepareCompaction with custom messages", () => {
-	it("should ignore excluded custom messages in compaction and branch-summary preparation", () => {
-		const excludedCustom = createCustomMessageEntry("x".repeat(1000), true);
-		const user = createMessageEntry(createUserMessage("keep"));
-		const simplePreparation = prepareCompaction([excludedCustom, user], {
-			enabled: true,
-			reserveTokens: 0,
-			keepRecentTokens: 1,
-		});
-		const branchPreparation = prepareBranchEntries([excludedCustom, user], 10);
-
-		expect(simplePreparation).toBeDefined();
-		expect(simplePreparation!.tokensBefore).toBe(1);
-		expect(simplePreparation!.messagesToSummarize).toEqual([]);
-		expect(branchPreparation.messages.map((message) => message.role)).toEqual(["user"]);
-		expect(branchPreparation.totalTokens).toBe(1);
-
-		resetEntryCounter();
-		const splitUser = createMessageEntry(createUserMessage("inspect file"));
-		const assistantWithToolCall = createMessageEntry({
-			...createAssistantMessage("calling tool"),
-			content: [{ type: "toolCall", id: "call-1", name: "read", arguments: { path: "file.ts" } }],
-		});
-		const splitExcludedCustom = createCustomMessageEntry("tool is running", true);
-		const toolResult = createMessageEntry({
-			role: "toolResult",
-			toolCallId: "call-1",
-			toolName: "read",
-			content: [{ type: "text", text: "x".repeat(1000) }],
-			isError: false,
-			timestamp: Date.now(),
-		});
-		const assistantFinal = createMessageEntry(createAssistantMessage("done"));
-		const splitPreparation = prepareCompaction(
-			[splitUser, assistantWithToolCall, splitExcludedCustom, toolResult, assistantFinal],
-			{
-				enabled: true,
-				reserveTokens: 0,
-				keepRecentTokens: 1,
-			},
-		);
-
-		expect(splitPreparation).toBeDefined();
-		expect(splitPreparation!.isSplitTurn).toBe(true);
-		expect(splitPreparation!.firstKeptEntryId).toBe(assistantFinal.id);
-		expect(splitPreparation!.turnPrefixMessages.map((message) => message.role)).toEqual([
-			"user",
-			"assistant",
-			"toolResult",
-		]);
-	});
-});
-
 describe("prepareCompaction with previous compaction", () => {
 	it("should preserve kept messages across repeated compactions when they still fit", () => {
 		const u1 = createMessageEntry(createUserMessage("user msg 1 (summarized by compaction1)"));

packages/coding-agent/test/suite/agent-session-queue.test.ts

@@ -258,108 +258,6 @@ describe("AgentSession queue characterization", () => {
 		expect(getAssistantTexts(harness)).toEqual(["", "original turn complete", "batched follow-up response"]);
 	});
 
-	it("records excluded custom messages immediately without starting provider turns", async () => {
-		const idleHarness = await createHarness();
-		harnesses.push(idleHarness);
-		let providerCalled = false;
-		idleHarness.setResponses([
-			() => {
-				providerCalled = true;
-				return fauxAssistantMessage("unexpected");
-			},
-		]);
-
-		await idleHarness.session.sendCustomMessage(
-			{ customType: "status", content: "display only", display: true, details: {}, excludeFromContext: true },
-			{ triggerTurn: true },
-		);
-
-		expect(providerCalled).toBe(false);
-		expect(idleHarness.session.messages[0]).toMatchObject({
-			role: "custom",
-			customType: "status",
-			excludeFromContext: true,
-		});
-		expect(idleHarness.getPendingResponseCount()).toBe(1);
-
-		for (const deliverAs of ["steer", "followUp"] as const) {
-			const waiting = await createWaitingHarness();
-			const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
-			harnesses.push(harness);
-			let providerCalledForQueuedMessage = false;
-			harness.setResponses([
-				fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
-				fauxAssistantMessage("done"),
-				() => {
-					providerCalledForQueuedMessage = true;
-					return fauxAssistantMessage("unexpected queued turn");
-				},
-			]);
-
-			await waitForToolStart;
-			await harness.session.sendCustomMessage(
-				{
-					customType: "status",
-					content: `${deliverAs} display only`,
-					display: true,
-					details: {},
-					excludeFromContext: true,
-				},
-				{ deliverAs },
-			);
-			const recordedBeforeRelease = harness.session.messages.some(
-				(message) => message.role === "custom" && message.customType === "status",
-			);
-			releaseToolExecution();
-			await promptPromise;
-
-			expect(recordedBeforeRelease).toBe(true);
-			expect(
-				harness.session.messages.filter((message) => message.role === "custom" && message.customType === "status"),
-			).toHaveLength(1);
-			expect(providerCalledForQueuedMessage).toBe(false);
-			expect(harness.getPendingResponseCount()).toBe(1);
-		}
-	});
-
-	it("persists excluded custom messages from message_end hooks after the triggering message", async () => {
-		const harness = await createHarness({
-			extensionFactories: [
-				(pi) => {
-					pi.on("message_end", (event) => {
-						if (event.message.role !== "assistant") return;
-						pi.sendMessage({
-							customType: "status",
-							content: "display only",
-							display: true,
-							details: {},
-							excludeFromContext: true,
-						});
-					});
-				},
-			],
-		});
-		harnesses.push(harness);
-		harness.setResponses([fauxAssistantMessage("reply")]);
-
-		await harness.session.prompt("hello");
-
-		const stateOrder = harness.session.messages.map((message) =>
-			message.role === "custom" ? `custom:${message.customType}` : message.role,
-		);
-		const branchOrder = harness.sessionManager.getBranch().map((entry) => {
-			if (entry.type === "message") {
-				return entry.message.role;
-			}
-			if (entry.type === "custom_message") {
-				return `custom:${entry.customType}`;
-			}
-			return entry.type;
-		});
-		expect(stateOrder).toEqual(["user", "assistant", "custom:status"]);
-		expect(branchOrder).toEqual(["user", "assistant", "custom:status"]);
-	});
-
 	it("queues custom messages with deliverAs steer while streaming", async () => {
 		const waiting = await createWaitingHarness();
 		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;

packages/coding-agent/test/suite/agent-session-retry-events.test.ts

@@ -72,6 +72,13 @@ describe("AgentSession retry and event characterization", () => {
 			harness.session.messages.some((message) => message.role === "custom" && message.customType === "status"),
 		).toBe(true);
 		expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("assistant");
+		expect(
+			harness.sessionManager.getBranch().map((entry) => {
+				if (entry.type === "message") return entry.message.role;
+				if (entry.type === "custom_message") return `custom:${entry.customType}`;
+				return entry.type;
+			}),
+		).toEqual(["user", "assistant", "custom:status", "assistant"]);
 	});
 
 	it("retries multiple transient failures and succeeds on the final attempt", async () => {