From 108eff32ea843baf4337e1de9d9ccb3cec6d97fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20M=C3=BCller?= Date: Fri, 21 Aug 2026 19:59:31 +0200 Subject: [PATCH] feat(coding-agent): session_before_refine extension hook (#1558) * feat(coding-agent): session_before_refine extension hook Let extensions customize continual-harness refinement the same way session_before_compact customizes compaction. The hook fires before the planning LLM call for /refine and auto-refine with the planning inputs (trigger, instructions, scope, planning harness state, refinement history, serialized conversation). An extension can return a RefinementProposal to replace the built-in planner (edits still pass apply-time validation and baseline conflict rejection), return { skip: true } to suppress the round, or return nothing to fall back to the default planner. Rollback refinements bypass the hook. Includes examples/extensions/custom-refinement.ts (planning with a cheaper model, covering discussion #1464) and documents the existing refine_complete event. * Update queue-test refine call assertions for the trigger argument * Propagate the auto trigger and skip semantics through serialized refine paths Serialized auto-refine now reaches session_before_refine with trigger "auto", an extension skip there stamps the cooldown without emitting refine_failed, and a skipped explicit refine.run surfaces the RefineSkippedError instead of passing as a silent reviewer decline. The example planner now labels entry ids with their scope and tells the model other-scope entries are read-only. * fix(coding-agent): consume skipped auto-refine rounds * fix(coding-agent): report skipped refine during disposal * fix(coding-agent): normalize extension refine proposals --------- Co-authored-by: Seth --- .../feat-session-before-refine-hook.md | 1 + packages/coding-agent/docs/extensions.md | 35 +++ .../examples/extensions/custom-refinement.ts | 109 +++++++ .../coding-agent/src/core/agent-session.ts | 96 +++++-- .../coding-agent/src/core/extensions/index.ts | 3 + .../src/core/extensions/runner.ts | 22 +- .../coding-agent/src/core/extensions/types.ts | 36 +++ .../src/core/refinement/refinement.ts | 34 ++- packages/coding-agent/src/index.ts | 9 + .../test/suite/agent-session-queue.test.ts | 2 + .../agent-session-refine-extension.test.ts | 270 ++++++++++++++++++ 11 files changed, 582 insertions(+), 35 deletions(-) create mode 100644 packages/coding-agent/.changes/feat-session-before-refine-hook.md create mode 100644 packages/coding-agent/examples/extensions/custom-refinement.ts create mode 100644 packages/coding-agent/test/suite/agent-session-refine-extension.test.ts diff --git a/packages/coding-agent/.changes/feat-session-before-refine-hook.md b/packages/coding-agent/.changes/feat-session-before-refine-hook.md new file mode 100644 index 000000000..31dd073e9 --- /dev/null +++ b/packages/coding-agent/.changes/feat-session-before-refine-hook.md @@ -0,0 +1 @@ +- Added a `session_before_refine` extension hook: extensions can replace `/refine` and auto-refine planning with their own proposal (for example using a cheaper model — see `examples/extensions/custom-refinement.ts`) or skip a refinement round; rollbacks bypass the hook and extension edits go through the normal apply-time validation. Also documents `refine_complete`. diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index a91f063a2..02804a445 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -319,6 +319,10 @@ user sends another prompt ◄───────────────── ├─► session_before_compact (can cancel or customize) └─► session_compact +/refine or auto-refine + ├─► session_before_refine (can skip or replace planning; rollbacks bypass it) + └─► refine_complete + /tree navigation ├─► session_before_tree (can cancel or customize) └─► session_tree @@ -432,6 +436,36 @@ pi.on("session_compact", async (event, ctx) => { }); ``` +#### session_before_refine / refine_complete + +Fired around continual-harness refinement (`/refine` and auto-refine). `session_before_refine` runs before the planning LLM call and can skip the round or replace the planner; rollback refinements bypass it. Proposal edits from extensions go through the same apply-time validation as built-in plans. + +```typescript +pi.on("session_before_refine", async (event, ctx) => { + const { trigger, instructions, scope, planningState, history, conversationText } = event.preparation; + // trigger - "manual" (/refine) or "auto" (auto-refine) + // scope - "global" or "local" harness target + + // Skip this refinement round: + return { skip: true }; + + // OR replace the built-in planner (e.g. with a cheaper model — see + // examples/extensions/custom-refinement.ts): + return { + proposal: { + summary: "...", + rationale: "...", + expectedOutcome: "...", + edits: [{ action: "create", kind: "memory", title: "...", content: "..." }], + }, + }; +}); + +pi.on("refine_complete", async (event, ctx) => { + // event.id, summary, appliedEdits, scope +}); +``` + #### session_before_tree / session_tree Fired on `/tree` navigation. See [Sessions](sessions.md) for tree navigation concepts. @@ -2544,6 +2578,7 @@ All examples in [examples/extensions/](../examples/extensions/). | `file-trigger.ts` | File watcher triggers messages | `sendMessage` | | **Compaction & Sessions** ||| | `custom-compaction.ts` | Custom compaction summary | `on("session_before_compact")` | +| `custom-refinement.ts` | Custom /refine planning (cheaper model) | `on("session_before_refine")` | | `trigger-compact.ts` | Trigger compaction manually | `compact()` | | `git-checkpoint.ts` | Git stash on turns | `on("turn_start")`, `on("session_before_fork")`, `exec` | | `auto-commit-on-exit.ts` | Commit on shutdown | `on("session_shutdown")`, `exec` | diff --git a/packages/coding-agent/examples/extensions/custom-refinement.ts b/packages/coding-agent/examples/extensions/custom-refinement.ts new file mode 100644 index 000000000..5e83287e4 --- /dev/null +++ b/packages/coding-agent/examples/extensions/custom-refinement.ts @@ -0,0 +1,109 @@ +/** + * Custom Refinement Extension + * + * Replaces the built-in /refine planner with a cheaper/faster model. + * The extension receives the planning inputs (harness state, refinement + * history, serialized conversation) and returns a RefinementProposal. + * Returned edits are still validated by the core apply path, so a bad + * proposal degrades to per-edit errors instead of corrupting the harness. + * + * Returning nothing falls back to the built-in planner; returning + * `{ skip: true }` suppresses the refinement round entirely. + * + * Usage: + * prime-agent -e examples/extensions/custom-refinement.ts + */ + +import { complete } from "@earendil-works/pi-ai"; +import type { ExtensionAPI, RefinementProposal } from "@earendil-works/pi-coding-agent"; + +export default function (pi: ExtensionAPI) { + pi.on("session_before_refine", async (event, ctx) => { + const { trigger, instructions, scope, planningState, history, conversationText } = event.preparation; + + // Use a cheaper model than the conversation model for planning. + const model = ctx.modelRegistry.find("google", "gemini-2.5-flash"); + if (!model) { + ctx.ui.notify("Custom refinement: model not found, using default planner", "warning"); + return; + } + const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model); + if (!auth.ok || !auth.apiKey) { + ctx.ui.notify("Custom refinement: no auth for model, using default planner", "warning"); + return; + } + + ctx.ui.notify(`Custom refinement (${trigger}, ${scope}) with ${model.id}...`, "info"); + + // Global entries are read-only during a local refinement; label scopes so + // the model never proposes edits against another scope's ids. + const existingIds = Object.entries(planningState.entries) + .flatMap(([kind, entries]) => + Object.entries(entries).map(([id, entry]) => `${entry.scope ?? scope}:${kind}:${id}`), + ) + .join(", "); + const recentRefinements = history + .slice(-5) + .map((item) => `- ${item.summary}`) + .join("\n"); + + const planningMessages = [ + { + role: "user" as const, + content: [ + { + type: "text" as const, + text: `You maintain a persistent "continual harness" of prompt notes, memories, skills, and subagent specs for a coding agent. +From the conversation, propose a SMALL set of create/update/delete edits that capture durable, reusable lessons. + +Existing entries as scope:kind:id (use the bare id in edits): ${existingIds || "(none)"} +This refinement targets the ${scope} scope. Entries from any other scope are read-only: never propose update or delete edits for them. +Recent refinements: +${recentRefinements || "(none)"} +${instructions ? `Focus: ${instructions}` : ""} + +Respond with ONLY a JSON object: +{"summary": "...", "rationale": "...", "expectedOutcome": "...", "edits": [{"action": "create"|"update"|"delete", "kind": "prompt"|"memory"|"skill"|"subagent", "id": "existing-id-for-update-or-delete", "title": "...", "content": "...", "reason": "..."}]} + + +${conversationText} +`, + }, + ], + timestamp: Date.now(), + }, + ]; + + try { + const response = await complete( + model, + { messages: planningMessages }, + { + apiKey: auth.apiKey, + headers: auth.headers, + maxTokens: 4096, + signal: event.signal, + }, + ); + + const text = response.content + .filter((block): block is { type: "text"; text: string } => block.type === "text") + .map((block) => block.text) + .join(""); + const start = text.indexOf("{"); + const end = text.lastIndexOf("}"); + if (start < 0 || end <= start) { + ctx.ui.notify("Custom refinement: unparseable plan, using default planner", "warning"); + return; + } + const proposal = JSON.parse(text.slice(start, end + 1)) as RefinementProposal; + return { proposal }; + } catch (error) { + if (!event.signal.aborted) { + const message = error instanceof Error ? error.message : String(error); + ctx.ui.notify(`Custom refinement failed (${message}), using default planner`, "warning"); + } + return; + } + }); +} diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 187253a4c..4bc8eff88 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -99,6 +99,7 @@ import { estimateContextTokens, generateBranchSummary, prepareCompaction, + serializeConversation, shouldCompact, } from "./compaction/index.js"; import { @@ -125,6 +126,7 @@ import { type MessageUpdateEvent, type ReplacedSessionContext, type SessionBeforeCompactResult, + type SessionBeforeRefineResult, type SessionBeforeTreeResult, type SessionStartEvent, type ShutdownHandler, @@ -165,6 +167,7 @@ import { type CompactionOutcome, type CompactionOutcomeReason, type CustomMessage, + convertToLlm, createCompactionOutcomeMessage, createHeartbeatPromptMessage, createRlmChildFailureMessage, @@ -186,6 +189,7 @@ import { type AutoRefineReview, appendGlobalRefinement, applyRefinementProposal, + generateRefinementId, getGlobalHarnessStateDir, getLocalHarnessStateDir, getRefinementHistory, @@ -195,6 +199,7 @@ import { loadHarnessState, mergeHarnessStates, mergeRefinementHistory, + normalizeRefinementProposal, planRefinement, REFINE_SKILL_NAME, type RefinementPlan, @@ -380,6 +385,9 @@ type UserBashEndDetails = { export class CompactionSkippedError extends Error {} +/** Thrown when a session_before_refine extension skips the refinement round. */ +export class RefineSkippedError extends Error {} + export interface AgentSessionConfig { agent: Agent; sessionManager: SessionManager; @@ -476,7 +484,7 @@ export type SerializedBackgroundPlanResult = abort: AbortController; branchVersion: number; } - | { status: "skip" } + | { status: "skip"; explicit?: boolean } | { status: "invalidated"; branchVersion: number } | { status: "failure"; @@ -2279,8 +2287,11 @@ export class AgentSession { } if (bgResult?.status === "skip") { - // Reviewer declined during background planning. + // Reviewer declined or an extension skipped during background planning. // Reset exactly once. Never retry the interval review; only fall through for a separate pending refine.run. + if (bgResult.explicit) { + this._emitRefineFailed(new RefineSkippedError("Refinement skipped by extension")); + } this._lastAutoRefineReviewAt = Date.now(); this._assistantTurnsSinceAutoRefine = 0; if (!this._pendingRequestedRefine) { @@ -2409,9 +2420,7 @@ export class AgentSession { this._assistantTurnsSinceAutoRefine = 0; return; } - await this._runSerializedRefine({ - instructions: autoRefineInstructions(reason, review), - }); + await this._runSerializedRefine({ instructions: autoRefineInstructions(reason, review) }, "auto"); if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { return; } @@ -2420,7 +2429,12 @@ export class AgentSession { } catch (error) { if (branchVersion === this._autoRefineBranchVersion) { this._lastAutoRefineReviewAt = Date.now(); - this._emitRefineFailed(error); + // An extension skip is an intentional non-round, not a failure. + if (error instanceof RefineSkippedError) { + this._assistantTurnsSinceAutoRefine = 0; + } else { + this._emitRefineFailed(error); + } } } finally { if (this._autoRefineReviewAbort === reviewAbort) { @@ -2582,7 +2596,7 @@ export class AgentSession { } // For explicit refine.run (skipReview=true), plan directly with // the user-provided options — no auto-review gate. - const plan = await this._planRefine(planOptions, refineAbort.signal); + const plan = await this._planRefine(planOptions, refineAbort.signal, skipReview ? "manual" : "auto"); if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { return { status: "invalidated", branchVersion }; } @@ -2593,10 +2607,13 @@ export class AgentSession { abort: refineAbort, branchVersion, }; - } catch { + } catch (error) { if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { return { status: "invalidated", branchVersion }; } + if (error instanceof RefineSkippedError) { + return { status: "skip", explicit: skipReview }; + } return { status: "failure", explicit: skipReview, @@ -2617,11 +2634,14 @@ export class AgentSession { * so the agent is between turns and _applyRefine's disconnect/reconnect * is safe. */ - private async _runSerializedRefine(options: { - instructions?: string; - rollbackId?: string; - global?: boolean; - }): Promise { + private async _runSerializedRefine( + options: { + instructions?: string; + rollbackId?: string; + global?: boolean; + }, + trigger: "manual" | "auto" = "manual", + ): Promise { if (this._disposed || this._disposing) { return; } @@ -2645,7 +2665,7 @@ export class AgentSession { const refineAbort = new AbortController(); this._refineAbortController = refineAbort; - const planRun = this._planRefine(options, refineAbort.signal); + const planRun = this._planRefine(options, refineAbort.signal, trigger); const planSettled = planRun.then( () => undefined, () => undefined, @@ -3904,6 +3924,9 @@ export class AgentSession { ) { this._pendingRequestedRefine = bgResult.options; } + if (bgResult?.status === "skip" && bgResult.explicit) { + this._emitRefineFailed(new RefineSkippedError("Refinement skipped by extension")); + } // For "skip" or "failure", stamp cooldown and reset counter // so the interval check below does not trigger a duplicate // terminal retry. @@ -7790,9 +7813,7 @@ export class AgentSession { private async _runApprovedRefine(reason: AutoRefineReason, review: AutoRefineReview): Promise { this._autoRefineInProgress = true; try { - await this.refine({ - instructions: autoRefineInstructions(reason, review), - }); + await this.refine({ instructions: autoRefineInstructions(reason, review) }, { trigger: "auto" }); this._pendingAutoRefineReview = undefined; this._turnIntervalAutoRefinePending = false; this._lastAutoRefineReviewAt = Date.now(); @@ -7800,11 +7821,18 @@ export class AgentSession { if (reason === "compact") { this._compactAutoRefinePending = false; } - } catch { + } catch (error) { // Auto-refine is opportunistic; manual /refine remains available. // Stamp the cooldown so a persistently failing refine doesn't retry // (via a retained pending review) on every agent end. this._lastAutoRefineReviewAt = Date.now(); + if (error instanceof RefineSkippedError) { + // A skipped round is consumed like a reviewer decline, not retained for retry. + this._pendingAutoRefineReview = undefined; + this._turnIntervalAutoRefinePending = false; + this._assistantTurnsSinceAutoRefine = 0; + if (reason === "compact") this._compactAutoRefinePending = false; + } } finally { this._autoRefineInProgress = false; this._scheduleDeferredAutoRefineIfIdle(); @@ -7863,7 +7891,7 @@ export class AgentSession { rollbackId?: string; global?: boolean; } = {}, - internal: { skipAbort?: boolean } = {}, + internal: { skipAbort?: boolean; trigger?: "manual" | "auto" } = {}, ): Promise { // Queued /refine executes from the session-input pump between turns; // refine never aborts the agent (planning is backgrounded and the apply @@ -7903,7 +7931,7 @@ export class AgentSession { const refineAbort = new AbortController(); this._refineAbortController = refineAbort; - const planRun = this._planRefine(options, refineAbort.signal); + const planRun = this._planRefine(options, refineAbort.signal, internal.trigger ?? "manual"); const planSettled = planRun.then( () => undefined, () => undefined, @@ -7989,6 +8017,7 @@ export class AgentSession { private async _planRefine( options: { instructions?: string; rollbackId?: string; global?: boolean }, signal: AbortSignal, + trigger: "manual" | "auto" = "manual", ): Promise { if (this._disposed) { throw new Error("Cannot refine a disposed session."); @@ -8030,6 +8059,33 @@ export class AgentSession { : baselineScope === "global" ? globalPlanningState : localPlanningState!; + if (!options.rollbackId && this._extensionRunner.hasHandlers("session_before_refine")) { + const result = (await this._extensionRunner.emit({ + type: "session_before_refine", + preparation: { + trigger, + instructions: options.instructions, + scope: requestedScope, + planningState, + history, + conversationText: serializeConversation(convertToLlm(this.agent.state.messages)).slice(-80_000), + }, + signal, + })) as SessionBeforeRefineResult | undefined; + if (this._disposed || signal.aborted) { + throw new Error("Refinement cancelled because the session was disposed."); + } + if (result?.skip) { + throw new RefineSkippedError("Refinement skipped by extension"); + } + if (result?.proposal !== undefined) { + return { + proposal: normalizeRefinementProposal(result.proposal), + id: generateRefinementId(), + baselineState, + }; + } + } const plan = await planRefinement( this.agent.state.messages, planningState, diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 26ea3b229..ba9631820 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -89,6 +89,7 @@ export type { ProviderConfig, ProviderModelConfig, RefineCompleteEvent, + RefinePreparation, RegisteredCommand, RegisteredTool, ReplacedSessionContext, @@ -102,6 +103,8 @@ export type { SessionBeforeCompactResult, SessionBeforeForkEvent, SessionBeforeForkResult, + SessionBeforeRefineEvent, + SessionBeforeRefineResult, SessionBeforeSwitchEvent, SessionBeforeSwitchResult, SessionBeforeTreeEvent, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 618aa1943..172f72500 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -46,6 +46,7 @@ import type { ResourcesDiscoverResult, SessionBeforeCompactResult, SessionBeforeForkResult, + SessionBeforeRefineResult, SessionBeforeSwitchResult, SessionBeforeTreeResult, SessionShutdownEvent, @@ -127,13 +128,21 @@ type RunnerEmitEvent = Exclude< type SessionBeforeEvent = Extract< RunnerEmitEvent, - { type: "session_before_switch" | "session_before_fork" | "session_before_compact" | "session_before_tree" } + { + type: + | "session_before_switch" + | "session_before_fork" + | "session_before_compact" + | "session_before_refine" + | "session_before_tree"; + } >; type SessionBeforeEventResult = | SessionBeforeSwitchResult | SessionBeforeForkResult | SessionBeforeCompactResult + | SessionBeforeRefineResult | SessionBeforeTreeResult; type RunnerEmitResult = TEvent extends { type: "session_before_switch" } @@ -142,9 +151,11 @@ type RunnerEmitResult = TEvent extends { type: " ? SessionBeforeForkResult | undefined : TEvent extends { type: "session_before_compact" } ? SessionBeforeCompactResult | undefined - : TEvent extends { type: "session_before_tree" } - ? SessionBeforeTreeResult | undefined - : undefined; + : TEvent extends { type: "session_before_refine" } + ? SessionBeforeRefineResult | undefined + : TEvent extends { type: "session_before_tree" } + ? SessionBeforeTreeResult | undefined + : undefined; export type ExtensionErrorListener = (error: ExtensionError) => void; @@ -667,6 +678,7 @@ export class ExtensionRunner { event.type === "session_before_switch" || event.type === "session_before_fork" || event.type === "session_before_compact" || + event.type === "session_before_refine" || event.type === "session_before_tree" ); } @@ -685,7 +697,7 @@ export class ExtensionRunner { if (this.isSessionBeforeEvent(event) && handlerResult) { result = handlerResult as SessionBeforeEventResult; - if (result.cancel) { + if (("cancel" in result && result.cancel) || ("skip" in result && result.skip)) { return result as RunnerEmitResult; } } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 289405fdf..4ff321fdf 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -49,6 +49,7 @@ import type { ReadonlyFooterDataProvider } from "../footer-data-provider.js"; import type { KeybindingsManager } from "../keybindings.js"; import type { CustomMessage } from "../messages.js"; import type { ModelRegistry } from "../model-registry.js"; +import type { HarnessState, RefinementProposal, RefinementResult } from "../refinement/index.js"; import type { BranchSummaryEntry, CompactionEntry, @@ -517,6 +518,36 @@ export interface SessionBeforeCompactEvent { signal: AbortSignal; } +/** Planning inputs for a refinement round. */ +export interface RefinePreparation { + /** Whether refinement was requested explicitly (/refine) or by auto-refine. */ + trigger: "manual" | "auto"; + /** Instructions passed to /refine or derived from the auto-refine review. */ + instructions?: string; + /** Whether the round targets the global or the session-local harness. */ + scope: "global" | "local"; + /** Harness state the planner would see (global overlaid with local for local scope). */ + planningState: HarnessState; + /** Prior refinement results, newest last. */ + history: RefinementResult[]; + /** Serialized conversation text the built-in planner would consume. */ + conversationText: string; +} + +/** Fired before refinement planning (can be skipped or replaced). Rollbacks bypass this event. */ +export interface SessionBeforeRefineEvent { + type: "session_before_refine"; + preparation: RefinePreparation; + signal: AbortSignal; +} + +export interface SessionBeforeRefineResult { + /** Skip this refinement round entirely. */ + skip?: boolean; + /** Replace the built-in planner. Edits are still validated at apply time. */ + proposal?: RefinementProposal; +} + /** Fired after context compaction */ export interface SessionCompactEvent { type: "session_compact"; @@ -569,6 +600,7 @@ export type SessionEvent = | SessionBeforeForkEvent | SessionBeforeCompactEvent | SessionCompactEvent + | SessionBeforeRefineEvent | SessionShutdownEvent | SessionBeforeTreeEvent | SessionTreeEvent; @@ -970,6 +1002,10 @@ export interface ExtensionAPI { handler: ExtensionHandler, ): void; on(event: "session_compact", handler: ExtensionHandler): void; + on( + event: "session_before_refine", + handler: ExtensionHandler, + ): void; on(event: "session_shutdown", handler: ExtensionHandler): void; on(event: "session_before_tree", handler: ExtensionHandler): void; on(event: "session_tree", handler: ExtensionHandler): void; diff --git a/packages/coding-agent/src/core/refinement/refinement.ts b/packages/coding-agent/src/core/refinement/refinement.ts index b19b33db6..8b0c2d89a 100644 --- a/packages/coding-agent/src/core/refinement/refinement.ts +++ b/packages/coding-agent/src/core/refinement/refinement.ts @@ -630,12 +630,13 @@ function extractJsonObject(text: string): unknown { throw new Error("Refiner did not return a JSON object"); } -function parseProposal(text: string): RefinementProposal { - const value = extractJsonObject(text); - if (typeof value !== "object" || value === null || Array.isArray(value)) { - throw new Error("Refiner JSON must be an object"); - } - const record = value as Record; +/** + * Normalizes an untrusted refinement proposal while preserving invalid edit + * fields for apply-time validation. + */ +export function normalizeRefinementProposal(value: unknown): RefinementProposal { + const record = + typeof value === "object" && value !== null && !Array.isArray(value) ? (value as Record) : {}; const edits = Array.isArray(record.edits) ? record.edits : []; return { summary: typeof record.summary === "string" ? record.summary : "Refined continual harness state", @@ -661,6 +662,14 @@ function parseProposal(text: string): RefinementProposal { }; } +function parseProposal(text: string): RefinementProposal { + const value = extractJsonObject(text); + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error("Refiner JSON must be an object"); + } + return normalizeRefinementProposal(value); +} + function validateEdit(edit: RefinementEdit, computedId?: string): string | undefined { if (!["create", "update", "delete"].includes(edit.action)) { return `unsupported action ${String(edit.action)}`; @@ -860,6 +869,14 @@ export interface RefinementPlan { * here can take many seconds, during which the kernel or another session may write * the shared `harness_state.json`. */ +/** Mint a refinement id in the canonical `refine_` format. */ +export function generateRefinementId(): string { + return `refine_${new Date() + .toISOString() + .replace(/[^0-9]/g, "") + .slice(0, 17)}`; +} + export async function planRefinement( messages: AgentMessage[], state: HarnessState, @@ -871,10 +888,7 @@ export async function planRefinement( signal?: AbortSignal, thinkingLevel?: ThinkingLevel, ): Promise { - const id = `refine_${new Date() - .toISOString() - .replace(/[^0-9]/g, "") - .slice(0, 17)}`; + const id = generateRefinementId(); if (options.rollbackId) { const target = history.find((item) => item.id === options.rollbackId); if (!target) { diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 093da58de..1eaa7a092 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -94,11 +94,14 @@ export type { ProviderConfig, ProviderModelConfig, RefineCompleteEvent, + RefinePreparation, RegisteredCommand, RegisteredTool, ResolvedCommand, SessionBeforeCompactEvent, SessionBeforeForkEvent, + SessionBeforeRefineEvent, + SessionBeforeRefineResult, SessionBeforeSwitchEvent, SessionBeforeTreeEvent, SessionCompactEvent, @@ -148,6 +151,12 @@ export type { ResolvedResource, } from "./core/package-manager.js"; export { DefaultPackageManager } from "./core/package-manager.js"; +export type { + HarnessState, + RefinementEdit, + RefinementProposal, + RefinementResult, +} from "./core/refinement/index.js"; export type { ResourceCollision, ResourceDiagnostic, ResourceLoader } from "./core/resource-loader.js"; export { DefaultResourceLoader, loadProjectContextFiles } from "./core/resource-loader.js"; // SDK for programmatic usage diff --git a/packages/coding-agent/test/suite/agent-session-queue.test.ts b/packages/coding-agent/test/suite/agent-session-queue.test.ts index 658db2182..bdf865a9c 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/packages/coding-agent/test/suite/agent-session-queue.test.ts @@ -232,6 +232,7 @@ describe("AgentSession queue characterization", () => { for (const fragment of refineFragments) { expect(refine).toHaveBeenCalledWith( expect.objectContaining({ instructions: expect.stringContaining(fragment) }), + { trigger: "auto" }, ); } } @@ -442,6 +443,7 @@ describe("AgentSession queue characterization", () => { expect(refine).toHaveBeenCalledWith( expect.objectContaining({ instructions: expect.stringContaining("durable lesson") }), + { trigger: "auto" }, ); expect(guardWasSetDuringRefine).toBe(true); expect(internals._autoRefineInProgress).toBe(false); diff --git a/packages/coding-agent/test/suite/agent-session-refine-extension.test.ts b/packages/coding-agent/test/suite/agent-session-refine-extension.test.ts new file mode 100644 index 000000000..98c3db2fb --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-refine-extension.test.ts @@ -0,0 +1,270 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { RefineSkippedError } from "../../src/core/agent-session.js"; +import type { SessionBeforeRefineEvent } from "../../src/core/extensions/index.js"; +import { loadHarnessState, type RefinementProposal } from "../../src/core/refinement/index.js"; +import { createHarness, type Harness } from "./harness.js"; + +describe("AgentSession session_before_refine extension hook", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("applies an extension-provided proposal without calling the built-in planner", async () => { + const events: SessionBeforeRefineEvent[] = []; + const harness = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async (event) => { + events.push(event); + return { + proposal: { + summary: "extension summary", + rationale: "extension rationale", + expectedOutcome: "extension outcome", + edits: [ + { + action: "create" as const, + kind: "memory" as const, + title: "Extension memory", + content: "Captured by the extension planner", + }, + ], + }, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([]); + await harness.session.prompt("hello").catch(() => {}); + + const result = await harness.session.refine({ instructions: "capture lessons" }); + + expect(result.summary).toBe("extension summary"); + expect(result.appliedEdits).toHaveLength(1); + expect(result.appliedEdits[0]?.applied).toBe(true); + expect(events).toHaveLength(1); + expect(events[0]?.preparation.trigger).toBe("manual"); + expect(events[0]?.preparation.scope).toBe("local"); + expect(events[0]?.preparation.instructions).toBe("capture lessons"); + expect(harness.getPendingResponseCount()).toBe(0); + + const state = loadHarnessState(result.harnessStatePath.replace(/\/[^/]+$/, ""), "local"); + const memories = Object.values(state.entries.memory); + expect(memories.some((entry) => entry.title === "Extension memory")).toBe(true); + }); + + it("rejects invalid extension edits at apply time", async () => { + const harness = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async () => ({ + proposal: { + summary: "bad plan", + rationale: "bad", + expectedOutcome: "bad", + edits: [ + // update without id is invalid + { action: "update" as const, kind: "memory" as const, title: "x", content: "y" }, + ], + }, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([]); + await harness.session.prompt("hello").catch(() => {}); + + const result = await harness.session.refine(); + + expect(result.appliedEdits[0]?.applied).toBe(false); + expect(result.appliedEdits[0]?.error).toContain("requires id"); + }); + + it("normalizes malformed runtime extension proposals before applying them", async () => { + const proposals: Array<{ proposal: unknown; expectedEdits: number }> = [ + { proposal: { summary: "missing edits" }, expectedEdits: 0 }, + { proposal: { edits: "not an array" }, expectedEdits: 0 }, + { + proposal: { + edits: [null, "not an edit", { action: "update", kind: "memory", title: "x", content: "y" }], + }, + expectedEdits: 1, + }, + ]; + + for (const { proposal, expectedEdits } of proposals) { + const harness = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async () => ({ + proposal: proposal as RefinementProposal, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([]); + await harness.session.prompt("hello").catch(() => {}); + + const result = await harness.session.refine(); + + expect(result.appliedEdits).toHaveLength(expectedEdits); + if (expectedEdits > 0) { + expect(result.appliedEdits[0]?.applied).toBe(false); + expect(result.appliedEdits[0]?.error).toContain("requires id"); + } + expect(harness.eventsOfType("refine_failed")).toHaveLength(0); + } + }); + + it("skips the refinement round when an extension returns skip", async () => { + const harness = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async () => ({ skip: true })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([]); + await harness.session.prompt("hello").catch(() => {}); + + await expect(harness.session.refine()).rejects.toThrow(RefineSkippedError); + }); + + it("falls back to the built-in planner when the handler returns nothing", async () => { + let handlerCalls = 0; + const harness = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async () => { + handlerCalls += 1; + return undefined; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([]); + await harness.session.prompt("hello").catch(() => {}); + + const internals = harness.session as unknown as { + _planRefine(options: unknown, signal: AbortSignal): Promise; + }; + // The handler runs but does not short-circuit: planning proceeds to the + // built-in planner LLM call, which fails here (no faux response queued) + // rather than being skipped. + const refineAbort = new AbortController(); + await expect(internals._planRefine({ instructions: "x" }, refineAbort.signal)).rejects.not.toThrow( + RefineSkippedError, + ); + expect(handlerCalls).toBe(1); + }); + + it("consumes a non-serialized auto-refine round when an extension skips it", async () => { + let handlerCalls = 0; + const harness = await createHarness({ + persistSession: true, + settings: { autoRefine: { enabled: true, turnInterval: 1, cooldownMs: 0 } }, + autoRefineReviewer: async () => ({ shouldRefine: true, rationale: "durable lesson" }), + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async () => { + handlerCalls += 1; + return { skip: true }; + }); + }, + ], + }); + harnesses.push(harness); + const internals = harness.session as unknown as { + _maybeAutoRefine(reason: "turn_interval"): Promise; + _assistantTurnsSinceAutoRefine: number; + _turnIntervalAutoRefinePending: boolean; + _pendingAutoRefineReview?: unknown; + }; + internals._assistantTurnsSinceAutoRefine = 1; + + await internals._maybeAutoRefine("turn_interval"); + + expect(handlerCalls).toBe(1); + expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._turnIntervalAutoRefinePending).toBe(false); + expect(internals._pendingAutoRefineReview).toBeUndefined(); + expect(harness.eventsOfType("refine_failed")).toHaveLength(0); + + await internals._maybeAutoRefine("turn_interval"); + expect(handlerCalls).toBe(1); + }); + + it("marks serialized auto-refine as trigger auto and treats extension skip as a non-failure", async () => { + const events: SessionBeforeRefineEvent[] = []; + const harness = await createHarness({ + persistSession: true, + serializedRefine: true, + autoRefineReviewer: async () => ({ + shouldRefine: true, + rationale: "durable lesson", + instructions: "capture the lesson", + }), + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async (event) => { + events.push(event); + return { skip: true }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([]); + await harness.session.prompt("hello").catch(() => {}); + + const internals = harness.session as unknown as { + _runSerializedAutoRefineReview(reason: "compact" | "turn_interval", branchVersion: number): Promise; + _autoRefineBranchVersion: number; + }; + await internals._runSerializedAutoRefineReview("turn_interval", internals._autoRefineBranchVersion); + + expect(events).toHaveLength(1); + expect(events[0]?.preparation.trigger).toBe("auto"); + expect(harness.eventsOfType("refine_failed")).toHaveLength(0); + }); + + it("emits refine_failed when disposal drains an extension-skipped explicit serialized refine.run", async () => { + const harness = await createHarness({ + persistSession: true, + serializedRefine: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", async () => ({ skip: true })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([]); + await harness.session.prompt("hello").catch(() => {}); + + (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; + harness.session.handleRefineHostRequest("refine.run", { instructions: "x" }); + (harness.session.agent.state as { isStreaming: boolean }).isStreaming = false; + + await harness.session.disposeAsync(); + + const failures = harness.eventsOfType("refine_failed"); + expect(failures).toHaveLength(1); + expect(failures[0]?.error).toBe("Refinement skipped by extension"); + }); +}); -- 2.51.2