mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
fix(mcp): preserve request-only peer log consumers
This commit is contained in:
parent
90d4357011
commit
86c932d155
3 changed files with 36 additions and 4 deletions
|
|
@ -44,13 +44,17 @@ class McpPeer:
|
|||
spec_path: Path | None = None
|
||||
consumed: list[int] = field(default_factory=lambda: [0])
|
||||
|
||||
def drain(self) -> tuple[dict[str, object], ...]:
|
||||
def drain(self, *, include_negotiation: bool = False) -> tuple[dict[str, object], ...]:
|
||||
if self.record is not None:
|
||||
lines: Final = self.record.read_text().splitlines() if self.record.exists() else []
|
||||
fresh: Final = tuple(json.loads(line) for line in lines[self.consumed[0] :])
|
||||
self.consumed[0] = len(lines)
|
||||
return fresh
|
||||
return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize()))
|
||||
return tuple(item for item in fresh if include_negotiation or "negotiation" not in item)
|
||||
return tuple(
|
||||
item
|
||||
for item in (self.calls.get_nowait() for _ in range(self.calls.qsize()))
|
||||
if include_negotiation or "negotiation" not in item
|
||||
)
|
||||
|
||||
def registration(self) -> dict[str, object]:
|
||||
if self.transport == "stdio":
|
||||
|
|
|
|||
|
|
@ -231,7 +231,7 @@ def test_pinned_revision_pairs_list_and_call_through_gateway(
|
|||
|
||||
peer.drain()
|
||||
negotiated: Final = asyncio.run(exercise())
|
||||
observed: Final = peer.drain()
|
||||
observed: Final = peer.drain(include_negotiation=True)
|
||||
negotiations: Final = tuple(
|
||||
(item["negotiation"]["requested"], item["negotiation"]["returned"])
|
||||
for item in observed
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from unittest.mock import Mock, patch
|
||||
|
|
@ -50,6 +51,33 @@ async def test_negotiation_records_preserve_the_peer_call_contract() -> None:
|
|||
assert records[0]["negotiation"] == {"requested": "2025-03-26", "returned": "2025-03-26"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("file_backed", (False, True))
|
||||
def test_peer_drain_preserves_received_requests(tmp_path: Path, file_backed: bool) -> None:
|
||||
from integration._support.mcp import McpPeer
|
||||
|
||||
request: Final = {"body": {"method": "tools/call"}, "headers": {"authorization": "synthetic"}}
|
||||
negotiation: Final = {"body": {}, "headers": {}, "negotiation": {"requested": "2025-03-26", "returned": "2025-03-26"}}
|
||||
records: Final = (request, negotiation)
|
||||
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
|
||||
path: Final = tmp_path / "peer.jsonl" if file_backed else None
|
||||
if path is not None:
|
||||
path.write_text("".join(json.dumps(record) + "\n" for record in records))
|
||||
else:
|
||||
for record in records:
|
||||
observed.put(record)
|
||||
peer: Final = McpPeer("http://peer", observed, record=path)
|
||||
assert peer.drain() == (request,)
|
||||
assert peer.drain() == ()
|
||||
if path is not None:
|
||||
with path.open("a") as sink:
|
||||
sink.write("".join(json.dumps(record) + "\n" for record in records))
|
||||
else:
|
||||
for record in records:
|
||||
observed.put(record)
|
||||
assert peer.drain(include_negotiation=True) == records
|
||||
assert peer.drain(include_negotiation=True) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("explicit", (False, True))
|
||||
def test_owned_proxy_isolates_automatic_coverage_unless_requested(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, explicit: bool
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue