fix(distributed): release staged request inputs

Ephemeral inputs remained on workers after inference completed. Release each exact key after synchronous and streaming calls, including partial staging failures.

Use a bounded cleanup context so caller cancellation cannot suppress release. Preserve caller requests and backend results when cleanup fails.

Assisted-by: Codex:gpt-6

core/services/nodes/file_staging_client.go

@@ -19,6 +19,8 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
+const stagedInputReleaseTimeout = 30 * time.Second
+
 // FileStagingClient wraps a grpc.Backend to transparently handle file transfer
 // for distributed mode. Input files are staged on the backend node before the
 // gRPC call. Output files are retrieved from the backend after the call.
@@ -49,21 +51,64 @@ func NewFileStagingClient(inner grpc.Backend, stager FileStager, nodeID string)
 
 // requestID generates a unique ID for ephemeral file keys.
 func requestID() string {
-	return uuid.New().String()[:8]
+	return uuid.NewString()
+}
+
+type stagedInputLifecycle struct {
+	client    *FileStagingClient
+	requestID string
+	keys      []string
+	seen      map[string]struct{}
+}
+
+func (f *FileStagingClient) newStagedInputLifecycle() *stagedInputLifecycle {
+	return &stagedInputLifecycle{
+		client:    f,
+		requestID: requestID(),
+		keys:      []string{},
+		seen:      map[string]struct{}{},
+	}
+}
+
+func (l *stagedInputLifecycle) track(key string) {
+	if _, ok := l.seen[key]; ok {
+		return
+	}
+	l.seen[key] = struct{}{}
+	l.keys = append(l.keys, key)
+}
+
+func (l *stagedInputLifecycle) release() {
+	if len(l.keys) == 0 {
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), stagedInputReleaseTimeout)
+	defer cancel()
+	for _, key := range l.keys {
+		if err := l.client.stager.ReleaseRemote(ctx, l.client.nodeID, key); err != nil {
+			xlog.Warn("Failed to release staged input", "node", l.client.nodeID, "key", key, "error", err)
+		}
+	}
 }
 
 // stageInputFile uploads a local file to the remote node via the FileStager.
-// Returns the remote-local path and the ephemeral key.
-func (f *FileStagingClient) stageInputFile(ctx context.Context, reqID, localPath, category string) (string, string, error) {
+func (f *FileStagingClient) stageInputFile(
+	ctx context.Context,
+	lifecycle *stagedInputLifecycle,
+	localPath,
+	category string,
+) (string, error) {
 	basename := filepath.Base(localPath)
-	key := storage.EphemeralKey(reqID, category, basename)
+	key := storage.EphemeralKey(lifecycle.requestID, category, basename)
+	lifecycle.track(key)
 
 	remotePath, err := f.stager.EnsureRemote(ctx, f.nodeID, localPath, key)
 	if err != nil {
-		return "", "", fmt.Errorf("staging input file: %w", err)
+		return "", fmt.Errorf("staging input file: %w", err)
 	}
 
-	return remotePath, key, nil
+	return remotePath, nil
 }
 
 // retrieveOutputFile retrieves an output file from the backend to a local path.
@@ -101,23 +146,29 @@ func (f *FileStagingClient) translateModelPath(frontendPath string) string {
 }
 
 func (f *FileStagingClient) Predict(ctx context.Context, in *pb.PredictOptions, opts ...ggrpc.CallOption) (*pb.Reply, error) {
-	reqID := requestID()
-	in, _ = f.stageMultimodalInputs(ctx, reqID, in)
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.PredictOptions)
+	in = f.stageMultimodalInputs(ctx, lifecycle, in)
 	return f.Backend.Predict(ctx, in, opts...)
 }
 
 func (f *FileStagingClient) PredictStream(ctx context.Context, in *pb.PredictOptions, fn func(reply *pb.Reply), opts ...ggrpc.CallOption) error {
-	reqID := requestID()
-	in, _ = f.stageMultimodalInputs(ctx, reqID, in)
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.PredictOptions)
+	in = f.stageMultimodalInputs(ctx, lifecycle, in)
 	return f.Backend.PredictStream(ctx, in, fn, opts...)
 }
 
 func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.GenerateImageRequest)
 
 	// Stage input source image if present
 	if in.Src != "" && isFilePath(in.Src) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.Src, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging image src: %w", err)
 		}
