@tokenbuddy/tokenbuddy 1.0.8 → 1.0.11

package/src/model-index.test.ts

@@ -0,0 +1,184 @@
+import { ModelIndex } from "../src/model-index.js";
+import type { RegistrySeller } from "../src/seller-catalog.js";
+
+function makeSeller(overrides: Partial<RegistrySeller> & { id: string; models?: string[] }): RegistrySeller {
+  return {
+    id: overrides.id,
+    name: overrides.name ?? overrides.id,
+    url: overrides.url ?? `https://${overrides.id}.example.com`,
+    supportedProtocols: overrides.supportedProtocols ?? ["chat_completions"],
+    paymentMethods: overrides.paymentMethods ?? ["clawtip"],
+    models: overrides.models
+  };
+}
+
+describe("ModelIndex", () => {
+  test("rebuilds an empty snapshot without throwing", () => {
+    const index = new ModelIndex();
+    index.rebuild([], { registryVersion: 1 });
+
+    expect(index.sellersFor("anything")).toEqual([]);
+    expect(index.knownModelIds()).toEqual([]);
+    const stats = index.stats();
+    expect(stats.sellerCount).toBe(0);
+    expect(stats.modelCount).toBe(0);
+    expect(stats.missingModelsCount).toBe(0);
+    expect(stats.registryVersion).toBe(1);
+  });
+
+  test("rebuild populates byModel and bySeller from a registry snapshot", () => {
+    const index = new ModelIndex();
+    const sellers: RegistrySeller[] = [
+      makeSeller({ id: "s1", models: ["claude-sonnet-4-5", "gpt-4o"] }),
+      makeSeller({ id: "s2", models: ["gpt-4o", "deepseek-v4-pro"] })
+    ];
+
+    index.rebuild(sellers, { registryVersion: 4, defaultSellerId: "s1" });
+
+    expect(index.sellersFor("claude-sonnet-4-5").map((s) => s.id)).toEqual(["s1"]);
+    expect(index.sellersFor("gpt-4o").map((s) => s.id)).toEqual(["s1", "s2"]);
+    expect(index.sellersFor("deepseek-v4-pro").map((s) => s.id)).toEqual(["s2"]);
+    expect(index.getSeller("s1")?.url).toBe("https://s1.example.com");
+    expect(index.getSeller("missing")).toBeUndefined();
+    expect(index.stats().defaultSellerId).toBe("s1");
+    expect(index.stats().registryVersion).toBe(4);
+  });
+
+  test("model id matching is case-insensitive and trims whitespace", () => {
+    const index = new ModelIndex();
+    index.rebuild([makeSeller({ id: "s1", models: ["claude-sonnet-4-5"] })]);
+
+    expect(index.sellersFor("Claude-Sonnet-4-5").map((s) => s.id)).toEqual(["s1"]);
+    expect(index.sellersFor("  claude-sonnet-4-5  ").map((s) => s.id)).toEqual(["s1"]);
+    expect(index.sellersFor("CLAUDE-SONNET-4-5").map((s) => s.id)).toEqual(["s1"]);
+  });
+
+  test("filters by protocol and payment method", () => {
+    const index = new ModelIndex();
+    const sellers: RegistrySeller[] = [
+      makeSeller({
+        id: "s-chat",
+        models: ["gpt-4o"],
+        supportedProtocols: ["chat_completions"],
+        paymentMethods: ["clawtip"]
+      }),
+      makeSeller({
+        id: "s-msg",
+        models: ["gpt-4o"],
+        supportedProtocols: ["messages"],
+        paymentMethods: ["mock"]
+      })
+    ];
+    index.rebuild(sellers);
+
+    const chatClawtip = index.sellersFor("gpt-4o", { protocol: "chat_completions", paymentMethod: "clawtip" });
+    expect(chatClawtip.map((s) => s.id)).toEqual(["s-chat"]);
+
+    const messages = index.sellersFor("gpt-4o", { protocol: "messages" });
+    expect(messages.map((s) => s.id)).toEqual(["s-msg"]);
+
+    const messagesClawtip = index.sellersFor("gpt-4o", { protocol: "messages", paymentMethod: "clawtip" });
+    expect(messagesClawtip).toEqual([]);
+
+    const noFilter = index.sellersFor("gpt-4o");
+    expect(noFilter.map((s) => s.id).sort()).toEqual(["s-chat", "s-msg"]);
+  });
+
+  test("sellers missing models field are excluded from lookups but still addressable by id", () => {
+    const index = new ModelIndex();
+    const sellers: RegistrySeller[] = [
+      makeSeller({ id: "ok", models: ["gpt-4o"] }),
+      makeSeller({ id: "broken" }) // no models
+    ];
+    index.rebuild(sellers);
+
+    expect(index.sellersFor("gpt-4o").map((s) => s.id)).toEqual(["ok"]);
+    expect(index.getSeller("broken")?.id).toBe("broken");
+    expect(index.stats().missingModelsCount).toBe(1);
+    expect(index.stats().sellerCount).toBe(2);
+  });
+
+  test("resolve returns matched flag and snapshot", () => {
+    const index = new ModelIndex();
+    index.rebuild([makeSeller({ id: "s1", models: ["gpt-4o"] })]);
+
+    const hit = index.resolve("gpt-4o");
+    expect(hit.matched).toBe(true);
+    expect(hit.sellers).toHaveLength(1);
+
+    const miss = index.resolve("unknown-model");
+    expect(miss.matched).toBe(false);
+    expect(miss.sellers).toEqual([]);
+  });
+
+  test("rebuild replaces the previous snapshot atomically", () => {
+    const index = new ModelIndex();
+    index.rebuild([makeSeller({ id: "s1", models: ["m1"] })], { registryVersion: 1 });
+    expect(index.sellersFor("m1")).toHaveLength(1);
+
+    index.rebuild([makeSeller({ id: "s2", models: ["m2"] })], { registryVersion: 2 });
+    expect(index.sellersFor("m1")).toEqual([]);
+    expect(index.sellersFor("m2")).toHaveLength(1);
+    expect(index.stats().registryVersion).toBe(2);
+  });
+
+  test("prune drops sellers that have not been seen within staleAfterMs", () => {
+    const index = new ModelIndex();
+    index.rebuild([
+      makeSeller({ id: "s1", models: ["gpt-4o", "m1"] }),
+      makeSeller({ id: "s2", models: ["m2"] })
+    ]);
+
+    const now = 1_700_000_000_000;
+    const lastSeenAt = new Map<string, number>([
+      ["s1", now - 1000], // fresh
+      ["s2", now - 10 * 24 * 60 * 60 * 1000] // 10 days old
+    ]);
+
+    const removed = index.prune(lastSeenAt, 7 * 24 * 60 * 60 * 1000, now);
+    expect(removed).toBe(1);
+    expect(index.getSeller("s1")).toBeDefined();
+    expect(index.getSeller("s2")).toBeUndefined();
+    // Models that belonged only to the pruned seller must be removed.
+    expect(index.sellersFor("m2")).toEqual([]);
+    // Models shared with the remaining seller must still resolve.
+    expect(index.sellersFor("gpt-4o")).toHaveLength(1);
+  });
+
+  test("prune is a no-op when staleAfterMs is non-positive", () => {
+    const index = new ModelIndex();
+    index.rebuild([makeSeller({ id: "s1", models: ["m1"] })]);
+    expect(index.prune(new Map([["s1", 0]]), 0)).toBe(0);
+    expect(index.sellersFor("m1")).toHaveLength(1);
+  });
+
+  test("ignores malformed model entries when populating the index", () => {
+    const index = new ModelIndex();
+    const sellers: RegistrySeller[] = [
+      {
+        id: "s1",
+        name: "s1",
+        url: "https://s1.example.com",
+        supportedProtocols: ["chat_completions"],
+        paymentMethods: ["clawtip"],
+        // Mixed valid / invalid entries; the index must skip the bad ones
+        // without throwing.
+        models: ["m1", "", "   ", "m2"] as string[]
+      }
+    ];
+    index.rebuild(sellers);
+
+    expect(index.sellersFor("m1")).toHaveLength(1);
+    expect(index.sellersFor("m2")).toHaveLength(1);
+    expect(index.knownModelIds().sort()).toEqual(["m1", "m2"]);
+  });
+
+  test("knownModelIds returns a snapshot that does not mutate internal state", () => {
+    const index = new ModelIndex();
+    index.rebuild([makeSeller({ id: "s1", models: ["m1"] })]);
+
+    const ids = index.knownModelIds();
+    ids.push("hacked");
+    expect(index.knownModelIds()).toEqual(["m1"]);
+  });
+});

