llmgw/internal/agent/agent_test.go

package agent

import (
	"context"
	"crypto/tls"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"sync"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
)

// statusRecorder captures what the control plane hears about shards, and what
// it is asked to adopt.
type statusRecorder struct {
	ch chan status

	mu      sync.Mutex
	adopted []reconciler.AdoptedShard
	unclaim []string // shard ids to report back as unclaimed
}

type status struct {
	replicaID, shardID string
	state              supervisor.State
	err                string
}

func newRecorder() *statusRecorder { return &statusRecorder{ch: make(chan status, 128)} }

func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
	select {
	case r.ch <- status{replicaID, shardID, st, errMsg}:
	default:
	}
}

func (r *statusRecorder) Adopt(shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.adopted = append(r.adopted, shards...)

	var unclaimed []reconciler.AdoptedShard
	for _, s := range shards {
		if slices.Contains(r.unclaim, s.ShardID) {
			unclaimed = append(unclaimed, s)
		}
	}
	return unclaimed
}

func (r *statusRecorder) adoptedShards() []reconciler.AdoptedShard {
	r.mu.Lock()
	defer r.mu.Unlock()
	return slices.Clone(r.adopted)
}

func (r *statusRecorder) markUnclaimed(ids ...string) {
	r.mu.Lock()
	r.unclaim = append(r.unclaim, ids...)
	r.mu.Unlock()
}

func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status {
	t.Helper()
	deadline := time.After(limit)
	var seen []status
	for {
		select {
		case s := <-r.ch:
			seen = append(seen, s)
			if pred(s) {
				return s
			}
		case <-deadline:
			t.Fatalf("no matching status within %v; saw %+v", limit, seen)
			return status{}
		}
	}
}

// fakeEngineScript writes a script that stands in for vLLM.
//
// It parses the real --port flag the vLLM adapter generates rather than being
// told out of band, so the adapter's argument construction is exercised too.
func fakeEngineScript(t *testing.T, body string) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "fake-vllm")
	script := `#!/bin/sh
port=""
while [ $# -gt 0 ]; do
  case "$1" in
    --port) port="$2"; shift 2 ;;
    *) shift ;;
  esac
done
` + body + "\n"
	if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
		t.Fatal(err)
	}
	return path
}

// servingEngine answers /health and /v1/models, which is what readiness needs.
func servingEngine(t *testing.T) string {
	t.Helper()
	py, err := exec.LookPath("python3")
	if err != nil {
		t.Skip("python3 is required for the fake engine")
	}
	return fakeEngineScript(t, `exec `+py+` -c '
import http.server, json, sys
port = int(sys.argv[1])
class H(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/health":
            self.send_response(200); self.end_headers(); return
        if self.path == "/v1/models":
            b = json.dumps({"object":"list","data":[{"id":"test-model"}]}).encode()
            self.send_response(200)
            self.send_header("content-type","application/json")
            self.send_header("content-length", str(len(b)))
            self.end_headers(); self.wfile.write(b); return
        self.send_response(404); self.end_headers()
    def log_message(self, *a): pass
http.server.HTTPServer(("127.0.0.1", port), H).serve_forever()
' "$port"`)
}

type harness struct {
	svc *agentsvc.Service
	inv *inventory.Store
	rec *statusRecorder
	srv *httptest.Server
	ag  *Agent
	ln  *dropListener
	// streams counts how many times an agent has opened the Connect stream,
	// which is the only honest way to tell a reconnect from a stale entry.
	streams *atomic.Int64
}

// dropListener tracks accepted connections so a test can sever them.
//
// httptest's CloseClientConnections does not work here: h2c with prior
// knowledge hijacks the connection, and httptest stops tracking hijacked
// connections. A test relying on it would pass without ever reconnecting.
type dropListener struct {
	net.Listener
	mu    sync.Mutex
	conns []net.Conn
}

func (l *dropListener) Accept() (net.Conn, error) {
	c, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	l.mu.Lock()
	l.conns = append(l.conns, c)
	l.mu.Unlock()
	return c, nil
}

// drop severs every connection, as a control plane restart would.
func (l *dropListener) drop() {
	l.mu.Lock()
	conns := l.conns
	l.conns = nil
	l.mu.Unlock()
	for _, c := range conns {
		_ = c.Close()
	}
}

// start wires a real agent to a real service over h2c, so this exercises the
// protocol rather than two structs calling each other.
func start(t *testing.T, cfgFns ...func(*Config)) *harness {
	t.Helper()
	return startWith(t, slog.New(slog.DiscardHandler), cfgFns...)
}

// startLogged is start with logging on, for diagnosing protocol behaviour.
func startLogged(t *testing.T, cfgFns ...func(*Config)) *harness {
	t.Helper()
	return startWith(t, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug})), cfgFns...)
}

