livekit/pkg/service/agentservice.go

// Copyright 2024 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"context"
	"errors"
	"math/rand"
	"net/http"
	"slices"
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/emptypb"

	"github.com/livekit/livekit-server/pkg/agent"
	"github.com/livekit/livekit-server/pkg/agent/endpoint"
	"github.com/livekit/livekit-server/pkg/config"
	"github.com/livekit/livekit-server/pkg/routing"
	"github.com/livekit/livekit-server/pkg/rtc"
	"github.com/livekit/livekit-server/pkg/rtc/types"
	"github.com/livekit/livekit-server/version"
	"github.com/livekit/protocol/auth"
	"github.com/livekit/protocol/livekit"
	"github.com/livekit/protocol/logger"
	"github.com/livekit/protocol/rpc"
	"github.com/livekit/protocol/utils"
	"github.com/livekit/psrpc"
)

type AgentSocketUpgrader struct {
	websocket.Upgrader
}

func (u AgentSocketUpgrader) Upgrade(
	w http.ResponseWriter,
	r *http.Request,
	responseHeader http.Header,
) (
	conn *websocket.Conn,
	registration agent.WorkerRegistration,
	ok bool,
) {
	if u.CheckOrigin == nil {
		// allow connections from any origin, since script may be hosted anywhere
		// security is enforced by access tokens
		u.CheckOrigin = func(r *http.Request) bool {
			return true
		}
	}

	// reject non websocket requests
	if !websocket.IsWebSocketUpgrade(r) {
		w.WriteHeader(404)
		return
	}

	// require a claim
	claims := GetGrants(r.Context())
	if claims == nil || claims.Video == nil || !claims.Video.Agent {
		HandleError(w, r, http.StatusUnauthorized, rtc.ErrPermissionDenied)
		return
	}

	registration = agent.MakeWorkerRegistration()
	registration.ClientIP = GetClientIP(r)

	// upgrade
	conn, err := u.Upgrader.Upgrade(w, r, responseHeader)
	if err != nil {
		HandleError(w, r, http.StatusInternalServerError, err)
		return
	}

	if pv, err := strconv.Atoi(r.FormValue("protocol")); err == nil {
		registration.Protocol = agent.WorkerProtocolVersion(pv)
	}

	return conn, registration, true
}

func DispatchAgentWorkerSignal(c agent.SignalConn, h agent.WorkerSignalHandler, l logger.Logger) bool {
	req, _, err := c.ReadWorkerMessage()
	if err != nil {
		if IsWebSocketCloseError(err) {
			l.Debugw("worker closed WS connection", "wsError", err)
		} else {
			l.Errorw("error reading from websocket", err)
		}
		return false
	}

	if err := agent.DispatchWorkerSignal(req, h); err != nil {
		l.Warnw("unable to handle worker signal", err, "req", logger.Proto(req))
		return false
	}

	return true
}

func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, registration agent.WorkerRegistration, l logger.Logger, opts ...func(*agent.WorkerRegisterer)) (r agent.WorkerRegistration, ok bool) {
	wr := agent.NewWorkerRegisterer(c, serverInfo, registration)
	for _, opt := range opts {
		opt(wr)
	}
	if err := c.SetReadDeadline(wr.Deadline()); err != nil {
		return
	}
	for !wr.Registered() {
		if ok = DispatchAgentWorkerSignal(c, wr, l); !ok {
			return
		}
	}
	if err := c.SetReadDeadline(time.Time{}); err != nil {
		return
	}
	return wr.Registration(), true
}

type AgentService struct {
	upgrader AgentSocketUpgrader

	signalMessageSizeLimit int64

	*AgentHandler
}

