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((resolve) => { lookupStarted = resolve; }); let releaseLookup!: () => void; const gate = new Promise((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((_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((_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: "status", 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 === "status")!; const ward = columns.find((candidate) => candidate.name === "ward")!; const sampleSecret = "ONLY_IN_TRANSIENT_SAMPLE_7f29c8"; 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-must-be-omitted" }] }, ], representativeValues: [{ column: ward.name, values: ["ward-1", "ward-2", "ward-3-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: [column.name] }, { 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(5); expect(representativeValues).toHaveLength(5); expect(context.targets[0].sourceSample.rows).toHaveLength(3); expect(context.targets[1].sourceSample.rows).toHaveLength(2); expect(context.targets[0].sourceSample.representativeValues).toEqual([{ column: column.name, values: [sampleSecret, "two", "three"], }]); expect(context.targets[1].sourceSample.representativeValues).toEqual([{ column: ward.name, values: ["ward-1", "ward-2"], }]); expect(userMessage).toContain(sampleSecret); expect(userMessage).not.toMatch( /row-6-must-be-omitted|ward-3-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("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); });