Squashed 'livekit-server/' content from commit 154b4d26

git-subtree-dir: livekit-server
git-subtree-split: 154b4d26b769c68a03c096124094b97bf61a996f

pkg/agent/agent_test.go

@@ -0,0 +1,183 @@
+package agent_test
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
+
+	"github.com/livekit/livekit-server/pkg/agent"
+	"github.com/livekit/livekit-server/pkg/agent/testutils"
+	"github.com/livekit/protocol/livekit"
+	"github.com/livekit/protocol/rpc"
+	"github.com/livekit/protocol/utils/guid"
+	"github.com/livekit/protocol/utils/must"
+	"github.com/livekit/psrpc"
+)
+
+func TestAgent(t *testing.T) {
+	t.Run("dispatched jobs are assigned to a worker", func(t *testing.T) {
+		bus := psrpc.NewLocalMessageBus()
+
+		client := must.Get(rpc.NewAgentInternalClient(bus))
+		server := testutils.NewTestServer(bus)
+		t.Cleanup(server.Close)
+
+		worker := server.SimulateAgentWorker()
+		worker.Register("test", livekit.JobType_JT_ROOM)
+		jobAssignments := worker.JobAssignments.Observe()
+
+		job := &livekit.Job{
+			Id:         guid.New(guid.AgentJobPrefix),
+			DispatchId: guid.New(guid.AgentDispatchPrefix),
+			Type:       livekit.JobType_JT_ROOM,
+			Room:       &livekit.Room{},
+			AgentName:  "test",
+		}
+		_, err := client.JobRequest(context.Background(), "test", agent.RoomAgentTopic, job)
+		require.NoError(t, err)
+
+		select {
+		case a := <-jobAssignments.Events():
+			require.EqualValues(t, job.Id, a.Job.Id)
+		case <-time.After(time.Second):
+			require.Fail(t, "job assignment timeout")
+		}
+	})
+}
+
+func testBatchJobRequest(t require.TestingT, batchSize int, totalJobs int, client rpc.AgentInternalClient, workers []*testutils.AgentWorker) <-chan struct{} {
+	var assigned atomic.Uint32
+	done := make(chan struct{})
+	for _, w := range workers {
+		assignments := w.JobAssignments.Observe()
+		go func() {
+			defer assignments.Stop()
+			for {
+				select {
+				case <-done:
+				case <-assignments.Events():
+					if assigned.Inc() == uint32(totalJobs) {
+						close(done)
+					}
+				}
+			}
+		}()
+	}
+
+	// wait for agent registration
+	time.Sleep(100 * time.Millisecond)
+
+	var wg sync.WaitGroup
+	for i := 0; i < totalJobs; i += batchSize {
+		wg.Add(1)
+		go func(start int) {
+			defer wg.Done()
+			for j := start; j < start+batchSize && j < totalJobs; j++ {
+				job := &livekit.Job{
+					Id:         guid.New(guid.AgentJobPrefix),
+					DispatchId: guid.New(guid.AgentDispatchPrefix),
+					Type:       livekit.JobType_JT_ROOM,
+					Room:       &livekit.Room{},
+					AgentName:  "test",
+				}
+				_, err := client.JobRequest(context.Background(), "test", agent.RoomAgentTopic, job)
+				require.NoError(t, err)
+			}
+		}(i)
+	}
+	wg.Wait()
+
+	return done
+}
+
+func TestAgentLoadBalancing(t *testing.T) {
+	t.Run("jobs are distributed normally with baseline worker load", func(t *testing.T) {
+		totalWorkers := 5
+		totalJobs := 100
+
+		bus := psrpc.NewLocalMessageBus()
+
+		client := must.Get(rpc.NewAgentInternalClient(bus))
+		t.Cleanup(client.Close)
+		server := testutils.NewTestServer(bus)
+		t.Cleanup(server.Close)
+
+		agents := make([]*testutils.AgentWorker, totalWorkers)
+		for i := 0; i < totalWorkers; i++ {
+			agents[i] = server.SimulateAgentWorker(
+				testutils.WithLabel(fmt.Sprintf("agent-%d", i)),
+				testutils.WithJobLoad(testutils.NewStableJobLoad(0.01)),
+			)
+			agents[i].Register("test", livekit.JobType_JT_ROOM)
+		}
+
+		select {
+		case <-testBatchJobRequest(t, 10, totalJobs, client, agents):
+		case <-time.After(time.Second):
+			require.Fail(t, "job assignment timeout")
+		}
+
+		jobCount := make(map[string]int)
+		for _, w := range agents {
+			jobCount[w.Label] = len(w.Jobs())
+		}
+
+		// check that jobs are distributed normally
+		for i := 0; i < totalWorkers; i++ {
+			label := fmt.Sprintf("agent-%d", i)
+			require.GreaterOrEqual(t, jobCount[label], 0)
+			require.Less(t, jobCount[label], 35) // three std deviations from the mean is 32
+		}
+	})
+
+	t.Run("jobs are distributed with variable and overloaded worker load", func(t *testing.T) {
+		totalWorkers := 4
+		totalJobs := 15
+
+		bus := psrpc.NewLocalMessageBus()
+
+		client := must.Get(rpc.NewAgentInternalClient(bus))
+		t.Cleanup(client.Close)
+		server := testutils.NewTestServer(bus)
+		t.Cleanup(server.Close)
+
+		agents := make([]*testutils.AgentWorker, totalWorkers)
+		for i := 0; i < totalWorkers; i++ {
+			label := fmt.Sprintf("agent-%d", i)
+			if i%2 == 0 {
+				// make sure we have some workers that can accept jobs
+				agents[i] = server.SimulateAgentWorker(testutils.WithLabel(label))
+			} else {
+				agents[i] = server.SimulateAgentWorker(testutils.WithLabel(label), testutils.WithDefaultWorkerLoad(0.9))
+			}
+			agents[i].Register("test", livekit.JobType_JT_ROOM)
+		}
+
+		select {
+		case <-testBatchJobRequest(t, 1, totalJobs, client, agents):
+		case <-time.After(time.Second):
+			require.Fail(t, "job assignment timeout")
+		}
+
+		jobCount := make(map[string]int)
+		for _, w := range agents {
+			jobCount[w.Label] = len(w.Jobs())
+		}
+
+		for i := 0; i < totalWorkers; i++ {
+			label := fmt.Sprintf("agent-%d", i)
+
+			if i%2 == 0 {
+				require.GreaterOrEqual(t, jobCount[label], 2)
+			} else {
+				require.Equal(t, 0, jobCount[label])
+			}
+			require.GreaterOrEqual(t, jobCount[label], 0)
+		}
+	})
+}

