@sixb/core 0.1.7 → 0.1.9

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"])
+}

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

@@ -0,0 +1,235 @@
+import { randomUUID } from "node:crypto"
+import type { ReadonlyJsonObject, RecordAiModelCallInput } from "../../storage"
+import type { ModelCallEndEvent, ModelUsage } from "../events"
+import type { LanguageModel, LanguageModelRequest } from "../language-model"
+import { ModelUsageRecordingError } from "./errors"
+import {
+  normalizeModelCallAccounting,
+  recordAiModelCallAccounting,
+  requireAccountingCapabilities,
+} from "./model-call-accounting"
+import {
+  aiModelCallOutputTokenAllowance,
+  type BeforeAiModelCall,
+  estimateAiModelCallInputTokens,
+  estimatedAiModelCallTotalTokens,
+} from "./model-call-admission"
+import type { AiModelCallLimitController } from "./model-call-limits"
+import { isPermanentAiUsageRecoveryError } from "./model-call-recovery"
+import type {
+  ModelCallAccountingStorage,
+  RecoverAiModelCall,
+  RecoverAiModelCallInput,
+} from "./types"
+import { aiModelCallUsageFromModel } from "./usage"
+
+const STORAGE_RETRY_DELAYS_MS = [50, 200, 600] as const
+
+export interface AiModelCallRecorderInput {
+  readonly storage: ModelCallAccountingStorage
+  readonly projectId: string
+  readonly executionId: string
+  readonly attempt: number
+  readonly requesterGroupIds: readonly string[]
+  readonly limits: AiModelCallLimitController
+  readonly recoverAiModelCall: RecoverAiModelCall
+}
+
+interface AiModelCallRecorderInternals {
+  readonly generateId?: () => string
+  readonly now?: () => Date
+  readonly retryDelaysMs?: readonly number[]
+  readonly sleep?: (ms: number) => Promise<void>
+  readonly recordAccounting?: (input: RecoverAiModelCallInput) => Promise<void>
+}
+
+/** Persist every completed provider call before another billable loop step can start. */
+export class AiModelCallRecorder {
+  private readonly generateId: () => string
+  private readonly now: () => Date
+  private readonly retryDelaysMs: readonly number[]
+  private readonly sleep: (ms: number) => Promise<void>
+  private readonly recordAccounting: (input: RecoverAiModelCallInput) => Promise<void>
+  private readonly reservations = new Set<string>()
+  private admissionError: unknown
+  private recordingError: ModelUsageRecordingError | undefined
+
+  constructor(
+    private readonly input: AiModelCallRecorderInput,
+    internals: AiModelCallRecorderInternals = {}
+  ) {
+    requireAccountingCapabilities(input.storage)
+    this.generateId = internals.generateId ?? (() => `ai_usage_${randomUUID()}`)
+    this.now = internals.now ?? (() => new Date())
+    this.retryDelaysMs = internals.retryDelaysMs ?? STORAGE_RETRY_DELAYS_MS
+    this.sleep = internals.sleep ?? sleep
+    this.recordAccounting =
+      internals.recordAccounting ??
+      (async (recovery) => {
+        await recordAiModelCallAccounting({ storage: this.input.storage, ...recovery })
+      })
+  }
+
+  /** Wrap the resolved model so all calls in an execution share admission. */
+  wrapModel(model: LanguageModel): LanguageModel {
+    return {
+      providerId: model.providerId,
+      modelId: model.modelId,
+      definition: model.definition,
+      costEstimator: model.costEstimator,
+      ...(model.validateResponseFormat === undefined
+        ? {}
+        : { validateResponseFormat: model.validateResponseFormat.bind(model) }),
+      stream: async (request) => {
+        this.assertHealthy()
+        const inputTokens = estimateAiModelCallInputTokens({
+          prompt: request.messages,
+          tools: request.tools,
+          responseFormat: request.responseFormat,
+        })
+        const outputTokenAllowance = aiModelCallOutputTokenAllowance(request.maxOutputTokens)
+        let decision: Awaited<ReturnType<BeforeAiModelCall>>
+        try {
+          decision = await this.input.limits.beforeModelCall({
+            ...this.identity(request),
+            requesterGroupIds: this.input.requesterGroupIds,
+            providerId: model.providerId,
+            modelId: model.modelId,
+            costEstimator: model.costEstimator,
+            inputTokens,
+            outputTokenAllowance,
+            estimatedTotalTokens: estimatedAiModelCallTotalTokens(
+              inputTokens,
+              outputTokenAllowance
+            ),
+          })
+        } catch (error) {
+          // Execution boundaries must retain the original coded admission failure.
+          this.admissionError = error
+          throw error
+        }
+        if (decision.reservation === "active") this.reservations.add(request.callId)
+        try {
+          return await model.stream(request)
+        } catch (error) {
+          if (this.reservations.has(request.callId)) {
+            await this.input.limits.markModelCallUnknown(this.identity(request))
+          }
+          throw error
+        }
+      },
+    }
+  }
+
+  private identity(request: Pick<LanguageModelRequest, "callId">) {
+    return {
+      projectId: this.input.projectId,
+      executionId: this.input.executionId,
+      attempt: this.input.attempt,
+      callId: request.callId,
+    }
+  }
+
+  readonly onModelCallEnd = async (event: ModelCallEndEvent): Promise<void> => {
+    // Calls already in flight still need their own accounting after another call fails.
+    // Only admission is gated by the execution's sticky failure.
+    try {
+      const occurredAt = this.now()
+      const record: RecordAiModelCallInput = {
+        id: this.generateId(),
+        projectId: this.input.projectId,
+        executionId: this.input.executionId,
+        attempt: this.input.attempt,
+        callId: event.callId,
+        requesterGroupIds: this.input.requesterGroupIds,
+        providerId: event.providerId,
+        ...(event.providerIds === undefined ? {} : { providerIds: event.providerIds }),
+        requestedModelId: event.modelId,
+        ...(event.requestedReasoning === undefined
+          ? {}
+          : { requestedReasoning: event.requestedReasoning }),
+        ...(event.responseModelId === undefined ? {} : { responseModelId: event.responseModelId }),
+        responseId: event.responseId,
+        usage: aiModelCallUsageFromModel(event.usage),
+        ...(event.usage.raw === undefined ? {} : { rawUsage: rawUsage(event.usage) }),
+        occurredAt,
+      }
+      const recoveryInput = normalizeModelCallAccounting({
+        usage: record,
+        cost: event.cost,
+        ...(event.estimate === undefined ? {} : { estimate: event.estimate }),
+        ...(event.route === undefined ? {} : { route: event.route }),
+        ratedAt: new Date(occurredAt),
+        ...(this.reservations.has(event.callId) ? { reconcileLimitReservation: true } : {}),
+      })
+
+      try {
+        await retryOperation(
+          () => this.recordAccounting(recoveryInput),
+          this.retryDelaysMs,
+          this.sleep,
+          (error) => !isPermanentAiUsageRecoveryError(error)
+        )
+      } catch (storageError) {
+        if (isPermanentAiUsageRecoveryError(storageError)) throw storageError
+
+        try {
+          await retryOperation(
+            () => this.input.recoverAiModelCall(recoveryInput),
+            this.retryDelaysMs,
+            this.sleep,
+            (error) => !isPermanentAiUsageRecoveryError(error)
+          )
+        } catch (recoveryError) {
+          throw new AggregateError(
+            [storageError, recoveryError],
+            "Direct AI accounting and durable recovery both failed."
+          )
+        }
+        throw new ModelUsageRecordingError(this.input.executionId, event.callId, true, {
+          cause: storageError,
+        })
+      }
+    } catch (error) {
+      const recordingError =
+        error instanceof ModelUsageRecordingError
+          ? error
+          : new ModelUsageRecordingError(this.input.executionId, event.callId, false, {
+              cause: error,
+            })
+      this.recordingError ??= recordingError
+      throw recordingError
+    }
+  }
+
+  assertHealthy(): void {
+    if (this.admissionError !== undefined) throw this.admissionError
+    if (this.recordingError) throw this.recordingError
+  }
+}
+
+function rawUsage(usage: ModelUsage): ReadonlyJsonObject {
+  return structuredClone(usage.raw ?? {})
+}
+
+async function retryOperation<TResult>(
+  operation: () => Promise<TResult>,
+  retryDelaysMs: readonly number[],
+  wait: (ms: number) => Promise<void>,
+  shouldRetry: (error: unknown) => boolean
+): Promise<TResult> {
+  for (let attempt = 0; ; attempt += 1) {
+    try {
+      return await operation()
+    } catch (error) {
+      if (!shouldRetry(error)) throw error
+      const delay = retryDelaysMs[attempt]
+      if (delay === undefined) throw error
+      await wait(delay)
+    }
+  }
+}
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}

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

