@sixb/pg 0.1.5 → 0.1.7

package/src/ai-cost-groups.ts

@@ -0,0 +1,102 @@
+import {
+  type AiModelCallGroupRow,
+  aiModelCallGroupFragmentFromRow,
+  buildAiModelCallGroupsFromFragments,
+  normalizeAiModelCallAccountingQuery,
+} from "@sixb/core/internal/ai-cost-storage-provider"
+import type { ListAiModelCallGroupsInput, ListAiModelCallGroupsResult } from "@sixb/core/storage"
+import type { PgStoreClient } from "./transactions"
+
+/** Paginate initiating executions, then aggregate their matching calls without loading usage payloads. */
+export async function listAiModelCallGroups(
+  sql: PgStoreClient,
+  input: ListAiModelCallGroupsInput
+): Promise<ListAiModelCallGroupsResult> {
+  const query = normalizeAiModelCallAccountingQuery(input)
+  const parameters = [
+    query.projectId,
+    query.from.toISOString(),
+    query.to.toISOString(),
+    query.providerId ?? null,
+    query.modelId ?? null,
+    query.valuationStatus ?? null,
+  ]
+  const [count] = await sql.unsafe<{ total: string }[]>(COUNT_SQL, parameters)
+  const rows = await sql.unsafe<AiModelCallGroupRow[]>(PAGE_SQL, [
+    ...parameters,
+    query.limit,
+    query.offset,
+  ])
+  const total = Number(count?.total)
+  if (!Number.isSafeInteger(total) || total < 0)
+    throw new RangeError("[Sixb] Invalid model-call group count.")
+  const items = buildAiModelCallGroupsFromFragments(rows.map(aiModelCallGroupFragmentFromRow))
+  return { items, total, hasMore: query.offset + items.length < total }
+}
+
+const FILTERED_SQL = `
+WITH filtered AS (
+  SELECT COALESCE(root_agent.execution_id, usage.execution_id) AS root_execution_id,
+    usage.execution_id, usage.provider_id, usage.requested_model_id, usage.occurred_at,
+    usage.total_tokens,
+    COALESCE(cost.status, 'unvalued') AS valuation_status,
+    cost.currency AS amount_currency, cost.amount_nanos,
+    root_agent.id AS root_agent_run_id,
+    root_agent.thread_id AS root_thread_id,
+    direct_agent.kind AS direct_kind,
+    direct_agent.id AS direct_run_id,
+    direct_agent.parent_run_id AS parent_run_id,
+    workflow_node.node_id AS workflow_agent_step_id,
+    workflow_agent.node_run_id AS workflow_node_run_id,
+    workflow_node.workflow_id AS workflow_id,
+    workflow_node.workflow_run_id AS workflow_run_id
+  FROM ai_model_call_usage AS usage
+  LEFT JOIN agent_runs AS direct_agent
+    ON direct_agent.project_id = usage.project_id AND direct_agent.execution_id = usage.execution_id
+  LEFT JOIN agent_runs AS root_agent
+    ON root_agent.project_id = usage.project_id
+    AND root_agent.id = COALESCE(direct_agent.parent_run_id, direct_agent.id)
+  LEFT JOIN workflow_agent_node_runs AS workflow_agent
+    ON workflow_agent.project_id = usage.project_id AND workflow_agent.execution_id = usage.execution_id
+  LEFT JOIN workflow_node_runs AS workflow_node
+    ON workflow_node.project_id = workflow_agent.project_id AND workflow_node.id = workflow_agent.node_run_id
+  LEFT JOIN ai_model_call_valuations AS cost
+    ON cost.project_id = usage.project_id AND cost.usage_record_id = usage.id
+  WHERE usage.project_id = $1
+    AND usage.occurred_at >= $2::timestamptz
+    AND usage.occurred_at < $3::timestamptz
+    AND ($4::text IS NULL OR usage.provider_id = $4)
+    AND ($5::text IS NULL OR usage.requested_model_id = $5)
+    AND ($6::text IS NULL OR COALESCE(cost.status, 'unvalued') = $6)
+), roots AS (
+  SELECT root_execution_id, MAX(occurred_at) AS last_call_at
+  FROM filtered GROUP BY root_execution_id
+)
+`
+const COUNT_SQL = `${FILTERED_SQL} SELECT CAST(COUNT(*) AS TEXT) AS total FROM roots`
+const PAGE_SQL = `${FILTERED_SQL}, page AS (
+  SELECT root_execution_id FROM roots
+  ORDER BY last_call_at DESC, root_execution_id COLLATE "C" ASC LIMIT $7 OFFSET $8
+)
+SELECT filtered.root_execution_id, execution_id, provider_id, requested_model_id,
+  MIN(occurred_at) AS first_call_at, MAX(occurred_at) AS last_call_at,
+  CAST(COUNT(*) AS TEXT) AS call_count,
+  CAST(COUNT(total_tokens) AS TEXT) AS token_count,
+  CAST(SUM(total_tokens) AS TEXT) AS total_tokens,
+  valuation_status, amount_currency,
+  CAST(SUM(amount_nanos / 1000000000000) AS TEXT) AS amount_high,
+  CAST(SUM((amount_nanos / 1000000) % 1000000) AS TEXT) AS amount_middle,
+  CAST(SUM(amount_nanos % 1000000) AS TEXT) AS amount_low,
+  MAX(root_agent_run_id) AS root_agent_run_id,
+  MAX(root_thread_id) AS root_thread_id,
+  MAX(direct_kind) AS direct_kind,
+  MAX(direct_run_id) AS direct_run_id,
+  MAX(parent_run_id) AS parent_run_id,
+  MAX(workflow_agent_step_id) AS workflow_agent_step_id,
+  MAX(workflow_node_run_id) AS workflow_node_run_id,
+  MAX(workflow_id) AS workflow_id,
+  MAX(workflow_run_id) AS workflow_run_id
+FROM filtered JOIN page ON page.root_execution_id = filtered.root_execution_id
+GROUP BY filtered.root_execution_id, execution_id, provider_id, requested_model_id,
+  valuation_status, amount_currency
+`

