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('computes evidence deterministically when a local model returns values only', async () => {
    const out = await structureDocument({
      ...base,
      gateway: {
        complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ value: { name: 'Acme' } }) })),
      },
    });
    expect(out).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } });
    if (out.status !== 'mechanically_validated') throw new Error('expected validated');
    expect(out.evidence).toEqual([
      { field: 'name', sourceId: 'page1', start: 6, end: 10, transform: 'exact' },
    ]);
  });
  it('accepts a flat object or a fenced JSON object from a small local model', async () => {
    const flat = await structureDocument({
      ...base,
      gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ name: 'Acme' }) })) },
    });
    expect(flat).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } });
    const fenced = await structureDocument({
      ...base,
      gateway: {
        complete: vi.fn(async () => ({ ...response, content: '```json\n{"value":{"name":"Acme"}}\n```' })),
      },
    });
    expect(fenced).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } });
  });
  it('refuses ambiguous or absent values when computing evidence itself', async () => {
    const twice = await structureDocument({
      ...base,
      sources: [{ id: 'p1', text: 'Acme AB and Acme AB' }],
      gateway: {
        complete: vi.fn(async () => ({
          ...response,
          content: JSON.stringify({ value: { name: 'Acme AB' } }),
        })),
      },
    });
    expect(twice.status).toBe('needs_review');
    const absent = await structureDocument({
      ...base,
      gateway: {
        complete: vi.fn(async () => ({
          ...response,
          content: JSON.stringify({ value: { name: 'Globex' } }),
        })),
      },
    });
    expect(absent.status).toBe('needs_review');
  });
  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]);
  });
});