type AgentHandler struct {
	agentServer rpc.AgentInternalServer
	// the server's only configured api key, when exactly one exists: the
	// unauthenticated scope for public endpoints regardless of which node holds
	// the workers
	singleAPIKey string
	mu           sync.Mutex
	logger       logger.Logger

	serverInfo  *livekit.ServerInfo
	workers     map[string]*agent.Worker
	jobToWorker map[livekit.JobID]*agent.Worker
	keyProvider auth.KeyProvider
	targetLoad  float32

	endpointRegistry *endpoint.Registry
	endpointsConfig  agent.EndpointsConfig

	namespaceWorkers    map[workerKey][]*agent.Worker
	roomKeyCount        int
	publisherKeyCount   int
	participantKeyCount int
	namespaces          []string // namespaces deprecated
	agentNames          []string

	roomTopic        string
	publisherTopic   string
	participantTopic string
}

type workerKey struct {
	agentName  string
	namespace  string
	jobType    livekit.JobType
	deployment string
}

func NewAgentService(
	conf *config.Config,
	currentNode routing.LocalNode,
	bus psrpc.MessageBus,
	keyProvider auth.KeyProvider,
) (*AgentService, error) {
	s := &AgentService{
		signalMessageSizeLimit: conf.Limit.AgentSignalMessageSizeLimit,
	}

	serverInfo := &livekit.ServerInfo{
		Edition:       livekit.ServerInfo_Standard,
		Version:       version.Version,
		Protocol:      types.CurrentProtocol,
		AgentProtocol: agent.CurrentProtocol,
		Region:        conf.Region,
		NodeId:        string(currentNode.NodeID()),
	}

	agentServer, err := rpc.NewAgentInternalServer(s, bus)
	if err != nil {
		return nil, err
	}
	s.AgentHandler = NewAgentHandler(
		agentServer,
		keyProvider,
		logger.GetLogger(),
		serverInfo,
		conf.Agents.TargetLoad,
		agent.RoomAgentTopic,
		agent.PublisherAgentTopic,
		agent.ParticipantAgentTopic,
	)
	s.AgentHandler.endpointsConfig = conf.Agents.Endpoints
	if len(conf.Keys) == 1 {
		for key := range conf.Keys {
			s.AgentHandler.singleAPIKey = key
		}
	}

	return s, nil
}

// EndpointFront is the /agents/{deployment}/{path...} handler backed by this
// node's attached workers. Project scope comes from validated grants when a
// token is present; unauthenticated requests reach public endpoints only.
func (s *AgentService) EndpointFront() http.Handler {
	front := endpoint.NewFront(s.endpointRegistry, func(r *http.Request) (string, bool) {
		if claims := GetGrants(r.Context()); claims != nil {
			return GetAPIKey(r.Context()), true
		}
		// unauthenticated: with a single configured key the scope is
		// unambiguous even when this node holds no registrations (multi-node)
		return s.singleAPIKey, false
	}, s.logger)
	front.WithSingleKeyFallback()
	return front
}

func (s *AgentService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	attach := r.URL.Query().Get("attach") != ""
	if conn, registration, ok := s.upgrader.Upgrade(w, r, nil); ok {
		// bound the size of a single signalling frame so an oversized message is
		// rejected by the transport before being fully buffered in memory. This
		// limits the compressed bytes read off the wire; the decompressed size is
		// bounded separately in WSSignalConnection.
		if s.signalMessageSizeLimit > 0 {
			conn.SetReadLimit(s.signalMessageSizeLimit)
		}

		if attach {
			// data wire: no registration handshake, no signal loop; the wire
			// speaks AgentHttp.Frame exclusively and the endpoint mux owns it
			// after a successful attach
			HandleEndpointAttach(s.endpointRegistry, NewEndpointWireConn(conn), s.wireParams(), nil)
			return
		}

		sigConn := NewWSSignalConnection(conn, s.signalMessageSizeLimit)
		defer sigConn.Close()
		s.HandleConnection(r.Context(), sigConn, registration)
	}
}

// endpointWireIdleTimeout is the read-side liveness bound on adopted wires:
// the SDK pings every 30s at the websocket level, so a wire silent for this
// long is dead and must release its pool slot instead of occupying it until a
// write happens to trip the write-stall deadline.
const endpointWireIdleTimeout = 2 * time.Minute

// endpointWireConn frames AgentHttp.Frame over a raw websocket: one frame per
// binary message.
type endpointWireConn struct {
	ws      *websocket.Conn
	writeMu sync.Mutex
}

