fix(home): implement home dispatch headers and enhance Gemini model handling

sdk/cliproxy/auth/conductor.go

@@ -3231,6 +3231,79 @@ func setHomeUserAPIKeyOnGinContext(ctx context.Context, apiKey string) {
 	ginCtx.Set("userApiKey", apiKey)
 }
 
+func homeDispatchHeaders(ctx context.Context, headers http.Header) http.Header {
+	apiKey, ok := homeQueryCredentialFromContext(ctx)
+	if !ok {
+		return headers
+	}
+	out := headers.Clone()
+	if out == nil {
+		out = http.Header{}
+	}
+	if out.Get("Authorization") != "" || out.Get("X-Goog-Api-Key") != "" || out.Get("X-Api-Key") != "" {
+		return out
+	}
+	out.Set("X-Goog-Api-Key", apiKey)
+	return out
+}
+
+func homeQueryCredentialFromContext(ctx context.Context) (string, bool) {
+	if ctx == nil {
+		return "", false
+	}
+	if queryCtx, ok := ctx.Value("gin").(interface{ Query(string) string }); ok && queryCtx != nil {
+		if apiKey := strings.TrimSpace(queryCtx.Query("key")); apiKey != "" {
+			return apiKey, true
+		}
+		if apiKey := strings.TrimSpace(queryCtx.Query("auth_token")); apiKey != "" {
+			return apiKey, true
+		}
+	}
+	ginCtx, ok := ctx.Value("gin").(interface{ Get(string) (any, bool) })
+	if !ok || ginCtx == nil {
+		return "", false
+	}
+	rawMetadata, ok := ginCtx.Get("accessMetadata")
+	if !ok {
+		return "", false
+	}
+	source := accessMetadataSource(rawMetadata)
+	if source != "query-key" && source != "query-auth-token" {
+		return "", false
+	}
+	rawAPIKey, ok := ginCtx.Get("userApiKey")
+	if !ok {
+		return "", false
+	}
+	apiKey := contextStringValue(rawAPIKey)
+	if apiKey == "" {
+		return "", false
+	}
+	return apiKey, true
+}
+
+func accessMetadataSource(raw any) string {
+	switch v := raw.(type) {
+	case map[string]string:
+		return strings.TrimSpace(v["source"])
+	case map[string]any:
+		return contextStringValue(v["source"])
+	default:
+		return ""
+	}
+}
+
+func contextStringValue(raw any) string {
+	switch v := raw.(type) {
+	case string:
+		return strings.TrimSpace(v)
+	case []byte:
+		return strings.TrimSpace(string(v))
+	default:
+		return ""
+	}
+}
+
 func homeExecutionSessionIDFromMetadata(meta map[string]any) string {
 	if len(meta) == 0 {
 		return ""
@@ -3352,8 +3425,9 @@ func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts clipro
 
 	requestedModel := requestedModelFromMetadata(opts.Metadata, model)
 	sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata)
+	dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers)
 
-	raw, err := client.RPopAuth(ctx, requestedModel, sessionID, opts.Headers, count)
+	raw, err := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, count)
 	if err != nil {
 		return nil, nil, "", &Error{Code: "auth_not_found", Message: err.Error(), HTTPStatus: http.StatusServiceUnavailable}
 	}

sdk/cliproxy/auth/home_dispatch_headers_test.go

@@ -0,0 +1,87 @@
+package auth
+
+import (
+	"context"
+	"net/http"
+	"testing"
+)
+
+type homeDispatchTestGinContext struct {
+	values map[string]any
+	query  map[string]string
+}
+
+func (c homeDispatchTestGinContext) Get(key string) (any, bool) {
+	v, ok := c.values[key]
+	return v, ok
+}
+
+func (c homeDispatchTestGinContext) Query(key string) string {
+	if c.query == nil {
+		return ""
+	}
+	return c.query[key]
+}
+
+func TestHomeDispatchHeadersAddsQueryKeyCredential(t *testing.T) {
+	ginCtx := homeDispatchTestGinContext{query: map[string]string{"key": "12345"}}
+	ctx := context.WithValue(context.Background(), "gin", ginCtx)
+	headers := http.Header{"User-Agent": {"client"}}
+
+	got := homeDispatchHeaders(ctx, headers)
+
+	if got.Get("X-Goog-Api-Key") != "12345" {
+		t.Fatalf("X-Goog-Api-Key = %q, want %q", got.Get("X-Goog-Api-Key"), "12345")
+	}
+	if headers.Get("X-Goog-Api-Key") != "" {
+		t.Fatalf("original headers were mutated: %v", headers)
+	}
+}
+
+func TestHomeDispatchHeadersAddsQueryCredentialFromAccessMetadata(t *testing.T) {
+	ginCtx := homeDispatchTestGinContext{values: map[string]any{
+		"accessMetadata": map[string]string{"source": "query-key"},
+		"userApiKey":     "12345",
+	}}
+	ctx := context.WithValue(context.Background(), "gin", ginCtx)
+	headers := http.Header{"User-Agent": {"client"}}
+
+	got := homeDispatchHeaders(ctx, headers)
+
+	if got.Get("X-Goog-Api-Key") != "12345" {
+		t.Fatalf("X-Goog-Api-Key = %q, want %q", got.Get("X-Goog-Api-Key"), "12345")
+	}
+	if headers.Get("X-Goog-Api-Key") != "" {
+		t.Fatalf("original headers were mutated: %v", headers)
+	}
+}
+
+func TestHomeDispatchHeadersKeepsExistingCredentialHeader(t *testing.T) {
+	ginCtx := homeDispatchTestGinContext{query: map[string]string{"key": "query-key"}}
+	ctx := context.WithValue(context.Background(), "gin", ginCtx)
+	headers := http.Header{"X-Goog-Api-Key": {"header-key"}}
+
+	got := homeDispatchHeaders(ctx, headers)
+
+	if got.Get("X-Goog-Api-Key") != "header-key" {
+		t.Fatalf("X-Goog-Api-Key = %q, want %q", got.Get("X-Goog-Api-Key"), "header-key")
+	}
+}
+
+func TestHomeDispatchHeadersIgnoresHeaderCredentialSource(t *testing.T) {
+	ginCtx := homeDispatchTestGinContext{values: map[string]any{
+		"accessMetadata": map[string]string{"source": "authorization"},
+		"userApiKey":     "12345",
+	}}
+	ctx := context.WithValue(context.Background(), "gin", ginCtx)
+	headers := http.Header{"Authorization": {"Bearer 12345"}}
+
+	got := homeDispatchHeaders(ctx, headers)
+
+	if got.Get("X-Goog-Api-Key") != "" {
+		t.Fatalf("X-Goog-Api-Key = %q, want empty", got.Get("X-Goog-Api-Key"))
+	}
+	if got.Get("Authorization") != "Bearer 12345" {
+		t.Fatalf("Authorization = %q, want %q", got.Get("Authorization"), "Bearer 12345")
+	}
+}