Showing the changes to internal/api/modules (2 of 74 files). Show all
Unified · Split
| @@ -3,8 +3,11 @@ package amp · 2 unmodified lines | ||
| 3 | 3 | import ( |
| 4 | 4 | "bytes" |
| 5 | 5 | "compress/gzip" |
| 6 | + "context" | |
| 7 | + "errors" | |
| 6 | 8 | "fmt" |
| 7 | 9 | "io" |
| 10 | + "net" | |
| 8 | 11 | "net/http" |
| 9 | 12 | "net/http/httputil" |
| 10 | 13 | "net/url" |
| @@ -102,7 +105,15 @@ func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputi · 91 unmodified lines | ||
| 102 | 105 | // Modify incoming responses to handle gzip without Content-Encoding |
| 103 | 106 | // This addresses the same issue as inline handler gzip handling, but at the proxy level |
| 104 | 107 | proxy.ModifyResponse = func(resp *http.Response) error { |
| 105 | − // Only process successful responses | |
| 108 | + // Log upstream error responses for diagnostics (502, 503, etc.) | |
| 109 | + // These are NOT proxy connection errors - the upstream responded with an error status | |
| 110 | + if resp.StatusCode >= 500 { | |
| 111 | + log.Errorf("amp upstream responded with error [%d] for %s %s", resp.StatusCode, resp.Request.Method, resp.Request.URL.Path) | |
| 112 | + } else if resp.StatusCode >= 400 { | |
| 113 | + log.Warnf("amp upstream responded with client error [%d] for %s %s", resp.StatusCode, resp.Request.Method, resp.Request.URL.Path) | |
| 114 | + } | |
| 115 | + | |
| 116 | + // Only process successful responses for gzip decompression | |
| 106 | 117 | if resp.StatusCode < 200 || resp.StatusCode >= 300 { |
| 107 | 118 | return nil |
| 108 | 119 | } |
| @@ -186,9 +197,29 @@ func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputi · 77 unmodified lines | ||
| 186 | 197 | return nil |
| 187 | 198 | } |
| 188 | 199 | |
| 189 | − // Error handler for proxy failures | |
| 200 | + // Error handler for proxy failures with detailed error classification for diagnostics | |
| 190 | 201 | proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) { |
| 191 | − log.Errorf("amp upstream proxy error for %s %s: %v", req.Method, req.URL.Path, err) | |
| 202 | + // Classify the error type for better diagnostics | |
| 203 | + var errType string | |
| 204 | + if errors.Is(err, context.DeadlineExceeded) { | |
| 205 | + errType = "timeout" | |
| 206 | + } else if errors.Is(err, context.Canceled) { | |
| 207 | + errType = "canceled" | |
| 208 | + } else if netErr, ok := err.(net.Error); ok && netErr.Timeout() { | |
| 209 | + errType = "dial_timeout" | |
| 210 | + } else if _, ok := err.(net.Error); ok { | |
| 211 | + errType = "network_error" | |
| 212 | + } else { | |
| 213 | + errType = "connection_error" | |
| 214 | + } | |
| 215 | + | |
| 216 | + // Don't log as error for context canceled - it's usually client closing connection | |
| 217 | + if errors.Is(err, context.Canceled) { | |
| 218 | + log.Debugf("amp upstream proxy [%s]: client canceled request for %s %s", errType, req.Method, req.URL.Path) | |
| 219 | + } else { | |
| 220 | + log.Errorf("amp upstream proxy error [%s] for %s %s: %v", errType, req.Method, req.URL.Path, err) | |
| 221 | + } | |
| 222 | + | |
| 192 | 223 | rw.Header().Set("Content-Type", "application/json") |
| 193 | 224 | rw.WriteHeader(http.StatusBadGateway) |
| 194 | 225 | _, _ = rw.Write([]byte(`{"error":"amp_upstream_proxy_error","message":"Failed to reach Amp upstream"}`)) |
| @@ -29,15 +29,71 @@ func NewResponseRewriter(w gin.ResponseWriter, originalModel string) *ResponseRe · 28 unmodified lines | ||
| 29 | 29 | } |
| 30 | 30 | } |
| 31 | 31 | |
| 32 | +const maxBufferedResponseBytes = 2 * 1024 * 1024 // 2MB safety cap | |
| 33 | + | |
| 34 | +func looksLikeSSEChunk(data []byte) bool { | |
| 35 | + // Fallback detection: some upstreams may omit/lie about Content-Type, causing SSE to be buffered. | |
| 36 | + // Heuristics are intentionally simple and cheap. | |
| 37 | + return bytes.Contains(data, []byte("data:")) || | |
| 38 | + bytes.Contains(data, []byte("event:")) || | |
| 39 | + bytes.Contains(data, []byte("message_start")) || | |
| 40 | + bytes.Contains(data, []byte("message_delta")) || | |
| 41 | + bytes.Contains(data, []byte("content_block_start")) || | |
| 42 | + bytes.Contains(data, []byte("content_block_delta")) || | |
| 43 | + bytes.Contains(data, []byte("content_block_stop")) || | |
| 44 | + bytes.Contains(data, []byte("\n\n")) | |
| 45 | +} | |
| 46 | + | |
| 47 | +func (rw *ResponseRewriter) enableStreaming(reason string) error { | |
| 48 | + if rw.isStreaming { | |
| 49 | + return nil | |
| 50 | + } | |
| 51 | + rw.isStreaming = true | |
| 52 | + | |
| 53 | + // Flush any previously buffered data to avoid reordering or data loss. | |
| 54 | + if rw.body != nil && rw.body.Len() > 0 { | |
| 55 | + buf := rw.body.Bytes() | |
| 56 | + // Copy before Reset() to keep bytes stable. | |
| 57 | + toFlush := make([]byte, len(buf)) | |
| 58 | + copy(toFlush, buf) | |
| 59 | + rw.body.Reset() | |
| 60 | + | |
| 61 | + if _, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(toFlush)); err != nil { | |
| 62 | + return err | |
| 63 | + } | |
| 64 | + if flusher, ok := rw.ResponseWriter.(http.Flusher); ok { | |
| 65 | + flusher.Flush() | |
| 66 | + } | |
| 67 | + } | |
| 68 | + | |
| 69 | + log.Debugf("amp response rewriter: switched to streaming (%s)", reason) | |
| 70 | + return nil | |
| 71 | +} | |
| 72 | + | |
| 32 | 73 | // Write intercepts response writes and buffers them for model name replacement |
| 33 | 74 | func (rw *ResponseRewriter) Write(data []byte) (int, error) { |
| 34 | − // Detect streaming on first write | |
| 35 | − if rw.body.Len() == 0 && !rw.isStreaming { | |
| 75 | + // Detect streaming on first write (header-based) | |
| 76 | + if !rw.isStreaming && rw.body.Len() == 0 { | |
| 36 | 77 | contentType := rw.Header().Get("Content-Type") |
| 37 | 78 | rw.isStreaming = strings.Contains(contentType, "text/event-stream") || |
| 38 | 79 | strings.Contains(contentType, "stream") |
| 39 | 80 | } |
| 40 | 81 | |
| 82 | + if !rw.isStreaming { | |
| 83 | + // Content-based fallback: detect SSE-like chunks even if Content-Type is missing/wrong. | |
| 84 | + if looksLikeSSEChunk(data) { | |
| 85 | + if err := rw.enableStreaming("sse heuristic"); err != nil { | |
| 86 | + return 0, err | |
| 87 | + } | |
| 88 | + } else if rw.body.Len()+len(data) > maxBufferedResponseBytes { | |
| 89 | + // Safety cap: avoid unbounded buffering on large responses. | |
| 90 | + log.Warnf("amp response rewriter: buffer exceeded %d bytes, switching to streaming", maxBufferedResponseBytes) | |
| 91 | + if err := rw.enableStreaming("buffer limit"); err != nil { | |
| 92 | + return 0, err | |
| 93 | + } | |
| 94 | + } | |
| 95 | + } | |
| 96 | + | |
| 41 | 97 | if rw.isStreaming { |
| 42 | 98 | n, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(data)) |
| 43 | 99 | if err == nil { |