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
This commit is contained in:
Mario Zechner 2026-01-08 19:54:34 +01:00
parent 0774db2e5a
commit 17cb328ca1
10 changed files with 65 additions and 27 deletions

View file

@ -2,6 +2,10 @@
## [Unreleased] ## [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 ### 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)) - `--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))

View file

@ -255,7 +255,7 @@ pi starts
user sends prompt ─────────────────────────────────────────┐ user sends prompt ─────────────────────────────────────────┐
│ │ │ │
├─► before_agent_start (can inject message, append to system prompt) ├─► before_agent_start (can inject message, modify system prompt)
├─► agent_start │ ├─► agent_start │
│ │ │ │
│ ┌─── turn (repeats while LLM calls tools) ───┐ │ │ ┌─── turn (repeats while LLM calls tools) ───┐ │
@ -414,12 +414,13 @@ pi.on("session_shutdown", async (_event, ctx) => {
#### before_agent_start #### 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 ```typescript
pi.on("before_agent_start", async (event, ctx) => { pi.on("before_agent_start", async (event, ctx) => {
// event.prompt - user's prompt text // event.prompt - user's prompt text
// event.images - attached images (if any) // event.images - attached images (if any)
// event.systemPrompt - current system prompt
return { return {
// Inject a persistent message (stored in session, sent to LLM) // 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", content: "Additional context for the LLM",
display: true, display: true,
}, },
// Append to system prompt for this turn only // Replace the system prompt for this turn (chained across extensions)
systemPromptAppend: "Extra instructions for this turn...", 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 #### agent_start / agent_end

View file

@ -61,7 +61,7 @@ export default function claudeRulesExtension(pi: ExtensionAPI) {
}); });
// Append available rules to system prompt // Append available rules to system prompt
pi.on("before_agent_start", async () => { pi.on("before_agent_start", async (event) => {
if (ruleFiles.length === 0) { if (ruleFiles.length === 0) {
return; return;
} }
@ -69,7 +69,10 @@ export default function claudeRulesExtension(pi: ExtensionAPI) {
const rulesList = ruleFiles.map((f) => `- .claude/rules/${f}`).join("\n"); const rulesList = ruleFiles.map((f) => `- .claude/rules/${f}`).join("\n");
return { return {
systemPromptAppend: ` systemPrompt:
event.systemPrompt +
`
## Project Rules ## Project Rules
The following project rules are available in .claude/rules/: The following project rules are available in .claude/rules/:

View file

@ -1,8 +1,8 @@
/** /**
* Pirate Extension * Pirate Extension
* *
* Demonstrates using systemPromptAppend in before_agent_start to dynamically * Demonstrates modifying the system prompt in before_agent_start to dynamically
* modify the system prompt based on extension state. * change agent behavior based on extension state.
* *
* Usage: * Usage:
* 1. Copy this file to ~/.pi/agent/extensions/ or your project's .pi/extensions/ * 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 // Append to system prompt when pirate mode is enabled
pi.on("before_agent_start", async () => { pi.on("before_agent_start", async (event) => {
if (pirateMode) { if (pirateMode) {
return { return {
systemPromptAppend: ` systemPrompt:
event.systemPrompt +
`
IMPORTANT: You are now in PIRATE MODE. You must: IMPORTANT: You are now in PIRATE MODE. You must:
- Speak like a stereotypical pirate in all responses - Speak like a stereotypical pirate in all responses
- Use phrases like "Arrr!", "Ahoy!", "Shiver me timbers!", "Avast!", "Ye scurvy dog!" - Use phrases like "Arrr!", "Ahoy!", "Shiver me timbers!", "Avast!", "Ye scurvy dog!"

View file

@ -345,10 +345,10 @@ export default function presetExtension(pi: ExtensionAPI) {
}); });
// Inject preset instructions into system prompt // Inject preset instructions into system prompt
pi.on("before_agent_start", async () => { pi.on("before_agent_start", async (event) => {
if (activePreset?.instructions) { if (activePreset?.instructions) {
return { return {
systemPromptAppend: activePreset.instructions, systemPrompt: `${event.systemPrompt}\n\n${activePreset.instructions}`,
}; };
} }
}); });

View file

@ -191,4 +191,16 @@ export default function (pi: ExtensionAPI) {
ctx.ui.notify(`SSH mode: ${ssh.remote}:${ssh.remoteCwd}`, "info"); 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 };
}
});
} }

View file

@ -613,7 +613,11 @@ export class AgentSession {
// Emit before_agent_start extension event // Emit before_agent_start extension event
if (this._extensionRunner) { 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 // Add all custom messages from extensions
if (result?.messages) { if (result?.messages) {
for (const msg of result.messages) { for (const msg of result.messages) {
@ -627,11 +631,11 @@ export class AgentSession {
}); });
} }
} }
// Apply extension systemPromptAppend on top of base prompt // Apply extension-modified system prompt, or reset to base
if (result?.systemPromptAppend) { if (result?.systemPrompt) {
this.agent.setSystemPrompt(`${this._baseSystemPrompt}\n\n${result.systemPromptAppend}`); this.agent.setSystemPrompt(result.systemPrompt);
} else { } 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); this.agent.setSystemPrompt(this._baseSystemPrompt);
} }
} }

View file

@ -38,7 +38,7 @@ import type {
/** Combined result from all before_agent_start handlers */ /** Combined result from all before_agent_start handlers */
interface BeforeAgentStartCombinedResult { interface BeforeAgentStartCombinedResult {
messages?: NonNullable<BeforeAgentStartEventResult["message"]>[]; messages?: NonNullable<BeforeAgentStartEventResult["message"]>[];
systemPromptAppend?: string; systemPrompt?: string;
} }
export type ExtensionErrorListener = (error: ExtensionError) => void; export type ExtensionErrorListener = (error: ExtensionError) => void;
@ -433,11 +433,13 @@ export class ExtensionRunner {
async emitBeforeAgentStart( async emitBeforeAgentStart(
prompt: string, prompt: string,
images?: ImageContent[], images: ImageContent[] | undefined,
systemPrompt: string,
): Promise<BeforeAgentStartCombinedResult | undefined> { ): Promise<BeforeAgentStartCombinedResult | undefined> {
const ctx = this.createContext(); const ctx = this.createContext();
const messages: NonNullable<BeforeAgentStartEventResult["message"]>[] = []; const messages: NonNullable<BeforeAgentStartEventResult["message"]>[] = [];
const systemPromptAppends: string[] = []; let currentSystemPrompt = systemPrompt;
let systemPromptModified = false;
for (const ext of this.extensions) { for (const ext of this.extensions) {
const handlers = ext.handlers.get("before_agent_start"); const handlers = ext.handlers.get("before_agent_start");
@ -445,7 +447,12 @@ export class ExtensionRunner {
for (const handler of handlers) { for (const handler of handlers) {
try { 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); const handlerResult = await handler(event, ctx);
if (handlerResult) { if (handlerResult) {
@ -453,8 +460,9 @@ export class ExtensionRunner {
if (result.message) { if (result.message) {
messages.push(result.message); messages.push(result.message);
} }
if (result.systemPromptAppend) { if (result.systemPrompt !== undefined) {
systemPromptAppends.push(result.systemPromptAppend); currentSystemPrompt = result.systemPrompt;
systemPromptModified = true;
} }
} }
} catch (err) { } catch (err) {
@ -470,10 +478,10 @@ export class ExtensionRunner {
} }
} }
if (messages.length > 0 || systemPromptAppends.length > 0) { if (messages.length > 0 || systemPromptModified) {
return { return {
messages: messages.length > 0 ? messages : undefined, messages: messages.length > 0 ? messages : undefined,
systemPromptAppend: systemPromptAppends.length > 0 ? systemPromptAppends.join("\n\n") : undefined, systemPrompt: systemPromptModified ? currentSystemPrompt : undefined,
}; };
} }

View file

@ -349,6 +349,7 @@ export interface BeforeAgentStartEvent {
type: "before_agent_start"; type: "before_agent_start";
prompt: string; prompt: string;
images?: ImageContent[]; images?: ImageContent[];
systemPrompt: string;
} }
/** Fired when an agent loop starts */ /** Fired when an agent loop starts */
@ -504,7 +505,8 @@ export interface ToolResultEventResult {
export interface BeforeAgentStartEventResult { export interface BeforeAgentStartEventResult {
message?: Pick<CustomMessage, "customType" | "content" | "display" | "details">; message?: Pick<CustomMessage, "customType" | "content" | "display" | "details">;
systemPromptAppend?: string; /** Replace the system prompt for this turn. If multiple extensions return this, they are chained. */
systemPrompt?: string;
} }
export interface SessionBeforeSwitchResult { export interface SessionBeforeSwitchResult {

View file

@ -264,6 +264,7 @@ export {
getMarkdownTheme, getMarkdownTheme,
getSelectListTheme, getSelectListTheme,
getSettingsListTheme, getSettingsListTheme,
initTheme,
Theme, Theme,
type ThemeColor, type ThemeColor,
} from "./modes/interactive/theme/theme.js"; } from "./modes/interactive/theme/theme.js";