Implement tunnel mode support in Tai service

- Introduce a new "tunnel" scheme for the Tai client, allowing connections through Yao's reverse proxy. - Enhance the Tai registry to manage tunnel-connected nodes and their ports. - Add WebSocket and reverse proxy routes for tunnel connections in the OpenAPI server. - Implement tunnel-specific proxy and VNC handling to facilitate communication with containerized environments. - Update gRPC environment variable handling to support tunnel connections. These changes improve the Tai service's flexibility and connectivity options, enabling better integration with remote and containerized environments.
2026-03-06 17:12:36 +08:00 · 2026-03-06 17:12:36 +08:00 · 6408c140fc
commit 6408c140fc
parent 7a2902a575
14 changed files with 2188 additions and 5 deletions
--- a/cmd/start.go
+++ b/cmd/start.go
@ -35,6 +35,7 @@ import (
 	"github.com/yaoapp/yao/service"
 	"github.com/yaoapp/yao/setup"
 	"github.com/yaoapp/yao/share"
+	tairegistry "github.com/yaoapp/yao/tai/registry"
 	itask "github.com/yaoapp/yao/task"
 )

@ -231,6 +232,10 @@ var startCmd = &cobra.Command{
 		ischedule.Start()
 		defer ischedule.Stop()

+		// Initialize the global Tai registry for tunnel and direct connections
+		// (must happen before HTTP/gRPC start so handlers can access it)
+		tairegistry.Init(nil)
+
 		// Start HTTP Server
 		srv, err := service.Start(config.Conf)
 		defer func() {
--- a/grpc/grpc.go
+++ b/grpc/grpc.go
@ -192,6 +192,15 @@ func Stop() {
 	addrs = nil
 }

+// GRPCServer returns the active gRPC server instance.
+// Used by the Tai tunnel server to serve data channel connections
+// on the existing gRPC server.
+func GRPCServer() *grpc.Server {
+	mu.Lock()
+	defer mu.Unlock()
+	return server
+}
+
 // Addr returns all addresses the gRPC server is listening on.
 func Addr() []string {
 	mu.Lock()
--- a/openapi/openapi.go
+++ b/openapi/openapi.go
@ -28,6 +28,7 @@ import (
 	"github.com/yaoapp/yao/openapi/team"
 	openapiTrace "github.com/yaoapp/yao/openapi/trace"
 	"github.com/yaoapp/yao/openapi/user"
+	taitunnel "github.com/yaoapp/yao/tai/tunnel"
 )

 // Server is the OpenAPI server
@ -175,6 +176,12 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) {
 	sandbox.SetPathPrefix(baseURL)
 	sandbox.Attach(group.Group("/sandbox"), openapi.OAuth)

+	// Tai tunnel WebSocket and reverse proxy routes
+	group.GET("/ws/tai", taitunnel.HandleControl)
+	group.GET("/ws/tai/data/:channel_id", taitunnel.HandleData)
+	group.Any("/tai/:taiID/proxy/*path", taitunnel.HandleProxy)
+	group.GET("/tai/:taiID/vnc/*path", taitunnel.HandleVNC)
+
 	// Custom handlers (Defined by developer)

 }
--- a/sandbox/v2/grpc.go
+++ b/sandbox/v2/grpc.go
@ -4,6 +4,7 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
+	"net/url"
 	"strconv"
 	"strings"
 )
@ -35,6 +36,7 @@ func RevokeContainerTokens(refresh string) error {
 }

 // BuildGRPCEnv builds the gRPC environment variables for a sandbox container.
+// Supports tai:// (direct), tunnel:// (NAT traversal), and local modes.
 func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) map[string]string {
 	portStr := strconv.Itoa(grpcPort)
 	env := map[string]string{
@ -42,12 +44,34 @@ func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) m
 		"YAO_TOKEN":         access,
 		"YAO_REFRESH_TOKEN": refresh,
 	}
-	if pool != nil && strings.Contains(pool.Addr, "tai://") {
-		taiHost := strings.TrimPrefix(pool.Addr, "tai://")
+
+	if pool == nil {
+		env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr)
+		return env
+	}
+
+	switch {
+	case strings.HasPrefix(pool.Addr, "tunnel://"):
 		env["YAO_GRPC_TAI"] = "enable"
-		env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:9100", taiHost)
+		env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%d", grpcPort)
 		env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr)
-	} else {
+
+	case strings.HasPrefix(pool.Addr, "tai://"):
+		u, err := url.Parse(pool.Addr)
+		if err != nil {
+			env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr)
+			return env
+		}
+		taiHost := u.Hostname()
+		taiPort := u.Port()
+		if taiPort == "" {
+			taiPort = "9100"
+		}
+		env["YAO_GRPC_TAI"] = "enable"
+		env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%s", taiHost, taiPort)
+		env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr)
+
+	default:
 		env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr)
 	}
 	return env
--- a/tai/proxy/connect.go
+++ b/tai/proxy/connect.go
@ -21,6 +21,16 @@ func (r *remoteProxy) Connect(ctx context.Context, containerID string, opts Conn
 	return connect(ctx, baseURL, opts.Protocol, r.client)
 }

