pmg/internal/pty/router.go

package pty

import (
	"bytes"
	"context"
	"io"
	"sync"
	"sync/atomic"
)

// OutputRouter manages buffered vs live output.
type OutputRouter struct {
	mu        sync.Mutex
	stdout    io.Writer
	buffer    bytes.Buffer
	buffering bool
}

func NewOutputRouter(out io.Writer) (*OutputRouter, error) {
	return &OutputRouter{
		stdout: out,
	}, nil
}

func (r *OutputRouter) Write(p []byte) (n int, err error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if r.buffering {
		// We are in "Prompt Mode", so save this output for later.
		// If we printed it now, it would mess up the confirmation prompt.
		return r.buffer.Write(p)
	}

	// Normal mode: just print it to stdout.
	return r.stdout.Write(p)
}

// Pause starts buffering output. Call this before showing a confirmation prompt.
func (r *OutputRouter) Pause() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.buffering = true
}

// Resume stops buffering, flushes any buffered output, and resumes live output.
// Call this after the confirmation prompt is complete.
func (r *OutputRouter) Resume() {
	r.mu.Lock()
	defer r.mu.Unlock()

	// Flush any buffered output
	if r.buffer.Len() > 0 {
		_, _ = io.Copy(r.stdout, &r.buffer)
		r.buffer.Reset()
	}

	r.buffering = false
}

// writerDest wraps io.Writer for use with atomic.Pointer
// (atomic.Value panics on nil interface stores)
type writerDest struct {
	w io.Writer
}

// InputRouter manages routing stdin to either PTY or a prompt pipe.
// Only ONE goroutine should call ReadLoop().
type InputRouter struct {
	dest       atomic.Pointer[writerDest]
	defaultDst io.Writer // PTY writer
}

func NewInputRouter(ptyWriter io.Writer) (*InputRouter, error) {
	return &InputRouter{
		defaultDst: ptyWriter,
	}, nil
}

// ReadLoop continuously reads from src and routes data to the current destination.
//
// IMPORTANT: Only ONE goroutine should call ReadLoop() because:
//  1. Multiple readers on the same source (e.g., stdin) cause data splitting -
//     one goroutine might read "hel" while another reads "lo\n"
//  2. Concurrent routing decisions create race conditions on the destination
//  3. User input becomes unpredictably interleaved between readers
//
// This function blocks until src returns an error (e.g., EOF).
func (r *InputRouter) ReadLoop(src io.Reader) {
	r.ReadLoopContext(context.Background(), src)
}

// ReadLoopContext continuously reads from src and routes data to the current
// destination until src returns an error or ctx is cancelled.
func (r *InputRouter) ReadLoopContext(ctx context.Context, src io.Reader) {
	buf := make([]byte, 1024)
	for {
		nr, err := readInput(ctx, src, buf)
		if err != nil {
			return
		}

		if nr == 0 {
			continue
		}

		if dest := r.dest.Load(); dest != nil {
			if _, err := dest.w.Write(buf[:nr]); err != nil {
				return
			}
		} else {
			if _, err := r.defaultDst.Write(buf[:nr]); err != nil {
				return
			}
		}
	}
}

// RouteToPrompt switches input to go to the given writer (prompt pipe)
func (r *InputRouter) RouteToPrompt(w io.Writer) {
	r.dest.Store(&writerDest{w: w})
}

// RouteToPTY switches input back to the PTY (default)
func (r *InputRouter) RouteToPTY() {
	r.dest.Store(nil)
}