Restore panic handler for Function stages

The version-2 Stage interface redesign dropped the panic recovery
from goStage. Restore it: add SetPanicHandler(), recoverPanic(),
and WithStagePanicHandler pipeline option.

The goroutine uses stacked defers (close(done) → close stdin →
close stdout → recoverPanic) so that when a Function panics,
recoverPanic fires first (sets s.err), then cleanup runs, then
done closes — allowing Wait() to return the caught panic error.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: znull

pipe/function.go

@@ -32,20 +32,26 @@ func Function(name string, f StageFunc) Stage {
 // goStage is a `Stage` that does its work by running an arbitrary
 // `stageFunc` in a goroutine.
 type goStage struct {
-	name string
-	f    StageFunc
-	done chan struct{}
-	err  error
+	name         string
+	f            StageFunc
+	done         chan struct{}
+	err          error
+	panicHandler StagePanicHandler
 }
 
 var (
-	_ Stage = (*goStage)(nil)
+	_ Stage                  = (*goStage)(nil)
+	_ StagePanicHandlerAware = (*goStage)(nil)
 )
 
 func (s *goStage) Name() string {
 	return s.name
 }
 
+func (s *goStage) SetPanicHandler(ph StagePanicHandler) {
+	s.panicHandler = ph
+}
+
 func (s *goStage) Preferences() StagePreferences {
 	return StagePreferences{
 		StdinPreference:  IOPreferenceUndefined,
@@ -67,21 +73,24 @@ func (s *goStage) Start(
 	}
 
 	go func() {
-		s.err = s.f(ctx, env, r, w)
-
-		if stdout != nil {
-			if err := stdout.Close(); err != nil && s.err == nil {
-				s.err = fmt.Errorf("error closing stdout for stage %q: %w", s.Name(), err)
+		defer close(s.done)
+		defer func() {
+			if stdin != nil {
+				if err := stdin.Close(); err != nil && s.err == nil {
+					s.err = fmt.Errorf("error closing stdin for stage %q: %w", s.Name(), err)
+				}
 			}
-		}
-
-		if stdin != nil {
-			if err := stdin.Close(); err != nil && s.err == nil {
-				s.err = fmt.Errorf("error closing stdin for stage %q: %w", s.Name(), err)
+		}()
+		defer func() {
+			if stdout != nil {
+				if err := stdout.Close(); err != nil && s.err == nil {
+					s.err = fmt.Errorf("error closing stdout for stage %q: %w", s.Name(), err)
+				}
 			}
-		}
+		}()
+		defer s.recoverPanic()
 
-		close(s.done)
+		s.err = s.f(ctx, env, r, w)
 	}()
 
 	return nil
@@ -91,3 +100,13 @@ func (s *goStage) Wait() error {
 	<-s.done
 	return s.err
 }
+
+func (s *goStage) recoverPanic() {
+	if s.panicHandler == nil {
+		return
+	}
+
+	if p := recover(); p != nil {
+		s.err = s.panicHandler(p)
+	}
+}

pipe/pipeline.go

@@ -65,6 +65,7 @@ type Pipeline struct {
 	started uint32
 
 	eventHandler func(e *Event)
+	panicHandler StagePanicHandler
 }
 
 var emptyEventHandler = func(e *Event) {}
@@ -216,6 +217,15 @@ func WithEventHandler(handler func(e *Event)) Option {
 	}
 }
 
+// WithStagePanicHandler sets a panic handler for the stages within a pipeline.
+// When a pipeline stage panics, the provided handler will be invoked, allowing
+// the client to handle the panic in whatever way they see fit.
+func WithStagePanicHandler(ph StagePanicHandler) Option {
+	return func(p *Pipeline) {
+		p.panicHandler = ph
+	}
+}
+
 func (p *Pipeline) hasStarted() bool {
 	return atomic.LoadUint32(&p.started) != 0
 }
@@ -315,6 +325,10 @@ func (p *Pipeline) Start(ctx context.Context) error {
 		ss := &stageStarters[i]
 		nextSS := &stageStarters[i+1]
 
+		if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil {
+			phs.SetPanicHandler(p.panicHandler)
+		}
+
 		// We need to generate a pipe pair for this stage to use
 		// to communicate with its successor:
 		if ss.prefs.StdoutPreference == IOPreferenceFile ||
@@ -343,6 +357,10 @@ func (p *Pipeline) Start(ctx context.Context) error {
 		s := p.stages[i]
 		ss := &stageStarters[i]
 
+		if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil {
+			phs.SetPanicHandler(p.panicHandler)
+		}
+
 		if err := s.Start(ctx, p.env, ss.stdin, ss.stdout); err != nil {
 			return abort(i, err)
 		}

pipe/pipeline_test.go

@@ -457,6 +457,27 @@ func TestFunction(t *testing.T) {
 		assert.NoError(t, err)
 		assert.EqualValues(t, "goodbye, cruel world", out)
 	})
+
+	t.Run("panic with handler", func(t *testing.T) {
+		p := pipe.New(
+			pipe.WithStagePanicHandler(func(p any) error {
+				return fmt.Errorf("panic handled: %v", p)
+			}),
+		)
+		p.Add(
+			pipe.Print("hello world"),
+			pipe.Function(
+				"farewell",
+				func(_ context.Context, _ pipe.Env, _ io.Reader, _ io.Writer) error {
+					panic("this is a panic")
+				},
+			),
+		)
+
+		out, err := p.Output(ctx)
+		assert.ErrorContains(t, err, "panic handled")
+		assert.Empty(t, out)
+	})
 }
 
 func TestPipelineWithFunction(t *testing.T) {