pkg/agent/client.go

@@ -0,0 +1,336 @@
+// 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 agent
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gammazero/workerpool"
+	"google.golang.org/protobuf/types/known/emptypb"
+
+	serverutils "github.com/livekit/livekit-server/pkg/utils"
+	"github.com/livekit/protocol/livekit"
+	"github.com/livekit/protocol/logger"
+	"github.com/livekit/protocol/rpc"
+	"github.com/livekit/protocol/utils"
+	"github.com/livekit/psrpc"
+)
+
+const (
+	EnabledCacheTTL         = 1 * time.Minute
+	RoomAgentTopic          = "room"
+	PublisherAgentTopic     = "publisher"
+	ParticipantAgentTopic   = "participant"
+	DefaultHandlerNamespace = ""
+
+	CheckEnabledTimeout = 5 * time.Second
+)
+
+var jobTypeTopics = map[livekit.JobType]string{
+	livekit.JobType_JT_ROOM:        RoomAgentTopic,
+	livekit.JobType_JT_PUBLISHER:   PublisherAgentTopic,
+	livekit.JobType_JT_PARTICIPANT: ParticipantAgentTopic,
+}
+
+type Client interface {
+	// LaunchJob starts a room or participant job on an agent.
+	// it will launch a job once for each worker in each namespace
+	LaunchJob(ctx context.Context, desc *JobRequest) *serverutils.IncrementalDispatcher[*livekit.Job]
+	TerminateJob(ctx context.Context, jobID string, reason rpc.JobTerminateReason) (*livekit.JobState, error)
+	Stop() error
+}
+
+type JobRequest struct {
+	DispatchId string
+	JobType    livekit.JobType
+	Room       *livekit.Room
+	// only set for participant jobs
+	Participant *livekit.ParticipantInfo
+	Metadata    string
+	AgentName   string
+}
+
+type agentClient struct {
+	client rpc.AgentInternalClient
+
+	mu sync.RWMutex
+
+	// cache response to avoid constantly checking with controllers
+	// cache is invalidated with AgentRegistered updates
+	roomNamespaces        *serverutils.IncrementalDispatcher[string] // deprecated
+	publisherNamespaces   *serverutils.IncrementalDispatcher[string] // deprecated
+	participantNamespaces *serverutils.IncrementalDispatcher[string] // deprecated
+	roomAgentNames        *serverutils.IncrementalDispatcher[string]
+	publisherAgentNames   *serverutils.IncrementalDispatcher[string]
+	participantAgentNames *serverutils.IncrementalDispatcher[string]
+
+	enabledExpiresAt time.Time
+
+	workers *workerpool.WorkerPool
+
+	invalidateSub psrpc.Subscription[*emptypb.Empty]
+	subDone       chan struct{}
+}
+
+func NewAgentClient(bus psrpc.MessageBus) (Client, error) {
+	client, err := rpc.NewAgentInternalClient(bus)
+	if err != nil {
+		return nil, err
+	}
+
+	c := &agentClient{
+		client:  client,
+		workers: workerpool.New(50),
+		subDone: make(chan struct{}),
+	}
+
+	sub, err := c.client.SubscribeWorkerRegistered(context.Background(), DefaultHandlerNamespace)
+	if err != nil {
+		return nil, err
+	}
+
+	c.invalidateSub = sub
+
+	go func() {
+		// invalidate cache
+		for range sub.Channel() {
+			c.mu.Lock()
+			c.roomNamespaces = nil
+			c.publisherNamespaces = nil
+			c.participantNamespaces = nil
+			c.roomAgentNames = nil
+			c.publisherAgentNames = nil
+			c.participantAgentNames = nil
+			c.mu.Unlock()
+		}
+
+		c.subDone <- struct{}{}
+	}()
+
+	return c, nil
+}
+
+func (c *agentClient) LaunchJob(ctx context.Context, desc *JobRequest) *serverutils.IncrementalDispatcher[*livekit.Job] {
+	var wg sync.WaitGroup
+	ret := serverutils.NewIncrementalDispatcher[*livekit.Job]()
+	defer func() {
+		c.workers.Submit(func() {
+			wg.Wait()
+			ret.Done()
+		})
+	}()
+
+	jobTypeTopic, ok := jobTypeTopics[desc.JobType]
+	if !ok {
+		return ret
+	}
+
+	dispatcher := c.getDispatcher(desc.AgentName, desc.JobType)
+
+	if dispatcher == nil {
+		logger.Infow("not dispatching agent job since no worker is available",
+			"agentName", desc.AgentName,
+			"jobType", desc.JobType,
+			"room", desc.Room.Name,
+			"roomID", desc.Room.Sid)
+		return ret
+	}
+
+	dispatcher.ForEach(func(curNs string) {
+		topic := GetAgentTopic(desc.AgentName, curNs)
+
+		wg.Add(1)
+		c.workers.Submit(func() {
+			defer wg.Done()
+			// The cached agent parameters do not provide the exact combination of available job type/agent name/namespace, so some of the JobRequest RPC may not trigger any worker
+			job := &livekit.Job{
+				Id:          utils.NewGuid(utils.AgentJobPrefix),
+				DispatchId:  desc.DispatchId,
+				Type:        desc.JobType,
+				Room:        desc.Room,
+				Participant: desc.Participant,
+				Namespace:   curNs,
+				AgentName:   desc.AgentName,
+				Metadata:    desc.Metadata,
+			}
+			resp, err := c.client.JobRequest(context.Background(), topic, jobTypeTopic, job)
+			if err != nil {
+				logger.Infow("failed to send job request", "error", err, "namespace", curNs, "jobType", desc.JobType, "agentName", desc.AgentName)
+				return
+			}
+			job.State = resp.State
+			ret.Add(job)
+		})
+	})
+
+	return ret
+}
+
+func (c *agentClient) TerminateJob(ctx context.Context, jobID string, reason rpc.JobTerminateReason) (*livekit.JobState, error) {
+	resp, err := c.client.JobTerminate(context.Background(), jobID, &rpc.JobTerminateRequest{
+		JobId:  jobID,
+		Reason: reason,
+	})
+	if err != nil {
+		logger.Infow("failed to send job request", "error", err, "jobID", jobID)
+		return nil, err
+	}
+
+	return resp.State, nil
+}
+
+func (c *agentClient) getDispatcher(agName string, jobType livekit.JobType) *serverutils.IncrementalDispatcher[string] {
+	c.mu.Lock()
+
+	if time.Since(c.enabledExpiresAt) > EnabledCacheTTL || c.roomNamespaces == nil ||
+		c.publisherNamespaces == nil || c.participantNamespaces == nil || c.roomAgentNames == nil || c.publisherAgentNames == nil || c.participantAgentNames == nil {
+		c.enabledExpiresAt = time.Now()
+		c.roomNamespaces = serverutils.NewIncrementalDispatcher[string]()
+		c.publisherNamespaces = serverutils.NewIncrementalDispatcher[string]()
+		c.participantNamespaces = serverutils.NewIncrementalDispatcher[string]()
+		c.roomAgentNames = serverutils.NewIncrementalDispatcher[string]()
+		c.publisherAgentNames = serverutils.NewIncrementalDispatcher[string]()
+		c.participantAgentNames = serverutils.NewIncrementalDispatcher[string]()
+
+		go c.checkEnabled(c.roomNamespaces, c.publisherNamespaces, c.participantNamespaces, c.roomAgentNames, c.publisherAgentNames, c.participantAgentNames)
+	}
+
+	var target *serverutils.IncrementalDispatcher[string]
+	var agentNames *serverutils.IncrementalDispatcher[string]
+	switch jobType {
+	case livekit.JobType_JT_ROOM:
+		target = c.roomNamespaces
+		agentNames = c.roomAgentNames
+	case livekit.JobType_JT_PUBLISHER:
+		target = c.publisherNamespaces
+		agentNames = c.publisherAgentNames
+	case livekit.JobType_JT_PARTICIPANT:
+		target = c.participantNamespaces
+		agentNames = c.participantAgentNames
+	}
+	c.mu.Unlock()
+
+	if agName == "" {
+		// if no agent name is given, we would need to dispatch backwards compatible mode
+		// which means dispatching to each of the namespaces
+		return target
+	}
+
+	done := make(chan *serverutils.IncrementalDispatcher[string], 1)
+	c.workers.Submit(func() {
+		agentNames.ForEach(func(ag string) {
+			if ag == agName {
+				select {
+				case done <- target:
+				default:
+				}
+			}
+		})
+		select {
+		case done <- nil:
+		default:
+		}
+	})
+
+	return <-done
+}
+
+func (c *agentClient) checkEnabled(roomNamespaces, publisherNamespaces, participantNamespaces, roomAgentNames, publisherAgentNames, participantAgentNames *serverutils.IncrementalDispatcher[string]) {
+	defer roomNamespaces.Done()
+	defer publisherNamespaces.Done()
+	defer participantNamespaces.Done()
+	defer roomAgentNames.Done()
+	defer publisherAgentNames.Done()
+	defer participantAgentNames.Done()
+
+	resChan, err := c.client.CheckEnabled(context.Background(), &rpc.CheckEnabledRequest{}, psrpc.WithRequestTimeout(CheckEnabledTimeout))
+	if err != nil {
+		logger.Errorw("failed to check enabled", err)
+		return
+	}
+
+	roomNSMap := make(map[string]bool)
+	publisherNSMap := make(map[string]bool)
+	participantNSMap := make(map[string]bool)
+	roomAgMap := make(map[string]bool)
+	publisherAgMap := make(map[string]bool)
+	participantAgMap := make(map[string]bool)
+
+	for r := range resChan {
+		if r.Result.GetRoomEnabled() {
+			for _, ns := range r.Result.GetNamespaces() {
+				if _, ok := roomNSMap[ns]; !ok {
+					roomNamespaces.Add(ns)
+					roomNSMap[ns] = true
+				}
+			}
+			for _, ag := range r.Result.GetAgentNames() {
+				if _, ok := roomAgMap[ag]; !ok {
+					roomAgentNames.Add(ag)
+					roomAgMap[ag] = true
+				}
+			}
+		}
+		if r.Result.GetPublisherEnabled() {
+			for _, ns := range r.Result.GetNamespaces() {
+				if _, ok := publisherNSMap[ns]; !ok {
+					publisherNamespaces.Add(ns)
+					publisherNSMap[ns] = true
+				}
+			}
+			for _, ag := range r.Result.GetAgentNames() {
+				if _, ok := publisherAgMap[ag]; !ok {
+					publisherAgentNames.Add(ag)
+					publisherAgMap[ag] = true
+				}
+			}
+		}
+		if r.Result.GetParticipantEnabled() {
+			for _, ns := range r.Result.GetNamespaces() {
+				if _, ok := participantNSMap[ns]; !ok {
+					participantNamespaces.Add(ns)
+					participantNSMap[ns] = true
+				}
+			}
+			for _, ag := range r.Result.GetAgentNames() {
+				if _, ok := participantAgMap[ag]; !ok {
+					participantAgentNames.Add(ag)
+					participantAgMap[ag] = true
+				}
+			}
+		}
+	}
+}
+
+func (c *agentClient) Stop() error {
+	_ = c.invalidateSub.Close()
+	<-c.subDone
+	return nil
+}
+
+func GetAgentTopic(agentName, namespace string) string {
+	if agentName == "" {
+		// Backward compatibility
+		return namespace
+	} else if namespace == "" {
+		// Forward compatibility once the namespace field is removed from the worker SDK
+		return agentName
+	} else {
+		return fmt.Sprintf("%s_%s", agentName, namespace)
+	}
+}

