lvmh/daemon/hub.go

package main

// hub.go — agent websocket hub (/agent/ws) and web websocket hub (/ws).

import (
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// Event types on the agent wire (and control frames on the web wire).
const (
	evHello            string = "hello"
	evWelcome          string = "welcome"
	evMessageStart     string = "message_start"
	evMessageUpdate    string = "message_update"
	evMessageEnd       string = "message_end"
	evToolExecStart    string = "tool_execution_start"
	evToolExecUpdate   string = "tool_execution_update"
	evToolExecEnd      string = "tool_execution_end"
	evAgentStart       string = "agent_start"
	evAgentEnd         string = "agent_end"
	evAgentSettled     string = "agent_settled"
	evSessionInfo      string = "session_info"
	evBye              string = "bye"
	evPrompt           string = "prompt"
	evAbort            string = "abort"
	evSetModel         string = "set_model"
	evRename           string = "rename"
	evErrorNotice      string = "error_notice"
	frameSessionList   string = "session_list"
	frameEvents        string = "events"
	frameSpawnStatus   string = "spawn_status"
	frameSubscribe     string = "subscribe"
	frameUnsubscribe   string = "unsubscribe"
	evPersistedSession string = "session" // payload key of hello/session_info
)

// Wire tuning constants.
const (
	maxFrameSize      int64         = 1 << 20 // 1 MiB per inbound frame
	agentSendQueue    int           = 32
	webSendQueue      int           = 128
	webFlushInterval  time.Duration = 40 * time.Millisecond
	webMaxPending     int           = 4096     // drop slow clients beyond this backlog
	webMaxPendingByte int           = 64 << 20 // ...or beyond this many queued bytes
	promptSendTimeout time.Duration = 3 * time.Second
	pongWait          time.Duration = 60 * time.Second
	pingPeriod        time.Duration = 30 * time.Second
	writeWait         time.Duration = 5 * time.Second
	daemonVersion     int           = 1 // envelope "v"
)

// ErrOffline is returned when a prompt/abort targets a session with no live agent WS.
var ErrOffline = errors.New("session offline")

// frame is one decoded agent wire frame (envelope fields flattened with payload).
type frame struct {
	raw       []byte
	typ       string
	sessionID string
	seq       int64
	ts        int64
}

func decodeFrame(data []byte) (frame, error) {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(data, &fields); err != nil {
		return frame{}, err
	}
	f := frame{raw: data}
	if raw, ok := fields["type"]; ok {
		_ = json.Unmarshal(raw, &f.typ)
	}
	if raw, ok := fields["sessionId"]; ok {
		_ = json.Unmarshal(raw, &f.sessionID)
	}
	if raw, ok := fields["seq"]; ok {
		_ = json.Unmarshal(raw, &f.seq)
	}
	if raw, ok := fields["ts"]; ok {
		_ = json.Unmarshal(raw, &f.ts)
	}
	return f, nil
}

// SessionView is the /api/sessions item shape.
type SessionView struct {
	ID          string  `json:"id"`
	Name        *string `json:"name"`
	Cwd         string  `json:"cwd"`
	Model       string  `json:"model"`
	Provider    string  `json:"provider"`
	Agent       bool    `json:"agent"`
	Repo        *string `json:"repo"`
	StartedAt   int64   `json:"startedAt"`
	Online      bool    `json:"online"`
	Busy        bool    `json:"busy"`
	LastEventAt int64   `json:"lastEventAt"`
}

// pendingEvent is one agent frame queued for a web subscriber.
type pendingEvent struct {
	sessionID string
	seq       int64
	raw       json.RawMessage
}

// webClient is one browser websocket with batched outbound delivery.
type webClient struct {
	hub  *Hub
	conn *websocket.Conn

	mu           sync.Mutex
	sub          string // subscribed sessionId, "" when none
	control      [][]byte
	events       []pendingEvent
	pendingBytes int
	dropped      bool

	closeOnce sync.Once
	done      chan struct{}
}

func (c *webClient) drop() {
	c.closeOnce.Do(func() {
		close(c.done)
		_ = c.conn.Close()
	})
}

// overflows reports whether one more item of size n would exceed the caps.
func (c *webClient) overflows(n int) bool {
	return len(c.control)+len(c.events) >= webMaxPending ||
		c.pendingBytes+n > webMaxPendingByte
}

func (c *webClient) deliverControl(b []byte) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.dropped {
		return
	}
	if c.overflows(len(b)) {
		c.dropped = true
		go c.drop()
		return
	}
	c.control = append(c.control, b)
	c.pendingBytes += len(b)
}

