// Package ws is the WebSocket transport to apps/server. It owns a reconnecting // connection and exposes decoded server messages plus connection-status changes // over channels the Bubble Tea loop drains via commands. package ws import ( "context" "net/url" "os" "time" "github.com/correx/tui-go/internal/protocol" "github.com/gorilla/websocket" ) const ( initialBackoff = 500 * time.Millisecond maxBackoff = 8 * time.Second ) // Status reports connection lifecycle transitions. type Status struct { Connected bool Reconnecting bool Attempt int } // Client is a reconnecting WebSocket client. Construct with New, then Run in a // goroutine. Incoming and Conn are drained by the UI. type Client struct { url string httpBase string send chan []byte in chan protocol.ServerMessage conn chan Status } // New builds a client targeting ws://host:port/stream. func New(host string, port int) *Client { hp := hostPort(host, port) u := url.URL{Scheme: "ws", Host: hp, Path: "/stream"} return &Client{ url: u.String(), httpBase: (&url.URL{Scheme: "http", Host: hp}).String(), send: make(chan []byte, 64), in: make(chan protocol.ServerMessage, 256), conn: make(chan Status, 16), } } // HTTPBase is the http://host:port origin the WS client connects to, for REST // calls (e.g. GET /sessions) that share the server's host/port. Empty for a nil // client (the test harness constructs Models with a nil client). func (c *Client) HTTPBase() string { if c == nil { return "" } return c.httpBase } func hostPort(host string, port int) string { h := host if h == "" { h = "localhost" } return h + ":" + itoa(port) } func itoa(n int) string { if n == 0 { return "0" } var b [20]byte i := len(b) for n > 0 { i-- b[i] = byte('0' + n%10) n /= 10 } return string(b[i:]) } // Incoming yields decoded server messages. func (c *Client) Incoming() <-chan protocol.ServerMessage { return c.in } // Conn yields connection-status changes. func (c *Client) Conn() <-chan Status { return c.conn } // Send queues a raw client frame. Non-blocking; drops if the buffer is full // (the server will be resynced on reconnect via snapshot). func (c *Client) Send(frame []byte) { select { case c.send <- frame: default: } } // Drain returns all queued-but-unsent client frames. Test-only: meaningful before // Run starts consuming the send channel, so a test can assert on what was Sent. func (c *Client) Drain() [][]byte { var out [][]byte for { select { case f := <-c.send: out = append(out, f) default: return out } } } // Run drives the reconnect loop until ctx is cancelled. func (c *Client) Run(ctx context.Context) { backoff := initialBackoff attempt := 0 for { if ctx.Err() != nil { return } conn, _, err := websocket.DefaultDialer.DialContext(ctx, c.url, nil) if err != nil { attempt++ c.emit(Status{Reconnecting: true, Attempt: attempt}) if !sleep(ctx, backoff) { return } backoff = min(backoff*2, maxBackoff) continue } backoff = initialBackoff attempt = 0 c.emit(Status{Connected: true}) if cwd, err := os.Getwd(); err == nil { _ = conn.WriteMessage(websocket.TextMessage, protocol.Hello(cwd)) } c.pump(ctx, conn) c.emit(Status{Reconnecting: true}) } } // pump runs the read/write loops for a single live connection and returns when // it drops. func (c *Client) pump(ctx context.Context, conn *websocket.Conn) { connCtx, cancel := context.WithCancel(ctx) defer cancel() go func() { for { select { case <-connCtx.Done(): return case frame := <-c.send: if err := conn.WriteMessage(websocket.TextMessage, frame); err != nil { cancel() return } } } }() for { _, raw, err := conn.ReadMessage() if err != nil { cancel() _ = conn.Close() return } if msg, derr := protocol.Decode(raw); derr == nil { select { case c.in <- msg: case <-connCtx.Done(): return } } } } func (c *Client) emit(s Status) { select { case c.conn <- s: default: } } func sleep(ctx context.Context, d time.Duration) bool { t := time.NewTimer(d) defer t.Stop() select { case <-ctx.Done(): return false case <-t.C: return true } }