From 3df56506c7da704e63b24129f4d48250c2eb49b3 Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 5 Aug 2026 21:28:37 +0000 Subject: [PATCH] fix(coding-agent): preserve before_agent_start prompt override across base rebuilds The per-turn systemPrompt returned by before_agent_start was applied only to the agent state, so any base-prompt rebuild firing in the prompt window (context-overflow compaction/promotion, memory promotion, MCP/RPC tool refresh, hindsight MM-TTL refresh) re-pushed the rebuilt base via setSystemPrompt and silently dropped the override before the request. SessionTools now tracks the active per-turn override and re-applies it on every base rebuild during the turn, clearing it when the turn ends. Fixes #7755 --- packages/coding-agent/CHANGELOG.md | 4 + .../coding-agent/src/session/agent-session.ts | 6 +- .../coding-agent/src/session/session-tools.ts | 40 +++++- ...before-agent-start-prompt-override.test.ts | 127 ++++++++++++++++++ 4 files changed, 173 insertions(+), 4 deletions(-) create mode 100644 packages/coding-agent/test/agent-session-before-agent-start-prompt-override.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 456e88609..cdc9393da 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed a per-turn `before_agent_start` system prompt override being silently dropped when a base-prompt rebuild fired between the hook and the provider request. The override lived only on the agent state, so `refreshBaseSystemPrompt`/`applyActiveToolsByName` re-pushing the rebuilt base (context-overflow compaction/promotion, memory promotion, MCP/RPC tool refresh, or the fire-and-forget hindsight MM-TTL refresh) clobbered it. The tools controller now tracks the active override and re-applies it on every base rebuild during the turn, clearing it when the turn ends ([#7755](https://github.com/can1357/oh-my-pi/issues/7755)). + ## [17.2.9] - 2026-08-05 ### Breaking Changes diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 9a34d5da4..a397923b5 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -5267,11 +5267,13 @@ export class AgentSession { if (result?.systemPrompt !== undefined) { baseXdevCatalogDelivered = false; - this.agent.setSystemPrompt(result.systemPrompt); + this.#tools.setTurnSystemPromptOverride(result.systemPrompt); } else { + this.#tools.clearTurnSystemPromptOverride(); this.agent.setSystemPrompt(beforeAgentStartSystemPrompt); } } else { + this.#tools.clearTurnSystemPromptOverride(); this.agent.setSystemPrompt(beforeAgentStartSystemPrompt); } @@ -5336,6 +5338,8 @@ export class AgentSession { await this.#waitForPostPromptRecovery(generation); } } finally { + // The per-turn before_agent_start override lives only for this turn. + this.#tools.clearTurnSystemPromptOverride(); this.#usagePreflightReadyForNextModelCall = false; this.#endInFlight(); } diff --git a/packages/coding-agent/src/session/session-tools.ts b/packages/coding-agent/src/session/session-tools.ts index e78ee3650..a1d06427e 100644 --- a/packages/coding-agent/src/session/session-tools.ts +++ b/packages/coding-agent/src/session/session-tools.ts @@ -198,6 +198,15 @@ export class SessionTools { #presentationPinnedToolNames: ReadonlySet | undefined; #runtimeSelectedToolNames: ReadonlySet | undefined; #baseSystemPrompt: string[]; + /** + * Per-turn system prompt returned by a `before_agent_start` extension hook + * ("replace the system prompt for this turn"). While set, base-prompt + * rebuilds keep this override on the agent instead of the rebuilt base, so a + * rebuild landing in the prompt window (compaction/promotion, memory + * promotion, MCP/RPC tool refresh, hindsight MM-TTL refresh) cannot silently + * drop it before the request. Cleared when the turn ends. + */ + #turnSystemPromptOverride: string[] | undefined; #lastAppliedToolSignature: string | undefined; /** * `xd://` device names the current base system prompt renders in its catalog @@ -262,6 +271,31 @@ export class SessionTools { this.#baseSystemPrompt = prompt; } + /** + * Pushes `base` to the agent as the effective system prompt, unless an active + * per-turn {@link #turnSystemPromptOverride} takes precedence. Every base + * rebuild applies its result through here so a mid-turn rebuild preserves the + * override. + */ + #applyAgentSystemPrompt(base: string[]): void { + this.#host.agent.setSystemPrompt(this.#turnSystemPromptOverride ?? base); + } + + /** + * Registers the per-turn `before_agent_start` system-prompt override and + * applies it to the agent. Base rebuilds during the turn preserve it until + * {@link clearTurnSystemPromptOverride}. + */ + setTurnSystemPromptOverride(prompt: string[]): void { + this.#turnSystemPromptOverride = prompt; + this.#host.agent.setSystemPrompt(prompt); + } + + /** Drops the active per-turn override; later rebuilds fall back to the base prompt. */ + clearTurnSystemPromptOverride(): void { + this.#turnSystemPromptOverride = undefined; + } + /** Skills currently rendered into the system prompt. */ get skills(): Skill[] { return this.#skills; @@ -663,7 +697,7 @@ export class SessionTools { if (this.#lastAppliedToolSignature !== undefined) this.#host.clearInheritedProviderPromptCacheKey(); this.#baseSystemPrompt = rebuiltSystemPrompt; this.#host.clearMemoryPromotionSnapshot(); - this.#host.agent.setSystemPrompt(this.#baseSystemPrompt); + this.#applyAgentSystemPrompt(this.#baseSystemPrompt); this.#lastAppliedToolSignature = rebuiltSignature; this.#promptModelKey = this.#currentPromptModelKey(); this.#basePromptXdevNames = new Set(rebuiltXdevCatalogNames); @@ -1103,7 +1137,7 @@ export class SessionTools { ) { this.#host.clearInheritedProviderPromptCacheKey(); } - this.#host.agent.setSystemPrompt(this.#baseSystemPrompt); + this.#applyAgentSystemPrompt(this.#baseSystemPrompt); this.#promptModelKey = this.#currentPromptModelKey(); // Refresh the cached signature so a subsequent `applyActiveToolsByName` with // the same tool set does not re-rebuild on top of the explicit refresh we @@ -1143,7 +1177,7 @@ export class SessionTools { this.#host.captureMemoryPromotionSnapshot(previousBaseSystemPrompt); const stablePrompt = [...previousBaseSystemPrompt, injected]; this.#baseSystemPrompt = stablePrompt; - this.#host.agent.setSystemPrompt(stablePrompt); + this.#applyAgentSystemPrompt(stablePrompt); return stablePrompt; } catch (err) { logger.debug("Memory backend beforeAgentStartPrompt failed", { diff --git a/packages/coding-agent/test/agent-session-before-agent-start-prompt-override.test.ts b/packages/coding-agent/test/agent-session-before-agent-start-prompt-override.test.ts new file mode 100644 index 000000000..27bf3f7a9 --- /dev/null +++ b/packages/coding-agent/test/agent-session-before-agent-start-prompt-override.test.ts @@ -0,0 +1,127 @@ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import { Agent } from "@oh-my-pi/pi-agent-core"; +import type { Model } from "@oh-my-pi/pi-ai"; +import { createMockModel, type MockResponseSource } from "@oh-my-pi/pi-ai/providers/mock"; +import { buildModel } from "@oh-my-pi/pi-catalog/build"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import type { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; + +// Contract: a per-turn system prompt returned by `before_agent_start` +// ("Replace the system prompt for this turn") must reach the provider for the +// turn. A base-prompt rebuild that fires in the prompt window — context-overflow +// compaction/promotion, memory promotion, MCP/RPC tool refresh, or the +// fire-and-forget hindsight MM-TTL refresh — re-sets the agent prompt to the +// rebuilt base. It must not clobber an active override. Regression for #7755. + +const OVERRIDE = "OVERRIDE-SYSTEM-PROMPT-LIFEOS_ROUTE"; +const REBUILT_BASE = "REBUILT-BASE-WITH-TOOL-CATALOG"; + +function createModel(): Model<"openai-responses"> { + return buildModel({ + id: "mock", + name: "mock", + api: "openai-responses", + provider: "openai", + baseUrl: "https://example.invalid", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 8192, + maxTokens: 2048, + }); +} + +describe("AgentSession before_agent_start system prompt override", () => { + let session: AgentSession | undefined; + + afterEach(async () => { + if (session) { + await session.dispose(); + session = undefined; + } + vi.restoreAllMocks(); + }); + + /** + * Builds a session whose `before_agent_start` replaces the prompt with + * {@link OVERRIDE} and whose base rebuild renders {@link REBUILT_BASE}. + * + * When `rebuildInWindow` is set, a base rebuild is fired from a + * `beforeModelCall` hook — which the agent loop runs immediately before it + * re-reads `state.systemPrompt` for the request — reproducing a rebuild that + * lands in the window between the hook and the provider request. + */ + function createSession( + responses: MockResponseSource, + options: { rebuildInWindow?: boolean } = {}, + ): { session: AgentSession; systemPrompts: string[][] } { + const mock = createMockModel({ responses }); + const systemPrompts: string[][] = []; + + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { + model: createModel(), + systemPrompt: ["initial-base"], + tools: [], + messages: [], + }, + convertToLlm, + streamFn: (model, context, streamOptions) => { + systemPrompts.push([...(context.systemPrompt ?? [])]); + return mock.stream(model, context, streamOptions); + }, + }); + + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings: Settings.isolated({ "compaction.enabled": false, "todo.enabled": false }), + modelRegistry: { getApiKey: async () => "test-key" } as never, + extensionRunner: { + emitBeforeAgentStart: async () => ({ systemPrompt: [OVERRIDE] }), + emit: async () => undefined, + } as unknown as ExtensionRunner, + rebuildSystemPrompt: async () => ({ systemPrompt: [REBUILT_BASE] }), + }); + const activeSession = session; + + if (options.rebuildInWindow) { + let fired = false; + agent.addBeforeModelCallHook(async () => { + if (fired) return; + fired = true; + await activeSession.refreshBaseSystemPrompt(); + }); + } + + return { session, systemPrompts }; + } + + it("keeps the override when a base rebuild fires in the prompt window", async () => { + const { session, systemPrompts } = createSession([{ content: ["Done"] }], { rebuildInWindow: true }); + + await session.prompt("hello"); + await session.waitForIdle(); + + // The rebuild ran right before the request re-read the agent prompt; the + // override must still reach the provider instead of the rebuilt base. + expect(systemPrompts).toHaveLength(1); + expect(systemPrompts[0]).toEqual([OVERRIDE]); + }); + + it("falls back to the rebuilt base once the turn ends", async () => { + const { session } = createSession([{ content: ["Done"] }]); + + await session.prompt("hello"); + await session.waitForIdle(); + + // The per-turn override is cleared when the turn completes, so a later + // rebuild applies the base prompt rather than leaking the stale override. + await session.refreshBaseSystemPrompt(); + expect(session.systemPrompt).toEqual([REBUILT_BASE]); + }); +});