@opengeni/subscriptions 2.0.1 → 2.0.2-canary.38789554071001

package/src/quota.ts

@@ -1,4 +1,4 @@
-import type { ModelId, SubscriptionConnection, SubscriptionQuota } from "./types";
+import type { ModelId, QuotaBlockClaim, SubscriptionConnection, SubscriptionQuota } from "./types";
 
 /**
  * What the shared quota model says about a connection's capacity now
@@ -73,6 +73,89 @@ export function modelCooldownUntil(
   return until !== undefined && until > now ? until : null;
 }
 
+/** One stored block a usage reading may end: a model cooldown, or (model null) the exhaustion deadline. */
+export type SubscriptionQuotaBlock = {
+  modelId: ModelId | null;
+  until: number;
+  /** Its claim, only while it still explains this exact deadline. */
+  claim: QuotaBlockClaim | null;
+};
+
+/**
+ * The claim of a stored block (a model's cooldown, or with `modelId` null
+ * the exhaustion deadline), only while it is tied to the block's current
+ * deadline; null otherwise (fail closed).
+ */
+export function quotaBlockClaim(
+  quota: Pick<
+    SubscriptionQuota,
+    "modelCooldowns" | "modelCooldownClaims" | "exhaustedUntil" | "exhaustedClaim"
+  >,
+  modelId: ModelId | null,
+): QuotaBlockClaim | null {
+  const until = modelId === null ? quota.exhaustedUntil : quota.modelCooldowns[modelId];
+  const claim = modelId === null ? quota.exhaustedClaim : quota.modelCooldownClaims?.[modelId];
+  return until !== null &&
+    until !== undefined &&
+    claim !== null &&
+    claim !== undefined &&
+    claim.until === until &&
+    claim.windowIds.length > 0
+    ? claim
+    : null;
+}
+
+/** The quota document a store writes: claims only while they explain a block. */
+export type StoredQuotaDocument = Pick<
+  SubscriptionQuota,
+  | "windows"
+  | "modelCooldowns"
+  | "modelCooldownClaims"
+  | "exhaustedUntil"
+  | "exhaustedKind"
+  | "exhaustedClaim"
+  | "source"
+>;
+
+/** The stored document of one quota (a claim that no longer explains its block is dropped). */
+export function storedQuotaDocument(quota: SubscriptionQuota): StoredQuotaDocument {
+  const modelCooldownClaims: Record<string, QuotaBlockClaim> = {};
+  for (const modelId of Object.keys(quota.modelCooldowns)) {
+    const claim = quotaBlockClaim(quota, modelId);
+    if (claim) modelCooldownClaims[modelId] = claim;
+  }
+  const exhaustedClaim = quotaBlockClaim(quota, null);
+  return {
+    windows: quota.windows,
+    modelCooldowns: quota.modelCooldowns,
+    ...(Object.keys(modelCooldownClaims).length > 0 ? { modelCooldownClaims } : {}),
+    exhaustedUntil: quota.exhaustedUntil,
+    exhaustedKind: quota.exhaustedKind,
+    ...(exhaustedClaim ? { exhaustedClaim } : {}),
+    source: quota.source,
+  };
+}
+
+/**
+ * The claim of a block merged from several observations at `until`: every
+ * source whose own block has that exact deadline must carry a claim tied to
+ * it (their window ids are joined); otherwise none (an unexplained limit
+ * with the same deadline keeps the block).
+ */
+function mergedClaim(
+  until: number,
+  sources: readonly { until: number | null | undefined; claim: QuotaBlockClaim | null }[],
+): QuotaBlockClaim | null {
+  const deciding = sources.filter((source) => source.until === until);
+  if (deciding.length === 0 || deciding.some((source) => source.claim === null)) return null;
+  const claims = deciding.map((source) => source.claim!);
+  return {
+    until,
+    observedAt: Math.max(...claims.map((claim) => claim.observedAt)),
+    windowIds: [...new Set(claims.flatMap((claim) => claim.windowIds))],
+  };
+}
+
 /**
  * A quota observation applies only to the credential generation it was made
  * with, so a refusal seen before a refresh cannot quarantine the renewed
@@ -140,6 +223,31 @@ export function applyQuotaObservation(
     current.exhaustedUntil !== null &&
     current.exhaustedUntil > observedAt &&
     (observation.exhaustedUntil === null || observation.exhaustedUntil < current.exhaustedUntil);
+  const exhaustedUntil = keepCurrentDeadline ? current.exhaustedUntil : observation.exhaustedUntil;
+  // A block keeps its evidence only where every observation that set its
+  // merged deadline explained it (`QuotaBlockClaim`).
+  const sources = [current, observation];
+  const modelCooldownClaims: Record<string, QuotaBlockClaim> = {};
+  for (const [modelId, until] of Object.entries(modelCooldowns)) {
+    const claim = mergedClaim(
+      until,
+      sources.map((source) => ({
+        until: source.modelCooldowns[modelId],
+        claim: quotaBlockClaim(source, modelId),
+      })),
+    );
+    if (claim) modelCooldownClaims[modelId] = claim;
+  }
+  const exhaustedClaim =
+    exhaustedUntil === null
+      ? null
+      : mergedClaim(
+          exhaustedUntil,
+          sources.map((source) => ({
+            until: source.exhaustedUntil,
+            claim: quotaBlockClaim(source, null),
+          })),
+        );
   const reported = new Set(observation.windows.map((window) => window.id));
   const omittedExhausted =
     observation.source === "usage_endpoint" && options.usageReadingsComplete === true
@@ -151,13 +259,18 @@ export function applyQuotaObservation(
             window.resetsAt !== null &&
             window.resetsAt > observedAt,
         );
-  return {
+  const merged: SubscriptionQuota = {
     ...observation,
     windows: omittedExhausted.length
       ? [...observation.windows, ...omittedExhausted]
       : observation.windows,
     modelCooldowns,
-    exhaustedUntil: keepCurrentDeadline ? current.exhaustedUntil : observation.exhaustedUntil,
+    exhaustedUntil,
     exhaustedKind: keepCurrentDeadline ? current.exhaustedKind : observation.exhaustedKind,
   };
+  delete merged.modelCooldownClaims;
+  delete merged.exhaustedClaim;
+  if (Object.keys(modelCooldownClaims).length > 0) merged.modelCooldownClaims = modelCooldownClaims;
+  if (exhaustedClaim) merged.exhaustedClaim = exhaustedClaim;
+  return merged;
 }

package/src/reference-model.ts

@@ -143,6 +143,7 @@ export type SwitchKind =
 
 export type WaitReason =
   | "pinned_account_unavailable"
+  | "pinned_account_disabled"
   /** The explicit choice can never serve this work until a person changes something (D-24). */
   | "pinned_account_ineligible"
   | "no_eligible_capacity"
