cleaned up deltaBuffer, added drain method

Author: pawbana

intercept/responses/base.go

@@ -38,10 +38,10 @@ type responsesInterceptionBase struct {
 	metrics metrics.Metrics
 }
 
-func NewResponsesService(baseURL string, key string, logger slog.Logger) responses.ResponseService {
+func NewResponsesService(baseURL string, apiKey string) responses.ResponseService {
 	opts := []option.RequestOption{
-		option.WithAPIKey(key),
 		option.WithBaseURL(baseURL),
+		option.WithAPIKey(apiKey),
 	}
 
 	return responses.NewResponseService(opts...)
@@ -99,7 +99,7 @@ func (i *responsesInterceptionBase) validateRequest(ctx context.Context, w http.
 	return nil
 }
 
-func (i *responsesInterceptionBase) requestOptions(respBody *deltaBuffer) []option.RequestOption {
+func (i *responsesInterceptionBase) requestOptions(payloadBuff *deltaBuffer) []option.RequestOption {
 	opts := []option.RequestOption{
 		// Sends original payload to solve json re-encoding issues
 		// eg. Codex CLI produces requests without ID set in reasoning items: https://platform.openai.com/docs/api-reference/responses/create#responses_create-input-input_item_list-item-reasoning-id
@@ -108,15 +108,15 @@ func (i *responsesInterceptionBase) requestOptions(respBody *deltaBuffer) []opti
 		option.WithRequestBody("application/json", i.reqPayload),
 
 		// Reads response body into given buffer
-		option.WithMiddleware(teeMiddleware(respBody)),
+		option.WithMiddleware(teeMiddleware(payloadBuff)),
 	}
 	if !i.req.Stream {
 		opts = append(opts, option.WithRequestTimeout(time.Second*60)) // TODO: configurable timeout
 	}
 	return opts
 }
 
-// handleUpstreamError checks error if it is an openAI error and if ProcessRequest should exit early.
+// handleUpstreamError checks if error is an openAI error and if caller should exit early.
 // If it is openAI responses error -> sets proper http response code and returns false to not exit early
 // response body will be sent using the same method as non-error response, using teeMiddleware -> deltaBuffer.
 // If it is a connection error or unnknown error -> returns given error + true to indicate early exit
@@ -142,22 +142,24 @@ func (i *responsesInterceptionBase) handleUpstreamError(ctx context.Context, ups
 }
 
 // teeMiddleware copies response body to given buffer leaving original response intact/consumable for openAI SDK
-func teeMiddleware(respBody io.Writer) func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) {
+func teeMiddleware(payloadBuff *deltaBuffer) func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) {
 	return func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) {
 		resp, err := next(req)
 		if err != nil || resp == nil || resp.Body == nil {
 			return resp, err
 		}
 
-		resp.Body = io.NopCloser(io.TeeReader(resp.Body, respBody))
+		payloadBuff.closer = resp.Body
+		resp.Body = io.NopCloser(io.TeeReader(resp.Body, payloadBuff))
 		return resp, nil
 	}
 }
 
-// deltaBuffer stores everything written to it and lets you read only the new bytes since last drain.
+// deltaBuffer is a thread safe byte buffer that supports reading incremental data (added after last read)
 type deltaBuffer struct {
-	mu  sync.Mutex
-	buf bytes.Buffer
+	mu     sync.Mutex
+	buf    bytes.Buffer
+	closer io.ReadCloser
 }
 
 func (d *deltaBuffer) Write(p []byte) (int, error) {
@@ -166,11 +168,16 @@ func (d *deltaBuffer) Write(p []byte) (int, error) {
 	return d.buf.Write(p)
 }
 
-func (d *deltaBuffer) Read(p []byte) (int, error) {
+// Reads all from original resqusts body so it is properly copied by TeeReader to buffer
+func (d *deltaBuffer) drain() error {
+	if d.closer == nil {
+		return nil
+	}
+
 	d.mu.Lock()
 	defer d.mu.Unlock()
-	c, err := d.buf.Read(p)
-	return c, err
+	_, err := io.ReadAll(d.closer)
+	return err
 }
 
 // readDelta returns only the bytes appended since the last readDelta call.

intercept/responses/blocking.go

@@ -3,7 +3,6 @@ package responses
 import (
 	"errors"
 	"fmt"
-	"io"
 	"net/http"
 
 	"cdr.dev/slog/v3"
@@ -45,25 +44,22 @@ func (i *BlockingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r *
 		return err
 	}
 
-	var respBody deltaBuffer
-	srv := NewResponsesService(i.baseURL, i.key, i.logger)
+	var respPayload deltaBuffer
+	srv := NewResponsesService(i.baseURL, i.key)
 
-	opts := i.requestOptions(&respBody)
-	_, err := srv.New(ctx, i.req.ResponseNewParams, opts...)
+	opts := i.requestOptions(&respPayload)
+	_, upstreamErr := srv.New(ctx, i.req.ResponseNewParams, opts...)
 
-	upstreamErr, earlyExit := i.handleUpstreamError(ctx, err, w)
+	upstreamErr, earlyExit := i.handleUpstreamError(ctx, upstreamErr, w)
 	if earlyExit {
 		return upstreamErr
 	}
 
-	w.Header().Set("Content-Type", "application/json")
-	out, err := io.ReadAll(&respBody)
-	if err != nil {
-		i.logger.Warn(ctx, "failed to read upstream response", slog.Error(err))
-		return errors.Join(fmt.Errorf("failed to read upstream response: %w", err), upstreamErr)
+	if err := respPayload.drain(); err != nil {
+		i.logger.Warn(ctx, "failed to drain original response body", slog.Error(err))
 	}
-
-	_, err = w.Write(out)
+	w.Header().Set("Content-Type", "application/json")
+	_, err := w.Write(respPayload.readDelta())
 	if err != nil {
 		i.logger.Warn(ctx, "failed to write response", slog.Error(err))
 		return errors.Join(fmt.Errorf("failed to write response: %w", err), upstreamErr)

intercept/responses/streaming.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"net/http"
 	"time"
 
@@ -63,15 +62,15 @@ func (i *StreamingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r
 		_ = events.Shutdown(streamCtx) // Catch-all in case it doesn't get shutdown after stream completes.
 	}()
 
-	var respBody deltaBuffer
+	var respPayload deltaBuffer
 
-	srv := NewResponsesService(i.baseURL, i.key, i.logger)
-	opts := i.requestOptions(&respBody)
+	srv := NewResponsesService(i.baseURL, i.key)
+	opts := i.requestOptions(&respPayload)
 	stream := srv.NewStreaming(ctx, i.req.ResponseNewParams, opts...)
 	defer stream.Close()
 
 	for stream.Next() {
-		if err := events.Send(ctx, respBody.readDelta()); err != nil {
+		if err := events.Send(ctx, respPayload.readDelta()); err != nil {
 			i.logger.Warn(ctx, "failed to relay chunk", slog.Error(err))
 			err = fmt.Errorf("relay chunk: %w", err)
 			stream.Close()
@@ -84,13 +83,10 @@ func (i *StreamingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r
 		return upstreamErr
 	}
 
-	// Sometimes stream.Next() returns before respBody buffer is filled
-	lastRead, err := io.ReadAll(&respBody)
-	if err != nil {
-		i.logger.Warn(ctx, "failed to read upstream response", slog.Error(err))
-		return fmt.Errorf("failed to read upstream response: %w", err)
+	if err := respPayload.drain(); err != nil {
+		i.logger.Warn(ctx, "failed to drain original response body", slog.Error(err))
 	}
-	events.Send(ctx, lastRead)
+	events.Send(ctx, respPayload.readDelta())
 
 	// Give the events stream 30 seconds (TODO: configurable) to gracefully shutdown.
 	shutdownCtx, shutdownCancel := context.WithTimeout(ctx, time.Second*30)