llmgw/internal/control/agentsvc/service.go

package agentsvc

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"maps"
	"net/http"
	"slices"
	"sort"
	"sync"
	"time"

	"connectrpc.com/connect"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
	"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
)

// StatusSink receives shard transitions and recovers replica state from what
// agents report they are already running. The reconciler implements it.
type StatusSink interface {
	OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string)
	// OnFetchProgress is separate because it is not a transition: the shard is
	// still starting, and treating each report as a state change would reset
	// the timers that decide when a slow start has become a failed one.
	OnFetchProgress(replicaID, shardID, source string, done, total int64)
	// OnShardCompute records where a shard's engine said it runs the model.
	OnShardCompute(replicaID, shardID string, c reconciler.Compute)
	Adopt(nodeID string, shards []reconciler.AdoptedShard) (unclaimed []reconciler.AdoptedShard)
}

// Service is the control plane's AgentService implementation, and at the same
// time the reconciler's Dispatcher: it holds the open streams, so it is the
// only thing that can deliver a command to a node.
type Service struct {
	// leader gates whether this process accepts agent streams at all.
	leader Leader

	inv    *inventory.Store
	status StatusSink
	log    *slog.Logger

	// requireIdentity refuses agents that present no certificate. It is off
	// only for the insecure development path.
	requireIdentity bool
	// revokedAt is when a node was removed, whose certificates issued before
	// then are refused; nil checks nothing.
	revokedAt func(ctx context.Context, nodeID string) (time.Time, error)

	mu    sync.RWMutex
	conns map[string]*conn
	// kv receives the KV cache events agents forward; see SetKVSink.
	kv func(kvevents.Update)

	// enginesFor is what a node is told to install, set once at start: the
	// cluster's engines, with any override that picks the node laid over.
	enginesFor func(node inventory.Node) map[string]string
	// choices are the engine builds chosen per node, and installableBy what
	// each connected node said it can install, for telling it again.
	choices       func(nodeID string) map[string]string
	installableBy map[string][]string
	// offline, when set, is where engines come from on a cluster with no
	// way to the internet; see SetOffline.
	offline func(engine, version string) (string, map[string]string)

	// controlPlanes lists what agents may dial; see SetControlPlanes.
	controlPlanes func(context.Context) ([]string, error)
	// firewall is what nodes are told of llmgw's firewall; see SetFirewall.
	firewall func() *llmgwv1.Firewall
	// agentSettings is what each node's agent is told to run on; see
	// SetAgentSettings.
	agentSettings func(node string) string

	// logWaits are the log requests awaiting a reply, by request id.
	logMu    sync.Mutex
	logWaits map[string]chan *llmgwv1.LogsReply
	logSeq   uint64
	// upgradeWaits are the upgrade requests awaiting a reply; guarded by logMu.
	upgradeWaits map[string]chan *llmgwv1.UpgradeReply
	// scanWaits are the fabric scans awaiting a reply, by request id.
	scanWaits map[string]chan *llmgwv1.FabricScanReply
	testWaits map[string]chan *llmgwv1.FabricTestReply
}

// SetEngines sets the engines every node installs, by name and version.
func (s *Service) SetEngines(e map[string]string) {
	e = maps.Clone(e)
	s.enginesFor = func(inventory.Node) map[string]string { return e }
}

// SetEnginesFor sets what each node installs, decided per node: a cluster
// of two GPU generations needs two vLLM builds, or vLLM on only one.
func (s *Service) SetEnginesFor(f func(node inventory.Node) map[string]string) { s.enginesFor = f }

// SetEngineChoices sets where the builds chosen for each node come from: by
// engine, a build ("cu129") or "none". An engine not named is left to the
// node's hardware.
func (s *Service) SetEngineChoices(f func(nodeID string) map[string]string) { s.choices = f }

// Reprovision tells a connected node again what to install, after a choice
// for it changed; one not connected is told when it next connects.
func (s *Service) Reprovision(nodeID string) error {
	if s.enginesFor == nil {
		return nil
	}
	s.mu.RLock()
	can, ok := s.installableBy[nodeID]
	s.mu.RUnlock()
	node, known := s.inv.Get(nodeID)
	if !ok || !known {
		return nil
	}
	return s.send(nodeID, &llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_Provision{Provision: s.provisionFor(node, can)},
	})
}

// conn is one agent's open stream.
type conn struct {
	nodeID string
	// send is serialized through a channel rather than a mutex around the
	// stream: connect's streams are not safe for concurrent writes, and the
	// reconciler dispatches to many nodes at once.
	send   chan *llmgwv1.ControlMessage
	closed chan struct{}
	once   sync.Once
	// kv is what this stream has said about KV caches.
	kv kvState
}

func (c *conn) close() { c.once.Do(func() { close(c.closed) }) }

// New builds the service.
//
// requireIdentity refuses any agent that has not proved who it is. It should be
// true whenever TLS is configured, and is false only for local runs without
// certificates.
func New(inv *inventory.Store, status StatusSink, requireIdentity bool, log *slog.Logger) *Service {
	if log == nil {
		log = slog.Default()
	}
	return &Service{
		inv: inv, status: status, log: log,
		requireIdentity: requireIdentity,
		conns:           map[string]*conn{},
	}
}

