diff --git a/plugins/provider-pi/src/bridge/bb-pi-extension.ts b/plugins/provider-pi/src/bridge/bb-pi-extension.ts index 67ca993ce5..1fcea52395 100644 --- a/plugins/provider-pi/src/bridge/bb-pi-extension.ts +++ b/plugins/provider-pi/src/bridge/bb-pi-extension.ts @@ -245,6 +245,8 @@ export default function bbExtension(pi) { } case "leaf": return { leafId: currentLeafId() }; + case "model-scope": + return currentModelScope(); default: throw new Error("unknown bridge request " + String(message.method)); } @@ -254,6 +256,17 @@ export default function bbExtension(pi) { return sessionContext?.sessionManager?.getLeafId?.() ?? null; } + function currentModelScope() { + return { + scopedModelIds: (sessionContext?.scopedModels ?? []).map( + ({ model }) => model.provider + "/" + model.id, + ), + defaultModelId: sessionContext?.model + ? sessionContext.model.provider + "/" + sessionContext.model.id + : null, + }; + } + for (const tool of tools) { pi.registerTool({ name: tool.name, @@ -291,6 +304,10 @@ export default function bbExtension(pi) { pi.on("session_start", async (_event, ctx) => { sessionContext = ctx; + writeLine(CHILD_TO_BRIDGE_FD, { + kind: "model-scope", + ...currentModelScope(), + }); // Pi's active-tool set is session state; a resumed or forked session can // predate the bb tools, so make sure every injected tool is active. if (tools.length > 0 && typeof pi.setActiveTools === "function") { diff --git a/plugins/provider-pi/src/bridge/catalog.test.ts b/plugins/provider-pi/src/bridge/catalog.test.ts new file mode 100644 index 0000000000..0695aba7a3 --- /dev/null +++ b/plugins/provider-pi/src/bridge/catalog.test.ts @@ -0,0 +1,54 @@ +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { BB_PI_EXTENSION_SOURCE } from "./bb-pi-extension.js"; +import { closeAllPiCatalogs, getPiCatalog } from "./catalog.js"; +import { PI_BRIDGE_ARGS_ENV, PI_BRIDGE_COMMAND_ENV } from "./rpc-child.js"; +import { fakePiPath } from "./test-support.js"; + +const originalEnv = { ...process.env }; +const tempDirs: string[] = []; + +afterEach(async () => { + await closeAllPiCatalogs(); + process.env = { ...originalEnv }; + for (const dir of tempDirs.splice(0)) { + rmSync(dir, { recursive: true, force: true }); + } +}); + +describe("pi catalog child generations", () => { + it("re-reads model scope after the catalog child restarts", async () => { + const workspace = mkdtempSync(join(tmpdir(), "bb-pi-catalog-")); + tempDirs.push(workspace); + const extensionPath = join(workspace, "bb-extension.mjs"); + const spawnCounterPath = join(workspace, "spawns.txt"); + const processLogPath = join(workspace, "processes.txt"); + writeFileSync(extensionPath, BB_PI_EXTENSION_SOURCE); + + process.env[PI_BRIDGE_COMMAND_ENV] = process.execPath; + process.env[PI_BRIDGE_ARGS_ENV] = JSON.stringify([fakePiPath]); + process.env.FAKE_PI_SPAWN_COUNTER_FILE = spawnCounterPath; + process.env.FAKE_PI_PROCESS_LOG = processLogPath; + process.env.FAKE_PI_SCOPE_BY_SPAWN = "1"; + process.env.FAKE_PI_EXIT_AFTER_FIRST_AVAILABLE = "1"; + + const catalog = await getPiCatalog(workspace, extensionPath); + const first = await catalog.listModels(); + expect(first.models.map((model) => model.id)).toEqual([ + "fake-provider/fake-model", + ]); + + await vi.waitFor(() => { + expect(readFileSync(processLogPath, "utf8")).toContain("exit:"); + }); + + const second = await catalog.listModels(); + expect(second.models.map((model) => model.id)).toEqual([ + "fake-provider/fake-mini", + ]); + expect(second.models[0]?.isDefault).toBe(true); + expect(readFileSync(spawnCounterPath, "utf8")).toBe("2"); + }); +}); diff --git a/plugins/provider-pi/src/bridge/catalog.ts b/plugins/provider-pi/src/bridge/catalog.ts index 1bd287f84f..67812e9410 100644 --- a/plugins/provider-pi/src/bridge/catalog.ts +++ b/plugins/provider-pi/src/bridge/catalog.ts @@ -5,7 +5,11 @@ import { type PiModelContextWindowResolver, } from "../delta-translation.js"; import { buildPiAvailableModels, type PiCatalogModel } from "../model-list.js"; -import { PiRpcChild, buildPiChildEnv } from "./rpc-child.js"; +import { + PiRpcChild, + PiRpcChildExitedError, + buildPiChildEnv, +} from "./rpc-child.js"; const EXTENDED_THINKING_LEVELS = [ "off", @@ -85,24 +89,90 @@ async function spawnCatalog( extensionPath: string, touch: () => void, ): Promise { - let child: PiRpcChild | null = null; - const spawnChild = (): PiRpcChild => { - if (child !== null && !child.exited) { - return child; - } - child = new PiRpcChild({ + interface CatalogChildGeneration { + child: PiRpcChild; + ready: Promise>; + getModelScope(): + | { scopedModelIds: string[]; defaultModelId?: string } + | undefined; + } + + let generation: CatalogChildGeneration | null = null; + const spawnGeneration = (): CatalogChildGeneration => { + let modelScope: + | { scopedModelIds: string[]; defaultModelId?: string } + | undefined; + let settleModelScopeRequest: (() => void) | undefined; + const acceptModelScope = (value: Record): void => { + const scopedModelIds = Array.isArray(value.scopedModelIds) + ? value.scopedModelIds.filter( + (id): id is string => typeof id === "string", + ) + : []; + modelScope = { + scopedModelIds, + defaultModelId: + typeof value.defaultModelId === "string" + ? value.defaultModelId + : undefined, + }; + }; + const child = new PiRpcChild({ cwd, env: buildPiChildEnv({}), args: ["--mode", "rpc", "--no-session", "--extension", extensionPath], onEvent: () => {}, - onChannelMessage: () => {}, + onChannelMessage: (message) => { + if (message.kind === "model-scope") { + acceptModelScope(message); + return; + } + if ( + message.kind === "reply" && + message.id === "catalog-model-scope" && + typeof message.result === "object" && + message.result !== null + ) { + acceptModelScope(message.result as Record); + settleModelScopeRequest?.(); + settleModelScopeRequest = undefined; + } + }, onExit: () => {}, recordThreadId: null, }); - return child; + const ready = (async (): Promise> => { + const data = await child.requestOk({ type: "get_state" }); + await new Promise((resolveScope) => { + const timeout = setTimeout(resolveScope, 2_000); + timeout.unref?.(); + settleModelScopeRequest = () => { + clearTimeout(timeout); + resolveScope(); + }; + child.sendChannel({ + kind: "request", + id: "catalog-model-scope", + method: "model-scope", + }); + }); + return typeof data === "object" && data !== null + ? (data as Record) + : {}; + })(); + return { child, ready, getModelScope: () => modelScope }; }; - const fetchRaw = async (): Promise => { - const data = (await spawnChild().requestOk({ + const activeGeneration = (): CatalogChildGeneration => { + if (generation === null || generation.child.exited) { + generation = spawnGeneration(); + } + return generation; + }; + const fetchRawFrom = async ( + active: CatalogChildGeneration, + ): Promise => { + await active.ready; + const data = (await active.child.requestOk({ type: "get_available_models", })) as { models?: unknown[] } | undefined; touch(); @@ -111,16 +181,29 @@ async function spawnCatalog( typeof entry === "object" && entry !== null, ); }; - const probe = async (): Promise> => { - const data = await spawnChild().requestOk({ type: "get_state" }); - return typeof data === "object" && data !== null - ? (data as Record) - : {}; + const fetchGeneration = async (): Promise<{ + active: CatalogChildGeneration; + raw: PiRpcModel[]; + }> => { + const first = activeGeneration(); + try { + return { active: first, raw: await fetchRawFrom(first) }; + } catch (error) { + if (!(error instanceof PiRpcChildExitedError)) { + throw error; + } + const active = activeGeneration(); + return { active, raw: await fetchRawFrom(active) }; + } }; + const fetchRaw = async (): Promise => + (await fetchGeneration()).raw; + const probe = async (): Promise> => + activeGeneration().ready; await probe(); return { async listModels() { - const raw = await fetchRaw(); + const { active, raw } = await fetchGeneration(); const models: PiCatalogModel[] = []; for (const model of raw) { const catalogModel = toCatalogModel(model); @@ -132,15 +215,18 @@ async function spawnCatalog( ); } } - return buildPiAvailableModels({ models }); + const modelScope = active.getModelScope(); + return buildPiAvailableModels({ + models, + scopedModelIds: modelScope?.scopedModelIds, + preferredDefaultId: modelScope?.defaultModelId, + }); }, rawModels: fetchRaw, probe, async close() { - const activeChild = child; - if (activeChild === null) { - return; - } + const activeChild = generation?.child; + if (activeChild === undefined) return; activeChild.kill(); await activeChild.waitForExit(); }, diff --git a/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs b/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs index 01eadca9f1..3ecc77a3fc 100644 --- a/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs +++ b/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs @@ -228,9 +228,15 @@ const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); const extensionTools = new Map(); const extensionHandlers = new Map(); let activeTools = ["read", "bash", "edit", "write"]; +const scopedModel = + process.env.FAKE_PI_SCOPE_BY_SPAWN === "1" + ? MODELS[spawnIndex === 1 ? 0 : 1] + : undefined; const extensionContext = { cwd: process.cwd(), sessionManager: { getLeafId: () => leafId }, + model: scopedModel, + scopedModels: scopedModel ? [{ model: scopedModel }] : [], }; async function emitExtensionEvent(type, payload = {}) { @@ -432,6 +438,12 @@ async function handle(command) { return; case "get_available_models": respond(id, "get_available_models", { models: MODELS }); + if ( + process.env.FAKE_PI_EXIT_AFTER_FIRST_AVAILABLE === "1" && + spawnIndex === 1 + ) { + setTimeout(exit, 25); + } return; case "set_model": { const found = MODELS.find( diff --git a/plugins/provider-pi/src/model-list.test.ts b/plugins/provider-pi/src/model-list.test.ts index 116031e873..0ef27fe2d3 100644 --- a/plugins/provider-pi/src/model-list.test.ts +++ b/plugins/provider-pi/src/model-list.test.ts @@ -125,6 +125,97 @@ describe("pi model list", () => { ); }); + it("restricts and orders the picker using Pi's workspace model scope", () => { + const { models, selectedOnlyModels } = buildPiAvailableModels({ + models: [ + { + id: "claude-sonnet-5", + name: "Claude Sonnet 5", + provider: "anthropic", + reasoning: true, + input: ["text"], + supportedThinkingLevels: ["low", "medium", "high"], + }, + { + id: "gpt-5.4", + name: "GPT-5.4", + provider: "openai", + reasoning: true, + input: ["text"], + supportedThinkingLevels: ["low", "medium", "high"], + }, + ], + scopedModelIds: ["openai/gpt-5.4", "anthropic/claude-sonnet-5"], + preferredDefaultId: "openai/gpt-5.4", + }); + + expect(models.map((model) => model.id)).toEqual([ + "openai/gpt-5.4", + "anthropic/claude-sonnet-5", + ]); + expect(models.find((model) => model.isDefault)?.id).toBe("openai/gpt-5.4"); + expect(selectedOnlyModels).toHaveLength(0); + }); + + it("inherits Pi's saved default when no scope is configured", () => { + const { models } = buildPiAvailableModels({ + models: [ + { + id: "claude-sonnet-5", + name: "Claude Sonnet 5", + provider: "anthropic", + reasoning: true, + input: ["text"], + supportedThinkingLevels: ["low", "medium", "high"], + }, + { + id: "gpt-5.4", + name: "GPT-5.4", + provider: "openai", + reasoning: true, + input: ["text"], + supportedThinkingLevels: ["low", "medium", "high"], + }, + ], + preferredDefaultId: "anthropic/claude-sonnet-5", + }); + + expect(models.find((model) => model.isDefault)?.id).toBe( + "anthropic/claude-sonnet-5", + ); + }); + + it("keeps a dated model explicitly included by Pi's scope", () => { + const { models, selectedOnlyModels } = buildPiAvailableModels({ + models: [ + { + id: "claude-opus-4-8", + name: "Claude Opus 4.8", + provider: "anthropic", + reasoning: true, + input: ["text"], + supportedThinkingLevels: ["low", "medium", "high"], + }, + { + id: "claude-opus-4-8-20260115", + name: "Claude Opus 4.8 (2026-01-15)", + provider: "anthropic", + reasoning: true, + input: ["text"], + supportedThinkingLevels: ["low", "medium", "high"], + }, + ], + scopedModelIds: ["anthropic/claude-opus-4-8-20260115"], + preferredDefaultId: "anthropic/claude-opus-4-8-20260115", + }); + + expect(models.map((model) => model.id)).toEqual([ + "anthropic/claude-opus-4-8-20260115", + ]); + expect(models[0]?.isDefault).toBe(true); + expect(selectedOnlyModels).toHaveLength(0); + }); + it("reads the context window of the provider that served the message", () => { const resolveContextWindow = createPiModelContextWindowResolverFrom([ { diff --git a/plugins/provider-pi/src/model-list.ts b/plugins/provider-pi/src/model-list.ts index 83e3ae5e75..9357676bcf 100644 --- a/plugins/provider-pi/src/model-list.ts +++ b/plugins/provider-pi/src/model-list.ts @@ -24,6 +24,8 @@ export interface PiCatalogModel { interface BuildPiAvailableModelsArgs { models: readonly PiCatalogModel[]; + scopedModelIds?: readonly string[]; + preferredDefaultId?: string; } interface BuildPiAvailableModelsResult { @@ -65,18 +67,36 @@ function buildPiAvailableModel(model: PiCatalogModel): AvailableModel { export function buildPiAvailableModels( args: BuildPiAvailableModelsArgs, ): BuildPiAvailableModelsResult { + const scopedModelIds = args.scopedModelIds; + const scopedIds = + scopedModelIds && scopedModelIds.length > 0 + ? new Set(scopedModelIds) + : undefined; + const sourceModels = scopedIds + ? [...scopedIds].flatMap((id) => { + const match = args.models.find( + (model) => toCanonicalPiModelId(model.provider, model.id) === id, + ); + return match ? [match] : []; + }) + : args.models; + const models: AvailableModel[] = []; const selectedOnlyModels: AvailableModel[] = []; - for (const model of args.models) { + for (const model of sourceModels) { const built = buildPiAvailableModel(model); - if (isModelAlias(model.id)) { + if (isModelAlias(model.id) || scopedIds?.has(built.id)) { models.push(built); } else { selectedOnlyModels.push(built); } } - const defaultId = resolveDefaultPiModelId(models); + const defaultId = + (args.preferredDefaultId && + models.some((model) => model.id === args.preferredDefaultId) + ? args.preferredDefaultId + : undefined) ?? resolveDefaultPiModelId(models); return { models: models.map((model) => model.id === defaultId ? { ...model, isDefault: true } : model,