@sixb/core 0.1.8 → 0.1.9

package/src/models/execution/index.ts

@@ -0,0 +1,26 @@
+export { resolveLanguageModel } from "../resolve"
+export { ModelUsageRecordingError } from "./errors"
+export type {
+  AiModelCallAdmissionDecision,
+  AiModelCallAdmissionInput,
+  AiModelCallInputTokenEstimate,
+  BeforeAiModelCall,
+  MarkAiModelCallUnknown,
+} from "./model-call-admission"
+export type { AiModelCallLimitController } from "./model-call-limits"
+export { createAiModelCallLimitController } from "./model-call-limits"
+export type { AiModelCallRecorderInput } from "./model-call-recorder"
+export { AiModelCallRecorder } from "./model-call-recorder"
+export {
+  isPermanentAiUsageRecoveryError,
+  recordRecoveredAiModelCall,
+} from "./model-call-recovery"
+export { enqueueAiModelCallRecovery } from "./recovery-queue"
+export type {
+  AiModelCallAccountingPayload,
+  AiModelCallRecordPayload,
+  AiModelCallRecoveryRecord,
+  ModelCallAccountingStorage,
+  RecoverAiModelCall,
+  RecoverAiModelCallInput,
+} from "./types"

package/src/models/execution/model-call-accounting.ts

