feat(openapi): integrate Tai forward handlers for proxy and VNC operations

- Replaced direct tunnel handling with a dedicated Tai integration for managing proxy and VNC requests.
- Updated routing to streamline the attachment of Tai handlers, enhancing modularity and maintainability.

Made-with: Cursor

openapi/openapi.go

@@ -27,12 +27,12 @@ import (
 	"github.com/yaoapp/yao/openapi/otp"
 	"github.com/yaoapp/yao/openapi/response"
 	"github.com/yaoapp/yao/openapi/sandbox"
+	openapiTai "github.com/yaoapp/yao/openapi/tai"
 	"github.com/yaoapp/yao/openapi/team"
 	openapiTrace "github.com/yaoapp/yao/openapi/trace"
 	"github.com/yaoapp/yao/openapi/user"
 	openapiWorkspace "github.com/yaoapp/yao/openapi/workspace"
 	taiapi "github.com/yaoapp/yao/tai/api"
-	taitunnel "github.com/yaoapp/yao/tai/tunnel"
 )
 
 // Server is the OpenAPI server
@@ -191,9 +191,8 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) {
 	// Tai nodes handlers
 	nodes.Attach(group.Group("/nodes"), openapi.OAuth)
 
-	// Tai tunnel: gRPC Forward-based HTTP/VNC transparent proxy
-	group.Any("/tai/:taiID/proxy/*path", taitunnel.HandleForwardLazy)
-	group.Any("/tai/:taiID/vnc/*path", taitunnel.HandleForwardLazy)
+	// Tai forward handlers (proxy + VNC, dispatches tunnel vs local)
+	openapiTai.Attach(group)
 
 	// Tai direct registration API (uses /tai-nodes/ prefix to avoid routing conflict with /tai/:taiID/)
 	group.POST("/tai-nodes/register", taiapi.HandleRegister)

openapi/tai/proxy.go

@@ -0,0 +1,60 @@
+package tai
+
+import (
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	yaoTai "github.com/yaoapp/yao/tai"
+)
+
+// handleLocalProxy resolves the container's HTTP address via Docker socket
+// and reverse-proxies the request.
+func handleLocalProxy(c *gin.Context, taiID string) {
+	res, ok := yaoTai.GetResources(taiID)
+	if !ok || res.Proxy == nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "proxy not available for node " + taiID})
+		return
+	}
+
+	// path format: /{containerID}:{port}/{rest...}
+	raw := strings.TrimPrefix(c.Param("path"), "/")
+	colonIdx := strings.Index(raw, ":")
+	if colonIdx < 0 {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid proxy path, expected /{containerID}:{port}/{path}"})
+		return
+	}
+
+	containerID := raw[:colonIdx]
+	rest := raw[colonIdx+1:]
+	slashIdx := strings.Index(rest, "/")
+	var portStr, subPath string
+	if slashIdx >= 0 {
+		portStr = rest[:slashIdx]
+		subPath = rest[slashIdx:]
+	} else {
+		portStr = rest
+		subPath = "/"
+	}
+
+	var port int
+	for _, ch := range portStr {
+		if ch < '0' || ch > '9' {
+			c.JSON(http.StatusBadRequest, gin.H{"error": "invalid port in proxy path"})
+			return
+		}
+		port = port*10 + int(ch-'0')
+	}
+	if port == 0 {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "missing port in proxy path"})
+		return
+	}
+
+	targetURL, err := res.Proxy.URL(c.Request.Context(), containerID, port, subPath)
+	if err != nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "resolve proxy target: " + err.Error()})
+		return
+	}
+
+	reverseProxy(c, targetURL)
+}

openapi/tai/tai.go

@@ -0,0 +1,39 @@
+package tai
+
+import (
+	"github.com/gin-gonic/gin"
+	yaoTai "github.com/yaoapp/yao/tai"
+	taitunnel "github.com/yaoapp/yao/tai/tunnel"
+)
+
+// Attach registers Tai forward routes on the given group.
+//
+//   - ANY /tai/:taiID/proxy/*path — HTTP forward (tunnel or local)
+//   - GET /tai/:taiID/vnc/*path   — VNC WebSocket forward (tunnel or local)
+func Attach(group *gin.RouterGroup) {
+	group.Any("/tai/:taiID/proxy/*path", handleProxy)
+	group.GET("/tai/:taiID/vnc/*path", handleVNC)
+}
+
+func handleProxy(c *gin.Context) {
+	taiID := c.Param("taiID")
+	if isLocalNode(taiID) {
+		handleLocalProxy(c, taiID)
+		return
+	}
+	taitunnel.HandleForwardLazy(c)
+}
+
+func handleVNC(c *gin.Context) {
+	taiID := c.Param("taiID")
+	if isLocalNode(taiID) {
+		handleLocalVNC(c, taiID)
+		return
+	}
+	taitunnel.HandleForwardLazy(c)
+}
+
+func isLocalNode(taiID string) bool {
+	meta, ok := yaoTai.GetNodeMeta(taiID)
+	return ok && meta.Mode == "local"
+}

