Two correctness defects from the deep-dive review.
B. register_agent left a stale in-memory "online" entry when the DB write
failed. The except block rolled back the session but kept the
ConnectedAgent in _agents, so the gateway still emitted auth_ok and
command routing targeted an agent with no backing AgentSession row
(online/offline desync). register_agent now returns Optional[str]: on DB
failure it removes its own in-memory entry (guarded on identity so a
concurrent reconnect isn't clobbered) and returns None. Both callers —
the WS gateway (on_auth) and the poll /connect endpoint — now emit
auth_fail / HTTP 500 instead of reporting a connected agent. A stale
'online' server row with no in-memory agent is self-corrected by the
heartbeat reaper.
C. send_command swallowed the final command-record DB update on the
timeout/result path: a failed commit left the row stuck in 'running'
with no log line, making it invisible. Now logs the exception (the row
is still left for a future sweep, but it's now diagnosable).
Adds a regression test that register_agent returns None and rolls back the
in-memory registration when the session write fails.
Full backend suite: 89 passed, 12 skipped.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Cluster of lower-severity agent-gateway defects from the deep-dive review.
agent_gateway.py:
- _check_auth_rate_limit was an unlocked read-modify-write of a module dict
under async_mode='threading' (on_auth + poll /connect race it), and the
dict grew without bound from one-off IPs. Now guarded by a threading.Lock
with a periodic sweep that drops aged-out buckets.
- on_stream crashed with AttributeError on a malformed/partial frame with a
null/non-string channel (precisely what the poll fallback exists to
tolerate). Guard the channel type and the container_id split; drop the
bad frame instead of killing the handler.
- on_credential_update_ack swallowed exceptions without db.session.rollback(),
leaving a poisoned SQLAlchemy session for the next user on the worker.
- Replaced all print() with module logging (info/warning/exception), so
auth/IP-block/rotation events are level-controllable and structured.
agent_registry.verify_agent_auth:
- Consume the replay nonce ONLY after the HMAC signature is verified.
Previously the nonce was recorded first, so a forged request (valid
agent_id+prefix+nonce, bogus signature) burned the nonce and locked the
real agent out when it next presented it. Reordered; added a regression
test (test_failed_signature_does_not_consume_nonce).
- Fixed the misleading "5 minute window" comment (the check is 60s) and
documented that the 60s window is the only replay guard for the nonce-less
HTTP-poll transport.
NOT changed (deliberately, to avoid breaking agents):
- Nonce is left optional. The HTTP-poll transport signs agent_id:timestamp
with no nonce (auth.Sign, "nonce":""); requiring it would lock out every
poll-fallback agent. Making it mandatory needs a coordinated Go-agent
change and is tracked in SECURITY.md.
- Per-message session-token validation remains a documented TODO; it
requires a protocol change on both the Go agent and backend.
Full backend suite: 84 passed, 12 skipped.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
When an agent reconnected, register_agent overwrote the in-memory
ConnectedAgent for the server but never disconnected the old socket or
drained its pending commands. Callers blocked in send_command() on the old
agent's result_queue were never signalled — they waited out the full
per-command timeout and returned a misleading "Command timeout" even though
the agent was back online (the late result would resolve against the new
agent and be unreachable). The agent was also briefly left holding two live
sockets.
On replacement, register_agent now (after releasing the lock to avoid
deadlocking with the disconnect handler):
- puts an AGENT_RECONNECTED error on each orphaned pending command's
result_queue so blocked callers return immediately, and
- disconnects the old WS socket (guarded: only for ws transport, only when
the socket id actually changed, only when socketio is initialized).
The old socket id is already removed from _socket_to_server before the
disconnect, so the resulting disconnect handler is a no-op for the freshly
registered agent. Adds two regression tests.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two defects in the agent control plane found during the deep-dive review.
1. verify_agent_auth() failed OPEN on a missing secret (critical). The HMAC
check was wrapped in `if api_secret:`, so when get_api_secret() returned
None (NULL ciphertext, or a decryption error it swallows) the signature
block was skipped and execution fell through to "authenticated" — reducing
agent auth to agent_id + api_key_prefix, both non-secret observable values.
Now mandatory: no decryptable secret => log + track_auth_attempt(False) =>
return None. Also added Server.get_pending_api_secret() and select the
secret matching whichever prefix matched, so signatures actually verify
during a key-rotation window instead of hitting the old bypass.
2. Heartbeat reaper evicted freshly-reconnected agents (TOCTOU race).
_check_heartbeats snapshotted dead server_ids under the lock then released
it; _handle_agent_timeout popped unconditionally. An agent that reconnected
in that window (new ConnectedAgent, new socket_id) got clobbered and flipped
offline right after reconnecting. The snapshot now carries socket_id, and
_handle_agent_timeout re-validates under the lock (same socket AND still
stale) before evicting. Extracted HEARTBEAT_TIMEOUT so the scan and the
re-check can't drift; added the missing rollback on error.
Adds tests/test_agent_registry_security.py (5 tests). Full suite: 80 passed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
