llmgw/internal/control/agentsvc/service.go

package agentsvc

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

	"connectrpc.com/connect"

	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/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
)

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

	mu    sync.RWMutex
	conns map[string]*conn

	// engines is what every node is told to install, set once at start.
	engines map[string]string

	// logWaits are the log requests awaiting a reply, by request id.
	logMu    sync.Mutex
	logWaits map[string]chan *llmgwv1.LogsReply
	logSeq   uint64
}

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

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

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{},
	}
}

// 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) {
	path, h := llmgwv1connect.NewAgentServiceHandler(s)
	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"))
	}

	// 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))
		}
	} else if s.requireIdentity {
		return connect.NewError(connect.CodeUnauthenticated,
			errors.New("no verified client certificate"))
	}

	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)

	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.engines != nil {
		p := &llmgwv1.Provision{}
		for _, e := range slices.Sorted(maps.Keys(s.engines)) {
			p.Engines = append(p.Engines, &llmgwv1.EngineWant{Engine: e, Version: s.engines[e]})
		}
		if err := stream.Send(&llmgwv1.ControlMessage{
			Msg: &llmgwv1.ControlMessage_Provision{Provision: p},
		}); 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()

	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 nil
			}
			s.log.Warn("agent stream failed", "node", nodeID, "err", err)
			return nil
		}
		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) {
	unclaimed := s.status.Adopt(nodeID, AdoptedFromRegister(nodeID, shards))
	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}},
		}:
		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
		}
		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_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()
	problems := hb.GetNewProblems()
	storage := hb.GetStorage()
	memMB := int(hb.GetMemoryTotalMb())
	if len(used) == 0 && len(problems) == 0 && len(storage) == 0 &&
		!hb.GetWeightsChanged() && memMB == 0 && !hb.GetEnginesChanged() &&
		!hb.GetProvisionChanged() && hb.GetLoad() == 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])
		}
	}
	if len(problems) > 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 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())
	}
	s.inv.Upsert(n)
}

// 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(),
		}
	}
	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 }

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

// StartShard implements reconciler.Dispatcher.
func (s *Service) StartShard(_ context.Context, cmd reconciler.ShardCommand) error {
	return s.send(cmd.NodeID, &llmgwv1.ControlMessage{
		Msg: &llmgwv1.ControlMessage_Start{Start: StartShardFromCommand(cmd)},
	})
}

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

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

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