@schlessera/brain-ui-server 0.19.0 → 0.20.0

package/src/routes/models.ts

@@ -8,27 +8,45 @@
 import type { Logger } from "@opentelemetry/api-logs";
 import { Hono } from "hono";
 import type { Database } from "bun:sqlite";
-import type {
-  ModelCatalogEntry,
-  ModelCatalogResponse,
+import {
+  isBillingMode,
+  type BillingMode,
+  type ModelCatalogEntry,
+  type ModelCatalogResponse,
 } from "@schlessera/brain-ui-sdk";
 import type { BackendRegistry } from "../agent/backend.js";
-import { getHiddenModelIds, setHiddenModelIds } from "../db/settings.js";
+import type { ModelPricingState } from "../pricing/model-pricing.js";
+import {
+  getBillingOverrides,
+  getHiddenModelIds,
+  setBillingOverrides,
+  setHiddenModelIds,
+} from "../db/settings.js";
 
 export function createModelRoutes(deps: {
   registry: BackendRegistry;
   db: Database;
+  /** The shared pricing instance — its state drives the client's staleness indicator. */
+  pricing?: { state(): ModelPricingState; ensureFresh(): Promise<void> };
   /** Where refresh failures are reported. */
   log?: Logger;
 }): Hono {
-  const { registry, db } = deps;
+  const { registry, db, pricing } = deps;
 
   async function buildCatalog(): Promise<ModelCatalogResponse> {
   const source = await registry.getModelSource();
   const hidden = new Set(getHiddenModelIds(db));
+  const overrides = getBillingOverrides(db);
+  // `billingMode` already rides each provider entry (the registry applies the
+  // override last); the catalog additionally tags WHICH rows carry an explicit
+  // override, so the settings screen can render auto vs forced.
   const models: ModelCatalogEntry[] = (
     await registry.listAllProviders({ includeHidden: true })
-  ).map((profile) => ({ ...profile, hidden: hidden.has(profile.id) }));
+  ).map((profile) => ({
+    ...profile,
+    hidden: hidden.has(profile.id),
+    ...(overrides[profile.id] ? { billingOverride: overrides[profile.id] } : {}),
+  }));
 
   const state = source?.state();
   return {
@@ -51,6 +69,17 @@ export function createModelRoutes(deps: {
     return c.json(await buildCatalog());
   })
 
+  .get("/models/pricing", async (c) => {
+    // Absent instance (a host that opted out) → 404; the client hides the
+    // indicator on any rejection, so old and opted-out servers degrade alike.
+    if (!pricing) return c.json({ error: "pricing not available" }, 404);
+    // Kick a refresh behind the response when stale — reading the state is
+    // also the natural moment to heal it. Never blocks: ensureFresh awaits
+    // only on a true cold start.
+    void pricing.ensureFresh().catch(() => {});
+    return c.json(pricing.state());
+  })
+
   .put("/models/hidden", async (c) => {
     const body = (await c.req.json().catch(() => null)) as unknown;
     const hidden = (body as { hidden?: unknown } | null)?.hidden;
@@ -67,6 +96,28 @@ export function createModelRoutes(deps: {
     return c.json(await buildCatalog());
   })
 
+  .put("/models/billing", async (c) => {
+    const body = (await c.req.json().catch(() => null)) as unknown;
+    const billing = (body as { billing?: unknown } | null)?.billing;
+    if (
+      typeof billing !== "object" ||
+      billing === null ||
+      Array.isArray(billing) ||
+      Object.values(billing).some((mode) => !isBillingMode(mode))
+    ) {
+      return c.json(
+        { error: 'billing must map profile ids to "subscription" or "api"' },
+        400
+      );
+    }
+
+    setBillingOverrides(db, billing as Record<string, BillingMode>);
+    // Same memo discipline as the hidden set: the next roster read (and the
+    // next run's classification) must see the override immediately.
+    registry.invalidateProfiles();
+    return c.json(await buildCatalog());
+  })
+
   .post("/models/refresh", async (c) => {
     const source = await registry.getModelSource();
     if (!source) {

package/src/ws/run-session.ts

@@ -1,4 +1,9 @@
-import type { ChatImageAttachment, ClientEnvironment } from "@schlessera/brain-ui-sdk/protocol";
+import type {
+  BillingMode,
+  ChatImageAttachment,
+  ClientEnvironment,
+} from "@schlessera/brain-ui-sdk/protocol";
+import type { BackendRegistry } from "../agent/backend.js";
 import type { WSContext } from "./clients.js";
 import { withSessionId, withTurnScope } from "./frames.js";
 import { makeBridge, emitTurnError } from "./bridge.js";
@@ -116,7 +121,14 @@ export async function runSession(
       turn.timeoutHandle = timeoutHandle;
 
       // One recorder per turn identity: a queued follow-up re-mints turnId
-      // and gets its own run in the activity record.
+      // and gets its own run in the activity record. The recorder is handed
+      // the RESOLVED profile — `profileId` here is post pin-drop fallback and
+      // carries any session_info re-pin forward between queued turns — plus
+      // its billing mode, so the run's root span classifies what actually ran,
+      // never what was requested.
+      const billing = host.activity
+        ? await resolveRunBilling(host.registry, backend.id, profileId)
+        : undefined;
       const recorder: TurnRecorder | undefined = host.activity
         ? createTurnRecorder(
             {
@@ -124,7 +136,7 @@ export async function runSession(
               onWrite: () => host.activity!.stream.pump(),
               log: host.log,
             },
-            { turnId: turn.turnId, sessionId: turn.sessionId }
+            { turnId: turn.turnId, sessionId: turn.sessionId, ...billing }
           )
         : undefined;
       const bridge = makeBridge(host, turn, text, backend.id, recorder);
@@ -194,6 +206,36 @@ export async function runSession(
   }
 }
 
+/**
+ * Profile identity + billing mode for a turn's root span. When the slot has no
+ * resolved profile id (a resumed session whose pin was dropped), the backend
+ * runs its default profile — the backend's first roster entry — so THAT
+ * profile's identity and billing are recorded, not the dead pin's. Hidden
+ * profiles are included: a session pinned to a profile the user later hid
+ * still runs on it and must classify as it. Never throws: an unreadable
+ * roster records no billing attrs, and the rollup falls back to env
+ * classification.
+ */
+async function resolveRunBilling(
+  registry: BackendRegistry,
+  backendId: string,
+  profileId: string | undefined
+): Promise<{ profileId: string; billingMode?: BillingMode } | undefined> {
+  try {
+    const providers = await registry.listAllProviders({ includeHidden: true });
+    const resolved = profileId
+      ? providers.find((provider) => provider.id === profileId)
+      : providers.find((provider) => provider.backendId === backendId);
+    if (!resolved) return undefined;
+    return {
+      profileId: resolved.id,
+      ...(resolved.billingMode ? { billingMode: resolved.billingMode } : {}),
+    };
+  } catch {
+    return undefined;
+  }
+}
+
 /** Queue sizes are only ever reported to a human, so one decimal of MB is plenty. */
 function formatMb(bytes: number): string {
   return `${(bytes / 1024 / 1024).toFixed(1)} MB`;