@crouton-kit/crouter 0.3.32 → 0.3.34

package/dist/builtin-pi-packages/pi-crtr-extensions/extensions/provider-rotation.ts

@@ -0,0 +1,574 @@
+import { spawn } from "node:child_process";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import {
+  createAssistantMessageEventStream,
+  type AssistantMessage,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type SimpleStreamOptions,
+} from "@earendil-works/pi-ai";
+import { streamSimpleAnthropic, streamSimpleOpenAICodexResponses } from "@earendil-works/pi-ai/compat";
+import {
+  loginAnthropic,
+  loginOpenAICodex,
+  refreshAnthropicToken,
+  refreshOpenAICodexToken,
+} from "@earendil-works/pi-ai/oauth";
+import {
+  ANTHROPIC_PROVIDER_ID,
+  OPENAI_CODEX_PROVIDER_ID,
+  formatStatusLine,
+  getProviderLabel,
+  isManagedProvider,
+  markSubscriptionAttempt,
+  markSubscriptionRateLimited,
+  markSubscriptionSuccess,
+  parseRetryAfterHeader,
+  promoteSubscription,
+  readRotationConfig,
+  readSubscriptionPool,
+  removeSubscription,
+  resolveFallbackTarget,
+  restorePreferredModelIfPossible,
+  switchToFallbackIfPossible,
+  upsertSubscription,
+  type FallbackTarget,
+  type ManagedProviderId,
+  type SubscriptionCredential,
+} from "../lib/subscription-state.js";
+
+const DEFAULT_RATE_LIMIT_BACKOFF_MS = 5 * 60 * 1000;
+const INVALID_REFRESH_TOKEN_BACKOFF_MS = 30 * 24 * 60 * 60 * 1000;
+
+let runtimeContext: ExtensionContext | undefined;
+let extensionAPI: ExtensionAPI | undefined;
+
+type AttemptResult = {
+  result: "success" | "retry" | "fatal";
+  retryAfterMs?: number;
+  attemptAt: number;
+};
+
+function isFallbackLike(error: unknown): boolean {
+  const value = error as { status?: unknown; errorMessage?: unknown; message?: unknown } | undefined;
+  const status = typeof value?.status === "number" ? value.status : Number.NaN;
+  const message = [value?.errorMessage, value?.message, error instanceof Error ? error.message : ""]
+    .filter((part): part is string => typeof part === "string" && part.length > 0)
+    .join(" ");
+  return (
+    status === 429 ||
+    status === 500 ||
+    status === 502 ||
+    status === 503 ||
+    status === 504 ||
+    /\b429\b|rate.?limit|too many requests|overloaded|service.?unavailable|temporarily unavailable|upstream|quota|usage|connection|network|fetch failed|socket|econnreset|econnrefused|enotfound|etimedout|timed? out|timeout|request was aborted/i.test(message)
+  );
+}
+
+function isInvalidRefreshTokenError(error: unknown): boolean {
+  const value = error as { status?: unknown; errorMessage?: unknown; message?: unknown } | undefined;
+  const status = typeof value?.status === "number" ? value.status : Number.NaN;
+  const message = [value?.errorMessage, value?.message, error instanceof Error ? error.message : ""]
+    .filter((part): part is string => typeof part === "string" && part.length > 0)
+    .join(" ");
+  return status === 400 && /invalid_grant|refresh token .*invalid|refresh token .*not found/i.test(message);
+}
+
+function setStatus(text: string): void {
+  runtimeContext?.ui.setStatus("provider-rotation", text);
+}
+
+function logRotation(message: string): void {
+  process.stderr.write(`[provider-rotation] ${message}\n`);
+}
+
+function getCurrentModelRef(model?: Model<any>): { providerId: ManagedProviderId; modelId: string } | undefined {
+  return isManagedProvider(model?.provider) ? { providerId: model.provider, modelId: model.id } : undefined;
+}
+
+function toAssistantError(model: Model<any>, errorMessage: string): AssistantMessage {
+  return {
+    role: "assistant",
+    content: [],
+    api: model.api,
+    provider: model.provider,
+    model: model.id,
+    usage: {
+      input: 0,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: 0,
+      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+    },
+    stopReason: "error",
+    errorMessage,
+    timestamp: Date.now(),
+  };
+}
+
+function retryAfterMs(headers: Record<string, string>): number | undefined {
+  const retryAfterMsHeader = headers["retry-after-ms"];
+  if (retryAfterMsHeader) {
+    const millis = Number(retryAfterMsHeader);
+    if (Number.isFinite(millis)) return Math.max(0, millis);
+  }
+  return parseRetryAfterHeader(headers["retry-after"]);
+}
+
+type RefreshedCredential = Pick<SubscriptionCredential, "refresh" | "access" | "expires">;
+type LoginCredential = Pick<SubscriptionCredential, "refresh" | "access" | "expires">;
+
+async function defaultRefreshForProvider(providerId: ManagedProviderId, refreshToken: string): Promise<RefreshedCredential> {
+  return providerId === ANTHROPIC_PROVIDER_ID ? refreshAnthropicToken(refreshToken) : refreshOpenAICodexToken(refreshToken);
+}
+
+// Indirection so tests can inject refresh failures without real network calls.
+let activeRefreshForProvider: typeof defaultRefreshForProvider = defaultRefreshForProvider;
+
+/** Test-only seam: override the per-provider refresh function. Pass undefined to restore the default. */
+export function __setRefreshForProviderForTest(fn: typeof defaultRefreshForProvider | undefined): void {
+  activeRefreshForProvider = fn ?? defaultRefreshForProvider;
+}
+
+function openBrowser(url: string): void {
+  const [cmd, args] =
+    process.platform === "darwin"
+      ? ["open", [url]]
+      : process.platform === "win32"
+        ? ["cmd", ["/c", "start", "", url]]
+        : ["xdg-open", [url]];
+  try {
+    const child = spawn(cmd, args, { stdio: "ignore", detached: true });
+    child.on("error", () => {
+      /* the notification still contains the URL */
+    });
+    child.unref();
+  } catch {
+    /* best effort */
+  }
+}
+
+async function refreshCredentialIfNeeded(providerId: ManagedProviderId, credential: SubscriptionCredential): Promise<SubscriptionCredential> {
+  if (credential.expires > Date.now()) return credential;
+  const refreshed = await activeRefreshForProvider(providerId, credential.refresh);
+  const next = {
+    ...credential,
+    refresh: refreshed.refresh,
+    access: refreshed.access,
+    expires: refreshed.expires,
+  };
+  return upsertSubscription(providerId, next).find((entry) => entry.label === next.label) ?? next;
+}
+
+async function defaultLoginForProvider(providerId: ManagedProviderId, ctx: ExtensionContext): Promise<LoginCredential> {
+  return loginProvider(providerId, ctx);
+}
+
+let activeLoginForProvider: typeof defaultLoginForProvider = defaultLoginForProvider;
+
+/** Test-only seam: override the per-provider login function. Pass undefined to restore the default. */
+export function __setLoginForProviderForTest(fn: typeof defaultLoginForProvider | undefined): void {
+  activeLoginForProvider = fn ?? defaultLoginForProvider;
+}
+
+async function reauthenticateCredential(providerId: ManagedProviderId, credential: SubscriptionCredential): Promise<SubscriptionCredential | undefined> {
+  const ctx = runtimeContext;
+  if (!ctx?.hasUI) return undefined;
+  ctx.ui.notify(`${getProviderLabel(providerId)} subscription "${credential.label}" needs re-authentication. Opening OAuth login...`, "warn");
+  const authenticated = await activeLoginForProvider(providerId, ctx);
+  const next: SubscriptionCredential = {
+    ...credential,
+    refresh: authenticated.refresh,
+    access: authenticated.access,
+    expires: authenticated.expires,
+    rateLimitedUntil: 0,
+    lastAttemptAt: Date.now(),
+    lastRateLimitedAt: credential.lastRateLimitedAt,
+  };
+  return upsertSubscription(providerId, next).find((entry) => entry.label === next.label) ?? next;
+}
+
+function defaultStreamForProvider(model: Model<any>, context: Context, options: SimpleStreamOptions | undefined): AssistantMessageEventStream {
+  if (model.provider === ANTHROPIC_PROVIDER_ID) {
+    return streamSimpleAnthropic(model as Model<"anthropic-messages">, context, options);
+  }
+  return streamSimpleOpenAICodexResponses(model as Model<"openai-codex-responses">, context, options);
+}
+
+// Indirection so tests can inject a fake per-provider stream without real network calls.
+let activeStreamForProvider: typeof defaultStreamForProvider = defaultStreamForProvider;
+
+function streamForProvider(model: Model<any>, context: Context, options: SimpleStreamOptions | undefined): AssistantMessageEventStream {
+  return activeStreamForProvider(model, context, options);
+}
+
+/** Test-only seam: override the per-provider stream factory. Pass undefined to restore the default. */
+export function __setStreamForProviderForTest(fn: typeof defaultStreamForProvider | undefined): void {
+  activeStreamForProvider = fn ?? defaultStreamForProvider;
+}
+
+async function streamProviderAttempt(
+  providerId: ManagedProviderId,
+  model: Model<any>,
+  context: Context,
+  options: SimpleStreamOptions | undefined,
+  credential: SubscriptionCredential,
+  outer: AssistantMessageEventStream,
+): Promise<AttemptResult> {
+  const attemptAt = Date.now();
+  markSubscriptionAttempt(providerId, credential.label, attemptAt);
+  let requestedRetryAfterMs: number | undefined;
+  const inner = streamForProvider(model, context, {
+    ...options,
+    apiKey: credential.access,
+    onResponse: async (response, responseModel) => {
+      if (response.status === 429 || response.status === 503) {
+        requestedRetryAfterMs = retryAfterMs(response.headers);
+      }
+      await options?.onResponse?.(response, responseModel);
+    },
+  });
+
+  const bufferedEvents: Parameters<AssistantMessageEventStream["push"]>[0][] = [];
+  let sawContent = false;
+
+  const flush = (): void => {
+    for (const event of bufferedEvents) outer.push(event);
+    bufferedEvents.length = 0;
+  };
+
+  try {
+    for await (const event of inner) {
+      if (event.type === "error") {
+        if (isFallbackLike(event.error)) {
+          return { result: sawContent ? "fatal" : "retry", retryAfterMs: requestedRetryAfterMs, attemptAt };
+        }
+        bufferedEvents.push(event);
+        flush();
+        return { result: "fatal", attemptAt };
+      }
+      if (event.type !== "start" && event.type !== "done") sawContent = true;
+      bufferedEvents.push(event);
+    }
+  } catch (error) {
+    if (isFallbackLike(error)) {
+      return { result: sawContent ? "fatal" : "retry", retryAfterMs: requestedRetryAfterMs, attemptAt };
+    }
+    bufferedEvents.push({
+      type: "error",
+      reason: "error",
+      error: toAssistantError(model, error instanceof Error ? error.message : String(error)),
+    });
+    flush();
+    return { result: "fatal", attemptAt };
+  }
+
+  if (requestedRetryAfterMs !== undefined && !sawContent) {
+    return { result: "retry", retryAfterMs: requestedRetryAfterMs, attemptAt };
+  }
+
+  flush();
+  return { result: "success", attemptAt };
+}
+
+async function switchToFallbackModel(originalModel: { providerId: ManagedProviderId; modelId: string }): Promise<FallbackTarget> {
+  const ctx = runtimeContext;
+  if (!extensionAPI || !ctx) {
+    throw new Error(`Provider fallback unavailable: runtime not initialized for ${originalModel.providerId}/${originalModel.modelId}`);
+  }
+  const target = await switchToFallbackIfPossible(extensionAPI, ctx, originalModel);
+  if (!target) {
+    const resolved = resolveFallbackTarget(originalModel.providerId, originalModel.modelId, readRotationConfig(), extensionAPI.getThinkingLevel?.());
+    throw new Error(`Provider fallback unavailable: failed to switch to ${resolved.label}`);
+  }
+  return target;
+}
+
+function emitTerminalError(stream: AssistantMessageEventStream, model: Model<any>, message: string): void {
+  stream.push({ type: "error", reason: "error", error: toAssistantError(model, message) });
+  stream.end();
+}
+
+// Drives one managed provider's pool into `stream`; on exhaustion it switches the
+// session model to the cross-provider fallback and CONTINUES the same turn on it,
+// rather than failing the turn with an error. `attempted` guards anthropic<->codex
+// ping-pong: once both pools are dry we emit a single terminal error.
+async function runManagedProvider(
+  model: Model<any>,
+  context: Context,
+  options: SimpleStreamOptions | undefined,
+  stream: AssistantMessageEventStream,
+  attempted: Set<ManagedProviderId>,
+): Promise<void> {
+  const providerId = model.provider as ManagedProviderId;
+  attempted.add(providerId);
+  const available = readSubscriptionPool(providerId).filter((entry) => !entry.rateLimitedUntil || entry.rateLimitedUntil <= Date.now());
+
+  for (const credential of available) {
+    let freshCredential: SubscriptionCredential;
+    try {
+      freshCredential = await refreshCredentialIfNeeded(providerId, credential);
+    } catch (error) {
+      const attemptAt = Date.now();
+      markSubscriptionAttempt(providerId, credential.label, attemptAt);
+      if (isInvalidRefreshTokenError(error)) {
+        let reauthenticated: SubscriptionCredential | undefined;
+        try {
+          reauthenticated = await reauthenticateCredential(providerId, credential);
+        } catch (loginError) {
+          runtimeContext?.ui.notify(
+            `${getProviderLabel(providerId)} re-authentication failed: ${loginError instanceof Error ? loginError.message : String(loginError)}`,
+            "error",
+          );
+        }
+        if (reauthenticated) {
+          freshCredential = reauthenticated;
+        } else {
+          markSubscriptionRateLimited(providerId, credential.label, INVALID_REFRESH_TOKEN_BACKOFF_MS, attemptAt);
+          setStatus(formatStatusLine(providerId));
+          continue;
+        }
+      } else {
+        throw error;
+      }
+    }
+    const result = await streamProviderAttempt(providerId, model, context, options, freshCredential, stream);
+
+    if (result.result === "success") {
+      markSubscriptionSuccess(providerId, freshCredential.label, result.attemptAt);
+      logRotation(`${getProviderLabel(providerId)} subscription "${freshCredential.label}" succeeded; cooldown cleared`);
+      setStatus(formatStatusLine(providerId));
+      stream.end();
+      return;
+    }
+
+    if (result.result === "retry") {
+      const cooldownMs = result.retryAfterMs ?? DEFAULT_RATE_LIMIT_BACKOFF_MS;
+      markSubscriptionRateLimited(providerId, freshCredential.label, cooldownMs, result.attemptAt);
+      logRotation(`${getProviderLabel(providerId)} subscription "${freshCredential.label}" cooled down for ${Math.ceil(cooldownMs / 1000)}s; trying next credential`);
+      continue;
+    }
+
+    // Fatal (or partial content already streamed): cannot safely switch providers mid-turn.
+    stream.end();
+    return;
+  }
+
+  // Pool exhausted — switch the session model to the fallback and continue this turn on it.
+  const originalModel = getCurrentModelRef(model);
+  if (!originalModel) throw new Error(`Provider fallback unavailable: missing current model`);
+
+  const plannedTarget = resolveFallbackTarget(originalModel.providerId, originalModel.modelId, readRotationConfig(), extensionAPI?.getThinkingLevel?.());
+  if (attempted.has(plannedTarget.providerId)) {
+    emitTerminalError(stream, model, `All managed provider pools exhausted (tried ${[...attempted].join(", ")})`);
+    return;
+  }
+
+  const target = await switchToFallbackModel(originalModel);
+  logRotation(`${getProviderLabel(providerId)} pool exhausted -> ${target.label}`);
+  setStatus(`${getProviderLabel(providerId)} -> ${target.label}`);
+
+  const fallbackModel = runtimeContext?.modelRegistry.find(target.providerId, target.modelId);
+  if (!fallbackModel) {
+    emitTerminalError(stream, model, `Provider fallback unavailable: ${target.label} not in model registry`);
+    return;
+  }
+
+  await runManagedProvider(fallbackModel, context, options, stream, attempted);
+}
+
+function streamManagedProvider(model: Model<any>, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream {
+  if (!isManagedProvider(model.provider)) {
+    return streamForProvider(model, context, options);
+  }
+
+  const stream = createAssistantMessageEventStream();
+
+  runManagedProvider(model, context, options, stream, new Set<ManagedProviderId>()).catch((error) => {
+    stream.push({
+      type: "error",
+      reason: "error",
+      error: toAssistantError(model, error instanceof Error ? error.message : String(error)),
+    });
+    stream.end();
+  });
+
+  return stream;
+}
+
+function parseCommand(args: string): string[] {
+  return args
+    .trim()
+    .split(/\s+/)
+    .map((part) => part.trim())
+    .filter(Boolean);
+}
+
+function parseProvider(value: string | undefined): ManagedProviderId | undefined {
+  return isManagedProvider(value) ? value : undefined;
+}
+
+function updateRotationStatus(ctx: ExtensionContext): void {
+  const current = getCurrentModelRef(ctx.model);
+  if (current) setStatus(formatStatusLine(current.providerId));
+}
+
+async function loginProvider(providerId: ManagedProviderId, ctx: ExtensionContext) {
+  if (providerId === ANTHROPIC_PROVIDER_ID) {
+    return loginAnthropic({
+      onAuth: ({ url, instructions }) => {
+        const lines = [`Open this URL to authenticate Claude: ${url}`];
+        if (instructions) lines.push(instructions);
+        ctx.ui.notify(lines.join("\n"), "info");
+        openBrowser(url);
+      },
+      onPrompt: async (prompt) => (await ctx.ui.input(prompt.message, prompt.placeholder ?? "")) ?? "",
+      onProgress: (message) => ctx.ui.notify(message, "info"),
+      onManualCodeInput: async () => (await ctx.ui.input("Paste Anthropic redirect URL or code", "http://localhost:53692/callback?...")) ?? "",
+    });
+  }
+
+  return loginOpenAICodex({
+    onAuth: ({ url, instructions }) => {
+      const lines = [`Open this URL to authenticate OpenAI Codex: ${url}`];
+      if (instructions) lines.push(instructions);
+      ctx.ui.notify(lines.join("\n"), "info");
+      openBrowser(url);
+    },
+    onPrompt: async (prompt) => (await ctx.ui.input(prompt.message, prompt.placeholder ?? "")) ?? "",
+    onProgress: (message) => ctx.ui.notify(message, "info"),
+    onManualCodeInput: async () => (await ctx.ui.input("Paste OpenAI redirect URL or code", "http://localhost:1455/auth/callback?...")) ?? "",
+  });
+}
+
+export default async function (pi: ExtensionAPI): Promise<void> {
+  extensionAPI = pi;
+
+  pi.registerProvider(ANTHROPIC_PROVIDER_ID, {
+    api: "anthropic-messages",
+    streamSimple: streamManagedProvider,
+  });
+
+  pi.registerProvider(OPENAI_CODEX_PROVIDER_ID, {
+    api: "openai-codex-responses",
+    streamSimple: streamManagedProvider,
+  });
+
+  pi.on("session_start", async (_event, ctx) => {
+    runtimeContext = ctx;
+    const restored = await restorePreferredModelIfPossible(pi, ctx);
+    if (restored && ctx.model) {
+      logRotation(`restored preferred ${ctx.model.provider}/${ctx.model.id}`);
+      setStatus(formatStatusLine(ctx.model.provider as ManagedProviderId));
+    }
+    updateRotationStatus(ctx);
+  });
+
+  pi.on("before_agent_start", async (_event, ctx) => {
+    runtimeContext = ctx;
+    const restored = await restorePreferredModelIfPossible(pi, ctx);
+    if (restored && ctx.model) {
+      logRotation(`restored preferred ${ctx.model.provider}/${ctx.model.id}`);
+      setStatus(formatStatusLine(ctx.model.provider as ManagedProviderId));
+    }
+  });
+
+  pi.on("model_select", async (_event, ctx) => {
+    runtimeContext = ctx;
+    updateRotationStatus(ctx);
+  });
+
+  pi.registerCommand("provider-sub", {
+    description: "Manage subscription credential pools for provider rotation",
+    handler: async (args, ctx) => {
+      const [providerArg, verb, ...rest] = parseCommand(args);
+      const providerId = parseProvider(providerArg);
+      if (!providerId) {
+        ctx.ui.notify("Usage: /provider-sub <anthropic|openai-codex> <list|add|select|rm>", "error");
+        return;
+      }
+
+      const action = verb === "rm" ? "remove" : verb || "list";
+      const pool = readSubscriptionPool(providerId);
+
+      if (action === "status" || action === "list") {
+        const lines = [formatStatusLine(providerId), ""];
+        for (const [index, entry] of pool.entries()) {
+          const cooldown = entry.rateLimitedUntil > Date.now() ? ` cooldown=${Math.ceil((entry.rateLimitedUntil - Date.now()) / 1000)}s` : "";
+          lines.push(`${index + 1}. ${entry.label}${cooldown}`);
+        }
+        ctx.ui.notify(lines.join("\n"), "info");
+        return;
+      }
+
+      if (action === "add") {
+        if (!ctx.hasUI) {
+          ctx.ui.notify(`${getProviderLabel(providerId)} login requires an interactive UI.`, "error");
+          return;
+        }
+        const label = rest.join(" ") || (await ctx.ui.input(`${getProviderLabel(providerId)} subscription label`, providerId));
+        if (!label) return;
+        if (pool.some((entry) => entry.label.toLowerCase() === label.toLowerCase())) {
+          ctx.ui.notify(`${getProviderLabel(providerId)} subscription already exists: ${label}`, "error");
+          return;
+        }
+        try {
+          const cred = await loginProvider(providerId, ctx);
+          upsertSubscription(providerId, {
+            label,
+            refresh: cred.refresh,
+            access: cred.access,
+            expires: cred.expires,
+            rateLimitedUntil: 0,
+            lastAttemptAt: 0,
+            lastRateLimitedAt: 0,
+          });
+          logRotation(`added ${getProviderLabel(providerId)} subscription "${label}"`);
+          ctx.ui.notify(`Added ${getProviderLabel(providerId)} subscription: ${label}`, "info");
+          updateRotationStatus(ctx);
+        } catch (error) {
+          ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
+        }
+        return;
+      }
+
+      if (action === "select") {
+        const ref = rest.join(" ");
+        if (!ref) {
+          ctx.ui.notify(`Usage: /provider-sub ${providerId} select <label|index>`, "error");
+          return;
+        }
+        try {
+          promoteSubscription(providerId, ref);
+          logRotation(`selected ${getProviderLabel(providerId)} subscription "${ref}"`);
+          ctx.ui.notify(`Preferred ${getProviderLabel(providerId)} subscription: ${ref}`, "info");
+          updateRotationStatus(ctx);
+        } catch (error) {
+          ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
+        }
+        return;
+      }
+
+      if (action === "remove") {
+        const ref = rest.join(" ");
+        if (!ref) {
+          ctx.ui.notify(`Usage: /provider-sub ${providerId} rm <label|index>`, "error");
+          return;
+        }
+        try {
+          removeSubscription(providerId, ref);
+          logRotation(`removed ${getProviderLabel(providerId)} subscription "${ref}"`);
+          ctx.ui.notify(`Removed ${getProviderLabel(providerId)} subscription: ${ref}`, "info");
+          updateRotationStatus(ctx);
+        } catch (error) {
+          ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
+        }
+        return;
+      }
+
+      ctx.ui.notify(`Unknown /provider-sub verb: ${verb}`, "error");
+    },
+  });
+}