fix(streaming): gate the interrupt worker join on live Relay managed execution
The unconditional 2s join before InterruptedError delayed interrupt detection when Relay managed execution was not active (CI: tests/run_agent/test_interrupt_propagation.py — detection took 2.34s against a <1.0s budget, because the mocked worker sleeps 5s and there is no Relay scope to unwind). Extract the join into _join_worker_for_relay_teardown(), which no-ops unless a Relay runtime exists AND managed execution consumers are registered — the only case where an orphaned physical scope can corrupt the LIFO stack (#81521). Applied at all three interrupt sites (streaming, non-streaming, Bedrock streaming). The regression test now simulates a live runtime so the join path stays covered.
This commit is contained in:
@@ -71,6 +71,35 @@ def _context_thread_target(callback):
|
||||
return lambda: context.run(callback)
|
||||
|
||||
|
||||
def _join_worker_for_relay_teardown(worker, *, label: str) -> None:
|
||||
"""Bounded worker join before raising InterruptedError (#81521).
|
||||
|
||||
Raising immediately lets turn teardown (finish_logical_calls /
|
||||
end_turn / close_session) race a still-open Relay physical LLM scope
|
||||
and corrupt the LIFO stack — "scope handle is not at the top of the
|
||||
stack" → CLI EIO / redraw storm. Only joins when Relay managed
|
||||
execution is actually live: when no Relay consumers are registered
|
||||
there is no scope to unwind, and the join would just delay interrupt
|
||||
detection (tests/run_agent/test_interrupt_propagation.py).
|
||||
"""
|
||||
try:
|
||||
from agent import relay_runtime
|
||||
|
||||
runtime = relay_runtime.get_runtime(create=False)
|
||||
if runtime is None or not runtime.managed_execution_enabled():
|
||||
return
|
||||
except Exception:
|
||||
return
|
||||
worker.join(timeout=2.0)
|
||||
if worker.is_alive():
|
||||
logger.warning(
|
||||
"%s worker still alive after interrupt abort (2.0s join "
|
||||
"timeout); Relay teardown will best-effort drain orphaned "
|
||||
"scopes (#81521).",
|
||||
label,
|
||||
)
|
||||
|
||||
|
||||
def _ra():
|
||||
"""Lazy ``run_agent`` reference.
|
||||
|
||||
@@ -1724,15 +1753,9 @@ def interruptible_api_call(agent, api_kwargs: dict):
|
||||
# #81521 (sibling of the streaming-path fix): wait for the worker
|
||||
# to unwind Relay-managed scopes before surfacing
|
||||
# InterruptedError, so turn teardown cannot race a still-open
|
||||
# physical scope and corrupt the LIFO stack.
|
||||
t.join(timeout=2.0)
|
||||
if t.is_alive():
|
||||
logger.warning(
|
||||
"Non-streaming worker still alive after interrupt abort "
|
||||
"(%.1fs join timeout); Relay teardown will best-effort "
|
||||
"drain orphaned scopes (#81521).",
|
||||
2.0,
|
||||
)
|
||||
# physical scope and corrupt the LIFO stack. No-op when Relay
|
||||
# managed execution is not live.
|
||||
_join_worker_for_relay_teardown(t, label="Non-streaming")
|
||||
raise InterruptedError("Agent interrupted during API call")
|
||||
if result["error"] is not None:
|
||||
raise result["error"]
|
||||
@@ -3426,16 +3449,9 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
||||
# #81521 (sibling of the main streaming-path fix): give
|
||||
# the Bedrock worker a bounded window to unwind its
|
||||
# Relay-managed stream scopes before surfacing
|
||||
# InterruptedError, so turn teardown cannot race a
|
||||
# still-open physical scope and corrupt the LIFO stack.
|
||||
t.join(timeout=2.0)
|
||||
if t.is_alive():
|
||||
logger.warning(
|
||||
"Bedrock streaming worker still alive after "
|
||||
"interrupt (%.1fs join timeout); Relay teardown "
|
||||
"will best-effort drain orphaned scopes (#81521).",
|
||||
2.0,
|
||||
)
|
||||
# InterruptedError. No-op when Relay managed execution
|
||||
# is not live.
|
||||
_join_worker_for_relay_teardown(t, label="Bedrock streaming")
|
||||
raise InterruptedError("Agent interrupted during Bedrock API call")
|
||||
# Liveness watchdog: no Bedrock event for longer than the stale
|
||||
# timeout means the stream has wedged (open socket, keep-alives but
|
||||
@@ -5098,15 +5114,9 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
||||
# (finish_logical_calls / end_turn / close_session) race a
|
||||
# still-open physical scope and corrupt the LIFO stack —
|
||||
# "scope handle is not at the top of the stack" → CLI EIO /
|
||||
# redraw storm (#81521).
|
||||
t.join(timeout=2.0)
|
||||
if t.is_alive():
|
||||
logger.warning(
|
||||
"Streaming worker still alive after interrupt abort "
|
||||
"(%.1fs join timeout); Relay teardown will best-effort "
|
||||
"drain orphaned scopes (#81521).",
|
||||
2.0,
|
||||
)
|
||||
# redraw storm (#81521). No-op when Relay managed execution
|
||||
# is not live.
|
||||
_join_worker_for_relay_teardown(t, label="Streaming")
|
||||
raise InterruptedError("Agent interrupted during streaming API call")
|
||||
# Worker thread exited before the main thread's poll loop could check
|
||||
# the interrupt flag. If the worker returned early due to an interrupt
|
||||
|
||||
@@ -272,6 +272,21 @@ class TestStreamInterruptJoinsWorkerBeforeRaise:
|
||||
|
||||
monkeypatch.setattr(threading.Thread, "join", spy_join)
|
||||
|
||||
# The join is gated on Relay managed execution being live (it is
|
||||
# pointless — and delays interrupt detection — when no Relay
|
||||
# consumers are registered). Simulate a live runtime.
|
||||
from agent import relay_runtime as rr
|
||||
|
||||
monkeypatch.setattr(
|
||||
rr,
|
||||
"get_runtime",
|
||||
lambda *a, **k: SimpleNamespace(
|
||||
managed_execution_enabled=lambda: True,
|
||||
get_session=lambda *a2, **k2: None,
|
||||
ensure_session=lambda *a2, **k2: None,
|
||||
),
|
||||
)
|
||||
|
||||
class HangUntilClosedStream:
|
||||
response = SimpleNamespace(headers={})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user