diff --git a/agent/caller/orchestrator.go b/agent/caller/orchestrator.go index fdf78655..f4b1ca54 100644 --- a/agent/caller/orchestrator.go +++ b/agent/caller/orchestrator.go @@ -277,26 +277,11 @@ func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Requ ) } - // Notify TUI of A2A call start (use parent requestID so it appears in parent panel) - parentRequestID := o.ctx.RequestID() - agentContext.SendTUI(agentContext.AgentEventMsg{ - RequestID: parentRequestID, - Event: agentContext.EventA2AStart, - Data: map[string]interface{}{"target": req.AgentID}, - }) - - // Execute the agent call with the provided context - // The agent.Stream method will use the context's Writer for output resp, err := agent.Stream(ctx, req.Messages, ctxOpts) if err != nil { if a2aNode != nil { a2aNode.Fail(err) } - agentContext.SendTUI(agentContext.AgentEventMsg{ - RequestID: parentRequestID, - Event: agentContext.EventA2ADone, - Data: map[string]interface{}{"target": req.AgentID, "error": err.Error()}, - }) return NewResult(req.AgentID, nil, fmt.Errorf("agent call failed: %w", err)) } @@ -306,11 +291,6 @@ func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Requ "status": "completed", }) } - agentContext.SendTUI(agentContext.AgentEventMsg{ - RequestID: parentRequestID, - Event: agentContext.EventA2ADone, - Data: map[string]interface{}{"target": req.AgentID}, - }) return NewResult(req.AgentID, resp, nil) } diff --git a/agent/context/log.go b/agent/context/log.go index 503fc3fe..af3ed978 100644 --- a/agent/context/log.go +++ b/agent/context/log.go @@ -192,11 +192,8 @@ func (l *RequestLogger) processEntry(entry LogEntry) { } } -// printDev sends to TUI if available, otherwise prints colored output to stdout +// printDev prints colored output to stdout in development mode func (l *RequestLogger) printDev(entry LogEntry) { - if GetTUIProgram() != nil { - return - } switch entry.Level { case LogLevelTrace: fmt.Printf("%s → %s%s\n", colorGray, entry.Message, colorReset) @@ -315,17 +312,6 @@ func (l *RequestLogger) Start() { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - ParentID: l.parentID, - AssistantID: l.currentAssistantID(), - Event: EventRequestStart, - }) - - if GetTUIProgram() != nil { - return - } - fmt.Println() fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("═", 60), colorReset) fmt.Printf("%s AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset) @@ -357,21 +343,6 @@ func (l *RequestLogger) End(success bool, err error) { return } - data := map[string]interface{}{"duration": duration.Round(time.Millisecond)} - if err != nil { - data["error"] = err.Error() - } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - AssistantID: l.currentAssistantID(), - Event: EventRequestEnd, - Data: data, - }) - - if GetTUIProgram() != nil { - return - } - fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) if success { fmt.Printf("%s REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset) @@ -400,15 +371,6 @@ func (l *RequestLogger) Phase(name string) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventPhase, - Data: map[string]interface{}{"name": name, "elapsed": elapsed}, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s > %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset) } @@ -425,15 +387,6 @@ func (l *RequestLogger) PhaseComplete(name string) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventPhaseDone, - Data: map[string]interface{}{"name": name, "elapsed": elapsed}, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s + %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset) } @@ -447,15 +400,6 @@ func (l *RequestLogger) PhaseSkip(name, reason string) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventPhaseSkip, - Data: map[string]interface{}{"name": name, "reason": reason}, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s - %s (%s)%s\n", colorGray, name, reason, colorReset) } @@ -472,19 +416,6 @@ func (l *RequestLogger) LLMStart(connector, model string, messageCount int) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventLLMCall, - Data: map[string]interface{}{ - "connector": connector, - "model": model, - "messages": messageCount, - }, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset) fmt.Printf("%s Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset) if model != "" { @@ -511,19 +442,6 @@ func (l *RequestLogger) LLMComplete(tokens int, hasToolCalls bool) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventLLMDone, - Data: map[string]interface{}{ - "detail": fmt.Sprintf("%s [tokens:%d, %v]", status, tokens, elapsed), - "tokens": tokens, - "status": status, - }, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s + LLM Response (%s)%s", colorGreen, status, colorReset) if tokens > 0 { fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset) @@ -543,15 +461,6 @@ func (l *RequestLogger) ToolStart(toolName string) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventToolCall, - Data: map[string]interface{}{"name": toolName}, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s Tool: %s%s\n", colorYellow, toolName, colorReset) } @@ -571,15 +480,6 @@ func (l *RequestLogger) ToolComplete(toolName string, success bool) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventToolDone, - Data: map[string]interface{}{"name": toolName, "success": success}, - }) - - if GetTUIProgram() != nil { - return - } if success { fmt.Printf("%s + %s completed%s\n", colorGreen, toolName, colorReset) } else { @@ -600,15 +500,6 @@ func (l *RequestLogger) HookStart(hookName string) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventHook, - Data: map[string]interface{}{"name": hookName}, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset) } @@ -624,15 +515,6 @@ func (l *RequestLogger) HookComplete(hookName string) { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventHookDone, - Data: map[string]interface{}{"name": hookName}, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s + %s done%s\n", colorGreen, hookName, colorReset) } @@ -644,7 +526,7 @@ func (l *RequestLogger) Cleanup(resource string) { kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource) - if !config.IsDevelopment() || GetTUIProgram() != nil { + if !config.IsDevelopment() { return } fmt.Printf("%s + %s%s\n", colorGray, resource, colorReset) @@ -658,7 +540,7 @@ func (l *RequestLogger) HistoryLoad(count, maxSize int) { kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize) - if !config.IsDevelopment() || GetTUIProgram() != nil { + if !config.IsDevelopment() { return } fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset) @@ -673,7 +555,7 @@ func (l *RequestLogger) HistoryOverlap(overlapCount int) { if overlapCount > 0 { kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount) - if !config.IsDevelopment() || GetTUIProgram() != nil { + if !config.IsDevelopment() { return } fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset) @@ -692,15 +574,6 @@ func (l *RequestLogger) Release() { return } - SendTUI(AgentEventMsg{ - RequestID: l.requestID, - Event: EventContextRelease, - Data: map[string]interface{}{"assistant": l.currentAssistantID()}, - }) - - if GetTUIProgram() != nil { - return - } fmt.Printf("%s RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.currentAssistantID(), colorReset) } diff --git a/agent/context/tui.go b/agent/context/tui.go deleted file mode 100644 index 5062bcc9..00000000 --- a/agent/context/tui.go +++ /dev/null @@ -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()) -} diff --git a/agent/context/tui_msg.go b/agent/context/tui_msg.go deleted file mode 100644 index f910df44..00000000 --- a/agent/context/tui_msg.go +++ /dev/null @@ -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 diff --git a/cmd/start.go b/cmd/start.go index 18d2e367..fb86503e 100644 --- a/cmd/start.go +++ b/cmd/start.go @@ -8,14 +8,11 @@ import ( "strings" "syscall" - tea "github.com/charmbracelet/bubbletea" "github.com/fatih/color" - "github.com/mattn/go-isatty" "github.com/spf13/cobra" "github.com/yaoapp/gou/api" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/fs" - "github.com/yaoapp/gou/helper" "github.com/yaoapp/gou/mcp" "github.com/yaoapp/gou/plugin" "github.com/yaoapp/gou/schedule" @@ -23,9 +20,7 @@ import ( "github.com/yaoapp/gou/store" "github.com/yaoapp/gou/task" "github.com/yaoapp/gou/websocket" - "github.com/yaoapp/kun/exception" "github.com/yaoapp/kun/log" - agentcontext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/engine" yaogrpc "github.com/yaoapp/yao/grpc" @@ -35,12 +30,12 @@ import ( "github.com/yaoapp/yao/service" "github.com/yaoapp/yao/setup" "github.com/yaoapp/yao/share" + tairegistry "github.com/yaoapp/yao/tai/registry" itask "github.com/yaoapp/yao/task" ) var startDebug = false var startDisableWatching = false -var startTUI = false var startCmd = &cobra.Command{ Use: "start", @@ -231,6 +226,10 @@ var startCmd = &cobra.Command{ ischedule.Start() defer ischedule.Stop() + // Initialize the global Tai registry for tunnel and direct connections + // (must happen before HTTP/gRPC start so handlers can access it) + tairegistry.Init(nil) + // Start HTTP Server srv, err := service.Start(config.Conf) defer func() { @@ -288,7 +287,6 @@ var startCmd = &cobra.Command{ case http.READY: fmt.Println(color.GreenString(L("Server is up and running..."))) fmt.Println(color.GreenString("Ctrl+C to stop")) - initAgentTUI() break case http.CLOSED: @@ -649,38 +647,7 @@ func colorMehtod(method string) string { } } -// initAgentTUI initializes the TUI for agent request visualization in dev mode. -// Must be called after HTTP READY to avoid interfering with startup messages. -func initAgentTUI() { - if !config.IsDevelopment() { - return - } - - if !startTUI && os.Getenv("YAO_TUI") != "on" { - return - } - - if !isatty.IsTerminal(os.Stdout.Fd()) { - return - } - - model := agentcontext.NewAgentTUIModel() - p := tea.NewProgram(model, tea.WithoutSignalHandler()) - - agentcontext.SetTUIProgram(p) - tuiWriter := &agentcontext.TUILogWriter{Program: p} - helper.SetDevWriter(tuiWriter) - exception.SetWriter(tuiWriter) - - go func() { - if _, err := p.Run(); err != nil { - log.Error("TUI error: %s", err.Error()) - } - }() -} - func init() { startCmd.PersistentFlags().BoolVarP(&startDebug, "debug", "", false, L("Development mode")) startCmd.PersistentFlags().BoolVarP(&startDisableWatching, "disable-watching", "", false, L("Disable watching")) - startCmd.PersistentFlags().BoolVarP(&startTUI, "tui", "", false, L("Enable TUI for agent request visualization")) } diff --git a/go.mod b/go.mod index 8290777e..522c0406 100644 --- a/go.mod +++ b/go.mod @@ -10,8 +10,6 @@ require ( github.com/blang/semver v3.5.1+incompatible github.com/bwmarrin/discordgo v0.29.0 github.com/caarlos0/env/v6 v6.10.1 - github.com/charmbracelet/bubbletea v1.3.10 - github.com/charmbracelet/lipgloss v1.1.0 github.com/dchest/captcha v1.1.0 github.com/docker/docker v28.5.2+incompatible github.com/docker/go-connections v0.5.0 @@ -34,7 +32,6 @@ require ( github.com/kaptinlin/jsonschema v0.6.1 github.com/larksuite/oapi-sdk-go/v3 v3.5.3 github.com/matoous/go-nanoid/v2 v2.1.0 - github.com/mattn/go-isatty v0.0.20 github.com/mozillazg/go-pinyin v0.20.0 github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1 github.com/pierrec/lz4/v4 v4.1.25 @@ -79,17 +76,12 @@ require ( github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.15 // indirect github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.15 // indirect github.com/aws/smithy-go v1.22.3 // indirect - github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/blang/semver/v4 v4.0.0 // indirect github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect github.com/bytedance/sonic v1.13.2 // indirect github.com/bytedance/sonic/loader v0.2.4 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect - github.com/charmbracelet/x/ansi v0.10.1 // indirect - github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect - github.com/charmbracelet/x/term v0.2.1 // indirect github.com/cloudwego/base64x v0.1.5 // indirect github.com/coder/websocket v1.8.14 // indirect github.com/containerd/errdefs v1.0.0 // indirect @@ -101,7 +93,6 @@ require ( github.com/docker/go-units v0.5.0 // indirect github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 // indirect github.com/emicklei/go-restful/v3 v3.13.0 // indirect - github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/fxamacker/cbor/v2 v2.9.0 // indirect github.com/gabriel-vasile/mimetype v1.4.9 // indirect @@ -153,11 +144,10 @@ require ( github.com/klauspost/cpuid/v2 v2.2.10 // indirect github.com/leodido/go-urn v1.4.0 // indirect github.com/lib/pq v1.10.9 // indirect - github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/mark3labs/mcp-go v0.32.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect - github.com/mattn/go-localereader v0.0.1 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-runewidth v0.0.16 // indirect github.com/mattn/go-sqlite3 v1.14.28 // indirect github.com/miekg/dns v1.1.66 // indirect @@ -167,9 +157,6 @@ require ( github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect github.com/montanaflynn/stats v0.7.1 // indirect - github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect - github.com/muesli/cancelreader v0.2.2 // indirect - github.com/muesli/termenv v0.16.0 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f // indirect github.com/neo4j/neo4j-go-driver/v5 v5.28.1 // indirect @@ -211,7 +198,6 @@ require ( github.com/xdg-go/pbkdf2 v1.0.0 // indirect github.com/xdg-go/scram v1.1.2 // indirect github.com/xdg-go/stringprep v1.0.4 // indirect - github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/xuri/efp v0.0.1 // indirect github.com/xuri/nfp v0.0.1 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/go.sum b/go.sum index 300d01e9..09ce1856 100644 --- a/go.sum +++ b/go.sum @@ -43,8 +43,6 @@ github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3 h1:BRXS0U76Z8wfF+bnkilA2QwpIch6U github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3/go.mod h1:bNXKFFyaiVvWuR6O16h/I1724+aXe/tAkA9/QS01t5k= github.com/aws/smithy-go v1.22.3 h1:Z//5NuZCSW6R4PhQ93hShNbyBbn8BWCmCVCt+Q8Io5k= github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= -github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= -github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ= github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk= github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM= @@ -70,18 +68,6 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw= -github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4= -github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs= -github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk= -github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY= -github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30= -github.com/charmbracelet/x/ansi v0.10.1 h1:rL3Koar5XvX0pHGfovN03f5cxLbCF2YvLeyz7D2jVDQ= -github.com/charmbracelet/x/ansi v0.10.1/go.mod h1:3RQDQ6lDnROptfpWuUVIUG64bD2g2BgntdxH0Ya5TeE= -github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8= -github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs= -github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= -github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= @@ -123,8 +109,6 @@ github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6/go.mod h1:iL2twTe github.com/emersion/go-textwrapper v0.0.0-20200911093747-65d896831594/go.mod h1:aqO8z8wPrjkscevZJFVE1wXJrLpC5LtJG7fqLOsPb2U= github.com/emicklei/go-restful/v3 v3.13.0 h1:C4Bl2xDndpU6nJ4bc1jXd+uTmYPVUwkD6bFY/oTyCes= github.com/emicklei/go-restful/v3 v3.13.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= -github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= -github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/evanw/esbuild v0.25.4 h1:k1bTSim+usBG27w7BfOCorhgx3tO+6bAfMj5pR+6SKg= github.com/evanw/esbuild v0.25.4/go.mod h1:D2vIQZqV/vIf/VRHtViaUtViZmG7o+kKmlBfVQuRi48= github.com/expr-lang/expr v1.17.7 h1:Q0xY/e/2aCIp8g9s/LGvMDCC5PxYlvHgDZRQ4y16JX8= @@ -304,8 +288,6 @@ github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= -github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= -github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= github.com/mark3labs/mcp-go v0.32.0 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8= @@ -320,8 +302,6 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= -github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= -github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= @@ -351,12 +331,6 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= github.com/mozillazg/go-pinyin v0.20.0 h1:BtR3DsxpApHfKReaPO1fCqF4pThRwH9uwvXzm+GnMFQ= github.com/mozillazg/go-pinyin v0.20.0/go.mod h1:iR4EnMMRXkfpFVV5FMi4FNB6wGq9NV6uDWbUuPhP4Yc= -github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= -github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= -github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= -github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= -github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= -github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f h1:y5//uYreIhSUg3J1GEMiLbxo1LJaP8RfCpH6pymGZus= @@ -488,8 +462,6 @@ github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY= github.com/xdg-go/scram v1.1.2/go.mod h1:RT/sEzTbU5y00aCK8UOx6R7YryM0iF1N2MOmC3kKLN4= github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8= github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM= -github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= -github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/xuri/efp v0.0.1 h1:fws5Rv3myXyYni8uwj2qKjVaRP30PdjeYe2Y6FDsCL8= github.com/xuri/efp v0.0.1/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI= github.com/xuri/excelize/v2 v2.9.1 h1:VdSGk+rraGmgLHGFaGG9/9IWu1nj4ufjJ7uwMDtj8Qw= @@ -612,7 +584,6 @@ golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= diff --git a/grpc/grpc.go b/grpc/grpc.go index 1f204034..b5cc6f1c 100644 --- a/grpc/grpc.go +++ b/grpc/grpc.go @@ -192,6 +192,15 @@ func Stop() { addrs = nil } +// GRPCServer returns the active gRPC server instance. +// Used by the Tai tunnel server to serve data channel connections +// on the existing gRPC server. +func GRPCServer() *grpc.Server { + mu.Lock() + defer mu.Unlock() + return server +} + // Addr returns all addresses the gRPC server is listening on. func Addr() []string { mu.Lock() diff --git a/openapi/openapi.go b/openapi/openapi.go index 9f5e6653..224546a8 100644 --- a/openapi/openapi.go +++ b/openapi/openapi.go @@ -28,6 +28,7 @@ import ( "github.com/yaoapp/yao/openapi/team" openapiTrace "github.com/yaoapp/yao/openapi/trace" "github.com/yaoapp/yao/openapi/user" + taitunnel "github.com/yaoapp/yao/tai/tunnel" ) // Server is the OpenAPI server @@ -175,6 +176,12 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) { sandbox.SetPathPrefix(baseURL) sandbox.Attach(group.Group("/sandbox"), openapi.OAuth) + // Tai tunnel WebSocket and reverse proxy routes + group.GET("/ws/tai", taitunnel.HandleControl) + group.GET("/ws/tai/data/:channel_id", taitunnel.HandleData) + group.Any("/tai/:taiID/proxy/*path", taitunnel.HandleProxy) + group.GET("/tai/:taiID/vnc/*path", taitunnel.HandleVNC) + // Custom handlers (Defined by developer) } diff --git a/sandbox/v2/grpc.go b/sandbox/v2/grpc.go index 522f0123..e676d5a7 100644 --- a/sandbox/v2/grpc.go +++ b/sandbox/v2/grpc.go @@ -4,6 +4,7 @@ import ( "crypto/rand" "encoding/hex" "fmt" + "net/url" "strconv" "strings" ) @@ -35,6 +36,7 @@ func RevokeContainerTokens(refresh string) error { } // BuildGRPCEnv builds the gRPC environment variables for a sandbox container. +// Supports tai:// (direct), tunnel:// (NAT traversal), and local modes. func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) map[string]string { portStr := strconv.Itoa(grpcPort) env := map[string]string{ @@ -42,12 +44,34 @@ func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) m "YAO_TOKEN": access, "YAO_REFRESH_TOKEN": refresh, } - if pool != nil && strings.Contains(pool.Addr, "tai://") { - taiHost := strings.TrimPrefix(pool.Addr, "tai://") + + if pool == nil { + env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr) + return env + } + + switch { + case strings.HasPrefix(pool.Addr, "tunnel://"): env["YAO_GRPC_TAI"] = "enable" - env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:9100", taiHost) + env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%d", grpcPort) env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr) - } else { + + case strings.HasPrefix(pool.Addr, "tai://"): + u, err := url.Parse(pool.Addr) + if err != nil { + env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr) + return env + } + taiHost := u.Hostname() + taiPort := u.Port() + if taiPort == "" { + taiPort = "9100" + } + env["YAO_GRPC_TAI"] = "enable" + env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%s", taiHost, taiPort) + env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr) + + default: env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr) } return env diff --git a/tai/proxy/connect.go b/tai/proxy/connect.go index ecd9f043..b8663d0b 100644 --- a/tai/proxy/connect.go +++ b/tai/proxy/connect.go @@ -21,6 +21,16 @@ func (r *remoteProxy) Connect(ctx context.Context, containerID string, opts Conn return connect(ctx, baseURL, opts.Protocol, r.client) } +// --- Tunnel Connect --- + +func (t *tunnelProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) { + baseURL, err := t.URL(ctx, containerID, opts.Port, opts.Path) + if err != nil { + return nil, err + } + return connect(ctx, baseURL, opts.Protocol, http.DefaultClient) +} + // --- Local Connect --- func (l *localProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) { diff --git a/tai/proxy/proxy.go b/tai/proxy/proxy.go index c79ca5d4..ba556259 100644 --- a/tai/proxy/proxy.go +++ b/tai/proxy/proxy.go @@ -73,6 +73,28 @@ func (r *remoteProxy) Healthz(ctx context.Context) error { return nil } +// --- Tunnel implementation --- + +type tunnelProxy struct { + taiID string + yaoBase string // e.g. "http://yao-host:5099" +} + +// NewTunnel creates a Proxy that routes through Yao's reverse proxy for +// tunnel-connected Tai instances. URLs point to {yaoBase}/tai/{taiID}/proxy/*. +func NewTunnel(taiID, yaoBase string) Proxy { + return &tunnelProxy{taiID: taiID, yaoBase: strings.TrimRight(yaoBase, "/")} +} + +func (t *tunnelProxy) URL(_ context.Context, containerID string, port int, path string) (string, error) { + path = strings.TrimPrefix(path, "/") + return fmt.Sprintf("%s/tai/%s/proxy/%s:%d/%s", t.yaoBase, t.taiID, containerID, port, path), nil +} + +func (t *tunnelProxy) Healthz(_ context.Context) error { + return nil +} + // --- Local implementation --- type localProxy struct { diff --git a/tai/registry/registry.go b/tai/registry/registry.go new file mode 100644 index 00000000..7bb15711 --- /dev/null +++ b/tai/registry/registry.go @@ -0,0 +1,401 @@ +package registry + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "io" + "log/slog" + "net" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +// TaiNode represents a registered Tai instance (direct or tunnel). +// Internal use only; external callers receive NodeSnapshot via Get()/List(). +type TaiNode struct { + TaiID string + MachineID string + Version string + Auth AuthInfo + Mode string // "direct" | "tunnel" + Addr string // direct mode: "tai-host"; tunnel mode: empty + YaoBase string // Yao server base URL reported by Tai (tunnel mode) + Ports map[string]int // {"grpc":9100, "http":8080, "vnc":6080, "docker":2375} + Capabilities map[string]bool + + ControlConn *websocket.Conn + connMu sync.Mutex // protects ControlConn writes + + Status string // "online" | "offline" | "connecting" + ConnectedAt time.Time + LastPing time.Time + PoolName string + + localListeners map[int]*tunnelListener +} + +// NodeSnapshot is a read-only copy of TaiNode fields safe to use outside locks. +type NodeSnapshot struct { + TaiID string + MachineID string + Version string + Auth AuthInfo + Mode string + Addr string + YaoBase string + Ports map[string]int + Capabilities map[string]bool + Status string + ConnectedAt time.Time + LastPing time.Time + PoolName string +} + +func (n *TaiNode) snapshot() NodeSnapshot { + ports := make(map[string]int, len(n.Ports)) + for k, v := range n.Ports { + ports[k] = v + } + caps := make(map[string]bool, len(n.Capabilities)) + for k, v := range n.Capabilities { + caps[k] = v + } + return NodeSnapshot{ + TaiID: n.TaiID, MachineID: n.MachineID, Version: n.Version, + Auth: n.Auth, Mode: n.Mode, Addr: n.Addr, YaoBase: n.YaoBase, + Ports: ports, Capabilities: caps, + Status: n.Status, ConnectedAt: n.ConnectedAt, LastPing: n.LastPing, + PoolName: n.PoolName, + } +} + +// AuthInfo holds Yao user authorization extracted from OAuth token. +type AuthInfo struct { + Subject string + UserID string + ClientID string + Scope string + TeamID string + TenantID string +} + +// pendingChannel represents a channel awaiting Tai's data WS connection. +type pendingChannel struct { + taiID string + result chan net.Conn + timer *time.Timer +} + +// tunnelListener wraps a TCP listener that bridges each accepted connection +// through the WS tunnel to a specific Tai port. +type tunnelListener struct { + listener net.Listener + taiID string + port int + cancel func() +} + +var ( + global *Registry + once sync.Once +) + +// Registry manages all Tai nodes (direct and tunnel). +type Registry struct { + mu sync.RWMutex + nodes map[string]*TaiNode + pending map[string]*pendingChannel + logger *slog.Logger +} + +// Init initializes the global registry singleton. +func Init(logger *slog.Logger) { + once.Do(func() { + if logger == nil { + logger = slog.Default() + } + global = &Registry{ + nodes: make(map[string]*TaiNode), + pending: make(map[string]*pendingChannel), + logger: logger, + } + }) +} + +// Global returns the global registry instance. +func Global() *Registry { + return global +} + +// Register adds or updates a Tai node in the registry. +func (r *Registry) Register(node *TaiNode) { + r.mu.Lock() + defer r.mu.Unlock() + + node.Status = "online" + node.ConnectedAt = time.Now() + node.LastPing = time.Now() + if node.localListeners == nil { + node.localListeners = make(map[int]*tunnelListener) + } + r.nodes[node.TaiID] = node + + r.logger.Info("tai node registered", + "tai_id", node.TaiID, "mode", node.Mode, "version", node.Version) +} + +// Unregister removes a Tai node and closes its local listeners and control connection. +func (r *Registry) Unregister(taiID string) { + r.mu.Lock() + node, ok := r.nodes[taiID] + if ok { + for _, tl := range node.localListeners { + tl.cancel() + tl.listener.Close() + } + node.connMu.Lock() + if node.ControlConn != nil { + node.ControlConn.Close() + node.ControlConn = nil + } + node.connMu.Unlock() + delete(r.nodes, taiID) + } + r.mu.Unlock() + + if ok { + r.logger.Info("tai node unregistered", "tai_id", taiID) + } +} + +// Get returns a snapshot of a Tai node by ID. Returns nil, false if not found. +func (r *Registry) Get(taiID string) (*NodeSnapshot, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + n, ok := r.nodes[taiID] + if !ok { + return nil, false + } + snap := n.snapshot() + return &snap, true +} + +// List returns snapshots of all registered Tai nodes. +func (r *Registry) List() []NodeSnapshot { + r.mu.RLock() + defer r.mu.RUnlock() + result := make([]NodeSnapshot, 0, len(r.nodes)) + for _, n := range r.nodes { + result = append(result, n.snapshot()) + } + return result +} + +// WriteControlJSON sends a JSON message on the node's control channel +// with proper serialization. Returns error if node not found or not tunnel. +func (r *Registry) WriteControlJSON(taiID string, v interface{}) error { + r.mu.RLock() + node := r.nodes[taiID] + r.mu.RUnlock() + + if node == nil { + return fmt.Errorf("tai node %s not found", taiID) + } + + node.connMu.Lock() + defer node.connMu.Unlock() + if node.ControlConn == nil { + return fmt.Errorf("tai node %s has no active control channel", taiID) + } + return node.ControlConn.WriteJSON(v) +} + +// UpdatePing records a heartbeat timestamp. +func (r *Registry) UpdatePing(taiID string) { + r.mu.Lock() + defer r.mu.Unlock() + if n, ok := r.nodes[taiID]; ok { + n.LastPing = time.Now() + } +} + +// RequestChannel sends an "open" command to a tunnel-connected Tai via its +// control channel. Returns a channel_id that Tai will use to connect back. +// Blocks until the data channel is established or timeout. +func (r *Registry) RequestChannel(taiID string, targetPort int) (string, chan net.Conn, error) { + r.mu.RLock() + node := r.nodes[taiID] + r.mu.RUnlock() + + if node == nil { + return "", nil, fmt.Errorf("tai node %s not found", taiID) + } + if node.Mode != "tunnel" { + return "", nil, fmt.Errorf("tai node %s is not a tunnel node", taiID) + } + node.connMu.Lock() + hasConn := node.ControlConn != nil + node.connMu.Unlock() + if !hasConn { + return "", nil, fmt.Errorf("tai node %s has no active control channel", taiID) + } + + channelID, err := generateChannelID() + if err != nil { + return "", nil, fmt.Errorf("generate channel_id: %w", err) + } + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(30*time.Second, func() { + r.mu.Lock() + if pc, ok := r.pending[channelID]; ok { + close(pc.result) + delete(r.pending, channelID) + } + r.mu.Unlock() + }) + + r.mu.Lock() + r.pending[channelID] = &pendingChannel{taiID: taiID, result: resultCh, timer: timer} + r.mu.Unlock() + + msg := map[string]interface{}{ + "type": "open", + "channel_id": channelID, + "target_port": targetPort, + } + if err := r.WriteControlJSON(taiID, msg); err != nil { + r.mu.Lock() + delete(r.pending, channelID) + r.mu.Unlock() + timer.Stop() + return "", nil, fmt.Errorf("send open command: %w", err) + } + + return channelID, resultCh, nil +} + +// AcceptDataChannel resolves a pending channel when Tai connects its data WS. +// The taiID must match the node that requested the channel via RequestChannel. +func (r *Registry) AcceptDataChannel(channelID, taiID string, conn net.Conn) error { + r.mu.Lock() + pc, ok := r.pending[channelID] + if ok { + delete(r.pending, channelID) + } + r.mu.Unlock() + + if !ok { + return fmt.Errorf("no pending channel for %s", channelID) + } + if pc.taiID != taiID { + pc.timer.Stop() + close(pc.result) + return fmt.Errorf("channel %s: tai_id mismatch (expected %s, got %s)", channelID, pc.taiID, taiID) + } + pc.timer.Stop() + pc.result <- conn + return nil +} + +// OpenLocalListener creates a localhost TCP listener that tunnels every +// accepted connection to the specified port on the given Tai node. +// Returns the listener address (e.g. "127.0.0.1:54321"). +func (r *Registry) OpenLocalListener(taiID string, targetPort int) (net.Listener, error) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("listen: %w", err) + } + + ctx, cancel := newContext() + tl := &tunnelListener{listener: ln, taiID: taiID, port: targetPort, cancel: cancel} + + r.mu.Lock() + node := r.nodes[taiID] + if node == nil { + r.mu.Unlock() + cancel() + ln.Close() + return nil, fmt.Errorf("tai node %s not found", taiID) + } + node.localListeners[targetPort] = tl + r.mu.Unlock() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + select { + case <-ctx.Done(): + return + default: + r.logger.Debug("tunnel listener accept error", "err", err) + return + } + } + go r.bridgeTunnelConn(taiID, targetPort, conn) + } + }() + + r.logger.Info("tunnel local listener started", + "tai_id", taiID, "target_port", targetPort, "local_addr", ln.Addr().String()) + return ln, nil +} + +func (r *Registry) bridgeTunnelConn(taiID string, targetPort int, localConn net.Conn) { + channelID, resultCh, err := r.RequestChannel(taiID, targetPort) + if err != nil { + localConn.Close() + r.logger.Error("request channel failed", "tai_id", taiID, "port", targetPort, "err", err) + return + } + + remoteConn, ok := <-resultCh + if !ok || remoteConn == nil { + localConn.Close() + r.logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID) + return + } + + bridgeTCP(localConn, remoteConn) +} + +// bridgeTCP copies bytes bidirectionally between two net.Conn, closing both when done. +func bridgeTCP(a, b net.Conn) { + var wg sync.WaitGroup + wg.Add(2) + + cp := func(dst, src net.Conn) { + defer wg.Done() + io.Copy(dst, src) + dst.Close() + } + + go cp(a, b) + go cp(b, a) + wg.Wait() +} + +func generateChannelID() (string, error) { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return "", err + } + return hex.EncodeToString(b), nil +} + +type contextCancel struct { + done chan struct{} +} + +func newContext() (*contextCancel, func()) { + cc := &contextCancel{done: make(chan struct{})} + return cc, func() { close(cc.done) } +} + +func (c *contextCancel) Done() <-chan struct{} { + return c.done +} diff --git a/tai/registry/registry_test.go b/tai/registry/registry_test.go new file mode 100644 index 00000000..25cf9a39 --- /dev/null +++ b/tai/registry/registry_test.go @@ -0,0 +1,499 @@ +package registry + +import ( + "log/slog" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +// newTestRegistry creates a standalone registry for testing (bypasses global singleton). +func newTestRegistry() *Registry { + return &Registry{ + nodes: make(map[string]*TaiNode), + pending: make(map[string]*pendingChannel), + logger: slog.Default(), + } +} + +func TestRegister_SetsFieldsAndOnline(t *testing.T) { + r := newTestRegistry() + node := &TaiNode{ + TaiID: "tai-001", + MachineID: "m-abc", + Version: "1.0.0", + Mode: "tunnel", + Ports: map[string]int{"grpc": 9100}, + } + r.Register(node) + + snap, ok := r.Get("tai-001") + if !ok { + t.Fatal("expected node to exist after Register") + } + if snap.Status != "online" { + t.Errorf("Status = %q, want online", snap.Status) + } + if snap.MachineID != "m-abc" { + t.Errorf("MachineID = %q, want m-abc", snap.MachineID) + } + if snap.ConnectedAt.IsZero() { + t.Error("ConnectedAt should be set") + } +} + +func TestRegister_Overwrite(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Version: "1.0"}) + r.Register(&TaiNode{TaiID: "tai-001", Version: "2.0"}) + + snap, ok := r.Get("tai-001") + if !ok { + t.Fatal("node should exist") + } + if snap.Version != "2.0" { + t.Errorf("Version = %q, want 2.0 after re-register", snap.Version) + } +} + +func TestUnregister_RemovesNode(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001"}) + r.Unregister("tai-001") + + if _, ok := r.Get("tai-001"); ok { + t.Error("expected node to be removed after Unregister") + } +} + +func TestUnregister_Nonexistent(t *testing.T) { + r := newTestRegistry() + r.Unregister("ghost") +} + +func TestGet_NotFound(t *testing.T) { + r := newTestRegistry() + if _, ok := r.Get("missing"); ok { + t.Error("expected false for missing node") + } +} + +func TestList_Empty(t *testing.T) { + r := newTestRegistry() + if got := r.List(); len(got) != 0 { + t.Errorf("List() = %d items, want 0", len(got)) + } +} + +func TestList_MultipleNodes(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "a"}) + r.Register(&TaiNode{TaiID: "b"}) + r.Register(&TaiNode{TaiID: "c"}) + + list := r.List() + if len(list) != 3 { + t.Errorf("List() = %d items, want 3", len(list)) + } + + ids := map[string]bool{} + for _, snap := range list { + ids[snap.TaiID] = true + } + for _, id := range []string{"a", "b", "c"} { + if !ids[id] { + t.Errorf("missing node %q in List()", id) + } + } +} + +func TestSnapshot_DeepCopy(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{ + TaiID: "tai-001", + Ports: map[string]int{"grpc": 9100, "http": 8080}, + }) + + snap, _ := r.Get("tai-001") + snap.Ports["grpc"] = 0 + + snap2, _ := r.Get("tai-001") + if snap2.Ports["grpc"] != 9100 { + t.Error("snapshot modification leaked into registry node") + } +} + +func TestUpdatePing(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001"}) + time.Sleep(10 * time.Millisecond) + + r.UpdatePing("tai-001") + snap, _ := r.Get("tai-001") + if snap.LastPing.Before(snap.ConnectedAt) { + t.Error("LastPing should be after ConnectedAt") + } +} + +func TestUpdatePing_NonexistentNode(t *testing.T) { + r := newTestRegistry() + r.UpdatePing("ghost") +} + +func TestWriteControlJSON_NoNode(t *testing.T) { + r := newTestRegistry() + err := r.WriteControlJSON("missing", map[string]string{"type": "test"}) + if err == nil { + t.Fatal("expected error for missing node") + } +} + +func TestWriteControlJSON_NilConn(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001"}) + err := r.WriteControlJSON("tai-001", map[string]string{"type": "test"}) + if err == nil { + t.Fatal("expected error for nil ControlConn") + } +} + +func TestRequestChannel_NotFound(t *testing.T) { + r := newTestRegistry() + _, _, err := r.RequestChannel("ghost", 9100) + if err == nil { + t.Fatal("expected error for missing node") + } +} + +func TestRequestChannel_DirectMode(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "direct"}) + _, _, err := r.RequestChannel("tai-001", 9100) + if err == nil { + t.Fatal("expected error for direct-mode node") + } +} + +func TestAcceptDataChannel_NotPending(t *testing.T) { + r := newTestRegistry() + pipe1, pipe2 := net.Pipe() + defer pipe1.Close() + defer pipe2.Close() + + err := r.AcceptDataChannel("unknown-channel", "tai-001", pipe1) + if err == nil { + t.Fatal("expected error for non-pending channel") + } +} + +func TestAcceptDataChannel_TaiIDMismatch(t *testing.T) { + r := newTestRegistry() + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(5*time.Second, func() {}) + r.mu.Lock() + r.pending["ch-001"] = &pendingChannel{taiID: "tai-owner", result: resultCh, timer: timer} + r.mu.Unlock() + + pipe1, pipe2 := net.Pipe() + defer pipe1.Close() + defer pipe2.Close() + + err := r.AcceptDataChannel("ch-001", "tai-intruder", pipe1) + if err == nil { + t.Fatal("expected error for tai_id mismatch") + } +} + +func TestAcceptDataChannel_Success(t *testing.T) { + r := newTestRegistry() + + resultCh := make(chan net.Conn, 1) + timer := time.AfterFunc(5*time.Second, func() {}) + r.mu.Lock() + r.pending["ch-002"] = &pendingChannel{taiID: "tai-001", result: resultCh, timer: timer} + r.mu.Unlock() + + pipe1, pipe2 := net.Pipe() + defer pipe2.Close() + + if err := r.AcceptDataChannel("ch-002", "tai-001", pipe1); err != nil { + t.Fatalf("AcceptDataChannel: %v", err) + } + + select { + case conn := <-resultCh: + if conn == nil { + t.Fatal("expected non-nil conn") + } + conn.Close() + case <-time.After(time.Second): + t.Fatal("timeout waiting for conn on resultCh") + } +} + +func TestGenerateChannelID_Unique(t *testing.T) { + seen := make(map[string]bool) + for i := 0; i < 100; i++ { + id, err := generateChannelID() + if err != nil { + t.Fatalf("generateChannelID: %v", err) + } + if len(id) != 64 { + t.Errorf("len = %d, want 64 hex chars", len(id)) + } + if seen[id] { + t.Fatalf("duplicate channel ID: %s", id) + } + seen[id] = true + } +} + +func TestBridgeTCP(t *testing.T) { + a1, a2 := net.Pipe() + b1, b2 := net.Pipe() + + go bridgeTCP(a2, b1) + + msg := []byte("hello tunnel") + go func() { + a1.Write(msg) + a1.Close() + }() + + buf := make([]byte, 64) + n, _ := b2.Read(buf) + if string(buf[:n]) != "hello tunnel" { + t.Errorf("got %q, want %q", buf[:n], "hello tunnel") + } + b2.Close() +} + +func TestConcurrentRegisterGet(t *testing.T) { + r := newTestRegistry() + var wg sync.WaitGroup + + for i := 0; i < 50; i++ { + wg.Add(2) + id := "tai-" + string(rune('A'+i%26)) + + go func() { + defer wg.Done() + r.Register(&TaiNode{TaiID: id, Mode: "tunnel"}) + }() + + go func() { + defer wg.Done() + r.Get(id) + r.List() + }() + } + + wg.Wait() +} + +func TestWriteControlJSON_Success(t *testing.T) { + done := make(chan map[string]string, 1) + + srv := newWSServer(func(conn *websocket.Conn) { + var msg map[string]string + conn.ReadJSON(&msg) + done <- msg + conn.Close() + }) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn}) + + payload := map[string]string{"type": "test", "data": "hello"} + if err := r.WriteControlJSON("tai-001", payload); err != nil { + t.Fatalf("WriteControlJSON: %v", err) + } + + select { + case got := <-done: + if got["type"] != "test" { + t.Errorf("type = %q, want test", got["type"]) + } + if got["data"] != "hello" { + t.Errorf("data = %q, want hello", got["data"]) + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for server to receive message") + } +} + +func TestRequestChannel_Success(t *testing.T) { + openCh := make(chan map[string]interface{}, 1) + + srv := newWSServer(func(conn *websocket.Conn) { + var msg map[string]interface{} + conn.ReadJSON(&msg) + openCh <- msg + time.Sleep(time.Second) + conn.Close() + }) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn}) + + channelID, resultCh, err := r.RequestChannel("tai-001", 9100) + if err != nil { + t.Fatalf("RequestChannel: %v", err) + } + if channelID == "" { + t.Fatal("channelID should not be empty") + } + if len(channelID) != 64 { + t.Errorf("channelID len = %d, want 64", len(channelID)) + } + if resultCh == nil { + t.Fatal("resultCh should not be nil") + } + + select { + case cmd := <-openCh: + if cmd["type"] != "open" { + t.Errorf("cmd type = %v, want open", cmd["type"]) + } + if cmd["channel_id"] != channelID { + t.Errorf("cmd channel_id = %v, want %s", cmd["channel_id"], channelID) + } + if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 { + t.Errorf("cmd target_port = %v, want 9100", cmd["target_port"]) + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for open command") + } +} + +func TestRequestChannel_NoControlConn(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel"}) + + _, _, err := r.RequestChannel("tai-001", 9100) + if err == nil { + t.Fatal("expected error for nil ControlConn") + } +} + +func TestOpenLocalListener_Success(t *testing.T) { + r := newTestRegistry() + + controlCh := make(chan map[string]interface{}, 1) + srv := newWSServer(func(conn *websocket.Conn) { + for { + var msg map[string]interface{} + if err := conn.ReadJSON(&msg); err != nil { + return + } + controlCh <- msg + } + }) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + + r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn}) + + ln, err := r.OpenLocalListener("tai-001", 9100) + if err != nil { + t.Fatalf("OpenLocalListener: %v", err) + } + defer ln.Close() + + addr := ln.Addr().String() + if addr == "" { + t.Fatal("listener address should not be empty") + } + if !strings.HasPrefix(addr, "127.0.0.1:") { + t.Errorf("addr = %q, want 127.0.0.1:*", addr) + } + + conn, err := net.DialTimeout("tcp", addr, time.Second) + if err != nil { + t.Fatalf("connect to local listener: %v", err) + } + defer conn.Close() + + select { + case cmd := <-controlCh: + if cmd["type"] != "open" { + t.Errorf("open cmd type = %v, want open", cmd["type"]) + } + if _, ok := cmd["channel_id"].(string); !ok { + t.Error("open cmd missing channel_id") + } + if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 { + t.Errorf("target_port = %v, want 9100", cmd["target_port"]) + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for open command from local listener") + } +} + +func TestOpenLocalListener_NodeNotFound(t *testing.T) { + r := newTestRegistry() + _, err := r.OpenLocalListener("ghost", 9100) + if err == nil { + t.Fatal("expected error for missing node") + } +} + +func newWSServer(handler func(*websocket.Conn)) *httptest.Server { + up := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + handler(conn) + })) +} + +func TestNodeSnapshot_AuthInfo(t *testing.T) { + r := newTestRegistry() + r.Register(&TaiNode{ + TaiID: "tai-001", + Auth: AuthInfo{ + Subject: "user123", + ClientID: "tai-001", + Scope: "tai:tunnel", + }, + }) + + snap, ok := r.Get("tai-001") + if !ok { + t.Fatal("node not found") + } + if snap.Auth.Subject != "user123" { + t.Errorf("Auth.Subject = %q, want user123", snap.Auth.Subject) + } + if snap.Auth.Scope != "tai:tunnel" { + t.Errorf("Auth.Scope = %q, want tai:tunnel", snap.Auth.Scope) + } +} diff --git a/tai/registry/testing.go b/tai/registry/testing.go new file mode 100644 index 00000000..6da804b0 --- /dev/null +++ b/tai/registry/testing.go @@ -0,0 +1,31 @@ +package registry + +import ( + "log/slog" + "net" + "time" +) + +// NewForTest creates a standalone Registry for use in tests. +// Not intended for production use. +func NewForTest() *Registry { + return &Registry{ + nodes: make(map[string]*TaiNode), + pending: make(map[string]*pendingChannel), + logger: slog.Default(), + } +} + +// SetGlobalForTest replaces the global registry singleton for testing. +// Not intended for production use. +func SetGlobalForTest(r *Registry) { + global = r +} + +// SetPendingForTest injects a pending channel entry for testing. +// Not intended for production use. +func (r *Registry) SetPendingForTest(channelID, taiID string, result chan net.Conn, timer *time.Timer) { + r.mu.Lock() + defer r.mu.Unlock() + r.pending[channelID] = &pendingChannel{taiID: taiID, result: result, timer: timer} +} diff --git a/tai/tai.go b/tai/tai.go index 1fc52ec0..76acf990 100644 --- a/tai/tai.go +++ b/tai/tai.go @@ -3,6 +3,7 @@ package tai import ( "context" "fmt" + "net" "net/http" "net/url" "strconv" @@ -10,6 +11,7 @@ import ( "time" "github.com/yaoapp/yao/tai/proxy" + "github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/sandbox" sipb "github.com/yaoapp/yao/tai/serverinfo/pb" "github.com/yaoapp/yao/tai/vnc" @@ -124,7 +126,7 @@ func mergedPorts(p Ports) Ports { // Client provides unified access to all Tai SDK sub-packages. type Client struct { - scheme string // "tai" or "docker" + scheme string // "tai", "docker", or "tunnel" host string addr string ports Ports @@ -135,6 +137,9 @@ type Client struct { prx proxy.Proxy vc vnc.VNC grpcConn *grpc.ClientConn + + // tunnel mode: local listeners that bridge to Tai via WS + tunnelListeners []net.Listener } // New creates a Client based on the address protocol: @@ -172,6 +177,8 @@ func New(addr string, opts ...Option) (*Client, error) { return c.initLocal(cfg) case "tai": return c.initRemote(cfg) + case "tunnel": + return c.initTunnel(cfg) default: return nil, fmt.Errorf("unsupported scheme: %s", scheme) } @@ -256,9 +263,97 @@ func (c *Client) initRemote(cfg *config) (*Client, error) { hc := cfg.httpClient c.prx = proxy.NewRemote(c.host, c.ports.HTTP, hc) c.vc = vnc.NewRemote(c.host, c.ports.VNC, hc) + + if reg := registry.Global(); reg != nil { + reg.Register(®istry.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 diff --git a/tai/tunnel/proxy.go b/tai/tunnel/proxy.go new file mode 100644 index 00000000..a02d92aa --- /dev/null +++ b/tai/tunnel/proxy.go @@ -0,0 +1,172 @@ +package tunnel + +import ( + "bufio" + "io" + "log/slog" + "net" + "net/http" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/yaoapp/yao/tai/registry" +) + +// HandleProxy handles HTTP reverse proxy requests for a tunnel-connected Tai: +// ANY /tai/:taiID/proxy/*path +// Opens a data channel to Tai's HTTP port, forwards the HTTP request, +// and streams the response back. +func HandleProxy(c *gin.Context) { + logger := slog.Default() + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"}) + return + } + + taiID := c.Param("taiID") + node, ok := reg.Get(taiID) + if !ok || node.Status != "online" { + c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"}) + return + } + + httpPort := node.Ports["http"] + if httpPort == 0 { + httpPort = 8080 + } + + channelID, resultCh, err := reg.RequestChannel(taiID, httpPort) + if err != nil { + logger.Error("request channel failed", "tai_id", taiID, "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"}) + return + } + + remoteConn, ok := <-resultCh + if !ok || remoteConn == nil { + logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID) + c.JSON(http.StatusGatewayTimeout, gin.H{"error": "data channel timeout"}) + return + } + defer remoteConn.Close() + + path := c.Param("path") + outReq, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, "http://tai-tunnel"+path, c.Request.Body) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "build request failed"}) + return + } + outReq.Header = c.Request.Header.Clone() + outReq.Host = c.Request.Host + + if err := outReq.Write(remoteConn); err != nil { + logger.Error("write request to tunnel", "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "write to tunnel failed"}) + return + } + + resp, err := http.ReadResponse(bufio.NewReader(remoteConn), outReq) + if err != nil { + logger.Error("read response from tunnel", "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "read from tunnel failed"}) + return + } + defer resp.Body.Close() + + for k, vv := range resp.Header { + for _, v := range vv { + c.Writer.Header().Add(k, v) + } + } + c.Writer.WriteHeader(resp.StatusCode) + io.Copy(c.Writer, resp.Body) +} + +// HandleVNC handles VNC WebSocket proxying for a tunnel-connected Tai: +// GET /tai/:taiID/vnc/*path +// Upgrades the client connection to WebSocket, opens a data channel to +// Tai's VNC port, and bridges the two WebSocket connections. +func HandleVNC(c *gin.Context) { + logger := slog.Default() + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"}) + return + } + + taiID := c.Param("taiID") + node, ok := reg.Get(taiID) + if !ok || node.Status != "online" { + c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"}) + return + } + + vncPort := node.Ports["vnc"] + if vncPort == 0 { + vncPort = 6080 + } + + channelID, resultCh, err := reg.RequestChannel(taiID, vncPort) + if err != nil { + logger.Error("request vnc channel failed", "tai_id", taiID, "err", err) + c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"}) + return + } + + clientConn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + logger.Error("ws upgrade client failed", "err", err) + return + } + + taiConn, ok := <-resultCh + if !ok || taiConn == nil { + logger.Error("vnc data channel timeout", "tai_id", taiID, "channel_id", channelID) + clientConn.Close() + return + } + + bridgeWSToConn(clientConn, taiConn) +} + +// bridgeWSToConn bridges a client WebSocket to a net.Conn (tunnel data channel). +func bridgeWSToConn(clientWS *websocket.Conn, taiConn net.Conn) { + done := make(chan struct{}, 2) + + // client WS -> tai conn + go func() { + defer func() { done <- struct{}{} }() + for { + _, data, err := clientWS.ReadMessage() + if err != nil { + return + } + if _, err := taiConn.Write(data); err != nil { + return + } + } + }() + + // tai conn -> client WS + go func() { + defer func() { done <- struct{}{} }() + buf := make([]byte, 32*1024) + for { + n, err := taiConn.Read(buf) + if n > 0 { + if wErr := clientWS.WriteMessage(websocket.BinaryMessage, buf[:n]); wErr != nil { + return + } + } + if err != nil { + return + } + } + }() + + <-done + clientWS.Close() + taiConn.Close() + <-done +} diff --git a/tai/tunnel/server.go b/tai/tunnel/server.go new file mode 100644 index 00000000..1910358c --- /dev/null +++ b/tai/tunnel/server.go @@ -0,0 +1,267 @@ +package tunnel + +import ( + "fmt" + "io" + "log/slog" + "net" + "net/http" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + oauth "github.com/yaoapp/yao/openapi/oauth" + "github.com/yaoapp/yao/tai/registry" +) + +var upgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, +} + +// HandleControl handles the Tai control channel WebSocket: GET /ws/tai. +// Authenticates via Bearer token, reads register + ping messages, +// and maintains the Tai node in the global registry. +func HandleControl(c *gin.Context) { + logger := slog.Default() + reg := registry.Global() + if reg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"}) + return + } + + bearer := extractBearer(c.Request) + if bearer == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"}) + return + } + + authInfo, err := authenticateBearerFunc(bearer) + if err != nil { + logger.Warn("tunnel auth failed", "err", err) + c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"}) + return + } + + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + logger.Error("ws upgrade failed", "err", err) + return + } + + // Read the register message + var regMsg registerMessage + if err := conn.ReadJSON(®Msg); 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 := ®istry.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) } diff --git a/tai/tunnel/server_test.go b/tai/tunnel/server_test.go new file mode 100644 index 00000000..186e13b5 --- /dev/null +++ b/tai/tunnel/server_test.go @@ -0,0 +1,604 @@ +package tunnel + +import ( + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/yaoapp/yao/tai/registry" +) + +func init() { + gin.SetMode(gin.TestMode) +} + +func setupTestRegistry() *registry.Registry { + r := registry.NewForTest() + registry.SetGlobalForTest(r) + return r +} + +func mockAuth(info registry.AuthInfo, authErr error) func() { + old := authenticateBearerFunc + authenticateBearerFunc = func(token string) (registry.AuthInfo, error) { + return info, authErr + } + return func() { authenticateBearerFunc = old } +} + +// --- extractBearer --- + +func TestExtractBearer(t *testing.T) { + tests := []struct { + name string + header string + want string + }{ + {"valid", "Bearer abc123", "abc123"}, + {"lowercase", "bearer xyz", "xyz"}, + {"empty", "", ""}, + {"no_scheme", "abc123", ""}, + {"only_bearer", "Bearer ", ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r := &http.Request{Header: http.Header{}} + if tt.header != "" { + r.Header.Set("Authorization", tt.header) + } + got := extractBearer(r) + if got != tt.want { + t.Errorf("extractBearer(%q) = %q, want %q", tt.header, got, tt.want) + } + }) + } +} + +// --- wsConn --- + +func TestWSConn_EchoRoundTrip(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + wc := newWSConn(conn) + buf := make([]byte, 256) + n, err := wc.Read(buf) + if err != nil { + return + } + wc.Write(buf[:n]) + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Errorf("handshake status = %d, want 101", resp.StatusCode) + } + + msg := []byte("hello tunnel") + if err := conn.WriteMessage(websocket.BinaryMessage, msg); err != nil { + t.Fatalf("write: %v", err) + } + + mt, reply, err := conn.ReadMessage() + if err != nil { + t.Fatalf("read: %v", err) + } + if mt != websocket.BinaryMessage { + t.Errorf("type = %d, want BinaryMessage(%d)", mt, websocket.BinaryMessage) + } + if string(reply) != "hello tunnel" { + t.Errorf("reply = %q, want %q", reply, "hello tunnel") + } +} + +func TestWSConn_MultipleMessages(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + wc := newWSConn(conn) + for i := 0; i < 3; i++ { + buf := make([]byte, 256) + n, err := wc.Read(buf) + if err != nil { + return + } + wc.Write(buf[:n]) + } + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + for i, msg := range []string{"one", "two", "three"} { + conn.WriteMessage(websocket.BinaryMessage, []byte(msg)) + _, reply, err := conn.ReadMessage() + if err != nil { + t.Fatalf("round %d read: %v", i, err) + } + if string(reply) != msg { + t.Errorf("round %d: got %q, want %q", i, reply, msg) + } + } +} + +func TestWSConn_ImplementsNetConn(t *testing.T) { + var _ net.Conn = (*wsConn)(nil) +} + +func TestWSConn_LocalRemoteAddr(t *testing.T) { + addrCh := make(chan [2]net.Addr, 1) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + wc := newWSConn(conn) + addrCh <- [2]net.Addr{wc.LocalAddr(), wc.RemoteAddr()} + wc.Close() + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + select { + case addrs := <-addrCh: + if addrs[0] == nil { + t.Error("LocalAddr should not be nil") + } + if addrs[1] == nil { + t.Error("RemoteAddr should not be nil") + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for addresses") + } +} + +// --- HandleControl --- + +func newGinRouter() *gin.Engine { + r := gin.New() + r.GET("/ws/tai", HandleControl) + r.GET("/ws/tai/data/:channel_id", HandleData) + return r +} + +func TestHandleControl_NoRegistry(t *testing.T) { + registry.SetGlobalForTest(nil) + defer setupTestRegistry() + + restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer test-token"}, + }) + if err == nil { + t.Fatal("expected dial to fail when registry is nil") + } + if resp != nil && resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable) + } +} + +func TestHandleControl_NoAuth(t *testing.T) { + setupTestRegistry() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err == nil { + t.Fatal("expected dial to fail without auth") + } + if resp != nil && resp.StatusCode != http.StatusUnauthorized { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized) + } +} + +func TestHandleControl_AuthFailed(t *testing.T) { + setupTestRegistry() + restore := mockAuth(registry.AuthInfo{}, fmt.Errorf("bad token")) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer bad-token"}, + }) + if err == nil { + t.Fatal("expected dial to fail with bad auth") + } + if resp != nil && resp.StatusCode != http.StatusUnauthorized { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized) + } +} + +func TestHandleControl_RegisterAndPing(t *testing.T) { + reg := setupTestRegistry() + restore := mockAuth(registry.AuthInfo{ + ClientID: "tai-001", + Subject: "user-test", + Scope: "tai:tunnel", + }, nil) + defer restore() + + srv := httptest.NewServer(newGinRouter()) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ + "Authorization": []string{"Bearer valid-token"}, + }) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + + if resp.StatusCode != http.StatusSwitchingProtocols { + t.Errorf("handshake = %d, want 101", resp.StatusCode) + } + + regMsg := registerMessage{ + Type: "register", + TaiID: "tai-001", + MachineID: "m-test", + Version: "2.0", + Ports: map[string]int{"grpc": 9100}, + } + if err := conn.WriteJSON(regMsg); err != nil { + t.Fatalf("write register: %v", err) + } + + var registered map[string]string + if err := conn.ReadJSON(®istered); 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(®istered); 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") + } +} diff --git a/tai/vnc/vnc.go b/tai/vnc/vnc.go index ec03348d..ba36b7bb 100644 --- a/tai/vnc/vnc.go +++ b/tai/vnc/vnc.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "strings" "github.com/yaoapp/yao/tai/sandbox" ) @@ -51,6 +52,29 @@ func (r *remoteVNC) Ping(ctx context.Context, containerID string) error { return nil } +// --- Tunnel implementation --- + +type tunnelVNC struct { + taiID string + yaoBase string // e.g. "http://yao-host:5099" +} + +// NewTunnel creates a VNC that routes through Yao's reverse proxy +// for tunnel-connected Tai instances. +func NewTunnel(taiID, yaoBase string) VNC { + return &tunnelVNC{taiID: taiID, yaoBase: strings.TrimRight(yaoBase, "/")} +} + +func (t *tunnelVNC) URL(_ context.Context, containerID string) (string, error) { + base := strings.Replace(t.yaoBase, "http://", "ws://", 1) + base = strings.Replace(base, "https://", "wss://", 1) + return fmt.Sprintf("%s/tai/%s/vnc/%s/ws", base, t.taiID, containerID), nil +} + +func (t *tunnelVNC) Ping(_ context.Context, _ string) error { + return nil +} + // --- Local implementation --- type localVNC struct {