Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 31 additions & 1 deletion src/operations_center/backends/worker_backend_selector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)",
Expand Down Expand Up @@ -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,)
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down
75 changes: 73 additions & 2 deletions tests/unit/backends/test_worker_backend_selector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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] = []

Expand Down Expand Up @@ -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"]
5 changes: 5 additions & 0 deletions tests/unit/backends/test_worker_backend_selector_cov.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
# --------------------------------------------------------------------------
Expand Down
Loading