pi-lilac-provider 1.1.0 → 1.2.0

package/index.ts

@@ -304,6 +304,8 @@ function cacheModels(models: JsonModel[]): void {
   } catch {
     // Cache write failure is non-fatal
   }
+  // Base model set changed; force getListModels() to rebuild list prices.
+  listModelsCache = null;
 }
 
 function mergeWithEmbedded(liveModels: JsonModel[], embeddedModels: JsonModel[]): JsonModel[] {
@@ -398,6 +400,41 @@ function applyDiscounts(models: JsonModel[], discounts: Map<string, JsonDiscount
   });
 }
 
+/**
+ * Apply the current discount to a single model's cost IN PLACE.
+ *
+ * applyDiscounts() returns discounted COPIES for the registry; those copies
+ * cannot reach the model pi already bound for the current turn. This mutates
+ * that bound object directly so the current turn's calculateCost() (in pi-ai)
+ * reads the discounted price. Cost is always recomputed from list price — the
+ * patch-applied, pre-discount value from buildModels() — so re-applying a
+ * changed discount never compounds a factor already present on the object.
+ * cacheWrite is left untouched (Lilac does not discount it, matching
+ * applyDiscounts).
+ *
+ * Targets the object pi bound for the turn, captured before any await or
+ * registerProvider() call. registerProvider() refreshes the session model for
+ * SUBSEQUENT turns only (via prepareNextTurn); it cannot affect the current
+ * turn's cost calc, which is why the in-place mutation is required.
+ */
+function applyDiscountInPlace(
+  model: { id: string; cost: { input: number; output: number; cacheRead: number; cacheWrite: number } } | undefined,
+  listModels: JsonModel[],
+  discounts: Map<string, JsonDiscount> | null,
+): void {
+  if (!model?.cost) return;
+  const list = listModels.find(m => m.id === model.id);
+  if (!list) return;
+  // A missing discount entry means list price (factor 1) — e.g. a model whose
+  // discount was removed since it was last priced.
+  const rawFactor = discounts?.get(model.id)?.creditMultiplier;
+  const factor = Number.isFinite(rawFactor) && rawFactor !== undefined ? rawFactor : 1;
+  const applyFactor = (n: number) => n > 0 ? Math.round(n * factor * 10000) / 10000 : n;
+  model.cost.input = applyFactor(list.cost.input);
+  model.cost.output = applyFactor(list.cost.output);
+  model.cost.cacheRead = applyFactor(list.cost.cacheRead);
+}
+
 function cacheDiscounts(discounts: Map<string, JsonDiscount>): void {
   try {
     fs.mkdirSync(CACHE_DIR, { recursive: true });
@@ -467,6 +504,10 @@ let revalidateAbort: AbortController | null = null;
 let latestDiscounts: Map<string, JsonDiscount> | null = null;
 let lastDiscountFetchTime = 0;
 const STATUS_CACHE_TTL_MS = 30000;
+// List-price (patch-applied, pre-discount) models, cached until the base set
+// changes. Reset in cacheModels() so the next getListModels() rebuilds from the
+// refreshed disk cache / embedded set.
+let listModelsCache: JsonModel[] | null = null;
 
 async function resolveApiKey(modelRegistry: ModelRegistry): Promise<void> {
   cachedApiKey = await modelRegistry.getApiKeyForProvider("lilac") ?? undefined;
@@ -479,6 +520,16 @@ export default function (pi: ExtensionAPI) {
   const customModels = customModelsData as JsonModel[];
   const patches = patchData as PatchData;
 
+  // List-price models (patch applied, pre-discount), cached at module scope and
+  // rebuilt only when the base set changes (see cacheModels). Used to recompute
+  // the in-flight model's cost without compounding an already-applied discount.
+  function getListModels(): JsonModel[] {
+    if (!listModelsCache) {
+      listModelsCache = buildModels(loadStaleModels(embeddedModels), customModels, patches);
+    }
+    return listModelsCache;
+  }
+
   const staleBase = loadStaleModels(embeddedModels);
   latestDiscounts = loadCachedDiscounts();
   const staleModels = applyDiscounts(buildModels(staleBase, customModels, patches), latestDiscounts);
@@ -582,10 +633,26 @@ export default function (pi: ExtensionAPI) {
   });
 
   pi.on("before_provider_request", async (_event, ctx) => {
-    if (ctx.model?.provider !== "lilac") return;
-
-    // Always show status for active lilac model
-    ctx.ui.setStatus("lilac", dimStatus(ctx, formatDiscountStatus(ctx.model.id)));
+    // Capture the in-flight model BEFORE any await or registerProvider() call.
+    // ctx.model is a lazy getter to the live session model; at hook start (before
+    // any refresh) it is the same object pi's agent loop bound as the model for
+    // THIS turn — the exact object whose .cost calculateCost() reads when the
+    // response arrives. registerProvider() can only refresh the session model
+    // for SUBSEQUENT turns (prepareNextTurn), so to affect the current turn's
+    // cost we mutate this object in place. Capturing before the await also guards
+    // against a concurrent session_start /status sync swapping the reference
+    // mid-hook.
+    const inFlightModel = ctx.model;
+    if (!inFlightModel || inFlightModel.provider !== "lilac") return;
+
+    const listModels = getListModels();
+
+    // Align the in-flight model with the best-known discount right away. This
+    // covers turns that don't trigger a fresh fetch (within TTL, or the fetch
+    // returned unchanged) and recomputes from list price so it never compounds
+    // a previously-applied factor.
+    applyDiscountInPlace(inFlightModel, listModels, latestDiscounts);
+    ctx.ui.setStatus("lilac", dimStatus(ctx, formatDiscountStatus(inFlightModel.id)));
 
     if (!cachedApiKey) return;
 
@@ -596,24 +663,31 @@ export default function (pi: ExtensionAPI) {
 
     const discounts = await fetchStatusDiscounts(cachedApiKey);
     if (!discounts) return;
+
+    lastDiscountFetchTime = now;
+
     if (!discountsChanged(latestDiscounts, discounts)) {
-      lastDiscountFetchTime = now;
-      ctx.ui.setStatus("lilac", dimStatus(ctx, formatDiscountStatus(ctx.model.id)));
+      ctx.ui.setStatus("lilac", dimStatus(ctx, formatDiscountStatus(inFlightModel.id)));
       return;
     }
 
-    lastDiscountFetchTime = now;
     cacheDiscounts(discounts);
     latestDiscounts = discounts;
 
-    const base = loadStaleModels(embeddedModels);
+    // Re-read list prices in case a concurrent /models sync (session_start)
+    // invalidated the list-price cache during the await above. Then re-apply the
+    // freshly-fetched discount so THIS turn (not just later ones) is costed at
+    // the new price, and re-register so other lilac models pick it up on their
+    // next request.
+    const freshList = getListModels();
+    applyDiscountInPlace(inFlightModel, freshList, discounts);
     pi.registerProvider("lilac", {
       baseUrl: BASE_URL,
       apiKey: "$LILAC_API_KEY",
       api: "openai-completions",
-      models: applyDiscounts(buildModels(base, customModels, patches), discounts),
+      models: applyDiscounts(freshList, discounts),
     });
-    ctx.ui.setStatus("lilac", dimStatus(ctx, formatDiscountStatus(ctx.model.id)));
+    ctx.ui.setStatus("lilac", dimStatus(ctx, formatDiscountStatus(inFlightModel.id)));
   });
 
   pi.on("model_select", async (event, ctx) => {
@@ -668,5 +742,5 @@ export default function (pi: ExtensionAPI) {
   });
 }
 
-export { fetchStatusDiscounts, applyDiscounts, loadCachedDiscounts, cacheDiscounts };
+export { fetchStatusDiscounts, applyDiscounts, applyDiscountInPlace, loadCachedDiscounts, cacheDiscounts };
 export type { JsonDiscount, JsonModel, PatchEntry, PatchData };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-lilac-provider",
-  "version": "1.1.0",
+  "version": "1.2.0",
   "description": "Lilac provider extension for pi - Access Kimi K2.6, GLM 5.1, and Gemma 4 models through Lilac's OpenAI-compatible API on idle GPUs",
   "type": "module",
   "main": "index.ts",

package/scripts/test-discounts.ts

@@ -13,6 +13,10 @@
  * 7. turn_end appends discount entry to session JSONL.
  * 8. before_provider_request refreshes discounts with a 30s cache.
  * 9. formatDiscountStatus returns fallbacks when data is missing.
+ * 10. applyDiscountInPlace mutates the in-flight model's cost in place,
+ *     recomputed from list price so re-applied discounts never compound.
+ * 11. before_provider_request mutates the bound (in-flight) model object so the
+ *     current turn's cost calc sees the discount in real time.
  */
 
 import type { ExtensionAPI, ModelRegistry } from "@earendil-works/pi-coding-agent";
@@ -43,6 +47,7 @@ const {
   default: registerLilac,
   fetchStatusDiscounts,
   applyDiscounts,
+  applyDiscountInPlace,
   loadCachedDiscounts,
   cacheDiscounts,
 } = await import("../index.ts");
@@ -446,6 +451,96 @@ const kimiFinal = afterSecond.config.models.find((m: any) => m.id === "moonshota
 assert(kimiFinal.discount.discountPercent === 25, "final registration uses correct discount (not stale from first re-register)");
 assert(kimiFinal.cost.input === 0.525, "final registration uses correct cost (0.70 * 0.75)");
 
+// ─── Test 12: applyDiscountInPlace mutates the in-flight model ───────────────
+
+console.log("\n--- Test 12: applyDiscountInPlace (in-flight mutation) ---");
+
+// applyDiscountInPlace closes the one-turn gap: registerProvider() refreshes
+// the registry for LATER turns, but the model pi already bound for the current
+// turn is a separate object whose .cost calculateCost() reads. We mutate THAT
+// object in place, recomputed from list price so a changed discount never
+// compounds on top of an already-applied factor.
+
+const listModels = [
+  { id: "moonshotai/kimi-k2.6", name: "Kimi K2.6", cost: { input: 0.70, output: 3.50, cacheRead: 0.20, cacheWrite: 0 } },
+] as any[];
+
+// A model object that is "already bound" for the turn — the same reference pi's
+// agent loop holds. It currently carries an OLD discount (factor 0.75).
+const inFlight = {
+  id: "moonshotai/kimi-k2.6",
+  cost: { input: 0.525, output: 2.625, cacheRead: 0.15, cacheWrite: 0 },
+};
+const inFlightRef = inFlight;
+
+// New discount: factor 0.50 (changed). Must recompute from LIST price (0.70),
+// NOT compound on top of the already-applied 0.75 (which would give 0.2625).
+const newDiscounts = new Map([
+  ["moonshotai/kimi-k2.6", { supplyState: "high", discountPercent: 50, creditMultiplier: 0.50 }],
+]);
+applyDiscountInPlace(inFlight, listModels, newDiscounts);
+assert(inFlight === inFlightRef, "mutates the same object in place (no new object created)");
+assert(inFlight.cost.input === 0.35, "recomputed from list price: 0.70 * 0.50 = 0.35 (not compounded 0.2625)");
+assert(inFlight.cost.output === 1.75, "output = 3.50 * 0.50 = 1.75");
+assert(inFlight.cost.cacheRead === 0.1, "cacheRead = 0.20 * 0.50 = 0.1");
+assert(inFlight.cost.cacheWrite === 0, "cacheWrite left untouched (Lilac does not discount it)");
+
+// Idempotency + no-compounding: applying the same factor again yields the same
+// values (always recomputed from list price, never compounding).
+applyDiscountInPlace(inFlight, listModels, newDiscounts);
+assert(inFlight.cost.input === 0.35, "re-applying same factor is idempotent (no compounding)");
+
+// Discount removed (model no longer in /status) → reverts to list price.
+const emptyDiscounts = new Map();
+applyDiscountInPlace(inFlight, listModels, emptyDiscounts);
+assert(inFlight.cost.input === 0.70, "removed discount reverts to list price (factor 1)");
+assert(inFlight.cost.output === 3.50, "removed discount: output reverts to list price");
+
+// null discounts → list price, no throw.
+applyDiscountInPlace(inFlight, listModels, null);
+assert(inFlight.cost.input === 0.70, "null discounts → list price");
+
+// Unknown model id (not in listModels) → no-op, no throw.
+const unknown = { id: "some/unknown-model", cost: { input: 9.99, output: 9.99, cacheRead: 9.99, cacheWrite: 0 } };
+applyDiscountInPlace(unknown, listModels, newDiscounts);
+assert(unknown.cost.input === 9.99, "unknown model id → no mutation (no list-price match)");
+
+// undefined model → no throw.
+applyDiscountInPlace(undefined, listModels, newDiscounts);
+assert(true, "undefined model does not throw");
+
+// ─── Test 13: before_provider_request mutates the bound (in-flight) model ────
+
+console.log("\n--- Test 13: before_provider_request mutates the bound model ---");
+
+// Simulate the object pi's agent loop bound for the turn. It holds an OUTDATED
+// cost (list price, as if no discount had been applied yet). This same reference
+// is what calculateCost() reads for the current turn — so the handler MUST mutate
+// it in place, not replace it. latestDiscounts currently holds kimi at 0.75 (from
+// Test 4), and the TTL is fresh, so the handler takes the within-TTL path (which
+// previously did NOT touch the in-flight model at all).
+const boundModel = {
+  id: "moonshotai/kimi-k2.6",
+  provider: "lilac",
+  cost: { input: 0.70, output: 3.50, cacheRead: 0.20, cacheWrite: 0 }, // list price (stale)
+};
+const boundRef = boundModel;
+
+for (const handler of handlers.get("before_provider_request") || []) {
+  await handler(
+    { type: "before_provider_request", payload: {} },
+    {
+      ui: mockUi,
+      model: boundModel,
+    }
+  );
+}
+
+assert(boundModel === boundRef, "handler did not replace the bound model object (same reference)");
+assert(boundModel.cost.input === 0.525, "bound model cost mutated in place to 0.70 * 0.75 = 0.525");
+assert(boundModel.cost.output === 2.625, "bound model output mutated to 3.50 * 0.75 = 2.625");
+assert(boundModel.cost.cacheRead === 0.15, "bound model cacheRead mutated to 0.20 * 0.75 = 0.15");
+
 // ─── Cleanup ──────────────────────────────────────────────────────────────────
 
 globalThis.fetch = originalFetch;