gpt-load/internal/gateway/dialects_integration_test.go

package gateway

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gin-gonic/gin"

	"gpt-load/internal/channel"
	"gpt-load/internal/dialect"
	"gpt-load/internal/execution"
	"gpt-load/internal/health"
	"gpt-load/internal/platform/config"
	"gpt-load/internal/platform/contentcoding"
	"gpt-load/internal/protocol"
	"gpt-load/internal/state"
	"gpt-load/internal/testutil/encryptiontest"
	"gpt-load/internal/testutil/fakeupstream"
)

type dialectGatewayGroup struct {
	id          uint
	name        string
	upstreamURL string
	channelID   channel.ID
	params      json.RawMessage
	apiKeys     []string
	models      []state.ModelConfig
	settings    config.Settings
	headerRules state.HeaderRules
	firstByte   time.Duration
}

func testUpstreamBaseURL(raw string, selected protocol.Protocol) string {
	parsed, err := url.Parse(raw)
	if err != nil {
		return raw
	}
	prefix := "/v1"
	if selected == protocol.Anthropic {
		return parsed.String()
	}
	if selected == protocol.Gemini {
		prefix = "/v1beta"
	}
	path := strings.TrimRight(parsed.Path, "/")
	if path != prefix && !strings.HasSuffix(path, prefix) {
		parsed.Path = path + prefix
		parsed.RawPath = ""
	}
	return parsed.String()
}

func newDialectGatewayEngine(
	t *testing.T,
	selectedProtocol protocol.Protocol,
	model string,
	dialects dialect.Set,
	groups ...dialectGatewayGroup,
) (*gin.Engine, *state.CredentialRegistry) {
	return newDialectGatewayEngineWithForwarder(
		t,
		selectedProtocol,
		model,
		dialects,
		newTestExecutionForwarder(t),
		groups...,
	)
}

func newDialectGatewayEngineWithForwarder(
	t *testing.T,
	selectedProtocol protocol.Protocol,
	model string,
	dialects dialect.Set,
	forwarder AttemptForwarder,
	groups ...dialectGatewayGroup,
) (*gin.Engine, *state.CredentialRegistry) {
	t.Helper()
	return newDialectGatewayEngineWithSystemSettings(t, selectedProtocol, model, dialects, forwarder, nil, groups...)
}

func newDialectGatewayEngineWithSystemSettings(
	t *testing.T,
	selectedProtocol protocol.Protocol,
	model string,
	dialects dialect.Set,
	forwarder AttemptForwarder,
	systemSettings config.Settings,
	groups ...dialectGatewayGroup,
) (*gin.Engine, *state.CredentialRegistry) {
	t.Helper()
	gin.SetMode(gin.TestMode)
	keyService := encryptiontest.Service(t, "dialect-gateway-test-master-key")

	configs := make([]state.GroupConfig, 0, len(groups))
	entries := make([]state.CredentialEntry, 0)
	credentialConfigs := make([]state.CredentialConfig, 0)
	credentialID := uint(1)
	for _, group := range groups {
		models := group.models
		if len(models) == 0 {
			models = []state.ModelConfig{{ID: model}}
		}
		channelID, params := group.channelID, group.params
		if channelID == "" {
			baseURL := testUpstreamBaseURL(group.upstreamURL, selectedProtocol)
			channelID, params = testChannelConfig(t, selectedProtocol, baseURL)
		}
		configs = append(configs, state.GroupConfig{ConnectionType: "api_key", ID: group.id, Name: group.name, ChannelID: channelID, Params: params,
			Models: models, Settings: group.settings, Enabled: true,
		})
		for _, apiKey := range group.apiKeys {
			entries = append(entries, testCredentialEntry(t, keyService, credentialID, group.id, apiKey))
			credentialConfigs = append(credentialConfigs, testCredentialConfig(credentialID, group.id))
			credentialID++
		}
	}

	manager := state.NewManager()
	snapshot, err := manager.Publish(state.CompileInput{
		SystemSettings:  systemSettings,
		ChannelRegistry: channel.NewRegistry(), Groups: configs,
		Credentials: credentialConfigs,
		AccessKeys: []state.AccessKeyConfig{{
			ID: 1, Name: "client", KeyHash: keyService.Hash("gl-client"),
			Status: state.AccessKeyStatusActive,
		}},
	})
	if err != nil {
		t.Fatalf("Publish() error = %v", err)
	}
	for _, group := range groups {
		view := snapshot.Groups[group.id]
		view.HeaderRules = group.headerRules
		if group.firstByte > 0 {
			view.Timeouts.FirstByte = group.firstByte
		}
		snapshot.Groups[group.id] = view
	}

	registry := state.NewCredentialRegistry()
	if err := registry.ReplaceCredentials(entries); err != nil {
		t.Fatalf("ReplaceCredentials() error = %v", err)
	}
	handler := NewHandler(
		manager,
		registry,
		keyService,
		forwarder,
		dialects,
		health.NewStatsStore(),
		health.NewMutationCoordinator(),
		nil,
		nil,
		nil,
	)

	engine := gin.New()
	bindGatewayRoutesForTest(t, engine, handler)
	return engine, registry
}

