Add reconnectable SSH exec sessions

api/openapi.yaml

@@ -429,11 +429,33 @@ paths:
       parameters:
         - in: query
           name: command
-          description: Command to execute
+          description: |
+            Command to execute.
+
+            Required when starting a new exec session. May be omitted when reconnecting to an
+            existing session identified by the `session` parameter.
           schema:
             type: string
             minLength: 1
-          required: true
+          required: false
+        - in: query
+          name: session
+          description: |
+            Optional stable exec session identifier. When present, websocket disconnects detach
+            from the command instead of terminating it, and later requests with the same VM name
+            and session id may reconnect and request buffered history.
+          schema:
+            type: string
+            minLength: 1
+          required: false
+        - in: query
+          name: cmux_session_id
+          description: |
+            Compatibility alias for `session`. Prefer `session` for new Orchard clients.
+          schema:
+            type: string
+            minLength: 1
+          required: false
         - in: query
           name: stdin
           description: |
@@ -483,7 +505,9 @@ paths:
         '400':
           description: Invalid parameters were supplied
         '404':
-          description: VM resource with the given name doesn't exist
+          description: VM resource with the given name or reconnectable exec session doesn't exist
+        '409':
+          description: Reconnectable exec session already exists with a different command
         '503':
           description: Controller failed to establish a connection with the VM
   /vms/{name}/ip:
@@ -800,10 +824,18 @@ components:
       description: WebSocket frame from Orchard Client to the Orchard Controller
       oneOf:
         - $ref: '#/components/schemas/ExecClientFrameStdin'
+        - $ref: '#/components/schemas/ExecClientFrameHistory'
+        - $ref: '#/components/schemas/ExecClientFrameAck'
+        - $ref: '#/components/schemas/ExecClientFrameDetach'
+        - $ref: '#/components/schemas/ExecClientFrameClose'
       discriminator:
         propertyName: type
         mapping:
           stdin: '#/components/schemas/ExecClientFrameStdin'
+          history: '#/components/schemas/ExecClientFrameHistory'
+          ack: '#/components/schemas/ExecClientFrameAck'
+          detach: '#/components/schemas/ExecClientFrameDetach'
+          close: '#/components/schemas/ExecClientFrameClose'
     ExecClientFrameStdin:
       description: Send bytes to the process standard input
       type: object
@@ -822,6 +854,56 @@ components:
       example:
         type: stdin
         data: aGVsbG8K
+    ExecClientFrameHistory:
+      description: Request buffered output strictly newer than the supplied watermark
+      type: object
+      required: [ type, watermark ]
+      properties:
+        type:
+          type: string
+          enum: [ history ]
+        watermark:
+          type: integer
+          format: int64
+          minimum: 0
+      example:
+        type: history
+        watermark: 42
+    ExecClientFrameAck:
+      description: Acknowledge that output has been durably consumed through the supplied watermark
+      type: object
+      required: [ type, watermark ]
+      properties:
+        type:
+          type: string
+          enum: [ ack ]
+        watermark:
+          type: integer
+          format: int64
+          minimum: 0
+      example:
+        type: ack
+        watermark: 42
+    ExecClientFrameDetach:
+      description: Detach this websocket while leaving the remote command running
+      type: object
+      required: [ type ]
+      properties:
+        type:
+          type: string
+          enum: [ detach ]
+      example:
+        type: detach
+    ExecClientFrameClose:
+      description: Close the reconnectable exec session and terminate the remote command
+      type: object
+      required: [ type ]
+      properties:
+        type:
+          type: string
+          enum: [ close ]
+      example:
+        type: close
     ExecControllerFrame:
       description: WebSocket frame from Orchard Controller to the Orchard Client
       oneOf:
@@ -829,6 +911,7 @@ components:
         - $ref: '#/components/schemas/ExecControllerFrameStderr'
         - $ref: '#/components/schemas/ExecControllerFrameExit'
         - $ref: '#/components/schemas/ExecControllerFrameError'
+        - $ref: '#/components/schemas/ExecControllerFrameNoMoreHistory'
       discriminator:
         propertyName: type
         mapping:
@@ -836,6 +919,7 @@ components:
           stderr: '#/components/schemas/ExecControllerFrameStderr'
           exit: '#/components/schemas/ExecControllerFrameExit'
           error: '#/components/schemas/ExecControllerFrameError'
+          no_more_history: '#/components/schemas/ExecControllerFrameNoMoreHistory'
     ExecControllerFrameStdout:
       description: Standard output from the process
       type: object
@@ -848,6 +932,10 @@ components:
           type: string
           format: byte
           description: Base64-encoded standard output bytes from the process
+        watermark:
+          type: integer
+          format: int64
+          description: Monotonic output watermark present on reconnectable sessions
       example:
         type: stdout
         data: aGVsbG8K
@@ -863,6 +951,10 @@ components:
           type: string
           format: byte
           description: Base64-encoded standard error bytes from the process
+        watermark:
+          type: integer
+          format: int64
+          description: Monotonic output watermark present on reconnectable sessions
       example:
         type: stderr
         data: aGVsbG8K
@@ -882,6 +974,10 @@ components:
               type: integer
               format: int32
               description: Process exit code
+        watermark:
+          type: integer
+          format: int64
+          description: Monotonic output watermark present on reconnectable sessions
       example:
         type: exit
         exit:
@@ -897,9 +993,28 @@ components:
         error:
           type: string
           description: Error message text
