feat(harness): add information gain stop and audit

2026-07-27 01:03:34 +08:00
parent de5a5b09d9
commit d0452184ee
92 changed files with 5019 additions and 122 deletions
@@ -0,0 +1,114 @@
+package com.superbiz.agent.harness.audit;
+
+import com.superbiz.agent.domain.entity.AgentStep;
+import com.superbiz.agent.harness.core.DiagnosisHarnessCore;
+import com.superbiz.agent.harness.core.RunBudgetLimits;
+import com.superbiz.agent.harness.core.RunContext;
+import com.superbiz.agent.harness.retry.HarnessRetryPolicies;
+import com.superbiz.agent.repository.AgentStepRepository;
+import org.junit.jupiter.api.Test;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+class ModelCallAuditorTest {
+
+    @Test
+    void recordsComponentUsageIntoRunBudgetTraceAndAgentStep() {
+        DiagnosisHarnessCore core = core();
+        RunContext context = core.startRun("session-token", "run-token");
+        AgentStepRepository steps = mock(AgentStepRepository.class);
+        AgentStep step = AgentStep.builder().id(7L).runId(context.runId()).stepIndex(0).build();
+        when(steps.findByRunIdAndStepIndex(context.runId(), 0)).thenReturn(Optional.of(step));
+        List<DiagnosisTraceAuditEvent> trace = new ArrayList<>();
+        ModelCallAuditor auditor = new ModelCallAuditor(core, trace::add, steps);
+
+        core.beforeModelCall(context);
+        ModelCallLedger.Call call = auditor.begin(context, ModelCallComponent.DIAGNOSIS_AGENT);
+        auditor.recordUsage(context, call, 13, 5, true);
+
+        assertEquals(18, context.budget().snapshot().totalTokens());
+        assertEquals(18, context.modelCalls().snapshot().totalTokens());
+        assertEquals(18, step.getTokenCount());
+        verify(steps).save(step);
+        DiagnosisTraceAuditEvent event = trace.get(0);
+        assertEquals(TraceEventType.MODEL_TOKEN_USAGE, event.eventType());
+        assertEquals("DIAGNOSIS_AGENT", event.details().get("component"));
+        assertEquals(1, event.details().get("component_round"));
+        assertEquals(13L, event.details().get("input_tokens"));
+        assertEquals(5L, event.details().get("output_tokens"));
+        assertEquals(18L, event.details().get("total_tokens"));
+        assertEquals(true, event.details().get("usage_available"));
+        String details = event.details().toString().toLowerCase();
+        assertFalse(details.contains("prompt"));
+        assertFalse(details.contains("content"));
+        assertFalse(details.contains("response"));
+        assertFalse(details.contains("argument"));
+    }
+
+    @Test
+    void unavailableUsageIsVisibleWithoutFabricatingTokens() {
+        DiagnosisHarnessCore core = core();
+        RunContext context = core.startRun("session-missing", "run-missing");
+        List<DiagnosisTraceAuditEvent> trace = new ArrayList<>();
+        ModelCallAuditor auditor = new ModelCallAuditor(
+                core, trace::add, mock(AgentStepRepository.class));
+
+        core.beforeModelCall(context);
+        ModelCallLedger.Call call = auditor.begin(context, ModelCallComponent.INTENT_ROUTER);
+        auditor.recordUsage(context, call, 0, 0, false);
+
+        ModelCallLedger.Snapshot snapshot = context.modelCalls().snapshot();
+        assertEquals(1, snapshot.auditedCallCount());
+        assertEquals(1, snapshot.usageUnavailableCount());
+        assertEquals(0, snapshot.totalTokens());
+        assertEquals(0, context.budget().snapshot().totalTokens());
+        assertEquals(false, trace.get(0).details().get("usage_available"));
+        assertFalse(trace.get(0).details().containsKey("input_tokens"));
+        assertFalse(trace.get(0).details().containsKey("output_tokens"));
+        assertFalse(trace.get(0).details().containsKey("total_tokens"));
+        DiagnosisTraceAuditEvent finished = TraceAuditEvents.runFinished(
+                context, com.superbiz.agent.harness.contract.IntentType.DIAGNOSIS,
+                com.superbiz.agent.harness.contract.ReleaseOutcome.FAILED, 1);
+        assertEquals(false, finished.details().get("tokens_reconciled"));
+    }
+
+    @Test
+    void runFinishedExposesTokenReconciliationWithoutContent() {
+        DiagnosisHarnessCore core = core();
+        RunContext context = core.startRun("session-finish", "run-finish");
+        ModelCallAuditor auditor = new ModelCallAuditor(
+                core, DiagnosisTraceRecorder.noop(), mock(AgentStepRepository.class));
+        core.beforeModelCall(context);
+        auditor.recordUsage(context,
+                auditor.begin(context, ModelCallComponent.INTENT_ROUTER), 8, 2, true);
+
+        DiagnosisTraceAuditEvent event = TraceAuditEvents.runFinished(
+                context, com.superbiz.agent.harness.contract.IntentType.DIAGNOSIS,
+                com.superbiz.agent.harness.contract.ReleaseOutcome.FALLBACK, 50);
+
+        assertEquals(10L, event.details().get("run_total_tokens"));
+        assertEquals(10L, event.details().get("audited_total_tokens"));
+        assertEquals(true, event.details().get("tokens_reconciled"));
+        assertEquals(0, event.details().get("usage_unavailable_count"));
+        assertTrue(event.details().containsKey("audited_model_call_count"));
+        assertFalse(event.details().toString().contains("secret"));
+    }
+
+    private DiagnosisHarnessCore core() {
+        return new DiagnosisHarnessCore(
+                Clock.systemUTC(), () -> "unused", Duration.ofMinutes(5),
+                new RunBudgetLimits(20, 10, 10, 100_000, 100_000, 200_000, 1_000_000),
+                HarnessRetryPolicies.strict());
+    }
+}