264 lines
12 KiB
Java
264 lines
12 KiB
Java
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;
|
|
|
|
import java.time.LocalDateTime;
|
|
import java.util.List;
|
|
import java.util.Optional;
|
|
|
|
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 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(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("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 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(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, "run-first");
|
|
|
|
assertEquals("run-first", response.getRunId());
|
|
assertEquals("first question", response.getRun().getQuery());
|
|
assertEquals(List.of("planner"),
|
|
response.getSteps().stream().map(DiagnosisTraceResponse.AgentStepTrace::getAgentName).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);
|
|
}
|
|
|
|
@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();
|
|
}
|
|
}
|