package/src/auth-storage/auth-storage.ts

@@ -1,8 +1,11 @@
 import type {
   AuthStorage,
+  CompleteDeviceAuthorizationInput,
+  CompleteDeviceAuthorizationResult,
   CompleteMagicLinkSignInInput,
   CompleteOidcSignInInput,
   CompleteSignInResult,
+  DeviceAuthorizationRecord,
   GroupMembershipRecord,
   InvitationRecord,
   SessionRecord,
@@ -19,6 +22,10 @@ import {
   runPgTransaction,
 } from "../transactions"
 import { PgAuthAccessTokenStore } from "./access-tokens"
+import {
+  getPgDeviceAuthorizationById,
+  PgAuthDeviceAuthorizationStore,
+} from "./device-authorizations"
 import { PgAuthGroupMembershipStore } from "./group-memberships"
 import { PgAuthUserIdentityStore } from "./identities"
 import { PgAuthInvitationStore } from "./invitations"
@@ -77,6 +84,7 @@ export class PgAuthStorage implements AuthStorage {
   readonly groupMemberships: PgAuthGroupMembershipStore
   readonly magicLinks: PgAuthMagicLinkStore
   readonly oidcAuthorizationAttempts: PgAuthOidcAuthorizationAttemptStore
+  readonly deviceAuthorizations: PgAuthDeviceAuthorizationStore
 
   constructor(options: PgAuthStorageOptions) {
     this.sql = options.sql
@@ -90,6 +98,43 @@ export class PgAuthStorage implements AuthStorage {
     this.groupMemberships = new PgAuthGroupMembershipStore(this.sql)
     this.magicLinks = new PgAuthMagicLinkStore(this.sql)
     this.oidcAuthorizationAttempts = new PgAuthOidcAuthorizationAttemptStore(this.sql)
+    this.deviceAuthorizations = new PgAuthDeviceAuthorizationStore(this.sql)
+  }
+
+  async completeDeviceAuthorization(
+    input: CompleteDeviceAuthorizationInput
+  ): Promise<CompleteDeviceAuthorizationResult> {
+    return runPgTransaction(this.sql, async (tx) => {
+      const authorization = await getPgDeviceAuthorizationById(tx, input.projectId, input.id)
+      if (!authorization) {
+        throw new AuthStorageError(
+          "missing_device_authorization",
+          `[Sixb] Device authorization '${input.id}' not found for project '${input.projectId}'.`
+        )
+      }
+      assertCompletableDeviceAuthorization(authorization, input)
+      const accessToken = await new PgAuthAccessTokenStore(tx).create(input.accessToken)
+      const [consumed] = await tx`
+        UPDATE auth_device_authorizations
+        SET status = 'consumed', consumed_at = ${input.completedAt}
+        WHERE project_id = ${input.projectId} AND id = ${input.id} AND status = 'approved'
+        RETURNING *
+      `
+      if (!consumed) {
+        throw new AuthStorageError(
+          "invalid_device_authorization",
+          "[Sixb] Device authorization was already consumed."
+        )
+      }
+      return {
+        authorization: {
+          ...authorization,
+          status: "consumed" as const,
+          consumedAt: new Date(input.completedAt),
+        },
+        accessToken,
+      }
+    })
   }
 
   async completeMagicLinkSignIn(
@@ -796,3 +841,42 @@ function assertSignInSessionAudience(
     `[Sixb] Sign-in session audience '${sessionAudience}' does not match stored auth audience '${storedAudience}' for project '${projectId}'.`
   )
 }
+
+function assertCompletableDeviceAuthorization(
+  authorization: DeviceAuthorizationRecord,
+  input: CompleteDeviceAuthorizationInput
+): void {
+  if (authorization.deviceCodeHash !== input.deviceCodeHash) {
+    throw new AuthStorageError(
+      "invalid_device_authorization",
+      "[Sixb] Invalid device authorization."
+    )
+  }
+  if (authorization.status === "pending") {
+    throw new AuthStorageError(
+      "pending_device_authorization",
+      "[Sixb] Device authorization is pending."
+    )
+  }
+  if (authorization.status === "denied") {
+    throw new AuthStorageError(
+      "device_authorization_denied",
+      "[Sixb] Device authorization was denied."
+    )
+  }
+  if (
+    authorization.status !== "approved" ||
+    authorization.expiresAt <= input.completedAt ||
+    !authorization.approvedUserId ||
+    !authorization.approvedSessionId ||
+    input.accessToken.projectId !== input.projectId ||
+    input.accessToken.subjectType !== "user" ||
+    input.accessToken.subjectId !== authorization.approvedUserId ||
+    input.accessToken.createdBySessionId !== authorization.approvedSessionId
+  ) {
+    throw new AuthStorageError(
+      "invalid_device_authorization",
+      "[Sixb] Device authorization is invalid."
+    )
+  }
+}

package/src/auth-storage/device-authorizations.ts

@@ -0,0 +1,131 @@
+import type {
+  AuthDeviceAuthorizationStore,
+  CreateDeviceAuthorizationInput,
+  DeviceAuthorizationRecord,
+} from "@sixb/core/storage"
+import { AuthStorageError, MAX_PENDING_DEVICE_AUTHORIZATIONS } from "@sixb/core/storage"
+import {
+  authLockKey,
+  lockAdvisoryKeys,
+  type PgStoreClient,
+  runPgTransaction,
+} from "../transactions"
+import type { PgAuthDeviceAuthorizationRow } from "./rows"
+import { rowToDeviceAuthorizationRecord } from "./rows"
+import { assertNonEmpty, mapUniqueConstraintError } from "./shared"
+
+export class PgAuthDeviceAuthorizationStore implements AuthDeviceAuthorizationStore {
+  constructor(private readonly sql: PgStoreClient) {}
+
+  async create(input: CreateDeviceAuthorizationInput): Promise<DeviceAuthorizationRecord> {
+    return runPgTransaction(this.sql, async (tx) => {
+      await lockAdvisoryKeys(tx, [authLockKey("device-authorizations", input.projectId)])
+      await tx`
+        DELETE FROM auth_device_authorizations
+        WHERE project_id = ${input.projectId} AND expires_at <= ${input.createdAt}
+      `
+      const [pending] = await tx<{ readonly count: string | number }[]>`
+        SELECT COUNT(*)::bigint AS count FROM auth_device_authorizations
+        WHERE project_id = ${input.projectId} AND status = 'pending'
+      `
+      if (Number(pending?.count ?? 0) >= MAX_PENDING_DEVICE_AUTHORIZATIONS) {
+        throw new AuthStorageError(
+          "device_authorization_limit_exceeded",
+          `[Sixb] Project '${input.projectId}' has too many pending device authorizations.`
+        )
+      }
+      try {
+        const [row] = await tx<PgAuthDeviceAuthorizationRow[]>`
+        INSERT INTO auth_device_authorizations (
+          project_id, id, device_code_hash, user_code, client_name, token_name,
+          token_expires_at, status, created_at, expires_at
+        ) VALUES (
+          ${assertNonEmpty(input.projectId, "Project id")},
+          ${assertNonEmpty(input.id, "Device authorization id")},
+          ${assertNonEmpty(input.deviceCodeHash, "Device code hash")},
+          ${assertNonEmpty(input.userCode, "User code")},
+          ${assertNonEmpty(input.clientName, "Client name")},
+          ${assertNonEmpty(input.tokenName, "Token name")},
+          ${input.tokenExpiresAt}, 'pending', ${input.createdAt}, ${input.expiresAt}
+        ) RETURNING *
+      `
+        return rowToDeviceAuthorizationRecord(row)
+      } catch (error) {
+        mapUniqueConstraintError(
+          error,
+          "duplicate_device_authorization",
+          `[Sixb] Device authorization '${input.id}' already exists for project '${input.projectId}'.`
+        )
+      }
+    })
+  }
+
+  async getById(params: { readonly projectId: string; readonly id: string }) {
+    return getPgDeviceAuthorizationById(this.sql, params.projectId, params.id)
+  }
+
+  async getByUserCode(params: { readonly projectId: string; readonly userCode: string }) {
+    const [row] = await this.sql<PgAuthDeviceAuthorizationRow[]>`
+      SELECT * FROM auth_device_authorizations
+      WHERE project_id = ${params.projectId} AND user_code = ${params.userCode}
+    `
+    return row ? rowToDeviceAuthorizationRecord(row) : null
+  }
+
+  async approve(params: {
+    readonly projectId: string
+    readonly id: string
+    readonly userId: string
+    readonly sessionId: string
+    readonly approvedAt: Date
+  }) {
+    const [row] = await this.sql<PgAuthDeviceAuthorizationRow[]>`
+      UPDATE auth_device_authorizations
+      SET status = 'approved', approved_user_id = ${params.userId},
+          approved_session_id = ${params.sessionId}, approved_at = ${params.approvedAt}
+      WHERE project_id = ${params.projectId} AND id = ${params.id}
+        AND status = 'pending' AND expires_at > ${params.approvedAt}
+        AND EXISTS (
+          SELECT 1 FROM auth_sessions
+          WHERE project_id = ${params.projectId} AND id = ${params.sessionId}
+            AND user_id = ${params.userId} AND revoked_at IS NULL
+            AND expires_at > ${params.approvedAt}
+            AND (absolute_expires_at IS NULL OR absolute_expires_at > ${params.approvedAt})
+        )
+      RETURNING *
+    `
+    if (!row) throw invalidTransition(params.id, params.projectId)
+    return rowToDeviceAuthorizationRecord(row)
+  }
+
+  async deny(params: { readonly projectId: string; readonly id: string; readonly deniedAt: Date }) {
+    const [row] = await this.sql<PgAuthDeviceAuthorizationRow[]>`
+      UPDATE auth_device_authorizations
+      SET status = 'denied', denied_at = ${params.deniedAt}
+      WHERE project_id = ${params.projectId} AND id = ${params.id}
+        AND status = 'pending' AND expires_at > ${params.deniedAt}
+      RETURNING *
+    `
+    if (!row) throw invalidTransition(params.id, params.projectId)
+    return rowToDeviceAuthorizationRecord(row)
+  }
+}
+
+export async function getPgDeviceAuthorizationById(
+  sql: PgStoreClient,
+  projectId: string,
+  id: string
+) {
+  const [row] = await sql<PgAuthDeviceAuthorizationRow[]>`
+    SELECT * FROM auth_device_authorizations
+    WHERE project_id = ${projectId} AND id = ${id}
+  `
+  return row ? rowToDeviceAuthorizationRecord(row) : null
+}
+
+function invalidTransition(id: string, projectId: string): AuthStorageError {
+  return new AuthStorageError(
+    "invalid_device_authorization",
+    `[Sixb] Device authorization '${id}' is no longer pending for project '${projectId}'.`
+  )
+}

package/src/auth-storage/rows.ts

@@ -1,6 +1,7 @@
 import type { Principal } from "@sixb/core"
 import type {
   AccessTokenRecord,
+  DeviceAuthorizationRecord,
   GroupMembershipRecord,
   GroupMembershipSource,
   InvitationRecord,
@@ -150,6 +151,24 @@ export interface PgAuthOidcAttemptRow {
   readonly consumed_at: PgDate | null
 }
 
+export interface PgAuthDeviceAuthorizationRow {
+  readonly project_id: string
+  readonly id: string
+  readonly device_code_hash: string
+  readonly user_code: string
+  readonly client_name: string
+  readonly token_name: string
+  readonly token_expires_at: PgDate
+  readonly status: DeviceAuthorizationRecord["status"]
+  readonly approved_user_id: string | null
+  readonly approved_session_id: string | null
+  readonly created_at: PgDate
+  readonly expires_at: PgDate
+  readonly approved_at: PgDate | null
+  readonly denied_at: PgDate | null
+  readonly consumed_at: PgDate | null
+}
+
 export function rowToUserRecord(row: PgAuthUserRow): UserRecord {
   return {
     id: row.id,
@@ -304,6 +323,28 @@ export function rowToOidcAuthorizationAttemptRecord(
   }
 }
 
+export function rowToDeviceAuthorizationRecord(
+  row: PgAuthDeviceAuthorizationRow
+): DeviceAuthorizationRecord {
+  return {
+    id: row.id,
+    projectId: row.project_id,
+    deviceCodeHash: row.device_code_hash,
+    userCode: row.user_code,
+    clientName: row.client_name,
+    tokenName: row.token_name,
+    tokenExpiresAt: toDate(row.token_expires_at),
+    status: row.status,
+    approvedUserId: row.approved_user_id ?? undefined,
+    approvedSessionId: row.approved_session_id ?? undefined,
+    createdAt: toDate(row.created_at),
+    expiresAt: toDate(row.expires_at),
+    approvedAt: row.approved_at ? toDate(row.approved_at) : undefined,
+    deniedAt: row.denied_at ? toDate(row.denied_at) : undefined,
+    consumedAt: row.consumed_at ? toDate(row.consumed_at) : undefined,
+  }
+}
+
 export function serializeOptionalRecord(
   value: Readonly<Record<string, unknown>> | undefined
 ): string | null {

package/src/auth-storage/shared.ts

@@ -761,6 +761,7 @@ export function mapUniqueConstraintError(
   error: unknown,
   code:
     | "duplicate_access_token"
+    | "duplicate_device_authorization"
     | "duplicate_identity"
     | "duplicate_invitation"
     | "duplicate_magic_link"

package/src/index.ts

@@ -20,6 +20,7 @@ import {
 } from "@sixb/core/internal/storage-operation-scope"
 import {
   type AiCostStorage,
+  type AiLimitStorage,
   type AiUsageStorage,
   createTransactionStorageProxy,
   type ObjectStorage,
@@ -33,6 +34,7 @@ import { createPostgresStorageMigrators, dropSchema } from "./migrations"
 import { PgOntologyStorage, type PgOntologyTransactionContext } from "./ontology-storage"
 import { PgActionRunStorage } from "./pg-action-run-storage"
 import { PgAiCostStorage } from "./pg-ai-cost-storage"
+import { PgAiLimitStorage } from "./pg-ai-limit-storage"
 import { PgAiUsageStorage } from "./pg-ai-usage-storage"
 import { createPgClient, type SQL, type SQLClient } from "./pg-client"
 import { PgExecutionStorage } from "./pg-execution-storage"
@@ -138,6 +140,7 @@ export class PostgresStorage implements MigrationCapableStorage {
   readonly agents: PgAgentStorage
   readonly aiUsage: AiUsageStorage
   readonly aiCosts: AiCostStorage
+  readonly aiLimits: AiLimitStorage
   readonly actionRuns: PgActionRunStorage
   readonly pipelineRuns: PgPipelineRunStorage
   readonly workflowRuns: PgWorkflowRunStorage
@@ -212,6 +215,7 @@ export class PostgresStorage implements MigrationCapableStorage {
     this.agents = createAgentOperationScope(stores.agents, scope)
     this.aiUsage = createOperationScopedFacade(stores.aiUsage, scope)
     this.aiCosts = createOperationScopedFacade(stores.aiCosts, scope)
+    this.aiLimits = createOperationScopedFacade(stores.aiLimits, scope)
     this.actionRuns = createOperationScopedFacade(stores.actionRuns, scope)
     this.pipelineRuns = createOperationScopedFacade(stores.pipelineRuns, scope)
     this.workflowRuns = createWorkflowRunOperationScope(stores.workflowRuns, scope)
@@ -258,6 +262,7 @@ export class PostgresStorage implements MigrationCapableStorage {
             transactionContext.materializations.assertCommittable()
             return result
           } finally {
+            await txStorage.aiLimits.settleOperations()
             transactionContext.active = false
             txStorage.ontology.deactivateSessions()
             transactionContext.materializations.deactivate()
@@ -305,7 +310,7 @@ export class PostgresStorage implements MigrationCapableStorage {
   private createTransactionStorage(
     client: PgStoreClient,
     transactionContext: PgOntologyTransactionContext
-  ): Storage & { readonly ontology: PgOntologyStorage } {
+  ): Storage & { readonly ontology: PgOntologyStorage; readonly aiLimits: PgAiLimitStorage } {
     return {
       ...createPostgresStores(client, {
         runOntologyOperation: async (run) => run(client),
@@ -342,6 +347,7 @@ function createPostgresStores(
     agents: new PgAgentStorage({ sql, executions }),
     aiUsage: new PgAiUsageStorage(sql),
     aiCosts: new PgAiCostStorage(sql),
+    aiLimits: new PgAiLimitStorage(sql),
     actionRuns: new PgActionRunStorage(sql, executions),
     pipelineRuns: new PgPipelineRunStorage(sql, executions),
     workflowRuns: new PgWorkflowRunStorage(sql, executions),
@@ -363,6 +369,7 @@ interface PostgresStoreSet {
   readonly agents: PgAgentStorage
   readonly aiUsage: PgAiUsageStorage
   readonly aiCosts: PgAiCostStorage
+  readonly aiLimits: PgAiLimitStorage
   readonly actionRuns: PgActionRunStorage
   readonly pipelineRuns: PgPipelineRunStorage
   readonly workflowRuns: PgWorkflowRunStorage

package/src/migrations/029-model-accounting.sql

@@ -0,0 +1,24 @@
+-- Call-time model configuration and provider identity. Historical identity remains unknown.
+ALTER TABLE ai_model_call_usage
+  ADD COLUMN requested_reasoning JSONB
+  CHECK (
+    requested_reasoning IS NULL OR
+    jsonb_typeof(requested_reasoning) IN ('string', 'object')
+  );
+ALTER TABLE ai_model_call_usage ADD COLUMN provider_ids JSONB;
+
+-- Upgrade the shipped valuation reason while retaining the selected-cost projection.
+ALTER TABLE ai_model_call_valuations
+  DROP CONSTRAINT ai_model_call_valuations_reason_check;
+
+UPDATE ai_model_call_valuations SET reason = 'missingRateCard'
+  WHERE reason = 'missingCatalogEntry';
+
+ALTER TABLE ai_model_call_valuations ADD CONSTRAINT ai_model_call_valuations_reason_check
+  CHECK (reason IS NULL OR reason IN (
+    'missingBillingIdentity',
+    'missingRateCard',
+    'missingUsageMeter',
+    'unsupportedPricingDimension',
+    'invalidUsageForFormula'
+  ));

package/src/migrations/030-ai-usage-limits.sql

@@ -0,0 +1,94 @@
+CREATE TABLE ai_usage_limits (
+  project_id TEXT NOT NULL CHECK (length(btrim(project_id)) > 0),
+  id TEXT NOT NULL CHECK (length(btrim(id)) > 0),
+  subject_type TEXT NOT NULL CHECK (
+    subject_type IN ('project', 'group', 'user', 'serviceAccount')
+  ),
+  subject_id TEXT NOT NULL CHECK (
+    (subject_type = 'project' AND subject_id = '')
+    OR (subject_type <> 'project' AND length(btrim(subject_id)) > 0)
+  ),
+  meter TEXT NOT NULL CHECK (meter IN ('tokens.total', 'cost.catalogEstimated')),
+  currency TEXT NOT NULL CHECK (
+    (meter = 'tokens.total' AND currency = '')
+    OR (meter = 'cost.catalogEstimated' AND currency = 'USD')
+  ),
+  limit_amount NUMERIC(78, 0) NOT NULL CHECK (limit_amount >= 0),
+  period_kind TEXT NOT NULL CHECK (period_kind = 'calendarMonth'),
+  enabled BOOLEAN NOT NULL,
+  created_at TIMESTAMPTZ NOT NULL,
+  updated_at TIMESTAMPTZ NOT NULL,
+  PRIMARY KEY (project_id, id),
+  UNIQUE (project_id, subject_type, subject_id, meter, currency)
+);
+
+CREATE INDEX idx_ai_usage_limits_project_enabled
+  ON ai_usage_limits (project_id, enabled, subject_type, subject_id, meter, currency);
+
+CREATE TABLE ai_usage_limit_periods (
+  project_id TEXT NOT NULL CHECK (length(btrim(project_id)) > 0),
+  subject_type TEXT NOT NULL CHECK (
+    subject_type IN ('project', 'group', 'user', 'serviceAccount')
+  ),
+  subject_id TEXT NOT NULL CHECK (
+    (subject_type = 'project' AND subject_id = '')
+    OR (subject_type <> 'project' AND length(btrim(subject_id)) > 0)
+  ),
+  meter TEXT NOT NULL CHECK (meter IN ('tokens.total', 'cost.catalogEstimated')),
+  currency TEXT NOT NULL CHECK (
+    (meter = 'tokens.total' AND currency = '')
+    OR (meter = 'cost.catalogEstimated' AND currency = 'USD')
+  ),
+  period_start TIMESTAMPTZ NOT NULL,
+  period_end TIMESTAMPTZ NOT NULL,
+  actual_amount NUMERIC(78, 0) NOT NULL DEFAULT 0 CHECK (actual_amount >= 0),
+  reserved_amount NUMERIC(78, 0) NOT NULL DEFAULT 0 CHECK (reserved_amount >= 0),
+  unknown_amount NUMERIC(78, 0) NOT NULL DEFAULT 0 CHECK (unknown_amount >= 0),
+  accounting_status TEXT NOT NULL DEFAULT 'complete' CHECK (
+    accounting_status IN ('complete', 'unavailable')
+  ),
+  updated_at TIMESTAMPTZ NOT NULL,
+  PRIMARY KEY (
+    project_id, subject_type, subject_id, meter, currency, period_start
+  )
+);
+
+CREATE INDEX idx_ai_usage_limit_periods_subject_period
+  ON ai_usage_limit_periods (
+    project_id, subject_type, subject_id, period_start, period_end, meter, currency
+  );
+
+CREATE TABLE ai_model_call_reservations (
+  project_id TEXT NOT NULL CHECK (length(btrim(project_id)) > 0),
+  execution_id TEXT NOT NULL CHECK (length(btrim(execution_id)) > 0),
+  attempt BIGINT NOT NULL CHECK (attempt >= 1),
+  call_id TEXT NOT NULL CHECK (length(btrim(call_id)) > 0),
+  buckets JSONB NOT NULL CHECK (jsonb_typeof(buckets) = 'array'),
+  request_key TEXT NOT NULL,
+  period_start TIMESTAMPTZ NOT NULL,
+  period_end TIMESTAMPTZ NOT NULL,
+  state TEXT NOT NULL CHECK (state IN ('active', 'reconciled', 'unknown')),
+  usage_record_id TEXT,
+  reserved_at TIMESTAMPTZ NOT NULL,
+  updated_at TIMESTAMPTZ NOT NULL,
+  PRIMARY KEY (project_id, execution_id, attempt, call_id),
+  FOREIGN KEY (project_id, execution_id)
+    REFERENCES executions(project_id, id)
+    ON DELETE RESTRICT,
+  FOREIGN KEY (project_id, usage_record_id)
+    REFERENCES ai_model_call_usage(project_id, id)
+    ON DELETE RESTRICT,
+  CHECK (
+    (
+      state = 'reconciled'
+      AND usage_record_id IS NOT NULL
+    )
+    OR (
+      state <> 'reconciled'
+      AND usage_record_id IS NULL
+    )
+  )
+);
+
+CREATE INDEX idx_ai_model_call_reservations_active
+  ON ai_model_call_reservations (project_id, state, period_end, execution_id, attempt, call_id);