func TestGatewayDecodesSupportedClientContentCodings(t *testing.T) {
	routes := []struct {
		name      string
		protocol  protocol.Protocol
		dialects  dialect.Set
		path      string
		plaintext []byte
	}{
		{
			name: "OpenAI completions", protocol: protocol.OpenAICompletions,
			dialects: dialect.NewSet(dialect.NewOpenAI()),
			path:     "/v1/chat/completions", plaintext: []byte(`{"model":"public-model","stream":false}`),
		},
		{
			name: "Anthropic messages", protocol: protocol.Anthropic,
			dialects: dialect.NewSet(dialect.NewAnthropic()),
			path:     "/v1/messages", plaintext: []byte(`{"model":"public-model","max_tokens":1,"messages":[]}`),
		},
		{
			name: "Gemini generation", protocol: protocol.Gemini,
			dialects: dialect.NewSet(dialect.NewGemini()),
			path:     "/v1beta/models/public-model:generateContent", plaintext: []byte(`{"contents":[{"role":"user","parts":[{"text":"ping"}]}]}`),
		},
		{
			name: "Responses create", protocol: protocol.OpenAIResponses,
			dialects: dialect.NewSet(dialect.NewOpenAIResponses()),
			path:     "/v1/responses", plaintext: []byte(`{"model":"public-model","input":"ping","store":false}`),
		},
	}
	encodings := []contentcoding.Encoding{
		contentcoding.Identity,
		contentcoding.Gzip,
		contentcoding.Brotli,
		contentcoding.Deflate,
		contentcoding.Zstd,
	}

	for _, route := range routes {
		t.Run(route.name, func(t *testing.T) {
			for _, encoding := range encodings {
				t.Run(string(encoding), func(t *testing.T) {
					forwarder := &scriptedForwarder{results: []UpstreamResult{{
						StatusCode:     http.StatusOK,
						Header:         http.Header{"Content-Type": {"application/json"}},
						Body:           []byte(`{"ok":true}`),
						RequestWritten: true,
						DispatchState:  execution.DispatchMaybeSent,
					}}}
					engine, _ := newDialectGatewayEngineWithForwarder(
						t,
						route.protocol,
						"public-model",
						route.dialects,
						forwarder,
						dialectGatewayGroup{
							id: 1, name: route.name, upstreamURL: "https://provider.example/v1",
							apiKeys: []string{"sk-provider"},
						},
					)
					request := httptest.NewRequest(
						http.MethodPost,
						route.path,
						bytes.NewReader(encodeContentCodingForGatewayTest(t, encoding, route.plaintext)),
					)
					request.Header.Set("Authorization", "Bearer gl-client")
					if encoding != contentcoding.Identity {
						request.Header.Set("Content-Encoding", string(encoding))
					}
					recorder := httptest.NewRecorder()
					engine.ServeHTTP(recorder, request)

					if recorder.Code != http.StatusOK {
						t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
					}
					if len(forwarder.inputs) != 1 || !bytes.Equal(forwarder.inputs[0].Request.Body, route.plaintext) {
						t.Fatalf("forwarded inputs = %#v, want plaintext %s", forwarder.inputs, route.plaintext)
					}
					if got := forwarder.inputs[0].Request.Header.Get("Content-Encoding"); got != "" {
						t.Fatalf("forwarded Content-Encoding = %q, want absent", got)
					}
				})
			}
		})
	}
}

func TestResponsesNamespaceRoutesOrdinaryMethodsAndRejectsDangerousMethodsLocally(t *testing.T) {
	tests := []struct {
		name            string
		method          string
		target          string
		body            string
		status          int
		wantBody        string
		wantContentType string
	}{
		{
			name:   "create",
			method: http.MethodPost,
			target: "/v1/responses?trace=true",
			body:   `{"model":"public-model","input":"ping"}`,
			status: http.StatusOK,
		},
		{
			name:   "retrieve",
			method: http.MethodGet,
			target: "/v1/responses/resp_123?stream=false&include=output",
			status: http.StatusOK,
		},
		{
			name:   "delete",
			method: http.MethodDelete,
			target: "/v1/responses/resp_123",
			status: http.StatusOK,
		},
		{
			name:   "replace extension",
			method: http.MethodPut,
			target: "/v1/responses/vendor-extension",
			body:   `{"enabled":true}`,
			status: http.StatusOK,
		},
		{
			name:   "head",
			method: http.MethodHead,
			target: "/v1/responses/resp_123",
			status: http.StatusOK,
		},
		{
			name:   "cancel",
			method: http.MethodPost,
			target: "/v1/responses/resp_123/cancel",
			status: http.StatusOK,
		},
		{
			name:   "input items",
			method: http.MethodGet,
			target: "/v1/responses/resp_123/input_items?limit=20",
			status: http.StatusOK,
		},
		{
			name:   "compact",
			method: http.MethodPost,
			target: "/v1/responses/compact",
			body:   `{"model":"public-model","input":"ping"}`,
			status: http.StatusOK,
		},
		{
			name:   "input tokens",
			method: http.MethodPost,
			target: "/v1/responses/input_tokens",
			body:   `{"model":"public-model","input":"ping"}`,
			status: http.StatusOK,
		},
		{
			name:   "unknown resource extension",
			method: http.MethodPatch,
			target: "/v1/responses/vendor-extension/nested",
			status: http.StatusOK,
		},
		{
			name:            "upstream resource error",
			method:          http.MethodGet,
			target:          "/v1/responses/resp_missing",
			status:          http.StatusNotFound,
			wantBody:        `{"error":{"message":"resource missing"}}`,
			wantContentType: "application/json",
		},
	}
	results := make([]UpstreamResult, 0, len(tests))
	for _, test := range tests {
		body := []byte(`{"id":"resp_123","object":"response"}`)
		if test.status == http.StatusNotFound {
			body = []byte(`{"error":{"message":"resource missing"}}`)
		}
		results = append(results, UpstreamResult{
			StatusCode:     test.status,
			Header:         http.Header{"Content-Type": {"application/json"}},
			Body:           body,
			RequestWritten: true,
			DispatchState:  execution.DispatchMaybeSent,
		})
	}
	forwarder := &scriptedForwarder{results: results}
	engine, _ := newDialectGatewayEngineWithForwarder(
		t,
		protocol.OpenAIResponses,
		"public-model",
		dialect.NewSet(dialect.NewOpenAIResponses()),
		forwarder,
		dialectGatewayGroup{
			id: 1, name: "responses", channelID: channel.OpenAI, params: json.RawMessage(`{}`),
			apiKeys: []string{"sk-responses"},
		},
	)
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			request := httptest.NewRequest(
				test.method,
				test.target,
				strings.NewReader(test.body),
			)
			request.Header.Set("Authorization", "Bearer gl-client")
			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, request)
			if recorder.Code != test.status {
				t.Fatalf(
					"response = %d %s, want %d",
					recorder.Code,
					recorder.Body.String(),
					test.status,
				)
			}
			if test.wantBody != "" && recorder.Body.String() != test.wantBody {
				t.Fatalf("response body = %q, want %q", recorder.Body.String(), test.wantBody)
			}
			if test.wantContentType != "" &&
				recorder.Header().Get("Content-Type") != test.wantContentType {
				t.Fatalf(
					"Content-Type = %q, want %q",
					recorder.Header().Get("Content-Type"),
					test.wantContentType,
				)
			}
		})
	}

	for _, test := range []struct {
		name       string
		method     string
		authorized bool
		wantStatus int
		wantCode   string
	}{
		{
			name:       "unauthenticated options keeps auth first",
			method:     http.MethodOptions,
			wantStatus: http.StatusUnauthorized,
			wantCode:   "invalid_access_key",
		},
		{
			name:       "authenticated options rejected locally",
			method:     http.MethodOptions,
			authorized: true,
			wantStatus: http.StatusNotFound,
			wantCode:   "protocol_endpoint_not_found",
		},
		{
			name:       "authenticated connect rejected locally",
			method:     http.MethodConnect,
			authorized: true,
			wantStatus: http.StatusNotFound,
			wantCode:   "protocol_endpoint_not_found",
		},
		{
			name:       "authenticated trace rejected locally",
			method:     http.MethodTrace,
			authorized: true,
			wantStatus: http.StatusNotFound,
			wantCode:   "protocol_endpoint_not_found",
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			request := httptest.NewRequest(
				test.method,
				"/v1/responses",
				nil,
			)
			if test.authorized {
				request.Header.Set("Authorization", "Bearer gl-client")
			}
			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, request)
			if recorder.Code != test.wantStatus ||
				!strings.Contains(
					recorder.Body.String(),
					`"code":"`+test.wantCode+`"`,
				) {
				t.Fatalf(
					"response = %d %s, want %d/%s",
					recorder.Code,
					recorder.Body.String(),
					test.wantStatus,
					test.wantCode,
				)
			}
		})
	}

	if len(forwarder.inputs) != len(tests) {
		t.Fatalf("forwarded requests = %d, want %d", len(forwarder.inputs), len(tests))
	}
	for index, input := range forwarder.inputs {
		requestURI := input.Request.Path
		if input.Request.RawQuery != "" {
			requestURI += "?" + input.Request.RawQuery
		}
		if input.Request.Method != tests[index].method || requestURI != tests[index].target ||
			input.APIKey != "sk-responses" {
			t.Fatalf("forwarded request %d = %s %s api_key=%q", index, input.Request.Method, requestURI, input.APIKey)
		}
	}
}

