@sixb/pg 0.1.1 → 0.1.3

package/src/pg-ai-usage-storage.ts

@@ -0,0 +1,290 @@
+import type { ReadonlyJsonObject } from "@sixb/core/storage"
+import {
+  type AiModelCallUsageInput,
+  type AiModelCallUsageRecord,
+  type AiUsageExecutionSummary,
+  type AiUsageStorage,
+  AiUsageStorageError,
+  aggregateAiModelCallUsage,
+  assertAiUsageExecutionId,
+  normalizeAiModelCallRecord,
+  type RecordAiModelCallInput,
+  type RecordAiModelCallResult,
+  type SummarizeAiUsageExecutionInput,
+  type SummarizeAiUsageExecutionsInput,
+} from "@sixb/core/storage"
+import type { SQLClient } from "./pg-client"
+import { type PgStoreClient, runPgTransaction } from "./transactions"
+
+/** PostgreSQL-backed append-only model-call usage ledger. */
+export class PgAiUsageStorage implements AiUsageStorage {
+  constructor(private readonly sql: PgStoreClient) {}
+
+  async recordModelCall(input: RecordAiModelCallInput): Promise<RecordAiModelCallResult> {
+    const record = normalizeAiModelCallRecord(input)
+    return runPgTransaction(this.sql, async (tx) => {
+      const existing = await this.findByIdempotencyKey(tx, record)
+      if (existing) return { record: existing, created: false }
+      await this.assertExecutionExists(tx, record.projectId, record.executionId)
+
+      const inserted = await tx<{ id: string }[]>`
+        INSERT INTO ai_model_call_usage (
+          project_id,
+          id,
+          execution_id,
+          attempt,
+          call_id,
+          provider_id,
+          requested_model_id,
+          response_model_id,
+          response_id,
+          input_tokens,
+          output_tokens,
+          total_tokens,
+          uncached_input_tokens,
+          cache_read_input_tokens,
+          cache_write_input_tokens,
+          text_output_tokens,
+          reasoning_output_tokens,
+          reporting_status,
+          raw_usage,
+          occurred_at,
+          recorded_at
+        ) VALUES (
+          ${record.projectId},
+          ${record.id},
+          ${record.executionId},
+          ${record.attempt},
+          ${record.callId},
+          ${record.providerId},
+          ${record.requestedModelId},
+          ${record.responseModelId ?? null},
+          ${record.responseId},
+          ${record.usage.inputTokens ?? null},
+          ${record.usage.outputTokens ?? null},
+          ${record.usage.totalTokens ?? null},
+          ${record.usage.uncachedInputTokens ?? null},
+          ${record.usage.cacheReadInputTokens ?? null},
+          ${record.usage.cacheWriteInputTokens ?? null},
+          ${record.usage.textOutputTokens ?? null},
+          ${record.usage.reasoningOutputTokens ?? null},
+          ${record.usage.reportingStatus},
+          ${record.rawUsage === undefined ? null : JSON.stringify(record.rawUsage)}::text::jsonb,
+          ${record.occurredAt},
+          ${record.recordedAt}
+        )
+        ON CONFLICT DO NOTHING
+        RETURNING id
+      `
+
+      if (inserted.length === 0) {
+        const concurrent = await this.findByIdempotencyKey(tx, record)
+        if (concurrent) return { record: concurrent, created: false }
+        throw new AiUsageStorageError(
+          "duplicate_id",
+          `[SixbPg] AI usage record '${record.id}' already exists for project '${record.projectId}'.`
+        )
+      }
+
+      for (const groupId of record.requesterGroupIds) {
+        await tx`
+          INSERT INTO ai_model_call_usage_groups (
+            project_id, usage_record_id, group_id, occurred_at
+          ) VALUES (
+            ${record.projectId}, ${record.id}, ${groupId}, ${record.occurredAt}
+          )
+        `
+      }
+
+      return { record: await this.requireById(tx, record.projectId, record.id), created: true }
+    })
+  }
+
+  async summarizeExecution(
+    input: SummarizeAiUsageExecutionInput
+  ): Promise<AiUsageExecutionSummary> {
+    const [summary] = await this.summarizeExecutions({
+      projectId: input.projectId,
+      executionIds: [input.executionId],
+    })
+    return summary!
+  }
+
+  async summarizeExecutions(
+    input: SummarizeAiUsageExecutionsInput
+  ): Promise<readonly AiUsageExecutionSummary[]> {
+    assertNonBlankProjectId(input.projectId)
+    for (const executionId of input.executionIds) assertAiUsageExecutionId(executionId)
+    if (input.executionIds.length === 0) return []
+
+    const rows = await this.executionRows(this.sql, input)
+    return aggregateExecutionRows(input.executionIds, rows)
+  }
+
+  private async findByIdempotencyKey(
+    sql: PgStoreClient,
+    record: Pick<
+      AiModelCallUsageRecord,
+      "projectId" | "executionId" | "attempt" | "callId" | "responseId"
+    >
+  ): Promise<AiModelCallUsageRecord | null> {
+    const rows = await sql<AiUsageRow[]>`
+      SELECT * FROM ai_model_call_usage
+      WHERE project_id = ${record.projectId}
+        AND execution_id = ${record.executionId}
+        AND attempt = ${record.attempt}
+        AND call_id = ${record.callId}
+        AND response_id = ${record.responseId}
+    `
+    return rows[0] ? this.rowToRecord(sql, rows[0]) : null
+  }
+
+  private async requireById(
+    sql: SQLClient,
+    projectId: string,
+    id: string
+  ): Promise<AiModelCallUsageRecord> {
+    const [row] = await sql<AiUsageRow[]>`
+      SELECT * FROM ai_model_call_usage WHERE project_id = ${projectId} AND id = ${id}
+    `
+    if (!row) throw new Error(`[SixbPg] AI usage record '${id}' disappeared after insert.`)
+    return this.rowToRecord(sql, row)
+  }
+
+  private async rowToRecord(sql: PgStoreClient, row: AiUsageRow): Promise<AiModelCallUsageRecord> {
+    const groupRows = await sql<Array<{ group_id: string }>>`
+      SELECT group_id FROM ai_model_call_usage_groups
+      WHERE project_id = ${row.project_id} AND usage_record_id = ${row.id}
+      ORDER BY group_id
+    `
+
+    return {
+      id: row.id,
+      projectId: row.project_id,
+      executionId: row.execution_id,
+      attempt: Number(row.attempt),
+      callId: row.call_id,
+      requesterGroupIds: groupRows.map((group) => group.group_id),
+      providerId: row.provider_id,
+      requestedModelId: row.requested_model_id,
+      ...(row.response_model_id === null ? {} : { responseModelId: row.response_model_id }),
+      responseId: row.response_id,
+      usage: {
+        ...usageFromRow(row),
+        ...(row.total_tokens === null ? {} : { totalTokens: Number(row.total_tokens) }),
+        reportingStatus: row.reporting_status,
+      },
+      ...(row.raw_usage === null ? {} : { rawUsage: rawUsageFromRow(row.raw_usage) }),
+      occurredAt: new Date(row.occurred_at),
+      recordedAt: new Date(row.recorded_at),
+    }
+  }
+
+  private executionRows(
+    sql: PgStoreClient,
+    input: SummarizeAiUsageExecutionsInput
+  ): Promise<AiUsageRow[]> {
+    const executionIds = [...new Set(input.executionIds)]
+
+    return sql<AiUsageRow[]>`
+      WITH requested(execution_id) AS (
+        SELECT * FROM unnest(${sql.array(executionIds)}::text[])
+      )
+      SELECT usage.* FROM requested
+      JOIN ai_model_call_usage AS usage
+        ON usage.project_id = ${input.projectId}
+        AND usage.execution_id = requested.execution_id
+    `
+  }
+
+  private async assertExecutionExists(
+    sql: PgStoreClient,
+    projectId: string,
+    executionId: string
+  ): Promise<void> {
+    const [execution] = await sql<{ exists: boolean }[]>`
+      SELECT EXISTS (
+        SELECT 1 FROM executions WHERE project_id = ${projectId} AND id = ${executionId}
+      ) AS exists
+    `
+    if (execution?.exists) return
+    throw new AiUsageStorageError(
+      "missing_execution",
+      `[SixbPg] AI usage execution '${executionId}' does not exist in project '${projectId}'.`
+    )
+  }
+}
+
+interface AiUsageRow {
+  readonly project_id: string
+  readonly id: string
+  readonly execution_id: string
+  readonly attempt: number | string
+  readonly call_id: string
+  readonly provider_id: string
+  readonly requested_model_id: string
+  readonly response_model_id: string | null
+  readonly response_id: string
+  readonly input_tokens: number | string | null
+  readonly output_tokens: number | string | null
+  readonly total_tokens: number | string | null
+  readonly uncached_input_tokens: number | string | null
+  readonly cache_read_input_tokens: number | string | null
+  readonly cache_write_input_tokens: number | string | null
+  readonly text_output_tokens: number | string | null
+  readonly reasoning_output_tokens: number | string | null
+  readonly reporting_status: AiModelCallUsageRecord["usage"]["reportingStatus"]
+  readonly raw_usage: ReadonlyJsonObject | string | null
+  readonly occurred_at: Date | string
+  readonly recorded_at: Date | string
+}
+
+function aggregateExecutionRows(
+  executionIds: readonly string[],
+  rows: readonly AiUsageRow[]
+): readonly AiUsageExecutionSummary[] {
+  const usageByExecutionId = new Map<string, AiModelCallUsageInput[]>()
+  for (const executionId of executionIds) usageByExecutionId.set(executionId, [])
+  for (const row of rows) {
+    usageByExecutionId.get(row.execution_id)?.push(usageFromRow(row))
+  }
+  return executionIds.map((executionId) => {
+    const usages = usageByExecutionId.get(executionId) ?? []
+    return {
+      modelCallCount: usages.length,
+      usage: aggregateAiModelCallUsage(usages),
+    }
+  })
+}
+
+function usageFromRow(row: AiUsageRow): AiModelCallUsageInput {
+  return {
+    ...(row.input_tokens === null ? {} : { inputTokens: Number(row.input_tokens) }),
+    ...(row.output_tokens === null ? {} : { outputTokens: Number(row.output_tokens) }),
+    ...(row.uncached_input_tokens === null
+      ? {}
+      : { uncachedInputTokens: Number(row.uncached_input_tokens) }),
+    ...(row.cache_read_input_tokens === null
+      ? {}
+      : { cacheReadInputTokens: Number(row.cache_read_input_tokens) }),
+    ...(row.cache_write_input_tokens === null
+      ? {}
+      : { cacheWriteInputTokens: Number(row.cache_write_input_tokens) }),
+    ...(row.text_output_tokens === null
+      ? {}
+      : { textOutputTokens: Number(row.text_output_tokens) }),
+    ...(row.reasoning_output_tokens === null
+      ? {}
+      : { reasoningOutputTokens: Number(row.reasoning_output_tokens) }),
+  }
+}
+
+function rawUsageFromRow(value: Exclude<AiUsageRow["raw_usage"], null>): ReadonlyJsonObject {
+  return typeof value === "string" ? (JSON.parse(value) as ReadonlyJsonObject) : value
+}
+
+function assertNonBlankProjectId(projectId: string): void {
+  if (typeof projectId !== "string" || projectId.trim().length === 0) {
+    throw new TypeError("[Sixb] AI usage projectId must be nonblank.")
+  }
+}

