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

package cli

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"os/exec"
	"regexp"
	"strings"
	"time"
	"unicode"

	"github.com/creack/pty"
	"go.uber.org/zap"

	"iop/apps/node/internal/adapters/cli/status"
	"iop/apps/node/internal/runtime"
	"iop/packages/go/config"
)

const (
	terminalInputDelay = 2 * time.Millisecond
	terminalRows       = 80
	terminalCols       = 240
)

func emitRuntimeError(ctx context.Context, sink runtime.EventSink, runID, msg string) error {
	_ = sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:     runID,
		Type:      runtime.EventTypeError,
		Error:     msg,
		Timestamp: time.Now(),
	})
	return fmt.Errorf("cli adapter: %s", msg)
}

type completionMatcher struct {
	line string
	re   *regexp.Regexp
}

func newCompletionMatcher(m config.CompletionMarkerConf) (completionMatcher, error) {
	var cm completionMatcher
	cm.line = m.Line
	if m.Regex != "" {
		re, err := regexp.Compile(m.Regex)
		if err != nil {
			return completionMatcher{}, fmt.Errorf("completion_marker regex: %w", err)
		}
		cm.re = re
	}
	return cm, nil
}

func (m completionMatcher) match(line string) bool {
	if m.line != "" && line == m.line {
		return true
	}
	if m.re != nil && m.re.MatchString(line) {
		return true
	}
	return false
}

func (c *CLI) executePersistent(ctx context.Context, spec runtime.ExecutionSpec, profile config.CLIProfileConf, sink runtime.EventSink) error {
	sess, err := c.resolveSession(ctx, spec, profile)
	if err != nil {
		return err
	}

	matcher, err := newCompletionMatcher(profile.CompletionMarker)
	if err != nil {
		return emitRuntimeError(ctx, sink, spec.RunID, err.Error())
	}

	sess.mu.Lock()
	defer sess.mu.Unlock()

	prompt := extractPrompt(spec.Input)
	idleTimeout := time.Duration(profile.ResponseIdleTimeoutMS) * time.Millisecond
	if idleTimeout <= 0 {
		idleTimeout = 1500 * time.Millisecond
	}
	targetName := cliTargetName(spec)
	waitForFilteredMessage := profile.Terminal && isClaudeTerminalProfile(targetName, profile)
	baselineAssistant := ""
	if waitForFilteredMessage {
		baselineAssistant, _ = latestClaudeAssistantMessageFromCleanOutput(cleanClaudeTerminalOutput(sess.getTail()))
	}
	outputFilter := newPersistentOutputFilter(targetName, profile, prompt, baselineAssistant)

	_ = sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:     spec.RunID,
		Type:      runtime.EventTypeStart,
		Timestamp: time.Now(),
	})

	if err := writePrompt(ctx, sess.input, prompt, profile); err != nil {
		return emitRuntimeError(ctx, sink, spec.RunID, fmt.Sprintf("write prompt: %v", err))
	}

	var idleTimer *time.Timer
	var idleC <-chan time.Time
	defer func() {
		if idleTimer != nil {
			idleTimer.Stop()
		}
	}()
	outputTokens := 0
	var markerBuf strings.Builder
	claudePromptReplayCount := 0

	for {
		select {
		case <-ctx.Done():
			// Drain output so the process remains usable for the next run.
			drainSessionUntilIdle(sess.output, idleTimeout, c.logger, sess.key)
			_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeCancelled,
				Message:   cancelEventForContext(ctx.Err()),
				Timestamp: time.Now(),
			})
			return runtime.ErrRunCancelled
		case out, ok := <-sess.output:
			if !ok {
				err := drainPersistentDone(sess)
				return c.emitPersistentExit(ctx, sink, spec.RunID, targetName, profile, sess, err)
			}
			if waitForFilteredMessage {
				if msg, cancelled := claudeTerminalCancelMessage(sess.getTail()); cancelled {
					c.removePersistentSession(sess)
					_ = closeProfileSession(context.Background(), sess)
					_ = sink.Emit(ctx, runtime.RuntimeEvent{
						RunID:     spec.RunID,
						Type:      runtime.EventTypeCancelled,
						Message:   msg,
						Timestamp: time.Now(),
					})
					return fmt.Errorf("%w: %s", runtime.ErrRunCancelled, msg)
				}
			}
			if delta := outputFilter.Filter(out.text); delta != "" {
				outputTokens += len(strings.Fields(delta))
				_ = sink.Emit(ctx, runtime.RuntimeEvent{
					RunID:     spec.RunID,
					Type:      runtime.EventTypeDelta,
					Delta:     delta,
					Timestamp: time.Now(),
				})
			}
			markerLines := []string(nil)
			if out.markerLine != "" {
				markerLines = append(markerLines, out.markerLine)
			} else {
				markerLines = consumeCompleteLines(&markerBuf, out.text)
			}
			if matcher.matchAny(markerLines) {
				if delta := outputFilter.Flush(); delta != "" {
					outputTokens += len(strings.Fields(delta))
					_ = sink.Emit(ctx, runtime.RuntimeEvent{
						RunID:     spec.RunID,
						Type:      runtime.EventTypeDelta,
						Delta:     delta,
						Timestamp: time.Now(),
					})
				}
				return sink.Emit(ctx, runtime.RuntimeEvent{
					RunID:   spec.RunID,
					Type:    runtime.EventTypeComplete,
					Message: "completion-marker",
					Usage: &runtime.UsageStats{
						InputTokens:  len(strings.Fields(prompt)),
						OutputTokens: outputTokens,
					},
					Timestamp: time.Now(),
				})
			}
			if idleTimer == nil {
				idleTimer = time.NewTimer(idleTimeout)
				idleC = idleTimer.C
			} else {
				if !idleTimer.Stop() {
					select {
					case <-idleTimer.C:
					default:
					}
				}
				idleTimer.Reset(idleTimeout)
			}
		case <-idleC:
			if delta := outputFilter.Flush(); delta != "" {
				outputTokens += len(strings.Fields(delta))
				_ = sink.Emit(ctx, runtime.RuntimeEvent{
					RunID:     spec.RunID,
					Type:      runtime.EventTypeDelta,
					Delta:     delta,
					Timestamp: time.Now(),
				})
			} else if waitForFilteredMessage && !outputFilter.HasOutput() {
				if claudePromptReplayCount == 0 && claudeTerminalReadyForInput(sess.getTail()) {
					claudePromptReplayCount++
					c.logger.Info("cli adapter: replaying claude prompt after ready screen", zap.String("target", targetName), zap.String("session", sess.key.sessionID))
					if err := writePrompt(ctx, sess.input, prompt, profile); err != nil {
						return emitRuntimeError(ctx, sink, spec.RunID, fmt.Sprintf("replay prompt: %v", err))
					}
				}
				idleTimer.Reset(idleTimeout)
				continue
			}
			completeMessage := "idle-timeout"
			if msg := outputFilter.CompletionMessage(); msg != "" {
				completeMessage = msg
			}
			return sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:   spec.RunID,
				Type:    runtime.EventTypeComplete,
				Message: completeMessage,
				Usage: &runtime.UsageStats{
					InputTokens:  len(strings.Fields(prompt)),
					OutputTokens: outputTokens,
				},
				Timestamp: time.Now(),
			})
		case err := <-sess.done:
			return c.emitPersistentExit(ctx, sink, spec.RunID, targetName, profile, sess, err)
		}
	}
}