// NewEndpointWireConn wraps an upgraded agent websocket as a data-plane wire.
func NewEndpointWireConn(ws *websocket.Conn) endpoint.WireConn {
	c := &endpointWireConn{ws: ws}
	ws.SetPingHandler(func(m string) error {
		_ = ws.SetReadDeadline(time.Now().Add(endpointWireIdleTimeout))
		return ws.WriteControl(websocket.PongMessage, []byte(m), time.Now().Add(10*time.Second))
	})
	return c
}

func (c *endpointWireConn) WriteFrame(f *livekit.AgentHttp_Frame) error {
	b, err := proto.Marshal(f)
	if err != nil {
		return err
	}
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.ws.WriteMessage(websocket.BinaryMessage, b)
}

func (c *endpointWireConn) ReadFrame() (*livekit.AgentHttp_Frame, error) {
	for {
		_ = c.ws.SetReadDeadline(time.Now().Add(endpointWireIdleTimeout))
		mt, b, err := c.ws.ReadMessage()
		if err != nil {
			return nil, err
		}
		if mt != websocket.BinaryMessage {
			continue
		}
		f := &livekit.AgentHttp_Frame{}
		if err := proto.Unmarshal(b, f); err != nil {
			return nil, err
		}
		return f, nil
	}
}

func (c *endpointWireConn) SetWriteDeadline(t time.Time) error {
	return c.ws.SetWriteDeadline(t)
}

func (c *endpointWireConn) SetReadDeadline(t time.Time) error {
	return c.ws.SetReadDeadline(t)
}

func (c *endpointWireConn) Close() error {
	return c.ws.Close()
}

// wireParams are this node's wire-level flow-control parameters, announced to
// the worker in each attach response.
func (s *AgentService) wireParams() endpoint.WireParams {
	p := endpoint.WireParams{
		CreditWindow:      s.endpointsConfig.CreditWindow,
		ConnectionWindow:  s.endpointsConfig.ConnectionWindow,
		MaxFrameSize:      s.endpointsConfig.MaxFrameSize,
		MaxStreamsPerConn: s.endpointsConfig.MaxStreamsPerConn,
	}.WithDefaults()
	// data frames ride the same websocket read limit as signalling: keep the
	// frame size safely under it so a full frame can never kill the wire
	if lim := s.signalMessageSizeLimit; lim > 0 && int64(p.MaxFrameSize) > lim/2 {
		p.MaxFrameSize = uint32(lim / 2)
	}
	return p
}

// AttachAdopter resolves an attach whose worker is unknown to the local
// registry - behind a load balancer the wire may land on a node other than the
// one holding the registration. The adopter may install a local registration
// (e.g. a satellite fetched from the holder); afterwards the attach is retried
// locally once. nil rejects unknown workers.
type AttachAdopter func(a *livekit.AgentHttp_AttachDataConnection) error

