0pt/tests/document-workbench/structured-batch.test.ts

import { describe, expect, it, vi } from 'vitest';
import { validateStructuredBatch, exportValidatedBatch } from '../../src/document-workbench/structured-batch';
import { structureDocumentBatch } from '../../src/document-workbench/structure-document-batch';
import { Inf0GatewayClient, type Inf0GatewayReceipt } from '../../src/document-workbench/inf0-gateway-client';

const fields = [{ name: 'name', type: 'string' as const, required: true }];
const sources = [{ id: 's1', text: 'Acme\nBeta' }];
const records = [
  { value: { name: 'Acme' }, evidence: [{ field: 'name', sourceId: 's1', start: 0, end: 4, transform: 'exact' as const }] },
  { value: { name: 'Beta' }, evidence: [{ field: 'name', sourceId: 's1', start: 5, end: 9, transform: 'exact' as const }] },
];
const request = {
  context: { tenantId: 'a', projectId: 'p', actorId: 'u', capability: 'inf0.complete', correlationId: 'c' },
  admission: { tenantId: 'a', projectId: 'p', jobId: 'j', modelId: 'local_only', dataClass: 'high', maxInputTokens: 200, maxOutputTokens: 100, maximumCredits: 5 },
  policy: { tenantId: 'a', projectId: 'p', allowedModelIds: ['local_only'], allowedDataClasses: ['high'], destination: 'internal', externalTransferAllowed: false, maxInputTokens: 200, maxOutputTokens: 100 },
  reservation: { tenantId: 'a', projectId: 'p', jobId: 'j', status: 'reserved', credits: 5 },
} satisfies Omit<Parameters<Inf0GatewayClient['complete']>[0], 'messages' | 'inputTokens'>;
const receipt: Inf0GatewayReceipt = { status: 'model_stopped', content: JSON.stringify({ records }), modelExecuted: 'm1', providerKind: 'local', promptTokens: 50, completionTokens: 20, totalTokens: 70 };

