Fix paid usage settlement across resets, retries and late pricing

docs/managed-usage-settlement.md

@@ -0,0 +1,67 @@
+# Managed usage settlement
+
+This follow-up accounts for work already admitted by managed inference. It does
+not introduce budgets, reservations, free-model allowances, price changes or a
+reconciliation scheduler. Response handling remains separate in #4471.
+
+## Durable identity and time
+
+Admission samples one timestamp while holding the organization policy lock and
+uses it to choose the allowance buckets. The proxy forwards that same timestamp
+in its server-owned trace, including when provider-key lookup crosses a reset.
+
+Every provider receipt resolves one existing `openrouter_usage` identity before
+pricing or settlement. Organization, membership, inference key and request must
+match. The first durable timestamp and priced cost remain authoritative on replay;
+changed delivery IDs, later timestamps and different reported costs cannot charge
+the request again or move it to a new window.
+
+Nullable provider facts distinguish `unpriced` from `priced`. Absent, null or blank
+prices are unknown, never fabricated zero costs. A complete later receipt promotes
+the same unpriced row under its lock; a priced row never becomes unpriced again.
+`priced` means monetary facts are known, not that every historical charge exists.
+
+## Settlement and reset
+
+Settlement uses the three original historical windows, even when exhausted. It
+never advances current window pointers, rewrites limits, provisions credentials or
+grants generation access. Revoked keys may settle attributable earlier work using
+the identity-matched durable timestamp. New requests still obey current DPA and
+access policy.
+
+Admission, Den's bucket writer, settlement and admin reset take policy locks in the
+same order. Per-bucket charge identities make concurrent retries idempotent and
+permit repair of a genuinely missing charge. Admin reset subtracts known usage but
+keeps zero-amount charge identities, so a receipt replay cannot undo forgiveness.
+Reset retains its existing scope: known charges in the current windows, not blanket
+amnesty for usage whose price has not arrived.
+
+Old reset implementations deleted charge identities. For legacy ledger entries
+without provider facts, a missing charge cannot safely be distinguished from an
+intentional reset. Those gaps are durably deferred rather than silently backfilled.
+Missing or overlapping historical windows are also deferred, never guessed.
+
+Webhook responses distinguish `ingested`, `deferred`, `skipped`, `invalid` and
+`failed`. Durably retained unpriced/ambiguous usage is acknowledged as deferred;
+explicitly invalid attribution is not accepted as another identity's charge.
+Malformed usage receives 400; transient persistence failures receive retryable 503.
+Resource-, scope- and span-level OTLP attributes use the same validation path.
+
+## Migration and proof
+
+Apply `0094_managed_usage_facts.sql` before deploying readers/writers. It adds one
+nullable JSON column; existing entries remain null. The generated snapshot is
+schema metadata, not a new ledger. Reconcile numbering against current dev before
+merge. The existing conversion factor, minimum-unit rounding and allowances remain
+unchanged.
+
+`pnpm evals:pr specs/paid-usage-settlement.test.ts` uses real Den/admin-reset and
+inference HTTP routes with disposable SQL and a controlled provider. It covers the
+real migration runner and no-op rerun, concurrent receipts and bucket writers,
+reset/replay races, unpriced promotion, ambiguous history, admission-time rollover,
+revocation/DPA boundaries and unchanged neighboring usage.
+
+Live provider broadcasts and production upgrades are not certified by the fixture.
+In-flight overspend and missing receipts remain separate concerns. This patch has
+shared integration surfaces with #4358, #3648 and #4620; it imports none of their
+unmerged gateway, status-rollover or free-allowance implementations.

ee/apps/den-api/src/inference.ts