+        watermark:
+          type: integer
+          format: int64
+          description: Monotonic output watermark present on reconnectable sessions
       example:
         type: error
         error: Failed to establish SSH connection to a VM
+    ExecControllerFrameNoMoreHistory:
+      description: Marker indicating that the requested replay range has been fully sent
+      type: object
+      required: [ type, watermark ]
+      properties:
+        type:
+          type: string
+          enum: [ no_more_history ]
+        watermark:
+          type: integer
+          format: int64
+          description: Highest watermark known to the exec session at the time of replay
+      example:
+        type: no_more_history
+        watermark: 42
     Event:
       title: Generic Resource Event
       type: object

internal/command/controller/run.go

@@ -35,6 +35,7 @@ var noExperimentalRPCV2 bool
 var experimentalPingInterval time.Duration
 var experimentalDisableDBCompression bool
 var workerOfflineTimeout time.Duration
+var execSessionExitTTL time.Duration
 var synthetic bool
 
 func newRunCommand() *cobra.Command {
@@ -87,6 +88,8 @@ func newRunCommand() *cobra.Command {
 		"duration (e.g. 60s or 5m30s) after which a worker is considered offline for the purposes "+
 			"of scheduling (no new VMs will be scheduled on such worker and already assigned VMs will be "+
 			"marked as failed)")
+	cmd.Flags().DurationVar(&execSessionExitTTL, "exec-session-exit-ttl", 10*time.Minute,
+		"duration to retain reconnectable exec session history after the command exits")
 
 	// Hidden flags
 	cmd.Flags().BoolVar(&synthetic, "synthetic", false, "")
@@ -147,6 +150,7 @@ func runController(cmd *cobra.Command, args []string) (err error) {
 		controller.WithListenAddr(address),
 		controller.WithDataDir(dataDir),
 		controller.WithWorkerOfflineTimeout(workerOfflineTimeout),
+		controller.WithExecSessionExitTTL(execSessionExitTTL),
 		controller.WithLogger(logger),
 	}
 

internal/controller/api_vms_exec.go

@@ -28,9 +28,13 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
 
 	// Retrieve and parse path and query parameters
 	name := ctx.Param("name")
+	sessionID := ctx.Query("session")
+	if sessionID == "" {
+		sessionID = ctx.Query("cmux_session_id")
+	}
 
 	command := ctx.Query("command")
-	if command == "" {
+	if sessionID == "" && command == "" {
 		return responder.JSON(http.StatusBadRequest,
 			NewErrorResponse("\"command\" parameter cannot be empty"))
 	}
@@ -43,6 +47,20 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
 		return responder.Code(http.StatusBadRequest)
 	}
 
+	if sessionID != "" {
+		return controller.execVMReconnectable(ctx, name, sessionID, command, stdin, wait)
+	}
+
+	return controller.execVMLegacy(ctx, name, command, stdin, wait)
+}
+
+func (controller *Controller) execVMLegacy(
+	ctx *gin.Context,
+	name string,
+	command string,
+	stdin bool,
+	wait uint64,
+) responder.Responder {
 	// Look-up the VM
 	waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
 	defer waitContextCancel()
@@ -151,6 +169,198 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
 	}
 }
 
