fix: scope and batch sensitive suggestions

2026-08-30 17:12:20 +02:00
parent 0736983bc5
commit cb40c09d9a
17 changed files with 1179 additions and 139 deletions
@@ -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;