fix(gateway): fail closed when a pinned route moves under a delegation completion

[wangtaotaotao95]

What does this PR do?

Fixes the race in #113690: a pinned async-delegation completion could overwrite a routing key that had already moved, because the resolution awaits the spawning session's DB row and then repoints the key unconditionally.

gateway/run_notifications.py::_resolve_async_delegation_session does:

pinned_row = await session_db.get_session(pinned_session_id)   # <-- await: the race window
...
switched = await self.async_session_store.switch_session(session_entry.session_key, target_session_id)

That await is the window. /new or /stop can revoke the in-flight run generation and a concurrent /resume can repoint the key while the lookup is suspended; the stale resolution then publishes the pin it snapshotted before the boundary. A revoked completion must not be able to write routing at all — the resolved entry is what every later delivery in that turn is addressed to.

The compression branch was already safe: it goes through advance_compression_session, which is a compare-and-swap on the expected current session. The non-compression branch had no such condition, and switch_session accepted none.

Related Issue

Fixes #113690

Type of Change

• 🐛 Bug fix (non-breaking change that fixes an issue)

Changes Made

• gateway/run_notifications.py — _resolve_async_delegation_session now takes the caller's run token (optional), snapshots it before the row lookup, and re-checks it after the await; a revoked run drops the injection instead of moving the route. The non-compression repoint passes the session it resolved against as a condition.
• gateway/session.py — switch_session takes an optional expected_session_id and refuses the repoint when the route no longer points there (returns None, i.e. fail closed). Existing two-argument callers (/resume, handoff, startup) are unchanged.
• gateway/run_turn.py — _hmwa_resolve_session (and its call from _handle_message_with_agent) forwards _quick_key / run_generation, which the resolver never received. The parameter is optional, so the eval harness caller is unaffected.
• gateway/run_agent_cache.py — adds _current_session_run_generation, so a caller that does not pass its own token still snapshots a comparable one.
• tests/gateway/test_pinned_route_race.py — new invariant test over a real SessionStore and the real generation primitives.
• tests/gateway/test_async_delegation_session_binding.py — the "live spawning session rebinds" assertion now expects the conditional repoint. The behaviour it asserts is the same; the contract got stricter.

Both boundaries fail closed: the route keeps whatever the newer decision set, and the completion is dropped.

How to Test

scripts/run_tests.sh -j 1 --file-retries 0 tests/gateway/test_pinned_route_race.py -p no:cacheprovider -q

The test suspends the session-row lookup on an asyncio.Event, applies a boundary, then releases it — no wall-clock sleeps:

• none — nothing moves; the pin repoints the route and is returned.
• revoke — the run generation is invalidated while the lookup is pending; the route must stay put and the result must be None.
• replace — the run is revoked and the key is repointed to a replacement session; the replacement must survive and the result must be None.

Proof it is red on base, at 6005aa1fd9:

[100.0% | 1/~1 | ✓1 | ✗2] ✗ tests/gateway/test_pinned_route_race.py (1✓ 2✗, 0.5s)
E   AssertionError: assert 'test-pinned' == 'test-replacement'
2 failed, 1 passed

With the fix: 3 passed.

No regressions. Full gateway suite (tests/gateway/, 901 files):

=== Summary: 901 files, 8862 tests passed, 6 failed, 50 skipped (100% complete) in 288.0s ===