// HandleEndpointAttach adopts a worker-dialed data wire: the first frame is
// Attach on stream 0, validated against the registration's epoch and token; the
// response carries this node's wire parameters. Shared by OSS and cloud
// servers.
func HandleEndpointAttach(registry *endpoint.Registry, wire endpoint.WireConn, params endpoint.WireParams, adopt AttachAdopter) {
	if err := wire.SetReadDeadline(time.Now().Add(agent.RegisterTimeout)); err != nil {
		_ = wire.Close()
		return
	}
	f, err := wire.ReadFrame()
	if err != nil {
		_ = wire.Close()
		return
	}
	att, ok := f.Message.(*livekit.AgentHttp_Frame_Attach)
	if !ok || f.StreamId != 0 {
		_ = wire.Close()
		return
	}

	effective := params.WithDefaults()
	respond := func(errStr string) error {
		resp := &livekit.AgentHttp_AttachDataConnectionResponse{Error: errStr}
		if errStr == "" {
			resp.CreditWindow = effective.CreditWindow
			resp.ConnectionWindow = effective.ConnectionWindow
			resp.MaxFrameSize = effective.MaxFrameSize
			resp.MaxStreamsPerConn = effective.MaxStreamsPerConn
		}
		// the only wire write outside the scheduler: bound it too
		_ = wire.SetWriteDeadline(time.Now().Add(agent.RegisterTimeout))
		return wire.WriteFrame(&livekit.AgentHttp_Frame{
			Message: &livekit.AgentHttp_Frame_AttachResponse{AttachResponse: resp},
		})
	}

	a := att.Attach
	// the slot is reserved before the ack is written (never a success ack for a
	// wire that then loses the cap race) and adopted only after it, so the
	// worker cannot observe stream frames ahead of the attach outcome
	ticket, err := registry.BeginAttach(a.GetWorkerId(), a.GetInstanceId(), a.GetAttachToken())
	if (errors.Is(err, endpoint.ErrUnknownWorker) || errors.Is(err, endpoint.ErrWrongEpoch)) && adopt != nil {
		if aerr := adopt(a); aerr != nil {
			_ = respond(aerr.Error())
			_ = wire.Close()
			return
		}
		ticket, err = registry.BeginAttach(a.GetWorkerId(), a.GetInstanceId(), a.GetAttachToken())
	}
	if err != nil {
		_ = respond(err.Error())
		_ = wire.Close()
		return
	}
	if err := respond(""); err != nil {
		ticket.Abort()
		_ = wire.Close()
		return
	}
	if !ticket.Complete(wire, effective) {
		_ = wire.Close()
	}
}

func NewAgentHandler(
	agentServer rpc.AgentInternalServer,
	keyProvider auth.KeyProvider,
	logger logger.Logger,
	serverInfo *livekit.ServerInfo,
	targetLoad float32,
	roomTopic string,
	publisherTopic string,
	participantTopic string,
) *AgentHandler {
	return &AgentHandler{
		agentServer:      agentServer,
		logger:           logger.WithComponent("agents"),
		workers:          make(map[string]*agent.Worker),
		jobToWorker:      make(map[livekit.JobID]*agent.Worker),
		namespaceWorkers: make(map[workerKey][]*agent.Worker),
		serverInfo:       serverInfo,
		keyProvider:      keyProvider,
		targetLoad:       targetLoad,
		roomTopic:        roomTopic,
		publisherTopic:   publisherTopic,
		participantTopic: participantTopic,
		endpointRegistry: endpoint.NewRegistry(),
	}
}

// endpointSettings validates a registration's manifest and negotiates the
// data-plane settings, minting the epoch's attach token.
func (h *AgentHandler) endpointSettings(req *livekit.RegisterWorkerRequest) (*livekit.AgentHttp_AgentEndpointSettings, error) {
	if h.endpointsConfig.Disabled {
		return nil, errors.New("agent HTTP endpoints are disabled on this server")
	}
	if _, err := endpoint.ParseManifest(req.GetEndpoints()); err != nil {
		return nil, err
	}
	if req.GetInstanceId() == "" {
		return nil, errors.New("registrations with endpoints require an instance_id")
	}
	settings := endpoint.Settings{
		Protocol:      endpoint.CurrentProtocol,
		AttachToken:   endpoint.NewAttachToken(),
		DataConnCount: h.endpointsConfig.DataConnCount,
	}
	if settings.DataConnCount == 0 {
		settings.DataConnCount = endpoint.DefaultDataConnCount
	}
	return settings.Proto(), nil
}

func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalConn, registration agent.WorkerRegistration) {
	registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger, func(wr *agent.WorkerRegisterer) {
		wr.WithEndpointSettings(h.endpointSettings)
	})
	if !ok {
		return
	}

	apiKey := GetAPIKey(ctx)
	apiSecret := h.keyProvider.GetSecret(apiKey)

	worker := agent.NewWorker(registration, apiKey, apiSecret, conn, h.logger)
	h.registerWorker(worker)

	endpointReg := h.registerEndpoints(worker)

	handlerWorker := &agentHandlerWorker{h, worker}
	for ok := true; ok; {
		ok = DispatchAgentWorkerSignal(conn, handlerWorker, worker.Logger())
	}

	if endpointReg != nil {
		h.endpointRegistry.Deregister(endpointReg)
	}
	h.deregisterWorker(worker)
	worker.Close()
}