func (m completionMatcher) matchAny(lines []string) bool {
	for _, line := range lines {
		if m.match(line) {
			return true
		}
	}
	return false
}

func consumeCompleteLines(buf *strings.Builder, text string) []string {
	if text == "" {
		return nil
	}
	buf.WriteString(text)
	raw := buf.String()
	start := 0
	var lines []string
	for i, r := range raw {
		if r != '\n' {
			continue
		}
		lines = append(lines, strings.TrimRight(raw[start:i], "\r"))
		start = i + 1
	}
	if start > 0 {
		buf.Reset()
		buf.WriteString(raw[start:])
	} else if len(raw) > 8192 {
		buf.Reset()
		buf.WriteString(raw[len(raw)-8192:])
	}
	return lines
}

func promptTerminator(profile config.CLIProfileConf) string {
	if profile.Terminal {
		return "\r"
	}
	return "\n"
}

func writePrompt(ctx context.Context, input io.Writer, prompt string, profile config.CLIProfileConf) error {
	if !profile.Terminal {
		_, err := io.WriteString(input, prompt+promptTerminator(profile))
		return err
	}
	for _, r := range prompt {
		if _, err := io.WriteString(input, string(r)); err != nil {
			return err
		}
		timer := time.NewTimer(terminalInputDelay)
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}
	}
	_, err := io.WriteString(input, promptTerminator(profile))
	return err
}

