From 3517be1bd1234b425a5e746398ec7d1ad1357457 Mon Sep 17 00:00:00 2001 From: wolo Date: Fri, 17 Jul 2026 22:15:56 +0000 Subject: [PATCH] fix(gcsartifact): prevent artifact version race in concurrent Save MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Save assigned versions non-atomically: it listed the existing versions, computed max+1, and wrote that blob. Two concurrent Save calls for the same artifact would read the same version set, pick the same next version, and the second write would silently overwrite the first — duplicate returned versions and a lost object. The code carried a TODO acknowledging this. Assign versions with optimistic concurrency instead: write the new blob with a "does not exist" generation precondition (x-goog-if-generation-match: 0). GCS has no transaction spanning the list and the write, so a colliding writer now fails the precondition (HTTP 412); Save re-reads the versions and retries with a fresh number, using a jittered backoff between attempts and bounded by maxSaveAttempts. It returns errVersionConflict if the budget is exhausted rather than overwriting an existing version. N concurrent saves therefore yield exactly versions 1..N with no lost writes. Testing Plan: - go build -mod=readonly ./... - go test -race -mod=readonly -count=1 -shuffle=on ./artifact/... - golangci-lint run ./artifact/... - go mod tidy -diff - TestGCSArtifactServiceConcurrentSave: concurrent savers assert distinct, contiguous versions and that nothing is overwritten (saver count kept within the retry budget). Verified it fails against the pre-fix code. - TestSaveReturnsWriteErrorWithoutRetry: a non-precondition write error is surfaced as-is and not retried. - TestSaveExhaustsRetriesOnPersistentConflict: persistent 412s give up with errVersionConflict after maxSaveAttempts. - TestBackoffDelayBounds: jittered backoff stays within bounds. --- artifact/gcsartifact/gcs_client.go | 8 + artifact/gcsartifact/gcs_test.go | 328 +++++++++++++++++++++++------ artifact/gcsartifact/service.go | 148 ++++++++++--- 3 files changed, 397 insertions(+), 87 deletions(-) diff --git a/artifact/gcsartifact/gcs_client.go b/artifact/gcsartifact/gcs_client.go index bf36435c6..8120ec079 100644 --- a/artifact/gcsartifact/gcs_client.go +++ b/artifact/gcsartifact/gcs_client.go @@ -39,6 +39,9 @@ type gcsObject interface { newReader(ctx context.Context) (io.ReadCloser, error) delete(ctx context.Context) error attrs(ctx context.Context) (*storage.ObjectAttrs, error) + // ifNotExist returns a handle whose write succeeds only if the object does + // not already exist; writing over an existing object fails with HTTP 412. + ifNotExist() gcsObject } // gcsObjectIterator @@ -97,6 +100,11 @@ func (w *gcsObjectWrapper) newWriter(ctx context.Context) gcsWriter { return &gcsWriterWrapper{w: w.object.NewWriter(ctx)} } +// ifNotExist implements the gcsObject interface for gcsObjectWrapper. +func (w *gcsObjectWrapper) ifNotExist() gcsObject { + return &gcsObjectWrapper{object: w.object.If(storage.Conditions{DoesNotExist: true})} +} + // NewReader implements the gcsObject interface for gcsObjectWrapper. func (w *gcsObjectWrapper) newReader(ctx context.Context) (io.ReadCloser, error) { return w.object.NewReader(ctx) diff --git a/artifact/gcsartifact/gcs_test.go b/artifact/gcsartifact/gcs_test.go index d2032536b..5bc7ad770 100644 --- a/artifact/gcsartifact/gcs_test.go +++ b/artifact/gcsartifact/gcs_test.go @@ -17,29 +17,44 @@ package gcsartifact import ( "bytes" "context" + "errors" + "fmt" "io" "io/fs" + "net/http" + "slices" "strings" "sync" "testing" "time" "cloud.google.com/go/storage" + "github.com/google/go-cmp/cmp" + "google.golang.org/api/googleapi" "google.golang.org/api/iterator" + "google.golang.org/genai" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/adk/v2/artifact" "google.golang.org/adk/v2/internal/artifact/tests" ) -// newGCSArtifactServiceForTesting creates a gcsService for the specified bucket using a mocked inmemory client -func newGCSArtifactServiceForTesting(bucketName string) (artifact.Service, error) { +// newGCSServiceForTesting creates a *gcsService backed by the in-memory fake. It +// stubs out the retry backoff so retry-heavy tests stay fast. +func newGCSServiceForTesting(bucketName string) *gcsService { client := newFakeClient() - s := &gcsService{ + return &gcsService{ bucketName: bucketName, storageClient: client, bucket: client.bucket(bucketName), + sleep: func(ctx context.Context, _ time.Duration) error { return ctx.Err() }, } - return s, nil +} + +// newGCSArtifactServiceForTesting creates a gcsService for the specified bucket using a mocked inmemory client +func newGCSArtifactServiceForTesting(bucketName string) (artifact.Service, error) { + return newGCSServiceForTesting(bucketName), nil } func TestGCSArtifactService(t *testing.T) { @@ -49,6 +64,159 @@ func TestGCSArtifactService(t *testing.T) { tests.TestArtifactService(t, "GCS", factory) } +// TestGCSArtifactServiceConcurrentSave checks that concurrent saves of the same +// artifact get distinct versions and none is overwritten. Pre-fix, colliding +// savers picked the same version and clobbered each other. +func TestGCSArtifactServiceConcurrentSave(t *testing.T) { + srv, err := newGCSArtifactServiceForTesting("concurrent") + if err != nil { + t.Fatalf("failed to set up service: %v", err) + } + + // Must stay <= maxSaveAttempts: a maximally-unlucky saver can lose one race + // per other writer, so it may need up to savers attempts to land a version. + const savers = 10 + + var wg sync.WaitGroup + gotVersions := make([]int64, savers) + errs := make([]error, savers) + for i := range savers { + wg.Add(1) + go func() { + defer wg.Done() + resp, err := srv.Save(t.Context(), saveReq()) + if err != nil { + errs[i] = err + return + } + gotVersions[i] = resp.Version + }() + } + wg.Wait() + + for i, err := range errs { + if err != nil { + // Fail fast: a saver error leaves zeroed versions, which would make + // the version diffs below noisy and misleading. + t.Fatalf("saver %d: Save() failed: %v", i, err) + } + } + + slices.Sort(gotVersions) + want := make([]int64, savers) + for i := range want { + want[i] = int64(i + 1) + } + if diff := cmp.Diff(want, gotVersions); diff != "" { + t.Errorf("returned versions mismatch (-want +got):\n%s", diff) + } + + // The stored versions must match what Save reported (nothing overwritten). + resp, err := srv.Versions(t.Context(), &artifact.VersionsRequest{ + AppName: "app", UserID: "user", SessionID: "session", FileName: "file", + }) + if err != nil { + t.Fatalf("Versions() failed: %v", err) + } + stored := resp.Versions + slices.Sort(stored) + if diff := cmp.Diff(want, stored); diff != "" { + t.Errorf("stored versions mismatch (-want +got):\n%s", diff) + } +} + +// TestSaveReturnsWriteErrorWithoutRetry checks that a non-precondition write +// error is surfaced as-is and not retried. +func TestSaveReturnsWriteErrorWithoutRetry(t *testing.T) { + svc := newGCSServiceForTesting("errtest") + fb := svc.bucket.(*fakeBucket) + wantErr := &googleapi.Error{Code: http.StatusInternalServerError, Message: "boom"} + fb.closeErr = wantErr + + if _, err := svc.Save(t.Context(), saveReq()); !errors.Is(err, wantErr) { + t.Fatalf("Save() err = %v, want %v", err, wantErr) + } + if got := fb.closeCalls; got != 1 { + t.Errorf("write attempts = %d, want 1 (no retry on non-precondition error)", got) + } +} + +// TestSaveExhaustsRetriesOnPersistentConflict checks that Save gives up with +// ErrVersionConflict after maxSaveAttempts when every write hits a precondition +// failure. +func TestSaveExhaustsRetriesOnPersistentConflict(t *testing.T) { + svc := newGCSServiceForTesting("errtest") + fb := svc.bucket.(*fakeBucket) + fb.closeErr = &googleapi.Error{Code: http.StatusPreconditionFailed} + + if _, err := svc.Save(t.Context(), saveReq()); !errors.Is(err, ErrVersionConflict) { + t.Fatalf("Save() err = %v, want ErrVersionConflict", err) + } + if got := fb.closeCalls; got != maxSaveAttempts { + t.Errorf("write attempts = %d, want %d", got, maxSaveAttempts) + } +} + +// TestBackoffDelayBounds checks the jittered backoff stays within [0, saveRetryMaxDelay]. +func TestBackoffDelayBounds(t *testing.T) { + for attempt := range maxSaveAttempts { + if d := backoffDelay(attempt); d < 0 || d > saveRetryMaxDelay { + t.Errorf("backoffDelay(%d) = %v, want within [0, %v]", attempt, d, saveRetryMaxDelay) + } + } +} + +// TestIsPreconditionFailed covers both transports (HTTP *googleapi.Error and +// gRPC status), raw and wrapped, since Save's retry hinges on this detection. +func TestIsPreconditionFailed(t *testing.T) { + for _, tc := range []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"generic", errors.New("boom"), false}, + {"http 412", &googleapi.Error{Code: http.StatusPreconditionFailed}, true}, + {"http 412 wrapped", fmt.Errorf("save: %w", &googleapi.Error{Code: http.StatusPreconditionFailed}), true}, + {"http 500", &googleapi.Error{Code: http.StatusInternalServerError}, false}, + {"grpc failed precondition", status.Error(codes.FailedPrecondition, "cond"), true}, + {"grpc failed precondition wrapped", fmt.Errorf("save: %w", status.Error(codes.FailedPrecondition, "cond")), true}, + {"grpc not found", status.Error(codes.NotFound, "nf"), false}, + } { + t.Run(tc.name, func(t *testing.T) { + if got := isPreconditionFailed(tc.err); got != tc.want { + t.Errorf("isPreconditionFailed(%v) = %v, want %v", tc.err, got, tc.want) + } + }) + } +} + +// TestSleepContext checks that sleepContext honors a non-positive duration, an +// already-cancelled context, and a normal short wait. +func TestSleepContext(t *testing.T) { + if err := sleepContext(t.Context(), 0); err != nil { + t.Errorf("sleepContext(_, 0) = %v, want nil", err) + } + + cancelled, cancel := context.WithCancel(t.Context()) + cancel() + if err := sleepContext(cancelled, time.Hour); !errors.Is(err, context.Canceled) { + t.Errorf("sleepContext(cancelled, 1h) = %v, want context.Canceled", err) + } + + if err := sleepContext(t.Context(), time.Millisecond); err != nil { + t.Errorf("sleepContext(_, 1ms) = %v, want nil", err) + } +} + +// saveReq returns a minimal valid SaveRequest for the fixed test artifact. +func saveReq() *artifact.SaveRequest { + return &artifact.SaveRequest{ + AppName: "app", UserID: "user", SessionID: "session", FileName: "file", + Part: genai.NewPartFromBytes([]byte("data"), "text/plain"), + } +} + // ---------------------------------- Mock Implementations ----------------------------------- // fakeClient implements the gcsClient interface for testing. type fakeClient struct { @@ -58,7 +226,7 @@ type fakeClient struct { func newFakeClient() gcsClient { return &fakeClient{ inMemoryBucket: &fakeBucket{ - objectsMap: make(map[string]*fakeObject), + blobs: make(map[string]*fakeBlob), }, } } @@ -70,45 +238,52 @@ func (c *fakeClient) bucket(name string) gcsBucket { // fakeBucket implements the gcsBucket interface for testing. type fakeBucket struct { - mu sync.Mutex - objectsMap map[string]*fakeObject + mu sync.Mutex + blobs map[string]*fakeBlob + + // Test hooks (guarded by mu): closeErr, when set, makes every writer Close + // return it (a simulated write failure); closeCalls counts Close calls. + closeErr error + closeCalls int } -// Object returns a fake object from the in-memory store. +// object returns a handle to the named blob, creating an empty backing store on +// first use so repeat handles to the same name share state. func (f *fakeBucket) object(name string) gcsObject { f.mu.Lock() defer f.mu.Unlock() - if _, ok := f.objectsMap[name]; !ok { - f.objectsMap[name] = &fakeObject{name: name} + b, ok := f.blobs[name] + if !ok { + b = &fakeBlob{name: name} + f.blobs[name] = b } - return f.objectsMap[name] + return &fakeObject{blob: b, bucket: f} } -// Objects simulates iterating over objects with a prefix. +// objects iterates the existing (written, not deleted) blobs matching the +// prefix, snapshotting attributes at call time like a real GCS listing so a +// later delete can't retroactively change what this iterator returns. func (f *fakeBucket) objects(ctx context.Context, q *storage.Query) gcsObjectIterator { f.mu.Lock() defer f.mu.Unlock() - var matchingObjects []*fakeObject - for name, obj := range f.objectsMap { + var attrs []*storage.ObjectAttrs + for name, b := range f.blobs { if q != nil && q.Prefix != "" && !strings.HasPrefix(name, q.Prefix) { continue } - if !obj.deleted { - matchingObjects = append(matchingObjects, obj) + b.mu.Lock() + if !b.deleted && b.data != nil { + attrs = append(attrs, &storage.ObjectAttrs{Name: b.name, ContentType: b.contentType}) } + b.mu.Unlock() } - - // This is the key change. We return a custom type that has a `Next` method - // that manages its own state and returns the correct values. - return &fakeObjectIterator{ - objects: matchingObjects, - index: 0, - } + return &fakeObjectIterator{attrs: attrs} } -// fakeObject implements the gcsObject interface for testing. -type fakeObject struct { +// fakeBlob is the shared backing store for one object name; handles reference it +// by pointer so concurrent writers to the same name exercise the precondition. +type fakeBlob struct { mu sync.Mutex name string data []byte @@ -116,44 +291,55 @@ type fakeObject struct { contentType string } -// NewWriter returns a fake writer that stores data in memory. -func (f *fakeObject) newWriter(ctx context.Context) gcsWriter { - f.mu.Lock() - defer f.mu.Unlock() - f.deleted = false // A write operation "undeletes" the object - f.data = nil // Clear existing data - return &fakeWriter{obj: f, buffer: &bytes.Buffer{}} +// fakeObject is a handle to a fakeBlob, optionally carrying a does-not-exist +// precondition (mirrors storage.ObjectHandle.If(Conditions{DoesNotExist: true})). +type fakeObject struct { + blob *fakeBlob + bucket *fakeBucket + mustNotExist bool } -// Attrs returns fake attributes for the object. -func (f *fakeObject) attrs(ctx context.Context) (*storage.ObjectAttrs, error) { - f.mu.Lock() - defer f.mu.Unlock() - if f.deleted || f.data == nil { +// newWriter returns a fake writer that commits data to the blob on Close. +func (o *fakeObject) newWriter(ctx context.Context) gcsWriter { + return &fakeWriter{obj: o, buffer: &bytes.Buffer{}} +} + +func (o *fakeObject) ifNotExist() gcsObject { + return &fakeObject{blob: o.blob, bucket: o.bucket, mustNotExist: true} +} + +// attrs returns fake attributes for the object. +func (o *fakeObject) attrs(ctx context.Context) (*storage.ObjectAttrs, error) { + b := o.blob + b.mu.Lock() + defer b.mu.Unlock() + if b.deleted || b.data == nil { return nil, storage.ErrObjectNotExist } - return &storage.ObjectAttrs{Name: f.name, Created: time.Now(), ContentType: f.contentType}, nil + return &storage.ObjectAttrs{Name: b.name, Created: time.Now(), ContentType: b.contentType}, nil } -// Delete marks the object as deleted in memory. -func (f *fakeObject) delete(ctx context.Context) error { - f.mu.Lock() - defer f.mu.Unlock() - f.deleted = true +// delete marks the object as deleted in memory. +func (o *fakeObject) delete(ctx context.Context) error { + b := o.blob + b.mu.Lock() + defer b.mu.Unlock() + b.deleted = true return nil } -// NewReader returns a reader for the in-memory data. -func (f *fakeObject) newReader(ctx context.Context) (io.ReadCloser, error) { - f.mu.Lock() - defer f.mu.Unlock() - if f.deleted || f.data == nil { +// newReader returns a reader for the in-memory data. +func (o *fakeObject) newReader(ctx context.Context) (io.ReadCloser, error) { + b := o.blob + b.mu.Lock() + defer b.mu.Unlock() + if b.deleted || b.data == nil { return nil, fs.ErrNotExist } - return io.NopCloser(bytes.NewReader(f.data)), nil + return io.NopCloser(bytes.NewReader(b.data)), nil } -// fakeWriter is a helper type to simulate an *storage.Writer +// fakeWriter is a helper type to simulate an *storage.Writer. type fakeWriter struct { obj *fakeObject buffer *bytes.Buffer @@ -164,11 +350,29 @@ func (w *fakeWriter) Write(p []byte) (n int, err error) { return w.buffer.Write(p) } +// Close commits the buffered data under the blob lock, so a conditional write +// that loses the race sees the winner's data and returns HTTP 412 like GCS. A +// bucket-level closeErr hook lets tests force a write failure. func (w *fakeWriter) Close() error { - w.obj.mu.Lock() - defer w.obj.mu.Unlock() - w.obj.data = w.buffer.Bytes() - w.obj.contentType = w.contentType + if bkt := w.obj.bucket; bkt != nil { + bkt.mu.Lock() + bkt.closeCalls++ + forced := bkt.closeErr + bkt.mu.Unlock() + if forced != nil { + return forced + } + } + + b := w.obj.blob + b.mu.Lock() + defer b.mu.Unlock() + if w.obj.mustNotExist && !b.deleted && b.data != nil { + return &googleapi.Error{Code: http.StatusPreconditionFailed, Message: "conditionNotMet"} + } + b.data = w.buffer.Bytes() + b.contentType = w.contentType + b.deleted = false return nil } @@ -177,22 +381,20 @@ func (w *fakeWriter) SetContentType(cType string) { w.contentType = cType } -// fakeObjectIterator is a fake iterator that returns attributes from a slice. -// This type is the key to solving the 'unknown field' error. +// fakeObjectIterator returns attribute snapshots taken at list time. type fakeObjectIterator struct { - objects []*fakeObject - index int + attrs []*storage.ObjectAttrs + index int } -// Next implements the iterator pattern. -// It returns the next object in the slice or an iterator.Done error. +// next returns the next object's attributes or an iterator.Done error. func (i *fakeObjectIterator) next() (*storage.ObjectAttrs, error) { - if i.index >= len(i.objects) { + if i.index >= len(i.attrs) { return nil, iterator.Done } - obj := i.objects[i.index] + a := i.attrs[i.index] i.index++ - return &storage.ObjectAttrs{Name: obj.name, ContentType: obj.contentType}, nil + return a, nil } var ( diff --git a/artifact/gcsartifact/service.go b/artifact/gcsartifact/service.go index cef3a6909..9c8c71641 100644 --- a/artifact/gcsartifact/service.go +++ b/artifact/gcsartifact/service.go @@ -21,20 +21,27 @@ package gcsartifact import ( "context" + "errors" "fmt" "io" "io/fs" "maps" + "math/rand/v2" + "net/http" "slices" "sort" "strconv" "strings" + "time" "cloud.google.com/go/storage" "golang.org/x/sync/errgroup" + "google.golang.org/api/googleapi" "google.golang.org/api/iterator" "google.golang.org/api/option" "google.golang.org/genai" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/adk/v2/artifact" ) @@ -44,6 +51,9 @@ type gcsService struct { bucketName string storageClient gcsClient bucket gcsBucket + // sleep waits for d or until ctx is done; a field so tests can stub out the + // Save retry backoff. Defaults to sleepContext. + sleep func(ctx context.Context, d time.Duration) error } // NewService creates a Google Cloud Storage service for the specified bucket. @@ -58,6 +68,7 @@ func NewService(ctx context.Context, bucketName string, opts ...option.ClientOpt bucketName: bucketName, storageClient: clientWrapper, bucket: clientWrapper.bucket(bucketName), + sleep: sleepContext, } return s, nil } @@ -90,50 +101,139 @@ func buildUserPrefix(appName, userID string) string { return fmt.Sprintf("%s/%s/user/", appName, userID) } -// Save implements [artifact.Service] -func (s *gcsService) Save(ctx context.Context, req *artifact.SaveRequest) (_ *artifact.SaveResponse, err error) { - err = req.Validate() - if err != nil { +const ( + // maxSaveAttempts caps Save's retries; each retry means a concurrent writer + // took the version we picked, so this is how many colliding writers we + // tolerate before returning [ErrVersionConflict]. + maxSaveAttempts = 16 + // saveRetryBaseDelay and saveRetryMaxDelay bound the jittered backoff Save + // waits between those retries. + saveRetryBaseDelay = 10 * time.Millisecond + saveRetryMaxDelay = 1 * time.Second +) + +// ErrVersionConflict is returned by Save when it cannot claim a new version +// within its retry budget because concurrent writers keep taking the version it +// picks. It is always wrapped, so test for it with [errors.Is]; a caller seeing +// it can safely retry the save. +var ErrVersionConflict = errors.New("artifact version conflict") + +// Save implements [artifact.Service]. +// +// GCS can't list-then-write atomically, so versions are assigned optimistically: +// Save writes version max+1 with a does-not-exist precondition and, if another +// writer already took that version, backs off and retries with a fresh one (up +// to [maxSaveAttempts]). It returns [ErrVersionConflict] if that budget is +// exhausted rather than overwriting an existing version. +func (s *gcsService) Save(ctx context.Context, req *artifact.SaveRequest) (*artifact.SaveResponse, error) { + if err := req.Validate(); err != nil { return nil, fmt.Errorf("request validation failed: %w", err) } appName, userID, sessionID, fileName := req.AppName, req.UserID, req.SessionID, req.FileName - newArtifact := req.Part - nextVersion := int64(1) + sleep := s.sleep + if sleep == nil { + sleep = sleepContext + } - // TODO race condition, could use mutex but it's a remote resource so the issue would still occurs - // with multiple consumers, and gcs does not have transactions spanning several operations - response, err := s.versions(ctx, &artifact.VersionsRequest{ - AppName: req.AppName, UserID: req.UserID, SessionID: req.SessionID, FileName: req.FileName, - }) - if err != nil { - return nil, fmt.Errorf("failed to list artifact versions: %w", err) + var lastErr error + for attempt := range maxSaveAttempts { + response, err := s.versions(ctx, &artifact.VersionsRequest{ + AppName: appName, UserID: userID, SessionID: sessionID, FileName: fileName, + }) + if err != nil { + return nil, fmt.Errorf("failed to list artifact versions: %w", err) + } + nextVersion := int64(1) + if len(response.Versions) > 0 { + nextVersion = slices.Max(response.Versions) + 1 + } + + blobName := buildBlobName(appName, userID, sessionID, fileName, nextVersion) + obj := s.bucket.object(blobName).ifNotExist() + err = writeArtifact(ctx, obj, req.Part) + if err == nil { + return &artifact.SaveResponse{Version: nextVersion}, nil + } + if !isPreconditionFailed(err) { + return nil, fmt.Errorf("failed to save artifact: %w", err) + } + lastErr = err + // Lost the race for this version. Back off (skip after the final attempt) + // and retry with a fresh version. + if attempt < maxSaveAttempts-1 { + if err := sleep(ctx, backoffDelay(attempt)); err != nil { + return nil, err + } + } + } + return nil, fmt.Errorf("failed to save artifact %q after %d attempts (last error: %v): %w", fileName, maxSaveAttempts, lastErr, ErrVersionConflict) +} + +// backoffDelay returns a full-jitter backoff for the given zero-based retry +// attempt, capped at saveRetryMaxDelay. +func backoffDelay(attempt int) time.Duration { + d := saveRetryMaxDelay + if attempt < 63 { + if scaled := saveRetryBaseDelay << attempt; scaled > 0 && scaled < d { + d = scaled + } + } + return time.Duration(rand.Int64N(int64(d) + 1)) +} + +// sleepContext waits for d or until ctx is done, returning ctx.Err() if ctx is +// done first. +func sleepContext(ctx context.Context, d time.Duration) error { + if err := ctx.Err(); err != nil { + return err } - if len(response.Versions) > 0 { - nextVersion = slices.Max(response.Versions) + 1 + if d <= 0 { + return nil + } + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + return nil } +} - blobName := buildBlobName(appName, userID, sessionID, fileName, nextVersion) - writer := s.bucket.object(blobName).newWriter(ctx) +// writeArtifact streams part to obj. A precondition on obj surfaces as an error +// from Close, not Write. +func writeArtifact(ctx context.Context, obj gcsObject, part *genai.Part) (err error) { + writer := obj.newWriter(ctx) defer func() { if closeErr := writer.Close(); closeErr != nil && err == nil { err = fmt.Errorf("failed to close blob writer: %w", closeErr) } }() - if newArtifact.InlineData != nil { - writer.SetContentType(newArtifact.InlineData.MIMEType) - if _, err := writer.Write(newArtifact.InlineData.Data); err != nil { - return nil, fmt.Errorf("failed to write blob to GCS: %w", err) + if part.InlineData != nil { + writer.SetContentType(part.InlineData.MIMEType) + if _, err := writer.Write(part.InlineData.Data); err != nil { + return fmt.Errorf("failed to write blob to GCS: %w", err) } } else { writer.SetContentType("text/plain") - if _, err := writer.Write([]byte(newArtifact.Text)); err != nil { - return nil, fmt.Errorf("failed to write text to GCS: %w", err) + if _, err := writer.Write([]byte(part.Text)); err != nil { + return fmt.Errorf("failed to write text to GCS: %w", err) } } + return nil +} - return &artifact.SaveResponse{Version: nextVersion}, nil +// isPreconditionFailed reports whether err is a GCS precondition failure: HTTP +// 412 (*googleapi.Error) on the JSON API, or codes.FailedPrecondition on the +// gRPC transport. Both status.FromError and errors.As unwrap wrapped errors. +func isPreconditionFailed(err error) bool { + var apiErr *googleapi.Error + if errors.As(err, &apiErr) { + return apiErr.Code == http.StatusPreconditionFailed + } + return status.Code(err) == codes.FailedPrecondition } // Delete implements [artifact.Service]