@ilya-lesikov/pi-pi 0.6.0 → 0.7.0

package/extensions/orchestrator/flant-infra.ts

@@ -2,6 +2,7 @@ import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { homedir } from "node:os";
 import { join } from "node:path";
 import lockfile from "proper-lockfile";
+import { refreshAnthropicToken } from "@earendil-works/pi-ai/oauth";
 import type { ExtensionAPI, ProviderModelConfig } from "@earendil-works/pi-coding-agent";
 import type { PiPiConfig } from "./config.js";
 import { updateRegistryFromAvailableModels } from "./model-registry.js";
@@ -28,6 +29,12 @@ export interface FlantSettings {
   lastUpdated: string | null;
   cachedFlantModels: string[] | null;
   cachedOpenRouterData: Record<string, OpenRouterModelData> | null;
+  /**
+   * When true, additionally register the `pp-flant-anthropic-sub` provider,
+   * which routes Claude requests through the gateway using the user's personal
+   * Claude OAuth token (billed against their personal Claude subscription).
+   */
+  subscription: boolean;
 }
 
 const GEMINI_MAP: Record<string, string> = {
@@ -56,11 +63,153 @@ const DEFAULT_SETTINGS: FlantSettings = {
   lastUpdated: null,
   cachedFlantModels: null,
   cachedOpenRouterData: null,
+  subscription: false,
 };
 
+/** Provider name for the personal-subscription Claude routing. */
+export const SUB_PROVIDER = "pp-flant-anthropic-sub";
+
+/** Prefix the gateway expects for personal-subscription Claude models. */
+export const SUB_MODEL_PREFIX = "sub/";
+
+/**
+ * Read the Claude OAuth access token persisted by pi's built-in `anthropic`
+ * OAuth provider. Returns null when absent or expired. The gateway uses this
+ * token (forwarded as `Authorization: Bearer ...`) to bill the user's personal
+ * Claude subscription; pi-ai automatically adds the Claude Code identity
+ * headers because the token has the `sk-ant-oat` prefix.
+ */
+export function readClaudeOAuthToken(): string | null {
+  const authPath = join(resolveAgentDir(), "auth.json");
+  if (!existsSync(authPath)) return null;
+  try {
+    const raw = JSON.parse(readFileSync(authPath, "utf-8")) as {
+      anthropic?: { access?: unknown; expires?: unknown };
+    };
+    const anthropic = raw.anthropic;
+    if (!anthropic || typeof anthropic.access !== "string" || !anthropic.access) return null;
+    if (typeof anthropic.expires === "number" && anthropic.expires <= Date.now()) return null;
+    return anthropic.access;
+  } catch {
+    return null;
+  }
+}
+
+interface AnthropicOAuthCreds {
+  type?: unknown;
+  access?: unknown;
+  refresh?: unknown;
+  expires?: unknown;
+}
+
+/**
+ * Ensure the persisted Claude OAuth access token is fresh, refreshing it via
+ * the stored refresh token when it has expired (or is within its safety
+ * margin). The refreshed credentials are written back to `auth.json` under the
+ * `anthropic` provider id in pi's own `{ type: "oauth", ... }` format, so both
+ * this extension and pi's built-in `anthropic` provider observe the new token.
+ *
+ * Unlike readClaudeOAuthToken, this is async (a refresh performs a network
+ * call) and returns the valid access token, or null when no usable
+ * credentials exist / a refresh fails. Async entry points call this before the
+ * synchronous readClaudeOAuthToken so downstream reads see a fresh token.
+ */
+export async function refreshClaudeOAuthToken(): Promise<string | null> {
+  const log = getLogger();
+  const authPath = join(resolveAgentDir(), "auth.json");
+  if (!existsSync(authPath)) return null;
+
+  let anthropic: AnthropicOAuthCreds | undefined;
+  try {
+    const raw = JSON.parse(readFileSync(authPath, "utf-8")) as { anthropic?: AnthropicOAuthCreds };
+    anthropic = raw.anthropic;
+  } catch {
+    return null;
+  }
+  if (!anthropic || typeof anthropic.access !== "string" || !anthropic.access) return null;
+
+  const expires = typeof anthropic.expires === "number" ? anthropic.expires : 0;
+  if (expires > Date.now()) return anthropic.access;
+
+  // Expired (or no expiry recorded): try to refresh.
+  if (typeof anthropic.refresh !== "string" || !anthropic.refresh) {
+    log.debug({ s: "flant" }, "claude oauth token expired and no refresh token available");
+    return null;
+  }
+
+  let refreshed: { refresh: string; access: string; expires: number };
+  try {
+    refreshed = await refreshAnthropicToken(anthropic.refresh);
+  } catch (err: any) {
+    log.debug({ s: "flant", err: err?.message }, "claude oauth token refresh failed");
+    return null;
+  }
+
+  // Persist refreshed credentials back under the `anthropic` provider id, using
+  // pi's { type: "oauth", ... } shape and the same file lock pi uses.
+  try {
+    const authDir = resolveAgentDir();
+    if (!existsSync(authDir)) mkdirSync(authDir, { recursive: true });
+    if (!existsSync(authPath)) writeFileSync(authPath, "{}\n", "utf-8");
+    const release = lockfile.lockSync(authPath, { stale: 10000 });
+    try {
+      let current: Record<string, unknown> = {};
+      try {
+        current = JSON.parse(readFileSync(authPath, "utf-8")) as Record<string, unknown>;
+      } catch {
+        current = {};
+      }
+      const existing = (current.anthropic && typeof current.anthropic === "object")
+        ? current.anthropic as Record<string, unknown>
+        : {};
+      // Another instance may have refreshed while we were waiting for the lock.
+      const existingExpires = typeof existing.expires === "number" ? existing.expires : 0;
+      if (existingExpires > Date.now() && typeof existing.access === "string" && existing.access) {
+        return existing.access;
+      }
+      current.anthropic = {
+        ...existing,
+        type: "oauth",
+        access: refreshed.access,
+        refresh: refreshed.refresh,
+        expires: refreshed.expires,
+      };
+      writeFileSync(authPath, JSON.stringify(current, null, 2) + "\n", "utf-8");
+    } finally {
+      release();
+    }
+  } catch (err: any) {
+    // If persistence fails we still return the freshly minted token so the
+    // current run can proceed; the next run will refresh again.
+    log.debug({ s: "flant", err: err?.message }, "failed to persist refreshed claude oauth token");
+  }
+
+  log.debug({ s: "flant" }, "refreshed claude oauth token");
+  return refreshed.access;
+}
+
+/** Resolve the gateway API key (LLM_API_KEY preferred, FLANT_API_KEY fallback). */
+export function readGatewayApiKey(): string | null {
+  return process.env.LLM_API_KEY ?? process.env.FLANT_API_KEY ?? null;
+}
+
 let piRef: ExtensionAPI | null = null;
 let generatedFlantConfig: Partial<PiPiConfig> | null = null;
 
+/**
+ * Inputs needed to (re)register the personal-subscription Claude provider.
+ * Cached at the last registerFlantProviders call so refreshSubProvider can
+ * rebuild the provider with a freshly-read OAuth token without re-discovering
+ * models. Null when subscription routing is not active.
+ */
+interface SubProviderContext {
+  anthropicModels: string[];
+  metadata: Record<string, OpenRouterModelData>;
+}
+let subProviderContext: SubProviderContext | null = null;
+/** The OAuth access token the sub provider was last registered with. */
+let lastSubToken: string | null = null;
+
 export function setPI(pi: ExtensionAPI): void {
   piRef = pi;
 }
@@ -74,6 +223,7 @@ export function unregisterFlantProviders(pi?: ExtensionAPI): void {
   if (!api) return;
   api.unregisterProvider("pp-flant-anthropic");
   api.unregisterProvider("pp-flant-openai");
+  api.unregisterProvider(SUB_PROVIDER);
 }
 
 function ensureSettingsDir(): void {
@@ -99,6 +249,7 @@ function normalizeSettings(raw: unknown): FlantSettings {
     enabled: !!value.enabled,
     autoUpdate: value.autoUpdate === undefined ? true : !!value.autoUpdate,
     cacheTTLDays,
+    subscription: !!value.subscription,
     lastUpdated: typeof value.lastUpdated === "string" ? value.lastUpdated : null,
     cachedFlantModels: Array.isArray(value.cachedFlantModels)
       ? value.cachedFlantModels.filter((m): m is string => typeof m === "string")
@@ -262,9 +413,24 @@ export async function fetchOpenRouterMetadata(modelIds: string[]): Promise<Recor
   return out;
 }
 
-function modelSpec(modelId: string): string {
-  const provider = modelId.startsWith("claude-") ? "pp-flant-anthropic" : "pp-flant-openai";
-  return `${provider}/${modelId}`;
+/**
+ * Returns true when the personal-subscription Claude path is fully usable:
+ * the setting is enabled AND both credentials (Claude OAuth token + gateway
+ * key) are present. Mirrors the gate in registerFlantProviders so we never
+ * generate `sub/` role assignments that cannot resolve to a real provider.
+ */
+export function isSubscriptionActive(settings?: FlantSettings): boolean {
+  const s = settings ?? loadFlantSettings();
+  return s.subscription && !!readClaudeOAuthToken() && !!readGatewayApiKey();
+}
+
+function modelSpec(modelId: string, subscriptionActive = false): string {
+  if (modelId.startsWith("claude-")) {
+    return subscriptionActive
+      ? `${SUB_PROVIDER}/${SUB_MODEL_PREFIX}${modelId}`
+      : `pp-flant-anthropic/${modelId}`;
+  }
+  return `pp-flant-openai/${modelId}`;
 }
 
 function buildProviderModelConfig(
@@ -289,16 +455,21 @@ function buildProviderModelConfig(
   };
 }
 
+export interface RegisterFlantOptions {
+  /** Whether to also register the personal-subscription provider. Defaults to loadFlantSettings().subscription. */
+  subscription?: boolean;
+}
+
 export function registerFlantProviders(
   pi: ExtensionAPI,
   models: string[],
   metadata: Record<string, OpenRouterModelData>,
+  options: RegisterFlantOptions = {},
 ): void {
   const log = getLogger();
   const uniqueModels = [...new Set(models)];
   const anthropicModels = uniqueModels.filter((m) => m.startsWith("claude-"));
   const openaiModels = uniqueModels.filter((m) => !m.startsWith("claude-"));
-  log.debug({ s: "flant", total: uniqueModels.length, anthropic: anthropicModels.length, openai: openaiModels.length }, "registering flant providers");
 
   unregisterFlantProviders(pi);
 
@@ -316,10 +487,97 @@ export function registerFlantProviders(
     models: openaiModels.map((m) => buildProviderModelConfig(m, metadata)),
   });
 
-  updateRegistryFromAvailableModels([
+  const availableSpecs = [
     ...anthropicModels.map((id) => `pp-flant-anthropic/${id}`),
     ...openaiModels.map((id) => `pp-flant-openai/${id}`),
-  ]);
+  ];
+
+  const subscription = options.subscription ?? loadFlantSettings().subscription;
+  let subModels: string[] = [];
+  if (subscription) {
+    // Remember the models/metadata so refreshSubProvider can rebuild the
+    // provider with a fresh OAuth token on each turn (see below).
+    subProviderContext = { anthropicModels, metadata };
+    subModels = registerSubProvider(pi, anthropicModels, metadata);
+    availableSpecs.push(...subModels.map((id) => `${SUB_PROVIDER}/${id}`));
+  } else {
+    subProviderContext = null;
+    lastSubToken = null;
+  }
+
+  log.debug({ s: "flant", total: uniqueModels.length, anthropic: anthropicModels.length, openai: openaiModels.length, sub: subModels.length }, "registering flant providers");
+
+  updateRegistryFromAvailableModels(availableSpecs);
+}
+
+/**
+ * Register (or re-register) the personal-subscription Claude provider using the
+ * OAuth access token currently persisted in auth.json. Returns the list of
+ * `sub/`-prefixed model ids registered (empty when credentials are missing).
+ *
+ * The provider is registered with a literal `apiKey` (the resolved token), so
+ * the value is a static snapshot for the lifetime of the registration. Because
+ * the OAuth token expires within a few hours, refreshSubProvider must be called
+ * periodically (on each turn) to rebuild the provider with a fresh token;
+ * otherwise long-lived sessions send a stale token and the gateway returns 401.
+ */
+function registerSubProvider(
+  pi: ExtensionAPI,
+  anthropicModels: string[],
+  metadata: Record<string, OpenRouterModelData>,
+): string[] {
+  const log = getLogger();
+  const oauthToken = readClaudeOAuthToken();
+  const gatewayKey = readGatewayApiKey();
+  if (!oauthToken || !gatewayKey) {
+    log.debug({ s: "flant", hasOAuth: !!oauthToken, hasGatewayKey: !!gatewayKey }, "subscription enabled but credentials missing; skipping sub provider");
+    lastSubToken = null;
+    return [];
+  }
+  pi.registerProvider(SUB_PROVIDER, {
+    name: "Flant (personal Claude subscription)",
+    api: "anthropic-messages",
+    baseUrl: "https://llm-api.flant.ru",
+    // The OAuth token (sk-ant-oat...) triggers pi-ai's Claude Code identity
+    // headers and is forwarded as `Authorization: Bearer ...`.
+    apiKey: oauthToken,
+    // Gateway key travels in a side header so it does not clobber the OAuth auth.
+    headers: { "x-litellm-api-key": `Bearer ${gatewayKey}` },
+    // Model id carries the `sub/` prefix the gateway expects, while pricing/
+    // metadata is looked up by the bare claude-* id.
+    models: anthropicModels.map((m) => {
+      const cfg = buildProviderModelConfig(m, metadata);
+      return { ...cfg, id: `${SUB_MODEL_PREFIX}${m}` };
+    }),
+  });
+  lastSubToken = oauthToken;
+  return anthropicModels.map((m) => `${SUB_MODEL_PREFIX}${m}`);
+}
+
+/**
+ * Ensure the personal-subscription Claude provider is registered with a
+ * non-expired OAuth token. Refreshes the token (persisting it to auth.json when
+ * needed) and, when it changed since the last registration, re-registers the
+ * provider so subsequent LLM calls pick up the fresh token.
+ *
+ * Called on each turn for the root session. No-op when subscription routing is
+ * not active. Cheap when the token is unchanged (only a token read + compare).
+ */
+export async function refreshSubProvider(pi?: ExtensionAPI): Promise<void> {
+  const ctx = subProviderContext;
+  if (!ctx) return;
+  const api = pi ?? piRef;
+  if (!api) return;
+
+  await refreshClaudeOAuthToken();
+  const token = readClaudeOAuthToken();
+  // Only re-register when the token actually changed; registerProvider takes
+  // effect immediately, so re-registering every turn would be wasteful churn.
+  if (token && token === lastSubToken) return;
+
+  const log = getLogger();
+  registerSubProvider(api, ctx.anthropicModels, ctx.metadata);
+  log.debug({ s: "flant", changed: token !== lastSubToken }, "refreshed sub provider oauth token");
 }
 
 function pickLatest(models: string[]): string | null {
@@ -339,18 +597,19 @@ function pickCheapestFastModel(models: string[]): string | null {
   return pickLatest(models.filter((m) => /^claude-haiku-/.test(m)));
 }
 
-function makeVariant(modelId: string | null, fallbackModelId: string): { enabled: boolean; model: string; thinking: string } {
-  if (!modelId) return { enabled: false, model: modelSpec(fallbackModelId), thinking: "high" };
-  return { enabled: true, model: modelSpec(modelId), thinking: "high" };
+function makeVariant(modelId: string | null, fallbackModelId: string, sub = false): { enabled: boolean; model: string; thinking: string } {
+  if (!modelId) return { enabled: false, model: modelSpec(fallbackModelId, sub), thinking: "high" };
+  return { enabled: true, model: modelSpec(modelId, sub), thinking: "high" };
 }
 
 function makeVariantWithThinking(
   modelId: string | null,
   fallbackModelId: string,
   thinking: string,
+  sub = false,
 ): { enabled: boolean; model: string; thinking: string } {
-  if (!modelId) return { enabled: false, model: modelSpec(fallbackModelId), thinking };
-  return { enabled: true, model: modelSpec(modelId), thinking };
+  if (!modelId) return { enabled: false, model: modelSpec(fallbackModelId, sub), thinking };
+  return { enabled: true, model: modelSpec(modelId, sub), thinking };
 }
 
 function buildPresetGroup(
@@ -360,9 +619,10 @@ function buildPresetGroup(
   return { default: defaultPreset, presets };
 }
 
-export function generateFlantConfig(models: string[]): Partial<PiPiConfig> {
+export function generateFlantConfig(models: string[], subscriptionActive = false): Partial<PiPiConfig> {
   const uniqueModels = [...new Set(models)];
   if (uniqueModels.length === 0) return {};
+  const sub = subscriptionActive;
 
   const latestOpus = pickLatest(uniqueModels.filter((m) => /^claude-opus-/.test(m)));
   const latestClaude = pickLatest(uniqueModels.filter((m) => /^claude-/.test(m)));
@@ -383,73 +643,74 @@ export function generateFlantConfig(models: string[]): Partial<PiPiConfig> {
   return {
     agents: {
       orchestrators: {
-        implement: { model: modelSpec(implementModel), thinking: "high" },
-        plan: { model: modelSpec(implementModel), thinking: "high" },
-        debug: { model: modelSpec(debugModel), thinking: "high" },
-        brainstorm: { model: modelSpec(brainstormModel), thinking: "high" },
-        review: { model: modelSpec(implementModel), thinking: "high" },
+        implement: { model: modelSpec(implementModel, sub), thinking: "high" },
+        plan: { model: modelSpec(implementModel, sub), thinking: "high" },
+        debug: { model: modelSpec(debugModel, sub), thinking: "high" },
+        brainstorm: { model: modelSpec(brainstormModel, sub), thinking: "high" },
+        review: { model: modelSpec(implementModel, sub), thinking: "high" },
+        quick: { model: modelSpec(implementModel, sub), thinking: "high" },
       },
       subagents: {
         simple: {
-          explore: { model: modelSpec(fastModel), thinking: "low" },
-          librarian: { model: modelSpec(fastModel), thinking: "medium" },
-          task: { model: modelSpec(taskModel), thinking: "medium" },
+          explore: { model: modelSpec(fastModel, sub), thinking: "low" },
+          librarian: { model: modelSpec(fastModel, sub), thinking: "medium" },
+          task: { model: modelSpec(taskModel, sub), thinking: "medium" },
         },
         presetGroups: {
           planners: buildPresetGroup({
             regular: {
               agents: {
-                opus: makeVariant(latestOpus, fallback),
-                gpt: makeVariant(latestGpt, fallback),
-                gemini: makeVariant(latestGeminiPro, fallback),
+                opus: makeVariant(latestOpus, fallback, sub),
+                gpt: makeVariant(latestGpt, fallback, sub),
+                gemini: makeVariant(latestGeminiPro, fallback, sub),
               },
             },
           }),
           planReviewers: buildPresetGroup({
             regular: {
               agents: {
-                opus: makeVariant(latestOpus, fallback),
-                gpt: makeVariant(latestGpt, fallback),
-                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh"),
+                opus: makeVariant(latestOpus, fallback, sub),
+                gpt: makeVariant(latestGpt, fallback, sub),
+                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh", sub),
               },
             },
             deep: {
               agents: {
-                opus: makeVariantWithThinking(latestOpus, fallback, "xhigh"),
-                gpt: makeVariantWithThinking(latestGpt, fallback, "xhigh"),
-                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh"),
+                opus: makeVariantWithThinking(latestOpus, fallback, "xhigh", sub),
+                gpt: makeVariantWithThinking(latestGpt, fallback, "xhigh", sub),
+                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh", sub),
               },
             },
           }),
           codeReviewers: buildPresetGroup({
             regular: {
               agents: {
-                opus: makeVariant(latestOpus, fallback),
-                gpt: makeVariant(latestGpt, fallback),
-                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh"),
+                opus: makeVariant(latestOpus, fallback, sub),
+                gpt: makeVariant(latestGpt, fallback, sub),
+                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh", sub),
               },
             },
             deep: {
               agents: {
-                opus: makeVariantWithThinking(latestOpus, fallback, "xhigh"),
-                gpt: makeVariantWithThinking(latestGpt, fallback, "xhigh"),
-                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh"),
+                opus: makeVariantWithThinking(latestOpus, fallback, "xhigh", sub),
+                gpt: makeVariantWithThinking(latestGpt, fallback, "xhigh", sub),
+                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh", sub),
               },
             },
           }),
           brainstormReviewers: buildPresetGroup({
             regular: {
               agents: {
-                opus: makeVariant(latestOpus, fallback),
-                gpt: makeVariant(latestGpt, fallback),
-                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh"),
+                opus: makeVariant(latestOpus, fallback, sub),
+                gpt: makeVariant(latestGpt, fallback, sub),
+                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh", sub),
               },
             },
             deep: {
               agents: {
-                opus: makeVariantWithThinking(latestOpus, fallback, "xhigh"),
-                gpt: makeVariantWithThinking(latestGpt, fallback, "xhigh"),
-                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh"),
+                opus: makeVariantWithThinking(latestOpus, fallback, "xhigh", sub),
+                gpt: makeVariantWithThinking(latestGpt, fallback, "xhigh", sub),
+                gemini: makeVariantWithThinking(latestGeminiPro, fallback, "xhigh", sub),
               },
             },
           }),
@@ -479,6 +740,12 @@ export async function updateFlantInfra(
   setPI(pi);
   const settings = loadFlantSettings();
 
+  // Refresh the personal-subscription Claude OAuth token before (re)registering
+  // providers so the sub provider is built with a valid, non-expired token.
+  if (settings.subscription) {
+    await refreshClaudeOAuthToken();
+  }
+
   let models = isCacheValid(settings) ? settings.cachedFlantModels : null;
   let metadata = isCacheValid(settings) ? settings.cachedOpenRouterData : null;
   let refreshed = false;
@@ -522,7 +789,7 @@ export async function updateFlantInfra(
 
   try {
     registerFlantProviders(pi, models, metadata);
-    generatedFlantConfig = generateFlantConfig(models);
+    generatedFlantConfig = generateFlantConfig(models, isSubscriptionActive(settings));
     if (!refreshed && settings.cachedFlantModels && settings.cachedOpenRouterData && !settings.lastUpdated) {
       settings.lastUpdated = new Date().toISOString();
       saveFlantSettings(settings);
@@ -544,7 +811,7 @@ export function initFlantSync(pi: ExtensionAPI): void {
   }
   if (settings.cachedFlantModels && settings.cachedOpenRouterData) {
     registerFlantProviders(pi, settings.cachedFlantModels, settings.cachedOpenRouterData);
-    generatedFlantConfig = generateFlantConfig(settings.cachedFlantModels);
+    generatedFlantConfig = generateFlantConfig(settings.cachedFlantModels, isSubscriptionActive(settings));
   }
 }
 
@@ -555,6 +822,11 @@ export async function initFlantOnStartup(pi: ExtensionAPI): Promise<void> {
     generatedFlantConfig = null;
     return;
   }
-  if (!settings.autoUpdate) return;
+  if (!settings.autoUpdate) {
+    // updateFlantInfra (which refreshes the token) is skipped when auto-update
+    // is off, but any registered sub provider still needs a fresh token.
+    if (settings.subscription) await refreshClaudeOAuthToken();
+    return;
+  }
   await updateFlantInfra(pi);
 }