返回 DeepSeek-Reasonix
protocol.go
根目录 / internal / persistentshell / protocol.go
1 package persistentshell
2
3 import (
4 "crypto/rand"
5 "encoding/hex"
6 "fmt"
7 "strconv"
8 "strings"
9 )
10
11 const (
12 readyToken = "REASONIX_SHELL_READY"
13 // markerOverlap is how much already-scanned output is re-examined with the
14 // next PTY read so a marker split across two reads is still found. It must
15 // exceed the longest marker plus its status digits and terminator.
16 markerOverlap = 96
17 // preStartCap bounds what is retained while waiting for the start marker.
18 // Only echoed wrapper source and stray output from a previous command's
19 // background child can appear there, and none of it is model-visible.
20 preStartCap = 64 << 10
21 )
22
23 func newMarkerID() string {
24 var b [8]byte
25 if _, err := rand.Read(b[:]); err != nil {
26 return fmt.Sprintf("%x", b[:])
27 }
28 return hex.EncodeToString(b[:])
29 }
30
31 func posixQuote(s string) string {
32 return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
33 }
34
35 // ansiCQuote renders s as an ASCII-only $'...' literal with no raw control bytes.
36 // A PTY line discipline is not a reliable carrier for arbitrary control bytes.
37 // Readline can interpret high bytes as editing keys in a C/unset locale before
38 // the shell parses the literal. Escape bytes, not runes, to preserve the command
39 // exactly without changing the user's locale or interactive editing settings.
40 func ansiCQuote(s string) string {
41 var b strings.Builder
42 b.Grow(len(s) + 8)
43 b.WriteString("$'")
44 for i := range len(s) {
45 c := s[i]
46 switch c {
47 case '\\':
48 b.WriteString(`\\`)
49 case '\'':
50 b.WriteString(`\'`)
51 case '\n':
52 b.WriteString(`\n`)
53 case '\r':
54 b.WriteString(`\r`)
55 case '\t':
56 b.WriteString(`\t`)
57 default:
58 if c < 0x20 || c >= 0x7f {
59 // Octal escapes work across supported POSIX shells; \xHH does not.
60 // Three digits prevent a following command digit joining the escape
61 // (e.g. byte 1 followed by "70").
62 b.WriteByte('\\')
63 b.WriteByte('0' + (c >> 6))
64 b.WriteByte('0' + ((c >> 3) & 7))
65 b.WriteByte('0' + (c & 7))
66 continue
67 }
68 b.WriteByte(c)
69 }
70 }
71 b.WriteString("'")
72 return b.String()
73 }
74
75 func posixSetupScript() string {
76 return strings.Join([]string{
77 // -onlcr stops the line discipline rewriting \n as \r\n and emitting a
78 // stray extra \r under output pressure, which normalisation would turn
79 // into a blank line. The sanitizer covers hosts that reject it.
80 "stty -echo -onlcr 2>/dev/null || stty -echo 2>/dev/null || true",
81 "unset PROMPT_COMMAND",
82 "PS1=",
83 "PS2=",
84 "PS4=",
85 "set +H 2>/dev/null || true",
86 "printf '%s\\n' " + posixQuote(readyToken),
87 }, "; ") + "\n"
88 }
89
90 // posixCommandScript frames a short command and detaches its stdin. Long
91 // commands are staged with acknowledgements by commandStages instead.
92 func posixCommandScript(command, start, end string) string {
93 return commandWordScript(ansiCQuote(command), start, end, "")
94 }
95
96 func commandWordScript(word, start, end, cleanup string) string {
97 return "printf '%s\\n' " + posixQuote(start) +
98 "; eval -- " + word + " </dev/null" +
99 "; __rx_status=$?" + cleanup +
100 "; printf '%s%s\\n' " + posixQuote(end) + ` "$__rx_status"` + "\n"
101 }
102
103 // normalizePTY collapses terminal line endings. A run of carriage returns
104 // before a newline is one line break: the line discipline can emit \r\r\n under
105 // output pressure, and mapping each \r to \n would inject blank lines into
106 // model-visible output. A standalone \r (progress bars) still becomes a break.
107 func normalizePTY(s string) string {
108 if !strings.ContainsRune(s, '\r') {
109 return s
110 }
111 var b strings.Builder
112 b.Grow(len(s))
113 for i := 0; i < len(s); i++ {
114 if s[i] != '\r' {
115 b.WriteByte(s[i])
116 continue
117 }
118 for i+1 < len(s) && s[i+1] == '\r' {
119 i++
120 }
121 if i+1 < len(s) && s[i+1] == '\n' {
122 continue // the \n itself is written on the next iteration
123 }
124 b.WriteByte('\n')
125 }
126 return b.String()
127 }
128
129 // readyLine reports whether the shell printed the ready token. The echoed setup
130 // source contains the token too, so completion requires the token to end its
131 // own line; the echo continues with a quote character.
132 func readyLine(buf string) bool {
133 return strings.Contains(normalizePTY(buf), readyToken+"\n")
134 }
135
136 // parseStatus reads the exit status that must follow an end marker. Digits
137 // terminated by a newline are required, which is what stops echoed wrapper
138 // source from fabricating a completion.
139 //
140 // pending distinguishes "the status line has not arrived yet" from "this marker
141 // is not a completion". Without it, a shell tracing its own input (set -x)
142 // parks on a marker whose trailer never becomes digits.
143 func parseStatus(after string) (status int, ok bool, pending bool) {
144 line, _, complete := strings.Cut(after, "\n")
145 if !complete {
146 return 0, false, true
147 }
148 n, err := strconv.Atoi(strings.TrimSuffix(line, "\r"))
149 if err != nil {
150 return 0, false, false
151 }
152 return n, true, false
153 }
154
155 // extractOutput finds a completed command in text. The markers are matched as
156 // substrings rather than whole lines: a command whose output does not end in a
157 // newline leaves the status marker mid-line, and requiring a line start there
158 // is what used to hang such a command until its deadline.
159 func extractOutput(buf, start, end string) (body string, code int, ok bool) {
160 text := normalizePTY(buf)
161 endIdx := strings.LastIndex(text, end)
162 if endIdx < 0 {
163 return "", 0, false
164 }
165 status, ok, _ := parseStatus(text[endIdx+len(end):])
166 if !ok {
167 return "", 0, false
168 }
169 body = text[:endIdx]
170 if startIdx := strings.LastIndex(body, start+"\n"); startIdx >= 0 {
171 body = body[startIdx+len(start)+1:]
172 }
173 return strings.TrimSuffix(body, "\n"), status, true
174 }
175
175 lines GO