From 6bd4e833a81841f5871f847ae838a45fb39b04b2 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Fri, 26 Jun 2026 13:33:28 +0300 Subject: [PATCH 01/16] feat: add sosana image provider --- README.md | 4 +- config.yaml.example | 15 + docs/getting-started/configuration.md | 20 +- docs/index.md | 2 +- docs/providers/index.md | 1 + docs/providers/sosana.md | 43 +++ internal/config/config.go | 8 +- internal/config/config_test.go | 81 +++++ internal/config/utils.go | 19 +- internal/converter/openai/types.go | 10 +- internal/converter/sosana/images.go | 218 +++++++++++++ internal/converter/sosana/images_test.go | 138 ++++++++ internal/litellmdb/model_table/model_table.go | 5 +- .../litellmdb/model_table/model_table_test.go | 16 + internal/models/manager.go | 2 + internal/models/price_calculator_test.go | 30 ++ internal/proxy/cometapi_test.go | 105 ++++-- internal/proxy/errors.go | 17 + internal/proxy/image_response.go | 102 ++++++ internal/proxy/proxy.go | 104 +++++- internal/proxy/proxy_log.go | 36 ++- internal/proxy/sosana.go | 296 +++++++++++++++++ internal/proxy/sosana_live_test.go | 63 ++++ internal/proxy/sosana_test.go | 304 ++++++++++++++++++ internal/proxy/upstream_masking.go | 36 +++ internal/proxy/upstream_masking_test.go | 296 +++++++++++++++++ 26 files changed, 1912 insertions(+), 59 deletions(-) create mode 100644 docs/providers/sosana.md create mode 100644 internal/converter/sosana/images.go create mode 100644 internal/converter/sosana/images_test.go create mode 100644 internal/proxy/image_response.go create mode 100644 internal/proxy/sosana.go create mode 100644 internal/proxy/sosana_live_test.go create mode 100644 internal/proxy/sosana_test.go create mode 100644 internal/proxy/upstream_masking.go create mode 100644 internal/proxy/upstream_masking_test.go diff --git a/README.md b/README.md index 472ad411..1c209054 100644 --- a/README.md +++ b/README.md @@ -7,11 +7,11 @@ [![GitHub Pages](https://img.shields.io/badge/docs-GitHub%20Pages-blue)](https://mixaill76.github.io/auto_ai_router/) [![license](https://img.shields.io/github/license/MiXaiLL76/auto_ai_router.svg)](https://github.com/MiXaiLL76/auto_ai_router/blob/main/LICENSE) -High-performance proxy router for LLM APIs with automatic load balancing, rate limiting, and fail2ban protection. Routes requests to OpenAI, Vertex AI, Gemini AI Studio, Anthropic, Comet API, and other Auto AI Router instances. +High-performance proxy router for LLM and image APIs with automatic load balancing, rate limiting, and fail2ban protection. Routes requests to OpenAI, Vertex AI, Gemini AI Studio, Anthropic, Comet API, Sosana.art images, and other Auto AI Router instances. ## Key Features -- **Multi-provider support** — OpenAI, Vertex AI, Gemini, Anthropic, Comet API, Proxy chains +- **Multi-provider support** — OpenAI, Vertex AI, Gemini, Anthropic, Comet API, Sosana.art images, Proxy chains - **Round-robin load balancing** — across multiple credentials per model - **Rate limiting** — per-credential and per-model RPM/TPM controls - **Fail2ban** — automatic provider banning on repeated errors diff --git a/config.yaml.example b/config.yaml.example index db73d944..bf4f70ec 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -53,6 +53,7 @@ credentials: rpm: 100 tpm: 50000 weight: 100 # Optional: weighted round-robin share (default 1). See docs/advanced/balancing.md + # mask_upstream_errors: true # Enable for OpenAI-compatible resellers whose raw errors must not reach clients. - name: "vertex_ai" type: "vertex-ai" @@ -77,6 +78,14 @@ credentials: rpm: 60 tpm: -1 + # Sosana.art image generation API (OpenAI Images-compatible through the router) + - name: "sosana_images" + type: "sosana" + api_key: "os.environ/SOSANA_API_KEY" + base_url: "https://sosana.art" + rpm: 60 + tpm: -1 + # CheapGPT / AIProductiv using the Anthropic-compatible Messages API - name: "cheapgpt_anthropic" type: "anthropic" @@ -107,6 +116,12 @@ models: rpm: 100 tpm: 50000 + # Sosana.art image model exposed through /v1/images/generations and /v1/images/edits. + - name: "nano-banana" + credential: sosana_images + rpm: 60 + tpm: -1 + # Comet API Claude aliases for public model names. - name: "anthropic/claude-haiku-4.5" model: "claude-haiku-4-5-20251001" diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md index 12549e40..b96e6310 100644 --- a/docs/getting-started/configuration.md +++ b/docs/getting-started/configuration.md @@ -154,16 +154,16 @@ Each credential defines a connection to an LLM provider. See [Providers](../prov Common fields for all credentials: -| Field | Type | Description | -| ------------------ | ------ | ------------------------------------------------------------------------------------------- | -| `name` | string | Unique credential identifier | -| `type` | string | Provider type: `openai`, `anthropic`, `cometapi`, `vertex-ai`, `gemini`, `bedrock`, `proxy` | -| `rpm` | int | Requests per minute limit (-1 = unlimited) | -| `tpm` | int | Tokens per minute limit (-1 = unlimited) | -| `is_fallback` | bool | Use as fallback when primary credentials are exhausted | -| `scopes` | list | Optional client scopes allowed to use and see this credential | -| `denied_scopes` | list | Optional client scopes that must not use or see this credential | -| `forbidden_scopes` | list | Alias for `denied_scopes` | +| Field | Type | Description | +| ------------------ | ------ | ----------------------------------------------------------------------------------------------------- | +| `name` | string | Unique credential identifier | +| `type` | string | Provider type: `openai`, `anthropic`, `cometapi`, `sosana`, `vertex-ai`, `gemini`, `bedrock`, `proxy` | +| `rpm` | int | Requests per minute limit (-1 = unlimited) | +| `tpm` | int | Tokens per minute limit (-1 = unlimited) | +| `is_fallback` | bool | Use as fallback when primary credentials are exhausted | +| `scopes` | list | Optional client scopes allowed to use and see this credential | +| `denied_scopes` | list | Optional client scopes that must not use or see this credential | +| `forbidden_scopes` | list | Alias for `denied_scopes` | ### Scoped credential visibility diff --git a/docs/index.md b/docs/index.md index dfb3be55..fdddfcbc 100644 --- a/docs/index.md +++ b/docs/index.md @@ -32,7 +32,7 @@ graph LR ## Features -- **Multi-provider routing** — OpenAI, Vertex AI, Gemini AI Studio, Anthropic, Comet API +- **Multi-provider routing** — OpenAI, Vertex AI, Gemini AI Studio, Anthropic, Comet API, Sosana.art images - **Proxy chains** — forward to other Auto AI Router instances as fallback - **Round-robin balancing** — distribute load across multiple credentials - **Two-level rate limiting** — per-credential RPM/TPM + per-model limits diff --git a/docs/providers/index.md b/docs/providers/index.md index 1c7d4d8c..69db9fbb 100644 --- a/docs/providers/index.md +++ b/docs/providers/index.md @@ -9,6 +9,7 @@ Auto AI Router supports multiple LLM providers. Each provider type has its own a | [OpenAI](openai.md) | `openai` | `api_key`, `base_url` | API Key | | [Anthropic](anthropic.md) | `anthropic` | `api_key`, `base_url` | API Key | | [Comet API](cometapi.md) | `cometapi` | `api_key`, `base_url` | API Key | +| [Sosana.art](sosana.md) | `sosana` | `api_key`, `base_url` | Bearer Token | | [AWS Bedrock](bedrock.md) | `bedrock` | `api_key`, `base_url` | Bearer Token | | [Vertex AI](vertex.md) | `vertex-ai` | `project_id`, `location`, `credentials_file` or `credentials_json` | OAuth2 / Service Account | | [Gemini AI Studio](gemini.md) | `gemini` | `api_key`, `base_url` | API Key | diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md new file mode 100644 index 00000000..aeca2fa0 --- /dev/null +++ b/docs/providers/sosana.md @@ -0,0 +1,43 @@ +# Sosana.art + +Sosana.art is supported as an image-only provider for the OpenAI-compatible +Images API. The router accepts `/v1/images/generations` and `/v1/images/edits`, +submits a Sosana Banana async task, polls it, and returns an OpenAI Images +response with `data[].url`. + +Chat Completions, Responses API, Embeddings, video, and slides are not routed to +Sosana in this integration. + +## Configuration + +```yaml +credentials: + - name: "sosana_images" + type: "sosana" + api_key: "os.environ/SOSANA_API_KEY" + base_url: "https://sosana.art" + rpm: 60 + tpm: -1 + +models: + - name: "nano-banana" + credential: sosana_images + rpm: 60 + tpm: -1 +``` + +## Behavior + +- `n` must be `1`. +- `response_format: "b64_json"` is accepted, but Sosana results are returned as + URLs because Sosana provides `result_file_url`. +- `/v1/images/edits` sends uploaded images as `data:image/...;base64,...` + values in Sosana `image_urls`. +- Mask images are not supported. + +## Error Masking + +Sosana upstream HTTP errors and terminal task errors are masked before they are +returned to clients or written to structured logs. The router preserves the +appropriate HTTP status but replaces provider details with neutral +OpenAI-compatible error bodies. diff --git a/internal/config/config.go b/internal/config/config.go index 453effb9..4a9384f2 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -26,6 +26,7 @@ const ( ProviderTypeGemini ProviderType = "gemini" ProviderTypeAnthropic ProviderType = "anthropic" ProviderTypeCometAPI ProviderType = "cometapi" + ProviderTypeSosana ProviderType = "sosana" ProviderTypeBedrock ProviderType = "bedrock" ProviderTypeProxy ProviderType = "proxy" ) @@ -40,7 +41,7 @@ func (p ProviderType) LogValue() slog.Value { // IsValid checks if the provider type is valid func (p ProviderType) IsValid() bool { switch p { - case ProviderTypeOpenAI, ProviderTypeVertexAI, ProviderTypeGemini, ProviderTypeAnthropic, ProviderTypeCometAPI, ProviderTypeBedrock, ProviderTypeProxy: + case ProviderTypeOpenAI, ProviderTypeVertexAI, ProviderTypeGemini, ProviderTypeAnthropic, ProviderTypeCometAPI, ProviderTypeSosana, ProviderTypeBedrock, ProviderTypeProxy: return true } return false @@ -50,6 +51,8 @@ func normalizeProviderType(raw string) ProviderType { switch strings.ToLower(strings.TrimSpace(raw)) { case "comet-api", "comet_api": return ProviderTypeCometAPI + case "sosana-art", "sosana_art": + return ProviderTypeSosana default: return ProviderType(strings.ToLower(strings.TrimSpace(raw))) } @@ -683,7 +686,6 @@ func (c *CredentialConfig) UnmarshalYAML(value *yaml.Node) error { if c.IsFallback, err = parseField(temp.IsFallback, false, strconv.ParseBool, "is_fallback for credential '"+c.Name+"'"); err != nil { return err } - // Copy models decoded via YAML anchors / inline definitions c.Models = temp.Models @@ -1292,7 +1294,7 @@ func (c *Config) Validate() error { // Validate provider type if !cred.Type.IsValid() { - return fmt.Errorf("credential %s: invalid type: %s (must be 'openai', 'vertex-ai', 'gemini', 'anthropic', 'cometapi', 'bedrock', or 'proxy')", cred.Name, cred.Type) + return fmt.Errorf("credential %s: invalid type: %s (must be 'openai', 'vertex-ai', 'gemini', 'anthropic', 'cometapi', 'sosana', 'bedrock', or 'proxy')", cred.Name, cred.Type) } if cred.AuthType != "" && cred.AuthType != "bearer" && cred.AuthType != "x-api-key" { return fmt.Errorf("credential %s: invalid auth_type: %s (must be 'bearer' or 'x-api-key')", cred.Name, cred.AuthType) diff --git a/internal/config/config_test.go b/internal/config/config_test.go index d54f3fff..1986aab0 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -651,6 +651,7 @@ func TestProviderType_IsValid(t *testing.T) { {"openai", ProviderTypeOpenAI, true}, {"vertex-ai", ProviderTypeVertexAI, true}, {"cometapi", ProviderTypeCometAPI, true}, + {"sosana", ProviderTypeSosana, true}, {"invalid", ProviderType("azure"), false}, {"empty", ProviderType(""), false}, } @@ -676,6 +677,86 @@ rpm: 60 assert.Equal(t, ProviderTypeCometAPI, cred.Type) } +func TestCredentialConfig_MaskUpstreamErrors(t *testing.T) { + var cred CredentialConfig + err := yaml.Unmarshal([]byte(` +name: sosana +type: openai +api_key: key +base_url: https://api.sosana.example/v1 +mask_upstream_errors: true +rpm: 60 +`), &cred) + + require.NoError(t, err) + assert.True(t, cred.MaskUpstreamErrors) +} + +func TestCredentialConfig_NormalizeSosanaProviderType(t *testing.T) { + tests := []struct { + name string + raw string + }{ + {"canonical", "sosana"}, + {"dash alias", "sosana-art"}, + {"underscore alias", "sosana_art"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var cred CredentialConfig + err := yaml.Unmarshal([]byte(` +name: sosana +type: `+tt.raw+` +api_key: key +base_url: https://sosana.art +rpm: 60 +`), &cred) + + require.NoError(t, err) + assert.Equal(t, ProviderTypeSosana, cred.Type) + }) + } +} + +func TestConfig_Validate_SosanaRequiresAPIKeyAndBaseURL(t *testing.T) { + tests := []struct { + name string + apiKey string + baseURL string + wantErr string + }{ + {name: "valid", apiKey: "key", baseURL: "https://sosana.art"}, + {name: "missing api key", baseURL: "https://sosana.art", wantErr: "api_key is required"}, + {name: "missing base url", apiKey: "key", wantErr: "base_url is required"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := &Config{ + Server: ServerConfig{ + Port: 8080, + MaxBodySizeMB: 10, + MasterKey: "test-key", + RequestTimeout: 30 * time.Second, + }, + Credentials: []CredentialConfig{ + {Name: "sosana", Type: ProviderTypeSosana, APIKey: tt.apiKey, BaseURL: tt.baseURL, RPM: 10}, + }, + Fail2Ban: Fail2BanConfig{MaxAttempts: 3}, + } + + err := cfg.Validate() + if tt.wantErr == "" { + assert.NoError(t, err) + } else { + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + } + }) + } +} + func TestConfig_Validate_VertexAI(t *testing.T) { tests := []struct { name string diff --git a/internal/config/utils.go b/internal/config/utils.go index 696821a5..0032ccb7 100644 --- a/internal/config/utils.go +++ b/internal/config/utils.go @@ -118,14 +118,15 @@ func PrintConfig(logger *slog.Logger, cfg *Config) { ) for i, cred := range cfg.Credentials { credLog := map[string]any{ - "name": cred.Name, - "type": cred.Type, - "base_url": cred.BaseURL, - "auth_type": cred.AuthType, - "rpm": rpmToString(cred.RPM), - "tpm": tpmToString(cred.TPM), - "is_fallback": cred.IsFallback, - "fallback_priority": cred.FallbackPriority, + "name": cred.Name, + "type": cred.Type, + "base_url": cred.BaseURL, + "auth_type": cred.AuthType, + "mask_upstream_errors": cred.MaskUpstreamErrors, + "rpm": rpmToString(cred.RPM), + "tpm": tpmToString(cred.TPM), + "is_fallback": cred.IsFallback, + "fallback_priority": cred.FallbackPriority, } // Add Vertex AI specific fields if present @@ -236,7 +237,7 @@ func banDurationToString(d time.Duration) string { func convertMapToArgs(m map[string]any) []any { // Define preferred order of keys keyOrder := []string{ - "name", "type", "base_url", "auth_type", "api_key", "project_id", "location", + "name", "type", "base_url", "auth_type", "mask_upstream_errors", "api_key", "project_id", "location", "credentials_file", "credentials_json", "rpm", "tpm", "is_fallback", } diff --git a/internal/converter/openai/types.go b/internal/converter/openai/types.go index 565d7c11..a621f065 100644 --- a/internal/converter/openai/types.go +++ b/internal/converter/openai/types.go @@ -225,9 +225,13 @@ type OpenAIImageUsage struct { // OpenAIImageResponse represents OpenAI image response type OpenAIImageResponse struct { - Created int64 `json:"created"` - Data []OpenAIImageData `json:"data"` - Usage *OpenAIImageUsage `json:"usage,omitempty"` + Created int64 `json:"created"` + Background string `json:"background,omitempty"` + Data []OpenAIImageData `json:"data"` + OutputFormat string `json:"output_format,omitempty"` + Quality string `json:"quality,omitempty"` + Size string `json:"size,omitempty"` + Usage *OpenAIImageUsage `json:"usage,omitempty"` } // Embedding types diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go new file mode 100644 index 00000000..899d0d88 --- /dev/null +++ b/internal/converter/sosana/images.go @@ -0,0 +1,218 @@ +package sosana + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "strconv" + "strings" + + "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" + "github.com/mixaill76/auto_ai_router/internal/converter/openai" +) + +const maxMultipartImageBytes = 20 * 1024 * 1024 + +const ( + StatusProcessing = "PROCESSING" + StatusCompleted = "COMPLETED" + StatusFailed = "FAILED" + StatusModerated = "MODERATED" +) + +type BananaCreateRequest struct { + Prompt string `json:"prompt"` + ImageURLs []string `json:"image_urls,omitempty"` + Model string `json:"model,omitempty"` + AspectRatio string `json:"aspect_ratio,omitempty"` +} + +type BananaTaskResponse struct { + UID string `json:"uid"` + Status string `json:"status"` + Prompt string `json:"prompt"` + ResultFileURL *string `json:"result_file_url"` + Error *string `json:"error"` +} + +func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, error) { + var req openai.OpenAIImageRequest + if err := json.Unmarshal(openAIBody, &req); err != nil { + return nil, fmt.Errorf("failed to parse OpenAI image request: %w", err) + } + if err := validateImageCount(req.N); err != nil { + return nil, err + } + prompt := strings.TrimSpace(req.Prompt) + if prompt == "" { + return nil, fmt.Errorf("image generation request missing prompt") + } + model := strings.TrimSpace(req.Model) + if model == "" { + model = modelID + } + return json.Marshal(BananaCreateRequest{ + Prompt: prompt, + Model: model, + AspectRatio: SizeToAspectRatio(req.Size), + }) +} + +func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, error) { + mediaType, params, err := mime.ParseMediaType(contentType) + if err != nil { + return nil, fmt.Errorf("failed to parse image edit content type: %w", err) + } + if !strings.HasPrefix(mediaType, "multipart/form-data") { + return nil, fmt.Errorf("image edits require multipart/form-data content type") + } + boundary := params["boundary"] + if boundary == "" { + return nil, fmt.Errorf("missing multipart boundary in content type") + } + + fields := make(map[string]string) + imageURLs := make([]string, 0, 1) + reader := multipart.NewReader(bytes.NewReader(openAIBody), boundary) + for { + part, err := reader.NextPart() + if err == io.EOF { + break + } + if err != nil { + return nil, fmt.Errorf("failed to read multipart image edit payload: %w", err) + } + + formName := part.FormName() + if formName == "" { + continue + } + data, err := readLimited(part, maxMultipartImageBytes) + if err != nil { + return nil, err + } + if part.FileName() == "" { + fields[formName] = strings.TrimSpace(string(data)) + continue + } + if formName == "mask" { + return nil, fmt.Errorf("sosana image edits do not support mask") + } + if formName != "image" && formName != "images" && formName != "image[]" { + continue + } + mimeType := detectImageMIMEType(part.Header.Get("Content-Type"), data) + imageURLs = append(imageURLs, "data:"+mimeType+";base64,"+base64.StdEncoding.EncodeToString(data)) + } + + if err := validateImageCountString(fields["n"]); err != nil { + return nil, err + } + prompt := strings.TrimSpace(fields["prompt"]) + if prompt == "" { + return nil, fmt.Errorf("image edit request missing prompt field") + } + model := strings.TrimSpace(fields["model"]) + if model == "" { + model = modelID + } + if len(imageURLs) == 0 { + return nil, fmt.Errorf("image edit request missing image") + } + return json.Marshal(BananaCreateRequest{ + Prompt: prompt, + ImageURLs: imageURLs, + Model: model, + AspectRatio: SizeToAspectRatio(fields["size"]), + }) +} + +func OpenAIImageResponse(task BananaTaskResponse) ([]byte, error) { + if task.ResultFileURL == nil || strings.TrimSpace(*task.ResultFileURL) == "" { + return nil, fmt.Errorf("sosana task completed without result_file_url") + } + resp := openai.OpenAIImageResponse{ + Created: converterutil.GetCurrentTimestamp(), + Data: []openai.OpenAIImageData{ + {URL: strings.TrimSpace(*task.ResultFileURL)}, + }, + } + return json.Marshal(resp) +} + +func CreateURL(baseURL string) string { + return strings.TrimSuffix(baseURL, "/") + "/api/banana/create-async" +} + +func PollURL(baseURL, uid string) string { + return strings.TrimSuffix(baseURL, "/") + "/api/banana/" + uid +} + +func SizeToAspectRatio(size string) string { + switch strings.TrimSpace(size) { + case "", "auto": + return "auto" + case "256x256", "512x512", "1024x1024", "4096x4096": + return "1:1" + case "1024x1536": + return "2:3" + case "1536x1024": + return "3:2" + case "1024x1792", "1080x1920", "4096x7168": + return "9:16" + case "1792x1024", "1920x1080", "7168x4096": + return "16:9" + default: + return "auto" + } +} + +func validateImageCount(n *int) error { + if n == nil || *n == 1 { + return nil + } + return fmt.Errorf("sosana supports n=1 only") +} + +func validateImageCountString(raw string) error { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil + } + n, err := strconv.Atoi(raw) + if err != nil { + return fmt.Errorf("invalid image count: %w", err) + } + if n != 1 { + return fmt.Errorf("sosana supports n=1 only") + } + return nil +} + +func readLimited(r io.Reader, limit int64) ([]byte, error) { + data, err := io.ReadAll(io.LimitReader(r, limit+1)) + if err != nil { + return nil, fmt.Errorf("failed to read multipart part: %w", err) + } + if int64(len(data)) > limit { + return nil, fmt.Errorf("multipart image exceeds %d bytes", limit) + } + return data, nil +} + +func detectImageMIMEType(header string, data []byte) string { + header = strings.ToLower(strings.TrimSpace(header)) + if strings.HasPrefix(header, "image/") { + return header + } + detected := http.DetectContentType(data) + if strings.HasPrefix(detected, "image/") { + return detected + } + return "application/octet-stream" +} diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go new file mode 100644 index 00000000..d17c8203 --- /dev/null +++ b/internal/converter/sosana/images_test.go @@ -0,0 +1,138 @@ +package sosana + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "mime/multipart" + "strings" + "testing" + + "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestImageGenerationRequest(t *testing.T) { + tests := []struct { + name string + size string + wantAspect string + wantErrPart string + }{ + {name: "square", size: "1024x1024", wantAspect: "1:1"}, + {name: "wide", size: "1792x1024", wantAspect: "16:9"}, + {name: "portrait", size: "1024x1792", wantAspect: "9:16"}, + {name: "unknown", size: "333x777", wantAspect: "auto"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := []byte(`{"model":"nano-banana","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) + got, err := ImageGenerationRequest(body, "fallback-model") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "draw a cat", req.Prompt) + assert.Equal(t, "nano-banana", req.Model) + assert.Equal(t, tt.wantAspect, req.AspectRatio) + assert.Empty(t, req.ImageURLs) + }) + } +} + +func TestImageGenerationRequestRejectsMultipleImages(t *testing.T) { + _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","n":2}`), "nano-banana") + require.Error(t, err) + assert.Contains(t, err.Error(), "n=1") +} + +func TestImageEditRequest(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "nano-banana", + "prompt": "make it blue", + "size": "1024x1024", + "n": "1", + }, map[string][]byte{ + "image": pngBytes(), + }) + + got, err := ImageEditRequest(body, contentType, "fallback-model") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "make it blue", req.Prompt) + assert.Equal(t, "nano-banana", req.Model) + assert.Equal(t, "1:1", req.AspectRatio) + require.Len(t, req.ImageURLs, 1) + assert.True(t, strings.HasPrefix(req.ImageURLs[0], "data:image/png;base64,")) + assert.Contains(t, req.ImageURLs[0], base64.StdEncoding.EncodeToString(pngBytes())) +} + +func TestImageEditRequestRejectsMask(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "nano-banana", + "prompt": "make it blue", + }, map[string][]byte{ + "image": pngBytes(), + "mask": pngBytes(), + }) + + _, err := ImageEditRequest(body, contentType, "fallback-model") + require.Error(t, err) + assert.Contains(t, err.Error(), "mask") +} + +func TestImageEditRequestRejectsMultipleImagesCount(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "nano-banana", + "prompt": "make it blue", + "n": "2", + }, map[string][]byte{ + "image": pngBytes(), + }) + + _, err := ImageEditRequest(body, contentType, "fallback-model") + require.Error(t, err) + assert.Contains(t, err.Error(), "n=1") +} + +func TestOpenAIImageResponse(t *testing.T) { + url := "https://cdn.sosana.art/result.png" + body, err := OpenAIImageResponse(BananaTaskResponse{ + Status: StatusCompleted, + ResultFileURL: &url, + }) + require.NoError(t, err) + + var resp openai.OpenAIImageResponse + require.NoError(t, json.Unmarshal(body, &resp)) + require.Len(t, resp.Data, 1) + assert.Equal(t, url, resp.Data[0].URL) + assert.Empty(t, resp.Data[0].B64JSON) + assert.NotZero(t, resp.Created) +} + +func multipartImageEditBody(t *testing.T, fields map[string]string, files map[string][]byte) ([]byte, string) { + t.Helper() + + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + for key, value := range fields { + require.NoError(t, writer.WriteField(key, value)) + } + for key, data := range files { + part, err := writer.CreateFormFile(key, key+".png") + require.NoError(t, err) + _, err = part.Write(data) + require.NoError(t, err) + } + require.NoError(t, writer.Close()) + return buf.Bytes(), writer.FormDataContentType() +} + +func pngBytes() []byte { + return []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} +} diff --git a/internal/litellmdb/model_table/model_table.go b/internal/litellmdb/model_table/model_table.go index f19e8f26..406c57ce 100644 --- a/internal/litellmdb/model_table/model_table.go +++ b/internal/litellmdb/model_table/model_table.go @@ -297,6 +297,8 @@ func mapProviderType(provider string) config.ProviderType { return config.ProviderTypeGemini case strings.Contains(p, "cometapi") || strings.Contains(p, "comet-api"): return config.ProviderTypeCometAPI + case strings.Contains(p, "sosana"): + return config.ProviderTypeSosana case strings.Contains(p, "xai"): return config.ProviderTypeOpenAI default: @@ -388,7 +390,8 @@ func convertPricingToModelPrice(p *queries.CustomPricingLiteLLMParams) *manager. if p == nil { return nil } - if p.InputCostPerToken == nil && p.OutputCostPerToken == nil { + if p.InputCostPerToken == nil && p.OutputCostPerToken == nil && + p.OutputCostPerImage == nil && p.OutputCostPerImageToken == nil { return nil } diff --git a/internal/litellmdb/model_table/model_table_test.go b/internal/litellmdb/model_table/model_table_test.go index 81f29dea..9ed9b678 100644 --- a/internal/litellmdb/model_table/model_table_test.go +++ b/internal/litellmdb/model_table/model_table_test.go @@ -6,6 +6,7 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/litellmdb/queries" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestMapProviderType(t *testing.T) { @@ -20,6 +21,8 @@ func TestMapProviderType(t *testing.T) { {"google", "GoogleAI", config.ProviderTypeGemini}, {"cometapi", "cometapi", config.ProviderTypeCometAPI}, {"comet-api", "comet-api", config.ProviderTypeCometAPI}, + {"sosana", "sosana", config.ProviderTypeSosana}, + {"sosana-art", "sosana-art", config.ProviderTypeSosana}, {"xai", "xAI", config.ProviderTypeOpenAI}, {"unknown", "some-other", ""}, } @@ -151,6 +154,19 @@ func TestConvertPricingToModelPrice(t *testing.T) { assert.Nil(t, convertPricingToModelPrice(nil)) } +func TestConvertPricingToModelPrice_ImageOnly(t *testing.T) { + outputImage := 0.5 + + price := convertPricingToModelPrice(&queries.CustomPricingLiteLLMParams{ + OutputCostPerImage: &outputImage, + }) + + require.NotNil(t, price) + assert.Equal(t, outputImage, price.OutputCostPerImage) + assert.Equal(t, 0.0, price.InputCostPerToken) + assert.Equal(t, 0.0, price.OutputCostPerToken) +} + func TestConvertPricingToModelPrice_AllFields(t *testing.T) { input := 0.01 output := 0.02 diff --git a/internal/models/manager.go b/internal/models/manager.go index 0dd50b96..3c98f2fc 100644 --- a/internal/models/manager.go +++ b/internal/models/manager.go @@ -350,6 +350,7 @@ var providerPassthroughDefaults = map[config.ProviderType]bool{ config.ProviderTypeGemini: false, config.ProviderTypeAnthropic: false, config.ProviderTypeCometAPI: false, + config.ProviderTypeSosana: false, config.ProviderTypeBedrock: false, } @@ -1677,6 +1678,7 @@ var providerTypeLiteLLMPrefix = map[config.ProviderType]string{ config.ProviderTypeGemini: "gemini", config.ProviderTypeAnthropic: "anthropic", config.ProviderTypeCometAPI: "cometapi", + config.ProviderTypeSosana: "sosana", config.ProviderTypeBedrock: "bedrock", config.ProviderTypeProxy: "openai", } diff --git a/internal/models/price_calculator_test.go b/internal/models/price_calculator_test.go index f9411618..bd5a5490 100644 --- a/internal/models/price_calculator_test.go +++ b/internal/models/price_calculator_test.go @@ -234,6 +234,36 @@ func TestCalculateTokenCosts_NilPrice(t *testing.T) { assert.Nil(t, costs) } +func TestCalculateTokenCosts_ImageCount(t *testing.T) { + usage := &converter.TokenUsage{ + ImageCount: 2, + } + price := &ModelPrice{ + OutputCostPerImage: 0.05, + } + + costs := CalculateTokenCosts(usage, price) + + assert.NotNil(t, costs) + assert.InDelta(t, 0.10, costs.ImageCost, 1e-9) + assert.InDelta(t, 0.10, costs.TotalCost, 1e-9) +} + +func TestCalculateTokenCosts_ImageCountUsesImageTokenFallback(t *testing.T) { + usage := &converter.TokenUsage{ + ImageCount: 3, + } + price := &ModelPrice{ + OutputCostPerImageToken: 0.02, + } + + costs := CalculateTokenCosts(usage, price) + + assert.NotNil(t, costs) + assert.InDelta(t, 0.06, costs.ImageCost, 1e-9) + assert.InDelta(t, 0.06, costs.TotalCost, 1e-9) +} + func TestModelPrice_CalculateCost(t *testing.T) { usage := &converter.TokenUsage{ PromptTokens: 100, diff --git a/internal/proxy/cometapi_test.go b/internal/proxy/cometapi_test.go index 38a7b0eb..a2028cfc 100644 --- a/internal/proxy/cometapi_test.go +++ b/internal/proxy/cometapi_test.go @@ -10,41 +10,106 @@ import ( func TestIsCometAPICredential(t *testing.T) { tests := []struct { - name string - cred *config.CredentialConfig - want bool + name string + cred *config.CredentialConfig + wantComet bool + wantMask bool }{ { - name: "dedicated provider type", - cred: &config.CredentialConfig{Type: config.ProviderTypeCometAPI}, - want: true, + name: "dedicated provider type", + cred: &config.CredentialConfig{Type: config.ProviderTypeCometAPI}, + wantComet: true, + wantMask: true, }, { - name: "comet host fallback", - cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, BaseURL: "https://api.cometapi.com/v1"}, - want: true, + name: "comet host fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, BaseURL: "https://api.cometapi.com/v1"}, + wantComet: true, + wantMask: true, }, { - name: "comet name fallback", - cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, Name: "comet-api-anthropic"}, - want: true, + name: "comet name fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, Name: "comet-api-anthropic"}, + wantComet: true, + wantMask: true, }, { - name: "regular anthropic", - cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, BaseURL: "https://api.anthropic.com"}, - want: false, + name: "regular anthropic", + cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, BaseURL: "https://api.anthropic.com"}, + wantComet: false, + wantMask: false, }, { - name: "nil credential", - cred: nil, - want: false, + name: "nil credential", + cred: nil, + wantComet: false, + wantMask: false, + }, + { + name: "explicit mask flag", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, MaskUpstreamErrors: true}, + wantComet: false, + wantMask: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.wantComet, isCometAPICredential(tt.cred)) + assert.Equal(t, tt.wantMask, shouldMaskUpstreamErrors(tt.cred)) + }) + } +} + +func TestIsSosanaCredential(t *testing.T) { + tests := []struct { + name string + cred *config.CredentialConfig + wantSosana bool + wantMask bool + }{ + { + name: "dedicated provider type", + cred: &config.CredentialConfig{Type: config.ProviderTypeSosana}, + wantSosana: true, + wantMask: true, + }, + { + name: "sosana host fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://sosana.art"}, + wantSosana: true, + wantMask: true, + }, + { + name: "sosana name fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, Name: "sosana-art-images"}, + wantSosana: true, + wantMask: true, + }, + { + name: "sasana host fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://api.sasana.example/v1"}, + wantSosana: true, + wantMask: true, + }, + { + name: "regular openai", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://api.openai.com"}, + wantSosana: false, + wantMask: false, + }, + { + name: "nil credential", + cred: nil, + wantSosana: false, + wantMask: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - assert.Equal(t, tt.want, isCometAPICredential(tt.cred)) - assert.Equal(t, tt.want, shouldMaskUpstreamErrors(tt.cred)) + assert.Equal(t, tt.wantSosana, isSosanaCredential(tt.cred)) + assert.Equal(t, tt.wantMask, shouldMaskUpstreamErrors(tt.cred)) }) } } diff --git a/internal/proxy/errors.go b/internal/proxy/errors.go index b86c6841..bee8e92f 100644 --- a/internal/proxy/errors.go +++ b/internal/proxy/errors.go @@ -87,6 +87,23 @@ func maskedUpstreamErrorBody(statusCode int) []byte { return append(body, '\n') } +func maskedContentPolicyBody() []byte { + code := "content_policy_violation" + resp := APIErrorResponse{ + Error: APIError{ + Message: "Content policy violation", + Type: errorTypeForStatus(http.StatusBadRequest), + Param: nil, + Code: &code, + }, + } + body, err := json.Marshal(resp) + if err != nil { + return []byte(`{"error":{"message":"Content policy violation","type":"invalid_request_error","param":null,"code":"content_policy_violation"}}`) + } + return append(body, '\n') +} + // WriteErrorBadRequest writes a 400 Bad Request JSON error. func WriteErrorBadRequest(w http.ResponseWriter, message string) { WriteJSONError(w, http.StatusBadRequest, message, errorTypeForStatus(http.StatusBadRequest), nil, nil) diff --git a/internal/proxy/image_response.go b/internal/proxy/image_response.go new file mode 100644 index 00000000..0df4f682 --- /dev/null +++ b/internal/proxy/image_response.go @@ -0,0 +1,102 @@ +package proxy + +import ( + "bytes" + "encoding/json" + "fmt" + "time" + + "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" + "github.com/mixaill76/auto_ai_router/internal/converter/openai" +) + +type sosanaTaskResponse struct { + Status string `json:"status"` + CreatedAt string `json:"created_at"` + ResultFileURL string `json:"result_file_url"` + OptimizedPrompt string `json:"optimized_prompt"` + Error json.RawMessage `json:"error"` +} + +func normalizeOpenAIImageResponseBody(body []byte) ([]byte, error) { + var envelope struct { + Error json.RawMessage `json:"error"` + } + if err := json.Unmarshal(body, &envelope); err != nil { + return nil, fmt.Errorf("parse image response: %w", err) + } + + var resp openai.OpenAIImageResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("parse OpenAI image response: %w", err) + } + if len(resp.Data) == 0 { + if sosanaResp, ok, err := openAIImageResponseFromSosanaTask(body); err != nil { + return nil, err + } else if ok { + resp = sosanaResp + } else { + if hasJSONValue(envelope.Error) { + return nil, fmt.Errorf("image response contains upstream error envelope") + } + return nil, fmt.Errorf("image response contains no image data") + } + } + for i, item := range resp.Data { + if item.B64JSON == "" && item.URL == "" { + return nil, fmt.Errorf("image response item %d contains neither b64_json nor url", i) + } + } + if resp.Created == 0 { + resp.Created = converterutil.GetCurrentTimestamp() + } + + normalized, err := json.Marshal(resp) + if err != nil { + return nil, fmt.Errorf("marshal OpenAI image response: %w", err) + } + return append(normalized, '\n'), nil +} + +func openAIImageResponseFromSosanaTask(body []byte) (openai.OpenAIImageResponse, bool, error) { + var task sosanaTaskResponse + if err := json.Unmarshal(body, &task); err != nil { + return openai.OpenAIImageResponse{}, false, nil + } + + switch task.Status { + case "": + return openai.OpenAIImageResponse{}, false, nil + case "COMPLETED": + if task.ResultFileURL == "" { + return openai.OpenAIImageResponse{}, true, fmt.Errorf("image task completed response contains no result URL") + } + resp := openai.OpenAIImageResponse{ + Created: parseSosanaCreatedAt(task.CreatedAt), + Data: []openai.OpenAIImageData{{ + URL: task.ResultFileURL, + RevisedPrompt: task.OptimizedPrompt, + }}, + } + return resp, true, nil + case "PROCESSING", "FAILED", "MODERATED": + return openai.OpenAIImageResponse{}, true, fmt.Errorf("image task response is not a completed image result") + default: + return openai.OpenAIImageResponse{}, true, fmt.Errorf("image task response contains an unknown status") + } +} + +func parseSosanaCreatedAt(value string) int64 { + if value == "" { + return converterutil.GetCurrentTimestamp() + } + if ts, err := time.Parse(time.RFC3339Nano, value); err == nil { + return ts.Unix() + } + return converterutil.GetCurrentTimestamp() +} + +func hasJSONValue(raw json.RawMessage) bool { + raw = bytes.TrimSpace(raw) + return len(raw) > 0 && !bytes.Equal(raw, []byte("null")) +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 77a8747e..e5e227f2 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -510,6 +510,11 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { } } + if cred.Type == config.ProviderTypeSosana { + p.handleSosanaRequest(w, r, body, cred, modelID, realModelID, isImageGeneration, isImageEdit, logCtx, start) + return + } + // Handle proxy credential type with same-type retry + fallback if cred.Type == config.ProviderTypeProxy { logCtx.Credential = cred @@ -632,9 +637,34 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { } // Write response (streaming or non-streaming) + maskProxyError := proxyResp.StatusCode >= 400 && shouldMaskProxyResponseErrors(cred, proxyResp) if proxyResp.IsStreaming { p.logger.DebugContext(r.Context(), "Response is streaming (no retry for streaming)", "credential", cred.Name, "status", proxyResp.StatusCode) + if maskProxyError { + if proxyResp.StreamBody != nil { + _ = proxyResp.StreamBody.Close() + } + body := maskedUpstreamErrorBody(proxyResp.StatusCode) + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body))) + if logCtx.IsProxyRequest && logCtx.ActualCredentialName != "" { + w.Header().Set("X-Credential-Name", logCtx.ActualCredentialName) + } + w.WriteHeader(proxyResp.StatusCode) + _, _ = w.Write(body) + logCtx.Status = "failure" + logCtx.HTTPStatus = proxyResp.StatusCode + logCtx.ErrorMsg = "Upstream provider error" + logCtx.TargetURL = cred.BaseURL + p.logUpstreamError(r.Context(), "Proxy request completed with masked streaming error status", proxyResp.StatusCode, cred, modelID, nil, + "url", cred.BaseURL, + "streaming", true, + "actual_credential", logCtx.ActualCredentialName, + "response_body_masked", true, + "request_id", logCtx.RequestID) + return + } streamCompleted := false if prepared.convertedResp { @@ -733,6 +763,28 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { p.setSessionBinding(logCtx.SessionID, modelID, cred.Name) } } else { + if proxyResp.StatusCode >= 200 && proxyResp.StatusCode < 300 && logCtx.IsImageGeneration && shouldMaskProxyResponseErrors(cred, proxyResp) { + normalizedBody, normErr := normalizeOpenAIImageResponseBody(proxyResp.Body) + if normErr != nil { + p.logger.ErrorContext(r.Context(), "Failed to normalize proxy image response to OpenAI format", + "credential", cred.Name, "model", modelID, "error", normErr, + "actual_credential", logCtx.ActualCredentialName, + "response_body_masked", true, + "request_id", logCtx.RequestID) + proxyResp.StatusCode = http.StatusBadGateway + maskProxyErrorResponse(proxyResp) + } else { + proxyResp.Body = normalizedBody + if proxyResp.Headers == nil { + proxyResp.Headers = http.Header{} + } + proxyResp.Headers.Set("Content-Type", "application/json") + } + } + if proxyResp.StatusCode >= 400 && shouldMaskProxyResponseErrors(cred, proxyResp) { + maskProxyErrorResponse(proxyResp) + } + // Save passthrough Responses API response or convert Chat Completions response if needed if prepared.passthroughResponses && proxyResp.StatusCode >= 200 && proxyResp.StatusCode < 300 { // Codex passthrough: body is already in Responses API format — just enrich and save. @@ -801,11 +853,18 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { // Final error returned to the client — single unified ERROR record. // For streaming responses the body was forwarded to the client and is // not available here (response_body is omitted). - p.logUpstreamError(r.Context(), "Proxy request completed with error status", proxyResp.StatusCode, cred, modelID, proxyResp.Body, + proxyErrorBody := proxyResp.Body + extra := []any{ "url", cred.BaseURL, "streaming", proxyResp.IsStreaming, "actual_credential", logCtx.ActualCredentialName, - "request_id", logCtx.RequestID) + "request_id", logCtx.RequestID, + } + if shouldMaskProxyResponseErrors(cred, proxyResp) { + proxyErrorBody = nil + extra = append(extra, "response_body_masked", true) + } + p.logUpstreamError(r.Context(), "Proxy request completed with error status", proxyResp.StatusCode, cred, modelID, proxyErrorBody, extra...) } logCtx.HTTPStatus = proxyResp.StatusCode logCtx.TargetURL = cred.BaseURL @@ -1288,7 +1347,13 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { } args = appendResponseBodyForLogs(args, cred, decodedBody) p.logger.ErrorContext(r.Context(), "Failed to transform provider response to OpenAI format", args...) - finalResponseBody = []byte(decodedBody) + if shouldMaskUpstreamErrors(cred) { + resp.StatusCode = http.StatusBadGateway + finalResponseBody = maskedUpstreamErrorBody(resp.StatusCode) + resp.Header.Set("Content-Type", "application/json") + } else { + finalResponseBody = []byte(decodedBody) + } } else { finalResponseBody = convertedBody p.logTransformedResponse(r.Context(), cred.Name, string(cred.Type), finalResponseBody) @@ -1317,7 +1382,12 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { } args = appendResponseBodyForLogs(args, cred, decodedBody) p.logger.ErrorContext(r.Context(), "Failed to convert native Responses API response", args...) - // finalResponseBody already holds decodedBody — return as-is + if shouldMaskUpstreamErrors(cred) { + resp.StatusCode = http.StatusBadGateway + finalResponseBody = maskedUpstreamErrorBody(resp.StatusCode) + bodyForTokenExtraction = finalResponseBody + resp.Header.Set("Content-Type", "application/json") + } } else { applyResponsesMetadata(nativeResp, prepared.responsesMetadata) if enriched, marshalErr := json.Marshal(nativeResp); marshalErr == nil { @@ -1350,7 +1420,12 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { p.logger.ErrorContext(r.Context(), "Failed to convert to Responses API format", "credential", cred.Name, "model", modelID, "error", convErr, "request_id", logCtx.RequestID) - // fallback: use Chat Completions body + if shouldMaskUpstreamErrors(cred) { + resp.StatusCode = http.StatusBadGateway + finalResponseBody = maskedUpstreamErrorBody(resp.StatusCode) + bodyForTokenExtraction = finalResponseBody + resp.Header.Set("Content-Type", "application/json") + } } else { // Enrich the response with request-echoed fields (store, previous_response_id, // metadata) for both the client payload and the store record. @@ -1373,6 +1448,25 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { } } + if resp.StatusCode >= 200 && resp.StatusCode < 300 && logCtx.IsImageGeneration && shouldMaskUpstreamErrors(cred) { + normalizedBody, normErr := normalizeOpenAIImageResponseBody(finalResponseBody) + if normErr != nil { + args := []any{ + "credential", cred.Name, "provider", string(cred.Type), + "model", modelID, "error", normErr, + "request_id", logCtx.RequestID, + } + args = appendResponseBodyForLogs(args, cred, string(finalResponseBody)) + p.logger.ErrorContext(r.Context(), "Failed to normalize image response to OpenAI format", args...) + resp.StatusCode = http.StatusBadGateway + finalResponseBody = maskedUpstreamErrorBody(resp.StatusCode) + } else { + finalResponseBody = normalizedBody + } + bodyForTokenExtraction = finalResponseBody + resp.Header.Set("Content-Type", "application/json") + } + rawErrorBody := finalResponseBody if resp.StatusCode >= 400 && shouldMaskUpstreamErrors(cred) { finalResponseBody = maskedUpstreamErrorBody(resp.StatusCode) diff --git a/internal/proxy/proxy_log.go b/internal/proxy/proxy_log.go index 5cdaf020..ac5228c1 100644 --- a/internal/proxy/proxy_log.go +++ b/internal/proxy/proxy_log.go @@ -49,7 +49,10 @@ func appendResponseBodyForLogs(args []any, cred *config.CredentialConfig, body s } func shouldMaskUpstreamErrors(cred *config.CredentialConfig) bool { - return isCometAPICredential(cred) + if cred == nil { + return false + } + return cred.MaskUpstreamErrors || isCometAPICredential(cred) || isSosanaCredential(cred) } func isCometAPICredential(cred *config.CredentialConfig) bool { @@ -66,19 +69,42 @@ func isCometAPICredential(cred *config.CredentialConfig) bool { } func isCometAPIHost(rawBaseURL string) bool { + host := normalizedHost(rawBaseURL) + return host == "cometapi.com" || strings.HasSuffix(host, ".cometapi.com") +} + +func normalizedHost(rawBaseURL string) string { baseURL := strings.TrimSpace(rawBaseURL) if baseURL == "" { - return false + return "" } u, err := url.Parse(baseURL) if err != nil || u.Hostname() == "" { u, err = url.Parse("https://" + baseURL) if err != nil { - return false + return "" } } - host := strings.TrimSuffix(strings.ToLower(u.Hostname()), ".") - return host == "cometapi.com" || strings.HasSuffix(host, ".cometapi.com") + return strings.TrimSuffix(strings.ToLower(u.Hostname()), ".") +} + +func isSosanaCredential(cred *config.CredentialConfig) bool { + if cred == nil { + return false + } + if cred.Type == config.ProviderTypeSosana { + return true + } + name := strings.ToLower(cred.Name) + host := normalizedHost(cred.BaseURL) + return isSosanaHost(cred.BaseURL) || + containsSosanaMarker(name) || + containsSosanaMarker(host) +} + +func isSosanaHost(rawBaseURL string) bool { + host := normalizedHost(rawBaseURL) + return host == "sosana.art" || strings.HasSuffix(host, ".sosana.art") } // logStreamHandlerError logs a streaming handler failure. Client disconnects are diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go new file mode 100644 index 00000000..c9e0045e --- /dev/null +++ b/internal/proxy/sosana.go @@ -0,0 +1,296 @@ +package proxy + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "math/rand" + "net/http" + "time" + + "github.com/mixaill76/auto_ai_router/internal/config" + "github.com/mixaill76/auto_ai_router/internal/converter" + "github.com/mixaill76/auto_ai_router/internal/converter/sosana" +) + +const sosanaPollInterval = 2 * time.Second + +type sosanaAttemptResult struct { + body []byte + statusCode int + retryable bool + retryReason RetryReason +} + +func (p *Proxy) handleSosanaRequest( + w http.ResponseWriter, + r *http.Request, + body []byte, + cred *config.CredentialConfig, + modelID string, + realModelID string, + isImageGeneration bool, + isImageEdit bool, + logCtx *RequestLogContext, + start time.Time, +) { + logCtx.Credential = cred + logCtx.TargetURL = cred.BaseURL + + if !isImageGeneration && !isImageEdit { + message := "sosana provider supports only image generation" + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = message + WriteErrorBadRequest(w, message) + return + } + + createBody, err := p.buildSosanaCreateBody(body, r.Header.Get("Content-Type"), realModelID, isImageEdit) + if err != nil { + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = err.Error() + WriteErrorBadRequest(w, err.Error()) + return + } + + ctx := r.Context() + var cancel context.CancelFunc + if p.requestTimeout > 0 { + ctx, cancel = context.WithTimeout(ctx, p.requestTimeout) + defer cancel() + } + + result := sosanaAttemptResult{statusCode: http.StatusBadGateway, body: maskedUpstreamErrorBody(http.StatusBadGateway)} + triedCreds := GetTried(r.Context()) + for attempt := 0; attempt <= p.maxProviderRetries; attempt++ { + if attempt > 0 { + nextCred, err := p.balancer.NextSameTypeForModelExcluding(modelID, config.ProviderTypeSosana, triedCreds) + if err != nil { + p.logger.DebugContext(r.Context(), "No more Sosana credentials for retry", + "model", modelID, "attempt", attempt, "error", err) + break + } + cred = nextCred + triedCreds[cred.Name] = true + logCtx.Credential = cred + logCtx.TargetURL = cred.BaseURL + + p.logger.InfoContext(r.Context(), "Retrying Sosana create request with next credential", + "credential", cred.Name, "model", modelID, + "attempt", attempt+1, "max_attempts", p.maxProviderRetries+1, + "retry_reason", result.retryReason) + time.Sleep(time.Duration(rand.Intn(50)) * time.Millisecond) + } + + result = p.createAndPollSosanaTask(ctx, cred, modelID, createBody, logCtx) + p.balancer.RecordResponse(cred.Name, modelID, result.statusCode) + p.metrics.RecordRequest(cred.Name, r.URL.Path, modelID, result.statusCode, time.Since(start)) + if !result.retryable { + break + } + + p.logger.WarnContext(r.Context(), "Sosana create request returned retryable error, will retry", + "error_code", result.statusCode, + "credential", cred.Name, + "reason", result.retryReason, + "model", modelID, + "attempt", attempt+1, + "max_attempts", p.maxProviderRetries+1, + "response_body_masked", true) + } + + if result.statusCode >= 400 { + logCtx.Status = "failure" + logCtx.HTTPStatus = result.statusCode + logCtx.ErrorMsg = "Upstream provider error" + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(result.statusCode) + _, _ = w.Write(result.body) + return + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(result.body) + + logCtx.Status = "success" + logCtx.HTTPStatus = http.StatusOK + logCtx.TokenUsage = &converter.TokenUsage{ImageCount: logCtx.ImageCount} + logCtx.RequestCompleted = true + logCtx.Logged = true + if err := p.logSpendToLiteLLMDB(logCtx); err != nil { + p.logger.WarnContext(r.Context(), "Failed to queue spend log", + "error", err, + "request_id", logCtx.RequestID) + } +} + +func (p *Proxy) buildSosanaCreateBody(body []byte, contentType, realModelID string, isImageEdit bool) ([]byte, error) { + if isImageEdit { + return sosana.ImageEditRequest(body, contentType, realModelID) + } + return sosana.ImageGenerationRequest(body, realModelID) +} + +func (p *Proxy) createAndPollSosanaTask( + ctx context.Context, + cred *config.CredentialConfig, + modelID string, + createBody []byte, + logCtx *RequestLogContext, +) sosanaAttemptResult { + task, rawBody, statusCode, err := p.doSosanaTaskRequest(ctx, http.MethodPost, sosana.CreateURL(cred.BaseURL), cred, createBody) + if err != nil { + body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosana.CreateURL(cred.BaseURL), logCtx) + return sosanaAttemptResult{ + body: body, + statusCode: code, + retryable: ctx.Err() == nil, + retryReason: RetryReasonNetErr, + } + } + if statusCode >= 400 { + p.logUpstreamError(ctx, "Sosana create request completed with error status", statusCode, cred, modelID, rawBody, + "url", sosana.CreateURL(cred.BaseURL), + "request_id", logCtx.RequestID) + retryable, reason := ShouldRetryWithFallback(statusCode, rawBody) + return sosanaAttemptResult{ + body: maskedUpstreamErrorBody(statusCode), + statusCode: statusCode, + retryable: retryable, + retryReason: reason, + } + } + + immediatePoll := true + for { + body, code, done := p.sosanaTaskBody(ctx, cred, modelID, task, rawBody, statusCode, logCtx) + if done { + return sosanaAttemptResult{body: body, statusCode: code} + } + + if !immediatePoll { + select { + case <-ctx.Done(): + p.logUpstreamError(context.Background(), "Sosana task polling timed out", http.StatusRequestTimeout, cred, modelID, rawBody, + "url", sosana.PollURL(cred.BaseURL, task.UID), + "request_id", logCtx.RequestID, + "error", ctx.Err()) + return sosanaAttemptResult{body: maskedUpstreamErrorBody(http.StatusRequestTimeout), statusCode: http.StatusRequestTimeout} + case <-time.After(sosanaPollInterval): + } + } + immediatePoll = false + + task, rawBody, statusCode, err = p.doSosanaTaskRequest(ctx, http.MethodGet, sosana.PollURL(cred.BaseURL, task.UID), cred, nil) + if err != nil { + body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosana.PollURL(cred.BaseURL, task.UID), logCtx) + return sosanaAttemptResult{body: body, statusCode: code} + } + if statusCode >= 400 { + p.logUpstreamError(ctx, "Sosana poll request completed with error status", statusCode, cred, modelID, rawBody, + "url", sosana.PollURL(cred.BaseURL, task.UID), + "request_id", logCtx.RequestID) + return sosanaAttemptResult{body: maskedUpstreamErrorBody(statusCode), statusCode: statusCode} + } + } +} + +func (p *Proxy) sosanaTaskBody( + ctx context.Context, + cred *config.CredentialConfig, + modelID string, + task sosana.BananaTaskResponse, + rawBody []byte, + statusCode int, + logCtx *RequestLogContext, +) ([]byte, int, bool) { + switch task.Status { + case sosana.StatusCompleted: + body, err := sosana.OpenAIImageResponse(task) + if err != nil { + p.logUpstreamError(ctx, "Sosana completed task missing result", http.StatusBadGateway, cred, modelID, rawBody, + "url", sosana.PollURL(cred.BaseURL, task.UID), + "request_id", logCtx.RequestID, + "error", err) + return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true + } + return body, http.StatusOK, true + case sosana.StatusFailed: + p.logUpstreamError(ctx, "Sosana task failed", http.StatusBadGateway, cred, modelID, rawBody, + "url", sosana.PollURL(cred.BaseURL, task.UID), + "request_id", logCtx.RequestID) + return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true + case sosana.StatusModerated: + p.logUpstreamError(ctx, "Sosana task moderated", http.StatusBadRequest, cred, modelID, rawBody, + "url", sosana.PollURL(cred.BaseURL, task.UID), + "request_id", logCtx.RequestID) + return maskedContentPolicyBody(), http.StatusBadRequest, true + case sosana.StatusProcessing: + if task.UID == "" { + p.logUpstreamError(ctx, "Sosana processing task missing uid", http.StatusBadGateway, cred, modelID, rawBody, + "url", cred.BaseURL, + "request_id", logCtx.RequestID) + return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true + } + return nil, statusCode, false + default: + p.logUpstreamError(ctx, "Sosana task returned unknown status", http.StatusBadGateway, cred, modelID, rawBody, + "url", sosana.PollURL(cred.BaseURL, task.UID), + "request_id", logCtx.RequestID, + "status", task.Status) + return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true + } +} + +func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosana.BananaTaskResponse, []byte, int, error) { + var reader *bytes.Reader + if body != nil { + reader = bytes.NewReader(body) + } else { + reader = bytes.NewReader(nil) + } + req, err := http.NewRequestWithContext(ctx, method, url, reader) + if err != nil { + return sosana.BananaTaskResponse{}, nil, http.StatusInternalServerError, err + } + req.Header.Set("Authorization", "Bearer "+cred.APIKey) + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + + resp, err := p.client.Do(req) + if err != nil { + return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err + } + defer func() { + if closeErr := resp.Body.Close(); closeErr != nil { + p.logger.WarnContext(ctx, "Failed to close Sosana response body", "error", closeErr) + } + }() + + rawBody, err := p.readLimitedResponseBody(resp.Body) + if err != nil { + return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err + } + var task sosana.BananaTaskResponse + if len(rawBody) > 0 { + _ = json.Unmarshal(rawBody, &task) + } + return task, rawBody, resp.StatusCode, nil +} + +func (p *Proxy) sosanaTransportError(ctx context.Context, err error, cred *config.CredentialConfig, modelID, url string, logCtx *RequestLogContext) ([]byte, int) { + statusCode := http.StatusBadGateway + if isTimeoutError(err) || errors.Is(ctx.Err(), context.DeadlineExceeded) { + statusCode = http.StatusRequestTimeout + } + p.logUpstreamError(context.Background(), "Sosana upstream request failed", statusCode, cred, modelID, nil, + "url", url, + "request_id", logCtx.RequestID, + "error", err) + return maskedUpstreamErrorBody(statusCode), statusCode +} diff --git a/internal/proxy/sosana_live_test.go b/internal/proxy/sosana_live_test.go new file mode 100644 index 00000000..32a52de3 --- /dev/null +++ b/internal/proxy/sosana_live_test.go @@ -0,0 +1,63 @@ +package proxy + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/mixaill76/auto_ai_router/internal/config" + "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/stretchr/testify/require" +) + +func TestProxyRequest_SosanaLiveAcceptance(t *testing.T) { + if os.Getenv("SOSANA_ACCEPTANCE") != "1" { + t.Skip("SOSANA_ACCEPTANCE=1 not set, skipping paid Sosana live acceptance test") + } + + apiKey := os.Getenv("SOSANA_API_KEY") + require.NotEmpty(t, apiKey, "SOSANA_API_KEY is required for Sosana live acceptance test") + + baseURL := os.Getenv("SOSANA_BASE_URL") + if baseURL == "" { + baseURL = "https://sosana.art" + } + model := os.Getenv("SOSANA_MODEL") + if model == "" { + model = "nano-banana" + } + prompt := os.Getenv("SOSANA_PROMPT") + if prompt == "" { + prompt = "a small blue cube on a white background" + } + + prx := NewTestProxyBuilder(). + WithSingleCredential("sosana-live", config.ProviderTypeSosana, baseURL, apiKey). + WithRequestTimeout(2 * time.Minute). + Build() + + body, err := json.Marshal(map[string]any{ + "model": model, + "prompt": prompt, + "size": "1024x1024", + "n": 1, + }) + require.NoError(t, err) + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(string(body))) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code, "response body: %s", w.Body.String()) + var resp openai.OpenAIImageResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp.Data, 1) + require.NotEmpty(t, resp.Data[0].URL) +} diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go new file mode 100644 index 00000000..27f53b26 --- /dev/null +++ b/internal/proxy/sosana_test.go @@ -0,0 +1,304 @@ +package proxy + +import ( + "bytes" + "encoding/json" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/mixaill76/auto_ai_router/internal/config" + "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { + var createSeen, pollSeen bool + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "Bearer sosana-key", r.Header.Get("Authorization")) + switch r.URL.Path { + case "/api/banana/create-async": + createSeen = true + var req map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + assert.Equal(t, "draw a fox", req["prompt"]) + assert.Equal(t, "nano-banana", req["model"]) + assert.Equal(t, "1:1", req["aspect_ratio"]) + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox"}`)) + case "/api/banana/task-1": + pollSeen = true + _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox","result_file_url":"https://cdn.sosana.art/fox.png"}`)) + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, nil) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw a fox","size":"1024x1024","n":1,"response_format":"b64_json"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + var resp openai.OpenAIImageResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp.Data, 1) + assert.Equal(t, "https://cdn.sosana.art/fox.png", resp.Data[0].URL) + assert.Empty(t, resp.Data[0].B64JSON) + assert.True(t, createSeen) + assert.True(t, pollSeen) +} + +func TestProxyRequest_SosanaRejectsNonImageEndpoint(t *testing.T) { + called := false + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, nil) + req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"nano-banana","messages":[]}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), "provider supports only image generation") + assert.False(t, called) +} + +func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { + var logBuf bytes.Buffer + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusPaymentRequired) + _, _ = w.Write([]byte(`{"detail":"sosana balance secret marker"}`)) + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusPaymentRequired, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "balance secret") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.NotContains(t, logBuf.String(), "balance secret") +} + +func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { + var createAuths []string + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + createAuths = append(createAuths, r.Header.Get("Authorization")) + if r.Header.Get("Authorization") == "Bearer sosana-key-a" { + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"detail":"first credential rate limited"}`)) + return + } + assert.Equal(t, "Bearer sosana-key-b", r.Header.Get("Authorization")) + _, _ = w.Write([]byte(`{"uid":"task-2","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-2": + assert.Equal(t, "Bearer sosana-key-b", r.Header.Get("Authorization")) + _, _ = w.Write([]byte(`{"uid":"task-2","status":"COMPLETED","prompt":"draw","result_file_url":"https://cdn.sosana.art/retry.png"}`)) + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials( + config.CredentialConfig{Name: "sosana-a", Type: config.ProviderTypeSosana, BaseURL: upstream.URL, APIKey: "sosana-key-a", RPM: 100, TPM: 10000}, + config.CredentialConfig{Name: "sosana-b", Type: config.ProviderTypeSosana, BaseURL: upstream.URL, APIKey: "sosana-key-b", RPM: 100, TPM: 10000}, + ). + WithMaxProviderRetries(1). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, []string{"Bearer sosana-key-a", "Bearer sosana-key-b"}, createAuths) + assert.Contains(t, w.Body.String(), "https://cdn.sosana.art/retry.png") +} + +func TestProxyRequest_SosanaPollHTTPErrorMasked(t *testing.T) { + var logBuf bytes.Buffer + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw"}`)) + case "/api/banana/task-1": + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"detail":"poll secret marker"}`)) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusInternalServerError, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "poll secret") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.NotContains(t, logBuf.String(), "poll secret") +} + +func TestProxyRequest_SosanaDoesNotRetryAfterTaskCreated(t *testing.T) { + createCalls := 0 + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + createCalls++ + assert.Equal(t, "Bearer sosana-key-a", r.Header.Get("Authorization")) + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + assert.Equal(t, "Bearer sosana-key-a", r.Header.Get("Authorization")) + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"detail":"poll failed after task was created"}`)) + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials( + config.CredentialConfig{Name: "sosana-a", Type: config.ProviderTypeSosana, BaseURL: upstream.URL, APIKey: "sosana-key-a", RPM: 100, TPM: 10000}, + config.CredentialConfig{Name: "sosana-b", Type: config.ProviderTypeSosana, BaseURL: upstream.URL, APIKey: "sosana-key-b", RPM: 100, TPM: 10000}, + ). + WithMaxProviderRetries(1). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusInternalServerError, w.Code) + assert.Equal(t, 1, createCalls) + assert.NotContains(t, w.Body.String(), "poll failed") + assert.Contains(t, w.Body.String(), "Upstream provider error") +} + +func TestProxyRequest_SosanaTaskFailedMasked(t *testing.T) { + var logBuf bytes.Buffer + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"FAILED","created_at":"2026-01-01T00:00:00Z","prompt":"draw","error":"failed secret marker"}`)) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "failed secret") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.NotContains(t, logBuf.String(), "failed secret") +} + +func TestProxyRequest_SosanaTaskModeratedMasked(t *testing.T) { + var logBuf bytes.Buffer + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"MODERATED","created_at":"2026-01-01T00:00:00Z","prompt":"draw","error":"moderation secret marker"}`)) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), "content_policy_violation") + assert.NotContains(t, w.Body.String(), "moderation secret") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.NotContains(t, logBuf.String(), "moderation secret") +} + +func TestProxyRequest_SosanaTimeoutMasked(t *testing.T) { + var logBuf bytes.Buffer + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw"}`)) + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + prx.requestTimeout = 5 * time.Millisecond + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusRequestTimeout, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.Contains(t, logBuf.String(), "response_body_masked=true") +} + +func newSosanaTestProxy(baseURL string, logBuf *bytes.Buffer) *Proxy { + logger := slog.New(slog.NewTextHandler(&bytes.Buffer{}, &slog.HandlerOptions{Level: slog.LevelDebug})) + if logBuf != nil { + logger = slog.New(slog.NewTextHandler(logBuf, &slog.HandlerOptions{Level: slog.LevelDebug})) + } + return NewTestProxyBuilder(). + WithSingleCredential("sosana", config.ProviderTypeSosana, baseURL, "sosana-key"). + WithRequestTimeout(30 * time.Second). + withLogger(logger). + Build() +} + +func (b *TestProxyBuilder) withLogger(logger *slog.Logger) *TestProxyBuilder { + b.config.Logger = logger + b.config.TokenManager = createTestTokenManager(logger) + b.config.ModelManager = createTestModelManager(logger) + return b +} diff --git a/internal/proxy/upstream_masking.go b/internal/proxy/upstream_masking.go new file mode 100644 index 00000000..7b83ad9c --- /dev/null +++ b/internal/proxy/upstream_masking.go @@ -0,0 +1,36 @@ +package proxy + +import ( + "net/http" + "strings" + + "github.com/mixaill76/auto_ai_router/internal/config" +) + +func shouldMaskProxyResponseErrors(cred *config.CredentialConfig, resp *ProxyResponse) bool { + if shouldMaskUpstreamErrors(cred) { + return true + } + if resp == nil { + return false + } + return containsSosanaMarker(resp.ActualCredentialName) +} + +func maskProxyErrorResponse(resp *ProxyResponse) { + if resp == nil { + return + } + resp.Body = maskedUpstreamErrorBody(resp.StatusCode) + if resp.Headers == nil { + resp.Headers = http.Header{} + } + resp.Headers.Set("Content-Type", "application/json") + resp.Headers.Del("Content-Encoding") + resp.Headers.Del("Content-Length") +} + +func containsSosanaMarker(value string) bool { + value = strings.ToLower(value) + return strings.Contains(value, "sosana") || strings.Contains(value, "sasana") +} diff --git a/internal/proxy/upstream_masking_test.go b/internal/proxy/upstream_masking_test.go new file mode 100644 index 00000000..9a2a9aae --- /dev/null +++ b/internal/proxy/upstream_masking_test.go @@ -0,0 +1,296 @@ +package proxy + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/mixaill76/auto_ai_router/internal/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestMaskedUpstreamError_DirectImageErrorDoesNotLeakProviderBody(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"message":"Sasana quota exhausted","type":"sasana_rate_limit","code":"SASANA_429"},"provider":"Sasana"}`)) + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials(config.CredentialConfig{ + Name: "sosana-art", + Type: config.ProviderTypeOpenAI, + BaseURL: upstream.URL, + APIKey: "upstream-key", + RPM: 100, + TPM: 10000, + MaskUpstreamErrors: true, + }). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-1","prompt":"cat"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusTooManyRequests, w.Code) + assert.NotContains(t, w.Body.String(), "Sasana") + assert.NotContains(t, w.Body.String(), "SASANA_429") + + var got APIErrorResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) + assert.Equal(t, "Upstream provider error", got.Error.Message) + assert.Equal(t, "rate_limit_error", got.Error.Type) + require.NotNil(t, got.Error.Code) + assert.Equal(t, "upstream_rate_limit", *got.Error.Code) +} + +func TestMaskedUpstreamImageSuccess_IsNormalizedToOpenAIShape(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "created":123, + "provider":"Sasana", + "background":"transparent", + "output_format":"png", + "quality":"high", + "size":"1024x1024", + "error":{"message":"Sasana warning that should not leak"}, + "data":[{"b64_json":"aW1hZ2U=","sasana_id":"vendor-image-id"}] + }`)) + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials(config.CredentialConfig{ + Name: "sosana-art", + Type: config.ProviderTypeOpenAI, + BaseURL: upstream.URL, + APIKey: "upstream-key", + RPM: 100, + TPM: 10000, + MaskUpstreamErrors: true, + }). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-1","prompt":"cat"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + assert.NotContains(t, w.Body.String(), "Sasana") + assert.NotContains(t, w.Body.String(), "sasana_id") + assert.NotContains(t, w.Body.String(), "warning") + + var got struct { + Created int64 `json:"created"` + Background string `json:"background"` + OutputFormat string `json:"output_format"` + Quality string `json:"quality"` + Size string `json:"size"` + Data []struct { + B64JSON string `json:"b64_json"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) + assert.Equal(t, int64(123), got.Created) + assert.Equal(t, "transparent", got.Background) + assert.Equal(t, "png", got.OutputFormat) + assert.Equal(t, "high", got.Quality) + assert.Equal(t, "1024x1024", got.Size) + require.Len(t, got.Data, 1) + assert.Equal(t, "aW1hZ2U=", got.Data[0].B64JSON) +} + +func TestMaskedUpstreamImageSuccess_WithErrorEnvelopeBecomesRouterError(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"error":{"message":"Sasana moderation failed","code":"SASANA_POLICY"}}`)) + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials(config.CredentialConfig{ + Name: "sosana-art", + Type: config.ProviderTypeOpenAI, + BaseURL: upstream.URL, + APIKey: "upstream-key", + RPM: 100, + TPM: 10000, + MaskUpstreamErrors: true, + }). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-1","prompt":"cat"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusBadGateway, w.Code) + assert.NotContains(t, w.Body.String(), "Sasana") + assert.NotContains(t, w.Body.String(), "SASANA_POLICY") + assert.Contains(t, w.Body.String(), "Upstream provider error") +} + +func TestMaskedUpstreamSosanaCompletedResponse_IsConvertedToOpenAIShape(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "uid":"01977614-95b0-7000-8000-example", + "status":"COMPLETED", + "created_at":"2026-06-26T12:34:56Z", + "prompt":"cat", + "optimized_prompt":"A detailed cat illustration", + "result_file_url":"https://cdn.sosana.art/results/cat.png", + "elapsed":12.3, + "provider":"Sosana" + }`)) + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials(config.CredentialConfig{ + Name: "vsellm-sosana-art", + Type: config.ProviderTypeOpenAI, + BaseURL: upstream.URL, + APIKey: "upstream-key", + RPM: 100, + TPM: 10000, + MaskUpstreamErrors: true, + }). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-2","prompt":"cat"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + assert.NotContains(t, w.Body.String(), "Sosana") + assert.NotContains(t, w.Body.String(), "uid") + assert.NotContains(t, w.Body.String(), "result_file_url") + assert.NotContains(t, w.Body.String(), "status") + + var got struct { + Created int64 `json:"created"` + Data []struct { + URL string `json:"url"` + RevisedPrompt string `json:"revised_prompt"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) + assert.Equal(t, int64(1782477296), got.Created) + require.Len(t, got.Data, 1) + assert.Equal(t, "https://cdn.sosana.art/results/cat.png", got.Data[0].URL) + assert.Equal(t, "A detailed cat illustration", got.Data[0].RevisedPrompt) +} + +func TestMaskedUpstreamSosanaProcessingResponse_BecomesRouterError(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "uid":"01977614-95b0-7000-8000-example", + "status":"PROCESSING", + "created_at":"2026-06-26T12:34:56Z", + "prompt":"cat", + "provider":"Sosana" + }`)) + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials(config.CredentialConfig{ + Name: "vsellm-sosana-art", + Type: config.ProviderTypeOpenAI, + BaseURL: upstream.URL, + APIKey: "upstream-key", + RPM: 100, + TPM: 10000, + MaskUpstreamErrors: true, + }). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-2","prompt":"cat"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusBadGateway, w.Code) + assert.NotContains(t, w.Body.String(), "Sosana") + assert.NotContains(t, w.Body.String(), "PROCESSING") + assert.NotContains(t, w.Body.String(), "uid") + assert.Contains(t, w.Body.String(), "Upstream provider error") +} + +func TestMaskedUpstreamError_ProxyChainActualCredentialDoesNotLeakProviderBody(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Credential-Name", "sosana-art-primary") + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":{"message":"Sasana internal failure","code":"SASANA_500"}}`)) + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithSingleCredential("proxy-hop", config.ProviderTypeProxy, upstream.URL, "proxy-key"). + Build() + + req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusInternalServerError, w.Code) + assert.NotContains(t, w.Body.String(), "Sasana") + assert.NotContains(t, w.Body.String(), "SASANA_500") + assert.Contains(t, w.Body.String(), "Upstream provider error") +} + +func TestMaskedUpstreamError_ProxyChainStreamingErrorDoesNotLeakProviderBody(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("X-Credential-Name", "sosana-art-primary") + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte("data: {\"error\":\"Sasana stream failed\"}\n\n")) + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithSingleCredential("proxy-hop", config.ProviderTypeProxy, upstream.URL, "proxy-key"). + Build() + + req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"hi"}]}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusInternalServerError, w.Code) + assert.Equal(t, "application/json", w.Header().Get("Content-Type")) + assert.NotContains(t, w.Body.String(), "Sasana") + assert.Contains(t, w.Body.String(), "Upstream provider error") +} + +func TestNormalizeOpenAIImageResponseBodyRejectsInvalidSuccess(t *testing.T) { + _, err := normalizeOpenAIImageResponseBody([]byte(`{"created":123,"data":[{"revised_prompt":"only text"}]}`)) + require.Error(t, err) + assert.Contains(t, err.Error(), "neither b64_json nor url") +} From 6682d1c8a4067ef2d51e636a054fc5f9e3df93ff Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Fri, 26 Jun 2026 14:13:18 +0300 Subject: [PATCH 02/16] fix: harden sosana image provider --- config.yaml.example | 1 + docs/providers/sosana.md | 8 ++ internal/converter/sosana/images.go | 49 ++++++++---- internal/converter/sosana/images_test.go | 42 ++++++++-- internal/proxy/sosana.go | 7 +- internal/proxy/sosana_test.go | 98 +++++++++++++++++++++++- 6 files changed, 178 insertions(+), 27 deletions(-) diff --git a/config.yaml.example b/config.yaml.example index bf4f70ec..ef27cecb 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -79,6 +79,7 @@ credentials: tpm: -1 # Sosana.art image generation API (OpenAI Images-compatible through the router) + # Async image tasks can run longer than chat requests; use request_timeout/write_timeout >= 2m in production. - name: "sosana_images" type: "sosana" api_key: "os.environ/SOSANA_API_KEY" diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md index aeca2fa0..097db5a0 100644 --- a/docs/providers/sosana.md +++ b/docs/providers/sosana.md @@ -26,6 +26,10 @@ models: tpm: -1 ``` +Sosana Banana tasks are asynchronous and can take longer than short chat +completion requests. For production Sosana credentials, set the router +`request_timeout` and HTTP `write_timeout` to at least `2m`. + ## Behavior - `n` must be `1`. @@ -41,3 +45,7 @@ Sosana upstream HTTP errors and terminal task errors are masked before they are returned to clients or written to structured logs. The router preserves the appropriate HTTP status but replaces provider details with neutral OpenAI-compatible error bodies. + +If Sosana is hidden behind another proxy credential, enable +`mask_upstream_errors: true` on that proxy unless the upstream router is known to +propagate the credential marker used by this router. diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go index 899d0d88..3d28d3ec 100644 --- a/internal/converter/sosana/images.go +++ b/internal/converter/sosana/images.go @@ -11,6 +11,7 @@ import ( "net/http" "strconv" "strings" + "time" "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" "github.com/mixaill76/auto_ai_router/internal/converter/openai" @@ -33,11 +34,13 @@ type BananaCreateRequest struct { } type BananaTaskResponse struct { - UID string `json:"uid"` - Status string `json:"status"` - Prompt string `json:"prompt"` - ResultFileURL *string `json:"result_file_url"` - Error *string `json:"error"` + UID string `json:"uid"` + Status string `json:"status"` + Prompt string `json:"prompt"` + CreatedAt string `json:"created_at"` + OptimizedPrompt string `json:"optimized_prompt"` + ResultFileURL *string `json:"result_file_url"` + Error *string `json:"error"` } func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, error) { @@ -52,13 +55,9 @@ func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, error) { if prompt == "" { return nil, fmt.Errorf("image generation request missing prompt") } - model := strings.TrimSpace(req.Model) - if model == "" { - model = modelID - } return json.Marshal(BananaCreateRequest{ Prompt: prompt, - Model: model, + Model: providerModel(modelID, req.Model), AspectRatio: SizeToAspectRatio(req.Size), }) } @@ -117,17 +116,13 @@ func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([] if prompt == "" { return nil, fmt.Errorf("image edit request missing prompt field") } - model := strings.TrimSpace(fields["model"]) - if model == "" { - model = modelID - } if len(imageURLs) == 0 { return nil, fmt.Errorf("image edit request missing image") } return json.Marshal(BananaCreateRequest{ Prompt: prompt, ImageURLs: imageURLs, - Model: model, + Model: providerModel(modelID, fields["model"]), AspectRatio: SizeToAspectRatio(fields["size"]), }) } @@ -137,14 +132,34 @@ func OpenAIImageResponse(task BananaTaskResponse) ([]byte, error) { return nil, fmt.Errorf("sosana task completed without result_file_url") } resp := openai.OpenAIImageResponse{ - Created: converterutil.GetCurrentTimestamp(), + Created: createdAtUnix(task.CreatedAt), Data: []openai.OpenAIImageData{ - {URL: strings.TrimSpace(*task.ResultFileURL)}, + { + URL: strings.TrimSpace(*task.ResultFileURL), + RevisedPrompt: strings.TrimSpace(task.OptimizedPrompt), + }, }, } return json.Marshal(resp) } +func providerModel(modelID, requestModel string) string { + if model := strings.TrimSpace(modelID); model != "" { + return model + } + return strings.TrimSpace(requestModel) +} + +func createdAtUnix(value string) int64 { + if value == "" { + return converterutil.GetCurrentTimestamp() + } + if ts, err := time.Parse(time.RFC3339Nano, value); err == nil { + return ts.Unix() + } + return converterutil.GetCurrentTimestamp() +} + func CreateURL(baseURL string) string { return strings.TrimSuffix(baseURL, "/") + "/api/banana/create-async" } diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go index d17c8203..fb43d2ec 100644 --- a/internal/converter/sosana/images_test.go +++ b/internal/converter/sosana/images_test.go @@ -7,6 +7,7 @@ import ( "mime/multipart" "strings" "testing" + "time" "github.com/mixaill76/auto_ai_router/internal/converter/openai" "github.com/stretchr/testify/assert" @@ -29,7 +30,7 @@ func TestImageGenerationRequest(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { body := []byte(`{"model":"nano-banana","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) - got, err := ImageGenerationRequest(body, "fallback-model") + got, err := ImageGenerationRequest(body, "nano-banana") require.NoError(t, err) var req BananaCreateRequest @@ -42,6 +43,15 @@ func TestImageGenerationRequest(t *testing.T) { } } +func TestImageGenerationRequestPrefersProviderModel(t *testing.T) { + got, err := ImageGenerationRequest([]byte(`{"model":"public-image","prompt":"draw","n":1}`), "nano-banana") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "nano-banana", req.Model) +} + func TestImageGenerationRequestRejectsMultipleImages(t *testing.T) { _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","n":2}`), "nano-banana") require.Error(t, err) @@ -58,7 +68,7 @@ func TestImageEditRequest(t *testing.T) { "image": pngBytes(), }) - got, err := ImageEditRequest(body, contentType, "fallback-model") + got, err := ImageEditRequest(body, contentType, "nano-banana") require.NoError(t, err) var req BananaCreateRequest @@ -71,6 +81,22 @@ func TestImageEditRequest(t *testing.T) { assert.Contains(t, req.ImageURLs[0], base64.StdEncoding.EncodeToString(pngBytes())) } +func TestImageEditRequestPrefersProviderModel(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "public-image", + "prompt": "make it blue", + }, map[string][]byte{ + "image": pngBytes(), + }) + + got, err := ImageEditRequest(body, contentType, "nano-banana") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "nano-banana", req.Model) +} + func TestImageEditRequestRejectsMask(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ "model": "nano-banana", @@ -101,9 +127,12 @@ func TestImageEditRequestRejectsMultipleImagesCount(t *testing.T) { func TestOpenAIImageResponse(t *testing.T) { url := "https://cdn.sosana.art/result.png" + createdAt := "2026-01-01T00:00:00Z" body, err := OpenAIImageResponse(BananaTaskResponse{ - Status: StatusCompleted, - ResultFileURL: &url, + Status: StatusCompleted, + CreatedAt: createdAt, + OptimizedPrompt: "A detailed result prompt", + ResultFileURL: &url, }) require.NoError(t, err) @@ -111,8 +140,11 @@ func TestOpenAIImageResponse(t *testing.T) { require.NoError(t, json.Unmarshal(body, &resp)) require.Len(t, resp.Data, 1) assert.Equal(t, url, resp.Data[0].URL) + assert.Equal(t, "A detailed result prompt", resp.Data[0].RevisedPrompt) assert.Empty(t, resp.Data[0].B64JSON) - assert.NotZero(t, resp.Created) + ts, err := time.Parse(time.RFC3339, createdAt) + require.NoError(t, err) + assert.Equal(t, ts.Unix(), resp.Created) } func multipartImageEditBody(t *testing.T, fields map[string]string, files map[string][]byte) ([]byte, string) { diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index c9e0045e..29c89afd 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -146,10 +146,9 @@ func (p *Proxy) createAndPollSosanaTask( if err != nil { body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosana.CreateURL(cred.BaseURL), logCtx) return sosanaAttemptResult{ - body: body, - statusCode: code, - retryable: ctx.Err() == nil, - retryReason: RetryReasonNetErr, + body: body, + statusCode: code, + retryable: false, } } if statusCode >= 400 { diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 27f53b26..67a06fb1 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "log/slog" + "mime/multipart" "net/http" "net/http/httptest" "strings" @@ -12,6 +13,7 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/mixaill76/auto_ai_router/internal/models" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -31,7 +33,7 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox"}`)) case "/api/banana/task-1": pollSeen = true - _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox","result_file_url":"https://cdn.sosana.art/fox.png"}`)) + _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox","optimized_prompt":"A detailed fox illustration","result_file_url":"https://cdn.sosana.art/fox.png"}`)) default: t.Fatalf("unexpected upstream path: %s", r.URL.Path) } @@ -51,11 +53,54 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) require.Len(t, resp.Data, 1) assert.Equal(t, "https://cdn.sosana.art/fox.png", resp.Data[0].URL) + assert.Equal(t, "A detailed fox illustration", resp.Data[0].RevisedPrompt) assert.Empty(t, resp.Data[0].B64JSON) + assert.Equal(t, int64(1767225600), resp.Created) assert.True(t, createSeen) assert.True(t, pollSeen) } +func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/api/banana/create-async", r.URL.Path) + assert.Equal(t, "Bearer sosana-key", r.Header.Get("Authorization")) + + var req map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + assert.Equal(t, "make it blue", req["prompt"]) + assert.Equal(t, "nano-banana", req["model"]) + imageURLs, ok := req["image_urls"].([]any) + require.True(t, ok) + require.Len(t, imageURLs, 1) + assert.True(t, strings.HasPrefix(imageURLs[0].(string), "data:image/png;base64,")) + + _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"make it blue","result_file_url":"https://cdn.sosana.art/edit.png"}`)) + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, nil) + prx.modelManager = models.New(prx.logger, 50, []config.ModelRPMConfig{ + {Name: "public-image", Model: "nano-banana", Credential: "sosana"}, + }) + + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "public-image", + "prompt": "make it blue", + "n": "1", + }, map[string][]byte{ + "image": {0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}, + }) + req := httptest.NewRequest("POST", "/v1/images/edits", bytes.NewReader(body)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", contentType) + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + assert.Contains(t, w.Body.String(), "https://cdn.sosana.art/edit.png") +} + func TestProxyRequest_SosanaRejectsNonImageEndpoint(t *testing.T) { called := false upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -142,6 +187,39 @@ func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { assert.Contains(t, w.Body.String(), "https://cdn.sosana.art/retry.png") } +func TestProxyRequest_SosanaDoesNotRetryCreateTransportError(t *testing.T) { + deadServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + deadURL := deadServer.URL + deadServer.Close() + + liveCalled := false + liveServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + liveCalled = true + _, _ = w.Write([]byte(`{"uid":"task-2","status":"COMPLETED","prompt":"draw","result_file_url":"https://cdn.sosana.art/unwanted.png"}`)) + })) + defer liveServer.Close() + + prx := NewTestProxyBuilder(). + WithCredentials( + config.CredentialConfig{Name: "sosana-a", Type: config.ProviderTypeSosana, BaseURL: deadURL, APIKey: "sosana-key-a", RPM: 100, TPM: 10000}, + config.CredentialConfig{Name: "sosana-b", Type: config.ProviderTypeSosana, BaseURL: liveServer.URL, APIKey: "sosana-key-b", RPM: 100, TPM: 10000}, + ). + WithMaxProviderRetries(1). + Build() + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusBadGateway, w.Code) + assert.False(t, liveCalled) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "unwanted.png") +} + func TestProxyRequest_SosanaPollHTTPErrorMasked(t *testing.T) { var logBuf bytes.Buffer upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -302,3 +380,21 @@ func (b *TestProxyBuilder) withLogger(logger *slog.Logger) *TestProxyBuilder { b.config.ModelManager = createTestModelManager(logger) return b } + +func sosanaMultipartEditBody(t *testing.T, fields map[string]string, files map[string][]byte) ([]byte, string) { + t.Helper() + + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + for key, value := range fields { + require.NoError(t, writer.WriteField(key, value)) + } + for key, data := range files { + part, err := writer.CreateFormFile(key, key+".png") + require.NoError(t, err) + _, err = part.Write(data) + require.NoError(t, err) + } + require.NoError(t, writer.Close()) + return buf.Bytes(), writer.FormDataContentType() +} From c9a1ede5bb4a632cf7cc84d18f88d9af5e226329 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Fri, 26 Jun 2026 14:27:27 +0300 Subject: [PATCH 03/16] test: cover sosana image billing --- config.yaml.example | 8 +- internal/config/config.go | 8 +- internal/config/config_test.go | 36 +++++++ internal/proxy/sosana_test.go | 171 ++++++++++++++++++++++++++++++++- 4 files changed, 213 insertions(+), 10 deletions(-) diff --git a/config.yaml.example b/config.yaml.example index ef27cecb..c84a9fba 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -2,9 +2,9 @@ server: port: 8080 max_body_size_mb: 100 # Maximum request body size in MB (default: 100) response_body_multiplier: 10 # Response body limit = max_body_size_mb * this value (default: 10) - request_timeout: 60s # Request timeout (default: 60s) - write_timeout: 60s # HTTP server write timeout (default: 60s) - idle_timeout: 2m # HTTP server idle timeout (default: 2*write_timeout) + request_timeout: 2m # Request timeout (default: 2m) + write_timeout: 2m # HTTP server write timeout (default: 2m) + idle_timeout: 4m # HTTP server idle timeout (default: 2*write_timeout) idle_conn_timeout: 120s # HTTP idle connection timeout (default: 120s) max_idle_conns: 200 # Maximum idle connections (default: 200) max_idle_conns_per_host: 20 # Maximum idle connections per host (default: 20) @@ -79,7 +79,7 @@ credentials: tpm: -1 # Sosana.art image generation API (OpenAI Images-compatible through the router) - # Async image tasks can run longer than chat requests; use request_timeout/write_timeout >= 2m in production. + # Async image tasks can run longer than short chat requests; keep request_timeout/write_timeout >= 2m. - name: "sosana_images" type: "sosana" api_key: "os.environ/SOSANA_API_KEY" diff --git a/internal/config/config.go b/internal/config/config.go index 4a9384f2..e5719662 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -432,7 +432,7 @@ type ServerConfig struct { MaxIdleConnsPerHost int `yaml:"max_idle_conns_per_host"` IdleConnTimeout time.Duration `yaml:"idle_conn_timeout"` ReadTimeout time.Duration `yaml:"-"` // HTTP server read timeout (equals request_timeout, not configurable via YAML) - WriteTimeout time.Duration `yaml:"write_timeout"` // HTTP server write timeout (default: 60s) + WriteTimeout time.Duration `yaml:"write_timeout"` // HTTP server write timeout (default: 2m) IdleTimeout time.Duration `yaml:"idle_timeout"` // HTTP server idle timeout (default: 2*write_timeout) MaxProviderRetries int `yaml:"max_provider_retries"` // Max same-type credential retries on provider errors (default: 2, meaning 3 total attempts) MaxFallbackAttempts int `yaml:"max_fallback_attempts"` // Max fallback proxy hops per request chain (default: 5) @@ -516,16 +516,16 @@ func (s *ServerConfig) UnmarshalYAML(value *yaml.Node) error { } // Duration fields - if s.RequestTimeout, err = parseField(temp.RequestTimeout, 60*time.Second, time.ParseDuration, "request_timeout"); err != nil { + if s.RequestTimeout, err = parseField(temp.RequestTimeout, 2*time.Minute, time.ParseDuration, "request_timeout"); err != nil { return err } if s.IdleConnTimeout, err = parseField(temp.IdleConnTimeout, 120*time.Second, time.ParseDuration, "idle_conn_timeout"); err != nil { return err } - if s.WriteTimeout, err = parseField(temp.WriteTimeout, 60*time.Second, time.ParseDuration, "write_timeout"); err != nil { + if s.WriteTimeout, err = parseField(temp.WriteTimeout, 2*time.Minute, time.ParseDuration, "write_timeout"); err != nil { return err } - if s.IdleTimeout, err = parseField(temp.IdleTimeout, 2*time.Minute, time.ParseDuration, "idle_timeout"); err != nil { + if s.IdleTimeout, err = parseField(temp.IdleTimeout, 4*time.Minute, time.ParseDuration, "idle_timeout"); err != nil { return err } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 1986aab0..b9e2e093 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1189,6 +1189,42 @@ monitoring: assert.Equal(t, 2, cfg.Server.MaxProviderRetries, "Default MaxProviderRetries should be 2") } +func TestLoad_ServerTimeoutDefaults(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "config.yaml") + + configContent := ` +server: + port: 8080 + max_body_size_mb: 10 + master_key: "sk-test" + +fail2ban: + max_attempts: 3 + ban_duration: permanent + error_codes: [401] + +credentials: + - name: "test" + type: "openai" + api_key: "sk-test" + base_url: "https://api.openai.com" + rpm: 10 + +monitoring: + prometheus_enabled: false +` + err := os.WriteFile(configPath, []byte(configContent), 0644) + require.NoError(t, err) + + cfg, err := Load(configPath) + require.NoError(t, err) + assert.Equal(t, 2*time.Minute, cfg.Server.RequestTimeout) + assert.Equal(t, 2*time.Minute, cfg.Server.ReadTimeout) + assert.Equal(t, 2*time.Minute, cfg.Server.WriteTimeout) + assert.Equal(t, 4*time.Minute, cfg.Server.IdleTimeout) +} + func TestLoad_MaxProviderRetries_Custom(t *testing.T) { tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, "config.yaml") diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 67a06fb1..5d4a4e09 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -2,18 +2,23 @@ package proxy import ( "bytes" + "context" "encoding/json" + "fmt" "log/slog" "mime/multipart" "net/http" "net/http/httptest" + "os" "strings" "testing" "time" "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter/openai" - "github.com/mixaill76/auto_ai_router/internal/models" + "github.com/mixaill76/auto_ai_router/internal/litellmdb" + litellmmodels "github.com/mixaill76/auto_ai_router/internal/litellmdb/models" + aimodels "github.com/mixaill76/auto_ai_router/internal/models" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -79,7 +84,7 @@ func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, nil) - prx.modelManager = models.New(prx.logger, 50, []config.ModelRPMConfig{ + prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ {Name: "public-image", Model: "nano-banana", Credential: "sosana"}, }) @@ -101,6 +106,146 @@ func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { assert.Contains(t, w.Body.String(), "https://cdn.sosana.art/edit.png") } +func TestProxyRequest_SosanaImageGenerationLogsLiteLLMImageSpend(t *testing.T) { + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":"https://cdn.sosana.art/spend.png"}`)) + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + spendManager := newCapturedSpendManager() + priceRegistry := aimodels.NewModelPriceRegistry() + priceRegistry.Update(map[string]*aimodels.ModelPrice{ + "nano-banana": {OutputCostPerImage: 0.07}, + }) + + prx := newSosanaTestProxy(upstream.URL, nil) + prx.LiteLLMDB = spendManager + prx.priceRegistry = priceRegistry + prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ + {Name: "public-image", Model: "nano-banana", Credential: "sosana"}, + }) + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"public-image","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Len(t, spendManager.entries, 1) + entry := spendManager.entries[0] + assert.Equal(t, "/v1/images/generations", entry.CallType) + assert.Equal(t, "public-image", entry.Model) + assert.Equal(t, "sosana:public-image", entry.ModelID) + assert.Equal(t, "sosana", entry.CustomLLMProvider) + assert.Equal(t, "success", entry.Status) + assert.Equal(t, 0, entry.TotalTokens) + assert.InDelta(t, 0.07, entry.Spend, 0.0000001) + var metadata map[string]any + require.NoError(t, json.Unmarshal([]byte(entry.Metadata), &metadata)) + costBreakdown, ok := metadata["cost_breakdown"].(map[string]any) + require.True(t, ok) + assert.InDelta(t, 0.07, costBreakdown["total_cost"].(float64), 0.0000001) +} + +func TestProxyRequest_SosanaImageGenerationWritesLiteLLMSpendLogIntegration(t *testing.T) { + dbURL := os.Getenv("LITELLM_DATABASE_URL") + if dbURL == "" { + t.Skip("LITELLM_DATABASE_URL not set, skipping LiteLLM spend-log integration test") + } + + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":"https://cdn.sosana.art/spend.png"}`)) + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + manager, err := litellmdb.New(&litellmmodels.Config{ + DatabaseURL: dbURL, + MaxConns: 5, + MinConns: 1, + AuthCacheSize: 100, + AuthCacheTTL: time.Second, + LogQueueSize: 100, + LogBatchSize: 1, + LogFlushInterval: 100 * time.Millisecond, + }) + require.NoError(t, err) + defer func() { + _ = manager.Shutdown(context.Background()) + }() + + alias := fmt.Sprintf("sosana-spend-test-%d", time.Now().UnixNano()) + priceRegistry := aimodels.NewModelPriceRegistry() + priceRegistry.Update(map[string]*aimodels.ModelPrice{ + "nano-banana": {OutputCostPerImage: 0.07}, + }) + + prx := newSosanaTestProxy(upstream.URL, nil) + prx.LiteLLMDB = manager + prx.priceRegistry = priceRegistry + prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ + {Name: alias, Model: "nano-banana", Credential: "sosana"}, + }) + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"`+alias+`","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + require.Equal(t, http.StatusOK, w.Code) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var requestID string + var spend float64 + var model string + var provider string + var status string + var totalTokens int + for { + err = manager.GetPool().QueryRow(ctx, ` + SELECT request_id, spend, model, custom_llm_provider, status, total_tokens + FROM "LiteLLM_SpendLogs" + WHERE model = $1 AND custom_llm_provider = 'sosana' + ORDER BY "startTime" DESC + LIMIT 1 + `, alias).Scan(&requestID, &spend, &model, &provider, &status, &totalTokens) + if err == nil { + break + } + if ctx.Err() != nil { + require.NoError(t, err) + } + time.Sleep(100 * time.Millisecond) + } + defer func() { + _, _ = manager.GetPool().Exec(context.Background(), `DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = $1`, requestID) + }() + + assert.Equal(t, alias, model) + assert.Equal(t, "sosana", provider) + assert.Equal(t, "success", status) + assert.Equal(t, 0, totalTokens) + assert.InDelta(t, 0.07, spend, 0.0000001) +} + func TestProxyRequest_SosanaRejectsNonImageEndpoint(t *testing.T) { called := false upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -398,3 +543,25 @@ func sosanaMultipartEditBody(t *testing.T, fields map[string]string, files map[s require.NoError(t, writer.Close()) return buf.Bytes(), writer.FormDataContentType() } + +type capturedSpendManager struct { + *litellmdb.NoopManager + entries []*litellmmodels.SpendLogEntry +} + +func newCapturedSpendManager() *capturedSpendManager { + return &capturedSpendManager{NoopManager: litellmdb.NewNoopManager()} +} + +func (m *capturedSpendManager) IsEnabled() bool { + return true +} + +func (m *capturedSpendManager) IsHealthy() bool { + return true +} + +func (m *capturedSpendManager) LogSpend(entry *litellmmodels.SpendLogEntry) error { + m.entries = append(m.entries, entry) + return nil +} From 5d4d3e6bb88b9f248cab067c63c979402b240cf5 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Fri, 26 Jun 2026 14:56:06 +0300 Subject: [PATCH 04/16] fix: log masked upstream error bodies --- internal/proxy/cometapi_test.go | 12 ++++++++++++ internal/proxy/sosana_test.go | 8 ++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/internal/proxy/cometapi_test.go b/internal/proxy/cometapi_test.go index a2028cfc..b48a853a 100644 --- a/internal/proxy/cometapi_test.go +++ b/internal/proxy/cometapi_test.go @@ -61,6 +61,18 @@ func TestIsCometAPICredential(t *testing.T) { } } +func TestAppendResponseBodyForLogs_MaskedProviderKeepsMaskedFlagAndLogsBody(t *testing.T) { + cred := &config.CredentialConfig{Type: config.ProviderTypeCometAPI} + body := `{"error":{"code":"permission_denied","message":"` + strings.Repeat("model access denied ", 50) + `","type":"comet_api_error"}}` + + args := appendResponseBodyForLogs([]any{}, cred, body) + + assert.Contains(t, args, "response_body_masked") + assert.Contains(t, args, true) + assert.Contains(t, args, "response_body") + assert.Contains(t, args, body) +} + func TestIsSosanaCredential(t *testing.T) { tests := []struct { name string diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 5d4a4e09..bc926747 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -287,7 +287,7 @@ func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { assert.Contains(t, w.Body.String(), "Upstream provider error") assert.NotContains(t, w.Body.String(), "balance secret") assert.Contains(t, logBuf.String(), "response_body_masked=true") - assert.NotContains(t, logBuf.String(), "balance secret") + assert.Contains(t, logBuf.String(), "balance secret") } func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { @@ -390,7 +390,7 @@ func TestProxyRequest_SosanaPollHTTPErrorMasked(t *testing.T) { assert.Contains(t, w.Body.String(), "Upstream provider error") assert.NotContains(t, w.Body.String(), "poll secret") assert.Contains(t, logBuf.String(), "response_body_masked=true") - assert.NotContains(t, logBuf.String(), "poll secret") + assert.Contains(t, logBuf.String(), "poll secret") } func TestProxyRequest_SosanaDoesNotRetryAfterTaskCreated(t *testing.T) { @@ -456,7 +456,7 @@ func TestProxyRequest_SosanaTaskFailedMasked(t *testing.T) { assert.Contains(t, w.Body.String(), "Upstream provider error") assert.NotContains(t, w.Body.String(), "failed secret") assert.Contains(t, logBuf.String(), "response_body_masked=true") - assert.NotContains(t, logBuf.String(), "failed secret") + assert.Contains(t, logBuf.String(), "failed secret") } func TestProxyRequest_SosanaTaskModeratedMasked(t *testing.T) { @@ -483,7 +483,7 @@ func TestProxyRequest_SosanaTaskModeratedMasked(t *testing.T) { assert.Contains(t, w.Body.String(), "content_policy_violation") assert.NotContains(t, w.Body.String(), "moderation secret") assert.Contains(t, logBuf.String(), "response_body_masked=true") - assert.NotContains(t, logBuf.String(), "moderation secret") + assert.Contains(t, logBuf.String(), "moderation secret") } func TestProxyRequest_SosanaTimeoutMasked(t *testing.T) { From 941733e9887b8949dc192339a58e72bb14c378b2 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Fri, 26 Jun 2026 17:16:20 +0300 Subject: [PATCH 05/16] feat: return sosana images as b64 json --- docs/providers/sosana-integration-readme.md | 455 ++++++++++++++++++++ docs/providers/sosana.md | 14 +- internal/converter/sosana/images.go | 21 +- internal/converter/sosana/images_test.go | 28 +- internal/proxy/sosana.go | 204 ++++++++- internal/proxy/sosana_live_test.go | 6 +- internal/proxy/sosana_test.go | 187 +++++++- 7 files changed, 883 insertions(+), 32 deletions(-) create mode 100644 docs/providers/sosana-integration-readme.md diff --git a/docs/providers/sosana-integration-readme.md b/docs/providers/sosana-integration-readme.md new file mode 100644 index 00000000..a18247ab --- /dev/null +++ b/docs/providers/sosana-integration-readme.md @@ -0,0 +1,455 @@ +# Sosana.art Integration Handoff + +Last reviewed: 2026-06-26 + +This document summarizes the implemented Sosana.art integration, the model +mapping research against VSELLM, and the remaining product decisions. It is +intended as an engineering handoff, not end-user documentation. + +## Current Scope + +Sosana.art is integrated as a minimal image-only provider: + +- Provider type: `sosana` +- Supported public endpoints: + - `POST /v1/images/generations` + - `POST /v1/images/edits` +- Unsupported endpoints for `sosana` credentials: + - Chat Completions + - Responses API + - Embeddings + - Audio + - Video + - Slides + +Unsupported endpoints return a local OpenAI-compatible 400 error and do not call +Sosana. + +Main implementation files: + +- `internal/config/config.go` - provider type, aliases, validation. +- `internal/converter/sosana/images.go` - pure OpenAI Images to Sosana Banana + conversion and Sosana task to OpenAI Images response conversion. +- `internal/proxy/sosana.go` - create + poll flow, retry behavior, error + mapping, LiteLLM spend logging. +- `internal/proxy/proxy_log.go` - unified upstream error logging and masking. +- `internal/proxy/upstream_masking.go` - proxy-chain Sosana masking helpers. +- `internal/proxy/image_response.go` - normalization for masked proxy image + responses. +- `docs/providers/sosana.md` - short public provider doc. +- `config.yaml.example` - example Sosana credential and model. + +## Provider Configuration + +The provider type accepts these forms: + +- `sosana` +- `sosana-art` +- `sosana_art` + +Config validation requires: + +- `api_key` +- `base_url` + +Recommended base URL: + +```yaml +credentials: + - name: "sosana_images" + type: "sosana" + api_key: "os.environ/SOSANA_API_KEY" + base_url: "https://sosana.art" +``` + +Production Sosana credentials should use request and server write timeouts of at +least 2 minutes. Sosana documents async image tasks as typically completing in +20-120 seconds. + +## Request Flow + +### Image Generation + +Client request: + +```http +POST /v1/images/generations +``` + +OpenAI-compatible JSON is converted to a Sosana Banana create request: + +```json +{ + "prompt": "...", + "model": "nano-banana", + "aspect_ratio": "1:1" +} +``` + +The model sent to Sosana is the resolved real model id from `models[].model` when +configured. This allows VSELLM public model ids to map to Sosana internal model +ids without changing the client-visible model name. + +### Image Edits + +Client request: + +```http +POST /v1/images/edits +Content-Type: multipart/form-data +``` + +The converter: + +- reads `prompt`, `model`, `size`, and `n` fields; +- reads uploaded `image`, `images`, or `image[]` file parts; +- encodes uploaded images as `data:image/...;base64,...`; +- sends them to Sosana as `image_urls`; +- rejects `mask` locally because Sosana Banana does not support masks in the + implemented flow. + +Each multipart image part is capped at 20 MB. + +### Async Sosana Flow + +The proxy does not expose Sosana async tasks to clients. It performs the async +flow synchronously inside the request: + +1. `POST {base_url}/api/banana/create-async` +2. Immediate first poll +3. `GET {base_url}/api/banana/{uid}` every 2 seconds +4. Stop on `COMPLETED`, `FAILED`, `MODERATED`, unknown terminal error, or + request deadline + +`request_timeout` is used as the polling deadline when positive. A timeout value +of `-1` means no extra deadline is added by the proxy. + +## Current Response Behavior + +On successful Sosana completion, the router downloads Sosana's +`result_file_url`, verifies that the body is an image, encodes it as base64, and +returns: + +```json +{ + "created": 1782478551, + "data": [ + { + "b64_json": "iVBORw0KGgo...", + "revised_prompt": "..." + } + ] +} +``` + +This matches the VSELLM public image generation examples, which decode +`data[].b64_json`. Sosana's public object URL remains an internal implementation +detail and is not returned to clients. + +Important implementation details: + +- The object download uses the request context and must fit into + `request_timeout`. +- The object download does not include the Sosana `Authorization` header. +- Raw image bytes are capped at 32 MiB and are never written to logs. +- `response_format: "url"` is rejected locally with 400 until VSELLM-owned + rehosting exists. + +## Local Validation + +The integration intentionally rejects unsupported OpenAI Images features before +calling Sosana: + +- `n > 1` returns local 400: `sosana supports n=1 only`. +- Missing prompt returns local 400. +- Missing image on edits returns local 400. +- `mask` on edits returns local 400. +- Non-image endpoints with a `sosana` credential return local 400: + `sosana provider supports only image generation`. +- `response_format: "url"` returns local 400 because direct Sosana storage URLs + would reveal the upstream provider. + +## Error Mapping + +Sosana details must not leak to clients. + +Client-facing mapping: + +| Condition | HTTP status | Client body | +| --- | ---: | --- | +| Local validation error | 400 | Specific local OpenAI-compatible error | +| Unsupported endpoint | 400 | `sosana provider supports only image generation` | +| Create HTTP error | Upstream status | Generic upstream OpenAI-compatible error | +| Poll HTTP error | Upstream status | Generic upstream OpenAI-compatible error | +| Transport error | 502 | Generic upstream OpenAI-compatible error | +| Transport timeout / deadline | 408 | Generic upstream OpenAI-compatible error | +| Task `FAILED` | 502 | Generic upstream OpenAI-compatible error | +| Task `MODERATED` | 400 | Neutral `content_policy_violation` | +| Task `COMPLETED` without result URL | 502 | Generic upstream OpenAI-compatible error | +| Result image download HTTP error | 502 | Generic upstream OpenAI-compatible error | +| Result image download timeout | 408 | Generic upstream OpenAI-compatible error | +| Result image too large / non-image / read error | 502 | Generic upstream OpenAI-compatible error | +| Unknown task status | 502 | Generic upstream OpenAI-compatible error | + +`FAILED` and upstream HTTP errors deliberately do not return Sosana's raw +`error`, `detail`, balance text, policy text, uid, or provider-specific codes. + +## Debug Logging And Masking + +Sosana is considered a masked upstream provider when any of these are true: + +- credential type is `sosana`; +- credential base URL host is `sosana.art` or a subdomain; +- credential name contains `sosana` or `sasana`; +- `mask_upstream_errors: true` is set. + +For masked upstream failures: + +- client receives only generic/neutral OpenAI-compatible errors; +- structured logs include `response_body_masked=true`; +- structured logs also include raw `response_body` for internal debugging. + +This mirrors the Comet API style: external users do not see provider internals, +while operators can still debug the actual upstream response. + +For proxy-chain setups, masking is reliable when: + +- the upstream router propagates `X-Credential-Name` with a Sosana marker; or +- the proxy credential has `mask_upstream_errors: true`. + +Older external proxies without a credential marker remain a configuration risk. + +## Retry And Fallback Behavior + +Sosana retry behavior is intentionally narrower than generic direct LLM +providers because image creation can be billable once a task is accepted. + +Implemented retry: + +- Retry same provider type (`sosana`) with another Sosana credential. +- Retry only when the create request returns a retryable HTTP status before a + task is accepted, for example 429 or 5xx according to the shared retry + classifier. + +No retry across credentials after task creation: + +- Poll HTTP errors are not retried against another credential. +- Task `FAILED` is not retried against another credential. +- Transport errors are not retried by the Sosana special-case because it can be + ambiguous whether the create request reached Sosana and became billable. + +This avoids duplicate image charges and duplicate generated assets. + +## Billing + +Successful Sosana image requests set: + +```go +logCtx.TokenUsage = &converter.TokenUsage{ImageCount: logCtx.ImageCount} +``` + +`ImageCount` is extracted from `n` and defaults to 1. Since Sosana currently +supports only `n=1`, successful spend should be one image per request. + +LiteLLM spend calculation uses: + +```text +spend = image_count * output_cost_per_image +``` + +Price lookup tries the real model name first, then the public alias. Production +pricing should include the public VSELLM model id and/or the Sosana real model +id used by `models[].model`. + +If future support for `n > 1` is added by spawning multiple Sosana tasks, billing +must keep `ImageCount` equal to the number of completed/generated images. + +## VSELLM Image Model Research + +VSELLM public docs and pricing list these image-generation models: + +- `google/gemini-3-pro-image-preview` +- `vertex_ai/imagen-4.0-fast-generate-001` +- `google/gemini-3.1-flash-image-preview` +- `vertex_ai/imagen-4.0-generate-001` +- `vertex_ai/imagen-4.0-ultra-generate-001` +- `google/gemini-2.5-flash-image` +- `openai/gpt-image-1-mini` +- `openai/gpt-image-1` + +VSELLM docs for `/v1/images/edits` mention these edit-capable models: + +- `openai/gpt-image-1-mini` +- `openai/gpt-image-1` +- `google/gemini-2.5-flash-image` +- `google/gemini-3-pro-image-preview` +- `google/gemini-3.1-flash-image-preview` + +Sources: + +- `https://vsellm.ru/docs` +- `https://vsellm.ru/provider/VSELLM` + +## Sosana Model Research + +Sosana Banana OpenAPI currently lists these image model ids: + +- `nano-banana` +- `nano-banana-2-1k` +- `nano-banana-2-2k` +- `nano-banana-2-2k-thinking` +- `nano-banana-pro-1k` +- `nano-banana-pro-2k` +- `nano-banana-pro-4k` +- `gpt-image-2` + +Sosana pricing also lists "Nano Banana 2 4K", but the OpenAPI enum does not +currently expose a confirmed `nano-banana-2-4k` id. Do not use that id until it +is live-tested or confirmed by Sosana. + +Sources: + +- `https://sosana.art/openapi.json` +- `https://sosana.art/` + +## Recommended VSELLM To Sosana Mapping + +Use `models[].model` for mapping. The `name` remains the public VSELLM model id, +and `model` is the Sosana model sent upstream. + +Recommended minimal production mapping: + +```yaml +models: + - name: "google/gemini-2.5-flash-image" + model: "nano-banana" + credential: sosana_images + + - name: "google/gemini-3.1-flash-image-preview" + model: "nano-banana-2-1k" + credential: sosana_images + + - name: "google/gemini-3-pro-image-preview" + model: "nano-banana-pro-1k" + credential: sosana_images +``` + +Why these map cleanly: + +- VSELLM describes `google/gemini-2.5-flash-image` as Gemini 2.5 Flash Image; + Sosana `nano-banana` is the matching Nano Banana family. +- VSELLM describes `google/gemini-3.1-flash-image-preview` as Nano Banana 2 / + Gemini 3.1 Flash Image Preview; Sosana `nano-banana-2-*` is the matching + family. +- VSELLM describes `google/gemini-3-pro-image-preview` as Nano Banana Pro / + Gemini 3 Pro Image Preview; Sosana `nano-banana-pro-*` is the matching family. + +Do not map these VSELLM models to Sosana Banana: + +- `vertex_ai/imagen-4.0-fast-generate-001` +- `vertex_ai/imagen-4.0-generate-001` +- `vertex_ai/imagen-4.0-ultra-generate-001` +- `openai/gpt-image-1-mini` +- `openai/gpt-image-1` + +Those are different model families with different quality, pricing, latency, and +behavior expectations. Substituting Sosana Banana behind those names would be a +model mismatch. + +`gpt-image-2` note: + +- Sosana exposes `gpt-image-2`. +- VSELLM did not show `openai/gpt-image-2` in the reviewed public image model + list. +- Do not expose it under an existing OpenAI model id. Add a new public model id + only as an explicit product decision. + +## 1K, 2K, 4K, And Thinking Variants + +The Sosana model suffixes should be treated as explicit cost/quality tiers, not +as hidden automatic upgrades. + +Recommended default behavior: + +- `google/gemini-3.1-flash-image-preview` -> `nano-banana-2-1k` +- `google/gemini-3-pro-image-preview` -> `nano-banana-pro-1k` + +Reasons: + +- 1K is the safest default for cost and latency. +- Silent upgrades to 2K, 4K, or thinking variants can surprise users and billing. +- Routing based on prompt complexity is subjective and hard to explain. + +If product needs higher tiers, prefer explicit public model aliases: + +```yaml +models: + - name: "google/gemini-3.1-flash-image-preview-2k" + model: "nano-banana-2-2k" + credential: sosana_images + + - name: "google/gemini-3.1-flash-image-preview-thinking" + model: "nano-banana-2-2k-thinking" + credential: sosana_images + + - name: "google/gemini-3-pro-image-preview-2k" + model: "nano-banana-pro-2k" + credential: sosana_images + + - name: "google/gemini-3-pro-image-preview-4k" + model: "nano-banana-pro-4k" + credential: sosana_images +``` + +A future `size`/`quality` based selector is possible, but it should be explicit, +tested, documented, and priced. It should not be inferred from prompt text. + +## Known Gaps Before Full VSELLM Parity + +1. VSELLM-owned URL responses: + - Default and `b64_json` responses hide Sosana storage. + - `response_format: "url"` is rejected locally. + - If URL responses are required, re-host in VSELLM-owned storage before + returning to clients. + +2. `mask` support: + - VSELLM edit docs mention optional masks. + - Current Sosana converter rejects masks. + - Do not route mask-dependent public models only to Sosana unless local 400 is + acceptable for that model. + +3. `n > 1` support: + - Current Sosana integration supports only `n=1`. + - VSELLM docs describe `n` as number of variants. + - Future support should create multiple tasks and bill by completed image + count. + +4. Tier selection: + - Current mapping is static via `models[].model`. + - 1K/2K/4K selection by `size` or `quality` is not implemented. + +## Suggested Verification + +Unit/regression tests: + +```bash +go test ./internal/converter/sosana ./internal/proxy ./internal/config ./internal/models ./internal/litellmdb/model_table -count=1 +go test ./... -count=1 +git diff --check +``` + +Live smoke test: + +1. Start router with a Sosana credential. +2. Check `/health` shows the Sosana credential and configured model. +3. Call `/v1/images/generations` with `n=1`. +4. Call `/v1/images/generations` with `n=2`; verify local 400 and no upstream + call. +5. Call `/v1/chat/completions` with a Sosana model; verify local 400 and no + upstream call. +6. Start with a bad Sosana key; verify the client response is masked and logs + contain `response_body_masked=true` plus raw `response_body` for debugging. + +For live bad-key tests, mount the config by absolute path. A missing host file in +`docker run -v` can be created as a directory by Docker, causing +`config.yaml: is a directory`. diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md index 097db5a0..3227614c 100644 --- a/docs/providers/sosana.md +++ b/docs/providers/sosana.md @@ -3,7 +3,7 @@ Sosana.art is supported as an image-only provider for the OpenAI-compatible Images API. The router accepts `/v1/images/generations` and `/v1/images/edits`, submits a Sosana Banana async task, polls it, and returns an OpenAI Images -response with `data[].url`. +response with `data[].b64_json`. Chat Completions, Responses API, Embeddings, video, and slides are not routed to Sosana in this integration. @@ -33,12 +33,20 @@ completion requests. For production Sosana credentials, set the router ## Behavior - `n` must be `1`. -- `response_format: "b64_json"` is accepted, but Sosana results are returned as - URLs because Sosana provides `result_file_url`. +- Default response format and `response_format: "b64_json"` return + `data[].b64_json`. +- `response_format: "url"` returns a local 400 because URL responses require + VSELLM-owned rehosting before they can hide Sosana storage. - `/v1/images/edits` sends uploaded images as `data:image/...;base64,...` values in Sosana `image_urls`. - Mask images are not supported. +On completion, Sosana returns a public object URL in `result_file_url`. The +router downloads that object without forwarding the Sosana `Authorization` +header, keeps the download bounded to 32 MiB, verifies the body is an image, and +base64-encodes it into the OpenAI-compatible JSON response. The upstream object +URL is not returned to clients. + ## Error Masking Sosana upstream HTTP errors and terminal task errors are masked before they are diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go index 3d28d3ec..ca138c95 100644 --- a/internal/converter/sosana/images.go +++ b/internal/converter/sosana/images.go @@ -51,6 +51,9 @@ func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, error) { if err := validateImageCount(req.N); err != nil { return nil, err } + if err := validateResponseFormat(req.ResponseFormat); err != nil { + return nil, err + } prompt := strings.TrimSpace(req.Prompt) if prompt == "" { return nil, fmt.Errorf("image generation request missing prompt") @@ -112,6 +115,9 @@ func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([] if err := validateImageCountString(fields["n"]); err != nil { return nil, err } + if err := validateResponseFormat(fields["response_format"]); err != nil { + return nil, err + } prompt := strings.TrimSpace(fields["prompt"]) if prompt == "" { return nil, fmt.Errorf("image edit request missing prompt field") @@ -127,15 +133,15 @@ func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([] }) } -func OpenAIImageResponse(task BananaTaskResponse) ([]byte, error) { - if task.ResultFileURL == nil || strings.TrimSpace(*task.ResultFileURL) == "" { - return nil, fmt.Errorf("sosana task completed without result_file_url") +func OpenAIImageResponse(task BananaTaskResponse, image []byte) ([]byte, error) { + if len(image) == 0 { + return nil, fmt.Errorf("sosana task completed without image bytes") } resp := openai.OpenAIImageResponse{ Created: createdAtUnix(task.CreatedAt), Data: []openai.OpenAIImageData{ { - URL: strings.TrimSpace(*task.ResultFileURL), + B64JSON: base64.StdEncoding.EncodeToString(image), RevisedPrompt: strings.TrimSpace(task.OptimizedPrompt), }, }, @@ -194,6 +200,13 @@ func validateImageCount(n *int) error { return fmt.Errorf("sosana supports n=1 only") } +func validateResponseFormat(format string) error { + if strings.EqualFold(strings.TrimSpace(format), "url") { + return fmt.Errorf("response_format=url is unsupported for this image model") + } + return nil +} + func validateImageCountString(raw string) error { raw = strings.TrimSpace(raw) if raw == "" { diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go index fb43d2ec..1d4c2872 100644 --- a/internal/converter/sosana/images_test.go +++ b/internal/converter/sosana/images_test.go @@ -58,6 +58,12 @@ func TestImageGenerationRequestRejectsMultipleImages(t *testing.T) { assert.Contains(t, err.Error(), "n=1") } +func TestImageGenerationRequestRejectsURLResponseFormat(t *testing.T) { + _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","response_format":"url"}`), "nano-banana") + require.Error(t, err) + assert.Contains(t, err.Error(), "response_format=url") +} + func TestImageEditRequest(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ "model": "nano-banana", @@ -125,23 +131,35 @@ func TestImageEditRequestRejectsMultipleImagesCount(t *testing.T) { assert.Contains(t, err.Error(), "n=1") } +func TestImageEditRequestRejectsURLResponseFormat(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "nano-banana", + "prompt": "make it blue", + "response_format": "url", + }, map[string][]byte{ + "image": pngBytes(), + }) + + _, err := ImageEditRequest(body, contentType, "fallback-model") + require.Error(t, err) + assert.Contains(t, err.Error(), "response_format=url") +} + func TestOpenAIImageResponse(t *testing.T) { - url := "https://cdn.sosana.art/result.png" createdAt := "2026-01-01T00:00:00Z" body, err := OpenAIImageResponse(BananaTaskResponse{ Status: StatusCompleted, CreatedAt: createdAt, OptimizedPrompt: "A detailed result prompt", - ResultFileURL: &url, - }) + }, pngBytes()) require.NoError(t, err) var resp openai.OpenAIImageResponse require.NoError(t, json.Unmarshal(body, &resp)) require.Len(t, resp.Data, 1) - assert.Equal(t, url, resp.Data[0].URL) + assert.Empty(t, resp.Data[0].URL) assert.Equal(t, "A detailed result prompt", resp.Data[0].RevisedPrompt) - assert.Empty(t, resp.Data[0].B64JSON) + assert.Equal(t, base64.StdEncoding.EncodeToString(pngBytes()), resp.Data[0].B64JSON) ts, err := time.Parse(time.RFC3339, createdAt) require.NoError(t, err) assert.Equal(t, ts.Unix(), resp.Created) diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index 29c89afd..f1a8501b 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -5,8 +5,11 @@ import ( "context" "encoding/json" "errors" + "io" "math/rand" "net/http" + "net/url" + "strings" "time" "github.com/mixaill76/auto_ai_router/internal/config" @@ -14,7 +17,11 @@ import ( "github.com/mixaill76/auto_ai_router/internal/converter/sosana" ) -const sosanaPollInterval = 2 * time.Second +const ( + sosanaPollInterval = 2 * time.Second + maxSosanaResultImageBytes int64 = 32 * 1024 * 1024 + maxSosanaResultErrorBytes int64 = 16 * 1024 +) type sosanaAttemptResult struct { body []byte @@ -209,15 +216,8 @@ func (p *Proxy) sosanaTaskBody( ) ([]byte, int, bool) { switch task.Status { case sosana.StatusCompleted: - body, err := sosana.OpenAIImageResponse(task) - if err != nil { - p.logUpstreamError(ctx, "Sosana completed task missing result", http.StatusBadGateway, cred, modelID, rawBody, - "url", sosana.PollURL(cred.BaseURL, task.UID), - "request_id", logCtx.RequestID, - "error", err) - return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true - } - return body, http.StatusOK, true + body, statusCode := p.sosanaCompletedImageBody(ctx, cred, modelID, task, rawBody, logCtx) + return body, statusCode, true case sosana.StatusFailed: p.logUpstreamError(ctx, "Sosana task failed", http.StatusBadGateway, cred, modelID, rawBody, "url", sosana.PollURL(cred.BaseURL, task.UID), @@ -245,6 +245,190 @@ func (p *Proxy) sosanaTaskBody( } } +func (p *Proxy) sosanaCompletedImageBody( + ctx context.Context, + cred *config.CredentialConfig, + modelID string, + task sosana.BananaTaskResponse, + rawBody []byte, + logCtx *RequestLogContext, +) ([]byte, int) { + image, contentType, statusCode, err := p.downloadSosanaResultImage(ctx, cred, modelID, task, rawBody, logCtx) + if err != nil { + return maskedUpstreamErrorBody(statusCode), statusCode + } + body, err := sosana.OpenAIImageResponse(task, image) + if err != nil { + p.logUpstreamError(ctx, "Sosana completed task could not be converted", http.StatusBadGateway, cred, modelID, nil, + "request_id", logCtx.RequestID, + "error", err) + return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway + } + p.logger.DebugContext(ctx, "Downloaded Sosana result image", + "credential", cred.Name, + "model", modelID, + "result_host", sosanaResultHost(task), + "image_bytes", len(image), + "content_type", contentType, + "request_id", logCtx.RequestID) + return body, http.StatusOK +} + +func (p *Proxy) downloadSosanaResultImage( + ctx context.Context, + cred *config.CredentialConfig, + modelID string, + task sosana.BananaTaskResponse, + rawTaskBody []byte, + logCtx *RequestLogContext, +) ([]byte, string, int, error) { + resultURL := "" + if task.ResultFileURL != nil { + resultURL = strings.TrimSpace(*task.ResultFileURL) + } + parsed, err := parseSosanaResultURL(resultURL) + if err != nil { + p.logUpstreamError(ctx, "Sosana completed task returned invalid result URL", http.StatusBadGateway, cred, modelID, rawTaskBody, + "request_id", logCtx.RequestID, + "error", err) + return nil, "", http.StatusBadGateway, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, resultURL, nil) + if err != nil { + p.logUpstreamError(ctx, "Failed to build Sosana result image request", http.StatusBadGateway, cred, modelID, nil, + "result_host", parsed.Hostname(), + "request_id", logCtx.RequestID, + "error", err) + return nil, "", http.StatusBadGateway, err + } + req.Header.Set("Accept", "image/*") + + resp, err := p.client.Do(req) + if err != nil { + statusCode := http.StatusBadGateway + if isTimeoutError(err) || errors.Is(ctx.Err(), context.DeadlineExceeded) { + statusCode = http.StatusRequestTimeout + } + p.logUpstreamError(context.Background(), "Sosana result image download failed", statusCode, cred, modelID, nil, + "result_host", parsed.Hostname(), + "request_id", logCtx.RequestID, + "error", err) + return nil, "", statusCode, err + } + defer func() { + if closeErr := resp.Body.Close(); closeErr != nil { + p.logger.WarnContext(ctx, "Failed to close Sosana result image body", "error", closeErr) + } + }() + + contentType := strings.TrimSpace(resp.Header.Get("Content-Type")) + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + errorBody := readTextBodyForSosanaResultLog(resp.Body, contentType) + p.logUpstreamError(ctx, "Sosana result image download returned error status", http.StatusBadGateway, cred, modelID, errorBody, + "upstream_status", resp.StatusCode, + "result_host", parsed.Hostname(), + "request_id", logCtx.RequestID) + return nil, "", http.StatusBadGateway, errors.New("sosana result image download returned error status") + } + + image, err := readLimitedSosanaResultImage(resp.Body) + if err != nil { + p.logUpstreamError(ctx, "Failed to read Sosana result image", http.StatusBadGateway, cred, modelID, nil, + "result_host", parsed.Hostname(), + "request_id", logCtx.RequestID, + "error", err) + return nil, "", http.StatusBadGateway, err + } + sniffedType := http.DetectContentType(image) + if !isImageContentType(contentType) && !isImageContentType(sniffedType) { + responseBody := textBodyPrefixForSosanaResultLog(image, contentType, sniffedType) + p.logUpstreamError(ctx, "Sosana result URL returned non-image content", http.StatusBadGateway, cred, modelID, responseBody, + "result_host", parsed.Hostname(), + "content_type", contentType, + "sniffed_content_type", sniffedType, + "request_id", logCtx.RequestID) + return nil, "", http.StatusBadGateway, errors.New("sosana result URL returned non-image content") + } + if !isImageContentType(contentType) { + contentType = sniffedType + } + return image, contentType, http.StatusOK, nil +} + +func parseSosanaResultURL(raw string) (*url.URL, error) { + if raw == "" { + return nil, errors.New("sosana task completed without result_file_url") + } + parsed, err := url.Parse(raw) + if err != nil { + return nil, err + } + if parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return nil, errors.New("sosana result_file_url must be an http or https URL") + } + return parsed, nil +} + +func sosanaResultHost(task sosana.BananaTaskResponse) string { + if task.ResultFileURL == nil { + return "" + } + parsed, err := url.Parse(strings.TrimSpace(*task.ResultFileURL)) + if err != nil { + return "" + } + return parsed.Hostname() +} + +func readLimitedSosanaResultImage(body io.Reader) ([]byte, error) { + data, err := io.ReadAll(io.LimitReader(body, maxSosanaResultImageBytes+1)) + if err != nil { + return nil, err + } + if int64(len(data)) > maxSosanaResultImageBytes { + return nil, ErrResponseBodyTooLarge + } + if len(data) == 0 { + return nil, errors.New("sosana result image body is empty") + } + return data, nil +} + +func readTextBodyForSosanaResultLog(body io.Reader, contentType string) []byte { + if !isTextContentType(contentType) { + return nil + } + data, err := io.ReadAll(io.LimitReader(body, maxSosanaResultErrorBytes)) + if err != nil { + return nil + } + return data +} + +func textBodyPrefixForSosanaResultLog(body []byte, contentTypes ...string) []byte { + for _, contentType := range contentTypes { + if isTextContentType(contentType) { + if int64(len(body)) > maxSosanaResultErrorBytes { + return body[:maxSosanaResultErrorBytes] + } + return body + } + } + return nil +} + +func isImageContentType(contentType string) bool { + return strings.HasPrefix(strings.ToLower(strings.TrimSpace(contentType)), "image/") +} + +func isTextContentType(contentType string) bool { + contentType = strings.ToLower(strings.TrimSpace(contentType)) + return strings.HasPrefix(contentType, "text/") || + strings.Contains(contentType, "json") || + strings.Contains(contentType, "xml") +} + func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosana.BananaTaskResponse, []byte, int, error) { var reader *bytes.Reader if body != nil { diff --git a/internal/proxy/sosana_live_test.go b/internal/proxy/sosana_live_test.go index 32a52de3..7aa6608a 100644 --- a/internal/proxy/sosana_live_test.go +++ b/internal/proxy/sosana_live_test.go @@ -1,6 +1,7 @@ package proxy import ( + "encoding/base64" "encoding/json" "net/http" "net/http/httptest" @@ -59,5 +60,8 @@ func TestProxyRequest_SosanaLiveAcceptance(t *testing.T) { var resp openai.OpenAIImageResponse require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) require.Len(t, resp.Data, 1) - require.NotEmpty(t, resp.Data[0].URL) + require.Empty(t, resp.Data[0].URL) + require.NotEmpty(t, resp.Data[0].B64JSON) + _, err = base64.StdEncoding.DecodeString(resp.Data[0].B64JSON) + require.NoError(t, err) } diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index bc926747..de5932ef 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -3,6 +3,7 @@ package proxy import ( "bytes" "context" + "encoding/base64" "encoding/json" "fmt" "log/slog" @@ -25,6 +26,10 @@ import ( func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { var createSeen, pollSeen bool + var imageAuths []string + imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/png", sosanaResultPNG, &imageAuths) + defer imageServer.Close() + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { assert.Equal(t, "Bearer sosana-key", r.Header.Get("Authorization")) switch r.URL.Path { @@ -38,7 +43,7 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox"}`)) case "/api/banana/task-1": pollSeen = true - _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox","optimized_prompt":"A detailed fox illustration","result_file_url":"https://cdn.sosana.art/fox.png"}`)) + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox","optimized_prompt":"A detailed fox illustration","result_file_url":%q}`, imageServer.URL+"/fox.png") default: t.Fatalf("unexpected upstream path: %s", r.URL.Path) } @@ -57,15 +62,22 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { var resp openai.OpenAIImageResponse require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) require.Len(t, resp.Data, 1) - assert.Equal(t, "https://cdn.sosana.art/fox.png", resp.Data[0].URL) + assert.Empty(t, resp.Data[0].URL) + assert.Equal(t, base64.StdEncoding.EncodeToString(sosanaResultPNG), resp.Data[0].B64JSON) assert.Equal(t, "A detailed fox illustration", resp.Data[0].RevisedPrompt) - assert.Empty(t, resp.Data[0].B64JSON) assert.Equal(t, int64(1767225600), resp.Created) + assert.NotContains(t, w.Body.String(), imageServer.URL) + assert.NotContains(t, w.Body.String(), "sosana") + assert.NotContains(t, w.Body.String(), "cdn") + assert.Equal(t, []string{""}, imageAuths) assert.True(t, createSeen) assert.True(t, pollSeen) } func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { + imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/png", sosanaResultPNG, nil) + defer imageServer.Close() + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, "/api/banana/create-async", r.URL.Path) assert.Equal(t, "Bearer sosana-key", r.Header.Get("Authorization")) @@ -79,7 +91,7 @@ func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { require.Len(t, imageURLs, 1) assert.True(t, strings.HasPrefix(imageURLs[0].(string), "data:image/png;base64,")) - _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"make it blue","result_file_url":"https://cdn.sosana.art/edit.png"}`)) + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","created_at":"2026-01-01T00:00:00Z","prompt":"make it blue","result_file_url":%q}`, imageServer.URL+"/edit.png") })) defer upstream.Close() @@ -103,16 +115,23 @@ func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { prx.ProxyRequest(w, req) require.Equal(t, http.StatusOK, w.Code) - assert.Contains(t, w.Body.String(), "https://cdn.sosana.art/edit.png") + var resp openai.OpenAIImageResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.Len(t, resp.Data, 1) + assert.Empty(t, resp.Data[0].URL) + assert.Equal(t, base64.StdEncoding.EncodeToString(sosanaResultPNG), resp.Data[0].B64JSON) } func TestProxyRequest_SosanaImageGenerationLogsLiteLLMImageSpend(t *testing.T) { + imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/png", sosanaResultPNG, nil) + defer imageServer.Close() + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/banana/create-async": _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) case "/api/banana/task-1": - _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":"https://cdn.sosana.art/spend.png"}`)) + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/spend.png") default: t.Fatalf("unexpected upstream path: %s", r.URL.Path) } @@ -162,12 +181,15 @@ func TestProxyRequest_SosanaImageGenerationWritesLiteLLMSpendLogIntegration(t *t t.Skip("LITELLM_DATABASE_URL not set, skipping LiteLLM spend-log integration test") } + imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/png", sosanaResultPNG, nil) + defer imageServer.Close() + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/banana/create-async": _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) case "/api/banana/task-1": - _, _ = w.Write([]byte(`{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":"https://cdn.sosana.art/spend.png"}`)) + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/spend.png") default: t.Fatalf("unexpected upstream path: %s", r.URL.Path) } @@ -267,6 +289,27 @@ func TestProxyRequest_SosanaRejectsNonImageEndpoint(t *testing.T) { assert.False(t, called) } +func TestProxyRequest_SosanaRejectsURLResponseFormat(t *testing.T) { + called := false + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, nil) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1,"response_format":"url"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), "response_format=url") + assert.False(t, called) +} + func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { var logBuf bytes.Buffer upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -292,6 +335,9 @@ func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { var createAuths []string + imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/png", sosanaResultPNG, nil) + defer imageServer.Close() + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/banana/create-async": @@ -305,7 +351,7 @@ func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { _, _ = w.Write([]byte(`{"uid":"task-2","status":"PROCESSING","prompt":"draw"}`)) case "/api/banana/task-2": assert.Equal(t, "Bearer sosana-key-b", r.Header.Get("Authorization")) - _, _ = w.Write([]byte(`{"uid":"task-2","status":"COMPLETED","prompt":"draw","result_file_url":"https://cdn.sosana.art/retry.png"}`)) + _, _ = fmt.Fprintf(w, `{"uid":"task-2","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/retry.png") default: t.Fatalf("unexpected upstream path: %s", r.URL.Path) } @@ -329,7 +375,8 @@ func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { require.Equal(t, http.StatusOK, w.Code) assert.Equal(t, []string{"Bearer sosana-key-a", "Bearer sosana-key-b"}, createAuths) - assert.Contains(t, w.Body.String(), "https://cdn.sosana.art/retry.png") + assert.Contains(t, w.Body.String(), base64.StdEncoding.EncodeToString(sosanaResultPNG)) + assert.NotContains(t, w.Body.String(), imageServer.URL) } func TestProxyRequest_SosanaDoesNotRetryCreateTransportError(t *testing.T) { @@ -393,6 +440,111 @@ func TestProxyRequest_SosanaPollHTTPErrorMasked(t *testing.T) { assert.Contains(t, logBuf.String(), "poll secret") } +func TestProxyRequest_SosanaImageResultHTTPErrorMasked(t *testing.T) { + var logBuf bytes.Buffer + imageServer := newSosanaResultImageServer(t, http.StatusInternalServerError, "text/plain", []byte("storage secret marker"), nil) + defer imageServer.Close() + + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/missing.png") + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "storage secret") + assert.NotContains(t, w.Body.String(), imageServer.URL) + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.Contains(t, logBuf.String(), "storage secret") + assert.Contains(t, logBuf.String(), "result_host=127.0.0.1") +} + +func TestProxyRequest_SosanaImageResultNonImageMasked(t *testing.T) { + var logBuf bytes.Buffer + imageServer := newSosanaResultImageServer(t, http.StatusOK, "text/plain", []byte("not an image secret marker"), nil) + defer imageServer.Close() + + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/text.txt") + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "not an image") + assert.NotContains(t, w.Body.String(), imageServer.URL) + assert.Contains(t, logBuf.String(), "non-image content") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.Contains(t, logBuf.String(), "not an image secret marker") +} + +func TestProxyRequest_SosanaImageResultTimeoutMasked(t *testing.T) { + var logBuf bytes.Buffer + imageServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(50 * time.Millisecond) + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(sosanaResultPNG) + })) + defer imageServer.Close() + + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/slow.png") + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + prx.requestTimeout = 5 * time.Millisecond + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusRequestTimeout, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "slow.png") + assert.Contains(t, logBuf.String(), "result image download failed") + assert.Contains(t, logBuf.String(), "result_host=127.0.0.1") +} + func TestProxyRequest_SosanaDoesNotRetryAfterTaskCreated(t *testing.T) { createCalls := 0 upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -507,6 +659,23 @@ func TestProxyRequest_SosanaTimeoutMasked(t *testing.T) { assert.Contains(t, logBuf.String(), "response_body_masked=true") } +var sosanaResultPNG = []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} + +func newSosanaResultImageServer(t *testing.T, status int, contentType string, body []byte, auths *[]string) *httptest.Server { + t.Helper() + + return newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if auths != nil { + *auths = append(*auths, r.Header.Get("Authorization")) + } + if contentType != "" { + w.Header().Set("Content-Type", contentType) + } + w.WriteHeader(status) + _, _ = w.Write(body) + })) +} + func newSosanaTestProxy(baseURL string, logBuf *bytes.Buffer) *Proxy { logger := slog.New(slog.NewTextHandler(&bytes.Buffer{}, &slog.HandlerOptions{Level: slog.LevelDebug})) if logBuf != nil { From c8d29a2fa61aab6ccfc77e40b385ab6201b1573e Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Fri, 3 Jul 2026 16:41:52 +0300 Subject: [PATCH 06/16] fix: minimize sosana image integration --- config.yaml.example | 6 +- docs/providers/sosana-integration-readme.md | 455 -------------------- docs/providers/sosana.md | 28 +- internal/config/config.go | 8 +- internal/config/config_test.go | 8 +- internal/converter/openai/types.go | 10 +- internal/converter/sosana/compatibility.go | 228 ++++++++++ internal/converter/sosana/images.go | 65 ++- internal/converter/sosana/images_test.go | 77 ++++ internal/proxy/cometapi_test.go | 29 +- internal/proxy/image_response.go | 102 ----- internal/proxy/proxy.go | 37 -- internal/proxy/proxy_log.go | 2 +- internal/proxy/sosana.go | 99 ++++- internal/proxy/sosana_test.go | 204 ++++++++- internal/proxy/upstream_masking_test.go | 192 --------- 16 files changed, 711 insertions(+), 839 deletions(-) delete mode 100644 docs/providers/sosana-integration-readme.md create mode 100644 internal/converter/sosana/compatibility.go delete mode 100644 internal/proxy/image_response.go diff --git a/config.yaml.example b/config.yaml.example index c84a9fba..441f5fcc 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -2,9 +2,9 @@ server: port: 8080 max_body_size_mb: 100 # Maximum request body size in MB (default: 100) response_body_multiplier: 10 # Response body limit = max_body_size_mb * this value (default: 10) - request_timeout: 2m # Request timeout (default: 2m) - write_timeout: 2m # HTTP server write timeout (default: 2m) - idle_timeout: 4m # HTTP server idle timeout (default: 2*write_timeout) + request_timeout: 60s # Request timeout (default: 60s) + write_timeout: 60s # HTTP server write timeout (default: 60s) + idle_timeout: 2m # HTTP server idle timeout (default: 2*write_timeout) idle_conn_timeout: 120s # HTTP idle connection timeout (default: 120s) max_idle_conns: 200 # Maximum idle connections (default: 200) max_idle_conns_per_host: 20 # Maximum idle connections per host (default: 20) diff --git a/docs/providers/sosana-integration-readme.md b/docs/providers/sosana-integration-readme.md deleted file mode 100644 index a18247ab..00000000 --- a/docs/providers/sosana-integration-readme.md +++ /dev/null @@ -1,455 +0,0 @@ -# Sosana.art Integration Handoff - -Last reviewed: 2026-06-26 - -This document summarizes the implemented Sosana.art integration, the model -mapping research against VSELLM, and the remaining product decisions. It is -intended as an engineering handoff, not end-user documentation. - -## Current Scope - -Sosana.art is integrated as a minimal image-only provider: - -- Provider type: `sosana` -- Supported public endpoints: - - `POST /v1/images/generations` - - `POST /v1/images/edits` -- Unsupported endpoints for `sosana` credentials: - - Chat Completions - - Responses API - - Embeddings - - Audio - - Video - - Slides - -Unsupported endpoints return a local OpenAI-compatible 400 error and do not call -Sosana. - -Main implementation files: - -- `internal/config/config.go` - provider type, aliases, validation. -- `internal/converter/sosana/images.go` - pure OpenAI Images to Sosana Banana - conversion and Sosana task to OpenAI Images response conversion. -- `internal/proxy/sosana.go` - create + poll flow, retry behavior, error - mapping, LiteLLM spend logging. -- `internal/proxy/proxy_log.go` - unified upstream error logging and masking. -- `internal/proxy/upstream_masking.go` - proxy-chain Sosana masking helpers. -- `internal/proxy/image_response.go` - normalization for masked proxy image - responses. -- `docs/providers/sosana.md` - short public provider doc. -- `config.yaml.example` - example Sosana credential and model. - -## Provider Configuration - -The provider type accepts these forms: - -- `sosana` -- `sosana-art` -- `sosana_art` - -Config validation requires: - -- `api_key` -- `base_url` - -Recommended base URL: - -```yaml -credentials: - - name: "sosana_images" - type: "sosana" - api_key: "os.environ/SOSANA_API_KEY" - base_url: "https://sosana.art" -``` - -Production Sosana credentials should use request and server write timeouts of at -least 2 minutes. Sosana documents async image tasks as typically completing in -20-120 seconds. - -## Request Flow - -### Image Generation - -Client request: - -```http -POST /v1/images/generations -``` - -OpenAI-compatible JSON is converted to a Sosana Banana create request: - -```json -{ - "prompt": "...", - "model": "nano-banana", - "aspect_ratio": "1:1" -} -``` - -The model sent to Sosana is the resolved real model id from `models[].model` when -configured. This allows VSELLM public model ids to map to Sosana internal model -ids without changing the client-visible model name. - -### Image Edits - -Client request: - -```http -POST /v1/images/edits -Content-Type: multipart/form-data -``` - -The converter: - -- reads `prompt`, `model`, `size`, and `n` fields; -- reads uploaded `image`, `images`, or `image[]` file parts; -- encodes uploaded images as `data:image/...;base64,...`; -- sends them to Sosana as `image_urls`; -- rejects `mask` locally because Sosana Banana does not support masks in the - implemented flow. - -Each multipart image part is capped at 20 MB. - -### Async Sosana Flow - -The proxy does not expose Sosana async tasks to clients. It performs the async -flow synchronously inside the request: - -1. `POST {base_url}/api/banana/create-async` -2. Immediate first poll -3. `GET {base_url}/api/banana/{uid}` every 2 seconds -4. Stop on `COMPLETED`, `FAILED`, `MODERATED`, unknown terminal error, or - request deadline - -`request_timeout` is used as the polling deadline when positive. A timeout value -of `-1` means no extra deadline is added by the proxy. - -## Current Response Behavior - -On successful Sosana completion, the router downloads Sosana's -`result_file_url`, verifies that the body is an image, encodes it as base64, and -returns: - -```json -{ - "created": 1782478551, - "data": [ - { - "b64_json": "iVBORw0KGgo...", - "revised_prompt": "..." - } - ] -} -``` - -This matches the VSELLM public image generation examples, which decode -`data[].b64_json`. Sosana's public object URL remains an internal implementation -detail and is not returned to clients. - -Important implementation details: - -- The object download uses the request context and must fit into - `request_timeout`. -- The object download does not include the Sosana `Authorization` header. -- Raw image bytes are capped at 32 MiB and are never written to logs. -- `response_format: "url"` is rejected locally with 400 until VSELLM-owned - rehosting exists. - -## Local Validation - -The integration intentionally rejects unsupported OpenAI Images features before -calling Sosana: - -- `n > 1` returns local 400: `sosana supports n=1 only`. -- Missing prompt returns local 400. -- Missing image on edits returns local 400. -- `mask` on edits returns local 400. -- Non-image endpoints with a `sosana` credential return local 400: - `sosana provider supports only image generation`. -- `response_format: "url"` returns local 400 because direct Sosana storage URLs - would reveal the upstream provider. - -## Error Mapping - -Sosana details must not leak to clients. - -Client-facing mapping: - -| Condition | HTTP status | Client body | -| --- | ---: | --- | -| Local validation error | 400 | Specific local OpenAI-compatible error | -| Unsupported endpoint | 400 | `sosana provider supports only image generation` | -| Create HTTP error | Upstream status | Generic upstream OpenAI-compatible error | -| Poll HTTP error | Upstream status | Generic upstream OpenAI-compatible error | -| Transport error | 502 | Generic upstream OpenAI-compatible error | -| Transport timeout / deadline | 408 | Generic upstream OpenAI-compatible error | -| Task `FAILED` | 502 | Generic upstream OpenAI-compatible error | -| Task `MODERATED` | 400 | Neutral `content_policy_violation` | -| Task `COMPLETED` without result URL | 502 | Generic upstream OpenAI-compatible error | -| Result image download HTTP error | 502 | Generic upstream OpenAI-compatible error | -| Result image download timeout | 408 | Generic upstream OpenAI-compatible error | -| Result image too large / non-image / read error | 502 | Generic upstream OpenAI-compatible error | -| Unknown task status | 502 | Generic upstream OpenAI-compatible error | - -`FAILED` and upstream HTTP errors deliberately do not return Sosana's raw -`error`, `detail`, balance text, policy text, uid, or provider-specific codes. - -## Debug Logging And Masking - -Sosana is considered a masked upstream provider when any of these are true: - -- credential type is `sosana`; -- credential base URL host is `sosana.art` or a subdomain; -- credential name contains `sosana` or `sasana`; -- `mask_upstream_errors: true` is set. - -For masked upstream failures: - -- client receives only generic/neutral OpenAI-compatible errors; -- structured logs include `response_body_masked=true`; -- structured logs also include raw `response_body` for internal debugging. - -This mirrors the Comet API style: external users do not see provider internals, -while operators can still debug the actual upstream response. - -For proxy-chain setups, masking is reliable when: - -- the upstream router propagates `X-Credential-Name` with a Sosana marker; or -- the proxy credential has `mask_upstream_errors: true`. - -Older external proxies without a credential marker remain a configuration risk. - -## Retry And Fallback Behavior - -Sosana retry behavior is intentionally narrower than generic direct LLM -providers because image creation can be billable once a task is accepted. - -Implemented retry: - -- Retry same provider type (`sosana`) with another Sosana credential. -- Retry only when the create request returns a retryable HTTP status before a - task is accepted, for example 429 or 5xx according to the shared retry - classifier. - -No retry across credentials after task creation: - -- Poll HTTP errors are not retried against another credential. -- Task `FAILED` is not retried against another credential. -- Transport errors are not retried by the Sosana special-case because it can be - ambiguous whether the create request reached Sosana and became billable. - -This avoids duplicate image charges and duplicate generated assets. - -## Billing - -Successful Sosana image requests set: - -```go -logCtx.TokenUsage = &converter.TokenUsage{ImageCount: logCtx.ImageCount} -``` - -`ImageCount` is extracted from `n` and defaults to 1. Since Sosana currently -supports only `n=1`, successful spend should be one image per request. - -LiteLLM spend calculation uses: - -```text -spend = image_count * output_cost_per_image -``` - -Price lookup tries the real model name first, then the public alias. Production -pricing should include the public VSELLM model id and/or the Sosana real model -id used by `models[].model`. - -If future support for `n > 1` is added by spawning multiple Sosana tasks, billing -must keep `ImageCount` equal to the number of completed/generated images. - -## VSELLM Image Model Research - -VSELLM public docs and pricing list these image-generation models: - -- `google/gemini-3-pro-image-preview` -- `vertex_ai/imagen-4.0-fast-generate-001` -- `google/gemini-3.1-flash-image-preview` -- `vertex_ai/imagen-4.0-generate-001` -- `vertex_ai/imagen-4.0-ultra-generate-001` -- `google/gemini-2.5-flash-image` -- `openai/gpt-image-1-mini` -- `openai/gpt-image-1` - -VSELLM docs for `/v1/images/edits` mention these edit-capable models: - -- `openai/gpt-image-1-mini` -- `openai/gpt-image-1` -- `google/gemini-2.5-flash-image` -- `google/gemini-3-pro-image-preview` -- `google/gemini-3.1-flash-image-preview` - -Sources: - -- `https://vsellm.ru/docs` -- `https://vsellm.ru/provider/VSELLM` - -## Sosana Model Research - -Sosana Banana OpenAPI currently lists these image model ids: - -- `nano-banana` -- `nano-banana-2-1k` -- `nano-banana-2-2k` -- `nano-banana-2-2k-thinking` -- `nano-banana-pro-1k` -- `nano-banana-pro-2k` -- `nano-banana-pro-4k` -- `gpt-image-2` - -Sosana pricing also lists "Nano Banana 2 4K", but the OpenAPI enum does not -currently expose a confirmed `nano-banana-2-4k` id. Do not use that id until it -is live-tested or confirmed by Sosana. - -Sources: - -- `https://sosana.art/openapi.json` -- `https://sosana.art/` - -## Recommended VSELLM To Sosana Mapping - -Use `models[].model` for mapping. The `name` remains the public VSELLM model id, -and `model` is the Sosana model sent upstream. - -Recommended minimal production mapping: - -```yaml -models: - - name: "google/gemini-2.5-flash-image" - model: "nano-banana" - credential: sosana_images - - - name: "google/gemini-3.1-flash-image-preview" - model: "nano-banana-2-1k" - credential: sosana_images - - - name: "google/gemini-3-pro-image-preview" - model: "nano-banana-pro-1k" - credential: sosana_images -``` - -Why these map cleanly: - -- VSELLM describes `google/gemini-2.5-flash-image` as Gemini 2.5 Flash Image; - Sosana `nano-banana` is the matching Nano Banana family. -- VSELLM describes `google/gemini-3.1-flash-image-preview` as Nano Banana 2 / - Gemini 3.1 Flash Image Preview; Sosana `nano-banana-2-*` is the matching - family. -- VSELLM describes `google/gemini-3-pro-image-preview` as Nano Banana Pro / - Gemini 3 Pro Image Preview; Sosana `nano-banana-pro-*` is the matching family. - -Do not map these VSELLM models to Sosana Banana: - -- `vertex_ai/imagen-4.0-fast-generate-001` -- `vertex_ai/imagen-4.0-generate-001` -- `vertex_ai/imagen-4.0-ultra-generate-001` -- `openai/gpt-image-1-mini` -- `openai/gpt-image-1` - -Those are different model families with different quality, pricing, latency, and -behavior expectations. Substituting Sosana Banana behind those names would be a -model mismatch. - -`gpt-image-2` note: - -- Sosana exposes `gpt-image-2`. -- VSELLM did not show `openai/gpt-image-2` in the reviewed public image model - list. -- Do not expose it under an existing OpenAI model id. Add a new public model id - only as an explicit product decision. - -## 1K, 2K, 4K, And Thinking Variants - -The Sosana model suffixes should be treated as explicit cost/quality tiers, not -as hidden automatic upgrades. - -Recommended default behavior: - -- `google/gemini-3.1-flash-image-preview` -> `nano-banana-2-1k` -- `google/gemini-3-pro-image-preview` -> `nano-banana-pro-1k` - -Reasons: - -- 1K is the safest default for cost and latency. -- Silent upgrades to 2K, 4K, or thinking variants can surprise users and billing. -- Routing based on prompt complexity is subjective and hard to explain. - -If product needs higher tiers, prefer explicit public model aliases: - -```yaml -models: - - name: "google/gemini-3.1-flash-image-preview-2k" - model: "nano-banana-2-2k" - credential: sosana_images - - - name: "google/gemini-3.1-flash-image-preview-thinking" - model: "nano-banana-2-2k-thinking" - credential: sosana_images - - - name: "google/gemini-3-pro-image-preview-2k" - model: "nano-banana-pro-2k" - credential: sosana_images - - - name: "google/gemini-3-pro-image-preview-4k" - model: "nano-banana-pro-4k" - credential: sosana_images -``` - -A future `size`/`quality` based selector is possible, but it should be explicit, -tested, documented, and priced. It should not be inferred from prompt text. - -## Known Gaps Before Full VSELLM Parity - -1. VSELLM-owned URL responses: - - Default and `b64_json` responses hide Sosana storage. - - `response_format: "url"` is rejected locally. - - If URL responses are required, re-host in VSELLM-owned storage before - returning to clients. - -2. `mask` support: - - VSELLM edit docs mention optional masks. - - Current Sosana converter rejects masks. - - Do not route mask-dependent public models only to Sosana unless local 400 is - acceptable for that model. - -3. `n > 1` support: - - Current Sosana integration supports only `n=1`. - - VSELLM docs describe `n` as number of variants. - - Future support should create multiple tasks and bill by completed image - count. - -4. Tier selection: - - Current mapping is static via `models[].model`. - - 1K/2K/4K selection by `size` or `quality` is not implemented. - -## Suggested Verification - -Unit/regression tests: - -```bash -go test ./internal/converter/sosana ./internal/proxy ./internal/config ./internal/models ./internal/litellmdb/model_table -count=1 -go test ./... -count=1 -git diff --check -``` - -Live smoke test: - -1. Start router with a Sosana credential. -2. Check `/health` shows the Sosana credential and configured model. -3. Call `/v1/images/generations` with `n=1`. -4. Call `/v1/images/generations` with `n=2`; verify local 400 and no upstream - call. -5. Call `/v1/chat/completions` with a Sosana model; verify local 400 and no - upstream call. -6. Start with a bad Sosana key; verify the client response is masked and logs - contain `response_body_masked=true` plus raw `response_body` for debugging. - -For live bad-key tests, mount the config by absolute path. A missing host file in -`docker run -v` can be created as a directory by Docker, causing -`config.yaml: is a directory`. diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md index 3227614c..ca1058b1 100644 --- a/docs/providers/sosana.md +++ b/docs/providers/sosana.md @@ -33,26 +33,36 @@ completion requests. For production Sosana credentials, set the router ## Behavior - `n` must be `1`. +- Requests selected to a Sosana credential return a local 400 when they require + controls that Sosana does not support. - Default response format and `response_format: "b64_json"` return `data[].b64_json`. - `response_format: "url"` returns a local 400 because URL responses require VSELLM-owned rehosting before they can hide Sosana storage. -- `/v1/images/edits` sends uploaded images as `data:image/...;base64,...` - values in Sosana `image_urls`. +- `/v1/images/edits` accepts PNG input images only, up to 14 files, and sends + them as `data:image/png;base64,...` values in Sosana `image_urls`. - Mask images are not supported. +- Output is PNG. `output_format` may be omitted or set to `png`; other formats + are not routed to Sosana. +- The router sends `prompt_optimization: false` so Sosana returns `MODERATED` + instead of rewriting moderated prompts into safe alternatives. On completion, Sosana returns a public object URL in `result_file_url`. The -router downloads that object without forwarding the Sosana `Authorization` -header, keeps the download bounded to 32 MiB, verifies the body is an image, and -base64-encodes it into the OpenAI-compatible JSON response. The upstream object -URL is not returned to clients. +router downloads that object only from allowed Sosana/CDN hosts, does not follow +redirects, does not forward the Sosana `Authorization` header, keeps the download +bounded to 32 MiB, verifies the body is PNG, and base64-encodes it into the +OpenAI-compatible JSON response. The upstream object URL is not returned to +clients. ## Error Masking Sosana upstream HTTP errors and terminal task errors are masked before they are -returned to clients or written to structured logs. The router preserves the -appropriate HTTP status but replaces provider details with neutral -OpenAI-compatible error bodies. +returned to clients. The router preserves the appropriate HTTP status but +replaces provider details with neutral OpenAI-compatible error bodies. + +For operator debugging, structured logs may include a truncated textual upstream +error body with `response_body_masked=true`. Raw image bytes and full result +URLs are not logged. If Sosana is hidden behind another proxy credential, enable `mask_upstream_errors: true` on that proxy unless the upstream router is known to diff --git a/internal/config/config.go b/internal/config/config.go index e5719662..4a9384f2 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -432,7 +432,7 @@ type ServerConfig struct { MaxIdleConnsPerHost int `yaml:"max_idle_conns_per_host"` IdleConnTimeout time.Duration `yaml:"idle_conn_timeout"` ReadTimeout time.Duration `yaml:"-"` // HTTP server read timeout (equals request_timeout, not configurable via YAML) - WriteTimeout time.Duration `yaml:"write_timeout"` // HTTP server write timeout (default: 2m) + WriteTimeout time.Duration `yaml:"write_timeout"` // HTTP server write timeout (default: 60s) IdleTimeout time.Duration `yaml:"idle_timeout"` // HTTP server idle timeout (default: 2*write_timeout) MaxProviderRetries int `yaml:"max_provider_retries"` // Max same-type credential retries on provider errors (default: 2, meaning 3 total attempts) MaxFallbackAttempts int `yaml:"max_fallback_attempts"` // Max fallback proxy hops per request chain (default: 5) @@ -516,16 +516,16 @@ func (s *ServerConfig) UnmarshalYAML(value *yaml.Node) error { } // Duration fields - if s.RequestTimeout, err = parseField(temp.RequestTimeout, 2*time.Minute, time.ParseDuration, "request_timeout"); err != nil { + if s.RequestTimeout, err = parseField(temp.RequestTimeout, 60*time.Second, time.ParseDuration, "request_timeout"); err != nil { return err } if s.IdleConnTimeout, err = parseField(temp.IdleConnTimeout, 120*time.Second, time.ParseDuration, "idle_conn_timeout"); err != nil { return err } - if s.WriteTimeout, err = parseField(temp.WriteTimeout, 2*time.Minute, time.ParseDuration, "write_timeout"); err != nil { + if s.WriteTimeout, err = parseField(temp.WriteTimeout, 60*time.Second, time.ParseDuration, "write_timeout"); err != nil { return err } - if s.IdleTimeout, err = parseField(temp.IdleTimeout, 4*time.Minute, time.ParseDuration, "idle_timeout"); err != nil { + if s.IdleTimeout, err = parseField(temp.IdleTimeout, 2*time.Minute, time.ParseDuration, "idle_timeout"); err != nil { return err } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index b9e2e093..2262bf09 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1219,10 +1219,10 @@ monitoring: cfg, err := Load(configPath) require.NoError(t, err) - assert.Equal(t, 2*time.Minute, cfg.Server.RequestTimeout) - assert.Equal(t, 2*time.Minute, cfg.Server.ReadTimeout) - assert.Equal(t, 2*time.Minute, cfg.Server.WriteTimeout) - assert.Equal(t, 4*time.Minute, cfg.Server.IdleTimeout) + assert.Equal(t, 60*time.Second, cfg.Server.RequestTimeout) + assert.Equal(t, 60*time.Second, cfg.Server.ReadTimeout) + assert.Equal(t, 60*time.Second, cfg.Server.WriteTimeout) + assert.Equal(t, 2*time.Minute, cfg.Server.IdleTimeout) } func TestLoad_MaxProviderRetries_Custom(t *testing.T) { diff --git a/internal/converter/openai/types.go b/internal/converter/openai/types.go index a621f065..565d7c11 100644 --- a/internal/converter/openai/types.go +++ b/internal/converter/openai/types.go @@ -225,13 +225,9 @@ type OpenAIImageUsage struct { // OpenAIImageResponse represents OpenAI image response type OpenAIImageResponse struct { - Created int64 `json:"created"` - Background string `json:"background,omitempty"` - Data []OpenAIImageData `json:"data"` - OutputFormat string `json:"output_format,omitempty"` - Quality string `json:"quality,omitempty"` - Size string `json:"size,omitempty"` - Usage *OpenAIImageUsage `json:"usage,omitempty"` + Created int64 `json:"created"` + Data []OpenAIImageData `json:"data"` + Usage *OpenAIImageUsage `json:"usage,omitempty"` } // Embedding types diff --git a/internal/converter/sosana/compatibility.go b/internal/converter/sosana/compatibility.go new file mode 100644 index 00000000..1e5144c7 --- /dev/null +++ b/internal/converter/sosana/compatibility.go @@ -0,0 +1,228 @@ +package sosana + +import ( + "bytes" + "encoding/json" + "fmt" + "mime" + "mime/multipart" + "strings" +) + +const maxInputImages = 14 + +var unsupportedImageFields = []string{ + "tools", + "tool_choice", + "google_search", + "thinking_level", + "thinking_budget", + "thinking_config", + "thinking", + "reasoning_effort", + "generation_config", + "temperature", + "top_p", + "top_k", + "seed", + "max_tokens", + "stop", + "stream", + "messages", + "extra_body", + "image_size", + "image", + "images", + "image_urls", + "reference_images", +} + +// UnsupportedRequest returns a short reason when a request needs image features +// that Sosana Banana does not expose in its public API. +func UnsupportedRequest(path string, body []byte, contentType string) string { + switch { + case strings.Contains(path, "/images/generations"): + return unsupportedGenerationRequest(body) + case strings.Contains(path, "/images/edits"): + return unsupportedEditRequest(body, contentType) + default: + return "sosana provider supports only image generation" + } +} + +func unsupportedGenerationRequest(body []byte) string { + var raw map[string]json.RawMessage + if err := json.Unmarshal(body, &raw); err != nil { + return "" + } + return unsupportedImageFieldsInJSON(raw) +} + +func unsupportedEditRequest(body []byte, contentType string) string { + mediaType, params, err := mime.ParseMediaType(contentType) + if err != nil || !strings.HasPrefix(mediaType, "multipart/form-data") { + return "" + } + boundary := params["boundary"] + if boundary == "" { + return "" + } + + fields := make(map[string]string) + imageCount := 0 + reader := multipart.NewReader(bytes.NewReader(body), boundary) + for { + part, err := reader.NextPart() + if err != nil { + break + } + + formName := part.FormName() + if formName == "" { + continue + } + data, err := readLimited(part, maxMultipartImageBytes) + if err != nil { + return err.Error() + } + if part.FileName() == "" { + fields[formName] = strings.TrimSpace(string(data)) + continue + } + if formName == "mask" { + return "mask is unsupported" + } + if formName != "image" && formName != "images" && formName != "image[]" { + continue + } + imageCount++ + if detectImageMIMEType(part.Header.Get("Content-Type"), data) != "image/png" { + return "only PNG input images are supported" + } + } + if imageCount > maxInputImages { + return "too many input images" + } + return unsupportedImageFieldsInForm(fields) +} + +func unsupportedImageFieldsInJSON(raw map[string]json.RawMessage) string { + if reason := unsupportedJSONImageCount(raw["n"]); reason != "" { + return reason + } + if reason := unsupportedJSONResponseFormat(raw["response_format"]); reason != "" { + return reason + } + if reason := unsupportedJSONOutputFormat(raw["output_format"]); reason != "" { + return reason + } + for _, field := range []string{"quality", "style", "background", "moderation"} { + if hasJSONValue(raw[field]) { + return field + " is unsupported" + } + } + if hasJSONValue(raw["output_compression"]) { + return "output_compression is unsupported" + } + for _, field := range unsupportedImageFields { + if hasJSONValue(raw[field]) { + return field + " is unsupported" + } + } + return "" +} + +func unsupportedImageFieldsInForm(fields map[string]string) string { + if err := validateImageCountString(fields["n"]); err != nil { + return err.Error() + } + if err := validateResponseFormat(fields["response_format"]); err != nil { + return err.Error() + } + if err := validateOutputFormat(fields["output_format"]); err != nil { + return err.Error() + } + for _, field := range []string{"quality", "style", "background", "moderation"} { + if strings.TrimSpace(fields[field]) != "" { + return field + " is unsupported" + } + } + if _, ok := fields["output_compression"]; ok { + return "output_compression is unsupported" + } + for _, field := range unsupportedImageFields { + if strings.TrimSpace(fields[field]) != "" { + return field + " is unsupported" + } + } + return "" +} + +func unsupportedJSONImageCount(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var n int + if err := json.Unmarshal(raw, &n); err == nil { + if n == 1 { + return "" + } + return "sosana supports n=1 only" + } + var f float64 + if err := json.Unmarshal(raw, &f); err == nil { + if f == 1 { + return "" + } + return "sosana supports n=1 only" + } + return "invalid image count" +} + +func unsupportedJSONResponseFormat(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "response_format is unsupported" + } + if strings.EqualFold(strings.TrimSpace(value), "b64_json") || strings.TrimSpace(value) == "" { + return "" + } + if strings.EqualFold(strings.TrimSpace(value), "url") { + return "response_format=url is unsupported for this image model" + } + return "response_format is unsupported" +} + +func unsupportedJSONOutputFormat(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "output_format is unsupported" + } + if outputFormatAllowed(value) { + return "" + } + return "output_format is unsupported" +} + +func validateOutputFormat(raw string) error { + if outputFormatAllowed(raw) { + return nil + } + return fmt.Errorf("output_format is unsupported") +} + +func outputFormatAllowed(raw string) bool { + value := strings.ToLower(strings.TrimSpace(raw)) + return value == "" || value == "png" +} + +func hasJSONValue(raw json.RawMessage) bool { + raw = bytes.TrimSpace(raw) + return len(raw) > 0 && !bytes.Equal(raw, []byte("null")) +} diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go index ca138c95..34df6355 100644 --- a/internal/converter/sosana/images.go +++ b/internal/converter/sosana/images.go @@ -27,10 +27,11 @@ const ( ) type BananaCreateRequest struct { - Prompt string `json:"prompt"` - ImageURLs []string `json:"image_urls,omitempty"` - Model string `json:"model,omitempty"` - AspectRatio string `json:"aspect_ratio,omitempty"` + Prompt string `json:"prompt"` + ImageURLs []string `json:"image_urls,omitempty"` + Model string `json:"model,omitempty"` + AspectRatio string `json:"aspect_ratio,omitempty"` + PromptOptimization bool `json:"prompt_optimization"` } type BananaTaskResponse struct { @@ -43,8 +44,17 @@ type BananaTaskResponse struct { Error *string `json:"error"` } +type openAIImageRequest struct { + openai.OpenAIImageRequest + AspectRatio string `json:"aspect_ratio,omitempty"` +} + func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, error) { - var req openai.OpenAIImageRequest + if reason := UnsupportedRequest("/v1/images/generations", openAIBody, "application/json"); reason != "" { + return nil, fmt.Errorf("%s", reason) + } + + var req openAIImageRequest if err := json.Unmarshal(openAIBody, &req); err != nil { return nil, fmt.Errorf("failed to parse OpenAI image request: %w", err) } @@ -54,18 +64,26 @@ func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, error) { if err := validateResponseFormat(req.ResponseFormat); err != nil { return nil, err } + if err := validateOutputFormat(req.OutputFormat); err != nil { + return nil, err + } prompt := strings.TrimSpace(req.Prompt) if prompt == "" { return nil, fmt.Errorf("image generation request missing prompt") } return json.Marshal(BananaCreateRequest{ - Prompt: prompt, - Model: providerModel(modelID, req.Model), - AspectRatio: SizeToAspectRatio(req.Size), + Prompt: prompt, + Model: providerModel(modelID, req.Model), + AspectRatio: aspectRatio(req.AspectRatio, req.Size), + PromptOptimization: false, }) } func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, error) { + if reason := UnsupportedRequest("/v1/images/edits", openAIBody, contentType); reason != "" { + return nil, fmt.Errorf("%s", reason) + } + mediaType, params, err := mime.ParseMediaType(contentType) if err != nil { return nil, fmt.Errorf("failed to parse image edit content type: %w", err) @@ -109,15 +127,24 @@ func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([] continue } mimeType := detectImageMIMEType(part.Header.Get("Content-Type"), data) + if mimeType != "image/png" { + return nil, fmt.Errorf("sosana image edits support PNG images only") + } imageURLs = append(imageURLs, "data:"+mimeType+";base64,"+base64.StdEncoding.EncodeToString(data)) } + if len(imageURLs) > maxInputImages { + return nil, fmt.Errorf("sosana supports up to %d input images", maxInputImages) + } if err := validateImageCountString(fields["n"]); err != nil { return nil, err } if err := validateResponseFormat(fields["response_format"]); err != nil { return nil, err } + if err := validateOutputFormat(fields["output_format"]); err != nil { + return nil, err + } prompt := strings.TrimSpace(fields["prompt"]) if prompt == "" { return nil, fmt.Errorf("image edit request missing prompt field") @@ -126,10 +153,11 @@ func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([] return nil, fmt.Errorf("image edit request missing image") } return json.Marshal(BananaCreateRequest{ - Prompt: prompt, - ImageURLs: imageURLs, - Model: providerModel(modelID, fields["model"]), - AspectRatio: SizeToAspectRatio(fields["size"]), + Prompt: prompt, + ImageURLs: imageURLs, + Model: providerModel(modelID, fields["model"]), + AspectRatio: aspectRatio(fields["aspect_ratio"], fields["size"]), + PromptOptimization: false, }) } @@ -193,6 +221,13 @@ func SizeToAspectRatio(size string) string { } } +func aspectRatio(explicit, size string) string { + if value := strings.TrimSpace(explicit); value != "" { + return value + } + return SizeToAspectRatio(size) +} + func validateImageCount(n *int) error { if n == nil || *n == 1 { return nil @@ -236,7 +271,11 @@ func readLimited(r io.Reader, limit int64) ([]byte, error) { func detectImageMIMEType(header string, data []byte) string { header = strings.ToLower(strings.TrimSpace(header)) if strings.HasPrefix(header, "image/") { - return header + mediaType, _, err := mime.ParseMediaType(header) + if err == nil { + return mediaType + } + return strings.TrimSpace(strings.Split(header, ";")[0]) } detected := http.DetectContentType(data) if strings.HasPrefix(detected, "image/") { diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go index 1d4c2872..636643f3 100644 --- a/internal/converter/sosana/images_test.go +++ b/internal/converter/sosana/images_test.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/base64" "encoding/json" + "fmt" "mime/multipart" "strings" "testing" @@ -38,6 +39,7 @@ func TestImageGenerationRequest(t *testing.T) { assert.Equal(t, "draw a cat", req.Prompt) assert.Equal(t, "nano-banana", req.Model) assert.Equal(t, tt.wantAspect, req.AspectRatio) + assert.False(t, req.PromptOptimization) assert.Empty(t, req.ImageURLs) }) } @@ -52,6 +54,15 @@ func TestImageGenerationRequestPrefersProviderModel(t *testing.T) { assert.Equal(t, "nano-banana", req.Model) } +func TestImageGenerationRequestUsesExplicitAspectRatio(t *testing.T) { + got, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","size":"1024x1024","aspect_ratio":"16:9"}`), "nano-banana") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "16:9", req.AspectRatio) +} + func TestImageGenerationRequestRejectsMultipleImages(t *testing.T) { _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","n":2}`), "nano-banana") require.Error(t, err) @@ -64,6 +75,36 @@ func TestImageGenerationRequestRejectsURLResponseFormat(t *testing.T) { assert.Contains(t, err.Error(), "response_format=url") } +func TestImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { + tests := []struct { + name string + body string + want string + }{ + {name: "tools", body: `{"model":"nano-banana","prompt":"draw","tools":[{"type":"google_search"}]}`, want: "tools"}, + {name: "thinking", body: `{"model":"nano-banana","prompt":"draw","thinking_level":"high"}`, want: "thinking_level"}, + {name: "output format", body: `{"model":"nano-banana","prompt":"draw","output_format":"jpeg"}`, want: "output_format"}, + {name: "output compression", body: `{"model":"nano-banana","prompt":"draw","output_compression":0}`, want: "output_compression"}, + {name: "quality auto", body: `{"model":"nano-banana","prompt":"draw","quality":"auto"}`, want: "quality"}, + {name: "messages", body: `{"model":"nano-banana","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, + {name: "image size", body: `{"model":"nano-banana","prompt":"draw","image_size":"2K"}`, want: "image_size"}, + {name: "reference images", body: `{"model":"nano-banana","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := ImageGenerationRequest([]byte(tt.body), "nano-banana") + require.Error(t, err) + assert.Contains(t, err.Error(), tt.want) + }) + } +} + +func TestImageGenerationRequestAllowsPNGOutputFormat(t *testing.T) { + _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","output_format":"png","response_format":"b64_json"}`), "nano-banana") + require.NoError(t, err) +} + func TestImageEditRequest(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ "model": "nano-banana", @@ -82,6 +123,7 @@ func TestImageEditRequest(t *testing.T) { assert.Equal(t, "make it blue", req.Prompt) assert.Equal(t, "nano-banana", req.Model) assert.Equal(t, "1:1", req.AspectRatio) + assert.False(t, req.PromptOptimization) require.Len(t, req.ImageURLs, 1) assert.True(t, strings.HasPrefix(req.ImageURLs[0], "data:image/png;base64,")) assert.Contains(t, req.ImageURLs[0], base64.StdEncoding.EncodeToString(pngBytes())) @@ -145,6 +187,37 @@ func TestImageEditRequestRejectsURLResponseFormat(t *testing.T) { assert.Contains(t, err.Error(), "response_format=url") } +func TestImageEditRequestRejectsJPEGInput(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "nano-banana", + "prompt": "make it blue", + }, map[string][]byte{ + "image": jpegBytes(), + }) + + _, err := ImageEditRequest(body, contentType, "fallback-model") + require.Error(t, err) + assert.Contains(t, err.Error(), "PNG") +} + +func TestImageEditRequestRejectsTooManyImages(t *testing.T) { + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("model", "nano-banana")) + require.NoError(t, writer.WriteField("prompt", "make it blue")) + for i := 0; i < maxInputImages+1; i++ { + part, err := writer.CreateFormFile("image", fmt.Sprintf("image-%02d.png", i)) + require.NoError(t, err) + _, err = part.Write(pngBytes()) + require.NoError(t, err) + } + require.NoError(t, writer.Close()) + + _, err := ImageEditRequest(buf.Bytes(), writer.FormDataContentType(), "fallback-model") + require.Error(t, err) + assert.Contains(t, err.Error(), "too many") +} + func TestOpenAIImageResponse(t *testing.T) { createdAt := "2026-01-01T00:00:00Z" body, err := OpenAIImageResponse(BananaTaskResponse{ @@ -186,3 +259,7 @@ func multipartImageEditBody(t *testing.T, fields map[string]string, files map[st func pngBytes() []byte { return []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} } + +func jpegBytes() []byte { + return []byte{0xff, 0xd8, 0xff, 0xdb, 0, 0x43, 0, 1, 2, 3} +} diff --git a/internal/proxy/cometapi_test.go b/internal/proxy/cometapi_test.go index b48a853a..fe2f9b2c 100644 --- a/internal/proxy/cometapi_test.go +++ b/internal/proxy/cometapi_test.go @@ -6,6 +6,7 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestIsCometAPICredential(t *testing.T) { @@ -69,8 +70,11 @@ func TestAppendResponseBodyForLogs_MaskedProviderKeepsMaskedFlagAndLogsBody(t *t assert.Contains(t, args, "response_body_masked") assert.Contains(t, args, true) - assert.Contains(t, args, "response_body") - assert.Contains(t, args, body) + loggedBody := responseBodyArg(t, args) + assert.Contains(t, loggedBody, "permission_denied") + assert.Contains(t, loggedBody, "model access denied") + assert.NotEqual(t, body, loggedBody) + assert.Less(t, len(loggedBody), len(body)) } func TestIsSosanaCredential(t *testing.T) { @@ -134,6 +138,23 @@ func TestAppendResponseBodyForLogs_CometKeepsMaskedFlagAndLogsBody(t *testing.T) assert.Contains(t, args, "response_body_masked") assert.Contains(t, args, true) - assert.Contains(t, args, "response_body") - assert.Contains(t, args, body) + loggedBody := responseBodyArg(t, args) + assert.Contains(t, loggedBody, "permission_denied") + assert.Contains(t, loggedBody, "model access denied") + assert.NotEqual(t, body, loggedBody) + assert.Less(t, len(loggedBody), len(body)) +} + +func responseBodyArg(t *testing.T, args []any) string { + t.Helper() + + for i := 0; i < len(args)-1; i++ { + if args[i] == "response_body" { + body, ok := args[i+1].(string) + require.True(t, ok) + return body + } + } + require.FailNow(t, "response_body arg not found") + return "" } diff --git a/internal/proxy/image_response.go b/internal/proxy/image_response.go deleted file mode 100644 index 0df4f682..00000000 --- a/internal/proxy/image_response.go +++ /dev/null @@ -1,102 +0,0 @@ -package proxy - -import ( - "bytes" - "encoding/json" - "fmt" - "time" - - "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" - "github.com/mixaill76/auto_ai_router/internal/converter/openai" -) - -type sosanaTaskResponse struct { - Status string `json:"status"` - CreatedAt string `json:"created_at"` - ResultFileURL string `json:"result_file_url"` - OptimizedPrompt string `json:"optimized_prompt"` - Error json.RawMessage `json:"error"` -} - -func normalizeOpenAIImageResponseBody(body []byte) ([]byte, error) { - var envelope struct { - Error json.RawMessage `json:"error"` - } - if err := json.Unmarshal(body, &envelope); err != nil { - return nil, fmt.Errorf("parse image response: %w", err) - } - - var resp openai.OpenAIImageResponse - if err := json.Unmarshal(body, &resp); err != nil { - return nil, fmt.Errorf("parse OpenAI image response: %w", err) - } - if len(resp.Data) == 0 { - if sosanaResp, ok, err := openAIImageResponseFromSosanaTask(body); err != nil { - return nil, err - } else if ok { - resp = sosanaResp - } else { - if hasJSONValue(envelope.Error) { - return nil, fmt.Errorf("image response contains upstream error envelope") - } - return nil, fmt.Errorf("image response contains no image data") - } - } - for i, item := range resp.Data { - if item.B64JSON == "" && item.URL == "" { - return nil, fmt.Errorf("image response item %d contains neither b64_json nor url", i) - } - } - if resp.Created == 0 { - resp.Created = converterutil.GetCurrentTimestamp() - } - - normalized, err := json.Marshal(resp) - if err != nil { - return nil, fmt.Errorf("marshal OpenAI image response: %w", err) - } - return append(normalized, '\n'), nil -} - -func openAIImageResponseFromSosanaTask(body []byte) (openai.OpenAIImageResponse, bool, error) { - var task sosanaTaskResponse - if err := json.Unmarshal(body, &task); err != nil { - return openai.OpenAIImageResponse{}, false, nil - } - - switch task.Status { - case "": - return openai.OpenAIImageResponse{}, false, nil - case "COMPLETED": - if task.ResultFileURL == "" { - return openai.OpenAIImageResponse{}, true, fmt.Errorf("image task completed response contains no result URL") - } - resp := openai.OpenAIImageResponse{ - Created: parseSosanaCreatedAt(task.CreatedAt), - Data: []openai.OpenAIImageData{{ - URL: task.ResultFileURL, - RevisedPrompt: task.OptimizedPrompt, - }}, - } - return resp, true, nil - case "PROCESSING", "FAILED", "MODERATED": - return openai.OpenAIImageResponse{}, true, fmt.Errorf("image task response is not a completed image result") - default: - return openai.OpenAIImageResponse{}, true, fmt.Errorf("image task response contains an unknown status") - } -} - -func parseSosanaCreatedAt(value string) int64 { - if value == "" { - return converterutil.GetCurrentTimestamp() - } - if ts, err := time.Parse(time.RFC3339Nano, value); err == nil { - return ts.Unix() - } - return converterutil.GetCurrentTimestamp() -} - -func hasJSONValue(raw json.RawMessage) bool { - raw = bytes.TrimSpace(raw) - return len(raw) > 0 && !bytes.Equal(raw, []byte("null")) -} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index e5e227f2..216679c8 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -763,24 +763,6 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { p.setSessionBinding(logCtx.SessionID, modelID, cred.Name) } } else { - if proxyResp.StatusCode >= 200 && proxyResp.StatusCode < 300 && logCtx.IsImageGeneration && shouldMaskProxyResponseErrors(cred, proxyResp) { - normalizedBody, normErr := normalizeOpenAIImageResponseBody(proxyResp.Body) - if normErr != nil { - p.logger.ErrorContext(r.Context(), "Failed to normalize proxy image response to OpenAI format", - "credential", cred.Name, "model", modelID, "error", normErr, - "actual_credential", logCtx.ActualCredentialName, - "response_body_masked", true, - "request_id", logCtx.RequestID) - proxyResp.StatusCode = http.StatusBadGateway - maskProxyErrorResponse(proxyResp) - } else { - proxyResp.Body = normalizedBody - if proxyResp.Headers == nil { - proxyResp.Headers = http.Header{} - } - proxyResp.Headers.Set("Content-Type", "application/json") - } - } if proxyResp.StatusCode >= 400 && shouldMaskProxyResponseErrors(cred, proxyResp) { maskProxyErrorResponse(proxyResp) } @@ -1448,25 +1430,6 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { } } - if resp.StatusCode >= 200 && resp.StatusCode < 300 && logCtx.IsImageGeneration && shouldMaskUpstreamErrors(cred) { - normalizedBody, normErr := normalizeOpenAIImageResponseBody(finalResponseBody) - if normErr != nil { - args := []any{ - "credential", cred.Name, "provider", string(cred.Type), - "model", modelID, "error", normErr, - "request_id", logCtx.RequestID, - } - args = appendResponseBodyForLogs(args, cred, string(finalResponseBody)) - p.logger.ErrorContext(r.Context(), "Failed to normalize image response to OpenAI format", args...) - resp.StatusCode = http.StatusBadGateway - finalResponseBody = maskedUpstreamErrorBody(resp.StatusCode) - } else { - finalResponseBody = normalizedBody - } - bodyForTokenExtraction = finalResponseBody - resp.Header.Set("Content-Type", "application/json") - } - rawErrorBody := finalResponseBody if resp.StatusCode >= 400 && shouldMaskUpstreamErrors(cred) { finalResponseBody = maskedUpstreamErrorBody(resp.StatusCode) diff --git a/internal/proxy/proxy_log.go b/internal/proxy/proxy_log.go index ac5228c1..773055a7 100644 --- a/internal/proxy/proxy_log.go +++ b/internal/proxy/proxy_log.go @@ -42,7 +42,7 @@ func appendResponseBodyForLogs(args []any, cred *config.CredentialConfig, body s if shouldMaskUpstreamErrors(cred) { return append(args, "response_body_masked", true, - "response_body", body, + "response_body", logger.TruncateLongFields(body, 500), ) } return append(args, "response_body", logger.TruncateLongFields(body, 500)) diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index f1a8501b..edae47b1 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -7,6 +7,7 @@ import ( "errors" "io" "math/rand" + "net" "net/http" "net/url" "strings" @@ -23,6 +24,8 @@ const ( maxSosanaResultErrorBytes int64 = 16 * 1024 ) +var allowPrivateSosanaResultURLForTests func(*url.URL) bool + type sosanaAttemptResult struct { body []byte statusCode int @@ -293,6 +296,13 @@ func (p *Proxy) downloadSosanaResultImage( "error", err) return nil, "", http.StatusBadGateway, err } + if err := validateSosanaResultURL(ctx, parsed); err != nil { + p.logUpstreamError(ctx, "Sosana completed task returned unsafe result URL", http.StatusBadGateway, cred, modelID, rawTaskBody, + "result_host", parsed.Hostname(), + "request_id", logCtx.RequestID, + "error", err) + return nil, "", http.StatusBadGateway, err + } req, err := http.NewRequestWithContext(ctx, http.MethodGet, resultURL, nil) if err != nil { @@ -304,7 +314,7 @@ func (p *Proxy) downloadSosanaResultImage( } req.Header.Set("Accept", "image/*") - resp, err := p.client.Do(req) + resp, err := p.doSosanaResultImageRequest(req) if err != nil { statusCode := http.StatusBadGateway if isTimeoutError(err) || errors.Is(ctx.Err(), context.DeadlineExceeded) { @@ -341,21 +351,29 @@ func (p *Proxy) downloadSosanaResultImage( return nil, "", http.StatusBadGateway, err } sniffedType := http.DetectContentType(image) - if !isImageContentType(contentType) && !isImageContentType(sniffedType) { + if !isPNGContentType(contentType) && !isPNGContentType(sniffedType) { responseBody := textBodyPrefixForSosanaResultLog(image, contentType, sniffedType) - p.logUpstreamError(ctx, "Sosana result URL returned non-image content", http.StatusBadGateway, cred, modelID, responseBody, + p.logUpstreamError(ctx, "Sosana result URL returned non-PNG content", http.StatusBadGateway, cred, modelID, responseBody, "result_host", parsed.Hostname(), "content_type", contentType, "sniffed_content_type", sniffedType, "request_id", logCtx.RequestID) - return nil, "", http.StatusBadGateway, errors.New("sosana result URL returned non-image content") + return nil, "", http.StatusBadGateway, errors.New("sosana result URL returned non-PNG content") } - if !isImageContentType(contentType) { + if !isPNGContentType(contentType) { contentType = sniffedType } return image, contentType, http.StatusOK, nil } +func (p *Proxy) doSosanaResultImageRequest(req *http.Request) (*http.Response, error) { + client := *p.client + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + } + return client.Do(req) +} + func parseSosanaResultURL(raw string) (*url.URL, error) { if raw == "" { return nil, errors.New("sosana task completed without result_file_url") @@ -370,6 +388,72 @@ func parseSosanaResultURL(raw string) (*url.URL, error) { return parsed, nil } +func validateSosanaResultURL(ctx context.Context, parsed *url.URL) error { + if allowPrivateSosanaResultURLForTests != nil && allowPrivateSosanaResultURLForTests(parsed) { + return nil + } + if parsed.Scheme != "https" { + return errors.New("sosana result_file_url must use https") + } + + host := parsed.Hostname() + if strings.EqualFold(host, "localhost") { + return errors.New("sosana result_file_url host is not allowed") + } + if !isAllowedSosanaResultHost(host) { + return errors.New("sosana result_file_url host is not allowed") + } + if ip := net.ParseIP(host); ip != nil { + if isUnsafeSosanaResultIP(ip) { + return errors.New("sosana result_file_url resolves to a private address") + } + return nil + } + + addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) + if err != nil { + return err + } + if len(addrs) == 0 { + return errors.New("sosana result_file_url host has no addresses") + } + for _, addr := range addrs { + if isUnsafeSosanaResultIP(addr.IP) { + return errors.New("sosana result_file_url resolves to a private address") + } + } + return nil +} + +func isAllowedSosanaResultHost(host string) bool { + host = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(host)), ".") + for _, suffix := range []string{ + "sosana.blog", + "sosana.art", + "storage.yandexcloud.net", + } { + if host == suffix || strings.HasSuffix(host, "."+suffix) { + return true + } + } + return false +} + +func isUnsafeSosanaResultIP(ip net.IP) bool { + if ip == nil { + return true + } + if v4 := ip.To4(); v4 != nil && v4[0] == 100 && v4[1]&0xc0 == 64 { + return true + } + return ip.IsLoopback() || + ip.IsPrivate() || + ip.IsLinkLocalUnicast() || + ip.IsLinkLocalMulticast() || + ip.IsMulticast() || + ip.IsUnspecified() +} + func sosanaResultHost(task sosana.BananaTaskResponse) string { if task.ResultFileURL == nil { return "" @@ -418,8 +502,9 @@ func textBodyPrefixForSosanaResultLog(body []byte, contentTypes ...string) []byt return nil } -func isImageContentType(contentType string) bool { - return strings.HasPrefix(strings.ToLower(strings.TrimSpace(contentType)), "image/") +func isPNGContentType(contentType string) bool { + contentType = strings.ToLower(strings.TrimSpace(contentType)) + return contentType == "image/png" || strings.HasPrefix(contentType, "image/png;") } func isTextContentType(contentType string) bool { diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index de5932ef..8ae9157b 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -8,8 +8,10 @@ import ( "fmt" "log/slog" "mime/multipart" + "net" "net/http" "net/http/httptest" + "net/url" "os" "strings" "testing" @@ -17,6 +19,7 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/mixaill76/auto_ai_router/internal/converter/sosana" "github.com/mixaill76/auto_ai_router/internal/litellmdb" litellmmodels "github.com/mixaill76/auto_ai_router/internal/litellmdb/models" aimodels "github.com/mixaill76/auto_ai_router/internal/models" @@ -40,6 +43,7 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { assert.Equal(t, "draw a fox", req["prompt"]) assert.Equal(t, "nano-banana", req["model"]) assert.Equal(t, "1:1", req["aspect_ratio"]) + assert.Equal(t, false, req["prompt_optimization"]) _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox"}`)) case "/api/banana/task-1": pollSeen = true @@ -86,6 +90,7 @@ func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) assert.Equal(t, "make it blue", req["prompt"]) assert.Equal(t, "nano-banana", req["model"]) + assert.Equal(t, false, req["prompt_optimization"]) imageURLs, ok := req["image_urls"].([]any) require.True(t, ok) require.Len(t, imageURLs, 1) @@ -310,6 +315,54 @@ func TestProxyRequest_SosanaRejectsURLResponseFormat(t *testing.T) { assert.False(t, called) } +func TestProxyRequest_IncompatibleImageRequestWithSosanaReturnsLocalError(t *testing.T) { + called := false + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, nil) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","tools":[{"type":"google_search"}]}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), "tools is unsupported") + assert.False(t, called) +} + +func TestProxyRequest_IncompatibleImageEditJPEGWithSosanaReturnsLocalError(t *testing.T) { + called := false + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "nano-banana", + "prompt": "make it blue", + }, map[string][]byte{ + "image": {0xff, 0xd8, 0xff, 0xdb, 0, 0x43}, + }) + prx := newSosanaTestProxy(upstream.URL, nil) + req := httptest.NewRequest("POST", "/v1/images/edits", bytes.NewReader(body)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", contentType) + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), "PNG input") + assert.False(t, called) +} + func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { var logBuf bytes.Buffer upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -503,12 +556,140 @@ func TestProxyRequest_SosanaImageResultNonImageMasked(t *testing.T) { assert.Contains(t, w.Body.String(), "Upstream provider error") assert.NotContains(t, w.Body.String(), "not an image") assert.NotContains(t, w.Body.String(), imageServer.URL) - assert.Contains(t, logBuf.String(), "non-image content") + assert.Contains(t, logBuf.String(), "non-PNG content") assert.Contains(t, logBuf.String(), "response_body_masked=true") assert.Contains(t, logBuf.String(), "not an image secret marker") } +func TestProxyRequest_SosanaImageResultJPEGMasked(t *testing.T) { + var logBuf bytes.Buffer + imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/jpeg", []byte{0xff, 0xd8, 0xff, 0xdb, 0, 0x43}, nil) + defer imageServer.Close() + + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/image.jpg") + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), imageServer.URL) + assert.Contains(t, logBuf.String(), "non-PNG content") + assert.NotContains(t, logBuf.String(), imageServer.URL) +} + +func TestProxyRequest_SosanaImageResultRedirectMaskedAndNotFollowed(t *testing.T) { + allowPrivateSosanaResultURLsForTest(t) + + var logBuf bytes.Buffer + targetCalled := false + targetServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + targetCalled = true + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(sosanaResultPNG) + })) + defer targetServer.Close() + + redirectServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, targetServer.URL+"/private.png", http.StatusFound) + })) + defer redirectServer.Close() + + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, redirectServer.URL+"/redirect") + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, &logBuf) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.False(t, targetCalled) + assert.NotContains(t, w.Body.String(), targetServer.URL) +} + +func TestDownloadSosanaResultImageRejectsUnsafeProductionURL(t *testing.T) { + var logBuf bytes.Buffer + called := false + imageServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + })) + defer imageServer.Close() + + resultURL := imageServer.URL + "/private.png" + prx := newSosanaTestProxy("https://sosana.art", &logBuf) + cred := &config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: "https://sosana.art"} + image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "nano-banana", sosana.BananaTaskResponse{ + Status: sosana.StatusCompleted, + ResultFileURL: &resultURL, + }, nil, &RequestLogContext{}) + + require.Error(t, err) + assert.Nil(t, image) + assert.Equal(t, http.StatusBadGateway, statusCode) + assert.False(t, called) + assert.Contains(t, logBuf.String(), "unsafe result URL") + assert.Contains(t, logBuf.String(), "result_host=127.0.0.1") +} + +func TestValidateSosanaResultURLBlocksUnsafeHosts(t *testing.T) { + tests := []string{ + "http://main-r2.sosana.blog/image.png", + "https://localhost/image.png", + "https://127.0.0.1/image.png", + "https://169.254.169.254/latest/meta-data", + "https://100.64.0.1/image.png", + "https://example.com/image.png", + } + + for _, rawURL := range tests { + t.Run(rawURL, func(t *testing.T) { + parsed, err := parseSosanaResultURL(rawURL) + require.NoError(t, err) + require.Error(t, validateSosanaResultURL(context.Background(), parsed)) + }) + } +} + +func TestValidateSosanaResultURLAllowsLocalOnlyWithTestHook(t *testing.T) { + allowPrivateSosanaResultURLsForTest(t) + + parsed, err := parseSosanaResultURL("http://127.0.0.1/image.png") + require.NoError(t, err) + require.NoError(t, validateSosanaResultURL(context.Background(), parsed)) +} + func TestProxyRequest_SosanaImageResultTimeoutMasked(t *testing.T) { + allowPrivateSosanaResultURLsForTest(t) + var logBuf bytes.Buffer imageServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(50 * time.Millisecond) @@ -663,6 +844,7 @@ var sosanaResultPNG = []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, func newSosanaResultImageServer(t *testing.T, status int, contentType string, body []byte, auths *[]string) *httptest.Server { t.Helper() + allowPrivateSosanaResultURLsForTest(t) return newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if auths != nil { @@ -676,6 +858,26 @@ func newSosanaResultImageServer(t *testing.T, status int, contentType string, bo })) } +func allowPrivateSosanaResultURLsForTest(t *testing.T) { + t.Helper() + + previous := allowPrivateSosanaResultURLForTests + allowPrivateSosanaResultURLForTests = func(parsed *url.URL) bool { + if parsed.Scheme != "http" { + return false + } + host := parsed.Hostname() + if strings.EqualFold(host, "localhost") { + return true + } + ip := net.ParseIP(host) + return ip != nil && isUnsafeSosanaResultIP(ip) + } + t.Cleanup(func() { + allowPrivateSosanaResultURLForTests = previous + }) +} + func newSosanaTestProxy(baseURL string, logBuf *bytes.Buffer) *Proxy { logger := slog.New(slog.NewTextHandler(&bytes.Buffer{}, &slog.HandlerOptions{Level: slog.LevelDebug})) if logBuf != nil { diff --git a/internal/proxy/upstream_masking_test.go b/internal/proxy/upstream_masking_test.go index 9a2a9aae..9edf980a 100644 --- a/internal/proxy/upstream_masking_test.go +++ b/internal/proxy/upstream_masking_test.go @@ -51,192 +51,6 @@ func TestMaskedUpstreamError_DirectImageErrorDoesNotLeakProviderBody(t *testing. assert.Equal(t, "upstream_rate_limit", *got.Error.Code) } -func TestMaskedUpstreamImageSuccess_IsNormalizedToOpenAIShape(t *testing.T) { - upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{ - "created":123, - "provider":"Sasana", - "background":"transparent", - "output_format":"png", - "quality":"high", - "size":"1024x1024", - "error":{"message":"Sasana warning that should not leak"}, - "data":[{"b64_json":"aW1hZ2U=","sasana_id":"vendor-image-id"}] - }`)) - })) - defer upstream.Close() - - prx := NewTestProxyBuilder(). - WithCredentials(config.CredentialConfig{ - Name: "sosana-art", - Type: config.ProviderTypeOpenAI, - BaseURL: upstream.URL, - APIKey: "upstream-key", - RPM: 100, - TPM: 10000, - MaskUpstreamErrors: true, - }). - Build() - - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-1","prompt":"cat"}`)) - req.Header.Set("Authorization", "Bearer master-key") - req.Header.Set("Content-Type", "application/json") - w := httptest.NewRecorder() - - prx.ProxyRequest(w, req) - - require.Equal(t, http.StatusOK, w.Code) - assert.NotContains(t, w.Body.String(), "Sasana") - assert.NotContains(t, w.Body.String(), "sasana_id") - assert.NotContains(t, w.Body.String(), "warning") - - var got struct { - Created int64 `json:"created"` - Background string `json:"background"` - OutputFormat string `json:"output_format"` - Quality string `json:"quality"` - Size string `json:"size"` - Data []struct { - B64JSON string `json:"b64_json"` - } `json:"data"` - } - require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) - assert.Equal(t, int64(123), got.Created) - assert.Equal(t, "transparent", got.Background) - assert.Equal(t, "png", got.OutputFormat) - assert.Equal(t, "high", got.Quality) - assert.Equal(t, "1024x1024", got.Size) - require.Len(t, got.Data, 1) - assert.Equal(t, "aW1hZ2U=", got.Data[0].B64JSON) -} - -func TestMaskedUpstreamImageSuccess_WithErrorEnvelopeBecomesRouterError(t *testing.T) { - upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"error":{"message":"Sasana moderation failed","code":"SASANA_POLICY"}}`)) - })) - defer upstream.Close() - - prx := NewTestProxyBuilder(). - WithCredentials(config.CredentialConfig{ - Name: "sosana-art", - Type: config.ProviderTypeOpenAI, - BaseURL: upstream.URL, - APIKey: "upstream-key", - RPM: 100, - TPM: 10000, - MaskUpstreamErrors: true, - }). - Build() - - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-1","prompt":"cat"}`)) - req.Header.Set("Authorization", "Bearer master-key") - req.Header.Set("Content-Type", "application/json") - w := httptest.NewRecorder() - - prx.ProxyRequest(w, req) - - require.Equal(t, http.StatusBadGateway, w.Code) - assert.NotContains(t, w.Body.String(), "Sasana") - assert.NotContains(t, w.Body.String(), "SASANA_POLICY") - assert.Contains(t, w.Body.String(), "Upstream provider error") -} - -func TestMaskedUpstreamSosanaCompletedResponse_IsConvertedToOpenAIShape(t *testing.T) { - upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{ - "uid":"01977614-95b0-7000-8000-example", - "status":"COMPLETED", - "created_at":"2026-06-26T12:34:56Z", - "prompt":"cat", - "optimized_prompt":"A detailed cat illustration", - "result_file_url":"https://cdn.sosana.art/results/cat.png", - "elapsed":12.3, - "provider":"Sosana" - }`)) - })) - defer upstream.Close() - - prx := NewTestProxyBuilder(). - WithCredentials(config.CredentialConfig{ - Name: "vsellm-sosana-art", - Type: config.ProviderTypeOpenAI, - BaseURL: upstream.URL, - APIKey: "upstream-key", - RPM: 100, - TPM: 10000, - MaskUpstreamErrors: true, - }). - Build() - - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-2","prompt":"cat"}`)) - req.Header.Set("Authorization", "Bearer master-key") - req.Header.Set("Content-Type", "application/json") - w := httptest.NewRecorder() - - prx.ProxyRequest(w, req) - - require.Equal(t, http.StatusOK, w.Code) - assert.NotContains(t, w.Body.String(), "Sosana") - assert.NotContains(t, w.Body.String(), "uid") - assert.NotContains(t, w.Body.String(), "result_file_url") - assert.NotContains(t, w.Body.String(), "status") - - var got struct { - Created int64 `json:"created"` - Data []struct { - URL string `json:"url"` - RevisedPrompt string `json:"revised_prompt"` - } `json:"data"` - } - require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got)) - assert.Equal(t, int64(1782477296), got.Created) - require.Len(t, got.Data, 1) - assert.Equal(t, "https://cdn.sosana.art/results/cat.png", got.Data[0].URL) - assert.Equal(t, "A detailed cat illustration", got.Data[0].RevisedPrompt) -} - -func TestMaskedUpstreamSosanaProcessingResponse_BecomesRouterError(t *testing.T) { - upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{ - "uid":"01977614-95b0-7000-8000-example", - "status":"PROCESSING", - "created_at":"2026-06-26T12:34:56Z", - "prompt":"cat", - "provider":"Sosana" - }`)) - })) - defer upstream.Close() - - prx := NewTestProxyBuilder(). - WithCredentials(config.CredentialConfig{ - Name: "vsellm-sosana-art", - Type: config.ProviderTypeOpenAI, - BaseURL: upstream.URL, - APIKey: "upstream-key", - RPM: 100, - TPM: 10000, - MaskUpstreamErrors: true, - }). - Build() - - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"gpt-image-2","prompt":"cat"}`)) - req.Header.Set("Authorization", "Bearer master-key") - req.Header.Set("Content-Type", "application/json") - w := httptest.NewRecorder() - - prx.ProxyRequest(w, req) - - require.Equal(t, http.StatusBadGateway, w.Code) - assert.NotContains(t, w.Body.String(), "Sosana") - assert.NotContains(t, w.Body.String(), "PROCESSING") - assert.NotContains(t, w.Body.String(), "uid") - assert.Contains(t, w.Body.String(), "Upstream provider error") -} - func TestMaskedUpstreamError_ProxyChainActualCredentialDoesNotLeakProviderBody(t *testing.T) { upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") @@ -288,9 +102,3 @@ func TestMaskedUpstreamError_ProxyChainStreamingErrorDoesNotLeakProviderBody(t * assert.NotContains(t, w.Body.String(), "Sasana") assert.Contains(t, w.Body.String(), "Upstream provider error") } - -func TestNormalizeOpenAIImageResponseBodyRejectsInvalidSuccess(t *testing.T) { - _, err := normalizeOpenAIImageResponseBody([]byte(`{"created":123,"data":[{"revised_prompt":"only text"}]}`)) - require.Error(t, err) - assert.Contains(t, err.Error(), "neither b64_json nor url") -} From 59d4c48de67e8447db4297b2f679ccdfe7dc6e1b Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Tue, 7 Jul 2026 17:25:37 +0300 Subject: [PATCH 07/16] fix: harden sosana routing and masking --- config.yaml.example | 4 +- docs/providers/sosana.md | 32 +++- internal/converter/sosana/compatibility.go | 90 ++++++++- internal/converter/sosana/images.go | 174 +++++++++++++---- internal/converter/sosana/images_test.go | 97 ++++++---- internal/proxy/proxy.go | 4 + internal/proxy/sosana.go | 101 +++++++++- internal/proxy/sosana_live_test.go | 2 +- internal/proxy/sosana_routing.go | 156 +++++++++++++++ internal/proxy/sosana_test.go | 211 ++++++++++++++++++--- 10 files changed, 742 insertions(+), 129 deletions(-) create mode 100644 internal/proxy/sosana_routing.go diff --git a/config.yaml.example b/config.yaml.example index 441f5fcc..e4bc1860 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -118,7 +118,9 @@ models: tpm: 50000 # Sosana.art image model exposed through /v1/images/generations and /v1/images/edits. - - name: "nano-banana" + # image_size selects banana-2-1k/2k/4k-compliant dynamically; default is 1K. + - name: "google/gemini-3.1-flash-image-preview" + model: "banana-2-{image_size}-compliant" credential: sosana_images rpm: 60 tpm: -1 diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md index ca1058b1..f4fa628e 100644 --- a/docs/providers/sosana.md +++ b/docs/providers/sosana.md @@ -20,12 +20,20 @@ credentials: tpm: -1 models: - - name: "nano-banana" + - name: "google/gemini-3.1-flash-image-preview" + model: "banana-2-{image_size}-compliant" credential: sosana_images rpm: 60 tpm: -1 ``` +The dynamic model template maps `image_size` to Sosana's concrete image models: +`banana-2-1k-compliant`, `banana-2-2k-compliant`, and +`banana-2-4k-compliant`. If `image_size` is omitted, the router uses `1K`. +This integration maps Sosana only for `google/gemini-3.1-flash-image-preview`; +other image families should be served by their native providers or fallback +proxies. + Sosana Banana tasks are asynchronous and can take longer than short chat completion requests. For production Sosana credentials, set the router `request_timeout` and HTTP `write_timeout` to at least `2m`. @@ -33,12 +41,17 @@ completion requests. For production Sosana credentials, set the router ## Behavior - `n` must be `1`. -- Requests selected to a Sosana credential return a local 400 when they require - controls that Sosana does not support. +- Requests selected to a Sosana credential are skipped when they require + controls that Sosana does not support. The router then tries another primary + credential for the same model and then the configured fallback proxy cascade. + If no compatible provider is available, the router returns a local 400. - Default response format and `response_format: "b64_json"` return `data[].b64_json`. -- `response_format: "url"` returns a local 400 because URL responses require - VSELLM-owned rehosting before they can hide Sosana storage. +- `response_format: "url"` is not routed to Sosana because URL responses require + VSELLM-owned rehosting before they can hide Sosana storage. Another provider + may handle it through normal fallback routing. +- `image_size` may be omitted or set to `1K`, `2K`, or `4K`; `0.5K` is not + routed to Sosana. - `/v1/images/edits` accepts PNG input images only, up to 14 files, and sends them as `data:image/png;base64,...` values in Sosana `image_urls`. - Mask images are not supported. @@ -54,6 +67,15 @@ bounded to 32 MiB, verifies the body is PNG, and base64-encodes it into the OpenAI-compatible JSON response. The upstream object URL is not returned to clients. +## Billing + +Sosana vendor prices are not used at request time and are not returned to +clients. Successful image requests log `ImageCount=1`; spend is calculated from +the internal price registry or LiteLLM model table using `output_cost_per_image`. +When `image_size` selects a concrete tier, the spend lookup uses the concrete +model first, for example `banana-2-2k-compliant`, and then falls back to the +public model name if no concrete price is configured. + ## Error Masking Sosana upstream HTTP errors and terminal task errors are masked before they are diff --git a/internal/converter/sosana/compatibility.go b/internal/converter/sosana/compatibility.go index 1e5144c7..453c1adc 100644 --- a/internal/converter/sosana/compatibility.go +++ b/internal/converter/sosana/compatibility.go @@ -30,7 +30,6 @@ var unsupportedImageFields = []string{ "stream", "messages", "extra_body", - "image_size", "image", "images", "image_urls", @@ -46,10 +45,17 @@ func UnsupportedRequest(path string, body []byte, contentType string) string { case strings.Contains(path, "/images/edits"): return unsupportedEditRequest(body, contentType) default: - return "sosana provider supports only image generation" + return "endpoint is unsupported" } } +func UnsupportedModel(modelID string) string { + if supportedSosanaModel(modelID) { + return "" + } + return "model is unsupported" +} + func unsupportedGenerationRequest(body []byte) string { var raw map[string]json.RawMessage if err := json.Unmarshal(body, &raw); err != nil { @@ -116,6 +122,12 @@ func unsupportedImageFieldsInJSON(raw map[string]json.RawMessage) string { if reason := unsupportedJSONOutputFormat(raw["output_format"]); reason != "" { return reason } + if reason := unsupportedJSONImageSize(raw["image_size"]); reason != "" { + return reason + } + if reason := unsupportedJSONExactSize(raw["size"]); reason != "" { + return reason + } for _, field := range []string{"quality", "style", "background", "moderation"} { if hasJSONValue(raw[field]) { return field + " is unsupported" @@ -142,6 +154,12 @@ func unsupportedImageFieldsInForm(fields map[string]string) string { if err := validateOutputFormat(fields["output_format"]); err != nil { return err.Error() } + if reason := unsupportedFormImageSize(fields["image_size"]); reason != "" { + return reason + } + if reason := unsupportedFormExactSize(fields["size"]); reason != "" { + return reason + } for _, field := range []string{"quality", "style", "background", "moderation"} { if strings.TrimSpace(fields[field]) != "" { return field + " is unsupported" @@ -167,14 +185,14 @@ func unsupportedJSONImageCount(raw json.RawMessage) string { if n == 1 { return "" } - return "sosana supports n=1 only" + return "image requests support n=1 only" } var f float64 if err := json.Unmarshal(raw, &f); err == nil { if f == 1 { return "" } - return "sosana supports n=1 only" + return "image requests support n=1 only" } return "invalid image count" } @@ -210,6 +228,70 @@ func unsupportedJSONOutputFormat(raw json.RawMessage) string { return "output_format is unsupported" } +func unsupportedJSONImageSize(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "image_size is unsupported" + } + if _, ok := normalizeImageSize(value); ok { + return "" + } + return "image_size is unsupported" +} + +func unsupportedJSONExactSize(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "size is unsupported" + } + if _, ok := imageSizeFromExactSize(value); ok { + return "" + } + return "size is unsupported" +} + +func unsupportedFormImageSize(raw string) string { + if strings.TrimSpace(raw) == "" { + return "" + } + if _, ok := normalizeImageSize(raw); ok { + return "" + } + return "image_size is unsupported" +} + +func unsupportedFormExactSize(raw string) string { + if strings.TrimSpace(raw) == "" { + return "" + } + if _, ok := imageSizeFromExactSize(raw); ok { + return "" + } + return "size is unsupported" +} + +func supportedSosanaModel(modelID string) bool { + model := strings.ToLower(strings.TrimSpace(modelID)) + if model == "" || model == "google/gemini-3.1-flash-image-preview" { + return true + } + if model == "banana-2-{image_size}-compliant" { + return true + } + switch model { + case "banana-2-1k-compliant", "banana-2-2k-compliant", "banana-2-4k-compliant": + return true + default: + return false + } +} + func validateOutputFormat(raw string) error { if outputFormatAllowed(raw) { return nil diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go index 34df6355..8aea3306 100644 --- a/internal/converter/sosana/images.go +++ b/internal/converter/sosana/images.go @@ -31,6 +31,7 @@ type BananaCreateRequest struct { ImageURLs []string `json:"image_urls,omitempty"` Model string `json:"model,omitempty"` AspectRatio string `json:"aspect_ratio,omitempty"` + ImageSize string `json:"image_size,omitempty"` PromptOptimization bool `json:"prompt_optimization"` } @@ -47,53 +48,68 @@ type BananaTaskResponse struct { type openAIImageRequest struct { openai.OpenAIImageRequest AspectRatio string `json:"aspect_ratio,omitempty"` + Ratio string `json:"ratio,omitempty"` + ImageSize string `json:"image_size,omitempty"` } -func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, error) { +func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, string, error) { if reason := UnsupportedRequest("/v1/images/generations", openAIBody, "application/json"); reason != "" { - return nil, fmt.Errorf("%s", reason) + return nil, "", fmt.Errorf("%s", reason) } var req openAIImageRequest if err := json.Unmarshal(openAIBody, &req); err != nil { - return nil, fmt.Errorf("failed to parse OpenAI image request: %w", err) + return nil, "", fmt.Errorf("failed to parse OpenAI image request: %w", err) } if err := validateImageCount(req.N); err != nil { - return nil, err + return nil, "", err } if err := validateResponseFormat(req.ResponseFormat); err != nil { - return nil, err + return nil, "", err } if err := validateOutputFormat(req.OutputFormat); err != nil { - return nil, err + return nil, "", err } prompt := strings.TrimSpace(req.Prompt) if prompt == "" { - return nil, fmt.Errorf("image generation request missing prompt") + return nil, "", fmt.Errorf("image generation request missing prompt") } - return json.Marshal(BananaCreateRequest{ + imageSize, err := imageSize(req.ImageSize, req.Size) + if err != nil { + return nil, "", err + } + concreteModel := providerModel(modelID, req.Model, imageSize) + if reason := UnsupportedModel(concreteModel); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + body, err := json.Marshal(BananaCreateRequest{ Prompt: prompt, - Model: providerModel(modelID, req.Model), - AspectRatio: aspectRatio(req.AspectRatio, req.Size), + Model: concreteModel, + AspectRatio: aspectRatio(req.AspectRatio, req.Ratio, req.Size), + ImageSize: imageSize, PromptOptimization: false, }) + if err != nil { + return nil, "", err + } + return body, concreteModel, nil } -func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, error) { +func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, string, error) { if reason := UnsupportedRequest("/v1/images/edits", openAIBody, contentType); reason != "" { - return nil, fmt.Errorf("%s", reason) + return nil, "", fmt.Errorf("%s", reason) } mediaType, params, err := mime.ParseMediaType(contentType) if err != nil { - return nil, fmt.Errorf("failed to parse image edit content type: %w", err) + return nil, "", fmt.Errorf("failed to parse image edit content type: %w", err) } if !strings.HasPrefix(mediaType, "multipart/form-data") { - return nil, fmt.Errorf("image edits require multipart/form-data content type") + return nil, "", fmt.Errorf("image edits require multipart/form-data content type") } boundary := params["boundary"] if boundary == "" { - return nil, fmt.Errorf("missing multipart boundary in content type") + return nil, "", fmt.Errorf("missing multipart boundary in content type") } fields := make(map[string]string) @@ -105,7 +121,7 @@ func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([] break } if err != nil { - return nil, fmt.Errorf("failed to read multipart image edit payload: %w", err) + return nil, "", fmt.Errorf("failed to read multipart image edit payload: %w", err) } formName := part.FormName() @@ -114,56 +130,69 @@ func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([] } data, err := readLimited(part, maxMultipartImageBytes) if err != nil { - return nil, err + return nil, "", err } if part.FileName() == "" { fields[formName] = strings.TrimSpace(string(data)) continue } if formName == "mask" { - return nil, fmt.Errorf("sosana image edits do not support mask") + return nil, "", fmt.Errorf("image edits do not support mask") } if formName != "image" && formName != "images" && formName != "image[]" { continue } mimeType := detectImageMIMEType(part.Header.Get("Content-Type"), data) if mimeType != "image/png" { - return nil, fmt.Errorf("sosana image edits support PNG images only") + return nil, "", fmt.Errorf("image edits support PNG images only") } imageURLs = append(imageURLs, "data:"+mimeType+";base64,"+base64.StdEncoding.EncodeToString(data)) } if len(imageURLs) > maxInputImages { - return nil, fmt.Errorf("sosana supports up to %d input images", maxInputImages) + return nil, "", fmt.Errorf("image edits support up to %d input images", maxInputImages) } if err := validateImageCountString(fields["n"]); err != nil { - return nil, err + return nil, "", err } if err := validateResponseFormat(fields["response_format"]); err != nil { - return nil, err + return nil, "", err } if err := validateOutputFormat(fields["output_format"]); err != nil { - return nil, err + return nil, "", err } prompt := strings.TrimSpace(fields["prompt"]) if prompt == "" { - return nil, fmt.Errorf("image edit request missing prompt field") + return nil, "", fmt.Errorf("image edit request missing prompt field") } if len(imageURLs) == 0 { - return nil, fmt.Errorf("image edit request missing image") + return nil, "", fmt.Errorf("image edit request missing image") + } + imageSize, err := imageSize(fields["image_size"], fields["size"]) + if err != nil { + return nil, "", err } - return json.Marshal(BananaCreateRequest{ + concreteModel := providerModel(modelID, fields["model"], imageSize) + if reason := UnsupportedModel(concreteModel); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + body, err := json.Marshal(BananaCreateRequest{ Prompt: prompt, ImageURLs: imageURLs, - Model: providerModel(modelID, fields["model"]), - AspectRatio: aspectRatio(fields["aspect_ratio"], fields["size"]), + Model: concreteModel, + AspectRatio: aspectRatio(fields["aspect_ratio"], fields["ratio"], fields["size"]), + ImageSize: imageSize, PromptOptimization: false, }) + if err != nil { + return nil, "", err + } + return body, concreteModel, nil } func OpenAIImageResponse(task BananaTaskResponse, image []byte) ([]byte, error) { if len(image) == 0 { - return nil, fmt.Errorf("sosana task completed without image bytes") + return nil, fmt.Errorf("image task completed without image bytes") } resp := openai.OpenAIImageResponse{ Created: createdAtUnix(task.CreatedAt), @@ -177,11 +206,18 @@ func OpenAIImageResponse(task BananaTaskResponse, image []byte) ([]byte, error) return json.Marshal(resp) } -func providerModel(modelID, requestModel string) string { - if model := strings.TrimSpace(modelID); model != "" { - return model +func providerModel(modelID, requestModel, imageSize string) string { + model := strings.TrimSpace(modelID) + if model == "" { + model = strings.TrimSpace(requestModel) + } + if model == "" { + return "" } - return strings.TrimSpace(requestModel) + if strings.EqualFold(model, "google/gemini-3.1-flash-image-preview") { + return "banana-2-" + strings.ToLower(imageSize) + "-compliant" + } + return strings.ReplaceAll(model, "{image_size}", strings.ToLower(imageSize)) } func createdAtUnix(value string) int64 { @@ -206,33 +242,89 @@ func SizeToAspectRatio(size string) string { switch strings.TrimSpace(size) { case "", "auto": return "auto" - case "256x256", "512x512", "1024x1024", "4096x4096": + case "256x256", "512x512", "1024x1024", "2048x2048", "4096x4096": return "1:1" - case "1024x1536": + case "1024x1536", "2048x3072": return "2:3" - case "1536x1024": + case "1536x1024", "3072x2048": return "3:2" - case "1024x1792", "1080x1920", "4096x7168": + case "1024x1792", "1080x1920", "2048x3584", "4096x7168": return "9:16" - case "1792x1024", "1920x1080", "7168x4096": + case "1792x1024", "1920x1080", "3584x2048", "7168x4096": return "16:9" + case "1024x768", "2048x1536", "4096x3072": + return "4:3" + case "768x1024", "1536x2048", "3072x4096": + return "3:4" + case "1024x819", "2048x1638", "4096x3276": + return "5:4" + case "819x1024", "1638x2048", "3276x4096": + return "4:5" + case "2016x864", "4032x1728": + return "21:9" default: return "auto" } } -func aspectRatio(explicit, size string) string { +func aspectRatio(explicit, ratio, size string) string { if value := strings.TrimSpace(explicit); value != "" { return value } + if value := strings.TrimSpace(ratio); value != "" { + return value + } return SizeToAspectRatio(size) } +func imageSize(explicit, size string) (string, error) { + if strings.TrimSpace(explicit) != "" { + if value, ok := normalizeImageSize(explicit); ok { + return value, nil + } + return "", fmt.Errorf("image_size is unsupported") + } + if value, ok := imageSizeFromExactSize(size); ok { + return value, nil + } + return "", fmt.Errorf("size is unsupported") +} + +func normalizeImageSize(raw string) (string, bool) { + switch strings.ToLower(strings.TrimSpace(raw)) { + case "", "auto", "1k": + return "1K", true + case "2k": + return "2K", true + case "4k": + return "4K", true + default: + return "", false + } +} + +func imageSizeFromExactSize(size string) (string, bool) { + switch strings.TrimSpace(size) { + case "", "auto", + "256x256", "512x512", "1024x1024", "1024x1536", "1536x1024", + "1024x1792", "1792x1024", "1080x1920", "1920x1080", + "1024x768", "768x1024", "1024x819", "819x1024", "2016x864": + return "1K", true + case "2048x2048", "2048x3072", "3072x2048", "2048x3584", "3584x2048", + "2048x1536", "1536x2048", "2048x1638", "1638x2048", "4032x1728": + return "2K", true + case "4096x4096", "4096x7168", "7168x4096", "4096x3072", "3072x4096", "4096x3276", "3276x4096": + return "4K", true + default: + return "", false + } +} + func validateImageCount(n *int) error { if n == nil || *n == 1 { return nil } - return fmt.Errorf("sosana supports n=1 only") + return fmt.Errorf("image requests support n=1 only") } func validateResponseFormat(format string) error { @@ -252,7 +344,7 @@ func validateImageCountString(raw string) error { return fmt.Errorf("invalid image count: %w", err) } if n != 1 { - return fmt.Errorf("sosana supports n=1 only") + return fmt.Errorf("image requests support n=1 only") } return nil } diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go index 636643f3..9b8dcea9 100644 --- a/internal/converter/sosana/images_test.go +++ b/internal/converter/sosana/images_test.go @@ -20,25 +20,29 @@ func TestImageGenerationRequest(t *testing.T) { name string size string wantAspect string + wantSize string wantErrPart string }{ - {name: "square", size: "1024x1024", wantAspect: "1:1"}, - {name: "wide", size: "1792x1024", wantAspect: "16:9"}, - {name: "portrait", size: "1024x1792", wantAspect: "9:16"}, - {name: "unknown", size: "333x777", wantAspect: "auto"}, + {name: "square", size: "1024x1024", wantAspect: "1:1", wantSize: "1K"}, + {name: "wide", size: "1792x1024", wantAspect: "16:9", wantSize: "1K"}, + {name: "portrait", size: "1024x1792", wantAspect: "9:16", wantSize: "1K"}, + {name: "two k", size: "2048x2048", wantAspect: "1:1", wantSize: "2K"}, + {name: "four k", size: "4096x4096", wantAspect: "1:1", wantSize: "4K"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - body := []byte(`{"model":"nano-banana","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) - got, err := ImageGenerationRequest(body, "nano-banana") + body := []byte(`{"model":"banana-2-1k-compliant","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) + got, concreteModel, err := ImageGenerationRequest(body, "banana-2-{image_size}-compliant") require.NoError(t, err) var req BananaCreateRequest require.NoError(t, json.Unmarshal(got, &req)) assert.Equal(t, "draw a cat", req.Prompt) - assert.Equal(t, "nano-banana", req.Model) + assert.Equal(t, "banana-2-"+strings.ToLower(tt.wantSize)+"-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) assert.Equal(t, tt.wantAspect, req.AspectRatio) + assert.Equal(t, tt.wantSize, req.ImageSize) assert.False(t, req.PromptOptimization) assert.Empty(t, req.ImageURLs) }) @@ -46,16 +50,31 @@ func TestImageGenerationRequest(t *testing.T) { } func TestImageGenerationRequestPrefersProviderModel(t *testing.T) { - got, err := ImageGenerationRequest([]byte(`{"model":"public-image","prompt":"draw","n":1}`), "nano-banana") + got, concreteModel, err := ImageGenerationRequest([]byte(`{"model":"public-image","prompt":"draw","n":1}`), "banana-2-1k-compliant") require.NoError(t, err) var req BananaCreateRequest require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "nano-banana", req.Model) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) +} + +func TestImageGenerationRequestMapsPublicGeminiModelToSosanaTier(t *testing.T) { + got, concreteModel, err := ImageGenerationRequest( + []byte(`{"model":"google/gemini-3.1-flash-image-preview","prompt":"draw","image_size":"2K","n":1}`), + "google/gemini-3.1-flash-image-preview", + ) + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "banana-2-2k-compliant", req.Model) + assert.Equal(t, "banana-2-2k-compliant", concreteModel) + assert.Equal(t, "2K", req.ImageSize) } func TestImageGenerationRequestUsesExplicitAspectRatio(t *testing.T) { - got, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","size":"1024x1024","aspect_ratio":"16:9"}`), "nano-banana") + got, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","size":"1024x1024","aspect_ratio":"16:9"}`), "banana-2-1k-compliant") require.NoError(t, err) var req BananaCreateRequest @@ -64,13 +83,13 @@ func TestImageGenerationRequestUsesExplicitAspectRatio(t *testing.T) { } func TestImageGenerationRequestRejectsMultipleImages(t *testing.T) { - _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","n":2}`), "nano-banana") + _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","n":2}`), "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), "n=1") } func TestImageGenerationRequestRejectsURLResponseFormat(t *testing.T) { - _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","response_format":"url"}`), "nano-banana") + _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","response_format":"url"}`), "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), "response_format=url") } @@ -81,19 +100,20 @@ func TestImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { body string want string }{ - {name: "tools", body: `{"model":"nano-banana","prompt":"draw","tools":[{"type":"google_search"}]}`, want: "tools"}, - {name: "thinking", body: `{"model":"nano-banana","prompt":"draw","thinking_level":"high"}`, want: "thinking_level"}, - {name: "output format", body: `{"model":"nano-banana","prompt":"draw","output_format":"jpeg"}`, want: "output_format"}, - {name: "output compression", body: `{"model":"nano-banana","prompt":"draw","output_compression":0}`, want: "output_compression"}, - {name: "quality auto", body: `{"model":"nano-banana","prompt":"draw","quality":"auto"}`, want: "quality"}, - {name: "messages", body: `{"model":"nano-banana","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, - {name: "image size", body: `{"model":"nano-banana","prompt":"draw","image_size":"2K"}`, want: "image_size"}, - {name: "reference images", body: `{"model":"nano-banana","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, + {name: "tools", body: `{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`, want: "tools"}, + {name: "thinking", body: `{"model":"banana-2-1k-compliant","prompt":"draw","thinking_level":"high"}`, want: "thinking_level"}, + {name: "output format", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"jpeg"}`, want: "output_format"}, + {name: "output compression", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_compression":0}`, want: "output_compression"}, + {name: "quality auto", body: `{"model":"banana-2-1k-compliant","prompt":"draw","quality":"auto"}`, want: "quality"}, + {name: "messages", body: `{"model":"banana-2-1k-compliant","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, + {name: "image size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_size":"0.5K"}`, want: "image_size"}, + {name: "unknown size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"333x777"}`, want: "size"}, + {name: "reference images", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - _, err := ImageGenerationRequest([]byte(tt.body), "nano-banana") + _, _, err := ImageGenerationRequest([]byte(tt.body), "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), tt.want) }) @@ -101,13 +121,13 @@ func TestImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { } func TestImageGenerationRequestAllowsPNGOutputFormat(t *testing.T) { - _, err := ImageGenerationRequest([]byte(`{"model":"nano-banana","prompt":"draw","output_format":"png","response_format":"b64_json"}`), "nano-banana") + _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"png","response_format":"b64_json"}`), "banana-2-1k-compliant") require.NoError(t, err) } func TestImageEditRequest(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "nano-banana", + "model": "banana-2-1k-compliant", "prompt": "make it blue", "size": "1024x1024", "n": "1", @@ -115,14 +135,16 @@ func TestImageEditRequest(t *testing.T) { "image": pngBytes(), }) - got, err := ImageEditRequest(body, contentType, "nano-banana") + got, concreteModel, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") require.NoError(t, err) var req BananaCreateRequest require.NoError(t, json.Unmarshal(got, &req)) assert.Equal(t, "make it blue", req.Prompt) - assert.Equal(t, "nano-banana", req.Model) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) assert.Equal(t, "1:1", req.AspectRatio) + assert.Equal(t, "1K", req.ImageSize) assert.False(t, req.PromptOptimization) require.Len(t, req.ImageURLs, 1) assert.True(t, strings.HasPrefix(req.ImageURLs[0], "data:image/png;base64,")) @@ -137,65 +159,66 @@ func TestImageEditRequestPrefersProviderModel(t *testing.T) { "image": pngBytes(), }) - got, err := ImageEditRequest(body, contentType, "nano-banana") + got, concreteModel, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") require.NoError(t, err) var req BananaCreateRequest require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "nano-banana", req.Model) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) } func TestImageEditRequestRejectsMask(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "nano-banana", + "model": "banana-2-1k-compliant", "prompt": "make it blue", }, map[string][]byte{ "image": pngBytes(), "mask": pngBytes(), }) - _, err := ImageEditRequest(body, contentType, "fallback-model") + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), "mask") } func TestImageEditRequestRejectsMultipleImagesCount(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "nano-banana", + "model": "banana-2-1k-compliant", "prompt": "make it blue", "n": "2", }, map[string][]byte{ "image": pngBytes(), }) - _, err := ImageEditRequest(body, contentType, "fallback-model") + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), "n=1") } func TestImageEditRequestRejectsURLResponseFormat(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "nano-banana", + "model": "banana-2-1k-compliant", "prompt": "make it blue", "response_format": "url", }, map[string][]byte{ "image": pngBytes(), }) - _, err := ImageEditRequest(body, contentType, "fallback-model") + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), "response_format=url") } func TestImageEditRequestRejectsJPEGInput(t *testing.T) { body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "nano-banana", + "model": "banana-2-1k-compliant", "prompt": "make it blue", }, map[string][]byte{ "image": jpegBytes(), }) - _, err := ImageEditRequest(body, contentType, "fallback-model") + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), "PNG") } @@ -203,7 +226,7 @@ func TestImageEditRequestRejectsJPEGInput(t *testing.T) { func TestImageEditRequestRejectsTooManyImages(t *testing.T) { var buf bytes.Buffer writer := multipart.NewWriter(&buf) - require.NoError(t, writer.WriteField("model", "nano-banana")) + require.NoError(t, writer.WriteField("model", "banana-2-1k-compliant")) require.NoError(t, writer.WriteField("prompt", "make it blue")) for i := 0; i < maxInputImages+1; i++ { part, err := writer.CreateFormFile("image", fmt.Sprintf("image-%02d.png", i)) @@ -213,7 +236,7 @@ func TestImageEditRequestRejectsTooManyImages(t *testing.T) { } require.NoError(t, writer.Close()) - _, err := ImageEditRequest(buf.Bytes(), writer.FormDataContentType(), "fallback-model") + _, _, err := ImageEditRequest(buf.Bytes(), writer.FormDataContentType(), "banana-2-1k-compliant") require.Error(t, err) assert.Contains(t, err.Error(), "too many") } diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 216679c8..7c914cea 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -510,6 +510,10 @@ func (p *Proxy) ProxyRequest(w http.ResponseWriter, r *http.Request) { } } + if !p.applySosanaCompatibilityRouting(w, r, prepared, modelID, &cred, &body, &proxyBody, &realModelID, isImageGeneration, isImageEdit, logCtx, start) { + return + } + if cred.Type == config.ProviderTypeSosana { p.handleSosanaRequest(w, r, body, cred, modelID, realModelID, isImageGeneration, isImageEdit, logCtx, start) return diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index edae47b1..fb07732c 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -49,7 +49,7 @@ func (p *Proxy) handleSosanaRequest( logCtx.TargetURL = cred.BaseURL if !isImageGeneration && !isImageEdit { - message := "sosana provider supports only image generation" + message := unsupportedProviderEndpointMessage logCtx.Status = "failure" logCtx.HTTPStatus = http.StatusBadRequest logCtx.ErrorMsg = message @@ -57,14 +57,7 @@ func (p *Proxy) handleSosanaRequest( return } - createBody, err := p.buildSosanaCreateBody(body, r.Header.Get("Content-Type"), realModelID, isImageEdit) - if err != nil { - logCtx.Status = "failure" - logCtx.HTTPStatus = http.StatusBadRequest - logCtx.ErrorMsg = err.Error() - WriteErrorBadRequest(w, err.Error()) - return - } + baseRealModelID := realModelID ctx := r.Context() var cancel context.CancelFunc @@ -95,6 +88,22 @@ func (p *Proxy) handleSosanaRequest( time.Sleep(time.Duration(rand.Intn(50)) * time.Millisecond) } + attemptRealModelID := p.sosanaRealModelIDForCredential(modelID, baseRealModelID, cred) + createBody, concreteModelID, err := p.buildSosanaCreateBody(body, r.Header.Get("Content-Type"), attemptRealModelID, isImageEdit) + if err != nil { + p.logger.DebugContext(r.Context(), "Failed to prepare provider image request", + "credential", cred.Name, + "model", modelID, + "real_model", attemptRealModelID, + "error", err) + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = unsupportedImageProviderRequestMessage + WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) + return + } + logCtx.RealModelID = concreteModelID + result = p.createAndPollSosanaTask(ctx, cred, modelID, createBody, logCtx) p.balancer.RecordResponse(cred.Name, modelID, result.statusCode) p.metrics.RecordRequest(cred.Name, r.URL.Path, modelID, result.statusCode, time.Since(start)) @@ -138,7 +147,19 @@ func (p *Proxy) handleSosanaRequest( } } -func (p *Proxy) buildSosanaCreateBody(body []byte, contentType, realModelID string, isImageEdit bool) ([]byte, error) { +func (p *Proxy) sosanaRealModelIDForCredential(modelID, fallbackRealModelID string, cred *config.CredentialConfig) string { + if p.modelManager != nil && cred != nil { + if realModelID, ok := p.modelManager.GetRealModelNameForCredential(modelID, cred.Name); ok { + return realModelID + } + } + if strings.TrimSpace(fallbackRealModelID) != "" { + return fallbackRealModelID + } + return modelID +} + +func (p *Proxy) buildSosanaCreateBody(body []byte, contentType, realModelID string, isImageEdit bool) ([]byte, string, error) { if isImageEdit { return sosana.ImageEditRequest(body, contentType, realModelID) } @@ -368,12 +389,72 @@ func (p *Proxy) downloadSosanaResultImage( func (p *Proxy) doSosanaResultImageRequest(req *http.Request) (*http.Response, error) { client := *p.client + client.Transport = sosanaResultImageTransport() client.CheckRedirect = func(req *http.Request, via []*http.Request) error { return http.ErrUseLastResponse } return client.Do(req) } +func sosanaResultImageTransport() http.RoundTripper { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.Proxy = nil + transport.DisableKeepAlives = true + transport.DialContext = dialSosanaResultAddress + return transport +} + +func dialSosanaResultAddress(ctx context.Context, network, address string) (net.Conn, error) { + host, port, err := net.SplitHostPort(address) + if err != nil { + return nil, err + } + + dialer := &net.Dialer{} + if allowPrivateSosanaResultHostForTests(host) { + return dialer.DialContext(ctx, network, address) + } + if ip := net.ParseIP(host); ip != nil { + if isUnsafeSosanaResultIP(ip) { + return nil, errors.New("sosana result_file_url resolves to a private address") + } + return dialer.DialContext(ctx, network, address) + } + + addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) + if err != nil { + return nil, err + } + if len(addrs) == 0 { + return nil, errors.New("sosana result_file_url host has no addresses") + } + for _, addr := range addrs { + if isUnsafeSosanaResultIP(addr.IP) { + return nil, errors.New("sosana result_file_url resolves to a private address") + } + } + + var dialErr error + for _, addr := range addrs { + conn, err := dialer.DialContext(ctx, network, net.JoinHostPort(addr.IP.String(), port)) + if err == nil { + return conn, nil + } + dialErr = err + } + if dialErr != nil { + return nil, dialErr + } + return nil, errors.New("sosana result_file_url host has no dialable addresses") +} + +func allowPrivateSosanaResultHostForTests(host string) bool { + if allowPrivateSosanaResultURLForTests == nil { + return false + } + return allowPrivateSosanaResultURLForTests(&url.URL{Scheme: "http", Host: host}) +} + func parseSosanaResultURL(raw string) (*url.URL, error) { if raw == "" { return nil, errors.New("sosana task completed without result_file_url") diff --git a/internal/proxy/sosana_live_test.go b/internal/proxy/sosana_live_test.go index 7aa6608a..6ea6da22 100644 --- a/internal/proxy/sosana_live_test.go +++ b/internal/proxy/sosana_live_test.go @@ -29,7 +29,7 @@ func TestProxyRequest_SosanaLiveAcceptance(t *testing.T) { } model := os.Getenv("SOSANA_MODEL") if model == "" { - model = "nano-banana" + model = "banana-2-1k-compliant" } prompt := os.Getenv("SOSANA_PROMPT") if prompt == "" { diff --git a/internal/proxy/sosana_routing.go b/internal/proxy/sosana_routing.go new file mode 100644 index 00000000..fc5484ce --- /dev/null +++ b/internal/proxy/sosana_routing.go @@ -0,0 +1,156 @@ +package proxy + +import ( + "net/http" + "time" + + "github.com/mixaill76/auto_ai_router/internal/config" + "github.com/mixaill76/auto_ai_router/internal/converter/sosana" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" +) + +const unsupportedImageProviderRequestMessage = "request parameters are not supported by available image providers" +const unsupportedProviderEndpointMessage = "request endpoint is not supported by available providers" + +func (p *Proxy) applySosanaCompatibilityRouting( + w http.ResponseWriter, + r *http.Request, + prepared *orchestratedRequest, + modelID string, + cred **config.CredentialConfig, + body *[]byte, + proxyBody *[]byte, + realModelID *string, + isImageGeneration bool, + isImageEdit bool, + logCtx *RequestLogContext, + start time.Time, +) bool { + if (*cred).Type != config.ProviderTypeSosana || (!isImageGeneration && !isImageEdit) { + return true + } + + reason := sosana.UnsupportedModel(*realModelID) + if reason == "" { + reason = sosana.UnsupportedRequest(r.URL.Path, *body, r.Header.Get("Content-Type")) + } + if reason == "" { + return true + } + + nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, reason) + if routed { + *cred = nextCred + *body = nextReq.body + *proxyBody = nextReq.proxyBody + *realModelID = nextReq.realModelID + r.URL.Path = nextReq.path + prepared.body = nextReq.body + prepared.proxyBody = nextReq.proxyBody + prepared.proxyPath = nextReq.proxyPath + prepared.realModelID = nextReq.realModelID + prepared.convertedResp = nextReq.convertedResp + prepared.passthroughResponses = nextReq.passthroughResponses + prepared.nativeResponses = nextReq.nativeResponses + logCtx.RealModelID = *realModelID + if span := trace.SpanFromContext(r.Context()); span.IsRecording() { + span.SetAttributes( + attribute.String("aar.real_model", *realModelID), + attribute.String("aar.credential", nextCred.Name), + attribute.String("aar.provider", string(nextCred.Type)), + attribute.Bool("aar.provider_compatibility_skip", true), + ) + } + return true + } + + success, fallbackReason := p.TryFallbackProxy( + w, + requestWithPath(r, prepared.proxyPath), + modelID, + (*cred).Name, + http.StatusBadRequest, + RetryReasonServerErr, + *proxyBody, + start, + logCtx, + ) + if success { + return false + } + p.logger.DebugContext(r.Context(), "No fallback handled unsupported image provider request", + "credential", (*cred).Name, + "model", modelID, + "reason", reason, + "fallback_reason", fallbackReason) + logCtx.Credential = *cred + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = unsupportedImageProviderRequestMessage + WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) + return false +} + +func (p *Proxy) nextPrimaryAfterUnsupportedSosana( + r *http.Request, + prepared *orchestratedRequest, + modelID string, + currentCred *config.CredentialConfig, + reason string, +) (*config.CredentialConfig, credentialPreparedRequest, bool) { + triedCreds := GetTried(r.Context()) + triedCreds[currentCred.Name] = true + + for attempts := 0; attempts < 128; attempts++ { + candidate, err := p.balancer.NextForModelExcluding(modelID, triedCreds) + if err != nil { + p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", + "model", modelID, + "credential", currentCred.Name, + "reason", reason, + "error", err) + return nil, credentialPreparedRequest{}, false + } + triedCreds[candidate.Name] = true + if candidate.Type == config.ProviderTypeSosana { + continue + } + + nextReq, prepErr := p.prepareRequestForCredential( + r, + prepared.baseBody, + prepared.baseProxyBody, + modelID, + prepared.baseRealModelID, + prepared.basePath, + prepared.streaming, + candidate, + prepared.isResponsesAPI, + prepared.responsesPrevHandled, + prepared.stickyCacheEligible, + ) + if prepErr != nil { + p.logger.WarnContext(r.Context(), "Failed to prepare alternate primary request after image compatibility skip", + "credential", candidate.Name, + "provider", string(candidate.Type), + "model", modelID, + "reason", reason, + "error", prepErr) + continue + } + + p.logger.InfoContext(r.Context(), "Skipping incompatible image credential for unsupported image request", + "credential", currentCred.Name, + "next_credential", candidate.Name, + "model", modelID, + "reason", reason) + return candidate, nextReq, true + } + + p.logger.WarnContext(r.Context(), "Image compatibility skip exhausted primary credential scan", + "credential", currentCred.Name, + "model", modelID, + "reason", reason) + return nil, credentialPreparedRequest{}, false +} diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 8ae9157b..3cee7ee0 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -41,7 +41,7 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { var req map[string]any require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) assert.Equal(t, "draw a fox", req["prompt"]) - assert.Equal(t, "nano-banana", req["model"]) + assert.Equal(t, "banana-2-1k-compliant", req["model"]) assert.Equal(t, "1:1", req["aspect_ratio"]) assert.Equal(t, false, req["prompt_optimization"]) _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","created_at":"2026-01-01T00:00:00Z","prompt":"draw a fox"}`)) @@ -55,7 +55,7 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, nil) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw a fox","size":"1024x1024","n":1,"response_format":"b64_json"}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw a fox","size":"1024x1024","n":1,"response_format":"b64_json"}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -71,6 +71,8 @@ func TestProxyRequest_SosanaImageGenerationSuccess(t *testing.T) { assert.Equal(t, "A detailed fox illustration", resp.Data[0].RevisedPrompt) assert.Equal(t, int64(1767225600), resp.Created) assert.NotContains(t, w.Body.String(), imageServer.URL) + assert.NotContains(t, w.Body.String(), "result_file_url") + assert.NotContains(t, w.Body.String(), "main-r2") assert.NotContains(t, w.Body.String(), "sosana") assert.NotContains(t, w.Body.String(), "cdn") assert.Equal(t, []string{""}, imageAuths) @@ -89,7 +91,7 @@ func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { var req map[string]any require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) assert.Equal(t, "make it blue", req["prompt"]) - assert.Equal(t, "nano-banana", req["model"]) + assert.Equal(t, "banana-2-1k-compliant", req["model"]) assert.Equal(t, false, req["prompt_optimization"]) imageURLs, ok := req["image_urls"].([]any) require.True(t, ok) @@ -102,7 +104,7 @@ func TestProxyRequest_SosanaImageEditUsesProviderModelAlias(t *testing.T) { prx := newSosanaTestProxy(upstream.URL, nil) prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ - {Name: "public-image", Model: "nano-banana", Credential: "sosana"}, + {Name: "public-image", Model: "banana-2-1k-compliant", Credential: "sosana"}, }) body, contentType := sosanaMultipartEditBody(t, map[string]string{ @@ -146,14 +148,14 @@ func TestProxyRequest_SosanaImageGenerationLogsLiteLLMImageSpend(t *testing.T) { spendManager := newCapturedSpendManager() priceRegistry := aimodels.NewModelPriceRegistry() priceRegistry.Update(map[string]*aimodels.ModelPrice{ - "nano-banana": {OutputCostPerImage: 0.07}, + "banana-2-1k-compliant": {OutputCostPerImage: 0.07}, }) prx := newSosanaTestProxy(upstream.URL, nil) prx.LiteLLMDB = spendManager prx.priceRegistry = priceRegistry prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ - {Name: "public-image", Model: "nano-banana", Credential: "sosana"}, + {Name: "public-image", Model: "banana-2-1k-compliant", Credential: "sosana"}, }) req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"public-image","prompt":"draw","n":1}`)) @@ -180,6 +182,55 @@ func TestProxyRequest_SosanaImageGenerationLogsLiteLLMImageSpend(t *testing.T) { assert.InDelta(t, 0.07, costBreakdown["total_cost"].(float64), 0.0000001) } +func TestProxyRequest_SosanaImageGenerationUsesConcreteTierPrice(t *testing.T) { + imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/png", sosanaResultPNG, nil) + defer imageServer.Close() + + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/banana/create-async": + var req map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + assert.Equal(t, "banana-2-2k-compliant", req["model"]) + assert.Equal(t, "2K", req["image_size"]) + _, _ = w.Write([]byte(`{"uid":"task-1","status":"PROCESSING","prompt":"draw"}`)) + case "/api/banana/task-1": + _, _ = fmt.Fprintf(w, `{"uid":"task-1","status":"COMPLETED","prompt":"draw","result_file_url":%q}`, imageServer.URL+"/spend.png") + default: + t.Fatalf("unexpected upstream path: %s", r.URL.Path) + } + })) + defer upstream.Close() + + spendManager := newCapturedSpendManager() + priceRegistry := aimodels.NewModelPriceRegistry() + priceRegistry.Update(map[string]*aimodels.ModelPrice{ + "google/gemini-3.1-flash-image-preview": {OutputCostPerImage: 9.99}, + "banana-2-2k-compliant": {OutputCostPerImage: 0.22}, + }) + + prx := newSosanaTestProxy(upstream.URL, nil) + prx.LiteLLMDB = spendManager + prx.priceRegistry = priceRegistry + prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ + {Name: "google/gemini-3.1-flash-image-preview", Model: "banana-2-{image_size}-compliant", Credential: "sosana"}, + }) + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"google/gemini-3.1-flash-image-preview","prompt":"draw","image_size":"2K","n":1}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + require.Len(t, spendManager.entries, 1) + assert.InDelta(t, 0.22, spendManager.entries[0].Spend, 0.0000001) + assert.Equal(t, "google/gemini-3.1-flash-image-preview", spendManager.entries[0].Model) + assert.NotContains(t, w.Body.String(), "cost") + assert.NotContains(t, w.Body.String(), "spend") +} + func TestProxyRequest_SosanaImageGenerationWritesLiteLLMSpendLogIntegration(t *testing.T) { dbURL := os.Getenv("LITELLM_DATABASE_URL") if dbURL == "" { @@ -219,14 +270,14 @@ func TestProxyRequest_SosanaImageGenerationWritesLiteLLMSpendLogIntegration(t *t alias := fmt.Sprintf("sosana-spend-test-%d", time.Now().UnixNano()) priceRegistry := aimodels.NewModelPriceRegistry() priceRegistry.Update(map[string]*aimodels.ModelPrice{ - "nano-banana": {OutputCostPerImage: 0.07}, + "banana-2-1k-compliant": {OutputCostPerImage: 0.07}, }) prx := newSosanaTestProxy(upstream.URL, nil) prx.LiteLLMDB = manager prx.priceRegistry = priceRegistry prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ - {Name: alias, Model: "nano-banana", Credential: "sosana"}, + {Name: alias, Model: "banana-2-1k-compliant", Credential: "sosana"}, }) req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"`+alias+`","prompt":"draw","n":1}`)) @@ -282,7 +333,7 @@ func TestProxyRequest_SosanaRejectsNonImageEndpoint(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, nil) - req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"nano-banana","messages":[]}`)) + req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"banana-2-1k-compliant","messages":[]}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -290,7 +341,8 @@ func TestProxyRequest_SosanaRejectsNonImageEndpoint(t *testing.T) { prx.ProxyRequest(w, req) assert.Equal(t, http.StatusBadRequest, w.Code) - assert.Contains(t, w.Body.String(), "provider supports only image generation") + assert.Contains(t, w.Body.String(), unsupportedProviderEndpointMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") assert.False(t, called) } @@ -303,7 +355,7 @@ func TestProxyRequest_SosanaRejectsURLResponseFormat(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, nil) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1,"response_format":"url"}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1,"response_format":"url"}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -311,7 +363,8 @@ func TestProxyRequest_SosanaRejectsURLResponseFormat(t *testing.T) { prx.ProxyRequest(w, req) assert.Equal(t, http.StatusBadRequest, w.Code) - assert.Contains(t, w.Body.String(), "response_format=url") + assert.Contains(t, w.Body.String(), unsupportedImageProviderRequestMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") assert.False(t, called) } @@ -324,7 +377,7 @@ func TestProxyRequest_IncompatibleImageRequestWithSosanaReturnsLocalError(t *tes defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, nil) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","tools":[{"type":"google_search"}]}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -332,7 +385,8 @@ func TestProxyRequest_IncompatibleImageRequestWithSosanaReturnsLocalError(t *tes prx.ProxyRequest(w, req) require.Equal(t, http.StatusBadRequest, w.Code) - assert.Contains(t, w.Body.String(), "tools is unsupported") + assert.Contains(t, w.Body.String(), unsupportedImageProviderRequestMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") assert.False(t, called) } @@ -345,7 +399,7 @@ func TestProxyRequest_IncompatibleImageEditJPEGWithSosanaReturnsLocalError(t *te defer upstream.Close() body, contentType := sosanaMultipartEditBody(t, map[string]string{ - "model": "nano-banana", + "model": "banana-2-1k-compliant", "prompt": "make it blue", }, map[string][]byte{ "image": {0xff, 0xd8, 0xff, 0xdb, 0, 0x43}, @@ -359,10 +413,89 @@ func TestProxyRequest_IncompatibleImageEditJPEGWithSosanaReturnsLocalError(t *te prx.ProxyRequest(w, req) require.Equal(t, http.StatusBadRequest, w.Code) - assert.Contains(t, w.Body.String(), "PNG input") + assert.Contains(t, w.Body.String(), unsupportedImageProviderRequestMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") assert.False(t, called) } +func TestProxyRequest_IncompatibleSosanaImageRequestRoutesToNextPrimary(t *testing.T) { + sosanaCalled := false + sosanaUpstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sosanaCalled = true + w.WriteHeader(http.StatusOK) + })) + defer sosanaUpstream.Close() + + nextPrimaryCalled := false + nextPrimary := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + nextPrimaryCalled = true + assert.Equal(t, "/v1/images/generations", r.URL.Path) + assert.Equal(t, "Bearer next-key", r.Header.Get("Authorization")) + var req map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + assert.Equal(t, "banana-2-1k-compliant", req["model"]) + assert.Contains(t, req, "tools") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"created":1782478551,"data":[{"b64_json":"fallback-primary-image"}]}`)) + })) + defer nextPrimary.Close() + + prx := NewTestProxyBuilder(). + WithCredentials( + config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: sosanaUpstream.URL, APIKey: "sosana-key", RPM: 100, TPM: 10000}, + config.CredentialConfig{Name: "next-primary", Type: config.ProviderTypeProxy, BaseURL: nextPrimary.URL, APIKey: "next-key", RPM: 100, TPM: 10000}, + ). + Build() + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + assert.True(t, nextPrimaryCalled) + assert.False(t, sosanaCalled) + assert.Contains(t, w.Body.String(), "fallback-primary-image") +} + +func TestProxyRequest_IncompatibleSosanaImageRequestRoutesToFallbackProxy(t *testing.T) { + sosanaCalled := false + sosanaUpstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sosanaCalled = true + w.WriteHeader(http.StatusOK) + })) + defer sosanaUpstream.Close() + + fallbackCalled := false + fallbackProxy := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fallbackCalled = true + assert.Equal(t, "/v1/images/generations", r.URL.Path) + assert.Equal(t, "Bearer fallback-key", r.Header.Get("Authorization")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"created":1782478551,"data":[{"b64_json":"fallback-proxy-image"}]}`)) + })) + defer fallbackProxy.Close() + + prx := NewTestProxyBuilder(). + WithCredentials( + config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: sosanaUpstream.URL, APIKey: "sosana-key", RPM: 100, TPM: 10000}, + config.CredentialConfig{Name: "fallback-proxy", Type: config.ProviderTypeProxy, BaseURL: fallbackProxy.URL, APIKey: "fallback-key", RPM: 100, TPM: 10000, IsFallback: true}, + ). + Build() + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","thinking_level":"high"}`)) + req.Header.Set("Authorization", "Bearer master-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusOK, w.Code) + assert.True(t, fallbackCalled) + assert.False(t, sosanaCalled) + assert.Contains(t, w.Body.String(), "fallback-proxy-image") +} + func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { var logBuf bytes.Buffer upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -372,7 +505,7 @@ func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -388,6 +521,7 @@ func TestProxyRequest_SosanaCreateHTTPErrorMasked(t *testing.T) { func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { var createAuths []string + var createModels []string imageServer := newSosanaResultImageServer(t, http.StatusOK, "image/png", sosanaResultPNG, nil) defer imageServer.Close() @@ -395,12 +529,16 @@ func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { switch r.URL.Path { case "/api/banana/create-async": createAuths = append(createAuths, r.Header.Get("Authorization")) + var req map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + createModels = append(createModels, req["model"].(string)) if r.Header.Get("Authorization") == "Bearer sosana-key-a" { w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"detail":"first credential rate limited"}`)) return } assert.Equal(t, "Bearer sosana-key-b", r.Header.Get("Authorization")) + assert.Equal(t, "banana-2-2k-compliant", req["model"]) _, _ = w.Write([]byte(`{"uid":"task-2","status":"PROCESSING","prompt":"draw"}`)) case "/api/banana/task-2": assert.Equal(t, "Bearer sosana-key-b", r.Header.Get("Authorization")) @@ -418,8 +556,12 @@ func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { ). WithMaxProviderRetries(1). Build() + prx.modelManager = aimodels.New(prx.logger, 50, []config.ModelRPMConfig{ + {Name: "public-image", Model: "banana-2-1k-compliant", Credential: "sosana-a"}, + {Name: "public-image", Model: "banana-2-2k-compliant", Credential: "sosana-b"}, + }) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"public-image","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -428,6 +570,7 @@ func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { require.Equal(t, http.StatusOK, w.Code) assert.Equal(t, []string{"Bearer sosana-key-a", "Bearer sosana-key-b"}, createAuths) + assert.Equal(t, []string{"banana-2-1k-compliant", "banana-2-2k-compliant"}, createModels) assert.Contains(t, w.Body.String(), base64.StdEncoding.EncodeToString(sosanaResultPNG)) assert.NotContains(t, w.Body.String(), imageServer.URL) } @@ -452,7 +595,7 @@ func TestProxyRequest_SosanaDoesNotRetryCreateTransportError(t *testing.T) { WithMaxProviderRetries(1). Build() - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -479,7 +622,7 @@ func TestProxyRequest_SosanaPollHTTPErrorMasked(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -511,7 +654,7 @@ func TestProxyRequest_SosanaImageResultHTTPErrorMasked(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -545,7 +688,7 @@ func TestProxyRequest_SosanaImageResultNonImageMasked(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -579,7 +722,7 @@ func TestProxyRequest_SosanaImageResultJPEGMasked(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -623,7 +766,7 @@ func TestProxyRequest_SosanaImageResultRedirectMaskedAndNotFollowed(t *testing.T defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -647,7 +790,7 @@ func TestDownloadSosanaResultImageRejectsUnsafeProductionURL(t *testing.T) { resultURL := imageServer.URL + "/private.png" prx := newSosanaTestProxy("https://sosana.art", &logBuf) cred := &config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: "https://sosana.art"} - image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "nano-banana", sosana.BananaTaskResponse{ + image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "banana-2-1k-compliant", sosana.BananaTaskResponse{ Status: sosana.StatusCompleted, ResultFileURL: &resultURL, }, nil, &RequestLogContext{}) @@ -687,6 +830,14 @@ func TestValidateSosanaResultURLAllowsLocalOnlyWithTestHook(t *testing.T) { require.NoError(t, validateSosanaResultURL(context.Background(), parsed)) } +func TestDialSosanaResultAddressRejectsPrivateIP(t *testing.T) { + conn, err := dialSosanaResultAddress(context.Background(), "tcp", net.JoinHostPort("127.0.0.1", "443")) + + require.Error(t, err) + assert.Nil(t, conn) + assert.Contains(t, err.Error(), "private address") +} + func TestProxyRequest_SosanaImageResultTimeoutMasked(t *testing.T) { allowPrivateSosanaResultURLsForTest(t) @@ -712,7 +863,7 @@ func TestProxyRequest_SosanaImageResultTimeoutMasked(t *testing.T) { prx := newSosanaTestProxy(upstream.URL, &logBuf) prx.requestTimeout = 5 * time.Millisecond - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -752,7 +903,7 @@ func TestProxyRequest_SosanaDoesNotRetryAfterTaskCreated(t *testing.T) { WithMaxProviderRetries(1). Build() - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -778,7 +929,7 @@ func TestProxyRequest_SosanaTaskFailedMasked(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -805,7 +956,7 @@ func TestProxyRequest_SosanaTaskModeratedMasked(t *testing.T) { defer upstream.Close() prx := newSosanaTestProxy(upstream.URL, &logBuf) - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() @@ -828,7 +979,7 @@ func TestProxyRequest_SosanaTimeoutMasked(t *testing.T) { prx := newSosanaTestProxy(upstream.URL, &logBuf) prx.requestTimeout = 5 * time.Millisecond - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"nano-banana","prompt":"draw","n":1}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() From 255f2dcec667321fcf204a0bd61c61dc5b5b6c1e Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Tue, 7 Jul 2026 17:44:28 +0300 Subject: [PATCH 08/16] fix: remove configurable upstream masking flag --- config.yaml.example | 1 - docs/providers/sosana.md | 6 +++--- internal/config/config_test.go | 15 ------------- internal/config/utils.go | 19 ++++++++--------- internal/proxy/cometapi_test.go | 28 +++++++------------------ internal/proxy/proxy_log.go | 2 +- internal/proxy/upstream_masking_test.go | 13 ++++++------ 7 files changed, 26 insertions(+), 58 deletions(-) diff --git a/config.yaml.example b/config.yaml.example index e4bc1860..d87696f4 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -53,7 +53,6 @@ credentials: rpm: 100 tpm: 50000 weight: 100 # Optional: weighted round-robin share (default 1). See docs/advanced/balancing.md - # mask_upstream_errors: true # Enable for OpenAI-compatible resellers whose raw errors must not reach clients. - name: "vertex_ai" type: "vertex-ai" diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md index f4fa628e..b0098016 100644 --- a/docs/providers/sosana.md +++ b/docs/providers/sosana.md @@ -86,6 +86,6 @@ For operator debugging, structured logs may include a truncated textual upstream error body with `response_body_masked=true`. Raw image bytes and full result URLs are not logged. -If Sosana is hidden behind another proxy credential, enable -`mask_upstream_errors: true` on that proxy unless the upstream router is known to -propagate the credential marker used by this router. +If Sosana is hidden behind another proxy credential, the upstream router should +propagate the credential marker used by this router so proxy-chain errors can be +masked as Sosana errors too. diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 2262bf09..379a41c7 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -677,21 +677,6 @@ rpm: 60 assert.Equal(t, ProviderTypeCometAPI, cred.Type) } -func TestCredentialConfig_MaskUpstreamErrors(t *testing.T) { - var cred CredentialConfig - err := yaml.Unmarshal([]byte(` -name: sosana -type: openai -api_key: key -base_url: https://api.sosana.example/v1 -mask_upstream_errors: true -rpm: 60 -`), &cred) - - require.NoError(t, err) - assert.True(t, cred.MaskUpstreamErrors) -} - func TestCredentialConfig_NormalizeSosanaProviderType(t *testing.T) { tests := []struct { name string diff --git a/internal/config/utils.go b/internal/config/utils.go index 0032ccb7..696821a5 100644 --- a/internal/config/utils.go +++ b/internal/config/utils.go @@ -118,15 +118,14 @@ func PrintConfig(logger *slog.Logger, cfg *Config) { ) for i, cred := range cfg.Credentials { credLog := map[string]any{ - "name": cred.Name, - "type": cred.Type, - "base_url": cred.BaseURL, - "auth_type": cred.AuthType, - "mask_upstream_errors": cred.MaskUpstreamErrors, - "rpm": rpmToString(cred.RPM), - "tpm": tpmToString(cred.TPM), - "is_fallback": cred.IsFallback, - "fallback_priority": cred.FallbackPriority, + "name": cred.Name, + "type": cred.Type, + "base_url": cred.BaseURL, + "auth_type": cred.AuthType, + "rpm": rpmToString(cred.RPM), + "tpm": tpmToString(cred.TPM), + "is_fallback": cred.IsFallback, + "fallback_priority": cred.FallbackPriority, } // Add Vertex AI specific fields if present @@ -237,7 +236,7 @@ func banDurationToString(d time.Duration) string { func convertMapToArgs(m map[string]any) []any { // Define preferred order of keys keyOrder := []string{ - "name", "type", "base_url", "auth_type", "mask_upstream_errors", "api_key", "project_id", "location", + "name", "type", "base_url", "auth_type", "api_key", "project_id", "location", "credentials_file", "credentials_json", "rpm", "tpm", "is_fallback", } diff --git a/internal/proxy/cometapi_test.go b/internal/proxy/cometapi_test.go index fe2f9b2c..34b01fee 100644 --- a/internal/proxy/cometapi_test.go +++ b/internal/proxy/cometapi_test.go @@ -14,50 +14,37 @@ func TestIsCometAPICredential(t *testing.T) { name string cred *config.CredentialConfig wantComet bool - wantMask bool }{ { name: "dedicated provider type", cred: &config.CredentialConfig{Type: config.ProviderTypeCometAPI}, wantComet: true, - wantMask: true, }, { name: "comet host fallback", cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, BaseURL: "https://api.cometapi.com/v1"}, wantComet: true, - wantMask: true, }, { name: "comet name fallback", cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, Name: "comet-api-anthropic"}, wantComet: true, - wantMask: true, }, { name: "regular anthropic", cred: &config.CredentialConfig{Type: config.ProviderTypeAnthropic, BaseURL: "https://api.anthropic.com"}, wantComet: false, - wantMask: false, }, { name: "nil credential", cred: nil, wantComet: false, - wantMask: false, - }, - { - name: "explicit mask flag", - cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, MaskUpstreamErrors: true}, - wantComet: false, - wantMask: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { assert.Equal(t, tt.wantComet, isCometAPICredential(tt.cred)) - assert.Equal(t, tt.wantMask, shouldMaskUpstreamErrors(tt.cred)) }) } } @@ -82,54 +69,53 @@ func TestIsSosanaCredential(t *testing.T) { name string cred *config.CredentialConfig wantSosana bool - wantMask bool }{ { name: "dedicated provider type", cred: &config.CredentialConfig{Type: config.ProviderTypeSosana}, wantSosana: true, - wantMask: true, }, { name: "sosana host fallback", cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://sosana.art"}, wantSosana: true, - wantMask: true, }, { name: "sosana name fallback", cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, Name: "sosana-art-images"}, wantSosana: true, - wantMask: true, }, { name: "sasana host fallback", cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://api.sasana.example/v1"}, wantSosana: true, - wantMask: true, }, { name: "regular openai", cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://api.openai.com"}, wantSosana: false, - wantMask: false, }, { name: "nil credential", cred: nil, wantSosana: false, - wantMask: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { assert.Equal(t, tt.wantSosana, isSosanaCredential(tt.cred)) - assert.Equal(t, tt.wantMask, shouldMaskUpstreamErrors(tt.cred)) }) } } +func TestShouldMaskUpstreamErrorsForKnownResellers(t *testing.T) { + assert.True(t, shouldMaskUpstreamErrors(&config.CredentialConfig{Type: config.ProviderTypeCometAPI})) + assert.True(t, shouldMaskUpstreamErrors(&config.CredentialConfig{Type: config.ProviderTypeSosana})) + assert.False(t, shouldMaskUpstreamErrors(&config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://api.openai.com"})) + assert.False(t, shouldMaskUpstreamErrors(nil)) +} + func TestAppendResponseBodyForLogs_CometKeepsMaskedFlagAndLogsBody(t *testing.T) { cred := &config.CredentialConfig{Type: config.ProviderTypeCometAPI} body := `{"error":{"code":"permission_denied","message":"` + strings.Repeat("model access denied ", 50) + `","type":"comet_api_error"}}` diff --git a/internal/proxy/proxy_log.go b/internal/proxy/proxy_log.go index 773055a7..fd35fd8e 100644 --- a/internal/proxy/proxy_log.go +++ b/internal/proxy/proxy_log.go @@ -52,7 +52,7 @@ func shouldMaskUpstreamErrors(cred *config.CredentialConfig) bool { if cred == nil { return false } - return cred.MaskUpstreamErrors || isCometAPICredential(cred) || isSosanaCredential(cred) + return isCometAPICredential(cred) || isSosanaCredential(cred) } func isCometAPICredential(cred *config.CredentialConfig) bool { diff --git a/internal/proxy/upstream_masking_test.go b/internal/proxy/upstream_masking_test.go index 9edf980a..4ad5dd77 100644 --- a/internal/proxy/upstream_masking_test.go +++ b/internal/proxy/upstream_masking_test.go @@ -22,13 +22,12 @@ func TestMaskedUpstreamError_DirectImageErrorDoesNotLeakProviderBody(t *testing. prx := NewTestProxyBuilder(). WithCredentials(config.CredentialConfig{ - Name: "sosana-art", - Type: config.ProviderTypeOpenAI, - BaseURL: upstream.URL, - APIKey: "upstream-key", - RPM: 100, - TPM: 10000, - MaskUpstreamErrors: true, + Name: "sosana-art", + Type: config.ProviderTypeOpenAI, + BaseURL: upstream.URL, + APIKey: "upstream-key", + RPM: 100, + TPM: 10000, }). Build() From b2e94a7530e8a3fe36b761ff29ad277a46ef6e38 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Tue, 7 Jul 2026 18:05:33 +0300 Subject: [PATCH 09/16] fix: align sosana image sizes with gemini table --- docs/providers/sosana.md | 4 +- internal/converter/sosana/images.go | 141 +++++++++++++++++------ internal/converter/sosana/images_test.go | 15 ++- internal/proxy/sosana_test.go | 6 +- 4 files changed, 121 insertions(+), 45 deletions(-) diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md index b0098016..c52b3ea0 100644 --- a/docs/providers/sosana.md +++ b/docs/providers/sosana.md @@ -51,7 +51,9 @@ completion requests. For production Sosana credentials, set the router VSELLM-owned rehosting before they can hide Sosana storage. Another provider may handle it through normal fallback routing. - `image_size` may be omitted or set to `1K`, `2K`, or `4K`; `0.5K` is not - routed to Sosana. + routed to Sosana. Pixel `size` values are accepted only when they match the + documented Gemini `image_size` + `aspect_ratio` table for `1K`, `2K`, or + `4K`. - `/v1/images/edits` accepts PNG input images only, up to 14 files, and sends them as `data:image/png;base64,...` values in Sosana `image_urls`. - Mask images are not supported. diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go index 8aea3306..c3a9a3cf 100644 --- a/internal/converter/sosana/images.go +++ b/internal/converter/sosana/images.go @@ -239,32 +239,10 @@ func PollURL(baseURL, uid string) string { } func SizeToAspectRatio(size string) string { - switch strings.TrimSpace(size) { - case "", "auto": - return "auto" - case "256x256", "512x512", "1024x1024", "2048x2048", "4096x4096": - return "1:1" - case "1024x1536", "2048x3072": - return "2:3" - case "1536x1024", "3072x2048": - return "3:2" - case "1024x1792", "1080x1920", "2048x3584", "4096x7168": - return "9:16" - case "1792x1024", "1920x1080", "3584x2048", "7168x4096": - return "16:9" - case "1024x768", "2048x1536", "4096x3072": - return "4:3" - case "768x1024", "1536x2048", "3072x4096": - return "3:4" - case "1024x819", "2048x1638", "4096x3276": - return "5:4" - case "819x1024", "1638x2048", "3276x4096": - return "4:5" - case "2016x864", "4032x1728": - return "21:9" - default: - return "auto" + if spec, ok := imageSpecFromExactSize(size); ok { + return spec.aspectRatio } + return "auto" } func aspectRatio(explicit, ratio, size string) string { @@ -304,19 +282,110 @@ func normalizeImageSize(raw string) (string, bool) { } func imageSizeFromExactSize(size string) (string, bool) { + if spec, ok := imageSpecFromExactSize(size); ok { + return spec.imageSize, true + } + return "", false +} + +type exactImageSpec struct { + imageSize string + aspectRatio string +} + +func imageSpecFromExactSize(size string) (exactImageSpec, bool) { switch strings.TrimSpace(size) { - case "", "auto", - "256x256", "512x512", "1024x1024", "1024x1536", "1536x1024", - "1024x1792", "1792x1024", "1080x1920", "1920x1080", - "1024x768", "768x1024", "1024x819", "819x1024", "2016x864": - return "1K", true - case "2048x2048", "2048x3072", "3072x2048", "2048x3584", "3584x2048", - "2048x1536", "1536x2048", "2048x1638", "1638x2048", "4032x1728": - return "2K", true - case "4096x4096", "4096x7168", "7168x4096", "4096x3072", "3072x4096", "4096x3276", "3276x4096": - return "4K", true + case "", "auto": + return exactImageSpec{imageSize: "1K", aspectRatio: "auto"}, true + + case "1024x1024": + return exactImageSpec{imageSize: "1K", aspectRatio: "1:1"}, true + case "512x2048": + return exactImageSpec{imageSize: "1K", aspectRatio: "1:4"}, true + case "384x3072": + return exactImageSpec{imageSize: "1K", aspectRatio: "1:8"}, true + case "848x1264": + return exactImageSpec{imageSize: "1K", aspectRatio: "2:3"}, true + case "1264x848": + return exactImageSpec{imageSize: "1K", aspectRatio: "3:2"}, true + case "896x1200": + return exactImageSpec{imageSize: "1K", aspectRatio: "3:4"}, true + case "2048x512": + return exactImageSpec{imageSize: "1K", aspectRatio: "4:1"}, true + case "1200x896": + return exactImageSpec{imageSize: "1K", aspectRatio: "4:3"}, true + case "928x1152": + return exactImageSpec{imageSize: "1K", aspectRatio: "4:5"}, true + case "1152x928": + return exactImageSpec{imageSize: "1K", aspectRatio: "5:4"}, true + case "3072x384": + return exactImageSpec{imageSize: "1K", aspectRatio: "8:1"}, true + case "768x1376": + return exactImageSpec{imageSize: "1K", aspectRatio: "9:16"}, true + case "1376x768": + return exactImageSpec{imageSize: "1K", aspectRatio: "16:9"}, true + case "1584x672": + return exactImageSpec{imageSize: "1K", aspectRatio: "21:9"}, true + + case "2048x2048": + return exactImageSpec{imageSize: "2K", aspectRatio: "1:1"}, true + case "1024x4096": + return exactImageSpec{imageSize: "2K", aspectRatio: "1:4"}, true + case "768x6144": + return exactImageSpec{imageSize: "2K", aspectRatio: "1:8"}, true + case "1696x2528": + return exactImageSpec{imageSize: "2K", aspectRatio: "2:3"}, true + case "2528x1696": + return exactImageSpec{imageSize: "2K", aspectRatio: "3:2"}, true + case "1792x2400": + return exactImageSpec{imageSize: "2K", aspectRatio: "3:4"}, true + case "4096x1024": + return exactImageSpec{imageSize: "2K", aspectRatio: "4:1"}, true + case "2400x1792": + return exactImageSpec{imageSize: "2K", aspectRatio: "4:3"}, true + case "1856x2304": + return exactImageSpec{imageSize: "2K", aspectRatio: "4:5"}, true + case "2304x1856": + return exactImageSpec{imageSize: "2K", aspectRatio: "5:4"}, true + case "6144x768": + return exactImageSpec{imageSize: "2K", aspectRatio: "8:1"}, true + case "1536x2752": + return exactImageSpec{imageSize: "2K", aspectRatio: "9:16"}, true + case "2752x1536": + return exactImageSpec{imageSize: "2K", aspectRatio: "16:9"}, true + case "3168x1344": + return exactImageSpec{imageSize: "2K", aspectRatio: "21:9"}, true + + case "4096x4096": + return exactImageSpec{imageSize: "4K", aspectRatio: "1:1"}, true + case "2048x8192": + return exactImageSpec{imageSize: "4K", aspectRatio: "1:4"}, true + case "1536x12288": + return exactImageSpec{imageSize: "4K", aspectRatio: "1:8"}, true + case "3392x5056": + return exactImageSpec{imageSize: "4K", aspectRatio: "2:3"}, true + case "5056x3392": + return exactImageSpec{imageSize: "4K", aspectRatio: "3:2"}, true + case "3584x4800": + return exactImageSpec{imageSize: "4K", aspectRatio: "3:4"}, true + case "8192x2048": + return exactImageSpec{imageSize: "4K", aspectRatio: "4:1"}, true + case "4800x3584": + return exactImageSpec{imageSize: "4K", aspectRatio: "4:3"}, true + case "3712x4608": + return exactImageSpec{imageSize: "4K", aspectRatio: "4:5"}, true + case "4608x3712": + return exactImageSpec{imageSize: "4K", aspectRatio: "5:4"}, true + case "12288x1536": + return exactImageSpec{imageSize: "4K", aspectRatio: "8:1"}, true + case "3072x5504": + return exactImageSpec{imageSize: "4K", aspectRatio: "9:16"}, true + case "5504x3072": + return exactImageSpec{imageSize: "4K", aspectRatio: "16:9"}, true + case "6336x2688": + return exactImageSpec{imageSize: "4K", aspectRatio: "21:9"}, true default: - return "", false + return exactImageSpec{}, false } } diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go index 9b8dcea9..b817ccfe 100644 --- a/internal/converter/sosana/images_test.go +++ b/internal/converter/sosana/images_test.go @@ -23,11 +23,14 @@ func TestImageGenerationRequest(t *testing.T) { wantSize string wantErrPart string }{ - {name: "square", size: "1024x1024", wantAspect: "1:1", wantSize: "1K"}, - {name: "wide", size: "1792x1024", wantAspect: "16:9", wantSize: "1K"}, - {name: "portrait", size: "1024x1792", wantAspect: "9:16", wantSize: "1K"}, - {name: "two k", size: "2048x2048", wantAspect: "1:1", wantSize: "2K"}, - {name: "four k", size: "4096x4096", wantAspect: "1:1", wantSize: "4K"}, + {name: "one k square", size: "1024x1024", wantAspect: "1:1", wantSize: "1K"}, + {name: "one k wide", size: "1376x768", wantAspect: "16:9", wantSize: "1K"}, + {name: "one k portrait", size: "768x1376", wantAspect: "9:16", wantSize: "1K"}, + {name: "one k tall", size: "512x2048", wantAspect: "1:4", wantSize: "1K"}, + {name: "two k square", size: "2048x2048", wantAspect: "1:1", wantSize: "2K"}, + {name: "two k wide", size: "2752x1536", wantAspect: "16:9", wantSize: "2K"}, + {name: "four k square", size: "4096x4096", wantAspect: "1:1", wantSize: "4K"}, + {name: "four k ultra wide", size: "6336x2688", wantAspect: "21:9", wantSize: "4K"}, } for _, tt := range tests { @@ -107,6 +110,8 @@ func TestImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { {name: "quality auto", body: `{"model":"banana-2-1k-compliant","prompt":"draw","quality":"auto"}`, want: "quality"}, {name: "messages", body: `{"model":"banana-2-1k-compliant","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, {name: "image size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_size":"0.5K"}`, want: "image_size"}, + {name: "exact size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"512x512"}`, want: "size"}, + {name: "legacy openai size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"1792x1024"}`, want: "size"}, {name: "unknown size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"333x777"}`, want: "size"}, {name: "reference images", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, } diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 3cee7ee0..50c55e26 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -418,7 +418,7 @@ func TestProxyRequest_IncompatibleImageEditJPEGWithSosanaReturnsLocalError(t *te assert.False(t, called) } -func TestProxyRequest_IncompatibleSosanaImageRequestRoutesToNextPrimary(t *testing.T) { +func TestProxyRequest_SosanaHalfKPixelSizeRoutesToNextPrimary(t *testing.T) { sosanaCalled := false sosanaUpstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { sosanaCalled = true @@ -434,7 +434,7 @@ func TestProxyRequest_IncompatibleSosanaImageRequestRoutesToNextPrimary(t *testi var req map[string]any require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) assert.Equal(t, "banana-2-1k-compliant", req["model"]) - assert.Contains(t, req, "tools") + assert.Equal(t, "512x512", req["size"]) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"created":1782478551,"data":[{"b64_json":"fallback-primary-image"}]}`)) })) @@ -446,7 +446,7 @@ func TestProxyRequest_IncompatibleSosanaImageRequestRoutesToNextPrimary(t *testi config.CredentialConfig{Name: "next-primary", Type: config.ProviderTypeProxy, BaseURL: nextPrimary.URL, APIKey: "next-key", RPM: 100, TPM: 10000}, ). Build() - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`)) + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","size":"512x512"}`)) req.Header.Set("Authorization", "Bearer master-key") req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() From 87a2c5a68d4efb462ac770e6f25f5caf1bc81a7b Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Tue, 7 Jul 2026 18:38:41 +0300 Subject: [PATCH 10/16] docs: clarify sosana api key auth --- docs/providers/index.md | 2 +- docs/providers/sosana.md | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/providers/index.md b/docs/providers/index.md index 69db9fbb..958a14fa 100644 --- a/docs/providers/index.md +++ b/docs/providers/index.md @@ -9,7 +9,7 @@ Auto AI Router supports multiple LLM providers. Each provider type has its own a | [OpenAI](openai.md) | `openai` | `api_key`, `base_url` | API Key | | [Anthropic](anthropic.md) | `anthropic` | `api_key`, `base_url` | API Key | | [Comet API](cometapi.md) | `cometapi` | `api_key`, `base_url` | API Key | -| [Sosana.art](sosana.md) | `sosana` | `api_key`, `base_url` | Bearer Token | +| [Sosana.art](sosana.md) | `sosana` | `api_key`, `base_url` | API Key via Bearer | | [AWS Bedrock](bedrock.md) | `bedrock` | `api_key`, `base_url` | Bearer Token | | [Vertex AI](vertex.md) | `vertex-ai` | `project_id`, `location`, `credentials_file` or `credentials_json` | OAuth2 / Service Account | | [Gemini AI Studio](gemini.md) | `gemini` | `api_key`, `base_url` | API Key | diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md index c52b3ea0..0c877bd3 100644 --- a/docs/providers/sosana.md +++ b/docs/providers/sosana.md @@ -27,6 +27,9 @@ models: tpm: -1 ``` +The credential value is configured as `api_key`. The router sends it to Sosana +as `Authorization: Bearer `, matching Sosana's API contract. + The dynamic model template maps `image_size` to Sosana's concrete image models: `banana-2-1k-compliant`, `banana-2-2k-compliant`, and `banana-2-4k-compliant`. If `image_size` is omitted, the router uses `1K`. From bbc8ddb3088b57d079549ab06aeeea7a240051f5 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Fri, 10 Jul 2026 18:23:51 +0300 Subject: [PATCH 11/16] feat: finalize Sosana billing and scoped routing --- internal/litellmdb/model_table/model_table.go | 5 +- .../litellmdb/model_table/model_table_test.go | 19 ++++ internal/models/manager.go | 1 + internal/models/price_calculator.go | 10 +- internal/models/price_calculator_test.go | 26 +++++ internal/models/price_loader_test.go | 16 ++++ internal/proxy/proxy_helpers.go | 2 + internal/proxy/proxy_helpers_test.go | 16 ++++ internal/proxy/sosana.go | 7 +- internal/proxy/sosana_routing.go | 6 +- internal/proxy/sosana_test.go | 96 ++++++++++++++++++- 11 files changed, 191 insertions(+), 13 deletions(-) diff --git a/internal/litellmdb/model_table/model_table.go b/internal/litellmdb/model_table/model_table.go index 406c57ce..cacd6db7 100644 --- a/internal/litellmdb/model_table/model_table.go +++ b/internal/litellmdb/model_table/model_table.go @@ -391,7 +391,7 @@ func convertPricingToModelPrice(p *queries.CustomPricingLiteLLMParams) *manager. return nil } if p.InputCostPerToken == nil && p.OutputCostPerToken == nil && - p.OutputCostPerImage == nil && p.OutputCostPerImageToken == nil { + p.InputCostPerImage == nil && p.OutputCostPerImage == nil && p.OutputCostPerImageToken == nil { return nil } @@ -435,6 +435,9 @@ func convertPricingToModelPrice(p *queries.CustomPricingLiteLLMParams) *manager. if p.CacheCreationInputTokenCostAbove272kTokens != nil { price.CacheCreationInputTokenCostAbove272k = *p.CacheCreationInputTokenCostAbove272kTokens } + if p.InputCostPerImage != nil { + price.InputCostPerImage = *p.InputCostPerImage + } if p.OutputCostPerImage != nil { price.OutputCostPerImage = *p.OutputCostPerImage } diff --git a/internal/litellmdb/model_table/model_table_test.go b/internal/litellmdb/model_table/model_table_test.go index 9ed9b678..a0f66db6 100644 --- a/internal/litellmdb/model_table/model_table_test.go +++ b/internal/litellmdb/model_table/model_table_test.go @@ -124,6 +124,7 @@ func TestConvertPricingToModelPrice(t *testing.T) { outputReasoning := 0.03 cacheRead := 0.04 cacheCreation := 0.05 + inputImage := 0.4 outputImage := 0.5 outputImageToken := 0.6 inputAbove200k := 0.07 @@ -134,6 +135,7 @@ func TestConvertPricingToModelPrice(t *testing.T) { OutputCostPerReasoningToken: &outputReasoning, CacheReadInputTokenCost: &cacheRead, CacheCreationInputTokenCost: &cacheCreation, + InputCostPerImage: &inputImage, OutputCostPerImage: &outputImage, OutputCostPerImageToken: &outputImageToken, InputCostPerTokenAbove200kTokens: &inputAbove200k, @@ -146,6 +148,7 @@ func TestConvertPricingToModelPrice(t *testing.T) { assert.Equal(t, outputReasoning, price.OutputCostPerReasoningToken) assert.Equal(t, cacheRead, price.InputCostPerCachedToken) assert.Equal(t, cacheCreation, price.CacheCreationInputTokenCost) + assert.Equal(t, inputImage, price.InputCostPerImage) assert.Equal(t, outputImage, price.OutputCostPerImage) assert.Equal(t, outputImageToken, price.OutputCostPerImageToken) assert.Equal(t, inputAbove200k, price.InputCostPerTokenAbove200k) @@ -167,6 +170,19 @@ func TestConvertPricingToModelPrice_ImageOnly(t *testing.T) { assert.Equal(t, 0.0, price.OutputCostPerToken) } +func TestConvertPricingToModelPrice_InputImageOnly(t *testing.T) { + inputImage := 0.088113 + + price := convertPricingToModelPrice(&queries.CustomPricingLiteLLMParams{ + InputCostPerImage: &inputImage, + }) + + require.NotNil(t, price) + assert.Equal(t, inputImage, price.InputCostPerImage) + assert.Equal(t, 0.0, price.InputCostPerToken) + assert.Equal(t, 0.0, price.OutputCostPerToken) +} + func TestConvertPricingToModelPrice_AllFields(t *testing.T) { input := 0.01 output := 0.02 @@ -181,6 +197,7 @@ func TestConvertPricingToModelPrice_AllFields(t *testing.T) { cacheCreation := 0.085 cacheReadAbove272k := 0.086 cacheCreationAbove272k := 0.087 + inputImage := 0.088 outputImage := 0.09 outputImageToken := 0.10 @@ -198,6 +215,7 @@ func TestConvertPricingToModelPrice_AllFields(t *testing.T) { CacheCreationInputTokenCost: &cacheCreation, CacheReadInputTokenCostAbove272kTokens: &cacheReadAbove272k, CacheCreationInputTokenCostAbove272kTokens: &cacheCreationAbove272k, + InputCostPerImage: &inputImage, OutputCostPerImage: &outputImage, OutputCostPerImageToken: &outputImageToken, }) @@ -216,6 +234,7 @@ func TestConvertPricingToModelPrice_AllFields(t *testing.T) { assert.Equal(t, cacheCreation, price.CacheCreationInputTokenCost) assert.Equal(t, cacheReadAbove272k, price.CacheReadInputTokenCostAbove272k) assert.Equal(t, cacheCreationAbove272k, price.CacheCreationInputTokenCostAbove272k) + assert.Equal(t, inputImage, price.InputCostPerImage) assert.Equal(t, outputImage, price.OutputCostPerImage) assert.Equal(t, outputImageToken, price.OutputCostPerImageToken) } diff --git a/internal/models/manager.go b/internal/models/manager.go index 3c98f2fc..41d58098 100644 --- a/internal/models/manager.go +++ b/internal/models/manager.go @@ -57,6 +57,7 @@ type ModelPrice struct { OutputCostPerPredictionToken float64 `json:"output_cost_per_prediction_token,omitempty"` // Vision/Images cost per image (not per token) + InputCostPerImage float64 `json:"input_cost_per_image,omitempty"` OutputCostPerImage float64 `json:"output_cost_per_image,omitempty"` } diff --git a/internal/models/price_calculator.go b/internal/models/price_calculator.go index 7e2bc712..d20d8752 100644 --- a/internal/models/price_calculator.go +++ b/internal/models/price_calculator.go @@ -139,15 +139,13 @@ func CalculateTokenCosts(usage *converter.TokenUsage, price *ModelPrice) *conver // Rejected prediction tokens count as regular output tokens costs.PredictionCost += float64(usage.RejectedPredictionTokens) * outputCostPerToken - // Image cost calculation: supports both per-image and per-image-token pricing - // Priority: 1) Per-image cost if available (typical for image generation APIs) - // 2) Per-image-token cost as fallback (rarely used for image generation) - // 3) Default: $0 if neither is configured + // Image cost calculation: prefer the explicit output price, then LiteLLM's + // input_cost_per_image convention for image generation, then token fallback. if usage.ImageCount > 0 && price.OutputCostPerImage > 0 { - // Per-image cost (e.g., $0.02 per image) costs.ImageCost = float64(usage.ImageCount) * price.OutputCostPerImage + } else if usage.ImageCount > 0 && price.InputCostPerImage > 0 { + costs.ImageCost = float64(usage.ImageCount) * price.InputCostPerImage } else if usage.ImageCount > 0 && price.OutputCostPerImageToken > 0 { - // Per-image-token cost fallback (rarely used for image generation) costs.ImageCost = float64(usage.ImageCount) * price.OutputCostPerImageToken } diff --git a/internal/models/price_calculator_test.go b/internal/models/price_calculator_test.go index bd5a5490..ede0b374 100644 --- a/internal/models/price_calculator_test.go +++ b/internal/models/price_calculator_test.go @@ -5,6 +5,7 @@ import ( "github.com/mixaill76/auto_ai_router/internal/converter" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestCalculateTokenCosts_RegularTokensOnly(t *testing.T) { @@ -249,6 +250,31 @@ func TestCalculateTokenCosts_ImageCount(t *testing.T) { assert.InDelta(t, 0.10, costs.TotalCost, 1e-9) } +func TestCalculateTokenCosts_ImageCountUsesInputImageFallback(t *testing.T) { + usage := &converter.TokenUsage{ImageCount: 1} + price := &ModelPrice{InputCostPerImage: 0.088113} + + costs := CalculateTokenCosts(usage, price) + + require.NotNil(t, costs) + assert.InDelta(t, 0.088113, costs.ImageCost, 1e-12) + assert.InDelta(t, 0.088113, costs.TotalCost, 1e-12) +} + +func TestCalculateTokenCosts_OutputImagePriceTakesPriority(t *testing.T) { + usage := &converter.TokenUsage{ImageCount: 1} + price := &ModelPrice{ + InputCostPerImage: 0.088113, + OutputCostPerImage: 0.09, + } + + costs := CalculateTokenCosts(usage, price) + + require.NotNil(t, costs) + assert.InDelta(t, 0.09, costs.ImageCost, 1e-12) + assert.InDelta(t, 0.09, costs.TotalCost, 1e-12) +} + func TestCalculateTokenCosts_ImageCountUsesImageTokenFallback(t *testing.T) { usage := &converter.TokenUsage{ ImageCount: 3, diff --git a/internal/models/price_loader_test.go b/internal/models/price_loader_test.go index 2db75968..d8a086b6 100644 --- a/internal/models/price_loader_test.go +++ b/internal/models/price_loader_test.go @@ -67,6 +67,22 @@ func TestLoadModelPrices_GPT56LongContext(t *testing.T) { assert.InDelta(t, 0.00001625, price.CacheCreationInputTokenCostAbove272k, 1e-12) } +func TestLoadModelPrices_InputCostPerImage(t *testing.T) { + filePath := filepath.Join(t.TempDir(), "prices.json") + pricesJSON := `{ + "gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.088113 + } + }` + require.NoError(t, os.WriteFile(filePath, []byte(pricesJSON), 0o600)) + + prices, err := LoadModelPrices(filePath) + require.NoError(t, err) + price := prices["gemini-3.1-flash-image-preview"] + require.NotNil(t, price) + assert.InDelta(t, 0.088113, price.InputCostPerImage, 1e-12) +} + func TestLoadModelPrices_FromFilePath(t *testing.T) { // Create a temporary file with valid JSON tmpDir := t.TempDir() diff --git a/internal/proxy/proxy_helpers.go b/internal/proxy/proxy_helpers.go index ad8e92b4..69db035b 100644 --- a/internal/proxy/proxy_helpers.go +++ b/internal/proxy/proxy_helpers.go @@ -172,6 +172,7 @@ func buildMetadata(hashedToken string, tokenInfo *litellmdb.TokenInfo, errorMsg "total_tokens": usage.Total(), "prompt_tokens": usage.PromptTokens, "completion_tokens": usage.CompletionTokens, + "image_count": usage.ImageCount, "prompt_tokens_details": promptTokensDetails, "completion_tokens_details": completionTokensDetails, } @@ -197,6 +198,7 @@ func buildMetadata(hashedToken string, tokenInfo *litellmdb.TokenInfo, errorMsg costBreakdown = map[string]interface{}{ "input_cost": costs.InputCost, "output_cost": costs.OutputCost, + "image_cost": costs.ImageCost, "cached_input_cost": costs.CachedInputCost, "cache_creation_cost": costs.CacheCreationCost, "total_cost": costs.TotalCost, diff --git a/internal/proxy/proxy_helpers_test.go b/internal/proxy/proxy_helpers_test.go index 87def0f3..f33381b5 100644 --- a/internal/proxy/proxy_helpers_test.go +++ b/internal/proxy/proxy_helpers_test.go @@ -12,6 +12,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/mixaill76/auto_ai_router/internal/converter" "github.com/mixaill76/auto_ai_router/internal/litellmdb" ) @@ -235,6 +236,21 @@ func TestBuildMetadata(t *testing.T) { assert.Equal(t, float64(429), errInfo["error_code"]) assert.Equal(t, "RateLimitError", errInfo["error_class"]) }) + + t.Run("with_image_usage_and_cost", func(t *testing.T) { + usage := &converter.TokenUsage{ImageCount: 1} + costs := &converter.TokenCosts{ImageCost: 0.088113, TotalCost: 0.088113} + result := buildMetadata("hashed-image", nil, "", http.StatusOK, usage, "", costs, "image-model", 0) + + var metadata map[string]interface{} + require.NoError(t, json.Unmarshal([]byte(result), &metadata)) + usageObject, ok := metadata["usage_object"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, float64(1), usageObject["image_count"]) + costBreakdown, ok := metadata["cost_breakdown"].(map[string]interface{}) + require.True(t, ok) + assert.InDelta(t, 0.088113, costBreakdown["image_cost"].(float64), 1e-12) + }) } func TestExtractEndUser(t *testing.T) { diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index fb07732c..528fdcf9 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -70,7 +70,12 @@ func (p *Proxy) handleSosanaRequest( triedCreds := GetTried(r.Context()) for attempt := 0; attempt <= p.maxProviderRetries; attempt++ { if attempt > 0 { - nextCred, err := p.balancer.NextSameTypeForModelExcluding(modelID, config.ProviderTypeSosana, triedCreds) + nextCred, err := p.balancer.NextSameTypeForModelExcludingScoped( + modelID, + config.ProviderTypeSosana, + triedCreds, + logCtx.Scope, + ) if err != nil { p.logger.DebugContext(r.Context(), "No more Sosana credentials for retry", "model", modelID, "attempt", attempt, "error", err) diff --git a/internal/proxy/sosana_routing.go b/internal/proxy/sosana_routing.go index fc5484ce..69900690 100644 --- a/internal/proxy/sosana_routing.go +++ b/internal/proxy/sosana_routing.go @@ -6,6 +6,7 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter/sosana" + "github.com/mixaill76/auto_ai_router/internal/scope" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" ) @@ -39,7 +40,7 @@ func (p *Proxy) applySosanaCompatibilityRouting( return true } - nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, reason) + nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, logCtx.Scope, reason) if routed { *cred = nextCred *body = nextReq.body @@ -97,13 +98,14 @@ func (p *Proxy) nextPrimaryAfterUnsupportedSosana( prepared *orchestratedRequest, modelID string, currentCred *config.CredentialConfig, + visibility scope.Context, reason string, ) (*config.CredentialConfig, credentialPreparedRequest, bool) { triedCreds := GetTried(r.Context()) triedCreds[currentCred.Name] = true for attempts := 0; attempts < 128; attempts++ { - candidate, err := p.balancer.NextForModelExcluding(modelID, triedCreds) + candidate, err := p.balancer.NextForModelExcludingScoped(modelID, triedCreds, visibility) if err != nil { p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", "model", modelID, diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 50c55e26..71f58182 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -148,7 +148,7 @@ func TestProxyRequest_SosanaImageGenerationLogsLiteLLMImageSpend(t *testing.T) { spendManager := newCapturedSpendManager() priceRegistry := aimodels.NewModelPriceRegistry() priceRegistry.Update(map[string]*aimodels.ModelPrice{ - "banana-2-1k-compliant": {OutputCostPerImage: 0.07}, + "banana-2-1k-compliant": {InputCostPerImage: 0.088113}, }) prx := newSosanaTestProxy(upstream.URL, nil) @@ -174,12 +174,16 @@ func TestProxyRequest_SosanaImageGenerationLogsLiteLLMImageSpend(t *testing.T) { assert.Equal(t, "sosana", entry.CustomLLMProvider) assert.Equal(t, "success", entry.Status) assert.Equal(t, 0, entry.TotalTokens) - assert.InDelta(t, 0.07, entry.Spend, 0.0000001) + assert.InDelta(t, 0.088113, entry.Spend, 1e-12) var metadata map[string]any require.NoError(t, json.Unmarshal([]byte(entry.Metadata), &metadata)) + usageObject, ok := metadata["usage_object"].(map[string]any) + require.True(t, ok) + assert.Equal(t, float64(1), usageObject["image_count"]) costBreakdown, ok := metadata["cost_breakdown"].(map[string]any) require.True(t, ok) - assert.InDelta(t, 0.07, costBreakdown["total_cost"].(float64), 0.0000001) + assert.InDelta(t, 0.088113, costBreakdown["image_cost"].(float64), 1e-12) + assert.InDelta(t, 0.088113, costBreakdown["total_cost"].(float64), 1e-12) } func TestProxyRequest_SosanaImageGenerationUsesConcreteTierPrice(t *testing.T) { @@ -459,6 +463,48 @@ func TestProxyRequest_SosanaHalfKPixelSizeRoutesToNextPrimary(t *testing.T) { assert.Contains(t, w.Body.String(), "fallback-primary-image") } +func TestProxyRequest_IncompatibleSosanaRequestDoesNotCrossCredentialScope(t *testing.T) { + sosanaCalled := false + sosanaUpstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sosanaCalled = true + w.WriteHeader(http.StatusOK) + })) + defer sosanaUpstream.Close() + + hiddenPrimaryCalled := false + hiddenPrimary := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hiddenPrimaryCalled = true + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"created":1782478551,"data":[{"b64_json":"hidden-image"}]}`)) + })) + defer hiddenPrimary.Close() + + prx := NewTestProxyBuilder(). + WithCredentials( + config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: sosanaUpstream.URL, APIKey: "sosana-key", RPM: 100, TPM: 10000, Scopes: []string{"team-a"}}, + config.CredentialConfig{Name: "hidden-primary", Type: config.ProviderTypeProxy, BaseURL: hiddenPrimary.URL, APIKey: "hidden-key", RPM: 100, TPM: 10000, Scopes: []string{"team-b"}}, + ). + Build() + prx.LiteLLMDB = scopeTestDB{ + NoopManager: litellmdb.NewNoopManager(), + info: &litellmmodels.TokenInfo{Metadata: map[string]interface{}{ + "air_scopes": []interface{}{"team-a"}, + }}, + } + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`)) + req.Header.Set("Authorization", "Bearer team-a-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), unsupportedImageProviderRequestMessage) + assert.False(t, sosanaCalled) + assert.False(t, hiddenPrimaryCalled) +} + func TestProxyRequest_IncompatibleSosanaImageRequestRoutesToFallbackProxy(t *testing.T) { sosanaCalled := false sosanaUpstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -575,6 +621,50 @@ func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { assert.NotContains(t, w.Body.String(), imageServer.URL) } +func TestProxyRequest_SosanaRetryDoesNotCrossCredentialScope(t *testing.T) { + hiddenCredentialCalled := false + upstream := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/api/banana/create-async", r.URL.Path) + switch r.Header.Get("Authorization") { + case "Bearer sosana-key-a": + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"detail":"rate limited"}`)) + case "Bearer sosana-key-b": + hiddenCredentialCalled = true + _, _ = w.Write([]byte(`{"uid":"hidden-task","status":"FAILED","error":"must not be called"}`)) + default: + t.Fatalf("unexpected authorization header: %q", r.Header.Get("Authorization")) + } + })) + defer upstream.Close() + + prx := NewTestProxyBuilder(). + WithCredentials( + config.CredentialConfig{Name: "sosana-a", Type: config.ProviderTypeSosana, BaseURL: upstream.URL, APIKey: "sosana-key-a", RPM: 100, TPM: 10000, Scopes: []string{"team-a"}}, + config.CredentialConfig{Name: "sosana-b", Type: config.ProviderTypeSosana, BaseURL: upstream.URL, APIKey: "sosana-key-b", RPM: 100, TPM: 10000, Scopes: []string{"team-b"}}, + ). + WithMaxProviderRetries(1). + Build() + prx.LiteLLMDB = scopeTestDB{ + NoopManager: litellmdb.NewNoopManager(), + info: &litellmmodels.TokenInfo{Metadata: map[string]interface{}{ + "air_scopes": []interface{}{"team-a"}, + }}, + } + + req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"model":"banana-2-1k-compliant","prompt":"draw","n":1}`)) + req.Header.Set("Authorization", "Bearer team-a-key") + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusTooManyRequests, w.Code) + assert.False(t, hiddenCredentialCalled) + assert.Contains(t, w.Body.String(), "Upstream provider error") + assert.NotContains(t, w.Body.String(), "rate limited") +} + func TestProxyRequest_SosanaDoesNotRetryCreateTransportError(t *testing.T) { deadServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) deadURL := deadServer.URL From 9ef927836803f0f0ad9c3290c4f487f4f69f3132 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Wed, 22 Jul 2026 18:29:13 +0300 Subject: [PATCH 12/16] fix: stabilize sosana timeout test --- internal/proxy/sosana_test.go | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index e3acaf0c..8afc2d75 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -1078,7 +1078,17 @@ func TestProxyRequest_SosanaTimeoutMasked(t *testing.T) { assert.Equal(t, http.StatusRequestTimeout, w.Code) assert.Contains(t, w.Body.String(), "Upstream provider error") - assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.NotContains(t, w.Body.String(), "task-1") + logText := logBuf.String() + assert.Contains(t, logText, "context deadline exceeded") + assert.True(t, + strings.Contains(logText, "Sosana task polling timed out") || + strings.Contains(logText, "Sosana upstream request failed"), + "unexpected timeout log: %s", logText, + ) + if strings.Contains(logText, "response_body=") { + assert.Contains(t, logText, "response_body_masked=true") + } } var sosanaResultPNG = []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} From afac9e9ca23730d87bdd9cad0ed3f9527a94e08d Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Thu, 23 Jul 2026 20:01:30 +0300 Subject: [PATCH 13/16] refactor: keep sosana image conversion in proxy --- internal/converter/sosana/compatibility.go | 310 ------- internal/converter/sosana/images.go | 446 ---------- internal/converter/sosana/images_test.go | 293 ------- internal/proxy/sosana.go | 939 ++++++++++++++++++++- internal/proxy/sosana_routing.go | 158 ---- internal/proxy/sosana_test.go | 255 +++++- 6 files changed, 1161 insertions(+), 1240 deletions(-) delete mode 100644 internal/converter/sosana/compatibility.go delete mode 100644 internal/converter/sosana/images.go delete mode 100644 internal/converter/sosana/images_test.go delete mode 100644 internal/proxy/sosana_routing.go diff --git a/internal/converter/sosana/compatibility.go b/internal/converter/sosana/compatibility.go deleted file mode 100644 index 453c1adc..00000000 --- a/internal/converter/sosana/compatibility.go +++ /dev/null @@ -1,310 +0,0 @@ -package sosana - -import ( - "bytes" - "encoding/json" - "fmt" - "mime" - "mime/multipart" - "strings" -) - -const maxInputImages = 14 - -var unsupportedImageFields = []string{ - "tools", - "tool_choice", - "google_search", - "thinking_level", - "thinking_budget", - "thinking_config", - "thinking", - "reasoning_effort", - "generation_config", - "temperature", - "top_p", - "top_k", - "seed", - "max_tokens", - "stop", - "stream", - "messages", - "extra_body", - "image", - "images", - "image_urls", - "reference_images", -} - -// UnsupportedRequest returns a short reason when a request needs image features -// that Sosana Banana does not expose in its public API. -func UnsupportedRequest(path string, body []byte, contentType string) string { - switch { - case strings.Contains(path, "/images/generations"): - return unsupportedGenerationRequest(body) - case strings.Contains(path, "/images/edits"): - return unsupportedEditRequest(body, contentType) - default: - return "endpoint is unsupported" - } -} - -func UnsupportedModel(modelID string) string { - if supportedSosanaModel(modelID) { - return "" - } - return "model is unsupported" -} - -func unsupportedGenerationRequest(body []byte) string { - var raw map[string]json.RawMessage - if err := json.Unmarshal(body, &raw); err != nil { - return "" - } - return unsupportedImageFieldsInJSON(raw) -} - -func unsupportedEditRequest(body []byte, contentType string) string { - mediaType, params, err := mime.ParseMediaType(contentType) - if err != nil || !strings.HasPrefix(mediaType, "multipart/form-data") { - return "" - } - boundary := params["boundary"] - if boundary == "" { - return "" - } - - fields := make(map[string]string) - imageCount := 0 - reader := multipart.NewReader(bytes.NewReader(body), boundary) - for { - part, err := reader.NextPart() - if err != nil { - break - } - - formName := part.FormName() - if formName == "" { - continue - } - data, err := readLimited(part, maxMultipartImageBytes) - if err != nil { - return err.Error() - } - if part.FileName() == "" { - fields[formName] = strings.TrimSpace(string(data)) - continue - } - if formName == "mask" { - return "mask is unsupported" - } - if formName != "image" && formName != "images" && formName != "image[]" { - continue - } - imageCount++ - if detectImageMIMEType(part.Header.Get("Content-Type"), data) != "image/png" { - return "only PNG input images are supported" - } - } - if imageCount > maxInputImages { - return "too many input images" - } - return unsupportedImageFieldsInForm(fields) -} - -func unsupportedImageFieldsInJSON(raw map[string]json.RawMessage) string { - if reason := unsupportedJSONImageCount(raw["n"]); reason != "" { - return reason - } - if reason := unsupportedJSONResponseFormat(raw["response_format"]); reason != "" { - return reason - } - if reason := unsupportedJSONOutputFormat(raw["output_format"]); reason != "" { - return reason - } - if reason := unsupportedJSONImageSize(raw["image_size"]); reason != "" { - return reason - } - if reason := unsupportedJSONExactSize(raw["size"]); reason != "" { - return reason - } - for _, field := range []string{"quality", "style", "background", "moderation"} { - if hasJSONValue(raw[field]) { - return field + " is unsupported" - } - } - if hasJSONValue(raw["output_compression"]) { - return "output_compression is unsupported" - } - for _, field := range unsupportedImageFields { - if hasJSONValue(raw[field]) { - return field + " is unsupported" - } - } - return "" -} - -func unsupportedImageFieldsInForm(fields map[string]string) string { - if err := validateImageCountString(fields["n"]); err != nil { - return err.Error() - } - if err := validateResponseFormat(fields["response_format"]); err != nil { - return err.Error() - } - if err := validateOutputFormat(fields["output_format"]); err != nil { - return err.Error() - } - if reason := unsupportedFormImageSize(fields["image_size"]); reason != "" { - return reason - } - if reason := unsupportedFormExactSize(fields["size"]); reason != "" { - return reason - } - for _, field := range []string{"quality", "style", "background", "moderation"} { - if strings.TrimSpace(fields[field]) != "" { - return field + " is unsupported" - } - } - if _, ok := fields["output_compression"]; ok { - return "output_compression is unsupported" - } - for _, field := range unsupportedImageFields { - if strings.TrimSpace(fields[field]) != "" { - return field + " is unsupported" - } - } - return "" -} - -func unsupportedJSONImageCount(raw json.RawMessage) string { - if !hasJSONValue(raw) { - return "" - } - var n int - if err := json.Unmarshal(raw, &n); err == nil { - if n == 1 { - return "" - } - return "image requests support n=1 only" - } - var f float64 - if err := json.Unmarshal(raw, &f); err == nil { - if f == 1 { - return "" - } - return "image requests support n=1 only" - } - return "invalid image count" -} - -func unsupportedJSONResponseFormat(raw json.RawMessage) string { - if !hasJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "response_format is unsupported" - } - if strings.EqualFold(strings.TrimSpace(value), "b64_json") || strings.TrimSpace(value) == "" { - return "" - } - if strings.EqualFold(strings.TrimSpace(value), "url") { - return "response_format=url is unsupported for this image model" - } - return "response_format is unsupported" -} - -func unsupportedJSONOutputFormat(raw json.RawMessage) string { - if !hasJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "output_format is unsupported" - } - if outputFormatAllowed(value) { - return "" - } - return "output_format is unsupported" -} - -func unsupportedJSONImageSize(raw json.RawMessage) string { - if !hasJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "image_size is unsupported" - } - if _, ok := normalizeImageSize(value); ok { - return "" - } - return "image_size is unsupported" -} - -func unsupportedJSONExactSize(raw json.RawMessage) string { - if !hasJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "size is unsupported" - } - if _, ok := imageSizeFromExactSize(value); ok { - return "" - } - return "size is unsupported" -} - -func unsupportedFormImageSize(raw string) string { - if strings.TrimSpace(raw) == "" { - return "" - } - if _, ok := normalizeImageSize(raw); ok { - return "" - } - return "image_size is unsupported" -} - -func unsupportedFormExactSize(raw string) string { - if strings.TrimSpace(raw) == "" { - return "" - } - if _, ok := imageSizeFromExactSize(raw); ok { - return "" - } - return "size is unsupported" -} - -func supportedSosanaModel(modelID string) bool { - model := strings.ToLower(strings.TrimSpace(modelID)) - if model == "" || model == "google/gemini-3.1-flash-image-preview" { - return true - } - if model == "banana-2-{image_size}-compliant" { - return true - } - switch model { - case "banana-2-1k-compliant", "banana-2-2k-compliant", "banana-2-4k-compliant": - return true - default: - return false - } -} - -func validateOutputFormat(raw string) error { - if outputFormatAllowed(raw) { - return nil - } - return fmt.Errorf("output_format is unsupported") -} - -func outputFormatAllowed(raw string) bool { - value := strings.ToLower(strings.TrimSpace(raw)) - return value == "" || value == "png" -} - -func hasJSONValue(raw json.RawMessage) bool { - raw = bytes.TrimSpace(raw) - return len(raw) > 0 && !bytes.Equal(raw, []byte("null")) -} diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go deleted file mode 100644 index c3a9a3cf..00000000 --- a/internal/converter/sosana/images.go +++ /dev/null @@ -1,446 +0,0 @@ -package sosana - -import ( - "bytes" - "encoding/base64" - "encoding/json" - "fmt" - "io" - "mime" - "mime/multipart" - "net/http" - "strconv" - "strings" - "time" - - "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" - "github.com/mixaill76/auto_ai_router/internal/converter/openai" -) - -const maxMultipartImageBytes = 20 * 1024 * 1024 - -const ( - StatusProcessing = "PROCESSING" - StatusCompleted = "COMPLETED" - StatusFailed = "FAILED" - StatusModerated = "MODERATED" -) - -type BananaCreateRequest struct { - Prompt string `json:"prompt"` - ImageURLs []string `json:"image_urls,omitempty"` - Model string `json:"model,omitempty"` - AspectRatio string `json:"aspect_ratio,omitempty"` - ImageSize string `json:"image_size,omitempty"` - PromptOptimization bool `json:"prompt_optimization"` -} - -type BananaTaskResponse struct { - UID string `json:"uid"` - Status string `json:"status"` - Prompt string `json:"prompt"` - CreatedAt string `json:"created_at"` - OptimizedPrompt string `json:"optimized_prompt"` - ResultFileURL *string `json:"result_file_url"` - Error *string `json:"error"` -} - -type openAIImageRequest struct { - openai.OpenAIImageRequest - AspectRatio string `json:"aspect_ratio,omitempty"` - Ratio string `json:"ratio,omitempty"` - ImageSize string `json:"image_size,omitempty"` -} - -func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, string, error) { - if reason := UnsupportedRequest("/v1/images/generations", openAIBody, "application/json"); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - - var req openAIImageRequest - if err := json.Unmarshal(openAIBody, &req); err != nil { - return nil, "", fmt.Errorf("failed to parse OpenAI image request: %w", err) - } - if err := validateImageCount(req.N); err != nil { - return nil, "", err - } - if err := validateResponseFormat(req.ResponseFormat); err != nil { - return nil, "", err - } - if err := validateOutputFormat(req.OutputFormat); err != nil { - return nil, "", err - } - prompt := strings.TrimSpace(req.Prompt) - if prompt == "" { - return nil, "", fmt.Errorf("image generation request missing prompt") - } - imageSize, err := imageSize(req.ImageSize, req.Size) - if err != nil { - return nil, "", err - } - concreteModel := providerModel(modelID, req.Model, imageSize) - if reason := UnsupportedModel(concreteModel); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - body, err := json.Marshal(BananaCreateRequest{ - Prompt: prompt, - Model: concreteModel, - AspectRatio: aspectRatio(req.AspectRatio, req.Ratio, req.Size), - ImageSize: imageSize, - PromptOptimization: false, - }) - if err != nil { - return nil, "", err - } - return body, concreteModel, nil -} - -func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, string, error) { - if reason := UnsupportedRequest("/v1/images/edits", openAIBody, contentType); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - - mediaType, params, err := mime.ParseMediaType(contentType) - if err != nil { - return nil, "", fmt.Errorf("failed to parse image edit content type: %w", err) - } - if !strings.HasPrefix(mediaType, "multipart/form-data") { - return nil, "", fmt.Errorf("image edits require multipart/form-data content type") - } - boundary := params["boundary"] - if boundary == "" { - return nil, "", fmt.Errorf("missing multipart boundary in content type") - } - - fields := make(map[string]string) - imageURLs := make([]string, 0, 1) - reader := multipart.NewReader(bytes.NewReader(openAIBody), boundary) - for { - part, err := reader.NextPart() - if err == io.EOF { - break - } - if err != nil { - return nil, "", fmt.Errorf("failed to read multipart image edit payload: %w", err) - } - - formName := part.FormName() - if formName == "" { - continue - } - data, err := readLimited(part, maxMultipartImageBytes) - if err != nil { - return nil, "", err - } - if part.FileName() == "" { - fields[formName] = strings.TrimSpace(string(data)) - continue - } - if formName == "mask" { - return nil, "", fmt.Errorf("image edits do not support mask") - } - if formName != "image" && formName != "images" && formName != "image[]" { - continue - } - mimeType := detectImageMIMEType(part.Header.Get("Content-Type"), data) - if mimeType != "image/png" { - return nil, "", fmt.Errorf("image edits support PNG images only") - } - imageURLs = append(imageURLs, "data:"+mimeType+";base64,"+base64.StdEncoding.EncodeToString(data)) - } - - if len(imageURLs) > maxInputImages { - return nil, "", fmt.Errorf("image edits support up to %d input images", maxInputImages) - } - if err := validateImageCountString(fields["n"]); err != nil { - return nil, "", err - } - if err := validateResponseFormat(fields["response_format"]); err != nil { - return nil, "", err - } - if err := validateOutputFormat(fields["output_format"]); err != nil { - return nil, "", err - } - prompt := strings.TrimSpace(fields["prompt"]) - if prompt == "" { - return nil, "", fmt.Errorf("image edit request missing prompt field") - } - if len(imageURLs) == 0 { - return nil, "", fmt.Errorf("image edit request missing image") - } - imageSize, err := imageSize(fields["image_size"], fields["size"]) - if err != nil { - return nil, "", err - } - concreteModel := providerModel(modelID, fields["model"], imageSize) - if reason := UnsupportedModel(concreteModel); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - body, err := json.Marshal(BananaCreateRequest{ - Prompt: prompt, - ImageURLs: imageURLs, - Model: concreteModel, - AspectRatio: aspectRatio(fields["aspect_ratio"], fields["ratio"], fields["size"]), - ImageSize: imageSize, - PromptOptimization: false, - }) - if err != nil { - return nil, "", err - } - return body, concreteModel, nil -} - -func OpenAIImageResponse(task BananaTaskResponse, image []byte) ([]byte, error) { - if len(image) == 0 { - return nil, fmt.Errorf("image task completed without image bytes") - } - resp := openai.OpenAIImageResponse{ - Created: createdAtUnix(task.CreatedAt), - Data: []openai.OpenAIImageData{ - { - B64JSON: base64.StdEncoding.EncodeToString(image), - RevisedPrompt: strings.TrimSpace(task.OptimizedPrompt), - }, - }, - } - return json.Marshal(resp) -} - -func providerModel(modelID, requestModel, imageSize string) string { - model := strings.TrimSpace(modelID) - if model == "" { - model = strings.TrimSpace(requestModel) - } - if model == "" { - return "" - } - if strings.EqualFold(model, "google/gemini-3.1-flash-image-preview") { - return "banana-2-" + strings.ToLower(imageSize) + "-compliant" - } - return strings.ReplaceAll(model, "{image_size}", strings.ToLower(imageSize)) -} - -func createdAtUnix(value string) int64 { - if value == "" { - return converterutil.GetCurrentTimestamp() - } - if ts, err := time.Parse(time.RFC3339Nano, value); err == nil { - return ts.Unix() - } - return converterutil.GetCurrentTimestamp() -} - -func CreateURL(baseURL string) string { - return strings.TrimSuffix(baseURL, "/") + "/api/banana/create-async" -} - -func PollURL(baseURL, uid string) string { - return strings.TrimSuffix(baseURL, "/") + "/api/banana/" + uid -} - -func SizeToAspectRatio(size string) string { - if spec, ok := imageSpecFromExactSize(size); ok { - return spec.aspectRatio - } - return "auto" -} - -func aspectRatio(explicit, ratio, size string) string { - if value := strings.TrimSpace(explicit); value != "" { - return value - } - if value := strings.TrimSpace(ratio); value != "" { - return value - } - return SizeToAspectRatio(size) -} - -func imageSize(explicit, size string) (string, error) { - if strings.TrimSpace(explicit) != "" { - if value, ok := normalizeImageSize(explicit); ok { - return value, nil - } - return "", fmt.Errorf("image_size is unsupported") - } - if value, ok := imageSizeFromExactSize(size); ok { - return value, nil - } - return "", fmt.Errorf("size is unsupported") -} - -func normalizeImageSize(raw string) (string, bool) { - switch strings.ToLower(strings.TrimSpace(raw)) { - case "", "auto", "1k": - return "1K", true - case "2k": - return "2K", true - case "4k": - return "4K", true - default: - return "", false - } -} - -func imageSizeFromExactSize(size string) (string, bool) { - if spec, ok := imageSpecFromExactSize(size); ok { - return spec.imageSize, true - } - return "", false -} - -type exactImageSpec struct { - imageSize string - aspectRatio string -} - -func imageSpecFromExactSize(size string) (exactImageSpec, bool) { - switch strings.TrimSpace(size) { - case "", "auto": - return exactImageSpec{imageSize: "1K", aspectRatio: "auto"}, true - - case "1024x1024": - return exactImageSpec{imageSize: "1K", aspectRatio: "1:1"}, true - case "512x2048": - return exactImageSpec{imageSize: "1K", aspectRatio: "1:4"}, true - case "384x3072": - return exactImageSpec{imageSize: "1K", aspectRatio: "1:8"}, true - case "848x1264": - return exactImageSpec{imageSize: "1K", aspectRatio: "2:3"}, true - case "1264x848": - return exactImageSpec{imageSize: "1K", aspectRatio: "3:2"}, true - case "896x1200": - return exactImageSpec{imageSize: "1K", aspectRatio: "3:4"}, true - case "2048x512": - return exactImageSpec{imageSize: "1K", aspectRatio: "4:1"}, true - case "1200x896": - return exactImageSpec{imageSize: "1K", aspectRatio: "4:3"}, true - case "928x1152": - return exactImageSpec{imageSize: "1K", aspectRatio: "4:5"}, true - case "1152x928": - return exactImageSpec{imageSize: "1K", aspectRatio: "5:4"}, true - case "3072x384": - return exactImageSpec{imageSize: "1K", aspectRatio: "8:1"}, true - case "768x1376": - return exactImageSpec{imageSize: "1K", aspectRatio: "9:16"}, true - case "1376x768": - return exactImageSpec{imageSize: "1K", aspectRatio: "16:9"}, true - case "1584x672": - return exactImageSpec{imageSize: "1K", aspectRatio: "21:9"}, true - - case "2048x2048": - return exactImageSpec{imageSize: "2K", aspectRatio: "1:1"}, true - case "1024x4096": - return exactImageSpec{imageSize: "2K", aspectRatio: "1:4"}, true - case "768x6144": - return exactImageSpec{imageSize: "2K", aspectRatio: "1:8"}, true - case "1696x2528": - return exactImageSpec{imageSize: "2K", aspectRatio: "2:3"}, true - case "2528x1696": - return exactImageSpec{imageSize: "2K", aspectRatio: "3:2"}, true - case "1792x2400": - return exactImageSpec{imageSize: "2K", aspectRatio: "3:4"}, true - case "4096x1024": - return exactImageSpec{imageSize: "2K", aspectRatio: "4:1"}, true - case "2400x1792": - return exactImageSpec{imageSize: "2K", aspectRatio: "4:3"}, true - case "1856x2304": - return exactImageSpec{imageSize: "2K", aspectRatio: "4:5"}, true - case "2304x1856": - return exactImageSpec{imageSize: "2K", aspectRatio: "5:4"}, true - case "6144x768": - return exactImageSpec{imageSize: "2K", aspectRatio: "8:1"}, true - case "1536x2752": - return exactImageSpec{imageSize: "2K", aspectRatio: "9:16"}, true - case "2752x1536": - return exactImageSpec{imageSize: "2K", aspectRatio: "16:9"}, true - case "3168x1344": - return exactImageSpec{imageSize: "2K", aspectRatio: "21:9"}, true - - case "4096x4096": - return exactImageSpec{imageSize: "4K", aspectRatio: "1:1"}, true - case "2048x8192": - return exactImageSpec{imageSize: "4K", aspectRatio: "1:4"}, true - case "1536x12288": - return exactImageSpec{imageSize: "4K", aspectRatio: "1:8"}, true - case "3392x5056": - return exactImageSpec{imageSize: "4K", aspectRatio: "2:3"}, true - case "5056x3392": - return exactImageSpec{imageSize: "4K", aspectRatio: "3:2"}, true - case "3584x4800": - return exactImageSpec{imageSize: "4K", aspectRatio: "3:4"}, true - case "8192x2048": - return exactImageSpec{imageSize: "4K", aspectRatio: "4:1"}, true - case "4800x3584": - return exactImageSpec{imageSize: "4K", aspectRatio: "4:3"}, true - case "3712x4608": - return exactImageSpec{imageSize: "4K", aspectRatio: "4:5"}, true - case "4608x3712": - return exactImageSpec{imageSize: "4K", aspectRatio: "5:4"}, true - case "12288x1536": - return exactImageSpec{imageSize: "4K", aspectRatio: "8:1"}, true - case "3072x5504": - return exactImageSpec{imageSize: "4K", aspectRatio: "9:16"}, true - case "5504x3072": - return exactImageSpec{imageSize: "4K", aspectRatio: "16:9"}, true - case "6336x2688": - return exactImageSpec{imageSize: "4K", aspectRatio: "21:9"}, true - default: - return exactImageSpec{}, false - } -} - -func validateImageCount(n *int) error { - if n == nil || *n == 1 { - return nil - } - return fmt.Errorf("image requests support n=1 only") -} - -func validateResponseFormat(format string) error { - if strings.EqualFold(strings.TrimSpace(format), "url") { - return fmt.Errorf("response_format=url is unsupported for this image model") - } - return nil -} - -func validateImageCountString(raw string) error { - raw = strings.TrimSpace(raw) - if raw == "" { - return nil - } - n, err := strconv.Atoi(raw) - if err != nil { - return fmt.Errorf("invalid image count: %w", err) - } - if n != 1 { - return fmt.Errorf("image requests support n=1 only") - } - return nil -} - -func readLimited(r io.Reader, limit int64) ([]byte, error) { - data, err := io.ReadAll(io.LimitReader(r, limit+1)) - if err != nil { - return nil, fmt.Errorf("failed to read multipart part: %w", err) - } - if int64(len(data)) > limit { - return nil, fmt.Errorf("multipart image exceeds %d bytes", limit) - } - return data, nil -} - -func detectImageMIMEType(header string, data []byte) string { - header = strings.ToLower(strings.TrimSpace(header)) - if strings.HasPrefix(header, "image/") { - mediaType, _, err := mime.ParseMediaType(header) - if err == nil { - return mediaType - } - return strings.TrimSpace(strings.Split(header, ";")[0]) - } - detected := http.DetectContentType(data) - if strings.HasPrefix(detected, "image/") { - return detected - } - return "application/octet-stream" -} diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go deleted file mode 100644 index b817ccfe..00000000 --- a/internal/converter/sosana/images_test.go +++ /dev/null @@ -1,293 +0,0 @@ -package sosana - -import ( - "bytes" - "encoding/base64" - "encoding/json" - "fmt" - "mime/multipart" - "strings" - "testing" - "time" - - "github.com/mixaill76/auto_ai_router/internal/converter/openai" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestImageGenerationRequest(t *testing.T) { - tests := []struct { - name string - size string - wantAspect string - wantSize string - wantErrPart string - }{ - {name: "one k square", size: "1024x1024", wantAspect: "1:1", wantSize: "1K"}, - {name: "one k wide", size: "1376x768", wantAspect: "16:9", wantSize: "1K"}, - {name: "one k portrait", size: "768x1376", wantAspect: "9:16", wantSize: "1K"}, - {name: "one k tall", size: "512x2048", wantAspect: "1:4", wantSize: "1K"}, - {name: "two k square", size: "2048x2048", wantAspect: "1:1", wantSize: "2K"}, - {name: "two k wide", size: "2752x1536", wantAspect: "16:9", wantSize: "2K"}, - {name: "four k square", size: "4096x4096", wantAspect: "1:1", wantSize: "4K"}, - {name: "four k ultra wide", size: "6336x2688", wantAspect: "21:9", wantSize: "4K"}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - body := []byte(`{"model":"banana-2-1k-compliant","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) - got, concreteModel, err := ImageGenerationRequest(body, "banana-2-{image_size}-compliant") - require.NoError(t, err) - - var req BananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "draw a cat", req.Prompt) - assert.Equal(t, "banana-2-"+strings.ToLower(tt.wantSize)+"-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) - assert.Equal(t, tt.wantAspect, req.AspectRatio) - assert.Equal(t, tt.wantSize, req.ImageSize) - assert.False(t, req.PromptOptimization) - assert.Empty(t, req.ImageURLs) - }) - } -} - -func TestImageGenerationRequestPrefersProviderModel(t *testing.T) { - got, concreteModel, err := ImageGenerationRequest([]byte(`{"model":"public-image","prompt":"draw","n":1}`), "banana-2-1k-compliant") - require.NoError(t, err) - - var req BananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "banana-2-1k-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) -} - -func TestImageGenerationRequestMapsPublicGeminiModelToSosanaTier(t *testing.T) { - got, concreteModel, err := ImageGenerationRequest( - []byte(`{"model":"google/gemini-3.1-flash-image-preview","prompt":"draw","image_size":"2K","n":1}`), - "google/gemini-3.1-flash-image-preview", - ) - require.NoError(t, err) - - var req BananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "banana-2-2k-compliant", req.Model) - assert.Equal(t, "banana-2-2k-compliant", concreteModel) - assert.Equal(t, "2K", req.ImageSize) -} - -func TestImageGenerationRequestUsesExplicitAspectRatio(t *testing.T) { - got, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","size":"1024x1024","aspect_ratio":"16:9"}`), "banana-2-1k-compliant") - require.NoError(t, err) - - var req BananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "16:9", req.AspectRatio) -} - -func TestImageGenerationRequestRejectsMultipleImages(t *testing.T) { - _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","n":2}`), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "n=1") -} - -func TestImageGenerationRequestRejectsURLResponseFormat(t *testing.T) { - _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","response_format":"url"}`), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "response_format=url") -} - -func TestImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { - tests := []struct { - name string - body string - want string - }{ - {name: "tools", body: `{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`, want: "tools"}, - {name: "thinking", body: `{"model":"banana-2-1k-compliant","prompt":"draw","thinking_level":"high"}`, want: "thinking_level"}, - {name: "output format", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"jpeg"}`, want: "output_format"}, - {name: "output compression", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_compression":0}`, want: "output_compression"}, - {name: "quality auto", body: `{"model":"banana-2-1k-compliant","prompt":"draw","quality":"auto"}`, want: "quality"}, - {name: "messages", body: `{"model":"banana-2-1k-compliant","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, - {name: "image size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_size":"0.5K"}`, want: "image_size"}, - {name: "exact size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"512x512"}`, want: "size"}, - {name: "legacy openai size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"1792x1024"}`, want: "size"}, - {name: "unknown size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"333x777"}`, want: "size"}, - {name: "reference images", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, _, err := ImageGenerationRequest([]byte(tt.body), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), tt.want) - }) - } -} - -func TestImageGenerationRequestAllowsPNGOutputFormat(t *testing.T) { - _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"png","response_format":"b64_json"}`), "banana-2-1k-compliant") - require.NoError(t, err) -} - -func TestImageEditRequest(t *testing.T) { - body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - "size": "1024x1024", - "n": "1", - }, map[string][]byte{ - "image": pngBytes(), - }) - - got, concreteModel, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.NoError(t, err) - - var req BananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "make it blue", req.Prompt) - assert.Equal(t, "banana-2-1k-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) - assert.Equal(t, "1:1", req.AspectRatio) - assert.Equal(t, "1K", req.ImageSize) - assert.False(t, req.PromptOptimization) - require.Len(t, req.ImageURLs, 1) - assert.True(t, strings.HasPrefix(req.ImageURLs[0], "data:image/png;base64,")) - assert.Contains(t, req.ImageURLs[0], base64.StdEncoding.EncodeToString(pngBytes())) -} - -func TestImageEditRequestPrefersProviderModel(t *testing.T) { - body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "public-image", - "prompt": "make it blue", - }, map[string][]byte{ - "image": pngBytes(), - }) - - got, concreteModel, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.NoError(t, err) - - var req BananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "banana-2-1k-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) -} - -func TestImageEditRequestRejectsMask(t *testing.T) { - body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - }, map[string][]byte{ - "image": pngBytes(), - "mask": pngBytes(), - }) - - _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "mask") -} - -func TestImageEditRequestRejectsMultipleImagesCount(t *testing.T) { - body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - "n": "2", - }, map[string][]byte{ - "image": pngBytes(), - }) - - _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "n=1") -} - -func TestImageEditRequestRejectsURLResponseFormat(t *testing.T) { - body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - "response_format": "url", - }, map[string][]byte{ - "image": pngBytes(), - }) - - _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "response_format=url") -} - -func TestImageEditRequestRejectsJPEGInput(t *testing.T) { - body, contentType := multipartImageEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - }, map[string][]byte{ - "image": jpegBytes(), - }) - - _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "PNG") -} - -func TestImageEditRequestRejectsTooManyImages(t *testing.T) { - var buf bytes.Buffer - writer := multipart.NewWriter(&buf) - require.NoError(t, writer.WriteField("model", "banana-2-1k-compliant")) - require.NoError(t, writer.WriteField("prompt", "make it blue")) - for i := 0; i < maxInputImages+1; i++ { - part, err := writer.CreateFormFile("image", fmt.Sprintf("image-%02d.png", i)) - require.NoError(t, err) - _, err = part.Write(pngBytes()) - require.NoError(t, err) - } - require.NoError(t, writer.Close()) - - _, _, err := ImageEditRequest(buf.Bytes(), writer.FormDataContentType(), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "too many") -} - -func TestOpenAIImageResponse(t *testing.T) { - createdAt := "2026-01-01T00:00:00Z" - body, err := OpenAIImageResponse(BananaTaskResponse{ - Status: StatusCompleted, - CreatedAt: createdAt, - OptimizedPrompt: "A detailed result prompt", - }, pngBytes()) - require.NoError(t, err) - - var resp openai.OpenAIImageResponse - require.NoError(t, json.Unmarshal(body, &resp)) - require.Len(t, resp.Data, 1) - assert.Empty(t, resp.Data[0].URL) - assert.Equal(t, "A detailed result prompt", resp.Data[0].RevisedPrompt) - assert.Equal(t, base64.StdEncoding.EncodeToString(pngBytes()), resp.Data[0].B64JSON) - ts, err := time.Parse(time.RFC3339, createdAt) - require.NoError(t, err) - assert.Equal(t, ts.Unix(), resp.Created) -} - -func multipartImageEditBody(t *testing.T, fields map[string]string, files map[string][]byte) ([]byte, string) { - t.Helper() - - var buf bytes.Buffer - writer := multipart.NewWriter(&buf) - for key, value := range fields { - require.NoError(t, writer.WriteField(key, value)) - } - for key, data := range files { - part, err := writer.CreateFormFile(key, key+".png") - require.NoError(t, err) - _, err = part.Write(data) - require.NoError(t, err) - } - require.NoError(t, writer.Close()) - return buf.Bytes(), writer.FormDataContentType() -} - -func pngBytes() []byte { - return []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} -} - -func jpegBytes() []byte { - return []byte{0xff, 0xd8, 0xff, 0xdb, 0, 0x43, 0, 1, 2, 3} -} diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index 528fdcf9..fc57183e 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -3,27 +3,48 @@ package proxy import ( "bytes" "context" + "encoding/base64" "encoding/json" "errors" + "fmt" "io" "math/rand" + "mime" + "mime/multipart" "net" "net/http" "net/url" + "strconv" "strings" "time" "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter" - "github.com/mixaill76/auto_ai_router/internal/converter/sosana" + "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" + "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/mixaill76/auto_ai_router/internal/scope" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" ) const ( - sosanaPollInterval = 2 * time.Second - maxSosanaResultImageBytes int64 = 32 * 1024 * 1024 - maxSosanaResultErrorBytes int64 = 16 * 1024 + sosanaPollInterval = 2 * time.Second + maxSosanaMultipartImageBytes = 20 * 1024 * 1024 + maxSosanaInputImages = 14 + maxSosanaResultImageBytes int64 = 32 * 1024 * 1024 + maxSosanaResultErrorBytes int64 = 16 * 1024 ) +const ( + sosanaStatusProcessing = "PROCESSING" + sosanaStatusCompleted = "COMPLETED" + sosanaStatusFailed = "FAILED" + sosanaStatusModerated = "MODERATED" +) + +const unsupportedImageProviderRequestMessage = "request parameters are not supported by available image providers" +const unsupportedProviderEndpointMessage = "request endpoint is not supported by available providers" + var allowPrivateSosanaResultURLForTests func(*url.URL) bool type sosanaAttemptResult struct { @@ -33,6 +54,175 @@ type sosanaAttemptResult struct { retryReason RetryReason } +type sosanaBananaCreateRequest struct { + Prompt string `json:"prompt"` + ImageURLs []string `json:"image_urls,omitempty"` + Model string `json:"model,omitempty"` + AspectRatio string `json:"aspect_ratio,omitempty"` + ImageSize string `json:"image_size,omitempty"` + PromptOptimization bool `json:"prompt_optimization"` +} + +type sosanaBananaTaskResponse struct { + UID string `json:"uid"` + Status string `json:"status"` + Prompt string `json:"prompt"` + CreatedAt string `json:"created_at"` + OptimizedPrompt string `json:"optimized_prompt"` + ResultFileURL *string `json:"result_file_url"` + Error *string `json:"error"` +} + +type sosanaOpenAIImageRequest struct { + openai.OpenAIImageRequest + AspectRatio string `json:"aspect_ratio,omitempty"` + Ratio string `json:"ratio,omitempty"` + ImageSize string `json:"image_size,omitempty"` +} + +func (p *Proxy) applySosanaCompatibilityRouting( + w http.ResponseWriter, + r *http.Request, + prepared *orchestratedRequest, + modelID string, + cred **config.CredentialConfig, + body *[]byte, + proxyBody *[]byte, + realModelID *string, + isImageGeneration bool, + isImageEdit bool, + logCtx *RequestLogContext, + start time.Time, +) bool { + if (*cred).Type != config.ProviderTypeSosana || (!isImageGeneration && !isImageEdit) { + return true + } + + reason := unsupportedSosanaModel(*realModelID) + if reason == "" { + reason = unsupportedSosanaRequest(r.URL.Path, *body, r.Header.Get("Content-Type")) + } + if reason == "" { + return true + } + + nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, logCtx.Scope, reason) + if routed { + *cred = nextCred + *body = nextReq.body + *proxyBody = nextReq.proxyBody + *realModelID = nextReq.realModelID + r.URL.Path = nextReq.path + prepared.body = nextReq.body + prepared.proxyBody = nextReq.proxyBody + prepared.proxyPath = nextReq.proxyPath + prepared.realModelID = nextReq.realModelID + prepared.convertedResp = nextReq.convertedResp + prepared.passthroughResponses = nextReq.passthroughResponses + prepared.nativeResponses = nextReq.nativeResponses + logCtx.RealModelID = *realModelID + if span := trace.SpanFromContext(r.Context()); span.IsRecording() { + span.SetAttributes( + attribute.String("aar.real_model", *realModelID), + attribute.String("aar.credential", nextCred.Name), + attribute.String("aar.provider", string(nextCred.Type)), + attribute.Bool("aar.provider_compatibility_skip", true), + ) + } + return true + } + + success, fallbackReason := p.TryFallbackProxy( + w, + requestWithPath(r, prepared.proxyPath), + modelID, + (*cred).Name, + http.StatusBadRequest, + RetryReasonServerErr, + *proxyBody, + start, + logCtx, + ) + if success { + return false + } + p.logger.DebugContext(r.Context(), "No fallback handled unsupported image provider request", + "credential", (*cred).Name, + "model", modelID, + "reason", reason, + "fallback_reason", fallbackReason) + logCtx.Credential = *cred + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = unsupportedImageProviderRequestMessage + WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) + return false +} + +func (p *Proxy) nextPrimaryAfterUnsupportedSosana( + r *http.Request, + prepared *orchestratedRequest, + modelID string, + currentCred *config.CredentialConfig, + visibility scope.Context, + reason string, +) (*config.CredentialConfig, credentialPreparedRequest, bool) { + triedCreds := GetTried(r.Context()) + triedCreds[currentCred.Name] = true + + for attempts := 0; attempts < 128; attempts++ { + candidate, err := p.balancer.NextForModelExcludingScoped(modelID, triedCreds, visibility) + if err != nil { + p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", + "model", modelID, + "credential", currentCred.Name, + "reason", reason, + "error", err) + return nil, credentialPreparedRequest{}, false + } + triedCreds[candidate.Name] = true + if candidate.Type == config.ProviderTypeSosana { + continue + } + + nextReq, prepErr := p.prepareRequestForCredential( + r, + prepared.baseBody, + prepared.baseProxyBody, + modelID, + prepared.baseRealModelID, + prepared.basePath, + prepared.streaming, + candidate, + prepared.isResponsesAPI, + prepared.responsesPrevHandled, + prepared.stickyCacheEligible, + ) + if prepErr != nil { + p.logger.WarnContext(r.Context(), "Failed to prepare alternate primary request after image compatibility skip", + "credential", candidate.Name, + "provider", string(candidate.Type), + "model", modelID, + "reason", reason, + "error", prepErr) + continue + } + + p.logger.InfoContext(r.Context(), "Skipping incompatible image credential for unsupported image request", + "credential", currentCred.Name, + "next_credential", candidate.Name, + "model", modelID, + "reason", reason) + return candidate, nextReq, true + } + + p.logger.WarnContext(r.Context(), "Image compatibility skip exhausted primary credential scan", + "credential", currentCred.Name, + "model", modelID, + "reason", reason) + return nil, credentialPreparedRequest{}, false +} + func (p *Proxy) handleSosanaRequest( w http.ResponseWriter, r *http.Request, @@ -166,9 +356,9 @@ func (p *Proxy) sosanaRealModelIDForCredential(modelID, fallbackRealModelID stri func (p *Proxy) buildSosanaCreateBody(body []byte, contentType, realModelID string, isImageEdit bool) ([]byte, string, error) { if isImageEdit { - return sosana.ImageEditRequest(body, contentType, realModelID) + return buildSosanaImageEditRequest(body, contentType, realModelID) } - return sosana.ImageGenerationRequest(body, realModelID) + return buildSosanaImageGenerationRequest(body, realModelID) } func (p *Proxy) createAndPollSosanaTask( @@ -178,9 +368,9 @@ func (p *Proxy) createAndPollSosanaTask( createBody []byte, logCtx *RequestLogContext, ) sosanaAttemptResult { - task, rawBody, statusCode, err := p.doSosanaTaskRequest(ctx, http.MethodPost, sosana.CreateURL(cred.BaseURL), cred, createBody) + task, rawBody, statusCode, err := p.doSosanaTaskRequest(ctx, http.MethodPost, sosanaCreateURL(cred.BaseURL), cred, createBody) if err != nil { - body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosana.CreateURL(cred.BaseURL), logCtx) + body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosanaCreateURL(cred.BaseURL), logCtx) return sosanaAttemptResult{ body: body, statusCode: code, @@ -189,7 +379,7 @@ func (p *Proxy) createAndPollSosanaTask( } if statusCode >= 400 { p.logUpstreamError(ctx, "Sosana create request completed with error status", statusCode, cred, modelID, rawBody, - "url", sosana.CreateURL(cred.BaseURL), + "url", sosanaCreateURL(cred.BaseURL), "request_id", logCtx.RequestID) retryable, reason := ShouldRetryWithFallback(statusCode, rawBody) return sosanaAttemptResult{ @@ -211,7 +401,7 @@ func (p *Proxy) createAndPollSosanaTask( select { case <-ctx.Done(): p.logUpstreamError(context.Background(), "Sosana task polling timed out", http.StatusRequestTimeout, cred, modelID, rawBody, - "url", sosana.PollURL(cred.BaseURL, task.UID), + "url", sosanaPollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID, "error", ctx.Err()) return sosanaAttemptResult{body: maskedUpstreamErrorBody(http.StatusRequestTimeout), statusCode: http.StatusRequestTimeout} @@ -220,14 +410,14 @@ func (p *Proxy) createAndPollSosanaTask( } immediatePoll = false - task, rawBody, statusCode, err = p.doSosanaTaskRequest(ctx, http.MethodGet, sosana.PollURL(cred.BaseURL, task.UID), cred, nil) + task, rawBody, statusCode, err = p.doSosanaTaskRequest(ctx, http.MethodGet, sosanaPollURL(cred.BaseURL, task.UID), cred, nil) if err != nil { - body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosana.PollURL(cred.BaseURL, task.UID), logCtx) + body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosanaPollURL(cred.BaseURL, task.UID), logCtx) return sosanaAttemptResult{body: body, statusCode: code} } if statusCode >= 400 { p.logUpstreamError(ctx, "Sosana poll request completed with error status", statusCode, cred, modelID, rawBody, - "url", sosana.PollURL(cred.BaseURL, task.UID), + "url", sosanaPollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID) return sosanaAttemptResult{body: maskedUpstreamErrorBody(statusCode), statusCode: statusCode} } @@ -238,26 +428,26 @@ func (p *Proxy) sosanaTaskBody( ctx context.Context, cred *config.CredentialConfig, modelID string, - task sosana.BananaTaskResponse, + task sosanaBananaTaskResponse, rawBody []byte, statusCode int, logCtx *RequestLogContext, ) ([]byte, int, bool) { switch task.Status { - case sosana.StatusCompleted: + case sosanaStatusCompleted: body, statusCode := p.sosanaCompletedImageBody(ctx, cred, modelID, task, rawBody, logCtx) return body, statusCode, true - case sosana.StatusFailed: + case sosanaStatusFailed: p.logUpstreamError(ctx, "Sosana task failed", http.StatusBadGateway, cred, modelID, rawBody, - "url", sosana.PollURL(cred.BaseURL, task.UID), + "url", sosanaPollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID) return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true - case sosana.StatusModerated: + case sosanaStatusModerated: p.logUpstreamError(ctx, "Sosana task moderated", http.StatusBadRequest, cred, modelID, rawBody, - "url", sosana.PollURL(cred.BaseURL, task.UID), + "url", sosanaPollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID) return maskedContentPolicyBody(), http.StatusBadRequest, true - case sosana.StatusProcessing: + case sosanaStatusProcessing: if task.UID == "" { p.logUpstreamError(ctx, "Sosana processing task missing uid", http.StatusBadGateway, cred, modelID, rawBody, "url", cred.BaseURL, @@ -267,7 +457,7 @@ func (p *Proxy) sosanaTaskBody( return nil, statusCode, false default: p.logUpstreamError(ctx, "Sosana task returned unknown status", http.StatusBadGateway, cred, modelID, rawBody, - "url", sosana.PollURL(cred.BaseURL, task.UID), + "url", sosanaPollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID, "status", task.Status) return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true @@ -278,7 +468,7 @@ func (p *Proxy) sosanaCompletedImageBody( ctx context.Context, cred *config.CredentialConfig, modelID string, - task sosana.BananaTaskResponse, + task sosanaBananaTaskResponse, rawBody []byte, logCtx *RequestLogContext, ) ([]byte, int) { @@ -286,7 +476,7 @@ func (p *Proxy) sosanaCompletedImageBody( if err != nil { return maskedUpstreamErrorBody(statusCode), statusCode } - body, err := sosana.OpenAIImageResponse(task, image) + body, err := buildSosanaOpenAIImageResponse(task, image) if err != nil { p.logUpstreamError(ctx, "Sosana completed task could not be converted", http.StatusBadGateway, cred, modelID, nil, "request_id", logCtx.RequestID, @@ -307,7 +497,7 @@ func (p *Proxy) downloadSosanaResultImage( ctx context.Context, cred *config.CredentialConfig, modelID string, - task sosana.BananaTaskResponse, + task sosanaBananaTaskResponse, rawTaskBody []byte, logCtx *RequestLogContext, ) ([]byte, string, int, error) { @@ -540,7 +730,7 @@ func isUnsafeSosanaResultIP(ip net.IP) bool { ip.IsUnspecified() } -func sosanaResultHost(task sosana.BananaTaskResponse) string { +func sosanaResultHost(task sosanaBananaTaskResponse) string { if task.ResultFileURL == nil { return "" } @@ -600,7 +790,7 @@ func isTextContentType(contentType string) bool { strings.Contains(contentType, "xml") } -func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosana.BananaTaskResponse, []byte, int, error) { +func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosanaBananaTaskResponse, []byte, int, error) { var reader *bytes.Reader if body != nil { reader = bytes.NewReader(body) @@ -609,7 +799,7 @@ func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cre } req, err := http.NewRequestWithContext(ctx, method, url, reader) if err != nil { - return sosana.BananaTaskResponse{}, nil, http.StatusInternalServerError, err + return sosanaBananaTaskResponse{}, nil, http.StatusInternalServerError, err } req.Header.Set("Authorization", "Bearer "+cred.APIKey) if body != nil { @@ -618,7 +808,7 @@ func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cre resp, err := p.client.Do(req) if err != nil { - return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err + return sosanaBananaTaskResponse{}, nil, http.StatusBadGateway, err } defer func() { if closeErr := resp.Body.Close(); closeErr != nil { @@ -628,9 +818,9 @@ func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cre rawBody, err := p.readLimitedResponseBody(resp.Body) if err != nil { - return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err + return sosanaBananaTaskResponse{}, nil, http.StatusBadGateway, err } - var task sosana.BananaTaskResponse + var task sosanaBananaTaskResponse if len(rawBody) > 0 { _ = json.Unmarshal(rawBody, &task) } @@ -648,3 +838,692 @@ func (p *Proxy) sosanaTransportError(ctx context.Context, err error, cred *confi "error", err) return maskedUpstreamErrorBody(statusCode), statusCode } + +func buildSosanaImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, string, error) { + if reason := unsupportedSosanaRequest("/v1/images/generations", openAIBody, "application/json"); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + + var req sosanaOpenAIImageRequest + if err := json.Unmarshal(openAIBody, &req); err != nil { + return nil, "", fmt.Errorf("failed to parse OpenAI image request: %w", err) + } + if err := validateSosanaImageCount(req.N); err != nil { + return nil, "", err + } + if err := validateSosanaResponseFormat(req.ResponseFormat); err != nil { + return nil, "", err + } + if err := validateSosanaOutputFormat(req.OutputFormat); err != nil { + return nil, "", err + } + prompt := strings.TrimSpace(req.Prompt) + if prompt == "" { + return nil, "", fmt.Errorf("image generation request missing prompt") + } + imageSize, err := sosanaImageSize(req.ImageSize, req.Size) + if err != nil { + return nil, "", err + } + concreteModel := sosanaProviderModel(modelID, req.Model, imageSize) + if reason := unsupportedSosanaModel(concreteModel); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + body, err := json.Marshal(sosanaBananaCreateRequest{ + Prompt: prompt, + Model: concreteModel, + AspectRatio: sosanaAspectRatio(req.AspectRatio, req.Ratio, req.Size), + ImageSize: imageSize, + PromptOptimization: false, + }) + if err != nil { + return nil, "", err + } + return body, concreteModel, nil +} + +func buildSosanaImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, string, error) { + if reason := unsupportedSosanaRequest("/v1/images/edits", openAIBody, contentType); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + + mediaType, params, err := mime.ParseMediaType(contentType) + if err != nil { + return nil, "", fmt.Errorf("failed to parse image edit content type: %w", err) + } + if !strings.HasPrefix(mediaType, "multipart/form-data") { + return nil, "", fmt.Errorf("image edits require multipart/form-data content type") + } + boundary := params["boundary"] + if boundary == "" { + return nil, "", fmt.Errorf("missing multipart boundary in content type") + } + + fields := make(map[string]string) + imageURLs := make([]string, 0, 1) + reader := multipart.NewReader(bytes.NewReader(openAIBody), boundary) + for { + part, err := reader.NextPart() + if err == io.EOF { + break + } + if err != nil { + return nil, "", fmt.Errorf("failed to read multipart image edit payload: %w", err) + } + + formName := part.FormName() + if formName == "" { + continue + } + data, err := readLimitedSosanaMultipartPart(part, maxSosanaMultipartImageBytes) + if err != nil { + return nil, "", err + } + if part.FileName() == "" { + fields[formName] = strings.TrimSpace(string(data)) + continue + } + if formName == "mask" { + return nil, "", fmt.Errorf("image edits do not support mask") + } + if formName != "image" && formName != "images" && formName != "image[]" { + continue + } + mimeType := detectSosanaImageMIMEType(part.Header.Get("Content-Type"), data) + if mimeType != "image/png" { + return nil, "", fmt.Errorf("image edits support PNG images only") + } + imageURLs = append(imageURLs, "data:"+mimeType+";base64,"+base64.StdEncoding.EncodeToString(data)) + } + + if len(imageURLs) > maxSosanaInputImages { + return nil, "", fmt.Errorf("image edits support up to %d input images", maxSosanaInputImages) + } + if err := validateSosanaImageCountString(fields["n"]); err != nil { + return nil, "", err + } + if err := validateSosanaResponseFormat(fields["response_format"]); err != nil { + return nil, "", err + } + if err := validateSosanaOutputFormat(fields["output_format"]); err != nil { + return nil, "", err + } + prompt := strings.TrimSpace(fields["prompt"]) + if prompt == "" { + return nil, "", fmt.Errorf("image edit request missing prompt field") + } + if len(imageURLs) == 0 { + return nil, "", fmt.Errorf("image edit request missing image") + } + imageSize, err := sosanaImageSize(fields["image_size"], fields["size"]) + if err != nil { + return nil, "", err + } + concreteModel := sosanaProviderModel(modelID, fields["model"], imageSize) + if reason := unsupportedSosanaModel(concreteModel); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + body, err := json.Marshal(sosanaBananaCreateRequest{ + Prompt: prompt, + ImageURLs: imageURLs, + Model: concreteModel, + AspectRatio: sosanaAspectRatio(fields["aspect_ratio"], fields["ratio"], fields["size"]), + ImageSize: imageSize, + PromptOptimization: false, + }) + if err != nil { + return nil, "", err + } + return body, concreteModel, nil +} + +func buildSosanaOpenAIImageResponse(task sosanaBananaTaskResponse, image []byte) ([]byte, error) { + if len(image) == 0 { + return nil, fmt.Errorf("image task completed without image bytes") + } + resp := openai.OpenAIImageResponse{ + Created: sosanaCreatedAtUnix(task.CreatedAt), + Data: []openai.OpenAIImageData{ + { + B64JSON: base64.StdEncoding.EncodeToString(image), + RevisedPrompt: strings.TrimSpace(task.OptimizedPrompt), + }, + }, + } + return json.Marshal(resp) +} + +var unsupportedSosanaImageFields = []string{ + "tools", + "tool_choice", + "google_search", + "thinking_level", + "thinking_budget", + "thinking_config", + "thinking", + "reasoning_effort", + "generation_config", + "temperature", + "top_p", + "top_k", + "seed", + "max_tokens", + "stop", + "stream", + "messages", + "extra_body", + "image", + "images", + "image_urls", + "reference_images", +} + +func unsupportedSosanaRequest(path string, body []byte, contentType string) string { + switch { + case strings.Contains(path, "/images/generations"): + return unsupportedSosanaGenerationRequest(body) + case strings.Contains(path, "/images/edits"): + return unsupportedSosanaEditRequest(body, contentType) + default: + return "endpoint is unsupported" + } +} + +func unsupportedSosanaModel(modelID string) string { + if supportedSosanaModel(modelID) { + return "" + } + return "model is unsupported" +} + +func unsupportedSosanaGenerationRequest(body []byte) string { + var raw map[string]json.RawMessage + if err := json.Unmarshal(body, &raw); err != nil { + return "" + } + return unsupportedSosanaImageFieldsInJSON(raw) +} + +func unsupportedSosanaEditRequest(body []byte, contentType string) string { + mediaType, params, err := mime.ParseMediaType(contentType) + if err != nil || !strings.HasPrefix(mediaType, "multipart/form-data") { + return "" + } + boundary := params["boundary"] + if boundary == "" { + return "" + } + + fields := make(map[string]string) + imageCount := 0 + reader := multipart.NewReader(bytes.NewReader(body), boundary) + for { + part, err := reader.NextPart() + if err != nil { + break + } + + formName := part.FormName() + if formName == "" { + continue + } + data, err := readLimitedSosanaMultipartPart(part, maxSosanaMultipartImageBytes) + if err != nil { + return err.Error() + } + if part.FileName() == "" { + fields[formName] = strings.TrimSpace(string(data)) + continue + } + if formName == "mask" { + return "mask is unsupported" + } + if formName != "image" && formName != "images" && formName != "image[]" { + continue + } + imageCount++ + if detectSosanaImageMIMEType(part.Header.Get("Content-Type"), data) != "image/png" { + return "only PNG input images are supported" + } + } + if imageCount > maxSosanaInputImages { + return "too many input images" + } + return unsupportedSosanaImageFieldsInForm(fields) +} + +func unsupportedSosanaImageFieldsInJSON(raw map[string]json.RawMessage) string { + if reason := unsupportedSosanaJSONImageCount(raw["n"]); reason != "" { + return reason + } + if reason := unsupportedSosanaJSONResponseFormat(raw["response_format"]); reason != "" { + return reason + } + if reason := unsupportedSosanaJSONOutputFormat(raw["output_format"]); reason != "" { + return reason + } + if reason := unsupportedSosanaJSONImageSize(raw["image_size"]); reason != "" { + return reason + } + if reason := unsupportedSosanaJSONExactSize(raw["size"]); reason != "" { + return reason + } + for _, field := range []string{"quality", "style", "background", "moderation"} { + if hasSosanaJSONValue(raw[field]) { + return field + " is unsupported" + } + } + if hasSosanaJSONValue(raw["output_compression"]) { + return "output_compression is unsupported" + } + for _, field := range unsupportedSosanaImageFields { + if hasSosanaJSONValue(raw[field]) { + return field + " is unsupported" + } + } + return "" +} + +func unsupportedSosanaImageFieldsInForm(fields map[string]string) string { + if err := validateSosanaImageCountString(fields["n"]); err != nil { + return err.Error() + } + if err := validateSosanaResponseFormat(fields["response_format"]); err != nil { + return err.Error() + } + if err := validateSosanaOutputFormat(fields["output_format"]); err != nil { + return err.Error() + } + if reason := unsupportedSosanaFormImageSize(fields["image_size"]); reason != "" { + return reason + } + if reason := unsupportedSosanaFormExactSize(fields["size"]); reason != "" { + return reason + } + for _, field := range []string{"quality", "style", "background", "moderation"} { + if strings.TrimSpace(fields[field]) != "" { + return field + " is unsupported" + } + } + if _, ok := fields["output_compression"]; ok { + return "output_compression is unsupported" + } + for _, field := range unsupportedSosanaImageFields { + if strings.TrimSpace(fields[field]) != "" { + return field + " is unsupported" + } + } + return "" +} + +func unsupportedSosanaJSONImageCount(raw json.RawMessage) string { + if !hasSosanaJSONValue(raw) { + return "" + } + var n int + if err := json.Unmarshal(raw, &n); err == nil { + if n == 1 { + return "" + } + return "image requests support n=1 only" + } + var f float64 + if err := json.Unmarshal(raw, &f); err == nil { + if f == 1 { + return "" + } + return "image requests support n=1 only" + } + return "invalid image count" +} + +func unsupportedSosanaJSONResponseFormat(raw json.RawMessage) string { + if !hasSosanaJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "response_format is unsupported" + } + if strings.EqualFold(strings.TrimSpace(value), "b64_json") || strings.TrimSpace(value) == "" { + return "" + } + if strings.EqualFold(strings.TrimSpace(value), "url") { + return "response_format=url is unsupported for this image model" + } + return "response_format is unsupported" +} + +func unsupportedSosanaJSONOutputFormat(raw json.RawMessage) string { + if !hasSosanaJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "output_format is unsupported" + } + if sosanaOutputFormatAllowed(value) { + return "" + } + return "output_format is unsupported" +} + +func unsupportedSosanaJSONImageSize(raw json.RawMessage) string { + if !hasSosanaJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "image_size is unsupported" + } + if _, ok := normalizeSosanaImageSize(value); ok { + return "" + } + return "image_size is unsupported" +} + +func unsupportedSosanaJSONExactSize(raw json.RawMessage) string { + if !hasSosanaJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "size is unsupported" + } + if _, ok := sosanaImageSizeFromExactSize(value); ok { + return "" + } + return "size is unsupported" +} + +func unsupportedSosanaFormImageSize(raw string) string { + if strings.TrimSpace(raw) == "" { + return "" + } + if _, ok := normalizeSosanaImageSize(raw); ok { + return "" + } + return "image_size is unsupported" +} + +func unsupportedSosanaFormExactSize(raw string) string { + if strings.TrimSpace(raw) == "" { + return "" + } + if _, ok := sosanaImageSizeFromExactSize(raw); ok { + return "" + } + return "size is unsupported" +} + +func supportedSosanaModel(modelID string) bool { + model := strings.ToLower(strings.TrimSpace(modelID)) + if model == "" || model == "google/gemini-3.1-flash-image-preview" { + return true + } + if model == "banana-2-{image_size}-compliant" { + return true + } + switch model { + case "banana-2-1k-compliant", "banana-2-2k-compliant", "banana-2-4k-compliant": + return true + default: + return false + } +} + +func validateSosanaOutputFormat(raw string) error { + if sosanaOutputFormatAllowed(raw) { + return nil + } + return fmt.Errorf("output_format is unsupported") +} + +func sosanaOutputFormatAllowed(raw string) bool { + value := strings.ToLower(strings.TrimSpace(raw)) + return value == "" || value == "png" +} + +func hasSosanaJSONValue(raw json.RawMessage) bool { + raw = bytes.TrimSpace(raw) + return len(raw) > 0 && !bytes.Equal(raw, []byte("null")) +} + +func sosanaProviderModel(modelID, requestModel, imageSize string) string { + model := strings.TrimSpace(modelID) + if model == "" { + model = strings.TrimSpace(requestModel) + } + if model == "" { + return "" + } + if strings.EqualFold(model, "google/gemini-3.1-flash-image-preview") { + return "banana-2-" + strings.ToLower(imageSize) + "-compliant" + } + return strings.ReplaceAll(model, "{image_size}", strings.ToLower(imageSize)) +} + +func sosanaCreatedAtUnix(value string) int64 { + if value == "" { + return converterutil.GetCurrentTimestamp() + } + if ts, err := time.Parse(time.RFC3339Nano, value); err == nil { + return ts.Unix() + } + return converterutil.GetCurrentTimestamp() +} + +func sosanaCreateURL(baseURL string) string { + return strings.TrimSuffix(baseURL, "/") + "/api/banana/create-async" +} + +func sosanaPollURL(baseURL, uid string) string { + return strings.TrimSuffix(baseURL, "/") + "/api/banana/" + uid +} + +func sosanaSizeToAspectRatio(size string) string { + if spec, ok := sosanaImageSpecFromExactSize(size); ok { + return spec.aspectRatio + } + return "auto" +} + +func sosanaAspectRatio(explicit, ratio, size string) string { + if value := strings.TrimSpace(explicit); value != "" { + return value + } + if value := strings.TrimSpace(ratio); value != "" { + return value + } + return sosanaSizeToAspectRatio(size) +} + +func sosanaImageSize(explicit, size string) (string, error) { + if strings.TrimSpace(explicit) != "" { + if value, ok := normalizeSosanaImageSize(explicit); ok { + return value, nil + } + return "", fmt.Errorf("image_size is unsupported") + } + if value, ok := sosanaImageSizeFromExactSize(size); ok { + return value, nil + } + return "", fmt.Errorf("size is unsupported") +} + +func normalizeSosanaImageSize(raw string) (string, bool) { + switch strings.ToLower(strings.TrimSpace(raw)) { + case "", "auto", "1k": + return "1K", true + case "2k": + return "2K", true + case "4k": + return "4K", true + default: + return "", false + } +} + +func sosanaImageSizeFromExactSize(size string) (string, bool) { + if spec, ok := sosanaImageSpecFromExactSize(size); ok { + return spec.imageSize, true + } + return "", false +} + +type sosanaExactImageSpec struct { + imageSize string + aspectRatio string +} + +func sosanaImageSpecFromExactSize(size string) (sosanaExactImageSpec, bool) { + switch strings.TrimSpace(size) { + case "", "auto": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "auto"}, true + + case "1024x1024": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "1:1"}, true + case "512x2048": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "1:4"}, true + case "384x3072": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "1:8"}, true + case "848x1264": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "2:3"}, true + case "1264x848": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "3:2"}, true + case "896x1200": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "3:4"}, true + case "2048x512": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "4:1"}, true + case "1200x896": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "4:3"}, true + case "928x1152": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "4:5"}, true + case "1152x928": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "5:4"}, true + case "3072x384": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "8:1"}, true + case "768x1376": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "9:16"}, true + case "1376x768": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "16:9"}, true + case "1584x672": + return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "21:9"}, true + + case "2048x2048": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "1:1"}, true + case "1024x4096": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "1:4"}, true + case "768x6144": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "1:8"}, true + case "1696x2528": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "2:3"}, true + case "2528x1696": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "3:2"}, true + case "1792x2400": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "3:4"}, true + case "4096x1024": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "4:1"}, true + case "2400x1792": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "4:3"}, true + case "1856x2304": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "4:5"}, true + case "2304x1856": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "5:4"}, true + case "6144x768": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "8:1"}, true + case "1536x2752": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "9:16"}, true + case "2752x1536": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "16:9"}, true + case "3168x1344": + return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "21:9"}, true + + case "4096x4096": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "1:1"}, true + case "2048x8192": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "1:4"}, true + case "1536x12288": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "1:8"}, true + case "3392x5056": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "2:3"}, true + case "5056x3392": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "3:2"}, true + case "3584x4800": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "3:4"}, true + case "8192x2048": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "4:1"}, true + case "4800x3584": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "4:3"}, true + case "3712x4608": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "4:5"}, true + case "4608x3712": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "5:4"}, true + case "12288x1536": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "8:1"}, true + case "3072x5504": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "9:16"}, true + case "5504x3072": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "16:9"}, true + case "6336x2688": + return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "21:9"}, true + default: + return sosanaExactImageSpec{}, false + } +} + +func validateSosanaImageCount(n *int) error { + if n == nil || *n == 1 { + return nil + } + return fmt.Errorf("image requests support n=1 only") +} + +func validateSosanaResponseFormat(format string) error { + if strings.EqualFold(strings.TrimSpace(format), "url") { + return fmt.Errorf("response_format=url is unsupported for this image model") + } + return nil +} + +func validateSosanaImageCountString(raw string) error { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil + } + n, err := strconv.Atoi(raw) + if err != nil { + return fmt.Errorf("invalid image count: %w", err) + } + if n != 1 { + return fmt.Errorf("image requests support n=1 only") + } + return nil +} + +func readLimitedSosanaMultipartPart(r io.Reader, limit int64) ([]byte, error) { + data, err := io.ReadAll(io.LimitReader(r, limit+1)) + if err != nil { + return nil, fmt.Errorf("failed to read multipart part: %w", err) + } + if int64(len(data)) > limit { + return nil, fmt.Errorf("multipart image exceeds %d bytes", limit) + } + return data, nil +} + +func detectSosanaImageMIMEType(header string, data []byte) string { + header = strings.ToLower(strings.TrimSpace(header)) + if strings.HasPrefix(header, "image/") { + mediaType, _, err := mime.ParseMediaType(header) + if err == nil { + return mediaType + } + return strings.TrimSpace(strings.Split(header, ";")[0]) + } + detected := http.DetectContentType(data) + if strings.HasPrefix(detected, "image/") { + return detected + } + return "application/octet-stream" +} diff --git a/internal/proxy/sosana_routing.go b/internal/proxy/sosana_routing.go deleted file mode 100644 index 69900690..00000000 --- a/internal/proxy/sosana_routing.go +++ /dev/null @@ -1,158 +0,0 @@ -package proxy - -import ( - "net/http" - "time" - - "github.com/mixaill76/auto_ai_router/internal/config" - "github.com/mixaill76/auto_ai_router/internal/converter/sosana" - "github.com/mixaill76/auto_ai_router/internal/scope" - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/trace" -) - -const unsupportedImageProviderRequestMessage = "request parameters are not supported by available image providers" -const unsupportedProviderEndpointMessage = "request endpoint is not supported by available providers" - -func (p *Proxy) applySosanaCompatibilityRouting( - w http.ResponseWriter, - r *http.Request, - prepared *orchestratedRequest, - modelID string, - cred **config.CredentialConfig, - body *[]byte, - proxyBody *[]byte, - realModelID *string, - isImageGeneration bool, - isImageEdit bool, - logCtx *RequestLogContext, - start time.Time, -) bool { - if (*cred).Type != config.ProviderTypeSosana || (!isImageGeneration && !isImageEdit) { - return true - } - - reason := sosana.UnsupportedModel(*realModelID) - if reason == "" { - reason = sosana.UnsupportedRequest(r.URL.Path, *body, r.Header.Get("Content-Type")) - } - if reason == "" { - return true - } - - nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, logCtx.Scope, reason) - if routed { - *cred = nextCred - *body = nextReq.body - *proxyBody = nextReq.proxyBody - *realModelID = nextReq.realModelID - r.URL.Path = nextReq.path - prepared.body = nextReq.body - prepared.proxyBody = nextReq.proxyBody - prepared.proxyPath = nextReq.proxyPath - prepared.realModelID = nextReq.realModelID - prepared.convertedResp = nextReq.convertedResp - prepared.passthroughResponses = nextReq.passthroughResponses - prepared.nativeResponses = nextReq.nativeResponses - logCtx.RealModelID = *realModelID - if span := trace.SpanFromContext(r.Context()); span.IsRecording() { - span.SetAttributes( - attribute.String("aar.real_model", *realModelID), - attribute.String("aar.credential", nextCred.Name), - attribute.String("aar.provider", string(nextCred.Type)), - attribute.Bool("aar.provider_compatibility_skip", true), - ) - } - return true - } - - success, fallbackReason := p.TryFallbackProxy( - w, - requestWithPath(r, prepared.proxyPath), - modelID, - (*cred).Name, - http.StatusBadRequest, - RetryReasonServerErr, - *proxyBody, - start, - logCtx, - ) - if success { - return false - } - p.logger.DebugContext(r.Context(), "No fallback handled unsupported image provider request", - "credential", (*cred).Name, - "model", modelID, - "reason", reason, - "fallback_reason", fallbackReason) - logCtx.Credential = *cred - logCtx.Status = "failure" - logCtx.HTTPStatus = http.StatusBadRequest - logCtx.ErrorMsg = unsupportedImageProviderRequestMessage - WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) - return false -} - -func (p *Proxy) nextPrimaryAfterUnsupportedSosana( - r *http.Request, - prepared *orchestratedRequest, - modelID string, - currentCred *config.CredentialConfig, - visibility scope.Context, - reason string, -) (*config.CredentialConfig, credentialPreparedRequest, bool) { - triedCreds := GetTried(r.Context()) - triedCreds[currentCred.Name] = true - - for attempts := 0; attempts < 128; attempts++ { - candidate, err := p.balancer.NextForModelExcludingScoped(modelID, triedCreds, visibility) - if err != nil { - p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", - "model", modelID, - "credential", currentCred.Name, - "reason", reason, - "error", err) - return nil, credentialPreparedRequest{}, false - } - triedCreds[candidate.Name] = true - if candidate.Type == config.ProviderTypeSosana { - continue - } - - nextReq, prepErr := p.prepareRequestForCredential( - r, - prepared.baseBody, - prepared.baseProxyBody, - modelID, - prepared.baseRealModelID, - prepared.basePath, - prepared.streaming, - candidate, - prepared.isResponsesAPI, - prepared.responsesPrevHandled, - prepared.stickyCacheEligible, - ) - if prepErr != nil { - p.logger.WarnContext(r.Context(), "Failed to prepare alternate primary request after image compatibility skip", - "credential", candidate.Name, - "provider", string(candidate.Type), - "model", modelID, - "reason", reason, - "error", prepErr) - continue - } - - p.logger.InfoContext(r.Context(), "Skipping incompatible image credential for unsupported image request", - "credential", currentCred.Name, - "next_credential", candidate.Name, - "model", modelID, - "reason", reason) - return candidate, nextReq, true - } - - p.logger.WarnContext(r.Context(), "Image compatibility skip exhausted primary credential scan", - "credential", currentCred.Name, - "model", modelID, - "reason", reason) - return nil, credentialPreparedRequest{}, false -} diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 8afc2d75..859c808b 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -19,7 +19,6 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter/openai" - "github.com/mixaill76/auto_ai_router/internal/converter/sosana" "github.com/mixaill76/auto_ai_router/internal/litellmdb" litellmmodels "github.com/mixaill76/auto_ai_router/internal/litellmdb/models" aimodels "github.com/mixaill76/auto_ai_router/internal/models" @@ -880,8 +879,8 @@ func TestDownloadSosanaResultImageRejectsUnsafeProductionURL(t *testing.T) { resultURL := imageServer.URL + "/private.png" prx := newSosanaTestProxy("https://sosana.art", &logBuf) cred := &config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: "https://sosana.art"} - image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "banana-2-1k-compliant", sosana.BananaTaskResponse{ - Status: sosana.StatusCompleted, + image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "banana-2-1k-compliant", sosanaBananaTaskResponse{ + Status: sosanaStatusCompleted, ResultFileURL: &resultURL, }, nil, &RequestLogContext{}) @@ -1091,6 +1090,256 @@ func TestProxyRequest_SosanaTimeoutMasked(t *testing.T) { } } +func TestSosanaImageGenerationRequest(t *testing.T) { + tests := []struct { + name string + size string + wantAspect string + wantSize string + }{ + {name: "one k square", size: "1024x1024", wantAspect: "1:1", wantSize: "1K"}, + {name: "one k wide", size: "1376x768", wantAspect: "16:9", wantSize: "1K"}, + {name: "one k portrait", size: "768x1376", wantAspect: "9:16", wantSize: "1K"}, + {name: "one k tall", size: "512x2048", wantAspect: "1:4", wantSize: "1K"}, + {name: "two k square", size: "2048x2048", wantAspect: "1:1", wantSize: "2K"}, + {name: "two k wide", size: "2752x1536", wantAspect: "16:9", wantSize: "2K"}, + {name: "four k square", size: "4096x4096", wantAspect: "1:1", wantSize: "4K"}, + {name: "four k ultra wide", size: "6336x2688", wantAspect: "21:9", wantSize: "4K"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := []byte(`{"model":"banana-2-1k-compliant","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) + got, concreteModel, err := buildSosanaImageGenerationRequest(body, "banana-2-{image_size}-compliant") + require.NoError(t, err) + + var req sosanaBananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "draw a cat", req.Prompt) + assert.Equal(t, "banana-2-"+strings.ToLower(tt.wantSize)+"-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) + assert.Equal(t, tt.wantAspect, req.AspectRatio) + assert.Equal(t, tt.wantSize, req.ImageSize) + assert.False(t, req.PromptOptimization) + assert.Empty(t, req.ImageURLs) + }) + } +} + +func TestSosanaImageGenerationRequestPrefersProviderModel(t *testing.T) { + got, concreteModel, err := buildSosanaImageGenerationRequest([]byte(`{"model":"public-image","prompt":"draw","n":1}`), "banana-2-1k-compliant") + require.NoError(t, err) + + var req sosanaBananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) +} + +func TestSosanaImageGenerationRequestMapsPublicGeminiModelToSosanaTier(t *testing.T) { + got, concreteModel, err := buildSosanaImageGenerationRequest( + []byte(`{"model":"google/gemini-3.1-flash-image-preview","prompt":"draw","image_size":"2K","n":1}`), + "google/gemini-3.1-flash-image-preview", + ) + require.NoError(t, err) + + var req sosanaBananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "banana-2-2k-compliant", req.Model) + assert.Equal(t, "banana-2-2k-compliant", concreteModel) + assert.Equal(t, "2K", req.ImageSize) +} + +func TestSosanaImageGenerationRequestUsesExplicitAspectRatio(t *testing.T) { + got, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","size":"1024x1024","aspect_ratio":"16:9"}`), "banana-2-1k-compliant") + require.NoError(t, err) + + var req sosanaBananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "16:9", req.AspectRatio) +} + +func TestSosanaImageGenerationRequestRejectsMultipleImages(t *testing.T) { + _, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","n":2}`), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "n=1") +} + +func TestSosanaImageGenerationRequestRejectsURLResponseFormat(t *testing.T) { + _, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","response_format":"url"}`), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "response_format=url") +} + +func TestSosanaImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { + tests := []struct { + name string + body string + want string + }{ + {name: "tools", body: `{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`, want: "tools"}, + {name: "thinking", body: `{"model":"banana-2-1k-compliant","prompt":"draw","thinking_level":"high"}`, want: "thinking_level"}, + {name: "output format", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"jpeg"}`, want: "output_format"}, + {name: "output compression", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_compression":0}`, want: "output_compression"}, + {name: "quality auto", body: `{"model":"banana-2-1k-compliant","prompt":"draw","quality":"auto"}`, want: "quality"}, + {name: "messages", body: `{"model":"banana-2-1k-compliant","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, + {name: "image size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_size":"0.5K"}`, want: "image_size"}, + {name: "exact size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"512x512"}`, want: "size"}, + {name: "legacy openai size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"1792x1024"}`, want: "size"}, + {name: "unknown size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"333x777"}`, want: "size"}, + {name: "reference images", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, _, err := buildSosanaImageGenerationRequest([]byte(tt.body), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), tt.want) + }) + } +} + +func TestSosanaImageGenerationRequestAllowsPNGOutputFormat(t *testing.T) { + _, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"png","response_format":"b64_json"}`), "banana-2-1k-compliant") + require.NoError(t, err) +} + +func TestSosanaImageEditRequest(t *testing.T) { + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + "size": "1024x1024", + "n": "1", + }, map[string][]byte{ + "image": sosanaResultPNG, + }) + + got, concreteModel, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.NoError(t, err) + + var req sosanaBananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "make it blue", req.Prompt) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) + assert.Equal(t, "1:1", req.AspectRatio) + assert.Equal(t, "1K", req.ImageSize) + assert.False(t, req.PromptOptimization) + require.Len(t, req.ImageURLs, 1) + assert.True(t, strings.HasPrefix(req.ImageURLs[0], "data:image/png;base64,")) + assert.Contains(t, req.ImageURLs[0], base64.StdEncoding.EncodeToString(sosanaResultPNG)) +} + +func TestSosanaImageEditRequestPrefersProviderModel(t *testing.T) { + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "public-image", + "prompt": "make it blue", + }, map[string][]byte{ + "image": sosanaResultPNG, + }) + + got, concreteModel, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.NoError(t, err) + + var req sosanaBananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) +} + +func TestSosanaImageEditRequestRejectsMask(t *testing.T) { + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + }, map[string][]byte{ + "image": sosanaResultPNG, + "mask": sosanaResultPNG, + }) + + _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "mask") +} + +func TestSosanaImageEditRequestRejectsMultipleImagesCount(t *testing.T) { + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + "n": "2", + }, map[string][]byte{ + "image": sosanaResultPNG, + }) + + _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "n=1") +} + +func TestSosanaImageEditRequestRejectsURLResponseFormat(t *testing.T) { + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + "response_format": "url", + }, map[string][]byte{ + "image": sosanaResultPNG, + }) + + _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "response_format=url") +} + +func TestSosanaImageEditRequestRejectsJPEGInput(t *testing.T) { + body, contentType := sosanaMultipartEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + }, map[string][]byte{ + "image": {0xff, 0xd8, 0xff, 0xdb, 0, 0x43, 0, 1, 2, 3}, + }) + + _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "PNG") +} + +func TestSosanaImageEditRequestRejectsTooManyImages(t *testing.T) { + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("model", "banana-2-1k-compliant")) + require.NoError(t, writer.WriteField("prompt", "make it blue")) + for i := 0; i < maxSosanaInputImages+1; i++ { + part, err := writer.CreateFormFile("image", fmt.Sprintf("image-%02d.png", i)) + require.NoError(t, err) + _, err = part.Write(sosanaResultPNG) + require.NoError(t, err) + } + require.NoError(t, writer.Close()) + + _, _, err := buildSosanaImageEditRequest(buf.Bytes(), writer.FormDataContentType(), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "too many") +} + +func TestSosanaOpenAIImageResponse(t *testing.T) { + createdAt := "2026-01-01T00:00:00Z" + body, err := buildSosanaOpenAIImageResponse(sosanaBananaTaskResponse{ + Status: sosanaStatusCompleted, + CreatedAt: createdAt, + OptimizedPrompt: "A detailed result prompt", + }, sosanaResultPNG) + require.NoError(t, err) + + var resp openai.OpenAIImageResponse + require.NoError(t, json.Unmarshal(body, &resp)) + require.Len(t, resp.Data, 1) + assert.Empty(t, resp.Data[0].URL) + assert.Equal(t, "A detailed result prompt", resp.Data[0].RevisedPrompt) + assert.Equal(t, base64.StdEncoding.EncodeToString(sosanaResultPNG), resp.Data[0].B64JSON) + ts, err := time.Parse(time.RFC3339, createdAt) + require.NoError(t, err) + assert.Equal(t, ts.Unix(), resp.Created) +} + var sosanaResultPNG = []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} func newSosanaResultImageServer(t *testing.T, status int, contentType string, body []byte, auths *[]string) *httptest.Server { From d52a62af7db6ac024464a55a55d8630b859ff19f Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Thu, 23 Jul 2026 20:12:11 +0300 Subject: [PATCH 14/16] Revert "refactor: keep sosana image conversion in proxy" This reverts commit afac9e9ca23730d87bdd9cad0ed3f9527a94e08d. --- internal/converter/sosana/compatibility.go | 310 +++++++ internal/converter/sosana/images.go | 446 ++++++++++ internal/converter/sosana/images_test.go | 293 +++++++ internal/proxy/sosana.go | 939 +-------------------- internal/proxy/sosana_routing.go | 158 ++++ internal/proxy/sosana_test.go | 255 +----- 6 files changed, 1240 insertions(+), 1161 deletions(-) create mode 100644 internal/converter/sosana/compatibility.go create mode 100644 internal/converter/sosana/images.go create mode 100644 internal/converter/sosana/images_test.go create mode 100644 internal/proxy/sosana_routing.go diff --git a/internal/converter/sosana/compatibility.go b/internal/converter/sosana/compatibility.go new file mode 100644 index 00000000..453c1adc --- /dev/null +++ b/internal/converter/sosana/compatibility.go @@ -0,0 +1,310 @@ +package sosana + +import ( + "bytes" + "encoding/json" + "fmt" + "mime" + "mime/multipart" + "strings" +) + +const maxInputImages = 14 + +var unsupportedImageFields = []string{ + "tools", + "tool_choice", + "google_search", + "thinking_level", + "thinking_budget", + "thinking_config", + "thinking", + "reasoning_effort", + "generation_config", + "temperature", + "top_p", + "top_k", + "seed", + "max_tokens", + "stop", + "stream", + "messages", + "extra_body", + "image", + "images", + "image_urls", + "reference_images", +} + +// UnsupportedRequest returns a short reason when a request needs image features +// that Sosana Banana does not expose in its public API. +func UnsupportedRequest(path string, body []byte, contentType string) string { + switch { + case strings.Contains(path, "/images/generations"): + return unsupportedGenerationRequest(body) + case strings.Contains(path, "/images/edits"): + return unsupportedEditRequest(body, contentType) + default: + return "endpoint is unsupported" + } +} + +func UnsupportedModel(modelID string) string { + if supportedSosanaModel(modelID) { + return "" + } + return "model is unsupported" +} + +func unsupportedGenerationRequest(body []byte) string { + var raw map[string]json.RawMessage + if err := json.Unmarshal(body, &raw); err != nil { + return "" + } + return unsupportedImageFieldsInJSON(raw) +} + +func unsupportedEditRequest(body []byte, contentType string) string { + mediaType, params, err := mime.ParseMediaType(contentType) + if err != nil || !strings.HasPrefix(mediaType, "multipart/form-data") { + return "" + } + boundary := params["boundary"] + if boundary == "" { + return "" + } + + fields := make(map[string]string) + imageCount := 0 + reader := multipart.NewReader(bytes.NewReader(body), boundary) + for { + part, err := reader.NextPart() + if err != nil { + break + } + + formName := part.FormName() + if formName == "" { + continue + } + data, err := readLimited(part, maxMultipartImageBytes) + if err != nil { + return err.Error() + } + if part.FileName() == "" { + fields[formName] = strings.TrimSpace(string(data)) + continue + } + if formName == "mask" { + return "mask is unsupported" + } + if formName != "image" && formName != "images" && formName != "image[]" { + continue + } + imageCount++ + if detectImageMIMEType(part.Header.Get("Content-Type"), data) != "image/png" { + return "only PNG input images are supported" + } + } + if imageCount > maxInputImages { + return "too many input images" + } + return unsupportedImageFieldsInForm(fields) +} + +func unsupportedImageFieldsInJSON(raw map[string]json.RawMessage) string { + if reason := unsupportedJSONImageCount(raw["n"]); reason != "" { + return reason + } + if reason := unsupportedJSONResponseFormat(raw["response_format"]); reason != "" { + return reason + } + if reason := unsupportedJSONOutputFormat(raw["output_format"]); reason != "" { + return reason + } + if reason := unsupportedJSONImageSize(raw["image_size"]); reason != "" { + return reason + } + if reason := unsupportedJSONExactSize(raw["size"]); reason != "" { + return reason + } + for _, field := range []string{"quality", "style", "background", "moderation"} { + if hasJSONValue(raw[field]) { + return field + " is unsupported" + } + } + if hasJSONValue(raw["output_compression"]) { + return "output_compression is unsupported" + } + for _, field := range unsupportedImageFields { + if hasJSONValue(raw[field]) { + return field + " is unsupported" + } + } + return "" +} + +func unsupportedImageFieldsInForm(fields map[string]string) string { + if err := validateImageCountString(fields["n"]); err != nil { + return err.Error() + } + if err := validateResponseFormat(fields["response_format"]); err != nil { + return err.Error() + } + if err := validateOutputFormat(fields["output_format"]); err != nil { + return err.Error() + } + if reason := unsupportedFormImageSize(fields["image_size"]); reason != "" { + return reason + } + if reason := unsupportedFormExactSize(fields["size"]); reason != "" { + return reason + } + for _, field := range []string{"quality", "style", "background", "moderation"} { + if strings.TrimSpace(fields[field]) != "" { + return field + " is unsupported" + } + } + if _, ok := fields["output_compression"]; ok { + return "output_compression is unsupported" + } + for _, field := range unsupportedImageFields { + if strings.TrimSpace(fields[field]) != "" { + return field + " is unsupported" + } + } + return "" +} + +func unsupportedJSONImageCount(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var n int + if err := json.Unmarshal(raw, &n); err == nil { + if n == 1 { + return "" + } + return "image requests support n=1 only" + } + var f float64 + if err := json.Unmarshal(raw, &f); err == nil { + if f == 1 { + return "" + } + return "image requests support n=1 only" + } + return "invalid image count" +} + +func unsupportedJSONResponseFormat(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "response_format is unsupported" + } + if strings.EqualFold(strings.TrimSpace(value), "b64_json") || strings.TrimSpace(value) == "" { + return "" + } + if strings.EqualFold(strings.TrimSpace(value), "url") { + return "response_format=url is unsupported for this image model" + } + return "response_format is unsupported" +} + +func unsupportedJSONOutputFormat(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "output_format is unsupported" + } + if outputFormatAllowed(value) { + return "" + } + return "output_format is unsupported" +} + +func unsupportedJSONImageSize(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "image_size is unsupported" + } + if _, ok := normalizeImageSize(value); ok { + return "" + } + return "image_size is unsupported" +} + +func unsupportedJSONExactSize(raw json.RawMessage) string { + if !hasJSONValue(raw) { + return "" + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "size is unsupported" + } + if _, ok := imageSizeFromExactSize(value); ok { + return "" + } + return "size is unsupported" +} + +func unsupportedFormImageSize(raw string) string { + if strings.TrimSpace(raw) == "" { + return "" + } + if _, ok := normalizeImageSize(raw); ok { + return "" + } + return "image_size is unsupported" +} + +func unsupportedFormExactSize(raw string) string { + if strings.TrimSpace(raw) == "" { + return "" + } + if _, ok := imageSizeFromExactSize(raw); ok { + return "" + } + return "size is unsupported" +} + +func supportedSosanaModel(modelID string) bool { + model := strings.ToLower(strings.TrimSpace(modelID)) + if model == "" || model == "google/gemini-3.1-flash-image-preview" { + return true + } + if model == "banana-2-{image_size}-compliant" { + return true + } + switch model { + case "banana-2-1k-compliant", "banana-2-2k-compliant", "banana-2-4k-compliant": + return true + default: + return false + } +} + +func validateOutputFormat(raw string) error { + if outputFormatAllowed(raw) { + return nil + } + return fmt.Errorf("output_format is unsupported") +} + +func outputFormatAllowed(raw string) bool { + value := strings.ToLower(strings.TrimSpace(raw)) + return value == "" || value == "png" +} + +func hasJSONValue(raw json.RawMessage) bool { + raw = bytes.TrimSpace(raw) + return len(raw) > 0 && !bytes.Equal(raw, []byte("null")) +} diff --git a/internal/converter/sosana/images.go b/internal/converter/sosana/images.go new file mode 100644 index 00000000..c3a9a3cf --- /dev/null +++ b/internal/converter/sosana/images.go @@ -0,0 +1,446 @@ +package sosana + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "strconv" + "strings" + "time" + + "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" + "github.com/mixaill76/auto_ai_router/internal/converter/openai" +) + +const maxMultipartImageBytes = 20 * 1024 * 1024 + +const ( + StatusProcessing = "PROCESSING" + StatusCompleted = "COMPLETED" + StatusFailed = "FAILED" + StatusModerated = "MODERATED" +) + +type BananaCreateRequest struct { + Prompt string `json:"prompt"` + ImageURLs []string `json:"image_urls,omitempty"` + Model string `json:"model,omitempty"` + AspectRatio string `json:"aspect_ratio,omitempty"` + ImageSize string `json:"image_size,omitempty"` + PromptOptimization bool `json:"prompt_optimization"` +} + +type BananaTaskResponse struct { + UID string `json:"uid"` + Status string `json:"status"` + Prompt string `json:"prompt"` + CreatedAt string `json:"created_at"` + OptimizedPrompt string `json:"optimized_prompt"` + ResultFileURL *string `json:"result_file_url"` + Error *string `json:"error"` +} + +type openAIImageRequest struct { + openai.OpenAIImageRequest + AspectRatio string `json:"aspect_ratio,omitempty"` + Ratio string `json:"ratio,omitempty"` + ImageSize string `json:"image_size,omitempty"` +} + +func ImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, string, error) { + if reason := UnsupportedRequest("/v1/images/generations", openAIBody, "application/json"); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + + var req openAIImageRequest + if err := json.Unmarshal(openAIBody, &req); err != nil { + return nil, "", fmt.Errorf("failed to parse OpenAI image request: %w", err) + } + if err := validateImageCount(req.N); err != nil { + return nil, "", err + } + if err := validateResponseFormat(req.ResponseFormat); err != nil { + return nil, "", err + } + if err := validateOutputFormat(req.OutputFormat); err != nil { + return nil, "", err + } + prompt := strings.TrimSpace(req.Prompt) + if prompt == "" { + return nil, "", fmt.Errorf("image generation request missing prompt") + } + imageSize, err := imageSize(req.ImageSize, req.Size) + if err != nil { + return nil, "", err + } + concreteModel := providerModel(modelID, req.Model, imageSize) + if reason := UnsupportedModel(concreteModel); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + body, err := json.Marshal(BananaCreateRequest{ + Prompt: prompt, + Model: concreteModel, + AspectRatio: aspectRatio(req.AspectRatio, req.Ratio, req.Size), + ImageSize: imageSize, + PromptOptimization: false, + }) + if err != nil { + return nil, "", err + } + return body, concreteModel, nil +} + +func ImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, string, error) { + if reason := UnsupportedRequest("/v1/images/edits", openAIBody, contentType); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + + mediaType, params, err := mime.ParseMediaType(contentType) + if err != nil { + return nil, "", fmt.Errorf("failed to parse image edit content type: %w", err) + } + if !strings.HasPrefix(mediaType, "multipart/form-data") { + return nil, "", fmt.Errorf("image edits require multipart/form-data content type") + } + boundary := params["boundary"] + if boundary == "" { + return nil, "", fmt.Errorf("missing multipart boundary in content type") + } + + fields := make(map[string]string) + imageURLs := make([]string, 0, 1) + reader := multipart.NewReader(bytes.NewReader(openAIBody), boundary) + for { + part, err := reader.NextPart() + if err == io.EOF { + break + } + if err != nil { + return nil, "", fmt.Errorf("failed to read multipart image edit payload: %w", err) + } + + formName := part.FormName() + if formName == "" { + continue + } + data, err := readLimited(part, maxMultipartImageBytes) + if err != nil { + return nil, "", err + } + if part.FileName() == "" { + fields[formName] = strings.TrimSpace(string(data)) + continue + } + if formName == "mask" { + return nil, "", fmt.Errorf("image edits do not support mask") + } + if formName != "image" && formName != "images" && formName != "image[]" { + continue + } + mimeType := detectImageMIMEType(part.Header.Get("Content-Type"), data) + if mimeType != "image/png" { + return nil, "", fmt.Errorf("image edits support PNG images only") + } + imageURLs = append(imageURLs, "data:"+mimeType+";base64,"+base64.StdEncoding.EncodeToString(data)) + } + + if len(imageURLs) > maxInputImages { + return nil, "", fmt.Errorf("image edits support up to %d input images", maxInputImages) + } + if err := validateImageCountString(fields["n"]); err != nil { + return nil, "", err + } + if err := validateResponseFormat(fields["response_format"]); err != nil { + return nil, "", err + } + if err := validateOutputFormat(fields["output_format"]); err != nil { + return nil, "", err + } + prompt := strings.TrimSpace(fields["prompt"]) + if prompt == "" { + return nil, "", fmt.Errorf("image edit request missing prompt field") + } + if len(imageURLs) == 0 { + return nil, "", fmt.Errorf("image edit request missing image") + } + imageSize, err := imageSize(fields["image_size"], fields["size"]) + if err != nil { + return nil, "", err + } + concreteModel := providerModel(modelID, fields["model"], imageSize) + if reason := UnsupportedModel(concreteModel); reason != "" { + return nil, "", fmt.Errorf("%s", reason) + } + body, err := json.Marshal(BananaCreateRequest{ + Prompt: prompt, + ImageURLs: imageURLs, + Model: concreteModel, + AspectRatio: aspectRatio(fields["aspect_ratio"], fields["ratio"], fields["size"]), + ImageSize: imageSize, + PromptOptimization: false, + }) + if err != nil { + return nil, "", err + } + return body, concreteModel, nil +} + +func OpenAIImageResponse(task BananaTaskResponse, image []byte) ([]byte, error) { + if len(image) == 0 { + return nil, fmt.Errorf("image task completed without image bytes") + } + resp := openai.OpenAIImageResponse{ + Created: createdAtUnix(task.CreatedAt), + Data: []openai.OpenAIImageData{ + { + B64JSON: base64.StdEncoding.EncodeToString(image), + RevisedPrompt: strings.TrimSpace(task.OptimizedPrompt), + }, + }, + } + return json.Marshal(resp) +} + +func providerModel(modelID, requestModel, imageSize string) string { + model := strings.TrimSpace(modelID) + if model == "" { + model = strings.TrimSpace(requestModel) + } + if model == "" { + return "" + } + if strings.EqualFold(model, "google/gemini-3.1-flash-image-preview") { + return "banana-2-" + strings.ToLower(imageSize) + "-compliant" + } + return strings.ReplaceAll(model, "{image_size}", strings.ToLower(imageSize)) +} + +func createdAtUnix(value string) int64 { + if value == "" { + return converterutil.GetCurrentTimestamp() + } + if ts, err := time.Parse(time.RFC3339Nano, value); err == nil { + return ts.Unix() + } + return converterutil.GetCurrentTimestamp() +} + +func CreateURL(baseURL string) string { + return strings.TrimSuffix(baseURL, "/") + "/api/banana/create-async" +} + +func PollURL(baseURL, uid string) string { + return strings.TrimSuffix(baseURL, "/") + "/api/banana/" + uid +} + +func SizeToAspectRatio(size string) string { + if spec, ok := imageSpecFromExactSize(size); ok { + return spec.aspectRatio + } + return "auto" +} + +func aspectRatio(explicit, ratio, size string) string { + if value := strings.TrimSpace(explicit); value != "" { + return value + } + if value := strings.TrimSpace(ratio); value != "" { + return value + } + return SizeToAspectRatio(size) +} + +func imageSize(explicit, size string) (string, error) { + if strings.TrimSpace(explicit) != "" { + if value, ok := normalizeImageSize(explicit); ok { + return value, nil + } + return "", fmt.Errorf("image_size is unsupported") + } + if value, ok := imageSizeFromExactSize(size); ok { + return value, nil + } + return "", fmt.Errorf("size is unsupported") +} + +func normalizeImageSize(raw string) (string, bool) { + switch strings.ToLower(strings.TrimSpace(raw)) { + case "", "auto", "1k": + return "1K", true + case "2k": + return "2K", true + case "4k": + return "4K", true + default: + return "", false + } +} + +func imageSizeFromExactSize(size string) (string, bool) { + if spec, ok := imageSpecFromExactSize(size); ok { + return spec.imageSize, true + } + return "", false +} + +type exactImageSpec struct { + imageSize string + aspectRatio string +} + +func imageSpecFromExactSize(size string) (exactImageSpec, bool) { + switch strings.TrimSpace(size) { + case "", "auto": + return exactImageSpec{imageSize: "1K", aspectRatio: "auto"}, true + + case "1024x1024": + return exactImageSpec{imageSize: "1K", aspectRatio: "1:1"}, true + case "512x2048": + return exactImageSpec{imageSize: "1K", aspectRatio: "1:4"}, true + case "384x3072": + return exactImageSpec{imageSize: "1K", aspectRatio: "1:8"}, true + case "848x1264": + return exactImageSpec{imageSize: "1K", aspectRatio: "2:3"}, true + case "1264x848": + return exactImageSpec{imageSize: "1K", aspectRatio: "3:2"}, true + case "896x1200": + return exactImageSpec{imageSize: "1K", aspectRatio: "3:4"}, true + case "2048x512": + return exactImageSpec{imageSize: "1K", aspectRatio: "4:1"}, true + case "1200x896": + return exactImageSpec{imageSize: "1K", aspectRatio: "4:3"}, true + case "928x1152": + return exactImageSpec{imageSize: "1K", aspectRatio: "4:5"}, true + case "1152x928": + return exactImageSpec{imageSize: "1K", aspectRatio: "5:4"}, true + case "3072x384": + return exactImageSpec{imageSize: "1K", aspectRatio: "8:1"}, true + case "768x1376": + return exactImageSpec{imageSize: "1K", aspectRatio: "9:16"}, true + case "1376x768": + return exactImageSpec{imageSize: "1K", aspectRatio: "16:9"}, true + case "1584x672": + return exactImageSpec{imageSize: "1K", aspectRatio: "21:9"}, true + + case "2048x2048": + return exactImageSpec{imageSize: "2K", aspectRatio: "1:1"}, true + case "1024x4096": + return exactImageSpec{imageSize: "2K", aspectRatio: "1:4"}, true + case "768x6144": + return exactImageSpec{imageSize: "2K", aspectRatio: "1:8"}, true + case "1696x2528": + return exactImageSpec{imageSize: "2K", aspectRatio: "2:3"}, true + case "2528x1696": + return exactImageSpec{imageSize: "2K", aspectRatio: "3:2"}, true + case "1792x2400": + return exactImageSpec{imageSize: "2K", aspectRatio: "3:4"}, true + case "4096x1024": + return exactImageSpec{imageSize: "2K", aspectRatio: "4:1"}, true + case "2400x1792": + return exactImageSpec{imageSize: "2K", aspectRatio: "4:3"}, true + case "1856x2304": + return exactImageSpec{imageSize: "2K", aspectRatio: "4:5"}, true + case "2304x1856": + return exactImageSpec{imageSize: "2K", aspectRatio: "5:4"}, true + case "6144x768": + return exactImageSpec{imageSize: "2K", aspectRatio: "8:1"}, true + case "1536x2752": + return exactImageSpec{imageSize: "2K", aspectRatio: "9:16"}, true + case "2752x1536": + return exactImageSpec{imageSize: "2K", aspectRatio: "16:9"}, true + case "3168x1344": + return exactImageSpec{imageSize: "2K", aspectRatio: "21:9"}, true + + case "4096x4096": + return exactImageSpec{imageSize: "4K", aspectRatio: "1:1"}, true + case "2048x8192": + return exactImageSpec{imageSize: "4K", aspectRatio: "1:4"}, true + case "1536x12288": + return exactImageSpec{imageSize: "4K", aspectRatio: "1:8"}, true + case "3392x5056": + return exactImageSpec{imageSize: "4K", aspectRatio: "2:3"}, true + case "5056x3392": + return exactImageSpec{imageSize: "4K", aspectRatio: "3:2"}, true + case "3584x4800": + return exactImageSpec{imageSize: "4K", aspectRatio: "3:4"}, true + case "8192x2048": + return exactImageSpec{imageSize: "4K", aspectRatio: "4:1"}, true + case "4800x3584": + return exactImageSpec{imageSize: "4K", aspectRatio: "4:3"}, true + case "3712x4608": + return exactImageSpec{imageSize: "4K", aspectRatio: "4:5"}, true + case "4608x3712": + return exactImageSpec{imageSize: "4K", aspectRatio: "5:4"}, true + case "12288x1536": + return exactImageSpec{imageSize: "4K", aspectRatio: "8:1"}, true + case "3072x5504": + return exactImageSpec{imageSize: "4K", aspectRatio: "9:16"}, true + case "5504x3072": + return exactImageSpec{imageSize: "4K", aspectRatio: "16:9"}, true + case "6336x2688": + return exactImageSpec{imageSize: "4K", aspectRatio: "21:9"}, true + default: + return exactImageSpec{}, false + } +} + +func validateImageCount(n *int) error { + if n == nil || *n == 1 { + return nil + } + return fmt.Errorf("image requests support n=1 only") +} + +func validateResponseFormat(format string) error { + if strings.EqualFold(strings.TrimSpace(format), "url") { + return fmt.Errorf("response_format=url is unsupported for this image model") + } + return nil +} + +func validateImageCountString(raw string) error { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil + } + n, err := strconv.Atoi(raw) + if err != nil { + return fmt.Errorf("invalid image count: %w", err) + } + if n != 1 { + return fmt.Errorf("image requests support n=1 only") + } + return nil +} + +func readLimited(r io.Reader, limit int64) ([]byte, error) { + data, err := io.ReadAll(io.LimitReader(r, limit+1)) + if err != nil { + return nil, fmt.Errorf("failed to read multipart part: %w", err) + } + if int64(len(data)) > limit { + return nil, fmt.Errorf("multipart image exceeds %d bytes", limit) + } + return data, nil +} + +func detectImageMIMEType(header string, data []byte) string { + header = strings.ToLower(strings.TrimSpace(header)) + if strings.HasPrefix(header, "image/") { + mediaType, _, err := mime.ParseMediaType(header) + if err == nil { + return mediaType + } + return strings.TrimSpace(strings.Split(header, ";")[0]) + } + detected := http.DetectContentType(data) + if strings.HasPrefix(detected, "image/") { + return detected + } + return "application/octet-stream" +} diff --git a/internal/converter/sosana/images_test.go b/internal/converter/sosana/images_test.go new file mode 100644 index 00000000..b817ccfe --- /dev/null +++ b/internal/converter/sosana/images_test.go @@ -0,0 +1,293 @@ +package sosana + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "mime/multipart" + "strings" + "testing" + "time" + + "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestImageGenerationRequest(t *testing.T) { + tests := []struct { + name string + size string + wantAspect string + wantSize string + wantErrPart string + }{ + {name: "one k square", size: "1024x1024", wantAspect: "1:1", wantSize: "1K"}, + {name: "one k wide", size: "1376x768", wantAspect: "16:9", wantSize: "1K"}, + {name: "one k portrait", size: "768x1376", wantAspect: "9:16", wantSize: "1K"}, + {name: "one k tall", size: "512x2048", wantAspect: "1:4", wantSize: "1K"}, + {name: "two k square", size: "2048x2048", wantAspect: "1:1", wantSize: "2K"}, + {name: "two k wide", size: "2752x1536", wantAspect: "16:9", wantSize: "2K"}, + {name: "four k square", size: "4096x4096", wantAspect: "1:1", wantSize: "4K"}, + {name: "four k ultra wide", size: "6336x2688", wantAspect: "21:9", wantSize: "4K"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := []byte(`{"model":"banana-2-1k-compliant","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) + got, concreteModel, err := ImageGenerationRequest(body, "banana-2-{image_size}-compliant") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "draw a cat", req.Prompt) + assert.Equal(t, "banana-2-"+strings.ToLower(tt.wantSize)+"-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) + assert.Equal(t, tt.wantAspect, req.AspectRatio) + assert.Equal(t, tt.wantSize, req.ImageSize) + assert.False(t, req.PromptOptimization) + assert.Empty(t, req.ImageURLs) + }) + } +} + +func TestImageGenerationRequestPrefersProviderModel(t *testing.T) { + got, concreteModel, err := ImageGenerationRequest([]byte(`{"model":"public-image","prompt":"draw","n":1}`), "banana-2-1k-compliant") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) +} + +func TestImageGenerationRequestMapsPublicGeminiModelToSosanaTier(t *testing.T) { + got, concreteModel, err := ImageGenerationRequest( + []byte(`{"model":"google/gemini-3.1-flash-image-preview","prompt":"draw","image_size":"2K","n":1}`), + "google/gemini-3.1-flash-image-preview", + ) + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "banana-2-2k-compliant", req.Model) + assert.Equal(t, "banana-2-2k-compliant", concreteModel) + assert.Equal(t, "2K", req.ImageSize) +} + +func TestImageGenerationRequestUsesExplicitAspectRatio(t *testing.T) { + got, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","size":"1024x1024","aspect_ratio":"16:9"}`), "banana-2-1k-compliant") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "16:9", req.AspectRatio) +} + +func TestImageGenerationRequestRejectsMultipleImages(t *testing.T) { + _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","n":2}`), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "n=1") +} + +func TestImageGenerationRequestRejectsURLResponseFormat(t *testing.T) { + _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","response_format":"url"}`), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "response_format=url") +} + +func TestImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { + tests := []struct { + name string + body string + want string + }{ + {name: "tools", body: `{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`, want: "tools"}, + {name: "thinking", body: `{"model":"banana-2-1k-compliant","prompt":"draw","thinking_level":"high"}`, want: "thinking_level"}, + {name: "output format", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"jpeg"}`, want: "output_format"}, + {name: "output compression", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_compression":0}`, want: "output_compression"}, + {name: "quality auto", body: `{"model":"banana-2-1k-compliant","prompt":"draw","quality":"auto"}`, want: "quality"}, + {name: "messages", body: `{"model":"banana-2-1k-compliant","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, + {name: "image size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_size":"0.5K"}`, want: "image_size"}, + {name: "exact size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"512x512"}`, want: "size"}, + {name: "legacy openai size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"1792x1024"}`, want: "size"}, + {name: "unknown size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"333x777"}`, want: "size"}, + {name: "reference images", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, _, err := ImageGenerationRequest([]byte(tt.body), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), tt.want) + }) + } +} + +func TestImageGenerationRequestAllowsPNGOutputFormat(t *testing.T) { + _, _, err := ImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"png","response_format":"b64_json"}`), "banana-2-1k-compliant") + require.NoError(t, err) +} + +func TestImageEditRequest(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + "size": "1024x1024", + "n": "1", + }, map[string][]byte{ + "image": pngBytes(), + }) + + got, concreteModel, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "make it blue", req.Prompt) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) + assert.Equal(t, "1:1", req.AspectRatio) + assert.Equal(t, "1K", req.ImageSize) + assert.False(t, req.PromptOptimization) + require.Len(t, req.ImageURLs, 1) + assert.True(t, strings.HasPrefix(req.ImageURLs[0], "data:image/png;base64,")) + assert.Contains(t, req.ImageURLs[0], base64.StdEncoding.EncodeToString(pngBytes())) +} + +func TestImageEditRequestPrefersProviderModel(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "public-image", + "prompt": "make it blue", + }, map[string][]byte{ + "image": pngBytes(), + }) + + got, concreteModel, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.NoError(t, err) + + var req BananaCreateRequest + require.NoError(t, json.Unmarshal(got, &req)) + assert.Equal(t, "banana-2-1k-compliant", req.Model) + assert.Equal(t, req.Model, concreteModel) +} + +func TestImageEditRequestRejectsMask(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + }, map[string][]byte{ + "image": pngBytes(), + "mask": pngBytes(), + }) + + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "mask") +} + +func TestImageEditRequestRejectsMultipleImagesCount(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + "n": "2", + }, map[string][]byte{ + "image": pngBytes(), + }) + + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "n=1") +} + +func TestImageEditRequestRejectsURLResponseFormat(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + "response_format": "url", + }, map[string][]byte{ + "image": pngBytes(), + }) + + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "response_format=url") +} + +func TestImageEditRequestRejectsJPEGInput(t *testing.T) { + body, contentType := multipartImageEditBody(t, map[string]string{ + "model": "banana-2-1k-compliant", + "prompt": "make it blue", + }, map[string][]byte{ + "image": jpegBytes(), + }) + + _, _, err := ImageEditRequest(body, contentType, "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "PNG") +} + +func TestImageEditRequestRejectsTooManyImages(t *testing.T) { + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("model", "banana-2-1k-compliant")) + require.NoError(t, writer.WriteField("prompt", "make it blue")) + for i := 0; i < maxInputImages+1; i++ { + part, err := writer.CreateFormFile("image", fmt.Sprintf("image-%02d.png", i)) + require.NoError(t, err) + _, err = part.Write(pngBytes()) + require.NoError(t, err) + } + require.NoError(t, writer.Close()) + + _, _, err := ImageEditRequest(buf.Bytes(), writer.FormDataContentType(), "banana-2-1k-compliant") + require.Error(t, err) + assert.Contains(t, err.Error(), "too many") +} + +func TestOpenAIImageResponse(t *testing.T) { + createdAt := "2026-01-01T00:00:00Z" + body, err := OpenAIImageResponse(BananaTaskResponse{ + Status: StatusCompleted, + CreatedAt: createdAt, + OptimizedPrompt: "A detailed result prompt", + }, pngBytes()) + require.NoError(t, err) + + var resp openai.OpenAIImageResponse + require.NoError(t, json.Unmarshal(body, &resp)) + require.Len(t, resp.Data, 1) + assert.Empty(t, resp.Data[0].URL) + assert.Equal(t, "A detailed result prompt", resp.Data[0].RevisedPrompt) + assert.Equal(t, base64.StdEncoding.EncodeToString(pngBytes()), resp.Data[0].B64JSON) + ts, err := time.Parse(time.RFC3339, createdAt) + require.NoError(t, err) + assert.Equal(t, ts.Unix(), resp.Created) +} + +func multipartImageEditBody(t *testing.T, fields map[string]string, files map[string][]byte) ([]byte, string) { + t.Helper() + + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + for key, value := range fields { + require.NoError(t, writer.WriteField(key, value)) + } + for key, data := range files { + part, err := writer.CreateFormFile(key, key+".png") + require.NoError(t, err) + _, err = part.Write(data) + require.NoError(t, err) + } + require.NoError(t, writer.Close()) + return buf.Bytes(), writer.FormDataContentType() +} + +func pngBytes() []byte { + return []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} +} + +func jpegBytes() []byte { + return []byte{0xff, 0xd8, 0xff, 0xdb, 0, 0x43, 0, 1, 2, 3} +} diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index fc57183e..528fdcf9 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -3,48 +3,27 @@ package proxy import ( "bytes" "context" - "encoding/base64" "encoding/json" "errors" - "fmt" "io" "math/rand" - "mime" - "mime/multipart" "net" "net/http" "net/url" - "strconv" "strings" "time" "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter" - "github.com/mixaill76/auto_ai_router/internal/converter/converterutil" - "github.com/mixaill76/auto_ai_router/internal/converter/openai" - "github.com/mixaill76/auto_ai_router/internal/scope" - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/trace" + "github.com/mixaill76/auto_ai_router/internal/converter/sosana" ) const ( - sosanaPollInterval = 2 * time.Second - maxSosanaMultipartImageBytes = 20 * 1024 * 1024 - maxSosanaInputImages = 14 - maxSosanaResultImageBytes int64 = 32 * 1024 * 1024 - maxSosanaResultErrorBytes int64 = 16 * 1024 + sosanaPollInterval = 2 * time.Second + maxSosanaResultImageBytes int64 = 32 * 1024 * 1024 + maxSosanaResultErrorBytes int64 = 16 * 1024 ) -const ( - sosanaStatusProcessing = "PROCESSING" - sosanaStatusCompleted = "COMPLETED" - sosanaStatusFailed = "FAILED" - sosanaStatusModerated = "MODERATED" -) - -const unsupportedImageProviderRequestMessage = "request parameters are not supported by available image providers" -const unsupportedProviderEndpointMessage = "request endpoint is not supported by available providers" - var allowPrivateSosanaResultURLForTests func(*url.URL) bool type sosanaAttemptResult struct { @@ -54,175 +33,6 @@ type sosanaAttemptResult struct { retryReason RetryReason } -type sosanaBananaCreateRequest struct { - Prompt string `json:"prompt"` - ImageURLs []string `json:"image_urls,omitempty"` - Model string `json:"model,omitempty"` - AspectRatio string `json:"aspect_ratio,omitempty"` - ImageSize string `json:"image_size,omitempty"` - PromptOptimization bool `json:"prompt_optimization"` -} - -type sosanaBananaTaskResponse struct { - UID string `json:"uid"` - Status string `json:"status"` - Prompt string `json:"prompt"` - CreatedAt string `json:"created_at"` - OptimizedPrompt string `json:"optimized_prompt"` - ResultFileURL *string `json:"result_file_url"` - Error *string `json:"error"` -} - -type sosanaOpenAIImageRequest struct { - openai.OpenAIImageRequest - AspectRatio string `json:"aspect_ratio,omitempty"` - Ratio string `json:"ratio,omitempty"` - ImageSize string `json:"image_size,omitempty"` -} - -func (p *Proxy) applySosanaCompatibilityRouting( - w http.ResponseWriter, - r *http.Request, - prepared *orchestratedRequest, - modelID string, - cred **config.CredentialConfig, - body *[]byte, - proxyBody *[]byte, - realModelID *string, - isImageGeneration bool, - isImageEdit bool, - logCtx *RequestLogContext, - start time.Time, -) bool { - if (*cred).Type != config.ProviderTypeSosana || (!isImageGeneration && !isImageEdit) { - return true - } - - reason := unsupportedSosanaModel(*realModelID) - if reason == "" { - reason = unsupportedSosanaRequest(r.URL.Path, *body, r.Header.Get("Content-Type")) - } - if reason == "" { - return true - } - - nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, logCtx.Scope, reason) - if routed { - *cred = nextCred - *body = nextReq.body - *proxyBody = nextReq.proxyBody - *realModelID = nextReq.realModelID - r.URL.Path = nextReq.path - prepared.body = nextReq.body - prepared.proxyBody = nextReq.proxyBody - prepared.proxyPath = nextReq.proxyPath - prepared.realModelID = nextReq.realModelID - prepared.convertedResp = nextReq.convertedResp - prepared.passthroughResponses = nextReq.passthroughResponses - prepared.nativeResponses = nextReq.nativeResponses - logCtx.RealModelID = *realModelID - if span := trace.SpanFromContext(r.Context()); span.IsRecording() { - span.SetAttributes( - attribute.String("aar.real_model", *realModelID), - attribute.String("aar.credential", nextCred.Name), - attribute.String("aar.provider", string(nextCred.Type)), - attribute.Bool("aar.provider_compatibility_skip", true), - ) - } - return true - } - - success, fallbackReason := p.TryFallbackProxy( - w, - requestWithPath(r, prepared.proxyPath), - modelID, - (*cred).Name, - http.StatusBadRequest, - RetryReasonServerErr, - *proxyBody, - start, - logCtx, - ) - if success { - return false - } - p.logger.DebugContext(r.Context(), "No fallback handled unsupported image provider request", - "credential", (*cred).Name, - "model", modelID, - "reason", reason, - "fallback_reason", fallbackReason) - logCtx.Credential = *cred - logCtx.Status = "failure" - logCtx.HTTPStatus = http.StatusBadRequest - logCtx.ErrorMsg = unsupportedImageProviderRequestMessage - WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) - return false -} - -func (p *Proxy) nextPrimaryAfterUnsupportedSosana( - r *http.Request, - prepared *orchestratedRequest, - modelID string, - currentCred *config.CredentialConfig, - visibility scope.Context, - reason string, -) (*config.CredentialConfig, credentialPreparedRequest, bool) { - triedCreds := GetTried(r.Context()) - triedCreds[currentCred.Name] = true - - for attempts := 0; attempts < 128; attempts++ { - candidate, err := p.balancer.NextForModelExcludingScoped(modelID, triedCreds, visibility) - if err != nil { - p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", - "model", modelID, - "credential", currentCred.Name, - "reason", reason, - "error", err) - return nil, credentialPreparedRequest{}, false - } - triedCreds[candidate.Name] = true - if candidate.Type == config.ProviderTypeSosana { - continue - } - - nextReq, prepErr := p.prepareRequestForCredential( - r, - prepared.baseBody, - prepared.baseProxyBody, - modelID, - prepared.baseRealModelID, - prepared.basePath, - prepared.streaming, - candidate, - prepared.isResponsesAPI, - prepared.responsesPrevHandled, - prepared.stickyCacheEligible, - ) - if prepErr != nil { - p.logger.WarnContext(r.Context(), "Failed to prepare alternate primary request after image compatibility skip", - "credential", candidate.Name, - "provider", string(candidate.Type), - "model", modelID, - "reason", reason, - "error", prepErr) - continue - } - - p.logger.InfoContext(r.Context(), "Skipping incompatible image credential for unsupported image request", - "credential", currentCred.Name, - "next_credential", candidate.Name, - "model", modelID, - "reason", reason) - return candidate, nextReq, true - } - - p.logger.WarnContext(r.Context(), "Image compatibility skip exhausted primary credential scan", - "credential", currentCred.Name, - "model", modelID, - "reason", reason) - return nil, credentialPreparedRequest{}, false -} - func (p *Proxy) handleSosanaRequest( w http.ResponseWriter, r *http.Request, @@ -356,9 +166,9 @@ func (p *Proxy) sosanaRealModelIDForCredential(modelID, fallbackRealModelID stri func (p *Proxy) buildSosanaCreateBody(body []byte, contentType, realModelID string, isImageEdit bool) ([]byte, string, error) { if isImageEdit { - return buildSosanaImageEditRequest(body, contentType, realModelID) + return sosana.ImageEditRequest(body, contentType, realModelID) } - return buildSosanaImageGenerationRequest(body, realModelID) + return sosana.ImageGenerationRequest(body, realModelID) } func (p *Proxy) createAndPollSosanaTask( @@ -368,9 +178,9 @@ func (p *Proxy) createAndPollSosanaTask( createBody []byte, logCtx *RequestLogContext, ) sosanaAttemptResult { - task, rawBody, statusCode, err := p.doSosanaTaskRequest(ctx, http.MethodPost, sosanaCreateURL(cred.BaseURL), cred, createBody) + task, rawBody, statusCode, err := p.doSosanaTaskRequest(ctx, http.MethodPost, sosana.CreateURL(cred.BaseURL), cred, createBody) if err != nil { - body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosanaCreateURL(cred.BaseURL), logCtx) + body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosana.CreateURL(cred.BaseURL), logCtx) return sosanaAttemptResult{ body: body, statusCode: code, @@ -379,7 +189,7 @@ func (p *Proxy) createAndPollSosanaTask( } if statusCode >= 400 { p.logUpstreamError(ctx, "Sosana create request completed with error status", statusCode, cred, modelID, rawBody, - "url", sosanaCreateURL(cred.BaseURL), + "url", sosana.CreateURL(cred.BaseURL), "request_id", logCtx.RequestID) retryable, reason := ShouldRetryWithFallback(statusCode, rawBody) return sosanaAttemptResult{ @@ -401,7 +211,7 @@ func (p *Proxy) createAndPollSosanaTask( select { case <-ctx.Done(): p.logUpstreamError(context.Background(), "Sosana task polling timed out", http.StatusRequestTimeout, cred, modelID, rawBody, - "url", sosanaPollURL(cred.BaseURL, task.UID), + "url", sosana.PollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID, "error", ctx.Err()) return sosanaAttemptResult{body: maskedUpstreamErrorBody(http.StatusRequestTimeout), statusCode: http.StatusRequestTimeout} @@ -410,14 +220,14 @@ func (p *Proxy) createAndPollSosanaTask( } immediatePoll = false - task, rawBody, statusCode, err = p.doSosanaTaskRequest(ctx, http.MethodGet, sosanaPollURL(cred.BaseURL, task.UID), cred, nil) + task, rawBody, statusCode, err = p.doSosanaTaskRequest(ctx, http.MethodGet, sosana.PollURL(cred.BaseURL, task.UID), cred, nil) if err != nil { - body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosanaPollURL(cred.BaseURL, task.UID), logCtx) + body, code := p.sosanaTransportError(ctx, err, cred, modelID, sosana.PollURL(cred.BaseURL, task.UID), logCtx) return sosanaAttemptResult{body: body, statusCode: code} } if statusCode >= 400 { p.logUpstreamError(ctx, "Sosana poll request completed with error status", statusCode, cred, modelID, rawBody, - "url", sosanaPollURL(cred.BaseURL, task.UID), + "url", sosana.PollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID) return sosanaAttemptResult{body: maskedUpstreamErrorBody(statusCode), statusCode: statusCode} } @@ -428,26 +238,26 @@ func (p *Proxy) sosanaTaskBody( ctx context.Context, cred *config.CredentialConfig, modelID string, - task sosanaBananaTaskResponse, + task sosana.BananaTaskResponse, rawBody []byte, statusCode int, logCtx *RequestLogContext, ) ([]byte, int, bool) { switch task.Status { - case sosanaStatusCompleted: + case sosana.StatusCompleted: body, statusCode := p.sosanaCompletedImageBody(ctx, cred, modelID, task, rawBody, logCtx) return body, statusCode, true - case sosanaStatusFailed: + case sosana.StatusFailed: p.logUpstreamError(ctx, "Sosana task failed", http.StatusBadGateway, cred, modelID, rawBody, - "url", sosanaPollURL(cred.BaseURL, task.UID), + "url", sosana.PollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID) return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true - case sosanaStatusModerated: + case sosana.StatusModerated: p.logUpstreamError(ctx, "Sosana task moderated", http.StatusBadRequest, cred, modelID, rawBody, - "url", sosanaPollURL(cred.BaseURL, task.UID), + "url", sosana.PollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID) return maskedContentPolicyBody(), http.StatusBadRequest, true - case sosanaStatusProcessing: + case sosana.StatusProcessing: if task.UID == "" { p.logUpstreamError(ctx, "Sosana processing task missing uid", http.StatusBadGateway, cred, modelID, rawBody, "url", cred.BaseURL, @@ -457,7 +267,7 @@ func (p *Proxy) sosanaTaskBody( return nil, statusCode, false default: p.logUpstreamError(ctx, "Sosana task returned unknown status", http.StatusBadGateway, cred, modelID, rawBody, - "url", sosanaPollURL(cred.BaseURL, task.UID), + "url", sosana.PollURL(cred.BaseURL, task.UID), "request_id", logCtx.RequestID, "status", task.Status) return maskedUpstreamErrorBody(http.StatusBadGateway), http.StatusBadGateway, true @@ -468,7 +278,7 @@ func (p *Proxy) sosanaCompletedImageBody( ctx context.Context, cred *config.CredentialConfig, modelID string, - task sosanaBananaTaskResponse, + task sosana.BananaTaskResponse, rawBody []byte, logCtx *RequestLogContext, ) ([]byte, int) { @@ -476,7 +286,7 @@ func (p *Proxy) sosanaCompletedImageBody( if err != nil { return maskedUpstreamErrorBody(statusCode), statusCode } - body, err := buildSosanaOpenAIImageResponse(task, image) + body, err := sosana.OpenAIImageResponse(task, image) if err != nil { p.logUpstreamError(ctx, "Sosana completed task could not be converted", http.StatusBadGateway, cred, modelID, nil, "request_id", logCtx.RequestID, @@ -497,7 +307,7 @@ func (p *Proxy) downloadSosanaResultImage( ctx context.Context, cred *config.CredentialConfig, modelID string, - task sosanaBananaTaskResponse, + task sosana.BananaTaskResponse, rawTaskBody []byte, logCtx *RequestLogContext, ) ([]byte, string, int, error) { @@ -730,7 +540,7 @@ func isUnsafeSosanaResultIP(ip net.IP) bool { ip.IsUnspecified() } -func sosanaResultHost(task sosanaBananaTaskResponse) string { +func sosanaResultHost(task sosana.BananaTaskResponse) string { if task.ResultFileURL == nil { return "" } @@ -790,7 +600,7 @@ func isTextContentType(contentType string) bool { strings.Contains(contentType, "xml") } -func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosanaBananaTaskResponse, []byte, int, error) { +func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosana.BananaTaskResponse, []byte, int, error) { var reader *bytes.Reader if body != nil { reader = bytes.NewReader(body) @@ -799,7 +609,7 @@ func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cre } req, err := http.NewRequestWithContext(ctx, method, url, reader) if err != nil { - return sosanaBananaTaskResponse{}, nil, http.StatusInternalServerError, err + return sosana.BananaTaskResponse{}, nil, http.StatusInternalServerError, err } req.Header.Set("Authorization", "Bearer "+cred.APIKey) if body != nil { @@ -808,7 +618,7 @@ func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cre resp, err := p.client.Do(req) if err != nil { - return sosanaBananaTaskResponse{}, nil, http.StatusBadGateway, err + return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err } defer func() { if closeErr := resp.Body.Close(); closeErr != nil { @@ -818,9 +628,9 @@ func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cre rawBody, err := p.readLimitedResponseBody(resp.Body) if err != nil { - return sosanaBananaTaskResponse{}, nil, http.StatusBadGateway, err + return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err } - var task sosanaBananaTaskResponse + var task sosana.BananaTaskResponse if len(rawBody) > 0 { _ = json.Unmarshal(rawBody, &task) } @@ -838,692 +648,3 @@ func (p *Proxy) sosanaTransportError(ctx context.Context, err error, cred *confi "error", err) return maskedUpstreamErrorBody(statusCode), statusCode } - -func buildSosanaImageGenerationRequest(openAIBody []byte, modelID string) ([]byte, string, error) { - if reason := unsupportedSosanaRequest("/v1/images/generations", openAIBody, "application/json"); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - - var req sosanaOpenAIImageRequest - if err := json.Unmarshal(openAIBody, &req); err != nil { - return nil, "", fmt.Errorf("failed to parse OpenAI image request: %w", err) - } - if err := validateSosanaImageCount(req.N); err != nil { - return nil, "", err - } - if err := validateSosanaResponseFormat(req.ResponseFormat); err != nil { - return nil, "", err - } - if err := validateSosanaOutputFormat(req.OutputFormat); err != nil { - return nil, "", err - } - prompt := strings.TrimSpace(req.Prompt) - if prompt == "" { - return nil, "", fmt.Errorf("image generation request missing prompt") - } - imageSize, err := sosanaImageSize(req.ImageSize, req.Size) - if err != nil { - return nil, "", err - } - concreteModel := sosanaProviderModel(modelID, req.Model, imageSize) - if reason := unsupportedSosanaModel(concreteModel); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - body, err := json.Marshal(sosanaBananaCreateRequest{ - Prompt: prompt, - Model: concreteModel, - AspectRatio: sosanaAspectRatio(req.AspectRatio, req.Ratio, req.Size), - ImageSize: imageSize, - PromptOptimization: false, - }) - if err != nil { - return nil, "", err - } - return body, concreteModel, nil -} - -func buildSosanaImageEditRequest(openAIBody []byte, contentType string, modelID string) ([]byte, string, error) { - if reason := unsupportedSosanaRequest("/v1/images/edits", openAIBody, contentType); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - - mediaType, params, err := mime.ParseMediaType(contentType) - if err != nil { - return nil, "", fmt.Errorf("failed to parse image edit content type: %w", err) - } - if !strings.HasPrefix(mediaType, "multipart/form-data") { - return nil, "", fmt.Errorf("image edits require multipart/form-data content type") - } - boundary := params["boundary"] - if boundary == "" { - return nil, "", fmt.Errorf("missing multipart boundary in content type") - } - - fields := make(map[string]string) - imageURLs := make([]string, 0, 1) - reader := multipart.NewReader(bytes.NewReader(openAIBody), boundary) - for { - part, err := reader.NextPart() - if err == io.EOF { - break - } - if err != nil { - return nil, "", fmt.Errorf("failed to read multipart image edit payload: %w", err) - } - - formName := part.FormName() - if formName == "" { - continue - } - data, err := readLimitedSosanaMultipartPart(part, maxSosanaMultipartImageBytes) - if err != nil { - return nil, "", err - } - if part.FileName() == "" { - fields[formName] = strings.TrimSpace(string(data)) - continue - } - if formName == "mask" { - return nil, "", fmt.Errorf("image edits do not support mask") - } - if formName != "image" && formName != "images" && formName != "image[]" { - continue - } - mimeType := detectSosanaImageMIMEType(part.Header.Get("Content-Type"), data) - if mimeType != "image/png" { - return nil, "", fmt.Errorf("image edits support PNG images only") - } - imageURLs = append(imageURLs, "data:"+mimeType+";base64,"+base64.StdEncoding.EncodeToString(data)) - } - - if len(imageURLs) > maxSosanaInputImages { - return nil, "", fmt.Errorf("image edits support up to %d input images", maxSosanaInputImages) - } - if err := validateSosanaImageCountString(fields["n"]); err != nil { - return nil, "", err - } - if err := validateSosanaResponseFormat(fields["response_format"]); err != nil { - return nil, "", err - } - if err := validateSosanaOutputFormat(fields["output_format"]); err != nil { - return nil, "", err - } - prompt := strings.TrimSpace(fields["prompt"]) - if prompt == "" { - return nil, "", fmt.Errorf("image edit request missing prompt field") - } - if len(imageURLs) == 0 { - return nil, "", fmt.Errorf("image edit request missing image") - } - imageSize, err := sosanaImageSize(fields["image_size"], fields["size"]) - if err != nil { - return nil, "", err - } - concreteModel := sosanaProviderModel(modelID, fields["model"], imageSize) - if reason := unsupportedSosanaModel(concreteModel); reason != "" { - return nil, "", fmt.Errorf("%s", reason) - } - body, err := json.Marshal(sosanaBananaCreateRequest{ - Prompt: prompt, - ImageURLs: imageURLs, - Model: concreteModel, - AspectRatio: sosanaAspectRatio(fields["aspect_ratio"], fields["ratio"], fields["size"]), - ImageSize: imageSize, - PromptOptimization: false, - }) - if err != nil { - return nil, "", err - } - return body, concreteModel, nil -} - -func buildSosanaOpenAIImageResponse(task sosanaBananaTaskResponse, image []byte) ([]byte, error) { - if len(image) == 0 { - return nil, fmt.Errorf("image task completed without image bytes") - } - resp := openai.OpenAIImageResponse{ - Created: sosanaCreatedAtUnix(task.CreatedAt), - Data: []openai.OpenAIImageData{ - { - B64JSON: base64.StdEncoding.EncodeToString(image), - RevisedPrompt: strings.TrimSpace(task.OptimizedPrompt), - }, - }, - } - return json.Marshal(resp) -} - -var unsupportedSosanaImageFields = []string{ - "tools", - "tool_choice", - "google_search", - "thinking_level", - "thinking_budget", - "thinking_config", - "thinking", - "reasoning_effort", - "generation_config", - "temperature", - "top_p", - "top_k", - "seed", - "max_tokens", - "stop", - "stream", - "messages", - "extra_body", - "image", - "images", - "image_urls", - "reference_images", -} - -func unsupportedSosanaRequest(path string, body []byte, contentType string) string { - switch { - case strings.Contains(path, "/images/generations"): - return unsupportedSosanaGenerationRequest(body) - case strings.Contains(path, "/images/edits"): - return unsupportedSosanaEditRequest(body, contentType) - default: - return "endpoint is unsupported" - } -} - -func unsupportedSosanaModel(modelID string) string { - if supportedSosanaModel(modelID) { - return "" - } - return "model is unsupported" -} - -func unsupportedSosanaGenerationRequest(body []byte) string { - var raw map[string]json.RawMessage - if err := json.Unmarshal(body, &raw); err != nil { - return "" - } - return unsupportedSosanaImageFieldsInJSON(raw) -} - -func unsupportedSosanaEditRequest(body []byte, contentType string) string { - mediaType, params, err := mime.ParseMediaType(contentType) - if err != nil || !strings.HasPrefix(mediaType, "multipart/form-data") { - return "" - } - boundary := params["boundary"] - if boundary == "" { - return "" - } - - fields := make(map[string]string) - imageCount := 0 - reader := multipart.NewReader(bytes.NewReader(body), boundary) - for { - part, err := reader.NextPart() - if err != nil { - break - } - - formName := part.FormName() - if formName == "" { - continue - } - data, err := readLimitedSosanaMultipartPart(part, maxSosanaMultipartImageBytes) - if err != nil { - return err.Error() - } - if part.FileName() == "" { - fields[formName] = strings.TrimSpace(string(data)) - continue - } - if formName == "mask" { - return "mask is unsupported" - } - if formName != "image" && formName != "images" && formName != "image[]" { - continue - } - imageCount++ - if detectSosanaImageMIMEType(part.Header.Get("Content-Type"), data) != "image/png" { - return "only PNG input images are supported" - } - } - if imageCount > maxSosanaInputImages { - return "too many input images" - } - return unsupportedSosanaImageFieldsInForm(fields) -} - -func unsupportedSosanaImageFieldsInJSON(raw map[string]json.RawMessage) string { - if reason := unsupportedSosanaJSONImageCount(raw["n"]); reason != "" { - return reason - } - if reason := unsupportedSosanaJSONResponseFormat(raw["response_format"]); reason != "" { - return reason - } - if reason := unsupportedSosanaJSONOutputFormat(raw["output_format"]); reason != "" { - return reason - } - if reason := unsupportedSosanaJSONImageSize(raw["image_size"]); reason != "" { - return reason - } - if reason := unsupportedSosanaJSONExactSize(raw["size"]); reason != "" { - return reason - } - for _, field := range []string{"quality", "style", "background", "moderation"} { - if hasSosanaJSONValue(raw[field]) { - return field + " is unsupported" - } - } - if hasSosanaJSONValue(raw["output_compression"]) { - return "output_compression is unsupported" - } - for _, field := range unsupportedSosanaImageFields { - if hasSosanaJSONValue(raw[field]) { - return field + " is unsupported" - } - } - return "" -} - -func unsupportedSosanaImageFieldsInForm(fields map[string]string) string { - if err := validateSosanaImageCountString(fields["n"]); err != nil { - return err.Error() - } - if err := validateSosanaResponseFormat(fields["response_format"]); err != nil { - return err.Error() - } - if err := validateSosanaOutputFormat(fields["output_format"]); err != nil { - return err.Error() - } - if reason := unsupportedSosanaFormImageSize(fields["image_size"]); reason != "" { - return reason - } - if reason := unsupportedSosanaFormExactSize(fields["size"]); reason != "" { - return reason - } - for _, field := range []string{"quality", "style", "background", "moderation"} { - if strings.TrimSpace(fields[field]) != "" { - return field + " is unsupported" - } - } - if _, ok := fields["output_compression"]; ok { - return "output_compression is unsupported" - } - for _, field := range unsupportedSosanaImageFields { - if strings.TrimSpace(fields[field]) != "" { - return field + " is unsupported" - } - } - return "" -} - -func unsupportedSosanaJSONImageCount(raw json.RawMessage) string { - if !hasSosanaJSONValue(raw) { - return "" - } - var n int - if err := json.Unmarshal(raw, &n); err == nil { - if n == 1 { - return "" - } - return "image requests support n=1 only" - } - var f float64 - if err := json.Unmarshal(raw, &f); err == nil { - if f == 1 { - return "" - } - return "image requests support n=1 only" - } - return "invalid image count" -} - -func unsupportedSosanaJSONResponseFormat(raw json.RawMessage) string { - if !hasSosanaJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "response_format is unsupported" - } - if strings.EqualFold(strings.TrimSpace(value), "b64_json") || strings.TrimSpace(value) == "" { - return "" - } - if strings.EqualFold(strings.TrimSpace(value), "url") { - return "response_format=url is unsupported for this image model" - } - return "response_format is unsupported" -} - -func unsupportedSosanaJSONOutputFormat(raw json.RawMessage) string { - if !hasSosanaJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "output_format is unsupported" - } - if sosanaOutputFormatAllowed(value) { - return "" - } - return "output_format is unsupported" -} - -func unsupportedSosanaJSONImageSize(raw json.RawMessage) string { - if !hasSosanaJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "image_size is unsupported" - } - if _, ok := normalizeSosanaImageSize(value); ok { - return "" - } - return "image_size is unsupported" -} - -func unsupportedSosanaJSONExactSize(raw json.RawMessage) string { - if !hasSosanaJSONValue(raw) { - return "" - } - var value string - if err := json.Unmarshal(raw, &value); err != nil { - return "size is unsupported" - } - if _, ok := sosanaImageSizeFromExactSize(value); ok { - return "" - } - return "size is unsupported" -} - -func unsupportedSosanaFormImageSize(raw string) string { - if strings.TrimSpace(raw) == "" { - return "" - } - if _, ok := normalizeSosanaImageSize(raw); ok { - return "" - } - return "image_size is unsupported" -} - -func unsupportedSosanaFormExactSize(raw string) string { - if strings.TrimSpace(raw) == "" { - return "" - } - if _, ok := sosanaImageSizeFromExactSize(raw); ok { - return "" - } - return "size is unsupported" -} - -func supportedSosanaModel(modelID string) bool { - model := strings.ToLower(strings.TrimSpace(modelID)) - if model == "" || model == "google/gemini-3.1-flash-image-preview" { - return true - } - if model == "banana-2-{image_size}-compliant" { - return true - } - switch model { - case "banana-2-1k-compliant", "banana-2-2k-compliant", "banana-2-4k-compliant": - return true - default: - return false - } -} - -func validateSosanaOutputFormat(raw string) error { - if sosanaOutputFormatAllowed(raw) { - return nil - } - return fmt.Errorf("output_format is unsupported") -} - -func sosanaOutputFormatAllowed(raw string) bool { - value := strings.ToLower(strings.TrimSpace(raw)) - return value == "" || value == "png" -} - -func hasSosanaJSONValue(raw json.RawMessage) bool { - raw = bytes.TrimSpace(raw) - return len(raw) > 0 && !bytes.Equal(raw, []byte("null")) -} - -func sosanaProviderModel(modelID, requestModel, imageSize string) string { - model := strings.TrimSpace(modelID) - if model == "" { - model = strings.TrimSpace(requestModel) - } - if model == "" { - return "" - } - if strings.EqualFold(model, "google/gemini-3.1-flash-image-preview") { - return "banana-2-" + strings.ToLower(imageSize) + "-compliant" - } - return strings.ReplaceAll(model, "{image_size}", strings.ToLower(imageSize)) -} - -func sosanaCreatedAtUnix(value string) int64 { - if value == "" { - return converterutil.GetCurrentTimestamp() - } - if ts, err := time.Parse(time.RFC3339Nano, value); err == nil { - return ts.Unix() - } - return converterutil.GetCurrentTimestamp() -} - -func sosanaCreateURL(baseURL string) string { - return strings.TrimSuffix(baseURL, "/") + "/api/banana/create-async" -} - -func sosanaPollURL(baseURL, uid string) string { - return strings.TrimSuffix(baseURL, "/") + "/api/banana/" + uid -} - -func sosanaSizeToAspectRatio(size string) string { - if spec, ok := sosanaImageSpecFromExactSize(size); ok { - return spec.aspectRatio - } - return "auto" -} - -func sosanaAspectRatio(explicit, ratio, size string) string { - if value := strings.TrimSpace(explicit); value != "" { - return value - } - if value := strings.TrimSpace(ratio); value != "" { - return value - } - return sosanaSizeToAspectRatio(size) -} - -func sosanaImageSize(explicit, size string) (string, error) { - if strings.TrimSpace(explicit) != "" { - if value, ok := normalizeSosanaImageSize(explicit); ok { - return value, nil - } - return "", fmt.Errorf("image_size is unsupported") - } - if value, ok := sosanaImageSizeFromExactSize(size); ok { - return value, nil - } - return "", fmt.Errorf("size is unsupported") -} - -func normalizeSosanaImageSize(raw string) (string, bool) { - switch strings.ToLower(strings.TrimSpace(raw)) { - case "", "auto", "1k": - return "1K", true - case "2k": - return "2K", true - case "4k": - return "4K", true - default: - return "", false - } -} - -func sosanaImageSizeFromExactSize(size string) (string, bool) { - if spec, ok := sosanaImageSpecFromExactSize(size); ok { - return spec.imageSize, true - } - return "", false -} - -type sosanaExactImageSpec struct { - imageSize string - aspectRatio string -} - -func sosanaImageSpecFromExactSize(size string) (sosanaExactImageSpec, bool) { - switch strings.TrimSpace(size) { - case "", "auto": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "auto"}, true - - case "1024x1024": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "1:1"}, true - case "512x2048": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "1:4"}, true - case "384x3072": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "1:8"}, true - case "848x1264": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "2:3"}, true - case "1264x848": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "3:2"}, true - case "896x1200": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "3:4"}, true - case "2048x512": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "4:1"}, true - case "1200x896": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "4:3"}, true - case "928x1152": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "4:5"}, true - case "1152x928": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "5:4"}, true - case "3072x384": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "8:1"}, true - case "768x1376": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "9:16"}, true - case "1376x768": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "16:9"}, true - case "1584x672": - return sosanaExactImageSpec{imageSize: "1K", aspectRatio: "21:9"}, true - - case "2048x2048": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "1:1"}, true - case "1024x4096": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "1:4"}, true - case "768x6144": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "1:8"}, true - case "1696x2528": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "2:3"}, true - case "2528x1696": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "3:2"}, true - case "1792x2400": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "3:4"}, true - case "4096x1024": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "4:1"}, true - case "2400x1792": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "4:3"}, true - case "1856x2304": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "4:5"}, true - case "2304x1856": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "5:4"}, true - case "6144x768": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "8:1"}, true - case "1536x2752": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "9:16"}, true - case "2752x1536": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "16:9"}, true - case "3168x1344": - return sosanaExactImageSpec{imageSize: "2K", aspectRatio: "21:9"}, true - - case "4096x4096": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "1:1"}, true - case "2048x8192": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "1:4"}, true - case "1536x12288": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "1:8"}, true - case "3392x5056": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "2:3"}, true - case "5056x3392": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "3:2"}, true - case "3584x4800": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "3:4"}, true - case "8192x2048": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "4:1"}, true - case "4800x3584": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "4:3"}, true - case "3712x4608": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "4:5"}, true - case "4608x3712": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "5:4"}, true - case "12288x1536": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "8:1"}, true - case "3072x5504": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "9:16"}, true - case "5504x3072": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "16:9"}, true - case "6336x2688": - return sosanaExactImageSpec{imageSize: "4K", aspectRatio: "21:9"}, true - default: - return sosanaExactImageSpec{}, false - } -} - -func validateSosanaImageCount(n *int) error { - if n == nil || *n == 1 { - return nil - } - return fmt.Errorf("image requests support n=1 only") -} - -func validateSosanaResponseFormat(format string) error { - if strings.EqualFold(strings.TrimSpace(format), "url") { - return fmt.Errorf("response_format=url is unsupported for this image model") - } - return nil -} - -func validateSosanaImageCountString(raw string) error { - raw = strings.TrimSpace(raw) - if raw == "" { - return nil - } - n, err := strconv.Atoi(raw) - if err != nil { - return fmt.Errorf("invalid image count: %w", err) - } - if n != 1 { - return fmt.Errorf("image requests support n=1 only") - } - return nil -} - -func readLimitedSosanaMultipartPart(r io.Reader, limit int64) ([]byte, error) { - data, err := io.ReadAll(io.LimitReader(r, limit+1)) - if err != nil { - return nil, fmt.Errorf("failed to read multipart part: %w", err) - } - if int64(len(data)) > limit { - return nil, fmt.Errorf("multipart image exceeds %d bytes", limit) - } - return data, nil -} - -func detectSosanaImageMIMEType(header string, data []byte) string { - header = strings.ToLower(strings.TrimSpace(header)) - if strings.HasPrefix(header, "image/") { - mediaType, _, err := mime.ParseMediaType(header) - if err == nil { - return mediaType - } - return strings.TrimSpace(strings.Split(header, ";")[0]) - } - detected := http.DetectContentType(data) - if strings.HasPrefix(detected, "image/") { - return detected - } - return "application/octet-stream" -} diff --git a/internal/proxy/sosana_routing.go b/internal/proxy/sosana_routing.go new file mode 100644 index 00000000..69900690 --- /dev/null +++ b/internal/proxy/sosana_routing.go @@ -0,0 +1,158 @@ +package proxy + +import ( + "net/http" + "time" + + "github.com/mixaill76/auto_ai_router/internal/config" + "github.com/mixaill76/auto_ai_router/internal/converter/sosana" + "github.com/mixaill76/auto_ai_router/internal/scope" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" +) + +const unsupportedImageProviderRequestMessage = "request parameters are not supported by available image providers" +const unsupportedProviderEndpointMessage = "request endpoint is not supported by available providers" + +func (p *Proxy) applySosanaCompatibilityRouting( + w http.ResponseWriter, + r *http.Request, + prepared *orchestratedRequest, + modelID string, + cred **config.CredentialConfig, + body *[]byte, + proxyBody *[]byte, + realModelID *string, + isImageGeneration bool, + isImageEdit bool, + logCtx *RequestLogContext, + start time.Time, +) bool { + if (*cred).Type != config.ProviderTypeSosana || (!isImageGeneration && !isImageEdit) { + return true + } + + reason := sosana.UnsupportedModel(*realModelID) + if reason == "" { + reason = sosana.UnsupportedRequest(r.URL.Path, *body, r.Header.Get("Content-Type")) + } + if reason == "" { + return true + } + + nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, logCtx.Scope, reason) + if routed { + *cred = nextCred + *body = nextReq.body + *proxyBody = nextReq.proxyBody + *realModelID = nextReq.realModelID + r.URL.Path = nextReq.path + prepared.body = nextReq.body + prepared.proxyBody = nextReq.proxyBody + prepared.proxyPath = nextReq.proxyPath + prepared.realModelID = nextReq.realModelID + prepared.convertedResp = nextReq.convertedResp + prepared.passthroughResponses = nextReq.passthroughResponses + prepared.nativeResponses = nextReq.nativeResponses + logCtx.RealModelID = *realModelID + if span := trace.SpanFromContext(r.Context()); span.IsRecording() { + span.SetAttributes( + attribute.String("aar.real_model", *realModelID), + attribute.String("aar.credential", nextCred.Name), + attribute.String("aar.provider", string(nextCred.Type)), + attribute.Bool("aar.provider_compatibility_skip", true), + ) + } + return true + } + + success, fallbackReason := p.TryFallbackProxy( + w, + requestWithPath(r, prepared.proxyPath), + modelID, + (*cred).Name, + http.StatusBadRequest, + RetryReasonServerErr, + *proxyBody, + start, + logCtx, + ) + if success { + return false + } + p.logger.DebugContext(r.Context(), "No fallback handled unsupported image provider request", + "credential", (*cred).Name, + "model", modelID, + "reason", reason, + "fallback_reason", fallbackReason) + logCtx.Credential = *cred + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = unsupportedImageProviderRequestMessage + WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) + return false +} + +func (p *Proxy) nextPrimaryAfterUnsupportedSosana( + r *http.Request, + prepared *orchestratedRequest, + modelID string, + currentCred *config.CredentialConfig, + visibility scope.Context, + reason string, +) (*config.CredentialConfig, credentialPreparedRequest, bool) { + triedCreds := GetTried(r.Context()) + triedCreds[currentCred.Name] = true + + for attempts := 0; attempts < 128; attempts++ { + candidate, err := p.balancer.NextForModelExcludingScoped(modelID, triedCreds, visibility) + if err != nil { + p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", + "model", modelID, + "credential", currentCred.Name, + "reason", reason, + "error", err) + return nil, credentialPreparedRequest{}, false + } + triedCreds[candidate.Name] = true + if candidate.Type == config.ProviderTypeSosana { + continue + } + + nextReq, prepErr := p.prepareRequestForCredential( + r, + prepared.baseBody, + prepared.baseProxyBody, + modelID, + prepared.baseRealModelID, + prepared.basePath, + prepared.streaming, + candidate, + prepared.isResponsesAPI, + prepared.responsesPrevHandled, + prepared.stickyCacheEligible, + ) + if prepErr != nil { + p.logger.WarnContext(r.Context(), "Failed to prepare alternate primary request after image compatibility skip", + "credential", candidate.Name, + "provider", string(candidate.Type), + "model", modelID, + "reason", reason, + "error", prepErr) + continue + } + + p.logger.InfoContext(r.Context(), "Skipping incompatible image credential for unsupported image request", + "credential", currentCred.Name, + "next_credential", candidate.Name, + "model", modelID, + "reason", reason) + return candidate, nextReq, true + } + + p.logger.WarnContext(r.Context(), "Image compatibility skip exhausted primary credential scan", + "credential", currentCred.Name, + "model", modelID, + "reason", reason) + return nil, credentialPreparedRequest{}, false +} diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 859c808b..8afc2d75 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -19,6 +19,7 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/mixaill76/auto_ai_router/internal/converter/sosana" "github.com/mixaill76/auto_ai_router/internal/litellmdb" litellmmodels "github.com/mixaill76/auto_ai_router/internal/litellmdb/models" aimodels "github.com/mixaill76/auto_ai_router/internal/models" @@ -879,8 +880,8 @@ func TestDownloadSosanaResultImageRejectsUnsafeProductionURL(t *testing.T) { resultURL := imageServer.URL + "/private.png" prx := newSosanaTestProxy("https://sosana.art", &logBuf) cred := &config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: "https://sosana.art"} - image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "banana-2-1k-compliant", sosanaBananaTaskResponse{ - Status: sosanaStatusCompleted, + image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "banana-2-1k-compliant", sosana.BananaTaskResponse{ + Status: sosana.StatusCompleted, ResultFileURL: &resultURL, }, nil, &RequestLogContext{}) @@ -1090,256 +1091,6 @@ func TestProxyRequest_SosanaTimeoutMasked(t *testing.T) { } } -func TestSosanaImageGenerationRequest(t *testing.T) { - tests := []struct { - name string - size string - wantAspect string - wantSize string - }{ - {name: "one k square", size: "1024x1024", wantAspect: "1:1", wantSize: "1K"}, - {name: "one k wide", size: "1376x768", wantAspect: "16:9", wantSize: "1K"}, - {name: "one k portrait", size: "768x1376", wantAspect: "9:16", wantSize: "1K"}, - {name: "one k tall", size: "512x2048", wantAspect: "1:4", wantSize: "1K"}, - {name: "two k square", size: "2048x2048", wantAspect: "1:1", wantSize: "2K"}, - {name: "two k wide", size: "2752x1536", wantAspect: "16:9", wantSize: "2K"}, - {name: "four k square", size: "4096x4096", wantAspect: "1:1", wantSize: "4K"}, - {name: "four k ultra wide", size: "6336x2688", wantAspect: "21:9", wantSize: "4K"}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - body := []byte(`{"model":"banana-2-1k-compliant","prompt":"draw a cat","size":"` + tt.size + `","n":1}`) - got, concreteModel, err := buildSosanaImageGenerationRequest(body, "banana-2-{image_size}-compliant") - require.NoError(t, err) - - var req sosanaBananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "draw a cat", req.Prompt) - assert.Equal(t, "banana-2-"+strings.ToLower(tt.wantSize)+"-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) - assert.Equal(t, tt.wantAspect, req.AspectRatio) - assert.Equal(t, tt.wantSize, req.ImageSize) - assert.False(t, req.PromptOptimization) - assert.Empty(t, req.ImageURLs) - }) - } -} - -func TestSosanaImageGenerationRequestPrefersProviderModel(t *testing.T) { - got, concreteModel, err := buildSosanaImageGenerationRequest([]byte(`{"model":"public-image","prompt":"draw","n":1}`), "banana-2-1k-compliant") - require.NoError(t, err) - - var req sosanaBananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "banana-2-1k-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) -} - -func TestSosanaImageGenerationRequestMapsPublicGeminiModelToSosanaTier(t *testing.T) { - got, concreteModel, err := buildSosanaImageGenerationRequest( - []byte(`{"model":"google/gemini-3.1-flash-image-preview","prompt":"draw","image_size":"2K","n":1}`), - "google/gemini-3.1-flash-image-preview", - ) - require.NoError(t, err) - - var req sosanaBananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "banana-2-2k-compliant", req.Model) - assert.Equal(t, "banana-2-2k-compliant", concreteModel) - assert.Equal(t, "2K", req.ImageSize) -} - -func TestSosanaImageGenerationRequestUsesExplicitAspectRatio(t *testing.T) { - got, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","size":"1024x1024","aspect_ratio":"16:9"}`), "banana-2-1k-compliant") - require.NoError(t, err) - - var req sosanaBananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "16:9", req.AspectRatio) -} - -func TestSosanaImageGenerationRequestRejectsMultipleImages(t *testing.T) { - _, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","n":2}`), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "n=1") -} - -func TestSosanaImageGenerationRequestRejectsURLResponseFormat(t *testing.T) { - _, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","response_format":"url"}`), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "response_format=url") -} - -func TestSosanaImageGenerationRequestRejectsUnsupportedControls(t *testing.T) { - tests := []struct { - name string - body string - want string - }{ - {name: "tools", body: `{"model":"banana-2-1k-compliant","prompt":"draw","tools":[{"type":"google_search"}]}`, want: "tools"}, - {name: "thinking", body: `{"model":"banana-2-1k-compliant","prompt":"draw","thinking_level":"high"}`, want: "thinking_level"}, - {name: "output format", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"jpeg"}`, want: "output_format"}, - {name: "output compression", body: `{"model":"banana-2-1k-compliant","prompt":"draw","output_compression":0}`, want: "output_compression"}, - {name: "quality auto", body: `{"model":"banana-2-1k-compliant","prompt":"draw","quality":"auto"}`, want: "quality"}, - {name: "messages", body: `{"model":"banana-2-1k-compliant","prompt":"draw","messages":[{"role":"user","content":"draw"}]}`, want: "messages"}, - {name: "image size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_size":"0.5K"}`, want: "image_size"}, - {name: "exact size 0.5k", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"512x512"}`, want: "size"}, - {name: "legacy openai size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"1792x1024"}`, want: "size"}, - {name: "unknown size", body: `{"model":"banana-2-1k-compliant","prompt":"draw","size":"333x777"}`, want: "size"}, - {name: "reference images", body: `{"model":"banana-2-1k-compliant","prompt":"draw","image_urls":["https://example.com/a.png"]}`, want: "image_urls"}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, _, err := buildSosanaImageGenerationRequest([]byte(tt.body), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), tt.want) - }) - } -} - -func TestSosanaImageGenerationRequestAllowsPNGOutputFormat(t *testing.T) { - _, _, err := buildSosanaImageGenerationRequest([]byte(`{"model":"banana-2-1k-compliant","prompt":"draw","output_format":"png","response_format":"b64_json"}`), "banana-2-1k-compliant") - require.NoError(t, err) -} - -func TestSosanaImageEditRequest(t *testing.T) { - body, contentType := sosanaMultipartEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - "size": "1024x1024", - "n": "1", - }, map[string][]byte{ - "image": sosanaResultPNG, - }) - - got, concreteModel, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.NoError(t, err) - - var req sosanaBananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "make it blue", req.Prompt) - assert.Equal(t, "banana-2-1k-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) - assert.Equal(t, "1:1", req.AspectRatio) - assert.Equal(t, "1K", req.ImageSize) - assert.False(t, req.PromptOptimization) - require.Len(t, req.ImageURLs, 1) - assert.True(t, strings.HasPrefix(req.ImageURLs[0], "data:image/png;base64,")) - assert.Contains(t, req.ImageURLs[0], base64.StdEncoding.EncodeToString(sosanaResultPNG)) -} - -func TestSosanaImageEditRequestPrefersProviderModel(t *testing.T) { - body, contentType := sosanaMultipartEditBody(t, map[string]string{ - "model": "public-image", - "prompt": "make it blue", - }, map[string][]byte{ - "image": sosanaResultPNG, - }) - - got, concreteModel, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.NoError(t, err) - - var req sosanaBananaCreateRequest - require.NoError(t, json.Unmarshal(got, &req)) - assert.Equal(t, "banana-2-1k-compliant", req.Model) - assert.Equal(t, req.Model, concreteModel) -} - -func TestSosanaImageEditRequestRejectsMask(t *testing.T) { - body, contentType := sosanaMultipartEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - }, map[string][]byte{ - "image": sosanaResultPNG, - "mask": sosanaResultPNG, - }) - - _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "mask") -} - -func TestSosanaImageEditRequestRejectsMultipleImagesCount(t *testing.T) { - body, contentType := sosanaMultipartEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - "n": "2", - }, map[string][]byte{ - "image": sosanaResultPNG, - }) - - _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "n=1") -} - -func TestSosanaImageEditRequestRejectsURLResponseFormat(t *testing.T) { - body, contentType := sosanaMultipartEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - "response_format": "url", - }, map[string][]byte{ - "image": sosanaResultPNG, - }) - - _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "response_format=url") -} - -func TestSosanaImageEditRequestRejectsJPEGInput(t *testing.T) { - body, contentType := sosanaMultipartEditBody(t, map[string]string{ - "model": "banana-2-1k-compliant", - "prompt": "make it blue", - }, map[string][]byte{ - "image": {0xff, 0xd8, 0xff, 0xdb, 0, 0x43, 0, 1, 2, 3}, - }) - - _, _, err := buildSosanaImageEditRequest(body, contentType, "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "PNG") -} - -func TestSosanaImageEditRequestRejectsTooManyImages(t *testing.T) { - var buf bytes.Buffer - writer := multipart.NewWriter(&buf) - require.NoError(t, writer.WriteField("model", "banana-2-1k-compliant")) - require.NoError(t, writer.WriteField("prompt", "make it blue")) - for i := 0; i < maxSosanaInputImages+1; i++ { - part, err := writer.CreateFormFile("image", fmt.Sprintf("image-%02d.png", i)) - require.NoError(t, err) - _, err = part.Write(sosanaResultPNG) - require.NoError(t, err) - } - require.NoError(t, writer.Close()) - - _, _, err := buildSosanaImageEditRequest(buf.Bytes(), writer.FormDataContentType(), "banana-2-1k-compliant") - require.Error(t, err) - assert.Contains(t, err.Error(), "too many") -} - -func TestSosanaOpenAIImageResponse(t *testing.T) { - createdAt := "2026-01-01T00:00:00Z" - body, err := buildSosanaOpenAIImageResponse(sosanaBananaTaskResponse{ - Status: sosanaStatusCompleted, - CreatedAt: createdAt, - OptimizedPrompt: "A detailed result prompt", - }, sosanaResultPNG) - require.NoError(t, err) - - var resp openai.OpenAIImageResponse - require.NoError(t, json.Unmarshal(body, &resp)) - require.Len(t, resp.Data, 1) - assert.Empty(t, resp.Data[0].URL) - assert.Equal(t, "A detailed result prompt", resp.Data[0].RevisedPrompt) - assert.Equal(t, base64.StdEncoding.EncodeToString(sosanaResultPNG), resp.Data[0].B64JSON) - ts, err := time.Parse(time.RFC3339, createdAt) - require.NoError(t, err) - assert.Equal(t, ts.Unix(), resp.Created) -} - var sosanaResultPNG = []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0} func newSosanaResultImageServer(t *testing.T, status int, contentType string, body []byte, auths *[]string) *httptest.Server { From b4e47ab12c9c63b18e536d28e77da42ec429f436 Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Thu, 23 Jul 2026 20:22:23 +0300 Subject: [PATCH 15/16] refactor: keep sosana client logic in converter --- internal/converter/sosana/client.go | 397 ++++++++++++++++++ internal/converter/sosana/client_test.go | 91 ++++ internal/converter/sosana/live_test.go | 82 ++++ internal/proxy/sosana.go | 507 +++++++++-------------- internal/proxy/sosana_live_test.go | 67 --- internal/proxy/sosana_routing.go | 158 ------- internal/proxy/sosana_test.go | 68 +-- 7 files changed, 765 insertions(+), 605 deletions(-) create mode 100644 internal/converter/sosana/client.go create mode 100644 internal/converter/sosana/client_test.go create mode 100644 internal/converter/sosana/live_test.go delete mode 100644 internal/proxy/sosana_live_test.go delete mode 100644 internal/proxy/sosana_routing.go diff --git a/internal/converter/sosana/client.go b/internal/converter/sosana/client.go new file mode 100644 index 00000000..8711d41d --- /dev/null +++ b/internal/converter/sosana/client.go @@ -0,0 +1,397 @@ +package sosana + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net" + "net/http" + "net/url" + "strings" + "time" +) + +const ( + PollInterval = 2 * time.Second + MaxResultImageBytes int64 = 32 * 1024 * 1024 + MaxResultErrorBytes int64 = 16 * 1024 +) + +var allowPrivateResultURLForTests func(*url.URL) bool + +type TaskHTTPResult struct { + Task BananaTaskResponse + RawBody []byte + StatusCode int +} + +type ResultImage struct { + Bytes []byte + ContentType string + Host string +} + +type ResultImageError struct { + StatusCode int + ResponseBody []byte + Host string + ContentType string + SniffedContentType string + UpstreamStatus int + Err error +} + +func (e *ResultImageError) Error() string { + if e.Err != nil { + return e.Err.Error() + } + return "sosana result image download failed" +} + +func (e *ResultImageError) Unwrap() error { + return e.Err +} + +func SetAllowPrivateResultURLForTests(fn func(*url.URL) bool) func() { + previous := allowPrivateResultURLForTests + allowPrivateResultURLForTests = fn + return func() { + allowPrivateResultURLForTests = previous + } +} + +func DoTaskRequest(ctx context.Context, client *http.Client, method, url string, apiKey string, body []byte) (TaskHTTPResult, error) { + var reader *bytes.Reader + if body != nil { + reader = bytes.NewReader(body) + } else { + reader = bytes.NewReader(nil) + } + req, err := http.NewRequestWithContext(ctx, method, url, reader) + if err != nil { + return TaskHTTPResult{StatusCode: http.StatusInternalServerError}, err + } + req.Header.Set("Authorization", "Bearer "+apiKey) + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + + if client == nil { + client = http.DefaultClient + } + resp, err := client.Do(req) + if err != nil { + return TaskHTTPResult{StatusCode: http.StatusBadGateway}, err + } + defer func() { + _ = resp.Body.Close() + }() + + rawBody, err := readLimitedResultBody(resp.Body, MaxResultImageBytes) + if err != nil { + return TaskHTTPResult{StatusCode: http.StatusBadGateway}, err + } + var task BananaTaskResponse + if len(rawBody) > 0 { + _ = json.Unmarshal(rawBody, &task) + } + return TaskHTTPResult{Task: task, RawBody: rawBody, StatusCode: resp.StatusCode}, nil +} + +func DownloadResultImage(ctx context.Context, client *http.Client, task BananaTaskResponse) (ResultImage, error) { + resultURL := "" + if task.ResultFileURL != nil { + resultURL = strings.TrimSpace(*task.ResultFileURL) + } + parsed, err := parseResultURL(resultURL) + if err != nil { + return ResultImage{}, resultImageError(http.StatusBadGateway, "", err) + } + host := parsed.Hostname() + if err := validateResultURL(ctx, parsed); err != nil { + return ResultImage{}, resultImageError(http.StatusBadGateway, host, err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, resultURL, nil) + if err != nil { + return ResultImage{}, resultImageError(http.StatusBadGateway, host, err) + } + req.Header.Set("Accept", "image/*") + + resp, err := doResultImageRequest(client, req) + if err != nil { + statusCode := http.StatusBadGateway + if isResultTimeout(ctx, err) { + statusCode = http.StatusRequestTimeout + } + return ResultImage{}, resultImageError(statusCode, host, err) + } + defer func() { + _ = resp.Body.Close() + }() + + contentType := strings.TrimSpace(resp.Header.Get("Content-Type")) + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return ResultImage{}, &ResultImageError{ + StatusCode: http.StatusBadGateway, + ResponseBody: readTextResultBody(resp.Body, contentType), + Host: host, + ContentType: contentType, + UpstreamStatus: resp.StatusCode, + Err: errors.New("sosana result image download returned error status"), + } + } + + image, err := readLimitedResultImage(resp.Body) + if err != nil { + return ResultImage{}, resultImageError(http.StatusBadGateway, host, err) + } + sniffedType := http.DetectContentType(image) + if !IsPNGContentType(contentType) && !IsPNGContentType(sniffedType) { + return ResultImage{}, &ResultImageError{ + StatusCode: http.StatusBadGateway, + ResponseBody: textResultBodyPrefix(image, contentType, sniffedType), + Host: host, + ContentType: contentType, + SniffedContentType: sniffedType, + Err: errors.New("sosana result URL returned non-PNG content"), + } + } + if !IsPNGContentType(contentType) { + contentType = sniffedType + } + return ResultImage{Bytes: image, ContentType: contentType, Host: host}, nil +} + +func ResultHost(task BananaTaskResponse) string { + if task.ResultFileURL == nil { + return "" + } + parsed, err := url.Parse(strings.TrimSpace(*task.ResultFileURL)) + if err != nil { + return "" + } + return parsed.Hostname() +} + +func IsUnsafeResultIP(ip net.IP) bool { + if ip == nil { + return true + } + if v4 := ip.To4(); v4 != nil && v4[0] == 100 && v4[1]&0xc0 == 64 { + return true + } + return ip.IsLoopback() || + ip.IsPrivate() || + ip.IsLinkLocalUnicast() || + ip.IsLinkLocalMulticast() || + ip.IsMulticast() || + ip.IsUnspecified() +} + +func IsPNGContentType(contentType string) bool { + contentType = strings.ToLower(strings.TrimSpace(contentType)) + return contentType == "image/png" || strings.HasPrefix(contentType, "image/png;") +} + +func isTextContentType(contentType string) bool { + contentType = strings.ToLower(strings.TrimSpace(contentType)) + return strings.HasPrefix(contentType, "text/") || + strings.Contains(contentType, "json") || + strings.Contains(contentType, "xml") +} + +func resultImageError(statusCode int, host string, err error) *ResultImageError { + return &ResultImageError{StatusCode: statusCode, Host: host, Err: err} +} + +func doResultImageRequest(client *http.Client, req *http.Request) (*http.Response, error) { + if client == nil { + client = http.DefaultClient + } + resultClient := *client + resultClient.Transport = resultImageTransport() + resultClient.CheckRedirect = func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + } + return resultClient.Do(req) +} + +func resultImageTransport() http.RoundTripper { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.Proxy = nil + transport.DisableKeepAlives = true + transport.DialContext = dialResultAddress + return transport +} + +func dialResultAddress(ctx context.Context, network, address string) (net.Conn, error) { + host, port, err := net.SplitHostPort(address) + if err != nil { + return nil, err + } + + dialer := &net.Dialer{} + if allowPrivateResultHostForTests(host) { + return dialer.DialContext(ctx, network, address) + } + if ip := net.ParseIP(host); ip != nil { + if IsUnsafeResultIP(ip) { + return nil, errors.New("sosana result_file_url resolves to a private address") + } + return dialer.DialContext(ctx, network, address) + } + + addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) + if err != nil { + return nil, err + } + if len(addrs) == 0 { + return nil, errors.New("sosana result_file_url host has no addresses") + } + for _, addr := range addrs { + if IsUnsafeResultIP(addr.IP) { + return nil, errors.New("sosana result_file_url resolves to a private address") + } + } + + var dialErr error + for _, addr := range addrs { + conn, err := dialer.DialContext(ctx, network, net.JoinHostPort(addr.IP.String(), port)) + if err == nil { + return conn, nil + } + dialErr = err + } + if dialErr != nil { + return nil, dialErr + } + return nil, errors.New("sosana result_file_url host has no dialable addresses") +} + +func allowPrivateResultHostForTests(host string) bool { + if allowPrivateResultURLForTests == nil { + return false + } + return allowPrivateResultURLForTests(&url.URL{Scheme: "http", Host: host}) +} + +func parseResultURL(raw string) (*url.URL, error) { + if raw == "" { + return nil, errors.New("sosana task completed without result_file_url") + } + parsed, err := url.Parse(raw) + if err != nil { + return nil, err + } + if parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return nil, errors.New("sosana result_file_url must be an http or https URL") + } + return parsed, nil +} + +func validateResultURL(ctx context.Context, parsed *url.URL) error { + if allowPrivateResultURLForTests != nil && allowPrivateResultURLForTests(parsed) { + return nil + } + if parsed.Scheme != "https" { + return errors.New("sosana result_file_url must use https") + } + + host := parsed.Hostname() + if strings.EqualFold(host, "localhost") { + return errors.New("sosana result_file_url host is not allowed") + } + if !isAllowedResultHost(host) { + return errors.New("sosana result_file_url host is not allowed") + } + if ip := net.ParseIP(host); ip != nil { + if IsUnsafeResultIP(ip) { + return errors.New("sosana result_file_url resolves to a private address") + } + return nil + } + + addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) + if err != nil { + return err + } + if len(addrs) == 0 { + return errors.New("sosana result_file_url host has no addresses") + } + for _, addr := range addrs { + if IsUnsafeResultIP(addr.IP) { + return errors.New("sosana result_file_url resolves to a private address") + } + } + return nil +} + +func isAllowedResultHost(host string) bool { + host = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(host)), ".") + for _, suffix := range []string{ + "sosana.blog", + "sosana.art", + "storage.yandexcloud.net", + } { + if host == suffix || strings.HasSuffix(host, "."+suffix) { + return true + } + } + return false +} + +func readLimitedResultImage(body io.Reader) ([]byte, error) { + data, err := readLimitedResultBody(body, MaxResultImageBytes) + if err != nil { + return nil, err + } + if len(data) == 0 { + return nil, errors.New("sosana result image body is empty") + } + return data, nil +} + +func readLimitedResultBody(body io.Reader, limit int64) ([]byte, error) { + data, err := io.ReadAll(io.LimitReader(body, limit+1)) + if err != nil { + return nil, err + } + if int64(len(data)) > limit { + return nil, errors.New("sosana response body is too large") + } + return data, nil +} + +func readTextResultBody(body io.Reader, contentType string) []byte { + if !isTextContentType(contentType) { + return nil + } + data, err := io.ReadAll(io.LimitReader(body, MaxResultErrorBytes)) + if err != nil { + return nil + } + return data +} + +func textResultBodyPrefix(body []byte, contentTypes ...string) []byte { + for _, contentType := range contentTypes { + if isTextContentType(contentType) { + if int64(len(body)) > MaxResultErrorBytes { + return body[:MaxResultErrorBytes] + } + return body + } + } + return nil +} + +func isResultTimeout(ctx context.Context, err error) bool { + if errors.Is(ctx.Err(), context.DeadlineExceeded) { + return true + } + var netErr net.Error + return errors.As(err, &netErr) && netErr.Timeout() +} diff --git a/internal/converter/sosana/client_test.go b/internal/converter/sosana/client_test.go new file mode 100644 index 00000000..5488354c --- /dev/null +++ b/internal/converter/sosana/client_test.go @@ -0,0 +1,91 @@ +package sosana + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateResultURLBlocksUnsafeHosts(t *testing.T) { + tests := []string{ + "http://main-r2.sosana.blog/image.png", + "https://localhost/image.png", + "https://127.0.0.1/image.png", + "https://169.254.169.254/latest/meta-data", + "https://100.64.0.1/image.png", + "https://example.com/image.png", + } + + for _, rawURL := range tests { + t.Run(rawURL, func(t *testing.T) { + parsed, err := parseResultURL(rawURL) + require.NoError(t, err) + require.Error(t, validateResultURL(context.Background(), parsed)) + }) + } +} + +func TestValidateResultURLAllowsLocalOnlyWithTestHook(t *testing.T) { + restore := SetAllowPrivateResultURLForTests(func(parsed *url.URL) bool { + return parsed.Scheme == "http" && parsed.Hostname() == "127.0.0.1" + }) + t.Cleanup(restore) + + parsed, err := parseResultURL("http://127.0.0.1/image.png") + require.NoError(t, err) + require.NoError(t, validateResultURL(context.Background(), parsed)) +} + +func TestDialResultAddressRejectsPrivateIP(t *testing.T) { + conn, err := dialResultAddress(context.Background(), "tcp", net.JoinHostPort("127.0.0.1", "443")) + + require.Error(t, err) + assert.Nil(t, conn) + assert.Contains(t, err.Error(), "private address") +} + +func TestDownloadResultImageRejectsUnsafeProductionURL(t *testing.T) { + called := false + imageServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + })) + defer imageServer.Close() + + resultURL := imageServer.URL + "/private.png" + image, err := DownloadResultImage(context.Background(), http.DefaultClient, BananaTaskResponse{ + Status: StatusCompleted, + ResultFileURL: &resultURL, + }) + + require.Error(t, err) + assert.Empty(t, image.Bytes) + assert.False(t, called) + var imageErr *ResultImageError + require.ErrorAs(t, err, &imageErr) + assert.Equal(t, http.StatusBadGateway, imageErr.StatusCode) + assert.Contains(t, err.Error(), "host is not allowed") +} + +func AllowPrivateResultURLsForTest(t *testing.T) { + t.Helper() + + restore := SetAllowPrivateResultURLForTests(func(parsed *url.URL) bool { + if parsed.Scheme != "http" { + return false + } + host := parsed.Hostname() + if strings.EqualFold(host, "localhost") { + return true + } + ip := net.ParseIP(host) + return ip != nil && IsUnsafeResultIP(ip) + }) + t.Cleanup(restore) +} diff --git a/internal/converter/sosana/live_test.go b/internal/converter/sosana/live_test.go new file mode 100644 index 00000000..58cefcf0 --- /dev/null +++ b/internal/converter/sosana/live_test.go @@ -0,0 +1,82 @@ +package sosana + +import ( + "context" + "encoding/base64" + "encoding/json" + "net/http" + "os" + "testing" + "time" + + "github.com/mixaill76/auto_ai_router/internal/converter/openai" + "github.com/stretchr/testify/require" +) + +func TestSosanaLiveAcceptance(t *testing.T) { + if os.Getenv("SOSANA_ACCEPTANCE") != "1" { + t.Skip("SOSANA_ACCEPTANCE=1 not set, skipping paid Sosana live acceptance test") + } + + apiKey := os.Getenv("SOSANA_API_KEY") + require.NotEmpty(t, apiKey, "SOSANA_API_KEY is required for Sosana live acceptance test") + + baseURL := os.Getenv("SOSANA_BASE_URL") + if baseURL == "" { + baseURL = "https://sosana.art" + } + model := os.Getenv("SOSANA_MODEL") + if model == "" { + model = "banana-2-1k-compliant" + } + prompt := os.Getenv("SOSANA_PROMPT") + if prompt == "" { + prompt = "a small blue cube on a white background" + } + + openAIBody, err := json.Marshal(map[string]any{ + "model": model, + "prompt": prompt, + "size": "1024x1024", + "n": 1, + }) + require.NoError(t, err) + + createBody, _, err := ImageGenerationRequest(openAIBody, model) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + taskResult, err := DoTaskRequest(ctx, http.DefaultClient, http.MethodPost, CreateURL(baseURL), apiKey, createBody) + require.NoError(t, err) + require.Less(t, taskResult.StatusCode, http.StatusBadRequest, "response body: %s", string(taskResult.RawBody)) + + task := taskResult.Task + for task.Status == StatusProcessing { + select { + case <-ctx.Done(): + t.Fatal(ctx.Err()) + case <-time.After(time.Duration(PollInterval)): + } + taskResult, err = DoTaskRequest(ctx, http.DefaultClient, http.MethodGet, PollURL(baseURL, task.UID), apiKey, nil) + require.NoError(t, err) + require.Less(t, taskResult.StatusCode, http.StatusBadRequest, "response body: %s", string(taskResult.RawBody)) + task = taskResult.Task + } + require.Equal(t, StatusCompleted, task.Status, "response body: %s", string(taskResult.RawBody)) + + image, err := DownloadResultImage(ctx, http.DefaultClient, task) + require.NoError(t, err) + require.NotEmpty(t, image.Bytes) + + body, err := OpenAIImageResponse(task, image.Bytes) + require.NoError(t, err) + var resp openai.OpenAIImageResponse + require.NoError(t, json.Unmarshal(body, &resp)) + require.Len(t, resp.Data, 1) + require.Empty(t, resp.Data[0].URL) + require.NotEmpty(t, resp.Data[0].B64JSON) + _, err = base64.StdEncoding.DecodeString(resp.Data[0].B64JSON) + require.NoError(t, err) +} diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go index 528fdcf9..9157268b 100644 --- a/internal/proxy/sosana.go +++ b/internal/proxy/sosana.go @@ -1,30 +1,27 @@ package proxy import ( - "bytes" "context" - "encoding/json" "errors" - "io" "math/rand" - "net" "net/http" - "net/url" "strings" "time" "github.com/mixaill76/auto_ai_router/internal/config" "github.com/mixaill76/auto_ai_router/internal/converter" "github.com/mixaill76/auto_ai_router/internal/converter/sosana" + "github.com/mixaill76/auto_ai_router/internal/scope" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" ) const ( - sosanaPollInterval = 2 * time.Second - maxSosanaResultImageBytes int64 = 32 * 1024 * 1024 - maxSosanaResultErrorBytes int64 = 16 * 1024 + sosanaPollInterval = time.Duration(sosana.PollInterval) ) -var allowPrivateSosanaResultURLForTests func(*url.URL) bool +const unsupportedImageProviderRequestMessage = "request parameters are not supported by available image providers" +const unsupportedProviderEndpointMessage = "request endpoint is not supported by available providers" type sosanaAttemptResult struct { body []byte @@ -33,6 +30,149 @@ type sosanaAttemptResult struct { retryReason RetryReason } +func (p *Proxy) applySosanaCompatibilityRouting( + w http.ResponseWriter, + r *http.Request, + prepared *orchestratedRequest, + modelID string, + cred **config.CredentialConfig, + body *[]byte, + proxyBody *[]byte, + realModelID *string, + isImageGeneration bool, + isImageEdit bool, + logCtx *RequestLogContext, + start time.Time, +) bool { + if (*cred).Type != config.ProviderTypeSosana || (!isImageGeneration && !isImageEdit) { + return true + } + + reason := sosana.UnsupportedModel(*realModelID) + if reason == "" { + reason = sosana.UnsupportedRequest(r.URL.Path, *body, r.Header.Get("Content-Type")) + } + if reason == "" { + return true + } + + nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, logCtx.Scope, reason) + if routed { + *cred = nextCred + *body = nextReq.body + *proxyBody = nextReq.proxyBody + *realModelID = nextReq.realModelID + r.URL.Path = nextReq.path + prepared.body = nextReq.body + prepared.proxyBody = nextReq.proxyBody + prepared.proxyPath = nextReq.proxyPath + prepared.realModelID = nextReq.realModelID + prepared.convertedResp = nextReq.convertedResp + prepared.passthroughResponses = nextReq.passthroughResponses + prepared.nativeResponses = nextReq.nativeResponses + logCtx.RealModelID = *realModelID + if span := trace.SpanFromContext(r.Context()); span.IsRecording() { + span.SetAttributes( + attribute.String("aar.real_model", *realModelID), + attribute.String("aar.credential", nextCred.Name), + attribute.String("aar.provider", string(nextCred.Type)), + attribute.Bool("aar.provider_compatibility_skip", true), + ) + } + return true + } + + success, fallbackReason := p.TryFallbackProxy( + w, + requestWithPath(r, prepared.proxyPath), + modelID, + (*cred).Name, + http.StatusBadRequest, + RetryReasonServerErr, + *proxyBody, + start, + logCtx, + ) + if success { + return false + } + p.logger.DebugContext(r.Context(), "No fallback handled unsupported image provider request", + "credential", (*cred).Name, + "model", modelID, + "reason", reason, + "fallback_reason", fallbackReason) + logCtx.Credential = *cred + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = unsupportedImageProviderRequestMessage + WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) + return false +} + +func (p *Proxy) nextPrimaryAfterUnsupportedSosana( + r *http.Request, + prepared *orchestratedRequest, + modelID string, + currentCred *config.CredentialConfig, + visibility scope.Context, + reason string, +) (*config.CredentialConfig, credentialPreparedRequest, bool) { + triedCreds := GetTried(r.Context()) + triedCreds[currentCred.Name] = true + + for attempts := 0; attempts < 128; attempts++ { + candidate, err := p.balancer.NextForModelExcludingScoped(modelID, triedCreds, visibility) + if err != nil { + p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", + "model", modelID, + "credential", currentCred.Name, + "reason", reason, + "error", err) + return nil, credentialPreparedRequest{}, false + } + triedCreds[candidate.Name] = true + if candidate.Type == config.ProviderTypeSosana { + continue + } + + nextReq, prepErr := p.prepareRequestForCredential( + r, + prepared.baseBody, + prepared.baseProxyBody, + modelID, + prepared.baseRealModelID, + prepared.basePath, + prepared.streaming, + candidate, + prepared.isResponsesAPI, + prepared.responsesPrevHandled, + prepared.stickyCacheEligible, + ) + if prepErr != nil { + p.logger.WarnContext(r.Context(), "Failed to prepare alternate primary request after image compatibility skip", + "credential", candidate.Name, + "provider", string(candidate.Type), + "model", modelID, + "reason", reason, + "error", prepErr) + continue + } + + p.logger.InfoContext(r.Context(), "Skipping incompatible image credential for unsupported image request", + "credential", currentCred.Name, + "next_credential", candidate.Name, + "model", modelID, + "reason", reason) + return candidate, nextReq, true + } + + p.logger.WarnContext(r.Context(), "Image compatibility skip exhausted primary credential scan", + "credential", currentCred.Name, + "model", modelID, + "reason", reason) + return nil, credentialPreparedRequest{}, false +} + func (p *Proxy) handleSosanaRequest( w http.ResponseWriter, r *http.Request, @@ -296,7 +436,7 @@ func (p *Proxy) sosanaCompletedImageBody( p.logger.DebugContext(ctx, "Downloaded Sosana result image", "credential", cred.Name, "model", modelID, - "result_host", sosanaResultHost(task), + "result_host", sosana.ResultHost(task), "image_bytes", len(image), "content_type", contentType, "request_id", logCtx.RequestID) @@ -311,330 +451,67 @@ func (p *Proxy) downloadSosanaResultImage( rawTaskBody []byte, logCtx *RequestLogContext, ) ([]byte, string, int, error) { - resultURL := "" - if task.ResultFileURL != nil { - resultURL = strings.TrimSpace(*task.ResultFileURL) - } - parsed, err := parseSosanaResultURL(resultURL) - if err != nil { - p.logUpstreamError(ctx, "Sosana completed task returned invalid result URL", http.StatusBadGateway, cred, modelID, rawTaskBody, - "request_id", logCtx.RequestID, - "error", err) - return nil, "", http.StatusBadGateway, err - } - if err := validateSosanaResultURL(ctx, parsed); err != nil { - p.logUpstreamError(ctx, "Sosana completed task returned unsafe result URL", http.StatusBadGateway, cred, modelID, rawTaskBody, - "result_host", parsed.Hostname(), - "request_id", logCtx.RequestID, - "error", err) - return nil, "", http.StatusBadGateway, err + image, err := sosana.DownloadResultImage(ctx, p.client, task) + if err == nil { + return image.Bytes, image.ContentType, http.StatusOK, nil } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, resultURL, nil) - if err != nil { - p.logUpstreamError(ctx, "Failed to build Sosana result image request", http.StatusBadGateway, cred, modelID, nil, - "result_host", parsed.Hostname(), - "request_id", logCtx.RequestID, - "error", err) - return nil, "", http.StatusBadGateway, err + statusCode := http.StatusBadGateway + resultHost := sosana.ResultHost(task) + var imageErr *sosana.ResultImageError + if errors.As(err, &imageErr) { + statusCode = imageErr.StatusCode + resultHost = imageErr.Host } - req.Header.Set("Accept", "image/*") - resp, err := p.doSosanaResultImageRequest(req) - if err != nil { - statusCode := http.StatusBadGateway - if isTimeoutError(err) || errors.Is(ctx.Err(), context.DeadlineExceeded) { - statusCode = http.StatusRequestTimeout - } - p.logUpstreamError(context.Background(), "Sosana result image download failed", statusCode, cred, modelID, nil, - "result_host", parsed.Hostname(), + switch { + case imageErr == nil: + p.logUpstreamError(ctx, "Sosana result image download failed", statusCode, cred, modelID, nil, + "result_host", resultHost, "request_id", logCtx.RequestID, "error", err) - return nil, "", statusCode, err - } - defer func() { - if closeErr := resp.Body.Close(); closeErr != nil { - p.logger.WarnContext(ctx, "Failed to close Sosana result image body", "error", closeErr) + case imageErr.ResponseBody != nil: + message := "Sosana result image download returned error status" + if imageErr.SniffedContentType != "" { + message = "Sosana result URL returned non-PNG content" } - }() - - contentType := strings.TrimSpace(resp.Header.Get("Content-Type")) - if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { - errorBody := readTextBodyForSosanaResultLog(resp.Body, contentType) - p.logUpstreamError(ctx, "Sosana result image download returned error status", http.StatusBadGateway, cred, modelID, errorBody, - "upstream_status", resp.StatusCode, - "result_host", parsed.Hostname(), - "request_id", logCtx.RequestID) - return nil, "", http.StatusBadGateway, errors.New("sosana result image download returned error status") - } - - image, err := readLimitedSosanaResultImage(resp.Body) - if err != nil { - p.logUpstreamError(ctx, "Failed to read Sosana result image", http.StatusBadGateway, cred, modelID, nil, - "result_host", parsed.Hostname(), + attrs := []any{ + "result_host", resultHost, + "content_type", imageErr.ContentType, "request_id", logCtx.RequestID, - "error", err) - return nil, "", http.StatusBadGateway, err - } - sniffedType := http.DetectContentType(image) - if !isPNGContentType(contentType) && !isPNGContentType(sniffedType) { - responseBody := textBodyPrefixForSosanaResultLog(image, contentType, sniffedType) - p.logUpstreamError(ctx, "Sosana result URL returned non-PNG content", http.StatusBadGateway, cred, modelID, responseBody, - "result_host", parsed.Hostname(), - "content_type", contentType, - "sniffed_content_type", sniffedType, - "request_id", logCtx.RequestID) - return nil, "", http.StatusBadGateway, errors.New("sosana result URL returned non-PNG content") - } - if !isPNGContentType(contentType) { - contentType = sniffedType - } - return image, contentType, http.StatusOK, nil -} - -func (p *Proxy) doSosanaResultImageRequest(req *http.Request) (*http.Response, error) { - client := *p.client - client.Transport = sosanaResultImageTransport() - client.CheckRedirect = func(req *http.Request, via []*http.Request) error { - return http.ErrUseLastResponse - } - return client.Do(req) -} - -func sosanaResultImageTransport() http.RoundTripper { - transport := http.DefaultTransport.(*http.Transport).Clone() - transport.Proxy = nil - transport.DisableKeepAlives = true - transport.DialContext = dialSosanaResultAddress - return transport -} - -func dialSosanaResultAddress(ctx context.Context, network, address string) (net.Conn, error) { - host, port, err := net.SplitHostPort(address) - if err != nil { - return nil, err - } - - dialer := &net.Dialer{} - if allowPrivateSosanaResultHostForTests(host) { - return dialer.DialContext(ctx, network, address) - } - if ip := net.ParseIP(host); ip != nil { - if isUnsafeSosanaResultIP(ip) { - return nil, errors.New("sosana result_file_url resolves to a private address") - } - return dialer.DialContext(ctx, network, address) - } - - addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) - if err != nil { - return nil, err - } - if len(addrs) == 0 { - return nil, errors.New("sosana result_file_url host has no addresses") - } - for _, addr := range addrs { - if isUnsafeSosanaResultIP(addr.IP) { - return nil, errors.New("sosana result_file_url resolves to a private address") + "error", imageErr.Err, } - } - - var dialErr error - for _, addr := range addrs { - conn, err := dialer.DialContext(ctx, network, net.JoinHostPort(addr.IP.String(), port)) - if err == nil { - return conn, nil + if imageErr.UpstreamStatus != 0 { + attrs = append(attrs, "upstream_status", imageErr.UpstreamStatus) } - dialErr = err - } - if dialErr != nil { - return nil, dialErr - } - return nil, errors.New("sosana result_file_url host has no dialable addresses") -} - -func allowPrivateSosanaResultHostForTests(host string) bool { - if allowPrivateSosanaResultURLForTests == nil { - return false - } - return allowPrivateSosanaResultURLForTests(&url.URL{Scheme: "http", Host: host}) -} - -func parseSosanaResultURL(raw string) (*url.URL, error) { - if raw == "" { - return nil, errors.New("sosana task completed without result_file_url") - } - parsed, err := url.Parse(raw) - if err != nil { - return nil, err - } - if parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { - return nil, errors.New("sosana result_file_url must be an http or https URL") - } - return parsed, nil -} - -func validateSosanaResultURL(ctx context.Context, parsed *url.URL) error { - if allowPrivateSosanaResultURLForTests != nil && allowPrivateSosanaResultURLForTests(parsed) { - return nil - } - if parsed.Scheme != "https" { - return errors.New("sosana result_file_url must use https") - } - - host := parsed.Hostname() - if strings.EqualFold(host, "localhost") { - return errors.New("sosana result_file_url host is not allowed") - } - if !isAllowedSosanaResultHost(host) { - return errors.New("sosana result_file_url host is not allowed") - } - if ip := net.ParseIP(host); ip != nil { - if isUnsafeSosanaResultIP(ip) { - return errors.New("sosana result_file_url resolves to a private address") + if imageErr.SniffedContentType != "" { + attrs = append(attrs, "sniffed_content_type", imageErr.SniffedContentType) } - return nil - } - - addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) - if err != nil { - return err - } - if len(addrs) == 0 { - return errors.New("sosana result_file_url host has no addresses") - } - for _, addr := range addrs { - if isUnsafeSosanaResultIP(addr.IP) { - return errors.New("sosana result_file_url resolves to a private address") + p.logUpstreamError(ctx, message, statusCode, cred, modelID, imageErr.ResponseBody, attrs...) + default: + message := "Sosana result image download failed" + if strings.Contains(err.Error(), "invalid") || strings.Contains(err.Error(), "without result_file_url") { + message = "Sosana completed task returned invalid result URL" } - } - return nil -} - -func isAllowedSosanaResultHost(host string) bool { - host = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(host)), ".") - for _, suffix := range []string{ - "sosana.blog", - "sosana.art", - "storage.yandexcloud.net", - } { - if host == suffix || strings.HasSuffix(host, "."+suffix) { - return true - } - } - return false -} - -func isUnsafeSosanaResultIP(ip net.IP) bool { - if ip == nil { - return true - } - if v4 := ip.To4(); v4 != nil && v4[0] == 100 && v4[1]&0xc0 == 64 { - return true - } - return ip.IsLoopback() || - ip.IsPrivate() || - ip.IsLinkLocalUnicast() || - ip.IsLinkLocalMulticast() || - ip.IsMulticast() || - ip.IsUnspecified() -} - -func sosanaResultHost(task sosana.BananaTaskResponse) string { - if task.ResultFileURL == nil { - return "" - } - parsed, err := url.Parse(strings.TrimSpace(*task.ResultFileURL)) - if err != nil { - return "" - } - return parsed.Hostname() -} - -func readLimitedSosanaResultImage(body io.Reader) ([]byte, error) { - data, err := io.ReadAll(io.LimitReader(body, maxSosanaResultImageBytes+1)) - if err != nil { - return nil, err - } - if int64(len(data)) > maxSosanaResultImageBytes { - return nil, ErrResponseBodyTooLarge - } - if len(data) == 0 { - return nil, errors.New("sosana result image body is empty") - } - return data, nil -} - -func readTextBodyForSosanaResultLog(body io.Reader, contentType string) []byte { - if !isTextContentType(contentType) { - return nil - } - data, err := io.ReadAll(io.LimitReader(body, maxSosanaResultErrorBytes)) - if err != nil { - return nil - } - return data -} - -func textBodyPrefixForSosanaResultLog(body []byte, contentTypes ...string) []byte { - for _, contentType := range contentTypes { - if isTextContentType(contentType) { - if int64(len(body)) > maxSosanaResultErrorBytes { - return body[:maxSosanaResultErrorBytes] - } - return body + if strings.Contains(err.Error(), "host is not allowed") || + strings.Contains(err.Error(), "private address") || + strings.Contains(err.Error(), "must use https") { + message = "Sosana completed task returned unsafe result URL" } + p.logUpstreamError(ctx, message, statusCode, cred, modelID, nil, + "result_host", resultHost, + "request_id", logCtx.RequestID, + "error", err) } - return nil -} - -func isPNGContentType(contentType string) bool { - contentType = strings.ToLower(strings.TrimSpace(contentType)) - return contentType == "image/png" || strings.HasPrefix(contentType, "image/png;") -} - -func isTextContentType(contentType string) bool { - contentType = strings.ToLower(strings.TrimSpace(contentType)) - return strings.HasPrefix(contentType, "text/") || - strings.Contains(contentType, "json") || - strings.Contains(contentType, "xml") + return nil, "", statusCode, err } func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosana.BananaTaskResponse, []byte, int, error) { - var reader *bytes.Reader - if body != nil { - reader = bytes.NewReader(body) - } else { - reader = bytes.NewReader(nil) - } - req, err := http.NewRequestWithContext(ctx, method, url, reader) + result, err := sosana.DoTaskRequest(ctx, p.client, method, url, cred.APIKey, body) if err != nil { - return sosana.BananaTaskResponse{}, nil, http.StatusInternalServerError, err - } - req.Header.Set("Authorization", "Bearer "+cred.APIKey) - if body != nil { - req.Header.Set("Content-Type", "application/json") - } - - resp, err := p.client.Do(req) - if err != nil { - return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err - } - defer func() { - if closeErr := resp.Body.Close(); closeErr != nil { - p.logger.WarnContext(ctx, "Failed to close Sosana response body", "error", closeErr) - } - }() - - rawBody, err := p.readLimitedResponseBody(resp.Body) - if err != nil { - return sosana.BananaTaskResponse{}, nil, http.StatusBadGateway, err - } - var task sosana.BananaTaskResponse - if len(rawBody) > 0 { - _ = json.Unmarshal(rawBody, &task) + return result.Task, result.RawBody, result.StatusCode, err } - return task, rawBody, resp.StatusCode, nil + return result.Task, result.RawBody, result.StatusCode, nil } func (p *Proxy) sosanaTransportError(ctx context.Context, err error, cred *config.CredentialConfig, modelID, url string, logCtx *RequestLogContext) ([]byte, int) { diff --git a/internal/proxy/sosana_live_test.go b/internal/proxy/sosana_live_test.go deleted file mode 100644 index 6ea6da22..00000000 --- a/internal/proxy/sosana_live_test.go +++ /dev/null @@ -1,67 +0,0 @@ -package proxy - -import ( - "encoding/base64" - "encoding/json" - "net/http" - "net/http/httptest" - "os" - "strings" - "testing" - "time" - - "github.com/mixaill76/auto_ai_router/internal/config" - "github.com/mixaill76/auto_ai_router/internal/converter/openai" - "github.com/stretchr/testify/require" -) - -func TestProxyRequest_SosanaLiveAcceptance(t *testing.T) { - if os.Getenv("SOSANA_ACCEPTANCE") != "1" { - t.Skip("SOSANA_ACCEPTANCE=1 not set, skipping paid Sosana live acceptance test") - } - - apiKey := os.Getenv("SOSANA_API_KEY") - require.NotEmpty(t, apiKey, "SOSANA_API_KEY is required for Sosana live acceptance test") - - baseURL := os.Getenv("SOSANA_BASE_URL") - if baseURL == "" { - baseURL = "https://sosana.art" - } - model := os.Getenv("SOSANA_MODEL") - if model == "" { - model = "banana-2-1k-compliant" - } - prompt := os.Getenv("SOSANA_PROMPT") - if prompt == "" { - prompt = "a small blue cube on a white background" - } - - prx := NewTestProxyBuilder(). - WithSingleCredential("sosana-live", config.ProviderTypeSosana, baseURL, apiKey). - WithRequestTimeout(2 * time.Minute). - Build() - - body, err := json.Marshal(map[string]any{ - "model": model, - "prompt": prompt, - "size": "1024x1024", - "n": 1, - }) - require.NoError(t, err) - - req := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(string(body))) - req.Header.Set("Authorization", "Bearer master-key") - req.Header.Set("Content-Type", "application/json") - w := httptest.NewRecorder() - - prx.ProxyRequest(w, req) - - require.Equal(t, http.StatusOK, w.Code, "response body: %s", w.Body.String()) - var resp openai.OpenAIImageResponse - require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) - require.Len(t, resp.Data, 1) - require.Empty(t, resp.Data[0].URL) - require.NotEmpty(t, resp.Data[0].B64JSON) - _, err = base64.StdEncoding.DecodeString(resp.Data[0].B64JSON) - require.NoError(t, err) -} diff --git a/internal/proxy/sosana_routing.go b/internal/proxy/sosana_routing.go deleted file mode 100644 index 69900690..00000000 --- a/internal/proxy/sosana_routing.go +++ /dev/null @@ -1,158 +0,0 @@ -package proxy - -import ( - "net/http" - "time" - - "github.com/mixaill76/auto_ai_router/internal/config" - "github.com/mixaill76/auto_ai_router/internal/converter/sosana" - "github.com/mixaill76/auto_ai_router/internal/scope" - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/trace" -) - -const unsupportedImageProviderRequestMessage = "request parameters are not supported by available image providers" -const unsupportedProviderEndpointMessage = "request endpoint is not supported by available providers" - -func (p *Proxy) applySosanaCompatibilityRouting( - w http.ResponseWriter, - r *http.Request, - prepared *orchestratedRequest, - modelID string, - cred **config.CredentialConfig, - body *[]byte, - proxyBody *[]byte, - realModelID *string, - isImageGeneration bool, - isImageEdit bool, - logCtx *RequestLogContext, - start time.Time, -) bool { - if (*cred).Type != config.ProviderTypeSosana || (!isImageGeneration && !isImageEdit) { - return true - } - - reason := sosana.UnsupportedModel(*realModelID) - if reason == "" { - reason = sosana.UnsupportedRequest(r.URL.Path, *body, r.Header.Get("Content-Type")) - } - if reason == "" { - return true - } - - nextCred, nextReq, routed := p.nextPrimaryAfterUnsupportedSosana(r, prepared, modelID, *cred, logCtx.Scope, reason) - if routed { - *cred = nextCred - *body = nextReq.body - *proxyBody = nextReq.proxyBody - *realModelID = nextReq.realModelID - r.URL.Path = nextReq.path - prepared.body = nextReq.body - prepared.proxyBody = nextReq.proxyBody - prepared.proxyPath = nextReq.proxyPath - prepared.realModelID = nextReq.realModelID - prepared.convertedResp = nextReq.convertedResp - prepared.passthroughResponses = nextReq.passthroughResponses - prepared.nativeResponses = nextReq.nativeResponses - logCtx.RealModelID = *realModelID - if span := trace.SpanFromContext(r.Context()); span.IsRecording() { - span.SetAttributes( - attribute.String("aar.real_model", *realModelID), - attribute.String("aar.credential", nextCred.Name), - attribute.String("aar.provider", string(nextCred.Type)), - attribute.Bool("aar.provider_compatibility_skip", true), - ) - } - return true - } - - success, fallbackReason := p.TryFallbackProxy( - w, - requestWithPath(r, prepared.proxyPath), - modelID, - (*cred).Name, - http.StatusBadRequest, - RetryReasonServerErr, - *proxyBody, - start, - logCtx, - ) - if success { - return false - } - p.logger.DebugContext(r.Context(), "No fallback handled unsupported image provider request", - "credential", (*cred).Name, - "model", modelID, - "reason", reason, - "fallback_reason", fallbackReason) - logCtx.Credential = *cred - logCtx.Status = "failure" - logCtx.HTTPStatus = http.StatusBadRequest - logCtx.ErrorMsg = unsupportedImageProviderRequestMessage - WriteErrorBadRequest(w, unsupportedImageProviderRequestMessage) - return false -} - -func (p *Proxy) nextPrimaryAfterUnsupportedSosana( - r *http.Request, - prepared *orchestratedRequest, - modelID string, - currentCred *config.CredentialConfig, - visibility scope.Context, - reason string, -) (*config.CredentialConfig, credentialPreparedRequest, bool) { - triedCreds := GetTried(r.Context()) - triedCreds[currentCred.Name] = true - - for attempts := 0; attempts < 128; attempts++ { - candidate, err := p.balancer.NextForModelExcludingScoped(modelID, triedCreds, visibility) - if err != nil { - p.logger.DebugContext(r.Context(), "No compatible primary credential available for unsupported image request", - "model", modelID, - "credential", currentCred.Name, - "reason", reason, - "error", err) - return nil, credentialPreparedRequest{}, false - } - triedCreds[candidate.Name] = true - if candidate.Type == config.ProviderTypeSosana { - continue - } - - nextReq, prepErr := p.prepareRequestForCredential( - r, - prepared.baseBody, - prepared.baseProxyBody, - modelID, - prepared.baseRealModelID, - prepared.basePath, - prepared.streaming, - candidate, - prepared.isResponsesAPI, - prepared.responsesPrevHandled, - prepared.stickyCacheEligible, - ) - if prepErr != nil { - p.logger.WarnContext(r.Context(), "Failed to prepare alternate primary request after image compatibility skip", - "credential", candidate.Name, - "provider", string(candidate.Type), - "model", modelID, - "reason", reason, - "error", prepErr) - continue - } - - p.logger.InfoContext(r.Context(), "Skipping incompatible image credential for unsupported image request", - "credential", currentCred.Name, - "next_credential", candidate.Name, - "model", modelID, - "reason", reason) - return candidate, nextReq, true - } - - p.logger.WarnContext(r.Context(), "Image compatibility skip exhausted primary credential scan", - "credential", currentCred.Name, - "model", modelID, - "reason", reason) - return nil, credentialPreparedRequest{}, false -} diff --git a/internal/proxy/sosana_test.go b/internal/proxy/sosana_test.go index 8afc2d75..c6613534 100644 --- a/internal/proxy/sosana_test.go +++ b/internal/proxy/sosana_test.go @@ -869,65 +869,6 @@ func TestProxyRequest_SosanaImageResultRedirectMaskedAndNotFollowed(t *testing.T assert.NotContains(t, w.Body.String(), targetServer.URL) } -func TestDownloadSosanaResultImageRejectsUnsafeProductionURL(t *testing.T) { - var logBuf bytes.Buffer - called := false - imageServer := newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - called = true - })) - defer imageServer.Close() - - resultURL := imageServer.URL + "/private.png" - prx := newSosanaTestProxy("https://sosana.art", &logBuf) - cred := &config.CredentialConfig{Name: "sosana", Type: config.ProviderTypeSosana, BaseURL: "https://sosana.art"} - image, _, statusCode, err := prx.downloadSosanaResultImage(context.Background(), cred, "banana-2-1k-compliant", sosana.BananaTaskResponse{ - Status: sosana.StatusCompleted, - ResultFileURL: &resultURL, - }, nil, &RequestLogContext{}) - - require.Error(t, err) - assert.Nil(t, image) - assert.Equal(t, http.StatusBadGateway, statusCode) - assert.False(t, called) - assert.Contains(t, logBuf.String(), "unsafe result URL") - assert.Contains(t, logBuf.String(), "result_host=127.0.0.1") -} - -func TestValidateSosanaResultURLBlocksUnsafeHosts(t *testing.T) { - tests := []string{ - "http://main-r2.sosana.blog/image.png", - "https://localhost/image.png", - "https://127.0.0.1/image.png", - "https://169.254.169.254/latest/meta-data", - "https://100.64.0.1/image.png", - "https://example.com/image.png", - } - - for _, rawURL := range tests { - t.Run(rawURL, func(t *testing.T) { - parsed, err := parseSosanaResultURL(rawURL) - require.NoError(t, err) - require.Error(t, validateSosanaResultURL(context.Background(), parsed)) - }) - } -} - -func TestValidateSosanaResultURLAllowsLocalOnlyWithTestHook(t *testing.T) { - allowPrivateSosanaResultURLsForTest(t) - - parsed, err := parseSosanaResultURL("http://127.0.0.1/image.png") - require.NoError(t, err) - require.NoError(t, validateSosanaResultURL(context.Background(), parsed)) -} - -func TestDialSosanaResultAddressRejectsPrivateIP(t *testing.T) { - conn, err := dialSosanaResultAddress(context.Background(), "tcp", net.JoinHostPort("127.0.0.1", "443")) - - require.Error(t, err) - assert.Nil(t, conn) - assert.Contains(t, err.Error(), "private address") -} - func TestProxyRequest_SosanaImageResultTimeoutMasked(t *testing.T) { allowPrivateSosanaResultURLsForTest(t) @@ -1112,8 +1053,7 @@ func newSosanaResultImageServer(t *testing.T, status int, contentType string, bo func allowPrivateSosanaResultURLsForTest(t *testing.T) { t.Helper() - previous := allowPrivateSosanaResultURLForTests - allowPrivateSosanaResultURLForTests = func(parsed *url.URL) bool { + restore := sosana.SetAllowPrivateResultURLForTests(func(parsed *url.URL) bool { if parsed.Scheme != "http" { return false } @@ -1122,11 +1062,9 @@ func allowPrivateSosanaResultURLsForTest(t *testing.T) { return true } ip := net.ParseIP(host) - return ip != nil && isUnsafeSosanaResultIP(ip) - } - t.Cleanup(func() { - allowPrivateSosanaResultURLForTests = previous + return ip != nil && sosana.IsUnsafeResultIP(ip) }) + t.Cleanup(restore) } func newSosanaTestProxy(baseURL string, logBuf *bytes.Buffer) *Proxy { From 9772fee673a818bcea0cde1eccf4ca6e56e7ab8f Mon Sep 17 00:00:00 2001 From: fello <7552494@gmail.com> Date: Thu, 6 Aug 2026 15:52:39 +0300 Subject: [PATCH 16/16] Format configuration docs --- docs/getting-started/configuration.md | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md index 3bc8f7ec..f5ca446a 100644 --- a/docs/getting-started/configuration.md +++ b/docs/getting-started/configuration.md @@ -172,17 +172,17 @@ Each credential defines a connection to an LLM provider. See [Providers](../prov Common fields for all credentials: -| Field | Type | Description | -| ------------------ | ------ | ------------------------------------------------------------------------------------------- | -| `name` | string | Unique credential identifier | +| Field | Type | Description | +| ------------------ | ------ | ---------------------------------------------------------------------------------------------------------------------- | +| `name` | string | Unique credential identifier | | `type` | string | Provider type: `openai`, `anthropic`, `cometapi`, `sosana`, `vertex-ai`, `gemini`, `bedrock`, `proxy`, `air`, `proman` | -| `rpm` | int | Requests per minute limit (-1 = unlimited) | -| `tpm` | int | Tokens per minute limit (-1 = unlimited) | -| `is_fallback` | bool | Use as fallback when primary credentials are exhausted | -| `reasoning_only` | bool | Route only requests that explicitly enable reasoning/thinking | -| `scopes` | list | Optional client scopes allowed to use and see this credential | -| `denied_scopes` | list | Optional client scopes that must not use or see this credential | -| `forbidden_scopes` | list | Alias for `denied_scopes` | +| `rpm` | int | Requests per minute limit (-1 = unlimited) | +| `tpm` | int | Tokens per minute limit (-1 = unlimited) | +| `is_fallback` | bool | Use as fallback when primary credentials are exhausted | +| `reasoning_only` | bool | Route only requests that explicitly enable reasoning/thinking | +| `scopes` | list | Optional client scopes allowed to use and see this credential | +| `denied_scopes` | list | Optional client scopes that must not use or see this credential | +| `forbidden_scopes` | list | Alias for `denied_scopes` | ### Scoped credential visibility