Skip to content

fix(pi): keep runtime context at system boundary #1002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Tarun-joy wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(pi): keep runtime context at system boundary #1002

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6d689b4
fix(pi): hide runtime context from chat
Tarun-joy Jul 27, 2026
4a6328a
fix(pi): keep runtime context at system boundary
Tarun-joy Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 4 additions & 37 deletions src/adapters/pi/extension.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,10 +160,6 @@ let _buildAutoInjection:
| null
| undefined = undefined;

// Pending context to inject via the 'context' hook (avoiding systemPrompt mutation
// which breaks prefix prompt cache on DeepSeek/Anthropic/OpenAI).
// See: https://github.com/mksglu/context-mode/issues/598
let _pendingContext = "";
async function getAutoInjection(
pluginRoot: string,
): Promise<((events: Array<{ category: string; data: string }>) => string) | null> {
Expand Down Expand Up @@ -606,7 +602,6 @@ export default function piExtension(pi: any): void {

pi.on("before_agent_start", async (event: any, ctx: any) => {
try {
_pendingContext = ""; // Reset — will be filled below if events exist
// Lazily start and await the MCP bridge only when Pi is about to
// dispatch a real agent turn. This is the non-brittle #534/#809 guard:
// help/version/package/config CLI paths may load the extension, but they
Expand Down Expand Up @@ -715,43 +710,15 @@ export default function piExtension(pi: any): void {
db.markResumeConsumed(_sessionId);
}

// Store extra context (routing anchor, active_memory, resume, behavioralDirective)
// for injection via the 'context' hook as a message, NOT as a systemPrompt
// modification. Mutating systemPrompt breaks prefix prompt caching on
// DeepSeek/Anthropic/OpenAI because the system message sits at messages[0]
// and any change invalidates the entire cache chain.
const baseLen = existingPrompt ? 1 : 0;
if (parts.length > baseLen) {
const extraParts = parts.slice(baseLen);
_pendingContext = extraParts.join("\n\n");
} else {
_pendingContext = "";
}
// Keep runtime instructions at Pi's host-owned model boundary. Injecting
// them as user/custom messages makes automation appear as user-authored
// chat in API consumers even when the TUI honors display: false.
return { systemPrompt: parts.join("\n\n") };
} catch {
_pendingContext = ""; // Reset — ensure no stale data escapes
// best effort — never break agent start
}
});

// ── 4a2. context — Inject active_memory + resume + behavioralDirective as message ──
// Uses the 'context' hook (like hindsight does) to append context at the END of
// messages rather than mutating systemPrompt at the beginning. This preserves
// prefix prompt cache for DeepSeek, Anthropic, and OpenAI.
pi.on("context", (event: any) => {
try {
if (!_pendingContext) return;
const ctx = _pendingContext;
_pendingContext = "";
event.messages.push({
role: "user",
content: ctx,
});
return { messages: event.messages };
} catch {
// best effort — never break context assembly
}
});

// ── 4b. before_provider_response — capture response metadata ───
// Pi-2: Register the missing event so providers can record latency,
// model, and token usage when Pi exposes them. Best-effort only;
Expand Down
117 changes: 39 additions & 78 deletions tests/pi-extension.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -593,15 +593,13 @@ describe("Pi Extension", () => {
{ sessionManager: { getSessionFile: () => sessionFile } }, // ctx (2nd arg)
);

// Verify the session was initialised with the file-derived ID by checking
// that before_agent_start doesn't blow up (it needs a valid _sessionId).
// In the new behavior, before_agent_start no longer returns systemPrompt;
// it stores context in _pendingContext for the context hook to inject.
// Verify the session was initialised with the file-derived ID and that
// routing instructions use Pi's model-only system boundary.
const result = await api._trigger("before_agent_start", {
systemPrompt: "Base.",
});
// before_agent_start may or may not return a value — the key is it doesn't throw
expect(result?.systemPrompt ?? null).toBe(null); // systemPrompt is no longer returned
expect(result?.systemPrompt).toContain("Base.");
expect(result?.systemPrompt).toContain("context-mode active");
});

it("handles session lifecycle in correct order", async () => {
Expand Down Expand Up @@ -637,7 +635,7 @@ describe("Pi Extension", () => {
// ═══════════════════════════════════════════════════════════

describe("Slice 5: Resume injection", () => {
it("delivers resume snapshot through context hook, not systemPrompt", async () => {
it("delivers resume snapshot through the model-only system prompt", async () => {
await registerPiExtension(api);

// Build up session state: capture events → compact → build resume
Expand All @@ -662,28 +660,16 @@ describe("Pi Extension", () => {
await api._trigger("session_before_compact", {});
await api._trigger("session_compact", {});

// before_agent_start should prepare context without mutating systemPrompt
// before_agent_start should deliver the resume at the model-only system boundary.
const result = await api._trigger("before_agent_start", {
systemPrompt: "You are a helpful assistant.",
});
expect(result?.systemPrompt ?? null).toBe(null);

// The context hook should deliver the resume as a trailing user message.
const messages = [{ role: "system", content: "You are a helpful assistant." }];
const ctxResult = await api._trigger("context", { messages });

expect(ctxResult?.messages).toBe(messages);
expect(messages).toHaveLength(2);
expect(messages[0]).toEqual({
role: "system",
content: "You are a helpful assistant.",
});
expect(messages[1].role).toBe("user");
expect(String(messages[1].content)).toContain("session_resume");
expect(String(messages[1].content)).not.toContain("You are a helpful assistant.");
expect(result?.systemPrompt).toContain("You are a helpful assistant.");
expect(result?.systemPrompt).toContain("session_resume");
});

it("appends resume and active context after existing messages", async () => {
it("appends resume and active context after the existing system prompt", async () => {
await registerPiExtension(api);

await api._trigger("session_start", {
Expand All @@ -710,25 +696,11 @@ describe("Pi Extension", () => {
systemPrompt: "Stable system prompt.",
prompt: "Continue with the refactor and avoid lodash.",
});
expect(result?.systemPrompt ?? null).toBe(null);

const messages = [
{ role: "system", content: "Stable system prompt." },
{ role: "user", content: "Continue with the refactor." },
];
const ctxResult = await api._trigger("context", { messages });

expect(ctxResult?.messages).toBe(messages);
expect(messages).toHaveLength(3);
expect(messages[0]).toEqual({ role: "system", content: "Stable system prompt." });
expect(messages[1]).toEqual({ role: "user", content: "Continue with the refactor." });
expect(messages[2].role).toBe("user");

const trailing = String(messages[2].content);
expect(trailing).toContain("context-mode active");
expect(trailing).toContain("session_resume");
expect(trailing).toContain("how_to_search");
expect(trailing).not.toContain("Stable system prompt.");

expect(result?.systemPrompt).toContain("Stable system prompt.");
expect(result?.systemPrompt).toContain("context-mode active");
expect(result?.systemPrompt).toContain("session_resume");
expect(result?.systemPrompt).toContain("how_to_search");
});

it("returns nothing when no resume exists", async () => {
Expand Down Expand Up @@ -901,29 +873,23 @@ describe("Pi Extension", () => {
// ═══════════════════════════════════════════════════════════

describe("Slice 7: Routing block injection", () => {
it("injects lightweight routing anchor via context hook on first before_agent_start", async () => {
it("injects lightweight routing anchor at the model-only system boundary", async () => {
await registerPiExtension(api);
await api._trigger("session_start", {}, {
sessionManager: { getSessionFile: () => `routing-1-${Date.now()}-${Math.random()}` },
});

// before_agent_start sets _pendingContext (was previously modifying systemPrompt)
await api._trigger("before_agent_start", {
const result = await api._trigger("before_agent_start", {
systemPrompt: "Base prompt.",
});

// context hook now injects the routing anchor as a user message at message end
const messages: any[] = [];
const ctxResult = await api._trigger("context", { messages });

expect(ctxResult?.messages).toBeDefined();
expect(ctxResult.messages.length).toBe(1);
expect(ctxResult.messages[0].role).toBe("user");
expect(ctxResult.messages[0].content).toContain("context-mode active");
expect(ctxResult.messages[0].content).toContain("ctx_batch_execute > ctx_execute > ctx_execute_file");
expect(result?.systemPrompt).toContain("Base prompt.");
expect(result?.systemPrompt).toContain("context-mode active");
expect(result?.systemPrompt).toContain("ctx_batch_execute > ctx_execute > ctx_execute_file");
expect(api._handlers.context).toBeUndefined();
});

it("re-injects the anchor via context hook on every subsequent call", async () => {
it("re-injects the anchor on every subsequent call", async () => {
await registerPiExtension(api);
await api._trigger("session_start", {}, {
sessionManager: { getSessionFile: () => `routing-2-${Date.now()}-${Math.random()}` },
Expand All @@ -935,10 +901,10 @@ describe("Pi Extension", () => {
const ANCHOR = "context-mode active";

for (let call = 0; call < 3; call++) {
await api._trigger("before_agent_start", { systemPrompt: "Base." });
const ctxResult = await api._trigger("context", { messages: [] });
expect(ctxResult?.messages).toBeDefined();
expect(ctxResult.messages[0]?.content).toContain(ANCHOR);
const result = await api._trigger("before_agent_start", {
systemPrompt: "Base.",
});
expect(result?.systemPrompt).toContain(ANCHOR);
}
});
});
Expand Down Expand Up @@ -984,7 +950,7 @@ describe("Pi Extension", () => {
// ═══════════════════════════════════════════════════════════

describe("Slice 9: active_memory injection", () => {
it("injects context every turn via context hook even when compact_count is 0", async () => {
it("injects context every turn even when compact_count is 0", async () => {
await registerPiExtension(api);
await api._trigger("session_start", {
sessionManager: { getSessionFile: () => `active-mem-1-${Date.now()}-${Math.random()}` },
Expand All @@ -997,25 +963,19 @@ describe("Pi Extension", () => {
});

// Second call rebuilds context (always-on, not just post-compaction).
await api._trigger("before_agent_start", {
const result = await api._trigger("before_agent_start", {
systemPrompt: "Base 2.",
});

// context hook injects the pending context as a user message
const ctxResult = await api._trigger("context", { messages: [] });

expect(ctxResult?.messages).toBeDefined();
expect(ctxResult.messages.length).toBe(1);
expect(ctxResult.messages[0].role).toBe("user");
// The always-on injection path fires every turn — the routing anchor
// proves context reaches the model even with compact_count 0.
const content = String(ctxResult.messages[0].content);
const content = String(result?.systemPrompt ?? "");
expect(content).toContain("context-mode active");
// Issue #856 — the role MUST NOT be pinned as a standing directive.
expect(content).not.toContain("<behavioral_directive>");
});

it("caps active_memory at ≤ 2000 characters (via context hook)", async () => {
it("caps active_memory at ≤ 2000 characters", async () => {
await registerPiExtension(api);
await api._trigger("session_start", {
sessionManager: { getSessionFile: () => `active-mem-2-${Date.now()}-${Math.random()}` },
Expand All @@ -1030,12 +990,11 @@ describe("Pi Extension", () => {
});
}

await api._trigger("before_agent_start", {
const result = await api._trigger("before_agent_start", {
systemPrompt: "Base final.",
});

const ctxResult = await api._trigger("context", { messages: [] });
const content = String(ctxResult?.messages?.[0]?.content ?? "");
const content = String(result?.systemPrompt ?? "");
// Issue #856 — flooding with role prompts must NOT accumulate any
// behavioral_directive, and the per-turn injection must stay bounded
// (it is now just the routing anchor; roles are filtered out entirely).
Expand Down Expand Up @@ -1068,9 +1027,10 @@ describe("Pi Extension", () => {
});

// Subsequent turn rebuilds context.
await api._trigger("before_agent_start", { systemPrompt: "Base 2." });
const ctxResult = await api._trigger("context", { messages: [] });
const content = String(ctxResult?.messages?.[0]?.content ?? "");
const result = await api._trigger("before_agent_start", {
systemPrompt: "Base 2.",
});
const content = String(result?.systemPrompt ?? "");

// The stale role must NOT be pinned as a standing behavioral_directive.
expect(content).not.toContain("<behavioral_directive>");
Expand All @@ -1087,9 +1047,10 @@ describe("Pi Extension", () => {
prompt: "You are a senior staff engineer reviewing this codebase.",
systemPrompt: "Base.",
});
await api._trigger("before_agent_start", { systemPrompt: "Base 2." });
const ctxResult = await api._trigger("context", { messages: [] });
const content = String(ctxResult?.messages?.[0]?.content ?? "");
const result = await api._trigger("before_agent_start", {
systemPrompt: "Base 2.",
});
const content = String(result?.systemPrompt ?? "");

// Role filtered out…
expect(content).not.toContain("<behavioral_directive>");
Expand Down