script-build-and-framework-tmp/agent/tests/test_orchestration_v2_control.py

import asyncio
from dataclasses import replace
from types import SimpleNamespace

import pytest

from agent import AgentExecutor
from agent.orchestration.config import OrchestrationConfig
from agent.orchestration.coordinator import TaskCoordinator
from agent.orchestration.executor import LocalAgentExecutor
from agent.orchestration.models import (
    AttemptStatus,
    AttemptSubmission,
    ExecutionStats,
    FailureCode,
    TaskAttempt,
    TaskStatus,
    ValidationVerdict,
    utc_now,
)
from agent.orchestration.protocols import AgentExecutor as ProtocolAgentExecutor
from agent.orchestration.protocols import WorkerRunResult
from agent.orchestration.store import (
    FileSystemArtifactStore,
    FileSystemTaskStore,
    TraceEventSink,
)

from test_coordinator_integration import FakeExecutor, create_task, make_coordinator


def test_v2_control_config_defaults_preserve_unbounded_v1_execution():
    config = OrchestrationConfig()

    assert config.worker_timeout_seconds is None
    assert config.validator_timeout_seconds is None
    assert config.stop_grace_seconds == 5.0


def test_agent_executor_port_is_available_from_public_package():
    assert AgentExecutor is ProtocolAgentExecutor


@pytest.mark.parametrize(
    ("field", "value"),
    [
        ("worker_timeout_seconds", 0),
        ("worker_timeout_seconds", -0.1),
        ("validator_timeout_seconds", 0),
        ("validator_timeout_seconds", -1),
        ("validator_timeout_seconds", float("nan")),
        ("worker_timeout_seconds", float("inf")),
        ("worker_timeout_seconds", True),
        ("stop_grace_seconds", -0.1),
        ("stop_grace_seconds", float("nan")),
        ("stop_grace_seconds", float("inf")),
        ("stop_grace_seconds", False),
    ],
)
def test_v2_control_config_rejects_invalid_durations(field, value):
    with pytest.raises(ValueError, match=field):
        OrchestrationConfig(**{field: value})


@pytest.mark.parametrize("field", ["max_parallel_tasks", "max_repair_continuations"])
def test_v2_control_config_rejects_boolean_integer_limits(field):
    with pytest.raises(ValueError, match=field):
        OrchestrationConfig(**{field: True})


def test_v2_control_config_allows_immediate_stop_and_explicit_timeouts():
    config = OrchestrationConfig(
        worker_timeout_seconds=1.5,
        validator_timeout_seconds=2,
        stop_grace_seconds=0,
    )

    assert config.worker_timeout_seconds == 1.5
    assert config.validator_timeout_seconds == 2
    assert config.stop_grace_seconds == 0


@pytest.mark.asyncio
async def test_local_executor_stop_delegates_every_request():
    class Runner:
        def __init__(self):
            self.calls = []

        async def stop(self, trace_id):
            self.calls.append(trace_id)
            return len(self.calls) == 1

    runner = Runner()
    executor = LocalAgentExecutor(runner)

    assert await executor.stop("worker-trace") is True
    assert await executor.stop("worker-trace") is False
    assert runner.calls == ["worker-trace", "worker-trace"]


@pytest.mark.asyncio
async def test_local_executor_stop_does_not_hide_runner_failure():
    class Runner:
        async def stop(self, trace_id):
            raise RuntimeError(f"cannot stop {trace_id}")

    with pytest.raises(RuntimeError, match="cannot stop validator-trace"):
        await LocalAgentExecutor(Runner()).stop("validator-trace")


@pytest.mark.asyncio
async def test_hard_cancellation_is_not_normalized_as_agent_failure():
    class Runner:
        async def run_result(self, **kwargs):
            raise asyncio.CancelledError

    with pytest.raises(asyncio.CancelledError):
        await LocalAgentExecutor(Runner()).run_worker({
            "worker_preset": "worker",
            "worker_trace_id": "worker-trace",
            "root_trace_id": "root",
            "task_id": "task",
            "spec_version": 1,
            "attempt_id": "attempt",
            "task_spec": {},
            "continue_trace_id": None,
        })