func startWith(t *testing.T, log *slog.Logger, cfgFns ...func(*Config)) *harness {
	t.Helper()

	inv := inventory.NewStore()
	rec := newRecorder()
	svc := agentsvc.New(inv, rec, log)

	path, handler := svc.Handler()
	mux := http.NewServeMux()
	mux.Handle(path, handler)

	var streams atomic.Int64
	counting := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		streams.Add(1)
		mux.ServeHTTP(w, r)
	})

	raw, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	ln := &dropListener{Listener: raw}
	// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
	srv := httptest.NewUnstartedServer(h2c.NewHandler(counting, &http2.Server{}))
	srv.Listener = ln
	srv.Start()

	cfg := DefaultConfig()
	cfg.NodeID = "n1"
	cfg.Server = srv.URL
	cfg.MgmtAddr = "127.0.0.1"
	cfg.HeartbeatInterval = 50 * time.Millisecond
	cfg.ReconnectBase = 20 * time.Millisecond
	cfg.ReconnectMax = 100 * time.Millisecond
	cfg.Discovery.SkipGPUs = true
	cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
	cfg.Superv = supervisor.Config{
		WorkDir: t.TempDir(), LogLines: 64,
		StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
	}
	for _, f := range cfgFns {
		f(&cfg)
	}

	client := &http.Client{Transport: &http2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, addr)
		},
	}}

	ag, err := New(cfg, client, log)
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() { _ = ag.Run(ctx) }()

	h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag, ln: ln, streams: &streams}
	t.Cleanup(func() {
		cancel()
		ag.Supervisor().StopAll(context.Background())
		srv.Close()
	})
	h.waitRegistered(t)
	return h
}

func (h *harness) waitRegistered(t *testing.T) {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if _, ok := h.inv.Get("n1"); ok && len(h.svc.ConnectedNodes()) == 1 {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("the agent never registered")
}

// shardCommand builds a single-node vLLM shard on a free port.
func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
	t.Helper()
	return reconciler.ShardCommand{
		ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1",
		Spec: engines.ShardSpec{
			ReplicaID: "rep-1", ShardID: shardID, Rank: 0, NNodes: 1,
			GPUIndices: []int{0}, HTTPPort: freePort(t),
			Profile: engines.Profile{
				Name: "p", Engine: "vllm", TP: 1,
				ModelPath: "/models/x", ServedModelName: "test-model",
			},
		},
		StartupTimeout: 15 * time.Second,
	}
}

func freePort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port
}

func TestRegistrationCarriesInventory(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Hostname = "gpu-a-01"
		c.Labels = map[string]string{"rack": "r12"}
		c.Fabrics = []inventory.FabricAttachment{
			{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200},
		}
	})

	n, ok := h.inv.Get("n1")
	if !ok {
		t.Fatal("node missing from inventory")
	}
	if n.Hostname != "gpu-a-01" || n.MgmtAddr != "127.0.0.1" {
		t.Errorf("identity wrong: %+v", n)
	}
	if n.Labels["rack"] != "r12" {
		t.Errorf("labels lost: %v", n.Labels)
	}
	// The fabric label cannot be inferred from the wire, so losing it in
	// transit would silently disable multi-node placement.
	if !n.HasFabric("ib-a") {
		t.Fatalf("fabric attachment lost: %+v", n.Fabrics)
	}
	if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
		t.Errorf("HCAs=%v", got)
	}
}

func TestHeartbeatKeepsNodeAlive(t *testing.T) {
	h := start(t)

	n, _ := h.inv.Get("n1")
	first := n.LastSeen

	time.Sleep(200 * time.Millisecond)
	n, _ = h.inv.Get("n1")
	if !n.LastSeen.After(first) {
		t.Error("heartbeats are not refreshing LastSeen; the node would be declared lost")
	}
}

func TestStartAndStopShardEndToEnd(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = servingEngine(t)
	})

	cmd := shardCommand(t, "shard-a")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}

	got := h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Ready
	})
	if got.replicaID != "rep-1" {
		t.Errorf("replica id lost in transit: %q", got.replicaID)
	}

	if err := h.svc.StopShard(t.Context(), "n1", "shard-a"); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Stopped
	})
}

func TestFailedShardCarriesDiagnosis(t *testing.T) {
	// When a model fails to load the engine's output is the only diagnosis, so
	// it has to survive the trip to the control plane.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `echo 'CUDA out of memory' >&2; exit 1`)
	})

	if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-bad")); err != nil {
		t.Fatal(err)
	}
	got := h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-bad" && s.state == supervisor.Failed
	})
	if got.err == "" {
		t.Error("the failure carried no error message")
	}
}

func TestEngineThatNeverBecomesReady(t *testing.T) {
	// A process that runs but never serves must time out rather than hang the
	// replica forever.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})

	cmd := shardCommand(t, "shard-stuck")
	cmd.StartupTimeout = 500 * time.Millisecond
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 20*time.Second, func(s status) bool {
		return s.shardID == "shard-stuck" && s.state == supervisor.Failed
	})
}

