feat(worker): self-update install, restart, health-check, and rollback (#890 C3) - #1910
feat(worker): self-update install, restart, health-check, and rollback (#890 C3)#1910hognek wants to merge 4 commits into
Conversation
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
📝 WalkthroughWalkthroughAdds worker self-update orchestration with checkpointing, dependency updates, draining, restart, health verification, rollback, and controller outcome reporting. Worker heartbeats expose drain completion, and update triggers run during heartbeat cycles. ChangesWorker self-update lifecycle
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant WorkerAgent
participant ClusterController
participant DeployHelper
WorkerAgent->>ClusterController: report update-available
ClusterController-->>WorkerAgent: return drain_complete heartbeat state
WorkerAgent->>DeployHelper: create checkpoint and restart worker
DeployHelper-->>WorkerAgent: return health-check result
WorkerAgent->>ClusterController: report success or rollback outcome
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
|
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| "output": f"git fetch origin {target_ref} failed", | ||
| "exit_code": rc, | ||
| } | ||
| branch = target_ref |
There was a problem hiding this comment.
WARNING: For a plain branch name (e.g. master), the new remote commits are fetched into origin/master but the worker then checks out the local master branch, which the fetch does NOT advance. The worker will therefore stay on the old commit and the update effectively no-ops. Also, for origin/<branch> refs the code checks out the detached remote-tracking ref rather than a local branch. Resolve by checking out FETCH_HEAD (or remote/branch) after fetch, e.g. git checkout --quiet FETCH_HEAD, or fast-forward the local branch with git checkout --quiet <branch> && git merge --ff-only <remote>/<branch>.
| logger.warning("self-update: drain heartbeat failed — continuing") | ||
|
|
||
| logger.info("self-update: drain wait complete (%.0fs elapsed)", elapsed) | ||
| return True |
There was a problem hiding this comment.
WARNING: _wait_for_drain unconditionally return True regardless of whether drain actually completed or heartbeats failed. As a result results["drain_wait"]["ok"] never reflects a real problem, and the if not drain_ok warning branch / the comment claiming "proceed anyway" is dead. The loop also never inspects the controller-reported drain/lease state (the args agent is only used to send heartbeats). Either detect completion (e.g. a heartbeat returning the worker's status/pending_leases or a controller endpoint) and return that, or drop the misleading ok semantics.
| import json | ||
| import logging | ||
| import os | ||
| import platform |
There was a problem hiding this comment.
SUGGESTION: import platform is unused (no other references in the module). Remove it, and also drop the unused from typing import Optional (line 19) and the unused import datetime (line 446) to keep imports honest.
| run health-check and signal the outcome.""" | ||
| marker = { | ||
| "checkpoint_tag": checkpoint_tag, | ||
| "from_sha": from_sha, |
There was a problem hiding this comment.
SUGGESTION: The marker writes "started_at": None with a comment "filled by caller with ISO timestamp", but no caller ever sets it —run_full_update calls _write_update_marker(...) without a timestamp. This is dead/misleading state. Either populate it (e.g. datetime.datetime.now().isoformat()) or remove the field and the comment.
| 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) |
There was a problem hiding this comment.
SUGGESTION: On outcome == "success" the endpoint does nothing beyond logging — the worker is expected to re-register via heartbeat. That's acceptable, but the worker's state on the controller is never explicitly reconciled/cleared here, and there is no record that the update succeeded (only informational logs + a rollback notification on failure). Consider recording the outcome on the worker object (or emitting a worker.update-success notification) so operators have an auditable trail and the controller can detect a stuck "update-available"/"draining" worker that never reports back.
| } | ||
|
|
||
| cmd_health_check() { | ||
| local timeout="${1:-10}" |
There was a problem hiding this comment.
SUGGESTION: cmd_health_check accepts a timeout argument (${1:-10}) but never uses it — no wait/retry/poll loop honors it. The caller run_health_check() passes POST_RESTART_HEALTH_TIMEOUT (30s), implying a timeout is expected, yet the shell check is instantaneous. Either implement a bounded retry/poll against the port/health endpoint using timeout, or remove the unused parameter to avoid implying a timeout is enforced.
Code Review SummaryStatus: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)SUGGESTION
Previously-reported findings — re-verified on this incrementAll findings from the prior review cycle (commit
Files Reviewed (incremental commit e6a751e)
Fix these issues in Kilo Cloud Previous Review Summaries (8 snapshots, latest commit f967846)Current summary above is authoritative. Previous snapshots are kept for context only. Previous review (commit f967846)Status: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)SUGGESTION
Previously-reported findings — re-verified on this incrementAll findings from the prior review cycle (commit
Files Reviewed (incremental commit e6a751e)
Fix these issues in Kilo Cloud Previous review (commit e6a751e)Status: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)SUGGESTION
Previously-reported findings — re-verified on this incrementAll findings from the prior review cycle (commit
Files Reviewed (incremental commit e6a751e)
Fix these issues in Kilo Cloud Previous review (commit 1ba39cf)Status: 2 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Previously-reported findings (re-verified)The two prior-cycle WARNINGs on
Files Reviewed (7 files, full PR diff vs merge-base)
Fix these issues in Kilo Cloud Previous review (commit b2433d3)Status: No new issues in incremental diff | 2 previous WARNINGs resolved | Recommendation: Merge incremental changes Overview
Resolved since last review (click to expand)WARNING (resolved)
Carried-forward suggestions (unchanged lines, not re-reported)
These remain open but are outside the lines changed in this incremental diff and were not duplicated. Files Reviewed (incremental diff, 2 files)
Previous review (commit a089e02)Status: 2 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Incremental changes reviewed since 6b583bc
Files Reviewed (3 files in incremental diff)
Fix these issues in Kilo Cloud Previous review (commit 6b583bc)Status: 1 Issue Found | Recommendation: Address before merge Overview
Issue Details (click to expand)SUGGESTION
Previously Reported Issues — Resolved in this update
Files Reviewed (3 files in incremental diff)
Fix these issues in Kilo Cloud Previous review (commit 04de66b)Status: 2 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)SUGGESTION
Previously Reported Issues — Resolved in this update
Files Reviewed (3 files in incremental diff)
Fix these issues in Kilo Cloud Previous review (commit b74d501)Status: 6 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)WARNING
SUGGESTION
Files Reviewed (8 files)
Reviewed by step-3.7-flash · Input: 162.9K · Output: 41K · Cached: 4.3M |
| log "worker service restarted" | ||
| } | ||
|
|
||
| cmd_health_check() { |
There was a problem hiding this comment.
SUGGESTION: Removing local timeout="${1:-10}" here leaves the run_health_check(timeout=...) caller's positional arg (_run_helper(["health-check", str(timeout)]) in self_update.py:341) silently ignored.
The timeout parameter is now dead: the shell function no longer reads $1, and the value is never used. The actual subprocess timeout is still enforced by _run_helper's default, so there is no functional break — but the Python-side contract that "the helper applies the health-check timeout" is now misleading. Either drop the timeout parameter from run_health_check, or make the shell actually honor it (e.g. bound the systemctl is-active/port checks with a deadline).
| "output": f"git fetch origin {target_ref} failed", | ||
| "exit_code": rc, | ||
| } | ||
| branch = f"origin/{target_ref}" |
There was a problem hiding this comment.
SUGGESTION: Checking out origin/<branch> fixes the stale-local-branch no-op, but it leaves the working tree in a detached HEAD state. The local <branch> ref is never fast-forwarded, so a subsequent plain-branch update will re-fetch but still check out the remote ref, and any later operation that assumes an attached branch (e.g. tag creation on the branch, future auto-pulls) may behave unexpectedly.
Consider updating the local branch ref instead, e.g. git fetch origin <branch>:<branch> (ff-only) followed by a non-detached checkout, or git branch -f <branch> origin/<branch> after checkout, so the update lands on the branch rather than a detached commit.
|
Two blockers plus the merge conflict (rebase onto post-#1903 dev — this is stacked on #1903 and duplicates its commit). BLOCKER 1 — restart bricks the worker. BLOCKER 2 — the orchestrator is dead code. Lower: |
Review-blocker resolution (6b583bc)BLOCKER 1 — restart bricks the worker ✓
BLOCKER 2 — orchestrator is dead code ✓
Fix — pull_update detached HEAD ✓
Fix — run_health_check unused timeout arg ✓
Rebase — onto current Tests: 68/68 pass ( |
| # after a direct origin/<ref> checkout. | ||
| if "/" in target_ref: | ||
| local_branch = target_ref.split("/", 1)[1] | ||
| await _run_git( |
There was a problem hiding this comment.
SUGGESTION: The git branch -f return code is ignored — await _run_git(...) discards both (rc, out). If the fast-forward fails (e.g. the local branch name collides with a tag or the remote-tracking ref isn't resolvable), the function still falls through to rev-parse HEAD and returns {"ok": True, ...}, silently reporting a successful checkout while the repo may remain in a detached/undesired state. Capture and log the rc (or fail the update) so a fast-forward failure isn't masked as success.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| path = "/api/cluster/heartbeat" | ||
| payload = _json.dumps({ | ||
| "name": name, | ||
| "load": 0.0, |
There was a problem hiding this comment.
WARNING: This drain-complete check fires a second, full heartbeat POST on every loop iteration, right after agent.heartbeat(...) already sent one. The first heartbeat's response body — which now carries drain_complete — is thrown away (because heartbeat() only returns the status code), so each 5s tick doubles heartbeat traffic to the controller and re-signs/re-sends a redundant request.
Additionally, this lightweight payload hardcodes "load": 0.0. The controller's ClusterManager.heartbeat() assigns worker.load = load unconditionally (no None guard, unlike backends/capabilities/vram), so every drain-window tick overwrites the worker's real load metric with 0.0, corrupting operator/routing views for the duration of the drain.
Prefer reading drain_complete from the response of the single heartbeat that is already being sent (e.g. have heartbeat() return the parsed body or expose the last response) instead of issuing this duplicate POST. If a separate request is unavoidable, omit load so the cached value is preserved.
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
…ranch -f - Replace redundant dual-POST heartbeat pattern in _wait_for_drain with a single POST that sends real CPU load (via psutil) and reads drain_complete from the response body. Eliminates heartbeat traffic doubling and load=0.0 clobber (Kilo WARNING). - Check git branch -f return code in pull_update and fail the update early when the fast-forward fails instead of silently masking the error (Kilo SUGGESTION). - Update test_pull_success mock to handle the branch git subcommand. Refs: jaylfc#1910
…ranch -f - Replace redundant dual-POST heartbeat pattern in _wait_for_drain with a single POST that sends real CPU load (via psutil) and reads drain_complete from the response body. Eliminates heartbeat traffic doubling and load=0.0 clobber (Kilo WARNING). - Check git branch -f return code in pull_update and fail the update early when the fast-forward fails instead of silently masking the error (Kilo SUGGESTION). - Update test_pull_success mock to handle the branch git subcommand. Refs: jaylfc#1910
b2433d3 to
1ba39cf
Compare
| # 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() |
There was a problem hiding this comment.
WARNING: The full self-update lifecycle runs inline in the single heartbeat loop, blocking heartbeat sends during long phases.
_check_update_trigger() awaits run_full_update(), whose phases run sequentially in this while self._running loop. The dependency-install phase (update_dependencies) can take well over the controller's HEARTBEAT_TIMEOUT = 30s (see tinyagentos/cluster/manager.py:15), and the drain wait is up to 120s. While these phases await subprocesses, the loop cannot send the next heartbeat, so the controller marks the worker offline mid-update (manager.py:767). A worker marked offline may be de-registered, and the post-restart hook depends on successful re-registration — so an eviction can break the whole self-update/rollback flow.
Consider moving the update into a detached asyncio.create_task(...) (fire-and-forget) instead of awaiting it inside the heartbeat loop, and/or continuing to send periodic heartbeats (status updating) while the update runs so the controller never sees a 30s gap.
Note also: the trigger is only checked after a successful heartbeat while _registered is True (line 886 is reached before the status == 404 -> _registered = False branch at line 887+). If the controller forgot the worker (404), the trigger will not fire until re-registration succeeds.
| worker = cluster._workers.get(body.name) | ||
| drain_complete = False | ||
| if worker is not None and worker.status == "draining" and hasattr(cluster, "_leases"): | ||
| active = [ |
There was a problem hiding this comment.
SUGGESTION: drain_complete counts every lease in cluster._leases, including expired ones not yet swept by _sweep_expired_leases.
The active-lease filter at lines 542-547 iterates cluster._leases.items() directly, but elsewhere the codebase treats expired leases as gone — e.g. get_leases() (manager.py:529-532) and find_existing_lease() (manager.py:432-438) both filter on lease.expires_at > now. If a lease has expired but has not yet been swept by the monitor loop, it is still counted here, so drain_complete stays False and the worker waits the full 120s _wait_for_drain timeout instead of proceeding early — defeating the purpose of the drain_complete optimization.
Filter by expiry to match the rest of the code:
| active = [ | |
| if worker is not None and worker.status == "draining" and hasattr(cluster, "_leases"): | |
| now = time.time() | |
| active = [ | |
| lid for lid, lease in cluster._leases.items() | |
| if getattr(lease, "expires_at", 0) > now | |
| and (parsed := cluster._parse_resource_id(lease.resource_id)) | |
| and parsed[0] == body.name | |
| ] | |
| drain_complete = len(active) == 0 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
scripts/taos-deploy-helper.sh (1)
348-348: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBrace the extras suffix to avoid array-subscript ambiguity (SC1087).
"$repo_dir[worker]"parses correctly in Bash today (variable expansion stops at[), but shellcheck flags it because[worker]reads like an array subscript. Use${repo_dir}to make the pip extras path unambiguous.♻️ Proposed tweak
- "$venv/bin/pip" install -q -e "$repo_dir[worker]" 2>/dev/null || \ + "$venv/bin/pip" install -q -e "${repo_dir}[worker]" 2>/dev/null || \🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/taos-deploy-helper.sh` at line 348, Update the pip install command around "$venv/bin/pip" to brace the repo_dir variable before appending the [worker] extras suffix, making the editable-install path unambiguous without changing its behavior.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tinyagentos/routes/cluster.py`:
- Around line 1491-1500: Update the handler containing require_worker_hmac and
the cluster.get_worker(name) lookup to compare request.state.hmac_worker_name
with the requested worker name before performing any lookup or outcome action.
Return the same 403 response used by worker_heartbeat for mismatched identities,
while preserving the existing HMAC failure and valid-worker behavior.
In `@tinyagentos/worker/self_update.py`:
- Around line 533-551: Reset the worker lifecycle status before returning from
both failure paths in the self-update flow: the `if not pull["ok"]` branch after
`pull_update(target_ref)` and the `if not deps["ok"]` branch after
`update_dependencies()`. Clear `_lifecycle_status` or send an `online` heartbeat
before `notify_drain_complete()` so the worker is routable again after aborting.
---
Nitpick comments:
In `@scripts/taos-deploy-helper.sh`:
- Line 348: Update the pip install command around "$venv/bin/pip" to brace the
repo_dir variable before appending the [worker] extras suffix, making the
editable-install path unambiguous without changing its behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: e71ddb3c-73fa-415a-bd54-b0f49c73ecfa
📒 Files selected for processing (7)
scripts/taos-deploy-helper.shtests/test_cluster.pytests/test_worker_self_update.pytinyagentos/routes/cluster.pytinyagentos/worker/agent.pytinyagentos/worker/deploy.pytinyagentos/worker/self_update.py
1ba39cf to
a089e02
Compare
…fc#1910) - agent.py: move self-update to detached asyncio.create_task so the heartbeat loop is never blocked during long update phases; add _update_in_progress guard to prevent duplicate launches (Kilo WARNING: heartbeat-loop blocking) - cluster.py (routes): filter drain_complete by lease expiry to match get_leases()/find_existing_lease() behavior, avoiding the full 120s timeout when expired leases haven't been swept yet (Kilo SUGGESTION: expired leases) - cluster.py (routes): add HMAC worker-name cross-check in report_update_outcome so a valid worker can't spoof another worker's outcome; matches existing pattern in worker_heartbeat (CodeRabbit Major: cross-worker spoofing guard) - self_update.py: reset _lifecycle_status to None and send a regular heartbeat instead of notify_drain_complete() in all three abort paths (pull failure, deps failure, restart failure) so the worker is routable again instead of stuck in 'updating' (CodeRabbit Major: lifecycle reset on abort) - deploy-helper.sh: brace $repo_dir variable to avoid SC1087 ambiguity (CodeRabbit Nitpick: shell variable bracing) Tests: 68/68 pass (test_worker_self_update, test_cluster, test_register_all_routers)
a089e02 to
e6a751e
Compare
| logger.error("self-update: pull failed — aborting update") | ||
| # Don't rollback yet — we haven't installed anything. | ||
| # Reset lifecycle so the worker is routable again instead of | ||
| # stuck in "updating" (CodeRabbit: cross-worker spoofing guard). |
There was a problem hiding this comment.
SUGGESTION: Misleading copy-paste comment.
This block resets the worker's lifecycle so it is routable again after an abort (the change referenced in the commit body as "lifecycle reset on abort"). The trailing (CodeRabbit: cross-worker spoofing guard) note is wrong here — the cross-worker spoofing guard is the HMAC name cross-check added in routes/cluster.py, unrelated to this abort path. Fix the comment so future maintainers are not misled about what this code does.
| # stuck in "updating" (CodeRabbit: cross-worker spoofing guard). | |
| # stuck in "updating" (CodeRabbit: lifecycle reset on abort). |
|
@hognek Rebase needed - this has conflicts against current dev (GitHub reports Everything else on it looks fine from my side; it is purely staleness. Rebase onto current dev and I will take another look. |
e6a751e to
90326d7
Compare
|
Flagging this from the periodic repo sweep: this PR is red on both required checks, and it is a real failure rather than infra.
Every one gets 403 where the test expects the endpoint to have been reached at all, including the two negative cases that expect 404 and 400. So the update-outcome endpoint is refusing the request before any of its own logic runs. That reads as an auth or permission gate applying to a route the tests exercise unauthenticated, rather than four separate bugs. Worth checking whether the tests need to authenticate now or the gate is wider than intended. The Not touching it since it is yours. |
|
Fixed the 4 failing Root cause: The tests patched Fix: Each test's All 42 cluster tests pass (including the 4 fixed). |
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tinyagentos/worker/agent.py (2)
733-743: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAborted update leaves the worker permanently
draining.
notify_drain_complete()sends one status-less heartbeat, butself._lifecycle_status/_lifecycle_reasonare still"draining"frominitiate_self_drain(). The run loop at Line 892 re-sendsstatus=self._lifecycle_statuson every subsequent tick, so after an aborted update (pull failed/dependency update failedinself_update.run_full_update) the controller keeps the worker out of routing forever. Clear the fields here.🔧 Proposed fix
logger.info("worker '%s': drain complete, ready for update", self.name) + self._lifecycle_status = None + self._lifecycle_reason = None return await self.heartbeat()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/agent.py` around lines 733 - 743, Clear the worker’s lifecycle state in notify_drain_complete() before sending the final heartbeat, resetting _lifecycle_status and _lifecycle_reason from the draining state so subsequent run-loop heartbeats no longer advertise the worker as draining after an aborted update.
117-119: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
WorkerUpdateServiceis undefined for the annotation (Ruff F821).The quoted annotation is runtime-safe, but Ruff flags the unresolved name and CI lint will fail. Add a
TYPE_CHECKINGimport.🔧 Proposed fix
+ # (add near the top of the module) + # if TYPE_CHECKING: + # from tinyagentos.worker.update_check import WorkerUpdateService self._update_service: "WorkerUpdateService | None" = None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/agent.py` around lines 117 - 119, Add a TYPE_CHECKING-only import for WorkerUpdateService so the annotation on Worker._update_service resolves for Ruff without introducing a runtime import; leave the existing lazy initialization and run() startup behavior unchanged.Source: Linters/SAST tools
🧹 Nitpick comments (3)
tests/test_worker_self_update.py (2)
438-439: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMock signature drifts from
run_health_check().
run_health_check()takes no parameters (self_update.pyLine 353), but these mocks declaretimeout=30. It passes today only because the default absorbs the missing argument, which means the tests won't catch future signature drift. Drop the parameter.♻️ Proposed tweak
- async def mock_health_check(timeout=30): + async def mock_health_check(): return {"ok": True, "output": "healthy", "exit_code": 0}Also applies to: 474-475
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_worker_self_update.py` around lines 438 - 439, Update the mock_health_check definitions in both referenced test locations to accept no parameters, matching the run_health_check signature while preserving their existing healthy response.
266-400: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo coverage for
_wait_for_drain(graceful=True) orsignal_update_outcome.Both tests use
graceful=Falseor abort before drain, so the drain-polling loop and the HMAC-signed outcome POST — the two most failure-prone paths in this module — are untested. Worth adding a case that drives_wait_for_drainwith a mocked controller returningdrain_complete: true.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_worker_self_update.py` around lines 266 - 400, The TestRunFullUpdate coverage should exercise the graceful drain path and signed update outcome reporting. Add a success test that calls run_full_update with graceful=True, mocks the controller interactions used by _wait_for_drain to return drain_complete=true, and verifies signal_update_outcome is invoked with the successful result while preserving the existing phase assertions.tinyagentos/worker/agent.py (1)
930-933: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
stop()reaches intoWorkerUpdateService._stop_event.Private-attribute access across a module boundary; the service already owns an async
stop(). Consider exposing a synchronousrequest_stop()onWorkerUpdateServiceand calling that.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/agent.py` around lines 930 - 933, Update Worker.stop to avoid directly accessing WorkerUpdateService._stop_event across the module boundary. Add a synchronous request_stop() method to WorkerUpdateService that signals its existing stop event, then call self._update_service.request_stop() from stop() while preserving the current None guard.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@scripts/taos-deploy-helper.sh`:
- Around line 320-330: Update the systemd-run dispatch branches in the worker
restart helper to use the default transient service mode instead of --scope,
while retaining asynchronous --no-block behavior. Ensure the dispatched
systemctl restart is independent of the worker service cgroup so it can outlive
the cgroup teardown, preserving the existing service fallback order and success
logging.
- Around line 272-281: Bound the suffix retry loop in the tag creation logic
within cmd_checkpoint so it makes only a finite number of attempts. Preserve the
existing counter-suffix behavior for tag collisions, but exit with a clear
failure when repeated git tag operations fail for other reasons instead of
allowing indefinite retries.
In `@tinyagentos/worker/agent.py`:
- Around line 851-873: The heartbeat loop must continue while the post-restart
hook and full update run. Update the startup hook around post_update_startup and
the _check_update_trigger flow so these long-running operations execute as
detached/background tasks, or otherwise emit an updating heartbeat throughout
their await periods, while preserving outcome logging and error handling.
In `@tinyagentos/worker/deploy.py`:
- Around line 33-37: Update the deploy helper contract between
rollback_to_checkpoint() and run_deploy() so checkpoint_tag values are handled
consistently: either allow and parse the tagged “rollback <tag>” form in
run_deploy() for controller-invoked rollbacks, or keep the allowlist restricted
to bare “rollback” and route explicit manual tags through scripts/rollback.sh.
Preserve bare self-update rollback behavior and ensure unsupported command forms
remain rejected.
In `@tinyagentos/worker/self_update.py`:
- Around line 243-255: The plain-branch update path must prevent git option
injection from target_ref. In tinyagentos/worker/self_update.py lines 243-255,
add an argument terminator before target_ref in git fetch and terminate the git
checkout arguments; in tinyagentos/worker/agent.py lines 768-793, validate
target_ref against a conservative ref pattern that rejects leading hyphens
before run_full_update, and delete the update trigger when validation fails.
- Around line 617-644: Update the self-update drain polling flow around
agent.heartbeat() to reuse the existing heartbeat response/body and read its
parsed drain_complete value, removing the second HTTP POST and duplicate request
construction. Ensure the heartbeat payload continues reporting the worker’s real
load rather than overwriting it with 0.0, while preserving the existing
successful drain-confirmation behavior.
- Around line 262-272: Consolidate the target_ref handling in the update flow
around _run_git: derive the branch name by removing the remote prefix when
target_ref contains "/", otherwise use target_ref directly, then perform one
branch -f operation. Capture and validate its return result, propagating failure
so pull_update does not return ok: True while HEAD remains detached.
---
Outside diff comments:
In `@tinyagentos/worker/agent.py`:
- Around line 733-743: Clear the worker’s lifecycle state in
notify_drain_complete() before sending the final heartbeat, resetting
_lifecycle_status and _lifecycle_reason from the draining state so subsequent
run-loop heartbeats no longer advertise the worker as draining after an aborted
update.
- Around line 117-119: Add a TYPE_CHECKING-only import for WorkerUpdateService
so the annotation on Worker._update_service resolves for Ruff without
introducing a runtime import; leave the existing lazy initialization and run()
startup behavior unchanged.
---
Nitpick comments:
In `@tests/test_worker_self_update.py`:
- Around line 438-439: Update the mock_health_check definitions in both
referenced test locations to accept no parameters, matching the run_health_check
signature while preserving their existing healthy response.
- Around line 266-400: The TestRunFullUpdate coverage should exercise the
graceful drain path and signed update outcome reporting. Add a success test that
calls run_full_update with graceful=True, mocks the controller interactions used
by _wait_for_drain to return drain_complete=true, and verifies
signal_update_outcome is invoked with the successful result while preserving the
existing phase assertions.
In `@tinyagentos/worker/agent.py`:
- Around line 930-933: Update Worker.stop to avoid directly accessing
WorkerUpdateService._stop_event across the module boundary. Add a synchronous
request_stop() method to WorkerUpdateService that signals its existing stop
event, then call self._update_service.request_stop() from stop() while
preserving the current None guard.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 010ad996-20ef-4d81-9c30-c5d244ab0bfd
📒 Files selected for processing (7)
scripts/taos-deploy-helper.shtests/test_cluster.pytests/test_worker_self_update.pytinyagentos/routes/cluster.pytinyagentos/worker/agent.pytinyagentos/worker/deploy.pytinyagentos/worker/self_update.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/test_cluster.py
- tinyagentos/routes/cluster.py
| # 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}" | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Unbounded retry loop can hang forever.
If git tag keeps failing for a reason other than "tag already exists" (disk full, index.lock present, permissions), the while loop spins indefinitely incrementing suffix, hanging cmd_checkpoint — and therefore the entire self-update flow that depends on it completing.
🔒 Bound the retry loop
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
+ while ! git -C "$repo_dir" tag "${tag}-${suffix}" HEAD 2>/dev/null; do
+ if (( suffix >= 20 )); then
+ die "failed to create checkpoint tag after ${suffix} attempts"
+ fi
((suffix++))
done
tag="${tag}-${suffix}"
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # 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}" | |
| } | |
| # 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 | |
| if (( suffix >= 20 )); then | |
| die "failed to create checkpoint tag after ${suffix} attempts" | |
| fi | |
| ((suffix++)) | |
| done | |
| tag="${tag}-${suffix}" | |
| } |
🧰 Tools
🪛 Shellcheck (0.11.0)
[warning] 273-273: Declare and assign separately to avoid masking return values.
(SC2155)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/taos-deploy-helper.sh` around lines 272 - 281, Bound the suffix retry
loop in the tag creation logic within cmd_checkpoint so it makes only a finite
number of attempts. Preserve the existing counter-suffix behavior for tag
collisions, but exit with a clear failure when repeated git tag operations fail
for other reasons instead of allowing indefinite retries.
| # 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 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does systemd-run --scope --no-block actually run in the background, or does --scope force synchronous execution regardless of --no-block?
💡 Result:
The short answer is yes, when you use systemd-run --scope, the command execution is inherently synchronous, meaning it will block the caller and return only when the command finishes, regardless of whether you include --no-block [1][2][3]. In the context of systemd-run, the --scope option fundamentally changes how the command is executed compared to the default transient service unit [1][2]. 1. Understanding --scope: When you run a command with --scope, the systemd-run process itself acts as the parent of the command [1][3]. The command inherits the execution environment of the caller and runs synchronously; systemd-run waits for the command to terminate [1][2][4]. This is different from the default service mode, where systemd acts as the parent and manages the process asynchronously in the background [1][2]. 2. The role of --no-block: The --no-block flag is designed to control whether systemd-run waits for the unit's start-up operation to be enqueued and completed [1][2]. While it is effective for default service units, it does not transform the synchronous nature of a scope unit into a background process [1][2]. 3. Why it behaves this way: The developers have noted that --scope is intended as a direct wrapper for the command execution within a transient control group [5]. Because the caller remains the parent process in this mode, it naturally remains attached to the lifecycle of the command [1][3]. If you need a process to run in the background while using --scope, you must handle the backgrounding yourself, for instance, by appending an ampersand (&) to your command string (e.g., systemd-run --scope my-command &) to place the entire systemd-run process in the background [5]. Top results for further reading: [1] https://man.archlinux.org/man/systemd-run.1.en [2] https://freedesktop.org/software/systemd/man/latest/systemd-run.html [5] https://lists.freedesktop.org/archives/systemd-devel/2018-February/040351.html [3] https://manpages.debian.org/trixie/systemd/systemd-run.1.en.html
Citations:
- 1: https://man.archlinux.org/man/systemd-run.1.en
- 2: https://freedesktop.org/software/systemd/man/latest/systemd-run.html
- 3: https://manpages.debian.org/trixie/systemd/systemd-run.1.en.html
- 4: https://docs.redhat.com/fr/documentation/red_hat_enterprise_linux/7/html/resource_management_guide/chap-using_control_groups
- 5: https://lists.freedesktop.org/archives/systemd-devel/2018-February/040351.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate file and relevant section =="
if [ -f scripts/taos-deploy-helper.sh ]; then
nl -ba scripts/taos-deploy-helper.sh | sed -n '280,355p'
else
echo "scripts/taos-deploy-helper.sh not found"
fi
echo
echo "== check systemd version availability if installed =="
if command -v systemd-run >/dev/null 2>&1; then
systemd-run --version | head -n 8 || true
else
echo "systemd-run not available in sandbox"
fi
echo
echo "== search for related docs/comments in repo =="
rg -n "systemd-run|--scope|--no-block|taos-worker|tinyagentos-worker" . --glob '!node_modules/**' --glob '!dist/**' --glob '!build/**' || trueRepository: jaylfc/taOS
Length of output: 228
🌐 Web query:
Red Hat systemd-run --scope synchronous no-block does not run in background systemd man page
💡 Result:
When you use systemd-run with the --scope option, the command is executed synchronously by design [1][2][3]. In this mode, systemd-run acts as the parent process and waits for the command to finish before it returns [1][4][3]. Because scope units are inherently tied to the lifecycle of the process that created them, they are not designed to be "backgrounded" in the way that transient service units (created without --scope) can be [1][4]. The --no-block option, which is intended to prevent systemd-run from waiting for the unit's initial start-up confirmation, does not override the fundamental synchronous nature of --scope [1][5][3]. While you may attempt to pass both, the manual page indicates that the synchronous execution behavior of --scope is a core aspect of how it manages the command as a child process of the caller [1][2][3]. If your goal is to run a command in the background, you should use the default transient service unit mode (by omitting --scope), which starts the service asynchronously by default [1][6][7].
Citations:
- 1: https://freedesktop.org/software/systemd/man/latest/systemd-run.html
- 2: https://manpages.debian.org/trixie/systemd/systemd-run.1.en.html
- 3: https://man.archlinux.org/man/systemd-run.1
- 4: https://deepwiki.com/systemd/systemd/9.2-systemd-run-and-transient-units
- 5: https://man7.org/linux/man-pages/man1/systemd-run.1.html
- 6: https://www.freedesktop.org/software/systemd/man/systemd-run
- 7: https://manpages.ubuntu.com/manpages/jammy/man1/systemd-run.1.html
Avoid relying on systemd-run --scope --no-block to stay prompt.
With --scope, the command is synchronous and --no-block does not make it backgrounded, so this dispatch path can still block the deploy script through the whole systemctl restart. Use the default transient service mode for async dispatch and ensure the restart can outlive the worker cgroup teardown.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/taos-deploy-helper.sh` around lines 320 - 330, Update the systemd-run
dispatch branches in the worker restart helper to use the default transient
service mode instead of --scope, while retaining asynchronous --no-block
behavior. Ensure the dispatched systemctl restart is independent of the worker
service cgroup so it can outlive the cgroup teardown, preserving the existing
service fallback order and success logging.
|
|
||
| # ── 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, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
The heartbeat loop stalls for the entire update / post-restart hook.
post_update_startup() blocks for POST_RESTART_GRACE_PERIOD (15s) plus the health check, and _check_update_trigger() awaits the whole run_full_update() (drain wait up to 120s, dependency install unbounded). No heartbeat is sent during either, so the controller's heartbeat timeout can mark the worker offline mid-update and de-register it, breaking the very re-registration the post-restart hook depends on. Run both as detached tasks, or keep an updating heartbeat ticking while they execute.
Also applies to: 919-925
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 869-869: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/agent.py` around lines 851 - 873, The heartbeat loop must
continue while the post-restart hook and full update run. Update the startup
hook around post_update_startup and the _check_update_trigger flow so these
long-running operations execute as detached/background tasks, or otherwise emit
an updating heartbeat throughout their await periods, while preserving outcome
logging and error handling.
| # Worker self-update subcommands (taOS #890 C3). | ||
| "checkpoint", | ||
| "rollback", | ||
| "restart-self", | ||
| "health-check", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether any caller passes a rollback tag through run_deploy().
rg -n 'run_deploy\(' -A2 tinyagentos | rg -n 'rollback'Repository: jaylfc/taOS
Length of output: 149
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate deploy.py and related helper files =="
fd -a 'deploy\.py$|self_update\.py$|.*deploy.*\.(py|sh)$' tinyagentos 2>/dev/null | sed 's#^\./##'
echo
echo "== tinyagentos/worker/deploy.py relevant lines =="
wc -l tinyagentos/worker/deploy.py
cat -n tinyagentos/worker/deploy.py | sed -n '1,120p'
echo
echo "== run_deploy call sites =="
rg -n 'run_deploy\(' tinyagentos -C 2 || true
echo
echo "== cmd_rollback definitions/usages =="
rg -n 'cmd_rollback|rollback|ALLOWED_COMMANDS' tinyagentos -C 2Repository: jaylfc/taOS
Length of output: 26964
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== helper rollback implementation =="
fd -a 'taos-deploy-helper|deploy.*helper|deploy-helper|deploy\.sh' /usr/local/bin /usr/local /tmp /home 2>/dev/null | sed 's#^\./##' || true
fd -a 'deploy-helper|deploy.*helper|deploy\.sh'| sed 's#^\./##' | head -50 || true
echo
echo "== local deployment scripts likely containing rollback argument parsing =="
for f in $(fd -a '.*deploy.*' tinyagentos tinyagentos | sed 's#^\./##'); do
if rg -q 'rollback|UPDATE_IN_PROGRESS|checkpoint|restart-self' "$f" 2>/dev/null; then
echo "--- $f"
wc -l "$f"
sed -n '1,240p' "$f"
fi
done
echo
echo "== read-only semantic probe of run_deploy allowlist and split behavior =="
python3 - <<'PY'
ALLOWED_COMMANDS = {
"install-ollama", "install-exo", "install-llama-cpp",
"install-llama-cpp --cuda", "install-vllm", "install-rknpu",
"update-worker", "status", "checkpoint", "rollback",
"restart-self", "health-check",
}
commands = ["rollback", "rollback abc", "rollback abc def", "install-llama-cpp --cuda"]
for command in commands:
print(command, "-- accepted:", command in ALLOWED_COMMANDS, "args:", command.split())
PYRepository: jaylfc/taOS
Length of output: 7645
Move rollback allowance and tag argument handling behind the deploy helper contract.
rollback_to_checkpoint() can pass checkpoint_tag, but run_deploy() only allows the exact string "rollback" and will reject "rollback <tag>". If run_deploy() is intended for controller-invoked rollbacks with a specific tag, add the tag argument path into the helper path; otherwise ensure the allowlist is stable for bare self-update rollbacks and explicit manual tags use the dedicated scripts/rollback.sh.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/deploy.py` around lines 33 - 37, Update the deploy helper
contract between rollback_to_checkpoint() and run_deploy() so checkpoint_tag
values are handled consistently: either allow and parse the tagged “rollback
<tag>” form in run_deploy() for controller-invoked rollbacks, or keep the
allowlist restricted to bare “rollback” and route explicit manual tags through
scripts/rollback.sh. Preserve bare self-update rollback behavior and ensure
unsupported command forms remain rejected.
| 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]) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Git-option injection via target_ref. A ref beginning with - flows from the update-trigger file into git argument positions that lack a -- separator, where git treats it as an option (--upload-pack=<cmd> is a command-execution vector).
tinyagentos/worker/self_update.py#L243-L255: add--beforetarget_refin the plain-branchgit fetch, and terminate thegit checkoutargument list.tinyagentos/worker/agent.py#L768-L793: validatetarget_refagainst a conservative ref pattern (and reject leading-) before invokingrun_full_update, deleting the trigger on rejection.
📍 Affects 2 files
tinyagentos/worker/self_update.py#L243-L255(this comment)tinyagentos/worker/agent.py#L768-L793
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/self_update.py` around lines 243 - 255, The plain-branch
update path must prevent git option injection from target_ref. In
tinyagentos/worker/self_update.py lines 243-255, add an argument terminator
before target_ref in git fetch and terminate the git checkout arguments; in
tinyagentos/worker/agent.py lines 768-793, validate target_ref against a
conservative ref pattern that rejects leading hyphens before run_full_update,
and delete the update trigger when validation fails.
| 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, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Redundant branch and still-ignored branch -f return code.
The elif condition's "/" not in target_ref is always true given the preceding if, so the two arms differ only in the local branch name. Both also discard the result: a failed fast-forward leaves a detached HEAD while pull_update returns ok: True.
♻️ Proposed fix
- 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,
- )
+ local_branch = target_ref.split("/", 1)[1] if "/" in target_ref else target_ref
+ if local_branch:
+ rc_b, out_b = await _run_git(["branch", "-f", local_branch, branch], timeout=30)
+ if rc_b != 0:
+ return {"ok": False, "output": f"branch -f {local_branch} failed: {out_b}", "exit_code": rc_b}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/self_update.py` around lines 262 - 272, Consolidate the
target_ref handling in the update flow around _run_git: derive the branch name
by removing the remote prefix when target_ref contains "/", otherwise use
target_ref directly, then perform one branch -f operation. Capture and validate
its return result, propagating failure so pull_update does not return ok: True
while HEAD remains detached.
| # 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 |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Second heartbeat POST per tick, and it zeroes the reported load.
agent.heartbeat() at Line 603 already sent a full heartbeat whose body carries drain_complete; this re-post doubles controller traffic every 5s and hardcodes "load": 0.0, which ClusterManager.heartbeat() assigns unconditionally — overwriting the worker's real load for the whole drain window. Have heartbeat() surface the parsed body (or the last response) and read drain_complete from it; at minimum omit load here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/self_update.py` around lines 617 - 644, Update the
self-update drain polling flow around agent.heartbeat() to reuse the existing
heartbeat response/body and read its parsed drain_complete value, removing the
second HTTP POST and duplicate request construction. Ensure the heartbeat
payload continues reporting the worker’s real load rather than overwriting it
with 0.0, while preserving the existing successful drain-confirmation behavior.
|
Confirming your fix resolved it, and that my earlier read was right rather than a flake: Not merging it on the sweep, and the reason is scope rather than anything wrong with it. This is +1815/-7 across 7 files implementing worker self-update, restart, health-check and rollback. No bot has commented on the current head, so there is no review covering this tree, and green CI is not a review. Self-update and rollback logic is close enough to "can take a worker down" that I want it read properly rather than merged by a 4-hourly automated pass. So it is blocked on review, not on CI. Nothing for you to do unless you want to flag anything specific you would like looked at first. |
43886e6 to
f967846
Compare
…ploy_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 jaylfc#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.
…orchestrator (jaylfc#890 C3) 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).
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.
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).
…ploy_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 jaylfc#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.
0d6e99b to
2bff5d0
Compare
There was a problem hiding this comment.
Actionable comments posted: 9
♻️ Duplicate comments (5)
scripts/taos-deploy-helper.sh (2)
272-281: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winBound the tag-collision retry loop.
git tagcan fail for reasons other than an existing tag, for example a staleindex.lock, a read-only filesystem, or a permission error. In that case thewhileloop never terminates andcmd_checkpointhangs, which blocks phase 1 ofrun_full_update. Add an attempt limit. The same diff also resolves the ShellCheck SC2155 warning on Line 273.🔒 Proposed fix
- local tag="taos-worker-pre-update-$(date -u +%Y%m%d-%H%M%S)" + local tag + 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 + if (( suffix >= 20 )); then + die "failed to create checkpoint tag after ${suffix} attempts" + fi ((suffix++)) done tag="${tag}-${suffix}" }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/taos-deploy-helper.sh` around lines 272 - 281, Update the tag creation logic in cmd_checkpoint to separate the tag variable assignment from command substitution to resolve ShellCheck SC2155, and bound the suffix retry loop with a finite attempt limit. Stop retrying when the limit is reached and propagate a clear failure instead of allowing run_full_update to hang on persistent git errors.Source: Linters/SAST tools
320-330: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
systemd-run --scope --no-blockdoes not detach the restart.With
--scope,systemd-runstays the parent of the command and waits for it to finish.--no-blockdoes not change that. This branch therefore runssystemctl restartinside the caller's cgroup, which is the exact failure this helper tries to avoid. Use the default transient service mode (drop--scope) so the restart runs undersystemdand survives the worker cgroup teardown.The PR objectives state that restart and rollback now use a single
systemctl restartoperation. The code still contains the three-tierat(1)/systemd-run --scope/systemctl --no-blockfallback chain. Confirm which behavior is intended.♻️ Proposed fix
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 + if systemd-run --no-block --unit="taos-worker-restart-$$" \ + --collect 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 + if systemd-run --no-block --unit="taos-worker-restart-alt-$$" \ + --collect systemctl restart taos-worker.service 2>/dev/null; then log "worker restart dispatched via systemd-run (taos-worker)" return 0 fi fi🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/taos-deploy-helper.sh` around lines 320 - 330, Update the systemd-run fallback in the worker restart flow to use transient service mode by removing --scope while retaining --no-block, and apply the same change to both tinyagentos-worker.service and taos-worker.service branches. Align the surrounding restart logic with the stated single systemctl restart approach, removing the obsolete three-tier at/systemd-run --scope/systemctl --no-block fallback behavior as appropriate.tinyagentos/worker/self_update.py (3)
262-272: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe
elifcondition is always true, and bothbranch -fresults are discarded.The
ifat Line 262 already handles every ref that contains/, so"/" not in target_refat Line 268 is redundant. The two arms differ only in the local branch name.Neither call checks the exit code. If the fast-forward fails,
pull_updatestill returns{"ok": True}while HEAD stays detached, and the next update repeats the detached checkout. This was raised on a previous commit.♻️ Proposed fix
- 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, - ) + local_branch = target_ref.split("/", 1)[1] if "/" in target_ref else target_ref + if local_branch: + rc_b, out_b = await _run_git(["branch", "-f", local_branch, branch], timeout=30) + if rc_b != 0: + return { + "ok": False, + "output": f"branch -f {local_branch} failed: {out_b}", + "exit_code": rc_b, + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/self_update.py` around lines 262 - 272, Update the target-branch handling in pull_update to use one branch-name selection path instead of the redundant elif, then check and propagate the result of the _run_git branch -f operation. Ensure a failed branch update prevents pull_update from returning success and preserves the existing failure behavior, avoiding a successful response while HEAD remains detached.
243-255: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
target_refstill reachesgitargument positions without a--terminator.A ref that begins with
-is parsed bygitas an option.git fetch origin --upload-pack=<cmd>executes a command. Two positions are affected:
- Line 244:
["fetch", "--quiet", "origin", target_ref].- Line 255:
["checkout", "--quiet", branch], wherebranch == target_refon theorigin/<branch>path.The comment on Lines 240-242 explains why
--is omitted forfetch. If that constraint holds, validatetarget_refagainst a conservative ref pattern at the entry point and reject a leading-. This was raised on a previous commit.🔒️ Proposed guard
+import re + +_SAFE_REF = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/-]*$") + async def pull_update(target_ref: str) -> dict: ... repo = _repo_dir() + if not _SAFE_REF.match(target_ref): + return { + "ok": False, + "output": f"refusing unsafe target_ref: {target_ref!r}", + "exit_code": -1, + } branch = target_ref🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/self_update.py` around lines 243 - 255, Validate target_ref at the self-update entry point before it reaches _run_git, rejecting refs that begin with "-" and enforcing a conservative valid-ref pattern. Apply this guard to both the fetch path using target_ref and the checkout path using origin/{target_ref}, while preserving the existing fetch argument structure and failure response conventions.
617-644: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftThe second heartbeat POST is still present and still reports
load: 0.0.
agent.heartbeat()at Line 603 already sent a heartbeat whose response body carriesdrain_complete. This block re-signs and re-sends a second heartbeat every 5s only to read that field. The payload hardcodes"load": 0.0, andClusterManager.heartbeat()assignsloadunconditionally, so the worker's real load is overwritten for the whole drain window. This was raised on a previous commit.Have
heartbeat()return the parsed body, then readdrain_completefrom it. If a separate request stays, omitload.Related:
elapsedadvances byintervalonly and ignores request latency, so the real wait exceedstimeout. Use a monotonic deadline instead.♻️ Deadline accounting
- elapsed = 0.0 interval = 5.0 + deadline = asyncio.get_running_loop().time() + timeout saw_ok = False - while elapsed < timeout: + while asyncio.get_running_loop().time() < deadline: await asyncio.sleep(interval) - elapsed += interval🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/self_update.py` around lines 617 - 644, Update the drain-wait logic in the self-update flow to use the response body returned by agent.heartbeat() and read drain_complete without issuing the duplicate heartbeat POST or overwriting load with 0.0. Replace elapsed-plus-interval accounting with a monotonic deadline based on the timeout, accounting for heartbeat/request latency while preserving the existing completion and timeout behavior.
🧹 Nitpick comments (7)
scripts/taos-deploy-helper.sh (1)
394-396: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBrace the variable so the extras suffix is unambiguous.
ShellCheck reports SC1087 on Line 394. The expansion
"$repo_dir[worker]"produces the intended pip target, but the syntax reads as an array index. Use"${repo_dir}[worker]"to make the intent explicit and clear the warning.♻️ Proposed fix
- "$venv/bin/pip" install -q -e "$repo_dir[worker]" 2>/dev/null || \ + "$venv/bin/pip" install -q -e "${repo_dir}[worker]" 2>/dev/null || \🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/taos-deploy-helper.sh` around lines 394 - 396, Update the pip install target in the deployment script to use braced expansion, changing the repo_dir reference in the worker extras expression to the unambiguous ${repo_dir}[worker] form while preserving the existing fallback chain.Source: Linters/SAST tools
tests/test_worker_self_update.py (2)
309-322: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePatch the sleep call on the module under test, not on
asyncio.
monkeypatch.setattr(asyncio, "sleep", ...)replaces the attribute on the stdlib module for the duration of the test. Every awaiting coroutine in the process observes the replacement, including pytest plugins and any background task. Patchsu.asyncio.sleepis the same object, so scope the patch to the awaited call instead.patch.objecton the module reference under test states the intent more precisely and keeps the blast radius small.♻️ Proposed change
- monkeypatch.setattr(asyncio, "sleep", short_sleep) + monkeypatch.setattr(su.asyncio, "sleep", short_sleep)Also applies to: 391-391
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_worker_self_update.py` around lines 309 - 322, Update the sleep monkeypatch in the test setup to target the module-under-test reference via patch.object on su.asyncio, rather than mutating the global asyncio module directly. Apply the same scoped change to the additional occurrence near the later test line, preserving short_sleep’s no-op behavior.
458-500: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winAdd coverage for
signal_update_outcomeand_wait_for_drain.The rollback test asserts only that
signal_update_outcomewas called withoutcome="rollback". The real function is stubbed in every test, so the HMAC signing of the outcome payload has no coverage in this file. That signing is the authentication boundary for the controller endpoint added by this PR. A signature-format regression would pass all current tests and fail only in production with HTTP 403.
_wait_for_drainis also untested.run_full_updateis exercised withgraceful=Falseat Line 329 and with the default at Line 357, wherepull_updatefails before the drain wait is reached. Add a test that drives_wait_for_drainto both the completed and the timeout outcome.Do you want me to draft these tests?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_worker_self_update.py` around lines 458 - 500, Add focused tests in tests/test_worker_self_update.py for signal_update_outcome that exercise the real implementation, verify the outcome payload is HMAC-signed with the agent signing key, and validate the request signature/format sent to the controller. Add coverage for _wait_for_drain covering both successful completion and timeout, using mocked drain/status behavior and a short timeout without stubbing the function under test. Keep existing rollback tests’ stubs unchanged.tinyagentos/routes/cluster.py (2)
1513-1520: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick win
body.nameis accepted but never validated or used.The handler authorises on the path
nameand logs the pathname.body.nameis never compared tonameand never read. A worker can post{"name": "other-worker", ...}and the request succeeds under its own identity. Either drop the field fromUpdateOutcomeBodyor reject a mismatch, so the payload cannot suggest an outcome for another worker in stored logs or future consumers.♻️ Proposed check
if getattr(request.state, "hmac_worker_name", None) != name: return JSONResponse( {"error": "Worker name in header does not match path"}, status_code=403, ) + if body.name != name: + return JSONResponse( + {"error": "Worker name in body does not match path"}, + status_code=400, + )Also applies to: 1541-1545
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/routes/cluster.py` around lines 1513 - 1520, Update UpdateOutcomeBody and the associated update-outcome handler so the payload cannot carry an unrelated worker identity: either remove the unused name field from UpdateOutcomeBody, or validate that body.name matches the path name and reject mismatches before processing or logging the outcome. Keep authorization and logging tied to the path name.
920-923: 🩺 Stability & Availability | 🔵 TrivialThe lock only serialises deploys inside one controller process.
_worker_deploy_locksis process-local. If the controller ever runs with more than one uvicorn worker process or more than one replica, two concurrent deploy requests can still reach the same worker. If a single-process deployment is guaranteed, state that in the comment. Otherwise, add an idempotency guard on the worker side (for example, reject/api/worker/deploywhile an install is already running).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/routes/cluster.py` around lines 920 - 923, Address the process-local scope of _worker_deploy_locks: either explicitly document and enforce that the controller runs as a single process, or add a worker-side idempotency guard in the /api/worker/deploy handling path that rejects deployments while an install is active. Ensure concurrent requests from separate controller processes cannot overlap on the same worker.tinyagentos/worker/self_update.py (2)
417-431: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a bounded retry for the outcome report.
This POST is the only signal the controller receives about the update result. One transient network error, or a controller restart during the worker restart window, drops the report permanently. On the rollback path, the operator then gets no
worker.update-rollbacknotification.
post_update_startupalso clears the marker right after this call, so nothing retries later. Add two or three attempts with backoff, and treat a non-200 status the same as a connection failure.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/self_update.py` around lines 417 - 431, Update the outcome-reporting POST in the self-update flow around post_update_startup to make two or three bounded attempts with backoff, retrying both request exceptions and non-200 responses. Return the successful status when received, and only log the failure and return 0 after all attempts are exhausted.
189-199: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winParse the checkpoint tag from a labelled line, not from the last output line.
_run_helpermerges stderr into stdout. Any trailing warning fromgitor from the helper becomes the last non-[taos-deploy]line, sotagtakes that value. The wrong tag then flows intorollback_to_checkpointand the rollback targets the wrong ref, or fails.Emit a labelled line from
taos-deploy-helper.sh(for exampleCHECKPOINT_TAG=<tag>) and match that prefix here. Also treat an emptytagas a checkpoint failure, because rollback without a tag depends on the manifest file being correct.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tinyagentos/worker/self_update.py` around lines 189 - 199, Update the checkpoint parsing in the self-update flow to extract the tag only from the labelled CHECKPOINT_TAG output emitted by taos-deploy-helper.sh, ignoring unrelated merged stderr or warning lines. Ensure the helper emits that label when creating a checkpoint, and treat a missing or empty parsed tag as checkpoint failure before passing it to rollback_to_checkpoint.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@scripts/taos-deploy-helper.sh`:
- Line 356: Update the checkpoint_tag extraction commands near checkpoint
manifest handling to pass manifest_file through Python’s command-line arguments
and read it via sys.argv, rather than interpolating it into the Python source.
Apply the same change to both affected lines while preserving the existing JSON
field lookup and fallback behavior.
- Around line 379-382: Update cmd_rollback to read the recorded git_branch from
the manifest and restore that branch at the checkpoint commit instead of
checking out checkpoint_tag directly. Ensure the repository ends on the recorded
branch, preserving the existing verification log and failure handling.
- Around line 437-453: Update the port detection logic in cmd_health_check to
avoid the Linux-only netstat -p invocation on macOS. Use a cross-platform
listening-port check, while preserving the existing ss/netstat fallback
structure and success/failure reporting so healthy macOS workers do not trigger
rollback.
- Around line 312-318: Update the at(1) branch in _detached_restart_worker to
verify that the atd daemon is active before treating a successfully queued job
as a valid restart. Only log success and return 0 when both atd is running and
the at command succeeds; otherwise continue to the existing fallback restart
paths.
In `@tinyagentos/routes/cluster.py`:
- Around line 1532-1535: Update the _HMACError handler in worker_heartbeat’s
surrounding request flow to return the exception’s structured response via
exc.response instead of constructing a generic 403 JSONResponse. Preserve
require_worker_hmac validation and ensure the worker can receive the original
status and error details.
In `@tinyagentos/worker/self_update.py`:
- Around line 116-133: Handle subprocess timeouts consistently across both
update sites: in tinyagentos/worker/self_update.py lines 116-133, update
_run_git to catch asyncio.TimeoutError around proc.communicate(), kill the
process, and return a non-zero status with a timeout message; in
tinyagentos/worker/self_update.py lines 291-322, apply the same timeout handling
to both wait_for calls in update_dependencies, killing each process and
returning {"ok": False, ...} so callers can continue rollback and completion
handling.
- Around line 517-531: Update both abort paths in the self-update flow around
pull failure and dependency failure to clear self_update_status after
notify_drain_complete(), allowing traffic to be re-routed. In the Phase 6
dependency failure path, restore the virtualenv/dependency environment before
calling rollback_to_checkpoint(checkpoint_tag), while preserving the existing
rollback and return behavior.
- Around line 682-691: Update the WorkerAgent.run flow so post_update_startup()
is scheduled as a background task rather than awaited inline, allowing the
worker loop to register and send heartbeats immediately. Preserve the existing
startup health-check behavior while ensuring the task is tracked or safely
handled during shutdown.
- Around line 693-720: The failed health-check path should signal the rollback
outcome and clear the update marker before invoking rollback_to_checkpoint(),
then terminate or delegate in a way that preserves the current post-rollback
process behavior. Add an attempt counter to update-in-progress.json, increment
it across restarts, and stop retrying after a fixed cap of failed post-restart
health checks.
---
Duplicate comments:
In `@scripts/taos-deploy-helper.sh`:
- Around line 272-281: Update the tag creation logic in cmd_checkpoint to
separate the tag variable assignment from command substitution to resolve
ShellCheck SC2155, and bound the suffix retry loop with a finite attempt limit.
Stop retrying when the limit is reached and propagate a clear failure instead of
allowing run_full_update to hang on persistent git errors.
- Around line 320-330: Update the systemd-run fallback in the worker restart
flow to use transient service mode by removing --scope while retaining
--no-block, and apply the same change to both tinyagentos-worker.service and
taos-worker.service branches. Align the surrounding restart logic with the
stated single systemctl restart approach, removing the obsolete three-tier
at/systemd-run --scope/systemctl --no-block fallback behavior as appropriate.
In `@tinyagentos/worker/self_update.py`:
- Around line 262-272: Update the target-branch handling in pull_update to use
one branch-name selection path instead of the redundant elif, then check and
propagate the result of the _run_git branch -f operation. Ensure a failed branch
update prevents pull_update from returning success and preserves the existing
failure behavior, avoiding a successful response while HEAD remains detached.
- Around line 243-255: Validate target_ref at the self-update entry point before
it reaches _run_git, rejecting refs that begin with "-" and enforcing a
conservative valid-ref pattern. Apply this guard to both the fetch path using
target_ref and the checkout path using origin/{target_ref}, while preserving the
existing fetch argument structure and failure response conventions.
- Around line 617-644: Update the drain-wait logic in the self-update flow to
use the response body returned by agent.heartbeat() and read drain_complete
without issuing the duplicate heartbeat POST or overwriting load with 0.0.
Replace elapsed-plus-interval accounting with a monotonic deadline based on the
timeout, accounting for heartbeat/request latency while preserving the existing
completion and timeout behavior.
---
Nitpick comments:
In `@scripts/taos-deploy-helper.sh`:
- Around line 394-396: Update the pip install target in the deployment script to
use braced expansion, changing the repo_dir reference in the worker extras
expression to the unambiguous ${repo_dir}[worker] form while preserving the
existing fallback chain.
In `@tests/test_worker_self_update.py`:
- Around line 309-322: Update the sleep monkeypatch in the test setup to target
the module-under-test reference via patch.object on su.asyncio, rather than
mutating the global asyncio module directly. Apply the same scoped change to the
additional occurrence near the later test line, preserving short_sleep’s no-op
behavior.
- Around line 458-500: Add focused tests in tests/test_worker_self_update.py for
signal_update_outcome that exercise the real implementation, verify the outcome
payload is HMAC-signed with the agent signing key, and validate the request
signature/format sent to the controller. Add coverage for _wait_for_drain
covering both successful completion and timeout, using mocked drain/status
behavior and a short timeout without stubbing the function under test. Keep
existing rollback tests’ stubs unchanged.
In `@tinyagentos/routes/cluster.py`:
- Around line 1513-1520: Update UpdateOutcomeBody and the associated
update-outcome handler so the payload cannot carry an unrelated worker identity:
either remove the unused name field from UpdateOutcomeBody, or validate that
body.name matches the path name and reject mismatches before processing or
logging the outcome. Keep authorization and logging tied to the path name.
- Around line 920-923: Address the process-local scope of _worker_deploy_locks:
either explicitly document and enforce that the controller runs as a single
process, or add a worker-side idempotency guard in the /api/worker/deploy
handling path that rejects deployments while an install is active. Ensure
concurrent requests from separate controller processes cannot overlap on the
same worker.
In `@tinyagentos/worker/self_update.py`:
- Around line 417-431: Update the outcome-reporting POST in the self-update flow
around post_update_startup to make two or three bounded attempts with backoff,
retrying both request exceptions and non-200 responses. Return the successful
status when received, and only log the failure and return 0 after all attempts
are exhausted.
- Around line 189-199: Update the checkpoint parsing in the self-update flow to
extract the tag only from the labelled CHECKPOINT_TAG output emitted by
taos-deploy-helper.sh, ignoring unrelated merged stderr or warning lines. Ensure
the helper emits that label when creating a checkpoint, and treat a missing or
empty parsed tag as checkpoint failure before passing it to
rollback_to_checkpoint.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: bef99d7e-ef83-4b44-86bd-c1d840ca46e6
📒 Files selected for processing (7)
scripts/taos-deploy-helper.shtests/test_cluster.pytests/test_worker_self_update.pytinyagentos/routes/cluster.pytinyagentos/worker/agent.pytinyagentos/worker/deploy.pytinyagentos/worker/self_update.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tinyagentos/worker/deploy.py
- tests/test_cluster.py
- tinyagentos/worker/agent.py
| # 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
The at(1) branch reports success even when atd is not running.
echo "$restart_cmd" | at now exits 0 when the job is queued. On many hosts the at binary is installed but the atd daemon is disabled or stopped. The job is then queued and never executed. _detached_restart_worker returns 0, restart_service() reports ok, and run_full_update returns {"ok": True} at phase 9 without a restart. The worker keeps running the old code from memory while the new code sits on disk, and the update marker is never processed, so post_update_startup() cannot detect or roll back the failure.
Check that atd is active before you accept this path.
🔒 Proposed fix
# 1 — at(1): schedules the restart after the helper exits (cleanest).
- if command -v at >/dev/null 2>&1; then
+ if command -v at >/dev/null 2>&1 && \
+ (systemctl is-active --quiet atd 2>/dev/null || systemctl is-active --quiet atd.service 2>/dev/null); then
if echo "$restart_cmd" | at now 2>/dev/null; then
log "worker restart scheduled via at(1)"
return 0
fi
fi📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # 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 | |
| # 1 — at(1): schedules the restart after the helper exits (cleanest). | |
| if command -v at >/dev/null 2>&1 && \ | |
| (systemctl is-active --quiet atd 2>/dev/null || systemctl is-active --quiet atd.service 2>/dev/null); then | |
| if echo "$restart_cmd" | at now 2>/dev/null; then | |
| log "worker restart scheduled via at(1)" | |
| return 0 | |
| fi | |
| fi |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/taos-deploy-helper.sh` around lines 312 - 318, Update the at(1)
branch in _detached_restart_worker to verify that the atd daemon is active
before treating a successfully queued job as a valid restart. Only log success
and return 0 when both atd is running and the at command succeeds; otherwise
continue to the existing fallback restart paths.
| 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)" |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Pass the manifest path as an argument instead of interpolating it into the Python source.
$manifest_file comes from TAOS_CHECKPOINT_MANIFEST. Both lines embed that value inside a single-quoted Python string literal. A path that contains a quote or a backslash breaks the literal, and the remainder is executed as Python code. Pass the path through sys.argv instead.
🔒 Proposed fix
- checkpoint_tag="$(python3 -c "import json; print(json.load(open('$manifest_file')).get('checkpoint_tag',''))" 2>/dev/null || true)"
+ checkpoint_tag="$(python3 -c "import json,sys; print(json.load(open(sys.argv[1])).get('checkpoint_tag',''))" "$manifest_file" 2>/dev/null || true)"- pkg_manager="$(python3 -c "import json; print(json.load(open('$manifest_file')).get('package_manager','pip'))" 2>/dev/null || echo "pip")"
+ pkg_manager="$(python3 -c "import json,sys; print(json.load(open(sys.argv[1])).get('package_manager','pip'))" "$manifest_file" 2>/dev/null || echo "pip")"Also applies to: 387-387
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/taos-deploy-helper.sh` at line 356, Update the checkpoint_tag
extraction commands near checkpoint manifest handling to pass manifest_file
through Python’s command-line arguments and read it via sys.argv, rather than
interpolating it into the Python source. Apply the same change to both affected
lines while preserving the existing JSON field lookup and fallback behavior.
| 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))" |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Rollback leaves the repository in detached HEAD, which breaks update-worker.
git checkout --quiet "$checkpoint_tag" checks out a tag, so HEAD becomes detached. cmd_update_worker (Line 217) then runs git pull --ff-based origin "$BRANCH" — git pull --ff-only origin "$BRANCH" — which fails in a detached HEAD state. Any later deploy through this helper therefore fails until an operator checks out a branch manually.
cmd_checkpoint already records git_branch in the manifest, but cmd_rollback never reads it. Restore the recorded branch to the checkpoint commit instead.
♻️ Proposed fix
+ local git_branch=""
+ if [[ -f "$manifest_file" ]]; then
+ git_branch="$(python3 -c "import json,sys; print(json.load(open(sys.argv[1])).get('git_branch',''))" "$manifest_file" 2>/dev/null || true)"
+ fi
+
git -C "$repo_dir" checkout --quiet "$checkpoint_tag" || {
die "git checkout $checkpoint_tag failed"
}
+ if [[ -n "$git_branch" && "$git_branch" != "detached" && "$git_branch" != "HEAD" ]]; then
+ git -C "$repo_dir" checkout --quiet -B "$git_branch" "$checkpoint_tag" \
+ || log "WARN: could not restore branch $git_branch — repo left detached at $checkpoint_tag"
+ fi
log "restored code to $checkpoint_tag ($(git -C "$repo_dir" rev-parse --short HEAD))"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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))" | |
| local git_branch="" | |
| if [[ -f "$manifest_file" ]]; then | |
| git_branch="$(python3 -c "import json,sys; print(json.load(open(sys.argv[1])).get('git_branch',''))" "$manifest_file" 2>/dev/null || true)" | |
| fi | |
| git -C "$repo_dir" checkout --quiet "$checkpoint_tag" || { | |
| die "git checkout $checkpoint_tag failed" | |
| } | |
| if [[ -n "$git_branch" && "$git_branch" != "detached" && "$git_branch" != "HEAD" ]]; then | |
| git -C "$repo_dir" checkout --quiet -B "$git_branch" "$checkpoint_tag" \ | |
| || log "WARN: could not restore branch $git_branch — repo left detached at $checkpoint_tag" | |
| fi | |
| log "restored code to $checkpoint_tag ($(git -C "$repo_dir" rev-parse --short HEAD))" |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/taos-deploy-helper.sh` around lines 379 - 382, Update cmd_rollback to
read the recorded git_branch from the manifest and restore that branch at the
checkpoint commit instead of checking out checkpoint_tag directly. Ensure the
repository ends on the recorded branch, preserving the existing verification log
and failure handling.
| # 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 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The port check fails on macOS and triggers a false rollback.
cmd_health_check supports macOS: Lines 426-434 fall back to launchctl. On macOS, ss is absent and BSD netstat does not accept -p in the Linux sense, so netstat -tlnp exits with a usage error. grep -q then finds no match, ok becomes false, and the command exits 1. post_update_startup() treats that as a failed health check and rolls back a healthy worker.
Detect the port with a tool that works on both platforms.
🔒 Proposed fix
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 lsof >/dev/null 2>&1; then
+ if lsof -nP -iTCP:"$port" -sTCP:LISTEN >/dev/null 2>&1; 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
+ if netstat -an 2>/dev/null | grep -qE "[.:]$port[[:space:]].*LISTEN"; then
log "health-check: port $port is listening"
else
ok=false
failures+=("port $port not listening")
fi
+ else
+ log "health-check: no tool available to check port $port — skipping"
fi📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # 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 | |
| # 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 lsof >/dev/null 2>&1; then | |
| if lsof -nP -iTCP:"$port" -sTCP:LISTEN >/dev/null 2>&1; 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 -an 2>/dev/null | grep -qE "[.:]$port[[:space:]].*LISTEN"; then | |
| log "health-check: port $port is listening" | |
| else | |
| ok=false | |
| failures+=("port $port not listening") | |
| fi | |
| else | |
| log "health-check: no tool available to check port $port — skipping" | |
| fi |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/taos-deploy-helper.sh` around lines 437 - 453, Update the port
detection logic in cmd_health_check to avoid the Linux-only netstat -p
invocation on macOS. Use a cross-platform listening-port check, while preserving
the existing ss/netstat fallback structure and success/failure reporting so
healthy macOS workers do not trigger rollback.
| try: | ||
| await require_worker_hmac(request) | ||
| except _HMACError: | ||
| return JSONResponse({"error": "hmac verification failed"}, status_code=403) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Return the HMAC error response instead of a flat 403.
require_worker_hmac raises _HMACError with a structured response and a 401 status for unpaired workers, stale timestamps, and bad signatures. This handler discards it and always returns 403 with a generic message. worker_heartbeat (Line 501) and deploy_backend (Line 944) return exc.response. Match them so the worker can distinguish a pairing problem from a signature problem.
🔧 Proposed fix
try:
await require_worker_hmac(request)
- except _HMACError:
- return JSONResponse({"error": "hmac verification failed"}, status_code=403)
+ except _HMACError as exc:
+ return exc.response📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| await require_worker_hmac(request) | |
| except _HMACError: | |
| return JSONResponse({"error": "hmac verification failed"}, status_code=403) | |
| try: | |
| await require_worker_hmac(request) | |
| except _HMACError as exc: | |
| return exc.response |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/routes/cluster.py` around lines 1532 - 1535, Update the
_HMACError handler in worker_heartbeat’s surrounding request flow to return the
exception’s structured response via exc.response instead of constructing a
generic 403 JSONResponse. Preserve require_worker_hmac validation and ensure the
worker can receive the original status and error details.
| 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 "") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Unhandled subprocess timeouts across the update phases. Three asyncio.wait_for(proc.communicate(), ...) calls have no except asyncio.TimeoutError. On timeout, wait_for cancels the wait but does not terminate the child, so git, uv, or pip keeps running while the exception escapes to run_full_update. run_full_update guards none of these calls, so the rollback and notify_drain_complete() steps are skipped and the worker stays unroutable. _run_helper already handles this case correctly; apply the same pattern, plus proc.kill().
tinyagentos/worker/self_update.py#L116-L133: wrap thewait_forin_run_git, kill the process, and return a non-zero exit code with a timeout message.tinyagentos/worker/self_update.py#L291-L322: wrap bothwait_forcalls inupdate_dependencies, kill the process, and return{"ok": False, ...}.
📍 Affects 1 file
tinyagentos/worker/self_update.py#L116-L133(this comment)tinyagentos/worker/self_update.py#L291-L322
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/self_update.py` around lines 116 - 133, Handle subprocess
timeouts consistently across both update sites: in
tinyagentos/worker/self_update.py lines 116-133, update _run_git to catch
asyncio.TimeoutError around proc.communicate(), kill the process, and return a
non-zero status with a timeout message; in tinyagentos/worker/self_update.py
lines 291-322, apply the same timeout handling to both wait_for calls in
update_dependencies, killing each process and returning {"ok": False, ...} so
callers can continue rollback and completion handling.
| 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} |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# 1) Does notify_drain_complete reset the worker lifecycle status?
fd -t f 'agent.py' -p tinyagentos/worker --exec ast-grep outline {} --items all
rg -nP -C 12 'def notify_drain_complete' tinyagentos/worker/agent.py
rg -nP -C 3 '_lifecycle_status\s*=' tinyagentos/worker/agent.py
# 2) What does the helper's rollback command restore?
fd -t f 'taos-deploy-helper.sh' --exec rg -nP -C 25 'rollback\)|^cmd_rollback|checkpoint\)'Repository: jaylfc/taOS
Length of output: 7001
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== agent lifecycle and heartbeat sections ==="
sed -n '280,355p;700,746p' tinyagentos/worker/agent.py | cat -n
echo "=== update_dependencies and rollback call sites ==="
sed -n '300,355p;370,410p;410,440p;517,531p' tinyagentos/worker/self_update.py | cat -n
echo "=== helper rollback complete section ==="
sed -n '347,410p' tinyagentos/worker/taos-deploy-helper.sh | cat -nRepository: jaylfc/taOS
Length of output: 12763
Reset the worker lifecycle status on abort, and restore the virtualenv before rollback.
notify_drain_complete() only sends a normal heartbeat; it does not reset the worker lifecycle status. Clear self_update_status or pass a status clear after the heartbeat in both abort paths so the controller can re-route traffic.
The Phase 6 rollback path runs update_dependencies() before rollback_to_checkpoint(checkpoint_tag). Restore the dependency environment as part of the rollback; otherwise the worker can run old application code with the partially updated dependencies.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/self_update.py` around lines 517 - 531, Update both abort
paths in the self-update flow around pull failure and dependency failure to
clear self_update_status after notify_drain_complete(), allowing traffic to be
re-routed. In the Phase 6 dependency failure path, restore the
virtualenv/dependency environment before calling
rollback_to_checkpoint(checkpoint_tag), while preserving the existing rollback
and return behavior.
| # 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() |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -nP -C 8 'post_update_startup' tinyagentos/worker/agent.pyRepository: jaylfc/taOS
Length of output: 1924
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## self_update.py sections"
sed -n '630,750p' tinyagentos/worker/self_update.py
echo
echo "## agent.py around WorkerAgent.run"
sed -n '780,885p' tinyagentos/worker/agent.py
echo
echo "## search heartbeat/register refresh tasks"
rg -n -C 4 'heartbeat|register|_refresh|_heartbeat|WorkerAgent\.run|run\(' tinyagentos/worker/agent.py tinyagentos/worker | head -n 220Repository: jaylfc/taOS
Length of output: 25222
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## self_update health-check helpers"
rg -n -C 6 'def run_health_check|_run_helper|POST_RESTART_GRACE_PERIOD|DEFAULT_.*TIMEOUT|timeout' tinyagentos/worker/self_update.py tinyagentos/worker | head -n 240
echo
echo "## exact self_update imports/constants"
sed -n '1,80p' tinyagentos/worker/self_update.pyRepository: jaylfc/taOS
Length of output: 18753
Schedule post_update_startup() outside the WorkerAgent.run() await path.
WorkerAgent.run() awaits post_update_startup() directly, and that coroutine sleeps for 15s before run_health_check() can add another default 600s wait. With an update-in-progress marker present, send the first post-restart heartbeat on a background task or use a short explicit timeout so the loop cannot miss controller registration/heartbeat deadlines.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/self_update.py` around lines 682 - 691, Update the
WorkerAgent.run flow so post_update_startup() is scheduled as a background task
rather than awaited inline, allowing the worker loop to register and send
heartbeats immediately. Preserve the existing startup health-check behavior
while ensuring the task is tracked or safely handled during shutdown.
| 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, | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -t f 'taos-deploy-helper.sh' --exec rg -nP -C 30 'rollback\)|cmd_rollback|systemctl (restart|start)|launchctl'Repository: jaylfc/taOS
Length of output: 11447
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Locate relevant helpers:\n'
fd -t f 'self_update.py|*deploy-helper*.sh|*.sh' | sed -n '1,120p'
printf '\nOutline self_update.py around rollback/signaling helper definitions:\n'
wc -l tinyagentos/worker/self_update.py
ast-grep outline tinyagentos/worker/self_update.py --match rollback_to_checkpoint --view expanded || true
ast-grep outline tinyagentos/worker/self_update.py --match signal_update_outcome --view expanded || true
printf '\nRelevant sections in tinyagentos/worker/self_update.py:\n'
sed -n '600,750p' tinyagentos/worker/self_update.py
printf '\nSearch rollback helper invocation in self_update.py:\n'
rg -n -C 6 'rollback_to_checkpoint|signal_update_outcome|clear_update_marker|POST_RESTART_GRACE_PERIOD' tinyagentos/worker/self_update.py
printf '\nInspect marker state and rollback helper invocation in git:\n'
rg -n -C 4 'class.*Update|marker|rollback|POST_RESTART|POST_RESTART_GRACE_PERIOD' tinyagentos || trueRepository: jaylfc/taOS
Length of output: 507
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Locate relevant files:\n'
fd -a 'self_update.py|taos-deploy-helper.sh|deploy.*(helper|-helper).*|install-worker.sh|worker.*\.sh' . 2>/dev/null | sed 's#^\./##' | sed -n '1,160p'
printf '\nGit-tracked candidates:\n'
git ls-files | rg '(^|/)self_update\.py$|(^|/)taos-deploy-helper\.sh$|(^|/)install-worker\.sh$|(^|/)worker.*\.sh$' || true
printf '\nRelevant sections in tinyagentos/worker/self_update.py:\n'
wc -l tinyagentos/worker/self_update.py
sed -n '1,120p' tinyagentos/worker/self_update.py
sed -n '600,750p' tinyagentos/worker/self_update.py
printf '\nSearch rollback helper invocation in self_update.py:\n'
rg -n -C 6 'rollback_to_checkpoint|signal_update_outcome|clear_update_marker|POST_RESTART_GRACE_PERIOD|update_marker' tinyagentos/worker/self_update.py || true
printf '\nMarker/rollback state search in tinyagentos:\n'
rg -n -C 4 'update[_A-Z0-9]*marker|marker|rollback|POST_RESTART|POST_RESTART_GRACE_PERIOD|attempts|attempt' tinyagentos || true
printf '\nRollback helper relevant section:\n'
sed -n '340,405p' scripts/taos-deploy-helper.sh 2>/dev/null || true
rg -n -C 8 'cmd_rollback|_detached_restart_worker|clear.*marker|marker|TAOS_CHECKPOINT|rollback-manifest|restore|checkout|git' scripts/taos-deploy-helper.sh 2>/dev/null || trueRepository: jaylfc/taOS
Length of output: 50368
Exit before rollback_to_checkpoint() and add a retry cap.
cmd_rollback() clears the worker after restoring the checkpoint, and rollback_to_checkpoint() is awaited before signal_update_outcome() / clear_update_marker(). This means the controller may not receive rollback and the marker may survive. Do the outcome signal and marker clear before the rollback; use an exit or helper that keeps the current post-rollback process alive. Store an attempt count in update-in-progress.json and stop after a fixed number of failed post-restart health checks.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tinyagentos/worker/self_update.py` around lines 693 - 720, The failed
health-check path should signal the rollback outcome and clear the update marker
before invoking rollback_to_checkpoint(), then terminate or delegate in a way
that preserves the current post-rollback process behavior. Add an attempt
counter to update-in-progress.json, increment it across restarts, and stop
retrying after a fixed cap of failed post-restart health checks.
2 participants
Summary
Implements the worker-side self-install, restart, re-register, and rollback flow for taOS Issue #890.
Builds on C2 (PR #1903, feat/worker-self-drain) drain protocol.
Changes
New module:
tinyagentos/worker/self_update.py(502 lines)post_update_startup()— reads in-progress marker, waits grace period, runs health-check, signals success or triggers rollbackExtended:
scripts/taos-deploy-helper.sh(+221 lines)checkpoint,rollback [<tag>],restart-self,health-check_stop_worker_service,_start_worker_serviceExtended:
tinyagentos/worker/deploy.py(+5 lines)ALLOWED_COMMANDSfor security allowlistNew endpoint:
routes/cluster.py(+82 lines)POST /api/cluster/workers/{name}/update-outcome— worker reports success or rollbackworker.update-rollbacknotification on failureTests (26 new, all pass)
tests/test_worker_self_update.py:tests/test_cluster.py:Depends on
Test results
All existing cluster + router tests also pass.
Bot review status (2026-07-18)
Kilo (
kilo-code-bot): ✅ All findings resolved across 5 review cycles.git branch -frc ignored (6b583bc)CodeRabbit: Rate-limited — no review produced.
Gitar: Working (pending).
Qodo: Paused for this user.
Summary by CodeRabbit
New Features
Bug Fixes
Tests