@pytest.mark.asyncio
async def test_failed_stats_lookup_does_not_mask_original_runner_error():
    class Store:
        async def get_trace(self, trace_id):
            raise RuntimeError("stats store unavailable")

    class Runner:
        trace_store = Store()

        async def run_result(self, **kwargs):
            raise RuntimeError("original runner failure")

    result = await LocalAgentExecutor(Runner()).run_worker({
        "worker_preset": "worker",
        "worker_trace_id": "worker-trace",
        "root_trace_id": "root",
        "task_id": "task",
        "spec_version": 1,
        "attempt_id": "attempt",
        "task_spec": {},
        "continue_trace_id": None,
    })

    assert result.error == "original runner failure"
    assert result.execution_stats.failure_code == FailureCode.EXECUTOR_ERROR


@pytest.mark.asyncio
async def test_worker_protects_operation_epoch_and_deadline():
    class Runner:
        def __init__(self):
            self.config = None

        async def run_result(self, *, messages, config):
            self.config = config
            return {"status": "completed", "summary": "done"}

    runner = Runner()
    executor = LocalAgentExecutor(runner)
    result = await executor.run_worker({
        "worker_preset": "worker",
        "worker_trace_id": "worker-trace",
        "root_trace_id": "root",
        "task_id": "task",
        "spec_version": 2,
        "attempt_id": "attempt",
        "task_spec": {"objective": "test"},
        "continue_trace_id": None,
        "operation_id": "operation",
        "execution_epoch": 7,
        "deadline": "2030-01-01T00:00:00+00:00",
        "untrusted": "must-not-leak",
    })

    assert result.trace_id == "worker-trace"
    assert result.status == "completed"
    assert runner.config.trace_id is None
    assert runner.config.new_trace_id == "worker-trace"
    assert runner.config.context == {
        "root_trace_id": "root",
        "task_id": "task",
        "spec_version": 2,
        "attempt_id": "attempt",
        "completion_policy": "explicit_validation",
        "operation_id": "operation",
        "execution_epoch": 7,
        "deadline": "2030-01-01T00:00:00+00:00",
    }


@pytest.mark.asyncio
async def test_validator_omits_absent_v2_context_and_normalizes_empty_result():
    class Runner:
        async def run_result(self, *, messages, config):
            self.config = config
            return {}

    runner = Runner()
    result = await LocalAgentExecutor(runner).run_validator({
        "validator_preset": "validator",
        "validator_trace_id": "validator-trace",
        "root_trace_id": "root",
        "task_id": "task",
        "spec_version": 1,
        "attempt_id": "attempt",
        "snapshot_id": "snapshot",
        "validation_id": "validation",
        "task_spec": {},
        "artifact_snapshot": {},
    })

    assert result.trace_id == "validator-trace"
    assert result.status == "failed"
    assert result.summary == ""
    assert result.error is None
    assert runner.config.trace_id is None
    assert runner.config.new_trace_id == "validator-trace"
    assert "operation_id" not in runner.config.context
    assert "execution_epoch" not in runner.config.context
    assert "deadline" not in runner.config.context


@pytest.mark.parametrize(
    "kwargs",
    [
        {"primary_model": " "},
        {"primary_model": 123},
        {"total_tokens": -1},
        {"total_tokens": True},
        {"total_cost": -0.1},
        {"total_cost": float("nan")},
        {"total_cost": float("inf")},
        {"total_cost": True},
        {"failure_code": "unknown"},
        {"failure_code": 123},
    ],
)
def test_execution_stats_reject_invalid_values(kwargs):
    with pytest.raises(ValueError, match="ExecutionStats"):
        ExecutionStats(**kwargs)


def test_execution_stats_normalizes_known_failure_code_string():
    assert ExecutionStats(failure_code="timeout").failure_code == FailureCode.TIMEOUT


def test_old_stage_without_stats_roundtrips_and_duration_is_derived():
    attempt = TaskAttempt.from_dict({
        "attempt_id": "attempt",
        "task_id": "task",
        "spec_version": 1,
        "worker_trace_id": "trace",
        "worker_preset": "worker",
        "execution_mode": "new",
        "started_at": "2026-07-18T00:00:00Z",
        "completed_at": "2026-07-18T00:00:01.250+00:00",
    })

    assert attempt.execution_stats is None
    assert attempt.duration_ms == 1250
    assert replace(attempt, completed_at="2026-07-17T23:59:59Z").duration_ms == 0
    assert replace(attempt, completed_at=None).duration_ms is None


