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
This commit is contained in:
tao.chen
2026-07-06 14:57:03 +08:00
parent f5a6ff8b0d
commit 13b0c4e53f
10 changed files with 479 additions and 168 deletions
+1
View File
@@ -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
+2
View File
@@ -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=
+1
View File
@@ -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)
+96 -22
View File
@@ -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)
})
}
+37 -2
View File
@@ -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"`
}
+41 -28
View File
@@ -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
}
+8 -3
View File
@@ -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)
+194 -85
View File
@@ -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()
+95 -28
View File
@@ -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)
}
}
+4
View File
@@ -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
}