The 6 failures are identical on unpatched 6005aa1fd9 (397 passed, 6 failed on the same six files) — they are host-environment failures (test_gateway_trust_env, test_slack, test_feishu, test_wecom, test_telegram_thread_fallback, test_multiplex_residue_parity try to reach the network / observe this host's system proxy). The 31 test files that reference any changed symbol pass: 356 passed.

ruff check and ruff format --check are clean on the new test file.

Checklist

Code

• I've read the Contributing Guide
• My commit messages follow Conventional Commits (fix(scope):, feat(scope):, etc.)
• I searched for existing PRs to make sure this isn't a duplicate
• My PR contains only changes related to this fix/feature (no unrelated commits)
• I've run the gateway suite and all tests that pass on base still pass
• I've added tests for my changes (required for bug fixes, strongly encouraged for features)
• I've tested on my platform: macOS 15 (darwin), Python 3.11.16

Documentation & Housekeeping

• I've updated relevant documentation — N/A (docstrings updated in place)
• I've updated cli-config.yaml.example if I added/changed config keys — N/A
• I've updated CONTRIBUTING.md or AGENTS.md if I changed architecture or workflows — N/A
• I've considered cross-platform impact — N/A (no platform-specific code; the test is OS-independent)
• I've updated tool descriptions/schemas if I changed tool behavior — N/A

Screenshots / Logs

Base (6005aa1fd9) vs. fix, same command:

$ scripts/run_tests.sh -j 1 --file-retries 0 tests/gateway/test_pinned_route_race.py -p no:cacheprovider -q --tb=line
✗ tests/gateway/test_pinned_route_race.py (1✓ 2✗, 0.5s)     # base
E   AssertionError: assert 'test-pinned' == '20260917_031405_073c3e20'
E   AssertionError: assert 'test-pinned' == 'test-replacement'

✓ tests/gateway/test_pinned_route_race.py (3✓, 0.6s)        # fix

[whyyagswhy]

Independent verification on the PR head (2ee68ad): the pinned-route race suite plus the delegation binding suite pass locally, 9/9 on Linux (canonical runner).

Cleaner than the sibling approach: the CAS lives inside switch_session as an optional expected_session_id rather than a parallel method, so all callers share one conditional path and stale resolutions fail closed with a logged refusal. Note for maintainers: #113692 covers adjacent ground with switch_session_if_current; worth converging on one shape.

[andrexibiza]

I traced this exact head through the real inbound run-generation claim, _hmwa_resolve_session, the async delegation resolver, AsyncSessionStore, the synchronous SessionStore mutation, the new event-synchronized regression, #113690, the competing #113692 carrier, and the established dispatch-time ownership lineage from merged #60871. The expected-session CAS is a real improvement and closes the concurrent-route-replacement half of the reporter's reproducer without weakening the captured owner. There is still one blocking authority race in the current shape.

P1 — the run-generation proof is still separated from the effect by another scheduling boundary.

The new check correctly catches a /stop//new that lands while session_db.get_session() is suspended. But the proof is consumed only once, immediately after that first await. The non-compression mutation then goes through AsyncSessionStore, whose generic wrapper is await asyncio.to_thread(attr, *args, **kwargs). That creates a second interference window after the generation check and before the synchronous SessionStore.switch_session() mutation.

A concrete execution is:

1. Turn owns generation G; the pinned-row lookup returns.
2. _is_session_run_current(key, G) succeeds.
3. Before the offloaded switch_session() acquires the store lock, /stop invalidates the run to G+1 without changing the route's current session id — the pure-revoke case from #113690.
4. switch_session(... expected_session_id=prior_session_id) still sees the expected route id and commits the stale repin.

The expected-session CAS cannot detect that boundary because route identity and run generation are different authority dimensions. The same root exists on the compression side: after the one generation check, _resolve_compression_lineage_target() can perform several more awaits, and the target_session_id == session_entry.session_id fast path can return a stale entry without any final generation validation at all. So the PR narrows the original race window, but it does not yet make the advertised invariant — a revoked completion cannot publish after the boundary — true for the whole mutation/delivery path.

The closure needs the generation proof to be effect-bound, not another preflight observation: validate the expected route and expected run generation under one coordination/commit boundary that cannot interleave with generation invalidation, then perform/authorize the route transition (or fail closed). A second check immediately before another await is still TOCTOU; post-hoc rollback is also unsafe because a newer route decision may have landed by then. Please add a deterministic witness that suspends the route commit after the new generation check, invalidates the generation while the commit is held, then releases it and proves both that the routing key stayed put and that the completion was not returned for injection. Mirror that witness across the compression/identity fast path as the other side of the same invariant.

The surrounding topology matters here:

• #113690 / tobific owns the current race report and its none / revoke / replace invariant.
• #113692 / KoNit-K is a competing partial carrier. Its route CAS closes replace, but the reporter's pure revoke case was independently shown still failing because the route can remain unchanged. This PR correctly adds generation to the proof; it should be the place where that proof is carried all the way to the effect rather than stopping one await early.
• #60871 / teknium1, salvaging #57535 / nankingjing with authorship preserved, established parent_session_id / gateway_session_id as the durable dispatch-time owner. This PR preserves that stronger identity invariant, which is important.
• #107812 / salch-cred is not an equivalent substitute: its current-session retargeting changes ownership and has an existing P1 for recreating the #57498 cross-session contamination class.
• Closed-unmerged #64530 / richkapp is useful historical compression-lineage/CAS design evidence, but it is not landed provenance by itself.

On verification: the author supplied useful RED→GREEN local evidence, and another contributor reports 9/9 focused Linux tests on this exact head. Hosted acceptance is nevertheless absent. This branch has one surviving commit (2ee68ade52ccf753dc8ba78e6991d9f465e9eb23), and its Nix run 35178430405, Docker run 35178430468, and CI run 35178430621 all concluded action_required; the CI run created zero jobs. With one surviving commit, exact-head and every-commit acceptance are the same object here, and that object is 0/1 hosted-green.

Current repository state at review: main / actual merge base is 6005aa1fd9aac8b1024ace50fec8cd1c85a04bae; this carrier is 1 ahead / 0 behind and GitHub reports it mergeable. The branch is nicely scoped and the expected-session CAS is worth keeping. The remaining fix is to make the generation proof survive to the actual effect boundary rather than merely checking it earlier in the async path.

[andrexibiza]

P1 — this generation check is still TOCTOU with respect to the route mutation. It proves G here, but the non-compression path later does await self.async_session_store.switch_session(...); AsyncSessionStore offloads the synchronous store call through asyncio.to_thread, so /stop can invalidate G -> G+1 after this check and before the store mutation. In the pure-revoke case the route id can remain unchanged, so expected_session_id still matches and the stale completion repins anyway. The compression path has the same root because several lineage awaits (and the target == current early return) occur after this lone generation check. Please bind expected generation + expected route to the effect/commit boundary itself, and add an event-synchronized regression that invalidates after this check but before route commit; another pre-await check is not sufficient.

[wangtaotaotao95]

Thanks — noted on #113692. I'm happy to converge on whichever shape maintainers prefer; folding the condition into switch_session was only to avoid a second copy of the repoint body, and I can switch to a dedicated method if that reads better to you.

One coverage difference worth flagging while both are open, so the choice is made on coverage rather than style:

A pure session-id CAS catches the case where the route itself moved. #113690's second case is different: /stop or /new revokes the run while the lookup is pending and leaves the routing key unchanged, so the CAS still matches its snapshot and would repoint. That case needs a check on the run token, not on the route.

tests/gateway/test_pinned_route_race.py exercises both boundaries against the real SessionStore and the real generation primitives, parametrized none / revoke / replace. On base it is 1 passed / 2 failed:

E   AssertionError: assert 'test-pinned' == '20260917_031405_073c3e20'   # revoke
E   AssertionError: assert 'test-pinned' == 'test-replacement'           # replace

So this PR carries the run-generation guard in addition to the CAS (_resolve_async_delegation_session re-checks _is_session_run_current after the await, using the token threaded down from _handle_message_with_agent, which the resolver previously never received). revoke is the row that a session-id-only check cannot turn green without mocking the store's return value.

Not blocking either way — just making the difference explicit.

[wangtaotaotao95]

You're right, and thanks for the line-level trace — I reproduced it before changing anything. The check was consumed once after get_session(); AsyncSessionStore.__getattr__ then wraps the mutation in await asyncio.to_thread(...), which is a second scheduling boundary. /stop landing there leaves the routing key untouched, so expected_session_id still matches and the stale pin commits. The compression branch was worse: _resolve_compression_lineage_target() adds three more awaits, and the target == session_entry.session_id fast path returned an entry for injection without sampling the token at all.

Fixed by making the proof effect-bound rather than preflight, in 280eb5da1b:

• switch_session_if_current and advance_compression_session now take an authorize predicate that is sampled under a new SessionStore.routing_authority() lock, immediately before the mutation.
• The run-generation bump takes that same lock (_begin_session_run_generation, _invalidate_session_run_generation), so a boundary cannot interleave with the commit. It lands either before the sample — and the transition refuses — or after it, where the commit was already legitimate.
• The identity fast path samples the token in the same block that returns it, with no await in between.
• switch_session is back to its original two-argument contract; the conditional path is a separate method, and the SQLite side is shared via finish_route_switch.

One thing worth recording, because it changed the shape of the fix: a synchronous commit on the loop is not available here. A repoint is structural, so _save_entry's single-entry fast path explicitly refuses it and the commit takes the full index rewrite (state.db replace + the multi-MB sessions.json mirror). Holding that on the event loop would stall the gateway on every completion delivery, so a shared authority is what makes the proof and the mutation one boundary rather than a non-yielding block.

I also hit a self-deadlock implementing it: _invalidate_session_run_generation nests _begin_session_run_generation, so with a plain threading.Lock the loop blocked on itself. The authority is an RLock.

Witness, as asked — test_revocation_is_sampled_where_the_transition_commits, parametrized over the two shapes that bypass the plain repoint (identity, compression). It injects the boundary at the exact instant the transition samples the token, then asserts the routing key stayed put, the completion was not returned for injection, and the token is no longer current. Red on base:

✗ tests/gateway/test_pinned_route_race.py (1✓ 4✗)
E   AssertionError: the transition never sampled the run token

Green after (5 passed). Full gateway suite: 8864 passed, 6 failed — the same six host-environment files fail on unpatched 6005aa1fd9, so there is no new failure. All 36 files referencing any changed symbol pass (389 passed).

On convergence with #113692: I kept the CAS inside the store rather than a parallel body and named the method switch_session_if_current to match, but it now carries the generation proof as well — a pure session-id CAS cannot close the revoke row, since that case leaves the route unchanged.

action_required on the hosted runs is still the remaining blocker on the acceptance side; nothing I can do from here until a maintainer approves the workflow runs.

[teknium1]

Thanks @wangtaotaotao95. Your change was salvaged into #114765 with your authorship preserved (co-authored); #114765 — fix(gateway): async completion no longer re-pins a route after /stop or /new won the race (#113690, salvage #113692, #113716) — is now merged on main at 3e408dcc2a74, closing issues #113690. Closing this PR in favour of the landed change; if you see a case it does not cover, please open a fresh issue with the repro and tag it.