iop/apps/node/internal/adapters/cli/persistent_process.go

package cli

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"strings"
	"time"

	"go.uber.org/zap"
	"iop/apps/node/internal/terminal"
	"iop/packages/go/config"
)

func startProfileSession(_ context.Context, key sessionKey, profile config.CLIProfileConf, workspace string, logger *zap.Logger) (*profileSession, error) {
	if profile.Command == "" {
		return nil, fmt.Errorf("target %q has no command", key.target)
	}

	dir, err := prepareWorkspaceDir(workspace)
	if err != nil {
		return nil, err
	}

	outputCh := make(chan cliOutput, 1024)
	doneCh := make(chan error, 1)

	sess := &profileSession{
		key:     key,
		name:    key.target,
		profile: profile,
		output:  outputCh,
		done:    doneCh,
	}

	if profile.Terminal {
		err = startTerminalProfileSession(sess, profile, dir, outputCh, doneCh)
	} else {
		err = startNonTerminalProfileSession(sess, profile, dir, outputCh, doneCh)
	}
	if err != nil {
		return nil, err
	}

	if profile.StartupIdleTimeoutMS > 0 {
		drainUntilIdle(sess.output, sess.input, time.Duration(profile.StartupIdleTimeoutMS)*time.Millisecond, logger, key, profile)
	}

	if err := checkEarlyExit(sess); err != nil {
		return nil, err
	}

	return sess, nil
}

func startTerminalProfileSession(sess *profileSession, profile config.CLIProfileConf, dir string, outputCh chan cliOutput, doneCh chan error) error {
	opts := terminal.Options{
		Command: profile.Command,
		Args:    profile.Args,
		Env:     profile.Env,
		Rows:    terminalRows,
		Cols:    terminalCols,
		Dir:     dir,
	}
	core, err := terminal.StartSession(context.Background(), opts)
	if err != nil {
		return err
	}
	sess.core = core
	sess.input = sessionWriter{sess: core}
	sess.closeFn = core.Close

	go func() {
		for out := range core.Output() {
			outputCh <- cliOutput{text: out.Text}
		}
		err := <-core.Done()
		doneCh <- err
		close(outputCh)
	}()
	return nil
}

func startNonTerminalProfileSession(sess *profileSession, profile config.CLIProfileConf, dir string, outputCh chan cliOutput, doneCh chan error) error {
	cmd, err := buildCmdWithoutContext(profile.Command, profile.Args, profile.Env, dir)
	if err != nil {
		return err
	}
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return fmt.Errorf("stdin pipe: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		_ = stdin.Close()
		return fmt.Errorf("stdout pipe: %w", err)
	}
	if err := cmd.Start(); err != nil {
		_ = stdin.Close()
		return fmt.Errorf("start: %w", err)
	}
	sess.cmd = cmd
	sess.input = stdin
	sess.closeFn = stdin.Close

	go func() {
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			line := scanner.Text()
			text := line + "\n"
			sess.appendTail(text)
			outputCh <- cliOutput{text: text, markerLine: line}
		}
		doneCh <- cmd.Wait()
		close(outputCh)
	}()
	return nil
}

func checkEarlyExit(sess *profileSession) error {
	select {
	case err := <-sess.done:
		_ = sess.closeFn()
		if sess.cmd != nil && sess.cmd.Process != nil {
			_ = sess.cmd.Process.Kill()
		}
		tail := sess.getTail()
		if err == nil {
			if tail != "" {
				return fmt.Errorf("process exited during startup: %s", tail)
			}
			return fmt.Errorf("process exited during startup")
		}
		if tail != "" {
			return fmt.Errorf("process exited during startup: %w (recent output: %q)", err, tail)
		}
		return fmt.Errorf("process exited during startup: %w", err)
	default:
		return nil
	}
}

func drainUntilIdle(outputCh <-chan cliOutput, input io.Writer, timeout time.Duration, logger *zap.Logger, key sessionKey, profile config.CLIProfileConf) {
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	var startupBuf strings.Builder
	acceptedClaudeWorkspaceTrust := false
	acceptedClaudeBypassWarning := false
	for {
		select {
		case out, ok := <-outputCh:
			if !ok {
				return
			}
			if profile.Terminal && out.text != "" {
				appendBounded(&startupBuf, out.text, 8192)
				rawStr := startupBuf.String()

				if !acceptedClaudeWorkspaceTrust && shouldAcceptClaudeWorkspaceTrust(rawStr) {
					if _, err := io.WriteString(input, "\r"); err != nil {
						logger.Warn("cli adapter: accept claude workspace trust", zap.String("target", key.target), zap.Error(err))
					} else {
						acceptedClaudeWorkspaceTrust = true
						logger.Info("cli adapter: accepted claude workspace trust", zap.String("target", key.target))
						startupBuf.Reset()
					}
				} else if !acceptedClaudeBypassWarning && shouldAcceptClaudeBypassWarning(rawStr) {
					if _, err := io.WriteString(input, "\x1b[B\r"); err != nil {
						logger.Warn("cli adapter: accept claude bypass warning", zap.String("target", key.target), zap.Error(err))
					} else {
						acceptedClaudeBypassWarning = true
						logger.Info("cli adapter: accepted claude bypass warning", zap.String("target", key.target))
					}
				}
			}
			if !timer.Stop() {
				select {
				case <-timer.C:
				default:
				}
			}
			timer.Reset(timeout)
			logger.Debug("cli adapter: startup drain", zap.String("target", key.target), zap.String("session", key.sessionID))
		case <-timer.C:
			return
		}
	}
}

func drainSessionUntilIdle(outputCh <-chan cliOutput, timeout time.Duration, logger *zap.Logger, key sessionKey) {
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	for {
		select {
		case _, ok := <-outputCh:
			if !ok {
				return
			}
			if !timer.Stop() {
				select {
				case <-timer.C:
				default:
				}
			}
			timer.Reset(timeout)
			logger.Debug("cli adapter: cancel drain", zap.String("target", key.target), zap.String("session", key.sessionID))
		case <-timer.C:
			return
		}
	}
}

func drainPersistentDone(sess *profileSession) error {
	select {
	case err := <-sess.done:
		return err
	default:
		return nil
	}
}

type sessionWriter struct {
	sess terminal.Session
}

func (w sessionWriter) Write(p []byte) (n int, err error) {
	err = w.sess.WriteInput(context.Background(), p)
	if err != nil {
		return 0, err
	}
	return len(p), nil
}