diff --git a/gateway/gateway.py b/gateway/gateway.py index d3ebf62227..79c468fb36 100644 --- a/gateway/gateway.py +++ b/gateway/gateway.py @@ -5543,6 +5543,33 @@ def session_delete_by_container(container_id: str) -> tuple[Response, int] | Res return make_success("Session deleted") +@app.route("/api/v1/sessions/by-container//heartbeat", methods=["POST"]) +@require_launcher_auth +def session_heartbeat_by_container(container_id: str) -> tuple[Response, int] | Response: + """ + Refresh a session's idle timer by container ID (orchestrator-only path). + + Used by the orchestrator to keep agent sessions alive while their + container is heartbeating on the BRC bus but not making gateway + requests — without this, the idle pruner evicts the session after + EGG_SESSION_IDLE_TIMEOUT_MINUTES even though the agent is still + working (see #2068). + + Auth: Bearer {launcher_secret} + + Returns 404 if no session exists for the container. No per-session + rate limit because the launcher secret already gates access — only + the orchestrator can call this. + """ + session_manager = get_session_manager() + refreshed = session_manager.heartbeat_session_by_container(container_id) + + if not refreshed: + return make_error("Session not found for container", status_code=404) + + return make_success("Heartbeat recorded") + + @app.route("/api/v1/sessions/", methods=["GET"]) @require_launcher_auth def session_get(session_token: str) -> tuple[Response, int] | Response: diff --git a/gateway/session_manager.py b/gateway/session_manager.py index c59b7059ee..fe22b70ab7 100644 --- a/gateway/session_manager.py +++ b/gateway/session_manager.py @@ -941,6 +941,39 @@ def delete_session_by_container(self, container_id: str) -> tuple[bool, threadin return False, None + def heartbeat_session_by_container(self, container_id: str) -> bool: + """Refresh a session's ``last_seen`` (and TTL) by container ID. + + Used by the orchestrator to keep gateway sessions alive while an + agent is heartbeating through the BRC bus but not making gateway + requests (e.g. a producer in ``WAITING_FOR_EVENT`` during a long + review cycle). Without this, the idle pruner evicts the session + after ``DEFAULT_SESSION_IDLE_TIMEOUT_MINUTES`` even though the + agent is still working — see #2068. + + Mirrors ``validate_session``'s in-flight pattern: the TTL extend + is in-memory only. We do **not** ``_save_to_disk()`` here — at + ~1 fan-out per agent every 60s a per-call atomic file write + would dominate gateway disk I/O for no benefit (worst-case loss + on a gateway crash is the session ages out and is re-registered + on the next gateway op, which is the same recovery path + ``validate_session`` already relies on). Disk persistence + happens at lifecycle events (registration, deletion, expiry). + + Args: + container_id: Container ID whose session to refresh. + + Returns: + True if a matching, non-expired session was refreshed; False + if no session exists for the container. + """ + with self._lock: + for session in self._sessions.values(): + if session.container_id == container_id and not session.is_expired(): + session.extend_ttl(self._ttl_hours) + return True + return False + def prune_expired_sessions(self) -> int: """ Remove all expired sessions. diff --git a/gateway/tests/test_gateway.py b/gateway/tests/test_gateway.py index b110a52f3c..8ffd8e215d 100644 --- a/gateway/tests/test_gateway.py +++ b/gateway/tests/test_gateway.py @@ -5118,6 +5118,46 @@ def test_session_delete_by_container_not_found(self, client, launcher_auth_heade assert response.status_code == 404 +class TestSessionHeartbeatByContainer: + """Tests for POST /api/v1/sessions/by-container//heartbeat (#2068).""" + + def test_refreshes_session(self, client, launcher_auth_headers): + mock_session_mgr = MagicMock() + mock_session_mgr.heartbeat_session_by_container.return_value = True + + with patch.object(gateway, "get_session_manager", return_value=mock_session_mgr): + response = client.post( + "/api/v1/sessions/by-container/egg-agent-pipeline-1-coder/heartbeat", + headers=launcher_auth_headers, + ) + + assert response.status_code == 200 + data = json.loads(response.data) + assert data["success"] is True + mock_session_mgr.heartbeat_session_by_container.assert_called_once_with( + "egg-agent-pipeline-1-coder" + ) + + def test_returns_404_when_session_missing(self, client, launcher_auth_headers): + mock_session_mgr = MagicMock() + mock_session_mgr.heartbeat_session_by_container.return_value = False + + with patch.object(gateway, "get_session_manager", return_value=mock_session_mgr): + response = client.post( + "/api/v1/sessions/by-container/nonexistent/heartbeat", + headers=launcher_auth_headers, + ) + + assert response.status_code == 404 + + def test_requires_launcher_auth(self, client): + """Endpoint must reject requests without the launcher secret.""" + response = client.post( + "/api/v1/sessions/by-container/egg-agent-pipeline-1-coder/heartbeat", + ) + assert response.status_code == 401 + + class TestBranchIsolation: """Tests for branch isolation enforcement in pipeline worktree sessions. diff --git a/gateway/tests/test_session_manager.py b/gateway/tests/test_session_manager.py index 05007eae3c..a546ea1114 100644 --- a/gateway/tests/test_session_manager.py +++ b/gateway/tests/test_session_manager.py @@ -779,6 +779,84 @@ def test_delete_clears_token_cache(self, manager): assert not result.valid +class TestHeartbeatByContainer: + """Tests for heartbeat_session_by_container (#2068).""" + + @pytest.fixture + def manager(self, tmp_path): + return SessionManager(persistence_file=tmp_path / "sessions.json") + + def test_heartbeat_refreshes_last_seen(self, manager): + """A heartbeat call advances last_seen so the idle pruner spares the session.""" + _token, session = manager.register_session( + container_id="egg-agent-pipeline-1-coder", + container_ip="172.18.0.5", + mode="private", + ) + + # Backdate last_seen past the idle timeout window. + session.last_seen = datetime.now(UTC) - timedelta(minutes=120) + + refreshed = manager.heartbeat_session_by_container("egg-agent-pipeline-1-coder") + assert refreshed is True + + # Session is now fresh -- not pruned by a 60-minute idle sweep. + assert manager.prune_idle_sessions(idle_timeout_minutes=60) == 0 + assert manager.get_session_by_container("egg-agent-pipeline-1-coder") is not None + + def test_heartbeat_unknown_container(self, manager): + """Heartbeat for an unknown container returns False without raising.""" + assert manager.heartbeat_session_by_container("not-a-real-container") is False + + def test_heartbeat_extends_ttl(self, manager): + """Heartbeat advances expires_at to ``now + ttl`` (mirrors validate_session). + + The 23h54m floor (rather than e.g. 23h59m) is a deliberate 6-minute + slack to absorb test-execution jitter between ``register_session`` + and the post-heartbeat ``datetime.now(UTC)`` call: a stalled + runner could in principle take a few seconds between the two, + and the assertion only needs to prove ``extend_ttl`` actually + fired (not that the clock was perfectly still). The strict + ``>`` against ``backdated_expiry`` (``now + 23.5h``) is what + proves the heartbeat moved the deadline forward — the 23h54m + floor is the *upper-bound* sanity check that it landed near + ``now + 24h`` rather than being clamped lower. + """ + _token, session = manager.register_session( + container_id="egg-agent-pipeline-1-coder", + container_ip="172.18.0.5", + mode="private", + ) + # Backdate both ``last_seen`` and ``expires_at`` 30 minutes into + # the past. ``extend_ttl`` should pull ``expires_at`` forward to + # roughly ``now + 24h`` — strictly greater than the backdated + # value of ``now + 23.5h``. + backdated_last_seen = datetime.now(UTC) - timedelta(minutes=30) + session.last_seen = backdated_last_seen + session.expires_at = backdated_last_seen + timedelta(hours=24) + backdated_expiry = session.expires_at + + manager.heartbeat_session_by_container("egg-agent-pipeline-1-coder") + refreshed_session = manager.get_session_by_container("egg-agent-pipeline-1-coder") + assert refreshed_session is not None + # Strict ``>`` proves the heartbeat moved ``expires_at`` forward; + # 23h54m floor proves it landed near ``now + 24h`` (sanity check + # — see docstring for the 6-minute slack rationale). + assert refreshed_session.expires_at > backdated_expiry + assert refreshed_session.expires_at > datetime.now(UTC) + timedelta(hours=23, minutes=54) + + def test_heartbeat_ignores_expired_session(self, manager): + """Already-expired sessions are not silently revived.""" + _token, session = manager.register_session( + container_id="egg-agent-pipeline-1-coder", + container_ip="172.18.0.5", + mode="private", + ) + session.expires_at = datetime.now(UTC) - timedelta(minutes=1) + + assert manager.heartbeat_session_by_container("egg-agent-pipeline-1-coder") is False + + class TestSessionModes: """Tests for session mode handling.""" diff --git a/orchestrator/gateway_client.py b/orchestrator/gateway_client.py index 56f5eb806a..ea01dfc1d5 100644 --- a/orchestrator/gateway_client.py +++ b/orchestrator/gateway_client.py @@ -573,6 +573,43 @@ def delete_session_by_container(self, container_id: str) -> bool: ) return False + def heartbeat_session_by_container(self, container_id: str) -> bool: + """Refresh a session's idle timer by container ID. + + Requires launcher secret authentication. Used to keep gateway + sessions alive while an agent is heartbeating on the BRC bus but + not making gateway requests — see #2068. + + Args: + container_id: Container ID whose session to refresh. + + Returns: + True if the session was refreshed; False if there is no + matching session or the gateway request failed. Best-effort + — callers should not fail on a False return. + """ + try: + result = self._make_request( + f"/api/v1/sessions/by-container/{quote(container_id, safe='')}/heartbeat", + method="POST", + use_launcher_auth=True, + ) + return result.get("success", False) + except GatewayError as e: + # Log the full container_id (not a secret — already shows up + # in k8s `get pods` output) so the failing pipeline+role is + # identifiable from #2068's exact failure mode. The sibling + # ``delete_session_by_container`` truncates to 12 chars + # (``egg-agent-is`` for realistic ids), which loses both + # pipeline and role; reviewer NB4 on #2076 flagged that as + # un-debuggable here even if it's pre-existing there. + logger.warning( + "Failed to heartbeat session by container", + container_id=container_id, + error=str(e), + ) + return False + def create_worktrees( self, container_id: str, diff --git a/orchestrator/heartbeat.py b/orchestrator/heartbeat.py index 0ecf1eb193..3bc4eda077 100644 --- a/orchestrator/heartbeat.py +++ b/orchestrator/heartbeat.py @@ -42,6 +42,8 @@ def __init__(self, window_seconds: int = 60) -> None: self._windows: dict[tuple[str, str], deque[float]] = {} # (pipeline_id, role) -> last-seen (state, waiting_on) self._last_state: dict[tuple[str, str], tuple[str, str]] = {} + # (pipeline_id, role) -> last gateway-session fan-out epoch-seconds + self._last_fan_out: dict[tuple[str, str], float] = {} def check_rate_limit( self, @@ -102,6 +104,48 @@ def record_state( with self._lock: self._last_state[key] = (state, waiting_on or "") + def should_fan_out_gateway_session( + self, + pipeline_id: str, + role: str, + min_interval_seconds: float, + ) -> bool: + """Throttle gateway-session fan-outs (issue #2076 NB2). + + The HEARTBEAT route fans out a gateway-session refresh both on + the dedup early-return path and after the rate-limit gate. The + dedup path bypasses the per-role rate limiter (by design — see + ``check_rate_limit``'s NB1 from #1897), so a misbehaving agent + hot-looping with identical state can amplify into the gateway + without burning rate budget. This per-role cooldown caps that + amplification independently of the heartbeat-acceptance rate. + + Returns ``True`` and records the new timestamp if at least + ``min_interval_seconds`` have passed since the previous fan-out + for this ``(pipeline_id, role)``; returns ``False`` (without + recording) if the caller should skip the fan-out this round. + Any non-positive ``min_interval_seconds`` (``<= 0``) disables + throttling — every call returns ``True`` without recording. + + Callers must pass a finite real number. ``float('nan')`` falls + through both branches (``nan <= 0`` and ``now - last < nan`` are + both False), which would silently behave as "always record, + never suppress" — not a meaningful state. The realistic call + site (``_GATEWAY_FANOUT_MIN_INTERVAL_SECONDS = 30.0``) is a + module constant, but if a future env-var-driven knob lands + (#2076 NB5) it should sanitize NaN/inf at parse time. + """ + if min_interval_seconds <= 0: + return True + key = (pipeline_id, role) + now = time.time() + with self._lock: + last = self._last_fan_out.get(key, 0.0) + if now - last < min_interval_seconds: + return False + self._last_fan_out[key] = now + return True + def clear(self, pipeline_id: str) -> None: """Drop all state for a pipeline (on phase transition).""" with self._lock: @@ -111,6 +155,9 @@ def clear(self, pipeline_id: str) -> None: for key in list(self._last_state): if key[0] == pipeline_id: del self._last_state[key] + for key in list(self._last_fan_out): + if key[0] == pipeline_id: + del self._last_fan_out[key] _coordinator: HeartbeatCoordinator | None = None diff --git a/orchestrator/routes/messages.py b/orchestrator/routes/messages.py index 221f9d5cda..8515fb5623 100644 --- a/orchestrator/routes/messages.py +++ b/orchestrator/routes/messages.py @@ -109,6 +109,16 @@ def _track_long_poll_end() -> None: # Rate-limit (20/min per role) still applies. _DEDUP_EXEMPT_HEARTBEAT_STATES: frozenset[str] = frozenset({"WAITING_FOR_EVENT"}) +# Minimum seconds between gateway-session fan-outs per (pipeline_id, +# role) (#2076 NB2). The dedup early-return path bypasses the per-role +# heartbeat rate limit by design (#1897 NB1: dedup'd heartbeats are +# no-ops and must not consume rate budget), so without a separate cap a +# misbehaving agent hot-looping with identical state could amplify into +# the gateway at the agent's emission rate. The gateway's idle window +# is 60 minutes, so fanning out every 30 s is far more than enough to +# keep the session alive; the cap exists purely to bound amplification. +_GATEWAY_FANOUT_MIN_INTERVAL_SECONDS: float = 30.0 + messages_bp = Blueprint("messages", __name__, url_prefix="/api/v1/pipelines") @@ -515,9 +525,21 @@ def post_heartbeat(pipeline_id: str) -> tuple[Response, int]: # budget (review NB1, issue #1897). States in # ``_DEDUP_EXEMPT_HEARTBEAT_STATES`` skip this check; see the # constant's docstring for the rationale. + # + # Note: the gateway-session fan-out below runs *after* dedup but + # *before* rate-limit. Dedup'd heartbeats still fan out so an agent + # stuck in a single state (e.g. ``WORKING`` through a slow + # ``make test``) keeps its gateway session alive even when its BRC + # state hasn't changed. Rate-limited heartbeats do not fan out: by + # definition the agent already got plenty of refreshes in the last + # minute, and a hot-looping agent shouldn't amplify into the + # gateway. ``_refresh_gateway_session`` itself applies a separate + # per-role cooldown (#2076 NB2) to bound dedup-path amplification + # without consuming rate budget. if state not in _DEDUP_EXEMPT_HEARTBEAT_STATES and coordinator.is_duplicate( pipeline_id, from_role, state, waiting_on ): + _refresh_gateway_session(pipeline_id, from_role) return _make_success( "HEARTBEAT deduped (unchanged state)", data={"deduped": True}, @@ -549,6 +571,14 @@ def post_heartbeat(pipeline_id: str) -> tuple[Response, int]: resp.headers["Retry-After"] = str(retry_after) return resp, 429 + # Refresh the agent's gateway session liveness (#2068). Runs after + # dedup and rate-limit gates: every accepted-or-deduped heartbeat + # fans out (dedup'd path above), but rate-limited ones do not. + # Best-effort: the gateway may be unreachable (tests, dev runs + # without a gateway) and a missing session is a 404; never fail the + # heartbeat on this path. + _refresh_gateway_session(pipeline_id, from_role) + # Emit as a normal HEARTBEAT message on the bus so downstream # consumers (HealthMonitor, overseer, UI) see it. metadata = {"state": state} @@ -589,6 +619,63 @@ def post_heartbeat(pipeline_id: str) -> tuple[Response, int]: ) +def _refresh_gateway_session(pipeline_id: str, from_role: str) -> None: + """Best-effort POST to the gateway so the BRC heartbeat counts as session liveness. + + Container-id normalization: k8s names are RFC-1123 labels (no + underscores), so ``kubernetes_spawner.JOB_NAME_FORMAT`` is filled + with ``agent_role.value.replace("_", "-")``. ``from_role`` arrives + from ``EGG_AGENT_ROLE`` which is the underscore form, so we mirror + the same normalization here — otherwise roles like + ``reviewer_refine`` build a container_id that never matches the + registered session and the gateway returns 404. See + ``orchestrator/kubernetes_spawner.py:370-375`` for the reference + pattern. + + Trust model: ``from_role`` is taken at face value from the request + body and is **not** correlated against the calling container's + session. This matches the existing message-bus trust model — any + agent in any container can already post messages claiming to be + another role — but with this fan-out a misbehaving agent can keep + a sibling's gateway session alive past the idle timeout. Tracked + as a follow-up; spoofing here doesn't grant any new capability, + only extends an existing session's lifetime. + + Throttle: per-role cooldown via + ``HeartbeatCoordinator.should_fan_out_gateway_session`` (#2076 NB2) + bounds amplification on the dedup early-return path, which bypasses + the heartbeat rate limiter by design. The cooldown is well below + the gateway's 60-minute idle window so it does not risk session + expiry under any realistic heartbeat cadence. + """ + coordinator = get_heartbeat_coordinator() + if not coordinator.should_fan_out_gateway_session( + pipeline_id, from_role, _GATEWAY_FANOUT_MIN_INTERVAL_SECONDS + ): + return + try: + try: + from gateway_client import get_gateway_client + except ImportError: # pragma: no cover + from ..gateway_client import ( + get_gateway_client, # type: ignore[no-redef,import-not-found] + ) + + # Mirror kubernetes_spawner.JOB_NAME_FORMAT's role normalization + # — k8s labels disallow underscores, so the registered + # container_id uses hyphens. + normalized_role = from_role.replace("_", "-") + container_id = f"egg-agent-{pipeline_id}-{normalized_role}" + get_gateway_client().heartbeat_session_by_container(container_id) + except Exception as exc: # pragma: no cover - logging only + logger.warning( + "Gateway session heartbeat fan-out failed", + pipeline_id=pipeline_id, + from_role=from_role, + error=str(exc), + ) + + @messages_bp.route("//messages/status", methods=["GET"]) def message_status(pipeline_id: str) -> tuple[Response, int]: """Get message bus status for a pipeline.""" diff --git a/orchestrator/tests/conftest.py b/orchestrator/tests/conftest.py index 3aa9d3549f..736f47aeca 100644 --- a/orchestrator/tests/conftest.py +++ b/orchestrator/tests/conftest.py @@ -222,6 +222,27 @@ def _skip_worktree_disk_check(monkeypatch): yield +@pytest.fixture(autouse=True) +def _reset_heartbeat_coordinator(): + """Reset the heartbeat coordinator singleton between every orchestrator test. + + The coordinator carries per-(pipeline, role) dedup, rate-limit, and + gateway-fan-out throttle state. Without resetting it, tests that + share role names can observe contaminated state across files. As + new coordinator surfaces are added, this single fixture covers them + all — see #2076 NB4 (the throttle was the third axis after dedup + and rate-limit, and the fixture had drifted across two test files). + """ + try: + from heartbeat import reset_heartbeat_coordinator + except ImportError: + yield + return + reset_heartbeat_coordinator() + yield + reset_heartbeat_coordinator() + + @pytest.fixture def lifecycle_secret() -> str: """The shared test bearer token. Useful for building explicit headers.""" diff --git a/orchestrator/tests/test_gateway_client.py b/orchestrator/tests/test_gateway_client.py index 87986122ff..527b79b644 100644 --- a/orchestrator/tests/test_gateway_client.py +++ b/orchestrator/tests/test_gateway_client.py @@ -74,6 +74,10 @@ def do_POST(self): self._handle_git_fetch(data) elif self.path == "/api/v1/gh/pr/create": self._handle_pr_create(data) + elif self.path.startswith("/api/v1/sessions/by-container/") and self.path.endswith( + "/heartbeat" + ): + self._handle_heartbeat_by_container() else: self._send_error(404, "Not found") @@ -169,6 +173,20 @@ def _handle_delete_by_container(self): self._send_json({"success": True}) + def _handle_heartbeat_by_container(self): + """Handle session heartbeat by container ID (POST .../by-container//heartbeat).""" + auth_header = self.headers.get("Authorization", "") + if not auth_header.startswith("Bearer ") or auth_header[7:] != "test-secret": + self._send_error(401, "Unauthorized") + return + + # Path: /api/v1/sessions/by-container//heartbeat + container_id = self.path.split("/")[-2] + if container_id == "missing": + self._send_json({"success": False, "message": "Session not found"}, status=404) + return + self._send_json({"success": True}) + def _handle_worktree_create(self, data): """Handle worktree creation (POST /api/v1/worktree/create).""" auth_header = self.headers.get("Authorization", "") @@ -487,6 +505,18 @@ def test_delete_session_by_container(self, gateway_client, mock_gateway_server): result = gateway_client.delete_session_by_container("container-123") assert result is True + def test_heartbeat_session_by_container(self, gateway_client, mock_gateway_server): + """Heartbeat-by-container returns True for an active session.""" + result = gateway_client.heartbeat_session_by_container("egg-agent-pipe-1-coder") + assert result is True + + def test_heartbeat_session_by_container_missing_returns_false( + self, gateway_client, mock_gateway_server + ): + """Heartbeat-by-container swallows 404 and returns False (best-effort path).""" + result = gateway_client.heartbeat_session_by_container("missing") + assert result is False + class TestSecurityBoundaryValidation: """Tests for security boundary validation.""" diff --git a/orchestrator/tests/test_heartbeat.py b/orchestrator/tests/test_heartbeat.py new file mode 100644 index 0000000000..2fb77a6323 --- /dev/null +++ b/orchestrator/tests/test_heartbeat.py @@ -0,0 +1,175 @@ +"""Unit tests for ``HeartbeatCoordinator``. + +Focused on ``should_fan_out_gateway_session`` (issue #2076 NB4) — the +route-level integration tests in ``test_messages.py`` exercise the +throttle through the HEARTBEAT endpoint, but a refactor of the +coordinator is hard to do safely without targeted unit coverage of the +disable case, the cooldown elapsed/not-elapsed branches, the ``clear()`` +reset, and concurrent access. +""" + +from __future__ import annotations + +import sys +import threading +import time +from pathlib import Path + +# Add orchestrator to path so bare imports resolve. +_orchestrator_path = Path(__file__).parent.parent +if str(_orchestrator_path) not in sys.path: + sys.path.insert(0, str(_orchestrator_path)) + +from heartbeat import ( # noqa: E402 + HeartbeatCoordinator, + get_heartbeat_coordinator, + reset_heartbeat_coordinator, +) + +# Singleton reset between tests is handled by ``_reset_heartbeat_coordinator`` +# in ``conftest.py`` (autouse-scoped to all orchestrator tests). + + +class TestShouldFanOutGatewaySession: + """Unit tests for the per-(pipeline, role) gateway-fan-out throttle.""" + + def test_zero_interval_disables_throttle(self): + """``min_interval_seconds == 0`` always returns True without recording.""" + coord = HeartbeatCoordinator() + + for _ in range(5): + assert coord.should_fan_out_gateway_session("p1", "coder", 0.0) is True + + # Nothing was recorded, so a subsequent positive-interval call + # sees no prior fan-out and fires. + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + + def test_negative_interval_disables_throttle(self): + """Any non-positive value disables (matches the docstring contract).""" + coord = HeartbeatCoordinator() + + for _ in range(3): + assert coord.should_fan_out_gateway_session("p1", "coder", -1.0) is True + + # Same as the zero case — no recording happened. + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + + def test_first_call_fires_and_records(self): + """First call fires AND records the new timestamp. + + The recording side-effect is verified directly: a second call + inside the cooldown window can only return False if the first + call wrote ``_last_fan_out[(p1, coder)]``. + """ + coord = HeartbeatCoordinator() + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + # Recording side-effect: second call within the window is suppressed. + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is False + + def test_second_call_within_cooldown_suppressed(self): + """A second call inside the cooldown window returns False.""" + coord = HeartbeatCoordinator() + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + # Immediately again — well inside 30s. + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is False + + def test_call_after_cooldown_fires(self): + """Once the window elapses, the next call fires again.""" + coord = HeartbeatCoordinator() + # 50ms cooldown + 200ms sleep — 150ms headroom keeps the test + # robust against GC pauses or noisy CI scheduling. ``time.sleep`` + # is a guaranteed lower bound, so the cooldown is always elapsed + # by the time the third assertion runs. + assert coord.should_fan_out_gateway_session("p1", "coder", 0.05) is True + assert coord.should_fan_out_gateway_session("p1", "coder", 0.05) is False + time.sleep(0.2) + assert coord.should_fan_out_gateway_session("p1", "coder", 0.05) is True + + def test_different_roles_throttled_independently(self): + """Throttle key includes the role — different roles don't collide.""" + coord = HeartbeatCoordinator() + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + assert coord.should_fan_out_gateway_session("p1", "tester", 30.0) is True + # Both are now suppressed independently. + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is False + assert coord.should_fan_out_gateway_session("p1", "tester", 30.0) is False + + def test_different_pipelines_throttled_independently(self): + """Throttle key includes the pipeline — different pipelines don't collide.""" + coord = HeartbeatCoordinator() + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + assert coord.should_fan_out_gateway_session("p2", "coder", 30.0) is True + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is False + assert coord.should_fan_out_gateway_session("p2", "coder", 30.0) is False + + def test_suppressed_call_does_not_advance_recorded_timestamp(self): + """Hot-looping must not push the cooldown forward indefinitely. + + Uses a 50ms cooldown + 200ms sleep — 150ms headroom keeps the + ``time.time()`` reads robust against scheduling jitter on CI. + """ + coord = HeartbeatCoordinator() + assert coord.should_fan_out_gateway_session("p1", "coder", 0.05) is True + # Hammer the coordinator inside the window — every call returns + # False, but the recorded timestamp stays at the initial fire. + for _ in range(10): + assert coord.should_fan_out_gateway_session("p1", "coder", 0.05) is False + time.sleep(0.2) + # Still fires once the original window elapses, which proves the + # suppressed calls didn't reset the clock. + assert coord.should_fan_out_gateway_session("p1", "coder", 0.05) is True + + def test_clear_drops_throttle_state_for_pipeline(self): + """``clear(pipeline)`` lets the next call fire immediately.""" + coord = HeartbeatCoordinator() + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is False + coord.clear("p1") + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + + def test_clear_only_affects_targeted_pipeline(self): + """``clear(p1)`` must not drop p2's throttle entries.""" + coord = HeartbeatCoordinator() + coord.should_fan_out_gateway_session("p1", "coder", 30.0) + coord.should_fan_out_gateway_session("p2", "coder", 30.0) + coord.clear("p1") + # p1 reset, fires again. + assert coord.should_fan_out_gateway_session("p1", "coder", 30.0) is True + # p2 untouched, still suppressed. + assert coord.should_fan_out_gateway_session("p2", "coder", 30.0) is False + + def test_concurrent_callers_only_one_wins(self): + """Under contention, exactly one thread should see True per window.""" + coord = HeartbeatCoordinator() + results: list[bool] = [] + results_lock = threading.Lock() + barrier = threading.Barrier(20) + + def worker(): + barrier.wait() + res = coord.should_fan_out_gateway_session("p1", "coder", 30.0) + with results_lock: + results.append(res) + + threads = [threading.Thread(target=worker) for _ in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + + # Exactly one thread crossed the throttle boundary first. + assert sum(results) == 1 + assert len(results) == 20 + + +class TestSingletonAccessor: + def test_get_returns_same_instance(self): + a = get_heartbeat_coordinator() + b = get_heartbeat_coordinator() + assert a is b + + def test_reset_replaces_instance(self): + a = get_heartbeat_coordinator() + reset_heartbeat_coordinator() + b = get_heartbeat_coordinator() + assert a is not b diff --git a/orchestrator/tests/test_messages.py b/orchestrator/tests/test_messages.py index 0eea831229..2343ff7667 100644 --- a/orchestrator/tests/test_messages.py +++ b/orchestrator/tests/test_messages.py @@ -6,6 +6,7 @@ import json import sys +import time from pathlib import Path from unittest.mock import MagicMock, patch @@ -39,7 +40,12 @@ def client(app): @pytest.fixture(autouse=True) def _reset_store(): - """Reset message store singleton between tests.""" + """Reset the message store between tests. + + Heartbeat coordinator reset is handled by ``_reset_heartbeat_coordinator`` + in ``conftest.py`` (autouse-scoped to all orchestrator tests) so the + cleanup is shared across files instead of drifting per-test-module. + """ reset_message_store() yield reset_message_store() @@ -1830,6 +1836,284 @@ def test_heartbeat_route_accepts_optional_since_field(self, client, app): data = json.loads(resp.data) assert data["data"]["message"]["metadata"]["since"] == ("2026-04-23T07:00:00Z") + def test_heartbeat_refreshes_gateway_session(self, client, app): + """A real heartbeat fans out to the gateway to refresh the agent's session. + + Regression guard for #2068: without this, an agent in + ``WAITING_FOR_EVENT`` for >60 min has its gateway session pruned + as idle even though it's actively heartbeating. + """ + with app.test_request_context(): + mock_gw_client = MagicMock() + with ( + patch( + "routes.messages.get_state_store_for_pipeline" + ) as mock_get_store_for_pipeline, + patch( + "gateway_client.get_gateway_client", + return_value=mock_gw_client, + ), + ): + mock_get_store_for_pipeline.return_value = ( + MagicMock(), + _make_pipeline_mock(), + ) + resp = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={ + "from_role": "fanout-role-coder", + "state": "WAITING_FOR_EVENT", + }, + ) + assert resp.status_code == 200 + mock_gw_client.heartbeat_session_by_container.assert_called_once_with( + "egg-agent-test-pipeline-fanout-role-coder" + ) + + def test_heartbeat_swallows_gateway_failure(self, client, app): + """Gateway fan-out failures must not fail the heartbeat call.""" + with app.test_request_context(): + mock_gw_client = MagicMock() + mock_gw_client.heartbeat_session_by_container.side_effect = RuntimeError("gateway down") + with ( + patch( + "routes.messages.get_state_store_for_pipeline" + ) as mock_get_store_for_pipeline, + patch( + "gateway_client.get_gateway_client", + return_value=mock_gw_client, + ), + ): + mock_get_store_for_pipeline.return_value = ( + MagicMock(), + _make_pipeline_mock(), + ) + resp = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={ + "from_role": "fanout-role-tester", + "state": "WORKING", + }, + ) + assert resp.status_code == 200 + + def test_heartbeat_fan_out_fires_on_deduped_state(self, client, app): + """Deduped heartbeats still refresh the gateway session. + + Reviewer NB3 (#2068): the original fix only ran the fan-out + *after* dedup, so an agent stuck in ``WORKING`` through a long + compute (e.g. a slow ``make test``) would emit identical + heartbeats that were dropped before refreshing the session. The + fan-out now runs *between* dedup and rate-limit (see + ``routes/messages.py``): the dedup early-return invokes it so + unchanged-state heartbeats still count as gateway-session + liveness, while rate-limited heartbeats do not amplify into the + gateway. + + The ``_GATEWAY_FANOUT_MIN_INTERVAL_SECONDS`` cooldown (#2076 NB2) + is patched to ``0`` here so the test stays focused on the + dedup-fan-out invariant; the throttle's caps are pinned by + ``test_heartbeat_fan_out_throttle_*`` below. + """ + with app.test_request_context(): + mock_gw_client = MagicMock() + with ( + patch( + "routes.messages.get_state_store_for_pipeline" + ) as mock_get_store_for_pipeline, + patch( + "gateway_client.get_gateway_client", + return_value=mock_gw_client, + ), + patch( + "routes.messages._GATEWAY_FANOUT_MIN_INTERVAL_SECONDS", + 0.0, + ), + ): + mock_get_store_for_pipeline.return_value = ( + MagicMock(), + _make_pipeline_mock(), + ) + # First WORKING heartbeat — recorded, not deduped. + resp1 = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={ + "from_role": "fanout-dedup-role", + "state": "WORKING", + }, + ) + assert resp1.status_code == 200 + # Second identical heartbeat — dedup gate should drop it. + resp2 = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={ + "from_role": "fanout-dedup-role", + "state": "WORKING", + }, + ) + assert resp2.status_code == 200 + assert json.loads(resp2.data)["data"]["deduped"] is True + # Both calls fan out — that's the bug fix. + assert mock_gw_client.heartbeat_session_by_container.call_count == 2 + + @pytest.mark.parametrize( + "from_role,expected_container_id", + [ + ("reviewer_refine", "egg-agent-test-pipeline-reviewer-refine"), + ("reviewer_code", "egg-agent-test-pipeline-reviewer-code"), + ("reviewer_agent_design", "egg-agent-test-pipeline-reviewer-agent-design"), + ("task_planner", "egg-agent-test-pipeline-task-planner"), + ("conflict_resolver", "egg-agent-test-pipeline-conflict-resolver"), + ("coder", "egg-agent-test-pipeline-coder"), + ], + ) + def test_heartbeat_fan_out_normalizes_underscores_to_hyphens( + self, client, app, from_role, expected_container_id + ): + """Fan-out container_id must mirror kubernetes_spawner's role hyphenation. + + Reviewer blocker (#2068 follow-up): k8s names are RFC-1123 + labels (no underscores), so ``kubernetes_spawner.JOB_NAME_FORMAT`` + is filled with ``role.replace("_", "-")``. ``from_role`` arrives + from ``EGG_AGENT_ROLE`` which is the underscore form, so the + fan-out must apply the same normalization — otherwise reviewer + roles like ``reviewer_refine`` build a container_id that never + matches the registered session and the gateway returns 404, + making the fan-out a silent no-op for exactly the BRC reviewer + roles that #2068 most acutely affects. + """ + with app.test_request_context(): + mock_gw_client = MagicMock() + with ( + patch( + "routes.messages.get_state_store_for_pipeline" + ) as mock_get_store_for_pipeline, + patch( + "gateway_client.get_gateway_client", + return_value=mock_gw_client, + ), + ): + mock_get_store_for_pipeline.return_value = ( + MagicMock(), + _make_pipeline_mock(), + ) + resp = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={ + "from_role": from_role, + "state": "WAITING_FOR_EVENT", + }, + ) + assert resp.status_code == 200 + mock_gw_client.heartbeat_session_by_container.assert_called_once_with( + expected_container_id + ) + + def test_heartbeat_fan_out_throttle_caps_dedup_amplification(self, client, app): + """#2076 NB2: dedup'd hot-loops cannot amplify into the gateway. + + The dedup early-return path bypasses the per-role rate limit by + design (#1897 NB1: dedup'd heartbeats are no-ops and must not + consume rate budget). Without a separate cap, a misbehaving + agent hot-looping with identical state could fan out a gateway + session refresh on every call. ``_refresh_gateway_session`` + applies a per-role cooldown via + ``HeartbeatCoordinator.should_fan_out_gateway_session`` to bound + amplification. + + Five back-to-back identical heartbeats (1 fresh + 4 dedup'd) + within the cooldown window MUST produce exactly one fan-out. + """ + with app.test_request_context(): + mock_gw_client = MagicMock() + with ( + patch( + "routes.messages.get_state_store_for_pipeline" + ) as mock_get_store_for_pipeline, + patch( + "gateway_client.get_gateway_client", + return_value=mock_gw_client, + ), + # Cooldown well above realistic test wall-clock so all + # five posts fall inside the same window. + patch( + "routes.messages._GATEWAY_FANOUT_MIN_INTERVAL_SECONDS", + 300.0, + ), + ): + mock_get_store_for_pipeline.return_value = ( + MagicMock(), + _make_pipeline_mock(), + ) + role = "fanout-throttle-hotloop-role" + for _ in range(5): + resp = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={"from_role": role, "state": "WORKING"}, + ) + assert resp.status_code == 200 + assert mock_gw_client.heartbeat_session_by_container.call_count == 1 + + def test_heartbeat_fan_out_throttle_resumes_after_window(self, client, app): + """#2076 NB2: cooldown is a throttle, not a one-shot mute. + + After the per-role cooldown elapses, the next heartbeat MUST + fan out again — otherwise a long-running agent in a single + state would fan out exactly once and then silently age out of + the gateway's 60-minute idle window. + + Uses a tiny cooldown + ``time.sleep`` rather than mocking + ``time.time`` so the patch doesn't ripple through unrelated + callers (Flask internals, gateway client) inside the request + scope. Uses ``WAITING_FOR_EVENT`` (dedup-exempt) so this + exercise also pins the throttle on the post-rate-limit fan-out + site, not just the dedup early-return. + """ + with app.test_request_context(): + mock_gw_client = MagicMock() + with ( + patch( + "routes.messages.get_state_store_for_pipeline" + ) as mock_get_store_for_pipeline, + patch( + "gateway_client.get_gateway_client", + return_value=mock_gw_client, + ), + patch( + "routes.messages._GATEWAY_FANOUT_MIN_INTERVAL_SECONDS", + 0.05, + ), + ): + mock_get_store_for_pipeline.return_value = ( + MagicMock(), + _make_pipeline_mock(), + ) + role = "fanout-throttle-window-role" + # First heartbeat — fans out. + resp1 = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={"from_role": role, "state": "WAITING_FOR_EVENT"}, + ) + assert resp1.status_code == 200 + assert mock_gw_client.heartbeat_session_by_container.call_count == 1 + + # Inside the 50 ms cooldown — no additional fan-out. + resp2 = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={"from_role": role, "state": "WAITING_FOR_EVENT"}, + ) + assert resp2.status_code == 200 + assert mock_gw_client.heartbeat_session_by_container.call_count == 1 + + # Past the cooldown — fan-out fires again. + time.sleep(0.07) + resp3 = client.post( + "/api/v1/pipelines/test-pipeline/heartbeat", + json={"from_role": role, "state": "WAITING_FOR_EVENT"}, + ) + assert resp3.status_code == 200 + assert mock_gw_client.heartbeat_session_by_container.call_count == 2 + class TestWaitTimeoutFloorRegression: """Plan non-blocking: ``timeout <= 0`` is silently coerced to 1s