Merge pull request #1488 from trheyi/main

Implement tunnel mode support in Tai service

agent/caller/orchestrator.go

@@ -277,26 +277,11 @@ func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Requ
 		)
 	}
 
-	// Notify TUI of A2A call start (use parent requestID so it appears in parent panel)
-	parentRequestID := o.ctx.RequestID()
-	agentContext.SendTUI(agentContext.AgentEventMsg{
-		RequestID: parentRequestID,
-		Event:     agentContext.EventA2AStart,
-		Data:      map[string]interface{}{"target": req.AgentID},
-	})
-
-	// Execute the agent call with the provided context
-	// The agent.Stream method will use the context's Writer for output
 	resp, err := agent.Stream(ctx, req.Messages, ctxOpts)
 	if err != nil {
 		if a2aNode != nil {
 			a2aNode.Fail(err)
 		}
-		agentContext.SendTUI(agentContext.AgentEventMsg{
-			RequestID: parentRequestID,
-			Event:     agentContext.EventA2ADone,
-			Data:      map[string]interface{}{"target": req.AgentID, "error": err.Error()},
-		})
 		return NewResult(req.AgentID, nil, fmt.Errorf("agent call failed: %w", err))
 	}
 
@@ -306,11 +291,6 @@ func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Requ
 			"status":   "completed",
 		})
 	}
-	agentContext.SendTUI(agentContext.AgentEventMsg{
-		RequestID: parentRequestID,
-		Event:     agentContext.EventA2ADone,
-		Data:      map[string]interface{}{"target": req.AgentID},
-	})
 
 	return NewResult(req.AgentID, resp, nil)
 }

agent/context/log.go

@@ -192,11 +192,8 @@ func (l *RequestLogger) processEntry(entry LogEntry) {
 	}
 }
 
