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");
+        }
+    }
+}