@animus-labs/cortex 0.2.2 → 0.2.4

package/src/provider-manager.ts

@@ -17,13 +17,14 @@
 import {
   PROVIDER_REGISTRY,
   OAUTH_PROVIDER_IDS,
-  UTILITY_MODEL_DEFAULTS,
+  UTILITY_MODEL_OVERRIDES,
 } from './provider-registry.js';
 import { createRequire } from 'node:module';
 import type { IncomingMessage, ServerResponse } from 'node:http';
 import type { ThinkingLevel } from './types.js';
 import type { ProviderInfo, ModelInfo } from './provider-registry.js';
 import { wrapModel } from './model-wrapper.js';
+import { inferUtilityModelId } from './utility-model-inference.js';
 import type { CortexModel } from './model-wrapper.js';
 
 const nodeRequire = createRequire(import.meta.url);
@@ -74,6 +75,16 @@ export interface OAuthCallbacks {
    * localhost callback routes and is restored immediately after the login flow.
    */
   renderCallbackPage?: OAuthCallbackPageRenderer | undefined;
+
+  /**
+   * Overall timeout for the OAuth flow, in milliseconds. pi-ai's
+   * callback-server flows (e.g. Anthropic) do not honor an abort signal and
+   * hang forever if the callback never arrives or arrives with an error, so
+   * Cortex enforces this timeout itself and rejects with an
+   * `OAuthError('timed_out')`. Defaults to 5 minutes. Pass `0` or a negative
+   * value to disable (not recommended).
+   */
+  timeoutMs?: number | undefined;
 }
 
 /** Status of the browser callback page produced by an OAuth flow. */
@@ -151,6 +162,48 @@ export interface OAuthRefreshResult {
   changed: boolean;
 }
 
