final cleanup

Author: angus-langchain

internal/aggregator/aggregator.go

@@ -91,11 +91,11 @@ func (a *Aggregator) worker(ctx context.Context, ch <-chan *model.Run) {
 		}
 		comp, boundary, _, err := sc.Close()
 		if err != nil {
-			slog.Error("serialize", "err", err)
+			slog.Error("Failed to flush runs", "err", err)
 			return
 		}
 		if len(comp) > 0 {
-			go a.up.Send(ctx, uploader.Batch{Data: comp, Boundary: boundary})
+			go a.up.Send(context.Background(), uploader.Batch{Data: comp, Boundary: boundary})
 		}
 	}
 
@@ -104,6 +104,7 @@ func (a *Aggregator) worker(ctx context.Context, ch <-chan *model.Run) {
 
 		scMu.Lock()
 		if err := sc.AddRun(r); err != nil {
+			slog.Error("failed to queue run", "err", err)
 			scMu.Unlock()
 			return
 		}

internal/handler/traces.go

@@ -66,12 +66,12 @@ func TracesHandler(maxBody int64, tr *translator.Translator, ch chan *model.Run)
 
 		switch {
 		case strings.Contains(ct, "json"):
-			if err := decodeJSON(r.Body, &req); err != nil {
+			if err := decodeJSON(reader, &req); err != nil {
 				handleError(w, r, http.StatusBadRequest, err)
 				return
 			}
 		case strings.Contains(ct, "protobuf"), strings.HasSuffix(ct, "+proto"):
-			if err := decodeProto(r.Body, &req); err != nil {
+			if err := decodeProto(reader, &req); err != nil {
 				handleError(w, r, http.StatusBadRequest, err)
 				return
 			}

internal/serializer/stream.go

@@ -5,9 +5,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"log/slog"
+	"math/rand/v2"
+	"strconv"
 
 	"github.com/DataDog/zstd"
-	"github.com/google/uuid"
 	"github.com/langchain-ai/langsmith-collector-proxy/internal/model"
 )
 
@@ -29,12 +31,12 @@ type StreamingCompressor struct {
 }
 
 func NewStreamingCompressor() *StreamingCompressor {
-	var buf bytes.Buffer
-	zw := zstd.NewWriter(&buf)
+	buf := &bytes.Buffer{}
+	zw := zstd.NewWriter(buf)
 	return &StreamingCompressor{
-		boundary: "BOUNDARY-" + uuid.NewString(),
+		boundary: "----LangSmithFormBoundary-" + strconv.FormatUint(rand.Uint64(), 36),
 		w:        zw,
-		buf:      &buf,
+		buf:      buf,
 	}
 }
 
@@ -78,24 +80,32 @@ func (sc *StreamingCompressor) AddRun(r *model.Run) error {
 }
 
 func (sc *StreamingCompressor) Close() ([]byte, string, int, error) {
-	// Write the final multipart boundary.
-	if sc.runCount == 0 || sc.uncompressed == 0 {
-		return nil, "", 0, nil
-	}
-	if _, err := sc.w.Write([]byte(fmt.Sprintf("--%s--\r\n", sc.boundary))); err != nil {
-		return nil, "", 0, err
-	}
-	// Close the zstd writer.
-	if err := sc.w.Close(); err != nil {
-		return nil, "", 0, err
+	// Write the final multipart boundary and reset state.
+	hasData := sc.runCount > 0 || sc.uncompressed > 0
+	var (
+		outBytes     []byte
+		boundary     string
+		uncompressed int
+		err          error
+	)
+
+	if hasData {
+		if _, err = sc.w.Write([]byte(fmt.Sprintf("--%s--\r\n", sc.boundary))); err == nil {
+			err = sc.w.Close()
+		}
+		if err == nil {
+			outBytes = sc.buf.Bytes()
+			boundary = sc.boundary
+			uncompressed = sc.uncompressed
+		}
+	} else {
+		slog.Error("Failed to write final multipart boundary", "err", err)
 	}
-	out := sc.buf.Bytes()
-	uncompressed := sc.uncompressed
-	// reset the buffer and writer for the next upload
-	sc.buf = bytes.NewBuffer(nil)
+	sc.buf = &bytes.Buffer{}
+	sc.boundary = "----LangSmithFormBoundary-" + strconv.FormatUint(rand.Uint64(), 36)
 	sc.w = zstd.NewWriter(sc.buf)
-	sc.runCount, sc.uncompressed = 0, 0
-	return out, sc.boundary, uncompressed, nil
+	sc.runCount = 0
+	return outBytes, boundary, uncompressed, err
 }
 
 func (sc *StreamingCompressor) Uncompressed() int {