agl 40.4.0 → 40.5.0

package/ts/classes.inferenceservice.ts

@@ -2,9 +2,10 @@ import * as plugins from './plugins.js';
 import { AuthorityState } from './classes.authoritystate.js';
 import type { IAuthoritySocketPaths } from './classes.authorityclient.js';
 import { FlexAccountAccess } from './classes.flexaccountaccess.js';
-import { InferenceRouter, InferenceRoutingError, type IInferenceAccount, type IInferenceAccountStatus,
-  type IInferenceRoutingPolicy } from './classes.inferencerouter.js';
+import { InferenceRouter, InferenceRoutingError, inferencePaymentReportOptions, type IInferenceAccount,
+  type IInferenceAccountStatus, type IInferenceRoutingPolicy } from './classes.inferencerouter.js';
 import { InferenceUsageCache } from './classes.inferenceusagecache.js';
+import { InferenceSpendOutbox, type IInferenceSpendStore } from './classes.inferencespend.js';
 import { InferenceUsageObserver, type IInferenceUsageObserverStatus } from './classes.inferenceusageobserver.js';
 import { InferenceGateway, type IInferenceGatewayClientOptions, type IInferenceGatewayGrant, type IInferenceGatewayLease,
   type IInferenceGatewayLimits, type IInferenceGatewayRestoredLease, type IInferenceGatewayStatus } from './classes.inferencegateway.js';
@@ -12,7 +13,8 @@ import type { IInferenceAffinityStore } from './classes.inferenceaffinity.js';
 import type { IInferencePlacementJournal, IInferencePlacementStatus } from './classes.inferenceplacement.js';
 import type { IInferenceUsageLedger } from './classes.inferenceusage.js';
 import type { TInferenceProvider } from './functions.proxycontract.js';
-import { findAuthorityGatewayLogin, findInferencePool, isAuthorityInstallationId, parseAuthorityScope } from '../ts_interfaces/index.js';
+import { findAuthorityApiKeyLogin, findAuthorityGatewayLogin, findInferencePool, isAuthorityInstallationId,
+  parseAuthorityScope } from '../ts_interfaces/index.js';
 
 export interface IInferenceServiceOptions {
   authorityPaths: IAuthoritySocketPaths;
@@ -25,6 +27,11 @@ export interface IInferenceServiceOptions {
   placements: IInferencePlacementJournal;
   /** Where the tokens of every answer are recorded for the usage metrics. */
   usageLedger: IInferenceUsageLedger;
+  /**
+   * Where the spend of the answers relayed on API keys is kept until the account service acknowledged it: the proxy's
+   * own database (`ProxyStore.spend`).
+   */
+  spendStore: IInferenceSpendStore;
   providerFetch?: typeof fetch;
   limits?: IInferenceGatewayLimits;
 }
@@ -45,6 +52,11 @@ export class InferenceService {
    * to the authority, whose answer the cache holds.
    */
   private readonly observer: InferenceUsageObserver;
+  /**
+   * The spend of the answers relayed on API keys until the account service acknowledged it: counted against each key's
+   * limit in the readings the cache hands out, and reported to the account service.
+   */
+  private readonly spend: InferenceSpendOutbox;
   private readonly abort = new AbortController();
   private ready = false;
   private started = false;
@@ -66,6 +78,21 @@ export class InferenceService {
         // Routing needs a reading: an account added a moment ago is routed on once its first read landed.
         return client.getUsage(accountIdArg, login.id, signalArg, { waitForFirstRead: true });
       },
+      // A key's reading counts the answers whose spend the account service has not acknowledged yet.
+      view: (accountIdArg, snapshotArg) => this.spend.judged(accountIdArg, snapshotArg),
+    });
+    this.spend = new InferenceSpendOutbox({
+      store: options.spendStore,
+      // Reported on the login each answer was served on, which the outbox keeps with it: also once the key was removed.
+      report: async (accountIdArg, loginIdArg, answersArg, signalArg) => {
+        const client = this.authority.readyClient();
+        if (!this.ready || !client) throw new Error('The account service is unavailable.');
+        return client.recordSpend(accountIdArg, loginIdArg, answersArg, signalArg);
+      },
+      // The account service's answer counts the report's answers: it replaces the held reading as they leave the outbox.
+      acknowledged: (accountIdArg, usageArg) => this.usage.accept(accountIdArg, usageArg),
+      // An answer recorded without a login is reported on the key's live login, or the removed key's.
+      attribute: (accountIdArg) => findAuthorityApiKeyLogin(this.authority.getProjection(), accountIdArg)?.id,
     });
     this.observer = new InferenceUsageObserver({
       send: async (accountIdArg, reportArg, signalArg) => {
@@ -89,15 +116,22 @@ export class InferenceService {
       },
       usage: (accountIdArg, signalArg) => this.usage.get(accountIdArg, signalArg),
       // The account service holds a payment problem on the account's usage, which the cache then routes on.
-      reportPaymentProblem: async (accountIdArg, observedAtArg) => {
+      reportPaymentProblem: async (accountIdArg, observedAtArg, causeArg) => {
         const client = this.authority.readyClient();
         const login = findAuthorityGatewayLogin(this.authority.getProjection(), accountIdArg);
         if (!this.ready || !client || typeof login === 'string') throw new Error('The account login is unavailable.');
-        const usage = await client.reportPaymentProblem(accountIdArg, login.id, observedAtArg, this.abort.signal);
+        const usage = await client.reportPaymentProblem(accountIdArg, login.id, observedAtArg, this.abort.signal,
+          inferencePaymentReportOptions(causeArg));
         this.usage.accept(accountIdArg, usage);
       },
       resolveAccess: (accountIdArg, signalArg, rejectedGenerationArg) => this.access.resolveGateway(
         accountIdArg, AbortSignal.any([signalArg, this.abort.signal]), rejectedGenerationArg),
+      // Recorded on the login the route was served on, as the pool was read when the request was placed.
+      spend: {
+        record: (accountIdArg, loginIdArg, recordArg) => this.spend.record(accountIdArg, loginIdArg, recordArg),
+        held: (accountIdArg) => this.spend.held(accountIdArg),
+        uncounted: (accountIdArg, modelArg) => this.spend.uncounted(accountIdArg, modelArg),
+      },
     });
     this.gateway = new InferenceGateway({ router: this.router, signatureKey: options.signatureKey,
       affinity: options.affinity, placements: options.placements, usageLedger: options.usageLedger, port: options.port,
@@ -114,8 +148,10 @@ export class InferenceService {
     if (!this.ready || projection.state !== 'ready' || !projection.snapshot) return undefined;
     return projection.snapshot.accounts.flatMap((account) => {
       const pool = findInferencePool(projection, account.id);
+      const login = findAuthorityGatewayLogin(projection, account.id);
       return pool ? [{ id: account.id, label: account.label, revision: String(account.revision), providerId: pool.provider,
-        credential: pool.credential, maxUsagePercent: account.maxUsagePercent, reserveLiftUntil: account.reserveLiftUntil, lapsesAt: account.lapsesAt }] : [];
+        credential: pool.credential, maxUsagePercent: account.maxUsagePercent, reserveLiftUntil: account.reserveLiftUntil, lapsesAt: account.lapsesAt,
+        spendLimitMicroUsd: account.spendLimitMicroUsd, loginId: typeof login === 'string' ? null : login.id }] : [];
     });
   }
 
@@ -143,6 +179,8 @@ export class InferenceService {
           incarnationId: binding.incarnationId, expectedRevision: binding.revision }, signal);
       }
       this.ready = true;
