From 47c3fcd3c821de618d1d4dabe7587c92eb1c829d Mon Sep 17 00:00:00 2001 From: nnNyx Date: Wed, 17 Jun 2026 15:08:39 +0200 Subject: [PATCH] fix(openai): process custom responses stream chunks --- core/providers/openai/openai.go | 164 +++++++------ core/providers/openai/responsesstream_test.go | 219 ++++++++++++++++++ 2 files changed, 309 insertions(+), 74 deletions(-) create mode 100644 core/providers/openai/responsesstream_test.go diff --git a/core/providers/openai/openai.go b/core/providers/openai/openai.go index 4230bf7c6a1..8235e93cb36 100644 --- a/core/providers/openai/openai.go +++ b/core/providers/openai/openai.go @@ -1857,18 +1857,26 @@ func HandleOpenAIResponsesStreaming( // Parse into bifrost response var response schemas.BifrostResponsesStreamResponse - // TODO fix this + var rawRequest interface{} + var rawResponse interface{} + var bifrostErr *schemas.BifrostError + preserveCustomRawErrorFields := func(err *schemas.BifrostError) *schemas.BifrostError { + if err == nil { + return nil + } + if sendBackRawRequest && rawRequest != nil { + err.ExtraFields.RawRequest = rawRequest + } + if sendBackRawResponse && rawResponse != nil { + err.ExtraFields.RawResponse = rawResponse + } + return err + } if customResponseHandler != nil { - rawRequest, rawResponse, bifrostErr := customResponseHandler([]byte(jsonData), &response, nil, false, false) + rawRequest, rawResponse, bifrostErr = customResponseHandler([]byte(jsonData), &response, jsonBody, sendBackRawRequest, sendBackRawResponse) if bifrostErr != nil { - if sendBackRawRequest { - bifrostErr.ExtraFields.RawRequest = rawRequest - } - if sendBackRawResponse { - bifrostErr.ExtraFields.RawResponse = rawResponse - } ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, providerUtils.EnrichError(ctx, bifrostErr, jsonBody, nil, sendBackRawRequest, sendBackRawResponse, latency), responseChan, logger, postHookSpanFinalizer) + providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, preserveCustomRawErrorFields(providerUtils.EnrichError(ctx, bifrostErr, jsonBody, nil, sendBackRawRequest, sendBackRawResponse, latency)), responseChan, logger, postHookSpanFinalizer) return } } else { @@ -1876,91 +1884,99 @@ func HandleOpenAIResponsesStreaming( logger.Warn("Failed to parse stream response: %v", err) continue } + } - if postResponseConverter != nil { - if converted := postResponseConverter(&response); converted != nil { - response = *converted - } else { - logger.Warn("postResponseConverter returned nil; leaving chunk unmodified") - } + if postResponseConverter != nil { + if converted := postResponseConverter(&response); converted != nil { + response = *converted + } else { + logger.Warn("postResponseConverter returned nil; leaving chunk unmodified") } + } - if sendBackRawResponse { + if sendBackRawResponse { + if rawResponse != nil { + response.ExtraFields.RawResponse = rawResponse + } else { response.ExtraFields.RawResponse = jsonData } + } - if response.Type == schemas.ResponsesStreamResponseTypeError { - bifrostErr := &schemas.BifrostError{ - Type: schemas.Ptr(string(schemas.ResponsesStreamResponseTypeError)), - IsBifrostError: false, - Error: &schemas.ErrorField{}, - } + if response.Type == schemas.ResponsesStreamResponseTypeError { + bifrostErr := &schemas.BifrostError{ + Type: schemas.Ptr(string(schemas.ResponsesStreamResponseTypeError)), + IsBifrostError: false, + Error: &schemas.ErrorField{}, + } - if response.Message != nil { - bifrostErr.Error.Message = *response.Message - } - if response.Param != nil { - bifrostErr.Error.Param = *response.Param + if response.Message != nil { + bifrostErr.Error.Message = *response.Message + } + if response.Param != nil { + bifrostErr.Error.Param = *response.Param + } + if response.Code != nil { + bifrostErr.Error.Code = response.Code + } + if response.Error != nil { + if response.Error.Message != "" && bifrostErr.Error.Message == "" { + bifrostErr.Error.Message = response.Error.Message } - if response.Code != nil { - bifrostErr.Error.Code = response.Code + if response.Error.Code != "" && (bifrostErr.Error.Code == nil || *bifrostErr.Error.Code == "") { + bifrostErr.Error.Code = &response.Error.Code } - if response.Error != nil { - if response.Error.Message != "" && bifrostErr.Error.Message == "" { - bifrostErr.Error.Message = response.Error.Message - } - if response.Error.Code != "" && (bifrostErr.Error.Code == nil || *bifrostErr.Error.Code == "") { - bifrostErr.Error.Code = &response.Error.Code - } + } + if response.Response != nil && response.Response.Error != nil { + if response.Response.Error.Message != "" && bifrostErr.Error.Message == "" { + bifrostErr.Error.Message = response.Response.Error.Message } - if response.Response != nil && response.Response.Error != nil { - if response.Response.Error.Message != "" && bifrostErr.Error.Message == "" { - bifrostErr.Error.Message = response.Response.Error.Message - } - if response.Response.Error.Code != "" && (bifrostErr.Error.Code == nil || *bifrostErr.Error.Code == "") { - bifrostErr.Error.Code = schemas.Ptr(response.Response.Error.Code) - } + if response.Response.Error.Code != "" && (bifrostErr.Error.Code == nil || *bifrostErr.Error.Code == "") { + bifrostErr.Error.Code = schemas.Ptr(response.Response.Error.Code) } - - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, providerUtils.EnrichError(ctx, bifrostErr, jsonBody, []byte(jsonData), sendBackRawRequest, sendBackRawResponse, latency), responseChan, logger, postHookSpanFinalizer) - return } - // Some providers (e.g. Fireworks) send response.failed on HTTP 200 streams - // instead of a pre-stream 4xx. Convert to BifrostError for consistent handling. - if response.Type == schemas.ResponsesStreamResponseTypeFailed { - bifrostErr := &schemas.BifrostError{ - Type: schemas.Ptr(string(schemas.ResponsesStreamResponseTypeFailed)), - IsBifrostError: false, - Error: &schemas.ErrorField{}, - } - if response.Response != nil && response.Response.Error != nil { - bifrostErr.Error.Message = response.Response.Error.Message - bifrostErr.Error.Code = &response.Response.Error.Code - } - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, providerUtils.EnrichError(ctx, bifrostErr, jsonBody, []byte(jsonData), sendBackRawRequest, sendBackRawResponse, latency), responseChan, logger, postHookSpanFinalizer) - return + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, preserveCustomRawErrorFields(providerUtils.EnrichError(ctx, bifrostErr, jsonBody, []byte(jsonData), sendBackRawRequest, sendBackRawResponse, latency)), responseChan, logger, postHookSpanFinalizer) + return + } + + // Some providers (e.g. Fireworks) send response.failed on HTTP 200 streams + // instead of a pre-stream 4xx. Convert to BifrostError for consistent handling. + if response.Type == schemas.ResponsesStreamResponseTypeFailed { + bifrostErr := &schemas.BifrostError{ + Type: schemas.Ptr(string(schemas.ResponsesStreamResponseTypeFailed)), + IsBifrostError: false, + Error: &schemas.ErrorField{}, } + if response.Response != nil && response.Response.Error != nil { + bifrostErr.Error.Message = response.Response.Error.Message + bifrostErr.Error.Code = &response.Response.Error.Code + } + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + providerUtils.ProcessAndSendBifrostError(ctx, postHookRunner, preserveCustomRawErrorFields(providerUtils.EnrichError(ctx, bifrostErr, jsonBody, []byte(jsonData), sendBackRawRequest, sendBackRawResponse, latency)), responseChan, logger, postHookSpanFinalizer) + return + } - response.ExtraFields.ChunkIndex = response.SequenceNumber - if response.Type == schemas.ResponsesStreamResponseTypeCompleted || response.Type == schemas.ResponsesStreamResponseTypeIncomplete { - // Set raw request if enabled - if sendBackRawRequest { + response.ExtraFields.ChunkIndex = response.SequenceNumber + if response.Type == schemas.ResponsesStreamResponseTypeCompleted || response.Type == schemas.ResponsesStreamResponseTypeIncomplete { + // Set raw request if enabled + if sendBackRawRequest { + if rawRequest != nil { + response.ExtraFields.RawRequest = rawRequest + } else { providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonBody) } - response.ExtraFields.Latency = time.Since(startTime).Milliseconds() - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, nil, &response, nil, nil, nil), responseChan, postHookSpanFinalizer) - return } - - response.ExtraFields.Latency = time.Since(lastChunkTime).Milliseconds() - lastChunkTime = time.Now() - + response.ExtraFields.Latency = time.Since(startTime).Milliseconds() + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, nil, &response, nil, nil, nil), responseChan, postHookSpanFinalizer) + return } + + response.ExtraFields.Latency = time.Since(lastChunkTime).Milliseconds() + lastChunkTime = time.Now() + + providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, nil, &response, nil, nil, nil), responseChan, postHookSpanFinalizer) } }() diff --git a/core/providers/openai/responsesstream_test.go b/core/providers/openai/responsesstream_test.go new file mode 100644 index 00000000000..ab6d9e29be5 --- /dev/null +++ b/core/providers/openai/responsesstream_test.go @@ -0,0 +1,219 @@ +package openai + +import ( + "context" + "net" + "testing" + "time" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/valyala/fasthttp" + "github.com/valyala/fasthttp/fasthttputil" +) + +func TestHandleOpenAIResponsesStreaming_CustomHandlerSendsCompletedChunkWithRawFields(t *testing.T) { + ln := fasthttputil.NewInmemoryListener() + server := &fasthttp.Server{ + Handler: func(ctx *fasthttp.RequestCtx) { + ctx.SetStatusCode(fasthttp.StatusOK) + ctx.Response.Header.SetContentType("text/event-stream") + ctx.SetBodyString("data: {\"provider\":\"chunk\"}\n\n") + }, + } + go func() { _ = server.Serve(ln) }() + defer ln.Close() + + client := &fasthttp.Client{ + Dial: func(addr string) (net.Conn, error) { + return ln.Dial() + }, + ReadTimeout: time.Second, + WriteTimeout: time.Second, + } + + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyTracer, &schemas.NoOpTracer{}) + + customHandlerCalled := false + customHandler := func(responseBody []byte, response *schemas.BifrostResponsesStreamResponse, requestBody []byte, sendBackRawRequest bool, sendBackRawResponse bool) (interface{}, interface{}, *schemas.BifrostError) { + customHandlerCalled = true + if !sendBackRawRequest { + t.Fatal("expected custom handler to receive sendBackRawRequest=true") + } + if !sendBackRawResponse { + t.Fatal("expected custom handler to receive sendBackRawResponse=true") + } + if len(requestBody) == 0 { + t.Fatal("expected serialized request body") + } + if string(responseBody) != `{"provider":"chunk"}` { + t.Fatalf("unexpected SSE payload: %s", responseBody) + } + + response.Type = schemas.ResponsesStreamResponseTypeCompleted + response.SequenceNumber = 7 + return map[string]interface{}{"request": "raw"}, map[string]interface{}{"response": "raw"}, nil + } + + stream, bifrostErr := HandleOpenAIResponsesStreaming( + ctx, + client, + "http://test/v1/responses", + &schemas.BifrostResponsesRequest{Provider: schemas.OpenAI, Model: "gpt-test"}, + nil, + nil, + 1, + true, + true, + schemas.OpenAI, + func(_ *schemas.BifrostContext, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) { + return result, err + }, + customHandler, + nil, + nil, + nil, + nil, + noopOpenAITestLogger{}, + nil, + ) + if bifrostErr != nil { + t.Fatalf("unexpected setup error: %v", bifrostErr.Error.Message) + } + + chunk, ok := <-stream + if !ok { + t.Fatal("expected completed stream chunk") + } + if !customHandlerCalled { + t.Fatal("expected custom response handler to be called") + } + if chunk.BifrostError != nil { + t.Fatalf("unexpected stream error: %v", chunk.BifrostError.Error.Message) + } + if chunk.BifrostResponsesStreamResponse == nil { + t.Fatal("expected responses stream response") + } + response := chunk.BifrostResponsesStreamResponse + if response.Type != schemas.ResponsesStreamResponseTypeCompleted { + t.Fatalf("expected completed chunk, got %q", response.Type) + } + if response.ExtraFields.ChunkIndex != 7 { + t.Fatalf("expected chunk index 7, got %d", response.ExtraFields.ChunkIndex) + } + if got := response.ExtraFields.RawRequest.(map[string]interface{})["request"]; got != "raw" { + t.Fatalf("expected custom raw request, got %#v", response.ExtraFields.RawRequest) + } + if got := response.ExtraFields.RawResponse.(map[string]interface{})["response"]; got != "raw" { + t.Fatalf("expected custom raw response, got %#v", response.ExtraFields.RawResponse) + } + if _, ok := <-stream; ok { + t.Fatal("expected stream to close after completed chunk") + } +} + +func TestHandleOpenAIResponsesStreaming_CustomHandlerPreservesRawFieldsOnErrorEvent(t *testing.T) { + ln := fasthttputil.NewInmemoryListener() + server := &fasthttp.Server{ + Handler: func(ctx *fasthttp.RequestCtx) { + ctx.SetStatusCode(fasthttp.StatusOK) + ctx.Response.Header.SetContentType("text/event-stream") + ctx.SetBodyString("data: {\"provider\":\"error\"}\n\n") + }, + } + go func() { _ = server.Serve(ln) }() + defer ln.Close() + + client := &fasthttp.Client{ + Dial: func(addr string) (net.Conn, error) { + return ln.Dial() + }, + ReadTimeout: time.Second, + WriteTimeout: time.Second, + } + + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyTracer, &schemas.NoOpTracer{}) + + customHandler := func(responseBody []byte, response *schemas.BifrostResponsesStreamResponse, requestBody []byte, sendBackRawRequest bool, sendBackRawResponse bool) (interface{}, interface{}, *schemas.BifrostError) { + if !sendBackRawRequest { + t.Fatal("expected custom handler to receive sendBackRawRequest=true") + } + if !sendBackRawResponse { + t.Fatal("expected custom handler to receive sendBackRawResponse=true") + } + if len(requestBody) == 0 { + t.Fatal("expected serialized request body") + } + if string(responseBody) != `{"provider":"error"}` { + t.Fatalf("unexpected SSE payload: %s", responseBody) + } + + response.Type = schemas.ResponsesStreamResponseTypeError + response.Message = schemas.Ptr("custom stream error") + response.Code = schemas.Ptr("custom_error") + return map[string]interface{}{"request": "raw"}, map[string]interface{}{"response": "raw"}, nil + } + + stream, bifrostErr := HandleOpenAIResponsesStreaming( + ctx, + client, + "http://test/v1/responses", + &schemas.BifrostResponsesRequest{Provider: schemas.OpenAI, Model: "gpt-test"}, + nil, + nil, + 1, + true, + true, + schemas.OpenAI, + func(_ *schemas.BifrostContext, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) { + return result, err + }, + customHandler, + nil, + nil, + nil, + nil, + noopOpenAITestLogger{}, + nil, + ) + if bifrostErr != nil { + t.Fatalf("unexpected setup error: %v", bifrostErr.Error.Message) + } + + chunk, ok := <-stream + if !ok { + t.Fatal("expected error stream chunk") + } + if chunk.BifrostError == nil { + t.Fatal("expected bifrost error") + } + if chunk.BifrostError.Error.Message != "custom stream error" { + t.Fatalf("expected custom stream error, got %q", chunk.BifrostError.Error.Message) + } + if chunk.BifrostError.Error.Code == nil || *chunk.BifrostError.Error.Code != "custom_error" { + t.Fatalf("expected custom_error code, got %#v", chunk.BifrostError.Error.Code) + } + if got := chunk.BifrostError.ExtraFields.RawRequest.(map[string]interface{})["request"]; got != "raw" { + t.Fatalf("expected custom raw request, got %#v", chunk.BifrostError.ExtraFields.RawRequest) + } + if got := chunk.BifrostError.ExtraFields.RawResponse.(map[string]interface{})["response"]; got != "raw" { + t.Fatalf("expected custom raw response, got %#v", chunk.BifrostError.ExtraFields.RawResponse) + } + if _, ok := <-stream; ok { + t.Fatal("expected stream to close after error chunk") + } +} + +type noopOpenAITestLogger struct{} + +func (noopOpenAITestLogger) Debug(string, ...any) {} +func (noopOpenAITestLogger) Info(string, ...any) {} +func (noopOpenAITestLogger) Warn(string, ...any) {} +func (noopOpenAITestLogger) Error(string, ...any) {} +func (noopOpenAITestLogger) Fatal(string, ...any) {} +func (noopOpenAITestLogger) SetLevel(schemas.LogLevel) {} +func (noopOpenAITestLogger) SetOutputType(schemas.LoggerOutputType) {} +func (noopOpenAITestLogger) LogHTTPRequest(schemas.LogLevel, string) schemas.LogEventBuilder { + return schemas.NoopLogEvent +}