iop/packages/go/agentprovider/cli/persistent_test_support_test.go

package cli_test

import (
	"fmt"
	"os"
	osexec "os/exec"
	"testing"
)

func TestRawTUIHelperProcess(t *testing.T) {
	switch {
	case os.Getenv("IOP_RAW_TUI_HELPER") == "1":
		runRawTUIHelper()
	case os.Getenv("IOP_CLAUDE_WARNING_HELPER") == "1":
		runClaudeWarningHelper()
	case os.Getenv("IOP_CLAUDE_TRUST_WARNING_HELPER") == "1":
		runClaudeTrustAndWarningHelper()
	case os.Getenv("IOP_CLAUDE_RATE_LIMIT_HELPER") == "1":
		runClaudeRateLimitHelper()
	case os.Getenv("IOP_CLAUDE_READY_REPLAY_HELPER") == "1":
		runClaudeReadyReplayHelper()
	default:
		return
	}
}

func setRawTerminal() {
	cmd := osexec.Command("stty", "raw", "-echo")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "stty raw: %v\n", err)
		os.Exit(2)
	}
}

// emitStartupReady announces that the helper has finished configuring its raw
// terminal and is now reading input. Re-exec'd test binaries pay a variable
// cold-start cost, so the adapter's startup drain must observe this marker (and
// idle after it) before it writes the prompt. Otherwise the prompt would race
// terminal setup and be echoed or line-translated. The marker is consumed by
// the startup drain and never surfaces as a delta.
func emitStartupReady() {
	fmt.Fprint(os.Stdout, "iop-tui-ready\r\n")
	_ = os.Stdout.Sync()
}

func runRawTUIHelper() {
	setRawTerminal()
	emitStartupReady()
	readPromptAndReply()
}

func runClaudeWarningHelper() {
	setRawTerminal()
	fmt.Fprint(os.Stdout, "WARNING: Claude Code running in BypassPermissions mode\r\n2. Yes, I accept\r\n")
	_ = os.Stdout.Sync()

	want := []byte{'\x1b', '[', 'B', '\r'}
	seen := make([]byte, 0, len(want))
	buf := make([]byte, 1)
	for {
		n, err := os.Stdin.Read(buf)
		if n > 0 {
			seen = append(seen, buf[0])
			if len(seen) > len(want) {
				seen = seen[len(seen)-len(want):]
			}
			if string(seen) == string(want) {
				fmt.Fprint(os.Stdout, "ready\r\n")
				_ = os.Stdout.Sync()
				break
			}
		}
		if err != nil {
			os.Exit(0)
		}
	}

	readPromptAndReply()
}

func runClaudeTrustAndWarningHelper() {
	setRawTerminal()
	fmt.Fprint(os.Stdout, "Quick safety check: Is this a project you created or one you trust?\r\n❯ 1. Yes, I trust this folder\r\n2. No, exit\r\n")
	_ = os.Stdout.Sync()

	buf := make([]byte, 1)
	for {
		n, err := os.Stdin.Read(buf)
		if n > 0 && buf[0] == '\r' {
			break
		}
		if err != nil {
			os.Exit(0)
		}
	}

	fmt.Fprint(os.Stdout, "WARNING: Claude Code running in BypassPermissions mode\r\n2. Yes, I accept\r\n")
	_ = os.Stdout.Sync()

	want := []byte{'\x1b', '[', 'B', '\r'}
	seen := make([]byte, 0, len(want))
	for {
		n, err := os.Stdin.Read(buf)
		if n > 0 {
			seen = append(seen, buf[0])
			if len(seen) > len(want) {
				seen = seen[len(seen)-len(want):]
			}
			if string(seen) == string(want) {
				fmt.Fprint(os.Stdout, "ready\r\n")
				_ = os.Stdout.Sync()
				break
			}
		}
		if err != nil {
			os.Exit(0)
		}
	}

	readPromptAndReply()
}

func runClaudeRateLimitHelper() {
	setRawTerminal()
	emitStartupReady()
	for {
		buf := make([]byte, 1)
		n, err := os.Stdin.Read(buf)
		// Trigger on either terminator: in raw mode the adapter's terminator is
		// "\r", but if the prompt raced terminal setup the line discipline may have
		// translated it to "\n". Accepting both keeps session-limit detection
		// independent of the re-exec'd helper's cold-start timing.
		if n > 0 && (buf[0] == '\r' || buf[0] == '\n') {
			fmt.Fprint(os.Stdout, "\u25cf \u23bf  You've hit your session limit \u00b7 resets 12:50pm (Asia/Seoul)\r\n")
			fmt.Fprint(os.Stdout, "What do you want to do?\r\n❯ 1. Stop and wait for limit to reset\r\n2. Upgrade your plan\r\n")
			_ = os.Stdout.Sync()
		}
		if err != nil {
			os.Exit(0)
		}
	}
}

func runClaudeReadyReplayHelper() {
	setRawTerminal()
	emitStartupReady()

	// Handshake via blocking reads instead of a fixed sleep: swallow the first
	// prompt delivery to model a CLI that is not ready for input yet, then
	// advertise the ready screen so the adapter replays the prompt exactly once.
	_ = readRawLine() // first prompt, intentionally dropped before the ready screen

	fmt.Fprint(os.Stdout, "\r\n❯ ")
	_ = os.Stdout.Sync()

	prompt := readRawLine() // replayed prompt
	fmt.Fprintf(os.Stdout, "\r\n\u25cf replay:%s\r\n\u276f ", prompt)
	_ = os.Stdout.Sync()

	buf := make([]byte, 1)
	for {
		if _, err := os.Stdin.Read(buf); err != nil {
			os.Exit(0)
		}
	}
}

func readRawLine() string {
	var line []byte
	buf := make([]byte, 1)
	for {
		n, err := os.Stdin.Read(buf)
		if n > 0 {
			if buf[0] == '\r' || buf[0] == '\n' {
				return string(line)
			}
			line = append(line, buf[0])
		}
		if err != nil {
			os.Exit(0)
		}
	}
}

func readPromptAndReply() {
	var prompt []byte
	buf := make([]byte, 1)
	for {
		n, err := os.Stdin.Read(buf)
		if n > 0 {
			if buf[0] == '\r' {
				fmt.Fprintf(os.Stdout, "reply:%s\n", string(prompt))
				_ = os.Stdout.Sync()
				prompt = prompt[:0]
			} else {
				prompt = append(prompt, buf[0])
			}
		}
		if err != nil {
			os.Exit(0)
		}
	}
}