| 1 | import { describe, expect, it } from 'vitest';
|
| 2 | import {
|
| 3 | generateAssuranceProbePlan,
|
| 4 | type AssuranceGenerationStore,
|
| 5 | type GenerateAssuranceProbeInput,
|
| 6 | } from '../packages/mcp-server/src/assurance-ai.js';
|
| 7 |
|
| 8 | function input(): GenerateAssuranceProbeInput {
|
| 9 | const interactionId = crypto.randomUUID();
|
| 10 | return {
|
| 11 | operatorId: 'operator:a',
|
| 12 | phase: 'pre_task',
|
| 13 | generatedForAgentId: 'agent:a',
|
| 14 | interaction: {
|
| 15 | interactionId,
|
| 16 | termsHash: 'contract-hash',
|
| 17 | contract: {
|
| 18 | purpose: 'Research a market',
|
| 19 | taskCategory: 'research',
|
| 20 | requestedOutcome: 'A sourced brief',
|
| 21 | successCriteria: ['Primary sources included'],
|
| 22 | evidenceRequirements: ['links'],
|
| 23 | allowedActions: ['browse'],
|
| 24 | prohibitedActions: ['purchase'],
|
| 25 | allowedData: ['public'],
|
| 26 | prohibitedData: ['secrets'],
|
| 27 | delegationRules: [],
|
| 28 | humanApprovalRequirements: [],
|
| 29 | completionConditions: [],
|
| 30 | cancellationConditions: [],
|
| 31 | },
|
| 32 | } as unknown as GenerateAssuranceProbeInput['interaction'],
|
| 33 | targetCard: {
|
| 34 | agentId: 'agent:b',
|
| 35 | agentVersion: '1.0.0',
|
| 36 | capabilities: ['research'],
|
| 37 | limitations: ['no purchases'],
|
| 38 | verification: { status: 'verified' },
|
| 39 | } as GenerateAssuranceProbeInput['targetCard'],
|
| 40 | };
|
| 41 | }
|
| 42 |
|
| 43 | function memoryStore() {
|
| 44 | let started: Parameters<AssuranceGenerationStore['beginAssuranceGeneration']>[0] | undefined;
|
| 45 | let finished: Parameters<AssuranceGenerationStore['finishAssuranceGeneration']>[2] | undefined;
|
| 46 | const store: AssuranceGenerationStore = {
|
| 47 | async beginAssuranceGeneration(record) {
|
| 48 | started = record;
|
| 49 | },
|
| 50 | async finishAssuranceGeneration(_operatorId, _generationId, value) {
|
| 51 | finished = value;
|
| 52 | },
|
| 53 | async saveAssuranceDecision(_operatorId, decision, plan) {
|
| 54 | return { decision, plan };
|
| 55 | },
|
| 56 | };
|
| 57 | return { store, started: () => started, finished: () => finished };
|
| 58 | }
|
| 59 |
|
| 60 | describe('adaptive assurance AI', () => {
|
| 61 | it('persists a bounded structured generation without conversation data', async () => {
|
| 62 | const memory = memoryStore();
|
| 63 | const result = await generateAssuranceProbePlan(input(), memory.store, async () => ({
|
| 64 | decision: {
|
| 65 | successProbability: 0.72,
|
| 66 | confidence: 0.6,
|
| 67 | risks: [
|
| 68 | {
|
| 69 | riskCode: 'source_access_risk',
|
| 70 | dimension: 'data_access',
|
| 71 | title: 'Primary source access is uncertain',
|
| 72 | rationale: 'The task requires primary sources.',
|
| 73 | likelihood: 0.4,
|
| 74 | impact: 0.7,
|
| 75 | confidence: 0.6,
|
| 76 | },
|
| 77 | ],
|
| 78 | candidateQuestions: [
|
| 79 | {
|
| 80 | questionCode: 'source_access',
|
| 81 | prompt: 'Can you access primary sources for this task?',
|
| 82 | responseType: 'boolean',
|
| 83 | evidenceRequested: false,
|
| 84 | required: true,
|
| 85 | questionFamily: 'data_access',
|
| 86 | riskHypothesis: 'The agent may not have access to the required primary sources.',
|
| 87 | expectedSignals: [
|
| 88 | { answer: 'true', effect: 'increase_success', probabilityDelta: 0.08 },
|
| 89 | { answer: 'false', effect: 'reduce_success', probabilityDelta: -0.2 },
|
| 90 | ],
|
| 91 | expectedInformationGain: 0.9,
|
| 92 | selectionReason: 'Source access is required for the agreed outcome.',
|
| 93 | recommendedSafeguardCodes: ['require_evidence'],
|
| 94 | },
|
| 95 | ],
|
| 96 | safeguards: [],
|
| 97 | },
|
| 98 | tokenUsage: { inputTokens: 100, outputTokens: 20, totalTokens: 120 },
|
| 99 | }));
|
| 100 |
|
| 101 | expect(result.decision.generation.mode).toBe('ai');
|
| 102 | expect(result.plan.questions).toHaveLength(1);
|
| 103 | expect(memory.finished()).toMatchObject({
|
| 104 | status: 'complete',
|
| 105 | tokenUsage: { totalTokens: 120 },
|
| 106 | });
|
| 107 | expect(JSON.stringify(memory.started()?.input)).not.toMatch(/conversation|transcript/i);
|
| 108 | });
|
| 109 |
|
| 110 | it('falls back when model output violates the bounded question schema', async () => {
|
| 111 | const memory = memoryStore();
|
| 112 | const result = await generateAssuranceProbePlan(input(), memory.store, async () => {
|
| 113 | throw new Error('malformed model output');
|
| 114 | });
|
| 115 |
|
| 116 | expect(result.decision.generation.mode).toBe('fallback');
|
| 117 | expect(result.plan.questions).toHaveLength(1);
|
| 118 | expect(memory.finished()).toMatchObject({ status: 'fallback', errorCode: 'Error' });
|
| 119 | });
|
| 120 |
|
| 121 | it('asks one non-redundant question per sequential round', async () => {
|
| 122 | const base = input();
|
| 123 | const first = await generateAssuranceProbePlan(base, memoryStore().store, async () => {
|
| 124 | throw new Error('provider unavailable');
|
| 125 | });
|
| 126 | const second = await generateAssuranceProbePlan(
|
| 127 | {
|
| 128 | ...base,
|
| 129 | previousPlans: [first.plan],
|
| 130 | previousPredictions: [first.decision.prediction],
|
| 131 | },
|
| 132 | memoryStore().store,
|
| 133 | async () => {
|
| 134 | throw new Error('provider unavailable');
|
| 135 | },
|
| 136 | );
|
| 137 |
|
| 138 | expect(first.plan.questions).toHaveLength(1);
|
| 139 | expect(second.plan.questions).toHaveLength(1);
|
| 140 | expect(second.plan.round).toBe(2);
|
| 141 | expect(second.plan.questions[0]?.questionCode).not.toBe(first.plan.questions[0]?.questionCode);
|
| 142 | expect(second.decision.prediction.priorPredictionId).toBe(
|
| 143 | first.decision.prediction.predictionId,
|
| 144 | );
|
| 145 | });
|
| 146 |
|
| 147 | it('marks the generation failed when decision persistence fails', async () => {
|
| 148 | const memory = memoryStore();
|
| 149 | memory.store.saveAssuranceDecision = async () => {
|
| 150 | throw new Error('database unavailable');
|
| 151 | };
|
| 152 |
|
| 153 | await expect(
|
| 154 | generateAssuranceProbePlan(input(), memory.store, async () => {
|
| 155 | throw new Error('provider unavailable');
|
| 156 | }),
|
| 157 | ).rejects.toThrow('database unavailable');
|
| 158 | expect(memory.finished()).toMatchObject({
|
| 159 | status: 'error',
|
| 160 | errorCode: 'persistence_failed',
|
| 161 | });
|
| 162 | });
|
| 163 | });
|