// registerEndpoints adopts a worker's endpoint manifest into the data-plane
// registry. The registration lives exactly as long as the control connection.
func (h *AgentHandler) registerEndpoints(w *agent.Worker) *endpoint.Registration {
	settings := w.EndpointSettings
	if settings == nil {
		return nil
	}
	manifest, err := endpoint.ParseManifest(w.Endpoints)
	if err != nil {
		// validated during the handshake; a failure here is a programming error
		w.Logger().Errorw("endpoint manifest failed to re-parse", err)
		return nil
	}
	reg := &endpoint.Registration{
		WorkerID:   w.ID,
		InstanceID: w.InstanceID,
		APIKey:     w.APIKey(),
		Deployment: w.Deployment,
		Manifest:   manifest,
		Endpoints:  w.Endpoints,
		Settings: endpoint.Settings{
			Protocol:      settings.GetProtocol(),
			AttachToken:   settings.GetAttachToken(),
			DataConnCount: settings.GetDataConnectionCount(),
		},
		Logger:   w.Logger(),
		Load:     w.Load,
		Draining: w.Draining,
	}
	if err := h.endpointRegistry.Register(reg); err != nil {
		w.Logger().Errorw("failed to register endpoints", err)
		return nil
	}
	w.Logger().Infow("endpoints registered",
		"deployment", w.Deployment, "routes", len(w.Endpoints))
	return reg
}

func (h *AgentHandler) registerWorker(w *agent.Worker) {
	h.mu.Lock()

	h.workers[w.ID] = w

	key := workerKey{w.AgentName, w.Namespace, w.JobType, w.Deployment}

	workers := h.namespaceWorkers[key]
	created := len(workers) == 0

	if created {
		nameTopic := agent.GetAgentTopic(w.AgentName, w.Namespace, w.Deployment)
		var typeTopic string
		switch w.JobType {
		case livekit.JobType_JT_ROOM:
			typeTopic = h.roomTopic
		case livekit.JobType_JT_PUBLISHER:
			typeTopic = h.publisherTopic
		case livekit.JobType_JT_PARTICIPANT:
			typeTopic = h.participantTopic
		}

		err := h.agentServer.RegisterJobRequestTopic(nameTopic, typeTopic)
		if err != nil {
			h.mu.Unlock()

			w.Logger().Errorw("failed to register job request topic", err)
			w.Close()
			return
		}

		switch w.JobType {
		case livekit.JobType_JT_ROOM:
			h.roomKeyCount++
		case livekit.JobType_JT_PUBLISHER:
			h.publisherKeyCount++
		case livekit.JobType_JT_PARTICIPANT:
			h.participantKeyCount++
		}

		h.namespaces = append(h.namespaces, w.Namespace)
		sort.Strings(h.namespaces)
		h.agentNames = append(h.agentNames, w.AgentName)
		sort.Strings(h.agentNames)
	}

	h.namespaceWorkers[key] = append(workers, w)
	h.mu.Unlock()

	h.logger.Infow("worker registered",
		"namespace", w.Namespace,
		"jobType", w.JobType,
		"agentName", w.AgentName,
		"workerID", w.ID,
	)
	if created {
		err := h.agentServer.PublishWorkerRegistered(context.Background(), agent.DefaultHandlerNamespace, &emptypb.Empty{})
		// TODO: when this happens, should we disconnect the worker so it'll retry?
		if err != nil {
			w.Logger().Errorw("failed to publish worker registered", err, "namespace", w.Namespace, "jobType", w.JobType, "agentName", w.AgentName)
		}
	}
}

