From 17cb328ca1e1210f2558fd3763faa817f83a4412 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 8 Jan 2026 19:54:34 +0100 Subject: [PATCH] Allow extensions to modify system prompt in before_agent_start - Add systemPrompt to BeforeAgentStartEvent so extensions can see current prompt - Change systemPromptAppend to systemPrompt in BeforeAgentStartEventResult for full replacement - Extensions can now chain modifications (each sees the result of previous) - Update ssh.ts to replace local cwd with remote cwd in system prompt - Update pirate.ts, claude-rules.ts, preset.ts to use new API fixes #575 --- packages/coding-agent/CHANGELOG.md | 4 ++++ packages/coding-agent/docs/extensions.md | 11 +++++---- .../examples/extensions/claude-rules.ts | 7 ++++-- .../examples/extensions/pirate.ts | 11 +++++---- .../examples/extensions/preset.ts | 4 ++-- .../coding-agent/examples/extensions/ssh.ts | 12 ++++++++++ .../coding-agent/src/core/agent-session.ts | 14 +++++++---- .../src/core/extensions/runner.ts | 24 ++++++++++++------- .../coding-agent/src/core/extensions/types.ts | 4 +++- packages/coding-agent/src/index.ts | 1 + 10 files changed, 65 insertions(+), 27 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 4b7e60cf..678e89cb 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Breaking Changes + +- `before_agent_start` event now receives `systemPrompt` in the event object and returns `systemPrompt` (full replacement) instead of `systemPromptAppend`. Extensions that were appending must now use `event.systemPrompt + extra` pattern. ([#575](https://github.com/badlogic/pi-mono/issues/575)) + ### Added - `--no-tools` flag to disable all built-in tools, allowing extension-only tool setups ([#557](https://github.com/badlogic/pi-mono/pull/557) by [@cv](https://github.com/cv)) diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index bb494ce8..7ec51839 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -255,7 +255,7 @@ pi starts ▼ user sends prompt ─────────────────────────────────────────┐ │ │ - ├─► before_agent_start (can inject message, append to system prompt) + ├─► before_agent_start (can inject message, modify system prompt) ├─► agent_start │ │ │ │ ┌─── turn (repeats while LLM calls tools) ───┐ │ @@ -414,12 +414,13 @@ pi.on("session_shutdown", async (_event, ctx) => { #### before_agent_start -Fired after user submits prompt, before agent loop. Can inject a message and/or append to the system prompt. +Fired after user submits prompt, before agent loop. Can inject a message and/or modify the system prompt. ```typescript pi.on("before_agent_start", async (event, ctx) => { // event.prompt - user's prompt text // event.images - attached images (if any) + // event.systemPrompt - current system prompt return { // Inject a persistent message (stored in session, sent to LLM) @@ -428,13 +429,13 @@ pi.on("before_agent_start", async (event, ctx) => { content: "Additional context for the LLM", display: true, }, - // Append to system prompt for this turn only - systemPromptAppend: "Extra instructions for this turn...", + // Replace the system prompt for this turn (chained across extensions) + systemPrompt: event.systemPrompt + "\n\nExtra instructions for this turn...", }; }); ``` -**Examples:** [claude-rules.ts](../examples/extensions/claude-rules.ts), [pirate.ts](../examples/extensions/pirate.ts), [plan-mode.ts](../examples/extensions/plan-mode.ts), [preset.ts](../examples/extensions/preset.ts) +**Examples:** [claude-rules.ts](../examples/extensions/claude-rules.ts), [pirate.ts](../examples/extensions/pirate.ts), [plan-mode.ts](../examples/extensions/plan-mode.ts), [preset.ts](../examples/extensions/preset.ts), [ssh.ts](../examples/extensions/ssh.ts) #### agent_start / agent_end diff --git a/packages/coding-agent/examples/extensions/claude-rules.ts b/packages/coding-agent/examples/extensions/claude-rules.ts index 747e9e7a..285bed42 100644 --- a/packages/coding-agent/examples/extensions/claude-rules.ts +++ b/packages/coding-agent/examples/extensions/claude-rules.ts @@ -61,7 +61,7 @@ export default function claudeRulesExtension(pi: ExtensionAPI) { }); // Append available rules to system prompt - pi.on("before_agent_start", async () => { + pi.on("before_agent_start", async (event) => { if (ruleFiles.length === 0) { return; } @@ -69,7 +69,10 @@ export default function claudeRulesExtension(pi: ExtensionAPI) { const rulesList = ruleFiles.map((f) => `- .claude/rules/${f}`).join("\n"); return { - systemPromptAppend: ` + systemPrompt: + event.systemPrompt + + ` + ## Project Rules The following project rules are available in .claude/rules/: diff --git a/packages/coding-agent/examples/extensions/pirate.ts b/packages/coding-agent/examples/extensions/pirate.ts index 559960f0..9231574b 100644 --- a/packages/coding-agent/examples/extensions/pirate.ts +++ b/packages/coding-agent/examples/extensions/pirate.ts @@ -1,8 +1,8 @@ /** * Pirate Extension * - * Demonstrates using systemPromptAppend in before_agent_start to dynamically - * modify the system prompt based on extension state. + * Demonstrates modifying the system prompt in before_agent_start to dynamically + * change agent behavior based on extension state. * * Usage: * 1. Copy this file to ~/.pi/agent/extensions/ or your project's .pi/extensions/ @@ -25,10 +25,13 @@ export default function pirateExtension(pi: ExtensionAPI) { }); // Append to system prompt when pirate mode is enabled - pi.on("before_agent_start", async () => { + pi.on("before_agent_start", async (event) => { if (pirateMode) { return { - systemPromptAppend: ` + systemPrompt: + event.systemPrompt + + ` + IMPORTANT: You are now in PIRATE MODE. You must: - Speak like a stereotypical pirate in all responses - Use phrases like "Arrr!", "Ahoy!", "Shiver me timbers!", "Avast!", "Ye scurvy dog!" diff --git a/packages/coding-agent/examples/extensions/preset.ts b/packages/coding-agent/examples/extensions/preset.ts index 32e02eb5..c6964cc6 100644 --- a/packages/coding-agent/examples/extensions/preset.ts +++ b/packages/coding-agent/examples/extensions/preset.ts @@ -345,10 +345,10 @@ export default function presetExtension(pi: ExtensionAPI) { }); // Inject preset instructions into system prompt - pi.on("before_agent_start", async () => { + pi.on("before_agent_start", async (event) => { if (activePreset?.instructions) { return { - systemPromptAppend: activePreset.instructions, + systemPrompt: `${event.systemPrompt}\n\n${activePreset.instructions}`, }; } }); diff --git a/packages/coding-agent/examples/extensions/ssh.ts b/packages/coding-agent/examples/extensions/ssh.ts index d269443e..76868434 100644 --- a/packages/coding-agent/examples/extensions/ssh.ts +++ b/packages/coding-agent/examples/extensions/ssh.ts @@ -191,4 +191,16 @@ export default function (pi: ExtensionAPI) { ctx.ui.notify(`SSH mode: ${ssh.remote}:${ssh.remoteCwd}`, "info"); } }); + + // Replace local cwd with remote cwd in system prompt + pi.on("before_agent_start", async (event) => { + const ssh = getSsh(); + if (ssh) { + const modified = event.systemPrompt.replace( + `Current working directory: ${localCwd}`, + `Current working directory: ${ssh.remoteCwd} (via SSH: ${ssh.remote})`, + ); + return { systemPrompt: modified }; + } + }); } diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index d02a5dd5..2d4444c7 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -613,7 +613,11 @@ export class AgentSession { // Emit before_agent_start extension event if (this._extensionRunner) { - const result = await this._extensionRunner.emitBeforeAgentStart(expandedText, options?.images); + const result = await this._extensionRunner.emitBeforeAgentStart( + expandedText, + options?.images, + this._baseSystemPrompt, + ); // Add all custom messages from extensions if (result?.messages) { for (const msg of result.messages) { @@ -627,11 +631,11 @@ export class AgentSession { }); } } - // Apply extension systemPromptAppend on top of base prompt - if (result?.systemPromptAppend) { - this.agent.setSystemPrompt(`${this._baseSystemPrompt}\n\n${result.systemPromptAppend}`); + // Apply extension-modified system prompt, or reset to base + if (result?.systemPrompt) { + this.agent.setSystemPrompt(result.systemPrompt); } else { - // Ensure we're using the base prompt (in case previous turn had appends) + // Ensure we're using the base prompt (in case previous turn had modifications) this.agent.setSystemPrompt(this._baseSystemPrompt); } } diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 9f7931af..733f7642 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -38,7 +38,7 @@ import type { /** Combined result from all before_agent_start handlers */ interface BeforeAgentStartCombinedResult { messages?: NonNullable[]; - systemPromptAppend?: string; + systemPrompt?: string; } export type ExtensionErrorListener = (error: ExtensionError) => void; @@ -433,11 +433,13 @@ export class ExtensionRunner { async emitBeforeAgentStart( prompt: string, - images?: ImageContent[], + images: ImageContent[] | undefined, + systemPrompt: string, ): Promise { const ctx = this.createContext(); const messages: NonNullable[] = []; - const systemPromptAppends: string[] = []; + let currentSystemPrompt = systemPrompt; + let systemPromptModified = false; for (const ext of this.extensions) { const handlers = ext.handlers.get("before_agent_start"); @@ -445,7 +447,12 @@ export class ExtensionRunner { for (const handler of handlers) { try { - const event: BeforeAgentStartEvent = { type: "before_agent_start", prompt, images }; + const event: BeforeAgentStartEvent = { + type: "before_agent_start", + prompt, + images, + systemPrompt: currentSystemPrompt, + }; const handlerResult = await handler(event, ctx); if (handlerResult) { @@ -453,8 +460,9 @@ export class ExtensionRunner { if (result.message) { messages.push(result.message); } - if (result.systemPromptAppend) { - systemPromptAppends.push(result.systemPromptAppend); + if (result.systemPrompt !== undefined) { + currentSystemPrompt = result.systemPrompt; + systemPromptModified = true; } } } catch (err) { @@ -470,10 +478,10 @@ export class ExtensionRunner { } } - if (messages.length > 0 || systemPromptAppends.length > 0) { + if (messages.length > 0 || systemPromptModified) { return { messages: messages.length > 0 ? messages : undefined, - systemPromptAppend: systemPromptAppends.length > 0 ? systemPromptAppends.join("\n\n") : undefined, + systemPrompt: systemPromptModified ? currentSystemPrompt : undefined, }; } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 168f6899..d2469e5c 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -349,6 +349,7 @@ export interface BeforeAgentStartEvent { type: "before_agent_start"; prompt: string; images?: ImageContent[]; + systemPrompt: string; } /** Fired when an agent loop starts */ @@ -504,7 +505,8 @@ export interface ToolResultEventResult { export interface BeforeAgentStartEventResult { message?: Pick; - systemPromptAppend?: string; + /** Replace the system prompt for this turn. If multiple extensions return this, they are chained. */ + systemPrompt?: string; } export interface SessionBeforeSwitchResult { diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index a8b1c70d..7bf128bb 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -264,6 +264,7 @@ export { getMarkdownTheme, getSelectListTheme, getSettingsListTheme, + initTheme, Theme, type ThemeColor, } from "./modes/interactive/theme/theme.js";