diff --git a/internal/channel/compiler.go b/internal/channel/compiler.go index b41334c5..8f6ed77d 100644 --- a/internal/channel/compiler.go +++ b/internal/channel/compiler.go @@ -395,7 +395,7 @@ func compileRoutes( func validProtocolOperation(clientProtocol protocol.Protocol, operation execution.Operation) bool { switch operation { case execution.OperationChatCompletion: - return clientProtocol != protocol.OpenAIResponses + return clientProtocol != protocol.OpenAIResponses && clientProtocol != protocol.OpenAIImages case execution.OperationCountTokens: return clientProtocol == protocol.Anthropic || clientProtocol == protocol.Gemini case execution.OperationResponsesCreate, @@ -407,10 +407,13 @@ func validProtocolOperation(clientProtocol protocol.Protocol, operation executio execution.OperationResponsesInputTokens, execution.OperationResponsesPassthrough: return clientProtocol == protocol.OpenAIResponses + case execution.OperationImagesGenerate, + execution.OperationImagesEdit: + return clientProtocol == protocol.OpenAIImages case execution.OperationListModels: - return clientProtocol != protocol.OpenAIResponses + return clientProtocol != protocol.OpenAIResponses && clientProtocol != protocol.OpenAIImages case execution.OperationProbe: - return true + return clientProtocol != protocol.OpenAIImages default: return false } diff --git a/internal/channel/compiler_test.go b/internal/channel/compiler_test.go index 0f1cb197..0cd1ff29 100644 --- a/internal/channel/compiler_test.go +++ b/internal/channel/compiler_test.go @@ -101,6 +101,26 @@ func TestCompilerRejectsOpenAIResponsesModelListRoute(t *testing.T) { } } +func TestValidProtocolOperationOpenAIImagesMatrix(t *testing.T) { + t.Parallel() + + for _, operation := range []execution.Operation{ + execution.OperationImagesGenerate, + execution.OperationImagesEdit, + } { + if !validProtocolOperation(protocol.OpenAIImages, operation) { + t.Fatalf("openai-images/%s must be valid", operation) + } + if validProtocolOperation(protocol.OpenAIResponses, operation) { + t.Fatalf("openai-responses/%s must be invalid", operation) + } + } + if validProtocolOperation(protocol.OpenAIImages, execution.OperationChatCompletion) || + validProtocolOperation(protocol.OpenAIImages, execution.OperationResponsesCreate) { + t.Fatal("openai-images accepted a non-images operation") + } +} + func TestCompilerDefaultsEmptyIconToChannelID(t *testing.T) { t.Parallel() diff --git a/internal/channel/modules/codex.go b/internal/channel/modules/codex.go index b9ca17d8..29279e92 100644 --- a/internal/channel/modules/codex.go +++ b/internal/channel/modules/codex.go @@ -38,6 +38,8 @@ func Codex() spec.Module { }, Routes: []spec.Route{ spec.NewRoute(protocol.OpenAICompletions, execution.OperationChatCompletion, execution.RouteConverted), + spec.NewRoute(protocol.OpenAIImages, execution.OperationImagesGenerate, execution.RouteNative), + spec.NewRoute(protocol.OpenAIImages, execution.OperationImagesEdit, execution.RouteNative), spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesCreate, execution.RouteNative), spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesInputTokens, execution.RouteNative), spec.NewRoute(protocol.Anthropic, execution.OperationChatCompletion, execution.RouteConverted), diff --git a/internal/channel/modules/openai.go b/internal/channel/modules/openai.go index f8758bfb..a816e29b 100644 --- a/internal/channel/modules/openai.go +++ b/internal/channel/modules/openai.go @@ -34,6 +34,8 @@ func OpenAI() spec.Module { }, Routes: []spec.Route{ spec.NewRoute(protocol.OpenAICompletions, execution.OperationChatCompletion, execution.RouteNative), + spec.NewRoute(protocol.OpenAIImages, execution.OperationImagesGenerate, execution.RouteNative), + spec.NewRoute(protocol.OpenAIImages, execution.OperationImagesEdit, execution.RouteNative), spec.NewRoute(protocol.OpenAICompletions, execution.OperationListModels, execution.RouteNative), spec.NewRoute(protocol.OpenAICompletions, execution.OperationProbe, execution.RouteNative), spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesCreate, execution.RouteNative), diff --git a/internal/channel/modules/openai_compatible.go b/internal/channel/modules/openai_compatible.go index 34de8957..3db66409 100644 --- a/internal/channel/modules/openai_compatible.go +++ b/internal/channel/modules/openai_compatible.go @@ -33,6 +33,8 @@ func OpenAICompatible() spec.Module { }, Routes: []spec.Route{ spec.NewRoute(protocol.OpenAICompletions, execution.OperationChatCompletion, execution.RouteNative), + spec.NewRoute(protocol.OpenAIImages, execution.OperationImagesGenerate, execution.RouteNative), + spec.NewRoute(protocol.OpenAIImages, execution.OperationImagesEdit, execution.RouteNative), spec.NewRoute(protocol.OpenAICompletions, execution.OperationListModels, execution.RouteNative), spec.NewRoute(protocol.OpenAICompletions, execution.OperationProbe, execution.RouteNative), spec.NewRoute(protocol.OpenAIResponses, execution.OperationResponsesCreate, execution.RouteConverted), diff --git a/internal/container/container.go b/internal/container/container.go index d1c6f68c..ee746b29 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -148,15 +148,17 @@ func BuildContainer() (*dig.Container, error) { }, dialect.NewOpenAI, dialect.NewOpenAIResponses, + dialect.NewOpenAIImages, dialect.NewAnthropic, dialect.NewGemini, func( openAI *dialect.OpenAI, openAIResponses *dialect.OpenAIResponses, + openAIImages *dialect.OpenAIImages, anthropic *dialect.Anthropic, gemini *dialect.Gemini, ) dialect.Set { - return dialect.NewSet(openAI, openAIResponses, anthropic, gemini) + return dialect.NewSet(openAI, openAIResponses, openAIImages, anthropic, gemini) }, func(registry *channel.Registry) (*bifrostexecutor.RuntimeManager, error) { return bifrostexecutor.NewManagedRuntime(registry) diff --git a/internal/container/container_test.go b/internal/container/container_test.go index 351405b3..6b2afaf4 100644 --- a/internal/container/container_test.go +++ b/internal/container/container_test.go @@ -445,6 +445,7 @@ func TestBuildContainerResolvesAllDialects(t *testing.T) { err = dependencyContainer.Invoke(func( openAI *dialect.OpenAI, openAIResponses *dialect.OpenAIResponses, + openAIImages *dialect.OpenAIImages, anthropic *dialect.Anthropic, gemini *dialect.Gemini, values dialect.Set, @@ -458,8 +459,9 @@ func TestBuildContainerResolvesAllDialects(t *testing.T) { }) if values[protocol.OpenAICompletions] != openAI || values[protocol.OpenAIResponses] != openAIResponses || + values[protocol.OpenAIImages] != openAIImages || values[protocol.Anthropic] != anthropic || - values[protocol.Gemini] != gemini || len(values) != 4 { + values[protocol.Gemini] != gemini || len(values) != 5 { t.Fatalf("dialect Set = %#v", values) } }) diff --git a/internal/control/access_keys_test.go b/internal/control/access_keys_test.go index 6113d8d5..7caf24d5 100644 --- a/internal/control/access_keys_test.go +++ b/internal/control/access_keys_test.go @@ -153,6 +153,7 @@ func TestAccessKeyCreateAcceptsAllEnabledProtocolsInCanonicalOrder(t *testing.T) Filters: &AccessKeyFilters{ Protocols: []protocol.Protocol{ protocol.Gemini, + protocol.OpenAIImages, protocol.OpenAIResponses, protocol.Anthropic, protocol.OpenAICompletions, diff --git a/internal/control/route_inspect_test.go b/internal/control/route_inspect_test.go index ce8a9ca7..d51ebea3 100644 --- a/internal/control/route_inspect_test.go +++ b/internal/control/route_inspect_test.go @@ -232,14 +232,16 @@ func TestRouteInspectDerivesStandardRequestMetadataFromProtocol(t *testing.T) { NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) tests := []struct { - protocol protocol.Protocol - operation execution.Operation - routeMode execution.RouteMode + protocol protocol.Protocol + operation execution.Operation + routeMode execution.RouteMode + routeRequirement execution.RouteRequirement }{ - {protocol: protocol.OpenAICompletions, operation: execution.OperationChatCompletion, routeMode: execution.RouteNative}, - {protocol: protocol.OpenAIResponses, operation: execution.OperationResponsesCreate, routeMode: execution.RouteNative}, - {protocol: protocol.Anthropic, operation: execution.OperationChatCompletion, routeMode: execution.RouteConverted}, - {protocol: protocol.Gemini, operation: execution.OperationChatCompletion, routeMode: execution.RouteConverted}, + {protocol: protocol.OpenAICompletions, operation: execution.OperationChatCompletion, routeMode: execution.RouteNative, routeRequirement: execution.RouteRequirementAny}, + {protocol: protocol.OpenAIResponses, operation: execution.OperationResponsesCreate, routeMode: execution.RouteNative, routeRequirement: execution.RouteRequirementAny}, + {protocol: protocol.OpenAIImages, operation: execution.OperationImagesGenerate, routeMode: execution.RouteNative, routeRequirement: execution.RouteRequirementNative}, + {protocol: protocol.Anthropic, operation: execution.OperationChatCompletion, routeMode: execution.RouteConverted, routeRequirement: execution.RouteRequirementAny}, + {protocol: protocol.Gemini, operation: execution.OperationChatCompletion, routeMode: execution.RouteConverted, routeRequirement: execution.RouteRequirementAny}, } for _, test := range tests { t.Run(string(test.protocol), func(t *testing.T) { @@ -250,7 +252,7 @@ func TestRouteInspectDerivesStandardRequestMetadataFromProtocol(t *testing.T) { ) got := decodeRouteInspectSuccess(t, recorder) if got.Protocol != test.protocol || got.Operation != test.operation || - got.RouteRequirement != execution.RouteRequirementAny || + got.RouteRequirement != test.routeRequirement || routeModelValue(got.ExternalModel) != "public" || len(got.Groups) != 1 || got.Groups[0].RouteMode != test.routeMode { t.Fatalf("route inspection = %#v", got) diff --git a/internal/dialect/openai_images.go b/internal/dialect/openai_images.go new file mode 100644 index 00000000..ca4d0d0e --- /dev/null +++ b/internal/dialect/openai_images.go @@ -0,0 +1,223 @@ +package dialect + +import ( + "bytes" + "encoding/json" + "fmt" + "mime" + "net/http" + "strings" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +const ( + openAIImagesGenerationsPath = "/v1/images/generations" + openAIImagesEditsPath = "/v1/images/edits" +) + +// OpenAIImages implements the OpenAI Images wire contract. +type OpenAIImages struct{} + +var ( + _ Dialect = (*OpenAIImages)(nil) + _ ModelRewriter = (*OpenAIImages)(nil) + _ StreamEventClassifier = (*OpenAIImages)(nil) +) + +func NewOpenAIImages() *OpenAIImages { + return &OpenAIImages{} +} + +func (*OpenAIImages) Protocol() protocol.Protocol { + return protocol.OpenAIImages +} + +func (d *OpenAIImages) InspectRequest(request *ParsedRequest) (RequestMetadata, error) { + if request == nil { + return RequestMetadata{}, fmt.Errorf("parsed request is required") + } + if request.Method != http.MethodPost { + return RequestMetadata{}, fmt.Errorf("%s only supports POST", d.Protocol()) + } + operation := execution.Operation("") + switch request.Path { + case openAIImagesGenerationsPath: + operation = execution.OperationImagesGenerate + case openAIImagesEditsPath: + operation = execution.OperationImagesEdit + default: + return RequestMetadata{}, fmt.Errorf("unsupported %s path %q", d.Protocol(), request.Path) + } + contentType := request.Header.Get("Content-Type") + mediaType := "" + if strings.TrimSpace(contentType) != "" { + parsedMediaType, _, parseErr := mime.ParseMediaType(contentType) + if parseErr != nil { + return RequestMetadata{}, fmt.Errorf("decode %s Content-Type: %w", d.Protocol(), parseErr) + } + mediaType = strings.ToLower(parsedMediaType) + } + if operation == execution.OperationImagesEdit && mediaType == "multipart/form-data" { + parsed, err := processOpenAIImagesMultipart(request.Body, contentType, nil) + if err != nil { + return RequestMetadata{}, fmt.Errorf("decode %s request: %w", d.Protocol(), err) + } + model := parsed.model + return RequestMetadata{ + Model: &model, Stream: parsed.stream, + Operation: operation, RouteRequirement: execution.RouteRequirementNative, + ObserveUsage: true, + }, nil + } + if mediaType != "" && mediaType != "application/json" { + return RequestMetadata{}, fmt.Errorf("unsupported %s Content-Type %q", d.Protocol(), contentType) + } + metadata, err := inspectJSONRequestFields(request.Body, true) + if err != nil { + return RequestMetadata{}, fmt.Errorf("decode %s request: %w", d.Protocol(), err) + } + metadata.Operation = operation + metadata.RouteRequirement = execution.RouteRequirementNative + metadata.ObserveUsage = true + return metadata, nil +} + +func (d *OpenAIImages) RewriteRequestModel(request *ParsedRequest, model string) (*ParsedRequest, error) { + if err := validateModelRewriteTarget(model, false); err != nil { + return nil, err + } + if request != nil && request.Path == openAIImagesEditsPath { + contentType := request.Header.Get("Content-Type") + mediaType, _, _ := mime.ParseMediaType(contentType) + if strings.EqualFold(mediaType, "multipart/form-data") { + result, err := processOpenAIImagesMultipart(request.Body, contentType, &model) + if err != nil { + return nil, fmt.Errorf("rewrite %s multipart request: %w", d.Protocol(), err) + } + return applyRebuiltImagesMultipart(request, result) + } + } + return rewriteJSONRequestModel(request, model, string(d.Protocol())) +} + +// SanitizeRequestForAttempt rebuilds one Images request for the selected +// upstream model and actual execution mode. The returned request is owned by +// the caller; the original client representation is never mutated. +func (d *OpenAIImages) SanitizeRequestForAttempt( + request *ParsedRequest, + model string, + stream bool, +) (*ParsedRequest, error) { + if err := validateModelRewriteTarget(model, false); err != nil { + return nil, err + } + if _, err := d.InspectRequest(request); err != nil { + return nil, err + } + contentType := request.Header.Get("Content-Type") + mediaType, _, _ := mime.ParseMediaType(contentType) + if request.Path == openAIImagesEditsPath && strings.EqualFold(mediaType, "multipart/form-data") { + result, err := processOpenAIImagesMultipartWithOptions( + request.Body, + contentType, + openAIImagesMultipartOptions{ + replacementModel: &model, + forceStream: &stream, + stripControls: true, + }, + ) + if err != nil { + return nil, fmt.Errorf("sanitize %s multipart request: %w", d.Protocol(), err) + } + return applyRebuiltImagesMultipart(request, result) + } + + object, err := decodeJSONObject(request.Body) + if err != nil { + return nil, fmt.Errorf("sanitize %s request: %w", d.Protocol(), err) + } + for field := range object { + if openAIImagesSecurityControlName(field) { + delete(object, field) + } + } + encodedModel, err := json.Marshal(model) + if err != nil { + return nil, fmt.Errorf("encode selected Images model: %w", err) + } + object["model"] = encodedModel + if _, exists := object["stream"]; exists || stream { + object["stream"] = json.RawMessage(fmt.Sprintf("%t", stream)) + } + body, err := json.Marshal(object) + if err != nil { + return nil, fmt.Errorf("encode sanitized %s request: %w", d.Protocol(), err) + } + clone, err := cloneParsedRequest(request) + if err != nil { + return nil, err + } + clone.Body = body + if clone.Header == nil { + clone.Header = make(http.Header) + } + if strings.TrimSpace(clone.Header.Get("Content-Type")) == "" { + clone.Header.Set("Content-Type", "application/json") + } + return clone, nil +} + +func (*OpenAIImages) RewriteResponseModel(body []byte, model string) ([]byte, error) { + if err := validateModelRewriteTarget(model, false); err != nil { + return nil, err + } + return rewriteOptionalJSONField(body, "model", model) +} + +func (*OpenAIImages) RequiresTerminalEvent() bool { + return true +} + +func (*OpenAIImages) ClassifyStreamEvent(event StreamEvent) (StreamEventClassification, error) { + payload := bytes.TrimSpace(event.Payload) + if bytes.Equal(payload, []byte("null")) { + return StreamEventClassification{}, fmt.Errorf("decode OpenAI Images stream event") + } + var object struct { + Type json.RawMessage `json:"type"` + Error json.RawMessage `json:"error"` + } + if err := json.Unmarshal(payload, &object); err != nil { + return StreamEventClassification{}, fmt.Errorf("decode OpenAI Images stream event") + } + + payloadType := "" + if len(bytes.TrimSpace(object.Type)) > 0 { + if err := json.Unmarshal(object.Type, &payloadType); err != nil || payloadType == "" { + return StreamEventClassification{}, fmt.Errorf("decode OpenAI Images stream event type") + } + } + if event.Name != "" && payloadType != "" && event.Name != payloadType { + return StreamEventClassification{}, fmt.Errorf( + "OpenAI Images stream event name %q conflicts with payload type %q", + event.Name, + payloadType, + ) + } + eventType := event.Name + if eventType == "" { + eventType = payloadType + } + if meaningfulJSONValue(object.Error) { + return StreamEventClassification{Disposition: StreamEventFailed}, nil + } + switch eventType { + case "image_generation.completed", "image_edit.completed": + return StreamEventClassification{Disposition: StreamEventCompleted}, nil + case "image_generation.failed", "image_edit.failed", "error": + return StreamEventClassification{Disposition: StreamEventFailed}, nil + } + return StreamEventClassification{Disposition: StreamEventContinue}, nil +} diff --git a/internal/dialect/openai_images_multipart.go b/internal/dialect/openai_images_multipart.go new file mode 100644 index 00000000..a4301f61 --- /dev/null +++ b/internal/dialect/openai_images_multipart.go @@ -0,0 +1,319 @@ +package dialect + +import ( + "bytes" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "strconv" + "strings" +) + +const ( + maxOpenAIImagesMultipartBodyBytes = 128 << 20 + maxOpenAIImagesMultipartParts = 128 + maxOpenAIImagesMultipartHeadersPerPart = 64 + maxOpenAIImagesMultipartHeaderBytes = 16 << 10 + maxOpenAIImagesMultipartTotalHeaderBytes = 256 << 10 + maxOpenAIImagesContentDispositionBytes = 4 << 10 + maxOpenAIImagesModelFieldBytes = 4 << 10 + maxOpenAIImagesStreamFieldBytes = 16 +) + +var rewrittenMultipartIntegrityHeaders = []string{ + "Content-Length", + "Content-Transfer-Encoding", + "ETag", + "Digest", + "Content-MD5", + "Content-Range", + "Content-Digest", + "Repr-Digest", + "Signature", + "Signature-Input", +} + +type openAIImagesMultipartResult struct { + model string + stream bool + body []byte + contentType string +} + +type openAIImagesMultipartOptions struct { + replacementModel *string + forceStream *bool + stripControls bool +} + +func processOpenAIImagesMultipart( + body []byte, + contentType string, + replacementModel *string, +) (openAIImagesMultipartResult, error) { + return processOpenAIImagesMultipartWithOptions(body, contentType, openAIImagesMultipartOptions{ + replacementModel: replacementModel, + }) +} + +func processOpenAIImagesMultipartWithOptions( + body []byte, + contentType string, + options openAIImagesMultipartOptions, +) (openAIImagesMultipartResult, error) { + return processOpenAIImagesMultipartWithOptionsAndLimit( + body, + contentType, + options, + maxOpenAIImagesMultipartBodyBytes, + ) +} + +func processOpenAIImagesMultipartWithOptionsAndLimit( + body []byte, + contentType string, + options openAIImagesMultipartOptions, + maxBodyBytes int, +) (openAIImagesMultipartResult, error) { + if maxBodyBytes <= 0 { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart body limit must be positive") + } + if len(body) > maxBodyBytes { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart body exceeds limit") + } + mediaType, parameters, err := mime.ParseMediaType(contentType) + if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") { + return openAIImagesMultipartResult{}, fmt.Errorf("decode multipart Content-Type") + } + boundary := strings.TrimSpace(parameters["boundary"]) + if boundary == "" { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart boundary is required") + } + + reader := multipart.NewReader(bytes.NewReader(body), boundary) + var output bytes.Buffer + var writer *multipart.Writer + if options.replacementModel != nil || options.forceStream != nil || options.stripControls { + writer = multipart.NewWriter(&output) + } + + result := openAIImagesMultipartResult{} + modelSeen := false + streamSeen := false + totalHeaderBytes := 0 + partCount := 0 + for { + part, nextErr := reader.NextRawPart() + if nextErr == io.EOF { + break + } + if nextErr != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("decode multipart part: %w", nextErr) + } + partCount++ + if partCount > maxOpenAIImagesMultipartParts { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart part count exceeds limit") + } + if err := validateOpenAIImagesPartHeaders(part.Header, &totalHeaderBytes); err != nil { + return openAIImagesMultipartResult{}, err + } + name, file, err := openAIImagesPartIdentity(part.Header) + if err != nil { + return openAIImagesMultipartResult{}, err + } + if options.stripControls && openAIImagesSecurityControlName(name) { + if _, err := io.Copy(io.Discard, part); err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("discard multipart control part: %w", err) + } + continue + } + + control := name == "model" || name == "stream" + var partBody []byte + if control { + if file { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart %s field must not be a file", name) + } + limit := maxOpenAIImagesModelFieldBytes + if name == "stream" { + limit = maxOpenAIImagesStreamFieldBytes + } + partBody, err = readOpenAIImagesControlPart(part, limit) + if err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("decode multipart %s field: %w", name, err) + } + } + + switch name { + case "model": + if modelSeen { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart model field must be unique") + } + modelSeen = true + result.model = string(partBody) + if result.model == "" || strings.TrimSpace(result.model) != result.model { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart model must be non-empty without boundary whitespace") + } + case "stream": + if streamSeen { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart stream field must be unique") + } + streamSeen = true + switch string(partBody) { + case "true": + result.stream = true + case "false": + result.stream = false + default: + return openAIImagesMultipartResult{}, fmt.Errorf("multipart stream must be true or false") + } + } + + if writer == nil { + if !control { + if _, err := io.Copy(io.Discard, part); err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("read multipart part: %w", err) + } + } + continue + } + + header := cloneImagesMIMEHeader(part.Header) + if name == "model" && options.replacementModel != nil { + for _, headerName := range rewrittenMultipartIntegrityHeaders { + header.Del(headerName) + } + partBody = []byte(*options.replacementModel) + } + if name == "stream" && options.forceStream != nil { + desired := strconv.FormatBool(*options.forceStream) + if string(partBody) != desired { + for _, headerName := range rewrittenMultipartIntegrityHeaders { + header.Del(headerName) + } + partBody = []byte(desired) + } + } + outputPart, err := writer.CreatePart(header) + if err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("rebuild multipart part: %w", err) + } + if control { + if _, err := outputPart.Write(partBody); err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("write multipart control part: %w", err) + } + } else if _, err := io.Copy(outputPart, part); err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("copy multipart part: %w", err) + } + } + if !modelSeen { + return openAIImagesMultipartResult{}, fmt.Errorf("multipart model field is required") + } + if writer != nil { + if options.forceStream != nil && *options.forceStream && !streamSeen { + part, err := writer.CreateFormField("stream") + if err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("add multipart stream field: %w", err) + } + if _, err := io.WriteString(part, "true"); err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("write multipart stream field: %w", err) + } + } + if err := writer.Close(); err != nil { + return openAIImagesMultipartResult{}, fmt.Errorf("close rebuilt multipart body: %w", err) + } + if output.Len() > maxBodyBytes { + return openAIImagesMultipartResult{}, fmt.Errorf("rebuilt multipart body exceeds limit") + } + result.body = output.Bytes() + result.contentType = writer.FormDataContentType() + } + return result, nil +} + +func openAIImagesSecurityControlName(name string) bool { + normalized := strings.ReplaceAll(strings.ToLower(name), "-", "_") + switch normalized { + case "provider", "fallback", "fallbacks", "authorization", "proxy_authorization", + "api_key", "apikey", "x_api_key", "x_goog_api_key": + return true + default: + return false + } +} + +func validateOpenAIImagesPartHeaders(header textproto.MIMEHeader, totalBytes *int) error { + count := 0 + partBytes := 0 + for name, values := range header { + count += len(values) + for _, value := range values { + partBytes += len(name) + len(value) + 4 + } + } + if count > maxOpenAIImagesMultipartHeadersPerPart { + return fmt.Errorf("multipart part header count exceeds limit") + } + if partBytes > maxOpenAIImagesMultipartHeaderBytes { + return fmt.Errorf("multipart part header bytes exceed limit") + } + *totalBytes += partBytes + if *totalBytes > maxOpenAIImagesMultipartTotalHeaderBytes { + return fmt.Errorf("multipart total header bytes exceed limit") + } + dispositionValues := header.Values("Content-Disposition") + if len(dispositionValues) != 1 || len(dispositionValues[0]) > maxOpenAIImagesContentDispositionBytes { + return fmt.Errorf("multipart Content-Disposition is missing, repeated, or exceeds limit") + } + return nil +} + +func openAIImagesPartIdentity(header textproto.MIMEHeader) (string, bool, error) { + disposition, parameters, err := mime.ParseMediaType(header.Get("Content-Disposition")) + if err != nil || !strings.EqualFold(disposition, "form-data") { + return "", false, fmt.Errorf("multipart Content-Disposition is invalid") + } + name := parameters["name"] + if name == "" { + return "", false, fmt.Errorf("multipart part name is required") + } + _, hasFilename := parameters["filename"] + return name, hasFilename, nil +} + +func readOpenAIImagesControlPart(part *multipart.Part, limit int) ([]byte, error) { + reader := io.LimitReader(part, int64(limit)+1) + body, err := io.ReadAll(reader) + if err != nil { + return nil, err + } + if len(body) > limit { + return nil, fmt.Errorf("field exceeds %d bytes", limit) + } + return body, nil +} + +func cloneImagesMIMEHeader(source textproto.MIMEHeader) textproto.MIMEHeader { + return textproto.MIMEHeader(http.Header(source).Clone()) +} + +func applyRebuiltImagesMultipart(request *ParsedRequest, result openAIImagesMultipartResult) (*ParsedRequest, error) { + if request == nil { + return nil, fmt.Errorf("parsed request is required") + } + clone := *request + clone.Header = request.Header.Clone() + clone.Body = result.body + if clone.Header == nil { + clone.Header = make(http.Header) + } + for _, headerName := range rewrittenMultipartIntegrityHeaders { + clone.Header.Del(headerName) + } + clone.Header.Set("Content-Type", result.contentType) + clone.Header.Set("Content-Length", strconv.Itoa(len(clone.Body))) + return &clone, nil +} diff --git a/internal/dialect/openai_images_multipart_test.go b/internal/dialect/openai_images_multipart_test.go new file mode 100644 index 00000000..cae4b97f --- /dev/null +++ b/internal/dialect/openai_images_multipart_test.go @@ -0,0 +1,245 @@ +package dialect + +import ( + "bytes" + "io" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "strconv" + "strings" + "testing" +) + +type multipartTestPart struct { + header textproto.MIMEHeader + body string +} + +func TestOpenAIImagesMultipartInspectAndRewrite(t *testing.T) { + t.Parallel() + + modelHeader := textproto.MIMEHeader{ + "Content-Disposition": {`form-data; name="model"`}, + "Content-Type": {"text/plain"}, + "Content-Length": {"12"}, + "Content-Transfer-Encoding": {"quoted-printable"}, + "ETag": {`"model-v1"`}, + "Digest": {"sha-256=model"}, + "Content-MD5": {"model-md5"}, + "Content-Range": {"bytes 0-11/12"}, + "Content-Digest": {"sha-256=:bW9kZWw=:"}, + "Repr-Digest": {"sha-256=:cmVwcg==:"}, + "Signature": {"sig=:c2ln:"}, + "Signature-Input": {`sig=("content-digest")`}, + } + body, contentType := buildMultipartForTest(t, []multipartTestPart{ + {header: textproto.MIMEHeader{ + "Content-Disposition": {`form-data; name="image"; filename="../original.png"`}, + "Content-Type": {"image/png"}, + "X-Future-Header": {"keep-me"}, + "Content-Transfer-Encoding": {"quoted-printable"}, + }, body: "=41=42=43"}, + {header: modelHeader, body: "public-image"}, + {header: textproto.MIMEHeader{ + "Content-Disposition": {`form-data; name="stream"`}, + }, body: "true"}, + {header: textproto.MIMEHeader{ + "Content-Disposition": {`form-data; name="future"`}, + "X-Unknown": {"one", "two"}, + }, body: "future-value"}, + }) + request := &ParsedRequest{ + Method: http.MethodPost, + Path: openAIImagesEditsPath, + Header: http.Header{ + "Content-Type": {contentType}, + "Content-Length": {strconv.Itoa(len(body))}, + "Digest": {"sha-256=outer"}, + }, + Body: body, + } + + metadata, err := NewOpenAIImages().InspectRequest(request) + if err != nil { + t.Fatalf("InspectRequest() error = %v", err) + } + if metadata.Model == nil || *metadata.Model != "public-image" || !metadata.Stream { + t.Fatalf("metadata = %#v", metadata) + } + + rewritten, err := NewOpenAIImages().RewriteRequestModel(request, "upstream-image") + if err != nil { + t.Fatalf("RewriteRequestModel() error = %v", err) + } + if bytes.Equal(rewritten.Body, request.Body) { + t.Fatal("multipart body was not rebuilt") + } + if rewritten.Header.Get("Content-Type") == contentType || + rewritten.Header.Get("Content-Length") != strconv.Itoa(len(rewritten.Body)) { + t.Fatalf("rewritten headers = %#v", rewritten.Header) + } + for _, name := range []string{ + "ETag", "Digest", "Content-MD5", "Content-Range", "Content-Digest", + "Repr-Digest", "Signature", "Signature-Input", + } { + if rewritten.Header.Get(name) != "" { + t.Fatalf("outer %s survived rewrite", name) + } + } + + parts := readMultipartForTest(t, rewritten.Body, rewritten.Header.Get("Content-Type")) + if len(parts) != 4 { + t.Fatalf("part count = %d", len(parts)) + } + if got := parts[0].header.Get("Content-Disposition"); got != `form-data; name="image"; filename="../original.png"` { + t.Fatalf("image Content-Disposition = %q", got) + } + if parts[0].header.Get("Content-Transfer-Encoding") != "quoted-printable" || + parts[0].header.Get("X-Future-Header") != "keep-me" || parts[0].body != "=41=42=43" { + t.Fatalf("image part = %#v", parts[0]) + } + if parts[1].body != "upstream-image" { + t.Fatalf("model body = %q", parts[1].body) + } + for _, name := range []string{ + "Content-Length", "Content-Transfer-Encoding", "ETag", "Digest", "Content-MD5", + "Content-Range", "Content-Digest", "Repr-Digest", "Signature", "Signature-Input", + } { + if parts[1].header.Get(name) != "" { + t.Fatalf("rewritten model %s survived", name) + } + } + if parts[2].body != "true" || parts[3].body != "future-value" || + len(parts[3].header.Values("X-Unknown")) != 2 { + t.Fatalf("unchanged parts = %#v", parts[2:]) + } + if string(request.Body) != string(body) || request.Header.Get("Content-Type") != contentType { + t.Fatal("rewrite mutated the original request") + } +} + +func TestOpenAIImagesMultipartRejectsRebuiltBodyAboveLimit(t *testing.T) { + body, contentType := buildMultipartForTest(t, []multipartTestPart{ + { + header: textproto.MIMEHeader{ + "Content-Disposition": {`form-data; name="model"`}, + }, + body: "public-image", + }, + { + header: textproto.MIMEHeader{ + "Content-Disposition": {`form-data; name="image"; filename="source.png"`}, + "Content-Type": {"image/png"}, + }, + body: "image-bytes", + }, + }) + forceStream := true + _, err := processOpenAIImagesMultipartWithOptionsAndLimit( + body, + contentType, + openAIImagesMultipartOptions{forceStream: &forceStream}, + len(body)+1, + ) + if err == nil || !strings.Contains(err.Error(), "rebuilt multipart body exceeds limit") { + t.Fatalf("processOpenAIImagesMultipartWithOptionsAndLimit() error = %v", err) + } +} + +func TestOpenAIImagesMultipartValidation(t *testing.T) { + t.Parallel() + + validModel := multipartTestPart{ + header: textproto.MIMEHeader{"Content-Disposition": {`form-data; name="model"`}}, + body: "gpt-image-2", + } + tests := []struct { + name string + parts []multipartTestPart + }{ + {name: "missing model", parts: []multipartTestPart{{header: textproto.MIMEHeader{"Content-Disposition": {`form-data; name="prompt"`}}, body: "draw"}}}, + {name: "duplicate model", parts: []multipartTestPart{validModel, validModel}}, + {name: "invalid stream", parts: []multipartTestPart{validModel, {header: textproto.MIMEHeader{"Content-Disposition": {`form-data; name="stream"`}}, body: "1"}}}, + {name: "model file", parts: []multipartTestPart{{header: textproto.MIMEHeader{"Content-Disposition": {`form-data; name="model"; filename="model.txt"`}}, body: "gpt-image-2"}}}, + {name: "oversized model", parts: []multipartTestPart{{header: textproto.MIMEHeader{"Content-Disposition": {`form-data; name="model"`}}, body: strings.Repeat("m", 4097)}}}, + {name: "too many parts", parts: repeatMultipartParts(validModel, 129)}, + {name: "oversized content disposition", parts: []multipartTestPart{{header: textproto.MIMEHeader{"Content-Disposition": {`form-data; name="` + strings.Repeat("m", 4096) + `"`}}, body: "gpt-image-2"}}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + body, contentType := buildMultipartForTest(t, test.parts) + metadata, err := NewOpenAIImages().InspectRequest(&ParsedRequest{ + Method: http.MethodPost, + Path: openAIImagesEditsPath, + Header: http.Header{"Content-Type": {contentType}}, + Body: body, + }) + if err == nil { + t.Fatalf("InspectRequest() = %#v, nil error", metadata) + } + }) + } +} + +func buildMultipartForTest(t *testing.T, parts []multipartTestPart) ([]byte, string) { + t.Helper() + var buffer bytes.Buffer + writer := multipart.NewWriter(&buffer) + if err := writer.SetBoundary("test-boundary"); err != nil { + t.Fatalf("SetBoundary() error = %v", err) + } + for _, part := range parts { + output, err := writer.CreatePart(part.header) + if err != nil { + t.Fatalf("CreatePart() error = %v", err) + } + if _, err := io.WriteString(output, part.body); err != nil { + t.Fatalf("write part: %v", err) + } + } + if err := writer.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + return bytes.Clone(buffer.Bytes()), writer.FormDataContentType() +} + +func readMultipartForTest(t *testing.T, body []byte, contentType string) []multipartTestPart { + t.Helper() + _, params, err := mime.ParseMediaType(contentType) + if err != nil { + t.Fatalf("ParseMediaType() error = %v", err) + } + reader := multipart.NewReader(bytes.NewReader(body), params["boundary"]) + parts := make([]multipartTestPart, 0) + for { + part, err := reader.NextRawPart() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("NextRawPart() error = %v", err) + } + data, err := io.ReadAll(part) + if err != nil { + t.Fatalf("ReadAll(part) error = %v", err) + } + parts = append(parts, multipartTestPart{ + header: textproto.MIMEHeader(http.Header(part.Header).Clone()), + body: string(data), + }) + } + return parts +} + +func repeatMultipartParts(part multipartTestPart, count int) []multipartTestPart { + parts := make([]multipartTestPart, count) + for index := range parts { + parts[index] = part + parts[index].header = textproto.MIMEHeader(http.Header(part.header).Clone()) + parts[index].header.Set("Content-Disposition", `form-data; name="part`+strconv.Itoa(index)+`"`) + } + return parts +} diff --git a/internal/dialect/openai_images_test.go b/internal/dialect/openai_images_test.go new file mode 100644 index 00000000..0a9a7727 --- /dev/null +++ b/internal/dialect/openai_images_test.go @@ -0,0 +1,178 @@ +package dialect + +import ( + "bytes" + "encoding/json" + "net/http" + "testing" + + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestOpenAIImagesInspectJSONRequests(t *testing.T) { + t.Parallel() + + selected := NewOpenAIImages() + tests := []struct { + name string + path string + body string + operation execution.Operation + stream bool + wantError bool + }{ + { + name: "generation", path: "/v1/images/generations", + body: `{"model":"gpt-image-2","prompt":"draw","quality":"auto"}`, + operation: execution.OperationImagesGenerate, + }, + { + name: "generation stream", path: "/v1/images/generations", + body: `{"model":"gpt-image-2","prompt":"draw","stream":true}`, + operation: execution.OperationImagesGenerate, stream: true, + }, + { + name: "JSON edit", path: "/v1/images/edits", + body: `{"model":"gpt-image-2","prompt":"edit","images":[{"image_url":"data:image/png;base64,AA=="}]}`, + operation: execution.OperationImagesEdit, + }, + { + name: "missing model", path: "/v1/images/generations", + body: `{"prompt":"draw"}`, wantError: true, + }, + { + name: "invalid stream", path: "/v1/images/generations", + body: `{"model":"gpt-image-2","stream":"true"}`, wantError: true, + }, + { + name: "trailing JSON", path: "/v1/images/generations", + body: `{"model":"gpt-image-2"}{}`, wantError: true, + }, + { + name: "unknown path", path: "/v1/images/variations", + body: `{"model":"gpt-image-2"}`, wantError: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + metadata, err := selected.InspectRequest(&ParsedRequest{ + Method: http.MethodPost, + Path: test.path, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: []byte(test.body), + }) + if test.wantError { + if err == nil { + t.Fatalf("InspectRequest() = %#v, nil error", metadata) + } + return + } + if err != nil { + t.Fatalf("InspectRequest() error = %v", err) + } + if metadata.Model == nil || *metadata.Model != "gpt-image-2" || + metadata.Operation != test.operation || metadata.Stream != test.stream || + metadata.RouteRequirement != execution.RouteRequirementNative || + !metadata.ObserveUsage { + t.Fatalf("metadata = %#v", metadata) + } + }) + } +} + +func TestOpenAIImagesModelRewritePreservesUnknownFields(t *testing.T) { + t.Parallel() + + selected := NewOpenAIImages() + original := &ParsedRequest{ + Method: http.MethodPost, + Path: "/v1/images/generations", + Header: http.Header{"Content-Type": {"application/json"}}, + Body: []byte(`{"model":"public-image","prompt":"draw","future":{"kept":true}}`), + } + rewritten, err := selected.RewriteRequestModel(original, "upstream-image") + if err != nil { + t.Fatalf("RewriteRequestModel() error = %v", err) + } + if bytes.Equal(rewritten.Body, original.Body) { + t.Fatal("RewriteRequestModel() did not rewrite the body") + } + var object map[string]json.RawMessage + if err := json.Unmarshal(rewritten.Body, &object); err != nil { + t.Fatalf("decode rewritten body: %v", err) + } + var model string + if err := json.Unmarshal(object["model"], &model); err != nil || model != "upstream-image" { + t.Fatalf("rewritten model = %q, %v", model, err) + } + if string(object["future"]) != `{"kept":true}` { + t.Fatalf("future field = %s", object["future"]) + } + if string(original.Body) != `{"model":"public-image","prompt":"draw","future":{"kept":true}}` { + t.Fatal("RewriteRequestModel() mutated the original") + } + + response := []byte(`{"created":1,"model":"upstream-image","data":[{"b64_json":"AA=="}]}`) + downstream, err := selected.RewriteResponseModel(response, "public-image") + if err != nil || !bytes.Contains(downstream, []byte(`"model":"public-image"`)) || + !bytes.Contains(downstream, []byte(`"b64_json":"AA=="`)) { + t.Fatalf("RewriteResponseModel() = %s, %v", downstream, err) + } +} + +func TestOpenAIImagesStreamLifecycle(t *testing.T) { + t.Parallel() + + selected := NewOpenAIImages() + if !selected.RequiresTerminalEvent() { + t.Fatal("RequiresTerminalEvent() = false") + } + tests := []struct { + name string + event StreamEvent + disposition StreamEventDisposition + wantError bool + }{ + {name: "generation partial", event: StreamEvent{Name: "image_generation.partial_image", Payload: []byte(`{"type":"image_generation.partial_image","b64_json":"AA=="}`)}}, + {name: "edit partial", event: StreamEvent{Payload: []byte(`{"type":"image_edit.partial_image","b64_json":"AA=="}`)}}, + {name: "generation completed", event: StreamEvent{Name: "image_generation.completed", Payload: []byte(`{"type":"image_generation.completed","b64_json":"AA=="}`)}, disposition: StreamEventCompleted}, + {name: "edit completed", event: StreamEvent{Payload: []byte(`{"type":"image_edit.completed","b64_json":"AA=="}`)}, disposition: StreamEventCompleted}, + {name: "completed with error", event: StreamEvent{Payload: []byte(`{"type":"image_generation.completed","error":{"message":"failed"}}`)}, disposition: StreamEventFailed}, + {name: "completed with null error", event: StreamEvent{Payload: []byte(`{"type":"image_generation.completed","error":null}`)}, disposition: StreamEventCompleted}, + {name: "completed with empty error", event: StreamEvent{Payload: []byte(`{"type":"image_generation.completed","error":{}}`)}, disposition: StreamEventCompleted}, + {name: "error event", event: StreamEvent{Payload: []byte(`{"type":"error","error":{"message":"failed"}}`)}, disposition: StreamEventFailed}, + {name: "null event", event: StreamEvent{Payload: []byte(`null`)}, wantError: true}, + {name: "name conflict", event: StreamEvent{Name: "image_generation.completed", Payload: []byte(`{"type":"image_edit.completed"}`)}, wantError: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + classification, err := selected.ClassifyStreamEvent(test.event) + if test.wantError { + if err == nil { + t.Fatalf("ClassifyStreamEvent() = %#v, nil error", classification) + } + return + } + if err != nil || classification.Disposition != test.disposition { + t.Fatalf("ClassifyStreamEvent() = %#v, %v", classification, err) + } + }) + } +} + +func TestOpenAIImagesStandardRequestUsesGeneration(t *testing.T) { + t.Parallel() + + metadata, err := InspectStandardRequest(protocol.OpenAIImages, "gpt-image-2") + if err != nil { + t.Fatalf("InspectStandardRequest() error = %v", err) + } + if metadata.Model == nil || *metadata.Model != "gpt-image-2" || + metadata.Operation != execution.OperationImagesGenerate || + metadata.RouteRequirement != execution.RouteRequirementNative { + t.Fatalf("metadata = %#v", metadata) + } +} diff --git a/internal/dialect/standard_request.go b/internal/dialect/standard_request.go index a8942fb2..f469466f 100644 --- a/internal/dialect/standard_request.go +++ b/internal/dialect/standard_request.go @@ -36,6 +36,10 @@ func standardRequest( selected = NewOpenAIResponses() request.Path = "/v1/responses" body = map[string]any{"model": model, "input": []any{}, "store": false} + case protocol.OpenAIImages: + selected = NewOpenAIImages() + request.Path = openAIImagesGenerationsPath + body = map[string]any{"model": model, "prompt": ""} case protocol.Anthropic: selected = NewAnthropic() request.Path = "/v1/messages" diff --git a/internal/execution/bifrost/capabilities.go b/internal/execution/bifrost/capabilities.go index 320b35d2..0a592bb6 100644 --- a/internal/execution/bifrost/capabilities.go +++ b/internal/execution/bifrost/capabilities.go @@ -61,12 +61,18 @@ func nativeRouteImplemented( ) bool { switch providerKind { case channel.ProviderOpenAI: + if clientProtocol == protocol.OpenAIImages { + return operation == execution.OperationImagesGenerate || operation == execution.OperationImagesEdit + } return nativeOpenAIProtocolOperation(clientProtocol, operation, true) case channel.ProviderAnthropic: return clientProtocol == protocol.Anthropic && standardProtocolOperation(clientProtocol, operation) case channel.ProviderGemini: return clientProtocol == protocol.Gemini && standardProtocolOperation(clientProtocol, operation) case channel.ProviderOpenAICompatible: + if clientProtocol == protocol.OpenAIImages { + return operation == execution.OperationImagesGenerate || operation == execution.OperationImagesEdit + } return clientProtocol == protocol.OpenAICompletions && standardProtocolOperation(clientProtocol, operation) case channel.ProviderDeepSeek: switch clientProtocol { diff --git a/internal/execution/bifrost/converted.go b/internal/execution/bifrost/converted.go index 927cabbe..5d88a1c7 100644 --- a/internal/execution/bifrost/converted.go +++ b/internal/execution/bifrost/converted.go @@ -647,6 +647,10 @@ func (r *Runtime) executeConvertedResponsesStream( return streamContextFailure(requestContext, true, true, headers, requestID, model, usageEvidence) case chunk, open := <-outcome.stream: idleTimer.pause() + if requestContext.Err() != nil { + callCancel() + return streamContextFailure(requestContext, false, true, headers, requestID, model, usageEvidence) + } if !open { if !sdkStreamEndedNormally(bifrostContext) { return terminatedStreamFailure(headers, requestID, model, usageEvidence) diff --git a/internal/execution/bifrost/executor.go b/internal/execution/bifrost/executor.go index f7038c66..c51dba56 100644 --- a/internal/execution/bifrost/executor.go +++ b/internal/execution/bifrost/executor.go @@ -57,6 +57,9 @@ type streamSDKResult struct { // Execute executes one non-streaming attempt. func (r *Runtime) Execute(parent context.Context, spec execution.AttemptSpec) (result execution.AttemptResult) { + defer func() { + normalizeImagesAttemptResult(spec, &result) + }() prepared, preflightError := r.prepare(spec, false) if preflightError != nil { return *preflightError @@ -164,6 +167,9 @@ func (r *Runtime) ExecuteStream( spec execution.AttemptSpec, sink execution.StreamSink, ) (result execution.StreamResult) { + defer func() { + normalizeImagesStreamResult(spec, &result) + }() if sink == nil { return notSentStreamFailure(execution.ErrorKindInvalidRequest, "stream sink is required") } @@ -258,6 +264,10 @@ func (r *Runtime) ExecuteStream( return streamContextFailure(requestContext, true, true, headers, requestID, model, usageEvidence) case chunk, open := <-outcome.stream: idleTimer.pause() + if requestContext.Err() != nil { + callCancel() + return streamContextFailure(requestContext, false, true, headers, requestID, model, usageEvidence) + } if !open { if !sdkStreamEndedNormally(bifrostContext) { return terminatedStreamFailure(headers, requestID, model, usageEvidence) @@ -418,9 +428,9 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte return preparedAttempt{}, &failure } provider := r.fixedConfig.provider - targetBaseURL := "" + customTargetBaseURL := "" if r.fixedConfig.custom { - targetBaseURL = r.fixedConfig.targetBaseURL + customTargetBaseURL = r.fixedConfig.targetBaseURL } if mode == channel.RouteNative && spec.Operation == execution.OperationListModels && !providerKindNativeForClient(providerKind, spec.ClientProtocol) { @@ -481,14 +491,14 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte responses := spec.ClientProtocol == protocol.OpenAIResponses typedURL, upstreamProtocol, targetErr := convertedTypedTarget( providerKind, - targetBaseURL, + customTargetBaseURL, spec.UpstreamModel, responses, false, "", ) if targetErr == nil && mode == channel.RouteNative && providerKind == channel.ProviderDeepSeek { - typedURL, upstreamProtocol, targetErr = deepSeekNativeTypedTarget(targetBaseURL, spec.ClientProtocol, "") + typedURL, upstreamProtocol, targetErr = deepSeekNativeTypedTarget(customTargetBaseURL, spec.ClientProtocol, "") } if targetErr != nil { failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid channel probe target") @@ -505,10 +515,15 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte } return prepared, nil } - if mode == channel.RouteNative && providerSupportsPassthrough(providerKind, targetBaseURL) { - body, err := sanitizeNativeRequestBody(spec, stream) + if mode == channel.RouteNative && providerSupportsPassthrough(providerKind, customTargetBaseURL, spec.ClientProtocol) { + body, sanitizedHeaders, err := sanitizeNativePassthroughRequest(spec, stream) if err != nil { failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid native request body") + if spec.ClientProtocol == protocol.OpenAIImages { + failure.Error.OriginHint = execution.ErrorOriginClient + failure.Error.ScopeHint = execution.ErrorScopeRequest + failure.Error.ReplaySafety = execution.ReplaySafetyUnknown + } return preparedAttempt{}, &failure } passthroughPath, err := nativePassthroughPath(spec, providerKind) @@ -516,7 +531,26 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid native request path") return preparedAttempt{}, &failure } - passthroughHeaders := safePassthroughHeaders(spec.Header) + passthroughHeaders := safePassthroughHeaders(sanitizedHeaders) + passthroughUpstreamURL := "" + if spec.ClientProtocol == protocol.OpenAIImages { + explicitPrefix, configured, prefixErr := targetBaseURL(resolved.TargetConfig) + if prefixErr != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid native request prefix") + failure.Error.OriginHint = execution.ErrorOriginClient + failure.Error.ScopeHint = execution.ErrorScopeRequest + failure.Error.ReplaySafety = execution.ReplaySafetyUnknown + return preparedAttempt{}, &failure + } + if providerKind == channel.ProviderOpenAICompatible || (providerKind == channel.ProviderOpenAI && configured) { + passthroughUpstreamURL = explicitPrefix + passthroughPath, err = openAIImagesPrefixPath(spec.Path) + if err != nil { + failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid native request path") + return preparedAttempt{}, &failure + } + } + } if providerKind == channel.ProviderGoogleVertex { if passthroughHeaders == nil { passthroughHeaders = make(map[string]string, 1) @@ -536,6 +570,7 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte Method: spec.Method, Path: passthroughPath, RawQuery: safeQuery, + UpstreamURL: passthroughUpstreamURL, Body: body, SafeHeaders: passthroughHeaders, }, @@ -544,7 +579,7 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte }, nil } if spec.Operation == execution.OperationListModels { - typedURL, upstreamProtocol, targetErr := convertedListModelsTarget(providerKind, targetBaseURL, safeQuery) + typedURL, upstreamProtocol, targetErr := convertedListModelsTarget(providerKind, customTargetBaseURL, safeQuery) if targetErr != nil { failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") return preparedAttempt{}, &failure @@ -574,7 +609,7 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte } typedURL, upstreamProtocol, targetErr := countTokensTypedTarget( providerKind, - targetBaseURL, + customTargetBaseURL, spec.UpstreamModel, safeQuery, ) @@ -604,9 +639,9 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, conversionErr.Error()) return preparedAttempt{}, &failure } - typedURL, upstreamProtocol, targetErr := convertedTypedTarget(providerKind, targetBaseURL, spec.UpstreamModel, true, stream, safeQuery) + typedURL, upstreamProtocol, targetErr := convertedTypedTarget(providerKind, customTargetBaseURL, spec.UpstreamModel, true, stream, safeQuery) if targetErr == nil && mode == channel.RouteNative && providerKind == channel.ProviderDeepSeek { - typedURL, upstreamProtocol, targetErr = deepSeekNativeTypedTarget(targetBaseURL, spec.ClientProtocol, safeQuery) + typedURL, upstreamProtocol, targetErr = deepSeekNativeTypedTarget(customTargetBaseURL, spec.ClientProtocol, safeQuery) } if targetErr != nil { failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") @@ -645,7 +680,7 @@ func (r *Runtime) prepare(spec execution.AttemptSpec, stream bool) (preparedAtte delete(request.Params.ExtraParams, "fallback") delete(request.Params.ExtraParams, "fallbacks") } - typedURL, upstreamProtocol, err := convertedTypedTarget(providerKind, targetBaseURL, spec.UpstreamModel, false, stream, safeQuery) + typedURL, upstreamProtocol, err := convertedTypedTarget(providerKind, customTargetBaseURL, spec.UpstreamModel, false, stream, safeQuery) if err != nil { failure := notSentUnaryFailure(execution.ErrorKindInvalidRequest, "invalid compatible channel target") return preparedAttempt{}, &failure @@ -680,11 +715,18 @@ func newProbeRequest( } } -func providerSupportsPassthrough(providerKind channel.ProviderKind, baseURL string) bool { +func providerSupportsPassthrough( + providerKind channel.ProviderKind, + baseURL string, + clientProtocol protocol.Protocol, +) bool { switch providerKind { case channel.ProviderOpenAI, channel.ProviderAnthropic, channel.ProviderGemini, channel.ProviderGoogleVertex: return true case channel.ProviderOpenAICompatible: + if clientProtocol == protocol.OpenAIImages { + return true + } // Bifrost v1.7.7's OpenAI passthrough always inserts /v1. A compatible // full API prefix with another suffix must use the typed custom path, // without changing the frozen route mode or channel capability. @@ -697,7 +739,8 @@ func providerSupportsPassthrough(providerKind channel.ProviderKind, baseURL stri func providerKindNativeForClient(providerKind channel.ProviderKind, clientProtocol protocol.Protocol) bool { switch providerKind { case channel.ProviderOpenAI, channel.ProviderOpenAICompatible: - return clientProtocol == protocol.OpenAICompletions || clientProtocol == protocol.OpenAIResponses + return clientProtocol == protocol.OpenAICompletions || clientProtocol == protocol.OpenAIResponses || + clientProtocol == protocol.OpenAIImages case channel.ProviderAnthropic: return clientProtocol == protocol.Anthropic case channel.ProviderGemini: @@ -760,6 +803,18 @@ func supportedRequestShape(spec execution.AttemptSpec, stream bool) bool { case execution.OperationResponsesPassthrough: return validResponsesPassthroughShape(spec, stream) } + case protocol.OpenAIImages: + if spec.RouteMode != execution.RouteNative || spec.Method != http.MethodPost { + return false + } + switch spec.Operation { + case execution.OperationImagesGenerate: + return spec.Path == "/v1/images/generations" + case execution.OperationImagesEdit: + return spec.Path == "/v1/images/edits" + default: + return false + } case protocol.Anthropic: return spec.Method == http.MethodPost && ((spec.Operation == execution.OperationChatCompletion && spec.Path == "/v1/messages") || @@ -783,6 +838,38 @@ func supportedRequestShape(spec execution.AttemptSpec, stream bool) bool { return false } +func openAIImagesPrefixPath(path string) (string, error) { + result, ok := strings.CutPrefix(path, "/v1") + if !ok || (result != "/images/generations" && result != "/images/edits") { + return "", fmt.Errorf("invalid OpenAI Images path") + } + return result, nil +} + +func normalizeImagesAttemptResult(spec execution.AttemptSpec, result *execution.AttemptResult) { + if result == nil || spec.ClientProtocol != protocol.OpenAIImages { + return + } + result.Usage = nil + if openAIResponseModel(result.Body, "") == "" { + result.Model = "" + } + if result.Error != nil && result.Error.ReplaySafety == "" { + result.Error.ReplaySafety = execution.ReplaySafetyUnknown + } +} + +func normalizeImagesStreamResult(spec execution.AttemptSpec, result *execution.StreamResult) { + if result == nil || spec.ClientProtocol != protocol.OpenAIImages { + return + } + result.Usage = nil + result.Model = "" + if result.Error != nil && result.Error.ReplaySafety == "" { + result.Error.ReplaySafety = execution.ReplaySafetyUnknown + } +} + func validResponsesPassthroughShape(spec execution.AttemptSpec, stream bool) bool { remainder, ok := strings.CutPrefix(spec.Path, "/v1/responses/") if !ok || remainder == "" { @@ -1081,7 +1168,10 @@ func (r *Runtime) newSDKContext(parent context.Context, spec execution.AttemptSp bifrostContext := schemas.NewBifrostContext(parent, schemas.NoDeadline) bifrostContext.SetValue(schemas.BifrostContextKeyRequestID, spec.RequestID) bifrostContext.SetValue(schemas.BifrostContextKeyDirectKey, directKey) - bifrostContext.SetValue(schemas.BifrostContextKeyLargeResponseThreshold, r.maxUnaryResponseBodyBytes) + bifrostContext.SetValue( + schemas.BifrostContextKeyLargeResponseThreshold, + r.unaryResponseBodyLimit(spec), + ) if spec.Operation == execution.OperationChatCompletion && r.providerKind(spec) == channel.ProviderDeepSeek && spec.ClientProtocol == protocol.Anthropic { diff --git a/internal/execution/bifrost/executor_test.go b/internal/execution/bifrost/executor_test.go index 965a4ce2..c1dcb4de 100644 --- a/internal/execution/bifrost/executor_test.go +++ b/internal/execution/bifrost/executor_test.go @@ -26,6 +26,17 @@ import ( const testAPIKey = "sk-test-do-not-leak" +func TestRuntimeUnaryResponseBodyLimitHonorsExplicitDefaultSizedOverride(t *testing.T) { + spec := execution.AttemptSpec{ClientProtocol: protocol.OpenAIImages} + + if got := (&Runtime{}).unaryResponseBodyLimit(spec); got != execution.OpenAIImagesUnaryResponseBodyLimitBytes { + t.Fatalf("unset Images unary limit = %d, want %d", got, execution.OpenAIImagesUnaryResponseBodyLimitBytes) + } + if got := (&Runtime{maxUnaryResponseBodyBytes: defaultMaxUnaryResponseBodyBytes}).unaryResponseBodyLimit(spec); got != defaultMaxUnaryResponseBodyBytes { + t.Fatalf("explicit Images unary limit = %d, want %d", got, defaultMaxUnaryResponseBodyBytes) + } +} + func TestStreamingSDKContextUsesLaterIdleGuard(t *testing.T) { runtime := &Runtime{} maximum := time.Duration(math.MaxInt64) diff --git a/internal/execution/bifrost/images_test.go b/internal/execution/bifrost/images_test.go new file mode 100644 index 00000000..19343a6a --- /dev/null +++ b/internal/execution/bifrost/images_test.go @@ -0,0 +1,476 @@ +package bifrost + +import ( + "bytes" + "context" + "encoding/json" + "io" + "mime" + "mime/multipart" + "net/http" + "net/http/httptest" + "net/textproto" + "strings" + "testing" + "time" + + "gpt-load/internal/channel" + "gpt-load/internal/execution" + "gpt-load/internal/protocol" +) + +func TestOpenAIImagesCompatibleUsesCompletePrefixAndSanitizesGeneration(t *testing.T) { + t.Parallel() + + const responseBody = `{"created":1,"data":[{"b64_json":"a2VlcC1tZQ==","url":"https://cdn.example/image.png?sig=sk-example-token"}],"unknown":{"keep":true}}` + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/tenant/api/v4/images/generations" { + t.Errorf("path = %q", request.URL.Path) + } + if request.URL.RawQuery != "vendor=one" { + t.Errorf("raw query = %q", request.URL.RawQuery) + } + if got := request.Header.Get("Authorization"); got != "Bearer "+testAPIKey { + t.Errorf("authorization = %q", got) + } + if got := request.Header.Get("X-Test-Header"); got != "keep" { + t.Errorf("X-Test-Header = %q", got) + } + for _, name := range []string{"Api-Key", "X-Api-Key", "Proxy-Authorization"} { + if got := request.Header.Get(name); got != "" { + t.Errorf("client credential header %s reached upstream: %q", name, got) + } + } + body, err := io.ReadAll(request.Body) + if err != nil { + t.Errorf("ReadAll(request.Body) error = %v", err) + return + } + var object map[string]json.RawMessage + if err := json.Unmarshal(body, &object); err != nil { + t.Errorf("request body = %s: %v", body, err) + return + } + if string(object["model"]) != `"provider-image"` { + t.Errorf("model = %s", object["model"]) + } + if raw, exists := object["stream"]; !exists || string(raw) != "false" { + t.Errorf("stream = %s, exists=%t", raw, exists) + } + for _, field := range []string{"provider", "fallback", "fallbacks", "authorization", "api_key", "x-api-key"} { + if _, exists := object[field]; exists { + t.Errorf("control field %q reached upstream", field) + } + } + if string(object["future_field"]) != `{"precise":1.2300}` { + t.Errorf("future_field = %s", object["future_field"]) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, responseBody) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, testRuntimeOptions{allowPrivateNetwork: true}) + spec := openAIImagesSpec( + channel.OpenAICompatible, + server.URL+"/tenant/api/v4", + execution.OperationImagesGenerate, + "/v1/images/generations", + "application/json", + []byte(`{"model":"public-image","stream":true,"prompt":"draw","provider":"client","fallback":"other","fallbacks":["other"],"authorization":"client","api_key":"client","x-api-key":"client","future_field":{"precise":1.2300}}`), + ) + spec.RawQuery = "vendor=one&api_key=client-secret" + spec.Query = nil + spec = execution.NewAttemptSpec(spec) + result := runtime.Execute(context.Background(), spec) + + if err := result.Validate(); err != nil { + t.Fatalf("result validation: %v; result=%+v", err, result) + } + if result.Error != nil || result.StatusCode != http.StatusOK || string(result.Body) != responseBody { + t.Fatalf("result = %+v body=%s", result, result.Body) + } + if result.Usage != nil { + t.Fatalf("Images usage = %#v, want nil", result.Usage) + } +} + +func TestNormalizeImagesResultsKeepOnlyObservedModel(t *testing.T) { + spec := execution.AttemptSpec{ClientProtocol: protocol.OpenAIImages} + + missing := execution.AttemptResult{ + Model: "provider-image", + Body: []byte(`{"created":1,"data":[{"b64_json":"AA=="}]}`), + } + normalizeImagesAttemptResult(spec, &missing) + if missing.Model != "" { + t.Fatalf("missing response model = %q, want empty", missing.Model) + } + + observed := execution.AttemptResult{ + Model: "provider-image", + Body: []byte(`{"model":"public-image","data":[{"b64_json":"AA=="}]}`), + } + normalizeImagesAttemptResult(spec, &observed) + if observed.Model != "provider-image" { + t.Fatalf("observed response model = %q, want provider-image", observed.Model) + } + + stream := execution.StreamResult{Model: "provider-image"} + normalizeImagesStreamResult(spec, &stream) + if stream.Model != "" { + t.Fatalf("unobserved stream model = %q, want empty", stream.Model) + } +} + +func TestOpenAIImagesJSONAddsMissingContentType(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if got := request.Header.Get("Content-Type"); got != "application/json" { + t.Errorf("Content-Type = %q, want application/json", got) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"created":1,"data":[{"b64_json":"AA=="}]}`) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, testRuntimeOptions{allowPrivateNetwork: true}) + spec := openAIImagesSpec( + channel.OpenAICompatible, + server.URL+"/v1", + execution.OperationImagesGenerate, + "/v1/images/generations", + "", + []byte(`{"model":"public-image","prompt":"draw"}`), + ) + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil || result.Error != nil || result.StatusCode != http.StatusOK { + t.Fatalf("result = %+v, validation=%v", result, err) + } +} + +func TestOpenAIImagesUnaryAllowsBodyAboveDefaultLimit(t *testing.T) { + prefix := []byte(`{"created":1,"data":[{"b64_json":"`) + suffix := []byte(`"}]}`) + body := make([]byte, 0, int(defaultMaxUnaryResponseBodyBytes)+1+len(prefix)+len(suffix)) + body = append(body, prefix...) + body = append(body, bytes.Repeat([]byte{'A'}, int(defaultMaxUnaryResponseBodyBytes)+1)...) + body = append(body, suffix...) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(body) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, testRuntimeOptions{allowPrivateNetwork: true}) + spec := openAIImagesSpec( + channel.OpenAICompatible, + server.URL+"/v1", + execution.OperationImagesGenerate, + "/v1/images/generations", + "application/json", + []byte(`{"model":"public-image","prompt":"draw"}`), + ) + spec.Timeouts = execution.AttemptTimeouts{ + FirstByte: 10 * time.Second, + Request: 30 * time.Second, + StreamIdle: 10 * time.Second, + } + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil || result.Error != nil || result.StatusCode != http.StatusOK { + t.Fatalf("result = %+v, validation=%v", result, err) + } + if !bytes.Equal(result.Body, body) { + t.Fatalf("result body length = %d, want %d", len(result.Body), len(body)) + } +} + +func TestOpenAIImagesMultipartEditPreservesDataAndRemovesControlParts(t *testing.T) { + t.Parallel() + + imageBytes := []byte{0x00, 0xff, 0x10, 0x20, '\r', '\n'} + body, contentType := imagesMultipartBody(t, imageBytes) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1/images/edits" { + t.Errorf("path = %q", request.URL.Path) + } + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + if err != nil || mediaType != "multipart/form-data" || parameters["boundary"] == "" { + t.Errorf("Content-Type = %q: %v", request.Header.Get("Content-Type"), err) + return + } + reader := multipart.NewReader(request.Body, parameters["boundary"]) + var names []string + for { + part, err := reader.NextRawPart() + if err == io.EOF { + break + } + if err != nil { + t.Errorf("NextRawPart() error = %v", err) + return + } + _, params, err := mime.ParseMediaType(part.Header.Get("Content-Disposition")) + if err != nil { + t.Errorf("Content-Disposition = %q: %v", part.Header.Get("Content-Disposition"), err) + return + } + name := params["name"] + names = append(names, name) + data, err := io.ReadAll(part) + if err != nil { + t.Errorf("ReadAll(part) error = %v", err) + return + } + switch name { + case "model": + if string(data) != "provider-image" { + t.Errorf("model = %q", data) + } + case "stream": + if string(data) != "false" { + t.Errorf("stream = %q", data) + } + case "image[]": + if !bytes.Equal(data, imageBytes) || part.Header.Get("X-Vendor-Part") != "keep" || params["filename"] != "../original.png" { + t.Errorf("image part changed: data=%x header=%q filename=%q", data, part.Header.Get("X-Vendor-Part"), params["filename"]) + } + case "api_key", "provider", "fallbacks": + t.Errorf("control part %q reached upstream", name) + } + } + if got := strings.Join(names, ","); got != "prompt,model,stream,image[],future" { + t.Errorf("part order/names = %q", got) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"created":1,"data":[{"b64_json":"b2s="}]}`) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, testRuntimeOptions{allowPrivateNetwork: true}) + spec := openAIImagesSpec( + channel.OpenAICompatible, + server.URL+"/v1", + execution.OperationImagesEdit, + "/v1/images/edits", + contentType, + body, + ) + result := runtime.Execute(context.Background(), spec) + if err := result.Validate(); err != nil || result.Error != nil || result.StatusCode != http.StatusOK { + t.Fatalf("result = %+v, validation=%v", result, err) + } +} + +func TestOpenAIImagesStreamForcesStreamAndForwardsSSEWithoutUsage(t *testing.T) { + t.Parallel() + + const streamBody = "event: image_generation.partial_image\ndata: {\"type\":\"image_generation.partial_image\",\"b64_json\":\"cGFydA==\"}\n\n" + + "event: image_generation.completed\ndata: {\"type\":\"image_generation.completed\",\"b64_json\":\"ZG9uZQ==\"}\n\n" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/v1/images/generations" { + t.Errorf("path = %q", request.URL.Path) + } + body, _ := io.ReadAll(request.Body) + var object map[string]json.RawMessage + if err := json.Unmarshal(body, &object); err != nil || string(object["stream"]) != "true" { + t.Errorf("stream request body = %s, err=%v", body, err) + } + writer.Header().Set("Content-Type", "text/event-stream") + writer.WriteHeader(http.StatusOK) + _, _ = io.WriteString(writer, streamBody) + })) + defer server.Close() + + runtime := newProtocolTestRuntime(t, testRuntimeOptions{allowPrivateNetwork: true}) + spec := openAIImagesSpec( + channel.OpenAICompatible, + server.URL+"/v1", + execution.OperationImagesGenerate, + "/v1/images/generations", + "application/json", + []byte(`{"model":"public-image","stream":false,"prompt":"draw"}`), + ) + var events []execution.StreamEvent + result := runtime.ExecuteStream(context.Background(), spec, func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + if err := result.Validate(); err != nil || result.Error != nil { + t.Fatalf("result = %+v, validation=%v", result, err) + } + var data bytes.Buffer + for _, event := range events { + if event.Kind == execution.StreamEventUsage { + t.Fatalf("unexpected Images usage event: %+v", event) + } + if event.Kind == execution.StreamEventData { + data.Write(event.Data) + } + } + if data.String() != streamBody || result.Usage != nil { + t.Fatalf("stream data/usage = %q/%#v", data.String(), result.Usage) + } +} + +func TestOpenAIImagesFailuresAreNotSentOrExplicitlyReplayUnknown(t *testing.T) { + t.Parallel() + + t.Run("invalid body is not sent", func(t *testing.T) { + runtime := newProtocolTestRuntime(t, testRuntimeOptions{allowPrivateNetwork: true}) + spec := openAIImagesSpec( + channel.OpenAICompatible, + "https://example.invalid/v1", + execution.OperationImagesGenerate, + "/v1/images/generations", + "application/json", + []byte(`{"model":`), + ) + result := runtime.Execute(context.Background(), spec) + if result.DispatchState != execution.DispatchNotSent || result.Error == nil || + result.Error.Kind != execution.ErrorKindInvalidRequest || + result.Error.OriginHint != execution.ErrorOriginClient || + result.Error.ScopeHint != execution.ErrorScopeRequest || + result.Error.ReplaySafety != execution.ReplaySafetyUnknown { + t.Fatalf("result = %+v", result) + } + }) + + t.Run("provider rejection is maybe sent and replay unknown", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(writer, `{"error":{"type":"invalid_request_error","code":"unsupported_model","message":"unsupported"}}`) + })) + defer server.Close() + runtime := newProtocolTestRuntime(t, testRuntimeOptions{allowPrivateNetwork: true}) + spec := openAIImagesSpec( + channel.OpenAICompatible, + server.URL+"/v1", + execution.OperationImagesGenerate, + "/v1/images/generations", + "application/json", + []byte(`{"model":"public-image","prompt":"draw"}`), + ) + result := runtime.Execute(context.Background(), spec) + if result.DispatchState != execution.DispatchMaybeSent || !result.ResponseStarted || result.Error == nil || + result.Error.ReplaySafety != execution.ReplaySafetyUnknown { + t.Fatalf("result = %+v", result) + } + }) +} + +func TestOpenAIImagesRequestShapeAndCapabilityAreExact(t *testing.T) { + t.Parallel() + + manager := &RuntimeManager{} + for _, providerKind := range []channel.ProviderKind{channel.ProviderOpenAI, channel.ProviderOpenAICompatible} { + for _, operation := range []execution.Operation{execution.OperationImagesGenerate, execution.OperationImagesEdit} { + route := channel.RouteDescriptor{ + ClientProtocol: protocol.OpenAIImages, + Operation: operation, + RouteMode: execution.RouteNative, + } + if err := manager.ValidateRouteCapability(providerKind, route); err != nil { + t.Errorf("ValidateRouteCapability(%q, %q) error = %v", providerKind, operation, err) + } + } + } + for _, providerKind := range []channel.ProviderKind{channel.ProviderAnthropic, channel.ProviderXAI} { + if err := manager.ValidateRouteCapability(providerKind, channel.RouteDescriptor{ + ClientProtocol: protocol.OpenAIImages, + Operation: execution.OperationImagesGenerate, + RouteMode: execution.RouteNative, + }); err == nil { + t.Errorf("provider %q unexpectedly supports Images", providerKind) + } + } + + valid := openAIImagesSpec( + channel.OpenAICompatible, + "https://example.com/v1", + execution.OperationImagesGenerate, + "/v1/images/generations", + "application/json", + []byte(`{"model":"public-image"}`), + ) + if !supportedRequestShape(valid, false) || !supportedRequestShape(valid, true) { + t.Fatal("valid Images generation shape was rejected") + } + for _, mutate := range []func(*execution.AttemptSpec){ + func(spec *execution.AttemptSpec) { spec.Method = http.MethodGet }, + func(spec *execution.AttemptSpec) { spec.Path = "/v1/images/variations" }, + func(spec *execution.AttemptSpec) { spec.Operation = execution.OperationImagesEdit }, + func(spec *execution.AttemptSpec) { spec.RouteMode = execution.RouteConverted }, + } { + invalid := valid.Clone() + mutate(&invalid) + if supportedRequestShape(invalid, false) { + t.Errorf("invalid Images shape was accepted: %+v", invalid) + } + } +} + +func openAIImagesSpec( + channelID channel.ID, + baseURL string, + operation execution.Operation, + path string, + contentType string, + body []byte, +) execution.AttemptSpec { + return freezeTestAttempt(execution.NewAttemptSpec(execution.AttemptSpec{ + RequestID: "images-request", AttemptID: "images-attempt", Sequence: 1, + ChannelID: string(channelID), RouteMode: execution.RouteNative, + ClientProtocol: protocol.OpenAIImages, Operation: operation, + RouteRequirement: execution.RouteRequirementNative, + ClientModel: "public-image", UpstreamModel: "provider-image", + Method: http.MethodPost, Path: path, Query: make(map[string][]string), + Header: http.Header{ + "Content-Type": {contentType}, + "Authorization": {"Bearer client"}, + "Proxy-Authorization": {"Basic client"}, + "Api-Key": {"client"}, + "X-Api-Key": {"client"}, + "X-Test-Header": {"keep"}, + }, + Body: body, TargetConfig: json.RawMessage(`{"base_url":"` + baseURL + `"}`), + Timeouts: execution.AttemptTimeouts{FirstByte: time.Second, Request: 2 * time.Second, StreamIdle: time.Second}, + Credential: execution.NewCredentialSnapshot(17, 1, 1, []byte(`{"api_key":"`+testAPIKey+`"}`)), + })) +} + +func imagesMultipartBody(t *testing.T, imageBytes []byte) ([]byte, string) { + t.Helper() + var body bytes.Buffer + writer := multipart.NewWriter(&body) + writeField := func(name, value string) { + t.Helper() + part, err := writer.CreateFormField(name) + if err != nil { + t.Fatal(err) + } + _, _ = io.WriteString(part, value) + } + writeField("prompt", "keep prompt") + writeField("api_key", "client-secret") + writeField("model", "public-image") + writeField("provider", "client") + writeField("stream", "true") + writeField("fallbacks", "other") + imageHeader := make(textproto.MIMEHeader) + imageHeader.Set("Content-Disposition", `form-data; name="image[]"; filename="../original.png"`) + imageHeader.Set("Content-Type", "image/png") + imageHeader.Set("X-Vendor-Part", "keep") + image, err := writer.CreatePart(imageHeader) + if err != nil { + t.Fatal(err) + } + _, _ = image.Write(imageBytes) + writeField("future", "keep") + if err := writer.Close(); err != nil { + t.Fatal(err) + } + return bytes.Clone(body.Bytes()), writer.FormDataContentType() +} diff --git a/internal/execution/bifrost/model_alias.go b/internal/execution/bifrost/model_alias.go index cf593985..c8b7d5f1 100644 --- a/internal/execution/bifrost/model_alias.go +++ b/internal/execution/bifrost/model_alias.go @@ -9,14 +9,28 @@ import ( "gpt-load/internal/protocol" ) -const maxNativeAliasSSEEventBytes = 10 << 20 +const maxNativeAliasSSEEventBytes = execution.DefaultSSEEventLimitBytes const maxNativeFirstSSEEventBytes = maxNativeAliasSSEEventBytes type nativeFirstSSEEventGate struct { - pending []byte - scanStart int - ready bool + pending []byte + scanStart int + ready bool + maxEventBytes int +} + +func newNativeFirstSSEEventGate(spec execution.AttemptSpec) *nativeFirstSSEEventGate { + return &nativeFirstSSEEventGate{ + maxEventBytes: execution.SSEEventLimit(spec.ClientProtocol), + } +} + +func (g *nativeFirstSSEEventGate) eventLimit() int { + if g != nil && g.maxEventBytes > 0 { + return g.maxEventBytes + } + return maxNativeFirstSSEEventBytes } func (g *nativeFirstSSEEventGate) push(chunk []byte) ([]byte, error) { @@ -24,7 +38,7 @@ func (g *nativeFirstSSEEventGate) push(chunk []byte) ([]byte, error) { return append([]byte(nil), chunk...), nil } g.pending = append(g.pending, chunk...) - if len(g.pending) > maxNativeFirstSSEEventBytes { + if len(g.pending) > g.eventLimit() { return nil, fmt.Errorf("first native SSE event exceeds limit") } for { @@ -97,6 +111,7 @@ type nativeAliasSSERewriter struct { clientProtocol protocol.Protocol clientModel string pending []byte + maxEventBytes int } func newNativeAliasSSERewriter(spec execution.AttemptSpec) *nativeAliasSSERewriter { @@ -106,9 +121,17 @@ func newNativeAliasSSERewriter(spec execution.AttemptSpec) *nativeAliasSSERewrit return &nativeAliasSSERewriter{ clientProtocol: spec.ClientProtocol, clientModel: spec.ClientModel, + maxEventBytes: execution.SSEEventLimit(spec.ClientProtocol), } } +func (r *nativeAliasSSERewriter) eventLimit() int { + if r != nil && r.maxEventBytes > 0 { + return r.maxEventBytes + } + return maxNativeAliasSSEEventBytes +} + func (r *nativeAliasSSERewriter) push(chunk []byte) ([]byte, error) { if r == nil || len(chunk) == 0 { return append([]byte(nil), chunk...), nil @@ -118,13 +141,13 @@ func (r *nativeAliasSSERewriter) push(chunk []byte) ([]byte, error) { for { index, delimiterLength := firstNativeSSEDelimiter(r.pending) if index < 0 { - if len(r.pending) > maxNativeAliasSSEEventBytes { + if len(r.pending) > r.eventLimit() { return nil, fmt.Errorf("native SSE event exceeds limit") } return output.Bytes(), nil } eventEnd := index + delimiterLength - if eventEnd > maxNativeAliasSSEEventBytes { + if eventEnd > r.eventLimit() { return nil, fmt.Errorf("native SSE event exceeds limit") } event := append([]byte(nil), r.pending[:eventEnd]...) @@ -141,7 +164,7 @@ func (r *nativeAliasSSERewriter) finish() ([]byte, error) { if r == nil || len(r.pending) == 0 { return nil, nil } - if len(r.pending) > maxNativeAliasSSEEventBytes { + if len(r.pending) > r.eventLimit() { return nil, fmt.Errorf("native SSE event exceeds limit") } event := append([]byte(nil), r.pending...) diff --git a/internal/execution/bifrost/model_alias_test.go b/internal/execution/bifrost/model_alias_test.go index a1d50dca..e581ad27 100644 --- a/internal/execution/bifrost/model_alias_test.go +++ b/internal/execution/bifrost/model_alias_test.go @@ -27,6 +27,32 @@ func TestNativeFirstSSEEventGateAcceptsTwoMiBEventAcrossChunks(t *testing.T) { } } +func TestNativeSSEBoundariesUseImagesEventLimit(t *testing.T) { + payload := `{"type":"image_generation.partial_image","model":"provider-model","b64_json":"` + + strings.Repeat("A", maxNativeFirstSSEEventBytes+1) + `"}` + wire := []byte("event: image_generation.partial_image\ndata: " + payload + "\n\n") + spec := execution.AttemptSpec{ + ClientProtocol: protocol.OpenAIImages, + ClientModel: "public-model", + UpstreamModel: "provider-model", + } + + gate := newNativeFirstSSEEventGate(spec) + output, err := gate.push(wire) + if err != nil || !bytes.Equal(output, wire) { + t.Fatalf("first-event gate output/error = %d bytes / %v", len(output), err) + } + + rewriter := newNativeAliasSSERewriter(spec) + rewritten, err := rewriter.push(wire) + if err != nil { + t.Fatalf("alias rewriter error = %v", err) + } + if !bytes.Contains(rewritten, []byte(`"model":"public-model"`)) { + t.Fatal("alias rewriter did not preserve the oversized Images event") + } +} + func TestRewriteNativeResponseModelCoversNativeProtocolShapes(t *testing.T) { t.Parallel() diff --git a/internal/execution/bifrost/passthrough.go b/internal/execution/bifrost/passthrough.go index 41ab47fb..bdbf11b7 100644 --- a/internal/execution/bifrost/passthrough.go +++ b/internal/execution/bifrost/passthrough.go @@ -12,6 +12,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" + "gpt-load/internal/dialect" "gpt-load/internal/execution" "gpt-load/internal/protocol" ) @@ -165,6 +166,28 @@ func sanitizeNativeRequestBody(spec execution.AttemptSpec, stream bool) ([]byte, } } +func sanitizeNativePassthroughRequest( + spec execution.AttemptSpec, + stream bool, +) ([]byte, http.Header, error) { + if spec.ClientProtocol == protocol.OpenAIImages && + (spec.Operation == execution.OperationImagesGenerate || spec.Operation == execution.OperationImagesEdit) { + request, err := dialect.NewOpenAIImages().SanitizeRequestForAttempt(&dialect.ParsedRequest{ + Method: spec.Method, Path: spec.Path, RawQuery: spec.RawQuery, + Header: spec.Header.Clone(), Body: bytes.Clone(spec.Body), + }, spec.UpstreamModel, stream) + if err != nil { + return nil, nil, err + } + return bytes.Clone(request.Body), request.Header.Clone(), nil + } + body, err := sanitizeNativeRequestBody(spec, stream) + if err != nil { + return nil, nil, err + } + return body, spec.Header.Clone(), nil +} + func decodeNativeJSONObject(body []byte) (map[string]json.RawMessage, error) { decoder := json.NewDecoder(bytes.NewReader(body)) var object map[string]json.RawMessage @@ -321,7 +344,7 @@ func (r *Runtime) executeNative( cancelCall() return startedUnaryFailure(status, headers, execution.ErrorKindInternal, "execution runtime changed response status") } - limit := r.maxUnaryResponseBodyBytes + limit := r.unaryResponseBodyLimit(spec) if status < http.StatusOK || status >= http.StatusMultipleChoices { limit = min(limit, int64(maxStreamErrorEvidenceBytes)) } @@ -331,7 +354,7 @@ func (r *Runtime) executeNative( } _, _ = body.Write(response.Body) if response.PassthroughUsage != nil { - chunkUsage, err := usageEvidenceFromPassthrough(response.PassthroughUsage) + chunkUsage, err := usageEvidenceFromPassthroughForSpec(spec, response.PassthroughUsage) if err != nil { cancelCall() return startedUnaryFailure(status, headers, execution.ErrorKindInternal, "normalize upstream usage") @@ -342,7 +365,10 @@ func (r *Runtime) executeNative( } complete: - bodyBytes := bytes.Clone(body.Bytes()) + bodyBytes := body.Bytes() + if spec.ClientProtocol != protocol.OpenAIImages { + bodyBytes = bytes.Clone(bodyBytes) + } if headers.Get("Content-Encoding") == "" && looksLikeEncodedResponse(bodyBytes) { return startedUnaryFailure(status, headers, execution.ErrorKindInternal, "encoded upstream response cannot be safely forwarded") } @@ -490,7 +516,7 @@ func (r *Runtime) executeNativeStream( model := spec.UpstreamModel var usageEvidence *execution.UsageEvidence var errorBody bytes.Buffer - firstEventGate := &nativeFirstSSEEventGate{} + firstEventGate := newNativeFirstSSEEventGate(spec) aliasRewriter := newNativeAliasSSERewriter(spec) idleTimer := newIdleTimer(spec.Timeouts.StreamIdle) defer idleTimer.stop() @@ -616,7 +642,7 @@ func (r *Runtime) executeNativeStream( } } if response.PassthroughUsage != nil { - chunkUsage, err := usageEvidenceFromPassthrough(response.PassthroughUsage) + chunkUsage, err := usageEvidenceFromPassthroughForSpec(spec, response.PassthroughUsage) if err != nil { cancelCall() return nativeStreamRuntimeFailure(bifrostContext, prepared.secrets, started, status, headers, model, usageEvidence) @@ -643,6 +669,16 @@ func (r *Runtime) executeNativeStream( } } +func usageEvidenceFromPassthroughForSpec( + spec execution.AttemptSpec, + source *schemas.BifrostPassthroughUsage, +) (*execution.UsageEvidence, error) { + if spec.ClientProtocol == protocol.OpenAIImages { + return nil, nil + } + return usageEvidenceFromPassthrough(source) +} + func nativeStreamRuntimeFailure( bifrostContext *schemas.BifrostContext, secrets []string, diff --git a/internal/execution/bifrost/runtime.go b/internal/execution/bifrost/runtime.go index 186fcdb8..45e1bf34 100644 --- a/internal/execution/bifrost/runtime.go +++ b/internal/execution/bifrost/runtime.go @@ -22,7 +22,7 @@ const ( // Keep the SDK transport timeout at the largest whole-second value that // fits time.Duration so each attempt context remains the effective owner. providerTimeoutSecs = 9_223_372_036 - defaultMaxUnaryResponseBodyBytes = int64(32 << 20) + defaultMaxUnaryResponseBodyBytes = execution.DefaultUnaryResponseBodyLimitBytes ) var errKeyPoolDisabled = errors.New("Bifrost key pool is disabled; DirectKey is required") @@ -42,6 +42,13 @@ type Runtime struct { logger schemas.Logger } +func (r *Runtime) unaryResponseBodyLimit(spec execution.AttemptSpec) int64 { + if r != nil && r.maxUnaryResponseBodyBytes > 0 { + return r.maxUnaryResponseBodyBytes + } + return execution.UnaryResponseBodyLimit(spec.ClientProtocol) +} + func newConfiguredRuntime( ctx context.Context, options runtimeOptions, @@ -99,15 +106,11 @@ func newRuntimeShell(options runtimeOptions, registry *channel.Registry) (*Runti if registry == nil { return nil, fmt.Errorf("initialize execution runtime: channel registry is required") } - maxUnaryResponseBodyBytes := options.maxUnaryResponseBodyBytes - if maxUnaryResponseBodyBytes <= 0 { - maxUnaryResponseBodyBytes = defaultMaxUnaryResponseBodyBytes - } return &Runtime{ account: newDirectAccount(), registry: registry, allowPrivate: options.allowPrivateNetwork, - maxUnaryResponseBodyBytes: maxUnaryResponseBodyBytes, + maxUnaryResponseBodyBytes: options.maxUnaryResponseBodyBytes, shutdownDone: make(chan struct{}), logger: options.logger, }, nil diff --git a/internal/execution/contracts.go b/internal/execution/contracts.go index bb56738c..ba097c69 100644 --- a/internal/execution/contracts.go +++ b/internal/execution/contracts.go @@ -40,6 +40,8 @@ const ( OperationResponsesInputTokens Operation = "responses_input_tokens" OperationCountTokens Operation = "count_tokens" OperationResponsesPassthrough Operation = "responses_passthrough" + OperationImagesGenerate Operation = "images_generate" + OperationImagesEdit Operation = "images_edit" OperationListModels Operation = "list_models" OperationProbe Operation = "probe" ) @@ -57,6 +59,8 @@ func (o Operation) Valid() bool { OperationResponsesInputTokens, OperationCountTokens, OperationResponsesPassthrough, + OperationImagesGenerate, + OperationImagesEdit, OperationListModels, OperationProbe: return true @@ -65,6 +69,30 @@ func (o Operation) Valid() bool { } } +// ReplayPolicy describes the proof required before GPT-Load may dispatch the +// same logical request to another upstream candidate after a request may have +// reached an upstream. +type ReplayPolicy uint8 + +const ( + // ReplayPolicyLegacy preserves the established behavior of existing + // operations. Individual health rules still decide whether retry is safe. + ReplayPolicyLegacy ReplayPolicy = iota + // ReplayPolicyRequireRejectedBeforeProcessing permits another dispatch only + // when the provider explicitly proves that processing never started. + ReplayPolicyRequireRejectedBeforeProcessing +) + +// ReplayPolicy returns the operation-level replay contract. +func (o Operation) ReplayPolicy() ReplayPolicy { + switch o { + case OperationImagesGenerate, OperationImagesEdit: + return ReplayPolicyRequireRejectedBeforeProcessing + default: + return ReplayPolicyLegacy + } +} + // RouteMode records the already-selected wire strategy for one attempt. // Executors must execute this decision or reject it; they must not reroute it. type RouteMode string diff --git a/internal/execution/contracts_test.go b/internal/execution/contracts_test.go index 0b3753ba..bb79c959 100644 --- a/internal/execution/contracts_test.go +++ b/internal/execution/contracts_test.go @@ -27,6 +27,8 @@ func TestOperationAndDispatchEnums(t *testing.T) { OperationResponsesCompact, OperationResponsesInputTokens, OperationResponsesPassthrough, + OperationImagesGenerate, + OperationImagesEdit, OperationListModels, OperationProbe, } @@ -38,6 +40,14 @@ func TestOperationAndDispatchEnums(t *testing.T) { if Operation("unknown").Valid() { t.Fatal("expected unknown operation to be invalid") } + for _, operation := range []Operation{OperationImagesGenerate, OperationImagesEdit} { + if !operationRequiresModel(operation) { + t.Fatalf("operation %q must require a model", operation) + } + if operation.ReplayPolicy() != ReplayPolicyRequireRejectedBeforeProcessing { + t.Fatalf("operation %q replay policy = %v", operation, operation.ReplayPolicy()) + } + } for _, mode := range []RouteMode{RouteNative, RouteConverted} { if !mode.Valid() { t.Fatalf("expected route mode %q to be valid", mode) diff --git a/internal/execution/cpa/adapter.go b/internal/execution/cpa/adapter.go index addb7a93..7a9b1600 100644 --- a/internal/execution/cpa/adapter.go +++ b/internal/execution/cpa/adapter.go @@ -79,8 +79,11 @@ func (a *Adapter) ValidateRouteCapability( return provider.ValidateRouteCapability(route) } -func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execution.AttemptResult { +func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) (result execution.AttemptResult) { spec = execution.NewAttemptSpec(spec) + defer func() { + normalizeCPAImagesAttemptResult(spec, &result) + }() provider, err := a.validateSpec(spec) if err != nil { return unaryNotSent(execution.ErrorKindInvalidRequest, "unsupported subscription request", "", err) @@ -99,7 +102,15 @@ func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execu Proxy: spec.Proxy, Fingerprint: spec.ProxyFingerprint, }) } - request := bridgeRequest(spec, proxySettings) + request, err := bridgeRequest(spec, proxySettings, false) + if err != nil { + return unaryNotSent( + execution.ErrorKindInvalidRequest, + "subscription request input is not supported", + "unsupported_subscription_input", + err, + ) + } if validator, ok := provider.(providerRequestValidator); ok { if err := validator.ValidateRequest(request); err != nil { return unaryNotSent( @@ -182,7 +193,7 @@ func (a *Adapter) Execute(ctx context.Context, spec execution.AttemptSpec) execu ) { result.Error.Hint = execution.FailureHintRequestRejected } - result.UpstreamProtocol = provider.UpstreamProtocol() + result.UpstreamProtocol = effectiveUpstreamProtocol(provider, response.UpstreamProtocol) result.AppliedReasoning = appliedReasoning(response.AppliedReasoningEffort) return result } @@ -205,14 +216,28 @@ func unaryProviderSuccess( } return execution.AttemptResult{ DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, - UpstreamProtocol: provider.UpstreamProtocol(), AppliedReasoning: appliedReasoning(response.AppliedReasoningEffort), StatusCode: http.StatusOK, + UpstreamProtocol: effectiveUpstreamProtocol(provider, response.UpstreamProtocol), AppliedReasoning: appliedReasoning(response.AppliedReasoningEffort), StatusCode: http.StatusOK, Header: headers, Body: body, Model: responseModel(body, spec.UpstreamModel), UpstreamRequestID: upstreamRequestID(headers), Usage: responseUsage(spec, body), } } -func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, sink execution.StreamSink) execution.StreamResult { +func effectiveUpstreamProtocol(provider providerBridge, observed protocol.Protocol) protocol.Protocol { + if observed != "" { + return observed + } + return provider.UpstreamProtocol() +} + +func (a *Adapter) ExecuteStream( + ctx context.Context, + spec execution.AttemptSpec, + sink execution.StreamSink, +) (result execution.StreamResult) { spec = execution.NewAttemptSpec(spec) + defer func() { + normalizeCPAImagesStreamResult(spec, &result) + }() if sink == nil { return streamNotSent(execution.ErrorKindInvalidRequest, "stream sink is required", "") } @@ -236,7 +261,14 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, Proxy: spec.Proxy, Fingerprint: spec.ProxyFingerprint, }) } - request := bridgeRequest(spec, proxySettings) + request, err := bridgeRequest(spec, proxySettings, true) + if err != nil { + return streamNotSent( + execution.ErrorKindInvalidRequest, + "subscription request input is not supported", + "unsupported_subscription_input", + ) + } if validator, ok := provider.(providerRequestValidator); ok { if err := validator.ValidateRequest(request); err != nil { return streamNotSent(execution.ErrorKindInvalidRequest, "subscription request input is not supported", "unsupported_subscription_input") @@ -259,6 +291,10 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, firstByte := startFirstByteGate(spec.Timeouts.FirstByte, cancelStream) defer firstByte.stop() response, err := provider.ExecuteStream(streamCtx, strconv.FormatUint(uint64(spec.Credential.ID), 10), credential, request) + upstreamProtocol := provider.UpstreamProtocol() + if response != nil { + upstreamProtocol = effectiveUpstreamProtocol(provider, response.UpstreamProtocol) + } if err != nil { result := unaryExecutionError(streamCtx, provider, err, credential) var applied *reasoning.Config @@ -267,7 +303,7 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, } return execution.StreamResult{ DispatchState: result.DispatchState, ResponseStarted: result.ResponseStarted, - UpstreamProtocol: provider.UpstreamProtocol(), AppliedReasoning: applied, + UpstreamProtocol: upstreamProtocol, AppliedReasoning: applied, StatusCode: result.StatusCode, Header: result.Header, UpstreamRequestID: result.UpstreamRequestID, Error: result.Error, } @@ -288,7 +324,7 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, payload := frameSSE(spec.ClientProtocol, unframed) payload, rewriteErr := rewriteStreamModelAlias(spec, payload) if rewriteErr != nil { - failure := streamInternalError(provider.UpstreamProtocol(), headers, applied, "rewrite subscription response model", ready) + failure := streamInternalError(upstreamProtocol, headers, applied, "rewrite subscription response model", ready) return &failure } if spec.ClientProtocol == protocol.OpenAICompletions && isOpenAIDone(payload) { @@ -299,14 +335,14 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, } if !ready { if err := sink(execution.StreamEvent{Kind: execution.StreamEventReady, Sequence: sequence, StatusCode: http.StatusOK, Header: headers}); err != nil { - failure := streamConsumerStopped(provider.UpstreamProtocol(), headers, applied, false) + failure := streamConsumerStopped(upstreamProtocol, headers, applied, false) return &failure } ready = true } sequence++ if err := sink(execution.StreamEvent{Kind: execution.StreamEventData, Sequence: sequence, Data: payload}); err != nil { - failure := streamConsumerStopped(provider.UpstreamProtocol(), headers, applied, ready) + failure := streamConsumerStopped(upstreamProtocol, headers, applied, ready) return &failure } } @@ -319,40 +355,40 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, } chunk, ok, idleErr := nextChunk(streamCtx, response.Chunks, idleTimeout) if idleErr != nil { - return streamExecutionError(streamCtx, provider, headers, idleErr, credential, applied, ready) + return streamExecutionError(streamCtx, provider, upstreamProtocol, headers, idleErr, credential, applied, ready) } if !ok { if err := context.Cause(streamCtx); err != nil { - return streamExecutionError(streamCtx, provider, headers, err, credential, applied, ready) + return streamExecutionError(streamCtx, provider, upstreamProtocol, headers, err, credential, applied, ready) } if nativeResponsesAssembler != nil { payloads, finishErr := nativeResponsesAssembler.finish() if finishErr != nil { - return streamInternalError(provider.UpstreamProtocol(), headers, applied, "invalid subscription SSE stream", ready) + return streamInternalError(upstreamProtocol, headers, applied, "invalid subscription SSE stream", ready) } if failure := emitPayloads(payloads); failure != nil { return *failure } } if !ready { - return streamInternalError(provider.UpstreamProtocol(), headers, applied, "subscription upstream stream ended without data", false) + return streamInternalError(upstreamProtocol, headers, applied, "subscription upstream stream ended without data", false) } if spec.ClientProtocol == protocol.OpenAICompletions && !openAIDone { sequence++ if err := sink(execution.StreamEvent{Kind: execution.StreamEventData, Sequence: sequence, Data: []byte("data: [DONE]\n\n")}); err != nil { - return streamConsumerStopped(provider.UpstreamProtocol(), headers, applied, ready) + return streamConsumerStopped(upstreamProtocol, headers, applied, ready) } } if spec.ClientProtocol == protocol.Gemini && !geminiTerminal { sequence++ if err := sink(execution.StreamEvent{Kind: execution.StreamEventData, Sequence: sequence, Data: []byte("data: {\"candidates\":[{\"finishReason\":\"STOP\"}]}\n\n")}); err != nil { - return streamConsumerStopped(provider.UpstreamProtocol(), headers, applied, ready) + return streamConsumerStopped(upstreamProtocol, headers, applied, ready) } } - return successfulStreamTerminal(provider.UpstreamProtocol(), spec, headers, applied) + return successfulStreamTerminal(upstreamProtocol, spec, headers, applied) } if chunk.Err != nil { - return streamExecutionError(streamCtx, provider, headers, chunk.Err, credential, applied, ready) + return streamExecutionError(streamCtx, provider, upstreamProtocol, headers, chunk.Err, credential, applied, ready) } if len(chunk.Payload) == 0 && nativeResponsesAssembler == nil { continue @@ -362,13 +398,13 @@ func (a *Adapter) ExecuteStream(ctx context.Context, spec execution.AttemptSpec, upstreamStarted = true } if err := context.Cause(streamCtx); err != nil { - return streamExecutionError(streamCtx, provider, headers, err, credential, applied, false) + return streamExecutionError(streamCtx, provider, upstreamProtocol, headers, err, credential, applied, false) } payloads := [][]byte{chunk.Payload} if nativeResponsesAssembler != nil { payloads, err = nativeResponsesAssembler.push(chunk.Payload) if err != nil { - return streamInternalError(provider.UpstreamProtocol(), headers, applied, "invalid subscription SSE stream", ready) + return streamInternalError(upstreamProtocol, headers, applied, "invalid subscription SSE stream", ready) } } if failure := emitPayloads(payloads); failure != nil { @@ -383,7 +419,7 @@ func countTokensOperation(operation execution.Operation) bool { } func responseUsage(spec execution.AttemptSpec, body []byte) *execution.UsageEvidence { - if countTokensOperation(spec.Operation) { + if countTokensOperation(spec.Operation) || spec.ClientProtocol == protocol.OpenAIImages { return nil } return usageEvidence(spec.ClientProtocol, body) @@ -416,27 +452,62 @@ func (a *Adapter) validateSpec(spec execution.AttemptSpec) (providerBridge, erro if !ok || execution.RouteMode(mode) != spec.RouteMode { return nil, fmt.Errorf("subscription route is not declared by the channel") } - if !json.Valid(spec.Body) { + if spec.ClientProtocol == protocol.OpenAIImages { + if _, err := canonicalCPAImagesRequestPath(spec); err != nil { + return nil, err + } + if _, err := dialect.NewOpenAIImages().InspectRequest(&dialect.ParsedRequest{ + Method: spec.Method, Path: spec.Path, RawQuery: spec.RawQuery, + Header: spec.Header.Clone(), Body: append([]byte(nil), spec.Body...), + }); err != nil { + return nil, fmt.Errorf("invalid Images request: %w", err) + } + } else if !json.Valid(spec.Body) { return nil, fmt.Errorf("subscription request body must be JSON") } return provider, nil } -func bridgeRequest(spec execution.AttemptSpec, proxySettings cpaProxySettings) providerRequest { +func bridgeRequest( + spec execution.AttemptSpec, + proxySettings cpaProxySettings, + stream bool, +) (providerRequest, error) { + payload := append([]byte(nil), spec.Body...) + headers := spec.Header.Clone() + requestPath := "" + if spec.ClientProtocol == protocol.OpenAIImages { + var err error + requestPath, err = canonicalCPAImagesRequestPath(spec) + if err != nil { + return providerRequest{}, err + } + rebuilt, err := dialect.NewOpenAIImages().SanitizeRequestForAttempt(&dialect.ParsedRequest{ + Method: spec.Method, Path: spec.Path, RawQuery: spec.RawQuery, + Header: headers, Body: payload, + }, spec.UpstreamModel, stream) + if err != nil { + return providerRequest{}, err + } + payload = append([]byte(nil), rebuilt.Body...) + headers = rebuilt.Header.Clone() + } return providerRequest{ - AttemptID: spec.AttemptID, Model: spec.UpstreamModel, Payload: append([]byte(nil), spec.Body...), - Format: formatFor(spec.ClientProtocol), Headers: spec.Header.Clone(), - OriginalRequest: append([]byte(nil), spec.Body...), + AttemptID: spec.AttemptID, Model: spec.UpstreamModel, Payload: payload, + Format: formatFor(spec.ClientProtocol), RequestPath: requestPath, Headers: headers, + OriginalRequest: append([]byte(nil), payload...), ContinuityKey: spec.ContinuityKey, ProxyURL: proxySettings.URL, ProxyFromEnvironment: proxySettings.FromEnvironment, - } + }, nil } func formatFor(clientProtocol protocol.Protocol) string { switch clientProtocol { case protocol.OpenAIResponses: return "openai-response" + case protocol.OpenAIImages: + return "openai-image" case protocol.Anthropic: return "claude" case protocol.Gemini: @@ -446,6 +517,50 @@ func formatFor(clientProtocol protocol.Protocol) string { } } +func canonicalCPAImagesRequestPath(spec execution.AttemptSpec) (string, error) { + if spec.ClientProtocol != protocol.OpenAIImages || spec.RouteMode != execution.RouteNative || + spec.Method != http.MethodPost { + return "", fmt.Errorf("unsupported Images route tuple") + } + expected := "" + switch spec.Operation { + case execution.OperationImagesGenerate: + expected = "/v1/images/generations" + case execution.OperationImagesEdit: + expected = "/v1/images/edits" + default: + return "", fmt.Errorf("unsupported Images operation") + } + if spec.Path != expected { + return "", fmt.Errorf("Images path does not match operation") + } + return expected, nil +} + +func normalizeCPAImagesAttemptResult(spec execution.AttemptSpec, result *execution.AttemptResult) { + if result == nil || spec.ClientProtocol != protocol.OpenAIImages { + return + } + result.Usage = nil + if responseModel(result.Body, "") == "" { + result.Model = "" + } + if result.Error != nil && result.Error.ReplaySafety == "" { + result.Error.ReplaySafety = execution.ReplaySafetyUnknown + } +} + +func normalizeCPAImagesStreamResult(spec execution.AttemptSpec, result *execution.StreamResult) { + if result == nil || spec.ClientProtocol != protocol.OpenAIImages { + return + } + result.Usage = nil + result.Model = "" + if result.Error != nil && result.Error.ReplaySafety == "" { + result.Error.ReplaySafety = execution.ReplaySafetyUnknown + } +} + func subscriptionResponseHeaders(source http.Header, contentType string) http.Header { headers := make(http.Header) for actualName, values := range source { @@ -508,6 +623,7 @@ func definitelyNotSentNetworkError(err error) bool { func streamExecutionError( ctx context.Context, provider providerBridge, + upstreamProtocol protocol.Protocol, headers http.Header, err error, credential providerCredential, @@ -521,7 +637,7 @@ func streamExecutionError( } return execution.StreamResult{ DispatchState: execution.DispatchMaybeSent, ResponseStarted: responseStarted, - UpstreamProtocol: provider.UpstreamProtocol(), AppliedReasoning: applied, StatusCode: status, + UpstreamProtocol: upstreamProtocol, AppliedReasoning: applied, StatusCode: status, Header: headers, UpstreamRequestID: upstreamRequestID(headers), Error: evidence, } } diff --git a/internal/execution/cpa/adapter_test.go b/internal/execution/cpa/adapter_test.go index 2fb76e5b..46bdd4b6 100644 --- a/internal/execution/cpa/adapter_test.go +++ b/internal/execution/cpa/adapter_test.go @@ -6,22 +6,29 @@ import ( "encoding/json" "errors" "io" + "mime" + "mime/multipart" "net" "net/http" + "net/http/httptest" "reflect" "strings" "sync" "testing" "time" + "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" "gorm.io/gorm" "gpt-load/internal/channel" + "gpt-load/internal/dialect" "gpt-load/internal/execution" + "gpt-load/internal/gateway" "gpt-load/internal/health" "gpt-load/internal/outboundproxy" "gpt-load/internal/platform/encryption" + "gpt-load/internal/platform/httproute" "gpt-load/internal/protocol" "gpt-load/internal/state" stateloader "gpt-load/internal/state/loader" @@ -272,6 +279,344 @@ func TestAdapterExecutesEverySupportedClientProtocolThroughCPA(t *testing.T) { } } +func TestAdapterExecutesCodexImagesWithFixedCanonicalPath(t *testing.T) { + t.Parallel() + + adapter, _, _, keyService, row := newAdapterFixture(t, credentialJSON("access", "refresh", time.Now().Add(time.Hour))) + fake := &fakeExecutor{result: codex.ExecuteResponse{ + Payload: []byte(`{"created":1,"data":[{"b64_json":"AA=="}],"usage":{"input_tokens":9}}`), + }} + setCodexExecutor(t, adapter, fake) + spec := validSpec(t, row, keyService) + spec.ClientProtocol = protocol.OpenAIImages + spec.Operation = execution.OperationImagesGenerate + spec.RouteMode = execution.RouteNative + spec.RouteRequirement = execution.RouteRequirementNative + spec.ClientModel = "public-image" + spec.UpstreamModel = "provider-image" + spec.Path = "/v1/images/generations" + spec.Header = http.Header{"Content-Type": {"application/json"}} + spec.Body = []byte(`{"model":"public-image","stream":true,"prompt":"draw","provider":"client","api_key":"client","future":{"keep":true}}`) + + result := adapter.Execute(t.Context(), spec) + if err := result.Validate(); err != nil || result.Error != nil || fake.calls != 1 { + t.Fatalf("result/calls = %+v/%d, validation=%v", result, fake.calls, err) + } + if fake.request.Format != "openai-image" || fake.request.RequestPath != "/v1/images/generations" { + t.Fatalf("CPA Images request = %#v", fake.request) + } + var body map[string]json.RawMessage + if err := json.Unmarshal(fake.request.Payload, &body); err != nil { + t.Fatal(err) + } + if string(body["model"]) != `"provider-image"` || string(body["stream"]) != "false" || + string(body["future"]) != `{"keep":true}` { + t.Fatalf("sanitized payload = %s", fake.request.Payload) + } + for _, field := range []string{"provider", "api_key"} { + if _, exists := body[field]; exists { + t.Fatalf("control field %q reached CPA: %s", field, fake.request.Payload) + } + } + if !bytes.Equal(fake.request.OriginalRequest, fake.request.Payload) { + t.Fatalf("CPA original request retained unsanitized body: %s / %s", fake.request.OriginalRequest, fake.request.Payload) + } + if result.Usage != nil { + t.Fatalf("Images usage = %#v, want nil", result.Usage) + } +} + +func TestNormalizeCPAImagesResultsKeepOnlyObservedModel(t *testing.T) { + spec := execution.AttemptSpec{ClientProtocol: protocol.OpenAIImages} + + missing := execution.AttemptResult{ + Model: "provider-image", + Body: []byte(`{"created":1,"data":[{"b64_json":"AA=="}]}`), + } + normalizeCPAImagesAttemptResult(spec, &missing) + if missing.Model != "" { + t.Fatalf("missing response model = %q, want empty", missing.Model) + } + + observed := execution.AttemptResult{ + Model: "provider-image", + Body: []byte(`{"model":"public-image","data":[{"b64_json":"AA=="}]}`), + } + normalizeCPAImagesAttemptResult(spec, &observed) + if observed.Model != "provider-image" { + t.Fatalf("observed response model = %q, want provider-image", observed.Model) + } + + stream := execution.StreamResult{Model: "provider-image"} + normalizeCPAImagesStreamResult(spec, &stream) + if stream.Model != "" { + t.Fatalf("unobserved stream model = %q, want empty", stream.Model) + } +} + +func TestAdapterReportsDirectCodexImagesUpstreamProtocol(t *testing.T) { + t.Parallel() + + adapter, _, _, keyService, row := newAdapterFixture( + t, + credentialJSON("access", "refresh", time.Now().Add(time.Hour)), + ) + var upstreamPath string + transport := roundTripperFunc(func(request *http.Request) (*http.Response, error) { + upstreamPath = request.URL.Path + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"created":1,"data":[{"b64_json":"AA=="}]}`)), + Request: request, + }, nil + }) + ctx := context.WithValue(t.Context(), "cliproxy.roundtripper", http.RoundTripper(transport)) + spec := validSpec(t, row, keyService) + spec.ClientProtocol = protocol.OpenAIImages + spec.Operation = execution.OperationImagesGenerate + spec.RouteMode = execution.RouteNative + spec.RouteRequirement = execution.RouteRequirementNative + spec.ClientModel = "gpt-image-2" + spec.UpstreamModel = "gpt-image-2" + spec.Path = "/v1/images/generations" + spec.Header = http.Header{"Content-Type": {"application/json"}} + spec.Body = []byte(`{"model":"gpt-image-2","prompt":"draw"}`) + + result := adapter.Execute(ctx, spec) + if err := result.Validate(); err != nil || result.Error != nil { + t.Fatalf("result = %+v, validation=%v", result, err) + } + if upstreamPath != "/backend-api/codex/images/generations" { + t.Fatalf("upstream path = %q, want Codex Images endpoint", upstreamPath) + } + if result.UpstreamProtocol != protocol.OpenAIImages { + t.Fatalf("upstream protocol = %q, want %q", result.UpstreamProtocol, protocol.OpenAIImages) + } +} + +func TestAdapterExecutesCodexImagesMultipartAndStream(t *testing.T) { + t.Parallel() + + t.Run("multipart edit", func(t *testing.T) { + adapter, _, _, keyService, row := newAdapterFixture(t, credentialJSON("access", "refresh", time.Now().Add(time.Hour))) + fake := &fakeExecutor{result: codex.ExecuteResponse{Payload: []byte(`{"created":1,"data":[{"b64_json":"AA=="}]}`)}} + setCodexExecutor(t, adapter, fake) + body, contentType := cpaImagesMultipart(t) + spec := validSpec(t, row, keyService) + spec.ClientProtocol = protocol.OpenAIImages + spec.Operation = execution.OperationImagesEdit + spec.RouteRequirement = execution.RouteRequirementNative + spec.ClientModel = "public-image" + spec.UpstreamModel = "provider-image" + spec.Path = "/v1/images/edits" + spec.Header = http.Header{"Content-Type": {contentType}} + spec.Body = body + + result := adapter.Execute(t.Context(), spec) + if err := result.Validate(); err != nil || result.Error != nil || fake.calls != 1 { + t.Fatalf("result/calls = %+v/%d, validation=%v", result, fake.calls, err) + } + if fake.request.Format != "openai-image" || fake.request.RequestPath != "/v1/images/edits" { + t.Fatalf("CPA multipart request = %#v", fake.request) + } + mediaType, params, err := mime.ParseMediaType(fake.request.Headers.Get("Content-Type")) + if err != nil || mediaType != "multipart/form-data" { + t.Fatalf("Content-Type = %q: %v", fake.request.Headers.Get("Content-Type"), err) + } + reader := multipart.NewReader(bytes.NewReader(fake.request.Payload), params["boundary"]) + seen := map[string][]byte{} + for { + part, err := reader.NextRawPart() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + _, disposition, _ := mime.ParseMediaType(part.Header.Get("Content-Disposition")) + seen[disposition["name"]], _ = io.ReadAll(part) + } + if string(seen["model"]) != "provider-image" || string(seen["stream"]) != "false" || + string(seen["image[]"]) != "png-data" || seen["api_key"] != nil { + t.Fatalf("multipart fields = %#v", seen) + } + }) + + t.Run("stream generation", func(t *testing.T) { + adapter, _, _, keyService, row := newAdapterFixture(t, credentialJSON("access", "refresh", time.Now().Add(time.Hour))) + chunks := make(chan codex.ExecuteStreamChunk, 2) + chunks <- codex.ExecuteStreamChunk{Payload: []byte("event: image_generation.partial_image\ndata: {\"type\":\"image_generation.partial_image\",\"b64_json\":\"AA==\"}\n")} + chunks <- codex.ExecuteStreamChunk{Payload: []byte("event: image_generation.completed\ndata: {\"type\":\"image_generation.completed\",\"b64_json\":\"BB==\"}\n")} + close(chunks) + fake := &fakeExecutor{stream: &codex.ExecuteStreamResponse{ + Chunks: chunks, + UpstreamRequestPath: "/backend-api/codex/images/generations", + }} + setCodexExecutor(t, adapter, fake) + spec := validSpec(t, row, keyService) + spec.ClientProtocol = protocol.OpenAIImages + spec.Operation = execution.OperationImagesGenerate + spec.RouteRequirement = execution.RouteRequirementNative + spec.ClientModel = "public-image" + spec.UpstreamModel = "provider-image" + spec.Path = "/v1/images/generations" + spec.Header = http.Header{"Content-Type": {"application/json"}} + spec.Body = []byte(`{"model":"public-image","prompt":"draw"}`) + var events []execution.StreamEvent + result := adapter.ExecuteStream(t.Context(), spec, func(event execution.StreamEvent) error { + events = append(events, event.Clone()) + return nil + }) + if err := result.Validate(); err != nil || result.Error != nil || fake.calls != 1 { + t.Fatalf("result/calls = %+v/%d, validation=%v", result, fake.calls, err) + } + if result.UpstreamProtocol != protocol.OpenAIImages { + t.Fatalf("upstream protocol = %q, want %q", result.UpstreamProtocol, protocol.OpenAIImages) + } + if fake.request.RequestPath != "/v1/images/generations" || fake.request.Format != "openai-image" || + !bytes.Contains(fake.request.Payload, []byte(`"stream":true`)) { + t.Fatalf("stream CPA request = %#v", fake.request) + } + for _, event := range events { + if event.Kind == execution.StreamEventUsage { + t.Fatalf("unexpected Images usage event: %+v", event) + } + } + }) +} + +func TestAdapterRejectsInvalidCodexImagesTuplesBeforeDispatch(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + mutate func(*execution.AttemptSpec) + }{ + {name: "wrong method", mutate: func(spec *execution.AttemptSpec) { spec.Method = http.MethodGet }}, + {name: "arbitrary path", mutate: func(spec *execution.AttemptSpec) { spec.Path = "/v1/images/variations" }}, + {name: "operation path mismatch", mutate: func(spec *execution.AttemptSpec) { spec.Operation = execution.OperationImagesEdit }}, + {name: "converted mode", mutate: func(spec *execution.AttemptSpec) { spec.RouteMode = execution.RouteConverted }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + adapter, _, _, keyService, row := newAdapterFixture(t, credentialJSON("access", "refresh", time.Now().Add(time.Hour))) + fake := &fakeExecutor{} + setCodexExecutor(t, adapter, fake) + spec := validSpec(t, row, keyService) + spec.ClientProtocol = protocol.OpenAIImages + spec.Operation = execution.OperationImagesGenerate + spec.RouteRequirement = execution.RouteRequirementNative + spec.Path = "/v1/images/generations" + spec.Header = http.Header{"Content-Type": {"application/json"}} + spec.Body = []byte(`{"model":"gpt-5","prompt":"draw"}`) + test.mutate(&spec) + result := adapter.Execute(t.Context(), spec) + if result.DispatchState != execution.DispatchNotSent || result.Error == nil || fake.calls != 0 { + t.Fatalf("result/calls = %+v/%d", result, fake.calls) + } + }) + } +} + +func TestCodexClientImageGenerationReachesSubscriptionExecutor(t *testing.T) { + gin.SetMode(gin.TestMode) + adapter, _, credentialRegistry, keyService, row := newAdapterFixture( + t, + credentialJSON("access", "refresh", time.Now().Add(time.Hour)), + ) + fake := &fakeExecutor{result: codex.ExecuteResponse{ + Payload: []byte(`{"created":1713833628,"data":[{"b64_json":"VEVTVA=="}]}`), + }} + setCodexExecutor(t, adapter, fake) + credentialRef, ok := credentialRegistry.CredentialRef(row.ID) + if !ok { + t.Fatal("credential ref is unavailable") + } + manager := state.NewManager() + if _, err := manager.Publish(state.CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []state.GroupConfig{{ + ID: row.GroupID, Name: "codex", ChannelID: channel.Codex, + ConnectionType: "subscription", Params: json.RawMessage(`{}`), + Models: []state.ModelConfig{{ID: "gpt-image-2"}}, Enabled: true, + }}, + Credentials: []state.CredentialConfig{{ + ID: row.ID, GroupID: row.GroupID, Status: state.CredentialStatusActive, + Version: credentialRef.Version, IdentityGeneration: credentialRef.IdentityGeneration, + Fingerprint: credentialRef.Fingerprint, + }}, + AccessKeys: []state.AccessKeyConfig{{ + ID: 1, Name: "codex-client", KeyHash: keyService.Hash("gl-codex-client"), + Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish() error = %v", err) + } + handler := gateway.NewHandler( + manager, + credentialRegistry, + keyService, + gateway.NewExecutionForwarder(adapter), + dialect.NewSet(dialect.NewOpenAIImages()), + health.NewStatsStore(), + health.NewMutationCoordinator(), + nil, + nil, + nil, + ) + engine := gin.New() + routes, err := httproute.NewRegistry(handler.HTTPModule()) + if err != nil { + t.Fatalf("NewRegistry() error = %v", err) + } + if err := routes.Bind(engine); err != nil { + t.Fatalf("Bind() error = %v", err) + } + + request := httptest.NewRequest( + http.MethodPost, + "/v1/images/generations", + strings.NewReader(`{"model":"gpt-image-2","prompt":"生成一张 T 字母图片","background":"auto","quality":"auto","size":"auto"}`), + ) + request.Header.Set("Authorization", "Bearer gl-codex-client") + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusOK || recorder.Body.String() != `{"created":1713833628,"data":[{"b64_json":"VEVTVA=="}]}` { + t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String()) + } + if fake.calls != 1 || fake.request.Format != "openai-image" || + fake.request.RequestPath != "/v1/images/generations" || + !bytes.Contains(fake.request.Payload, []byte(`"model":"gpt-image-2"`)) { + t.Fatalf("Codex executor request = calls=%d %#v", fake.calls, fake.request) + } +} + +func cpaImagesMultipart(t *testing.T) ([]byte, string) { + t.Helper() + var body bytes.Buffer + writer := multipart.NewWriter(&body) + for name, value := range map[string]string{ + "model": "public-image", "prompt": "edit", "stream": "true", "api_key": "client", + } { + if err := writer.WriteField(name, value); err != nil { + t.Fatal(err) + } + } + part, err := writer.CreateFormFile("image[]", "source.png") + if err != nil { + t.Fatal(err) + } + _, _ = io.WriteString(part, "png-data") + if err := writer.Close(); err != nil { + t.Fatal(err) + } + return bytes.Clone(body.Bytes()), writer.FormDataContentType() +} + func TestSubscriptionResponseHeadersUseAllowlist(t *testing.T) { source := http.Header{ "Content-Type": {"application/upstream"}, diff --git a/internal/execution/cpa/codex_provider.go b/internal/execution/cpa/codex_provider.go index 6497502f..b5cde8b2 100644 --- a/internal/execution/cpa/codex_provider.go +++ b/internal/execution/cpa/codex_provider.go @@ -43,6 +43,19 @@ func (*codexProviderBridge) UpstreamProtocol() protocol.Protocol { return protocol.OpenAIResponses } +func codexUpstreamProtocol(requestPath string) protocol.Protocol { + requestPath = strings.TrimSpace(requestPath) + switch { + case strings.HasSuffix(requestPath, "/images/generations"), + strings.HasSuffix(requestPath, "/images/edits"): + return protocol.OpenAIImages + case strings.HasSuffix(requestPath, "/responses"): + return protocol.OpenAIResponses + default: + return "" + } +} + func (*codexProviderBridge) ValidateRouteCapability(route channel.RouteDescriptor) error { valid := route.ClientProtocol == protocol.OpenAIResponses && (route.Operation == execution.OperationResponsesCreate || @@ -58,6 +71,11 @@ func (*codexProviderBridge) ValidateRouteCapability(route channel.RouteDescripto route.RouteMode == execution.RouteConverted) } } + if route.ClientProtocol == protocol.OpenAIImages { + valid = (route.Operation == execution.OperationImagesGenerate || + route.Operation == execution.OperationImagesEdit) && + route.RouteMode == execution.RouteNative + } if !valid { return fmt.Errorf("route is not implemented by Codex") } @@ -73,7 +91,8 @@ func (bridge *codexProviderBridge) CountTokensLocal( } response, err := bridge.executor.CountTokens(ctx, "local-token-count", codex.Credential{}, codex.ExecuteRequest{ Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: request.Format, - Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), + RequestPath: request.RequestPath, + Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), ProxyURL: request.ProxyURL, ProxyFromEnvironment: request.ProxyFromEnvironment, }) headers := response.Headers.Clone() @@ -274,12 +293,14 @@ func (bridge *codexProviderBridge) Execute( } response, err := bridge.executor.Execute(ctx, credentialID, codexCredential.value, codex.ExecuteRequest{ Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: request.Format, - Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), + RequestPath: request.RequestPath, + Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), ProxyURL: request.ProxyURL, ProxyFromEnvironment: request.ProxyFromEnvironment, }) return providerResponse{ Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), AppliedReasoningEffort: response.AppliedReasoningEffort, + UpstreamProtocol: codexUpstreamProtocol(response.UpstreamRequestPath), }, err } @@ -295,7 +316,8 @@ func (bridge *codexProviderBridge) ExecuteStream( } response, err := bridge.executor.ExecuteStream(ctx, credentialID, codexCredential.value, codex.ExecuteRequest{ Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: request.Format, - Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), + RequestPath: request.RequestPath, + Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), ProxyURL: request.ProxyURL, ProxyFromEnvironment: request.ProxyFromEnvironment, }) if response == nil { @@ -316,6 +338,7 @@ func (bridge *codexProviderBridge) ExecuteStream( return &providerStreamResponse{ Headers: response.Headers.Clone(), Chunks: chunks, AppliedReasoningEffort: response.AppliedReasoningEffort, + UpstreamProtocol: codexUpstreamProtocol(response.UpstreamRequestPath), }, err } diff --git a/internal/execution/cpa/codex_provider_test.go b/internal/execution/cpa/codex_provider_test.go index 5a86cbe4..7b3cd7f6 100644 --- a/internal/execution/cpa/codex_provider_test.go +++ b/internal/execution/cpa/codex_provider_test.go @@ -7,6 +7,7 @@ import ( "time" "gpt-load/internal/execution" + "gpt-load/internal/protocol" ) type codexClassifiedTestError struct { @@ -15,6 +16,25 @@ type codexClassifiedTestError struct { retryAfter time.Duration } +func TestCodexUpstreamProtocolUsesObservedRequestPath(t *testing.T) { + t.Parallel() + + tests := []struct { + path string + want protocol.Protocol + }{ + {path: "/backend-api/codex/images/generations", want: protocol.OpenAIImages}, + {path: "/backend-api/codex/images/edits", want: protocol.OpenAIImages}, + {path: "/backend-api/codex/responses", want: protocol.OpenAIResponses}, + {path: "/unknown"}, + } + for _, test := range tests { + if got := codexUpstreamProtocol(test.path); got != test.want { + t.Errorf("codexUpstreamProtocol(%q) = %q, want %q", test.path, got, test.want) + } + } +} + func (err *codexClassifiedTestError) Error() string { return err.payload } func (err *codexClassifiedTestError) StatusCode() int { return err.status } func (err *codexClassifiedTestError) RetryAfter() *time.Duration { diff --git a/internal/execution/cpa/provider.go b/internal/execution/cpa/provider.go index 4141992f..915b7199 100644 --- a/internal/execution/cpa/provider.go +++ b/internal/execution/cpa/provider.go @@ -94,6 +94,7 @@ type providerRequest struct { Model string Payload []byte Format string + RequestPath string Headers http.Header OriginalRequest []byte // ContinuityKey is a private, tenant-scoped key used only by providers @@ -138,6 +139,7 @@ type providerResponse struct { Payload []byte Headers http.Header AppliedReasoningEffort string + UpstreamProtocol protocol.Protocol Local bool } @@ -145,6 +147,7 @@ type providerStreamResponse struct { Headers http.Header Chunks <-chan providerStreamChunk AppliedReasoningEffort string + UpstreamProtocol protocol.Protocol } type providerStreamChunk struct { diff --git a/internal/execution/resource_limits.go b/internal/execution/resource_limits.go new file mode 100644 index 00000000..37200b82 --- /dev/null +++ b/internal/execution/resource_limits.go @@ -0,0 +1,28 @@ +package execution + +import "gpt-load/internal/protocol" + +const ( + DefaultUnaryResponseBodyLimitBytes = int64(32 << 20) + OpenAIImagesUnaryResponseBodyLimitBytes = int64(64 << 20) + DefaultSSEEventLimitBytes = 10 << 20 + OpenAIImagesSSEEventLimitBytes = 32 << 20 +) + +// UnaryResponseBodyLimit returns the internal buffered success-response limit +// for one client wire protocol. +func UnaryResponseBodyLimit(clientProtocol protocol.Protocol) int64 { + if clientProtocol == protocol.OpenAIImages { + return OpenAIImagesUnaryResponseBodyLimitBytes + } + return DefaultUnaryResponseBodyLimitBytes +} + +// SSEEventLimit returns the internal maximum size of one complete SSE event +// for one client wire protocol. +func SSEEventLimit(clientProtocol protocol.Protocol) int { + if clientProtocol == protocol.OpenAIImages { + return OpenAIImagesSSEEventLimitBytes + } + return DefaultSSEEventLimitBytes +} diff --git a/internal/execution/resource_limits_test.go b/internal/execution/resource_limits_test.go new file mode 100644 index 00000000..fdfe90ac --- /dev/null +++ b/internal/execution/resource_limits_test.go @@ -0,0 +1,22 @@ +package execution + +import ( + "testing" + + "gpt-load/internal/protocol" +) + +func TestProtocolResourceLimits(t *testing.T) { + if got := UnaryResponseBodyLimit(protocol.OpenAICompletions); got != 32<<20 { + t.Fatalf("default unary response body limit = %d", got) + } + if got := UnaryResponseBodyLimit(protocol.OpenAIImages); got != 64<<20 { + t.Fatalf("Images unary response body limit = %d", got) + } + if got := SSEEventLimit(protocol.OpenAICompletions); got != 10<<20 { + t.Fatalf("default SSE event limit = %d", got) + } + if got := SSEEventLimit(protocol.OpenAIImages); got != 32<<20 { + t.Fatalf("Images SSE event limit = %d", got) + } +} diff --git a/internal/execution/responsealias/response_alias.go b/internal/execution/responsealias/response_alias.go index f0ac4986..3ea8c4f0 100644 --- a/internal/execution/responsealias/response_alias.go +++ b/internal/execution/responsealias/response_alias.go @@ -63,6 +63,8 @@ func modelRewriter(clientProtocol protocol.Protocol) (dialect.ModelRewriter, err return dialect.NewOpenAI(), nil case protocol.OpenAIResponses: return dialect.NewOpenAIResponses(), nil + case protocol.OpenAIImages: + return dialect.NewOpenAIImages(), nil case protocol.Anthropic: return dialect.NewAnthropic(), nil case protocol.Gemini: diff --git a/internal/execution/validation.go b/internal/execution/validation.go index d6fc711c..9fb0dd3b 100644 --- a/internal/execution/validation.go +++ b/internal/execution/validation.go @@ -305,6 +305,8 @@ func operationRequiresModel(operation Operation) bool { operation == OperationResponsesCompact || operation == OperationResponsesInputTokens || operation == OperationCountTokens || + operation == OperationImagesGenerate || + operation == OperationImagesEdit || operation == OperationProbe } diff --git a/internal/gateway/execution_forward.go b/internal/gateway/execution_forward.go index a678f77d..52c2cc93 100644 --- a/internal/gateway/execution_forward.go +++ b/internal/gateway/execution_forward.go @@ -108,10 +108,18 @@ func (forwarder *ExecutionForwarder) ForwardStream( input.Group.HeaderRules, input.CredentialSecrets..., ) - streamBuffer := newSSEEventObservationBuffer(func( + credentialSecrets := append([]string(nil), input.CredentialSecrets...) + if input.APIKey != "" { + credentialSecrets = append(credentialSecrets, input.APIKey) + } + streamBuffer := newSSEEventObservationBuffer(execution.SSEEventLimit(input.ClientProtocol), func( event dialect.StreamEvent, genericProviderError bool, ) (bool, error) { + if input.ClientProtocol == protocol.OpenAIImages && + imagesCredentialLiteralsRemain(event.Payload, credentialSecrets) { + return false, fmt.Errorf("%w: credential remains in response event", ErrUpstreamProtocol) + } wasTerminal := streamEvents.sawTerminal providerError, err := streamEvents.classify(event, genericProviderError) if err != nil { @@ -176,11 +184,18 @@ func (forwarder *ExecutionForwarder) ForwardStream( errorBody = appendExecutionErrorBody(errorBody, event.Data) return nil } - terminalInChunk, err := streamBuffer.push(event.Data) + completeData, terminalInChunk, err := streamBuffer.push(event.Data) if err != nil { downstreamErr = executionStreamProtocolFailure(err) return downstreamErr } + forwardData := event.Data + if input.ClientProtocol == protocol.OpenAIImages { + forwardData = completeData + if len(forwardData) == 0 { + return nil + } + } if !committed { if !firstResponse { firstResponse = true @@ -189,11 +204,11 @@ func (forwarder *ExecutionForwarder) ForwardStream( } } if streamEvents.firstEventWasProviderError() { - errorBody = appendExecutionErrorBody(errorBody, event.Data) + errorBody = appendExecutionErrorBody(errorBody, forwardData) return nil } committed = true - if err := commitStream(controller, ready.StatusCode, ready.Header, event.Data); err != nil { + if err := commitStream(controller, ready.StatusCode, ready.Header, forwardData); err != nil { downstreamErr = err return err } @@ -205,7 +220,7 @@ func (forwarder *ExecutionForwarder) ForwardStream( } return nil } - written, err := controller.write(event.Data) + written, err := controller.write(forwardData) if err != nil { downstreamErr = &streamFailure{ kind: streamFailureDownstreamWrite, @@ -213,7 +228,7 @@ func (forwarder *ExecutionForwarder) ForwardStream( } return downstreamErr } - if written != len(event.Data) { + if written != len(forwardData) { downstreamErr = &streamFailure{ kind: streamFailureDownstreamWrite, err: fmt.Errorf("write execution stream: %w", io.ErrShortWrite), @@ -703,8 +718,16 @@ func upstreamFromExecutionResult( upstream.AppliedReasoning = result.AppliedReasoning.Clone() } upstream.UpstreamProtocol = result.UpstreamProtocol - upstream.Body = append([]byte(nil), result.Body...) - upstream.ClassificationBody = append([]byte(nil), result.Body...) + if input.ClientProtocol == protocol.OpenAIImages { + // AttemptResult owns Body after the executor returns. The buffered Images + // representation consumes it synchronously, so move that ownership across + // the internal boundary instead of cloning a large base64 payload twice. + upstream.Body = result.Body + upstream.ClassificationBody = result.Body + } else { + upstream.Body = append([]byte(nil), result.Body...) + upstream.ClassificationBody = append([]byte(nil), result.Body...) + } if !result.ResponseStarted && result.Error != nil { upstream.Err = executionFailureError(ctx, result.Error) } diff --git a/internal/gateway/execution_forward_test.go b/internal/gateway/execution_forward_test.go index ed840d09..a5dc6378 100644 --- a/internal/gateway/execution_forward_test.go +++ b/internal/gateway/execution_forward_test.go @@ -306,6 +306,227 @@ func TestExecutionForwarderCommitsSuccessfulStreamOnlyOnFirstData(t *testing.T) } } +func TestExecutionForwarderAllowsImagesEventAboveDefaultLimit(t *testing.T) { + partial := "event: image_generation.partial_image\n" + + "data: {\"type\":\"image_generation.partial_image\",\"b64_json\":\"" + + strings.Repeat("A", maxSSEEventBytes+1) + "\"}\n\n" + completed := "event: image_generation.completed\n" + + "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"AA==\"}\n\n" + executor := fakeExecutionExecutor{stream: func( + _ context.Context, + _ execution.AttemptSpec, + sink execution.StreamSink, + ) execution.StreamResult { + for _, event := range []execution.StreamEvent{ + { + Sequence: 1, Kind: execution.StreamEventReady, StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + }, + {Sequence: 2, Kind: execution.StreamEventData, Data: []byte(partial)}, + {Sequence: 3, Kind: execution.StreamEventData, Data: []byte(completed)}, + } { + if err := sink(event); err != nil { + t.Fatalf("stream sink: %v", err) + } + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, + } + }} + input := executionForwardInput() + input.Dialect = dialect.NewOpenAIImages() + input.ClientProtocol = protocol.OpenAIImages + input.Operation = execution.OperationImagesGenerate + input.Request.Path = "/v1/images/generations" + input.Request.RawQuery = "" + input.Request.Body = []byte(`{"model":"public","stream":true}`) + recorder := httptest.NewRecorder() + result := NewExecutionForwarder(executor).ForwardStream(context.Background(), input, recorder) + if result.Err != nil || !result.Committed || result.Stream.EndReason != StreamEndCleanEOF { + t.Fatalf("ForwardStream() = %#v", result) + } + if !bytes.Equal(recorder.Body.Bytes(), []byte(partial+completed)) { + t.Fatalf("stream body length = %d, want %d", recorder.Body.Len(), len(partial)+len(completed)) + } +} + +func TestExecutionForwarderChecksImagesStreamCredentialsBeforeForwarding(t *testing.T) { + const ( + partial = "event: image_generation.partial_image\n" + + "data: {\"type\":\"image_generation.partial_image\",\"b64_json\":\"AAAA\",\"A\":1}\n\n" + completedStart = "event: image_generation.completed\n" + + "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"AA" + completedEnd = "AA\",\"A\":1}\n\n" + completed = completedStart + completedEnd + apiKeyLeakStart = "event: image_generation.completed\n" + + "data: {\"type\":\"image_generation.completed\",\"revised_prompt\":\"A" + credentialLeakStart = "event: image_generation.completed\n" + + "data: {\"type\":\"image_generation.completed\",\"metadata\":{\"token\":\"oauth-secret" + leakEnd = "\"}}\n\n" + ) + tests := []struct { + name string + chunks []string + wantSinkError bool + wantCommitted bool + wantEndReason StreamEndReason + wantBodies []string + }{ + { + name: "split media and structure collisions remain valid", + chunks: []string{completedStart, completedEnd}, + wantCommitted: true, + wantEndReason: StreamEndCleanEOF, + wantBodies: []string{"", completed}, + }, + { + name: "split credential leak before commit fails closed", + chunks: []string{credentialLeakStart, leakEnd}, + wantSinkError: true, + wantBodies: []string{"", ""}, + }, + { + name: "split credential leak after commit drops offending event", + chunks: []string{partial, apiKeyLeakStart, "\"}\n\n"}, + wantSinkError: true, + wantCommitted: true, + wantEndReason: StreamEndUpstreamProtocolError, + wantBodies: []string{partial, partial, partial}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var sinkErr error + recorder := httptest.NewRecorder() + bodies := make([]string, 0, len(test.chunks)) + executor := fakeExecutionExecutor{stream: func( + _ context.Context, + _ execution.AttemptSpec, + sink execution.StreamSink, + ) execution.StreamResult { + if err := sink(execution.StreamEvent{ + Sequence: 1, Kind: execution.StreamEventReady, StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + }); err != nil { + t.Fatalf("ready sink: %v", err) + } + for index, chunk := range test.chunks { + sinkErr = sink(execution.StreamEvent{ + Sequence: uint64(index + 2), Kind: execution.StreamEventData, Data: []byte(chunk), + }) + bodies = append(bodies, recorder.Body.String()) + if sinkErr != nil { + break + } + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, + } + }} + input := executionForwardInput() + input.Dialect = dialect.NewOpenAIImages() + input.ClientProtocol = protocol.OpenAIImages + input.Operation = execution.OperationImagesGenerate + input.Request.Path = "/v1/images/generations" + input.Request.RawQuery = "" + input.Request.Body = []byte(`{"model":"public","stream":true}`) + input.APIKey = "A" + input.CredentialSecrets = []string{"oauth-secret"} + + result := NewExecutionForwarder(executor).ForwardStream(context.Background(), input, recorder) + if test.wantSinkError { + if !errors.Is(sinkErr, ErrUpstreamProtocol) || !errors.Is(result.Err, ErrUpstreamProtocol) { + t.Fatalf("sink/result errors = %v / %v, want upstream protocol errors", sinkErr, result.Err) + } + } else if sinkErr != nil || result.Err != nil { + t.Fatalf("sink/result errors = %v / %v", sinkErr, result.Err) + } + if result.Committed != test.wantCommitted { + t.Fatalf("Committed = %t, want %t", result.Committed, test.wantCommitted) + } + if test.wantCommitted && result.Stream.EndReason != test.wantEndReason { + t.Fatalf("EndReason = %q, want %q", result.Stream.EndReason, test.wantEndReason) + } + if !reflect.DeepEqual(bodies, test.wantBodies) { + t.Fatalf("response bodies = %#v, want %#v", bodies, test.wantBodies) + } + }) + } +} + +func TestExecutionForwarderClassifiesImagesCompletedErrorObject(t *testing.T) { + tests := []struct { + name string + frames []string + wantEndReason StreamEndReason + }{ + { + name: "empty error completes", + frames: []string{ + "event: image_generation.completed\n" + + "data: {\"type\":\"image_generation.completed\",\"error\":{}}\n\n", + }, + wantEndReason: StreamEndCleanEOF, + }, + { + name: "non-empty error fails", + frames: []string{ + "event: image_generation.partial_image\n" + + "data: {\"type\":\"image_generation.partial_image\",\"b64_json\":\"AA==\"}\n\n", + "event: image_generation.completed\n" + + "data: {\"type\":\"image_generation.completed\",\"error\":{\"message\":\"failed\"}}\n\n", + }, + wantEndReason: StreamEndSSEError, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + executor := fakeExecutionExecutor{stream: func( + _ context.Context, + _ execution.AttemptSpec, + sink execution.StreamSink, + ) execution.StreamResult { + if err := sink(execution.StreamEvent{ + Sequence: 1, Kind: execution.StreamEventReady, StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + }); err != nil { + t.Fatalf("ready sink: %v", err) + } + for index, frame := range test.frames { + if err := sink(execution.StreamEvent{ + Sequence: uint64(index + 2), Kind: execution.StreamEventData, Data: []byte(frame), + }); err != nil { + t.Fatalf("data sink: %v", err) + } + } + return execution.StreamResult{ + DispatchState: execution.DispatchMaybeSent, ResponseStarted: true, + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, + } + }} + input := executionForwardInput() + input.Dialect = dialect.NewOpenAIImages() + input.ClientProtocol = protocol.OpenAIImages + input.Operation = execution.OperationImagesGenerate + input.Request.Path = "/v1/images/generations" + input.Request.RawQuery = "" + input.Request.Body = []byte(`{"model":"public","stream":true}`) + recorder := httptest.NewRecorder() + + result := NewExecutionForwarder(executor).ForwardStream(context.Background(), input, recorder) + if result.Err != nil || !result.Committed || result.Stream.EndReason != test.wantEndReason { + t.Fatalf("ForwardStream() = %#v", result) + } + if recorder.Body.String() != strings.Join(test.frames, "") { + t.Fatalf("response body = %q", recorder.Body.String()) + } + }) + } +} + func TestExecutionForwarderIgnoresOpenAIDoneForUsageCapture(t *testing.T) { executor := fakeExecutionExecutor{stream: func( _ context.Context, diff --git a/internal/gateway/forward.go b/internal/gateway/forward.go index 351d509a..918fcb71 100644 --- a/internal/gateway/forward.go +++ b/internal/gateway/forward.go @@ -87,7 +87,7 @@ func (result UpstreamResult) HasResponse() bool { } const ( - maxNonStreamingResponseBodyBytes = int64(32 << 20) + maxNonStreamingResponseBodyBytes = execution.DefaultUnaryResponseBodyLimitBytes maxErrorResponseBodyBytes = int64(64 << 10) maxDecompressedErrorBodyBytes = int64(1 << 20) maxStreamingErrorBodyBytes = int(maxErrorResponseBodyBytes) diff --git a/internal/gateway/images_response_representation_test.go b/internal/gateway/images_response_representation_test.go new file mode 100644 index 00000000..77542568 --- /dev/null +++ b/internal/gateway/images_response_representation_test.go @@ -0,0 +1,153 @@ +package gateway + +import ( + "bytes" + "net/http" + "strconv" + "testing" + + "gpt-load/internal/dialect" + "gpt-load/internal/execution" + "gpt-load/internal/platform/redact" + "gpt-load/internal/protocol" +) + +func TestImagesSuccessRepresentationPreservesSignedURL(t *testing.T) { + t.Parallel() + + body := []byte(`{"created":1,"data":[{"url":"https://cdn.example/image.png?signature=abc123&expires=999"}]}`) + headers := http.Header{ + "Content-Type": {"application/json"}, + "Content-Length": {strconv.Itoa(len(body))}, + } + prepared, err := (&responseProcessor{redactor: redact.New()}).prepareSuccessRepresentation( + ForwardInput{ + Dialect: dialect.NewOpenAIResponses(), + ClientProtocol: protocol.Protocol("openai-images"), + Operation: execution.Operation("images_generate"), + ExternalModel: "gpt-image-2", + UpstreamModelID: "gpt-image-2", + }, + http.StatusOK, + headers, + body, + nil, + ) + if err != nil { + t.Fatalf("prepareSuccessRepresentation() error = %v", err) + } + if !bytes.Equal(prepared.downstream, body) { + t.Fatalf("downstream body = %s, want exact %s", prepared.downstream, body) + } +} + +func TestImagesSuccessRepresentationFailsClosedOnKnownCredential(t *testing.T) { + t.Parallel() + + const secret = "A" + body := []byte(`{"created":1,"data":[{"revised_prompt":"` + secret + `"}]}`) + _, err := (&responseProcessor{redactor: redact.New()}).prepareSuccessRepresentation( + ForwardInput{ + Dialect: dialect.NewOpenAIResponses(), + ClientProtocol: protocol.Protocol("openai-images"), + Operation: execution.Operation("images_generate"), + ExternalModel: "gpt-image-2", + UpstreamModelID: "gpt-image-2", + }, + http.StatusOK, + http.Header{"Content-Type": {"application/json"}}, + body, + []string{secret}, + ) + if err == nil { + t.Fatal("prepareSuccessRepresentation() error = nil, want fail closed") + } +} + +func TestImagesSuccessRepresentationIgnoresCredentialCollisionInMediaAndStructure(t *testing.T) { + t.Parallel() + + secrets := []string{"A", "data", "1"} + body := []byte(`{"created":1,"model":"provider","data":[{"b64_json":"AAAA","url":"https://cdn.example/image"}]}`) + prepared, err := (&responseProcessor{redactor: redact.New()}).prepareSuccessRepresentation( + ForwardInput{ + Dialect: dialect.NewOpenAIImages(), + ClientProtocol: protocol.OpenAIImages, + Operation: execution.OperationImagesGenerate, + ExternalModel: "public", + UpstreamModelID: "provider", + }, + http.StatusOK, + http.Header{"Content-Type": {"application/json"}}, + body, + secrets, + ) + if err != nil { + t.Fatalf("prepareSuccessRepresentation() error = %v", err) + } + for _, value := range [][]byte{ + []byte(`"model":"public"`), + []byte(`"b64_json":"AAAA"`), + []byte(`"url":"https://cdn.example/image"`), + } { + if !bytes.Contains(prepared.downstream, value) { + t.Fatalf("downstream body = %s, want %s", prepared.downstream, value) + } + } +} + +func TestImagesSuccessRepresentationAllowsBodyAboveDefaultLimit(t *testing.T) { + prefix := []byte(`{"created":1,"data":[{"b64_json":"`) + suffix := []byte(`"}]}`) + payloadBytes := int(execution.OpenAIImagesUnaryResponseBodyLimitBytes) - len(prefix) - len(suffix) - 1 + body := make([]byte, 0, payloadBytes+len(prefix)+len(suffix)) + body = append(body, prefix...) + body = append(body, bytes.Repeat([]byte{'A'}, payloadBytes)...) + body = append(body, suffix...) + if got, want := len(body), int(execution.OpenAIImagesUnaryResponseBodyLimitBytes)-1; got != want { + t.Fatalf("capacity fixture length = %d, want %d", got, want) + } + + prepared, err := (&responseProcessor{redactor: redact.New()}).prepareSuccessRepresentation( + ForwardInput{ + Dialect: dialect.NewOpenAIImages(), + ClientProtocol: protocol.OpenAIImages, + Operation: execution.OperationImagesGenerate, + ExternalModel: "gpt-image-2", + UpstreamModelID: "gpt-image-2", + }, + http.StatusOK, + http.Header{"Content-Type": {"application/json"}}, + body, + []string{"sk-capacity-secret"}, + ) + if err != nil { + t.Fatalf("prepareSuccessRepresentation() error = %v", err) + } + if !bytes.Equal(prepared.downstream, body) { + t.Fatalf("downstream body length = %d, want %d", len(prepared.downstream), len(body)) + } +} + +func TestImagesSuccessRepresentationRejectsBodyAboveImagesLimit(t *testing.T) { + body := bytes.Repeat( + []byte{'A'}, + int(execution.OpenAIImagesUnaryResponseBodyLimitBytes)+1, + ) + _, err := (&responseProcessor{redactor: redact.New()}).prepareSuccessRepresentation( + ForwardInput{ + Dialect: dialect.NewOpenAIImages(), + ClientProtocol: protocol.OpenAIImages, + Operation: execution.OperationImagesGenerate, + ExternalModel: "gpt-image-2", + UpstreamModelID: "gpt-image-2", + }, + http.StatusOK, + http.Header{"Content-Type": {"application/json"}}, + body, + nil, + ) + if err == nil { + t.Fatal("prepareSuccessRepresentation() error = nil, want Images size rejection") + } +} diff --git a/internal/gateway/models.go b/internal/gateway/models.go index 6ed096dc..499ff27d 100644 --- a/internal/gateway/models.go +++ b/internal/gateway/models.go @@ -135,6 +135,7 @@ func modelListProtocols(value protocol.Protocol) []protocol.Protocol { return []protocol.Protocol{ protocol.OpenAICompletions, protocol.OpenAIResponses, + protocol.OpenAIImages, } } return []protocol.Protocol{value} diff --git a/internal/gateway/models_test.go b/internal/gateway/models_test.go index 8f7429d6..83587c7c 100644 --- a/internal/gateway/models_test.go +++ b/internal/gateway/models_test.go @@ -119,6 +119,15 @@ func TestVisibleOpenAIModelIDsUnionsChatAndResponses(t *testing.T) { }}, want: []string{"both", "chat-only", "responses-only", "shared"}, }, + { + name: "Images protocol filter", + accessKey: state.AccessKeyView{Filters: state.FilterSet{ + Protocols: map[protocol.Protocol]struct{}{ + protocol.OpenAIImages: {}, + }, + }}, + want: []string{"chat-only", "shared"}, + }, } for _, test := range tests { diff --git a/internal/gateway/request_log.go b/internal/gateway/request_log.go index d24bbd40..d5a3207b 100644 --- a/internal/gateway/request_log.go +++ b/internal/gateway/request_log.go @@ -116,6 +116,7 @@ func (recorder *requestRecorder) emit() { return } recorder.emitted = true + recorder.freezeImagesErrorSummaries() completedAt := recorder.now() duration := completedAt.Sub(recorder.startedAt) if duration < 0 { @@ -146,6 +147,22 @@ func (recorder *requestRecorder) emit() { }) } +func (recorder *requestRecorder) freezeImagesErrorSummaries() { + if recorder == nil || recorder.protocol != protocol.OpenAIImages { + return + } + if recorder.outcome.errorCode != "" { + recorder.outcome.errorSummary = fixedErrorSummary(recorder.outcome.errorCode) + } + for index := range recorder.attempts { + if recorder.attempts[index].ErrorCode != "" { + recorder.attempts[index].ErrorSummary = fixedErrorSummary( + recorder.attempts[index].ErrorCode, + ) + } + } +} + func (recorder *requestRecorder) estimatedCostNanoUSD() int64 { if recorder == nil { return 0 diff --git a/internal/gateway/request_log_test.go b/internal/gateway/request_log_test.go index ae6baada..3f8ada72 100644 --- a/internal/gateway/request_log_test.go +++ b/internal/gateway/request_log_test.go @@ -611,6 +611,103 @@ func TestRequestRecorderNon2xxKeepsAttributionButNotApplicable(t *testing.T) { } } +func TestOpenAIImagesRequestLogUsesFixedErrorsAndMissingUnpricedUsage(t *testing.T) { + t.Parallel() + + t.Run("provider diagnostics are never persisted", func(t *testing.T) { + sink := &recordingRequestLogSink{} + recorder := newRequestRecorder( + sink, + "req-images-error", + time.Unix(100, 0), + 9, + protocol.OpenAIImages, + func() time.Time { return time.Unix(101, 0) }, + ) + recorder.setOperation(execution.OperationImagesGenerate) + selection := requestLogSelection(12, 22, "images") + providerSummary := "prompt=private prompt url=https://cdn.example/private.png filename=secret.png" + index := recorder.appendAttempt( + selection, + UpstreamResult{StatusCode: http.StatusBadRequest}, + telemetry.FailureCategoryClientError, + telemetry.ActionTerminate, + "upstream_client_error", + providerSummary, + time.Unix(100, 0), + time.Unix(100, 0), + ) + recorder.completeResponse( + UpstreamResult{StatusCode: http.StatusBadRequest, ErrorSummary: providerSummary}, + health.Decision{Category: health.FailureCategoryClientError}, + "provider-image", + index, + ) + recorder.emit() + + if len(sink.events) != 1 { + t.Fatalf("events = %#v", sink.events) + } + event := sink.events[0] + want := fixedErrorSummary("upstream_client_error") + if event.ErrorSummary != want || len(event.Attempts) != 1 || + event.Attempts[0].ErrorSummary != want { + t.Fatalf("Images summaries = %q / %#v", event.ErrorSummary, event.Attempts) + } + encoded, _ := json.Marshal(event) + for _, forbidden := range []string{"private prompt", "cdn.example", "secret.png"} { + if bytes.Contains(encoded, []byte(forbidden)) { + t.Fatalf("request log retained %q: %s", forbidden, encoded) + } + } + }) + + t.Run("successful request is missing and unpriced", func(t *testing.T) { + sink := &recordingRequestLogSink{} + recorder := newRequestRecorder( + sink, + "req-images-success", + time.Unix(100, 0), + 9, + protocol.OpenAIImages, + func() time.Time { return time.Unix(101, 0) }, + ) + recorder.setOperation(execution.OperationImagesGenerate) + model := "provider-image" + recorder.freezeNextAttemptPricing(frozenAttemptPricing{ + channelID: string(channel.OpenAI), groupID: 12, + upstreamModel: model, applicable: true, + }) + selection := requestLogSelection(12, 22, "images") + selection.UpstreamModelID = &model + index := recorder.appendAttempt( + selection, + UpstreamResult{StatusCode: http.StatusOK}, + telemetry.FailureCategoryOK, + telemetry.ActionTerminate, + "", + "", + time.Unix(100, 0), + time.Unix(100, 0), + ) + recorder.completeResponse( + UpstreamResult{StatusCode: http.StatusOK, Usage: usage.Result{State: usage.StateMissing}}, + health.Decision{}, + model, + index, + ) + recorder.emit() + + event := sink.events[0] + if event.Usage.Result.State != usage.StateMissing || + event.Usage.Pricing.CostState != string(pricing.CostStateUnpriced) || + event.Usage.Pricing.PricingCompleteness != string(pricing.CompletenessUnavailable) || + event.Usage.Pricing.EstimatedCostNanoUSD != 0 || event.Usage.Pricing.ReceiptJSON != "" { + t.Fatalf("Images usage/pricing = %#v", event.Usage) + } + }) +} + func TestRequestRecorderInvalidAttemptIndexDoesNotForgeUsageAttribution(t *testing.T) { for _, index := range []int{-1, 1} { t.Run(fmt.Sprintf("index_%d", index), func(t *testing.T) { diff --git a/internal/gateway/response_representation.go b/internal/gateway/response_representation.go index f5a3ec40..7754bf76 100644 --- a/internal/gateway/response_representation.go +++ b/internal/gateway/response_representation.go @@ -12,8 +12,10 @@ import ( "strings" "gpt-load/internal/dialect" + "gpt-load/internal/execution" "gpt-load/internal/platform/contentcoding" "gpt-load/internal/platform/redact" + "gpt-load/internal/protocol" ) type preparedSuccessRepresentation struct { @@ -42,7 +44,8 @@ func (forwarder *responseProcessor) prepareSuccessRepresentation( wire []byte, secrets []string, ) (preparedSuccessRepresentation, error) { - if int64(len(wire)) > maxNonStreamingResponseBodyBytes { + bodyLimit := execution.UnaryResponseBodyLimit(input.ClientProtocol) + if int64(len(wire)) > bodyLimit { return preparedSuccessRepresentation{}, successRepresentationProtocolError("wire body exceeds limit") } @@ -52,34 +55,56 @@ func (forwarder *responseProcessor) prepareSuccessRepresentation( if err != nil { return preparedSuccessRepresentation{}, successRepresentationProtocolError("unsupported or malformed Content-Encoding") } - originalPlain, err := contentcoding.DecodeLimited( - encoding, - wire, - maxNonStreamingResponseBodyBytes, - ) - if err != nil { - return preparedSuccessRepresentation{}, successRepresentationProtocolError("decompress response body") + imagesRepresentation := input.ClientProtocol == protocol.OpenAIImages + var originalPlain []byte + if imagesRepresentation && encoding == contentcoding.Identity { + // The buffered attempt result owns wire for the duration of this terminal + // representation step. Keep the identity bytes instead of copying a large + // base64 payload before exact credential inspection. + originalPlain = wire + } else { + originalPlain, err = contentcoding.DecodeLimited( + encoding, + wire, + bodyLimit, + ) + if err != nil { + return preparedSuccessRepresentation{}, successRepresentationProtocolError("decompress response body") + } } modelTracker := newResponseModelTracker(input.Dialect, input.UpstreamModelID) modelTracker.observe(originalPlain) - patternSafePlain := forwarder.redactor.Bytes(originalPlain) - if int64(len(patternSafePlain)) > maxNonStreamingResponseBodyBytes { - return preparedSuccessRepresentation{}, successRepresentationProtocolError("redacted response body exceeds limit") - } - safePlain, residualCredential, ok := redactCredentialLiterals( - patternSafePlain, - secrets, - maxNonStreamingResponseBodyBytes, - ) - if !ok { - return preparedSuccessRepresentation{}, successRepresentationProtocolError("redact response body") - } - if residualCredential { - return preparedSuccessRepresentation{}, successRepresentationProtocolError("credential remains in response body") + var safePlain []byte + if imagesRepresentation { + if imagesCredentialLiteralsRemain(originalPlain, secrets) { + return preparedSuccessRepresentation{}, successRepresentationProtocolError("credential remains in response body") + } + safePlain = originalPlain + } else { + patternSafePlain := forwarder.redactor.Bytes(originalPlain) + if int64(len(patternSafePlain)) > bodyLimit { + return preparedSuccessRepresentation{}, successRepresentationProtocolError("redacted response body exceeds limit") + } + var residualCredential bool + var ok bool + safePlain, residualCredential, ok = redactCredentialLiterals( + patternSafePlain, + secrets, + bodyLimit, + ) + if !ok { + return preparedSuccessRepresentation{}, successRepresentationProtocolError("redact response body") + } + if residualCredential { + return preparedSuccessRepresentation{}, successRepresentationProtocolError("credential remains in response body") + } } - inspectablePlain := bytes.Clone(safePlain) + inspectablePlain := safePlain + if !imagesRepresentation { + inspectablePlain = bytes.Clone(safePlain) + } downstreamPlain := safePlain if needsModelRewrite(input) { rewriter, supported := input.Dialect.(dialect.ModelRewriter) @@ -90,10 +115,16 @@ func (forwarder *responseProcessor) prepareSuccessRepresentation( if err != nil { return preparedSuccessRepresentation{}, successRepresentationProtocolError("rewrite response model") } - if int64(len(downstreamPlain)) > maxNonStreamingResponseBodyBytes { + if int64(len(downstreamPlain)) > bodyLimit { return preparedSuccessRepresentation{}, successRepresentationProtocolError("rewritten response body exceeds limit") } - if credentialLiteralsRemain(downstreamPlain, secrets) { + var credentialRemains bool + if imagesRepresentation { + credentialRemains = imagesCredentialLiteralsRemain(downstreamPlain, secrets) + } else { + credentialRemains = credentialLiteralsRemain(downstreamPlain, secrets) + } + if credentialRemains { return preparedSuccessRepresentation{}, successRepresentationProtocolError("credential remains after model rewrite") } } @@ -569,6 +600,50 @@ func credentialLiteralsRemain(body []byte, secrets []string) bool { return jsonCredentialLiteralsRemain(value, replacers.residual) } +func imagesCredentialLiteralsRemain(body []byte, secrets []string) bool { + replacers, exists := newCredentialLiteralReplacers(secrets) + if !exists { + return false + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return credentialLiteralRemains(string(body), replacers.residual) + } + var trailing any + if err := decoder.Decode(&trailing); err != io.EOF { + return credentialLiteralRemains(string(body), replacers.residual) + } + return imagesJSONCredentialLiteralsRemain(value, replacers.residual) +} + +func imagesJSONCredentialLiteralsRemain(value any, residual *strings.Replacer) bool { + switch typed := value.(type) { + case string: + return credentialLiteralRemains(typed, residual) + case []any: + for _, item := range typed { + if imagesJSONCredentialLiteralsRemain(item, residual) { + return true + } + } + case map[string]any: + for key, item := range typed { + if key == "b64_json" { + if _, isString := item.(string); isString { + continue + } + } + if imagesJSONCredentialLiteralsRemain(item, residual) { + return true + } + } + } + return false +} + func jsonCredentialLiteralsRemain(value any, residual *strings.Replacer) bool { switch typed := value.(type) { case string: diff --git a/internal/gateway/router.go b/internal/gateway/router.go index fb9261a0..267018e0 100644 --- a/internal/gateway/router.go +++ b/internal/gateway/router.go @@ -8,13 +8,15 @@ import ( ) const ( - openAICompletionsPath = "/v1/chat/completions" - anthropicMessagesPath = "/v1/messages" - geminiModelsPath = "/v1beta/models" - geminiGenerationPattern = geminiModelsPath + "/:model_action" - modelsPath = "/v1/models" - openAIResponsesPath = "/v1/responses" - responsesResourcePattern = openAIResponsesPath + "/*resource_path" + openAICompletionsPath = "/v1/chat/completions" + anthropicMessagesPath = "/v1/messages" + geminiModelsPath = "/v1beta/models" + geminiGenerationPattern = geminiModelsPath + "/:model_action" + modelsPath = "/v1/models" + openAIResponsesPath = "/v1/responses" + responsesResourcePattern = openAIResponsesPath + "/*resource_path" + openAIImagesGenerationsPath = "/v1/images/generations" + openAIImagesEditsPath = "/v1/images/edits" ) type endpointKind uint8 @@ -91,6 +93,18 @@ func dataPlaneEndpointCatalog() []dataPlaneEndpoint { rejectAfterAuth: rejectResponsesAfterAuthentication, resolve: staticRoute(protocol.OpenAIResponses, endpointForward), }, + { + name: "data.openai.images.generations", + methods: []string{http.MethodPost}, + path: openAIImagesGenerationsPath, + resolve: staticRoute(protocol.OpenAIImages, endpointForward), + }, + { + name: "data.openai.images.edits", + methods: []string{http.MethodPost}, + path: openAIImagesEditsPath, + resolve: staticRoute(protocol.OpenAIImages, endpointForward), + }, } } diff --git a/internal/gateway/router_test.go b/internal/gateway/router_test.go index f0718945..faf48f9d 100644 --- a/internal/gateway/router_test.go +++ b/internal/gateway/router_test.go @@ -66,6 +66,16 @@ func TestDataPlaneEndpointCatalogDeclaresCompleteHTTPRoutes(t *testing.T) { methods: registeredResponsesMethods, path: "/v1/responses/*resource_path", }, + { + name: "data.openai.images.generations", + methods: []string{http.MethodPost}, + path: "/v1/images/generations", + }, + { + name: "data.openai.images.edits", + methods: []string{http.MethodPost}, + path: "/v1/images/edits", + }, } catalog := dataPlaneEndpointCatalog() @@ -190,6 +200,18 @@ func TestDataPlaneEndpointCatalogResolvesProtocolAndKind(t *testing.T) { want: route{Protocol: protocol.OpenAIResponses, Kind: endpointForward}, validPath: true, }, + { + name: "Images generation", endpoint: "data.openai.images.generations", + method: http.MethodPost, path: "/v1/images/generations", + want: route{Protocol: protocol.OpenAIImages, Kind: endpointForward}, + validPath: true, + }, + { + name: "Images edit", endpoint: "data.openai.images.edits", + method: http.MethodPost, path: "/v1/images/edits", + want: route{Protocol: protocol.OpenAIImages, Kind: endpointForward}, + validPath: true, + }, } catalog := dataPlaneEndpointCatalog() diff --git a/internal/gateway/stream_observation.go b/internal/gateway/stream_observation.go index 1a493c01..0b4960b6 100644 --- a/internal/gateway/stream_observation.go +++ b/internal/gateway/stream_observation.go @@ -63,36 +63,44 @@ type streamEventObserver struct { } // sseEventObservationBuffer frames arbitrary executor data chunks without -// changing their wire bytes. A terminal boundary returned by push belongs to -// the supplied chunk, so callers may mark it forwarded only after that chunk -// has been written and flushed successfully. +// changing their wire bytes. push returns the complete, successfully observed +// wire prefix; callers that use it may mark a returned terminal boundary +// forwarded only after that prefix has been written and flushed successfully. type sseEventObservationBuffer struct { pending []byte pendingTerminal bool + maxEventBytes int scanner sseRewriteBoundaryScanner observe func(dialect.StreamEvent, bool) (bool, error) } func newSSEEventObservationBuffer( + maxEventBytes int, observe func(dialect.StreamEvent, bool) (bool, error), ) *sseEventObservationBuffer { - return &sseEventObservationBuffer{observe: observe} + return &sseEventObservationBuffer{ + maxEventBytes: normalizedSSEEventLimit(maxEventBytes), + observe: observe, + } } -func (buffer *sseEventObservationBuffer) push(chunk []byte) (bool, error) { +func (buffer *sseEventObservationBuffer) push(chunk []byte) ([]byte, bool, error) { if buffer == nil || buffer.observe == nil { - return false, fmt.Errorf("SSE observation callback is required") + return nil, false, fmt.Errorf("SSE observation callback is required") } buffer.pending = append(buffer.pending, chunk...) + wire := buffer.pending + consumed := 0 terminal := false for { waitingTerminal := buffer.pendingTerminal optionalLF, overflow := buffer.scanner.ConsumeOptionalLineFeed( buffer.pending, false, + buffer.maxEventBytes, ) if overflow { - return false, errSSEEventTooLarge + return nil, false, errSSEEventTooLarge } if waitingTerminal && !buffer.scanner.optionalLineFeed { buffer.pendingTerminal = false @@ -100,25 +108,27 @@ func (buffer *sseEventObservationBuffer) push(chunk []byte) (bool, error) { } if optionalLF > 0 { buffer.discard(optionalLF) + consumed += optionalLF continue } eventEnd, complete := buffer.scanner.Find(buffer.pending) if !complete { - if len(buffer.pending) > maxSSEEventBytes { - return false, errSSEEventTooLarge + if len(buffer.pending) > buffer.maxEventBytes { + return nil, false, errSSEEventTooLarge } - return terminal, nil + return wire[:consumed], terminal, nil } - if eventEnd > maxSSEEventBytes { - return false, errSSEEventTooLarge + if eventEnd > buffer.maxEventBytes { + return nil, false, errSSEEventTooLarge } terminalEvent, err := buffer.observeEvent(buffer.pending[:eventEnd]) if err != nil { - return false, err + return nil, false, err } buffer.discard(eventEnd) + consumed += eventEnd buffer.scanner.AfterEvent(eventEnd, eventEnd) if terminalEvent && buffer.scanner.optionalLineFeed { buffer.pendingTerminal = true @@ -160,6 +170,7 @@ func (buffer *sseEventObservationBuffer) finish() error { optionalLF, overflow := buffer.scanner.ConsumeOptionalLineFeed( buffer.pending, true, + buffer.maxEventBytes, ) if overflow { return errSSEEventTooLarge diff --git a/internal/gateway/stream_observation_test.go b/internal/gateway/stream_observation_test.go index a4ab1bdd..8117a567 100644 --- a/internal/gateway/stream_observation_test.go +++ b/internal/gateway/stream_observation_test.go @@ -1,12 +1,29 @@ package gateway import ( + "bytes" "context" + "errors" "testing" + "gpt-load/internal/dialect" + "gpt-load/internal/execution" "gpt-load/internal/httplifecycle" + "gpt-load/internal/protocol" ) +func TestSSEObservationRejectsEventAboveImagesLimit(t *testing.T) { + limit := execution.SSEEventLimit(protocol.OpenAIImages) + buffer := newSSEEventObservationBuffer( + limit, + func(dialect.StreamEvent, bool) (bool, error) { return false, nil }, + ) + _, _, err := buffer.push(bytes.Repeat([]byte{'x'}, limit+1)) + if !errors.Is(err, errSSEEventTooLarge) { + t.Fatalf("push() error = %v, want %v", err, errSSEEventTooLarge) + } +} + func TestPrioritizeStreamObservationIdentifiesServerShutdown(t *testing.T) { requestContext, cancel := context.WithCancelCause(context.Background()) cancel(httplifecycle.ErrServerShutdown) diff --git a/internal/gateway/stream_rewrite.go b/internal/gateway/stream_rewrite.go index a2ffd8ef..d9a8d901 100644 --- a/internal/gateway/stream_rewrite.go +++ b/internal/gateway/stream_rewrite.go @@ -9,9 +9,10 @@ import ( "sync" "gpt-load/internal/dialect" + "gpt-load/internal/execution" ) -const maxSSEEventBytes = 10 << 20 +const maxSSEEventBytes = execution.DefaultSSEEventLimitBytes var ( errSSEEventTooLarge = errors.New("SSE event exceeds size limit") @@ -32,8 +33,9 @@ type sseRewriteStream struct { readMu sync.Mutex mu sync.Mutex - body io.ReadCloser - rewrite sseEventPayloadRewriter + body io.ReadCloser + rewrite sseEventPayloadRewriter + maxEventBytes int pending []byte output []byte @@ -49,14 +51,23 @@ type sseRewriteStream struct { func newSSEEventRewriteStream( body io.ReadCloser, + maxEventBytes int, rewrite sseEventPayloadRewriter, ) *sseRewriteStream { return &sseRewriteStream{ body: body, rewrite: rewrite, - scratch: make([]byte, streamReadBufferSize), + maxEventBytes: normalizedSSEEventLimit(maxEventBytes), + scratch: make([]byte, streamReadBufferSize), } } +func normalizedSSEEventLimit(limit int) int { + if limit > 0 { + return limit + } + return maxSSEEventBytes +} + func (stream *sseRewriteStream) Read(target []byte) (int, error) { if len(target) == 0 { return 0, nil @@ -95,6 +106,7 @@ func (stream *sseRewriteStream) Read(target []byte) (int, error) { optionalLF, overflow := stream.scanner.ConsumeOptionalLineFeed( stream.pending, stream.deferredErr != nil, + stream.maxEventBytes, ) if overflow { stream.finishLocked() @@ -115,7 +127,7 @@ func (stream *sseRewriteStream) Read(target []byte) (int, error) { eventEnd, complete := stream.scanner.Find(stream.pending) if complete { - if eventEnd > maxSSEEventBytes { + if eventEnd > stream.maxEventBytes { stream.finishLocked() stream.mu.Unlock() return 0, errSSEEventTooLarge @@ -135,7 +147,7 @@ func (stream *sseRewriteStream) Read(target []byte) (int, error) { stream.mu.Unlock() return 0, err } - if len(rewritten.body) > maxSSEEventBytes { + if len(rewritten.body) > stream.maxEventBytes { stream.mu.Lock() stream.finishLocked() stream.mu.Unlock() @@ -153,7 +165,7 @@ func (stream *sseRewriteStream) Read(target []byte) (int, error) { stream.mu.Unlock() continue } - if len(stream.pending) > maxSSEEventBytes { + if len(stream.pending) > stream.maxEventBytes { stream.finishLocked() stream.mu.Unlock() return 0, errSSEEventTooLarge @@ -181,7 +193,7 @@ func (stream *sseRewriteStream) Read(target []byte) (int, error) { return 0, err } body := stream.body - readLimit := min(streamReadBufferSize, maxSSEEventBytes+1-len(stream.pending)) + readLimit := min(streamReadBufferSize, stream.maxEventBytes+1-len(stream.pending)) chunk := stream.scratch[:readLimit] stream.mu.Unlock() @@ -320,7 +332,11 @@ func (scanner *sseRewriteBoundaryScanner) AfterEvent(inputBytes, outputBytes int scanner.skipLineFeed = false } -func (scanner *sseRewriteBoundaryScanner) ConsumeOptionalLineFeed(data []byte, final bool) (int, bool) { +func (scanner *sseRewriteBoundaryScanner) ConsumeOptionalLineFeed( + data []byte, + final bool, + maxEventBytes int, +) (int, bool) { if !scanner.optionalLineFeed { return 0, false } @@ -336,8 +352,8 @@ func (scanner *sseRewriteBoundaryScanner) ConsumeOptionalLineFeed(data []byte, f scanner.previousOutputBytes = 0 return 0, false } - overflow := scanner.previousInputBytes >= maxSSEEventBytes || - scanner.previousOutputBytes >= maxSSEEventBytes + overflow := scanner.previousInputBytes >= maxEventBytes || + scanner.previousOutputBytes >= maxEventBytes scanner.previousInputBytes = 0 scanner.previousOutputBytes = 0 return 1, overflow @@ -449,7 +465,8 @@ func isSSEErrorPayload(payload []byte) bool { } errorValue := bytes.TrimSpace(envelope.Error) return envelope.Type == "error" || - (len(errorValue) > 0 && !bytes.Equal(errorValue, []byte("null"))) + (len(errorValue) > 0 && !bytes.Equal(errorValue, []byte("null")) && + !bytes.Equal(errorValue, []byte("{}"))) } func splitSSEEventLines(event []byte) []sseEventLine { diff --git a/internal/gateway/stream_rewrite_test.go b/internal/gateway/stream_rewrite_test.go index acc21419..73e3991b 100644 --- a/internal/gateway/stream_rewrite_test.go +++ b/internal/gateway/stream_rewrite_test.go @@ -16,10 +16,11 @@ func newSSERewriteStream( rewrite func([]byte, bool) ([]byte, error), ) io.ReadCloser { if rewrite == nil { - return newSSEEventRewriteStream(body, nil) + return newSSEEventRewriteStream(body, maxSSEEventBytes, nil) } return newSSEEventRewriteStream( body, + maxSSEEventBytes, func(event dialect.StreamEvent, providerError bool) (sseEventRewriteResult, error) { body, err := rewrite(event.Payload, providerError) return sseEventRewriteResult{body: body}, err diff --git a/internal/health/execution_judge.go b/internal/health/execution_judge.go index 680ea50e..74188770 100644 --- a/internal/health/execution_judge.go +++ b/internal/health/execution_judge.go @@ -183,6 +183,7 @@ func JudgeExecution(attempt ExecutionAttempt, decisionContext DecisionContext) D result.RuleID == "fallback.ambiguous" { result.RuleID = "safety.replay_unknown" } + result = constrainOperationReplay(result, attempt, decisionContext) return constrainCommittedDecision(result, attempt) } @@ -322,6 +323,16 @@ func decisionForExecutionCategory( return rateLimitDecision(attempt, decisionContext) case FailureCategoryModelUnavailable: retry := retryUnlessExplicitlyUnknown(attempt.Evidence) + if decisionContext.Operation.ReplayPolicy() == execution.ReplayPolicyRequireRejectedBeforeProcessing { + return decision( + category, + origin, + scopeOrDefault(scope, execution.ErrorScopeModel), + retry, + EffectNone, + "images.model_unavailable", + ) + } result := decision( category, origin, @@ -387,6 +398,25 @@ func decisionForExecutionCategory( } } +func constrainOperationReplay( + result Decision, + attempt ExecutionAttempt, + decisionContext DecisionContext, +) Decision { + if result.Retry == RetryNone || + attempt.DispatchState != execution.DispatchMaybeSent || + decisionContext.Operation.ReplayPolicy() != execution.ReplayPolicyRequireRejectedBeforeProcessing { + return result + } + if attempt.Evidence != nil && + attempt.Evidence.ReplaySafety == execution.ReplaySafetyRejectedBeforeProcessing { + return result + } + result.Retry = RetryNone + result.RuleID = "safety.operation_replay_unsafe" + return result +} + func ambiguousRuleID(evidence *execution.ErrorEvidence) RuleID { if evidence == nil { return "fallback.ambiguous" diff --git a/internal/health/execution_judge_v2_test.go b/internal/health/execution_judge_v2_test.go index df09af0f..bb8e8794 100644 --- a/internal/health/execution_judge_v2_test.go +++ b/internal/health/execution_judge_v2_test.go @@ -412,6 +412,90 @@ func TestJudgeExecutionStableFallbackRuleMatrix(t *testing.T) { } } +func TestJudgeExecutionImagesRequireExplicitReplaySafety(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 8, 26, 12, 0, 0, 0, time.UTC) + imageOperation := execution.Operation("images_generate") + tests := []struct { + name string + status int + evidence execution.ErrorEvidence + wantRetry RetryDirective + wantEffect Effect + wantCooldown bool + }{ + { + name: "model unavailable without rejection proof", + status: http.StatusNotFound, + evidence: execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, Hint: execution.FailureHintModelUnavailable, + ScopeHint: execution.ErrorScopeModel, StatusCode: http.StatusNotFound, + Summary: "model unavailable", + }, + wantRetry: RetryNone, wantEffect: EffectNone, + }, + { + name: "candidate unavailable without rejection proof", + status: http.StatusBadRequest, + evidence: execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, Hint: execution.FailureHintCandidateUnavailable, + ScopeHint: execution.ErrorScopeModel, StatusCode: http.StatusBadRequest, + Summary: "operation unsupported", + }, + wantRetry: RetryNone, wantEffect: EffectNone, + }, + { + name: "rate limit without rejection proof", + status: http.StatusTooManyRequests, + evidence: execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, Hint: execution.FailureHintRateLimited, + StatusCode: http.StatusTooManyRequests, Summary: "rate limited", + }, + wantRetry: RetryNone, wantEffect: EffectCooldownCredential, wantCooldown: true, + }, + { + name: "invalid credential without rejection proof", + status: http.StatusUnauthorized, + evidence: execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, Hint: execution.FailureHintInvalidCredential, + ScopeHint: execution.ErrorScopeCredential, StatusCode: http.StatusUnauthorized, + Summary: "credential rejected", + }, + wantRetry: RetryNone, wantEffect: EffectRecordCredentialFailure, + }, + { + name: "explicit pre-processing rejection may advance", + status: http.StatusTooManyRequests, + evidence: execution.ErrorEvidence{ + Kind: execution.ErrorKindHTTP, Hint: execution.FailureHintRateLimited, + StatusCode: http.StatusTooManyRequests, Summary: "request rejected", + ReplaySafety: execution.ReplaySafetyRejectedBeforeProcessing, + }, + wantRetry: RetryNextCandidate, wantEffect: EffectCooldownCredential, wantCooldown: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result := JudgeExecution(ExecutionAttempt{ + DispatchState: execution.DispatchMaybeSent, + StatusCode: test.status, Evidence: &test.evidence, Now: now, + }, DecisionContext{ + Operation: imageOperation, + Method: http.MethodPost, + DefaultRateLimitCooldown: time.Minute, + }) + if result.Retry != test.wantRetry || result.Effect != test.wantEffect { + t.Fatalf("JudgeExecution() = %#v, want retry=%q effect=%q", result, test.wantRetry, test.wantEffect) + } + if result.CooldownUntil.IsZero() == test.wantCooldown { + t.Fatalf("cooldown = %v, want present=%t", result.CooldownUntil, test.wantCooldown) + } + }) + } +} + func TestJudgeExecutionRejectsFinalInformationalStatusRange(t *testing.T) { for _, status := range []int{http.StatusSwitchingProtocols, 199} { t.Run(http.StatusText(status), func(t *testing.T) { diff --git a/internal/protocol/protocol.go b/internal/protocol/protocol.go index 33e88745..5c0f266d 100644 --- a/internal/protocol/protocol.go +++ b/internal/protocol/protocol.go @@ -6,13 +6,14 @@ type Protocol string const ( OpenAICompletions Protocol = "openai-completions" OpenAIResponses Protocol = "openai-responses" + OpenAIImages Protocol = "openai-images" Anthropic Protocol = "anthropic" Gemini Protocol = "gemini" ) func (p Protocol) Valid() bool { switch p { - case OpenAICompletions, OpenAIResponses, Anthropic, Gemini: + case OpenAICompletions, OpenAIResponses, OpenAIImages, Anthropic, Gemini: return true default: return false @@ -21,7 +22,7 @@ func (p Protocol) Valid() bool { func (p Protocol) DataPlaneEnabled() bool { switch p { - case OpenAICompletions, OpenAIResponses, Anthropic, Gemini: + case OpenAICompletions, OpenAIResponses, OpenAIImages, Anthropic, Gemini: return true default: return false @@ -36,6 +37,7 @@ func DataPlaneProtocols() []Protocol { return []Protocol{ OpenAICompletions, OpenAIResponses, + OpenAIImages, Anthropic, Gemini, } diff --git a/internal/protocol/protocol_test.go b/internal/protocol/protocol_test.go index 065b0ba7..0fc0ba4c 100644 --- a/internal/protocol/protocol_test.go +++ b/internal/protocol/protocol_test.go @@ -49,6 +49,7 @@ func TestProtocolKnownAndDataPlaneEnabled(t *testing.T) { }{ {value: OpenAICompletions, known: true, enabled: true}, {value: OpenAIResponses, known: true, enabled: true}, + {value: OpenAIImages, known: true, enabled: true}, {value: Anthropic, known: true, enabled: true}, {value: Gemini, known: true, enabled: true}, {value: Protocol("openai"), known: false, enabled: false}, @@ -106,6 +107,7 @@ func TestDataPlaneProtocolsReturnsCanonicalOrderAndIndependentCopies(t *testing. want := []Protocol{ OpenAICompletions, OpenAIResponses, + OpenAIImages, Anthropic, Gemini, } diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 5e44581a..42710205 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -283,7 +283,10 @@ func normalizeQuery(query Query) normalizedQuery { clientProtocol := query.ClientProtocol operation := query.Operation if operation == "" { - if clientProtocol == protocol.OpenAIResponses { + if clientProtocol == protocol.OpenAIImages { + // Images endpoints always select generate or edit explicitly. Keep an + // omitted operation invalid instead of silently changing the action. + } else if clientProtocol == protocol.OpenAIResponses { if query.ExternalModel == nil { operation = execution.OperationResponsesRetrieve } else { diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 1459afc5..31337362 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -681,6 +681,18 @@ func TestIteratorExhaustsForNilOrEmptyDependencies(t *testing.T) { } } +func TestNormalizeQueryDoesNotDefaultOpenAIImagesOperation(t *testing.T) { + t.Parallel() + + normalized := normalizeQuery(Query{ + ClientProtocol: protocol.OpenAIImages, + ExternalModel: modelPointer("gpt-image-2"), + }) + if normalized.operation != "" { + t.Fatalf("operation = %q, want empty", normalized.operation) + } +} + func schedulerSnapshot() *state.ConfigSnapshot { snapshot, err := state.Compile(state.CompileInput{ ChannelRegistry: channel.NewRegistry(), diff --git a/internal/state/snapshot.go b/internal/state/snapshot.go index 37be2ff3..3efb1395 100644 --- a/internal/state/snapshot.go +++ b/internal/state/snapshot.go @@ -328,7 +328,9 @@ func appendExecutionTargets( execution.OperationResponsesCreate, execution.OperationResponsesCompact, execution.OperationResponsesInputTokens, - execution.OperationCountTokens: + execution.OperationCountTokens, + execution.OperationImagesGenerate, + execution.OperationImagesEdit: for _, model := range group.Models { modelMode, supported := target.ModeForModel(clientProtocol, operation, model.ID) if !supported { diff --git a/internal/state/snapshot_test.go b/internal/state/snapshot_test.go index 313fd146..5b746dd7 100644 --- a/internal/state/snapshot_test.go +++ b/internal/state/snapshot_test.go @@ -55,6 +55,31 @@ func TestCompileIndexesExternalModelsAndPreservesUpstreamIDs(t *testing.T) { } } +func TestCompileIndexesOpenAIImagesOperationsForAllConfiguredModels(t *testing.T) { + t.Parallel() + + snapshot, err := Compile(CompileInput{ + ChannelRegistry: channel.NewRegistry(), + Groups: []GroupConfig{ + {ConnectionType: "api_key", ID: 1, ChannelID: channel.OpenAI, Params: json.RawMessage(`{}`), + Models: []ModelConfig{{ID: "official-image", Alias: "public"}}, Enabled: true}, + {ConnectionType: "api_key", ID: 2, ChannelID: channel.OpenAICompatible, + Params: json.RawMessage(`{"base_url":"https://proxy.example/api/v4"}`), + Models: []ModelConfig{{ID: "compatible-image", Alias: "public"}}, Enabled: true}, + }, + }) + if err != nil { + t.Fatalf("Compile() error = %v", err) + } + for _, operation := range []execution.Operation{execution.OperationImagesGenerate, execution.OperationImagesEdit} { + got := snapshot.ExecutionCandidates[protocol.OpenAIImages][operation]["public"] + if len(got) != 2 || got[0].GroupID != 1 || got[1].GroupID != 2 || + got[0].Mode != channel.RouteNative || got[1].Mode != channel.RouteNative { + t.Errorf("Images candidates for %q = %#v", operation, got) + } + } +} + func TestCompileSubscriptionPublishesOnlyVerifiedCodexOperations(t *testing.T) { t.Parallel() snapshot, err := Compile(CompileInput{ diff --git a/internal/subscription/providers/codex/codex.go b/internal/subscription/providers/codex/codex.go index 3acac27c..85693488 100644 --- a/internal/subscription/providers/codex/codex.go +++ b/internal/subscription/providers/codex/codex.go @@ -253,6 +253,7 @@ type ExecuteRequest struct { Model string Payload []byte Format string + RequestPath string Headers http.Header OriginalRequest []byte ProxyURL string @@ -264,6 +265,7 @@ type ExecuteResponse struct { Payload []byte Headers http.Header AppliedReasoningEffort string + UpstreamRequestPath string } // ExecuteStreamResponse contains converted streaming chunks and response metadata. @@ -271,6 +273,7 @@ type ExecuteStreamResponse struct { Headers http.Header Chunks <-chan ExecuteStreamChunk AppliedReasoningEffort string + UpstreamRequestPath string } // ExecuteStreamChunk contains one converted payload or terminal bridge error. @@ -311,6 +314,7 @@ func (e *executor) Execute( Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), AppliedReasoningEffort: response.AppliedReasoningEffort, + UpstreamRequestPath: response.UpstreamRequestPath, }, err } @@ -330,6 +334,7 @@ func (e *executor) CountTokens( Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), AppliedReasoningEffort: response.AppliedReasoningEffort, + UpstreamRequestPath: response.UpstreamRequestPath, }, err } @@ -364,6 +369,7 @@ func (e *executor) ExecuteStream( Headers: response.Headers.Clone(), Chunks: chunks, AppliedReasoningEffort: response.AppliedReasoningEffort, + UpstreamRequestPath: response.UpstreamRequestPath, }, err } @@ -372,6 +378,7 @@ func executeRequestToBridge(value ExecuteRequest) cpaembedded.ExecuteRequest { Model: value.Model, Payload: append([]byte(nil), value.Payload...), Format: value.Format, + RequestPath: value.RequestPath, Headers: value.Headers.Clone(), OriginalRequest: append([]byte(nil), value.OriginalRequest...), ProxyURL: value.ProxyURL, diff --git a/internal/subscription/providers/codex/codex_test.go b/internal/subscription/providers/codex/codex_test.go index 6a4164a2..165327e5 100644 --- a/internal/subscription/providers/codex/codex_test.go +++ b/internal/subscription/providers/codex/codex_test.go @@ -72,3 +72,12 @@ func TestExecuteRequestToBridgePreservesEnvironmentProxyPolicy(t *testing.T) { t.Fatalf("bridge proxy request = %#v", request) } } + +func TestExecuteRequestToBridgePreservesFixedRequestPath(t *testing.T) { + t.Parallel() + + request := executeRequestToBridge(ExecuteRequest{RequestPath: "/v1/images/generations"}) + if request.RequestPath != "/v1/images/generations" { + t.Fatalf("bridge request path = %q", request.RequestPath) + } +} diff --git a/third_party/cpaembedded/embedded/embedded.go b/third_party/cpaembedded/embedded/embedded.go index 94a0b535..d22f02b9 100644 --- a/third_party/cpaembedded/embedded/embedded.go +++ b/third_party/cpaembedded/embedded/embedded.go @@ -145,6 +145,7 @@ type ExecuteRequest struct { Model string Payload []byte Format string + RequestPath string Headers http.Header OriginalRequest []byte ContinuityKey string @@ -156,12 +157,14 @@ type ExecuteResponse struct { Payload []byte Headers http.Header AppliedReasoningEffort string + UpstreamRequestPath string } type ExecuteStreamResponse struct { Headers http.Header Chunks <-chan ExecuteStreamChunk AppliedReasoningEffort string + UpstreamRequestPath string } type ExecuteStreamChunk struct { @@ -377,16 +380,17 @@ func (e *CodexHTTPExecutor) ExecuteCanonical(ctx context.Context, credentialID s executionCtx := e.executionContext(ctx, auth, observation, request.ProxyFromEnvironment) response, err := e.inner.Execute(executionCtx, authWithoutProxyURL(auth), cliproxyexecutor.Request{ Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, - }, cliproxyexecutor.Options{ - Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), - SourceFormat: format, ResponseFormat: format, - }) + }, codexExecutionOptions(request, format, false)) if err != nil { - return ExecuteResponse{AppliedReasoningEffort: observation.reasoningEffort()}, err + return ExecuteResponse{ + AppliedReasoningEffort: observation.reasoningEffort(), + UpstreamRequestPath: observation.upstreamRequestPath(), + }, err } return ExecuteResponse{ Payload: append([]byte(nil), response.Payload...), Headers: response.Headers.Clone(), AppliedReasoningEffort: observation.reasoningEffort(), + UpstreamRequestPath: observation.upstreamRequestPath(), }, nil } @@ -398,23 +402,27 @@ func (e *CodexHTTPExecutor) CountTokensCanonical(ctx context.Context, credential executionCtx := e.executionContext(ctx, auth, observation, request.ProxyFromEnvironment) response, err := e.inner.CountTokens(executionCtx, authWithoutProxyURL(auth), cliproxyexecutor.Request{ Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, - }, cliproxyexecutor.Options{ - Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), - SourceFormat: format, ResponseFormat: format, - }) + }, codexExecutionOptions(request, format, false)) if err != nil { - return ExecuteResponse{AppliedReasoningEffort: observation.reasoningEffort()}, err + return ExecuteResponse{ + AppliedReasoningEffort: observation.reasoningEffort(), + UpstreamRequestPath: observation.upstreamRequestPath(), + }, err } payload := append([]byte(nil), response.Payload...) if format == sdktranslator.FormatOpenAIResponse { payload, err = normalizeCodexResponsesTokenCount(payload) if err != nil { - return ExecuteResponse{AppliedReasoningEffort: observation.reasoningEffort()}, err + return ExecuteResponse{ + AppliedReasoningEffort: observation.reasoningEffort(), + UpstreamRequestPath: observation.upstreamRequestPath(), + }, err } } return ExecuteResponse{ Payload: payload, Headers: response.Headers.Clone(), AppliedReasoningEffort: observation.reasoningEffort(), + UpstreamRequestPath: observation.upstreamRequestPath(), }, nil } @@ -449,12 +457,12 @@ func (e *CodexHTTPExecutor) ExecuteStreamCanonical(ctx context.Context, credenti executionCtx := e.executionContext(ctx, auth, observation, request.ProxyFromEnvironment) response, err := e.inner.ExecuteStream(executionCtx, authWithoutProxyURL(auth), cliproxyexecutor.Request{ Model: request.Model, Payload: append([]byte(nil), request.Payload...), Format: format, - }, cliproxyexecutor.Options{ - Stream: true, Headers: request.Headers.Clone(), OriginalRequest: append([]byte(nil), request.OriginalRequest...), - SourceFormat: format, ResponseFormat: format, - }) + }, codexExecutionOptions(request, format, true)) if err != nil { - return &ExecuteStreamResponse{AppliedReasoningEffort: observation.reasoningEffort()}, err + return &ExecuteStreamResponse{ + AppliedReasoningEffort: observation.reasoningEffort(), + UpstreamRequestPath: observation.upstreamRequestPath(), + }, err } chunks := make(chan ExecuteStreamChunk) go func() { @@ -470,9 +478,28 @@ func (e *CodexHTTPExecutor) ExecuteStreamCanonical(ctx context.Context, credenti return &ExecuteStreamResponse{ Headers: response.Headers.Clone(), Chunks: chunks, AppliedReasoningEffort: observation.reasoningEffort(), + UpstreamRequestPath: observation.upstreamRequestPath(), }, nil } +func codexExecutionOptions( + request ExecuteRequest, + format sdktranslator.Format, + stream bool, +) cliproxyexecutor.Options { + options := cliproxyexecutor.Options{ + Stream: stream, Headers: request.Headers.Clone(), + OriginalRequest: append([]byte(nil), request.OriginalRequest...), + SourceFormat: format, ResponseFormat: format, + } + if request.RequestPath != "" { + options.Metadata = map[string]any{ + cliproxyexecutor.RequestPathMetadataKey: request.RequestPath, + } + } + return options +} + func (e *CodexHTTPExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { executionCtx := e.executionContext(ctx, auth, nil, false) return e.inner.Execute(executionCtx, authWithoutProxyURL(auth), req, opts) @@ -720,14 +747,18 @@ func (t noRedirectRoundTripper) RoundTrip(req *http.Request) (*http.Response, er } type executionObservation struct { - capture bool - provider string - mu sync.RWMutex - effort string + capture bool + captureRequestPath bool + provider string + mu sync.RWMutex + effort string + observedRequestPath string } func newExecutionObservation(request ExecuteRequest) *executionObservation { - return newProviderExecutionObservation(request, ProviderCodex) + observation := newProviderExecutionObservation(request, ProviderCodex) + observation.captureRequestPath = true + return observation } func newProviderExecutionObservation(request ExecuteRequest, provider string) *executionObservation { @@ -742,7 +773,19 @@ func newProviderExecutionObservation(request ExecuteRequest, provider string) *e } func (o *executionObservation) observe(request *http.Request) { - if o == nil || !o.capture || request == nil || request.GetBody == nil || request.ContentLength > maxObservedBodyBytes { + if o == nil || request == nil { + return + } + if o.captureRequestPath { + requestPath := "" + if request.URL != nil { + requestPath = request.URL.Path + } + o.mu.Lock() + o.observedRequestPath = requestPath + o.mu.Unlock() + } + if !o.capture || request.GetBody == nil || request.ContentLength > maxObservedBodyBytes { return } bodyReader, err := request.GetBody() @@ -774,6 +817,15 @@ func (o *executionObservation) reasoningEffort() string { return o.effort } +func (o *executionObservation) upstreamRequestPath() string { + if o == nil { + return "" + } + o.mu.RLock() + defer o.mu.RUnlock() + return o.observedRequestPath +} + func exchangeToken(ctx context.Context, values url.Values, options Options) (CodexCredential, error) { if ctx == nil { ctx = context.Background() diff --git a/third_party/cpaembedded/embedded/embedded_test.go b/third_party/cpaembedded/embedded/embedded_test.go index 67df6783..8ce92730 100644 --- a/third_party/cpaembedded/embedded/embedded_test.go +++ b/third_party/cpaembedded/embedded/embedded_test.go @@ -75,6 +75,48 @@ func TestParseCodexCredentialJSONRejectsMalformedCPAControlMetadata(t *testing.T } } +func TestCodexExecutionOptionsCarryOnlyExplicitRequestPath(t *testing.T) { + t.Parallel() + + format := sdktranslator.FromString("openai-image") + options := codexExecutionOptions(ExecuteRequest{ + RequestPath: "/v1/images/generations", + Headers: http.Header{"Content-Type": {"application/json"}}, + }, format, true) + if options.SourceFormat != format || options.ResponseFormat != format || !options.Stream { + t.Fatalf("options formats/stream = %#v", options) + } + if got := options.Metadata[cliproxyexecutor.RequestPathMetadataKey]; got != "/v1/images/generations" { + t.Fatalf("request path metadata = %#v", options.Metadata) + } + + withoutPath := codexExecutionOptions(ExecuteRequest{}, format, false) + if _, exists := withoutPath.Metadata[cliproxyexecutor.RequestPathMetadataKey]; exists { + t.Fatalf("empty request path produced metadata: %#v", withoutPath.Metadata) + } +} + +func TestExecutionObservationCapturesRequestPathWithoutReasoning(t *testing.T) { + t.Parallel() + + observation := newExecutionObservation(ExecuteRequest{ + Format: "openai-image", + Payload: []byte(`{"model":"gpt-image-2","prompt":"draw"}`), + }) + if observation.capture { + t.Fatal("image request unexpectedly enabled reasoning observation") + } + request := httptest.NewRequest( + http.MethodPost, + "https://chatgpt.com/backend-api/codex/images/generations", + strings.NewReader(`{"model":"gpt-image-2"}`), + ) + observation.observe(request) + if got := observation.upstreamRequestPath(); got != "/backend-api/codex/images/generations" { + t.Fatalf("upstream request path = %q", got) + } +} + func TestParseCodexCredentialJSONRejectsInvalidOrDangerousInput(t *testing.T) { t.Parallel() diff --git a/web/src/api/control/protocols.ts b/web/src/api/control/protocols.ts index 9f882f97..4f82553b 100644 --- a/web/src/api/control/protocols.ts +++ b/web/src/api/control/protocols.ts @@ -7,6 +7,10 @@ export const protocolCatalog = [ value: 'openai-responses', supportsProtocolOnlyRouting: true, }, + { + value: 'openai-images', + supportsProtocolOnlyRouting: false, + }, { value: 'anthropic', supportsProtocolOnlyRouting: false, diff --git a/web/src/app/resources/channels.ts b/web/src/app/resources/channels.ts index bc7c4966..32e50edc 100644 --- a/web/src/app/resources/channels.ts +++ b/web/src/app/resources/channels.ts @@ -54,6 +54,8 @@ export type ChannelOperation = | 'responses_input_tokens' | 'count_tokens' | 'responses_passthrough' + | 'images_generate' + | 'images_edit' | 'list_models' | 'probe' @@ -154,6 +156,8 @@ const operations = [ 'responses_input_tokens', 'count_tokens', 'responses_passthrough', + 'images_generate', + 'images_edit', 'list_models', 'probe', ] as const diff --git a/web/src/app/resources/request-logs.ts b/web/src/app/resources/request-logs.ts index 83e50272..df03d93e 100644 --- a/web/src/app/resources/request-logs.ts +++ b/web/src/app/resources/request-logs.ts @@ -50,6 +50,8 @@ export type RequestLogOperation = | 'responses_input_tokens' | 'count_tokens' | 'responses_passthrough' + | 'images_generate' + | 'images_edit' | 'list_models' | 'probe' export type RequestLogRouteMode = 'native' | 'converted' @@ -234,6 +236,8 @@ const operations = [ 'responses_input_tokens', 'count_tokens', 'responses_passthrough', + 'images_generate', + 'images_edit', 'list_models', 'probe', ] as const diff --git a/web/src/app/resources/route-inspection.ts b/web/src/app/resources/route-inspection.ts index ff3bae7b..e86dc9fd 100644 --- a/web/src/app/resources/route-inspection.ts +++ b/web/src/app/resources/route-inspection.ts @@ -53,6 +53,8 @@ export type RouteInspectOperation = | 'responses_input_tokens' | 'count_tokens' | 'responses_passthrough' + | 'images_generate' + | 'images_edit' export type RouteInspectRequirement = 'any' | 'native' export type RouteInspectMode = 'native' | 'converted' @@ -109,6 +111,8 @@ export const routeInspectOperations = [ 'responses_input_tokens', 'count_tokens', 'responses_passthrough', + 'images_generate', + 'images_edit', ] as const export const routeInspectRequirements = ['any', 'native'] as const const routeModes = ['native', 'converted'] as const diff --git a/web/src/i18n/locales/en-US/access-keys.ts b/web/src/i18n/locales/en-US/access-keys.ts index fa2119a0..90e440cb 100644 --- a/web/src/i18n/locales/en-US/access-keys.ts +++ b/web/src/i18n/locales/en-US/access-keys.ts @@ -183,7 +183,7 @@ export default { remove: 'Remove cost limit rule', }, allGroupsAllowed: 'Currently allows routing to all Groups', - allProtocolsAllowed: 'Currently allows all four canonical protocols', + allProtocolsAllowed: 'Currently allows all five canonical protocols', allModelsAllowed: 'Currently allows every model routable by the target Groups', scopeExpansionWarning: 'Changing Specified to All expands client permissions. Confirm that the current scope is intended.', diff --git a/web/src/i18n/locales/en-US/monitor.ts b/web/src/i18n/locales/en-US/monitor.ts index 53578315..b5030040 100644 --- a/web/src/i18n/locales/en-US/monitor.ts +++ b/web/src/i18n/locales/en-US/monitor.ts @@ -388,7 +388,7 @@ export default { title: 'Route inspector', description: 'Inspect candidate Groups and credentials by protocol, model, and access key.', boundary: - 'Simulates a standard stateless request without conversation content or an affinity hit; Responses uses store:false. Read-only, with no upstream request or token usage.', + 'Simulates a standard stateless request without conversation content or an affinity hit; Responses uses store:false, while Images checks generation only. Read-only, with no upstream request or token usage.', routeModes: { native: 'Native', converted: 'Protocol conversion', @@ -404,6 +404,8 @@ export default { responses_input_tokens: 'Count input tokens', count_tokens: 'Count tokens', responses_passthrough: 'Responses extension operation', + images_generate: 'Generate image', + images_edit: 'Edit image', }, routeRequirements: { any: 'Allow protocol conversion (possibly lossy)', @@ -783,6 +785,8 @@ export default { responses_input_tokens: 'Count input tokens', count_tokens: 'Count tokens', responses_passthrough: 'Responses resource passthrough', + images_generate: 'Generate image', + images_edit: 'Edit image', list_models: 'List models', probe: 'Health probe', }, diff --git a/web/src/i18n/locales/ja-JP/access-keys.ts b/web/src/i18n/locales/ja-JP/access-keys.ts index bcd41e41..2a69dfb5 100644 --- a/web/src/i18n/locales/ja-JP/access-keys.ts +++ b/web/src/i18n/locales/ja-JP/access-keys.ts @@ -182,7 +182,7 @@ export default { remove: '費用上限ルールを削除', }, allGroupsAllowed: '現在はすべてのグループへのルーティングを許可', - allProtocolsAllowed: '現在は 4 つの正規プロトコルをすべて許可', + allProtocolsAllowed: '現在は 5 つの正規プロトコルをすべて許可', allModelsAllowed: '現在は対象グループがルーティングできるすべてのモデルを許可', scopeExpansionWarning: '「指定」から「すべて」への変更はクライアント権限を拡大します。現在の範囲が意図どおりか確認してください。', diff --git a/web/src/i18n/locales/ja-JP/monitor.ts b/web/src/i18n/locales/ja-JP/monitor.ts index a1f45fc9..44eb7d0f 100644 --- a/web/src/i18n/locales/ja-JP/monitor.ts +++ b/web/src/i18n/locales/ja-JP/monitor.ts @@ -388,7 +388,7 @@ export default { title: 'ルート検査', description: 'プロトコル、モデル、アクセスキーから候補グループと認証情報を確認します。', boundary: - '会話内容やアフィニティヒットを含まない標準的なステートレスリクエストを模擬します。Responses は store:false を使用します。読み取り専用で、上流送信や Token 消費はありません。', + '会話内容やアフィニティヒットを含まない標準的なステートレスリクエストを模擬します。Responses は store:false を使用し、Images は画像生成のみを確認します。読み取り専用で、上流送信や Token 消費はありません。', routeModes: { native: 'ネイティブ', converted: 'プロトコル変換', @@ -404,6 +404,8 @@ export default { responses_input_tokens: '入力 Token を計算', count_tokens: 'Token を計算', responses_passthrough: 'Responses 拡張操作', + images_generate: '画像を生成', + images_edit: '画像を編集', }, routeRequirements: { any: 'プロトコル変換を許可（損失の可能性あり）', @@ -782,6 +784,8 @@ export default { responses_input_tokens: '入力 Token を計算', count_tokens: 'Token を計算', responses_passthrough: 'Responses リソースの透過転送', + images_generate: '画像を生成', + images_edit: '画像を編集', list_models: 'モデル一覧', probe: 'ヘルスプローブ', }, diff --git a/web/src/i18n/locales/zh-CN/access-keys.ts b/web/src/i18n/locales/zh-CN/access-keys.ts index 18856918..a5eb1454 100644 --- a/web/src/i18n/locales/zh-CN/access-keys.ts +++ b/web/src/i18n/locales/zh-CN/access-keys.ts @@ -174,7 +174,7 @@ export default { remove: '删除费用规则', }, allGroupsAllowed: '当前允许路由到全部分组', - allProtocolsAllowed: '当前允许全部 4 种规范协议', + allProtocolsAllowed: '当前允许全部 5 种规范协议', allModelsAllowed: '当前允许目标分组可路由的全部模型', scopeExpansionWarning: '把“指定”改为“全部”会扩大客户端权限，请确认当前范围符合预期。', save: '保存访问密钥', diff --git a/web/src/i18n/locales/zh-CN/monitor.ts b/web/src/i18n/locales/zh-CN/monitor.ts index eb5a2f69..249d130b 100644 --- a/web/src/i18n/locales/zh-CN/monitor.ts +++ b/web/src/i18n/locales/zh-CN/monitor.ts @@ -372,7 +372,7 @@ export default { title: '路由检查', description: '按协议、模型与访问密钥检查当前候选分组和凭据。', boundary: - '按所选协议模拟不含会话内容与亲和命中的标准无状态请求；Responses 使用 store:false。只读检查，不会发送上游请求或消耗 Token。', + '按所选协议模拟不含会话内容与亲和命中的标准无状态请求；Responses 使用 store:false，Images 只检查图片生成。只读检查，不会发送上游请求或消耗 Token。', routeModes: { native: '原生', converted: '协议转换', @@ -388,6 +388,8 @@ export default { responses_input_tokens: '计算输入 Token', count_tokens: '计算 Token', responses_passthrough: 'Responses 扩展操作', + images_generate: '生成图片', + images_edit: '编辑图片', }, routeRequirements: { any: '允许协议转换（可有损）', @@ -764,6 +766,8 @@ export default { responses_input_tokens: '计算输入 Tokens', count_tokens: '计算 Tokens', responses_passthrough: 'Responses 资源透传', + images_generate: '生成图片', + images_edit: '编辑图片', list_models: '列出模型', probe: '健康探测', },