-// printDev sends to TUI if available, otherwise prints colored output to stdout
+// printDev prints colored output to stdout in development mode
 func (l *RequestLogger) printDev(entry LogEntry) {
-	if GetTUIProgram() != nil {
-		return
-	}
 	switch entry.Level {
 	case LogLevelTrace:
 		fmt.Printf("%s  → %s%s\n", colorGray, entry.Message, colorReset)
@@ -315,17 +312,6 @@ func (l *RequestLogger) Start() {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID:   l.requestID,
-		ParentID:    l.parentID,
-		AssistantID: l.currentAssistantID(),
-		Event:       EventRequestStart,
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
-
 	fmt.Println()
 	fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("═", 60), colorReset)
 	fmt.Printf("%s  AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset)
@@ -357,21 +343,6 @@ func (l *RequestLogger) End(success bool, err error) {
 		return
 	}
 
-	data := map[string]interface{}{"duration": duration.Round(time.Millisecond)}
-	if err != nil {
-		data["error"] = err.Error()
-	}
-	SendTUI(AgentEventMsg{
-		RequestID:   l.requestID,
-		AssistantID: l.currentAssistantID(),
-		Event:       EventRequestEnd,
-		Data:        data,
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
-
 	fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
 	if success {
 		fmt.Printf("%s  REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset)
@@ -400,15 +371,6 @@ func (l *RequestLogger) Phase(name string) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventPhase,
-		Data:      map[string]interface{}{"name": name, "elapsed": elapsed},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  > %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset)
 }
 
@@ -425,15 +387,6 @@ func (l *RequestLogger) PhaseComplete(name string) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventPhaseDone,
-		Data:      map[string]interface{}{"name": name, "elapsed": elapsed},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  + %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset)
 }
 
@@ -447,15 +400,6 @@ func (l *RequestLogger) PhaseSkip(name, reason string) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventPhaseSkip,
-		Data:      map[string]interface{}{"name": name, "reason": reason},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  - %s (%s)%s\n", colorGray, name, reason, colorReset)
 }
 
@@ -472,19 +416,6 @@ func (l *RequestLogger) LLMStart(connector, model string, messageCount int) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventLLMCall,
-		Data: map[string]interface{}{
-			"connector": connector,
-			"model":     model,
-			"messages":  messageCount,
-		},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset)
 	fmt.Printf("%s    Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset)
 	if model != "" {
@@ -511,19 +442,6 @@ func (l *RequestLogger) LLMComplete(tokens int, hasToolCalls bool) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventLLMDone,
-		Data: map[string]interface{}{
-			"detail": fmt.Sprintf("%s [tokens:%d, %v]", status, tokens, elapsed),
-			"tokens": tokens,
-			"status": status,
-		},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  + LLM Response (%s)%s", colorGreen, status, colorReset)
 	if tokens > 0 {
 		fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset)
@@ -543,15 +461,6 @@ func (l *RequestLogger) ToolStart(toolName string) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventToolCall,
-		Data:      map[string]interface{}{"name": toolName},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  Tool: %s%s\n", colorYellow, toolName, colorReset)
 }
 
@@ -571,15 +480,6 @@ func (l *RequestLogger) ToolComplete(toolName string, success bool) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventToolDone,
-		Data:      map[string]interface{}{"name": toolName, "success": success},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	if success {
 		fmt.Printf("%s    + %s completed%s\n", colorGreen, toolName, colorReset)
 	} else {
@@ -600,15 +500,6 @@ func (l *RequestLogger) HookStart(hookName string) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventHook,
-		Data:      map[string]interface{}{"name": hookName},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset)
 }
 
@@ -624,15 +515,6 @@ func (l *RequestLogger) HookComplete(hookName string) {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventHookDone,
-		Data:      map[string]interface{}{"name": hookName},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s    + %s done%s\n", colorGreen, hookName, colorReset)
 }
 
@@ -644,7 +526,7 @@ func (l *RequestLogger) Cleanup(resource string) {
 
 	kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource)
 
-	if !config.IsDevelopment() || GetTUIProgram() != nil {
+	if !config.IsDevelopment() {
 		return
 	}
 	fmt.Printf("%s    + %s%s\n", colorGray, resource, colorReset)
@@ -658,7 +540,7 @@ func (l *RequestLogger) HistoryLoad(count, maxSize int) {
 
 	kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize)
 
-	if !config.IsDevelopment() || GetTUIProgram() != nil {
+	if !config.IsDevelopment() {
 		return
 	}
 	fmt.Printf("%s    Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset)
@@ -673,7 +555,7 @@ func (l *RequestLogger) HistoryOverlap(overlapCount int) {
 	if overlapCount > 0 {
 		kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount)
 
-		if !config.IsDevelopment() || GetTUIProgram() != nil {
+		if !config.IsDevelopment() {
 			return
 		}
 		fmt.Printf("%s    Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset)
@@ -692,15 +574,6 @@ func (l *RequestLogger) Release() {
 		return
 	}
 
-	SendTUI(AgentEventMsg{
-		RequestID: l.requestID,
-		Event:     EventContextRelease,
-		Data:      map[string]interface{}{"assistant": l.currentAssistantID()},
-	})
-
-	if GetTUIProgram() != nil {
-		return
-	}
 	fmt.Printf("%s  RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.currentAssistantID(), colorReset)
 }
 

agent/context/tui.go

@@ -1,799 +0,0 @@
-package context
-
-import (
-	"fmt"
-	"strings"
-	"sync"
-	"time"
-
-	tea "github.com/charmbracelet/bubbletea"
-	"github.com/charmbracelet/lipgloss"
-)
-
-var (
-	tuiProgram   *tea.Program
-	tuiProgramMu sync.RWMutex
-)
-
-// SetTUIProgram sets the global TUI program (called from start.go after HTTP READY)
-func SetTUIProgram(p *tea.Program) {
-	tuiProgramMu.Lock()
-	tuiProgram = p
-	tuiProgramMu.Unlock()
-}
-
-// GetTUIProgram returns the global TUI program (nil if not in TUI mode)
-func GetTUIProgram() *tea.Program {
-	tuiProgramMu.RLock()
-	defer tuiProgramMu.RUnlock()
-	return tuiProgram
-}
-
-// SendTUI sends a message to the TUI program if available
-func SendTUI(msg tea.Msg) {
-	if p := GetTUIProgram(); p != nil {
-		p.Send(msg)
-	}
-}
-
-// TUILogWriter implements io.Writer to bridge gou DevWriter -> TUI AppLogMsg
-type TUILogWriter struct {
-	Program *tea.Program
-}
-
-func (w *TUILogWriter) Write(p []byte) (n int, err error) {
-	content := strings.TrimRight(string(p), "\n")
-	if content == "" {
-		return len(p), nil
-	}
-	w.Program.Send(AppLogMsg{Content: content})
-	return len(p), nil
-}
-
-// ─── Styles ───────────────────────────────────────────────────────────────────
-
-var (
-	boxRunning = lipgloss.NewStyle().
-			Border(lipgloss.RoundedBorder()).
-			BorderForeground(lipgloss.Color("33")).
-			PaddingLeft(1).PaddingRight(1)
-
-	boxDone = lipgloss.NewStyle().
-		Border(lipgloss.RoundedBorder()).
-		BorderForeground(lipgloss.Color("240")).
-		PaddingLeft(1).PaddingRight(1)
-
-	boxFailed = lipgloss.NewStyle().
-			Border(lipgloss.RoundedBorder()).
-			BorderForeground(lipgloss.Color("31")).
-			PaddingLeft(1).PaddingRight(1)
-
-	boxAppLog = lipgloss.NewStyle().
-			Border(lipgloss.RoundedBorder()).
-			BorderForeground(lipgloss.Color("240")).
-			PaddingLeft(1).PaddingRight(1)
-
-	sRunning = lipgloss.NewStyle().Foreground(lipgloss.Color("33"))
-	sDone    = lipgloss.NewStyle().Foreground(lipgloss.Color("34"))
-	sFailed  = lipgloss.NewStyle().Foreground(lipgloss.Color("31"))
-	sDim     = lipgloss.NewStyle().Foreground(lipgloss.Color("240"))
-	sBold    = lipgloss.NewStyle().Bold(true)
-	sYellow  = lipgloss.NewStyle().Foreground(lipgloss.Color("33"))
-	sRed     = lipgloss.NewStyle().Foreground(lipgloss.Color("31"))
-	sBlue    = lipgloss.NewStyle().Foreground(lipgloss.Color("34"))
-	sMagenta = lipgloss.NewStyle().Foreground(lipgloss.Color("35"))
-	sTree    = lipgloss.NewStyle().Foreground(lipgloss.Color("240"))
-)
-
-// ─── Data ─────────────────────────────────────────────────────────────────────
-
-// RequestPanel represents a single top-level agent request
-type RequestPanel struct {
-	RequestID   string
-	ShortID     string
-	AssistantID string
-	StartTime   time.Time
-	EndTime     time.Time // set when done/failed, freezes elapsed display
-	Status      PanelStatus
-	Nodes       []TreeNode
-	ParentID    string
-	Collapsed   bool
-	viewRow     int // Y offset of the header line (for mouse click)
-}
-
-// TreeNode represents a step within a request panel
-type TreeNode struct {
-	Kind      NodeKind
-	Label     string
-	Status    NodeStatus
-	Detail    string
-	Children  []*TreeNode
-	StartTime time.Time
-	EndTime   time.Time
-	Collapsed bool
-}
-
-// AgentTUIModel is the bubbletea Model for agent request visualization
-type AgentTUIModel struct {
-	panels       []*RequestPanel
-	panelIndex   map[string]int // requestID -> index in panels (first registration wins)
-	appLogs      []AppLogEntry
-	appLogExpand bool
-	appLogRow    int // Y offset of app log header
-	cursor       int
-	width        int
-	height       int
-	scrollOffset int
-	autoFollow   bool // auto-scroll to bottom when new content arrives
-	mouseOn      bool
-	quitting     bool
-}
-
-// NewAgentTUIModel creates a new TUI model
-func NewAgentTUIModel() AgentTUIModel {
-	return AgentTUIModel{
-		panels:     []*RequestPanel{},
-		panelIndex: map[string]int{},
-		appLogs:    []AppLogEntry{},
-		width:      80,
-		height:     24,
-		autoFollow: true,
-	}
-}
-
-func (m AgentTUIModel) Init() tea.Cmd {
-	return tickCmd()
-}
-
-func tickCmd() tea.Cmd {
-	return tea.Tick(500*time.Millisecond, func(t time.Time) tea.Msg {
-		return TickMsg(t)
-	})
-}
-
-// ─── Update ───────────────────────────────────────────────────────────────────
-
-func (m AgentTUIModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
-	switch msg := msg.(type) {
-	case tea.WindowSizeMsg:
-		m.width = msg.Width
-		m.height = msg.Height
-		return m, nil
-
-	case tea.KeyMsg:
-		return m.handleKey(msg)
-
-	case tea.MouseMsg:
-		return m.handleMouse(msg)
-
-	case AgentEventMsg:
-		return m.handleAgentEvent(msg), nil
-
-	case AppLogMsg:
-		m.appLogs = append(m.appLogs, AppLogEntry{
-			Content: msg.Content,
-			Time:    time.Now(),
-		})
-		return m, nil
-
-	case TickMsg:
-		return m, tickCmd()
-	}
-	return m, nil
-}
-
-func (m AgentTUIModel) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
-	topPanels := m.topLevelPanels()
-	total := len(topPanels) + 1 // +1 for app log
-	viewH := m.viewHeight()
-
-	switch msg.String() {
-	case "q", "ctrl+c":
-		m.quitting = true
-		return m, tea.Quit
-
-	// Scrolling
-	case "j", "down":
-		m.scrollOffset++
-		m.autoFollow = false
-	case "k", "up":
-		if m.scrollOffset > 0 {
-			m.scrollOffset--
-		}
-		m.autoFollow = false
-	case "pgdown", "ctrl+d":
-		m.scrollOffset += viewH / 2
-		m.autoFollow = false
-	case "pgup", "ctrl+u":
-		m.scrollOffset -= viewH / 2
-		if m.scrollOffset < 0 {
-			m.scrollOffset = 0
-		}
-		m.autoFollow = false
-	case "G", "end":
-		m.autoFollow = true
-	case "g", "home":
-		m.scrollOffset = 0
-		m.autoFollow = false
-
-	// Cursor navigation for panel selection (wraps around)
-	case "tab":
-		m.cursor = (m.cursor + 1) % total
-		m.scrollToCursor(topPanels)
-	case "shift+tab":
-		m.cursor = (m.cursor - 1 + total) % total
-		m.scrollToCursor(topPanels)
-
-	case "enter", " ":
-		if m.cursor < len(topPanels) {
-			topPanels[m.cursor].Collapsed = !topPanels[m.cursor].Collapsed
-		} else {
-			m.appLogExpand = !m.appLogExpand
-		}
-	case "c":
-		m.appLogExpand = !m.appLogExpand
-	case "a":
-		for _, p := range m.panels {
-			p.Collapsed = false
-		}
-		m.appLogExpand = true
-	case "A":
-		for _, p := range m.panels {
-			p.Collapsed = true
-		}
-		m.appLogExpand = false
-	case "m":
-		m.mouseOn = !m.mouseOn
-		if m.mouseOn {
-			return m, tea.EnableMouseCellMotion
-		}
-		return m, tea.DisableMouse
-	}
-	return m, nil
-}
-
-func (m AgentTUIModel) viewHeight() int {
-	h := m.height - 2 // reserve for status bar
-	if h < 4 {
-		h = 4
-	}
-	return h
-}
-
-func (m *AgentTUIModel) scrollToCursor(topPanels []*RequestPanel) {
-	targetRow := 0
-	if m.cursor < len(topPanels) {
-		targetRow = topPanels[m.cursor].viewRow
-	} else {
-		targetRow = m.appLogRow
-	}
-	viewH := m.viewHeight()
-	if targetRow < m.scrollOffset {
-		m.scrollOffset = targetRow
-	} else if targetRow >= m.scrollOffset+viewH {
-		m.scrollOffset = targetRow - viewH + 3
-	}
-}
-
-func (m AgentTUIModel) handleMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
-	switch {
-	case msg.Button == tea.MouseButtonWheelUp:
-		m.scrollOffset -= 3
-		if m.scrollOffset < 0 {
-			m.scrollOffset = 0
-		}
-		m.autoFollow = false
-		return m, nil
-	case msg.Button == tea.MouseButtonWheelDown:
-		m.scrollOffset += 3
-		m.autoFollow = false
-		return m, nil
-	}
-
-	if msg.Button != tea.MouseButtonLeft || msg.Action != tea.MouseActionRelease {
-		return m, nil
-	}
-	y := msg.Y + m.scrollOffset
-
-	// Check app log header
-	if y == m.appLogRow {
-		m.appLogExpand = !m.appLogExpand
-		return m, nil
-	}
-
-	// Check panel headers
-	for _, p := range m.panels {
-		if p.ParentID != "" {
-			continue
-		}
-		if y == p.viewRow {
-			p.Collapsed = !p.Collapsed
-			return m, nil
-		}
-	}
-	return m, nil
-}
-
-// ─── Agent Events ─────────────────────────────────────────────────────────────
-
-func (m *AgentTUIModel) handleAgentEvent(msg AgentEventMsg) tea.Model {
-	switch msg.Event {
-	case EventRequestStart:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			// Delegate sub-call: same requestID, different assistantID.
-			// Add as a tree node inside the existing panel instead of creating a new one.
-			p := m.panels[idx]
-			p.Nodes = append(p.Nodes, TreeNode{
-				Kind:      NodeA2A,
-				Label:     msg.AssistantID,
-				Status:    NodeRunning,
-				StartTime: time.Now(),
-			})
-			return m
-		}
-
-		panel := &RequestPanel{
-			RequestID:   msg.RequestID,
-			ShortID:     shortID(msg.RequestID),
-			AssistantID: msg.AssistantID,
-			StartTime:   time.Now(),
-			Status:      PanelRunning,
-			ParentID:    msg.ParentID,
-		}
-		m.panelIndex[msg.RequestID] = len(m.panels)
-		m.panels = append(m.panels, panel)
-
-	case EventRequestEnd:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			p := m.panels[idx]
-
-			// Only mark panel done if the ending assistantID matches the panel's original assistantID
-			// (delegate sub-calls End with a different assistantID, they update their tree node instead)
-			if msg.AssistantID == p.AssistantID || msg.AssistantID == "" {
-				if errVal, has := msg.Data["error"]; has && errVal != nil {
-					p.Status = PanelFailed
-				} else {
-					p.Status = PanelSuccess
-				}
-				p.EndTime = time.Now()
-				p.Collapsed = true
-
-				// Finalize any still-running child nodes (e.g. hook interrupted mid-execution)
-				finalStatus := NodeDone
-				if p.Status == PanelFailed {
-					finalStatus = NodeFailed
-				}
-				for i := range p.Nodes {
-					if p.Nodes[i].Status == NodeRunning {
-						p.Nodes[i].Status = finalStatus
-						p.Nodes[i].EndTime = p.EndTime
-					}
-				}
-			} else {
-				// Delegate sub-call finished: mark its tree node as done
-				for i := len(p.Nodes) - 1; i >= 0; i-- {
-					if p.Nodes[i].Kind == NodeA2A && p.Nodes[i].Label == msg.AssistantID && p.Nodes[i].Status == NodeRunning {
-						p.Nodes[i].Status = NodeDone
-						p.Nodes[i].EndTime = time.Now()
-						break
-					}
-				}
-			}
-		}
-
-	case EventLLMCall:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{
-				Kind: NodeLLM, Label: "LLM", Status: NodeRunning, StartTime: time.Now(),
-			})
-		}
-
-	case EventLLMDone:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			p := m.panels[idx]
-			for i := len(p.Nodes) - 1; i >= 0; i-- {
-				if p.Nodes[i].Kind == NodeLLM && p.Nodes[i].Status == NodeRunning {
-					p.Nodes[i].Status = NodeDone
-					p.Nodes[i].EndTime = time.Now()
-					if d, has := msg.Data["detail"]; has {
-						p.Nodes[i].Detail = fmt.Sprintf("%v", d)
-					}
-					break
-				}
-			}
-		}
-
-	case EventToolCall:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			name := dataStr(msg.Data, "name")
-			m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{
-				Kind: NodeTool, Label: name, Status: NodeRunning, StartTime: time.Now(),
-			})
-		}
-
-	case EventToolDone:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			name := dataStr(msg.Data, "name")
-			p := m.panels[idx]
-			for i := len(p.Nodes) - 1; i >= 0; i-- {
-				if p.Nodes[i].Kind == NodeTool && p.Nodes[i].Label == name && p.Nodes[i].Status == NodeRunning {
-					p.Nodes[i].Status = NodeDone
-					p.Nodes[i].EndTime = time.Now()
-					break
-				}
-			}
-		}
-
-	case EventHook:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			name := dataStr(msg.Data, "name")
-			m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{
-				Kind: NodeHook, Label: name, Status: NodeRunning, StartTime: time.Now(),
-			})
-		}
-
-	case EventHookDone:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			p := m.panels[idx]
-			for i := len(p.Nodes) - 1; i >= 0; i-- {
-				if p.Nodes[i].Kind == NodeHook && p.Nodes[i].Status == NodeRunning {
-					p.Nodes[i].Status = NodeDone
-					p.Nodes[i].EndTime = time.Now()
-					break
-				}
-			}
-		}
-
-	case EventA2AStart:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			target := dataStr(msg.Data, "target")
-			m.panels[idx].Nodes = append(m.panels[idx].Nodes, TreeNode{
-				Kind: NodeA2A, Label: target, Status: NodeRunning, StartTime: time.Now(),
-			})
-		}
-
-	case EventA2ADone:
-		if idx, ok := m.panelIndex[msg.RequestID]; ok {
-			target := dataStr(msg.Data, "target")
-			p := m.panels[idx]
-			for i := len(p.Nodes) - 1; i >= 0; i-- {
-				if p.Nodes[i].Kind == NodeA2A && p.Nodes[i].Status == NodeRunning && (target == "" || p.Nodes[i].Label == target) {
-					p.Nodes[i].Status = NodeDone
-					p.Nodes[i].EndTime = time.Now()
-					break
-				}
-			}
-		}
-	}
-	return m
-}
-
-// ─── View ─────────────────────────────────────────────────────────────────────
-
-func (m AgentTUIModel) View() string {
-	if m.quitting {
-		return ""
-	}
-
-	boxW := m.width - 2
-	if boxW < 40 {
-		boxW = 40
-	}
-
-	// Render full content
-	var sb strings.Builder
-	row := 0
-	topIdx := 0
-
-	for _, panel := range m.panels {
-		if panel.ParentID != "" {
-			continue
-		}
-		selected := (topIdx == m.cursor)
-		rendered := m.renderPanelBox(panel, boxW, selected, &row)
-		sb.WriteString(rendered)
-		sb.WriteString("\n")
-		row++
-		topIdx++
-	}
-
-	// App Log
-	m.appLogRow = row
-	sb.WriteString(m.renderAppLogBox(boxW, topIdx == m.cursor, &row))
-
-	fullContent := sb.String()
-	lines := strings.Split(fullContent, "\n")
-	totalLines := len(lines)
-	viewH := m.viewHeight()
-
-	// Auto-follow: snap to bottom
-	if m.autoFollow {
-		m.scrollOffset = totalLines - viewH
-	}
-
-	// Clamp scroll offset
-	maxScroll := totalLines - viewH
-	if maxScroll < 0 {
-		maxScroll = 0
-	}
-	if m.scrollOffset > maxScroll {
-		m.scrollOffset = maxScroll
-	}
-	if m.scrollOffset < 0 {
-		m.scrollOffset = 0
-	}
-
-	// Slice visible lines
-	end := m.scrollOffset + viewH
-	if end > totalLines {
-		end = totalLines
-	}
-	visible := lines[m.scrollOffset:end]
-
-	// Build output
-	var out strings.Builder
-	out.WriteString(strings.Join(visible, "\n"))
-
-	// Status bar with scroll indicator
-	mouseLabel := "off"
-	if m.mouseOn {
-		mouseLabel = "on"
-	}
-	scrollInfo := ""
-	if totalLines > viewH {
-		pct := 100
-		if maxScroll > 0 {
-			pct = m.scrollOffset * 100 / maxScroll
-		}
-		scrollInfo = fmt.Sprintf(" [%d%%]", pct)
-	}
-	followLabel := ""
-	if m.autoFollow {
-		followLabel = " AUTO"
-	}
-	hint := sDim.Render(fmt.Sprintf("  j/k:scroll  tab:select  space:toggle  a/A:all  G:bottom  g:top  m:mouse(%s)%s%s  q:quit",
-		mouseLabel, scrollInfo, followLabel))
-	out.WriteString("\n" + hint)
-
-	return out.String()
-}
-
-func (m AgentTUIModel) renderPanelBox(panel *RequestPanel, boxW int, selected bool, row *int) string {
-	// Record header row for mouse
-	panel.viewRow = *row
-
-	elapsed := m.panelElapsed(panel)
-	icon, statusText, style := panelStatusDisplay(panel.Status, elapsed)
-
-	// Title line
-	collapser := "▾"
-	if panel.Collapsed {
-		collapser = "▸"
-	}
-	cursor := " "
-	if selected {
-		cursor = "›"
-	}
-	title := fmt.Sprintf("%s %s %s %s  %s",
-		sDim.Render(cursor),
-		sDim.Render(collapser),
-		sBold.Render(panel.ShortID),
-		panel.AssistantID,
-		style.Render(icon+" "+statusText),
-	)
-
-	if panel.Collapsed {
-		box := boxForStatus(panel.Status).Width(boxW)
-		result := box.Render(title)
-		*row += strings.Count(result, "\n") + 1
-		return result
-	}
-
-	// Build body
-	var body strings.Builder
-	body.WriteString(title + "\n")
-
-	for _, node := range panel.Nodes {
-		body.WriteString(m.renderTreeNode(node, "  ", false, panel))
-	}
-
-	// Render fork children (different requestID, parentID matches)
-	children := m.childPanels(panel.RequestID)
-	for i, child := range children {
-		isLast := (i == len(children)-1)
-		body.WriteString(m.renderChildSummary(child, "  ", isLast))
-	}
-
-	box := boxForStatus(panel.Status).Width(boxW)
-	result := box.Render(body.String())
-	*row += strings.Count(result, "\n") + 1
-	return result
-}
-
-func (m AgentTUIModel) renderTreeNode(node TreeNode, prefix string, isChild bool, panel *RequestPanel) string {
-	panelEnded := panel != nil && panel.Status != PanelRunning
-	displayNode := node
-	if panelEnded && displayNode.Status == NodeRunning {
-		displayNode.Status = NodeFailed
-	}
-	icon, statusText := nodeStatusDisplay(displayNode)
-	elapsed := m.nodeElapsed(node, panelEnded, panel.EndTime)
-
-	label := ""
-	switch node.Kind {
-	case NodeHook:
-		label = sMagenta.Render("Hook: "+node.Label) + " " + statusText
-	case NodeLLM:
-		detail := ""
-		if node.Detail != "" {
-			detail = " " + sDim.Render("["+node.Detail+"]")
-		}
-		label = sBlue.Render("LLM") + " " + statusText + detail
-	case NodeTool:
-		label = sTree.Render("├ ") + sYellow.Render(node.Label) + " " + statusText
-	case NodeA2A:
-		label = sTree.Render("⤷ ") + sBold.Render(node.Label) + " " + statusText
-	case NodePhase:
-		label = node.Label + " " + statusText
-	default:
-		label = node.Label + " " + statusText
-	}
-
-	_ = icon
-	line := prefix + label
-	if elapsed != "" {
-		line += " " + sDim.Render(elapsed)
-	}
-	return line + "\n"
-}
-
-func (m AgentTUIModel) renderChildSummary(panel *RequestPanel, prefix string, isLast bool) string {
-	elapsed := m.panelElapsed(panel)
-	icon, statusText, style := panelStatusDisplay(panel.Status, elapsed)
-
-	branch := sTree.Render("├─ ")
-	if isLast {
-		branch = sTree.Render("└─ ")
-	}
-	return fmt.Sprintf("%s%s%s %s %s\n",
-		prefix, branch,
-		sBold.Render(panel.ShortID+" "+panel.AssistantID),
-		style.Render(icon+" "+statusText),
-		sDim.Render(elapsed),
-	)
-}
-
-func (m AgentTUIModel) renderAppLogBox(boxW int, selected bool, row *int) string {
-	cursor := " "
-	if selected {
-		cursor = "›"
-	}
-	collapser := "▸"
-	if m.appLogExpand {
-		collapser = "▾"
-	}
-
-	count := len(m.appLogs)
-	title := fmt.Sprintf("%s %s %s (%d)",
-		sDim.Render(cursor),
-		sDim.Render(collapser),
-		sBold.Render("App Output"),
-		count,
-	)
-
-	if !m.appLogExpand || count == 0 {
-		result := boxAppLog.Width(boxW).Render(title)
-		*row += strings.Count(result, "\n") + 1
-		return result
-	}
-
-	var body strings.Builder
-	body.WriteString(title + "\n")
-
-	start := 0
-	if count > 50 {
-		start = count - 50
-	}
-	for _, entry := range m.appLogs[start:] {
-		body.WriteString("  " + entry.Content + "\n")
-	}
-
-	result := boxAppLog.Width(boxW).Render(body.String())
-	*row += strings.Count(result, "\n") + 1
-	return result
-}
-
-// ─── Helpers ──────────────────────────────────────────────────────────────────
-
-func (m AgentTUIModel) topLevelPanels() []*RequestPanel {
-	var result []*RequestPanel
-	for _, p := range m.panels {
-		if p.ParentID == "" {
-			result = append(result, p)
-		}
-	}
-	return result
-}
-
-func (m AgentTUIModel) childPanels(parentRequestID string) []*RequestPanel {
-	var result []*RequestPanel
-	for _, p := range m.panels {
-		if p.ParentID == parentRequestID {
-			result = append(result, p)
-		}
-	}
-	return result
-}
-
-func (m AgentTUIModel) panelElapsed(p *RequestPanel) string {
-	if p.Status != PanelRunning && !p.EndTime.IsZero() {
-		return fmtDuration(p.EndTime.Sub(p.StartTime))
-	}
-	return fmtDuration(time.Since(p.StartTime))
-}
-
-func (m AgentTUIModel) nodeElapsed(n TreeNode, panelEnded bool, panelEndTime time.Time) string {
-	if n.Status == NodeDone || n.Status == NodeFailed {
-		if !n.EndTime.IsZero() {
-			return fmtDuration(n.EndTime.Sub(n.StartTime))
-		}
-	}
-	if n.Status == NodeRunning {
-		if panelEnded && !panelEndTime.IsZero() {
-			return fmtDuration(panelEndTime.Sub(n.StartTime))
-		}
-		return fmtDuration(time.Since(n.StartTime))
-	}
-	return ""
-}
-
-func panelStatusDisplay(status PanelStatus, elapsed string) (icon string, text string, style lipgloss.Style) {
-	switch status {
-	case PanelRunning:
-		return "⟳", "running " + elapsed, sRunning
-	case PanelSuccess:
-		return "✓", "done " + elapsed, sDone
-	case PanelFailed:
-		return "✗", "failed " + elapsed, sFailed
-	}
-	return "", "", sDim
-}
-
-func nodeStatusDisplay(n TreeNode) (icon string, text string) {
-	switch n.Status {
-	case NodePending:
-		return "…", sDim.Render("…")
-	case NodeRunning:
-		return "⟳", sRunning.Render("⟳")
-	case NodeDone:
-		return "✓", sDone.Render("✓")
-	case NodeFailed:
-		return "✗", sFailed.Render("✗")
-	}
-	return "", ""
-}
-
-func boxForStatus(status PanelStatus) lipgloss.Style {
-	switch status {
-	case PanelRunning:
-		return boxRunning
-	case PanelFailed:
-		return boxFailed
-	default:
-		return boxDone
-	}
-}
-
-func dataStr(data map[string]interface{}, key string) string {
-	if v, ok := data[key]; ok {
-		return fmt.Sprintf("%v", v)
-	}
-	return ""
-}
-
-func fmtDuration(d time.Duration) string {
-	if d < time.Second {
-		return fmt.Sprintf("%dms", d.Milliseconds())
-	}
-	return fmt.Sprintf("%.1fs", d.Seconds())
-}

