fix(memory): preserve keepalives and retry delays across tool rounds

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

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

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

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,
     )
 

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(