func TestValidationRejectsBeforeLaunch(t *testing.T) {
	// An engine that refuses a profile does so after minutes of loading;
	// catching it at dispatch costs nothing.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = servingEngine(t)
	})

	cmd := shardCommand(t, "shard-invalid")
	cmd.Spec.Profile.ModelPath = "" // validation requires one

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	got := h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-invalid" && s.state == supervisor.Failed
	})
	if got.err == "" {
		t.Error("validation failure carried no reason")
	}
}

func TestDispatchToUnconnectedNode(t *testing.T) {
	h := start(t)

	err := h.svc.StartShard(t.Context(), reconciler.ShardCommand{NodeID: "nowhere", ShardID: "x"})
	if err == nil {
		t.Fatal("dispatching to an unconnected node succeeded")
	}
	if !errors.Is(err, agentsvc.ErrNotConnected) {
		t.Errorf("err=%v, expected ErrNotConnected", err)
	}
}

func TestAgentReconnects(t *testing.T) {
	// A control plane restart must not require touching every GPU node.
	h := start(t)

	before := h.streams.Load()
	h.ln.drop()

	// Counting streams, not connected nodes: a stale entry in the connection
	// map would make a node look connected without any reconnect having
	// happened, and an assertion on it would pass vacuously.
	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if h.streams.Load() > before && len(h.svc.ConnectedNodes()) == 1 {
			if err := h.svc.DrainShard(t.Context(), "n1", "nothing"); err != nil {
				t.Fatalf("the reconnected stream does not accept commands: %v", err)
			}
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("the agent never reconnected (streams %d -> %d)", before, h.streams.Load())
}

func TestIdempotentStart(t *testing.T) {
	// The control plane retries commands; a repeat must not start a second
	// process (plan §5.6).
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 30`)
	})

	cmd := shardCommand(t, "shard-idem")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-idem" && s.state == supervisor.Starting
	})

	sup := h.ag.Supervisor()
	sh, err := sup.Get("shard-idem")
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	time.Sleep(400 * time.Millisecond)

	sh2, err := sup.Get("shard-idem")
	if err != nil {
		t.Fatal(err)
	}
	if sh2.Status().PID != pid {
		t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID)
	}
}

// TestRunningShardsReportedOnReconnect covers the control plane restart hole
// (plan §6.5): a reconnecting agent has to describe what it is already running,
// or the control plane concludes the cluster is empty and places duplicates
// while the originals keep holding GPUs.
func TestRunningShardsReportedOnReconnect(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})

	cmd := shardCommand(t, "shard-live")
	cmd.Pool = "ib-a/h100-80g"
	cmd.Fabric = "ib-a"
	cmd.Spec.HCAs = []string{"mlx5_0"}

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-live" && s.state == supervisor.Starting
	})

	// Drop the stream, as a control plane restart would.
	h.ln.drop()

	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if got := h.rec.adoptedShards(); len(got) > 0 {
			a := got[0]
			if a.ShardID != "shard-live" || a.ReplicaID != "rep-1" {
				t.Fatalf("wrong shard reported: %+v", a)
			}
			// Everything the control plane needs to rebuild the replica.
			if a.Model != "test-model" || a.Profile != "p" {
				t.Errorf("catalog identity lost: model=%q profile=%q", a.Model, a.Profile)
			}
			if a.Pool != "ib-a/h100-80g" || a.Fabric != "ib-a" {
				t.Errorf("pool/fabric lost: %q %q", a.Pool, a.Fabric)
			}
			if !slices.Equal(a.HCAs, []string{"mlx5_0"}) {
				t.Errorf("HCAs lost: %v", a.HCAs)
			}
			if a.Endpoint == "" {
				t.Error("rank 0 reported no endpoint; the router would have nowhere to send")
			}
			if len(a.GPUIndices) == 0 {
				t.Error("GPU indices lost; the inventory could not reclaim them")
			}
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("the reconnecting agent never reported its running shard")
}

// TestUnclaimedShardIsStopped covers the other half: a shard whose model left
// the catalog is holding GPUs nothing will ever route to.
func TestUnclaimedShardIsStopped(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})

	cmd := shardCommand(t, "shard-orphaned")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-orphaned" && s.state == supervisor.Starting
	})

	sh, err := h.ag.Supervisor().Get("shard-orphaned")
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	// The control plane will refuse to claim it on reconnect.
	h.rec.markUnclaimed("shard-orphaned")
	h.ln.drop()

	deadline := time.Now().Add(20 * time.Second)
	for time.Now().Before(deadline) {
		if !processAliveForTest(pid) {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Error("an unclaimed shard kept running and holding its GPUs")
}

// processAliveForTest mirrors the supervisor's liveness check.
func processAliveForTest(pid int) bool {
	return syscall.Kill(pid, 0) == nil
}