diff --git a/packages/coding-agent/test/agent-session-message-pipeline.test.ts b/packages/coding-agent/test/agent-session-message-pipeline.test.ts index 96ba0e79e..4c1d4db3a 100644 --- a/packages/coding-agent/test/agent-session-message-pipeline.test.ts +++ b/packages/coding-agent/test/agent-session-message-pipeline.test.ts @@ -14,12 +14,15 @@ import { } from "@oh-my-pi/pi-ai"; import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; import { buildModel } from "@oh-my-pi/pi-catalog/build"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import * as memoryBackend from "@oh-my-pi/pi-coding-agent/memory-backend"; import type { MemoryBackend } from "@oh-my-pi/pi-coding-agent/memory-backend/types"; import { type MnemopiSessionState, setMnemopiSessionState } from "@oh-my-pi/pi-coding-agent/mnemopi/state"; +import { createAgentSession, type ExtensionFactory } from "@oh-my-pi/pi-coding-agent/sdk"; import { obfuscateProviderContext, SecretObfuscator } from "@oh-my-pi/pi-coding-agent/secrets"; 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 { convertToLlm, wrapSteeringForModel } from "@oh-my-pi/pi-coding-agent/session/messages"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { TempDir } from "@oh-my-pi/pi-utils"; @@ -689,6 +692,94 @@ describe("AgentSession message pipeline", () => { expect(contexts[1]!.systemPrompt).toEqual(firstSystemPrompt); }); + it("preserves append-only prefixes in subagent sessions when context handlers rewrite prior turns", async () => { + using tempDir = TempDir.createSync("@pi-subagent-append-only-"); + const api = "test-subagent-append-only-cache"; + const contexts: Context[] = []; + registerCustomApi(api, (_model, context) => { + contexts.push(context); + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + const message = createAssistantMessage(`ok-${contexts.length}`); + stream.push({ type: "text_delta", contentIndex: 0, delta: "ok", partial: message }); + stream.push({ type: "done", reason: "stop", message }); + }); + return stream; + }); + const model = buildModel({ + id: "local-subagent-model", + name: "Local Subagent Model", + api, + provider: "llama.cpp", + baseUrl: "http://127.0.0.1:8080/v1", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 4096, + maxTokens: 1024, + } as ModelSpec) as Model; + const rewritePriorAssistant: ExtensionFactory = pi => { + pi.on("context", async event => { + const hasSecondTurn = event.messages.some(message => { + if (message.role !== "user") return false; + const content = message.content; + if (typeof content === "string") return content.includes("second"); + return content.some(part => part.type === "text" && part.text.includes("second")); + }); + if (!hasSecondTurn) return undefined; + return { + messages: event.messages.map(message => + message.role === "assistant" + ? { ...message, content: [{ type: "text" as const, text: "rewritten assistant" }] } + : message, + ), + }; + }); + }; + const authStorage = await AuthStorage.create(tempDir.join("auth.db")); + const modelRegistry = new ModelRegistry(authStorage, tempDir.join("models.yml")); + const { session } = await createAgentSession({ + cwd: tempDir.path(), + agentDir: tempDir.path(), + sessionManager: SessionManager.inMemory(tempDir.path()), + authStorage, + modelRegistry, + settings: Settings.isolated({ + "compaction.enabled": false, + "provider.appendOnlyContext": "auto", + }), + model, + disableExtensionDiscovery: true, + extensions: [rewritePriorAssistant], + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + skipPythonPreflight: true, + taskDepth: 1, + agentId: "SubAgent", + }); + try { + expect(session.agent.appendOnlyContext).toBeDefined(); + + await session.sendUserMessage("first"); + await session.sendUserMessage("second"); + + expect(contexts).toHaveLength(2); + expect(contexts[0]!.messages).toHaveLength(1); + expect(contexts[1]!.messages).toHaveLength(3); + expect(contexts[1]!.messages[0]).toBe(contexts[0]!.messages[0]); + expect((contexts[1]!.messages[1] as { content: unknown }).content).toEqual([ + { type: "text", text: "rewritten assistant" }, + ]); + } finally { + await session.dispose(); + authStorage.close(); + } + }); + it("clears promoted memory from the base prompt when switching sessions", async () => { using tempDir = TempDir.createSync("@pi-injected-memory-switch-"); const sessionManager = SessionManager.create(tempDir.path(), tempDir.join("sessions"));