fix(credentials): close resolver switch and guard missing provider/auth

- credential-resolver.resolve: throw explicit CREDENTIAL_INVALID when the connector has no auth schema, and add a default branch so the switch has a definitive return path for every auth kind - credential-resolver api_key branch: guard definition.provider being unset before calling pi-ai's getEnvApiKey - cost-writer: only attempt cost lookup when definition.provider is defined; adapter-only connectors still write an unattributed row
2026-04-17 01:28:37 -04:00
1 parent bee5535349
commit cad17b4d45
@@ -0,0 +1,66 @@
+import { calculateCost, getModel, type Usage } from "@mariozechner/pi-ai";
+import type { AppDatabase } from "@main/database";
+import { getConnectorDefinition } from "@main/connectors";
+
+export interface CostWriterDeps {
+  db: AppDatabase;
+  logger: { warn: (...args: unknown[]) => void };
+}
+
+interface ObserveInput {
+  runId: string;
+  companyId: string;
+  agentId: string;
+  taskId: string | null;
+  connectorId: string;
+  modelId: string;
+  usage: Usage;
+}
+
+/**
+ * Writes a `cost_entries` row for an assistant `message_end` and updates the
+ * cumulative `runs.cost_usd` for the active run.
+ */
+export class CostWriter {
+  constructor(private readonly deps: CostWriterDeps) {}
+
+  observeAssistantMessage(input: ObserveInput): void {
+    let amountUsd: number | null = null;
+    try {
+      const definition = getConnectorDefinition(input.connectorId);
+      if (definition.provider) {
+        const model = getModel(definition.provider as never, input.modelId as never);
+        if (model) {
+          const cost = calculateCost(model, input.usage);
+          amountUsd = Number.isFinite(cost.total) ? cost.total : null;
+        }
+      }
+    } catch (error) {
+      this.deps.logger.warn(
+        `[cost-writer] cost calculation failed for ${input.connectorId}/${input.modelId}: ${error instanceof Error ? error.message : String(error)}`,
+      );
+      amountUsd = null;
+    }
+
+    this.deps.db.queryAll(
+      "insert into cost_entries (id, company_id, run_id, task_id, agent_id, connector_id, amount_usd, unattributed, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+      cryptoRandomId(),
+      input.companyId,
+      input.runId,
+      input.taskId,
+      input.agentId,
+      input.connectorId,
+      amountUsd,
+      amountUsd === null ? 1 : 0,
+      new Date().toISOString(),
+    );
+  }
+}
+
+function cryptoRandomId(): string {
+  // Avoid pulling node:crypto at module top-level so this file can be
+  // imported from non-Node contexts during type checks.
+  return globalThis.crypto?.randomUUID
+    ? globalThis.crypto.randomUUID()
+    : `${Date.now()}-${Math.random().toString(36).slice(2)}`;
+}
@@ -0,0 +1,165 @@
+import { getOAuthProvider } from "@mariozechner/pi-ai/oauth";
+import { getEnvApiKey } from "@mariozechner/pi-ai";
+import { getConnectorDefinition } from "@main/connectors";
+import { CredentialStore } from "./credential-store";
+import type { OAuthCredentialPayload, StoredCredentialPayload } from "./types";
+
+const REFRESH_WINDOW_MS = 60_000;
+
+export class CredentialError extends Error {
+  constructor(
+    public readonly code: "CREDENTIAL_MISSING" | "CREDENTIAL_EXPIRED" | "CREDENTIAL_INVALID" | "CREDENTIAL_REFRESH_FAILED",
+    public readonly connectorId: string,
+    message?: string,
+  ) {
+    super(message ?? code);
+    this.name = "CredentialError";
+  }
+}
+
+export interface ResolvedCredential {
+  apiKey: string;
+  extraEnv?: Record<string, string>;
+}
+
+export class CredentialResolver {
+  private readonly refreshLocks = new Map<string, Promise<void>>();
+
+  constructor(private readonly store: CredentialStore) {}
+
+  async resolve(connectorId: string): Promise<ResolvedCredential> {
+    const definition = getConnectorDefinition(connectorId);
+    if (!definition.auth) {
+      throw new CredentialError("CREDENTIAL_INVALID", connectorId, `Connector ${connectorId} has no auth schema.`);
+    }
+    const auth = definition.auth;
+    const stored = this.store.read(connectorId);
+
+    switch (auth.type) {
+      case "oauth": {
+        if (!auth.oauthId) {
+          throw new CredentialError("CREDENTIAL_INVALID", connectorId, `OAuth auth is missing oauthId.`);
+        }
+        if (stored?.kind !== "oauth") {
+          throw new CredentialError("CREDENTIAL_MISSING", connectorId, `OAuth not yet completed for ${connectorId}.`);
+        }
+        await this.refreshIfNeeded(connectorId, auth.oauthId, stored.value);
+        const fresh = this.store.read(connectorId);
+        if (fresh?.kind !== "oauth") {
+          throw new CredentialError("CREDENTIAL_INVALID", connectorId);
+        }
+        const provider = getOAuthProvider(auth.oauthId);
+        if (!provider) throw new CredentialError("CREDENTIAL_INVALID", connectorId, `Unknown OAuth provider: ${auth.oauthId}`);
+        return { apiKey: provider.getApiKey(fresh.value as never) };
+      }
+      case "api_key": {
+        if (stored?.kind === "api_key") {
+          if (stored.value.source === "literal") {
+            return { apiKey: stored.value.value };
+          }
+          const fromEnv = process.env[stored.value.value];
+          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}.`);
+        }
+        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 };
+      }
+      case "azure": {
+        if (stored?.kind !== "azure") {
+          throw new CredentialError("CREDENTIAL_MISSING", connectorId);
+        }
+        const cfg = stored.value;
+        const env: Record<string, string> = { AZURE_OPENAI_API_KEY: cfg.apiKey };
+        if (cfg.baseUrl) env.AZURE_OPENAI_BASE_URL = cfg.baseUrl;
+        if (cfg.resourceName) env.AZURE_OPENAI_RESOURCE_NAME = cfg.resourceName;
+        if (cfg.apiVersion) env.AZURE_OPENAI_API_VERSION = cfg.apiVersion;
+        if (cfg.deploymentMap && Object.keys(cfg.deploymentMap).length > 0) {
+          env.AZURE_OPENAI_DEPLOYMENT_NAME_MAP = Object.entries(cfg.deploymentMap)
+            .map(([k, v]) => `${k}=${v}`)
+            .join(",");
+        }
+        return { apiKey: cfg.apiKey, extraEnv: env };
+      }
+      case "aws": {
+        if (stored?.kind !== "aws") {
+          throw new CredentialError("CREDENTIAL_MISSING", connectorId);
+        }
+        const cfg = stored.value;
+        const env: Record<string, string> = { AWS_REGION: cfg.region ?? "us-east-1" };
+        if (cfg.authMode === "profile" && cfg.awsProfile) env.AWS_PROFILE = cfg.awsProfile;
+        if (cfg.authMode === "iam") {
+          if (!cfg.accessKeyId || !cfg.secretAccessKey) {
+            throw new CredentialError("CREDENTIAL_INVALID", connectorId, "IAM auth requires both accessKeyId and secretAccessKey.");
+          }
+          env.AWS_ACCESS_KEY_ID = cfg.accessKeyId;
+          env.AWS_SECRET_ACCESS_KEY = cfg.secretAccessKey;
+          if (cfg.sessionToken) env.AWS_SESSION_TOKEN = cfg.sessionToken;
+        }
+        if (cfg.authMode === "bearer") {
+          if (!cfg.bearerToken) {
+            throw new CredentialError("CREDENTIAL_INVALID", connectorId, "Bearer auth requires a bearerToken.");
+          }
+          env.AWS_BEARER_TOKEN_BEDROCK = cfg.bearerToken;
+        }
+        if (cfg.forceCache) env.AWS_BEDROCK_FORCE_CACHE = "1";
+        if (cfg.endpointUrl) env.AWS_ENDPOINT_URL_BEDROCK_RUNTIME = cfg.endpointUrl;
+        if (cfg.skipAuth) env.AWS_BEDROCK_SKIP_AUTH = "1";
+        if (cfg.forceHttp1) env.AWS_BEDROCK_FORCE_HTTP1 = "1";
+        return { apiKey: "", extraEnv: env };
+      }
+      case "gcp_adc": {
+        if (stored?.kind !== "gcp_adc") {
+          throw new CredentialError("CREDENTIAL_MISSING", connectorId);
+        }
+        const cfg = stored.value;
+        const env: Record<string, string> = { GOOGLE_CLOUD_PROJECT: cfg.project, GOOGLE_CLOUD_LOCATION: cfg.location };
+        if (cfg.credentialsPath) env.GOOGLE_APPLICATION_CREDENTIALS = cfg.credentialsPath;
+        return { apiKey: "", extraEnv: env };
+      }
+      default:
+        throw new CredentialError("CREDENTIAL_INVALID", connectorId, `Unsupported auth kind: ${String(auth.type)}`);
+    }
+  }
+
+  private async refreshIfNeeded(connectorId: string, oauthId: string, creds: OAuthCredentialPayload): Promise<void> {
+    if (typeof creds.expires !== "number" || !Number.isFinite(creds.expires)) return;
+    if (Date.now() < creds.expires - REFRESH_WINDOW_MS) return;
+
+    const existing = this.refreshLocks.get(connectorId);
+    if (existing) {
+      await existing;
+      return;
+    }
+
+    const promise = (async () => {
+      const provider = getOAuthProvider(oauthId);
+      if (!provider) {
+        throw new CredentialError("CREDENTIAL_INVALID", connectorId, `Unknown OAuth provider: ${oauthId}`);
+      }
+      try {
+        const refreshed = await provider.refreshToken(creds as never);
+        const payload: StoredCredentialPayload = { kind: "oauth", value: refreshed as OAuthCredentialPayload };
+        const expiresAt = typeof refreshed.expires === "number" && Number.isFinite(refreshed.expires)
+          ? new Date(refreshed.expires).toISOString()
+          : null;
+        this.store.saveRefreshedOAuth(connectorId, payload, expiresAt);
+      } catch (error) {
+        this.store.setStatus(connectorId, "expired", error instanceof Error ? error.message : String(error));
+        throw new CredentialError("CREDENTIAL_REFRESH_FAILED", connectorId, error instanceof Error ? error.message : String(error));
+      }
+    })();
+
+    this.refreshLocks.set(connectorId, promise);
+    try {
+      await promise;
+    } finally {
+      this.refreshLocks.delete(connectorId);
+    }
+  }
+}