sproutclaw-web/backend/internal/services/web_terminal.go

package services

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"sync"

	"github.com/Kodecable/crosspty"
	"github.com/gorilla/websocket"
)

type terminalWSMessage struct {
	Type string `json:"type"`
	Data string `json:"data,omitempty"`
	Cols uint16 `json:"cols,omitempty"`
	Rows uint16 `json:"rows,omitempty"`
}

// ServeWebTerminal attaches a PTY shell session to the websocket connection.
func ServeWebTerminal(conn *websocket.Conn, workDir, agentDir string) error {
	abs, err := filepath.Abs(workDir)
	if err != nil {
		return fmt.Errorf("解析目录失败: %w", err)
	}
	info, err := os.Stat(abs)
	if err != nil {
		return fmt.Errorf("目录不存在: %s", abs)
	}
	if !info.IsDir() {
		return fmt.Errorf("不是有效目录: %s", abs)
	}

	envInject := map[string]string{
		"TERM":      "xterm-256color",
		"COLORTERM": "truecolor",
	}
	if agentDir != "" {
		envInject["SPROUTCLAW_CODING_AGENT_DIR"] = agentDir
		envInject["PI_CODING_AGENT_DIR"] = agentDir
	}

	ptmx, err := crosspty.Start(crosspty.CommandConfig{
		Argv:      defaultShellArgv(),
		Dir:       abs,
		Env:       os.Environ(),
		EnvInject: envInject,
		Size:      crosspty.TermSize{Rows: 24, Cols: 80},
	})
	if err != nil {
		if err == crosspty.ErrConPTYNotSupported {
			return fmt.Errorf("当前 Windows 版本不支持 ConPTY，需 Windows 10 1809 或更高版本")
		}
		return fmt.Errorf("启动终端失败: %w", err)
	}
	defer ptmx.Close()

	var writeMu sync.Mutex
	writeJSON := func(v any) error {
		writeMu.Lock()
		defer writeMu.Unlock()
		return conn.WriteJSON(v)
	}

	go func() {
		buf := make([]byte, 4096)
		for {
			n, readErr := ptmx.Read(buf)
			if n > 0 {
				writeMu.Lock()
				werr := conn.WriteMessage(websocket.BinaryMessage, buf[:n])
				writeMu.Unlock()
				if werr != nil {
					return
				}
			}
			if readErr != nil {
				return
			}
		}
	}()

	for {
		msgType, payload, readErr := conn.ReadMessage()
		if readErr != nil {
			if websocket.IsCloseError(readErr, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				return nil
			}
			if readErr == io.EOF {
				return nil
			}
			return readErr
		}

		switch msgType {
		case websocket.BinaryMessage:
			if _, werr := ptmx.Write(payload); werr != nil {
				return werr
			}
		case websocket.TextMessage:
			var msg terminalWSMessage
			if err := json.Unmarshal(payload, &msg); err != nil {
				if _, werr := ptmx.Write(payload); werr != nil {
					return werr
				}
				continue
			}
			switch msg.Type {
			case "input":
				if _, werr := ptmx.Write([]byte(msg.Data)); werr != nil {
					return werr
				}
			case "resize":
				if msg.Cols > 0 && msg.Rows > 0 {
					_ = ptmx.Resize(crosspty.TermSize{Rows: msg.Rows, Cols: msg.Cols})
				}
			case "ping":
				if err := writeJSON(map[string]string{"type": "pong"}); err != nil {
					return err
				}
			default:
				if msg.Data != "" {
					if _, werr := ptmx.Write([]byte(msg.Data)); werr != nil {
						return werr
					}
				}
			}
		}
	}
}

func defaultShellArgv() []string {
	if runtime.GOOS == "windows" {
		if comspec := os.Getenv("COMSPEC"); comspec != "" {
			return []string{comspec}
		}
		return []string{"cmd.exe"}
	}
	if shell := os.Getenv("SHELL"); shell != "" {
		return []string{shell, "-l"}
	}
	return []string{"/bin/bash", "-l"}
}