@pytest.mark.asyncio
async def test_repair_continuation_records_only_trace_usage_delta():
    trace = SimpleNamespace(
        context={}, total_tokens=100, total_cost=1.0, model="repair-model"
    )

    class Coordinator:
        async def validate_continue_from(self, *args):
            return "trace"

    class Store:
        async def get_trace(self, trace_id):
            return trace

        async def update_trace(self, trace_id, **updates):
            trace.context = updates["context"]

    class Runner:
        task_coordinator = Coordinator()
        trace_store = Store()

        async def run_result(self, **kwargs):
            trace.total_tokens = 130
            trace.total_cost = 1.4
            return {
                "trace_id": "trace",
                "status": "completed",
                "stats": {"total_tokens": 130, "total_cost": 1.4},
            }

    result = await LocalAgentExecutor(Runner()).run_worker({
        "worker_preset": "worker",
        "worker_trace_id": "trace",
        "root_trace_id": "root",
        "task_id": "task",
        "spec_version": 1,
        "attempt_id": "new-attempt",
        "prior_attempt_id": "old-attempt",
        "task_spec": {},
        "continue_trace_id": "trace",
    })

    assert result.execution_stats.primary_model == "repair-model"
    assert result.execution_stats.total_tokens == 30
    assert result.execution_stats.total_cost == pytest.approx(0.4)


class StatsExecutor(FakeExecutor):
    async def run_worker(self, context):
        result = await super().run_worker(context)
        return replace(
            result,
            execution_stats=ExecutionStats(
                primary_model="worker-model", total_tokens=12, total_cost=0.12
            ),
        )

    async def run_validator(self, context):
        result = await super().run_validator(context)
        return replace(
            result,
            execution_stats=ExecutionStats(
                primary_model="validator-model", total_tokens=7, total_cost=0.07
            ),
        )


@pytest.mark.asyncio
async def test_successful_stage_stats_persist_once_and_failed_verdict_is_not_run_failure(
    tmp_path,
):
    executor = StatsExecutor([ValidationVerdict.FAILED])
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    task_id = await create_task(coordinator, "persist execution stats")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    ledger = await store.load("root")
    attempt = ledger.attempts[result.attempt_id]
    validation = ledger.validations[result.validation_id]

    assert attempt.execution_stats.primary_model == "worker-model"
    assert attempt.execution_stats.total_tokens == 12
    assert validation.verdict == ValidationVerdict.FAILED
    assert validation.execution_stats.primary_model == "validator-model"
    assert validation.execution_stats.failure_code is None
    assert attempt.duration_ms is not None
    assert validation.duration_ms is not None


@pytest.mark.asyncio
async def test_completed_worker_without_submit_is_protocol_failure(tmp_path):
    executor = StatsExecutor([], submit_worker=False)
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    task_id = await create_task(coordinator, "protocol failure")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    attempt = (await store.load("root")).attempts[result.attempt_id]

    assert attempt.execution_stats.failure_code == FailureCode.PROTOCOL_VIOLATION
    assert attempt.execution_stats.total_tokens == 12


@pytest.mark.asyncio
async def test_protocol_failure_overrides_executor_supplied_failure_code(tmp_path):
    class MisclassifiedExecutor(StatsExecutor):
        async def run_worker(self, context):
            result = await super().run_worker(context)
            return replace(
                result,
                execution_stats=replace(
                    result.execution_stats,
                    failure_code=FailureCode.AGENT_FAILED,
                ),
            )

    executor = MisclassifiedExecutor([], submit_worker=False)
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    task_id = await create_task(coordinator, "protocol code wins")
    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    attempt = (await store.load("root")).attempts[result.attempt_id]

    assert attempt.execution_stats.failure_code == FailureCode.PROTOCOL_VIOLATION


@pytest.mark.asyncio
async def test_late_failure_cannot_overwrite_concurrent_submission(tmp_path):
    executor = FakeExecutor([])
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    task_id = await create_task(coordinator, "concurrent submission")
    operation = await coordinator.start_operation(
        "root", "dispatch", task_ids=[task_id]
    )
    completed = await coordinator.await_operation("root", operation.operation_id)
    ledger = await store.load("root")
    attempt_id = completed.attempt_ids[0]
    original = ledger.attempts[attempt_id]

    await coordinator._mark_worker_failure(
        "root",
        task_id,
        attempt_id,
        "failed",
        "late failure",
        ExecutionStats(failure_code=FailureCode.AGENT_FAILED),
        operation.operation_id,
        completed.execution_epoch,
    )
    after = (await store.load("root")).attempts[attempt_id]

    assert after.status == original.status
    assert after.error == original.error
    assert after.execution_stats == original.execution_stats


