SuperBizAgent-java/src/test/java/com/superbiz/agent/harness/audit/JpaToolInvocationAuditSinkTest.java

package com.superbiz.agent.harness.audit;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.domain.entity.ToolInvocation;
import com.superbiz.agent.harness.contract.EvidenceStatus;
import com.superbiz.agent.harness.contract.InvocationStatus;
import com.superbiz.agent.repository.ToolInvocationRepository;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;

class JpaToolInvocationAuditSinkTest {

    @Test
    void persistsOnlyBoundedStableMetadataForGenericTools() {
        ToolInvocationRepository repository = mock(ToolInvocationRepository.class);
        DiagnosisTraceRecorder traceRecorder = mock(DiagnosisTraceRecorder.class);
        JpaToolInvocationAuditSink sink = JpaToolInvocationAuditSink.forTest(
                repository, new ObjectMapper(), traceRecorder, null);

        sink.record(new ToolInvocationAuditEvent(
                "session-1", "run-1", "call-1", "query_logs",
                InvocationStatus.ERROR, EvidenceStatus.ERROR, "TOOL_EXECUTION_ERROR",
                17, 83, 0));

        ArgumentCaptor<ToolInvocation> captor = ArgumentCaptor.forClass(ToolInvocation.class);
        verify(repository).save(captor.capture());
        ToolInvocation saved = captor.getValue();
        assertEquals("session-1", saved.getSessionId());
        assertEquals("run-1", saved.getRunId());
        assertEquals("query_logs", saved.getToolName());
        assertEquals("{\"tool_call_id\":\"call-1\",\"request_bytes\":83}", saved.getInputParams());
        assertEquals("status=ERROR,evidence_status=ERROR", saved.getOutputPreview());
        assertEquals("TOOL_EXECUTION_ERROR", saved.getErrorMessage());
        // generic tools must not put evidence_status into relevance_level column
        assertNull(saved.getRelevanceLevel());
        assertNull(saved.getStepId());
        String serialized = saved.getInputParams() + saved.getOutputPreview() + saved.getRetrievalDetails();
        assertFalse(serialized.contains("raw_response"));
        assertTrue(saved.getRetrievalDetails().contains("evidence_status"));
        ArgumentCaptor<DiagnosisTraceAuditEvent> traceCaptor =
                ArgumentCaptor.forClass(DiagnosisTraceAuditEvent.class);
        verify(traceRecorder).record(traceCaptor.capture());
        assertEquals(TraceEventType.TOOL_INVOCATION, traceCaptor.getValue().eventType());
        assertEquals("call-1", traceCaptor.getValue().details().get("tool_call_id"));
    }

    @Test
    void persistsStepIdAndSafeQueryPreviewInInputParams() {
        ToolInvocationRepository repository = mock(ToolInvocationRepository.class);
        DiagnosisTraceRecorder traceRecorder = mock(DiagnosisTraceRecorder.class);
        JpaToolInvocationAuditSink sink = JpaToolInvocationAuditSink.forTest(
                repository, new ObjectMapper(), traceRecorder, null);

        String longQuery = "q".repeat(200);
        sink.record(new ToolInvocationAuditEvent(
                "session-1", "run-1", "call-q", "lookup_knowledge",
                InvocationStatus.READY, EvidenceStatus.EVIDENCE_FOUND, null,
                12, 50, 80, null, null, 99L,
                "{\"query\":\"" + longQuery + "\",\"password\":\"should-not-store\",\"limit\":5}"));

        ArgumentCaptor<ToolInvocation> captor = ArgumentCaptor.forClass(ToolInvocation.class);
        verify(repository).save(captor.capture());
        ToolInvocation saved = captor.getValue();
        assertEquals(99L, saved.getStepId());
        assertTrue(saved.getInputParams().contains("\"step_id\":99"));
        assertTrue(saved.getInputParams().contains("\"query\":\"" + "q".repeat(160) + "\""));
        assertTrue(saved.getInputParams().contains("\"query_truncated\":true"));
        assertTrue(saved.getInputParams().contains("\"query_chars\":200"));
        assertTrue(saved.getInputParams().contains("\"limit\":5"));
        assertFalse(saved.getInputParams().contains("should-not-store"));
        assertFalse(saved.getInputParams().contains("password"));

        ArgumentCaptor<DiagnosisTraceAuditEvent> traceCaptor =
                ArgumentCaptor.forClass(DiagnosisTraceAuditEvent.class);
        verify(traceRecorder).record(traceCaptor.capture());
        assertEquals(99L, traceCaptor.getValue().details().get("step_id"));
    }

    @Test
    void enrichesLookupKnowledgeWithRagFieldsAndTrueRelevanceLevel() {
        ToolInvocationRepository repository = mock(ToolInvocationRepository.class);
        DiagnosisTraceRecorder traceRecorder = mock(DiagnosisTraceRecorder.class);
        RagLookupAuditEnricher enricher = new RagLookupAuditEnricher(new ObjectMapper(), "hybrid");
        JpaToolInvocationAuditSink sink = JpaToolInvocationAuditSink.forTest(
                repository, new ObjectMapper(), traceRecorder, enricher);

        String lookup = """
                {
                  "found": true,
                  "relevanceLevel": "REFERENCE",
                  "evidenceBlockCount": 1,
                  "evidenceBlocks": [
                    {"evidenceKey":"doc#chunk-0","source":"doc","retrievalLayer":"L1","content":"body"}
                  ],
                  "retrievalTrace": {
                    "selectedAttempt": "FILTERED_VECTOR",
                    "fallbackReason": null,
                    "queryHints": {"l0_match_count": 2, "domains": ["mysql"]},
                    "attempts": [{"name":"FILTERED_VECTOR","candidateCount":3,"usable":true,"topSimilarity":0.8}]
                  }
                }
                """;
        String agent = """
                {"evidence_status":"EVIDENCE_FOUND","returned_count":1,"truncated":false,"relevance_level":"REFERENCE"}
                """;

        sink.record(new ToolInvocationAuditEvent(
                "session-1", "run-1", "call-rag", "lookup_knowledge",
                InvocationStatus.READY, EvidenceStatus.EVIDENCE_FOUND, null,
                42, 10, 100, lookup, agent, 7L,
                "{\"query\":\"MySQL HikariCP pool exhausted\"}"));

        ArgumentCaptor<ToolInvocation> captor = ArgumentCaptor.forClass(ToolInvocation.class);
        verify(repository).save(captor.capture());
        ToolInvocation saved = captor.getValue();
        assertEquals(7L, saved.getStepId());
        assertEquals("L1", saved.getRetrievalLayer());
        assertEquals("REFERENCE", saved.getRelevanceLevel());
        assertEquals(2, saved.getL0MatchCount());
        assertEquals(1, saved.getL1MatchCount());
        assertTrue(saved.getRetrievalDetails().contains("FILTERED_VECTOR"));
        assertTrue(saved.getRetrievalDetails().contains("search_mode"));
        assertTrue(saved.getRetrievalDetails().contains("EVIDENCE_FOUND"));
        assertTrue(saved.getOutputPreview().contains("REFERENCE"));
        assertTrue(saved.getInputParams().contains("MySQL HikariCP pool exhausted"));
        assertTrue(saved.getInputParams().contains("\"step_id\":7"));
        // must not store full excerpt dump as sole content; body may appear only if tiny — ensure attempt present
        assertTrue(saved.getRetrievalDetails().contains("attempts"));
    }
}