openapi/tai/util.go

@@ -0,0 +1,95 @@
+package tai
+
+import (
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+)
+
+// extractContainerID parses container ID from *path param.
+// /{containerID}/ws → containerID
+func extractContainerID(path string) string {
+	path = strings.TrimPrefix(path, "/")
+	path = strings.TrimSuffix(path, "/ws")
+	path = strings.TrimSuffix(path, "/")
+	if path == "" || path == "__host__" {
+		return "__host__"
+	}
+	return path
+}
+
+// bridgeWebSocket copies messages bidirectionally between two WebSocket connections.
+func bridgeWebSocket(client, target *websocket.Conn) {
+	done := make(chan struct{}, 2)
+
+	go func() {
+		defer func() { done <- struct{}{} }()
+		for {
+			mt, data, err := client.ReadMessage()
+			if err != nil {
+				return
+			}
+			if err := target.WriteMessage(mt, data); err != nil {
+				return
+			}
+		}
+	}()
+
+	go func() {
+		defer func() { done <- struct{}{} }()
+		for {
+			mt, data, err := target.ReadMessage()
+			if err != nil {
+				return
+			}
+			if err := client.WriteMessage(mt, data); err != nil {
+				return
+			}
+		}
+	}()
+
+	<-done
+}
+
+// reverseProxy forwards an HTTP request to targetURL and streams the response back.
+func reverseProxy(c *gin.Context, targetURL string) {
+	req, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, targetURL, c.Request.Body)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "create proxy request: " + err.Error()})
+		return
+	}
+	for k, vv := range c.Request.Header {
+		for _, v := range vv {
+			req.Header.Add(k, v)
+		}
+	}
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "proxy request failed: " + err.Error()})
+		return
+	}
+	defer resp.Body.Close()
+
+	for k, vv := range resp.Header {
+		for _, v := range vv {
+			c.Writer.Header().Add(k, v)
+		}
+	}
+	c.Writer.WriteHeader(resp.StatusCode)
+	c.Writer.Flush()
+
+	buf := make([]byte, 32*1024)
+	for {
+		n, readErr := resp.Body.Read(buf)
+		if n > 0 {
+			c.Writer.Write(buf[:n])
+			c.Writer.Flush()
+		}
+		if readErr != nil {
+			return
+		}
+	}
+}

openapi/tai/vnc.go

@@ -0,0 +1,57 @@
+package tai
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	yaoTai "github.com/yaoapp/yao/tai"
+)
+
+var wsUpgrader = websocket.Upgrader{
+	CheckOrigin:  func(r *http.Request) bool { return true },
+	Subprotocols: []string{"binary"},
+}
+
+// handleLocalVNC resolves the container's VNC address via Docker socket
+// and proxies the WebSocket connection.
+func handleLocalVNC(c *gin.Context, taiID string) {
+	containerID := extractContainerID(c.Param("path"))
+	if containerID == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "missing container ID in path"})
+		return
+	}
+
+	res, ok := yaoTai.GetResources(taiID)
+	if !ok || res.VNC == nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "VNC not available for node " + taiID})
+		return
+	}
+
+	targetURL, err := res.VNC.URL(c.Request.Context(), containerID)
+	if err != nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "resolve VNC target: " + err.Error()})
+		return
+	}
+
+	clientConn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		return
+	}
+	defer clientConn.Close()
+
+	dialer := websocket.Dialer{
+		Subprotocols:     []string{"binary"},
+		HandshakeTimeout: 5 * time.Second,
+	}
+	targetConn, _, err := dialer.Dial(targetURL, nil)
+	if err != nil {
+		clientConn.WriteMessage(websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "VNC connection failed"))
+		return
+	}
+	defer targetConn.Close()
+
+	bridgeWebSocket(clientConn, targetConn)
+}