Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 28 additions & 2 deletions agent/context_compressor.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,15 @@
- Richer tool call/result detail in summarizer input
"""

import concurrent.futures
import hashlib
import json
import logging
import re
import time
from typing import Any, Dict, List, Optional

from agent.auxiliary_client import call_llm, _is_connection_error
from agent.auxiliary_client import call_llm, _get_task_timeout, _is_connection_error
from agent.context_engine import ContextEngine
from agent.model_metadata import (
MINIMUM_CONTEXT_LENGTH,
Expand Down Expand Up @@ -1310,6 +1311,31 @@ def _fallback_to_main_for_compression(self, e: Exception, reason: str) -> None:
self.summary_model = "" # empty = use main model
self._summary_failure_cooldown_until = 0.0 # no cooldown — retry immediately

def _call_summary_llm_with_hard_timeout(self, call_kwargs: Dict[str, Any]) -> Any:
"""Bound how long compression summary generation can block the caller.

``call_llm()`` already forwards ``auxiliary.compression.timeout`` to
the provider client, but the dashboard/gateway still waits
synchronously for that call to return. Run it in a worker thread and
stop waiting after the configured compression timeout so the normal
fallback summary path can keep the process responsive.
"""
hard_timeout = max(1.0, float(_get_task_timeout("compression")))
pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
future = pool.submit(call_llm, **call_kwargs)
try:
return future.result(timeout=hard_timeout)
except concurrent.futures.TimeoutError as exc:
future.cancel()
raise TimeoutError(
f"Context compression hard-timed out after {hard_timeout:g}s"
) from exc
finally:
# Do not wait for a stuck network stack to unwind in the caller's
# thread. The timed-out summary attempt is being abandoned in
# favor of the existing deterministic compression fallback.
pool.shutdown(wait=False, cancel_futures=True)

def _generate_summary(
self,
turns_to_summarize: List[Dict[str, Any]],
Expand Down Expand Up @@ -1519,7 +1545,7 @@ def _generate_summary(
}
if self.summary_model:
call_kwargs["model"] = self.summary_model
response = call_llm(**call_kwargs)
response = self._call_summary_llm_with_hard_timeout(call_kwargs)
content = response.choices[0].message.content
# Handle cases where content is not a string (e.g., dict from llama.cpp)
if not isinstance(content, str):
Expand Down
57 changes: 57 additions & 0 deletions tests/agent/test_context_compressor.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
"""Tests for agent/context_compressor.py — compression logic, thresholds, truncation fallback."""

import concurrent.futures

import pytest
from unittest.mock import patch, MagicMock

Expand Down Expand Up @@ -1069,6 +1071,61 @@ def test_summary_failure_fallback_is_bounded(self):
assert "deterministic fallback" in fallback
assert "important detail" in fallback

def test_compress_hard_timeout_uses_existing_fallback_path(self):
events = {"shutdown_calls": []}

class _TimeoutFuture:
def cancel(self):
events["cancel_called"] = True
return False

def result(self, timeout=None):
events["result_timeout"] = timeout
raise concurrent.futures.TimeoutError("timed out")

class _FakeExecutor:
def __init__(self, max_workers=1):
events["max_workers"] = max_workers

def submit(self, fn, *args, **kwargs):
events["submitted_fn"] = getattr(fn, "__name__", repr(fn))
events["submitted_kwargs"] = kwargs
return _TimeoutFuture()

def shutdown(self, wait=True, cancel_futures=False):
events["shutdown_calls"].append((wait, cancel_futures))

with patch("agent.context_compressor.get_model_context_length", return_value=100000):
c = ContextCompressor(model="test", quiet_mode=True, protect_first_n=2, protect_last_n=2)

msgs = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "msg 1"},
{"role": "assistant", "content": "msg 2"},
{"role": "user", "content": "msg 3"},
{"role": "assistant", "content": "msg 4"},
{"role": "user", "content": "msg 5"},
{"role": "assistant", "content": "msg 6"},
{"role": "user", "content": "msg 7"},
]

with patch("agent.context_compressor._get_task_timeout", return_value=12.5), \
patch("agent.context_compressor.concurrent.futures.ThreadPoolExecutor", _FakeExecutor):
result = c.compress(msgs)

assert events["max_workers"] == 1
assert events["submitted_fn"] == "call_llm"
assert events["submitted_kwargs"]["task"] == "compression"
assert events["result_timeout"] == 12.5
assert events["cancel_called"] is True
assert events["shutdown_calls"] == [(False, True)]
assert c._last_summary_fallback_used is True
assert "hard-timed out after 12.5s" in (c._last_summary_error or "")
assert any(
isinstance(m.get("content"), str) and "Summary generation was unavailable" in m["content"]
for m in result
)

def test_compress_clears_fallback_flag_on_subsequent_success(self):
mock_response = MagicMock()
mock_response.choices = [MagicMock()]
Expand Down
Loading