From 7793f2da35d5896fc43ecbc4c2778d0e0d777207 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 22 Apr 2026 18:58:30 -0400 Subject: [PATCH] fix(connectors): guard pi provider resolution --- src/main/connector-definitions.ts | 6 ++- src/main/connector-health-service.ts | 4 +- src/main/connectors.ts | 30 +++++++++-- src/main/credentials/credential-resolver.ts | 13 ++--- src/main/model-resolver.ts | 15 ++++-- tests/unit/connector-provider-guard.test.ts | 56 +++++++++++++++++++++ 6 files changed, 107 insertions(+), 17 deletions(-) create mode 100644 tests/unit/connector-provider-guard.test.ts diff --git a/src/main/connector-definitions.ts b/src/main/connector-definitions.ts index b7230406..f906672f 100644 --- a/src/main/connector-definitions.ts +++ b/src/main/connector-definitions.ts @@ -23,6 +23,7 @@ * - `connector-definitions/cloud.ts` — 3 Cloud providers */ +import type { KnownProvider } from "@mariozechner/pi-ai"; import type { ConnectorId } from "@shared/types"; import { OAUTH_CONNECTORS, API_KEY_CONNECTORS, CLOUD_CONNECTORS } from "./connector-definitions/index"; @@ -67,11 +68,14 @@ export type ConnectorAuth = | ConnectorAuthAzure | ConnectorAuthGcpAdc; +export type LocalModelResolverProvider = "fireworks" | "fireworks-firepass"; +export type ConnectorProvider = KnownProvider | LocalModelResolverProvider; + export interface ConnectorDefinition { /** Stable AgentCompany id; the `connectors` table primary key. */ id: ConnectorId; /** pi-ai `KnownProvider` id (or a locally-resolved provider for Fireworks). */ - provider: string; + provider: ConnectorProvider; /** Human-readable label for the Settings panel and onboarding rail. */ label: string; /** One-line explanation shown under the label. */ diff --git a/src/main/connector-health-service.ts b/src/main/connector-health-service.ts index 829f96f0..73fd9830 100644 --- a/src/main/connector-health-service.ts +++ b/src/main/connector-health-service.ts @@ -1,6 +1,6 @@ import { getEnvApiKey } from "@mariozechner/pi-ai"; import type { AppDatabase } from "./database"; -import { getConnectorDefinition, parseEnvBindingText } from "./connectors"; +import { getConnectorDefinition, isPiKnownProvider, parseEnvBindingText } from "./connectors"; import type { ConnectorId } from "@shared/types"; import type { CredentialStore } from "./credentials"; @@ -51,7 +51,7 @@ export function createConnectorHealthService({ if (stored?.kind === "api_key") { return { status: "ready", authState: "ok", lastError: null }; } - if (definition.provider && getEnvApiKey(definition.provider as never)) { + if (isPiKnownProvider(definition.provider) && getEnvApiKey(definition.provider)) { return { status: "ready", authState: "ok", lastError: null }; } // Re-parse envBindingText on every probe: it is user-editable and must diff --git a/src/main/connectors.ts b/src/main/connectors.ts index a65309e2..20a1bc4e 100644 --- a/src/main/connectors.ts +++ b/src/main/connectors.ts @@ -1,8 +1,9 @@ -import type { CapabilityMatrix, ConnectorRecord } from "@shared/types"; -import { CONNECTORS, type ConnectorDefinition } from "./connector-definitions"; +import { getProviders, type KnownProvider } from "@mariozechner/pi-ai"; +import type { CapabilityMatrix, ConnectorId, ConnectorRecord } from "@shared/types"; +import { CONNECTORS, type ConnectorDefinition, type ConnectorProvider, type LocalModelResolverProvider } from "./connector-definitions"; import { buildApiKeyDoc, buildCloudDoc, buildOAuthDoc } from "./connector-definitions-docs"; export { CONNECTORS } from "./connector-definitions"; -export type { ConnectorDefinition } from "./connector-definitions"; +export type { ConnectorDefinition, ConnectorProvider, LocalModelResolverProvider } from "./connector-definitions"; export { findConnectorDefinition, listFeaturedOnboardingConnectors } from "./connector-definitions-lookup"; const RESERVED_ENV_BINDING_KEYS = new Set([ @@ -17,6 +18,12 @@ const RESERVED_ENV_BINDING_KEYS = new Set([ "PWD", ]); +const LOCAL_MODEL_RESOLVER_PROVIDERS: ReadonlySet = new Set([ + "fireworks", + "fireworks-firepass", +]); +const PI_KNOWN_PROVIDER_IDS: ReadonlySet = new Set(getProviders()); + /** * Pi-mono drives all provider dispatch through pi-ai, so the adapter-era * capability matrix now collapses to a uniform baseline. These values are @@ -39,6 +46,23 @@ export function getConnectorDefinition(id: string) { return definition; } +export function isPiKnownProvider(provider: string): provider is KnownProvider { + return PI_KNOWN_PROVIDER_IDS.has(provider); +} + +export function isLocalModelResolverProvider(provider: string): provider is LocalModelResolverProvider { + return LOCAL_MODEL_RESOLVER_PROVIDERS.has(provider as LocalModelResolverProvider); +} + +export function getConnectorProvider(id: ConnectorId): ConnectorProvider { + return getConnectorDefinition(id).provider; +} + +export function getConnectorKnownProvider(id: ConnectorId): KnownProvider | null { + const provider = getConnectorProvider(id); + return isPiKnownProvider(provider) ? provider : null; +} + export function parseEnvBindingText(text: string): Record { const bindings: Record = {}; for (const rawLine of text.split(/\r?\n/)) { diff --git a/src/main/credentials/credential-resolver.ts b/src/main/credentials/credential-resolver.ts index 383ffbf5..483480c1 100644 --- a/src/main/credentials/credential-resolver.ts +++ b/src/main/credentials/credential-resolver.ts @@ -1,6 +1,6 @@ import { getOAuthProvider } from "@mariozechner/pi-ai/oauth"; import { getEnvApiKey } from "@mariozechner/pi-ai"; -import { getConnectorDefinition } from "@main/connectors"; +import { getConnectorDefinition, isPiKnownProvider } from "@main/connectors"; import { CredentialStore } from "./credential-store"; import type { OAuthCredentialPayload, StoredCredentialPayload } from "./types"; @@ -72,14 +72,11 @@ export class CredentialResolver { if (!fromEnv) throw new CredentialError("CREDENTIAL_MISSING", connectorId, `Env var ${stored.value.value} is not set.`); return { apiKey: fromEnv }; } - if (!definition.provider) { - throw new CredentialError("CREDENTIAL_MISSING", connectorId, `No credential or env API key for ${connectorId}.`); + if (isPiKnownProvider(definition.provider)) { + const envKey = getEnvApiKey(definition.provider); + if (envKey) return { apiKey: envKey }; } - const envKey = getEnvApiKey(definition.provider as never); - if (!envKey) { - throw new CredentialError("CREDENTIAL_MISSING", connectorId, `No credential or env API key for ${connectorId}.`); - } - return { apiKey: envKey }; + throw new CredentialError("CREDENTIAL_MISSING", connectorId, `No credential or env API key for ${connectorId}.`); } case "azure": { if (stored?.kind !== "azure") { diff --git a/src/main/model-resolver.ts b/src/main/model-resolver.ts index 81c9d3b9..7d1c1464 100644 --- a/src/main/model-resolver.ts +++ b/src/main/model-resolver.ts @@ -17,7 +17,8 @@ */ import log from "electron-log/main.js"; -import { getModel, getModels, type Api, type Model } from "@mariozechner/pi-ai"; +import { getModels, type Api, type Model } from "@mariozechner/pi-ai"; +import { isPiKnownProvider } from "./connectors"; // Fireworks' OpenAI-compatible root. The OpenAI SDK appends // `/chat/completions`, `/models`, etc. to this baseURL. @@ -67,8 +68,12 @@ export function resolveModel(provider: string, modelId: string): Model | nu const curated = FIREWORKS_CURATED_MODELS.find((entry) => entry.id === modelId); return buildFireworksModel(modelId, curated?.name, curated?.contextWindow); } + if (!isPiKnownProvider(provider)) { + log.warn(`[model-resolver] Unknown pi provider: ${provider}`); + return null; + } try { - return getModel(provider as never, modelId as never) ?? null; + return getModels(provider).find((model) => model.id === modelId) ?? null; } catch (err) { log.warn(`[model-resolver] getModel failed for ${provider}/${modelId}:`, err instanceof Error ? err.message : String(err)); return null; @@ -82,8 +87,12 @@ export function resolveModels(provider: string): Array<{ id: string; name: strin if (provider === "fireworks") { return FIREWORKS_CURATED_MODELS.map(({ id, name }) => ({ id, name })); } + if (!isPiKnownProvider(provider)) { + log.warn(`[model-resolver] Unknown pi provider: ${provider}`); + return []; + } try { - const models = getModels(provider as never) ?? []; + const models = getModels(provider); return models.map((model: { id: string; name?: string }) => ({ id: model.id, name: model.name ?? model.id })); } catch (err) { log.warn(`[model-resolver] getModels failed for ${provider}:`, err instanceof Error ? err.message : String(err)); diff --git a/tests/unit/connector-provider-guard.test.ts b/tests/unit/connector-provider-guard.test.ts new file mode 100644 index 00000000..a4aba8cb --- /dev/null +++ b/tests/unit/connector-provider-guard.test.ts @@ -0,0 +1,56 @@ +import { readFileSync } from "node:fs"; +import { join } from "node:path"; +import { getProviders } from "@mariozechner/pi-ai"; +import { describe, expect, it } from "vitest"; +import { + CONNECTORS, + getConnectorKnownProvider, + getConnectorProvider, + isLocalModelResolverProvider, + isPiKnownProvider, +} from "../../src/main/connectors"; +import { FIREWORKS_FIREPASS_MODEL_ID, resolveModel, resolveModels } from "../../src/main/model-resolver"; + +describe("connector provider guard", () => { + it("classifies catalog providers before they enter pi-ai APIs", () => { + const piProviders = new Set(getProviders()); + + for (const connector of CONNECTORS) { + const provider = getConnectorProvider(connector.id); + expect(provider).toBe(connector.provider); + + if (isLocalModelResolverProvider(provider)) { + expect(getConnectorKnownProvider(connector.id)).toBeNull(); + } else { + expect(isPiKnownProvider(provider)).toBe(true); + expect(piProviders.has(provider)).toBe(true); + expect(getConnectorKnownProvider(connector.id)).toBe(provider); + } + } + }); + + it("uses the provider guard in model discovery", () => { + const anthropicModels = resolveModels("anthropic"); + const firstAnthropicModel = anthropicModels[0]; + if (!firstAnthropicModel) throw new Error("Anthropic model catalog is empty."); + + expect(resolveModel("anthropic", firstAnthropicModel.id)?.id).toBe(firstAnthropicModel.id); + expect(resolveModels("not-a-pi-provider")).toEqual([]); + expect(resolveModel("not-a-pi-provider", "model-id")).toBeNull(); + expect(resolveModel("fireworks-firepass", "ignored")?.id).toBe(FIREWORKS_FIREPASS_MODEL_ID); + }); + + it("keeps provider calls free of erased never casts", () => { + const files = [ + "src/main/model-resolver.ts", + "src/main/connector-health-service.ts", + "src/main/credentials/credential-resolver.ts", + ]; + + for (const file of files) { + const source = readFileSync(join(process.cwd(), file), "utf8"); + expect(source).not.toContain("provider as never"); + expect(source).not.toContain("definition.provider as never"); + } + }); +});