+/**
+ * Discriminant for OAuth flow failures, so consumers can render specific
+ * UX instead of parsing error strings.
+ *
+ * - `unsupported_provider`: provider has no OAuth support.
+ * - `callback_port_in_use`: the provider's fixed loopback callback port is
+ *   already bound (e.g. another Anthropic app on 53692, or a leftover flow).
+ *   Detected before the browser opens.
+ * - `cancelled`: the flow was cancelled via `cancelOAuth()`.
+ * - `timed_out`: the flow exceeded its timeout (pi-ai's callback servers do
+ *   not honor an abort signal, so this is the backstop against hangs).
+ * - `callback_failed`: the browser callback fired but the provider reported
+ *   an error (e.g. state mismatch). Surfaced immediately instead of hanging.
+ */
+export type OAuthErrorCode =
+  | 'unsupported_provider'
+  | 'callback_port_in_use'
+  | 'cancelled'
+  | 'timed_out'
+  | 'callback_failed';
+
+/** Structured error thrown by initiateOAuth. */
+export class OAuthError extends Error {
+  readonly code: OAuthErrorCode;
+  readonly provider: string;
+  /** The fixed callback port, when relevant (`callback_port_in_use`). */
+  readonly port?: number | undefined;
+
+  constructor(
+    code: OAuthErrorCode,
+    provider: string,
+    message: string,
+    options?: { port?: number | undefined; cause?: unknown },
+  ) {
+    super(message, options?.cause !== undefined ? { cause: options.cause } : undefined);
+    this.name = 'OAuthError';
+    this.code = code;
+    this.provider = provider;
+    this.port = options?.port;
+  }
+}
+
 /** Configuration for creating a custom model endpoint. */
 export interface CustomModelConfig {
   /** Base URL of the OpenAI-compatible API (e.g., 'http://localhost:11434/v1'). */
@@ -255,7 +308,13 @@ interface ActiveOAuthCallbackPageShim {
   readonly provider: string;
   readonly providerName: string;
   readonly route: OAuthCallbackRoute;
-  readonly render: OAuthCallbackPageRenderer;
+  readonly render: OAuthCallbackPageRenderer | undefined;
+  /**
+   * Notified exactly once when the browser callback fires and its status
+   * (success/error) is known. Lets the flow react immediately instead of
+   * waiting on pi-ai (which hangs on non-success callbacks).
+   */
+  readonly onResult?: ((status: OAuthCallbackPageStatus, context: OAuthCallbackPageContext) => void) | undefined;
 }
 
 type ServerResponseEnd = ServerResponse['end'];
@@ -266,30 +325,88 @@ const OAUTH_CALLBACK_ROUTES: Record<string, OAuthCallbackRoute> = {
 };
 
 let activeOAuthCallbackPageShim: ActiveOAuthCallbackPageShim | null = null;
+/** Ensures `onResult` fires at most once per installed shim. */
+let oauthCallbackResultNotified = false;
+
+/** Default overall OAuth flow timeout (pi-ai hangs without this). */
+const DEFAULT_OAUTH_FLOW_TIMEOUT_MS = 5 * 60_000;
+
+/**
+ * Probe whether something is already listening on a loopback port. Used to
+ * fail an OAuth flow fast (before opening a browser) when the provider's
+ * fixed callback port is occupied — otherwise pi-ai binds the other stack /
+ * the browser hits the wrong listener and the user gets a dead page while
+ * pi-ai waits forever.
+ */
+function probeCallbackPortInUse(port: number, host: string): Promise<boolean> {
+  const net = nodeRequire('node:net') as typeof import('node:net');
+  return new Promise((resolve) => {
+    let settled = false;
+    const finish = (inUse: boolean) => {
+      if (settled) return;
+      settled = true;
+      socket.destroy();
+      resolve(inUse);
+    };
+    const socket = net.connect({ port, host });
+    socket.once('connect', () => finish(true));
+    socket.once('error', () => finish(false));
+    socket.setTimeout(600, () => finish(false));
+  });
+}
+
+/**
+ * Throw an `OAuthError('callback_port_in_use')` if the provider's fixed
+ * callback port is occupied on either IPv4 or IPv6 loopback. No-op for
+ * providers without a known callback route (manual/device-code flows).
+ */
+async function assertOAuthCallbackPortAvailable(provider: string): Promise<void> {
+  const route = OAUTH_CALLBACK_ROUTES[provider];
+  if (!route) return;
+
+  for (const host of ['127.0.0.1', '::1']) {
+    if (await probeCallbackPortInUse(route.port, host)) {
+      throw new OAuthError(
+        'callback_port_in_use',
+        provider,
+        `OAuth callback port ${route.port} for "${provider}" is already in use ` +
+        `(detected on ${host}). This is a fixed port: another application is ` +
+        `holding it, or a previous sign-in did not finish. Close that ` +
+        `application (or restart the host process), then try again.`,
+        { port: route.port },
+      );
+    }
+  }
+}
 
-async function withOAuthCallbackPageShim<T>(
+/**
+ * Install the callback-page shim for a flow when there is a known callback
+ * route AND something to do with it (a custom renderer and/or a result
+ * observer). Returns a release function (a no-op when no shim was installed).
+ *
+ * Unlike a `try/finally` wrapper around pi-ai's `login()`, the caller owns
+ * the release lifecycle: pi-ai's callback-server flows hang forever on a
+ * non-success callback, so cleanup must be tied to the flow's own
+ * race/timeout, not to awaiting the (possibly never-settling) login promise.
+ */
+function maybeInstallOAuthCallbackShim(
   provider: string,
   providerName: string,
   render: OAuthCallbackPageRenderer | undefined,
-  run: () => Promise<T>,
-): Promise<T> {
+  onResult: ActiveOAuthCallbackPageShim['onResult'],
+): () => void {
   const route = OAUTH_CALLBACK_ROUTES[provider];
-  if (!render || !route) {
-    return run();
+  if (!route || (!render && !onResult)) {
+    return () => {};
   }
 
-  const release = installOAuthCallbackPageShim({
+  return installOAuthCallbackPageShim({
     provider,
     providerName,
     route,
     render,
+    onResult,
   });
-
-  try {
-    return await run();
-  } finally {
-    release();
-  }
 }
 
 function installOAuthCallbackPageShim(shim: ActiveOAuthCallbackPageShim): () => void {
@@ -303,6 +420,7 @@ function installOAuthCallbackPageShim(shim: ActiveOAuthCallbackPageShim): () =>
   const prototype = http.ServerResponse.prototype;
   const previousEnd = prototype.end;
   activeOAuthCallbackPageShim = shim;
+  oauthCallbackResultNotified = false;
 
   const patchedEnd = function patchedOAuthCallbackEnd(this: ServerResponse, ...args: unknown[]) {
     const replacement = maybeRenderOAuthCallbackPage(this, args[0]);
@@ -373,6 +491,19 @@ function maybeRenderOAuthCallbackPage(response: ServerResponse, chunk: unknown):
     context.details = details;
   }
 
+  // Notify the flow that the browser callback fired (once). This lets
+  // initiateOAuth react to a failed callback immediately rather than
+  // waiting on pi-ai, which hangs on non-success callbacks.
+  if (!oauthCallbackResultNotified && shim.onResult) {
+    oauthCallbackResultNotified = true;
+    try {
+      shim.onResult(status, context);
+    } catch {
+      // An observer must never break the callback response.
+    }
+  }
+
+  if (!shim.render) return null;
   try {
     const rendered = shim.render(context);
     return typeof rendered === 'string' && rendered.trim().length > 0 ? rendered : null;
@@ -639,10 +770,12 @@ function mapRawToModelInfo(
     supportsThinking: supportedThinkingLevels.some(level => level !== 'off')
       || !!(raw['supportsThinking'] || raw['reasoning']),
     supportedThinkingLevels,
-    supportsImages: !!raw['supportsImages'],
+    supportsImages: Array.isArray(raw['input'])
+      ? raw['input'].includes('image')
+      : !!raw['supportsImages'],
   };
 
-  const rawPricing = raw['pricing'];
+  const rawPricing = raw['pricing'] ?? raw['cost'];
   if (rawPricing && typeof rawPricing === 'object') {
     const pricing = rawPricing as Record<string, unknown>;
     const inputPrice = pricing['input'];
@@ -729,41 +862,112 @@ export class ProviderManager implements IProviderManager {
    * @param provider - OAuth provider identifier
    * @param callbacks - UI callbacks for auth URL, prompts, and progress
    * @returns The OAuth credentials and display metadata
-   * @throws Error if the provider does not support OAuth or pi-ai is not installed
+   * @throws {OAuthError} `unsupported_provider`, `callback_port_in_use`,
+   *   `cancelled`, `timed_out`, or `callback_failed`. Other errors (e.g.
+   *   network/token-exchange failures from pi-ai) propagate as-is.
    */
   async initiateOAuth(provider: string, callbacks: OAuthCallbacks): Promise<OAuthResult> {
     const oauthModule = await loadPiAiOAuth();
     const oauthProvider = oauthModule.getOAuthProvider?.(provider);
     if (!oauthProvider) {
-      throw new Error(`Provider "${provider}" does not support OAuth`);
+      throw new OAuthError(
+        'unsupported_provider',
+        provider,
+        `Provider "${provider}" does not support OAuth`,
+      );
     }
 
-    this.activeOAuthAbort = new AbortController();
+    // (A) Fail fast — before opening a browser — if the provider's fixed
+    // callback port is already taken. Otherwise pi-ai binds the other
+    // stack, the browser hits the wrong listener, and pi-ai waits forever.
+    await assertOAuthCallbackPortAvailable(provider);
 
-    try {
-      const rawCredentials = await withOAuthCallbackPageShim(
+    const abort = new AbortController();
+    this.activeOAuthAbort = abort;
+
+    // (C) pi-ai only settles its callback wait on success; on a failed
+    // callback (e.g. state mismatch) it hangs. The render shim already sees
+    // that response — use it to fail the flow immediately with the reason.
+    let failFromCallback!: (err: OAuthError) => void;
+    const callbackFailure = new Promise<never>((_, reject) => {
+      failFromCallback = reject;
+    });
+    const handleCallbackResult = (
+      status: OAuthCallbackPageStatus,
+      ctx: OAuthCallbackPageContext,
+    ): void => {
+      if (status !== 'error') return;
+      const detail = ctx.details ? ` (${ctx.details})` : '';
+      failFromCallback(new OAuthError(
+        'callback_failed',
         provider,
-        oauthProvider.name,
-        callbacks.renderCallbackPage,
-        () => oauthProvider.login({
-          onAuth: callbacks.onAuth,
-          onPrompt: callbacks.onPrompt,
-          onProgress: callbacks.onProgress,
-          onManualCodeInput: callbacks.onManualCodeInput,
-          onSelect: callbacks.onSelect,
-          signal: this.activeOAuthAbort!.signal,
-        }),
-      );
+        `OAuth callback for "${provider}" reported a failure: ${ctx.message}${detail}`,
+      ));
+    };
 
-      this.activeOAuthAbort = null;
+    const releaseShim = maybeInstallOAuthCallbackShim(
+      provider,
+      oauthProvider.name,
+      callbacks.renderCallbackPage,
+      handleCallbackResult,
+    );
+
+    // (B) pi-ai callback servers ignore the abort signal, so cancellation
+    // and timeout are enforced here. Without this the flow hangs forever.
+    const timeoutMs = callbacks.timeoutMs ?? DEFAULT_OAUTH_FLOW_TIMEOUT_MS;
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    const timeout = new Promise<never>((_, reject) => {
+      if (timeoutMs > 0 && Number.isFinite(timeoutMs)) {
+        timer = setTimeout(() => reject(new OAuthError(
+          'timed_out',
+          provider,
+          `OAuth flow for "${provider}" timed out after ${timeoutMs}ms.`,
+        )), timeoutMs);
+        timer.unref?.();
+      }
+    });
+    const cancelled = new Promise<never>((_, reject) => {
+      abort.signal.addEventListener('abort', () => reject(new OAuthError(
+        'cancelled',
+        provider,
+        `OAuth flow for "${provider}" was cancelled.`,
+      )), { once: true });
+    });
+
+    const login = oauthProvider.login({
+      onAuth: callbacks.onAuth,
+      onPrompt: callbacks.onPrompt,
+      onProgress: callbacks.onProgress,
+      onManualCodeInput: callbacks.onManualCodeInput,
+      onSelect: callbacks.onSelect,
+      signal: abort.signal,
+    }) as Promise<Record<string, unknown>>;
+    // Whichever promise loses the race may still settle later (pi-ai's
+    // login can hang or settle late; the aux promises can reject after the
+    // race is decided). Attach inert handlers so a late rejection never
+    // surfaces as an unhandled rejection. Promise.race still observes the
+    // first settlement independently.
+    login.catch(() => {});
+    cancelled.catch(() => {});
+    timeout.catch(() => {});
+    callbackFailure.catch(() => {});
+
+    try {
+      const rawCredentials = await Promise.race([
+        login,
+        cancelled,
+        timeout,
+        callbackFailure,
+      ]);
 
       const credentials = JSON.stringify(rawCredentials);
       const meta = buildOAuthMeta(provider, rawCredentials);
 
       return { credentials, meta };
-    } catch (err) {
-      this.activeOAuthAbort = null;
-      throw err;
+    } finally {
+      if (timer) clearTimeout(timer);
+      releaseShim();
+      if (this.activeOAuthAbort === abort) this.activeOAuthAbort = null;
     }
   }
 
@@ -830,32 +1034,31 @@ export class ProviderManager implements IProviderManager {
   async validateApiKey(provider: string, apiKey: string): Promise<ApiKeyValidationResult> {
     const piAi = await loadPiAi();
 
-    // Find the cheapest model for this provider to minimize validation cost
-    const cheapestModelId = this.getSmallestModelId(provider);
-    if (!cheapestModelId) {
-      // No known model, try a generic test with the provider's first model
-      const models = piAi.getModels(provider);
-      if (models.length === 0) {
-        return {
-          provider,
-          modelId: null,
-          valid: false,
-          retryable: false,
-          status: 'resolution_error',
-          message: `No models found for provider "${provider}"`,
-        };
-      }
-      const firstRawId = models[0]!['id'];
-      const firstRawName = models[0]!['name'];
-      const firstModelId = typeof firstRawId === 'string'
-        ? firstRawId
-        : typeof firstRawName === 'string'
-          ? firstRawName
-          : String(firstRawId ?? firstRawName);
-      return this.tryValidation(piAi, provider, firstModelId, apiKey);
+    const models = piAi.getModels(provider) ?? [];
+    if (models.length === 0) {
+      return {
+        provider,
+        modelId: null,
+        valid: false,
+        retryable: false,
+        status: 'resolution_error',
+        message: `No models found for provider "${provider}"`,
+      };
+    }
+
+    const modelId = this.getSmallestModelId(provider, models);
+    if (!modelId) {
+      return {
+        provider,
+        modelId: null,
+        valid: false,
+        retryable: false,
+        status: 'resolution_error',
+        message: `No usable models found for provider "${provider}"`,
+      };
     }
 
-    return this.tryValidation(piAi, provider, cheapestModelId, apiKey);
+    return this.tryValidation(piAi, provider, modelId, apiKey);
   }
 
   /**
@@ -948,11 +1151,10 @@ export class ProviderManager implements IProviderManager {
   // -----------------------------------------------------------------------
 
   /**
-   * Get the cheapest known model ID for a provider.
-   * Uses the UTILITY_MODEL_DEFAULTS as a proxy for "smallest model."
+   * Get the cheapest likely utility model ID for a provider.
    */
-  private getSmallestModelId(provider: string): string | null {
-    return UTILITY_MODEL_DEFAULTS[provider] ?? null;
+  private getSmallestModelId(provider: string, models: Array<Record<string, unknown>>): string | null {
+    return UTILITY_MODEL_OVERRIDES[provider] ?? inferUtilityModelId(models);
   }
 
   /**

package/src/provider-registry.ts

@@ -4,7 +4,7 @@
  * This module contains:
  * 1. PROVIDER_REGISTRY: metadata for all known providers (auth methods, env vars, key prefixes)
  * 2. OAUTH_PROVIDER_IDS: the subset of providers that support OAuth
- * 3. UTILITY_MODEL_DEFAULTS: per-provider cheapest-capable model for utility operations
+ * 3. UTILITY_MODEL_OVERRIDES: per-provider utility model overrides for inference exceptions
  *
  * OAuth flows are resolved through pi-ai's OAuth provider registry at runtime.
  *
@@ -270,17 +270,9 @@ export const OAUTH_PROVIDER_IDS: string[] = [
 ];
 
 // ---------------------------------------------------------------------------
-// Utility Model Defaults
+// Model Defaults
 // ---------------------------------------------------------------------------
 
-/**
- * Default utility model IDs per provider.
- * Used when utilityModel is 'default' or undefined.
- *
- * These are the cheapest capable models for each provider,
- * suitable for internal operations like WebFetch summarization
- * and safety classification.
- */
 /**
  * Default primary model IDs per provider.
  * Used when a user first connects a provider and no model is explicitly selected.
@@ -289,21 +281,22 @@ export const OAUTH_PROVIDER_IDS: string[] = [
 export const PRIMARY_MODEL_DEFAULTS: Record<string, string> = {
   anthropic: 'claude-sonnet-4-6',
   openai: 'gpt-5.4',
+  'openai-codex': 'gpt-5.5',
   google: 'gemini-3.1-pro-preview',
+  xai: 'grok-4',
   groq: 'openai/gpt-oss-120b',
   cerebras: 'gpt-oss-120b',
   mistral: 'mistral-large-2512',
 };
 
-export const UTILITY_MODEL_DEFAULTS: Record<string, string> = {
-  anthropic: 'claude-haiku-4-5-20251001',     // $1.00/$5.00 per 1M tokens
-  openai: 'gpt-4.1-nano',                     // $0.10/$0.40 per 1M tokens
-  'openai-codex': 'gpt-5.4-mini',            // Current small Codex-capable model
-  google: 'gemini-2.5-flash-lite',            // $0.10/$0.40 per 1M tokens
-  groq: 'llama-3.1-8b-instant',              // ~$0.05/$0.08 per 1M tokens
-  cerebras: 'llama3.1-8b',                    // ~$0.10/$0.10 per 1M tokens
-  mistral: 'mistral-small-2506',             // $0.06/$0.18 per 1M tokens
-};
+/**
+ * Per-provider utility model overrides for inference exceptions.
+ * Leave empty unless dynamic inference picks a bad utility model for a provider.
+ */
+export const UTILITY_MODEL_OVERRIDES: Record<string, string> = {};
+
+/** Backwards-compatible alias. Prefer UTILITY_MODEL_OVERRIDES for new code. */
+export const UTILITY_MODEL_DEFAULTS = UTILITY_MODEL_OVERRIDES;
 
 // ---------------------------------------------------------------------------
 // Cache Retention