Agent Evaluation & Reliability

package com.interviewpreparator.agents;
import java.time.*;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.*;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.DriverManagerDataSource;
import org.testcontainers.postgresql.PostgreSQLContainer;
import org.testcontainers.utility.DockerImageName;
import static com.interviewpreparator.agents.AgentContracts.*;
import static org.junit.jupiter.api.Assertions.*;
class AgentCheckpointIT {
    static final PostgreSQLContainer DATABASE = new PostgreSQLContainer(DockerImageName.parse("pgvector/pgvector:0.8.0-pg17").asCompatibleSubstituteFor("postgres"));
    static final Clock CLOCK = Clock.fixed(Instant.parse("2026-10-30T00:00:00Z"), ZoneOffset.UTC);
    static JdbcTemplate jdbc;
    @BeforeAll static void start() {
        DATABASE.start();
        jdbc = new JdbcTemplate(new DriverManagerDataSource(DATABASE.getJdbcUrl(), DATABASE.getUsername(), DATABASE.getPassword()));
        new JdbcRunStore(jdbc).initializeLocalSchema();
        new JdbcPlanTools(jdbc).initializeLocalSchema();
    }
    @AfterAll static void stop() { DATABASE.stop(); }
    static Scope scope() { return new Scope("tenant-a", "alpha", "conversation-1", UUID.randomUUID().toString()); }
    static State initial() { return new State(Status.RUNNING, 0, 0, 0, 0, CLOCK.millis() + 60_000, null, List.of(), Map.of(), ""); }
    static BoundedAgent.Budget budget() { return new BoundedAgent.Budget(4, 4, 100, 1000, 20, 100, 1, Duration.ofSeconds(60), Duration.ofSeconds(5)); }
    static Model model() { return model("PLAN-100"); }
    static Model model(String plan) { return state -> new ModelReply(state.observations().isEmpty()
            ? new ToolCall("write-1", "completePlan", plan) : new Finish("Completed."), 10, 50); }
    @Test void restoresPendingApprovalAndStateAfterRepositoryRecreation() {
        Scope scope = scope();
        var tools = new JdbcPlanTools(jdbc);
        tools.seedPlan("tenant-a", "PLAN-100");
        try (var first = new BoundedAgent(new JdbcRunStore(jdbc), CLOCK)) {
            assertEquals(Status.WAITING_APPROVAL, first.run(scope, budget(), model(), tools, Set.of()).state().status());
        }
        try (var resumed = new BoundedAgent(new JdbcRunStore(jdbc), CLOCK)) {
            var result = resumed.run(scope, budget(), model(), new JdbcPlanTools(jdbc), Set.of("PLAN-100"));
            assertEquals(Status.COMPLETED, result.state().status());
            assertEquals(2, result.state().steps());
            assertEquals(1, result.state().toolAttempts());
        }
    }
    @Test void usesOptimisticConcurrencyAndSeparatesEveryScopeDimension() {
        Scope scope = scope();
        var first = new JdbcRunStore(jdbc);
        var snapshot = first.createOrLoad(scope, initial());
        State privateState = new State(Status.COMPLETED, 1, 0, 10, 50, initial().expiresAtMillis(), null, List.of("private-note"), Map.of(), "private-result");
        first.save(scope, snapshot.version(), privateState);
        assertThrows(Conflict.class, () -> new JdbcRunStore(jdbc).save(scope, snapshot.version(), initial()));
        for (Scope other : List.of(new Scope("tenant-b", scope.user(), scope.conversation(), scope.runId()),
                new Scope(scope.tenant(), "beta", scope.conversation(), scope.runId()),
                new Scope(scope.tenant(), scope.user(), "conversation-2", scope.runId()))) {
            var isolated = first.createOrLoad(other, initial());
            assertEquals(0, isolated.version());
            assertTrue(isolated.state().observations().isEmpty());
        }
    }
    @Test void survivesTheGapBetweenToolCommitAndCheckpointSave() {
        Scope scope = scope();
        var durable = new JdbcRunStore(jdbc);
        var tools = new JdbcPlanTools(jdbc);
        tools.seedPlan("tenant-a", "PLAN-102");
        int before = tools.mutations("tenant-a", "PLAN-102");
        RunStore broken = new RunStore() {
            public Snapshot createOrLoad(Scope s, State initial) { return durable.createOrLoad(s, initial); }
            public Snapshot save(Scope s, long expected, State next) {
                if (!next.receipts().isEmpty()) throw new IllegalStateException("Simulated worker crash after tool commit");
                return durable.save(s, expected, next);
            }
        };
        try (var first = new BoundedAgent(broken, CLOCK)) {
            assertThrows(IllegalStateException.class, () -> first.run(scope, budget(), model("PLAN-102"), tools, Set.of("PLAN-102")));
        }
        try (var restarted = new BoundedAgent(new JdbcRunStore(jdbc), CLOCK)) {
            var result = restarted.run(scope, budget(), model("PLAN-102"), new JdbcPlanTools(jdbc), Set.of("PLAN-102"));
            assertEquals(Status.COMPLETED, result.state().status());
            assertEquals(2, result.state().toolAttempts());
            assertEquals(1, tools.mutations("tenant-a", "PLAN-102") - before);
        }
    }
    @Test void concurrentDeliveryOfTheSameWriteHasOneReceiptAndEffect() throws Exception {
        Scope scope = scope();
        var tools = new JdbcPlanTools(jdbc);
        tools.seedPlan("tenant-a", "PLAN-101");
        var call = new ToolCall("same-call", "completePlan", "PLAN-101");
        try (var executor = Executors.newVirtualThreadPerTaskExecutor()) {
            var tasks = List.of(executor.submit(() -> tools.execute(scope, call, Set.of("PLAN-101"))),
                    executor.submit(() -> new JdbcPlanTools(jdbc).execute(scope, call, Set.of("PLAN-101"))));
            for (var task : tasks) assertEquals("COMPLETED", task.get(5, TimeUnit.SECONDS));
        }
        assertEquals(1, tools.mutations("tenant-a", "PLAN-101"));
        assertThrows(SecurityException.class, () -> tools.execute(scope, new ToolCall("same-call", "readPlan", "PLAN-101"), Set.of()));
        assertThrows(SecurityException.class, () -> tools.execute(new Scope("tenant-b", "beta", scope.conversation(), scope.runId()), call, Set.of("PLAN-101")));
    }
    @Test void persistedDeadlineCannotBeExtendedByRestartingAWorker() {
        Scope scope = scope();
        var store = new JdbcRunStore(jdbc);
        store.createOrLoad(scope, new State(Status.RUNNING, 0, 0, 0, 0, CLOCK.millis() - 1, null, List.of(), Map.of(), ""));
        try (var agent = new BoundedAgent(new JdbcRunStore(jdbc), CLOCK)) {
            var result = agent.run(scope, budget(), state -> { throw new AssertionError("Expired model call"); }, new JdbcPlanTools(jdbc), Set.of());
            assertEquals(Status.TIMED_OUT, result.state().status());
        }
    }
}