correx/apps/tui-go/internal/ws/client.go

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