func TestResponsesNamespaceRejectsNormalizationEscapesBeforeForward(t *testing.T) {
	var upstreamCalls atomic.Int64
	upstream := httptest.NewServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		_ *http.Request,
	) {
		upstreamCalls.Add(1)
		writer.Header().Set("Location", "/v1/models")
		writer.WriteHeader(http.StatusTemporaryRedirect)
	}))
	defer upstream.Close()

	engine, _ := newDialectGatewayEngine(
		t,
		protocol.OpenAIResponses,
		"public-model",
		dialect.NewSet(dialect.NewOpenAIResponses()),
		dialectGatewayGroup{
			id: 1, name: "responses", upstreamURL: upstream.URL,
			apiKeys: []string{"sk-responses"},
		},
	)

	for _, target := range []string{
		"/v1/responses/../models",
		"/v1/responses/./resp_123",
		"/v1/responses/%2e%2e/models",
		"/v1/responses//%2e%2e/models",
	} {
		t.Run(target, func(t *testing.T) {
			request := httptest.NewRequest(http.MethodGet, target, nil)
			request.Header.Set("Authorization", "Bearer gl-client")
			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, request)
			if recorder.Code != http.StatusNotFound ||
				!strings.Contains(
					recorder.Body.String(),
					`"code":"protocol_endpoint_not_found"`,
				) {
				t.Fatalf(
					"response = %d %s, want local endpoint rejection",
					recorder.Code,
					recorder.Body.String(),
				)
			}
		})
	}
	if calls := upstreamCalls.Load(); calls != 0 {
		t.Fatalf("upstream calls = %d, want 0", calls)
	}
}

func TestResponsesNamespaceRejectsDecodedDuplicateStreamQueryBeforeForward(t *testing.T) {
	var upstreamCalls atomic.Int64
	upstream := httptest.NewServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		_ *http.Request,
	) {
		upstreamCalls.Add(1)
		writer.WriteHeader(http.StatusOK)
	}))
	defer upstream.Close()

	engine, _ := newDialectGatewayEngine(
		t,
		protocol.OpenAIResponses,
		"public-model",
		dialect.NewSet(dialect.NewOpenAIResponses()),
		dialectGatewayGroup{
			id: 1, name: "responses", upstreamURL: upstream.URL,
			apiKeys: []string{"sk-responses"},
		},
	)
	request := httptest.NewRequest(
		http.MethodGet,
		"/v1/responses/resp_123?stream=true&%73tream=false",
		nil,
	)
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusBadRequest ||
		!strings.Contains(
			recorder.Body.String(),
			`"code":"invalid_protocol_request"`,
		) ||
		upstreamCalls.Load() != 0 {
		t.Fatalf(
			"response/upstream calls = %d %s / %d",
			recorder.Code,
			recorder.Body.String(),
			upstreamCalls.Load(),
		)
	}
}

