feat(harness): add run context and retry core

2026-07-21 18:36:19 +08:00
parent 4274f3350b
commit 6b74990f86
47 changed files with 1910 additions and 1 deletions
@@ -0,0 +1,27 @@
+package com.superbiz.agent.harness.core;
+
+public final class BudgetExceededException extends RuntimeException {
+
+    private final BudgetKind kind;
+    private final long limit;
+    private final long attempted;
+
+    public BudgetExceededException(BudgetKind kind, long limit, long attempted) {
+        super("Budget exceeded: kind=" + kind + ", limit=" + limit + ", attempted=" + attempted);
+        this.kind = kind;
+        this.limit = limit;
+        this.attempted = attempted;
+    }
+
+    public BudgetKind kind() {
+        return kind;
+    }
+
+    public long limit() {
+        return limit;
+    }
+
+    public long attempted() {
+        return attempted;
+    }
+}
@@ -0,0 +1,11 @@
+package com.superbiz.agent.harness.core;
+
+public enum BudgetKind {
+    MODEL_CALLS,
+    TOOL_CALLS,
+    TOOL_CALLS_PER_TOOL,
+    INPUT_TOKENS,
+    OUTPUT_TOKENS,
+    TOTAL_TOKENS,
+    RUN_BYTES
+}
@@ -0,0 +1,147 @@
+package com.superbiz.agent.harness.core;
+
+import com.superbiz.agent.harness.retry.HarnessRetryPolicies;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Objects;
+import java.util.function.Supplier;
+
+public final class DiagnosisHarnessCore {
+
+    private final Clock clock;
+    private final Supplier<String> runIdSupplier;
+    private final Duration maxRunDuration;
+    private final RunBudgetLimits budgetLimits;
+    private final HarnessRetryPolicies retryPolicies;
+
+    public DiagnosisHarnessCore(Clock clock,
+                                Supplier<String> runIdSupplier,
+                                Duration maxRunDuration,
+                                RunBudgetLimits budgetLimits,
+                                HarnessRetryPolicies retryPolicies) {
+        this.clock = Objects.requireNonNull(clock, "clock must not be null");
+        this.runIdSupplier = Objects.requireNonNull(runIdSupplier, "runIdSupplier must not be null");
+        this.maxRunDuration = requirePositive(maxRunDuration, "maxRunDuration");
+        this.budgetLimits = Objects.requireNonNull(budgetLimits, "budgetLimits must not be null");
+        this.retryPolicies = Objects.requireNonNull(retryPolicies, "retryPolicies must not be null");
+    }
+
+    public RunContext startRun(String sessionId) {
+        return startRun(sessionId, runIdSupplier.get());
+    }
+
+    public RunContext startRun(String sessionId, String runId) {
+        Instant deadline = clock.instant().plus(maxRunDuration);
+        RunCancellation cancellation = new RunCancellation();
+        RunLifecycle lifecycle = new RunLifecycle(clock);
+        RunContext context = new RunContext(
+                sessionId,
+                runId,
+                deadline,
+                cancellation,
+                new RunBudget(budgetLimits),
+                retryPolicies,
+                lifecycle);
+        cancellation.onCancel(reason -> lifecycle.finish(terminalState(reason), reason.name()));
+        return context;
+    }
+
+    public void checkActive(RunContext context) {
+        Objects.requireNonNull(context, "context must not be null");
+        context.lifecycle().termination().ifPresent(termination -> {
+            throw new RunAbortedException(termination);
+        });
+
+        if (!clock.instant().isBefore(context.deadline())) {
+            context.lifecycle().finish(RunState.TIMED_OUT, RunCancellationReason.DEADLINE_EXCEEDED.name());
+            context.cancellation().cancel(RunCancellationReason.DEADLINE_EXCEEDED);
+            throw new RunAbortedException(context.lifecycle().termination().orElseThrow());
+        }
+
+        if (context.cancellation().isCancelled()) {
+            throw new RunAbortedException(context.lifecycle().termination().orElseThrow());
+        }
+    }
+
+    public void beforeModelCall(RunContext context) {
+        checkActive(context);
+        applyBudget(context, context.budget()::reserveModelCall);
+    }
+
+    public void beforeToolCall(RunContext context, String toolName) {
+        checkActive(context);
+        applyBudget(context, () -> context.budget().reserveToolCall(toolName));
+    }
+
+    public void recordTokens(RunContext context, long inputTokens, long outputTokens) {
+        Objects.requireNonNull(context, "context must not be null");
+        applyBudget(context, () -> context.budget().recordTokens(inputTokens, outputTokens));
+    }
+
+    public long reserveRunBytes(RunContext context, long bytes) {
+        checkActive(context);
+        try {
+            return context.budget().reserveRunBytes(bytes);
+        } catch (BudgetExceededException e) {
+            exhaustBudget(context, e);
+            throw e;
+        }
+    }
+
+    public boolean cancel(RunContext context, RunCancellationReason reason) {
+        Objects.requireNonNull(context, "context must not be null");
+        Objects.requireNonNull(reason, "reason must not be null");
+        if (context.lifecycle().termination().isPresent()) {
+            return false;
+        }
+        return context.cancellation().cancel(reason);
+    }
+
+    public boolean completeSuccess(RunContext context) {
+        Objects.requireNonNull(context, "context must not be null");
+        return context.lifecycle().finish(RunState.SUCCESS, "COMPLETED");
+    }
+
+    public boolean completeFailure(RunContext context, String reason) {
+        Objects.requireNonNull(context, "context must not be null");
+        String safeReason = reason == null || reason.isBlank() ? "INTERNAL_FAILURE" : reason;
+        boolean completed = context.lifecycle().finish(RunState.FAILED, safeReason);
+        if (completed) {
+            context.cancellation().cancel(RunCancellationReason.INTERNAL_FAILURE);
+        }
+        return completed;
+    }
+
+    private void applyBudget(RunContext context, Runnable operation) {
+        try {
+            operation.run();
+        } catch (BudgetExceededException e) {
+            exhaustBudget(context, e);
+            throw e;
+        }
+    }
+
+    private void exhaustBudget(RunContext context, BudgetExceededException exception) {
+        context.lifecycle().finish(RunState.BUDGET_EXHAUSTED, exception.getMessage());
+        context.cancellation().cancel(RunCancellationReason.BUDGET_EXHAUSTED);
+    }
+
+    private static RunState terminalState(RunCancellationReason reason) {
+        return switch (reason) {
+            case DEADLINE_EXCEEDED -> RunState.TIMED_OUT;
+            case BUDGET_EXHAUSTED -> RunState.BUDGET_EXHAUSTED;
+            case INTERNAL_FAILURE -> RunState.FAILED;
+            case CLIENT_DISCONNECTED, USER_REQUESTED -> RunState.CANCELLED;
+        };
+    }
+
+    private static Duration requirePositive(Duration duration, String name) {
+        Objects.requireNonNull(duration, name + " must not be null");
+        if (duration.isZero() || duration.isNegative()) {
+            throw new IllegalArgumentException(name + " must be positive");
+        }
+        return duration;
+    }
+}
@@ -0,0 +1,18 @@
+package com.superbiz.agent.harness.core;
+
+import java.util.Objects;
+
+public final class RunAbortedException extends RuntimeException {
+
+    private final RunTermination termination;
+
+    public RunAbortedException(RunTermination termination) {
+        super("Run is no longer active: state=" + Objects.requireNonNull(termination).state()
+                + ", reason=" + termination.reason());
+        this.termination = termination;
+    }
+
+    public RunTermination termination() {
+        return termination;
+    }
+}
@@ -0,0 +1,98 @@
+package com.superbiz.agent.harness.core;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+public final class RunBudget {
+
+    private final RunBudgetLimits limits;
+    private final RunCapacityCounter capacity;
+    private final Map<String, Integer> toolCallsByName = new HashMap<>();
+
+    private int modelCalls;
+    private int toolCalls;
+    private long inputTokens;
+    private long outputTokens;
+    private long totalTokens;
+
+    public RunBudget(RunBudgetLimits limits) {
+        this.limits = Objects.requireNonNull(limits, "limits must not be null");
+        this.capacity = new RunCapacityCounter(limits.maxRunBytes());
+    }
+
+    public synchronized void reserveModelCall() {
+        int attempted = modelCalls + 1;
+        if (attempted > limits.maxModelCalls()) {
+            throw new BudgetExceededException(BudgetKind.MODEL_CALLS, limits.maxModelCalls(), attempted);
+        }
+        modelCalls = attempted;
+    }
+
+    public synchronized void reserveToolCall(String toolName) {
+        if (toolName == null || toolName.isBlank()) {
+            throw new IllegalArgumentException("toolName must not be blank");
+        }
+        int attemptedTotal = toolCalls + 1;
+        int attemptedForTool = toolCallsByName.getOrDefault(toolName, 0) + 1;
+        if (attemptedTotal > limits.maxToolCalls()) {
+            throw new BudgetExceededException(BudgetKind.TOOL_CALLS, limits.maxToolCalls(), attemptedTotal);
+        }
+        if (attemptedForTool > limits.maxCallsPerTool()) {
+            throw new BudgetExceededException(
+                    BudgetKind.TOOL_CALLS_PER_TOOL, limits.maxCallsPerTool(), attemptedForTool);
+        }
+        toolCalls = attemptedTotal;
+        toolCallsByName.put(toolName, attemptedForTool);
+    }
+
+    public synchronized void recordTokens(long input, long output) {
+        if (input < 0 || output < 0) {
+            throw new IllegalArgumentException("token counts must not be negative");
+        }
+        inputTokens = safeAdd(inputTokens, input);
+        outputTokens = safeAdd(outputTokens, output);
+        totalTokens = safeAdd(totalTokens, safeAdd(input, output));
+
+        if (inputTokens > limits.maxInputTokens()) {
+            throw new BudgetExceededException(BudgetKind.INPUT_TOKENS, limits.maxInputTokens(), inputTokens);
+        }
+        if (outputTokens > limits.maxOutputTokens()) {
+            throw new BudgetExceededException(BudgetKind.OUTPUT_TOKENS, limits.maxOutputTokens(), outputTokens);
+        }
+        if (totalTokens > limits.maxTotalTokens()) {
+            throw new BudgetExceededException(BudgetKind.TOTAL_TOKENS, limits.maxTotalTokens(), totalTokens);
+        }
+    }
+
+    public long reserveRunBytes(long bytes) {
+        return capacity.reserve(bytes);
+    }
+
+    public synchronized RunBudgetUsage snapshot() {
+        return new RunBudgetUsage(
+                modelCalls,
+                toolCalls,
+                toolCallsByName,
+                inputTokens,
+                outputTokens,
+                totalTokens,
+                capacity.usedBytes());
+    }
+
+    public RunBudgetLimits limits() {
+        return limits;
+    }
+
+    public RunCapacityCounter capacity() {
+        return capacity;
+    }
+
+    private static long safeAdd(long left, long right) {
+        try {
+            return Math.addExact(left, right);
+        } catch (ArithmeticException e) {
+            return Long.MAX_VALUE;
+        }
+    }
+}
@@ -0,0 +1,27 @@
+package com.superbiz.agent.harness.core;
+
+public record RunBudgetLimits(
+        int maxModelCalls,
+        int maxToolCalls,
+        int maxCallsPerTool,
+        long maxInputTokens,
+        long maxOutputTokens,
+        long maxTotalTokens,
+        long maxRunBytes) {
+
+    public RunBudgetLimits {
+        requirePositive(maxModelCalls, "maxModelCalls");
+        requirePositive(maxToolCalls, "maxToolCalls");
+        requirePositive(maxCallsPerTool, "maxCallsPerTool");
+        requirePositive(maxInputTokens, "maxInputTokens");
+        requirePositive(maxOutputTokens, "maxOutputTokens");
+        requirePositive(maxTotalTokens, "maxTotalTokens");
+        requirePositive(maxRunBytes, "maxRunBytes");
+    }
+
+    private static void requirePositive(long value, String name) {
+        if (value <= 0) {
+            throw new IllegalArgumentException(name + " must be positive");
+        }
+    }
+}
@@ -0,0 +1,17 @@
+package com.superbiz.agent.harness.core;
+
+import java.util.Map;
+
+public record RunBudgetUsage(
+        int modelCalls,
+        int toolCalls,
+        Map<String, Integer> toolCallsByName,
+        long inputTokens,
+        long outputTokens,
+        long totalTokens,
+        long runBytes) {
+
+    public RunBudgetUsage {
+        toolCallsByName = toolCallsByName == null ? Map.of() : Map.copyOf(toolCallsByName);
+    }
+}
@@ -0,0 +1,64 @@
+package com.superbiz.agent.harness.core;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+public final class RunCancellation {
+
+    private static final Logger log = LoggerFactory.getLogger(RunCancellation.class);
+
+    private final AtomicReference<RunCancellationReason> reason = new AtomicReference<>();
+    private final CopyOnWriteArrayList<Consumer<RunCancellationReason>> callbacks =
+            new CopyOnWriteArrayList<>();
+
+    public boolean cancel(RunCancellationReason cancellationReason) {
+        Objects.requireNonNull(cancellationReason, "cancellationReason must not be null");
+        if (!reason.compareAndSet(null, cancellationReason)) {
+            return false;
+        }
+        callbacks.forEach(callback -> notifyCallback(callback, cancellationReason));
+        callbacks.clear();
+        return true;
+    }
+
+    public boolean isCancelled() {
+        return reason.get() != null;
+    }
+
+    public Optional<RunCancellationReason> reason() {
+        return Optional.ofNullable(reason.get());
+    }
+
+    public void onCancel(Consumer<RunCancellationReason> callback) {
+        Objects.requireNonNull(callback, "callback must not be null");
+        AtomicBoolean invoked = new AtomicBoolean();
+        Consumer<RunCancellationReason> once = cancellationReason -> {
+            if (invoked.compareAndSet(false, true)) {
+                callback.accept(cancellationReason);
+            }
+        };
+        callbacks.add(once);
+
+        RunCancellationReason current = reason.get();
+        if (current != null) {
+            notifyCallback(once, current);
+            callbacks.remove(once);
+        }
+    }
+
+    private void notifyCallback(Consumer<RunCancellationReason> callback,
+                                RunCancellationReason cancellationReason) {
+        try {
+            callback.accept(cancellationReason);
+        } catch (RuntimeException e) {
+            log.warn("Run cancellation callback failed: reason={}", cancellationReason, e);
+        }
+    }
+}
@@ -0,0 +1,9 @@
+package com.superbiz.agent.harness.core;
+
+public enum RunCancellationReason {
+    CLIENT_DISCONNECTED,
+    USER_REQUESTED,
+    DEADLINE_EXCEEDED,
+    BUDGET_EXHAUSTED,
+    INTERNAL_FAILURE
+}
@@ -0,0 +1,48 @@
+package com.superbiz.agent.harness.core;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+public final class RunCapacityCounter {
+
+    private final long maxBytes;
+    private final AtomicLong usedBytes = new AtomicLong();
+
+    public RunCapacityCounter(long maxBytes) {
+        if (maxBytes <= 0) {
+            throw new IllegalArgumentException("maxBytes must be positive");
+        }
+        this.maxBytes = maxBytes;
+    }
+
+    public long reserve(long bytes) {
+        if (bytes <= 0) {
+            throw new IllegalArgumentException("bytes must be positive");
+        }
+        while (true) {
+            long current = usedBytes.get();
+            long attempted = safeAdd(current, bytes);
+            if (attempted > maxBytes) {
+                throw new BudgetExceededException(BudgetKind.RUN_BYTES, maxBytes, attempted);
+            }
+            if (usedBytes.compareAndSet(current, attempted)) {
+                return attempted;
+            }
+        }
+    }
+
+    public long usedBytes() {
+        return usedBytes.get();
+    }
+
+    public long maxBytes() {
+        return maxBytes;
+    }
+
+    private static long safeAdd(long left, long right) {
+        try {
+            return Math.addExact(left, right);
+        } catch (ArithmeticException e) {
+            return Long.MAX_VALUE;
+        }
+    }
+}
@@ -0,0 +1,35 @@
+package com.superbiz.agent.harness.core;
+
+import com.superbiz.agent.harness.retry.HarnessRetryPolicies;
+
+import java.time.Instant;
+import java.util.Objects;
+
+/**
+ * Structurally immutable context. Mutable per-run state lives in thread-safe handles.
+ */
+public record RunContext(
+        String sessionId,
+        String runId,
+        Instant deadline,
+        RunCancellation cancellation,
+        RunBudget budget,
+        HarnessRetryPolicies retryPolicies,
+        RunLifecycle lifecycle) {
+
+    public RunContext {
+        requireText(sessionId, "sessionId");
+        requireText(runId, "runId");
+        Objects.requireNonNull(deadline, "deadline must not be null");
+        Objects.requireNonNull(cancellation, "cancellation must not be null");
+        Objects.requireNonNull(budget, "budget must not be null");
+        Objects.requireNonNull(retryPolicies, "retryPolicies must not be null");
+        Objects.requireNonNull(lifecycle, "lifecycle must not be null");
+    }
+
+    private static void requireText(String value, String name) {
+        if (value == null || value.isBlank()) {
+            throw new IllegalArgumentException(name + " must not be blank");
+        }
+    }
+}
@@ -0,0 +1,29 @@
+package com.superbiz.agent.harness.core;
+
+import java.time.Clock;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+
+public final class RunLifecycle {
+
+    private final Clock clock;
+    private final AtomicReference<RunTermination> termination = new AtomicReference<>();
+
+    public RunLifecycle(Clock clock) {
+        this.clock = Objects.requireNonNull(clock, "clock must not be null");
+    }
+
+    public RunState state() {
+        RunTermination current = termination.get();
+        return current == null ? RunState.RUNNING : current.state();
+    }
+
+    public Optional<RunTermination> termination() {
+        return Optional.ofNullable(termination.get());
+    }
+
+    public boolean finish(RunState state, String reason) {
+        return termination.compareAndSet(null, new RunTermination(state, reason, clock.instant()));
+    }
+}
@@ -0,0 +1,20 @@
+package com.superbiz.agent.harness.core;
+
+public enum RunState {
+    RUNNING(false),
+    SUCCESS(true),
+    FAILED(true),
+    CANCELLED(true),
+    TIMED_OUT(true),
+    BUDGET_EXHAUSTED(true);
+
+    private final boolean terminal;
+
+    RunState(boolean terminal) {
+        this.terminal = terminal;
+    }
+
+    public boolean isTerminal() {
+        return terminal;
+    }
+}
@@ -0,0 +1,18 @@
+package com.superbiz.agent.harness.core;
+
+import java.time.Instant;
+import java.util.Objects;
+
+public record RunTermination(RunState state, String reason, Instant completedAt) {
+
+    public RunTermination {
+        Objects.requireNonNull(state, "state must not be null");
+        Objects.requireNonNull(completedAt, "completedAt must not be null");
+        if (!state.isTerminal()) {
+            throw new IllegalArgumentException("state must be terminal");
+        }
+        if (reason == null || reason.isBlank()) {
+            throw new IllegalArgumentException("reason must not be blank");
+        }
+    }
+}
@@ -0,0 +1,60 @@
+package com.superbiz.agent.harness.retry;
+
+import com.superbiz.agent.harness.core.DiagnosisHarnessCore;
+import com.superbiz.agent.harness.core.BudgetExceededException;
+import com.superbiz.agent.harness.core.RunAbortedException;
+import com.superbiz.agent.harness.core.RunContext;
+import com.superbiz.agent.harness.core.RunState;
+
+import java.util.Objects;
+import java.util.function.Consumer;
+
+public final class HarnessRetryExecutor {
+
+    private final DiagnosisHarnessCore core;
+
+    public HarnessRetryExecutor(DiagnosisHarnessCore core) {
+        this.core = Objects.requireNonNull(core, "core must not be null");
+    }
+
+    public <T> T execute(RunContext context,
+                         RetryPolicy policy,
+                         RetryOperation<T> operation,
+                         RetryFailureClassifier classifier,
+                         Consumer<RetryAttempt> recorder) {
+        Objects.requireNonNull(context, "context must not be null");
+        Objects.requireNonNull(policy, "policy must not be null");
+        Objects.requireNonNull(operation, "operation must not be null");
+        Objects.requireNonNull(classifier, "classifier must not be null");
+        Objects.requireNonNull(recorder, "recorder must not be null");
+
+        for (int attempt = 1; attempt <= policy.maxAttempts(); attempt++) {
+            core.checkActive(context);
+            try {
+                T result = operation.execute();
+                recorder.accept(RetryAttempt.succeeded(attempt));
+                return result;
+            } catch (RunAbortedException exception) {
+                RetryFailure failure = exception.termination().state() == RunState.BUDGET_EXHAUSTED
+                        ? RetryFailure.BUDGET_EXHAUSTED
+                        : RetryFailure.CANCELLED;
+                recorder.accept(RetryAttempt.failed(attempt, failure));
+                throw new RetryExecutionException(attempt, failure, exception);
+            } catch (BudgetExceededException exception) {
+                recorder.accept(RetryAttempt.failed(attempt, RetryFailure.BUDGET_EXHAUSTED));
+                throw new RetryExecutionException(
+                        attempt, RetryFailure.BUDGET_EXHAUSTED, exception);
+            } catch (Exception exception) {
+                RetryFailure failure = classifier.classify(exception);
+                if (failure == null) {
+                    failure = RetryFailure.UNKNOWN;
+                }
+                recorder.accept(RetryAttempt.failed(attempt, failure));
+                if (!policy.allowsRetry(attempt, failure)) {
+                    throw new RetryExecutionException(attempt, failure, exception);
+                }
+            }
+        }
+        throw new IllegalStateException("retry loop exited without a result");
+    }
+}
@@ -0,0 +1,37 @@
+package com.superbiz.agent.harness.retry;
+
+import java.util.Objects;
+import java.util.Set;
+
+public record HarnessRetryPolicies(
+        RetryPolicy intentRouter,
+        RetryPolicy diagnosisAgent,
+        RetryPolicy toolCall,
+        RetryPolicy semanticGuard,
+        RetryPolicy evidenceRepair) {
+
+    public HarnessRetryPolicies {
+        Objects.requireNonNull(intentRouter, "intentRouter must not be null");
+        Objects.requireNonNull(diagnosisAgent, "diagnosisAgent must not be null");
+        Objects.requireNonNull(toolCall, "toolCall must not be null");
+        Objects.requireNonNull(semanticGuard, "semanticGuard must not be null");
+        Objects.requireNonNull(evidenceRepair, "evidenceRepair must not be null");
+    }
+
+    public static HarnessRetryPolicies strict() {
+        RetryPolicy oneAttempt = new RetryPolicy(1, Set.of());
+        return new HarnessRetryPolicies(
+                new RetryPolicy(2, Set.of(
+                        RetryFailure.TIMEOUT,
+                        RetryFailure.TRANSPORT,
+                        RetryFailure.INVALID_OUTPUT)),
+                oneAttempt,
+                oneAttempt,
+                new RetryPolicy(2, Set.of(
+                        RetryFailure.TIMEOUT,
+                        RetryFailure.TRANSPORT,
+                        RetryFailure.PARSE_ERROR,
+                        RetryFailure.SCHEMA_INVALID)),
+                oneAttempt);
+    }
+}
@@ -0,0 +1,24 @@
+package com.superbiz.agent.harness.retry;
+
+public record RetryAttempt(int attemptNumber, boolean success, RetryFailure failure) {
+
+    public RetryAttempt {
+        if (attemptNumber <= 0) {
+            throw new IllegalArgumentException("attemptNumber must be positive");
+        }
+        if (success && failure != null) {
+            throw new IllegalArgumentException("successful attempt must not contain a failure");
+        }
+        if (!success && failure == null) {
+            throw new IllegalArgumentException("failed attempt must contain a failure");
+        }
+    }
+
+    public static RetryAttempt succeeded(int attemptNumber) {
+        return new RetryAttempt(attemptNumber, true, null);
+    }
+
+    public static RetryAttempt failed(int attemptNumber, RetryFailure failure) {
+        return new RetryAttempt(attemptNumber, false, failure);
+    }
+}
@@ -0,0 +1,21 @@
+package com.superbiz.agent.harness.retry;
+
+public final class RetryExecutionException extends RuntimeException {
+
+    private final int attempts;
+    private final RetryFailure failure;
+
+    public RetryExecutionException(int attempts, RetryFailure failure, Exception cause) {
+        super("Operation failed after " + attempts + " attempt(s): " + failure, cause);
+        this.attempts = attempts;
+        this.failure = failure;
+    }
+
+    public int attempts() {
+        return attempts;
+    }
+
+    public RetryFailure failure() {
+        return failure;
+    }
+}
@@ -0,0 +1,14 @@
+package com.superbiz.agent.harness.retry;
+
+public enum RetryFailure {
+    TIMEOUT,
+    TRANSPORT,
+    INVALID_OUTPUT,
+    PARSE_ERROR,
+    SCHEMA_INVALID,
+    NO_EVIDENCE,
+    BUSINESS_REJECTION,
+    CANCELLED,
+    BUDGET_EXHAUSTED,
+    UNKNOWN
+}
@@ -0,0 +1,6 @@
+package com.superbiz.agent.harness.retry;
+
+@FunctionalInterface
+public interface RetryFailureClassifier {
+    RetryFailure classify(Exception exception);
+}
@@ -0,0 +1,6 @@
+package com.superbiz.agent.harness.retry;
+
+@FunctionalInterface
+public interface RetryOperation<T> {
+    T execute() throws Exception;
+}
@@ -0,0 +1,17 @@
+package com.superbiz.agent.harness.retry;
+
+import java.util.Set;
+
+public record RetryPolicy(int maxAttempts, Set<RetryFailure> retryableFailures) {
+
+    public RetryPolicy {
+        if (maxAttempts < 1 || maxAttempts > 2) {
+            throw new IllegalArgumentException("maxAttempts must be 1 or 2");
+        }
+        retryableFailures = retryableFailures == null ? Set.of() : Set.copyOf(retryableFailures);
+    }
+
+    public boolean allowsRetry(int completedAttempts, RetryFailure failure) {
+        return completedAttempts < maxAttempts && retryableFailures.contains(failure);
+    }
+}
@@ -0,0 +1,40 @@
+package com.superbiz.agent.harness.tool.store;
+
+import java.util.regex.Pattern;
+
+public final class ToolCallKeyFactory {
+
+    private static final int MAX_SEGMENT_LENGTH = 128;
+    private static final Pattern SAFE_SEGMENT = Pattern.compile("[A-Za-z0-9][A-Za-z0-9._-]*");
+
+    private final String keyPrefix;
+
+    public ToolCallKeyFactory(String keyPrefix) {
+        if (keyPrefix == null || keyPrefix.isBlank()) {
+            throw new IllegalArgumentException("keyPrefix must not be blank");
+        }
+        if (keyPrefix.startsWith(":") || keyPrefix.endsWith(":") || keyPrefix.contains("::")) {
+            throw new IllegalArgumentException("keyPrefix contains an empty segment");
+        }
+        this.keyPrefix = keyPrefix;
+    }
+
+    public String create(String runId, String toolCallId) {
+        requireSafeSegment(runId, "runId");
+        requireSafeSegment(toolCallId, "toolCallId");
+        return keyPrefix + ":" + runId + ":" + toolCallId;
+    }
+
+    public String keyPrefix() {
+        return keyPrefix;
+    }
+
+    private static void requireSafeSegment(String value, String name) {
+        if (value == null || value.isBlank()) {
+            throw new IllegalArgumentException(name + " must not be blank");
+        }
+        if (value.length() > MAX_SEGMENT_LENGTH || !SAFE_SEGMENT.matcher(value).matches()) {
+            throw new IllegalArgumentException(name + " contains unsafe characters or is too long");
+        }
+    }
+}
@@ -93,6 +93,10 @@ spring:
          min-idle: 0

  ai:
