Skip to content

Commit 02ed575

Browse files
Pin LiteLLM to patched 1.84.3 release
Upgrade the SDK and proxy to address published vulnerabilities while keeping async logging safe across short-lived event loops. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 2c08671 commit 02ed575

9 files changed

Lines changed: 446 additions & 146 deletions

File tree

eval_protocol/litellm_compat.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import asyncio
2+
3+
4+
async def allow_litellm_logging_to_start() -> None:
5+
"""Let LiteLLM claim queued logging callbacks before the event loop can close.
6+
7+
LiteLLM 1.80+ queues coroutine objects in a process-global logging worker.
8+
When pytest replaces a function-scoped event loop before the worker claims a
9+
callback, LiteLLM drops that coroutine while rebinding the queue. Two event
10+
loop turns let the worker wrap the callback in a task so normal loop
11+
shutdown can cancel it cleanly.
12+
"""
13+
await asyncio.sleep(0)
14+
await asyncio.sleep(0)

eval_protocol/mcp/execution/policy.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
from litellm.caching.in_memory_cache import InMemoryCache
1919
from litellm.caching.redis_cache import RedisCache
2020

21+
from eval_protocol.litellm_compat import allow_litellm_logging_to_start
22+
2123
from .base_policy import LLMBasePolicy
2224

2325
logger = logging.getLogger(__name__)
@@ -205,6 +207,8 @@ async def _make_llm_call(self, messages: List[Dict[str, Any]], tools: List[Dict[
205207
else:
206208
response = await acompletion(model=self.model_id, **request_params)
207209

210+
await allow_litellm_logging_to_start()
211+
208212
assert response is not None, "Response is None"
209213
assert isinstance(response, ModelResponse), "Response should be ModelResponse"
210214

eval_protocol/proxy/docker-compose.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ services:
1515

1616
# LiteLLM Backend - Handles actual LLM proxying
1717
litellm-backend:
18-
image: litellm/litellm:v1.77.3-stable
18+
image: litellm/litellm:v1.84.3
1919
platform: linux/amd64
2020
container_name: litellm-backend
2121
command: ["--config", "/app/config.yaml", "--port", "4000", "--host", "0.0.0.0"]

eval_protocol/pytest/default_single_turn_rollout_process.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
1414

1515
from eval_protocol.dataset_logger import default_logger
16+
from eval_protocol.litellm_compat import allow_litellm_logging_to_start
1617
from eval_protocol.models import EvaluationRow, Message
1718
from openai.types import CompletionUsage
1819
from eval_protocol.pytest.rollout_processor import RolloutProcessor
@@ -129,6 +130,8 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
129130
else:
130131
response = await acompletion(**request_params)
131132

133+
await allow_litellm_logging_to_start()
134+
132135
assert response is not None, "Response is None"
133136
assert isinstance(response, ModelResponse), "Response should be ModelResponse"
134137
assert isinstance(response.choices[0], Choices), "Response choice should be a Choices"

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ dependencies = [
3131
"omegaconf>=2.3.0",
3232
"httpx>=0.24.0",
3333
"anthropic>=0.59.0",
34-
"litellm<1.75.0",
34+
"litellm==1.84.3",
3535
"pytest>=6.0.0",
3636
"pytest-asyncio>=0.21.0",
3737
"peewee>=3.18.2",
Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
import asyncio
2+
3+
import litellm
4+
import pytest
5+
6+
from eval_protocol.dataset_logger import default_logger
7+
from eval_protocol.litellm_compat import allow_litellm_logging_to_start
8+
from eval_protocol.mcp.execution.policy import LiteLLMPolicy
9+
from eval_protocol.models import EvaluationRow, Message
10+
from eval_protocol.pytest.default_single_turn_rollout_process import SingleTurnRolloutProcessor
11+
from eval_protocol.pytest.exception_config import get_default_exception_handler_config
12+
from eval_protocol.pytest.types import RolloutProcessorConfig
13+
from vendor.tau2.data_model.message import UserMessage
14+
from vendor.tau2.utils.llm_utils import generate
15+
16+
17+
@pytest.mark.parametrize("index", range(4))
18+
@pytest.mark.asyncio
19+
async def test_acompletion_across_pytest_event_loops(index: int) -> None:
20+
response = await litellm.acompletion(
21+
model="openai/gpt-4o-mini",
22+
messages=[{"role": "user", "content": "ping"}],
23+
api_key="test",
24+
mock_response=f"ok-{index}",
25+
)
26+
await allow_litellm_logging_to_start()
27+
28+
assert response.choices[0].message.content == f"ok-{index}"
29+
30+
31+
@pytest.mark.parametrize(("index", "stream"), [(0, False), (1, True), (2, False), (3, True)])
32+
@pytest.mark.asyncio
33+
async def test_single_turn_processor_across_pytest_event_loops(index: int, stream: bool) -> None:
34+
config = RolloutProcessorConfig(
35+
completion_params={
36+
"model": "openai/gpt-4o-mini",
37+
"api_key": "test",
38+
"mock_response": f"single-turn-{index}",
39+
"stream": stream,
40+
},
41+
mcp_config_path="",
42+
semaphore=asyncio.Semaphore(1),
43+
server_script_path=None,
44+
steps=1,
45+
logger=default_logger,
46+
exception_handler_config=get_default_exception_handler_config(),
47+
)
48+
row = EvaluationRow(messages=[Message(role="user", content="ping")])
49+
50+
result = await SingleTurnRolloutProcessor()([row], config)[0]
51+
52+
assert result.messages[-1].content == f"single-turn-{index}"
53+
54+
55+
@pytest.mark.parametrize(("index", "stream"), [(0, False), (1, True), (2, False), (3, True)])
56+
@pytest.mark.asyncio
57+
async def test_litellm_policy_across_pytest_event_loops(index: int, stream: bool) -> None:
58+
policy = LiteLLMPolicy(
59+
model_id="openai/gpt-4o-mini",
60+
use_caching=False,
61+
api_key="test",
62+
mock_response=f"policy-{index}",
63+
stream=stream,
64+
)
65+
66+
result = await policy._make_llm_call([{"role": "user", "content": "ping"}], tools=[])
67+
68+
assert result["choices"][0]["message"]["content"] == f"policy-{index}"
69+
70+
71+
@pytest.mark.parametrize("index", range(2))
72+
@pytest.mark.asyncio
73+
async def test_tau2_generate_across_pytest_event_loops(index: int) -> None:
74+
result = await generate(
75+
model="openai/gpt-4o-mini",
76+
messages=[UserMessage(role="user", content="ping")],
77+
api_key="test",
78+
mock_response=f"tau2-{index}",
79+
)
80+
81+
assert result.content == f"tau2-{index}"

tests/test_litellm_policy_provider_fields.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ async def test_litellm_policy_surfaces_provider_specific_reasoning_details(monke
1616
# Define a fake ModelResponse base class and patch the module's ModelResponse
1717
class FakeModelResponseBase: ...
1818

19-
policy_mod.ModelResponse = FakeModelResponseBase
19+
monkeypatch.setattr(policy_mod, "ModelResponse", FakeModelResponseBase)
2020

2121
async def fake_acompletion(*args, **kwargs):
2222
# This mimics the LiteLLM Message object shape we rely on in policy._make_llm_call

0 commit comments

Comments
 (0)