7242811de4
test_stale_abort_is_atomic_with_holder_read asserts the post-stale-kill death is a no-retry partial (a single stream_error_cleanup close). Its agent had no stream consumer, so the "hello" delta was never recorded as delivered; with the parent commit a death after undelivered deltas is reclassified as an undelivered failure and retried on a fresh client, which added stream_retry_cleanup + a second attempt's close. The test's premise is "died AFTER tokens were delivered", so register a consumer instead of weakening the new classification.
413 lines
15 KiB
Python
413 lines
15 KiB
Python
"""Races between the request-client reuse cache and the abort machinery.
|
|
|
|
Invariants pinned here:
|
|
|
|
1. A worker-side interrupt that breaks out of the SSE chunk loop must close
|
|
the half-read stream first (owner thread). Abandoning it leaves its
|
|
connection permanently checked out of the httpx pool while the partial
|
|
response makes the worker's finally report a reuse-reason close, caching
|
|
the client with the leaked connection — one more per interrupt until the
|
|
pool hits ``max_connections`` and every request dies with PoolTimeout.
|
|
If the close fails, the slot must be poisoned so the finally really
|
|
closes the pool.
|
|
|
|
2. The stranger-thread abort (stale detector / interrupt loop) must read
|
|
the holder and fire the abort atomically, under ``request_client_lock``.
|
|
An abort landing after the lock is released races the worker's finally,
|
|
which can pop + cache the client and let the NEXT call check it out —
|
|
the late abort would then poison the slot and shut down an innocent
|
|
in-flight request's sockets.
|
|
|
|
3. The same atomicity contract holds at the third holder-abort site:
|
|
``direct_api_call``'s ``_abort_active_request`` (the cron inline path,
|
|
reached cross-thread via ``AIAgent.interrupt()``).
|
|
|
|
4. ``run_codex_stream``'s finally must poison the reuse slot when
|
|
``event_stream.close()`` fails; otherwise a failed close caches the
|
|
client with a connection still checked out of the pool.
|
|
|
|
5. Relay's managed stream wrapper must preserve the same close failure and
|
|
request-client identity rather than masking the signal used to poison the
|
|
slot.
|
|
"""
|
|
from contextlib import contextmanager
|
|
import threading
|
|
import time
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
def _make_agent():
|
|
from run_agent import AIAgent
|
|
|
|
agent = AIAgent(
|
|
api_key="test-key",
|
|
base_url="https://example.com/v1",
|
|
model="test/model",
|
|
quiet_mode=True,
|
|
skip_context_files=True,
|
|
skip_memory=True,
|
|
)
|
|
agent.api_mode = "chat_completions"
|
|
return agent
|
|
|
|
|
|
@contextmanager
|
|
def _managed_relay_turn(agent, tmp_path, monkeypatch):
|
|
pytest.importorskip("nemo_relay")
|
|
from agent import relay_runtime
|
|
|
|
monkeypatch.setenv("HERMES_HOME", str(tmp_path / "profile"))
|
|
relay_runtime._reset_for_tests()
|
|
lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation(
|
|
profile_key=relay_runtime.current_profile_key(),
|
|
session_id=agent.session_id,
|
|
platform="cli",
|
|
)
|
|
turn = relay_runtime.SESSION_COORDINATOR.begin_turn(
|
|
lease,
|
|
turn_id="request-client-reuse-turn",
|
|
task_id="request-client-reuse-task",
|
|
)
|
|
lease.host.retain_managed_execution("test.request_client_reuse")
|
|
try:
|
|
yield
|
|
finally:
|
|
lease.host.release_managed_execution("test.request_client_reuse")
|
|
relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success")
|
|
relay_runtime.SESSION_COORDINATOR.release_conversation(lease)
|
|
relay_runtime._reset_for_tests()
|
|
|
|
|
|
def _chunk(content=None, finish_reason=None):
|
|
return SimpleNamespace(
|
|
choices=[
|
|
SimpleNamespace(
|
|
index=0,
|
|
delta=SimpleNamespace(
|
|
content=content,
|
|
tool_calls=None,
|
|
reasoning_content=None,
|
|
reasoning=None,
|
|
),
|
|
finish_reason=finish_reason,
|
|
)
|
|
],
|
|
model="test/model",
|
|
usage=None,
|
|
)
|
|
|
|
|
|
class _FakeStream:
|
|
"""SSE stream stand-in.
|
|
|
|
Deliberately has NO ``choices`` attribute so ``_call_chat_completions``
|
|
treats it as a genuine token stream (a MagicMock would auto-create
|
|
``choices`` and get misread as a completed response object).
|
|
"""
|
|
|
|
response = None
|
|
|
|
def __init__(self, chunk_iter_factory, close_raises=False):
|
|
self._factory = chunk_iter_factory
|
|
self._close_raises = close_raises
|
|
self.close_calls = 0
|
|
|
|
def __iter__(self):
|
|
return self._factory()
|
|
|
|
def close(self):
|
|
self.close_calls += 1
|
|
if self._close_raises:
|
|
raise RuntimeError("close failed")
|
|
|
|
|
|
def _mock_wire_client(stream):
|
|
client = MagicMock()
|
|
client.chat.completions.create.return_value = stream
|
|
return client
|
|
|
|
|
|
@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
|
|
def test_worker_interrupt_break_closes_stream():
|
|
"""Interrupt noticed between chunks must close the half-read stream.
|
|
|
|
Without the close, the connection stays checked out of the httpx pool
|
|
while the partial-response finally caches the client for reuse — the
|
|
leak that eventually exhausted the pool (PoolTimeout on every request).
|
|
"""
|
|
agent = _make_agent()
|
|
|
|
def chunks():
|
|
yield _chunk(content="partial ")
|
|
# /stop arrives while the provider is still streaming.
|
|
agent._interrupt_requested = True
|
|
yield _chunk(content="never processed")
|
|
|
|
stream = _FakeStream(chunks)
|
|
|
|
with patch.object(
|
|
agent, "_create_request_openai_client", return_value=_mock_wire_client(stream)
|
|
), patch.object(agent, "_close_request_openai_client"):
|
|
with pytest.raises(InterruptedError):
|
|
agent._interruptible_streaming_api_call({})
|
|
|
|
assert stream.close_calls == 1
|
|
|
|
|
|
|
|
|
|
def test_relay_managed_close_failure_poisons_request_client(tmp_path, monkeypatch):
|
|
"""Relay must preserve close failures needed by the reuse-slot guard."""
|
|
agent = _make_agent()
|
|
|
|
def chunks():
|
|
yield _chunk(content="partial ")
|
|
agent._interrupt_requested = True
|
|
yield _chunk(content="never processed")
|
|
|
|
stream = _FakeStream(chunks, close_raises=True)
|
|
request_client = _mock_wire_client(stream)
|
|
abort_reasons = []
|
|
|
|
with _managed_relay_turn(agent, tmp_path, monkeypatch), patch.object(
|
|
agent, "_create_request_openai_client", return_value=request_client
|
|
), patch.object(agent, "_close_request_openai_client"), patch.object(
|
|
agent,
|
|
"_abort_request_openai_client",
|
|
side_effect=lambda client, *, reason: abort_reasons.append(
|
|
(client, reason)
|
|
),
|
|
):
|
|
with pytest.raises(InterruptedError):
|
|
agent._interruptible_streaming_api_call(
|
|
{
|
|
"model": "test/model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
}
|
|
)
|
|
|
|
assert stream.close_calls == 1
|
|
assert abort_reasons == [(request_client, "interrupt_stream_close_failed")]
|
|
|
|
|
|
def test_stale_abort_is_atomic_with_holder_read(monkeypatch):
|
|
"""The stranger-thread abort must complete before the worker's finally
|
|
can pop + cache the client.
|
|
|
|
The abort runs under ``request_client_lock``; a worker that finishes
|
|
while the abort is in flight must block in its finally until the abort
|
|
returns. Without that, the worker could cache the client (and the next
|
|
call check it out) between the holder read and the abort — the abort
|
|
then killed an innocent request's sockets.
|
|
"""
|
|
monkeypatch.setenv("HERMES_STREAM_STALE_TIMEOUT", "0.05")
|
|
agent = _make_agent()
|
|
# A display consumer, so "hello" counts as DELIVERED text: the death after
|
|
# the stale kill is then a no-retry partial (one stream_error_cleanup close),
|
|
# not an undelivered failure that retries on a fresh client.
|
|
agent.stream_delta_callback = lambda text: None
|
|
|
|
allow_finish = threading.Event()
|
|
worker_close_reasons = []
|
|
abort_reasons = []
|
|
observed = {"worker_finished_during_abort": None}
|
|
|
|
def chunks():
|
|
yield _chunk(content="hello")
|
|
# Stall long enough for the stale detector to fire, then finish
|
|
# cleanly the moment the abort (below) unblocks us.
|
|
allow_finish.wait(timeout=5.0)
|
|
yield _chunk(finish_reason="stop")
|
|
|
|
stream = _FakeStream(chunks)
|
|
|
|
def fake_abort(client, *, reason):
|
|
abort_reasons.append(reason)
|
|
# Let the worker race toward its finally while the abort is still
|
|
# in flight. Under the fix it must block on the holder lock, so the
|
|
# owner-side close cannot land until this abort returns.
|
|
allow_finish.set()
|
|
deadline = time.time() + 0.6
|
|
while time.time() < deadline and not worker_close_reasons:
|
|
time.sleep(0.02)
|
|
observed["worker_finished_during_abort"] = bool(worker_close_reasons)
|
|
|
|
with patch.object(
|
|
agent, "_create_request_openai_client", return_value=_mock_wire_client(stream)
|
|
), patch.object(
|
|
agent,
|
|
"_close_request_openai_client",
|
|
side_effect=lambda client, *, reason: worker_close_reasons.append(reason),
|
|
), patch.object(
|
|
agent, "_abort_request_openai_client", side_effect=fake_abort
|
|
), patch.object(
|
|
agent, "_replace_primary_openai_client"
|
|
):
|
|
response = agent._interruptible_streaming_api_call({})
|
|
|
|
assert response is not None
|
|
assert "stale_stream_kill" in abort_reasons
|
|
# The atomicity contract: no owner-side close slipped in mid-abort.
|
|
assert observed["worker_finished_during_abort"] is False
|
|
# ...and the worker's own finally still performed its close afterwards.
|
|
# The stale kill cancels the stream attempt before aborting, so the
|
|
# worker's late clean finish is treated as a superseded stream and the
|
|
# finally reports the error-cleanup reason — really closing the
|
|
# socket-aborted (poisoned) client instead of caching it for reuse.
|
|
assert worker_close_reasons == ["stream_error_cleanup"]
|
|
|
|
|
|
def test_direct_api_call_abort_is_atomic_with_holder_read():
|
|
"""Same atomicity contract for the cron inline path's holder abort.
|
|
|
|
``direct_api_call``'s ``_abort_active_request`` (fired cross-thread via
|
|
``AIAgent.interrupt()`` -> ``_active_request_abort``) must complete
|
|
before the inline finally can pop + cache the client — otherwise the
|
|
delayed abort poisons the slot and kills the NEXT request's sockets.
|
|
"""
|
|
from agent.chat_completion_helpers import direct_api_call
|
|
|
|
agent = _make_agent()
|
|
|
|
allow_finish = threading.Event()
|
|
close_reasons = []
|
|
abort_reasons = []
|
|
observed = {"close_during_abort": None}
|
|
|
|
client = MagicMock()
|
|
|
|
def blocking_create(**kwargs):
|
|
# Block until the aborting thread lets us race toward the finally.
|
|
allow_finish.wait(timeout=5.0)
|
|
return SimpleNamespace(choices=[])
|
|
|
|
client.chat.completions.create.side_effect = blocking_create
|
|
|
|
def fake_abort(aborted_client, *, reason):
|
|
abort_reasons.append(reason)
|
|
# Unblock the owner mid-abort. Under the fix it must block in its
|
|
# finally on the holder lock until this abort returns.
|
|
allow_finish.set()
|
|
deadline = time.time() + 0.6
|
|
while time.time() < deadline and not close_reasons:
|
|
time.sleep(0.02)
|
|
observed["close_during_abort"] = bool(close_reasons)
|
|
|
|
with patch.object(
|
|
agent, "_create_request_openai_client", return_value=client
|
|
), patch.object(
|
|
agent,
|
|
"_close_request_openai_client",
|
|
side_effect=lambda c, *, reason: close_reasons.append(reason),
|
|
), patch.object(
|
|
agent, "_abort_request_openai_client", side_effect=fake_abort
|
|
):
|
|
result = {}
|
|
|
|
def owner():
|
|
result["response"] = direct_api_call(agent, {})
|
|
|
|
owner_thread = threading.Thread(target=owner, daemon=True)
|
|
owner_thread.start()
|
|
deadline = time.time() + 5.0
|
|
while time.time() < deadline and getattr(agent, "_active_request_abort", None) is None:
|
|
time.sleep(0.01)
|
|
abort_fn = getattr(agent, "_active_request_abort", None)
|
|
assert abort_fn is not None
|
|
# Stranger-thread abort (this thread) racing the inline finally.
|
|
abort_fn("test_stranger_abort")
|
|
owner_thread.join(timeout=5.0)
|
|
assert not owner_thread.is_alive()
|
|
|
|
assert abort_reasons == ["test_stranger_abort"]
|
|
# The atomicity contract: no owner-side close slipped in mid-abort.
|
|
assert observed["close_during_abort"] is False
|
|
# ...and the inline finally still performed its close afterwards.
|
|
assert close_reasons == ["request_complete"]
|
|
assert result["response"] is not None
|
|
|
|
|
|
class _FakeCodexEventStream:
|
|
"""Codex Responses SSE stream stand-in (iterable, no ``output`` attr)."""
|
|
|
|
def __init__(self, events, close_raises=False):
|
|
self._events = events
|
|
self._close_raises = close_raises
|
|
self.close_calls = 0
|
|
|
|
def __iter__(self):
|
|
return iter(self._events)
|
|
|
|
def close(self):
|
|
self.close_calls += 1
|
|
if self._close_raises:
|
|
raise RuntimeError("close failed")
|
|
|
|
|
|
def _codex_delta_event(text):
|
|
return SimpleNamespace(type="response.output_text.delta", delta=text)
|
|
|
|
|
|
def _codex_completed_event():
|
|
return SimpleNamespace(
|
|
type="response.completed",
|
|
response=SimpleNamespace(usage=None, id="resp_1", status="completed"),
|
|
)
|
|
|
|
|
|
def test_codex_stream_close_failure_poisons_slot():
|
|
"""A failed ``event_stream.close()`` must poison the reuse slot.
|
|
|
|
Otherwise the connection stays checked out of the pool while the
|
|
worker's finally reports ``request_complete`` and caches the client
|
|
with the leaked connection (the codex sibling of the chat-streaming
|
|
interrupt-break leak).
|
|
"""
|
|
from agent.codex_runtime import run_codex_stream
|
|
|
|
agent = _make_agent()
|
|
stream = _FakeCodexEventStream(
|
|
[_codex_delta_event("partial "), _codex_completed_event()],
|
|
close_raises=True,
|
|
)
|
|
client = MagicMock()
|
|
client.responses.create.return_value = stream
|
|
abort_reasons = []
|
|
|
|
with patch.object(
|
|
agent,
|
|
"_abort_request_openai_client",
|
|
side_effect=lambda c, *, reason: abort_reasons.append(reason),
|
|
):
|
|
final = run_codex_stream(agent, {"model": "test/model"}, client=client)
|
|
|
|
assert final.output_text == "partial "
|
|
assert stream.close_calls == 1
|
|
assert abort_reasons == ["codex_stream_close_failed"]
|
|
|
|
|
|
def test_codex_stream_close_failure_on_primary_client_does_not_abort():
|
|
"""``client=None`` means the shared primary client — never reuse-cached,
|
|
and its sockets must not be force-shut by the close-failure handler."""
|
|
from agent.codex_runtime import run_codex_stream
|
|
|
|
agent = _make_agent()
|
|
stream = _FakeCodexEventStream(
|
|
[_codex_delta_event("hello"), _codex_completed_event()],
|
|
close_raises=True,
|
|
)
|
|
primary = MagicMock()
|
|
primary.responses.create.return_value = stream
|
|
|
|
with patch.object(
|
|
agent, "_ensure_primary_openai_client", return_value=primary
|
|
), patch.object(agent, "_abort_request_openai_client") as abort_mock:
|
|
final = run_codex_stream(agent, {"model": "test/model"}, client=None)
|
|
|
|
assert final.output_text == "hello"
|
|
assert stream.close_calls == 1
|
|
abort_mock.assert_not_called()
|