+    # Harness owns allowed retries; provider SDK calls execute once per recorded attempt.
+    retry:
+      max-attempts: 1
+
    vectorstore:
      type: milvus
      milvus:
@@ -0,0 +1,33 @@
+package com.superbiz.agent.config;
+
+import org.junit.jupiter.api.Test;
+import org.yaml.snakeyaml.Yaml;
+
+import java.io.InputStream;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+class SpringAiRetryConfigurationTest {
+
+    @Test
+    void disablesHiddenSpringAiRetries() {
+        try (InputStream input = getClass().getClassLoader().getResourceAsStream("application.yml")) {
+            assertNotNull(input);
+            Map<String, Object> root = new Yaml().load(input);
+            Map<String, Object> spring = map(root.get("spring"));
+            Map<String, Object> ai = map(spring.get("ai"));
+            Map<String, Object> retry = map(ai.get("retry"));
+
+            assertEquals(1, ((Number) retry.get("max-attempts")).intValue());
+        } catch (Exception e) {
+            throw new AssertionError("Failed to read application.yml", e);
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private Map<String, Object> map(Object value) {
+        return (Map<String, Object>) value;
+    }
+}
@@ -0,0 +1,92 @@
+package com.superbiz.agent.harness.core;
+
+import com.superbiz.agent.harness.retry.HarnessRetryPolicies;
+import org.junit.jupiter.api.Test;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.concurrent.atomic.AtomicBoolean;
+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.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class DiagnosisHarnessCoreTest {
+
+    @Test
+    void deadlinePreventsFakeToolInvocation() {
+        MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
+        DiagnosisHarnessCore core = HarnessCoreFixtures.core(clock);
+        RunContext context = core.startRun("session-1", "run-1");
+        AtomicInteger toolCalls = new AtomicInteger();
+        clock.advance(Duration.ofMinutes(2));
+
+        RunAbortedException failure = assertThrows(RunAbortedException.class, () -> {
+            core.beforeToolCall(context, "query_logs");
+            toolCalls.incrementAndGet();
+        });
+
+        assertEquals(0, toolCalls.get());
+        assertEquals(RunState.TIMED_OUT, failure.termination().state());
+        assertEquals(RunCancellationReason.DEADLINE_EXCEEDED, context.cancellation().reason().orElseThrow());
+        assertThrows(RunAbortedException.class, () -> core.beforeModelCall(context));
+    }
+
+    @Test
+    void clientCancellationStopsSubsequentFakeModel() {
+        DiagnosisHarnessCore core = HarnessCoreFixtures.core(new MutableClock(Instant.parse("2026-07-21T10:00:00Z")));
+        RunContext context = core.startRun("session-1", "run-1");
+        AtomicBoolean resourceCancelled = new AtomicBoolean();
+        context.cancellation().onCancel(reason -> resourceCancelled.set(true));
+
+        assertTrue(core.cancel(context, RunCancellationReason.CLIENT_DISCONNECTED));
+
+        assertTrue(resourceCancelled.get());
+        assertThrows(RunAbortedException.class, () -> core.beforeModelCall(context));
+        assertEquals(RunState.CANCELLED, context.lifecycle().state());
+    }
+
+    @Test
+    void toolBudgetExhaustionWinsOverLaterFailure() {
+        MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
+        DiagnosisHarnessCore core = new DiagnosisHarnessCore(
+                clock,
+                () -> "run-1",
+                Duration.ofMinutes(1),
+                new RunBudgetLimits(2, 1, 1, 100, 100, 200, 100),
+                HarnessRetryPolicies.strict());
+        RunContext context = core.startRun("session-1");
+        core.beforeToolCall(context, "query_logs");
+
+        BudgetExceededException failure = assertThrows(
+                BudgetExceededException.class,
+                () -> core.beforeToolCall(context, "query_mysql"));
+        RunTermination termination = context.lifecycle().termination().orElseThrow();
+
+        assertEquals(BudgetKind.TOOL_CALLS, failure.kind());
+        assertEquals(RunState.BUDGET_EXHAUSTED, termination.state());
+        assertFalse(core.completeFailure(context, "late failure"));
+        assertEquals(termination, context.lifecycle().termination().orElseThrow());
+    }
+
+    @Test
+    void actualTokenExhaustionIsRecordedAndStopsNextOperation() {
+        MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
+        DiagnosisHarnessCore core = new DiagnosisHarnessCore(
+                clock,
+                () -> "run-1",
+                Duration.ofMinutes(1),
+                new RunBudgetLimits(2, 2, 2, 10, 10, 15, 100),
+                HarnessRetryPolicies.strict());
+        RunContext context = core.startRun("session-1");
+        core.beforeModelCall(context);
+
+        assertThrows(BudgetExceededException.class, () -> core.recordTokens(context, 9, 8));
+
+        assertEquals(17, context.budget().snapshot().totalTokens());
+        assertEquals(RunState.BUDGET_EXHAUSTED, context.lifecycle().state());
+        assertThrows(RunAbortedException.class, () -> core.beforeToolCall(context, "query_logs"));
+    }
+}
@@ -0,0 +1,24 @@
+package com.superbiz.agent.harness.core;
+
+import com.superbiz.agent.harness.retry.HarnessRetryPolicies;
+
+import java.time.Duration;
+
+public final class HarnessCoreFixtures {
+
+    private HarnessCoreFixtures() {
+    }
+
+    public static RunBudgetLimits generousLimits() {
+        return new RunBudgetLimits(10, 10, 5, 10_000, 10_000, 20_000, 1_000_000);
+    }
+
+    public static DiagnosisHarnessCore core(MutableClock clock) {
+        return new DiagnosisHarnessCore(
+                clock,
+                () -> "run-generated-1",
+                Duration.ofMinutes(2),
+                generousLimits(),
+                HarnessRetryPolicies.strict());
+    }
+}
@@ -0,0 +1,42 @@
+package com.superbiz.agent.harness.core;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicReference;
+
+public final class MutableClock extends Clock {
+
+    private final AtomicReference<Instant> instant;
+    private final ZoneId zone;
+
+    public MutableClock(Instant instant) {
+        this(instant, ZoneId.of("UTC"));
+    }
+
+    private MutableClock(Instant instant, ZoneId zone) {
+        this.instant = new AtomicReference<>(Objects.requireNonNull(instant));
+        this.zone = Objects.requireNonNull(zone);
+    }
+
+    public void advance(Duration duration) {
+        instant.updateAndGet(current -> current.plus(duration));
+    }
+
+    @Override
+    public ZoneId getZone() {
+        return zone;
+    }
+
+    @Override
+    public Clock withZone(ZoneId zone) {
+        return new MutableClock(instant(), zone);
+    }
+
+    @Override
+    public Instant instant() {
+        return instant.get();
+    }
+}
@@ -0,0 +1,74 @@
+package com.superbiz.agent.harness.core;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class RunBudgetTest {
+
+    @Test
+    void rejectsInvalidLimits() {
+        assertThrows(IllegalArgumentException.class,
+                () -> new RunBudgetLimits(0, 1, 1, 1, 1, 1, 1));
+    }
+
+    @Test
+    void perToolExhaustionDoesNotPartiallyIncrementCounters() {
+        RunBudget budget = new RunBudget(new RunBudgetLimits(2, 3, 1, 100, 100, 200, 100));
+        budget.reserveToolCall("query_logs");
+
+        BudgetExceededException failure = assertThrows(
+                BudgetExceededException.class,
+                () -> budget.reserveToolCall("query_logs"));
+
+        RunBudgetUsage usage = budget.snapshot();
+        assertEquals(BudgetKind.TOOL_CALLS_PER_TOOL, failure.kind());
+        assertEquals(1, usage.toolCalls());
+        assertEquals(1, usage.toolCallsByName().get("query_logs"));
+    }
+
+    @Test
+    void tokenExhaustionKeepsActualUsage() {
+        RunBudget budget = new RunBudget(new RunBudgetLimits(2, 2, 2, 10, 10, 15, 100));
+
+        BudgetExceededException failure = assertThrows(
+                BudgetExceededException.class,
+                () -> budget.recordTokens(9, 8));
+
+        assertEquals(BudgetKind.TOTAL_TOKENS, failure.kind());
+        assertEquals(9, budget.snapshot().inputTokens());
+        assertEquals(8, budget.snapshot().outputTokens());
+        assertEquals(17, budget.snapshot().totalTokens());
+    }
+
+    @Test
+    void concurrentCapacityReservationsAreAtomic() throws Exception {
+        RunCapacityCounter counter = new RunCapacityCounter(100);
+        ExecutorService executor = Executors.newFixedThreadPool(5);
+        try {
+            List<Future<Long>> reservations = new ArrayList<>();
+            for (int i = 0; i < 10; i++) {
+                reservations.add(executor.submit(() -> counter.reserve(10)));
+            }
+            for (Future<Long> reservation : reservations) {
+                reservation.get();
+            }
+        } finally {
+            executor.shutdownNow();
+        }
+
+        assertEquals(100, counter.usedBytes());
+        BudgetExceededException failure = assertThrows(
+                BudgetExceededException.class,
+                () -> counter.reserve(1));
+        assertEquals(BudgetKind.RUN_BYTES, failure.kind());
+        assertEquals(100, counter.usedBytes());
+    }
+}
@@ -0,0 +1,79 @@
+package com.superbiz.agent.harness.core;
+
+import org.junit.jupiter.api.Test;
+
+import java.time.Instant;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class RunContextTest {
+
+    @Test
+    void rejectsBlankIdentityBeforePublishingContext() {
+        DiagnosisHarnessCore core = HarnessCoreFixtures.core(new MutableClock(Instant.parse("2026-07-21T10:00:00Z")));
+
+        assertThrows(IllegalArgumentException.class, () -> core.startRun(" ", "run-1"));
+        assertThrows(IllegalArgumentException.class, () -> core.startRun("session-1", " "));
+    }
+
+    @Test
+    void explicitlyPropagatesSameHandlesAcrossAsyncBoundary() throws Exception {
+        DiagnosisHarnessCore core = HarnessCoreFixtures.core(new MutableClock(Instant.parse("2026-07-21T10:00:00Z")));
+        RunContext context = core.startRun("session-1", "run-1");
+        AtomicReference<Thread> worker = new AtomicReference<>();
+
+        RunContext observed = CompletableFuture.supplyAsync(() -> {
+            worker.set(Thread.currentThread());
+            return fakeTool(context);
+        }).get(5, TimeUnit.SECONDS);
+
+        assertEquals("session-1", observed.sessionId());
+        assertEquals("run-1", observed.runId());
+        assertSame(context.cancellation(), observed.cancellation());
+        assertSame(context.budget(), observed.budget());
+        assertSame(context.lifecycle(), observed.lifecycle());
+        assertFalse(worker.get().equals(Thread.currentThread()));
+    }
+
+    @Test
+    void cancellationUsesFirstReasonAndNotifiesRegisteredCallbacks() {
+        DiagnosisHarnessCore core = HarnessCoreFixtures.core(new MutableClock(Instant.parse("2026-07-21T10:00:00Z")));
+        RunContext context = core.startRun("session-1", "run-1");
+        AtomicInteger callbacks = new AtomicInteger();
+        context.cancellation().onCancel(reason -> callbacks.incrementAndGet());
+        context.cancellation().onCancel(reason -> callbacks.incrementAndGet());
+
+        assertTrue(core.cancel(context, RunCancellationReason.CLIENT_DISCONNECTED));
+        assertFalse(context.cancellation().cancel(RunCancellationReason.USER_REQUESTED));
+
+        assertEquals(RunCancellationReason.CLIENT_DISCONNECTED, context.cancellation().reason().orElseThrow());
+        assertEquals(2, callbacks.get());
+        assertEquals(RunState.CANCELLED, context.lifecycle().state());
+    }
+
+    @Test
+    void firstTerminalOutcomeCannotBeOverwritten() {
+        DiagnosisHarnessCore core = HarnessCoreFixtures.core(new MutableClock(Instant.parse("2026-07-21T10:00:00Z")));
+        RunContext context = core.startRun("session-1", "run-1");
+
+        assertTrue(core.completeSuccess(context));
+        RunTermination success = context.lifecycle().termination().orElseThrow();
+        assertFalse(core.completeFailure(context, "late failure"));
+        assertFalse(core.cancel(context, RunCancellationReason.CLIENT_DISCONNECTED));
+
+        assertSame(success, context.lifecycle().termination().orElseThrow());
+        assertEquals(RunState.SUCCESS, context.lifecycle().state());
+    }
+
+    private RunContext fakeTool(RunContext context) {
+        return context;
+    }
+}
@@ -0,0 +1,134 @@
+package com.superbiz.agent.harness.retry;
+
+import com.superbiz.agent.harness.core.DiagnosisHarnessCore;
+import com.superbiz.agent.harness.core.BudgetExceededException;
+import com.superbiz.agent.harness.core.BudgetKind;
+import com.superbiz.agent.harness.core.HarnessCoreFixtures;
+import com.superbiz.agent.harness.core.MutableClock;
+import com.superbiz.agent.harness.core.RunAbortedException;
+import com.superbiz.agent.harness.core.RunCancellationReason;
+import com.superbiz.agent.harness.core.RunContext;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class HarnessRetryExecutorTest {
+
+    private DiagnosisHarnessCore core;
+    private RunContext context;
+    private HarnessRetryExecutor executor;
+
+    @BeforeEach
+    void setUp() {
+        core = HarnessCoreFixtures.core(new MutableClock(Instant.parse("2026-07-21T10:00:00Z")));
+        context = core.startRun("session-1", "run-1");
+        executor = new HarnessRetryExecutor(core);
+    }
+
+    @Test
+    void retriesRouterTransportFailureOnceAndRecordsBothAttempts() {
+        AtomicInteger calls = new AtomicInteger();
+        List<RetryAttempt> attempts = new ArrayList<>();
+
+        String result = executor.execute(
+                context,
+                context.retryPolicies().intentRouter(),
+                () -> {
+                    if (calls.incrementAndGet() == 1) {
+                        throw new IOException("temporary transport failure");
+                    }
+                    return "DIAGNOSIS";
+                },
+                exception -> RetryFailure.TRANSPORT,
+                attempts::add);
+
+        assertEquals("DIAGNOSIS", result);
+        assertEquals(2, calls.get());
+        assertEquals(List.of(
+                RetryAttempt.failed(1, RetryFailure.TRANSPORT),
+                RetryAttempt.succeeded(2)), attempts);
+    }
+
+    @Test
+    void toolPolicyRunsOnlyOneAttempt() {
+        AtomicInteger calls = new AtomicInteger();
+        List<RetryAttempt> attempts = new ArrayList<>();
+
+        RetryExecutionException failure = assertThrows(RetryExecutionException.class, () -> executor.execute(
+                context,
+                context.retryPolicies().toolCall(),
+                () -> {
+                    calls.incrementAndGet();
+                    throw new IOException("tool failed");
+                },
+                exception -> RetryFailure.TRANSPORT,
+                attempts::add));
+
+        assertEquals(1, calls.get());
+        assertEquals(1, failure.attempts());
+        assertEquals(List.of(RetryAttempt.failed(1, RetryFailure.TRANSPORT)), attempts);
+    }
+
+    @Test
+    void cancellationBetweenAttemptsPreventsSecondInvocation() {
+        AtomicInteger calls = new AtomicInteger();
+
+        assertThrows(RunAbortedException.class, () -> executor.execute(
+                context,
+                context.retryPolicies().semanticGuard(),
+                () -> {
+                    calls.incrementAndGet();
+                    throw new IOException("temporary transport failure");
+                },
+                exception -> RetryFailure.TRANSPORT,
+                attempt -> core.cancel(context, RunCancellationReason.CLIENT_DISCONNECTED)));
+
+        assertEquals(1, calls.get());
+    }
+
+    @Test
+    void noEvidenceAndBusinessRejectionAreNotRetryable() {
+        AtomicInteger calls = new AtomicInteger();
+
+        RetryExecutionException failure = assertThrows(RetryExecutionException.class, () -> executor.execute(
+                context,
+                context.retryPolicies().intentRouter(),
+                () -> {
+                    calls.incrementAndGet();
+                    throw new IllegalStateException("valid empty result");
+                },
+                exception -> RetryFailure.NO_EVIDENCE,
+                attempt -> { }));
+
+        assertEquals(1, calls.get());
+        assertEquals(RetryFailure.NO_EVIDENCE, failure.failure());
+        assertEquals(1, context.retryPolicies().diagnosisAgent().maxAttempts());
+        assertEquals(1, context.retryPolicies().evidenceRepair().maxAttempts());
+    }
+
+    @Test
+    void budgetFailureCannotBeMisclassifiedAsRetryable() {
+        AtomicInteger calls = new AtomicInteger();
+
+        RetryExecutionException failure = assertThrows(RetryExecutionException.class, () -> executor.execute(
+                context,
+                context.retryPolicies().intentRouter(),
+                () -> {
+                    calls.incrementAndGet();
+                    throw new BudgetExceededException(BudgetKind.MODEL_CALLS, 1, 2);
+                },
+                exception -> RetryFailure.TRANSPORT,
+                attempt -> { }));
+
+        assertEquals(1, calls.get());
+        assertEquals(RetryFailure.BUDGET_EXHAUSTED, failure.failure());
+    }
+}
@@ -0,0 +1,26 @@
+package com.superbiz.agent.harness.tool.store;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class ToolCallKeyFactoryTest {
+
+    private final ToolCallKeyFactory factory = new ToolCallKeyFactory("superbiz:harness:tool-call");
+
+    @Test
+    void preservesExactFrameworkToolCallId() {
+        assertEquals(
+                "superbiz:harness:tool-call:run-123:call_ABC-123",
+                factory.create("run-123", "call_ABC-123"));
+    }
+
+    @Test
+    void rejectsBlankUnsafeAndOversizedSegmentsWithoutFallbackId() {
+        assertThrows(IllegalArgumentException.class, () -> factory.create("run-1", ""));
+        assertThrows(IllegalArgumentException.class, () -> factory.create("run:1", "call-1"));
+        assertThrows(IllegalArgumentException.class, () -> factory.create("run-1", "call/1"));
+        assertThrows(IllegalArgumentException.class, () -> factory.create("run-1", "a".repeat(129)));
+    }
+}