+func (controller *Controller) execVMReconnectable(
+	ctx *gin.Context,
+	name string,
+	sessionID string,
+	command string,
+	stdin bool,
+	wait uint64,
+) responder.Responder {
+	key := execSessionKey{
+		vmName:    name,
+		sessionID: sessionID,
+	}
+
+	session, ok := controller.execSessions.get(key)
+	if ok {
+		if !session.commandMatches(command) {
+			return responder.JSON(http.StatusConflict,
+				NewErrorResponse("exec session %q is already running a different command", sessionID))
+		}
+	} else {
+		if command == "" {
+			return responder.JSON(http.StatusNotFound,
+				NewErrorResponse("exec session %q does not exist", sessionID))
+		}
+
+		waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
+		defer waitContextCancel()
+
+		vm, responderImpl := controller.waitForVM(waitContext, name)
+		if responderImpl != nil {
+			return responderImpl
+		}
+
+		var err error
+		session, _, err = controller.execSessions.getOrCreate(waitContext, key, func() (*execSession, error) {
+			portForwardConn, err := retry.NewWithData[net.Conn](
+				retry.Context(waitContext),
+				retry.DelayType(retry.FixedDelay),
+				retry.Delay(time.Second),
+				retry.Attempts(0),
+				retry.LastErrorOnly(true),
+			).Do(func() (net.Conn, error) {
+				return controller.portForwardConnection(ctx, waitContext, vm.Worker, vm.UID, 22)
+			})
+			if err != nil {
+				return nil, err
+			}
+
+			exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin)
+			if err != nil {
+				_ = portForwardConn.Close()
+
+				return nil, err
+			}
+
+			return newExecSession(
+				key,
+				command,
+				exec,
+				portForwardConn,
+				controller.execSessions,
+				controller.execSessionExitTTL,
+			), nil
+		})
+		if err != nil {
+			return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err))
+		}
+
+		if !session.commandMatches(command) {
+			return responder.JSON(http.StatusConflict,
+				NewErrorResponse("exec session %q is already running a different command", sessionID))
+		}
+	}
+
+	wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
+		OriginPatterns: []string{"*"},
+	})
+	if err != nil {
+		return responder.Error(err)
+	}
+	defer func() {
+		_ = wsConn.CloseNow()
+	}()
+
+	subscriber, err := session.attach()
+	if err != nil {
+		_ = wsConn.Close(websocket.StatusNormalClosure, err.Error())
+
+		return responder.Empty()
+	}
+	defer session.detach(subscriber)
+
+	readFramesErrCh := make(chan error, 1)
+	go func() {
+		readFramesErrCh <- controller.readReconnectableFrames(ctx, wsConn, session, subscriber)
+	}()
+
+	for {
+		select {
+		case readFramesErr := <-readFramesErrCh:
+			if readFramesErr != nil &&
+				!errors.Is(readFramesErr, errExecSessionDetached) &&
+				!errors.Is(readFramesErr, errExecSessionClosed) {
+				controller.logger.Warnf("failed to read and process reconnectable exec frames from WebSocket: %v",
+					readFramesErr)
+			}
+
+			return responder.Empty()
+		case outgoingFrame, ok := <-subscriber.frames:
+			if !ok {
+				if err := wsConn.Close(websocket.StatusNormalClosure, "Command finished"); err != nil {
+					controller.logger.Warnf("exec: failed to close WebSocket cleanly: %v", err)
+				}
+
+				return responder.Empty()
+			}
+
+			if err := execstream.WriteFrame(ctx, wsConn, outgoingFrame); err != nil {
+				controller.logger.Warnf("failed to write reconnectable exec frame to the client: %v", err)
+
+				return responder.Empty()
+			}
+		case <-time.After(controller.pingInterval):
+			pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second)
+
+			if err := wsConn.Ping(pingCtx); err != nil {
+				controller.logger.Warnf("reconnectable exec: failed to ping the client, "+
+					"connection might time out: %v", err)
+			}
+
+			pingCtxCancel()
+		case <-ctx.Done():
+			controller.logger.Warnf("client disconnected prematurely")
+
+			return responder.Empty()
+		}
+	}
+}
+
+var (
+	errExecSessionDetached = errors.New("exec session detached")
+	errExecSessionClosed   = errors.New("exec session closed")
+)
+
+func (controller *Controller) readReconnectableFrames(
+	ctx context.Context,
+	wsConn *websocket.Conn,
+	session *execSession,
+	subscriber *execSessionSubscriber,
+) error {
+	for {
+		var frame execstream.Frame
+
+		messageType, payloadBytes, err := wsConn.Read(ctx)
+		if err != nil {
+			var closeErr websocket.CloseError
+			if errors.As(err, &closeErr) && closeErr.Code == websocket.StatusNormalClosure {
+				return errExecSessionDetached
+			}
+
+			return fmt.Errorf("failed to read next frame from WebSocket: %w", err)
+		}
+
+		if messageType != websocket.MessageText {
+			continue
+		}
+
+		if err := json.Unmarshal(payloadBytes, &frame); err != nil {
+			return err
+		}
+
+		switch frame.Type {
+		case execstream.FrameTypeStdin:
+			if err := session.writeStdin(frame.Data); err != nil {
+				return fmt.Errorf("failed to handle %q frame: %w", frame.Type, err)
+			}
+		case execstream.FrameTypeHistory:
+			session.sendHistory(subscriber, frame.Watermark)
+		case execstream.FrameTypeAck:
+			session.ack(frame.Watermark)
+		case execstream.FrameTypeDetach:
+			return errExecSessionDetached
+		case execstream.FrameTypeClose:
+			session.close()
+
+			return errExecSessionClosed
+		default:
+			return fmt.Errorf("unexpected frame type received: %q", frame.Type)
+		}
+	}
+}
+
 func (controller *Controller) readFrames(
 	ctx context.Context,
 	wsConn *websocket.Conn,

internal/controller/controller.go

@@ -55,6 +55,7 @@ type Controller struct {
 	ipRendezvous         *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]]
 	enableSwaggerDocs    bool
 	workerOfflineTimeout time.Duration
+	execSessionExitTTL   time.Duration
 	experimentalRPCV2    bool
 	disableDBCompression bool
 	pingInterval         time.Duration
