feat: refine metadata catalog workflows

2026-09-02 11:38:47 +02:00
parent 076c9742c5
commit 4531746038
48 changed files with 1345 additions and 493 deletions
@@ -59,7 +59,7 @@ import { DescriptionGenerationWorker } from "./catalog/description-generation-wo
 import { SensitiveDataSuggester } from "./catalog/sensitive-data-suggester.js";
 import { SensitiveDataSuggestionRunner } from "./catalog/sensitive-data-suggestion-runner.js";
 import {
-  PostgresDescriptionSourceSampler,
+  ConcreteDescriptionSourceSampler,
  type DescriptionSourceSampler,
 } from "./catalog/description-source-sampler.js";
 import { catalogDescriptionGenerationRoutes } from "./routes/catalog-description-generation.js";
@@ -177,7 +177,7 @@ export function buildApp(config: AppConfig, deps?: BuildAppDeps): FastifyInstanc
    { connectTimeoutMs: config.workspaceDiagnosticTimeoutMs },
  );
  const descriptionSourceSampler = deps?.descriptionSourceSampler
-    ?? new PostgresDescriptionSourceSampler(catalogPostgresAccess);
+    ?? new ConcreteDescriptionSourceSampler(catalogPostgresAccess, workspaceSecretStore);
  const descriptionGenerationWorker = new DescriptionGenerationWorker(
    catalogRepository,
    workspaceRegistry,
@@ -1,7 +1,7 @@
 import { z } from "zod";
 import type { WorkspaceRegistry } from "../workspaces/registry.js";
 import type { MetadataGenerationModels, ResolvedMetadataGenerationModel } from "./metadata-generation-models.js";
-import type { ModelCompleter, ModelCompletionMessage } from "./model-completer.js";
+import type { ModelCompleter, ModelCompletionMessage, ModelCompletionResult } from "./model-completer.js";
 import {
  ModelCompletionCancelledError,
  ModelCompletionProviderError,
@@ -38,6 +38,7 @@ const MAX_SAMPLE_FIELDS_PER_ROW = 4;
 const MAX_SAMPLE_COLUMNS = 4;
 const MAX_REPRESENTATIVE_VALUES_PER_REQUEST = 5;
 const MAX_TARGET_SAMPLE_JSON_BYTES = 8 * 1024;
+const MAX_COMPLETION_ATTEMPTS_PER_BATCH = 2;
 const generatedOutcomeSchema = z.object({
  targetId: z.uuid(),
  outcome: z.literal("generated"),
@@ -55,7 +56,9 @@ const completionResponseSchema = z.object({
  results: z.array(outcomeSchema).min(1).max(MAX_TARGETS_PER_BATCH),
 }).strict();

-class InvalidModelOutcomeError extends Error {}
+class InvalidModelJsonError extends Error {}
+class InvalidModelSchemaError extends Error {}
+class MissingModelTargetsError extends Error {}
 class DescriptionGenerationBatchError extends Error {
  constructor(
    readonly failure: unknown,
@@ -143,6 +146,9 @@ interface DescriptionGenerationCounters {
  generated: number;
  nonGeneratable: number;
  failed: number;
+  inputTokens: number;
+  cacheReadTokens: number;
+  outputTokens: number;
  consecutiveTechnicalFailures: number;
 }

@@ -152,6 +158,9 @@ function persistedCounters(counters: DescriptionGenerationCounters) {
    generated: counters.generated,
    nonGeneratable: counters.nonGeneratable,
    failed: counters.failed,
+    inputTokens: counters.inputTokens,
+    cacheReadTokens: counters.cacheReadTokens,
+    outputTokens: counters.outputTokens,
  };
 }

@@ -625,36 +634,47 @@ function messagesFor(
 }

 function parseOutcomes(content: string, expectedTargetIds: readonly string[]): Map<string, ParsedOutcome> {
+  const trimmed = content.trim();
+  const fenced = /^```(?:json)?[ \t]*\r?\n([\s\S]*?)\r?\n```$/iu.exec(trimmed);
+  let parsed: unknown;
  try {
-    const trimmed = content.trim();
-    const fenced = /^```(?:json)?[ \t]*\r?\n([\s\S]*?)\r?\n```$/iu.exec(trimmed);
-    const outcomes = completionResponseSchema.parse(JSON.parse(fenced?.[1] ?? trimmed)).results;
-    const expected = new Set(expectedTargetIds);
-    if (outcomes.length !== expectedTargetIds.length || expected.size !== expectedTargetIds.length) {
-      throw new InvalidModelOutcomeError();
-    }
-    const mapped = new Map<string, ParsedOutcome>();
-    for (const outcome of outcomes) {
-      if (!expected.has(outcome.targetId) || mapped.has(outcome.targetId)) {
-        throw new InvalidModelOutcomeError();
-      }
-      mapped.set(outcome.targetId, outcome.outcome === "generated"
-        ? { ...outcome, description: outcome.description.trim() }
-        : outcome);
-    }
-    if (mapped.size !== expected.size) throw new InvalidModelOutcomeError();
-    return mapped;
-  } catch (error) {
-    if (error instanceof InvalidModelOutcomeError) throw error;
-    throw new InvalidModelOutcomeError();
+    parsed = JSON.parse(fenced?.[1] ?? trimmed);
+  } catch {
+    throw new InvalidModelJsonError();
  }
+
+  const response = completionResponseSchema.safeParse(parsed);
+  if (!response.success) throw new InvalidModelSchemaError();
+
+  const outcomes = response.data.results;
+  const expected = new Set(expectedTargetIds);
+  if (outcomes.length !== expectedTargetIds.length || expected.size !== expectedTargetIds.length) {
+    throw new MissingModelTargetsError();
+  }
+  const mapped = new Map<string, ParsedOutcome>();
+  for (const outcome of outcomes) {
+    if (!expected.has(outcome.targetId) || mapped.has(outcome.targetId)) {
+      throw new MissingModelTargetsError();
+    }
+    mapped.set(outcome.targetId, outcome.outcome === "generated"
+      ? { ...outcome, description: outcome.description.trim() }
+      : outcome);
+  }
+  if (mapped.size !== expected.size) throw new MissingModelTargetsError();
+  return mapped;
 }

 function safeFailure(error: unknown): string {
  if (error instanceof DescriptionGenerationFailureStreakError) return error.message;
  const failure = error instanceof DescriptionGenerationBatchError ? error.failure : error;
  if (failure instanceof ModelCompletionProviderError) return "The model provider request failed.";
-  if (failure instanceof InvalidModelOutcomeError) return "The model response was invalid.";
+  if (failure instanceof InvalidModelJsonError) return "The model response was not valid JSON.";
+  if (failure instanceof InvalidModelSchemaError) {
+    return "The model response did not match the required schema.";
+  }
+  if (failure instanceof MissingModelTargetsError) {
+    return "The model response was missing one or more requested targets.";
+  }
  return "Description generation failed.";
 }

@@ -907,6 +927,9 @@ export class DescriptionGenerationWorker {
      generated: 0,
      nonGeneratable: 0,
      failed: 0,
+      inputTokens: 0,
+      cacheReadTokens: 0,
+      outputTokens: 0,
      consecutiveTechnicalFailures: 0,
    };
    await this.processTargets(run, database, plan.columnTargets, model, counters, signal);
@@ -973,20 +996,41 @@ export class DescriptionGenerationWorker {
        );
      }
      throwIfCancelled(signal);
-      let outcomes: Map<string, ParsedOutcome>;
-      try {
-        const content = await this.completer.complete({
-          model,
-          messages: messagesFor(database, batch, run.language, sourceSamples),
-          signal,
-        });
-        throwIfCancelled(signal);
-        outcomes = parseOutcomes(content, batch.map((target) => (
-          target.kind === "column" ? target.column.id : target.table.id
-        )));
-      } catch (error) {
-        if (error instanceof ModelCompletionCancelledError) throw error;
-        const batchError = new DescriptionGenerationBatchError(error, batch.map(failureTarget));
+      const expectedTargetIds = batch.map((target) => (
+        target.kind === "column" ? target.column.id : target.table.id
+      ));
+      const messages = messagesFor(database, batch, run.language, sourceSamples);
+      let outcomes: Map<string, ParsedOutcome> | undefined;
+      let terminalFailure: unknown;
+      for (let attempt = 1; attempt <= MAX_COMPLETION_ATTEMPTS_PER_BATCH; attempt += 1) {
+        try {
+          const completion = await this.completer.complete({ model, messages, signal });
+          const result: ModelCompletionResult = typeof completion === "string"
+            ? { content: completion, usage: { input: 0, cacheRead: 0, output: 0 } }
+            : completion;
+          counters.inputTokens += result.usage.input;
+          counters.cacheReadTokens += result.usage.cacheRead;
+          counters.outputTokens += result.usage.output;
+          throwIfCancelled(signal);
+          outcomes = parseOutcomes(result.content, expectedTargetIds);
+          break;
+        } catch (error) {
+          if (error instanceof ModelCompletionCancelledError) throw error;
+          terminalFailure = error;
+          if (attempt < MAX_COMPLETION_ATTEMPTS_PER_BATCH) {
+            await this.appendEvent(
+              run.id,
+              "warning",
+              `${safeFailure(error)} Retrying batch (attempt ${attempt + 1} of ${MAX_COMPLETION_ATTEMPTS_PER_BATCH}).`,
+            );
+          }
+        }
+      }
+      if (!outcomes) {
+        const batchError = new DescriptionGenerationBatchError(
+          terminalFailure,
+          batch.map(failureTarget),
+        );
        counters.processed += batch.length;
        counters.failed += batch.length;
        counters.consecutiveTechnicalFailures += 1;
@@ -1009,7 +1053,10 @@ export class DescriptionGenerationWorker {
        const targetId = target.kind === "column" ? target.column.id : target.table.id;
        const outcome = outcomes.get(targetId);
        if (!outcome) {
-          throw new DescriptionGenerationBatchError(new InvalidModelOutcomeError(), [failureTarget(target)]);
+          throw new DescriptionGenerationBatchError(
+            new MissingModelTargetsError(),
+            [failureTarget(target)],
+          );
        }
        const generatedDescription = outcome.outcome === "generated"
          ? outcome.description
@@ -1,5 +1,8 @@
+import { readFile } from "node:fs/promises";
+import type { WorkspaceSecretStore } from "../workspaces/secret-store.js";
 import type { CatalogPostgresAccess } from "./postgres-access.js";
-import type { WorkspaceDatabase } from "./types.js";
+import { CATALOG_SECRET_IDS } from "./secrets.js";
+import { CatalogConnectorError, type WorkspaceDatabase } from "./types.js";

 const MAX_SOURCE_ROWS = 5;
 const MAX_REPRESENTATIVE_VALUES = 5;
@@ -79,15 +82,82 @@ function distinctKey(value: Exclude<DescriptionSourceSampleValue, null>): string
  return `${typeof value}:${String(value)}`;
 }

-/** PostgreSQL-wire sampler. REST bindings remain unsupported by CatalogPostgresAccess. */
-export class PostgresDescriptionSourceSampler implements DescriptionSourceSampler {
-  constructor(private readonly access: CatalogPostgresAccess) {}
+function columnsFor(target: DescriptionSourceSamplingTarget): string[] {
+  return [...new Set(target.columnNames)].slice(0, MAX_SOURCE_COLUMNS_PER_TARGET);
+}
+
+function normalizedSample(
+  target: DescriptionSourceSamplingTarget,
+  columnNames: readonly string[],
+  sourceRows: readonly Record<string, unknown>[],
+): DescriptionTargetSourceSample {
+  const rows = sourceRows.slice(0, MAX_SOURCE_ROWS).map((row) => ({
+    fields: columnNames.flatMap((name) => {
+      const value = normalizeValue(row[name]);
+      return value === undefined ? [] : [{ name, value }];
+    }),
+  }));
+  const valuesByColumn = new Map<string, Exclude<DescriptionSourceSampleValue, null>[]>();
+  const seenByColumn = new Map<string, Set<string>>();
+  let representativeValueCount = 0;
+  for (const row of rows) {
+    for (const field of row.fields) {
+      if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
+      if (field.value === null) continue;
+      const seen = seenByColumn.get(field.name) ?? new Set<string>();
+      const key = distinctKey(field.value);
+      if (seen.has(key)) continue;
+      seen.add(key);
+      seenByColumn.set(field.name, seen);
+      const values = valuesByColumn.get(field.name) ?? [];
+      values.push(field.value);
+      valuesByColumn.set(field.name, values);
+      representativeValueCount += 1;
+    }
+    if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
+  }
+  return {
+    targetId: target.targetId,
+    tableName: target.tableName,
+    rows,
+    representativeValues: columnNames.flatMap((column) => {
+      const values = valuesByColumn.get(column);
+      return values && values.length > 0 ? [{ column, values }] : [];
+    }),
+  };
+}
+
+function samplingSql(
+  database: WorkspaceDatabase,
+  target: DescriptionSourceSamplingTarget,
+  columns: readonly string[],
+): string {
+  const projections = columns.map((columnName) => {
+    const identifier = quoteIdentifier(columnName);
+    return `LEFT((${identifier})::text, ${MAX_SOURCE_VALUE_BYTES}) AS ${identifier}`;
+  });
+  return [
+    `SELECT ${projections.join(", ")}`,
+    `FROM ${quoteIdentifier(database.schema)}.${quoteIdentifier(target.tableName)}`,
+    `LIMIT ${MAX_SOURCE_ROWS}`,
+  ].join(" ");
+}
+
+/** Bounded source sampler that follows the database's PostgreSQL-wire or REST binding. */
+export class ConcreteDescriptionSourceSampler implements DescriptionSourceSampler {
+  constructor(
+    private readonly access: CatalogPostgresAccess,
+    private readonly secretStore?: Pick<WorkspaceSecretStore, "materialize">,
+  ) {}

  async sample(
    database: WorkspaceDatabase,
    targets: readonly DescriptionSourceSamplingTarget[],
    signal: AbortSignal,
  ): Promise<readonly DescriptionTargetSourceSample[]> {
+    if (database.binding.transport === "rest_api") {
+      return await this.sampleRest(database, targets, signal);
+    }
    const client = await this.access.connect(database, signal);
    let transactionOpen = false;
    try {
@@ -95,7 +165,7 @@ export class PostgresDescriptionSourceSampler implements DescriptionSourceSample
      transactionOpen = true;
      const samples: DescriptionTargetSourceSample[] = [];
      for (const target of targets) {
-        const columnNames = [...new Set(target.columnNames)].slice(0, MAX_SOURCE_COLUMNS_PER_TARGET);
+        const columnNames = columnsFor(target);
        if (columnNames.length === 0) {
          samples.push({
            targetId: target.targetId,
@@ -115,44 +185,7 @@ export class PostgresDescriptionSourceSampler implements DescriptionSourceSample
          "LIMIT $2",
        ].join(" ");
        const result = await client.query(sql, [MAX_SOURCE_VALUE_BYTES, MAX_SOURCE_ROWS]);
-        const rows = result.rows.slice(0, MAX_SOURCE_ROWS).map((row) => ({
-          fields: columnNames.flatMap((name) => {
-            const value = normalizeValue(row[name]);
-            return value === undefined ? [] : [{ name, value }];
-          }),
-        }));
-        const valuesByColumn = new Map<
-          string,
-          Exclude<DescriptionSourceSampleValue, null>[]
-        >();
-        const seenByColumn = new Map<string, Set<string>>();
-        let representativeValueCount = 0;
-        for (const row of rows) {
-          for (const field of row.fields) {
-            if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
-            if (field.value === null) continue;
-            const seen = seenByColumn.get(field.name) ?? new Set<string>();
-            const key = distinctKey(field.value);
-            if (seen.has(key)) continue;
-            seen.add(key);
-            seenByColumn.set(field.name, seen);
-            const values = valuesByColumn.get(field.name) ?? [];
-            values.push(field.value);
-            valuesByColumn.set(field.name, values);
-            representativeValueCount += 1;
-          }
-          if (representativeValueCount === MAX_REPRESENTATIVE_VALUES) break;
-        }
-        const representativeValues = columnNames.flatMap((column) => {
-          const values = valuesByColumn.get(column);
-          return values && values.length > 0 ? [{ column, values }] : [];
-        });
-        samples.push({
-          targetId: target.targetId,
-          tableName: target.tableName,
-          rows,
-          representativeValues,
-        });
+        samples.push(normalizedSample(target, columnNames, result.rows));
      }
      return samples;
    } finally {
@@ -160,4 +193,60 @@ export class PostgresDescriptionSourceSampler implements DescriptionSourceSample
      await client.end().catch(() => undefined);
    }
  }
+
+  private async sampleRest(
+    database: WorkspaceDatabase,
+    targets: readonly DescriptionSourceSamplingTarget[],
+    signal: AbortSignal,
+  ): Promise<readonly DescriptionTargetSourceSample[]> {
+    if (!this.secretStore) throw new CatalogConnectorError("REST source sampling is not configured");
+    const auth = database.binding.restAuth ?? "bearer";
+    const materialized = this.secretStore.materialize(
+      database.workspaceId,
+      auth === "none" ? [] : [CATALOG_SECRET_IDS.apiKey],
+    );
+    try {
+      const headers: Record<string, string> = { "content-type": "application/json" };
+      if (auth !== "none") {
+        const credentialFile = materialized.files.get(CATALOG_SECRET_IDS.apiKey);
+        if (!credentialFile) throw new CatalogConnectorError("REST API key is not configured");
+        const credential = (await readFile(credentialFile, "utf8")).trim();
+        if (auth === "bearer") headers.authorization = `Bearer ${credential}`;
+        else headers["x-api-key"] = credential;
+      }
+      const baseUrl = database.binding.baseUrl?.replace(/\/+$/, "");
+      if (!baseUrl) throw new CatalogConnectorError("Database binding is incomplete");
+      const samples: DescriptionTargetSourceSample[] = [];
+      for (const target of targets) {
+        const columnNames = columnsFor(target);
+        if (columnNames.length === 0) {
+          samples.push(normalizedSample(target, columnNames, []));
+          continue;
+        }
+        const response = await fetch(`${baseUrl}/rpc/run_query`, {
+          method: "POST",
+          headers,
+          body: JSON.stringify({ query_text: samplingSql(database, target, columnNames) }),
+          signal,
+        });
+        if (!response.ok) throw new CatalogConnectorError("REST source sampling failed");
+        const body: unknown = await response.json();
+        if (!Array.isArray(body)
+          || body.some((row) => !row || typeof row !== "object" || Array.isArray(row))) {
+          throw new CatalogConnectorError("REST source sampling response is invalid");
+        }
+        samples.push(normalizedSample(
+          target,
+          columnNames,
+          body as Array<Record<string, unknown>>,
+        ));
+      }
+      return samples;
+    } catch (error) {
+      if (error instanceof CatalogConnectorError) throw error;
+      throw new CatalogConnectorError("REST source sampling failed");
+    } finally {
+      materialized.release();
+    }
+  }
 }
@@ -335,7 +335,10 @@ export class MemoryCatalogRepository implements CatalogRepository {
      processed: 0,
      generated: 0,
      nonGeneratable: 0,
-      failed: 0,
+        failed: 0,
+        inputTokens: 0,
+        cacheReadTokens: 0,
+        outputTokens: 0,
      createdAt: now,
      startedAt: null,
      updatedAt: now,
@@ -444,7 +447,10 @@ export class MemoryCatalogRepository implements CatalogRepository {
      status: "running",
      total: 0,
      suggestedSensitive: 0,
-      suggestedNonSensitive: 0,
+        suggestedNonSensitive: 0,
+        inputTokens: 0,
+        cacheReadTokens: 0,
+        outputTokens: 0,
      createdAt: now,
      startedAt: now,
      updatedAt: now,
@@ -11,6 +11,7 @@ import * as descriptionGenerationRunsMigration from "./migrations/005_descriptio
 import * as sensitiveDataFlagMigration from "./migrations/006_sensitive_data_flag.js";
 import * as sensitiveDataSuggestionRunsMigration from "./migrations/007_sensitive_data_suggestion_runs.js";
 import * as catalogLogicalRelationshipsMigration from "./migrations/008_catalog_logical_relationships.js";
+import * as aiTokenUsageMigration from "./migrations/009_ai_token_usage.js";

 const connectionString = process.env.THT_CATALOG_MIGRATOR_DATABASE_URL;
 const host = process.env.THT_CATALOG_DB_HOST;
@@ -44,6 +45,7 @@ const provider: MigrationProvider = {
      "006_sensitive_data_flag": sensitiveDataFlagMigration,
      "007_sensitive_data_suggestion_runs": sensitiveDataSuggestionRunsMigration,
      "008_catalog_logical_relationships": catalogLogicalRelationshipsMigration,
+      "009_ai_token_usage": aiTokenUsageMigration,
    };
  },
 };
@@ -0,0 +1,20 @@
+import type { Kysely } from "kysely";
+import type { CatalogDatabase } from "../repository.js";
+
+export async function up(db: Kysely<CatalogDatabase>): Promise<void> {
+  for (const table of ["description_generation_runs", "sensitive_data_suggestion_runs"] as const) {
+    await db.schema.alterTable(table)
+      .addColumn("input_tokens", "integer", (col) => col.notNull().defaultTo(0))
+      .addColumn("cache_read_tokens", "integer", (col) => col.notNull().defaultTo(0))
+      .addColumn("output_tokens", "integer", (col) => col.notNull().defaultTo(0))
+      .execute();
+  }
+}
+
+export async function down(db: Kysely<CatalogDatabase>): Promise<void> {
+  for (const table of ["sensitive_data_suggestion_runs", "description_generation_runs"] as const) {
+    await db.schema.alterTable(table)
+      .dropColumn("input_tokens").dropColumn("cache_read_tokens").dropColumn("output_tokens")
+      .execute();
+  }
+}
@@ -4,7 +4,7 @@ import type { ResolvedMetadataGenerationModel } from "./metadata-generation-mode

 const MAX_HELPER_OUTPUT_BYTES = 64 * 1024;
 const helperOutputSchema = z.discriminatedUnion("ok", [
-  z.object({ ok: z.literal(true), content: z.string() }).strict(),
+  z.object({ ok: z.literal(true), content: z.string(), usage: z.object({ input: z.number().int().nonnegative(), cacheRead: z.number().int().nonnegative(), output: z.number().int().nonnegative() }).strict().optional() }).strict(),
  z.object({ ok: z.literal(false), error: z.literal("provider_failure") }).strict(),
 ]);

@@ -18,10 +18,12 @@ export interface ModelCompletionRequest {
  messages: readonly ModelCompletionMessage[];
  signal: AbortSignal;
 }
+export interface ModelCompletionUsage { input: number; cacheRead: number; output: number; }
+export interface ModelCompletionResult { content: string; usage: ModelCompletionUsage; }

 /** The provider boundary used by Description Generation. */
 export interface ModelCompleter {
-  complete(request: ModelCompletionRequest): Promise<string>;
+  complete(request: ModelCompletionRequest): Promise<string | ModelCompletionResult>;
 }

 export class ModelCompletionProviderError extends Error {
@@ -47,7 +49,7 @@ export class PythonModelCompleter implements ModelCompleter {
    terminationGraceMs?: number;
  }) {}

-  async complete(request: ModelCompletionRequest): Promise<string> {
+  async complete(request: ModelCompletionRequest): Promise<ModelCompletionResult> {
    if (request.signal.aborted) throw new ModelCompletionCancelledError();
    const payload = {
      model: `${request.model.provider}/${request.model.model}`,
@@ -62,7 +64,7 @@ export class PythonModelCompleter implements ModelCompleter {
      ...(request.model.disableThinking === true ? { disable_thinking: true } : {}),
    };

-    return await new Promise<string>((resolve, reject) => {
+    return await new Promise<ModelCompletionResult>((resolve, reject) => {
      const child = spawn(
        this.options.pythonExecutable,
        ["-m", this.options.helperModule ?? "tht.internal.litellm_completion"],
@@ -135,7 +137,7 @@ export class PythonModelCompleter implements ModelCompleter {
          if (!output.ok) return fail();
          settled = true;
          cleanup();
-          resolve(output.content);
+          resolve({ content: output.content, usage: output.usage ?? { input: 0, cacheRead: 0, output: 0 } });
        } catch {
          fail();
        }
@@ -171,6 +171,9 @@ interface DescriptionGenerationRunTable {
  generated: number;
  nonGeneratable: number;
  failed: number;
+  inputTokens: number;
+  cacheReadTokens: number;
+  outputTokens: number;
  createdAt: Timestamp;
  startedAt: Timestamp | null;
  updatedAt: Timestamp;
@@ -195,6 +198,9 @@ interface SensitiveDataSuggestionRunTable {
  total: number;
  suggestedSensitive: number;
  suggestedNonSensitive: number;
+  inputTokens: number;
+  cacheReadTokens: number;
+  outputTokens: number;
  createdAt: Timestamp;
  startedAt: Timestamp;
  updatedAt: Timestamp;
@@ -816,6 +822,9 @@ export class KyselyCatalogRepository implements CatalogRepository {
        generated: 0,
        nonGeneratable: 0,
        failed: 0,
+        inputTokens: 0,
+        cacheReadTokens: 0,
+        outputTokens: 0,
        startedAt: null,
        finishedAt: null,
        errorSummary: null,
@@ -943,6 +952,9 @@ export class KyselyCatalogRepository implements CatalogRepository {
      total: 0,
      suggestedSensitive: 0,
      suggestedNonSensitive: 0,
+      inputTokens: 0,
+      cacheReadTokens: 0,
+      outputTokens: 0,
      finishedAt: null,
      errorSummary: null,
    }).returningAll().executeTakeFirstOrThrow();
@@ -1,6 +1,6 @@
 import { z } from "zod";
 import type { MetadataGenerationModels } from "./metadata-generation-models.js";
-import type { ModelCompleter, ModelCompletionMessage } from "./model-completer.js";
+import type { ModelCompleter, ModelCompletionMessage, ModelCompletionResult, ModelCompletionUsage } from "./model-completer.js";
 import type {
  CatalogColumn,
  CatalogRepository,
@@ -201,6 +201,8 @@ export class SensitiveDataSuggester {
    targetIds: readonly string[],
    signal: AbortSignal,
    onPrepared?: (total: number) => void | Promise<void>,
+    onProgress?: (processed: number, suggestions: readonly SensitiveDataSuggestion[]) => void | Promise<void>,
+    onUsage?: (usage: ModelCompletionUsage) => void | Promise<void>,
  ): Promise<readonly SensitiveDataSuggestion[]> {
    const database = await this.repository.get(databaseId);
    if (!database) throw new SensitiveDataSuggestionTargetNotFoundError("database");
@@ -212,11 +214,16 @@ export class SensitiveDataSuggester {
    for (const batch of batchesFor(database, columns)) {
      let received: Map<string, { columnId: string; sensitive: boolean }> | undefined;
      for (let attempt = 0; attempt < 2 && !received; attempt += 1) {
-        const content = await this.completer.complete({
+        const completion = await this.completer.complete({
          model,
          signal,
          messages: [systemMessage, { role: "user", content: userContent(database, batch) }],
        });
+        const result: ModelCompletionResult = typeof completion === "string"
+          ? { content: completion, usage: { input: 0, cacheRead: 0, output: 0 } }
+          : completion;
+        await onUsage?.(result.usage);
+        const content = result.content;
        try {
          const parsed = responseSchema.parse(JSON.parse(content));
          const expected = new Set(batch.map((column) => column.columnId));
@@ -240,6 +247,7 @@ export class SensitiveDataSuggester {
        currentSensitive: column.currentSensitive,
        sensitive: received!.get(column.columnId)!.sensitive,
      })));
+      await onProgress?.(suggestions.length, suggestions.slice(-batch.length));
    }
    return suggestions;
  }
@@ -10,6 +10,7 @@ import type {
  SensitiveDataSuggestionRun,
  SensitiveDataSuggestionScope,
 } from "./types.js";
+import type { ModelCompletionUsage } from "./model-completer.js";

 const interruptedMessage = "Sensitive-field suggestion generation was interrupted by backend restart.";
 const failedMessage = "Sensitive-field suggestion generation failed.";
@@ -73,6 +74,31 @@ export class SensitiveDataSuggestionRunner {
          });
          if (!prepared) throw new Error("Sensitive Data Suggestion Run disappeared");
        },
+        async (processed, batch) => {
+          const suggestedSensitive = batch.filter((suggestion) => suggestion.sensitive).length;
+          const suggestedNonSensitive = batch.length - suggestedSensitive;
+          const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
+          if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
+          const progress = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
+            suggestedSensitive: current.suggestedSensitive + suggestedSensitive,
+            suggestedNonSensitive: current.suggestedNonSensitive + suggestedNonSensitive,
+          });
+          if (!progress) throw new Error("Sensitive Data Suggestion Run disappeared");
+          await this.repository.appendSensitiveDataSuggestionEvent(
+            started.id,
+            "info",
+            `Classified ${processed} of ${progress.total} columns.`,
+          );
+        },
+        async (usage: ModelCompletionUsage) => {
+          const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
+          if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
+          await this.repository.updateSensitiveDataSuggestionRun(started.id, {
+            inputTokens: current.inputTokens + usage.input,
+            cacheReadTokens: current.cacheReadTokens + usage.cacheRead,
+            outputTokens: current.outputTokens + usage.output,
+          });
+        },
      );
      const suggestedSensitive = suggestions.filter((suggestion) => suggestion.sensitive).length;
      const suggestedNonSensitive = suggestions.length - suggestedSensitive;
@@ -234,6 +234,9 @@ export interface DescriptionGenerationRun {
  generated: number;
  nonGeneratable: number;
  failed: number;
+  inputTokens: number;
+  cacheReadTokens: number;
+  outputTokens: number;
  createdAt: string;
  startedAt: string | null;
  updatedAt: string;
@@ -250,6 +253,9 @@ export interface DescriptionGenerationRunUpdate {
  startedAt?: string | null;
  finishedAt?: string | null;
  errorSummary?: string | null;
+  inputTokens?: number;
+  cacheReadTokens?: number;
+  outputTokens?: number;
 }

 export interface DescriptionGenerationEvent {
@@ -272,6 +278,9 @@ export interface SensitiveDataSuggestionRun {
  total: number;
  suggestedSensitive: number;
  suggestedNonSensitive: number;
+  inputTokens: number;
+  cacheReadTokens: number;
+  outputTokens: number;
  createdAt: string;
  startedAt: string;
  updatedAt: string;
@@ -286,6 +295,9 @@ export interface SensitiveDataSuggestionRunUpdate {
  suggestedNonSensitive?: number;
  finishedAt?: string | null;
  errorSummary?: string | null;
+  inputTokens?: number;
+  cacheReadTokens?: number;
+  outputTokens?: number;
 }

 export interface SensitiveDataSuggestionEvent {
@@ -101,6 +101,9 @@ function publicRun(run: DescriptionGenerationRun) {
    generated: run.generated,
    nonGeneratable: run.nonGeneratable,
    failed: run.failed,
+    inputTokens: run.inputTokens,
+    cacheReadTokens: run.cacheReadTokens,
+    outputTokens: run.outputTokens,
    createdAt: run.createdAt,
    startedAt: run.startedAt,
    updatedAt: run.updatedAt,
@@ -129,6 +132,9 @@ function publicSensitiveDataSuggestionRun(run: SensitiveDataSuggestionRun) {
    total: run.total,
    suggestedSensitive: run.suggestedSensitive,
    suggestedNonSensitive: run.suggestedNonSensitive,
+    inputTokens: run.inputTokens,
+    cacheReadTokens: run.cacheReadTokens,
+    outputTokens: run.outputTokens,
    createdAt: run.createdAt,
    startedAt: run.startedAt,
    updatedAt: run.updatedAt,