@@ -1,4 +1,4 @@
-import { and, eq, inArray, isNull, sql } from "@openwork-ee/den-db/drizzle"
+import { and, asc, eq, inArray, isNull, sql } from "@openwork-ee/den-db/drizzle"
 import {
   InferenceKeyTable,
   InferenceOrgLimitPolicyTable,
@@ -354,72 +354,77 @@ export async function syncInferenceAfterMemberChange(input: {
 }
 
 async function syncInferenceLimitPolicies(input: { organizationId: OrgId; tier: InferenceTier; memberCount: number }) {
-  const now = new Date()
-  for (const windowType of Object.keys(INFERENCE_TIER_LIMITS[input.tier])) {
-    await db
-      .insert(InferenceOrgLimitPolicyTable)
-      .values({
-        id: createDenTypeId("inferenceOrgLimitPolicy"),
-        organization_id: input.organizationId,
-        window_type: windowType as keyof typeof INFERENCE_TIER_LIMITS[InferenceTier],
-        reset_strategy: INFERENCE_RESET_STRATEGY_BY_WINDOW_TYPE[windowType as keyof typeof INFERENCE_TIER_LIMITS[InferenceTier]],
-        anchor_at: now,
-      })
-      .onDuplicateKeyUpdate({
-        set: {
+  await db.transaction(async (tx) => {
+    const anchorAt = new Date()
+    for (const windowType of Object.keys(INFERENCE_TIER_LIMITS[input.tier])) {
+      await tx
+        .insert(InferenceOrgLimitPolicyTable)
+        .values({
+          id: createDenTypeId("inferenceOrgLimitPolicy"),
+          organization_id: input.organizationId,
+          window_type: windowType as keyof typeof INFERENCE_TIER_LIMITS[InferenceTier],
           reset_strategy: INFERENCE_RESET_STRATEGY_BY_WINDOW_TYPE[windowType as keyof typeof INFERENCE_TIER_LIMITS[InferenceTier]],
-        },
-      })
-  }
-
-  const policies = await db
-    .select({
-      id: InferenceOrgLimitPolicyTable.id,
-      windowType: InferenceOrgLimitPolicyTable.window_type,
-      resetStrategy: InferenceOrgLimitPolicyTable.reset_strategy,
-      anchorAt: InferenceOrgLimitPolicyTable.anchor_at,
-      currentBucketId: InferenceOrgLimitPolicyTable.current_bucket_id,
-    })
-    .from(InferenceOrgLimitPolicyTable)
-    .where(eq(InferenceOrgLimitPolicyTable.organization_id, input.organizationId))
+          anchor_at: anchorAt,
+        })
+        .onDuplicateKeyUpdate({
+          set: {
+            reset_strategy: INFERENCE_RESET_STRATEGY_BY_WINDOW_TYPE[windowType as keyof typeof INFERENCE_TIER_LIMITS[InferenceTier]],
+          },
+        })
+    }
 
-  for (const policy of policies) {
-    const limitAmount = INFERENCE_TIER_LIMITS[input.tier][policy.windowType] * input.memberCount
-    const currentBucket = policy.currentBucketId
-      ? (await db.select().from(InferenceOrgUsageBucketTable).where(eq(InferenceOrgUsageBucketTable.id, policy.currentBucketId)).limit(1))[0]
-      : null
+    const policies = await tx
+      .select({
+        id: InferenceOrgLimitPolicyTable.id,
+        windowType: InferenceOrgLimitPolicyTable.window_type,
+        resetStrategy: InferenceOrgLimitPolicyTable.reset_strategy,
+        anchorAt: InferenceOrgLimitPolicyTable.anchor_at,
+        currentBucketId: InferenceOrgLimitPolicyTable.current_bucket_id,
+      })
+      .from(InferenceOrgLimitPolicyTable)
+      .where(eq(InferenceOrgLimitPolicyTable.organization_id, input.organizationId))
+      .orderBy(asc(InferenceOrgLimitPolicyTable.window_type))
+      .for("update")
+    const now = new Date()
+
+    for (const policy of policies) {
+      const limitAmount = INFERENCE_TIER_LIMITS[input.tier][policy.windowType] * input.memberCount
+      const currentBucket = policy.currentBucketId
+        ? (await tx.select().from(InferenceOrgUsageBucketTable).where(eq(InferenceOrgUsageBucketTable.id, policy.currentBucketId)).limit(1).for("update"))[0]
+        : null
+
+      if (currentBucket && currentBucket.window_start_at <= now && currentBucket.window_end_at > now) {
+        await tx
+          .update(InferenceOrgUsageBucketTable)
+          .set({ limit_amount: limitAmount })
+          .where(eq(InferenceOrgUsageBucketTable.id, currentBucket.id))
+        continue
+      }
 
-    if (currentBucket && currentBucket.window_start_at <= now && currentBucket.window_end_at > now) {
-      await db
-        .update(InferenceOrgUsageBucketTable)
-        .set({ limit_amount: limitAmount })
-        .where(eq(InferenceOrgUsageBucketTable.id, currentBucket.id))
-      continue
+      const window = policy.resetStrategy === "anchored"
+        ? currentWindow({
+            anchorAt: policy.anchorAt,
+            currentEnd: currentBucket?.window_end_at ?? null,
+            windowType: policy.windowType,
+            now,
+          })
+        : { start: now, end: addWindow(now, policy.windowType) }
+      const bucketId = createDenTypeId("inferenceOrgUsageBucket")
+      await tx.insert(InferenceOrgUsageBucketTable).values({
+        id: bucketId,
+        organization_id: input.organizationId,
+        policy_id: policy.id,
+        window_start_at: window.start,
+        window_end_at: window.end,
+        limit_amount: limitAmount,
+        used_amount: 0,
+      })
+      await tx
+        .update(InferenceOrgLimitPolicyTable)
+        .set({ current_bucket_id: bucketId })
+        .where(eq(InferenceOrgLimitPolicyTable.id, policy.id))
     }
-
-    const window = policy.resetStrategy === "anchored"
-      ? currentWindow({
-          anchorAt: policy.anchorAt,
-          currentEnd: currentBucket?.window_end_at ?? null,
-          windowType: policy.windowType,
-          now,
-        })
-      : { start: now, end: addWindow(now, policy.windowType) }
-    const bucketId = createDenTypeId("inferenceOrgUsageBucket")
-    await db.insert(InferenceOrgUsageBucketTable).values({
-      id: bucketId,
-      organization_id: input.organizationId,
-      policy_id: policy.id,
-      window_start_at: window.start,
-      window_end_at: window.end,
-      limit_amount: limitAmount,
-      used_amount: 0,
-    })
-    await db
-      .update(InferenceOrgLimitPolicyTable)
-      .set({ current_bucket_id: bucketId })
-      .where(eq(InferenceOrgLimitPolicyTable.id, policy.id))
-  }
+  })
 }
 
 type OpenRouterKeyCreateResponse = {

ee/apps/den-api/src/routes/admin/index.ts

@@ -1428,6 +1428,13 @@ export function registerAdminRoutes<T extends { Variables: AuthContextVariables
       }
 
       const resetAmount = await db.transaction(async (tx) => {
+        const memberships = await tx.select({ organizationId: MemberTable.organizationId }).from(MemberTable)
+          .where(and(eq(MemberTable.userId, userId), isNull(MemberTable.removedAt)))
+        if (memberships.length === 0) return 0
+        // Same order as admission/settlement: policies before ledger, buckets and charges.
+        await tx.select({ id: InferenceOrgLimitPolicyTable.id }).from(InferenceOrgLimitPolicyTable)
+          .where(inArray(InferenceOrgLimitPolicyTable.organization_id, memberships.map((member) => member.organizationId)))
+          .orderBy(asc(InferenceOrgLimitPolicyTable.organization_id), asc(InferenceOrgLimitPolicyTable.window_type)).for("update")
         const charges = await tx
           .select({
             id: InferenceUsageLedgerBucketChargeTable.id,
@@ -1461,7 +1468,8 @@ export function registerAdminRoutes<T extends { Variables: AuthContextVariables
             .where(eq(InferenceOrgUsageBucketTable.id, bucketId))
         }
         if (charges.length > 0) {
-          await tx.delete(InferenceUsageLedgerBucketChargeTable).where(inArray(
+          // Keep the identity so retries cannot resurrect intentionally forgiven usage.
+          await tx.update(InferenceUsageLedgerBucketChargeTable).set({ amount: 0 }).where(inArray(
             InferenceUsageLedgerBucketChargeTable.id,
             charges.map((charge) => charge.id),
           ))

ee/apps/inference/src/limits.ts

@@ -1,4 +1,4 @@
-import { and, eq, isNull, sql } from "@openwork-ee/den-db/drizzle"
+import { and, asc, eq, isNull, sql } from "@openwork-ee/den-db/drizzle"
 import { InferenceOrgLimitPolicyTable, InferenceOrgUsageBucketTable, MemberTable, OrganizationTable } from "@openwork-ee/den-db"
 import { createDenTypeId, normalizeDenTypeId, type DenTypeId } from "@openwork-ee/utils/typeid"
 import { INFERENCE_TIER_LIMITS, INFERENCE_WINDOW_DURATIONS_MS } from "@openwork/types/den/inference"
@@ -7,6 +7,7 @@ import { db } from "./db.js"
 
 export type BucketMetadata = Partial<Record<string, DenTypeId<"inferenceOrgUsageBucket">>>
 export type BucketLimitMetadata = Partial<Record<string, number>>
+type Transaction = Parameters<Parameters<typeof db.transaction>[0]>[0]
 
 function addWindow(start: Date, windowType: InferenceWindowType) {
   return new Date(start.getTime() + INFERENCE_WINDOW_DURATIONS_MS[windowType])
@@ -32,16 +33,16 @@ function readInferenceTier(metadata: Record<string, unknown> | null): InferenceT
   return tier === "tier1" || tier === "tier2" ? tier : null
 }
 
-async function getEffectiveLimits(organizationId: DenTypeId<"organization">) {
-  const [organization] = await db
+async function getEffectiveLimits(tx: Transaction, organizationId: DenTypeId<"organization">) {
+  const [organization] = await tx
     .select({ metadata: OrganizationTable.metadata })
     .from(OrganizationTable)
     .where(eq(OrganizationTable.id, organizationId))
     .limit(1)
   const tier = readInferenceTier(organization?.metadata ?? null)
   if (!tier) return null
 
-  const [memberCountRow] = await db
+  const [memberCountRow] = await tx
     .select({ count: sql<number>`count(*)` })
     .from(MemberTable)
     .where(and(eq(MemberTable.organizationId, organizationId), isNull(MemberTable.removedAt)))
@@ -52,13 +53,13 @@ async function getEffectiveLimits(organizationId: DenTypeId<"organization">) {
   ) as Record<InferenceWindowType, number>
 }
 
-async function ensureBucket(policy: typeof InferenceOrgLimitPolicyTable.$inferSelect, now: Date, effectiveLimit: number) {
+async function ensureBucket(tx: Transaction, policy: typeof InferenceOrgLimitPolicyTable.$inferSelect, now: Date, effectiveLimit: number) {
   const current = policy.current_bucket_id
-    ? (await db.select().from(InferenceOrgUsageBucketTable).where(eq(InferenceOrgUsageBucketTable.id, policy.current_bucket_id)).limit(1))[0]
+    ? (await tx.select().from(InferenceOrgUsageBucketTable).where(eq(InferenceOrgUsageBucketTable.id, policy.current_bucket_id)).limit(1).for("update"))[0]
     : null
   if (current && current.window_start_at <= now && current.window_end_at > now) {
     if (current.limit_amount !== effectiveLimit) {
-      await db.update(InferenceOrgUsageBucketTable).set({ limit_amount: effectiveLimit }).where(eq(InferenceOrgUsageBucketTable.id, current.id))
+      await tx.update(InferenceOrgUsageBucketTable).set({ limit_amount: effectiveLimit }).where(eq(InferenceOrgUsageBucketTable.id, current.id))
       current.limit_amount = effectiveLimit
     }
     return current
@@ -74,7 +75,7 @@ async function ensureBucket(policy: typeof InferenceOrgLimitPolicyTable.$inferSe
     : { start: now, end: addWindow(now, policy.window_type) }
   const id = createDenTypeId("inferenceOrgUsageBucket")
 
-  await db.insert(InferenceOrgUsageBucketTable).values({
+  await tx.insert(InferenceOrgUsageBucketTable).values({
     id,
     organization_id: policy.organization_id,
     policy_id: policy.id,
@@ -83,46 +84,52 @@ async function ensureBucket(policy: typeof InferenceOrgLimitPolicyTable.$inferSe
     limit_amount: effectiveLimit,
     used_amount: 0,
   })
-  await db.update(InferenceOrgLimitPolicyTable).set({ current_bucket_id: id }).where(eq(InferenceOrgLimitPolicyTable.id, policy.id))
+  await tx.update(InferenceOrgLimitPolicyTable).set({ current_bucket_id: id }).where(eq(InferenceOrgLimitPolicyTable.id, policy.id))
 
-  return (await db.select().from(InferenceOrgUsageBucketTable).where(eq(InferenceOrgUsageBucketTable.id, id)).limit(1))[0]
+  return (await tx.select().from(InferenceOrgUsageBucketTable).where(eq(InferenceOrgUsageBucketTable.id, id)).limit(1))[0]
 }
 
-export async function ensureUsableBuckets(organizationId: string, now = new Date()) {
+export async function ensureUsableBuckets(organizationId: string) {
   const orgId = normalizeDenTypeId("organization", organizationId)
-  const effectiveLimits = await getEffectiveLimits(orgId)
-  if (!effectiveLimits) {
-    return { ok: false as const, bucketIds: {}, bucketLimits: {}, limitedBy: "inference_metadata", windowType: "monthly" as const }
-  }
+  return db.transaction(async (tx) => {
+    const policies = await tx.select().from(InferenceOrgLimitPolicyTable)
+      .where(eq(InferenceOrgLimitPolicyTable.organization_id, orgId))
+      .orderBy(asc(InferenceOrgLimitPolicyTable.window_type)).for("update")
+    // This timestamp owns attribution, including time spent waiting for provider credentials.
+    const admittedAt = new Date()
+    const effectiveLimits = await getEffectiveLimits(tx, orgId)
+    if (!effectiveLimits) {
+      return { ok: false as const, bucketIds: {}, bucketLimits: {}, limitedBy: "inference_metadata", windowType: "monthly" as const }
+    }
 
-  const policies = await db.select().from(InferenceOrgLimitPolicyTable).where(eq(InferenceOrgLimitPolicyTable.organization_id, orgId))
-  const bucketIds: BucketMetadata = {}
-  const bucketLimits: BucketLimitMetadata = {}
+    const bucketIds: BucketMetadata = {}
+    const bucketLimits: BucketLimitMetadata = {}
 
-  for (const policy of policies) {
-    const effectiveLimit = effectiveLimits[policy.window_type]
-    const bucket = await ensureBucket(policy, now, effectiveLimit)
-    if (!bucket) {
-      continue
-    }
-    const remaining = effectiveLimit - bucket.used_amount
-    if (remaining <= 0) {
-      return {
-        ok: false as const,
-        bucketIds,
-        bucketLimits,
-        limitedBy: bucket.id,
-        windowType: policy.window_type,
-        limitedBucket: {
-          limitAmount: effectiveLimit,
-          usedAmount: bucket.used_amount,
-          windowEndAt: bucket.window_end_at,
-        },
+    for (const policy of policies) {
+      const effectiveLimit = effectiveLimits[policy.window_type]
+      const bucket = await ensureBucket(tx, policy, admittedAt, effectiveLimit)
+      if (!bucket) {
+        continue
       }
+      const remaining = effectiveLimit - bucket.used_amount
+      if (remaining <= 0) {
+        return {
+          ok: false as const,
+          bucketIds,
+          bucketLimits,
+          limitedBy: bucket.id,
+          windowType: policy.window_type,
+          limitedBucket: {
+            limitAmount: effectiveLimit,
+            usedAmount: bucket.used_amount,
+            windowEndAt: bucket.window_end_at,
+          },
+        }
+      }
+      bucketIds[policy.window_type] = bucket.id
+      bucketLimits[policy.window_type] = effectiveLimit
     }
-    bucketIds[policy.window_type] = bucket.id
-    bucketLimits[policy.window_type] = effectiveLimit
-  }
 
-  return { ok: true as const, bucketIds, bucketLimits }
+    return { ok: true as const, bucketIds, bucketLimits, admittedAt }
+  })
 }