From f280927452d96c5b36aef92a9e78a51dd67d9f77 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 12 Sep 2026 15:20:43 -0700 Subject: [PATCH] fix(memory): preserve keepalives and retry delays across tool rounds --- deploy/memory-pilot/README.md | 10 ++ deploy/memory-pilot/proxy_config.yaml | 1 + deploy/memory-pilot/start.sh | 2 +- litellm/proxy/memory/gateway.py | 26 ++++- .../proxy/memory/test_memory_v2_boundaries.py | 110 ++++++++++++++++++ 5 files changed, 144 insertions(+), 5 deletions(-) diff --git a/deploy/memory-pilot/README.md b/deploy/memory-pilot/README.md index 3beda2eac3a..255ae15c67c 100644 --- a/deploy/memory-pilot/README.md +++ b/deploy/memory-pilot/README.md @@ -33,6 +33,8 @@ pilot must not be connected to an older gateway's production database. 3. Set the build command to `bash deploy/memory-pilot/build.sh`, the start command to `bash deploy/memory-pilot/start.sh`, and the health path to `/health/readiness`. The build includes the dashboard from this branch. + Set the service's maximum shutdown delay to 300 seconds so active requests + can drain during a deployment. Uvicorn allows 290 seconds before cleanup 4. Set these environment variables in Render: | Variable | Value | @@ -79,6 +81,14 @@ correct, or delete entries in Memory; callers can use the self-service API. ## Behavior and limits +- Streaming keeps LiteLLM's configured SSE keepalives across silent memory + rounds. The pilot sends comments every 15 seconds of silence and disables + proxy buffering. A failure after streaming starts arrives as a native SSE + error; before streaming starts, HTTP errors retain their retry delay +- Model calls retain LiteLLM's normal timeout and retry settings. The separate + upstream credential check has a 20-second timeout. Deployments drain existing + requests for up to five minutes; requests still running after that can be + interrupted. Schedule pilot updates outside active office usage - Supported surfaces: Chat Completions, Responses, and Anthropic Messages, including their native streaming responses and client tool continuation. - The selected model must support function calling. The actual answering model diff --git a/deploy/memory-pilot/proxy_config.yaml b/deploy/memory-pilot/proxy_config.yaml index a595915067a..160f98e1d0b 100644 --- a/deploy/memory-pilot/proxy_config.yaml +++ b/deploy/memory-pilot/proxy_config.yaml @@ -10,6 +10,7 @@ litellm_settings: - hooks.forward_credential drop_params: true turn_off_message_logging: true + sse_keepalive_ping_interval_seconds: 15 general_settings: master_key: os.environ/LITELLM_MASTER_KEY store_model_in_db: true diff --git a/deploy/memory-pilot/start.sh b/deploy/memory-pilot/start.sh index bfce87a3c35..9eac610730b 100755 --- a/deploy/memory-pilot/start.sh +++ b/deploy/memory-pilot/start.sh @@ -6,4 +6,4 @@ export PRISMA_CLI_PATH="$PRISMA_BINARY_CACHE_DIR/node_modules/.bin/prisma" prisma migrate deploy --schema litellm-proxy-extras/litellm_proxy_extras/schema.prisma export WORKER_CONFIG="$PWD/deploy/memory-pilot/proxy_config.yaml" export PYTHONPATH="$PWD/deploy/memory-pilot${PYTHONPATH:+:$PYTHONPATH}" -exec uvicorn pilot:create_app --factory --host 0.0.0.0 --port "${PORT:-4000}" +exec uvicorn pilot:create_app --factory --host 0.0.0.0 --port "${PORT:-4000}" --timeout-graceful-shutdown 290 diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index a794e7b2b9e..fea44ba2f19 100644 --- a/litellm/proxy/memory/gateway.py +++ b/litellm/proxy/memory/gateway.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import ( inject_server_tools, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.sse_keepalive import wrap_passthrough_sse_bytes_with_keepalive_pings from litellm.proxy.memory.continuation import MemoryContinuation, MemoryContinuations, prefix_hashes, transcript_items from litellm.proxy.memory.knowledge import ( MEMORY_FUNCTIONS, @@ -156,7 +157,15 @@ class GatewayMemoryLoop: start: Final = await call.started status: Final = start.status if status >= 400: - raise HTTPException(status_code=status, detail="The authenticated gateway model call failed") + raise HTTPException( + status_code=status, + detail="The authenticated gateway model call failed", + headers={ # mutable-ok: FastAPI's HTTPException accepts a native header dictionary. + name.decode("latin-1"): value.decode("latin-1") + for name, value in start.headers + if name.lower() == b"retry-after" + }, + ) self.headers = MappingProxyType( { name.decode("latin-1"): value.decode("latin-1") @@ -338,7 +347,7 @@ async def process_gateway_memory( ) if data.get("background") is True or data.get("n", 1) != 1: raise HTTPException(status_code=400, detail="Gateway memory requires a foreground request with one completion") - from litellm.proxy.proxy_server import app + from litellm.proxy.proxy_server import app, llm_router loop: Final = GatewayMemoryLoop(app, request, data, route, store) iterator: Final = loop.run() @@ -364,12 +373,21 @@ async def process_gateway_memory( from litellm.proxy.common_request_processing import ( _UpstreamClosingStreamingResponse, # pyright: ignore[reportPrivateUsage] # Reuse cleanup when a client disconnects before consuming the prefetched stream. + ttft_keepalive_interval, ) return _UpstreamClosingStreamingResponse( - stream(), + wrap_passthrough_sse_bytes_with_keepalive_pings( + stream(), + ping_interval_seconds=ttft_keepalive_interval(data, llm_router), + upstream_headers=MappingProxyType({"content-type": "text/event-stream"}), + ), media_type="text/event-stream", - headers=loop.response_headers(), + headers={ # mutable-ok: Native ASGI response headers. + **loop.response_headers(), + "cache-control": "no-cache", + "x-accel-buffering": "no", + }, upstream_generator=iterator, ) diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py index 2ad2554a25b..c821521edd5 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -662,6 +662,116 @@ async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client assert "8347" in json.dumps(observed[1][2]["messages"]) and "8347" not in json.dumps(original) +@pytest.mark.asyncio +@pytest.mark.parametrize("between_rounds", [False, True]) +@pytest.mark.parametrize("disconnect", [False, True]) +async def test_silent_memory_rounds_keep_the_client_alive_and_cancel_upstream( + prisma_edge: MagicMock, between_rounds: bool, disconnect: bool +) -> None: + import asyncio + from unittest.mock import patch + + from starlette.responses import StreamingResponse + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.memory.gateway import process_gateway_memory + + waiting = asyncio.Event() + cancelled = asyncio.Event() + release = asyncio.Event() + calls = [] + + async def provider(scope: Scope, receive: Receive, send: Send) -> None: + calls.append(await receive()) + await send({"type": "http.response.start", "status": 200, "headers": []}) + frames = ( + { + "type": "message_start", + "message": { + "id": "msg_slow", + "role": "assistant", + "model": "test", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "search", + "name": "litellm_memory_search", + "input": {"query": "demo"}, + }, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 20}}, + {"type": "message_stop"}, + ) + if len(calls) == 1: + for frame in frames if between_rounds else frames[:1]: + await send( + { + "type": "http.response.body", + "body": b"data: " + json.dumps(frame).encode() + b"\n\n", + "more_body": True, + } + ) + if between_rounds: + await send({"type": "http.response.body", "body": b"", "more_body": False}) + return + waiting.set() + try: + await release.wait() + raise RuntimeError("private upstream failure") + finally: + cancelled.set() + + with ( + patch("litellm.sse_keepalive_ping_interval_seconds", 0.01), # test-quality-ok: Set the real operator configuration. + patch.multiple( # test-quality-ok: Replace the model HTTP boundary, preserving the real internal ASGI transport. + "litellm.proxy.proxy_server", app=provider, llm_router=None + ), + patch( # test-quality-ok: Inject authorized database edge; execute the real loop, SSE serialization and teardown. + "litellm.proxy.memory.gateway.gateway_memory_store", new=AsyncMock(return_value=store(prisma_edge)) + ), + ): + response = await process_gateway_memory( + {"model": "test", "stream": True, "messages": []}, request(), UserAPIKeyAuth(), "anthropic_messages" + ) + assert isinstance(response, StreamingResponse) + public = response.body_iterator + assert b"message_start" in await anext(public) + next_chunk = asyncio.create_task(anext(public)) + await asyncio.wait_for(waiting.wait(), timeout=1) + assert await asyncio.wait_for(next_chunk, timeout=0.5) == b": ping\n\n" + if disconnect: + await public.aclose() + else: + release.set() + remaining = b"".join([chunk async for chunk in public]) + assert remaining.count(b"event: error") == 1 + assert b"private upstream failure" not in remaining and b"message_stop" not in remaining + assert cancelled.is_set() and len(calls) == (2 if between_rounds else 1) + prisma_edge.db.litellm_memorycontinuation.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_gateway_preserves_upstream_retry_delay_without_exposing_provider_details(prisma_edge: MagicMock) -> None: + async def provider(scope: Scope, receive: Receive, send: Send) -> None: + await JSONResponse({"error": "private provider detail"}, status_code=429, headers={"Retry-After": "17"})( + scope, receive, send + ) + + loop = GatewayMemoryLoop(provider, request(), {"messages": []}, "anthropic_messages", store(prisma_edge)) + with pytest.raises(HTTPException) as exc: + async for _ in loop.run(): + pass + assert exc.value.status_code == 429 and exc.value.headers == {"retry-after": "17"} + assert "private provider detail" not in str(exc.value.detail) + + @pytest.mark.asyncio @pytest.mark.parametrize("share_auth_cache", [False, True]) async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pubsub(