mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
[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:
@@ -972,6 +972,7 @@ class AgentOrchestrator:
|
|||||||
# the spawn gate + reaper see them as live immediately (no double-spawn,
|
# the spawn gate + reaper see them as live immediately (no double-spawn,
|
||||||
# no over-reap). Inert when nothing is running. Must run before the
|
# no over-reap). Inert when nothing is running. Must run before the
|
||||||
# dispatcher/reaper loops launch below.
|
# dispatcher/reaper loops launch below.
|
||||||
|
await self._heal_stale_agent_tokens()
|
||||||
await self._readopt_running_agents()
|
await self._readopt_running_agents()
|
||||||
|
|
||||||
# Orphan sandbox sweep: a sandbox whose owning agent container didn't
|
# Orphan sandbox sweep: a sandbox whose owning agent container didn't
|
||||||
@@ -9783,6 +9784,97 @@ Start now: evidence(task_id="{task_id}")
|
|||||||
)
|
)
|
||||||
return None
|
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:
|
async def _readopt_running_agents(self) -> int:
|
||||||
"""Re-adopt still-running agent containers into ``_instances`` at startup.
|
"""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()
|
||||||
Reference in New Issue
Block a user