describe('evidence-bound document batches', () => {
  it('serializes every supported row in order to JSON and CSV', () => {
    const json = exportValidatedBatch({ format: 'json', fields, sources, records });
    expect(json.status).toBe('ready');
    if (json.status === 'ready') expect(JSON.parse(json.content)).toEqual([{ name: 'Acme' }, { name: 'Beta' }]);
    expect(exportValidatedBatch({ format: 'csv', fields, sources, records })).toEqual({
      status: 'ready', mediaType: 'text/csv', content: '"name"\r\n"Acme"\r\n"Beta"\r\n',
    });
  });
  it('rejects the entire output when one row is unsupported or malformed', () => {
    const altered = [records[0], { ...records[1]!, value: { name: 'Invented' } }];
    expect(exportValidatedBatch({ format: 'json', fields, sources, records: altered })).toMatchObject({ status: 'needs_review', issues: [{ row: 1, field: 'name', code: 'evidence_mismatch' }] });
    expect(validateStructuredBatch({ fields, sources, records: [null] })).toMatchObject({ status: 'needs_review', issues: [{ row: 0, code: 'invalid_record' }] });
    expect(validateStructuredBatch({ fields, sources, records: [] })).toMatchObject({ status: 'needs_review' });
    expect(validateStructuredBatch({ fields, sources, records: Array(201).fill(records[0]) })).toMatchObject({ status: 'needs_review' });
  });
  it('blocks spreadsheet formulas in any row while JSON preserves the literal value', () => {
    const source = [{ id: 's1', text: 'Acme\n=1+1' }];
    const unsafe = [records[0]!, { value: { name: '=1+1' }, evidence: [{ field: 'name', sourceId: 's1', start: 5, end: 9, transform: 'exact' }] }];
    expect(exportValidatedBatch({ format: 'csv', fields, sources: source, records: unsafe })).toMatchObject({ status: 'needs_review', issues: [{ row: 1, code: 'csv_unsafe_cell' }] });
    expect(exportValidatedBatch({ format: 'json', fields, sources: source, records: unsafe }).status).toBe('ready');
  });
  it('preserves quoted multiline headers and cells in a multi-row CSV', () => {
    const f = [{ name: 'company\r\nname', type: 'string' as const, required: true }];
    const s = [{ id: 's', text: 'A\r\nB,C' }];
    const r = [{ value: { 'company\r\nname': 'A\r\nB' }, evidence: [{ field: 'company\r\nname', sourceId: 's', start: 0, end: 4, transform: 'exact' }] },
      { value: { 'company\r\nname': ',C' }, evidence: [{ field: 'company\r\nname', sourceId: 's', start: 4, end: 6, transform: 'exact' }] }];
    expect(exportValidatedBatch({ format: 'csv', fields: f, sources: s, records: r })).toMatchObject({ status: 'ready', content: '"company\r\nname"\r\n"A\r\nB"\r\n",C"\r\n' });
  });
  it('extracts multiple records and keeps their evidence tied to the original source', async () => {
    const mutable = structuredClone(sources);
    const complete = vi.fn<Inf0GatewayClient['complete']>(async () => receipt);
    const countTokens = vi.fn(async (messages: Parameters<Inf0GatewayClient['complete']>[0]['messages']) => {
      expect(Object.isFrozen(messages[1])).toBe(true);
      mutable[0]!.text = 'Evil\nBeta';
      return 50;
    });
    const result = await structureDocumentBatch({ request, fields, sources: mutable, instructions: 'Extract companies', countTokens, gateway: { complete } });
    expect(result).toMatchObject({ status: 'mechanically_validated', records });
    expect(result.model).not.toHaveProperty('content');
    expect(complete.mock.calls[0]?.[0]?.messages).toEqual(countTokens.mock.calls[0]?.[0]);
  });
  it('runs a batch through the real INF0 client with scoped policy and usage checks', async () => {
    const fetchImpl = vi.fn(async () => Response.json({ content: receipt.content, modelExecuted: 'm1', finishReason: 'stop',
      routeDecision: { route: 'local_only', providerKind: 'local' }, cost: { promptTokens: 50, completionTokens: 20, totalTokens: 70 } }));
    const gateway = new Inf0GatewayClient({ baseUrl: 'https://inf0.example.test', serviceToken: 'test', fetchImpl });
    const result = await structureDocumentBatch({ request, fields, sources, instructions: 'Extract companies', countTokens: async () => 50, gateway });
    expect(result).toMatchObject({ status: 'mechanically_validated', records });
    if (result.status !== 'mechanically_validated') throw new Error('Expected a verified batch');
    const exported = exportValidatedBatch({ format: 'json', fields, sources, records: result.records });
    expect(exported.status).toBe('ready');
    if (exported.status === 'ready') expect(JSON.parse(exported.content)).toEqual([{ name: 'Acme' }, { name: 'Beta' }]);
    expect(fetchImpl).toHaveBeenCalledOnce();
  });
  it('requires review for CSV if any optional field is absent in any row', () => {
    const optional = [...fields, { name: 'reference', type: 'string' as const, required: false }];
    expect(exportValidatedBatch({ format: 'json', fields: optional, sources, records }).status).toBe('ready');
    expect(exportValidatedBatch({ format: 'csv', fields: optional, sources, records })).toMatchObject({
      status: 'needs_review', issues: [{ row: 0, field: 'reference', code: 'csv_missing_value' }],
    });
  });
  it('routes truncated, empty, malformed and unsupported model results to review', async () => {
    for (const r of [
      { ...receipt, status: 'needs_review' as const }, { ...receipt, content: '{' },
      { ...receipt, content: JSON.stringify({ records: [] }) },
      { ...receipt, content: JSON.stringify({ records: [{ ...records[0], value: { name: 'Wrong' } }] }) },
    ]) {
      const result = await structureDocumentBatch({ request, fields, sources, instructions: 'Extract companies', countTokens: async () => 50, gateway: { complete: async () => r } });
      expect(result.status).toBe('needs_review');
    }
  });
});