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