From 13b0c4e53f37cf318ea3559057f51d9a0ab3eee2 Mon Sep 17 00:00:00 2001 From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com> Date: Mon, 6 Jul 2026 14:57:03 +0800 Subject: [PATCH] feat(shell): multi-shell per workspace + fix WS disconnect killing bash BREAKING: every shell operation now requires a shellId. The Manager previously keyed by workspaceID alone (one bash per workspace). It now keys by (workspaceID, shellID) where shellID is a UUID returned by Start. Fixes the long-standing bug where the WS handler's closeAll called stdin.Close() and killed bash when a WS client disconnected. The Session owns the pty file; the WS handler no longer closes it. The pty is only closed by Manager.Stop (explicit) or by captureOutput when the process naturally exits (EOF). - internal/shell/manager.go: Manager interface gains List and every method takes shellID; storage becomes map[workspaceID]map[shellID]*Session; Start returns (shellID, err) via uuid.NewString; Resize/Status/ExitStatus/Subscribe/Stdin/Stop route by shellID; new List(workspaceID) returns ShellInfo[] in creation order. - internal/shell/session.go: Session gains ShellID + CreatedAt; Status type gains ShellID. - internal/shell/manager_test.go: updated existing tests for new signatures; added TestShellMultiInstance (two shells in one workspace, no output cross-talk, independent stop, List behavior). - internal/service/shell_service.go: wrappers carry shellID; new List method. - internal/service/workspace_service.go: auto-start captures/logs shellID; Delete iterates and stops all workspace shells. - internal/api/shell_handler.go: WS closeAll drops stdin.Close(); start/restart return 201 with {shellId, pid, ...}; new list handler; stop/resize take shellId in body. - internal/api/router.go: GET /api/workspaces/:id/shell (list). - internal/model/shell.go: new ShellStartResponse, ShellInfo, ShellListResponse, ShellStopRequest, ShellRestartRequest; updated ShellStatusResponse + ShellResizeRequest to carry shellId. - go.mod/go.sum: github.com/google/uuid. E2E: - workspace create -> 1 auto shell - start 2 more -> 3 shells in list - stop 1 -> 2 shells in list - WS connect -> send cmd -> disconnect -> WS reconnect -> send cmd -> response OK, no [process exited] banner --- go.mod | 1 + go.sum | 2 + internal/api/router.go | 1 + internal/api/shell_handler.go | 118 +++++++++-- internal/model/shell.go | 39 +++- internal/service/shell_service.go | 69 ++++--- internal/service/workspace_service.go | 11 +- internal/shell/manager.go | 279 ++++++++++++++++++-------- internal/shell/manager_test.go | 123 +++++++++--- internal/shell/session.go | 4 + 10 files changed, 479 insertions(+), 168 deletions(-) diff --git a/go.mod b/go.mod index 8f86493..1d23ba4 100644 --- a/go.mod +++ b/go.mod @@ -21,6 +21,7 @@ require ( github.com/go-playground/validator/v10 v10.30.3 // indirect github.com/goccy/go-json v0.10.6 // indirect github.com/goccy/go-yaml v1.19.2 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/cpuid/v2 v2.4.0 // indirect github.com/leodido/go-urn v1.4.0 // indirect diff --git a/go.sum b/go.sum index 4c21633..0bc918e 100644 --- a/go.sum +++ b/go.sum @@ -32,6 +32,8 @@ github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7Lk github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= diff --git a/internal/api/router.go b/internal/api/router.go index dca89db..0013631 100644 --- a/internal/api/router.go +++ b/internal/api/router.go @@ -47,6 +47,7 @@ func NewRouter(workspaces *service.WorkspaceService, files *service.FileService, api.POST("/workspaces/:id/shell/start", shellHandler.start) api.POST("/workspaces/:id/shell/stop", shellHandler.stop) api.POST("/workspaces/:id/shell/restart", shellHandler.restart) + api.GET("/workspaces/:id/shell", shellHandler.list) api.GET("/workspaces/:id/shell/status", shellHandler.status) api.GET("/workspaces/:id/shell/ws", shellHandler.ws) api.POST("/workspaces/:id/shell/resize", shellHandler.resize) diff --git a/internal/api/shell_handler.go b/internal/api/shell_handler.go index 3f2df1e..887c2ad 100644 --- a/internal/api/shell_handler.go +++ b/internal/api/shell_handler.go @@ -1,15 +1,14 @@ package api import ( - "fmt" "net/http" + "strconv" "sync" "time" "codespace/internal/model" "codespace/internal/service" "codespace/internal/shell" - "codespace/internal/util" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" @@ -21,23 +20,43 @@ type shellHandler struct { func shellExitBanner(exit shell.ExitInfo) string { if exit.Signal != "" { - return fmt.Sprintf("\r\n[process exited: signal %s]\r\n", exit.Signal) + return "\r\n[process exited: signal " + exit.Signal + "]\r\n" } - return fmt.Sprintf("\r\n[process exited with code %d]\r\n", exit.Code) + return "\r\n[process exited with code " + strconv.Itoa(exit.Code) + "]\r\n" } func (h *shellHandler) start(c *gin.Context) { id := c.Param("id") - if err := h.svc.Start(id); err != nil { + shellID, err := h.svc.Start(id) + if err != nil { writeError(c, err) return } - c.Status(http.StatusNoContent) + status, err := h.svc.Status(id, shellID) + if err != nil { + writeError(c, err) + return + } + c.JSON(http.StatusCreated, model.ShellStartResponse{ + WorkspaceID: status.WorkspaceID, + ShellID: status.ShellID, + Running: status.Running, + PID: status.PID, + }) } func (h *shellHandler) stop(c *gin.Context) { id := c.Param("id") - if err := h.svc.Stop(id); err != nil { + var req model.ShellStopRequest + if err := c.ShouldBindJSON(&req); err != nil { + writeBadRequest(c, "invalid json") + return + } + if req.ShellID == "" { + writeBadRequest(c, "shellId is required") + return + } + if err := h.svc.Stop(id, req.ShellID); err != nil { writeError(c, err) return } @@ -46,11 +65,31 @@ func (h *shellHandler) stop(c *gin.Context) { func (h *shellHandler) restart(c *gin.Context) { id := c.Param("id") - if err := h.svc.Restart(id); err != nil { + var req model.ShellRestartRequest + if err := c.ShouldBindJSON(&req); err != nil { + writeBadRequest(c, "invalid json") + return + } + if req.ShellID == "" { + writeBadRequest(c, "shellId is required") + return + } + newShellID, err := h.svc.Restart(id, req.ShellID) + if err != nil { writeError(c, err) return } - c.Status(http.StatusNoContent) + status, err := h.svc.Status(id, newShellID) + if err != nil { + writeError(c, err) + return + } + c.JSON(http.StatusCreated, model.ShellStartResponse{ + WorkspaceID: status.WorkspaceID, + ShellID: status.ShellID, + Running: status.Running, + PID: status.PID, + }) } func (h *shellHandler) resize(c *gin.Context) { @@ -60,11 +99,15 @@ func (h *shellHandler) resize(c *gin.Context) { writeBadRequest(c, "invalid json") return } + if req.ShellID == "" { + writeBadRequest(c, "shellId is required") + return + } if req.Cols <= 0 || req.Rows <= 0 { writeBadRequest(c, "cols and rows must be positive") return } - if err := h.svc.Resize(id, req.Cols, req.Rows); err != nil { + if err := h.svc.Resize(id, req.ShellID, req.Cols, req.Rows); err != nil { writeError(c, err) return } @@ -73,38 +116,70 @@ func (h *shellHandler) resize(c *gin.Context) { func (h *shellHandler) status(c *gin.Context) { id := c.Param("id") - status, err := h.svc.Status(id) + shellID := c.Query("shellId") + if shellID == "" { + writeBadRequest(c, "shellId query parameter is required") + return + } + status, err := h.svc.Status(id, shellID) if err != nil { writeError(c, err) return } c.JSON(http.StatusOK, model.ShellStatusResponse{ WorkspaceID: status.WorkspaceID, + ShellID: status.ShellID, Running: status.Running, PID: status.PID, }) } +func (h *shellHandler) list(c *gin.Context) { + id := c.Param("id") + shells, err := h.svc.List(id) + if err != nil { + writeError(c, err) + return + } + + infos := make([]model.ShellInfo, 0, len(shells)) + for _, sh := range shells { + infos = append(infos, model.ShellInfo{ + WorkspaceID: sh.WorkspaceID, + ShellID: sh.ShellID, + Running: sh.Running, + PID: sh.PID, + CreatedAt: sh.CreatedAt.Format(time.RFC3339Nano), + }) + } + + c.JSON(http.StatusOK, model.ShellListResponse{ + WorkspaceID: id, + Shells: infos, + }) +} + func (h *shellHandler) ws(c *gin.Context) { id := c.Param("id") - - stdin, err := h.svc.Input(id) - if err != nil { - if util.CodeOf(err) == util.CodeNotFound { - c.JSON(409, gin.H{"error": "shell not running"}) - return - } - writeError(c, err) + shellID := c.Query("shellId") + if shellID == "" { + writeBadRequest(c, "shellId query parameter is required") return } - sub, err := h.svc.Subscribe(id) + sub, err := h.svc.Subscribe(id, shellID) if err != nil { writeError(c, err) return } - exit, _ := h.svc.ExitStatus(id) + stdin, err := h.svc.Input(id, shellID) + if err != nil { + writeError(c, err) + return + } + + exit, _ := h.svc.ExitStatus(id, shellID) upgrader := websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, @@ -123,7 +198,6 @@ func (h *shellHandler) ws(c *gin.Context) { closeOnce.Do(func() { sub.Close() conn.Close() - stdin.Close() close(done) }) } diff --git a/internal/model/shell.go b/internal/model/shell.go index 06d6f9a..9cd792f 100644 --- a/internal/model/shell.go +++ b/internal/model/shell.go @@ -3,12 +3,47 @@ package model // ShellStatusResponse is the API response for shell status. type ShellStatusResponse struct { WorkspaceID string `json:"workspaceId"` + ShellID string `json:"shellId,omitempty"` Running bool `json:"running"` PID int `json:"pid,omitempty"` } +// ShellStartResponse is the API response for POST /shell/start. +type ShellStartResponse struct { + WorkspaceID string `json:"workspaceId"` + ShellID string `json:"shellId"` + Running bool `json:"running"` + PID int `json:"pid,omitempty"` +} + +// ShellInfo is a single shell entry in a list response. +type ShellInfo struct { + WorkspaceID string `json:"workspaceId"` + ShellID string `json:"shellId"` + Running bool `json:"running"` + PID int `json:"pid,omitempty"` + CreatedAt string `json:"createdAt,omitempty"` +} + +// ShellListResponse is the API response for GET /shell. +type ShellListResponse struct { + WorkspaceID string `json:"workspaceId"` + Shells []ShellInfo `json:"shells"` +} + +// ShellStopRequest is the body for POST /shell/stop. +type ShellStopRequest struct { + ShellID string `json:"shellId"` +} + // ShellResizeRequest is the body for POST /shell/resize. type ShellResizeRequest struct { - Cols int `json:"cols"` - Rows int `json:"rows"` + ShellID string `json:"shellId"` + Cols int `json:"cols"` + Rows int `json:"rows"` +} + +// ShellRestartRequest is the body for POST /shell/restart. +type ShellRestartRequest struct { + ShellID string `json:"shellId"` } diff --git a/internal/service/shell_service.go b/internal/service/shell_service.go index 9bd24c8..13b2390 100644 --- a/internal/service/shell_service.go +++ b/internal/service/shell_service.go @@ -26,86 +26,99 @@ func NewShellService(workspaces workspace.Manager, shells shell.Manager, lg *slo } // Start starts an interactive shell for the given workspace. -func (s *ShellService) Start(workspaceID string) error { +func (s *ShellService) Start(workspaceID string) (string, error) { ws, err := s.workspaces.Get(workspaceID) if err != nil { - return err + return "", err } - if err := s.shells.Start(workspaceID, ws.Root); err != nil { + shellID, err := s.shells.Start(workspaceID, ws.Root) + if err != nil { s.logger.Error("shell start failed", "workspace_id", workspaceID, "error", err) - return err + return "", err } - status := s.shells.Status(workspaceID) - s.logger.Info("shell started", "workspace_id", workspaceID, "pid", status.PID) - return nil + status := s.shells.Status(workspaceID, shellID) + s.logger.Info("shell started", "workspace_id", workspaceID, "shell_id", shellID, "pid", status.PID) + return shellID, nil } // Stop stops the interactive shell for the given workspace. -func (s *ShellService) Stop(workspaceID string) error { - if err := s.shells.Stop(workspaceID); err != nil { +func (s *ShellService) Stop(workspaceID string, shellID string) error { + if _, err := s.workspaces.Get(workspaceID); err != nil { + return err + } + if err := s.shells.Stop(workspaceID, shellID); err != nil { s.logger.Error("shell stop failed", "workspace_id", workspaceID, "error", err) return err } - s.logger.Info("shell stopped", "workspace_id", workspaceID) + s.logger.Info("shell stopped", "workspace_id", workspaceID, "shell_id", shellID) return nil } // Restart restarts the interactive shell for the given workspace. -func (s *ShellService) Restart(workspaceID string) error { +func (s *ShellService) Restart(workspaceID string, shellID string) (string, error) { ws, err := s.workspaces.Get(workspaceID) if err != nil { - return err + return "", err } - if err := s.shells.Restart(workspaceID, ws.Root); err != nil { + newShellID, err := s.shells.Restart(workspaceID, shellID, ws.Root) + if err != nil { s.logger.Error("shell restart failed", "workspace_id", workspaceID, "error", err) - return err + return "", err } - status := s.shells.Status(workspaceID) - s.logger.Info("shell restarted", "workspace_id", workspaceID, "pid", status.PID) - return nil + status := s.shells.Status(workspaceID, newShellID) + s.logger.Info("shell restarted", "workspace_id", workspaceID, "old_shell_id", shellID, "new_shell_id", newShellID, "pid", status.PID) + return newShellID, nil } // Resize resizes the PTY for the given workspace. -func (s *ShellService) Resize(workspaceID string, cols, rows int) error { +func (s *ShellService) Resize(workspaceID string, shellID string, cols, rows int) error { if _, err := s.workspaces.Get(workspaceID); err != nil { return err } - if err := s.shells.Resize(workspaceID, cols, rows); err != nil { + if err := s.shells.Resize(workspaceID, shellID, cols, rows); err != nil { s.logger.Error("shell resize failed", "workspace_id", workspaceID, "error", err) return err } - s.logger.Info("shell resized", "workspace_id", workspaceID, "cols", cols, "rows", rows) + s.logger.Info("shell resized", "workspace_id", workspaceID, "shell_id", shellID, "cols", cols, "rows", rows) return nil } // Status returns the shell status for the given workspace. -func (s *ShellService) Status(workspaceID string) (shell.Status, error) { +func (s *ShellService) Status(workspaceID string, shellID string) (shell.Status, error) { if _, err := s.workspaces.Get(workspaceID); err != nil { return shell.Status{}, err } - return s.shells.Status(workspaceID), nil + return s.shells.Status(workspaceID, shellID), nil } // Subscribe subscribes to output events for the workspace shell. -func (s *ShellService) Subscribe(workspaceID string) (shell.Subscription, error) { +func (s *ShellService) Subscribe(workspaceID string, shellID string) (shell.Subscription, error) { if _, err := s.workspaces.Get(workspaceID); err != nil { return nil, err } - return s.shells.Subscribe(workspaceID) + return s.shells.Subscribe(workspaceID, shellID) } // Input returns the stdin writer for the workspace shell. -func (s *ShellService) Input(workspaceID string) (io.WriteCloser, error) { +func (s *ShellService) Input(workspaceID string, shellID string) (io.WriteCloser, error) { if _, err := s.workspaces.Get(workspaceID); err != nil { return nil, err } - return s.shells.Stdin(workspaceID) + return s.shells.Stdin(workspaceID, shellID) } // ExitStatus returns the exit status for the workspace shell. -func (s *ShellService) ExitStatus(workspaceID string) (shell.ExitInfo, error) { +func (s *ShellService) ExitStatus(workspaceID string, shellID string) (shell.ExitInfo, error) { if _, err := s.workspaces.Get(workspaceID); err != nil { return shell.ExitInfo{}, err } - return s.shells.ExitStatus(workspaceID) + return s.shells.ExitStatus(workspaceID, shellID) +} + +// List returns all shells for the workspace. +func (s *ShellService) List(workspaceID string) ([]shell.ShellInfo, error) { + if _, err := s.workspaces.Get(workspaceID); err != nil { + return nil, err + } + return s.shells.List(workspaceID), nil } diff --git a/internal/service/workspace_service.go b/internal/service/workspace_service.go index a369ba9..be901ea 100644 --- a/internal/service/workspace_service.go +++ b/internal/service/workspace_service.go @@ -33,8 +33,10 @@ func (s *WorkspaceService) Create(id string) (*workspace.Workspace, error) { } s.logger.Info("workspace created", "workspace_id", id) - if err := s.shells.Start(ws.ID, ws.Root); err != nil { + if shellID, err := s.shells.Start(ws.ID, ws.Root); err != nil { s.logger.Warn("failed to auto-start shell for workspace", "workspace_id", id, "error", err) + } else { + s.logger.Info("auto-started default shell", "workspace_id", id, "shell_id", shellID) } return ws, nil @@ -52,8 +54,11 @@ func (s *WorkspaceService) List() ([]workspace.Workspace, error) { // Delete stops any running process for the workspace, then removes it. func (s *WorkspaceService) Delete(id string) error { - if err := s.shells.Stop(id); err != nil { - s.logger.Warn("failed to stop workspace shell before delete", "workspace_id", id, "error", err) + shells := s.shells.List(id) + for _, info := range shells { + if err := s.shells.Stop(id, info.ShellID); err != nil { + s.logger.Warn("failed to stop workspace shell before delete", "workspace_id", id, "shell_id", info.ShellID, "error", err) + } } status := s.processes.Status(id) diff --git a/internal/shell/manager.go b/internal/shell/manager.go index 192a040..4aa9f9f 100644 --- a/internal/shell/manager.go +++ b/internal/shell/manager.go @@ -7,10 +7,12 @@ import ( "os/exec" "sync" "syscall" + "time" "codespace/internal/util" "github.com/creack/pty" + "github.com/google/uuid" ) const ( @@ -20,36 +22,50 @@ const ( // DefaultShellCommand is the default command used to launch a shell. const DefaultShellCommand = "bash" +// ShellInfo describes a shell instance for listing. +type ShellInfo struct { + ShellID string + WorkspaceID string + Running bool + PID int + CreatedAt time.Time + ExitCode int + Signal string +} + // Manager manages interactive shells per workspace. type Manager interface { - Start(workspaceID string, workspaceRoot string) error - Stop(workspaceID string) error - Restart(workspaceID string, workspaceRoot string) error - Status(workspaceID string) Status - Subscribe(workspaceID string) (Subscription, error) - Stdin(workspaceID string) (io.WriteCloser, error) - ExitStatus(workspaceID string) (ExitInfo, error) - Resize(workspaceID string, cols, rows int) error + Start(workspaceID string, workspaceRoot string) (string, error) + Stop(workspaceID string, shellID string) error + Restart(workspaceID string, shellID string, workspaceRoot string) (string, error) + Status(workspaceID string, shellID string) Status + Subscribe(workspaceID string, shellID string) (Subscription, error) + Stdin(workspaceID string, shellID string) (io.WriteCloser, error) + ExitStatus(workspaceID string, shellID string) (ExitInfo, error) + Resize(workspaceID string, shellID string, cols, rows int) error + List(workspaceID string) []ShellInfo } // LocalManager implements Manager using local OS processes. type LocalManager struct { - command string - args []string - mu sync.Mutex - sessions map[string]*Session - exitedMu sync.Mutex - exited map[string]struct{} + command string + args []string + mu sync.Mutex + sessions map[string]map[string]*Session // workspaceID -> shellID -> Session + exitedMu sync.Mutex + exited map[string]map[string]struct{} // workspaceID -> shellID -> exited + shellOrder map[string][]string // workspaceID -> shellIDs in creation order } // NewManager creates a LocalManager with the given command and arguments. // If command is empty, it defaults to "bash". func NewManager(command string, args []string) *LocalManager { return &LocalManager{ - command: normalizeCommand(command), - args: args, - sessions: make(map[string]*Session), - exited: make(map[string]struct{}), + command: normalizeCommand(command), + args: args, + sessions: make(map[string]map[string]*Session), + exited: make(map[string]map[string]struct{}), + shellOrder: make(map[string][]string), } } @@ -60,20 +76,12 @@ func normalizeCommand(command string) string { return command } -// Start launches a shell for the given workspace. -// Returns CodeConflict if a session is already running. -func (m *LocalManager) Start(workspaceID string, workspaceRoot string) error { +// Start launches a new shell for the given workspace and returns its shellID. +func (m *LocalManager) Start(workspaceID string, workspaceRoot string) (string, error) { m.mu.Lock() defer m.mu.Unlock() - if _, ok := m.sessions[workspaceID]; ok { - m.exitedMu.Lock() - _, exited := m.exited[workspaceID] - m.exitedMu.Unlock() - if !exited { - return util.New(util.CodeConflict, "shell already running") - } - } + shellID := uuid.NewString() cmd := exec.Command(m.command, m.args...) cmd.Dir = workspaceRoot @@ -84,27 +92,36 @@ func (m *LocalManager) Start(workspaceID string, workspaceRoot string) error { ptyF, err := pty.StartWithSize(cmd, &pty.Winsize{Rows: 24, Cols: 80, X: 0, Y: 0}) if err != nil { - return util.Wrap(util.CodeInternal, "failed to start shell on pty", err) + return "", util.Wrap(util.CodeInternal, "failed to start shell on pty", err) } sess := &Session{ + ShellID: shellID, WorkspaceID: workspaceID, Root: workspaceRoot, Cmd: cmd, Stdin: ptyF, Subscribers: make(map[Subscription]struct{}), + CreatedAt: time.Now().UTC(), } - m.sessions[workspaceID] = sess - m.exitedMu.Lock() - delete(m.exited, workspaceID) - m.exitedMu.Unlock() + if m.sessions[workspaceID] == nil { + m.sessions[workspaceID] = make(map[string]*Session) + } + m.sessions[workspaceID][shellID] = sess + + if m.exited[workspaceID] == nil { + m.exited[workspaceID] = make(map[string]struct{}) + } + delete(m.exited[workspaceID], shellID) + + m.shellOrder[workspaceID] = append(m.shellOrder[workspaceID], shellID) outputDone := make(chan struct{}) go m.captureOutput(sess, ptyF, outputDone) go m.waitExit(sess, outputDone) - return nil + return shellID, nil } // captureOutput reads from the PTY master and fans out each chunk to all @@ -144,7 +161,10 @@ func (m *LocalManager) waitExit(sess *Session, outputDone <-chan struct{}) { <-outputDone m.exitedMu.Lock() - m.exited[sess.WorkspaceID] = struct{}{} + if m.exited[sess.WorkspaceID] == nil { + m.exited[sess.WorkspaceID] = make(map[string]struct{}) + } + m.exited[sess.WorkspaceID][sess.ShellID] = struct{}{} m.exitedMu.Unlock() sess.mu.Lock() @@ -162,56 +182,85 @@ func (m *LocalManager) waitExit(sess *Session, outputDone <-chan struct{}) { } } -// Stop kills the shell for the given workspace. +// Stop kills the shell for the given workspace and shellID. // Returns CodeNotFound if no session exists. -func (m *LocalManager) Stop(workspaceID string) error { +func (m *LocalManager) Stop(workspaceID string, shellID string) error { m.mu.Lock() defer m.mu.Unlock() - sess, ok := m.sessions[workspaceID] + ws, ok := m.sessions[workspaceID] if !ok { - return util.New(util.CodeNotFound, "no running shell for workspace") + return util.New(util.CodeNotFound, "shell not found") + } + sess, ok := ws[shellID] + if !ok { + return util.New(util.CodeNotFound, "shell not found") } if sess.Cmd.Process != nil { _ = sess.Cmd.Process.Kill() } - delete(m.sessions, workspaceID) + delete(ws, shellID) + if len(ws) == 0 { + delete(m.sessions, workspaceID) + } + m.exitedMu.Lock() - delete(m.exited, workspaceID) + if m.exited[workspaceID] != nil { + delete(m.exited[workspaceID], shellID) + if len(m.exited[workspaceID]) == 0 { + delete(m.exited, workspaceID) + } + } m.exitedMu.Unlock() return nil } -// Restart stops (if running) then starts the shell. -func (m *LocalManager) Restart(workspaceID string, workspaceRoot string) error { +// Restart stops the identified shell (if running) then starts a new one. +func (m *LocalManager) Restart(workspaceID string, shellID string, workspaceRoot string) (string, error) { m.mu.Lock() - if sess, ok := m.sessions[workspaceID]; ok { - m.exitedMu.Lock() - _, exited := m.exited[workspaceID] - m.exitedMu.Unlock() - if !exited && sess.Cmd.Process != nil { - _ = sess.Cmd.Process.Kill() + ws, ok := m.sessions[workspaceID] + if ok { + if sess, ok := ws[shellID]; ok { + m.exitedMu.Lock() + _, exited := m.exited[workspaceID][shellID] + m.exitedMu.Unlock() + if !exited && sess.Cmd.Process != nil { + _ = sess.Cmd.Process.Kill() + } + delete(ws, shellID) + if len(ws) == 0 { + delete(m.sessions, workspaceID) + } + m.exitedMu.Lock() + if m.exited[workspaceID] != nil { + delete(m.exited[workspaceID], shellID) + if len(m.exited[workspaceID]) == 0 { + delete(m.exited, workspaceID) + } + } + m.exitedMu.Unlock() } - delete(m.sessions, workspaceID) - m.exitedMu.Lock() - delete(m.exited, workspaceID) - m.exitedMu.Unlock() } m.mu.Unlock() return m.Start(workspaceID, workspaceRoot) } -// Resize resizes the PTY for the given workspace. +// Resize resizes the PTY for the given shell. // Returns CodeNotFound if no session exists, or CodeBadRequest for invalid dimensions. -func (m *LocalManager) Resize(workspaceID string, cols, rows int) error { +func (m *LocalManager) Resize(workspaceID string, shellID string, cols, rows int) error { m.mu.Lock() - sess, ok := m.sessions[workspaceID] - m.mu.Unlock() + defer m.mu.Unlock() + + ws, ok := m.sessions[workspaceID] if !ok { - return util.New(util.CodeNotFound, "no running shell for workspace") + return util.New(util.CodeNotFound, "shell not found") + } + sess, ok := ws[shellID] + if !ok { + return util.New(util.CodeNotFound, "shell not found") } if cols <= 0 || rows <= 0 || cols > 10000 || rows > 10000 { return util.New(util.CodeBadRequest, "invalid cols/rows") @@ -224,38 +273,47 @@ func (m *LocalManager) Resize(workspaceID string, cols, rows int) error { } // Status returns the current shell status for the workspace. -func (m *LocalManager) Status(workspaceID string) Status { +func (m *LocalManager) Status(workspaceID string, shellID string) Status { m.mu.Lock() - sess, ok := m.sessions[workspaceID] - m.mu.Unlock() + defer m.mu.Unlock() + ws, ok := m.sessions[workspaceID] if !ok { - return Status{WorkspaceID: workspaceID, Running: false} + return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false} + } + sess, ok := ws[shellID] + if !ok { + return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false} } m.exitedMu.Lock() - _, exited := m.exited[workspaceID] + _, exited := m.exited[workspaceID][shellID] m.exitedMu.Unlock() if exited { - return Status{WorkspaceID: workspaceID, Running: false} + return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false} } return Status{ WorkspaceID: workspaceID, + ShellID: shellID, Running: true, PID: sess.Cmd.Process.Pid, } } -// Subscribe creates a new output subscription for the workspace. -// Returns CodeNotFound if the workspace has no session. -func (m *LocalManager) Subscribe(workspaceID string) (Subscription, error) { +// Subscribe creates a new output subscription for the identified shell. +// Returns CodeNotFound if the shell has no session. +func (m *LocalManager) Subscribe(workspaceID string, shellID string) (Subscription, error) { m.mu.Lock() - sess, ok := m.sessions[workspaceID] - m.mu.Unlock() + defer m.mu.Unlock() + ws, ok := m.sessions[workspaceID] if !ok { - return nil, util.New(util.CodeNotFound, "no running shell") + return nil, util.New(util.CodeNotFound, "shell not found") + } + sess, ok := ws[shellID] + if !ok { + return nil, util.New(util.CodeNotFound, "shell not found") } sess.mu.Lock() @@ -264,7 +322,7 @@ func (m *LocalManager) Subscribe(workspaceID string) (Subscription, error) { sess.mu.Unlock() m.exitedMu.Lock() - _, exited := m.exited[workspaceID] + _, exited := m.exited[workspaceID][shellID] m.exitedMu.Unlock() if exited { sub.closeChan() @@ -273,34 +331,85 @@ func (m *LocalManager) Subscribe(workspaceID string) (Subscription, error) { return sub, nil } -// Stdin returns the stdin writer for the workspace shell. -func (m *LocalManager) Stdin(workspaceID string) (io.WriteCloser, error) { - return m.stdinOf(workspaceID) +// Stdin returns the stdin writer for the identified shell. +func (m *LocalManager) Stdin(workspaceID string, shellID string) (io.WriteCloser, error) { + return m.stdinOf(workspaceID, shellID) } -// ExitStatus returns the exit status for the workspace shell. -func (m *LocalManager) ExitStatus(workspaceID string) (ExitInfo, error) { - return m.exitOf(workspaceID) +// ExitStatus returns the exit status for the identified shell. +func (m *LocalManager) ExitStatus(workspaceID string, shellID string) (ExitInfo, error) { + return m.exitOf(workspaceID, shellID) +} + +// List returns all shells for the workspace in creation order. +func (m *LocalManager) List(workspaceID string) []ShellInfo { + m.mu.Lock() + defer m.mu.Unlock() + + var infos []ShellInfo + for _, shellID := range m.shellOrder[workspaceID] { + ws, ok := m.sessions[workspaceID] + if !ok { + continue + } + sess, ok := ws[shellID] + if !ok { + continue + } + + info := ShellInfo{ + ShellID: shellID, + WorkspaceID: workspaceID, + CreatedAt: sess.CreatedAt, + } + + m.exitedMu.Lock() + _, exited := m.exited[workspaceID][shellID] + m.exitedMu.Unlock() + + if !exited && sess.Cmd.Process != nil { + info.Running = true + info.PID = sess.Cmd.Process.Pid + } + + sess.mu.Lock() + info.ExitCode = sess.Exit.Code + info.Signal = sess.Exit.Signal + sess.mu.Unlock() + + infos = append(infos, info) + } + return infos } // stdinOf returns the session's stdin writer or CodeNotFound. -func (m *LocalManager) stdinOf(workspaceID string) (io.WriteCloser, error) { +func (m *LocalManager) stdinOf(workspaceID string, shellID string) (io.WriteCloser, error) { m.mu.Lock() defer m.mu.Unlock() - sess, ok := m.sessions[workspaceID] + + ws, ok := m.sessions[workspaceID] if !ok { - return nil, util.New(util.CodeNotFound, "no running shell") + return nil, util.New(util.CodeNotFound, "shell not found") + } + sess, ok := ws[shellID] + if !ok { + return nil, util.New(util.CodeNotFound, "shell not found") } return sess.Stdin, nil } // exitOf returns the session's recorded exit info or CodeNotFound. -func (m *LocalManager) exitOf(workspaceID string) (ExitInfo, error) { +func (m *LocalManager) exitOf(workspaceID string, shellID string) (ExitInfo, error) { m.mu.Lock() - sess, ok := m.sessions[workspaceID] - m.mu.Unlock() + defer m.mu.Unlock() + + ws, ok := m.sessions[workspaceID] if !ok { - return ExitInfo{}, util.New(util.CodeNotFound, "no running shell") + return ExitInfo{}, util.New(util.CodeNotFound, "shell not found") + } + sess, ok := ws[shellID] + if !ok { + return ExitInfo{}, util.New(util.CodeNotFound, "shell not found") } sess.mu.Lock() diff --git a/internal/shell/manager_test.go b/internal/shell/manager_test.go index 45874b2..b9640dc 100644 --- a/internal/shell/manager_test.go +++ b/internal/shell/manager_test.go @@ -10,37 +10,34 @@ import ( ) import ( - "os" - "codespace/internal/util" - - "github.com/creack/pty" ) -func startTestShell(t *testing.T) (*LocalManager, io.WriteCloser, Subscription) { +func startTestShell(t *testing.T) (*LocalManager, string, io.WriteCloser, Subscription) { t.Helper() mgr := NewManager("bash", []string{"-i"}) root := t.TempDir() - if err := mgr.Start("test-ws", root); err != nil { + shellID, err := mgr.Start("test-ws", root) + if err != nil { t.Fatalf("Start failed: %v", err) } - stdin, err := mgr.Stdin("test-ws") + stdin, err := mgr.Stdin("test-ws", shellID) if err != nil { t.Fatalf("Stdin failed: %v", err) } - sub, err := mgr.Subscribe("test-ws") + sub, err := mgr.Subscribe("test-ws", shellID) if err != nil { t.Fatalf("Subscribe failed: %v", err) } t.Cleanup(func() { sub.Close() - _ = mgr.Stop("test-ws") + _ = mgr.Stop("test-ws", shellID) }) - return mgr, stdin, sub + return mgr, shellID, stdin, sub } func waitForMarker(t *testing.T, sub Subscription, stdin io.Writer, input, marker string) string { @@ -71,8 +68,9 @@ func waitForMarker(t *testing.T, sub Subscription, stdin io.Writer, input, marke } func TestShellPTYIsRealTerminal(t *testing.T) { - mgr, stdin, sub := startTestShell(t) + mgr, shellID, stdin, sub := startTestShell(t) _ = mgr + _ = shellID out := waitForMarker(t, sub, stdin, "echo HELLO_PTY_TEST\n", "HELLO_PTY_TEST") @@ -85,7 +83,7 @@ func TestShellPTYIsRealTerminal(t *testing.T) { } func TestShellPTYTermEnv(t *testing.T) { - _, stdin, sub := startTestShell(t) + _, _, stdin, sub := startTestShell(t) out := waitForMarker(t, sub, stdin, "echo \"$TERM\"\n", "xterm-256color") if !strings.Contains(out, "xterm-256color") { @@ -94,31 +92,23 @@ func TestShellPTYTermEnv(t *testing.T) { } func TestShellResizeUpdatesPTYSize(t *testing.T) { - mgr, stdin, sub := startTestShell(t) + mgr, shellID, stdin, sub := startTestShell(t) _ = sub - f, ok := stdin.(*os.File) - if !ok { - t.Fatal("stdin is not a pty file") - } - - if err := mgr.Resize("test-ws", 120, 40); err != nil { + if err := mgr.Resize("test-ws", shellID, 120, 40); err != nil { t.Fatalf("Resize failed: %v", err) } - rows, cols, err := pty.Getsize(f) - if err != nil { - t.Fatalf("Getsize failed: %v", err) - } - if cols != 120 || rows != 40 { - t.Errorf("expected cols=120 rows=40, got cols=%d rows=%d", cols, rows) + out := waitForMarker(t, sub, stdin, "stty size\n", "40 120") + if !strings.Contains(out, "40 120") { + t.Errorf("expected stty size to report 40 120, got:\n%s", out) } } func TestShellResizeRejectsInvalidSize(t *testing.T) { - mgr, _, _ := startTestShell(t) + mgr, shellID, _, _ := startTestShell(t) - if err := mgr.Resize("test-ws", 0, 40); err == nil { + if err := mgr.Resize("test-ws", shellID, 0, 40); err == nil { t.Fatal("expected error for invalid size") } else if util.CodeOf(err) != util.CodeBadRequest { t.Errorf("expected CodeBadRequest, got %v", util.CodeOf(err)) @@ -128,9 +118,86 @@ func TestShellResizeRejectsInvalidSize(t *testing.T) { func TestShellResizeRejectsMissingSession(t *testing.T) { mgr := NewManager("bash", []string{"-i"}) - if err := mgr.Resize("missing-ws", 80, 24); err == nil { + if err := mgr.Resize("missing-ws", "missing-shell", 80, 24); err == nil { t.Fatal("expected error for missing session") } else if util.CodeOf(err) != util.CodeNotFound { t.Errorf("expected CodeNotFound, got %v", util.CodeOf(err)) } } + +func TestShellMultiInstance(t *testing.T) { + mgr := NewManager("bash", []string{"-i"}) + root := t.TempDir() + + shellID1, err := mgr.Start("test-ws", root) + if err != nil { + t.Fatalf("Start shell1 failed: %v", err) + } + shellID2, err := mgr.Start("test-ws", root) + if err != nil { + t.Fatalf("Start shell2 failed: %v", err) + } + + if shellID1 == shellID2 { + t.Fatalf("expected distinct shell IDs, got %q and %q", shellID1, shellID2) + } + + stdin1, err := mgr.Stdin("test-ws", shellID1) + if err != nil { + t.Fatalf("Stdin shell1 failed: %v", err) + } + stdin2, err := mgr.Stdin("test-ws", shellID2) + if err != nil { + t.Fatalf("Stdin shell2 failed: %v", err) + } + + sub1, err := mgr.Subscribe("test-ws", shellID1) + if err != nil { + t.Fatalf("Subscribe shell1 failed: %v", err) + } + sub2, err := mgr.Subscribe("test-ws", shellID2) + if err != nil { + t.Fatalf("Subscribe shell2 failed: %v", err) + } + + t.Cleanup(func() { + sub1.Close() + sub2.Close() + _ = mgr.Stop("test-ws", shellID1) + _ = mgr.Stop("test-ws", shellID2) + }) + + list := mgr.List("test-ws") + if len(list) != 2 { + t.Fatalf("List returned %d shells, want 2", len(list)) + } + + out1 := waitForMarker(t, sub1, stdin1, "echo SHELL_ONE\n", "SHELL_ONE") + if strings.Contains(out1, "SHELL_TWO") { + t.Errorf("shell1 output leaked shell2 output: %s", out1) + } + + out2 := waitForMarker(t, sub2, stdin2, "echo SHELL_TWO\n", "SHELL_TWO") + if strings.Contains(out2, "SHELL_ONE") { + t.Errorf("shell2 output leaked shell1 output: %s", out2) + } + + if err := mgr.Stop("test-ws", shellID1); err != nil { + t.Fatalf("Stop shell1 failed: %v", err) + } + + if _, err := mgr.Stdin("test-ws", shellID1); err == nil { + t.Fatal("expected shell1 stdin to be unavailable after stop") + } + + // shell2 should still work. + _ = waitForMarker(t, sub2, stdin2, "echo STILL_ALIVE\n", "STILL_ALIVE") + + list = mgr.List("test-ws") + if len(list) != 1 { + t.Fatalf("List returned %d shells after stop, want 1", len(list)) + } + if list[0].ShellID != shellID2 { + t.Fatalf("List returned shell %q, want %q", list[0].ShellID, shellID2) + } +} diff --git a/internal/shell/session.go b/internal/shell/session.go index fc8844d..78cc536 100644 --- a/internal/shell/session.go +++ b/internal/shell/session.go @@ -4,11 +4,13 @@ import ( "io" "os/exec" "sync" + "time" ) // Status represents the status of a shell for a workspace. type Status struct { WorkspaceID string `json:"workspaceId"` + ShellID string `json:"shellId,omitempty"` Running bool `json:"running"` PID int `json:"pid,omitempty"` } @@ -21,12 +23,14 @@ type ExitInfo struct { // Session holds the runtime state of a started shell. type Session struct { + ShellID string WorkspaceID string Root string Cmd *exec.Cmd Stdin io.WriteCloser Subscribers map[Subscription]struct{} Exit ExitInfo + CreatedAt time.Time mu sync.Mutex }