Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 55 additions & 2 deletions .pi/extensions/fm-branch-supervision.ts
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ import {
getAgentDir,
keyHint,
ModelRuntime,
type ModelRegistry,
SessionManager,
ToolExecutionComponent,
type AgentSession,
Expand Down Expand Up @@ -644,6 +645,11 @@ export default function (pi: ExtensionAPI) {
// extension plus its model_select event, because createBranch runs at wake
// time with no context of its own. It is what "follow main" applies.
let mainModel: { provider: string; id: string } | null = null;
// Main's own model registry, captured from the contexts Pi hands this
// extension the same way mainModel is. It is the ONLY read path to
// providers an extension registered at runtime (pi-devin-auth's "devin"),
// which the branch's isolated ModelRuntime cannot see on its own.
let mainModelRegistry: ModelRegistry | null = null;

// Main's own current effort needs no such tracking: Pi answers it directly
// on demand, including at wake time. It throws only when the extension
Expand All @@ -657,8 +663,9 @@ export default function (pi: ExtensionAPI) {
}
}

function rememberMainModel(ctx?: { model?: { provider: string; id: string } }): void {
function rememberMainModel(ctx?: { model?: { provider: string; id: string }; modelRegistry?: ModelRegistry }): void {
if (ctx?.model) mainModel = { provider: ctx.model.provider, id: ctx.model.id };
if (ctx?.modelRegistry) mainModelRegistry = ctx.modelRegistry;
}

function deliverBranchHealthNote(text: string): void {
Expand Down Expand Up @@ -708,10 +715,55 @@ export default function (pi: ExtensionAPI) {
// and same user as main, so stored credentials keep their own semantics
// (OAuth stays OAuth, an API key stays an API key) and nothing is ever
// installed, converted, derived, or overwritten here.
// A provider that exists only because an extension registered it into
// main's runtime (pi-devin-auth's "devin", whose streamSimple is the custom
// gRPC path no static catalog can express) is invisible to an isolated
// branch runtime until its registration is copied across. The config object
// carries that streamSimple and oauth wiring by reference, so copying it
// reuses the provider's own registration rather than reimplementing its
// wire protocol; the copy is never persisted and stays scoped to this one
// runtime. One registration that fails to compose must not blind the rest,
// so each copy is isolated. A just-registered provider's auth check has not
// run yet, so the copied providers are refreshed here and every caller's
// hasConfiguredAuth verdict is real rather than the provisional entry
// registration leaves behind.
async function copyExtensionProviders(modelRuntime: ModelRuntime): Promise<void> {
if (!mainModelRegistry) return;
let providerIds: readonly string[];
try {
providerIds = mainModelRegistry.getRegisteredProviderIds();
} catch {
return;
}
const copied: string[] = [];
for (const providerId of providerIds) {
try {
const config = mainModelRegistry.getRegisteredProviderConfig(providerId);
if (config) {
modelRuntime.registerProvider(providerId, config);
copied.push(providerId);
}
} catch {
// A registration that fails to compose in the isolated runtime leaves
// that provider unavailable, exactly as if it were never copied.
}
}
if (copied.length === 0) return;
try {
await modelRuntime.refresh({ providers: copied, allowNetwork: false });
} catch {
// A failed availability refresh is answered by hasConfiguredAuth.
}
}

async function resolveBranchModel(provider: string, modelId: string): Promise<BranchModelResolution> {
const label = `${provider}/${modelId}`;
const modelRuntime = await ModelRuntime.create();
const model = modelRuntime.getModel(provider, modelId) as BranchModel | undefined;
let model = modelRuntime.getModel(provider, modelId) as BranchModel | undefined;
if (!model) {
await copyExtensionProviders(modelRuntime);
model = modelRuntime.getModel(provider, modelId) as BranchModel | undefined;
}
if (!model) return { ok: false, reason: `${label} is unavailable to the isolated branch runtime` };
if (!modelRuntime.hasConfiguredAuth(provider)) {
return { ok: false, reason: `${label} has no configured credentials in the isolated branch runtime` };
Expand Down Expand Up @@ -1659,6 +1711,7 @@ ${context.command}
let available: string[];
try {
const modelRuntime = await ModelRuntime.create();
await copyExtensionProviders(modelRuntime);
available = ctx.modelRegistry
.getAvailable()
.filter((model) => modelRuntime.getModel(model.provider, model.id) && modelRuntime.hasConfiguredAuth(model.provider))
Expand Down
3 changes: 2 additions & 1 deletion docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,8 @@ The effort list is a handful of levels and stays on Pi's plain selector dialog.
Both picks change the supervision branch alone and never the captain's own conversation model or effort.
It persists the model pick in gitignored `config/supervision-branch-model` and the effort pick in gitignored `config/supervision-branch-effort`, both under the effective Firstmate home, resolved from `FM_HOME`, then `FM_ROOT_OVERRIDE`, then the tracked code root derived from the extension path, or under `FM_CONFIG_OVERRIDE` when that test and specialized-setup override is present.
Firstmate keeps no model catalog of its own; the list is the intersection of what Pi reports when the picker opens and what a fresh isolated branch runtime can run.
A provider that exists only because an extension registered it inside the captain's session is not offered, while stored OAuth and API-key credentials retain their native credential type because Firstmate never copies, converts, installs, or overwrites credentials for the branch runtime.
A provider that exists only because an extension registered it inside the captain's session, such as pi-devin-auth's `devin`, is offered and can be pinned or followed like any other; [pi-supervision-branch.md](pi-supervision-branch.md#cost-model-and-the-byte-stable-prefix) owns how that registration reaches the isolated branch runtime.
Stored OAuth and API-key credentials retain their native credential type because Firstmate never copies, converts, installs, or overwrites credentials for the branch runtime.
The file holds one `<provider>/<model-id>` line followed by one newline, split at the first `/` so a provider-qualified model id such as `openrouter/anthropic/claude-sonnet-4-5` survives intact.
An absent, unreadable, or unparseable file means no pin, and the branch then follows main's own current model, applied explicitly and live whenever main changes models mid-session.
A valid pin wins over main and remains unaffected by main's model changes.
Expand Down
2 changes: 2 additions & 0 deletions docs/pi-supervision-branch.md
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,8 @@ Every other fleet-wide or unresolvable wake - including watcher-failure alarms,
The captain accepted the normal provider prompt-caching strategy: a byte-identical branch prefix generated once per firstmate version, the same tool set in the same order on every request, and one shared `prompt_cache_key` per home for all branch sessions (set in a `before_provider_request` hook, and only for providers whose requests already carry that field); main keeps its own per-session key.
Budget roughly 60% cache hits on a new branch conversation's first call and 95% on later calls within that conversation; the shared per-home key is what carries the byte-identical prefix across the conversation each main session start opens, and reuse is best-effort, never guaranteed.
The branch can also run on a cheaper model and a shallower reasoning effort than main, both pinned with the Pi `/supervision-model` command; [configuration.md](configuration.md#pi-supervision-branch-model-and-effort-configsupervision-branch-model-configsupervision-branch-effort) owns those pins' operator-facing schema and unpinned behavior.
A provider an extension registered only into main's runtime, such as pi-devin-auth's `devin`, reaches the isolated branch runtime by copying its provider config from main's captured `ModelRegistry` into the branch `ModelRuntime` at model-resolution time and in the `/supervision-model` picker, so the provider's own `streamSimple` transport and OAuth wiring are reused by reference rather than reimplemented.
That carve-out is scoped to provider registration alone: the branch keeps its `noExtensions`, `noSkills`, and `noContextFiles` isolation, the copy is never persisted, a provider whose registration fails to compose is simply unavailable, and `tests/fm-pi-branch-extension.test.sh` pins the pin-and-fallthrough behavior.
No caching machinery beyond this exists, deliberately: any later dynamic content in the branch prefix silently removes most of the cache benefit, which is why `bin/fm-branch-prompt.sh`'s header is the contract's single owner and `tests/fm-branch-supervision.test.sh` pins the output to byte identity.

## Away mode
Expand Down
111 changes: 111 additions & 0 deletions tests/fm-pi-branch-extension.test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,10 @@ export class ModelRuntime {
constructor() {
this.models = (globalThis.__fmBranchStaticModels?.() ?? []).map((model) => ({ ...model }));
this.authenticated = new Set(this.models.filter((model) => model.storedAuth !== false).map((model) => model.provider));
this.registeredProviderConfigs = new Map();
// Like the real runtime, a registered provider's credentials are only
// known once refresh() has run for it; registration alone is provisional.
this.pendingAuth = new Set();
}
static async create() {
const queuedError = globalThis.__fmModelRuntimeErrors?.shift();
Expand All @@ -115,6 +119,18 @@ export class ModelRuntime {
(globalThis.__fmModelRuntimes ??= []).push(runtime);
return runtime;
}
registerProvider(providerId, config) {
this.registeredProviderConfigs.set(providerId, config);
for (const model of config.models ?? []) {
this.models.push({ ...model, provider: providerId });
}
if (config.oauth || config.apiKey) this.pendingAuth.add(providerId);
}
async refresh(options) {
for (const providerId of options?.providers ?? this.pendingAuth) {
if (this.pendingAuth.delete(providerId)) this.authenticated.add(providerId);
}
}
getModel(provider, id) {
return this.models.find((model) => model.provider === provider && model.id === id);
}
Expand Down Expand Up @@ -446,6 +462,8 @@ const modelRegistry = {
getAvailable: () => registryModels.filter((model) => model.mainAvailable !== false).slice(),
find: (provider, id) => registryModels.find((model) => model.provider === provider && model.id === id),
hasConfiguredAuth: (model) => model.mainAvailable !== false,
getRegisteredProviderConfig: (providerId) => globalThis.__fmExtensionProviderConfigs?.get(providerId),
getRegisteredProviderIds: () => [...(globalThis.__fmExtensionProviderConfigs?.keys() ?? [])],
};
function makeCtx(extra) {
return {
Expand Down Expand Up @@ -4735,6 +4753,98 @@ EOF
pass "a failed cursor write re-delivers a routine note exactly once more while a captain outcome stays deduplicated"
}

test_extension_registered_provider_resolves_in_the_branch() {
local repo home out status
repo="$TMP_ROOT/extprov-root"
home="$TMP_ROOT/extprov-home"
mkdir -p "$home/state" "$home/config"
install_pi_branch_extension_fixture "$repo"
PLUGIN="$repo/.pi/extensions/fm-branch-supervision.ts" FM_HOME="$home" FM_ROOT_OVERRIDE="$ROOT" \
DRIVER_PRELUDE="$DRIVER_PRELUDE" node --input-type=module > "$TMP_ROOT/node-output" 2>&1 <<'EOF'
const prelude = process.env.DRIVER_PRELUDE;
await eval(`(async () => { ${prelude}; globalThis.__t = { fire, dispatch, settle, makeCtx, registryModels, uiSelections, uiPrompts, notices, commands, home }; })()`);
const { fire, dispatch, settle, makeCtx, registryModels, uiSelections, uiPrompts, notices, commands, home } = globalThis.__t;
import { readFileSync, writeFileSync } from "node:fs";

// An extension-registered provider exists only in main's registry, never in
// the isolated branch runtime's static catalog. Registering its config on
// main's registry is what makes it resolvable for the branch.
registryModels.push(
{ provider: "anthropic", id: "main-model" },
// Available in main's registry but absent from the branch runtime's static
// catalog, exactly like a provider an extension registered at runtime.
{ provider: "devin", id: "swe-1-7", branchAvailable: false },
);
globalThis.__fmExtensionProviderConfigs = new Map([
[
"devin",
{
name: "Devin (Cognition)",
api: "devin-cloud",
baseUrl: "https://server.codeium.com",
models: [{ id: "swe-1-7", name: "SWE 1.7", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 200000, maxTokens: 8192 }],
oauth: { name: "Devin (Cognition / Windsurf)", login: async () => ({}), refreshToken: async (c) => c, getApiKey: (c) => c.access },
streamSimple: () => {},
},
],
]);

await fire("session_start", {}, makeCtx());

// The picker must offer the extension-registered model: it is available in
// main's registry and resolvable in the branch runtime once its registration
// is copied across.
const command = commands.get("supervision-model");
if (!command) throw new Error("the supervision-model command was not registered");
uiSelections.push("devin/swe-1-7");
await command.handler("", makeCtx());
const offered = uiPrompts[0];
if (!offered.options.includes("devin/swe-1-7")) {
throw new Error(`the picker must offer an extension-registered provider the branch can run: ${JSON.stringify(offered.options)}`);
}
if (readFileSync(`${home}/config/supervision-branch-model`, "utf8") !== "devin/swe-1-7\n") {
throw new Error("the extension-registered pick was not persisted");
}
dispatch("signal: extension provider pin");
await settle(() => (globalThis.__fmSessions ?? []).length === 1, "pinned extension-provider branch build");
const pinned = globalThis.__fmSessions[0].options.model;
if (!pinned || pinned.provider !== "devin" || pinned.id !== "swe-1-7") {
throw new Error(`the extension-registered pin did not bind the branch: ${JSON.stringify(pinned)}`);
}
// Copying the provider registration must not loosen the branch's isolation:
// the devin-pinned session still loads no extensions, skills, or context files.
const pinnedLoader = globalThis.__fmLoaders.at(-1);
for (const key of ["noExtensions", "noSkills", "noContextFiles"]) {
if (pinnedLoader.options[key] !== true) throw new Error(`devin-pinned branch loader must keep ${key}`);
}

// Without the registration, the same pin is unavailable and the branch
// refuses to build rather than silently downgrading.
globalThis.__fmExtensionProviderConfigs = new Map();
await fire("session_shutdown", {});
await fire("session_start", {}, makeCtx());
const unregisteredOffer = dispatch("signal: unregistered provider pin");
if (!unregisteredOffer.accepted) throw new Error("unregistered-pin wake was not initially accepted");
const unregisteredFailure = await unregisteredOffer.settlement.then(
() => null,
(error) => error,
);
if (
!(unregisteredFailure instanceof Error) ||
!unregisteredFailure.message.includes("devin/swe-1-7") ||
!unregisteredFailure.message.includes("supervision model pin")
) {
throw new Error(`the unregistered pin did not reject with its own name: ${String(unregisteredFailure)}`);
}
if ((globalThis.__fmSessions ?? []).length !== 1) throw new Error("an unregistered pin must not build a second branch session");
process.exit(0);
EOF
status=$?
out=$(cat "$TMP_ROOT/node-output")
expect_code 0 "$status" "an extension-registered provider must resolve in the isolated branch runtime: $out"
pass "an extension-registered provider resolves in the isolated branch runtime"
}

test_outcomes_tool_uses_stock_execution_and_export_consumers
test_real_pi_picker_primitives_stay_bounded_and_searchable
test_branch_dispatch_two_stage_filter_and_prefix_contract
Expand Down Expand Up @@ -4763,6 +4873,7 @@ test_supervision_model_picker_is_bounded_searchable_and_branch_only
test_branch_model_picker_keeps_follow_main_first_under_ranking
test_branch_effort_pin_applies_and_absent_pin_follows_main
test_unpinned_branch_follows_main_effort_changes_live
test_extension_registered_provider_resolves_in_the_branch
test_supervision_model_command_picks_effort_after_the_model
test_unusable_model_pin_falls_back_to_main
test_replacement_activation_cleans_leases_and_retries_failure
Expand Down
Loading