diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 18f0d91cd..0a7aae43a 100644 --- a/packages/agent/test/agent.test.ts +++ b/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 { @@ -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"); }); diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index 19826e013..57d636279 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/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", diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 3a9294aa8..929d06ec5 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/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)")); 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 129f4c4bc..56f690a96 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/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; diff --git a/packages/coding-agent/test/suite/agent-session-retry-events.test.ts b/packages/coding-agent/test/suite/agent-session-retry-events.test.ts index a8d3ae553..74da5de66 100644 --- a/packages/coding-agent/test/suite/agent-session-retry-events.test.ts +++ b/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 () => {