Max e633640998 refactor(sandbox/v2): transition from pool to node configuration

- Updated benchmark and test functions to utilize node configurations instead of pool configurations for improved clarity and consistency.
- Refactored related setup functions and test cases to align with the new node-based architecture.
- Adjusted error messages and documentation to reflect the transition from pool to node terminology.

Made-with: Cursor

2026-03-09 03:20:09 +08:00

5.2 KiB

Raw Blame History

Package `registry`

In-memory registry for Tai nodes. Manages both direct (network-reachable) and tunnel (WebSocket-bridged) connections. Used server-side by Yao to track all connected Tai instances.

Architecture

Direct Mode:     Yao  ── TCP ──>  Tai (gRPC/HTTP/Docker/VNC)
Tunnel Mode:     Yao  <── WS ──  Tai (control channel)
                 Yao  <── WS ──  Tai (data channels, on-demand)

Types

TaiNode

type TaiNode struct {
    TaiID        string
    MachineID    string
    Version      string
    Auth         AuthInfo
    System       SystemInfo
    Mode         string            // "direct" | "tunnel"
    Addr         string            // direct: "tai-host"; tunnel: empty
    YaoBase      string            // tunnel: Yao server base URL
    Ports        map[string]int    // {"grpc":19100, "http":8099, ...}
    Capabilities map[string]bool   // {"docker":true, "host_exec":true}
    ControlConn  *websocket.Conn   // tunnel: WS control channel
    Status       string            // "online" | "offline" | "connecting"
    ConnectedAt  time.Time
    LastPing     time.Time
    DisplayName  string
}

NodeSnapshot

Read-only copy of TaiNode safe to use outside locks. Returned by Get() and List().

type NodeSnapshot struct {
    TaiID, MachineID, Version string
    Auth         AuthInfo
    System       SystemInfo
    Mode, Addr, YaoBase string
    Ports        map[string]int
    Capabilities map[string]bool
    Status       string
    ConnectedAt, LastPing time.Time
    DisplayName  string
}

AuthInfo

type AuthInfo struct {
    Subject  string
    UserID   string
    ClientID string
    Scope    string
    TeamID   string
    TenantID string
}

SystemInfo

type SystemInfo struct {
    OS       string `json:"os"`
    Arch     string `json:"arch"`
    Hostname string `json:"hostname"`
    NumCPU   int    `json:"num_cpu"`
    TotalMem int64  `json:"total_mem,omitempty"`
}

Registry API

Init / Global

func Init(logger *slog.Logger)
func Global() *Registry

Init creates the global singleton (once). Global returns it (nil before Init).

Register / Unregister

func (r *Registry) Register(node *TaiNode)
func (r *Registry) Unregister(taiID string)

Register adds or replaces a node, setting Status="online" and recording timestamps. Unregister closes all tunnel listeners and the control connection, then removes the node.

Query

func (r *Registry) Get(taiID string) (*NodeSnapshot, bool)
func (r *Registry) List() []NodeSnapshot
func (r *Registry) ListByTeam(teamID string) []NodeSnapshot

Heartbeat

func (r *Registry) UpdatePing(taiID string)

Health Check

func (r *Registry) StartHealthCheck(done <-chan struct{}, interval, timeout, cleanupAfter time.Duration)

Runs a background goroutine that:

Marks direct-mode nodes as "offline" if LastPing exceeds timeout
Auto-unregisters nodes that stay offline longer than timeout + cleanupAfter

Tunnel API

For tunnel-mode nodes, the registry manages on-demand TCP-over-WebSocket channels.

RequestChannel

func (r *Registry) RequestChannel(taiID string, targetPort int) (channelID string, result chan net.Conn, err error)

Sends an {"type":"open", "channel_id":"...", "target_port":...} command to the node's control WebSocket. Returns a channel that receives the net.Conn when Tai connects back with the data channel. Times out after 30 seconds.

AcceptDataChannel

func (r *Registry) AcceptDataChannel(channelID, taiID string, conn net.Conn) error

Called when Tai establishes a data WebSocket for a pending channel. Validates taiID ownership and delivers the connection to the waiting RequestChannel caller.

WriteControlJSON

func (r *Registry) WriteControlJSON(taiID string, v interface{}) error

Thread-safe JSON write to a node's control WebSocket.

OpenLocalListener

func (r *Registry) OpenLocalListener(taiID string, targetPort int) (net.Listener, error)

Creates a 127.0.0.1:0 TCP listener. Every accepted connection is automatically bridged through the tunnel to targetPort on the Tai node. Returns the listener so the caller can read ln.Addr() to get the ephemeral port.

Connection Flow (Tunnel)

1. Tai → Yao: WebSocket upgrade to GET /ws/tai (Bearer auth)
2. Tai → Yao: sends {"type":"register", "tai_id":"xxx", ...} on WS
3. Yao: Register(node) with Mode="tunnel", ControlConn=ws
4. Yao → Tai: sends {"type":"registered", "tai_id":"xxx"}
5. Client → Yao: tai.New("tunnel://tai-abc123")
6. Yao: OpenLocalListener("tai-abc123", 19100) → 127.0.0.1:54321
7. Yao: grpc.Dial("passthrough:///127.0.0.1:54321") → triggers accept
8. Yao: RequestChannel("tai-abc123", 19100) → sends {"type":"open"} on control WS
9. Tai: receives "open", dials localhost:19100, connects data WS to GET /ws/tai/data/:channel_id
10. Yao: AcceptDataChannel(channelID, taiID, conn) → bridges local TCP ↔ data WS
11. gRPC traffic flows transparently through the tunnel

Keep-alive

Tai sends {"type":"ping"} periodically on the control channel. Yao replies {"type":"pong"} and updates LastPing.

5.2 KiB Raw Blame History

Package registry