feat(compression): progress-aware timeouts — stop punishing slow summary models

The gateway's pre-agent session-hygiene compression killed the summary
call at a fixed 30s wall-clock deadline (compression.hygiene_timeout_seconds),
regardless of whether the summary model was hung or merely slow. A reasoning
model happily streaming a large summary was cut off mid-generation, the user
got '⚠️ Context compression timed out after 30.0s', and a 300s failure
cooldown left the session oversized — a doom loop for slow-but-healthy
auxiliary models.

Timeouts are now liveness-based instead of wall-clock-based:

- agent/auxiliary_client.py: new thread-local aux_progress_hook. When
  installed (only by context compression today), the primary call_llm
  attempt streams (stream=True) and aggregates chunks back into a complete
  response, ticking the hook per chunk. The configured timeout then acts
  per stream read (idle) instead of as a total budget. Providers that
  reject streaming fall back to the plain non-streaming call; auth/payment/
  rate-limit/transport errors propagate unchanged into the existing
  recovery chains. Codex Responses (per SSE event) and Anthropic Messages
  (per stream event, via the new create_anthropic_message on_stream_event
  callback) tick the same hook from inside their wire adapters.

- agent/conversation_compression.py: CompressionCommitFence gains
  touch_progress()/seconds_since_progress(); compress_context() installs
  fence.touch_progress as the progress hook around the compress call.

- gateway/run.py: the hygiene wait loop treats hygiene_timeout_seconds as
  an INACTIVITY budget — while the fence reports fresh progress the wait
  extends, bounded by the new compression.hygiene_total_ceiling_seconds
  (default 600s, clamped >= the idle budget) so a degenerate trickle
  stream still dies. The timeout warning now says the summary model
  produced no output, which is the only case that still triggers it.

- config/docs: hygiene_total_ceiling_seconds added to DEFAULT_CONFIG and
  configuration.md; hygiene_timeout_seconds documented as inactivity-based.

Tests: tests/agent/test_aux_progress_streaming.py (hook plumbing, stream
aggregation incl. tool-call deltas and reasoning deltas, rejection
fallback, ceiling kill, fence progress surface); two new gateway tests
prove a slow-but-streaming worker survives past the fixed timeout
(sabotage-verified: fails with the old fixed deadline) and a
forever-trickling worker is still cut off at the ceiling.
This commit is contained in:
Teknium 2026-07-25 11:33:20 -07:00
parent ac327dfa38
commit 32fd9d65cf
9 changed files with 884 additions and 15 deletions

View file

@ -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):

View file

@ -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

View file

@ -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

View file

@ -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 "

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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()

View file

@ -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.