pi-better-openai 0.1.19 → 0.1.21

package/src/pets.ts

@@ -2,7 +2,7 @@ import { constants, type Dirent } from "node:fs";
 import { access, lstat, readdir, readFile, realpath } from "node:fs/promises";
 import { homedir } from "node:os";
 import { join, resolve, sep } from "node:path";
-import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 import {
   calculateImageRows,
   getCapabilities,
@@ -12,7 +12,7 @@ import {
   truncateToWidth,
   type AutocompleteItem,
   type ImageTheme,
-} from "@mariozechner/pi-tui";
+} from "@earendil-works/pi-tui";
 import sharp from "sharp";
 import { isRecord, type PetState } from "./config.ts";
 
@@ -25,6 +25,7 @@ const DEFAULT_SPRITESHEET_PATH = "spritesheet.webp";
 const EXPECTED_ATLAS_WIDTH = DEFAULT_CELL_WIDTH * PET_COLUMNS;
 const EXPECTED_ATLAS_HEIGHT = DEFAULT_CELL_HEIGHT * PET_ROWS;
 const PET_CATALOG_CACHE_TTL_MS = 1500;
+const PET_CATALOG_READ_CONCURRENCY = 4;
 
 type ListCodexPetsOptions = { refresh?: boolean };
 
@@ -61,8 +62,8 @@ export type CodexPetPackage = {
 export type PetFrame = {
   data?: string;
   rawRgbaData?: string;
+  kittyUpload?: string;
   kittyImageId?: number;
-  kittyUploaded?: boolean;
   mimeType: "image/png";
   durationMs: number;
   widthPx: number;
@@ -265,6 +266,25 @@ async function readCodexPetPackage(
   };
 }
 
+async function readCodexPetPackages(
+  petsDir: string,
+  entries: Dirent[],
+  options: { validateSpritesheet: boolean },
+): Promise<CodexPetPackage[]> {
+  const pets: Array<CodexPetPackage | undefined> = Array.from({ length: entries.length });
+  let nextIndex = 0;
+  const readNext = async (): Promise<void> => {
+    while (nextIndex < entries.length) {
+      const index = nextIndex++;
+      pets[index] = await readCodexPetPackage(petsDir, entries[index], options);
+    }
+  };
+  await Promise.all(
+    Array.from({ length: Math.min(PET_CATALOG_READ_CONCURRENCY, entries.length) }, readNext),
+  );
+  return pets.filter((pet): pet is CodexPetPackage => pet !== undefined);
+}
+
 export async function listCodexPets(
   home = codexHome(),
   options: ListCodexPetsOptions = {},
@@ -274,22 +294,22 @@ export async function listCodexPets(
   if (!options.refresh && cached && cached.expiresAt > Date.now()) return cached.pets;
 
   const { dir, entries } = await readPetDirectoryEntries(home);
-  const pets: CodexPetPackage[] = [];
-  for (const entry of entries) {
-    const pet = await readCodexPetPackage(dir, entry, { validateSpritesheet: true });
-    if (pet) pets.push(pet);
-  }
+  const pets = await readCodexPetPackages(dir, entries, { validateSpritesheet: true });
   pets.sort((a, b) => a.name.localeCompare(b.name));
   petCatalogCache.set(cacheKey, { expiresAt: Date.now() + PET_CATALOG_CACHE_TTL_MS, pets });
   return pets;
 }
 
-function petLookupKey(value: string): string {
-  return value.toLowerCase().replace(/[^a-z0-9]+/g, "");
+export function petLookupKey(value: string): string {
+  return value
+    .normalize("NFKC")
+    .toLowerCase()
+    .replace(/[^\p{L}\p{N}\p{M}]+/gu, "");
 }
 
 type PetLookupRequest = {
   hasText: boolean;
+  exact?: string;
   key?: string;
 };
 
@@ -297,7 +317,11 @@ function petLookupRequest(value?: string): PetLookupRequest {
   const requestedText = value?.trim();
   if (!requestedText) return { hasText: false };
 
-  return { hasText: true, key: petLookupKey(requestedText) || undefined };
+  return {
+    hasText: true,
+    exact: requestedText.normalize("NFKC").toLowerCase(),
+    key: petLookupKey(requestedText) || undefined,
+  };
 }
 
 function petMatchesLookup(pet: CodexPetPackage, requested: string): boolean {
@@ -308,11 +332,35 @@ function petMatchesLookup(pet: CodexPetPackage, requested: string): boolean {
   );
 }
 
+function exactPetMatches(
+  pets: CodexPetPackage[],
+  requested: string,
+): CodexPetPackage[] | undefined {
+  const fields: Array<(pet: CodexPetPackage) => string | undefined> = [
+    (pet) => pet.slug,
+    (pet) => pet.id,
+    (pet) => pet.name,
+  ];
+  for (const field of fields) {
+    const matches = pets.filter((pet) => field(pet)?.normalize("NFKC").toLowerCase() === requested);
+    if (matches.length > 0) return matches;
+  }
+  return undefined;
+}
+
+function uniquePetMatch(pets: CodexPetPackage[]): CodexPetPackage | undefined {
+  return pets.length === 1 ? pets[0] : undefined;
+}
+
 export function findCodexPet(pets: CodexPetPackage[], value?: string): CodexPetPackage | undefined {
   const request = petLookupRequest(value);
   const { key } = request;
   if (!key) return undefined;
-  return pets.find((pet) => petMatchesLookup(pet, key));
+  if (request.exact) {
+    const exactMatches = exactPetMatches(pets, request.exact);
+    if (exactMatches) return uniquePetMatch(exactMatches);
+  }
+  return uniquePetMatch(pets.filter((pet) => petMatchesLookup(pet, key)));
 }
 
 export function findReadyCodexPet(
@@ -324,7 +372,14 @@ export function findReadyCodexPet(
 
   const { key } = request;
   if (!key) return undefined;
-  return pets.find((pet) => pet.hasSpritesheet && petMatchesLookup(pet, key));
+  if (request.exact) {
+    const exactMatches = exactPetMatches(pets, request.exact);
+    if (exactMatches) {
+      const exact = uniquePetMatch(exactMatches);
+      return exact?.hasSpritesheet ? exact : undefined;
+    }
+  }
+  return uniquePetMatch(pets.filter((pet) => pet.hasSpritesheet && petMatchesLookup(pet, key)));
 }
 
 function selectPet(pets: CodexPetPackage[], slug?: string): CodexPetPackage | undefined {
@@ -431,18 +486,15 @@ async function findCodexPetDirect(
   const request = petLookupRequest(slug);
   if (!request.key) return undefined;
   const { dir, entries } = await readPetDirectoryEntries(home);
-  for (const entry of entries) {
-    const candidate = await readCodexPetPackage(dir, entry, { validateSpritesheet: false });
-    if (!candidate || !petMatchesLookup(candidate, request.key)) continue;
-    if (candidate.spritesheetIssue) return candidate;
-    const spritesheetIssue = await validatePetSpritesheet(candidate.dir, candidate.spritesheetPath);
-    return {
-      ...candidate,
-      hasSpritesheet: spritesheetIssue === undefined,
-      spritesheetIssue,
-    };
-  }
-  return undefined;
+  const candidates = await readCodexPetPackages(dir, entries, { validateSpritesheet: false });
+  const candidate = findCodexPet(candidates, slug);
+  if (!candidate || candidate.spritesheetIssue) return candidate;
+  const spritesheetIssue = await validatePetSpritesheet(candidate.dir, candidate.spritesheetPath);
+  return {
+    ...candidate,
+    hasSpritesheet: spritesheetIssue === undefined,
+    spritesheetIssue,
+  };
 }
 
 export async function loadCodexPet(
@@ -513,12 +565,20 @@ export async function loadCodexPet(
           kernel: sharp.kernel.nearest,
         });
       }
+      const kittyImageId = useKitty ? kittyImageBase + kittyFrameOffset++ : undefined;
       const encoded = useKitty
-        ? { rawRgbaData: (await frame.clone().ensureAlpha().raw().toBuffer()).toString("base64") }
-        : { data: (await frame.clone().png().toBuffer()).toString("base64") };
+        ? {
+            kittyUpload: encodeKittyRawRgbaData(
+              (await frame.ensureAlpha().raw().toBuffer()).toString("base64"),
+              outputWidthPx,
+              outputHeightPx,
+              kittyImageId!,
+            ),
+          }
+        : { data: (await frame.png().toBuffer()).toString("base64") };
       states[state].push({
         ...encoded,
-        kittyImageId: useKitty ? kittyImageBase + kittyFrameOffset++ : undefined,
+        kittyImageId,
         mimeType: "image/png",
         durationMs,
         widthPx: outputWidthPx,
@@ -575,12 +635,6 @@ function forEachCodexPetFrame(
   }
 }
 
-function markCodexPetKittyFramesUnuploaded(pet?: LoadedCodexPet): void {
-  forEachCodexPetFrame(pet, (frame) => {
-    frame.kittyUploaded = false;
-  });
-}
-
 function codexPetKittyImageIds(pet?: LoadedCodexPet): number[] {
   const imageIds = new Set<number>();
   forEachCodexPetFrame(pet, (frame) => {
@@ -601,18 +655,15 @@ function placeKittyImage(imageId: number, columns: number, rows: number): string
   return `\x1b_Ga=p,i=${imageId},p=1,c=${columns},r=${rows},q=2,C=1\x1b\\`;
 }
 
-function encodeKittyRawRgba(frame: PetFrame, imageId: number): string {
-  const rawRgbaData = frame.rawRgbaData;
+function encodeKittyRawRgbaData(
+  rawRgbaData: string,
+  widthPx: number,
+  heightPx: number,
+  imageId: number,
+): string {
   if (!rawRgbaData) return "";
   const chunkSize = 4096;
-  const params = [
-    "a=t",
-    "f=32",
-    `s=${frame.widthPx}`,
-    `v=${frame.heightPx}`,
-    `i=${imageId}`,
-    "q=2",
-  ].join(",");
+  const params = ["a=t", "f=32", `s=${widthPx}`, `v=${heightPx}`, `i=${imageId}`, "q=2"].join(",");
   if (rawRgbaData.length <= chunkSize) {
     return `\x1b_G${params};${rawRgbaData}\x1b\\`;
   }
@@ -628,6 +679,15 @@ function encodeKittyRawRgba(frame: PetFrame, imageId: number): string {
   return chunks.join("");
 }
 
+function encodeKittyRawRgba(frame: PetFrame, imageId: number): string {
+  return (
+    frame.kittyUpload ??
+    (frame.rawRgbaData
+      ? encodeKittyRawRgbaData(frame.rawRgbaData, frame.widthPx, frame.heightPx, imageId)
+      : "")
+  );
+}
+
 export class CodexPetKittyManager {
   private previousFrameImageId: number | undefined;
   private readonly pendingCleanupImageIds = new Set<number>();
@@ -641,7 +701,6 @@ export class CodexPetKittyManager {
   invalidate(pet?: LoadedCodexPet): void {
     this.queueCleanup(pet);
     this.previousFrameImageId = undefined;
-    markCodexPetKittyFramesUnuploaded(pet);
   }
 
   resetForResize(pet?: LoadedCodexPet): void {
@@ -683,7 +742,6 @@ export class CodexPetKittyManager {
     // occasionally invisible pet. The TUI diff still avoids writing this payload
     // again when the rendered line is unchanged.
     const upload = encodeKittyRawRgba(frame, frameImageId);
-    frame.kittyUploaded = true;
     this.previousFrameImageId = frameImageId;
     // Recent pi-tui versions automatically free Kitty image IDs they own on changed lines.
     // Pet frames are managed here and intentionally reused across animation loops, so keep
@@ -716,7 +774,6 @@ function defaultKittyManager(placementImageId: number): CodexPetKittyManager {
 export function resetCodexPetKittyCache(pet?: LoadedCodexPet, placementImageId?: number): void {
   if (placementImageId !== undefined) defaultKittyManagers.get(placementImageId)?.invalidate(pet);
   else for (const manager of defaultKittyManagers.values()) manager.invalidate(pet);
-  markCodexPetKittyFramesUnuploaded(pet);
 }
 
 export function renderCodexPetFrame(

package/src/usage-controller.ts

@@ -1,6 +1,6 @@
-import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import type { ResolvedConfig } from "./config.ts";
-import { maskIdentifier } from "./format.ts";
+import { maskIdentifier, sanitizeDiagnosticError } from "./format.ts";
 import {
   AUTH_FILE,
   type UsageSnapshot,
@@ -9,19 +9,34 @@ import {
   parseUsageSnapshot,
   readCodexAuth,
   requestCodexUsage,
+  usageScopeForModel,
 } from "./usage.ts";
 import { currentModelKey } from "./fast-controller.ts";
 
-export function isOpenAISubscriptionModel(ctx: ExtensionContext, cfg: ResolvedConfig): boolean {
-  if (!ctx.model || (ctx.model.provider !== "openai" && ctx.model.provider !== "openai-codex"))
-    return false;
-  return !cfg.usage.showOnlyOnSubscriptionModels || ctx.modelRegistry.isUsingOAuth(ctx.model);
+export function isOpenAISubscriptionModel(
+  ctx: ExtensionContext,
+  cfg: ResolvedConfig,
+  isUsingOAuth?: boolean,
+): boolean {
+  const model = ctx.model;
+  if (!model || (model.provider !== "openai" && model.provider !== "openai-codex")) return false;
+  return (
+    !cfg.usage.showOnlyOnSubscriptionModels ||
+    (isUsingOAuth ?? ctx.modelRegistry.isUsingOAuth(model))
+  );
+}
+
+const STALE_EXTENSION_CONTEXT_MESSAGE = "This extension ctx is stale";
+
+function isStaleExtensionContextError(error: unknown): boolean {
+  return error instanceof Error && error.message.includes(STALE_EXTENSION_CONTEXT_MESSAGE);
 }
 
 type UsageRefreshOptions = { notify?: boolean; force?: boolean };
 
 type QueuedUsageRefresh = {
   ctx: ExtensionContext;
+  generation: number;
   modelId?: string;
   notify?: boolean;
   force?: boolean;
@@ -39,6 +54,7 @@ export class UsageController {
   private usageAbortController: AbortController | undefined;
   private sessionAbortSignal: AbortSignal | undefined;
   private sessionAbortHandler: (() => void) | undefined;
+  private sessionGeneration = 0;
 
   constructor(
     private readonly getConfig: (ctx: ExtensionContext) => ResolvedConfig,
@@ -49,8 +65,16 @@ export class UsageController {
     return this.usageSnapshot;
   }
 
-  statusLine(ctx: ExtensionContext, cfg = this.getConfig(ctx)): string | undefined {
-    return this.usageSnapshot && cfg.usage.enabled && isOpenAISubscriptionModel(ctx, cfg)
+  statusLine(
+    ctx: ExtensionContext,
+    cfg = this.getConfig(ctx),
+    isUsingOAuth?: boolean,
+  ): string | undefined {
+    return this.usageSnapshot &&
+      !this.usageError &&
+      this.usageSnapshot.scope === usageScopeForModel(ctx.model?.id) &&
+      cfg.usage.enabled &&
+      isOpenAISubscriptionModel(ctx, cfg, isUsingOAuth)
       ? formatUsageSnapshot(this.usageSnapshot, cfg.usage)
       : undefined;
   }
@@ -60,8 +84,9 @@ export class UsageController {
     if (!cfg.usage.enabled) return "Usage display is disabled.";
     if (!isOpenAISubscriptionModel(ctx, cfg))
       return "Usage hidden: current model is not an OpenAI subscription model.";
-    if (!this.usageSnapshot)
-      return `Usage unavailable${this.usageError ? `: ${this.usageError}` : "."}`;
+    if (this.usageError) return `Usage unavailable: ${this.usageError}`;
+    if (!this.usageSnapshot || this.usageSnapshot.scope !== usageScopeForModel(ctx.model?.id))
+      return "Usage unavailable.";
     const stale =
       this.usageUpdatedAt && Date.now() - this.usageUpdatedAt > cfg.usage.refreshIntervalMs * 2
         ? ` | stale ${formatResetCountdown((Date.now() - this.usageUpdatedAt) / 1000)}`
@@ -87,74 +112,124 @@ export class UsageController {
     ].join("\n");
   }
 
+  private isGenerationCurrent(generation: number): boolean {
+    return !this.shuttingDown && generation === this.sessionGeneration;
+  }
+
+  private deactivateGeneration(generation: number): void {
+    if (generation !== this.sessionGeneration) return;
+    this.shuttingDown = true;
+    this.sessionGeneration++;
+    this.queuedUsageRefresh = undefined;
+    this.usageAbortController?.abort();
+    this.usageAbortController = undefined;
+    this.stopTimer();
+  }
+
+  private handleStaleContextError(error: unknown, generation: number): boolean {
+    if (!isStaleExtensionContextError(error)) return false;
+    this.deactivateGeneration(generation);
+    return true;
+  }
+
   async refresh(
     ctx: ExtensionContext,
-    modelId = ctx.model?.id,
+    modelId?: string,
     options?: UsageRefreshOptions,
+    generation = this.sessionGeneration,
   ): Promise<void> {
-    if (this.shuttingDown || !ctx.hasUI) return;
+    if (!this.isGenerationCurrent(generation)) return;
+
+    let resolvedModelId = modelId;
+    try {
+      if (!ctx.hasUI) return;
+      resolvedModelId ??= ctx.model?.id;
+    } catch (error) {
+      this.handleStaleContextError(error, generation);
+      return;
+    }
+
     if (this.usageRefreshInFlight) {
+      const queued =
+        this.queuedUsageRefresh?.generation === generation ? this.queuedUsageRefresh : undefined;
       this.queuedUsageRefresh = {
         ctx,
-        modelId,
-        notify: this.queuedUsageRefresh?.notify || options?.notify,
-        force: this.queuedUsageRefresh?.force || options?.force,
+        generation,
+        modelId: resolvedModelId,
+        notify: queued?.notify || options?.notify,
+        force: queued?.force || options?.force,
       };
       return;
     }
+
     this.usageRefreshInFlight = true;
-    const cfg = this.getConfig(ctx);
     try {
+      const cfg = this.getConfig(ctx);
+      if (!this.isGenerationCurrent(generation)) return;
+
       if (!cfg.usage.enabled) {
         this.usageSnapshot = undefined;
         this.usageError = "Usage display is disabled.";
-        if (!this.shuttingDown) this.updateFooter(ctx);
-        if (!this.shuttingDown && options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
+        this.updateFooter(ctx);
+        if (options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
         return;
       }
       if (!isOpenAISubscriptionModel(ctx, cfg)) {
-        if (!this.shuttingDown) this.updateFooter(ctx);
-        if (!this.shuttingDown && options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
+        this.updateFooter(ctx);
+        if (options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
         return;
       }
       const shouldThrottle =
         !options?.force &&
         !options?.notify &&
         this.usageLastFetchAt !== undefined &&
-        Date.now() - this.usageLastFetchAt < cfg.usage.refreshIntervalMs &&
-        this.usageSnapshot !== undefined &&
-        this.usageError === undefined;
-      if (shouldThrottle) {
-        if (!this.shuttingDown) this.updateFooter(ctx);
-        return;
-      }
+        Date.now() - this.usageLastFetchAt < cfg.usage.refreshIntervalMs;
+      if (shouldThrottle) return;
+      this.usageLastFetchAt = Date.now();
       this.usageAbortController = new AbortController();
       const timeoutSignal = AbortSignal.timeout(10_000);
       const signal = ctx.signal
         ? AbortSignal.any([ctx.signal, timeoutSignal, this.usageAbortController.signal])
         : AbortSignal.any([timeoutSignal, this.usageAbortController.signal]);
       const data = await requestCodexUsage(ctx, signal);
-      this.usageLastFetchAt = Date.now();
-      this.usageSnapshot = data ? parseUsageSnapshot(data, modelId) : undefined;
+      if (!this.isGenerationCurrent(generation)) return;
+
+      this.usageSnapshot = data ? parseUsageSnapshot(data, resolvedModelId) : undefined;
       this.usageUpdatedAt = this.usageSnapshot ? Date.now() : undefined;
       this.usageError = data
         ? undefined
         : `Missing openai-codex OAuth credentials in ${AUTH_FILE}.`;
-      if (!this.shuttingDown) this.updateFooter(ctx);
-      if (!this.shuttingDown && options?.notify)
+      this.updateFooter(ctx);
+      if (options?.notify)
         ctx.ui.notify(this.formatStatus(ctx), this.usageSnapshot ? "info" : "warning");
     } catch (error) {
-      if (this.shuttingDown) return;
-      this.usageError = error instanceof Error ? error.message : String(error);
-      this.updateFooter(ctx);
-      if (options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
+      if (this.handleStaleContextError(error, generation) || !this.isGenerationCurrent(generation))
+        return;
+      this.usageError = sanitizeDiagnosticError(
+        error instanceof Error ? error.message : String(error),
+      );
+      try {
+        this.updateFooter(ctx);
+        if (options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
+      } catch (secondaryError) {
+        if (!this.handleStaleContextError(secondaryError, generation)) {
+          this.usageError = sanitizeDiagnosticError(
+            secondaryError instanceof Error ? secondaryError.message : String(secondaryError),
+          );
+        }
+      }
     } finally {
       this.usageAbortController = undefined;
       this.usageRefreshInFlight = false;
-      if (!this.shuttingDown && this.queuedUsageRefresh) {
-        const next = this.queuedUsageRefresh;
-        this.queuedUsageRefresh = undefined;
-        void this.refresh(next.ctx, next.modelId, { notify: next.notify, force: next.force });
+      const next = this.queuedUsageRefresh;
+      this.queuedUsageRefresh = undefined;
+      if (next && !this.shuttingDown && next.generation === this.sessionGeneration) {
+        void this.refresh(
+          next.ctx,
+          next.modelId,
+          { notify: next.notify, force: next.force },
+          next.generation,
+        );
       }
     }
   }
@@ -170,33 +245,42 @@ export class UsageController {
   }
 
   start(ctx: ExtensionContext): void {
-    this.shuttingDown = false;
+    this.usageAbortController?.abort();
+    this.queuedUsageRefresh = undefined;
     this.stopTimer();
+    const generation = ++this.sessionGeneration;
+    this.shuttingDown = false;
+
     const cfg = this.getConfig(ctx);
     if (!cfg.usage.enabled) return;
     const sessionSignal = ctx.signal;
-    if (sessionSignal?.aborted) return;
+    if (sessionSignal?.aborted) {
+      this.deactivateGeneration(generation);
+      return;
+    }
     this.sessionAbortSignal = sessionSignal;
     this.sessionAbortHandler = () => {
-      this.stopTimer();
-      this.usageAbortController?.abort();
-      this.usageAbortController = undefined;
-      this.queuedUsageRefresh = undefined;
+      this.deactivateGeneration(generation);
     };
     sessionSignal?.addEventListener("abort", this.sessionAbortHandler, { once: true });
-    void this.refresh(ctx, undefined, { force: true });
+    void this.refresh(ctx, undefined, { force: true }, generation);
     this.usageTimer = setInterval(() => {
+      if (!this.isGenerationCurrent(generation)) return;
       if (sessionSignal?.aborted) {
-        this.stopTimer();
+        this.deactivateGeneration(generation);
         return;
       }
-      void this.refresh(ctx);
+      void this.refresh(ctx, undefined, undefined, generation);
     }, cfg.usage.refreshIntervalMs);
     this.usageTimer.unref?.();
   }
 
   restartAfterSettingsChange(ctx: ExtensionContext, cfg: ResolvedConfig): void {
+    this.usageAbortController?.abort();
+    this.queuedUsageRefresh = undefined;
     this.stopTimer();
+    this.sessionGeneration++;
+    this.shuttingDown = false;
     if (cfg.usage.enabled) this.start(ctx);
     else {
       this.usageSnapshot = undefined;
@@ -206,6 +290,7 @@ export class UsageController {
 
   shutdown(): void {
     this.shuttingDown = true;
+    this.sessionGeneration++;
     this.queuedUsageRefresh = undefined;
     this.usageAbortController?.abort();
     this.usageAbortController = undefined;