@@ -0,0 +1,211 @@
+import { createHash } from "node:crypto"
+import { assertJsonObject } from "../../json"
+import type {
+  AiBillableMeter,
+  AiCostStorage,
+  AiLimitStorage,
+  AiModelCallCostRecord,
+  AiModelCallUsageRecord,
+  AiPricingContext,
+  AiUsageStorage,
+  RecordAiModelCallResult,
+  Storage,
+} from "../../storage"
+import {
+  aiModelCallCostMatchesUsage,
+  normalizeAiModelCallCostRecord,
+} from "../../storage/ai-cost/provider"
+import { normalizeAiModelCallRecord } from "../../storage/ai-usage/record"
+import type { ModelCostEstimate } from "../pricing"
+import type { ModelCallAccountingStorage, RecoverAiModelCallInput } from "./types"
+
+interface RecordAiModelCallAccountingInput extends RecoverAiModelCallInput {
+  readonly storage: Pick<ModelCallAccountingStorage, "transaction">
+}
+
+interface AiAccountingCapabilities {
+  readonly aiUsage: AiUsageStorage
+  readonly aiCosts: AiCostStorage
+  readonly aiLimits: AiLimitStorage
+}
+
+/** Atomically append one provider call's usage and the valuation captured by the model runtime. */
+export async function recordAiModelCallAccounting(
+  input: RecordAiModelCallAccountingInput
+): Promise<RecordAiModelCallResult> {
+  return input.storage.transaction(async (tx) => {
+    const { aiUsage, aiCosts, aiLimits } = requireAccountingCapabilities(tx)
+    const usage = await aiUsage.recordModelCall(input.usage)
+    // Usage deduplicates provider lifecycle replays by execution/call identity. Always attach the
+    // valuation to the canonical row it returns, not the fresh candidate ID from a replay.
+    const canonical = normalizeModelCallAccounting({
+      ...input,
+      usage: usage.record,
+      ratedAt: usage.record.recordedAt,
+    })
+    const record = modelCostRecord(canonical, usage.record.id)
+    const estimate =
+      canonical.cost.status !== "reported" || canonical.estimate === undefined
+        ? undefined
+        : modelCostRecord({ ...canonical, cost: canonical.estimate }, usage.record.id)
+    await aiCosts.recordModelCallCost({
+      ...record,
+      ...(estimate === undefined
+        ? {}
+        : {
+            estimate:
+              estimate.status === "rated"
+                ? { status: "rated", money: estimate.money, components: estimate.components }
+                : {
+                    status: "unpriceable",
+                    reason: estimate.reason,
+                    ...(estimate.missingMeters === undefined
+                      ? {}
+                      : { missingMeters: estimate.missingMeters }),
+                  },
+          }),
+    })
+    if (usage.created) {
+      await aiLimits.recordModelCallActuals({
+        projectId: usage.record.projectId,
+        usageRecordId: usage.record.id,
+        recordedAt: input.ratedAt,
+      })
+    }
+    if (input.reconcileLimitReservation) {
+      await aiLimits.reconcileModelCall({
+        projectId: usage.record.projectId,
+        executionId: usage.record.executionId,
+        attempt: usage.record.attempt,
+        callId: usage.record.callId,
+        usageRecordId: usage.record.id,
+        reconciledAt: input.ratedAt,
+      })
+    }
+    return usage
+  })
+}
+
+/** Sanitize optional enrichment before either persistence or a JSON-safe recovery handoff. */
+export function normalizeModelCallAccounting(
+  input: RecoverAiModelCallInput
+): RecoverAiModelCallInput {
+  const usage = normalizeAiModelCallRecord(input.usage)
+  return {
+    ...input,
+    cost:
+      input.cost.status === "reported" ? input.cost : validatedEstimate(input.cost, input, usage),
+    ...(input.estimate === undefined
+      ? {}
+      : { estimate: validatedEstimate(input.estimate, input, usage) }),
+  }
+}
+
+function validatedEstimate(
+  cost: ModelCostEstimate,
+  input: RecoverAiModelCallInput,
+  usage: AiModelCallUsageRecord
+): ModelCostEstimate {
+  try {
+    assertJsonObject(cost, "model cost estimate")
+    const record = normalizeAiModelCallCostRecord(modelCostRecord({ ...input, cost }, usage.id))
+    if (!aiModelCallCostMatchesUsage(record, usage)) {
+      throw new TypeError("Estimate quantities do not match recorded usage.")
+    }
+    return structuredClone(cost)
+  } catch {
+    console.warn(
+      "[SixbModels] Invalid model cost estimate; preserving call usage and provider cost."
+    )
+    return { status: "unpriceable", reason: "inconsistent-usage" }
+  }
+}
+
+function modelCostRecord(
+  input: RecoverAiModelCallInput,
+  usageRecordId: string
+): AiModelCallCostRecord {
+  const billingIdentity = {
+    providerId: input.usage.providerId,
+    modelId: input.usage.requestedModelId,
+  }
+  const pricingContext: AiPricingContext = {
+    ...(input.route?.providerId === undefined ? {} : { routedProviderId: input.route.providerId }),
+    ...(input.route?.modelId === undefined ? {} : { routedModelId: input.route.modelId }),
+  }
+  if (input.cost.status === "reported" || input.cost.status === "rated") {
+    return {
+      projectId: input.usage.projectId,
+      usageRecordId,
+      status: "rated",
+      billingIdentity,
+      pricingContext,
+      priceSource: costPriceSource(input.usage, input.cost),
+      money: input.cost.money,
+      components: input.cost.status === "rated" ? input.cost.components : [],
+      ratedAt: new Date(input.ratedAt),
+    }
+  }
+
+  const reason =
+    input.cost.reason === "missing-rate-card"
+      ? "missingRateCard"
+      : input.cost.reason === "missing-usage"
+        ? "missingUsageMeter"
+        : "invalidUsageForFormula"
+  const missingMeters =
+    reason === "missingUsageMeter" ? normalizeMissingMeters(input.cost.missingMeters) : undefined
+  return {
+    projectId: input.usage.projectId,
+    usageRecordId,
+    status: "unpriceable",
+    billingIdentity,
+    pricingContext,
+    reason,
+    ...(missingMeters === undefined ? {} : { missingMeters }),
+    ratedAt: new Date(input.ratedAt),
+  }
+}
+
+function costPriceSource(
+  usage: RecoverAiModelCallInput["usage"],
+  cost: Extract<RecoverAiModelCallInput["cost"], { status: "reported" | "rated" }>
+) {
+  if (cost.status === "reported") {
+    return {
+      sourceId: "provider-reported",
+      sourceEntryId: `${usage.providerId}/${usage.responseId}`,
+      sourceVersion: "response-v1",
+      observedAt: new Date(usage.occurredAt),
+    }
+  }
+  const rates = cost.components.map(({ meter, rateAmountNanosPerMillion }) => ({
+    meter,
+    rateAmountNanosPerMillion,
+  }))
+  const version = createHash("sha256").update(JSON.stringify(rates)).digest("hex")
+  return {
+    sourceId: "model-rate-card",
+    sourceEntryId: `${usage.providerId}/${usage.requestedModelId}`,
+    sourceVersion: `sha256:${version}`,
+    observedAt: new Date(usage.occurredAt),
+  }
+}
+
+function normalizeMissingMeters(
+  meters: Extract<RecoverAiModelCallInput["cost"], { status: "unpriceable" }>["missingMeters"]
+): readonly AiBillableMeter[] {
+  const fallback: readonly AiBillableMeter[] = ["tokens.input.total", "tokens.output.total"]
+  return [...new Set<AiBillableMeter>(meters ?? fallback)]
+}
+
+export function requireAccountingCapabilities(
+  storage: Pick<Storage, "aiUsage" | "aiCosts" | "aiLimits">
+): AiAccountingCapabilities {
+  if (!storage.aiUsage || !storage.aiCosts || !storage.aiLimits) {
+    throw new Error(
+      "[SixbModels] AI model-call accounting requires storage.aiUsage, storage.aiCosts, and storage.aiLimits."
+    )
+  }
+  return { aiUsage: storage.aiUsage, aiCosts: storage.aiCosts, aiLimits: storage.aiLimits }
+}

