osiman/SuperBizAgent-java
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(trace): add run-scoped trace reads

Browse Source
This commit is contained in:
zhuyongxin
2026-07-10 20:07:23 +08:00
parent 26d5529280
commit 027aed1eeb
11 changed files with 620 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main/java/com/superbiz/agent/controller/ChatController.java
+11
View File
@@ -5,8 +5,10 @@ import lombok.Getter;
import lombok.Setter;
import com.superbiz.agent.domain.model.SessionContext;
import com.superbiz.agent.dto.AIOpsRequest;
import com.superbiz.agent.dto.DiagnosisTraceResponse;
import com.superbiz.agent.service.AiOpsService;
import com.superbiz.agent.service.ChatService;
import com.superbiz.agent.service.DiagnosisTraceService;
import com.superbiz.agent.service.session.SessionManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -43,6 +45,9 @@ public class ChatController {
@Autowired
private ChatService chatService;
@Autowired
private DiagnosisTraceService diagnosisTraceService;
@Autowired
private SessionManager sessionManager;
@@ -325,6 +330,12 @@ public class ChatController {
}
}
@GetMapping("/chat/session/{sessionId}/runs")
public ResponseEntity<ApiResponse<List<DiagnosisTraceResponse.RunSummary>>> listSessionRuns(
@PathVariable String sessionId) {
return ResponseEntity.ok(ApiResponse.success(diagnosisTraceService.listRunSummaries(sessionId)));
}
// ==================== 辅助方法 ====================
private SessionContext getOrCreateSession(String sessionId) {
src/main/java/com/superbiz/agent/controller/DiagnosisTraceController.java
+6 -2
View File
@@ -7,6 +7,7 @@ import lombok.RequiredArgsConstructor;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
@@ -18,7 +19,10 @@ public class DiagnosisTraceController {
private final DiagnosisTraceService diagnosisTraceService;
@GetMapping("/{sessionId}/trace")
public ResponseEntity<Result<DiagnosisTraceResponse>> getTrace(@PathVariable String sessionId) {
return ResponseEntity.ok(Result.success(diagnosisTraceService.getTrace(sessionId)));
public ResponseEntity<Result<DiagnosisTraceResponse>> getTrace(
@PathVariable String sessionId,
@RequestParam(required = false) String runId
) {
return ResponseEntity.ok(Result.success(diagnosisTraceService.getTrace(sessionId, runId)));
}
}
src/main/java/com/superbiz/agent/dto/DiagnosisTraceResponse.java
+60
View File
@@ -15,11 +15,28 @@ import java.util.Map;
@AllArgsConstructor
public class DiagnosisTraceResponse {
private String runId;
private ChatSessionTrace chatSession;
private SessionTrace session;
private RunTrace run;
private List<AgentStepTrace> steps;
private List<ToolInvocationTrace> toolInvocations;
private TraceSummary summary;
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public static class ChatSessionTrace {
private Long id;
private String sessionId;
private String status;
private Integer messagePairCount;
private LocalDateTime createdAt;
private LocalDateTime lastActiveAt;
private LocalDateTime expiresAt;
}
@Data
@Builder
@NoArgsConstructor
@@ -42,6 +59,29 @@ public class DiagnosisTraceResponse {
private LocalDateTime updatedAt;
}
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public static class RunTrace {
private Long id;
private String runId;
private String sessionId;
private String query;
private String status;
private String agentFlow;
private Integer totalDurationMs;
private Integer totalTokenCount;
private Integer stepCount;
private Integer toolCallCount;
private String answer;
private String selfEvaluationRaw;
private Map<String, Object> selfEvaluation;
private String feedback;
private LocalDateTime createdAt;
private LocalDateTime updatedAt;
}
@Data
@Builder
@NoArgsConstructor
@@ -49,6 +89,7 @@ public class DiagnosisTraceResponse {
public static class AgentStepTrace {
private Long id;
private String sessionId;
private String runId;
private Integer stepIndex;
private String agentName;
private String modelInput;
@@ -67,6 +108,7 @@ public class DiagnosisTraceResponse {
public static class ToolInvocationTrace {
private Long id;
private String sessionId;
private String runId;
private Long stepId;
private String toolName;
private String inputParamsRaw;
@@ -92,6 +134,7 @@ public class DiagnosisTraceResponse {
@NoArgsConstructor
@AllArgsConstructor
public static class TraceSummary {
private String resolvedRunId;
private int persistedStepCount;
private int returnedStepCount;
private int persistedToolCallCount;
@@ -100,4 +143,21 @@ public class DiagnosisTraceResponse {
private boolean hasAiOpsRuleEvaluation;
private boolean hasFeedback;
}
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public static class RunSummary {
private String runId;
private String sessionId;
private String query;
private String status;
private String agentFlow;
private String answerPreview;
private Integer stepCount;
private Integer toolCallCount;
private LocalDateTime createdAt;
private LocalDateTime updatedAt;
}
}
src/main/java/com/superbiz/agent/service/DiagnosisTraceService.java
+161
View File
@@ -3,11 +3,15 @@ package com.superbiz.agent.service;
import com.fasterxml.jackson.core.type.TypeReference;
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 lombok.RequiredArgsConstructor;
@@ -25,18 +29,76 @@ public class DiagnosisTraceService {
};
private final DiagnosisSessionRepository diagnosisSessionRepository;
private final ChatSessionRepository chatSessionRepository;
private final DiagnosisRunRepository diagnosisRunRepository;
private final AgentStepRepository agentStepRepository;
private final ToolInvocationRepository toolInvocationRepository;
private final ObjectMapper objectMapper;
public DiagnosisTraceResponse getTrace(String sessionId) {
return getTrace(sessionId, null);
}
public DiagnosisTraceResponse getTrace(String sessionId, String runId) {
if (runId != null && !runId.isBlank()) {
DiagnosisRun run = diagnosisRunRepository.findBySessionIdAndRunId(sessionId, runId)
.orElseThrow(() -> buildRunLookupException(sessionId, runId));
return buildRunTraceResponse(sessionId, run);
}
return diagnosisRunRepository.findFirstBySessionIdOrderByCreatedAtDescIdDesc(sessionId)
.map(run -> buildRunTraceResponse(sessionId, run))
.orElseGet(() -> buildLegacyTraceResponse(sessionId));
}
public List<DiagnosisTraceResponse.RunSummary> listRunSummaries(String sessionId) {
List<DiagnosisRun> runs = diagnosisRunRepository.findBySessionIdOrderByCreatedAtDescIdDesc(sessionId);
if (runs.isEmpty() && chatSessionRepository.findBySessionId(sessionId).isEmpty()) {
throw new SessionNotFoundException(sessionId);
}
return runs.stream()
.map(this::toRunSummary)
.toList();
}
private RuntimeException buildRunLookupException(String sessionId, String runId) {
if (diagnosisRunRepository.findByRunId(runId).isPresent()) {
return new IllegalArgumentException("runId does not belong to sessionId: " + runId);
}
return new SessionNotFoundException(sessionId, "Run not found: " + runId);
}
private DiagnosisTraceResponse buildRunTraceResponse(String sessionId, DiagnosisRun run) {
List<AgentStep> steps = orderStepsForTrace(agentStepRepository.findByRunIdOrderByStepIndex(run.getRunId()));
List<ToolInvocation> toolInvocations = toolInvocationRepository.findByRunIdOrderByIdAsc(run.getRunId());
DiagnosisTraceResponse.RunTrace runTrace = toRunTrace(run);
return DiagnosisTraceResponse.builder()
.runId(run.getRunId())
.chatSession(chatSessionRepository.findBySessionId(sessionId)
.map(this::toChatSessionTrace)
.orElse(null))
.session(toSessionTrace(runTrace))
.run(runTrace)
.steps(steps.stream().map(this::toAgentStepTrace).toList())
.toolInvocations(toolInvocations.stream().map(this::toToolInvocationTrace).toList())
.summary(toSummary(runTrace, steps, toolInvocations))
.build();
}
private DiagnosisTraceResponse buildLegacyTraceResponse(String sessionId) {
DiagnosisSession session = diagnosisSessionRepository.findBySessionId(sessionId)
.orElseThrow(() -> new SessionNotFoundException(sessionId));
List<AgentStep> steps = orderStepsForTrace(agentStepRepository.findBySessionId(sessionId));
List<ToolInvocation> toolInvocations = toolInvocationRepository.findBySessionIdOrderByIdAsc(sessionId);
return DiagnosisTraceResponse.builder()
.runId(null)
.chatSession(chatSessionRepository.findBySessionId(sessionId)
.map(this::toChatSessionTrace)
.orElse(null))
.session(toSessionTrace(session))
.run(null)
.steps(steps.stream().map(this::toAgentStepTrace).toList())
.toolInvocations(toolInvocations.stream().map(this::toToolInvocationTrace).toList())
.summary(toSummary(session, steps, toolInvocations))
@@ -52,6 +114,18 @@ public class DiagnosisTraceService {
.toList();
}
private DiagnosisTraceResponse.ChatSessionTrace toChatSessionTrace(ChatSession chatSession) {
return DiagnosisTraceResponse.ChatSessionTrace.builder()
.id(chatSession.getId())
.sessionId(chatSession.getSessionId())
.status(chatSession.getStatus())
.messagePairCount(chatSession.getMessagePairCount())
.createdAt(chatSession.getCreatedAt())
.lastActiveAt(chatSession.getLastActiveAt())
.expiresAt(chatSession.getExpiresAt())
.build();
}
private DiagnosisTraceResponse.SessionTrace toSessionTrace(DiagnosisSession session) {
return DiagnosisTraceResponse.SessionTrace.builder()
.id(session.getId())
@@ -72,10 +146,52 @@ public class DiagnosisTraceService {
.build();
}
private DiagnosisTraceResponse.SessionTrace toSessionTrace(DiagnosisTraceResponse.RunTrace run) {
return DiagnosisTraceResponse.SessionTrace.builder()
.id(run.getId())
.sessionId(run.getSessionId())
.query(run.getQuery())
.status(run.getStatus())
.agentFlow(run.getAgentFlow())
.totalDurationMs(run.getTotalDurationMs())
.totalTokenCount(run.getTotalTokenCount())
.stepCount(run.getStepCount())
.toolCallCount(run.getToolCallCount())
.answer(run.getAnswer())
.selfEvaluationRaw(run.getSelfEvaluationRaw())
.selfEvaluation(run.getSelfEvaluation())
.feedback(run.getFeedback())
.createdAt(run.getCreatedAt())
.updatedAt(run.getUpdatedAt())
.build();
}
private DiagnosisTraceResponse.RunTrace toRunTrace(DiagnosisRun run) {
return DiagnosisTraceResponse.RunTrace.builder()
.id(run.getId())
.runId(run.getRunId())
.sessionId(run.getSessionId())
.query(run.getQuery())
.status(run.getStatus())
.agentFlow(run.getAgentFlow())
.totalDurationMs(run.getTotalDurationMs())
.totalTokenCount(run.getTotalTokenCount())
.stepCount(run.getStepCount())
.toolCallCount(run.getToolCallCount())
.answer(run.getAnswer())
.selfEvaluationRaw(run.getSelfEvaluation())
.selfEvaluation(parseJsonObject(run.getSelfEvaluation()))
.feedback(run.getFeedback())
.createdAt(run.getCreatedAt())
.updatedAt(run.getUpdatedAt())
.build();
}
private DiagnosisTraceResponse.AgentStepTrace toAgentStepTrace(AgentStep step) {
return DiagnosisTraceResponse.AgentStepTrace.builder()
.id(step.getId())
.sessionId(step.getSessionId())
.runId(step.getRunId())
.stepIndex(step.getStepIndex())
.agentName(step.getAgentName())
.modelInput(step.getModelInput())
@@ -92,6 +208,7 @@ public class DiagnosisTraceService {
return DiagnosisTraceResponse.ToolInvocationTrace.builder()
.id(invocation.getId())
.sessionId(invocation.getSessionId())
.runId(invocation.getRunId())
.stepId(invocation.getStepId())
.toolName(invocation.getToolName())
.inputParamsRaw(invocation.getInputParams())
@@ -113,6 +230,21 @@ public class DiagnosisTraceService {
.build();
}
private DiagnosisTraceResponse.RunSummary toRunSummary(DiagnosisRun run) {
return DiagnosisTraceResponse.RunSummary.builder()
.runId(run.getRunId())
.sessionId(run.getSessionId())
.query(run.getQuery())
.status(run.getStatus())
.agentFlow(run.getAgentFlow())
.answerPreview(preview(run.getAnswer(), 160))
.stepCount(run.getStepCount())
.toolCallCount(run.getToolCallCount())
.createdAt(run.getCreatedAt())
.updatedAt(run.getUpdatedAt())
.build();
}
private DiagnosisTraceResponse.TraceSummary toSummary(
DiagnosisSession session,
List<AgentStep> steps,
@@ -120,6 +252,7 @@ public class DiagnosisTraceService {
) {
Map<String, Object> selfEvaluation = parseJsonObject(session.getSelfEvaluation());
return DiagnosisTraceResponse.TraceSummary.builder()
.resolvedRunId(null)
.persistedStepCount(defaultInt(session.getStepCount()))
.returnedStepCount(steps.size())
.persistedToolCallCount(defaultInt(session.getToolCallCount()))
@@ -130,6 +263,24 @@ public class DiagnosisTraceService {
.build();
}
private DiagnosisTraceResponse.TraceSummary toSummary(
DiagnosisTraceResponse.RunTrace run,
List<AgentStep> steps,
List<ToolInvocation> toolInvocations
) {
Map<String, Object> selfEvaluation = run.getSelfEvaluation();
return DiagnosisTraceResponse.TraceSummary.builder()
.resolvedRunId(run.getRunId())
.persistedStepCount(defaultInt(run.getStepCount()))
.returnedStepCount(steps.size())
.persistedToolCallCount(defaultInt(run.getToolCallCount()))
.returnedToolCallCount(toolInvocations.size())
.hasVerifierEvaluation(selfEvaluation != null && selfEvaluation.containsKey("verifier_evaluation"))
.hasAiOpsRuleEvaluation(selfEvaluation != null && selfEvaluation.containsKey("aiops_rule_evaluation"))
.hasFeedback(run.getFeedback() != null && !run.getFeedback().isBlank())
.build();
}
private Map<String, Object> parseJsonObject(String json) {
if (json == null || json.isBlank()) {
return null;
@@ -144,4 +295,14 @@ public class DiagnosisTraceService {
private int defaultInt(Integer value) {
return value == null ? 0 : value;
}
private String preview(String text, int maxLength) {
if (text == null) {
return null;
}
if (text.length() <= maxLength) {
return text;
}
return text.substring(0, maxLength);
}
}
src/test/java/com/superbiz/agent/service/DiagnosisTraceServiceTest.java
+199 -98
View File
@@ -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();
}
}