From 0264730fbcfea817068c9994f5f260485abf37eb Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Fri, 12 Jun 2026 23:53:12 +0200 Subject: [PATCH] feat(coding-agent): exclude custom messages from context --- packages/agent/src/harness/agent-harness.ts | 8 +- .../compaction/branch-summarization.ts | 10 +- .../src/harness/compaction/compaction.ts | 15 ++- packages/agent/src/harness/messages.ts | 6 ++ packages/agent/src/harness/session/session.ts | 3 + packages/agent/src/harness/types.ts | 1 + .../agent/test/harness/compaction.test.ts | 49 ++++++++++ .../coding-agent/src/core/agent-session.ts | 35 ++++--- .../core/compaction/branch-summarization.ts | 10 +- .../src/core/compaction/compaction.ts | 23 ++++- .../coding-agent/src/core/extensions/types.ts | 8 +- packages/coding-agent/src/core/messages.ts | 7 ++ .../coding-agent/src/core/session-manager.ts | 14 ++- .../test/branch-summarization.test.ts | 47 ++++++++++ packages/coding-agent/test/compaction.test.ts | 55 +++++++++++ .../agent-session-bash-persistence.test.ts | 36 +++++++- .../test/suite/agent-session-queue.test.ts | 92 +++++++++++++++++++ 17 files changed, 392 insertions(+), 27 deletions(-) create mode 100644 packages/coding-agent/test/branch-summarization.test.ts diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 965634653..57fe29c14 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -495,7 +495,13 @@ export class AgentHarness< } else if (write.type === "custom") { await this.session.appendCustomEntry(write.customType, write.data); } else if (write.type === "custom_message") { - await this.session.appendCustomMessageEntry(write.customType, write.content, write.display, write.details); + await this.session.appendCustomMessageEntry( + write.customType, + write.content, + write.display, + write.details, + write.excludeFromContext, + ); } else if (write.type === "label") { await this.session.appendLabel(write.targetId, write.label); } else if (write.type === "session_info") { diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts index c1824ebf8..4389aa6ab 100644 --- a/packages/agent/src/harness/compaction/branch-summarization.ts +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -103,7 +103,14 @@ function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined return entry.message; case "custom_message": - return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + return createCustomMessage( + entry.customType, + entry.content, + entry.display, + entry.details, + entry.timestamp, + entry.excludeFromContext, + ); case "branch_summary": return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); @@ -143,6 +150,7 @@ export function prepareBranchEntries(entries: SessionTreeEntry[], tokenBudget: n const entry = entries[i]; const message = getMessageFromEntry(entry); if (!message) continue; + if (message.role === "custom" && message.excludeFromContext) continue; extractFileOpsFromMessage(message, fileOps); const tokens = estimateTokens(message); diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index dba753d78..7582b696b 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -68,6 +68,7 @@ function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined entry.display, entry.details, entry.timestamp, + entry.excludeFromContext, ); } if (entry.type === "branch_summary") { @@ -83,7 +84,11 @@ function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage if (entry.type === "compaction") { return undefined; } - return getMessageFromEntry(entry); + const message = getMessageFromEntry(entry); + if (message?.role === "custom" && message.excludeFromContext) { + return undefined; + } + return message; } /** Generated compaction data ready to be persisted as a compaction entry. */ @@ -240,7 +245,13 @@ export function estimateTokens(message: AgentMessage): number { } return Math.ceil(chars / 4); } - case "custom": + case "custom": { + if (message.excludeFromContext) { + return 0; + } + chars = estimateTextAndImageContentChars(message.content); + return Math.ceil(chars / 4); + } case "toolResult": { chars = estimateTextAndImageContentChars(message.content); return Math.ceil(chars / 4); diff --git a/packages/agent/src/harness/messages.ts b/packages/agent/src/harness/messages.ts index 36ce96a1e..34f2f76b4 100644 --- a/packages/agent/src/harness/messages.ts +++ b/packages/agent/src/harness/messages.ts @@ -34,6 +34,7 @@ export interface CustomMessage { content: string | (TextContent | ImageContent)[]; display: boolean; details?: T; + excludeFromContext?: boolean; timestamp: number; } @@ -106,6 +107,7 @@ export function createCustomMessage( display: boolean, details: unknown | undefined, timestamp: string, + excludeFromContext?: boolean, ): CustomMessage { return { role: "custom", @@ -113,6 +115,7 @@ export function createCustomMessage( content, display, details, + excludeFromContext, timestamp: new Date(timestamp).getTime(), }; } @@ -131,6 +134,9 @@ export function convertToLlm(messages: AgentMessage[]): Message[] { timestamp: m.timestamp, }; case "custom": { + if (m.excludeFromContext) { + return undefined; + } const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; return { role: "user", diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index 6f136208a..4c79b4bae 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -51,6 +51,7 @@ export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionCon entry.display, entry.details, entry.timestamp, + entry.excludeFromContext, ), ); } else if (entry.type === "branch_summary" && entry.summary) { @@ -206,6 +207,7 @@ export class Session { content: string | (TextContent | ImageContent)[], display: boolean, details?: T, + excludeFromContext?: boolean, ): Promise { return this.appendTypedEntry({ type: "custom_message", @@ -216,6 +218,7 @@ export class Session { content, display, details, + excludeFromContext, } satisfies CustomMessageEntry); } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 4756ca841..d4a44b0e9 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -388,6 +388,7 @@ export interface CustomMessageEntry extends SessionTreeEntryBase { content: string | (TextContent | ImageContent)[]; details?: T; display: boolean; + excludeFromContext?: boolean; } export interface LabelEntry extends SessionTreeEntryBase { diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index b694d9f5b..abcf1b5fe 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -8,6 +8,7 @@ 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, @@ -289,9 +290,12 @@ describe("harness compaction", () => { timestamp: Date.now(), }; + const excludedCustom: AgentMessage = { ...customString, content: "x".repeat(1000), excludeFromContext: true }; + expect(estimateTokens({ role: "user", content: "plain user", timestamp: Date.now() })).toBeGreaterThan(0); expect(estimateTokens(assistantWithThinkingAndTool)).toBeGreaterThan(0); expect(estimateTokens(customString)).toBeGreaterThan(0); + expect(estimateTokens(excludedCustom)).toBe(0); expect(estimateTokens(toolResultWithImage)).toBeGreaterThan(1000); expect(estimateTokens(bashExecution)).toBeGreaterThan(0); expect(estimateTokens(branchSummaryMessage)).toBeGreaterThan(0); @@ -311,6 +315,11 @@ describe("harness compaction", () => { usageTokens: 20, lastUsageIndex: 0, }); + expect(estimateContextTokens([assistant, excludedCustom])).toMatchObject({ + tokens: 20, + usageTokens: 20, + trailingTokens: 0, + }); }); it("builds session context with a compaction entry", () => { @@ -380,6 +389,46 @@ describe("harness compaction", () => { expect([...preparation!.fileOps.written]).toContain("written.ts"); }); + it("skips excluded custom messages during compaction token estimates", () => { + const customMessage: 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"), customMessage.id); + + const preparation = getOrThrow( + prepareCompaction([customMessage, user], { enabled: true, reserveTokens: 0, keepRecentTokens: 1 }), + ); + + expect(preparation?.tokensBefore).toBe(1); + expect(preparation?.messagesToSummarize).toEqual([]); + }); + + it("skips excluded custom messages before branch summary token budgeting", () => { + const user = createMessageEntry(createUserMessage("keep")); + const customMessage: CustomMessageEntry = { + type: "custom_message", + id: createId(), + parentId: user.id, + timestamp: new Date().toISOString(), + customType: "status", + content: "x".repeat(1000), + display: true, + excludeFromContext: true, + }; + + const preparation = prepareBranchEntries([user, customMessage], 10); + + expect(preparation.messages.map((message) => message.role)).toEqual(["user"]); + expect(preparation.totalTokens).toBe(1); + }); + it("prepares custom and branch summary entries for summarization", () => { const branchSummary: BranchSummaryEntry = { type: "branch_summary", diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index af9d32d7b..3a6ad0874 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -512,6 +512,7 @@ export class AgentSession { event.message.content, event.message.display, event.message.details, + event.message.excludeFromContext, ); } else if ( event.message.role === "user" || @@ -1112,6 +1113,7 @@ export class AgentSession { content: msg.content, display: msg.display, details: msg.details, + excludeFromContext: msg.excludeFromContext, timestamp: Date.now(), }); } @@ -1289,7 +1291,8 @@ export class AgentSession { /** * Send a custom message to the session. Creates a CustomMessageEntry. * - * Handles three cases: + * Handles four cases: + * - Excluded from context: appends to state/session, no turn * - Streaming: queues message, processed when loop pulls from queue * - Not streaming + triggerTurn: appends to state/session, starts new turn * - Not streaming + no trigger: appends to state/session, no turn @@ -1299,7 +1302,7 @@ export class AgentSession { * @param options.deliverAs Delivery mode: "steer", "followUp", or "nextTurn" */ async sendCustomMessage( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "excludeFromContext">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, ): Promise { const appMessage = { @@ -1308,9 +1311,12 @@ export class AgentSession { content: message.content, display: message.display, details: message.details, + excludeFromContext: message.excludeFromContext, timestamp: Date.now(), } satisfies CustomMessage; - if (options?.deliverAs === "nextTurn") { + if (appMessage.excludeFromContext) { + this._recordCustomMessage(appMessage); + } else if (options?.deliverAs === "nextTurn") { this._pendingNextTurnMessages.push(appMessage); } else if (this.isStreaming) { if (options?.deliverAs === "followUp") { @@ -1321,18 +1327,23 @@ export class AgentSession { } else if (options?.triggerTurn) { await this._runAgentPrompt(appMessage); } else { - this.agent.state.messages.push(appMessage); - this.sessionManager.appendCustomMessageEntry( - message.customType, - message.content, - message.display, - message.details, - ); - this._emit({ type: "message_start", message: appMessage }); - this._emit({ type: "message_end", message: appMessage }); + this._recordCustomMessage(appMessage); } } + private _recordCustomMessage(message: CustomMessage): void { + this.agent.state.messages.push(message); + this.sessionManager.appendCustomMessageEntry( + message.customType, + message.content, + message.display, + message.details, + message.excludeFromContext, + ); + this._emit({ type: "message_start", message }); + this._emit({ type: "message_end", message }); + } + /** * Send a user message to the agent. Always triggers a turn. * When the agent is streaming, use deliverAs to specify how to queue the message. diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index f6ce3d677..ec76220a9 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -153,7 +153,14 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { return entry.message; case "custom_message": - return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + return createCustomMessage( + entry.customType, + entry.content, + entry.display, + entry.details, + entry.timestamp, + entry.excludeFromContext, + ); case "branch_summary": return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); @@ -212,6 +219,7 @@ export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: numbe const entry = entries[i]; const message = getMessageFromEntry(entry); if (!message) continue; + if (message.role === "custom" && message.excludeFromContext) continue; // Extract file ops from assistant messages (tool calls) extractFileOpsFromMessage(message, fileOps); diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index eae19794e..46f802b70 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -81,7 +81,14 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { return entry.message; } if (entry.type === "custom_message") { - return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + return createCustomMessage( + entry.customType, + entry.content, + entry.display, + entry.details, + entry.timestamp, + entry.excludeFromContext, + ); } if (entry.type === "branch_summary") { return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); @@ -96,7 +103,11 @@ function getMessageFromEntryForCompaction(entry: SessionEntry): AgentMessage | u if (entry.type === "compaction") { return undefined; } - return getMessageFromEntry(entry); + const message = getMessageFromEntry(entry); + if (message?.role === "custom" && message.excludeFromContext) { + return undefined; + } + return message; } /** Result from compact() - SessionManager adds uuid/parentUuid when saving */ @@ -270,7 +281,13 @@ export function estimateTokens(message: AgentMessage): number { } return Math.ceil(chars / 4); } - case "custom": + case "custom": { + if (message.excludeFromContext) { + return 0; + } + chars = estimateTextAndImageContentChars(message.content); + return Math.ceil(chars / 4); + } case "toolResult": { chars = estimateTextAndImageContentChars(message.content); return Math.ceil(chars / 4); diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index a869a55d1..c09a1d6d5 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -379,7 +379,7 @@ export interface ExtensionCommandContext extends ExtensionContext { */ export interface ReplacedSessionContext extends ExtensionCommandContext { sendMessage( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "excludeFromContext">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, ): Promise; @@ -1043,7 +1043,7 @@ export interface MessageEndEventResult { } export interface BeforeAgentStartEventResult { - message?: Pick; + message?: Pick; /** Replace the system prompt for this turn. If multiple extensions return this, they are chained. */ systemPrompt?: string; } @@ -1213,7 +1213,7 @@ export interface ExtensionAPI { /** Send a custom message to the session. */ sendMessage( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "excludeFromContext">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, ): void; @@ -1442,7 +1442,7 @@ export interface ExtensionShortcut { type HandlerFn = (...args: unknown[]) => Promise; export type SendMessageHandler = ( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "excludeFromContext">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, ) => void; diff --git a/packages/coding-agent/src/core/messages.ts b/packages/coding-agent/src/core/messages.ts index 81ffabb58..57c4c8644 100644 --- a/packages/coding-agent/src/core/messages.ts +++ b/packages/coding-agent/src/core/messages.ts @@ -49,6 +49,8 @@ export interface CustomMessage { content: string | (TextContent | ImageContent)[]; display: boolean; details?: T; + /** If true, this message is excluded from LLM context */ + excludeFromContext?: boolean; timestamp: number; } @@ -126,6 +128,7 @@ export function createCustomMessage( display: boolean, details: unknown | undefined, timestamp: string, + excludeFromContext?: boolean, ): CustomMessage { return { role: "custom", @@ -133,6 +136,7 @@ export function createCustomMessage( content, display, details, + excludeFromContext, timestamp: new Date(timestamp).getTime(), }; } @@ -160,6 +164,9 @@ export function convertToLlm(messages: AgentMessage[]): Message[] { timestamp: m.timestamp, }; case "custom": { + if (m.excludeFromContext) { + return undefined; + } const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; return { role: "user", diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 62942480c..b74c4c477 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -120,7 +120,7 @@ export interface SessionInfoEntry extends SessionEntryBase { * Custom message entry for extensions to inject messages into LLM context. * Use customType to identify your extension's entries. * - * Unlike CustomEntry, this DOES participate in LLM context. + * Unlike CustomEntry, this DOES participate in LLM context unless excludeFromContext is true. * The content is converted to a user message in buildSessionContext(). * Use details for extension-specific metadata (not sent to LLM). * @@ -134,6 +134,7 @@ export interface CustomMessageEntry extends SessionEntryBase { content: string | (TextContent | ImageContent)[]; details?: T; display: boolean; + excludeFromContext?: boolean; } /** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */ @@ -390,7 +391,14 @@ export function buildSessionContext( messages.push(entry.message); } else if (entry.type === "custom_message") { messages.push( - createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp), + createCustomMessage( + entry.customType, + entry.content, + entry.display, + entry.details, + entry.timestamp, + entry.excludeFromContext, + ), ); } else if (entry.type === "branch_summary" && entry.summary) { messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)); @@ -1063,6 +1071,7 @@ export class SessionManager { content: string | (TextContent | ImageContent)[], display: boolean, details?: T, + excludeFromContext?: boolean, ): string { const entry: CustomMessageEntry = { type: "custom_message", @@ -1070,6 +1079,7 @@ export class SessionManager { content, display, details, + excludeFromContext, id: generateId(this.byId), parentId: this.leafId, timestamp: new Date().toISOString(), diff --git a/packages/coding-agent/test/branch-summarization.test.ts b/packages/coding-agent/test/branch-summarization.test.ts new file mode 100644 index 000000000..8050ddac9 --- /dev/null +++ b/packages/coding-agent/test/branch-summarization.test.ts @@ -0,0 +1,47 @@ +import { describe, expect, it } from "vitest"; +import { prepareBranchEntries } from "../src/core/compaction/branch-summarization.ts"; +import type { CustomMessageEntry, SessionMessageEntry } from "../src/core/session-manager.ts"; + +function userEntry(id: string, parentId: string | null, text: string): SessionMessageEntry { + return { + type: "message", + id, + parentId, + timestamp: "2025-01-01T00:00:00Z", + message: { + role: "user", + content: [{ type: "text", text }], + timestamp: 1, + }, + }; +} + +function customEntry( + id: string, + parentId: string | null, + content: string, + excludeFromContext: boolean, +): CustomMessageEntry { + return { + type: "custom_message", + id, + parentId, + timestamp: "2025-01-01T00:00:00Z", + customType: "status", + content, + display: true, + excludeFromContext, + }; +} + +describe("branch summarization", () => { + it("skips excluded custom messages before token budgeting", () => { + const user = userEntry("user", null, "keep"); + const excluded = customEntry("custom", user.id, "x".repeat(1000), true); + + const preparation = prepareBranchEntries([user, excluded], 10); + + expect(preparation.messages.map((message) => message.role)).toEqual(["user"]); + expect(preparation.totalTokens).toBe(1); + }); +}); diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 929d06ec5..037ff5f85 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -15,9 +15,11 @@ import { prepareCompaction, shouldCompact, } from "../src/core/compaction/index.ts"; +import type { CustomMessage } from "../src/core/messages.ts"; import { buildSessionContext, type CompactionEntry, + type CustomMessageEntry, type ModelChangeEntry, migrateSessionEntries, parseSessionEntries, @@ -92,6 +94,22 @@ 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 = { @@ -184,6 +202,27 @@ describe("Token calculation", () => { const usage = createMockUsage(0, 0, 0, 0); expect(calculateContextTokens(usage)).toBe(0); }); + + it("should ignore excluded custom messages in context token estimates", () => { + const excludedCustom: CustomMessage = { + role: "custom", + customType: "status", + content: "x".repeat(1000), + display: true, + excludeFromContext: true, + timestamp: Date.now(), + }; + const visibleCustom: CustomMessage = { ...excludedCustom, excludeFromContext: false }; + const assistant = createAssistantMessage("assistant", createMockUsage(10, 5)); + + expect(estimateContextTokens([excludedCustom])).toMatchObject({ tokens: 0, trailingTokens: 0 }); + expect(estimateContextTokens([visibleCustom]).tokens).toBeGreaterThan(0); + expect(estimateContextTokens([assistant, excludedCustom])).toMatchObject({ + tokens: 15, + usageTokens: 15, + trailingTokens: 0, + }); + }); }); describe("getLastAssistantUsage", () => { @@ -395,6 +434,22 @@ describe("buildSessionContext", () => { }); }); +describe("prepareCompaction with custom messages", () => { + it("should ignore excluded custom messages in token estimates and summarized messages", () => { + const excludedCustom = createCustomMessageEntry("x".repeat(1000), true); + const user = createMessageEntry(createUserMessage("keep")); + const preparation = prepareCompaction([excludedCustom, user], { + enabled: true, + reserveTokens: 0, + keepRecentTokens: 1, + }); + + expect(preparation).toBeDefined(); + expect(preparation!.tokensBefore).toBe(1); + expect(preparation!.messagesToSummarize).toEqual([]); + }); +}); + 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)")); diff --git a/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts b/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts index efd51cfcf..419364f03 100644 --- a/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts +++ b/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts @@ -4,7 +4,7 @@ import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it } from "vitest"; import type { BashOperations } from "../../src/core/tools/bash.ts"; -import { createHarness, type Harness } from "./harness.ts"; +import { createHarness, getMessageText, type Harness } from "./harness.ts"; function getEntryTypes(harness: Harness): string[] { return harness.sessionManager.getEntries().map((entry) => entry.type); @@ -175,6 +175,40 @@ describe("AgentSession bash and persistence characterization", () => { ]); }); + it("excludes flagged custom messages from LLM context while preserving them", async () => { + const harness = await createHarness(); + harnesses.push(harness); + let userTexts: string[] = []; + + await harness.session.sendCustomMessage({ + customType: "status", + content: "status panel", + display: true, + details: { a: 1 }, + excludeFromContext: true, + }); + harness.setResponses([ + (context) => { + userTexts = context.messages + .filter((message) => message.role === "user") + .map((message) => getMessageText(message)); + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("next prompt"); + + expect(userTexts).toEqual(["next prompt"]); + const entries = harness.sessionManager.getEntries(); + expect(entries[0]?.type).toBe("custom_message"); + if (entries[0]?.type !== "custom_message") return; + expect(entries[0].excludeFromContext).toBe(true); + expect(harness.sessionManager.buildSessionContext().messages[0]).toMatchObject({ + role: "custom", + excludeFromContext: true, + }); + }); + it("does not emit message_end for bash execution messages", async () => { const harness = await createHarness(); harnesses.push(harness); diff --git a/packages/coding-agent/test/suite/agent-session-queue.test.ts b/packages/coding-agent/test/suite/agent-session-queue.test.ts index a1c3fd8a1..ed29425b7 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/packages/coding-agent/test/suite/agent-session-queue.test.ts @@ -258,6 +258,98 @@ describe("AgentSession queue characterization", () => { expect(getAssistantTexts(harness)).toEqual(["", "original turn complete", "batched follow-up response"]); }); + it("records excluded custom messages with triggerTurn without starting a provider turn", async () => { + const harness = await createHarness(); + harnesses.push(harness); + let providerCalled = false; + harness.setResponses([ + () => { + providerCalled = true; + return fauxAssistantMessage("unexpected"); + }, + ]); + + await harness.session.sendCustomMessage( + { customType: "status", content: "display only", display: true, details: {}, excludeFromContext: true }, + { triggerTurn: true }, + ); + + expect(providerCalled).toBe(false); + expect(harness.session.messages).toHaveLength(1); + expect(harness.session.messages[0]).toMatchObject({ + role: "custom", + customType: "status", + excludeFromContext: true, + }); + expect(harness.getPendingResponseCount()).toBe(1); + }); + + it("records excluded custom messages with deliverAs steer while streaming", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + let recordedBeforeRelease = false; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + await waitForToolStart; + await harness.session.sendCustomMessage( + { customType: "status", content: "steer display only", display: true, details: {}, excludeFromContext: true }, + { deliverAs: "steer" }, + ); + 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); + }); + + it("records excluded custom messages with deliverAs followUp while streaming without starting another turn", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + let providerCalledForFollowUp = false; + let recordedBeforeRelease = false; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + () => { + providerCalledForFollowUp = true; + return fauxAssistantMessage("unexpected follow-up"); + }, + ]); + + await waitForToolStart; + await harness.session.sendCustomMessage( + { + customType: "status", + content: "follow-up display only", + display: true, + details: {}, + excludeFromContext: true, + }, + { deliverAs: "followUp" }, + ); + recordedBeforeRelease = harness.session.messages.some( + (message) => message.role === "custom" && message.customType === "status", + ); + releaseToolExecution(); + await promptPromise; + + expect(recordedBeforeRelease).toBe(true); + expect(providerCalledForFollowUp).toBe(false); + expect(harness.getPendingResponseCount()).toBe(1); + }); + it("queues custom messages with deliverAs steer while streaming", async () => { const waiting = await createWaitingHarness(); const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;