// admit decides whether the peer on ctx may be nodeID.
func (s *Service) admit(ctx context.Context, nodeID string) error {
	// The certificate decides who this is, not the message.
	//
	// Without this check an agent could register as any node in the fleet and
	// receive that node's shards — which is the same exposure the stream had
	// before TLS, just with encryption on top. mTLS is only worth having if the
	// identity it proves is the one actually used.
	if peer := peerFrom(ctx); peer.Verified {
		if peer.CommonName != nodeID {
			s.log.Warn("rejecting an agent claiming a node id its certificate does not authorize",
				"certificate_cn", peer.CommonName, "claimed", nodeID)
			return connect.NewError(connect.CodePermissionDenied,
				fmt.Errorf("certificate authorizes node %q, not %q", peer.CommonName, nodeID))
		}
		if s.revokedAt != nil {
			at, err := s.revokedAt(ctx, nodeID)
			if err != nil {
				// Unknown is not "not revoked": the agent retries.
				return connect.NewError(connect.CodeUnavailable, fmt.Errorf("checking revocation: %w", err))
			}
			if !at.IsZero() && !peer.IssuedAt.After(at) {
				s.log.Warn("rejecting a removed node's certificate", "node", nodeID,
					"issued", peer.IssuedAt, "removed", at)
				return connect.NewError(connect.CodePermissionDenied, fmt.Errorf(
					"node %q was removed from the cluster, and its certificate with it; onboard it again", nodeID))
			}
		}
	} else if s.requireIdentity {
		return connect.NewError(connect.CodeUnauthenticated,
			errors.New("no verified client certificate"))
	}
	return nil
}

// SetRevocations has agents refused whose certificate was issued before
// their node was removed.
func (s *Service) SetRevocations(f func(ctx context.Context, nodeID string) (time.Time, error)) {
	s.revokedAt = f
}

// Handler returns the HTTP path and handler to mount.
//
// The handler is wrapped so the verified client identity reaches the stream:
// connect exposes the peer's address but not its certificate.
func (s *Service) Handler() (string, http.Handler) {
	// A message is a heartbeat, a report or a log; nothing an agent says
	// needs more, and a bound keeps one connection from taking the memory.
	path, h := llmgwv1connect.NewAgentServiceHandler(s, connect.WithReadMaxBytes(16<<20))
	return path, AuthMiddleware(h)
}

// Connect serves one agent's stream for the lifetime of that agent.
func (s *Service) Connect(
	ctx context.Context,
	stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage],
) error {
	// A standby must not hold an agent's stream. Commands travel down the
	// stream the agent opened, so whichever control plane holds it is the only
	// one that can start or stop anything on that node — and if that is not
	// the leader, the leader's placements fail with "no agent connected" for a
	// node that is plainly connected.
	//
	// Refused before the opening message is read, so the agent spends nothing
	// on a connection that cannot serve it.
	if s.leader != nil && !s.leader.IsLeader() {
		return connect.NewError(connect.CodeUnavailable,
			errors.New("this control plane is a standby; connect to the leader"))
	}

	first, err := stream.Receive()
	if err != nil {
		return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("no opening message: %w", err))
	}
	reg := first.GetRegister()
	if reg == nil {
		// Registration first, always: the control plane has to know which node
		// it is talking to before it can route anything to it.
		return connect.NewError(connect.CodeInvalidArgument,
			errors.New("the first message must be a Register"))
	}
	nodeID := reg.GetNodeId()
	if nodeID == "" {
		return connect.NewError(connect.CodeInvalidArgument, errors.New("node_id is required"))
	}

	if err := s.admit(ctx, nodeID); err != nil {
		return err
	}

	node := NodeFromRegister(reg)
	s.inv.Upsert(node)
	s.log.Info("agent registered",
		"node", nodeID, "gpus", len(node.GPUs), "fabrics", node.FabricNames(),
		"running_shards", len(reg.GetRunningShards()),
		"agent_version", reg.GetAgentVersion())

	c := s.attach(ctx, nodeID)
	defer s.detach(nodeID, c)
	for _, rs := range reg.GetRunningShards() {
		c.own(rs.GetReplicaId())
	}
	defer s.kvGone(c)

	if err := stream.Send(&llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_Ack{Ack: &llmgwv1.Ack{MessageId: "register", Ok: true}},
	}); err != nil {
		return err
	}

	// What this node should have installed. Sent on every registration, so a
	// node rebuilt from nothing is told the same as one that has been here a
	// month, and a changed version reaches every node on its next connect.
	if s.enginesFor != nil {
		can := installable(reg)
		s.mu.Lock()
		if s.installableBy == nil {
			s.installableBy = map[string][]string{}
		}
		s.installableBy[nodeID] = can
		s.mu.Unlock()
		p := s.provisionFor(node, can)
		if err := stream.Send(&llmgwv1.ControlMessage{
			Msg: &llmgwv1.ControlMessage_Provision{Provision: p},
		}); err != nil {
			return err
		}
	}
	// Which control planes there are, so an agent configured with one learns
	// of a second added since.
	if msg := s.controlPlanesMsg(ctx); msg != nil {
		if err := stream.Send(msg); err != nil {
			return err
		}
	}
	// Whether to keep llmgw's firewall, as the cluster's settings say now.
	if msg := s.firewallMsg(); msg != nil {
		if err := stream.Send(msg); err != nil {
			return err
		}
	}
	// What the cluster's settings say this agent runs on. An agent already
	// on them ignores it; one that is not starts again on them.
	if msg := s.agentSettingsMsg(nodeID); msg != nil {
		if err := stream.Send(msg); err != nil {
			return err
		}
	}

	// Recover replica state from what this node is already running, before any
	// reconciliation pass can conclude the cluster is empty and place
	// duplicates (plan §6.5).
	s.adopt(ctx, nodeID, reg.GetRunningShards(), c)

	// One goroutine writes, this one reads. Splitting them is what lets a
	// command be dispatched while the agent is mid-heartbeat.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		s.writeLoop(ctx, stream, c)
	}()
	defer wg.Wait()

	// Read in a goroutine of its own, so the stream can also be ended from
	// this side: Receive blocks until the agent sends, and an agent that is
	// only heartbeating would keep a stream this process should no longer
	// hold for as long as it likes. Returning ends it; the reader's Receive
	// then fails and it follows.
	read := make(chan struct{})
	go func() {
		defer close(read)
		s.readLoop(ctx, stream, c, nodeID)
	}()
	tick := time.NewTicker(leaderCheck)
	defer tick.Stop()
	for {
		select {
		case <-read:
			return nil
		case <-c.closed:
			// Replaced by a newer stream from the same node, or the writer
			// failed.
			return nil
		case <-tick.C:
			// Leadership lost while the process lives on — the store was
			// slow, or another took the lease. A standby holding agents is
			// a cluster whose leader cannot reach its nodes, so let them go
			// and they find the leader.
			if s.leader != nil && !s.leader.IsLeader() {
				s.log.Info("no longer leading; letting the agent go", "node", nodeID)
				c.close()
				return connect.NewError(connect.CodeUnavailable,
					errors.New("this control plane is a standby; connect to the leader"))
			}
		}
	}
}

