From c7cd9705b78d5bf1636b3144466349692c5a47ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9o=20Monnom?= <theo.monnom@outlook.com>
Date: Thu, 20 Aug 2026 12:48:57 -0700
Subject: [PATCH] agent: attach adopter hook for wires landing on foreign nodes

Behind a load balancer a data wire may reach a node other than the one
holding the worker's registration; the adopter lets that layer (cloud)
install a local satellite registration before the attach is retried locally.
---
 pkg/service/agentservice.go | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/pkg/service/agentservice.go b/pkg/service/agentservice.go
index 9d338c947..2e7fcf118 100644
--- a/pkg/service/agentservice.go
+++ b/pkg/service/agentservice.go
@@ -252,7 +252,7 @@ func (s *AgentService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			// 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())
+			HandleEndpointAttach(s.endpointRegistry, NewEndpointWireConn(conn), s.wireParams(), nil)
 			return
 		}
 
@@ -342,11 +342,18 @@ func (s *AgentService) wireParams() endpoint.WireParams {
 	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) {
+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
@@ -383,6 +390,14 @@ func HandleEndpointAttach(registry *endpoint.Registry, wire endpoint.WireConn, p
 	// 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) && 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()