test(coding-agent): cover subagent append-only context rewrites

Add a createAgentSession-based subagent regression for issue #3406. The test
runs a taskDepth=1 session against a loopback llama.cpp-style model so
provider.appendOnlyContext auto-enables through the SDK path used by task
subagents.

A real context extension rewrites the prior assistant turn on the second
subagent request. The captured provider contexts assert the first request's user
message object is reused in the second request, proving the append-only log kept
the stable prefix instead of clearing and re-rendering it for subagents.

Fixes #3406
This commit is contained in:
roboomp
2026-06-24 23:00:19 +00:00
parent cce627ee4b
commit a2471339a7
@@ -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<Api>) as Model<Api>;
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"));