@narumitw/pi-usage 0.61.1 → 0.62.0

package/src/usage-settings-ui.ts

@@ -9,7 +9,7 @@ const ON = "On";
 const REMAINING = "Remaining";
 const USED = "Used";
 
-type UsageSettingId = "codexFastMode" | "codexStatusResetCountdown" | "codexStatusPercentage";
+type UsageSettingId = "codexFastMode" | "codexStatusResetCountdown" | "codexStatusPercentage" | "openaiCompanionUsage";
 
 function settingValueLabel(id: UsageSettingId, value: UsageSettings[UsageSettingId]): string {
   if (id === "codexStatusPercentage") return value === "used" ? USED : REMAINING;
@@ -29,11 +29,12 @@ export async function showUsageSettings(
   }
   const { HorizontalRule, renderBoundedFrame } = await import("@narumitw/pi-tui-kit");
   if (parentSignal.aborted || !isCurrent()) return false;
-  return (
-    (await ctx.ui.custom<boolean>((tui, theme, _keybindings, done) => {
+  let changed = false;
+  let selectedId: UsageSettingId | undefined;
+  for (;;) {
+    const result = await ctx.ui.custom<boolean | "consent">((tui, theme, _keybindings, done) => {
       const localController = new AbortController();
       const signal = AbortSignal.any([parentSignal, localController.signal]);
-      let changed = false;
       let closing = false;
       let saveQueue = Promise.resolve();
       const state = settingsRuntime.get();
@@ -59,18 +60,28 @@ export async function showUsageSettings(
           currentValue: settingValueLabel("codexStatusPercentage", state.settings.codexStatusPercentage),
           values: [REMAINING, USED],
         },
+        {
+          id: "openaiCompanionUsage",
+          label: "Experimental ChatGPT companion usage",
+          description:
+            "Read undocumented plan/app usage using a same-account Codex login; matching registration does not prove identity.",
+          currentValue: state.settings.openaiCompanionUsage ? ON : OFF,
+          values: [OFF, ON],
+        },
       ];
       const rule = new HorizontalRule({ ruleStyle: (text) => theme.fg("border", text) });
 
       let settingsList: SettingsList;
-      const cancel = () => {
+      const close = (result: boolean | "consent") => {
         if (closing) return;
         closing = true;
         localController.abort();
-        done(changed);
+        done(result);
       };
+      const cancel = () => close(changed);
       const queueUpdate = <Id extends UsageSettingId>(id: Id, requested: UsageSettings[Id], display: string) => {
         saveQueue = saveQueue.then(async () => {
+          if (signal.aborted || !isCurrent()) return;
           const previous = settingsRuntime.get().settings[id];
           if (settingsRuntime.get().kind === "invalid") {
             settingsList.updateValue(id, settingValueLabel(id, previous));
@@ -81,6 +92,14 @@ export async function showUsageSettings(
             return;
           }
           try {
+            if (id === "openaiCompanionUsage" && requested === true && previous !== true) {
+              // Pi's confirm dialog replaces a non-overlay custom screen and restores
+              // the editor, not that screen. Complete/dispose Settings first; after
+              // consent, mount a fresh instance with the initiating row selected.
+              selectedId = id;
+              close("consent");
+              return;
+            }
             const patch: Partial<UsageSettings> = {};
             patch[id] = requested;
             await settingsRuntime.update(patch, signal);
@@ -91,6 +110,8 @@ export async function showUsageSettings(
             tui.requestRender();
             return;
           }
+          // A rename can win a cancellation race. Apply owned lifecycle cleanup after a
+          // durable save, but let the owner guard its session before touching UI.
           if (previous !== requested) {
             changed = true;
             onApplied(id);
@@ -106,15 +127,17 @@ export async function showUsageSettings(
         getSettingsListTheme(),
         (id, value) => {
           if (closing || signal.aborted || !isCurrent()) return;
+          selectedId = id as UsageSettingId;
           if (id === "codexStatusPercentage") {
             queueUpdate(id, value === USED ? "used" : "remaining", value);
-          } else if (id === "codexFastMode" || id === "codexStatusResetCountdown") {
+          } else if (id === "codexFastMode" || id === "codexStatusResetCountdown" || id === "openaiCompanionUsage") {
             queueUpdate(id, value === ON, value);
           }
         },
         cancel,
       );
 
+      if (selectedId) settingsList.selectItem(selectedId);
       parentSignal.addEventListener("abort", cancel, { once: true });
       return {
         render(width: number) {
@@ -147,6 +170,25 @@ export async function showUsageSettings(
           parentSignal.removeEventListener("abort", cancel);
         },
       };
-    })) ?? false
-  );
+    });
+    if (result !== "consent" || parentSignal.aborted || !isCurrent()) return changed;
+    try {
+      const accepted = await ctx.ui.confirm(
+        "Enable experimental ChatGPT companion usage?",
+        "Requires /login openai-codex with the same ChatGPT account/workspace as native OpenAI. Only the companion token and its matching account ID are sent to undocumented ChatGPT usage endpoints. Registration matching is not independent proof of the same user/workspace. No reset or allowance mutations are performed.",
+        { signal: parentSignal },
+      );
+      if (parentSignal.aborted || !isCurrent()) return changed;
+      if (accepted && !settingsRuntime.get().settings.openaiCompanionUsage) {
+        await settingsRuntime.update({ openaiCompanionUsage: true }, parentSignal);
+        changed = true;
+        // A durable rename may win cancellation; the owner guards lifecycle cleanup.
+        onApplied("openaiCompanionUsage");
+      }
+    } catch (error) {
+      if (parentSignal.aborted || !isCurrent()) return changed;
+      ctx.ui.notify(`Could not enable ChatGPT companion usage: ${errorMessage(error)}`, "error");
+    }
+    if (parentSignal.aborted || !isCurrent()) return changed;
+  }
 }

package/src/usage.ts

@@ -21,7 +21,13 @@ import {
 import { abortError, awaitWithDeadline, errorMessage, runWithConcurrency, UsageCache } from "./core.js";
 import { formatProviderStates, formatUsageStatusline } from "./format.js";
 import { createOAuthCredentialCandidateReader } from "./oauth-credential-source.js";
-import { adapterForProvider, isStaleExtensionContextError, queryProviderUsage, resolveUsageAuth } from "./query.js";
+import { UnsupportedOpenAIUsageAuthError } from "./providers/openai-chatgpt.js";
+import {
+  adapterForProvider,
+  isStaleExtensionContextError,
+  queryProviderUsage,
+  resolveUsageAuth as resolveRuntimeUsageAuth,
+} from "./query.js";
 import { createUsageSettingsRuntime, type UsageSettingsRuntime, type UsageSettingsState } from "./settings.js";
 import type {
   PiModel,
@@ -65,7 +71,7 @@ type UsageExtensionDependencies = {
 type QueryOutcome = {
   state: ProviderUsageState;
   fingerprint?: string;
-  authState?: "unavailable";
+  authState?: "unavailable" | "unsupported" | "failed";
   rememberedTargetId?: string;
 };
 
@@ -83,10 +89,20 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
   const credentialCandidates = createOAuthCredentialCandidateReader(pi, credentialReader);
   const createRedemptionId = dependencies.createRedemptionId ?? randomUUID;
   const settingsRuntime = dependencies.settingsRuntime ?? createUsageSettingsRuntime();
+  const resolveUsageAuth: typeof resolveRuntimeUsageAuth = (ctx, adapter, salt, reader, candidates) =>
+    resolveRuntimeUsageAuth(
+      ctx,
+      adapter,
+      salt,
+      reader,
+      candidates,
+      settingsRuntime.get().settings.openaiCompanionUsage,
+    );
   const cache = new UsageCache(CACHE_TTL_MS);
   const failureBackoff = new Map<string, { until: number; message: string }>();
   const latestQueries = new Map<string, number>();
   const activeControllers = new Set<AbortController>();
+  const nativeQueryControllers = new Set<AbortController>();
   let querySequence = 0;
   let activeCurrentIdentity: string | undefined;
   let sessionActive = false;
@@ -182,7 +198,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
     if (shouldSchedule && sessionActive) scheduleStatusRefresh(ctx, model);
     if (
       sessionActive &&
-      showCodexResetCountdown &&
+      (showCodexResetCountdown || outcome.state.report.source === "openai-chatgpt-companion") &&
       outcome.state.report.buckets.some(
         (bucket) => bucket.resetsAt !== undefined && Number.isFinite(bucket.resetsAt) && bucket.resetsAt * 1_000 > now,
       )
@@ -219,7 +235,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
     activeCurrentIdentity = nextIdentity;
   };
 
-  const queryAdapterState = async (
+  const queryAdapterStateInternal = async (
     ctx: ExtensionContext,
     adapter: UsageProviderAdapter,
     displayState: UsageDisplayState,
@@ -232,6 +248,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
     const expectedSessionId = ctx.sessionManager.getSessionId();
     const expectedModelIdentity = modelIdentity(ctx.model);
     const expectedTargetId = adapter.targets ? settingsRuntime.get().settings.selectedTargets[adapter.id] : undefined;
+    const expectedCompanionUsage = settingsRuntime.get().settings.openaiCompanionUsage;
     const providerName = providerDisplayName(ctx, adapter.id);
     let auth: ResolvedUsageAuth | undefined;
     try {
@@ -243,6 +260,13 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
       );
     } catch (error) {
       if (isStaleExtensionContextError(error) || isAbortError(error)) throw error;
+      if (adapter.id === "openai") invalidateProviderState(adapter.id);
+      const authState =
+        error instanceof UnsupportedOpenAIUsageAuthError
+          ? "unsupported"
+          : adapter.id === "openai" && !isTimeoutError(error)
+            ? "failed"
+            : undefined;
       if (displayState === "current") {
         transitionCurrentIdentity(`${adapter.id}:auth-error`, adapter.id);
       }
@@ -251,9 +275,11 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
           providerId: adapter.id,
           providerName,
           displayState,
-          status: isTimeoutError(error) ? "query-failed" : "auth-unavailable",
+          status:
+            authState === "unsupported" ? "unsupported" : isTimeoutError(error) ? "query-failed" : "auth-unavailable",
           message: errorMessage(error),
         },
+        ...(authState ? { authState } : {}),
       };
     }
     const requiresRequestBoundaryGuard =
@@ -265,6 +291,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
         "minimax-cn",
         "moonshotai",
         "moonshotai-cn",
+        "openai",
         "vercel-ai-gateway",
         "xai",
         "zai",
@@ -274,10 +301,12 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
       expectedSessionGeneration !== sessionGeneration ||
       ctx.sessionManager.getSessionId() !== expectedSessionId ||
       modelIdentity(ctx.model) !== expectedModelIdentity ||
+      (adapter.id === "openai" && settingsRuntime.get().settings.openaiCompanionUsage !== expectedCompanionUsage) ||
       (adapter.targets !== undefined &&
         settingsRuntime.get().settings.selectedTargets[adapter.id] !== expectedTargetId);
     if (requiresRequestBoundaryGuard && requestContextChanged()) throw abortError();
     if (!auth) {
+      if (adapter.id === "openai") invalidateProviderState(adapter.id);
       if (displayState === "current") {
         transitionCurrentIdentity(`${adapter.id}:unavailable`, adapter.id);
       }
@@ -292,6 +321,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
         authState: "unavailable",
       };
     }
+    if (adapter.id === "openai" && !auth.openaiCompanion) invalidateProviderState(adapter.id);
     let retryableAuthChanged = false;
     const guard = async () => {
       if (signal.aborted || requestContextChanged()) throw abortError();
@@ -456,6 +486,25 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
     }
   };
 
+  const queryAdapterState = async (...args: Parameters<typeof queryAdapterStateInternal>): Promise<QueryOutcome> => {
+    const [ctx, adapter, displayState, force, signal, ...rest] = args;
+    if (adapter.id !== "openai") return queryAdapterStateInternal(...args);
+    const controller = new AbortController();
+    nativeQueryControllers.add(controller);
+    try {
+      return await queryAdapterStateInternal(
+        ctx,
+        adapter,
+        displayState,
+        force,
+        AbortSignal.any([signal, controller.signal]),
+        ...rest,
+      );
+    } finally {
+      nativeQueryControllers.delete(controller);
+    }
+  };
+
   const loadTargetChoices = async (ctx: ExtensionContext, adapter: UsageProviderAdapter, signal: AbortSignal) => {
     if (!adapter.targets) throw new Error("Provider does not support usage targets.");
     const expectedSessionGeneration = sessionGeneration;
@@ -598,11 +647,11 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
     if (generation !== statusGeneration || modelIdentity(ctx.model) !== modelIdentity(model)) {
       return false;
     }
-    const adapter = adapterForProvider(model?.provider);
+    const adapter = adapterForProvider(outcome.state.providerId);
     const selectionStillCurrent =
       !adapter?.targets || settingsRuntime.get().settings.selectedTargets[adapter.id] === outcome.rememberedTargetId;
     if (!selectionStillCurrent) return false;
-    if (outcome.authState === "unavailable") {
+    if (outcome.authState) {
       if (!adapter) return false;
       try {
         const auth = await awaitWithDeadline(
@@ -616,11 +665,27 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
           modelIdentity(ctx.model) === modelIdentity(model) &&
           (!adapter.targets ||
             settingsRuntime.get().settings.selectedTargets[adapter.id] === outcome.rememberedTargetId) &&
+          outcome.authState === "unavailable" &&
           auth === undefined
         );
       } catch (error) {
         if (isAbortError(error) || isStaleExtensionContextError(error)) throw error;
-        return false;
+        // Stable failures are still useful. A recovered credential or changed reason must
+        // retry, but confirming the same native failure must not create a refresh loop.
+        const unchangedFailure =
+          (outcome.authState === "unsupported" && error instanceof UnsupportedOpenAIUsageAuthError) ||
+          (outcome.authState === "failed" &&
+            outcome.state.status === "auth-unavailable" &&
+            !(error instanceof UnsupportedOpenAIUsageAuthError) &&
+            !isTimeoutError(error) &&
+            errorMessage(error) === outcome.state.message);
+        return (
+          unchangedFailure &&
+          generation === statusGeneration &&
+          modelIdentity(ctx.model) === modelIdentity(model) &&
+          (!adapter.targets ||
+            settingsRuntime.get().settings.selectedTargets[adapter.id] === outcome.rememberedTargetId)
+        );
       }
     }
     if (!outcome.fingerprint) return true;
@@ -645,6 +710,27 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
     }
   };
 
+  const queryStableAdapterState = async (
+    ctx: ExtensionContext,
+    adapter: UsageProviderAdapter,
+    displayState: UsageDisplayState,
+    force: boolean,
+    signal: AbortSignal,
+  ): Promise<QueryOutcome> => {
+    if (adapter.id !== "openai") return queryAdapterState(ctx, adapter, displayState, force, signal);
+    // Configured/all-provider views must confirm native auth failures too, independently
+    // of the selected provider's subsequent revalidation.
+    for (let attempt = 0; attempt < 3; attempt += 1) {
+      const model = ctx.model;
+      const generation = statusGeneration;
+      const outcome = await queryAdapterState(ctx, adapter, displayState, force, signal);
+      if (signal.aborted) throw abortError();
+      if (await outcomeStillCurrent(ctx, model, generation, outcome, signal)) return outcome;
+      force = false;
+    }
+    throw new Error("OpenAI runtime authentication kept changing; reopen /usage to retry.");
+  };
+
   const queryStableCurrent = async (
     ctx: ExtensionCommandContext,
     force: boolean,
@@ -976,7 +1062,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
                 ctx,
                 `Checking ${providerDisplayName(ctx, adapter.id)} usage…`,
                 controller.signal,
-                (signal) => queryAdapterState(ctx, adapter, "configured", true, signal),
+                (signal) => queryStableAdapterState(ctx, adapter, "configured", true, signal),
               );
               if (!outcome) return { kind: "stay" };
               const revalidated = await queryStableCurrent(ctx, false, controller, "Revalidating current usage…");
@@ -1004,6 +1090,20 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
               controller.signal,
               () => statusGeneration === menuGeneration && !controller.signal.aborted,
               (id) => {
+                if (
+                  id === "openaiCompanionUsage" &&
+                  statusGeneration === menuGeneration &&
+                  !controller.signal.aborted
+                ) {
+                  for (const active of nativeQueryControllers) active.abort();
+                  invalidateProviderState("openai");
+                  if (ctx.model?.provider === "openai") {
+                    statusController?.abort();
+                    statusController = undefined;
+                    clearStatusTimers();
+                    safeSetStatus(ctx, undefined);
+                  }
+                }
                 if (
                   (id === "codexStatusResetCountdown" || id === "codexStatusPercentage") &&
                   stableCurrent &&
@@ -1189,7 +1289,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
               ctx,
               `Checking ${providerDisplayName(ctx, adapter.id)} usage…`,
               controller.signal,
-              (signal) => queryAdapterState(ctx, adapter, "configured", false, signal),
+              (signal) => queryStableAdapterState(ctx, adapter, "configured", false, signal),
             );
             if (!outcome) return { kind: "back" };
             if (outcome.state.status === "selection-required") {
@@ -1200,12 +1300,23 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
                 ctx,
                 `Checking ${providerDisplayName(ctx, adapter.id)} usage…`,
                 controller.signal,
-                (signal) => queryAdapterState(ctx, adapter, "configured", true, signal),
+                (signal) => queryStableAdapterState(ctx, adapter, "configured", true, signal),
               );
               if (!outcome) return { kind: "back" };
             }
             const revalidated = await queryStableCurrent(ctx, false, controller, "Revalidating current usage…");
             if (!revalidated) return { kind: "back" };
+            // Current-provider revalidation can await unrelated work after the configured
+            // native report was checked. Refresh its identity at the publication boundary.
+            if (adapter.id === "openai") {
+              outcome = await runMenuOperation(
+                ctx,
+                "Revalidating ChatGPT companion usage…",
+                controller.signal,
+                (signal) => queryStableAdapterState(ctx, adapter, "configured", false, signal),
+              );
+              if (!outcome) return { kind: "back" };
+            }
             stableCurrent = revalidated;
             current = revalidated.outcome;
             visibleStates = [
@@ -1227,7 +1338,7 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
                   adapters,
                   ALL_PROVIDER_CONCURRENCY,
                   (adapter, _index, workerSignal) =>
-                    queryAdapterState(
+                    queryStableAdapterState(
                       ctx,
                       adapter,
                       adapter.id === currentProviderId ? "current" : "configured",
@@ -1253,6 +1364,22 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
             });
             const revalidated = await queryStableCurrent(ctx, false, controller, "Revalidating current usage…");
             if (!revalidated) return { kind: "stay" };
+            const nativeIndex = adapters.findIndex(
+              (adapter) => adapter.id === "openai" && adapter.id !== currentProviderId,
+            );
+            if (nativeIndex >= 0) {
+              const nativeAdapter = adapters[nativeIndex];
+              if (nativeAdapter) {
+                const native = await runMenuOperation(
+                  ctx,
+                  "Revalidating ChatGPT companion usage…",
+                  controller.signal,
+                  (signal) => queryStableAdapterState(ctx, nativeAdapter, "configured", false, signal),
+                );
+                if (!native) return { kind: "stay" };
+                queriedStates[nativeIndex] = native.state;
+              }
+            }
             stableCurrent = revalidated;
             current = revalidated.outcome;
             visibleStates = [
@@ -1309,6 +1436,8 @@ export default function usageExtension(pi: ExtensionAPI, dependencies: UsageExte
     startStatusRefresh(ctx, ctx.model, false);
   });
   pi.on("model_select", (event, ctx) => {
+    // Cancel read-only native queries, not confirmed legacy reset mutations.
+    for (const controller of nativeQueryControllers) controller.abort();
     startStatusRefresh(ctx, event.model, false);
   });
   pi.on("turn_start", (_event, ctx) => {