oira666_pi-limits-wait 0.4.13 → 0.4.14

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "oira666_pi-limits-wait",
-  "version": "0.4.13",
+  "version": "0.4.14",
   "description": "Extension for PI that allows to automatically wait for limits to reset and retry the request",
   "type": "module",
   "license": "MIT",

package/src/index.ts

@@ -1,4 +1,4 @@
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
 import { isInternalSyntheticModel } from "./models.js";
 import { loadFallbackSettings } from "./settings.js";
 import { state } from "./state.js";
@@ -29,6 +29,13 @@ export {
 export { freezingEnabled, waitForRateLimit, waitForRetry } from "./ui.js";
 export type { FallbackModel } from "./types.js";
 
+function registerContextApis(pi: ExtensionAPI, ctx: ExtensionContext): void {
+  for (const model of ctx.modelRegistry.getAll()) {
+    if (!isInternalSyntheticModel(model)) registerWrappedApi(pi, model.api);
+  }
+  if (ctx.model && !isInternalSyntheticModel(ctx.model)) registerWrappedApi(pi, ctx.model.api);
+}
+
 export default function (pi: ExtensionAPI) {
   state.extensionApi = pi;
 
@@ -50,14 +57,20 @@ export default function (pi: ExtensionAPI) {
 
     // Some extensions may register providers after this extension loads. Wrap
     // any APIs that exist by the time an agent starts too.
-    for (const model of ctx.modelRegistry.getAll()) {
-      if (!isInternalSyntheticModel(model)) registerWrappedApi(pi, model.api);
-    }
-    if (ctx.model && !isInternalSyntheticModel(ctx.model)) registerWrappedApi(pi, ctx.model.api);
+    registerContextApis(pi, ctx);
 
     // Anthropic subscription/OAuth identity is applied per provider request in
     // streamWithLimitsRetry(). Doing it there ensures every retry/fallback
     // attempt gets the right prompt for its actual target model, while
     // non-Anthropic providers never inherit the Claude Code identity.
   });
+
+  // Run once more after all before_agent_start handlers have completed and the
+  // agent loop is about to make its first provider request. This catches
+  // providers that another extension registered during before_agent_start after
+  // our handler had already run.
+  pi.on("agent_start", (_event, ctx) => {
+    state.sharedCtx = ctx;
+    registerContextApis(pi, ctx);
+  });
 }

package/src/models.ts

@@ -1,3 +1,4 @@
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
 import type { Api, Model, SimpleStreamOptions } from "@mariozechner/pi-ai";
 import { DEFAULT_NON_RETRYABLE_FREEZE_MS } from "./constants.js";
 import { isRateLimitError, parseRetryDelayMs } from "./retry-errors.js";