@@ -0,0 +1,125 @@
+import type { RecordAiModelCallInput, RecordAiModelCallResult } from "../../storage"
+import { AiCostStorageError, AiLimitStorageError, AiUsageStorageError } from "../../storage"
+import { normalizeAiModelCallRecord } from "../../storage/ai-usage/record"
+import { recordAiModelCallAccounting } from "./model-call-accounting"
+import type {
+  AiModelCallAccountingPayload,
+  AiModelCallRecordPayload,
+  AiModelCallRecoveryRecord,
+  ModelCallAccountingStorage,
+  RecoverAiModelCallInput,
+} from "./types"
+import { aiModelCallUsageFromModel } from "./usage"
+
+class InvalidAiUsageRecoveryJobError extends Error {
+  readonly name = "InvalidAiUsageRecoveryJobError"
+
+  constructor(message: string, options?: ErrorOptions) {
+    super(`[SixbModels] ${message}`, options)
+  }
+}
+
+/** Encode one failed accounting transaction for a durable, JSON-safe handoff. */
+export function modelCallRecoveryPayload(input: RecoverAiModelCallInput): {
+  readonly record: AiModelCallRecordPayload
+  readonly accounting: AiModelCallAccountingPayload
+} {
+  normalizeAiModelCallRecord(input.usage)
+  return {
+    record: toQueuePayload(input.usage),
+    accounting: toAccountingPayload(input),
+  }
+}
+
+/** Replay one durable handoff through the idempotent atomic accounting boundary. */
+export async function recordRecoveredAiModelCall(
+  storage: Pick<ModelCallAccountingStorage, "transaction">,
+  job: AiModelCallRecoveryRecord
+): Promise<RecordAiModelCallResult> {
+  const usage = fromQueuePayload(job)
+  const accounting = accountingFromQueuePayload(job)
+  return recordAiModelCallAccounting({ storage, usage, ...accounting })
+}
+
+/** Validation and referential-integrity failures cannot become valid through queue redelivery. */
+export function isPermanentAiUsageRecoveryError(error: unknown): boolean {
+  return (
+    error instanceof InvalidAiUsageRecoveryJobError ||
+    error instanceof TypeError ||
+    (error instanceof AiUsageStorageError &&
+      (error.code === "duplicate_id" || error.code === "missing_execution")) ||
+    (error instanceof AiCostStorageError &&
+      (error.code === "missing_usage" || error.code === "cost_mismatch")) ||
+    (error instanceof AiLimitStorageError &&
+      (error.code === "missing_execution" ||
+        error.code === "missing_reservation" ||
+        error.code === "missing_usage_record" ||
+        error.code === "usage_mismatch" ||
+        error.code === "reservation_conflict" ||
+        error.code === "reconciliation_conflict" ||
+        error.code === "invalid_reservation_state"))
+  )
+}
+
+function toQueuePayload(record: RecordAiModelCallInput): AiModelCallRecordPayload {
+  return {
+    id: record.id,
+    executionId: record.executionId,
+    attempt: record.attempt,
+    callId: record.callId,
+    requesterGroupIds: [...record.requesterGroupIds],
+    providerId: record.providerId,
+    ...(record.providerIds === undefined
+      ? {}
+      : { providerIds: structuredClone(record.providerIds) }),
+    requestedModelId: record.requestedModelId,
+    ...(record.requestedReasoning === undefined
+      ? {}
+      : { requestedReasoning: structuredClone(record.requestedReasoning) }),
+    ...(record.responseModelId === undefined ? {} : { responseModelId: record.responseModelId }),
+    responseId: record.responseId,
+    usage: aiModelCallUsageFromModel(record.usage),
+    ...(record.rawUsage === undefined ? {} : { rawUsage: structuredClone(record.rawUsage) }),
+    occurredAt: record.occurredAt.toISOString(),
+  }
+}
+
+function toAccountingPayload(input: RecoverAiModelCallInput): AiModelCallAccountingPayload {
+  return {
+    cost: structuredClone(input.cost),
+    ...(input.estimate === undefined ? {} : { estimate: structuredClone(input.estimate) }),
+    ...(input.route === undefined ? {} : { route: structuredClone(input.route) }),
+    ratedAt: input.ratedAt.toISOString(),
+    ...(input.reconcileLimitReservation ? { reconcileLimitReservation: true } : {}),
+  }
+}
+
+function fromQueuePayload(job: AiModelCallRecoveryRecord): RecordAiModelCallInput {
+  const occurredAt = parseDate(job.payload.record.occurredAt, job.id, "occurredAt")
+  return { ...job.payload.record, projectId: job.projectId, occurredAt }
+}
+
+function accountingFromQueuePayload(
+  job: AiModelCallRecoveryRecord
+): Omit<RecoverAiModelCallInput, "usage"> {
+  const accounting = job.payload.accounting
+  return {
+    cost: structuredClone(accounting.cost),
+    ...(accounting.estimate === undefined
+      ? {}
+      : { estimate: structuredClone(accounting.estimate) }),
+    ...(accounting.route === undefined ? {} : { route: structuredClone(accounting.route) }),
+    ratedAt: parseDate(accounting.ratedAt, job.id, "ratedAt"),
+    reconcileLimitReservation: accounting.reconcileLimitReservation === true,
+  }
+}
+
+function parseDate(value: string, jobId: string, field: string): Date {
+  const date = new Date(value)
+  if (!Number.isFinite(date.getTime())) {
+    throw new InvalidAiUsageRecoveryJobError(
+      `AI usage recovery job '${jobId}' has an invalid ${field} timestamp.`
+    )
+  }
+  return date
+}