// resolveSession returns an existing session for the given key or creates one
// when SessionMode allows it.
func (c *CLI) resolveSession(ctx context.Context, spec runtime.ExecutionSpec, profile config.CLIProfileConf) (*profileSession, error) {
	target := cliTargetName(spec)
	key := sessionKey{target: target, sessionID: normalizeSessionID(spec.SessionID)}
	c.mu.Lock()
	defer c.mu.Unlock()

	if sess, ok := c.sessions[key]; ok {
		return sess, nil
	}
	if spec.SessionMode == runtime.SessionModeRequireExisting {
		return nil, fmt.Errorf("cli adapter: no persistent session for target %q session %q", target, key.sessionID)
	}
	sess, err := startProfileSession(ctx, key, profile, c.logger)
	if err != nil {
		return nil, err
	}
	c.sessions[key] = sess
	return sess, nil
}

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

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

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

	var input io.Writer
	var closeFn func() error
	var cmd *exec.Cmd

	if profile.Terminal {
		cmd = exec.Command(profile.Command, profile.Args...)
		if len(profile.Env) > 0 {
			cmd.Env = append(cmd.Environ(), profile.Env...)
		}
		ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{
			Rows: terminalRows,
			Cols: terminalCols,
		})
		if err != nil {
			return nil, fmt.Errorf("pty start: %w", err)
		}
		input = ptmx
		closeFn = ptmx.Close
		sess.cmd = cmd
		sess.input = ptmx
		sess.closeFn = ptmx.Close

		go func() {
			buf := make([]byte, 4096)
			for {
				n, err := ptmx.Read(buf)
				if n > 0 {
					text := string(buf[:n])
					sess.appendTail(text)
					outputCh <- cliOutput{text: text}
				}
				if err != nil {
					doneCh <- cmd.Wait()
					close(outputCh)
					return
				}
			}
		}()
	} else {
		cmd = exec.Command(profile.Command, profile.Args...)
		if len(profile.Env) > 0 {
			cmd.Env = append(cmd.Environ(), profile.Env...)
		}
		stdin, err := cmd.StdinPipe()
		if err != nil {
			return nil, fmt.Errorf("stdin pipe: %w", err)
		}
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			_ = stdin.Close()
			return nil, fmt.Errorf("stdout pipe: %w", err)
		}
		if err := cmd.Start(); err != nil {
			_ = stdin.Close()
			return nil, fmt.Errorf("start: %w", err)
		}
		input = stdin
		closeFn = stdin.Close
		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)
		}()
	}

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

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

	return sess, 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 appendBounded(buf *strings.Builder, s string, max int) {
	buf.WriteString(s)
	raw := buf.String()
	if len(raw) <= max {
		return
	}
	buf.Reset()
	buf.WriteString(raw[len(raw)-max:])
}

func shouldAcceptClaudeWorkspaceTrust(raw string) bool {
	compact := compactTerminalText(raw)
	return strings.Contains(compact, "quicksafetycheckisthisaproject") &&
		strings.Contains(compact, "yesitrustthisfolder")
}

func shouldAcceptClaudeBypassWarning(raw string) bool {
	compact := compactTerminalText(raw)
	return strings.Contains(compact, "claudecoderunninginbypasspermissionsmode") &&
		strings.Contains(compact, "yesiaccept")
}

func claudeTerminalCancelMessage(raw string) (string, bool) {
	compact := compactTerminalText(raw)
	switch {
	case strings.Contains(compact, "youvehityoursessionlimit"):
		if line := findClaudeTerminalLine(raw, "session limit"); line != "" {
			return line, true
		}
		return "Claude session limit reached", true
	case strings.Contains(compact, "stopandwaitforlimittoreset") && strings.Contains(compact, "upgradeyourplan"):
		return "Claude session limit reached; upgrade was not selected", true
	default:
		return "", false
	}
}

func claudeTerminalReadyForInput(raw string) bool {
	return claudeScreenHasInputPrompt(status.RenderVisibleScreen(raw, terminalRows, terminalCols))
}

func findClaudeTerminalLine(raw, needle string) string {
	needle = strings.ToLower(needle)
	for _, line := range strings.Split(cleanClaudeTerminalOutput(raw), "\n") {
		line = strings.Join(strings.Fields(line), " ")
		if strings.Contains(strings.ToLower(line), needle) {
			return line
		}
	}
	return ""
}

func compactTerminalText(s string) string {
	s = cleanClaudeTerminalOutput(s)
	var b strings.Builder
	inEscape := false
	for _, r := range s {
		if r == '\x1b' {
			inEscape = true
			continue
		}
		if inEscape {
			if r >= '@' && r <= '~' {
				inEscape = false
			}
			continue
		}
		if unicode.IsLetter(r) || unicode.IsDigit(r) {
			b.WriteRune(unicode.ToLower(r))
		}
	}
	return b.String()
}

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
	}
}

func (c *CLI) emitPersistentExit(ctx context.Context, sink runtime.EventSink, runID, targetName string, profile config.CLIProfileConf, sess *profileSession, err error) error {
	c.removePersistentSession(sess)

	cmdSummary := fmt.Sprintf("%s %s", profile.Command, strings.Join(profile.Args, " "))
	tail := sess.getTail()

	var msg string
	if err != nil {
		msg = fmt.Sprintf("persistent session process exited: %s, target=%s, session=%s, command=%q", err.Error(), targetName, sess.key.sessionID, cmdSummary)
	} else {
		msg = fmt.Sprintf("persistent session process exited unexpectedly: target=%s, session=%s, command=%q", targetName, sess.key.sessionID, cmdSummary)
	}

	if tail != "" {
		msg = fmt.Sprintf("%s, recent output: %q", msg, tail)
	}
	return emitRuntimeError(ctx, sink, runID, msg)
}

func (c *CLI) removePersistentSession(sess *profileSession) {
	c.mu.Lock()
	if s, found := c.sessions[sess.key]; found && s == sess {
		delete(c.sessions, sess.key)
	}
	c.mu.Unlock()
}