easyai-ai-gateway/apps/api/internal/clients/wire_response_test.go

package clients

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestDecodeHTTPResponseForProtocolCapturesOfficialErrorWire(t *testing.T) {
	response := &http.Response{
		StatusCode: http.StatusTooManyRequests,
		Status:     "429 Too Many Requests",
		Header: http.Header{
			"Content-Type":               {"application/json"},
			"X-Request-Id":               {"req_1"},
			"X-Ratelimit-Reset-Requests": {"1s"},
			"Set-Cookie":                 {"secret=must-not-leak"},
			"Connection":                 {"keep-alive"},
		},
		Body: io.NopCloser(strings.NewReader(`{"error":{"message":"slow down","future_field":true}}`)),
	}
	_, wire, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
	if err == nil {
		t.Fatal("expected upstream error")
	}
	if wire == nil || wire.StatusCode != http.StatusTooManyRequests || wire.Protocol != ProtocolOpenAIResponses {
		t.Fatalf("unexpected wire response: %+v", wire)
	}
	if wire.Headers["X-Request-Id"][0] != "req_1" || wire.Headers["X-Ratelimit-Reset-Requests"][0] != "1s" {
		t.Fatalf("official response headers were lost: %+v", wire.Headers)
	}
	if _, ok := wire.Headers["Set-Cookie"]; ok {
		t.Fatalf("sensitive header leaked: %+v", wire.Headers)
	}
	if _, ok := wire.Headers["Connection"]; ok {
		t.Fatalf("hop-by-hop header leaked: %+v", wire.Headers)
	}
	if ErrorWireResponse(err) != wire {
		t.Fatal("client error did not retain its wire response")
	}
}

func TestDecodeHTTPResponseForProtocolAllowsLargeImagePayload(t *testing.T) {
	payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
	response := &http.Response{
		StatusCode: http.StatusOK,
		Status:     "200 OK",
		Header:     http.Header{"Content-Type": {"application/json"}},
		Body:       io.NopCloser(strings.NewReader(payload)),
	}

	result, wire, err := decodeHTTPResponseForProtocol(response, ProtocolGeminiGenerateContent)
	if err != nil {
		t.Fatalf("large Gemini image response failed: %v", err)
	}
	if got, _ := result["data"].(string); len(got) != int(defaultMaxJSONResponseBytes) {
		t.Fatalf("large Gemini image payload length = %d", len(got))
	}
	if wire == nil || len(wire.RawJSON) != len(payload) {
		t.Fatalf("large Gemini wire payload was truncated: %+v", wire)
	}
}

func TestDecodeHTTPResponseForProtocolRejectsOversizedDefaultPayload(t *testing.T) {
	payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
	response := &http.Response{
		StatusCode: http.StatusOK,
		Status:     "200 OK",
		Header:     http.Header{"Content-Type": {"application/json"}},
		Body:       io.NopCloser(strings.NewReader(payload)),
	}

	_, _, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
	if err == nil {
		t.Fatal("expected oversized default response to fail")
	}
	clientErr, ok := err.(*ClientError)
	if !ok || clientErr.Code != "response_too_large" {
		t.Fatalf("unexpected oversized response error: %T %v", err, err)
	}
}

func TestGeminiNativeStreamPreservesOfficialEvents(t *testing.T) {
	var requestedPath string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestedPath = r.URL.RequestURI()
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("X-Goog-Request-Id", "goog-stream-1")
		_, _ = io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}]}}],\"futureOfficialField\":true}\n\n")
	}))
	defer upstream.Close()

	var events []StreamDeltaEvent
	response, err := (GeminiClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{
		Kind: "chat.completions", ModelType: "text", Model: "gemini-test", Stream: true,
		Body: map[string]any{"prompt": "hello"},
		Candidate: store.RuntimeModelCandidate{
			Provider: "gemini", BaseURL: upstream.URL, ProviderModelName: "gemini-test",
			Credentials: map[string]any{"apiKey": "test-key"},
		},
		StreamDelta: func(event StreamDeltaEvent) error {
			events = append(events, event)
			return nil
		},
	})
	if err != nil {
		t.Fatalf("Gemini stream failed: %v", err)
	}
	if !strings.Contains(requestedPath, ":streamGenerateContent") || !strings.Contains(requestedPath, "alt=sse") {
		t.Fatalf("unexpected Gemini stream endpoint: %s", requestedPath)
	}
	if len(events) != 1 || events[0].WireProtocol != ProtocolGeminiGenerateContent || events[0].Event["futureOfficialField"] != true {
		raw, _ := json.Marshal(events)
		t.Fatalf("official Gemini event was not preserved: %s", raw)
	}
	if response.Wire == nil || response.Wire.Headers["X-Goog-Request-Id"][0] != "goog-stream-1" {
		t.Fatalf("Gemini stream wire metadata was lost: %+v", response.Wire)
	}
}