diff --git a/README.md b/README.md index 472ad41..1c20905 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 5e659af..8beeea1 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -83,6 +83,15 @@ credentials: rpm: 60 tpm: -1 + # Sosana.art image generation API (OpenAI Images-compatible through the router) + # 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" + base_url: "https://sosana.art" + rpm: 60 + tpm: -1 + # CheapGPT / AIProductiv using the Anthropic-compatible Messages API - name: "cheapgpt_anthropic" type: "anthropic" @@ -113,6 +122,14 @@ models: rpm: 100 tpm: 50000 + # Sosana.art image model exposed through /v1/images/generations and /v1/images/edits. + # 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 + # 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 d20060d..f5ca446 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 | -| `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 | -| `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` | +| 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` | ### Scoped credential visibility diff --git a/docs/index.md b/docs/index.md index 73399d6..3e3c4c8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -34,7 +34,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 - **AIR 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 13bf71e..6877b7f 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` | API Key via Bearer | | [ProMan](proman.md) | `proman` | `api_key`, `base_url` | API Key | | [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 | diff --git a/docs/providers/sosana.md b/docs/providers/sosana.md new file mode 100644 index 0000000..0c877bd --- /dev/null +++ b/docs/providers/sosana.md @@ -0,0 +1,96 @@ +# 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[].b64_json`. + +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: "google/gemini-3.1-flash-image-preview" + model: "banana-2-{image_size}-compliant" + credential: sosana_images + rpm: 60 + 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`. +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`. + +## Behavior + +- `n` must be `1`. +- 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"` 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. 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. +- 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 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. + +## 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 +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, 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.go b/internal/config/config.go index 50565db..d0d5cba 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" ProviderTypeProMan ProviderType = "proman" ProviderTypeBedrock ProviderType = "bedrock" ProviderTypeProxy ProviderType = "proxy" @@ -42,7 +43,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, ProviderTypeProMan, ProviderTypeBedrock, ProviderTypeProxy, ProviderTypeAIR: + case ProviderTypeOpenAI, ProviderTypeVertexAI, ProviderTypeGemini, ProviderTypeAnthropic, ProviderTypeCometAPI, ProviderTypeSosana, ProviderTypeProMan, ProviderTypeBedrock, ProviderTypeProxy, ProviderTypeAIR: return true } return false @@ -59,6 +60,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 case "aar", "auto-ai-router", "auto_ai_router": return ProviderTypeAIR case "pro-man", "pro_man": @@ -791,7 +794,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 @@ -1646,7 +1648,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', 'proman', 'bedrock', 'proxy', or 'air')", cred.Name, cred.Type) + return fmt.Errorf("credential %s: invalid type: %s (must be 'openai', 'vertex-ai', 'gemini', 'anthropic', 'cometapi', 'sosana', 'proman', 'bedrock', 'proxy', or 'air')", 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 3dfa909..8575bac 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -830,6 +830,7 @@ func TestProviderType_IsValid(t *testing.T) { {"openai", ProviderTypeOpenAI, true}, {"vertex-ai", ProviderTypeVertexAI, true}, {"cometapi", ProviderTypeCometAPI, true}, + {"sosana", ProviderTypeSosana, true}, {"air", ProviderTypeAIR, true}, {"proman", ProviderTypeProMan, true}, {"invalid", ProviderType("azure"), false}, @@ -865,6 +866,33 @@ rpm: 60 assert.Equal(t, ProviderTypeCometAPI, cred.Type) } +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 TestCredentialConfig_NormalizeAIRProviderTypeAliases(t *testing.T) { tests := []struct { name string @@ -893,6 +921,44 @@ rpm: 60 } } +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 TestCredentialConfig_NormalizeProManProviderType(t *testing.T) { var cred CredentialConfig err := yaml.Unmarshal([]byte(` @@ -1340,6 +1406,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, 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) { tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, "config.yaml") diff --git a/internal/converter/sosana/client.go b/internal/converter/sosana/client.go new file mode 100644 index 0000000..8711d41 --- /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 0000000..5488354 --- /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/compatibility.go b/internal/converter/sosana/compatibility.go new file mode 100644 index 0000000..453c1ad --- /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 0000000..c3a9a3c --- /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 0000000..b817ccf --- /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/converter/sosana/live_test.go b/internal/converter/sosana/live_test.go new file mode 100644 index 0000000..58cefcf --- /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/litellmdb/model_table/model_table.go b/internal/litellmdb/model_table/model_table.go index d920690..c6f2bf7 100644 --- a/internal/litellmdb/model_table/model_table.go +++ b/internal/litellmdb/model_table/model_table.go @@ -300,6 +300,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, "proman") || strings.Contains(p, "pro-man") || strings.Contains(p, "pro_man"): return config.ProviderTypeProMan case strings.Contains(p, "xai"): @@ -394,7 +396,9 @@ func convertPricingToModelPrice(p *queries.CustomPricingLiteLLMParams) *manager. if p == nil { return nil } - if p.InputCostPerToken == nil && p.OutputCostPerToken == nil && len(p.SearchContextCostPerQuery) == 0 { + if p.InputCostPerToken == nil && p.OutputCostPerToken == nil && + p.InputCostPerImage == nil && p.OutputCostPerImage == nil && p.OutputCostPerImageToken == nil && + len(p.SearchContextCostPerQuery) == 0 { return nil } @@ -450,6 +454,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.CacheReadInputAudioTokenCost != nil { price.CacheReadInputAudioTokenCost = *p.CacheReadInputAudioTokenCost } diff --git a/internal/litellmdb/model_table/model_table_test.go b/internal/litellmdb/model_table/model_table_test.go index 12b418e..cdd0088 100644 --- a/internal/litellmdb/model_table/model_table_test.go +++ b/internal/litellmdb/model_table/model_table_test.go @@ -24,6 +24,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}, {"proman", "proman", config.ProviderTypeProMan}, {"pro-man", "pro-man", config.ProviderTypeProMan}, {"xai", "xAI", config.ProviderTypeOpenAI}, @@ -127,6 +129,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 @@ -137,6 +140,7 @@ func TestConvertPricingToModelPrice(t *testing.T) { OutputCostPerReasoningToken: &outputReasoning, CacheReadInputTokenCost: &cacheRead, CacheCreationInputTokenCost: &cacheCreation, + InputCostPerImage: &inputImage, OutputCostPerImage: &outputImage, OutputCostPerImageToken: &outputImageToken, InputCostPerTokenAbove200kTokens: &inputAbove200k, @@ -149,6 +153,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) @@ -157,6 +162,32 @@ 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_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 @@ -175,7 +206,8 @@ func TestConvertPricingToModelPrice_AllFields(t *testing.T) { cacheCreationAbove1hrAbove200k := 0.0857 cacheReadAbove272k := 0.086 cacheCreationAbove272k := 0.087 - cacheReadAudio := 0.088 + inputImage := 0.088 + cacheReadAudio := 0.089 outputImage := 0.09 outputImageToken := 0.10 searchContextCost := map[string]float64{ @@ -203,6 +235,7 @@ func TestConvertPricingToModelPrice_AllFields(t *testing.T) { CacheCreationInputTokenCostAbove1hrAbove200kTokens: &cacheCreationAbove1hrAbove200k, CacheReadInputTokenCostAbove272kTokens: &cacheReadAbove272k, CacheCreationInputTokenCostAbove272kTokens: &cacheCreationAbove272k, + InputCostPerImage: &inputImage, CacheReadInputAudioTokenCost: &cacheReadAudio, OutputCostPerImage: &outputImage, OutputCostPerImageToken: &outputImageToken, @@ -228,6 +261,7 @@ func TestConvertPricingToModelPrice_AllFields(t *testing.T) { assert.Equal(t, cacheCreationAbove1hrAbove200k, price.CacheCreationInputTokenCostAbove1hrAbove200k) assert.Equal(t, cacheReadAbove272k, price.CacheReadInputTokenCostAbove272k) assert.Equal(t, cacheCreationAbove272k, price.CacheCreationInputTokenCostAbove272k) + assert.Equal(t, inputImage, price.InputCostPerImage) assert.Equal(t, cacheReadAudio, price.CacheReadInputAudioTokenCost) 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 e1cf631..e4ec6ff 100644 --- a/internal/models/manager.go +++ b/internal/models/manager.go @@ -62,6 +62,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"` // Built-in web search tool pricing. Values are per query/call, keyed by @@ -455,6 +456,7 @@ var providerPassthroughDefaults = map[config.ProviderType]bool{ config.ProviderTypeGemini: false, config.ProviderTypeAnthropic: false, config.ProviderTypeCometAPI: false, + config.ProviderTypeSosana: false, config.ProviderTypeProMan: false, config.ProviderTypeBedrock: false, } @@ -2362,6 +2364,7 @@ var providerTypeLiteLLMPrefix = map[config.ProviderType]string{ config.ProviderTypeGemini: "gemini", config.ProviderTypeAnthropic: "anthropic", config.ProviderTypeCometAPI: "cometapi", + config.ProviderTypeSosana: "sosana", config.ProviderTypeProMan: "proman", config.ProviderTypeBedrock: "bedrock", config.ProviderTypeProxy: "openai", diff --git a/internal/models/price_calculator.go b/internal/models/price_calculator.go index d97dcb9..dc75e6b 100644 --- a/internal/models/price_calculator.go +++ b/internal/models/price_calculator.go @@ -227,6 +227,10 @@ func CalculateTokenCosts(usage *converter.TokenUsage, price *ModelPrice) *conver costs.ImageCost += float64(outputImageTokens) * outputImageCost } else if imageCount > 0 && price.OutputCostPerImage > 0 { costs.ImageCost += float64(imageCount) * price.OutputCostPerImage + } else if imageCount > 0 && price.InputCostPerImage > 0 { + costs.ImageCost += float64(imageCount) * price.InputCostPerImage + } else if imageCount > 0 && price.OutputCostPerImageToken > 0 { + costs.ImageCost += float64(imageCount) * price.OutputCostPerImageToken } webSearchRequests := converterutil.NonNegativeTokenCount(usage.WebSearchRequests) diff --git a/internal/models/price_calculator_test.go b/internal/models/price_calculator_test.go index 06f6481..45b3a0e 100644 --- a/internal/models/price_calculator_test.go +++ b/internal/models/price_calculator_test.go @@ -359,6 +359,61 @@ 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_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, + } + 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/models/price_loader_test.go b/internal/models/price_loader_test.go index 79f2dc3..734ffb6 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_CachePricingExtensions(t *testing.T) { filePath := filepath.Join(t.TempDir(), "prices.json") pricesJSON := `{ diff --git a/internal/proxy/cometapi_test.go b/internal/proxy/cometapi_test.go index 38a7b0e..34b01fe 100644 --- a/internal/proxy/cometapi_test.go +++ b/internal/proxy/cometapi_test.go @@ -6,49 +6,116 @@ import ( "github.com/mixaill76/auto_ai_router/internal/config" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestIsCometAPICredential(t *testing.T) { tests := []struct { - name string - cred *config.CredentialConfig - want bool + name string + cred *config.CredentialConfig + wantComet 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, }, { - 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, }, { - 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, }, { - 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, }, { - name: "nil credential", - cred: nil, - want: false, + name: "nil credential", + cred: nil, + wantComet: 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.wantComet, isCometAPICredential(tt.cred)) }) } } +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) + 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) { + tests := []struct { + name string + cred *config.CredentialConfig + wantSosana bool + }{ + { + name: "dedicated provider type", + cred: &config.CredentialConfig{Type: config.ProviderTypeSosana}, + wantSosana: true, + }, + { + name: "sosana host fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://sosana.art"}, + wantSosana: true, + }, + { + name: "sosana name fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, Name: "sosana-art-images"}, + wantSosana: true, + }, + { + name: "sasana host fallback", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://api.sasana.example/v1"}, + wantSosana: true, + }, + { + name: "regular openai", + cred: &config.CredentialConfig{Type: config.ProviderTypeOpenAI, BaseURL: "https://api.openai.com"}, + wantSosana: false, + }, + { + name: "nil credential", + cred: nil, + wantSosana: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.wantSosana, isSosanaCredential(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"}}` @@ -57,6 +124,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/errors.go b/internal/proxy/errors.go index 86616b7..9b899df 100644 --- a/internal/proxy/errors.go +++ b/internal/proxy/errors.go @@ -90,6 +90,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/proxy.go b/internal/proxy/proxy.go index 338a036..2b8180e 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -24,6 +24,7 @@ import ( anthropicconv "github.com/mixaill76/auto_ai_router/internal/converter/anthropic" promanutils "github.com/mixaill76/auto_ai_router/internal/converter/proman/utils" "github.com/mixaill76/auto_ai_router/internal/converter/responses" + "github.com/mixaill76/auto_ai_router/internal/converter/sosana" "github.com/mixaill76/auto_ai_router/internal/httputil" "github.com/mixaill76/auto_ai_router/internal/kafkalog" "github.com/mixaill76/auto_ai_router/internal/litellmdb" @@ -81,7 +82,7 @@ func (p *Proxy) applyCredentialCompatibilityRouting( return true } - reason, provider := unsupportedCredentialRequest(cred, prepared.body, modelID) + reason, provider := unsupportedCredentialRequest(cred, prepared.body, modelID, prepared.basePath, r.Header.Get("Content-Type"), prepared.realModelID) if reason == "" { return true } @@ -179,9 +180,6 @@ func (p *Proxy) nextPrimaryAfterUnsupportedCredential( return nil, credentialPreparedRequest{}, false } triedCreds[candidate.Name] = true - if candidateReason, _ := unsupportedCredentialRequest(candidate, prepared.baseBody, modelID); candidateReason != "" { - continue - } nextReq, prepErr := p.prepareRequestForCredential( r, @@ -205,6 +203,9 @@ func (p *Proxy) nextPrimaryAfterUnsupportedCredential( "error", prepErr) continue } + if candidateReason, _ := unsupportedCredentialRequest(candidate, nextReq.body, modelID, prepared.basePath, r.Header.Get("Content-Type"), nextReq.realModelID); candidateReason != "" { + continue + } p.logger.InfoContext(r.Context(), "Skipping incompatible credential for unsupported request", "credential", currentCred.Name, @@ -223,7 +224,13 @@ func (p *Proxy) nextPrimaryAfterUnsupportedCredential( return nil, credentialPreparedRequest{}, false } -func unsupportedCredentialRequest(cred *config.CredentialConfig, body []byte, selectedModel string) (string, string) { +func unsupportedCredentialRequest(cred *config.CredentialConfig, body []byte, selectedModel string, path string, contentType string, realModelID string) (string, string) { + if cred != nil && cred.Type == config.ProviderTypeSosana { + if reason := sosana.UnsupportedModel(realModelID); reason != "" { + return reason, "Sosana" + } + return sosana.UnsupportedRequest(path, body, contentType), "Sosana" + } if promanutils.IsCredential(cred) { return promanutils.UnsupportedRequest(body, selectedModel), "ProMan" } @@ -928,10 +935,22 @@ func (p *Proxy) proxyRequest(w http.ResponseWriter, r *http.Request) { proxyBody = prepared.proxyBody realModelID = prepared.realModelID cred = prepared.cred + if cred.Type == config.ProviderTypeSosana && (isImageGeneration || isImageEdit) { + if _, concreteModelID, err := p.buildSosanaCreateBody(body, r.Header.Get("Content-Type"), realModelID, isImageEdit); err == nil && concreteModelID != "" { + realModelID = concreteModelID + prepared.realModelID = concreteModelID + logCtx.RealModelID = concreteModelID + } + } if !p.enforceBudgetAndRateLimits(w, r, logCtx, modelID, realModelID, body) { return } + if cred.Type == config.ProviderTypeSosana { + p.handleSosanaRequest(w, r, body, cred, modelID, realModelID, isImageGeneration, isImageEdit, logCtx, start) + return + } + // Handle proxy/AIR credential types with exact same-type retry + fallback. if cred.IsProxyLike() { logCtx.Credential = cred @@ -1080,6 +1099,7 @@ func (p *Proxy) proxyRequest(w http.ResponseWriter, r *http.Request) { p.metrics.RecordRequest(cred.Name, r.URL.Path, modelID, proxyResp.StatusCode, time.Since(start)) // Write response (streaming or non-streaming) + maskProxyError := proxyResp.StatusCode >= 400 && shouldMaskProxyResponseErrors(cred, proxyResp) tokenUsageOptions := tokenUsageExtractionOptionsForResponse(cred, proxyResp.Headers) // Populated in the non-streaming branch below via extractOpenAITokensAndUsage // (plan item G — one shared decode of proxyResp.Body instead of the @@ -1090,6 +1110,30 @@ func (p *Proxy) proxyRequest(w http.ResponseWriter, r *http.Request) { 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.convertedMessages && proxyResp.StatusCode >= 200 && proxyResp.StatusCode < 300 { @@ -1214,6 +1258,10 @@ func (p *Proxy) proxyRequest(w http.ResponseWriter, r *http.Request) { p.setSessionBinding(logCtx.SessionID, modelID, cred.Name) } } else { + if proxyResp.StatusCode >= 400 && shouldMaskProxyResponseErrors(cred, proxyResp) { + maskProxyErrorResponse(proxyResp) + } + // Save passthrough Responses API response or convert Chat Completions response if needed if prepared.convertedMessages && proxyResp.StatusCode >= 200 && proxyResp.StatusCode < 300 { messagesBody, convErr := anthropicconv.ChatToMessages(proxyResp.Body, prepared.messagesMetadata) @@ -1304,11 +1352,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 @@ -1865,7 +1920,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 tokenUsageOptions.AudioInputIncludesCachedAudio = false @@ -1908,7 +1969,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 { @@ -1945,7 +2011,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. diff --git a/internal/proxy/proxy_helpers.go b/internal/proxy/proxy_helpers.go index 1662200..3cef529 100644 --- a/internal/proxy/proxy_helpers.go +++ b/internal/proxy/proxy_helpers.go @@ -208,6 +208,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, "server_tool_use": serverToolUse, @@ -240,6 +241,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, "reasoning_cost": costs.ReasoningCost, "cached_input_cost": costs.CachedInputCost, "cache_creation_cost": costs.CacheCreationCost, diff --git a/internal/proxy/proxy_helpers_test.go b/internal/proxy/proxy_helpers_test.go index 3624b3c..a0e5102 100644 --- a/internal/proxy/proxy_helpers_test.go +++ b/internal/proxy/proxy_helpers_test.go @@ -237,6 +237,21 @@ func TestBuildMetadata(t *testing.T) { 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) + }) + t.Run("normalizes_usage", func(t *testing.T) { usage := &converter.TokenUsage{ PromptTokens: -1, diff --git a/internal/proxy/proxy_log.go b/internal/proxy/proxy_log.go index 727bf09..8d9afb7 100644 --- a/internal/proxy/proxy_log.go +++ b/internal/proxy/proxy_log.go @@ -46,14 +46,17 @@ 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)) } func shouldMaskUpstreamErrors(cred *config.CredentialConfig) bool { - return isCometAPICredential(cred) || promanutils.IsCredential(cred) + if cred == nil { + return false + } + return isCometAPICredential(cred) || isSosanaCredential(cred) || promanutils.IsCredential(cred) } func isCometAPICredential(cred *config.CredentialConfig) bool { @@ -86,6 +89,19 @@ func isProviderHost(rawBaseURL, domain string) bool { return host == domain || strings.HasSuffix(host, "."+domain) } +func isSosanaCredential(cred *config.CredentialConfig) bool { + if cred == nil { + return false + } + if cred.Type == config.ProviderTypeSosana { + return true + } + name := strings.ToLower(cred.Name) + return isProviderHost(cred.BaseURL, "sosana.art") || + containsSosanaMarker(name) || + containsSosanaMarker(cred.BaseURL) +} + // logStreamHandlerError logs a streaming handler failure. Client disconnects are // expected during normal operation and go to DEBUG; real failures go to ERROR. func (p *Proxy) logStreamHandlerError(ctx context.Context, msg string, err error, extra ...any) { diff --git a/internal/proxy/sosana.go b/internal/proxy/sosana.go new file mode 100644 index 0000000..f97cfc1 --- /dev/null +++ b/internal/proxy/sosana.go @@ -0,0 +1,388 @@ +package proxy + +import ( + "context" + "errors" + "math/rand" + "net/http" + "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" +) + +const ( + sosanaPollInterval = time.Duration(sosana.PollInterval) +) + +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 := unsupportedCredentialRequestMessage + logCtx.Status = "failure" + logCtx.HTTPStatus = http.StatusBadRequest + logCtx.ErrorMsg = message + WriteErrorBadRequest(w, message) + return + } + + baseRealModelID := realModelID + + 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.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) + 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) + } + + 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 = unsupportedCredentialRequestMessage + WriteErrorBadRequest(w, unsupportedCredentialRequestMessage) + return + } + logCtx.RealModelID = concreteModelID + if concreteModelID != "" { + priceModelID, modelPrice := lookupBillingModelPrice(p.priceRegistry, logCtx.PublicModelID, modelID, concreteModelID) + if modelPrice != nil { + logCtx.PriceModelID = priceModelID + logCtx.ModelPrice = modelPrice + logCtx.billingPriceResolved = true + logCtx.billingPriceModelID = priceModelID + logCtx.billingPrice = modelPrice + } + } + + 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) 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) + } + 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: false, + } + } + 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, 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), + "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) 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", sosana.ResultHost(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) { + image, err := sosana.DownloadResultImage(ctx, p.client, task) + if err == nil { + return image.Bytes, image.ContentType, http.StatusOK, nil + } + + statusCode := http.StatusBadGateway + resultHost := sosana.ResultHost(task) + var imageErr *sosana.ResultImageError + if errors.As(err, &imageErr) { + statusCode = imageErr.StatusCode + resultHost = imageErr.Host + } + + 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) + case imageErr.ResponseBody != nil: + message := "Sosana result image download returned error status" + if imageErr.SniffedContentType != "" { + message = "Sosana result URL returned non-PNG content" + } + attrs := []any{ + "result_host", resultHost, + "content_type", imageErr.ContentType, + "request_id", logCtx.RequestID, + "error", imageErr.Err, + } + if imageErr.UpstreamStatus != 0 { + attrs = append(attrs, "upstream_status", imageErr.UpstreamStatus) + } + if imageErr.SniffedContentType != "" { + attrs = append(attrs, "sniffed_content_type", imageErr.SniffedContentType) + } + 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" + } + 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, "", statusCode, err +} + +func (p *Proxy) doSosanaTaskRequest(ctx context.Context, method, url string, cred *config.CredentialConfig, body []byte) (sosana.BananaTaskResponse, []byte, int, error) { + result, err := sosana.DoTaskRequest(ctx, p.client, method, url, cred.APIKey, body) + if err != nil { + return result.Task, result.RawBody, result.StatusCode, err + } + 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) { + 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_test.go b/internal/proxy/sosana_test.go new file mode 100644 index 0000000..7c8236d --- /dev/null +++ b/internal/proxy/sosana_test.go @@ -0,0 +1,1145 @@ +package proxy + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "log/slog" + "mime/multipart" + "net" + "net/http" + "net/http/httptest" + "net/url" + "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/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" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +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 { + 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, "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"}`)) + case "/api/banana/task-1": + pollSeen = true + _, _ = 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) + } + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, nil) + 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() + + 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.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.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) + 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")) + + 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, "banana-2-1k-compliant", req["model"]) + assert.Equal(t, false, req["prompt_optimization"]) + 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,")) + + _, _ = 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() + + prx := newSosanaTestProxy(upstream.URL, nil) + setSosanaTestModels(prx, []config.ModelRPMConfig{ + {Name: "public-image", Model: "banana-2-1k-compliant", 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) + 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": + _, _ = 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{ + "banana-2-1k-compliant": {InputCostPerImage: 0.088113}, + }) + + prx := newSosanaTestProxy(upstream.URL, nil) + prx.LiteLLMDB = spendManager + prx.priceRegistry = priceRegistry + setSosanaTestModels(prx, []config.ModelRPMConfig{ + {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}`)) + 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, "aimage_generation", 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.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.088113, costBreakdown["image_cost"].(float64), 1e-12) + assert.InDelta(t, 0.088113, costBreakdown["total_cost"].(float64), 1e-12) +} + +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 + setSosanaTestModels(prx, []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 == "" { + 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": + _, _ = 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() + + 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{ + "banana-2-1k-compliant": {OutputCostPerImage: 0.07}, + }) + + prx := newSosanaTestProxy(upstream.URL, nil) + prx.LiteLLMDB = manager + prx.priceRegistry = priceRegistry + setSosanaTestModels(prx, []config.ModelRPMConfig{ + {Name: alias, Model: "banana-2-1k-compliant", 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) { + called = true + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + prx := newSosanaTestProxy(upstream.URL, nil) + 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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), unsupportedCredentialRequestMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), unsupportedCredentialRequestMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") + 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":"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.StatusBadRequest, w.Code) + assert.Contains(t, w.Body.String(), unsupportedCredentialRequestMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") + 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": "banana-2-1k-compliant", + "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(), unsupportedCredentialRequestMessage) + assert.NotContains(t, strings.ToLower(w.Body.String()), "sosana") + assert.False(t, called) +} + +func TestProxyRequest_SosanaHalfKPixelSizeRoutesToNextPrimary(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.Equal(t, "512x512", req["size"]) + 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","size":"512x512"}`)) + 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_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() + setSosanaTestModels(prx, []config.ModelRPMConfig{ + {Name: "banana-2-1k-compliant", Credential: "sosana"}, + {Name: "banana-2-1k-compliant", Credential: "hidden-primary"}, + }) + 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(), unsupportedCredentialRequestMessage) + 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) { + 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) { + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusPaymentRequired, w.Code) + assert.Contains(t, w.Body.String(), "Request failed") + assert.NotContains(t, w.Body.String(), "balance secret") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.Contains(t, logBuf.String(), "balance secret") +} + +func TestProxyRequest_SosanaRetriesCreateWithNextCredential(t *testing.T) { + var createAuths []string + var createModels []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": + 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")) + _, _ = 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) + } + })) + 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() + setSosanaTestModels(prx, []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":"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) + 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) +} + +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() + setSosanaTestModels(prx, []config.ModelRPMConfig{ + {Name: "banana-2-1k-compliant", Credential: "sosana-a"}, + {Name: "banana-2-1k-compliant", Credential: "sosana-b"}, + }) + 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(), "Rate limit exceeded") + 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 + 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":"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() + + prx.ProxyRequest(w, req) + + require.Equal(t, http.StatusBadGateway, w.Code) + assert.False(t, liveCalled) + assert.Contains(t, w.Body.String(), "Request failed") + 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) { + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusInternalServerError, w.Code) + assert.Contains(t, w.Body.String(), "Request failed") + assert.NotContains(t, w.Body.String(), "poll secret") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Request failed") + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Request failed") + assert.NotContains(t, w.Body.String(), "not an image") + assert.NotContains(t, w.Body.String(), imageServer.URL) + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Request failed") + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Request failed") + assert.False(t, targetCalled) + assert.NotContains(t, w.Body.String(), targetServer.URL) +} + +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) + 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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusRequestTimeout, w.Code) + assert.Contains(t, w.Body.String(), "Request timed out") + 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) { + 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":"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() + + 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(), "Request failed") +} + +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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusBadGateway, w.Code) + assert.Contains(t, w.Body.String(), "Request failed") + assert.NotContains(t, w.Body.String(), "failed secret") + assert.Contains(t, logBuf.String(), "response_body_masked=true") + assert.Contains(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":"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() + + 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.Contains(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":"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() + + prx.ProxyRequest(w, req) + + assert.Equal(t, http.StatusRequestTimeout, w.Code) + assert.Contains(t, w.Body.String(), "Request timed out") + 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} + +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 { + *auths = append(*auths, r.Header.Get("Authorization")) + } + if contentType != "" { + w.Header().Set("Content-Type", contentType) + } + w.WriteHeader(status) + _, _ = w.Write(body) + })) +} + +func allowPrivateSosanaResultURLsForTest(t *testing.T) { + t.Helper() + + restore := sosana.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 && sosana.IsUnsafeResultIP(ip) + }) + t.Cleanup(restore) +} + +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 setSosanaTestModels(prx *Proxy, modelConfigs []config.ModelRPMConfig) { + prx.modelManager = aimodels.New(prx.logger, 50, modelConfigs) + prx.modelManager.LoadModelsFromConfig(prx.balancer.GetCredentialsSnapshot()) + prx.balancer.SetModelChecker(prx.modelManager) +} + +func (b *TestProxyBuilder) withLogger(logger *slog.Logger) *TestProxyBuilder { + b.config.Logger = logger + b.config.TokenManager = createTestTokenManager(logger) + 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() +} + +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) SpendLoggingEnabled() bool { + return true +} + +func (m *capturedSpendManager) LogSpend(entry *litellmmodels.SpendLogEntry) error { + m.entries = append(m.entries, entry) + return nil +} diff --git a/internal/proxy/upstream_masking.go b/internal/proxy/upstream_masking.go new file mode 100644 index 0000000..7b83ad9 --- /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 0000000..7efe990 --- /dev/null +++ b/internal/proxy/upstream_masking_test.go @@ -0,0 +1,103 @@ +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, + }). + 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, "Rate limit exceeded", got.Error.Message) + assert.Equal(t, "rate_limit_error", got.Error.Type) + require.NotNil(t, got.Error.Code) + assert.Equal(t, "rate_limit_error", *got.Error.Code) +} + +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(), "Request failed") +} + +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(), "Request failed") +}