feat(trace): add session run isolation schema

2026-07-10 17:47:56 +08:00
parent 841437fa06
commit 52bf0302c6
21 changed files with 1718 additions and 1 deletions
@@ -0,0 +1,112 @@
+package com.superbiz.agent.repository;
+
+import com.superbiz.agent.domain.entity.AgentStep;
+import com.superbiz.agent.domain.entity.ChatSession;
+import com.superbiz.agent.domain.entity.DiagnosisRun;
+import com.superbiz.agent.domain.entity.ToolInvocation;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.autoconfigure.jdbc.AutoConfigureTestDatabase;
+import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
+import org.springframework.test.context.TestPropertySource;
+
+import java.util.List;
+import java.util.UUID;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+@DataJpaTest
+@AutoConfigureTestDatabase(replace = AutoConfigureTestDatabase.Replace.NONE)
+@TestPropertySource(properties = {
+    "spring.flyway.enabled=true",
+    "spring.jpa.hibernate.ddl-auto=validate",
+    "spring.jpa.show-sql=true"
+})
+class DiagnosisRunRepositoryTest {
+
+    @Autowired
+    private ChatSessionRepository chatSessionRepository;
+
+    @Autowired
+    private DiagnosisRunRepository diagnosisRunRepository;
+
+    @Autowired
+    private AgentStepRepository agentStepRepository;
+
+    @Autowired
+    private ToolInvocationRepository toolInvocationRepository;
+
+    @Test
+    void saveAndFindLatestRunBySessionId() {
+        String sessionId = "test-session-" + UUID.randomUUID();
+        DiagnosisRun firstRun = saveRun(sessionId, "run-" + UUID.randomUUID(), "first question");
+        DiagnosisRun secondRun = saveRun(sessionId, "run-" + UUID.randomUUID(), "second question");
+
+        DiagnosisRun latest = diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId)
+                .orElseThrow();
+
+        assertEquals(secondRun.getRunId(), latest.getRunId());
+        assertEquals(firstRun.getRunId(), diagnosisRunRepository.findByRunId(firstRun.getRunId()).orElseThrow().getRunId());
+        assertTrue(diagnosisRunRepository.findBySessionIdAndRunId(sessionId, secondRun.getRunId()).isPresent());
+        assertEquals(2, diagnosisRunRepository.findBySessionIdOrderByCreatedAtDescIdDesc(sessionId).size());
+    }
+
+    @Test
+    void stepAndToolCanBeQueriedByRunId() {
+        String sessionId = "test-session-" + UUID.randomUUID();
+        String runId = "run-" + UUID.randomUUID();
+        saveRun(sessionId, runId, "run scoped trace");
+
+        agentStepRepository.save(AgentStep.builder()
+                .sessionId(sessionId)
+                .runId(runId)
+                .stepIndex(1)
+                .agentName("executor")
+                .hasToolCall(true)
+                .build());
+        agentStepRepository.save(AgentStep.builder()
+                .sessionId(sessionId)
+                .runId(runId)
+                .stepIndex(0)
+                .agentName("planner")
+                .hasToolCall(false)
+                .build());
+
+        toolInvocationRepository.save(ToolInvocation.builder()
+                .sessionId(sessionId)
+                .runId(runId)
+                .toolName("lookup_knowledge")
+                .inputParams("{\"query\":\"payment timeout\"}")
+                .success(true)
+                .build());
+
+        List<AgentStep> steps = agentStepRepository.findByRunIdOrderByStepIndex(runId);
+        List<ToolInvocation> tools = toolInvocationRepository.findByRunIdOrderByIdAsc(runId);
+
+        assertEquals(2, steps.size());
+        assertEquals("planner", steps.get(0).getAgentName());
+        assertEquals(1, tools.size());
+        assertEquals(runId, tools.get(0).getRunId());
+        assertEquals(2, agentStepRepository.countByRunId(runId));
+        assertEquals(1, toolInvocationRepository.countByRunId(runId));
+    }
+
+    private DiagnosisRun saveRun(String sessionId, String runId, String query) {
+        chatSessionRepository.findBySessionId(sessionId)
+                .orElseGet(() -> chatSessionRepository.save(ChatSession.builder()
+                        .sessionId(sessionId)
+                        .status("ACTIVE")
+                        .messagePairCount(0)
+                        .build()));
+
+        return diagnosisRunRepository.save(DiagnosisRun.builder()
+                .sessionId(sessionId)
+                .runId(runId)
+                .query(query)
+                .status("SUCCESS")
+                .agentFlow("CHAT")
+                .answer("answer for " + query)
+                .build());
+    }
+}
+