feat(trace): add run-scoped trace reads
This commit is contained in:
zhuyongxin
@@ -2,11 +2,15 @@ package com.superbiz.agent.service;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
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.DiagnosisSession;
|
||||
import com.superbiz.agent.domain.entity.ToolInvocation;
|
||||
import com.superbiz.agent.dto.DiagnosisTraceResponse;
|
||||
import com.superbiz.agent.exception.SessionNotFoundException;
|
||||
import com.superbiz.agent.repository.AgentStepRepository;
|
||||
import com.superbiz.agent.repository.ChatSessionRepository;
|
||||
import com.superbiz.agent.repository.DiagnosisRunRepository;
|
||||
import com.superbiz.agent.repository.DiagnosisSessionRepository;
|
||||
import com.superbiz.agent.repository.ToolInvocationRepository;
|
||||
import org.junit.jupiter.api.Test;
|
||||
@@ -15,148 +19,245 @@ import java.time.LocalDateTime;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
import static org.mockito.Mockito.*;
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertNull;
|
||||
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.verifyNoInteractions;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
class DiagnosisTraceServiceTest {
|
||||
|
||||
private final DiagnosisSessionRepository diagnosisSessionRepository = mock(DiagnosisSessionRepository.class);
|
||||
private final ChatSessionRepository chatSessionRepository = mock(ChatSessionRepository.class);
|
||||
private final DiagnosisRunRepository diagnosisRunRepository = mock(DiagnosisRunRepository.class);
|
||||
private final AgentStepRepository agentStepRepository = mock(AgentStepRepository.class);
|
||||
private final ToolInvocationRepository toolInvocationRepository = mock(ToolInvocationRepository.class);
|
||||
private final DiagnosisTraceService service = new DiagnosisTraceService(
|
||||
diagnosisSessionRepository,
|
||||
chatSessionRepository,
|
||||
diagnosisRunRepository,
|
||||
agentStepRepository,
|
||||
toolInvocationRepository,
|
||||
new ObjectMapper()
|
||||
);
|
||||
|
||||
@Test
|
||||
void getTraceAggregatesSessionStepsAndTools() {
|
||||
String sessionId = "trace-session-001";
|
||||
LocalDateTime now = LocalDateTime.of(2026, 7, 3, 14, 30);
|
||||
DiagnosisSession session = DiagnosisSession.builder()
|
||||
.id(1L)
|
||||
.sessionId(sessionId)
|
||||
.query("payment timeout")
|
||||
.status("SUCCESS")
|
||||
.agentFlow("COMPLEX")
|
||||
.totalDurationMs(1200)
|
||||
.totalTokenCount(300)
|
||||
.stepCount(2)
|
||||
.toolCallCount(1)
|
||||
.answer("restart payment gateway pool")
|
||||
.selfEvaluation("{\"verifier_evaluation\":{\"verdict\":\"PASS\"},\"aiops_rule_evaluation\":{\"verdict\":\"WARN\"}}")
|
||||
.feedback("useful")
|
||||
.createdAt(now)
|
||||
.updatedAt(now)
|
||||
.build();
|
||||
AgentStep step = AgentStep.builder()
|
||||
.id(10L)
|
||||
.sessionId(sessionId)
|
||||
.stepIndex(1)
|
||||
.agentName("chat_executor")
|
||||
.modelInput("input")
|
||||
.modelOutput("output")
|
||||
.thought("executor finished")
|
||||
.hasToolCall(true)
|
||||
.durationMs(500)
|
||||
.tokenCount(100)
|
||||
.createdAt(now)
|
||||
.build();
|
||||
ToolInvocation invocation = ToolInvocation.builder()
|
||||
.id(20L)
|
||||
.sessionId(sessionId)
|
||||
.stepId(10L)
|
||||
.toolName("lookup_knowledge")
|
||||
.inputParams("{\"query\":\"ERR_TIMEOUT\"}")
|
||||
.outputPreview("payment timeout doc")
|
||||
.outputLength(19)
|
||||
.retrievalLayer("L0")
|
||||
.l0MatchCount(1)
|
||||
.l1MatchCount(0)
|
||||
.isTruncated(false)
|
||||
.relevanceLevel("HIGHLY_RELEVANT")
|
||||
.dedupReason("FIRST_HIT")
|
||||
.retrievalDetails("{\"documents\":[\"payment-errors.md\"]}")
|
||||
.durationMs(80)
|
||||
.success(true)
|
||||
.createdAt(now)
|
||||
.build();
|
||||
void getTraceWithoutRunIdResolvesLatestRun() {
|
||||
String sessionId = "trace-session-latest";
|
||||
LocalDateTime base = LocalDateTime.of(2026, 7, 10, 10, 0);
|
||||
DiagnosisRun latest = run(2L, sessionId, "run-latest", "second question", base.plusMinutes(1));
|
||||
AgentStep step = step(20L, sessionId, "run-latest", 0, "composer", base.plusMinutes(1));
|
||||
ToolInvocation invocation = invocation(30L, sessionId, "run-latest", "query_metrics", base.plusMinutes(1));
|
||||
|
||||
when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(session));
|
||||
when(agentStepRepository.findBySessionId(sessionId)).thenReturn(List.of(step));
|
||||
when(toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId)).thenReturn(List.of(invocation));
|
||||
when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
|
||||
.thenReturn(Optional.of(latest));
|
||||
when(chatSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(chatSession(sessionId)));
|
||||
when(agentStepRepository.findByRunIdOrderByStepIndex("run-latest")).thenReturn(List.of(step));
|
||||
when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-latest")).thenReturn(List.of(invocation));
|
||||
|
||||
DiagnosisTraceResponse response = service.getTrace(sessionId);
|
||||
|
||||
assertEquals(sessionId, response.getSession().getSessionId());
|
||||
assertEquals("payment timeout", response.getSession().getQuery());
|
||||
assertEquals("PASS", ((java.util.Map<?, ?>) response.getSession()
|
||||
.getSelfEvaluation()
|
||||
.get("verifier_evaluation")).get("verdict"));
|
||||
assertEquals(1, response.getSteps().size());
|
||||
assertEquals("chat_executor", response.getSteps().get(0).getAgentName());
|
||||
assertEquals(1, response.getToolInvocations().size());
|
||||
assertEquals("ERR_TIMEOUT", response.getToolInvocations().get(0).getInputParams().get("query"));
|
||||
assertEquals(2, response.getSummary().getPersistedStepCount());
|
||||
assertEquals(1, response.getSummary().getReturnedStepCount());
|
||||
assertEquals(1, response.getSummary().getPersistedToolCallCount());
|
||||
assertEquals(1, response.getSummary().getReturnedToolCallCount());
|
||||
assertTrue(response.getSummary().isHasVerifierEvaluation());
|
||||
assertTrue(response.getSummary().isHasAiOpsRuleEvaluation());
|
||||
assertTrue(response.getSummary().isHasFeedback());
|
||||
assertEquals("run-latest", response.getRunId());
|
||||
assertEquals("run-latest", response.getRun().getRunId());
|
||||
assertEquals("run-latest", response.getSummary().getResolvedRunId());
|
||||
assertEquals("second question", response.getSession().getQuery());
|
||||
assertEquals("run-latest", response.getSteps().get(0).getRunId());
|
||||
assertEquals("run-latest", response.getToolInvocations().get(0).getRunId());
|
||||
assertEquals(1, response.getChatSession().getMessagePairCount());
|
||||
}
|
||||
|
||||
@Test
|
||||
void getTraceOrdersStepsByCreationTimeAndIdNotPerAgentStepIndex() {
|
||||
String sessionId = "trace-session-ordered";
|
||||
LocalDateTime base = LocalDateTime.of(2026, 7, 7, 0, 0);
|
||||
DiagnosisSession session = DiagnosisSession.builder()
|
||||
.id(1L)
|
||||
.sessionId(sessionId)
|
||||
.query("mysql pool issue")
|
||||
.status("SUCCESS")
|
||||
.stepCount(4)
|
||||
.toolCallCount(0)
|
||||
.createdAt(base)
|
||||
.updatedAt(base)
|
||||
.build();
|
||||
AgentStep planner = step(1L, sessionId, 0, "planner", base.plusSeconds(1));
|
||||
AgentStep executor0 = step(2L, sessionId, 0, "executor", base.plusSeconds(2));
|
||||
AgentStep executor1 = step(3L, sessionId, 1, "executor", base.plusSeconds(3));
|
||||
AgentStep verifier = step(4L, sessionId, 0, "verifier", base.plusSeconds(4));
|
||||
void getTraceWithRunIdReturnsExactFirstRun() {
|
||||
String sessionId = "trace-session-exact";
|
||||
DiagnosisRun first = run(1L, sessionId, "run-first", "first question",
|
||||
LocalDateTime.of(2026, 7, 10, 10, 0));
|
||||
AgentStep firstStep = step(10L, sessionId, "run-first", 0, "planner", first.getCreatedAt());
|
||||
|
||||
when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(session));
|
||||
when(agentStepRepository.findBySessionId(sessionId))
|
||||
.thenReturn(List.of(planner, executor0, verifier, executor1));
|
||||
when(toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId)).thenReturn(List.of());
|
||||
when(diagnosisRunRepository.findBySessionIdAndRunId(sessionId, "run-first"))
|
||||
.thenReturn(Optional.of(first));
|
||||
when(agentStepRepository.findByRunIdOrderByStepIndex("run-first")).thenReturn(List.of(firstStep));
|
||||
when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-first")).thenReturn(List.of());
|
||||
|
||||
DiagnosisTraceResponse response = service.getTrace(sessionId);
|
||||
DiagnosisTraceResponse response = service.getTrace(sessionId, "run-first");
|
||||
|
||||
assertEquals(List.of("planner", "executor", "executor", "verifier"),
|
||||
assertEquals("run-first", response.getRunId());
|
||||
assertEquals("first question", response.getRun().getQuery());
|
||||
assertEquals(List.of("planner"),
|
||||
response.getSteps().stream().map(DiagnosisTraceResponse.AgentStepTrace::getAgentName).toList());
|
||||
assertEquals(List.of(0, 0, 1, 0),
|
||||
response.getSteps().stream().map(DiagnosisTraceResponse.AgentStepTrace::getStepIndex).toList());
|
||||
verify(diagnosisRunRepository, never()).findFirstBySessionIdOrderByCreatedAtDescIdDesc(any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void getTraceWithRunIdReturnsExactSecondRun() {
|
||||
String sessionId = "trace-session-exact";
|
||||
DiagnosisRun second = run(2L, sessionId, "run-second", "second question",
|
||||
LocalDateTime.of(2026, 7, 10, 10, 1));
|
||||
ToolInvocation secondTool = invocation(20L, sessionId, "run-second", "query_logs", second.getCreatedAt());
|
||||
|
||||
when(diagnosisRunRepository.findBySessionIdAndRunId(sessionId, "run-second"))
|
||||
.thenReturn(Optional.of(second));
|
||||
when(agentStepRepository.findByRunIdOrderByStepIndex("run-second")).thenReturn(List.of());
|
||||
when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-second")).thenReturn(List.of(secondTool));
|
||||
|
||||
DiagnosisTraceResponse response = service.getTrace(sessionId, "run-second");
|
||||
|
||||
assertEquals("run-second", response.getRunId());
|
||||
assertEquals("second question", response.getSession().getQuery());
|
||||
assertEquals(List.of("query_logs"),
|
||||
response.getToolInvocations().stream().map(DiagnosisTraceResponse.ToolInvocationTrace::getToolName).toList());
|
||||
}
|
||||
|
||||
@Test
|
||||
void getTraceRejectsRunFromAnotherSession() {
|
||||
String runId = "run-other-session";
|
||||
when(diagnosisRunRepository.findBySessionIdAndRunId("path-session", runId)).thenReturn(Optional.empty());
|
||||
when(diagnosisRunRepository.findByRunId(runId)).thenReturn(Optional.of(run(
|
||||
1L, "actual-session", runId, "query", LocalDateTime.of(2026, 7, 10, 10, 0))));
|
||||
|
||||
IllegalArgumentException error = assertThrows(IllegalArgumentException.class,
|
||||
() -> service.getTrace("path-session", runId));
|
||||
|
||||
assertTrue(error.getMessage().contains("does not belong"));
|
||||
verifyNoInteractions(agentStepRepository, toolInvocationRepository);
|
||||
}
|
||||
|
||||
@Test
|
||||
void getTraceThrowsWhenSessionMissing() {
|
||||
String sessionId = "missing-session";
|
||||
when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
|
||||
.thenReturn(Optional.empty());
|
||||
when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.empty());
|
||||
|
||||
assertThrows(SessionNotFoundException.class, () -> service.getTrace(sessionId));
|
||||
|
||||
verify(diagnosisRunRepository).findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId);
|
||||
verify(diagnosisSessionRepository).findBySessionId(sessionId);
|
||||
verifyNoInteractions(agentStepRepository, toolInvocationRepository);
|
||||
}
|
||||
|
||||
private AgentStep step(Long id, String sessionId, int stepIndex, String agentName, LocalDateTime createdAt) {
|
||||
@Test
|
||||
void getTraceIsReadOnly() {
|
||||
String sessionId = "trace-session-readonly";
|
||||
DiagnosisRun latest = run(3L, sessionId, "run-readonly", "readonly question",
|
||||
LocalDateTime.of(2026, 7, 10, 10, 2));
|
||||
|
||||
when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
|
||||
.thenReturn(Optional.of(latest));
|
||||
when(agentStepRepository.findByRunIdOrderByStepIndex("run-readonly")).thenReturn(List.of());
|
||||
when(toolInvocationRepository.findByRunIdOrderByIdAsc("run-readonly")).thenReturn(List.of());
|
||||
|
||||
service.getTrace(sessionId);
|
||||
|
||||
verify(chatSessionRepository, never()).save(any());
|
||||
verify(diagnosisRunRepository, never()).save(any());
|
||||
verify(diagnosisSessionRepository, never()).save(any());
|
||||
verify(agentStepRepository, never()).save(any());
|
||||
verify(toolInvocationRepository, never()).save(any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void listRunSummariesDoesNotExpandTraceDetails() {
|
||||
String sessionId = "trace-session-runs";
|
||||
DiagnosisRun second = run(2L, sessionId, "run-second", "second question",
|
||||
LocalDateTime.of(2026, 7, 10, 10, 1));
|
||||
second.setAnswer("answer ".repeat(40));
|
||||
DiagnosisRun first = run(1L, sessionId, "run-first", "first question",
|
||||
LocalDateTime.of(2026, 7, 10, 10, 0));
|
||||
|
||||
when(diagnosisRunRepository.findBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
|
||||
.thenReturn(List.of(second, first));
|
||||
|
||||
List<DiagnosisTraceResponse.RunSummary> summaries = service.listRunSummaries(sessionId);
|
||||
|
||||
assertEquals(List.of("run-second", "run-first"),
|
||||
summaries.stream().map(DiagnosisTraceResponse.RunSummary::getRunId).toList());
|
||||
assertEquals("second question", summaries.get(0).getQuery());
|
||||
assertTrue(summaries.get(0).getAnswerPreview().length() <= 160);
|
||||
verifyNoInteractions(agentStepRepository, toolInvocationRepository);
|
||||
}
|
||||
|
||||
@Test
|
||||
void legacyTraceFallbackKeepsHistoricalSessionReadable() {
|
||||
String sessionId = "legacy-session";
|
||||
DiagnosisSession legacy = DiagnosisSession.builder()
|
||||
.id(1L)
|
||||
.sessionId(sessionId)
|
||||
.query("legacy question")
|
||||
.status("SUCCESS")
|
||||
.stepCount(0)
|
||||
.toolCallCount(0)
|
||||
.createdAt(LocalDateTime.of(2026, 7, 10, 9, 0))
|
||||
.updatedAt(LocalDateTime.of(2026, 7, 10, 9, 1))
|
||||
.build();
|
||||
|
||||
when(diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId))
|
||||
.thenReturn(Optional.empty());
|
||||
when(diagnosisSessionRepository.findBySessionId(sessionId)).thenReturn(Optional.of(legacy));
|
||||
when(agentStepRepository.findBySessionId(sessionId)).thenReturn(List.of());
|
||||
when(toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId)).thenReturn(List.of());
|
||||
|
||||
DiagnosisTraceResponse response = service.getTrace(sessionId);
|
||||
|
||||
assertNull(response.getRunId());
|
||||
assertNull(response.getRun());
|
||||
assertEquals("legacy question", response.getSession().getQuery());
|
||||
}
|
||||
|
||||
private DiagnosisRun run(Long id, String sessionId, String runId, String query, LocalDateTime createdAt) {
|
||||
return DiagnosisRun.builder()
|
||||
.id(id)
|
||||
.sessionId(sessionId)
|
||||
.runId(runId)
|
||||
.query(query)
|
||||
.status("SUCCESS")
|
||||
.agentFlow("CHAT")
|
||||
.answer("answer for " + runId)
|
||||
.selfEvaluation("{\"verifier_evaluation\":{\"verdict\":\"PASS\"}}")
|
||||
.feedback("useful")
|
||||
.stepCount(1)
|
||||
.toolCallCount(1)
|
||||
.createdAt(createdAt)
|
||||
.updatedAt(createdAt.plusSeconds(1))
|
||||
.build();
|
||||
}
|
||||
|
||||
private ChatSession chatSession(String sessionId) {
|
||||
return ChatSession.builder()
|
||||
.id(99L)
|
||||
.sessionId(sessionId)
|
||||
.status("ACTIVE")
|
||||
.messagePairCount(1)
|
||||
.createdAt(LocalDateTime.of(2026, 7, 10, 9, 0))
|
||||
.lastActiveAt(LocalDateTime.of(2026, 7, 10, 10, 0))
|
||||
.build();
|
||||
}
|
||||
|
||||
private AgentStep step(Long id, String sessionId, String runId, int stepIndex, String agentName, LocalDateTime createdAt) {
|
||||
return AgentStep.builder()
|
||||
.id(id)
|
||||
.sessionId(sessionId)
|
||||
.runId(runId)
|
||||
.stepIndex(stepIndex)
|
||||
.agentName(agentName)
|
||||
.createdAt(createdAt)
|
||||
.build();
|
||||
}
|
||||
|
||||
private ToolInvocation invocation(Long id, String sessionId, String runId, String toolName, LocalDateTime createdAt) {
|
||||
return ToolInvocation.builder()
|
||||
.id(id)
|
||||
.sessionId(sessionId)
|
||||
.runId(runId)
|
||||
.toolName(toolName)
|
||||
.inputParams("{\"query\":\"timeout\"}")
|
||||
.retrievalDetails("{\"evidence_refs\":[]}")
|
||||
.success(true)
|
||||
.createdAt(createdAt)
|
||||
.build();
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user