150 lines
3.3 KiB
Go
150 lines
3.3 KiB
Go
package infra
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"os/exec"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/creack/pty"
|
|
)
|
|
|
|
// PTYCaptureIdle launches a command inside a pseudo-terminal (PTY) and captures
|
|
// all output until the terminal has been idle for at least idle duration, or
|
|
// maxDur has elapsed. Before sending inputs it waits warmup to let the process
|
|
// initialize. Between each input step it waits stepDelay.
|
|
//
|
|
// The returned string is the raw PTY output, ANSI escape sequences included.
|
|
// To turn it into plain text: use vt_render_go_tui to reconstruct the 2D screen
|
|
// layout for TUIs with absolute cursor positioning (claude, htop), or
|
|
// strip_ansi_go_core for sequential, log-like output.
|
|
func PTYCaptureIdle(
|
|
ctx context.Context,
|
|
name string,
|
|
args []string,
|
|
warmup time.Duration,
|
|
inputs []string,
|
|
stepDelay time.Duration,
|
|
idle time.Duration,
|
|
maxDur time.Duration,
|
|
) (string, error) {
|
|
cmd := exec.CommandContext(ctx, name, args...)
|
|
|
|
ptmx, err := pty.Start(cmd)
|
|
if err != nil {
|
|
return "", fmt.Errorf("pty_capture_idle: pty.Start: %w", err)
|
|
}
|
|
|
|
// Set a reasonable terminal size so TUIs render without truncating.
|
|
if szErr := pty.Setsize(ptmx, &pty.Winsize{Rows: 40, Cols: 120}); szErr != nil {
|
|
// Non-fatal: continue even if resize fails.
|
|
_ = szErr
|
|
}
|
|
|
|
var (
|
|
mu sync.Mutex
|
|
buf bytes.Buffer
|
|
lastByte = time.Now()
|
|
)
|
|
|
|
// Reader goroutine: copy PTY output into buf and track last-byte time.
|
|
readDone := make(chan struct{})
|
|
go func() {
|
|
defer close(readDone)
|
|
tmp := make([]byte, 4096)
|
|
for {
|
|
n, rerr := ptmx.Read(tmp)
|
|
if n > 0 {
|
|
mu.Lock()
|
|
buf.Write(tmp[:n])
|
|
lastByte = time.Now()
|
|
mu.Unlock()
|
|
}
|
|
if rerr != nil {
|
|
// EIO/EOF is normal on Linux when the PTY master is closed
|
|
// after the child exits. Not a real error.
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
start := time.Now()
|
|
|
|
// Wait for warmup so the TUI/CLI has time to initialize.
|
|
select {
|
|
case <-time.After(warmup):
|
|
case <-ctx.Done():
|
|
_ = ptmx.Close()
|
|
<-readDone
|
|
mu.Lock()
|
|
out := buf.String()
|
|
mu.Unlock()
|
|
return out, fmt.Errorf("pty_capture_idle: context cancelled during warmup: %w", ctx.Err())
|
|
}
|
|
|
|
// Send inputs one by one with stepDelay between them.
|
|
for _, in := range inputs {
|
|
if _, werr := ptmx.Write([]byte(in)); werr != nil {
|
|
// PTY may have closed already; stop sending.
|
|
break
|
|
}
|
|
select {
|
|
case <-time.After(stepDelay):
|
|
case <-ctx.Done():
|
|
goto done
|
|
}
|
|
}
|
|
|
|
done:
|
|
// Poll until idle or maxDur.
|
|
ticker := time.NewTicker(100 * time.Millisecond)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
mu.Lock()
|
|
sinceLastByte := time.Since(lastByte)
|
|
mu.Unlock()
|
|
|
|
elapsed := time.Since(start)
|
|
if sinceLastByte >= idle || elapsed >= maxDur {
|
|
goto shutdown
|
|
}
|
|
case <-ctx.Done():
|
|
goto shutdown
|
|
}
|
|
}
|
|
|
|
shutdown:
|
|
// Close the PTY master. This signals EOF to the reader goroutine.
|
|
_ = ptmx.Close()
|
|
|
|
// Graceful shutdown: SIGTERM first, then SIGKILL after 2s.
|
|
if cmd.Process != nil {
|
|
_ = cmd.Process.Signal(syscall.SIGTERM)
|
|
killTimer := time.NewTimer(2 * time.Second)
|
|
waitCh := make(chan error, 1)
|
|
go func() { waitCh <- cmd.Wait() }()
|
|
select {
|
|
case <-waitCh:
|
|
// Process exited cleanly.
|
|
case <-killTimer.C:
|
|
_ = cmd.Process.Kill()
|
|
<-waitCh
|
|
}
|
|
killTimer.Stop()
|
|
}
|
|
|
|
<-readDone
|
|
|
|
mu.Lock()
|
|
out := buf.String()
|
|
mu.Unlock()
|
|
|
|
return out, nil
|
|
}
|