@@ -127,7 +178,7 @@ func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateIm
 	// Stage reference images
 	for i, img := range in.RefImages {
 		if isFilePath(img) {
-			backendPath, _, err := f.stageInputFile(ctx, reqID, img, "inputs")
+			backendPath, err := f.stageInputFile(ctx, lifecycle, img, "inputs")
 			if err != nil {
 				return nil, fmt.Errorf("staging ref image: %w", err)
 			}
@@ -161,25 +212,27 @@ func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateIm
 }
 
 func (f *FileStagingClient) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.GenerateVideoRequest)
 
 	// Stage start/end images and optional audio conditioning.
 	if in.StartImage != "" && isFilePath(in.StartImage) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.StartImage, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.StartImage, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging start image: %w", err)
 		}
 		in.StartImage = backendPath
 	}
 	if in.EndImage != "" && isFilePath(in.EndImage) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.EndImage, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.EndImage, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging end image: %w", err)
 		}
 		in.EndImage = backendPath
 	}
 	if in.Audio != "" && isFilePath(in.Audio) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.Audio, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Audio, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging video audio: %w", err)
 		}
@@ -211,11 +264,13 @@ func (f *FileStagingClient) GenerateVideo(ctx context.Context, in *pb.GenerateVi
 }
 
 func (f *FileStagingClient) Generate3D(ctx context.Context, in *pb.Generate3DRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.Generate3DRequest)
 
 	// Stage the conditioning image or existing GLB used by 3D post-processing.
 	if in.Src != "" && isFilePath(in.Src) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.Src, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging 3D input asset: %w", err)
 		}
@@ -247,7 +302,9 @@ func (f *FileStagingClient) Generate3D(ctx context.Context, in *pb.Generate3DReq
 }
 
 func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.TTSRequest)
 
 	// Translate model path from frontend to remote worker path.
 	// The model and its companion files (e.g. .onnx.json) were already staged
@@ -258,7 +315,7 @@ func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ...
 	// Voice may be a named backend speaker or a request-scoped reference WAV.
 	// Only path-shaped values are staged; speaker IDs pass through unchanged.
 	if in.Voice != "" && isFilePath(in.Voice) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.Voice, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Voice, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging TTS voice reference: %w", err)
 		}
@@ -290,14 +347,16 @@ func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ...
 }
 
 func (f *FileStagingClient) TTSStream(ctx context.Context, in *pb.TTSRequest, fn func(*pb.Reply), opts ...ggrpc.CallOption) error {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.TTSRequest)
 
 	// Translate model path from frontend to remote worker path (same as TTS above)
 	if in.Model != "" && isFilePath(in.Model) {
 		in.Model = f.translateModelPath(in.Model)
 	}
 	if in.Voice != "" && isFilePath(in.Voice) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.Voice, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Voice, "inputs")
 		if err != nil {
 			return fmt.Errorf("staging streaming TTS voice reference: %w", err)
 		}
@@ -308,11 +367,13 @@ func (f *FileStagingClient) TTSStream(ctx context.Context, in *pb.TTSRequest, fn
 }
 
 func (f *FileStagingClient) SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.SoundGenerationRequest)
 
 	// Stage input source
 	if in.Src != nil && *in.Src != "" && isFilePath(*in.Src) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, *in.Src, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, *in.Src, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging sound src: %w", err)
 		}
