feat(harness): add canonical tool invocation boundary
This commit is contained in:
zhuyongxin
@@ -0,0 +1,215 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
import com.superbiz.agent.harness.contract.InvocationStatus;
|
||||
import com.superbiz.agent.harness.core.DiagnosisHarnessCore;
|
||||
import com.superbiz.agent.harness.core.HarnessCoreFixtures;
|
||||
import com.superbiz.agent.harness.core.MutableClock;
|
||||
import com.superbiz.agent.harness.core.RunContext;
|
||||
import com.superbiz.agent.harness.tool.store.CanonicalInvocationLimits;
|
||||
import com.superbiz.agent.harness.tool.store.CanonicalInvocationStore;
|
||||
import com.superbiz.agent.harness.tool.store.CanonicalToolInvocation;
|
||||
import com.superbiz.agent.harness.tool.store.DuplicateInvocationException;
|
||||
import com.superbiz.agent.harness.tool.store.ToolCallKeyFactory;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.time.Clock;
|
||||
import java.time.Duration;
|
||||
import java.time.Instant;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
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.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
class ToolBoundaryTest {
|
||||
|
||||
private final ObjectMapper objectMapper = new ObjectMapper().findAndRegisterModules();
|
||||
|
||||
@Test
|
||||
void storesRawAndReturnsOnlyProjectedReadyResult() {
|
||||
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
|
||||
FakeStore store = new FakeStore(new CanonicalInvocationLimits(Duration.ofHours(2), 4096, 1024));
|
||||
ToolBoundary boundary = boundary(clock, store);
|
||||
RunContext context = HarnessCoreFixtures.core(clock).startRun("session-1", "run-1");
|
||||
AtomicInteger executions = new AtomicInteger();
|
||||
|
||||
ToolBoundaryResult result = boundary.execute(
|
||||
context,
|
||||
request("run-1", "call-1", true, true),
|
||||
json -> {
|
||||
executions.incrementAndGet();
|
||||
return "{\"secret_raw\":true}";
|
||||
},
|
||||
raw -> new ProjectedToolResult("{\"evidence\":true}", EvidenceStatus.EVIDENCE_FOUND));
|
||||
|
||||
assertEquals(InvocationStatus.READY, result.status());
|
||||
assertEquals("call-1", result.toolCallId());
|
||||
assertNull(result.errorCode());
|
||||
assertFalse(String.valueOf(result.agentResult()).contains("secret_raw"));
|
||||
assertEquals(1, executions.get());
|
||||
CanonicalToolInvocation saved = store.find("superbiz:harness:tool-call:run-1:call-1").orElseThrow();
|
||||
assertEquals("{\"secret_raw\":true}", saved.rawResponse());
|
||||
assertTrue(saved.isReferencableBy("run-1"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void preservesNoEvidenceAndRejectsDuplicateAndCrossRun() {
|
||||
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
|
||||
FakeStore store = new FakeStore(new CanonicalInvocationLimits(Duration.ofHours(2), 4096, 1024));
|
||||
ToolBoundary boundary = boundary(clock, store);
|
||||
RunContext context = HarnessCoreFixtures.core(clock).startRun("session-1", "run-1");
|
||||
AtomicInteger executions = new AtomicInteger();
|
||||
ToolExecutor executor = request -> {
|
||||
executions.incrementAndGet();
|
||||
return "raw";
|
||||
};
|
||||
ToolResultProjector projector = raw -> new ProjectedToolResult(
|
||||
"{\"scope\":\"none\"}", EvidenceStatus.NO_EVIDENCE);
|
||||
|
||||
ToolBoundaryResult first = boundary.execute(context, request("run-1", "call-2", true, true), executor, projector);
|
||||
ToolBoundaryResult duplicate = boundary.execute(context, request("run-1", "call-2", true, true), executor, projector);
|
||||
ToolBoundaryResult crossRun = boundary.execute(context, request("other-run", "call-3", true, true), executor, projector);
|
||||
|
||||
assertEquals(EvidenceStatus.NO_EVIDENCE, first.evidenceStatus());
|
||||
assertEquals(ToolBoundaryErrorCode.DUPLICATE_TOOL_CALL.name(), duplicate.errorCode());
|
||||
assertEquals(ToolBoundaryErrorCode.RUN_MISMATCH.name(), crossRun.errorCode());
|
||||
assertEquals(1, executions.get());
|
||||
assertTrue(store.find("superbiz:harness:tool-call:run-1:call-2").orElseThrow().isReferencableBy("run-1"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void rejectsUnauthorizedWritableAndInvalidIdBeforeTool() {
|
||||
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
|
||||
FakeStore store = new FakeStore(new CanonicalInvocationLimits(Duration.ofHours(2), 4096, 1024));
|
||||
ToolBoundary boundary = boundary(clock, store);
|
||||
RunContext context = HarnessCoreFixtures.core(clock).startRun("session-1", "run-1");
|
||||
AtomicInteger executions = new AtomicInteger();
|
||||
ToolExecutor executor = request -> {
|
||||
executions.incrementAndGet();
|
||||
return "raw";
|
||||
};
|
||||
ToolResultProjector projector = raw -> new ProjectedToolResult("agent", EvidenceStatus.EVIDENCE_FOUND);
|
||||
|
||||
ToolBoundaryResult unauthorized = boundary.execute(context, request("run-1", "call-4", false, true), executor, projector);
|
||||
ToolBoundaryResult writable = boundary.execute(context, request("run-1", "call-5", true, false), executor, projector);
|
||||
ToolBoundaryResult invalidId = boundary.execute(context, request("run-1", "bad:id", true, true), executor, projector);
|
||||
|
||||
assertEquals(ToolBoundaryErrorCode.UNAUTHORIZED.name(), unauthorized.errorCode());
|
||||
assertEquals(ToolBoundaryErrorCode.NOT_READ_ONLY.name(), writable.errorCode());
|
||||
assertEquals(ToolBoundaryErrorCode.INVALID_TOOL_CALL_ID.name(), invalidId.errorCode());
|
||||
assertEquals(0, executions.get());
|
||||
}
|
||||
|
||||
@Test
|
||||
void rawOverflowSkipsProjectorAndAgentOverflowIsNotReturned() {
|
||||
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
|
||||
RunContext context = HarnessCoreFixtures.core(clock).startRun("session-1", "run-1");
|
||||
AtomicInteger projectorCalls = new AtomicInteger();
|
||||
FakeStore rawStore = new FakeStore(new CanonicalInvocationLimits(Duration.ofHours(2), 20, 10));
|
||||
ToolBoundary rawBoundary = boundary(clock, rawStore);
|
||||
ToolBoundaryResult rawOverflow = rawBoundary.execute(
|
||||
context, request("run-1", "call-6", true, true),
|
||||
requestJson -> "x".repeat(100),
|
||||
raw -> {
|
||||
projectorCalls.incrementAndGet();
|
||||
return new ProjectedToolResult("agent", EvidenceStatus.EVIDENCE_FOUND);
|
||||
});
|
||||
|
||||
FakeStore agentStore = new FakeStore(new CanonicalInvocationLimits(Duration.ofHours(2), 4096, 4));
|
||||
ToolBoundary agentBoundary = boundary(clock, agentStore);
|
||||
ToolBoundaryResult agentOverflow = agentBoundary.execute(
|
||||
context, request("run-1", "call-7", true, true),
|
||||
requestJson -> "raw",
|
||||
raw -> new ProjectedToolResult("too-large", EvidenceStatus.EVIDENCE_FOUND));
|
||||
|
||||
assertEquals(ToolBoundaryErrorCode.RESULT_TOO_LARGE.name(), rawOverflow.errorCode());
|
||||
assertEquals(0, projectorCalls.get());
|
||||
assertEquals(ToolBoundaryErrorCode.RESULT_TOO_LARGE.name(), agentOverflow.errorCode());
|
||||
assertNull(agentOverflow.agentResult());
|
||||
assertEquals(InvocationStatus.ERROR, agentStore.find("superbiz:harness:tool-call:run-1:call-7").orElseThrow().status());
|
||||
}
|
||||
|
||||
@Test
|
||||
void executionAndProjectionErrorsAreCanonicalErrors() {
|
||||
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
|
||||
FakeStore store = new FakeStore(new CanonicalInvocationLimits(Duration.ofHours(2), 4096, 1024));
|
||||
ToolBoundary boundary = boundary(clock, store);
|
||||
RunContext context = HarnessCoreFixtures.core(clock).startRun("session-1", "run-1");
|
||||
|
||||
ToolBoundaryResult executionError = boundary.execute(
|
||||
context, request("run-1", "call-8", true, true),
|
||||
requestJson -> { throw new IllegalStateException("internal raw error"); },
|
||||
raw -> new ProjectedToolResult("agent", EvidenceStatus.EVIDENCE_FOUND));
|
||||
ToolBoundaryResult projectionError = boundary.execute(
|
||||
context, request("run-1", "call-9", true, true),
|
||||
requestJson -> "raw",
|
||||
raw -> { throw new IllegalStateException("projection error"); });
|
||||
|
||||
assertEquals(ToolBoundaryErrorCode.TOOL_EXECUTION_ERROR.name(), executionError.errorCode());
|
||||
assertEquals(ToolBoundaryErrorCode.PROJECTION_ERROR.name(), projectionError.errorCode());
|
||||
assertFalse(store.find("superbiz:harness:tool-call:run-1:call-8").orElseThrow().isReferencableBy("run-1"));
|
||||
assertFalse(store.find("superbiz:harness:tool-call:run-1:call-9").orElseThrow().isReferencableBy("run-1"));
|
||||
}
|
||||
|
||||
private ToolBoundary boundary(MutableClock clock, FakeStore store) {
|
||||
DiagnosisHarnessCore core = HarnessCoreFixtures.core(clock);
|
||||
return new ToolBoundary(core, new ToolCallKeyFactory("superbiz:harness:tool-call"),
|
||||
store, objectMapper, clock);
|
||||
}
|
||||
|
||||
private ToolCallRequestEnvelope request(String runId, String toolCallId,
|
||||
boolean authorized, boolean readOnly) {
|
||||
return new ToolCallRequestEnvelope(runId, toolCallId, "query_logs", "{\"query\":\"timeout\"}",
|
||||
authorized, readOnly);
|
||||
}
|
||||
|
||||
private static final class FakeStore implements CanonicalInvocationStore {
|
||||
private final CanonicalInvocationLimits limits;
|
||||
private final Map<String, CanonicalToolInvocation> records = new HashMap<>();
|
||||
|
||||
private FakeStore(CanonicalInvocationLimits limits) {
|
||||
this.limits = limits;
|
||||
}
|
||||
|
||||
@Override
|
||||
public CanonicalInvocationLimits limits() {
|
||||
return limits;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void begin(String key, CanonicalToolInvocation invocation) {
|
||||
if (records.putIfAbsent(key, invocation) != null) {
|
||||
throw new DuplicateInvocationException();
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public Optional<CanonicalToolInvocation> find(String key) {
|
||||
return Optional.ofNullable(records.get(key));
|
||||
}
|
||||
|
||||
@Override
|
||||
public CanonicalToolInvocation markReady(String key, String rawResponse, String agentResult,
|
||||
EvidenceStatus evidenceStatus, Instant completedAt) {
|
||||
CanonicalToolInvocation current = records.get(key);
|
||||
CanonicalToolInvocation updated = current.markReady(rawResponse, agentResult, evidenceStatus, completedAt);
|
||||
records.put(key, updated);
|
||||
return updated;
|
||||
}
|
||||
|
||||
@Override
|
||||
public CanonicalToolInvocation markError(String key, String rawResponse, String errorCode,
|
||||
Instant completedAt) {
|
||||
CanonicalToolInvocation current = records.get(key);
|
||||
CanonicalToolInvocation updated = current.markError(rawResponse, errorCode, completedAt);
|
||||
records.put(key, updated);
|
||||
return updated;
|
||||
}
|
||||
}
|
||||
}
|
||||
+100
@@ -0,0 +1,100 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
import com.superbiz.agent.harness.contract.InvocationStatus;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.data.redis.core.RedisTemplate;
|
||||
import org.springframework.data.redis.core.ValueOperations;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.time.Instant;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertFalse;
|
||||
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
class CanonicalInvocationStoreTest {
|
||||
|
||||
private RedisTemplate<String, Object> redisTemplate;
|
||||
private ValueOperations<String, Object> values;
|
||||
private RedisCanonicalInvocationStore store;
|
||||
private CanonicalToolInvocation projecting;
|
||||
private final ObjectMapper objectMapper = new ObjectMapper().findAndRegisterModules();
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
redisTemplate = mock(RedisTemplate.class);
|
||||
values = mock(ValueOperations.class);
|
||||
when(redisTemplate.opsForValue()).thenReturn(values);
|
||||
store = new RedisCanonicalInvocationStore(
|
||||
redisTemplate,
|
||||
objectMapper,
|
||||
new CanonicalInvocationLimits(Duration.ofHours(2), 4096, 1024));
|
||||
projecting = CanonicalToolInvocation.projecting(
|
||||
"call-1", "run-1", "query_logs", "{\"query\":\"timeout\"}",
|
||||
Instant.parse("2026-07-21T10:00:00Z"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void createsCompleteProjectingRecordAndRejectsDuplicate() throws Exception {
|
||||
when(values.setIfAbsent(any(), any(), any(Long.class), any())).thenReturn(true, false);
|
||||
store.begin("prefix:run-1:call-1", projecting);
|
||||
|
||||
assertEquals(InvocationStatus.PROJECTING, projecting.status());
|
||||
assertFalse(projecting.request().isBlank());
|
||||
assertThrows(DuplicateInvocationException.class,
|
||||
() -> store.begin("prefix:run-1:call-1", projecting));
|
||||
verify(values, times(2)).setIfAbsent(
|
||||
eq("prefix:run-1:call-1"), any(), eq(Duration.ofHours(2).toMillis()), any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void updatesSameRecordToReadyWithRemainingTtlAndReadDoesNotRefresh() throws Exception {
|
||||
String projectingJson = objectMapper.writeValueAsString(projecting);
|
||||
when(values.get("key")).thenReturn(projectingJson, projectingJson);
|
||||
when(redisTemplate.getExpire("key", java.util.concurrent.TimeUnit.MILLISECONDS)).thenReturn(3210L);
|
||||
|
||||
CanonicalToolInvocation ready = store.markReady(
|
||||
"key", "{\"raw\":true}", "{\"evidence\":true}",
|
||||
EvidenceStatus.EVIDENCE_FOUND, Instant.parse("2026-07-21T10:00:01Z"));
|
||||
|
||||
assertEquals(InvocationStatus.READY, ready.status());
|
||||
assertEquals("call-1", ready.toolCallId());
|
||||
assertTrueJson(ready.agentResult());
|
||||
verify(values).set(eq("key"), any(), eq(3210L), eq(java.util.concurrent.TimeUnit.MILLISECONDS));
|
||||
|
||||
when(values.get("read-only")).thenReturn(projectingJson);
|
||||
store.find("read-only");
|
||||
verify(redisTemplate, never()).getExpire("read-only", java.util.concurrent.TimeUnit.MILLISECONDS);
|
||||
}
|
||||
|
||||
@Test
|
||||
void rejectsInvalidReadyEvidenceAndWritesErrorState() throws Exception {
|
||||
String projectingJson = objectMapper.writeValueAsString(projecting);
|
||||
when(values.get("key")).thenReturn(projectingJson, projectingJson);
|
||||
when(redisTemplate.getExpire("key", java.util.concurrent.TimeUnit.MILLISECONDS)).thenReturn(3210L);
|
||||
|
||||
assertThrows(InvocationStateException.class, () -> store.markReady(
|
||||
"key", "raw", "agent", EvidenceStatus.ERROR,
|
||||
Instant.parse("2026-07-21T10:00:01Z")));
|
||||
|
||||
CanonicalToolInvocation error = store.markError(
|
||||
"key", "raw", "PROJECTION_ERROR", Instant.parse("2026-07-21T10:00:02Z"));
|
||||
assertEquals(InvocationStatus.ERROR, error.status());
|
||||
assertEquals(EvidenceStatus.ERROR, error.evidenceStatus());
|
||||
assertFalse(error.isReferencableBy("run-1"));
|
||||
}
|
||||
|
||||
private void assertTrueJson(String json) throws Exception {
|
||||
assertEquals(true, objectMapper.readTree(json).path("evidence").asBoolean());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user