| 1 | package persistentshell |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "crypto/sha256" |
| 6 | "encoding/hex" |
| 7 | "errors" |
| 8 | "fmt" |
| 9 | "io" |
| 10 | "runtime" |
| 11 | "strings" |
| 12 | "sync" |
| 13 | "time" |
| 14 | |
| 15 | "reasonix/internal/sandbox" |
| 16 | "reasonix/internal/tool" |
| 17 | ) |
| 18 | |
| 19 | var ( |
| 20 | // ErrUnavailable reports that the manager can no longer start or reuse a |
| 21 | // persistent shell (sealed after the last controller owner released). |
| 22 | ErrUnavailable = errors.New("persistent shell unavailable") |
| 23 | errEmptyArgv = errors.New("empty argv") |
| 24 | ) |
| 25 | |
| 26 | const ( |
| 27 | startupTimeout = 10 * time.Second |
| 28 | powerShellStartupTimeout = 30 * time.Second |
| 29 | readChunk = 4096 |
| 30 | ) |
| 31 | |
| 32 | // Request is one foreground command to run in the session-scoped PTY. |
| 33 | type Request struct { |
| 34 | Argv []string |
| 35 | Dir string |
| 36 | Env []string |
| 37 | Command string |
| 38 | Timeout time.Duration |
| 39 | Shell sandbox.Shell |
| 40 | // Progress receives live output chunks. Callers pass the shared capped |
| 41 | // writer; this package does not re-implement the live-output bound. |
| 42 | Progress io.Writer |
| 43 | } |
| 44 | |
| 45 | // Result is the structured outcome of one persistent-shell command. |
| 46 | type Result struct { |
| 47 | Output string |
| 48 | ExitCode int |
| 49 | ExitCodeKnown bool |
| 50 | TimedOut bool |
| 51 | Canceled bool |
| 52 | ShellDied bool |
| 53 | Started bool |
| 54 | // Reset reports that the shell was retired, so the next command starts from |
| 55 | // the workspace with a fresh directory and environment. The model is told, |
| 56 | // because it otherwise keeps reasoning about a cwd that no longer exists. |
| 57 | Reset bool |
| 58 | State string |
| 59 | FailurePhase string |
| 60 | Err error |
| 61 | } |
| 62 | |
| 63 | // Manager owns one logical session's persistent PTY. Controllers Retain/Release |
| 64 | // it so hot rebuilds share the live shell; Rotate closes it so a new logical |
| 65 | // session cannot inherit cwd or environment. |
| 66 | type Manager struct { |
| 67 | metrics shellMetrics |
| 68 | mu sync.Mutex |
| 69 | runMu sync.Mutex |
| 70 | owners int |
| 71 | sealed bool |
| 72 | live *session |
| 73 | startupFailure *StartupError |
| 74 | startupFingerprint string |
| 75 | startupFailedAt time.Time |
| 76 | } |
| 77 | |
| 78 | // StartupError preserves stderr from a PTY that opened but never became ready. |
| 79 | // It must not be discarded in favor of retrying the same runtime one-shot. |
| 80 | type StartupError struct { |
| 81 | Output string |
| 82 | Err error |
| 83 | } |
| 84 | |
| 85 | func (e *StartupError) Error() string { |
| 86 | return fmt.Sprintf("persistent shell startup failed; requested command was not run: %v", e.Err) |
| 87 | } |
| 88 | func (e *StartupError) Unwrap() error { return e.Err } |
| 89 | |
| 90 | type session struct { |
| 91 | powershell *powershellProcess |
| 92 | mu sync.Mutex |
| 93 | conn ptyConn |
| 94 | fp string |
| 95 | closed bool |
| 96 | san sanitizer |
| 97 | pendingRead chan readChunkResult |
| 98 | readerDone chan struct{} |
| 99 | } |
| 100 | |
| 101 | type readChunkResult struct { |
| 102 | data []byte |
| 103 | err error |
| 104 | } |
| 105 | |
| 106 | // New returns a Manager with zero controller owners. Callers must Retain |
| 107 | // before Run. |
| 108 | func New() *Manager { |
| 109 | return &Manager{} |
| 110 | } |
| 111 | |
| 112 | // OrNew returns m when it is non-nil, otherwise a fresh Manager. |
| 113 | func OrNew(m *Manager) *Manager { |
| 114 | if m != nil { |
| 115 | return m |
| 116 | } |
| 117 | return New() |
| 118 | } |
| 119 | |
| 120 | // Retain adds a Controller owner reference. Hot rebuilds Retain the shared |
| 121 | // Manager before publishing the replacement Controller. |
| 122 | func (m *Manager) Retain() { |
| 123 | if m == nil { |
| 124 | return |
| 125 | } |
| 126 | m.mu.Lock() |
| 127 | if !m.sealed { |
| 128 | m.owners++ |
| 129 | } |
| 130 | m.mu.Unlock() |
| 131 | } |
| 132 | |
| 133 | // Release drops a Controller owner reference. The last owner seals the manager |
| 134 | // and closes the live PTY. |
| 135 | func (m *Manager) Release() { |
| 136 | if m == nil { |
| 137 | return |
| 138 | } |
| 139 | m.mu.Lock() |
| 140 | if m.owners > 0 { |
| 141 | m.owners-- |
| 142 | } |
| 143 | if m.owners == 0 { |
| 144 | m.sealed = true |
| 145 | live := m.live |
| 146 | m.live = nil |
| 147 | m.mu.Unlock() |
| 148 | if live != nil { |
| 149 | live.close() |
| 150 | } |
| 151 | return |
| 152 | } |
| 153 | m.mu.Unlock() |
| 154 | } |
| 155 | |
| 156 | // Sealed reports whether the last Controller owner has released the Manager. |
| 157 | func (m *Manager) Sealed() bool { |
| 158 | if m == nil { |
| 159 | return true |
| 160 | } |
| 161 | m.mu.Lock() |
| 162 | defer m.mu.Unlock() |
| 163 | return m.sealed |
| 164 | } |
| 165 | |
| 166 | // Rotate closes the live PTY so the next Run starts a fresh shell. Rotate on a |
| 167 | // sealed Manager is a no-op. |
| 168 | func (m *Manager) Rotate() { |
| 169 | if m == nil { |
| 170 | return |
| 171 | } |
| 172 | m.mu.Lock() |
| 173 | if m.sealed { |
| 174 | m.mu.Unlock() |
| 175 | return |
| 176 | } |
| 177 | live := m.live |
| 178 | m.live = nil |
| 179 | m.startupFailure = nil |
| 180 | m.mu.Unlock() |
| 181 | if live != nil { |
| 182 | live.close() |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | // Close is Rotate plus an explicit shutdown of the live PTY without sealing. |
| 187 | // Tests call it to reap the process; controllers use Release. |
| 188 | func (m *Manager) Close() { |
| 189 | m.Rotate() |
| 190 | } |
| 191 | |
| 192 | // Run executes command in the session-scoped PTY, creating the shell on first |
| 193 | // use. Commands with a matching launch fingerprint reuse cwd and environment. |
| 194 | func (m *Manager) Run(ctx context.Context, req Request) Result { |
| 195 | if m == nil { |
| 196 | return failResult(fmt.Errorf("%w: manager is nil", ErrUnavailable), tool.ShellPhaseLaunch) |
| 197 | } |
| 198 | if !Supports(req.Shell) { |
| 199 | return failResult(fmt.Errorf("%w: %s", ErrUnavailable, unsupportedShellReason), tool.ShellPhaseLaunch) |
| 200 | } |
| 201 | m.runMu.Lock() |
| 202 | defer m.runMu.Unlock() |
| 203 | sess, err := m.sessionFor(req) |
| 204 | if err != nil { |
| 205 | m.metrics.startupFailed.Add(1) |
| 206 | return failResult(err, tool.ShellPhaseLaunch) |
| 207 | } |
| 208 | res := sess.run(ctx, req) |
| 209 | if res.Started { |
| 210 | m.metrics.started.Add(1) |
| 211 | } |
| 212 | if res.Started && !res.ExitCodeKnown { |
| 213 | m.metrics.completionMissing.Add(1) |
| 214 | } |
| 215 | if res.TimedOut { |
| 216 | m.metrics.timedOut.Add(1) |
| 217 | } |
| 218 | if res.ShellDied || res.TimedOut || res.Canceled { |
| 219 | m.metrics.reset.Add(1) |
| 220 | m.drop(sess) |
| 221 | res.Reset = true |
| 222 | } |
| 223 | return res |
| 224 | } |
| 225 | |
| 226 | func (m *Manager) sessionFor(req Request) (*session, error) { |
| 227 | fp := fingerprint(req) |
| 228 | m.mu.Lock() |
| 229 | if m.sealed || m.owners == 0 { |
| 230 | m.mu.Unlock() |
| 231 | return nil, fmt.Errorf("%w: manager closed", ErrUnavailable) |
| 232 | } |
| 233 | if m.startupFailure != nil && m.startupFingerprint == fp && time.Since(m.startupFailedAt) < 30*time.Second { |
| 234 | err := m.startupFailure |
| 235 | m.mu.Unlock() |
| 236 | return nil, err |
| 237 | } |
| 238 | if m.live != nil && m.live.fp != fp { |
| 239 | old := m.live |
| 240 | m.live = nil |
| 241 | m.mu.Unlock() |
| 242 | old.close() |
| 243 | m.mu.Lock() |
| 244 | if m.sealed || m.owners == 0 { |
| 245 | m.mu.Unlock() |
| 246 | return nil, fmt.Errorf("%w: manager closed", ErrUnavailable) |
| 247 | } |
| 248 | } |
| 249 | if m.live != nil { |
| 250 | sess := m.live |
| 251 | m.mu.Unlock() |
| 252 | return sess, nil |
| 253 | } |
| 254 | m.mu.Unlock() |
| 255 | |
| 256 | sess, err := startSession(req, fp) |
| 257 | if err != nil { |
| 258 | var startup *StartupError |
| 259 | if errors.As(err, &startup) { |
| 260 | m.mu.Lock() |
| 261 | m.startupFailure, m.startupFingerprint, m.startupFailedAt = startup, fp, time.Now() |
| 262 | m.mu.Unlock() |
| 263 | } |
| 264 | return nil, err |
| 265 | } |
| 266 | |
| 267 | m.mu.Lock() |
| 268 | if m.sealed || m.owners == 0 { |
| 269 | m.mu.Unlock() |
| 270 | sess.close() |
| 271 | return nil, fmt.Errorf("%w: manager closed", ErrUnavailable) |
| 272 | } |
| 273 | if m.live != nil { |
| 274 | // Another caller won the start race; keep the existing shell. |
| 275 | existing := m.live |
| 276 | m.mu.Unlock() |
| 277 | sess.close() |
| 278 | return existing, nil |
| 279 | } |
| 280 | m.live = sess |
| 281 | m.mu.Unlock() |
| 282 | return sess, nil |
| 283 | } |
| 284 | |
| 285 | func (m *Manager) drop(sess *session) { |
| 286 | if m == nil || sess == nil { |
| 287 | return |
| 288 | } |
| 289 | m.mu.Lock() |
| 290 | if m.live == sess { |
| 291 | m.live = nil |
| 292 | } |
| 293 | m.mu.Unlock() |
| 294 | sess.close() |
| 295 | } |
| 296 | |
| 297 | func fingerprint(req Request) string { |
| 298 | h := sha256.New() |
| 299 | for _, a := range req.Argv { |
| 300 | h.Write([]byte(a)) |
| 301 | h.Write([]byte{0}) |
| 302 | } |
| 303 | h.Write([]byte{1}) |
| 304 | h.Write([]byte(req.Dir)) |
| 305 | h.Write([]byte{1}) |
| 306 | for _, e := range req.Env { |
| 307 | h.Write([]byte(e)) |
| 308 | h.Write([]byte{0}) |
| 309 | } |
| 310 | h.Write([]byte{1}) |
| 311 | h.Write([]byte(req.Shell.Kind.String())) |
| 312 | return hex.EncodeToString(h.Sum(nil)) |
| 313 | } |
| 314 | |
| 315 | func failResult(err error, phase string) Result { |
| 316 | result := Result{ |
| 317 | Err: err, |
| 318 | State: tool.ShellStateFailed, |
| 319 | FailurePhase: phase, |
| 320 | } |
| 321 | var startup *StartupError |
| 322 | if errors.As(err, &startup) { |
| 323 | result.Output = startup.Output |
| 324 | } |
| 325 | return result |
| 326 | } |
| 327 | |
| 328 | func startSession(req Request, fp string) (*session, error) { |
| 329 | if req.Shell.Kind == sandbox.ShellPowerShell { |
| 330 | return startPowerShell(req, fp) |
| 331 | } |
| 332 | return startPOSIXSession(req, fp, startPTY) |
| 333 | } |
| 334 | |
| 335 | func startPOSIXSession(req Request, fp string, start func([]string, string, []string) (ptyConn, error)) (*session, error) { |
| 336 | conn, err := start(req.Argv, req.Dir, req.Env) |
| 337 | if err != nil { |
| 338 | return nil, err |
| 339 | } |
| 340 | s := &session{ |
| 341 | conn: conn, |
| 342 | fp: fp, |
| 343 | } |
| 344 | s.startReader() |
| 345 | ctx, cancel := context.WithTimeout(context.Background(), startupTimeout) |
| 346 | defer cancel() |
| 347 | if err := s.writeScript(posixSetupScript()); err != nil { |
| 348 | s.close() |
| 349 | return nil, err |
| 350 | } |
| 351 | var buf []byte |
| 352 | if err := s.pump(ctx, func(text string) bool { |
| 353 | buf = append(buf, text...) |
| 354 | if len(buf) > tool.OutputTailMaxBytes { |
| 355 | buf = buf[len(buf)-tool.OutputTailMaxBytes:] |
| 356 | } |
| 357 | return readyLine(string(buf)) |
| 358 | }); err != nil { |
| 359 | s.close() |
| 360 | return nil, &StartupError{Output: strings.ToValidUTF8(string(buf), "\uFFFD"), Err: err} |
| 361 | } |
| 362 | return s, nil |
| 363 | } |
| 364 | |
| 365 | func (s *session) startReader() { |
| 366 | s.pendingRead = make(chan readChunkResult, 1) |
| 367 | s.readerDone = make(chan struct{}) |
| 368 | go func() { |
| 369 | buf := make([]byte, readChunk) |
| 370 | for { |
| 371 | n, err := s.conn.Read(buf) |
| 372 | chunk := make([]byte, n) |
| 373 | copy(chunk, buf[:n]) |
| 374 | select { |
| 375 | case s.pendingRead <- readChunkResult{data: chunk, err: err}: |
| 376 | case <-s.readerDone: |
| 377 | return |
| 378 | } |
| 379 | if err != nil { |
| 380 | return |
| 381 | } |
| 382 | } |
| 383 | }() |
| 384 | } |
| 385 | |
| 386 | func (s *session) writeScript(script string) error { |
| 387 | _, err := s.conn.Write([]byte(script)) |
| 388 | return err |
| 389 | } |
| 390 | |
| 391 | func (s *session) run(ctx context.Context, req Request) Result { |
| 392 | if s.powershell != nil { |
| 393 | return s.runPowerShell(ctx, req) |
| 394 | } |
| 395 | s.mu.Lock() |
| 396 | if s.closed { |
| 397 | s.mu.Unlock() |
| 398 | return Result{ |
| 399 | ShellDied: true, |
| 400 | State: tool.ShellStateFailed, |
| 401 | FailurePhase: tool.ShellPhaseLaunch, |
| 402 | Err: errors.New("persistent shell closed"), |
| 403 | } |
| 404 | } |
| 405 | s.mu.Unlock() |
| 406 | |
| 407 | runCtx := ctx |
| 408 | var cancel context.CancelFunc |
| 409 | if req.Timeout > 0 { |
| 410 | runCtx, cancel = context.WithTimeout(ctx, req.Timeout) |
| 411 | defer cancel() |
| 412 | } |
| 413 | |
| 414 | id := newMarkerID() |
| 415 | start := "REASONIX_START_" + id |
| 416 | // The status digits must follow the end marker immediately, so echoed |
| 417 | // wrapper source can never fabricate a completion. |
| 418 | end := "REASONIX_END_" + id + ":" |
| 419 | if err := s.writeCommand(runCtx, req.Command, start, end); err != nil { |
| 420 | s.markClosed() |
| 421 | return Result{ |
| 422 | ShellDied: true, |
| 423 | TimedOut: errors.Is(err, context.DeadlineExceeded), |
| 424 | Canceled: errors.Is(err, context.Canceled), |
| 425 | State: tool.ShellStateFailed, |
| 426 | FailurePhase: tool.ShellPhasePreflight, |
| 427 | Err: err, |
| 428 | } |
| 429 | } |
| 430 | |
| 431 | capt := newCapture(start, end, req.Progress) |
| 432 | err := s.pump(runCtx, func(text string) bool { |
| 433 | capt.push(text) |
| 434 | return capt.done |
| 435 | }) |
| 436 | if capt.done { |
| 437 | res := Result{Output: capt.body(), ExitCode: capt.exitCode, ExitCodeKnown: true, Started: true} |
| 438 | if capt.exitCode != 0 { |
| 439 | res.State = tool.ShellStateFailed |
| 440 | res.FailurePhase = tool.ShellPhaseExecution |
| 441 | res.Err = fmt.Errorf("exit status %d", capt.exitCode) |
| 442 | } else { |
| 443 | res.State = tool.ShellStateCompleted |
| 444 | } |
| 445 | return res |
| 446 | } |
| 447 | // No status marker: whatever the command printed before it stopped is the |
| 448 | // only evidence the model gets, so it is reported rather than discarded. |
| 449 | res := Result{Output: capt.partial(), Started: true} |
| 450 | switch { |
| 451 | case ctx.Err() != nil && errors.Is(ctx.Err(), context.Canceled): |
| 452 | res.Canceled = true |
| 453 | res.State = tool.ShellStateCancelled |
| 454 | res.FailurePhase = tool.ShellPhaseCancellation |
| 455 | res.Err = ctx.Err() |
| 456 | case errors.Is(runCtx.Err(), context.DeadlineExceeded): |
| 457 | res.TimedOut = true |
| 458 | res.State = tool.ShellStateTimedOut |
| 459 | res.FailurePhase = tool.ShellPhaseTimeout |
| 460 | res.Err = fmt.Errorf("command timed out (> %s)", req.Timeout) |
| 461 | case err != nil: |
| 462 | res.ShellDied = true |
| 463 | res.State = tool.ShellStateFailed |
| 464 | res.FailurePhase = tool.ShellPhaseExecution |
| 465 | res.Err = err |
| 466 | default: |
| 467 | res.ShellDied = true |
| 468 | res.State = tool.ShellStateFailed |
| 469 | res.FailurePhase = tool.ShellPhaseExecution |
| 470 | res.Err = errors.New("persistent shell exited before command completed") |
| 471 | } |
| 472 | s.markClosed() |
| 473 | return res |
| 474 | } |
| 475 | |
| 476 | func (s *session) markClosed() { |
| 477 | s.mu.Lock() |
| 478 | already := s.closed |
| 479 | s.closed = true |
| 480 | conn := s.conn |
| 481 | s.mu.Unlock() |
| 482 | if !already && conn != nil { |
| 483 | if s.readerDone != nil { |
| 484 | close(s.readerDone) |
| 485 | } |
| 486 | _ = conn.Close() |
| 487 | } |
| 488 | } |
| 489 | |
| 490 | // pump feeds sanitized PTY reads to step until it reports completion. Each read |
| 491 | // is handed over once, so the cost of a command is linear in its output rather |
| 492 | // than quadratic in a re-scanned transcript. |
| 493 | func (s *session) pump(ctx context.Context, step func(string) bool) error { |
| 494 | for { |
| 495 | select { |
| 496 | case <-ctx.Done(): |
| 497 | return ctx.Err() |
| 498 | case chunk, ok := <-s.pendingRead: |
| 499 | if !ok { |
| 500 | return errors.New("persistent shell reader closed") |
| 501 | } |
| 502 | if len(chunk.data) > 0 && step(s.san.push(chunk.data)) { |
| 503 | return nil |
| 504 | } |
| 505 | if chunk.err != nil { |
| 506 | if text := s.san.flush(); text != "" && step(text) { |
| 507 | return nil |
| 508 | } |
| 509 | return chunk.err |
| 510 | } |
| 511 | } |
| 512 | } |
| 513 | } |
| 514 | |
| 515 | func (s *session) close() { |
| 516 | s.markClosed() |
| 517 | } |
| 518 | |
| 519 | const unsupportedShellReason = "unsupported persistent shell dialect" |
| 520 | |
| 521 | // Supports includes native PowerShell's framed transport and POSIX PTYs. |
| 522 | func Supports(sh sandbox.Shell) bool { |
| 523 | return sh.Kind == sandbox.ShellPowerShell || sh.Kind.IsPOSIX() |
| 524 | } |
| 525 | |
| 526 | // InteractiveArgv is the long-lived interpreter argv (no -c / -Command), |
| 527 | // before sandbox wrapping. Windows POSIX interpreters consume a pipe without |
| 528 | // interactive prompts or line editing; POSIX hosts retain their native PTY. |
| 529 | func InteractiveArgv(sh sandbox.Shell) []string { |
| 530 | return interactiveArgvForOS(sh, runtime.GOOS) |
| 531 | } |
| 532 | |
| 533 | func interactiveArgvForOS(sh sandbox.Shell, goos string) []string { |
| 534 | path := sh.Path |
| 535 | if path == "" { |
| 536 | path = sh.Kind.String() |
| 537 | } |
| 538 | switch sh.Kind { |
| 539 | case sandbox.ShellPowerShell: |
| 540 | return []string{path, "-NoLogo", "-NoProfile", "-NonInteractive", "-OutputFormat", "Text", "-EncodedCommand", encodedPowerShell(powershellBootstrap)} |
| 541 | case sandbox.ShellZsh: |
| 542 | if goos == "windows" { |
| 543 | return []string{path, "-f"} |
| 544 | } |
| 545 | return []string{path, "-f", "-i"} |
| 546 | case sandbox.ShellSh: |
| 547 | if goos == "windows" { |
| 548 | return []string{path} |
| 549 | } |
| 550 | return []string{path, "-i"} |
| 551 | default: |
| 552 | if goos == "windows" { |
| 553 | return []string{path, "--noprofile", "--norc"} |
| 554 | } |
| 555 | return []string{path, "--noprofile", "--norc", "-i"} |
| 556 | } |
| 557 | } |
| 558 |