feat(images): 接入 OpenAI 图片生成与编辑 (#476)
* feat(images): 接入 OpenAI 图片生成与编辑 * fix(images): 补齐响应容量与边界处理 * fix(images): 修正协议观测与响应上限覆盖 * fix(stream): 保证取消优先于流关闭 * fix(images): 修正空 SSE 错误对象判定 * fix(images): 避免媒体数据触发凭据误判 * fix(images): 阻止流式响应泄漏凭据 * fix(images): 校验完整流事件后再转发 * fix(images): 修正响应模型观测
This commit is contained in:
1 parent
69483c7109
commit
e9e88dce99
71 files changed
+3432
-177
No files matched your search
@@ -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
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
|
||||
@@ -153,6 +153,7 @@ func TestAccessKeyCreateAcceptsAllEnabledProtocolsInCanonicalOrder(t *testing.T)
|
||||
Filters: &AccessKeyFilters{
|
||||
Protocols: []protocol.Protocol{
|
||||
protocol.Gemini,
|
||||
protocol.OpenAIImages,
|
||||
protocol.OpenAIResponses,
|
||||
protocol.Anthropic,
|
||||
protocol.OpenAICompletions,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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...)
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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:
|
||||
|
||||
@@ -305,6 +305,8 @@ func operationRequiresModel(operation Operation) bool {
|
||||
operation == OperationResponsesCompact ||
|
||||
operation == OperationResponsesInputTokens ||
|
||||
operation == OperationCountTokens ||
|
||||
operation == OperationImagesGenerate ||
|
||||
operation == OperationImagesEdit ||
|
||||
operation == OperationProbe
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -135,6 +135,7 @@ func modelListProtocols(value protocol.Protocol) []protocol.Protocol {
|
||||
return []protocol.Protocol{
|
||||
protocol.OpenAICompletions,
|
||||
protocol.OpenAIResponses,
|
||||
protocol.OpenAIImages,
|
||||
}
|
||||
}
|
||||
return []protocol.Protocol{value}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+74
-22
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -7,6 +7,10 @@ export const protocolCatalog = [
|
||||
value: 'openai-responses',
|
||||
supportsProtocolOnlyRouting: true,
|
||||
},
|
||||
{
|
||||
value: 'openai-images',
|
||||
supportsProtocolOnlyRouting: false,
|
||||
},
|
||||
{
|
||||
value: 'anthropic',
|
||||
supportsProtocolOnlyRouting: false,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.',
|
||||
|
||||
@@ -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',
|
||||
},
|
||||
|
||||
@@ -182,7 +182,7 @@ export default {
|
||||
remove: '費用上限ルールを削除',
|
||||
},
|
||||
allGroupsAllowed: '現在はすべてのグループへのルーティングを許可',
|
||||
allProtocolsAllowed: '現在は 4 つの正規プロトコルをすべて許可',
|
||||
allProtocolsAllowed: '現在は 5 つの正規プロトコルをすべて許可',
|
||||
allModelsAllowed: '現在は対象グループがルーティングできるすべてのモデルを許可',
|
||||
scopeExpansionWarning:
|
||||
'「指定」から「すべて」への変更はクライアント権限を拡大します。現在の範囲が意図どおりか確認してください。',
|
||||
|
||||
@@ -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: 'ヘルスプローブ',
|
||||
},
|
||||
|
||||
@@ -174,7 +174,7 @@ export default {
|
||||
remove: '删除费用规则',
|
||||
},
|
||||
allGroupsAllowed: '当前允许路由到全部分组',
|
||||
allProtocolsAllowed: '当前允许全部 4 种规范协议',
|
||||
allProtocolsAllowed: '当前允许全部 5 种规范协议',
|
||||
allModelsAllowed: '当前允许目标分组可路由的全部模型',
|
||||
scopeExpansionWarning: '把“指定”改为“全部”会扩大客户端权限,请确认当前范围符合预期。',
|
||||
save: '保存访问密钥',
|
||||
|
||||
@@ -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: '健康探测',
|
||||
},
|
||||
|
||||
Reference in new issue
Block a user