func (c *webClient) deliverEvent(e pendingEvent) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.dropped {
		return
	}
	if c.overflows(len(e.raw)) {
		c.dropped = true
		go c.drop()
		return
	}
	c.events = append(c.events, e)
	c.pendingBytes += len(e.raw)
}

// writeEventsFrame flushes one batched events frame for a single session.
func (c *webClient) writeEventsFrame(batch []pendingEvent) error {
	items := make([]json.RawMessage, 0, len(batch))
	for _, e := range batch {
		items = append(items, e.raw)
	}
	payload, err := json.Marshal(map[string]any{
		"type":      frameEvents,
		"sessionId": batch[0].sessionID,
		"after":     batch[0].seq - 1,
		"events":    items,
	})
	if err != nil {
		return nil // unmarshalable raw JSON cannot happen; drop the batch
	}
	_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
	return c.conn.WriteMessage(websocket.TextMessage, payload)
}

// writePump flushes queued frames every webFlushInterval, batching events of
// the subscribed session into single frames (one frame per contiguous
// sessionID run — a mid-queue resubscribe must never mix sessions).
// Never blocks the hub.
func (c *webClient) writePump() {
	ticker := time.NewTicker(webFlushInterval)
	pings := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	defer pings.Stop()
	for {
		select {
		case <-c.done:
			return
		case <-ticker.C:
		case <-pings.C:
			if !c.sendPing() {
				c.hub.dropWeb(c)
				return
			}
			continue
		}
		c.mu.Lock()
		control := c.control
		c.control = nil
		events := c.events
		c.events = nil
		drained := 0
		for _, b := range control {
			drained += len(b)
		}
		for _, e := range events {
			drained += len(e.raw)
		}
		c.pendingBytes -= drained
		c.mu.Unlock()
		for _, b := range control {
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.TextMessage, b); err != nil {
				c.hub.dropWeb(c)
				return
			}
		}
		for i := 0; i < len(events); {
			j := i
			for j < len(events) && events[j].sessionID == events[i].sessionID {
				j++
			}
			if err := c.writeEventsFrame(events[i:j]); err != nil {
				c.hub.dropWeb(c)
				return
			}
			i = j
		}
	}
}

// sendPing writes one ping frame honoring the write deadline.
func (c *webClient) sendPing() bool {
	_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
	return c.conn.WriteMessage(websocket.PingMessage, nil) == nil
}

// readPump consumes subscribe/unsubscribe frames; malformed input never kills the server.
func (c *webClient) readPump() {
	defer c.hub.dropWeb(c)
	c.conn.SetReadLimit(maxFrameSize)
	_ = c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error {
		return c.conn.SetReadDeadline(time.Now().Add(pongWait))
	})
	for {
		_, data, err := c.conn.ReadMessage()
		if err != nil {
			return
		}
		var msg map[string]any
		if err := json.Unmarshal(data, &msg); err != nil {
			continue
		}
		typ, _ := msg["type"].(string)
		sessionID, _ := msg["sessionId"].(string)
		c.mu.Lock()
		switch typ {
		case frameSubscribe:
			c.sub = sessionID
		case frameUnsubscribe:
			c.sub = ""
		}
		c.mu.Unlock()
	}
}

func (c *webClient) subscription() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.sub
}

// agentConn is one plugin websocket bound to a session after hello.
type agentConn struct {
	hub       *Hub
	conn      *websocket.Conn
	sessionID string
	send      chan []byte
	done      chan struct{}
	closeOnce sync.Once
}

func (a *agentConn) drop() {
	a.closeOnce.Do(func() {
		close(a.done)
		_ = a.conn.Close()
	})
}

// writePump serializes daemon→plugin frames (prompt/abort/welcome) and
// keeps the conn alive with periodic pings.
func (a *agentConn) writePump() {
	pings := time.NewTicker(pingPeriod)
	defer pings.Stop()
	for {
		select {
		case <-a.done:
			return
		case b := <-a.send:
			_ = a.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := a.conn.WriteMessage(websocket.TextMessage, b); err != nil {
				a.drop()
				return
			}
		case <-pings.C:
			if !a.sendPing() {
				a.drop()
				return
			}
		}
	}
}

