diff --git a/apps/web/src/components/settings/SettingsPanels.tsx b/apps/web/src/components/settings/SettingsPanels.tsx index c4e3d27a07fa7..7ea4d8ddc1498 100644 --- a/apps/web/src/components/settings/SettingsPanels.tsx +++ b/apps/web/src/components/settings/SettingsPanels.tsx @@ -59,7 +59,7 @@ import { ThreadTitleInstructionsSettings } from "./ThreadTitleInstructionsSettin import { SnoozePresetsRow } from "./SnoozePresetsSettings"; import { useDesktopUpdateState } from "../../state/desktopUpdate"; import { - getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, resolveAppModelSelectionState, selectsPlanAgent, } from "../../modelSelection"; @@ -1516,7 +1516,7 @@ export function ThreadsSettingsPanel() { ); const textGenProvider: ProviderDriverKind = textGenInstanceEntry?.driverKind ?? DEFAULT_DRIVER_KIND; - const textGenerationModelOptionsByInstance = getCustomModelOptionsByInstance( + const textGenerationModelOptionsByInstance = getTextGenerationModelOptionsByInstance( settings, textGenerationProviders, textGenInstanceId, @@ -1529,6 +1529,7 @@ export function ThreadsSettingsPanel() { const textGenerationModelDisabledReason = useScopedModelDisabledReason( settings, textGenerationModelInstanceEntries, + true, ); const mixedAddProjectBaseDirectory = useScopedSettingsMixed(["addProjectBaseDirectory"]); const mixedTextGenerationModel = useScopedSettingsMixed(["textGenerationModelSelection"]); diff --git a/apps/web/src/components/settings/SourceControlWritingSettings.browser.test.tsx b/apps/web/src/components/settings/SourceControlWritingSettings.browser.test.tsx index 5a44beb50ac82..d04e929f2f22e 100644 --- a/apps/web/src/components/settings/SourceControlWritingSettings.browser.test.tsx +++ b/apps/web/src/components/settings/SourceControlWritingSettings.browser.test.tsx @@ -1,5 +1,13 @@ import { render } from "../../../test/render"; import { DEFAULT_UNIFIED_SETTINGS } from "@t3tools/contracts/settings"; +import { + EnvironmentId, + ProviderDriverKind, + ProviderInstanceId, + type ServerProvider, +} from "@t3tools/contracts"; +import { useScopedModelDisabledReason } from "./useScopedModelAvailability"; +import { deriveProviderInstanceEntries } from "../../providerInstances"; import { StrictMode, type ReactNode } from "react"; import { describe, expect, it, vi } from "vite-plus/test"; import { page, userEvent } from "vite-plus/test/browser"; @@ -10,6 +18,8 @@ import type { ScopedSettingsPatch } from "./scopedSettings"; type WritingStyle = typeof DEFAULT_UNIFIED_SETTINGS.sourceControlWritingStyle; const state = vi.hoisted(() => ({ styles: [] as WritingStyle[], + providerInstances: {} as typeof DEFAULT_UNIFIED_SETTINGS.providerInstances, + providers: new Map(), updateSettings: vi.fn<(patch: ScopedSettingsPatch) => void>(), })); @@ -22,6 +32,7 @@ vi.mock(import("./useScopedSettings"), async (original) => ({ useScopedSettings: (() => ({ ...DEFAULT_UNIFIED_SETTINGS, sourceControlWritingStyle: state.styles[0], + providerInstances: state.providerInstances, })) as never, useScopedSettingsMixed: (keys) => !keys.includes("reviewGuidesEnabled") && @@ -32,17 +43,33 @@ vi.mock(import("./SettingsScopeContext"), async (original) => ({ ...(await original()), useSettingsScope: (() => ({ scope: { kind: "all", environmentIds: [] }, - environment: null, - connectedEnvironments: [], - targets: state.styles.map((style) => ({ - settings: { ...DEFAULT_UNIFIED_SETTINGS, sourceControlWritingStyle: style }, + environment: { environmentId: EnvironmentId.make("representative"), serverConfig: {} }, + connectedEnvironments: [...state.providers.keys()].map((environmentId) => ({ + environmentId, + serverConfig: {}, + })), + representative: { environmentId: EnvironmentId.make("representative") }, + targets: state.styles.map((style, index) => ({ + environmentId: EnvironmentId.make(`environment-${index}`), + settings: { + ...DEFAULT_UNIFIED_SETTINGS, + providerInstances: state.providerInstances, + sourceControlWritingStyle: style, + }, })), target: null, })) as never, })); vi.mock(import("../../state/environments"), async (original) => ({ ...(await original()), - useEnvironmentProviders: () => new Map(), + useEnvironmentProviders: () => state.providers, + useEnvironments: (() => ({ + environments: [...state.providers.keys()].map((environmentId) => ({ + environmentId, + label: environmentId, + serverConfig: {}, + })), + })) as never, })); vi.mock(import("./settingsLayout"), async (original) => ({ ...(await original()), @@ -69,6 +96,8 @@ const content = () => ( ); async function setup() { + state.providers = new Map(); + state.providerInstances = {}; state.styles = [ { ...DEFAULT_UNIFIED_SETTINGS.sourceControlWritingStyle, @@ -112,3 +141,85 @@ describe("mixed source control instructions", () => { ).toEqual(initialInstructions); }); }); + +it("rejects a scoped model when another environment requires extensions for it", async () => { + state.updateSettings.mockReset(); + state.styles = [DEFAULT_UNIFIED_SETTINGS.sourceControlWritingStyle]; + const instanceId = ProviderInstanceId.make("pi"); + state.providerInstances = { + [instanceId]: { driver: ProviderDriverKind.make("pi"), enabled: true }, + }; + const provider: ServerProvider = { + instanceId, + driver: ProviderDriverKind.make("pi"), + enabled: true, + installed: true, + version: null, + status: "ready", + auth: { status: "authenticated" }, + checkedAt: "2026-01-01T00:00:00.000Z", + slashCommands: [], + skills: [], + models: [{ slug: "demo/model", name: "Demo", isCustom: false, capabilities: {} }], + }; + state.providers = new Map([ + [EnvironmentId.make("representative"), [provider]], + [ + EnvironmentId.make("environment-0"), + [ + { + ...provider, + models: provider.models.map((model) => ({ ...model, supportsTextGeneration: false })), + }, + ], + ], + ]); + function Availability() { + const backgroundReason = useScopedModelDisabledReason( + { ...DEFAULT_UNIFIED_SETTINGS, providerInstances: state.providerInstances }, + deriveProviderInstanceEntries([provider]), + true, + ); + const threadReason = useScopedModelDisabledReason( + { ...DEFAULT_UNIFIED_SETTINGS, providerInstances: state.providerInstances }, + deriveProviderInstanceEntries([provider]), + ); + return ( + <> +

