mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-21 16:18:55 +00:00
fix(streaming): block stale stream deltas
This commit is contained in:
davidb73-hub
kshitij
parent
3ec4c9ce4d
commit
4fa67d2014
2 changed files with 158 additions and 2 deletions
|
|
@ -2429,6 +2429,67 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
# resolved, so the builder degrades to its plain default if it ever runs
|
||||
# first.
|
||||
_stream_stale_timeout = None
|
||||
stream_attempt_lock = threading.Lock()
|
||||
stream_attempt_state = {
|
||||
"current": 0,
|
||||
"cancelled": set(),
|
||||
"discarded_chunks": 0,
|
||||
"discarded_bytes": 0,
|
||||
}
|
||||
|
||||
def _start_stream_attempt() -> int:
|
||||
with stream_attempt_lock:
|
||||
stream_attempt_state["current"] += 1
|
||||
return int(stream_attempt_state["current"])
|
||||
|
||||
def _cancel_current_stream_attempt(reason: str) -> None:
|
||||
with stream_attempt_lock:
|
||||
current = int(stream_attempt_state.get("current") or 0)
|
||||
if current:
|
||||
stream_attempt_state["cancelled"].add(current)
|
||||
if current:
|
||||
logger.debug(
|
||||
"Marked stream attempt %s cancelled: %s",
|
||||
current,
|
||||
reason,
|
||||
)
|
||||
|
||||
def _stream_attempt_is_active(stream_attempt_id: int) -> bool:
|
||||
with stream_attempt_lock:
|
||||
return (
|
||||
stream_attempt_id == int(stream_attempt_state.get("current") or 0)
|
||||
and stream_attempt_id not in stream_attempt_state["cancelled"]
|
||||
)
|
||||
|
||||
def _stream_attempt_was_cancelled(stream_attempt_id: int) -> bool:
|
||||
with stream_attempt_lock:
|
||||
return stream_attempt_id in stream_attempt_state["cancelled"]
|
||||
|
||||
def _discard_stale_stream_chunk(stream_attempt_id: int, chunk) -> None:
|
||||
try:
|
||||
chunk_bytes = len(repr(chunk))
|
||||
except Exception:
|
||||
chunk_bytes = 0
|
||||
with stream_attempt_lock:
|
||||
stream_attempt_state["discarded_chunks"] += 1
|
||||
stream_attempt_state["discarded_bytes"] += chunk_bytes
|
||||
discarded_chunks = stream_attempt_state["discarded_chunks"]
|
||||
discarded_bytes = stream_attempt_state["discarded_bytes"]
|
||||
if discarded_chunks == 1:
|
||||
logger.warning(
|
||||
"Discarding chunk from superseded stream attempt %s "
|
||||
"(discarded_chunks=%s discarded_bytes=%s)",
|
||||
stream_attempt_id,
|
||||
discarded_chunks,
|
||||
discarded_bytes,
|
||||
)
|
||||
logger.debug(
|
||||
"Discarded stale stream chunk from attempt %s "
|
||||
"(discarded_chunks=%s discarded_bytes=%s)",
|
||||
stream_attempt_id,
|
||||
discarded_chunks,
|
||||
discarded_bytes,
|
||||
)
|
||||
|
||||
def _fire_first_delta():
|
||||
if not first_delta_fired["done"] and on_first_delta:
|
||||
|
|
@ -2438,7 +2499,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def _call_chat_completions():
|
||||
def _call_chat_completions(stream_attempt_id: int):
|
||||
"""Stream a chat completions response."""
|
||||
import httpx as _httpx
|
||||
# Per-provider / per-model request_timeout_seconds (from config.yaml)
|
||||
|
|
@ -2635,6 +2696,10 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
if agent._interrupt_requested:
|
||||
break
|
||||
|
||||
if not _stream_attempt_is_active(stream_attempt_id):
|
||||
_discard_stale_stream_chunk(stream_attempt_id, chunk)
|
||||
continue
|
||||
|
||||
if not chunk.choices:
|
||||
if hasattr(chunk, "model") and chunk.model:
|
||||
model_name = chunk.model
|
||||
|
|
@ -2760,6 +2825,11 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
if hasattr(chunk, "usage") and chunk.usage:
|
||||
usage_obj = chunk.usage
|
||||
|
||||
if _stream_attempt_was_cancelled(stream_attempt_id):
|
||||
raise _httpx.RemoteProtocolError(
|
||||
f"stream attempt {stream_attempt_id} was superseded"
|
||||
)
|
||||
|
||||
# Build mock response matching non-streaming shape
|
||||
full_content = "".join(content_parts) or None
|
||||
mock_tool_calls = None
|
||||
|
|
@ -3048,6 +3118,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
|
||||
try:
|
||||
for _stream_attempt in range(_max_stream_retries + 1):
|
||||
stream_attempt_id = _start_stream_attempt()
|
||||
# Check for interrupt before each retry attempt. Without
|
||||
# this, /stop closes the HTTP connection (outer poll loop),
|
||||
# but the retry loop opens a FRESH connection — negating the
|
||||
|
|
@ -3055,13 +3126,14 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
# retry can block for the full stream-read timeout (120s+),
|
||||
# causing multi-minute delays between /stop and response.
|
||||
if agent._interrupt_requested:
|
||||
_cancel_current_stream_attempt("interrupt_before_stream_retry")
|
||||
raise InterruptedError("Agent interrupted before stream retry")
|
||||
try:
|
||||
if agent.api_mode == "anthropic_messages":
|
||||
agent._try_refresh_anthropic_client_credentials()
|
||||
result["response"] = _call_anthropic()
|
||||
else:
|
||||
result["response"] = _call_chat_completions()
|
||||
result["response"] = _call_chat_completions(stream_attempt_id)
|
||||
return # success
|
||||
except Exception as e:
|
||||
# If the main poll loop force-closed this request because
|
||||
|
|
@ -3183,6 +3255,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
mid_tool_call=True,
|
||||
diag=request_client_holder.get("diag"),
|
||||
)
|
||||
_cancel_current_stream_attempt("stream_mid_tool_retry_cleanup")
|
||||
_close_request_client_once("stream_mid_tool_retry_cleanup")
|
||||
if agent.api_mode == "anthropic_messages":
|
||||
try:
|
||||
|
|
@ -3247,6 +3320,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
diag=request_client_holder.get("diag"),
|
||||
)
|
||||
# Close the stale request client before retry
|
||||
_cancel_current_stream_attempt("stream_retry_cleanup")
|
||||
_close_request_client_once("stream_retry_cleanup")
|
||||
# Also rebuild the primary client to purge
|
||||
# any dead connections from the pool.
|
||||
|
|
@ -3486,6 +3560,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
f"Reconnecting..."
|
||||
)
|
||||
try:
|
||||
_cancel_current_stream_attempt("stale_stream_kill")
|
||||
_close_request_client_once("stale_stream_kill")
|
||||
except Exception:
|
||||
pass
|
||||
|
|
@ -3527,6 +3602,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
|
|||
"(not a network error)."
|
||||
)
|
||||
try:
|
||||
_cancel_current_stream_attempt("stream_interrupt_abort")
|
||||
if agent.api_mode == "anthropic_messages":
|
||||
agent._anthropic_client.close()
|
||||
agent._rebuild_anthropic_client()
|
||||
|
|
|
|||
|
|
@ -160,3 +160,83 @@ class TestStreamInterruptBeforeRetry:
|
|||
result = agent._interruptible_streaming_api_call({})
|
||||
assert result is not None
|
||||
assert attempts[0] == 3
|
||||
|
||||
@pytest.mark.filterwarnings(
|
||||
"ignore::pytest.PytestUnhandledThreadExceptionWarning"
|
||||
)
|
||||
@patch("run_agent.AIAgent._replace_primary_openai_client")
|
||||
@patch("run_agent.AIAgent._abort_request_openai_client")
|
||||
@patch("run_agent.AIAgent._create_request_openai_client")
|
||||
@patch("run_agent.AIAgent._close_request_openai_client")
|
||||
def test_stale_stream_attempt_cannot_emit_late_chunks_after_retry(
|
||||
self,
|
||||
mock_close,
|
||||
mock_create,
|
||||
mock_abort,
|
||||
mock_replace,
|
||||
monkeypatch,
|
||||
):
|
||||
"""A stale attempt must not keep writing deltas after it is killed.
|
||||
|
||||
This reproduces the race where the outer stale detector aborts an SSE
|
||||
connection, but the old iterator still yields one more chunk before
|
||||
surfacing the connection error that triggers the retry.
|
||||
"""
|
||||
import httpx
|
||||
import time
|
||||
|
||||
from tests.run_agent.test_streaming import (
|
||||
_make_stream_chunk,
|
||||
_make_tool_call_delta,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("HERMES_STREAM_STALE_TIMEOUT", "0.05")
|
||||
monkeypatch.setenv("HERMES_STREAM_RETRIES", "1")
|
||||
|
||||
class LateChunkAfterStaleStream:
|
||||
response = SimpleNamespace(headers={})
|
||||
|
||||
def __iter__(self):
|
||||
yield _make_stream_chunk(content="old start ")
|
||||
yield _make_stream_chunk(
|
||||
tool_calls=[
|
||||
_make_tool_call_delta(
|
||||
index=0,
|
||||
tc_id="call_1",
|
||||
name="terminal",
|
||||
)
|
||||
]
|
||||
)
|
||||
time.sleep(0.45)
|
||||
yield _make_stream_chunk(content="old late ")
|
||||
raise httpx.RemoteProtocolError("peer closed connection")
|
||||
|
||||
retry_chunks = [
|
||||
_make_stream_chunk(content="new final"),
|
||||
_make_stream_chunk(finish_reason="stop", model="test/model"),
|
||||
]
|
||||
class RetryStream:
|
||||
response = SimpleNamespace(headers={})
|
||||
|
||||
def __iter__(self):
|
||||
return iter(retry_chunks)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.chat.completions.create.side_effect = [
|
||||
LateChunkAfterStaleStream(),
|
||||
RetryStream(),
|
||||
]
|
||||
mock_create.return_value = mock_client
|
||||
|
||||
agent = _make_agent()
|
||||
agent._interrupt_requested = False
|
||||
deltas = []
|
||||
agent.stream_delta_callback = deltas.append
|
||||
|
||||
response = agent._interruptible_streaming_api_call({})
|
||||
|
||||
delivered = "".join(deltas)
|
||||
assert "old late" not in delivered
|
||||
assert "new final" in delivered
|
||||
assert response.choices[0].message.content == "new final"
|
||||
assert mock_abort.called
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue