diff --git a/go.mod b/go.mod index a8225e5..c821c5f 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.25.0 require ( github.com/gin-gonic/gin v1.12.0 + github.com/gorilla/websocket v1.5.3 gopkg.in/yaml.v3 v3.0.1 ) diff --git a/go.sum b/go.sum index bb72904..f3880d6 100644 --- a/go.sum +++ b/go.sum @@ -30,6 +30,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/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= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw= diff --git a/internal/api/process_handler.go b/internal/api/process_handler.go index 82008a9..9b0e8c6 100644 --- a/internal/api/process_handler.go +++ b/internal/api/process_handler.go @@ -1,18 +1,31 @@ package api import ( + "fmt" "net/http" + "sync" + "time" "codespace/internal/model" + "codespace/internal/process" "codespace/internal/service" + "codespace/internal/util" "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" ) type processHandler struct { svc *service.ProcessService } +func exitBanner(exit process.ExitInfo) string { + if exit.Signal != "" { + return fmt.Sprintf("\r\n[process exited: signal %s]\r\n", exit.Signal) + } + return fmt.Sprintf("\r\n[process exited with code %d]\r\n", exit.Code) +} + func (h *processHandler) start(c *gin.Context) { id := c.Param("id") if err := h.svc.Start(id); err != nil { @@ -53,3 +66,92 @@ func (h *processHandler) status(c *gin.Context) { PID: status.PID, }) } + +func (h *processHandler) 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": "process not running"}) + return + } + writeError(c, err) + return + } + + sub, err := h.svc.Subscribe(id) + if err != nil { + writeError(c, err) + return + } + + exit, _ := h.svc.ExitStatus(id) + + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + ReadBufferSize: 4096, + WriteBufferSize: 4096, + } + + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + return + } + + done := make(chan struct{}) + var closeOnce sync.Once + closeAll := func() { + closeOnce.Do(func() { + sub.Close() + conn.Close() + stdin.Close() + close(done) + }) + } + + go func() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ticker.C: + if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(5*time.Second)); err != nil { + closeAll() + return + } + case chunk, ok := <-sub.Output(): + if !ok { + banner := exitBanner(exit) + conn.WriteMessage(websocket.TextMessage, []byte(banner)) + conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) + closeAll() + return + } + if err := conn.WriteMessage(websocket.TextMessage, chunk); err != nil { + closeAll() + return + } + } + } + }() + + conn.SetReadLimit(1 << 20) + conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + conn.SetPongHandler(func(string) error { + conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + return nil + }) + for { + mt, data, err := conn.ReadMessage() + if err != nil { + break + } + if mt == websocket.TextMessage { + stdin.Write(data) + } + } + closeAll() +} diff --git a/internal/api/process_ws_test.go b/internal/api/process_ws_test.go new file mode 100644 index 0000000..a84d2ab --- /dev/null +++ b/internal/api/process_ws_test.go @@ -0,0 +1,90 @@ +package api + +import ( + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "codespace/internal/process" + "codespace/internal/service" + "codespace/internal/workspace" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +func TestProcessWS(t *testing.T) { + tmpDir := t.TempDir() + + opencodePath := filepath.Join(tmpDir, "opencode") + script := []byte("#!/bin/sh\nread line\necho \"got: $line\"\nsleep 0.1\nexit 0\n") + if err := os.WriteFile(opencodePath, script, 0o755); err != nil { + t.Fatalf("write fake opencode: %v", err) + } + + wsRoot := filepath.Join(tmpDir, "workspaces") + wsMgr := workspace.NewLocalManager(wsRoot) + procMgr := process.NewManager(opencodePath) + wsSvc := service.NewWorkspaceService(wsMgr, procMgr, nil) + fileSvc := service.NewFileService(wsMgr, 1<<20) + procSvc := service.NewProcessService(wsMgr, procMgr, nil) + + r := NewRouter(wsSvc, fileSvc, procSvc, nil, gin.TestMode) + srv := httptest.NewServer(r) + defer srv.Close() + + ws, err := wsSvc.Create("wstest") + if err != nil { + t.Fatalf("create workspace: %v", err) + } + t.Cleanup(func() { + if err := wsSvc.Delete(ws.ID); err != nil { + t.Logf("cleanup delete workspace: %v", err) + } + }) + + if err := procSvc.Start(ws.ID); err != nil { + t.Fatalf("start process: %v", err) + } + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/api/workspaces/" + ws.ID + "/process/ws" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer conn.Close() + + if err := conn.WriteMessage(websocket.TextMessage, []byte("hello\n")); err != nil { + t.Fatalf("write message: %v", err) + } + + if err := conn.SetReadDeadline(time.Now().Add(3 * time.Second)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + + var gotEcho, gotExit bool + for { + _, data, err := conn.ReadMessage() + if err != nil { + break + } + text := string(data) + if strings.Contains(text, "got: hello") { + gotEcho = true + } + if strings.Contains(text, "[process exited") { + gotExit = true + break + } + } + + if !gotEcho { + t.Fatal("did not receive echo frame containing \"got: hello\"") + } + if !gotExit { + t.Fatal("did not receive exit banner frame") + } +} diff --git a/internal/api/router.go b/internal/api/router.go index e5ad6d2..6ba9183 100644 --- a/internal/api/router.go +++ b/internal/api/router.go @@ -40,6 +40,7 @@ func NewRouter(workspaces *service.WorkspaceService, files *service.FileService, api.POST("/workspaces/:id/process/stop", procHandler.stop) api.POST("/workspaces/:id/process/restart", procHandler.restart) api.GET("/workspaces/:id/process/status", procHandler.status) + api.GET("/workspaces/:id/process/ws", procHandler.ws) return r }