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