pkg/agent/testutils/server.go

@@ -0,0 +1,483 @@
+package testutils
+
+import (
+	"context"
+	"errors"
+	"io"
+	"math"
+	"math/rand/v2"
+	"sync"
+	"time"
+
+	"github.com/frostbyte73/core"
+	"github.com/gammazero/deque"
+	"golang.org/x/exp/maps"
+
+	"github.com/livekit/livekit-server/pkg/agent"
+	"github.com/livekit/livekit-server/pkg/config"
+	"github.com/livekit/livekit-server/pkg/routing"
+	"github.com/livekit/livekit-server/pkg/service"
+	"github.com/livekit/protocol/auth"
+	"github.com/livekit/protocol/livekit"
+	"github.com/livekit/protocol/utils/events"
+	"github.com/livekit/protocol/utils/guid"
+	"github.com/livekit/protocol/utils/must"
+	"github.com/livekit/protocol/utils/options"
+	"github.com/livekit/psrpc"
+)
+
+type AgentService interface {
+	HandleConnection(context.Context, agent.SignalConn, agent.WorkerProtocolVersion)
+	DrainConnections(time.Duration)
+}
+
+type TestServer struct {
+	AgentService
+}
+
+func NewTestServer(bus psrpc.MessageBus) *TestServer {
+	localNode, _ := routing.NewLocalNode(nil)
+	return NewTestServerWithService(must.Get(service.NewAgentService(
+		&config.Config{Region: "test"},
+		localNode,
+		bus,
+		auth.NewSimpleKeyProvider("test", "verysecretsecret"),
+	)))
+}
+
+func NewTestServerWithService(s AgentService) *TestServer {
+	return &TestServer{s}
+}
+
+type SimulatedWorkerOptions struct {
+	Context            context.Context
+	Label              string
+	SupportResume      bool
+	DefaultJobLoad     float32
+	JobLoadThreshold   float32
+	DefaultWorkerLoad  float32
+	HandleAvailability func(AgentJobRequest)
+	HandleAssignment   func(*livekit.Job) JobLoad
+}
+
+type SimulatedWorkerOption func(*SimulatedWorkerOptions)
+
+func WithContext(ctx context.Context) SimulatedWorkerOption {
+	return func(o *SimulatedWorkerOptions) {
+		o.Context = ctx
+	}
+}
+
+func WithLabel(label string) SimulatedWorkerOption {
+	return func(o *SimulatedWorkerOptions) {
+		o.Label = label
+	}
+}
+
+func WithJobAvailabilityHandler(h func(AgentJobRequest)) SimulatedWorkerOption {
+	return func(o *SimulatedWorkerOptions) {
+		o.HandleAvailability = h
+	}
+}
+
+func WithJobAssignmentHandler(h func(*livekit.Job) JobLoad) SimulatedWorkerOption {
+	return func(o *SimulatedWorkerOptions) {
+		o.HandleAssignment = h
+	}
+}
+
+func WithJobLoad(l JobLoad) SimulatedWorkerOption {
+	return WithJobAssignmentHandler(func(j *livekit.Job) JobLoad { return l })
+}
+
+func WithDefaultWorkerLoad(load float32) SimulatedWorkerOption {
+	return func(o *SimulatedWorkerOptions) {
+		o.DefaultWorkerLoad = load
+	}
+}
+
+func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWorker {
+	o := &SimulatedWorkerOptions{
+		Context:            context.Background(),
+		Label:              guid.New("TEST_AGENT_"),
+		DefaultJobLoad:     0.1,
+		JobLoadThreshold:   0.8,
+		DefaultWorkerLoad:  0.0,
+		HandleAvailability: func(r AgentJobRequest) { r.Accept() },
+		HandleAssignment:   func(j *livekit.Job) JobLoad { return nil },
+	}
+	options.Apply(o, opts)
+
+	w := &AgentWorker{
+		workerMessages:         make(chan *livekit.WorkerMessage, 1),
+		jobs:                   map[string]*AgentJob{},
+		SimulatedWorkerOptions: o,
+
+		RegisterWorkerResponses: events.NewObserverList[*livekit.RegisterWorkerResponse](),
+		AvailabilityRequests:    events.NewObserverList[*livekit.AvailabilityRequest](),
+		JobAssignments:          events.NewObserverList[*livekit.JobAssignment](),
+		JobTerminations:         events.NewObserverList[*livekit.JobTermination](),
+		WorkerPongs:             events.NewObserverList[*livekit.WorkerPong](),
+	}
+	w.ctx, w.cancel = context.WithCancel(context.Background())
+
+	if o.DefaultWorkerLoad > 0.0 {
+		w.sendStatus()
+	}
+
+	ctx := service.WithAPIKey(o.Context, &auth.ClaimGrants{}, "test")
+	go h.HandleConnection(ctx, w, agent.CurrentProtocol)
+
+	return w
+}
+
+func (h *TestServer) Close() {
+	h.DrainConnections(1)
+}
+
+var _ agent.SignalConn = (*AgentWorker)(nil)
+
+type JobLoad interface {
+	Load() float32
+}
+
+type AgentJob struct {
+	*livekit.Job
+	JobLoad
+}
+
+type AgentJobRequest struct {
+	w *AgentWorker
+	*livekit.AvailabilityRequest
+}
+
+func (r AgentJobRequest) Accept() {
+	identity := guid.New("PI_")
+	r.w.SendAvailability(&livekit.AvailabilityResponse{
+		JobId:               r.Job.Id,
+		Available:           true,
+		SupportsResume:      r.w.SupportResume,
+		ParticipantName:     identity,
+		ParticipantIdentity: identity,
+	})
+}
+
+func (r AgentJobRequest) Reject() {
+	r.w.SendAvailability(&livekit.AvailabilityResponse{
+		JobId:     r.Job.Id,
+		Available: false,
+	})
+}
+
+type AgentWorker struct {
+	*SimulatedWorkerOptions
+
+	fuse           core.Fuse
+	mu             sync.Mutex
+	ctx            context.Context
+	cancel         context.CancelFunc
+	workerMessages chan *livekit.WorkerMessage
+	serverMessages deque.Deque[*livekit.ServerMessage]
+	jobs           map[string]*AgentJob
+
+	RegisterWorkerResponses *events.ObserverList[*livekit.RegisterWorkerResponse]
+	AvailabilityRequests    *events.ObserverList[*livekit.AvailabilityRequest]
+	JobAssignments          *events.ObserverList[*livekit.JobAssignment]
+	JobTerminations         *events.ObserverList[*livekit.JobTermination]
+	WorkerPongs             *events.ObserverList[*livekit.WorkerPong]
+}
+
+func (w *AgentWorker) statusWorker() {
+	t := time.NewTicker(2 * time.Second)
+	defer t.Stop()
+
+	for !w.fuse.IsBroken() {
+		w.sendStatus()
+		<-t.C
+	}
+}
+
+func (w *AgentWorker) Close() error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.fuse.Break()
+	return nil
+}
+
+func (w *AgentWorker) SetReadDeadline(t time.Time) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if !w.fuse.IsBroken() {
+		cancel := w.cancel
+		if t.IsZero() {
+			w.ctx, w.cancel = context.WithCancel(context.Background())
+		} else {
+			w.ctx, w.cancel = context.WithDeadline(context.Background(), t)
+		}
+		cancel()
+	}
+	return nil
+}
+
+func (w *AgentWorker) ReadWorkerMessage() (*livekit.WorkerMessage, int, error) {
+	for {
+		w.mu.Lock()
+		ctx := w.ctx
+		w.mu.Unlock()
+
+		select {
+		case <-w.fuse.Watch():
+			return nil, 0, io.EOF
+		case <-ctx.Done():
+			if err := ctx.Err(); errors.Is(err, context.DeadlineExceeded) {
+				return nil, 0, err
+			}
+		case m := <-w.workerMessages:
+			return m, 0, nil
+		}
+	}
+}
+
+func (w *AgentWorker) WriteServerMessage(m *livekit.ServerMessage) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.serverMessages.PushBack(m)
+	if w.serverMessages.Len() == 1 {
+		go w.handleServerMessages()
+	}
+	return 0, nil
+}
+
+func (w *AgentWorker) handleServerMessages() {
+	w.mu.Lock()
+	for w.serverMessages.Len() != 0 {
+		m := w.serverMessages.Front()
+		w.mu.Unlock()
+
+		switch m := m.Message.(type) {
+		case *livekit.ServerMessage_Register:
+			w.handleRegister(m.Register)
+		case *livekit.ServerMessage_Availability:
+			w.handleAvailability(m.Availability)
+		case *livekit.ServerMessage_Assignment:
+			w.handleAssignment(m.Assignment)
+		case *livekit.ServerMessage_Termination:
+			w.handleTermination(m.Termination)
+		case *livekit.ServerMessage_Pong:
+			w.handlePong(m.Pong)
+		}
+
+		w.mu.Lock()
+		w.serverMessages.PopFront()
+	}
+	w.mu.Unlock()
+}
+
+func (w *AgentWorker) handleRegister(m *livekit.RegisterWorkerResponse) {
+	w.RegisterWorkerResponses.Emit(m)
+}
+
+func (w *AgentWorker) handleAvailability(m *livekit.AvailabilityRequest) {
+	w.AvailabilityRequests.Emit(m)
+	if w.HandleAvailability != nil {
+		w.HandleAvailability(AgentJobRequest{w, m})
+	} else {
+		AgentJobRequest{w, m}.Accept()
+	}
+}
+
+func (w *AgentWorker) handleAssignment(m *livekit.JobAssignment) {
+	w.JobAssignments.Emit(m)
+
+	var load JobLoad
+	if w.HandleAssignment != nil {
+		load = w.HandleAssignment(m.Job)
+	}
+
+	if load == nil {
+		load = NewStableJobLoad(w.DefaultJobLoad)
+	}
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.jobs[m.Job.Id] = &AgentJob{m.Job, load}
+}
+
+func (w *AgentWorker) handleTermination(m *livekit.JobTermination) {
+	w.JobTerminations.Emit(m)
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	delete(w.jobs, m.JobId)
+}
+
+func (w *AgentWorker) handlePong(m *livekit.WorkerPong) {
+	w.WorkerPongs.Emit(m)
+}
+
+func (w *AgentWorker) sendMessage(m *livekit.WorkerMessage) {
+	select {
+	case <-w.fuse.Watch():
+	case w.workerMessages <- m:
+	}
+}
+
+func (w *AgentWorker) SendRegister(m *livekit.RegisterWorkerRequest) {
+	w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Register{
+		Register: m,
+	}})
+}
+
+func (w *AgentWorker) SendAvailability(m *livekit.AvailabilityResponse) {
+	w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Availability{
+		Availability: m,
+	}})
+}
+
+func (w *AgentWorker) SendUpdateWorker(m *livekit.UpdateWorkerStatus) {
+	w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateWorker{
+		UpdateWorker: m,
+	}})
+}
+
+func (w *AgentWorker) SendUpdateJob(m *livekit.UpdateJobStatus) {
+	w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateJob{
+		UpdateJob: m,
+	}})
+}
+
+func (w *AgentWorker) SendPing(m *livekit.WorkerPing) {
+	w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Ping{
+		Ping: m,
+	}})
+}
+
+func (w *AgentWorker) SendSimulateJob(m *livekit.SimulateJobRequest) {
+	w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_SimulateJob{
+		SimulateJob: m,
+	}})
+}
+
+func (w *AgentWorker) SendMigrateJob(m *livekit.MigrateJobRequest) {
+	w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_MigrateJob{
+		MigrateJob: m,
+	}})
+}
+
+func (w *AgentWorker) sendStatus() {
+	w.mu.Lock()
+	var load float32
+	jobCount := len(w.jobs)
+
+	if len(w.jobs) == 0 {
+		load = w.DefaultWorkerLoad
+	} else {
+		for _, j := range w.jobs {
+			load += j.Load()
+		}
+	}
+	w.mu.Unlock()
+
+	status := livekit.WorkerStatus_WS_AVAILABLE
+	if load > w.JobLoadThreshold {
+		status = livekit.WorkerStatus_WS_FULL
+	}
+
+	w.SendUpdateWorker(&livekit.UpdateWorkerStatus{
+		Status:   &status,
+		Load:     load,
+		JobCount: uint32(jobCount),
+	})
+}
+
+func (w *AgentWorker) Register(agentName string, jobType livekit.JobType) {
+	w.SendRegister(&livekit.RegisterWorkerRequest{
+		Type:      jobType,
+		AgentName: agentName,
+	})
+	go w.statusWorker()
+}
+
+func (w *AgentWorker) SimulateRoomJob(roomName string) {
+	w.SendSimulateJob(&livekit.SimulateJobRequest{
+		Type: livekit.JobType_JT_ROOM,
+		Room: &livekit.Room{
+			Sid:  guid.New(guid.RoomPrefix),
+			Name: roomName,
+		},
+	})
+}
+
+func (w *AgentWorker) Jobs() []*AgentJob {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return maps.Values(w.jobs)
+}
+
+type stableJobLoad struct {
+	load float32
+}
+
+func NewStableJobLoad(load float32) JobLoad {
+	return stableJobLoad{load}
+}
+
+func (s stableJobLoad) Load() float32 {
+	return s.load
+}
+
+type periodicJobLoad struct {
+	amplitude float64
+	period    time.Duration
+	epoch     time.Time
+}
+
+func NewPeriodicJobLoad(max float32, period time.Duration) JobLoad {
+	return periodicJobLoad{
+		amplitude: float64(max / 2),
+		period:    period,
+		epoch:     time.Now().Add(-time.Duration(rand.Int64N(int64(period)))),
+	}
+}
+
+func (s periodicJobLoad) Load() float32 {
+	a := math.Sin(time.Since(s.epoch).Seconds() / s.period.Seconds() * math.Pi * 2)
+	return float32(s.amplitude + a*s.amplitude)
+}
+
+type uniformRandomJobLoad struct {
+	min, max float32
+	rng      func() float64
+}
+
+func NewUniformRandomJobLoad(min, max float32) JobLoad {
+	return uniformRandomJobLoad{min, max, rand.Float64}
+}
+
+func NewUniformRandomJobLoadWithRNG(min, max float32, rng *rand.Rand) JobLoad {
+	return uniformRandomJobLoad{min, max, rng.Float64}
+}
+
+func (s uniformRandomJobLoad) Load() float32 {
+	return rand.Float32()*(s.max-s.min) + s.min
+}
+
+type normalRandomJobLoad struct {
+	mean, stddev float64
+	rng          func() float64
+}
+
+func NewNormalRandomJobLoad(mean, stddev float64) JobLoad {
+	return normalRandomJobLoad{mean, stddev, rand.Float64}
+}
+
+func NewNormalRandomJobLoadWithRNG(mean, stddev float64, rng *rand.Rand) JobLoad {
+	return normalRandomJobLoad{mean, stddev, rng.Float64}
+}
+
+func (s normalRandomJobLoad) Load() float32 {
+	u := 1 - s.rng()
+	v := s.rng()
+	z := math.Sqrt(-2*math.Log(u)) * math.Cos(2*math.Pi*v)
+	return float32(max(0, z*s.stddev+s.mean))
+}

