feat(harness): add canonical tool invocation boundary
This commit is contained in:
zhuyongxin
@@ -0,0 +1,19 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
|
||||
import java.util.Objects;
|
||||
|
||||
public record ProjectedToolResult(String agentResult, EvidenceStatus evidenceStatus) {
|
||||
|
||||
public ProjectedToolResult {
|
||||
if (agentResult == null || agentResult.isBlank()) {
|
||||
throw new IllegalArgumentException("agentResult must not be blank");
|
||||
}
|
||||
Objects.requireNonNull(evidenceStatus, "evidenceStatus must not be null");
|
||||
if (evidenceStatus != EvidenceStatus.EVIDENCE_FOUND
|
||||
&& evidenceStatus != EvidenceStatus.NO_EVIDENCE) {
|
||||
throw new IllegalArgumentException("projected result must be evidence or no-evidence");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,219 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.JsonNode;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
import com.superbiz.agent.harness.core.BudgetExceededException;
|
||||
import com.superbiz.agent.harness.core.DiagnosisHarnessCore;
|
||||
import com.superbiz.agent.harness.core.RunAbortedException;
|
||||
import com.superbiz.agent.harness.core.RunContext;
|
||||
import com.superbiz.agent.harness.tool.store.CanonicalInvocationStore;
|
||||
import com.superbiz.agent.harness.tool.store.CanonicalStoreException;
|
||||
import com.superbiz.agent.harness.tool.store.CanonicalToolInvocation;
|
||||
import com.superbiz.agent.harness.tool.store.DuplicateInvocationException;
|
||||
import com.superbiz.agent.harness.tool.store.ResultTooLargeException;
|
||||
import com.superbiz.agent.harness.tool.store.ToolCallKeyFactory;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import java.time.Clock;
|
||||
import java.util.Objects;
|
||||
|
||||
public final class ToolBoundary {
|
||||
|
||||
private static final Logger log = LoggerFactory.getLogger(ToolBoundary.class);
|
||||
|
||||
private final DiagnosisHarnessCore core;
|
||||
private final ToolCallKeyFactory keyFactory;
|
||||
private final CanonicalInvocationStore store;
|
||||
private final ObjectMapper objectMapper;
|
||||
private final Clock clock;
|
||||
|
||||
public ToolBoundary(DiagnosisHarnessCore core,
|
||||
ToolCallKeyFactory keyFactory,
|
||||
CanonicalInvocationStore store,
|
||||
ObjectMapper objectMapper,
|
||||
Clock clock) {
|
||||
this.core = Objects.requireNonNull(core, "core must not be null");
|
||||
this.keyFactory = Objects.requireNonNull(keyFactory, "keyFactory must not be null");
|
||||
this.store = Objects.requireNonNull(store, "store must not be null");
|
||||
this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null");
|
||||
this.clock = Objects.requireNonNull(clock, "clock must not be null");
|
||||
}
|
||||
|
||||
public ToolBoundaryResult execute(RunContext context,
|
||||
ToolCallRequestEnvelope request,
|
||||
ToolExecutor executor,
|
||||
ToolResultProjector projector) {
|
||||
String toolCallId = request == null ? null : request.toolCallId();
|
||||
String key;
|
||||
try {
|
||||
key = preflight(context, request);
|
||||
core.beforeToolCall(context, request.toolName());
|
||||
long requestBytes = store.limits().utf8Bytes(request.requestJson());
|
||||
if (requestBytes > store.limits().maxRecordBytes()) {
|
||||
return errorAndNoRecord(toolCallId, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
|
||||
}
|
||||
core.reserveRunBytes(context, requestBytes);
|
||||
store.begin(key, CanonicalToolInvocation.projecting(
|
||||
request.toolCallId(), request.runId(), request.toolName(),
|
||||
request.requestJson(), clock.instant()));
|
||||
} catch (DuplicateInvocationException e) {
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.DUPLICATE_TOOL_CALL);
|
||||
} catch (RunAbortedException | BudgetExceededException e) {
|
||||
return ToolBoundaryResult.error(toolCallId,
|
||||
e instanceof BudgetExceededException
|
||||
? ToolBoundaryErrorCode.BUDGET_EXHAUSTED
|
||||
: ToolBoundaryErrorCode.RUN_INACTIVE);
|
||||
} catch (IllegalArgumentException e) {
|
||||
return ToolBoundaryResult.error(toolCallId, classifyPreflightError(e));
|
||||
} catch (CanonicalStoreException e) {
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.STORE_ERROR);
|
||||
}
|
||||
|
||||
String rawResponse;
|
||||
try {
|
||||
rawResponse = Objects.requireNonNull(executor, "executor must not be null")
|
||||
.execute(request.requestJson());
|
||||
if (rawResponse == null) {
|
||||
throw new IllegalArgumentException("executor returned null");
|
||||
}
|
||||
} catch (Exception e) {
|
||||
markErrorSafely(key, null, ToolBoundaryErrorCode.TOOL_EXECUTION_ERROR);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.TOOL_EXECUTION_ERROR);
|
||||
}
|
||||
|
||||
try {
|
||||
store.limits().validateRawCandidate(request.requestJson(), rawResponse);
|
||||
core.reserveRunBytes(context, store.limits().utf8Bytes(rawResponse));
|
||||
} catch (ResultTooLargeException e) {
|
||||
markErrorSafely(key, null, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
|
||||
} catch (BudgetExceededException e) {
|
||||
markErrorSafely(key, null, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
|
||||
} catch (RunAbortedException e) {
|
||||
markErrorSafely(key, null, ToolBoundaryErrorCode.RUN_INACTIVE);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RUN_INACTIVE);
|
||||
}
|
||||
|
||||
ProjectedToolResult projected;
|
||||
try {
|
||||
projected = Objects.requireNonNull(projector, "projector must not be null")
|
||||
.project(rawResponse);
|
||||
if (projected == null) {
|
||||
throw new IllegalArgumentException("projector returned null");
|
||||
}
|
||||
} catch (Exception e) {
|
||||
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.PROJECTION_ERROR);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.PROJECTION_ERROR);
|
||||
}
|
||||
|
||||
try {
|
||||
store.limits().validateAgentResult(projected.agentResult());
|
||||
core.reserveRunBytes(context, store.limits().utf8Bytes(projected.agentResult()));
|
||||
store.markReady(key, rawResponse, projected.agentResult(), projected.evidenceStatus(), clock.instant());
|
||||
return ToolBoundaryResult.ready(toolCallId, projected.agentResult(), projected.evidenceStatus());
|
||||
} catch (ResultTooLargeException e) {
|
||||
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
|
||||
} catch (BudgetExceededException e) {
|
||||
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
|
||||
} catch (RunAbortedException e) {
|
||||
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.RUN_INACTIVE);
|
||||
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RUN_INACTIVE);
|
||||
} catch (CanonicalStoreException e) {
|
||||
ToolBoundaryErrorCode code = e instanceof ResultTooLargeException
|
||||
? ToolBoundaryErrorCode.RESULT_TOO_LARGE
|
||||
: ToolBoundaryErrorCode.PROJECTION_ERROR;
|
||||
markErrorSafely(key, rawResponse, code);
|
||||
return ToolBoundaryResult.error(toolCallId, code);
|
||||
}
|
||||
}
|
||||
|
||||
private String preflight(RunContext context, ToolCallRequestEnvelope request) {
|
||||
if (context == null || request == null) {
|
||||
throw new IllegalArgumentException("request/context must not be null");
|
||||
}
|
||||
if (!context.runId().equals(request.runId())) {
|
||||
throw new RunMismatchException();
|
||||
}
|
||||
if (!request.authorized()) {
|
||||
throw new UnauthorizedException();
|
||||
}
|
||||
if (!request.readOnly()) {
|
||||
throw new NotReadOnlyException();
|
||||
}
|
||||
if (request.toolName() == null || request.toolName().isBlank()
|
||||
|| request.requestJson() == null || request.requestJson().isBlank()) {
|
||||
throw new IllegalArgumentException("tool name and request must not be blank");
|
||||
}
|
||||
try {
|
||||
JsonNode root = objectMapper.readTree(request.requestJson());
|
||||
if (root == null || !root.isObject()) {
|
||||
throw new IllegalArgumentException("request must be a JSON object");
|
||||
}
|
||||
} catch (JsonProcessingException e) {
|
||||
throw new IllegalArgumentException("request must be valid JSON", e);
|
||||
}
|
||||
try {
|
||||
return keyFactory.create(request.runId(), request.toolCallId());
|
||||
} catch (IllegalArgumentException e) {
|
||||
throw new InvalidToolCallIdException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private void markErrorSafely(String key, String rawResponse, ToolBoundaryErrorCode code) {
|
||||
try {
|
||||
store.markError(key, rawResponse, code.name(), clock.instant());
|
||||
} catch (RuntimeException e) {
|
||||
log.warn("Failed to persist canonical Tool error: code={}", code, e);
|
||||
}
|
||||
}
|
||||
|
||||
private ToolBoundaryResult errorAndNoRecord(String toolCallId, ToolBoundaryErrorCode code) {
|
||||
return ToolBoundaryResult.error(toolCallId, code);
|
||||
}
|
||||
|
||||
private ToolBoundaryErrorCode classifyPreflightError(IllegalArgumentException exception) {
|
||||
if (exception instanceof InvalidToolCallIdException) {
|
||||
return ToolBoundaryErrorCode.INVALID_TOOL_CALL_ID;
|
||||
}
|
||||
if (exception instanceof RunMismatchException) {
|
||||
return ToolBoundaryErrorCode.RUN_MISMATCH;
|
||||
}
|
||||
if (exception instanceof UnauthorizedException) {
|
||||
return ToolBoundaryErrorCode.UNAUTHORIZED;
|
||||
}
|
||||
if (exception instanceof NotReadOnlyException) {
|
||||
return ToolBoundaryErrorCode.NOT_READ_ONLY;
|
||||
}
|
||||
return ToolBoundaryErrorCode.INVALID_REQUEST;
|
||||
}
|
||||
|
||||
private static final class InvalidToolCallIdException extends IllegalArgumentException {
|
||||
private InvalidToolCallIdException(Throwable cause) {
|
||||
super("Invalid tool call ID", cause);
|
||||
}
|
||||
}
|
||||
|
||||
private static final class RunMismatchException extends IllegalArgumentException {
|
||||
private RunMismatchException() {
|
||||
super("Run ID does not match context");
|
||||
}
|
||||
}
|
||||
|
||||
private static final class UnauthorizedException extends IllegalArgumentException {
|
||||
private UnauthorizedException() {
|
||||
super("Tool call is not authorized");
|
||||
}
|
||||
}
|
||||
|
||||
private static final class NotReadOnlyException extends IllegalArgumentException {
|
||||
private NotReadOnlyException() {
|
||||
super("Tool call is not read-only");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
public enum ToolBoundaryErrorCode {
|
||||
INVALID_REQUEST,
|
||||
INVALID_TOOL_CALL_ID,
|
||||
RUN_MISMATCH,
|
||||
UNAUTHORIZED,
|
||||
NOT_READ_ONLY,
|
||||
RUN_INACTIVE,
|
||||
BUDGET_EXHAUSTED,
|
||||
DUPLICATE_TOOL_CALL,
|
||||
RESULT_TOO_LARGE,
|
||||
TOOL_EXECUTION_ERROR,
|
||||
PROJECTION_ERROR,
|
||||
INVALID_EVIDENCE_STATUS,
|
||||
STORE_ERROR
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
import com.superbiz.agent.harness.contract.InvocationStatus;
|
||||
|
||||
public record ToolBoundaryResult(
|
||||
@JsonProperty("status") InvocationStatus status,
|
||||
@JsonProperty("evidence_status") EvidenceStatus evidenceStatus,
|
||||
@JsonProperty("tool_call_id") String toolCallId,
|
||||
@JsonProperty("agent_result") String agentResult,
|
||||
@JsonProperty("error_code") String errorCode) {
|
||||
|
||||
public ToolBoundaryResult {
|
||||
if (status == InvocationStatus.READY) {
|
||||
if (agentResult == null || evidenceStatus == null
|
||||
|| (evidenceStatus != EvidenceStatus.EVIDENCE_FOUND
|
||||
&& evidenceStatus != EvidenceStatus.NO_EVIDENCE)) {
|
||||
throw new IllegalArgumentException("READY result requires bounded evidence result");
|
||||
}
|
||||
if (errorCode != null) {
|
||||
throw new IllegalArgumentException("READY result must not contain errorCode");
|
||||
}
|
||||
} else if (status == InvocationStatus.ERROR) {
|
||||
if (evidenceStatus != EvidenceStatus.ERROR || errorCode == null || errorCode.isBlank()) {
|
||||
throw new IllegalArgumentException("ERROR result requires errorCode and ERROR evidence status");
|
||||
}
|
||||
if (agentResult != null) {
|
||||
throw new IllegalArgumentException("ERROR result must not contain agent result");
|
||||
}
|
||||
} else {
|
||||
throw new IllegalArgumentException("ToolBoundaryResult must be READY or ERROR");
|
||||
}
|
||||
}
|
||||
|
||||
public static ToolBoundaryResult ready(String toolCallId,
|
||||
String agentResult,
|
||||
EvidenceStatus evidenceStatus) {
|
||||
return new ToolBoundaryResult(
|
||||
InvocationStatus.READY, evidenceStatus, toolCallId, agentResult, null);
|
||||
}
|
||||
|
||||
public static ToolBoundaryResult error(String toolCallId, ToolBoundaryErrorCode errorCode) {
|
||||
return new ToolBoundaryResult(
|
||||
InvocationStatus.ERROR, EvidenceStatus.ERROR, toolCallId, null, errorCode.name());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
|
||||
public record ToolCallRequestEnvelope(
|
||||
@JsonProperty("run_id") String runId,
|
||||
@JsonProperty("tool_call_id") String toolCallId,
|
||||
@JsonProperty("tool_name") String toolName,
|
||||
@JsonProperty("request") String requestJson,
|
||||
@JsonProperty("authorized") boolean authorized,
|
||||
@JsonProperty("read_only") boolean readOnly) {
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
@FunctionalInterface
|
||||
public interface ToolExecutor {
|
||||
String execute(String requestJson) throws Exception;
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package com.superbiz.agent.harness.tool.boundary;
|
||||
|
||||
@FunctionalInterface
|
||||
public interface ToolResultProjector {
|
||||
ProjectedToolResult project(String rawResponse) throws Exception;
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.time.Duration;
|
||||
import java.util.Objects;
|
||||
|
||||
public record CanonicalInvocationLimits(
|
||||
Duration ttl,
|
||||
long maxRecordBytes,
|
||||
long maxAgentResultBytes) {
|
||||
|
||||
public CanonicalInvocationLimits {
|
||||
Objects.requireNonNull(ttl, "ttl must not be null");
|
||||
if (ttl.isZero() || ttl.isNegative()) {
|
||||
throw new IllegalArgumentException("ttl must be positive");
|
||||
}
|
||||
if (maxRecordBytes <= 0 || maxAgentResultBytes <= 0) {
|
||||
throw new IllegalArgumentException("byte limits must be positive");
|
||||
}
|
||||
if (maxAgentResultBytes > maxRecordBytes) {
|
||||
throw new IllegalArgumentException("maxAgentResultBytes must not exceed maxRecordBytes");
|
||||
}
|
||||
ttl.toMillis();
|
||||
}
|
||||
|
||||
public void validateRawCandidate(String request, String rawResponse) {
|
||||
long actual = utf8Bytes(request) + utf8Bytes(rawResponse);
|
||||
if (actual > maxRecordBytes) {
|
||||
throw new ResultTooLargeException("raw_response", maxRecordBytes, actual);
|
||||
}
|
||||
}
|
||||
|
||||
public void validateAgentResult(String agentResult) {
|
||||
long actual = utf8Bytes(agentResult);
|
||||
if (actual > maxAgentResultBytes) {
|
||||
throw new ResultTooLargeException("agent_result", maxAgentResultBytes, actual);
|
||||
}
|
||||
}
|
||||
|
||||
public void validateSerializedRecord(String json) {
|
||||
long actual = utf8Bytes(json);
|
||||
if (actual > maxRecordBytes) {
|
||||
throw new ResultTooLargeException("canonical_record", maxRecordBytes, actual);
|
||||
}
|
||||
}
|
||||
|
||||
public static long utf8Bytes(String value) {
|
||||
return value == null ? 0 : value.getBytes(StandardCharsets.UTF_8).length;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
|
||||
import java.time.Instant;
|
||||
import java.util.Optional;
|
||||
|
||||
public interface CanonicalInvocationStore {
|
||||
|
||||
CanonicalInvocationLimits limits();
|
||||
|
||||
void begin(String key, CanonicalToolInvocation invocation);
|
||||
|
||||
Optional<CanonicalToolInvocation> find(String key);
|
||||
|
||||
CanonicalToolInvocation markReady(String key,
|
||||
String rawResponse,
|
||||
String agentResult,
|
||||
EvidenceStatus evidenceStatus,
|
||||
Instant completedAt);
|
||||
|
||||
CanonicalToolInvocation markError(String key,
|
||||
String rawResponse,
|
||||
String errorCode,
|
||||
Instant completedAt);
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
public class CanonicalStoreException extends RuntimeException {
|
||||
|
||||
public CanonicalStoreException(String message) {
|
||||
super(message);
|
||||
}
|
||||
|
||||
public CanonicalStoreException(String message, Throwable cause) {
|
||||
super(message, cause);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
import com.superbiz.agent.harness.contract.InvocationStatus;
|
||||
|
||||
import java.time.Instant;
|
||||
import java.util.Objects;
|
||||
|
||||
public record CanonicalToolInvocation(
|
||||
@JsonProperty("tool_call_id") String toolCallId,
|
||||
@JsonProperty("run_id") String runId,
|
||||
@JsonProperty("tool_name") String toolName,
|
||||
@JsonProperty("request") String request,
|
||||
@JsonProperty("raw_response") String rawResponse,
|
||||
@JsonProperty("agent_result") String agentResult,
|
||||
@JsonProperty("status") InvocationStatus status,
|
||||
@JsonProperty("evidence_status") EvidenceStatus evidenceStatus,
|
||||
@JsonProperty("error_code") String errorCode,
|
||||
@JsonProperty("started_at") Instant startedAt,
|
||||
@JsonProperty("completed_at") Instant completedAt) {
|
||||
|
||||
public CanonicalToolInvocation {
|
||||
requireText(toolCallId, "toolCallId");
|
||||
requireText(runId, "runId");
|
||||
requireText(toolName, "toolName");
|
||||
requireText(request, "request");
|
||||
Objects.requireNonNull(status, "status must not be null");
|
||||
Objects.requireNonNull(startedAt, "startedAt must not be null");
|
||||
validateState(status, evidenceStatus, rawResponse, agentResult, errorCode, completedAt);
|
||||
}
|
||||
|
||||
public static CanonicalToolInvocation projecting(String toolCallId,
|
||||
String runId,
|
||||
String toolName,
|
||||
String request,
|
||||
Instant startedAt) {
|
||||
return new CanonicalToolInvocation(
|
||||
toolCallId, runId, toolName, request, null, null,
|
||||
InvocationStatus.PROJECTING, null, null, startedAt, null);
|
||||
}
|
||||
|
||||
public CanonicalToolInvocation markReady(String rawResponse,
|
||||
String agentResult,
|
||||
EvidenceStatus evidenceStatus,
|
||||
Instant completedAt) {
|
||||
requireProjecting();
|
||||
return new CanonicalToolInvocation(
|
||||
toolCallId, runId, toolName, request, rawResponse, agentResult,
|
||||
InvocationStatus.READY, evidenceStatus, null, startedAt, completedAt);
|
||||
}
|
||||
|
||||
public CanonicalToolInvocation markError(String rawResponse,
|
||||
String errorCode,
|
||||
Instant completedAt) {
|
||||
requireProjecting();
|
||||
return new CanonicalToolInvocation(
|
||||
toolCallId, runId, toolName, request, rawResponse, null,
|
||||
InvocationStatus.ERROR, EvidenceStatus.ERROR, errorCode, startedAt, completedAt);
|
||||
}
|
||||
|
||||
public boolean isReferencableBy(String expectedRunId) {
|
||||
return status == InvocationStatus.READY
|
||||
&& Objects.equals(runId, expectedRunId)
|
||||
&& agentResult != null
|
||||
&& (evidenceStatus == EvidenceStatus.EVIDENCE_FOUND
|
||||
|| evidenceStatus == EvidenceStatus.NO_EVIDENCE);
|
||||
}
|
||||
|
||||
private void requireProjecting() {
|
||||
if (status != InvocationStatus.PROJECTING) {
|
||||
throw new InvocationStateException("Only PROJECTING invocation can transition");
|
||||
}
|
||||
}
|
||||
|
||||
private static void validateState(InvocationStatus status,
|
||||
EvidenceStatus evidenceStatus,
|
||||
String rawResponse,
|
||||
String agentResult,
|
||||
String errorCode,
|
||||
Instant completedAt) {
|
||||
switch (status) {
|
||||
case PROJECTING -> {
|
||||
if (evidenceStatus != null || agentResult != null || errorCode != null || completedAt != null) {
|
||||
throw new InvocationStateException("PROJECTING invocation contains terminal fields");
|
||||
}
|
||||
}
|
||||
case READY -> {
|
||||
if (rawResponse == null || agentResult == null || completedAt == null) {
|
||||
throw new InvocationStateException("READY invocation requires raw, agent result and completion");
|
||||
}
|
||||
if (evidenceStatus != EvidenceStatus.EVIDENCE_FOUND
|
||||
&& evidenceStatus != EvidenceStatus.NO_EVIDENCE) {
|
||||
throw new InvocationStateException("READY invocation has invalid evidence status");
|
||||
}
|
||||
if (errorCode != null) {
|
||||
throw new InvocationStateException("READY invocation must not contain errorCode");
|
||||
}
|
||||
}
|
||||
case ERROR -> {
|
||||
if (evidenceStatus != EvidenceStatus.ERROR || completedAt == null) {
|
||||
throw new InvocationStateException("ERROR invocation requires ERROR evidence status and completion");
|
||||
}
|
||||
requireText(errorCode, "errorCode");
|
||||
if (agentResult != null) {
|
||||
throw new InvocationStateException("ERROR invocation must not contain agent result");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static void requireText(String value, String name) {
|
||||
if (value == null || value.isBlank()) {
|
||||
throw new IllegalArgumentException(name + " must not be blank");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
public final class DuplicateInvocationException extends CanonicalStoreException {
|
||||
|
||||
public DuplicateInvocationException() {
|
||||
super("Canonical invocation already exists");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
public final class InvocationStateException extends CanonicalStoreException {
|
||||
|
||||
public InvocationStateException(String message) {
|
||||
super(message);
|
||||
}
|
||||
}
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.superbiz.agent.harness.contract.EvidenceStatus;
|
||||
import com.superbiz.agent.harness.contract.InvocationStatus;
|
||||
import org.springframework.data.redis.core.RedisTemplate;
|
||||
import org.springframework.data.redis.core.ValueOperations;
|
||||
|
||||
import java.time.Instant;
|
||||
import java.util.Objects;
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.function.UnaryOperator;
|
||||
|
||||
public final class RedisCanonicalInvocationStore implements CanonicalInvocationStore {
|
||||
|
||||
private final RedisTemplate<String, Object> redisTemplate;
|
||||
private final ValueOperations<String, Object> values;
|
||||
private final ObjectMapper objectMapper;
|
||||
private final CanonicalInvocationLimits limits;
|
||||
|
||||
public RedisCanonicalInvocationStore(RedisTemplate<String, Object> redisTemplate,
|
||||
ObjectMapper objectMapper,
|
||||
CanonicalInvocationLimits limits) {
|
||||
this.redisTemplate = Objects.requireNonNull(redisTemplate, "redisTemplate must not be null");
|
||||
this.values = redisTemplate.opsForValue();
|
||||
this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null");
|
||||
this.limits = Objects.requireNonNull(limits, "limits must not be null");
|
||||
}
|
||||
|
||||
@Override
|
||||
public CanonicalInvocationLimits limits() {
|
||||
return limits;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void begin(String key, CanonicalToolInvocation invocation) {
|
||||
requireKey(key);
|
||||
Objects.requireNonNull(invocation, "invocation must not be null");
|
||||
if (invocation.status() != InvocationStatus.PROJECTING) {
|
||||
throw new InvocationStateException("begin requires PROJECTING invocation");
|
||||
}
|
||||
String json = serialize(invocation);
|
||||
limits.validateSerializedRecord(json);
|
||||
Boolean created = values.setIfAbsent(
|
||||
key, json, limits.ttl().toMillis(), TimeUnit.MILLISECONDS);
|
||||
if (!Boolean.TRUE.equals(created)) {
|
||||
throw new DuplicateInvocationException();
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public Optional<CanonicalToolInvocation> find(String key) {
|
||||
requireKey(key);
|
||||
Object stored = values.get(key);
|
||||
if (stored == null) {
|
||||
return Optional.empty();
|
||||
}
|
||||
if (!(stored instanceof String json)) {
|
||||
throw new CanonicalStoreException("Canonical invocation value is not JSON text");
|
||||
}
|
||||
return Optional.of(deserialize(json));
|
||||
}
|
||||
|
||||
@Override
|
||||
public CanonicalToolInvocation markReady(String key,
|
||||
String rawResponse,
|
||||
String agentResult,
|
||||
EvidenceStatus evidenceStatus,
|
||||
Instant completedAt) {
|
||||
CanonicalToolInvocation existing = find(key)
|
||||
.orElseThrow(() -> new InvocationStateException("Canonical invocation is missing or expired"));
|
||||
limits.validateRawCandidate(existing.request(), requireValue(rawResponse, "rawResponse"));
|
||||
limits.validateAgentResult(requireValue(agentResult, "agentResult"));
|
||||
return update(key, current -> current.markReady(
|
||||
rawResponse, agentResult, evidenceStatus, completedAt));
|
||||
}
|
||||
|
||||
@Override
|
||||
public CanonicalToolInvocation markError(String key,
|
||||
String rawResponse,
|
||||
String errorCode,
|
||||
Instant completedAt) {
|
||||
try {
|
||||
return update(key, current -> current.markError(rawResponse, errorCode, completedAt));
|
||||
} catch (ResultTooLargeException e) {
|
||||
if (rawResponse == null) {
|
||||
throw e;
|
||||
}
|
||||
return update(key, current -> current.markError(
|
||||
null, ResultTooLargeException.ERROR_CODE, completedAt));
|
||||
}
|
||||
}
|
||||
|
||||
private CanonicalToolInvocation update(String key,
|
||||
UnaryOperator<CanonicalToolInvocation> transition) {
|
||||
CanonicalToolInvocation current = find(key)
|
||||
.orElseThrow(() -> new InvocationStateException("Canonical invocation is missing or expired"));
|
||||
long remainingTtl = remainingTtlMillis(key);
|
||||
CanonicalToolInvocation updated = transition.apply(current);
|
||||
String json = serialize(updated);
|
||||
limits.validateSerializedRecord(json);
|
||||
values.set(key, json, remainingTtl, TimeUnit.MILLISECONDS);
|
||||
return updated;
|
||||
}
|
||||
|
||||
private long remainingTtlMillis(String key) {
|
||||
Long remaining = redisTemplate.getExpire(key, TimeUnit.MILLISECONDS);
|
||||
if (remaining == null || remaining <= 0) {
|
||||
throw new InvocationStateException("Canonical invocation is missing or expired");
|
||||
}
|
||||
return Math.min(remaining, limits.ttl().toMillis());
|
||||
}
|
||||
|
||||
private String serialize(CanonicalToolInvocation invocation) {
|
||||
try {
|
||||
return objectMapper.writeValueAsString(invocation);
|
||||
} catch (JsonProcessingException e) {
|
||||
throw new CanonicalStoreException("Failed to serialize canonical invocation", e);
|
||||
}
|
||||
}
|
||||
|
||||
private CanonicalToolInvocation deserialize(String json) {
|
||||
try {
|
||||
return objectMapper.readValue(json, CanonicalToolInvocation.class);
|
||||
} catch (JsonProcessingException e) {
|
||||
throw new CanonicalStoreException("Failed to deserialize canonical invocation", e);
|
||||
}
|
||||
}
|
||||
|
||||
private static String requireValue(String value, String name) {
|
||||
if (value == null) {
|
||||
throw new IllegalArgumentException(name + " must not be null");
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
private static void requireKey(String key) {
|
||||
if (key == null || key.isBlank()) {
|
||||
throw new IllegalArgumentException("key must not be blank");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
package com.superbiz.agent.harness.tool.store;
|
||||
|
||||
public final class ResultTooLargeException extends CanonicalStoreException {
|
||||
|
||||
public static final String ERROR_CODE = "RESULT_TOO_LARGE";
|
||||
|
||||
public ResultTooLargeException(String field, long limit, long actual) {
|
||||
super(field + " exceeds byte limit: limit=" + limit + ", actual=" + actual);
|
||||
}
|
||||
}
|
||||
@@ -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