From 323d7b0abe897ed22ec213c1d1fddc5800726717 Mon Sep 17 00:00:00 2001 From: Kiran Kashalkar Date: Thu, 20 Aug 2026 20:42:20 +0530 Subject: [PATCH] fix(needlepath): bound per-request selection fan-out A request carrying arbitrarily many eligible tool messages previously opened one concurrent outbound selection call per message. Cap targets at 16 per request (largest messages first, where selection pays off most; the rest are forwarded untouched, fail-open as ever) and bound in-flight calls with a semaphore of 4 so one request cannot monopolise the proxy's shared HTTP pool or burn selection quota. Adds tests for the cap (only the largest messages selected, one call per selected message) and the concurrency high-water mark. --- .../guardrail_hooks/needlepath/needlepath.py | 59 ++++++++++++---- .../guardrail_hooks/test_needlepath.py | 68 +++++++++++++++++++ 2 files changed, 115 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py index b97259c3f6b..13c998fc163 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py +++ b/litellm/proxy/guardrails/guardrail_hooks/needlepath/needlepath.py @@ -72,6 +72,16 @@ DEFAULT_MIN_CHARS_TO_SELECT: Final = 500 # to hold an inbound LLM request behind an optional optimisation. A stall past # this bound is a decline, and the original message is forwarded. _SELECT_TIMEOUT_SECONDS: Final = 30.0 +# A single proxy request can carry arbitrarily many eligible messages, and each +# one becomes an outbound selection call. These two bounds keep a pathological +# request (say, thousands of minimally qualifying tool outputs) from +# monopolising the proxy's shared HTTP pool or burning selection quota: at most +# _MAX_TARGETS_PER_REQUEST messages are selected per request -- the largest +# first, where selection pays off most -- and at most +# _MAX_CONCURRENT_SELECTIONS calls are in flight at once. Messages past the +# cap are forwarded untouched, the same fail-open outcome as any decline. +_MAX_TARGETS_PER_REQUEST: Final = 16 +_MAX_CONCURRENT_SELECTIONS: Final = 4 # The service reports a deliberate no-op through the gate. Any reason under this # prefix means "the engine chose not to select"; the original content is what # the caller should send. @@ -516,31 +526,56 @@ class NeedlepathGuardrail(CustomGuardrail): # One (idx, query) pair per eligible message whose query is non-blank. # Selection is conditioned on a query; without one there is nothing to # select against, so that message is left exactly as it arrived. - selected_pairs: Final = tuple( + eligible_pairs: Final = tuple( (idx, query) for idx in self._select_targets(messages, query_idx) if (query := _query_for_target(messages, idx, fallback_query)).strip() ) - if not selected_pairs: + if not eligible_pairs: verbose_proxy_logger.debug("Needlepath: no messages eligible for selection") return inputs - start_time: Final = time.monotonic() - # One request per message: each carries its own query, and the service - # renders one block per call. Running them concurrently keeps the added - # latency at roughly one round trip rather than one per message. - selections: Final = await asyncio.gather( - *( - self._selected_text( + # Cap the fan-out: the largest messages are kept because that is where + # selection saves the most, and the choice must not depend on message + # order in the request. Everything past the cap is forwarded untouched. + selected_pairs: Final = ( + eligible_pairs + if len(eligible_pairs) <= _MAX_TARGETS_PER_REQUEST + else tuple( + sorted( + eligible_pairs, + key=lambda pair: len(content_to_text(messages[pair[0]].get("content"))), + reverse=True, + )[:_MAX_TARGETS_PER_REQUEST] + ) + ) + if len(selected_pairs) < len(eligible_pairs): + verbose_proxy_logger.debug( + "Needlepath: %d eligible messages, selecting only the %d largest", + len(eligible_pairs), + _MAX_TARGETS_PER_REQUEST, + ) + + semaphore: Final = asyncio.Semaphore(_MAX_CONCURRENT_SELECTIONS) + + async def _bounded_selected_text(idx: int, query: str) -> str | None: + # The semaphore bounds how many selection calls this one proxy + # request holds open at a time; see _MAX_CONCURRENT_SELECTIONS. + async with semaphore: + return await self._selected_text( text=content_to_text(messages[idx].get("content")), query=query, title=_title_for(messages, idx), source=messages[idx].get("tool_call_id"), kind=_record_kind(messages[idx].get("role")), ) - for idx, query in selected_pairs - ) - ) + + start_time: Final = time.monotonic() + # One request per message: each carries its own query, and the service + # renders one block per call. Running them concurrently (up to the + # semaphore's bound) keeps the added latency near one round trip + # rather than one per message. + selections: Final = await asyncio.gather(*(_bounded_selected_text(idx, query) for idx, query in selected_pairs)) end_time: Final = time.monotonic() # Needs in-place index assignment below to build the edited copy without diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py index b8621d2c43b..1a019595042 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_needlepath.py @@ -13,6 +13,7 @@ Tests cover: the messages come back byte-identical. """ +import asyncio import copy from unittest.mock import AsyncMock, MagicMock, patch @@ -20,6 +21,8 @@ import httpx import pytest from litellm.proxy.guardrails.guardrail_hooks.needlepath.needlepath import ( + _MAX_CONCURRENT_SELECTIONS, + _MAX_TARGETS_PER_REQUEST, DEFAULT_MAX_CONTEXT_TOKENS, DEFAULT_MIN_CHARS_TO_SELECT, DEFAULT_OPERATING_POINT, @@ -395,3 +398,68 @@ async def test_one_message_declining_does_not_block_another(): assert out[TOOL_INDEX]["content"] == SELECTED_BLOCK # The system message stood down, so it is byte-identical. assert out[0] == AGENT_MESSAGES[0] + + +# ── fan-out bounds ──────────────────────────────────────────────────── + + +def _many_tool_messages(count: int) -> list: + """A user question plus `count` eligible tool outputs of strictly increasing size.""" + messages = [{"role": "user", "content": USER_QUESTION}] + for i in range(count): + messages.append( + { + "role": "tool", + "tool_call_id": f"call_{i}", + # All above the min-chars threshold; index i is the (i+1)-th smallest. + "content": "a" * (DEFAULT_MIN_CHARS_TO_SELECT + 20 + i), + } + ) + return messages + + +@pytest.mark.asyncio +async def test_target_cap_selects_only_the_largest_messages(guardrail: NeedlepathGuardrail): + """Past the per-request cap, only the largest messages are selected. + + The smallest overflow messages come back byte-identical and no service + call is made for them. + """ + overflow = 4 + messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST + overflow) + + post = AsyncMock(return_value=_select_response()) + result = await _run(guardrail, post, messages=messages) + + assert post.await_count == _MAX_TARGETS_PER_REQUEST + out = result["structured_messages"] + # Tool messages start at index 1 and grow with index: the first `overflow` + # are the smallest, so they are the ones left untouched. + for idx in range(1, 1 + overflow): + assert out[idx] == messages[idx] + for idx in range(1 + overflow, len(messages)): + assert out[idx]["content"] == SELECTED_BLOCK + + +@pytest.mark.asyncio +async def test_concurrent_selections_are_bounded(guardrail: NeedlepathGuardrail): + """No more than _MAX_CONCURRENT_SELECTIONS service calls are in flight at once.""" + in_flight = 0 + high_water = 0 + + async def _tracking_post(*args, **kwargs): + nonlocal in_flight, high_water + in_flight += 1 + high_water = max(high_water, in_flight) + # Yield twice so every scheduled call gets a chance to start before + # this one finishes; without the semaphore all of them would overlap. + await asyncio.sleep(0) + await asyncio.sleep(0) + in_flight -= 1 + return _select_response() + + messages = _many_tool_messages(_MAX_TARGETS_PER_REQUEST) + result = await _run(guardrail, _tracking_post, messages=messages) + + assert high_water <= _MAX_CONCURRENT_SELECTIONS + assert all(m["content"] == SELECTED_BLOCK for m in result["structured_messages"][1:])