From 0cd4d0203645f8f23f7ceca716dbc7e4799bb4ad Mon Sep 17 00:00:00 2001 From: Christian-Sidak <61099993+Christian-Sidak@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:36:50 -0700 Subject: [PATCH] fix: heartbeat is_current_turn comparison against lock, not durable row When a steer happens, `_start_turn()` overwrites the durable row's `turn_id` with the new turn before the old turn's next heartbeat fires. The original code compared the durable row's `turn_id` to decide if a lock loss was a cancel vs. a first-beat race. For the steer case, the row already shows the new turn_id, so `turn_was_established` is False and `is_current_turn` stays True -- the old turn thinks it is still current even though it has been displaced. Fix: when `refresh_alive` fails, read the alive lock's current holder directly. If the lock is held by a different turn_id, the old turn was displaced by a steer and `is_current_turn` is set to False. The durable row check (`turn_was_established`) is kept as the fallback for the cancel case (lock absent, row matches). Adds a regression test that drives the real steer write path via `command()` with `force=True` rather than manually simulating the lock cancellation, matching the exact sequence from issue #5790. Fixes #5790 Signed-off-by: Christian-Sidak <61099993+Christian-Sidak@users.noreply.github.com> --- api/oss/src/core/sessions/streams/service.py | 41 ++++++--- .../test_heartbeat_is_current_turn.py | 86 ++++++++++++++++++- 2 files changed, 114 insertions(+), 13 deletions(-) diff --git a/api/oss/src/core/sessions/streams/service.py b/api/oss/src/core/sessions/streams/service.py index 39298d3e32..7ce009a665 100644 --- a/api/oss/src/core/sessions/streams/service.py +++ b/api/oss/src/core/sessions/streams/service.py @@ -30,6 +30,7 @@ clear_running, force_cancel_alive, force_clear_owner, + get_alive_owner, get_session_liveness, refresh_alive, refresh_running, @@ -284,16 +285,23 @@ async def heartbeat( ) # True only when this turn_id still (or again, uninterrupted) owns the alive lock at - # the moment of this heartbeat. A cancel/steer/kill deletes the alive key entirely, - # which the nx=True re-acquire below would otherwise silently re-establish under the - # SAME turn_id, masking the interruption from the runner's watchdog (W7.4 — this is - # what `is_current_turn` exists to surface). An absent key is ambiguous by itself: it - # is also the normal state before this turn's VERY FIRST heartbeat (the API's - # `_start_turn` acquire may not have landed yet, or this beat wins a race with it), and - # that is NOT an interruption. Disambiguate with the durable row's `turn_id`: if it - # already recorded THIS turn_id as established (a prior heartbeat's write), the key - # being gone now is something else's doing; if the row shows no turn yet, or a - # different one, this is establishment. + # the moment of this heartbeat. A cancel/steer/kill deletes the alive key (cancel) + # or replaces it with a new turn_id (steer), which the nx=True re-acquire below would + # otherwise silently re-establish under the SAME turn_id, masking the interruption from + # the runner's watchdog (W7.4 -- this is what `is_current_turn` exists to surface). + # + # An absent key is ambiguous by itself: it is also the normal state before this turn's + # VERY FIRST heartbeat (the API's `_start_turn` acquire may not have landed yet, or + # this beat wins a race with it), and that is NOT an interruption. Disambiguate via the + # alive lock's current owner: + # - lock held by a different turn_id --> steer displaced us (is_current_turn=False) + # - lock absent + row already has our turn_id --> cancel cleared it (is_current_turn=False) + # - lock absent + row shows no/other turn_id --> first-beat race, treat as establishment + # + # Comparing against the lock (not the durable row) is correct here: a steer calls + # `_start_turn`, which overwrites the durable row's turn_id with the new turn before + # the old turn's next heartbeat fires, so the row-based check always sees a "different" + # turn_id and cannot distinguish "displaced" from "first beat of new turn". prior_stream = await self._dao.get_by_session_id( project_id=project_id, session_id=request.session_id, @@ -305,14 +313,23 @@ async def heartbeat( if request.turn_id and request.is_running: # Acquire-then-refresh: the first heartbeat must establish the nest locks - # itself (acquire_* is nx=True — a no-op if _start_turn already holds them). + # itself (acquire_* is nx=True -- a no-op if _start_turn already holds them). if not await refresh_alive( self._lock, project_id=str(project_id), session_id=request.session_id, turn_id=request.turn_id, ): - if turn_was_established: + alive_holder = await get_alive_owner( + self._lock, + project_id=str(project_id), + session_id=request.session_id, + ) + if alive_holder is not None and alive_holder != request.turn_id: + # Another turn holds the lock -- we were displaced by a steer. + is_current_turn = False + elif turn_was_established: + # Lock is absent and this turn was previously established -- cancelled. is_current_turn = False await acquire_alive( self._lock, diff --git a/api/oss/tests/pytest/unit/sessions/test_heartbeat_is_current_turn.py b/api/oss/tests/pytest/unit/sessions/test_heartbeat_is_current_turn.py index 4f72336691..155395c472 100644 --- a/api/oss/tests/pytest/unit/sessions/test_heartbeat_is_current_turn.py +++ b/api/oss/tests/pytest/unit/sessions/test_heartbeat_is_current_turn.py @@ -1,7 +1,7 @@ """WP7 (W7.4): the control signal from cancel/steer/kill must reach the runner's heartbeat. Before this, `heartbeat()`'s acquire-then-refresh fallback silently re-acquired a lost alive -lock under the SAME turn_id (nx=True is a no-op only when the key is gone) — a cancel/steer/ +lock under the SAME turn_id (nx=True is a no-op only when the key is gone) -- a cancel/steer/ kill that raced a heartbeat was invisible to the runner: the beat still looked like a normal `ok` heartbeat. `is_current_turn` on `SessionHeartbeatResult` surfaces the interruption so the runner's watchdog can abort the in-flight run (`services/runner/src/sessions/alive.ts`'s @@ -13,6 +13,8 @@ is_current_turn to False (the lock was gone, then silently re-acquired); - a steer (different turn_id takes the lock) also reports the OLD turn's next beat as is_current_turn=False, and does not steal the lock back for the old turn; + - a steer via command() (the real write path) flips the old turn's heartbeat: regression + for issue #5790 where the durable-row comparison masked the displacement; - a replica that lost the owner claim entirely reports is_current_turn=False. """ @@ -23,9 +25,12 @@ import pytest import pytest_asyncio +from agenta.sdk.models.workflows import WorkflowServiceRequestData + from oss.src.core.sessions.streams.dtos import ( SessionHeartbeatRequest, SessionStream, + SessionStreamCommandRequest, ) from oss.src.core.sessions.streams.service import SessionStreamsService from oss.src.dbs.redis.sessions.locks import force_cancel_alive, get_alive_owner @@ -34,6 +39,7 @@ _PROJECT = uuid4() +_USER = uuid4() _SESSION = "session_interrupt" @@ -175,3 +181,81 @@ async def test_losing_owner_claim_reports_not_current(lock_engine): assert result.is_current_turn is False assert result.replica_id == "replica-a" + + +@pytest.mark.asyncio +async def test_steer_via_command_flips_old_turn_heartbeat_to_not_current(lock_engine): + """Regression for issue #5790. + + The old code compared against the durable row's turn_id to detect whether a lock loss + was an interruption or a first-beat race. A steer calls _start_turn(), which overwrites + the durable row with the new turn_id before the old turn's next heartbeat fires. The old + turn therefore saw `prior_stream.turn_id != request.turn_id`, treated the lock loss as a + first-beat race, and kept reporting is_current_turn=True even though it had been displaced. + + The fix compares against the alive lock's current holder (not the durable row): if the + lock is held by a different turn_id, the old turn is displaced regardless of what the row + says. + + This test drives the real steer write path (command() with force=True) instead of + manually simulating it, so it catches the exact sequence that triggered the bug. + """ + session_id = f"session_steer_cmd_{uuid4().hex[:8]}" + svc = _service(lock_engine) + + # Start turn-1 via the command endpoint (send path). + send_resp = await svc.command( + project_id=_PROJECT, + user_id=_USER, + request=SessionStreamCommandRequest( + session_id=session_id, + data=WorkflowServiceRequestData(inputs={"messages": ["hello"]}), + force=False, + ), + ) + old_turn_id = send_resp.turn_id + assert old_turn_id is not None + + # Runner's heartbeat establishes turn-1 in the durable row. + first_beat = await svc.heartbeat( + project_id=_PROJECT, + request=SessionHeartbeatRequest( + session_id=session_id, + replica_id="replica-a", + turn_id=old_turn_id, + is_running=True, + ), + ) + assert first_beat.is_current_turn is True + + # Steer: a new message displaces turn-1 via the command endpoint's write path. + # This calls _start_turn(), which overwrites the durable row with the new turn_id -- + # the exact condition that caused the bug. + steer_resp = await svc.command( + project_id=_PROJECT, + user_id=_USER, + request=SessionStreamCommandRequest( + session_id=session_id, + data=WorkflowServiceRequestData(inputs={"messages": ["steer"]}), + force=True, + ), + ) + new_turn_id = steer_resp.turn_id + assert new_turn_id is not None + assert new_turn_id != old_turn_id + + # Old turn's next heartbeat must report is_current_turn=False now that the durable row + # shows new_turn_id and the alive lock is held by new_turn_id, not old_turn_id. + old_turn_next_beat = await svc.heartbeat( + project_id=_PROJECT, + request=SessionHeartbeatRequest( + session_id=session_id, + replica_id="replica-a", + turn_id=old_turn_id, + is_running=True, + ), + ) + assert old_turn_next_beat.is_current_turn is False, ( + "displaced turn's heartbeat must report is_current_turn=False after a steer " + "via command(), not just after a manually simulated lock cancellation" + )