sundynix-agentix/sundynix-gateway/internal/handler/local_runner.go

package handler

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"

	"github.com/sundynix/sundynix-shared/contract"
)

// 本地执行 runner 接入（JARVIS P4「本地的手」，LOCAL_AGENT_DESIGN 档 A）：
// 桌面端 Go host 登录后连本 WS，把自己注册成该用户的本地执行器。local_* 工具调用
// 经 NATS sundynix.local.exec.<uid>（request-reply，本连接队列组订阅）路由到这里，
// 转发给 runner 执行、结果原路回。runner 掉线即退订——调用方拿到"不在线"而非干等。
//
// 安全：执行发生在**用户自己的机器**、以用户 OS 身份、锁在用户自选的工作目录（沙箱在
// runner 侧强制）；服务端只做路由。归属天然成立：subject 按 uid 分桶，工具 inject user_id。

var runnerUpgrader = websocket.Upgrader{
	ReadBufferSize: 4096, WriteBufferSize: 4096,
	CheckOrigin: func(*http.Request) bool { return true }, // 鉴权由 token 把关
}

const (
	// 等 runner 回执行结果。写文件/执行命令要等用户在桌面端点原生确认框（最长 60s）
	// 再真跑（最长 60s），所以给到 140s；只读操作瞬时返回，不受影响。
	runnerCallTimeout = 140 * time.Second
	runnerWriteWait   = 10 * time.Second
)

// runnerReq / runnerResp 是 gateway ↔ 桌面 runner 的 WS 消息（JSON 文本帧）。
type runnerReq struct {
	ID   string         `json:"id"`
	Tool string         `json:"tool"`
	Args map[string]any `json:"args,omitempty"`
}

type runnerResp struct {
	ID      string `json:"id"`
	OK      bool   `json:"ok"`
	Content string `json:"content,omitempty"`
	Error   string `json:"error,omitempty"`
	Workdir string `json:"workdir,omitempty"` // hello 帧带：runner 声明的沙箱根（仅展示/日志）
}

// LocalRunner: GET /api/v1/local/runner —— 桌面端注册本地执行器（WS，?token= 鉴权）。
func (h *Handler) LocalRunner(c *gin.Context) {
	uid := userID(c)
	if uid == "" {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "需要登录"})
		return
	}
	conn, err := runnerUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	r := &runnerConn{conn: conn, pending: map[string]chan *runnerResp{}}

	// 队列组订阅本用户执行主题：多台桌面在线时 NATS 队列组自动挑一台执行。
	drain, err := h.bus.ServeTool(contract.LocalExecSubject(uid), "local-"+uid, r.forward)
	if err != nil {
		log.Printf("[local] runner 订阅失败 uid=%s: %v", uid, err)
		return
	}
	log.Printf("[local] 本地执行器上线 uid=%s", uid)

	// 读循环：分发 runner 的执行回执；连接断开即退出。
	for {
		_, data, rerr := conn.ReadMessage()
		if rerr != nil {
			break
		}
		var resp runnerResp
		if json.Unmarshal(data, &resp) != nil {
			continue
		}
		if resp.ID == "hello" {
			log.Printf("[local] runner 声明工作目录 uid=%s dir=%s", uid, resp.Workdir)
			continue
		}
		r.deliver(&resp)
	}

	dctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	drain(dctx)
	cancel()
	r.failAll("本地执行器已断开")
	log.Printf("[local] 本地执行器下线 uid=%s", uid)
}

// runnerConn 一条 runner 连接：串行化写 + 按请求 id 关联回执。
type runnerConn struct {
	conn    *websocket.Conn
	writeMu sync.Mutex
	mu      sync.Mutex
	pending map[string]chan *runnerResp
	seq     atomic.Int64
}

// forward 把一次 local_* 工具调用转发给 runner，等回执（NATS ServeTool 的 handler）。
func (r *runnerConn) forward(ctx context.Context, call *contract.ToolCall) *contract.ToolResult {
	id := fmt.Sprintf("r%d", r.seq.Add(1))
	ch := make(chan *runnerResp, 1)
	r.mu.Lock()
	r.pending[id] = ch
	r.mu.Unlock()
	defer func() {
		r.mu.Lock()
		delete(r.pending, id)
		r.mu.Unlock()
	}()

	data, _ := json.Marshal(runnerReq{ID: id, Tool: call.Tool, Args: call.Args})
	r.writeMu.Lock()
	_ = r.conn.SetWriteDeadline(time.Now().Add(runnerWriteWait))
	werr := r.conn.WriteMessage(websocket.TextMessage, data)
	r.writeMu.Unlock()
	if werr != nil {
		return &contract.ToolResult{OK: false, Error: "转发本地执行失败: " + werr.Error()}
	}

	select {
	case resp := <-ch:
		return &contract.ToolResult{OK: resp.OK, Content: resp.Content, Error: resp.Error}
	case <-time.After(runnerCallTimeout):
		return &contract.ToolResult{OK: false, Error: "本地执行超时（桌面端未响应）"}
	case <-ctx.Done():
		return &contract.ToolResult{OK: false, Error: "本地执行取消: " + ctx.Err().Error()}
	}
}

// deliver 把 runner 回执投递给等待中的 forward。
func (r *runnerConn) deliver(resp *runnerResp) {
	r.mu.Lock()
	ch := r.pending[resp.ID]
	r.mu.Unlock()
	if ch != nil {
		select {
		case ch <- resp:
		default:
		}
	}
}

// failAll 连接断开时让所有在途调用立即失败（别干等超时）。
func (r *runnerConn) failAll(msg string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	for id, ch := range r.pending {
		select {
		case ch <- &runnerResp{ID: id, OK: false, Error: msg}:
		default:
		}
	}
}