fix(interrupt): extend post-worker /stop guard to Bedrock streaming path
The salvaged fix added a post-worker _interrupt_requested re-check to the main OpenAI/Anthropic streaming poll loop. The Bedrock Converse poll loop (interruptible_streaming_api_call, api_mode='bedrock_converse') has the same bug class: its worker calls stream_converse_with_callbacks(on_interrupt_check= ...), which breaks out of the event loop on interrupt and returns a PARTIAL response WITHOUT raising (bedrock_adapter.py). The worker sets result[ 'response'] and exits with _interrupt_requested still True, so the in-loop raise never fires and the poll loop returns the partial — silently swallowing /stop on Bedrock exactly as it was on the paths the salvaged commit fixed. Add the identical post-worker re-check before the Bedrock loop's return. The non-streaming loop (interruptible_api_call) is structurally immune: its worker's only early return fires off _request_cancelled, which is set by the main loop immediately before it raises in-loop, so no swallow window exists. Guard test flips _interrupt_requested True mid-stream (after the pre-flight check) and asserts InterruptedError is raised; verified RED without the fix (DID NOT RAISE) and GREEN with it.
This commit is contained in:
parent
c2c73605e0
commit
144457d801
2 changed files with 98 additions and 0 deletions
|
|
@ -1881,6 +1881,16 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
t.join(timeout=0.3)
|
||||
if agent._interrupt_requested:
|
||||
raise InterruptedError("Agent interrupted during Bedrock API call")
|
||||
# Worker exited before the poll loop observed the interrupt flag. The
|
||||
# Bedrock stream callback breaks out and returns a PARTIAL response
|
||||
# without raising on interrupt (see bedrock_adapter.py
|
||||
# stream_converse_with_callbacks / on_interrupt_check), so result[
|
||||
# "response"] is populated with error=None and the in-loop raise above
|
||||
# never fires. Re-check here so /stop is not silently swallowed on the
|
||||
# Bedrock path — mirrors the post-worker guard on the main streaming
|
||||
# loop. (#59999 area)
|
||||
if agent._interrupt_requested:
|
||||
raise InterruptedError("Agent interrupted during Bedrock API call (post-worker)")
|
||||
if result["error"] is not None:
|
||||
raise result["error"]
|
||||
return result["response"]
|
||||
|
|
|
|||
88
tests/agent/test_bedrock_interrupt_post_worker.py
Normal file
88
tests/agent/test_bedrock_interrupt_post_worker.py
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
"""Regression: /stop must not be swallowed on the Bedrock streaming path.
|
||||
|
||||
Companion to the OpenAI/Anthropic streaming post-worker guard. The Bedrock
|
||||
Converse stream callback (bedrock_adapter.stream_converse_with_callbacks) breaks
|
||||
out of its event loop on interrupt and returns a PARTIAL response WITHOUT
|
||||
raising. The worker thread then sets result["response"] and exits cleanly with
|
||||
agent._interrupt_requested still True. Without a post-worker re-check in the
|
||||
poll loop, interruptible_streaming_api_call would return that partial response
|
||||
and silently swallow the /stop signal.
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agent import chat_completion_helpers as cch
|
||||
|
||||
|
||||
class _FakeAgent:
|
||||
api_mode = "bedrock_converse"
|
||||
_interrupt_requested = False # not interrupted at entry (passes pre-flight)
|
||||
_disable_streaming = False
|
||||
reasoning_callback = None
|
||||
stream_delta_callback = None
|
||||
|
||||
def _has_stream_consumers(self):
|
||||
return False
|
||||
|
||||
def _fire_stream_delta(self, text):
|
||||
pass
|
||||
|
||||
def _fire_tool_gen_started(self, name):
|
||||
pass
|
||||
|
||||
def _fire_reasoning_delta(self, text):
|
||||
pass
|
||||
|
||||
def _safe_print(self, *a, **k):
|
||||
pass
|
||||
|
||||
|
||||
def test_bedrock_stream_interrupt_not_swallowed_post_worker():
|
||||
"""A /stop arriving MID-stream: the pre-flight check (top of function) has
|
||||
already passed, the worker's stream callback breaks and returns a partial
|
||||
response WITHOUT raising, leaving _interrupt_requested True. The post-worker
|
||||
re-check must raise InterruptedError instead of returning the partial."""
|
||||
agent = _FakeAgent()
|
||||
|
||||
partial = SimpleNamespace(choices=[], usage=None, stop_reason="interrupted")
|
||||
|
||||
# Simulate the real adapter: on interrupt it breaks out and returns a
|
||||
# partial response WITHOUT raising. Flip the interrupt flag here to model
|
||||
# /stop arriving mid-stream (after the pre-flight check, during the worker).
|
||||
def _fake_stream(*args, **kwargs):
|
||||
agent._interrupt_requested = True
|
||||
return partial
|
||||
|
||||
fake_client = SimpleNamespace(converse_stream=lambda **kw: {"stream": []})
|
||||
|
||||
with patch("agent.bedrock_adapter._get_bedrock_runtime_client", return_value=fake_client), \
|
||||
patch("agent.bedrock_adapter.stream_converse_with_callbacks", side_effect=_fake_stream), \
|
||||
patch("agent.bedrock_adapter.normalize_converse_response", side_effect=lambda r: r), \
|
||||
patch("agent.bedrock_adapter.is_stale_connection_error", return_value=False), \
|
||||
patch("agent.bedrock_adapter.is_streaming_access_denied_error", return_value=False), \
|
||||
patch("agent.bedrock_adapter.invalidate_runtime_client", lambda *a, **k: None):
|
||||
api_kwargs = {"__bedrock_region__": "us-east-1", "__bedrock_converse__": True}
|
||||
with pytest.raises(InterruptedError):
|
||||
cch.interruptible_streaming_api_call(agent, api_kwargs)
|
||||
|
||||
|
||||
def test_bedrock_stream_returns_normally_when_not_interrupted():
|
||||
"""Sanity: with no interrupt, the same path returns the response (guard
|
||||
must not fire spuriously)."""
|
||||
agent = _FakeAgent()
|
||||
agent._interrupt_requested = False
|
||||
|
||||
resp = SimpleNamespace(choices=[], usage=None, stop_reason="end_turn")
|
||||
fake_client = SimpleNamespace(converse_stream=lambda **kw: {"stream": []})
|
||||
|
||||
with patch("agent.bedrock_adapter._get_bedrock_runtime_client", return_value=fake_client), \
|
||||
patch("agent.bedrock_adapter.stream_converse_with_callbacks", return_value=resp), \
|
||||
patch("agent.bedrock_adapter.normalize_converse_response", side_effect=lambda r: r), \
|
||||
patch("agent.bedrock_adapter.is_stale_connection_error", return_value=False), \
|
||||
patch("agent.bedrock_adapter.is_streaming_access_denied_error", return_value=False), \
|
||||
patch("agent.bedrock_adapter.invalidate_runtime_client", lambda *a, **k: None):
|
||||
api_kwargs = {"__bedrock_region__": "us-east-1", "__bedrock_converse__": True}
|
||||
out = cch.interruptible_streaming_api_call(agent, api_kwargs)
|
||||
assert out is resp
|
||||
Loading…
Add table
Add a link
Reference in a new issue