// leaderCheck is how often an open stream checks that this process still
// leads.
var leaderCheck = time.Second

func (s *Service) readLoop(ctx context.Context,
	stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], c *conn, nodeID string,
) {
	for {
		msg, err := stream.Receive()
		if err != nil {
			c.close()
			if ctx.Err() != nil || errors.Is(err, io.EOF) {
				s.log.Info("agent disconnected", "node", nodeID)
				return
			}
			s.log.Warn("agent stream failed", "node", nodeID, "err", err)
			return
		}
		if kv := msg.GetKvEvents(); kv != nil {
			s.kvEvents(c, kv)
			continue
		}
		if sh := msg.GetShard(); sh != nil {
			c.own(sh.GetReplicaId())
		}
		s.handle(nodeID, msg)
	}
}

// adopt hands the agent's running shards to the reconciler and stops whatever
// nothing in the catalog claims.
//
// An empty list is not skipped. "I am running nothing" is information — it is
// how the control plane learns that an engine died while its agent survived —
// and returning early on it made the reconciler's half of that exchange
// unreachable in production while its unit tests, which call the reconciler
// directly, kept passing.
func (s *Service) adopt(ctx context.Context, nodeID string, shards []*llmgwv1.RunningShard, c *conn) {
	// A trace of its own, and only for a node that reported something: a
	// node running nothing registering is not an event anyone looks for.
	tctx, span := ctx, trace.SpanFromContext(context.Background()) // the no-op span
	if len(shards) > 0 {
		tctx, span = obs.Tracer(tracerName).Start(ctx, "adopt",
			trace.WithNewRoot(),
			trace.WithAttributes(
				attribute.String("llmgw.node", nodeID),
				attribute.Int("llmgw.shards", len(shards)),
			))
	}
	defer span.End()
	unclaimed := s.status.Adopt(nodeID, AdoptedFromRegister(nodeID, shards))
	span.SetAttributes(attribute.Int("llmgw.unclaimed", len(unclaimed)))
	for _, u := range unclaimed {
		// The model or profile left the catalog while this was running. Nobody
		// will ever route to it, and it is holding GPUs.
		s.log.Warn("stopping an unclaimed shard",
			"node", nodeID, "shard", u.ShardID, "model", u.Model, "profile", u.Profile)
		select {
		case c.send <- &llmgwv1.ControlMessage{
			Msg:   &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: u.ShardID}},
			Trace: obs.Inject(tctx),
		}:
		case <-ctx.Done():
			return
		}
	}
}

func (s *Service) writeLoop(ctx context.Context,
	stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], c *conn,
) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-c.closed:
			return
		case msg := <-c.send:
			if err := stream.Send(msg); err != nil {
				s.log.Warn("sending to agent failed", "node", c.nodeID, "err", err)
				c.close()
				return
			}
		}
	}
}

func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) {
	switch m := msg.Msg.(type) {
	case *llmgwv1.AgentMessage_Heartbeat:
		s.inv.Touch(nodeID)
		s.applyHeartbeat(nodeID, m.Heartbeat)

	case *llmgwv1.AgentMessage_Shard:
		// Progress rather than a transition. Falling through to OnShardStatus
		// would re-mark the shard as starting on every chunk, which is how a
		// download that takes an hour would keep a start timer at zero.
		if f := m.Shard.GetFetch(); f != nil {
			s.status.OnFetchProgress(m.Shard.GetReplicaId(), m.Shard.GetShardId(),
				f.GetSource(), f.GetBytesDone(), f.GetBytesTotal())
			return
		}
		if c := m.Shard.GetCompute(); c != nil {
			s.status.OnShardCompute(m.Shard.GetReplicaId(), m.Shard.GetShardId(), reconciler.Compute{
				Backend: c.GetBackend(), Devices: c.GetDevices(),
				Offloaded: int(c.GetLayersOffloaded()), Total: int(c.GetLayersTotal()),
			})
			// Compute alone, with no state: not a transition.
			if m.Shard.GetState() == llmgwv1.ShardStatus_STATE_UNSPECIFIED {
				return
			}
		}
		st := StateFromWire(m.Shard.GetState())
		s.status.OnShardStatus(m.Shard.GetReplicaId(), m.Shard.GetShardId(), st, m.Shard.GetError())
		if m.Shard.GetState() == llmgwv1.ShardStatus_STATE_FAILED {
			s.log.Error("shard failed",
				"node", nodeID, "shard", m.Shard.GetShardId(),
				"err", m.Shard.GetError(), "exit_code", m.Shard.GetExitCode(),
				"last_output", m.Shard.GetLastOutput())
		}

	case *llmgwv1.AgentMessage_FabricTest:
		s.logMu.Lock()
		ch, ok := s.testWaits[m.FabricTest.GetRequestId()]
		delete(s.testWaits, m.FabricTest.GetRequestId())
		s.logMu.Unlock()
		if ok {
			ch <- m.FabricTest
		}

	case *llmgwv1.AgentMessage_FabricScan:
		s.logMu.Lock()
		ch, ok := s.scanWaits[m.FabricScan.GetRequestId()]
		delete(s.scanWaits, m.FabricScan.GetRequestId())
		s.logMu.Unlock()
		if ok {
			ch <- m.FabricScan
		}

	case *llmgwv1.AgentMessage_Upgrade:
		s.logMu.Lock()
		ch, ok := s.upgradeWaits[m.Upgrade.GetRequestId()]
		delete(s.upgradeWaits, m.Upgrade.GetRequestId())
		s.logMu.Unlock()
		if ok {
			ch <- m.Upgrade
		}

	case *llmgwv1.AgentMessage_Logs:
		s.logMu.Lock()
		ch, ok := s.logWaits[m.Logs.GetRequestId()]
		delete(s.logWaits, m.Logs.GetRequestId())
		s.logMu.Unlock()
		if ok {
			ch <- m.Logs // buffered; the asker may have given up
		}

	case *llmgwv1.AgentMessage_Register:
		// A re-registration on a live stream means the agent's view of its own
		// hardware changed — a GPU fell out, say. Take it.
		s.inv.Upsert(NodeFromRegister(m.Register))
		s.log.Info("agent re-registered", "node", nodeID)
	}
}

