fix: tighten gateway interrupt salvage follow-ups
Follow-up on top of the helix4u #12388 cherry-picks: - make deferred post-delivery callbacks generation-aware end-to-end so stale runs cannot clear callbacks registered by a fresher run for the same session - bind callback ownership to the active session event at run start and snapshot that generation inside base adapter processing so later event mutation cannot retarget cleanup - pass run_generation through proxy mode and drop stale proxy streams / final results the same way local runs are dropped - centralize stop/new interrupt cleanup into one helper and replace the open-coded branches with shared logic - unify internal control interrupt reason strings via shared constants - remove the return from base.py's finally block so cleanup no longer swallows cancellation/exception flow - add focused regressions for generation forwarding, proxy stale suppression, and newer-callback preservation This addresses all review findings from the initial #12388 review while keeping the fix scoped to stale-output/typing-loop interrupt handling.
This commit is contained in:
@@ -19,6 +19,7 @@ def _make_runner(proxy_url=None):
|
||||
runner.config = MagicMock()
|
||||
runner.config.streaming = StreamingConfig()
|
||||
runner._running_agents = {}
|
||||
runner._session_run_generation = {}
|
||||
runner._session_model_overrides = {}
|
||||
runner._agent_cache = {}
|
||||
runner._agent_cache_lock = None
|
||||
@@ -160,10 +161,12 @@ class TestRunAgentProxyDispatch:
|
||||
source=source,
|
||||
session_id="test-session-123",
|
||||
session_key="test-key",
|
||||
run_generation=7,
|
||||
)
|
||||
|
||||
assert result["final_response"] == "Hello from remote!"
|
||||
runner._run_agent_via_proxy.assert_called_once()
|
||||
assert runner._run_agent_via_proxy.call_args.kwargs["run_generation"] == 7
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_agent_skips_proxy_when_not_configured(self, monkeypatch):
|
||||
@@ -370,6 +373,40 @@ class TestRunAgentViaProxy:
|
||||
assert "session_id" in result
|
||||
assert result["session_id"] == "sess-123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_stale_generation_returns_empty_result(self, monkeypatch):
|
||||
monkeypatch.setenv("GATEWAY_PROXY_URL", "http://host:8642")
|
||||
monkeypatch.delenv("GATEWAY_PROXY_KEY", raising=False)
|
||||
runner = _make_runner()
|
||||
source = _make_source()
|
||||
runner._session_run_generation["test-key"] = 2
|
||||
|
||||
resp = _FakeSSEResponse(
|
||||
status=200,
|
||||
sse_chunks=[
|
||||
'data: {"choices":[{"delta":{"content":"stale"}}]}\n\n',
|
||||
"data: [DONE]\n\n",
|
||||
],
|
||||
)
|
||||
session = _FakeSession(resp)
|
||||
|
||||
with patch("gateway.run._load_gateway_config", return_value={}):
|
||||
with _patch_aiohttp(session):
|
||||
with patch("aiohttp.ClientTimeout"):
|
||||
result = await runner._run_agent_via_proxy(
|
||||
message="hi",
|
||||
context_prompt="",
|
||||
history=[],
|
||||
source=source,
|
||||
session_id="sess-123",
|
||||
session_key="test-key",
|
||||
run_generation=1,
|
||||
)
|
||||
|
||||
assert result["final_response"] == ""
|
||||
assert result["messages"] == []
|
||||
assert result["api_calls"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_auth_header_without_key(self, monkeypatch):
|
||||
monkeypatch.setenv("GATEWAY_PROXY_URL", "http://host:8642")
|
||||
|
||||
@@ -270,6 +270,75 @@ async def test_handle_message_discards_stale_result_after_session_invalidation(m
|
||||
assert session_key not in runner.adapters[Platform.TELEGRAM]._post_delivery_callbacks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_message_stale_result_keeps_newer_generation_callback(monkeypatch):
|
||||
import gateway.run as gateway_run
|
||||
|
||||
class _Adapter:
|
||||
def __init__(self):
|
||||
self._post_delivery_callbacks = {}
|
||||
|
||||
async def send(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
def pop_post_delivery_callback(self, session_key, *, generation=None):
|
||||
entry = self._post_delivery_callbacks.get(session_key)
|
||||
if entry is None:
|
||||
return None
|
||||
if isinstance(entry, tuple):
|
||||
entry_generation, callback = entry
|
||||
if generation is not None and entry_generation != generation:
|
||||
return None
|
||||
self._post_delivery_callbacks.pop(session_key, None)
|
||||
return callback
|
||||
if generation is not None:
|
||||
return None
|
||||
return self._post_delivery_callbacks.pop(session_key, None)
|
||||
|
||||
session_entry = SessionEntry(
|
||||
session_key=build_session_key(_make_source()),
|
||||
session_id="sess-1",
|
||||
created_at=datetime.now(),
|
||||
updated_at=datetime.now(),
|
||||
platform=Platform.TELEGRAM,
|
||||
chat_type="dm",
|
||||
)
|
||||
runner = _make_runner(session_entry)
|
||||
runner.session_store.load_transcript.return_value = [{"role": "user", "content": "earlier"}]
|
||||
session_key = session_entry.session_key
|
||||
adapter = _Adapter()
|
||||
runner.adapters[Platform.TELEGRAM] = adapter
|
||||
|
||||
async def _stale_result(**kwargs):
|
||||
# Simulate a newer run claiming the callback slot before the stale run unwinds.
|
||||
runner._session_run_generation[session_key] = 2
|
||||
adapter._post_delivery_callbacks[session_key] = (2, lambda: None)
|
||||
return {
|
||||
"final_response": "late reply",
|
||||
"messages": [],
|
||||
"tools": [],
|
||||
"history_offset": 0,
|
||||
"last_prompt_tokens": 80,
|
||||
"input_tokens": 120,
|
||||
"output_tokens": 45,
|
||||
"model": "openai/test-model",
|
||||
}
|
||||
|
||||
runner._run_agent = AsyncMock(side_effect=_stale_result)
|
||||
|
||||
monkeypatch.setattr(gateway_run, "_resolve_runtime_agent_kwargs", lambda: {"api_key": "***"})
|
||||
monkeypatch.setattr(
|
||||
"agent.model_metadata.get_model_context_length",
|
||||
lambda *_args, **_kwargs: 100000,
|
||||
)
|
||||
|
||||
result = await runner._handle_message(_make_event("hello"))
|
||||
|
||||
assert result is None
|
||||
assert session_key in adapter._post_delivery_callbacks
|
||||
assert adapter._post_delivery_callbacks[session_key][0] == 2
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_status_command_bypasses_active_session_guard():
|
||||
|
||||
Reference in New Issue
Block a user