package/src/models/execution/recovery-queue.ts

@@ -0,0 +1,31 @@
+import { createSixbError } from "../../errors/internal"
+import type { AgentQueueJob, Queue } from "../../queues"
+import { modelCallRecoveryPayload } from "./model-call-recovery"
+import type { RecoverAiModelCallInput } from "./types"
+
+/** Stable queue identity: a lost enqueue response can safely retry the same accounting handoff. */
+function agentAiUsageRecoveryJobId(recordId: string): string {
+  return `agt_usage_job_${recordId}`
+}
+
+/** Transport adapter for the agent worker's existing durable recovery lane. */
+export async function enqueueAiModelCallRecovery(
+  queue: Pick<Queue<AgentQueueJob>, "enqueue">,
+  input: RecoverAiModelCallInput
+): Promise<void> {
+  const payload = modelCallRecoveryPayload(input)
+  const record = input.usage
+  const jobId = agentAiUsageRecoveryJobId(record.id)
+  const [job] = await queue.enqueue({
+    projectId: record.projectId,
+    jobs: [{ id: jobId, type: "agent.ai-usage.record.requested", payload }],
+  })
+
+  if (job?.id !== jobId || job.type !== "agent.ai-usage.record.requested") {
+    throw createSixbError(
+      "internal.unexpected",
+      `[SixbModels] Agent queue did not confirm AI usage recovery job '${jobId}'.`,
+      { details: { jobId } }
+    )
+  }
+}