// applyHeartbeat folds in the cheap facts a heartbeat carries, so a foreign
// process appearing on a card after registration is noticed without waiting for
// a reconnect.
func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) {
	used := hb.GetGpuMemoryUsedMb()
	foreign := hb.GetGpuForeignMemoryMb()
	problems := hb.GetNewProblems()
	storage := hb.GetStorage()
	memMB := int(hb.GetMemoryTotalMb())
	cpus := int(hb.GetCpuCount())
	if len(used) == 0 && len(foreign) == 0 && len(problems) == 0 && len(storage) == 0 &&
		!hb.GetWeightsChanged() && memMB == 0 && cpus == 0 && !hb.GetEnginesChanged() &&
		!hb.GetProvisionChanged() && hb.GetLoad() == nil && len(hb.GetIbPorts()) == 0 &&
		!hb.GetFabricsChanged() && len(hb.GetGpuFabric()) == 0 && hb.GetImex() == nil {
		return
	}

	n, ok := s.inv.Get(nodeID)
	if !ok {
		return
	}
	for i := range n.GPUs {
		if i < len(used) {
			n.GPUs[i].MemoryUsedMB = int(used[i])
		}
	}
	// Another tenant on a card, judged afresh: the reading at registration is
	// replaced, both ways. Without this a card that had something on it when
	// its agent started stayed out of placement after it left, and said so,
	// until the agent was restarted. Only a reading for every card replaces
	// it; an agent that sends none, or could not measure, leaves it standing.
	if len(foreign) == len(n.GPUs) {
		for i := range n.GPUs {
			ps := slices.DeleteFunc(slices.Clone(n.GPUs[i].Health.Problems), discovery.IsForeignMemoryProblem)
			if p := discovery.ForeignMemoryProblem(int(foreign[i])); p != "" {
				ps = append(ps, p)
			}
			n.GPUs[i].Health.Problems = ps
		}
	}
	if len(problems) > 0 && len(n.GPUs) > 0 {
		// Problems are reported per node, not per GPU, so attach them to the
		// node's first GPU rather than guessing which one is affected. Any
		// problem is enough to keep that GPU out of placement, which is the
		// conservative direction.
		n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...)
	}

	// Disk is a fact about a moment, not about a machine. Sampled only at
	// registration it is wrong on exactly the node worth watching — the one
	// busy filling up — and it stays wrong until that agent restarts.
	if len(storage) > 0 {
		n.Storage = n.Storage[:0]
		for _, st := range storage {
			n.Storage = append(n.Storage, discovery.Storage{
				Path:       st.GetPath(),
				FreeBytes:  st.GetFreeBytes(),
				TotalBytes: st.GetTotalBytes(),
				Shared:     st.GetShared(),
			})
		}
	}

	// Only when the agent says the list is meant. An empty list without the
	// flag is silence, and treating it as "this node holds nothing" would
	// empty a node's inventory on every heartbeat.
	if hb.GetWeightsChanged() {
		n.Weights = n.Weights[:0]
		for _, w := range hb.GetWeights() {
			n.Weights = append(n.Weights, discovery.LocalWeights{
				Path: w.GetPath(), SizeBytes: w.GetSizeBytes(), Shared: w.GetShared(),
			})
		}
	}

	// RAM and engines change under a running agent too: memory added to a VM,
	// an engine installed this afternoon. Said in the log when they do, because
	// "the node now qualifies" is the answer to why a model placed an hour after
	// it was asked for, and "the node no longer does" is the answer to why one
	// stopped placing.
	if memMB > 0 && memMB != n.MemoryTotalMB {
		s.log.Info("node memory changed", "node", nodeID,
			"from_mb", n.MemoryTotalMB, "to_mb", memMB)
		n.MemoryTotalMB = memMB
	}
	if cpus > 0 && cpus != n.CPUCount {
		if n.CPUCount > 0 {
			s.log.Info("node CPUs changed", "node", nodeID, "from", n.CPUCount, "to", cpus)
		}
		n.CPUCount = cpus
	}
	if hb.GetEnginesChanged() {
		vers := maps.Clone(hb.GetEngineVersions())
		if !maps.Equal(vers, n.EngineVersions) {
			s.log.Info("node engines changed", "node", nodeID,
				"from", slices.Sorted(maps.Keys(n.EngineVersions)),
				"to", slices.Sorted(maps.Keys(vers)))
		}
		n.EngineVersions = vers
	}
	if hb.GetProvisionChanged() {
		prev := map[string]inventory.EngineInstall{}
		for _, p := range n.Provision {
			prev[p.Engine] = p
		}
		n.Provision = n.Provision[:0]
		for _, p := range hb.GetProvision() {
			e := inventory.EngineInstall{
				Engine: p.GetEngine(), Version: p.GetVersion(), Variant: p.GetVariant(),
				State: p.GetState(), Detail: p.GetDetail(),
			}
			// A state change is worth a line; a download's progress is not.
			if old, ok := prev[e.Engine]; !ok || old.State != e.State {
				lvl := slog.LevelInfo
				if e.State == "failed" || e.State == "unsupported" {
					lvl = slog.LevelWarn
				}
				s.log.Log(context.Background(), lvl, "engine installation on a node",
					"node", nodeID, "engine", e.Engine, "version", e.Version,
					"variant", e.Variant, "state", e.State, "detail", e.Detail)
			}
			n.Provision = append(n.Provision, e)
		}
	}
	if l := hb.GetLoad(); l != nil {
		n.Load = loadFrom(l, time.Now())
	}
	if ports := hb.GetIbPorts(); len(ports) > 0 {
		before := map[string]string{}
		prev := map[string]inventory.IBPort{}
		for _, p := range n.IBPorts {
			before[fmt.Sprintf("%s/%d", p.HCA, p.Port)] = p.SMGUID
			prev[fmt.Sprintf("%s/%d", p.HCA, p.Port)] = p
		}
		now := time.Now()
		secs := now.Sub(n.IBAt).Seconds()
		n.IBPorts = nil
		for _, p := range ports {
			// A different master than last time is a failover, or a port
			// moved: the fabric's own event, kept for the page to show.
			where := fmt.Sprintf("%s/%d", p.GetHca(), p.GetPort())
			if old := before[where]; old != "" && p.GetSmGuid() != "" && old != p.GetSmGuid() {
				s.log.Warn("a port sees a different master subnet manager", "node", nodeID,
					"port", where, "from", old, "to", p.GetSmGuid())
				s.inv.RecordSMChange(inventory.SMChange{At: time.Now(), Node: nodeID, Port: where,
					From: old, To: p.GetSmGuid()})
			}
			var sms []inventory.SubnetManager
			for _, r := range p.GetSubnetManagers() {
				sms = append(sms, inventory.SubnetManager{GUID: r.GetGuid(), LID: int(r.GetLid()),
					Priority: int(r.GetPriority()), State: r.GetState()})
			}
			n.IBPorts = append(n.IBPorts, inventory.IBPort{
				SMs: sms, SMsError: p.GetSubnetManagersError(),
				TxBps: rate(prev[where].Counters["port_xmit_data"], p.GetCounters()["port_xmit_data"], secs),
				RxBps: rate(prev[where].Counters["port_rcv_data"], p.GetCounters()["port_rcv_data"], secs),
				HCA:   p.GetHca(), Port: int(p.GetPort()), State: p.GetState(),
				PhysState: p.GetPhysState(), RateGbps: int(p.GetRateGbps()),
				LinkLayer: p.GetLinkLayer(), Counters: maps.Clone(p.GetCounters()),
				NodeGUID: p.GetNodeGuid(), PortGUID: p.GetPortGuid(), GIDPrefix: p.GetGidPrefix(),
				SMGUID: p.GetSmGuid(), SMLID: int(p.GetSmLid()), SMPriority: int(p.GetSmPriority()),
				SMState: p.GetSmState(), SMError: p.GetSmError(), Network: p.GetNetwork(),
			})
		}
	}
	if len(hb.GetIbPorts()) > 0 {
		n.IBAt = time.Now()
	}
	// NVLink domain membership, which a fabric manager restarting takes
	// away from every card at once.
	if fab := hb.GetGpuFabric(); len(fab) > 0 {
		before := n.NVLinkDomain()
		for _, f := range fab {
			for i := range n.GPUs {
				if n.GPUs[i].Index == int(f.GetIndex()) {
					n.GPUs[i].Fabric = fabricFrom(f.GetState(), f.GetStatus(), f.GetClusterUuid(), f.GetCliqueId(), f.GetBandwidth())
				}
			}
		}
		if after := n.NVLinkDomain(); after != before {
			s.log.Warn("a node's NVLink domain changed", "node", nodeID, "from", before, "to", after)
		}
	}
	if m := hb.GetImex(); m != nil {
		n.IMEX = discovery.IMEX{State: m.GetState(), Channels: int(m.GetChannels())}
	}
	// A port went down, or came back. Placement reads the fabrics, so this is
	// what stops a node being offered multi-node work it can no longer carry.
	if hb.GetFabricsChanged() {
		var fabs []inventory.FabricAttachment
		for _, f := range hb.GetFabrics() {
			fabs = append(fabs, inventory.FabricAttachment{
				Fabric: f.GetFabric(), HCAs: slices.Clone(f.GetHcas()), RateGbps: int(f.GetRateGbps()),
			})
		}
		s.log.Warn("a node's fabrics changed", "node", nodeID,
			"from", n.FabricNames(), "to", fabricNames(fabs))
		n.Fabrics = fabs
	}
	s.inv.Upsert(n)
}