+// --- Tunnel Connect ---
+
+func (t *tunnelProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) {
+	baseURL, err := t.URL(ctx, containerID, opts.Port, opts.Path)
+	if err != nil {
+		return nil, err
+	}
+	return connect(ctx, baseURL, opts.Protocol, http.DefaultClient)
+}
+
 // --- Local Connect ---

 func (l *localProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) {
--- a/tai/proxy/proxy.go
+++ b/tai/proxy/proxy.go
@ -73,6 +73,28 @@ func (r *remoteProxy) Healthz(ctx context.Context) error {
 	return nil
 }

+// --- Tunnel implementation ---
+
+type tunnelProxy struct {
+	taiID   string
+	yaoBase string // e.g. "http://yao-host:5099"
+}
+
+// NewTunnel creates a Proxy that routes through Yao's reverse proxy for
+// tunnel-connected Tai instances. URLs point to {yaoBase}/tai/{taiID}/proxy/*.
+func NewTunnel(taiID, yaoBase string) Proxy {
+	return &tunnelProxy{taiID: taiID, yaoBase: strings.TrimRight(yaoBase, "/")}
+}
+
+func (t *tunnelProxy) URL(_ context.Context, containerID string, port int, path string) (string, error) {
+	path = strings.TrimPrefix(path, "/")
+	return fmt.Sprintf("%s/tai/%s/proxy/%s:%d/%s", t.yaoBase, t.taiID, containerID, port, path), nil
+}
+
+func (t *tunnelProxy) Healthz(_ context.Context) error {
+	return nil
+}
+
 // --- Local implementation ---

 type localProxy struct {
--- a/tai/registry/registry.go
+++ b/tai/registry/registry.go
@ -0,0 +1,401 @@
+package registry
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"log/slog"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// TaiNode represents a registered Tai instance (direct or tunnel).
+// Internal use only; external callers receive NodeSnapshot via Get()/List().
+type TaiNode struct {
+	TaiID        string
+	MachineID    string
+	Version      string
+	Auth         AuthInfo
+	Mode         string         // "direct" | "tunnel"
+	Addr         string         // direct mode: "tai-host"; tunnel mode: empty
+	YaoBase      string         // Yao server base URL reported by Tai (tunnel mode)
+	Ports        map[string]int // {"grpc":9100, "http":8080, "vnc":6080, "docker":2375}
+	Capabilities map[string]bool
+
+	ControlConn *websocket.Conn
+	connMu      sync.Mutex // protects ControlConn writes
+
+	Status      string // "online" | "offline" | "connecting"
+	ConnectedAt time.Time
+	LastPing    time.Time
+	PoolName    string
+
+	localListeners map[int]*tunnelListener
+}
+
+// NodeSnapshot is a read-only copy of TaiNode fields safe to use outside locks.
+type NodeSnapshot struct {
+	TaiID        string
+	MachineID    string
+	Version      string
+	Auth         AuthInfo
+	Mode         string
+	Addr         string
+	YaoBase      string
+	Ports        map[string]int
+	Capabilities map[string]bool
+	Status       string
+	ConnectedAt  time.Time
+	LastPing     time.Time
+	PoolName     string
+}
+
+func (n *TaiNode) snapshot() NodeSnapshot {
+	ports := make(map[string]int, len(n.Ports))
+	for k, v := range n.Ports {
+		ports[k] = v
+	}
+	caps := make(map[string]bool, len(n.Capabilities))
+	for k, v := range n.Capabilities {
+		caps[k] = v
+	}
+	return NodeSnapshot{
+		TaiID: n.TaiID, MachineID: n.MachineID, Version: n.Version,
+		Auth: n.Auth, Mode: n.Mode, Addr: n.Addr, YaoBase: n.YaoBase,
+		Ports: ports, Capabilities: caps,
+		Status: n.Status, ConnectedAt: n.ConnectedAt, LastPing: n.LastPing,
+		PoolName: n.PoolName,
+	}
+}
+
+// AuthInfo holds Yao user authorization extracted from OAuth token.
+type AuthInfo struct {
+	Subject  string
+	UserID   string
+	ClientID string
+	Scope    string
+	TeamID   string
+	TenantID string
+}
+
+// pendingChannel represents a channel awaiting Tai's data WS connection.
+type pendingChannel struct {
+	taiID  string
+	result chan net.Conn
+	timer  *time.Timer
+}
+
+// tunnelListener wraps a TCP listener that bridges each accepted connection
+// through the WS tunnel to a specific Tai port.
+type tunnelListener struct {
+	listener net.Listener
+	taiID    string
+	port     int
+	cancel   func()
+}
+
+var (
+	global *Registry
+	once   sync.Once
+)
+
+// Registry manages all Tai nodes (direct and tunnel).
+type Registry struct {
+	mu      sync.RWMutex
+	nodes   map[string]*TaiNode
+	pending map[string]*pendingChannel
+	logger  *slog.Logger
+}
+
+// Init initializes the global registry singleton.
+func Init(logger *slog.Logger) {
+	once.Do(func() {
+		if logger == nil {
+			logger = slog.Default()
+		}
+		global = &Registry{
+			nodes:   make(map[string]*TaiNode),
+			pending: make(map[string]*pendingChannel),
+			logger:  logger,
+		}
+	})
+}
+
+// Global returns the global registry instance.
+func Global() *Registry {
+	return global
+}
+
+// Register adds or updates a Tai node in the registry.
+func (r *Registry) Register(node *TaiNode) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	node.Status = "online"
+	node.ConnectedAt = time.Now()
+	node.LastPing = time.Now()
+	if node.localListeners == nil {
+		node.localListeners = make(map[int]*tunnelListener)
+	}
+	r.nodes[node.TaiID] = node
+
+	r.logger.Info("tai node registered",
+		"tai_id", node.TaiID, "mode", node.Mode, "version", node.Version)
+}
+
+// Unregister removes a Tai node and closes its local listeners and control connection.
+func (r *Registry) Unregister(taiID string) {
+	r.mu.Lock()
+	node, ok := r.nodes[taiID]
+	if ok {
+		for _, tl := range node.localListeners {
+			tl.cancel()
+			tl.listener.Close()
+		}
+		node.connMu.Lock()
+		if node.ControlConn != nil {
+			node.ControlConn.Close()
+			node.ControlConn = nil
+		}
+		node.connMu.Unlock()
+		delete(r.nodes, taiID)
+	}
+	r.mu.Unlock()
+
+	if ok {
+		r.logger.Info("tai node unregistered", "tai_id", taiID)
+	}
+}
+
+// Get returns a snapshot of a Tai node by ID. Returns nil, false if not found.
+func (r *Registry) Get(taiID string) (*NodeSnapshot, bool) {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	n, ok := r.nodes[taiID]
+	if !ok {
+		return nil, false
+	}
+	snap := n.snapshot()
+	return &snap, true
+}
+
+// List returns snapshots of all registered Tai nodes.
+func (r *Registry) List() []NodeSnapshot {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	result := make([]NodeSnapshot, 0, len(r.nodes))
+	for _, n := range r.nodes {
+		result = append(result, n.snapshot())
+	}
+	return result
+}
+
+// WriteControlJSON sends a JSON message on the node's control channel
+// with proper serialization. Returns error if node not found or not tunnel.
+func (r *Registry) WriteControlJSON(taiID string, v interface{}) error {
+	r.mu.RLock()
+	node := r.nodes[taiID]
+	r.mu.RUnlock()
+
+	if node == nil {
+		return fmt.Errorf("tai node %s not found", taiID)
+	}
+
+	node.connMu.Lock()
+	defer node.connMu.Unlock()
+	if node.ControlConn == nil {
+		return fmt.Errorf("tai node %s has no active control channel", taiID)
+	}
+	return node.ControlConn.WriteJSON(v)
+}
+
+// UpdatePing records a heartbeat timestamp.
+func (r *Registry) UpdatePing(taiID string) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if n, ok := r.nodes[taiID]; ok {
+		n.LastPing = time.Now()
+	}
+}
+
+// RequestChannel sends an "open" command to a tunnel-connected Tai via its
+// control channel. Returns a channel_id that Tai will use to connect back.
+// Blocks until the data channel is established or timeout.
+func (r *Registry) RequestChannel(taiID string, targetPort int) (string, chan net.Conn, error) {
+	r.mu.RLock()
+	node := r.nodes[taiID]
+	r.mu.RUnlock()
+
+	if node == nil {
+		return "", nil, fmt.Errorf("tai node %s not found", taiID)
+	}
+	if node.Mode != "tunnel" {
+		return "", nil, fmt.Errorf("tai node %s is not a tunnel node", taiID)
+	}
+	node.connMu.Lock()
+	hasConn := node.ControlConn != nil
+	node.connMu.Unlock()
+	if !hasConn {
+		return "", nil, fmt.Errorf("tai node %s has no active control channel", taiID)
+	}
+
+	channelID, err := generateChannelID()
+	if err != nil {
+		return "", nil, fmt.Errorf("generate channel_id: %w", err)
+	}
+
+	resultCh := make(chan net.Conn, 1)
+	timer := time.AfterFunc(30*time.Second, func() {
+		r.mu.Lock()
+		if pc, ok := r.pending[channelID]; ok {
+			close(pc.result)
+			delete(r.pending, channelID)
+		}
+		r.mu.Unlock()
+	})
+
+	r.mu.Lock()
+	r.pending[channelID] = &pendingChannel{taiID: taiID, result: resultCh, timer: timer}
+	r.mu.Unlock()
+
+	msg := map[string]interface{}{
+		"type":        "open",
+		"channel_id":  channelID,
+		"target_port": targetPort,
+	}
+	if err := r.WriteControlJSON(taiID, msg); err != nil {
+		r.mu.Lock()
+		delete(r.pending, channelID)
+		r.mu.Unlock()
+		timer.Stop()
+		return "", nil, fmt.Errorf("send open command: %w", err)
+	}
+
+	return channelID, resultCh, nil
+}
+
+// AcceptDataChannel resolves a pending channel when Tai connects its data WS.
+// The taiID must match the node that requested the channel via RequestChannel.
+func (r *Registry) AcceptDataChannel(channelID, taiID string, conn net.Conn) error {
+	r.mu.Lock()
+	pc, ok := r.pending[channelID]
+	if ok {
+		delete(r.pending, channelID)
+	}
+	r.mu.Unlock()
+
+	if !ok {
+		return fmt.Errorf("no pending channel for %s", channelID)
+	}
+	if pc.taiID != taiID {
+		pc.timer.Stop()
+		close(pc.result)
+		return fmt.Errorf("channel %s: tai_id mismatch (expected %s, got %s)", channelID, pc.taiID, taiID)
+	}
+	pc.timer.Stop()
+	pc.result <- conn
+	return nil
+}
+
+// OpenLocalListener creates a localhost TCP listener that tunnels every
+// accepted connection to the specified port on the given Tai node.
+// Returns the listener address (e.g. "127.0.0.1:54321").
+func (r *Registry) OpenLocalListener(taiID string, targetPort int) (net.Listener, error) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		return nil, fmt.Errorf("listen: %w", err)
+	}
+
+	ctx, cancel := newContext()
+	tl := &tunnelListener{listener: ln, taiID: taiID, port: targetPort, cancel: cancel}
+
+	r.mu.Lock()
+	node := r.nodes[taiID]
+	if node == nil {
+		r.mu.Unlock()
+		cancel()
+		ln.Close()
+		return nil, fmt.Errorf("tai node %s not found", taiID)
+	}
+	node.localListeners[targetPort] = tl
+	r.mu.Unlock()
+
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					r.logger.Debug("tunnel listener accept error", "err", err)
+					return
+				}
+			}
+			go r.bridgeTunnelConn(taiID, targetPort, conn)
+		}
+	}()
+
+	r.logger.Info("tunnel local listener started",
+		"tai_id", taiID, "target_port", targetPort, "local_addr", ln.Addr().String())
+	return ln, nil
+}
+
+func (r *Registry) bridgeTunnelConn(taiID string, targetPort int, localConn net.Conn) {
+	channelID, resultCh, err := r.RequestChannel(taiID, targetPort)
+	if err != nil {
+		localConn.Close()
+		r.logger.Error("request channel failed", "tai_id", taiID, "port", targetPort, "err", err)
+		return
+	}
+
+	remoteConn, ok := <-resultCh
+	if !ok || remoteConn == nil {
+		localConn.Close()
+		r.logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID)
+		return
+	}
+
+	bridgeTCP(localConn, remoteConn)
+}
+
+// bridgeTCP copies bytes bidirectionally between two net.Conn, closing both when done.
+func bridgeTCP(a, b net.Conn) {
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	cp := func(dst, src net.Conn) {
+		defer wg.Done()
+		io.Copy(dst, src)
+		dst.Close()
+	}
+
+	go cp(a, b)
+	go cp(b, a)
+	wg.Wait()
+}
+
+func generateChannelID() (string, error) {
+	b := make([]byte, 32)
+	if _, err := rand.Read(b); err != nil {
+		return "", err
+	}
+	return hex.EncodeToString(b), nil
+}
+
+type contextCancel struct {
+	done chan struct{}
+}
+
+func newContext() (*contextCancel, func()) {
+	cc := &contextCancel{done: make(chan struct{})}
+	return cc, func() { close(cc.done) }
+}
+
+func (c *contextCancel) Done() <-chan struct{} {
+	return c.done
+}
--- a/tai/registry/registry_test.go
+++ b/tai/registry/registry_test.go
@ -0,0 +1,499 @@
+package registry
+
+import (
+	"log/slog"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// newTestRegistry creates a standalone registry for testing (bypasses global singleton).
+func newTestRegistry() *Registry {
+	return &Registry{
+		nodes:   make(map[string]*TaiNode),
+		pending: make(map[string]*pendingChannel),
+		logger:  slog.Default(),
+	}
+}
+
+func TestRegister_SetsFieldsAndOnline(t *testing.T) {
+	r := newTestRegistry()
+	node := &TaiNode{
+		TaiID:     "tai-001",
+		MachineID: "m-abc",
+		Version:   "1.0.0",
+		Mode:      "tunnel",
+		Ports:     map[string]int{"grpc": 9100},
+	}
+	r.Register(node)
+
+	snap, ok := r.Get("tai-001")
+	if !ok {
+		t.Fatal("expected node to exist after Register")
+	}
+	if snap.Status != "online" {
+		t.Errorf("Status = %q, want online", snap.Status)
+	}
+	if snap.MachineID != "m-abc" {
+		t.Errorf("MachineID = %q, want m-abc", snap.MachineID)
+	}
+	if snap.ConnectedAt.IsZero() {
+		t.Error("ConnectedAt should be set")
+	}
+}
+
+func TestRegister_Overwrite(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001", Version: "1.0"})
+	r.Register(&TaiNode{TaiID: "tai-001", Version: "2.0"})
+
+	snap, ok := r.Get("tai-001")
+	if !ok {
+		t.Fatal("node should exist")
+	}
+	if snap.Version != "2.0" {
+		t.Errorf("Version = %q, want 2.0 after re-register", snap.Version)
+	}
+}
+
+func TestUnregister_RemovesNode(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001"})
+	r.Unregister("tai-001")
+
+	if _, ok := r.Get("tai-001"); ok {
+		t.Error("expected node to be removed after Unregister")
+	}
+}
+
+func TestUnregister_Nonexistent(t *testing.T) {
+	r := newTestRegistry()
+	r.Unregister("ghost")
+}
+
+func TestGet_NotFound(t *testing.T) {
+	r := newTestRegistry()
+	if _, ok := r.Get("missing"); ok {
+		t.Error("expected false for missing node")
+	}
+}
+
+func TestList_Empty(t *testing.T) {
+	r := newTestRegistry()
+	if got := r.List(); len(got) != 0 {
+		t.Errorf("List() = %d items, want 0", len(got))
+	}
+}
+
+func TestList_MultipleNodes(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "a"})
+	r.Register(&TaiNode{TaiID: "b"})
+	r.Register(&TaiNode{TaiID: "c"})
+
+	list := r.List()
+	if len(list) != 3 {
+		t.Errorf("List() = %d items, want 3", len(list))
+	}
+
+	ids := map[string]bool{}
+	for _, snap := range list {
+		ids[snap.TaiID] = true
+	}
+	for _, id := range []string{"a", "b", "c"} {
+		if !ids[id] {
+			t.Errorf("missing node %q in List()", id)
+		}
+	}
+}
+
+func TestSnapshot_DeepCopy(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{
+		TaiID: "tai-001",
+		Ports: map[string]int{"grpc": 9100, "http": 8080},
+	})
+
+	snap, _ := r.Get("tai-001")
+	snap.Ports["grpc"] = 0
+
+	snap2, _ := r.Get("tai-001")
+	if snap2.Ports["grpc"] != 9100 {
+		t.Error("snapshot modification leaked into registry node")
+	}
+}
+
+func TestUpdatePing(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001"})
+	time.Sleep(10 * time.Millisecond)
+
+	r.UpdatePing("tai-001")
+	snap, _ := r.Get("tai-001")
+	if snap.LastPing.Before(snap.ConnectedAt) {
+		t.Error("LastPing should be after ConnectedAt")
+	}
+}
+
+func TestUpdatePing_NonexistentNode(t *testing.T) {
+	r := newTestRegistry()
+	r.UpdatePing("ghost")
+}
+
+func TestWriteControlJSON_NoNode(t *testing.T) {
+	r := newTestRegistry()
+	err := r.WriteControlJSON("missing", map[string]string{"type": "test"})
+	if err == nil {
+		t.Fatal("expected error for missing node")
+	}
+}
+
+func TestWriteControlJSON_NilConn(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001"})
+	err := r.WriteControlJSON("tai-001", map[string]string{"type": "test"})
+	if err == nil {
+		t.Fatal("expected error for nil ControlConn")
+	}
+}
+
+func TestRequestChannel_NotFound(t *testing.T) {
+	r := newTestRegistry()
+	_, _, err := r.RequestChannel("ghost", 9100)
+	if err == nil {
+		t.Fatal("expected error for missing node")
+	}
+}
+
+func TestRequestChannel_DirectMode(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001", Mode: "direct"})
+	_, _, err := r.RequestChannel("tai-001", 9100)
+	if err == nil {
+		t.Fatal("expected error for direct-mode node")
+	}
+}
+
+func TestAcceptDataChannel_NotPending(t *testing.T) {
+	r := newTestRegistry()
+	pipe1, pipe2 := net.Pipe()
+	defer pipe1.Close()
+	defer pipe2.Close()
+
+	err := r.AcceptDataChannel("unknown-channel", "tai-001", pipe1)
+	if err == nil {
+		t.Fatal("expected error for non-pending channel")
+	}
+}
+
+func TestAcceptDataChannel_TaiIDMismatch(t *testing.T) {
+	r := newTestRegistry()
+
+	resultCh := make(chan net.Conn, 1)
+	timer := time.AfterFunc(5*time.Second, func() {})
+	r.mu.Lock()
+	r.pending["ch-001"] = &pendingChannel{taiID: "tai-owner", result: resultCh, timer: timer}
+	r.mu.Unlock()
+
+	pipe1, pipe2 := net.Pipe()
+	defer pipe1.Close()
+	defer pipe2.Close()
+
+	err := r.AcceptDataChannel("ch-001", "tai-intruder", pipe1)
+	if err == nil {
+		t.Fatal("expected error for tai_id mismatch")
+	}
+}
+
+func TestAcceptDataChannel_Success(t *testing.T) {
+	r := newTestRegistry()
+
+	resultCh := make(chan net.Conn, 1)
+	timer := time.AfterFunc(5*time.Second, func() {})
+	r.mu.Lock()
+	r.pending["ch-002"] = &pendingChannel{taiID: "tai-001", result: resultCh, timer: timer}
+	r.mu.Unlock()
+
+	pipe1, pipe2 := net.Pipe()
+	defer pipe2.Close()
+
+	if err := r.AcceptDataChannel("ch-002", "tai-001", pipe1); err != nil {
+		t.Fatalf("AcceptDataChannel: %v", err)
+	}
+
+	select {
+	case conn := <-resultCh:
+		if conn == nil {
+			t.Fatal("expected non-nil conn")
+		}
+		conn.Close()
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for conn on resultCh")
+	}
+}
+
+func TestGenerateChannelID_Unique(t *testing.T) {
+	seen := make(map[string]bool)
+	for i := 0; i < 100; i++ {
+		id, err := generateChannelID()
+		if err != nil {
+			t.Fatalf("generateChannelID: %v", err)
+		}
+		if len(id) != 64 {
+			t.Errorf("len = %d, want 64 hex chars", len(id))
+		}
+		if seen[id] {
+			t.Fatalf("duplicate channel ID: %s", id)
+		}
+		seen[id] = true
+	}
+}
+
+func TestBridgeTCP(t *testing.T) {
+	a1, a2 := net.Pipe()
+	b1, b2 := net.Pipe()
+
+	go bridgeTCP(a2, b1)
+
+	msg := []byte("hello tunnel")
+	go func() {
+		a1.Write(msg)
+		a1.Close()
+	}()
+
+	buf := make([]byte, 64)
+	n, _ := b2.Read(buf)
+	if string(buf[:n]) != "hello tunnel" {
+		t.Errorf("got %q, want %q", buf[:n], "hello tunnel")
+	}
+	b2.Close()
+}
+
+func TestConcurrentRegisterGet(t *testing.T) {
+	r := newTestRegistry()
+	var wg sync.WaitGroup
+
+	for i := 0; i < 50; i++ {
+		wg.Add(2)
+		id := "tai-" + string(rune('A'+i%26))
+
+		go func() {
+			defer wg.Done()
+			r.Register(&TaiNode{TaiID: id, Mode: "tunnel"})
+		}()
+
+		go func() {
+			defer wg.Done()
+			r.Get(id)
+			r.List()
+		}()
+	}
+
+	wg.Wait()
+}
+
+func TestWriteControlJSON_Success(t *testing.T) {
+	done := make(chan map[string]string, 1)
+
+	srv := newWSServer(func(conn *websocket.Conn) {
+		var msg map[string]string
+		conn.ReadJSON(&msg)
+		done <- msg
+		conn.Close()
+	})
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
+
+	payload := map[string]string{"type": "test", "data": "hello"}
+	if err := r.WriteControlJSON("tai-001", payload); err != nil {
+		t.Fatalf("WriteControlJSON: %v", err)
+	}
+
+	select {
+	case got := <-done:
+		if got["type"] != "test" {
+			t.Errorf("type = %q, want test", got["type"])
+		}
+		if got["data"] != "hello" {
+			t.Errorf("data = %q, want hello", got["data"])
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout waiting for server to receive message")
+	}
+}
+
+func TestRequestChannel_Success(t *testing.T) {
+	openCh := make(chan map[string]interface{}, 1)
+
+	srv := newWSServer(func(conn *websocket.Conn) {
+		var msg map[string]interface{}
+		conn.ReadJSON(&msg)
+		openCh <- msg
+		time.Sleep(time.Second)
+		conn.Close()
+	})
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
+
+	channelID, resultCh, err := r.RequestChannel("tai-001", 9100)
+	if err != nil {
+		t.Fatalf("RequestChannel: %v", err)
+	}
+	if channelID == "" {
+		t.Fatal("channelID should not be empty")
+	}
+	if len(channelID) != 64 {
+		t.Errorf("channelID len = %d, want 64", len(channelID))
+	}
+	if resultCh == nil {
+		t.Fatal("resultCh should not be nil")
+	}
+
+	select {
+	case cmd := <-openCh:
+		if cmd["type"] != "open" {
+			t.Errorf("cmd type = %v, want open", cmd["type"])
+		}
+		if cmd["channel_id"] != channelID {
+			t.Errorf("cmd channel_id = %v, want %s", cmd["channel_id"], channelID)
+		}
+		if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 {
+			t.Errorf("cmd target_port = %v, want 9100", cmd["target_port"])
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout waiting for open command")
+	}
+}
+
+func TestRequestChannel_NoControlConn(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel"})
+
+	_, _, err := r.RequestChannel("tai-001", 9100)
+	if err == nil {
+		t.Fatal("expected error for nil ControlConn")
+	}
+}
+
+func TestOpenLocalListener_Success(t *testing.T) {
+	r := newTestRegistry()
+
+	controlCh := make(chan map[string]interface{}, 1)
+	srv := newWSServer(func(conn *websocket.Conn) {
+		for {
+			var msg map[string]interface{}
+			if err := conn.ReadJSON(&msg); err != nil {
+				return
+			}
+			controlCh <- msg
+		}
+	})
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+
+	r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
+
+	ln, err := r.OpenLocalListener("tai-001", 9100)
+	if err != nil {
+		t.Fatalf("OpenLocalListener: %v", err)
+	}
+	defer ln.Close()
+
+	addr := ln.Addr().String()
+	if addr == "" {
+		t.Fatal("listener address should not be empty")
+	}
+	if !strings.HasPrefix(addr, "127.0.0.1:") {
+		t.Errorf("addr = %q, want 127.0.0.1:*", addr)
+	}
+
+	conn, err := net.DialTimeout("tcp", addr, time.Second)
+	if err != nil {
+		t.Fatalf("connect to local listener: %v", err)
+	}
+	defer conn.Close()
+
+	select {
+	case cmd := <-controlCh:
+		if cmd["type"] != "open" {
+			t.Errorf("open cmd type = %v, want open", cmd["type"])
+		}
+		if _, ok := cmd["channel_id"].(string); !ok {
+			t.Error("open cmd missing channel_id")
+		}
+		if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 {
+			t.Errorf("target_port = %v, want 9100", cmd["target_port"])
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout waiting for open command from local listener")
+	}
+}
+
+func TestOpenLocalListener_NodeNotFound(t *testing.T) {
+	r := newTestRegistry()
+	_, err := r.OpenLocalListener("ghost", 9100)
+	if err == nil {
+		t.Fatal("expected error for missing node")
+	}
+}
+
+func newWSServer(handler func(*websocket.Conn)) *httptest.Server {
+	up := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := up.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		handler(conn)
+	}))
+}
+
+func TestNodeSnapshot_AuthInfo(t *testing.T) {
+	r := newTestRegistry()
+	r.Register(&TaiNode{
+		TaiID: "tai-001",
+		Auth: AuthInfo{
+			Subject:  "user123",
+			ClientID: "tai-001",
+			Scope:    "tai:tunnel",
+		},
+	})
+
+	snap, ok := r.Get("tai-001")
+	if !ok {
+		t.Fatal("node not found")
+	}
+	if snap.Auth.Subject != "user123" {
+		t.Errorf("Auth.Subject = %q, want user123", snap.Auth.Subject)
+	}
+	if snap.Auth.Scope != "tai:tunnel" {
+		t.Errorf("Auth.Scope = %q, want tai:tunnel", snap.Auth.Scope)
+	}
+}
--- a/tai/registry/testing.go
+++ b/tai/registry/testing.go
@ -0,0 +1,31 @@
+package registry
+
+import (
+	"log/slog"
+	"net"
+	"time"
+)
+
+// NewForTest creates a standalone Registry for use in tests.
+// Not intended for production use.
+func NewForTest() *Registry {
+	return &Registry{
+		nodes:   make(map[string]*TaiNode),
+		pending: make(map[string]*pendingChannel),
+		logger:  slog.Default(),
+	}
+}
+
+// SetGlobalForTest replaces the global registry singleton for testing.
+// Not intended for production use.
+func SetGlobalForTest(r *Registry) {
+	global = r
+}
+
+// SetPendingForTest injects a pending channel entry for testing.
+// Not intended for production use.
+func (r *Registry) SetPendingForTest(channelID, taiID string, result chan net.Conn, timer *time.Timer) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.pending[channelID] = &pendingChannel{taiID: taiID, result: result, timer: timer}
+}
--- a/tai/tai.go
+++ b/tai/tai.go
@ -3,6 +3,7 @@ package tai
 import (
 	"context"
 	"fmt"
+	"net"
 	"net/http"
 	"net/url"
 	"strconv"
@ -10,6 +11,7 @@ import (
 	"time"

 	"github.com/yaoapp/yao/tai/proxy"
+	"github.com/yaoapp/yao/tai/registry"
 	"github.com/yaoapp/yao/tai/sandbox"
 	sipb "github.com/yaoapp/yao/tai/serverinfo/pb"
 	"github.com/yaoapp/yao/tai/vnc"
@ -124,7 +126,7 @@ func mergedPorts(p Ports) Ports {

 // Client provides unified access to all Tai SDK sub-packages.
 type Client struct {
-	scheme   string // "tai" or "docker"
+	scheme   string // "tai", "docker", or "tunnel"
 	host     string
 	addr     string
 	ports    Ports
@ -135,6 +137,9 @@ type Client struct {
 	prx      proxy.Proxy
 	vc       vnc.VNC
 	grpcConn *grpc.ClientConn
+
+	// tunnel mode: local listeners that bridge to Tai via WS
+	tunnelListeners []net.Listener
 }

 // New creates a Client based on the address protocol:
@ -172,6 +177,8 @@ func New(addr string, opts ...Option) (*Client, error) {
 		return c.initLocal(cfg)
 	case "tai":
 		return c.initRemote(cfg)
+	case "tunnel":
+		return c.initTunnel(cfg)
 	default:
 		return nil, fmt.Errorf("unsupported scheme: %s", scheme)
 	}
@ -256,9 +263,97 @@ func (c *Client) initRemote(cfg *config) (*Client, error) {
 	hc := cfg.httpClient
 	c.prx = proxy.NewRemote(c.host, c.ports.HTTP, hc)
 	c.vc = vnc.NewRemote(c.host, c.ports.VNC, hc)
+
+	if reg := registry.Global(); reg != nil {
+		reg.Register(&registry.TaiNode{
+			TaiID: c.host,
+			Mode:  "direct",
+			Addr:  c.host,
+			Ports: map[string]int{
+				"grpc":   c.ports.GRPC,
+				"http":   c.ports.HTTP,
+				"vnc":    c.ports.VNC,
+				"docker": c.ports.Docker,
+				"k8s":    c.ports.K8s,
+			},
+		})
+	}
+
 	return c, nil
 }

+func (c *Client) initTunnel(cfg *config) (*Client, error) {
+	reg := registry.Global()
+	if reg == nil {
+		return nil, fmt.Errorf("tai registry not initialized")
+	}
+
+	taiID := c.host // for tunnel:// scheme, host stores the taiID
+	node, ok := reg.Get(taiID)
+	if !ok || node.Status != "online" {
+		return nil, fmt.Errorf("tai node %s not online", taiID)
+	}
+
+	c.ports = Ports{
+		GRPC:   nodePort(node.Ports, "grpc", 9100),
+		HTTP:   nodePort(node.Ports, "http", 8080),
+		VNC:    nodePort(node.Ports, "vnc", 6080),
+		Docker: nodePort(node.Ports, "docker", 2375),
+	}
+
+	grpcLn, err := reg.OpenLocalListener(taiID, c.ports.GRPC)
+	if err != nil {
+		return nil, fmt.Errorf("open grpc tunnel listener: %w", err)
+	}
+	c.tunnelListeners = append(c.tunnelListeners, grpcLn)
+
+	grpcAddr := grpcLn.Addr().String()
+	conn, err := grpc.NewClient("passthrough:///"+grpcAddr,
+		grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		grpcLn.Close()
+		return nil, fmt.Errorf("grpc dial tunnel %s: %w", grpcAddr, err)
+	}
+	c.grpcConn = conn
+	c.vol = volume.NewRemote(conn)
+
+	dockerLn, err := reg.OpenLocalListener(taiID, c.ports.Docker)
+	if err != nil {
+		conn.Close()
+		grpcLn.Close()
+		return nil, fmt.Errorf("open docker tunnel listener: %w", err)
+	}
+	c.tunnelListeners = append(c.tunnelListeners, dockerLn)
+
+	sbAddr := fmt.Sprintf("tcp://%s", dockerLn.Addr().String())
+	sb, err := sandbox.NewDocker(sbAddr)
+	if err != nil {
+		c.closeTunnelListeners()
+		conn.Close()
+		return nil, err
+	}
+	c.sb = sb
+	c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
+
+	c.prx = proxy.NewTunnel(taiID, node.YaoBase)
+	c.vc = vnc.NewTunnel(taiID, node.YaoBase)
+	return c, nil
+}
+
+func (c *Client) closeTunnelListeners() {
+	for _, ln := range c.tunnelListeners {
+		ln.Close()
+	}
+	c.tunnelListeners = nil
+}
+
+func nodePort(ports map[string]int, key string, fallback int) int {
+	if p, ok := ports[key]; ok && p > 0 {
+		return p
+	}
+	return fallback
+}
+
 // Close releases all resources.
 func (c *Client) Close() error {
 	var errs []error
@ -277,6 +372,12 @@ func (c *Client) Close() error {
 			errs = append(errs, err)
 		}
 	}
+	c.closeTunnelListeners()
+	if c.scheme == "tai" {
+		if reg := registry.Global(); reg != nil {
+			reg.Unregister(c.host)
+		}
+	}
 	if len(errs) > 0 {
 		return fmt.Errorf("close: %v", errs)
 	}
@ -355,6 +456,13 @@ func parseAddr(addr string) (scheme, host, dockerAddr string, grpcPort int, err
 		}
 		return "tai", hostname, "", grpcPort, nil

+	case "tunnel":
+		taiID := u.Host
+		if taiID == "" {
+			return "", "", "", 0, fmt.Errorf("tunnel:// requires a tai ID")
+		}
+		return "tunnel", taiID, "", 0, nil
+
 	case "docker":
 		return "docker", "", addr, 0, nil

--- a/tai/tunnel/proxy.go
+++ b/tai/tunnel/proxy.go
@ -0,0 +1,172 @@
+package tunnel
+
+import (
+	"bufio"
+	"io"
+	"log/slog"
+	"net"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	"github.com/yaoapp/yao/tai/registry"
+)
+
+// HandleProxy handles HTTP reverse proxy requests for a tunnel-connected Tai:
+// ANY /tai/:taiID/proxy/*path
+// Opens a data channel to Tai's HTTP port, forwards the HTTP request,
+// and streams the response back.
+func HandleProxy(c *gin.Context) {
+	logger := slog.Default()
+	reg := registry.Global()
+	if reg == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"})
+		return
+	}
+
+	taiID := c.Param("taiID")
+	node, ok := reg.Get(taiID)
+	if !ok || node.Status != "online" {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"})
+		return
+	}
+
+	httpPort := node.Ports["http"]
+	if httpPort == 0 {
+		httpPort = 8080
+	}
+
+	channelID, resultCh, err := reg.RequestChannel(taiID, httpPort)
+	if err != nil {
+		logger.Error("request channel failed", "tai_id", taiID, "err", err)
+		c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"})
+		return
+	}
+
+	remoteConn, ok := <-resultCh
+	if !ok || remoteConn == nil {
+		logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID)
+		c.JSON(http.StatusGatewayTimeout, gin.H{"error": "data channel timeout"})
+		return
+	}
+	defer remoteConn.Close()
+
+	path := c.Param("path")
+	outReq, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, "http://tai-tunnel"+path, c.Request.Body)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "build request failed"})
+		return
+	}
+	outReq.Header = c.Request.Header.Clone()
+	outReq.Host = c.Request.Host
+
+	if err := outReq.Write(remoteConn); err != nil {
+		logger.Error("write request to tunnel", "err", err)
+		c.JSON(http.StatusBadGateway, gin.H{"error": "write to tunnel failed"})
+		return
+	}
+
+	resp, err := http.ReadResponse(bufio.NewReader(remoteConn), outReq)
+	if err != nil {
+		logger.Error("read response from tunnel", "err", err)
+		c.JSON(http.StatusBadGateway, gin.H{"error": "read from tunnel failed"})
+		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)
+	io.Copy(c.Writer, resp.Body)
+}
+
+// HandleVNC handles VNC WebSocket proxying for a tunnel-connected Tai:
+// GET /tai/:taiID/vnc/*path
+// Upgrades the client connection to WebSocket, opens a data channel to
+// Tai's VNC port, and bridges the two WebSocket connections.
+func HandleVNC(c *gin.Context) {
+	logger := slog.Default()
+	reg := registry.Global()
+	if reg == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"})
+		return
+	}
+
+	taiID := c.Param("taiID")
+	node, ok := reg.Get(taiID)
+	if !ok || node.Status != "online" {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"})
+		return
+	}
+
+	vncPort := node.Ports["vnc"]
+	if vncPort == 0 {
+		vncPort = 6080
+	}
+
+	channelID, resultCh, err := reg.RequestChannel(taiID, vncPort)
+	if err != nil {
+		logger.Error("request vnc channel failed", "tai_id", taiID, "err", err)
+		c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"})
+		return
+	}
+
+	clientConn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error("ws upgrade client failed", "err", err)
+		return
+	}
+
+	taiConn, ok := <-resultCh
+	if !ok || taiConn == nil {
+		logger.Error("vnc data channel timeout", "tai_id", taiID, "channel_id", channelID)
+		clientConn.Close()
+		return
+	}
+
+	bridgeWSToConn(clientConn, taiConn)
+}
+
+// bridgeWSToConn bridges a client WebSocket to a net.Conn (tunnel data channel).
+func bridgeWSToConn(clientWS *websocket.Conn, taiConn net.Conn) {
+	done := make(chan struct{}, 2)
+
+	// client WS -> tai conn
+	go func() {
+		defer func() { done <- struct{}{} }()
+		for {
+			_, data, err := clientWS.ReadMessage()
+			if err != nil {
+				return
+			}
+			if _, err := taiConn.Write(data); err != nil {
+				return
+			}
+		}
+	}()
+
+	// tai conn -> client WS
+	go func() {
+		defer func() { done <- struct{}{} }()
+		buf := make([]byte, 32*1024)
+		for {
+			n, err := taiConn.Read(buf)
+			if n > 0 {
+				if wErr := clientWS.WriteMessage(websocket.BinaryMessage, buf[:n]); wErr != nil {
+					return
+				}
+			}
+			if err != nil {
+				return
+			}
+		}
+	}()
+
+	<-done
+	clientWS.Close()
+	taiConn.Close()
+	<-done
+}
--- a/tai/tunnel/server.go
+++ b/tai/tunnel/server.go
@ -0,0 +1,267 @@
+package tunnel
+
+import (
+	"fmt"
+	"io"
+	"log/slog"
+	"net"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	oauth "github.com/yaoapp/yao/openapi/oauth"
+	"github.com/yaoapp/yao/tai/registry"
+)
+
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool { return true },
+}
+
+// HandleControl handles the Tai control channel WebSocket: GET /ws/tai.
+// Authenticates via Bearer token, reads register + ping messages,
+// and maintains the Tai node in the global registry.
+func HandleControl(c *gin.Context) {
+	logger := slog.Default()
+	reg := registry.Global()
+	if reg == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"})
+		return
+	}
+
+	bearer := extractBearer(c.Request)
+	if bearer == "" {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
+		return
+	}
+
+	authInfo, err := authenticateBearerFunc(bearer)
+	if err != nil {
+		logger.Warn("tunnel auth failed", "err", err)
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
+		return
+	}
+
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error("ws upgrade failed", "err", err)
+		return
+	}
+
+	// Read the register message
+	var regMsg registerMessage
+	if err := conn.ReadJSON(&regMsg); err != nil {
+		logger.Error("read register message", "err", err)
+		conn.Close()
+		return
+	}
+	if regMsg.Type != "register" {
+		logger.Error("expected register message", "got", regMsg.Type)
+		conn.Close()
+		return
+	}
+	if regMsg.TaiID == "" {
+		logger.Error("register message missing tai_id")
+		conn.Close()
+		return
+	}
+
+	node := &registry.TaiNode{
+		TaiID:        regMsg.TaiID,
+		MachineID:    regMsg.MachineID,
+		Version:      regMsg.Version,
+		Auth:         authInfo,
+		Mode:         "tunnel",
+		YaoBase:      regMsg.Server,
+		Ports:        regMsg.Ports,
+		Capabilities: regMsg.Capabilities,
+		ControlConn:  conn,
+	}
+	reg.Register(node)
+	defer func() {
+		reg.Unregister(regMsg.TaiID)
+		logger.Info("tai tunnel disconnected", "tai_id", regMsg.TaiID)
+	}()
+
+	if err := reg.WriteControlJSON(regMsg.TaiID, map[string]string{"type": "registered", "tai_id": regMsg.TaiID}); err != nil {
+		logger.Error("write registered response", "err", err)
+		return
+	}
+
+	logger.Info("tai tunnel connected", "tai_id", regMsg.TaiID, "version", regMsg.Version)
+
+	for {
+		var msg controlMsg
+		if err := conn.ReadJSON(&msg); err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+				logger.Debug("control channel read error", "err", err)
+			}
+			return
+		}
+
+		switch msg.Type {
+		case "ping":
+			reg.UpdatePing(regMsg.TaiID)
+			if err := reg.WriteControlJSON(regMsg.TaiID, map[string]string{"type": "pong"}); err != nil {
+				logger.Debug("pong write failed", "err", err)
+				return
+			}
+		default:
+			logger.Debug("unknown control message", "type", msg.Type)
+		}
+	}
+}
+
+// HandleData handles a Tai data channel WebSocket: GET /ws/tai/data/:channel_id.
+// Authenticates via Bearer token, verifies the caller matches the pending
+// channel's owner, then wraps the WS as a net.Conn for bidirectional bridging.
+func HandleData(c *gin.Context) {
+	logger := slog.Default()
+	reg := registry.Global()
+	if reg == nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"})
+		return
+	}
+
+	bearer := extractBearer(c.Request)
+	if bearer == "" {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
+		return
+	}
+	authInfo, err := authenticateBearerFunc(bearer)
+	if err != nil {
+		logger.Warn("data channel auth failed", "err", err)
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
+		return
+	}
+
+	channelID := c.Param("channel_id")
+	if channelID == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "missing channel_id"})
+		return
+	}
+
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error("ws data upgrade failed", "err", err)
+		return
+	}
+
+	wsConn := newWSConn(conn)
+	if err := reg.AcceptDataChannel(channelID, authInfo.ClientID, wsConn); err != nil {
+		logger.Debug("accept data channel failed", "channel_id", channelID, "err", err)
+		conn.Close()
+		return
+	}
+}
+
+// registerMessage is the JSON structure for Tai's register message.
+type registerMessage struct {
+	Type         string          `json:"type"`
+	TaiID        string          `json:"tai_id"`
+	MachineID    string          `json:"machine_id"`
+	Version      string          `json:"version"`
+	Server       string          `json:"server"`
+	Ports        map[string]int  `json:"ports"`
+	Capabilities map[string]bool `json:"capabilities"`
+}
+
+// controlMsg is a generic control channel message.
+type controlMsg struct {
+	Type string `json:"type"`
+}
+
+func extractBearer(r *http.Request) string {
+	auth := r.Header.Get("Authorization")
+	if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") {
+		return auth[7:]
+	}
+	return ""
+}
+
+var authenticateBearerFunc = authenticateBearerDefault
+
+func authenticateBearerDefault(token string) (registry.AuthInfo, error) {
+	svc := oauth.OAuth
+	if svc == nil {
+		return registry.AuthInfo{}, fmt.Errorf("oauth service not initialized")
+	}
+
+	result, err := svc.AuthenticateToken(oauth.AuthInput{
+		AccessToken: token,
+	})
+	if err != nil {
+		return registry.AuthInfo{}, err
+	}
+
+	info := registry.AuthInfo{}
+	if result.Info != nil {
+		info.Subject = result.Info.Subject
+		info.UserID = result.Info.UserID
+		info.ClientID = result.Info.ClientID
+		info.Scope = result.Info.Scope
+		info.TeamID = result.Info.TeamID
+		info.TenantID = result.Info.TenantID
+	}
+	return info, nil
+}
+
+// wsConn wraps a gorilla/websocket.Conn to implement net.Conn for raw byte bridging.
+type wsConn struct {
+	ws     *websocket.Conn
+	reader io.Reader
+	mu     sync.Mutex
+}
+
+func newWSConn(ws *websocket.Conn) *wsConn {
+	return &wsConn{ws: ws}
+}
+
+func (c *wsConn) Read(p []byte) (int, error) {
+	for {
+		if c.reader != nil {
+			n, err := c.reader.Read(p)
+			if n > 0 {
+				return n, nil
+			}
+			c.reader = nil
+			if err != nil && err != io.EOF {
+				return 0, err
+			}
+		}
+		_, reader, err := c.ws.NextReader()
+		if err != nil {
+			return 0, err
+		}
+		c.reader = reader
+	}
+}
+
+func (c *wsConn) Write(p []byte) (int, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	err := c.ws.WriteMessage(websocket.BinaryMessage, p)
+	if err != nil {
+		return 0, err
+	}
+	return len(p), nil
+}
+
+func (c *wsConn) Close() error {
+	return c.ws.Close()
+}
+
+func (c *wsConn) LocalAddr() net.Addr  { return c.ws.LocalAddr() }
+func (c *wsConn) RemoteAddr() net.Addr { return c.ws.RemoteAddr() }
+
+func (c *wsConn) SetDeadline(t time.Time) error {
+	if err := c.ws.SetReadDeadline(t); err != nil {
+		return err
+	}
+	return c.ws.SetWriteDeadline(t)
+}
+
+func (c *wsConn) SetReadDeadline(t time.Time) error  { return c.ws.SetReadDeadline(t) }
+func (c *wsConn) SetWriteDeadline(t time.Time) error { return c.ws.SetWriteDeadline(t) }
--- a/tai/tunnel/server_test.go
+++ b/tai/tunnel/server_test.go
@ -0,0 +1,604 @@
+package tunnel
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	"github.com/yaoapp/yao/tai/registry"
+)
+
+func init() {
+	gin.SetMode(gin.TestMode)
+}
+
+func setupTestRegistry() *registry.Registry {
+	r := registry.NewForTest()
+	registry.SetGlobalForTest(r)
+	return r
+}
+
+func mockAuth(info registry.AuthInfo, authErr error) func() {
+	old := authenticateBearerFunc
+	authenticateBearerFunc = func(token string) (registry.AuthInfo, error) {
+		return info, authErr
+	}
+	return func() { authenticateBearerFunc = old }
+}
+
+// --- extractBearer ---
+
+func TestExtractBearer(t *testing.T) {
+	tests := []struct {
+		name   string
+		header string
+		want   string
+	}{
+		{"valid", "Bearer abc123", "abc123"},
+		{"lowercase", "bearer xyz", "xyz"},
+		{"empty", "", ""},
+		{"no_scheme", "abc123", ""},
+		{"only_bearer", "Bearer ", ""},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := &http.Request{Header: http.Header{}}
+			if tt.header != "" {
+				r.Header.Set("Authorization", tt.header)
+			}
+			got := extractBearer(r)
+			if got != tt.want {
+				t.Errorf("extractBearer(%q) = %q, want %q", tt.header, got, tt.want)
+			}
+		})
+	}
+}
+
+// --- wsConn ---
+
+func TestWSConn_EchoRoundTrip(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+
+		wc := newWSConn(conn)
+		buf := make([]byte, 256)
+		n, err := wc.Read(buf)
+		if err != nil {
+			return
+		}
+		wc.Write(buf[:n])
+	}))
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		t.Errorf("handshake status = %d, want 101", resp.StatusCode)
+	}
+
+	msg := []byte("hello tunnel")
+	if err := conn.WriteMessage(websocket.BinaryMessage, msg); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+
+	mt, reply, err := conn.ReadMessage()
+	if err != nil {
+		t.Fatalf("read: %v", err)
+	}
+	if mt != websocket.BinaryMessage {
+		t.Errorf("type = %d, want BinaryMessage(%d)", mt, websocket.BinaryMessage)
+	}
+	if string(reply) != "hello tunnel" {
+		t.Errorf("reply = %q, want %q", reply, "hello tunnel")
+	}
+}
+
+func TestWSConn_MultipleMessages(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+
+		wc := newWSConn(conn)
+		for i := 0; i < 3; i++ {
+			buf := make([]byte, 256)
+			n, err := wc.Read(buf)
+			if err != nil {
+				return
+			}
+			wc.Write(buf[:n])
+		}
+	}))
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	for i, msg := range []string{"one", "two", "three"} {
+		conn.WriteMessage(websocket.BinaryMessage, []byte(msg))
+		_, reply, err := conn.ReadMessage()
+		if err != nil {
+			t.Fatalf("round %d read: %v", i, err)
+		}
+		if string(reply) != msg {
+			t.Errorf("round %d: got %q, want %q", i, reply, msg)
+		}
+	}
+}
+
+func TestWSConn_ImplementsNetConn(t *testing.T) {
+	var _ net.Conn = (*wsConn)(nil)
+}
+
+func TestWSConn_LocalRemoteAddr(t *testing.T) {
+	addrCh := make(chan [2]net.Addr, 1)
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		wc := newWSConn(conn)
+		addrCh <- [2]net.Addr{wc.LocalAddr(), wc.RemoteAddr()}
+		wc.Close()
+	}))
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	select {
+	case addrs := <-addrCh:
+		if addrs[0] == nil {
+			t.Error("LocalAddr should not be nil")
+		}
+		if addrs[1] == nil {
+			t.Error("RemoteAddr should not be nil")
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout waiting for addresses")
+	}
+}
+
+// --- HandleControl ---
+
+func newGinRouter() *gin.Engine {
+	r := gin.New()
+	r.GET("/ws/tai", HandleControl)
+	r.GET("/ws/tai/data/:channel_id", HandleData)
+	return r
+}
+
+func TestHandleControl_NoRegistry(t *testing.T) {
+	registry.SetGlobalForTest(nil)
+	defer setupTestRegistry()
+
+	restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
+	defer restore()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
+	_, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer test-token"},
+	})
+	if err == nil {
+		t.Fatal("expected dial to fail when registry is nil")
+	}
+	if resp != nil && resp.StatusCode != http.StatusServiceUnavailable {
+		t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable)
+	}
+}
+
+func TestHandleControl_NoAuth(t *testing.T) {
+	setupTestRegistry()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
+	_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err == nil {
+		t.Fatal("expected dial to fail without auth")
+	}
+	if resp != nil && resp.StatusCode != http.StatusUnauthorized {
+		t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized)
+	}
+}
+
+func TestHandleControl_AuthFailed(t *testing.T) {
+	setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{}, fmt.Errorf("bad token"))
+	defer restore()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
+	_, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer bad-token"},
+	})
+	if err == nil {
+		t.Fatal("expected dial to fail with bad auth")
+	}
+	if resp != nil && resp.StatusCode != http.StatusUnauthorized {
+		t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized)
+	}
+}
+
+func TestHandleControl_RegisterAndPing(t *testing.T) {
+	reg := setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{
+		ClientID: "tai-001",
+		Subject:  "user-test",
+		Scope:    "tai:tunnel",
+	}, nil)
+	defer restore()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
+	conn, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		t.Errorf("handshake = %d, want 101", resp.StatusCode)
+	}
+
+	regMsg := registerMessage{
+		Type:      "register",
+		TaiID:     "tai-001",
+		MachineID: "m-test",
+		Version:   "2.0",
+		Ports:     map[string]int{"grpc": 9100},
+	}
+	if err := conn.WriteJSON(regMsg); err != nil {
+		t.Fatalf("write register: %v", err)
+	}
+
+	var registered map[string]string
+	if err := conn.ReadJSON(&registered); err != nil {
+		t.Fatalf("read registered: %v", err)
+	}
+	if registered["type"] != "registered" {
+		t.Errorf("response type = %q, want registered", registered["type"])
+	}
+	if registered["tai_id"] != "tai-001" {
+		t.Errorf("response tai_id = %q, want tai-001", registered["tai_id"])
+	}
+
+	snap, ok := reg.Get("tai-001")
+	if !ok {
+		t.Fatal("node not found in registry after register")
+	}
+	if snap.Status != "online" {
+		t.Errorf("Status = %q, want online", snap.Status)
+	}
+	if snap.MachineID != "m-test" {
+		t.Errorf("MachineID = %q, want m-test", snap.MachineID)
+	}
+	if snap.Version != "2.0" {
+		t.Errorf("Version = %q, want 2.0", snap.Version)
+	}
+	if snap.Mode != "tunnel" {
+		t.Errorf("Mode = %q, want tunnel", snap.Mode)
+	}
+	if snap.Auth.ClientID != "tai-001" {
+		t.Errorf("Auth.ClientID = %q, want tai-001", snap.Auth.ClientID)
+	}
+	if snap.Auth.Subject != "user-test" {
+		t.Errorf("Auth.Subject = %q, want user-test", snap.Auth.Subject)
+	}
+	if snap.Ports["grpc"] != 9100 {
+		t.Errorf("Ports[grpc] = %d, want 9100", snap.Ports["grpc"])
+	}
+
+	time.Sleep(10 * time.Millisecond)
+	if err := conn.WriteJSON(map[string]string{"type": "ping"}); err != nil {
+		t.Fatalf("write ping: %v", err)
+	}
+
+	var pong map[string]string
+	if err := conn.ReadJSON(&pong); err != nil {
+		t.Fatalf("read pong: %v", err)
+	}
+	if pong["type"] != "pong" {
+		t.Errorf("pong type = %q, want pong", pong["type"])
+	}
+
+	snap2, _ := reg.Get("tai-001")
+	if !snap2.LastPing.After(snap.LastPing) {
+		t.Error("LastPing should be updated after ping")
+	}
+
+	conn.WriteMessage(websocket.CloseMessage,
+		websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+	time.Sleep(100 * time.Millisecond)
+
+	if _, ok := reg.Get("tai-001"); ok {
+		t.Error("node should be unregistered after connection close")
+	}
+}
+
+func TestHandleControl_BadRegisterType(t *testing.T) {
+	setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
+	defer restore()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	conn.WriteJSON(map[string]string{"type": "not-register"})
+	_, _, readErr := conn.ReadMessage()
+	if readErr == nil {
+		t.Error("expected connection to close for bad register type")
+	}
+}
+
+func TestHandleControl_MissingTaiID(t *testing.T) {
+	setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
+	defer restore()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	conn.WriteJSON(map[string]string{"type": "register"})
+	_, _, readErr := conn.ReadMessage()
+	if readErr == nil {
+		t.Error("expected connection to close for missing tai_id")
+	}
+}
+
+// --- HandleData ---
+
+func TestHandleData_NoAuth(t *testing.T) {
+	setupTestRegistry()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-001"
+	_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err == nil {
+		t.Fatal("expected dial to fail without auth")
+	}
+	if resp != nil && resp.StatusCode != http.StatusUnauthorized {
+		t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized)
+	}
+}
+
+func TestHandleData_AcceptSuccess(t *testing.T) {
+	reg := setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
+	defer restore()
+
+	resultCh := make(chan net.Conn, 1)
+	timer := time.AfterFunc(5*time.Second, func() {})
+	reg.SetPendingForTest("ch-test-123", "tai-001", resultCh, timer)
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-test-123"
+	conn, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		t.Errorf("status = %d, want 101", resp.StatusCode)
+	}
+
+	select {
+	case c := <-resultCh:
+		if c == nil {
+			t.Fatal("expected non-nil conn from resultCh")
+		}
+		c.Close()
+	case <-time.After(2 * time.Second):
+		t.Fatal("timeout waiting for conn on resultCh")
+	}
+}
+
+func TestHandleData_ChannelNotPending(t *testing.T) {
+	setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
+	defer restore()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/nonexistent"
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		return
+	}
+	defer conn.Close()
+
+	_, _, readErr := conn.ReadMessage()
+	if readErr == nil {
+		t.Error("expected connection to close for non-pending channel")
+	}
+}
+
+func TestHandleData_TaiIDMismatch(t *testing.T) {
+	reg := setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{ClientID: "tai-intruder"}, nil)
+	defer restore()
+
+	resultCh := make(chan net.Conn, 1)
+	timer := time.AfterFunc(5*time.Second, func() {})
+	reg.SetPendingForTest("ch-mismatch", "tai-owner", resultCh, timer)
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-mismatch"
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		return
+	}
+	defer conn.Close()
+
+	_, _, readErr := conn.ReadMessage()
+	if readErr == nil {
+		t.Error("expected connection to close for tai_id mismatch")
+	}
+}
+
+// --- Full open-channel flow ---
+
+func TestHandleControl_OpenChannelAndBridge(t *testing.T) {
+	reg := setupTestRegistry()
+	restore := mockAuth(registry.AuthInfo{
+		ClientID: "tai-001",
+		Subject:  "user-test",
+	}, nil)
+	defer restore()
+
+	srv := httptest.NewServer(newGinRouter())
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
+	ctrlConn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		t.Fatalf("dial control: %v", err)
+	}
+	defer ctrlConn.Close()
+
+	ctrlConn.WriteJSON(registerMessage{
+		Type:  "register",
+		TaiID: "tai-001",
+		Ports: map[string]int{"grpc": 9100},
+	})
+	var registered map[string]string
+	if err := ctrlConn.ReadJSON(&registered); err != nil {
+		t.Fatalf("read registered: %v", err)
+	}
+	if registered["type"] != "registered" {
+		t.Fatalf("expected registered, got %v", registered)
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	var requestErr error
+	var channelConn net.Conn
+	go func() {
+		defer wg.Done()
+		_, resultCh, err := reg.RequestChannel("tai-001", 9100)
+		if err != nil {
+			requestErr = err
+			return
+		}
+		channelConn = <-resultCh
+	}()
+
+	time.Sleep(50 * time.Millisecond)
+
+	var openCmd map[string]interface{}
+	if err := ctrlConn.ReadJSON(&openCmd); err != nil {
+		t.Fatalf("read open cmd: %v", err)
+	}
+	if openCmd["type"] != "open" {
+		t.Errorf("open type = %v, want open", openCmd["type"])
+	}
+	channelID, ok := openCmd["channel_id"].(string)
+	if !ok || channelID == "" {
+		t.Fatalf("missing channel_id: %v", openCmd)
+	}
+	if tp, ok := openCmd["target_port"].(float64); !ok || int(tp) != 9100 {
+		t.Errorf("target_port = %v, want 9100", openCmd["target_port"])
+	}
+
+	dataURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/" + channelID
+	dataConn, _, err := websocket.DefaultDialer.Dial(dataURL, http.Header{
+		"Authorization": []string{"Bearer valid-token"},
+	})
+	if err != nil {
+		t.Fatalf("dial data: %v", err)
+	}
+	defer dataConn.Close()
+
+	wg.Wait()
+	if requestErr != nil {
+		t.Fatalf("RequestChannel: %v", requestErr)
+	}
+	if channelConn == nil {
+		t.Fatal("expected non-nil conn from RequestChannel")
+	}
+	defer channelConn.Close()
+
+	payload := []byte("grpc-payload-test")
+	dataConn.WriteMessage(websocket.BinaryMessage, payload)
+
+	buf := make([]byte, 256)
+	n, err := channelConn.Read(buf)
+	if err != nil && err != io.EOF {
+		t.Fatalf("read bridged: %v", err)
+	}
+	if string(buf[:n]) != "grpc-payload-test" {
+		t.Errorf("bridged data = %q, want %q", buf[:n], "grpc-payload-test")
+	}
+}
--- a/tai/vnc/vnc.go
+++ b/tai/vnc/vnc.go
@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strings"

 	"github.com/yaoapp/yao/tai/sandbox"
 )
@ -51,6 +52,29 @@ func (r *remoteVNC) Ping(ctx context.Context, containerID string) error {
 	return nil
 }

+// --- Tunnel implementation ---
+
+type tunnelVNC struct {
+	taiID   string
+	yaoBase string // e.g. "http://yao-host:5099"
+}
+
+// NewTunnel creates a VNC that routes through Yao's reverse proxy
+// for tunnel-connected Tai instances.
+func NewTunnel(taiID, yaoBase string) VNC {
+	return &tunnelVNC{taiID: taiID, yaoBase: strings.TrimRight(yaoBase, "/")}
+}
+
+func (t *tunnelVNC) URL(_ context.Context, containerID string) (string, error) {
+	base := strings.Replace(t.yaoBase, "http://", "ws://", 1)
+	base = strings.Replace(base, "https://", "wss://", 1)
+	return fmt.Sprintf("%s/tai/%s/vnc/%s/ws", base, t.taiID, containerID), nil
+}
+
+func (t *tunnelVNC) Ping(_ context.Context, _ string) error {
+	return nil
+}
+
 // --- Local implementation ---

 type localVNC struct {