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

551 lines
14 KiB
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/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)
outputFilter := newPersistentOutputFilter(targetName, profile, prompt)
waitForFilteredMessage := profile.Terminal && isClaudeTerminalProfile(targetName, profile)
_ = 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
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 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() {
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 compactTerminalText(s string) string {
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.mu.Lock()
if s, found := c.sessions[sess.key]; found && s == sess {
delete(c.sessions, sess.key)
}
c.mu.Unlock()
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)
}