opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #1488 from trheyi/main

Browse source 
Implement tunnel mode support in Tai service
This commit is contained in:
Max 2026-03-06 20:57:30 +08:00 • committed by GitHub
commit 61548fff10
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 2193 additions and 1127 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
agent/caller/orchestrator.go
View file

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

135
agent/context/log.go
View file

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

799
agent/context/tui.go
View file

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

90
agent/context/tui_msg.go
View file

@ -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

43
cmd/start.go
View file

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

16
go.mod
View file

@ -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

29
go.sum
View file

@ -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=

9
grpc/grpc.go
View file

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

7
openapi/openapi.go
View file

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

32
sandbox/v2/grpc.go
View file

@ -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

10
tai/proxy/connect.go
View file

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

22
tai/proxy/proxy.go
View file

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

401
tai/registry/registry.go Normal file
View file

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

499
tai/registry/registry_test.go Normal file
View file

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

31
tai/registry/testing.go Normal file
View file

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

110
tai/tai.go
View file

@ -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

172
tai/tunnel/proxy.go Normal file
View file

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

267
tai/tunnel/server.go Normal file
View file

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

604
tai/tunnel/server_test.go Normal file
View file

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

24
tai/vnc/vnc.go
View file

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