agent/context/tui_msg.go

@@ -1,90 +0,0 @@
-package context
-
-import "time"
-
-// EventType represents the type of agent lifecycle event
-type EventType int
-
-const (
-	EventRequestStart EventType = iota
-	EventPhase
-	EventPhaseDone
-	EventPhaseSkip
-	EventLLMCall
-	EventLLMDone
-	EventToolCall
-	EventToolDone
-	EventHook
-	EventHookDone
-	EventA2AStart
-	EventA2ADone
-	EventRequestEnd
-	EventContextFork
-	EventContextRelease
-)
-
-// AgentEventMsg is sent from RequestLogger to the TUI Program
-type AgentEventMsg struct {
-	RequestID   string
-	ParentID    string
-	AssistantID string
-	Event       EventType
-	Data        map[string]interface{}
-}
-
-// AppLogLevel represents the severity of application-side output
-type AppLogLevel int
-
-const (
-	AppLogLevelLog AppLogLevel = iota
-	AppLogLevelInfo
-	AppLogLevelWarn
-	AppLogLevelError
-	AppLogLevelException
-)
-
-// AppLogMsg is sent from the DevWriter (gou layer) to the TUI Program
-type AppLogMsg struct {
-	Level   AppLogLevel
-	Content string
-}
-
-// AppLogEntry stores a single application output entry
-type AppLogEntry struct {
-	Level   AppLogLevel
-	Content string
-	Time    time.Time
-}
-
-// PanelStatus represents the lifecycle state of a request panel
-type PanelStatus int
-
-const (
-	PanelRunning PanelStatus = iota
-	PanelSuccess
-	PanelFailed
-)
-
-// NodeKind represents the type of a tree node within a request panel
-type NodeKind int
-
-const (
-	NodePhase NodeKind = iota
-	NodeLLM
-	NodeTool
-	NodeHook
-	NodeA2A
-)
-
-// NodeStatus represents the state of a tree node
-type NodeStatus int
-
-const (
-	NodePending NodeStatus = iota
-	NodeRunning
-	NodeDone
-	NodeFailed
-)
-
-// TickMsg triggers periodic UI refresh for elapsed time display
-type TickMsg time.Time