@@ -64,6 +65,7 @@ type Controller struct {
 	sshSigner       ssh.Signer
 	sshNoClientAuth bool
 	sshServer       *sshserver.SSHServer
+	execSessions    *execSessionRegistry
 
 	single singleflight.Group
 
@@ -75,7 +77,9 @@ func New(opts ...Option) (*Controller, error) {
 		connRendezvous:       rendezvous.New[rendezvous.ResultWithErrorMessage[net.Conn]](),
 		ipRendezvous:         rendezvous.New[rendezvous.ResultWithErrorMessage[string]](),
 		workerOfflineTimeout: 3 * time.Minute,
+		execSessionExitTTL:   10 * time.Minute,
 		pingInterval:         30 * time.Second,
+		execSessions:         newExecSessionRegistry(),
 		single:               singleflight.Group{},
 	}
 
@@ -308,6 +312,8 @@ func (controller *Controller) Run(ctx context.Context) error {
 	go func() {
 		<-ctx.Done()
 
+		controller.execSessions.closeAll()
+
 		if err := controller.httpServer.Shutdown(ctx); err != nil {
 			controller.logger.Errorf("failed to cleanly shutdown the HTTP server: %v", err)
 		}

internal/controller/exec_sessions.go

@@ -0,0 +1,462 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/cirruslabs/orchard/internal/execstream"
+)
+
+const execSessionReplayBufferBytes = 4 * 1024 * 1024
+
+type sshExecRunner interface {
+	Stdin() io.WriteCloser
+	Run(ctx context.Context, command string, outgoingFrames chan<- *execstream.Frame) error
+	Close() error
+}
+
+type execSessionKey struct {
+	vmName    string
+	sessionID string
+}
+
+type execSessionCreation struct {
+	done    chan struct{}
+	session *execSession
+	err     error
+}
+
+type execSessionRegistry struct {
+	mu       sync.Mutex
+	sessions map[execSessionKey]*execSession
+	creating map[execSessionKey]*execSessionCreation
+}
+
+func newExecSessionRegistry() *execSessionRegistry {
+	return &execSessionRegistry{
+		sessions: map[execSessionKey]*execSession{},
+		creating: map[execSessionKey]*execSessionCreation{},
+	}
+}
+
+func (registry *execSessionRegistry) get(key execSessionKey) (*execSession, bool) {
+	registry.mu.Lock()
+	defer registry.mu.Unlock()
+
+	session, ok := registry.sessions[key]
+
+	return session, ok
+}
+
+func (registry *execSessionRegistry) getOrCreate(
+	ctx context.Context,
+	key execSessionKey,
+	create func() (*execSession, error),
+) (*execSession, bool, error) {
+	for {
+		registry.mu.Lock()
+
+		if session, ok := registry.sessions[key]; ok {
+			registry.mu.Unlock()
+
+			return session, false, nil
+		}
+
+		if creation, ok := registry.creating[key]; ok {
+			registry.mu.Unlock()
+
+			select {
+			case <-ctx.Done():
+				return nil, false, ctx.Err()
+			case <-creation.done:
+				if creation.err != nil {
+					return nil, false, creation.err
+				}
+
+				return creation.session, false, nil
+			}
+		}
+
+		creation := &execSessionCreation{done: make(chan struct{})}
+		registry.creating[key] = creation
+		registry.mu.Unlock()
+
+		session, err := create()
+
+		registry.mu.Lock()
+		delete(registry.creating, key)
+		if err == nil {
+			registry.sessions[key] = session
+		}
+		creation.session = session
+		creation.err = err
+		close(creation.done)
+		registry.mu.Unlock()
+
+		return session, true, err
+	}
+}
+
+func (registry *execSessionRegistry) remove(key execSessionKey, expected *execSession) {
+	registry.mu.Lock()
+	defer registry.mu.Unlock()
+
+	if registry.sessions[key] == expected {
+		delete(registry.sessions, key)
+	}
+}
+
+func (registry *execSessionRegistry) closeAll() {
+	registry.mu.Lock()
+	sessions := make([]*execSession, 0, len(registry.sessions))
+	for _, session := range registry.sessions {
+		sessions = append(sessions, session)
+	}
+	registry.mu.Unlock()
+
+	for _, session := range sessions {
+		session.close()
+	}
+}
+
+type execReplayFrame struct {
+	frame *execstream.Frame
+	size  int
+}
+
+type execSessionSubscriber struct {
+	frames chan *execstream.Frame
+}
+
+func newExecSessionSubscriber() *execSessionSubscriber {
+	return &execSessionSubscriber{
+		frames: make(chan *execstream.Frame, 128),
+	}
+}
+
+func (subscriber *execSessionSubscriber) enqueue(frame *execstream.Frame) bool {
+	select {
+	case subscriber.frames <- cloneExecFrame(frame):
+		return true
+	default:
+		return false
+	}
+}
+
+type execSession struct {
+	key       execSessionKey
+	command   string
+	exec      sshExecRunner
+	transport net.Conn
+	registry  *execSessionRegistry
+	exitTTL   time.Duration
+
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	mu             sync.Mutex
+	stdin          io.WriteCloser
+	stdinClosed    bool
+	subscribers    map[*execSessionSubscriber]struct{}
+	frames         []execReplayFrame
+	bufferBytes    int
+	nextWatermark  uint64
+	ackedWatermark uint64
+	finished       bool
+	closed         bool
+	expiryTimer    *time.Timer
+
+	done     chan struct{}
+	doneOnce sync.Once
+}
+
+func newExecSession(
+	key execSessionKey,
+	command string,
+	exec sshExecRunner,
+	transport net.Conn,
+	registry *execSessionRegistry,
+	exitTTL time.Duration,
+) *execSession {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	session := &execSession{
+		key:         key,
+		command:     command,
+		exec:        exec,
+		transport:   transport,
+		registry:    registry,
+		exitTTL:     exitTTL,
+		ctx:         ctx,
+		cancel:      cancel,
+		stdin:       exec.Stdin(),
+		subscribers: map[*execSessionSubscriber]struct{}{},
+		done:        make(chan struct{}),
+	}
+
+	go session.run()
+
+	return session
+}
+
+func (session *execSession) commandMatches(command string) bool {
+	return command == "" || session.command == command
+}
+
+func (session *execSession) attach() (*execSessionSubscriber, error) {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+
+	if session.closed {
+		return nil, errors.New("exec session is closed")
+	}
+
+	subscriber := newExecSessionSubscriber()
+	session.subscribers[subscriber] = struct{}{}
+
+	return subscriber, nil
+}
+
+func (session *execSession) detach(subscriber *execSessionSubscriber) {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+
+	session.detachLocked(subscriber)
+}
+
+func (session *execSession) detachLocked(subscriber *execSessionSubscriber) {
+	if _, ok := session.subscribers[subscriber]; !ok {
+		return
+	}
+
+	delete(session.subscribers, subscriber)
+	close(subscriber.frames)
+}
+
+func (session *execSession) writeStdin(data []byte) error {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+
+	if session.stdin == nil || session.stdinClosed {
+		return errors.New("this exec session has no stdin enabled or it is already closed")
+	}
+
+	if len(data) == 0 {
+		if err := session.stdin.Close(); err != nil {
+			return err
+		}
+
+		session.stdinClosed = true
+
+		return nil
+	}
+
+	_, err := session.stdin.Write(data)
+
+	return err
+}
+
+func (session *execSession) ack(watermark uint64) {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+
+	if watermark <= session.ackedWatermark {
+		return
+	}
+
+	session.ackedWatermark = watermark
+	session.trimAcknowledgedLocked()
+}
+
+func (session *execSession) sendHistory(
+	subscriber *execSessionSubscriber,
+	watermark uint64,
+) {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+
+	if _, ok := session.subscribers[subscriber]; !ok {
+		return
+	}
+
+	for _, record := range session.frames {
+		if record.frame.Watermark <= watermark {
+			continue
+		}
+
+		if !subscriber.enqueue(record.frame) {
+			session.detachLocked(subscriber)
+
+			return
+		}
+	}
+
+	if !subscriber.enqueue(&execstream.Frame{
+		Type:      execstream.FrameTypeNoMoreHistory,
+		Watermark: session.nextWatermark,
+	}) {
+		session.detachLocked(subscriber)
+	}
+}
+
+func (session *execSession) close() {
+	session.mu.Lock()
+	if session.closed {
+		session.mu.Unlock()
+
+		return
+	}
+
+	session.closed = true
+	if session.expiryTimer != nil {
+		session.expiryTimer.Stop()
+		session.expiryTimer = nil
+	}
+
+	subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers))
+	for subscriber := range session.subscribers {
+		subscribers = append(subscribers, subscriber)
+	}
+	session.subscribers = map[*execSessionSubscriber]struct{}{}
+	session.mu.Unlock()
+
+	for _, subscriber := range subscribers {
+		close(subscriber.frames)
+	}
+
+	session.cancel()
+	_ = session.exec.Close()
+	if session.transport != nil {
+		_ = session.transport.Close()
+	}
+	session.registry.remove(session.key, session)
+}
+
+func (session *execSession) run() {
+	outgoingFrames := make(chan *execstream.Frame)
+	runErrCh := make(chan error, 1)
+
+	go func() {
+		runErrCh <- session.exec.Run(session.ctx, session.command, outgoingFrames)
+		close(outgoingFrames)
+	}()
+
+	for frame := range outgoingFrames {
+		session.recordFrame(frame)
+	}
+
+	runErr := <-runErrCh
+	if runErr != nil && !errors.Is(runErr, context.Canceled) {
+		session.recordFrame(&execstream.Frame{
+			Type:  execstream.FrameTypeError,
+			Error: runErr.Error(),
+		})
+	}
+
+	session.markFinished()
+}
+
+func (session *execSession) recordFrame(frame *execstream.Frame) {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+
+	if session.closed {
+		return
+	}
+
+	session.nextWatermark++
+	frame = cloneExecFrame(frame)
+	frame.Watermark = session.nextWatermark
+
+	session.frames = append(session.frames, execReplayFrame{
+		frame: frame,
+		size:  execFrameSize(frame),
+	})
+	session.bufferBytes += execFrameSize(frame)
+	session.trimAcknowledgedLocked()
+	session.trimToLimitLocked()
+
+	for subscriber := range session.subscribers {
+		if subscriber.enqueue(frame) {
+			continue
+		}
+
+		session.detachLocked(subscriber)
+	}
+}
+
+func (session *execSession) markFinished() {
+	session.mu.Lock()
+	if session.finished {
+		session.mu.Unlock()
+
+		return
+	}
+
+	session.finished = true
+	if !session.closed {
+		session.expiryTimer = time.AfterFunc(session.exitTTL, session.expire)
+	}
+
+	subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers))
+	for subscriber := range session.subscribers {
+		subscribers = append(subscribers, subscriber)
+	}
+	session.subscribers = map[*execSessionSubscriber]struct{}{}
+	session.mu.Unlock()
+
+	for _, subscriber := range subscribers {
+		close(subscriber.frames)
+	}
+
+	session.doneOnce.Do(func() {
+		close(session.done)
+	})
+}
+
+func (session *execSession) expire() {
+	session.close()
+}
+
+func (session *execSession) trimAcknowledgedLocked() {
+	for len(session.frames) > 0 && session.frames[0].frame.Watermark <= session.ackedWatermark {
+		session.bufferBytes -= session.frames[0].size
+		session.frames = session.frames[1:]
+	}
+}
+
+func (session *execSession) trimToLimitLocked() {
+	for session.bufferBytes > execSessionReplayBufferBytes && len(session.frames) > 0 {
+		session.bufferBytes -= session.frames[0].size
+		session.frames = session.frames[1:]
+	}
+}
+
+func cloneExecFrame(frame *execstream.Frame) *execstream.Frame {
+	if frame == nil {
+		return nil
+	}
+
+	clone := *frame
+	if frame.Data != nil {
+		clone.Data = append([]byte(nil), frame.Data...)
+	}
+	if frame.Exit != nil {
+		exit := *frame.Exit
+		clone.Exit = &exit
+	}
+
+	return &clone
+}
+
+func execFrameSize(frame *execstream.Frame) int {
+	if frame == nil {
+		return 0
+	}
+
+	return len(frame.Data) + len(frame.Error) + 16
+}

