diff --git a/src/operations_center/backends/worker_backend_selector.py b/src/operations_center/backends/worker_backend_selector.py index 87cd1606..a4ea69ef 100644 --- a/src/operations_center/backends/worker_backend_selector.py +++ b/src/operations_center/backends/worker_backend_selector.py @@ -2,6 +2,7 @@ # Copyright (C) 2026 ProtocolWarden from __future__ import annotations +import os import re from dataclasses import dataclass from datetime import UTC, datetime, timedelta @@ -36,6 +37,12 @@ _SUPPORTED_WORKER_BACKENDS = ("claude_code", "codex_cli", "aider_local", "direct_local") _REMOTE_WORKER_BACKENDS = ("claude_code", "codex_cli") _LOCAL_WORKER_BACKENDS = ("aider_local", "direct_local") +_PROVIDER_TO_WORKER_BACKEND = { + "claude": "claude_code", + "anthropic": "claude_code", + "codex": "codex_cli", + "openai": "codex_cli", +} _TIMEZONE_RESET_RE = re.compile(r"resets\s+(\d{1,2}:\d{2}(?:am|pm))\s+\(([^)]+)\)", re.IGNORECASE) _ISO_RESET_RE = re.compile( r"resets?(?:\s+at)?\s+(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}(?::\d{2})?Z)", @@ -86,6 +93,18 @@ def worker_backend_observed_runtime( } +def _allowed_remote_worker_backends() -> tuple[str, ...]: + raw = os.environ.get("OPERATIONS_CENTER_ALLOWED_PROVIDERS", "").strip() + if not raw: + return _REMOTE_WORKER_BACKENDS + allowed = { + _PROVIDER_TO_WORKER_BACKEND[provider.strip().lower()] + for provider in raw.split(",") + if provider.strip().lower() in _PROVIDER_TO_WORKER_BACKEND + } + return tuple(backend for backend in _REMOTE_WORKER_BACKENDS if backend in allowed) + + def worker_backend_candidates(preferred_backend: str) -> tuple[str, ...]: if preferred_backend not in _SUPPORTED_WORKER_BACKENDS: return (preferred_backend,) @@ -94,9 +113,13 @@ def worker_backend_candidates(preferred_backend: str) -> tuple[str, ...]: # themselves (currently a set of one, so no fallback until executor services # support them natively). if preferred_backend in _REMOTE_WORKER_BACKENDS: - pool = _REMOTE_WORKER_BACKENDS + pool = _allowed_remote_worker_backends() else: pool = _LOCAL_WORKER_BACKENDS + if not pool: + return () + if preferred_backend not in pool: + return pool alternates = tuple(backend for backend in pool if backend != preferred_backend) return (preferred_backend, *alternates) @@ -143,6 +166,13 @@ def select_worker_backend( ) -> WorkerBackendSelection: current = now or datetime.now(UTC) candidates = worker_backend_candidates(preferred_backend) + if not candidates: + return WorkerBackendSelection( + preferred_backend=preferred_backend, + selected_backend=None, + cooldowns={}, + reason=f"no allowed worker backends for preferred backend {preferred_backend}", + ) cooldowns = { backend: _read_worker_backend_cooldown(usage_store, backend, now=current) for backend in candidates diff --git a/tests/unit/backends/test_worker_backend_selector.py b/tests/unit/backends/test_worker_backend_selector.py index 2e436489..72da6762 100644 --- a/tests/unit/backends/test_worker_backend_selector.py +++ b/tests/unit/backends/test_worker_backend_selector.py @@ -19,7 +19,10 @@ def _usage_store() -> SimpleNamespace: ) -def test_select_worker_backend_prefers_alternate_when_preferred_cooling_down() -> None: +def test_select_worker_backend_prefers_alternate_when_preferred_cooling_down( + monkeypatch, +) -> None: + monkeypatch.setenv("OPERATIONS_CENTER_ALLOWED_PROVIDERS", "claude,codex") now = datetime(2026, 5, 25, 16, 0, tzinfo=UTC) def _cooldown(worker_backend: str, *, now): @@ -40,6 +43,31 @@ def _cooldown(worker_backend: str, *, now): assert selection.selected_backend == "codex_cli" +def test_select_worker_backend_respects_allowed_provider_env(monkeypatch) -> None: + now = datetime(2026, 5, 25, 16, 0, tzinfo=UTC) + monkeypatch.setenv("OPERATIONS_CENTER_ALLOWED_PROVIDERS", "claude") + + def _cooldown(worker_backend: str, *, now): + if worker_backend == "claude_code": + return datetime(2026, 5, 25, 17, 0, tzinfo=UTC) + return None + + usage_store = _usage_store() + usage_store.worker_backend_cooldown_until = _cooldown + + selection = select_worker_backend( + preferred_backend="claude_code", + usage_store=usage_store, + dynamic_enabled=True, + now=now, + ) + + assert selection.selected_backend is None + assert selection.reason is not None + assert "claude_code until" in selection.reason + assert "codex_cli" not in selection.reason + + def test_select_worker_backend_not_blocked_by_lone_model_weekly(monkeypatch, tmp_path) -> None: # End-to-end with a real UsageStore: a burnt sonnet weekly must leave the # claude_code backend selectable (haiku/opus still runnable) rather than @@ -79,7 +107,10 @@ def test_parse_worker_backend_reset_handles_relative_message() -> None: assert reset_at == datetime(2026, 5, 25, 21, 0, tzinfo=UTC) -def test_execute_with_worker_backend_round_robin_retries_on_capacity_limit() -> None: +def test_execute_with_worker_backend_round_robin_retries_on_capacity_limit( + monkeypatch, +) -> None: + monkeypatch.setenv("OPERATIONS_CENTER_ALLOWED_PROVIDERS", "claude,codex") cooldowns: dict[str, datetime | None] = {"claude_code": None, "codex_cli": None} calls: list[str] = [] @@ -114,3 +145,43 @@ def _run_once(worker_backend: str): assert executed.selected_backend == "codex_cli" assert executed.fallback_used is True assert calls == ["claude_code", "codex_cli"] + + +def test_round_robin_does_not_fallback_to_disallowed_backend(monkeypatch) -> None: + monkeypatch.setenv("OPERATIONS_CENTER_ALLOWED_PROVIDERS", "claude") + cooldowns: dict[str, datetime | None] = {"claude_code": None, "codex_cli": None} + calls: list[str] = [] + + def _cooldown_until(worker_backend: str, *, now): + return cooldowns.get(worker_backend) + + def _record(worker_backend: str, reset_at, now) -> None: + cooldowns[worker_backend] = reset_at + + usage_store = _usage_store() + usage_store.worker_backend_cooldown_until = _cooldown_until + usage_store.record_worker_backend_cooldown = _record + + def _run_once(worker_backend: str): + calls.append(worker_backend) + return { + "status": "failed", + "error_summary": "usage limit hit, please try again in 5h 0m", + } + + executed = execute_with_worker_backend_round_robin( + preferred_backend="claude_code", + usage_store=usage_store, + dynamic_enabled=True, + execute_once=_run_once, + failed=lambda payload: payload["status"] != "succeeded", + failure_text=lambda payload: payload.get("error_summary"), + ) + + assert executed.selected_backend == "claude_code" + assert executed.fallback_used is False + assert executed.selection.selected_backend is None + assert executed.selection.reason is not None + assert "claude_code until" in executed.selection.reason + assert "codex_cli" not in executed.selection.reason + assert calls == ["claude_code"] diff --git a/tests/unit/backends/test_worker_backend_selector_cov.py b/tests/unit/backends/test_worker_backend_selector_cov.py index f84b7332..68d25213 100644 --- a/tests/unit/backends/test_worker_backend_selector_cov.py +++ b/tests/unit/backends/test_worker_backend_selector_cov.py @@ -23,6 +23,11 @@ NOW = datetime(2026, 6, 2, 12, 0, 0, tzinfo=UTC) +@pytest.fixture(autouse=True) +def _allow_all_remote_worker_backends(monkeypatch) -> None: + monkeypatch.setenv("OPERATIONS_CENTER_ALLOWED_PROVIDERS", "claude,codex") + + # -------------------------------------------------------------------------- # worker_backend_candidates # --------------------------------------------------------------------------