Files
homelab/k8s/apps/agent-pod/hub-configmap.yaml
T

210 lines
5.9 KiB
YAML
Raw Normal View History

apiVersion: v1
kind: ConfigMap
metadata:
name: hub-src
namespace: agent-pod
data:
go.mod: |
module agent-hub
go 1.25.0
require github.com/gorilla/websocket v1.5.3
go.sum: |
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
main.go: |
// agent-hub sits inside agent-pod alongside pi. Spawned agent runs report
// their lifecycle over localhost HTTP, relaying pi's own `--mode json`
// session protocol (session/agent_start/turn_start/message_start/
// message_update/message_end/turn_end/agent_end -- the same shape Claude
// Code sessions use) verbatim. agent-console (outside the pod, via kubectl
// port-forward) watches all concurrent agents live over one WebSocket and
// can render each session as a real transcript, not a log tail.
package main
import (
"encoding/json"
"log"
"net/http"
"os"
"sync"
"time"
"github.com/gorilla/websocket"
)
type Session struct {
ID string `json:"id"`
Agent string `json:"agent"`
Status string `json:"status"` // running | done | error
Events []json.RawMessage `json:"events"` // raw pi --mode json events, in order
StartedAt time.Time `json:"startedAt"`
EndedAt time.Time `json:"endedAt,omitempty"`
}
type Notification struct {
Type string `json:"type"` // snapshot | start | event | end
Session Session `json:"session"`
}
type Hub struct {
mu sync.Mutex
sessions map[string]*Session
viewers map[*websocket.Conn]struct{}
}
func newHub() *Hub {
return &Hub{
sessions: make(map[string]*Session),
viewers: make(map[*websocket.Conn]struct{}),
}
}
func (h *Hub) broadcast(n Notification) {
h.mu.Lock()
defer h.mu.Unlock()
for conn := range h.viewers {
if err := conn.WriteJSON(n); err != nil {
conn.Close()
delete(h.viewers, conn)
}
}
}
func (h *Hub) start(id, agent string) {
h.mu.Lock()
s := &Session{ID: id, Agent: agent, Status: "running", StartedAt: time.Now()}
h.sessions[id] = s
h.mu.Unlock()
h.broadcast(Notification{Type: "start", Session: *s})
}
func (h *Hub) event(id string, raw json.RawMessage) {
h.mu.Lock()
s, ok := h.sessions[id]
if !ok {
h.mu.Unlock()
return
}
s.Events = append(s.Events, raw)
snapshot := *s
h.mu.Unlock()
h.broadcast(Notification{Type: "event", Session: snapshot})
}
func (h *Hub) end(id, status string) {
h.mu.Lock()
s, ok := h.sessions[id]
if !ok {
h.mu.Unlock()
return
}
s.Status = status
s.EndedAt = time.Now()
snapshot := *s
h.mu.Unlock()
h.broadcast(Notification{Type: "end", Session: snapshot})
}
func (h *Hub) snapshot() []Session {
h.mu.Lock()
defer h.mu.Unlock()
out := make([]Session, 0, len(h.sessions))
for _, s := range h.sessions {
out = append(out, *s)
}
return out
}
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool { return true },
}
func main() {
addr := os.Getenv("HUB_ADDR")
if addr == "" {
addr = ":9090"
}
h := newHub()
http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
})
// Plain JSON snapshot for curl/debugging -- the WebSocket below is the
// live path agent-console actually uses.
http.HandleFunc("/sessions", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(h.snapshot())
})
// Agent side: fire-and-forget POSTs, no persistent connection needed
// since each pi run is short-lived.
http.HandleFunc("/agent/start", func(w http.ResponseWriter, r *http.Request) {
var body struct{ ID, Agent string }
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
h.start(body.ID, body.Agent)
w.WriteHeader(http.StatusOK)
})
http.HandleFunc("/agent/event", func(w http.ResponseWriter, r *http.Request) {
var body struct {
ID string
Event json.RawMessage
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
h.event(body.ID, body.Event)
w.WriteHeader(http.StatusOK)
})
http.HandleFunc("/agent/end", func(w http.ResponseWriter, r *http.Request) {
var body struct{ ID, Status string }
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
h.end(body.ID, body.Status)
w.WriteHeader(http.StatusOK)
})
// Console side: connect, get a snapshot of everything running, then
// stream live start/event/end notifications as they happen.
http.HandleFunc("/console", func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
h.mu.Lock()
for _, s := range h.sessions {
conn.WriteJSON(Notification{Type: "snapshot", Session: *s})
}
h.viewers[conn] = struct{}{}
h.mu.Unlock()
// Drain reads to detect disconnect; console never sends anything.
go func() {
defer func() {
h.mu.Lock()
delete(h.viewers, conn)
h.mu.Unlock()
conn.Close()
}()
for {
if _, _, err := conn.ReadMessage(); err != nil {
return
}
}
}()
})
log.Printf("agent-hub listening on %s", addr)
log.Fatal(http.ListenAndServe(addr, nil))
}