oira666_pi-limits-wait 0.4.13 → 0.4.15

package/package.json

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

package/src/constants.ts

@@ -10,6 +10,13 @@ export const PI_REMOVAL_ANCHORS = [
 export const PI_IDENTITY_SENTENCE_PATTERN =
   /(?:^|\n)\s*You are pi\b[^.!?\n]*(?:[.!?](?=\s|$)|(?=\n|$))/gi;
 
+// Strips the host-harness identity clause pi prepends to its system prompt, e.g.
+// "You are an expert coding assistant operating inside pi, a coding agent harness."
+// The Claude Code identity block already establishes who the assistant is, so
+// leaving this clause in tells the model (and Anthropic) it is running inside pi.
+export const PI_HARNESS_IDENTITY_PATTERN =
+  /\s+operating inside pi,\s*a coding agent harness/gi;
+
 export const CLAUDE_CODE_IDENTITY_PATTERN =
   /(?:^|\n)\s*You are Claude Code, Anthropic's official CLI for Claude\.\s*/gi;
 

package/src/index.ts

@@ -1,10 +1,11 @@
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
 import { isInternalSyntheticModel } from "./models.js";
+import { sanitiseAnthropicPayloadSystem } from "./prompt.js";
 import { loadFallbackSettings } from "./settings.js";
 import { state } from "./state.js";
 import { registerWrappedApi } from "./stream.js";
 
-export { sanitiseSystemPrompt } from "./prompt.js";
+export { sanitiseAnthropicPayloadSystem, sanitiseSystemPrompt } from "./prompt.js";
 export {
   getRetryableError,
   isAuthenticationRefreshError,
@@ -29,9 +30,26 @@ 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;
 
+  // Anthropic now rejects Claude Pro/Max (OAuth) requests whose system prompt
+  // still carries the host agent's fingerprints (pi-coding-agent paths, "You
+  // are pi", etc.). pi-ai prepends the Claude Code identity block for OAuth
+  // tokens but passes the host system prompt through as a second block, which
+  // trips the check. Sanitise that block on the outbound payload right before
+  // it is sent. This hook fires for every provider request and only mutates
+  // Anthropic OAuth payloads (identified by the Claude Code identity block), so
+  // API-key Anthropic calls and non-Anthropic providers are untouched.
+  pi.on("before_provider_request", (event) => sanitiseAnthropicPayloadSystem(event.payload));
+
   pi.on("model_select", (event) => {
     if (state.suppressNextModelSelect) {
       state.suppressNextModelSelect = false;
@@ -50,14 +68,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,6 +1,6 @@
 import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
-import type { Api, Context, Model } from "@mariozechner/pi-ai";
-import { CLAUDE_CODE_IDENTITY, CLAUDE_CODE_IDENTITY_PATTERN, PI_IDENTITY_SENTENCE_PATTERN, PI_REMOVAL_ANCHORS } from "./constants.js";
+import type { Api, Context, Model, SimpleStreamOptions } from "@mariozechner/pi-ai";
+import { CLAUDE_CODE_IDENTITY, CLAUDE_CODE_IDENTITY_PATTERN, PI_HARNESS_IDENTITY_PATTERN, PI_IDENTITY_SENTENCE_PATTERN, PI_REMOVAL_ANCHORS } from "./constants.js";
 import { state } from "./state.js";
 
 export function sanitiseSystemPrompt(raw: string): string {
@@ -13,6 +13,7 @@ export function sanitiseSystemPrompt(raw: string): string {
     .join("\n\n")
     .replace(CLAUDE_CODE_IDENTITY_PATTERN, "")
     .replace(PI_IDENTITY_SENTENCE_PATTERN, "")
+    .replace(PI_HARNESS_IDENTITY_PATTERN, "")
     .replace(/\n{3,}/g, "\n\n")
     .trim();
 }
@@ -28,8 +29,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 {
@@ -39,3 +88,61 @@ export function anthropicSubscriptionContext(model: Model<Api>, context: Context
       : CLAUDE_CODE_IDENTITY,
   };
 }
+
+interface SystemTextBlock {
+  type: string;
+  text?: unknown;
+  [key: string]: unknown;
+}
+
+function isClaudeCodeIdentityBlock(block: unknown): block is SystemTextBlock {
+  if (!block || typeof block !== "object") return false;
+  const b = block as SystemTextBlock;
+  return b.type === "text" && typeof b.text === "string" && b.text.includes(CLAUDE_CODE_IDENTITY);
+}
+
+/**
+ * Sanitise the `system` blocks of an Anthropic provider request payload.
+ *
+ * Anthropic now rejects Claude Pro/Max (OAuth) requests whose system prompt
+ * still carries third-party-agent fingerprints (e.g. pi-coding-agent paths or
+ * "You are pi" identity). pi-ai already prepends the Claude Code identity as
+ * the first `system` block for OAuth tokens, but passes the host agent's full
+ * system prompt through as a second block unchanged, which trips the check.
+ *
+ * This runs from the `before_provider_request` extension hook, which fires for
+ * every provider request and whose return value replaces the outbound payload.
+ * It only touches payloads that already carry the Claude Code identity block
+ * (i.e. Anthropic OAuth requests), so API-key Anthropic calls and non-Anthropic
+ * providers are left untouched. Each remaining text block is run through
+ * `sanitiseSystemPrompt`; blocks that become empty are dropped.
+ */
+export function sanitiseAnthropicPayloadSystem(payload: unknown): unknown {
+  if (!payload || typeof payload !== "object") return payload;
+  const system = (payload as { system?: unknown }).system;
+  if (!Array.isArray(system)) return payload;
+
+  const identityIndex = system.findIndex(isClaudeCodeIdentityBlock);
+  if (identityIndex === -1) return payload;
+
+  const newSystem: unknown[] = [];
+  for (let i = 0; i < system.length; i++) {
+    const block = system[i];
+    if (i === identityIndex) {
+      newSystem.push(block);
+      continue;
+    }
+    if (
+      block && typeof block === "object"
+      && (block as SystemTextBlock).type === "text"
+      && typeof (block as SystemTextBlock).text === "string"
+    ) {
+      const sanitised = sanitiseSystemPrompt((block as SystemTextBlock).text as string);
+      if (sanitised) newSystem.push({ ...block, text: sanitised });
+    } else {
+      newSystem.push(block);
+    }
+  }
+
+  return { ...payload, system: newSystem };
+}

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;
 }