| 1 | package persistentshell |
| 2 | |
| 3 | import ( |
| 4 | "io" |
| 5 | "strings" |
| 6 | |
| 7 | "reasonix/internal/shellrun" |
| 8 | ) |
| 9 | |
| 10 | // maxIncompleteEscape bounds the escape-sequence fragment carried between PTY |
| 11 | // reads. A stream that never terminates a sequence is malformed, not a reason |
| 12 | // to grow memory. |
| 13 | const maxIncompleteEscape = 8 << 10 |
| 14 | |
| 15 | // sanitizer turns raw PTY bytes into model-visible text. The command now runs |
| 16 | // on a real terminal, so programs that probe isatty emit colour and cursor |
| 17 | // control that the one-shot path never produced; those sequences are dropped |
| 18 | // rather than shown. Incomplete sequences and a trailing carriage return carry |
| 19 | // into the next read. |
| 20 | type sanitizer struct { |
| 21 | pending []byte |
| 22 | carryCR bool |
| 23 | } |
| 24 | |
| 25 | func (s *sanitizer) push(chunk []byte) string { |
| 26 | if len(chunk) == 0 { |
| 27 | return "" |
| 28 | } |
| 29 | s.pending = append(s.pending, chunk...) |
| 30 | text, rest := stripTerminalControls(s.pending) |
| 31 | if len(rest) > maxIncompleteEscape { |
| 32 | rest = nil |
| 33 | } |
| 34 | s.pending = append(s.pending[:0], rest...) |
| 35 | return s.lineEndings(text) |
| 36 | } |
| 37 | |
| 38 | // flush returns any carried text once the PTY is done producing bytes. |
| 39 | func (s *sanitizer) flush() string { |
| 40 | s.pending = nil |
| 41 | if !s.carryCR { |
| 42 | return "" |
| 43 | } |
| 44 | s.carryCR = false |
| 45 | return "\n" |
| 46 | } |
| 47 | |
| 48 | func (s *sanitizer) lineEndings(text string) string { |
| 49 | if s.carryCR { |
| 50 | text = "\r" + text |
| 51 | s.carryCR = false |
| 52 | } |
| 53 | // A trailing carriage return cannot be classified yet: the next read decides |
| 54 | // whether it precedes a newline (one line break) or stands alone. |
| 55 | for strings.HasSuffix(text, "\r") { |
| 56 | text = text[:len(text)-1] |
| 57 | s.carryCR = true |
| 58 | } |
| 59 | return normalizePTY(text) |
| 60 | } |
| 61 | |
| 62 | // stripTerminalControls removes CSI/OSC/two-byte escape sequences and BEL, |
| 63 | // returning the printable text plus any trailing incomplete sequence. |
| 64 | func stripTerminalControls(b []byte) (string, []byte) { |
| 65 | if idx := indexControl(b); idx < 0 { |
| 66 | return string(b), nil |
| 67 | } |
| 68 | var out strings.Builder |
| 69 | out.Grow(len(b)) |
| 70 | i := 0 |
| 71 | for i < len(b) { |
| 72 | c := b[i] |
| 73 | if c == 0x07 { // BEL is a notification, not text |
| 74 | i++ |
| 75 | continue |
| 76 | } |
| 77 | if c != 0x1b { |
| 78 | out.WriteByte(c) |
| 79 | i++ |
| 80 | continue |
| 81 | } |
| 82 | if i+1 >= len(b) { |
| 83 | return out.String(), b[i:] |
| 84 | } |
| 85 | switch b[i+1] { |
| 86 | case '[': |
| 87 | end := i + 2 |
| 88 | for end < len(b) && (b[end] < 0x40 || b[end] > 0x7e) { |
| 89 | end++ |
| 90 | } |
| 91 | if end >= len(b) { |
| 92 | return out.String(), b[i:] |
| 93 | } |
| 94 | i = end + 1 |
| 95 | case ']': |
| 96 | end, ok := indexOSCTerminator(b, i+2) |
| 97 | if !ok { |
| 98 | return out.String(), b[i:] |
| 99 | } |
| 100 | i = end |
| 101 | default: |
| 102 | i += 2 |
| 103 | } |
| 104 | } |
| 105 | return out.String(), nil |
| 106 | } |
| 107 | |
| 108 | func indexControl(b []byte) int { |
| 109 | for i, c := range b { |
| 110 | if c == 0x1b || c == 0x07 { |
| 111 | return i |
| 112 | } |
| 113 | } |
| 114 | return -1 |
| 115 | } |
| 116 | |
| 117 | func indexOSCTerminator(b []byte, from int) (int, bool) { |
| 118 | for i := from; i < len(b); i++ { |
| 119 | if b[i] == 0x07 { |
| 120 | return i + 1, true |
| 121 | } |
| 122 | if b[i] == 0x1b { |
| 123 | if i+1 >= len(b) { |
| 124 | return 0, false |
| 125 | } |
| 126 | if b[i+1] == '\\' { |
| 127 | return i + 2, true |
| 128 | } |
| 129 | } |
| 130 | } |
| 131 | return 0, false |
| 132 | } |
| 133 | |
| 134 | // capture extracts one command's body from the sanitized stream. It scans only |
| 135 | // newly arrived text plus a marker-width overlap, so a command that prints |
| 136 | // megabytes costs time linear in its output instead of rescanning the whole |
| 137 | // transcript on every read. |
| 138 | type capture struct { |
| 139 | start, end string |
| 140 | out *shellrun.BoundedOutput |
| 141 | progress io.Writer |
| 142 | |
| 143 | started bool |
| 144 | done bool |
| 145 | exitCode int |
| 146 | |
| 147 | pre string // scratch while waiting for the start marker |
| 148 | hold string // post-start text withheld until it cannot be a marker prefix |
| 149 | } |
| 150 | |
| 151 | func newCapture(start, end string, progress io.Writer) *capture { |
| 152 | return &capture{start: start, end: end, out: shellrun.NewBoundedOutput(), progress: progress} |
| 153 | } |
| 154 | |
| 155 | func (c *capture) push(text string) { |
| 156 | if text == "" || c.done { |
| 157 | return |
| 158 | } |
| 159 | if !c.started { |
| 160 | c.pre += text |
| 161 | marker := c.start + "\n" |
| 162 | idx := strings.Index(c.pre, marker) |
| 163 | if idx < 0 { |
| 164 | if len(c.pre) > preStartCap { |
| 165 | c.pre = c.pre[len(c.pre)-preStartCap:] |
| 166 | } |
| 167 | return |
| 168 | } |
| 169 | rest := c.pre[idx+len(marker):] |
| 170 | c.started, c.pre = true, "" |
| 171 | c.consume(rest) |
| 172 | return |
| 173 | } |
| 174 | c.consume(text) |
| 175 | } |
| 176 | |
| 177 | func (c *capture) consume(text string) { |
| 178 | c.hold += text |
| 179 | for { |
| 180 | idx := strings.Index(c.hold, c.end) |
| 181 | if idx < 0 { |
| 182 | break |
| 183 | } |
| 184 | status, ok, pending := parseStatus(c.hold[idx+len(c.end):]) |
| 185 | if pending { |
| 186 | // The status line has not arrived yet; withhold from the marker on. |
| 187 | c.emit(c.hold[:idx]) |
| 188 | c.hold = c.hold[idx:] |
| 189 | return |
| 190 | } |
| 191 | if !ok { |
| 192 | // A complete trailer that is not digits is not this command's |
| 193 | // completion (a shell tracing its own input prints one). Treat it |
| 194 | // as output and keep scanning. |
| 195 | c.emit(c.hold[:idx+len(c.end)]) |
| 196 | c.hold = c.hold[idx+len(c.end):] |
| 197 | continue |
| 198 | } |
| 199 | c.emit(c.hold[:idx]) |
| 200 | c.hold, c.done, c.exitCode = "", true, status |
| 201 | return |
| 202 | } |
| 203 | if len(c.hold) > markerOverlap { |
| 204 | keep := len(c.hold) - markerOverlap |
| 205 | c.emit(c.hold[:keep]) |
| 206 | c.hold = c.hold[keep:] |
| 207 | } |
| 208 | } |
| 209 | |
| 210 | func (c *capture) emit(text string) { |
| 211 | if text == "" { |
| 212 | return |
| 213 | } |
| 214 | c.out.WriteString(text) |
| 215 | if c.progress != nil { |
| 216 | _, _ = io.WriteString(c.progress, text) |
| 217 | } |
| 218 | } |
| 219 | |
| 220 | // partial releases text still withheld for marker matching. A command that |
| 221 | // timed out or died never printed its status marker, and its output is the only |
| 222 | // evidence the model gets about what ran. |
| 223 | func (c *capture) partial() string { |
| 224 | if c.hold != "" { |
| 225 | c.emit(c.hold) |
| 226 | c.hold = "" |
| 227 | } |
| 228 | return c.out.String() |
| 229 | } |
| 230 | |
| 231 | func (c *capture) body() string { return c.out.String() } |
| 232 |