takomi 2.0.6 → 2.1.0

package/.pi/extensions/oauth-router/provider.ts

@@ -0,0 +1,492 @@
+import {
+  createAssistantMessageEventStream,
+  type Api,
+  type AssistantMessage,
+  type AssistantMessageEvent,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type SimpleStreamOptions,
+  streamSimpleOpenAICodexResponses,
+  streamSimpleOpenAICompletions,
+  streamSimpleOpenAIResponses,
+} from "@mariozechner/pi-ai";
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { loadRouterConfig } from "./config.ts";
+import { refreshAccountCredentials, getApiKeyForAccount } from "./oauth-flow.ts";
+import { RouterAccountStore } from "./oauth-store.ts";
+import { chooseEligibleAccount } from "./policies.ts";
+import { RouterStateStore } from "./state.ts";
+import type {
+  DelegateSelection,
+  EligibleAccount,
+  FailureClassification,
+  RouterConfig,
+  RouterErrorMessageInput,
+  RouterModelConfig,
+  RouterStatusRow,
+  RouterUpstreamConfig,
+  RoutingPolicyName,
+  StoredRouterAccount,
+} from "./types.ts";
+
+function now() {
+  return Date.now();
+}
+
+function parseRetryAfterMs(value: string | undefined): number | undefined {
+  if (!value) return undefined;
+  const seconds = Number(value);
+  if (Number.isFinite(seconds)) return Math.max(0, seconds * 1000);
+  const timestamp = Date.parse(value);
+  if (Number.isFinite(timestamp)) return Math.max(0, timestamp - now());
+  return undefined;
+}
+
+function isMeaningfulEvent(event: AssistantMessageEvent): boolean {
+  switch (event.type) {
+    case "text_delta":
+    case "thinking_delta":
+    case "toolcall_delta":
+      return event.delta.length > 0;
+    case "text_end":
+    case "thinking_end":
+      return event.content.length > 0;
+    case "toolcall_end":
+      return true;
+    case "done":
+      return event.message.content.length > 0;
+    default:
+      return false;
+  }
+}
+
+function createErrorMessage({ model, message, stopReason = "error" }: RouterErrorMessageInput): 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,
+    errorMessage: message,
+    timestamp: Date.now(),
+  };
+}
+
+function classifyFailure(
+  status: number | undefined,
+  headers: Record<string, string>,
+  message: string,
+  config: RouterConfig,
+): FailureClassification {
+  const retryAfterMs = parseRetryAfterMs(headers["retry-after"]) ?? config.rateLimitCooldownMs;
+  const lower = message.toLowerCase();
+
+  if (status === 429 || lower.includes("rate limit") || lower.includes("too many requests")) {
+    return { kind: "rate-limit", status: status ?? 429, retryAfterMs, message };
+  }
+
+  if (status === 401 || status === 403 || lower.includes("unauthorized") || lower.includes("forbidden")) {
+    return { kind: "auth", status: status ?? 401, message };
+  }
+
+  if (
+    (status !== undefined && status >= 500) ||
+    lower.includes("network") ||
+    lower.includes("fetch failed") ||
+    lower.includes("connection") ||
+    lower.includes("timeout")
+  ) {
+    return { kind: "transient", status: status ?? 500, message };
+  }
+
+  return { kind: "fatal", status, message };
+}
+
+function delegateStream(
+  upstream: RouterUpstreamConfig,
+  model: Model<Api>,
+  context: Context,
+  options: SimpleStreamOptions | undefined,
+): AssistantMessageEventStream {
+  switch (upstream.api) {
+    case "openai-completions":
+      return streamSimpleOpenAICompletions(model as Model<"openai-completions">, context, options);
+    case "openai-responses":
+      return streamSimpleOpenAIResponses(model as Model<"openai-responses">, context, options);
+    case "openai-codex-responses":
+      return streamSimpleOpenAICodexResponses(model as Model<"openai-codex-responses">, context, options);
+    default:
+      throw new Error(`Unsupported upstream api: ${upstream.api}`);
+  }
+}
+
+export class RouterRuntime {
+  private config: RouterConfig;
+  private accounts: RouterAccountStore;
+  private state: RouterStateStore;
+
+  constructor() {
+    this.config = loadRouterConfig();
+    this.accounts = new RouterAccountStore();
+    this.state = new RouterStateStore(this.config.policy);
+    this.state.pruneAccountIds(this.accounts.list().map((account) => account.id));
+  }
+
+  reloadConfig() {
+    this.config = loadRouterConfig();
+    this.accounts.reload();
+    this.state.pruneAccountIds(this.accounts.list().map((account) => account.id));
+  }
+
+  getConfig(): RouterConfig {
+    this.reloadConfig();
+    return this.config;
+  }
+
+  listUpstreams(): RouterUpstreamConfig[] {
+    return this.getConfig().upstreams;
+  }
+
+  listAccounts(): StoredRouterAccount[] {
+    this.reloadConfig();
+    return this.accounts.list();
+  }
+
+  getAccount(id: string): StoredRouterAccount | undefined {
+    this.reloadConfig();
+    return this.accounts.get(id);
+  }
+
+  addAccount(account: StoredRouterAccount) {
+    this.accounts.add(account);
+    this.state.ensureAccount(account.id);
+  }
+
+  updateAccount(account: StoredRouterAccount) {
+    this.accounts.update(account);
+    this.state.ensureAccount(account.id);
+  }
+
+  removeAccount(id: string) {
+    this.accounts.remove(id);
+    this.state.pruneAccountIds(this.accounts.list().map((account) => account.id));
+  }
+
+  setEnabled(id: string, enabled: boolean) {
+    this.accounts.setEnabled(id, enabled);
+    if (enabled) this.state.clearHealth(id);
+  }
+
+  setWeight(id: string, weight: number) {
+    this.accounts.setWeight(id, weight);
+  }
+
+  setPolicy(policy: RoutingPolicyName) {
+    this.state.setPolicy(policy);
+  }
+
+  getPolicy(): RoutingPolicyName {
+    return this.state.getPolicy(this.config.policy);
+  }
+
+  async refreshAccount(id: string): Promise<StoredRouterAccount> {
+    const account = this.accounts.get(id);
+    if (!account) throw new Error(`Unknown account: ${id}`);
+    const refreshed = await refreshAccountCredentials(account);
+    this.accounts.update(refreshed);
+    this.state.clearHealth(id);
+    return refreshed;
+  }
+
+  getProviderModels() {
+    return this.getConfig().models.map((model) => ({
+      id: model.id,
+      name: model.name,
+      reasoning: model.reasoning,
+      input: model.input,
+      cost: model.cost,
+      contextWindow: model.contextWindow,
+      maxTokens: model.maxTokens,
+      compat: model.compat,
+    }));
+  }
+
+  getStatusRows(): RouterStatusRow[] {
+    this.reloadConfig();
+    return this.accounts.list().map((account) => {
+      const state = this.state.ensureAccount(account.id);
+      return {
+        id: account.id,
+        label: account.label,
+        upstream: account.upstreamId,
+        provider: account.provider,
+        enabled: account.enabled,
+        weight: account.weight,
+        authHealth: state.authHealth,
+        cooldownUntil: state.cooldownUntil,
+        penaltyUntil: state.penaltyUntil,
+        lastUsedAt: state.lastUsedAt,
+        lastStatus: state.lastStatus,
+        failures: state.failures,
+        rateLimitCount: state.rateLimitCount,
+        authFailureCount: state.authFailureCount,
+        successCount: state.successCount,
+        expires: account.expires,
+      };
+    });
+  }
+
+  private getModelConfig(modelId: string): RouterModelConfig {
+    const model = this.config.models.find((entry) => entry.id === modelId);
+    if (!model) throw new Error(`Model is not configured for oauth-router: ${modelId}`);
+    return model;
+  }
+
+  private getUpstream(upstreamId: string): RouterUpstreamConfig | undefined {
+    return this.config.upstreams.find((upstream) => upstream.id === upstreamId);
+  }
+
+  private getEligibleAccounts(modelId: string, excludedIds: Set<string>): EligibleAccount[] {
+    const modelConfig = this.getModelConfig(modelId);
+    const allowedUpstreams = new Set(modelConfig.route?.upstreamIds ?? []);
+    const hasRouteFilter = allowedUpstreams.size > 0;
+    const currentTime = now();
+
+    return this.accounts
+      .list()
+      .filter((account) => !excludedIds.has(account.id))
+      .map((account) => ({ account, upstream: this.getUpstream(account.upstreamId), state: this.state.ensureAccount(account.id) }))
+      .filter((entry): entry is { account: StoredRouterAccount; upstream: RouterUpstreamConfig; state: ReturnType<RouterStateStore["ensureAccount"]> } => Boolean(entry.upstream))
+      .filter(({ account, upstream, state }) => {
+        if (!account.enabled) return false;
+        if (!upstream.enabled) return false;
+        if (!upstream.modelIds.includes(modelId)) return false;
+        if (hasRouteFilter && !allowedUpstreams.has(upstream.id)) return false;
+        if (state.authHealth === "invalid") return false;
+        if (state.cooldownUntil && state.cooldownUntil > currentTime) return false;
+        if (state.penaltyUntil && state.penaltyUntil > currentTime) return false;
+        return true;
+      })
+      .map(({ account, upstream, state }) => ({ account, upstream, modelConfig, state }));
+  }
+
+  private async prepareSelection(
+    eligible: EligibleAccount,
+    model: Model<Api>,
+    options?: SimpleStreamOptions,
+  ): Promise<DelegateSelection> {
+    let account = eligible.account;
+
+    if (account.provider !== "api-key" && account.expires - this.config.tokenRefreshSkewMs <= now()) {
+      account = await refreshAccountCredentials(account);
+      this.accounts.update(account);
+      this.state.clearHealth(account.id);
+    }
+
+    const apiKey = await getApiKeyForAccount(account);
+    const headers = { ...(eligible.upstream.headers ?? {}) };
+    const delegatedModel = {
+      ...model,
+      id: eligible.modelConfig.id,
+      name: eligible.modelConfig.name,
+      reasoning: eligible.modelConfig.reasoning,
+      input: eligible.modelConfig.input,
+      cost: eligible.modelConfig.cost,
+      contextWindow: eligible.modelConfig.contextWindow,
+      maxTokens: eligible.modelConfig.maxTokens,
+      compat: eligible.modelConfig.compat,
+      api: eligible.upstream.api,
+      baseUrl: eligible.upstream.baseUrl,
+      headers,
+    } as Model<Api>;
+
+    return {
+      account,
+      upstream: eligible.upstream,
+      modelConfig: eligible.modelConfig,
+      apiKey,
+      headers,
+      delegatedModel,
+    };
+  }
+
+  private markFailure(accountId: string, failure: FailureClassification) {
+    switch (failure.kind) {
+      case "rate-limit":
+        this.state.markRateLimit(accountId, failure.retryAfterMs ?? this.config.rateLimitCooldownMs, failure.status, failure.message);
+        break;
+      case "auth":
+        this.state.markAuthFailure(accountId, failure.status, failure.message);
+        break;
+      case "transient":
+        this.state.markTransientFailure(
+          accountId,
+          this.config.transientPenaltyMs,
+          failure.status,
+          failure.message,
+        );
+        break;
+      case "fatal":
+        this.state.markTransientFailure(accountId, this.config.transientPenaltyMs, failure.status ?? 500, failure.message);
+        break;
+    }
+  }
+
+  private async trySingleAccount(
+    selection: DelegateSelection,
+    outerStream: ReturnType<typeof createAssistantMessageEventStream>,
+    context: Context,
+    options?: SimpleStreamOptions,
+  ): Promise<{ completed: boolean; emittedMeaningfulOutput: boolean; failure?: FailureClassification }> {
+    let responseStatus: number | undefined;
+    let responseHeaders: Record<string, string> = {};
+    const buffered: AssistantMessageEvent[] = [];
+    let emittedMeaningfulOutput = false;
+
+    const streamOptions: SimpleStreamOptions = {
+      ...options,
+      apiKey: selection.apiKey,
+      headers: {
+        ...(options?.headers ?? {}),
+        ...selection.headers,
+      },
+      onResponse: async (response, responseModel) => {
+        responseStatus = response.status;
+        responseHeaders = response.headers;
+        await options?.onResponse?.(response, responseModel);
+      },
+      onPayload: options?.onPayload,
+    };
+
+    const inner = delegateStream(selection.upstream, selection.delegatedModel, context, streamOptions);
+
+    for await (const event of inner) {
+      if (event.type === "error") {
+        const failure = classifyFailure(responseStatus, responseHeaders, event.error.errorMessage || "Upstream request failed", this.config);
+        this.markFailure(selection.account.id, failure);
+
+        if (!emittedMeaningfulOutput) {
+          return { completed: false, emittedMeaningfulOutput: false, failure };
+        }
+
+        outerStream.push(event);
+        return { completed: false, emittedMeaningfulOutput: true, failure };
+      }
+
+      if (!emittedMeaningfulOutput && isMeaningfulEvent(event)) {
+        for (const pending of buffered) outerStream.push(pending);
+        buffered.length = 0;
+        emittedMeaningfulOutput = true;
+      }
+
+      if (emittedMeaningfulOutput) {
+        outerStream.push(event);
+      } else {
+        buffered.push(event);
+      }
+
+      if (event.type === "done") {
+        if (!emittedMeaningfulOutput) {
+          for (const pending of buffered) outerStream.push(pending);
+          buffered.length = 0;
+        }
+        this.state.markSuccess(selection.account.id, responseStatus ?? 200);
+        return { completed: true, emittedMeaningfulOutput };
+      }
+    }
+
+    return {
+      completed: false,
+      emittedMeaningfulOutput,
+      failure: {
+        kind: "transient",
+        status: responseStatus,
+        message: "Upstream stream ended unexpectedly",
+      },
+    };
+  }
+
+  stream(model: Model<Api>, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream {
+    this.reloadConfig();
+    const outer = createAssistantMessageEventStream();
+
+    (async () => {
+      const tried = new Set<string>();
+      let lastFailure: FailureClassification | undefined;
+      let emittedMeaningfulOutput = false;
+
+      try {
+        while (true) {
+          const eligible = this.getEligibleAccounts(model.id, tried);
+          const policy = this.getPolicy();
+          const cursor = this.state.getCursor(policy);
+          const picked = chooseEligibleAccount(policy, eligible, cursor);
+
+          if (!picked.selected) {
+            const message = lastFailure
+              ? `No healthy oauth-router accounts remaining after failover. Last error: ${lastFailure.message}`
+              : `No healthy oauth-router accounts are configured for model ${model.id}. Add accounts with /router-login add.`;
+            outer.push({ type: "error", reason: "error", error: createErrorMessage({ model, message }) });
+            outer.end();
+            return;
+          }
+
+          this.state.advanceCursor(policy, picked.nextCursor);
+          tried.add(picked.selected.account.id);
+          this.state.markAttempt(picked.selected.account.id, model.id);
+
+          let selection: DelegateSelection;
+          try {
+            selection = await this.prepareSelection(picked.selected, model, options);
+          } catch (error) {
+            const message = error instanceof Error ? error.message : String(error);
+            lastFailure = { kind: "auth", status: 401, message };
+            this.markFailure(picked.selected.account.id, lastFailure);
+            continue;
+          }
+
+          const result = await this.trySingleAccount(selection, outer, context, options);
+          if (result.completed) {
+            outer.end();
+            return;
+          }
+
+          emittedMeaningfulOutput ||= result.emittedMeaningfulOutput;
+          lastFailure = result.failure;
+
+          if (emittedMeaningfulOutput) {
+            outer.end();
+            return;
+          }
+        }
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        outer.push({ type: "error", reason: "error", error: createErrorMessage({ model, message }) });
+        outer.end();
+      }
+    })();
+
+    return outer;
+  }
+}
+
+export function registerRouterProvider(pi: ExtensionAPI, runtime: RouterRuntime) {
+  const config = runtime.getConfig();
+  pi.registerProvider(config.providerName, {
+    baseUrl: "https://oauth-router.local",
+    apiKey: "OAUTH_ROUTER_DISABLED",
+    api: "oauth-router-api",
+    models: runtime.getProviderModels(),
+    streamSimple: (model, context, options) => runtime.stream(model, context, options),
+  });
+}

package/.pi/extensions/oauth-router/scripts/vibe-verify.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+REQUIRED = [
+    ROOT / "index.ts",
+    ROOT / "provider.ts",
+    ROOT / "commands.ts",
+    ROOT / "oauth-flow.ts",
+    ROOT / "oauth-store.ts",
+    ROOT / "state.ts",
+    ROOT / "policies.ts",
+    ROOT / "config.ts",
+    ROOT / "types.ts",
+    ROOT / "README.md",
+    ROOT / "docs" / "Project_Requirements.md",
+    ROOT / "docs" / "Coding_Guidelines.md",
+    ROOT / "docs" / "Builder_Prompt.md",
+    ROOT / "docs" / "issues" / "FR-001.md",
+]
+EXPECTED_MODELS = ["oauth-router", "gpt-4o", "gpt-4.1", "o4-mini", "gpt-5.4"]
+PI_CANDIDATES = [
+    os.environ.get("PI_BIN"),
+    shutil.which("pi"),
+    shutil.which("pi.cmd"),
+    str(Path.home() / "AppData" / "Roaming" / "npm" / "pi"),
+    str(Path.home() / "AppData" / "Roaming" / "npm" / "pi.cmd"),
+]
+
+
+def resolve_pi() -> str | None:
+    for candidate in PI_CANDIDATES:
+        if not candidate:
+            continue
+        path = Path(candidate)
+        if path.exists():
+            return str(path)
+    return None
+
+
+def run(cmd: list[str], timeout: int = 90) -> tuple[int, str, str]:
+    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+    return proc.returncode, proc.stdout, proc.stderr
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--quick", action="store_true")
+    args = parser.parse_args()
+
+    failed = False
+
+    for path in REQUIRED:
+        if not path.exists():
+            print(f"[FAIL] missing: {path}")
+            failed = True
+        else:
+            print(f"[ OK ] found: {path.relative_to(ROOT)}")
+
+    if args.quick:
+        return 1 if failed else 0
+
+    pi_bin = resolve_pi()
+    if not pi_bin:
+        print("[FAIL] unable to locate `pi` binary")
+        return 1
+
+    try:
+        code, stdout, stderr = run([pi_bin, "--list-models"])
+        output = f"{stdout}\n{stderr}"
+        if code != 0:
+            print("[FAIL] `pi --list-models` failed")
+            print(output.strip())
+            failed = True
+        else:
+            print("[ OK ] `pi --list-models` executed")
+            for token in EXPECTED_MODELS:
+                if token not in output:
+                    print(f"[FAIL] model token not found in list output: {token}")
+                    failed = True
+                else:
+                    print(f"[ OK ] model token present: {token}")
+    except Exception as exc:
+        print(f"[FAIL] unable to execute `pi --list-models`: {exc}")
+        failed = True
+
+    return 1 if failed else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())