internal/controller/exec_sessions_test.go

@@ -0,0 +1,172 @@
+package controller
+
+import (
+	"context"
+	"io"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/cirruslabs/orchard/internal/execstream"
+	"github.com/stretchr/testify/require"
+)
+
+type fakeExec struct {
+	stdin      io.WriteCloser
+	run        func(context.Context, string, chan<- *execstream.Frame) error
+	closeCalls atomic.Int32
+}
+
+func (exec *fakeExec) Stdin() io.WriteCloser {
+	return exec.stdin
+}
+
+func (exec *fakeExec) Run(
+	ctx context.Context,
+	command string,
+	outgoingFrames chan<- *execstream.Frame,
+) error {
+	if exec.run != nil {
+		return exec.run(ctx, command, outgoingFrames)
+	}
+
+	return nil
+}
+
+func (exec *fakeExec) Close() error {
+	exec.closeCalls.Add(1)
+
+	return nil
+}
+
+func newManualExecSessionForTest(
+	key execSessionKey,
+	registry *execSessionRegistry,
+) *execSession {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	return &execSession{
+		key:         key,
+		command:     "echo test",
+		exec:        &fakeExec{},
+		registry:    registry,
+		exitTTL:     time.Minute,
+		ctx:         ctx,
+		cancel:      cancel,
+		subscribers: map[*execSessionSubscriber]struct{}{},
+		done:        make(chan struct{}),
+	}
+}
+
+func TestExecSessionRegistryGetOrCreateReusesInflightCreation(t *testing.T) {
+	registry := newExecSessionRegistry()
+	key := execSessionKey{vmName: "vm", sessionID: "session"}
+
+	createStarted := make(chan struct{})
+	releaseCreate := make(chan struct{})
+	var createCalls atomic.Int32
+
+	create := func() (*execSession, error) {
+		createCalls.Add(1)
+		close(createStarted)
+		<-releaseCreate
+
+		return newManualExecSessionForTest(key, registry), nil
+	}
+
+	firstDone := make(chan struct{})
+	go func() {
+		defer close(firstDone)
+		_, _, err := registry.getOrCreate(context.Background(), key, create)
+		require.NoError(t, err)
+	}()
+
+	<-createStarted
+
+	secondDone := make(chan struct{})
+	go func() {
+		defer close(secondDone)
+		_, created, err := registry.getOrCreate(context.Background(), key, create)
+		require.NoError(t, err)
+		require.False(t, created)
+	}()
+
+	close(releaseCreate)
+
+	<-firstDone
+	<-secondDone
+	require.EqualValues(t, 1, createCalls.Load())
+}
+
+func TestExecSessionHistoryReplayAndAck(t *testing.T) {
+	registry := newExecSessionRegistry()
+	session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
+
+	session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")})
+	session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStderr, Data: []byte("err")})
+	session.recordFrame(&execstream.Frame{
+		Type: execstream.FrameTypeExit,
+		Exit: &execstream.Exit{Code: 7},
+	})
+
+	subscriber, err := session.attach()
+	require.NoError(t, err)
+
+	session.sendHistory(subscriber, 0)
+
+	require.Equal(t, execstream.FrameTypeStdout, (<-subscriber.frames).Type)
+	require.Equal(t, execstream.FrameTypeStderr, (<-subscriber.frames).Type)
+	require.Equal(t, execstream.FrameTypeExit, (<-subscriber.frames).Type)
+	noMoreHistory := <-subscriber.frames
+	require.Equal(t, execstream.FrameTypeNoMoreHistory, noMoreHistory.Type)
+	require.EqualValues(t, 3, noMoreHistory.Watermark)
+
+	session.ack(2)
+	require.Len(t, session.frames, 1)
+	require.EqualValues(t, 3, session.frames[0].frame.Watermark)
+}
+
+func TestExecSessionDetachKeepsProcessAlive(t *testing.T) {
+	registry := newExecSessionRegistry()
+	session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
+	exec := session.exec.(*fakeExec)
+
+	subscriber, err := session.attach()
+	require.NoError(t, err)
+
+	session.detach(subscriber)
+
+	require.False(t, session.closed)
+	require.EqualValues(t, 0, exec.closeCalls.Load())
+}
+
+func TestExecSessionCloseStopsProcessAndRemovesRegistryEntry(t *testing.T) {
+	registry := newExecSessionRegistry()
+	key := execSessionKey{vmName: "vm", sessionID: "session"}
+	session := newManualExecSessionForTest(key, registry)
+	exec := session.exec.(*fakeExec)
+	registry.sessions[key] = session
+
+	session.close()
+
+	require.True(t, session.closed)
+	require.EqualValues(t, 1, exec.closeCalls.Load())
+	_, ok := registry.get(key)
+	require.False(t, ok)
+}
+
+func TestExecSessionFinishedEntryExpiresAfterTTL(t *testing.T) {
+	registry := newExecSessionRegistry()
+	key := execSessionKey{vmName: "vm", sessionID: "session"}
+	session := newManualExecSessionForTest(key, registry)
+	session.exitTTL = 10 * time.Millisecond
+	registry.sessions[key] = session
+
+	session.markFinished()
+
+	require.Eventually(t, func() bool {
+		_, ok := registry.get(key)
+
+		return !ok
+	}, time.Second, 10*time.Millisecond)
+}

