2 files changed, 77 insertions, 4 deletions

diff --git a/internal/probe/remux.go b/internal/probe/remux.go
index 7f3e3d5..35e87d3 100644
--- a/internal/probe/remux.go
+++ b/internal/probe/remux.go
@@ -2,6 +2,7 @@ package probe
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -47,17 +48,33 @@ func (f *FFRemuxer) Remux(ctx context.Context, inputPath string, w io.Writer) er
 		return fmt.Errorf("remux stdout pipe: %w", err)
 	}
 	if err := cmd.Start(); err != nil {
+		if isContextError(err) {
+			return nil
+		}
 		return fmt.Errorf("remux start: %w", err)
 	}
-	if _, err := io.Copy(w, stdout); err != nil && ctx.Err() == nil {
-		slog.Error("copy remuxed media", "file", inputPath, "err", err)
+	var firstErr error
+	if _, err := io.Copy(w, stdout); err != nil {
+		if !isContextError(err) {
+			slog.Error("copy remuxed media", "file", inputPath, "err", err)
+			firstErr = fmt.Errorf("copy remuxed media: %w", err)
+		}
+	}
+	if err := cmd.Wait(); err != nil {
+		if !isContextError(err) && firstErr == nil {
+			firstErr = fmt.Errorf("remux wait: %w", err)
+		}
 	}
-	if err := cmd.Wait(); err != nil && ctx.Err() == nil {
-		return fmt.Errorf("remux wait: %w", err)
+	if firstErr != nil {
+		return firstErr
 	}
 	return nil
 }
 
+func isContextError(err error) bool {
+	return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)
+}
+
 // LooksLikeMPEGTS inspects the first bytes of a file for MPEG-TS sync
 // markers (0x47 at 188- or 192-byte intervals).
 func LooksLikeMPEGTS(path string) bool {
diff --git a/internal/probe/remux_test.go b/internal/probe/remux_test.go
index 6e11561..0a5a611 100644
--- a/internal/probe/remux_test.go
+++ b/internal/probe/remux_test.go
@@ -1,11 +1,56 @@
 package probe
 
 import (
+	"context"
+	"errors"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 )
 
+type cancelingWriter struct {
+	cancel context.CancelFunc
+	err    error
+}
+
+func (w cancelingWriter) Write(p []byte) (int, error) {
+	w.cancel()
+	if w.err != nil {
+		return 0, w.err
+	}
+	return len(p), nil
+}
+
+func TestFFRemuxerReturnsCopyErrorWhenContextCanceled(t *testing.T) {
+	installFakeFFmpeg(t, "printf remuxed-output\nsleep 10\n")
+	errRemuxWrite := errors.New("remux write failed")
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	err := NewFFRemuxer().Remux(ctx, "input.ts", cancelingWriter{
+		cancel: cancel,
+		err:    errRemuxWrite,
+	})
+	if !errors.Is(err, errRemuxWrite) {
+		t.Fatalf("expected copy error, got %v", err)
+	}
+}
+
+func TestFFRemuxerReturnsWaitErrorWhenContextCanceled(t *testing.T) {
+	installFakeFFmpeg(t, "printf remuxed-output\nexit 7\n")
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	err := NewFFRemuxer().Remux(ctx, "input.ts", cancelingWriter{cancel: cancel})
+	if err == nil {
+		t.Fatal("expected wait error")
+	}
+	if !strings.Contains(err.Error(), "remux wait") {
+		t.Fatalf("expected remux wait error, got %v", err)
+	}
+}
+
 func TestLooksLikeMPEGTS(t *testing.T) {
 	ts := make([]byte, 188*5)
 	for i := 0; i < len(ts); i += 188 {
@@ -27,3 +72,14 @@ func TestLooksLikeMPEGTS(t *testing.T) {
 		t.Fatal("did not expect MP4 header to be detected as MPEG-TS")
 	}
 }
+
+func installFakeFFmpeg(t *testing.T, body string) {
+	t.Helper()
+	dir := t.TempDir()
+	path := filepath.Join(dir, "ffmpeg")
+	script := "#!/bin/sh\n" + body
+	if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH"))
+}