pi-lilac-provider 1.6.2 → 1.7.0

package/README.md

@@ -23,6 +23,7 @@ Access Kimi K2.6, GLM 5.1, MiniMax M2.7, and Gemma 4 models through Lilac's Open
 - **Reasoning Models** — Chain-of-thought via `chat_template_kwargs` (all models)
 - **Vision Support** — Image input on Kimi K2.6 and Gemma 4
 - **Context Caching** — Cache read pricing on Kimi K2.6 and GLM 5.1
+- **Flex (Discount Gating)** — Only let the LLM respond when the active model's discount meets a threshold you set (`/lilac-flex`)
 - **Idle GPU Scheduling** — Lilac leverages idle GPU capacity for cost-efficient inference
 
 ## Installation
@@ -218,6 +219,9 @@ Create `~/.pi/agent/extensions/lilac.json` (auto-populated with defaults on firs
 
 ```jsonc
 {
+  // Only respond when the active model's discount is >= this percent. null = off.
+  // See "Flex (Discount Gating)" below. Set interactively with /lilac-flex.
+  "flexThreshold": null,
   "modelOverrides": {
     // Disable full-history reasoning for kimi-k2.6 (e.g. to save tokens):
     "moonshotai/kimi-k2.6": { "compat": { "chatTemplateKwargs": { "preserve_thinking": false } } },
@@ -231,6 +235,29 @@ Create `~/.pi/agent/extensions/lilac.json` (auto-populated with defaults on firs
 
 The full set of overridable fields matches the model schema (`compat`, `thinkingLevelMap`, `cost`, `contextWindow`, `maxTokens`, `reasoning`, `input`). See [Compat Settings](#compat-settings) for the catalog of compat flags and what `chatTemplateKwargs` values mean per family. An invalid JSON file is left untouched (defaults are used) so a typo isn't silently wiped — fix the file and restart pi.
 
+### Flex (Discount Gating)
+
+Lilac's per-model discount fluctuates with idle-GPU supply. **Flex** lets you set a discount threshold so pi **only sends a prompt to the LLM when the active model's current discount is at or above it** — e.g. "only respond when the discount is ≥ 75%". Below the threshold, the prompt is blocked (dropped with a warning) until the next discount poll brings the discount back up. This is a spend-control feature: you only spend when supply is cheap.
+
+Set it interactively with the `/lilac-flex` command:
+
+```
+/lilac-flex          # picker: Off / ≥50% / ≥75% / Custom…
+/lilac-flex 75       # set threshold directly (only respond at ≥75% discount)
+/lilac-flex 50%      # trailing % accepted on the command line
+/lilac-flex off      # disable flex (allow all discounts)
+```
+
+The threshold persists in `~/.pi/agent/extensions/lilac.json` as `flexThreshold` (a number `0`–`100`, or `null` for off) alongside `modelOverrides`, so it survives restarts. `/lilac-flex` updates it live — no restart needed.
+
+Behavior notes:
+
+- **Gating point:** flex checks at prompt-submission time. When blocked, the prompt is dropped (you get a warning notification) and you re-submit once the discount improves. It does **not** queue the prompt.
+- **Scope:** only **interactive** (TUI-typed) prompts are gated. `rpc`/`print` (automation) and extension-injected messages are not gated, so flex never causes a silent failure in a pipeline or an extension loop. Flex only applies to the active **lilac** model; non-lilac models always pass.
+- **No data / no discount entry = 0%.** A lilac model with no discount (list price), or before the first discount poll has data, counts as 0% and is blocked when flex is on. This matches how discounts are priced elsewhere in the extension.
+- **Freshness:** when a prompt is blocked, the extension triggers an immediate `/status` refresh (throttled to once per ~5s) so you're not stuck on a stale low value from the 5-minute idle poll. The next submission sees the fresh discount. The footer status reflects the gate: `… · flex ≥75% ok` or `… · flex ≥75% blocked`.
+- **Discount lock-in:** per Lilac, a discount is locked in when a request starts. Flex gates on the best-known discount at submit time, which is what gets locked in for that turn.
+
 ## Updating Models
 
 Run the update script to fetch the latest models from Lilac's API:

package/index.ts

@@ -192,10 +192,14 @@ interface ModelOverride {
 
 interface LilacConfig {
   modelOverrides?: Record<string, ModelOverride>;
+  // Flex discount threshold: only allow interactive prompts to reach the LLM
+  // when the active lilac model's discountPercent is >= this. null/undefined =
+  // off (allow all). Set via /lilac-flex; persisted in lilac.json.
+  flexThreshold?: number | null;
 }
 
 const CONFIG_PATH = path.join(os.homedir(), ".pi", "agent", "extensions", "lilac.json");
-const DEFAULT_CONFIG: LilacConfig = { modelOverrides: {} };
+const DEFAULT_CONFIG: LilacConfig = { modelOverrides: {}, flexThreshold: null };
 
 // Validate user-supplied modelOverrides from the config file. Non-object ids and
 // non-object overrides are dropped silently so a malformed file doesn't crash
@@ -220,6 +224,23 @@ function parseModelOverrides(raw: unknown): Record<string, ModelOverride> | unde
   return Object.keys(result).length > 0 ? result : undefined;
 }
 
+// Validate a flex discount threshold from the config file. A number in [0,100] is
+// returned as-is; null / "off" / "" map to null (flex disabled); anything else
+// returns undefined (treated as disabled by the gate, but not normalized so an
+// invalid value isn't silently rewritten). Mirrors parseModelOverrides' lenient
+// parsing so a malformed file doesn't crash the gate or model registration.
+function parseFlexThreshold(raw: unknown): number | null | undefined {
+  if (raw === null) return null;
+  if (typeof raw === "number" && Number.isFinite(raw) && raw >= 0 && raw <= 100) return raw;
+  if (typeof raw === "string") {
+    const s = raw.trim().toLowerCase();
+    if (s === "" || s === "off" || s === "none") return null;
+    const n = Number(s);
+    if (Number.isFinite(n) && n >= 0 && n <= 100) return n;
+  }
+  return undefined;
+}
+
 // Reads ~/.pi/agent/extensions/lilac.json. Missing file → populate with defaults
 // so the user can discover it, then return defaults. An existing-but-invalid file
 // is left untouched (defaults returned) so a user's typo isn't silently wiped —
@@ -241,7 +262,10 @@ function loadConfig(): LilacConfig {
   }
   try {
     const raw = JSON.parse(rawText);
-    return { modelOverrides: parseModelOverrides(raw.modelOverrides) };
+    return {
+      modelOverrides: parseModelOverrides(raw.modelOverrides),
+      flexThreshold: parseFlexThreshold(raw.flexThreshold),
+    };
   } catch {
     // File exists but is invalid JSON — return defaults WITHOUT overwriting.
     return { ...DEFAULT_CONFIG };
@@ -288,6 +312,29 @@ function getConfig(): LilacConfig {
   return config;
 }
 
+// Read-modify-write the config file and refresh the in-memory cache. Used by
+// /lilac-flex so a threshold change takes effect immediately (the input gate and
+// footer read getConfig()) without a restart, and without clobbering the user's
+// modelOverrides. Reads via loadConfig (which validates), so modelOverrides the
+// user hand-edited since startup survive the spread. The file is normalized to a
+// discoverable shape (modelOverrides: {}, flexThreshold: null always present) so
+// /lilac-flex never strips the modelOverrides scaffold from the file.
+function updateConfig(mutator: (cfg: LilacConfig) => LilacConfig): LilacConfig {
+  const next = mutator(loadConfig());
+  const toWrite: LilacConfig = {
+    modelOverrides: next.modelOverrides ?? {},
+    flexThreshold: next.flexThreshold ?? null,
+  };
+  try {
+    fs.mkdirSync(path.dirname(CONFIG_PATH), { recursive: true });
+    fs.writeFileSync(CONFIG_PATH, JSON.stringify(toWrite, null, 2) + "\n");
+  } catch {
+    // Write failure is non-fatal — the in-memory cache below still updates.
+  }
+  config = next;
+  return next;
+}
+
 function activeOverrides(): Record<string, ModelOverride> {
   return getConfig().modelOverrides ?? {};
 }
@@ -656,10 +703,19 @@ function loadCachedDiscounts(): Map<string, JsonDiscount> | null {
 
 function formatDiscountStatus(modelId?: string): string {
   if (!modelId) return "supply: —";
-  if (!latestDiscounts) return "supply: checking…";
-  const discount = latestDiscounts.get(modelId);
-  if (!discount) return "supply: —";
-  return `supply: ${discount.supplyState} · sub-discount: ${discount.discountPercent}%`;
+  const threshold = getConfig().flexThreshold ?? null;
+  const discount = latestDiscounts?.get(modelId);
+  let base: string;
+  if (discount) {
+    base = `supply: ${discount.supplyState} · sub-discount: ${discount.discountPercent}%`;
+  } else if (latestDiscounts) {
+    base = "supply: —";
+  } else {
+    base = "supply: checking…";
+  }
+  if (threshold == null) return base;
+  const pct = discount?.discountPercent ?? 0;
+  return `${base} · flex ≥${threshold}% ${pct >= threshold ? "ok" : "blocked"}`;
 }
 
 function dimStatus(ctx: any, text: string): string {
@@ -705,6 +761,50 @@ function syncStatus(ctx: any): void {
   }
 }
 
+// Parse a /lilac-flex argument or custom-threshold input: "off"/"none"/"" → null
+// (disabled), a number in [0,100] (optionally with a trailing %) → that number,
+// anything else → undefined (invalid).
+function parseFlexArg(raw: string): number | null | undefined {
+  const s = raw.trim().toLowerCase().replace(/%$/, "");
+  if (s === "" || s === "off" || s === "none") return null;
+  const n = Number(s);
+  if (!Number.isFinite(n) || n < 0 || n > 100) return undefined;
+  return Math.round(n);
+}
+
+// Persist a flex threshold via updateConfig and report the resulting state against
+// the live model's current discount. The footer is re-painted (syncStatus) so the
+// flex indicator appears immediately. Used by the /lilac-flex command.
+function applyFlexThreshold(value: number | null, ctx: any): void {
+  updateConfig((cfg) => ({ ...cfg, flexThreshold: value }));
+  const threshold = getConfig().flexThreshold ?? null;
+  const model = ctx.model;
+  const discount = model?.provider === "lilac" ? latestDiscounts?.get(model.id) : undefined;
+  const pct = discount?.discountPercent;
+
+  let msg: string;
+  let level: "info" | "warning";
+  if (threshold == null) {
+    msg = "lilac-flex: off — all discounts allowed";
+    level = "info";
+  } else if (pct == null) {
+    msg = `lilac-flex: ≥ ${threshold}% — no discount data yet; will block until the next poll`;
+    level = "warning";
+  } else if (pct >= threshold) {
+    msg = `lilac-flex: ≥ ${threshold}% — current ${pct}% discount allowed`;
+    level = "info";
+  } else {
+    msg = `lilac-flex: ≥ ${threshold}% — current ${pct}% discount blocked until it improves`;
+    level = "warning";
+  }
+  try {
+    ctx.ui.notify(msg, level);
+  } catch {
+    // notify is a no-op without a UI runner
+  }
+  syncStatus(ctx);
+}
+
 function discountsChanged(
   a: Map<string, JsonDiscount> | null,
   b: Map<string, JsonDiscount> | null,
@@ -735,6 +835,10 @@ let lastDiscountFetchTime = 0;
 // they cooperate: a poll that just ran lets the next turn skip its own fetch
 // within the TTL.
 const STATUS_CACHE_TTL_MS = 60000;
+// Throttle for the immediate /status refresh fired when flex blocks a prompt, so
+// a user retrying repeatedly doesn't hammer the endpoint while still getting fresh
+// data within seconds (vs. waiting on the 5-min idle poll).
+const BLOCK_REFRESH_THROTTLE_MS = 5000;
 // Lilac refreshes discounts ~every 10 minutes (per their docs: "Discounts refresh
 // approximately every 10 minutes and are locked in when a request starts"). Poll
 // every 5 min during idle — half the refresh window — so a long-idle session
@@ -803,6 +907,22 @@ export default function (pi: ExtensionAPI) {
     }
   }
 
+  // Cache + hot-swap freshly fetched discounts: update latestDiscounts, re-register
+  // the provider so other lilac models pick up the new price on their next request,
+  // and re-paint the footer from the LIVE model. Shared by the idle poll and the
+  // flex block refresh so the apply logic can't drift between them.
+  function applyFreshDiscounts(discounts: Map<string, JsonDiscount>, ctx: any): void {
+    cacheDiscounts(discounts);
+    latestDiscounts = discounts;
+    pi.registerProvider("lilac", {
+      baseUrl: BASE_URL,
+      apiKey: "$LILAC_API_KEY",
+      api: "openai-completions",
+      models: applyDiscounts(getListModels(), discounts),
+    });
+    syncStatus(ctx);
+  }
+
   /**
    * Background /status poll, fired every STATUS_POLL_INTERVAL_MS (5 min) from
    * session_start to cover idle sessions. Mirrors the discount half of
@@ -821,16 +941,25 @@ export default function (pi: ExtensionAPI) {
         syncStatus(ctx);
         return;
       }
-      cacheDiscounts(discounts);
-      latestDiscounts = discounts;
-      const freshList = getListModels();
-      pi.registerProvider("lilac", {
-        baseUrl: BASE_URL,
-        apiKey: "$LILAC_API_KEY",
-        api: "openai-completions",
-        models: applyDiscounts(freshList, discounts),
-      });
-      syncStatus(ctx);
+      applyFreshDiscounts(discounts, ctx);
+    }).catch(() => { /* network errors are non-fatal */ });
+  }
+
+  // Fire-and-forget /status refresh used when flex blocks a prompt, so the user
+  // isn't stuck on a stale low value from the 5-min idle poll. Throttled by
+  // BLOCK_REFRESH_THROTTLE_MS so repeated retries don't hammer the endpoint. Not
+  // signal-guarded (user-initiated, not session-scoped); network errors are non-fatal.
+  function triggerDiscountRefresh(ctx: any): void {
+    if (!cachedApiKey) return;
+    if (Date.now() - lastDiscountFetchTime < BLOCK_REFRESH_THROTTLE_MS) return;
+    fetchStatusDiscounts(cachedApiKey).then(discounts => {
+      if (!discounts) return;
+      lastDiscountFetchTime = Date.now();
+      if (!discountsChanged(latestDiscounts, discounts)) {
+        syncStatus(ctx);
+        return;
+      }
+      applyFreshDiscounts(discounts, ctx);
     }).catch(() => { /* network errors are non-fatal */ });
   }
 
@@ -1025,6 +1154,84 @@ export default function (pi: ExtensionAPI) {
     };
   });
 
+  // lilac-flex: gate interactive prompts on the active lilac model's discount.
+  // Only interactive prompts are gated — extension-injected messages are skipped
+  // (programmatic; would loop) and rpc/print are skipped (a silent block would be
+  // a confusing failure in automation). A missing discount entry or no data yet
+  // counts as 0% (list price), consistent with applyDiscounts/applyDiscountInPlace.
+  // When blocked, the prompt is dropped (handled) and an immediate /status refresh
+  // is triggered so the next submission isn't waiting on the 5-min idle poll.
+  pi.on("input", async (event, ctx) => {
+    if (event.source !== "interactive") return { action: "continue" };
+
+    const threshold = getConfig().flexThreshold ?? null;
+    if (threshold == null) return { action: "continue" };
+
+    const model = ctx.model;
+    if (!model || model.provider !== "lilac") return { action: "continue" };
+
+    const discount = latestDiscounts?.get(model.id);
+    const discountPercent = discount?.discountPercent ?? 0;
+    if (discountPercent >= threshold) return { action: "continue" };
+
+    const desc = discount
+      ? `${discountPercent}% discount`
+      : (latestDiscounts ? "no discount on this model (list price)" : "no discount data yet");
+    ctx.ui.notify(
+      `lilac-flex: ${desc} < ${threshold}% threshold — blocked until the next poll. Re-submit once the discount improves, or run /lilac-flex to adjust.`,
+      "warning",
+    );
+    triggerDiscountRefresh(ctx);
+    return { action: "handled" };
+  });
+
+  pi.registerCommand("lilac-flex", {
+    description: "Set lilac flex discount threshold — only respond at/above this discount",
+    async handler(args, ctx) {
+      if (!ctx.hasUI) {
+        ctx.ui.notify("/lilac-flex requires interactive mode.", "error");
+        return;
+      }
+
+      const arg = (args ?? "").trim();
+      if (arg) {
+        const value = parseFlexArg(arg);
+        if (value === undefined) {
+          ctx.ui.notify("Usage: /lilac-flex [off | 0-100]", "warning");
+          return;
+        }
+        applyFlexThreshold(value, ctx);
+        return;
+      }
+
+      const current = getConfig().flexThreshold ?? null;
+      const labels = [
+        "Off — allow all discounts",
+        "≥ 50% discount",
+        "≥ 75% discount",
+        "Custom…",
+      ];
+      const choice = await ctx.ui.select("Lilac flex: only respond at/above this discount", labels);
+      if (choice === undefined) return; // cancelled
+
+      let value: number | null;
+      if (choice === labels[0]) value = null;
+      else if (choice === labels[1]) value = 50;
+      else if (choice === labels[2]) value = 75;
+      else {
+        const input = await ctx.ui.input("Threshold (0–100, or 'off'):", String(current ?? ""));
+        if (input === undefined) return; // cancelled
+        const parsed = parseFlexArg(input);
+        if (parsed === undefined) {
+          ctx.ui.notify("Invalid threshold. Use a number 0–100 or 'off'.", "warning");
+          return;
+        }
+        value = parsed;
+      }
+      applyFlexThreshold(value, ctx);
+    },
+  });
+
   pi.on("session_shutdown", () => {
     revalidateAbort?.abort();
     if (pollInterval) {
@@ -1034,5 +1241,5 @@ export default function (pi: ExtensionAPI) {
   });
 }
 
-export { fetchStatusDiscounts, applyDiscounts, applyDiscountInPlace, loadCachedDiscounts, cacheDiscounts, buildModels, applyModelOverride, parseModelOverrides, loadConfig, getConfig };
+export { fetchStatusDiscounts, applyDiscounts, applyDiscountInPlace, loadCachedDiscounts, cacheDiscounts, buildModels, applyModelOverride, parseModelOverrides, parseFlexThreshold, updateConfig, loadConfig, getConfig };
 export type { JsonDiscount, JsonModel, PatchEntry, PatchData, ModelOverride, LilacConfig };

package/models.json

@@ -102,7 +102,7 @@
     "cost": {
       "input": 0.3,
       "output": 1.2,
-      "cacheRead": 0.06,
+      "cacheRead": 0.055,
       "cacheWrite": 0
     },
     "contextWindow": 204800,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-lilac-provider",
-  "version": "1.6.2",
+  "version": "1.7.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",
@@ -34,10 +34,11 @@
     "clean": "echo 'nothing to clean'",
     "build": "echo 'nothing to build'",
     "check": "echo 'nothing to check'",
-    "test": "node scripts/test-discounts.ts && node scripts/test-preserved-thinking.ts && node scripts/test-model-overrides.ts",
+    "test": "node scripts/test-discounts.ts && node scripts/test-preserved-thinking.ts && node scripts/test-model-overrides.ts && node scripts/test-flex.ts",
     "test:discounts": "node scripts/test-discounts.ts",
     "test:thinking": "node scripts/test-preserved-thinking.ts",
     "test:overrides": "node scripts/test-model-overrides.ts",
+    "test:flex": "node scripts/test-flex.ts",
     "update-models": "node scripts/update-models.js"
   }
 }

package/scripts/test-flex.ts

@@ -0,0 +1,370 @@
+#!/usr/bin/env node
+/**
+ * Tests for the lilac-flex feature: discount-threshold gating.
+ *
+ * Verifies, against the REAL exported helpers and registered handlers/commands
+ * from index.ts (not a re-implementation):
+ *   - parseFlexThreshold: number in [0,100] passes; null/"off"/"" -> null;
+ *     out-of-range / non-numeric / object -> undefined.
+ *   - loadConfig: parses flexThreshold from an existing file; missing key ->
+ *     undefined (off); invalid value -> undefined while modelOverrides survive.
+ *   - updateConfig: sets flexThreshold, persists to disk, and PRESERVES the
+ *     user's modelOverrides (doesn't clobber them); normalizes the file shape.
+ *   - input gate (integration via the registered handler): blocks (handled) when
+ *     discount < threshold; allows (continue) when >= threshold, when flex is
+ *     off, for non-lilac models, for non-interactive sources, and for models
+ *     with no discount entry (treated as 0% -> blocked when flex on).
+ *   - /lilac-flex command (integration via the registered handler): direct arg
+ *     ("/lilac-flex 75", "off", "50%"), picker presets, custom input, cancel,
+ *     invalid arg, and non-interactive guard.
+ *
+ * Config FS + discount cache are isolated to a temp HOME so nothing touches the
+ * real ~/.pi.
+ */
+
+import fs from "fs";
+import os from "os";
+import path from "path";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+// Isolate config + cache to a temp HOME so loadConfig/cacheDiscounts never touch
+// the real ~/.pi. Must be set before importing index.ts, which computes
+// CONFIG_PATH / CACHE_PATH at module scope.
+const tmpHome = `/tmp/pi-lilac-flex-test-${Date.now()}`;
+fs.mkdirSync(tmpHome, { recursive: true });
+process.env.HOME = tmpHome;
+
+const {
+  default: registerLilac,
+  parseFlexThreshold,
+  loadConfig,
+  getConfig,
+  updateConfig,
+  cacheDiscounts,
+} = await import("../index.ts");
+
+let passed = 0;
+let failed = 0;
+function assert(condition: boolean, message: string) {
+  if (condition) {
+    console.log(`  ✓ ${message}`);
+    passed++;
+  } else {
+    console.error(`  ✗ ${message}`);
+    failed++;
+  }
+}
+function eq<T>(actual: T, expected: T, message: string) {
+  const ok = JSON.stringify(actual) === JSON.stringify(expected);
+  if (ok) {
+    console.log(`  ✓ ${message}`);
+    passed++;
+  } else {
+    console.error(`  ✗ ${message}\n     expected: ${JSON.stringify(expected)}\n     actual:   ${JSON.stringify(actual)}`);
+    failed++;
+  }
+}
+
+const cfgPath = path.join(os.homedir(), ".pi", "agent", "extensions", "lilac.json");
+const KIMI = "moonshotai/kimi-k2.6";
+
+// ─── parseFlexThreshold ───────────────────────────────────────────────────────
+
+console.log("\n--- parseFlexThreshold ---");
+eq(parseFlexThreshold(null), null, "null -> null (off)");
+eq(parseFlexThreshold(0), 0, "0 -> 0");
+eq(parseFlexThreshold(100), 100, "100 -> 100");
+eq(parseFlexThreshold(75), 75, "75 -> 75");
+eq(parseFlexThreshold(50.5), 50.5, "fractional number preserved");
+eq(parseFlexThreshold(-1), undefined, "negative -> undefined");
+eq(parseFlexThreshold(101), undefined, ">100 -> undefined");
+eq(parseFlexThreshold(NaN), undefined, "NaN -> undefined");
+eq(parseFlexThreshold("75"), 75, "numeric string -> number");
+// Note: parseFlexThreshold is strict for the JSON file (number or off/none/null).
+// Trailing-% stripping lives in parseFlexArg for typed command/input args.
+eq(parseFlexThreshold("50%"), undefined, "string with trailing % -> undefined (strict; parseFlexArg handles %)");
+eq(parseFlexThreshold("off"), null, "'off' -> null");
+eq(parseFlexThreshold("OFF"), null, "'OFF' (case-insensitive) -> null");
+eq(parseFlexThreshold("none"), null, "'none' -> null");
+eq(parseFlexThreshold(""), null, "empty string -> null");
+eq(parseFlexThreshold("abc"), undefined, "non-numeric string -> undefined");
+eq(parseFlexThreshold(true), undefined, "boolean -> undefined");
+eq(parseFlexThreshold({ a: 1 }), undefined, "object -> undefined");
+eq(parseFlexThreshold([75]), undefined, "array -> undefined");
+
+// ─── loadConfig: flexThreshold parsing ────────────────────────────────────────
+
+console.log("\n--- loadConfig: flexThreshold ---");
+
+// The config dir doesn't exist yet in the temp HOME; direct writes below bypass
+// loadConfig's missing-file scaffold, so create it first.
+fs.mkdirSync(path.dirname(cfgPath), { recursive: true });
+
+{
+  fs.writeFileSync(cfgPath, JSON.stringify({ flexThreshold: 75 }));
+  const cfg = loadConfig();
+  assert(cfg.flexThreshold === 75, "file with flexThreshold:75 -> 75");
+}
+
+{
+  fs.writeFileSync(cfgPath, JSON.stringify({ flexThreshold: null }));
+  const cfg = loadConfig();
+  assert(cfg.flexThreshold === null, "file with flexThreshold:null -> null");
+}
+
+{
+  fs.writeFileSync(cfgPath, JSON.stringify({ unrelated: true }));
+  const cfg = loadConfig();
+  assert(cfg.flexThreshold === undefined, "file without flexThreshold key -> undefined (off)");
+}
+
+{
+  // Invalid flexThreshold falls back to undefined, but modelOverrides still parse.
+  fs.writeFileSync(cfgPath, JSON.stringify({
+    flexThreshold: "not-a-number",
+    modelOverrides: { [KIMI]: { compat: { chatTemplateKwargs: { preserve_thinking: false } } } },
+  }));
+  const cfg = loadConfig();
+  assert(cfg.flexThreshold === undefined, "invalid flexThreshold -> undefined (off)");
+  assert((cfg.modelOverrides as any)?.[KIMI]?.compat?.chatTemplateKwargs?.preserve_thinking === false, "modelOverrides still parsed alongside invalid flexThreshold");
+}
+
+// ─── updateConfig: round-trip + modelOverrides preservation ───────────────────
+
+console.log("\n--- updateConfig ---");
+
+{
+  // Start from a file with real modelOverrides, then set flex via updateConfig:
+  // the overrides must survive (read-merge-write, not clobber).
+  fs.writeFileSync(cfgPath, JSON.stringify({
+    modelOverrides: { [KIMI]: { compat: { chatTemplateKwargs: { preserve_thinking: false } } } },
+  }));
+  updateConfig((c) => ({ ...c, flexThreshold: 75 }));
+  const cfg = getConfig();
+  assert(cfg.flexThreshold === 75, "updateConfig sets flexThreshold to 75");
+  assert((cfg.modelOverrides as any)?.[KIMI]?.compat?.chatTemplateKwargs?.preserve_thinking === false, "updateConfig preserves existing modelOverrides");
+  const onDisk = JSON.parse(fs.readFileSync(cfgPath, "utf8"));
+  assert(onDisk.flexThreshold === 75, "file on disk has flexThreshold 75");
+  assert(onDisk.modelOverrides[KIMI].compat.chatTemplateKwargs.preserve_thinking === false, "file on disk preserves modelOverrides");
+}
+
+{
+  // Setting to null persists null and keeps the modelOverrides scaffold shape.
+  updateConfig((c) => ({ ...c, flexThreshold: null }));
+  const onDisk = JSON.parse(fs.readFileSync(cfgPath, "utf8"));
+  assert(onDisk.flexThreshold === null, "updateConfig null -> file has flexThreshold null");
+  assert(onDisk.modelOverrides[KIMI].compat.chatTemplateKwargs.preserve_thinking === false, "modelOverrides still preserved after setting flex off");
+}
+
+// ─── Setup: register extension + seed latestDiscounts from cache ──────────────
+
+console.log("\n--- input gate ---");
+
+// Seed the discount cache BEFORE registering so the extension's init picks it up
+// via loadCachedDiscounts() (latestDiscounts). Kimi at 25%, glm-5.1 uncached.
+cacheDiscounts(new Map([
+  [KIMI, { supplyState: "medium", discountPercent: 25, creditMultiplier: 0.75 }],
+]));
+
+const handlers = new Map<string, ((...args: any[]) => any)[]>();
+const commands = new Map<string, { handler: (...args: any[]) => any }>();
+
+const mockApi: ExtensionAPI = {
+  registerProvider: () => {},
+  on: (event: string, handler: (...args: any[]) => any) => {
+    if (!handlers.has(event)) handlers.set(event, []);
+    handlers.get(event)!.push(handler);
+  },
+  registerCommand: (name: string, opts: any) => {
+    commands.set(name, opts);
+  },
+  appendEntry: () => {},
+  exec: async () => ({ exitCode: 0, stdout: "", stderr: "" }),
+} as any;
+
+registerLilac(mockApi);
+
+const inputHandlers = handlers.get("input") ?? [];
+assert(inputHandlers.length === 1, "exactly one input handler registered");
+
+function runInput(source: string, model: any) {
+  const notifications: { msg: string; level: string }[] = [];
+  const ctx = {
+    model,
+    ui: {
+      notify: (msg: string, level: string) => notifications.push({ msg, level }),
+      setStatus: () => {},
+      theme: { fg: (_c: string, t: string) => t },
+    },
+  };
+  return Promise.all(inputHandlers.map((h) => h({ source }, ctx))).then((results) => ({
+    result: results[0],
+    notifications,
+  }));
+}
+
+// Reset flex to a known state for the gate suite.
+updateConfig((c) => ({ ...c, flexThreshold: null }));
+
+// flex off -> always allow
+{
+  const { result, notifications } = await runInput("interactive", { id: KIMI, provider: "lilac" });
+  eq(result, { action: "continue" }, "flex off -> continue");
+  assert(notifications.length === 0, "flex off -> no notify");
+}
+
+// flex 75, kimi at 25% -> blocked
+updateConfig((c) => ({ ...c, flexThreshold: 75 }));
+{
+  const { result, notifications } = await runInput("interactive", { id: KIMI, provider: "lilac" });
+  eq(result, { action: "handled" }, "25% < 75% threshold -> handled (blocked)");
+  assert(notifications.length === 1, "blocked -> one notify");
+  assert(notifications[0].level === "warning", "blocked notify is a warning");
+  assert(notifications[0].msg.includes("25%") && notifications[0].msg.includes("75%"), "blocked notify mentions current and threshold");
+}
+
+// flex 20, kimi at 25% -> allowed (>= threshold)
+updateConfig((c) => ({ ...c, flexThreshold: 20 }));
+{
+  const { result, notifications } = await runInput("interactive", { id: KIMI, provider: "lilac" });
+  eq(result, { action: "continue" }, "25% >= 20% threshold -> continue");
+  assert(notifications.length === 0, "allowed -> no notify");
+}
+
+// flex 75, non-lilac model -> allowed (gate only applies to lilac)
+updateConfig((c) => ({ ...c, flexThreshold: 75 }));
+{
+  const { result, notifications } = await runInput("interactive", { id: "anthropic/claude", provider: "anthropic" });
+  eq(result, { action: "continue" }, "non-lilac model -> continue");
+  assert(notifications.length === 0, "non-lilac model -> no notify");
+}
+
+// flex 75, rpc source -> allowed (only interactive is gated)
+{
+  const { result } = await runInput("rpc", { id: KIMI, provider: "lilac" });
+  eq(result, { action: "continue" }, "rpc source -> continue (not gated)");
+}
+
+// flex 75, extension source -> allowed (no loop)
+{
+  const { result } = await runInput("extension", { id: KIMI, provider: "lilac" });
+  eq(result, { action: "continue" }, "extension source -> continue (no loop)");
+}
+
+// flex 75, lilac model with no discount entry -> 0% -> blocked, mentions list price
+{
+  const { result, notifications } = await runInput("interactive", { id: "zai-org/glm-5.1", provider: "lilac" });
+  eq(result, { action: "handled" }, "uncached lilac model -> handled (0% < 75%)");
+  assert(notifications.length === 1, "uncached model blocked -> one notify");
+  assert(notifications[0].msg.includes("no discount on this model"), "uncached model notify mentions list price");
+}
+
+// flex exactly equal to discount -> allowed (>= is inclusive)
+updateConfig((c) => ({ ...c, flexThreshold: 25 }));
+{
+  const { result } = await runInput("interactive", { id: KIMI, provider: "lilac" });
+  eq(result, { action: "continue" }, "25% >= 25% threshold (inclusive) -> continue");
+}
+
+// ─── /lilac-flex command ──────────────────────────────────────────────────────
+
+console.log("\n--- /lilac-flex command ---");
+
+assert(commands.has("lilac-flex"), "/lilac-flex command registered");
+const cmd = commands.get("lilac-flex")!;
+
+function runCmd(args: string, opts: { select?: string; input?: string; hasUI?: boolean; model?: any } = {}) {
+  const notifications: { msg: string; level: string }[] = [];
+  const ctx = {
+    hasUI: opts.hasUI ?? true,
+    model: opts.model ?? { id: KIMI, provider: "lilac" },
+    ui: {
+      notify: (msg: string, level: string) => notifications.push({ msg, level }),
+      setStatus: () => {},
+      theme: { fg: (_c: string, t: string) => t },
+      select: async (_title: string, _labels: string[]) => opts.select,
+      input: async (_title: string, _placeholder: string) => opts.input,
+    },
+  };
+  return cmd.handler(args, ctx).then(() => ({ notifications, flexThreshold: getConfig().flexThreshold ?? null }));
+}
+
+// direct numeric arg
+{
+  const { flexThreshold } = await runCmd("75");
+  assert(flexThreshold === 75, "/lilac-flex 75 -> flexThreshold 75");
+  assert(JSON.parse(fs.readFileSync(cfgPath, "utf8")).flexThreshold === 75, "/lilac-flex 75 persisted to disk");
+}
+
+// trailing % accepted
+{
+  const { flexThreshold } = await runCmd("50%");
+  assert(flexThreshold === 50, "/lilac-flex 50% -> flexThreshold 50");
+}
+
+// off keyword
+{
+  const { flexThreshold } = await runCmd("off");
+  assert(flexThreshold === null, "/lilac-flex off -> flexThreshold null");
+}
+
+// invalid arg -> no change, usage notify
+{
+  const { flexThreshold, notifications } = await runCmd("bogus");
+  assert(flexThreshold === null, "/lilac-flex bogus -> no change (still off)");
+  assert(notifications.some((n) => n.msg.includes("Usage")), "/lilac-flex bogus -> usage notify");
+}
+
+// picker: "≥ 75% discount" preset
+{
+  const { flexThreshold } = await runCmd("", { select: "≥ 75% discount" });
+  assert(flexThreshold === 75, "picker '≥ 75% discount' -> 75");
+}
+
+// picker: "Off" preset
+{
+  const { flexThreshold } = await runCmd("", { select: "Off — allow all discounts" });
+  assert(flexThreshold === null, "picker 'Off' -> null");
+}
+
+// picker: "Custom…" then input "60"
+{
+  const { flexThreshold } = await runCmd("", { select: "Custom…", input: "60" });
+  assert(flexThreshold === 60, "picker Custom + input 60 -> 60");
+}
+
+// picker: custom input invalid -> no change, warning
+{
+  const { flexThreshold, notifications } = await runCmd("", { select: "Custom…", input: "abc" });
+  assert(flexThreshold === 60, "invalid custom input -> no change (still 60)");
+  assert(notifications.some((n) => n.level === "warning" && n.msg.includes("Invalid")), "invalid custom input -> warning notify");
+}
+
+// picker cancelled (select returns undefined) -> no change, no notify
+{
+  const { flexThreshold, notifications } = await runCmd("", { select: undefined });
+  assert(flexThreshold === 60, "picker cancelled -> no change (still 60)");
+  assert(notifications.length === 0, "picker cancelled -> no notify");
+}
+
+// non-interactive guard -> error notify, no change
+{
+  const { flexThreshold, notifications } = await runCmd("75", { hasUI: false });
+  assert(flexThreshold === 60, "non-interactive -> no change (still 60)");
+  assert(notifications.some((n) => n.level === "error" && n.msg.includes("interactive")), "non-interactive -> error notify");
+}
+
+// setting flex via command is visible to the gate (config cache in sync)
+{
+  await runCmd("75");
+  const { result } = await runInput("interactive", { id: KIMI, provider: "lilac" });
+  eq(result, { action: "handled" }, "after /lilac-flex 75, gate blocks kimi@25%");
+  await runCmd("off");
+  const { result: afterOff } = await runInput("interactive", { id: KIMI, provider: "lilac" });
+  eq(afterOff, { action: "continue" }, "after /lilac-flex off, gate allows kimi@25%");
+}
+
+// ─── Summary ───────────────────────────────────────────────────────────────────
+
+console.log(`\n${failed === 0 ? "ALL PASS" : `${failed} FAILED`}`);
+process.exit(failed === 0 ? 0 : 1);

package/scripts/test-model-overrides.ts

@@ -166,9 +166,9 @@ const cfgPath = path.join(os.homedir(), ".pi", "agent", "extensions", "lilac.jso
   // Fresh tmpHome: no config file -> loadConfig auto-populates the scaffold and returns defaults
   assert(!fs.existsSync(cfgPath), "scaffold not present before first loadConfig");
   const cfg = loadConfig();
-  eq(cfg, { modelOverrides: {} }, "missing file -> defaults (empty modelOverrides)");
+  eq(cfg, { modelOverrides: {}, flexThreshold: null }, "missing file -> defaults (empty modelOverrides, flex off)");
   assert(fs.existsSync(cfgPath), "loadConfig auto-populates the scaffold file on missing file");
-  eq(JSON.parse(fs.readFileSync(cfgPath, "utf8")), { modelOverrides: {} }, "scaffold file contains the default shape");
+  eq(JSON.parse(fs.readFileSync(cfgPath, "utf8")), { modelOverrides: {}, flexThreshold: null }, "scaffold file contains the default shape");
 }
 
 {
@@ -192,7 +192,7 @@ const cfgPath = path.join(os.homedir(), ".pi", "agent", "extensions", "lilac.jso
   // Existing file with invalid JSON -> defaults returned, file left UNTOUCHED (typo not wiped)
   fs.writeFileSync(cfgPath, "not json {{{");
   const cfg = loadConfig();
-  eq(cfg, { modelOverrides: {} }, "invalid JSON -> defaults");
+  eq(cfg, { modelOverrides: {}, flexThreshold: null }, "invalid JSON -> defaults");
   assert(fs.readFileSync(cfgPath, "utf8") === "not json {{{", "invalid file is not overwritten (typo preserved)");
 }