fix: 修复 PTY Websocket 客户端断开连接后命令仍在后台运行的问题 (#1222)

* Initial plan

* fix: 修复 PTY Websocket 客户端断开连接后命令仍在后台运行的问题

Co-authored-by: devhaozi <115467771+devhaozi@users.noreply.github.com>

* fix: 解决连接中断后命令不被杀死的问题

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: devhaozi <115467771+devhaozi@users.noreply.github.com>
Co-authored-by: 耗子 <haozi@loli.email>

pkg/docker/turn.go

@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"syscall"
+	"time"
 
 	"github.com/coder/websocket"
 	"github.com/moby/moby/client"
@@ -19,12 +21,11 @@ type MessageResize struct {
 
 // Turn 容器终端转发器
 type Turn struct {
-	ctx       context.Context
-	ws        *websocket.Conn
-	client    *client.Client
-	execID    string
-	hijack    client.ExecAttachResult
-	closeOnce bool
+	ctx    context.Context
+	ws     *websocket.Conn
+	client *client.Client
+	execID string
+	hijack client.ExecAttachResult
 }
 
 // NewTurn 创建容器终端转发器
@@ -77,11 +78,26 @@ func (t *Turn) Write(p []byte) (n int, err error) {
 
 // Close 关闭连接
 func (t *Turn) Close() {
-	if t.closeOnce {
-		return
+	// 检查进程是否仍在运行
+	inspectResp, err := t.client.ExecInspect(t.ctx, t.execID, client.ExecInspectOptions{})
+	if err == nil && inspectResp.Running {
+		_ = syscall.Kill(inspectResp.PID, syscall.SIGTERM)
+		// 等待最多 10 秒
+		for i := 0; i < 10; i++ {
+			time.Sleep(1 * time.Second)
+			inspectResp, err = t.client.ExecInspect(t.ctx, t.execID, client.ExecInspectOptions{})
+			if err != nil || !inspectResp.Running {
+				break
+			}
+		}
+		// 如果仍在运行，KILL
+		if err == nil && inspectResp.Running {
+			_ = syscall.Kill(inspectResp.PID, syscall.SIGKILL)
+		}
 	}
-	t.closeOnce = true
+
 	t.hijack.Close()
+	_ = t.hijack.CloseWrite()
 	_ = t.client.Close()
 }
 

pkg/shell/pty.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"os/exec"
 	"syscall"
+	"time"
 
 	"github.com/coder/websocket"
 	"github.com/creack/pty"
@@ -65,8 +66,25 @@ func (t *Turn) Wait() {
 	_ = t.cmd.Wait()
 }
 
-// Close 关闭 PTY
+// Close 关闭 PTY 并终止子进程
 func (t *Turn) Close() {
+	_ = t.cmd.Process.Signal(syscall.SIGTERM)
+
+	// 等待最多 10 秒
+	done := make(chan struct{})
+	go func() {
+		_ = t.cmd.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// 进程已退出
+	case <-time.After(10 * time.Second):
+		// 超时，KILL
+		_ = t.cmd.Process.Kill()
+	}
+
 	_ = t.ptmx.Close()
 }
 

pkg/ssh/turn.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"time"
 
 	"github.com/coder/websocket"
 	"golang.org/x/crypto/ssh"
@@ -63,6 +64,22 @@ func (t *Turn) Write(p []byte) (n int, err error) {
 func (t *Turn) Close() {
 	_ = t.stdin.Close()
 	_ = t.session.Signal(ssh.SIGTERM)
+
+	// 等待最多 10 秒
+	done := make(chan struct{})
+	go func() {
+		_ = t.session.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// 会话已退出
+	case <-time.After(10 * time.Second):
+		// 超时，KILL
+		_ = t.session.Signal(ssh.SIGKILL)
+	}
+
 	_ = t.session.Close()
 }