// sendPing writes one ping frame honoring the write deadline.
func (a *agentConn) sendPing() bool {
	_ = a.conn.SetWriteDeadline(time.Now().Add(writeWait))
	return a.conn.WriteMessage(websocket.PingMessage, nil) == nil
}

// Hub tracks live agent conns and web subscribers; it is the only component
// that mutates online state.
type Hub struct {
	store *Store

	mu     sync.Mutex
	agents map[string]*agentConn
	webs   map[*webClient]struct{}

	// busy tracks sessions mid-turn (agent_start without agent_settled),
	// surfaced on session rows so the UI can show a live activity pulse.
	busy map[string]bool

	// SpawnStatus lets the API push spawn job snapshots to web clients.
	SpawnStatus func() []SpawnJob

	upgrader websocket.Upgrader
}

func NewHub(store *Store) *Hub {
	return &Hub{
		store:  store,
		agents: make(map[string]*agentConn),
		webs:   make(map[*webClient]struct{}),
		busy:   make(map[string]bool),
		upgrader: websocket.Upgrader{
			ReadBufferSize:  4096,
			WriteBufferSize: 4096,
			CheckOrigin:     func(*http.Request) bool { return true },
		},
	}
}

// IsOnline reports whether the session currently has a live agent WS.
func (h *Hub) IsOnline(sessionID string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	_, ok := h.agents[sessionID]
	return ok
}

