feat(harness,rag): dual LLM audit fields, run conclusion, and hybrid quality

Persist provider reasoning and assistant text separately on agent_reasoning_audit
(DeepSeekAssistantMessage path), extract diagnosis_run.conclusion, enrich RAG
tool audit (step_id/query/qualityScore), gate empty mysql tools, drop devtools,
and align MVP docs after live E2E verification.
This commit is contained in:
zhuyongxin
2026-07-28 19:43:13 +08:00
parent 2f40536248
commit 7ae9707a3b
116 changed files with 8364 additions and 1141 deletions
@@ -31,6 +31,10 @@ public final class HarnessEvidenceTools {
private final List<ToolCallback> callbacks;
private final Map<String, EvidenceToolInvoker> invokers;
/**
* @param mysqlInvoker optional; when null, {@code query_mysql} is not registered
* (no logical datasource configured / tool unavailable).
*/
public HarnessEvidenceTools(EvidenceToolInvoker ragInvoker,
EvidenceToolInvoker logsInvoker,
EvidenceToolInvoker mysqlInvoker) {
@@ -39,28 +43,38 @@ public final class HarnessEvidenceTools {
Objects.requireNonNull(ragInvoker, "ragInvoker must not be null"));
registered.put(AgentToolContracts.QUERY_LOGS,
Objects.requireNonNull(logsInvoker, "logsInvoker must not be null"));
registered.put(AgentToolContracts.QUERY_MYSQL,
Objects.requireNonNull(mysqlInvoker, "mysqlInvoker must not be null"));
if (mysqlInvoker != null) {
registered.put(AgentToolContracts.QUERY_MYSQL, mysqlInvoker);
}
this.invokers = Map.copyOf(registered);
this.callbacks = List.of(
definition(AgentToolContracts.LOOKUP_KNOWLEDGE,
AgentToolContracts.LOOKUP_KNOWLEDGE_DESCRIPTION, RagToolCall.class),
definition(AgentToolContracts.QUERY_LOGS,
AgentToolContracts.QUERY_LOGS_DESCRIPTION, QueryLogsToolCall.class),
definition(AgentToolContracts.QUERY_MYSQL,
AgentToolContracts.QUERY_MYSQL_DESCRIPTION, MysqlToolCall.class));
List<ToolCallback> built = new java.util.ArrayList<>();
built.add(definition(AgentToolContracts.LOOKUP_KNOWLEDGE,
AgentToolContracts.LOOKUP_KNOWLEDGE_DESCRIPTION, RagToolCall.class));
built.add(definition(AgentToolContracts.QUERY_LOGS,
AgentToolContracts.QUERY_LOGS_DESCRIPTION, QueryLogsToolCall.class));
if (mysqlInvoker != null) {
built.add(definition(AgentToolContracts.QUERY_MYSQL,
AgentToolContracts.QUERY_MYSQL_DESCRIPTION, MysqlToolCall.class));
}
this.callbacks = List.copyOf(built);
}
/**
* @param mysqlAdapter optional; omit registration when null or when no datasources are wired
*/
public static HarnessEvidenceTools fromAdapters(RagToolAdapter ragAdapter,
QueryLogsToolAdapter logsAdapter,
MysqlToolAdapter mysqlAdapter) {
Objects.requireNonNull(ragAdapter, "ragAdapter must not be null");
Objects.requireNonNull(logsAdapter, "logsAdapter must not be null");
Objects.requireNonNull(mysqlAdapter, "mysqlAdapter must not be null");
EvidenceToolInvoker mysql = mysqlAdapter == null
? null
: bridge(AgentToolContracts.QUERY_MYSQL, mysqlAdapter::execute);
return new HarnessEvidenceTools(
bridge(AgentToolContracts.LOOKUP_KNOWLEDGE, ragAdapter::execute),
bridge(AgentToolContracts.QUERY_LOGS, logsAdapter::execute),
bridge(AgentToolContracts.QUERY_MYSQL, mysqlAdapter::execute));
mysql);
}
public List<ToolCallback> callbacks() {
@@ -6,18 +6,21 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectReader;
import com.superbiz.agent.domain.entity.ChatSession;
import com.superbiz.agent.domain.entity.DiagnosisRun;
import com.superbiz.agent.harness.audit.AgentStepAuditTracker;
import com.superbiz.agent.harness.audit.ModelCallComponent;
import com.superbiz.agent.harness.audit.RunConclusionExtractor;
import com.superbiz.agent.harness.contract.IntentType;
import com.superbiz.agent.harness.contract.PreviousTurn;
import com.superbiz.agent.harness.contract.PublishedResult;
import com.superbiz.agent.harness.contract.ReleaseOutcome;
import com.superbiz.agent.harness.core.RunBudgetUsage;
import com.superbiz.agent.harness.core.RunContext;
import com.superbiz.agent.harness.audit.ModelCallComponent;
import com.superbiz.agent.repository.ChatSessionRepository;
import com.superbiz.agent.repository.DiagnosisRunRepository;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.beans.factory.annotation.Autowired;
import java.time.LocalDateTime;
import java.util.List;
@@ -32,19 +35,28 @@ public class JpaChatRunStore implements ChatRunStore {
private final ObjectMapper objectMapper;
private final ObjectReader publishedReader;
private final PublishedResultPolicy publishedPolicy;
private final AgentStepAuditTracker stepTracker;
public JpaChatRunStore(ChatSessionRepository chatSessions,
DiagnosisRunRepository runs,
ObjectMapper objectMapper) {
this(chatSessions, runs, objectMapper,
new PublishedResultPolicy(PreviousTurnLimits.defaults()));
new PublishedResultPolicy(PreviousTurnLimits.defaults()), null);
}
public JpaChatRunStore(ChatSessionRepository chatSessions,
DiagnosisRunRepository runs,
ObjectMapper objectMapper,
PublishedResultPolicy publishedPolicy) {
this(chatSessions, runs, objectMapper, publishedPolicy, null);
}
@Autowired
public JpaChatRunStore(ChatSessionRepository chatSessions,
DiagnosisRunRepository runs,
ObjectMapper objectMapper,
PublishedResultPolicy publishedPolicy) {
PublishedResultPolicy publishedPolicy,
@Nullable AgentStepAuditTracker stepTracker) {
this.chatSessions = Objects.requireNonNull(chatSessions, "chatSessions must not be null");
this.runs = Objects.requireNonNull(runs, "runs must not be null");
this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null");
@@ -52,6 +64,7 @@ public class JpaChatRunStore implements ChatRunStore {
.with(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
.with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS);
this.publishedPolicy = Objects.requireNonNull(publishedPolicy, "publishedPolicy must not be null");
this.stepTracker = stepTracker;
}
@Override
@@ -110,27 +123,34 @@ public class JpaChatRunStore implements ChatRunStore {
String safeContentJson, PublishedResult publishedResult, int durationMs) {
Objects.requireNonNull(context, "context must not be null");
Objects.requireNonNull(outcome, "outcome must not be null");
DiagnosisRun run = requiredRun(context.runId());
run.setIntent(intent);
run.setReleaseOutcome(outcome);
run.setStatus(status(outcome));
run.setAnswer(safeContentJson);
run.setPublishedResult(outcome == ReleaseOutcome.SUCCESS && intent == IntentType.DIAGNOSIS
&& publishedResult != null ? write(publishedResult) : null);
run.setTotalDurationMs(Math.max(0, durationMs));
RunBudgetUsage usage = context.budget().snapshot();
run.setTotalTokenCount(saturatingInt(usage.totalTokens()));
run.setStepCount(context.modelCalls().componentCallCount(
ModelCallComponent.DIAGNOSIS_AGENT));
run.setToolCallCount(saturatingInt(usage.toolCalls()));
runs.save(run);
if (outcome == ReleaseOutcome.SUCCESS || outcome == ReleaseOutcome.FALLBACK) {
chatSessions.findBySessionId(context.sessionId()).ifPresent(session -> {
session.setLastActiveAt(LocalDateTime.now());
session.setMessagePairCount(
session.getMessagePairCount() == null ? 1 : session.getMessagePairCount() + 1);
chatSessions.save(session);
});
try {
DiagnosisRun run = requiredRun(context.runId());
run.setIntent(intent);
run.setReleaseOutcome(outcome);
run.setStatus(status(outcome));
run.setAnswer(safeContentJson);
run.setConclusion(RunConclusionExtractor.extract(objectMapper, safeContentJson));
run.setPublishedResult(outcome == ReleaseOutcome.SUCCESS && intent == IntentType.DIAGNOSIS
&& publishedResult != null ? write(publishedResult) : null);
run.setTotalDurationMs(Math.max(0, durationMs));
RunBudgetUsage usage = context.budget().snapshot();
run.setTotalTokenCount(saturatingInt(usage.totalTokens()));
run.setStepCount(context.modelCalls().componentCallCount(
ModelCallComponent.DIAGNOSIS_AGENT));
run.setToolCallCount(saturatingInt(usage.toolCalls()));
runs.save(run);
if (outcome == ReleaseOutcome.SUCCESS || outcome == ReleaseOutcome.FALLBACK) {
chatSessions.findBySessionId(context.sessionId()).ifPresent(session -> {
session.setLastActiveAt(LocalDateTime.now());
session.setMessagePairCount(
session.getMessagePairCount() == null ? 1 : session.getMessagePairCount() + 1);
chatSessions.save(session);
});
}
} finally {
if (stepTracker != null) {
stepTracker.clear(context.runId());
}
}
}
@@ -0,0 +1,39 @@
package com.superbiz.agent.harness.audit;
import org.springframework.stereotype.Component;
import java.util.concurrent.ConcurrentHashMap;
/**
* Binds the in-flight {@code agent_step.id} for a run so tool audits can set {@code step_id}.
*
* <p>Lifecycle: {@link #bind} on Agent {@code beforeModel} (step row created). Binding stays
* until the next {@code beforeModel} for the same run, covering tool execution that happens
* after {@code afterModel} emits tool_calls. {@link #clear} on run finish is optional cleanup.</p>
*/
@Component
public final class AgentStepAuditTracker {
private final ConcurrentHashMap<String, Long> currentStepIdByRun = new ConcurrentHashMap<>();
public void bind(String runId, Long stepId) {
if (runId == null || runId.isBlank() || stepId == null) {
return;
}
currentStepIdByRun.put(runId, stepId);
}
public Long currentStepId(String runId) {
if (runId == null || runId.isBlank()) {
return null;
}
return currentStepIdByRun.get(runId);
}
public void clear(String runId) {
if (runId == null || runId.isBlank()) {
return;
}
currentStepIdByRun.remove(runId);
}
}
@@ -7,32 +7,53 @@ import com.alibaba.cloud.ai.graph.agent.hook.messages.AgentCommand;
import com.alibaba.cloud.ai.graph.agent.hook.messages.MessagesModelHook;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.domain.entity.AgentStep;
import com.superbiz.agent.domain.entity.AgentReasoningAudit;
import com.superbiz.agent.domain.entity.AgentStep;
import com.superbiz.agent.repository.AgentReasoningAuditRepository;
import com.superbiz.agent.repository.AgentStepRepository;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.deepseek.DeepSeekAssistantMessage;
import java.nio.charset.StandardCharsets;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
/**
* Persists per-model-step audit for the diagnosis agent.
*
* <p>LLM-facing content policy for {@link AgentReasoningAudit}:</p>
* <ul>
* <li>{@code reasoning_content}: provider thinking/CoT when present</li>
* <li>{@code assistant_text}: assistant visible text and/or tool-call plan</li>
* <li>Tool <em>result</em> payloads are never stored here (use {@code tool_invocation})</li>
* </ul>
*/
@HookPositions({HookPosition.BEFORE_MODEL, HookPosition.AFTER_MODEL})
public final class HarnessAgentAuditHook extends MessagesModelHook {
private static final Logger log = LoggerFactory.getLogger(HarnessAgentAuditHook.class);
private static final int MAX_TEXT_CHARS = 32_000;
public static final String SOURCE_BOTH = "PROVIDER_REASONING+ASSISTANT_TEXT";
public static final String SOURCE_PROVIDER = "PROVIDER_REASONING";
public static final String SOURCE_ASSISTANT = "ASSISTANT_TEXT";
public static final String SOURCE_TOOL_PLAN = "TOOL_CALL_PLAN";
public static final String SOURCE_NONE = "NONE";
private final AgentStepRepository repository;
private final ObjectMapper objectMapper;
private final String agentName;
private final DiagnosisTraceRecorder traceRecorder;
private final AgentReasoningAuditRepository reasoningRepository;
private final AgentStepAuditTracker stepTracker;
private final ConcurrentHashMap<String, Integer> stepCounters = new ConcurrentHashMap<>();
private final ConcurrentHashMap<String, PendingStep> pendingSteps = new ConcurrentHashMap<>();
@@ -42,12 +63,19 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
public HarnessAgentAuditHook(AgentStepRepository repository, ObjectMapper objectMapper,
String agentName, DiagnosisTraceRecorder traceRecorder) {
this(repository, objectMapper, agentName, traceRecorder, null);
this(repository, objectMapper, agentName, traceRecorder, null, null);
}
public HarnessAgentAuditHook(AgentStepRepository repository, ObjectMapper objectMapper,
String agentName, DiagnosisTraceRecorder traceRecorder,
AgentReasoningAuditRepository reasoningRepository) {
this(repository, objectMapper, agentName, traceRecorder, reasoningRepository, null);
}
public HarnessAgentAuditHook(AgentStepRepository repository, ObjectMapper objectMapper,
String agentName, DiagnosisTraceRecorder traceRecorder,
AgentReasoningAuditRepository reasoningRepository,
AgentStepAuditTracker stepTracker) {
this.repository = Objects.requireNonNull(repository, "repository must not be null");
this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null");
if (agentName == null || agentName.isBlank()) {
@@ -56,6 +84,7 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
this.agentName = agentName;
this.traceRecorder = Objects.requireNonNull(traceRecorder, "traceRecorder must not be null");
this.reasoningRepository = reasoningRepository;
this.stepTracker = stepTracker;
}
@Override
@@ -82,6 +111,9 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
.hasToolCall(false)
.build());
pending = new PendingStep(saved.getId(), pending.startedNanos(), pending.input());
if (stepTracker != null && saved.getId() != null) {
stepTracker.bind(identity.runId(), saved.getId());
}
} catch (RuntimeException exception) {
log.warn("Failed to persist AgentStep audit before model: agent={}", agentName);
}
@@ -106,14 +138,15 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
? List.of()
: assistant.getToolCalls().stream().map(AssistantMessage.ToolCall::name)
.distinct().sorted().toList();
ReasoningContent reasoning = reasoningContent(assistant);
Map<String, Object> output = outputMetadata(assistant, toolNames, reasoning);
LlmTurnContent turn = llmTurnContent(assistant, toolNames);
Map<String, Object> output = outputMetadata(assistant, toolNames, turn);
int durationMs = durationMillis(pending.startedNanos());
try {
AgentStep step = pending.id() == null ? null : repository.findById(pending.id()).orElse(null);
if (step != null) {
step.setModelOutput(write(output));
step.setThought(null);
// Keep thought aligned with provider reasoning when present; else assistant text.
step.setThought(firstNonBlank(turn.reasoningContent(), turn.assistantText()));
step.setHasToolCall(!toolNames.isEmpty());
step.setDurationMs(durationMs);
repository.save(step);
@@ -121,7 +154,7 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
} catch (RuntimeException exception) {
log.warn("Failed to complete AgentStep audit: agent={}", agentName);
}
persistReasoning(identity, stepIndex, reasoning);
persistReasoning(identity, stepIndex, turn);
traceRecorder.record(TraceAuditEvents.agentModelStep(
identity.sessionId(), identity.runId(), agentName, stepIndex,
durationMs, pending.input(), output));
@@ -139,32 +172,148 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
}
private Map<String, Object> outputMetadata(AssistantMessage assistant, List<String> toolNames,
ReasoningContent reasoning) {
LlmTurnContent turn) {
Map<String, Object> metadata = new LinkedHashMap<>();
metadata.put("has_text", assistant != null
&& assistant.getText() != null && !assistant.getText().isBlank());
metadata.put("tool_names", toolNames);
metadata.put("reasoning_available", reasoning.available());
metadata.put("reasoning_bytes", reasoning.bytes());
metadata.put("reasoning_available", turn.reasoningAvailable());
metadata.put("reasoning_bytes", turn.reasoningBytes());
metadata.put("assistant_bytes", turn.assistantBytes());
metadata.put("content_source", turn.contentSource());
return metadata;
}
private ReasoningContent reasoningContent(AssistantMessage assistant) {
if (assistant == null || assistant.getMetadata() == null) {
return ReasoningContent.empty();
/**
* Captures both provider reasoning and assistant-visible text.
* Tool <em>results</em> are never included.
*/
static LlmTurnContent llmTurnContent(AssistantMessage assistant, List<String> toolNames) {
String reasoning = providerReasoning(assistant);
String assistantBody = assistantText(assistant);
String toolPlan = toolCallPlan(assistant, toolNames);
String assistantCombined = joinNonBlank("\n", assistantBody, toolPlan);
boolean hasReasoning = hasText(reasoning);
boolean hasAssistant = hasText(assistantCombined);
String source;
if (hasReasoning && hasAssistant) {
source = SOURCE_BOTH;
} else if (hasReasoning) {
source = SOURCE_PROVIDER;
} else if (hasText(assistantBody)) {
source = SOURCE_ASSISTANT;
} else if (hasText(toolPlan)) {
source = SOURCE_TOOL_PLAN;
} else {
source = SOURCE_NONE;
}
for (String key : List.of("reasoning_content", "reasoningContent", "reasoning", "thinking")) {
Object value = assistant.getMetadata().get(key);
if (value instanceof CharSequence text && !text.toString().isBlank()) {
String content = bounded(text.toString());
return new ReasoningContent(content, true,
content.getBytes(StandardCharsets.UTF_8).length);
}
}
return ReasoningContent.empty();
String reasoningBound = bound(reasoning);
String assistantBound = bound(assistantCombined);
int bytes = utf8Bytes(reasoningBound) + utf8Bytes(assistantBound);
return new LlmTurnContent(
hasReasoning,
reasoningBound,
assistantBound,
source,
utf8Bytes(reasoningBound),
utf8Bytes(assistantBound),
bytes);
}
private void persistReasoning(AuditIdentity identity, int stepIndex, ReasoningContent reasoning) {
/**
* DeepSeek puts CoT on {@link DeepSeekAssistantMessage#getReasoningContent()},
* <em>not</em> on {@link AssistantMessage#getMetadata()}. Older docs/tests used metadata keys;
* keep those as fallback for mocks and non-DeepSeek providers.
*/
static String providerReasoning(AssistantMessage assistant) {
if (assistant == null) {
return null;
}
// 1) Native DeepSeek message field (primary path in production)
if (assistant instanceof DeepSeekAssistantMessage deepSeek) {
String nativeReasoning = blankToNull(deepSeek.getReasoningContent());
if (nativeReasoning != null) {
return nativeReasoning;
}
}
// 2) Reflective getReasoningContent() for subclasses / reloaded types
String reflective = invokeReasoningGetter(assistant);
if (reflective != null) {
return reflective;
}
// 3) Metadata keys (tests / other providers)
Map<String, Object> metadata = assistant.getMetadata();
if (metadata != null) {
for (String key : List.of(
"reasoning_content", "reasoningContent", "reasoning", "thinking",
"reasoning_text", "reasoningText")) {
Object value = metadata.get(key);
if (value instanceof CharSequence text) {
String trimmed = blankToNull(text.toString());
if (trimmed != null) {
return trimmed;
}
}
}
}
return null;
}
private static String invokeReasoningGetter(AssistantMessage assistant) {
try {
Method method = assistant.getClass().getMethod("getReasoningContent");
Object value = method.invoke(assistant);
return value instanceof CharSequence text ? blankToNull(text.toString()) : null;
} catch (ReflectiveOperationException ignored) {
return null;
}
}
private static String blankToNull(String value) {
if (value == null) {
return null;
}
String trimmed = value.trim();
return trimmed.isEmpty() ? null : trimmed;
}
private static String assistantText(AssistantMessage assistant) {
if (assistant == null || assistant.getText() == null || assistant.getText().isBlank()) {
return null;
}
return assistant.getText().trim();
}
/**
* Records which tools the model decided to call (names + arg preview), not tool outputs.
*/
private static String toolCallPlan(AssistantMessage assistant, List<String> toolNames) {
if (assistant == null || assistant.getToolCalls() == null || assistant.getToolCalls().isEmpty()) {
return null;
}
List<String> lines = new ArrayList<>();
lines.add("tool_calls:");
for (AssistantMessage.ToolCall call : assistant.getToolCalls()) {
if (call == null) {
continue;
}
String name = call.name() == null ? "?" : call.name();
String args = call.arguments() == null ? "" : call.arguments().trim();
if (args.length() > 500) {
args = args.substring(0, 500) + "...";
}
lines.add("- " + name + (args.isEmpty() ? "" : " args=" + args));
}
if (lines.size() == 1 && toolNames != null && !toolNames.isEmpty()) {
lines.add("- " + String.join(", ", toolNames));
}
return lines.size() <= 1 ? null : String.join("\n", lines);
}
private void persistReasoning(AuditIdentity identity, int stepIndex, LlmTurnContent turn) {
if (reasoningRepository == null) {
return;
}
@@ -174,18 +323,51 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
.runId(identity.runId())
.stepIndex(stepIndex)
.agentName(agentName)
.reasoningAvailable(reasoning.available())
.reasoningContent(reasoning.content())
.contentBytes(reasoning.bytes())
.reasoningAvailable(turn.reasoningAvailable())
.reasoningContent(turn.reasoningContent())
.assistantText(turn.assistantText())
.contentSource(turn.contentSource())
.contentBytes(turn.totalBytes())
.build());
} catch (RuntimeException exception) {
log.warn("Failed to persist reasoning audit: agent={}, step={}", agentName, stepIndex);
}
}
private static String bounded(String value) {
int maxChars = 32_000;
return value.length() <= maxChars ? value : value.substring(0, maxChars);
private static String bound(String value) {
if (value == null) {
return null;
}
return value.length() <= MAX_TEXT_CHARS ? value : value.substring(0, MAX_TEXT_CHARS);
}
private static int utf8Bytes(String value) {
return value == null ? 0 : value.getBytes(StandardCharsets.UTF_8).length;
}
private static String joinNonBlank(String sep, String a, String b) {
boolean ha = hasText(a);
boolean hb = hasText(b);
if (ha && hb) {
return a + sep + b;
}
if (ha) {
return a;
}
if (hb) {
return b;
}
return null;
}
private static String firstNonBlank(String a, String b) {
if (hasText(a)) {
return a;
}
if (hasText(b)) {
return b;
}
return null;
}
private AuditIdentity identity(RunnableConfig config) {
@@ -233,10 +415,14 @@ public final class HarnessAgentAuditHook extends MessagesModelHook {
private record AuditIdentity(String sessionId, String runId) {
}
private record ReasoningContent(String content, boolean available, int bytes) {
private static ReasoningContent empty() {
return new ReasoningContent(null, false, 0);
}
record LlmTurnContent(
boolean reasoningAvailable,
String reasoningContent,
String assistantText,
String contentSource,
int reasoningBytes,
int assistantBytes,
int totalBytes) {
}
private record PendingStep(Long id, long startedNanos, Map<String, Object> input) {
@@ -1,66 +1,247 @@
package com.superbiz.agent.harness.audit;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.domain.entity.ToolInvocation;
import com.superbiz.agent.harness.contract.InvocationStatus;
import com.superbiz.agent.repository.ToolInvocationRepository;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
/**
* Persists tool invocation audit rows.
*
* <p>For {@code lookup_knowledge}, enriches RAG-specific columns from internal LookupResult
* JSON when present. {@code relevance_level} stores RAG PRECISE/REFERENCE (not evidence_status).
* {@code evidence_status} remains in {@code retrieval_details} for all tools.</p>
*/
@Component
public final class JpaToolInvocationAuditSink implements ToolInvocationAuditSink {
/** Max chars stored for free-text query / sql / topic style fields. */
private static final int MAX_TEXT_PREVIEW = 160;
private static final int MAX_REQUEST_KEYS = 12;
private static final Set<String> SENSITIVE_KEY_FRAGMENTS = Set.of(
"password", "passwd", "secret", "token", "apikey", "api_key",
"authorization", "credential", "private_key", "access_key");
private final ToolInvocationRepository repository;
private final ObjectMapper objectMapper;
private final DiagnosisTraceRecorder traceRecorder;
public JpaToolInvocationAuditSink(ToolInvocationRepository repository, ObjectMapper objectMapper) {
this(repository, objectMapper, DiagnosisTraceRecorder.noop());
}
private final RagLookupAuditEnricher ragEnricher;
@Autowired
public JpaToolInvocationAuditSink(ToolInvocationRepository repository, ObjectMapper objectMapper,
DiagnosisTraceRecorder traceRecorder) {
public JpaToolInvocationAuditSink(ToolInvocationRepository repository,
ObjectMapper objectMapper,
DiagnosisTraceRecorder traceRecorder,
ObjectProvider<RagLookupAuditEnricher> ragEnricherProvider) {
this.repository = Objects.requireNonNull(repository, "repository must not be null");
this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null");
this.traceRecorder = Objects.requireNonNull(traceRecorder, "traceRecorder must not be null");
this.ragEnricher = ragEnricherProvider == null ? null : ragEnricherProvider.getIfAvailable();
}
/** Test helper with explicit enricher (may be null). */
static JpaToolInvocationAuditSink forTest(ToolInvocationRepository repository,
ObjectMapper objectMapper,
DiagnosisTraceRecorder traceRecorder,
RagLookupAuditEnricher ragEnricher) {
return new JpaToolInvocationAuditSink(repository, objectMapper, traceRecorder, ragEnricher, true);
}
private JpaToolInvocationAuditSink(ToolInvocationRepository repository,
ObjectMapper objectMapper,
DiagnosisTraceRecorder traceRecorder,
RagLookupAuditEnricher ragEnricher,
boolean testMarker) {
this.repository = Objects.requireNonNull(repository, "repository must not be null");
this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null");
this.traceRecorder = Objects.requireNonNull(traceRecorder, "traceRecorder must not be null");
this.ragEnricher = ragEnricher;
}
@Override
public void record(ToolInvocationAuditEvent event) {
Objects.requireNonNull(event, "event must not be null");
traceRecorder.record(TraceAuditEvents.toolInvocation(event));
Map<String, Object> details = baseResultMetadata(event);
String retrievalLayer = "HARNESS";
Integer l0 = null;
Integer l1 = null;
String relevanceLevel = null;
Boolean truncated = false;
String dedupReason = null;
String outputPreview = "status=%s,evidence_status=%s".formatted(
event.status(), event.evidenceStatus());
if (ragEnricher != null && ragEnricher.supports(event.toolName())) {
RagLookupAuditEnricher.Enrichment enrichment =
ragEnricher.enrich(event.rawResultJson(), event.agentResultJson());
if (enrichment.retrievalDetails() != null) {
details.putAll(enrichment.retrievalDetails());
}
if (enrichment.retrievalLayer() != null && !enrichment.retrievalLayer().isBlank()) {
retrievalLayer = enrichment.retrievalLayer();
}
l0 = enrichment.l0MatchCount();
l1 = enrichment.l1MatchCount();
// RAG semantic level only — never overwrite with evidence_status enum names
relevanceLevel = enrichment.relevanceLevel();
if (enrichment.truncated() != null) {
truncated = enrichment.truncated();
}
dedupReason = enrichment.dedupReason();
if (enrichment.outputPreview() != null && !enrichment.outputPreview().isBlank()) {
outputPreview = enrichment.outputPreview()
+ ",status=" + event.status()
+ ",evidence_status=" + event.evidenceStatus();
}
}
// Always keep harness evidence_status in details (column relevance_level is RAG-only when set)
details.put("evidence_status", event.evidenceStatus().name());
details.put("invocation_status", event.status().name());
repository.save(ToolInvocation.builder()
.sessionId(event.sessionId())
.runId(event.runId())
.stepId(event.stepId())
.toolName(event.toolName())
.inputParams(write(inputMetadata(event)))
.outputPreview("status=%s,evidence_status=%s".formatted(
event.status(), event.evidenceStatus()))
.outputPreview(outputPreview)
.outputLength(event.agentResultBytes())
.retrievalLayer("HARNESS")
.isTruncated(false)
.relevanceLevel(event.evidenceStatus().name())
.retrievalDetails(write(resultMetadata(event)))
.retrievalLayer(retrievalLayer)
.l0MatchCount(l0)
.l1MatchCount(l1)
.isTruncated(Boolean.TRUE.equals(truncated))
.relevanceLevel(relevanceLevel)
.dedupReason(dedupReason)
.retrievalDetails(write(details))
.durationMs(event.durationMs())
.success(event.status() == InvocationStatus.READY)
.errorMessage(event.errorCode())
.build());
}
private Map<String, Object> inputMetadata(ToolInvocationAuditEvent event) {
/**
* Bounded request audit: always tool_call_id + request_bytes; optionally step_id and
* safe scalar fields from request JSON (query/topic/region/limit/...). Never stores
* password/token-like keys or nested blobs wholesale.
*/
Map<String, Object> inputMetadata(ToolInvocationAuditEvent event) {
Map<String, Object> metadata = new LinkedHashMap<>();
metadata.put("tool_call_id", event.toolCallId());
metadata.put("request_bytes", event.requestBytes());
if (event.stepId() != null) {
metadata.put("step_id", event.stepId());
}
appendSafeRequestFields(metadata, event.requestJson());
return metadata;
}
private Map<String, Object> resultMetadata(ToolInvocationAuditEvent event) {
private void appendSafeRequestFields(Map<String, Object> metadata, String requestJson) {
if (requestJson == null || requestJson.isBlank()) {
return;
}
try {
JsonNode root = objectMapper.readTree(requestJson);
if (root == null || !root.isObject()) {
return;
}
int added = 0;
Iterator<Map.Entry<String, JsonNode>> fields = root.fields();
while (fields.hasNext() && added < MAX_REQUEST_KEYS) {
Map.Entry<String, JsonNode> entry = fields.next();
String key = entry.getKey();
if (key == null || key.isBlank() || isSensitiveKey(key)) {
continue;
}
JsonNode value = entry.getValue();
if (value == null || value.isNull()) {
continue;
}
if (value.isTextual()) {
String text = value.asText();
if (text == null || text.isBlank()) {
continue;
}
metadata.put(key, truncate(text, MAX_TEXT_PREVIEW));
if (text.length() > MAX_TEXT_PREVIEW) {
metadata.put(key + "_truncated", true);
metadata.put(key + "_chars", text.length());
}
added++;
} else if (value.isNumber()) {
metadata.put(key, value.numberValue());
added++;
} else if (value.isBoolean()) {
metadata.put(key, value.booleanValue());
added++;
} else if (value.isArray() && isStringArray(value)) {
List<String> items = new ArrayList<>();
for (int i = 0; i < value.size() && items.size() < 8; i++) {
JsonNode item = value.get(i);
if (item != null && item.isTextual() && !item.asText().isBlank()) {
items.add(truncate(item.asText(), 64));
}
}
if (!items.isEmpty()) {
metadata.put(key, items);
added++;
}
}
// objects / mixed arrays intentionally omitted
}
} catch (Exception ignored) {
metadata.put("request_parse", "failed");
}
}
private static boolean isStringArray(JsonNode value) {
if (value == null || !value.isArray() || value.isEmpty()) {
return false;
}
for (JsonNode n : value) {
if (n == null || !n.isTextual()) {
return false;
}
}
return true;
}
private static boolean isSensitiveKey(String key) {
String normalized = key.toLowerCase(Locale.ROOT);
for (String fragment : SENSITIVE_KEY_FRAGMENTS) {
if (normalized.contains(fragment)) {
return true;
}
}
return false;
}
private static String truncate(String value, int maxChars) {
if (value == null) {
return null;
}
if (value.length() <= maxChars) {
return value;
}
return value.substring(0, maxChars);
}
private Map<String, Object> baseResultMetadata(ToolInvocationAuditEvent event) {
Map<String, Object> metadata = new LinkedHashMap<>();
metadata.put("tool_call_id", event.toolCallId());
metadata.put("status", event.status().name());
@@ -0,0 +1,335 @@
package com.superbiz.agent.harness.audit;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
/**
* Builds durable RAG audit fields from internal {@code LookupResult} JSON
* (and optional projected agent result) without dumping full excerpts/traces.
*/
@Component
public class RagLookupAuditEnricher {
public static final String TOOL_LOOKUP_KNOWLEDGE = "lookup_knowledge";
private static final int MAX_EVIDENCE_KEYS = 12;
private static final int MAX_HINT_VALUES = 8;
private static final int PREVIEW_CHARS = 160;
private final ObjectMapper objectMapper;
private final String searchMode;
public RagLookupAuditEnricher(
ObjectMapper objectMapper,
@Value("${retrieval.search.mode:hybrid}") String searchMode) {
this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper");
this.searchMode = searchMode == null || searchMode.isBlank()
? "hybrid"
: searchMode.trim().toLowerCase(Locale.ROOT);
}
public boolean supports(String toolName) {
return TOOL_LOOKUP_KNOWLEDGE.equals(toolName);
}
/**
* @param rawLookupJson internal LookupResult JSON (may be null on hard failures)
* @param agentResultJson projected RagToolResult JSON (may be null)
*/
public Enrichment enrich(String rawLookupJson, String agentResultJson) {
Map<String, Object> details = new LinkedHashMap<>();
details.put("audit_schema", "rag_lookup_v1");
details.put("search_mode", searchMode);
String retrievalLayer = "L1";
Integer l0MatchCount = null;
Integer l1MatchCount = null;
String relevanceLevel = null;
Boolean truncated = null;
String dedupReason = null;
String outputPreview = null;
try {
if (rawLookupJson != null && !rawLookupJson.isBlank()) {
JsonNode root = objectMapper.readTree(rawLookupJson);
if (root != null && root.isObject()) {
relevanceLevel = text(root, "relevanceLevel");
if (relevanceLevel == null) {
relevanceLevel = text(root, "relevance_level");
}
JsonNode blocks = root.get("evidenceBlocks");
if (blocks != null && blocks.isArray()) {
l1MatchCount = blocks.size();
List<String> keys = new ArrayList<>();
List<String> sources = new ArrayList<>();
String layer = null;
for (int i = 0; i < blocks.size() && keys.size() < MAX_EVIDENCE_KEYS; i++) {
JsonNode b = blocks.get(i);
if (b == null || !b.isObject()) {
continue;
}
String key = firstText(b, "evidenceKey", "evidence_key");
if (key != null) {
keys.add(key);
}
String source = text(b, "source");
if (source != null && sources.size() < MAX_EVIDENCE_KEYS) {
sources.add(source);
}
if (layer == null) {
layer = text(b, "retrievalLayer");
}
}
if (!keys.isEmpty()) {
details.put("evidence_keys", keys);
}
if (!sources.isEmpty()) {
details.put("sources", sources);
}
if (layer != null && !layer.isBlank()) {
retrievalLayer = layer;
}
}
Integer candidateCount = intVal(root, "evidenceCandidateCount");
if (candidateCount != null) {
details.put("evidence_candidate_count", candidateCount);
if (l1MatchCount == null) {
l1MatchCount = candidateCount;
}
}
Integer blockCount = intVal(root, "evidenceBlockCount");
if (blockCount != null) {
details.put("evidence_block_count", blockCount);
}
String completenessHint = text(root, "completenessHint");
if (completenessHint != null) {
details.put("completeness_hint", truncate(completenessHint, PREVIEW_CHARS));
}
JsonNode trace = root.get("retrievalTrace");
if (trace != null && trace.isObject()) {
putText(details, "selected_attempt", text(trace, "selectedAttempt"));
putText(details, "fallback_reason", text(trace, "fallbackReason"));
putText(details, "retrieval_evidence_status", text(trace, "evidenceStatus"));
putText(details, "category_filter", text(trace, "categoryFilter"));
// query texts intentionally omitted from durable audit by default (PII/size)
JsonNode hints = trace.get("queryHints");
if (hints != null && hints.isObject()) {
Integer l0 = intVal(hints, "l0_match_count");
if (l0 == null) {
l0 = arraySize(hints.get("domains"));
}
l0MatchCount = l0;
Map<String, Object> hintSnap = new LinkedHashMap<>();
putLimitedList(hintSnap, "domains", hints.get("domains"));
putLimitedList(hintSnap, "matched_keywords", hints.get("matched_keywords"));
if (hintSnap.isEmpty()) {
putLimitedList(hintSnap, "matched_keywords", hints.get("matchedKeywords"));
}
if (!hintSnap.isEmpty()) {
details.put("l0_hints", hintSnap);
}
}
JsonNode attempts = trace.get("attempts");
if (attempts != null && attempts.isArray()) {
List<Map<String, Object>> attemptSnap = new ArrayList<>();
for (JsonNode a : attempts) {
if (a == null || !a.isObject()) {
continue;
}
Map<String, Object> row = new LinkedHashMap<>();
putText(row, "name", text(a, "name"));
putText(row, "category_filter", text(a, "categoryFilter"));
if (a.has("candidateCount") && a.get("candidateCount").canConvertToInt()) {
row.put("candidate_count", a.get("candidateCount").asInt());
}
if (a.has("usable") && a.get("usable").isBoolean()) {
row.put("usable", a.get("usable").asBoolean());
}
if (a.has("topSimilarity") && a.get("topSimilarity").isNumber()) {
row.put("top_similarity", a.get("topSimilarity").asDouble());
}
if (a.has("durationMs") && a.get("durationMs").canConvertToInt()) {
row.put("duration_ms", a.get("durationMs").asInt());
}
if (!row.isEmpty()) {
attemptSnap.add(row);
}
}
if (!attemptSnap.isEmpty()) {
details.put("attempts", attemptSnap);
}
}
}
JsonNode pack = root.get("contextPack");
if (pack != null && pack.isObject()) {
if (pack.has("usedChars") && pack.get("usedChars").canConvertToInt()) {
details.put("context_used_chars", pack.get("usedChars").asInt());
}
if (pack.has("charBudget") && pack.get("charBudget").canConvertToInt()) {
details.put("context_char_budget", pack.get("charBudget").asInt());
}
}
// compact preview for list UIs
outputPreview = buildPreview(relevanceLevel, details);
}
}
} catch (Exception ignored) {
details.put("enrich_error", "lookup_result_parse_failed");
}
try {
if (agentResultJson != null && !agentResultJson.isBlank()) {
JsonNode agent = objectMapper.readTree(agentResultJson);
if (agent != null && agent.isObject()) {
if (agent.has("truncated") && agent.get("truncated").isBoolean()) {
truncated = agent.get("truncated").asBoolean();
details.put("truncated", truncated);
}
if (relevanceLevel == null) {
relevanceLevel = text(agent, "relevance_level");
if (relevanceLevel == null) {
relevanceLevel = text(agent, "relevanceLevel");
}
}
if (agent.has("returned_count") && agent.get("returned_count").canConvertToInt()) {
details.put("returned_count", agent.get("returned_count").asInt());
}
}
}
} catch (Exception ignored) {
details.put("agent_result_parse", "failed");
}
if (outputPreview == null) {
outputPreview = buildPreview(relevanceLevel, details);
}
return new Enrichment(
retrievalLayer,
l0MatchCount,
l1MatchCount,
relevanceLevel,
truncated,
dedupReason,
details,
outputPreview
);
}
private static String buildPreview(String relevanceLevel, Map<String, Object> details) {
String attempt = details.get("selected_attempt") == null ? null : String.valueOf(details.get("selected_attempt"));
String fallback = details.get("fallback_reason") == null ? null : String.valueOf(details.get("fallback_reason"));
StringBuilder sb = new StringBuilder("lookup_knowledge");
if (relevanceLevel != null) {
sb.append(" level=").append(relevanceLevel);
}
if (attempt != null) {
sb.append(" attempt=").append(attempt);
}
if (fallback != null) {
sb.append(" fallback=").append(fallback);
}
Object keys = details.get("evidence_keys");
if (keys instanceof List<?> list) {
sb.append(" keys=").append(list.size());
}
return truncate(sb.toString(), PREVIEW_CHARS);
}
private static void putLimitedList(Map<String, Object> target, String key, JsonNode node) {
if (node == null || !node.isArray() || node.isEmpty()) {
return;
}
List<String> values = new ArrayList<>();
for (int i = 0; i < node.size() && values.size() < MAX_HINT_VALUES; i++) {
JsonNode n = node.get(i);
if (n != null && n.isTextual() && !n.asText().isBlank()) {
values.add(n.asText());
}
}
if (!values.isEmpty()) {
target.put(key, values);
}
}
private static void putText(Map<String, Object> map, String key, String value) {
if (value != null && !value.isBlank()) {
map.put(key, value);
}
}
private static String firstText(JsonNode node, String... fields) {
for (String f : fields) {
String v = text(node, f);
if (v != null) {
return v;
}
}
return null;
}
private static String text(JsonNode node, String field) {
if (node == null || field == null || !node.has(field) || node.get(field).isNull()) {
return null;
}
String v = node.get(field).asText(null);
return v == null || v.isBlank() ? null : v;
}
private static Integer intVal(JsonNode node, String field) {
if (node == null || !node.has(field) || node.get(field).isNull()) {
return null;
}
JsonNode n = node.get(field);
if (n.isIntegralNumber() || n.canConvertToInt()) {
return n.asInt();
}
return null;
}
private static Integer arraySize(JsonNode node) {
if (node != null && node.isArray()) {
return node.size();
}
return null;
}
private static String truncate(String value, int max) {
if (value == null) {
return null;
}
if (value.length() <= max) {
return value;
}
return value.substring(0, max) + "...";
}
public record Enrichment(
String retrievalLayer,
Integer l0MatchCount,
Integer l1MatchCount,
String relevanceLevel,
Boolean truncated,
String dedupReason,
Map<String, Object> retrievalDetails,
String outputPreview
) {
}
}
@@ -0,0 +1,114 @@
package com.superbiz.agent.harness.audit;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
/**
* Extracts a short human-readable conclusion string from released safe content JSON
* for {@code diagnosis_run.conclusion}.
*/
public final class RunConclusionExtractor {
private static final int MAX_CHARS = 4_000;
private RunConclusionExtractor() {
}
public static String extract(ObjectMapper objectMapper, String safeContentJson) {
if (safeContentJson == null || safeContentJson.isBlank()) {
return null;
}
try {
JsonNode root = objectMapper.readTree(safeContentJson);
if (root == null || !root.isObject()) {
return bound(safeContentJson.trim());
}
// DIAGNOSIS_REPORT payload shape stored in answer
String fromReport = text(root.path("report").path("conclusion").path("text"));
if (fromReport != null) {
return bound(fromReport);
}
// nested content envelope (defensive)
String nested = text(root.path("payload").path("report").path("conclusion").path("text"));
if (nested != null) {
return bound(nested);
}
// SAFE_FALLBACK envelope: {"fallback":{type,message,...}}
String fromFallback = fallbackConclusion(root.path("fallback"));
if (fromFallback != null) {
return bound(fromFallback);
}
String nestedFb = fallbackConclusion(root.path("payload").path("fallback"));
if (nestedFb != null) {
return bound(nestedFb);
}
// bare SafeFallback object (defensive)
if (root.hasNonNull("type") || root.has("message") || root.has("conclusion")) {
String bare = fallbackConclusion(root);
if (bare != null) {
return bound(bare);
}
}
// knowledge / plain answer shapes
String plain = text(root.path("answer"));
if (plain != null) {
return bound(plain);
}
return null;
} catch (Exception ignored) {
return bound(safeContentJson.trim());
}
}
/**
* Builds a short conclusion from a SafeFallback-shaped node.
* Prefer explicit conclusion text, else "type: message", else either alone.
*/
private static String fallbackConclusion(JsonNode node) {
if (node == null || node.isMissingNode() || node.isNull() || !node.isObject()) {
return null;
}
String conclusion = text(node.path("conclusion"));
if (conclusion != null) {
return conclusion;
}
String type = enumOrText(node.path("type"));
String message = text(node.path("message"));
if (type != null && message != null) {
return type + ": " + message;
}
if (message != null) {
return message;
}
return type;
}
private static String enumOrText(JsonNode node) {
if (node == null || node.isMissingNode() || node.isNull()) {
return null;
}
if (node.isTextual() || node.isNumber() || node.isBoolean()) {
String value = node.asText();
return value == null || value.isBlank() ? null : value.trim();
}
return null;
}
private static String text(JsonNode node) {
if (node == null || node.isMissingNode() || node.isNull()) {
return null;
}
if (!node.isTextual()) {
return null;
}
String value = node.asText();
return value == null || value.isBlank() ? null : value.trim();
}
private static String bound(String value) {
if (value == null) {
return null;
}
return value.length() <= MAX_CHARS ? value : value.substring(0, MAX_CHARS);
}
}
@@ -5,6 +5,17 @@ import com.superbiz.agent.harness.contract.InvocationStatus;
import java.util.Objects;
/**
* Durable tool-invocation audit event.
*
* <p>{@code rawResultJson} is optional internal executor output (e.g. full {@code LookupResult}
* before Agent projection). It is used only to enrich durable RAG fields and must not be
* echoed wholesale into agent-facing views.</p>
*
* <p>{@code stepId} links to {@code agent_step.id} when the in-flight Agent step is known.
* {@code requestJson} is the tool request envelope body used only to derive bounded
* {@code input_params} (e.g. query preview) — never dump secrets wholesale.</p>
*/
public record ToolInvocationAuditEvent(
String sessionId,
String runId,
@@ -15,7 +26,11 @@ public record ToolInvocationAuditEvent(
String errorCode,
int durationMs,
int requestBytes,
int agentResultBytes) {
int agentResultBytes,
String rawResultJson,
String agentResultJson,
Long stepId,
String requestJson) {
public ToolInvocationAuditEvent {
requireText(sessionId, "sessionId");
@@ -35,6 +50,38 @@ public record ToolInvocationAuditEvent(
}
}
/** Backward-compatible constructor without raw/agent JSON / step / request body. */
public ToolInvocationAuditEvent(String sessionId,
String runId,
String toolCallId,
String toolName,
InvocationStatus status,
EvidenceStatus evidenceStatus,
String errorCode,
int durationMs,
int requestBytes,
int agentResultBytes) {
this(sessionId, runId, toolCallId, toolName, status, evidenceStatus, errorCode,
durationMs, requestBytes, agentResultBytes, null, null, null, null);
}
/** Backward-compatible constructor with raw/agent JSON only. */
public ToolInvocationAuditEvent(String sessionId,
String runId,
String toolCallId,
String toolName,
InvocationStatus status,
EvidenceStatus evidenceStatus,
String errorCode,
int durationMs,
int requestBytes,
int agentResultBytes,
String rawResultJson,
String agentResultJson) {
this(sessionId, runId, toolCallId, toolName, status, evidenceStatus, errorCode,
durationMs, requestBytes, agentResultBytes, rawResultJson, agentResultJson, null, null);
}
private static void requireText(String value, String name) {
if (value == null || value.isBlank()) {
throw new IllegalArgumentException(name + " must not be blank");
@@ -129,9 +129,14 @@ public final class TraceAuditEvents {
details.put("evidence_status", tool.evidenceStatus().name());
details.put("request_bytes", tool.requestBytes());
details.put("agent_result_bytes", tool.agentResultBytes());
details.put("has_raw_result", tool.rawResultJson() != null && !tool.rawResultJson().isBlank());
if (tool.stepId() != null) {
details.put("step_id", tool.stepId());
}
if (tool.errorCode() != null) {
details.put("error_code", tool.errorCode());
}
// Do not embed raw LookupResult / agent JSON here — durable RAG fields go to tool_invocation.
return new DiagnosisTraceAuditEvent(
tool.sessionId(), tool.runId(), TracePhase.TOOL, TraceEventType.TOOL_INVOCATION,
tool.status() == com.superbiz.agent.harness.contract.InvocationStatus.READY
@@ -3,9 +3,9 @@ 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.audit.AgentStepAuditTracker;
import com.superbiz.agent.harness.audit.ToolInvocationAuditEvent;
import com.superbiz.agent.harness.audit.ToolInvocationAuditSink;
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;
@@ -19,10 +19,10 @@ import com.superbiz.agent.harness.tool.store.ToolCallKeyFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.nio.charset.StandardCharsets;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
public final class ToolBoundary {
@@ -35,13 +35,14 @@ public final class ToolBoundary {
private final ObjectMapper objectMapper;
private final Clock clock;
private final ToolInvocationAuditSink auditSink;
private final AgentStepAuditTracker stepTracker;
public ToolBoundary(DiagnosisHarnessCore core,
ToolCallKeyFactory keyFactory,
CanonicalInvocationStore store,
ObjectMapper objectMapper,
Clock clock) {
this(core, keyFactory, store, objectMapper, clock, ToolInvocationAuditSink.noop());
this(core, keyFactory, store, objectMapper, clock, ToolInvocationAuditSink.noop(), null);
}
public ToolBoundary(DiagnosisHarnessCore core,
@@ -50,12 +51,23 @@ public final class ToolBoundary {
ObjectMapper objectMapper,
Clock clock,
ToolInvocationAuditSink auditSink) {
this(core, keyFactory, store, objectMapper, clock, auditSink, null);
}
public ToolBoundary(DiagnosisHarnessCore core,
ToolCallKeyFactory keyFactory,
CanonicalInvocationStore store,
ObjectMapper objectMapper,
Clock clock,
ToolInvocationAuditSink auditSink,
AgentStepAuditTracker stepTracker) {
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");
this.auditSink = Objects.requireNonNull(auditSink, "auditSink must not be null");
this.stepTracker = stepTracker;
}
public ToolBoundaryResult execute(RunContext context,
@@ -63,15 +75,15 @@ public final class ToolBoundary {
ToolExecutor executor,
ToolResultProjector projector) {
Instant startedAt = clock.instant();
ToolBoundaryResult result = executeCanonical(context, request, executor, projector);
auditSafely(context, request, result, startedAt);
return result;
ExecutionOutcome outcome = executeCanonical(context, request, executor, projector);
auditSafely(context, request, outcome.result(), outcome.rawResponse(), startedAt);
return outcome.result();
}
private ToolBoundaryResult executeCanonical(RunContext context,
ToolCallRequestEnvelope request,
ToolExecutor executor,
ToolResultProjector projector) {
private ExecutionOutcome executeCanonical(RunContext context,
ToolCallRequestEnvelope request,
ToolExecutor executor,
ToolResultProjector projector) {
String toolCallId = request == null ? null : request.toolCallId();
String key;
try {
@@ -79,23 +91,23 @@ public final class ToolBoundary {
core.beforeToolCall(context, request.toolName());
long requestBytes = store.limits().utf8Bytes(request.requestJson());
if (requestBytes > store.limits().maxRecordBytes()) {
return errorAndNoRecord(toolCallId, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
return ExecutionOutcome.of(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);
return ExecutionOutcome.of(ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.DUPLICATE_TOOL_CALL));
} catch (RunAbortedException | BudgetExceededException e) {
return ToolBoundaryResult.error(toolCallId,
return ExecutionOutcome.of(ToolBoundaryResult.error(toolCallId,
e instanceof BudgetExceededException
? ToolBoundaryErrorCode.BUDGET_EXHAUSTED
: ToolBoundaryErrorCode.RUN_INACTIVE);
: ToolBoundaryErrorCode.RUN_INACTIVE));
} catch (IllegalArgumentException e) {
return ToolBoundaryResult.error(toolCallId, classifyPreflightError(e));
return ExecutionOutcome.of(ToolBoundaryResult.error(toolCallId, classifyPreflightError(e)));
} catch (CanonicalStoreException e) {
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.STORE_ERROR);
return ExecutionOutcome.of(ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.STORE_ERROR));
}
String rawResponse;
@@ -107,7 +119,7 @@ public final class ToolBoundary {
}
} catch (Exception e) {
markErrorSafely(key, null, ToolBoundaryErrorCode.TOOL_EXECUTION_ERROR);
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.TOOL_EXECUTION_ERROR);
return ExecutionOutcome.of(ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.TOOL_EXECUTION_ERROR));
}
try {
@@ -115,13 +127,16 @@ public final class ToolBoundary {
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);
return new ExecutionOutcome(
ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RESULT_TOO_LARGE), rawResponse);
} catch (BudgetExceededException e) {
markErrorSafely(key, null, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
return new ExecutionOutcome(
ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.BUDGET_EXHAUSTED), rawResponse);
} catch (RunAbortedException e) {
markErrorSafely(key, null, ToolBoundaryErrorCode.RUN_INACTIVE);
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RUN_INACTIVE);
return new ExecutionOutcome(
ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RUN_INACTIVE), rawResponse);
}
ProjectedToolResult projected;
@@ -133,29 +148,35 @@ public final class ToolBoundary {
}
} catch (Exception e) {
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.PROJECTION_ERROR);
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.PROJECTION_ERROR);
return new ExecutionOutcome(
ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.PROJECTION_ERROR), rawResponse);
}
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());
return new ExecutionOutcome(
ToolBoundaryResult.ready(toolCallId, projected.agentResult(), projected.evidenceStatus()),
rawResponse);
} catch (ResultTooLargeException e) {
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RESULT_TOO_LARGE);
return new ExecutionOutcome(
ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RESULT_TOO_LARGE), rawResponse);
} catch (BudgetExceededException e) {
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.BUDGET_EXHAUSTED);
return new ExecutionOutcome(
ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.BUDGET_EXHAUSTED), rawResponse);
} catch (RunAbortedException e) {
markErrorSafely(key, rawResponse, ToolBoundaryErrorCode.RUN_INACTIVE);
return ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RUN_INACTIVE);
return new ExecutionOutcome(
ToolBoundaryResult.error(toolCallId, ToolBoundaryErrorCode.RUN_INACTIVE), rawResponse);
} catch (CanonicalStoreException e) {
ToolBoundaryErrorCode code = e instanceof ResultTooLargeException
? ToolBoundaryErrorCode.RESULT_TOO_LARGE
: ToolBoundaryErrorCode.PROJECTION_ERROR;
markErrorSafely(key, rawResponse, code);
return ToolBoundaryResult.error(toolCallId, code);
return new ExecutionOutcome(ToolBoundaryResult.error(toolCallId, code), rawResponse);
}
}
@@ -222,6 +243,7 @@ public final class ToolBoundary {
private void auditSafely(RunContext context,
ToolCallRequestEnvelope request,
ToolBoundaryResult result,
String rawResponse,
Instant startedAt) {
if (context == null || request == null || result == null
|| !context.runId().equals(request.runId())
@@ -230,11 +252,16 @@ public final class ToolBoundary {
return;
}
try {
Long stepId = stepTracker == null ? null : stepTracker.currentStepId(context.runId());
auditSink.record(new ToolInvocationAuditEvent(
context.sessionId(), context.runId(), request.toolCallId(), request.toolName(),
result.status(), result.evidenceStatus(), result.errorCode(),
saturatingInt(Math.max(0L, Duration.between(startedAt, clock.instant()).toMillis())),
utf8Bytes(request.requestJson()), utf8Bytes(result.agentResult())));
utf8Bytes(request.requestJson()), utf8Bytes(result.agentResult()),
rawResponse,
result.agentResult(),
stepId,
request.requestJson()));
} catch (RuntimeException exception) {
log.warn("Failed to persist Tool durable audit: tool={}, status={}",
request.toolName(), result.status());
@@ -249,6 +276,12 @@ public final class ToolBoundary {
return value >= Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) value;
}
private record ExecutionOutcome(ToolBoundaryResult result, String rawResponse) {
static ExecutionOutcome of(ToolBoundaryResult result) {
return new ExecutionOutcome(result, null);
}
}
private static final class InvalidToolCallIdException extends IllegalArgumentException {
private InvalidToolCallIdException(Throwable cause) {
super("Invalid tool call ID", cause);