func (h *AgentHandler) deregisterWorker(w *agent.Worker) {
	h.mu.Lock()
	defer h.mu.Unlock()

	delete(h.workers, w.ID)

	key := workerKey{w.AgentName, w.Namespace, w.JobType, w.Deployment}

	workers, ok := h.namespaceWorkers[key]
	if !ok {
		return
	}
	index := slices.Index(workers, w)
	if index == -1 {
		return
	}

	if len(workers) > 1 {
		h.namespaceWorkers[key] = slices.Delete(workers, index, index+1)
	} else {
		h.logger.Infow("last worker deregistered",
			"namespace", w.Namespace,
			"jobType", w.JobType,
			"agentName", w.AgentName,
			"workerID", w.ID,
		)
		delete(h.namespaceWorkers, key)

		topic := agent.GetAgentTopic(w.AgentName, w.Namespace, w.Deployment)

		switch w.JobType {
		case livekit.JobType_JT_ROOM:
			h.roomKeyCount--
			h.agentServer.DeregisterJobRequestTopic(topic, h.roomTopic)
		case livekit.JobType_JT_PUBLISHER:
			h.publisherKeyCount--
			h.agentServer.DeregisterJobRequestTopic(topic, h.publisherTopic)
		case livekit.JobType_JT_PARTICIPANT:
			h.participantKeyCount--
			h.agentServer.DeregisterJobRequestTopic(topic, h.participantTopic)
		}

		// agentNames and namespaces contains repeated entries for each agentNames/namespaces combinations
		if i := slices.Index(h.namespaces, w.Namespace); i != -1 {
			h.namespaces = slices.Delete(h.namespaces, i, i+1)
		}
		if i := slices.Index(h.agentNames, w.AgentName); i != -1 {
			h.agentNames = slices.Delete(h.agentNames, i, i+1)
		}
	}

	jobs := w.RunningJobs()
	for jobID := range jobs {
		h.deregisterJob(jobID)
	}
}

func (h *AgentHandler) deregisterJob(jobID livekit.JobID) {
	h.agentServer.DeregisterJobTerminateTopic(string(jobID))

	delete(h.jobToWorker, jobID)

	// TODO update dispatch state
}

func (h *AgentHandler) JobRequest(ctx context.Context, job *livekit.Job) (*rpc.JobRequestResponse, error) {
	logger := h.logger.WithUnlikelyValues(
		"jobID", job.Id,
		"namespace", job.Namespace,
		"agentName", job.AgentName,
		"jobType", job.Type.String(),
	)
	if job.Room != nil {
		logger = logger.WithValues("room", job.Room.Name, "roomID", job.Room.Sid)
	}
	if job.Participant != nil {
		logger = logger.WithValues("participant", job.Participant.Identity)
	}

	key := workerKey{job.AgentName, job.Namespace, job.Type, job.Deployment}
	attempted := make(map[*agent.Worker]struct{})
	for {
		selected, err := h.selectWorkerWeightedByLoad(key, attempted)
		if err != nil {
			logger.Warnw("no worker available to handle job", err)
			return nil, psrpc.NewError(psrpc.ResourceExhausted, err)
		}

		logger := logger.WithValues("workerID", selected.ID)
		attempted[selected] = struct{}{}

		state, err := selected.AssignJob(ctx, job, nil)
		switch state.GetStatus() {
		case livekit.JobStatus_JS_RUNNING:
			logger.Infow("assigned job to worker", "apiKey", selected.APIKey())
			h.mu.Lock()
			h.jobToWorker[livekit.JobID(job.Id)] = selected
			h.mu.Unlock()

			err = h.agentServer.RegisterJobTerminateTopic(job.Id)
			if err != nil {
				logger.Errorw("failed to register JobTerminate handler", err)
			}
			fallthrough
		case livekit.JobStatus_JS_SUCCESS:
			return &rpc.JobRequestResponse{
				State: state,
			}, nil
		default:
			retry := utils.ErrorIsOneOf(err, agent.ErrWorkerNotAvailable, agent.ErrWorkerClosed)
			logger.Warnw("failed to assign job to worker", err, "retry", retry)
			if !retry {
				return nil, err
			}
		}
	}
}

func (h *AgentHandler) JobRequestAffinity(ctx context.Context, job *livekit.Job) float32 {
	h.mu.Lock()
	defer h.mu.Unlock()

	var affinity float32
	for _, w := range h.workers {
		if w.AgentName != job.AgentName || w.Namespace != job.Namespace || w.JobType != job.Type || w.Deployment != job.Deployment {
			continue
		}

		if w.Status() == livekit.WorkerStatus_WS_AVAILABLE {
			affinity += max(0, h.targetLoad-w.Load())
		}
	}

	return affinity
}

