fix(gateway): harden AI provider HTTP error handling and context propagation

gateway/internal/ai/error_handling_test.go

@@ -0,0 +1,116 @@
+package ai
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+// TestOpenRouterErrorBodyLimit verifies that the OpenRouter provider caps error
+// response body reads at maxErrorResponseBytes to prevent OOM from a malicious
+// upstream.
+func TestOpenRouterErrorBodyLimit(t *testing.T) {
+	// Serve a body larger than the limit (8KB > 4KB limit)
+	oversizedBody := strings.Repeat("X", 8192)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(oversizedBody))
+	}))
+	defer srv.Close()
+
+	p := &OpenRouterProvider{
+		apiKey: "test-key",
+		model:  "test-model",
+		url:    srv.URL,
+	}
+
+	_, err := p.Generate(context.Background(), "test prompt")
+	if err == nil {
+		t.Fatal("expected error from 500 response, got nil")
+	}
+
+	// The error message should NOT contain the full 8KB body
+	errMsg := err.Error()
+	if len(errMsg) > int(maxErrorResponseBytes)+200 {
+		t.Errorf("error message too long (%d bytes), expected capped at ~%d", len(errMsg), maxErrorResponseBytes)
+	}
+}
+
+// TestOllamaErrorBodyLimit verifies that the Ollama provider caps error
+// response body reads at maxOllamaErrorResponseBytes.
+func TestOllamaErrorBodyLimit(t *testing.T) {
+	oversizedBody := strings.Repeat("Y", 8192)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(oversizedBody))
+	}))
+	defer srv.Close()
+
+	p := &OllamaProvider{
+		url:   srv.URL,
+		model: "test-model",
+	}
+
+	_, err := p.Generate(context.Background(), "test prompt")
+	if err == nil {
+		t.Fatal("expected error from 500 response, got nil")
+	}
+
+	errMsg := err.Error()
+	if len(errMsg) > int(maxOllamaErrorResponseBytes)+200 {
+		t.Errorf("error message too long (%d bytes), expected capped at ~%d", len(errMsg), maxOllamaErrorResponseBytes)
+	}
+}
+
+// TestOpenRouterContextCanceled verifies that context.Canceled is returned
+// (not wrapped in a connection error) when the context is canceled.
+func TestOpenRouterContextCanceled(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Never respond — let the context cancel
+		<-r.Context().Done()
+	}))
+	defer srv.Close()
+
+	p := &OpenRouterProvider{
+		apiKey: "test-key",
+		model:  "test-model",
+		url:    srv.URL,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // Cancel immediately
+
+	_, err := p.Generate(ctx, "test prompt")
+	if err == nil {
+		t.Fatal("expected error from canceled context, got nil")
+	}
+	if err != context.Canceled {
+		t.Errorf("expected context.Canceled, got: %v", err)
+	}
+}
+
+// TestOllamaContextCanceled verifies the same for the Ollama provider.
+func TestOllamaContextCanceled(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		<-r.Context().Done()
+	}))
+	defer srv.Close()
+
+	p := &OllamaProvider{
+		url:   srv.URL,
+		model: "test-model",
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err := p.Generate(ctx, "test prompt")
+	if err == nil {
+		t.Fatal("expected error from canceled context, got nil")
+	}
+	if err != context.Canceled {
+		t.Errorf("expected context.Canceled, got: %v", err)
+	}
+}

gateway/internal/ai/ollama.go

@@ -11,6 +11,11 @@ import (
 	"os"
 )
 
+// maxErrorResponseBytes caps the amount of upstream error response body we
+// read on non-200 responses. Prevents OOM if a misbehaving Ollama instance
+// streams an unbounded error payload.
+const maxOllamaErrorResponseBytes int64 = 4096
+
 // OllamaProvider implements the Provider interface for Ollama API
 type OllamaProvider struct {
 	url   string
@@ -53,15 +58,16 @@ func (p *OllamaProvider) Generate(ctx context.Context, text string) (string, err
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded {
-			return "", context.DeadlineExceeded
+		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) ||
+			ctx.Err() == context.DeadlineExceeded || ctx.Err() == context.Canceled {
+			return "", ctx.Err()
 		}
 		return "", fmt.Errorf("failed to connect to Ollama: %w", err)
 	}
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		body, _ := io.ReadAll(resp.Body)
+		body, _ := io.ReadAll(io.LimitReader(resp.Body, maxOllamaErrorResponseBytes))
 		return "", fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, string(body))
 	}
 

gateway/internal/ai/openrouter.go

@@ -12,6 +12,12 @@ import (
 	"os"
 )
 
+// maxErrorResponseBytes caps the amount of upstream error response body we are
+// willing to read. Without this limit, a misbehaving or malicious upstream
+// could stream an unbounded response on an error status, exhausting gateway
+// memory.
+const maxErrorResponseBytes int64 = 4096
+
 // OpenRouterProvider implements the Provider interface for OpenRouter API
 type OpenRouterProvider struct {
 	apiKey string
@@ -59,16 +65,17 @@ func (p *OpenRouterProvider) Generate(ctx context.Context, text string) (string,
 
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded {
-			return "", context.DeadlineExceeded
+		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) ||
+			ctx.Err() == context.DeadlineExceeded || ctx.Err() == context.Canceled {
+			return "", ctx.Err()
 		}
 		return "", err
 	}
 	defer resp.Body.Close()
 
 	// Check status code before decoding
 	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		body, _ := io.ReadAll(resp.Body)
+		body, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorResponseBytes))
 		return "", fmt.Errorf("openrouter returned status %d: %s", resp.StatusCode, string(body))
 	}
 

gateway/main.go

@@ -523,7 +523,9 @@ func handleSummarize(c *gin.Context) {
 	// 3. Call AI Service
 	summary, err := aiProvider.Generate(c.Request.Context(), req.Text)
 	if err != nil {
-		if errors.Is(err, context.DeadlineExceeded) || errors.Is(c.Request.Context().Err(), context.DeadlineExceeded) {
+		ctxErr := c.Request.Context().Err()
+		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) ||
+			ctxErr == context.DeadlineExceeded || ctxErr == context.Canceled {
 			respondError(c, 504, "upstream_timeout", err)
 		} else {
 			respondError(c, 502, "upstream_unavailable", err)