yao/tai/docs/registry.md

# 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

```go
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()`.

```go
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

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

### SystemInfo

```go
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

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

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

### Register / Unregister

```go
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

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

### Heartbeat

```go
func (r *Registry) UpdatePing(taiID string)
```

### Health Check

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

Runs a background goroutine that:
1. Marks direct-mode nodes as `"offline"` if `LastPing` exceeds `timeout`
2. 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

```go
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

```go
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

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

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

### OpenLocalListener

```go
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`.