func (h *AgentHandler) JobTerminate(ctx context.Context, req *rpc.JobTerminateRequest) (*rpc.JobTerminateResponse, error) {
	h.mu.Lock()
	w := h.jobToWorker[livekit.JobID(req.JobId)]
	h.mu.Unlock()

	if w == nil {
		return nil, psrpc.NewErrorf(psrpc.NotFound, "no worker for jobID")
	}

	state, err := w.TerminateJob(livekit.JobID(req.JobId), req.Reason)
	if err != nil {
		return nil, err
	}

	return &rpc.JobTerminateResponse{
		State: state,
	}, nil
}

func (h *AgentHandler) CheckEnabled(ctx context.Context, req *rpc.CheckEnabledRequest) (*rpc.CheckEnabledResponse, error) {
	h.mu.Lock()
	defer h.mu.Unlock()

	// This doesn't return the full agentName -> namespace mapping, which can cause some unnecessary RPC.
	// namespaces are however deprecated.
	return &rpc.CheckEnabledResponse{
		Namespaces:         slices.Compact(slices.Clone(h.namespaces)),
		AgentNames:         slices.Compact(slices.Clone(h.agentNames)),
		RoomEnabled:        h.roomKeyCount != 0,
		PublisherEnabled:   h.publisherKeyCount != 0,
		ParticipantEnabled: h.participantKeyCount != 0,
	}, nil
}

func (h *AgentHandler) DrainConnections(interval time.Duration, force bool) {
	// Snapshot workers and release the lock before Close. Worker.Close closes the
	// signal connection, which unblocks HandleConnection's read loop so it can call
	// deregisterWorker — that needs h.mu. Holding the lock across Close deadlocks drain.
	h.mu.Lock()
	workers := make([]*agent.Worker, 0, len(h.workers))
	for _, w := range h.workers {
		workers = append(workers, w)
	}
	h.mu.Unlock()

	if !force {
		// jitter drain start
		time.Sleep(time.Duration(rand.Int63n(int64(interval))))

		t := time.NewTicker(interval)
		defer t.Stop()

		for _, w := range workers {
			w.Close()
			<-t.C
		}
	} else {
		// drain as quickly as possible when forced
		for _, w := range workers {
			w.Close()
		}
	}
}

func (h *AgentHandler) selectWorkerWeightedByLoad(key workerKey, ignore map[*agent.Worker]struct{}) (*agent.Worker, error) {
	h.mu.Lock()
	defer h.mu.Unlock()

	workers, ok := h.namespaceWorkers[key]
	if !ok {
		return nil, errors.New("no workers available")
	}

	normalizedLoads := make(map[*agent.Worker]float32)
	var availableSum float32
	for _, w := range workers {
		if _, ok := ignore[w]; !ok && w.Status() == livekit.WorkerStatus_WS_AVAILABLE {
			normalizedLoads[w] = max(0, 1-w.Load())
			availableSum += normalizedLoads[w]
		}
	}

	if availableSum == 0 {
		return nil, errors.New("no workers with sufficient capacity")
	}

	currentSum := rand.Float32() * availableSum
	for w, load := range normalizedLoads {
		if currentSum -= load; currentSum <= 0 {
			return w, nil
		}
	}
	return workers[0], nil
}

var _ agent.WorkerSignalHandler = (*agentHandlerWorker)(nil)

type agentHandlerWorker struct {
	h *AgentHandler
	*agent.Worker
}

func (w *agentHandlerWorker) HandleUpdateJob(update *livekit.UpdateJobStatus) error {
	if err := w.Worker.HandleUpdateJob(update); err != nil {
		return err
	}

	if agent.JobStatusIsEnded(update.Status) {
		w.h.mu.Lock()
		w.h.deregisterJob(livekit.JobID(update.JobId))
		w.h.mu.Unlock()
	}
	return nil
}