func TestGatewayRewritesEachAttemptFromOriginal(t *testing.T) {
	first := fakeupstream.New(fakeupstream.Step{Status: http.StatusUnauthorized, Fixture: "401.json"})
	defer first.Close()
	second := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"})
	defer second.Close()

	engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "public",
		dialect.NewSet(dialect.NewOpenAI()),
		dialectGatewayGroup{
			id: 1, name: "first", upstreamURL: first.URL, apiKeys: []string{"sk-first"},
			models: []state.ModelConfig{{ID: "provider-one", Alias: "public"}},
		},
		dialectGatewayGroup{
			id: 2, name: "second", upstreamURL: second.URL, apiKeys: []string{"sk-second"},
			models: []state.ModelConfig{{ID: "provider-two", Alias: "public"}},
		},
	)
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"public"}`))
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "2" {
		t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
	}
	requests := append(first.Requests(), second.Requests()...)
	if len(requests) != 2 {
		t.Fatalf("upstream requests = %d, want 2", len(requests))
	}
	wantModels := []string{"provider-one", "provider-two"}
	for index, received := range requests {
		var body struct {
			Model string `json:"model"`
		}
		if err := json.Unmarshal(received.Body, &body); err != nil {
			t.Fatalf("decode upstream request %d: %v", index+1, err)
		}
		if body.Model != wantModels[index] {
			t.Fatalf("upstream request %d model = %q, want %q; body=%s", index+1, body.Model, wantModels[index], received.Body)
		}
	}
}

func TestGatewayAppliesEachGroupsParameterOverridesFromOriginal(t *testing.T) {
	first := fakeupstream.New(fakeupstream.Step{Status: http.StatusUnauthorized, Fixture: "401.json"})
	defer first.Close()
	second := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"})
	defer second.Close()

	rules := func(marker string) config.Settings {
		return config.Settings{state.SettingParameterOverrides: []any{
			map[string]any{
				"match": map[string]any{
					"protocol": string(protocol.OpenAICompletions),
					"model":    "public*",
				},
				"remove": []any{"/nested/remove"},
				"set": map[string]any{
					"marker": marker,
					"nested": map[string]any{"group": marker},
				},
			},
		}}
	}
	engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "public",
		dialect.NewSet(dialect.NewOpenAI()),
		dialectGatewayGroup{
			id: 1, name: "first", upstreamURL: first.URL, apiKeys: []string{"sk-first"},
			models:   []state.ModelConfig{{ID: "provider-one", Alias: "public"}},
			settings: rules("first"),
		},
		dialectGatewayGroup{
			id: 2, name: "second", upstreamURL: second.URL, apiKeys: []string{"sk-second"},
			models:   []state.ModelConfig{{ID: "provider-two", Alias: "public"}},
			settings: rules("second"),
		},
	)
	request := httptest.NewRequest(
		http.MethodPost,
		"/v1/chat/completions",
		strings.NewReader(`{"model":"public","nested":{"remove":true,"original":true}}`),
	)
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK {
		t.Fatalf("response = %d body=%s", recorder.Code, recorder.Body.String())
	}
	requests := append(first.Requests(), second.Requests()...)
	if len(requests) != 2 {
		t.Fatalf("upstream requests = %d, want 2", len(requests))
	}
	for index, received := range requests {
		var body map[string]any
		if err := json.Unmarshal(received.Body, &body); err != nil {
			t.Fatal(err)
		}
		wantMarker := []string{"first", "second"}[index]
		if body["marker"] != wantMarker {
			t.Fatalf("request %d marker = %#v, want %q", index+1, body["marker"], wantMarker)
		}
		nested, ok := body["nested"].(map[string]any)
		if !ok || nested["remove"] != nil || nested["original"] != true || nested["group"] != wantMarker {
			t.Fatalf("request %d nested = %#v", index+1, body["nested"])
		}
	}
}

func TestHandlerHostFailureRetriesGenerationAcrossGroups(t *testing.T) {
	primary := fakeupstream.New(
		fakeupstream.Step{Status: http.StatusInternalServerError, Fixture: "openai/500.json"},
		fakeupstream.Step{Status: http.StatusInternalServerError, Fixture: "openai/500.json"},
		fakeupstream.Step{Status: http.StatusInternalServerError, Fixture: "openai/500.json"},
		fakeupstream.Step{Status: http.StatusInternalServerError, Fixture: "openai/500.json"},
	)
	defer primary.Close()
	backup := fakeupstream.New(
		fakeupstream.Step{Status: http.StatusOK, Fixture: "openai/success.json"},
		fakeupstream.Step{Status: http.StatusOK, Fixture: "openai/success.json"},
	)
	defer backup.Close()

	engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "gpt-4o",
		dialect.NewSet(dialect.NewOpenAI()),
		dialectGatewayGroup{id: 1, name: "primary", upstreamURL: primary.URL,
			apiKeys: []string{"sk-primary-one", "sk-primary-two"}},
		dialectGatewayGroup{id: 2, name: "backup", upstreamURL: backup.URL,
			apiKeys: []string{"sk-backup"}},
	)
	for range 2 {
		request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions",
			bytes.NewBufferString(`{"model":"gpt-4o"}`))
		request.Header.Set("Authorization", "Bearer gl-client")
		recorder := httptest.NewRecorder()
		engine.ServeHTTP(recorder, request)
		if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "2" {
			t.Fatalf("response = %d attempts=%s body=%s",
				recorder.Code, recorder.Header().Get(debugHeaderAttempts), recorder.Body.String())
		}
	}
	if got := len(primary.Requests()); got != 2 {
		t.Fatalf("primary requests = %d, want one per downstream request", got)
	}
	if got := len(backup.Requests()); got != 2 {
		t.Fatalf("backup requests = %d, want one per downstream request", got)
	}
}

func TestHandlerReturnsLastHostErrorWhenSkippedGroupHasNoBackup(t *testing.T) {
	upstream := fakeupstream.New(
		fakeupstream.Step{Status: http.StatusInternalServerError, Fixture: "openai/500.json"},
		fakeupstream.Step{Status: http.StatusOK, Fixture: "openai/success.json"},
	)
	defer upstream.Close()

	engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "gpt-4o",
		dialect.NewSet(dialect.NewOpenAI()),
		dialectGatewayGroup{id: 1, name: "only", upstreamURL: upstream.URL,
			apiKeys: []string{"sk-one", "sk-two"}},
	)
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions",
		bytes.NewBufferString(`{"model":"gpt-4o"}`))
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusInternalServerError || len(upstream.Requests()) != 1 ||
		!bytes.Contains(recorder.Body.Bytes(), []byte("internal_error")) {
		t.Fatalf("response/attempts = %d/%d body=%s, want safe 500/1",
			recorder.Code, len(upstream.Requests()), recorder.Body.String())
	}
}

func TestForwarderRewritesAliasedNonStreamingResponses(t *testing.T) {
	tests := []struct {
		name             string
		value            protocol.Protocol
		dialects         dialect.Set
		path             string
		requestBody      string
		upstreamResponse string
		responseField    string
	}{
		{
			name: "OpenAI", value: protocol.OpenAICompletions,
			dialects: dialect.NewSet(dialect.NewOpenAI()),
			path:     "/v1/chat/completions", requestBody: `{"model":"public-model"}`,
			upstreamResponse: `{"id":"chatcmpl-1","model":"provider-response"}`,
			responseField:    "model",
		},
		{
			name: "OpenAI Responses", value: protocol.OpenAIResponses,
			dialects: dialect.NewSet(dialect.NewOpenAIResponses()),
			path:     "/v1/responses", requestBody: `{"model":"public-model","input":"ping","store":false}`,
			upstreamResponse: `{"id":"chatcmpl-1","object":"chat.completion","created":123,"model":"provider-response","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`,
			responseField:    "model",
		},
		{
			name: "Anthropic", value: protocol.Anthropic,
			dialects: dialect.NewSet(dialect.NewAnthropic()),
			path:     "/v1/messages", requestBody: `{"model":"public-model"}`,
			upstreamResponse: `{"type":"message","model":"provider-response"}`,
			responseField:    "model",
		},
		{
			name: "Gemini", value: protocol.Gemini,
			dialects: dialect.NewSet(dialect.NewGemini()),
			path:     "/v1beta/models/public-model:generateContent", requestBody: `{}`,
			upstreamResponse: `{"modelVersion":"provider-response","candidates":[]}`,
			responseField:    "modelVersion",
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var receivedPath string
			var receivedHeader http.Header
			var receivedBody []byte
			upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
				receivedPath = request.URL.Path
				receivedHeader = request.Header.Clone()
				receivedBody, _ = io.ReadAll(request.Body)
				writer.Header().Set("Content-Type", "application/json")
				writer.Header().Set("Content-Length", strconv.Itoa(len(test.upstreamResponse)))
				setRepresentationMetadata(writer.Header())
				_, _ = writer.Write([]byte(test.upstreamResponse))
			}))
			defer upstream.Close()

			engine, _ := newDialectGatewayEngine(t, test.value, "public-model", test.dialects,
				dialectGatewayGroup{
					id: 1, name: test.name, upstreamURL: upstream.URL, apiKeys: []string{"provider-key"},
					models: []state.ModelConfig{{ID: "provider-model", Alias: "public-model"}},
				},
			)
			request := httptest.NewRequest(http.MethodPost, test.path, strings.NewReader(test.requestBody))
			request.Header.Set("Authorization", "Bearer gl-client")
			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, request)

			if recorder.Code != http.StatusOK {
				t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
			}
			if receivedHeader.Get("Accept-Encoding") != "identity" {
				t.Fatalf("upstream Accept-Encoding = %q, want identity", receivedHeader.Get("Accept-Encoding"))
			}
			if test.value == protocol.Gemini {
				if receivedPath != "/v1beta/models/provider-model:generateContent" {
					t.Fatalf("upstream path = %q", receivedPath)
				}
			} else {
				var upstreamBody struct {
					Model string `json:"model"`
				}
				if err := json.Unmarshal(receivedBody, &upstreamBody); err != nil || upstreamBody.Model != "provider-model" {
					t.Fatalf("upstream request body/model = %s / %q / %v", receivedBody, upstreamBody.Model, err)
				}
			}
			var downstreamBody map[string]json.RawMessage
			if err := json.Unmarshal(recorder.Body.Bytes(), &downstreamBody); err != nil {
				t.Fatalf("decode downstream response: %v", err)
			}
			var responseModel string
			if err := json.Unmarshal(downstreamBody[test.responseField], &responseModel); err != nil || responseModel != "public-model" {
				t.Fatalf("downstream model = %q, %v; body=%s", responseModel, err, recorder.Body.String())
			}
			if got := recorder.Header().Get("Content-Length"); got != strconv.Itoa(recorder.Body.Len()) {
				t.Fatalf("Content-Length = %q, want %d", got, recorder.Body.Len())
			}
			assertRepresentationMetadata(t, recorder.Header(), false)
		})
	}
}

