Files
ThothII/backend/test/catalog-description-generation-worker.test.ts
T

644 lines
22 KiB
TypeScript

import { expect, test, vi } from "vitest";
import { DescriptionGenerationWorker } from "../src/catalog/description-generation-worker.js";
import type { DescriptionSourceSampler } from "../src/catalog/description-source-sampler.js";
import { MemoryCatalogRepository } from "../src/catalog/memory-repository.js";
import type { MetadataGenerationModels } from "../src/catalog/metadata-generation-models.js";
import { ModelCompletionCancelledError } from "../src/catalog/model-completer.js";
import type { ModelCompleter, ModelCompletionRequest } from "../src/catalog/model-completer.js";
import { CatalogOperationCoordinator } from "../src/catalog/operation-coordinator.js";
import type { WorkspaceRegistry } from "../src/workspaces/registry.js";
test("serializes Unlock with Start so stale recovery cannot release a new reservation", async () => {
let lookupStarted!: () => void;
const started = new Promise<void>((resolve) => { lookupStarted = resolve; });
let releaseLookup!: () => void;
const gate = new Promise<void>((resolve) => { releaseLookup = resolve; });
const repository = {
getActiveDescriptionGenerationRun: vi.fn(async () => {
lookupStarted();
await gate;
return undefined;
}),
} as unknown as MemoryCatalogRepository;
const resolveModel = vi.fn();
const worker = new DescriptionGenerationWorker(
repository,
{} as WorkspaceRegistry,
{
catalog: () => ({ models: [], default: "" }),
resolve: resolveModel,
} as MetadataGenerationModels,
{} as ModelCompleter,
new CatalogOperationCoordinator(),
{ sample: vi.fn(async () => []) },
);
const unlocking = worker.unlock();
await started;
await expect(worker.start(
"11111111-1111-4111-8111-111111111111",
"openai-mini",
"missing",
[],
)).rejects.toThrow("already active");
expect(resolveModel).not.toHaveBeenCalled();
releaseLookup();
await expect(unlocking).resolves.toBeUndefined();
});
test("exposes an awaitable background job and absorbs provider promise rejection", async () => {
const repository = new MemoryCatalogRepository();
const database = await repository.create({
workspaceId: "psd-clinical",
engine: "postgres",
databaseName: "warehouse",
schema: "datawarehouse",
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
});
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables: [{ name: "patients", sourceComment: null }],
columns: [{
tableName: "patients",
name: "birth_date",
ordinalPosition: 1,
dataType: "date",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}],
relationships: [],
});
const table = (await repository.listTables(database.id))[0]!;
const column = (await repository.listColumns(database.id, table.id))[0]!;
let rejectCompletion!: (error: Error) => void;
const pendingCompletion = new Promise<string>((_resolve, reject) => { rejectCompletion = reject; });
const completer: ModelCompleter = {
complete: vi.fn(async () => await pendingCompletion),
};
const models: MetadataGenerationModels = {
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
resolve: () => ({
id: "openai-mini",
provider: "openai",
model: "gpt-4.1-mini",
apiKeyEnv: "OPENAI_API_KEY",
apiKey: "test-provider-secret",
}),
};
const operations = new CatalogOperationCoordinator();
const sourceSampler: DescriptionSourceSampler = {
sample: vi.fn(async () => []),
};
const worker = new DescriptionGenerationWorker(
repository,
{
read: vi.fn(async () => ({
workspace: { workspace: { language: "it" } },
revision: {},
})),
} as unknown as WorkspaceRegistry,
models,
completer,
operations,
sourceSampler,
);
const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
let settled = false;
const waiting = worker.waitForRun(run.id).then(() => { settled = true; });
await new Promise((resolve) => setTimeout(resolve, 0));
expect(settled).toBe(false);
rejectCompletion(new Error("test-provider-secret private prompt raw response"));
await expect(waiting).resolves.toBeUndefined();
expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
status: "completed_with_errors",
failed: 1,
errorSummary: "Description generation completed with errors.",
});
const events = await repository.listDescriptionGenerationEvents(run.id);
expect(JSON.stringify(events)).not.toMatch(/test-provider-secret|private prompt|raw response/);
const release = operations.reserve(database.id);
release();
await expect(worker.start(
database.id,
"openai-mini",
"selected_columns",
[],
)).rejects.toThrow("at least one target ID is required");
});
test("marks an active run interrupted when the backend worker stops", async () => {
const repository = new MemoryCatalogRepository();
const database = await repository.create({
workspaceId: "psd-clinical",
engine: "postgres",
databaseName: "warehouse",
schema: "datawarehouse",
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
});
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables: [{ name: "patients", sourceComment: null }],
columns: [{
tableName: "patients",
name: "status",
ordinalPosition: 1,
dataType: "text",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}],
relationships: [],
});
const table = (await repository.listTables(database.id))[0]!;
const column = (await repository.listColumns(database.id, table.id))[0]!;
const completer: ModelCompleter = {
complete: vi.fn(async (request) => await new Promise<string>((_resolve, reject) => {
const cancel = () => reject(new ModelCompletionCancelledError());
if (request.signal.aborted) cancel();
else request.signal.addEventListener("abort", cancel, { once: true });
})),
};
const worker = new DescriptionGenerationWorker(
repository,
{
read: vi.fn(async () => ({
workspace: { workspace: { language: "it" } },
revision: {},
})),
} as unknown as WorkspaceRegistry,
{
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
resolve: () => ({
id: "openai-mini",
provider: "openai",
model: "gpt-4.1-mini",
apiKeyEnv: "OPENAI_API_KEY",
apiKey: "test-provider-secret",
}),
},
completer,
new CatalogOperationCoordinator(),
{ sample: vi.fn(async () => []) },
);
const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
await vi.waitFor(() => expect(completer.complete).toHaveBeenCalledOnce());
await worker.stop();
expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
status: "interrupted",
errorSummary: "Description generation was interrupted by backend shutdown.",
});
expect(await repository.listDescriptionGenerationEvents(run.id)).toContainEqual(
expect.objectContaining({
level: "warning",
message: "Description generation was interrupted by backend shutdown.",
}),
);
});
test("adds only bounded transient source samples to the model request", async () => {
const repository = new MemoryCatalogRepository();
const database = await repository.create({
workspaceId: "psd-clinical",
engine: "postgres",
databaseName: "warehouse",
schema: "datawarehouse",
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
});
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables: [{ name: "patients", sourceComment: null }],
columns: [{
tableName: "patients",
name: "patient_email",
ordinalPosition: 1,
dataType: "text",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}, {
tableName: "patients",
name: "ward",
ordinalPosition: 2,
dataType: "text",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}],
relationships: [],
});
const table = (await repository.listTables(database.id))[0]!;
const columns = await repository.listColumns(database.id, table.id);
const column = columns.find((candidate) => candidate.name === "patient_email")!;
const ward = columns.find((candidate) => candidate.name === "ward")!;
await repository.updateColumnMetadata(
database.id,
table.id,
column.id,
column.version,
column.description,
column.generatedDescription,
true,
);
const sampleSecret = "real.patient@hospital.invalid";
const sourceSampler: DescriptionSourceSampler = {
sample: vi.fn(async () => [{
targetId: column.id,
tableName: table.name,
rows: [
{ fields: [{ name: column.name, value: sampleSecret }] },
{ fields: [{ name: column.name, value: "row-2" }] },
{ fields: [{ name: column.name, value: "row-3" }] },
],
representativeValues: [{
column: column.name,
values: [sampleSecret, sampleSecret, "two", "three"],
}],
}, {
targetId: ward.id,
tableName: table.name,
rows: [
{ fields: [{ name: ward.name, value: "row-4" }] },
{ fields: [{ name: ward.name, value: "row-5" }] },
{ fields: [{ name: ward.name, value: "row-6" }] },
{ fields: [{ name: ward.name, value: "row-7" }] },
{ fields: [{ name: ward.name, value: "row-8" }] },
{ fields: [{ name: ward.name, value: "row-9-must-be-omitted" }] },
],
representativeValues: [{
column: ward.name,
values: ["ward-1", "ward-2", "ward-3", "ward-4", "ward-5", "ward-6-must-be-omitted"],
}],
}]),
};
const completer: ModelCompleter = {
complete: vi.fn(async () => JSON.stringify({
results: [{
targetId: column.id,
outcome: "generated",
description: "Stato amministrativo del paziente.",
}, {
targetId: ward.id,
outcome: "generated",
description: "Reparto associato al paziente.",
}],
})),
};
const models: MetadataGenerationModels = {
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
resolve: () => ({
id: "openai-mini",
provider: "openai",
model: "gpt-4.1-mini",
apiKeyEnv: "OPENAI_API_KEY",
apiKey: "test-provider-secret",
}),
};
const worker = new DescriptionGenerationWorker(
repository,
{
read: vi.fn(async () => ({
workspace: { workspace: { language: "it" } },
revision: {},
})),
} as unknown as WorkspaceRegistry,
models,
completer,
new CatalogOperationCoordinator(),
sourceSampler,
);
const run = await worker.start(
database.id,
"openai-mini",
"selected_columns",
[column.id, ward.id],
);
await worker.waitForRun(run.id);
expect(sourceSampler.sample).toHaveBeenCalledWith(
expect.objectContaining({ id: database.id, binding: database.binding }),
[
{ targetId: column.id, tableName: table.name, columnNames: [] },
{ targetId: ward.id, tableName: table.name, columnNames: [ward.name] },
],
expect.any(AbortSignal),
);
const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
expect(request.messages[0]?.content).toContain("untrusted");
const userMessage = request.messages[1]!.content;
const context = JSON.parse(userMessage.slice(userMessage.indexOf("\n") + 1));
const sampledRows = context.targets.flatMap(
(targetContext: { sourceSample?: { rows: unknown[] } }) => targetContext.sourceSample?.rows ?? [],
);
const representativeValues = context.targets.flatMap(
(targetContext: { sourceSample?: { representativeValues: Array<{ values: unknown[] }> } }) => (
targetContext.sourceSample?.representativeValues.flatMap((entry) => entry.values) ?? []
),
);
expect(sampledRows).toHaveLength(10);
expect(representativeValues).toHaveLength(10);
expect(context.targets[0].sourceSample.rows).toHaveLength(5);
expect(context.targets[1].sourceSample.rows).toHaveLength(5);
expect(context.targets[0].sourceSample.representativeValues).toEqual([{
column: column.name,
values: [
"marta.rossi@example.com",
"luca.bianchi@example.com",
"elena.conti@example.com",
"paolo.romano@example.com",
"giulia.ferrari@example.com",
],
}]);
expect(context.targets[1].sourceSample.representativeValues).toEqual([{
column: ward.name,
values: ["ward-1", "ward-2", "ward-3", "ward-4", "ward-5"],
}]);
expect(userMessage).not.toContain(sampleSecret);
expect(userMessage).not.toMatch(/synthetic|fake|fittizi/i);
expect(userMessage).toContain("marta.rossi@example.com");
expect(userMessage).not.toMatch(
/row-9-must-be-omitted|ward-6-must-be-omitted/,
);
const persisted = JSON.stringify({
run: await repository.getDescriptionGenerationRun(run.id),
events: await repository.listDescriptionGenerationEvents(run.id),
database: await repository.get(database.id),
table: await repository.getTable(database.id, table.id),
column: await repository.getColumn(database.id, table.id, column.id),
ward: await repository.getColumn(database.id, table.id, ward.id),
});
expect(persisted).not.toContain(sampleSecret);
});
test("gives every sensitive column synthetic context without consuming the real sample budget", async () => {
const repository = new MemoryCatalogRepository();
const database = await repository.create({
workspaceId: "psd-clinical",
engine: "postgres",
databaseName: "warehouse",
schema: "datawarehouse",
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
});
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables: [{ name: "patients", sourceComment: null }],
columns: [{
tableName: "patients",
name: "patient_email",
ordinalPosition: 1,
dataType: "text",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}, {
tableName: "patients",
name: "patient_phone",
ordinalPosition: 2,
dataType: "text",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}, {
tableName: "patients",
name: "ward",
ordinalPosition: 3,
dataType: "text",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}],
relationships: [],
});
const table = (await repository.listTables(database.id))[0]!;
const columns = await repository.listColumns(database.id, table.id);
const email = columns.find((column) => column.name === "patient_email")!;
const phone = columns.find((column) => column.name === "patient_phone")!;
const ward = columns.find((column) => column.name === "ward")!;
for (const column of [email, phone]) {
await repository.updateColumnMetadata(
database.id,
table.id,
column.id,
column.version,
column.description,
column.generatedDescription,
true,
);
}
const wardValues = ["ward-a", "ward-b", "ward-c", "ward-d", "ward-e"];
const sourceSampler: DescriptionSourceSampler = {
sample: vi.fn(async (_database, targets) => targets.map((target) => {
if (target.columnNames.length === 0) {
return {
targetId: target.targetId,
tableName: target.tableName,
rows: [],
representativeValues: [],
};
}
const columnName = target.columnNames[0]!;
return {
targetId: target.targetId,
tableName: target.tableName,
rows: wardValues.map((value) => ({ fields: [{ name: columnName, value }] })),
representativeValues: [{ column: columnName, values: wardValues }],
};
})),
};
const completer: ModelCompleter = {
complete: vi.fn(async (request) => {
const context = JSON.parse(request.messages[1]!.content.split("\n").slice(1).join("\n"));
return JSON.stringify({
results: context.targets.map((target: { targetId: string }) => ({
targetId: target.targetId,
outcome: "generated",
description: "Descrizione generata.",
})),
});
}),
};
const models: MetadataGenerationModels = {
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
resolve: () => ({
id: "openai-mini",
provider: "openai",
model: "gpt-4.1-mini",
apiKeyEnv: "OPENAI_API_KEY",
apiKey: "test-provider-secret",
}),
};
const worker = new DescriptionGenerationWorker(
repository,
{
read: vi.fn(async () => ({
workspace: { workspace: { language: "it" } },
revision: {},
})),
} as unknown as WorkspaceRegistry,
models,
completer,
new CatalogOperationCoordinator(),
sourceSampler,
);
const run = await worker.start(
database.id,
"openai-mini",
"selected_columns",
[email.id, phone.id, ward.id],
);
await worker.waitForRun(run.id);
expect(sourceSampler.sample).toHaveBeenCalledWith(
expect.objectContaining({ id: database.id }),
[
{ targetId: email.id, tableName: table.name, columnNames: [] },
{ targetId: phone.id, tableName: table.name, columnNames: [] },
{ targetId: ward.id, tableName: table.name, columnNames: [ward.name] },
],
expect.any(AbortSignal),
);
const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
const context = JSON.parse(request.messages[1]!.content.split("\n").slice(1).join("\n"));
const targets = new Map(
context.targets.map((target: { targetId: string }) => [target.targetId, target]),
);
expect(targets.get(email.id)).toMatchObject({
sourceSample: {
rows: expect.arrayContaining([
{ fields: [{ name: email.name, value: "marta.rossi@example.com" }] },
]),
representativeValues: [{
column: email.name,
values: expect.arrayContaining(["marta.rossi@example.com"]),
}],
},
});
expect(targets.get(phone.id)).toMatchObject({
sourceSample: {
rows: expect.arrayContaining([
{ fields: [{ name: phone.name, value: "+39 02 5550 1001" }] },
]),
representativeValues: [{
column: phone.name,
values: expect.arrayContaining(["+39 02 5550 1001"]),
}],
},
});
expect(targets.get(ward.id)).toMatchObject({
sourceSample: {
rows: wardValues.map((value) => ({ fields: [{ name: ward.name, value }] })),
representativeValues: [{ column: ward.name, values: wardValues }],
},
});
});
test("continues metadata-only with one safe warning when source sampling is unavailable", async () => {
const repository = new MemoryCatalogRepository();
const database = await repository.create({
workspaceId: "psd-clinical",
engine: "postgres",
databaseName: "warehouse",
schema: "datawarehouse",
binding: {
transport: "rest_api",
baseUrl: "https://dwh.example.test",
restPath: "/rpc/run_query",
restAuth: "none",
},
});
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables: [{ name: "patients", sourceComment: null }],
columns: [{
tableName: "patients",
name: "status",
ordinalPosition: 1,
dataType: "text",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
}],
relationships: [],
});
const table = (await repository.listTables(database.id))[0]!;
const column = (await repository.listColumns(database.id, table.id))[0]!;
const samplingFailureSecret = "UNAVAILABLE_SAMPLE_DETAIL_48b1f1";
const sourceSampler: DescriptionSourceSampler = {
sample: vi.fn(async () => { throw new Error(samplingFailureSecret); }),
};
const completer: ModelCompleter = {
complete: vi.fn(async () => JSON.stringify({
results: [{
targetId: column.id,
outcome: "generated",
description: "Stato del paziente.",
}],
})),
};
const models: MetadataGenerationModels = {
catalog: () => ({ models: [{ id: "openai-mini", label: "OpenAI Mini" }], default: "openai-mini" }),
resolve: () => ({
id: "openai-mini",
provider: "openai",
model: "gpt-4.1-mini",
apiKeyEnv: "OPENAI_API_KEY",
apiKey: "test-provider-secret",
}),
};
const worker = new DescriptionGenerationWorker(
repository,
{
read: vi.fn(async () => ({
workspace: { workspace: { language: "it" } },
revision: {},
})),
} as unknown as WorkspaceRegistry,
models,
completer,
new CatalogOperationCoordinator(),
sourceSampler,
);
const run = await worker.start(database.id, "openai-mini", "selected_columns", [column.id]);
await worker.waitForRun(run.id);
expect(await repository.getDescriptionGenerationRun(run.id)).toMatchObject({
status: "completed",
processed: 1,
generated: 1,
failed: 0,
});
const request = vi.mocked(completer.complete).mock.calls[0]![0] as ModelCompletionRequest;
expect(request.messages[1]?.content).not.toMatch(/sourceSample|UNAVAILABLE_SAMPLE_DETAIL/);
const events = await repository.listDescriptionGenerationEvents(run.id);
expect(events.filter((event) => event.level === "warning")).toEqual([
expect.objectContaining({
message: "Source samples unavailable for this batch; generation continued with catalog metadata only.",
}),
]);
expect(JSON.stringify(events)).not.toContain(samplingFailureSecret);
});