@@ -566,6 +567,21 @@ export function pinnedNeverServes(
   return !isAuthorized(world, session, { ...pinned, healthy: true, allocatorEnabled: true });
 }
 
+function pinnedAllocationDisabled(world: World, session: Session, pinned: Connection): boolean {
+  if (!pinned.allocatorEnabled) return true;
+  if (pinned.assignmentPolicies === undefined) return false;
+  const source =
+    effectiveSettings(world.settings, session.workspaceId).values.inferenceSource?.[
+      pinned.provider
+    ] ?? "automatic";
+  return !pinned.assignmentPolicies.some(
+    (policy) =>
+      policy.workspaceId === session.workspaceId &&
+      (source === "automatic" || policy.inferencePool === source) &&
+      policy.allocatorEnabled,
+  );
+}
+
 /** Future times at which some quota or model cooldown in the world resets, ascending. */
 function futureResetTimes(world: World, now: number): number[] {
   const times = new Set<number>();
@@ -652,6 +668,9 @@ export function decide(world: World, sessionId: string, now: number): Decision {
     if (pinnedNeverServes(world, session, pinned)) {
       return { kind: "wait", reason: "pinned_account_ineligible", earliestResetAt: null };
     }
+    if (pinned && pinnedAllocationDisabled(world, session, pinned)) {
+      return { kind: "wait", reason: "pinned_account_disabled", earliestResetAt: null };
+    }
     // Otherwise the earliest future time the chosen account can serve (D-22).
     return {
       kind: "wait",
@@ -906,7 +925,9 @@ export function checkDecision(
         : session.pinnedConnectionId
           ? never
             ? "pinned_account_ineligible"
-            : "pinned_account_unavailable"
+            : pinned && pinnedAllocationDisabled(world, session, pinned)
+              ? "pinned_account_disabled"
+              : "pinned_account_unavailable"
           : "no_eligible_capacity";
     if (decision.reason !== expectedReason) {
       fail(

package/src/types.ts

@@ -101,12 +101,34 @@ export type QuotaWindow = {
 
 export type QuotaSource = "usage_endpoint" | "response_headers" | "refusal";
 
+/**
+ * Which reported limits one stored block (a model cooldown, or the quota
+ * exhaustion deadline) came from: the provider's own window ids, as the
+ * observation that set the block named them, tied to that block's exact
+ * deadline. A block whose deadline no longer equals `until` (a later write
+ * extended it without evidence) has no claim. A block without a claim (an
+ * unrecognized or missing limit, or a refusal without usage evidence) never
+ * ends before its own deadline.
+ */
+export type QuotaBlockClaim = {
+  /** The deadline of the block this claim explains. */
+  until: number;
+  /** When the observation that set the block was made. */
+  observedAt: number;
+  /** The provider window ids whose limit set the block (never empty). */
+  windowIds: readonly string[];
+};
+
 export type SubscriptionQuota = {
   windows: readonly QuotaWindow[];
   /** Per-model cooldowns (model-specific limits some providers report). */
   modelCooldowns: Readonly<Record<ModelId, number>>;
+  /** Optional evidence per model cooldown (see `QuotaBlockClaim`); absent: none. */
+  modelCooldownClaims?: Readonly<Record<ModelId, QuotaBlockClaim>>;
   exhaustedUntil: number | null;
   exhaustedKind: "quota" | "rate_limit" | null;
+  /** Optional evidence for `exhaustedUntil` (see `QuotaBlockClaim`); absent: none. */
+  exhaustedClaim?: QuotaBlockClaim | null;
   revision: number;
   observedAt: number | null;
   observedRefreshGeneration: number | null;
@@ -315,6 +337,8 @@ export type AccountSwitchReason =
 
 export type WaitReason =
   | "pinned_account_unavailable"
+  /** Allocation was explicitly switched off; a person must enable it or change the pin. */
+  | "pinned_account_disabled"
   /** The explicit choice can never serve this work until a person changes something (D-24). */
   | "pinned_account_ineligible"
   | "no_eligible_capacity"

package/src/wait-presentation.ts

@@ -0,0 +1,27 @@
+import type { WaitReason } from "./types";
+
+/** Current wait guidance shared by initial parking and later reclassification.
+ * These are new audit events only; never rewrite an earlier model request. */
+export function subscriptionCapacityWaitReasonText(reason: string, name: string): string | null {
+  const messages: Record<WaitReason, string> = {
+    pinned_account_unavailable: `The ${name} account chosen for this session is unavailable right now. This turn continues automatically once it is available again.`,
+    pinned_account_disabled: `The ${name} account pinned to this chat is turned off. Turn it back on, choose another account, or switch the chat to automatic. Waiting alone will not fix this.`,
+    pinned_account_ineligible: `The ${name} account chosen for this session can no longer serve this work. This turn continues once another account is chosen or the session is switched back to automatic.`,
+    no_eligible_capacity: `No ${name} subscription has capacity for this turn right now. It continues automatically when an account is available.`,
+    model_not_allowed: `This model is not allowed for the ${name} subscriptions available here. This turn continues if that changes.`,
+    compaction_provider_locked:
+      "This session is tied to a provider that has no capacity right now. It continues automatically when an account is available.",
+    accepted_authority_unavailable:
+      "This work was accepted before the account move and its account access could not be carried. Send it again to continue.",
+  };
+  return Object.hasOwn(messages, reason) ? messages[reason as WaitReason] : null;
+}
+
+export function subscriptionCapacityWaitNeedsAction(reason: string): boolean {
+  return [
+    "pinned_account_disabled",
+    "pinned_account_ineligible",
+    "model_not_allowed",
+    "accepted_authority_unavailable",
+  ].includes(reason);
+}