{backgroundReason(instanceId, "demo/model") ?? "Available"}

+

{threadReason(instanceId, "demo/model") ?? "Available for threads"}

+ + ); + } + const availability = () => ( + <> + + + + ); + await using screen = await render(availability()); + await userEvent.click(screen.getByRole("button", { name: "Advanced" })); + await userEvent.click( + screen.getByRole("switch", { name: "Use a separate source control writer model" }), + ); + expect(state.updateSettings).not.toHaveBeenCalled(); + await expect.element(screen.getByText("Available for threads")).toBeVisible(); + await expect + .element( + screen.getByText( + "This model is unavailable on environment-0. Select that environment to choose its model separately.", + ), + ) + .toBeVisible(); + state.providers = new Map([ + [EnvironmentId.make("representative"), [provider]], + [EnvironmentId.make("environment-0"), [provider]], + ]); + await screen.rerender(availability()); + await userEvent.click( + screen.getByRole("switch", { name: "Use a separate source control writer model" }), + ); + expect(state.updateSettings).toHaveBeenCalledWith({ + sourceControlWriterModelSelection: { instanceId, model: "demo/model" }, + }); +}); diff --git a/apps/web/src/components/settings/SourceControlWritingSettings.tsx b/apps/web/src/components/settings/SourceControlWritingSettings.tsx index dd3618902f8b5..205d6aa1d0ca1 100644 --- a/apps/web/src/components/settings/SourceControlWritingSettings.tsx +++ b/apps/web/src/components/settings/SourceControlWritingSettings.tsx @@ -21,7 +21,7 @@ import { useScopedModelDisabledReason } from "./useScopedModelAvailability"; import { useSettingsScope } from "./SettingsScopeContext"; import { providerInstanceEntriesFor } from "../../providerInstances"; import { - getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, resolveAppModelSelectionState, selectsPlanAgent, } from "../../modelSelection"; @@ -139,13 +139,13 @@ export function SourceControlWritingSettingsSection() { (entry) => entry.instanceId === defaultModelSelection.instanceId && entry.enabled && entry.isAvailable, ); - const modelOptionsByInstance = getCustomModelOptionsByInstance( + const modelOptionsByInstance = getTextGenerationModelOptionsByInstance( settings, projectEffectiveTextGenerationProviders, activeSelection.instanceId, activeSelection.model, ); - const writerModelDisabledReason = useScopedModelDisabledReason(settings, instanceEntries); + const writerModelDisabledReason = useScopedModelDisabledReason(settings, instanceEntries, true); const writerInstanceEntry = instanceEntries.find( (entry) => entry.instanceId === activeSelection.instanceId, ); @@ -161,7 +161,7 @@ export function SourceControlWritingSettingsSection() { }, textGenerationProviders, ); - const reviewGuideModelOptionsByInstance = getCustomModelOptionsByInstance( + const reviewGuideModelOptionsByInstance = getTextGenerationModelOptionsByInstance( settings, textGenerationProviders, reviewGuideSelection.instanceId, @@ -577,7 +577,21 @@ export function SourceControlWritingSettingsSection() { + onCheckedChange={(checked) => { + if (checked) { + const reason = writerModelDisabledReason( + defaultModelSelection.instanceId, + defaultModelSelection.model, + ); + if (reason) { + toastManager.add({ + type: "error", + title: "Writer model not saved", + description: reason, + }); + return; + } + } updateSettings({ sourceControlWriterModelSelection: checked ? createModelSelection( @@ -586,8 +600,8 @@ export function SourceControlWritingSettingsSection() { defaultModelSelection.options, ) : null, - }) - } + }); + }} aria-label="Use a separate source control writer model" /> diff --git a/apps/web/src/components/settings/useScopedModelAvailability.ts b/apps/web/src/components/settings/useScopedModelAvailability.ts index abf1e1e1afc03..9fa4eeb938b1b 100644 --- a/apps/web/src/components/settings/useScopedModelAvailability.ts +++ b/apps/web/src/components/settings/useScopedModelAvailability.ts @@ -1,7 +1,10 @@ import type { ProviderInstanceId, UnifiedSettings } from "@t3tools/contracts"; import { mergeEnvironmentSettings } from "../../hooks/useSettings"; -import { getCustomModelOptionsByInstance } from "../../modelSelection"; +import { + getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, +} from "../../modelSelection"; import { applyProviderInstanceSettings, deriveProviderInstanceEntries, @@ -13,6 +16,7 @@ import { useSettingsScope } from "./SettingsScopeContext"; export function useScopedModelDisabledReason( settings: UnifiedSettings, entries: readonly ProviderInstanceEntry[], + textGeneration = false, ) { const { targets } = useSettingsScope(); const { environments } = useEnvironments(); @@ -30,7 +34,10 @@ export function useScopedModelDisabledReason( candidate.settings, candidate.projectId, ).find((option) => option.instanceId === instanceId); - const options = getCustomModelOptionsByInstance( + const optionsForInstance = textGeneration + ? getTextGenerationModelOptionsByInstance + : getCustomModelOptionsByInstance; + const options = optionsForInstance( mergeEnvironmentSettings(candidate.settings, settings), providers, ).get(instanceId); diff --git a/apps/web/src/modelSelection.test.ts b/apps/web/src/modelSelection.test.ts index 05a7e76b6b25b..00b333f468991 100644 --- a/apps/web/src/modelSelection.test.ts +++ b/apps/web/src/modelSelection.test.ts @@ -13,6 +13,8 @@ import { getComposerProviderState } from "./components/chat/composerProviderStat import { deriveProviderInstanceEntries, NO_PROVIDER_MODEL_SELECTION } from "./providerInstances"; import { getAppModelOptionsForInstance, + getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, resolveAppModelSelectionForInstance, resolveAppModelSelectionState, resolveEffectiveDefaultModelSelection, @@ -670,6 +672,80 @@ describe("instance-scoped model selection", () => { NO_PROVIDER_MODEL_SELECTION, ); }); + + it("does not select a model that the provider cannot use for text generation", () => { + const codex = provider({ instanceId: "codex", models: ["gpt-5.6-sol", "extension-model"] }); + const providers = [ + { + ...codex, + models: codex.models.map((m) => + m.slug === "extension-model" ? { ...m, supportsTextGeneration: false } : m, + ), + }, + ]; + const settings: UnifiedSettings = { + ...settingsWithProviderInstances(), + textGenerationModelSelection: createModelSelection(codex.instanceId, "extension-model"), + }; + + expect(resolveAppModelSelectionState(settings, providers)).toEqual( + createModelSelection(codex.instanceId, "gpt-5.6-sol"), + ); + }); + + it("excludes customized extension models while preserving unknown custom and thread models", () => { + const pi = provider({ + provider: ProviderDriverKind.make("pi"), + instanceId: "pi", + models: ["default", "demo/demo-model"], + }); + const providers = [ + { + ...pi, + models: pi.models.map((model) => + model.slug === "demo/demo-model" ? { ...model, supportsTextGeneration: false } : model, + ), + }, + ]; + const settings: UnifiedSettings = { + ...DEFAULT_UNIFIED_SETTINGS, + providerInstances: { + [pi.instanceId]: { + driver: ProviderDriverKind.make("pi"), + config: { + customModels: [ + { slug: "demo/demo-model", name: "My extension model" }, + { slug: "unknown/model", name: "My custom model" }, + ], + }, + }, + }, + textGenerationModelSelection: createModelSelection(pi.instanceId, "demo/demo-model"), + }; + const supported = providers.filter((provider) => provider.supportsTextGeneration !== false); + expect( + getTextGenerationModelOptionsByInstance(settings, supported, pi.instanceId, "demo/demo-model") + .get(pi.instanceId) + ?.map((option) => option.slug), + ).toEqual(["default", "unknown/model"]); + expect(resolveAppModelSelectionState(settings, supported)).toEqual( + createModelSelection(pi.instanceId, "default"), + ); + expect( + resolveAppModelSelectionState( + { + ...settings, + textGenerationModelSelection: createModelSelection(pi.instanceId, "unknown/model"), + }, + supported, + ).model, + ).toBe("unknown/model"); + expect( + getCustomModelOptionsByInstance(settings, providers) + .get(pi.instanceId) + ?.map((option) => option.slug), + ).toEqual(["default", "demo/demo-model", "unknown/model"]); + }); }); describe("resolveAppModelSelectionState with the opencode plan agent", () => { diff --git a/apps/web/src/modelSelection.ts b/apps/web/src/modelSelection.ts index db52028c397e0..d1aac480846a9 100644 --- a/apps/web/src/modelSelection.ts +++ b/apps/web/src/modelSelection.ts @@ -328,6 +328,42 @@ export function getCustomModelOptionsByInstance( return out; } +function getTextGenerationModelOptionsForInstance( + settings: UnifiedSettings, + entry: ProviderInstanceEntry, + selectedModel?: string | null, +): AppModelOption[] { + const unsupported = new Set( + entry.models + .filter((model) => model.supportsTextGeneration === false) + .map((model) => model.slug), + ); + // Filter after custom entries are rebuilt so a display-name override cannot restore a model. + return getAppModelOptionsForInstance(settings, entry, selectedModel).filter( + (option) => !unsupported.has(option.slug), + ); +} + +export function getTextGenerationModelOptionsByInstance( + settings: UnifiedSettings, + providers: ReadonlyArray, + selectedInstanceId?: ProviderInstanceId | null, + selectedModel?: string | null, +): ReadonlyMap> { + return new Map( + deriveProviderInstanceEntries( + providers.filter((provider) => provider.supportsTextGeneration !== false), + ).map((entry) => [ + entry.instanceId, + getTextGenerationModelOptionsForInstance( + settings, + entry, + entry.instanceId === selectedInstanceId ? selectedModel : null, + ), + ]), + ); +} + export function selectsPlanAgent(options: ModelSelection["options"]): boolean { return options?.some((option) => option.id === "agent" && option.value === "plan") ?? false; } @@ -352,7 +388,7 @@ export function resolveAppModelSelectionState( if (entry) { const selectedModel = selectedEntry ? selection.model : null; // Use the cheap generation default; the catalog default is the expensive chat flagship. - const options = getAppModelOptionsForInstance(settings, entry); + const options = getTextGenerationModelOptionsForInstance(settings, entry); const model = resolveSelectableModel(entry.driverKind, selectedModel, options) ?? resolveSelectableModel( @@ -360,13 +396,9 @@ export function resolveAppModelSelectionState( DEFAULT_TEXT_GENERATION_MODEL_BY_PROVIDER[entry.driverKind], options, ) ?? - resolveAppModelSelectionForInstance( - entry.instanceId, - settings, - supportedProviders, - selectedModel, - ) ?? - entry.models[0]?.slug ?? + options.find((option) => option.isDefault)?.slug ?? + options[0]?.slug ?? + entry.models.find((model) => model.supportsTextGeneration !== false)?.slug ?? DEFAULT_TEXT_GENERATION_MODEL_BY_PROVIDER[entry.driverKind]; if (!model) { return createModelSelection(entry.instanceId, "", []); diff --git a/packages/contracts/src/server.ts b/packages/contracts/src/server.ts index ab79a9dea60f5..9a0f6878991df 100644 --- a/packages/contracts/src/server.ts +++ b/packages/contracts/src/server.ts @@ -82,6 +82,9 @@ export const ServerProviderModel = Schema.Struct({ isCustom: Schema.Boolean, isDefault: Schema.optional(Schema.Boolean), isLegacy: Schema.optional(Schema.Boolean), + // False when the provider can run this model in sessions but not for + // application text generation. Absent means supported. + supportsTextGeneration: Schema.optional(Schema.Boolean), capabilities: Schema.NullOr(ModelCapabilities), }); export type ServerProviderModel = typeof ServerProviderModel.Type; diff --git a/packages/provider-pi/src/server/status.test.ts b/packages/provider-pi/src/server/status.test.ts index b8559d1ab7d8a..083d1a70003a7 100644 --- a/packages/provider-pi/src/server/status.test.ts +++ b/packages/provider-pi/src/server/status.test.ts @@ -2,6 +2,9 @@ import * as NodeServices from "@effect/platform-node/NodeServices"; import { assert, describe, it } from "@effect/vitest"; import * as Cause from "effect/Cause"; import * as Effect from "effect/Effect"; +import * as Deferred from "effect/Deferred"; +import * as Fiber from "effect/Fiber"; +import * as TestClock from "effect/testing/TestClock"; import * as Queue from "effect/Queue"; import * as Schema from "effect/Schema"; import * as Sink from "effect/Sink"; @@ -52,14 +55,10 @@ function piProbeSpawner(version: string) { }); } -/** Fails the test if status runs `pi auth check` without `authCheck` or with other arguments. */ -const piDiscoverySpawner = ( +const piRpcProcess = ( inventory: unknown, - options: { - readonly version?: string; - readonly state?: unknown; - readonly authCheck?: string | undefined; - } = {}, + onModelRequest: Effect.Effect = Effect.void, + state: unknown = {}, ) => Effect.gen(function* () { const stdout = yield* Queue.unbounded(); @@ -74,13 +73,17 @@ const piDiscoverySpawner = ( stdinBuffer = stdinBuffer.slice(newline + 1); if (line.length === 0) continue; const request = decodeRpcRequest(line); + if (request.type === "get_available_models") { + yield* onModelRequest; + if (inventory === null) continue; + } const data = request.type === "get_available_models" ? inventory : request.type === "get_commands" ? { commands: [] } : request.type === "get_state" - ? (options.state ?? {}) + ? state : {}; yield* Queue.offer( stdout, @@ -91,7 +94,28 @@ const piDiscoverySpawner = ( } }), ); - return ChildProcessSpawner.make((command) => { + return ChildProcessSpawner.makeHandle({ + ...processHandle({}), + exitCode: Effect.never, + isRunning: Effect.succeed(true), + stdin, + stdout: Stream.fromQueue(stdout), + }); + }); + +// Text generation lists models with --no-extensions, so its inventory can differ. +const piDiscoverySpawner = ( + inventory: unknown, + extensionFreeInventory: unknown = inventory, + onExtensionFreeModelRequest: Effect.Effect = Effect.void, + options: { + readonly version?: string; + readonly state?: unknown; + readonly authCheck?: string | undefined; + } = {}, +) => + Effect.succeed( + ChildProcessSpawner.make((command) => { const args = ChildProcess.isStandardCommand(command) ? command.args : []; if (args.includes("--version")) { return Effect.succeed(processHandle({ stdout: `pi ${options.version ?? "0.84.3"}\n` })); @@ -102,17 +126,13 @@ const piDiscoverySpawner = ( ? Effect.die(`Unexpected pi ${args.join(" ")}`) : Effect.succeed(processHandle({ stdout: options.authCheck, exitCode: 1 })); } - return Effect.succeed( - ChildProcessSpawner.makeHandle({ - ...processHandle({}), - exitCode: Effect.never, - isRunning: Effect.succeed(true), - stdin, - stdout: Stream.fromQueue(stdout), - }), + return piRpcProcess( + args.includes("--no-extensions") ? extensionFreeInventory : inventory, + args.includes("--no-extensions") ? onExtensionFreeModelRequest : Effect.void, + options.state, ); - }); - }); + }), + ); const settings = { enabled: true, @@ -216,6 +236,8 @@ describe("PiProvider", () => { { provider: "router", id: "auto" }, ], }, + undefined, + Effect.void, { version: row.version, state: { model: row.activeModel }, authCheck: row.authCheck }, ); const snapshot = yield* checkPiProviderStatus({ @@ -229,6 +251,62 @@ describe("PiProvider", () => { }).pipe(Effect.provide(NodeServices.layer)), ); + it.effect("marks models that only load with Pi extensions as unusable for text generation", () => + Effect.gen(function* () { + const spawner = yield* piDiscoverySpawner( + { + models: [ + { provider: "anthropic", id: "claude-sonnet-4" }, + { provider: "demo", id: "demo-model" }, + ], + }, + { models: [{ provider: "anthropic", id: "claude-sonnet-4" }] }, + ); + const snapshot = yield* checkPiProviderStatus(settings).pipe( + Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner), + ); + assert.deepEqual( + snapshot.models.map((model) => [model.slug, model.supportsTextGeneration]), + [ + ["default", undefined], + ["anthropic/claude-sonnet-4", undefined], + ["demo/demo-model", false], + ], + ); + }).pipe(Effect.provide(NodeServices.layer)), + ); + + it.effect("preserves normal discovery when the extension-free model listing times out", () => + Effect.gen(function* () { + const modelRequested = yield* Deferred.make(); + const spawner = yield* piDiscoverySpawner( + { models: [{ provider: "demo", id: "demo-model" }] }, + null, + Deferred.succeed(modelRequested, undefined).pipe(Effect.asVoid), + ); + const fiber = yield* checkPiProviderStatus(settings).pipe( + Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner), + Effect.forkChild, + ); + yield* Deferred.await(modelRequested); + yield* TestClock.adjust("5 seconds"); + const snapshot = yield* Fiber.join(fiber); + assert.equal(snapshot.status, "ready"); + assert.equal(snapshot.auth.status, "authenticated"); + assert.deepEqual( + snapshot.models.map((model) => [model.slug, model.supportsTextGeneration]), + [ + ["default", undefined], + ["demo/demo-model", undefined], + ], + ); + assert.include( + snapshot.slashCommands.map((command) => command.name), + "compact", + ); + }).pipe(Effect.provide(NodeServices.layer)), + ); + // regression: #18498 a malformed model inventory must not restore stale Pi models it.effect.each([ ["missing models", {}], diff --git a/packages/provider-pi/src/server/status.ts b/packages/provider-pi/src/server/status.ts index a180ab9c41a3b..6065fd1caffdd 100644 --- a/packages/provider-pi/src/server/status.ts +++ b/packages/provider-pi/src/server/status.ts @@ -11,6 +11,7 @@ import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Stream from "effect/Stream"; import * as Exit from "effect/Exit"; +import * as Fiber from "effect/Fiber"; import * as Option from "effect/Option"; import * as Result from "effect/Result"; import * as Schema from "effect/Schema"; @@ -58,6 +59,7 @@ const PI_PRESENTATION = { const VERSION_PROBE_TIMEOUT_MS = 4_000; const PI_RPC_DISCOVERY_TIMEOUT_MS = 30_000; +const PI_TEXT_GENERATION_DISCOVERY_TIMEOUT_MS = 5_000; /** * get_entries arrived in 0.80.3 and agent_settled landed in source at 0.80.4. * Version 0.80.5 was the first published package containing both hooks. @@ -142,6 +144,7 @@ const makePiDiscoveryConnection = Effect.fnUntraced(function* ( environment: NodeJS.ProcessEnv, launchArgs: ReadonlyArray, cwd?: string, + textGenerationLaunch = false, ) { const launch = buildPiRpcLaunch({ launchArgs, @@ -149,6 +152,8 @@ const makePiDiscoveryConnection = Effect.fnUntraced(function* ( mcpSession: undefined, extensionPath: undefined, ephemeral: true, + disableExtensions: textGenerationLaunch, + disableTools: textGenerationLaunch, }); const connection = yield* makePiRpcConnection({ command: piSettings.binaryPath || "pi", @@ -164,12 +169,49 @@ const makePiDiscoveryConnection = Effect.fnUntraced(function* ( return connection; }); +// Text generation runs Pi with --no-extensions, so it cannot load models that an +// extension or package registers. Listing models with that launch tells them apart. +const discoverPiTextGenerationModelSlugs = Effect.fn( + "PiProviderStatus.discoverPiTextGenerationModelSlugs", +)(function* ( + piSettings: PiSettings, + environment: NodeJS.ProcessEnv, + launchArgs: ReadonlyArray, + cwd?: string, +) { + const connection = yield* makePiDiscoveryConnection( + piSettings, + environment, + launchArgs, + cwd, + true, + ); + const modelsData = yield* connection.request({ type: "get_available_models" }); + const models = yield* Effect.fromResult(parseDiscoveredModels(modelsData, undefined)); + return new Set(models.map((model) => model.slug)); +}, Effect.scoped); + const discoverPiViaRpc = Effect.fn("PiProviderStatus.discoverPiViaRpc")(function* ( piSettings: PiSettings, environment: NodeJS.ProcessEnv, launchArgs: ReadonlyArray, cwd?: string, ) { + // This optional probe must finish before the normal discovery deadline so + // it cannot discard an otherwise usable catalog and authentication state. + const textGenerationModelSlugsFiber = yield* discoverPiTextGenerationModelSlugs( + piSettings, + environment, + launchArgs, + cwd, + ).pipe( + Effect.orElseSucceed(() => undefined), + Effect.timeoutOrElse({ + duration: PI_TEXT_GENERATION_DISCOVERY_TIMEOUT_MS, + orElse: () => Effect.void, + }), + Effect.forkScoped, + ); const connection = yield* makePiDiscoveryConnection(piSettings, environment, launchArgs, cwd); const stateData = yield* connection.request({ type: "get_state" }); const modelsData = yield* connection.request({ type: "get_available_models" }); @@ -178,13 +220,21 @@ const discoverPiViaRpc = Effect.fn("PiProviderStatus.discoverPiViaRpc")(function .pipe(Effect.orElseSucceed(() => undefined)); const discoveredModels = yield* Effect.fromResult(parseDiscoveredModels(modelsData, stateData)); const { slashCommands, skills } = parsePiDiscoveredCommands(commandsData); + const textGenerationModelSlugs = yield* Fiber.join(textGenerationModelSlugsFiber); // Pi lists only models of providers it holds credentials for, including extension, virtual and // --api-key ones. Turns need only the provider's credentials, so uncatalogued model ids still run. const activeModel = recordField(stateData, "model"); const activeProvider = recordString(activeModel, "provider"); const activeModelId = recordString(activeModel, "id"); return { - models: discoveredModels, + models: + textGenerationModelSlugs === undefined + ? discoveredModels + : discoveredModels.map((model) => + textGenerationModelSlugs.has(model.slug) + ? model + : { ...model, supportsTextGeneration: false }, + ), slashCommands: withPiBuiltinSlashCommands(slashCommands), skills, authenticated: discoveredModels.length > 0,