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