cmd/start.go

@@ -8,14 +8,11 @@ import (
 	"strings"
 	"syscall"
 
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/fatih/color"
-	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
 	"github.com/yaoapp/gou/api"
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/fs"
-	"github.com/yaoapp/gou/helper"
 	"github.com/yaoapp/gou/mcp"
 	"github.com/yaoapp/gou/plugin"
 	"github.com/yaoapp/gou/schedule"
@@ -23,9 +20,7 @@ import (
 	"github.com/yaoapp/gou/store"
 	"github.com/yaoapp/gou/task"
 	"github.com/yaoapp/gou/websocket"
-	"github.com/yaoapp/kun/exception"
 	"github.com/yaoapp/kun/log"
-	agentcontext "github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/engine"
 	yaogrpc "github.com/yaoapp/yao/grpc"
@@ -35,12 +30,12 @@ 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"
 )
 
 var startDebug = false
 var startDisableWatching = false
-var startTUI = false
 
 var startCmd = &cobra.Command{
 	Use:   "start",
@@ -231,6 +226,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() {
@@ -288,7 +287,6 @@ var startCmd = &cobra.Command{
 				case http.READY:
 					fmt.Println(color.GreenString(L("Server is up and running...")))
 					fmt.Println(color.GreenString("Ctrl+C to stop"))
-					initAgentTUI()
 					break
 
 				case http.CLOSED:
@@ -649,38 +647,7 @@ func colorMehtod(method string) string {
 	}
 }
 
-// initAgentTUI initializes the TUI for agent request visualization in dev mode.
-// Must be called after HTTP READY to avoid interfering with startup messages.
-func initAgentTUI() {
-	if !config.IsDevelopment() {
-		return
-	}
-
-	if !startTUI && os.Getenv("YAO_TUI") != "on" {
-		return
-	}
-
-	if !isatty.IsTerminal(os.Stdout.Fd()) {
-		return
-	}
-
-	model := agentcontext.NewAgentTUIModel()
-	p := tea.NewProgram(model, tea.WithoutSignalHandler())
-
-	agentcontext.SetTUIProgram(p)
-	tuiWriter := &agentcontext.TUILogWriter{Program: p}
-	helper.SetDevWriter(tuiWriter)
-	exception.SetWriter(tuiWriter)
-
-	go func() {
-		if _, err := p.Run(); err != nil {
-			log.Error("TUI error: %s", err.Error())
-		}
-	}()
-}
-
 func init() {
 	startCmd.PersistentFlags().BoolVarP(&startDebug, "debug", "", false, L("Development mode"))
 	startCmd.PersistentFlags().BoolVarP(&startDisableWatching, "disable-watching", "", false, L("Disable watching"))
-	startCmd.PersistentFlags().BoolVarP(&startTUI, "tui", "", false, L("Enable TUI for agent request visualization"))
 }

go.mod

@@ -10,8 +10,6 @@ require (
 	github.com/blang/semver v3.5.1+incompatible
 	github.com/bwmarrin/discordgo v0.29.0
 	github.com/caarlos0/env/v6 v6.10.1
-	github.com/charmbracelet/bubbletea v1.3.10
-	github.com/charmbracelet/lipgloss v1.1.0
 	github.com/dchest/captcha v1.1.0
 	github.com/docker/docker v28.5.2+incompatible
 	github.com/docker/go-connections v0.5.0
@@ -34,7 +32,6 @@ require (
 	github.com/kaptinlin/jsonschema v0.6.1
 	github.com/larksuite/oapi-sdk-go/v3 v3.5.3
 	github.com/matoous/go-nanoid/v2 v2.1.0
-	github.com/mattn/go-isatty v0.0.20
 	github.com/mozillazg/go-pinyin v0.20.0
 	github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
 	github.com/pierrec/lz4/v4 v4.1.25
@@ -79,17 +76,12 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.15 // indirect
 	github.com/aws/smithy-go v1.22.3 // indirect
-	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/blang/semver/v4 v4.0.0 // indirect
 	github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect
 	github.com/bytedance/sonic v1.13.2 // indirect
 	github.com/bytedance/sonic/loader v0.2.4 // indirect
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
-	github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect
-	github.com/charmbracelet/x/ansi v0.10.1 // indirect
-	github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect
-	github.com/charmbracelet/x/term v0.2.1 // indirect
 	github.com/cloudwego/base64x v0.1.5 // indirect
 	github.com/coder/websocket v1.8.14 // indirect
 	github.com/containerd/errdefs v1.0.0 // indirect
@@ -101,7 +93,6 @@ require (
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 // indirect
 	github.com/emicklei/go-restful/v3 v3.13.0 // indirect
-	github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/fxamacker/cbor/v2 v2.9.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.9 // indirect
@@ -153,11 +144,10 @@ require (
 	github.com/klauspost/cpuid/v2 v2.2.10 // indirect
 	github.com/leodido/go-urn v1.4.0 // indirect
 	github.com/lib/pq v1.10.9 // indirect
-	github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mark3labs/mcp-go v0.32.0 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect
-	github.com/mattn/go-localereader v0.0.1 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-runewidth v0.0.16 // indirect
 	github.com/mattn/go-sqlite3 v1.14.28 // indirect
 	github.com/miekg/dns v1.1.66 // indirect
@@ -167,9 +157,6 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
 	github.com/montanaflynn/stats v0.7.1 // indirect
-	github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
-	github.com/muesli/cancelreader v0.2.2 // indirect
-	github.com/muesli/termenv v0.16.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f // indirect
 	github.com/neo4j/neo4j-go-driver/v5 v5.28.1 // indirect
@@ -211,7 +198,6 @@ require (
 	github.com/xdg-go/pbkdf2 v1.0.0 // indirect
 	github.com/xdg-go/scram v1.1.2 // indirect
 	github.com/xdg-go/stringprep v1.0.4 // indirect
-	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
 	github.com/xuri/efp v0.0.1 // indirect
 	github.com/xuri/nfp v0.0.1 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect

go.sum

@@ -43,8 +43,6 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3 h1:BRXS0U76Z8wfF+bnkilA2QwpIch6U
 github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3/go.mod h1:bNXKFFyaiVvWuR6O16h/I1724+aXe/tAkA9/QS01t5k=
 github.com/aws/smithy-go v1.22.3 h1:Z//5NuZCSW6R4PhQ93hShNbyBbn8BWCmCVCt+Q8Io5k=
 github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=
-github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
-github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
 github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ=
 github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk=
 github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM=
@@ -70,18 +68,6 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3
 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
-github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
-github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4=
-github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs=
-github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk=
-github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
-github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
-github.com/charmbracelet/x/ansi v0.10.1 h1:rL3Koar5XvX0pHGfovN03f5cxLbCF2YvLeyz7D2jVDQ=
-github.com/charmbracelet/x/ansi v0.10.1/go.mod h1:3RQDQ6lDnROptfpWuUVIUG64bD2g2BgntdxH0Ya5TeE=
-github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8=
-github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs=
-github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ=
-github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg=
 github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
 github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
 github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
@@ -123,8 +109,6 @@ github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6/go.mod h1:iL2twTe
 github.com/emersion/go-textwrapper v0.0.0-20200911093747-65d896831594/go.mod h1:aqO8z8wPrjkscevZJFVE1wXJrLpC5LtJG7fqLOsPb2U=
 github.com/emicklei/go-restful/v3 v3.13.0 h1:C4Bl2xDndpU6nJ4bc1jXd+uTmYPVUwkD6bFY/oTyCes=
 github.com/emicklei/go-restful/v3 v3.13.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc=
-github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
-github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
 github.com/evanw/esbuild v0.25.4 h1:k1bTSim+usBG27w7BfOCorhgx3tO+6bAfMj5pR+6SKg=
 github.com/evanw/esbuild v0.25.4/go.mod h1:D2vIQZqV/vIf/VRHtViaUtViZmG7o+kKmlBfVQuRi48=
 github.com/expr-lang/expr v1.17.7 h1:Q0xY/e/2aCIp8g9s/LGvMDCC5PxYlvHgDZRQ4y16JX8=
@@ -304,8 +288,6 @@ github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
 github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
 github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
 github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
-github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
-github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
 github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
 github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
 github.com/mark3labs/mcp-go v0.32.0 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8=
@@ -320,8 +302,6 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky
 github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
-github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
-github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
 github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
 github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
@@ -351,12 +331,6 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
 github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
 github.com/mozillazg/go-pinyin v0.20.0 h1:BtR3DsxpApHfKReaPO1fCqF4pThRwH9uwvXzm+GnMFQ=
 github.com/mozillazg/go-pinyin v0.20.0/go.mod h1:iR4EnMMRXkfpFVV5FMi4FNB6wGq9NV6uDWbUuPhP4Yc=
-github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
-github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
-github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
-github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
-github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc=
-github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
 github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f h1:y5//uYreIhSUg3J1GEMiLbxo1LJaP8RfCpH6pymGZus=
@@ -488,8 +462,6 @@ github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY=
 github.com/xdg-go/scram v1.1.2/go.mod h1:RT/sEzTbU5y00aCK8UOx6R7YryM0iF1N2MOmC3kKLN4=
 github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8=
 github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM=
-github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
-github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
 github.com/xuri/efp v0.0.1 h1:fws5Rv3myXyYni8uwj2qKjVaRP30PdjeYe2Y6FDsCL8=
 github.com/xuri/efp v0.0.1/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI=
 github.com/xuri/excelize/v2 v2.9.1 h1:VdSGk+rraGmgLHGFaGG9/9IWu1nj4ufjJ7uwMDtj8Qw=
@@ -612,7 +584,6 @@ golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

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()

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)
 
 }

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

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) {

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 {

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
+}

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)
+	}
+}

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}
+}

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
 

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
+}

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) }

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")
+	}
+}

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 {