dandri/livekit
1
0 Fork 0
mirror of https://github.com/livekit/livekit.git synced 2026-10-23 13:02:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

agent endpoints: WebTransport data plane

Browse Source 
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.
This commit is contained in:
Théo Monnom committed 2026-08-30 18:59:30 -07:00
1 parent 1652c0d889
commit 751d85c9bf
25 files changed
+1373 -2485
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

pkg/agent/config.go
+14 -7
View File
@@ -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"`
}
pkg/agent/endpoint/conformance/worker.go
+156 -697
View File
@@ -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()
}
pkg/agent/endpoint/conn.go
-386
View File
@@ -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
}
}
}
pkg/agent/endpoint/control.go
+64
View File
@@ -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)
}
pkg/agent/endpoint/front.go
+33 -37
View File
@@ -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
}
pkg/agent/endpoint/manifest.go
+19
View File
@@ -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{
pkg/agent/endpoint/registry.go
+72 -218
View File
@@ -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
pkg/agent/endpoint/registry_test.go
+59 -19
View File
@@ -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"))
}
pkg/agent/endpoint/routematch.go
+36 -20
View File
@@ -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
}
pkg/agent/endpoint/routematch_test.go
+76 -44
View File
@@ -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))
}
pkg/agent/endpoint/scheduler.go
-296
View File
@@ -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
}
}
}
pkg/agent/endpoint/session.go
+78
View File
@@ -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
}
pkg/agent/endpoint/stream.go
-366
View File
@@ -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()
}
pkg/agent/endpoint/template.go
+2 -7
View File
@@ -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
}
pkg/agent/endpoint/template_test.go
+9
View File
@@ -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"}},
})
pkg/agent/endpoint/webtransport.go
+160
View File
@@ -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
}
pkg/agent/endpoint/webtransport_test.go
+184
View File
@@ -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)
})
}
pkg/agent/endpoint/wire.go
-122
View File
@@ -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,
}
}
pkg/agent/worker.go
+8 -2
View File
@@ -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
pkg/service/agentendpoint_test.go
+110 -27
View File
@@ -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)
pkg/service/agentservice.go
+28 -230
View File
@@ -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
pkg/service/agentwebtransport.go
+234
View File
@@ -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)
}
pkg/service/server.go
+12
View File
@@ -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