func fabricNames(fs []inventory.FabricAttachment) []string {
	out := make([]string, len(fs))
	for i, f := range fs {
		out[i] = f.Fabric
	}
	return out
}

// loadFrom is a heartbeat's load as the inventory keeps it.
func loadFrom(l *llmgwv1.NodeLoad, at time.Time) *inventory.Load {
	out := &inventory.Load{
		At: at, CPUPercent: l.GetCpuPercent(), CPUCount: int(l.GetCpuCount()),
		Load1: l.GetLoad1(), Load5: l.GetLoad5(), Load15: l.GetLoad15(),
		MemUsedBytes: l.GetMemoryUsedBytes(), MemAvailBytes: l.GetMemoryAvailableBytes(),
		Shards: map[string]inventory.ShardLoad{},
	}
	for _, sh := range l.GetShards() {
		out.Shards[sh.GetShardId()] = inventory.ShardLoad{
			CPUCores: sh.GetCpuCores(), RSSBytes: sh.GetMemoryRssBytes(),
		}
	}
	out.CPUTempC, out.CPUPowerW, out.InletTempC = l.GetCpuTempC(), l.GetCpuPowerW(), l.GetInletTempC()
	for _, g := range l.GetGpus() {
		out.GPUs = append(out.GPUs, inventory.GPUSensor{
			Index: int(g.GetIndex()), TempC: g.GetTempC(), MemoryTempC: g.GetMemoryTempC(),
			PowerW: g.GetPowerW(), PowerLimitW: g.GetPowerLimitW(), Throttle: g.GetThrottle(),
			UtilPercent: g.GetUtilizationPercent(), MemUtilPercent: g.GetMemoryUtilizationPercent(),
			MemoryUsedMB: g.GetMemoryUsedMb(), MemoryTotalMB: g.GetMemoryTotalMb(),
			SMClockMHz: g.GetSmClockMhz(), ECCUncorrected: g.GetEccUncorrected(),
			Wear: wearFrom(g),
		})
	}
	return out
}