package/src/models/execution/types.ts

@@ -0,0 +1,58 @@
+import type {
+  AiCostStorage,
+  AiLimitStorage,
+  AiUsageStorage,
+  RecordAiModelCallInput,
+  Storage,
+} from "../../storage"
+import type { ModelRoute } from "../events"
+import type { ModelCallCost, ModelCostEstimate } from "../pricing"
+
+/** Only the stores and transaction boundary needed by model-call accounting. */
+export interface ModelCallAccountingStorage extends Pick<Storage, "transaction"> {
+  readonly aiUsage: AiUsageStorage
+  readonly aiCosts: AiCostStorage
+  readonly aiLimits: AiLimitStorage
+}
+
+export interface RecoverAiModelCallInput {
+  readonly usage: RecordAiModelCallInput
+  readonly cost: ModelCallCost
+  readonly estimate?: ModelCostEstimate
+  readonly route?: ModelRoute
+  readonly ratedAt: Date
+  /** True only when admission created an aggregate-budget reservation for this provider attempt. */
+  readonly reconcileLimitReservation?: boolean
+}
+
+export type RecoverAiModelCall = (input: RecoverAiModelCallInput) => Promise<void>
+
+/** JSON-safe ledger input shared by durable recovery transports. */
+export type AiModelCallRecordPayload = Omit<
+  RecordAiModelCallInput,
+  "projectId" | "usage" | "occurredAt" | "recordedAt"
+> & {
+  readonly usage: {
+    readonly [Field in keyof RecordAiModelCallInput["usage"]]: RecordAiModelCallInput["usage"][Field]
+  }
+  readonly occurredAt: string
+}
+
+export interface AiModelCallAccountingPayload {
+  readonly estimate?: ModelCostEstimate
+  readonly cost: ModelCallCost
+  readonly route?: ModelRoute
+  readonly ratedAt: string
+  /** True when admission created a reservation that recovery must reconcile. */
+  readonly reconcileLimitReservation?: boolean
+}
+
+/** Transport identity is used only for recovery diagnostics, not accounting identity. */
+export interface AiModelCallRecoveryRecord {
+  readonly id: string
+  readonly projectId: string
+  readonly payload: {
+    readonly record: AiModelCallRecordPayload
+    readonly accounting: AiModelCallAccountingPayload
+  }
+}