+      // The answers a run before left unacknowledged count from the start, before any request is placed, and are sent.
+      await this.spend.start(signal);
       return this.baseUrl;
     } catch (errorArg) {
       await this.stop();
@@ -165,20 +203,21 @@ export class InferenceService {
   /**
    * The gateway's state, each pool account judged on the usage reading a request would route on now -- the account
    * service's answer from its own cache (`InferenceUsageCache.current`), so a pool no request used since the proxy
-   * started is judged as its next request would judge it -- the subscription account each provider's next new client
-   * session goes to, and the latest placements the journal holds.
+   * started is judged as its next request would judge it -- the account each provider's next new session of a client
+   * that names no credential goes to (a subscription while one takes it, else an API key), and the latest placements the
+   * journal holds.
    */
   public async status(signalArg: AbortSignal): Promise<IInferenceGatewayStatus & { baseUrl: string;
     authority: 'ready' | 'unavailable'; accounts: IInferenceAccountStatus[]; usageObservations: IInferenceUsageObserverStatus;
     nextPlacement: Record<TInferenceProvider, string | null>; recentPlacements: IInferencePlacementStatus[] }> {
     const pool = this.poolAccounts() ?? [];
-    const readings = new Map(await Promise.all(pool.filter((accountArg) => accountArg.credential === 'subscription')
+    // A key's usage is its month's spend, with the payment problem the account service holds on it.
+    const readings = new Map(await Promise.all(pool
       .map(async (accountArg) => [accountArg.id, await this.usage.current(accountArg.id, signalArg)] as const)));
     const accounts = this.router.accountStatus(pool, (accountIdArg) => readings.get(accountIdArg));
     const recent = await this.options.placements.recent(recentPlacementCount, signalArg);
     const next = (providerArg: TInferenceProvider): string | null =>
-      accounts.find((accountArg) => accountArg.provider === providerArg && accountArg.credential === 'subscription'
-        && accountArg.placementRank === 1)?.id ?? null;
+      accounts.find((accountArg) => accountArg.provider === providerArg && accountArg.placementRank === 1)?.id ?? null;
     return { ...this.gateway.status(), baseUrl: this.baseUrl ?? '', usageObservations: this.observer.status(),
       authority: this.ready && this.authority.getProjection().state === 'ready' ? 'ready' : 'unavailable',
       accounts,
@@ -193,9 +232,13 @@ export class InferenceService {
       this.usage.stop();
       try { await this.gateway.stop(); }
       finally {
-        // Hand-overs in flight end before the authority connection closes under them.
+        // Hand-overs in flight end before the authority connection closes under them; the answers the gateway ended are
+        // in the outbox, which keeps what it holds for the next start.
         try { await this.observer.stop(); }
-        finally { try { await this.access.stop(); } finally { await this.authority.stop(); } }
+        finally {
+          try { await this.spend.stop(); }
+          finally { try { await this.access.stop(); } finally { await this.authority.stop(); } }
+        }
       }
     })();
   }