mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(anthropic_messages): dispatch deferred spend logging when the client disconnects mid-relay
When the pump finishes draining while the client is still connected, billing is deferred to the proxy's post-response hook, which only fires on a normally completed response. A client disconnect before the relay consumed the queued tail tore the generator down past that hook, so the request logged no spend at all. The relay teardown now dispatches the stored deferred billing whenever it never reached the end-of-stream sentinel. Also drops the live pass_through_tests script: that CI job runs against a fixed config with no Bedrock model or AWS credentials, so it could only fail there. The scenario is covered by unit tests on the relay/pump seam.
This commit is contained in:
mateo-berri
parent
6e59ce1773
commit
0e78c5bff7
3 changed files with 75 additions and 156 deletions
|
|
@ -480,16 +480,37 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
_UPSTREAM_PUMP_TASKS.add(pump_task)
|
||||
pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard)
|
||||
|
||||
reached_end = False # rebind-ok: flipped once the relay consumes the end-of-stream sentinel
|
||||
try:
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is None:
|
||||
reached_end = True
|
||||
break
|
||||
if isinstance(item, BaseException):
|
||||
raise item
|
||||
yield item
|
||||
finally:
|
||||
client_detached.set()
|
||||
if not reached_end:
|
||||
self._dispatch_pending_deferred_logging()
|
||||
|
||||
def _dispatch_pending_deferred_logging(self) -> None:
|
||||
"""Fire deferred billing that a torn-down response would otherwise drop.
|
||||
|
||||
When the pump finishes draining while the client is still connected it
|
||||
stores the logging coroutine for ProxyLogging._fire_deferred_stream_logging,
|
||||
which the proxy only fires on a normally completed response: a client
|
||||
disconnect (GeneratorExit / CancelledError) re-raises past it. Without
|
||||
this dispatch that window loses the spend row entirely.
|
||||
"""
|
||||
deferred_cb: Final = getattr(self.litellm_logging_obj, "_on_deferred_stream_complete", None)
|
||||
deferred_args: Final = getattr(self.litellm_logging_obj, "_deferred_stream_complete_args", None)
|
||||
if deferred_cb is None or deferred_args is None:
|
||||
return
|
||||
self.litellm_logging_obj._on_deferred_stream_complete = None
|
||||
self.litellm_logging_obj._deferred_stream_complete_args = None
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=deferred_cb(*deferred_args))
|
||||
|
||||
async def _bill_collected_chunks(
|
||||
self,
|
||||
|
|
@ -541,9 +562,10 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
return False
|
||||
try:
|
||||
queue.put_nowait(item)
|
||||
return True
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
else:
|
||||
return True
|
||||
put_task: Final = asyncio.ensure_future(queue.put(item))
|
||||
detached_task: Final = asyncio.ensure_future(client_detached.wait())
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1,155 +0,0 @@
|
|||
"""
|
||||
Regression test: /v1/messages streaming interrupted mid-stream must still
|
||||
produce a spend-log entry.
|
||||
|
||||
On v1.79.1 the proxy records spend for the partially-streamed request.
|
||||
A refactor on `main` broke that path, so the same scenario now produces
|
||||
zero spend-log rows.
|
||||
|
||||
Run against a live proxy (e.g. ``litellm --config proxy_server_config.yaml``):
|
||||
|
||||
pytest tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py -s
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
|
||||
|
||||
BASE_URL = "http://127.0.0.1:4000" # change appropriately
|
||||
ADMIN_KEY = "sk-1234"
|
||||
|
||||
|
||||
async def _generate_key(session: aiohttp.ClientSession) -> str:
|
||||
"""Create a fresh virtual key so spend is isolated."""
|
||||
url = f"{BASE_URL}/key/generate"
|
||||
headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"}
|
||||
async with session.post(url, headers=headers, json={"models": []}) as resp:
|
||||
assert resp.status == 200, f"key/generate failed: {await resp.text()}"
|
||||
data = await resp.json()
|
||||
return data["key"]
|
||||
|
||||
|
||||
async def _get_spend_logs_by_spend_id(session: aiohttp.ClientSession, api_key: str, spend_id: str):
|
||||
"""Query /spend/logs by api_key then filter by spend_id in metadata."""
|
||||
url = f"{BASE_URL}/spend/logs?api_key={api_key}"
|
||||
headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"}
|
||||
async with session.get(url, headers=headers) as resp:
|
||||
assert resp.status == 200, f"spend/logs failed: {await resp.text()}"
|
||||
all_logs = await resp.json()
|
||||
if not isinstance(all_logs, list):
|
||||
return []
|
||||
matched = []
|
||||
for log in all_logs:
|
||||
meta = log.get("metadata")
|
||||
if isinstance(meta, str):
|
||||
meta = json.loads(meta)
|
||||
if isinstance(meta, dict):
|
||||
slm = meta.get("spend_logs_metadata") or {}
|
||||
if slm.get("spend_id") == spend_id:
|
||||
matched.append(log)
|
||||
return matched
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.flaky(retries=3, delay=2)
|
||||
async def test_v1_messages_streaming_disconnect_has_spend_log():
|
||||
"""
|
||||
1. Send a streaming POST to /v1/messages.
|
||||
2. Read a few SSE chunks, then close the connection (simulating a client
|
||||
disconnect / interruption).
|
||||
3. Wait for the proxy's async spend-tracking pipeline to flush.
|
||||
4. Assert that at least one spend-log row exists for the request.
|
||||
|
||||
This PASSES on v1.79.1 and FAILS on the latest main branch.
|
||||
"""
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(total=60)
|
||||
) as session:
|
||||
key = await _generate_key(session)
|
||||
|
||||
spend_id = str(uuid.uuid4())
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
"x-litellm-spend-logs-metadata": '{"spend_id": "' + spend_id + '"}',
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"max_tokens": 3000,
|
||||
"stream": True,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
f"Write several detailed paragraphs (at least 500 words) about the "
|
||||
f"history of the Roman Empire. Unique id: {uuid.uuid4()}"
|
||||
),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
chunks_read = 0
|
||||
|
||||
async with session.post(
|
||||
f"{BASE_URL}/v1/messages", json=payload, headers=headers
|
||||
) as resp:
|
||||
assert resp.status == 200, f"/v1/messages failed: {await resp.text()}"
|
||||
|
||||
async for raw_line in resp.content:
|
||||
line = raw_line.decode("utf-8", errors="replace").strip()
|
||||
if not line:
|
||||
continue
|
||||
chunks_read += 1
|
||||
print(f" chunk #{chunks_read}: {line[:120]}")
|
||||
if chunks_read >= 5:
|
||||
break
|
||||
|
||||
assert chunks_read >= 3, (
|
||||
f"Expected at least 3 chunks before disconnect, got {chunks_read}"
|
||||
)
|
||||
|
||||
print(
|
||||
f"\nDisconnected after {chunks_read} chunks. "
|
||||
f"Waiting for spend pipeline to flush …"
|
||||
)
|
||||
|
||||
spend_data = None
|
||||
max_retries = 4
|
||||
for attempt in range(1, max_retries + 1):
|
||||
await asyncio.sleep(10)
|
||||
print(f" spend-log poll attempt {attempt}/{max_retries}")
|
||||
spend_data = await _get_spend_logs_by_spend_id(session, key, spend_id)
|
||||
if spend_data and len(spend_data) > 0:
|
||||
print(f" ✓ found {len(spend_data)} spend-log row(s)")
|
||||
break
|
||||
print(" … not found yet")
|
||||
|
||||
assert spend_data is not None and len(spend_data) > 0, (
|
||||
f"No spend-log entry found for spend_id={spend_id} "
|
||||
f"after streaming disconnect. "
|
||||
f"This is the regression: interrupted /v1/messages streams must "
|
||||
f"still record spend."
|
||||
)
|
||||
|
||||
log_entry = spend_data[0]
|
||||
print(
|
||||
f"\nSpend-log entry:\n{json.dumps(log_entry, indent=2, default=str)}"
|
||||
)
|
||||
|
||||
prompt_tokens = log_entry.get("prompt_tokens", 0)
|
||||
completion_tokens = log_entry.get("completion_tokens", 0)
|
||||
assert prompt_tokens > 0, (
|
||||
"Spend-log row exists but has zero prompt tokens, so usage was not recorded."
|
||||
)
|
||||
assert completion_tokens >= 100, (
|
||||
f"Spend-log completion_tokens={completion_tokens} is far below the full "
|
||||
f"response Bedrock generated and billed. The interrupted stream was billed "
|
||||
f"on the few chunks the client drained, not the full upstream output. "
|
||||
f"chunks_read={chunks_read}, prompt_tokens={prompt_tokens}"
|
||||
)
|
||||
|
|
@ -472,6 +472,58 @@ async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all():
|
|||
assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_dispatches_deferred_logging_when_client_disconnects_mid_tail():
|
||||
"""
|
||||
Regression: when the pump finishes draining while the client is still
|
||||
connected, ``_handle_streaming_logging`` defers billing for the proxy's
|
||||
post-response hook (``ProxyLogging._fire_deferred_stream_logging``), which
|
||||
only fires on a normally completed response. If the client then disconnects
|
||||
before consuming the queued tail, the response generator tears down via
|
||||
GeneratorExit and that hook never runs. The relay teardown must dispatch
|
||||
the stored deferred billing itself, or the request logs no spend at all.
|
||||
"""
|
||||
dispatched = []
|
||||
deferred_fired = asyncio.Event()
|
||||
|
||||
def _deferred_stream_complete(logging_coroutine):
|
||||
dispatched.append(logging_coroutine)
|
||||
|
||||
async def _consume():
|
||||
logging_coroutine.close()
|
||||
deferred_fired.set()
|
||||
|
||||
return _consume()
|
||||
|
||||
logging_obj = _make_logging_obj("test_deferred_dispatch_on_disconnect_mid_tail")
|
||||
logging_obj._on_deferred_stream_complete = _deferred_stream_complete
|
||||
iterator = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body={})
|
||||
|
||||
async def _full_stream():
|
||||
for event in (*_STREAM_PREFIX, *_STREAM_TAIL):
|
||||
yield event
|
||||
|
||||
gen = iterator.async_sse_wrapper(_full_stream())
|
||||
client_chunks = []
|
||||
async for chunk in gen:
|
||||
client_chunks.append(chunk)
|
||||
if len(client_chunks) == len(_STREAM_PREFIX):
|
||||
break
|
||||
|
||||
for _ in range(100):
|
||||
if getattr(logging_obj, "_deferred_stream_complete_args", None) is not None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert getattr(logging_obj, "_deferred_stream_complete_args", None) is not None, "pump never deferred billing"
|
||||
|
||||
await gen.aclose()
|
||||
|
||||
assert len(dispatched) == 1, "relay teardown did not dispatch the deferred billing"
|
||||
assert logging_obj._on_deferred_stream_complete is None
|
||||
assert logging_obj._deferred_stream_complete_args is None
|
||||
await asyncio.wait_for(deferred_fired.wait(), timeout=5)
|
||||
|
||||
|
||||
class _ProviderStreamError(Exception):
|
||||
"""Stand-in for a provider-specific streaming failure carrying a status code."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue