diff --git a/agent/anthropic_adapter.py b/agent/anthropic_adapter.py index 87ee4631d71d..7fa3ba391fdc 100644 --- a/agent/anthropic_adapter.py +++ b/agent/anthropic_adapter.py @@ -2893,6 +2893,7 @@ def create_anthropic_message( *, log_prefix: str = "", prefer_stream: bool = True, + on_stream_event=None, ) -> Any: """Create an Anthropic message, aggregating via stream when available. @@ -2902,6 +2903,13 @@ def create_anthropic_message( crash on ``.content``. Prefer ``messages.stream().get_final_message()`` to match the main turn path, falling back to ``create()`` only for providers that explicitly do not support streaming, such as restricted Bedrock roles. + + ``on_stream_event``: optional callable invoked once per streamed event + (best-effort, exceptions swallowed). Lets callers report forward progress + to liveness watchdogs — e.g. the auxiliary compression path ticking its + progress hook so a slow-but-generating summary model isn't treated as + hung. Only fires on the streaming path; the ``create()`` fallback has no + events to report. """ sanitize_anthropic_kwargs(api_kwargs, log_prefix=log_prefix) @@ -2912,6 +2920,18 @@ def create_anthropic_message( stream_kwargs.pop("stream", None) try: with stream_fn(**stream_kwargs) as stream: + if callable(on_stream_event): + # Consume the event stream manually so each event can + # tick the caller's progress callback; get_final_message + # then returns the accumulated snapshot. + for _event in stream: + try: + on_stream_event(_event) + except Exception: + logger.debug( + "%son_stream_event callback failed", + log_prefix, exc_info=True, + ) return stream.get_final_message() except Exception as exc: if not _is_stream_unavailable_error(exc): diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index bd852d54df1a..7e4b9ae43b32 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -247,6 +247,50 @@ def aux_interrupt_protection(active: bool = True): _aux_interrupt_protection.active = prev +# ── Forward-progress hook for streamed auxiliary calls ─────────────────── +# Long auxiliary calls (context compression is the prime case) are watched by +# wall-clock deadlines in their hosts (gateway session hygiene). A fixed +# deadline punishes SLOW summary models exactly as hard as HUNG ones: a +# reasoning model happily streaming a large summary is killed mid-generation. +# This thread-local hook lets the host observe liveness instead: the wire +# consumers below tick it on every streamed token/SSE event, and the host +# extends its deadline while tokens are moving (see gateway/run.py session +# hygiene + CompressionCommitFence.touch_progress). Thread-local matches the +# call topology — the aux call and its stream consumption run synchronously +# on the thread that installed the hook. +_aux_progress = threading.local() + + +def _notify_aux_progress() -> None: + """Tick the installed forward-progress hook, if any. Never raises.""" + hook = getattr(_aux_progress, "hook", None) + if hook is None: + return + try: + hook() + except Exception: + logger.debug("aux progress hook failed", exc_info=True) + + +def _aux_progress_active() -> bool: + return getattr(_aux_progress, "hook", None) is not None + + +@contextlib.contextmanager +def aux_progress_hook(hook): + """Install *hook* as the current thread's aux forward-progress callback. + + ``hook=None`` is a no-op passthrough so callers can wire it + unconditionally. Re-entrant-safe: restores the previous hook on exit. + """ + prev = getattr(_aux_progress, "hook", None) + _aux_progress.hook = hook if callable(hook) else prev + try: + yield + finally: + _aux_progress.hook = prev + + def _safe_isinstance(obj: Any, maybe_type: Any) -> bool: """Return False instead of raising when a patched symbol is not a type.""" try: @@ -1162,6 +1206,10 @@ class _CodexCompletionsAdapter: def _on_each_event(_event: Any) -> None: # Re-check timeout/cancellation per event, matching the # cadence the old in-line ``_check_cancelled()`` used. + # Each SSE event is also forward progress for hosts watching + # a progress hook (gateway session hygiene): a reasoning + # model streaming a long summary must not look hung. + _notify_aux_progress() _check_cancelled() event_stream = self._client.responses.create(**stream_kwargs) @@ -1403,7 +1451,18 @@ class _AnthropicCompletionsAdapter: existing = {} anthropic_kwargs["extra_body"] = {**existing, **passthrough} - response = create_anthropic_message(self._client, anthropic_kwargs) + response = create_anthropic_message( + self._client, + anthropic_kwargs, + # Tick the aux forward-progress hook per streamed event so hosts + # watching liveness (gateway session hygiene) don't kill a + # slow-but-generating summary model. No-op when no hook is + # installed (None keeps the fast get_final_message path). + on_stream_event=( + (lambda _event: _notify_aux_progress()) + if _aux_progress_active() else None + ), + ) _transport = get_transport("anthropic_messages") _nr = _transport.normalize_response( response, strip_tool_prefix=self._is_oauth @@ -7117,6 +7176,228 @@ def _obj_get(obj: Any, key: str, default: Any = None) -> Any: return value +# ── Streamed aggregation for progress-hooked auxiliary calls ───────────── +# When a forward-progress hook is installed (aux_progress_hook — today only +# by context compression), the primary chat.completions attempt is upgraded +# to a streamed request that is aggregated back into a complete response. +# Two effects, both deliberate: +# 1. The configured ``timeout`` becomes an INTER-CHUNK idle timeout instead +# of a total budget (httpx applies the read timeout per stream read), so +# a slow-but-generating summary model is never killed mid-generation +# while tokens are moving — only a genuinely silent connection dies. +# 2. Every arriving chunk ticks the progress hook, letting outer watchdogs +# (gateway session hygiene) extend their deadlines on liveness instead +# of guessing with a fixed wall clock. +# A total ceiling still bounds the pathological 1-token-per-idle-window +# stream; see _aux_stream_total_ceiling(). + +_AUX_STREAM_CEILING_FLOOR_SECONDS = 600.0 +_AUX_STREAM_CEILING_MULTIPLIER = 4.0 + + +def _aux_stream_total_ceiling(effective_timeout: Optional[float]) -> float: + """Absolute wall-clock bound for a progress-hooked streamed aux call. + + Generous by design — the idle timeout is the real guard; this only stops + a degenerate stream that trickles one token per idle window forever. + """ + try: + timeout = float(effective_timeout) if effective_timeout is not None else 0.0 + except (TypeError, ValueError): + timeout = 0.0 + return max(_AUX_STREAM_CEILING_FLOOR_SECONDS, + _AUX_STREAM_CEILING_MULTIPLIER * timeout) + + +def _client_streams_internally(client: Any) -> bool: + """Wire adapters that consume a stream inside .create() already tick the + progress hook themselves (Codex per SSE event, Anthropic per stream + event); Bedrock's Converse shim cannot stream at all. None of them + accept chat-completions ``stream=True`` semantics from us.""" + return isinstance(client, ( + CodexAuxiliaryClient, + AnthropicAuxiliaryClient, + BedrockAuxiliaryClient, + )) + + +def _is_streaming_rejected_error(exc: Exception) -> bool: + """Provider explicitly refused a streamed chat.completions request.""" + err = str(exc).lower() + if "stream_options" in err: + return True + return "stream" in err and ( + "not supported" in err + or "unsupported" in err + or "not allowed" in err + or "disabled" in err + ) + + +def _create_with_progress(client: Any, kwargs: Dict[str, Any], task: Optional[str] = None) -> Any: + """chat.completions.create() that streams when a progress hook is active. + + Behavior is byte-for-byte identical to a plain ``create(**kwargs)`` when + no hook is installed (every existing caller/task) or when the client's + wire adapter streams internally. With a hook + a chunk-capable client, + the request is sent with ``stream=True`` and aggregated, ticking the hook + per chunk — so the configured ``timeout`` acts per stream read (idle) + rather than as a total budget, and outer liveness watchdogs see tokens + moving. Providers that reject the streamed request fall back to the + plain non-streaming call. + """ + _notify_aux_progress() # request dispatched counts as progress + if not _aux_progress_active() or _client_streams_internally(client): + return client.chat.completions.create(**kwargs) + + total_ceiling = _aux_stream_total_ceiling(kwargs.get("timeout")) + stream_kwargs = dict(kwargs) + stream_kwargs["stream"] = True + stream_kwargs["stream_options"] = {"include_usage": True} + try: + chunks = client.chat.completions.create(**stream_kwargs) + except Exception as exc: + # Genuine provider failures (auth, credit, rate limit, network) are + # not streaming's fault — surface them unchanged so the existing + # recovery chains (credential refresh, pool rotation, provider + # fallback) see the same error they would on a plain call. + if ( + _is_transient_transport_error(exc) + or _is_auth_error(exc) + or _is_payment_error(exc) + or _is_rate_limit_error(exc) + ): + raise + # Anything else may be a streaming-specific rejection (explicit + # "stream not supported", stream_options 400, or an idiosyncratic + # 4xx). Retry non-streaming once; if the request itself is bad the + # plain call reproduces the real error for the normal except-chains. + logger.debug( + "Auxiliary %s: streamed request failed (%s); retrying " + "non-streaming", task or "call", exc, + ) + return client.chat.completions.create(**kwargs) + + # Some shims (MoA virtual provider under quiet mode, defensive adapters) + # return a complete response even when stream=True was requested. + if hasattr(chunks, "choices"): + _notify_aux_progress() + return chunks + return _aggregate_chat_stream( + chunks, model=str(kwargs.get("model") or ""), total_ceiling=total_ceiling, + ) + + +def _aggregate_chat_stream( + chunks: Any, + *, + model: str = "", + total_ceiling: Optional[float] = None, +) -> Any: + """Consume a chat.completions chunk stream into a complete response. + + Ticks the thread-local aux progress hook on every chunk. Raises + TimeoutError when *total_ceiling* seconds elapse before the stream + finishes — phrased with "timed out" so existing timeout classification + (``_is_timeout_error``) treats it exactly like a request timeout. + """ + started = time.monotonic() + content_parts: List[str] = [] + reasoning_parts: List[str] = [] + tool_calls_acc: Dict[int, Dict[str, Any]] = {} + finish_reason = None + usage = None + resp_id = "" + resp_model = model or "" + try: + for chunk in chunks: + _notify_aux_progress() + if total_ceiling is not None and (time.monotonic() - started) >= total_ceiling: + raise TimeoutError( + f"Auxiliary streamed call timed out after {total_ceiling:.0f}s " + "total ceiling (stream still open but over budget)" + ) + resp_id = getattr(chunk, "id", None) or resp_id + resp_model = getattr(chunk, "model", None) or resp_model + chunk_usage = getattr(chunk, "usage", None) + if chunk_usage: + usage = chunk_usage + choices = getattr(chunk, "choices", None) or [] + if not choices: + continue + choice = choices[0] + finish_reason = getattr(choice, "finish_reason", None) or finish_reason + delta = getattr(choice, "delta", None) + if delta is None: + continue + piece = getattr(delta, "content", None) + if piece: + content_parts.append(piece) + # Thinking models stream reasoning deltas before any content — + # they count as forward progress (the tick above) and are kept + # out of .content, mirroring non-streaming responses where + # reasoning arrives in a separate field. + reasoning_piece = ( + getattr(delta, "reasoning", None) + or getattr(delta, "reasoning_content", None) + ) + if reasoning_piece and isinstance(reasoning_piece, str): + reasoning_parts.append(reasoning_piece) + for tc in (getattr(delta, "tool_calls", None) or []): + idx = getattr(tc, "index", 0) or 0 + acc = tool_calls_acc.setdefault( + idx, {"id": "", "name": "", "arguments": []} + ) + if getattr(tc, "id", None): + acc["id"] = tc.id + fn = getattr(tc, "function", None) + if fn is not None: + if getattr(fn, "name", None): + acc["name"] = fn.name + if getattr(fn, "arguments", None): + acc["arguments"].append(fn.arguments) + finally: + close_fn = getattr(chunks, "close", None) + if callable(close_fn): + try: + close_fn() + except Exception: + pass + + tool_calls = None + if tool_calls_acc: + tool_calls = [ + SimpleNamespace( + id=acc["id"], + type="function", + function=SimpleNamespace( + name=acc["name"], + arguments="".join(acc["arguments"]), + ), + ) + for _idx, acc in sorted(tool_calls_acc.items()) + ] + + message = SimpleNamespace( + role="assistant", + content="".join(content_parts), + tool_calls=tool_calls, + reasoning="".join(reasoning_parts) or None, + ) + choice = SimpleNamespace( + index=0, + message=message, + finish_reason=finish_reason or "stop", + ) + return SimpleNamespace( + id=resp_id, + model=resp_model, + object="chat.completion", + choices=[choice], + usage=usage, + ) + + def call_llm( task: str = None, *, @@ -7317,7 +7598,7 @@ def call_llm( # for the transient retry every auxiliary task shares. (PR #16587) try: return _validate_llm_response( - client.chat.completions.create(**kwargs), task, + _create_with_progress(client, kwargs, task), task, provider=resolved_provider, base_url=_base_info) except Exception as transient_err: if not _is_transient_transport_error(transient_err): @@ -7350,7 +7631,7 @@ def call_llm( time.sleep(_backoff) try: return _validate_llm_response( - client.chat.completions.create(**kwargs), task) + _create_with_progress(client, kwargs, task), task) except Exception as retry_transient: if not _is_transient_transport_error(retry_transient): raise diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 723f50d2ea07..78d4a146e762 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -221,6 +221,25 @@ class CompressionCommitFence: self._lock = threading.Lock() self._cancelled = False self._commit_started = False + # Forward-progress telemetry: the compression worker touches this + # whenever the streamed summary call produces a token (see + # ContextCompressor._call_summary_llm). Waiters use it to distinguish + # a SLOW-but-alive summary model from a HUNG one, so slow models are + # not killed by a fixed wall-clock deadline while tokens are moving. + self._last_progress = time.monotonic() + + def touch_progress(self) -> None: + """Record forward progress (e.g. a streamed summary token arriving). + + Called from the compression worker thread; read by async waiters via + :meth:`seconds_since_progress`. A bare float store is atomic in + CPython, so no lock is needed. + """ + self._last_progress = time.monotonic() + + def seconds_since_progress(self) -> float: + """Seconds since the worker last reported forward progress.""" + return max(0.0, time.monotonic() - self._last_progress) def cancel_before_commit(self) -> bool: """Cancel a pending commit, or wait for an active commit to finish. @@ -1741,7 +1760,17 @@ def compress_context( messages_before_compression = copy.deepcopy(messages) _activity_heartbeat = _CompressionActivityHeartbeat(agent).start() - compressed = compress_fn(messages, **compress_kwargs) + # Publish forward progress to the commit fence while the summary LLM + # call streams. Async hosts (gateway session hygiene) poll + # ``commit_fence.seconds_since_progress()`` to extend their deadline + # while tokens are moving — so a SLOW summary model is only killed + # when it is actually silent, not merely thorough. The hook is + # thread-local and the compress call is synchronous on this thread, + # so it cannot leak into unrelated auxiliary calls. + from agent.auxiliary_client import aux_progress_hook + _progress_hook = commit_fence.touch_progress if commit_fence is not None else None + with aux_progress_hook(_progress_hook): + compressed = compress_fn(messages, **compress_kwargs) except BaseException as _compress_exc: # ANY exception after lock acquisition — memory hook, capability # inspection, engine lookup, or compress() — must release the lock so diff --git a/gateway/run.py b/gateway/run.py index 22d7c58e9e11..aac6a192555f 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -13191,6 +13191,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew _hyg_compression_enabled = True _hyg_hard_msg_limit = 5000 _hyg_timeout_seconds = 30.0 + _hyg_total_ceiling_seconds = 600.0 _hyg_failure_cooldown_seconds = 300.0 _hyg_config_context_length = None _hyg_provider = None @@ -13245,6 +13246,19 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew _hyg_timeout_seconds = _parsed except (TypeError, ValueError): pass + _raw_ceiling = _comp_cfg.get("hygiene_total_ceiling_seconds") + if _raw_ceiling is not None: + try: + _parsed = float(_raw_ceiling) + if _parsed > 0: + _hyg_total_ceiling_seconds = _parsed + except (TypeError, ValueError): + pass + # The ceiling can never be tighter than one idle + # window, or the extension loop would be dead code. + _hyg_total_ceiling_seconds = max( + _hyg_total_ceiling_seconds, _hyg_timeout_seconds, + ) _raw_cooldown = _comp_cfg.get("hygiene_failure_cooldown_seconds") if _raw_cooldown is not None: try: @@ -13465,10 +13479,43 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew ), ) try: - _compressed, _ = await asyncio.wait_for( - asyncio.shield(_hyg_future), - timeout=_hyg_timeout_seconds, - ) + # Progress-aware wait: the timeout is an + # INACTIVITY budget, not a total one. The + # compression worker streams its summary + # call and ticks the fence per token + # (CompressionCommitFence.touch_progress), + # so a slow reasoning model that is still + # generating keeps extending the deadline; + # only a genuinely silent worker times out. + # A hard ceiling bounds the total wait so + # a degenerate trickle stream can't hold + # the turn forever. + _hyg_wait_started = time.monotonic() + while True: + try: + _compressed, _ = await asyncio.wait_for( + asyncio.shield(_hyg_future), + timeout=_hyg_timeout_seconds, + ) + break + except asyncio.TimeoutError: + _hyg_waited = time.monotonic() - _hyg_wait_started + _idle = _hyg_commit_fence.seconds_since_progress() + if ( + _idle < _hyg_timeout_seconds + and _hyg_waited < _hyg_total_ceiling_seconds + ): + logger.info( + "Session hygiene compression for " + "session %s still streaming after " + "%.0fs (last progress %.1fs ago) — " + "extending wait (ceiling %.0fs)", + session_entry.session_id, + _hyg_waited, _idle, + _hyg_total_ceiling_seconds, + ) + continue + raise except asyncio.TimeoutError: _cancelled = None while _cancelled is None: @@ -13497,14 +13544,18 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew ] = time.time() + _hyg_failure_cooldown_seconds logger.warning( "Session hygiene compression for session %s " - "timed out after %.1fs; continuing without " - "compression", + "made no progress for %.1fs " + "(total wait %.1fs, ceiling %.1fs); " + "continuing without compression", session_entry.session_id, - _hyg_timeout_seconds, + _hyg_commit_fence.seconds_since_progress(), + time.monotonic() - _hyg_wait_started, + _hyg_total_ceiling_seconds, ) _timeout_msg = ( "⚠️ Context compression timed out " - f"after {_hyg_timeout_seconds:.1f}s. " + f"after {_hyg_timeout_seconds:.1f}s " + "with no output from the summary model. " "No messages were dropped — continuing without " "compression. Run /compress to retry, /reset for " "a clean session, or check your " diff --git a/hermes_cli/config.py b/hermes_cli/config.py index 2f2e22268089..a82b328de5d9 100644 --- a/hermes_cli/config.py +++ b/hermes_cli/config.py @@ -1469,6 +1469,13 @@ DEFAULT_CONFIG = { # tool iteration. 0 = commit any non-zero prune. "hygiene_hard_message_limit": 5000, # gateway session-hygiene force-compress threshold by message count "hygiene_timeout_seconds": 30, # max seconds gateway waits for pre-agent hygiene compression + # WITHOUT forward progress. The summary call streams, so + # this is an inactivity budget: a slow model still + # producing tokens keeps extending the wait; only a + # silent/hung call is cut off. + "hygiene_total_ceiling_seconds": 600, # absolute cap on the hygiene compression wait even + # while tokens are still moving — bounds a degenerate + # trickle stream. Clamped to >= hygiene_timeout_seconds. "hygiene_failure_cooldown_seconds": 300, # skip repeated failed hygiene attempts for this session "protect_first_n": 3, # non-system head messages always preserved # verbatim, in ADDITION to the system prompt diff --git a/tests/agent/test_aux_progress_streaming.py b/tests/agent/test_aux_progress_streaming.py new file mode 100644 index 000000000000..05e1ccb47d90 --- /dev/null +++ b/tests/agent/test_aux_progress_streaming.py @@ -0,0 +1,273 @@ +"""Tests for the auxiliary forward-progress streaming layer. + +Slow summary models must not be punished like hung ones (#see PR): when a +forward-progress hook is installed (context compression), the primary +auxiliary call streams and ticks the hook per chunk, so outer watchdogs +(gateway session hygiene) can extend their deadline on liveness. Without a +hook, behavior is byte-for-byte the old non-streaming call. +""" + +import threading +import time +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from agent.auxiliary_client import ( + _aggregate_chat_stream, + _aux_stream_total_ceiling, + _create_with_progress, + _notify_aux_progress, + aux_progress_hook, +) +from agent.conversation_compression import CompressionCommitFence + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _chunk(content=None, reasoning=None, finish_reason=None, usage=None, + tool_calls=None, model="m1", chunk_id="c1"): + delta = SimpleNamespace( + content=content, + reasoning=reasoning, + reasoning_content=None, + tool_calls=tool_calls, + ) + choice = SimpleNamespace(delta=delta, finish_reason=finish_reason) + return SimpleNamespace( + id=chunk_id, model=model, choices=[choice], usage=usage, + ) + + +class _FakeClient: + """OpenAI-shaped client whose create() returns a canned value or stream.""" + + def __init__(self, response=None, stream_chunks=None, stream_error=None): + self.calls = [] + self._response = response + self._stream_chunks = stream_chunks + self._stream_error = stream_error + completions = SimpleNamespace(create=self._create) + self.chat = SimpleNamespace(completions=completions) + + def _create(self, **kwargs): + self.calls.append(kwargs) + if kwargs.get("stream"): + if self._stream_error is not None: + raise self._stream_error + return iter(self._stream_chunks or []) + return self._response + + +_COMPLETE = SimpleNamespace( + id="r1", model="m1", object="chat.completion", + choices=[SimpleNamespace( + index=0, + message=SimpleNamespace(role="assistant", content="non-streamed"), + finish_reason="stop", + )], + usage=None, +) + + +# --------------------------------------------------------------------------- +# aux_progress_hook plumbing +# --------------------------------------------------------------------------- + +class TestAuxProgressHook: + def test_hook_installed_and_restored(self): + ticks = [] + with aux_progress_hook(lambda: ticks.append(1)): + _notify_aux_progress() + _notify_aux_progress() # outside — must not tick + assert ticks == [1] + + def test_none_hook_is_noop_passthrough(self): + with aux_progress_hook(None): + _notify_aux_progress() # must not raise + + def test_hook_exception_is_swallowed(self): + def _boom(): + raise RuntimeError("hook blew up") + with aux_progress_hook(_boom): + _notify_aux_progress() # must not raise + + def test_hook_is_thread_local(self): + ticks = [] + seen_in_thread = [] + + def _other_thread(): + # No hook installed on this thread. + _notify_aux_progress() + seen_in_thread.append(len(ticks)) + + with aux_progress_hook(lambda: ticks.append(1)): + t = threading.Thread(target=_other_thread) + t.start() + t.join() + assert seen_in_thread == [0] + + +# --------------------------------------------------------------------------- +# _create_with_progress +# --------------------------------------------------------------------------- + +class TestCreateWithProgress: + def test_no_hook_means_plain_nonstreaming_call(self): + client = _FakeClient(response=_COMPLETE) + result = _create_with_progress(client, {"model": "m1", "messages": []}) + assert result is _COMPLETE + assert len(client.calls) == 1 + assert "stream" not in client.calls[0] + + def test_hook_upgrades_to_streaming_and_ticks_per_chunk(self): + chunks = [ + _chunk(reasoning="thinking..."), + _chunk(content="Hello "), + _chunk(content="world", finish_reason="stop", + usage=SimpleNamespace(prompt_tokens=5, completion_tokens=2, + total_tokens=7)), + ] + client = _FakeClient(stream_chunks=chunks) + ticks = [] + with aux_progress_hook(lambda: ticks.append(1)): + result = _create_with_progress( + client, {"model": "m1", "messages": [], "timeout": 30}, + ) + assert client.calls[0]["stream"] is True + assert result.choices[0].message.content == "Hello world" + assert result.choices[0].message.reasoning == "thinking..." + assert result.choices[0].finish_reason == "stop" + assert result.usage.total_tokens == 7 + # 1 dispatch tick + 1 per chunk + assert len(ticks) >= len(chunks) + 1 + + def test_streaming_rejected_falls_back_to_plain_call(self): + client = _FakeClient( + response=_COMPLETE, + stream_error=RuntimeError("stream is not supported by this model"), + ) + with aux_progress_hook(lambda: None): + result = _create_with_progress( + client, {"model": "m1", "messages": []}, + ) + assert result is _COMPLETE + # streamed attempt + non-streaming fallback + assert len(client.calls) == 2 + assert client.calls[0].get("stream") is True + assert "stream" not in client.calls[1] + + def test_auth_error_propagates_without_nonstreaming_retry(self): + class _FakeAuthError(Exception): + status_code = 401 + client = _FakeClient(stream_error=_FakeAuthError("Error code: 401 - unauthorized")) + with aux_progress_hook(lambda: None): + with pytest.raises(_FakeAuthError): + _create_with_progress(client, {"model": "m1", "messages": []}) + assert len(client.calls) == 1 # no silent non-streaming retry + + def test_shim_returning_complete_response_passes_through(self): + # Adapters may ignore stream=True and hand back a full response. + class _ShimClient(_FakeClient): + def _create(self, **kwargs): + self.calls.append(kwargs) + return _COMPLETE + client = _ShimClient() + with aux_progress_hook(lambda: None): + result = _create_with_progress(client, {"model": "m1", "messages": []}) + assert result is _COMPLETE + + +# --------------------------------------------------------------------------- +# _aggregate_chat_stream +# --------------------------------------------------------------------------- + +class TestAggregateChatStream: + def test_tool_call_deltas_are_reassembled(self): + tc0 = SimpleNamespace( + index=0, id="call_1", + function=SimpleNamespace(name="do_thing", arguments='{"a"'), + ) + tc1 = SimpleNamespace( + index=0, id=None, + function=SimpleNamespace(name=None, arguments=': 1}'), + ) + chunks = [ + _chunk(tool_calls=[tc0]), + _chunk(tool_calls=[tc1], finish_reason="tool_calls"), + ] + result = _aggregate_chat_stream(iter(chunks)) + tool_calls = result.choices[0].message.tool_calls + assert len(tool_calls) == 1 + assert tool_calls[0].id == "call_1" + assert tool_calls[0].function.name == "do_thing" + assert tool_calls[0].function.arguments == '{"a": 1}' + assert result.choices[0].finish_reason == "tool_calls" + + def test_total_ceiling_kills_trickle_stream_as_timeout(self): + def _trickle(): + while True: + time.sleep(0.01) + yield _chunk(content="x") + with pytest.raises(TimeoutError, match="timed out"): + _aggregate_chat_stream(_trickle(), total_ceiling=0.05) + + def test_stream_close_is_called(self): + closed = [] + + class _Stream: + def __iter__(self): + return iter([_chunk(content="ok", finish_reason="stop")]) + + def close(self): + closed.append(True) + + result = _aggregate_chat_stream(_Stream()) + assert result.choices[0].message.content == "ok" + assert closed == [True] + + def test_empty_choices_chunks_are_skipped(self): + empty = SimpleNamespace(id="c", model="m", choices=[], usage=None) + chunks = [empty, _chunk(content="ok", finish_reason="stop")] + result = _aggregate_chat_stream(iter(chunks)) + assert result.choices[0].message.content == "ok" + + +# --------------------------------------------------------------------------- +# Ceiling arithmetic +# --------------------------------------------------------------------------- + +class TestStreamCeiling: + def test_floor_applies_to_small_timeouts(self): + assert _aux_stream_total_ceiling(30) == 600.0 + + def test_multiplier_wins_for_large_timeouts(self): + assert _aux_stream_total_ceiling(300) == 1200.0 + + def test_none_timeout_gets_floor(self): + assert _aux_stream_total_ceiling(None) == 600.0 + + +# --------------------------------------------------------------------------- +# CompressionCommitFence progress surface +# --------------------------------------------------------------------------- + +class TestFenceProgress: + def test_touch_progress_resets_idle_clock(self): + fence = CompressionCommitFence() + time.sleep(0.05) + assert fence.seconds_since_progress() >= 0.04 + fence.touch_progress() + assert fence.seconds_since_progress() < 0.05 + + def test_fence_hook_wiring_matches_compressor_usage(self): + # conversation_compression installs fence.touch_progress as the hook; + # verify the pair works end-to-end through _notify_aux_progress. + fence = CompressionCommitFence() + time.sleep(0.05) + with aux_progress_hook(fence.touch_progress): + _notify_aux_progress() + assert fence.seconds_since_progress() < 0.05 diff --git a/tests/agent/test_auxiliary_client.py b/tests/agent/test_auxiliary_client.py index fae90d2a89e4..abe1e7d23bdc 100644 --- a/tests/agent/test_auxiliary_client.py +++ b/tests/agent/test_auxiliary_client.py @@ -1741,7 +1741,7 @@ class TestVisionClientFallback: real_client = SimpleNamespace(messages=messages_api) captured_kwargs = {} - def fake_create_anthropic_message(_client, kwargs): + def fake_create_anthropic_message(_client, kwargs, **_kw): captured_kwargs.update(kwargs) return final_message diff --git a/tests/gateway/test_session_hygiene.py b/tests/gateway/test_session_hygiene.py index b41b30378443..9254d1b2663c 100644 --- a/tests/gateway/test_session_hygiene.py +++ b/tests/gateway/test_session_hygiene.py @@ -1412,3 +1412,208 @@ async def test_session_hygiene_default_hard_message_limit_does_not_fire_at_12_me assert FakeCompressAgent.last_instance is None, ( "Compression should NOT fire at 12 messages with default hard_limit=5000" ) + + +# --------------------------------------------------------------------------- +# Progress-aware hygiene wait: slow-but-streaming models are not punished +# --------------------------------------------------------------------------- + +def _make_progress_runner(monkeypatch, tmp_path, agent_cls, cfg_text): + """Shared scaffolding for the progress-aware hygiene wait tests.""" + fake_dotenv = types.ModuleType("dotenv") + fake_dotenv.load_dotenv = lambda *args, **kwargs: None + monkeypatch.setitem(sys.modules, "dotenv", fake_dotenv) + + fake_run_agent = types.ModuleType("run_agent") + fake_run_agent.AIAgent = agent_cls + monkeypatch.setitem(sys.modules, "run_agent", fake_run_agent) + + cfg_path = tmp_path / "config.yaml" + cfg_path.write_text(cfg_text) + + gateway_run = importlib.import_module("gateway.run") + GatewayRunner = gateway_run.GatewayRunner + + adapter = HygieneCaptureAdapter() + runner = object.__new__(GatewayRunner) + runner.config = GatewayConfig( + platforms={Platform.TELEGRAM: PlatformConfig(enabled=True, token="fake-token")} + ) + runner.adapters = {Platform.TELEGRAM: adapter} + runner._voice_mode = {} + runner.hooks = SimpleNamespace(emit=AsyncMock(), loaded_hooks=False) + runner.session_store = MagicMock() + runner.session_store.get_or_create_session.return_value = SessionEntry( + session_key="agent:main:telegram:dm:12345", + session_id="sess-progress", + created_at=datetime.now(), + updated_at=datetime.now(), + platform=Platform.TELEGRAM, + chat_type="dm", + ) + runner.session_store.load_transcript.return_value = _make_history(6, content_size=400) + runner.session_store.has_any_sessions.return_value = True + runner.session_store.rewrite_transcript = MagicMock() + runner.session_store.append_to_transcript = MagicMock() + runner._running_agents = {} + runner._pending_messages = {} + runner._pending_approvals = {} + runner._session_db = None + runner._is_user_authorized = lambda _source: True + runner._set_session_env = lambda _context: None + runner._run_agent = AsyncMock( + return_value={ + "final_response": "ok", + "messages": [], + "tools": [], + "history_offset": 0, + "last_prompt_tokens": 0, + } + ) + + monkeypatch.setattr(gateway_run, "_hermes_home", tmp_path) + monkeypatch.setattr(gateway_run, "_resolve_runtime_agent_kwargs", lambda: {"api_key": "fake"}) + monkeypatch.setattr( + "agent.model_metadata.get_model_context_length", + lambda *_args, **_kwargs: 100, + ) + + event = MessageEvent( + text="hello", + source=SessionSource( + platform=Platform.TELEGRAM, + chat_id="12345", + chat_type="dm", + user_id="12345", + ), + message_id="1", + ) + return runner, adapter, event + + +@pytest.mark.asyncio +async def test_hygiene_slow_but_streaming_worker_survives_past_timeout( + monkeypatch, tmp_path +): + """A summary model still streaming tokens must NOT be killed at the fixed + hygiene timeout. The worker here takes several idle windows to finish but + ticks fence.touch_progress() continually (as the streamed summary call + does per chunk) — the wait must extend and consume the completed result. + """ + class SlowStreamingCompressAgent: + last_instance = None + + def __init__(self, **kwargs): + self.session_id = kwargs.get("session_id", "fake-session") + self._print_fn = None + self.shutdown_memory_provider = MagicMock() + self.close = MagicMock() + type(self).last_instance = self + + def _compress_context(self, messages, *_args, commit_fence=None, **_kwargs): + # 6 idle windows of work, ticking progress the whole way. + deadline = time.monotonic() + 0.6 + while time.monotonic() < deadline: + if commit_fence is not None: + commit_fence.touch_progress() + time.sleep(0.02) + if commit_fence is not None and not commit_fence.begin_commit(): + return (messages, None) + try: + self.session_id = f"{self.session_id}_compressed" + return ([{"role": "assistant", "content": "compressed"}], None) + finally: + if commit_fence is not None: + commit_fence.finish_commit() + + runner, adapter, event = _make_progress_runner( + monkeypatch, tmp_path, SlowStreamingCompressAgent, + "compression:\n" + " enabled: true\n" + " hygiene_timeout_seconds: 0.1\n" # << worker runtime (0.6s) + " hygiene_total_ceiling_seconds: 10\n" + " hygiene_failure_cooldown_seconds: 120\n", + ) + runner._hygiene_compression_failure_cooldowns = {} + + result = await runner._handle_message(event) + + assert result == "ok" + assert SlowStreamingCompressAgent.last_instance is not None + # The slow-but-alive worker completed: rotation happened, no timeout + # warning was delivered, and no failure cooldown was recorded. + assert SlowStreamingCompressAgent.last_instance.session_id.endswith("_compressed") + timeout_warnings = [ + s for s in adapter.sent if "Context compression timed out" in s["content"] + ] + assert timeout_warnings == [] + assert "sess-progress" not in runner._hygiene_compression_failure_cooldowns + + +@pytest.mark.asyncio +async def test_hygiene_trickle_stream_is_bounded_by_total_ceiling( + monkeypatch, tmp_path +): + """A worker that keeps ticking progress forever must still be cut off at + hygiene_total_ceiling_seconds — liveness extends the wait, but not + indefinitely.""" + release_worker = threading.Event() + + class TrickleForeverCompressAgent: + last_instance = None + + def __init__(self, **kwargs): + self.session_id = kwargs.get("session_id", "fake-session") + self._print_fn = None + self._last_compaction_in_place = False + self.context_compressor = SimpleNamespace( + bind_session_state=MagicMock(), + _last_compress_aborted=False, + _last_aux_model_failure_model=None, + ) + self.shutdown_memory_provider = MagicMock() + self.close = MagicMock() + type(self).last_instance = self + + def _compress_context(self, messages, *_args, commit_fence=None, **_kwargs): + # Ticks progress on every iteration but never finishes until + # the test releases it — models a degenerate trickle stream. + while not release_worker.wait(0.02): + if commit_fence is not None: + commit_fence.touch_progress() + if commit_fence is not None and not commit_fence.begin_commit(): + return (messages, None) + try: + return (messages, None) + finally: + if commit_fence is not None: + commit_fence.finish_commit() + + runner, adapter, event = _make_progress_runner( + monkeypatch, tmp_path, TrickleForeverCompressAgent, + "compression:\n" + " enabled: true\n" + " hygiene_timeout_seconds: 0.1\n" + " hygiene_total_ceiling_seconds: 0.3\n" + " hygiene_failure_cooldown_seconds: 120\n", + ) + runner._hygiene_compression_failure_cooldowns = {} + runner._session_db = SimpleNamespace(_db=MagicMock()) + + started = time.monotonic() + try: + result = await runner._handle_message(event) + finally: + release_worker.set() + elapsed = time.monotonic() - started + + assert result == "ok" + # Loose wall-clock bound per flake policy: asserts the ceiling stopped + # the wait (would otherwise spin until release_worker), not precise + # latency. + assert elapsed < 5.0 + timeout_warnings = [ + s for s in adapter.sent if "Context compression timed out" in s["content"] + ] + assert len(timeout_warnings) == 1 + assert runner._hygiene_compression_failure_cooldowns["sess-progress"] > time.time() diff --git a/website/docs/user-guide/configuration.md b/website/docs/user-guide/configuration.md index 4498d1852988..a5d20f6f8539 100644 --- a/website/docs/user-guide/configuration.md +++ b/website/docs/user-guide/configuration.md @@ -751,7 +751,8 @@ compression: protect_first_n: 3 # Non-system head messages pinned across compactions (0 = pin nothing) idle_compact_after_seconds: 0 # Opt-in idle compaction (0 = disabled) — see below hygiene_hard_message_limit: 5000 # Gateway safety valve — see below - hygiene_timeout_seconds: 30 # Max seconds gateway waits for pre-agent hygiene compression + hygiene_timeout_seconds: 30 # Max seconds of NO summary-model output before hygiene compression is cut off + hygiene_total_ceiling_seconds: 600 # Absolute cap on the hygiene wait even while tokens are still streaming hygiene_failure_cooldown_seconds: 300 # Skip repeated failed hygiene attempts for this session proactive_prune_tokens: 0 # Opt-in tokens trigger for the no-LLM tool-result prune (0 = off; see below) proactive_prune_min_result_chars: 8000 # Prune's summarize pass only touches tool results larger than this (clamped >= 200) @@ -773,7 +774,9 @@ Older configs with `compression.summary_model`, `compression.summary_provider`, `hygiene_hard_message_limit` is a gateway-only **pre-compression safety valve**. It exists to break a death spiral: when API calls keep disconnecting on an oversized session, the gateway never receives token-usage data, so the token-based threshold can't fire, so the transcript keeps growing and disconnects get worse. This count-based floor fires on message count alone (always known, regardless of API failures) to force compression and recover the session. Default `5000` — far above any normal session, including large-context (1M+) models doing thousands of short turns, which compress on the token threshold long before this. Raise it further for unusual platforms, lower it to force more aggressive compression. Editing this value on a running gateway takes effect on the next message (see below). -`hygiene_timeout_seconds` caps how long the gateway waits for this pre-agent compression pass. If the auxiliary compression backend is down or very slow, the gateway warns the user, continues the incoming message without compression, and records a temporary per-session failure cooldown instead of appearing stuck. +`hygiene_timeout_seconds` is the gateway's **inactivity budget** for this pre-agent compression pass — not a total wall-clock cap. The compression summary call streams from the model, and each arriving token counts as forward progress: a slow reasoning model that is still generating keeps extending its own deadline, so slow-but-healthy summary models are never cut off mid-generation. Only when the summary model produces **no output** for this many seconds (backend down, hung connection, silent provider) does the gateway warn the user, continue the incoming message without compression, and record a temporary per-session failure cooldown instead of appearing stuck. + +`hygiene_total_ceiling_seconds` (default `600`) bounds the total wait even while tokens are still moving, so a degenerate trickle stream can't hold a turn hostage indefinitely. It is clamped to at least `hygiene_timeout_seconds`. `hygiene_failure_cooldown_seconds` controls that per-session cooldown after a hygiene compression timeout or abort. During the cooldown, the gateway skips repeated hygiene attempts for the same oversized session so every incoming message does not block on the same broken auxiliary backend. `/compress`, `/reset`, or a healthy later turn can still recover the session.