pkg/agent/worker.go

@@ -0,0 +1,559 @@
+// 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 agent
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"google.golang.org/protobuf/proto"
+
+	pagent "github.com/livekit/protocol/agent"
+	"github.com/livekit/protocol/livekit"
+	"github.com/livekit/protocol/logger"
+	"github.com/livekit/protocol/rpc"
+	"github.com/livekit/protocol/utils"
+	"github.com/livekit/protocol/utils/guid"
+	"github.com/livekit/psrpc"
+)
+
+var (
+	ErrUnimplementedWrorkerSignal = errors.New("unimplemented worker signal")
+	ErrUnknownWorkerSignal        = errors.New("unknown worker signal")
+	ErrUnknownJobType             = errors.New("unknown job type")
+	ErrJobNotFound                = psrpc.NewErrorf(psrpc.NotFound, "no running job for given jobID")
+	ErrWorkerClosed               = errors.New("worker closed")
+	ErrWorkerNotAvailable         = errors.New("worker not available")
+	ErrAvailabilityTimeout        = errors.New("agent worker availability timeout")
+	ErrDuplicateJobAssignment     = errors.New("duplicate job assignment")
+)
+
+type WorkerProtocolVersion int
+
+const CurrentProtocol = 1
+
+const (
+	RegisterTimeout  = 10 * time.Second
+	AssignJobTimeout = 10 * time.Second
+)
+
+type SignalConn interface {
+	WriteServerMessage(msg *livekit.ServerMessage) (int, error)
+	ReadWorkerMessage() (*livekit.WorkerMessage, int, error)
+	SetReadDeadline(time.Time) error
+	Close() error
+}
+
+func JobStatusIsEnded(s livekit.JobStatus) bool {
+	return s == livekit.JobStatus_JS_SUCCESS || s == livekit.JobStatus_JS_FAILED
+}
+
+type WorkerSignalHandler interface {
+	HandleRegister(*livekit.RegisterWorkerRequest) error
+	HandleAvailability(*livekit.AvailabilityResponse) error
+	HandleUpdateJob(*livekit.UpdateJobStatus) error
+	HandleSimulateJob(*livekit.SimulateJobRequest) error
+	HandlePing(*livekit.WorkerPing) error
+	HandleUpdateWorker(*livekit.UpdateWorkerStatus) error
+	HandleMigrateJob(*livekit.MigrateJobRequest) error
+}
+
+func DispatchWorkerSignal(req *livekit.WorkerMessage, h WorkerSignalHandler) error {
+	switch m := req.Message.(type) {
+	case *livekit.WorkerMessage_Register:
+		return h.HandleRegister(m.Register)
+	case *livekit.WorkerMessage_Availability:
+		return h.HandleAvailability(m.Availability)
+	case *livekit.WorkerMessage_UpdateJob:
+		return h.HandleUpdateJob(m.UpdateJob)
+	case *livekit.WorkerMessage_SimulateJob:
+		return h.HandleSimulateJob(m.SimulateJob)
+	case *livekit.WorkerMessage_Ping:
+		return h.HandlePing(m.Ping)
+	case *livekit.WorkerMessage_UpdateWorker:
+		return h.HandleUpdateWorker(m.UpdateWorker)
+	case *livekit.WorkerMessage_MigrateJob:
+		return h.HandleMigrateJob(m.MigrateJob)
+	default:
+		return ErrUnknownWorkerSignal
+	}
+}
+
+var _ WorkerSignalHandler = (*UnimplementedWorkerSignalHandler)(nil)
+
+type UnimplementedWorkerSignalHandler struct{}
+
+func (UnimplementedWorkerSignalHandler) HandleRegister(*livekit.RegisterWorkerRequest) error {
+	return fmt.Errorf("%w: Register", ErrUnimplementedWrorkerSignal)
+}
+func (UnimplementedWorkerSignalHandler) HandleAvailability(*livekit.AvailabilityResponse) error {
+	return fmt.Errorf("%w: Availability", ErrUnimplementedWrorkerSignal)
+}
+func (UnimplementedWorkerSignalHandler) HandleUpdateJob(*livekit.UpdateJobStatus) error {
+	return fmt.Errorf("%w: UpdateJob", ErrUnimplementedWrorkerSignal)
+}
+func (UnimplementedWorkerSignalHandler) HandleSimulateJob(*livekit.SimulateJobRequest) error {
+	return fmt.Errorf("%w: SimulateJob", ErrUnimplementedWrorkerSignal)
+}
+func (UnimplementedWorkerSignalHandler) HandlePing(*livekit.WorkerPing) error {
+	return fmt.Errorf("%w: Ping", ErrUnimplementedWrorkerSignal)
+}
+func (UnimplementedWorkerSignalHandler) HandleUpdateWorker(*livekit.UpdateWorkerStatus) error {
+	return fmt.Errorf("%w: UpdateWorker", ErrUnimplementedWrorkerSignal)
+}
+func (UnimplementedWorkerSignalHandler) HandleMigrateJob(*livekit.MigrateJobRequest) error {
+	return fmt.Errorf("%w: MigrateJob", ErrUnimplementedWrorkerSignal)
+}
+
+type WorkerPingHandler struct {
+	UnimplementedWorkerSignalHandler
+	conn SignalConn
+}
+
+func (h WorkerPingHandler) HandlePing(ping *livekit.WorkerPing) error {
+	_, err := h.conn.WriteServerMessage(&livekit.ServerMessage{
+		Message: &livekit.ServerMessage_Pong{
+			Pong: &livekit.WorkerPong{
+				LastTimestamp: ping.Timestamp,
+				Timestamp:     time.Now().UnixMilli(),
+			},
+		},
+	})
+	return err
+}
+
+type WorkerRegistration struct {
+	Protocol    WorkerProtocolVersion
+	ID          string
+	Version     string
+	AgentName   string
+	Namespace   string
+	JobType     livekit.JobType
+	Permissions *livekit.ParticipantPermission
+}
+
+var _ WorkerSignalHandler = (*WorkerRegisterer)(nil)
+
+type WorkerRegisterer struct {
+	WorkerPingHandler
+	serverInfo *livekit.ServerInfo
+	protocol   WorkerProtocolVersion
+	deadline   time.Time
+
+	registration WorkerRegistration
+	registered   bool
+}
+
+func NewWorkerRegisterer(conn SignalConn, serverInfo *livekit.ServerInfo, protocol WorkerProtocolVersion) *WorkerRegisterer {
+	return &WorkerRegisterer{
+		WorkerPingHandler: WorkerPingHandler{conn: conn},
+		serverInfo:        serverInfo,
+		protocol:          protocol,
+		deadline:          time.Now().Add(RegisterTimeout),
+	}
+}
+
+func (h *WorkerRegisterer) Deadline() time.Time {
+	return h.deadline
+}
+
+func (h *WorkerRegisterer) Registration() WorkerRegistration {
+	return h.registration
+}
+
+func (h *WorkerRegisterer) Registered() bool {
+	return h.registered
+}
+
+func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) error {
+	if !livekit.IsJobType(req.GetType()) {
+		return ErrUnknownJobType
+	}
+
+	permissions := req.AllowedPermissions
+	if permissions == nil {
+		permissions = &livekit.ParticipantPermission{
+			CanSubscribe:      true,
+			CanPublish:        true,
+			CanPublishData:    true,
+			CanUpdateMetadata: true,
+		}
+	}
+
+	h.registration = WorkerRegistration{
+		Protocol:    h.protocol,
+		ID:          guid.New(guid.AgentWorkerPrefix),
+		Version:     req.Version,
+		AgentName:   req.AgentName,
+		Namespace:   req.GetNamespace(),
+		JobType:     req.GetType(),
+		Permissions: permissions,
+	}
+	h.registered = true
+
+	_, err := h.conn.WriteServerMessage(&livekit.ServerMessage{
+		Message: &livekit.ServerMessage_Register{
+			Register: &livekit.RegisterWorkerResponse{
+				WorkerId:   h.registration.ID,
+				ServerInfo: h.serverInfo,
+			},
+		},
+	})
+	return err
+}
+
+var _ WorkerSignalHandler = (*Worker)(nil)
+
+type Worker struct {
+	WorkerPingHandler
+	WorkerRegistration
+
+	apiKey    string
+	apiSecret string
+	logger    logger.Logger
+
+	ctx    context.Context
+	cancel context.CancelFunc
+	closed chan struct{}
+
+	mu     sync.Mutex
+	load   float32
+	status livekit.WorkerStatus
+
+	runningJobs  map[livekit.JobID]*livekit.Job
+	availability map[livekit.JobID]chan *livekit.AvailabilityResponse
+}
+
+func NewWorker(
+	registration WorkerRegistration,
+	apiKey string,
+	apiSecret string,
+	conn SignalConn,
+	logger logger.Logger,
+) *Worker {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	return &Worker{
+		WorkerPingHandler:  WorkerPingHandler{conn: conn},
+		WorkerRegistration: registration,
+		apiKey:             apiKey,
+		apiSecret:          apiSecret,
+		logger: logger.WithValues(
+			"workerID", registration.ID,
+			"agentName", registration.AgentName,
+			"jobType", registration.JobType.String(),
+		),
+
+		ctx:    ctx,
+		cancel: cancel,
+		closed: make(chan struct{}),
+
+		runningJobs:  make(map[livekit.JobID]*livekit.Job),
+		availability: make(map[livekit.JobID]chan *livekit.AvailabilityResponse),
+	}
+}
+
+func (w *Worker) sendRequest(req *livekit.ServerMessage) {
+	if _, err := w.conn.WriteServerMessage(req); err != nil {
+		w.logger.Warnw("error writing to websocket", err)
+	}
+}
+
+func (w *Worker) Status() livekit.WorkerStatus {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.status
+}
+
+func (w *Worker) Load() float32 {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.load
+}
+
+func (w *Worker) Logger() logger.Logger {
+	return w.logger
+}
+
+func (w *Worker) RunningJobs() map[livekit.JobID]*livekit.Job {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	jobs := make(map[livekit.JobID]*livekit.Job, len(w.runningJobs))
+	for k, v := range w.runningJobs {
+		jobs[k] = v
+	}
+	return jobs
+}
+
+func (w *Worker) RunningJobCount() int {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return len(w.runningJobs)
+}
+
+func (w *Worker) GetJobState(jobID livekit.JobID) (*livekit.JobState, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	j, ok := w.runningJobs[jobID]
+	if !ok {
+		return nil, ErrJobNotFound
+	}
+	return utils.CloneProto(j.State), nil
+}
+
+func (w *Worker) AssignJob(ctx context.Context, job *livekit.Job) (*livekit.JobState, error) {
+	availCh := make(chan *livekit.AvailabilityResponse, 1)
+	job = utils.CloneProto(job)
+	jobID := livekit.JobID(job.Id)
+
+	w.mu.Lock()
+	if _, ok := w.availability[jobID]; ok {
+		w.mu.Unlock()
+		return nil, ErrDuplicateJobAssignment
+	}
+
+	w.availability[jobID] = availCh
+	w.mu.Unlock()
+
+	defer func() {
+		w.mu.Lock()
+		delete(w.availability, jobID)
+		w.mu.Unlock()
+	}()
+
+	if job.State == nil {
+		job.State = &livekit.JobState{}
+	}
+	now := time.Now()
+	job.State.UpdatedAt = now.UnixNano()
+	job.State.StartedAt = now.UnixNano()
+	job.State.Status = livekit.JobStatus_JS_RUNNING
+
+	w.sendRequest(&livekit.ServerMessage{Message: &livekit.ServerMessage_Availability{
+		Availability: &livekit.AvailabilityRequest{Job: job},
+	}})
+
+	timeout := time.NewTimer(AssignJobTimeout)
+	defer timeout.Stop()
+
+	// See handleAvailability for the response
+	select {
+	case res := <-availCh:
+		if !res.Available {
+			return nil, ErrWorkerNotAvailable
+		}
+
+		job.State.ParticipantIdentity = res.ParticipantIdentity
+
+		token, err := pagent.BuildAgentToken(
+			w.apiKey,
+			w.apiSecret,
+			job.Room.Name,
+			res.ParticipantIdentity,
+			res.ParticipantName,
+			res.ParticipantMetadata,
+			res.ParticipantAttributes,
+			w.Permissions,
+		)
+		if err != nil {
+			w.logger.Errorw("failed to build agent token", err)
+			return nil, err
+		}
+
+		// In OSS, Url is nil, and the used API Key is the same as the one used to connect the worker
+		w.sendRequest(&livekit.ServerMessage{Message: &livekit.ServerMessage_Assignment{
+			Assignment: &livekit.JobAssignment{Job: job, Url: nil, Token: token},
+		}})
+
+		state := utils.CloneProto(job.State)
+
+		w.mu.Lock()
+		w.runningJobs[jobID] = job
+		w.mu.Unlock()
+
+		// TODO sweep jobs that are never started. We can't do this until all SDKs actually update the the JOB state
+
+		return state, nil
+	case <-timeout.C:
+		return nil, ErrAvailabilityTimeout
+	case <-w.ctx.Done():
+		return nil, ErrWorkerClosed
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
+func (w *Worker) TerminateJob(jobID livekit.JobID, reason rpc.JobTerminateReason) (*livekit.JobState, error) {
+	w.mu.Lock()
+	_, ok := w.runningJobs[jobID]
+	w.mu.Unlock()
+
+	if !ok {
+		return nil, ErrJobNotFound
+	}
+
+	w.sendRequest(&livekit.ServerMessage{Message: &livekit.ServerMessage_Termination{
+		Termination: &livekit.JobTermination{
+			JobId: string(jobID),
+		},
+	}})
+
+	status := livekit.JobStatus_JS_SUCCESS
+	errorStr := ""
+	if reason == rpc.JobTerminateReason_AGENT_LEFT_ROOM {
+		status = livekit.JobStatus_JS_FAILED
+		errorStr = "agent worker left the room"
+	}
+
+	return w.UpdateJobStatus(&livekit.UpdateJobStatus{
+		JobId:  string(jobID),
+		Status: status,
+		Error:  errorStr,
+	})
+}
+
+func (w *Worker) UpdateMetadata(metadata string) {
+	w.logger.Debugw("worker metadata updated", nil, "metadata", metadata)
+}
+
+func (w *Worker) IsClosed() bool {
+	select {
+	case <-w.closed:
+		return true
+	default:
+		return false
+	}
+}
+
+func (w *Worker) Close() {
+	w.mu.Lock()
+	if w.IsClosed() {
+		w.mu.Unlock()
+		return
+	}
+
+	w.logger.Infow("closing worker", "workerID", w.ID, "jobType", w.JobType, "agentName", w.AgentName)
+
+	close(w.closed)
+	w.cancel()
+	_ = w.conn.Close()
+	w.mu.Unlock()
+}
+
+func (w *Worker) HandleAvailability(res *livekit.AvailabilityResponse) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	jobID := livekit.JobID(res.JobId)
+	availCh, ok := w.availability[jobID]
+	if !ok {
+		w.logger.Warnw("received availability response for unknown job", nil, "jobID", jobID)
+		return nil
+	}
+
+	availCh <- res
+	delete(w.availability, jobID)
+
+	return nil
+}
+
+func (w *Worker) HandleUpdateJob(update *livekit.UpdateJobStatus) error {
+	_, err := w.UpdateJobStatus(update)
+	if err != nil {
+		// treating this as a debug message only
+		// this can happen if the Room closes first, which would delete the agent dispatch
+		// that would mark the job as successful. subsequent updates from the same worker
+		// would not be able to find the same jobID.
+		w.logger.Debugw("received job update for unknown job", "jobID", update.JobId)
+	}
+	return nil
+}
+
+func (w *Worker) UpdateJobStatus(update *livekit.UpdateJobStatus) (*livekit.JobState, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	jobID := livekit.JobID(update.JobId)
+	job, ok := w.runningJobs[jobID]
+	if !ok {
+		return nil, psrpc.NewErrorf(psrpc.NotFound, "received job update for unknown job")
+	}
+
+	now := time.Now()
+	job.State.UpdatedAt = now.UnixNano()
+
+	if job.State.Status == livekit.JobStatus_JS_PENDING && update.Status != livekit.JobStatus_JS_PENDING {
+		job.State.StartedAt = now.UnixNano()
+	}
+
+	job.State.Status = update.Status
+	job.State.Error = update.Error
+
+	if JobStatusIsEnded(update.Status) {
+		job.State.EndedAt = now.UnixNano()
+		delete(w.runningJobs, jobID)
+
+		w.logger.Infow("job ended", "jobID", update.JobId, "status", update.Status, "error", update.Error)
+	}
+
+	return proto.Clone(job.State).(*livekit.JobState), nil
+}
+
+func (w *Worker) HandleSimulateJob(simulate *livekit.SimulateJobRequest) error {
+	jobType := livekit.JobType_JT_ROOM
+	if simulate.Participant != nil {
+		jobType = livekit.JobType_JT_PUBLISHER
+	}
+
+	job := &livekit.Job{
+		Id:          guid.New(guid.AgentJobPrefix),
+		Type:        jobType,
+		Room:        simulate.Room,
+		Participant: simulate.Participant,
+		Namespace:   w.Namespace,
+		AgentName:   w.AgentName,
+	}
+
+	go func() {
+		_, err := w.AssignJob(w.ctx, job)
+		if err != nil {
+			w.logger.Errorw("unable to simulate job", err, "jobID", job.Id)
+		}
+	}()
+
+	return nil
+}
+
+func (w *Worker) HandleUpdateWorker(update *livekit.UpdateWorkerStatus) error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if status := update.Status; status != nil && w.status != *status {
+		w.status = *status
+		w.Logger().Debugw("worker status changed", "status", w.status)
+	}
+	w.load = update.GetLoad()
+
+	return nil
+}
+
+func (w *Worker) HandleMigrateJob(req *livekit.MigrateJobRequest) error {
+	// TODO(theomonnom): On OSS this is not implemented
+	// We could maybe just move a specific job to another worker
+	return nil
+}