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 | import ( | 3 | import ( |
| 4 | "bytes" | 4 | "bytes" |
| 5 | "compress/gzip" | 5 | "compress/gzip" |
| 6 | "context" | ||
| 7 | "errors" | ||
| 6 | "fmt" | 8 | "fmt" |
| 7 | "io" | 9 | "io" |
| 10 | "net" | ||
| 8 | "net/http" | 11 | "net/http" |
| 9 | "net/http/httputil" | 12 | "net/http/httputil" |
| 10 | "net/url" | 13 | "net/url" |
| @@ -102,7 +105,15 @@ func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputi · 91 unmodified lines | |||
| 102 | // Modify incoming responses to handle gzip without Content-Encoding | 105 | // Modify incoming responses to handle gzip without Content-Encoding |
| 103 | // This addresses the same issue as inline handler gzip handling, but at the proxy level | 106 | // This addresses the same issue as inline handler gzip handling, but at the proxy level |
| 104 | proxy.ModifyResponse = func(resp *http.Response) error { | 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 | if resp.StatusCode < 200 || resp.StatusCode >= 300 { | 117 | if resp.StatusCode < 200 || resp.StatusCode >= 300 { |
| 107 | return nil | 118 | return nil |
| 108 | } | 119 | } |
| @@ -186,9 +197,29 @@ func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputi · 77 unmodified lines | |||
| 186 | return nil | 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 | proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) { | 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 | rw.Header().Set("Content-Type", "application/json") | 223 | rw.Header().Set("Content-Type", "application/json") |
| 193 | rw.WriteHeader(http.StatusBadGateway) | 224 | rw.WriteHeader(http.StatusBadGateway) |
| 194 | _, _ = rw.Write([]byte(`{"error":"amp_upstream_proxy_error","message":"Failed to reach Amp upstream"}`)) | 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 | // Write intercepts response writes and buffers them for model name replacement | 73 | // Write intercepts response writes and buffers them for model name replacement |
| 33 | func (rw *ResponseRewriter) Write(data []byte) (int, error) { | 74 | func (rw *ResponseRewriter) Write(data []byte) (int, error) { |
| 34 | // Detect streaming on first write | 75 | // Detect streaming on first write (header-based) |
| 35 | if rw.body.Len() == 0 && !rw.isStreaming { | 76 | if !rw.isStreaming && rw.body.Len() == 0 { |
| 36 | contentType := rw.Header().Get("Content-Type") | 77 | contentType := rw.Header().Get("Content-Type") |
| 37 | rw.isStreaming = strings.Contains(contentType, "text/event-stream") || | 78 | rw.isStreaming = strings.Contains(contentType, "text/event-stream") || |
| 38 | strings.Contains(contentType, "stream") | 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 | if rw.isStreaming { | 97 | if rw.isStreaming { |
| 42 | n, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(data)) | 98 | n, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(data)) |
| 43 | if err == nil { | 99 | if err == nil { |