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

package cli

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

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

	"iop/apps/node/internal/runtime"
	"iop/packages/config"
)

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
	}

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

	if _, err := fmt.Fprintf(sess.input, "%s\n", prompt); 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

	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 {
				// Process exited — remove the dead session.
				c.mu.Lock()
				if s, found := c.sessions[sess.key]; found && s == sess {
					delete(c.sessions, sess.key)
				}
				c.mu.Unlock()
				return emitRuntimeError(ctx, sink, spec.RunID, "persistent session process exited unexpectedly")
			}
			outputTokens += len(strings.Fields(out.line))
			_ = sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeDelta,
				Delta:     out.line + "\n",
				Timestamp: time.Now(),
			})
			if matcher.match(out.line) {
				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:
			return sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:   spec.RunID,
				Type:    runtime.EventTypeComplete,
				Message: "idle-timeout",
				Usage: &runtime.UsageStats{
					InputTokens:  len(strings.Fields(prompt)),
					OutputTokens: outputTokens,
				},
				Timestamp: time.Now(),
			})
		case err := <-sess.done:
			c.mu.Lock()
			if s, found := c.sessions[sess.key]; found && s == sess {
				delete(c.sessions, sess.key)
			}
			c.mu.Unlock()
			msg := "persistent session process exited"
			if err != nil {
				msg = fmt.Sprintf("persistent session process exited: %v", err)
			}
			return emitRuntimeError(ctx, sink, spec.RunID, msg)
		}
	}
}

// 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)

	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.Start(cmd)
		if err != nil {
			return nil, fmt.Errorf("pty start: %w", err)
		}
		input = ptmx
		closeFn = ptmx.Close
		go func() {
			scanner := bufio.NewScanner(ptmx)
			for scanner.Scan() {
				line := strings.TrimRight(scanner.Text(), "\r")
				outputCh <- cliOutput{line: line}
			}
			doneCh <- cmd.Wait()
			close(outputCh)
		}()
	} 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
		go func() {
			scanner := bufio.NewScanner(stdout)
			for scanner.Scan() {
				outputCh <- cliOutput{line: scanner.Text()}
			}
			doneCh <- cmd.Wait()
			close(outputCh)
		}()
	}

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

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

	return &profileSession{
		key:     key,
		name:    key.target,
		profile: profile,
		cmd:     cmd,
		input:   input,
		output:  outputCh,
		done:    doneCh,
		closeFn: closeFn,
	}, nil
}

func drainUntilIdle(outputCh <-chan cliOutput, timeout time.Duration, logger *zap.Logger, name string) {
	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: startup drain", zap.String("profile", name))
		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
		}
	}
}