From ff40282a5e28b778800bb2f4054a1dbd99ed3d3a Mon Sep 17 00:00:00 2001
From: Jason Lernerman <43508353+u9g@users.noreply.github.com>
Date: Wed, 21 Oct 2026 22:47:05 -0400
Subject: [PATCH] agent: track jobs migrated onto a new worker connection

HandleMigrateJob was a no-op, so a job reported by a reconnected worker
was missing from runningJobs. TerminateJob then returned ErrJobNotFound
without sending JobTermination, and UpdateJobStatus rejected the
worker's own updates for that job.
---
 pkg/agent/worker.go      | 22 ++++++++++++++++++++--
 pkg/agent/worker_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/pkg/agent/worker.go b/pkg/agent/worker.go
index bdd5f5c40..7b7411803 100644
--- a/pkg/agent/worker.go
+++ b/pkg/agent/worker.go
@@ -662,8 +662,26 @@ func (w *Worker) Draining() bool {
 	return w.draining
 }
 
+// MigrateJobRequest carries only job IDs, so a migrated job is tracked without its room or dispatch details.
 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
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	now := time.Now().UnixNano()
+	for _, id := range req.JobIds {
+		jobID := livekit.JobID(id)
+		if _, ok := w.runningJobs[jobID]; ok {
+			continue
+		}
+		w.runningJobs[jobID] = &livekit.Job{
+			Id: id,
+			State: &livekit.JobState{
+				Status:    livekit.JobStatus_JS_RUNNING,
+				UpdatedAt: now,
+				WorkerId:  w.ID,
+				AgentId:   w.AgentID,
+			},
+		}
+	}
 	return nil
 }
diff --git a/pkg/agent/worker_test.go b/pkg/agent/worker_test.go
index 5b9947892..6167f742d 100644
--- a/pkg/agent/worker_test.go
+++ b/pkg/agent/worker_test.go
@@ -12,6 +12,8 @@ import (
 	"github.com/livekit/livekit-server/pkg/agent"
 	"github.com/livekit/livekit-server/pkg/agent/endpoint/wire"
 	"github.com/livekit/protocol/livekit"
+	"github.com/livekit/protocol/logger"
+	"github.com/livekit/protocol/rpc"
 )
 
 type nopSignalConn struct{}
@@ -76,3 +78,33 @@ func TestHandleRegisterWithoutEndpointsIgnoresReservedNames(t *testing.T) {
 		}), name)
 	}
 }
+
+type recordingSignalConn struct {
+	nopSignalConn
+	sent []*livekit.ServerMessage
+}
+
+func (c *recordingSignalConn) WriteServerMessage(m *livekit.ServerMessage) (int, error) {
+	c.sent = append(c.sent, m)
+	return 0, nil
+}
+
+// jobs reported on a new connection can be terminated and updated there.
+func TestHandleMigrateJob(t *testing.T) {
+	conn := &recordingSignalConn{}
+	w := agent.NewWorker(agent.WorkerRegistration{}, "key", "secret", conn, logger.GetLogger())
+
+	require.NoError(t, w.HandleMigrateJob(&livekit.MigrateJobRequest{JobIds: []string{"AJ_terminated", "AJ_updated"}}))
+	require.Equal(t, 2, w.RunningJobCount())
+
+	state, err := w.TerminateJob("AJ_terminated", rpc.JobTerminateReason_TERMINATION_REQUESTED)
+	require.NoError(t, err)
+	require.Equal(t, livekit.JobStatus_JS_SUCCESS, state.Status)
+	require.Len(t, conn.sent, 1)
+	require.Equal(t, "AJ_terminated", conn.sent[0].GetTermination().GetJobId())
+
+	state, err = w.UpdateJobStatus(&livekit.UpdateJobStatus{JobId: "AJ_updated", Status: livekit.JobStatus_JS_SUCCESS})
+	require.NoError(t, err)
+	require.Equal(t, livekit.JobStatus_JS_SUCCESS, state.Status)
+	require.Zero(t, w.RunningJobCount())
+}