func TestTransparentModelRoutePreservesWire(t *testing.T) {
	rawResponse := []byte("{\n  \"model\" : \"same-model\", \"n\": 9007199254740993\n}\n")
	var receivedHeader http.Header
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		receivedHeader = request.Header.Clone()
		writer.Header().Set("Content-Type", "application/json")
		writer.Header().Set("Content-Length", strconv.Itoa(len(rawResponse)))
		setRepresentationMetadata(writer.Header())
		_, _ = writer.Write(rawResponse)
	}))
	defer upstream.Close()

	engine, _ := newDialectGatewayEngine(t, protocol.OpenAICompletions, "same-model",
		dialect.NewSet(dialect.NewOpenAI()),
		dialectGatewayGroup{
			id: 1, name: "transparent", upstreamURL: upstream.URL, apiKeys: []string{"provider-key"},
			models:      []state.ModelConfig{{ID: "same-model"}},
			headerRules: state.HeaderRules{Set: map[string]string{"Accept-Encoding": "gzip"}},
		},
	)
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"same-model"}`))
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusOK || !bytes.Equal(recorder.Body.Bytes(), rawResponse) {
		t.Fatalf("transparent response = %d %q, want %q", recorder.Code, recorder.Body.Bytes(), rawResponse)
	}
	if receivedHeader.Get("Accept-Encoding") != "identity" {
		t.Fatalf("upstream Accept-Encoding = %q, want final identity normalization", receivedHeader.Get("Accept-Encoding"))
	}
	for name, want := range map[string]string{
		"ETag":           `"wire-v1"`,
		"Digest":         "sha-256=wire-digest",
		"Content-MD5":    "d2lyZQ==",
		"Content-Digest": "sha-256=:d2lyZQ==:",
		"Repr-Digest":    "sha-256=:cmVwcg==:",
	} {
		if got := recorder.Header().Get(name); got != want {
			t.Errorf("%s = %q, want preserved value %q", name, got, want)
		}
	}
	if got := recorder.Header().Get("Content-Range"); got != "" {
		t.Errorf("Content-Range = %q, want absent for 200 response", got)
	}
	if recorder.Header().Get("Signature") != "" || recorder.Header().Get("Signature-Input") != "" {
		t.Errorf("downstream response retained signatures: %#v", recorder.Header())
	}
}

func TestGatewayRewritesAliasedStreams(t *testing.T) {
	tests := []struct {
		name        string
		value       protocol.Protocol
		dialects    dialect.Set
		path        string
		requestBody string
		streamBody  string
		want        string
		unchanged   string
	}{
		{
			name: "OpenAI", value: protocol.OpenAICompletions,
			dialects: dialect.NewSet(dialect.NewOpenAI()),
			path:     "/v1/chat/completions", requestBody: `{"model":"public-model","stream":true}`,
			streamBody: "data: {\"id\":\"1\",\"model\":\"provider-model\",\"choices\":[]}\n\ndata: [DONE]\n\n",
			want:       `"model":"public-model"`, unchanged: "data: [DONE]\n\n",
		},
		{
			name: "Anthropic", value: protocol.Anthropic,
			dialects: dialect.NewSet(dialect.NewAnthropic()),
			path:     "/v1/messages", requestBody: `{"model":"public-model","stream":true}`,
			streamBody: "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"provider-model\"}}\n\n" +
				"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"unchanged\"}}\n\n",
			want: `"model":"public-model"`, unchanged: "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"unchanged\"}}\n\n",
		},
		{
			name: "Gemini", value: protocol.Gemini,
			dialects: dialect.NewSet(dialect.NewGemini()),
			path:     "/v1beta/models/public-model:streamGenerateContent", requestBody: `{}`,
			streamBody: "data: {\"modelVersion\":\"provider-model\",\"candidates\":[]}\n\n",
			want:       `"modelVersion":"public-model"`,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var receivedPath string
			var receivedHeader http.Header
			var receivedBody []byte
			upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
				receivedPath = request.URL.Path
				receivedHeader = request.Header.Clone()
				receivedBody, _ = io.ReadAll(request.Body)
				writer.Header().Set("Content-Type", "text/event-stream")
				writer.Header().Set("Content-Length", strconv.Itoa(len(test.streamBody)))
				setRepresentationMetadata(writer.Header())
				_, _ = writer.Write([]byte(test.streamBody))
				writer.(http.Flusher).Flush()
			}))
			defer upstream.Close()

			engine, _ := newDialectGatewayEngine(t, test.value, "public-model", test.dialects,
				dialectGatewayGroup{
					id: 1, name: test.name, upstreamURL: upstream.URL, apiKeys: []string{"provider-key"},
					models: []state.ModelConfig{{ID: "provider-model", Alias: "public-model"}},
				},
			)
			request := httptest.NewRequest(http.MethodPost, test.path, strings.NewReader(test.requestBody))
			request.Header.Set("Authorization", "Bearer gl-client")
			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, request)

			if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), test.want) {
				t.Fatalf("stream response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
			}
			if test.unchanged != "" && !strings.Contains(recorder.Body.String(), test.unchanged) {
				t.Fatalf("stream lost unchanged event: %q", recorder.Body.String())
			}
			if receivedHeader.Get("Accept-Encoding") != "identity" {
				t.Fatalf("upstream Accept-Encoding = %q", receivedHeader.Get("Accept-Encoding"))
			}
			if test.value == protocol.Gemini {
				if receivedPath != "/v1beta/models/provider-model:streamGenerateContent" {
					t.Fatalf("upstream path = %q", receivedPath)
				}
			} else {
				var upstreamBody struct {
					Model string `json:"model"`
				}
				if err := json.Unmarshal(receivedBody, &upstreamBody); err != nil || upstreamBody.Model != "provider-model" {
					t.Fatalf("upstream request model/body = %q / %s / %v", upstreamBody.Model, receivedBody, err)
				}
			}
			if recorder.Header().Get("Content-Length") != "" {
				t.Fatalf("stream Content-Length = %q, want removed", recorder.Header().Get("Content-Length"))
			}
			assertRepresentationMetadata(t, recorder.Header(), false)
		})
	}
}

func TestAnthropicGatewayNonStreamAuthAndForwarding(t *testing.T) {
	for _, test := range []struct {
		name        string
		authHeader  string
		apiKey      string
		version     string
		wantVersion string
	}{
		{name: "Bearer remains Anthropic", authHeader: "Bearer gl-client", wantVersion: anthropicDefaultVersionForTest},
		{name: "x-api-key carrier", apiKey: "gl-client", version: "2024-01-01", wantVersion: anthropicDefaultVersionForTest},
	} {
		t.Run(test.name, func(t *testing.T) {
			upstream := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"})
			defer upstream.Close()
			engine, _ := newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet",
				dialect.NewSet(dialect.NewAnthropic()),
				dialectGatewayGroup{id: 1, name: "anthropic", upstreamURL: upstream.URL, apiKeys: []string{"sk-anthropic-upstream"}},
			)
			body := `{"model":"claude-3-5-sonnet","messages":[{"role":"user","content":"ping"}]}`
			request := httptest.NewRequest(http.MethodPost, "/v1/messages?trace=true", strings.NewReader(body))
			if test.authHeader != "" {
				request.Header.Set("Authorization", test.authHeader)
			}
			if test.apiKey != "" {
				request.Header.Set("X-Api-Key", test.apiKey)
			}
			if test.version != "" {
				request.Header.Set("Anthropic-Version", test.version)
			}
			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, request)
			if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "1" ||
				!strings.Contains(recorder.Body.String(), `"type":"message"`) {
				t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
			}
			requests := upstream.Requests()
			if len(requests) != 1 {
				t.Fatalf("upstream requests = %d", len(requests))
			}
			got := requests[0]
			if got.Headers.Get("X-Api-Key") != "sk-anthropic-upstream" ||
				got.Headers.Get("Authorization") != "" ||
				got.Headers.Get("Anthropic-Version") != test.wantVersion ||
				got.RawQuery != "trace=true" || string(got.Body) != body {
				t.Fatalf("upstream request = %#v", got)
			}
		})
	}
}

const anthropicDefaultVersionForTest = "2023-06-01"

func TestAnthropicGatewayFailover(t *testing.T) {
	upstream := fakeupstream.New(
		fakeupstream.Step{Status: http.StatusUnauthorized, Fixture: "401.json"},
		fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"},
	)
	defer upstream.Close()
	engine, _ := newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet",
		dialect.NewSet(dialect.NewAnthropic()),
		dialectGatewayGroup{id: 1, name: "anthropic", upstreamURL: upstream.URL, apiKeys: []string{"sk-anthropic-one", "sk-anthropic-two"}},
	)
	request := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-3-5-sonnet"}`))
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)
	if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "2" {
		t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
	}
	requests := upstream.Requests()
	if len(requests) != 2 || requests[0].Headers.Get("X-Api-Key") != "sk-anthropic-one" ||
		requests[1].Headers.Get("X-Api-Key") != "sk-anthropic-two" {
		t.Fatalf("upstream requests = %#v", requests)
	}
	for _, secret := range []string{"sk-anthropic-one", "sk-anthropic-two"} {
		if strings.Contains(recorder.Body.String(), secret) || strings.Contains(recorder.Header().Get(debugHeaderKey), secret) {
			t.Fatalf("response exposes %q", secret)
		}
	}

	var clientErrorRequests atomic.Int64
	clientError := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		clientErrorRequests.Add(1)
		writer.Header().Set("Content-Type", "application/json")
		writer.WriteHeader(http.StatusBadRequest)
		_, _ = writer.Write([]byte(`{"type":"error","error":{"type":"invalid_request_error","message":"function calling is not supported with this model"}}`))
	}))
	defer clientError.Close()
	engine, _ = newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet",
		dialect.NewSet(dialect.NewAnthropic()),
		dialectGatewayGroup{id: 1, name: "anthropic", upstreamURL: clientError.URL, apiKeys: []string{"sk-one", "sk-two"}},
	)
	request = httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-3-5-sonnet"}`))
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder = httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)
	if recorder.Code != http.StatusBadRequest || clientErrorRequests.Load() != 2 || recorder.Header().Get(debugHeaderAttempts) != "2" {
		t.Fatalf("client error response = %d attempts=%s requests=%d body=%s", recorder.Code, recorder.Header().Get(debugHeaderAttempts), clientErrorRequests.Load(), recorder.Body.String())
	}
}

