From a4d71ae42e2bd80090acfe3c18d98f5886ea6c69 Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:00:39 +0200 Subject: [PATCH 1/9] fix(worker): cgroup-safe restart via at(1) detach + wire self-update orchestrator (#890 C3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BLOCKER 1: _detached_restart_worker() uses at(1) → systemd-run --scope → --no-block fallback chain so the restart survives systemd KillMode=control-group teardown. Previously the deploy-helper (child of the worker's cgroup) was SIGTERM'd mid-restart, permanently bricking the worker. BLOCKER 2: Wire _check_update_trigger() into the agent heartbeat loop (reads update-trigger.json), and post_update_startup() into the post-registration path (reads update-in-progress.json marker). The orchestrator now has real callers. Also: drain_complete in heartbeat response, update-outcome HMAC endpoint, C3 deploy-helper subcommands (checkpoint/rollback/ restart-self/health-check), and deploy.py allowlist update. Tests: 42 pass (22 self_update + 20 cluster unit). App-fixture endpoint tests deferred to CI (pre-existing conftest hang). --- scripts/taos-deploy-helper.sh | 237 ++++++++++ tests/test_cluster.py | 111 +++++ tests/test_worker_self_update.py | 500 ++++++++++++++++++++ tinyagentos/routes/cluster.py | 114 ++++- tinyagentos/worker/agent.py | 96 +++- tinyagentos/worker/deploy.py | 5 + tinyagentos/worker/self_update.py | 747 ++++++++++++++++++++++++++++++ 7 files changed, 1803 insertions(+), 7 deletions(-) create mode 100644 tests/test_worker_self_update.py create mode 100644 tinyagentos/worker/self_update.py diff --git a/scripts/taos-deploy-helper.sh b/scripts/taos-deploy-helper.sh index fa0e19389..71ae22e38 100755 --- a/scripts/taos-deploy-helper.sh +++ b/scripts/taos-deploy-helper.sh @@ -237,6 +237,238 @@ cmd_status() { echo '}' } +# ── Worker self-update subcommands (taOS #890 C3) ───────────────────────── + +cmd_checkpoint() { + local manifest_file="${TAOS_CHECKPOINT_MANIFEST:-$INSTALL_DIR/rollback-manifest.json}" + local repo_dir="$INSTALL_DIR/tinyagentos" + local venv="${TAOS_VENV:-$INSTALL_DIR/.venv}" + + log "creating pre-update checkpoint at $manifest_file" + + if [[ ! -d "$repo_dir/.git" ]]; then + die "worker repo not found at $repo_dir" + fi + + local git_sha + git_sha="$(git -C "$repo_dir" rev-parse HEAD)" || die "failed to get current git SHA" + + local git_branch + git_branch="$(git -C "$repo_dir" rev-parse --abbrev-ref HEAD)" || git_branch="detached" + + local deps_snapshot="" + local pkg_manager="unknown" + if [[ -x "$venv/bin/pip" ]]; then + deps_snapshot="$("$venv/bin/pip" freeze 2>/dev/null || true)" + pkg_manager="pip" + fi + + # Detect if uv is in use (uv.lock present in repo root) + if [[ -f "$repo_dir/uv.lock" ]]; then + pkg_manager="uv" + deps_snapshot="uv-lock:$(sha256sum "$repo_dir/uv.lock" 2>/dev/null | awk '{print $1}' || echo "unknown")" + fi + + # Tag the current HEAD so it survives a checkout (detached or branch) + local tag="taos-worker-pre-update-$(date -u +%Y%m%d-%H%M%S)" + git -C "$repo_dir" tag "$tag" HEAD 2>/dev/null || { + # Tag already exists — add a counter suffix + local suffix=1 + while ! git -C "$repo_dir" tag "${tag}-${suffix}" HEAD 2>/dev/null; do + ((suffix++)) + done + tag="${tag}-${suffix}" + } + log "tagged current HEAD as $tag" + + # Build the manifest + cat > "$manifest_file" </dev/null || echo '""'), + "created_at": "$(date -u -Iseconds)", + "hostname": "$(hostname)" +} +MANIFEST + + log "checkpoint saved: tag=$tag sha=${git_sha:0:8} pkg=$pkg_manager" + echo "$tag" # stdout = checkpoint tag for the caller to record +} + +# ── Detached restart helper (taOS #890 C3) ───────────────────────────────── +# The worker service runs with systemd KillMode=control-group, so the default +# ``systemctl stop`` (and the stop phase of ``systemctl restart``) sends +# SIGTERM to every process in the cgroup — including this deploy-helper +# script when it is invoked as a child of the worker. A detached restart +# (via at(1) or systemd-run) survives the teardown because it runs outside +# the service cgroup. + +_detached_restart_worker() { + local restart_cmd="systemctl restart tinyagentos-worker.service 2>/dev/null || systemctl restart taos-worker.service 2>/dev/null || true" + + # 1 — at(1): schedules the restart after the helper exits (cleanest). + if command -v at >/dev/null 2>&1; then + if echo "$restart_cmd" | at now 2>/dev/null; then + log "worker restart scheduled via at(1)" + return 0 + fi + fi + + # 2 — systemd-run --scope: runs outside the service cgroup. + if command -v systemd-run >/dev/null 2>&1; then + if systemd-run --scope --no-block systemctl restart tinyagentos-worker.service 2>/dev/null; then + log "worker restart dispatched via systemd-run" + return 0 + fi + if systemd-run --scope --no-block systemctl restart taos-worker.service 2>/dev/null; then + log "worker restart dispatched via systemd-run (taos-worker)" + return 0 + fi + fi + + # 3 — Last resort: --no-block may race with cgroup teardown, but the + # helper returns immediately so it often wins. + if systemctl restart --no-block tinyagentos-worker.service 2>/dev/null; then + log "worker restart dispatched via systemctl --no-block (may race with cgroup teardown)" + return 0 + fi + if systemctl restart --no-block taos-worker.service 2>/dev/null; then + log "worker restart dispatched via systemctl --no-block (taos-worker)" + return 0 + fi + + log "WARN: could not restart worker service via any mechanism" + return 1 +} + +cmd_rollback() { + local checkpoint_tag="${1:-}" + local manifest_file="${TAOS_CHECKPOINT_MANIFEST:-$INSTALL_DIR/rollback-manifest.json}" + local repo_dir="$INSTALL_DIR/tinyagentos" + local venv="${TAOS_VENV:-$INSTALL_DIR/.venv}" + + if [[ -z "$checkpoint_tag" ]]; then + # Read the tag from the manifest if not provided + if [[ -f "$manifest_file" ]]; then + checkpoint_tag="$(python3 -c "import json; print(json.load(open('$manifest_file')).get('checkpoint_tag',''))" 2>/dev/null || true)" + fi + if [[ -z "$checkpoint_tag" ]]; then + die "no checkpoint tag provided and no manifest found at $manifest_file" + fi + fi + + log "rolling back to checkpoint tag $checkpoint_tag" + + if [[ ! -d "$repo_dir/.git" ]]; then + die "worker repo not found at $repo_dir" + fi + + # Verify the tag exists + if ! git -C "$repo_dir" rev-parse --verify "$checkpoint_tag^{commit}" >/dev/null 2>&1; then + die "checkpoint tag $checkpoint_tag not found" + fi + + # Restore the code. We do NOT stop the worker service first because + # this helper runs inside the worker's systemd cgroup and systemctl + # stop would kill us before the rollback completes. The worker was + # drained before this phase; a restart at the end picks up the + # restored code. + git -C "$repo_dir" checkout --quiet "$checkpoint_tag" || { + die "git checkout $checkpoint_tag failed" + } + log "restored code to $checkpoint_tag ($(git -C "$repo_dir" rev-parse --short HEAD))" + + # Reinstall dependencies + if [[ -f "$manifest_file" ]]; then + local pkg_manager + pkg_manager="$(python3 -c "import json; print(json.load(open('$manifest_file')).get('package_manager','pip'))" 2>/dev/null || echo "pip")" + if [[ "$pkg_manager" == "uv" ]] && command -v uv >/dev/null 2>&1; then + log "reinstalling deps with uv sync" + cd "$repo_dir" && uv sync --frozen 2>/dev/null || uv sync || log "WARN: uv sync had errors — continuing" + else + if [[ -x "$venv/bin/pip" ]]; then + log "reinstalling deps with pip" + "$venv/bin/pip" install -q -e "$repo_dir[worker]" 2>/dev/null || \ + "$venv/bin/pip" install -q -e "$repo_dir" || \ + log "WARN: pip install had errors — continuing" + fi + fi + fi + + # Detached restart so the cgroup teardown doesn't kill us mid-rollback. + _detached_restart_worker + log "rollback complete — worker restart initiated from $checkpoint_tag" +} + +cmd_restart_self() { + log "restarting worker service" + _detached_restart_worker +} + +cmd_health_check() { + local manifest_file="${TAOS_CHECKPOINT_MANIFEST:-$INSTALL_DIR/rollback-manifest.json}" + + local ok=true + local failures=() + + # 1. Check service is active + if systemctl is-active --quiet tinyagentos-worker.service 2>/dev/null; then + log "health-check: tinyagentos-worker.service is active" + elif systemctl is-active --quiet taos-worker.service 2>/dev/null; then + log "health-check: taos-worker.service is active" + else + ok=false + failures+=("worker service not active") + # Try launchd (macOS) + if command -v launchctl >/dev/null 2>&1; then + if launchctl list | grep -q tinyagentos-worker 2>/dev/null; then + log "health-check: tinyagentos-worker found in launchd" + ok=true + failures=() + else + failures+=("worker not found in systemd or launchd") + fi + fi + fi + + # 2. Check worker port is listening (default 9898) + local port="${TAOS_WORKER_PORT:-9898}" + if command -v ss >/dev/null 2>&1; then + if ss -tlnp 2>/dev/null | grep -q ":$port "; then + log "health-check: port $port is listening" + else + ok=false + failures+=("port $port not listening") + fi + elif command -v netstat >/dev/null 2>&1; then + if netstat -tlnp 2>/dev/null | grep -q ":$port "; then + log "health-check: port $port is listening" + else + ok=false + failures+=("port $port not listening") + fi + fi + + # 3. Check the checkpoint manifest exists + if [[ -f "$manifest_file" ]]; then + log "health-check: checkpoint manifest present" + else + log "health-check: no checkpoint manifest (not an error during normal operation)" + fi + + if $ok; then + log "health-check: PASS" + echo '{"healthy": true}' + else + log "health-check: FAIL — ${failures[*]}" + echo "{\"healthy\": false, \"failures\": $(python3 -c "import sys,json; print(json.dumps(sys.argv[1:]))" "${failures[@]}" 2>/dev/null || echo '[]')}" + exit 1 + fi +} + # --- dispatch --------------------------------------------------------------- case "${1:-help}" in install-ollama) cmd_install_ollama ;; @@ -246,10 +478,15 @@ case "${1:-help}" in install-rknpu) cmd_install_rknpu ;; update-worker) cmd_update_worker ;; status) cmd_status ;; + checkpoint) cmd_checkpoint ;; + rollback) shift; cmd_rollback "$@" ;; + restart-self) cmd_restart_self ;; + health-check) shift; cmd_health_check "$@" ;; help|*) echo "usage: taos-deploy-helper.sh " echo "commands: install-ollama, install-exo, install-llama-cpp [--cuda]," echo " install-vllm, install-rknpu, update-worker, status" + echo " checkpoint, rollback [], restart-self, health-check" exit 1 ;; esac diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 7fcdf5d2d..02b01efeb 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -578,3 +578,114 @@ async def test_updating_workers_excluded_from_routing(self): result = mgr.get_workers_for_capability("chat") assert len(result) == 1 assert result[0].name == "online-gpu" + + +# ── Update-outcome endpoint (taOS #890 C3) ─────────────────────────── + + +@pytest.mark.asyncio +class TestUpdateOutcomeEndpoint: + async def test_update_outcome_success(self, client, app): + """Worker reports successful self-update.""" + from unittest.mock import patch + + # Register a worker so the endpoint can find it + mgr = app.state.cluster_manager + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + payload = { + "name": "gpu-box", + "outcome": "success", + "from_version": "abc1234def", + "to_version": "def5678abc", + } + + # Bypass HMAC for the test — we test HMAC separately + with patch( + "tinyagentos.routes.cluster.require_worker_hmac", + side_effect=lambda r: None, + ): + resp = await client.post( + "/api/cluster/workers/gpu-box/update-outcome", + json=payload, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["worker"] == "gpu-box" + assert data["outcome"] == "success" + assert data["acknowledged"] is True + + async def test_update_outcome_rollback(self, client, app): + """Worker reports a rollback after failed update.""" + from unittest.mock import patch + + mgr = app.state.cluster_manager + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + payload = { + "name": "gpu-box", + "outcome": "rollback", + "from_version": "abc1234def", + "to_version": "def5678abc", + "failure_reason": "health-check: port not listening", + "rollback_to": "abc1234def", + } + + with patch( + "tinyagentos.routes.cluster.require_worker_hmac", + side_effect=lambda r: None, + ): + resp = await client.post( + "/api/cluster/workers/gpu-box/update-outcome", + json=payload, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["outcome"] == "rollback" + assert data["acknowledged"] is True + + async def test_update_outcome_worker_not_found(self, client): + """404 when the worker is not registered.""" + from unittest.mock import patch + + payload = { + "name": "nonexistent", + "outcome": "success", + "from_version": "aaa", + "to_version": "bbb", + } + + with patch( + "tinyagentos.routes.cluster.require_worker_hmac", + side_effect=lambda r: None, + ): + resp = await client.post( + "/api/cluster/workers/nonexistent/update-outcome", + json=payload, + ) + assert resp.status_code == 404 + + async def test_update_outcome_unknown_outcome(self, client, app): + """400 when the outcome is not 'success' or 'rollback'.""" + from unittest.mock import patch + + mgr = app.state.cluster_manager + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + payload = { + "name": "gpu-box", + "outcome": "unknown-status", + } + + with patch( + "tinyagentos.routes.cluster.require_worker_hmac", + side_effect=lambda r: None, + ): + resp = await client.post( + "/api/cluster/workers/gpu-box/update-outcome", + json=payload, + ) + assert resp.status_code == 400 diff --git a/tests/test_worker_self_update.py b/tests/test_worker_self_update.py new file mode 100644 index 000000000..b38bac7dd --- /dev/null +++ b/tests/test_worker_self_update.py @@ -0,0 +1,500 @@ +"""Tests for the worker self-update orchestrator (taOS #890 C3).""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from tinyagentos.worker.self_update import ( + clear_update_marker, + read_update_marker, + _write_update_marker, + _detect_package_manager, +) + + +# ── Marker file helpers ─────────────────────────────────────────────── + + +class TestUpdateMarker: + """Tests for the in-progress update marker (filesystem I/O).""" + + def test_write_and_read_marker(self, tmp_path: Path): + state_dir = tmp_path / "worker-state" + _write_update_marker( + state_dir, + checkpoint_tag="taos-worker-pre-update-20260717-120000", + from_sha="abc1234", + to_sha="def5678", + ) + + marker = read_update_marker(state_dir) + assert marker is not None + assert marker["checkpoint_tag"] == "taos-worker-pre-update-20260717-120000" + assert marker["from_sha"] == "abc1234" + assert marker["to_sha"] == "def5678" + + def test_read_marker_nonexistent(self, tmp_path: Path): + state_dir = tmp_path / "nonexistent" + assert read_update_marker(state_dir) is None + + def test_clear_marker(self, tmp_path: Path): + state_dir = tmp_path / "worker-state" + _write_update_marker(state_dir, "tag", "a", "b") + assert read_update_marker(state_dir) is not None + + clear_update_marker(state_dir) + assert read_update_marker(state_dir) is None + + def test_clear_marker_idempotent(self, tmp_path: Path): + """Clearing a non-existent marker should not raise.""" + state_dir = tmp_path / "no-marker" + clear_update_marker(state_dir) # should not raise + + def test_read_marker_invalid_json(self, tmp_path: Path): + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + (state_dir / "update-in-progress.json").write_text("not json") + + assert read_update_marker(state_dir) is None + + def test_write_marker_creates_directory(self, tmp_path: Path): + state_dir = tmp_path / "deeply" / "nested" / "state" + _write_update_marker(state_dir, "tag", "a", "b") + assert read_update_marker(state_dir) is not None + + +# ── Package manager detection ──────────────────────────────────────── + + +class TestDetectPackageManager: + """Tests for _detect_package_manager().""" + + def test_pip_default(self, monkeypatch, tmp_path: Path): + """When no uv.lock exists, should return 'pip'.""" + # Override _repo_dir and _install_dir to use tmp_path + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_install_dir", lambda: tmp_path) + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + # No uv.lock in tmp_path — should detect pip + assert su._detect_package_manager() == "pip" + + def test_uv_detected_when_lockfile_present(self, monkeypatch, tmp_path: Path): + """When uv.lock exists, should return 'uv'.""" + import tinyagentos.worker.self_update as su + + (tmp_path / "uv.lock").write_text("") + monkeypatch.setattr(su, "_install_dir", lambda: tmp_path) + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + assert su._detect_package_manager() == "uv" + + +# ── create_checkpoint ───────────────────────────────────────────────── + + +class TestCreateCheckpoint: + @pytest.mark.asyncio + async def test_checkpoint_success(self, monkeypatch, tmp_path: Path): + """create_checkpoint should return tag and SHA on success.""" + import tinyagentos.worker.self_update as su + + # Mock _run_helper to return success + async def mock_run_helper(args, timeout=600): + return {"ok": True, "output": "taos-worker-pre-update-20260717-120000", "exit_code": 0} + + # Mock _run_git to return a known SHA + async def mock_run_git(args, cwd=None, timeout=120): + return 0, "abc1234def5678\n" + + monkeypatch.setattr(su, "_run_helper", mock_run_helper) + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.create_checkpoint() + assert result["ok"] is True + assert result["checkpoint_tag"] == "taos-worker-pre-update-20260717-120000" + assert result["git_sha"] == "abc1234def5678" + + @pytest.mark.asyncio + async def test_checkpoint_helper_fails(self, monkeypatch, tmp_path: Path): + """create_checkpoint should return ok=False when helper fails.""" + import tinyagentos.worker.self_update as su + + async def mock_run_helper(args, timeout=600): + return {"ok": False, "output": "helper not found", "exit_code": -1} + + async def mock_run_git(args, cwd=None, timeout=120): + return 0, "sha1234\n" + + monkeypatch.setattr(su, "_run_helper", mock_run_helper) + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.create_checkpoint() + assert result["ok"] is False + assert result["checkpoint_tag"] == "" + + +# ── pull_update ─────────────────────────────────────────────────────── + + +class TestPullUpdate: + @pytest.mark.asyncio + async def test_pull_success(self, monkeypatch, tmp_path: Path): + """pull_update should fetch and checkout successfully.""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + + async def mock_run_git(args, cwd=None, timeout=120): + if args[0] == "fetch": + return 0, "" + if args[0] == "checkout": + return 0, "" + if args[0] == "rev-parse": + return 0, "def5678\n" + return 1, "unknown" + + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.pull_update("origin/master") + assert result["ok"] is True + + @pytest.mark.asyncio + async def test_pull_fetch_fails(self, monkeypatch, tmp_path: Path): + """pull_update should fail when git fetch fails.""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + + async def mock_run_git(args, cwd=None, timeout=120): + if args[0] == "fetch": + return 128, "fatal: could not fetch" + return 1, "unknown" + + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.pull_update("origin/master") + assert result["ok"] is False + + @pytest.mark.asyncio + async def test_pull_checkout_fails(self, monkeypatch, tmp_path: Path): + """pull_update should fail when git checkout fails.""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + + async def mock_run_git(args, cwd=None, timeout=120): + if args[0] == "fetch": + return 0, "" + if args[0] == "checkout": + return 128, "fatal: could not checkout" + return 1, "unknown" + + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.pull_update("origin/master") + assert result["ok"] is False + + +# ── update_dependencies ─────────────────────────────────────────────── + + +class TestUpdateDependencies: + @pytest.mark.asyncio + async def test_pip_install(self, monkeypatch, tmp_path: Path): + """update_dependencies should run pip install -e .[worker].""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_detect_package_manager", lambda: "pip") + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + venv_dir = tmp_path / ".venv" + venv_dir.mkdir() + (venv_dir / "bin").mkdir() + (venv_dir / "bin" / "pip").write_text("#!/bin/sh\necho fake pip") + (venv_dir / "bin" / "pip").chmod(0o755) + monkeypatch.setattr(su, "_venv_dir", lambda: venv_dir) + + # Mock subprocess for pip + async def mock_communicate(): + return b"installed ok", b"" + + mock_proc = MagicMock() + mock_proc.returncode = 0 + mock_proc.communicate = mock_communicate + + with patch("asyncio.create_subprocess_exec", AsyncMock(return_value=mock_proc)) as mock_exec: + result = await su.update_dependencies() + assert result["ok"] is True + assert result["package_manager"] == "pip" + mock_exec.assert_called_once() + + @pytest.mark.asyncio + async def test_pip_not_found(self, monkeypatch, tmp_path: Path): + """update_dependencies should fail when pip is not found.""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_detect_package_manager", lambda: "pip") + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + venv_dir = tmp_path / ".venv" + venv_dir.mkdir() + monkeypatch.setattr(su, "_venv_dir", lambda: venv_dir) + + result = await su.update_dependencies() + assert result["ok"] is False + + +# ── run_migrations ──────────────────────────────────────────────────── + + +class TestRunMigrations: + @pytest.mark.asyncio + async def test_migrations_deferred(self): + """run_migrations should return ok (deferred to post-restart).""" + import tinyagentos.worker.self_update as su + + result = await su.run_migrations() + assert result["ok"] is True + + +# ── run_full_update integration test ────────────────────────────────── + + +class TestRunFullUpdate: + @pytest.mark.asyncio + async def test_full_update_success_flow(self, monkeypatch, tmp_path: Path): + """run_full_update should orchestrate all phases and return ok=True.""" + import tinyagentos.worker.self_update as su + + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + + # Mock agent + agent = MagicMock() + agent.name = "test-worker" + agent._signing_key = b"fake-key" + + async def mock_heartbeat_ok(*args, **kwargs): + return 200 + + agent.report_update_available = AsyncMock(side_effect=mock_heartbeat_ok) + agent.initiate_self_drain = AsyncMock(side_effect=mock_heartbeat_ok) + agent.notify_drain_complete = AsyncMock(side_effect=mock_heartbeat_ok) + agent.heartbeat = AsyncMock(side_effect=mock_heartbeat_ok) + + # Mock all the sub-operations + async def mock_checkpoint(): + return {"ok": True, "checkpoint_tag": "tag-123", "git_sha": "from-sha", "output": "tag-123", "exit_code": 0} + + async def mock_pull(target_ref): + return {"ok": True, "output": "ok", "exit_code": 0} + + async def mock_deps(): + return {"ok": True, "output": "ok", "exit_code": 0, "package_manager": "pip"} + + async def mock_migrations(): + return {"ok": True, "output": "ok", "exit_code": 0} + + async def mock_restart(): + return {"ok": True, "output": "restarted", "exit_code": 0} + + async def mock_git(args, cwd=None, timeout=120): + if args[0] == "rev-parse": + return 0, "to-sha-5678\n" + return 0, "" + + async def short_sleep(*args, **kwargs): + pass # Don't actually sleep in tests + + monkeypatch.setattr(su, "create_checkpoint", mock_checkpoint) + monkeypatch.setattr(su, "pull_update", mock_pull) + monkeypatch.setattr(su, "update_dependencies", mock_deps) + monkeypatch.setattr(su, "run_migrations", mock_migrations) + monkeypatch.setattr(su, "restart_service", mock_restart) + monkeypatch.setattr(su, "_run_git", mock_git) + monkeypatch.setattr(su, "_install_dir", lambda: tmp_path) + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + monkeypatch.setattr(su, "_venv_dir", lambda: tmp_path) + monkeypatch.setattr(su, "_write_update_marker", lambda *args: None) + monkeypatch.setattr(asyncio, "sleep", short_sleep) + + result = await su.run_full_update( + target_ref="origin/master", + controller_url="http://localhost:9898", + agent=agent, + state_dir=state_dir, + graceful=False, # Skip drain wait for test speed + ) + + assert result["ok"] is True + assert "checkpoint" in result["phases"] + assert result["phases"]["checkpoint"]["ok"] is True + assert "pull" in result["phases"] + assert result["phases"]["pull"]["ok"] is True + assert "dependencies" in result["phases"] + assert result["phases"]["dependencies"]["ok"] is True + assert "restart" in result["phases"] + + @pytest.mark.asyncio + async def test_full_update_checkpoint_fails(self, monkeypatch, tmp_path: Path): + """run_full_update should abort when checkpoint fails.""" + import tinyagentos.worker.self_update as su + + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + + agent = MagicMock() + agent.name = "test-worker" + + async def mock_checkpoint(): + return {"ok": False, "checkpoint_tag": "", "git_sha": "", "output": "failed", "exit_code": 1} + + monkeypatch.setattr(su, "create_checkpoint", mock_checkpoint) + + result = await su.run_full_update( + target_ref="origin/master", + controller_url="http://localhost:9898", + agent=agent, + state_dir=state_dir, + ) + assert result["ok"] is False + assert "checkpoint failed" in result["error"] + + @pytest.mark.asyncio + async def test_full_update_pull_fails(self, monkeypatch, tmp_path: Path): + """run_full_update should abort when pull fails (no rollback needed).""" + import tinyagentos.worker.self_update as su + + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + + agent = MagicMock() + agent.name = "test-worker" + agent.report_update_available = AsyncMock(return_value=200) + agent.initiate_self_drain = AsyncMock(return_value=200) + agent.notify_drain_complete = AsyncMock(return_value=200) + + async def mock_checkpoint(): + return {"ok": True, "checkpoint_tag": "tag-123", "git_sha": "from-sha", "output": "tag-123", "exit_code": 0} + + async def mock_pull(target_ref): + return {"ok": False, "output": "checkout failed", "exit_code": 1} + + monkeypatch.setattr(su, "create_checkpoint", mock_checkpoint) + monkeypatch.setattr(su, "pull_update", mock_pull) + monkeypatch.setattr(su, "_install_dir", lambda: tmp_path) + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + monkeypatch.setattr(su, "_venv_dir", lambda: tmp_path) + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + + result = await su.run_full_update( + target_ref="origin/master", + controller_url="http://localhost:9898", + agent=agent, + state_dir=state_dir, + ) + assert result["ok"] is False + assert "pull failed" in result["error"] + + +# ── post_update_startup ─────────────────────────────────────────────── + + +class TestPostUpdateStartup: + @pytest.mark.asyncio + async def test_no_marker_returns_none(self, monkeypatch, tmp_path: Path): + """post_update_startup should return None when no marker exists.""" + import tinyagentos.worker.self_update as su + + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + + agent = MagicMock() + + result = await su.post_update_startup( + controller_url="http://controller:9898", + agent=agent, + state_dir=state_dir, + ) + assert result is None + + @pytest.mark.asyncio + async def test_health_check_pass(self, monkeypatch, tmp_path: Path): + """post_update_startup should report success on healthy restart.""" + import tinyagentos.worker.self_update as su + + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + + _write_update_marker(state_dir, "tag-123", "from-sha", "to-sha") + + agent = MagicMock() + agent.name = "test-worker" + agent._signing_key = b"fake-key" + + async def mock_health_check(timeout=30): + return {"ok": True, "output": "healthy", "exit_code": 0} + + async def mock_signal(*args, **kwargs): + return 200 + + monkeypatch.setattr(su, "run_health_check", mock_health_check) + monkeypatch.setattr(su, "signal_update_outcome", mock_signal) + monkeypatch.setattr(su, "POST_RESTART_GRACE_PERIOD", 0) + monkeypatch.setattr(su, "clear_update_marker", lambda d: None) + + result = await su.post_update_startup( + controller_url="http://controller:9898", + agent=agent, + state_dir=state_dir, + ) + assert result is not None + assert result["ok"] is True + assert result["outcome"] == "success" + + @pytest.mark.asyncio + async def test_health_check_fail_triggers_rollback(self, monkeypatch, tmp_path: Path): + """post_update_startup should rollback on failed health check.""" + import tinyagentos.worker.self_update as su + + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + + _write_update_marker(state_dir, "tag-123", "from-sha", "to-sha") + + agent = MagicMock() + agent.name = "test-worker" + agent._signing_key = b"fake-key" + + call_log = [] + + async def mock_health_check(timeout=30): + return {"ok": False, "output": "port not listening", "exit_code": 1} + + async def mock_rollback(checkpoint_tag=None): + call_log.append(("rollback", checkpoint_tag)) + return {"ok": True, "output": "rolled back", "exit_code": 0} + + async def mock_signal(*args, **kwargs): + call_log.append(("signal", kwargs.get("outcome"))) + return 200 + + monkeypatch.setattr(su, "run_health_check", mock_health_check) + monkeypatch.setattr(su, "rollback_to_checkpoint", mock_rollback) + monkeypatch.setattr(su, "signal_update_outcome", mock_signal) + monkeypatch.setattr(su, "POST_RESTART_GRACE_PERIOD", 0) + monkeypatch.setattr(su, "clear_update_marker", lambda d: None) + + result = await su.post_update_startup( + controller_url="http://controller:9898", + agent=agent, + state_dir=state_dir, + ) + assert result is not None + assert result["ok"] is False + assert result["outcome"] == "rollback" + assert ("rollback", "tag-123") in call_log + assert ("signal", "rollback") in call_log diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index 92bd67cda..c6d02d492 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -563,7 +563,21 @@ async def worker_heartbeat(request: Request, body: HeartbeatBody): if not ok: return JSONResponse({"error": "Worker not registered"}, status_code=404) cluster = request.app.state.cluster_manager - return {"status": "ok", "generation": cluster.generation} + + # Include drain status in the response so the worker's self-update + # orchestrator can detect when in-flight leases are all released and + # proceed with the update without waiting the full timeout (taOS #890 C3). + worker_obj = cluster.get_worker(body.name) + drain_complete = False + if worker_obj is not None and worker_obj.status == "draining": + active = [ + lid for lid, lease in cluster._leases.items() + if (parsed := cluster._parse_resource_id(lease.resource_id)) + and parsed[0] == body.name + ] + drain_complete = len(active) == 0 + + return {"status": "ok", "generation": cluster.generation, "drain_complete": drain_complete} @router.delete("/api/cluster/workers/{name}") @@ -1560,3 +1574,101 @@ async def update_all_workers(request: Request): "skipped": skipped, "total_targets": len(targets), } + + +# ── Worker self-update outcome reporting (taOS #890 C3) ────────────────── + + +class UpdateOutcomeBody(BaseModel): + """Payload for the worker to report its update outcome.""" + name: str + outcome: str # "success" | "rollback" + from_version: str = "" + to_version: str = "" + failure_reason: str = "" + rollback_to: str = "" + + +@router.post("/api/cluster/workers/{name}/update-outcome") +async def report_update_outcome(request: Request, name: str, body: UpdateOutcomeBody): + """Record the outcome of a worker self-update. + + Called by the worker after restart + health-check, or on rollback. + + HMAC-signed: the worker's signing key validates that the outcome + came from the actual worker, not a spoofed request. + """ + try: + await require_worker_hmac(request) + except _HMACError: + return JSONResponse({"error": "hmac verification failed"}, status_code=403) + + cluster = request.app.state.cluster_manager + worker = cluster.get_worker(name) + if not worker: + return JSONResponse({"error": f"Worker '{name}' not found"}, status_code=404) + + outcome = body.outcome + logger.info( + "worker '%s' reported update outcome: %s (from=%s to=%s)", + name, outcome, body.from_version[:8], body.to_version[:8], + ) + + notifications = getattr(request.app.state, "notifications", None) + + if outcome == "success": + # Worker is already back online via re-registration. + # Nothing to do — the heartbeat loop handles it. + logger.info("worker '%s' self-update SUCCESS", name) + # Emit a success notification so the operator has an audit trail. + if notifications: + try: + await notifications.emit_event( + "worker.update-success", + f"Worker '{name}' self-update succeeded", + ( + f"Updated from {body.from_version[:8] or 'unknown'} " + f"to {body.to_version[:8] or 'unknown'}." + ), + level="info", + ) + except Exception: + logger.exception( + "notification emit failed for update success" + ) + + elif outcome == "rollback": + # Worker rolled back after a failed health-check. + logger.warning( + "worker '%s' self-update ROLLBACK: %s (from %s to %s, rolled back to %s)", + name, + body.failure_reason or "unknown failure", + body.from_version[:8], + body.to_version[:8], + body.rollback_to[:8] or "checkpoint", + ) + if notifications: + try: + await notifications.emit_event( + "worker.update-rollback", + f"Worker '{name}' rolled back after failed update", + ( + f"Update from {body.from_version[:8]} to {body.to_version[:8]} " + f"failed: {body.failure_reason or 'health check failed'}. " + f"Rolled back to {body.rollback_to[:8] or 'checkpoint'}." + ), + level="warning", + ) + except Exception: + logger.exception("notification emit failed for update rollback") + + else: + return JSONResponse( + {"error": f"unknown outcome: {outcome}"}, status_code=400 + ) + + return { + "worker": name, + "outcome": outcome, + "acknowledged": True, + } diff --git a/tinyagentos/worker/agent.py b/tinyagentos/worker/agent.py index 2a0168489..5c6f1146b 100644 --- a/tinyagentos/worker/agent.py +++ b/tinyagentos/worker/agent.py @@ -762,16 +762,71 @@ async def notify_drain_complete(self) -> int: """Notify the controller that the worker's drain is complete. After in-flight work finishes, the worker sends one final - heartbeat with status="updating" to signal readiness for - the update. The controller can then proceed with the update - deploy. + heartbeat to confirm readiness for the update. The controller + can then proceed with the update deploy. Returns the HTTP status code from the controller. """ logger.info("worker '%s': drain complete, ready for update", self.name) - self._lifecycle_status = "updating" - self._lifecycle_reason = "drain-complete" - return await self.heartbeat(status="updating") + return await self.heartbeat() + + # ── Self-update trigger (taOS #890 C3) ───────────────────────────── + + async def _check_update_trigger(self) -> bool: + """Check for a pending self-update trigger file and execute it. + + The trigger file is written atomically by an external mechanism + (e.g. the controller via a deploy command or a cron-managed poller). + Its presence signals that the worker should run the full self-update + lifecycle (checkpoint → drain → pull → deps → restart). + + Returns True if a trigger was found and processed; False otherwise. + """ + trigger_path = self._state_dir / "update-trigger.json" + if not trigger_path.exists(): + return False + + try: + trigger = json.loads(trigger_path.read_text()) + except (json.JSONDecodeError, OSError): + logger.warning("update trigger file unreadable — removing") + trigger_path.unlink(missing_ok=True) + return False + + target_ref = trigger.get("target_ref", "") + if not target_ref: + logger.warning("update trigger missing target_ref — removing") + trigger_path.unlink(missing_ok=True) + return False + + graceful = trigger.get("graceful", True) + + logger.info( + "self-update: trigger received — target=%s graceful=%s", + target_ref, graceful, + ) + + # Remove the trigger file BEFORE running the update so a + # duplicate run is never launched. + trigger_path.unlink(missing_ok=True) + + try: + from tinyagentos.worker.self_update import run_full_update + result = await run_full_update( + target_ref=target_ref, + controller_url=self.controller_url, + agent=self, + state_dir=self._state_dir, + graceful=graceful, + ) + if not result.get("ok"): + logger.error( + "self-update: update failed: %s", result.get("error", "unknown") + ) + except Exception: + logger.error("self-update: update threw exception", exc_info=True) + + return True def _log_repair_instruction(self) -> None: logger.error( @@ -821,6 +876,29 @@ async def run(self): if result is True: logger.info(f"worker '{self.name}' registered with {self.controller_url}") _in_repair = False + + # ── Post-update startup hook (taOS #890 C3) ───────── + # If an update-in-progress marker exists from a + # pre-restart checkpoint, run the health-check and + # signal the outcome (success or rollback) to the + # controller. + try: + from tinyagentos.worker.self_update import post_update_startup + outcome = await post_update_startup( + controller_url=self.controller_url, + agent=self, + state_dir=self._state_dir, + ) + if outcome is not None: + logger.info( + "self-update: post-restart outcome=%s", + outcome.get("outcome", "unknown"), + ) + except Exception: + logger.warning( + "post_update_startup hook failed — continuing", + exc_info=True, + ) continue if result == _NEEDS_REPAIR: # Controller rejected our key -- enter needs-re-pair state. @@ -866,6 +944,12 @@ async def run(self): # registered flag yet; the controller may still know # us when it comes back. Just retry on next tick. pass + + # Check for pending self-update trigger on each heartbeat + # cycle. The check is cheap (a file stat) and the trigger + # file is removed atomically before the update runs. + await self._check_update_trigger() + await asyncio.sleep(5) finally: if self._update_service is not None: diff --git a/tinyagentos/worker/deploy.py b/tinyagentos/worker/deploy.py index 9000cecc3..eb7a387d6 100644 --- a/tinyagentos/worker/deploy.py +++ b/tinyagentos/worker/deploy.py @@ -30,6 +30,11 @@ "install-rknpu", "update-worker", "status", + # Worker self-update subcommands (taOS #890 C3). + "checkpoint", + "rollback", + "restart-self", + "health-check", } diff --git a/tinyagentos/worker/self_update.py b/tinyagentos/worker/self_update.py new file mode 100644 index 000000000..82954f872 --- /dev/null +++ b/tinyagentos/worker/self_update.py @@ -0,0 +1,747 @@ +"""Worker self-update orchestrator — checkpoint, install, restart, rollback. + +Part of taOS #890: worker auto-update lifecycle. Coordinates with the +deploy helper (via passwordless sudo) for privileged operations and with +the worker agent for controller signaling. + +Flow: + checkpoint → drain → pull → install deps → restart + (post-restart) → health-check → outcome signal | rollback +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import shutil +from pathlib import Path + +logger = logging.getLogger(__name__) + +# Path to the deploy helper installed by install-worker.sh. +DEPLOY_HELPER = "/usr/local/bin/taos-deploy-helper" + +# How long to wait for the worker port to come back after restart (seconds). +POST_RESTART_HEALTH_TIMEOUT = 30 + +# How long after restart to wait before the first health probe (seconds). +# Gives backends (Ollama, llama.cpp, etc.) time to stabilise. +POST_RESTART_GRACE_PERIOD = 15 + +# Marker file written by the pre-restart phase so the post-restart +# startup hook knows an update was in progress. +_UPDATE_IN_PROGRESS_MARKER = "update-in-progress.json" + + +def _install_dir() -> Path: + """Return the worker install directory (TAOS_INSTALL_DIR or default).""" + env = os.environ.get("TAOS_INSTALL_DIR", "") + if env.strip(): + return Path(env.strip()) + return Path.home() / ".local" / "share" / "tinyagentos-worker" + + +def _repo_dir() -> Path: + """Return the taOS git checkout directory on this worker.""" + return _install_dir() / "tinyagentos" + + +def _venv_dir() -> Path: + """Return the worker's virtualenv directory.""" + return _install_dir() / ".venv" + + +async def _run_helper( + args: list[str], + timeout: float = 600, +) -> dict: + """Run a deploy-helper command via passwordless sudo. + + Uses ``asyncio.create_subprocess_exec`` (no shell) with a fixed + binary path — same security pattern as ``deploy.py``. + + Returns a dict with keys: ok (bool), output (str), exit_code (int). + """ + if not shutil.which(DEPLOY_HELPER): + return { + "ok": False, + "output": f"deploy helper not found at {DEPLOY_HELPER}", + "exit_code": -1, + } + + cmd = ["sudo", DEPLOY_HELPER] + args + logger.info("self-update: running %s", " ".join(cmd)) + + try: + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout) + output = stdout.decode("utf-8", errors="replace") if stdout else "" + ok = proc.returncode == 0 + + if ok: + logger.info("self-update: '%s' completed", args[0]) + else: + logger.error( + "self-update: '%s' failed (exit %d): %s", + args[0], proc.returncode, output[-500:], + ) + + return { + "ok": ok, + "output": output.strip(), + "exit_code": proc.returncode or 0, + } + except asyncio.TimeoutError: + logger.error("self-update: '%s' timed out after %.0fs", args[0], timeout) + return { + "ok": False, + "output": f"timed out after {timeout:.0f}s", + "exit_code": -1, + } + except Exception as exc: + logger.error("self-update: '%s' failed: %s", args[0], exc) + return { + "ok": False, + "output": str(exc), + "exit_code": -1, + } + + +async def _run_git( + args: list[str], + cwd: Path | None = None, + timeout: float = 120, +) -> tuple[int, str]: + """Run a git command safely (list of args, no shell). + + Returns (returncode, stdout_or_stderr). + """ + repo = cwd or _repo_dir() + proc = await asyncio.create_subprocess_exec( + "git", *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + cwd=str(repo), + ) + stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout) + return proc.returncode or 0, (stdout.decode("utf-8", errors="replace") if stdout else "") + + +def _detect_package_manager() -> str: + """Return 'uv' if uv.lock exists in the repo, else 'pip'.""" + if (_repo_dir() / "uv.lock").exists(): + return "uv" + return "pip" + + +def _write_update_marker( + state_dir: Path, + checkpoint_tag: str, + from_sha: str, + to_sha: str, +) -> None: + """Write the in-progress marker so the post-restart hook knows to + run health-check and signal the outcome.""" + import datetime + marker = { + "checkpoint_tag": checkpoint_tag, + "from_sha": from_sha, + "to_sha": to_sha, + "started_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + marker_path = state_dir / _UPDATE_IN_PROGRESS_MARKER + marker_path.parent.mkdir(parents=True, exist_ok=True) + marker_path.write_text(json.dumps(marker, indent=2)) + + +def read_update_marker(state_dir: Path) -> dict | None: + """Read the in-progress update marker, if it exists. Returns None + if no update is in progress.""" + marker_path = state_dir / _UPDATE_IN_PROGRESS_MARKER + if not marker_path.exists(): + return None + try: + return json.loads(marker_path.read_text()) + except (json.JSONDecodeError, OSError): + return None + + +def clear_update_marker(state_dir: Path) -> None: + """Delete the in-progress update marker — update complete (success + or rollback handled).""" + marker_path = state_dir / _UPDATE_IN_PROGRESS_MARKER + marker_path.unlink(missing_ok=True) + + +async def create_checkpoint() -> dict: + """Create a pre-update checkpoint via the deploy helper. + + Returns a dict with keys: + ok (bool), checkpoint_tag (str), git_sha (str), + output (str), exit_code (int). + """ + result = await _run_helper(["checkpoint"]) + tag = "" + if result["ok"]: + # The deploy helper prints the tag to stdout. + lines = result["output"].splitlines() + # Last non-log line is the tag. + for line in reversed(lines): + line = line.strip() + if line and not line.startswith("[taos-deploy]"): + tag = line + break + + # Get current SHA for the marker. + rc, sha = await _run_git(["rev-parse", "HEAD"]) + current_sha = sha.strip() if rc == 0 else "" + + return { + **result, + "checkpoint_tag": tag, + "git_sha": current_sha, + } + + +async def pull_update(target_ref: str) -> dict: + """Fetch and checkout the target ref (branch or tag). + + Args: + target_ref: The git ref to check out (e.g. 'origin/master'). + + Returns: dict with ok, output, exit_code. + """ + repo = _repo_dir() + branch = target_ref + + # Parse "origin/branch" into fetch + checkout. + if "/" in target_ref: + remote, remote_branch = target_ref.split("/", 1) + rc, _ = await _run_git( + ["fetch", "--quiet", remote, "--", remote_branch], + timeout=120, + ) + if rc != 0: + return { + "ok": False, + "output": f"git fetch {remote} {remote_branch} failed", + "exit_code": rc, + } + branch = target_ref + else: + # Plain branch name — fetch from origin, then check out + # origin/ so we get the remote's version, not a stale + # local tracking branch. No ``--`` separator: ``git fetch + # origin refspec`` updates origin/refspec; ``git fetch origin + # -- refspec`` may not in some git versions. + rc, _ = await _run_git( + ["fetch", "--quiet", "origin", target_ref], + timeout=120, + ) + if rc != 0: + return { + "ok": False, + "output": f"git fetch origin {target_ref} failed", + "exit_code": rc, + } + branch = f"origin/{target_ref}" + + rc, out = await _run_git(["checkout", "--quiet", branch]) + if rc != 0: + return {"ok": False, "output": out, "exit_code": rc} + + # Fast-forward the local tracking branch so it matches the remote + # we just checked out — avoids leaving the repo in detached HEAD + # after a direct origin/ checkout. + if "/" in target_ref: + local_branch = target_ref.split("/", 1)[1] + await _run_git( + ["branch", "-f", local_branch, branch], + timeout=30, + ) + elif target_ref and "/" not in target_ref: + await _run_git( + ["branch", "-f", target_ref, branch], + timeout=30, + ) + + rc, sha_out = await _run_git(["rev-parse", "HEAD"]) + return { + "ok": True, + "output": f"checked out {branch} ({sha_out.strip()[:8]})", + "exit_code": 0, + } + + +async def update_dependencies() -> dict: + """Install/update Python dependencies using the detected package manager. + + Detects uv vs pip and runs the appropriate install command. + """ + pkg = _detect_package_manager() + repo = _repo_dir() + + if pkg == "uv" and shutil.which("uv"): + proc = await asyncio.create_subprocess_exec( + "uv", "sync", "--frozen", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + cwd=str(repo), + ) + stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=300) + output = stdout.decode("utf-8", errors="replace") if stdout else "" + return { + "ok": proc.returncode == 0, + "output": output, + "exit_code": proc.returncode or 0, + "package_manager": "uv", + } + + # Default: pip + venv = _venv_dir() + pip = str(venv / "bin" / "pip") + if not os.path.isfile(pip): + return { + "ok": False, + "output": f"pip not found at {pip}", + "exit_code": -1, + "package_manager": "pip", + } + + proc = await asyncio.create_subprocess_exec( + pip, "install", "-q", "-e", f"{repo}[worker]", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=300) + output = stdout.decode("utf-8", errors="replace") if stdout else "" + return { + "ok": proc.returncode == 0, + "output": output, + "exit_code": proc.returncode or 0, + "package_manager": "pip", + } + + +async def run_migrations() -> dict: + """Execute any pending DB migrations. + + On the worker, migrations are typically handled by the main + application startup. This is a no-op placeholder — worker-side + DB schema changes are applied when the new code runs its lifespan + after restart. + """ + logger.info("self-update: migrations handled by post-restart startup") + return {"ok": True, "output": "migrations deferred to post-restart startup", "exit_code": 0} + + +async def restart_service() -> dict: + """Restart the worker service via the deploy helper. + + WARNING: This kills the current process. The return value is + best-effort — if the restart succeeds, we never see the result. + """ + return await _run_helper(["restart-self"]) + + +async def run_health_check() -> dict: + """Run the deploy-helper health-check command. + + Called post-restart to verify the worker is healthy. + """ + return await _run_helper(["health-check"]) + + +async def rollback_to_checkpoint(checkpoint_tag: str | None = None) -> dict: + """Restore the worker to the pre-update checkpoint. + + Args: + checkpoint_tag: Git tag from the checkpoint, or None to read + from the manifest file. + + Returns: dict with ok, output, exit_code. + """ + args = ["rollback"] + if checkpoint_tag: + args.append(checkpoint_tag) + return await _run_helper(args) + + +async def signal_update_outcome( + controller_url: str, + worker_name: str, + outcome: str, + from_version: str, + to_version: str, + failure_reason: str = "", + rollback_to: str = "", + signing_key: bytes | None = None, +) -> int: + """POST the update outcome to the controller. + + Sends ``POST /api/cluster/workers/{name}/update-outcome`` with + the outcome payload. HMAC-signed if a signing key is provided. + + Returns the HTTP status code, or 0 on connection failure. + """ + import httpx + from tinyagentos.worker.pairing import sign_request_headers + + path = f"/api/cluster/workers/{worker_name}/update-outcome" + body_data = { + "name": worker_name, + "outcome": outcome, + "from_version": from_version, + "to_version": to_version, + } + if failure_reason: + body_data["failure_reason"] = failure_reason + if rollback_to: + body_data["rollback_to"] = rollback_to + + body = json.dumps(body_data).encode() + headers = {"content-type": "application/json"} + if signing_key: + auth_headers = sign_request_headers( + signing_key, worker_name, "POST", path, body + ) + headers.update(auth_headers) + headers["content-type"] = "application/json" + + try: + async with httpx.AsyncClient(timeout=10) as client: + resp = await client.post( + f"{controller_url.rstrip('/')}{path}", + content=body, + headers=headers, + ) + logger.info( + "self-update: outcome=%s reported to controller (status=%d)", + outcome, resp.status_code, + ) + return resp.status_code + except Exception as exc: + logger.error("self-update: failed to signal outcome: %s", exc) + return 0 + + +async def run_full_update( + target_ref: str, + controller_url: str, + agent, # WorkerAgent instance (for signaling) + state_dir: Path, + graceful: bool = True, +) -> dict: + """Run the full worker self-update lifecycle. + + Sequence: + 1. Create pre-update checkpoint (git tag + manifest). + 2. Signal update-available to controller. + 3. Initiate self-drain (stop accepting new work). + 4. Wait for in-flight leases to complete (if graceful). + 5. Pull new code. + 6. Update dependencies. + 7. Run migrations. + 8. Write in-progress marker for post-restart hook. + 9. Restart the service (this kills us). + + Args: + target_ref: Git ref to update to (e.g. 'origin/master'). + controller_url: The controller's base URL. + agent: WorkerAgent instance for heartbeat/signaling. + state_dir: Worker state directory for the update marker. + graceful: If True, wait for in-flight work to drain. + + Returns a dict describing each phase result. The restart phase + is fire-and-forget — if it succeeds we never return. + """ + import datetime + + results: dict[str, dict] = {} + worker_name = agent.name + + # ── Phase 1: Checkpoint ─────────────────────────────────────── + logger.info("self-update: phase 1 — creating checkpoint") + cp = await create_checkpoint() + results["checkpoint"] = cp + if not cp["ok"]: + logger.error("self-update: checkpoint failed — aborting update") + return {"ok": False, "error": "checkpoint failed", "phases": results} + + checkpoint_tag = cp.get("checkpoint_tag", "") + from_sha = cp.get("git_sha", "") + + # ── Phase 2: Signal update-available ────────────────────────── + logger.info("self-update: phase 2 — signaling update-available") + status = await agent.report_update_available(reason=f"target={target_ref}") + results["signal_update_available"] = {"status_code": status} + if status not in (200, 0): + logger.warning( + "self-update: update-available signal returned %d — continuing anyway", + status, + ) + + # ── Phase 3: Initiate self-drain ────────────────────────────── + logger.info("self-update: phase 3 — initiating self-drain") + status = await agent.initiate_self_drain(reason=f"update to {target_ref}") + results["initiate_drain"] = {"status_code": status} + if status not in (200, 0): + logger.warning( + "self-update: self-drain signal returned %d — continuing anyway", + status, + ) + + # ── Phase 4: Wait for drain (if graceful) ───────────────────── + if graceful: + logger.info("self-update: phase 4 — waiting for in-flight work to drain") + drain_ok = await _wait_for_drain(agent, timeout=120) + results["drain_wait"] = {"ok": drain_ok} + if not drain_ok: + logger.warning( + "self-update: drain wait incomplete — proceeding with update " + "(leases will be released by monitor loop timeout)" + ) + else: + results["drain_wait"] = {"ok": True, "forced": True} + + # ── Phase 5: Pull new code ──────────────────────────────────── + logger.info("self-update: phase 5 — pulling %s", target_ref) + pull = await pull_update(target_ref) + results["pull"] = pull + if not pull["ok"]: + logger.error("self-update: pull failed — aborting update") + # Don't rollback yet — we haven't installed anything. + await agent.notify_drain_complete() + return {"ok": False, "error": "pull failed", "phases": results} + + # ── Phase 6: Update dependencies ────────────────────────────── + logger.info("self-update: phase 6 — updating dependencies") + deps = await update_dependencies() + results["dependencies"] = deps + if not deps["ok"]: + logger.error("self-update: dependency update failed — rolling back") + await rollback_to_checkpoint(checkpoint_tag) + await agent.notify_drain_complete() + return {"ok": False, "error": "dependency update failed", "phases": results} + + # ── Phase 7: Run migrations ─────────────────────────────────── + logger.info("self-update: phase 7 — running migrations") + migs = await run_migrations() + results["migrations"] = migs + + # ── Phase 8: Write update marker ────────────────────────────── + logger.info("self-update: phase 8 — writing update marker") + to_sha = "" + rc, sha_out = await _run_git(["rev-parse", "HEAD"]) + if rc == 0: + to_sha = sha_out.strip() + + _write_update_marker(state_dir, checkpoint_tag, from_sha, to_sha) + results["marker"] = { + "ok": True, + "checkpoint_tag": checkpoint_tag, + "from_sha": from_sha, + "to_sha": to_sha, + } + + # ── Phase 9: Restart ────────────────────────────────────────── + logger.info("self-update: phase 9 — restarting service") + restart = await restart_service() + results["restart"] = restart + # If we reach here, restart failed or returned synchronously. + # In normal operation, the process is killed by the restart. + if not restart["ok"]: + logger.error("self-update: restart failed — rolling back") + await rollback_to_checkpoint(checkpoint_tag) + clear_update_marker(state_dir) + await agent.notify_drain_complete() + return {"ok": False, "error": "restart failed", "phases": results} + + return {"ok": True, "phases": results} + + +async def _wait_for_drain(agent, timeout: float = 120) -> bool: + """Wait for in-flight work to complete before updating. + + Polls the heartbeat response; the controller stops routing new + work once the worker is in draining status. Each heartbeat + response now includes a ``drain_complete`` field (taOS #890 C3) + so the worker can detect when all leases are released and proceed + without waiting the full timeout. + + Returns True when the drain is confirmed complete by the + controller. Returns False if the controller was never reachable + during the wait window — the caller should decide whether to + proceed. + """ + logger.info("self-update: waiting up to %.0fs for drain", timeout) + import httpx + from tinyagentos.worker.pairing import sign_request_headers + import json as _json + + controller = agent.controller_url + name = agent.name + key = agent._signing_key + + elapsed = 0.0 + interval = 5.0 + saw_ok = False + while elapsed < timeout: + await asyncio.sleep(interval) + elapsed += interval + + # Send a heartbeat to keep the controller updated on our drain + # status — this also lets us detect if the controller is still + # reachable. + try: + status = await agent.heartbeat(status="draining", drain_reason="update") + if status == 200: + logger.debug( + "self-update: drain heartbeat ok (%.0fs elapsed)", elapsed + ) + saw_ok = True + else: + logger.warning( + "self-update: drain heartbeat returned %d", status + ) + except Exception: + logger.warning("self-update: drain heartbeat failed — continuing") + continue + + # Check the heartbeat response body for drain_complete. + # We re-post a lightweight request so we can read the response + # payload (heartbeat() only returns the status code). + try: + path = "/api/cluster/heartbeat" + payload = _json.dumps({ + "name": name, + "load": 0.0, + "status": "draining", + "drain_reason": "update", + }).encode() + headers = sign_request_headers(key, name, "POST", path, payload) if key else {} + headers["content-type"] = "application/json" + + async with httpx.AsyncClient(timeout=5) as client: + resp = await client.post( + f"{controller.rstrip('/')}{path}", + content=payload, + headers=headers, + ) + if resp.status_code == 200: + body = resp.json() + if body.get("drain_complete"): + logger.info( + "self-update: drain confirmed complete by controller " + "(%.0fs elapsed)", elapsed, + ) + return True + except Exception: + logger.debug("self-update: drain-complete check failed — continuing") + + logger.info( + "self-update: drain wait complete (%.0fs elapsed, controller_reachable=%s)", + elapsed, saw_ok, + ) + return saw_ok + + +async def post_update_startup( + controller_url: str, + agent, # WorkerAgent + state_dir: Path, +) -> dict | None: + """Run the post-restart health-check and outcome signaling. + + Called during worker startup when an update-in-progress marker + is found. Returns the outcome dict or None if no update was + in progress. + + On health-check failure, initiates a rollback. + """ + marker = read_update_marker(state_dir) + if marker is None: + return None + + logger.info( + "self-update: post-restart hook — update in progress: %s -> %s", + marker.get("from_sha", "?")[:8], + marker.get("to_sha", "?")[:8], + ) + + checkpoint_tag = marker.get("checkpoint_tag", "") + from_sha = marker.get("from_sha", "") + to_sha = marker.get("to_sha", "") + + # Wait for grace period to let backends stabilise. + logger.info( + "self-update: waiting %ds grace period for backends", + POST_RESTART_GRACE_PERIOD, + ) + await asyncio.sleep(POST_RESTART_GRACE_PERIOD) + + # ── Health check ────────────────────────────────────────────── + logger.info("self-update: running post-restart health check") + health = await run_health_check() + + if not health["ok"]: + logger.error( + "self-update: health check FAILED — rolling back to %s", + checkpoint_tag, + ) + # Attempt rollback + rollback_result = await rollback_to_checkpoint(checkpoint_tag) + + # Signal rollback outcome (best-effort — we may not reach the + # controller if networking is the problem). + await signal_update_outcome( + controller_url=controller_url, + worker_name=agent.name, + outcome="rollback", + from_version=from_sha, + to_version=to_sha, + failure_reason=f"health-check: {health.get('output', 'unknown')}", + rollback_to=from_sha, + signing_key=agent._signing_key, + ) + + clear_update_marker(state_dir) + return { + "ok": False, + "outcome": "rollback", + "rollback": rollback_result, + "health": health, + } + + # ── Health check passed ─────────────────────────────────────── + logger.info("self-update: health check PASSED") + clear_update_marker(state_dir) + + # Re-registering happens naturally via the agent's run loop — + # the agent calls register() after heartbeat 404s are resolved. + # We just need to signal the outcome. + status = await signal_update_outcome( + controller_url=controller_url, + worker_name=agent.name, + outcome="success", + from_version=from_sha, + to_version=to_sha, + signing_key=agent._signing_key, + ) + + logger.info( + "self-update: outcome=success reported (status=%d)", status + ) + return { + "ok": True, + "outcome": "success", + "from_sha": from_sha, + "to_sha": to_sha, + "signal_status": status, + } From 87d3a65f428a2f4f058b50bba464df23d0840d57 Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:02:51 +0200 Subject: [PATCH 2/9] fix(security): add HMAC worker-name check to update-outcome endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The HMAC gate only proves the caller is a paired worker, not which worker. Without this check, worker A could spoof an outcome report for worker B. Now matches the pattern used by incus-enroll and heartbeat endpoints. QA-Impacted: TestUpdateOutcomeEndpoint tests may need hmac_worker_name set on request.state — CI will catch. --- tinyagentos/routes/cluster.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index c6d02d492..fffa69caa 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -1603,6 +1603,16 @@ async def report_update_outcome(request: Request, name: str, body: UpdateOutcome except _HMACError: return JSONResponse({"error": "hmac verification failed"}, status_code=403) + # Verify the HMAC-authenticated worker matches the body name. + # Without this, worker A could spoof an outcome report for + # worker B (the HMAC only proves the caller IS a paired worker, + # not WHICH worker). + if getattr(request.state, "hmac_worker_name", None) != name: + return JSONResponse( + {"error": "Worker name in header does not match path"}, + status_code=403, + ) + cluster = request.app.state.cluster_manager worker = cluster.get_worker(name) if not worker: From 0a5efedc880f41591de6512af1686b63348461b9 Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Tue, 28 Jul 2026 07:00:13 +0200 Subject: [PATCH 3/9] fix(test): set hmac_worker_name in update-outcome test patches The route-level name cross-check (report_update_outcome line 1500) requires request.state.hmac_worker_name to match the path {name} parameter. The previous side_effect=lambda r: None left it unset, causing all 4 TestUpdateOutcomeEndpoint tests to return 403 instead of their expected status codes (200/404/400). --- tests/test_cluster.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 02b01efeb..39bfeb553 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -601,10 +601,12 @@ async def test_update_outcome_success(self, client, app): "to_version": "def5678abc", } - # Bypass HMAC for the test — we test HMAC separately + # Bypass HMAC for the test — we test HMAC separately. + # The side_effect must also set hmac_worker_name so the route-level + # name cross-check in report_update_outcome passes. with patch( "tinyagentos.routes.cluster.require_worker_hmac", - side_effect=lambda r: None, + side_effect=lambda r: setattr(r.state, "hmac_worker_name", "gpu-box"), ): resp = await client.post( "/api/cluster/workers/gpu-box/update-outcome", @@ -635,7 +637,7 @@ async def test_update_outcome_rollback(self, client, app): with patch( "tinyagentos.routes.cluster.require_worker_hmac", - side_effect=lambda r: None, + side_effect=lambda r: setattr(r.state, "hmac_worker_name", "gpu-box"), ): resp = await client.post( "/api/cluster/workers/gpu-box/update-outcome", @@ -659,7 +661,7 @@ async def test_update_outcome_worker_not_found(self, client): with patch( "tinyagentos.routes.cluster.require_worker_hmac", - side_effect=lambda r: None, + side_effect=lambda r: setattr(r.state, "hmac_worker_name", "nonexistent"), ): resp = await client.post( "/api/cluster/workers/nonexistent/update-outcome", @@ -682,7 +684,7 @@ async def test_update_outcome_unknown_outcome(self, client, app): with patch( "tinyagentos.routes.cluster.require_worker_hmac", - side_effect=lambda r: None, + side_effect=lambda r: setattr(r.state, "hmac_worker_name", "gpu-box"), ): resp = await client.post( "/api/cluster/workers/gpu-box/update-outcome", From 057ab5415c54f365fe185d16a8a4e24bb212101f Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:01:43 +0200 Subject: [PATCH 4/9] fix(security): add worker-name auth check and concurrency guard to deploy_backend 1. Worker-name authorization: deploy_backend now calls require_worker_hmac and verifies the authenticated worker matches the path name. Previously the endpoint had no auth at all - any caller could trigger deploys on any worker. Now only the paired worker can deploy on itself (CodeRabbit, PR #1910). 2. Concurrency guard: added per-worker asyncio.Lock (_worker_deploy_locks) to both deploy_backend and _do_single_worker_update. Prevents concurrent API calls from double-installing or double-draining the same worker. --- tinyagentos/routes/cluster.py | 61 +++++++++++++++++++++++++++++------ 1 file changed, 51 insertions(+), 10 deletions(-) diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index fffa69caa..17a74c3db 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -1027,6 +1027,12 @@ class WorkerRemoteRequest(BaseModel): ] +# Per-worker deploy locks -- prevents concurrent install/restart on the +# same worker (CodeRabbit finding on PR #1910). +import collections +_worker_deploy_locks: dict[str, asyncio.Lock] = collections.defaultdict(asyncio.Lock) + + @router.post("/api/cluster/workers/{name}/deploy") async def deploy_backend(request: Request, name: str, body: DeployRequest): """Trigger a backend install on a remote worker. @@ -1034,7 +1040,26 @@ async def deploy_backend(request: Request, name: str, body: DeployRequest): The controller proxies this to the worker's deploy endpoint. The worker runs taos-deploy-helper.sh via passwordless sudo. Only commands in the fixed allowlist are accepted. + + HMAC-gated: the worker signs the request with its pairing key, and + the authenticated worker name must match the URL path. This prevents + worker A from triggering a deploy on worker B (CodeRabbit, PR #1910). + + Per-worker asyncio.Lock prevents concurrent install/restart calls on + the same worker from double-installing. """ + # HMAC gate -- only paired workers may trigger deploys, and only on themselves. + try: + await require_worker_hmac(request) + except _HMACError as exc: + return exc.response + # Verify the HMAC-authenticated worker matches the path name. + if getattr(request.state, "hmac_worker_name", None) != name: + return JSONResponse( + {"error": "Worker name in header does not match path"}, + status_code=403, + ) + cluster = request.app.state.cluster_manager worker = cluster.get_worker(name) if not worker: @@ -1047,16 +1072,19 @@ async def deploy_backend(request: Request, name: str, body: DeployRequest): status_code=400, ) - import httpx - try: - async with httpx.AsyncClient(timeout=620) as client: - resp = await client.post( - f"{worker.url}/api/worker/deploy", - json={"command": body.command}, - ) - return resp.json() - except Exception as exc: - return JSONResponse({"error": str(exc)}, status_code=502) + # Serialise deploys per worker -- concurrent calls can double-install. + lock = _worker_deploy_locks[name] + async with lock: + import httpx + try: + async with httpx.AsyncClient(timeout=620) as client: + resp = await client.post( + f"{worker.url}/api/worker/deploy", + json={"command": body.command}, + ) + return resp.json() + except Exception as exc: + return JSONResponse({"error": str(exc)}, status_code=502) @router.post("/api/cluster/workers/{name}/remote") @@ -1349,9 +1377,22 @@ async def _do_single_worker_update(cluster, worker) -> dict: On success: ``{"success": True, "worker": ..., "status": "updating", ...}``. On failure: ``{"success": False, "worker": ..., "error": "..."}``. Never raises -- all exceptions are caught and converted into error dicts. + + Per-worker asyncio.Lock prevents concurrent update calls on the same + worker from double-draining or double-deploying (CodeRabbit, PR #1910). """ name = worker.name + # Serialise updates per worker -- concurrent calls can double-drain/deploy. + lock = _worker_deploy_locks[name] + async with lock: + return await _do_single_worker_update_locked(cluster, worker) + + +async def _do_single_worker_update_locked(cluster, worker) -> dict: + """Inner implementation of _do_single_worker_update (lock held).""" + name = worker.name + # Step 1: Begin draining (with exception isolation -- drain_worker # may raise from notification or background-task failures). try: From 3d9124c0c70554f245e1746c834b750f7575b01f Mon Sep 17 00:00:00 2001 From: Hogne <227774406+hognek@users.noreply.github.com> Date: Sun, 2 Aug 2026 19:08:27 +0200 Subject: [PATCH 5/9] fix(worker): address PR #1910 blockers + CodeRabbit 5 findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BLOCKER 1 — restart bricks worker: - Add _restart_service_by_name() helper that tries systemctl --user scope (via su -l $SUDO_USER) when system-level restart fails. - at(1) script now includes --user fallback commands. - systemd-run --scope tries --user scope when system scope fails. - Verify atd daemon is active before trusting at(1) (CodeRabbit). BLOCKER 2 — unstack from #1903: - Rebase onto clean origin/dev (ad7e4fea). CodeRabbit findings (Jul 31): - Subprocess timeouts: catch asyncio.TimeoutError in _run_git and update_dependencies, kill the process, return structured error. - HMAC error: return exc.response instead of flat 403 in report_update_outcome, matching worker_heartbeat pattern. - post_update_startup: schedule via asyncio.ensure_future in background task instead of awaiting inline, so the heartbeat loop is never blocked by the grace period + health-check delay. --- scripts/taos-deploy-helper.sh | 50 +++++++++++++++++++++++---- tinyagentos/routes/cluster.py | 4 +-- tinyagentos/worker/agent.py | 57 ++++++++++++++++++++++--------- tinyagentos/worker/self_update.py | 41 ++++++++++++++++++++-- 4 files changed, 123 insertions(+), 29 deletions(-) diff --git a/scripts/taos-deploy-helper.sh b/scripts/taos-deploy-helper.sh index 71ae22e38..ad1ab4653 100755 --- a/scripts/taos-deploy-helper.sh +++ b/scripts/taos-deploy-helper.sh @@ -306,19 +306,46 @@ MANIFEST # (via at(1) or systemd-run) survives the teardown because it runs outside # the service cgroup. +# Try restarting a worker service by name, first as a system service then +# as a user service (systemctl --user). The deploy helper runs via sudo +# (as root), so plain systemctl cannot reach the invoking user's --user +# session; we try both scopes to cover both system-level and user-level +# worker installations. +_restart_service_by_name() { + local svc_name="$1" + # System service (systemctl as root). + systemctl restart "$svc_name" 2>/dev/null && return 0 + # User service — use the original user if sudo preserved SUDO_USER, + # otherwise try a raw ``systemctl --user`` (works when the helper is + # called without sudo). + if [[ -n "${SUDO_USER:-}" ]]; then + su -l "$SUDO_USER" -c "systemctl --user restart '$svc_name'" 2>/dev/null && return 0 + fi + systemctl --user restart "$svc_name" 2>/dev/null && return 0 + return 1 +} + _detached_restart_worker() { - local restart_cmd="systemctl restart tinyagentos-worker.service 2>/dev/null || systemctl restart taos-worker.service 2>/dev/null || true" + # Build an at(1) script that restarts via both system and user scope, + # because atd runs as the original user and can reach --user services. + local restart_script + restart_script="systemctl restart tinyagentos-worker.service 2>/dev/null || systemctl restart taos-worker.service 2>/dev/null || systemctl --user restart tinyagentos-worker.service 2>/dev/null || systemctl --user restart taos-worker.service 2>/dev/null || true" # 1 — at(1): schedules the restart after the helper exits (cleanest). if command -v at >/dev/null 2>&1; then - if echo "$restart_cmd" | at now 2>/dev/null; then - log "worker restart scheduled via at(1)" - return 0 + # Only trust at(1) if the atd daemon is actually running — + # ``at now`` exits 0 even when atd is stopped (CodeRabbit, Jul 31). + if systemctl is-active --quiet atd 2>/dev/null || systemctl is-active --quiet atd.service 2>/dev/null; then + if echo "$restart_script" | at now 2>/dev/null; then + log "worker restart scheduled via at(1)" + return 0 + fi fi fi # 2 — systemd-run --scope: runs outside the service cgroup. if command -v systemd-run >/dev/null 2>&1; then + # Direct systemctl via systemd-run (system scope). if systemd-run --scope --no-block systemctl restart tinyagentos-worker.service 2>/dev/null; then log "worker restart dispatched via systemd-run" return 0 @@ -327,15 +354,24 @@ _detached_restart_worker() { log "worker restart dispatched via systemd-run (taos-worker)" return 0 fi + # Try --user scope via systemd-run. + if systemd-run --scope --no-block systemctl --user restart tinyagentos-worker.service 2>/dev/null; then + log "worker restart dispatched via systemd-run --user" + return 0 + fi + if systemd-run --scope --no-block systemctl --user restart taos-worker.service 2>/dev/null; then + log "worker restart dispatched via systemd-run --user (taos-worker)" + return 0 + fi fi # 3 — Last resort: --no-block may race with cgroup teardown, but the # helper returns immediately so it often wins. - if systemctl restart --no-block tinyagentos-worker.service 2>/dev/null; then - log "worker restart dispatched via systemctl --no-block (may race with cgroup teardown)" + if _restart_service_by_name "tinyagentos-worker.service"; then + log "worker restart dispatched via systemctl --no-block" return 0 fi - if systemctl restart --no-block taos-worker.service 2>/dev/null; then + if _restart_service_by_name "taos-worker.service"; then log "worker restart dispatched via systemctl --no-block (taos-worker)" return 0 fi diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index 17a74c3db..d583a0c2f 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -1641,8 +1641,8 @@ async def report_update_outcome(request: Request, name: str, body: UpdateOutcome """ try: await require_worker_hmac(request) - except _HMACError: - return JSONResponse({"error": "hmac verification failed"}, status_code=403) + except _HMACError as exc: + return exc.response # Verify the HMAC-authenticated worker matches the body name. # Without this, worker A could spoof an outcome report for diff --git a/tinyagentos/worker/agent.py b/tinyagentos/worker/agent.py index 5c6f1146b..22a3453a2 100644 --- a/tinyagentos/worker/agent.py +++ b/tinyagentos/worker/agent.py @@ -87,6 +87,37 @@ def _is_repair_rejection(resp) -> bool: return False +async def _run_post_update_hook( + controller_url: str, + agent, # WorkerAgent + state_dir: Path, +) -> None: + """Background wrapper around post_update_startup. + + Runs the post-restart health-check and outcome-signalling hook as + a fire-and-forget task so the worker heartbeat loop is never blocked + by the grace period + health-check delay (CodeRabbit, Jul 31). + """ + try: + from tinyagentos.worker.self_update import post_update_startup + + outcome = await post_update_startup( + controller_url=controller_url, + agent=agent, + state_dir=state_dir, + ) + if outcome is not None: + logger.info( + "self-update: post-restart outcome=%s", + outcome.get("outcome", "unknown"), + ) + except Exception: + logger.warning( + "post_update_startup hook failed — continuing", + exc_info=True, + ) + + class WorkerAgent: def __init__( self, @@ -882,23 +913,15 @@ async def run(self): # pre-restart checkpoint, run the health-check and # signal the outcome (success or rollback) to the # controller. - try: - from tinyagentos.worker.self_update import post_update_startup - outcome = await post_update_startup( - controller_url=self.controller_url, - agent=self, - state_dir=self._state_dir, - ) - if outcome is not None: - logger.info( - "self-update: post-restart outcome=%s", - outcome.get("outcome", "unknown"), - ) - except Exception: - logger.warning( - "post_update_startup hook failed — continuing", - exc_info=True, - ) + # + # Scheduled as a background task so the worker loop + # can register and send heartbeats immediately; the + # hook's grace period + health check would otherwise + # block the loop long enough for the controller to + # mark the worker offline (CodeRabbit, Jul 31). + asyncio.ensure_future(_run_post_update_hook( + self.controller_url, self, self._state_dir, + )) continue if result == _NEEDS_REPAIR: # Controller rejected our key -- enter needs-re-pair state. diff --git a/tinyagentos/worker/self_update.py b/tinyagentos/worker/self_update.py index 82954f872..4765a49eb 100644 --- a/tinyagentos/worker/self_update.py +++ b/tinyagentos/worker/self_update.py @@ -129,7 +129,16 @@ async def _run_git( stderr=asyncio.subprocess.STDOUT, cwd=str(repo), ) - stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout) + try: + stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout) + except asyncio.TimeoutError: + try: + proc.kill() + await proc.wait() + except ProcessLookupError: + pass + logger.error("self-update: git %s timed out after %.0fs", args[0], timeout) + return -1, f"git {args[0]} timed out after {timeout:.0f}s" return proc.returncode or 0, (stdout.decode("utf-8", errors="replace") if stdout else "") @@ -294,7 +303,20 @@ async def update_dependencies() -> dict: stderr=asyncio.subprocess.STDOUT, cwd=str(repo), ) - stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=300) + try: + stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=300) + except asyncio.TimeoutError: + try: + proc.kill() + await proc.wait() + except ProcessLookupError: + pass + return { + "ok": False, + "output": "uv sync timed out after 300s", + "exit_code": -1, + "package_manager": "uv", + } output = stdout.decode("utf-8", errors="replace") if stdout else "" return { "ok": proc.returncode == 0, @@ -319,7 +341,20 @@ async def update_dependencies() -> dict: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, ) - stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=300) + try: + stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=300) + except asyncio.TimeoutError: + try: + proc.kill() + await proc.wait() + except ProcessLookupError: + pass + return { + "ok": False, + "output": "pip install timed out after 300s", + "exit_code": -1, + "package_manager": "pip", + } output = stdout.decode("utf-8", errors="replace") if stdout else "" return { "ok": proc.returncode == 0, From 8184095c98a543106c096bee3f7e7c316d1db92e Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:10:37 +0200 Subject: [PATCH 6/9] docs: add changelog fragment for worker self-update (#1910) --- .github/workflows/ci.yml | 2 +- .github/workflows/deleted-symbols-gate.yml | 3 -- .github/workflows/secret-ignores-gate.yml | 35 ---------------- .github/workflows/security.yml | 5 +-- .github/workflows/store-wiring-gate.yml | 46 ---------------------- changelog.d/1910-worker-self-update.md | 3 ++ 6 files changed, 5 insertions(+), 89 deletions(-) delete mode 100644 .github/workflows/secret-ignores-gate.yml delete mode 100644 .github/workflows/store-wiring-gate.yml create mode 100644 changelog.d/1910-worker-self-update.md diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 16ce2f875..81b95d5dd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -184,7 +184,7 @@ jobs: - name: Set up Node uses: actions/setup-node@v7 with: - node-version: "22" + node-version: "20" cache: "npm" cache-dependency-path: desktop/package-lock.json diff --git a/.github/workflows/deleted-symbols-gate.yml b/.github/workflows/deleted-symbols-gate.yml index 8e5e1a8ea..15c746a9c 100644 --- a/.github/workflows/deleted-symbols-gate.yml +++ b/.github/workflows/deleted-symbols-gate.yml @@ -14,9 +14,6 @@ name: Deleted symbols gate on: pull_request: - # "edited" so a waiver trailer added by editing the PR body retriggers the - # gate (a re-run replays the stale event payload with the old body). - types: [opened, synchronize, reopened, edited] branches: [master, dev] jobs: diff --git a/.github/workflows/secret-ignores-gate.yml b/.github/workflows/secret-ignores-gate.yml deleted file mode 100644 index 20034fcf3..000000000 --- a/.github/workflows/secret-ignores-gate.yml +++ /dev/null @@ -1,35 +0,0 @@ -name: Secret-ignores gate - -# Verifies that the committed .gitignore still protects known secret-shaped -# paths (data/hub/identity.json, foo.key, creds.json, x.p8, ...) on every -# promotion target. A .gitignore is the kind of file a rebase conflict can -# quietly drop during a dev->master promotion while every test still passes and -# nothing builds red, so promotion is verified here, not assumed. -# -# Trigger scope is deliberate: -# - push to master/dev/release/* : a dropped pattern fails the branch it -# lands on (this is the post-promotion check from tsk-laezfg step 1). -# - pull_request to master/dev/release/* : a conflict-resolution loss fails -# BEFORE the merge, since the merge commit's .gitignore is what is checked. -# The run is pure stdlib (~1s), so it carries no shard-timeout risk. -# -# See scripts/check_secret_ignores.py for REQUIRED_PATTERNS and SECRET_PATHS. - -on: - push: - branches: [master, dev, release/*] - pull_request: - branches: [master, dev, release/*] - -jobs: - secret-ignores-gate: - runs-on: ubuntu-latest - permissions: - contents: read - steps: - - uses: actions/checkout@v7 - - uses: actions/setup-python@v7 - with: - python-version: "3.12" - - name: Assert secret-shaped paths are ignored - run: python scripts/check_secret_ignores.py diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml index 12fc87551..2e0576e06 100644 --- a/.github/workflows/security.yml +++ b/.github/workflows/security.yml @@ -41,9 +41,6 @@ jobs: - name: Install dependencies run: pip install -e ".[dev,proxy,worker]" - - name: Install uv - run: pip install uv - - name: Check for known vulnerabilities # CVE-2026-3219 affects pip itself with no fix released yet; revisit once # a patched pip is available on PyPI. @@ -52,4 +49,4 @@ jobs: run: | python -m pip install --upgrade "pip>=26.1" pip install pip-audit - python scripts/check_dependency_audit_ignores.py + pip-audit --ignore-vuln CVE-2026-3219 diff --git a/.github/workflows/store-wiring-gate.yml b/.github/workflows/store-wiring-gate.yml deleted file mode 100644 index 258191365..000000000 --- a/.github/workflows/store-wiring-gate.yml +++ /dev/null @@ -1,46 +0,0 @@ -name: Store wiring gate - -# Detects PRs that add a new BaseStore subclass without wiring it into -# tinyagentos/app.py. Routes reach stores ONLY via request.app.state, so a -# store that is never assigned to app.state is unreachable. -# -# Start with a NAME-LEVEL check (class name appears in the lifespan file). -# Only newly added classes are policed; pre-existing orphans are skipped so -# the check is mergeable on any branch. -# -# A "Store-Unwired-Intentionally: , " trailer in the PR body -# waives a named class and logs it, for stores genuinely constructed elsewhere -# (tests, CLI, workers). -# -# See scripts/check_store_wiring.py for the implementation. - -on: - pull_request: - # "edited" so a waiver trailer added by editing the PR body retriggers the - # gate (a re-run replays the stale event payload with the old body). - types: [opened, synchronize, reopened, edited] - branches: [master, dev] - -jobs: - store-wiring-gate: - runs-on: ubuntu-latest - permissions: - contents: read - env: - BASE_REF: ${{ github.base_ref }} - steps: - - uses: actions/checkout@v7 - with: - fetch-depth: 0 - - - uses: actions/setup-python@v7 - with: - python-version: "3.12" - - - name: Fetch base branch - run: git fetch origin "$BASE_REF" - - - name: Check for unwired BaseStore subclasses - env: - PR_BODY: ${{ github.event.pull_request.body }} - run: python scripts/check_store_wiring.py --base "origin/$BASE_REF" diff --git a/changelog.d/1910-worker-self-update.md b/changelog.d/1910-worker-self-update.md new file mode 100644 index 000000000..29a0cde4c --- /dev/null +++ b/changelog.d/1910-worker-self-update.md @@ -0,0 +1,3 @@ +### Added + +- Worker self-update: `scripts/taos-deploy-helper.sh` gains install, restart, health-check, and rollback subcommands, with a cgroup-safe detached restart (#1910). From 21e5eb69938531838a747b17d4fd7d53aa482541 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:11:53 +0200 Subject: [PATCH 7/9] =?UTF-8?q?docs:=20doc-gate=20=E2=80=94=20Docs-Reviewe?= =?UTF-8?q?d=20for=20worker=20self-update?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Docs-Reviewed: worker self-update adds internal worker subcommands (install/restart/health-check/rollback) and HMAC-authenticated deploy endpoints; no agent-facing API surface or contributor workflow change. Subcommand usage is documented in scripts/taos-deploy-helper.sh --help; worker/README.md covers build/install and the deploy runbook will document the ops subcommands. From 106664b66ce4a883b75b6f36ba4cc4bd9f9b8021 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:51:17 +0200 Subject: [PATCH 8/9] chore: restore dev workflow baseline (drop stale-branch workflow regressions) --- .github/workflows/ci.yml | 2 +- .github/workflows/deleted-symbols-gate.yml | 3 ++ .github/workflows/secret-ignores-gate.yml | 35 ++++++++++++++++ .github/workflows/security.yml | 5 ++- .github/workflows/store-wiring-gate.yml | 46 ++++++++++++++++++++++ 5 files changed, 89 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/secret-ignores-gate.yml create mode 100644 .github/workflows/store-wiring-gate.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 81b95d5dd..16ce2f875 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -184,7 +184,7 @@ jobs: - name: Set up Node uses: actions/setup-node@v7 with: - node-version: "20" + node-version: "22" cache: "npm" cache-dependency-path: desktop/package-lock.json diff --git a/.github/workflows/deleted-symbols-gate.yml b/.github/workflows/deleted-symbols-gate.yml index 15c746a9c..8e5e1a8ea 100644 --- a/.github/workflows/deleted-symbols-gate.yml +++ b/.github/workflows/deleted-symbols-gate.yml @@ -14,6 +14,9 @@ name: Deleted symbols gate on: pull_request: + # "edited" so a waiver trailer added by editing the PR body retriggers the + # gate (a re-run replays the stale event payload with the old body). + types: [opened, synchronize, reopened, edited] branches: [master, dev] jobs: diff --git a/.github/workflows/secret-ignores-gate.yml b/.github/workflows/secret-ignores-gate.yml new file mode 100644 index 000000000..20034fcf3 --- /dev/null +++ b/.github/workflows/secret-ignores-gate.yml @@ -0,0 +1,35 @@ +name: Secret-ignores gate + +# Verifies that the committed .gitignore still protects known secret-shaped +# paths (data/hub/identity.json, foo.key, creds.json, x.p8, ...) on every +# promotion target. A .gitignore is the kind of file a rebase conflict can +# quietly drop during a dev->master promotion while every test still passes and +# nothing builds red, so promotion is verified here, not assumed. +# +# Trigger scope is deliberate: +# - push to master/dev/release/* : a dropped pattern fails the branch it +# lands on (this is the post-promotion check from tsk-laezfg step 1). +# - pull_request to master/dev/release/* : a conflict-resolution loss fails +# BEFORE the merge, since the merge commit's .gitignore is what is checked. +# The run is pure stdlib (~1s), so it carries no shard-timeout risk. +# +# See scripts/check_secret_ignores.py for REQUIRED_PATTERNS and SECRET_PATHS. + +on: + push: + branches: [master, dev, release/*] + pull_request: + branches: [master, dev, release/*] + +jobs: + secret-ignores-gate: + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-python@v7 + with: + python-version: "3.12" + - name: Assert secret-shaped paths are ignored + run: python scripts/check_secret_ignores.py diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml index 2e0576e06..12fc87551 100644 --- a/.github/workflows/security.yml +++ b/.github/workflows/security.yml @@ -41,6 +41,9 @@ jobs: - name: Install dependencies run: pip install -e ".[dev,proxy,worker]" + - name: Install uv + run: pip install uv + - name: Check for known vulnerabilities # CVE-2026-3219 affects pip itself with no fix released yet; revisit once # a patched pip is available on PyPI. @@ -49,4 +52,4 @@ jobs: run: | python -m pip install --upgrade "pip>=26.1" pip install pip-audit - pip-audit --ignore-vuln CVE-2026-3219 + python scripts/check_dependency_audit_ignores.py diff --git a/.github/workflows/store-wiring-gate.yml b/.github/workflows/store-wiring-gate.yml new file mode 100644 index 000000000..258191365 --- /dev/null +++ b/.github/workflows/store-wiring-gate.yml @@ -0,0 +1,46 @@ +name: Store wiring gate + +# Detects PRs that add a new BaseStore subclass without wiring it into +# tinyagentos/app.py. Routes reach stores ONLY via request.app.state, so a +# store that is never assigned to app.state is unreachable. +# +# Start with a NAME-LEVEL check (class name appears in the lifespan file). +# Only newly added classes are policed; pre-existing orphans are skipped so +# the check is mergeable on any branch. +# +# A "Store-Unwired-Intentionally: , " trailer in the PR body +# waives a named class and logs it, for stores genuinely constructed elsewhere +# (tests, CLI, workers). +# +# See scripts/check_store_wiring.py for the implementation. + +on: + pull_request: + # "edited" so a waiver trailer added by editing the PR body retriggers the + # gate (a re-run replays the stale event payload with the old body). + types: [opened, synchronize, reopened, edited] + branches: [master, dev] + +jobs: + store-wiring-gate: + runs-on: ubuntu-latest + permissions: + contents: read + env: + BASE_REF: ${{ github.base_ref }} + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + + - uses: actions/setup-python@v7 + with: + python-version: "3.12" + + - name: Fetch base branch + run: git fetch origin "$BASE_REF" + + - name: Check for unwired BaseStore subclasses + env: + PR_BODY: ${{ github.event.pull_request.body }} + run: python scripts/check_store_wiring.py --base "origin/$BASE_REF" From ef12411bf9a6404888dc8be8ca2ad77dc135f7e2 Mon Sep 17 00:00:00 2001 From: hognek <227774406+hognek@users.noreply.github.com> Date: Sun, 30 Aug 2026 11:47:46 +0200 Subject: [PATCH 9/9] =?UTF-8?q?fix(worker):=20address=20PR=20#1910=20lead?= =?UTF-8?q?=20review=20=E2=80=94=203=20blockers=20+=20argument=20injection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BLOCKER 1 — drop the worker HMAC gate on POST /api/cluster/workers/{name}/deploy. It is a session-gated operator action; the HMAC+name cross-check locked out both the operator (session, no HMAC) and the worker (HMAC, no session). BLOCKER 2 — add /update-outcome to the auth middleware session-exempt list alongside /incus-enroll so the worker's HMAC-only caller (no session cookie) reaches the route-level HMAC gate instead of dying at AuthMiddleware. BLOCKER 3 — roll back on a stale update marker in post_update_startup. A healthy update clears the marker within ~1-2 min; a marker far older than STALE_UPDATE_MARKER_SECONDS means the new code never booted. Staleness alone triggers rollback, because the health check can only pass when the worker is already running. Security (non-blocking) — validate target_ref in pull_update against ^[A-Za-z0-9._/-]+$ and reject a leading '-'; keep the '--' separator in both fetch branches to block option/remote-protocol injection. --- tests/test_cluster.py | 95 +++++++++++++++++ tests/test_worker_self_update.py | 167 ++++++++++++++++++++++++++++++ tinyagentos/auth_middleware.py | 12 ++- tinyagentos/routes/cluster.py | 19 +--- tinyagentos/worker/self_update.py | 108 ++++++++++++++++++- 5 files changed, 378 insertions(+), 23 deletions(-) diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 39bfeb553..4cd608e6c 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -691,3 +691,98 @@ async def test_update_outcome_unknown_outcome(self, client, app): json=payload, ) assert resp.status_code == 400 + + +@pytest.mark.asyncio +class TestDeployEndpointOperatorAccess: + async def test_operator_can_trigger_deploy_without_hmac(self, client, app): + """BLOCKER 1 regression: deploy is a session-gated operator action. + + An operator (session cookie, no worker HMAC headers) must reach the + route. The old code layered ``require_worker_hmac`` plus a name + cross-check on top of the session gate, so the operator (session, no + HMAC) 401'd at the HMAC gate and the worker (HMAC, no session) 401'd + at the middleware — nobody could call it. + """ + from unittest.mock import patch + + mgr = app.state.cluster_manager + w = _make_worker("gpu-box") + await mgr.register_worker(w) + + fake_resp = MagicMock() + fake_resp.json.return_value = {"status": "deployed"} + + class _FakeClientCtx: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + async def post(self, *args, **kwargs): + return fake_resp + + with patch("httpx.AsyncClient", return_value=_FakeClientCtx()): + resp = await client.post( + "/api/cluster/workers/gpu-box/deploy", + json={"command": "status"}, + ) + + # 200 (proxied through to the worker), not 401/403 — the operator's + # session alone satisfies the gate and the request reached the proxy. + assert resp.status_code == 200, resp.text + assert resp.json() == {"status": "deployed"} + + +@pytest.mark.asyncio +class TestUpdateOutcomeRealCaller: + async def test_worker_can_report_outcome_without_session( + self, client, app, pair_and_register_worker, + ): + """BLOCKER 2 regression: the worker's real caller is HMAC + no cookie. + + ``signal_update_outcome()`` sends the three HMAC headers and no session + cookie. Before the fix, ``/update-outcome`` was not session-exempt, so + the request died at AuthMiddleware (401 ``Authentication required``) + before the route's own HMAC gate ever ran — outcomes fell into a black + hole. The route-level HMAC gate is now the only auth, matching + heartbeat / incus-enroll. + """ + import json as _json + + from tinyagentos.worker.pairing import sign_request_headers + + # Pair + register a worker; this stores its signing key on the controller. + await pair_and_register_worker( + client, app, + {"name": "gpu-box", "url": "http://localhost:9000"}, + ) + signing_key = await app.state.cluster_pairing.get_signing_key("gpu-box") + assert signing_key is not None + + payload = { + "name": "gpu-box", + "outcome": "success", + "from_version": "abc1234def", + "to_version": "def5678abc", + } + body = _json.dumps(payload).encode() + path = "/api/cluster/workers/gpu-box/update-outcome" + headers = sign_request_headers(signing_key, "gpu-box", "POST", path, body) + headers["content-type"] = "application/json" + + # A *separate* cookie-less client: ``client.post(cookies={})`` does NOT + # clear the fixture's cookie jar (httpx deprecates per-request cookies), + # so reusing the fixture would carry the admin session and prove nothing. + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://test", + ) as c: + assert not c.cookies # control: the worker holds no session + resp = await c.post(path, content=body, headers=headers) + + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["outcome"] == "success" + assert data["acknowledged"] is True diff --git a/tests/test_worker_self_update.py b/tests/test_worker_self_update.py index b38bac7dd..fdab13cdf 100644 --- a/tests/test_worker_self_update.py +++ b/tests/test_worker_self_update.py @@ -498,3 +498,170 @@ async def mock_signal(*args, **kwargs): assert result["outcome"] == "rollback" assert ("rollback", "tag-123") in call_log assert ("signal", "rollback") in call_log + + @pytest.mark.asyncio + async def test_stale_marker_triggers_rollback_even_when_health_check_passes( + self, monkeypatch, tmp_path: Path, + ): + """A stale marker means the new code never booted — roll back without + trusting a health check that can only pass when the process is running. + + Regression guard for BLOCKER 3: the health check asserting "service + active + port listening" is tautologically true whenever the hook runs, + so it can never catch a failed boot. The marker's staleness is the + signal that the update never completed. + """ + import datetime + import tinyagentos.worker.self_update as su + + state_dir = tmp_path / "worker-state" + state_dir.mkdir(parents=True, exist_ok=True) + + # Write a marker whose started_at is long in the past — the new code + # never came up, and the worker has only now recovered (e.g. after a + # systemd start-limit give-up and a later restart). + _write_update_marker(state_dir, "tag-123", "from-sha", "to-sha") + marker_path = state_dir / "update-in-progress.json" + marker = json.loads(marker_path.read_text()) + marker["started_at"] = ( + datetime.datetime.now(datetime.timezone.utc) + - datetime.timedelta(seconds=su.STALE_UPDATE_MARKER_SECONDS + 60) + ).isoformat() + marker_path.write_text(json.dumps(marker)) + + agent = MagicMock() + agent.name = "test-worker" + agent._signing_key = b"fake-key" + + call_log = [] + + async def mock_health_check(timeout=30): + # Health check PASSES — worker is active and listening. On the old + # code this would have reported "success" and swallowed the failure. + call_log.append("health-check") + return {"ok": True, "output": "healthy", "exit_code": 0} + + async def mock_rollback(checkpoint_tag=None): + call_log.append(("rollback", checkpoint_tag)) + return {"ok": True, "output": "rolled back", "exit_code": 0} + + async def mock_signal(*args, **kwargs): + call_log.append(("signal", kwargs.get("outcome"))) + return 200 + + monkeypatch.setattr(su, "run_health_check", mock_health_check) + monkeypatch.setattr(su, "rollback_to_checkpoint", mock_rollback) + monkeypatch.setattr(su, "signal_update_outcome", mock_signal) + monkeypatch.setattr(su, "POST_RESTART_GRACE_PERIOD", 0) + monkeypatch.setattr(su, "clear_update_marker", lambda d: None) + + result = await su.post_update_startup( + controller_url="http://controller:9898", + agent=agent, + state_dir=state_dir, + ) + + assert result is not None + assert result["ok"] is False + assert result["outcome"] == "rollback" + assert result.get("stale_marker") is True + assert ("rollback", "tag-123") in call_log + assert ("signal", "rollback") in call_log + # The health check was never consulted — staleness alone triggered it. + assert "health-check" not in call_log + + +# ── pull_update argument validation ─────────────────────────────────── + + +class TestPullUpdateValidation: + @pytest.mark.asyncio + async def test_rejects_option_like_ref(self, monkeypatch, tmp_path: Path): + """A ref that begins with '-' must be rejected before reaching git.""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + calls = [] + + async def mock_run_git(args, cwd=None, timeout=120): + calls.append(list(args)) + return 0, "" + + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.pull_update("--upload-pack=evil") + assert result["ok"] is False + assert "invalid target_ref" in result["output"] + assert calls == [] # git was never invoked + + @pytest.mark.asyncio + async def test_rejects_remote_protocol_injection(self, monkeypatch, tmp_path: Path): + """A remote of the form 'ext::sh -c …' must be rejected (command exec).""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + calls = [] + + async def mock_run_git(args, cwd=None, timeout=120): + calls.append(list(args)) + return 0, "" + + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.pull_update("ext::sh -c evil/branch") + assert result["ok"] is False + assert calls == [] + + @pytest.mark.asyncio + async def test_plain_branch_fetch_uses_double_dash(self, monkeypatch, tmp_path: Path): + """The no-slash branch must keep the '--' separator before the ref.""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + calls = [] + + async def mock_run_git(args, cwd=None, timeout=120): + calls.append(list(args)) + if args[0] == "fetch": + return 0, "" + if args[0] == "checkout": + return 0, "" + if args[0] == "branch": + return 0, "" + if args[0] == "rev-parse": + return 0, "def5678\n" + return 1, "unknown" + + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.pull_update("mybranch") + assert result["ok"] is True + fetch = [c for c in calls if c[0] == "fetch"][0] + assert fetch == ["fetch", "--quiet", "origin", "--", "mybranch"] + + @pytest.mark.asyncio + async def test_slash_branch_fetch_uses_double_dash(self, monkeypatch, tmp_path: Path): + """The slash branch must keep the '--' separator before the branch.""" + import tinyagentos.worker.self_update as su + + monkeypatch.setattr(su, "_repo_dir", lambda: tmp_path) + calls = [] + + async def mock_run_git(args, cwd=None, timeout=120): + calls.append(list(args)) + if args[0] == "fetch": + return 0, "" + if args[0] == "checkout": + return 0, "" + if args[0] == "branch": + return 0, "" + if args[0] == "rev-parse": + return 0, "def5678\n" + return 1, "unknown" + + monkeypatch.setattr(su, "_run_git", mock_run_git) + + result = await su.pull_update("origin/master") + assert result["ok"] is True + fetch = [c for c in calls if c[0] == "fetch"][0] + assert fetch == ["fetch", "--quiet", "origin", "--", "master"] diff --git a/tinyagentos/auth_middleware.py b/tinyagentos/auth_middleware.py index 6f9145897..22bf90ef7 100644 --- a/tinyagentos/auth_middleware.py +++ b/tinyagentos/auth_middleware.py @@ -412,6 +412,7 @@ def _is_exempt(method: str, path: str) -> bool: GET /api/cluster/workers — public worker list POST /api/cluster/workers — session-exempt, HMAC gate at route level POST /api/cluster/workers/{n}/incus-enroll — session-exempt, HMAC gate at route level + POST /api/cluster/workers/{n}/update-outcome — session-exempt, HMAC gate at route level POST /api/cluster/heartbeat — session-exempt, HMAC gate at route level """ if path in EXEMPT_PATHS or any(path.startswith(p) for p in EXEMPT_PREFIXES): @@ -439,12 +440,17 @@ def _is_exempt(method: str, path: str) -> bool: return True if method == "POST" and path == _CLUSTER_HEARTBEAT: return True - # POST /api/cluster/workers//incus-enroll — session-exempt; the route - # verifies the worker's HMAC signature (see tinyagentos.worker.enroll). + # POST /api/cluster/workers//incus-enroll and + # POST /api/cluster/workers//update-outcome — session-exempt; the + # route verifies the worker's HMAC signature itself (see + # tinyagentos.worker.enroll and the worker self-update outcome signal). + # A worker holds no session cookie, so the session gate would otherwise + # refuse these before the route's HMAC gate could run — the same pattern + # as /api/cluster/heartbeat. if ( method == "POST" and path.startswith(_CLUSTER_WORKERS + "/") - and path.endswith("/incus-enroll") + and path.endswith(("/incus-enroll", "/update-outcome")) ): return True # Project-invite redeem: POST /api/projects/invites/redeem is diff --git a/tinyagentos/routes/cluster.py b/tinyagentos/routes/cluster.py index d583a0c2f..0a897c7e3 100644 --- a/tinyagentos/routes/cluster.py +++ b/tinyagentos/routes/cluster.py @@ -1041,25 +1041,14 @@ async def deploy_backend(request: Request, name: str, body: DeployRequest): worker runs taos-deploy-helper.sh via passwordless sudo. Only commands in the fixed allowlist are accepted. - HMAC-gated: the worker signs the request with its pairing key, and - the authenticated worker name must match the URL path. This prevents - worker A from triggering a deploy on worker B (CodeRabbit, PR #1910). + Operator action: gated by the operator session in AuthMiddleware (this + route is not session-exempt). A worker holds no session cookie, so it is + already refused at the middleware; there is no HMAC gate here because the + only legitimate caller is the session-authenticated operator. Per-worker asyncio.Lock prevents concurrent install/restart calls on the same worker from double-installing. """ - # HMAC gate -- only paired workers may trigger deploys, and only on themselves. - try: - await require_worker_hmac(request) - except _HMACError as exc: - return exc.response - # Verify the HMAC-authenticated worker matches the path name. - if getattr(request.state, "hmac_worker_name", None) != name: - return JSONResponse( - {"error": "Worker name in header does not match path"}, - status_code=403, - ) - cluster = request.app.state.cluster_manager worker = cluster.get_worker(name) if not worker: diff --git a/tinyagentos/worker/self_update.py b/tinyagentos/worker/self_update.py index 4765a49eb..96c3d0e09 100644 --- a/tinyagentos/worker/self_update.py +++ b/tinyagentos/worker/self_update.py @@ -12,9 +12,11 @@ from __future__ import annotations import asyncio +import datetime import json import logging import os +import re import shutil from pathlib import Path @@ -30,10 +32,26 @@ # Gives backends (Ollama, llama.cpp, etc.) time to stabilise. POST_RESTART_GRACE_PERIOD = 15 +# If the post-restart hook has not cleared the update marker within this many +# seconds of the marker being written, the update is presumed to have failed +# to boot and the worker rolls back to the checkpoint on recovery. A healthy +# update clears the marker within ~1-2 minutes (grace period + health check); +# 5 minutes leaves ample headroom without leaving a dead worker un-rolled-back +# for long. This is the trigger that survives "new code does not start": the +# health check below can only pass when the worker process is already running, +# so a failed boot would otherwise leave the marker on disk forever. +STALE_UPDATE_MARKER_SECONDS = 300 + # Marker file written by the pre-restart phase so the post-restart # startup hook knows an update was in progress. _UPDATE_IN_PROGRESS_MARKER = "update-in-progress.json" +# A target ref may only contain characters that can appear in a git ref or +# remote name. ``:`` and whitespace are excluded, which blocks ``ext::…`` +# remote-protocol injection; a leading ``-`` is rejected separately (in +# pull_update) so git never parses the ref as an option. +_VALID_TARGET_REF = re.compile(r"^[A-Za-z0-9._/-]+$") + def _install_dir() -> Path: """Return the worker install directory (TAOS_INSTALL_DIR or default).""" @@ -188,6 +206,33 @@ def clear_update_marker(state_dir: Path) -> None: marker_path.unlink(missing_ok=True) +def _marker_age_seconds( + marker: dict, + now: datetime.datetime | None = None, +) -> float | None: + """Return the marker's age in seconds, or None if it has no parseable + ``started_at`` timestamp. + + Used to detect a stale marker: if the post-restart hook has not cleared + the marker within the expected window, the new code most likely never + came up (a failed boot) and the worker has only now recovered. + """ + started_at = marker.get("started_at", "") + if not started_at: + return None + try: + started = datetime.datetime.fromisoformat(started_at) + except (ValueError, TypeError): + return None + if started.tzinfo is None: + started = started.replace(tzinfo=datetime.timezone.utc) + if now is None: + now = datetime.datetime.now(datetime.timezone.utc) + if now.tzinfo is None: + now = now.replace(tzinfo=datetime.timezone.utc) + return (now - started).total_seconds() + + async def create_checkpoint() -> dict: """Create a pre-update checkpoint via the deploy helper. @@ -226,6 +271,20 @@ async def pull_update(target_ref: str) -> dict: Returns: dict with ok, output, exit_code. """ + # Validate the ref before it reaches git. A ref beginning with ``-`` + # would be parsed by git as an option (e.g. ``--upload-pack=…``), and a + # ``remote`` of the form ``ext::sh -c …`` in the slash branch would be + # straight command execution. The allowlist permits the characters a + # real ref/remote name can contain and rejects everything else, + # including ``:`` and whitespace. Not reachable today (nothing writes + # the update trigger yet), but the trigger writer lands next. + if not _VALID_TARGET_REF.match(target_ref) or target_ref.startswith("-"): + return { + "ok": False, + "output": f"invalid target_ref: {target_ref!r}", + "exit_code": 1, + } + repo = _repo_dir() branch = target_ref @@ -246,11 +305,11 @@ async def pull_update(target_ref: str) -> dict: else: # Plain branch name — fetch from origin, then check out # origin/ so we get the remote's version, not a stale - # local tracking branch. No ``--`` separator: ``git fetch - # origin refspec`` updates origin/refspec; ``git fetch origin - # -- refspec`` may not in some git versions. + # local tracking branch. ``--`` separates the ref from any + # option-looking tokens (the ref is validated above, but the + # separator is cheap defence in depth). rc, _ = await _run_git( - ["fetch", "--quiet", "origin", target_ref], + ["fetch", "--quiet", "origin", "--", target_ref], timeout=120, ) if rc != 0: @@ -698,7 +757,10 @@ async def post_update_startup( is found. Returns the outcome dict or None if no update was in progress. - On health-check failure, initiates a rollback. + On health-check failure, initiates a rollback. It also rolls back when + the marker is *stale* — i.e. the update was written long enough ago that + the new code clearly never came up — regardless of the health check, + because the health check can only pass when the worker is already running. """ marker = read_update_marker(state_dir) if marker is None: @@ -714,6 +776,42 @@ async def post_update_startup( from_sha = marker.get("from_sha", "") to_sha = marker.get("to_sha", "") + # ── Stale-marker guard ───────────────────────────────────────── + # A healthy update clears this marker within ~1-2 minutes of the restart + # (grace period + health check). If the marker is still here long after + # that, the new code did not come up in time — the classic "new build + # fails to boot" failure — and the worker has only now recovered (e.g. + # after a systemd start-limit give-up and a later restart). In that + # state the health check would pass trivially (the process is running + # now), so staleness alone must trigger the rollback. + marker_age = _marker_age_seconds(marker) + if marker_age is not None and marker_age > STALE_UPDATE_MARKER_SECONDS: + logger.error( + "self-update: update marker is stale (%.0fs old) — new code did " + "not come up; rolling back to %s", + marker_age, checkpoint_tag, + ) + rollback_result = await rollback_to_checkpoint(checkpoint_tag) + await signal_update_outcome( + controller_url=controller_url, + worker_name=agent.name, + outcome="rollback", + from_version=from_sha, + to_version=to_sha, + failure_reason=( + f"new code did not start (stale marker, {marker_age:.0f}s)" + ), + rollback_to=from_sha, + signing_key=agent._signing_key, + ) + clear_update_marker(state_dir) + return { + "ok": False, + "outcome": "rollback", + "rollback": rollback_result, + "stale_marker": True, + } + # Wait for grace period to let backends stabilise. logger.info( "self-update: waiting %ds grace period for backends",