func (s *Service) attach(ctx context.Context, nodeID string) *conn {
	c := &conn{
		nodeID: nodeID,
		// Buffered so a burst of commands does not block the reconciler on a
		// slow agent; if it fills, the agent is too far behind to be useful.
		send:   make(chan *llmgwv1.ControlMessage, 64),
		closed: make(chan struct{}),
	}

	s.mu.Lock()
	if prev, ok := s.conns[nodeID]; ok {
		// A node reconnecting before the old stream was reaped. The new one
		// wins: the old is by definition the dead one.
		s.log.Warn("agent reconnected while an old stream was open", "node", nodeID)
		prev.close()
	}
	s.conns[nodeID] = c
	s.mu.Unlock()

	go func() {
		<-ctx.Done()
		c.close()
	}()
	return c
}

func (s *Service) detach(nodeID string, c *conn) {
	c.close()
	s.mu.Lock()
	if s.conns[nodeID] == c {
		delete(s.conns, nodeID)
	}
	s.mu.Unlock()

	// The node is not marked lost here. Losing a stream is not losing a node,
	// and the reconciler's two-stage expiry already handles silence with the
	// right delays (plan §6.3).
	s.log.Info("agent stream closed", "node", nodeID)
}

// ConnectedNodes lists the nodes with an open stream.
func (s *Service) ConnectedNodes() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make([]string, 0, len(s.conns))
	for id := range s.conns {
		out = append(out, id)
	}
	slices.Sort(out)
	return out
}

// --- Dispatcher ---

// Leader reports whether this process is the one that should be holding agent
// streams.
type Leader interface{ IsLeader() bool }

// SetLeadership makes this service refuse agents while it is a standby.
//
// Nil accepts them unconditionally, which is what a single control plane wants
// and what every deployment before leadership had.
func (s *Service) SetLeadership(l Leader) { s.leader = l }

// tracerName names the spans this package opens.
const tracerName = "gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"

// ErrNotConnected means no agent stream is open for that node.
var ErrNotConnected = errors.New("no agent connected")

// StartShard implements reconciler.Dispatcher.
//
// The command carries the trace of ctx, so the agent's spans for this start
// — the weights, the engine's launch, its readiness — are in the trace of
// the replica they are starting rather than in one of their own.
func (s *Service) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error {
	return s.send(cmd.NodeID, &llmgwv1.ControlMessage{
		Msg:   &llmgwv1.ControlMessage_Start{Start: StartShardFromCommand(cmd)},
		Trace: obs.Inject(ctx),
	})
}

// StopShard implements reconciler.Dispatcher.
func (s *Service) StopShard(ctx context.Context, nodeID, shardID string) error {
	return s.send(nodeID, &llmgwv1.ControlMessage{
		Msg:   &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: shardID}},
		Trace: obs.Inject(ctx),
	})
}

// DrainShard implements reconciler.Dispatcher.
func (s *Service) DrainShard(ctx context.Context, nodeID, shardID string) error {
	return s.send(nodeID, &llmgwv1.ControlMessage{
		Msg:   &llmgwv1.ControlMessage_Drain{Drain: &llmgwv1.DrainShard{ShardId: shardID}},
		Trace: obs.Inject(ctx),
	})
}

// RemoveWeights asks a node to delete cached files no model needs.
//
// The node checks every path again before unlinking it, so this is a proposal
// rather than an instruction — which is the right shape for a message that
// deletes data the control plane cannot see.
func (s *Service) RemoveWeights(_ context.Context, nodeID string, paths []string) error {
	if len(paths) == 0 {
		return nil
	}
	return s.send(nodeID, &llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_RemoveWeights{
			RemoveWeights: &llmgwv1.RemoveWeights{Paths: paths},
		},
	})
}