func TestAnthropicGatewayStream(t *testing.T) {
	firstEventSent := make(chan struct{})
	release := make(chan struct{})
	requestHeaders := make(chan http.Header, 1)
	var releaseOnce sync.Once
	defer releaseOnce.Do(func() { close(release) })
	primary := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		_, _ = io.Copy(io.Discard, request.Body)
		requestHeaders <- request.Header.Clone()
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\"}\n\n"))
		writer.(http.Flusher).Flush()
		close(firstEventSent)
		<-release
		_, _ = writer.Write([]byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"))
		writer.(http.Flusher).Flush()
	}))
	defer primary.Close()
	backup := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "stream.sse", Stream: true})
	defer backup.Close()

	engine, _ := newDialectGatewayEngine(t, protocol.Anthropic, "claude-3-5-sonnet",
		dialect.NewSet(dialect.NewAnthropic()),
		dialectGatewayGroup{
			id: 1, name: "primary", upstreamURL: primary.URL, apiKeys: []string{"sk-primary"},
			headerRules: state.HeaderRules{Set: map[string]string{"Accept-Encoding": "gzip"}},
		},
		dialectGatewayGroup{id: 2, name: "backup", upstreamURL: backup.URL, apiKeys: []string{"sk-backup"}},
	)
	gatewayServer := httptest.NewServer(engine)
	defer gatewayServer.Close()

	request, err := http.NewRequest(http.MethodPost, gatewayServer.URL+"/v1/messages", strings.NewReader(`{"model":"claude-3-5-sonnet","stream":true}`))
	if err != nil {
		t.Fatalf("create request: %v", err)
	}
	request.Header.Set("Authorization", "Bearer gl-client")
	response, err := gatewayServer.Client().Do(request)
	if err != nil {
		t.Fatalf("stream request: %v", err)
	}
	defer response.Body.Close()
	select {
	case <-firstEventSent:
	case <-time.After(time.Second):
		t.Fatal("first upstream event was not sent")
	}
	reader := bufio.NewReader(response.Body)
	firstLine, err := reader.ReadString('\n')
	if err != nil || firstLine != "event: message_start\n" {
		t.Fatalf("first streamed line = %q, %v", firstLine, err)
	}
	headers := <-requestHeaders
	if headers.Get("Accept-Encoding") != "identity" || headers.Get("X-Api-Key") != "sk-primary" {
		t.Fatalf("upstream headers = %#v", headers)
	}
	if len(backup.Requests()) != 0 || response.Header.Get(debugHeaderAttempts) != "1" {
		t.Fatalf("backup requests=%d headers=%v", len(backup.Requests()), response.Header)
	}
	releaseOnce.Do(func() { close(release) })
	rest, err := io.ReadAll(reader)
	if err != nil || !strings.Contains(string(rest), "message_stop") || strings.Contains(string(rest), `"code":`) {
		t.Fatalf("remaining stream = %q, %v", rest, err)
	}
	if values := response.Header.Values(debugHeaderKey); len(values) != 0 {
		t.Fatalf("debug key = %#v, want absent", values)
	}
}

