From 751d85c9bf5b07d4d64bc89d925e569a21228d2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9o=20Monnom?= Date: Sun, 30 Aug 2026 18:59:30 -0700 Subject: [PATCH] agent endpoints: WebTransport data plane Replace the capsule mux / credit / attach-pool machinery with one WebTransport (QUIC) session per worker: a control stream carries registration and status, and each HTTP exchange rides its own QUIC stream (native mux, per-stream flow control, half-close, and reset). Encode the HTTP method into the cross-node route filter so candidate selection is method-aware. Keep the legacy WebSocket control transport for backward compatibility. --- go.mod | 6 +- go.sum | 20 +- pkg/agent/config.go | 21 +- pkg/agent/endpoint/conformance/worker.go | 853 +++++------------------ pkg/agent/endpoint/conn.go | 386 ---------- pkg/agent/endpoint/control.go | 64 ++ pkg/agent/endpoint/front.go | 70 +- pkg/agent/endpoint/manifest.go | 19 + pkg/agent/endpoint/registry.go | 290 ++------ pkg/agent/endpoint/registry_test.go | 78 ++- pkg/agent/endpoint/routematch.go | 56 +- pkg/agent/endpoint/routematch_test.go | 120 ++-- pkg/agent/endpoint/scheduler.go | 296 -------- pkg/agent/endpoint/session.go | 78 +++ pkg/agent/endpoint/stream.go | 366 ---------- pkg/agent/endpoint/template.go | 9 +- pkg/agent/endpoint/template_test.go | 9 + pkg/agent/endpoint/webtransport.go | 160 +++++ pkg/agent/endpoint/webtransport_test.go | 184 +++++ pkg/agent/endpoint/wire.go | 122 ---- pkg/agent/worker.go | 10 +- pkg/service/agentendpoint_test.go | 137 +++- pkg/service/agentservice.go | 258 +------ pkg/service/agentwebtransport.go | 234 +++++++ pkg/service/server.go | 12 + 25 files changed, 1373 insertions(+), 2485 deletions(-) delete mode 100644 pkg/agent/endpoint/conn.go create mode 100644 pkg/agent/endpoint/control.go delete mode 100644 pkg/agent/endpoint/scheduler.go create mode 100644 pkg/agent/endpoint/session.go delete mode 100644 pkg/agent/endpoint/stream.go create mode 100644 pkg/agent/endpoint/webtransport.go create mode 100644 pkg/agent/endpoint/webtransport_test.go delete mode 100644 pkg/agent/endpoint/wire.go create mode 100644 pkg/service/agentwebtransport.go diff --git a/go.mod b/go.mod index 72471b2a2..5b31f4014 100644 --- a/go.mod +++ b/go.mod @@ -21,7 +21,7 @@ require ( github.com/jxskiss/base62 v1.1.0 github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731 github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0 - github.com/livekit/protocol v1.50.5-0.20260822193149-0c325211bedd + github.com/livekit/protocol v1.50.5-0.20260831015717-ba22cc1ce840 github.com/livekit/psrpc v0.7.3 github.com/mackerelio/go-osstat v0.2.8 github.com/magefile/mage v1.17.2 @@ -42,6 +42,8 @@ require ( github.com/pion/webrtc/v4 v4.2.18 github.com/pkg/errors v0.9.1 github.com/prometheus/client_golang v1.24.1 + github.com/quic-go/quic-go v0.61.0 + github.com/quic-go/webtransport-go v0.12.0 github.com/redis/go-redis/v9 v9.22.0 github.com/rs/cors v1.11.1 github.com/stretchr/testify v1.11.1 @@ -67,6 +69,7 @@ require ( github.com/containerd/errdefs v1.0.0 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect github.com/distribution/reference v0.6.0 // indirect + github.com/dunglas/httpsfv v1.1.0 // indirect github.com/fatih/color v1.19.0 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/stdr v1.2.2 // indirect @@ -84,6 +87,7 @@ require ( github.com/olekukonko/ll v0.1.8 // indirect github.com/petermattis/goid v0.0.0-20260725062400-500c67a39b75 // indirect github.com/puzpuzpuz/xsync/v4 v4.5.0 // indirect + github.com/quic-go/qpack v0.6.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 // indirect go.opentelemetry.io/otel v1.45.0 // indirect diff --git a/go.sum b/go.sum index 83aad9e43..e9db5cc97 100644 --- a/go.sum +++ b/go.sum @@ -54,6 +54,8 @@ github.com/docker/go-connections v0.7.0 h1:6SsRfJddP22WMrCkj19x9WKjEDTB+ahsdiGYf github.com/docker/go-connections v0.7.0/go.mod h1:no1qkHdjq7kLMGUXYAduOhYPSJxxvgWBh7ogVvptn3Q= github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= +github.com/dunglas/httpsfv v1.1.0 h1:Jw76nAyKWKZKFrpMMcL76y35tOpYHqQPzHQiwDvpe54= +github.com/dunglas/httpsfv v1.1.0/go.mod h1:zID2mqw9mFsnt7YC3vYQ9/cjq30q41W+1AnDwH8TiMg= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/elliotchance/orderedmap/v3 v3.1.1 h1:eV7lfZ5fVL8d36b8Wogqi/eqm7R/kZcftA9Yiyj+63M= @@ -166,12 +168,8 @@ github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731 h1:9x+U2HGLrSw5AT github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731/go.mod h1:Rs3MhFwutWhGwmY1VQsygw28z5bWcnEYmS1OG9OxjOQ= github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0 h1:XHNNzebIKZRkLimla/hFGrAIX5EMWHctrgt3hLw7s+I= github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0/go.mod h1:o8CFmAdrVwzJNOCsQCLUzXRjokkufNshnQHOe4fRaqU= -github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e h1:AjYyvyLEC+rB2HKit2RbwMnqUwujH2b7OSUaRUcYbVA= -github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34= -github.com/livekit/protocol v1.50.5-0.20260822191934-f1de69bed00e h1:SvFZZf/XyBPt4KsoxUa+lq/Y+F7XDcMji4GEiIzfbxg= -github.com/livekit/protocol v1.50.5-0.20260822191934-f1de69bed00e/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34= -github.com/livekit/protocol v1.50.5-0.20260822193149-0c325211bedd h1:ajR2EF9jKlVLW9AqHz46tMdLaFQsF60ZHqHImM4falo= -github.com/livekit/protocol v1.50.5-0.20260822193149-0c325211bedd/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34= +github.com/livekit/protocol v1.50.5-0.20260831015717-ba22cc1ce840 h1:VGdD4zhgRk5o/J+CPX/qT9lU/7l78eVwW8SPvrKBP8I= +github.com/livekit/protocol v1.50.5-0.20260831015717-ba22cc1ce840/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34= github.com/livekit/psrpc v0.7.3 h1:bekuZt/ZQzg8+/M8G6G5jq7bvV9fAKdPHSOZeTwrIIc= github.com/livekit/psrpc v0.7.3/go.mod h1:rAI+m2+/cb4x9RXhLRtUx5ZwdfjjXOl4zi46IjEetaw= github.com/livekit/webrtc-pion/v4 v4.2.18-warp.1 h1:fH+v4W+NFp9FfPzON6FaUFNmazGcctaAhb2P+Ksf+1s= @@ -286,6 +284,14 @@ github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+ github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY= github.com/puzpuzpuz/xsync/v4 v4.5.0 h1:vOSWu6b57/emh+L/Cw0BeQfvxa/cogFywXHeGUxQxAg= github.com/puzpuzpuz/xsync/v4 v4.5.0/go.mod h1:VJDmTCJMBt8igNxnkQd86r+8KUeN1quSfNKu5bLYFQo= +github.com/quic-go/go-ossfuzz-seeds v0.1.0 h1:APacT+iIaNF6fd8AGEiN3bT/Jtkd2jz4v4TzM7MFjy0= +github.com/quic-go/go-ossfuzz-seeds v0.1.0/go.mod h1:3IOHRbJIc+L6YKMwfDtJAM9Vj9k0YY4muhuyUYk5tbk= +github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= +github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= +github.com/quic-go/quic-go v0.61.0 h1:ui88A53s8MSVYLC56en0KQ17HARk+9986Dn0SBfKNvA= +github.com/quic-go/quic-go v0.61.0/go.mod h1:9So2anK4Tp22URSQq00k+Vo2PNkle96ycDPDHL4s9vs= +github.com/quic-go/webtransport-go v0.12.0 h1:CpnKNwZvdV0LD73xoHO8QaR0NI3llqpWRwnazdZS0sE= +github.com/quic-go/webtransport-go v0.12.0/go.mod h1:GHne8aRFJ24h73pAMrcywXtuaz/ShBXCLXLvG/NPFdU= github.com/redis/go-redis/v9 v9.22.0 h1:laDvpYXTJtZLloinw1fA5Kqd6HAEH2XKxOkG/PDq2F0= github.com/redis/go-redis/v9 v9.22.0/go.mod h1:y2g0Wj8rQvuK0ELM+oxSudcLtC09JScs98I/X9gRWY4= github.com/rodaine/protogofakeit v0.1.1 h1:ZKouljuRM3A+TArppfBqnH8tGZHOwM/pjvtXe9DaXH8= @@ -345,6 +351,8 @@ go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.uber.org/mock v0.5.2 h1:LbtPTcP8A5k9WPXj54PPPbjcI4Y6lhyOZXn+VS7wNko= +go.uber.org/mock v0.5.2/go.mod h1:wLlUxC2vVTPTaE3UD51E0BGOAElKrILxhVSDYQLld5o= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/zap v1.28.0 h1:IZzaP1Fv73/T/pBMLk4VutPl36uNC+OSUh3JLG3FIjo= diff --git a/pkg/agent/config.go b/pkg/agent/config.go index 8fd32b41e..d8e82ee9d 100644 --- a/pkg/agent/config.go +++ b/pkg/agent/config.go @@ -13,13 +13,20 @@ type Config struct { } type EndpointsConfig struct { - // Disabled turns off the /agents/{deployment}/... front and rejects - // registrations that declare endpoints. + // Disabled turns off the /agents/{agent_name}/{deployment}/... front and + // rejects registrations that declare endpoints. Disabled bool `yaml:"disabled,omitempty"` - DataConnCount uint32 `yaml:"data_conn_count,omitempty"` - CreditWindow uint32 `yaml:"credit_window,omitempty"` - ConnectionWindow uint32 `yaml:"connection_window,omitempty"` - MaxFrameSize uint32 `yaml:"max_frame_size,omitempty"` - MaxStreamsPerConn uint32 `yaml:"max_streams_per_conn,omitempty"` + // WebTransportPort is the UDP port for the worker WebTransport listener that + // carries the control stream and HTTP exchanges. 0 disables the listener, so + // no worker can serve endpoints (the front then only relays, in cloud). + WebTransportPort uint32 `yaml:"webtransport_port,omitempty"` + // TLSCertFile / TLSKeyFile is the certificate the WebTransport (QUIC) listener + // presents. QUIC has no plaintext mode, so a listener needs one; in dev mode a + // self-signed cert is generated when these are unset. + TLSCertFile string `yaml:"tls_cert_file,omitempty"` + TLSKeyFile string `yaml:"tls_key_file,omitempty"` + // MaxStreams is the soft per-session concurrent-stream cap used for capacity + // weighting; 0 takes the endpoint package default. + MaxStreams uint32 `yaml:"max_streams,omitempty"` } diff --git a/pkg/agent/endpoint/conformance/worker.go b/pkg/agent/endpoint/conformance/worker.go index e81cb72b9..944f99eb6 100644 --- a/pkg/agent/endpoint/conformance/worker.go +++ b/pkg/agent/endpoint/conformance/worker.go @@ -13,25 +13,27 @@ // limitations under the License. // Package conformance is a reference worker-side implementation of the agent -// HTTP endpoints data plane, used as test infrastructure: it registers a -// manifest over the control connection, attaches the fixed wire pool, and -// bridges each stream to a local HTTP server. The acceptance suite drives it to +// HTTP endpoints data plane over WebTransport, used as test infrastructure: it +// opens one WebTransport session to /agent, registers its manifest on the +// control stream, then serves each node-opened stream by bridging the opaque +// HTTP exchange to a local HTTP server. The acceptance suite drives it to // exercise the server end to end in Go, without a Python SDK worker. package conformance import ( "context" - "errors" + "crypto/tls" "fmt" "io" "net" - "net/url" + "net/http" "sync" "time" - "github.com/gorilla/websocket" - "google.golang.org/protobuf/proto" + "github.com/quic-go/quic-go/http3" + "github.com/quic-go/webtransport-go" + "github.com/livekit/livekit-server/pkg/agent/endpoint" "github.com/livekit/protocol/auth" "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" @@ -39,7 +41,7 @@ import ( ) type Config struct { - // ServerURL is the ws(s):// URL of the /agent endpoint + // ServerURL is the https:// URL of the /agent endpoint (WebTransport). ServerURL string APIKey string APISecret string @@ -48,26 +50,26 @@ type Config struct { Deployment string Endpoints []*livekit.AgentHttp_AgentEndpoint - // TargetAddr is the host:port of the local HTTP server streams bridge into + // TargetAddr is the host:port of the local HTTP server streams bridge into. TargetAddr string + // Insecure skips TLS verification (self-signed certs in tests). + Insecure bool + Logger logger.Logger } -// Worker is one registration epoch: a control connection plus the fixed -// data-connection pool. +// Worker is one registration epoch: a single WebTransport session carrying the +// control stream and every HTTP-exchange stream. type Worker struct { cfg Config instanceID string mu sync.Mutex + sess *webtransport.Session workerID string - settings *livekit.AgentHttp_AgentEndpointSettings - control *wsConn - dataConns []*dataConn closed bool registered chan struct{} - statusSeq uint64 } func New(cfg Config) *Worker { @@ -81,68 +83,171 @@ func New(cfg Config) *Worker { } } -// Start registers and attaches the pool; it returns once the data plane is up. +// Start dials the WebTransport session, registers on the control stream, and +// begins serving node-opened streams. It returns once registration completes. func (w *Worker) Start(ctx context.Context) error { token, err := w.mintToken() if err != nil { return err } - control, err := dialWS(ctx, w.cfg.ServerURL, token, false) + tr := &webtransport.Transport{ + TLSClientConfig: &tls.Config{ + InsecureSkipVerify: w.cfg.Insecure, //nolint:gosec // test/self-signed only + NextProtos: []string{http3.NextProtoH3}, + }, + } + hdr := http.Header{} + hdr.Set("Authorization", "Bearer "+token) + _, sess, err := tr.Dial(ctx, w.cfg.ServerURL, hdr) if err != nil { - return err + return fmt.Errorf("dial webtransport: %w", err) } w.mu.Lock() - w.control = control + w.sess = sess w.mu.Unlock() - if err := control.writeWorker(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Register{ - Register: &livekit.RegisterWorkerRequest{ - Type: livekit.JobType_JT_ROOM, - AgentName: w.cfg.AgentName, - Version: "endpoint-conformance-client", - PingInterval: 30, - Deployment: w.cfg.Deployment, - Endpoints: w.cfg.Endpoints, - InstanceId: w.instanceID, - EndpointProtocol: 1, + // control stream: the worker opens it first, then registers + control, err := sess.OpenStreamSync(ctx) + if err != nil { + return fmt.Errorf("open control stream: %w", err) + } + if err := endpoint.WriteControlMessage(control, &livekit.WorkerMessage{ + Message: &livekit.WorkerMessage_Register{ + Register: &livekit.RegisterWorkerRequest{ + Type: livekit.JobType_JT_ROOM, + AgentName: w.cfg.AgentName, + Version: "endpoint-conformance-client", + PingInterval: 30, + Deployment: w.cfg.Deployment, + Endpoints: w.cfg.Endpoints, + InstanceId: w.instanceID, + EndpointProtocol: 1, + }, }, - }}); err != nil { - return err + }); err != nil { + return fmt.Errorf("write register: %w", err) } - go w.controlLoop(ctx, control) + var resp livekit.ServerMessage + if err := endpoint.ReadControlMessage(control, &resp); err != nil { + return fmt.Errorf("read register response: %w", err) + } + reg := resp.GetRegister() + if reg == nil { + return fmt.Errorf("expected register response, got %T", resp.GetMessage()) + } + w.mu.Lock() + w.workerID = reg.GetWorkerId() + close(w.registered) + w.mu.Unlock() + go w.controlLoop(control) + // the serve loop lives as long as the session, not the dial context (which + // the caller may cancel as soon as Start returns) + go w.serveLoop(sess) + return nil +} + +// controlLoop drains further control messages (availability requests are +// declined; the conformance worker takes no jobs). +func (w *Worker) controlLoop(control *webtransport.Stream) { + for { + var msg livekit.ServerMessage + if err := endpoint.ReadControlMessage(control, &msg); err != nil { + return + } + if a := msg.GetAvailability(); a != nil { + _ = endpoint.WriteControlMessage(control, &livekit.WorkerMessage{ + Message: &livekit.WorkerMessage_Availability{ + Availability: &livekit.AvailabilityResponse{ + JobId: a.GetJob().GetId(), + Available: false, + }, + }, + }) + } + } +} + +// serveLoop accepts node-opened streams (one HTTP exchange each) and bridges +// them to the local target. It runs until the session ends. +func (w *Worker) serveLoop(sess *webtransport.Session) { + ctx := sess.Context() + for { + stream, err := sess.AcceptStream(ctx) + if err != nil { + return + } + go w.serve(stream) + } +} + +// serve bridges one HTTP exchange: the opaque request bytes on the stream are +// piped to a fresh TCP connection to the target, and the target's response +// bytes are piped back. Half-closes are propagated in both directions so +// streaming responses (SSE) flush incrementally. +func (w *Worker) serve(stream *webtransport.Stream) { + tcp, err := net.Dial("tcp", w.cfg.TargetAddr) + if err != nil { + // the request was never dispatched: refuse so the front may retry + stream.CancelWrite(endpoint.StreamCodeRefused) + stream.CancelRead(endpoint.StreamCodeRefused) + return + } + + var wg sync.WaitGroup + wg.Add(2) + // request: stream -> target, then half-close the target's write side + go func() { + defer wg.Done() + _, _ = io.Copy(tcp, stream) + if c, ok := tcp.(*net.TCPConn); ok { + _ = c.CloseWrite() + } + }() + // response: target -> stream, then FIN the stream's response side + go func() { + defer wg.Done() + _, _ = io.Copy(stream, tcp) + _ = stream.Close() + }() + wg.Wait() + _ = tcp.Close() +} + +// WaitRegistered blocks until registration completes or ctx is done. +func (w *Worker) WaitRegistered(ctx context.Context) error { select { case <-w.registered: + return nil case <-ctx.Done(): return ctx.Err() case <-time.After(10 * time.Second): - return errors.New("registration timeout") + return fmt.Errorf("timed out waiting for registration") } +} +// WorkerID returns the id assigned by the server (valid after registration). +func (w *Worker) WorkerID() string { w.mu.Lock() - settings := w.settings - workerID := w.workerID - w.mu.Unlock() - if settings == nil { - return errors.New("server did not negotiate endpoint settings") - } + defer w.mu.Unlock() + return w.workerID +} - for i := uint32(0); i < settings.GetDataConnectionCount(); i++ { - dc, err := w.attachDataConn(ctx, token, workerID, settings) - if err != nil { - return fmt.Errorf("attach %d: %w", i, err) - } - w.mu.Lock() - w.dataConns = append(w.dataConns, dc) +// Close tears down the session. +func (w *Worker) Close() { + w.mu.Lock() + if w.closed { w.mu.Unlock() + return + } + w.closed = true + sess := w.sess + w.mu.Unlock() + if sess != nil { + _ = sess.CloseWithError(0, "worker closed") } - - // report availability so the front's load weighting sees the worker - _ = w.UpdateStatus(livekit.WorkerStatus_WS_AVAILABLE, 0, false) - - return nil } func (w *Worker) mintToken() (string, error) { @@ -151,649 +256,3 @@ func (w *Worker) mintToken() (string, error) { SetValidFor(24 * time.Hour) return at.ToJWT() } - -func (w *Worker) WorkerID() string { - w.mu.Lock() - defer w.mu.Unlock() - return w.workerID -} - -func (w *Worker) UpdateStatus(status livekit.WorkerStatus, load float32, draining bool) error { - w.mu.Lock() - w.statusSeq++ - seq := w.statusSeq - control := w.control - w.mu.Unlock() - if control == nil { - return errors.New("not started") - } - return control.writeWorker(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateWorker{ - UpdateWorker: &livekit.UpdateWorkerStatus{ - Status: &status, - Load: load, - Draining: draining, - Seq: seq, - }, - }}) -} - -func (w *Worker) Close() { - w.mu.Lock() - if w.closed { - w.mu.Unlock() - return - } - w.closed = true - control := w.control - conns := w.dataConns - w.mu.Unlock() - - if control != nil { - _ = control.close() - } - for _, dc := range conns { - dc.close(errors.New("worker closed")) - } -} - -func (w *Worker) controlLoop(ctx context.Context, control *wsConn) { - for { - msg, err := control.readServer() - if err != nil { - return - } - switch m := msg.Message.(type) { - case *livekit.ServerMessage_Register: - w.mu.Lock() - w.workerID = m.Register.GetWorkerId() - w.settings = m.Register.GetEndpointSettings() - w.mu.Unlock() - close(w.registered) - case *livekit.ServerMessage_Availability: - // decline room jobs: the conformance client only serves endpoints - _ = control.writeWorker(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Availability{ - Availability: &livekit.AvailabilityResponse{ - JobId: m.Availability.GetJob().GetId(), - Available: false, - }, - }}) - case *livekit.ServerMessage_GoAway: - w.cfg.Logger.Infow("server draining, closing", "reason", m.GoAway.GetReason()) - w.Close() - return - case *livekit.ServerMessage_Pong: - default: - } - } -} - -// --- data connections --- - -type dataConn struct { - ws *wsConn - params *livekit.AgentHttp_AttachDataConnectionResponse - target string - log logger.Logger - - writeMu sync.Mutex // client side serializes writes; correctness over throughput - - mu sync.Mutex - streams map[uint32]*clientStream - closed bool - - // connection-level send window (response bytes toward the server), refilled - // by server credit frames on stream 0 - sendMu sync.Mutex - sendCond *sync.Cond - connSendCredit int64 - - // connection-level receive accounting (request bytes from the server) - recvMu sync.Mutex - connRecvUnacked int64 -} - -func (w *Worker) attachDataConn(ctx context.Context, token, workerID string, settings *livekit.AgentHttp_AgentEndpointSettings) (*dataConn, error) { - ws, err := dialWS(ctx, w.cfg.ServerURL, token, true) - if err != nil { - return nil, err - } - if err := ws.writeFrame(&livekit.AgentHttp_Frame{ - Message: &livekit.AgentHttp_Frame_Attach{Attach: &livekit.AgentHttp_AttachDataConnection{ - WorkerId: workerID, - InstanceId: w.instanceID, - AttachToken: settings.GetAttachToken(), - }}, - }); err != nil { - _ = ws.close() - return nil, err - } - f, err := ws.readFrame() - if err != nil { - _ = ws.close() - return nil, err - } - resp, ok := f.Message.(*livekit.AgentHttp_Frame_AttachResponse) - if !ok { - _ = ws.close() - return nil, errors.New("expected attach response") - } - if e := resp.AttachResponse.GetError(); e != "" { - _ = ws.close() - return nil, fmt.Errorf("attach rejected: %s", e) - } - - dc := &dataConn{ - ws: ws, - params: resp.AttachResponse, - target: w.cfg.TargetAddr, - log: w.cfg.Logger, - streams: make(map[uint32]*clientStream), - connSendCredit: int64(resp.AttachResponse.GetConnectionWindow()), - } - dc.sendCond = sync.NewCond(&dc.sendMu) - ws.enableWireLiveness() - go dc.readLoop() - go dc.pingLoop() - return dc, nil -} - -// pingLoop keeps the wire's liveness visible in both directions: the server -// refreshes its read deadline on our pings, and our pong-refreshed deadline -// detects a dead server. -func (dc *dataConn) pingLoop() { - t := time.NewTicker(30 * time.Second) - defer t.Stop() - for range t.C { - dc.mu.Lock() - closed := dc.closed - dc.mu.Unlock() - if closed { - return - } - if err := dc.ws.ping(); err != nil { - return - } - } -} - -// writeFrame serializes all frames FIFO: unlike the server's scheduler there is -// no control/data priority here, so a credit can queue behind sibling data. Fine -// for a conformance client; bounded by the per-write deadline. -func (dc *dataConn) writeFrame(f *livekit.AgentHttp_Frame) error { - dc.writeMu.Lock() - defer dc.writeMu.Unlock() - return dc.ws.writeFrame(f) -} - -// reserveConnSend takes up to want bytes from the shared window. -func (dc *dataConn) reserveConnSend(want int64, failed func() bool) (int64, error) { - dc.sendMu.Lock() - defer dc.sendMu.Unlock() - for dc.connSendCredit <= 0 { - dc.mu.Lock() - closed := dc.closed - dc.mu.Unlock() - if closed || failed() { - return 0, errors.New("connection or stream closed") - } - dc.sendCond.Wait() - } - n := want - if n > dc.connSendCredit { - n = dc.connSendCredit - } - dc.connSendCredit -= n - return n, nil -} - -// connConsumed replenishes the shared receive window once request bytes were -// consumed, threshold-acked at half the window on stream 0. -func (dc *dataConn) connConsumed(n int64) { - if n <= 0 { - return - } - dc.recvMu.Lock() - dc.connRecvUnacked += n - var credit int64 - if dc.connRecvUnacked >= int64(dc.params.GetConnectionWindow())/2 { - credit = dc.connRecvUnacked - dc.connRecvUnacked = 0 - } - dc.recvMu.Unlock() - if credit > 0 { - _ = dc.writeFrame(&livekit.AgentHttp_Frame{ - StreamId: 0, - Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(credit)}, - }) - } -} - -func (dc *dataConn) close(err error) { - dc.mu.Lock() - if dc.closed { - dc.mu.Unlock() - return - } - dc.closed = true - streams := make([]*clientStream, 0, len(dc.streams)) - for _, s := range dc.streams { - streams = append(streams, s) - } - dc.streams = map[uint32]*clientStream{} - dc.mu.Unlock() - - for _, s := range streams { - s.fail(err) - } - dc.sendMu.Lock() - dc.sendCond.Broadcast() - dc.sendMu.Unlock() - _ = dc.ws.close() -} - -func (dc *dataConn) readLoop() { - for { - f, err := dc.ws.readFrame() - if err != nil { - dc.close(err) - return - } - switch m := f.Message.(type) { - case *livekit.AgentHttp_Frame_Open: - dc.handleOpen(f.StreamId, m.Open) - case *livekit.AgentHttp_Frame_Data: - delivered := false - dc.withStream(f.StreamId, func(s *clientStream) { - s.onData(m.Data) - delivered = true - }) - if !delivered { - // unknown stream: protocol contract says ignore, but the bytes - // consumed the shared window and must be credited back - dc.connConsumed(int64(len(m.Data))) - } - case *livekit.AgentHttp_Frame_Eof: - dc.withStream(f.StreamId, func(s *clientStream) { - s.onEOF() - }) - case *livekit.AgentHttp_Frame_Reset_: - dc.withStream(f.StreamId, func(s *clientStream) { - s.fail(fmt.Errorf("reset: %s", m.Reset_.GetError())) - }) - case *livekit.AgentHttp_Frame_Credit: - if f.StreamId == 0 { - dc.sendMu.Lock() - dc.connSendCredit += int64(m.Credit) - dc.sendCond.Broadcast() - dc.sendMu.Unlock() - } else { - dc.withStream(f.StreamId, func(s *clientStream) { - s.onCredit(m.Credit) - }) - } - default: - // attach frames after the handshake are a protocol violation - dc.close(errors.New("unexpected frame on data connection")) - return - } - } -} - -func (dc *dataConn) withStream(id uint32, f func(*clientStream)) { - dc.mu.Lock() - s := dc.streams[id] - dc.mu.Unlock() - if s != nil { - f(s) - } -} - -func (dc *dataConn) handleOpen(id uint32, open *livekit.AgentHttp_HttpStreamOpen) { - s := &clientStream{ - id: id, - dc: dc, - sendCredit: int64(dc.params.GetCreditWindow()), - failCh: make(chan struct{}), - } - s.sendCond = sync.NewCond(&s.sendMu) - s.recvCond = sync.NewCond(&s.recvMu) - _ = open - - dc.mu.Lock() - if dc.closed || uint32(len(dc.streams)) >= dc.params.GetMaxStreamsPerConn() { - dc.mu.Unlock() - _ = dc.writeFrame(resetFrame(id, livekit.AgentHttp_HSR_REFUSED, "no stream capacity")) - return - } - dc.streams[id] = s - dc.mu.Unlock() - - go s.run() -} - -func resetFrame(id uint32, code livekit.AgentHttp_HttpStreamResetCode, reason string) *livekit.AgentHttp_Frame { - return &livekit.AgentHttp_Frame{ - StreamId: id, - Message: &livekit.AgentHttp_Frame_Reset_{ - Reset_: &livekit.AgentHttp_HttpStreamReset{Code: code, Error: reason}, - }, - } -} - -// clientStream bridges one stream to one TCP connection to the local HTTP -// server. The transport parses nothing. -type clientStream struct { - id uint32 - dc *dataConn - - sendMu sync.Mutex - sendCond *sync.Cond - sendCredit int64 - sendFailed bool - - // recv side: a cond-guarded buffer, never a blocking channel - the wire - // read loop must not stall behind one slow local app (bytes in flight are - // already bounded by the credit window) - recvMu sync.Mutex - recvCond *sync.Cond - recvBuf [][]byte - recvEOF bool - recvDone bool // removed or failed: arriving bytes are settled immediately - recvUnacked int64 - - failCh chan struct{} - failOnce sync.Once - err error -} - -func (s *clientStream) fail(err error) { - s.failOnce.Do(func() { - s.recvMu.Lock() - s.err = err - s.recvDone = true - // bytes the app will never read are settled against the shared window - for _, p := range s.recvBuf { - s.dc.connConsumed(int64(len(p))) - } - s.recvBuf = nil - s.recvCond.Broadcast() - s.recvMu.Unlock() - s.sendMu.Lock() - s.sendFailed = true - s.sendCond.Broadcast() - s.sendMu.Unlock() - close(s.failCh) - }) -} - -func (s *clientStream) onData(payload []byte) { - s.recvMu.Lock() - if s.recvDone { - s.recvMu.Unlock() - // never consumed by the app: release the shared window immediately - s.dc.connConsumed(int64(len(payload))) - return - } - s.recvBuf = append(s.recvBuf, payload) - s.recvCond.Broadcast() - s.recvMu.Unlock() -} - -func (s *clientStream) onEOF() { - s.recvMu.Lock() - s.recvEOF = true - s.recvCond.Broadcast() - s.recvMu.Unlock() -} - -func (s *clientStream) onCredit(inc uint32) { - s.sendMu.Lock() - s.sendCredit += int64(inc) - s.sendCond.Broadcast() - s.sendMu.Unlock() -} - -func (s *clientStream) remove() { - // mark done BEFORE unlinking so a delivery racing the removal settles the - // shared window instead of vanishing into an orphaned buffer - s.recvMu.Lock() - s.recvDone = true - for _, p := range s.recvBuf { - s.dc.connConsumed(int64(len(p))) - } - s.recvBuf = nil - s.recvCond.Broadcast() - s.recvMu.Unlock() - - s.dc.mu.Lock() - delete(s.dc.streams, s.id) - s.dc.mu.Unlock() -} - -// nextRecv blocks for the next request chunk; ok=false means EOF (done=false) -// or stream failure (done=true). -func (s *clientStream) nextRecv() (payload []byte, ok bool, failed bool) { - s.recvMu.Lock() - defer s.recvMu.Unlock() - for len(s.recvBuf) == 0 && !s.recvEOF && !s.recvDone { - s.recvCond.Wait() - } - if len(s.recvBuf) > 0 { - payload = s.recvBuf[0] - s.recvBuf = s.recvBuf[1:] - return payload, true, false - } - if s.recvDone { - return nil, false, true - } - return nil, false, false -} - -// run executes the pump: dial the local app, copy stream->app and app->stream -// with two-level credit accounting on both directions. This is deliberately the -// report's five-line tunnel plus flow control. -func (s *clientStream) run() { - defer s.remove() - - app, err := net.DialTimeout("tcp", s.dc.target, 10*time.Second) - if err != nil { - // nothing was dispatched: REFUSED tells the server a retry is safe - _ = s.dc.writeFrame(resetFrame(s.id, livekit.AgentHttp_HSR_REFUSED, "local app unreachable")) - return - } - defer app.Close() - - done := make(chan struct{}, 2) - - // stream -> app (request bytes) - go func() { - defer func() { done <- struct{}{} }() - for { - payload, ok, failed := s.nextRecv() - if failed { - return - } - if !ok { - // EOF after all buffered chunks drained - if tc, ok := app.(*net.TCPConn); ok { - _ = tc.CloseWrite() - } - return - } - // enforce, consume, and replenish both windows as bytes reach the app - if _, err := app.Write(payload); err != nil { - s.fail(err) - return - } - s.dc.connConsumed(int64(len(payload))) - s.recvUnacked += int64(len(payload)) - if s.recvUnacked >= int64(s.dc.params.GetCreditWindow())/2 { - _ = s.dc.writeFrame(&livekit.AgentHttp_Frame{ - StreamId: s.id, - Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(s.recvUnacked)}, - }) - s.recvUnacked = 0 - } - } - }() - - // app -> stream (response bytes), under the stream window and the wire's - // shared connection window - go func() { - defer func() { done <- struct{}{} }() - failed := func() bool { - select { - case <-s.failCh: - return true - default: - return false - } - } - buf := make([]byte, int(s.dc.params.GetMaxFrameSize())) - for { - n, rerr := app.Read(buf) - if n > 0 { - remaining := buf[:n] - for len(remaining) > 0 { - s.sendMu.Lock() - for s.sendCredit <= 0 && !s.sendFailed { - s.sendCond.Wait() - } - if s.sendFailed { - s.sendMu.Unlock() - return - } - want := int64(len(remaining)) - if want > s.sendCredit { - want = s.sendCredit - } - s.sendMu.Unlock() - - c, err := s.dc.reserveConnSend(want, failed) - if err != nil { - return - } - s.sendMu.Lock() - s.sendCredit -= c - s.sendMu.Unlock() - - chunk := make([]byte, c) - copy(chunk, remaining[:c]) - if err := s.dc.writeFrame(&livekit.AgentHttp_Frame{ - StreamId: s.id, - Message: &livekit.AgentHttp_Frame_Data{Data: chunk}, - }); err != nil { - s.fail(err) - return - } - remaining = remaining[c:] - } - } - if rerr != nil { - if rerr != io.EOF { - s.fail(rerr) - } - _ = s.dc.writeFrame(&livekit.AgentHttp_Frame{ - StreamId: s.id, - Message: &livekit.AgentHttp_Frame_Eof{Eof: &livekit.AgentHttp_HttpStreamEof{}}, - }) - return - } - } - }() - - <-done - <-done -} - -// --- websocket transport --- - -type wsConn struct { - ws *websocket.Conn - mu sync.Mutex -} - -func dialWS(ctx context.Context, serverURL, token string, attach bool) (*wsConn, error) { - u, err := url.Parse(serverURL) - if err != nil { - return nil, err - } - q := u.Query() - q.Set("access_token", token) - if attach { - q.Set("attach", "1") - } - u.RawQuery = q.Encode() - - ws, _, err := websocket.DefaultDialer.DialContext(ctx, u.String(), nil) - if err != nil { - return nil, err - } - return &wsConn{ws: ws}, nil -} - -func (c *wsConn) writeWorker(msg *livekit.WorkerMessage) error { - payload, err := proto.Marshal(msg) - if err != nil { - return err - } - c.mu.Lock() - defer c.mu.Unlock() - _ = c.ws.SetWriteDeadline(time.Now().Add(30 * time.Second)) - return c.ws.WriteMessage(websocket.BinaryMessage, payload) -} - -func (c *wsConn) readServer() (*livekit.ServerMessage, error) { - _, payload, err := c.ws.ReadMessage() - if err != nil { - return nil, err - } - msg := &livekit.ServerMessage{} - if err := proto.Unmarshal(payload, msg); err != nil { - return nil, err - } - return msg, nil -} - -func (c *wsConn) writeFrame(f *livekit.AgentHttp_Frame) error { - payload, err := proto.Marshal(f) - if err != nil { - return err - } - c.mu.Lock() - defer c.mu.Unlock() - _ = c.ws.SetWriteDeadline(time.Now().Add(30 * time.Second)) - return c.ws.WriteMessage(websocket.BinaryMessage, payload) -} - -func (c *wsConn) readFrame() (*livekit.AgentHttp_Frame, error) { - _ = c.ws.SetReadDeadline(time.Now().Add(wireIdleTimeout)) - _, payload, err := c.ws.ReadMessage() - if err != nil { - return nil, err - } - f := &livekit.AgentHttp_Frame{} - if err := proto.Unmarshal(payload, f); err != nil { - return nil, err - } - return f, nil -} - -const wireIdleTimeout = 2 * time.Minute - -// enableWireLiveness arms a read deadline refreshed by any inbound traffic and -// by pongs to our pings. -func (c *wsConn) enableWireLiveness() { - _ = c.ws.SetReadDeadline(time.Now().Add(wireIdleTimeout)) - c.ws.SetPongHandler(func(string) error { - return c.ws.SetReadDeadline(time.Now().Add(wireIdleTimeout)) - }) -} - -func (c *wsConn) ping() error { - return c.ws.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)) -} - -func (c *wsConn) close() error { - return c.ws.Close() -} diff --git a/pkg/agent/endpoint/conn.go b/pkg/agent/endpoint/conn.go deleted file mode 100644 index 26d8936d0..000000000 --- a/pkg/agent/endpoint/conn.go +++ /dev/null @@ -1,386 +0,0 @@ -// Copyright 2026 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 endpoint - -import ( - "sync" - "sync/atomic" - "time" - - "github.com/livekit/protocol/livekit" - "github.com/livekit/protocol/logger" -) - -// DataConn is one adopted worker wire carrying multiplexed streams. Only the -// server opens streams; ids are odd and scoped to the wire (even ids are -// reserved for worker-opened streams). Flow control is two-level: per-stream -// windows plus a shared connection window (stream 0 credit), so one stalled -// consumer can neither starve its siblings nor pin unbounded memory. -type DataConn struct { - wire WireConn - params WireParams - logger logger.Logger - - sched *scheduler - - mu sync.Mutex - streams map[uint32]*Stream - nextID uint32 - closed bool - closeErr error - onClose func(*DataConn) - - // connection-level send window (request bytes toward the worker), refilled - // by worker credit frames on stream 0 - sendMu sync.Mutex - sendCond *sync.Cond - connSendCredit int64 - - // connection-level receive accounting (response bytes from the worker): - // recvConnAvail is what the worker may still put in flight, replenished on - // consumption via connConsumed - recvMu sync.Mutex - recvConnAvail int64 - recvConnUnacked int64 - - // activity is an exponentially decayed byte counter used by stream - // placement: a connection busy moving a heavy transfer should not receive - // new streams while a lighter sibling exists. - activity atomic.Int64 - lastDecay atomic.Int64 // unix nanos - decayHalfLife time.Duration -} - -func NewDataConn(wire WireConn, params WireParams, onClose func(*DataConn), log logger.Logger) *DataConn { - params = params.WithDefaults() - c := &DataConn{ - wire: wire, - params: params, - logger: log, - streams: make(map[uint32]*Stream), - nextID: 1, // odd ids; even reserved for worker-opened streams - onClose: onClose, - connSendCredit: int64(params.ConnectionWindow), - recvConnAvail: int64(params.ConnectionWindow), - decayHalfLife: 500 * time.Millisecond, - } - c.sendCond = sync.NewCond(&c.sendMu) - c.lastDecay.Store(time.Now().UnixNano()) - c.sched = newScheduler(wire, int(params.MaxFrameSize)) - go c.readLoop() - return c -} - -// OpenStream opens a stream and sends its OPEN frame. It fails fast when the -// per-wire stream cap is reached; the caller picks another wire or worker. -func (c *DataConn) OpenStream(meta *livekit.AgentHttp_HttpStreamOpen) (*Stream, error) { - c.mu.Lock() - if c.closed { - err := c.closeErr - c.mu.Unlock() - if err == nil { - err = ErrConnClosed - } - return nil, err - } - if uint32(len(c.streams)) >= c.params.MaxStreamsPerConn { - c.mu.Unlock() - return nil, ErrTooManyStreams - } - id := c.nextID - c.nextID += 2 - s := newStream(id, c, int64(c.params.CreditWindow)) - c.streams[id] = s - c.mu.Unlock() - - if err := c.sched.enqueueControl(&livekit.AgentHttp_Frame{ - StreamId: id, - Message: &livekit.AgentHttp_Frame_Open{Open: meta}, - }); err != nil { - c.removeStream(id) - return nil, err - } - return s, nil -} - -// OpenStreams reports open streams (active and parked). -func (c *DataConn) OpenStreams() int { - c.mu.Lock() - defer c.mu.Unlock() - return len(c.streams) -} - -// SpareStreams reports how many more streams this conn can accept before its -// per-conn cap; never negative. -func (c *DataConn) SpareStreams() int { - c.mu.Lock() - defer c.mu.Unlock() - if n := int(c.params.MaxStreamsPerConn) - len(c.streams); n > 0 { - return n - } - return 0 -} - -// HasCapacity reports whether a new stream may be opened. -func (c *DataConn) HasCapacity() bool { - c.mu.Lock() - defer c.mu.Unlock() - return !c.closed && uint32(len(c.streams)) < c.params.MaxStreamsPerConn -} - -// --- connection-level send window --- - -// reserveConnSend blocks until part of the shared window is available and -// reserves up to want bytes. cancelled lets a caller whose stream died stop -// waiting for credit that may never come. -func (c *DataConn) reserveConnSend(want int64, cancelled func() bool) (int64, error) { - c.sendMu.Lock() - defer c.sendMu.Unlock() - for c.connSendCredit <= 0 { - if c.closedErr() != nil { - return 0, c.closedErr() - } - if cancelled() { - return 0, ErrStreamClosed - } - c.sendCond.Wait() - } - if err := c.closedErr(); err != nil { - return 0, err - } - n := want - if n > c.connSendCredit { - n = c.connSendCredit - } - c.connSendCredit -= n - return n, nil -} - -func (c *DataConn) returnConnSend(n int64) { - if n <= 0 { - return - } - c.sendMu.Lock() - c.connSendCredit += n - c.sendCond.Broadcast() - c.sendMu.Unlock() -} - -func (c *DataConn) addConnSend(increment uint32) { - c.sendMu.Lock() - c.connSendCredit += int64(increment) - c.sendCond.Broadcast() - c.sendMu.Unlock() -} - -// wakeSendWaiters unblocks reserveConnSend callers so a reset stream's writer -// can observe its cancellation. -func (c *DataConn) wakeSendWaiters() { - c.sendMu.Lock() - c.sendCond.Broadcast() - c.sendMu.Unlock() -} - -func (c *DataConn) closedErr() error { - c.mu.Lock() - defer c.mu.Unlock() - if !c.closed { - return nil - } - if c.closeErr != nil { - return c.closeErr - } - return ErrConnClosed -} - -// --- connection-level receive window --- - -// consumeConnRecv accounts an arriving payload against the shared window. -func (c *DataConn) consumeConnRecv(n int64) error { - c.recvMu.Lock() - defer c.recvMu.Unlock() - if n > c.recvConnAvail { - return ErrProtocol - } - c.recvConnAvail -= n - return nil -} - -// connConsumed replenishes the shared window once bytes are consumed (read by -// the client, or drained by a torn-down stream), threshold-acked at half the -// window on stream 0. -func (c *DataConn) connConsumed(n int64) { - if n <= 0 { - return - } - c.recvMu.Lock() - c.recvConnUnacked += n - var credit int64 - if c.recvConnUnacked >= int64(c.params.ConnectionWindow)/2 { - credit = c.recvConnUnacked - c.recvConnUnacked = 0 - c.recvConnAvail += credit - } - c.recvMu.Unlock() - if credit > 0 { - _ = c.sched.enqueueControl(&livekit.AgentHttp_Frame{ - StreamId: 0, - Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(credit)}, - }) - } -} - -// noteActivity feeds the placement score: every payload byte in either -// direction counts, decayed over time so idle (parked) streams stop weighing. -func (c *DataConn) noteActivity(n int64) { - c.decay() - c.activity.Add(n) -} - -func (c *DataConn) decay() { - now := time.Now().UnixNano() - last := c.lastDecay.Load() - elapsed := time.Duration(now - last) - if elapsed < c.decayHalfLife { - return - } - if !c.lastDecay.CompareAndSwap(last, now) { - return - } - halvings := int(elapsed / c.decayHalfLife) - v := c.activity.Load() - for i := 0; i < halvings && v != 0; i++ { - v /= 2 - } - c.activity.Store(v) -} - -// Score is the placement weight: recent payload activity plus queued backlog. -// Lower is lighter. Heavy returns true when the connection is actively moving a -// large transfer and should not be co-located onto. -func (c *DataConn) Score() (score int64, heavy bool) { - c.decay() - a := c.activity.Load() - c.sched.mu.Lock() - q := c.sched.queuedBytes - c.sched.mu.Unlock() - score = a + q + int64(c.OpenStreams())*1024 - heavy = a+q >= int64(c.params.CreditWindow)/2 - return score, heavy -} - -func (c *DataConn) Close(err error) { - c.mu.Lock() - if c.closed { - c.mu.Unlock() - return - } - c.closed = true - c.closeErr = err - streams := make([]*Stream, 0, len(c.streams)) - for _, s := range c.streams { - streams = append(streams, s) - } - c.streams = map[uint32]*Stream{} - onClose := c.onClose - c.mu.Unlock() - - if err == nil { - err = ErrConnClosed - } - for _, s := range streams { - s.mu.Lock() - if s.err == nil { - s.err = err - } - s.closed = true - s.cond.Broadcast() - s.mu.Unlock() - } - c.wakeSendWaiters() - c.sched.close(err) - if onClose != nil { - onClose(c) - } -} - -func (c *DataConn) removeStream(id uint32) { - c.mu.Lock() - delete(c.streams, id) - c.mu.Unlock() -} - -func (c *DataConn) stream(id uint32) *Stream { - c.mu.Lock() - defer c.mu.Unlock() - return c.streams[id] -} - -func (c *DataConn) readLoop() { - for { - f, err := c.wire.ReadFrame() - if err != nil { - c.Close(err) - return - } - switch m := f.Message.(type) { - case *livekit.AgentHttp_Frame_Data: - if err := c.consumeConnRecv(int64(len(m.Data))); err != nil { - c.logger.Warnw("data plane protocol violation", err) - c.Close(err) - return - } - s := c.stream(f.StreamId) - if s == nil { - // reset locally; the payload will never be read - c.connConsumed(int64(len(m.Data))) - continue - } - c.noteActivity(int64(len(m.Data))) - if err := s.onData(m.Data); err != nil { - if err == errStreamGone { - c.connConsumed(int64(len(m.Data))) - continue - } - c.logger.Warnw("data plane protocol violation", err) - c.Close(err) - return - } - case *livekit.AgentHttp_Frame_Eof: - if s := c.stream(f.StreamId); s != nil { - s.onEOF() - } - case *livekit.AgentHttp_Frame_Reset_: - if s := c.stream(f.StreamId); s != nil { - s.onReset(m.Reset_.GetCode(), m.Reset_.GetError()) - c.removeStream(f.StreamId) - c.sched.dropStream(f.StreamId) - c.wakeSendWaiters() - } - case *livekit.AgentHttp_Frame_Credit: - if f.StreamId == 0 { - c.addConnSend(m.Credit) - } else if s := c.stream(f.StreamId); s != nil { - s.onCredit(m.Credit) - } - default: - // attach handshakes ended before adoption; open only flows toward - // the worker - c.logger.Warnw("unexpected frame on data connection", ErrProtocol, "frame", logger.Proto(f)) - c.Close(ErrProtocol) - return - } - } -} diff --git a/pkg/agent/endpoint/control.go b/pkg/agent/endpoint/control.go new file mode 100644 index 000000000..48a97920f --- /dev/null +++ b/pkg/agent/endpoint/control.go @@ -0,0 +1,64 @@ +// Copyright 2026 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 endpoint + +import ( + "encoding/binary" + "fmt" + "io" + + "google.golang.org/protobuf/proto" +) + +// MaxControlMessageSize bounds one framed control message so a malformed length +// prefix cannot make the reader allocate without bound. +const MaxControlMessageSize = 1 << 20 + +// WriteControlMessage writes a length-delimited protobuf on the control stream: +// a 4-byte big-endian length followed by the marshaled message. The control +// stream carries the same WorkerMessage/ServerMessage exchange the WebSocket +// control connection used; only the framing (QUIC stream instead of WS message) +// differs. +func WriteControlMessage(w io.Writer, m proto.Message) error { + b, err := proto.Marshal(m) + if err != nil { + return err + } + var hdr [4]byte + binary.BigEndian.PutUint32(hdr[:], uint32(len(b))) + if _, err := w.Write(hdr[:]); err != nil { + return err + } + _, err = w.Write(b) + return err +} + +// ReadControlMessage reads one length-delimited protobuf written by +// WriteControlMessage into m. +func ReadControlMessage(r io.Reader, m proto.Message) error { + var hdr [4]byte + if _, err := io.ReadFull(r, hdr[:]); err != nil { + return err + } + n := binary.BigEndian.Uint32(hdr[:]) + if n > MaxControlMessageSize { + return fmt.Errorf("control message too large: %d bytes", n) + } + b := make([]byte, n) + if _, err := io.ReadFull(r, b); err != nil { + return err + } + return proto.Unmarshal(b, m) +} diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go index ab0863f89..88c023188 100644 --- a/pkg/agent/endpoint/front.go +++ b/pkg/agent/endpoint/front.go @@ -17,7 +17,6 @@ package endpoint import ( "bufio" "context" - "errors" "fmt" "io" "math/rand/v2" @@ -28,9 +27,7 @@ import ( "sync" "time" - "github.com/livekit/protocol/livekit" "github.com/livekit/protocol/logger" - "github.com/livekit/protocol/utils/guid" ) // HeaderEndpointMiss marks a response the front produced ITSELF - a routing @@ -52,7 +49,7 @@ const ( ) const ( - // PathPrefix is the public route namespace: /agents/{deployment}/{path...} + // PathPrefix is the public route namespace: /agents/{agent_name}/{deployment}/{path...} PathPrefix = "/agents/" // responseHeadTimeout bounds the wait for the worker's response head. Bodies @@ -109,7 +106,12 @@ type FallbackRequest struct { // unauthenticated) APIKey string Authenticated bool - Deployment string + AgentName string + Deployment string + // Method is the request's HTTP method, so a multi-node layer can select + // candidates method-aware (a node serving the path under a different method + // is not a candidate). + Method string // Path within the deployment, '/'-rooted Path string } @@ -160,6 +162,11 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { http.NotFound(w, r) return } + agentName, rest, found := strings.Cut(rest, "/") + if !found || agentName == "" { + http.NotFound(w, r) + return + } deployment, path, found := strings.Cut(rest, "/") if !found { path = "" @@ -182,7 +189,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - candidates := f.registry.Candidates(apiKey, deployment) + candidates := f.registry.Candidates(apiKey, agentName, deployment) if len(candidates) == 0 && f.fallback == nil { w.Header().Set("Retry-After", "1") f.writeMiss(w, http.StatusServiceUnavailable, MissUnavailable, "no workers available for deployment") @@ -215,7 +222,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { // nothing local can serve: hand off before the local status mapping if f.fallback(w, r, &FallbackRequest{ APIKey: apiKey, Authenticated: authenticated, - Deployment: deployment, Path: path, + AgentName: agentName, Deployment: deployment, Method: r.Method, Path: path, }) { return } @@ -236,7 +243,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { for _, reg := range candidates { if alt, ok := reg.Manifest.RedirectSlashes(path, r.Method); ok { u := *r.URL - u.Path = PathPrefix + deployment + alt + u.Path = PathPrefix + agentName + "/" + deployment + alt http.Redirect(w, r, u.String(), http.StatusTemporaryRedirect) return } @@ -271,7 +278,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { if bodyConsumed == 0 && f.fallback != nil { if f.fallback(w, r, &FallbackRequest{ APIKey: apiKey, Authenticated: authenticated, - Deployment: deployment, Path: path, + AgentName: agentName, Deployment: deployment, Method: r.Method, Path: path, }) { return } @@ -285,13 +292,12 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { // registrations at random and take the one with fewer in-flight streams (least // outstanding requests). This approximates optimal load spreading without global // coordination or the herding of exact least-loaded, and - unlike a -// self-reported load - the in-flight count is observed here on this node's own -// data conns, so it is accurate for adopted (satellite) registrations too. -// Eligible = not already attempted, has attached conns, not draining. +// self-reported load - the in-flight count is observed here from the worker's +// own session. Eligible = not already attempted, has a live session, not draining. func pickWorker(regs []*Registration, ignore map[*Registration]bool) *Registration { eligible := regs[:0:0] for _, reg := range regs { - if ignore[reg] || reg.AttachedConns() == 0 { + if ignore[reg] || !reg.HasSession() { continue } if reg.Draining != nil && reg.Draining() { @@ -335,23 +341,15 @@ func (f *Front) bridge( body io.Reader, bodyConsumedBefore int64, ) (done bool, retryable bool) { - conn := reg.PickConn() - if conn == nil { - return false, true // no capacity here; try another worker - } - - stream, err := conn.OpenStream(&livekit.AgentHttp_HttpStreamOpen{ - RequestId: guid.New("AER_"), - ClientAddr: r.RemoteAddr, - }) + ctx := r.Context() + stream, err := reg.OpenStream(ctx) if err != nil { - return false, true + return false, true // no session/capacity here; try another worker } defer stream.Close() - ctx := r.Context() stop := context.AfterFunc(ctx, func() { - stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client disconnected") + stream.Reset(ResetCancel, "client disconnected") }) defer stop() @@ -365,7 +363,7 @@ func (f *Front) bridge( err = stream.CloseWrite() } else { // fail fast: the worker is waiting for bytes that will never come - stream.Reset(livekit.AgentHttp_HSR_CANCEL, "request write failed") + stream.Reset(ResetCancel, "request write failed") } writeErrCh <- err }() @@ -389,7 +387,7 @@ func (f *Front) bridge( } // join the request writer before another attempt touches the shared // body reader (retries are bodyless per the table, so this is prompt) - stream.Reset(livekit.AgentHttp_HSR_CANCEL, "retrying elsewhere") + stream.Reset(ResetCancel, "retrying elsewhere") <-writeErrCh return false, true } @@ -406,7 +404,7 @@ func (f *Front) bridge( n, rerr := resp.Body.Read(buf) if n > 0 { if _, werr := w.Write(buf[:n]); werr != nil { - stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client write failed") + stream.Reset(ResetCancel, "client write failed") return true, false } _ = rc.Flush() @@ -427,16 +425,18 @@ func (f *Front) bridge( return true, false } -// readResponseHead reads the worker's response head, relaying informational -// responses (1xx except 101) to the client. -func (f *Front) readResponseHead(w http.ResponseWriter, br *bufio.Reader, outReq *http.Request, stream *Stream) (*http.Response, error) { +// readResponseHead reads the worker's response head, relaying 1xx informational +// responses to the client and returning the final head. A 101 is not treated as +// informational (there is no protocol upgrade over an opaque request stream): it +// is returned as-is as the final response rather than looped past. +func (f *Front) readResponseHead(w http.ResponseWriter, br *bufio.Reader, outReq *http.Request, stream Stream) (*http.Response, error) { deadline := time.NewTimer(responseHeadTimeout) defer deadline.Stop() headCh := make(chan struct{}) go func() { select { case <-deadline.C: - stream.Reset(livekit.AgentHttp_HSR_CANCEL, "response head timeout") + stream.Reset(ResetCancel, "response head timeout") case <-headCh: } }() @@ -463,7 +463,7 @@ func (f *Front) readResponseHead(w http.ResponseWriter, br *bufio.Reader, outReq } // classifyRetry implements the retry table. -func (f *Front) classifyRetry(r *http.Request, stream *Stream, responseBytes, bodyConsumedBefore int64, err error) bool { +func (f *Front) classifyRetry(r *http.Request, stream Stream, responseBytes, bodyConsumedBefore int64, err error) bool { if responseBytes > 0 || stream.BytesRead() > 0 { return false } @@ -472,9 +472,6 @@ func (f *Front) classifyRetry(r *http.Request, stream *Stream, responseBytes, bo // request body can be replayed (nothing consumed yet) return bodyConsumedBefore == 0 && r.ContentLength == 0 } - if errors.Is(err, ErrStreamRefused) { - return bodyConsumedBefore == 0 && r.ContentLength == 0 - } switch r.Method { case http.MethodGet, http.MethodHead, http.MethodOptions: return r.ContentLength == 0 @@ -549,4 +546,3 @@ func (c *countingReader) Read(p []byte) (int, error) { *c.n += int64(n) return n, err } - diff --git a/pkg/agent/endpoint/manifest.go b/pkg/agent/endpoint/manifest.go index 961ad23d9..ed867a9ca 100644 --- a/pkg/agent/endpoint/manifest.go +++ b/pkg/agent/endpoint/manifest.go @@ -16,6 +16,7 @@ package endpoint import ( "fmt" + "net/http" "slices" "strings" @@ -24,6 +25,21 @@ import ( const MaxManifestRoutes = 256 +// allowedMethods is the set of HTTP methods a route may declare: exactly the +// verbs FastAPI (starlette) can route. Registration rejects anything else, so a +// typo ("GTE") or a garbage token can't sit in a manifest silently never +// matching. Widen this only if the worker side ever routes beyond FastAPI. +var allowedMethods = map[string]struct{}{ + http.MethodGet: {}, + http.MethodHead: {}, + http.MethodPost: {}, + http.MethodPut: {}, + http.MethodPatch: {}, + http.MethodDelete: {}, + http.MethodOptions: {}, + http.MethodTrace: {}, +} + // Route is one validated manifest entry. type Route struct { Template *Template @@ -77,6 +93,9 @@ func ParseManifest(endpoints []*livekit.AgentHttp_AgentEndpoint) (*Manifest, err if u != method { return nil, fmt.Errorf("endpoint %q method %q must be uppercase", ep.GetPath(), method) } + if _, ok := allowedMethods[u]; !ok { + return nil, fmt.Errorf("endpoint %q declares unsupported method %q", ep.GetPath(), method) + } methods = append(methods, u) } m.routes = append(m.routes, Route{ diff --git a/pkg/agent/endpoint/registry.go b/pkg/agent/endpoint/registry.go index 49a0b70fa..ad0e27d6b 100644 --- a/pkg/agent/endpoint/registry.go +++ b/pkg/agent/endpoint/registry.go @@ -15,29 +15,25 @@ package endpoint import ( - "crypto/subtle" + "context" "errors" "slices" "sync" "github.com/livekit/protocol/livekit" - "github.com/livekit/protocol/logger" - "github.com/livekit/protocol/utils/guid" ) // DefaultDeployment is the URL segment that addresses workers registered with an // empty deployment name (self-hosted workers typically set none). const DefaultDeployment = "default" -var ( - ErrAttachRejected = errors.New("attach rejected") - ErrUnknownWorker = errors.New("unknown worker registration") - // ErrWrongEpoch means a registration exists for the worker but the wire's - // credentials name a different epoch: a multi-node adopter re-validates - // against the registration holder and refreshes, since the holder is the - // authority on the current epoch. - ErrWrongEpoch = errors.New("wrong registration epoch") -) +// DefaultMaxStreams is the soft per-session concurrency cap used only for +// capacity weighting; QUIC's own stream limit is the hard bound. +const DefaultMaxStreams = 256 + +// ErrNoSession means the registration has no live data-plane session, so no +// stream can be opened toward the worker. +var ErrNoSession = errors.New("registration has no data-plane session") func normalizeDeployment(d string) string { if d == "" { @@ -46,187 +42,84 @@ func normalizeDeployment(d string) string { return d } -// Registration is one worker's data-plane state: its manifest, negotiated -// settings, and attached data connections. It lives exactly -// as long as the worker's control connection (epoch fencing: a reconnecting -// worker forms a new registration with a fresh attach token; connections of the -// old epoch die with it). +// Registration is one worker's data-plane state: its manifest and the single +// WebTransport session that carries both its control stream and the HTTP +// exchanges the node opens toward it. It lives exactly as long as that session +// (epoch fencing: a reconnecting worker forms a new registration, and the old +// session dies with it). type Registration struct { WorkerID string - InstanceID string APIKey string + AgentName string Deployment string Manifest *Manifest // Endpoints is the raw manifest as declared, kept so a multi-node layer can // replicate the route set for remote path matching without re-deriving it. Endpoints []*livekit.AgentHttp_AgentEndpoint - Settings Settings - Logger logger.Logger // Draining is provided by the control-plane layer that owns the worker; a // draining worker takes no new streams. Worker selection uses live in-flight - // streams, not a reported load, so no load hook is needed here. + // streams, not a reported load. Draining func() bool - // pendingAttaches counts slots reserved by validated-but-unadopted wires so - // two racing attaches at the cap cannot both be acked - pendingAttaches int - - mu sync.Mutex - conns []*DataConn - closed bool + mu sync.Mutex + session Session + closed bool } -func (r *Registration) validateAttach(instanceID, token string) error { - if subtle.ConstantTimeCompare([]byte(token), []byte(r.Settings.AttachToken)) != 1 { - return ErrAttachRejected - } - if instanceID != r.InstanceID { - return ErrAttachRejected - } +// SetSession attaches the worker's live data-plane session. One session per +// worker: the WebTransport session that also carries its control stream. +func (r *Registration) SetSession(s Session) { r.mu.Lock() - defer r.mu.Unlock() - if r.closed { - return ErrAttachRejected - } - if len(r.conns) >= int(r.Settings.DataConnCount) { - return ErrAttachRejected - } - return nil -} - -// beginAttach validates the wire's credentials and atomically reserves a pool -// slot, so the attach ack can be written before adoption without two racing -// wires both being acked at the cap. -func (r *Registration) beginAttach(instanceID, token string) (*AttachTicket, error) { - if subtle.ConstantTimeCompare([]byte(token), []byte(r.Settings.AttachToken)) != 1 || - instanceID != r.InstanceID { - return nil, ErrWrongEpoch - } - r.mu.Lock() - defer r.mu.Unlock() - if r.closed { - return nil, ErrAttachRejected - } - // the pool is fixed size: tolerate redials, but never more live conns than - // negotiated - if len(r.conns)+r.pendingAttaches >= int(r.Settings.DataConnCount) { - return nil, ErrAttachRejected - } - r.pendingAttaches++ - return &AttachTicket{r: r}, nil -} - -// AttachTicket is a reserved pool slot: Complete adopts the wire into it, Abort -// releases it. -type AttachTicket struct { - r *Registration - used bool -} - -// Complete adopts the wire; it reports false when the registration closed while -// the ack was in flight (the caller must close the wire). -func (t *AttachTicket) Complete(wire WireConn, params WireParams) bool { - r := t.r - r.mu.Lock() - defer r.mu.Unlock() - if t.used { - return false - } - t.used = true - r.pendingAttaches-- - if r.closed { - return false - } - conn := NewDataConn(wire, params, r.removeConn, r.Logger) - r.conns = append(r.conns, conn) - return true -} - -func (t *AttachTicket) Abort() { - r := t.r - r.mu.Lock() - defer r.mu.Unlock() - if t.used { - return - } - t.used = true - r.pendingAttaches-- -} - -func (r *Registration) removeConn(c *DataConn) { - r.mu.Lock() - if i := slices.Index(r.conns, c); i != -1 { - r.conns = slices.Delete(r.conns, i, i+1) - } + r.session = s r.mu.Unlock() } -// PickConn places a new stream: the lightest non-heavy connection with capacity, -// falling back to spread-evenly when everything is heavy. -func (r *Registration) PickConn() *DataConn { - r.mu.Lock() - conns := slices.Clone(r.conns) - r.mu.Unlock() - - var best, bestHeavy *DataConn - var bestScore, bestHeavyScore int64 - for _, c := range conns { - if !c.HasCapacity() { - continue - } - score, heavy := c.Score() - if heavy { - if bestHeavy == nil || score < bestHeavyScore { - bestHeavy, bestHeavyScore = c, score - } - continue - } - if best == nil || score < bestScore { - best, bestScore = c, score - } - } - if best != nil { - return best - } - return bestHeavy -} - -// AttachedConns reports live data connections. -func (r *Registration) AttachedConns() int { +func (r *Registration) getSession() Session { r.mu.Lock() defer r.mu.Unlock() - return len(r.conns) + if r.closed { + return nil + } + return r.session } -// InflightStreams reports the open streams across this registration's data -// conns - the requests currently being served from this node. It is the -// least-outstanding-requests signal for worker selection: a per-node view (each -// node counts only the streams it opened), which is exactly what a node needs -// to spread its own traffic, and it needs no load reported by the worker. +// OpenStream opens a data-plane stream toward the worker for one HTTP exchange. +func (r *Registration) OpenStream(ctx context.Context) (Stream, error) { + s := r.getSession() + if s == nil { + return nil, ErrNoSession + } + return s.OpenStream(ctx) +} + +// HasSession reports whether the worker has a live data-plane session, i.e. can +// serve HTTP requests right now. +func (r *Registration) HasSession() bool { + return r.getSession() != nil +} + +// InflightStreams reports the open streams on this worker's session - the +// least-outstanding-requests signal for worker selection. func (r *Registration) InflightStreams() int { - r.mu.Lock() - defer r.mu.Unlock() - n := 0 - for _, c := range r.conns { - n += c.OpenStreams() + s := r.getSession() + if s == nil { + return 0 } - return n + return s.OpenStreams() } -// SpareStreams reports the total spare stream capacity across this -// registration's data conns on this node - how many more concurrent requests it -// can accept. It is the node's live, locally-observed serving headroom for the -// worker; being measured on the conns this node holds, it is accurate for -// adopted (satellite) registrations too, which report no load. +// SpareStreams reports remaining stream capacity - the node's live serving +// headroom for the worker, used to weight node selection. func (r *Registration) SpareStreams() int { - r.mu.Lock() - defer r.mu.Unlock() - n := 0 - for _, c := range r.conns { - n += c.SpareStreams() + s := r.getSession() + if s == nil { + return 0 } - return n + if spare := s.MaxStreams() - s.OpenStreams(); spare > 0 { + return spare + } + return 0 } func (r *Registration) close() { @@ -236,17 +129,16 @@ func (r *Registration) close() { return } r.closed = true - conns := slices.Clone(r.conns) - r.conns = nil + s := r.session + r.session = nil r.mu.Unlock() - - for _, c := range conns { - c.Close(ErrConnClosed) + if s != nil { + s.Close("registration closed") } } // Registry tracks data-plane registrations on this node, keyed by -// (api key, deployment). The api key is the project identity in OSS. +// (api key, agent name, deployment). The api key is the project identity in OSS. type Registry struct { mu sync.Mutex regs map[string]*Registration // by worker id @@ -255,6 +147,7 @@ type Registry struct { type regKey struct { apiKey string + agentName string deployment string } @@ -265,17 +158,12 @@ func NewRegistry() *Registry { } } -// NewAttachToken mints a registration's attach token. -func NewAttachToken() string { - return guid.New("ATT_") -} - -// Register adopts a registration. A worker id already present is superseded: -// worker ids are stable across reconnects, and the retiring control connection -// must not be able to strand the new epoch (its own Deregister is a no-op once -// replaced). The superseded epoch's connections die with it. +// Register records a registration. A worker id already present is superseded: +// worker ids are stable across reconnects, and the retiring session must not be +// able to strand the new epoch (its own Deregister is a no-op once replaced). +// The superseded epoch's session is closed. func (g *Registry) Register(r *Registration) error { - key := regKey{r.APIKey, normalizeDeployment(r.Deployment)} + key := regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)} g.mu.Lock() old := g.regs[r.WorkerID] if old != nil { @@ -293,7 +181,7 @@ func (g *Registry) Register(r *Registration) error { // removeLocked unlinks a registration from all indexes. Callers hold g.mu. func (g *Registry) removeLocked(r *Registration) { delete(g.regs, r.WorkerID) - key := regKey{r.APIKey, normalizeDeployment(r.Deployment)} + key := regKey{r.APIKey, r.AgentName, normalizeDeployment(r.Deployment)} if regs := g.byKey[key]; len(regs) > 0 { if i := slices.Index(regs, r); i != -1 { regs = slices.Delete(regs, i, i+1) @@ -319,45 +207,11 @@ func (g *Registry) Deregister(r *Registration) { r.close() } -// ValidateAttach checks an attach without adopting the connection, so the -// attach response can be written before stream frames may flow. apiKey is the -// project the connecting grant is authorized for; it must match the -// registration's, so a leaked attach token is useless without a grant for the -// worker's own project (defense in depth over the token secret). -func (g *Registry) ValidateAttach(workerID, instanceID, apiKey, token string) error { - g.mu.Lock() - r, ok := g.regs[workerID] - g.mu.Unlock() - if !ok { - return ErrUnknownWorker - } - if apiKey != r.APIKey { - return ErrAttachRejected // wrong project: terminal, never adopt-retried - } - return r.validateAttach(instanceID, token) -} - -// BeginAttach validates an attach and reserves a pool slot; the caller writes -// the ack and then Completes (or Aborts) the ticket. apiKey binds the attach to -// the connecting grant's project (see ValidateAttach). -func (g *Registry) BeginAttach(workerID, instanceID, apiKey, token string) (*AttachTicket, error) { - g.mu.Lock() - r, ok := g.regs[workerID] - g.mu.Unlock() - if !ok { - return nil, ErrUnknownWorker - } - if apiKey != r.APIKey { - return nil, ErrAttachRejected // wrong project: terminal, never adopt-retried - } - return r.beginAttach(instanceID, token) -} - -// Candidates returns the registrations for (api key, deployment segment). -func (g *Registry) Candidates(apiKey, deployment string) []*Registration { +// Candidates returns the registrations for (api key, agent name, deployment segment). +func (g *Registry) Candidates(apiKey, agentName, deployment string) []*Registration { g.mu.Lock() defer g.mu.Unlock() - return slices.Clone(g.byKey[regKey{apiKey, normalizeDeployment(deployment)}]) + return slices.Clone(g.byKey[regKey{apiKey, agentName, normalizeDeployment(deployment)}]) } // SingleAPIKey returns the api key when every registration shares one - the OSS diff --git a/pkg/agent/endpoint/registry_test.go b/pkg/agent/endpoint/registry_test.go index 559e51e74..daac7daaa 100644 --- a/pkg/agent/endpoint/registry_test.go +++ b/pkg/agent/endpoint/registry_test.go @@ -1,17 +1,31 @@ +// Copyright 2026 LiveKit, Inc. + package endpoint import ( + "context" "testing" "github.com/stretchr/testify/require" "github.com/livekit/protocol/livekit" - "github.com/livekit/protocol/logger" ) +// fakeSession is a no-op Session for registry tests: it records whether it was +// closed so supersede/deregister can be asserted. +type fakeSession struct { + open int + closed bool +} + +func (s *fakeSession) OpenStream(context.Context) (Stream, error) { return nil, ErrNoSession } +func (s *fakeSession) OpenStreams() int { return s.open } +func (s *fakeSession) MaxStreams() int { return DefaultMaxStreams } +func (s *fakeSession) Close(string) { s.closed = true } + // worker ids are stable across reconnects: a re-registration must supersede the -// old epoch, and the retiring connection's Deregister must not strand the new -// one. +// old epoch and close its session, and the retiring session's Deregister must +// not strand the new one. func TestRegistrySupersede(t *testing.T) { g := NewRegistry() manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{ @@ -19,32 +33,58 @@ func TestRegistrySupersede(t *testing.T) { }}) require.NoError(t, err) - mk := func(instance string) *Registration { - return &Registration{ - WorkerID: "AW_1", InstanceID: instance, APIKey: "key", - Deployment: "production", Manifest: manifest, - Settings: Settings{DataConnCount: 1, AttachToken: "ATT_" + instance}, - Logger: logger.GetLogger(), + mk := func() (*Registration, *fakeSession) { + r := &Registration{ + WorkerID: "AW_1", APIKey: "key", + AgentName: "agent", Deployment: "production", Manifest: manifest, } + s := &fakeSession{} + r.SetSession(s) + return r, s } - oldReg := mk("i-1") + oldReg, oldSess := mk() require.NoError(t, g.Register(oldReg)) - newReg := mk("i-2") + newReg, newSess := mk() require.NoError(t, g.Register(newReg)) - require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production")) - require.Error(t, g.ValidateAttach("AW_1", "i-1", "key", "ATT_i-1"), "old epoch must not attach") - require.NoError(t, g.ValidateAttach("AW_1", "i-2", "key", "ATT_i-2")) - // a grant for another project cannot attach even with the right token - require.ErrorIs(t, g.ValidateAttach("AW_1", "i-2", "other", "ATT_i-2"), ErrAttachRejected, - "wrong-project grant must not attach") + require.Equal(t, []*Registration{newReg}, g.Candidates("key", "agent", "production")) + require.True(t, oldSess.closed, "superseded epoch's session must be closed") + require.False(t, newSess.closed) // the old control connection tears down after the new one registered g.Deregister(oldReg) - require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production"), + require.Equal(t, []*Registration{newReg}, g.Candidates("key", "agent", "production"), "the retiring epoch must not deregister its successor") g.Deregister(newReg) - require.Empty(t, g.Candidates("key", "production")) + require.Empty(t, g.Candidates("key", "agent", "production")) + require.True(t, newSess.closed, "deregistered session must be closed") +} + +// Candidates keys on (apiKey, agentName, deployment): a request for one agent +// must not see another agent's workers in the same project/deployment. +func TestRegistryAgentScoping(t *testing.T) { + g := NewRegistry() + manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{ + Path: "/x", Methods: []string{"GET"}, Public: true, + }}) + require.NoError(t, err) + + mk := func(workerID, agentName string) *Registration { + r := &Registration{ + WorkerID: workerID, APIKey: "key", + AgentName: agentName, Deployment: "production", Manifest: manifest, + } + r.SetSession(&fakeSession{}) + return r + } + a := mk("AW_a", "alpha") + b := mk("AW_b", "beta") + require.NoError(t, g.Register(a)) + require.NoError(t, g.Register(b)) + + require.Equal(t, []*Registration{a}, g.Candidates("key", "alpha", "production")) + require.Equal(t, []*Registration{b}, g.Candidates("key", "beta", "production")) + require.Empty(t, g.Candidates("key", "gamma", "production")) } diff --git a/pkg/agent/endpoint/routematch.go b/pkg/agent/endpoint/routematch.go index 7dac75901..2b2053995 100644 --- a/pkg/agent/endpoint/routematch.go +++ b/pkg/agent/endpoint/routematch.go @@ -22,18 +22,20 @@ import ( // Route matching for the multi-node edge. Templates are canonicalized by // replacing each path parameter with a typed wildcard token ({id:int} -> the // int token, {name} or {name:str} -> the str token, {rest:path} -> the greedy -// glob), then indexed by ALL prefixes plus the full route as a terminal. A -// concrete request path is matched by a left-to-right pruned walk: at each -// segment we branch over the literal value and the wildcard tokens the value is +// glob), then indexed by ALL prefixes (path-only) plus, for each declared HTTP +// method, a method-qualified terminal route key. A concrete request (method + +// path) is matched by a left-to-right pruned walk over the path: at each segment +// we branch over the literal value and the wildcard tokens the value is // type-compatible with, keeping only branches that are a prefix of some route - // so the cost is linear in the path depth, never 2^segments, and the surviving -// set never exceeds the route count. +// set never exceeds the route count - then the terminal check is method-aware, +// so a node that serves the path under a different method is not a candidate. // -// This is a SHAPE match only: the exact convertor (int/uuid/...) and the method -// are re-checked by the worker's own router, which stays authoritative. A shape -// match that the worker then rejects costs one wasted relay, never a wrong -// route; a false positive from a probabilistic backing (see FilterIndex) is the -// same bounded cost. +// This is a SHAPE match only: the exact convertor (int/uuid/...) is re-checked +// by the worker's own router, which stays authoritative (the method is now part +// of the match, but the worker still enforces it). A shape match the worker then +// rejects costs one wasted relay, never a wrong route; a false positive from a +// probabilistic backing (see FilterIndex) is the same bounded cost. // canonical wildcard tokens; the control-char prefix keeps them from colliding // with any real path segment @@ -158,6 +160,15 @@ func extendKey(prefix, tok string) string { return prefix + segSep + tok } +// RouteKey qualifies a canonical path terminal with its HTTP method, so a +// terminal is a route only for the method(s) the worker declared. Methods are +// uppercase ASCII (allowlisted at registration) and canonical tokens start with +// a control char, so method and path can't collide. Prefix keys stay path-only +// (the pruned walk is method-agnostic until the terminal check). +func RouteKey(method, canonicalRoute string) string { + return method + segSep + canonicalRoute +} + // PrefixIndex is the backing an edge matches against: whether a canonical prefix // key is a prefix of some route, and whether it is a complete route. An exact // map-backed index (ExactIndex) gives precise answers; a cuckoo-filter backing @@ -167,10 +178,12 @@ type PrefixIndex interface { IsRoute(key string) bool } -// Matches reports whether the request path matches any route in the index, via -// the pruned walk. A path deeper than maxMatchDepth returns true so the caller -// relays and lets the worker decide (never a spurious 404 on a deep path). -func Matches(idx PrefixIndex, path string) bool { +// Matches reports whether the request (method + path) matches any route in the +// index, via the pruned walk. The path walk is method-agnostic (prefix keys are +// path-only); only the terminal check is method-qualified. A path deeper than +// maxMatchDepth returns true so the caller relays and lets the worker decide +// (never a spurious 404 on a deep path). +func Matches(idx PrefixIndex, method, path string) bool { segs := splitSegments(path) if len(segs) > maxMatchDepth { return true @@ -180,7 +193,7 @@ func Matches(idx PrefixIndex, path string) bool { for _, seg := range segs { // a glob at any currently-live prefix eats this segment and the rest for p := range live { - if idx.IsRoute(extendKey(p, tokGlob)) { + if idx.IsRoute(RouteKey(method, extendKey(p, tokGlob))) { return true } } @@ -198,10 +211,10 @@ func Matches(idx PrefixIndex, path string) bool { } live = next } - // consumed every segment: a live prefix that is a complete route matches, - // and a glob at a live prefix matches the empty remainder too + // consumed every segment: a live prefix that is a complete route for this + // method matches, and a glob at a live prefix matches the empty remainder too for p := range live { - if idx.IsRoute(p) || idx.IsRoute(extendKey(p, tokGlob)) { + if idx.IsRoute(RouteKey(method, p)) || idx.IsRoute(RouteKey(method, extendKey(p, tokGlob))) { return true } } @@ -224,8 +237,9 @@ func NewExactIndex() *ExactIndex { } } -// Add canonicalizes a template and inserts its prefixes and terminal. -func (x *ExactIndex) Add(path string) error { +// Add canonicalizes a template and inserts its path prefixes plus a +// method-qualified terminal per declared method. +func (x *ExactIndex) Add(methods []string, path string) error { segs, err := canonicalizeTemplate(path) if err != nil { return err @@ -235,7 +249,9 @@ func (x *ExactIndex) Add(path string) error { key = extendKey(key, s) x.prefixes[key] = struct{}{} } - x.routes[key] = struct{}{} + for _, m := range methods { + x.routes[RouteKey(m, key)] = struct{}{} + } return nil } diff --git a/pkg/agent/endpoint/routematch_test.go b/pkg/agent/endpoint/routematch_test.go index d51b8aaa6..b3f2d2a50 100644 --- a/pkg/agent/endpoint/routematch_test.go +++ b/pkg/agent/endpoint/routematch_test.go @@ -3,49 +3,57 @@ package endpoint import ( + "net/http" "testing" "github.com/stretchr/testify/require" ) +// idx builds an ExactIndex whose routes all serve GET; the path-shape tests +// below match with GET via the matches helper. Method-awareness has its own test. func idx(t *testing.T, routes ...string) *ExactIndex { x := NewExactIndex() for _, r := range routes { - require.NoError(t, x.Add(r), r) + require.NoError(t, x.Add([]string{http.MethodGet}, r), r) } return x } +// matches is a GET-method shorthand so the path-shape assertions read cleanly. +func matches(idx PrefixIndex, path string) bool { + return Matches(idx, http.MethodGet, path) +} + func TestRouteMatchStaticAndParams(t *testing.T) { x := idx(t, "/sms", "/orders/{id}/items", "/orders/{id}/list", "/health") // static - require.True(t, Matches(x, "/sms")) - require.True(t, Matches(x, "/health")) - require.False(t, Matches(x, "/unknown")) + require.True(t, matches(x, "/sms")) + require.True(t, matches(x, "/health")) + require.False(t, matches(x, "/unknown")) // single-segment params, and the two sub-routes are distinguished - require.True(t, Matches(x, "/orders/42/items")) - require.True(t, Matches(x, "/orders/abc/list")) - require.False(t, Matches(x, "/orders/42/refund")) // neither sub-route exists - require.False(t, Matches(x, "/orders/42")) // prefix only, not a full route - require.False(t, Matches(x, "/orders/42/items/x")) + require.True(t, matches(x, "/orders/42/items")) + require.True(t, matches(x, "/orders/abc/list")) + require.False(t, matches(x, "/orders/42/refund")) // neither sub-route exists + require.False(t, matches(x, "/orders/42")) // prefix only, not a full route + require.False(t, matches(x, "/orders/42/items/x")) // trailing slash matches the same shape (worker 307s) - require.True(t, Matches(x, "/orders/42/items/")) + require.True(t, matches(x, "/orders/42/items/")) } func TestRouteMatchTypedParams(t *testing.T) { x := idx(t, "/orders/{id:int}", "/u/{u:uuid}") - require.True(t, Matches(x, "/orders/42")) // int - require.False(t, Matches(x, "/orders/abc")) // not an int -> edge miss, no relay - require.False(t, Matches(x, "/orders/3.14")) // float is not int + require.True(t, matches(x, "/orders/42")) // int + require.False(t, matches(x, "/orders/abc")) // not an int -> edge miss, no relay + require.False(t, matches(x, "/orders/3.14")) // float is not int uuid := "550e8400-e29b-41d4-a716-446655440000" - require.True(t, Matches(x, "/u/"+uuid)) - require.True(t, Matches(x, "/u/"+"550e8400e29b41d4a716446655440000")) // hyphenless - require.False(t, Matches(x, "/u/not-a-uuid")) + require.True(t, matches(x, "/u/"+uuid)) + require.True(t, matches(x, "/u/"+"550e8400e29b41d4a716446655440000")) // hyphenless + require.False(t, matches(x, "/u/not-a-uuid")) } func TestRouteMatchUUIDvsStr(t *testing.T) { @@ -53,45 +61,45 @@ func TestRouteMatchUUIDvsStr(t *testing.T) { uuid := "550e8400-e29b-41d4-a716-446655440000" strOnly := idx(t, "/x/{v}") // str - require.True(t, Matches(strOnly, "/x/"+uuid)) // uuid value matches a str route - require.True(t, Matches(strOnly, "/x/anything")) + require.True(t, matches(strOnly, "/x/"+uuid)) // uuid value matches a str route + require.True(t, matches(strOnly, "/x/anything")) uuidOnly := idx(t, "/x/{v:uuid}") - require.True(t, Matches(uuidOnly, "/x/"+uuid)) - require.False(t, Matches(uuidOnly, "/x/anything")) // non-uuid never matches a uuid route + require.True(t, matches(uuidOnly, "/x/"+uuid)) + require.False(t, matches(uuidOnly, "/x/anything")) // non-uuid never matches a uuid route } func TestRouteMatchGlob(t *testing.T) { x := idx(t, "/files/{rest:path}", "/static") - require.True(t, Matches(x, "/files/a")) - require.True(t, Matches(x, "/files/a/b/c/d.txt")) // spans segments + require.True(t, matches(x, "/files/a")) + require.True(t, matches(x, "/files/a/b/c/d.txt")) // spans segments // /files with no trailing slash is a tolerated FALSE POSITIVE (Starlette // requires the separating slash); the worker returns the real status - require.True(t, Matches(x, "/files")) - require.False(t, Matches(x, "/other/a/b")) - require.True(t, Matches(x, "/static")) + require.True(t, matches(x, "/files")) + require.False(t, matches(x, "/other/a/b")) + require.True(t, matches(x, "/static")) } // a path convertor that is not the last segment must still match (the glob is // treated as terminal, over-approximating the suffix) func TestRouteMatchNonTerminalGlob(t *testing.T) { x := idx(t, "/files/{rest:path}/edit") - require.True(t, Matches(x, "/files/a/b/edit")) - require.True(t, Matches(x, "/files/a/edit")) + require.True(t, matches(x, "/files/a/b/edit")) + require.True(t, matches(x, "/files/a/edit")) // suffix is over-approximated: a path under /files matches even without // /edit (the worker's router returns the real 404) - require.True(t, Matches(x, "/files/a/b")) - require.False(t, Matches(x, "/other/a/edit")) + require.True(t, matches(x, "/files/a/b")) + require.False(t, matches(x, "/other/a/edit")) } // a path convertor mixed with literal text in a segment spans slashes and must // not be narrowed to a single-segment str func TestRouteMatchMixedGlob(t *testing.T) { x := idx(t, "/files/pre{rest:path}") - require.True(t, Matches(x, "/files/prea")) - require.True(t, Matches(x, "/files/prea/b/c")) // spans segments (would drop under str) - require.False(t, Matches(x, "/other/x")) + require.True(t, matches(x, "/files/prea")) + require.True(t, matches(x, "/files/prea/b/c")) // spans segments (would drop under str) + require.False(t, matches(x, "/other/x")) } func TestRouteMatchHeterogeneousShapes(t *testing.T) { @@ -99,18 +107,41 @@ func TestRouteMatchHeterogeneousShapes(t *testing.T) { ints := idx(t, "/item/{id:int}") strs := idx(t, "/item/{slug}") - require.True(t, Matches(ints, "/item/42")) - require.False(t, Matches(ints, "/item/foo")) // routes only to the str worker - require.True(t, Matches(strs, "/item/foo")) - require.True(t, Matches(strs, "/item/42")) // str accepts a numeric slug too + require.True(t, matches(ints, "/item/42")) + require.False(t, matches(ints, "/item/foo")) // routes only to the str worker + require.True(t, matches(strs, "/item/foo")) + require.True(t, matches(strs, "/item/42")) // str accepts a numeric slug too } func TestRouteMatchLiteralBeatsWildcardCoexist(t *testing.T) { // a literal segment and a param at the same position coexist x := idx(t, "/users/me", "/users/{id:int}") - require.True(t, Matches(x, "/users/me")) // literal - require.True(t, Matches(x, "/users/42")) // param - require.False(t, Matches(x, "/users/abc")) // neither: not "me", not an int + require.True(t, matches(x, "/users/me")) // literal + require.True(t, matches(x, "/users/42")) // param + require.False(t, matches(x, "/users/abc")) // neither: not "me", not an int +} + +// the terminal is method-qualified: the SAME path shape served under different +// methods (e.g. two workers, POST vs PUT) matches only for the declared method, +// so the cross-node presence filter is method-aware. +func TestRouteMatchMethodAware(t *testing.T) { + postOnly := NewExactIndex() + require.NoError(t, postOnly.Add([]string{http.MethodPost}, "/test")) + putOnly := NewExactIndex() + require.NoError(t, putOnly.Add([]string{http.MethodPut}, "/test")) + + require.True(t, Matches(postOnly, http.MethodPost, "/test")) + require.False(t, Matches(postOnly, http.MethodPut, "/test")) + require.True(t, Matches(putOnly, http.MethodPut, "/test")) + require.False(t, Matches(putOnly, http.MethodPost, "/test")) + + // a route declaring several methods matches each, and only those + multi := NewExactIndex() + require.NoError(t, multi.Add([]string{http.MethodGet, http.MethodHead}, "/thing/{id:int}")) + require.True(t, Matches(multi, http.MethodGet, "/thing/5")) + require.True(t, Matches(multi, http.MethodHead, "/thing/5")) + require.False(t, Matches(multi, http.MethodDelete, "/thing/5")) + require.False(t, Matches(multi, http.MethodGet, "/thing/abc")) // still shape-checked } func TestRouteMatchDepthCapFallsBackToRelay(t *testing.T) { @@ -119,16 +150,17 @@ func TestRouteMatchDepthCapFallsBackToRelay(t *testing.T) { for i := 0; i < maxMatchDepth+5; i++ { deep += "x/" } - require.True(t, Matches(x, deep)) // over cap -> relay and let the worker decide + require.True(t, matches(x, deep)) // over cap -> relay and let the worker decide } -func TestRouteDepthRejectedAtAdd(t *testing.T) { +// depth enforcement lives in the manifest layer (ParseManifest), not in the +// index: Add canonicalizes any depth without error, so the two layers stay +// decoupled. +func TestRouteDepthNotEnforcedAtAdd(t *testing.T) { x := NewExactIndex() long := "/" for i := 0; i < MaxRouteDepth+2; i++ { long += "s/" } - // canonicalization succeeds (depth enforcement lives in the manifest layer); - // this test just documents that very deep templates still canonicalize - require.NoError(t, x.Add(long)) + require.NoError(t, x.Add([]string{http.MethodGet}, long)) } diff --git a/pkg/agent/endpoint/scheduler.go b/pkg/agent/endpoint/scheduler.go deleted file mode 100644 index eab70c234..000000000 --- a/pkg/agent/endpoint/scheduler.go +++ /dev/null @@ -1,296 +0,0 @@ -// Copyright 2026 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 endpoint - -import ( - "slices" - "sync" - "time" - - "github.com/livekit/protocol/livekit" -) - -// scheduler serializes all writes on one wire through a single writer goroutine -// with two classes: control frames (OPEN/CREDIT/RESET) preempt DATA, and DATA -// drains one chunk per stream in round-robin order; a stream's EOF rides its -// data queue so it can never overtake the stream's own bytes. Without the class -// split, credits queue behind a deep data backlog and the peer's send windows -// never refill; without the round-robin, one heavy -// stream starves its siblings; without the budget, many streams times a full -// credit window pin unbounded memory; without the write deadline, a peer that -// stops reading freezes cancellations too. -// -// Locking rule: the wire write happens OUTSIDE the scheduler lock. Holding the -// lock across a blocking write would let one stalled connection stop enqueues -// from every stream sharing it. -// chunkPoolMinSize is the payload size at or above which a frame's buffer is -// drawn from the recycle pool. Below it (webhook-sized bodies, headers) a -// right-sized make is cheaper than borrowing a full MaxFrameSize buffer, so -// small/average requests never pay the pooled buffer's footprint. -const chunkPoolMinSize = 16 << 10 - -type scheduler struct { - wire WireConn - - mu sync.Mutex - cond *sync.Cond - closed bool - err error - - control []*livekit.AgentHttp_Frame - - queues map[uint32]*streamQueue - rr []uint32 // stream ids with queued data, round-robin order - rrIdx int - - queuedBytes int64 // aggregate queued DATA payload bytes, capped by connBufferBudget - - // chunkPool recycles MaxFrameSize payload buffers for large frames (>= - // chunkPoolMinSize), so a big upload does not allocate one per frame. A - // buffer is drawn in Stream.Write, held in a streamQueue until its frame is - // written, then returned (in writeLoop after WriteFrame, or in dropStream). - chunkPool sync.Pool -} - -type streamQueue struct { - chunks []queuedChunk - eof bool // send an EOF frame after the last chunk drains -} - -// queuedChunk is one pending DATA payload. data is the bytes to send; bp is the -// pooled buffer backing them and is non-nil only for pooled (large) chunks, -// which are the ones recycled once their frame is written. -type queuedChunk struct { - data []byte - bp *[]byte -} - -func newScheduler(wire WireConn, maxFrameSize int) *scheduler { - s := &scheduler{ - wire: wire, - queues: make(map[uint32]*streamQueue), - chunkPool: sync.Pool{New: func() any { - b := make([]byte, maxFrameSize) - return &b - }}, - } - s.cond = sync.NewCond(&s.mu) - go s.writeLoop() - return s -} - -// newChunk copies p into a queued chunk, drawing a pooled MaxFrameSize buffer -// for large payloads and a right-sized make for small ones (so webhook-sized -// requests never borrow a full pooled buffer). The bytes are copied because the -// scheduler is async: the caller's buffer is reused before the frame is written. -func (s *scheduler) newChunk(p []byte) queuedChunk { - if len(p) >= chunkPoolMinSize { - bp := s.chunkPool.Get().(*[]byte) - n := copy(*bp, p) - return queuedChunk{data: (*bp)[:n], bp: bp} - } - data := make([]byte, len(p)) - copy(data, p) - return queuedChunk{data: data} -} - -// recycle returns a chunk's pooled buffer if it has one. -func (s *scheduler) recycle(c queuedChunk) { - if c.bp != nil { - s.chunkPool.Put(c.bp) - } -} - -// enqueueControl never blocks: control frames are small and bounded by the -// number of live streams. -func (s *scheduler) enqueueControl(f *livekit.AgentHttp_Frame) error { - s.mu.Lock() - defer s.mu.Unlock() - if s.closed { - return s.errLocked() - } - s.control = append(s.control, f) - s.cond.Broadcast() - return nil -} - -// enqueueData blocks while the connection buffer budget is exhausted. Stream -// and connection credit are enforced by the caller (Stream.Write); the budget -// is the local memory cap on top of them. -func (s *scheduler) enqueueData(streamID uint32, chunk queuedChunk) error { - s.mu.Lock() - defer s.mu.Unlock() - for !s.closed && s.queuedBytes+int64(len(chunk.data)) > connBufferBudget && s.queuedBytes > 0 { - s.cond.Wait() - } - if s.closed { - s.recycle(chunk) // never enqueued; return its buffer to the pool - return s.errLocked() - } - q := s.queues[streamID] - if q == nil { - q = &streamQueue{} - s.queues[streamID] = q - } - if len(q.chunks) == 0 && !q.eof { - s.rr = append(s.rr, streamID) - } - q.chunks = append(q.chunks, chunk) - s.queuedBytes += int64(len(chunk.data)) - s.cond.Broadcast() - return nil -} - -// enqueueEOF marks the stream's send side done; the EOF frame is emitted after -// its queued chunks, preserving stream order. -func (s *scheduler) enqueueEOF(streamID uint32) error { - s.mu.Lock() - defer s.mu.Unlock() - if s.closed { - return s.errLocked() - } - q := s.queues[streamID] - if q == nil { - q = &streamQueue{} - s.queues[streamID] = q - s.rr = append(s.rr, streamID) - } else if len(q.chunks) == 0 && !q.eof { - s.rr = append(s.rr, streamID) - } - q.eof = true - s.cond.Broadcast() - return nil -} - -// dropStream discards any queued data for a reset stream. A writer blocked on -// the budget can still enqueue one late chunk after the reset; receivers -// tolerate frames for unknown stream ids by protocol contract. -func (s *scheduler) dropStream(streamID uint32) { - s.mu.Lock() - defer s.mu.Unlock() - if q, ok := s.queues[streamID]; ok { - for _, c := range q.chunks { - s.queuedBytes -= int64(len(c.data)) - s.recycle(c) - } - delete(s.queues, streamID) - if i := slices.Index(s.rr, streamID); i != -1 { - s.rr = slices.Delete(s.rr, i, i+1) - if s.rrIdx > i { - s.rrIdx-- - } - } - s.cond.Broadcast() - } -} - -func (s *scheduler) close(err error) { - s.mu.Lock() - if s.closed { - s.mu.Unlock() - return - } - s.closed = true - s.err = err - s.cond.Broadcast() - s.mu.Unlock() - _ = s.wire.Close() -} - -func (s *scheduler) errLocked() error { - if s.err != nil { - return s.err - } - return ErrConnClosed -} - -// next pops the next frame to write: all pending control first, then one DATA -// chunk (or the EOF sentinel) from the round-robin ring. recycle is the pooled -// buffer backing a DATA frame, to be returned once the frame is written (nil -// for control/EOF/unpooled frames). -func (s *scheduler) next() (f *livekit.AgentHttp_Frame, recycle *[]byte, ok bool) { - s.mu.Lock() - defer s.mu.Unlock() - for { - if s.closed { - return nil, nil, false - } - if len(s.control) > 0 { - f = s.control[0] - s.control = s.control[1:] - return f, nil, true - } - if len(s.rr) > 0 { - if s.rrIdx >= len(s.rr) { - s.rrIdx = 0 - } - id := s.rr[s.rrIdx] - q := s.queues[id] - if len(q.chunks) > 0 { - chunk := q.chunks[0] - q.chunks = q.chunks[1:] - s.queuedBytes -= int64(len(chunk.data)) - recycle = chunk.bp - f = &livekit.AgentHttp_Frame{ - StreamId: id, - Message: &livekit.AgentHttp_Frame_Data{Data: chunk.data}, - } - } - if len(q.chunks) == 0 { - if q.eof { - if f == nil { - f = &livekit.AgentHttp_Frame{ - StreamId: id, - Message: &livekit.AgentHttp_Frame_Eof{Eof: &livekit.AgentHttp_HttpStreamEof{}}, - } - } else { - // keep the stream in the ring so the EOF drains next round - s.rrIdx++ - s.cond.Broadcast() - return f, recycle, true - } - } - delete(s.queues, id) - s.rr = slices.Delete(s.rr, s.rrIdx, s.rrIdx+1) - } else { - s.rrIdx++ - } - s.cond.Broadcast() // budget freed - return f, recycle, true - } - s.cond.Wait() - } -} - -func (s *scheduler) writeLoop() { - for { - f, recycle, ok := s.next() - if !ok { - return - } - _ = s.wire.SetWriteDeadline(time.Now().Add(writeStallTimeout)) - err := s.wire.WriteFrame(f) - // WriteFrame marshalled the payload into its own buffer and gorilla does - // not retain it, so the pooled buffer is free to recycle here whether or - // not the write succeeded. - if recycle != nil { - s.chunkPool.Put(recycle) - } - if err != nil { - s.close(err) - return - } - } -} diff --git a/pkg/agent/endpoint/session.go b/pkg/agent/endpoint/session.go new file mode 100644 index 000000000..1ccab8431 --- /dev/null +++ b/pkg/agent/endpoint/session.go @@ -0,0 +1,78 @@ +// Copyright 2026 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 endpoint + +import ( + "context" + "io" +) + +// CurrentProtocol is the data-plane protocol version this server implements, +// negotiated in RegisterWorkerResponse. +const CurrentProtocol uint32 = 1 + +// ResetCode is why a stream was aborted. +type ResetCode int + +const ( + // ResetCancel: the node/client abandoned the exchange (client disconnect, + // timeout, retry elsewhere). + ResetCancel ResetCode = iota + // ResetRefused: the worker aborted before dispatching the request to the + // application, so the request was never applied and is safe to retry. + ResetRefused +) + +// Session is the node's handle to one worker's data plane: a single +// WebTransport (QUIC) session that also carries the worker's control stream. +// The node opens a fresh Stream per HTTP exchange; QUIC provides the +// multiplexing and per-stream flow control, so there is no capsule layer. +type Session interface { + // OpenStream opens a bidirectional stream toward the worker for one HTTP + // exchange. It blocks until a stream can be opened (QUIC stream-limit + // backpressure) or ctx is done. + OpenStream(ctx context.Context) (Stream, error) + // OpenStreams reports the streams currently open on this session - the + // least-outstanding-requests signal for worker selection. + OpenStreams() int + // MaxStreams is the soft cap on concurrent streams, used to derive spare + // capacity for node-selection weighting. + MaxStreams() int + // Close tears the session down. + Close(reason string) +} + +// Stream is one HTTP exchange over a Session: opaque HTTP/1.1 request bytes are +// written toward the worker and opaque response bytes are read back. CloseWrite +// half-closes the request side (the worker then sees EOF); Reset aborts both +// directions. +type Stream interface { + io.Reader + io.Writer + // CloseWrite half-closes the send side once the request is fully written + // (QUIC stream FIN). + CloseWrite() error + // Reset aborts the stream in both directions (QUIC RESET_STREAM / + // STOP_SENDING) with the given code. + Reset(code ResetCode, reason string) + // Close releases the stream after a completed exchange. + Close() error + // BytesRead reports response bytes consumed so far; the retry boundary is + // "no response byte arrived". + BytesRead() int64 + // Refused reports whether the worker aborted the stream with ResetRefused, + // i.e. the request was never dispatched and is safe to retry elsewhere. + Refused() bool +} diff --git a/pkg/agent/endpoint/stream.go b/pkg/agent/endpoint/stream.go deleted file mode 100644 index 8acb5b86f..000000000 --- a/pkg/agent/endpoint/stream.go +++ /dev/null @@ -1,366 +0,0 @@ -// Copyright 2026 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 endpoint - -import ( - "errors" - "fmt" - "io" - "sync" - - "github.com/livekit/protocol/livekit" -) - -// Stream is the server side of one multiplexed stream: one opaque HTTP -// exchange. Write carries request bytes toward the worker under the -// peer-granted stream window AND the wire's shared connection window; Read -// consumes response bytes and replenishes both windows as they are consumed, so -// a slow HTTP client suspends the worker's send without touching sibling -// streams. -type Stream struct { - id uint32 - conn *DataConn - - mu sync.Mutex - cond *sync.Cond - - // send side (request bytes) - sendCredit int64 - writeClosed bool - - // recv side (response bytes) - recvChunks [][]byte - recvOffset int // read offset into recvChunks[0] - recvWindow int64 // bytes the peer may still send; enforced, not advisory - recvUnacked int64 // consumed bytes not yet credited back - recvEOF bool - refused bool - err error - closed bool - bytesRead int64 - bytesWritten int64 -} - -func newStream(id uint32, conn *DataConn, window int64) *Stream { - s := &Stream{ - id: id, - conn: conn, - sendCredit: window, - recvWindow: window, - } - s.cond = sync.NewCond(&s.mu) - return s -} - -func (s *Stream) ID() uint32 { return s.id } - -// BytesRead reports response bytes consumed so far; the retry boundary is -// "no response byte arrived". -func (s *Stream) BytesRead() int64 { - s.mu.Lock() - defer s.mu.Unlock() - return s.bytesRead -} - -// Refused reports whether the worker reset the stream with HSR_REFUSED, i.e. -// the request was never dispatched and is safe to retry elsewhere. -func (s *Stream) Refused() bool { - s.mu.Lock() - defer s.mu.Unlock() - return s.refused -} - -// Write sends request bytes toward the worker, blocking on the stream window, -// the wire's shared connection window and the local buffer budget. -func (s *Stream) Write(p []byte) (int, error) { - cancelled := func() bool { - s.mu.Lock() - defer s.mu.Unlock() - return s.closed || s.err != nil || s.writeClosed - } - total := 0 - for len(p) > 0 { - s.mu.Lock() - for s.sendCredit <= 0 && s.err == nil && !s.closed && !s.writeClosed { - s.cond.Wait() - } - if err := s.writeErrLocked(); err != nil { - s.mu.Unlock() - return total, err - } - want := int64(len(p)) - if want > s.sendCredit { - want = s.sendCredit - } - if max := int64(s.conn.params.MaxFrameSize); want > max { - want = max - } - s.mu.Unlock() - - // the shared window is reserved after the stream window and outside - // s.mu (the cancelled callback re-locks it) - n, err := s.conn.reserveConnSend(want, cancelled) - if err != nil { - return total, err - } - - s.mu.Lock() - if err := s.writeErrLocked(); err != nil { - s.mu.Unlock() - s.conn.returnConnSend(n) - return total, err - } - var overshoot int64 - if n > s.sendCredit { - // single-writer streams never hit this; stay safe regardless - overshoot = n - s.sendCredit - n = s.sendCredit - } - s.sendCredit -= n - s.bytesWritten += n - s.mu.Unlock() - s.conn.returnConnSend(overshoot) - if n == 0 { - // a concurrent writer drained the stream window between the peek - // and the reservation: never emit an empty frame - continue - } - - if err := s.conn.sched.enqueueData(s.id, s.conn.sched.newChunk(p[:n])); err != nil { - return total, err - } - s.conn.noteActivity(n) - p = p[n:] - total += int(n) - } - return total, nil -} - -// ReadFrom streams r to the worker in MaxFrameSize frames. Without it, -// http.Request.Write wraps this stream in a 4KiB bufio.Writer (Stream is not an -// io.ByteWriter) and a large body would be refragmented into ~4KiB frames; -// bufio.Writer.ReadFrom delegates here once its buffer drains, so the body -// flows at the negotiated frame size instead. Implements io.ReaderFrom. Write -// still handles headers and small bodies (bufio never delegates for those). -func (s *Stream) ReadFrom(r io.Reader) (int64, error) { - buf := make([]byte, s.conn.params.MaxFrameSize) - var total int64 - for { - n, rerr := r.Read(buf) - if n > 0 { - w, werr := s.Write(buf[:n]) - total += int64(w) - if werr != nil { - return total, werr - } - } - if rerr != nil { - if rerr == io.EOF { - return total, nil - } - return total, rerr - } - } -} - -func (s *Stream) writeErrLocked() error { - switch { - case s.err != nil: - return s.err - case s.closed: - return ErrStreamClosed - case s.writeClosed: - return fmt.Errorf("%w: write after CloseWrite", ErrStreamClosed) - } - return nil -} - -// CloseWrite half-closes the send side; the worker sees EOF after the queued -// request bytes drain. -func (s *Stream) CloseWrite() error { - s.mu.Lock() - if s.writeClosed || s.closed { - s.mu.Unlock() - return nil - } - s.writeClosed = true - s.mu.Unlock() - return s.conn.sched.enqueueEOF(s.id) -} - -// Read consumes response bytes. Stream and connection credit are granted back -// to the worker on consumption (threshold acking at half the respective -// window) - receipt alone never replenishes either window. -func (s *Stream) Read(p []byte) (int, error) { - s.mu.Lock() - for len(s.recvChunks) == 0 && !s.recvEOF && s.err == nil && !s.closed { - s.cond.Wait() - } - if len(s.recvChunks) == 0 { - err := s.err - if err == nil && s.recvEOF { - err = io.EOF - } else if err == nil { - err = ErrStreamClosed - } - s.mu.Unlock() - return 0, err - } - - chunk := s.recvChunks[0] - n := copy(p, chunk[s.recvOffset:]) - s.recvOffset += n - if s.recvOffset == len(chunk) { - s.recvChunks = s.recvChunks[1:] - s.recvOffset = 0 - } - s.bytesRead += int64(n) - s.recvUnacked += int64(n) - - var credit int64 - if s.recvUnacked >= int64(s.conn.params.CreditWindow)/2 { - credit = s.recvUnacked - s.recvUnacked = 0 - s.recvWindow += credit - } - s.mu.Unlock() - - if credit > 0 { - _ = s.conn.sched.enqueueControl(&livekit.AgentHttp_Frame{ - StreamId: s.id, - Message: &livekit.AgentHttp_Frame_Credit{Credit: uint32(credit)}, - }) - } - s.conn.connConsumed(int64(n)) - s.conn.noteActivity(int64(n)) - return n, nil -} - -// drainUnreadLocked releases response bytes that arrived but will never be -// Read, so a torn-down stream cannot leak the wire's connection window. -// Callers hold s.mu. -func (s *Stream) drainUnreadLocked() int64 { - var pending int64 - for i, c := range s.recvChunks { - pending += int64(len(c)) - if i == 0 { - pending -= int64(s.recvOffset) - } - } - s.recvChunks = nil - s.recvOffset = 0 - return pending -} - -// Reset aborts the stream in both directions. -func (s *Stream) Reset(code livekit.AgentHttp_HttpStreamResetCode, reason string) { - s.mu.Lock() - if s.closed { - s.mu.Unlock() - return - } - s.closed = true - if s.err == nil { - s.err = ErrStreamClosed - } - pending := s.drainUnreadLocked() - s.cond.Broadcast() - s.mu.Unlock() - - s.conn.connConsumed(pending) - s.conn.wakeSendWaiters() - s.conn.sched.dropStream(s.id) - _ = s.conn.sched.enqueueControl(&livekit.AgentHttp_Frame{ - StreamId: s.id, - Message: &livekit.AgentHttp_Frame_Reset_{ - Reset_: &livekit.AgentHttp_HttpStreamReset{Code: code, Error: reason}, - }, - }) - s.conn.removeStream(s.id) -} - -// Close releases the stream after a completed exchange. -func (s *Stream) Close() error { - s.mu.Lock() - if s.closed { - s.mu.Unlock() - return nil - } - s.closed = true - if s.err == nil { - s.err = ErrStreamClosed - } - pending := s.drainUnreadLocked() - s.cond.Broadcast() - s.mu.Unlock() - - s.conn.connConsumed(pending) - s.conn.wakeSendWaiters() - s.conn.sched.dropStream(s.id) - s.conn.removeStream(s.id) - return nil -} - -// --- frame delivery from the wire read loop --- - -func (s *Stream) onData(payload []byte) error { - s.mu.Lock() - defer s.mu.Unlock() - if s.closed { - // stream was reset locally; the conn window is released by the caller - return errStreamGone - } - if int64(len(payload)) > s.recvWindow { - return fmt.Errorf("%w: peer exceeded credit window on stream %d", ErrProtocol, s.id) - } - s.recvWindow -= int64(len(payload)) - s.recvChunks = append(s.recvChunks, payload) - s.cond.Broadcast() - return nil -} - -// errStreamGone tells the read loop the payload was not adopted and its -// connection-window share must be released immediately. -var errStreamGone = errors.New("stream gone") - -func (s *Stream) onEOF() { - s.mu.Lock() - s.recvEOF = true - s.cond.Broadcast() - s.mu.Unlock() -} - -func (s *Stream) onReset(code livekit.AgentHttp_HttpStreamResetCode, reason string) { - s.mu.Lock() - if s.err == nil { - if code == livekit.AgentHttp_HSR_REFUSED { - s.refused = true - s.err = ErrStreamRefused - } else { - s.err = fmt.Errorf("stream reset by worker: %s (%s)", code, reason) - } - } - pending := s.drainUnreadLocked() - s.cond.Broadcast() - s.mu.Unlock() - s.conn.connConsumed(pending) -} - -func (s *Stream) onCredit(increment uint32) { - s.mu.Lock() - s.sendCredit += int64(increment) - s.cond.Broadcast() - s.mu.Unlock() -} diff --git a/pkg/agent/endpoint/template.go b/pkg/agent/endpoint/template.go index 62008de02..974e5d044 100644 --- a/pkg/agent/endpoint/template.go +++ b/pkg/agent/endpoint/template.go @@ -27,8 +27,7 @@ import ( // starlette's semantics would make the server route requests the worker-side // router then refuses. type Template struct { - raw string - re *regexp.Regexp + re *regexp.Regexp } // convertor patterns copied verbatim from starlette's convertors.py. The uuid @@ -86,13 +85,9 @@ func ParseTemplate(path string) (*Template, error) { if err != nil { return nil, fmt.Errorf("invalid path template %q: %w", path, err) } - return &Template{raw: path, re: re}, nil + return &Template{re: re}, nil } func (t *Template) Match(path string) bool { return t.re.MatchString(path) } - -func (t *Template) String() string { - return t.raw -} diff --git a/pkg/agent/endpoint/template_test.go b/pkg/agent/endpoint/template_test.go index a12d0404a..24c3d6246 100644 --- a/pkg/agent/endpoint/template_test.go +++ b/pkg/agent/endpoint/template_test.go @@ -101,6 +101,15 @@ func TestManifestValidation(t *testing.T) { _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"get"}, false)}) require.Error(t, err, "lowercase method") + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"GTE"}, false)}) + require.Error(t, err, "unsupported method rejected") + + // the full FastAPI/starlette verb set is accepted + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ + ep("/x", []string{"GET", "HEAD", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "TRACE"}, false), + }) + require.NoError(t, err, "all FastAPI verbs accepted") + _, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ {Path: "/ws", Kind: livekit.AgentHttp_AEK_TEXT, Methods: []string{"GET"}}, }) diff --git a/pkg/agent/endpoint/webtransport.go b/pkg/agent/endpoint/webtransport.go new file mode 100644 index 000000000..ceb8a1757 --- /dev/null +++ b/pkg/agent/endpoint/webtransport.go @@ -0,0 +1,160 @@ +// Copyright 2026 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 endpoint + +import ( + "context" + "errors" + "sync" + "sync/atomic" + + "github.com/quic-go/webtransport-go" +) + +// Stream reset codes on the wire, shared by the server and the Go conformance +// worker so the two never drift. QUIC carries the numeric code on +// RESET_STREAM/STOP_SENDING; a worker uses Refused to mark a request it never +// dispatched (the front then treats it as safe to retry elsewhere). +const ( + StreamCodeCancel webtransport.StreamErrorCode = 0 + StreamCodeRefused webtransport.StreamErrorCode = 2 +) + +func resetToWTCode(c ResetCode) webtransport.StreamErrorCode { + if c == ResetRefused { + return StreamCodeRefused + } + return StreamCodeCancel +} + +// wtSession adapts a WebTransport session to Session: the node opens one stream +// per HTTP exchange and QUIC multiplexes them. +type wtSession struct { + sess *webtransport.Session + maxStreams int + open atomic.Int64 +} + +// NewWebTransportSession wraps a live WebTransport session as a data-plane +// Session. maxStreams is the soft concurrency cap used only for capacity +// weighting; QUIC's own stream limit is the hard bound. +func NewWebTransportSession(sess *webtransport.Session, maxStreams int) Session { + return &wtSession{sess: sess, maxStreams: maxStreams} +} + +func (s *wtSession) OpenStream(ctx context.Context) (Stream, error) { + qs, err := s.sess.OpenStreamSync(ctx) + if err != nil { + return nil, err + } + s.open.Add(1) + return &wtStream{sess: s, qs: qs}, nil +} + +func (s *wtSession) OpenStreams() int { return int(s.open.Load()) } +func (s *wtSession) MaxStreams() int { return s.maxStreams } +func (s *wtSession) Close(reason string) { _ = s.sess.CloseWithError(0, reason) } + +// wtStream is one HTTP exchange over a WebTransport bidi stream. +type wtStream struct { + sess *wtSession + qs *webtransport.Stream + + mu sync.Mutex + bytesRead int64 + refused bool + done bool + sendClosed bool // request side FIN'd (CloseWrite) or reset +} + +func (s *wtStream) Read(p []byte) (int, error) { + n, err := s.qs.Read(p) + if n > 0 { + s.mu.Lock() + s.bytesRead += int64(n) + s.mu.Unlock() + } + if err != nil { + var se *webtransport.StreamError + if errors.As(err, &se) && se.ErrorCode == StreamCodeRefused { + s.mu.Lock() + s.refused = true + s.mu.Unlock() + } + } + return n, err +} + +func (s *wtStream) Write(p []byte) (int, error) { return s.qs.Write(p) } + +// CloseWrite sends STREAM FIN on the request side; the worker then reads EOF. +func (s *wtStream) CloseWrite() error { + s.mu.Lock() + s.sendClosed = true + s.mu.Unlock() + return s.qs.Close() +} + +func (s *wtStream) Reset(code ResetCode, _ string) { + c := resetToWTCode(code) + s.mu.Lock() + s.sendClosed = true + s.mu.Unlock() + s.qs.CancelWrite(c) + s.qs.CancelRead(c) + s.release() +} + +func (s *wtStream) Close() error { + // If the request side was never cleanly FIN'd (CloseWrite) or reset, a writer + // goroutine may still be blocked in Write because the worker stopped reading: + // cancel the send side to unblock it (else the goroutine and QUIC stream leak). + // After a clean FIN we must NOT reset - that would turn a completed request + // into an abort on the wire. + s.mu.Lock() + cancelWrite := !s.sendClosed + s.sendClosed = true + s.mu.Unlock() + if cancelWrite { + s.qs.CancelWrite(StreamCodeCancel) + } + s.qs.CancelRead(StreamCodeCancel) + s.release() + return nil +} + +// release decrements the session's open-stream count exactly once. +func (s *wtStream) release() { + s.mu.Lock() + if s.done { + s.mu.Unlock() + return + } + s.done = true + s.mu.Unlock() + s.sess.open.Add(-1) +} + +func (s *wtStream) BytesRead() int64 { + s.mu.Lock() + defer s.mu.Unlock() + return s.bytesRead +} + +func (s *wtStream) Refused() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.refused +} diff --git a/pkg/agent/endpoint/webtransport_test.go b/pkg/agent/endpoint/webtransport_test.go new file mode 100644 index 000000000..6744e7c82 --- /dev/null +++ b/pkg/agent/endpoint/webtransport_test.go @@ -0,0 +1,184 @@ +// Copyright 2026 LiveKit, Inc. + +package endpoint_test + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "io" + "math/big" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/quic-go/quic-go/http3" + "github.com/quic-go/webtransport-go" + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/agent/endpoint/conformance" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +func selfSignedTLS(t *testing.T) *tls.Config { + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "localhost"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + require.NoError(t, err) + return &tls.Config{ + Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}}, + NextProtos: []string{http3.NextProtoH3}, + } +} + +// startWTServer runs a WebTransport /agent server that registers each session's +// worker (read off the control stream) into reg and keeps the session alive. +func startWTServer(t *testing.T, reg *endpoint.Registry) string { + udp, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + require.NoError(t, err) + + mux := http.NewServeMux() + srv := &webtransport.Server{H3: &http3.Server{TLSConfig: selfSignedTLS(t), Handler: mux}} + mux.HandleFunc("/agent", func(w http.ResponseWriter, r *http.Request) { + sess, err := srv.Upgrade(w, r) + if err != nil { + return + } + go handleSession(reg, sess) + }) + go func() { _ = srv.Serve(udp) }() + t.Cleanup(func() { _ = srv.Close(); _ = udp.Close() }) + return "https://" + udp.LocalAddr().String() + "/agent" +} + +func handleSession(reg *endpoint.Registry, sess *webtransport.Session) { + ctx := sess.Context() + control, err := sess.AcceptStream(ctx) + if err != nil { + return + } + var msg livekit.WorkerMessage + if err := endpoint.ReadControlMessage(control, &msg); err != nil { + return + } + rw := msg.GetRegister() + if rw == nil { + return + } + manifest, err := endpoint.ParseManifest(rw.GetEndpoints()) + if err != nil { + return + } + registration := &endpoint.Registration{ + WorkerID: rw.GetInstanceId(), + APIKey: "test", + AgentName: rw.GetAgentName(), + Deployment: rw.GetDeployment(), + Manifest: manifest, + Endpoints: rw.GetEndpoints(), + } + registration.SetSession(endpoint.NewWebTransportSession(sess, endpoint.DefaultMaxStreams)) + _ = reg.Register(registration) + _ = endpoint.WriteControlMessage(control, &livekit.ServerMessage{ + Message: &livekit.ServerMessage_Register{ + Register: &livekit.RegisterWorkerResponse{WorkerId: registration.WorkerID}, + }, + }) + // keep the session alive; drain further control messages until it dies + for { + var m livekit.WorkerMessage + if err := endpoint.ReadControlMessage(control, &m); err != nil { + return + } + } +} + +func TestWebTransportEndpointRoundTrip(t *testing.T) { + // the worker's local app + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/hello": + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"ok":true}`) + case "/echo": + body, _ := io.ReadAll(r.Body) + _, _ = w.Write(body) + default: + http.NotFound(w, r) + } + })) + defer target.Close() + + reg := endpoint.NewRegistry() + base := startWTServer(t, reg) + + front := endpoint.NewFront(reg, func(*http.Request) (string, bool) { return "", false }, logger.GetLogger()). + WithSingleKeyFallback() + ts := httptest.NewServer(front) + defer ts.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + w := conformance.New(conformance.Config{ + ServerURL: base, + APIKey: "APIkey", + APISecret: "secret-that-is-long-enough-to-sign", + AgentName: "myagent", + Deployment: "production", + TargetAddr: strings.TrimPrefix(target.URL, "http://"), + Insecure: true, + Endpoints: []*livekit.AgentHttp_AgentEndpoint{ + {Path: "/hello", Methods: []string{"GET"}, Public: true}, + {Path: "/echo", Methods: []string{"POST"}, Public: true}, + }, + }) + require.NoError(t, w.Start(ctx)) + t.Cleanup(w.Close) + require.NoError(t, w.WaitRegistered(ctx)) + + agentBase := ts.URL + "/agents/myagent/production" + + t.Run("GET json", func(t *testing.T) { + resp, err := http.Get(agentBase + "/hello") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.JSONEq(t, `{"ok":true}`, string(body)) + }) + + t.Run("POST echo", func(t *testing.T) { + payload := strings.Repeat("x", 64<<10) + resp, err := http.Post(agentBase+"/echo", "application/octet-stream", strings.NewReader(payload)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.Equal(t, len(payload), len(body)) + require.Equal(t, payload, string(body)) + }) + + t.Run("unknown path 404", func(t *testing.T) { + resp, err := http.Get(agentBase + "/nope") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusNotFound, resp.StatusCode) + }) +} diff --git a/pkg/agent/endpoint/wire.go b/pkg/agent/endpoint/wire.go deleted file mode 100644 index f6ee99edb..000000000 --- a/pkg/agent/endpoint/wire.go +++ /dev/null @@ -1,122 +0,0 @@ -// Copyright 2026 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 endpoint implements the server side of the agent HTTP endpoints data -// plane: worker-dialed wires carrying multiplexed streams, each stream one -// opaque HTTP exchange, flow-controlled per stream and per wire with a -// prioritized per-wire write scheduler. Wires speak AgentHttp.Frame -// exclusively: one frame per websocket binary message. -package endpoint - -import ( - "errors" - "time" - - "github.com/livekit/protocol/livekit" -) - -// CurrentProtocol is the data-plane protocol version negotiated in -// RegisterWorkerResponse.endpoint_settings. -const CurrentProtocol = 1 - -var ( - ErrConnClosed = errors.New("data connection closed") - ErrStreamClosed = errors.New("stream closed") - ErrStreamRefused = errors.New("stream refused by worker") - ErrTooManyStreams = errors.New("too many open streams on connection") - ErrProtocol = errors.New("data plane protocol violation") -) - -// WireConn is the transport under one data wire: whole binary websocket -// messages in, one Frame each. The write deadline is what turns a peer that -// stopped reading into a dead connection instead of a stuck writer goroutine. -type WireConn interface { - WriteFrame(f *livekit.AgentHttp_Frame) error - ReadFrame() (*livekit.AgentHttp_Frame, error) - SetWriteDeadline(t time.Time) error - SetReadDeadline(t time.Time) error - Close() error -} - -// Settings are the registration-level parameters negotiated on the control -// connection. Wire-level parameters ride each wire's attach response instead: -// wires may be adopted by any node, and the adopting node's parameters govern. -type Settings struct { - Protocol uint32 - AttachToken string - DataConnCount uint32 -} - -// WireParams are one wire's flow-control parameters, chosen by the node that -// adopted it and announced in AttachDataConnectionResponse. -type WireParams struct { - CreditWindow uint32 - ConnectionWindow uint32 - MaxFrameSize uint32 - MaxStreamsPerConn uint32 -} - -const ( - DefaultDataConnCount = 8 - DefaultCreditWindow = 1 << 20 // 1MiB per stream - DefaultConnectionWindow = 4 << 20 // shared by all streams on a wire - DefaultMaxFrameSize = 64 << 10 - DefaultMaxStreamsPerConn = 64 // total open; parked streams included - - // writeStallTimeout bounds a single wire write; a peer that stops reading - // kills the connection instead of freezing its sibling streams. - writeStallTimeout = 30 * time.Second - - // connBufferBudget caps the aggregate queued-but-unwritten bytes per - // connection so many streams times a full window cannot pin unbounded memory. - connBufferBudget = 4 << 20 -) - -// WithDefaults fills zero fields with the package defaults. -func (p WireParams) WithDefaults() WireParams { - if p.CreditWindow == 0 { - p.CreditWindow = DefaultCreditWindow - } - if p.ConnectionWindow == 0 { - p.ConnectionWindow = DefaultConnectionWindow - } - if p.MaxFrameSize == 0 { - p.MaxFrameSize = DefaultMaxFrameSize - } - if p.MaxStreamsPerConn == 0 { - p.MaxStreamsPerConn = DefaultMaxStreamsPerConn - } - return p -} - -// frameOverhead bounds the AgentHttp.Frame envelope around a payload (stream id, -// oneof tag, length varints); a data frame on the wire is at most -// MaxFrameSize + frameOverhead. -const frameOverhead = 1024 - -// MaxMessageSize is the largest websocket message a wire may legitimately carry -// (a full data frame). A reader caps its read limit at this so an oversized -// frame is rejected before it can exhaust memory, even when no signalling limit -// is configured. -func (p WireParams) MaxMessageSize() int64 { - return int64(p.MaxFrameSize) + frameOverhead -} - -func (s Settings) Proto() *livekit.AgentHttp_AgentEndpointSettings { - return &livekit.AgentHttp_AgentEndpointSettings{ - Protocol: s.Protocol, - AttachToken: s.AttachToken, - DataConnectionCount: s.DataConnCount, - } -} diff --git a/pkg/agent/worker.go b/pkg/agent/worker.go index ab8f50d66..00116528c 100644 --- a/pkg/agent/worker.go +++ b/pkg/agent/worker.go @@ -159,7 +159,6 @@ type WorkerRegistration struct { // agent HTTP endpoints data plane Endpoints []*livekit.AgentHttp_AgentEndpoint InstanceID string - EndpointProtocol uint32 EndpointSettings *livekit.AgentHttp_AgentEndpointSettings } @@ -243,12 +242,19 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er h.registration.Deployment = req.GetDeployment() h.registration.Endpoints = req.GetEndpoints() h.registration.InstanceID = req.GetInstanceId() - h.registration.EndpointProtocol = req.GetEndpointProtocol() if len(req.GetEndpoints()) > 0 && req.GetEndpointProtocol() > 0 { if h.endpointSettings == nil { return errors.New("agent HTTP endpoints are not supported by this server") } + // endpoints are addressed at /agents/{agent_name}/{deployment}/...; a + // worker serving them needs a non-empty, URL-safe agent name + if req.GetAgentName() == "" { + return errors.New("agent HTTP endpoints require an agent name") + } + if err := protoagent.ValidateAgentName(req.GetAgentName()); err != nil { + return err + } settings, err := h.endpointSettings(req) if err != nil { return err diff --git a/pkg/service/agentendpoint_test.go b/pkg/service/agentendpoint_test.go index b6d9212a9..b04947082 100644 --- a/pkg/service/agentendpoint_test.go +++ b/pkg/service/agentendpoint_test.go @@ -18,10 +18,17 @@ import ( "bufio" "bytes" "context" + "crypto/ecdsa" + "crypto/elliptic" "crypto/rand" "crypto/sha256" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" "fmt" "io" + "math/big" + "net" "net/http" "net/http/httptest" "strings" @@ -29,6 +36,7 @@ import ( "testing" "time" + "github.com/quic-go/quic-go/http3" "github.com/stretchr/testify/require" "github.com/livekit/livekit-server/pkg/agent" @@ -48,9 +56,30 @@ const ( ) type endpointStack struct { - t *testing.T - ts *httptest.Server - svc *service.AgentService + t *testing.T + ts *httptest.Server + svc *service.AgentService + wtURL string // https://host:port/agent (WebTransport control+data) +} + +// selfSignedTLS mints an in-memory cert for 127.0.0.1 with the h3 ALPN, for the +// node's WebTransport listener. +func selfSignedTLS(t *testing.T) *tls.Config { + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "localhost"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + require.NoError(t, err) + return &tls.Config{ + Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}}, + NextProtos: []string{http3.NextProtoH3}, + } } func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpointStack { @@ -70,10 +99,9 @@ func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpoin ) require.NoError(t, err) + // public front (client-facing HTTP/1.1), unchanged mux := http.NewServeMux() - mux.Handle("/agent", svc) mux.Handle(endpoint.PathPrefix, svc.EndpointFront()) - authMW := service.NewAPIKeyAuthMiddleware(keyProvider) ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { authMW.ServeHTTP(w, r, mux.ServeHTTP) @@ -81,26 +109,32 @@ func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpoin t.Cleanup(ts.Close) t.Cleanup(func() { svc.DrainConnections(time.Millisecond, true) }) - return &endpointStack{t: t, ts: ts, svc: svc} -} + // workers connect over WebTransport (control + data on one session) + udp, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}) + require.NoError(t, err) + wt := service.NewAgentWebTransportServer(svc, keyProvider, selfSignedTLS(t)) + go func() { _ = wt.Serve(udp) }() + t.Cleanup(func() { _ = wt.Close(); _ = udp.Close() }) + wtURL := "https://" + udp.LocalAddr().String() + "/agent" -func (s *endpointStack) wsURL() string { - return "ws" + strings.TrimPrefix(s.ts.URL, "http") + "/agent" + return &endpointStack{t: t, ts: ts, svc: svc, wtURL: wtURL} } func (s *endpointStack) startWorker(target string, deployment string, endpoints []*livekit.AgentHttp_AgentEndpoint) *conformance.Worker { w := conformance.New(conformance.Config{ - ServerURL: s.wsURL(), + ServerURL: s.wtURL, APIKey: testKey, APISecret: testSecret, AgentName: "test-agent", Deployment: deployment, Endpoints: endpoints, TargetAddr: strings.TrimPrefix(target, "http://"), + Insecure: true, }) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() require.NoError(s.t, w.Start(ctx)) + require.NoError(s.t, w.WaitRegistered(ctx)) s.t.Cleanup(w.Close) return w } @@ -164,7 +198,7 @@ func TestAgentEndpointsCorrectnessGate(t *testing.T) { httpEP("/sse", []string{"GET"}, true), }) - base := stack.ts.URL + "/agents/production" + base := stack.ts.URL + "/agents/test-agent/production" t.Run("json round trip", func(t *testing.T) { resp, err := http.Get(base + "/json") @@ -250,7 +284,7 @@ func TestAgentEndpointsStatusMapping(t *testing.T) { httpEP("/private", []string{"GET"}, false), }) - base := stack.ts.URL + "/agents/production" + base := stack.ts.URL + "/agents/test-agent/production" t.Run("404 unknown path", func(t *testing.T) { resp, _ := http.Get(base + "/nope") @@ -276,7 +310,7 @@ func TestAgentEndpointsStatusMapping(t *testing.T) { require.Equal(t, 200, resp.StatusCode) }) t.Run("503 unknown deployment", func(t *testing.T) { - resp, _ := http.Get(stack.ts.URL + "/agents/staging/hook") + resp, _ := http.Get(stack.ts.URL + "/agents/test-agent/staging/hook") resp.Body.Close() require.Equal(t, 503, resp.StatusCode) }) @@ -288,15 +322,64 @@ func TestAgentEndpointsStatusMapping(t *testing.T) { require.NoError(t, err) resp.Body.Close() require.Equal(t, 307, resp.StatusCode) - require.Equal(t, "/agents/production/hook", resp.Header.Get("Location")) + require.Equal(t, "/agents/test-agent/production/hook", resp.Header.Get("Location")) + }) +} + +// two workers of the same agent+deployment serve the SAME path with DIFFERENT +// methods; each request must route to the worker whose manifest serves its +// method (method-aware candidate selection: FULL = path+method). +func TestAgentEndpointsMethodAcrossWorkers(t *testing.T) { + postMux := http.NewServeMux() + postMux.HandleFunc("POST /thing", func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "post-worker") + }) + postApp := newTargetApp(t, postMux) + + getMux := http.NewServeMux() + getMux.HandleFunc("GET /thing", func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "get-worker") + }) + getApp := newTargetApp(t, getMux) + + stack := newEndpointStack(t, agent.EndpointsConfig{}) + stack.startWorker(postApp.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/thing", []string{"POST"}, true), + }) + stack.startWorker(getApp.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ + httpEP("/thing", []string{"GET"}, true), + }) + base := stack.ts.URL + "/agents/test-agent/production" + + t.Run("GET routes to the GET worker", func(t *testing.T) { + resp, err := http.Get(base + "/thing") + require.NoError(t, err) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + require.Equal(t, "get-worker", string(body)) + }) + t.Run("POST routes to the POST worker", func(t *testing.T) { + resp, err := http.Post(base+"/thing", "text/plain", nil) + require.NoError(t, err) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + require.Equal(t, "post-worker", string(body)) + }) + t.Run("unserved method is 405", func(t *testing.T) { + req, _ := http.NewRequest(http.MethodDelete, base+"/thing", nil) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + require.Equal(t, 405, resp.StatusCode) }) } func TestAgentEndpointsHOL(t *testing.T) { - // one data conn forces every stream onto the same socket: the credit windows - // and the write scheduler are the only things standing between a stalled - // reader and its siblings - blocked := make(chan struct{}) + // every exchange rides its own QUIC stream on the one session: QUIC's + // per-stream flow control is what keeps a stalled reader from blocking its + // siblings (there is no shared socket, no credit window, no write scheduler). mux := http.NewServeMux() mux.HandleFunc("GET /drip", func(w http.ResponseWriter, r *http.Request) { f := w.(http.Flusher) @@ -318,12 +401,12 @@ func TestAgentEndpointsHOL(t *testing.T) { }) app := newTargetApp(t, mux) - stack := newEndpointStack(t, agent.EndpointsConfig{DataConnCount: 1}) + stack := newEndpointStack(t, agent.EndpointsConfig{}) stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{ httpEP("/drip", []string{"GET"}, true), httpEP("/quick", []string{"GET"}, true), }) - base := stack.ts.URL + "/agents/production" + base := stack.ts.URL + "/agents/test-agent/production" // a stalled client: open /drip, read a little, then stop reading entirely resp, err := http.Get(base + "/drip") @@ -332,8 +415,7 @@ func TestAgentEndpointsHOL(t *testing.T) { small := make([]byte, 1024) _, err = io.ReadFull(resp.Body, small) require.NoError(t, err) - // do not read further; the stream's credit window fills and stays full - close(blocked) + // do not read further; the drip stream's QUIC flow-control window fills time.Sleep(500 * time.Millisecond) // sibling streams on the SAME connection must proceed at full speed @@ -365,13 +447,14 @@ func TestAgentEndpointsRetrySafety(t *testing.T) { eps := []*livekit.AgentHttp_AgentEndpoint{httpEP("/json", []string{"GET"}, true)} broken := conformance.New(conformance.Config{ - ServerURL: stack.wsURL(), APIKey: testKey, APISecret: testSecret, + ServerURL: stack.wtURL, APIKey: testKey, APISecret: testSecret, AgentName: "test-agent", Deployment: "production", - Endpoints: eps, TargetAddr: deadTarget, + Endpoints: eps, TargetAddr: deadTarget, Insecure: true, }) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() require.NoError(t, broken.Start(ctx)) + require.NoError(t, broken.WaitRegistered(ctx)) t.Cleanup(broken.Close) stack.startWorker(app.URL, "production", eps) @@ -379,7 +462,7 @@ func TestAgentEndpointsRetrySafety(t *testing.T) { // run enough requests that both workers get picked first sometimes okCount := 0 for i := 0; i < 12; i++ { - resp, err := http.Get(stack.ts.URL + "/agents/production/json") + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production/json") require.NoError(t, err) body, _ := io.ReadAll(resp.Body) resp.Body.Close() @@ -409,7 +492,7 @@ func TestAgentEndpointsTruncationAborts(t *testing.T) { httpEP("/partial", []string{"GET"}, true), }) - resp, err := http.Get(stack.ts.URL + "/agents/production/partial") + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production/partial") require.NoError(t, err) defer resp.Body.Close() @@ -437,7 +520,7 @@ func TestAgentEndpointsNoLocalListenerContract(t *testing.T) { }) for _, path := range []string{"/", "/worker", "/undeclared"} { - resp, err := http.Get(stack.ts.URL + "/agents/production" + path) + resp, err := http.Get(stack.ts.URL + "/agents/test-agent/production" + path) require.NoError(t, err) resp.Body.Close() require.Equal(t, 404, resp.StatusCode, path) diff --git a/pkg/service/agentservice.go b/pkg/service/agentservice.go index 15859b52c..64fdb1a57 100644 --- a/pkg/service/agentservice.go +++ b/pkg/service/agentservice.go @@ -17,7 +17,6 @@ package service import ( "context" "errors" - "io" "math/rand" "net/http" "slices" @@ -27,7 +26,6 @@ import ( "time" "github.com/gorilla/websocket" - "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/emptypb" "github.com/livekit/livekit-server/pkg/agent" @@ -239,26 +237,12 @@ func (s *AgentService) EndpointFront() http.Handler { } 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 { - 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. The read is bounded by the data wire's - // own max message size (a full frame), independent of the signalling - // limit, so a peer cannot OOM the server with an oversized frame. The - // connecting grant's api key binds the attach to its own project (a - // leaked attach token alone cannot bind a data conn to another - // project's worker). - params := s.wireParams() - conn.SetReadLimit(params.MaxMessageSize()) - HandleEndpointAttach(s.endpointRegistry, NewEndpointWireConn(conn), params, nil, GetAPIKey(r.Context())) - return - } - // control/signal conn: bound a single signalling frame (0 = operator // disabled) before it is buffered; decompressed size is bounded in - // WSSignalConnection. + // WSSignalConnection. HTTP endpoints are served over a worker's + // WebTransport session (see the WebTransport listener), not this + // WebSocket, which now carries control/job-dispatch traffic only. if s.signalMessageSizeLimit > 0 { conn.SetReadLimit(s.signalMessageSizeLimit) } @@ -268,193 +252,6 @@ func (s *AgentService) ServeHTTP(w http.ResponseWriter, r *http.Request) { } } -// 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 - // scratch holds the last frame's marshalled bytes, reused across writes so a - // high-rate stream does not allocate a buffer per frame. Guarded by writeMu. - scratch []byte - // readBuf holds the last frame's raw bytes, reused across reads. Only the - // single reader (the conn's readLoop) touches it, and proto.Unmarshal copies - // the payload out, so it is safe to recycle each frame with no lock. - readBuf []byte -} - -// 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 { - c.writeMu.Lock() - defer c.writeMu.Unlock() - // marshal into the reused buffer and write it before returning; gorilla's - // WriteMessage does not retain the payload, so the next frame can reuse it. - b, err := proto.MarshalOptions{}.MarshalAppend(c.scratch[:0], f) - if err != nil { - return err - } - c.scratch = b - return c.ws.WriteMessage(websocket.BinaryMessage, b) -} - -func (c *endpointWireConn) ReadFrame() (*livekit.AgentHttp_Frame, error) { - for { - _ = c.ws.SetReadDeadline(time.Now().Add(endpointWireIdleTimeout)) - mt, r, err := c.ws.NextReader() - if err != nil { - return nil, err - } - if mt != websocket.BinaryMessage { - continue // gorilla discards this reader on the next NextReader - } - // read into the reused buffer instead of letting ReadMessage allocate one - // per frame; proto.Unmarshal copies the payload into Frame.Data, so the - // buffer is free to recycle on the next read - c.readBuf, err = readAllInto(c.readBuf, r) - if err != nil { - return nil, err - } - f := &livekit.AgentHttp_Frame{} - if err := proto.Unmarshal(c.readBuf, f); err != nil { - return nil, err - } - return f, nil - } -} - -// readAllInto reads r to EOF, reusing buf's capacity, and returns the filled -// slice (a re-slice of buf, or a larger buffer if the message outgrew it). -func readAllInto(buf []byte, r io.Reader) ([]byte, error) { - buf = buf[:0] - for { - if len(buf) == cap(buf) { - buf = append(buf, 0)[:len(buf)] // grow capacity, keep length - } - n, err := r.Read(buf[len(buf):cap(buf)]) - buf = buf[:len(buf)+n] - if err != nil { - if err == io.EOF { - err = nil - } - return buf, err - } - } -} - -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. -// wireParams are the data-plane flow-control parameters, from config with -// package defaults. They are independent of the signalling message limit, which -// governs the separate control websocket. -func (s *AgentService) wireParams() endpoint.WireParams { - return endpoint.WireParams{ - CreditWindow: s.endpointsConfig.CreditWindow, - ConnectionWindow: s.endpointsConfig.ConnectionWindow, - MaxFrameSize: s.endpointsConfig.MaxFrameSize, - MaxStreamsPerConn: s.endpointsConfig.MaxStreamsPerConn, - }.WithDefaults() -} - -// 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. apiKey is the identity the connecting grant is authorized for: -// the adopter must refuse to install a registration for any other project, so -// an unentitled grant cannot adopt (or even probe) another project's worker. -// nil rejects unknown workers. -type AttachAdopter func(a *livekit.AgentHttp_AttachDataConnection, apiKey string) 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, apiKey string) { - 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(), apiKey, a.GetAttachToken()) - if (errors.Is(err, endpoint.ErrUnknownWorker) || errors.Is(err, endpoint.ErrWrongEpoch)) && adopt != nil { - if aerr := adopt(a, apiKey); aerr != nil { - _ = respond(aerr.Error()) - _ = wire.Close() - return - } - ticket, err = registry.BeginAttach(a.GetWorkerId(), a.GetInstanceId(), apiKey, 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, @@ -482,7 +279,7 @@ func NewAgentHandler( } // endpointSettings validates a registration's manifest and negotiates the -// data-plane settings, minting the epoch's attach token. +// data-plane protocol version. 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") @@ -493,18 +290,22 @@ func (h *AgentHandler) endpointSettings(req *livekit.RegisterWorkerRequest) (*li 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 + // the data plane is a WebTransport session (no attach token, no fixed conn + // pool); only the protocol version is negotiated here. + return &livekit.AgentHttp_AgentEndpointSettings{Protocol: endpoint.CurrentProtocol}, nil } +// HandleConnection serves a worker's control connection with no data-plane +// session (the WebSocket path: control and job dispatch only). func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalConn, registration agent.WorkerRegistration) { + h.handleConnection(ctx, conn, registration, nil) +} + +// handleConnection serves a worker's control connection. sess is the worker's +// data-plane session when it registered over WebTransport (control and HTTP +// exchanges share it); nil for a WebSocket control connection, which serves no +// endpoints. +func (h *AgentHandler) handleConnection(ctx context.Context, conn agent.SignalConn, registration agent.WorkerRegistration, sess endpoint.Session) { registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger, func(wr *agent.WorkerRegisterer) { wr.WithEndpointSettings(h.endpointSettings) }) @@ -518,7 +319,7 @@ func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalCo worker := agent.NewWorker(registration, apiKey, apiSecret, conn, h.logger) h.registerWorker(worker) - endpointReg := h.registerEndpoints(worker) + endpointReg := h.registerEndpoints(worker, sess) handlerWorker := &agentHandlerWorker{h, worker} for ok := true; ok; { @@ -532,11 +333,13 @@ func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalCo 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 { +// registerEndpoints registers a worker's endpoint manifest into the data-plane +// registry, binding it to the worker's data-plane session. The registration +// lives exactly as long as the control connection. A nil session (WebSocket +// control path) registers nothing, since HTTP endpoints require a WebTransport +// session to serve them. +func (h *AgentHandler) registerEndpoints(w *agent.Worker, sess endpoint.Session) *endpoint.Registration { + if sess == nil || w.EndpointSettings == nil { return nil } manifest, err := endpoint.ParseManifest(w.Endpoints) @@ -547,19 +350,14 @@ func (h *AgentHandler) registerEndpoints(w *agent.Worker) *endpoint.Registration } reg := &endpoint.Registration{ WorkerID: w.ID, - InstanceID: w.InstanceID, APIKey: w.APIKey(), + AgentName: w.AgentName, Deployment: w.Deployment, Manifest: manifest, Endpoints: w.Endpoints, - Settings: endpoint.Settings{ - Protocol: settings.GetProtocol(), - AttachToken: settings.GetAttachToken(), - DataConnCount: settings.GetDataConnectionCount(), - }, - Logger: w.Logger(), - Draining: w.Draining, + Draining: w.Draining, } + reg.SetSession(sess) if err := h.endpointRegistry.Register(reg); err != nil { w.Logger().Errorw("failed to register endpoints", err) return nil diff --git a/pkg/service/agentwebtransport.go b/pkg/service/agentwebtransport.go new file mode 100644 index 000000000..f246bff09 --- /dev/null +++ b/pkg/service/agentwebtransport.go @@ -0,0 +1,234 @@ +// Copyright 2026 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" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "errors" + "math/big" + "net" + "net/http" + "strconv" + "sync" + "time" + + "github.com/quic-go/quic-go/http3" + "github.com/quic-go/webtransport-go" + + "github.com/livekit/livekit-server/pkg/agent" + "github.com/livekit/livekit-server/pkg/agent/endpoint" + "github.com/livekit/livekit-server/pkg/rtc" + "github.com/livekit/protocol/auth" + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" +) + +// StartWebTransport starts the worker WebTransport listener (control + data) on +// the configured UDP port, if endpoints are enabled and a port is set. It +// returns a stop function (always non-nil). QUIC has no plaintext mode, so a +// TLS certificate is required: from tls_cert_file/tls_key_file, or a generated +// self-signed cert in dev mode. +func (s *AgentService) StartWebTransport(dev bool) (func(), error) { + noop := func() {} + cfg := s.endpointsConfig + if cfg.Disabled || cfg.WebTransportPort == 0 { + return noop, nil + } + tlsConf, err := WebTransportTLS(cfg.TLSCertFile, cfg.TLSKeyFile, dev) + if err != nil { + return noop, err + } + udp, err := net.ListenUDP("udp", &net.UDPAddr{Port: int(cfg.WebTransportPort)}) + if err != nil { + return noop, err + } + wt := NewAgentWebTransportServer(s, s.keyProvider, tlsConf) + go func() { + if err := wt.Serve(udp); err != nil { + logger.Infow("agent webtransport listener stopped", "error", err) + } + }() + logger.Infow("agent webtransport listener started", "port", cfg.WebTransportPort) + return func() { _ = wt.Close(); _ = udp.Close() }, nil +} + +// WebTransportTLS builds the listener's TLS config from cert files, or a +// generated self-signed cert in dev mode. Shared by the OSS and cloud servers. +func WebTransportTLS(certFile, keyFile string, dev bool) (*tls.Config, error) { + if certFile != "" && keyFile != "" { + cert, err := tls.LoadX509KeyPair(certFile, keyFile) + if err != nil { + return nil, err + } + return &tls.Config{Certificates: []tls.Certificate{cert}, NextProtos: []string{http3.NextProtoH3}}, nil + } + if !dev { + return nil, errors.New("agent endpoints WebTransport requires agents.endpoints.tls_cert_file/tls_key_file") + } + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return nil, err + } + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "livekit-agent-webtransport-dev"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(365 * 24 * time.Hour), + IPAddresses: []net.IP{net.IPv4(127, 0, 0, 1), net.IPv6loopback}, + DNSNames: []string{"localhost"}, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + if err != nil { + return nil, err + } + logger.Warnw("agent webtransport using a generated self-signed cert (dev mode); workers must skip verification", nil) + return &tls.Config{ + Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}}, + NextProtos: []string{http3.NextProtoH3}, + }, nil +} + +// NewWebTransportServer wraps an HTTP/3 WebTransport server around a handler. +// register is called with the constructed server so the caller can mount routes +// that Upgrade on it (Upgrade needs the *webtransport.Server); it returns the +// HTTP/3 handler. tlsConf must be usable for HTTP/3 (the h3 ALPN is set here if +// absent). Shared by the OSS and cloud agent servers. +func NewWebTransportServer(tlsConf *tls.Config, register func(*webtransport.Server) http.Handler) *webtransport.Server { + tlsConf = tlsConf.Clone() + if len(tlsConf.NextProtos) == 0 { + tlsConf.NextProtos = []string{http3.NextProtoH3} + } + wt := &webtransport.Server{H3: &http3.Server{TLSConfig: tlsConf}} + wt.H3.Handler = register(wt) + return wt +} + +// NewAgentWebTransportServer builds the WebTransport server that terminates a +// worker's unified session on /agent: the control stream carries the same +// WorkerMessage/ServerMessage exchange as the WebSocket control connection, and +// each HTTP exchange rides a node-opened QUIC stream on the same session. The +// handler runs behind the same api-key auth middleware as the rest of the +// agent surface, so the agent grant is enforced identically. +func NewAgentWebTransportServer(svc *AgentService, keyProvider auth.KeyProvider, tlsConf *tls.Config) *webtransport.Server { + return NewWebTransportServer(tlsConf, func(wt *webtransport.Server) http.Handler { + authMW := NewAPIKeyAuthMiddleware(keyProvider) + mux := http.NewServeMux() + mux.HandleFunc("/agent", func(w http.ResponseWriter, r *http.Request) { + svc.ServeWebTransport(wt, w, r) + }) + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + authMW.ServeHTTP(w, r, mux.ServeHTTP) + }) + }) +} + +// ServeWebTransport verifies the agent grant, upgrades the request to a +// WebTransport session, and serves it. It assumes an outer auth middleware has +// already placed the grant in the request context (as the WebSocket path does). +func (s *AgentService) ServeWebTransport(wt *webtransport.Server, w http.ResponseWriter, r *http.Request) { + claims := GetGrants(r.Context()) + if claims == nil || claims.Video == nil || !claims.Video.Agent { + HandleError(w, r, http.StatusUnauthorized, rtc.ErrPermissionDenied) + return + } + apiKey := GetAPIKey(r.Context()) + + registration := agent.MakeWorkerRegistration() + registration.ClientIP = GetClientIP(r) + if pv, err := strconv.Atoi(r.FormValue("protocol")); err == nil { + registration.Protocol = agent.WorkerProtocolVersion(pv) + } + + sess, err := wt.Upgrade(w, r) + if err != nil { + s.logger.Warnw("agent webtransport upgrade failed", err) + return + } + + // a detached context carries the worker's identity for the whole session; + // the control loop exits on the control stream's read error when the session + // dies, so it needs no cancellation from the request. + ctx := WithGrants(context.Background(), claims, apiKey) + go s.serveWebTransportSession(ctx, sess, registration) +} + +func (s *AgentService) serveWebTransportSession(ctx context.Context, sess *webtransport.Session, registration agent.WorkerRegistration) { + // the worker opens the control stream first (see the conformance worker and + // the SDK); every other stream the node opens is an HTTP exchange. Bound the + // wait so a session that upgrades but never opens a control stream can't hold + // a goroutine until the QUIC idle timeout. + acceptCtx, cancel := context.WithTimeout(sess.Context(), agent.RegisterTimeout) + defer cancel() + control, err := sess.AcceptStream(acceptCtx) + if err != nil { + _ = sess.CloseWithError(0, "no control stream") + return + } + sigConn := NewWTSignalConn(sess, control) + defer sigConn.Close() + maxStreams := int(s.endpointsConfig.MaxStreams) + if maxStreams <= 0 { + maxStreams = endpoint.DefaultMaxStreams + } + epSession := endpoint.NewWebTransportSession(sess, maxStreams) + s.handleConnection(ctx, sigConn, registration, epSession) +} + +// wtSignalConn adapts a WebTransport control stream to agent.SignalConn: the +// same length-delimited WorkerMessage/ServerMessage exchange the WebSocket +// control connection carries, framed on one QUIC bidirectional stream. +type wtSignalConn struct { + sess *webtransport.Session + control *webtransport.Stream + writeMu sync.Mutex +} + +// NewWTSignalConn adapts a WebTransport session's control stream to +// agent.SignalConn. Shared by the OSS and cloud agent servers. +func NewWTSignalConn(sess *webtransport.Session, control *webtransport.Stream) agent.SignalConn { + return &wtSignalConn{sess: sess, control: control} +} + +func (c *wtSignalConn) WriteServerMessage(msg *livekit.ServerMessage) (int, error) { + c.writeMu.Lock() + defer c.writeMu.Unlock() + return 0, endpoint.WriteControlMessage(c.control, msg) +} + +func (c *wtSignalConn) ReadWorkerMessage() (*livekit.WorkerMessage, int, error) { + var msg livekit.WorkerMessage + if err := endpoint.ReadControlMessage(c.control, &msg); err != nil { + return nil, 0, err + } + return &msg, 0, nil +} + +func (c *wtSignalConn) SetReadDeadline(t time.Time) error { + return c.control.SetReadDeadline(t) +} + +func (c *wtSignalConn) Close() error { + return c.sess.CloseWithError(0, "") +} + +func (c *wtSignalConn) CloseWithReason(reason string) error { + return c.sess.CloseWithError(0, reason) +} diff --git a/pkg/service/server.go b/pkg/service/server.go index 38aaf52e2..8cfadd56c 100644 --- a/pkg/service/server.go +++ b/pkg/service/server.go @@ -60,6 +60,7 @@ type LivekitServer struct { signalServer *SignalServer turnServer *turn.Server currentNode routing.LocalNode + agentWTStop func() running atomic.Bool doneChan chan struct{} closedChan chan struct{} @@ -236,6 +237,14 @@ func (s *LivekitServer) Start() error { return err } + if s.agentService != nil { + stop, err := s.agentService.StartWebTransport(s.config.Development) + if err != nil { + return err + } + s.agentWTStop = stop + } + addresses := s.config.BindAddresses if addresses == nil { addresses = []string{""} @@ -374,6 +383,9 @@ func (s *LivekitServer) Stop(force bool) { } s.router.Stop() + if s.agentWTStop != nil { + s.agentWTStop() + } close(s.doneChan) // wait for fully closed