agent: route JobTerminate to the connection a job migrated to

pkg/agent/agent_test.go

@@ -279,6 +279,87 @@ func TestJobTerminateReleasesJob(t *testing.T) {
 	})
 }
 
+// a reconnected worker reports its running jobs on its new connection, which must then receive their terminations.
+func TestMigratedJobTerminate(t *testing.T) {
+	const agentName = "test_agent"
+
+	// the server reads the ping only after handling the messages sent before it.
+	waitHandled := func(t *testing.T, w *testutils.AgentWorker) {
+		pongs := w.WorkerPongs.Observe()
+		defer pongs.Stop()
+		w.SendPing(&livekit.WorkerPing{Timestamp: time.Now().UnixMilli()})
+		select {
+		case <-pongs.Events():
+		case <-time.After(5 * time.Second):
+			require.Fail(t, "pong timeout")
+		}
+	}
+
+	for _, tc := range []struct {
+		name               string
+		closeBeforeMigrate bool
+	}{
+		{"old connection closes before migration", true},
+		{"old connection closes after migration", false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			bus := psrpc.NewLocalMessageBus()
+			client := must.Get(rpc.NewAgentInternalClient(bus))
+			server := testutils.NewTestServer(bus)
+			t.Cleanup(server.Close)
+
+			old := server.SimulateAgentWorker()
+			old.Register(agentName, livekit.JobType_JT_ROOM)
+			waitHandled(t, old)
+
+			job := &livekit.Job{
+				Id:         guid.New(guid.AgentJobPrefix),
+				DispatchId: guid.New(guid.AgentDispatchPrefix),
+				Type:       livekit.JobType_JT_ROOM,
+				Room:       &livekit.Room{},
+				AgentName:  agentName,
+			}
+			_, err := client.JobRequest(context.Background(), agentName, agent.RoomAgentTopic, job)
+			require.NoError(t, err)
+
+			closeOld := func() {
+				require.NoError(t, old.Close())
+				select {
+				case <-old.ConnectionDone():
+				case <-time.After(5 * time.Second):
+					require.Fail(t, "old connection teardown timeout")
+				}
+			}
+
+			if tc.closeBeforeMigrate {
+				closeOld()
+			}
+			reconnected := server.SimulateAgentWorker()
+			reconnected.Register(agentName, livekit.JobType_JT_ROOM)
+			reconnected.SendMigrateJob(&livekit.MigrateJobRequest{JobIds: []string{job.Id}})
+			waitHandled(t, reconnected)
+			if !tc.closeBeforeMigrate {
+				closeOld()
+			}
+
+			terminations := reconnected.JobTerminations.Observe()
+			defer terminations.Stop()
+			res, err := client.JobTerminate(context.Background(), job.Id, &rpc.JobTerminateRequest{
+				JobId:  job.Id,
+				Reason: rpc.JobTerminateReason_TERMINATION_REQUESTED,
+			}, psrpc.WithRequestTimeout(time.Second))
+			require.NoError(t, err)
+			require.Equal(t, livekit.JobStatus_JS_SUCCESS, res.State.Status)
+			select {
+			case m := <-terminations.Events():
+				require.Equal(t, job.Id, m.JobId)
+			case <-time.After(5 * time.Second):
+				require.Fail(t, "job termination timeout")
+			}
+		})
+	}
+}
+
 func TestConnectionClosedOnDispatchError(t *testing.T) {
 	t.Run("connection closed when unknown message type received", func(t *testing.T) {
 		bus := psrpc.NewLocalMessageBus()

pkg/agent/testutils/server.go

@@ -129,6 +129,7 @@ func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWo
 		JobAssignments:          events.NewObserverList[*livekit.JobAssignment](),
 		JobTerminations:         events.NewObserverList[*livekit.JobTermination](),
 		WorkerPongs:             events.NewObserverList[*livekit.WorkerPong](),
+		connectionDone:          make(chan struct{}),
 	}
 	w.ctx, w.cancel = context.WithCancel(context.Background())
 
@@ -137,7 +138,10 @@ func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWo
 	}
 
 	ctx := service.WithAPIKey(o.Context, &auth.ClaimGrants{}, "test")
-	go h.HandleConnection(ctx, w, agent.MakeWorkerRegistration())
+	go func() {
+		h.HandleConnection(ctx, w, agent.MakeWorkerRegistration())
+		close(w.connectionDone)
+	}()
 
 	return w
 }
@@ -190,6 +194,7 @@ type AgentWorker struct {
 	workerMessages chan *livekit.WorkerMessage
 	serverMessages deque.Deque[*livekit.ServerMessage]
 	jobs           map[string]*AgentJob
+	connectionDone chan struct{}
 
 	RegisterWorkerResponses *events.ObserverList[*livekit.RegisterWorkerResponse]
 	AvailabilityRequests    *events.ObserverList[*livekit.AvailabilityRequest]
@@ -224,6 +229,11 @@ func (w *AgentWorker) Closed() <-chan struct{} {
 	return w.fuse.Watch()
 }
 
+// ConnectionDone returns a channel that is closed when the server has finished handling the connection
+func (w *AgentWorker) ConnectionDone() <-chan struct{} {
+	return w.connectionDone
+}
+
 func (w *AgentWorker) SetReadDeadline(t time.Time) error {
 	w.mu.Lock()
 	defer w.mu.Unlock()

pkg/service/agenthandler.go

@@ -375,7 +375,10 @@ func (h *AgentHandler) deregisterWorker(w *agent.Worker) {
 
 	jobs := w.RunningJobs()
 	for jobID := range jobs {
-		h.deregisterJob(jobID)
+		// a job migrated to another connection is routed there.
+		if h.jobToWorker[jobID] == w {
+			h.deregisterJob(jobID)
+		}
 	}
 }
 
@@ -579,6 +582,19 @@ type agentHandlerWorker struct {
 	*agent.Worker
 }
 
+func (w *agentHandlerWorker) HandleMigrateJob(req *livekit.MigrateJobRequest) error {
+	if err := w.Worker.HandleMigrateJob(req); err != nil {
+		return err
+	}
+
+	for _, id := range req.JobIds {
+		if err := w.h.registerJob(w.Worker, livekit.JobID(id)); err != nil {
+			w.Logger().Errorw("failed to register JobTerminate handler", err, "jobID", id)
+		}
+	}
+	return nil
+}
+
 func (w *agentHandlerWorker) HandleUpdateJob(update *livekit.UpdateJobStatus) error {
 	if err := w.Worker.HandleUpdateJob(update); err != nil {
 		return err