// SessionsView merges persisted session rows with live online flags.
func (h *Hub) SessionsView() []SessionView {
	rows, err := h.store.Sessions()
	if err != nil {
		log.Printf("hub: list sessions: %v", err)
		return []SessionView{}
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	out := make([]SessionView, 0, len(rows))
	for _, row := range rows {
		v := SessionView{
			ID:          row.Info.ID,
			Name:        row.Info.Name,
			Cwd:         row.Info.Cwd,
			Model:       row.Info.Model,
			Provider:    row.Info.Provider,
			Agent:       row.Info.Agent,
			Repo:        row.Info.Repo,
			StartedAt:   row.Info.StartedAt,
			LastEventAt: row.LastEventAt,
		}
		_, v.Online = h.agents[row.ID]
		v.Busy = h.busy[row.ID]
		out = append(out, v)
	}
	return out
}

func (h *Hub) dropWeb(c *webClient) {
	h.mu.Lock()
	if _, ok := h.webs[c]; !ok {
		h.mu.Unlock()
		return
	}
	delete(h.webs, c)
	h.mu.Unlock()
	c.drop()
	h.BroadcastSessionList()
}

// BroadcastSessionList pushes a full session list to every web client.
func (h *Hub) BroadcastSessionList() {
	payload, err := json.Marshal(map[string]any{
		"type":     frameSessionList,
		"sessions": h.SessionsView(),
	})
	if err != nil {
		return
	}
	h.mu.Lock()
	clients := make([]*webClient, 0, len(h.webs))
	for c := range h.webs {
		clients = append(clients, c)
	}
	h.mu.Unlock()
	for _, c := range clients {
		c.deliverControl(payload)
	}
}

// BroadcastSpawnStatus pushes the spawn job snapshot (if a provider is set).
func (h *Hub) BroadcastSpawnStatus() {
	if h.SpawnStatus == nil {
		return
	}
	jobs := h.SpawnStatus()
	payload, err := json.Marshal(map[string]any{
		"type": frameSpawnStatus,
		"jobs": jobs,
	})
	if err != nil {
		return
	}
	h.mu.Lock()
	clients := make([]*webClient, 0, len(h.webs))
	for c := range h.webs {
		clients = append(clients, c)
	}
	h.mu.Unlock()
	for _, c := range clients {
		c.deliverControl(payload)
	}
}

// publishEvent fans one agent frame out to web subscribers of that session.
func (h *Hub) publishEvent(sessionID string, seq int64, raw json.RawMessage) {
	h.mu.Lock()
	clients := make([]*webClient, 0, len(h.webs))
	for c := range h.webs {
		clients = append(clients, c)
	}
	h.mu.Unlock()
	ev := pendingEvent{sessionID: sessionID, seq: seq, raw: raw}
	for _, c := range clients {
		if c.subscription() == sessionID {
			c.deliverEvent(ev)
		}
	}
}

// ServeAgentWS handles GET /agent/ws (bearer auth already enforced by middleware).
func (h *Hub) ServeAgentWS(w http.ResponseWriter, r *http.Request) {
	ws, err := h.upgrader.Upgrade(w, r, nil)
	if err != nil {
		return // Upgrade already wrote the HTTP error
	}
	ac := &agentConn{hub: h, conn: ws, send: make(chan []byte, agentSendQueue), done: make(chan struct{})}
	go ac.writePump()
	h.readAgentLoop(ac)
}

func (h *Hub) readAgentLoop(ac *agentConn) {
	defer ac.drop()
	ac.conn.SetReadLimit(maxFrameSize)
	_ = ac.conn.SetReadDeadline(time.Now().Add(pongWait))
	ac.conn.SetPongHandler(func(string) error {
		return ac.conn.SetReadDeadline(time.Now().Add(pongWait))
	})
	registered := false
	defer func() {
		if registered {
			h.unregister(ac)
		}
	}()
	for {
		_, data, err := ac.conn.ReadMessage()
		if err != nil {
			return
		}
		f, err := decodeFrame(data)
		if err != nil {
			log.Printf("hub: malformed frame from agent: %v", err)
			continue // never kill the conn loop on a bad frame
		}
		switch f.typ {
		case evHello:
			registered = true
			h.handleHello(ac, f)
		case evSessionInfo:
			h.handleSessionInfo(f)
		case evBye:
			return
		default:
			h.handleEvent(f)
		}
	}
}

// handleHello persists the session snapshot, replaces any older conn, replies welcome.
func (h *Hub) handleHello(ac *agentConn, f frame) {
	var fields struct {
		Session SessionInfo `json:"session"`
	}
	if err := json.Unmarshal(f.raw, &fields); err != nil || fields.Session.ID == "" {
		log.Printf("hub: hello without session payload from %s", f.sessionID)
		return
	}
	sessionID := f.sessionID
	if sessionID == "" {
		sessionID = fields.Session.ID
	}
	fields.Session.ID = sessionID
	ac.sessionID = sessionID

	h.mu.Lock()
	if old, ok := h.agents[sessionID]; ok && old != ac {
		old.drop() // reconnect: new conn replaces old
	}
	h.agents[sessionID] = ac
	h.mu.Unlock()

	if err := h.store.UpsertSession(fields.Session); err != nil {
		log.Printf("hub: persist session %s: %v", sessionID, err)
	}
	if err := h.store.SetOnline(sessionID, true); err != nil {
		log.Printf("hub: mark online %s: %v", sessionID, err)
	}
	lastSeq, err := h.store.LastSeq(sessionID)
	if err != nil {
		log.Printf("hub: lastSeq %s: %v", sessionID, err)
		lastSeq = 0
	}
	welcome, err := json.Marshal(map[string]any{
		"v":         daemonVersion,
		"type":      evWelcome,
		"sessionId": sessionID,
		"seq":       0,
		"ts":        time.Now().UnixMilli(),
		"lastSeq":   lastSeq,
	})
	if err != nil {
		return
	}
	select {
	case ac.send <- welcome:
	default:
		ac.drop()
	}
	h.BroadcastSessionList()
}

func (h *Hub) handleSessionInfo(f frame) {
	var fields struct {
		Session SessionInfo `json:"session"`
	}
	if err := json.Unmarshal(f.raw, &fields); err != nil || fields.Session.ID == "" {
		return
	}
	sessionID := f.sessionID
	if sessionID == "" {
		sessionID = fields.Session.ID
	}
	fields.Session.ID = sessionID
	if err := h.store.UpsertSession(fields.Session); err != nil {
		log.Printf("hub: persist session_info %s: %v", sessionID, err)
	}
	h.BroadcastSessionList()
}

// handleEvent persists (except message_update), updates bookkeeping, fans out live.
func (h *Hub) handleEvent(f frame) {
	if f.sessionID == "" {
		return
	}
	if f.typ != evMessageUpdate {
		if err := h.store.AppendEvent(Event{
			SessionID: f.sessionID,
			Seq:       f.seq,
			TS:        f.ts,
			Type:      f.typ,
			Payload:   json.RawMessage(f.raw),
		}); err != nil {
			log.Printf("hub: persist event %s#%d: %v", f.sessionID, f.seq, err)
		}
		if err := h.store.TouchSession(f.sessionID, f.seq, f.ts); err != nil {
			log.Printf("hub: touch session %s: %v", f.sessionID, err)
		}
	}
	// Track mid-turn state for the session-list activity pulse.
	if f.typ == evAgentStart || f.typ == evAgentSettled {
		busy := f.typ == evAgentStart
		h.mu.Lock()
		if h.busy[f.sessionID] != busy {
			h.busy[f.sessionID] = busy
			h.mu.Unlock()
			go h.BroadcastSessionList()
		} else {
			h.mu.Unlock()
		}
	}
	h.publishEvent(f.sessionID, f.seq, json.RawMessage(f.raw))
}

func (h *Hub) unregister(ac *agentConn) {
	h.mu.Lock()
	removed := false
	if cur, ok := h.agents[ac.sessionID]; ok && cur == ac {
		delete(h.agents, ac.sessionID)
		removed = true
	}
	h.mu.Unlock()
	// Only flip the persisted flag when no newer conn replaced this one;
	// a reconnect races this unregister path.
	if removed && ac.sessionID != "" {
		if err := h.store.SetOnline(ac.sessionID, false); err != nil {
			log.Printf("hub: mark offline %s: %v", ac.sessionID, err)
		}
		h.mu.Lock()
		delete(h.busy, ac.sessionID)
		h.mu.Unlock()
	}
	h.BroadcastSessionList()
}

// Prompt routes a prompt frame to the live plugin conn for sessionID.
func (h *Hub) Prompt(sessionID, message string) error {
	return h.sendToAgent(sessionID, map[string]any{
		"v":         daemonVersion,
		"type":      evPrompt,
		"sessionId": sessionID,
		"ts":        time.Now().UnixMilli(),
		"promptId":  newUUID(),
		"message":   message,
	})
}

// Abort routes an abort frame to the live plugin conn for sessionID.
func (h *Hub) Abort(sessionID string) error {
	return h.sendToAgent(sessionID, map[string]any{
		"v":         daemonVersion,
		"type":      evAbort,
		"sessionId": sessionID,
		"ts":        time.Now().UnixMilli(),
	})
}

// SetModel routes a set_model frame to the live plugin conn for sessionID.
func (h *Hub) SetModel(sessionID, provider, modelID string) error {
	return h.sendToAgent(sessionID, map[string]any{
		"v":         daemonVersion,
		"type":      evSetModel,
		"sessionId": sessionID,
		"ts":        time.Now().UnixMilli(),
		"provider":  provider,
		"modelId":   modelID,
	})
}

// Rename routes a rename frame to the live plugin conn for sessionID.
func (h *Hub) Rename(sessionID, name string) error {
	return h.sendToAgent(sessionID, map[string]any{
		"v":         daemonVersion,
		"type":      evRename,
		"sessionId": sessionID,
		"ts":        time.Now().UnixMilli(),
		"name":      name,
	})
}

func (h *Hub) sendToAgent(sessionID string, frameBody map[string]any) error {
	b, err := json.Marshal(frameBody)
	if err != nil {
		return err
	}
	h.mu.Lock()
	ac := h.agents[sessionID]
	h.mu.Unlock()
	if ac == nil {
		return ErrOffline
	}
	select {
	case ac.send <- b:
		return nil
	case <-ac.done:
		return ErrOffline
	case <-time.After(promptSendTimeout):
		return ErrOffline
	}
}

// ServeWebWS handles GET /ws?token=... — token arrives as query param.
func (h *Hub) ServeWebWS(w http.ResponseWriter, r *http.Request) {
	ws, err := h.upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	c := &webClient{hub: h, conn: ws, done: make(chan struct{})}
	h.mu.Lock()
	h.webs[c] = struct{}{}
	h.mu.Unlock()
	go c.writePump()
	go c.readPump()

	payload, err := json.Marshal(map[string]any{
		"type":     frameSessionList,
		"sessions": h.SessionsView(),
	})
	if err == nil {
		c.deliverControl(payload)
	}
}

// OnlineCount reports how many agent WS connections are currently live.
func (h *Hub) OnlineCount() int {
	h.mu.Lock()
	defer h.mu.Unlock()
	return len(h.agents)
}