package/src/models/execution/model-call-admission.ts

@@ -0,0 +1,127 @@
+import type { AiModelCallReservationIdentity } from "../../storage"
+import type { ModelCostEstimator } from "../pricing"
+
+/** Default output ceiling when a call and its resolved model supply no limit. */
+export const AI_MODEL_CALL_OUTPUT_TOKEN_ALLOWANCE = 4_096
+
+export type AiModelCallInputTokenEstimate =
+  | {
+      readonly status: "estimated"
+      readonly tokens: number
+      readonly method: "utf8BytesDividedByFour"
+    }
+  | {
+      readonly status: "unavailable"
+      readonly reason: "nonTextInput" | "unserializableInput"
+    }
+
+/** Provider-neutral snapshot passed to aggregate-budget admission before a provider request. */
+export interface AiModelCallAdmissionInput {
+  readonly projectId: string
+  readonly executionId: string
+  readonly attempt: number
+  readonly requesterGroupIds: readonly string[]
+  readonly callId: string
+  readonly providerId: string
+  readonly modelId: string
+  readonly costEstimator?: ModelCostEstimator
+  readonly inputTokens: AiModelCallInputTokenEstimate
+  readonly outputTokenAllowance: number
+  readonly estimatedTotalTokens?: number
+}
+
+export interface AiModelCallAdmissionDecision {
+  /** Active means accounting must reconcile the reservation for this provider attempt. */
+  readonly reservation: "active" | "none"
+}
+
+export type BeforeAiModelCall = (
+  input: AiModelCallAdmissionInput
+) => AiModelCallAdmissionDecision | Promise<AiModelCallAdmissionDecision>
+
+export type MarkAiModelCallUnknown = (input: AiModelCallReservationIdentity) => void | Promise<void>
+
+/**
+ * Estimate the exact, prepared provider request without requiring a provider-specific tokenizer.
+ *
+ * The heuristic is intentionally explicit so reservations are reproducible. Non-text inputs are
+ * reported as unavailable; an enforced token or cost policy can then fail closed instead of
+ * pretending that an image, file, or provider-defined payload has no token cost.
+ */
+export function estimateAiModelCallInputTokens(input: {
+  readonly prompt: unknown
+  readonly tools: unknown
+  readonly responseFormat: unknown
+}): AiModelCallInputTokenEstimate {
+  if (containsNonTextInput(input.prompt)) {
+    return { status: "unavailable", reason: "nonTextInput" }
+  }
+
+  let serialized: string | undefined
+  try {
+    serialized = JSON.stringify(input)
+  } catch {
+    return { status: "unavailable", reason: "unserializableInput" }
+  }
+  if (serialized === undefined) {
+    return { status: "unavailable", reason: "unserializableInput" }
+  }
+
+  const bytes = new TextEncoder().encode(serialized).byteLength
+  return {
+    status: "estimated",
+    tokens: Math.ceil(bytes / 4),
+    method: "utf8BytesDividedByFour",
+  }
+}
+
+export function aiModelCallOutputTokenAllowance(maxOutputTokens: number | undefined): number {
+  if (maxOutputTokens === undefined) return AI_MODEL_CALL_OUTPUT_TOKEN_ALLOWANCE
+  if (!Number.isSafeInteger(maxOutputTokens) || maxOutputTokens <= 0) {
+    throw new TypeError(
+      "[SixbModels] AI model-call maxOutputTokens must be a positive safe integer."
+    )
+  }
+  return maxOutputTokens
+}
+
+export function estimatedAiModelCallTotalTokens(
+  inputTokens: AiModelCallInputTokenEstimate,
+  outputTokenAllowance: number
+): number | undefined {
+  if (inputTokens.status !== "estimated") return undefined
+  const total = inputTokens.tokens + outputTokenAllowance
+  if (!Number.isSafeInteger(total)) {
+    throw new TypeError("[SixbModels] AI model-call token estimate exceeds the safe range.")
+  }
+  return total
+}
+
+function containsNonTextInput(value: unknown, seen = new WeakSet<object>()): boolean {
+  if (value === null || typeof value !== "object") return false
+  if (
+    value instanceof URL ||
+    value instanceof ArrayBuffer ||
+    ArrayBuffer.isView(value) ||
+    (typeof Blob !== "undefined" && value instanceof Blob)
+  ) {
+    return true
+  }
+  if (seen.has(value)) return false
+  seen.add(value)
+
+  if (!Array.isArray(value)) {
+    const type = Reflect.get(value, "type")
+    if (
+      type === "image" ||
+      type === "file" ||
+      type === "reasoning-file" ||
+      type === "custom" ||
+      type === "provider-state"
+    ) {
+      return true
+    }
+  }
+
+  return Object.values(value).some((entry) => containsNonTextInput(entry, seen))
+}

