diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index 9b75dc38..e4ab0d0d 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -972,6 +972,7 @@ class AgentOrchestrator: # the spawn gate + reaper see them as live immediately (no double-spawn, # no over-reap). Inert when nothing is running. Must run before the # dispatcher/reaper loops launch below. + await self._heal_stale_agent_tokens() await self._readopt_running_agents() # Orphan sandbox sweep: a sandbox whose owning agent container didn't @@ -9783,6 +9784,97 @@ Start now: evidence(task_id="{task_id}") ) return None + async def _read_container_auth_env( + self, container_name: str + ) -> tuple[str, str, str] | None: + """Read (token, agent_id, role) from a running agent container's env. + + Returns None on any probe failure or missing var so the caller can + skip the container (best-effort — the reaper still covers it). + """ + try: + proc = await asyncio.create_subprocess_exec( + "docker", + "exec", + container_name, + "printenv", + "ROBOCO_AGENT_TOKEN", + "ROBOCO_AGENT_ID", + "ROBOCO_AGENT_ROLE", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + ) + stdout, _ = await proc.communicate() + except Exception: + return None + if proc.returncode != 0: + return None + token: str | None = None + agent_id_env: str | None = None + role_env: str | None = None + for line in stdout.decode("utf-8", "replace").splitlines(): + if line.startswith("ROBOCO_AGENT_TOKEN="): + token = line[len("ROBOCO_AGENT_TOKEN=") :] + elif line.startswith("ROBOCO_AGENT_ID="): + agent_id_env = line[len("ROBOCO_AGENT_ID=") :] + elif line.startswith("ROBOCO_AGENT_ROLE="): + role_env = line[len("ROBOCO_AGENT_ROLE=") :] + if not token or not agent_id_env or not role_env: + return None + return token, agent_id_env, role_env + + async def _heal_stale_agent_tokens(self) -> int: + """Kill running agent containers whose ROBOCO_AGENT_TOKEN no longer + verifies against the current ``ROBOCO_AGENT_AUTH_SECRET``. + + A token is baked into the container env at spawn (``_append_agent_auth_env`` + signs with the orchestrator's secret at that moment). If the secret later + drifts — a `.env` change, a compose recreate that reloads the + orchestrator's env without recreating the agent containers, an image + redeploy — the surviving agent keeps sending its old token and the + middleware 401s every verb with "signature mismatch". The container stays + alive (heartbeating), so the reaper never reclaims it and no fresh agent + spawns: the fleet stalls. This self-heals it at startup by killing each + stale-token container so the normal dispatch re-spawns it with a freshly + signed token. + + Inert when the secret is unset (dev): ``verify_agent_token`` fails for + every token without a secret, so the heal would kill the whole fleet — + gated to prod-only. Best-effort: a probe failure leaves the container + alone (the reaper's own liveness path still covers it). + """ + from roboco.agents_config import _auth_secret, verify_agent_token + + if not _auth_secret(): + return 0 + killed = 0 + for slug in AGENT_IMAGES: + try: + is_running, _ = await self._inspect_container_state( + f"roboco-agent-{slug}" + ) + except Exception: + continue + if not is_running: + continue + env = await self._read_container_auth_env(f"roboco-agent-{slug}") + if env is None: + continue + token, agent_id_env, role_env = env + team = get_agent_team(agent_id_env) or "" + if verify_agent_token(token, agent_id_env, role_env, team): + continue + logger.warning( + "Killing agent with a stale auth token at startup; the reaper " + "will re-spawn it with a freshly signed token", + slug=slug, + ) + await self._remove_container(f"roboco-agent-{slug}", teardown_sandbox=False) + killed += 1 + if killed: + logger.info("Healed stale agent tokens at startup", count=killed) + return killed + async def _readopt_running_agents(self) -> int: """Re-adopt still-running agent containers into ``_instances`` at startup. diff --git a/tests/unit/runtime/test_heal_stale_agent_tokens.py b/tests/unit/runtime/test_heal_stale_agent_tokens.py new file mode 100644 index 00000000..291a1f3d --- /dev/null +++ b/tests/unit/runtime/test_heal_stale_agent_tokens.py @@ -0,0 +1,124 @@ +"""Startup self-heal: kill agent containers whose baked-in ROBOCO_AGENT_TOKEN +no longer verifies against the current ``ROBOCO_AGENT_AUTH_SECRET``. + +A token is signed once at spawn. If the secret drifts afterwards (a `.env` +change, a compose recreate that reloads the orchestrator's env without +recreating the agent containers), the surviving agent keeps sending its stale +token and the middleware 401s every verb with "signature mismatch". The +container stays alive (heartbeating) so the reaper never reclaims it and no +fresh agent spawns — the fleet stalls. ``_heal_stale_agent_tokens`` runs at +startup and kills each stale-token container so normal dispatch re-spawns it +with a freshly signed token. +""" + +from __future__ import annotations + +import secrets +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from roboco.agents_config import issue_agent_token +from roboco.runtime.orchestrator import AgentOrchestrator + + +def _orch() -> Any: + orch = AgentOrchestrator.__new__(AgentOrchestrator) # bypass __init__ + orch._instances = {} + return orch + + +@pytest.mark.asyncio +async def test_heal_kills_stale_token_containers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + secret = secrets.token_hex(32) + monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", secret) + orch = _orch() + # be-dev-1 holds a token signed with a DIFFERENT secret (rotated after spawn). + stale_token = issue_agent_token( + "00000000-0000-0000-0001-000000000001", "developer", "backend" + ) + monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", secrets.token_hex(32)) + + async def inspect(name: str) -> tuple[bool, int | None]: + return (name == "roboco-agent-be-dev-1", 0) + + orch._inspect_container_state = AsyncMock(side_effect=inspect) + orch._read_container_auth_env = AsyncMock( + return_value=( + stale_token, + "00000000-0000-0000-0001-000000000001", + "developer", + ) + ) + removed: list[str] = [] + orch._remove_container = AsyncMock( + side_effect=lambda name, **_: removed.append(name) + ) + + n = await orch._heal_stale_agent_tokens() + + assert n == 1 + assert removed == ["roboco-agent-be-dev-1"] + + +@pytest.mark.asyncio +async def test_heal_leaves_valid_token_containers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", secrets.token_hex(32)) + orch = _orch() + be_dev_1 = "00000000-0000-0000-0001-000000000001" + valid_token = issue_agent_token(be_dev_1, "developer", "backend") + + async def inspect(name: str) -> tuple[bool, int | None]: + return (name == "roboco-agent-be-dev-1", 0) + + orch._inspect_container_state = AsyncMock(side_effect=inspect) + orch._read_container_auth_env = AsyncMock( + return_value=(valid_token, be_dev_1, "developer") + ) + orch._remove_container = AsyncMock() + + n = await orch._heal_stale_agent_tokens() + + assert n == 0 + orch._remove_container.assert_not_called() + + +@pytest.mark.asyncio +async def test_heal_inert_when_secret_unset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Dev mode (no secret): verify_agent_token fails for every token, so an + # ungated heal would kill the whole fleet. The heal must short-circuit. + monkeypatch.delenv("ROBOCO_AGENT_AUTH_SECRET", raising=False) + orch = _orch() + orch._inspect_container_state = AsyncMock(return_value=(True, 0)) + orch._read_container_auth_env = AsyncMock( + return_value=("UNSIGNED", "be-dev-1", "developer") + ) + orch._remove_container = AsyncMock() + + n = await orch._heal_stale_agent_tokens() + + assert n == 0 + orch._remove_container.assert_not_called() + + +@pytest.mark.asyncio +async def test_heal_skips_when_env_probe_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", secrets.token_hex(32)) + orch = _orch() + orch._inspect_container_state = AsyncMock(return_value=(True, 0)) + # docker exec fails (container mid-shutdown, etc.) → None → skip, don't kill. + orch._read_container_auth_env = AsyncMock(return_value=None) + orch._remove_container = AsyncMock() + + n = await orch._heal_stale_agent_tokens() + + assert n == 0 + orch._remove_container.assert_not_called()