package/src/models/execution/usage.ts

@@ -0,0 +1,23 @@
+import type { AiModelCallUsageInput } from "../../storage"
+import type { ModelUsage } from "../events"
+
+/** Map one provider-neutral model call into Sixb's durable accounting vocabulary. */
+export function aiModelCallUsageFromModel(usage: ModelUsage): AiModelCallUsageInput {
+  return {
+    ...(usage.inputTokens === undefined ? {} : { inputTokens: usage.inputTokens }),
+    ...(usage.outputTokens === undefined ? {} : { outputTokens: usage.outputTokens }),
+    ...(usage.uncachedInputTokens === undefined
+      ? {}
+      : { uncachedInputTokens: usage.uncachedInputTokens }),
+    ...(usage.cacheReadInputTokens === undefined
+      ? {}
+      : { cacheReadInputTokens: usage.cacheReadInputTokens }),
+    ...(usage.cacheWriteInputTokens === undefined
+      ? {}
+      : { cacheWriteInputTokens: usage.cacheWriteInputTokens }),
+    ...(usage.textOutputTokens === undefined ? {} : { textOutputTokens: usage.textOutputTokens }),
+    ...(usage.reasoningOutputTokens === undefined
+      ? {}
+      : { reasoningOutputTokens: usage.reasoningOutputTokens }),
+  }
+}