551 lines
14 KiB
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)
|
|
}
|