Attach team_id metadata to uploaded storage objects (#2533)

* Attach team / build lineage labels to uploaded storage objects

Adds custom object metadata to GCS/S3 uploads from snapshot/template
builds so artifacts can be filtered after upload. Local filesystem
backend is a no-op.

Per-object labels (every artifact in {buildID}/):
  - team_id
  - root_build_id

Per-build labels (only on {buildID}/metadata.json):
  - build_kind  (template_layer | sandbox_pause | sandbox_checkpoint)
  - parent_build_id

The Blob.Put / Seekable.StoreFile interfaces gain a variadic
PutOption, with WithMetadata as the first option. Snapshot.Upload
and TemplateBuild take a SnapshotUploadMetadata{Common, MetadataOnly}
to control what goes on every object vs. only on metadata.json.

Option types live in a new storageopts subpackage so generated mocks
can reference them without an import cycle.

* Address review: BaseBuildId for snapshot root, log source-meta error, fix nlreturn

- root_build_id for sandbox pause/checkpoint now uses
  snapshot.MemfileDiffHeader.Metadata.BaseBuildId (the diff-chain root /
  originating template), preserved across pause -> resume -> pause.
  Falls back to the snapshot's own build ID if the header is missing.
  Layer builds keep BuildContext.Template.BuildID since it's the
  user-requested top-level build (not redundant with the path).
- Log a warning when localTemplate.Metadata() fails in BuildLayer
  instead of silently dropping parent_build_id.
- Restore blank lines before returns in template_build.go to satisfy
  the nlreturn linter.

* Use BaseBuildId from header for root_build_id everywhere; omit when unknown

root_build_id should be the very first build ID in the lineage chain.
Now derived consistently from the memfile header's Metadata.BaseBuildId
in all three upload paths:

- Sandbox pause/checkpoint: read from snapshot.MemfileDiffHeader
- Template layer build: read from snapshot.MemfileDiffHeader after pause
- Optimize phase rewrite: read from the cached template's memfile header

If the header doesn't carry a usable base build ID, root_build_id is
omitted from the labels rather than falling back to a wrong value.

* Strip lineage labels: keep only team_id on uploaded objects

Author: ValentaTomas

packages/orchestrator/cmd/copy-build/main.go

@@ -108,7 +108,7 @@ func (o *osFileBlob) Exists(_ context.Context) (bool, error) {
 	return true, nil
 }
 
-func (o *osFileBlob) Put(_ context.Context, _ []byte) error {
+func (o *osFileBlob) Put(_ context.Context, _ []byte, _ ...storage.PutOption) error {
 	return errors.New("not implemented")
 }
 

packages/orchestrator/cmd/resume-build/main.go

@@ -641,13 +641,13 @@ func (r *runner) pauseOnce(ctx context.Context, opts pauseOptions, verbose bool)
 		paths := storage.Paths{BuildID: opts.newBuildID}
 		if opts.isRemoteStorage {
 			fmt.Println("📤 Uploading snapshot...")
-			if err := snapshot.Upload(ctx, r.storage, paths); err != nil {
+			if err := snapshot.Upload(ctx, r.storage, paths, nil); err != nil {
 				return timings, fmt.Errorf("failed to upload snapshot: %w", err)
 			}
 			fmt.Println("✅ Snapshot uploaded successfully")
 		} else {
 			fmt.Println("💾 Saving snapshot to local storage...")
-			if err := snapshot.Upload(ctx, r.storage, paths); err != nil {
+			if err := snapshot.Upload(ctx, r.storage, paths, nil); err != nil {
 				return timings, fmt.Errorf("failed to save snapshot: %w", err)
 			}
 			fmt.Println("✅ Snapshot saved successfully")
@@ -864,7 +864,7 @@ func (r *runner) collectAndUploadPrefetch(ctx context.Context, opts pauseOptions
 		Memory: mapping,
 	})
 
-	if err := metadata.UploadMetadata(ctx, r.storage, updatedMeta); err != nil {
+	if err := metadata.UploadMetadata(ctx, r.storage, updatedMeta, nil); err != nil {
 		return fmt.Errorf("upload metadata: %w", err)
 	}
 

packages/orchestrator/pkg/sandbox/snapshot.go

@@ -21,10 +21,13 @@ type Snapshot struct {
 	cleanup *Cleanup
 }
 
+// Upload writes snapshot artifacts to persistence under paths. objectMetadata
+// is attached to every uploaded object; pass nil to skip.
 func (s *Snapshot) Upload(
 	ctx context.Context,
 	persistence storage.StorageProvider,
 	paths storage.Paths,
+	objectMetadata storage.ObjectMetadata,
 ) error {
 	var memfilePath *string
 	switch r := s.MemfileDiff.(type) {
@@ -55,6 +58,7 @@ func (s *Snapshot) Upload(
 		s.RootfsDiffHeader,
 		persistence,
 		paths,
+		objectMetadata,
 	)
 
 	if err := templateBuild.Upload(

packages/orchestrator/pkg/sandbox/template/peerclient/blob.go

@@ -76,14 +76,14 @@ func (b *peerBlob) Exists(ctx context.Context) (bool, error) {
 	)
 }
 
-func (b *peerBlob) Put(ctx context.Context, data []byte) error {
+func (b *peerBlob) Put(ctx context.Context, data []byte, opts ...storage.PutOption) error {
 	// Writes always go to the base provider (GCS/S3); the peer is read-only.
 	fallback, err := b.getOrOpenBase(ctx)
 	if err != nil {
 		return err
 	}
 
-	return fallback.Put(ctx, data)
+	return fallback.Put(ctx, data, opts...)
 }
 
 // openPeerBlobStream opens a GetBuildBlob stream, checks peer availability,

packages/orchestrator/pkg/sandbox/template/peerclient/seekable.go

@@ -120,14 +120,14 @@ func (s *peerSeekable) OpenRangeReader(ctx context.Context, off, length int64) (
 	)
 }
 
-func (s *peerSeekable) StoreFile(ctx context.Context, path string) error {
+func (s *peerSeekable) StoreFile(ctx context.Context, path string, opts ...storage.PutOption) error {
 	// Writes always go to the base provider (GCS/S3); the peer is read-only.
 	fallback, err := s.getOrOpenBase(ctx)
 	if err != nil {
 		return err
 	}
 
-	return fallback.StoreFile(ctx, path)
+	return fallback.StoreFile(ctx, path, opts...)
 }
 
 // openPeerSeekableStream opens a ReadAtBuildSeekable stream, checks peer availability,

packages/orchestrator/pkg/sandbox/template_build.go

@@ -13,17 +13,25 @@ import (
 )
 
 type TemplateBuild struct {
-	paths       storage.Paths
-	persistence storage.StorageProvider
+	paths          storage.Paths
+	persistence    storage.StorageProvider
+	objectMetadata storage.ObjectMetadata
 
 	memfileHeader *headers.Header
 	rootfsHeader  *headers.Header
 }
 
-func NewTemplateBuild(memfileHeader *headers.Header, rootfsHeader *headers.Header, persistence storage.StorageProvider, paths storage.Paths) *TemplateBuild {
+func NewTemplateBuild(
+	memfileHeader *headers.Header,
+	rootfsHeader *headers.Header,
+	persistence storage.StorageProvider,
+	paths storage.Paths,
+	objectMetadata storage.ObjectMetadata,
+) *TemplateBuild {
 	return &TemplateBuild{
-		persistence: persistence,
-		paths:       paths,
+		persistence:    persistence,
+		paths:          paths,
+		objectMetadata: objectMetadata,
 
 		memfileHeader: memfileHeader,
 		rootfsHeader:  rootfsHeader,
@@ -39,6 +47,14 @@ func (t *TemplateBuild) Remove(ctx context.Context) error {
 	return nil
 }
 
+func (t *TemplateBuild) putOpts() []storage.PutOption {
+	if len(t.objectMetadata) == 0 {
+		return nil
+	}
+
+	return []storage.PutOption{storage.WithMetadata(t.objectMetadata)}
+}
+
 func (t *TemplateBuild) uploadMemfileHeader(ctx context.Context, h *headers.Header) error {
 	object, err := t.persistence.OpenBlob(ctx, t.paths.MemfileHeader(), storage.MemfileHeaderObjectType)
 	if err != nil {
@@ -50,7 +66,7 @@ func (t *TemplateBuild) uploadMemfileHeader(ctx context.Context, h *headers.Head
 		return fmt.Errorf("error when serializing memfile header: %w", err)
 	}
 
-	err = object.Put(ctx, serialized)
+	err = object.Put(ctx, serialized, t.putOpts()...)
 	if err != nil {
 		return fmt.Errorf("error when uploading memfile header: %w", err)
 	}
@@ -64,7 +80,7 @@ func (t *TemplateBuild) uploadMemfile(ctx context.Context, memfilePath string) e
 		return err
 	}
 
-	err = object.StoreFile(ctx, memfilePath)
+	err = object.StoreFile(ctx, memfilePath, t.putOpts()...)
 	if err != nil {
 		return fmt.Errorf("error when uploading memfile: %w", err)
 	}
@@ -83,7 +99,7 @@ func (t *TemplateBuild) uploadRootfsHeader(ctx context.Context, h *headers.Heade
 		return fmt.Errorf("error when serializing memfile header: %w", err)
 	}
 
-	err = object.Put(ctx, serialized)
+	err = object.Put(ctx, serialized, t.putOpts()...)
 	if err != nil {
 		return fmt.Errorf("error when uploading memfile header: %w", err)
 	}
@@ -97,7 +113,7 @@ func (t *TemplateBuild) uploadRootfs(ctx context.Context, rootfsPath string) err
 		return err
 	}
 
-	err = object.StoreFile(ctx, rootfsPath)
+	err = object.StoreFile(ctx, rootfsPath, t.putOpts()...)
 	if err != nil {
 		return fmt.Errorf("error when uploading rootfs: %w", err)
 	}
@@ -112,7 +128,7 @@ func (t *TemplateBuild) uploadSnapfile(ctx context.Context, path string) error {
 		return err
 	}
 
-	if err = uploadFileAsBlob(ctx, object, path); err != nil {
+	if err = uploadFileAsBlob(ctx, object, path, t.putOpts()...); err != nil {
 		return fmt.Errorf("error when uploading snapfile: %w", err)
 	}
 
@@ -126,14 +142,14 @@ func (t *TemplateBuild) uploadMetadata(ctx context.Context, path string) error {
 		return err
 	}
 
-	if err := uploadFileAsBlob(ctx, object, path); err != nil {
+	if err := uploadFileAsBlob(ctx, object, path, t.putOpts()...); err != nil {
 		return fmt.Errorf("error when uploading metadata: %w", err)
 	}
 
 	return nil
 }
 
-func uploadFileAsBlob(ctx context.Context, b storage.Blob, path string) error {
+func uploadFileAsBlob(ctx context.Context, b storage.Blob, path string, opts ...storage.PutOption) error {
 	f, err := os.Open(path)
 	if err != nil {
 		return fmt.Errorf("failed to open file %s: %w", path, err)
@@ -145,7 +161,7 @@ func uploadFileAsBlob(ctx context.Context, b storage.Blob, path string) error {
 		return fmt.Errorf("failed to read file %s: %w", path, err)
 	}
 
-	err = b.Put(ctx, data)
+	err = b.Put(ctx, data, opts...)
 	if err != nil {
 		return fmt.Errorf("failed to write data to object: %w", err)
 	}

packages/orchestrator/pkg/server/sandboxes.go

@@ -667,7 +667,7 @@ func (s *Server) Checkpoint(ctx context.Context, in *orchestrator.SandboxCheckpo
 		defer cancel()
 		defer res.completeUpload(uploadCtx)
 
-		if err := res.snapshot.Upload(uploadCtx, s.persistence, res.paths); err != nil {
+		if err := res.snapshot.Upload(uploadCtx, s.persistence, res.paths, res.objectMetadata); err != nil {
 			telemetry.ReportCriticalError(ctx, "error uploading snapshot for checkpoint", err, telemetry.WithSandboxID(in.GetSandboxId()))
 
 			s.sandboxFactory.Sandboxes.MarkStopping(ctx, resumedSbx.Runtime.SandboxID, resumedSbx.LifecycleID)
@@ -721,6 +721,7 @@ type snapshotResult struct {
 	meta           metadata.Template
 	snapshot       *sandbox.Snapshot
 	paths          storage.Paths
+	objectMetadata storage.ObjectMetadata
 	completeUpload func(ctx context.Context)
 }
 
@@ -765,6 +766,9 @@ func (s *Server) snapshotAndCacheSandbox(
 	telemetry.ReportEvent(ctx, "added snapshot to template cache")
 
 	paths := storage.Paths{BuildID: meta.Template.BuildID}
+	objectMetadata := storage.ObjectMetadata{
+		storage.ObjectMetadataTeamID: sbx.Runtime.TeamID,
+	}
 
 	// Register in Redis so other orchestrators can find us for peer routing.
 	if s.featureFlags.BoolFlag(ctx, featureflags.PeerToPeerChunkTransferFlag) {
@@ -786,6 +790,7 @@ func (s *Server) snapshotAndCacheSandbox(
 			meta:           meta,
 			snapshot:       snapshot,
 			paths:          paths,
+			objectMetadata: objectMetadata,
 			completeUpload: completeUpload,
 		}, nil
 	}
@@ -794,6 +799,7 @@ func (s *Server) snapshotAndCacheSandbox(
 		meta:           meta,
 		snapshot:       snapshot,
 		paths:          paths,
+		objectMetadata: objectMetadata,
 		completeUpload: func(context.Context) {},
 	}, nil
 }
@@ -808,7 +814,7 @@ func (s *Server) uploadSnapshotAsync(ctx context.Context, sbx *sandbox.Sandbox,
 		defer cancel()
 		defer res.completeUpload(ctx)
 
-		err := res.snapshot.Upload(ctx, s.persistence, res.paths)
+		err := res.snapshot.Upload(ctx, s.persistence, res.paths, res.objectMetadata)
 		if err != nil {
 			sbxlogger.I(sbx).Error(ctx, "error uploading snapshot files", zap.Error(err))
 

packages/orchestrator/pkg/template/build/layer/layer_executor.go

@@ -293,6 +293,9 @@ func (lb *LayerExecutor) PauseAndUpload(
 			ctx,
 			lb.templateStorage,
 			storage.Paths{BuildID: meta.Template.BuildID},
+			storage.ObjectMetadata{
+				storage.ObjectMetadataTeamID: lb.BuildContext.Config.TeamID,
+			},
 		)
 		if err != nil {
 			return fmt.Errorf("error uploading snapshot: %w", err)

packages/orchestrator/pkg/template/build/phases/optimize/builder.go

@@ -256,7 +256,9 @@ func (pb *OptimizeBuilder) runSandboxAndCollectPrefetch(
 
 // updateMetadata updates the template metadata in storages.
 func (pb *OptimizeBuilder) updateMetadata(ctx context.Context, t metadata.Template) error {
-	err := metadata.UploadMetadata(ctx, pb.templateStorage, t)
+	err := metadata.UploadMetadata(ctx, pb.templateStorage, t, storage.ObjectMetadata{
+		storage.ObjectMetadataTeamID: pb.BuildContext.Config.TeamID,
+	})
 	if err != nil {
 		return err
 	}

packages/orchestrator/pkg/template/metadata/prefetch.go

@@ -43,8 +43,9 @@ func PrefetchEntriesToMapping(entries []block.PrefetchBlockEntry, blockSize int6
 	}
 }
 
-// UploadMetadata uploads the template metadata to storage.
-func UploadMetadata(ctx context.Context, persistence storage.StorageProvider, t Template) error {
+// UploadMetadata uploads the template metadata to storage. objectMetadata
+// is attached to the object; pass nil to skip.
+func UploadMetadata(ctx context.Context, persistence storage.StorageProvider, t Template, objectMetadata storage.ObjectMetadata) error {
 	ctx, span := tracer.Start(ctx, "upload-metadata")
 	defer span.End()
 
@@ -60,7 +61,12 @@ func UploadMetadata(ctx context.Context, persistence storage.StorageProvider, t
 		return fmt.Errorf("failed to serialize metadata: %w", err)
 	}
 
-	err = object.Put(ctx, metaBytes)
+	var opts []storage.PutOption
+	if len(objectMetadata) > 0 {
+		opts = append(opts, storage.WithMetadata(objectMetadata))
+	}
+
+	err = object.Put(ctx, metaBytes, opts...)
 	if err != nil {
 		return fmt.Errorf("failed to write metadata: %w", err)
 	}