func (s *Service) send(nodeID string, msg *llmgwv1.ControlMessage) error {
	s.mu.RLock()
	c, ok := s.conns[nodeID]
	s.mu.RUnlock()
	if !ok {
		return fmt.Errorf("%w: node %s", ErrNotConnected, nodeID)
	}

	select {
	case c.send <- msg:
		return nil
	case <-c.closed:
		return fmt.Errorf("%w: node %s stream closed", ErrNotConnected, nodeID)
	default:
		// The queue is full, so this agent is far enough behind that queueing
		// more would only delay discovering it is broken.
		return fmt.Errorf("node %s is not keeping up with commands", nodeID)
	}
}

// Logs asks a node for the recent lines of one of its logs: "agent", or a
// shard id for that engine's output.
//
// It returns the lines, and every source the node can show.
func (s *Service) Logs(ctx context.Context, nodeID, source string, lines int) ([]string, []string, error) {
	s.logMu.Lock()
	if s.logWaits == nil {
		s.logWaits = map[string]chan *llmgwv1.LogsReply{}
	}
	s.logSeq++
	id := fmt.Sprintf("logs-%d", s.logSeq)
	ch := make(chan *llmgwv1.LogsReply, 1)
	s.logWaits[id] = ch
	s.logMu.Unlock()
	defer func() {
		s.logMu.Lock()
		delete(s.logWaits, id)
		s.logMu.Unlock()
	}()

	if source == "" {
		source = "agent"
	}
	if err := s.send(nodeID, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_Logs{
		Logs: &llmgwv1.LogsRequest{RequestId: id, Source: source, Lines: int32(lines)}, //nolint:gosec // bounded by the caller
	}}); err != nil {
		return nil, nil, err
	}
	// Bounded, because a page is waiting: an agent older than this request
	// never answers it at all.
	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	select {
	case r := <-ch:
		if r.GetError() != "" {
			return nil, r.GetSources(), errors.New(r.GetError())
		}
		return r.GetLines(), r.GetSources(), nil
	case <-ctx.Done():
		return nil, nil, fmt.Errorf("node %s did not answer; an agent older than v0.8.4 cannot", nodeID)
	}
}

// ScanFabric asks a node to map the fabric one of its ports is on.
func (s *Service) ScanFabric(ctx context.Context, nodeID, hca string, port int) (*llmgwv1.FabricScanReply, error) {
	s.logMu.Lock()
	if s.scanWaits == nil {
		s.scanWaits = map[string]chan *llmgwv1.FabricScanReply{}
	}
	s.logSeq++
	id := fmt.Sprintf("scan-%d", s.logSeq)
	ch := make(chan *llmgwv1.FabricScanReply, 1)
	s.scanWaits[id] = ch
	s.logMu.Unlock()
	defer func() {
		s.logMu.Lock()
		delete(s.scanWaits, id)
		s.logMu.Unlock()
	}()
	if err := s.send(nodeID, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_FabricScan{
		FabricScan: &llmgwv1.FabricScanRequest{RequestId: id, Hca: hca, Port: int32(port)}, //nolint:gosec // a port number
	}}); err != nil {
		return nil, err
	}
	ctx, cancel := context.WithTimeout(ctx, 3*time.Minute)
	defer cancel()
	select {
	case r := <-ch:
		return r, nil
	case <-ctx.Done():
		return nil, fmt.Errorf("node %s did not answer the fabric scan", nodeID)
	}
}

// FabricTest asks a node for one side of a bandwidth test, and waits for
// what it measured. The request's id is filled in here.
func (s *Service) FabricTest(ctx context.Context, nodeID string, req *llmgwv1.FabricTestRequest) (*llmgwv1.FabricTestReply, error) {
	s.logMu.Lock()
	if s.testWaits == nil {
		s.testWaits = map[string]chan *llmgwv1.FabricTestReply{}
	}
	s.logSeq++
	id := fmt.Sprintf("test-%d", s.logSeq)
	ch := make(chan *llmgwv1.FabricTestReply, 1)
	s.testWaits[id] = ch
	s.logMu.Unlock()
	defer func() {
		s.logMu.Lock()
		delete(s.testWaits, id)
		s.logMu.Unlock()
	}()
	req.RequestId = id
	if err := s.send(nodeID, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_FabricTest{FabricTest: req}}); err != nil {
		return nil, err
	}
	// The test's own length, and a package install the first time.
	ctx, cancel := context.WithTimeout(ctx, time.Duration(req.GetSeconds())*time.Second+5*time.Minute)
	defer cancel()
	select {
	case r := <-ch:
		return r, nil
	case <-ctx.Done():
		return nil, fmt.Errorf("node %s did not answer the bandwidth test", nodeID)
	}
}

// ScanFabrics maps every physical InfiniBand fabric every so often, asking
// one node on each: the first, by name, whose port on it is up. A fabric
// looks the same from all its ports, so one is enough, and the rest of its
// nodes are left to their work.
func (s *Service) ScanFabrics(ctx context.Context, every time.Duration) {
	if every <= 0 {
		every = 5 * time.Minute
	}
	// Soon after a start, once agents have sent a heartbeat or two.
	timer := time.NewTimer(30 * time.Second)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-timer.C:
		}
		s.scanOnce(ctx)
		timer.Reset(every)
	}
}

