[auth] Self-heal stale agent tokens at orchestrator startup

A token is signed once at spawn. If ROBOCO_AGENT_AUTH_SECRET drifts
afterwards (a .env change, a compose recreate that reloads the
orchestrator's env without recreating 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.

_heal_stale_agent_tokens runs at startup (before _readopt_running_agents)
and kills each running agent container whose baked-in token no longer
verifies against the current secret, so normal dispatch re-spawns it
with a freshly signed token. Inert when the secret is unset (dev):
verify fails for every token without a secret, so the heal would kill
the whole fleet without this gate. Best-effort: a probe failure leaves
the container alone (the reaper still covers it).
This commit is contained in:
Renn F
2026-07-06 07:52:17 +02:00
parent 2b8bc10d8d
commit c0328971f2
2 changed files with 216 additions and 0 deletions
+92
View File
@@ -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.
@@ -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()