@@ -344,24 +405,27 @@ func (f *FileStagingClient) SoundGeneration(ctx context.Context, in *pb.SoundGen
 }
 
 func (f *FileStagingClient) SoundDetection(ctx context.Context, in *pb.SoundDetectionRequest, opts ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) {
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.SoundDetectionRequest)
 	if in.Src != "" && isFilePath(in.Src) {
-		backendPath, _, err := f.stageInputFile(ctx, requestID(), in.Src, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging audio for sound detection: %w", err)
 		}
-		// Keep the frontend path available if the caller retries on another node.
-		in = proto.Clone(in).(*pb.SoundDetectionRequest)
 		in.Src = backendPath
 	}
 	return f.Backend.SoundDetection(ctx, in, opts...)
 }
 
 func (f *FileStagingClient) AudioTranscription(ctx context.Context, in *pb.TranscriptRequest, opts ...ggrpc.CallOption) (*pb.TranscriptResult, error) {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.TranscriptRequest)
 
 	// Stage input audio file
 	if in.Dst != "" && isFilePath(in.Dst) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.Dst, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Dst, "inputs")
 		if err != nil {
 			return nil, fmt.Errorf("staging audio for transcription: %w", err)
 		}
@@ -372,11 +436,13 @@ func (f *FileStagingClient) AudioTranscription(ctx context.Context, in *pb.Trans
 }
 
 func (f *FileStagingClient) AudioTranscriptionStream(ctx context.Context, in *pb.TranscriptRequest, fn func(chunk *pb.TranscriptStreamResponse), opts ...ggrpc.CallOption) error {
-	reqID := requestID()
+	lifecycle := f.newStagedInputLifecycle()
+	defer lifecycle.release()
+	in = proto.Clone(in).(*pb.TranscriptRequest)
 
 	// Stage input audio file
 	if in.Dst != "" && isFilePath(in.Dst) {
-		backendPath, _, err := f.stageInputFile(ctx, reqID, in.Dst, "inputs")
+		backendPath, err := f.stageInputFile(ctx, lifecycle, in.Dst, "inputs")
 		if err != nil {
 			return fmt.Errorf("staging audio for transcription stream: %w", err)
 		}
@@ -458,26 +524,33 @@ func (f *FileStagingClient) QuantizationProgress(ctx context.Context, in *pb.Qua
 
 // stageMultimodalInputs stages Images, Videos, Audios fields in PredictOptions
 // if they are file paths (not base64 or URLs).
-func (f *FileStagingClient) stageMultimodalInputs(ctx context.Context, reqID string, in *pb.PredictOptions) (*pb.PredictOptions, []string) {
-	var keys []string
-	in.Images = f.stagePathSlice(ctx, reqID, in.Images, "inputs", &keys)
-	in.Videos = f.stagePathSlice(ctx, reqID, in.Videos, "inputs", &keys)
-	in.Audios = f.stagePathSlice(ctx, reqID, in.Audios, "inputs", &keys)
-	return in, keys
+func (f *FileStagingClient) stageMultimodalInputs(
+	ctx context.Context,
+	lifecycle *stagedInputLifecycle,
+	in *pb.PredictOptions,
+) *pb.PredictOptions {
+	in.Images = f.stagePathSlice(ctx, lifecycle, in.Images, "inputs")
+	in.Videos = f.stagePathSlice(ctx, lifecycle, in.Videos, "inputs")
+	in.Audios = f.stagePathSlice(ctx, lifecycle, in.Audios, "inputs")
+	return in
 }
 
-func (f *FileStagingClient) stagePathSlice(ctx context.Context, reqID string, paths []string, category string, keys *[]string) []string {
+func (f *FileStagingClient) stagePathSlice(
+	ctx context.Context,
+	lifecycle *stagedInputLifecycle,
+	paths []string,
+	category string,
+) []string {
 	result := make([]string, len(paths))
 	for i, p := range paths {
 		if isFilePath(p) {
-			backendPath, key, err := f.stageInputFile(ctx, reqID, p, category)
+			backendPath, err := f.stageInputFile(ctx, lifecycle, p, category)
 			if err != nil {
 				xlog.Warn("Failed to stage multimodal file, passing through", "path", p, "error", err)
 				result[i] = p
 				continue
 			}
 			result[i] = backendPath
-			*keys = append(*keys, key)
 		} else {
 			result[i] = p
 		}

core/services/nodes/file_staging_lifecycle_test.go

@@ -0,0 +1,288 @@
+package nodes
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	grpc "github.com/mudler/LocalAI/pkg/grpc"
+	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
+	ggrpc "google.golang.org/grpc"
+	"google.golang.org/protobuf/proto"
+)
+
+const fullUUIDPattern = `[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`
+
+type lifecycleStager struct {
+	fakeFileStager
+	ensureErr          error
+	releaseErr         error
+	releasedKeys       []string
+	releaseCtxErr      []error
+	releaseHasDeadline []bool
+	releaseDeadlines   []time.Time
+}
+
+func (s *lifecycleStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
+	s.fakeFileStager.EnsureRemote(ctx, nodeID, localPath, key)
+	if s.ensureErr != nil {
+		return "", s.ensureErr
+	}
+	return "/remote/" + key, nil
+}
+
+func (s *lifecycleStager) ReleaseRemote(ctx context.Context, _ string, key string) error {
+	s.releasedKeys = append(s.releasedKeys, key)
+	s.releaseCtxErr = append(s.releaseCtxErr, ctx.Err())
+	deadline, ok := ctx.Deadline()
+	s.releaseHasDeadline = append(s.releaseHasDeadline, ok)
+	s.releaseDeadlines = append(s.releaseDeadlines, deadline)
+	return s.releaseErr
+}
+
+type lifecycleBackend struct {
+	grpc.Backend
+	predictResult *pb.Reply
+	predictErr    error
+	streamBlock   <-chan struct{}
+	streamStarted chan<- struct{}
+}
+
+func (b *lifecycleBackend) Predict(_ context.Context, _ *pb.PredictOptions, _ ...ggrpc.CallOption) (*pb.Reply, error) {
+	if b.predictResult == nil {
+		b.predictResult = &pb.Reply{}
+	}
+	return b.predictResult, b.predictErr
+}
+
+func (b *lifecycleBackend) PredictStream(_ context.Context, _ *pb.PredictOptions, _ func(*pb.Reply), _ ...ggrpc.CallOption) error {
+	if b.streamStarted != nil {
+		b.streamStarted <- struct{}{}
+	}
+	if b.streamBlock != nil {
+		<-b.streamBlock
+	}
+	return nil
+}
+
+func (b *lifecycleBackend) GenerateImage(_ context.Context, _ *pb.GenerateImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
+	return &pb.Result{Success: true}, nil
+}
+
+func (b *lifecycleBackend) GenerateVideo(_ context.Context, _ *pb.GenerateVideoRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
+	return &pb.Result{Success: true}, nil
+}
+
+func (b *lifecycleBackend) Generate3D(_ context.Context, _ *pb.Generate3DRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
+	return &pb.Result{Success: true}, nil
+}
+
+func (b *lifecycleBackend) TTS(_ context.Context, _ *pb.TTSRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
+	return &pb.Result{Success: true}, nil
+}
+
+func (b *lifecycleBackend) TTSStream(_ context.Context, _ *pb.TTSRequest, _ func(*pb.Reply), _ ...ggrpc.CallOption) error {
+	return nil
+}
+
+func (b *lifecycleBackend) SoundGeneration(_ context.Context, _ *pb.SoundGenerationRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
+	return &pb.Result{Success: true}, nil
+}
+
+func (b *lifecycleBackend) SoundDetection(_ context.Context, _ *pb.SoundDetectionRequest, _ ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) {
+	return &pb.SoundDetectionResponse{}, nil
+}
+
+func (b *lifecycleBackend) AudioTranscription(_ context.Context, _ *pb.TranscriptRequest, _ ...ggrpc.CallOption) (*pb.TranscriptResult, error) {
+	return &pb.TranscriptResult{}, nil
+}
+
+func (b *lifecycleBackend) AudioTranscriptionStream(_ context.Context, _ *pb.TranscriptRequest, _ func(*pb.TranscriptStreamResponse), _ ...ggrpc.CallOption) error {
+	return nil
+}
+
+var _ = Describe("FileStagingClient request lifecycle", func() {
+	It("uses a full UUID for ephemeral request keys", func() {
+		Expect(requestID()).To(MatchRegexp(`^` + fullUUIDPattern + `$`))
+	})
+
+	It("releases every staged key and preserves caller requests", func(ctx SpecContext) {
+		tests := []struct {
+			name     string
+			keyCount int
+			invoke   func(*FileStagingClient) proto.Message
+		}{
+			{name: "predict", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}, Videos: []string{"/tmp/video.mp4"}, Audios: []string{"/tmp/audio.wav"}}
+				original := proto.Clone(request)
+				_, err := client.Predict(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "predict stream", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}}
+				original := proto.Clone(request)
+				Expect(client.PredictStream(ctx, request, func(*pb.Reply) {})).To(Succeed())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "image generation", keyCount: 2, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.GenerateImageRequest{Src: "/tmp/source.png", RefImages: []string{"/tmp/reference.png"}}
+				original := proto.Clone(request)
+				_, err := client.GenerateImage(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "video generation", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.GenerateVideoRequest{StartImage: "/tmp/start.png", EndImage: "/tmp/end.png", Audio: "/tmp/audio.wav"}
+				original := proto.Clone(request)
+				_, err := client.GenerateVideo(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "3D generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.Generate3DRequest{Src: "/tmp/source.glb"}
+				original := proto.Clone(request)
+				_, err := client.Generate3D(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.TTSRequest{Voice: "/tmp/voice.wav"}
+				original := proto.Clone(request)
+				_, err := client.TTS(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "streaming TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.TTSRequest{Voice: "/tmp/voice.wav"}
+				original := proto.Clone(request)
+				Expect(client.TTSStream(ctx, request, func(*pb.Reply) {})).To(Succeed())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "sound generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				source := "/tmp/source.wav"
+				request := &pb.SoundGenerationRequest{Src: &source}
+				original := proto.Clone(request)
+				_, err := client.SoundGeneration(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "sound detection", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.SoundDetectionRequest{Src: "/tmp/source.wav"}
+				original := proto.Clone(request)
+				_, err := client.SoundDetection(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"}
+				original := proto.Clone(request)
+				_, err := client.AudioTranscription(ctx, request)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+			{name: "streaming transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
+				request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"}
+				original := proto.Clone(request)
+				Expect(client.AudioTranscriptionStream(ctx, request, func(*pb.TranscriptStreamResponse) {})).To(Succeed())
+				Expect(proto.Equal(request, original)).To(BeTrue())
+				return request
+			}},
+		}
+
+		for _, test := range tests {
+			By(test.name)
+			stager := &lifecycleStager{}
+			client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
+			test.invoke(client)
+			Expect(stager.ensureCalls).To(HaveLen(test.keyCount))
+			Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
+			Expect(stager.releaseCtxErr).To(ConsistOf(make([]error, test.keyCount)))
+			Expect(stager.releaseHasDeadline).To(HaveLen(test.keyCount))
+			for _, hasDeadline := range stager.releaseHasDeadline {
+				Expect(hasDeadline).To(BeTrue())
+			}
+			for _, deadline := range stager.releaseDeadlines {
+				Expect(time.Until(deadline)).To(BeNumerically(">", 0))
+				Expect(time.Until(deadline)).To(BeNumerically("<=", time.Minute))
+			}
+		}
+	})
+
+	It("tracks a key before staging so a partial upload failure is released", func(ctx SpecContext) {
+		uploadErr := errors.New("upload failed")
+		stager := &lifecycleStager{ensureErr: uploadErr, releaseErr: errors.New("release failed")}
+		client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
+
+		_, err := client.GenerateImage(ctx, &pb.GenerateImageRequest{Src: "/tmp/source.png"})
+
+		Expect(err).To(MatchError(ContainSubstring("upload failed")))
+		Expect(stager.ensureCalls).To(HaveLen(1))
+		Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
+	})
+
+	It("uses an active bounded cleanup context after caller cancellation", func() {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		stager := &lifecycleStager{}
+		client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
+
+		_, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}})
+
+		Expect(err).NotTo(HaveOccurred())
+		Expect(stager.releaseCtxErr).To(Equal([]error{nil}))
+	})
+
+	It("does not release streaming inputs before the backend completes", func(ctx SpecContext) {
+		block := make(chan struct{})
+		started := make(chan struct{}, 1)
+		backend := &lifecycleBackend{streamBlock: block, streamStarted: started}
+		stager := &lifecycleStager{}
+		client := NewFileStagingClient(backend, stager, "worker-1")
+		done := make(chan error, 1)
+
+		go func() {
+			done <- client.PredictStream(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}}, func(*pb.Reply) {})
+		}()
+
+		Eventually(started).Should(Receive())
+		Expect(stager.ensureCalls).To(HaveLen(1))
+		Expect(stager.releasedKeys).To(BeEmpty())
+		close(block)
+		Eventually(done).Should(Receive(Succeed()))
+		Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
+	})
+
+	It("does not replace a backend result when cleanup fails", func(ctx SpecContext) {
+		reply := &pb.Reply{Message: []byte("ok")}
+		backend := &lifecycleBackend{predictResult: reply}
+		stager := &lifecycleStager{releaseErr: errors.New("release failed")}
+		client := NewFileStagingClient(backend, stager, "worker-1")
+
+		result, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}})
+
+		Expect(err).NotTo(HaveOccurred())
+		Expect(result).To(BeIdenticalTo(reply))
+	})
+})
+
+func keysFromEnsureCalls(calls []ensureCall) []string {
+	keys := make([]string, len(calls))
+	for i, call := range calls {
+		keys[i] = call.key
+	}
+	return keys
+}

core/services/nodes/file_staging_sound_detection_test.go

@@ -28,6 +28,10 @@ func (s *soundStagingFailure) EnsureRemote(context.Context, string, string, stri
 	return "", errors.New("upload failed")
 }
 
+func (s *soundStagingFailure) ReleaseRemote(context.Context, string, string) error {
+	return nil
+}
+
 type soundRouteFactory struct{ client grpc.Backend }
 
 func (f *soundRouteFactory) NewClient(string, bool) grpc.Backend { return f.client }

core/services/nodes/file_staging_tts_test.go

@@ -39,7 +39,8 @@ var _ = Describe("FileStagingClient TTS references", func() {
 		Expect(stager.ensureCalls).To(HaveLen(1))
 		Expect(stager.ensureCalls[0].localPath).To(Equal("/data/voice-profiles/profile/reference.wav"))
 		Expect(backend.ttsRequest.Voice).To(HavePrefix("/remote/ephemeral/"))
-		Expect(backend.ttsRequest.Voice).To(MatchRegexp(`/inputs/[0-9a-f]{8}/reference\.wav$`))
+		voicePathPattern := `/inputs/` + fullUUIDPattern + `/reference\.wav$`
+		Expect(backend.ttsRequest.Voice).To(MatchRegexp(voicePathPattern))
 	})
 
 	It("stages a reference WAV before streaming synthesis", func(ctx SpecContext) {