@@ -147,13 +148,14 @@ export async function optionsForModel(
   originalModel: Model<Api>,
   target: FallbackModel,
   options?: SimpleStreamOptions,
+  ctx: ExtensionContext | undefined = state.sharedCtx,
 ): Promise<SimpleStreamOptions | undefined> {
   const level = target.reasoningEffort ?? state.primaryThinkingLevel;
   const reasoning = level && level !== "off" ? level : undefined;
   if (modelKey(originalModel) === modelKey(target.model)) {
     return reasoning ? { ...options, reasoning } : options;
   }
-  const auth = await state.sharedCtx?.modelRegistry.getApiKeyAndHeaders(target.model);
+  const auth = await ctx?.modelRegistry.getApiKeyAndHeaders(target.model);
   if (!auth?.ok) throw new Error(auth ? auth.error : "Model registry is unavailable.");
   return {
     ...options,

package/src/prompt.ts

@@ -1,5 +1,5 @@
 import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
-import type { Api, Context, Model } from "@mariozechner/pi-ai";
+import type { Api, Context, Model, SimpleStreamOptions } from "@mariozechner/pi-ai";
 import { CLAUDE_CODE_IDENTITY, CLAUDE_CODE_IDENTITY_PATTERN, PI_IDENTITY_SENTENCE_PATTERN, PI_REMOVAL_ANCHORS } from "./constants.js";
 import { state } from "./state.js";
 
@@ -28,8 +28,56 @@ export function isAnthropicOAuthSession(ctx: ExtensionContext): boolean {
   );
 }
 
-export function anthropicSubscriptionContext(model: Model<Api>, context: Context): Context {
-  if (model.provider !== "anthropic" || !state.sharedCtx?.modelRegistry.isUsingOAuth(model)) return context;
+export function isAnthropicOAuthToken(value: string | undefined): boolean {
+  return Boolean(value?.includes("sk-ant-oat"));
+}
+
+function headerValue(headers: Record<string, string> | undefined, name: string): string | undefined {
+  if (!headers) return undefined;
+  const direct = headers[name];
+  if (direct !== undefined) return direct;
+  const lower = name.toLowerCase();
+  return Object.entries(headers).find(([key]) => key.toLowerCase() === lower)?.[1];
+}
+
+function requestAuthIsAnthropicOAuth(options?: SimpleStreamOptions): boolean | undefined {
+  if (typeof options?.apiKey === "string" && options.apiKey.length > 0) {
+    return isAnthropicOAuthToken(options.apiKey);
+  }
+
+  const authorization = headerValue(options?.headers, "authorization");
+  if (authorization) {
+    const bearer = authorization.match(/^\s*Bearer\s+(.+?)\s*$/i)?.[1] ?? authorization;
+    return isAnthropicOAuthToken(bearer);
+  }
+
+  return undefined;
+}
+
+export function isAnthropicSubscriptionRequest(
+  model: Model<Api>,
+  options?: SimpleStreamOptions,
+  ctx: ExtensionContext | undefined = state.sharedCtx,
+): boolean {
+  if (model.provider !== "anthropic") return false;
+
+  // Prefer the actual per-request auth. This avoids both false negatives
+  // (OAuth token supplied via options/fallback auth while modelRegistry does not
+  // report OAuth) and false positives (registry has OAuth, but this request was
+  // explicitly made with a normal API key).
+  const requestAuth = requestAuthIsAnthropicOAuth(options);
+  if (requestAuth !== undefined) return requestAuth;
+
+  return Boolean(ctx?.modelRegistry.isUsingOAuth(model));
+}
+
+export function anthropicSubscriptionContext(
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+  ctx: ExtensionContext | undefined = state.sharedCtx,
+): Context {
+  if (!isAnthropicSubscriptionRequest(model, options, ctx)) return context;
 
   const sanitised = sanitiseSystemPrompt(context.systemPrompt ?? "");
   return {

package/src/stream.ts

@@ -89,6 +89,7 @@ export function streamWithLimitsRetry(
   const output = createAssistantMessageEventStream();
 
   void (async () => {
+    const requestCtx = state.sharedCtx;
     let committed = false;
     const allowFallback = fallbackEnabled() && isFallbackEligibleModel(model);
     let attempt: FallbackModel = allowFallback ? initialAttempt(model) : { model, reasoningEffort: state.primaryThinkingLevel };
@@ -124,14 +125,14 @@ export function streamWithLimitsRetry(
       try {
         try {
           const attemptOptions = allowFallback
-            ? await optionsForModel(model, attempt, options)
+            ? await optionsForModel(model, attempt, options, requestCtx)
             : options;
           const attemptDelegate = attempt.model.api === model.api
             ? delegate
             : state.builtinStreamSimpleByApi.get(attempt.model.api)
               ?? getApiProviders().find((provider) => provider.api === attempt.model.api)?.streamSimple;
           if (!attemptDelegate) throw new Error(`No stream handler registered for API ${attempt.model.api}.`);
-          const attemptContext = anthropicSubscriptionContext(attempt.model, context);
+          const attemptContext = anthropicSubscriptionContext(attempt.model, context, attemptOptions, requestCtx);
           const inner = await attemptDelegate(attempt.model, attemptContext, attemptOptions);
           for await (const event of inner) {
             if (!committed) {
@@ -338,17 +339,34 @@ export function streamWithLimitsRetry(
 
 export const streamWithRateLimitRetry = streamWithLimitsRetry;
 
-export function registerWrappedApi(pi: ExtensionAPI, api: Api): void {
-  if (state.wrappedApis.has(api)) return;
+const LIMITS_WAIT_WRAPPER_FLAG = "__piLimitsWaitWrapper";
+type LimitsWaitStreamSimpleFn = ((
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+) => AssistantMessageEventStream) & { [LIMITS_WAIT_WRAPPER_FLAG]?: true };
 
-  const builtinStreamSimple = getApiProviders().find((provider) => provider.api === api)?.streamSimple;
-  if (!builtinStreamSimple) return;
+export function registerWrappedApi(pi: ExtensionAPI, api: Api): void {
+  const currentStreamSimple = getApiProviders().find((provider) => provider.api === api)?.streamSimple;
+  if (!currentStreamSimple) return;
+  if ((currentStreamSimple as LimitsWaitStreamSimpleFn)[LIMITS_WAIT_WRAPPER_FLAG]) return;
 
-  state.builtinStreamSimpleByApi.set(api, builtinStreamSimple);
+  state.builtinStreamSimpleByApi.set(api, currentStreamSimple);
   state.wrappedApis.add(api);
+
+  const wrappedStreamSimple = ((model: Model<Api>, context: Context, options?: SimpleStreamOptions) =>
+    streamWithLimitsRetry(currentStreamSimple, model, context, options)) as LimitsWaitStreamSimpleFn;
+  wrappedStreamSimple[LIMITS_WAIT_WRAPPER_FLAG] = true;
+
   pi.registerProvider(`limits-wait-${api}`, {
     api,
-    streamSimple: (model: Model<Api>, context: Context, options?: SimpleStreamOptions) =>
-      streamWithLimitsRetry(builtinStreamSimple, model, context, options),
+    streamSimple: wrappedStreamSimple,
   });
+
+  // pi-ai wraps provider functions when registering them, so the function later
+  // returned by getApiProviders() is not the same object as wrappedStreamSimple.
+  // Tag the registered wrapper too; otherwise repeated agent_start/model loops
+  // would wrap an already wrapped API again.
+  const registeredStreamSimple = getApiProviders().find((provider) => provider.api === api)?.streamSimple as LimitsWaitStreamSimpleFn | undefined;
+  if (registeredStreamSimple) registeredStreamSimple[LIMITS_WAIT_WRAPPER_FLAG] = true;
 }