From e604606409f887761e7b74f6dbf6fe251a306f68 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Tue, 18 Aug 2026 10:43:29 -0700 Subject: [PATCH] gateway_proxy: shorten upstream read timeout and surface stream interruptions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lower the upstream `read` timeout from 600s to 120s. On the relayed path the pings the proxy sees are AIGW's keep-alive frames (~10s) covering only healthy streams, and AIGW aborts a genuine stall at ~60s and closes cleanly. The proxy read timeout then only needs to catch a dark upstream (pod crash / network partition, where neither tokens nor pings arrive) — 120s catches that in ~2 min while staying well above the 10s keep-alive so slow-but-healthy streams never trip it. connect/write/pool unchanged. When the upstream drops mid-stream after the head is sent, emit a terminal Anthropic `event: error` frame so Claude Code retries cleanly instead of rendering a silent truncation as "incomplete". Gate this on the body being uncompressed text/event-stream — this path also relays non-streaming JSON and possibly-gzipped bodies byte-for-byte, and injecting plaintext SSE framing into either would corrupt it (strictly worse than a clean truncation). Lead with a double CRLF to force an SSE event boundary in case the drop landed mid-line. Note: B1/B2 are necessary but not sufficient — a keep-alive-fed stall is still bounded only by AIGW's ~600s webClientRequestTimeoutMs until #2390858 lands. Co-authored-by: Isaac --- src/ucode/gateway_proxy.py | 37 ++++++++++++++++++++++++++++++------- 1 file changed, 30 insertions(+), 7 deletions(-) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 55a9d8e..e794f59 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -52,10 +52,14 @@ # plus the swap header (replaced with a freshly-minted value per request). _STRIP_ON_FORWARD = _HOP_BY_HOP | {_SWAP_HEADER.lower()} _STREAM_CHUNK = 8192 -# Per-operation upstream timeouts. `read` is generous because model turns stream -# over a single response and Anthropic emits SSE pings, so inter-chunk gaps stay -# small; `connect`/`pool` fail fast when the gateway is unreachable. -_UPSTREAM_TIMEOUT = httpx.Timeout(connect=10.0, read=600.0, write=600.0, pool=10.0) +# Per-operation upstream timeouts. On the relayed path the pings the proxy sees +# are AIGW's keep-alive frames (~every 10s), which cover only healthy streams; +# AIGW itself aborts a genuine stall at ~60s and closes cleanly, so `read` only +# needs to catch a dark upstream (pod crash / network partition, where neither +# tokens nor pings arrive) — 120s catches that in ~2 min while staying well above +# the 10s keep-alive so slow-but-healthy streams never trip it. `connect`/`pool` +# fail fast when the gateway is unreachable. +_UPSTREAM_TIMEOUT = httpx.Timeout(connect=10.0, read=120.0, write=600.0, pool=10.0) # Refresh once the token has less than this many seconds of life left. Databricks # access tokens live ~1h; a 10-min buffer leaves ample headroom for a retry. _REFRESH_BUFFER_S = 600 @@ -243,9 +247,28 @@ def _relay_response(self, resp: httpx.Response) -> None: # quietly rather than crashing the handler thread. return except httpx.HTTPError: - # Upstream dropped mid-stream. Headers (and status) may already be - # sent, so we can't reliably signal a fresh error — stop and let the - # client see a truncated stream rather than corrupt the framing. + # Upstream dropped mid-stream after the head was already sent. For an + # *uncompressed SSE* body we can emit a terminal Anthropic error frame + # so the client retries cleanly instead of rendering a silent + # truncation as "incomplete". But this path also relays non-streaming + # JSON and (since Accept-Encoding is forwarded) possibly-gzipped + # bodies byte-for-byte — injecting plaintext SSE framing into either + # would corrupt it, which is strictly worse than a clean truncation. + # So gate on content-type + content-encoding, and lead with a double + # CRLF to force an SSE event boundary in case the drop landed + # mid-line. Also guard for the client already being gone. + ctype = resp.headers.get("content-type", "") + cenc = resp.headers.get("content-encoding", "") + if "text/event-stream" in ctype and cenc in ("", "identity"): + try: + self.wfile.write( + b"\r\n\r\nevent: error\r\n" + b'data: {"type":"error","error":{"type":"overloaded_error",' + b'"message":"gateway proxy: upstream stream interrupted"}}\r\n\r\n' + ) + self.wfile.flush() + except OSError: + pass return # Forward every method: this is a transparent pass-through, so routing any