/**
* 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": "..."}]}
<conversation>
${conversationText}
</conversation>`,
},
],
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;
}
});
}