internal/controller/option.go

@@ -60,6 +60,12 @@ func WithWorkerOfflineTimeout(workerOfflineTimeout time.Duration) Option {
 	}
 }
 
+func WithExecSessionExitTTL(execSessionExitTTL time.Duration) Option {
+	return func(controller *Controller) {
+		controller.execSessionExitTTL = execSessionExitTTL
+	}
+}
+
 func WithExperimentalRPCV2() Option {
 	return func(controller *Controller) {
 		controller.experimentalRPCV2 = true

internal/execstream/frame.go

@@ -15,6 +15,11 @@ const (
 	FrameTypeStderr        FrameType = "stderr"
 	FrameTypeExit          FrameType = "exit"
 	FrameTypeError         FrameType = "error"
+	FrameTypeHistory       FrameType = "history"
+	FrameTypeNoMoreHistory FrameType = "no_more_history"
+	FrameTypeAck           FrameType = "ack"
+	FrameTypeDetach        FrameType = "detach"
+	FrameTypeClose         FrameType = "close"
 )
 
 type Frame struct {
@@ -22,6 +27,7 @@ type Frame struct {
 	Data      []byte    `json:"data,omitempty"`
 	Exit      *Exit     `json:"exit,omitempty"`
 	Error     string    `json:"error,omitempty"`
+	Watermark uint64    `json:"watermark,omitempty"`
 }
 
 type Exit struct {

internal/execstream/frame_test.go

@@ -0,0 +1,23 @@
+package execstream
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFrameRoundTripsWatermark(t *testing.T) {
+	frame := Frame{
+		Type:      FrameTypeHistory,
+		Watermark: 42,
+	}
+
+	payload, err := json.Marshal(frame)
+	require.NoError(t, err)
+
+	var decoded Frame
+	err = json.Unmarshal(payload, &decoded)
+	require.NoError(t, err)
+	require.Equal(t, frame, decoded)
+}

internal/tests/exec_test.go

@@ -132,6 +132,167 @@ func TestVMExecScript(t *testing.T) {
 	require.Equal(t, websocket.StatusNormalClosure, closeError.Code)
 }
 
+func TestVMExecSessionReconnectHistory(t *testing.T) {
+	devClient, vmName := prepareForExec(t)
+	sessionID := uuid.NewString()
+
+	wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		Command:     "sh -c 'echo first; sleep 1; echo second'",
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+
+	firstFrame := readFrame(t, wsConn)
+	require.Equal(t, execstream.FrameTypeStdout, firstFrame.Type)
+	require.Equal(t, "first\n", string(firstFrame.Data))
+	require.EqualValues(t, 1, firstFrame.Watermark)
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
+	require.NoError(t, err)
+	_ = wsConn.CloseNow()
+
+	wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+	defer wsConn.CloseNow()
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
+		Type:      execstream.FrameTypeHistory,
+		Watermark: firstFrame.Watermark,
+	})
+	require.NoError(t, err)
+
+	frames := readFramesUntilExit(t, wsConn)
+	require.Len(t, framesByType(frames, execstream.FrameTypeStdout), 1)
+	require.Equal(t, "second\n", string(framesByType(frames, execstream.FrameTypeStdout)[0].Data))
+	require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code)
+}
+
+func TestVMExecSessionReconnectAfterExit(t *testing.T) {
+	devClient, vmName := prepareForExec(t)
+	sessionID := uuid.NewString()
+
+	wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		Command:     "sh -c 'echo replay-me'",
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
+	require.NoError(t, err)
+	_ = wsConn.CloseNow()
+
+	time.Sleep(time.Second)
+
+	wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+	defer wsConn.CloseNow()
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory})
+	require.NoError(t, err)
+
+	frames := readFramesUntilExit(t, wsConn)
+	require.Equal(t, "replay-me\n", string(framesByType(frames, execstream.FrameTypeStdout)[0].Data))
+	require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code)
+}
+
+func TestVMExecSessionReplayPreservesStreams(t *testing.T) {
+	devClient, vmName := prepareForExec(t)
+	sessionID := uuid.NewString()
+
+	wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		Command:     "sh -c 'echo out1; sleep 1; echo err1 >&2; sleep 1; echo out2; sleep 1; echo err2 >&2'",
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
+	require.NoError(t, err)
+	_ = wsConn.CloseNow()
+
+	time.Sleep(4 * time.Second)
+
+	wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+	defer wsConn.CloseNow()
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory})
+	require.NoError(t, err)
+
+	frames := readFramesUntilExit(t, wsConn)
+	require.Equal(t, []execstream.FrameType{
+		execstream.FrameTypeStdout,
+		execstream.FrameTypeStderr,
+		execstream.FrameTypeStdout,
+		execstream.FrameTypeStderr,
+		execstream.FrameTypeExit,
+	}, frameTypes(frames))
+	require.Equal(t, "out1\n", string(frames[0].Data))
+	require.Equal(t, "err1\n", string(frames[1].Data))
+	require.Equal(t, "out2\n", string(frames[2].Data))
+	require.Equal(t, "err2\n", string(frames[3].Data))
+}
+
+func TestVMExecSessionStdinSurvivesReconnect(t *testing.T) {
+	devClient, vmName := prepareForExec(t)
+	sessionID := uuid.NewString()
+
+	wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		Command:     "/bin/cat",
+		Stdin:       true,
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
+		Type: execstream.FrameTypeStdin,
+		Data: []byte("one\n"),
+	})
+	require.NoError(t, err)
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
+	require.NoError(t, err)
+	_ = wsConn.CloseNow()
+
+	wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
+		WaitSeconds: 30,
+		Session:     sessionID,
+	})
+	require.NoError(t, err)
+	defer wsConn.CloseNow()
+
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
+		Type: execstream.FrameTypeStdin,
+		Data: []byte("two\n"),
+	})
+	require.NoError(t, err)
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
+		Type: execstream.FrameTypeStdin,
+		Data: []byte{},
+	})
+	require.NoError(t, err)
+	err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory})
+	require.NoError(t, err)
+
+	frames := readFramesUntilExit(t, wsConn)
+	stdoutFrames := framesByType(frames, execstream.FrameTypeStdout)
+	require.Len(t, stdoutFrames, 2)
+	require.Equal(t, "one\n", string(stdoutFrames[0].Data))
+	require.Equal(t, "two\n", string(stdoutFrames[1].Data))
+	require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code)
+}
+
 func prepareForExec(t *testing.T) (*client.Client, string) {
 	devClient, _, _ := devcontroller.StartIntegrationTestEnvironment(t)
 
@@ -164,3 +325,43 @@ func readFrame(t *testing.T, wsConn *websocket.Conn) *execstream.Frame {
 
 	return &frame
 }
+
+func readFramesUntilExit(t *testing.T, wsConn *websocket.Conn) []*execstream.Frame {
+	t.Helper()
+
+	var frames []*execstream.Frame
+
+	for {
+		frame := readFrame(t, wsConn)
+		if frame.Type == execstream.FrameTypeNoMoreHistory {
+			continue
+		}
+
+		frames = append(frames, frame)
+		if frame.Type == execstream.FrameTypeExit {
+			return frames
+		}
+	}
+}
+
+func framesByType(frames []*execstream.Frame, frameType execstream.FrameType) []*execstream.Frame {
+	var result []*execstream.Frame
+
+	for _, frame := range frames {
+		if frame.Type == frameType {
+			result = append(result, frame)
+		}
+	}
+
+	return result
+}
+
+func frameTypes(frames []*execstream.Frame) []execstream.FrameType {
+	var result []execstream.FrameType
+
+	for _, frame := range frames {
+		result = append(result, frame.Type)
+	}
+
+	return result
+}

pkg/client/vms.go

@@ -47,6 +47,13 @@ type EventsPageOptions struct {
 	Cursor string
 }
 
+type ExecSessionOptions struct {
+	Command     string
+	Stdin       bool
+	WaitSeconds uint16
+	Session     string
+}
+
 func (service *VMsService) Create(ctx context.Context, vm *v1.VM) error {
 	err := service.client.request(ctx, http.MethodPost, "vms",
 		vm, nil, nil)
@@ -164,14 +171,35 @@ func (service *VMsService) Exec(
 	stdin bool,
 	waitSeconds uint16,
 ) (*websocket.Conn, error) {
-	return service.client.wsRequestRaw(ctx, fmt.Sprintf("vms/%s/exec", url.PathEscape(name)),
+	return service.ExecSession(ctx, name, ExecSessionOptions{
-		map[string]string{
+		Command:     command,
-			"command": command,
+		Stdin:       stdin,
-			"stdin":   strconv.FormatBool(stdin),
+		WaitSeconds: waitSeconds,
-			"wait":    strconv.FormatUint(uint64(waitSeconds), 10),
 	})
 }
 
+func (service *VMsService) ExecSession(
+	ctx context.Context,
+	name string,
+	options ExecSessionOptions,
+) (*websocket.Conn, error) {
+	params := map[string]string{
+		"wait": strconv.FormatUint(uint64(options.WaitSeconds), 10),
+	}
+	if options.Command != "" {
+		params["command"] = options.Command
+	}
+	if options.Stdin {
+		params["stdin"] = strconv.FormatBool(true)
+	}
+	if options.Session != "" {
+		params["session"] = options.Session
+	}
+
+	return service.client.wsRequestRaw(ctx, fmt.Sprintf("vms/%s/exec", url.PathEscape(name)),
+		params)
+}
+
 func (service *VMsService) IP(ctx context.Context, name string, waitSeconds uint16) (string, error) {
 	result := struct {
 		IP string `json:"ip"`

pkg/client/vms_test.go

@@ -0,0 +1,40 @@
+package client
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/coder/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+func TestExecSessionBuildsReconnectableQuery(t *testing.T) {
+	var query map[string][]string
+
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		query = request.URL.Query()
+
+		conn, err := websocket.Accept(writer, request, nil)
+		require.NoError(t, err)
+		defer conn.CloseNow()
+	}))
+	defer server.Close()
+
+	devClient, err := New(WithAddress(server.URL))
+	require.NoError(t, err)
+
+	conn, err := devClient.VMs().ExecSession(t.Context(), "vm", ExecSessionOptions{
+		Command:     "echo hello",
+		Stdin:       true,
+		WaitSeconds: 7,
+		Session:     "resume-me",
+	})
+	require.NoError(t, err)
+	defer conn.CloseNow()
+
+	require.Equal(t, []string{"echo hello"}, query["command"])
+	require.Equal(t, []string{"true"}, query["stdin"])
+	require.Equal(t, []string{"7"}, query["wait"])
+	require.Equal(t, []string{"resume-me"}, query["session"])
+}