daemon: golang WS hub, REST, gitlab, docker spawner, sqlite (18/18 tests)
This commit is contained in:
Raphael Westphal
+621
@@ -0,0 +1,621 @@
|
||||
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"
|
||||
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
|
||||
promptSendTimeout time.Duration = 3 * 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"`
|
||||
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
|
||||
dropped bool
|
||||
|
||||
closeOnce sync.Once
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func (c *webClient) drop() {
|
||||
c.closeOnce.Do(func() {
|
||||
close(c.done)
|
||||
_ = c.conn.Close()
|
||||
})
|
||||
}
|
||||
|
||||
func (c *webClient) deliverControl(b []byte) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.dropped {
|
||||
return
|
||||
}
|
||||
if len(c.control)+len(c.events) >= webMaxPending {
|
||||
c.dropped = true
|
||||
go c.drop()
|
||||
return
|
||||
}
|
||||
c.control = append(c.control, b)
|
||||
}
|
||||
|
||||
func (c *webClient) deliverEvent(e pendingEvent) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.dropped {
|
||||
return
|
||||
}
|
||||
if len(c.control)+len(c.events) >= webMaxPending {
|
||||
c.dropped = true
|
||||
go c.drop()
|
||||
return
|
||||
}
|
||||
c.events = append(c.events, e)
|
||||
}
|
||||
|
||||
// writePump flushes queued frames every webFlushInterval, batching events of
|
||||
// the subscribed session into single frames. Never blocks the hub.
|
||||
func (c *webClient) writePump() {
|
||||
ticker := time.NewTicker(webFlushInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-c.done:
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
c.mu.Lock()
|
||||
control := c.control
|
||||
c.control = nil
|
||||
events := c.events
|
||||
c.events = nil
|
||||
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
|
||||
}
|
||||
}
|
||||
if len(events) == 0 {
|
||||
continue
|
||||
}
|
||||
items := make([]json.RawMessage, 0, len(events))
|
||||
after := events[0].seq - 1
|
||||
for _, e := range events {
|
||||
items = append(items, e.raw)
|
||||
}
|
||||
payload, err := json.Marshal(map[string]any{
|
||||
"type": frameEvents,
|
||||
"sessionId": events[0].sessionID,
|
||||
"after": after,
|
||||
"events": items,
|
||||
})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if err := c.conn.WriteMessage(websocket.TextMessage, payload); err != nil {
|
||||
c.hub.dropWeb(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// readPump consumes subscribe/unsubscribe frames; malformed input never kills the server.
|
||||
func (c *webClient) readPump() {
|
||||
defer c.hub.dropWeb(c)
|
||||
c.conn.SetReadLimit(maxFrameSize)
|
||||
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).
|
||||
func (a *agentConn) writePump() {
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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{}
|
||||
|
||||
// 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{}),
|
||||
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]
|
||||
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)
|
||||
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)
|
||||
}
|
||||
}
|
||||
h.publishEvent(f.sessionID, f.seq, json.RawMessage(f.raw))
|
||||
}
|
||||
|
||||
func (h *Hub) unregister(ac *agentConn) {
|
||||
h.mu.Lock()
|
||||
if cur, ok := h.agents[ac.sessionID]; ok && cur == ac {
|
||||
delete(h.agents, ac.sessionID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
if ac.sessionID != "" {
|
||||
if err := h.store.SetOnline(ac.sessionID, false); err != nil {
|
||||
log.Printf("hub: mark offline %s: %v", ac.sessionID, err)
|
||||
}
|
||||
}
|
||||
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(),
|
||||
})
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user