fix(gateway): interrupt every in-flight API turn on shutdown, not just /v1/runs - #79881
Conversation
…t /v1/runs The shutdown drain ACCOUNTS for API-server work but never INTERRUPTS it. `_drain_active_agents()` folds `_active_api_run_count()` into both its wait loop and its `timed_out` verdict, while `_interrupt_running_agents()` iterates `self._running_agents` only -- a dict no API turn ever enters, because the API server owns its own agent lifecycle. `gateway/run.py` states the gap against itself: "API-server / desk sessions have the same structural gap (NousResearch#63529)." The user-visible result is that every gateway restart with a live API or desktop turn burns the full drain timeout and then runs `_kill_tool_subprocesses("post-interrupt")`, which amputates the turn's tool subprocesses with no cooperative interrupt and no resume marker. There are seven API agent-entry points. Six funnel through `_run_agent()` (both session-chat routes, and `/v1/chat/completions` + `/v1/responses` in streaming and non-streaming form) and are counted by `_inflight_agent_runs`; the seventh, `/v1/runs`, runs its own lifecycle and is counted through `_active_run_tasks`. None of the six has a run_id, so the run_id-keyed `_active_run_agents` cannot reach them, and only two pass `agent_ref` -- which lands in a caller-local list, not a registry. So register once at the single unconditional creation site inside `_run_agent`, beside the existing `_publish_turn_process_ownership()` call, and unregister in the same `finally` that already clears it. That one symmetric pair covers all six callers. The registry is adapter-owned and keyed by object identity, kept separate from `_active_run_agents` because that dict is run_id-keyed and scoped to the public `/v1/runs` stop API. `interrupt_active_runs()` then walks both registries, deduped by identity, so the interrupt set matches the set the drain waits on. The settle window after the interrupt now polls API work as well: the interrupt is cooperative, and without this the window closes the instant `_running_agents` is empty -- which it always is for API turns -- and the tool kill lands on a turn that was asked to stop microseconds earlier.
There was a problem hiding this comment.
Pull request overview
This PR fixes a gateway shutdown/restart hang-and-amputate failure mode by ensuring all API-server in-flight agent turns (not only POST /v1/runs) are cooperatively interrupted when the shutdown drain times out, and by extending the post-interrupt settle window to wait for API work to unwind before tool subprocesses are killed.
Changes:
- Add an adapter-owned registry for
_run_agent()-created agents and exposeAPIServerAdapter.interrupt_active_runs()to interrupt/v1/runsagents plus all other API turn shapes. - Add a duck-typed
GatewayRunner._interrupt_api_server_runs()hook and call it from_interrupt_running_agents(). - Expand the settle window to poll
_active_api_run_count()and add end-to-end regression tests covering interruption + settle behavior.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
gateway/platforms/api_server.py |
Track _run_agent() agents for shutdown interruption and implement interrupt_active_runs() across API entrypoints. |
gateway/run.py |
Add API-server interrupt hook and include API work in the post-interrupt settle window. |
tests/gateway/test_api_server_active_work_drain.py |
Add regression tests proving non-/v1/runs API turns are interrupted and the settle window waits for API work to clear. |
Suppressed comments (3)
gateway/platforms/api_server.py:1520
interrupt_active_runs()snapshots_shutdown_interruptible_agentswithlist(self._shutdown_interruptible_agents.values()). Because registrations/unregistrations happen in an executor thread, this can throw at runtime if the dict is mutated concurrently. Snapshot the values under a lock before iterating (and ensure writers also take the same lock).
for agent in list(self._shutdown_interruptible_agents.values()):
if agent is not None:
# Dedupe by object identity — the two registries are disjoint
# today (/v1/runs runs its own lifecycle, not _run_agent), but
# an agent published to both must still be interrupted once.
gateway/platforms/api_server.py:6105
- Registration into
_shutdown_interruptible_agentshappens inside the_run_agent()executor thread. This needs to take the same lock used byinterrupt_active_runs()snapshots to avoid concurrent-mutation errors during shutdown.
# the _ProviderAuthResolutionError handler below lives here
# rather than in each route. Only two callers pass
# ``agent_ref``, and only /v1/runs has a run_id, so neither
# is a usable hook for the rest.
self._shutdown_interruptible_agents[id(agent)] = agent
gateway/platforms/api_server.py:6239
- Unregistration from
_shutdown_interruptible_agentshappens in the_run_agent()executor thread. This should also take the lock so shutdown-time snapshots/interrupts cannot race with removal and throwRuntimeError.
# Symmetric with the registration above: the turn is
# over, so it must not be interrupted by a later
# shutdown. pop() is a no-op when _create_agent
# succeeded but the turn never reached registration.
self._shutdown_interruptible_agents.pop(id(agent), None)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # the dict holds a strong reference for the life of the turn, so an | ||
| # id() can never be recycled while it is still registered. | ||
| self._shutdown_interruptible_agents: Dict[int, Any] = {} | ||
| # Back-reference to the owning GatewayRunner (set by gateway/run.py) | ||
| # so /api/platforms/{platform}/events can resolve sibling adapters. |
There was a problem hiding this comment.
The threading premise is right — _run() is submitted via loop.run_in_executor(None, _run), so registration and removal genuinely happen on an executor thread while interrupt_active_runs() reads from the event loop thread. The conclusion doesn't follow, though, and the contrast with the lock two lines away is the reason.
_publish_turn_process_ownership (api_server.py:767) holds _TURN_PROCESS_EPOCH_LOCK because it does a compound operation — epoch = next(_TURN_PROCESS_EPOCH_COUNTER) then _TURN_PROCESS_EPOCHS[task_id] = epoch — and _clear_turn_process_ownership (:787) does a check-then-act: if _TURN_PROCESS_EPOCHS.get(task_id) == epoch: del .... Both span two operations and would tear without a lock.
This registry has no compound invariant. Every access is a single operation:
self._shutdown_interruptible_agents[id(agent)] = agent— one atomic store.self._shutdown_interruptible_agents.pop(id(agent), None)— one atomic delete.list(self._shutdown_interruptible_agents.values())— one atomic snapshot.
That last one is the specific concern raised, so to be concrete about why it can't raise RuntimeError: dictionary changed size during iteration: list(view) is a single C-level call that never re-enters the interpreter, so the eval loop has no bytecode boundary at which to switch threads inside it. The values are borrowed PyObject* refs and the key type is int, whose __hash__/__eq__ are C-level and cannot call back into Python either. This is why list(...) snapshotting is the established idiom here rather than an oversight — _interrupt_running_agents (run.py:9189) does list(self._running_agents.items()), and _sweep_orphaned_runs_once does the same with list(self._run_streams_created.items()) (api_server.py:6971 on this branch). I also ran 6 writer threads against 3 snapshotting threads on this exact access pattern for 6s: 0 RuntimeErrors.
Two threads also never contend for the same key: id(agent) is unique per live object, and the dict holds a strong reference for the life of the turn, so an id can't be recycled while it's registered.
On "miss registrations mid-shutdown" — a turn that registers after the snapshot is outside what a lock can fix; it's ordered by the drain, not by mutual exclusion. By the time _interrupt_running_agents() runs, _draining is set and _draining_response() is already rejecting new API turns with 503 (test_drain_refuses_every_agent_start_endpoint covers all five entry routes), and _drain_active_agents() has waited on _active_api_run_count() first.
So I'd rather not add the lock: held across request_hard_interrupt() it would block registering/unregistering executor threads from a shutdown path for the duration of an agent callback, and held only around the snapshot it would be equivalent to what list() already does. Happy to reconsider if there's a specific interleaving I've missed.
alt-glitch
added
type/bug
…window exit Review follow-up for the salvaged NousResearch#79881/NousResearch#63963 stack: the shutdown interrupt fires exactly once, but work can materialize AFTER that one shot on BOTH sibling paths: - a /v1/runs task admitted before the drain populates _active_run_agents only once _create_agent returns (queued-before-agent window); - a _running_agents entry claimed as _AGENT_PENDING_SENTINEL is promoted to the real agent by track_agent() on its own schedule, after the one-shot walk skipped the sentinel. Either way the settle loop waited on work nothing signaled, and the turn went straight to the post-interrupt tool-subprocess kill — the exact amputation the fix exists to avoid, in a rarer window. If any work is still live when the settle loop exits, re-invoke _interrupt_running_agents (which already skips sentinels and folds in the API-server helper) so late-materializing agents on either path get the cooperative interrupt. Regression test drives the real stop() path with an accelerated loop clock and asserts exactly two interrupt signals.
|
Superseded by #80917 (merging via auto-rebase) — your commit is cherry-picked there with authorship preserved, stacked on @necoweb3's #63963 base commit so both contributions survive in git history. The salvage re-resolved the cherry-pick conflicts hunk-by-hunk against current main (your branch predated the pinned-sessions/hygiene/heartbeat work, so a whole-file resolution would have reverted them), kept necoweb3's five tests alongside yours, and added one follow-up: re-signaling interrupts at settle-window exit for agents that materialize after the one-shot interrupt (the queued-before-agent window your PR's docstring called out, plus the |
…window exit Review follow-up for the salvaged #79881/#63963 stack: the shutdown interrupt fires exactly once, but work can materialize AFTER that one shot on BOTH sibling paths: - a /v1/runs task admitted before the drain populates _active_run_agents only once _create_agent returns (queued-before-agent window); - a _running_agents entry claimed as _AGENT_PENDING_SENTINEL is promoted to the real agent by track_agent() on its own schedule, after the one-shot walk skipped the sentinel. Either way the settle loop waited on work nothing signaled, and the turn went straight to the post-interrupt tool-subprocess kill — the exact amputation the fix exists to avoid, in a rarer window. If any work is still live when the settle loop exits, re-invoke _interrupt_running_agents (which already skips sentinels and folds in the API-server helper) so late-materializing agents on either path get the cooperative interrupt. Regression test drives the real stop() path with an accelerated loop clock and asserts exactly two interrupt signals.
…window exit Review follow-up for the salvaged NousResearch#79881/NousResearch#63963 stack: the shutdown interrupt fires exactly once, but work can materialize AFTER that one shot on BOTH sibling paths: - a /v1/runs task admitted before the drain populates _active_run_agents only once _create_agent returns (queued-before-agent window); - a _running_agents entry claimed as _AGENT_PENDING_SENTINEL is promoted to the real agent by track_agent() on its own schedule, after the one-shot walk skipped the sentinel. Either way the settle loop waited on work nothing signaled, and the turn went straight to the post-interrupt tool-subprocess kill — the exact amputation the fix exists to avoid, in a rarer window. If any work is still live when the settle loop exits, re-invoke _interrupt_running_agents (which already skips sentinels and folds in the API-server helper) so late-materializing agents on either path get the cooperative interrupt. Regression test drives the real stop() path with an accelerated loop clock and asserts exactly two interrupt signals.
…window exit Review follow-up for the salvaged NousResearch#79881/NousResearch#63963 stack: the shutdown interrupt fires exactly once, but work can materialize AFTER that one shot on BOTH sibling paths: - a /v1/runs task admitted before the drain populates _active_run_agents only once _create_agent returns (queued-before-agent window); - a _running_agents entry claimed as _AGENT_PENDING_SENTINEL is promoted to the real agent by track_agent() on its own schedule, after the one-shot walk skipped the sentinel. Either way the settle loop waited on work nothing signaled, and the turn went straight to the post-interrupt tool-subprocess kill — the exact amputation the fix exists to avoid, in a rarer window. If any work is still live when the settle loop exits, re-invoke _interrupt_running_agents (which already skips sentinels and folds in the API-server helper) so late-materializing agents on either path get the cooperative interrupt. Regression test drives the real stop() path with an accelerated loop clock and asserts exactly two interrupt signals.
…window exit Review follow-up for the salvaged NousResearch#79881/NousResearch#63963 stack: the shutdown interrupt fires exactly once, but work can materialize AFTER that one shot on BOTH sibling paths: - a /v1/runs task admitted before the drain populates _active_run_agents only once _create_agent returns (queued-before-agent window); - a _running_agents entry claimed as _AGENT_PENDING_SENTINEL is promoted to the real agent by track_agent() on its own schedule, after the one-shot walk skipped the sentinel. Either way the settle loop waited on work nothing signaled, and the turn went straight to the post-interrupt tool-subprocess kill — the exact amputation the fix exists to avoid, in a rarer window. If any work is still live when the settle loop exits, re-invoke _interrupt_running_agents (which already skips sentinels and folds in the API-server helper) so late-materializing agents on either path get the cooperative interrupt. Regression test drives the real stop() path with an accelerated loop clock and asserts exactly two interrupt signals.
supersedes #63963
@necoweb3 found this bug and got the architecture right. Two parts of #63963 are kept here essentially as written, with credit:
GatewayRunner._interrupt_api_server_runs()hook, called from_interrupt_running_agents()— an adapter without the hook is skipped rather than raising mid-shutdown;_active_api_run_count()into the 5s settle window after the interrupt.What this PR adds is the coverage teknium1's review asked for. #63963 has been open 23 days with that one review unanswered, so rather than let the fix sit, this ships the superset.
What does this PR do?
The shutdown drain accounts for API-server work but never interrupts it.
_drain_active_agents()folds_active_api_run_count()into both its wait loop and itstimed_outverdict._interrupt_running_agents()iteratesself._running_agentsonly — a dict no API turn ever enters, because the API server owns its own agent lifecycle. The repo states the gap against itself, ingateway/run.pydirectly above the drain loop:# API-server / desk sessions have the same structural gap (#63529).User-visible symptom: every gateway restart (
/restart,/update, SIGUSR1 — they all funnel through the samestop()) with a live API or desktop turn burns the full drain timeout, then runs_kill_tool_subprocesses("post-interrupt"), which amputates the turn's tool subprocesses. No cooperative interrupt, no chance to unwind. "It hung, then I lost my work."Seven entry points, not one
POST /v1/runs(_handle_runs)_active_run_tasks_active_run_agents[run_id]POST /api/sessions/{id}/chat(_handle_session_chat)_inflight_agent_runsPOST /api/sessions/{id}/chat/stream(_handle_session_chat_stream)_inflight_agent_runsPOST /v1/chat/completions, streaming (_handle_chat_completions)_inflight_agent_runsagent_reflistPOST /v1/chat/completions, non-streaming_inflight_agent_runsPOST /v1/responses, streaming (_handle_responses)_inflight_agent_runsagent_reflistPOST /v1/responses, non-streaming_inflight_agent_runsRows 2–7 are the six
_run_agent()callers._inflight_agent_runs += 1 / -= 1wraps every one of them, so they all hold the drain open — and none of them has arun_id, so the run_id-keyed_active_run_agentsstructurally cannot reach them. Only rows 4 and 6 passagent_ref, and that lands in a caller-local one-element list, never a registry, so it is not a usable hook either.That is broader than the review note, which named the synchronous session-chat path specifically.
The approach: one registration site
Rather than touch six call sites, register at the single unconditional agent-creation point inside
_run_agent's inner_run, immediately beside the existing_publish_turn_process_ownership(agent, effective_task_id)call, and unregister in the same existingfinallythat already calls_clear_turn_process_ownership(agent). One symmetric pair, six callers covered, mirroring a lifecycle idiom already accepted in that exact function — and the same "handle it here, once, covers every_run_agent()caller" reasoning the_ProviderAuthResolutionErrorbranch a few lines below already uses.The new registry is deliberately not
_active_run_agents:_active_run_agentsisrun_id-keyed and scoped to the publicPOST /v1/runs/{run_id}/stopAPI; shutdown needs everything, and six of the seven paths have norun_id._active_run_agents[run_id] = agentinto_run_agentfor its own purposes.interrupt_active_runs()dedupes by object identity, so if that lands and an agent ends up in both registries, it is still interrupted exactly once._pending_agent_requestsis deliberately not covered, and the docstring says so: it counts admitted requests that have not constructed an agent yet, so there is no object to interrupt.Cron is out of scope here — cron agents live in
cron/scheduler.py's own module-level registry keyed byjob_id, entirely outside this adapter, and #43085 covers that surface.teknium1's
### Suggested changeson #63963, as an acceptance checklist✅ One registration inside
_run_agentcovers all six callers by construction —_inflight_agent_runsand the registry are incremented/populated around the same body, so they cannot drift.interrupt_active_runs()walks that registry unioned with_active_run_agents, deduped by identity.✅ Two of them, both end-to-end through the real aiohttp handlers with the turn genuinely parked inside
run_conversationon an executor thread:test_chat_completions_turn_is_interruptedandtest_session_chat_sse_turn_is_interrupted. Each asserts the drain sees the turn (_active_api_run_count() == 1), that it is not in_running_agents, and then thatagent.interruptis called exactly once with the shutdown reason.(
request_hard_interrupt()is used rather than a bareagent.interrupt(reason), matching_interrupt_running_agentsand the/v1/runsstop handler. For an unspecced mock it falls through toagent.interrupt(reason), so the assertion above is on the exact ABI named in the review.)Related Issue
Context: #63529 (closed — its accounting half landed; this is the interrupt half). Not auto-closing it.
Type of Change
Changes Made
gateway/platforms/api_server.py__init__: add_shutdown_interruptible_agents, an adapter-owned registry of agents currently inside_run_agent(), keyed byid(). It holds a strong reference for the life of the turn, so anid()cannot be recycled while registered._run_agent()inner_run(): register beside_publish_turn_process_ownership(); unregister in the existingfinallybeside_clear_turn_process_ownership().pop(..., None)is a no-op if_create_agentsucceeded but the turn never reached registration.interrupt_active_runs(reason) -> int: walks_active_run_agents∪_shutdown_interruptible_agents, dedupes by object identity, callsrequest_hard_interruptper agent, swallows and logs per-agent failures so one torn-down agent cannot strand the rest, and returns the number that accepted an interrupt.gateway/run.py_interrupt_api_server_runs(reason)— duck-typed adapter hook (from fix(gateway): interrupt api_server runs on shutdown timeout #63963)._interrupt_running_agents()calls it after its_running_agentsloop (from fix(gateway): interrupt api_server runs on shutdown timeout #63963)._active_api_run_count()as well as_running_agents(from fix(gateway): interrupt api_server runs on shutdown timeout #63963).tests/gateway/test_api_server_active_work_drain.py— 9 new tests.How to Test
14 passed (5 pre-existing + 9 new).
Red-before / green-after. Every hunk was reverted independently against clean
origin/mainand the tests re-run in their final location:origin/mainapi_server.pyfixed,run.pyatorigin/mainrun.pyfixed,api_server.pyatorigin/mainThe last row but one is the point: the settle-window change is independently load-bearing, not a cosmetic tag-along.
test_settle_window_waits_for_interrupted_api_workdrives the realrunner.stop()down its drain-timeout path and asserts, via a deterministic poll counter rather than any wall-clock measurement, that the post-interrupt tool kill ran only after the interrupted API work had settled.Adjacent suites, all green on this branch:
Manual: start the gateway with the API server enabled, send a long-running turn to
POST /v1/chat/completions(non-streaming) or to the desktop app, then/restart. Before: the drain runs to its full timeout and the turn's tool subprocesses are killed from under it. After: the turn is interrupted cooperatively as soon as the drain times out, and the settle window lets it unwind before the tool kill.Checklist
Code
fix(scope):,feat(scope):, etc.)pytest tests/ -qand all tests pass — ran the touched file plus the adjacent gateway/API/shutdown suites (179 tests, all green); the full suite is not green on a stock local checkout, so I did not claim itDocumentation & Housekeeping
docs/, docstrings) — docstrings only; no user-facing docs describe this internal shutdown pathcli-config.yaml.exampleif I added/changed config keys — N/A, no new configCONTRIBUTING.mdorAGENTS.mdif I changed architecture or workflows — N/ARelated / Positioning
_active_run_agents[run_id]into_run_agentfor that. This PR deliberately does not touch_active_run_agents, and its identity dedupe means the two compose without conflict whichever lands first.cron/scheduler.py. Disjoint registry, disjoint population; intentionally not folded in here.