func TestGeminiGatewayNonStreamAuthAndQuery(t *testing.T) {
	for _, test := range []struct {
		name       string
		target     string
		authHeader string
		googKey    string
		body       string
	}{
		{
			name:   "query key authenticates without forwarding",
			target: "/v1beta/models/gemini-2.5-pro:generateContent?key=gl-client&alt=json&trace=true",
			body:   "{",
		},
		{
			name:    "x-goog carrier remains Gemini",
			target:  "/v1beta/models/gemini-2.5-pro:generateContent?alt=json&trace=true",
			googKey: "gl-client",
			body:    `{"model":"different-model"}`,
		},
		{
			name:       "Bearer carrier remains Gemini",
			target:     "/v1beta/models/gemini-2.5-pro:generateContent?alt=json&trace=true",
			authHeader: "Bearer gl-client",
			body:       `{}`,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			upstream := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"})
			defer upstream.Close()
			engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro",
				dialect.NewSet(dialect.NewGemini()),
				dialectGatewayGroup{
					id: 1, name: "gemini", upstreamURL: upstream.URL,
					apiKeys: []string{"gemini-upstream-key"},
				},
			)
			request := httptest.NewRequest(http.MethodPost, test.target, strings.NewReader(test.body))
			if test.authHeader != "" {
				request.Header.Set("Authorization", test.authHeader)
			}
			if test.googKey != "" {
				request.Header.Set("X-Goog-Api-Key", test.googKey)
			}
			recorder := httptest.NewRecorder()
			engine.ServeHTTP(recorder, request)
			if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "1" {
				t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
			}
			requests := upstream.Requests()
			if len(requests) != 1 {
				t.Fatalf("upstream requests = %d", len(requests))
			}
			got := requests[0]
			query := gotQuery(t, got.RawQuery)
			if got.Path != "/v1beta/models/gemini-2.5-pro:generateContent" ||
				got.Headers.Get("X-Goog-Api-Key") != "gemini-upstream-key" ||
				got.Headers.Get("Authorization") != "" ||
				query.Get("key") != "" || query.Get("alt") != "" ||
				query.Get("tenant") != "" || query.Get("trace") != "true" ||
				string(got.Body) != test.body {
				t.Fatalf("upstream request = %#v query=%v", got, query)
			}
		})
	}
}

