From e5d460d936c58aa39924c233703ce2ccaa75cb14 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Sat, 15 Aug 2026 07:11:16 +0000 Subject: [PATCH 1/9] fix: stream multipart uploads without buffering --- pkg/cmd/flagoptions.go | 35 ++-- pkg/cmd/multipartbody.go | 142 +++++++++++++++ pkg/cmd/multipartbody_test.go | 322 ++++++++++++++++++++++++++++++++++ 3 files changed, 483 insertions(+), 16 deletions(-) create mode 100644 pkg/cmd/multipartbody.go create mode 100644 pkg/cmd/multipartbody_test.go diff --git a/pkg/cmd/flagoptions.go b/pkg/cmd/flagoptions.go index d257e389..c3b30d92 100644 --- a/pkg/cmd/flagoptions.go +++ b/pkg/cmd/flagoptions.go @@ -1,14 +1,13 @@ package cmd import ( - "bytes" "encoding/base64" "encoding/json" + "errors" "fmt" "io" "maps" "mime" - "mime/multipart" "net/http" "os" "path/filepath" @@ -121,7 +120,7 @@ func embedFilesValue(v reflect.Value, embedStyle FileEmbedStyle, stdin *onceStdi val := iter.Value() newVal, err := embedFilesValue(val, embedStyle, stdin) if err != nil { - return reflect.Value{}, err + return reflect.Value{}, errors.Join(err, closeFileUploads(result.Interface())) } result.SetMapIndex(key, newVal) } @@ -136,7 +135,7 @@ func embedFilesValue(v reflect.Value, embedStyle FileEmbedStyle, stdin *onceStdi for i := 0; i < v.Len(); i++ { newVal, err := embedFilesValue(v.Index(i), embedStyle, stdin) if err != nil { - return reflect.Value{}, err + return reflect.Value{}, errors.Join(err, closeFileUploads(result.Slice(0, i).Interface())) } result.Index(i).Set(newVal) } @@ -324,8 +323,16 @@ func flagOptions( // This parameter is true if stdin is already in use to pass a binary parameter by using the special value // "-". In this case, we won't attempt to read it as a JSON/YAML blob for options setting. ignoreStdin bool, -) ([]option.RequestOption, error) { - var options []option.RequestOption +) (options []option.RequestOption, err error) { + // Once multipart files have been opened, close them on every error path + // until ownership is transferred to multipartRequestBody below. + var pendingMultipartUploads any + defer func() { + if err != nil && pendingMultipartUploads != nil { + err = errors.Join(err, closeFileUploads(pendingMultipartUploads)) + } + }() + if cmd.Bool("debug") { options = append(options, option.WithMiddleware(debugmiddleware.NewRequestLogger().Middleware())) } @@ -434,6 +441,9 @@ func flagOptions( return nil, err } else { requestContents.Body = embedded + if bodyType == MultipartFormEncoded { + pendingMultipartUploads = embedded + } } if headersWithFiles, err := embedFiles(requestContents.Headers, EmbedText, &stdinReader); err != nil { @@ -492,22 +502,15 @@ func flagOptions( case EmptyBody: break case MultipartFormEncoded: - buf := new(bytes.Buffer) - writer := multipart.NewWriter(buf) - // For multipart/form-encoded, we need a map structure bodyMap, ok := requestContents.Body.(map[string]any) if !ok { return nil, fmt.Errorf("Cannot send a non-map value to a form-encoded endpoint: %v\n", requestContents.Body) } encodingFormat := apiform.FormatBrackets - if err := apiform.MarshalWithSettings(bodyMap, writer, encodingFormat); err != nil { - return nil, err - } - if err := writer.Close(); err != nil { - return nil, err - } - options = append(options, option.WithRequestBody(writer.FormDataContentType(), buf)) + body := newMultipartRequestBody(bodyMap, encodingFormat) + pendingMultipartUploads = nil // multipartRequestBody owns them now. + options = append(options, multipartRequestOptions(body)...) case ApplicationJSON: bodyBytes, err := json.Marshal(requestContents.Body) diff --git a/pkg/cmd/multipartbody.go b/pkg/cmd/multipartbody.go new file mode 100644 index 00000000..f33562d5 --- /dev/null +++ b/pkg/cmd/multipartbody.go @@ -0,0 +1,142 @@ +package cmd + +import ( + "errors" + "io" + "mime/multipart" + "reflect" + "sync" + + "github.com/openai/openai-cli/internal/apiform" + "github.com/openai/openai-go/v3/option" +) + +// multipartRequestBody streams multipart encoding through an io.Pipe. Encoding +// starts on the first Read so constructing request options never reads upload +// data or leaves an encoder goroutine blocked before the HTTP client takes +// ownership of the body. +type multipartRequestBody struct { + reader *io.PipeReader + writer *io.PipeWriter + multipartWriter *multipart.Writer + bodyMap map[string]any + encodingFormat apiform.FormFormat + contentType string + + startOnce sync.Once + cleanupOnce sync.Once + done chan struct{} + cleanupErr error +} + +var _ io.ReadCloser = (*multipartRequestBody)(nil) + +func newMultipartRequestBody(bodyMap map[string]any, encodingFormat apiform.FormFormat) *multipartRequestBody { + reader, writer := io.Pipe() + multipartWriter := multipart.NewWriter(writer) + return &multipartRequestBody{ + reader: reader, + writer: writer, + multipartWriter: multipartWriter, + bodyMap: bodyMap, + encodingFormat: encodingFormat, + contentType: multipartWriter.FormDataContentType(), + done: make(chan struct{}), + } +} + +func (b *multipartRequestBody) Read(p []byte) (int, error) { + b.start() + return b.reader.Read(p) +} + +func (b *multipartRequestBody) Close() error { + // Closing the read side releases an encoder blocked on a pipe write. Closing + // owned uploads also releases an encoder blocked while reading a local file. + _ = b.reader.Close() + cleanupErr := b.cleanup() + + // Starting after the read side is closed is safe and ensures the writer and + // done channel are finalized even when the body was closed before its first + // Read. The encoder's first pipe write fails immediately in that case. + b.start() + return cleanupErr +} + +func (b *multipartRequestBody) ContentType() string { + return b.contentType +} + +func (b *multipartRequestBody) start() { + b.startOnce.Do(func() { + go b.encode() + }) +} + +func (b *multipartRequestBody) encode() { + defer close(b.done) + + encodeErr := apiform.MarshalWithSettings(b.bodyMap, b.multipartWriter, b.encodingFormat) + writerCloseErr := b.multipartWriter.Close() + cleanupErr := b.cleanup() + + _ = b.writer.CloseWithError(errors.Join(encodeErr, writerCloseErr, cleanupErr)) +} + +func (b *multipartRequestBody) cleanup() error { + b.cleanupOnce.Do(func() { + b.cleanupErr = closeFileUploads(b.bodyMap) + }) + return b.cleanupErr +} + +func multipartRequestOptions(body *multipartRequestBody) []option.RequestOption { + return []option.RequestOption{ + option.WithRequestBody(body.ContentType(), body), + // A streamed multipart body cannot be replayed. Make the behavior + // explicit even though the SDK also avoids retrying bodies without + // an http.Request.GetBody function. + option.WithMaxRetries(0), + } +} + +// closeFileUploads closes every file opened by openFileUpload, including files +// nested inside maps and arrays. Other io.ReadClosers (notably stdin wrappers) +// are intentionally left alone because this code does not own them. +func closeFileUploads(value any) error { + return closeFileUploadsValue(reflect.ValueOf(value)) +} + +func closeFileUploadsValue(value reflect.Value) error { + if !value.IsValid() { + return nil + } + + for value.Kind() == reflect.Interface || value.Kind() == reflect.Pointer { + if value.IsNil() { + return nil + } + value = value.Elem() + } + + if value.CanInterface() { + if upload, ok := value.Interface().(fileUpload); ok { + return upload.Close() + } + } + + var errs []error + switch value.Kind() { + case reflect.Map: + iter := value.MapRange() + for iter.Next() { + errs = append(errs, closeFileUploadsValue(iter.Value())) + } + case reflect.Slice, reflect.Array: + for i := 0; i < value.Len(); i++ { + errs = append(errs, closeFileUploadsValue(value.Index(i))) + } + } + + return errors.Join(errs...) +} diff --git a/pkg/cmd/multipartbody_test.go b/pkg/cmd/multipartbody_test.go new file mode 100644 index 00000000..3b92be40 --- /dev/null +++ b/pkg/cmd/multipartbody_test.go @@ -0,0 +1,322 @@ +package cmd + +import ( + "context" + "errors" + "io" + "mime" + "mime/multipart" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/openai/openai-cli/internal/apiform" + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/option" + "github.com/stretchr/testify/require" +) + +const multipartTestTimeout = 5 * time.Second + +func TestMultipartRequestBodyStreamsWithBackpressure(t *testing.T) { + const uploadSize = int64(64 << 20) + source := &countingReader{remaining: uploadSize} + upload := &recordingReadCloser{reader: source} + body := newMultipartRequestBody(map[string]any{ + "file": fileUpload{ + Reader: upload, + filename: "large.bin", + contentType: "application/octet-stream", + }, + }, apiform.FormatBrackets) + + // Construction is fully lazy: it neither reads the upload nor starts a + // producer that can be left blocked before an HTTP request consumes it. + require.Zero(t, source.bytesRead.Load()) + + mediaType, params, err := mime.ParseMediaType(body.ContentType()) + require.NoError(t, err) + require.Equal(t, "multipart/form-data", mediaType) + + part, err := multipart.NewReader(body, params["boundary"]).NextPart() + require.NoError(t, err) + require.Equal(t, "file", part.FormName()) + require.Equal(t, "large.bin", part.FileName()) + + buf := make([]byte, 1) + n, err := part.Read(buf) + require.NoError(t, err) + require.Equal(t, 1, n) + require.Positive(t, source.bytesRead.Load()) + require.Less(t, source.bytesRead.Load(), uploadSize, + "the producer must not read the complete upload ahead of the consumer") + + require.NoError(t, body.Close()) + waitForMultipartBody(t, body) + require.Equal(t, int32(1), upload.closeCount.Load()) +} + +func TestMultipartRequestBodyPropagatesProducerErrors(t *testing.T) { + readErr := errors.New("upload read failed") + upload := &recordingReadCloser{reader: errorReader{err: readErr}} + body := newMultipartRequestBody(map[string]any{ + "file": fileUpload{Reader: upload, filename: "broken.bin"}, + }, apiform.FormatBrackets) + + _, err := io.ReadAll(body) + require.ErrorIs(t, err, readErr) + waitForMultipartBody(t, body) + require.Equal(t, int32(1), upload.closeCount.Load()) + require.NoError(t, body.Close()) +} + +func TestMultipartRequestBodyClosesOpenedFiles(t *testing.T) { + path := filepath.Join(t.TempDir(), "upload.bin") + require.NoError(t, os.WriteFile(path, []byte("payload"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + file := upload.Reader.(*os.File) + body := newMultipartRequestBody(map[string]any{"file": upload}, apiform.FormatBrackets) + + _, err = io.ReadAll(body) + require.NoError(t, err) + waitForMultipartBody(t, body) + _, err = file.Stat() + require.ErrorIs(t, err, os.ErrClosed) + require.NoError(t, body.Close()) +} + +func TestMultipartRequestBodyClosesEveryUpload(t *testing.T) { + firstCloseErr := errors.New("first close failed") + secondCloseErr := errors.New("second close failed") + first := &recordingReadCloser{reader: strings.NewReader("first"), closeErr: firstCloseErr} + second := &recordingReadCloser{reader: strings.NewReader("second"), closeErr: secondCloseErr} + third := &recordingReadCloser{reader: strings.NewReader("third")} + body := newMultipartRequestBody(map[string]any{ + "files": []any{ + fileUpload{Reader: first, filename: "first.txt"}, + map[string]any{"nested": fileUpload{Reader: second, filename: "second.txt"}}, + }, + "typed": map[string]fileUpload{ + "nested": {Reader: third, filename: "third.txt"}, + }, + }, apiform.FormatBrackets) + + _, err := io.ReadAll(body) + require.ErrorIs(t, err, firstCloseErr) + require.ErrorIs(t, err, secondCloseErr) + waitForMultipartBody(t, body) + require.Equal(t, int32(1), first.closeCount.Load()) + require.Equal(t, int32(1), second.closeCount.Load()) + require.Equal(t, int32(1), third.closeCount.Load()) + + err = body.Close() + require.ErrorIs(t, err, firstCloseErr) + require.ErrorIs(t, err, secondCloseErr) + require.Equal(t, int32(1), first.closeCount.Load()) + require.Equal(t, int32(1), second.closeCount.Load()) + require.Equal(t, int32(1), third.closeCount.Load()) +} + +func TestMultipartRequestBodyCloseCancelsEncoding(t *testing.T) { + readErr := errors.New("upload closed while reading") + upload := newCloseUnblocksReader(readErr) + body := newMultipartRequestBody(map[string]any{ + "file": fileUpload{Reader: upload, filename: "blocked.bin"}, + }, apiform.FormatBrackets) + + readDone := make(chan error, 1) + go func() { + _, err := io.Copy(io.Discard, body) + readDone <- err + }() + + select { + case <-upload.readStarted: + case <-time.After(multipartTestTimeout): + t.Fatal("multipart encoder did not start reading the upload") + } + + require.NoError(t, body.Close()) + + select { + case err := <-readDone: + require.Error(t, err) + case <-time.After(multipartTestTimeout): + t.Fatal("request body read did not stop after Close") + } + waitForMultipartBody(t, body) + require.Equal(t, int32(1), upload.closeCount.Load()) +} + +func TestMultipartRequestBodyCloseBeforeRead(t *testing.T) { + source := &countingReader{remaining: 1024} + upload := &recordingReadCloser{reader: source} + body := newMultipartRequestBody(map[string]any{ + "file": fileUpload{Reader: upload, filename: "unused.bin"}, + }, apiform.FormatBrackets) + + require.NoError(t, body.Close()) + waitForMultipartBody(t, body) + require.Zero(t, source.bytesRead.Load()) + require.Equal(t, int32(1), upload.closeCount.Load()) +} + +func TestMultipartRequestBodyDoesNotCloseUnownedReaders(t *testing.T) { + stdinLikeReader := &recordingReadCloser{reader: strings.NewReader("stdin payload")} + embedded, err := embedFiles( + map[string]any{"file": "@-"}, + EmbedIOReader, + &onceStdinReader{stdinReader: stdinLikeReader}, + ) + require.NoError(t, err) + body := newMultipartRequestBody(embedded.(map[string]any), apiform.FormatBrackets) + + data, err := io.ReadAll(body) + require.NoError(t, err) + require.Contains(t, string(data), "stdin payload") + waitForMultipartBody(t, body) + require.NoError(t, body.Close()) + require.Zero(t, stdinLikeReader.closeCount.Load()) +} + +func TestMultipartRequestOptionsDisableRetries(t *testing.T) { + type requestInfo struct { + contentLength int64 + transferEncoding []string + } + requestInfoCh := make(chan requestInfo, 4) + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + requestInfoCh <- requestInfo{ + contentLength: r.ContentLength, + transferEncoding: r.TransferEncoding, + } + _, _ = io.Copy(io.Discard, r.Body) + _ = r.Body.Close() + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusInternalServerError) + _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"server_error"}}`) + })) + t.Cleanup(server.Close) + + upload := &recordingReadCloser{reader: strings.NewReader("payload")} + body := newMultipartRequestBody(map[string]any{ + "file": fileUpload{Reader: upload, filename: "upload.txt"}, + }, apiform.FormatBrackets) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err := client.Files.New(context.Background(), openai.FileNewParams{}, multipartRequestOptions(body)...) + require.Error(t, err) + require.Equal(t, int32(1), requestCount.Load()) + info := <-requestInfoCh + require.Equal(t, int64(-1), info.contentLength) + require.Contains(t, info.transferEncoding, "chunked") + waitForMultipartBody(t, body) + require.Equal(t, int32(1), upload.closeCount.Load()) +} + +func TestEmbedFilesClosesPartialUploadsOnError(t *testing.T) { + upload := &recordingReadCloser{reader: strings.NewReader("already opened")} + missingPath := FilePathValue(t.TempDir() + "/missing.bin") + + _, err := embedFiles([]any{ + fileUpload{Reader: upload, filename: "opened.bin"}, + missingPath, + }, EmbedIOReader, &onceStdinReader{}) + require.Error(t, err) + require.Equal(t, int32(1), upload.closeCount.Load()) +} + +func waitForMultipartBody(t *testing.T, body *multipartRequestBody) { + t.Helper() + select { + case <-body.done: + case <-time.After(multipartTestTimeout): + t.Fatal("multipart encoder did not stop") + } +} + +type countingReader struct { + remaining int64 + bytesRead atomic.Int64 +} + +func (r *countingReader) Read(p []byte) (int, error) { + if r.remaining == 0 { + return 0, io.EOF + } + if int64(len(p)) > r.remaining { + p = p[:r.remaining] + } + for i := range p { + p[i] = 0 + } + n := len(p) + r.remaining -= int64(n) + r.bytesRead.Add(int64(n)) + return n, nil +} + +type errorReader struct { + err error +} + +func (r errorReader) Read([]byte) (int, error) { + return 0, r.err +} + +type recordingReadCloser struct { + reader io.Reader + closeErr error + closeCount atomic.Int32 +} + +func (r *recordingReadCloser) Read(p []byte) (int, error) { + return r.reader.Read(p) +} + +func (r *recordingReadCloser) Close() error { + r.closeCount.Add(1) + return r.closeErr +} + +type closeUnblocksReader struct { + err error + readStarted chan struct{} + closed chan struct{} + startOnce sync.Once + closeOnce sync.Once + closeCount atomic.Int32 +} + +func newCloseUnblocksReader(err error) *closeUnblocksReader { + return &closeUnblocksReader{ + err: err, + readStarted: make(chan struct{}), + closed: make(chan struct{}), + } +} + +func (r *closeUnblocksReader) Read([]byte) (int, error) { + r.startOnce.Do(func() { close(r.readStarted) }) + <-r.closed + return 0, r.err +} + +func (r *closeUnblocksReader) Close() error { + r.closeCount.Add(1) + r.closeOnce.Do(func() { close(r.closed) }) + return nil +} From e48d568edf35128cdfd804e3b15ffa5ae0214dbf Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Mon, 17 Aug 2026 17:51:56 +0000 Subject: [PATCH 2/9] fix: preserve multipart request semantics --- pkg/cmd/flagoptions.go | 46 ++++++- pkg/cmd/multipartbody.go | 252 ++++++++++++++++++++++++++++++++-- pkg/cmd/multipartbody_test.go | 203 ++++++++++++++++++++++++++- 3 files changed, 486 insertions(+), 15 deletions(-) diff --git a/pkg/cmd/flagoptions.go b/pkg/cmd/flagoptions.go index c3b30d92..2f17534f 100644 --- a/pkg/cmd/flagoptions.go +++ b/pkg/cmd/flagoptions.go @@ -508,9 +508,12 @@ func flagOptions( return nil, fmt.Errorf("Cannot send a non-map value to a form-encoded endpoint: %v\n", requestContents.Body) } encodingFormat := apiform.FormatBrackets - body := newMultipartRequestBody(bodyMap, encodingFormat) + multipartOptions, err := multipartRequestOptions(bodyMap, encodingFormat) + if err != nil { + return nil, err + } pendingMultipartUploads = nil // multipartRequestBody owns them now. - options = append(options, multipartRequestOptions(body)...) + options = append(options, multipartOptions...) case ApplicationJSON: bodyBytes, err := json.Marshal(requestContents.Body) @@ -554,6 +557,9 @@ type fileUpload struct { io.Reader // apiform checks for reader and reads its contents during encode filename string contentType string + sourcePath string + size int64 + knownSize bool } func (f fileUpload) Filename() string { return f.filename } @@ -565,6 +571,32 @@ func (f fileUpload) Close() error { return nil } +func (f fileUpload) canReplay() bool { + return f.sourcePath != "" && f.knownSize +} + +func (f fileUpload) reopen() (fileUpload, error) { + if !f.canReplay() { + return fileUpload{}, errors.New("multipart upload is not replayable") + } + file, err := os.Open(f.sourcePath) + if err != nil { + return fileUpload{}, err + } + info, err := file.Stat() + if err != nil { + return fileUpload{}, errors.Join(err, file.Close()) + } + if !info.Mode().IsRegular() || info.Size() != f.size { + return fileUpload{}, errors.Join( + fmt.Errorf("upload %q changed while following a redirect", f.sourcePath), + file.Close(), + ) + } + f.Reader = file + return f, nil +} + // openFileUpload opens the file at path and returns a fileUpload whose filename // is the path's basename and whose content type is derived from the file // extension (falling back to application/octet-stream when unknown). @@ -573,6 +605,13 @@ func openFileUpload(path string) (fileUpload, error) { if err != nil { return fileUpload{}, err } + info, err := file.Stat() + if err != nil { + return fileUpload{}, errors.Join(err, file.Close()) + } + if info.IsDir() { + return fileUpload{}, errors.Join(fmt.Errorf("read %s: is a directory", path), file.Close()) + } contentType := mime.TypeByExtension(filepath.Ext(path)) if contentType == "" { contentType = "application/octet-stream" @@ -581,6 +620,9 @@ func openFileUpload(path string) (fileUpload, error) { Reader: file, filename: filepath.Base(path), contentType: contentType, + sourcePath: path, + size: info.Size(), + knownSize: info.Mode().IsRegular(), }, nil } diff --git a/pkg/cmd/multipartbody.go b/pkg/cmd/multipartbody.go index f33562d5..ac4817e9 100644 --- a/pkg/cmd/multipartbody.go +++ b/pkg/cmd/multipartbody.go @@ -1,10 +1,15 @@ package cmd import ( + "bytes" "errors" + "fmt" "io" + "math" "mime/multipart" + "net/http" "reflect" + "strings" "sync" "github.com/openai/openai-cli/internal/apiform" @@ -34,6 +39,31 @@ var _ io.ReadCloser = (*multipartRequestBody)(nil) func newMultipartRequestBody(bodyMap map[string]any, encodingFormat apiform.FormFormat) *multipartRequestBody { reader, writer := io.Pipe() multipartWriter := multipart.NewWriter(writer) + return makeMultipartRequestBody(reader, writer, multipartWriter, bodyMap, encodingFormat) +} + +func newMultipartRequestBodyWithBoundary( + bodyMap map[string]any, + encodingFormat apiform.FormFormat, + boundary string, +) (*multipartRequestBody, error) { + reader, writer := io.Pipe() + multipartWriter := multipart.NewWriter(writer) + if err := multipartWriter.SetBoundary(boundary); err != nil { + _ = reader.Close() + _ = writer.Close() + return nil, err + } + return makeMultipartRequestBody(reader, writer, multipartWriter, bodyMap, encodingFormat), nil +} + +func makeMultipartRequestBody( + reader *io.PipeReader, + writer *io.PipeWriter, + multipartWriter *multipart.Writer, + bodyMap map[string]any, + encodingFormat apiform.FormFormat, +) *multipartRequestBody { return &multipartRequestBody{ reader: reader, writer: writer, @@ -67,6 +97,10 @@ func (b *multipartRequestBody) ContentType() string { return b.contentType } +func (b *multipartRequestBody) Boundary() string { + return b.multipartWriter.Boundary() +} + func (b *multipartRequestBody) start() { b.startOnce.Do(func() { go b.encode() @@ -76,11 +110,18 @@ func (b *multipartRequestBody) start() { func (b *multipartRequestBody) encode() { defer close(b.done) - encodeErr := apiform.MarshalWithSettings(b.bodyMap, b.multipartWriter, b.encodingFormat) + if err := apiform.MarshalWithSettings(b.bodyMap, b.multipartWriter, b.encodingFormat); err != nil { + // A final multipart boundary asserts that every part completed. Abort the + // pipe immediately on a source error instead of making a truncated file + // look like a valid zero-byte upload to the receiver. + _ = b.writer.CloseWithError(err) + _ = b.cleanup() + return + } + writerCloseErr := b.multipartWriter.Close() cleanupErr := b.cleanup() - - _ = b.writer.CloseWithError(errors.Join(encodeErr, writerCloseErr, cleanupErr)) + _ = b.writer.CloseWithError(errors.Join(writerCloseErr, cleanupErr)) } func (b *multipartRequestBody) cleanup() error { @@ -90,13 +131,208 @@ func (b *multipartRequestBody) cleanup() error { return b.cleanupErr } -func multipartRequestOptions(body *multipartRequestBody) []option.RequestOption { - return []option.RequestOption{ +// multipartRequestOptions preserves the SDK's buffered, replayable behavior +// when a multipart form contains only scalar fields. Uploads are streamed. A +// form containing only regular files is still replayable for HTTP redirects and +// carries an exact Content-Length; forms containing stdin, pipes, or arbitrary +// readers use chunked transfer and fail explicitly if a 307/308 requires replay. +func multipartRequestOptions( + bodyMap map[string]any, + encodingFormat apiform.FormFormat, +) ([]option.RequestOption, error) { + info := inspectMultipartBody(bodyMap) + if !info.hasUpload { + return bufferedMultipartRequestOptions(bodyMap, encodingFormat) + } + + boundary := multipart.NewWriter(io.Discard).Boundary() + var contentLength int64 + if info.replayable { + var err error + contentLength, err = multipartContentLength(bodyMap, encodingFormat, boundary) + if err != nil { + return nil, err + } + } + + body, err := newMultipartRequestBodyWithBoundary(bodyMap, encodingFormat, boundary) + if err != nil { + return nil, err + } + options := []option.RequestOption{ option.WithRequestBody(body.ContentType(), body), - // A streamed multipart body cannot be replayed. Make the behavior - // explicit even though the SDK also avoids retrying bodies without - // an http.Request.GetBody function. + } + + if info.replayable { + newBody := func() (io.ReadCloser, error) { + reopened, err := reopenMultipartBody(bodyMap) + if err != nil { + return nil, err + } + return newMultipartRequestBodyWithBoundary(reopened, encodingFormat, boundary) + } + options = append(options, option.WithMiddleware(func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { + req.ContentLength = contentLength + req.GetBody = newBody + return next(req) + })) + return options, nil + } + + options = append(options, option.WithMaxRetries(0), + option.WithMiddleware(rejectUnreplayableRedirect), + ) + return options, nil +} + +func bufferedMultipartRequestOptions( + bodyMap map[string]any, + encodingFormat apiform.FormFormat, +) ([]option.RequestOption, error) { + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + if err := apiform.MarshalWithSettings(bodyMap, writer, encodingFormat); err != nil { + return nil, err + } + if err := writer.Close(); err != nil { + return nil, err + } + return []option.RequestOption{option.WithRequestBody(writer.FormDataContentType(), buf.Bytes())}, nil +} + +func rejectUnreplayableRedirect(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { + res, err := next(req) + if err != nil || res == nil || (res.StatusCode != http.StatusTemporaryRedirect && res.StatusCode != http.StatusPermanentRedirect) { + return res, err + } + + location := res.Header.Get("Location") + if res.Body != nil { + _ = res.Body.Close() + } + return res, fmt.Errorf( + "cannot follow HTTP %d redirect to %q: multipart upload contains a non-replayable source", + res.StatusCode, + location, + ) +} + +type multipartBodyInfo struct { + hasUpload bool + replayable bool +} + +func inspectMultipartBody(value any) multipartBodyInfo { + switch value := value.(type) { + case map[string]any: + info := multipartBodyInfo{replayable: true} + for _, child := range value { + childInfo := inspectMultipartBody(child) + info.hasUpload = info.hasUpload || childInfo.hasUpload + info.replayable = info.replayable && childInfo.replayable + } + return info + case []any: + info := multipartBodyInfo{replayable: true} + for _, child := range value { + childInfo := inspectMultipartBody(child) + info.hasUpload = info.hasUpload || childInfo.hasUpload + info.replayable = info.replayable && childInfo.replayable + } + return info + case fileUpload: + return multipartBodyInfo{hasUpload: true, replayable: value.canReplay()} + default: + _, isReader := value.(io.Reader) + return multipartBodyInfo{hasUpload: isReader, replayable: !isReader} + } +} + +func multipartContentLength( + bodyMap map[string]any, + encodingFormat apiform.FormFormat, + boundary string, +) (int64, error) { + var uploadBytes int64 + emptyValue, err := transformFileUploads(bodyMap, func(upload fileUpload) (fileUpload, error) { + if !upload.canReplay() { + return fileUpload{}, errors.New("cannot determine multipart length for a non-replayable upload") + } + if upload.size > math.MaxInt64-uploadBytes { + return fileUpload{}, errors.New("multipart content length overflows int64") + } + uploadBytes += upload.size + upload.Reader = strings.NewReader("") + return upload, nil + }) + if err != nil { + return 0, err + } + + var framing bytes.Buffer + writer := multipart.NewWriter(&framing) + if err := writer.SetBoundary(boundary); err != nil { + return 0, err + } + if err := apiform.MarshalWithSettings(emptyValue, writer, encodingFormat); err != nil { + return 0, err + } + if err := writer.Close(); err != nil { + return 0, err + } + if int64(framing.Len()) > math.MaxInt64-uploadBytes { + return 0, errors.New("multipart content length overflows int64") + } + return int64(framing.Len()) + uploadBytes, nil +} + +func reopenMultipartBody(bodyMap map[string]any) (map[string]any, error) { + value, err := transformFileUploads(bodyMap, func(upload fileUpload) (fileUpload, error) { + return upload.reopen() + }) + if err != nil { + return nil, err + } + reopened, ok := value.(map[string]any) + if !ok { + return nil, errors.New("multipart body must be a map") + } + return reopened, nil +} + +func transformFileUploads( + value any, + transform func(fileUpload) (fileUpload, error), +) (any, error) { + switch value := value.(type) { + case fileUpload: + return transform(value) + case map[string]any: + result := make(map[string]any, len(value)) + for key, child := range value { + transformed, err := transformFileUploads(child, transform) + if err != nil { + return nil, errors.Join(err, closeFileUploads(result)) + } + result[key] = transformed + } + return result, nil + case []any: + result := make([]any, 0, len(value)) + for _, child := range value { + transformed, err := transformFileUploads(child, transform) + if err != nil { + return nil, errors.Join(err, closeFileUploads(result)) + } + result = append(result, transformed) + } + return result, nil + default: + if _, isReader := value.(io.Reader); isReader { + return nil, errors.New("multipart body contains a non-replayable reader") + } + return value, nil } } diff --git a/pkg/cmd/multipartbody_test.go b/pkg/cmd/multipartbody_test.go index 3b92be40..323e5798 100644 --- a/pkg/cmd/multipartbody_test.go +++ b/pkg/cmd/multipartbody_test.go @@ -3,6 +3,7 @@ package cmd import ( "context" "errors" + "fmt" "io" "mime" "mime/multipart" @@ -19,6 +20,7 @@ import ( "github.com/openai/openai-cli/internal/apiform" "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/option" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -69,13 +71,44 @@ func TestMultipartRequestBodyPropagatesProducerErrors(t *testing.T) { "file": fileUpload{Reader: upload, filename: "broken.bin"}, }, apiform.FormatBrackets) - _, err := io.ReadAll(body) + data, err := io.ReadAll(body) require.ErrorIs(t, err, readErr) + require.NotContains(t, string(data), "--"+body.Boundary()+"--", + "a failed file part must not be followed by a successful terminal boundary") waitForMultipartBody(t, body) require.Equal(t, int32(1), upload.closeCount.Load()) require.NoError(t, body.Close()) } +func TestMultipartRequestBodySourceErrorCannotBeParsedAsCompleteUpload(t *testing.T) { + readErr := errors.New("upload read failed") + parseErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + parseErrCh <- r.ParseMultipartForm(1 << 20) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": fileUpload{Reader: errorReader{err: readErr}, filename: "broken.bin"}, + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.ErrorIs(t, err, readErr) + select { + case parseErr := <-parseErrCh: + require.Error(t, parseErr, "the receiver must reject a file part without a terminal boundary") + case <-time.After(multipartTestTimeout): + t.Fatal("server did not receive the failed streaming request") + } +} + func TestMultipartRequestBodyClosesOpenedFiles(t *testing.T) { path := filepath.Join(t.TempDir(), "upload.bin") require.NoError(t, os.WriteFile(path, []byte("payload"), 0o600)) @@ -186,7 +219,162 @@ func TestMultipartRequestBodyDoesNotCloseUnownedReaders(t *testing.T) { require.Zero(t, stdinLikeReader.closeCount.Load()) } -func TestMultipartRequestOptionsDisableRetries(t *testing.T) { +func TestMultipartRequestOptionsFollow307And308ForRegularFiles(t *testing.T) { + for _, status := range []int{http.StatusTemporaryRedirect, http.StatusPermanentRedirect} { + t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { + path := filepath.Join(t.TempDir(), "redirect-upload.txt") + require.NoError(t, os.WriteFile(path, []byte("redirect payload"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + originalFile := upload.Reader.(*os.File) + + var requestCount atomic.Int32 + var uploaded string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + switch r.URL.Path { + case "/files": + http.Redirect(w, r, "/upload", status) + case "/upload": + assert.Positive(t, r.ContentLength) + assert.Empty(t, r.TransferEncoding) + file, _, parseErr := r.FormFile("file") + if !assert.NoError(t, parseErr) { + http.Error(w, parseErr.Error(), http.StatusBadRequest) + return + } + contents, readErr := io.ReadAll(file) + assert.NoError(t, readErr) + assert.NoError(t, file.Close()) + uploaded = string(contents) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + var response []byte + options = append(options, option.WithResponseBodyInto(&response)) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.NoError(t, err) + require.Equal(t, int32(2), requestCount.Load()) + require.Equal(t, "redirect payload", uploaded) + _, err = originalFile.Stat() + require.ErrorIs(t, err, os.ErrClosed) + }) + } +} + +func TestMultipartRequestOptionsReject307And308ForUnknownSources(t *testing.T) { + for _, status := range []int{http.StatusTemporaryRedirect, http.StatusPermanentRedirect} { + t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + http.Redirect(w, r, "/upload", status) + })) + t.Cleanup(server.Close) + + upload := &recordingReadCloser{reader: strings.NewReader("stdin payload")} + options, err := multipartRequestOptions(map[string]any{ + "file": fileUpload{Reader: upload, filename: "stdin"}, + }, apiform.FormatBrackets) + require.NoError(t, err) + var response []byte + options = append(options, option.WithResponseBodyInto(&response)) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.ErrorContains(t, err, "non-replayable source") + require.Equal(t, int32(1), requestCount.Load()) + require.Equal(t, int32(1), upload.closeCount.Load()) + }) + } +} + +func TestMultipartRequestOptionsRetryScalarOnlyForms(t *testing.T) { + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + if !assert.NoError(t, r.ParseMultipartForm(1<<20)) { + http.Error(w, "invalid multipart form", http.StatusBadRequest) + return + } + assert.Equal(t, "hello", r.FormValue("prompt")) + w.Header().Set("Content-Type", "application/json") + if requestCount.Load() == 1 { + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"rate_limit_error"}}`) + return + } + _, _ = io.WriteString(w, `{"id":"video_123","object":"video","status":"queued"}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{"prompt": "hello"}, apiform.FormatBrackets) + require.NoError(t, err) + var response []byte + options = append(options, option.WithResponseBodyInto(&response)) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Videos.New(context.Background(), openai.VideoNewParams{}, options...) + require.NoError(t, err) + require.Equal(t, int32(2), requestCount.Load()) +} + +func TestMultipartRequestOptionsSetContentLengthForRegularFiles(t *testing.T) { + path := filepath.Join(t.TempDir(), "upload.txt") + require.NoError(t, os.WriteFile(path, []byte("known-size payload"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Positive(t, r.ContentLength) + assert.Empty(t, r.TransferEncoding) + contents, readErr := io.ReadAll(r.Body) + assert.NoError(t, readErr) + assert.Equal(t, r.ContentLength, int64(len(contents))) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + var response []byte + options = append(options, option.WithResponseBodyInto(&response)) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.NoError(t, err) +} + +func TestMultipartRequestOptionsDisableRetriesForUnknownSources(t *testing.T) { type requestInfo struct { contentLength int64 transferEncoding []string @@ -209,24 +397,29 @@ func TestMultipartRequestOptionsDisableRetries(t *testing.T) { t.Cleanup(server.Close) upload := &recordingReadCloser{reader: strings.NewReader("payload")} - body := newMultipartRequestBody(map[string]any{ + options, err := multipartRequestOptions(map[string]any{ "file": fileUpload{Reader: upload, filename: "upload.txt"}, }, apiform.FormatBrackets) + require.NoError(t, err) client := openai.NewClient( option.WithAPIKey("test-key"), option.WithBaseURL(server.URL+"/"), ) - _, err := client.Files.New(context.Background(), openai.FileNewParams{}, multipartRequestOptions(body)...) + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) require.Error(t, err) require.Equal(t, int32(1), requestCount.Load()) info := <-requestInfoCh require.Equal(t, int64(-1), info.contentLength) require.Contains(t, info.transferEncoding, "chunked") - waitForMultipartBody(t, body) require.Equal(t, int32(1), upload.closeCount.Load()) } +func TestOpenFileUploadRejectsDirectoriesBeforeRequest(t *testing.T) { + _, err := openFileUpload(t.TempDir()) + require.ErrorContains(t, err, "is a directory") +} + func TestEmbedFilesClosesPartialUploadsOnError(t *testing.T) { upload := &recordingReadCloser{reader: strings.NewReader("already opened")} missingPath := FilePathValue(t.TempDir() + "/missing.bin") From 521f01ca48148e3cbff968d26faa71cb472ea804 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Mon, 17 Aug 2026 18:51:41 +0000 Subject: [PATCH 3/9] Preserve multipart replay semantics --- pkg/cmd/flagoptions.go | 89 +++++++++--- pkg/cmd/multipartbody.go | 195 ++++++++++++++++++++++---- pkg/cmd/multipartbody_test.go | 252 +++++++++++++++++++++++++++++++++- 3 files changed, 495 insertions(+), 41 deletions(-) diff --git a/pkg/cmd/flagoptions.go b/pkg/cmd/flagoptions.go index 2f17534f..a90a2d37 100644 --- a/pkg/cmd/flagoptions.go +++ b/pkg/cmd/flagoptions.go @@ -13,6 +13,7 @@ import ( "path/filepath" "reflect" "strings" + "sync" "unicode/utf8" "github.com/openai/openai-cli/internal/apiform" @@ -557,14 +558,18 @@ type fileUpload struct { io.Reader // apiform checks for reader and reads its contents during encode filename string contentType string - sourcePath string size int64 knownSize bool + source *replayableFileSource + ownsSource bool } func (f fileUpload) Filename() string { return f.filename } func (f fileUpload) ContentType() string { return f.contentType } func (f fileUpload) Close() error { + if f.ownsSource && f.source != nil { + return f.source.Close() + } if c, ok := f.Reader.(io.Closer); ok { return c.Close() } @@ -572,29 +577,68 @@ func (f fileUpload) Close() error { } func (f fileUpload) canReplay() bool { - return f.sourcePath != "" && f.knownSize + return f.source != nil && f.knownSize } -func (f fileUpload) reopen() (fileUpload, error) { +func (f fileUpload) replay() (fileUpload, error) { if !f.canReplay() { return fileUpload{}, errors.New("multipart upload is not replayable") } - file, err := os.Open(f.sourcePath) + reader, err := f.source.Reader() if err != nil { return fileUpload{}, err } - info, err := file.Stat() + f.Reader = reader + f.ownsSource = false + return f, nil +} + +type replayableFileSource struct { + file *os.File + size int64 + + closeOnce sync.Once + closeErr error +} + +func (s *replayableFileSource) Reader() (io.Reader, error) { + info, err := s.file.Stat() if err != nil { - return fileUpload{}, errors.Join(err, file.Close()) + return nil, err } - if !info.Mode().IsRegular() || info.Size() != f.size { - return fileUpload{}, errors.Join( - fmt.Errorf("upload %q changed while following a redirect", f.sourcePath), - file.Close(), - ) + if !info.Mode().IsRegular() || info.Size() != s.size { + return nil, fmt.Errorf("upload %q changed while being replayed", s.file.Name()) } - f.Reader = file - return f, nil + return io.NewSectionReader(s.file, 0, s.size), nil +} + +func (s *replayableFileSource) Close() error { + s.closeOnce.Do(func() { + s.closeErr = s.file.Close() + }) + return s.closeErr +} + +func hasTrustworthyFileSize(file *os.File, info os.FileInfo) bool { + if !info.Mode().IsRegular() || info.Size() < 0 { + return false + } + + // Regular-looking virtual files can report a size unrelated to their + // readable contents. Verify both sides of the reported EOF without moving + // the descriptor offset; ordinary and sparse files support these ReadAt + // probes, while procfs/sysfs-style sources fail or expose extra bytes. + var probe [1]byte + if info.Size() == 0 { + n, err := file.ReadAt(probe[:], 0) + return n == 0 && errors.Is(err, io.EOF) + } + n, err := file.ReadAt(probe[:], info.Size()-1) + if n != 1 || err != nil { + return false + } + n, err = file.ReadAt(probe[:], info.Size()) + return n == 0 && errors.Is(err, io.EOF) } // openFileUpload opens the file at path and returns a fileUpload whose filename @@ -616,13 +660,26 @@ func openFileUpload(path string) (fileUpload, error) { if contentType == "" { contentType = "application/octet-stream" } + if !hasTrustworthyFileSize(file, info) { + return fileUpload{ + Reader: file, + filename: filepath.Base(path), + contentType: contentType, + }, nil + } + source := &replayableFileSource{file: file, size: info.Size()} + reader, err := source.Reader() + if err != nil { + return fileUpload{}, errors.Join(err, source.Close()) + } return fileUpload{ - Reader: file, + Reader: reader, filename: filepath.Base(path), contentType: contentType, - sourcePath: path, size: info.Size(), - knownSize: info.Mode().IsRegular(), + knownSize: true, + source: source, + ownsSource: true, }, nil } diff --git a/pkg/cmd/multipartbody.go b/pkg/cmd/multipartbody.go index ac4817e9..bead5b7f 100644 --- a/pkg/cmd/multipartbody.go +++ b/pkg/cmd/multipartbody.go @@ -155,35 +155,183 @@ func multipartRequestOptions( } } + if info.replayable { + return replayableMultipartRequestOptions(bodyMap, encodingFormat, boundary, contentLength) + } + body, err := newMultipartRequestBodyWithBoundary(bodyMap, encodingFormat, boundary) if err != nil { return nil, err } - options := []option.RequestOption{ + return []option.RequestOption{ option.WithRequestBody(body.ContentType(), body), + option.WithMaxRetries(0), + option.WithMiddleware(rejectUnreplayableRedirect), + }, nil +} + +// openai-go v3.51 only installs GetBody for its concrete buffered body types. +// A streamed body therefore needs a replay marker on the SDK-owned request so +// its native retry loop runs. The middleware replaces that empty marker with a +// real multipart stream on every SDK attempt and supplies the same factory to +// net/http for 307/308 redirects. The CLI does not expose retry configuration, +// so the request fixes the SDK's current default of two retries; that bound also +// lets the middleware deterministically close its stable file descriptors. +const replayableMultipartMaxRetries = 2 + +type replayableMultipartRequest struct { + bodyMap map[string]any + encodingFormat apiform.FormFormat + boundary string + contentLength int64 + + mu sync.Mutex + attempts int + closed bool + closeErr error +} + +func replayableMultipartRequestOptions( + bodyMap map[string]any, + encodingFormat apiform.FormFormat, + boundary string, + contentLength int64, +) ([]option.RequestOption, error) { + writer := multipart.NewWriter(io.Discard) + if err := writer.SetBoundary(boundary); err != nil { + return nil, err + } + request := &replayableMultipartRequest{ + bodyMap: bodyMap, + encodingFormat: encodingFormat, + boundary: boundary, + contentLength: contentLength, } - if info.replayable { - newBody := func() (io.ReadCloser, error) { - reopened, err := reopenMultipartBody(bodyMap) - if err != nil { - return nil, err - } - return newMultipartRequestBodyWithBoundary(reopened, encodingFormat, boundary) - } - options = append(options, option.WithMiddleware(func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { - req.ContentLength = contentLength - req.GetBody = newBody - return next(req) - })) - return options, nil + // The empty *bytes.Reader is deliberate: it makes openai-go install GetBody + // on its canonical request before retry decisions. middleware replaces the + // marker on every attempt, so no marker bytes reach the network. + return []option.RequestOption{ + option.WithRequestBody(writer.FormDataContentType(), bytes.NewReader(nil)), + option.WithMaxRetries(replayableMultipartMaxRetries), + option.WithMiddleware(request.middleware), + }, nil +} + +func (r *replayableMultipartRequest) middleware( + req *http.Request, + next option.MiddlewareNext, +) (*http.Response, error) { + attemptIndex := r.beginAttempt() + attempt := &multipartReplayAttempt{request: r} + + // Discard the SDK's empty replay marker before installing the actual body. + if req.Body != nil { + _ = req.Body.Close() } + body, err := attempt.newBody() + if err != nil { + return nil, errors.Join(err, r.close()) + } + req.Body = body + req.ContentLength = r.contentLength + req.GetBody = attempt.newBody + + res, requestErr := next(req) + attemptErr := attempt.close() + if req.Context().Err() != nil || + !shouldRetryReplayableMultipart(res) || + attemptIndex >= replayableMultipartMaxRetries { + return res, errors.Join(requestErr, attemptErr, r.close()) + } + return res, errors.Join(requestErr, attemptErr) +} - options = append(options, - option.WithMaxRetries(0), - option.WithMiddleware(rejectUnreplayableRedirect), - ) - return options, nil +func (r *replayableMultipartRequest) beginAttempt() int { + r.mu.Lock() + defer r.mu.Unlock() + attempt := r.attempts + r.attempts++ + return attempt +} + +func (r *replayableMultipartRequest) newBody() (*multipartRequestBody, error) { + bodyMap, err := replayMultipartBody(r.bodyMap) + if err != nil { + return nil, err + } + return newMultipartRequestBodyWithBoundary(bodyMap, r.encodingFormat, r.boundary) +} + +func (r *replayableMultipartRequest) close() error { + r.mu.Lock() + defer r.mu.Unlock() + if !r.closed { + r.closed = true + r.closeErr = closeFileUploads(r.bodyMap) + } + return r.closeErr +} + +type multipartReplayAttempt struct { + request *replayableMultipartRequest + + mu sync.Mutex + bodies []*multipartRequestBody + closed bool +} + +func (a *multipartReplayAttempt) newBody() (io.ReadCloser, error) { + body, err := a.request.newBody() + if err != nil { + return nil, err + } + + a.mu.Lock() + defer a.mu.Unlock() + if a.closed { + _ = body.Close() + return nil, errors.New("multipart replay attempt is closed") + } + a.bodies = append(a.bodies, body) + return body, nil +} + +func (a *multipartReplayAttempt) close() error { + a.mu.Lock() + if a.closed { + a.mu.Unlock() + return nil + } + a.closed = true + bodies := append([]*multipartRequestBody(nil), a.bodies...) + a.mu.Unlock() + + var errs []error + for _, body := range bodies { + errs = append(errs, body.Close()) + <-body.done + } + return errors.Join(errs...) +} + +// Keep this predicate aligned with openai-go's retry policy. It is only used +// to decide when the request-level file descriptors can be closed; the SDK +// still owns retry timing, response cleanup, and the retry decision itself. +func shouldRetryReplayableMultipart(res *http.Response) bool { + if res == nil { + return true + } + if res.Header.Get("x-should-retry") == "true" { + return true + } + if res.Header.Get("x-should-retry") == "false" { + return false + } + return res.StatusCode == http.StatusRequestTimeout || + res.StatusCode == http.StatusConflict || + res.StatusCode == http.StatusTooManyRequests || + res.StatusCode >= http.StatusInternalServerError } func bufferedMultipartRequestOptions( @@ -264,6 +412,7 @@ func multipartContentLength( } uploadBytes += upload.size upload.Reader = strings.NewReader("") + upload.ownsSource = false return upload, nil }) if err != nil { @@ -287,9 +436,9 @@ func multipartContentLength( return int64(framing.Len()) + uploadBytes, nil } -func reopenMultipartBody(bodyMap map[string]any) (map[string]any, error) { +func replayMultipartBody(bodyMap map[string]any) (map[string]any, error) { value, err := transformFileUploads(bodyMap, func(upload fileUpload) (fileUpload, error) { - return upload.reopen() + return upload.replay() }) if err != nil { return nil, err @@ -336,7 +485,7 @@ func transformFileUploads( } } -// closeFileUploads closes every file opened by openFileUpload, including files +// closeFileUploads closes every source owned by a fileUpload, including values // nested inside maps and arrays. Other io.ReadClosers (notably stdin wrappers) // are intentionally left alone because this code does not own them. func closeFileUploads(value any) error { diff --git a/pkg/cmd/multipartbody_test.go b/pkg/cmd/multipartbody_test.go index 323e5798..ed673a4c 100644 --- a/pkg/cmd/multipartbody_test.go +++ b/pkg/cmd/multipartbody_test.go @@ -11,6 +11,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "runtime" "strings" "sync" "sync/atomic" @@ -114,7 +115,7 @@ func TestMultipartRequestBodyClosesOpenedFiles(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte("payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - file := upload.Reader.(*os.File) + file := upload.source.file body := newMultipartRequestBody(map[string]any{"file": upload}, apiform.FormatBrackets) _, err = io.ReadAll(body) @@ -226,7 +227,7 @@ func TestMultipartRequestOptionsFollow307And308ForRegularFiles(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte("redirect payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - originalFile := upload.Reader.(*os.File) + originalFile := upload.source.file var requestCount atomic.Int32 var uploaded string @@ -341,6 +342,177 @@ func TestMultipartRequestOptionsRetryScalarOnlyForms(t *testing.T) { require.Equal(t, int32(2), requestCount.Load()) } +func TestMultipartRequestOptionsRetryRegularFiles(t *testing.T) { + path := filepath.Join(t.TempDir(), "retry-upload.txt") + require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + sourceFile := upload.source.file + + var requestCount atomic.Int32 + var uploadsMu sync.Mutex + var uploads []string + var contentLengths []int64 + var retryCounts []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + file, _, parseErr := r.FormFile("file") + if !assert.NoError(t, parseErr) { + http.Error(w, parseErr.Error(), http.StatusBadRequest) + return + } + contents, readErr := io.ReadAll(file) + assert.NoError(t, readErr) + assert.NoError(t, file.Close()) + uploadsMu.Lock() + uploads = append(uploads, string(contents)) + contentLengths = append(contentLengths, r.ContentLength) + retryCounts = append(retryCounts, r.Header.Get("X-Stainless-Retry-Count")) + uploadsMu.Unlock() + + w.Header().Set("Content-Type", "application/json") + if requestCount.Load() == 1 { + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"rate_limit_error"}}`) + return + } + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.NoError(t, err) + require.Equal(t, int32(2), requestCount.Load()) + uploadsMu.Lock() + require.Equal(t, []string{"retry payload", "retry payload"}, uploads) + require.Len(t, contentLengths, 2) + require.Positive(t, contentLengths[0]) + require.Equal(t, contentLengths[0], contentLengths[1]) + require.Equal(t, []string{"0", "1"}, retryCounts) + uploadsMu.Unlock() + _, err = sourceFile.Stat() + require.ErrorIs(t, err, os.ErrClosed) +} + +func TestMultipartRequestOptionsCloseRegularFilesAfterFinalRetry(t *testing.T) { + path := filepath.Join(t.TempDir(), "failed-retry-upload.txt") + require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + sourceFile := upload.source.file + + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + file, _, parseErr := r.FormFile("file") + if !assert.NoError(t, parseErr) { + http.Error(w, parseErr.Error(), http.StatusBadRequest) + return + } + contents, readErr := io.ReadAll(file) + assert.NoError(t, readErr) + assert.Equal(t, "retry payload", string(contents)) + assert.NoError(t, file.Close()) + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusInternalServerError) + _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"server_error"}}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.Error(t, err) + require.Equal(t, int32(3), requestCount.Load()) + _, err = sourceFile.Stat() + require.ErrorIs(t, err, os.ErrClosed) +} + +func TestMultipartRequestOptionsRedirectReplayUsesOpenedFile(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Windows does not permit atomically replacing this open test file") + } + + dir := t.TempDir() + path := filepath.Join(dir, "redirect-upload.txt") + replacement := filepath.Join(dir, "replacement.txt") + require.NoError(t, os.WriteFile(path, []byte("benign-data"), 0o600)) + require.NoError(t, os.WriteFile(replacement, []byte("secret-data"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + + var requestCount atomic.Int32 + var uploadsMu sync.Mutex + var uploads []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + file, _, parseErr := r.FormFile("file") + if !assert.NoError(t, parseErr) { + http.Error(w, parseErr.Error(), http.StatusBadRequest) + return + } + contents, readErr := io.ReadAll(file) + assert.NoError(t, readErr) + assert.NoError(t, file.Close()) + uploadsMu.Lock() + uploads = append(uploads, string(contents)) + uploadsMu.Unlock() + + if r.URL.Path == "/files" { + if !assert.NoError(t, os.Rename(replacement, path)) { + http.Error(w, "could not replace upload path", http.StatusInternalServerError) + return + } + http.Redirect(w, r, "/upload", http.StatusTemporaryRedirect) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.NoError(t, err) + require.Equal(t, int32(2), requestCount.Load()) + uploadsMu.Lock() + require.Equal(t, []string{"benign-data", "benign-data"}, uploads) + uploadsMu.Unlock() + require.FileExists(t, path) + replacedContents, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, "secret-data", string(replacedContents)) +} + func TestMultipartRequestOptionsSetContentLengthForRegularFiles(t *testing.T) { path := filepath.Join(t.TempDir(), "upload.txt") require.NoError(t, os.WriteFile(path, []byte("known-size payload"), 0o600)) @@ -420,6 +592,82 @@ func TestOpenFileUploadRejectsDirectoriesBeforeRequest(t *testing.T) { require.ErrorContains(t, err, "is a directory") } +func TestOpenFileUploadStreamsProcfsWithUnknownLength(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("procfs reproduction is Linux-specific") + } + const path = "/proc/version" + upload, err := openFileUpload(path) + if errors.Is(err, os.ErrNotExist) || errors.Is(err, os.ErrPermission) { + t.Skipf("%s is unavailable: %v", path, err) + } + require.NoError(t, err) + require.False(t, upload.knownSize) + require.False(t, upload.canReplay()) + sourceFile := upload.Reader.(*os.File) + + type requestInfo struct { + contentLength int64 + transferEncoding []string + contents string + } + infoCh := make(chan requestInfo, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + file, _, parseErr := r.FormFile("file") + if parseErr != nil { + infoCh <- requestInfo{contentLength: r.ContentLength, transferEncoding: r.TransferEncoding} + http.Error(w, parseErr.Error(), http.StatusBadRequest) + return + } + contents, readErr := io.ReadAll(file) + _ = file.Close() + if readErr != nil { + infoCh <- requestInfo{contentLength: r.ContentLength, transferEncoding: r.TransferEncoding} + http.Error(w, readErr.Error(), http.StatusBadRequest) + return + } + infoCh <- requestInfo{ + contentLength: r.ContentLength, + transferEncoding: append([]string(nil), r.TransferEncoding...), + contents: string(contents), + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.NoError(t, err) + info := <-infoCh + require.Equal(t, int64(-1), info.contentLength) + require.Contains(t, info.transferEncoding, "chunked") + require.NotEmpty(t, info.contents) + _, err = sourceFile.Stat() + require.ErrorIs(t, err, os.ErrClosed) +} + +func TestOpenFileUploadTreatsEmptyRegularFileAsKnownLength(t *testing.T) { + path := filepath.Join(t.TempDir(), "empty.txt") + require.NoError(t, os.WriteFile(path, nil, 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, upload.Close()) }) + + require.True(t, upload.knownSize) + require.True(t, upload.canReplay()) + require.Zero(t, upload.size) +} + func TestEmbedFilesClosesPartialUploadsOnError(t *testing.T) { upload := &recordingReadCloser{reader: strings.NewReader("already opened")} missingPath := FilePathValue(t.TempDir() + "/missing.bin") From 872425944c457cd1bbc71d4869e84e6b5aaff431 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Mon, 17 Aug 2026 19:16:38 +0000 Subject: [PATCH 4/9] Fix multipart replay lifecycle --- pkg/cmd/flagoptions.go | 30 +++-- pkg/cmd/multipartbody.go | 95 ++++++++++--- pkg/cmd/multipartbody_test.go | 221 +++++++++++++++++++++++++++++++ pkg/cmd/multipartfile_unix.go | 18 +++ pkg/cmd/multipartfile_windows.go | 29 ++++ 5 files changed, 366 insertions(+), 27 deletions(-) create mode 100644 pkg/cmd/multipartfile_unix.go create mode 100644 pkg/cmd/multipartfile_windows.go diff --git a/pkg/cmd/flagoptions.go b/pkg/cmd/flagoptions.go index a90a2d37..70947e4f 100644 --- a/pkg/cmd/flagoptions.go +++ b/pkg/cmd/flagoptions.go @@ -584,7 +584,7 @@ func (f fileUpload) replay() (fileUpload, error) { if !f.canReplay() { return fileUpload{}, errors.New("multipart upload is not replayable") } - reader, err := f.source.Reader() + reader, err := f.source.ReplayReader() if err != nil { return fileUpload{}, err } @@ -594,6 +594,8 @@ func (f fileUpload) replay() (fileUpload, error) { } type replayableFileSource struct { + // file is a stable identity anchor. Each request body reads from a duplicate + // handle so closing the anchor never truncates a transport-owned upload. file *os.File size int64 @@ -601,7 +603,7 @@ type replayableFileSource struct { closeErr error } -func (s *replayableFileSource) Reader() (io.Reader, error) { +func (s *replayableFileSource) ReplayReader() (io.ReadCloser, error) { info, err := s.file.Stat() if err != nil { return nil, err @@ -609,7 +611,14 @@ func (s *replayableFileSource) Reader() (io.Reader, error) { if !info.Mode().IsRegular() || info.Size() != s.size { return nil, fmt.Errorf("upload %q changed while being replayed", s.file.Name()) } - return io.NewSectionReader(s.file, 0, s.size), nil + file, err := duplicateFile(s.file) + if err != nil { + return nil, fmt.Errorf("duplicate upload %q: %w", s.file.Name(), err) + } + return §ionReadCloser{ + SectionReader: io.NewSectionReader(file, 0, s.size), + file: file, + }, nil } func (s *replayableFileSource) Close() error { @@ -619,6 +628,15 @@ func (s *replayableFileSource) Close() error { return s.closeErr } +type sectionReadCloser struct { + *io.SectionReader + file *os.File +} + +func (r *sectionReadCloser) Close() error { + return r.file.Close() +} + func hasTrustworthyFileSize(file *os.File, info os.FileInfo) bool { if !info.Mode().IsRegular() || info.Size() < 0 { return false @@ -668,12 +686,8 @@ func openFileUpload(path string) (fileUpload, error) { }, nil } source := &replayableFileSource{file: file, size: info.Size()} - reader, err := source.Reader() - if err != nil { - return fileUpload{}, errors.Join(err, source.Close()) - } return fileUpload{ - Reader: reader, + Reader: io.NewSectionReader(file, 0, info.Size()), filename: filepath.Base(path), contentType: contentType, size: info.Size(), diff --git a/pkg/cmd/multipartbody.go b/pkg/cmd/multipartbody.go index bead5b7f..f2427355 100644 --- a/pkg/cmd/multipartbody.go +++ b/pkg/cmd/multipartbody.go @@ -2,6 +2,7 @@ package cmd import ( "bytes" + "context" "errors" "fmt" "io" @@ -114,8 +115,7 @@ func (b *multipartRequestBody) encode() { // A final multipart boundary asserts that every part completed. Abort the // pipe immediately on a source error instead of making a truncated file // look like a valid zero-byte upload to the receiver. - _ = b.writer.CloseWithError(err) - _ = b.cleanup() + _ = b.writer.CloseWithError(errors.Join(err, b.cleanup())) return } @@ -189,6 +189,8 @@ type replayableMultipartRequest struct { attempts int closed bool closeErr error + + pendingCancel func() bool } func replayableMultipartRequestOptions( @@ -224,6 +226,10 @@ func (r *replayableMultipartRequest) middleware( ) (*http.Response, error) { attemptIndex := r.beginAttempt() attempt := &multipartReplayAttempt{request: r} + stopCancellation := context.AfterFunc(req.Context(), func() { + _ = attempt.abort() + _ = r.close() + }) // Discard the SDK's empty replay marker before installing the actual body. if req.Body != nil { @@ -231,6 +237,7 @@ func (r *replayableMultipartRequest) middleware( } body, err := attempt.newBody() if err != nil { + stopCancellation() return nil, errors.Join(err, r.close()) } req.Body = body @@ -238,23 +245,57 @@ func (r *replayableMultipartRequest) middleware( req.GetBody = attempt.newBody res, requestErr := next(req) - attemptErr := attempt.close() - if req.Context().Err() != nil || - !shouldRetryReplayableMultipart(res) || - attemptIndex >= replayableMultipartMaxRetries { + requestCanceled := req.Context().Err() != nil + terminal := requestCanceled || !shouldRetryReplayableMultipart(res) || + attemptIndex >= replayableMultipartMaxRetries + if terminal { + stopCancellation() + var attemptErr error + if requestErr != nil || requestCanceled || res == nil || + res.StatusCode >= http.StatusBadRequest { + attemptErr = attempt.abort() + } else { + attempt.release() + } return res, errors.Join(requestErr, attemptErr, r.close()) } + + // A retry must not race an earlier body. Abort its duplicate handles without + // waiting for an uninterruptible read, while retaining the stable request + // anchors needed to construct the next SDK attempt. + attemptErr := attempt.abort() + r.armPendingCancellation(stopCancellation) return res, errors.Join(requestErr, attemptErr) } func (r *replayableMultipartRequest) beginAttempt() int { r.mu.Lock() - defer r.mu.Unlock() + stopCancellation := r.pendingCancel + r.pendingCancel = nil attempt := r.attempts r.attempts++ + r.mu.Unlock() + if stopCancellation != nil { + stopCancellation() + } return attempt } +func (r *replayableMultipartRequest) armPendingCancellation(stopCancellation func() bool) { + r.mu.Lock() + if r.closed { + r.mu.Unlock() + stopCancellation() + return + } + previous := r.pendingCancel + r.pendingCancel = stopCancellation + r.mu.Unlock() + if previous != nil { + previous() + } +} + func (r *replayableMultipartRequest) newBody() (*multipartRequestBody, error) { bodyMap, err := replayMultipartBody(r.bodyMap) if err != nil { @@ -265,12 +306,18 @@ func (r *replayableMultipartRequest) newBody() (*multipartRequestBody, error) { func (r *replayableMultipartRequest) close() error { r.mu.Lock() - defer r.mu.Unlock() + stopCancellation := r.pendingCancel + r.pendingCancel = nil if !r.closed { r.closed = true r.closeErr = closeFileUploads(r.bodyMap) } - return r.closeErr + closeErr := r.closeErr + r.mu.Unlock() + if stopCancellation != nil { + stopCancellation() + } + return closeErr } type multipartReplayAttempt struct { @@ -297,24 +344,34 @@ func (a *multipartReplayAttempt) newBody() (io.ReadCloser, error) { return body, nil } -func (a *multipartReplayAttempt) close() error { - a.mu.Lock() - if a.closed { - a.mu.Unlock() - return nil - } - a.closed = true - bodies := append([]*multipartRequestBody(nil), a.bodies...) - a.mu.Unlock() +func (a *multipartReplayAttempt) release() { + _ = a.seal() +} +func (a *multipartReplayAttempt) abort() error { + bodies := a.seal() var errs []error for _, body := range bodies { errs = append(errs, body.Close()) - <-body.done } return errors.Join(errs...) } +// seal prevents net/http from creating another redirect body after an attempt +// returns. The caller either leaves the returned transport-owned bodies to +// finish naturally or closes them to abort an errored/canceled attempt. +func (a *multipartReplayAttempt) seal() []*multipartRequestBody { + a.mu.Lock() + defer a.mu.Unlock() + if a.closed { + return nil + } + a.closed = true + bodies := append([]*multipartRequestBody(nil), a.bodies...) + a.bodies = nil + return bodies +} + // Keep this predicate aligned with openai-go's retry policy. It is only used // to decide when the request-level file descriptors can be closed; the SDK // still owns retry timing, response cleanup, and the retry decision itself. diff --git a/pkg/cmd/multipartbody_test.go b/pkg/cmd/multipartbody_test.go index ed673a4c..610a7dc2 100644 --- a/pkg/cmd/multipartbody_test.go +++ b/pkg/cmd/multipartbody_test.go @@ -81,6 +81,24 @@ func TestMultipartRequestBodyPropagatesProducerErrors(t *testing.T) { require.NoError(t, body.Close()) } +func TestMultipartRequestBodyJoinsCleanupErrorAfterSourceError(t *testing.T) { + readErr := errors.New("upload read failed") + closeErr := errors.New("upload close failed") + upload := &recordingReadCloser{ + reader: errorReader{err: readErr}, + closeErr: closeErr, + } + body := newMultipartRequestBody(map[string]any{ + "file": fileUpload{Reader: upload, filename: "broken.bin"}, + }, apiform.FormatBrackets) + + _, err := io.ReadAll(body) + require.ErrorIs(t, err, readErr) + require.ErrorIs(t, err, closeErr) + waitForMultipartBody(t, body) + require.Equal(t, int32(1), upload.closeCount.Load()) +} + func TestMultipartRequestBodySourceErrorCannotBeParsedAsCompleteUpload(t *testing.T) { readErr := errors.New("upload read failed") parseErrCh := make(chan error, 1) @@ -448,6 +466,145 @@ func TestMultipartRequestOptionsCloseRegularFilesAfterFinalRetry(t *testing.T) { require.ErrorIs(t, err, os.ErrClosed) } +func TestMultipartRequestOptionsCloseRegularFilesWhenBackoffCanceled(t *testing.T) { + path := filepath.Join(t.TempDir(), "canceled-retry-upload.txt") + require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + sourceFile := upload.source.file + + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + _ = r.Body.Close() + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "10") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"rate_limit_error"}}`) + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + attemptReturned := make(chan struct{}) + observeAttempt := option.WithMiddleware(func( + req *http.Request, + next option.MiddlewareNext, + ) (*http.Response, error) { + res, requestErr := next(req) + close(attemptReturned) + return res, requestErr + }) + options = append([]option.RequestOption{observeAttempt}, options...) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + ctx, cancel := context.WithCancel(context.Background()) + requestDone := make(chan error, 1) + go func() { + _, requestErr := client.Files.New(ctx, openai.FileNewParams{}, options...) + requestDone <- requestErr + }() + + select { + case <-attemptReturned: + case <-time.After(multipartTestTimeout): + t.Fatal("first retryable attempt did not return") + } + cancel() + select { + case requestErr := <-requestDone: + require.ErrorIs(t, requestErr, context.Canceled) + case <-time.After(multipartTestTimeout): + t.Fatal("request did not stop after cancellation during retry backoff") + } + require.Equal(t, int32(1), requestCount.Load()) + require.Eventually(t, func() bool { + _, statErr := sourceFile.Stat() + return errors.Is(statErr, os.ErrClosed) + }, time.Second, time.Millisecond, "cancellation cleanup must close the replay source") +} + +func TestMultipartRequestOptionsLetEarlySuccessfulUploadFinish(t *testing.T) { + path := filepath.Join(t.TempDir(), "early-response-upload.txt") + require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("payload", 1024)), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + sourceFile := upload.source.file + + doer := newEarlyResponseHTTPDoer() + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL("https://example.com/"), + option.WithHTTPClient(doer), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.NoError(t, err) + close(doer.release) + select { + case <-doer.done: + case <-time.After(multipartTestTimeout): + t.Fatal("transport did not finish consuming the request body") + } + require.NoError(t, doer.err) + require.Positive(t, doer.contentLength) + require.Equal(t, doer.contentLength, doer.bytesRead) + _, err = sourceFile.Stat() + require.ErrorIs(t, err, os.ErrClosed) +} + +func TestMultipartReplayAttemptAbortDoesNotWaitForBlockedReader(t *testing.T) { + source := newManuallyReleasedReader() + body := newMultipartRequestBody(map[string]any{ + "file": fileUpload{Reader: source, filename: "blocked.bin"}, + }, apiform.FormatBrackets) + attempt := &multipartReplayAttempt{bodies: []*multipartRequestBody{body}} + readDone := make(chan error, 1) + go func() { + _, err := io.Copy(io.Discard, body) + readDone <- err + }() + select { + case <-source.started: + case <-time.After(multipartTestTimeout): + t.Fatal("multipart encoder did not start the blocked source read") + } + + closeDone := make(chan error, 1) + go func() { closeDone <- attempt.abort() }() + returnedPromptly := false + select { + case <-closeDone: + returnedPromptly = true + case <-time.After(100 * time.Millisecond): + } + close(source.release) + if !returnedPromptly { + select { + case <-closeDone: + case <-time.After(multipartTestTimeout): + t.Fatal("attempt cleanup did not return after releasing the source") + } + } + select { + case <-readDone: + case <-time.After(multipartTestTimeout): + t.Fatal("request body read did not stop") + } + require.True(t, returnedPromptly, "attempt cleanup must not wait for an uninterruptible source read") +} + func TestMultipartRequestOptionsRedirectReplayUsesOpenedFile(t *testing.T) { if runtime.GOOS == "windows" { t.Skip("Windows does not permit atomically replacing this open test file") @@ -724,6 +881,70 @@ type recordingReadCloser struct { closeCount atomic.Int32 } +type earlyResponseHTTPDoer struct { + release chan struct{} + done chan struct{} + + contentLength int64 + bytesRead int64 + err error +} + +func newEarlyResponseHTTPDoer() *earlyResponseHTTPDoer { + return &earlyResponseHTTPDoer{ + release: make(chan struct{}), + done: make(chan struct{}), + } +} + +func (d *earlyResponseHTTPDoer) Do(req *http.Request) (*http.Response, error) { + readStarted := make(chan struct{}) + d.contentLength = req.ContentLength + go func() { + defer close(d.done) + buf := make([]byte, 1) + n, readErr := req.Body.Read(buf) + d.bytesRead += int64(n) + close(readStarted) + <-d.release + if readErr == nil { + copied, copyErr := io.Copy(io.Discard, req.Body) + d.bytesRead += copied + readErr = copyErr + } + if errors.Is(readErr, io.EOF) { + readErr = nil + } + d.err = errors.Join(readErr, req.Body.Close()) + }() + <-readStarted + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"file_123","object":"file"}`)), + Request: req, + }, nil +} + +type manuallyReleasedReader struct { + started chan struct{} + release chan struct{} + once sync.Once +} + +func newManuallyReleasedReader() *manuallyReleasedReader { + return &manuallyReleasedReader{ + started: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (r *manuallyReleasedReader) Read([]byte) (int, error) { + r.once.Do(func() { close(r.started) }) + <-r.release + return 0, io.EOF +} + func (r *recordingReadCloser) Read(p []byte) (int, error) { return r.reader.Read(p) } diff --git a/pkg/cmd/multipartfile_unix.go b/pkg/cmd/multipartfile_unix.go new file mode 100644 index 00000000..43b5682e --- /dev/null +++ b/pkg/cmd/multipartfile_unix.go @@ -0,0 +1,18 @@ +//go:build !windows + +package cmd + +import ( + "os" + + "golang.org/x/sys/unix" +) + +func duplicateFile(file *os.File) (*os.File, error) { + fd, err := unix.Dup(int(file.Fd())) + if err != nil { + return nil, err + } + unix.CloseOnExec(fd) + return os.NewFile(uintptr(fd), file.Name()), nil +} diff --git a/pkg/cmd/multipartfile_windows.go b/pkg/cmd/multipartfile_windows.go new file mode 100644 index 00000000..219012ce --- /dev/null +++ b/pkg/cmd/multipartfile_windows.go @@ -0,0 +1,29 @@ +//go:build windows + +package cmd + +import ( + "os" + + "golang.org/x/sys/windows" +) + +func duplicateFile(file *os.File) (*os.File, error) { + process, err := windows.GetCurrentProcess() + if err != nil { + return nil, err + } + var duplicate windows.Handle + if err := windows.DuplicateHandle( + process, + windows.Handle(file.Fd()), + process, + &duplicate, + 0, + false, + windows.DUPLICATE_SAME_ACCESS, + ); err != nil { + return nil, err + } + return os.NewFile(uintptr(duplicate), file.Name()), nil +} From af8013d49bdc4dd06563160b9eb33edc60513e77 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Mon, 17 Aug 2026 22:04:48 +0000 Subject: [PATCH 5/9] Preserve multipart upload content across replays --- pkg/cmd/flagoptions.go | 150 --------------- pkg/cmd/multipartfile.go | 232 ++++++++++++++++++++++ pkg/cmd/multipartfile_unix.go | 12 ++ pkg/cmd/multipartfile_windows.go | 24 +++ pkg/cmd/multipartintegrity_test.go | 296 +++++++++++++++++++++++++++++ 5 files changed, 564 insertions(+), 150 deletions(-) create mode 100644 pkg/cmd/multipartfile.go create mode 100644 pkg/cmd/multipartintegrity_test.go diff --git a/pkg/cmd/flagoptions.go b/pkg/cmd/flagoptions.go index 70947e4f..459ca567 100644 --- a/pkg/cmd/flagoptions.go +++ b/pkg/cmd/flagoptions.go @@ -7,13 +7,10 @@ import ( "fmt" "io" "maps" - "mime" "net/http" "os" - "path/filepath" "reflect" "strings" - "sync" "unicode/utf8" "github.com/openai/openai-cli/internal/apiform" @@ -550,153 +547,6 @@ func flagOptions( // as a file path without needing the "@" prefix. type FilePathValue string -// fileUpload wraps an io.Reader with filename and content-type metadata for -// use as a multipart form part. The apiform encoder detects the Filename and -// ContentType methods and uses them to populate the Content-Disposition -// filename and the Content-Type header on the part. -type fileUpload struct { - io.Reader // apiform checks for reader and reads its contents during encode - filename string - contentType string - size int64 - knownSize bool - source *replayableFileSource - ownsSource bool -} - -func (f fileUpload) Filename() string { return f.filename } -func (f fileUpload) ContentType() string { return f.contentType } -func (f fileUpload) Close() error { - if f.ownsSource && f.source != nil { - return f.source.Close() - } - if c, ok := f.Reader.(io.Closer); ok { - return c.Close() - } - return nil -} - -func (f fileUpload) canReplay() bool { - return f.source != nil && f.knownSize -} - -func (f fileUpload) replay() (fileUpload, error) { - if !f.canReplay() { - return fileUpload{}, errors.New("multipart upload is not replayable") - } - reader, err := f.source.ReplayReader() - if err != nil { - return fileUpload{}, err - } - f.Reader = reader - f.ownsSource = false - return f, nil -} - -type replayableFileSource struct { - // file is a stable identity anchor. Each request body reads from a duplicate - // handle so closing the anchor never truncates a transport-owned upload. - file *os.File - size int64 - - closeOnce sync.Once - closeErr error -} - -func (s *replayableFileSource) ReplayReader() (io.ReadCloser, error) { - info, err := s.file.Stat() - if err != nil { - return nil, err - } - if !info.Mode().IsRegular() || info.Size() != s.size { - return nil, fmt.Errorf("upload %q changed while being replayed", s.file.Name()) - } - file, err := duplicateFile(s.file) - if err != nil { - return nil, fmt.Errorf("duplicate upload %q: %w", s.file.Name(), err) - } - return §ionReadCloser{ - SectionReader: io.NewSectionReader(file, 0, s.size), - file: file, - }, nil -} - -func (s *replayableFileSource) Close() error { - s.closeOnce.Do(func() { - s.closeErr = s.file.Close() - }) - return s.closeErr -} - -type sectionReadCloser struct { - *io.SectionReader - file *os.File -} - -func (r *sectionReadCloser) Close() error { - return r.file.Close() -} - -func hasTrustworthyFileSize(file *os.File, info os.FileInfo) bool { - if !info.Mode().IsRegular() || info.Size() < 0 { - return false - } - - // Regular-looking virtual files can report a size unrelated to their - // readable contents. Verify both sides of the reported EOF without moving - // the descriptor offset; ordinary and sparse files support these ReadAt - // probes, while procfs/sysfs-style sources fail or expose extra bytes. - var probe [1]byte - if info.Size() == 0 { - n, err := file.ReadAt(probe[:], 0) - return n == 0 && errors.Is(err, io.EOF) - } - n, err := file.ReadAt(probe[:], info.Size()-1) - if n != 1 || err != nil { - return false - } - n, err = file.ReadAt(probe[:], info.Size()) - return n == 0 && errors.Is(err, io.EOF) -} - -// openFileUpload opens the file at path and returns a fileUpload whose filename -// is the path's basename and whose content type is derived from the file -// extension (falling back to application/octet-stream when unknown). -func openFileUpload(path string) (fileUpload, error) { - file, err := os.Open(path) - if err != nil { - return fileUpload{}, err - } - info, err := file.Stat() - if err != nil { - return fileUpload{}, errors.Join(err, file.Close()) - } - if info.IsDir() { - return fileUpload{}, errors.Join(fmt.Errorf("read %s: is a directory", path), file.Close()) - } - contentType := mime.TypeByExtension(filepath.Ext(path)) - if contentType == "" { - contentType = "application/octet-stream" - } - if !hasTrustworthyFileSize(file, info) { - return fileUpload{ - Reader: file, - filename: filepath.Base(path), - contentType: contentType, - }, nil - } - source := &replayableFileSource{file: file, size: info.Size()} - return fileUpload{ - Reader: io.NewSectionReader(file, 0, info.Size()), - filename: filepath.Base(path), - contentType: contentType, - size: info.Size(), - knownSize: true, - source: source, - ownsSource: true, - }, nil -} - // applyDataAliases rewrites keys in a body map based on flag `DataAliases` metadata. For top-level flags, // `{alias: value}` becomes `{canonical: value}`. For inner flags (those registered under an outer flag // via WithInnerFlags), the alias translation is also applied to the nested map under the outer flag's diff --git a/pkg/cmd/multipartfile.go b/pkg/cmd/multipartfile.go new file mode 100644 index 00000000..70643329 --- /dev/null +++ b/pkg/cmd/multipartfile.go @@ -0,0 +1,232 @@ +package cmd + +import ( + "errors" + "fmt" + "io" + "mime" + "os" + "path/filepath" + "sync" +) + +// fileUpload wraps an io.Reader with filename and content-type metadata for +// use as a multipart form part. The apiform encoder detects the Filename and +// ContentType methods and uses them to populate the Content-Disposition +// filename and the Content-Type header on the part. +type fileUpload struct { + io.Reader // apiform checks for reader and reads its contents during encode + filename string + contentType string + size int64 + knownSize bool + source *replayableFileSource + ownsSource bool +} + +func (f fileUpload) Filename() string { return f.filename } +func (f fileUpload) ContentType() string { return f.contentType } +func (f fileUpload) Close() error { + if f.ownsSource && f.source != nil { + return f.source.Close() + } + if closer, ok := f.Reader.(io.Closer); ok { + return closer.Close() + } + return nil +} + +func (f fileUpload) canReplay() bool { + return f.source != nil && f.knownSize +} + +func (f fileUpload) replay() (fileUpload, error) { + if !f.canReplay() { + return fileUpload{}, errors.New("multipart upload is not replayable") + } + reader, err := f.source.ReplayReader() + if err != nil { + return fileUpload{}, err + } + f.Reader = reader + f.ownsSource = false + return f, nil +} + +// replayableFileSource owns an immutable, bounded on-disk snapshot. Each body +// reads from a duplicate handle, so redirects and retries always send the same +// bytes and closing the anchor cannot truncate a transport-owned upload. +type replayableFileSource struct { + file *os.File + size int64 + sourceName string + + closeOnce sync.Once + closeErr error +} + +func (s *replayableFileSource) ReplayReader() (io.ReadCloser, error) { + info, err := s.file.Stat() + if err != nil { + return nil, err + } + if !info.Mode().IsRegular() || info.Size() != s.size { + return nil, fmt.Errorf("upload snapshot for %q changed while being replayed", s.displayName()) + } + file, err := duplicateFile(s.file) + if err != nil { + return nil, fmt.Errorf("duplicate upload snapshot for %q: %w", s.displayName(), err) + } + return &exactLengthReadCloser{ + exactLengthReader: exactLengthReader{ + reader: io.NewSectionReader(file, 0, s.size), + remaining: s.size, + }, + closer: file, + }, nil +} + +func (s *replayableFileSource) displayName() string { + if s.sourceName != "" { + return s.sourceName + } + return s.file.Name() +} + +func (s *replayableFileSource) Close() error { + s.closeOnce.Do(func() { + s.closeErr = s.file.Close() + }) + return s.closeErr +} + +type exactLengthReader struct { + reader io.Reader + remaining int64 +} + +func (r *exactLengthReader) Read(p []byte) (int, error) { + if r.remaining == 0 { + return 0, io.EOF + } + if int64(len(p)) > r.remaining { + p = p[:r.remaining] + } + n, err := r.reader.Read(p) + r.remaining -= int64(n) + if errors.Is(err, io.EOF) { + if r.remaining > 0 { + return n, io.ErrUnexpectedEOF + } + return n, nil + } + return n, err +} + +type exactLengthReadCloser struct { + exactLengthReader + closer io.Closer +} + +func (r *exactLengthReadCloser) Close() error { + return r.closer.Close() +} + +func hasTrustworthyFileSize(file *os.File, info os.FileInfo) bool { + if !info.Mode().IsRegular() || info.Size() < 0 { + return false + } + + // Regular-looking virtual files can report a size unrelated to their + // readable contents. Verify both sides of the reported EOF without moving + // the descriptor offset; ordinary and sparse files support these ReadAt + // probes, while procfs/sysfs-style sources fail or expose extra bytes. + var probe [1]byte + if info.Size() == 0 { + n, err := file.ReadAt(probe[:], 0) + return n == 0 && errors.Is(err, io.EOF) + } + n, err := file.ReadAt(probe[:], info.Size()-1) + if n != 1 || err != nil { + return false + } + n, err = file.ReadAt(probe[:], info.Size()) + return n == 0 && errors.Is(err, io.EOF) +} + +// openFileUpload opens path once. Ordinary files are copied into a secure, +// size-bounded temporary snapshot so multipart redirects and retries are +// replayable without retaining the caller's mutable file or buffering it in +// memory. Virtual and unknown-length files remain single-shot streams. +func openFileUpload(path string) (fileUpload, error) { + file, err := os.Open(path) + if err != nil { + return fileUpload{}, err + } + info, err := file.Stat() + if err != nil { + return fileUpload{}, errors.Join(err, file.Close()) + } + if info.IsDir() { + return fileUpload{}, errors.Join(fmt.Errorf("read %s: is a directory", path), file.Close()) + } + contentType := mime.TypeByExtension(filepath.Ext(path)) + if contentType == "" { + contentType = "application/octet-stream" + } + if !hasTrustworthyFileSize(file, info) { + return fileUpload{ + Reader: file, + filename: filepath.Base(path), + contentType: contentType, + }, nil + } + + source, err := snapshotReplayableFile(file, info.Size(), path) + if err != nil { + return fileUpload{}, err + } + return fileUpload{ + Reader: &exactLengthReader{ + reader: io.NewSectionReader(source.file, 0, source.size), + remaining: source.size, + }, + filename: filepath.Base(path), + contentType: contentType, + size: source.size, + knownSize: true, + source: source, + ownsSource: true, + }, nil +} + +func snapshotReplayableFile(source *os.File, size int64, sourceName string) (*replayableFileSource, error) { + snapshot, err := createReplaySnapshotFile() + if err != nil { + return nil, errors.Join(fmt.Errorf("snapshot upload %q: %w", sourceName, err), source.Close()) + } + fail := func(cause error) (*replayableFileSource, error) { + return nil, errors.Join(cause, source.Close(), snapshot.Close()) + } + + exactSource := &exactLengthReader{reader: source, remaining: size} + if _, err := io.Copy(snapshot, exactSource); err != nil { + return fail(fmt.Errorf("snapshot upload %q: %w", sourceName, err)) + } + var extra [1]byte + n, readErr := source.Read(extra[:]) + if n != 0 || !errors.Is(readErr, io.EOF) { + if readErr == nil { + readErr = errors.New("source grew while being snapshotted") + } + return fail(fmt.Errorf("snapshot upload %q: %w", sourceName, readErr)) + } + if err := source.Close(); err != nil { + return nil, errors.Join(fmt.Errorf("close upload %q after snapshot: %w", sourceName, err), snapshot.Close()) + } + return &replayableFileSource{ + file: snapshot, + size: size, + sourceName: sourceName, + }, nil +} diff --git a/pkg/cmd/multipartfile_unix.go b/pkg/cmd/multipartfile_unix.go index 43b5682e..9e5daf31 100644 --- a/pkg/cmd/multipartfile_unix.go +++ b/pkg/cmd/multipartfile_unix.go @@ -3,11 +3,23 @@ package cmd import ( + "errors" "os" "golang.org/x/sys/unix" ) +func createReplaySnapshotFile() (*os.File, error) { + file, err := os.CreateTemp("", "openai-multipart-upload-*") + if err != nil { + return nil, err + } + if err := os.Remove(file.Name()); err != nil { + return nil, errors.Join(err, file.Close()) + } + return file, nil +} + func duplicateFile(file *os.File) (*os.File, error) { fd, err := unix.Dup(int(file.Fd())) if err != nil { diff --git a/pkg/cmd/multipartfile_windows.go b/pkg/cmd/multipartfile_windows.go index 219012ce..f2c8a6a6 100644 --- a/pkg/cmd/multipartfile_windows.go +++ b/pkg/cmd/multipartfile_windows.go @@ -3,11 +3,35 @@ package cmd import ( + "errors" "os" "golang.org/x/sys/windows" ) +func createReplaySnapshotFile() (*os.File, error) { + placeholder, err := os.CreateTemp("", "openai-multipart-upload-*") + if err != nil { + return nil, err + } + path := placeholder.Name() + if err := placeholder.Close(); err != nil { + return nil, errors.Join(err, os.Remove(path)) + } + if err := os.Remove(path); err != nil { + return nil, err + } + file, err := os.OpenFile( + path, + os.O_RDWR|os.O_CREATE|os.O_EXCL|windows.O_FILE_FLAG_DELETE_ON_CLOSE, + 0o600, + ) + if err != nil { + return nil, errors.Join(err, os.Remove(path)) + } + return file, nil +} + func duplicateFile(file *os.File) (*os.File, error) { process, err := windows.GetCurrentProcess() if err != nil { diff --git a/pkg/cmd/multipartintegrity_test.go b/pkg/cmd/multipartintegrity_test.go new file mode 100644 index 00000000..480ea39d --- /dev/null +++ b/pkg/cmd/multipartintegrity_test.go @@ -0,0 +1,296 @@ +package cmd + +import ( + "context" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/openai/openai-cli/internal/apiform" + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/option" + "github.com/stretchr/testify/require" + "github.com/urfave/cli/v3" +) + +func TestFilesCreateCLIUsesImmutableReplayContent(t *testing.T) { + for _, status := range []int{ + http.StatusTemporaryRedirect, + http.StatusPermanentRedirect, + http.StatusTooManyRequests, + } { + t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { + path := filepath.Join(t.TempDir(), "upload.txt") + require.NoError(t, os.WriteFile(path, []byte("benign-data"), 0o600)) + + var mu sync.Mutex + var uploads []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + contents, readErr := readMultipartUpload(r) + if readErr != nil { + http.Error(w, readErr.Error(), http.StatusBadRequest) + return + } + mu.Lock() + uploads = append(uploads, string(contents)) + requestNumber := len(uploads) + mu.Unlock() + + if requestNumber == 1 { + if writeErr := os.WriteFile(path, []byte("secret-data"), 0o600); writeErr != nil { + http.Error(w, writeErr.Error(), http.StatusInternalServerError) + return + } + if status == http.StatusTooManyRequests { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "0") + w.WriteHeader(status) + _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"rate_limit_error"}}`) + return + } + http.Redirect(w, r, "/upload", status) + return + } + + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + require.NoError(t, runFilesCreateCLI(t.Context(), server.URL+"/", path)) + mu.Lock() + require.Equal(t, []string{"benign-data", "benign-data"}, uploads) + mu.Unlock() + }) + } +} + +func TestFilesCreateCLISnapshotsBeforeSourceShrinks(t *testing.T) { + const uploadSize = 16 << 20 + path := filepath.Join(t.TempDir(), "shrinking.bin") + require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("a", uploadSize)), 0o600)) + + type receiverResult struct { + bytes int64 + committed bool + err error + } + resultCh := make(chan receiverResult, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mediaType, params, parseErr := mime.ParseMediaType(r.Header.Get("Content-Type")) + if parseErr != nil || mediaType != "multipart/form-data" { + resultCh <- receiverResult{err: parseErr} + http.Error(w, "invalid content type", http.StatusBadRequest) + return + } + reader := multipart.NewReader(r.Body, params["boundary"]) + part, partErr := nextFilePart(reader) + if partErr != nil { + resultCh <- receiverResult{err: partErr} + http.Error(w, partErr.Error(), http.StatusBadRequest) + return + } + buf := make([]byte, 1) + n, readErr := part.Read(buf) + if readErr == nil { + readErr = os.Truncate(path, 0) + } + copied, copyErr := io.Copy(io.Discard, part) + if readErr == nil { + readErr = copyErr + } + finalErr := consumeMultipartRemainder(reader) + resultCh <- receiverResult{ + bytes: int64(n) + copied, + committed: errors.Is(finalErr, io.EOF), + err: errors.Join(readErr, incompleteMultipartError(finalErr)), + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + require.NoError(t, runFilesCreateCLI(t.Context(), server.URL+"/", path)) + result := <-resultCh + require.NoError(t, result.err) + require.True(t, result.committed) + require.Equal(t, int64(uploadSize), result.bytes) +} + +func TestMultipartReplayPrematureEOFDoesNotCommitFinalBoundary(t *testing.T) { + const uploadSize = 8 << 20 + path := filepath.Join(t.TempDir(), "replay-source.bin") + require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("a", uploadSize)), 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + source := &replayableFileSource{file: file, size: uploadSize} + upload := fileUpload{ + Reader: io.NewSectionReader(file, 0, uploadSize), + filename: filepath.Base(path), + contentType: "application/octet-stream", + size: uploadSize, + knownSize: true, + source: source, + ownsSource: true, + } + + type receiverResult struct { + committed bool + err error + } + resultCh := make(chan receiverResult, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mediaType, params, parseErr := mime.ParseMediaType(r.Header.Get("Content-Type")) + if parseErr != nil || mediaType != "multipart/form-data" { + resultCh <- receiverResult{err: parseErr} + return + } + reader := multipart.NewReader(r.Body, params["boundary"]) + part, partErr := nextFilePart(reader) + if partErr != nil { + resultCh <- receiverResult{err: partErr} + return + } + buf := make([]byte, 1) + _, readErr := part.Read(buf) + if readErr == nil { + readErr = os.Truncate(path, 0) + } + if readErr == nil { + _, readErr = io.Copy(io.Discard, part) + } + finalErr := consumeMultipartRemainder(reader) + resultCh <- receiverResult{ + committed: errors.Is(finalErr, io.EOF), + err: errors.Join(readErr, incompleteMultipartError(finalErr)), + } + })) + t.Cleanup(server.Close) + + options, err := multipartRequestOptions(map[string]any{ + "file": upload, + "purpose": "assistants", + }, apiform.FormatBrackets) + require.NoError(t, err) + options = append(options, option.WithMaxRetries(0)) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/"), + ) + + _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) + require.ErrorIs(t, err, io.ErrUnexpectedEOF) + result := <-resultCh + require.False(t, result.committed, "an incomplete replay source must not emit a terminal multipart boundary") + require.Error(t, result.err) +} + +func TestOpenFileUploadRemovesSnapshotOnClose(t *testing.T) { + path := filepath.Join(t.TempDir(), "upload.txt") + require.NoError(t, os.WriteFile(path, []byte("snapshot payload"), 0o600)) + upload, err := openFileUpload(path) + require.NoError(t, err) + snapshotPath := upload.source.file.Name() + + require.NoError(t, upload.Close()) + _, err = os.Stat(snapshotPath) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestSnapshotReplayableFileRejectsPrematureEOF(t *testing.T) { + path := filepath.Join(t.TempDir(), "short.txt") + require.NoError(t, os.WriteFile(path, []byte("short"), 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + + _, err = snapshotReplayableFile(file, int64(len("longer than short")), path) + require.ErrorIs(t, err, io.ErrUnexpectedEOF) +} + +func readMultipartUpload(r *http.Request) ([]byte, error) { + if err := r.ParseMultipartForm(1 << 20); err != nil { + return nil, err + } + file, _, err := r.FormFile("file") + if err != nil { + return nil, err + } + defer file.Close() + return io.ReadAll(file) +} + +func nextFilePart(reader *multipart.Reader) (*multipart.Part, error) { + for { + part, err := reader.NextPart() + if err != nil { + return nil, err + } + if part.FormName() == "file" { + return part, nil + } + if _, err := io.Copy(io.Discard, part); err != nil { + return nil, err + } + if err := part.Close(); err != nil { + return nil, err + } + } +} + +func consumeMultipartRemainder(reader *multipart.Reader) error { + for { + part, err := reader.NextPart() + if err != nil { + return err + } + if _, err := io.Copy(io.Discard, part); err != nil { + return err + } + if err := part.Close(); err != nil { + return err + } + } +} + +func incompleteMultipartError(err error) error { + if errors.Is(err, io.EOF) { + return nil + } + return err +} + +func runFilesCreateCLI(ctx context.Context, baseURL, path string) error { + create := filesCreate + command := &cli.Command{ + Name: "openai", + Flags: []cli.Flag{ + &cli.BoolFlag{Name: "debug"}, + &cli.StringFlag{Name: "base-url"}, + &cli.StringFlag{Name: "format", Value: "json"}, + &cli.StringFlag{Name: "transform"}, + &cli.BoolFlag{Name: "raw-output"}, + &cli.StringFlag{Name: "api-key"}, + }, + Commands: []*cli.Command{{ + Name: "files", + Commands: []*cli.Command{&create}, + }}, + } + return command.Run(ctx, []string{ + "openai", + "--api-key", "test-key", + "--base-url", baseURL, + "files", "create", + "--file", path, + "--purpose", "assistants", + }) +} From 602f07376bd06349e9a2264f55ca2fa56d3f391d Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 18 Aug 2026 00:45:31 +0000 Subject: [PATCH 6/9] fix: stream multipart before optional replay capture --- pkg/cmd/flagoptions.go | 3 +- pkg/cmd/multipartbody.go | 132 +++++++++++------ pkg/cmd/multipartbody_test.go | 27 ++-- pkg/cmd/multipartfallback_test.go | 196 +++++++++++++++++++++++++ pkg/cmd/multipartfile.go | 142 ++++-------------- pkg/cmd/multipartfile_windows.go | 18 ++- pkg/cmd/multipartfile_windows_test.go | 26 ++++ pkg/cmd/multipartintegrity_test.go | 46 +++--- pkg/cmd/multipartreplay.go | 160 ++++++++++++++++++++ pkg/cmd/multipartresource_unix_test.go | 117 +++++++++++++++ 10 files changed, 674 insertions(+), 193 deletions(-) create mode 100644 pkg/cmd/multipartfallback_test.go create mode 100644 pkg/cmd/multipartfile_windows_test.go create mode 100644 pkg/cmd/multipartreplay.go create mode 100644 pkg/cmd/multipartresource_unix_test.go diff --git a/pkg/cmd/flagoptions.go b/pkg/cmd/flagoptions.go index 459ca567..8b20ff85 100644 --- a/pkg/cmd/flagoptions.go +++ b/pkg/cmd/flagoptions.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "io/fs" "maps" "net/http" "os" @@ -272,7 +273,7 @@ func embedFilesValue(v reflect.Value, embedStyle FileEmbedStyle, stdin *onceStdi upload, err := openFileUpload(filename) if err != nil { - if !expectsFile { + if !expectsFile && errors.Is(err, fs.ErrNotExist) { // For strings that start with "@" and don't look like a filename, return the string return v, nil } diff --git a/pkg/cmd/multipartbody.go b/pkg/cmd/multipartbody.go index f2427355..6e6d3199 100644 --- a/pkg/cmd/multipartbody.go +++ b/pkg/cmd/multipartbody.go @@ -133,9 +133,10 @@ func (b *multipartRequestBody) cleanup() error { // multipartRequestOptions preserves the SDK's buffered, replayable behavior // when a multipart form contains only scalar fields. Uploads are streamed. A -// form containing only regular files is still replayable for HTTP redirects and -// carries an exact Content-Length; forms containing stdin, pipes, or arbitrary -// readers use chunked transfer and fail explicitly if a 307/308 requires replay. +// form containing only regular files carries an exact Content-Length and +// opportunistically captures the first encoded stream for redirects and +// retries. Forms containing stdin, pipes, or arbitrary readers use chunked +// transfer and fail explicitly if a 307/308 requires replay. func multipartRequestOptions( bodyMap map[string]any, encodingFormat apiform.FormFormat, @@ -147,7 +148,7 @@ func multipartRequestOptions( boundary := multipart.NewWriter(io.Discard).Boundary() var contentLength int64 - if info.replayable { + if info.knownLength { var err error contentLength, err = multipartContentLength(bodyMap, encodingFormat, boundary) if err != nil { @@ -155,7 +156,7 @@ func multipartRequestOptions( } } - if info.replayable { + if info.knownLength { return replayableMultipartRequestOptions(bodyMap, encodingFormat, boundary, contentLength) } @@ -184,6 +185,7 @@ type replayableMultipartRequest struct { encodingFormat apiform.FormFormat boundary string contentLength int64 + source *multipartReplaySource mu sync.Mutex attempts int @@ -208,6 +210,7 @@ func replayableMultipartRequestOptions( encodingFormat: encodingFormat, boundary: boundary, contentLength: contentLength, + source: newMultipartReplaySource(contentLength), } // The empty *bytes.Reader is deliberate: it makes openai-go install GetBody @@ -225,7 +228,7 @@ func (r *replayableMultipartRequest) middleware( next option.MiddlewareNext, ) (*http.Response, error) { attemptIndex := r.beginAttempt() - attempt := &multipartReplayAttempt{request: r} + attempt := &multipartReplayAttempt{request: r, firstAttempt: attemptIndex == 0} stopCancellation := context.AfterFunc(req.Context(), func() { _ = attempt.abort() _ = r.close() @@ -246,8 +249,15 @@ func (r *replayableMultipartRequest) middleware( res, requestErr := next(req) requestCanceled := req.Context().Err() != nil + replayAvailable := r.source.Ready() + if res != nil && shouldRetryReplayableMultipart(res) && !replayAvailable { + // The first upload is still valid when the optional replay snapshot + // cannot be created. Tell openai-go not to issue another network + // attempt that this request cannot reproduce faithfully. + res.Header.Set("x-should-retry", "false") + } terminal := requestCanceled || !shouldRetryReplayableMultipart(res) || - attemptIndex >= replayableMultipartMaxRetries + attemptIndex >= replayableMultipartMaxRetries || !replayAvailable if terminal { stopCancellation() var attemptErr error @@ -255,7 +265,8 @@ func (r *replayableMultipartRequest) middleware( res.StatusCode >= http.StatusBadRequest { attemptErr = attempt.abort() } else { - attempt.release() + attemptErr = attempt.release() + return res, errors.Join(requestErr, attemptErr) } return res, errors.Join(requestErr, attemptErr, r.close()) } @@ -296,12 +307,12 @@ func (r *replayableMultipartRequest) armPendingCancellation(stopCancellation fun } } -func (r *replayableMultipartRequest) newBody() (*multipartRequestBody, error) { - bodyMap, err := replayMultipartBody(r.bodyMap) +func (r *replayableMultipartRequest) newFirstBody() (io.ReadCloser, error) { + body, err := newMultipartRequestBodyWithBoundary(r.bodyMap, r.encodingFormat, r.boundary) if err != nil { return nil, err } - return newMultipartRequestBodyWithBoundary(bodyMap, r.encodingFormat, r.boundary) + return &capturingMultipartBody{body: body, source: r.source}, nil } func (r *replayableMultipartRequest) close() error { @@ -310,7 +321,7 @@ func (r *replayableMultipartRequest) close() error { r.pendingCancel = nil if !r.closed { r.closed = true - r.closeErr = closeFileUploads(r.bodyMap) + r.closeErr = errors.Join(closeFileUploads(r.bodyMap), r.source.Close()) } closeErr := r.closeErr r.mu.Unlock() @@ -321,15 +332,32 @@ func (r *replayableMultipartRequest) close() error { } type multipartReplayAttempt struct { - request *replayableMultipartRequest + request *replayableMultipartRequest + firstAttempt bool - mu sync.Mutex - bodies []*multipartRequestBody - closed bool + mu sync.Mutex + bodies []io.ReadCloser + bodyCount int + closed bool } func (a *multipartReplayAttempt) newBody() (io.ReadCloser, error) { - body, err := a.request.newBody() + a.mu.Lock() + if a.closed { + a.mu.Unlock() + return nil, errors.New("multipart replay attempt is closed") + } + bodyIndex := a.bodyCount + a.bodyCount++ + a.mu.Unlock() + + var body io.ReadCloser + var err error + if a.firstAttempt && bodyIndex == 0 { + body, err = a.request.newFirstBody() + } else { + body, err = a.request.source.ReplayReader() + } if err != nil { return nil, err } @@ -344,8 +372,35 @@ func (a *multipartReplayAttempt) newBody() (io.ReadCloser, error) { return body, nil } -func (a *multipartReplayAttempt) release() { - _ = a.seal() +func (a *multipartReplayAttempt) release() error { + bodies := a.seal() + var waiters []<-chan struct{} + for _, body := range bodies { + if waiter, ok := body.(interface{ Done() <-chan struct{} }); ok { + waiters = append(waiters, waiter.Done()) + } + } + if len(waiters) == 0 { + return a.request.close() + } + allDone := true + for _, done := range waiters { + select { + case <-done: + default: + allDone = false + } + } + if allDone { + return a.request.close() + } + go func() { + for _, done := range waiters { + <-done + } + _ = a.request.close() + }() + return nil } func (a *multipartReplayAttempt) abort() error { @@ -360,14 +415,14 @@ func (a *multipartReplayAttempt) abort() error { // seal prevents net/http from creating another redirect body after an attempt // returns. The caller either leaves the returned transport-owned bodies to // finish naturally or closes them to abort an errored/canceled attempt. -func (a *multipartReplayAttempt) seal() []*multipartRequestBody { +func (a *multipartReplayAttempt) seal() []io.ReadCloser { a.mu.Lock() defer a.mu.Unlock() if a.closed { return nil } a.closed = true - bodies := append([]*multipartRequestBody(nil), a.bodies...) + bodies := append([]io.ReadCloser(nil), a.bodies...) a.bodies = nil return bodies } @@ -424,33 +479,33 @@ func rejectUnreplayableRedirect(req *http.Request, next option.MiddlewareNext) ( } type multipartBodyInfo struct { - hasUpload bool - replayable bool + hasUpload bool + knownLength bool } func inspectMultipartBody(value any) multipartBodyInfo { switch value := value.(type) { case map[string]any: - info := multipartBodyInfo{replayable: true} + info := multipartBodyInfo{knownLength: true} for _, child := range value { childInfo := inspectMultipartBody(child) info.hasUpload = info.hasUpload || childInfo.hasUpload - info.replayable = info.replayable && childInfo.replayable + info.knownLength = info.knownLength && childInfo.knownLength } return info case []any: - info := multipartBodyInfo{replayable: true} + info := multipartBodyInfo{knownLength: true} for _, child := range value { childInfo := inspectMultipartBody(child) info.hasUpload = info.hasUpload || childInfo.hasUpload - info.replayable = info.replayable && childInfo.replayable + info.knownLength = info.knownLength && childInfo.knownLength } return info case fileUpload: - return multipartBodyInfo{hasUpload: true, replayable: value.canReplay()} + return multipartBodyInfo{hasUpload: true, knownLength: value.hasKnownSize()} default: _, isReader := value.(io.Reader) - return multipartBodyInfo{hasUpload: isReader, replayable: !isReader} + return multipartBodyInfo{hasUpload: isReader, knownLength: !isReader} } } @@ -461,15 +516,14 @@ func multipartContentLength( ) (int64, error) { var uploadBytes int64 emptyValue, err := transformFileUploads(bodyMap, func(upload fileUpload) (fileUpload, error) { - if !upload.canReplay() { - return fileUpload{}, errors.New("cannot determine multipart length for a non-replayable upload") + if !upload.hasKnownSize() { + return fileUpload{}, errors.New("cannot determine multipart length for an unknown-size upload") } if upload.size > math.MaxInt64-uploadBytes { return fileUpload{}, errors.New("multipart content length overflows int64") } uploadBytes += upload.size upload.Reader = strings.NewReader("") - upload.ownsSource = false return upload, nil }) if err != nil { @@ -493,20 +547,6 @@ func multipartContentLength( return int64(framing.Len()) + uploadBytes, nil } -func replayMultipartBody(bodyMap map[string]any) (map[string]any, error) { - value, err := transformFileUploads(bodyMap, func(upload fileUpload) (fileUpload, error) { - return upload.replay() - }) - if err != nil { - return nil, err - } - reopened, ok := value.(map[string]any) - if !ok { - return nil, errors.New("multipart body must be a map") - } - return reopened, nil -} - func transformFileUploads( value any, transform func(fileUpload) (fileUpload, error), diff --git a/pkg/cmd/multipartbody_test.go b/pkg/cmd/multipartbody_test.go index 610a7dc2..c4d858fb 100644 --- a/pkg/cmd/multipartbody_test.go +++ b/pkg/cmd/multipartbody_test.go @@ -133,7 +133,7 @@ func TestMultipartRequestBodyClosesOpenedFiles(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte("payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - file := upload.source.file + file := uploadSourceFile(t, upload) body := newMultipartRequestBody(map[string]any{"file": upload}, apiform.FormatBrackets) _, err = io.ReadAll(body) @@ -245,7 +245,7 @@ func TestMultipartRequestOptionsFollow307And308ForRegularFiles(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte("redirect payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - originalFile := upload.source.file + originalFile := uploadSourceFile(t, upload) var requestCount atomic.Int32 var uploaded string @@ -365,7 +365,7 @@ func TestMultipartRequestOptionsRetryRegularFiles(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - sourceFile := upload.source.file + sourceFile := uploadSourceFile(t, upload) var requestCount atomic.Int32 var uploadsMu sync.Mutex @@ -428,7 +428,7 @@ func TestMultipartRequestOptionsCloseRegularFilesAfterFinalRetry(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - sourceFile := upload.source.file + sourceFile := uploadSourceFile(t, upload) var requestCount atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -471,7 +471,7 @@ func TestMultipartRequestOptionsCloseRegularFilesWhenBackoffCanceled(t *testing. require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - sourceFile := upload.source.file + sourceFile := uploadSourceFile(t, upload) var requestCount atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -535,7 +535,7 @@ func TestMultipartRequestOptionsLetEarlySuccessfulUploadFinish(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("payload", 1024)), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - sourceFile := upload.source.file + sourceFile := uploadSourceFile(t, upload) doer := newEarlyResponseHTTPDoer() options, err := multipartRequestOptions(map[string]any{ @@ -569,7 +569,7 @@ func TestMultipartReplayAttemptAbortDoesNotWaitForBlockedReader(t *testing.T) { body := newMultipartRequestBody(map[string]any{ "file": fileUpload{Reader: source, filename: "blocked.bin"}, }, apiform.FormatBrackets) - attempt := &multipartReplayAttempt{bodies: []*multipartRequestBody{body}} + attempt := &multipartReplayAttempt{bodies: []io.ReadCloser{body}} readDone := make(chan error, 1) go func() { _, err := io.Copy(io.Discard, body) @@ -760,7 +760,7 @@ func TestOpenFileUploadStreamsProcfsWithUnknownLength(t *testing.T) { } require.NoError(t, err) require.False(t, upload.knownSize) - require.False(t, upload.canReplay()) + require.False(t, upload.hasKnownSize()) sourceFile := upload.Reader.(*os.File) type requestInfo struct { @@ -821,7 +821,7 @@ func TestOpenFileUploadTreatsEmptyRegularFileAsKnownLength(t *testing.T) { t.Cleanup(func() { require.NoError(t, upload.Close()) }) require.True(t, upload.knownSize) - require.True(t, upload.canReplay()) + require.True(t, upload.hasKnownSize()) require.Zero(t, upload.size) } @@ -846,6 +846,15 @@ func waitForMultipartBody(t *testing.T, body *multipartRequestBody) { } } +func uploadSourceFile(t *testing.T, upload fileUpload) *os.File { + t.Helper() + reader, ok := upload.Reader.(*exactLengthReadCloser) + require.True(t, ok, "regular upload must use an exact-length reader") + file, ok := reader.closer.(*os.File) + require.True(t, ok, "regular upload must own its opened file") + return file +} + type countingReader struct { remaining int64 bytesRead atomic.Int64 diff --git a/pkg/cmd/multipartfallback_test.go b/pkg/cmd/multipartfallback_test.go new file mode 100644 index 00000000..7ae2149d --- /dev/null +++ b/pkg/cmd/multipartfallback_test.go @@ -0,0 +1,196 @@ +package cmd + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + "github.com/urfave/cli/v3" +) + +func TestFilesCreateCLIStreamsWhenSnapshotDirectoryUnavailable(t *testing.T) { + for _, payload := range [][]byte{[]byte("streamed payload"), {}} { + name := "non_empty" + if len(payload) == 0 { + name = "zero_byte" + } + t.Run(name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "upload.bin") + require.NoError(t, os.WriteFile(path, payload, 0o600)) + setUnavailableTempDir(t) + + received := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + contents, err := readMultipartUpload(r) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + received <- contents + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + require.NoError(t, runFilesCreateCLI(t.Context(), server.URL+"/", path)) + require.Equal(t, payload, <-received) + }) + } +} + +func TestFilesCreateCLIDisablesReplayWhenSnapshotDirectoryUnavailable(t *testing.T) { + for _, status := range []int{http.StatusTemporaryRedirect, http.StatusTooManyRequests} { + t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { + path := filepath.Join(t.TempDir(), "upload.bin") + require.NoError(t, os.WriteFile(path, []byte("first attempt only"), 0o600)) + setUnavailableTempDir(t) + + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + contents, err := readMultipartUpload(r) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + require.Equal(t, []byte("first attempt only"), contents) + if status == http.StatusTemporaryRedirect { + http.Redirect(w, r, "/upload", status) + return + } + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "0") + w.WriteHeader(status) + _, _ = io.WriteString(w, `{"error":{"message":"retry unavailable","type":"rate_limit_error"}}`) + })) + t.Cleanup(server.Close) + + err := runFilesCreateCLI(t.Context(), server.URL+"/", path) + require.Error(t, err) + require.Equal(t, int32(1), requests.Load(), "snapshot failure must disable network replay") + }) + } +} + +func TestVideosCreateCLIPropagatesExtensionlessDirectoryError(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + require.NoError(t, os.Mkdir("artifactdir", 0o700)) + + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requests.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{}`) + })) + t.Cleanup(server.Close) + + err := runVideosCreateCLI(t.Context(), server.URL+"/", "@artifactdir") + require.ErrorContains(t, err, "is a directory") + require.Zero(t, requests.Load(), "an existing invalid upload must not fall back to a scalar") +} + +func TestVideosCreateCLIUploadsExtensionlessFileWhenSnapshotDirectoryUnavailable(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + require.NoError(t, os.WriteFile("artifact", []byte("extensionless payload"), 0o600)) + setUnavailableTempDir(t) + + received := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseMultipartForm(1 << 20); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + file, _, err := r.FormFile("input_reference") + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + defer file.Close() + contents, err := io.ReadAll(file) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + received <- contents + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{}`) + })) + t.Cleanup(server.Close) + + require.NoError(t, runVideosCreateCLI(t.Context(), server.URL+"/", "@artifact")) + require.Equal(t, []byte("extensionless payload"), <-received) +} + +func setUnavailableTempDir(t *testing.T) { + t.Helper() + missing := filepath.Join(t.TempDir(), "missing") + t.Setenv("TMPDIR", missing) + t.Setenv("TMP", missing) + t.Setenv("TEMP", missing) +} + +func runVideosCreateCLI(ctx context.Context, baseURL, inputReference string) error { + create := videosCreate + command := &cli.Command{ + Name: "openai", + Flags: []cli.Flag{ + &cli.BoolFlag{Name: "debug"}, + &cli.StringFlag{Name: "base-url"}, + &cli.StringFlag{Name: "format", Value: "json"}, + &cli.StringFlag{Name: "transform"}, + &cli.BoolFlag{Name: "raw-output"}, + &cli.StringFlag{Name: "api-key"}, + }, + Commands: []*cli.Command{{ + Name: "videos", + Commands: []*cli.Command{&create}, + }}, + } + return command.Run(ctx, []string{ + "openai", + "--api-key", "test-key", + "--base-url", baseURL, + "videos", "create", + "--prompt", "hello", + "--input-reference", inputReference, + }) +} + +func runImagesEditCLI(ctx context.Context, baseURL string, paths []string) error { + edit := imagesEdit + command := &cli.Command{ + Name: "openai", + Flags: []cli.Flag{ + &cli.BoolFlag{Name: "debug"}, + &cli.StringFlag{Name: "base-url"}, + &cli.StringFlag{Name: "format", Value: "json"}, + &cli.StringFlag{Name: "transform"}, + &cli.BoolFlag{Name: "raw-output"}, + &cli.StringFlag{Name: "api-key"}, + }, + Commands: []*cli.Command{{ + Name: "images", + Commands: []*cli.Command{&edit}, + }}, + } + args := []string{ + "openai", + "--api-key", "test-key", + "--base-url", baseURL, + "images", "edit", + "--prompt", "hello", + } + for _, path := range paths { + args = append(args, "--image", path) + } + return command.Run(ctx, args) +} diff --git a/pkg/cmd/multipartfile.go b/pkg/cmd/multipartfile.go index 70643329..bca30247 100644 --- a/pkg/cmd/multipartfile.go +++ b/pkg/cmd/multipartfile.go @@ -20,84 +20,19 @@ type fileUpload struct { contentType string size int64 knownSize bool - source *replayableFileSource - ownsSource bool } func (f fileUpload) Filename() string { return f.filename } func (f fileUpload) ContentType() string { return f.contentType } func (f fileUpload) Close() error { - if f.ownsSource && f.source != nil { - return f.source.Close() - } if closer, ok := f.Reader.(io.Closer); ok { return closer.Close() } return nil } -func (f fileUpload) canReplay() bool { - return f.source != nil && f.knownSize -} - -func (f fileUpload) replay() (fileUpload, error) { - if !f.canReplay() { - return fileUpload{}, errors.New("multipart upload is not replayable") - } - reader, err := f.source.ReplayReader() - if err != nil { - return fileUpload{}, err - } - f.Reader = reader - f.ownsSource = false - return f, nil -} - -// replayableFileSource owns an immutable, bounded on-disk snapshot. Each body -// reads from a duplicate handle, so redirects and retries always send the same -// bytes and closing the anchor cannot truncate a transport-owned upload. -type replayableFileSource struct { - file *os.File - size int64 - sourceName string - - closeOnce sync.Once - closeErr error -} - -func (s *replayableFileSource) ReplayReader() (io.ReadCloser, error) { - info, err := s.file.Stat() - if err != nil { - return nil, err - } - if !info.Mode().IsRegular() || info.Size() != s.size { - return nil, fmt.Errorf("upload snapshot for %q changed while being replayed", s.displayName()) - } - file, err := duplicateFile(s.file) - if err != nil { - return nil, fmt.Errorf("duplicate upload snapshot for %q: %w", s.displayName(), err) - } - return &exactLengthReadCloser{ - exactLengthReader: exactLengthReader{ - reader: io.NewSectionReader(file, 0, s.size), - remaining: s.size, - }, - closer: file, - }, nil -} - -func (s *replayableFileSource) displayName() string { - if s.sourceName != "" { - return s.sourceName - } - return s.file.Name() -} - -func (s *replayableFileSource) Close() error { - s.closeOnce.Do(func() { - s.closeErr = s.file.Close() - }) - return s.closeErr +func (f fileUpload) hasKnownSize() bool { + return f.knownSize } type exactLengthReader struct { @@ -125,11 +60,24 @@ func (r *exactLengthReader) Read(p []byte) (int, error) { type exactLengthReadCloser struct { exactLengthReader - closer io.Closer + closer io.Closer + closeOnce sync.Once + closeErr error +} + +func (r *exactLengthReadCloser) Read(p []byte) (int, error) { + n, err := r.exactLengthReader.Read(p) + if r.remaining == 0 || err != nil { + _ = r.Close() + } + return n, err } func (r *exactLengthReadCloser) Close() error { - return r.closer.Close() + r.closeOnce.Do(func() { + r.closeErr = r.closer.Close() + }) + return r.closeErr } func hasTrustworthyFileSize(file *os.File, info os.FileInfo) bool { @@ -154,10 +102,10 @@ func hasTrustworthyFileSize(file *os.File, info os.FileInfo) bool { return n == 0 && errors.Is(err, io.EOF) } -// openFileUpload opens path once. Ordinary files are copied into a secure, -// size-bounded temporary snapshot so multipart redirects and retries are -// replayable without retaining the caller's mutable file or buffering it in -// memory. Virtual and unknown-length files remain single-shot streams. +// openFileUpload opens path once. Ordinary files retain their exact stat size +// so the first request can stream directly with Content-Length while the +// multipart body opportunistically captures replay bytes. Virtual and +// unknown-length files remain single-shot streams. func openFileUpload(path string) (fileUpload, error) { file, err := os.Open(path) if err != nil { @@ -182,51 +130,17 @@ func openFileUpload(path string) (fileUpload, error) { }, nil } - source, err := snapshotReplayableFile(file, info.Size(), path) - if err != nil { - return fileUpload{}, err - } return fileUpload{ - Reader: &exactLengthReader{ - reader: io.NewSectionReader(source.file, 0, source.size), - remaining: source.size, + Reader: &exactLengthReadCloser{ + exactLengthReader: exactLengthReader{ + reader: io.NewSectionReader(file, 0, info.Size()), + remaining: info.Size(), + }, + closer: file, }, filename: filepath.Base(path), contentType: contentType, - size: source.size, + size: info.Size(), knownSize: true, - source: source, - ownsSource: true, - }, nil -} - -func snapshotReplayableFile(source *os.File, size int64, sourceName string) (*replayableFileSource, error) { - snapshot, err := createReplaySnapshotFile() - if err != nil { - return nil, errors.Join(fmt.Errorf("snapshot upload %q: %w", sourceName, err), source.Close()) - } - fail := func(cause error) (*replayableFileSource, error) { - return nil, errors.Join(cause, source.Close(), snapshot.Close()) - } - - exactSource := &exactLengthReader{reader: source, remaining: size} - if _, err := io.Copy(snapshot, exactSource); err != nil { - return fail(fmt.Errorf("snapshot upload %q: %w", sourceName, err)) - } - var extra [1]byte - n, readErr := source.Read(extra[:]) - if n != 0 || !errors.Is(readErr, io.EOF) { - if readErr == nil { - readErr = errors.New("source grew while being snapshotted") - } - return fail(fmt.Errorf("snapshot upload %q: %w", sourceName, readErr)) - } - if err := source.Close(); err != nil { - return nil, errors.Join(fmt.Errorf("close upload %q after snapshot: %w", sourceName, err), snapshot.Close()) - } - return &replayableFileSource{ - file: snapshot, - size: size, - sourceName: sourceName, }, nil } diff --git a/pkg/cmd/multipartfile_windows.go b/pkg/cmd/multipartfile_windows.go index f2c8a6a6..a891fe80 100644 --- a/pkg/cmd/multipartfile_windows.go +++ b/pkg/cmd/multipartfile_windows.go @@ -21,15 +21,23 @@ func createReplaySnapshotFile() (*os.File, error) { if err := os.Remove(path); err != nil { return nil, err } - file, err := os.OpenFile( - path, - os.O_RDWR|os.O_CREATE|os.O_EXCL|windows.O_FILE_FLAG_DELETE_ON_CLOSE, - 0o600, + pathPtr, err := windows.UTF16PtrFromString(path) + if err != nil { + return nil, errors.Join(err, os.Remove(path)) + } + handle, err := windows.CreateFile( + pathPtr, + windows.GENERIC_READ|windows.GENERIC_WRITE, + windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE|windows.FILE_SHARE_DELETE, + nil, + windows.CREATE_NEW, + windows.FILE_ATTRIBUTE_TEMPORARY|windows.FILE_FLAG_DELETE_ON_CLOSE, + 0, ) if err != nil { return nil, errors.Join(err, os.Remove(path)) } - return file, nil + return os.NewFile(uintptr(handle), path), nil } func duplicateFile(file *os.File) (*os.File, error) { diff --git a/pkg/cmd/multipartfile_windows_test.go b/pkg/cmd/multipartfile_windows_test.go new file mode 100644 index 00000000..71a503af --- /dev/null +++ b/pkg/cmd/multipartfile_windows_test.go @@ -0,0 +1,26 @@ +//go:build windows + +package cmd + +import ( + "os" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestReplaySnapshotIsDeletedAfterFinalHandleCloses(t *testing.T) { + snapshot, err := createReplaySnapshotFile() + require.NoError(t, err) + path := snapshot.Name() + _, err = snapshot.Write([]byte("sensitive upload")) + require.NoError(t, err) + + duplicate, err := duplicateFile(snapshot) + require.NoError(t, err) + require.NoError(t, snapshot.Close()) + require.NoError(t, duplicate.Close()) + + _, err = os.Stat(path) + require.ErrorIs(t, err, os.ErrNotExist) +} diff --git a/pkg/cmd/multipartintegrity_test.go b/pkg/cmd/multipartintegrity_test.go index 480ea39d..497c5c4d 100644 --- a/pkg/cmd/multipartintegrity_test.go +++ b/pkg/cmd/multipartintegrity_test.go @@ -74,7 +74,7 @@ func TestFilesCreateCLIUsesImmutableReplayContent(t *testing.T) { } } -func TestFilesCreateCLISnapshotsBeforeSourceShrinks(t *testing.T) { +func TestFilesCreateCLIAbortsWhenSourceShrinksDuringFirstSend(t *testing.T) { const uploadSize = 16 << 20 path := filepath.Join(t.TempDir(), "shrinking.bin") require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("a", uploadSize)), 0o600)) @@ -119,11 +119,12 @@ func TestFilesCreateCLISnapshotsBeforeSourceShrinks(t *testing.T) { })) t.Cleanup(server.Close) - require.NoError(t, runFilesCreateCLI(t.Context(), server.URL+"/", path)) + err := runFilesCreateCLI(t.Context(), server.URL+"/", path) + require.ErrorIs(t, err, io.ErrUnexpectedEOF) result := <-resultCh - require.NoError(t, result.err) - require.True(t, result.committed) - require.Equal(t, int64(uploadSize), result.bytes) + require.Error(t, result.err) + require.False(t, result.committed, "a truncated first stream must not emit a terminal multipart boundary") + require.Less(t, result.bytes, int64(uploadSize)) } func TestMultipartReplayPrematureEOFDoesNotCommitFinalBoundary(t *testing.T) { @@ -132,15 +133,18 @@ func TestMultipartReplayPrematureEOFDoesNotCommitFinalBoundary(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("a", uploadSize)), 0o600)) file, err := os.Open(path) require.NoError(t, err) - source := &replayableFileSource{file: file, size: uploadSize} upload := fileUpload{ - Reader: io.NewSectionReader(file, 0, uploadSize), + Reader: &exactLengthReadCloser{ + exactLengthReader: exactLengthReader{ + reader: io.NewSectionReader(file, 0, uploadSize), + remaining: uploadSize, + }, + closer: file, + }, filename: filepath.Base(path), contentType: "application/octet-stream", size: uploadSize, knownSize: true, - source: source, - ownsSource: true, } type receiverResult struct { @@ -194,25 +198,31 @@ func TestMultipartReplayPrematureEOFDoesNotCommitFinalBoundary(t *testing.T) { require.Error(t, result.err) } -func TestOpenFileUploadRemovesSnapshotOnClose(t *testing.T) { +func TestOpenFileUploadDoesNotRequireSnapshotBeforeRead(t *testing.T) { path := filepath.Join(t.TempDir(), "upload.txt") - require.NoError(t, os.WriteFile(path, []byte("snapshot payload"), 0o600)) + require.NoError(t, os.WriteFile(path, []byte("streamed payload"), 0o600)) + setUnavailableTempDir(t) upload, err := openFileUpload(path) require.NoError(t, err) - snapshotPath := upload.source.file.Name() - + contents, err := io.ReadAll(upload.Reader) + require.NoError(t, err) + require.Equal(t, "streamed payload", string(contents)) require.NoError(t, upload.Close()) - _, err = os.Stat(snapshotPath) - require.ErrorIs(t, err, os.ErrNotExist) } -func TestSnapshotReplayableFileRejectsPrematureEOF(t *testing.T) { +func TestExactLengthUploadRejectsPrematureEOF(t *testing.T) { path := filepath.Join(t.TempDir(), "short.txt") require.NoError(t, os.WriteFile(path, []byte("short"), 0o600)) file, err := os.Open(path) require.NoError(t, err) - - _, err = snapshotReplayableFile(file, int64(len("longer than short")), path) + reader := &exactLengthReadCloser{ + exactLengthReader: exactLengthReader{ + reader: file, + remaining: int64(len("longer than short")), + }, + closer: file, + } + _, err = io.ReadAll(reader) require.ErrorIs(t, err, io.ErrUnexpectedEOF) } diff --git a/pkg/cmd/multipartreplay.go b/pkg/cmd/multipartreplay.go new file mode 100644 index 00000000..5ad917a5 --- /dev/null +++ b/pkg/cmd/multipartreplay.go @@ -0,0 +1,160 @@ +package cmd + +import ( + "errors" + "fmt" + "io" + "os" + "sync" +) + +// multipartReplaySource opportunistically captures the exact encoded bytes of +// the first request body. Capture failures never fail that first stream; they +// only make redirects and retries unavailable. +type multipartReplaySource struct { + mu sync.Mutex + expected int64 + captured int64 + file *os.File + ready bool + closed bool + err error + closeErr error +} + +func newMultipartReplaySource(expected int64) *multipartReplaySource { + return &multipartReplaySource{expected: expected} +} + +func (s *multipartReplaySource) Capture(p []byte) { + if len(p) == 0 { + return + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.ready || s.closed || s.err != nil { + return + } + if s.file == nil { + file, err := createReplaySnapshotFile() + if err != nil { + s.failLocked(fmt.Errorf("create multipart replay snapshot: %w", err)) + return + } + s.file = file + } + + n, err := s.file.Write(p) + s.captured += int64(n) + if err == nil && n != len(p) { + err = io.ErrShortWrite + } + if err != nil { + s.failLocked(fmt.Errorf("write multipart replay snapshot: %w", err)) + return + } + if s.captured > s.expected { + s.failLocked(fmt.Errorf("multipart replay snapshot exceeded expected length %d", s.expected)) + return + } + if s.captured == s.expected { + s.ready = true + } +} + +func (s *multipartReplaySource) Fail(err error) { + if err == nil { + err = errors.New("multipart replay capture did not complete") + } + s.mu.Lock() + defer s.mu.Unlock() + if !s.ready { + s.failLocked(err) + } +} + +func (s *multipartReplaySource) failLocked(err error) { + if s.err == nil { + s.err = err + } + if s.file != nil { + s.closeErr = errors.Join(s.closeErr, s.file.Close()) + s.file = nil + } +} + +func (s *multipartReplaySource) Ready() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.ready && !s.closed +} + +func (s *multipartReplaySource) ReplayReader() (io.ReadCloser, error) { + s.mu.Lock() + if !s.ready || s.closed || s.file == nil { + err := s.err + if err == nil { + err = errors.New("first multipart stream did not complete") + } + s.mu.Unlock() + return nil, fmt.Errorf("multipart upload is not replayable: %w", err) + } + file, err := duplicateFile(s.file) + s.mu.Unlock() + if err != nil { + return nil, fmt.Errorf("duplicate multipart replay snapshot: %w", err) + } + return &exactLengthReadCloser{ + exactLengthReader: exactLengthReader{ + reader: io.NewSectionReader(file, 0, s.expected), + remaining: s.expected, + }, + closer: file, + }, nil +} + +func (s *multipartReplaySource) Close() error { + s.mu.Lock() + if s.closed { + err := s.closeErr + s.mu.Unlock() + return err + } + s.closed = true + file := s.file + s.file = nil + if file != nil { + s.closeErr = errors.Join(s.closeErr, file.Close()) + } + err := s.closeErr + s.mu.Unlock() + return err +} + +type capturingMultipartBody struct { + body *multipartRequestBody + source *multipartReplaySource +} + +func (b *capturingMultipartBody) Read(p []byte) (int, error) { + n, err := b.body.Read(p) + if n > 0 { + b.source.Capture(p[:n]) + } + if err != nil && !errors.Is(err, io.EOF) { + b.source.Fail(err) + } + return n, err +} + +func (b *capturingMultipartBody) Close() error { + if !b.source.Ready() { + b.source.Fail(errors.New("first multipart stream closed before replay capture completed")) + } + return b.body.Close() +} + +func (b *capturingMultipartBody) Done() <-chan struct{} { + return b.body.done +} diff --git a/pkg/cmd/multipartresource_unix_test.go b/pkg/cmd/multipartresource_unix_test.go new file mode 100644 index 00000000..2e010c61 --- /dev/null +++ b/pkg/cmd/multipartresource_unix_test.go @@ -0,0 +1,117 @@ +//go:build !windows + +package cmd + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/sys/unix" +) + +const multipartResourceHelperEnv = "OPENAI_MULTIPART_RESOURCE_HELPER" + +func TestFilesCreateCLIStreamsSparseFileWhenSnapshotQuotaExceeded(t *testing.T) { + const uploadSize = 256 << 10 + path := filepath.Join(t.TempDir(), "sparse.bin") + file, err := os.Create(path) + require.NoError(t, err) + require.NoError(t, file.Truncate(uploadSize)) + require.NoError(t, file.Close()) + + received := make(chan int64, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + contents, err := readMultipartUpload(r) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + received <- int64(len(contents)) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + + output, err := runMultipartResourceHelper(t, map[string]string{ + multipartResourceHelperEnv: "file-size", + "OPENAI_HELPER_BASE_URL": server.URL + "/", + "OPENAI_HELPER_PATH": path, + "TMPDIR": t.TempDir(), + }) + require.NoError(t, err, output) + require.Equal(t, int64(uploadSize), <-received) +} + +func TestImagesEditCLIDialsBeforeOpeningReplayDescriptors(t *testing.T) { + paths := make([]string, 16) + for i := range paths { + paths[i] = filepath.Join(t.TempDir(), "image.png") + require.NoError(t, os.WriteFile(paths[i], []byte{byte(i)}, 0o600)) + } + encodedPaths, err := json.Marshal(paths) + require.NoError(t, err) + + received := make(chan int, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseMultipartForm(1 << 20); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + var fileCount int + for _, files := range r.MultipartForm.File { + fileCount += len(files) + } + received <- fileCount + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{}`) + })) + t.Cleanup(server.Close) + + output, err := runMultipartResourceHelper(t, map[string]string{ + multipartResourceHelperEnv: "file-descriptors", + "OPENAI_HELPER_BASE_URL": server.URL + "/", + "OPENAI_HELPER_PATHS": string(encodedPaths), + }) + require.NoError(t, err, output) + require.Equal(t, 16, <-received) +} + +func TestMultipartResourceHelper(t *testing.T) { + mode := os.Getenv(multipartResourceHelperEnv) + if mode == "" { + t.Skip("helper process") + } + + switch mode { + case "file-size": + require.NoError(t, unix.Setrlimit(unix.RLIMIT_FSIZE, &unix.Rlimit{Cur: 64 << 10, Max: 64 << 10})) + require.NoError(t, runFilesCreateCLI(t.Context(), os.Getenv("OPENAI_HELPER_BASE_URL"), os.Getenv("OPENAI_HELPER_PATH"))) + case "file-descriptors": + require.NoError(t, unix.Setrlimit(unix.RLIMIT_NOFILE, &unix.Rlimit{Cur: 23, Max: 23})) + var paths []string + require.NoError(t, json.Unmarshal([]byte(os.Getenv("OPENAI_HELPER_PATHS")), &paths)) + require.NoError(t, runImagesEditCLI(t.Context(), os.Getenv("OPENAI_HELPER_BASE_URL"), paths)) + default: + t.Fatalf("unknown helper mode %q", mode) + } +} + +func runMultipartResourceHelper(t *testing.T, env map[string]string) (string, error) { + t.Helper() + executable, err := os.Executable() + require.NoError(t, err) + cmd := exec.Command(executable, "-test.run=^TestMultipartResourceHelper$") + cmd.Env = os.Environ() + for key, value := range env { + cmd.Env = append(cmd.Env, key+"="+value) + } + output, err := cmd.CombinedOutput() + return string(output), err +} From 7e4a444fad016106c162dd4e9f794cb4dd16caa1 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 18 Aug 2026 01:16:11 +0000 Subject: [PATCH 7/9] fix: protect Windows replay snapshots --- pkg/cmd/multipartfile_windows.go | 31 ++++++++++++++++++++-- pkg/cmd/multipartfile_windows_test.go | 38 +++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/pkg/cmd/multipartfile_windows.go b/pkg/cmd/multipartfile_windows.go index a891fe80..f663d528 100644 --- a/pkg/cmd/multipartfile_windows.go +++ b/pkg/cmd/multipartfile_windows.go @@ -4,7 +4,9 @@ package cmd import ( "errors" + "fmt" "os" + "unsafe" "golang.org/x/sys/windows" ) @@ -25,11 +27,15 @@ func createReplaySnapshotFile() (*os.File, error) { if err != nil { return nil, errors.Join(err, os.Remove(path)) } + securityAttributes, err := replaySnapshotSecurityAttributes() + if err != nil { + return nil, errors.Join(err, os.Remove(path)) + } handle, err := windows.CreateFile( pathPtr, windows.GENERIC_READ|windows.GENERIC_WRITE, - windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE|windows.FILE_SHARE_DELETE, - nil, + windows.FILE_SHARE_DELETE, + securityAttributes, windows.CREATE_NEW, windows.FILE_ATTRIBUTE_TEMPORARY|windows.FILE_FLAG_DELETE_ON_CLOSE, 0, @@ -40,6 +46,27 @@ func createReplaySnapshotFile() (*os.File, error) { return os.NewFile(uintptr(handle), path), nil } +func replaySnapshotSecurityAttributes() (*windows.SecurityAttributes, error) { + user, err := windows.GetCurrentProcessToken().GetTokenUser() + if err != nil { + return nil, fmt.Errorf("get current user for multipart replay snapshot: %w", err) + } + userSID := user.User.Sid.String() + if userSID == "" { + return nil, errors.New("get current user SID for multipart replay snapshot") + } + securityDescriptor, err := windows.SecurityDescriptorFromString( + "D:P(A;;GA;;;" + userSID + ")", + ) + if err != nil { + return nil, fmt.Errorf("create multipart replay snapshot security descriptor: %w", err) + } + return &windows.SecurityAttributes{ + Length: uint32(unsafe.Sizeof(windows.SecurityAttributes{})), + SecurityDescriptor: securityDescriptor, + }, nil +} + func duplicateFile(file *os.File) (*os.File, error) { process, err := windows.GetCurrentProcess() if err != nil { diff --git a/pkg/cmd/multipartfile_windows_test.go b/pkg/cmd/multipartfile_windows_test.go index 71a503af..4feee1eb 100644 --- a/pkg/cmd/multipartfile_windows_test.go +++ b/pkg/cmd/multipartfile_windows_test.go @@ -5,10 +5,48 @@ package cmd import ( "os" "testing" + "unsafe" "github.com/stretchr/testify/require" + "golang.org/x/sys/windows" ) +func TestReplaySnapshotHasProtectedCurrentUserDACL(t *testing.T) { + snapshot, err := createReplaySnapshotFile() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, snapshot.Close()) }) + + securityDescriptor, err := windows.GetSecurityInfo( + windows.Handle(snapshot.Fd()), + windows.SE_FILE_OBJECT, + windows.DACL_SECURITY_INFORMATION, + ) + require.NoError(t, err) + control, _, err := securityDescriptor.Control() + require.NoError(t, err) + require.NotZero(t, control&windows.SE_DACL_PROTECTED) + + dacl, _, err := securityDescriptor.DACL() + require.NoError(t, err) + require.NotNil(t, dacl) + require.Equal(t, uint16(1), dacl.AceCount) + var ace *windows.ACCESS_ALLOWED_ACE + require.NoError(t, windows.GetAce(dacl, 0, &ace)) + require.Equal(t, uint8(windows.ACCESS_ALLOWED_ACE_TYPE), ace.Header.AceType) + require.Zero(t, ace.Header.AceFlags&windows.INHERITED_ACE) + + currentUser, err := windows.GetCurrentProcessToken().GetTokenUser() + require.NoError(t, err) + aceSID := (*windows.SID)(unsafe.Pointer(&ace.SidStart)) + require.True(t, windows.EqualSid(currentUser.User.Sid, aceSID)) + + otherHandle, err := os.Open(snapshot.Name()) + if otherHandle != nil { + require.NoError(t, otherHandle.Close()) + } + require.Error(t, err) +} + func TestReplaySnapshotIsDeletedAfterFinalHandleCloses(t *testing.T) { snapshot, err := createReplaySnapshotFile() require.NoError(t, err) From 27f98fe3b2092916e4117bc7a4c9f47baf6b3d5e Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 18 Aug 2026 02:39:46 +0000 Subject: [PATCH 8/9] fix: make multipart uploads single-attempt streams --- go.mod | 2 +- go.sum | 4 +- pkg/cmd/multipartbody.go | 342 ++---------------- pkg/cmd/multipartbody_test.go | 481 +++---------------------- pkg/cmd/multipartfallback_test.go | 8 +- pkg/cmd/multipartfile.go | 5 +- pkg/cmd/multipartfile_unix.go | 30 -- pkg/cmd/multipartfile_windows.go | 88 ----- pkg/cmd/multipartfile_windows_test.go | 64 ---- pkg/cmd/multipartintegrity_test.go | 62 +--- pkg/cmd/multipartreplay.go | 160 -------- pkg/cmd/multipartresource_unix_test.go | 4 +- 12 files changed, 100 insertions(+), 1150 deletions(-) delete mode 100644 pkg/cmd/multipartfile_unix.go delete mode 100644 pkg/cmd/multipartfile_windows.go delete mode 100644 pkg/cmd/multipartfile_windows_test.go delete mode 100644 pkg/cmd/multipartreplay.go diff --git a/go.mod b/go.mod index 18b005ca..c7486542 100644 --- a/go.mod +++ b/go.mod @@ -10,7 +10,7 @@ require ( github.com/goccy/go-yaml v1.18.0 github.com/itchyny/json2yaml v0.1.4 github.com/muesli/reflow v0.3.0 - github.com/openai/openai-go/v3 v3.51.0 + github.com/openai/openai-go/v3 v3.52.0 github.com/stretchr/testify v1.10.0 github.com/tidwall/gjson v1.19.0 github.com/tidwall/pretty v1.2.1 diff --git a/go.sum b/go.sum index 4159d8fd..2fd4f728 100644 --- a/go.sum +++ b/go.sum @@ -45,8 +45,8 @@ github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s= github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8= github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= -github.com/openai/openai-go/v3 v3.51.0 h1:+ys88LqUflSr0nRM37aWxkMMpHn+zqzVJGIK89eumdM= -github.com/openai/openai-go/v3 v3.51.0/go.mod h1:Vy3y2/I2H/MbqvJGXEK8VbN5+avZV6zxux4I3eBdvaA= +github.com/openai/openai-go/v3 v3.52.0 h1:VDSjIvI5Sr2/AzGJI6219sM2Il+zBWuopvluMy6KdjE= +github.com/openai/openai-go/v3 v3.52.0/go.mod h1:Vy3y2/I2H/MbqvJGXEK8VbN5+avZV6zxux4I3eBdvaA= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= diff --git a/pkg/cmd/multipartbody.go b/pkg/cmd/multipartbody.go index 6e6d3199..07130b71 100644 --- a/pkg/cmd/multipartbody.go +++ b/pkg/cmd/multipartbody.go @@ -2,7 +2,6 @@ package cmd import ( "bytes" - "context" "errors" "fmt" "io" @@ -98,10 +97,6 @@ func (b *multipartRequestBody) ContentType() string { return b.contentType } -func (b *multipartRequestBody) Boundary() string { - return b.multipartWriter.Boundary() -} - func (b *multipartRequestBody) start() { b.startOnce.Do(func() { go b.encode() @@ -112,9 +107,8 @@ func (b *multipartRequestBody) encode() { defer close(b.done) if err := apiform.MarshalWithSettings(b.bodyMap, b.multipartWriter, b.encodingFormat); err != nil { - // A final multipart boundary asserts that every part completed. Abort the - // pipe immediately on a source error instead of making a truncated file - // look like a valid zero-byte upload to the receiver. + // A final boundary asserts that every part completed. Abort immediately + // instead of making a truncated source look like a complete upload. _ = b.writer.CloseWithError(errors.Join(err, b.cleanup())) return } @@ -131,12 +125,11 @@ func (b *multipartRequestBody) cleanup() error { return b.cleanupErr } -// multipartRequestOptions preserves the SDK's buffered, replayable behavior -// when a multipart form contains only scalar fields. Uploads are streamed. A -// form containing only regular files carries an exact Content-Length and -// opportunistically captures the first encoded stream for redirects and -// retries. Forms containing stdin, pipes, or arbitrary readers use chunked -// transfer and fail explicitly if a 307/308 requires replay. +// multipartRequestOptions keeps scalar-only forms buffered and replayable. +// Forms containing an upload stream exactly once: the SDK retry loop is +// disabled because the CLI cannot promise that a reader will produce identical +// bytes a second time. Regular files retain an exact Content-Length; unknown +// sources use chunked transfer encoding. func multipartRequestOptions( bodyMap map[string]any, encodingFormat apiform.FormFormat, @@ -147,303 +140,47 @@ func multipartRequestOptions( } boundary := multipart.NewWriter(io.Discard).Boundary() - var contentLength int64 + contentLength := int64(-1) if info.knownLength { var err error contentLength, err = multipartContentLength(bodyMap, encodingFormat, boundary) if err != nil { - return nil, err + return nil, errors.Join(err, closeFileUploads(bodyMap)) } } - if info.knownLength { - return replayableMultipartRequestOptions(bodyMap, encodingFormat, boundary, contentLength) - } - body, err := newMultipartRequestBodyWithBoundary(bodyMap, encodingFormat, boundary) if err != nil { - return nil, err + return nil, errors.Join(err, closeFileUploads(bodyMap)) } return []option.RequestOption{ option.WithRequestBody(body.ContentType(), body), option.WithMaxRetries(0), - option.WithMiddleware(rejectUnreplayableRedirect), - }, nil -} - -// openai-go v3.51 only installs GetBody for its concrete buffered body types. -// A streamed body therefore needs a replay marker on the SDK-owned request so -// its native retry loop runs. The middleware replaces that empty marker with a -// real multipart stream on every SDK attempt and supplies the same factory to -// net/http for 307/308 redirects. The CLI does not expose retry configuration, -// so the request fixes the SDK's current default of two retries; that bound also -// lets the middleware deterministically close its stable file descriptors. -const replayableMultipartMaxRetries = 2 - -type replayableMultipartRequest struct { - bodyMap map[string]any - encodingFormat apiform.FormFormat - boundary string - contentLength int64 - source *multipartReplaySource - - mu sync.Mutex - attempts int - closed bool - closeErr error - - pendingCancel func() bool -} - -func replayableMultipartRequestOptions( - bodyMap map[string]any, - encodingFormat apiform.FormFormat, - boundary string, - contentLength int64, -) ([]option.RequestOption, error) { - writer := multipart.NewWriter(io.Discard) - if err := writer.SetBoundary(boundary); err != nil { - return nil, err - } - request := &replayableMultipartRequest{ - bodyMap: bodyMap, - encodingFormat: encodingFormat, - boundary: boundary, - contentLength: contentLength, - source: newMultipartReplaySource(contentLength), - } - - // The empty *bytes.Reader is deliberate: it makes openai-go install GetBody - // on its canonical request before retry decisions. middleware replaces the - // marker on every attempt, so no marker bytes reach the network. - return []option.RequestOption{ - option.WithRequestBody(writer.FormDataContentType(), bytes.NewReader(nil)), - option.WithMaxRetries(replayableMultipartMaxRetries), - option.WithMiddleware(request.middleware), + option.WithMiddleware(streamedMultipartMiddleware(contentLength)), }, nil } -func (r *replayableMultipartRequest) middleware( - req *http.Request, - next option.MiddlewareNext, -) (*http.Response, error) { - attemptIndex := r.beginAttempt() - attempt := &multipartReplayAttempt{request: r, firstAttempt: attemptIndex == 0} - stopCancellation := context.AfterFunc(req.Context(), func() { - _ = attempt.abort() - _ = r.close() - }) - - // Discard the SDK's empty replay marker before installing the actual body. - if req.Body != nil { - _ = req.Body.Close() - } - body, err := attempt.newBody() - if err != nil { - stopCancellation() - return nil, errors.Join(err, r.close()) - } - req.Body = body - req.ContentLength = r.contentLength - req.GetBody = attempt.newBody - - res, requestErr := next(req) - requestCanceled := req.Context().Err() != nil - replayAvailable := r.source.Ready() - if res != nil && shouldRetryReplayableMultipart(res) && !replayAvailable { - // The first upload is still valid when the optional replay snapshot - // cannot be created. Tell openai-go not to issue another network - // attempt that this request cannot reproduce faithfully. - res.Header.Set("x-should-retry", "false") - } - terminal := requestCanceled || !shouldRetryReplayableMultipart(res) || - attemptIndex >= replayableMultipartMaxRetries || !replayAvailable - if terminal { - stopCancellation() - var attemptErr error - if requestErr != nil || requestCanceled || res == nil || - res.StatusCode >= http.StatusBadRequest { - attemptErr = attempt.abort() - } else { - attemptErr = attempt.release() - return res, errors.Join(requestErr, attemptErr) - } - return res, errors.Join(requestErr, attemptErr, r.close()) - } - - // A retry must not race an earlier body. Abort its duplicate handles without - // waiting for an uninterruptible read, while retaining the stable request - // anchors needed to construct the next SDK attempt. - attemptErr := attempt.abort() - r.armPendingCancellation(stopCancellation) - return res, errors.Join(requestErr, attemptErr) -} - -func (r *replayableMultipartRequest) beginAttempt() int { - r.mu.Lock() - stopCancellation := r.pendingCancel - r.pendingCancel = nil - attempt := r.attempts - r.attempts++ - r.mu.Unlock() - if stopCancellation != nil { - stopCancellation() - } - return attempt -} - -func (r *replayableMultipartRequest) armPendingCancellation(stopCancellation func() bool) { - r.mu.Lock() - if r.closed { - r.mu.Unlock() - stopCancellation() - return - } - previous := r.pendingCancel - r.pendingCancel = stopCancellation - r.mu.Unlock() - if previous != nil { - previous() - } -} - -func (r *replayableMultipartRequest) newFirstBody() (io.ReadCloser, error) { - body, err := newMultipartRequestBodyWithBoundary(r.bodyMap, r.encodingFormat, r.boundary) - if err != nil { - return nil, err - } - return &capturingMultipartBody{body: body, source: r.source}, nil -} - -func (r *replayableMultipartRequest) close() error { - r.mu.Lock() - stopCancellation := r.pendingCancel - r.pendingCancel = nil - if !r.closed { - r.closed = true - r.closeErr = errors.Join(closeFileUploads(r.bodyMap), r.source.Close()) - } - closeErr := r.closeErr - r.mu.Unlock() - if stopCancellation != nil { - stopCancellation() - } - return closeErr -} - -type multipartReplayAttempt struct { - request *replayableMultipartRequest - firstAttempt bool - - mu sync.Mutex - bodies []io.ReadCloser - bodyCount int - closed bool -} - -func (a *multipartReplayAttempt) newBody() (io.ReadCloser, error) { - a.mu.Lock() - if a.closed { - a.mu.Unlock() - return nil, errors.New("multipart replay attempt is closed") - } - bodyIndex := a.bodyCount - a.bodyCount++ - a.mu.Unlock() - - var body io.ReadCloser - var err error - if a.firstAttempt && bodyIndex == 0 { - body, err = a.request.newFirstBody() - } else { - body, err = a.request.source.ReplayReader() - } - if err != nil { - return nil, err - } - - a.mu.Lock() - defer a.mu.Unlock() - if a.closed { - _ = body.Close() - return nil, errors.New("multipart replay attempt is closed") - } - a.bodies = append(a.bodies, body) - return body, nil -} - -func (a *multipartReplayAttempt) release() error { - bodies := a.seal() - var waiters []<-chan struct{} - for _, body := range bodies { - if waiter, ok := body.(interface{ Done() <-chan struct{} }); ok { - waiters = append(waiters, waiter.Done()) - } - } - if len(waiters) == 0 { - return a.request.close() - } - allDone := true - for _, done := range waiters { - select { - case <-done: - default: - allDone = false +func streamedMultipartMiddleware(contentLength int64) option.Middleware { + return func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { + if contentLength >= 0 { + req.ContentLength = contentLength } - } - if allDone { - return a.request.close() - } - go func() { - for _, done := range waiters { - <-done + res, err := next(req) + if err != nil || res == nil || + (res.StatusCode != http.StatusTemporaryRedirect && res.StatusCode != http.StatusPermanentRedirect) { + return res, err } - _ = a.request.close() - }() - return nil -} -func (a *multipartReplayAttempt) abort() error { - bodies := a.seal() - var errs []error - for _, body := range bodies { - errs = append(errs, body.Close()) - } - return errors.Join(errs...) -} - -// seal prevents net/http from creating another redirect body after an attempt -// returns. The caller either leaves the returned transport-owned bodies to -// finish naturally or closes them to abort an errored/canceled attempt. -func (a *multipartReplayAttempt) seal() []io.ReadCloser { - a.mu.Lock() - defer a.mu.Unlock() - if a.closed { - return nil - } - a.closed = true - bodies := append([]io.ReadCloser(nil), a.bodies...) - a.bodies = nil - return bodies -} - -// Keep this predicate aligned with openai-go's retry policy. It is only used -// to decide when the request-level file descriptors can be closed; the SDK -// still owns retry timing, response cleanup, and the retry decision itself. -func shouldRetryReplayableMultipart(res *http.Response) bool { - if res == nil { - return true - } - if res.Header.Get("x-should-retry") == "true" { - return true - } - if res.Header.Get("x-should-retry") == "false" { - return false + location := res.Header.Get("Location") + if res.Body != nil { + _ = res.Body.Close() + } + return nil, fmt.Errorf( + "cannot follow HTTP %d redirect to %q: streamed multipart uploads are not replayable", + res.StatusCode, + location, + ) } - return res.StatusCode == http.StatusRequestTimeout || - res.StatusCode == http.StatusConflict || - res.StatusCode == http.StatusTooManyRequests || - res.StatusCode >= http.StatusInternalServerError } func bufferedMultipartRequestOptions( @@ -461,23 +198,6 @@ func bufferedMultipartRequestOptions( return []option.RequestOption{option.WithRequestBody(writer.FormDataContentType(), buf.Bytes())}, nil } -func rejectUnreplayableRedirect(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { - res, err := next(req) - if err != nil || res == nil || (res.StatusCode != http.StatusTemporaryRedirect && res.StatusCode != http.StatusPermanentRedirect) { - return res, err - } - - location := res.Header.Get("Location") - if res.Body != nil { - _ = res.Body.Close() - } - return res, fmt.Errorf( - "cannot follow HTTP %d redirect to %q: multipart upload contains a non-replayable source", - res.StatusCode, - location, - ) -} - type multipartBodyInfo struct { hasUpload bool knownLength bool @@ -559,7 +279,7 @@ func transformFileUploads( for key, child := range value { transformed, err := transformFileUploads(child, transform) if err != nil { - return nil, errors.Join(err, closeFileUploads(result)) + return nil, err } result[key] = transformed } @@ -569,14 +289,14 @@ func transformFileUploads( for _, child := range value { transformed, err := transformFileUploads(child, transform) if err != nil { - return nil, errors.Join(err, closeFileUploads(result)) + return nil, err } result = append(result, transformed) } return result, nil default: if _, isReader := value.(io.Reader); isReader { - return nil, errors.New("multipart body contains a non-replayable reader") + return nil, errors.New("multipart body contains an unknown-size reader") } return value, nil } diff --git a/pkg/cmd/multipartbody_test.go b/pkg/cmd/multipartbody_test.go index c4d858fb..942258ea 100644 --- a/pkg/cmd/multipartbody_test.go +++ b/pkg/cmd/multipartbody_test.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "math" "mime" "mime/multipart" "net/http" @@ -71,10 +72,12 @@ func TestMultipartRequestBodyPropagatesProducerErrors(t *testing.T) { body := newMultipartRequestBody(map[string]any{ "file": fileUpload{Reader: upload, filename: "broken.bin"}, }, apiform.FormatBrackets) + _, params, err := mime.ParseMediaType(body.ContentType()) + require.NoError(t, err) data, err := io.ReadAll(body) require.ErrorIs(t, err, readErr) - require.NotContains(t, string(data), "--"+body.Boundary()+"--", + require.NotContains(t, string(data), "--"+params["boundary"]+"--", "a failed file part must not be followed by a successful terminal boundary") waitForMultipartBody(t, body) require.Equal(t, int32(1), upload.closeCount.Load()) @@ -84,10 +87,7 @@ func TestMultipartRequestBodyPropagatesProducerErrors(t *testing.T) { func TestMultipartRequestBodyJoinsCleanupErrorAfterSourceError(t *testing.T) { readErr := errors.New("upload read failed") closeErr := errors.New("upload close failed") - upload := &recordingReadCloser{ - reader: errorReader{err: readErr}, - closeErr: closeErr, - } + upload := &recordingReadCloser{reader: errorReader{err: readErr}, closeErr: closeErr} body := newMultipartRequestBody(map[string]any{ "file": fileUpload{Reader: upload, filename: "broken.bin"}, }, apiform.FormatBrackets) @@ -196,7 +196,6 @@ func TestMultipartRequestBodyCloseCancelsEncoding(t *testing.T) { } require.NoError(t, body.Close()) - select { case err := <-readDone: require.Error(t, err) @@ -238,39 +237,18 @@ func TestMultipartRequestBodyDoesNotCloseUnownedReaders(t *testing.T) { require.Zero(t, stdinLikeReader.closeCount.Load()) } -func TestMultipartRequestOptionsFollow307And308ForRegularFiles(t *testing.T) { +func TestMultipartRequestOptionsReject307And308ForUploads(t *testing.T) { for _, status := range []int{http.StatusTemporaryRedirect, http.StatusPermanentRedirect} { t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { path := filepath.Join(t.TempDir(), "redirect-upload.txt") require.NoError(t, os.WriteFile(path, []byte("redirect payload"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) - originalFile := uploadSourceFile(t, upload) var requestCount atomic.Int32 - var uploaded string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requestCount.Add(1) - switch r.URL.Path { - case "/files": - http.Redirect(w, r, "/upload", status) - case "/upload": - assert.Positive(t, r.ContentLength) - assert.Empty(t, r.TransferEncoding) - file, _, parseErr := r.FormFile("file") - if !assert.NoError(t, parseErr) { - http.Error(w, parseErr.Error(), http.StatusBadRequest) - return - } - contents, readErr := io.ReadAll(file) - assert.NoError(t, readErr) - assert.NoError(t, file.Close()) - uploaded = string(contents) - w.Header().Set("Content-Type", "application/json") - _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) - default: - http.NotFound(w, r) - } + http.Redirect(w, r, "/upload", status) })) t.Cleanup(server.Close) @@ -279,49 +257,14 @@ func TestMultipartRequestOptionsFollow307And308ForRegularFiles(t *testing.T) { "purpose": "assistants", }, apiform.FormatBrackets) require.NoError(t, err) - var response []byte - options = append(options, option.WithResponseBodyInto(&response)) - client := openai.NewClient( - option.WithAPIKey("test-key"), - option.WithBaseURL(server.URL+"/"), - ) - - _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) - require.NoError(t, err) - require.Equal(t, int32(2), requestCount.Load()) - require.Equal(t, "redirect payload", uploaded) - _, err = originalFile.Stat() - require.ErrorIs(t, err, os.ErrClosed) - }) - } -} - -func TestMultipartRequestOptionsReject307And308ForUnknownSources(t *testing.T) { - for _, status := range []int{http.StatusTemporaryRedirect, http.StatusPermanentRedirect} { - t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { - var requestCount atomic.Int32 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - requestCount.Add(1) - http.Redirect(w, r, "/upload", status) - })) - t.Cleanup(server.Close) - - upload := &recordingReadCloser{reader: strings.NewReader("stdin payload")} - options, err := multipartRequestOptions(map[string]any{ - "file": fileUpload{Reader: upload, filename: "stdin"}, - }, apiform.FormatBrackets) - require.NoError(t, err) - var response []byte - options = append(options, option.WithResponseBodyInto(&response)) client := openai.NewClient( option.WithAPIKey("test-key"), option.WithBaseURL(server.URL+"/"), ) _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) - require.ErrorContains(t, err, "non-replayable source") + require.ErrorContains(t, err, "streamed multipart uploads are not replayable") require.Equal(t, int32(1), requestCount.Load()) - require.Equal(t, int32(1), upload.closeCount.Load()) }) } } @@ -348,8 +291,6 @@ func TestMultipartRequestOptionsRetryScalarOnlyForms(t *testing.T) { options, err := multipartRequestOptions(map[string]any{"prompt": "hello"}, apiform.FormatBrackets) require.NoError(t, err) - var response []byte - options = append(options, option.WithResponseBodyInto(&response)) client := openai.NewClient( option.WithAPIKey("test-key"), option.WithBaseURL(server.URL+"/"), @@ -360,72 +301,9 @@ func TestMultipartRequestOptionsRetryScalarOnlyForms(t *testing.T) { require.Equal(t, int32(2), requestCount.Load()) } -func TestMultipartRequestOptionsRetryRegularFiles(t *testing.T) { - path := filepath.Join(t.TempDir(), "retry-upload.txt") - require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) - upload, err := openFileUpload(path) - require.NoError(t, err) - sourceFile := uploadSourceFile(t, upload) - - var requestCount atomic.Int32 - var uploadsMu sync.Mutex - var uploads []string - var contentLengths []int64 - var retryCounts []string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - requestCount.Add(1) - file, _, parseErr := r.FormFile("file") - if !assert.NoError(t, parseErr) { - http.Error(w, parseErr.Error(), http.StatusBadRequest) - return - } - contents, readErr := io.ReadAll(file) - assert.NoError(t, readErr) - assert.NoError(t, file.Close()) - uploadsMu.Lock() - uploads = append(uploads, string(contents)) - contentLengths = append(contentLengths, r.ContentLength) - retryCounts = append(retryCounts, r.Header.Get("X-Stainless-Retry-Count")) - uploadsMu.Unlock() - - w.Header().Set("Content-Type", "application/json") - if requestCount.Load() == 1 { - w.Header().Set("Retry-After", "0") - w.WriteHeader(http.StatusTooManyRequests) - _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"rate_limit_error"}}`) - return - } - _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) - })) - t.Cleanup(server.Close) - - options, err := multipartRequestOptions(map[string]any{ - "file": upload, - "purpose": "assistants", - }, apiform.FormatBrackets) - require.NoError(t, err) - client := openai.NewClient( - option.WithAPIKey("test-key"), - option.WithBaseURL(server.URL+"/"), - ) - - _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) - require.NoError(t, err) - require.Equal(t, int32(2), requestCount.Load()) - uploadsMu.Lock() - require.Equal(t, []string{"retry payload", "retry payload"}, uploads) - require.Len(t, contentLengths, 2) - require.Positive(t, contentLengths[0]) - require.Equal(t, contentLengths[0], contentLengths[1]) - require.Equal(t, []string{"0", "1"}, retryCounts) - uploadsMu.Unlock() - _, err = sourceFile.Stat() - require.ErrorIs(t, err, os.ErrClosed) -} - -func TestMultipartRequestOptionsCloseRegularFilesAfterFinalRetry(t *testing.T) { - path := filepath.Join(t.TempDir(), "failed-retry-upload.txt") - require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) +func TestMultipartRequestOptionsDoNotRetryUploads(t *testing.T) { + path := filepath.Join(t.TempDir(), "single-attempt.txt") + require.NoError(t, os.WriteFile(path, []byte("single attempt"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) sourceFile := uploadSourceFile(t, upload) @@ -433,19 +311,12 @@ func TestMultipartRequestOptionsCloseRegularFilesAfterFinalRetry(t *testing.T) { var requestCount atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requestCount.Add(1) - file, _, parseErr := r.FormFile("file") - if !assert.NoError(t, parseErr) { - http.Error(w, parseErr.Error(), http.StatusBadRequest) - return - } - contents, readErr := io.ReadAll(file) - assert.NoError(t, readErr) - assert.Equal(t, "retry payload", string(contents)) - assert.NoError(t, file.Close()) + assert.Positive(t, r.ContentLength) + _, _ = io.Copy(io.Discard, r.Body) w.Header().Set("Content-Type", "application/json") w.Header().Set("Retry-After", "0") w.WriteHeader(http.StatusInternalServerError) - _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"server_error"}}`) + _, _ = io.WriteString(w, `{"error":{"message":"do not retry","type":"server_error"}}`) })) t.Cleanup(server.Close) @@ -461,83 +332,20 @@ func TestMultipartRequestOptionsCloseRegularFilesAfterFinalRetry(t *testing.T) { _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) require.Error(t, err) - require.Equal(t, int32(3), requestCount.Load()) + require.Equal(t, int32(1), requestCount.Load()) _, err = sourceFile.Stat() require.ErrorIs(t, err, os.ErrClosed) } -func TestMultipartRequestOptionsCloseRegularFilesWhenBackoffCanceled(t *testing.T) { - path := filepath.Join(t.TempDir(), "canceled-retry-upload.txt") - require.NoError(t, os.WriteFile(path, []byte("retry payload"), 0o600)) +func TestMultipartRequestOptionsPreserveIncompleteTransportError(t *testing.T) { + path := filepath.Join(t.TempDir(), "transport-error.txt") + require.NoError(t, os.WriteFile(path, []byte("never consumed"), 0o600)) upload, err := openFileUpload(path) require.NoError(t, err) sourceFile := uploadSourceFile(t, upload) - var requestCount atomic.Int32 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - requestCount.Add(1) - _, _ = io.Copy(io.Discard, r.Body) - _ = r.Body.Close() - w.Header().Set("Content-Type", "application/json") - w.Header().Set("Retry-After", "10") - w.WriteHeader(http.StatusTooManyRequests) - _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"rate_limit_error"}}`) - })) - t.Cleanup(server.Close) - - options, err := multipartRequestOptions(map[string]any{ - "file": upload, - "purpose": "assistants", - }, apiform.FormatBrackets) - require.NoError(t, err) - attemptReturned := make(chan struct{}) - observeAttempt := option.WithMiddleware(func( - req *http.Request, - next option.MiddlewareNext, - ) (*http.Response, error) { - res, requestErr := next(req) - close(attemptReturned) - return res, requestErr - }) - options = append([]option.RequestOption{observeAttempt}, options...) - client := openai.NewClient( - option.WithAPIKey("test-key"), - option.WithBaseURL(server.URL+"/"), - ) - ctx, cancel := context.WithCancel(context.Background()) - requestDone := make(chan error, 1) - go func() { - _, requestErr := client.Files.New(ctx, openai.FileNewParams{}, options...) - requestDone <- requestErr - }() - - select { - case <-attemptReturned: - case <-time.After(multipartTestTimeout): - t.Fatal("first retryable attempt did not return") - } - cancel() - select { - case requestErr := <-requestDone: - require.ErrorIs(t, requestErr, context.Canceled) - case <-time.After(multipartTestTimeout): - t.Fatal("request did not stop after cancellation during retry backoff") - } - require.Equal(t, int32(1), requestCount.Load()) - require.Eventually(t, func() bool { - _, statErr := sourceFile.Stat() - return errors.Is(statErr, os.ErrClosed) - }, time.Second, time.Millisecond, "cancellation cleanup must close the replay source") -} - -func TestMultipartRequestOptionsLetEarlySuccessfulUploadFinish(t *testing.T) { - path := filepath.Join(t.TempDir(), "early-response-upload.txt") - require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("payload", 1024)), 0o600)) - upload, err := openFileUpload(path) - require.NoError(t, err) - sourceFile := uploadSourceFile(t, upload) - - doer := newEarlyResponseHTTPDoer() + transportErr := errors.New("dial failed before reading request body") + doer := &failingHTTPDoer{err: transportErr} options, err := multipartRequestOptions(map[string]any{ "file": upload, "purpose": "assistants", @@ -550,126 +358,12 @@ func TestMultipartRequestOptionsLetEarlySuccessfulUploadFinish(t *testing.T) { ) _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) - require.NoError(t, err) - close(doer.release) - select { - case <-doer.done: - case <-time.After(multipartTestTimeout): - t.Fatal("transport did not finish consuming the request body") - } - require.NoError(t, doer.err) - require.Positive(t, doer.contentLength) - require.Equal(t, doer.contentLength, doer.bytesRead) + require.ErrorIs(t, err, transportErr) + require.Equal(t, int32(1), doer.calls.Load()) _, err = sourceFile.Stat() require.ErrorIs(t, err, os.ErrClosed) } -func TestMultipartReplayAttemptAbortDoesNotWaitForBlockedReader(t *testing.T) { - source := newManuallyReleasedReader() - body := newMultipartRequestBody(map[string]any{ - "file": fileUpload{Reader: source, filename: "blocked.bin"}, - }, apiform.FormatBrackets) - attempt := &multipartReplayAttempt{bodies: []io.ReadCloser{body}} - readDone := make(chan error, 1) - go func() { - _, err := io.Copy(io.Discard, body) - readDone <- err - }() - select { - case <-source.started: - case <-time.After(multipartTestTimeout): - t.Fatal("multipart encoder did not start the blocked source read") - } - - closeDone := make(chan error, 1) - go func() { closeDone <- attempt.abort() }() - returnedPromptly := false - select { - case <-closeDone: - returnedPromptly = true - case <-time.After(100 * time.Millisecond): - } - close(source.release) - if !returnedPromptly { - select { - case <-closeDone: - case <-time.After(multipartTestTimeout): - t.Fatal("attempt cleanup did not return after releasing the source") - } - } - select { - case <-readDone: - case <-time.After(multipartTestTimeout): - t.Fatal("request body read did not stop") - } - require.True(t, returnedPromptly, "attempt cleanup must not wait for an uninterruptible source read") -} - -func TestMultipartRequestOptionsRedirectReplayUsesOpenedFile(t *testing.T) { - if runtime.GOOS == "windows" { - t.Skip("Windows does not permit atomically replacing this open test file") - } - - dir := t.TempDir() - path := filepath.Join(dir, "redirect-upload.txt") - replacement := filepath.Join(dir, "replacement.txt") - require.NoError(t, os.WriteFile(path, []byte("benign-data"), 0o600)) - require.NoError(t, os.WriteFile(replacement, []byte("secret-data"), 0o600)) - upload, err := openFileUpload(path) - require.NoError(t, err) - - var requestCount atomic.Int32 - var uploadsMu sync.Mutex - var uploads []string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - requestCount.Add(1) - file, _, parseErr := r.FormFile("file") - if !assert.NoError(t, parseErr) { - http.Error(w, parseErr.Error(), http.StatusBadRequest) - return - } - contents, readErr := io.ReadAll(file) - assert.NoError(t, readErr) - assert.NoError(t, file.Close()) - uploadsMu.Lock() - uploads = append(uploads, string(contents)) - uploadsMu.Unlock() - - if r.URL.Path == "/files" { - if !assert.NoError(t, os.Rename(replacement, path)) { - http.Error(w, "could not replace upload path", http.StatusInternalServerError) - return - } - http.Redirect(w, r, "/upload", http.StatusTemporaryRedirect) - return - } - w.Header().Set("Content-Type", "application/json") - _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) - })) - t.Cleanup(server.Close) - - options, err := multipartRequestOptions(map[string]any{ - "file": upload, - "purpose": "assistants", - }, apiform.FormatBrackets) - require.NoError(t, err) - client := openai.NewClient( - option.WithAPIKey("test-key"), - option.WithBaseURL(server.URL+"/"), - ) - - _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) - require.NoError(t, err) - require.Equal(t, int32(2), requestCount.Load()) - uploadsMu.Lock() - require.Equal(t, []string{"benign-data", "benign-data"}, uploads) - uploadsMu.Unlock() - require.FileExists(t, path) - replacedContents, err := os.ReadFile(path) - require.NoError(t, err) - require.Equal(t, "secret-data", string(replacedContents)) -} - func TestMultipartRequestOptionsSetContentLengthForRegularFiles(t *testing.T) { path := filepath.Join(t.TempDir(), "upload.txt") require.NoError(t, os.WriteFile(path, []byte("known-size payload"), 0o600)) @@ -692,8 +386,6 @@ func TestMultipartRequestOptionsSetContentLengthForRegularFiles(t *testing.T) { "purpose": "assistants", }, apiform.FormatBrackets) require.NoError(t, err) - var response []byte - options = append(options, option.WithResponseBodyInto(&response)) client := openai.NewClient( option.WithAPIKey("test-key"), option.WithBaseURL(server.URL+"/"), @@ -703,25 +395,21 @@ func TestMultipartRequestOptionsSetContentLengthForRegularFiles(t *testing.T) { require.NoError(t, err) } -func TestMultipartRequestOptionsDisableRetriesForUnknownSources(t *testing.T) { +func TestMultipartRequestOptionsUseChunkedEncodingForUnknownSources(t *testing.T) { type requestInfo struct { contentLength int64 transferEncoding []string } - requestInfoCh := make(chan requestInfo, 4) + infoCh := make(chan requestInfo, 1) var requestCount atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requestCount.Add(1) - requestInfoCh <- requestInfo{ - contentLength: r.ContentLength, - transferEncoding: r.TransferEncoding, - } + infoCh <- requestInfo{r.ContentLength, append([]string(nil), r.TransferEncoding...)} _, _ = io.Copy(io.Discard, r.Body) - _ = r.Body.Close() w.Header().Set("Content-Type", "application/json") w.Header().Set("Retry-After", "0") w.WriteHeader(http.StatusInternalServerError) - _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"server_error"}}`) + _, _ = io.WriteString(w, `{"error":{"message":"do not retry","type":"server_error"}}`) })) t.Cleanup(server.Close) @@ -738,12 +426,24 @@ func TestMultipartRequestOptionsDisableRetriesForUnknownSources(t *testing.T) { _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) require.Error(t, err) require.Equal(t, int32(1), requestCount.Load()) - info := <-requestInfoCh + info := <-infoCh require.Equal(t, int64(-1), info.contentLength) require.Contains(t, info.transferEncoding, "chunked") require.Equal(t, int32(1), upload.closeCount.Load()) } +func TestMultipartRequestOptionsCloseUploadsOnLengthFailure(t *testing.T) { + first := &recordingReadCloser{reader: strings.NewReader("first")} + second := &recordingReadCloser{reader: strings.NewReader("second")} + _, err := multipartRequestOptions(map[string]any{ + "first": fileUpload{Reader: first, filename: "first.bin", knownSize: true, size: math.MaxInt64}, + "second": fileUpload{Reader: second, filename: "second.bin", knownSize: true, size: 1}, + }, apiform.FormatBrackets) + require.ErrorContains(t, err, "overflows int64") + require.Equal(t, int32(1), first.closeCount.Load()) + require.Equal(t, int32(1), second.closeCount.Load()) +} + func TestOpenFileUploadRejectsDirectoriesBeforeRequest(t *testing.T) { _, err := openFileUpload(t.TempDir()) require.ErrorContains(t, err, "is a directory") @@ -760,34 +460,20 @@ func TestOpenFileUploadStreamsProcfsWithUnknownLength(t *testing.T) { } require.NoError(t, err) require.False(t, upload.knownSize) - require.False(t, upload.hasKnownSize()) sourceFile := upload.Reader.(*os.File) - type requestInfo struct { - contentLength int64 - transferEncoding []string - contents string - } - infoCh := make(chan requestInfo, 1) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, int64(-1), r.ContentLength) + assert.Contains(t, r.TransferEncoding, "chunked") file, _, parseErr := r.FormFile("file") - if parseErr != nil { - infoCh <- requestInfo{contentLength: r.ContentLength, transferEncoding: r.TransferEncoding} + if !assert.NoError(t, parseErr) { http.Error(w, parseErr.Error(), http.StatusBadRequest) return } contents, readErr := io.ReadAll(file) + assert.NoError(t, readErr) + assert.NotEmpty(t, contents) _ = file.Close() - if readErr != nil { - infoCh <- requestInfo{contentLength: r.ContentLength, transferEncoding: r.TransferEncoding} - http.Error(w, readErr.Error(), http.StatusBadRequest) - return - } - infoCh <- requestInfo{ - contentLength: r.ContentLength, - transferEncoding: append([]string(nil), r.TransferEncoding...), - contents: string(contents), - } w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) })) @@ -805,10 +491,6 @@ func TestOpenFileUploadStreamsProcfsWithUnknownLength(t *testing.T) { _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) require.NoError(t, err) - info := <-infoCh - require.Equal(t, int64(-1), info.contentLength) - require.Contains(t, info.transferEncoding, "chunked") - require.NotEmpty(t, info.contents) _, err = sourceFile.Stat() require.ErrorIs(t, err, os.ErrClosed) } @@ -821,7 +503,6 @@ func TestOpenFileUploadTreatsEmptyRegularFileAsKnownLength(t *testing.T) { t.Cleanup(func() { require.NoError(t, upload.Close()) }) require.True(t, upload.knownSize) - require.True(t, upload.hasKnownSize()) require.Zero(t, upload.size) } @@ -855,6 +536,16 @@ func uploadSourceFile(t *testing.T, upload fileUpload) *os.File { return file } +type failingHTTPDoer struct { + err error + calls atomic.Int32 +} + +func (d *failingHTTPDoer) Do(*http.Request) (*http.Response, error) { + d.calls.Add(1) + return nil, d.err +} + type countingReader struct { remaining int64 bytesRead atomic.Int64 @@ -890,70 +581,6 @@ type recordingReadCloser struct { closeCount atomic.Int32 } -type earlyResponseHTTPDoer struct { - release chan struct{} - done chan struct{} - - contentLength int64 - bytesRead int64 - err error -} - -func newEarlyResponseHTTPDoer() *earlyResponseHTTPDoer { - return &earlyResponseHTTPDoer{ - release: make(chan struct{}), - done: make(chan struct{}), - } -} - -func (d *earlyResponseHTTPDoer) Do(req *http.Request) (*http.Response, error) { - readStarted := make(chan struct{}) - d.contentLength = req.ContentLength - go func() { - defer close(d.done) - buf := make([]byte, 1) - n, readErr := req.Body.Read(buf) - d.bytesRead += int64(n) - close(readStarted) - <-d.release - if readErr == nil { - copied, copyErr := io.Copy(io.Discard, req.Body) - d.bytesRead += copied - readErr = copyErr - } - if errors.Is(readErr, io.EOF) { - readErr = nil - } - d.err = errors.Join(readErr, req.Body.Close()) - }() - <-readStarted - return &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: io.NopCloser(strings.NewReader(`{"id":"file_123","object":"file"}`)), - Request: req, - }, nil -} - -type manuallyReleasedReader struct { - started chan struct{} - release chan struct{} - once sync.Once -} - -func newManuallyReleasedReader() *manuallyReleasedReader { - return &manuallyReleasedReader{ - started: make(chan struct{}), - release: make(chan struct{}), - } -} - -func (r *manuallyReleasedReader) Read([]byte) (int, error) { - r.once.Do(func() { close(r.started) }) - <-r.release - return 0, io.EOF -} - func (r *recordingReadCloser) Read(p []byte) (int, error) { return r.reader.Read(p) } diff --git a/pkg/cmd/multipartfallback_test.go b/pkg/cmd/multipartfallback_test.go index 7ae2149d..bf302929 100644 --- a/pkg/cmd/multipartfallback_test.go +++ b/pkg/cmd/multipartfallback_test.go @@ -15,7 +15,7 @@ import ( "github.com/urfave/cli/v3" ) -func TestFilesCreateCLIStreamsWhenSnapshotDirectoryUnavailable(t *testing.T) { +func TestFilesCreateCLIStreamsWithoutTempDirectory(t *testing.T) { for _, payload := range [][]byte{[]byte("streamed payload"), {}} { name := "non_empty" if len(payload) == 0 { @@ -45,7 +45,7 @@ func TestFilesCreateCLIStreamsWhenSnapshotDirectoryUnavailable(t *testing.T) { } } -func TestFilesCreateCLIDisablesReplayWhenSnapshotDirectoryUnavailable(t *testing.T) { +func TestFilesCreateCLIDoesNotReplayUploads(t *testing.T) { for _, status := range []int{http.StatusTemporaryRedirect, http.StatusTooManyRequests} { t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { path := filepath.Join(t.TempDir(), "upload.bin") @@ -74,7 +74,7 @@ func TestFilesCreateCLIDisablesReplayWhenSnapshotDirectoryUnavailable(t *testing err := runFilesCreateCLI(t.Context(), server.URL+"/", path) require.Error(t, err) - require.Equal(t, int32(1), requests.Load(), "snapshot failure must disable network replay") + require.Equal(t, int32(1), requests.Load(), "streamed uploads must not be replayed") }) } } @@ -97,7 +97,7 @@ func TestVideosCreateCLIPropagatesExtensionlessDirectoryError(t *testing.T) { require.Zero(t, requests.Load(), "an existing invalid upload must not fall back to a scalar") } -func TestVideosCreateCLIUploadsExtensionlessFileWhenSnapshotDirectoryUnavailable(t *testing.T) { +func TestVideosCreateCLIUploadsExtensionlessFileWithoutTempDirectory(t *testing.T) { dir := t.TempDir() t.Chdir(dir) require.NoError(t, os.WriteFile("artifact", []byte("extensionless payload"), 0o600)) diff --git a/pkg/cmd/multipartfile.go b/pkg/cmd/multipartfile.go index bca30247..bc8d85e9 100644 --- a/pkg/cmd/multipartfile.go +++ b/pkg/cmd/multipartfile.go @@ -103,9 +103,8 @@ func hasTrustworthyFileSize(file *os.File, info os.FileInfo) bool { } // openFileUpload opens path once. Ordinary files retain their exact stat size -// so the first request can stream directly with Content-Length while the -// multipart body opportunistically captures replay bytes. Virtual and -// unknown-length files remain single-shot streams. +// so the request can stream with Content-Length. Virtual and unknown-length +// files remain single-shot streams with chunked transfer encoding. func openFileUpload(path string) (fileUpload, error) { file, err := os.Open(path) if err != nil { diff --git a/pkg/cmd/multipartfile_unix.go b/pkg/cmd/multipartfile_unix.go deleted file mode 100644 index 9e5daf31..00000000 --- a/pkg/cmd/multipartfile_unix.go +++ /dev/null @@ -1,30 +0,0 @@ -//go:build !windows - -package cmd - -import ( - "errors" - "os" - - "golang.org/x/sys/unix" -) - -func createReplaySnapshotFile() (*os.File, error) { - file, err := os.CreateTemp("", "openai-multipart-upload-*") - if err != nil { - return nil, err - } - if err := os.Remove(file.Name()); err != nil { - return nil, errors.Join(err, file.Close()) - } - return file, nil -} - -func duplicateFile(file *os.File) (*os.File, error) { - fd, err := unix.Dup(int(file.Fd())) - if err != nil { - return nil, err - } - unix.CloseOnExec(fd) - return os.NewFile(uintptr(fd), file.Name()), nil -} diff --git a/pkg/cmd/multipartfile_windows.go b/pkg/cmd/multipartfile_windows.go deleted file mode 100644 index f663d528..00000000 --- a/pkg/cmd/multipartfile_windows.go +++ /dev/null @@ -1,88 +0,0 @@ -//go:build windows - -package cmd - -import ( - "errors" - "fmt" - "os" - "unsafe" - - "golang.org/x/sys/windows" -) - -func createReplaySnapshotFile() (*os.File, error) { - placeholder, err := os.CreateTemp("", "openai-multipart-upload-*") - if err != nil { - return nil, err - } - path := placeholder.Name() - if err := placeholder.Close(); err != nil { - return nil, errors.Join(err, os.Remove(path)) - } - if err := os.Remove(path); err != nil { - return nil, err - } - pathPtr, err := windows.UTF16PtrFromString(path) - if err != nil { - return nil, errors.Join(err, os.Remove(path)) - } - securityAttributes, err := replaySnapshotSecurityAttributes() - if err != nil { - return nil, errors.Join(err, os.Remove(path)) - } - handle, err := windows.CreateFile( - pathPtr, - windows.GENERIC_READ|windows.GENERIC_WRITE, - windows.FILE_SHARE_DELETE, - securityAttributes, - windows.CREATE_NEW, - windows.FILE_ATTRIBUTE_TEMPORARY|windows.FILE_FLAG_DELETE_ON_CLOSE, - 0, - ) - if err != nil { - return nil, errors.Join(err, os.Remove(path)) - } - return os.NewFile(uintptr(handle), path), nil -} - -func replaySnapshotSecurityAttributes() (*windows.SecurityAttributes, error) { - user, err := windows.GetCurrentProcessToken().GetTokenUser() - if err != nil { - return nil, fmt.Errorf("get current user for multipart replay snapshot: %w", err) - } - userSID := user.User.Sid.String() - if userSID == "" { - return nil, errors.New("get current user SID for multipart replay snapshot") - } - securityDescriptor, err := windows.SecurityDescriptorFromString( - "D:P(A;;GA;;;" + userSID + ")", - ) - if err != nil { - return nil, fmt.Errorf("create multipart replay snapshot security descriptor: %w", err) - } - return &windows.SecurityAttributes{ - Length: uint32(unsafe.Sizeof(windows.SecurityAttributes{})), - SecurityDescriptor: securityDescriptor, - }, nil -} - -func duplicateFile(file *os.File) (*os.File, error) { - process, err := windows.GetCurrentProcess() - if err != nil { - return nil, err - } - var duplicate windows.Handle - if err := windows.DuplicateHandle( - process, - windows.Handle(file.Fd()), - process, - &duplicate, - 0, - false, - windows.DUPLICATE_SAME_ACCESS, - ); err != nil { - return nil, err - } - return os.NewFile(uintptr(duplicate), file.Name()), nil -} diff --git a/pkg/cmd/multipartfile_windows_test.go b/pkg/cmd/multipartfile_windows_test.go deleted file mode 100644 index 4feee1eb..00000000 --- a/pkg/cmd/multipartfile_windows_test.go +++ /dev/null @@ -1,64 +0,0 @@ -//go:build windows - -package cmd - -import ( - "os" - "testing" - "unsafe" - - "github.com/stretchr/testify/require" - "golang.org/x/sys/windows" -) - -func TestReplaySnapshotHasProtectedCurrentUserDACL(t *testing.T) { - snapshot, err := createReplaySnapshotFile() - require.NoError(t, err) - t.Cleanup(func() { require.NoError(t, snapshot.Close()) }) - - securityDescriptor, err := windows.GetSecurityInfo( - windows.Handle(snapshot.Fd()), - windows.SE_FILE_OBJECT, - windows.DACL_SECURITY_INFORMATION, - ) - require.NoError(t, err) - control, _, err := securityDescriptor.Control() - require.NoError(t, err) - require.NotZero(t, control&windows.SE_DACL_PROTECTED) - - dacl, _, err := securityDescriptor.DACL() - require.NoError(t, err) - require.NotNil(t, dacl) - require.Equal(t, uint16(1), dacl.AceCount) - var ace *windows.ACCESS_ALLOWED_ACE - require.NoError(t, windows.GetAce(dacl, 0, &ace)) - require.Equal(t, uint8(windows.ACCESS_ALLOWED_ACE_TYPE), ace.Header.AceType) - require.Zero(t, ace.Header.AceFlags&windows.INHERITED_ACE) - - currentUser, err := windows.GetCurrentProcessToken().GetTokenUser() - require.NoError(t, err) - aceSID := (*windows.SID)(unsafe.Pointer(&ace.SidStart)) - require.True(t, windows.EqualSid(currentUser.User.Sid, aceSID)) - - otherHandle, err := os.Open(snapshot.Name()) - if otherHandle != nil { - require.NoError(t, otherHandle.Close()) - } - require.Error(t, err) -} - -func TestReplaySnapshotIsDeletedAfterFinalHandleCloses(t *testing.T) { - snapshot, err := createReplaySnapshotFile() - require.NoError(t, err) - path := snapshot.Name() - _, err = snapshot.Write([]byte("sensitive upload")) - require.NoError(t, err) - - duplicate, err := duplicateFile(snapshot) - require.NoError(t, err) - require.NoError(t, snapshot.Close()) - require.NoError(t, duplicate.Close()) - - _, err = os.Stat(path) - require.ErrorIs(t, err, os.ErrNotExist) -} diff --git a/pkg/cmd/multipartintegrity_test.go b/pkg/cmd/multipartintegrity_test.go index 497c5c4d..fe06c79f 100644 --- a/pkg/cmd/multipartintegrity_test.go +++ b/pkg/cmd/multipartintegrity_test.go @@ -3,7 +3,6 @@ package cmd import ( "context" "errors" - "fmt" "io" "mime" "mime/multipart" @@ -12,7 +11,6 @@ import ( "os" "path/filepath" "strings" - "sync" "testing" "github.com/openai/openai-cli/internal/apiform" @@ -22,58 +20,6 @@ import ( "github.com/urfave/cli/v3" ) -func TestFilesCreateCLIUsesImmutableReplayContent(t *testing.T) { - for _, status := range []int{ - http.StatusTemporaryRedirect, - http.StatusPermanentRedirect, - http.StatusTooManyRequests, - } { - t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) { - path := filepath.Join(t.TempDir(), "upload.txt") - require.NoError(t, os.WriteFile(path, []byte("benign-data"), 0o600)) - - var mu sync.Mutex - var uploads []string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - contents, readErr := readMultipartUpload(r) - if readErr != nil { - http.Error(w, readErr.Error(), http.StatusBadRequest) - return - } - mu.Lock() - uploads = append(uploads, string(contents)) - requestNumber := len(uploads) - mu.Unlock() - - if requestNumber == 1 { - if writeErr := os.WriteFile(path, []byte("secret-data"), 0o600); writeErr != nil { - http.Error(w, writeErr.Error(), http.StatusInternalServerError) - return - } - if status == http.StatusTooManyRequests { - w.Header().Set("Content-Type", "application/json") - w.Header().Set("Retry-After", "0") - w.WriteHeader(status) - _, _ = io.WriteString(w, `{"error":{"message":"retry me","type":"rate_limit_error"}}`) - return - } - http.Redirect(w, r, "/upload", status) - return - } - - w.Header().Set("Content-Type", "application/json") - _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) - })) - t.Cleanup(server.Close) - - require.NoError(t, runFilesCreateCLI(t.Context(), server.URL+"/", path)) - mu.Lock() - require.Equal(t, []string{"benign-data", "benign-data"}, uploads) - mu.Unlock() - }) - } -} - func TestFilesCreateCLIAbortsWhenSourceShrinksDuringFirstSend(t *testing.T) { const uploadSize = 16 << 20 path := filepath.Join(t.TempDir(), "shrinking.bin") @@ -127,9 +73,9 @@ func TestFilesCreateCLIAbortsWhenSourceShrinksDuringFirstSend(t *testing.T) { require.Less(t, result.bytes, int64(uploadSize)) } -func TestMultipartReplayPrematureEOFDoesNotCommitFinalBoundary(t *testing.T) { +func TestMultipartPrematureEOFDoesNotCommitFinalBoundary(t *testing.T) { const uploadSize = 8 << 20 - path := filepath.Join(t.TempDir(), "replay-source.bin") + path := filepath.Join(t.TempDir(), "short-source.bin") require.NoError(t, os.WriteFile(path, []byte(strings.Repeat("a", uploadSize)), 0o600)) file, err := os.Open(path) require.NoError(t, err) @@ -194,11 +140,11 @@ func TestMultipartReplayPrematureEOFDoesNotCommitFinalBoundary(t *testing.T) { _, err = client.Files.New(context.Background(), openai.FileNewParams{}, options...) require.ErrorIs(t, err, io.ErrUnexpectedEOF) result := <-resultCh - require.False(t, result.committed, "an incomplete replay source must not emit a terminal multipart boundary") + require.False(t, result.committed, "an incomplete source must not emit a terminal multipart boundary") require.Error(t, result.err) } -func TestOpenFileUploadDoesNotRequireSnapshotBeforeRead(t *testing.T) { +func TestOpenFileUploadDoesNotRequireTempStorageBeforeRead(t *testing.T) { path := filepath.Join(t.TempDir(), "upload.txt") require.NoError(t, os.WriteFile(path, []byte("streamed payload"), 0o600)) setUnavailableTempDir(t) diff --git a/pkg/cmd/multipartreplay.go b/pkg/cmd/multipartreplay.go deleted file mode 100644 index 5ad917a5..00000000 --- a/pkg/cmd/multipartreplay.go +++ /dev/null @@ -1,160 +0,0 @@ -package cmd - -import ( - "errors" - "fmt" - "io" - "os" - "sync" -) - -// multipartReplaySource opportunistically captures the exact encoded bytes of -// the first request body. Capture failures never fail that first stream; they -// only make redirects and retries unavailable. -type multipartReplaySource struct { - mu sync.Mutex - expected int64 - captured int64 - file *os.File - ready bool - closed bool - err error - closeErr error -} - -func newMultipartReplaySource(expected int64) *multipartReplaySource { - return &multipartReplaySource{expected: expected} -} - -func (s *multipartReplaySource) Capture(p []byte) { - if len(p) == 0 { - return - } - - s.mu.Lock() - defer s.mu.Unlock() - if s.ready || s.closed || s.err != nil { - return - } - if s.file == nil { - file, err := createReplaySnapshotFile() - if err != nil { - s.failLocked(fmt.Errorf("create multipart replay snapshot: %w", err)) - return - } - s.file = file - } - - n, err := s.file.Write(p) - s.captured += int64(n) - if err == nil && n != len(p) { - err = io.ErrShortWrite - } - if err != nil { - s.failLocked(fmt.Errorf("write multipart replay snapshot: %w", err)) - return - } - if s.captured > s.expected { - s.failLocked(fmt.Errorf("multipart replay snapshot exceeded expected length %d", s.expected)) - return - } - if s.captured == s.expected { - s.ready = true - } -} - -func (s *multipartReplaySource) Fail(err error) { - if err == nil { - err = errors.New("multipart replay capture did not complete") - } - s.mu.Lock() - defer s.mu.Unlock() - if !s.ready { - s.failLocked(err) - } -} - -func (s *multipartReplaySource) failLocked(err error) { - if s.err == nil { - s.err = err - } - if s.file != nil { - s.closeErr = errors.Join(s.closeErr, s.file.Close()) - s.file = nil - } -} - -func (s *multipartReplaySource) Ready() bool { - s.mu.Lock() - defer s.mu.Unlock() - return s.ready && !s.closed -} - -func (s *multipartReplaySource) ReplayReader() (io.ReadCloser, error) { - s.mu.Lock() - if !s.ready || s.closed || s.file == nil { - err := s.err - if err == nil { - err = errors.New("first multipart stream did not complete") - } - s.mu.Unlock() - return nil, fmt.Errorf("multipart upload is not replayable: %w", err) - } - file, err := duplicateFile(s.file) - s.mu.Unlock() - if err != nil { - return nil, fmt.Errorf("duplicate multipart replay snapshot: %w", err) - } - return &exactLengthReadCloser{ - exactLengthReader: exactLengthReader{ - reader: io.NewSectionReader(file, 0, s.expected), - remaining: s.expected, - }, - closer: file, - }, nil -} - -func (s *multipartReplaySource) Close() error { - s.mu.Lock() - if s.closed { - err := s.closeErr - s.mu.Unlock() - return err - } - s.closed = true - file := s.file - s.file = nil - if file != nil { - s.closeErr = errors.Join(s.closeErr, file.Close()) - } - err := s.closeErr - s.mu.Unlock() - return err -} - -type capturingMultipartBody struct { - body *multipartRequestBody - source *multipartReplaySource -} - -func (b *capturingMultipartBody) Read(p []byte) (int, error) { - n, err := b.body.Read(p) - if n > 0 { - b.source.Capture(p[:n]) - } - if err != nil && !errors.Is(err, io.EOF) { - b.source.Fail(err) - } - return n, err -} - -func (b *capturingMultipartBody) Close() error { - if !b.source.Ready() { - b.source.Fail(errors.New("first multipart stream closed before replay capture completed")) - } - return b.body.Close() -} - -func (b *capturingMultipartBody) Done() <-chan struct{} { - return b.body.done -} diff --git a/pkg/cmd/multipartresource_unix_test.go b/pkg/cmd/multipartresource_unix_test.go index 2e010c61..12257906 100644 --- a/pkg/cmd/multipartresource_unix_test.go +++ b/pkg/cmd/multipartresource_unix_test.go @@ -18,7 +18,7 @@ import ( const multipartResourceHelperEnv = "OPENAI_MULTIPART_RESOURCE_HELPER" -func TestFilesCreateCLIStreamsSparseFileWhenSnapshotQuotaExceeded(t *testing.T) { +func TestFilesCreateCLIStreamsSparseFileWithLowFileSizeLimit(t *testing.T) { const uploadSize = 256 << 10 path := filepath.Join(t.TempDir(), "sparse.bin") file, err := os.Create(path) @@ -49,7 +49,7 @@ func TestFilesCreateCLIStreamsSparseFileWhenSnapshotQuotaExceeded(t *testing.T) require.Equal(t, int64(uploadSize), <-received) } -func TestImagesEditCLIDialsBeforeOpeningReplayDescriptors(t *testing.T) { +func TestImagesEditCLIStreamsMultipleFilesWithLowDescriptorLimit(t *testing.T) { paths := make([]string, 16) for i := range paths { paths[i] = filepath.Join(t.TempDir(), "image.png") From 1a74fd7630d7deb37dd3b17f4577cbf5dae8d870 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 18 Aug 2026 16:37:24 +0000 Subject: [PATCH 9/9] fix: cancel stalled multipart uploads --- pkg/cmd/multipartbody.go | 16 ++++ pkg/cmd/multipartbody_test.go | 80 +++++++++++++++++ pkg/cmd/multipartresource_unix_test.go | 115 +++++++++++++++++++++++++ 3 files changed, 211 insertions(+) diff --git a/pkg/cmd/multipartbody.go b/pkg/cmd/multipartbody.go index 07130b71..a9983e60 100644 --- a/pkg/cmd/multipartbody.go +++ b/pkg/cmd/multipartbody.go @@ -2,6 +2,7 @@ package cmd import ( "bytes" + "context" "errors" "fmt" "io" @@ -165,7 +166,22 @@ func streamedMultipartMiddleware(contentLength int64) option.Middleware { if contentLength >= 0 { req.ContentLength = contentLength } + + // The SDK closes a request body after next returns, but a transport can + // remain blocked in a streamed source read after the request is canceled. + // Close the body independently to release owned files, and wait for an + // already-started callback so it cannot outlive this middleware attempt. + cancelDone := make(chan struct{}) + stopCancellation := context.AfterFunc(req.Context(), func() { + defer close(cancelDone) + if req.Body != nil { + _ = req.Body.Close() + } + }) res, err := next(req) + if !stopCancellation() { + <-cancelDone + } if err != nil || res == nil || (res.StatusCode != http.StatusTemporaryRedirect && res.StatusCode != http.StatusPermanentRedirect) { return res, err diff --git a/pkg/cmd/multipartbody_test.go b/pkg/cmd/multipartbody_test.go index 942258ea..446f006d 100644 --- a/pkg/cmd/multipartbody_test.go +++ b/pkg/cmd/multipartbody_test.go @@ -364,6 +364,65 @@ func TestMultipartRequestOptionsPreserveIncompleteTransportError(t *testing.T) { require.ErrorIs(t, err, os.ErrClosed) } +func TestMultipartRequestOptionsCancellationClosesSourceOnce(t *testing.T) { + upload := newCloseUnblocksReader(errors.New("source closed")) + options, err := multipartRequestOptions(map[string]any{ + "file": fileUpload{Reader: upload, filename: "blocked.bin"}, + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL("https://example.com/"), + option.WithHTTPClient(bodyReadingHTTPDoer{}), + ) + ctx, cancel := context.WithCancel(context.Background()) + requestDone := make(chan error, 1) + go func() { + _, requestErr := client.Files.New(ctx, openai.FileNewParams{}, options...) + requestDone <- requestErr + }() + + select { + case <-upload.readStarted: + case <-time.After(multipartTestTimeout): + cancel() + t.Fatal("transport did not start reading the upload") + } + cancel() + select { + case requestErr := <-requestDone: + require.ErrorIs(t, requestErr, context.Canceled) + case <-time.After(multipartTestTimeout): + t.Fatal("canceled request did not return") + } + require.Equal(t, int32(1), upload.closeCount.Load()) +} + +func TestMultipartRequestOptionsStopCancellationWatcherAfterAttempt(t *testing.T) { + upload := &recordingReadCloser{reader: strings.NewReader("not consumed")} + doer := &retainingHTTPDoer{} + options, err := multipartRequestOptions(map[string]any{ + "file": fileUpload{Reader: upload, filename: "retained.bin"}, + }, apiform.FormatBrackets) + require.NoError(t, err) + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL("https://example.com/"), + option.WithHTTPClient(doer), + ) + ctx, cancel := context.WithCancel(context.Background()) + + _, err = client.Files.New(ctx, openai.FileNewParams{}, options...) + require.NoError(t, err) + cancel() + require.Never(t, func() bool { return upload.closeCount.Load() != 0 }, 100*time.Millisecond, time.Millisecond, + "a stopped watcher must not close a body after the attempt returns") + require.NotNil(t, doer.body) + require.NoError(t, doer.body.Close()) + require.NoError(t, doer.body.Close()) + require.Equal(t, int32(1), upload.closeCount.Load()) +} + func TestMultipartRequestOptionsSetContentLengthForRegularFiles(t *testing.T) { path := filepath.Join(t.TempDir(), "upload.txt") require.NoError(t, os.WriteFile(path, []byte("known-size payload"), 0o600)) @@ -541,6 +600,27 @@ type failingHTTPDoer struct { calls atomic.Int32 } +type bodyReadingHTTPDoer struct{} + +func (bodyReadingHTTPDoer) Do(req *http.Request) (*http.Response, error) { + _, err := io.Copy(io.Discard, req.Body) + return nil, err +} + +type retainingHTTPDoer struct { + body io.ReadCloser +} + +func (d *retainingHTTPDoer) Do(req *http.Request) (*http.Response, error) { + d.body = req.Body + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"file_123","object":"file"}`)), + Request: req, + }, nil +} + func (d *failingHTTPDoer) Do(*http.Request) (*http.Response, error) { d.calls.Add(1) return nil, d.err diff --git a/pkg/cmd/multipartresource_unix_test.go b/pkg/cmd/multipartresource_unix_test.go index 12257906..2199b752 100644 --- a/pkg/cmd/multipartresource_unix_test.go +++ b/pkg/cmd/multipartresource_unix_test.go @@ -3,14 +3,20 @@ package cmd import ( + "context" "encoding/json" + "fmt" "io" + "mime" + "mime/multipart" "net/http" "net/http/httptest" "os" "os/exec" "path/filepath" + "sync" "testing" + "time" "github.com/stretchr/testify/require" "golang.org/x/sys/unix" @@ -18,6 +24,115 @@ import ( const multipartResourceHelperEnv = "OPENAI_MULTIPART_RESOURCE_HELPER" +func TestFilesCreateCLICancelClosesStalledFIFO(t *testing.T) { + fifoPath := filepath.Join(t.TempDir(), "upload.fifo") + require.NoError(t, unix.Mkfifo(fifoPath, 0o600)) + + writerRelease := make(chan struct{}) + writerProbe := make(chan chan<- error) + writerDone := make(chan error, 1) + go func() { + writer, err := os.OpenFile(fifoPath, os.O_WRONLY, 0) + if err != nil { + writerDone <- err + return + } + defer writer.Close() + if _, err := writer.Write([]byte("partial upload")); err != nil { + writerDone <- err + return + } + select { + case result := <-writerProbe: + _, err := writer.Write([]byte("still open")) + result <- err + <-writerRelease + case <-writerRelease: + } + writerDone <- nil + }() + var releaseWriterOnce sync.Once + releaseWriter := func() { releaseWriterOnce.Do(func() { close(writerRelease) }) } + + requestStarted := make(chan struct{}) + receiverDone := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + receiverDone <- fmt.Errorf("authorization header = %q", got) + close(requestStarted) + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + mediaType, params, err := mime.ParseMediaType(r.Header.Get("Content-Type")) + if err != nil { + receiverDone <- fmt.Errorf("parse content type: %w", err) + close(requestStarted) + http.Error(w, "invalid content type", http.StatusBadRequest) + return + } + if mediaType != "multipart/form-data" { + receiverDone <- fmt.Errorf("content type = %q", mediaType) + close(requestStarted) + http.Error(w, "invalid content type", http.StatusBadRequest) + return + } + part, err := nextFilePart(multipart.NewReader(r.Body, params["boundary"])) + if err != nil { + receiverDone <- err + close(requestStarted) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + buf := make([]byte, 1) + _, err = part.Read(buf) + close(requestStarted) + if err == nil { + _, err = io.Copy(io.Discard, part) + } + receiverDone <- err + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"file_123","object":"file"}`) + })) + t.Cleanup(server.Close) + t.Cleanup(releaseWriter) + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + requestDone := make(chan error, 1) + go func() { + requestDone <- runFilesCreateCLI(ctx, server.URL+"/", fifoPath) + }() + + select { + case <-requestStarted: + case <-time.After(multipartTestTimeout): + releaseWriter() + t.Fatal("authenticated receiver did not observe the FIFO upload") + } + cancel() + + select { + case err := <-requestDone: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + releaseWriter() + err := <-requestDone + t.Fatalf("canceled FIFO upload waited for its producer to close: %v", err) + } + select { + case err := <-writerDone: + t.Fatalf("FIFO producer exited before the test released it: %v", err) + default: + } + probeResult := make(chan error, 1) + writerProbe <- probeResult + require.ErrorIs(t, <-probeResult, unix.EPIPE, "cancellation must close the owned FIFO reader") + + releaseWriter() + require.NoError(t, <-writerDone) + require.Error(t, <-receiverDone, "canceled multipart stream must end without a terminal boundary") +} + func TestFilesCreateCLIStreamsSparseFileWithLowFileSizeLimit(t *testing.T) { const uploadSize = 256 << 10 path := filepath.Join(t.TempDir(), "sparse.bin")