fix: scope and batch sensitive suggestions

backend/src/catalog/sensitive-data-suggester.ts

@@ -1,28 +1,74 @@
 import { z } from "zod";
 import type { MetadataGenerationModels } from "./metadata-generation-models.js";
-import type { ModelCompleter } from "./model-completer.js";
-import type { CatalogRepository } from "./types.js";
+import type { ModelCompleter, ModelCompletionMessage } from "./model-completer.js";
+import type { CatalogColumn, CatalogRepository, CatalogTable } from "./types.js";
 
-const MAX_COLUMNS = 10_000;
+// The helper accepts at most 64 KiB per message. Keep the same safety margin used by
+// Description Generation so UTF-8 structural metadata never reaches that hard limit.
+const MAX_USER_MESSAGE_BYTES = 60 * 1024;
+// Preserve ThothAI's proven completion granularity: small batches keep generation time and
+// structured-output accuracy predictable even when the helper byte limit would allow much more.
+const MAX_COLUMNS_PER_BATCH = 10;
 const responseSchema = z.object({
   suggestions: z.array(z.object({
     columnId: z.uuid(),
     sensitive: z.boolean(),
-  }).strict()).max(MAX_COLUMNS),
+  }).strict()),
 }).strict();
 
+export type SensitiveDataSuggestionScope = "all" | "selected_tables" | "selected_columns";
+
+interface StructuralColumn {
+  columnId: string;
+  tableId: string;
+  table: string;
+  column: string;
+  dataType: string;
+  nullable: boolean;
+  primaryKey: boolean;
+  foreignKey: boolean;
+  version: number;
+  currentSensitive: boolean;
+}
+
 export interface SensitiveDataSuggestion {
   columnId: string;
+  tableId: string;
+  tableName: string;
+  columnName: string;
+  version: number;
+  currentSensitive: boolean;
   sensitive: boolean;
 }
 
 export class SensitiveDataSuggestionTargetNotFoundError extends Error {
-  constructor() {
-    super("database not found");
+  constructor(readonly target: "database" | "table" | "column") {
+    super(`${target} not found`);
     this.name = "SensitiveDataSuggestionTargetNotFoundError";
   }
 }
 
+export class SensitiveDataSuggestionDuplicateTargetIdsError extends Error {
+  constructor() {
+    super("sensitive-data suggestion target IDs must be unique");
+    this.name = "SensitiveDataSuggestionDuplicateTargetIdsError";
+  }
+}
+
+export class SensitiveDataSuggestionNoEligibleColumnsError extends Error {
+  constructor(readonly scope: SensitiveDataSuggestionScope) {
+    super("selected scope has no catalog columns");
+    this.name = "SensitiveDataSuggestionNoEligibleColumnsError";
+  }
+}
+
+export class SensitiveDataSuggestionPayloadTooLargeError extends Error {
+  constructor() {
+    super("sensitive-data suggestion structural metadata is too large");
+    this.name = "SensitiveDataSuggestionPayloadTooLargeError";
+  }
+}
+
 export class SensitiveDataSuggestionInvalidResponseError extends Error {
   constructor() {
     super("sensitive-data suggestion response is invalid");
@@ -30,6 +76,77 @@ export class SensitiveDataSuggestionInvalidResponseError extends Error {
   }
 }
 
+function userContent(
+  database: { databaseName: string; schema: string },
+  columns: readonly StructuralColumn[],
+): string {
+  return JSON.stringify({
+    database: database.databaseName,
+    schema: database.schema,
+    columns: columns.map((column) => ({
+      columnId: column.columnId,
+      table: column.table,
+      column: column.column,
+      dataType: column.dataType,
+      nullable: column.nullable,
+      primaryKey: column.primaryKey,
+      foreignKey: column.foreignKey,
+    })),
+  });
+}
+
+function batchesFor(
+  database: { databaseName: string; schema: string },
+  columns: readonly StructuralColumn[],
+): StructuralColumn[][] {
+  const batches: StructuralColumn[][] = [];
+  let current: StructuralColumn[] = [];
+  for (const column of columns) {
+    if (current.length === MAX_COLUMNS_PER_BATCH) {
+      batches.push(current);
+      current = [];
+    }
+    const candidate = [...current, column];
+    if (Buffer.byteLength(userContent(database, candidate), "utf8") <= MAX_USER_MESSAGE_BYTES) {
+      current = candidate;
+      continue;
+    }
+    if (current.length === 0) throw new SensitiveDataSuggestionPayloadTooLargeError();
+    batches.push(current);
+    current = [column];
+    if (Buffer.byteLength(userContent(database, current), "utf8") > MAX_USER_MESSAGE_BYTES) {
+      throw new SensitiveDataSuggestionPayloadTooLargeError();
+    }
+  }
+  if (current.length > 0) batches.push(current);
+  return batches;
+}
+
+function structuralColumn(table: CatalogTable, column: CatalogColumn): StructuralColumn {
+  return {
+    columnId: column.id,
+    tableId: table.id,
+    table: table.name,
+    column: column.name,
+    dataType: column.dataType,
+    nullable: column.isNullable,
+    primaryKey: column.isPrimaryKey,
+    foreignKey: column.isForeignKey,
+    version: column.version,
+    currentSensitive: column.sensitive,
+  };
+}
+
+const systemMessage: ModelCompletionMessage = {
+  role: "system",
+  content: [
+    "Classify whether each database column is likely to contain sensitive source values.",
+    "Use only the supplied structural metadata. Return strict JSON with this exact shape:",
+    '{"suggestions":[{"columnId":"uuid","sensitive":true}]}',
+    "Return every supplied column exactly once. Do not add explanations or markdown.",
+  ].join("\n"),
+};
+
 export class SensitiveDataSuggester {
   constructor(
     private readonly repository: CatalogRepository,
@@ -37,67 +154,86 @@ export class SensitiveDataSuggester {
     private readonly completer: ModelCompleter,
   ) {}
 
+  private async selectColumns(
+    databaseId: string,
+    scope: SensitiveDataSuggestionScope,
+    targetIds: readonly string[],
+  ): Promise<StructuralColumn[]> {
+    if (new Set(targetIds).size !== targetIds.length) {
+      throw new SensitiveDataSuggestionDuplicateTargetIdsError();
+    }
+    const tables = await this.repository.listTables(databaseId);
+    const tableIds = new Set(targetIds);
+    const selectedTables = scope === "selected_tables"
+      ? tables.filter((table) => tableIds.has(table.id))
+      : tables;
+    if (scope === "selected_tables" && selectedTables.length !== targetIds.length) {
+      throw new SensitiveDataSuggestionTargetNotFoundError("table");
+    }
+
+    const columns = (await Promise.all(selectedTables.map(async (table) => (
+      (await this.repository.listColumns(databaseId, table.id)).map((column) => (
+        structuralColumn(table, column)
+      ))
+    )))).flat();
+    const columnIds = new Set(targetIds);
+    const selectedColumns = scope === "selected_columns"
+      ? columns.filter((column) => columnIds.has(column.columnId))
+      : columns;
+    if (scope === "selected_columns" && selectedColumns.length !== targetIds.length) {
+      throw new SensitiveDataSuggestionTargetNotFoundError("column");
+    }
+    if (selectedColumns.length === 0) {
+      throw new SensitiveDataSuggestionNoEligibleColumnsError(scope);
+    }
+    return selectedColumns;
+  }
+
   async suggest(
     databaseId: string,
     modelId: string,
+    scope: SensitiveDataSuggestionScope,
+    targetIds: readonly string[],
     signal: AbortSignal,
   ): Promise<readonly SensitiveDataSuggestion[]> {
     const database = await this.repository.get(databaseId);
-    if (!database) throw new SensitiveDataSuggestionTargetNotFoundError();
+    if (!database) throw new SensitiveDataSuggestionTargetNotFoundError("database");
+    const columns = await this.selectColumns(databaseId, scope, targetIds);
+    const model = this.models.resolve(modelId);
+    const suggestions: SensitiveDataSuggestion[] = [];
 
-    const tables = await this.repository.listTables(databaseId);
-    const columns = (await Promise.all(tables.map(async (table) => ({
-      table,
-      columns: await this.repository.listColumns(databaseId, table.id),
-    })))).flatMap(({ table, columns: tableColumns }) => tableColumns.map((column) => ({
-      columnId: column.id,
-      table: table.name,
-      column: column.name,
-      dataType: column.dataType,
-      nullable: column.isNullable,
-      primaryKey: column.isPrimaryKey,
-      foreignKey: column.isForeignKey,
-    })));
-    if (columns.length === 0) return [];
-    if (columns.length > MAX_COLUMNS) throw new SensitiveDataSuggestionInvalidResponseError();
-
-    const content = await this.completer.complete({
-      model: this.models.resolve(modelId),
-      signal,
-      messages: [
-        {
-          role: "system",
-          content: [
-            "Classify whether each database column is likely to contain sensitive source values.",
-            "Use only the supplied structural metadata. Return strict JSON with this exact shape:",
-            '{"suggestions":[{"columnId":"uuid","sensitive":true}]}',
-            "Return every supplied column exactly once. Do not add explanations or markdown.",
-          ].join("\n"),
-        },
-        {
-          role: "user",
-          content: JSON.stringify({
-            database: database.databaseName,
-            schema: database.schema,
-            columns,
-          }),
-        },
-      ],
-    });
-
-    try {
-      const parsed = responseSchema.parse(JSON.parse(content));
-      const expected = new Set(columns.map((column) => column.columnId));
-      const received = new Set(parsed.suggestions.map((suggestion) => suggestion.columnId));
-      if (received.size !== parsed.suggestions.length
-        || received.size !== expected.size
-        || [...received].some((columnId) => !expected.has(columnId))) {
-        throw new SensitiveDataSuggestionInvalidResponseError();
+    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({
+          model,
+          signal,
+          messages: [systemMessage, { role: "user", content: userContent(database, batch) }],
+        });
+        try {
+          const parsed = responseSchema.parse(JSON.parse(content));
+          const expected = new Set(batch.map((column) => column.columnId));
+          const candidate = new Map(parsed.suggestions.map((suggestion) => [suggestion.columnId, suggestion]));
+          if (candidate.size !== parsed.suggestions.length
+            || candidate.size !== expected.size
+            || [...candidate.keys()].some((columnId) => !expected.has(columnId))) {
+            throw new SensitiveDataSuggestionInvalidResponseError();
+          }
+          received = candidate;
+        } catch {
+          if (attempt === 1) throw new SensitiveDataSuggestionInvalidResponseError();
+        }
       }
-      return parsed.suggestions;
-    } catch (error) {
-      if (error instanceof SensitiveDataSuggestionInvalidResponseError) throw error;
-      throw new SensitiveDataSuggestionInvalidResponseError();
+      suggestions.push(...batch.map((column) => ({
+        columnId: column.columnId,
+        tableId: column.tableId,
+        tableName: column.table,
+        columnName: column.column,
+        version: column.version,
+        currentSensitive: column.currentSensitive,
+        sensitive: received!.get(column.columnId)!.sensitive,
+      })));
     }
+    return suggestions;
   }
 }

backend/src/routes/catalog-description-generation.ts

@@ -14,7 +14,10 @@ import { MetadataGenerationModelUnavailableError } from "../catalog/metadata-gen
 import { ModelCompletionProviderError } from "../catalog/model-completer.js";
 import {
   SensitiveDataSuggester,
+  SensitiveDataSuggestionDuplicateTargetIdsError,
   SensitiveDataSuggestionInvalidResponseError,
+  SensitiveDataSuggestionNoEligibleColumnsError,
+  SensitiveDataSuggestionPayloadTooLargeError,
   SensitiveDataSuggestionTargetNotFoundError,
 } from "../catalog/sensitive-data-suggester.js";
 import {
@@ -28,8 +31,20 @@ import {
 
 const idSchema = z.uuid();
 const modelIdSchema = z.string().regex(/^[a-z][a-z0-9._-]{0,63}$/);
-const suggestionSchema = z.object({ modelId: modelIdSchema }).strict();
 const selectedTargetIdsSchema = z.array(idSchema).min(1);
+const suggestionSchema = z.discriminatedUnion("scope", [
+  z.object({ modelId: modelIdSchema, scope: z.literal("all") }).strict(),
+  z.object({
+    modelId: modelIdSchema,
+    scope: z.literal("selected_tables"),
+    targetIds: selectedTargetIdsSchema,
+  }).strict(),
+  z.object({
+    modelId: modelIdSchema,
+    scope: z.literal("selected_columns"),
+    targetIds: selectedTargetIdsSchema,
+  }).strict(),
+]);
 const startSchema = z.discriminatedUnion("scope", [
   z.object({
     modelId: modelIdSchema,
@@ -123,19 +138,6 @@ function safeError(reply: FastifyReply, error: unknown) {
       message: "The selected metadata-generation model is unavailable.",
     });
   }
-  if (error instanceof SensitiveDataSuggestionTargetNotFoundError) {
-    return reply.code(404).send({
-      code: "database_not_found",
-      message: "Database configuration was not found.",
-    });
-  }
-  if (error instanceof SensitiveDataSuggestionInvalidResponseError
-    || error instanceof ModelCompletionProviderError) {
-    return reply.code(502).send({
-      code: "sensitive_data_suggestion_failed",
-      message: "Sensitive-data suggestions could not be prepared.",
-    });
-  }
   if (error instanceof DescriptionGenerationDuplicateTargetIdsError) {
     return reply.code(400).send({
       code: "description_generation_target_ids_duplicate",
@@ -190,6 +192,74 @@ function safeError(reply: FastifyReply, error: unknown) {
   });
 }
 
+function safeSuggestionError(reply: FastifyReply, error: unknown) {
+  if (error instanceof CatalogUnavailableError) {
+    return reply.code(503).send({
+      code: "catalog_unavailable",
+      message: "The database catalog is unavailable, so no sensitive-field suggestions were prepared.",
+    });
+  }
+  if (error instanceof MetadataGenerationModelUnavailableError) {
+    return reply.code(409).send({
+      code: "metadata_generation_model_unavailable",
+      message: "The selected metadata-generation model is unavailable.",
+    });
+  }
+  if (error instanceof SensitiveDataSuggestionTargetNotFoundError) {
+    const code = error.target === "database"
+      ? "database_not_found"
+      : error.target === "table"
+        ? "catalog_table_not_found"
+        : "catalog_column_not_found";
+    const message = error.target === "database"
+      ? "The database configuration was not found."
+      : error.target === "table"
+        ? "One or more selected Catalog Tables were not found in this database."
+        : "One or more selected Catalog Columns were not found in this database.";
+    return reply.code(404).send({ code, message });
+  }
+  if (error instanceof SensitiveDataSuggestionDuplicateTargetIdsError) {
+    return reply.code(400).send({
+      code: "sensitive_data_suggestion_target_ids_duplicate",
+      message: "Each selected table or column must appear only once.",
+    });
+  }
+  if (error instanceof SensitiveDataSuggestionNoEligibleColumnsError) {
+    return reply.code(409).send({
+      code: "sensitive_data_suggestion_no_columns",
+      message: "The selected scope contains no Catalog Columns to classify.",
+    });
+  }
+  if (error instanceof SensitiveDataSuggestionPayloadTooLargeError) {
+    return reply.code(413).send({
+      code: "sensitive_data_suggestion_payload_too_large",
+      message: "The selected structural metadata cannot be divided into safe LLM requests.",
+    });
+  }
+  if (error instanceof SensitiveDataSuggestionInvalidResponseError) {
+    return reply.code(502).send({
+      code: "sensitive_data_suggestion_invalid_response",
+      message: "The LLM returned an incomplete or invalid classification. No suggestions were applied.",
+    });
+  }
+  if (error instanceof ModelCompletionProviderError) {
+    return reply.code(502).send({
+      code: "sensitive_data_suggestion_provider_unavailable",
+      message: "The selected LLM service could not complete the request. No suggestions were applied.",
+    });
+  }
+  if (error instanceof z.ZodError) {
+    return reply.code(400).send({
+      code: "sensitive_data_suggestion_request_invalid",
+      message: "Choose a database, one or more tables, or one or more columns to classify.",
+    });
+  }
+  return reply.code(500).send({
+    code: "sensitive_data_suggestion_failed",
+    message: "Sensitive-field suggestions failed before review. No changes were applied.",
+  });
+}
+
 export function catalogDescriptionGenerationRoutes(
   app: FastifyInstance,
   deps: {
@@ -206,11 +276,13 @@ export function catalogDescriptionGenerationRoutes(
       const suggestions = await deps.sensitiveDataSuggester.suggest(
         databaseId,
         input.modelId,
+        input.scope,
+        "targetIds" in input ? input.targetIds : [],
         new AbortController().signal,
       );
       return { suggestions };
     } catch (error) {
-      return safeError(reply, error);
+      return safeSuggestionError(reply, error);
     }
   });
 

backend/src/routes/catalog-schema.ts

@@ -16,10 +16,12 @@ import {
 const idSchema = z.uuid();
 const metadataSchema = z.object({
   version: z.number().int().positive(),
-  description: z.string().max(20_000).nullable(),
-  generatedDescription: z.string().max(20_000).nullable(),
+  description: z.string().max(20_000).nullable().optional(),
+  generatedDescription: z.string().max(20_000).nullable().optional(),
   sensitive: z.boolean().optional(),
-}).strict();
+}).strict().refine((value) => (
+  "description" in value || "generatedDescription" in value || "sensitive" in value
+));
 const createRunSchema = z.object({
   version: z.number().int().positive(),
   scope: z.enum(["tables", "columns", "relationships", "all"]),
@@ -116,8 +118,10 @@ export function catalogSchemaRoutes(
           tableId,
           columnId,
           input.version,
-          normalized(input.description),
-          normalized(input.generatedDescription),
+          "description" in input ? normalized(input.description ?? null) : current.description,
+          "generatedDescription" in input
+            ? normalized(input.generatedDescription ?? null)
+            : current.generatedDescription,
           input.sensitive,
         );
         if (!updated) return reply.code(409).send({ code: "column_stale", message: "Column metadata changed. Reload and try again." });

backend/test/catalog-description-generation-routes.test.ts

@@ -147,7 +147,7 @@ test("suggests sensitive flags from structural metadata without persisting them"
       suggestions: [{ columnId: expect.any(String), sensitive: true }],
     })),
   };
-  const { app, repository, database, column } = await setup(modelCompleter);
+  const { app, repository, database, table, column } = await setup(modelCompleter);
   modelCompleter.complete.mockResolvedValueOnce(JSON.stringify({
     suggestions: [{ columnId: column.id, sensitive: true }],
   }));
@@ -156,12 +156,20 @@ test("suggests sensitive flags from structural metadata without persisting them"
     const response = await app.inject({
       method: "POST",
       url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
-      payload: { modelId: configuredModel.id },
+      payload: { modelId: configuredModel.id, scope: "all" },
     });
 
     expect(response.statusCode).toBe(200);
     expect(response.json()).toEqual({
-      suggestions: [{ columnId: column.id, sensitive: true }],
+      suggestions: [{
+        columnId: column.id,
+        tableId: table.id,
+        tableName: table.name,
+        columnName: column.name,
+        version: column.version,
+        currentSensitive: false,
+        sensitive: true,
+      }],
     });
     expect(await repository.getColumn(database.id, column.tableId, column.id))
       .toMatchObject({ sensitive: false });
@@ -178,6 +186,262 @@ test("suggests sensitive flags from structural metadata without persisting them"
   }
 });
 
+test("limits sensitive-data suggestions to the selected tables or columns", async () => {
+  const modelCompleter: ModelCompleter = {
+    complete: vi.fn(async (request) => {
+      const payload = JSON.parse(request.messages.find((message) => message.role === "user")!.content) as {
+        columns: Array<{ columnId: string; column: string }>;
+      };
+      return JSON.stringify({
+        suggestions: payload.columns.map((column) => ({
+          columnId: column.columnId,
+          sensitive: column.column.includes("name") || column.column.includes("note"),
+        })),
+      });
+    }),
+  };
+  const { app, repository, database } = await setup(modelCompleter);
+  await repository.applySchemaSync(database.id, database.version, "all", [], {
+    schemaVersion: 1,
+    capabilities: { tables: "available", columns: "available", relationships: "available" },
+    tables: [
+      { name: "patients", sourceComment: null },
+      { name: "visits", sourceComment: null },
+      { name: "billing", sourceComment: null },
+    ],
+    columns: [
+      { tableName: "patients", name: "patient_name", ordinalPosition: 1, dataType: "text", isNullable: false, defaultExpression: null, primaryKeyPosition: null, sourceComment: null },
+      { tableName: "patients", name: "status", ordinalPosition: 2, dataType: "text", isNullable: true, defaultExpression: null, primaryKeyPosition: null, sourceComment: null },
+      { tableName: "visits", name: "clinical_note", ordinalPosition: 1, dataType: "text", isNullable: true, defaultExpression: null, primaryKeyPosition: null, sourceComment: null },
+      { tableName: "billing", name: "invoice_total", ordinalPosition: 1, dataType: "numeric", isNullable: false, defaultExpression: null, primaryKeyPosition: null, sourceComment: null },
+    ],
+    relationships: [],
+  });
+  const tables = await repository.listTables(database.id);
+  const patients = tables.find((table) => table.name === "patients")!;
+  const visits = tables.find((table) => table.name === "visits")!;
+  const billing = tables.find((table) => table.name === "billing")!;
+  const patientColumns = await repository.listColumns(database.id, patients.id);
+  const visitColumns = await repository.listColumns(database.id, visits.id);
+  const billingColumns = await repository.listColumns(database.id, billing.id);
+  const status = patientColumns.find((column) => column.name === "status")!;
+  const clinicalNote = visitColumns.find((column) => column.name === "clinical_note")!;
+
+  try {
+    const tableResponse = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: {
+        modelId: configuredModel.id,
+        scope: "selected_tables",
+        targetIds: [visits.id, patients.id],
+      },
+    });
+    expect(tableResponse.statusCode).toBe(200);
+    expect(tableResponse.json().suggestions).toHaveLength(3);
+    expect(tableResponse.json().suggestions).toEqual(expect.arrayContaining([
+      expect.objectContaining({ tableId: patients.id, columnName: "patient_name", sensitive: true }),
+      expect.objectContaining({ tableId: patients.id, columnName: "status", sensitive: false }),
+      expect.objectContaining({ tableId: visits.id, columnName: "clinical_note", sensitive: true }),
+    ]));
+    expect(tableResponse.json().suggestions).not.toEqual(expect.arrayContaining([
+      expect.objectContaining({ tableId: billing.id }),
+    ]));
+
+    const columnResponse = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: {
+        modelId: configuredModel.id,
+        scope: "selected_columns",
+        targetIds: [clinicalNote.id, status.id],
+      },
+    });
+    expect(columnResponse.statusCode).toBe(200);
+    expect(columnResponse.json().suggestions).toHaveLength(2);
+    expect(columnResponse.json().suggestions).toEqual(expect.arrayContaining([
+      expect.objectContaining({ tableId: patients.id, columnId: status.id, sensitive: false }),
+      expect.objectContaining({ tableId: visits.id, columnId: clinicalNote.id, sensitive: true }),
+    ]));
+
+    const prompts = vi.mocked(modelCompleter.complete).mock.calls.map(([request]) => (
+      JSON.parse(request.messages.find((message) => message.role === "user")!.content) as {
+        columns: Array<{ columnId: string }>;
+      }
+    ));
+    expect(prompts[0]!.columns.map((column) => column.columnId).sort()).toEqual(
+      [...patientColumns, ...visitColumns].map((column) => column.id).sort(),
+    );
+    expect(prompts[0]!.columns.map((column) => column.columnId)).not.toContain(billingColumns[0]!.id);
+    expect(prompts[1]!.columns.map((column) => column.columnId).sort()).toEqual(
+      [status.id, clinicalNote.id].sort(),
+    );
+  } finally {
+    await app.close();
+  }
+});
+
+test("explains invalid sensitive-data suggestion selections without calling the model", async () => {
+  const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") };
+  const { app, database, table } = await setup(modelCompleter);
+
+  try {
+    const empty = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: { modelId: configuredModel.id, scope: "selected_tables", targetIds: [] },
+    });
+    expect(empty.statusCode).toBe(400);
+    expect(empty.json()).toEqual({
+      code: "sensitive_data_suggestion_request_invalid",
+      message: "Choose a database, one or more tables, or one or more columns to classify.",
+    });
+
+    const duplicate = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: {
+        modelId: configuredModel.id,
+        scope: "selected_tables",
+        targetIds: [table.id, table.id],
+      },
+    });
+    expect(duplicate.statusCode).toBe(400);
+    expect(duplicate.json()).toEqual({
+      code: "sensitive_data_suggestion_target_ids_duplicate",
+      message: "Each selected table or column must appear only once.",
+    });
+
+    const missingTable = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: {
+        modelId: configuredModel.id,
+        scope: "selected_tables",
+        targetIds: ["00000000-0000-4000-8000-000000000001"],
+      },
+    });
+    expect(missingTable.statusCode).toBe(404);
+    expect(missingTable.json()).toEqual({
+      code: "catalog_table_not_found",
+      message: "One or more selected Catalog Tables were not found in this database.",
+    });
+
+    const missingColumn = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: {
+        modelId: configuredModel.id,
+        scope: "selected_columns",
+        targetIds: ["00000000-0000-4000-8000-000000000002"],
+      },
+    });
+    expect(missingColumn.statusCode).toBe(404);
+    expect(missingColumn.json()).toEqual({
+      code: "catalog_column_not_found",
+      message: "One or more selected Catalog Columns were not found in this database.",
+    });
+    expect(modelCompleter.complete).not.toHaveBeenCalled();
+  } finally {
+    await app.close();
+  }
+});
+
+test("batches sensitive-data suggestions for schemas larger than one helper message", async () => {
+  const maxHelperMessageBytes = 64 * 1024;
+  const seenColumnIds: string[] = [];
+  const modelCompleter: ModelCompleter = {
+    complete: vi.fn(async (request) => {
+      const userMessage = request.messages.find((message) => message.role === "user")!;
+      expect(Buffer.byteLength(userMessage.content, "utf8")).toBeLessThanOrEqual(maxHelperMessageBytes);
+      const payload = JSON.parse(userMessage.content) as {
+        columns: Array<{ columnId: string; column: string }>;
+      };
+      expect(payload.columns.length).toBeLessThanOrEqual(10);
+      seenColumnIds.push(...payload.columns.map((column) => column.columnId));
+      return JSON.stringify({
+        suggestions: payload.columns.map((column) => ({
+          columnId: column.columnId,
+          sensitive: column.column.endsWith("_private"),
+        })),
+      });
+    }),
+  };
+  const { app, repository, database } = await setup(modelCompleter);
+  const columnCount = 900;
+  await repository.applySchemaSync(database.id, database.version, "all", [], {
+    schemaVersion: 1,
+    capabilities: { tables: "available", columns: "available", relationships: "available" },
+    tables: [{ name: "wide_table", sourceComment: null }],
+    columns: Array.from({ length: columnCount }, (_, index) => ({
+      tableName: "wide_table",
+      name: `field_${index.toString().padStart(4, "0")}${index % 10 === 0 ? "_private" : ""}`,
+      ordinalPosition: index + 1,
+      dataType: "character varying(255)",
+      isNullable: true,
+      defaultExpression: null,
+      primaryKeyPosition: null,
+      sourceComment: null,
+    })),
+    relationships: [],
+  });
+  const wideTable = (await repository.listTables(database.id)).find((table) => table.name === "wide_table")!;
+  const expectedColumnIds = (await repository.listColumns(database.id, wideTable.id)).map((column) => column.id);
+
+  try {
+    const response = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: { modelId: configuredModel.id, scope: "all" },
+    });
+
+    expect(response.statusCode).toBe(200);
+    const suggestions = response.json().suggestions as Array<{
+      columnName: string;
+      currentSensitive: boolean;
+      sensitive: boolean;
+    }>;
+    expect(suggestions).toHaveLength(columnCount);
+    expect(suggestions).toEqual(expect.arrayContaining([
+      expect.objectContaining({ columnName: "field_0000_private", currentSensitive: false, sensitive: true }),
+      expect.objectContaining({ columnName: "field_0001", currentSensitive: false, sensitive: false }),
+    ]));
+    expect(vi.mocked(modelCompleter.complete).mock.calls.length).toBeGreaterThan(1);
+    expect(seenColumnIds.slice().sort()).toEqual(expectedColumnIds.slice().sort());
+    expect(new Set(seenColumnIds).size).toBe(columnCount);
+  } finally {
+    await app.close();
+  }
+});
+
+test("retries one invalid sensitive-data classification before returning the review draft", async () => {
+  const modelCompleter: ModelCompleter = {
+    complete: vi.fn(async () => "unused"),
+  };
+  const { app, database, column } = await setup(modelCompleter);
+  vi.mocked(modelCompleter.complete)
+    .mockResolvedValueOnce("not-json")
+    .mockResolvedValueOnce(JSON.stringify({
+      suggestions: [{ columnId: column.id, sensitive: true }],
+    }));
+
+  try {
+    const response = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: { modelId: configuredModel.id, scope: "all" },
+    });
+
+    expect(response.statusCode).toBe(200);
+    expect(response.json().suggestions).toEqual([
+      expect.objectContaining({ columnId: column.id, sensitive: true }),
+    ]);
+    expect(modelCompleter.complete).toHaveBeenCalledTimes(2);
+  } finally {
+    await app.close();
+  }
+});
+
 test.each(["malformed", "incomplete", "duplicate"] as const)(
   "fails safely when sensitive-data suggestions are %s",
   async (kind) => {
@@ -201,13 +465,13 @@ test.each(["malformed", "incomplete", "duplicate"] as const)(
       const response = await app.inject({
         method: "POST",
         url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
-        payload: { modelId: configuredModel.id },
+        payload: { modelId: configuredModel.id, scope: "all" },
       });
 
       expect(response.statusCode).toBe(502);
       expect(response.json()).toEqual({
-        code: "sensitive_data_suggestion_failed",
-        message: "Sensitive-data suggestions could not be prepared.",
+        code: "sensitive_data_suggestion_invalid_response",
+        message: "The LLM returned an incomplete or invalid classification. No suggestions were applied.",
       });
       expect(response.body).not.toContain(rawResponse);
       expect(await repository.getColumn(database.id, column.tableId, column.id))
@@ -218,6 +482,34 @@ test.each(["malformed", "incomplete", "duplicate"] as const)(
   },
 );
 
+test("explains a sensitive-data suggestion provider failure without exposing provider details", async () => {
+  const modelCompleter: ModelCompleter = {
+    complete: vi.fn(async () => {
+      throw new ModelCompletionProviderError();
+    }),
+  };
+  const { app, repository, database, column } = await setup(modelCompleter);
+
+  try {
+    const response = await app.inject({
+      method: "POST",
+      url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
+      payload: { modelId: configuredModel.id, scope: "all" },
+    });
+
+    expect(response.statusCode).toBe(502);
+    expect(response.json()).toEqual({
+      code: "sensitive_data_suggestion_provider_unavailable",
+      message: "The selected LLM service could not complete the request. No suggestions were applied.",
+    });
+    expect(response.body).not.toContain("model completion failed");
+    expect(await repository.getColumn(database.id, column.tableId, column.id))
+      .toMatchObject({ sensitive: false });
+  } finally {
+    await app.close();
+  }
+});
+
 interface SseFrame {
   id?: string;
   event?: string;

backend/test/catalog-schema-routes.test.ts

@@ -241,14 +241,29 @@ test("keeps generated descriptions editable and preserves them across synchroniz
       version: idColumn.version,
       description: "Reviewed key",
       generatedDescription: "Generated key draft",
-      sensitive: true,
+      sensitive: false,
     },
   });
   expect(editedColumn.json()).toMatchObject({
+    description: "Reviewed key",
+    generatedDescription: "Generated key draft",
+    sensitive: false,
+  });
+  const sensitiveOnly = await app.inject({
+    method: "PATCH", url: `/catalog/databases/${database.id}/tables/${patients.id}/columns/${idColumn.id}`,
+    payload: { version: editedColumn.json().version, sensitive: true },
+  });
+  expect(sensitiveOnly.statusCode).toBe(200);
+  expect(sensitiveOnly.json()).toMatchObject({
     description: "Reviewed key",
     generatedDescription: "Generated key draft",
     sensitive: true,
   });
+  const emptyPatch = await app.inject({
+    method: "PATCH", url: `/catalog/databases/${database.id}/tables/${patients.id}/columns/${idColumn.id}`,
+    payload: { version: sensitiveOnly.json().version },
+  });
+  expect(emptyPatch.statusCode).toBe(400);
 
   const second = await app.inject({ method: "POST", url: `/catalog/databases/${database.id}/sync-runs`, payload: { version: database.version, scope: "all", tableIds: [] } });
   await waitFor(repository, second.json().id, "succeeded");