0pt/tests/document-workbench/structure-document.test.ts

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

type Req = Omit<Parameters<Inf0GatewayClient['complete']>[0], 'messages' | 'inputTokens'>;
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 Req;
const response: Inf0GatewayReceipt = { status: 'model_stopped', content: JSON.stringify({ value: { name: 'Acme' }, evidence: [{ field: 'name', sourceId: 'page1', start: 6, end: 10, transform: 'exact' }] }), modelExecuted: 'm1', providerKind: 'local', promptTokens: 50, completionTokens: 20, totalTokens: 70 };
const base = { request, countTokens: vi.fn(async () => 50), instructions: 'Extract the customer name', fields: [{ name: 'name', type: 'string' as const, required: true }], sources: [{ id: 'page1', text: 'Name: Acme' }] };
describe('structure document workflow', () => {
  it('returns only source-supported values and carries model usage without raw prompt', async () => {
    const complete = vi.fn(async () => response);
    const out = await structureDocument({ ...base, gateway: { complete } });
    expect(out).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' }, model: { modelExecuted: 'm1' } });
    expect(out.model).not.toHaveProperty('content');
    expect(complete).toHaveBeenCalledOnce();
  });
  it('sends hallucinated data, invalid JSON and truncated output to review', async () => {
    const fake = { ...response, content: JSON.stringify({ value: { name: 'Other' }, evidence: [{ field: 'name', sourceId: 'page1', start: 6, end: 10, transform: 'exact' }] }) };
    expect((await structureDocument({ ...base, gateway: { complete: vi.fn(async () => fake) } })).status).toBe('needs_review');
    expect((await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, content: 'not json' })) } })).status).toBe('needs_review');
    expect((await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, status: 'needs_review' as const })) } })).status).toBe('needs_review');
  });
  it('runs through the real INF0 adapter and rejects an ungrounded response', async () => {
    const fetchMock = vi.fn(async () => Response.json({
      content: response.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: fetchMock });
    const out = await structureDocument({ ...base, gateway });
    expect(out).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } });
    const call = fetchMock.mock.calls[0] as unknown as [string, RequestInit];
    const body = JSON.parse(String(call[1].body));
    expect(body.messages[1].content).toContain('Name: Acme');
    expect(body.data_classification).toBe('high');
  });

  it('rejects malformed contracts before a model call', async () => {
    const complete = vi.fn(async () => response);
    await expect(structureDocument({ ...base, fields: [], gateway: { complete } })).rejects.toThrow('Invalid structure');
    expect(complete).not.toHaveBeenCalled();
  });
  it('validates against the exact source and field contract supplied before async tokenization', async () => {
    const sources = [{ id: 'page1', text: 'Name: Evil' }];
    const fields = [{ name: 'name', type: 'string' as const, required: true }];
    const countTokens = vi.fn(async () => { sources[0]!.text = 'Name: Acme'; fields[0]!.required = false; return 50; });
    const complete = vi.fn<Inf0GatewayClient['complete']>(async () => response);
    const result = await structureDocument({ ...base, sources, fields, countTokens, gateway: { complete } });
    expect(result).toMatchObject({ status: 'needs_review', issues: [{ field: 'name', code: 'evidence_mismatch' }] });
    expect(complete.mock.calls[0]?.[0]?.messages[1]?.content).toContain('Name: Evil');
  });
  it('uses the original admission and policy even if the caller mutates its request during tokenization', async () => {
    const mutableRequest = structuredClone(request);
    const countTokens = vi.fn(async () => { mutableRequest.admission.modelId = 'changed'; mutableRequest.policy.allowedModelIds.push('changed'); return 50; });
    const complete = vi.fn<Inf0GatewayClient['complete']>(async () => response);
    await structureDocument({ ...base, request: mutableRequest, countTokens, gateway: { complete } });
    expect(complete.mock.calls[0]?.[0]?.admission.modelId).toBe('local_only');
    expect(complete.mock.calls[0]?.[0]?.policy.allowedModelIds).toEqual(['local_only']);
  });

  it('keeps tokenizer input identical to the message sent to the gateway', async () => {
    const countTokens = vi.fn(async (messages: Parameters<Inf0GatewayClient['complete']>[0]['messages']) => {
      expect(Object.isFrozen(messages)).toBe(true);
      expect(Object.isFrozen(messages[1])).toBe(true);
      return 50;
    });
    const complete = vi.fn<Inf0GatewayClient['complete']>(async () => response);
    await structureDocument({ ...base, countTokens, gateway: { complete } });
    expect(complete.mock.calls[0]?.[0]?.messages).toEqual(countTokens.mock.calls[0]?.[0]);
  });

});