@pytest.mark.asyncio
async def test_timeout_persists_machine_failure_code(tmp_path):
    class SlowExecutor(FakeExecutor):
        async def run_worker(self, context):
            await asyncio.sleep(1)
            return WorkerRunResult(context["worker_trace_id"], "completed")

    executor = SlowExecutor([])
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    coordinator.config = OrchestrationConfig(worker_timeout_seconds=0.001)
    task_id = await create_task(coordinator, "timeout stats")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    attempt = (await store.load("root")).attempts[result.attempt_id]

    assert attempt.status == AttemptStatus.EXPIRED
    assert attempt.execution_stats.failure_code == FailureCode.TIMEOUT


@pytest.mark.asyncio
async def test_remote_stop_wins_over_late_worker_failure(tmp_path):
    class LateFailureExecutor:
        def __init__(self):
            self.started = asyncio.Event()
            self.release = asyncio.Event()

        async def run_worker(self, context):
            self.started.set()
            await self.release.wait()
            return WorkerRunResult(
                context["worker_trace_id"],
                "failed",
                error="late failure",
                execution_stats=ExecutionStats(
                    total_tokens=3, failure_code=FailureCode.AGENT_FAILED
                ),
            )

        async def run_validator(self, context):
            raise AssertionError("validator must not run")

        async def stop(self, trace_id):
            return True

    executor = LateFailureExecutor()
    coordinator_a, store, trace_store = await make_coordinator(tmp_path, executor)
    task_id = await create_task(coordinator_a, "remote stop stats")
    operation = await coordinator_a.start_operation("root", "dispatch", task_ids=[task_id])
    await asyncio.wait_for(executor.started.wait(), timeout=1)

    observer_executor = FakeExecutor([])
    coordinator_b = TaskCoordinator(
        store,
        FileSystemArtifactStore(str(tmp_path)),
        trace_store,
        OrchestrationConfig(stop_grace_seconds=0),
        TraceEventSink(str(tmp_path)),
        observer_executor,
    )
    observer_executor.coordinator = coordinator_b
    await coordinator_b.stop_operation("root", operation.operation_id)
    executor.release.set()
    await coordinator_a.runtime.wait("root", operation.operation_id)

    ledger = await store.load("root")
    attempt = ledger.attempts[ledger.operations[operation.operation_id].attempt_ids[0]]
    assert attempt.status == AttemptStatus.STOPPED
    assert attempt.execution_stats.failure_code == FailureCode.STOPPED
    assert attempt.error == "Operation stopped"
    assert ledger.tasks[task_id].status == TaskStatus.NEEDS_REPLAN


@pytest.mark.asyncio
async def test_cas_retry_does_not_overwrite_concurrent_attempt_submission(tmp_path):
    class ConcurrentSubmissionStore:
        def __init__(self, inner):
            self.inner = inner
            self.base_path = inner.base_path
            self.injected = False

        async def load(self, root_trace_id):
            return await self.inner.load(root_trace_id)

        async def commit(
            self,
            ledger,
            expected_revision,
            idempotency_key=None,
            event=None,
        ):
            if (
                event is not None
                and event.event_type == "worker_failed"
                and not self.injected
            ):
                self.injected = True
                winner = await self.inner.load(ledger.root_trace_id)
                attempt = next(
                    item
                    for item in winner.attempts.values()
                    if item.status == AttemptStatus.RUNNING
                )
                attempt.status = AttemptStatus.SUBMITTED
                attempt.submission = AttemptSubmission(summary="concurrent winner")
                attempt.completed_at = utc_now()
                winner.tasks[attempt.task_id].status = TaskStatus.AWAITING_VALIDATION
                await self.inner.commit(winner, winner.revision)
            return await self.inner.commit(
                ledger,
                expected_revision,
                idempotency_key=idempotency_key,
                event=event,
            )

        async def list_events(self, *args, **kwargs):
            return await self.inner.list_events(*args, **kwargs)

        async def list_recoverable(self, *args, **kwargs):
            return await self.inner.list_recoverable(*args, **kwargs)

    task_store = ConcurrentSubmissionStore(FileSystemTaskStore(str(tmp_path)))
    executor = StatsExecutor([], submit_worker=False)
    coordinator, store, _ = await make_coordinator(
        tmp_path, executor, task_store=task_store
    )
    task_id = await create_task(coordinator, "CAS submit winner")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    ledger = await store.load("root")
    attempt = ledger.attempts[result.attempt_id]

    assert task_store.injected is True
    assert attempt.status == AttemptStatus.SUBMITTED
    assert attempt.submission.summary == "concurrent winner"
    assert attempt.error is None
    assert attempt.execution_stats is None
    assert ledger.tasks[task_id].status == TaskStatus.AWAITING_VALIDATION