func (s *Service) scanOnce(ctx context.Context) {
	nodes := s.inv.Nodes()
	sort.Slice(nodes, func(i, j int) bool { return nodes[i].ID < nodes[j].ID })
	canon := inventory.CanonicalSMs(nodes)
	chosen := map[string]bool{}
	for _, n := range nodes {
		if n.State != inventory.Ready {
			continue
		}
		for _, p := range n.IBPorts {
			key := p.FabricKeyIn(canon)
			if key == "" || p.RoCE() || chosen[key] {
				continue
			}
			chosen[key] = true
			where := fmt.Sprintf("%s/%d", p.HCA, p.Port)
			r, err := s.ScanFabric(ctx, n.ID, p.HCA, p.Port)
			sc := inventory.FabricScan{Key: key, At: time.Now(), Node: n.ID, Port: where}
			switch {
			case err != nil:
				sc.Err = err.Error()
			case r.GetError() != "":
				sc.Err = r.GetError()
			default:
				sc.Output = r.GetOutput()
			}
			if r != nil {
				sc.Took = time.Duration(r.GetTookMs()) * time.Millisecond
			}
			if sc.Err != "" {
				s.log.Warn("could not map a fabric", "fabric", key, "node", n.ID, "port", where, "err", sc.Err)
			}
			s.inv.SetScan(sc)
		}
	}
}

// rate is a counter's increase a second, zero when there is no earlier
// reading or the counter went backwards — an adapter reset starts it again.
func rate(before, now uint64, secs float64) float64 {
	if before == 0 || now < before || secs <= 0 || secs > 120 {
		return 0
	}
	return float64(now-before) / secs
}

// provisionFor is what a node is told to install: the engines decided for
// it, those it can install.
func (s *Service) provisionFor(node inventory.Node, can []string) *llmgwv1.Provision {
	p := &llmgwv1.Provision{}
	wants := s.enginesFor(node)
	var chosen map[string]string
	if s.choices != nil {
		chosen = s.choices(node.ID)
	}
	for _, e := range slices.Sorted(maps.Keys(wants)) {
		if !slices.Contains(can, e) || chosen[e] == "none" {
			continue
		}
		w := &llmgwv1.EngineWant{Engine: e, Version: wants[e], Variant: chosen[e]}
		if s.offline != nil {
			w.OfflineBase, w.OfflineSha256 = s.offline(e, wants[e])
		}
		p.Engines = append(p.Engines, w)
	}
	return p
}

// SetOffline makes nodes install engines from the cluster's mirror rather
// than from their publishers: f gives where an engine version's files are
// there, and each file's checksum from the bundle the control plane trusts.
// A version no bundle carried has no checksums, and the node says which
// bundle to import rather than reaching for the internet.
func (s *Service) SetOffline(f func(engine, version string) (base string, sha256 map[string]string)) {
	s.offline = f
}

// installable are the engines a node can be asked to install. An agent older
// than the list said nothing, and could install llama.cpp alone: asked for
// vLLM too, it reported it unsupported, and a cluster half way through an
// upgrade showed every such node as failing.
func installable(reg *llmgwv1.Register) []string {
	if can := reg.GetInstallableEngines(); len(can) > 0 {
		return can
	}
	return []string{"llamacpp"}
}

// The admin API finds the bandwidth test by asserting this; a signature that
// drifted would make it quietly nil.
var _ fabrictest.Asker = (*Service)(nil)

// Upgrade tells a node's agent to replace its binary with a release and
// restart, and returns once the agent says the new binary is in place — or
// why it is not. The agent then exits and comes back on the new version; the
// caller watches for that in the inventory.
func (s *Service) Upgrade(ctx context.Context, nodeID, version, releases string) error {
	s.logMu.Lock()
	if s.upgradeWaits == nil {
		s.upgradeWaits = map[string]chan *llmgwv1.UpgradeReply{}
	}
	s.logSeq++
	id := fmt.Sprintf("upgrade-%d", s.logSeq)
	ch := make(chan *llmgwv1.UpgradeReply, 1)
	s.upgradeWaits[id] = ch
	s.logMu.Unlock()
	defer func() {
		s.logMu.Lock()
		delete(s.upgradeWaits, id)
		s.logMu.Unlock()
	}()

	if err := s.send(nodeID, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_Upgrade{
		Upgrade: &llmgwv1.UpgradeAgent{RequestId: id, Version: version, Releases: releases},
	}}); err != nil {
		return err
	}
	s.log.Info("told a node to upgrade", "node", nodeID, "version", version)
	// A download of some tens of megabytes on a slow link, and no longer: an
	// agent too old to know this message never answers.
	ctx, cancel := context.WithTimeout(ctx, 5*time.Minute)
	defer cancel()
	select {
	case r := <-ch:
		if r.GetError() != "" {
			return errors.New(r.GetError())
		}
		return nil
	case <-ctx.Done():
		return fmt.Errorf("node %s did not answer the upgrade; an agent older than v0.10.4 cannot upgrade itself", nodeID)
	}
}

// wearFrom is a card's wear, when the agent read it.
func wearFrom(g *llmgwv1.GPUSensor) *inventory.GPUWear {
	if !g.GetHasWear() {
		return nil
	}
	w := &inventory.GPUWear{
		PCIeReplays: g.GetPcieReplays(), RetiredSBE: g.GetRetiredPagesSbe(), RetiredDBE: g.GetRetiredPagesDbe(),
		RetirePending: g.GetRetirementPending(), RemappedCorrectable: g.GetRemappedRowsCorrectable(),
		RemappedUncorrectable: g.GetRemappedRowsUncorrectable(), RemapPending: g.GetRemapPending(),
		RemapFailure: g.GetRemapFailure(), NVLinkCRC: g.GetNvlinkCrcErrors(), NVLinkReplay: g.GetNvlinkReplayErrors(),
		NVLinkRecovery: g.GetNvlinkRecoveryErrors(), XIDs: g.GetXidErrors(), LastXID: int(g.GetLastXid()),
		LastXIDText: g.GetLastXidText(),
	}
	if u := g.GetLastXidUnix(); u > 0 {
		w.LastXIDAt = time.Unix(u, 0)
	}
	return w
}