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:
@@ -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"));
|
||||
|
||||
Reference in New Issue
Block a user