72 lines
3.3 KiB
Java
72 lines
3.3 KiB
Java
package com.superbiz.agent.harness.audit;
|
|
|
|
import com.alibaba.cloud.ai.graph.RunnableConfig;
|
|
import com.fasterxml.jackson.databind.ObjectMapper;
|
|
import com.superbiz.agent.domain.entity.AgentStep;
|
|
import com.superbiz.agent.harness.agent.DiagnosisAgentFactory;
|
|
import com.superbiz.agent.repository.AgentStepRepository;
|
|
import org.junit.jupiter.api.Test;
|
|
import org.mockito.ArgumentCaptor;
|
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
|
import org.springframework.ai.chat.messages.UserMessage;
|
|
|
|
import java.util.List;
|
|
import java.util.Optional;
|
|
|
|
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.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.when;
|
|
|
|
class HarnessAgentAuditHookTest {
|
|
|
|
@Test
|
|
void persistsExactIdentityAndMetadataWithoutContentOrArguments() {
|
|
AgentStepRepository repository = mock(AgentStepRepository.class);
|
|
AgentStep persisted = AgentStep.builder().id(7L).build();
|
|
when(repository.save(any(AgentStep.class))).thenReturn(persisted);
|
|
when(repository.findById(7L)).thenReturn(Optional.of(persisted));
|
|
HarnessAgentAuditHook hook = new HarnessAgentAuditHook(
|
|
repository, new ObjectMapper(), DiagnosisAgentFactory.AGENT_NAME);
|
|
RunnableConfig config = RunnableConfig.builder()
|
|
.addMetadata("sessionId", "session-audit")
|
|
.addMetadata("runId", "run-audit")
|
|
.build();
|
|
|
|
hook.beforeModel(List.of(new UserMessage("secret-query")), config);
|
|
AssistantMessage response = AssistantMessage.builder()
|
|
.content("secret-model-output")
|
|
.toolCalls(List.of(new AssistantMessage.ToolCall(
|
|
"call-1", "function", "query_logs", "{\"query\":\"secret-argument\"}")))
|
|
.build();
|
|
hook.afterModel(List.of(response), config);
|
|
|
|
ArgumentCaptor<AgentStep> captor = ArgumentCaptor.forClass(AgentStep.class);
|
|
verify(repository, org.mockito.Mockito.times(2)).save(captor.capture());
|
|
AgentStep started = captor.getAllValues().get(0);
|
|
AgentStep completed = captor.getAllValues().get(1);
|
|
assertEquals("session-audit", started.getSessionId());
|
|
assertEquals("run-audit", started.getRunId());
|
|
assertFalse(started.getModelInput().contains("secret-query"));
|
|
assertFalse(completed.getModelOutput().contains("secret-model-output"));
|
|
assertFalse(completed.getModelOutput().contains("secret-argument"));
|
|
assertEquals("{\"has_text\":true,\"tool_names\":[\"query_logs\"]}", completed.getModelOutput());
|
|
assertNull(completed.getThought());
|
|
}
|
|
|
|
@Test
|
|
void missingIdentitySkipsPersistence() {
|
|
AgentStepRepository repository = mock(AgentStepRepository.class);
|
|
HarnessAgentAuditHook hook = new HarnessAgentAuditHook(
|
|
repository, new ObjectMapper(), DiagnosisAgentFactory.AGENT_NAME);
|
|
|
|
hook.beforeModel(List.of(new UserMessage("secret")), RunnableConfig.builder().build());
|
|
|
|
verify(repository, never()).save(any());
|
|
}
|
|
}
|