package/src/pg-execution-storage.ts

@@ -0,0 +1,147 @@
+import {
+  type ExecutionStorageRow,
+  executionRecordFromStorageRow,
+  executionRecordToStorageRow,
+  normalizeExecutionRecord,
+  validateExecutionRecordReferences,
+} from "@sixb/core/internal/execution-storage"
+import type { CreateExecutionInput, ExecutionRecord, ExecutionStorage } from "@sixb/core/storage"
+import { ExecutionStorageError } from "@sixb/core/storage"
+import type { PgAuthStorage } from "./auth-storage"
+import { isUniqueViolation } from "./storage-errors"
+import type { PgStoreClient } from "./transactions"
+
+export class PgExecutionStorage implements ExecutionStorage {
+  constructor(
+    private readonly sql: PgStoreClient,
+    private readonly auth: PgAuthStorage
+  ) {}
+
+  async create(input: CreateExecutionInput): Promise<ExecutionRecord> {
+    const record = normalizeExecutionRecord(input)
+    await validateExecutionRecordReferences(record, {
+      auth: this.auth,
+      getExecution: (params) => this.getById(params),
+    })
+    const row = executionRecordToStorageRow(record)
+
+    try {
+      const [created] = await this.sql<PgExecutionRow[]>`
+        INSERT INTO executions (
+          project_id,
+          id,
+          executor_kind,
+          executor_id,
+          source_kind,
+          source_id,
+          requested_by_user_id,
+          requested_by_service_account_id,
+          correlation_id,
+          parent_execution_id,
+          authority_kind,
+          authority_user_id,
+          authority_service_account_id,
+          authority_session_id,
+          authority_access_token_id,
+          authority_primitive_kind,
+          authority_primitive_id,
+          authority_kernel_operation,
+          created_at
+        ) VALUES (
+          ${row.projectId},
+          ${row.id},
+          ${row.executorKind},
+          ${row.executorId},
+          ${row.sourceKind},
+          ${row.sourceId},
+          ${row.requestedByUserId},
+          ${row.requestedByServiceAccountId},
+          ${row.correlationId},
+          ${row.parentExecutionId},
+          ${row.authorityKind},
+          ${row.authorityUserId},
+          ${row.authorityServiceAccountId},
+          ${row.authoritySessionId},
+          ${row.authorityAccessTokenId},
+          ${row.authorityPrimitiveKind},
+          ${row.authorityPrimitiveId},
+          ${row.authorityKernelOperation},
+          ${row.createdAt}
+        )
+        RETURNING *
+      `
+
+      if (!created) {
+        throw new Error("[SixbPg] Execution insert returned no row.")
+      }
+      return executionRecordFromStorageRow(toStorageRow(created))
+    } catch (error) {
+      if (isUniqueViolation(error)) {
+        throw new ExecutionStorageError(
+          "duplicate_execution",
+          `[SixbPg] Execution '${record.id}' already exists in project '${record.projectId}'.`,
+          { cause: error }
+        )
+      }
+      throw error
+    }
+  }
+
+  async getById(params: {
+    readonly projectId: string
+    readonly id: string
+  }): Promise<ExecutionRecord | null> {
+    const [row] = await this.sql<PgExecutionRow[]>`
+      SELECT *
+      FROM executions
+      WHERE project_id = ${params.projectId} AND id = ${params.id}
+    `
+    return row ? executionRecordFromStorageRow(toStorageRow(row)) : null
+  }
+}
+
+function toStorageRow(row: PgExecutionRow): ExecutionStorageRow {
+  return {
+    projectId: row.project_id,
+    id: row.id,
+    executorKind: row.executor_kind,
+    executorId: row.executor_id,
+    sourceKind: row.source_kind,
+    sourceId: row.source_id,
+    requestedByUserId: row.requested_by_user_id,
+    requestedByServiceAccountId: row.requested_by_service_account_id,
+    correlationId: row.correlation_id,
+    parentExecutionId: row.parent_execution_id,
+    authorityKind: row.authority_kind,
+    authorityUserId: row.authority_user_id,
+    authorityServiceAccountId: row.authority_service_account_id,
+    authoritySessionId: row.authority_session_id,
+    authorityAccessTokenId: row.authority_access_token_id,
+    authorityPrimitiveKind: row.authority_primitive_kind,
+    authorityPrimitiveId: row.authority_primitive_id,
+    authorityKernelOperation: row.authority_kernel_operation,
+    createdAt: new Date(row.created_at),
+  }
+}
+
+interface PgExecutionRow {
+  readonly project_id: string
+  readonly id: string
+  readonly executor_kind: ExecutionStorageRow["executorKind"]
+  readonly executor_id: string
+  readonly source_kind: ExecutionStorageRow["sourceKind"]
+  readonly source_id: string
+  readonly requested_by_user_id: string | null
+  readonly requested_by_service_account_id: string | null
+  readonly correlation_id: string
+  readonly parent_execution_id: string | null
+  readonly authority_kind: ExecutionStorageRow["authorityKind"]
+  readonly authority_user_id: string | null
+  readonly authority_service_account_id: string | null
+  readonly authority_session_id: string | null
+  readonly authority_access_token_id: string | null
+  readonly authority_primitive_kind: ExecutionStorageRow["authorityPrimitiveKind"]
+  readonly authority_primitive_id: string | null
+  readonly authority_kernel_operation: ExecutionStorageRow["authorityKernelOperation"]
+  readonly created_at: Date | string
+}