diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 22bf9c165c..5a33d25cc7 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -7601,6 +7601,77 @@ def _get_task_extra_body(task: str) -> Dict[str, Any]: return result +# --------------------------------------------------------------------------- +# Per-task concurrency limiting (#23324) +# --------------------------------------------------------------------------- +# Background auxiliary work (title generation, context compression, etc.) can +# spawn unbounded concurrent LLM calls when many sessions are active. During +# provider incidents each call also retries / fans out across the fallback +# chain, multiplying request volume on already-degraded endpoints. A per-task +# semaphore caps in-flight calls so retry amplification stays bounded. + +_aux_sync_semaphores: Dict[str, Tuple[int, threading.BoundedSemaphore]] = {} +_aux_async_semaphores: Dict[Tuple[str, int], Tuple[int, Any]] = {} +_aux_sem_lock = threading.Lock() + + +def _get_task_max_concurrency(task: Optional[str]) -> Optional[int]: + """Return ``auxiliary..max_concurrency`` as a positive int, or None.""" + if not task or task == "vision": + # Vision already uses this key for its encode/resize CPU worker pool; + # its LLM calls deliberately remain concurrent. + return None + raw = _get_auxiliary_task_config(task).get("max_concurrency") + if raw is None: + return None + try: + value = int(raw) + except (TypeError, ValueError): + return None + return value if value > 0 else None + + +def _acquire_sync_aux_semaphore(task: Optional[str]) -> Optional[threading.BoundedSemaphore]: + """Get a per-task sync semaphore, rebuilding it after a config change.""" + limit = _get_task_max_concurrency(task) + if limit is None: + return None + with _aux_sem_lock: + entry = _aux_sync_semaphores.get(task) + if entry is None or entry[0] != limit: + semaphore = threading.BoundedSemaphore(limit) + _aux_sync_semaphores[task] = (limit, semaphore) + return semaphore + return entry[1] + + +def _acquire_async_aux_semaphore(task: Optional[str]): + """Get a per-task, per-event-loop async semaphore after config lookup.""" + limit = _get_task_max_concurrency(task) + if limit is None: + return None + import asyncio + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return None + key = (task, id(loop)) + with _aux_sem_lock: + entry = _aux_async_semaphores.get(key) + if entry is None or entry[0] != limit: + semaphore = asyncio.Semaphore(limit) + _aux_async_semaphores[key] = (limit, semaphore) + return semaphore + return entry[1] + + +def _reset_aux_semaphores() -> None: + """Drop cached semaphores (test helper).""" + with _aux_sem_lock: + _aux_sync_semaphores.clear() + _aux_async_semaphores.clear() + + # --------------------------------------------------------------------------- # Anthropic-compatible endpoint detection + image block conversion # --------------------------------------------------------------------------- @@ -8455,6 +8526,73 @@ def call_llm( api_mode: str = None, stream: bool = False, stream_options: dict = None, +) -> Any: + """Run an auxiliary LLM request, applying the configured task limit.""" + semaphore = _acquire_sync_aux_semaphore(task) + if semaphore is not None: + semaphore.acquire() + try: + response = _call_llm_impl( + task=task, + provider=provider, + model=model, + base_url=base_url, + api_key=api_key, + main_runtime=main_runtime, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + tools=tools, + timeout=timeout, + extra_body=extra_body, + reasoning_config=reasoning_config, + api_mode=api_mode, + stream=stream, + stream_options=stream_options, + ) + if stream and semaphore is not None: + stream_semaphore = semaphore + semaphore = None + return _release_sync_semaphore_after_stream(response, stream_semaphore) + return response + finally: + if semaphore is not None: + semaphore.release() + + +def _release_sync_semaphore_after_stream( + stream: Any, semaphore: threading.BoundedSemaphore, +): + """Release a permit only after a streaming response is consumed or closed.""" + try: + yield from stream + finally: + try: + close = getattr(stream, "close", None) + if callable(close): + close() + finally: + semaphore.release() + + +def _call_llm_impl( + task: str = None, + *, + provider: str = None, + model: str = None, + base_url: str = None, + api_key: str = None, + main_runtime: Optional[Dict[str, Any]] = None, + messages: list, + temperature: Optional[float] = None, + max_tokens: int = None, + tools: list = None, + timeout: float = None, + extra_body: dict = None, + reasoning_config: Optional[dict] = None, + api_mode: str = None, + stream: bool = False, + stream_options: dict = None, ) -> Any: """Centralized synchronous LLM call. @@ -9218,6 +9356,47 @@ async def async_call_llm( timeout: float = None, extra_body: dict = None, reasoning_config: Optional[dict] = None, +) -> Any: + """Run an asynchronous auxiliary LLM request under the configured limit.""" + semaphore = _acquire_async_aux_semaphore(task) + if semaphore is not None: + await semaphore.acquire() + try: + return await _async_call_llm_impl( + task=task, + provider=provider, + model=model, + base_url=base_url, + api_key=api_key, + main_runtime=main_runtime, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + tools=tools, + timeout=timeout, + extra_body=extra_body, + reasoning_config=reasoning_config, + ) + finally: + if semaphore is not None: + semaphore.release() + + +async def _async_call_llm_impl( + task: str = None, + *, + provider: str = None, + model: str = None, + base_url: str = None, + api_key: str = None, + main_runtime: Optional[Dict[str, Any]] = None, + messages: list, + temperature: Optional[float] = None, + max_tokens: int = None, + tools: list = None, + timeout: float = None, + extra_body: dict = None, + reasoning_config: Optional[dict] = None, ) -> Any: """Centralized asynchronous LLM call. diff --git a/cli-config.yaml.example b/cli-config.yaml.example index 04ccfb6da4..6b7a4c327a 100644 --- a/cli-config.yaml.example +++ b/cli-config.yaml.example @@ -652,6 +652,25 @@ prompt_caching: # # extra_body: # # chat_template_kwargs: # # enable_thinking: false +# +# # Auto-generated short session titles after the first exchange. +# # Each active Discord/Telegram channel can spawn a background title +# # call. Cap concurrency to keep retries during provider incidents +# # from amplifying the request burst. Leave unset for legacy behavior +# # (unlimited). +# title_generation: +# provider: "auto" +# model: "" +# # max_concurrency: 2 # Optional: cap simultaneous title calls +# +# # Context compression — summarizes long sessions to shrink the prompt. +# # Heavy and often hits the slowest provider chain. Setting a small +# # cap prevents many sessions from compressing simultaneously during +# # provider degradation. Leave unset for legacy behavior (unlimited). +# compression: +# provider: "auto" +# model: "" +# # max_concurrency: 2 # Optional: cap simultaneous compression calls # ============================================================================= # Persistent Memory diff --git a/tests/agent/test_auxiliary_concurrency.py b/tests/agent/test_auxiliary_concurrency.py new file mode 100644 index 0000000000..27fccd97ec --- /dev/null +++ b/tests/agent/test_auxiliary_concurrency.py @@ -0,0 +1,403 @@ +"""Tests for per-task concurrency limiting on auxiliary LLM calls (#23324).""" + +import asyncio +import threading +import time +from unittest.mock import MagicMock, AsyncMock, patch + +import pytest + +from agent.auxiliary_client import ( + call_llm, + async_call_llm, + _acquire_sync_aux_semaphore, + _acquire_async_aux_semaphore, + _get_task_max_concurrency, + _reset_aux_semaphores, +) + + +@pytest.fixture(autouse=True) +def _clean_semaphore_cache(): + _reset_aux_semaphores() + yield + _reset_aux_semaphores() + + +class TestGetTaskMaxConcurrency: + def test_returns_none_for_missing_task(self): + assert _get_task_max_concurrency(None) is None + assert _get_task_max_concurrency("") is None + + def test_returns_none_when_unset(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", return_value={} + ): + assert _get_task_max_concurrency("title_generation") is None + + def test_does_not_reuse_vision_cpu_limit_for_llm_calls(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 1}, + ): + assert _get_task_max_concurrency("vision") is None + + def test_returns_int_when_configured(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 3}, + ): + assert _get_task_max_concurrency("compression") == 3 + + def test_returns_none_for_non_numeric(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": "not-a-number"}, + ): + assert _get_task_max_concurrency("compression") is None + + def test_returns_none_for_zero_or_negative(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 0}, + ): + assert _get_task_max_concurrency("compression") is None + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": -2}, + ): + assert _get_task_max_concurrency("compression") is None + + +class TestSemaphoreCache: + def test_sync_returns_none_when_unset(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", return_value={} + ): + assert _acquire_sync_aux_semaphore("title_generation") is None + + def test_sync_reuses_semaphore_for_same_limit(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 2}, + ): + sem1 = _acquire_sync_aux_semaphore("compression") + sem2 = _acquire_sync_aux_semaphore("compression") + assert sem1 is sem2 + + def test_sync_rebuilds_when_limit_changes(self): + cfg = {"max_concurrency": 2} + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value=cfg, + ): + sem1 = _acquire_sync_aux_semaphore("compression") + cfg["max_concurrency"] = 5 + sem2 = _acquire_sync_aux_semaphore("compression") + assert sem1 is not sem2 + + @pytest.mark.asyncio + async def test_async_reuses_semaphore_within_same_loop(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 2}, + ): + sem1 = _acquire_async_aux_semaphore("compression") + sem2 = _acquire_async_aux_semaphore("compression") + assert sem1 is sem2 + + def test_async_returns_none_with_no_running_loop(self): + with patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 2}, + ): + # Called outside an asyncio loop — should bail rather than crash. + assert _acquire_async_aux_semaphore("compression") is None + + +class TestSyncCallEnforcesLimit: + def test_call_llm_caps_concurrent_inflight(self): + limit = 2 + n_callers = 6 + + active = 0 + max_active = 0 + lock = threading.Lock() + + def fake_create(**kwargs): + nonlocal active, max_active + with lock: + active += 1 + if active > max_active: + max_active = active + try: + time.sleep(0.05) + finally: + with lock: + active -= 1 + return MagicMock() + + client = MagicMock() + client.base_url = "https://example.test/v1" + client.chat.completions.create.side_effect = fake_create + + with ( + patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=("openrouter", "test-model", None, None, None), + ), + patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "test-model"), + ), + patch( + "agent.auxiliary_client._validate_llm_response", + side_effect=lambda resp, _task, **_kwargs: resp, + ), + patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": limit}, + ), + ): + threads = [ + threading.Thread( + target=lambda: call_llm( + task="title_generation", + messages=[{"role": "user", "content": "hi"}], + ) + ) + for _ in range(n_callers) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5) + + assert max_active <= limit, f"observed {max_active} > limit {limit}" + assert client.chat.completions.create.call_count == n_callers + + def test_call_llm_unlimited_when_not_configured(self): + client = MagicMock() + client.base_url = "https://example.test/v1" + client.chat.completions.create.return_value = MagicMock() + + with ( + patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=("openrouter", "test-model", None, None, None), + ), + patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "test-model"), + ), + patch( + "agent.auxiliary_client._validate_llm_response", + side_effect=lambda resp, _task, **_kwargs: resp, + ), + patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={}, + ), + ): + # With no max_concurrency in config, no semaphore is acquired. + call_llm( + task="title_generation", + messages=[{"role": "user", "content": "hi"}], + ) + + assert client.chat.completions.create.call_count == 1 + + def test_semaphore_released_on_exception(self): + """Errors inside call_llm must release the semaphore so the next call proceeds.""" + client = MagicMock() + client.base_url = "https://example.test/v1" + client.chat.completions.create.side_effect = RuntimeError("boom") + + with ( + patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=("openrouter", "test-model", None, None, None), + ), + patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "test-model"), + ), + patch( + "agent.auxiliary_client._validate_llm_response", + side_effect=lambda resp, _task, **_kwargs: resp, + ), + patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 1}, + ), + ): + for _ in range(3): + with pytest.raises(RuntimeError, match="boom"): + call_llm( + task="title_generation", + messages=[{"role": "user", "content": "hi"}], + ) + + def test_stream_holds_permit_until_consumed_and_preserves_options(self): + client = MagicMock() + client.base_url = "https://example.test/v1" + client.chat.completions.create.side_effect = [iter(["chunk"]), MagicMock()] + second_call_started = threading.Event() + + def make_second_call(): + second_call_started.set() + call_llm( + task="compression", + messages=[{"role": "user", "content": "second"}], + ) + + with ( + patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=("openrouter", "test-model", None, None, None), + ), + patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "test-model"), + ), + patch( + "agent.auxiliary_client._validate_llm_response", + side_effect=lambda response, _task, **_kwargs: response, + ), + patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 1}, + ), + ): + stream = call_llm( + task="compression", + messages=[{"role": "user", "content": "first"}], + stream=True, + stream_options={"include_usage": True}, + ) + thread = threading.Thread(target=make_second_call) + thread.start() + assert second_call_started.wait(timeout=1) + time.sleep(0.05) + assert client.chat.completions.create.call_count == 1 + assert list(stream) == ["chunk"] + thread.join(timeout=1) + + assert not thread.is_alive() + assert client.chat.completions.create.call_count == 2 + assert client.chat.completions.create.call_args_list[0].kwargs["stream"] is True + assert client.chat.completions.create.call_args_list[0].kwargs["stream_options"] == { + "include_usage": True + } + + def test_api_mode_is_forwarded_to_client_resolution(self): + client = MagicMock() + client.base_url = "https://example.test/v1" + client.chat.completions.create.return_value = MagicMock() + + with ( + patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=("openrouter", "test-model", None, None, None), + ), + patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "test-model"), + ) as get_client, + patch( + "agent.auxiliary_client._validate_llm_response", + side_effect=lambda response, _task, **_kwargs: response, + ), + ): + call_llm( + task="title_generation", + messages=[{"role": "user", "content": "hi"}], + api_mode="codex_responses", + ) + + assert get_client.call_args.kwargs["api_mode"] == "codex_responses" + + +class TestAsyncCallEnforcesLimit: + @pytest.mark.asyncio + async def test_async_call_llm_caps_concurrent_inflight(self): + limit = 2 + n_callers = 6 + + active = 0 + max_active = 0 + + async def fake_create(**kwargs): + nonlocal active, max_active + active += 1 + if active > max_active: + max_active = active + try: + await asyncio.sleep(0.05) + finally: + active -= 1 + return MagicMock() + + client = MagicMock() + client.base_url = "https://example.test/v1" + client.chat.completions.create = AsyncMock(side_effect=fake_create) + + with ( + patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=("openrouter", "test-model", None, None, None), + ), + patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "test-model"), + ), + patch( + "agent.auxiliary_client._validate_llm_response", + side_effect=lambda resp, _task, **_kwargs: resp, + ), + patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": limit}, + ), + ): + await asyncio.gather(*[ + async_call_llm( + task="compression", + messages=[{"role": "user", "content": "hi"}], + ) + for _ in range(n_callers) + ]) + + assert max_active <= limit, f"observed {max_active} > limit {limit}" + assert client.chat.completions.create.await_count == n_callers + + @pytest.mark.asyncio + async def test_async_semaphore_released_on_exception(self): + client = MagicMock() + client.base_url = "https://example.test/v1" + client.chat.completions.create = AsyncMock(side_effect=RuntimeError("boom")) + + with ( + patch( + "agent.auxiliary_client._resolve_task_provider_model", + return_value=("openrouter", "test-model", None, None, None), + ), + patch( + "agent.auxiliary_client._get_cached_client", + return_value=(client, "test-model"), + ), + patch( + "agent.auxiliary_client._validate_llm_response", + side_effect=lambda resp, _task, **_kwargs: resp, + ), + patch( + "agent.auxiliary_client._get_auxiliary_task_config", + return_value={"max_concurrency": 1}, + ), + ): + for _ in range(3): + with pytest.raises(RuntimeError, match="boom"): + await async_call_llm( + task="compression", + messages=[{"role": "user", "content": "hi"}], + ) diff --git a/website/docs/user-guide/configuration.md b/website/docs/user-guide/configuration.md index b1fddd57c3..cfe85dbd55 100644 --- a/website/docs/user-guide/configuration.md +++ b/website/docs/user-guide/configuration.md @@ -1189,6 +1189,8 @@ auxiliary: # model: google/gemini-2.5-flash # base_url: "" # api_key: "" + # max_concurrency: 2 # Optional: cap simultaneous compression LLM calls so + # multiple sessions don't pile retries on a degraded provider # Auto-generated session titles. Empty language follows the conversation; # set e.g. "English" or "Japanese" to pin titles to one language. @@ -1217,6 +1219,15 @@ auxiliary: api_key: "" timeout: 30 + # Auto-generated short session titles after the first exchange + title_generation: + provider: "auto" + model: "" + base_url: "" + api_key: "" + timeout: 30 + # max_concurrency: 2 # Optional: cap simultaneous title-generation calls + # Kanban triage specifier — `hermes kanban specify ` (or the # dashboard's ✨ Specify button on Triage-column cards) uses this # slot to expand a one-liner into a concrete spec and promote the @@ -1266,6 +1277,25 @@ Each entry supports the same three knobs as any auxiliary task config: `fallback_chain` is available on any auxiliary task — `compression`, `vision`, `web_extract`, `approval`, `skills_hub`, `mcp`, etc. +### Limiting auxiliary concurrency + +`max_concurrency` caps in-flight LLM calls for auxiliary tasks such as `compression` and `title_generation` across the whole process. `auxiliary.vision.max_concurrency` is excluded: it already controls only vision's CPU-bound image encode/resize workers, not LLM requests. This is most useful when: + +- Many sessions can spawn background work simultaneously (Discord/Telegram channels, multiple terminals) +- Your provider is rate-limited or going through an incident and retries would amplify the burst + +The default is unlimited. A typical safety cap is `2`: + +```yaml +auxiliary: + title_generation: + max_concurrency: 2 + compression: + max_concurrency: 2 +``` + +The semaphore wraps the entire call including retries and fallbacks, so a single slow call counts only once toward the limit. + ### OpenRouter routing & Pareto Code for auxiliary tasks When an auxiliary task resolves to OpenRouter (either explicitly or via `provider: "main"` while your main agent is on OpenRouter), the main agent's `provider_routing` and `openrouter.min_coding_score` settings **do not propagate** — by design, each auxiliary task is independent. To set OpenRouter provider preferences or use the [Pareto Code router](/integrations/providers#openrouter-pareto-code-router) for a specific aux task, set them per-task via `extra_body`: