diff --git a/apps/server/src/textGeneration/CodexTextGeneration.test.ts b/apps/server/src/textGeneration/CodexTextGeneration.test.ts index 389a6458677e7..ae5c32bf57432 100644 --- a/apps/server/src/textGeneration/CodexTextGeneration.test.ts +++ b/apps/server/src/textGeneration/CodexTextGeneration.test.ts @@ -546,6 +546,8 @@ const PREFLIGHT_PRIMARY_STATES = [ "not installed, fallback disabled for project", "signed out", "signed out, fallback failed", + "cannot generate text", + "model cannot generate text", ] as const; it.effect.each(PREFLIGHT_PRIMARY_STATES)("text generation preflight: %s", (primaryState) => @@ -590,6 +592,22 @@ it.effect.each(PREFLIGHT_PRIMARY_STATES)("text generation preflight: %s", (prima ...(primaryState === "signed out, fallback failed" && instanceId === fallback ? { status: "error", auth: { status: "unknown" } } : {}), + ...(primaryState === "cannot generate text" && instanceId === primary + ? { supportsTextGeneration: false } + : {}), + ...(primaryState === "model cannot generate text" && instanceId === primary + ? { + models: [ + { + slug: selection.model, + name: selection.model, + isCustom: false, + supportsTextGeneration: false, + capabilities: null, + }, + ], + } + : {}), ...(primaryState === "exhausted" && instanceId === primary ? { usageLimits: { diff --git a/apps/server/src/textGeneration/TextGeneration.ts b/apps/server/src/textGeneration/TextGeneration.ts index b89e264db8d57..818afb4fae553 100644 --- a/apps/server/src/textGeneration/TextGeneration.ts +++ b/apps/server/src/textGeneration/TextGeneration.ts @@ -108,8 +108,11 @@ export const make = Effect.gen(function* () { const unusable = textGeneratorCannotRun || (selected?.enabled === true && - selected.status === "error" && - (!selected.installed || selected.auth.status === "unauthenticated")); + (!isProviderTextGenerationCapable(selected) || + selected.models.find((model) => model.slug === selection.model) + ?.supportsTextGeneration === false || + (selected.status === "error" && + (!selected.installed || selected.auth.status === "unauthenticated")))); if (unusable || exhaustedUsageWindow(selected, now, selection.model)) { const settings = yield* serverSettings.getSettings.pipe( Effect.mapError( diff --git a/apps/web/src/components/settings/SettingsPanels.tsx b/apps/web/src/components/settings/SettingsPanels.tsx index 82d6a35df6816..b20295847f934 100644 --- a/apps/web/src/components/settings/SettingsPanels.tsx +++ b/apps/web/src/components/settings/SettingsPanels.tsx @@ -67,7 +67,7 @@ import { SnoozePresetsRow } from "./SnoozePresetsSettings"; import { useDesktopUpdateState } from "../../state/desktopUpdate"; import { managedTelemetryOffMessage, usePrimaryManagedPolicy } from "../../state/managedPolicy"; import { - getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, resolveAppModelSelectionState, selectsPlanAgent, } from "../../modelSelection"; @@ -1743,7 +1743,7 @@ export function ThreadsSettingsPanel() { ); const textGenProvider: ProviderDriverKind = textGenInstanceEntry?.driverKind ?? DEFAULT_DRIVER_KIND; - const textGenerationModelOptionsByInstance = getCustomModelOptionsByInstance( + const textGenerationModelOptionsByInstance = getTextGenerationModelOptionsByInstance( settings, textGenerationProviders, textGenInstanceId, @@ -1756,6 +1756,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.tsx b/apps/web/src/components/settings/SourceControlWritingSettings.tsx index 9d65ffc83ea08..26cf6eb3dcbb2 100644 --- a/apps/web/src/components/settings/SourceControlWritingSettings.tsx +++ b/apps/web/src/components/settings/SourceControlWritingSettings.tsx @@ -23,7 +23,7 @@ import { useSettingsScope } from "./SettingsScopeContext"; import { providerInstanceEntriesFor } from "../../providerInstances"; import { useTextGenerationTargets } from "../../textGenerators"; import { - getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, resolveAppModelSelectionState, selectsPlanAgent, } from "../../modelSelection"; @@ -151,13 +151,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, writerEntries); + const writerModelDisabledReason = useScopedModelDisabledReason(settings, writerEntries, true); const writerInstanceEntry = writerEntries.find( (entry) => entry.instanceId === activeSelection.instanceId, ); @@ -179,7 +179,7 @@ export function SourceControlWritingSettingsSection() { }, projectAgentProviders, ); - const reviewGuideModelOptionsByInstance = getCustomModelOptionsByInstance( + const reviewGuideModelOptionsByInstance = getTextGenerationModelOptionsByInstance( settings, projectAgentProviders, reviewGuideSelection.instanceId, @@ -595,7 +595,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( @@ -604,8 +618,8 @@ export function SourceControlWritingSettingsSection() { defaultModelSelection.options, ) : null, - }) - } + }); + }} aria-label="Use a separate source control writer model" /> diff --git a/apps/web/src/components/settings/TextGeneratorsSettings.tsx b/apps/web/src/components/settings/TextGeneratorsSettings.tsx index 4dd163ae0e7ee..ce0e729309bf2 100644 --- a/apps/web/src/components/settings/TextGeneratorsSettings.tsx +++ b/apps/web/src/components/settings/TextGeneratorsSettings.tsx @@ -28,7 +28,7 @@ import { textGeneratorBrand, textGeneratorDisplayName, } from "../../textGenerators"; -import { getCustomModelOptionsByInstance } from "../../modelSelection"; +import { getTextGenerationModelOptionsByInstance } from "../../modelSelection"; import { ProviderInstanceIcon } from "../chat/ProviderInstanceIcon"; import { ProviderModelPicker } from "../chat/ProviderModelPicker"; import { @@ -563,7 +563,7 @@ function TextGenerationTaskPicker({ model={selection.model} lockedProvider={null} instanceEntries={targets.entries} - modelOptionsByInstance={getCustomModelOptionsByInstance( + modelOptionsByInstance={getTextGenerationModelOptionsByInstance( settings, targets.providers, selection.instanceId, diff --git a/apps/web/src/components/settings/VoiceSettingsPanel.tsx b/apps/web/src/components/settings/VoiceSettingsPanel.tsx index 2875a6edbe355..22ae59e55d31e 100644 --- a/apps/web/src/components/settings/VoiceSettingsPanel.tsx +++ b/apps/web/src/components/settings/VoiceSettingsPanel.tsx @@ -72,7 +72,7 @@ import { import { updateClientSettings, useClientSettings } from "../../hooks/useSettings"; import { runEnvironmentSpeech } from "../../speech/runEnvironmentSpeech"; import { - getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, resolveAppModelSelectionState, } from "../../modelSelection"; import { useTextGenerationTargets } from "../../textGenerators"; @@ -531,14 +531,14 @@ function CleanupModelRow() { }, textGenerationProviders, ); - const modelOptionsByInstance = getCustomModelOptionsByInstance( + const modelOptionsByInstance = getTextGenerationModelOptionsByInstance( settings, textGenerationProviders, selection.instanceId, selection.model, ); const instanceEntry = instanceEntries.find((entry) => entry.instanceId === selection.instanceId); - const modelDisabledReason = useScopedModelDisabledReason(settings, instanceEntries); + const modelDisabledReason = useScopedModelDisabledReason(settings, instanceEntries, true); const setModel = (next: ReturnType) => { const reason = modelDisabledReason(next.instanceId, next.model); if (reason) { diff --git a/apps/web/src/components/settings/useScopedModelAvailability.ts b/apps/web/src/components/settings/useScopedModelAvailability.ts index 2471825cb3d16..a5d89813cae46 100644 --- a/apps/web/src/components/settings/useScopedModelAvailability.ts +++ b/apps/web/src/components/settings/useScopedModelAvailability.ts @@ -5,7 +5,10 @@ import { } from "@t3tools/contracts"; import { mergeEnvironmentSettings } from "../../hooks/useSettings"; -import { getCustomModelOptionsByInstance } from "../../modelSelection"; +import { + getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, +} from "../../modelSelection"; import { applyProviderInstanceSettings, deriveProviderInstanceEntries, @@ -17,6 +20,7 @@ import { useSettingsScope } from "./SettingsScopeContext"; export function useScopedModelDisabledReason( settings: UnifiedSettings, entries: readonly ProviderInstanceEntry[], + textGeneration = false, ) { const { targets } = useSettingsScope(); const { environments } = useEnvironments(); @@ -48,7 +52,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 1f112c48e01d7..17bd7a4a24183 100644 --- a/apps/web/src/modelSelection.test.ts +++ b/apps/web/src/modelSelection.test.ts @@ -25,6 +25,8 @@ import { import { defaultModelSelectionsAreMixed, getAppModelOptionsForInstance, + getCustomModelOptionsByInstance, + getTextGenerationModelOptionsByInstance, resolveAppModelSelectionForInstance, resolveAppModelSelectionState, resolveEffectiveDefaultModelSelection, @@ -917,6 +919,60 @@ describe("instance-scoped model selection", () => { NO_PROVIDER_MODEL_SELECTION, ); }); + + 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 61be2ece7632f..bf03e216fc123 100644 --- a/apps/web/src/modelSelection.ts +++ b/apps/web/src/modelSelection.ts @@ -366,6 +366,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; } @@ -390,7 +426,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( @@ -398,13 +434,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 3449b01b99054..a983e7c0826ff 100644 --- a/packages/contracts/src/server.ts +++ b/packages/contracts/src/server.ts @@ -83,6 +83,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 2b887500f3023..78d16ebd31c62 100644 --- a/packages/provider-pi/src/server/status.test.ts +++ b/packages/provider-pi/src/server/status.test.ts @@ -2,14 +2,15 @@ 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"; import * as Stream from "effect/Stream"; import * as ChildProcess from "effect/process/ChildProcess"; import * as ChildProcessSpawner from "effect/process/ChildProcessSpawner"; -import * as Fiber from "effect/Fiber"; -import * as TestClock from "effect/testing/TestClock"; import { checkPiProviderStatus, MINIMUM_PI_VERSION } from "./status.ts"; @@ -57,14 +58,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(); @@ -79,13 +76,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, @@ -96,7 +97,31 @@ const piDiscoverySpawner = ( } }), ); - return ChildProcessSpawner.make((command) => { + return ChildProcessSpawner.makeHandle({ + ...processHandle({}), + exitCode: Effect.never, + isRunning: Effect.succeed(true), + stdin, + stdout: Stream.fromQueue(stdout), + }); + }); + +/** Makes the extension-free launch exit before answering, like launch args naming an extension provider. */ +const exitsWithoutExtensions = Symbol("exitsWithoutExtensions"); + +// 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` })); @@ -107,17 +132,18 @@ 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), - }), + if (args.includes("--no-extensions") && extensionFreeInventory === exitsWithoutExtensions) { + return Effect.succeed( + processHandle({ stderr: 'Error: Unknown provider "demo".', exitCode: 1 }), + ); + } + return piRpcProcess( + args.includes("--no-extensions") ? extensionFreeInventory : inventory, + args.includes("--no-extensions") ? onExtensionFreeModelRequest : Effect.void, + options.state, ); - }); - }); + }), + ); const settings = { enabled: true, @@ -263,6 +289,8 @@ describe("PiProvider", () => { { provider: "router", id: "auto" }, ], }, + undefined, + Effect.void, { version: row.version, state: { model: row.activeModel }, authCheck: row.authCheck }, ); const snapshot = yield* checkPiProviderStatus({ @@ -276,6 +304,89 @@ 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], + ], + ); + assert.equal(snapshot.supportsTextGeneration, undefined); + }).pipe(Effect.provide(NodeServices.layer)), + ); + + it.effect.each([ + { label: "exits before answering", extensionFreeInventory: exitsWithoutExtensions }, + { label: "lists no models", extensionFreeInventory: { models: [] } }, + ])("makes Pi unusable for text generation when Pi without extensions $label", (row) => + Effect.gen(function* () { + const spawner = yield* piDiscoverySpawner( + { models: [{ provider: "demo", id: "extension-model" }] }, + row.extensionFreeInventory, + ); + const snapshot = yield* checkPiProviderStatus({ + ...settings, + launchArgs: "--provider demo --model extension-model", + customModels: ["local/custom-alias"], + }).pipe(Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner)); + assert.equal(snapshot.status, "ready"); + assert.equal(snapshot.auth.status, "authenticated"); + // Pi default and custom models still run in chat, so the picker keeps them. + assert.deepEqual( + snapshot.models.map((model) => model.slug), + ["default", "demo/extension-model", "local/custom-alias"], + ); + assert.equal(snapshot.supportsTextGeneration, 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.equal(snapshot.supportsTextGeneration, undefined); + assert.include( + snapshot.slashCommands.map((command) => command.name), + "compact", + ); + }).pipe(Effect.provide(NodeServices.layer)), + ); + it.effect("advertises a model's service tiers beside its thinking levels", () => Effect.gen(function* () { const spawner = yield* piDiscoverySpawner({ diff --git a/packages/provider-pi/src/server/status.ts b/packages/provider-pi/src/server/status.ts index 5dfce10b81060..01047b00eaf26 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"; @@ -57,6 +58,7 @@ const PI_PRESENTATION = { const VERSION_PROBE_TIMEOUT_MS = 10_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. @@ -81,6 +83,8 @@ interface PiDiscovery extends PiDiscoveredCommands { readonly authenticated: boolean; /** Pi's active model when Pi lists no model of its provider, so it cannot start a turn. */ readonly unavailableActiveModel: { readonly provider: string; readonly id: string } | undefined; + /** Pi without extensions can load no model, so text generation fails for every entry. */ + readonly textGenerationUnavailable: boolean; } function piModelsFromSettings( @@ -141,6 +145,7 @@ const makePiDiscoveryConnection = Effect.fnUntraced(function* ( environment: NodeJS.ProcessEnv, launchArgs: ReadonlyArray, cwd?: string, + textGenerationLaunch = false, ) { const launch = buildPiRpcLaunch({ launchArgs, @@ -148,6 +153,8 @@ const makePiDiscoveryConnection = Effect.fnUntraced(function* ( mcpSession: undefined, extensionPath: undefined, ephemeral: true, + disableExtensions: textGenerationLaunch, + disableTools: textGenerationLaunch, shortLived: true, }); const connection = yield* makePiRpcConnection({ @@ -164,12 +171,55 @@ 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( + // Pi exits before answering when launch args name a provider only an extension registers. + Effect.catchIf( + (error) => + error._tag === "PiRpcError" && (error.operation === "read" || error.operation === "exit"), + () => Effect.succeed(new Set()), + ), + 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 +228,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, @@ -194,6 +252,7 @@ const discoverPiViaRpc = Effect.fn("PiProviderStatus.discoverPiViaRpc")(function !discoveredModels.some((model) => model.subProvider === activeProvider) ? { provider: activeProvider, id: activeModelId } : undefined, + textGenerationUnavailable: textGenerationModelSlugs?.size === 0, } satisfies PiDiscovery; }, Effect.scoped); @@ -514,22 +573,25 @@ export const checkPiProviderStatus = Effect.fn("checkPiProviderStatus")(function resolvedLaunchArgs.args, discovery.unavailableActiveModel, ); - return buildServerProvider({ - presentation: PI_PRESENTATION, - enabled: piSettings.enabled, - checkedAt, - models, - slashCommands: discovery.slashCommands, - skills: discovery.skills, - probe: { - installed: true, - version, - // Other Pi models still work when only the default is unavailable, so keep the picker usable. - status: discovery.authenticated ? "ready" : "warning", - auth: { status: discovery.authenticated ? "authenticated" : "unauthenticated", type: "pi" }, - ...(message === undefined ? {} : { message }), - }, - }); + return { + ...buildServerProvider({ + presentation: PI_PRESENTATION, + enabled: piSettings.enabled, + checkedAt, + models, + slashCommands: discovery.slashCommands, + skills: discovery.skills, + probe: { + installed: true, + version, + // Other Pi models still work when only the default is unavailable, so keep the picker usable. + status: discovery.authenticated ? "ready" : "warning", + auth: { status: discovery.authenticated ? "authenticated" : "unauthenticated", type: "pi" }, + ...(message === undefined ? {} : { message }), + }, + }), + ...(discovery.textGenerationUnavailable ? { supportsTextGeneration: false } : {}), + }; }); export const enrichPiSnapshot = (input: {