package/src/models/execution/model-call-limits.ts

@@ -0,0 +1,135 @@
+import { isSixbError } from "../../errors/internal"
+import type { AuthorizablePrincipal } from "../../execution"
+import type {
+  AiLimitPolicy,
+  AiLimitQuantity,
+  AiModelCallReservationIdentity,
+  ReserveAiModelCallResult,
+} from "../../storage"
+import {
+  aiLimitSubjectsFromAttribution,
+  aiUsageLimitExceededError,
+  aiUsageLimitUnavailableError,
+  applicableAiLimitPolicies,
+} from "../../storage/ai-limits/enforcement"
+import type {
+  AiModelCallAdmissionDecision,
+  AiModelCallAdmissionInput,
+  BeforeAiModelCall,
+  MarkAiModelCallUnknown,
+} from "./model-call-admission"
+import type { ModelCallAccountingStorage } from "./types"
+
+export interface AiModelCallLimitController {
+  readonly beforeModelCall: BeforeAiModelCall
+  readonly markModelCallUnknown: MarkAiModelCallUnknown
+}
+
+/** Build the reservation lifecycle for model calls from any execution. */
+export function createAiModelCallLimitController(input: {
+  readonly storage: Pick<ModelCallAccountingStorage, "aiLimits">
+  readonly projectId: string
+  readonly requestedBy?: AuthorizablePrincipal
+  readonly requesterGroupIds: readonly string[]
+}): AiModelCallLimitController {
+  const subjects = aiLimitSubjectsFromAttribution(input.requestedBy, input.requesterGroupIds)
+
+  return {
+    beforeModelCall: async (admission) => {
+      let applicable: readonly AiLimitPolicy[]
+      try {
+        const policies = await input.storage.aiLimits.listPolicies({
+          projectId: input.projectId,
+        })
+        applicable = applicableAiLimitPolicies(policies, subjects)
+      } catch (error) {
+        throw limitStorageUnavailable(error)
+      }
+      // This read is the fast no-policy path. A concurrently created policy linearizes after it;
+      // every later call will see it, while existing policies always continue to atomic reserve.
+      if (applicable.length === 0) return { reservation: "none" }
+
+      const estimates = reservationEstimates(
+        admission,
+        applicable.some((policy) => policy.limit.meter === "cost.catalogEstimated")
+      )
+      if (estimates.length === 0) {
+        throw aiUsageLimitUnavailableError(["missingEstimate"])
+      }
+
+      let result: ReserveAiModelCallResult
+      try {
+        result = await input.storage.aiLimits.reserveModelCall({
+          ...reservationIdentity(admission),
+          subjects,
+          estimates,
+        })
+      } catch (error) {
+        throw limitStorageUnavailable(error)
+      }
+      switch (result.status) {
+        case "reserved":
+          return { reservation: "active" } satisfies AiModelCallAdmissionDecision
+        case "notRequired":
+          return { reservation: "none" } satisfies AiModelCallAdmissionDecision
+        case "denied":
+          throw aiUsageLimitExceededError(result.resetAt)
+        case "unavailable":
+          throw aiUsageLimitUnavailableError(result.reasons)
+        case "terminal":
+          throw aiUsageLimitUnavailableError(["terminalReservationReplay"])
+      }
+    },
+    markModelCallUnknown: async (identity) => {
+      try {
+        await input.storage.aiLimits.markReservationUnknown(identity)
+      } catch (error) {
+        throw limitStorageUnavailable(error)
+      }
+    },
+  }
+}
+
+function reservationEstimates(
+  input: AiModelCallAdmissionInput,
+  needsCost: boolean
+): readonly AiLimitQuantity[] {
+  if (input.inputTokens.status !== "estimated" || input.estimatedTotalTokens === undefined) {
+    return []
+  }
+  const estimates: AiLimitQuantity[] = [
+    { meter: "tokens.total", amount: input.estimatedTotalTokens },
+  ]
+  if (needsCost) {
+    try {
+      const money = input.costEstimator?.estimateReservation?.({
+        inputTokens: input.inputTokens.tokens,
+        outputTokens: input.outputTokenAllowance,
+      })
+      if (money) estimates.push({ meter: "cost.catalogEstimated", amount: money })
+    } catch {
+      // An unavailable price must not prevent a token-only policy from being evaluated.
+      return estimates
+    }
+  }
+  return estimates
+}
+
+function reservationIdentity(input: AiModelCallAdmissionInput): AiModelCallReservationIdentity {
+  return {
+    projectId: input.projectId,
+    executionId: input.executionId,
+    attempt: input.attempt,
+    callId: input.callId,
+  }
+}
+
+function limitStorageUnavailable(error: unknown): unknown {
+  if (
+    isSixbError(error) &&
+    (error.code === "ai.usage_limit_exceeded" || error.code === "ai.usage_limit_unavailable")
+  ) {
+    return error
+  }
+  return aiUsageLimitUnavailableError(["limitStorageUnavailable"])
+}