func TestGeminiGatewayFailover(t *testing.T) {
	upstream := fakeupstream.New(
		fakeupstream.Step{Status: http.StatusUnauthorized, Fixture: "401.json"},
		fakeupstream.Step{Status: http.StatusOK, Fixture: "success.json"},
	)
	defer upstream.Close()
	engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro",
		dialect.NewSet(dialect.NewGemini()),
		dialectGatewayGroup{id: 1, name: "gemini", upstreamURL: upstream.URL, apiKeys: []string{"gemini-key-one", "gemini-key-two"}},
	)
	request := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent?trace=true", strings.NewReader(`{}`))
	request.Header.Set("Authorization", "Bearer gl-client")
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)
	if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "2" {
		t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
	}
	requests := upstream.Requests()
	if len(requests) != 2 || requests[0].Headers.Get("X-Goog-Api-Key") != "gemini-key-one" ||
		requests[1].Headers.Get("X-Goog-Api-Key") != "gemini-key-two" ||
		requests[0].Path != requests[1].Path || requests[0].RawQuery != requests[1].RawQuery {
		t.Fatalf("upstream requests = %#v", requests)
	}
}

func TestGeminiGatewayStream(t *testing.T) {
	upstream := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "stream.sse", Stream: true})
	defer upstream.Close()
	engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro",
		dialect.NewSet(dialect.NewGemini()),
		dialectGatewayGroup{
			id: 1, name: "gemini-stream", upstreamURL: upstream.URL,
			apiKeys:     []string{"gemini-stream-key"},
			headerRules: state.HeaderRules{Set: map[string]string{"Accept-Encoding": "gzip"}},
		},
	)
	request := httptest.NewRequest(
		http.MethodPost,
		"/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=gl-client&alt=xml&trace=true",
		strings.NewReader(`{}`),
	)
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, request)
	if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "1" ||
		!strings.Contains(recorder.Body.String(), "data:") || strings.Contains(recorder.Body.String(), `"code":`) {
		t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String())
	}
	requests := upstream.Requests()
	if len(requests) != 1 {
		t.Fatalf("upstream requests = %d", len(requests))
	}
	query := gotQuery(t, requests[0].RawQuery)
	if got := query["alt"]; len(got) != 1 || got[0] != "sse" {
		t.Fatalf("alt = %#v", got)
	}
	if query.Get("key") != "" || query.Get("trace") != "true" || query.Get("tenant") != "" ||
		requests[0].Headers.Get("Accept-Encoding") != "identity" ||
		requests[0].Headers.Get("X-Goog-Api-Key") != "gemini-stream-key" {
		t.Fatalf("upstream request = %#v query=%v", requests[0], query)
	}
}

func TestGeminiGatewayCompressed(t *testing.T) {
	t.Run("SDK-normalized response remains one logical attempt", func(t *testing.T) {
		compressed := fakeupstream.New(fakeupstream.Step{
			Status: http.StatusOK, Fixture: "stream.sse", Stream: true,
			Headers: http.Header{"Content-Encoding": {"gzip"}},
		})
		defer compressed.Close()
		backup := fakeupstream.New(fakeupstream.Step{Status: http.StatusOK, Fixture: "stream.sse", Stream: true})
		defer backup.Close()
		engine, _ := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro",
			dialect.NewSet(dialect.NewGemini()),
			dialectGatewayGroup{id: 1, name: "compressed", upstreamURL: compressed.URL, apiKeys: []string{"bad-key"}},
			dialectGatewayGroup{id: 2, name: "backup", upstreamURL: backup.URL, apiKeys: []string{"good-key"}},
		)
		request := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:streamGenerateContent", strings.NewReader(`{}`))
		request.Header.Set("Authorization", "Bearer gl-client")
		recorder := httptest.NewRecorder()
		engine.ServeHTTP(recorder, request)
		if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "1" ||
			!strings.Contains(recorder.Body.String(), "data:") ||
			len(compressed.Requests()) != 1 || len(backup.Requests()) != 0 {
			t.Fatalf("response = %d headers=%v compressed=%d backup=%d body=%s", recorder.Code, recorder.Header(), len(compressed.Requests()), len(backup.Requests()), recorder.Body.String())
		}
	})

	t.Run("transport encoding metadata does not trigger candidate retry", func(t *testing.T) {
		first := fakeupstream.New(fakeupstream.Step{
			Status: http.StatusOK, Fixture: "stream.sse", Stream: true,
			Headers: http.Header{"Content-Encoding": {"gzip"}},
		})
		defer first.Close()
		second := fakeupstream.New(fakeupstream.Step{
			Status: http.StatusOK, Fixture: "stream.sse", Stream: true,
			Headers: http.Header{"Content-Encoding": {"br"}},
		})
		defer second.Close()
		engine, registry := newDialectGatewayEngine(t, protocol.Gemini, "gemini-2.5-pro",
			dialect.NewSet(dialect.NewGemini()),
			dialectGatewayGroup{id: 1, name: "first", upstreamURL: first.URL, apiKeys: []string{"secret-one"}},
			dialectGatewayGroup{id: 2, name: "second", upstreamURL: second.URL, apiKeys: []string{"secret-two"}},
		)
		request := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:streamGenerateContent", strings.NewReader(`{}`))
		request.Header.Set("Authorization", "Bearer gl-client")
		recorder := httptest.NewRecorder()
		engine.ServeHTTP(recorder, request)
		if recorder.Code != http.StatusOK || recorder.Header().Get(debugHeaderAttempts) != "1" ||
			!strings.Contains(recorder.Body.String(), "data:") ||
			len(first.Requests()) != 1 || len(second.Requests()) != 0 ||
			len(registry.CollectCredentialCandidates([]uint{1, 2}, nil, time.Time{})) != 2 {
			t.Fatalf("response = %d headers=%v candidates=%d body=%s", recorder.Code, recorder.Header(), len(registry.CollectCredentialCandidates([]uint{1, 2}, nil, time.Time{})), recorder.Body.String())
		}
		for _, forbidden := range []string{"secret-one", "secret-two"} {
			if strings.Contains(recorder.Body.String(), forbidden) {
				t.Fatalf("response exposes %q: %s", forbidden, recorder.Body.String())
			}
		}
	})
}

func gotQuery(t *testing.T, raw string) url.Values {
	t.Helper()
	values, err := url.ParseQuery(raw)
	if err != nil {
		t.Fatalf("parse query %q: %v", raw, err)
	}
	return values
}