package/src/model-index.ts

@@ -0,0 +1,266 @@
+import { createModuleLogger } from "@tokenbuddy/logging";
+import type { RegistrySeller } from "./seller-catalog.js";
+
+const logger = createModuleLogger("tb-proxyd:model-index");
+
+/**
+ * Normalize a model id for index lookup. Trims whitespace and lowercases
+ * the value so that `claude-sonnet-4-5` and `Claude-Sonnet-4-5 ` resolve to
+ * the same entry. v1.2 model matching is case-insensitive.
+ */
+function normalizeModelId(modelId: string): string {
+  return modelId.trim().toLowerCase();
+}
+
+/**
+ * Result of resolving a model id against the index.
+ *
+ * - `empty` is returned when the index has been built but no seller serves
+ *   the requested model. The caller should treat this as "no compatible
+ *   seller" and surface the empty result to the user.
+ * - `sellers` is a snapshot array of the matching registry entries. Order
+ *   follows the registry's `defaultSeller` preference then declaration order.
+ */
+export interface ModelIndexResolution {
+  modelId: string;
+  matched: boolean;
+  sellers: RegistrySeller[];
+  missingModelsFlag: number;
+}
+
+interface ModelIndexInternals {
+  byModel: Map<string, RegistrySeller[]>;
+  bySeller: Map<string, RegistrySeller>;
+  registryVersion: number;
+  registryFetchedAt: number;
+  defaultSellerId?: string;
+  missingModelsCount: number;
+  totalSellers: number;
+}
+
+/**
+ * In-memory index mapping `modelId -> sellers[]` from a fetched registry
+ * snapshot. The index is rebuilt atomically on every `rebuild()` call so
+ * callers always observe a consistent snapshot (Node is single-threaded, but
+ * the rebuild path copies data before swapping the internal maps).
+ */
+export class ModelIndex {
+  private internals: ModelIndexInternals = emptyInternals();
+
+  /**
+   * Atomically replace the index contents with a new registry snapshot.
+   * Sellers that do not declare `models` are kept in the `bySeller` reverse
+   * map so they can still be addressed by id, but they are excluded from
+   * `sellersFor()` results. The count of such sellers is exposed via
+   * `missingModelsCount` and via the `models_refresh.seller_missing_models`
+   * log event so operators can fix upstream registry payloads.
+   */
+  rebuild(sellers: RegistrySeller[], opts: { registryVersion?: number; defaultSellerId?: string } = {}): void {
+    const byModel = new Map<string, RegistrySeller[]>();
+    const bySeller = new Map<string, RegistrySeller>();
+    let missingModels = 0;
+
+    for (const seller of sellers) {
+      if (!seller || !seller.id) {
+        continue;
+      }
+      // v1.2 registry schema 用 "anthropic_messages" 作为协议名（OpenAI / ClawTip
+      // / 外部 client 都用这个），但 buyer 内部 `endpointProtocol` 对 /v1/messages
+      // 返 "messages"（更短、更易读）。在 modelIndex 重建时做 alias 映射，让两边
+      // 都能匹配到同一 seller。alias 只在内存里生效，不回写 registry。
+      if (
+        Array.isArray(seller.supportedProtocols) &&
+        seller.supportedProtocols.includes("anthropic_messages") &&
+        !seller.supportedProtocols.includes("messages")
+      ) {
+        seller.supportedProtocols = [...seller.supportedProtocols, "messages"];
+      }
+      bySeller.set(seller.id, seller);
+      if (!Array.isArray(seller.models) || seller.models.length === 0) {
+        missingModels += 1;
+        logger.warn("models_refresh.seller_missing_models", "registry seller entry missing models array; excluded from model index", {
+          sellerId: seller.id,
+          sellerUrl: seller.url
+        });
+        continue;
+      }
+      for (const raw of seller.models) {
+        if (typeof raw !== "string" || raw.trim() === "") {
+          continue;
+        }
+        const key = normalizeModelId(raw);
+        const bucket = byModel.get(key);
+        if (bucket) {
+          bucket.push(seller);
+        } else {
+          byModel.set(key, [seller]);
+        }
+      }
+    }
+
+    this.internals = {
+      byModel,
+      bySeller,
+      registryVersion: opts.registryVersion ?? 0,
+      registryFetchedAt: Date.now(),
+      defaultSellerId: opts.defaultSellerId,
+      missingModelsCount: missingModels,
+      totalSellers: bySeller.size
+    };
+
+    logger.info("models_refresh.rebuilt", "model index rebuilt from registry snapshot", {
+      registryVersion: this.internals.registryVersion,
+      sellerCount: this.internals.totalSellers,
+      modelCount: byModel.size,
+      missingModels: missingModels,
+      defaultSellerId: opts.defaultSellerId ?? null
+    });
+  }
+
+  /**
+   * Resolve sellers for a given model id, optionally filtered by protocol
+   * and payment method. Returns a snapshot; mutating the result does not
+   * affect the index.
+   */
+  sellersFor(modelId: string, filter?: { protocol?: string; paymentMethod?: string }): RegistrySeller[] {
+    const key = normalizeModelId(modelId);
+    const bucket = this.internals.byModel.get(key);
+    if (!bucket) {
+      return [];
+    }
+    if (!filter) {
+      return bucket.slice();
+    }
+    return bucket.filter((seller) => matchesFilter(seller, filter));
+  }
+
+  /**
+   * Resolve a single model id with diagnostic metadata. Used by route-failover
+   * to log structured "no compatible seller" events without losing the
+   * missing-models count.
+   */
+  resolve(modelId: string, filter?: { protocol?: string; paymentMethod?: string }): ModelIndexResolution {
+    return {
+      modelId,
+      matched: this.internals.byModel.has(normalizeModelId(modelId)),
+      sellers: this.sellersFor(modelId, filter),
+      missingModelsFlag: this.internals.missingModelsCount
+    };
+  }
+
+  /**
+   * Returns the registry seller entry by id, or `undefined` if the seller is
+   * not present in the latest snapshot. Used by the route-failover to look up
+   * a seller for token bookkeeping even when the requested model is unknown.
+   */
+  getSeller(sellerId: string): RegistrySeller | undefined {
+    return this.internals.bySeller.get(sellerId);
+  }
+
+  /**
+   * List every model id known to the index. Intended for diagnostics
+   * (`tb doctor`) and CLI completion; not used on the hot path.
+   */
+  knownModelIds(): string[] {
+    return Array.from(this.internals.byModel.keys());
+  }
+
+  /**
+   * Drop stale seller entries from the reverse map when the registry has not
+   * surfaced them within `staleAfterMs`. Model-keyed buckets are derived from
+   * the same `bySeller` set so dropping sellers implicitly drops their
+   * model associations. Returns the number of removed entries.
+   */
+  prune(lastSeenAt: Map<string, number>, staleAfterMs: number, now: number = Date.now()): number {
+    if (staleAfterMs <= 0) {
+      return 0;
+    }
+    const cutoff = now - staleAfterMs;
+    const toRemove: string[] = [];
+    for (const [sellerId, seen] of lastSeenAt.entries()) {
+      if (seen < cutoff && this.internals.bySeller.has(sellerId)) {
+        toRemove.push(sellerId);
+      }
+    }
+    if (toRemove.length === 0) {
+      return 0;
+    }
+
+    const byModel = new Map<string, RegistrySeller[]>();
+    for (const [key, sellers] of this.internals.byModel.entries()) {
+      const filtered = sellers.filter((seller) => !toRemove.includes(seller.id));
+      if (filtered.length > 0) {
+        byModel.set(key, filtered);
+      }
+    }
+    const bySeller = new Map<string, RegistrySeller>();
+    for (const [sellerId, seller] of this.internals.bySeller.entries()) {
+      if (!toRemove.includes(sellerId)) {
+        bySeller.set(sellerId, seller);
+      }
+    }
+
+    this.internals = {
+      ...this.internals,
+      byModel,
+      bySeller,
+      registryFetchedAt: now
+    };
+
+    logger.info("models_refresh.pruned", "stale sellers pruned from model index", {
+      removedCount: toRemove.length,
+      remainingSellers: bySeller.size,
+      remainingModels: byModel.size
+    });
+    return toRemove.length;
+  }
+
+  /**
+   * Snapshot of the internal counters for diagnostics. Cheap to call; does
+   * not copy the maps.
+   */
+  stats(): {
+    sellerCount: number;
+    modelCount: number;
+    missingModelsCount: number;
+    registryVersion: number;
+    registryFetchedAt: number;
+    defaultSellerId?: string;
+  } {
+    return {
+      sellerCount: this.internals.totalSellers,
+      modelCount: this.internals.byModel.size,
+      missingModelsCount: this.internals.missingModelsCount,
+      registryVersion: this.internals.registryVersion,
+      registryFetchedAt: this.internals.registryFetchedAt,
+      defaultSellerId: this.internals.defaultSellerId
+    };
+  }
+}
+
+function emptyInternals(): ModelIndexInternals {
+  return {
+    byModel: new Map(),
+    bySeller: new Map(),
+    registryVersion: 0,
+    registryFetchedAt: 0,
+    missingModelsCount: 0,
+    totalSellers: 0
+  };
+}
+
+function matchesFilter(seller: RegistrySeller, filter: { protocol?: string; paymentMethod?: string }): boolean {
+  if (filter.protocol) {
+    const protocols = seller.supportedProtocols ?? [];
+    if (!protocols.includes(filter.protocol)) {
+      return false;
+    }
+  }
+  if (filter.paymentMethod) {
+    const methods = seller.paymentMethods ?? [];
+    if (!methods.includes(filter.paymentMethod)) {
+      return false;
+    }
+  }
+  return true;
+}