agent: HTTP endpoints data plane

Serve worker-declared FastAPI routes at /agents/{deployment}/{path} without
any worker-side listener: workers dial a fixed pool of wires speaking
AgentHttp.Frame, the server opens multiplexed streams carrying one opaque
HTTP/1.1 exchange each (two-level credit flow control, prioritized write
scheduler, attach epoch fencing). The front does starlette-exact manifest
matching with per-endpoint public access, typed 401/404/405, a retry table,
and SSE/WebSocket passthrough; a pluggable fallback hook lets multi-node
deployments resolve misses elsewhere. Includes a conformance client, a
manual sidecar, and the acceptance suite.

The auth middleware now reads access_token from the query string only:
FormValue consumed the bodies of proxied url-encoded POSTs.
This commit is contained in:
Théo Monnom
2026-08-19 18:25:35 -07:00
parent 636214a6b0
commit 8c7011cb7a
21 changed files with 4207 additions and 20 deletions
+111
View File
@@ -0,0 +1,111 @@
// 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.
// agent-endpoint-client is the reference sidecar for the agent HTTP endpoints
// data plane: it registers a manifest against a livekit-server and bridges
// tunnel streams to any local HTTP server, standing in for the SDK's tunnel
// client until it lands.
//
// Example, against a dev server:
//
// livekit-server --dev &
// python app.py # any local HTTP server on :8080
// agent-endpoint-client -url ws://localhost:7880/agent -api-key devkey \
// -api-secret secret -deployment production -target 127.0.0.1:8080 \
// -route "GET /json" -route "public GET /sse"
// curl http://localhost:7880/agents/production/json
package main
import (
"context"
"flag"
"fmt"
"os"
"os/signal"
"strings"
"syscall"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/livekit-server/pkg/agent/endpoint/client"
)
type routeFlags []string
func (r *routeFlags) String() string { return strings.Join(*r, ",") }
func (r *routeFlags) Set(v string) error { *r = append(*r, v); return nil }
func main() {
var (
url = flag.String("url", "ws://localhost:7880/agent", "livekit-server /agent URL")
apiKey = flag.String("api-key", "devkey", "API key")
apiSecret = flag.String("api-secret", "secret", "API secret")
agentName = flag.String("agent-name", "endpoint-sidecar", "agent name")
deployment = flag.String("deployment", "", "deployment name (empty = default)")
target = flag.String("target", "127.0.0.1:8080", "local HTTP server to bridge into")
routes routeFlags
)
flag.Var(&routes, "route", `route to expose, e.g. "GET /json", "public POST /sms", repeatable`)
flag.Parse()
if len(routes) == 0 {
fmt.Fprintln(os.Stderr, "at least one -route is required")
os.Exit(1)
}
var endpoints []*livekit.AgentHttp_AgentEndpoint
for _, r := range routes {
parts := strings.Fields(r)
public := false
if len(parts) > 0 && parts[0] == "public" {
public = true
parts = parts[1:]
}
if len(parts) != 2 {
fmt.Fprintf(os.Stderr, "invalid -route %q, want \"[public] METHOD /path\"\n", r)
os.Exit(1)
}
endpoints = append(endpoints, &livekit.AgentHttp_AgentEndpoint{
Path: parts[1],
Methods: []string{strings.ToUpper(parts[0])},
Public: public,
})
}
logger.InitFromConfig(&logger.Config{Level: "info"}, "agent-endpoint-client")
w := client.New(client.Config{
ServerURL: *url,
APIKey: *apiKey,
APISecret: *apiSecret,
AgentName: *agentName,
Deployment: *deployment,
Endpoints: endpoints,
TargetAddr: *target,
Logger: logger.GetLogger(),
})
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer cancel()
if err := w.Start(ctx); err != nil {
fmt.Fprintln(os.Stderr, "start failed:", err)
os.Exit(1)
}
logger.Infow("endpoint sidecar attached", "workerID", w.WorkerID(), "routes", len(endpoints))
<-ctx.Done()
w.Close()
}
+1 -1
View File
@@ -21,7 +21,7 @@ require (
github.com/jxskiss/base62 v1.1.0
github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731
github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0
github.com/livekit/protocol v1.50.5-0.20260814120900-8b1ab81c7d00
github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e
github.com/livekit/psrpc v0.7.3
github.com/mackerelio/go-osstat v0.2.8
github.com/magefile/mage v1.17.2
+2 -2
View File
@@ -166,8 +166,8 @@ github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731 h1:9x+U2HGLrSw5AT
github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731/go.mod h1:Rs3MhFwutWhGwmY1VQsygw28z5bWcnEYmS1OG9OxjOQ=
github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0 h1:XHNNzebIKZRkLimla/hFGrAIX5EMWHctrgt3hLw7s+I=
github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0/go.mod h1:o8CFmAdrVwzJNOCsQCLUzXRjokkufNshnQHOe4fRaqU=
github.com/livekit/protocol v1.50.5-0.20260814120900-8b1ab81c7d00 h1:g4Rdg7gAqPX/CrV7ptAERPtDX1YbjpqAF67G/6pd2Zs=
github.com/livekit/protocol v1.50.5-0.20260814120900-8b1ab81c7d00/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34=
github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e h1:AjYyvyLEC+rB2HKit2RbwMnqUwujH2b7OSUaRUcYbVA=
github.com/livekit/protocol v1.50.5-0.20260820011432-2480d740da5e/go.mod h1:/kYxa0dlTuH981LaBFHG/Swyr969d0+2+/6Lm7fFc34=
github.com/livekit/psrpc v0.7.3 h1:bekuZt/ZQzg8+/M8G6G5jq7bvV9fAKdPHSOZeTwrIIc=
github.com/livekit/psrpc v0.7.3/go.mod h1:rAI+m2+/cb4x9RXhLRtUx5ZwdfjjXOl4zi46IjEetaw=
github.com/livekit/webrtc-pion/v4 v4.2.18-warp.1 h1:fH+v4W+NFp9FfPzON6FaUFNmazGcctaAhb2P+Ksf+1s=
+3 -3
View File
@@ -62,9 +62,9 @@ type JobRequest struct {
// only set for participant jobs
Participant *livekit.ParticipantInfo
Metadata string
AgentName string
Deployment string
Attributes map[string]string
AgentName string
Deployment string
Attributes map[string]string
}
type agentClient struct {
+16
View File
@@ -6,4 +6,20 @@ type Config struct {
EnableUserDataRecording bool `yaml:"enable_user_data_recording"`
EnableUserDataRedaction bool `yaml:"enable_user_data_redaction"`
TargetLoad float32 `yaml:"target_load,omitempty"`
// agent HTTP endpoints data plane; zero values take the endpoint package
// defaults
Endpoints EndpointsConfig `yaml:"endpoints,omitempty"`
}
type EndpointsConfig struct {
// Disabled turns off the /agents/{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"`
}
+799
View File
@@ -0,0 +1,799 @@
// 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 client is the reference worker-side implementation of the agent HTTP
// endpoints data plane: it registers a manifest over the control connection,
// attaches the fixed wire pool, and bridges each stream to a local HTTP server.
// It doubles as the protocol conformance harness and as a manual sidecar for
// non-SDK workers.
package client
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/url"
"sync"
"time"
"github.com/gorilla/websocket"
"google.golang.org/protobuf/proto"
"github.com/livekit/protocol/auth"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/utils/guid"
)
type Config struct {
// ServerURL is the ws(s):// URL of the /agent endpoint
ServerURL string
APIKey string
APISecret string
AgentName string
Deployment string
Endpoints []*livekit.AgentHttp_AgentEndpoint
// TargetAddr is the host:port of the local HTTP server streams bridge into
TargetAddr string
Logger logger.Logger
}
// Worker is one registration epoch: a control connection plus the fixed
// data-connection pool.
type Worker struct {
cfg Config
instanceID string
mu sync.Mutex
workerID string
settings *livekit.AgentHttp_AgentEndpointSettings
control *wsConn
dataConns []*dataConn
closed bool
registered chan struct{}
statusSeq uint64
}
func New(cfg Config) *Worker {
if cfg.Logger == nil {
cfg.Logger = logger.GetLogger()
}
return &Worker{
cfg: cfg,
instanceID: guid.New("AEI_"),
registered: make(chan struct{}),
}
}
// Start registers and attaches the pool; it returns once the data plane is up.
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)
if err != nil {
return err
}
w.mu.Lock()
w.control = control
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,
},
}}); err != nil {
return err
}
go w.controlLoop(ctx, control)
select {
case <-w.registered:
case <-ctx.Done():
return ctx.Err()
case <-time.After(10 * time.Second):
return errors.New("registration timeout")
}
w.mu.Lock()
settings := w.settings
workerID := w.workerID
w.mu.Unlock()
if settings == nil {
return errors.New("server did not negotiate endpoint settings")
}
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)
w.mu.Unlock()
}
// 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) {
at := auth.NewAccessToken(w.cfg.APIKey, w.cfg.APISecret).
SetVideoGrant(&auth.VideoGrant{Agent: true}).
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()
}
+375
View File
@@ -0,0 +1,375 @@
// 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)
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)
}
// 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
}
}
}
+579
View File
@@ -0,0 +1,579 @@
// 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 (
"bufio"
"context"
"errors"
"fmt"
"io"
"math/rand/v2"
"net"
"net/http"
"net/url"
"strings"
"time"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/utils/guid"
)
const (
// PathPrefix is the public route namespace: /agents/{deployment}/{path...}
PathPrefix = "/agents/"
// responseHeadTimeout bounds the wait for the worker's response head. Bodies
// (SSE, long streams) are unbounded; the head never legitimately takes this
// long.
responseHeadTimeout = 90 * time.Second
// maxAttempts bounds worker retries per request
maxAttempts = 3
)
// ScopeResolver maps an inbound request to its project scope. It returns the
// api key the request is authorized for (empty when unauthenticated) - the
// service layer implements it from validated grants.
type ScopeResolver func(r *http.Request) (apiKey string, authenticated bool)
type Front struct {
registry *Registry
resolveScope ScopeResolver
logger logger.Logger
// fallback is consulted when nothing local can serve the request (no
// candidates, no route match, or every match without capacity); a
// multi-node deployment plugs its resolve-and-relay here. nil means local
// misses are final.
fallback Fallback
// see WithSingleKeyFallback
singleKeyFallback bool
}
func NewFront(registry *Registry, resolveScope ScopeResolver, log logger.Logger) *Front {
return &Front{
registry: registry,
resolveScope: resolveScope,
logger: log.WithComponent("agents.endpoint"),
}
}
// FallbackRequest describes a request nothing local could serve. The request
// body is untouched when the fallback runs.
type FallbackRequest struct {
// Scope is the project identity the front resolved (api key in OSS)
Scope string
Authenticated bool
Deployment string
// Path within the deployment, '/'-rooted
Path string
WebSocket bool
}
// Fallback serves a request elsewhere (e.g. a multi-node relay); it reports
// whether a response was written. Returning false falls back to the local
// status mapping.
type Fallback func(w http.ResponseWriter, r *http.Request, req *FallbackRequest) bool
// WithFallback installs the miss handler consulted when nothing local can
// serve a request.
func (f *Front) WithFallback(fb Fallback) *Front {
f.fallback = fb
return f
}
// WithSingleKeyFallback resolves unauthenticated requests to the registry's
// single api key when the scope resolver yields none. Self-hosted convenience
// only: a multi-tenant front must never guess a scope from what happens to be
// registered.
func (f *Front) WithSingleKeyFallback() *Front {
f.singleKeyFallback = true
return f
}
func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) {
rest, ok := strings.CutPrefix(r.URL.Path, PathPrefix)
if !ok {
http.NotFound(w, r)
return
}
deployment, path, found := strings.Cut(rest, "/")
if !found {
path = ""
}
path = "/" + path
if deployment == "" {
http.NotFound(w, r)
return
}
isWS := isWebSocketUpgrade(r)
apiKey, authenticated := f.resolveScope(r)
if apiKey == "" && f.singleKeyFallback {
// unauthenticated: OSS serves public routes when the worker fleet
// belongs to a single key
apiKey, _ = f.registry.SingleAPIKey()
}
if apiKey == "" {
w.Header().Set("WWW-Authenticate", "Bearer")
http.Error(w, "authentication required", http.StatusUnauthorized)
return
}
candidates := f.registry.Candidates(apiKey, deployment)
if len(candidates) == 0 && f.fallback == nil {
w.Header().Set("Retry-After", "1")
http.Error(w, "no workers available for deployment", http.StatusServiceUnavailable)
return
}
// manifest match across the deployment's workers: FULL wins; PARTIAL only
// yields 405 when nothing matches fully; slash-redirect mirrors FastAPI
var matched []*Registration
var route *Route
partial := false
restricted := false
for _, reg := range candidates {
rt, res := reg.Manifest.Match(path, r.Method, isWS)
switch res {
case MatchFull:
if !authenticated && !rt.Public {
restricted = true
continue
}
if route == nil {
route = rt
}
matched = append(matched, reg)
case MatchPartial:
partial = true
}
}
if route == nil && f.fallback != nil {
// nothing local can serve: hand off before the local status mapping
if f.fallback(w, r, &FallbackRequest{
Scope: apiKey, Authenticated: authenticated,
Deployment: deployment, Path: path, WebSocket: isWS,
}) {
return
}
if len(candidates) == 0 && !restricted && !partial {
w.Header().Set("Retry-After", "1")
http.Error(w, "no workers available for deployment", http.StatusServiceUnavailable)
return
}
}
if route == nil {
switch {
case restricted:
w.Header().Set("WWW-Authenticate", "Bearer")
http.Error(w, "authentication required", http.StatusUnauthorized)
case partial:
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
default:
for _, reg := range candidates {
if alt, ok := reg.Manifest.RedirectSlashes(path, r.Method, isWS); ok {
u := *r.URL
u.Path = PathPrefix + deployment + alt
http.Redirect(w, r, u.String(), http.StatusTemporaryRedirect)
return
}
}
http.NotFound(w, r)
}
return
}
bodyConsumed := int64(0)
countingBody := &countingReader{r: r.Body, n: &bodyConsumed}
attempted := make(map[*Registration]bool)
for attempt := 0; attempt < maxAttempts; attempt++ {
reg := pickWeighted(matched, attempted)
if reg == nil {
break
}
attempted[reg] = true
done, retryable := f.bridge(w, r, reg, path, countingBody, bodyConsumed, isWS)
if done || !retryable {
return
}
}
// the route matched locally but nothing served it (matches draining or
// conn-less, or every attempt failed before writing): the fallback may hold
// capacity elsewhere. Safe exactly while no request bytes were consumed -
// reaching this point implies it, since consuming attempts are never
// retryable.
if bodyConsumed == 0 && f.fallback != nil {
if f.fallback(w, r, &FallbackRequest{
Scope: apiKey, Authenticated: authenticated,
Deployment: deployment, Path: path, WebSocket: isWS,
}) {
return
}
}
w.Header().Set("Retry-After", "1")
http.Error(w, "no worker could serve the request", http.StatusServiceUnavailable)
}
// pickWeighted is capacity-weighted random over non-draining workers with
// attached data connections.
func pickWeighted(regs []*Registration, ignore map[*Registration]bool) *Registration {
var sum float32
weights := make([]float32, len(regs))
for i, reg := range regs {
if ignore[reg] || reg.AttachedConns() == 0 {
continue
}
if reg.Draining != nil && reg.Draining() {
continue
}
w := float32(1)
if reg.Load != nil {
w = max(0.01, 1-reg.Load())
}
weights[i] = w
sum += w
}
if sum == 0 {
return nil
}
target := rand.Float32() * sum
for i, reg := range regs {
if target -= weights[i]; weights[i] > 0 && target <= 0 {
return reg
}
}
return nil
}
// bridge runs one attempt against one worker. done means a response (or abort)
// reached the client; retryable reports whether another attempt is safe per the
// retry table: idempotent/bodyless until any response byte arrived,
// anything on HSR_REFUSED, nothing once bytes were consumed otherwise.
func (f *Front) bridge(
w http.ResponseWriter,
r *http.Request,
reg *Registration,
path string,
body io.Reader,
bodyConsumedBefore int64,
isWS bool,
) (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,
})
if err != nil {
return false, true
}
defer stream.Close()
ctx := r.Context()
stop := context.AfterFunc(ctx, func() {
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client disconnected")
})
defer stop()
// serialize the request into the stream concurrently with response reading:
// directions are independent (full duplex within the stream)
outReq := f.outboundRequest(r, path, body, isWS)
writeErrCh := make(chan error, 1)
if isWS {
// upgrade requests have no body: write the head inline and do NOT
// half-close - the client->worker direction carries the session, and
// the upgrade pump must be the stream's only writer
if err := outReq.Write(stream); err != nil {
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "request write failed")
return false, true
}
writeErrCh <- nil
} else {
go func() {
err := outReq.Write(stream)
if err == nil {
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")
}
writeErrCh <- err
}()
}
counted := &countingReader{r: stream, n: new(int64)}
br := bufio.NewReader(counted)
resp, err := f.readResponseHead(w, br, outReq, stream)
if err != nil {
retryable = f.classifyRetry(r, stream, *counted.n, bodyConsumedBefore, err)
if !retryable {
f.logger.Warnw("agent endpoint request failed", err,
"workerID", reg.WorkerID, "path", path)
http.Error(w, "bad gateway", http.StatusBadGateway)
return true, false
}
// 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")
<-writeErrCh
return false, true
}
// a response byte arrived: from here every failure is surfaced, never retried
if resp.StatusCode == http.StatusSwitchingProtocols {
f.bridgeUpgrade(w, resp, br, stream)
return true, false
}
copyResponseHeaders(w.Header(), resp)
w.WriteHeader(resp.StatusCode)
rc := http.NewResponseController(w)
buf := make([]byte, 32<<10)
for {
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")
return true, false
}
_ = rc.Flush()
}
if rerr == io.EOF {
break
}
if rerr != nil {
// never expose a clean-looking short body
select {
case werr := <-writeErrCh:
f.logger.Debugw("request write result after response failure", "error", werr)
default:
}
panic(http.ErrAbortHandler)
}
}
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) {
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")
case <-headCh:
}
}()
defer close(headCh)
for {
resp, err := http.ReadResponse(br, outReq)
if err != nil {
return nil, err
}
if resp.StatusCode >= 100 && resp.StatusCode < 200 && resp.StatusCode != http.StatusSwitchingProtocols {
// informational: relay and keep reading
for k, vv := range resp.Header {
for _, v := range vv {
w.Header().Add(k, v)
}
}
w.WriteHeader(resp.StatusCode)
clear(w.Header())
continue
}
return resp, nil
}
}
// bridgeUpgrade hijacks the client connection after a 101 and pumps raw bytes in
// both directions; the stream carries the rest of the WebSocket session.
func (f *Front) bridgeUpgrade(w http.ResponseWriter, resp *http.Response, br *bufio.Reader, stream *Stream) {
hj, ok := w.(http.Hijacker)
if !ok {
// e.g. HTTP/2 client conns cannot be upgraded
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "client does not support upgrade")
http.Error(w, "upgrade not supported on this connection", http.StatusBadGateway)
return
}
clientConn, clientRW, err := hj.Hijack()
if err != nil {
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "hijack failed")
return
}
defer clientConn.Close()
if err := resp.Write(clientRW); err != nil {
return
}
if err := clientRW.Flush(); err != nil {
return
}
errCh := make(chan error, 2)
go func() {
// worker -> client, including bytes the bufio reader already buffered
_, err := io.Copy(clientConn, br)
errCh <- err
}()
go func() {
// client -> worker
buf := make([]byte, 32<<10)
for {
n, rerr := clientRW.Read(buf)
if n > 0 {
if _, werr := stream.Write(buf[:n]); werr != nil {
errCh <- werr
return
}
}
if rerr != nil {
_ = stream.CloseWrite()
errCh <- rerr
return
}
}
}()
<-errCh
stream.Reset(livekit.AgentHttp_HSR_CANCEL, "upgrade session ended")
}
// classifyRetry implements the retry table.
func (f *Front) classifyRetry(r *http.Request, stream *Stream, responseBytes, bodyConsumedBefore int64, err error) bool {
if responseBytes > 0 || stream.BytesRead() > 0 {
return false
}
if stream.Refused() {
// the worker proved non-dispatch; safe for any method, but only when the
// 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
}
return false
}
// outboundRequest builds the request serialized into the stream: the path the
// worker's router sees (deployment prefix stripped), hop-by-hop headers removed,
// forwarding headers appended.
func (f *Front) outboundRequest(r *http.Request, path string, body io.Reader, isWS bool) *http.Request {
out := r.Clone(r.Context())
out.RequestURI = ""
out.URL = &url.URL{Path: path, RawQuery: r.URL.RawQuery}
out.Host = r.Host
out.Body = io.NopCloser(body)
// one exchange per stream: closing the worker-local app connection after the
// response is what lets the opaque pump observe the end of the exchange and
// free the stream slot
out.Close = !isWS
removeHopByHopHeaders(out.Header, isWS)
out.Header.Del("Expect") // the front owns 100-continue semantics client-side
if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
prior := out.Header.Get("X-Forwarded-For")
if prior != "" {
out.Header.Set("X-Forwarded-For", prior+", "+host)
} else {
out.Header.Set("X-Forwarded-For", host)
}
}
return out
}
// hop-by-hop headers per RFC 9110; Connection-nominated headers are dropped too.
// For WebSocket upgrades Connection/Upgrade survive so the worker-side bridge
// sees a real upgrade request.
func removeHopByHopHeaders(h http.Header, isWS bool) {
for _, f := range h.Values("Connection") {
for _, sf := range strings.Split(f, ",") {
if sf = strings.TrimSpace(sf); sf != "" && !strings.EqualFold(sf, "upgrade") {
h.Del(sf)
}
}
}
for _, k := range []string{
"Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization",
"Te", "Trailer", "Transfer-Encoding",
} {
h.Del(k)
}
if !isWS {
h.Del("Connection")
h.Del("Upgrade")
} else {
h.Set("Connection", "Upgrade")
}
}
func copyResponseHeaders(dst http.Header, resp *http.Response) {
for k, vv := range resp.Header {
for _, v := range vv {
dst.Add(k, v)
}
}
removeHopByHopHeaders(dst, false)
if resp.ContentLength >= 0 && dst.Get("Content-Length") == "" {
dst.Set("Content-Length", fmt.Sprintf("%d", resp.ContentLength))
}
}
type countingReader struct {
r io.Reader
n *int64
}
func (c *countingReader) Read(p []byte) (int, error) {
n, err := c.r.Read(p)
*c.n += int64(n)
return n, err
}
func isWebSocketUpgrade(r *http.Request) bool {
return strings.EqualFold(r.Header.Get("Upgrade"), "websocket") &&
httpHeaderContainsToken(r.Header, "Connection", "upgrade")
}
func httpHeaderContainsToken(h http.Header, name, token string) bool {
for _, v := range h.Values(name) {
for _, f := range strings.Split(v, ",") {
if strings.EqualFold(strings.TrimSpace(f), token) {
return true
}
}
}
return false
}
+133
View File
@@ -0,0 +1,133 @@
// 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 (
"fmt"
"slices"
"strings"
"github.com/livekit/protocol/livekit"
)
const MaxManifestRoutes = 256
// Route is one validated manifest entry.
type Route struct {
Template *Template
Methods []string // uppercase; empty for websocket routes
Kind livekit.AgentHttp_AgentEndpointKind
Public bool
}
// Manifest is a worker's ordered route table.
type Manifest struct {
routes []Route
}
// MatchResult mirrors starlette's Match enum: a FULL match selects the route, a
// PARTIAL match (path matched, method didn't) yields 405 only after the whole
// table has been scanned, so a later route with the right method still wins.
type MatchResult int
const (
MatchNone MatchResult = iota
MatchPartial
MatchFull
)
// ParseManifest validates a registration's endpoint list.
func ParseManifest(endpoints []*livekit.AgentHttp_AgentEndpoint) (*Manifest, error) {
if len(endpoints) > MaxManifestRoutes {
return nil, fmt.Errorf("manifest exceeds %d routes", MaxManifestRoutes)
}
m := &Manifest{routes: make([]Route, 0, len(endpoints))}
for _, ep := range endpoints {
tpl, err := ParseTemplate(ep.GetPath())
if err != nil {
return nil, err
}
var methods []string
switch ep.GetKind() {
case livekit.AgentHttp_AEK_HTTP:
if len(ep.GetMethods()) == 0 {
return nil, fmt.Errorf("endpoint %q declares no methods", ep.GetPath())
}
for _, method := range ep.GetMethods() {
u := strings.ToUpper(method)
if u != method {
return nil, fmt.Errorf("endpoint %q method %q must be uppercase", ep.GetPath(), method)
}
methods = append(methods, u)
}
case livekit.AgentHttp_AEK_WEBSOCKET:
if len(ep.GetMethods()) != 0 {
return nil, fmt.Errorf("websocket endpoint %q must not declare methods", ep.GetPath())
}
default:
return nil, fmt.Errorf("endpoint %q has unsupported kind %s", ep.GetPath(), ep.GetKind())
}
m.routes = append(m.routes, Route{
Template: tpl,
Methods: methods,
Kind: ep.GetKind(),
Public: ep.GetPublic(),
})
}
return m, nil
}
// Match resolves a request path against the table. websocket selects the
// websocket route class (upgrade requests); otherwise the HTTP class with
// method matching.
func (m *Manifest) Match(path, method string, websocket bool) (*Route, MatchResult) {
partial := false
for i := range m.routes {
r := &m.routes[i]
if websocket != (r.Kind == livekit.AgentHttp_AEK_WEBSOCKET) {
continue
}
if !r.Template.Match(path) {
continue
}
if websocket || slices.Contains(r.Methods, method) {
return r, MatchFull
}
partial = true
}
if partial {
return nil, MatchPartial
}
return nil, MatchNone
}
// RedirectSlashes reports whether the alternate-slash form of path would match,
// mirroring starlette's redirect_slashes default (FastAPI 307s /x/ to /x and
// vice versa when only the alternate form matches).
func (m *Manifest) RedirectSlashes(path, method string, websocket bool) (string, bool) {
if path == "/" {
return "", false
}
var alt string
if strings.HasSuffix(path, "/") {
alt = strings.TrimSuffix(path, "/")
} else {
alt = path + "/"
}
if _, res := m.Match(alt, method, websocket); res == MatchFull {
return alt, true
}
return "", false
}
+409
View File
@@ -0,0 +1,409 @@
// 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 (
"crypto/subtle"
"errors"
"slices"
"sync"
"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")
)
func normalizeDeployment(d string) string {
if d == "" {
return DefaultDeployment
}
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).
type Registration struct {
WorkerID string
InstanceID string
APIKey string
Deployment string
Manifest *Manifest
Settings Settings
Logger logger.Logger
// Load and Draining are provided by the control-plane layer that owns the
// worker (reported load rides UpdateWorkerStatus on the control connection).
Load func() float32
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
}
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
}
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 {
return nil, ErrAttachRejected
}
if instanceID != r.InstanceID {
return nil, ErrAttachRejected
}
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.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 {
r.mu.Lock()
defer r.mu.Unlock()
return len(r.conns)
}
func (r *Registration) close() {
r.mu.Lock()
if r.closed {
r.mu.Unlock()
return
}
r.closed = true
conns := slices.Clone(r.conns)
r.conns = nil
r.mu.Unlock()
for _, c := range conns {
c.Close(ErrConnClosed)
}
}
// Registry tracks data-plane registrations on this node, keyed by
// (api key, deployment). The api key is the project identity in OSS.
type Registry struct {
mu sync.Mutex
regs map[string]*Registration // by worker id
byKey map[regKey][]*Registration
// scope hooks fire when a scope (api key / project) gains its first or
// loses its last registration; the Remote uses them to (de)register the
// resolve topic on the bus
scopeCounts map[string]int
onScopeActive func(scope string)
onScopeIdle func(scope string)
// hookMu serializes scope hook invocations. Transitions are decided against
// the registry's CURRENT state under hookMu, never from values computed
// earlier: a reconnect's activation racing the old connection's idle must
// not leave a live scope without its resolve topic.
hookMu sync.Mutex
scopeNotified map[string]bool
}
type regKey struct {
apiKey string
deployment string
}
func NewRegistry() *Registry {
return &Registry{
regs: make(map[string]*Registration),
byKey: make(map[regKey][]*Registration),
scopeCounts: make(map[string]int),
scopeNotified: make(map[string]bool),
}
}
// notifyScope reconciles the scope's hook state with the registry's current
// truth. Racing register/deregister notifications converge instead of
// interleaving.
func (g *Registry) notifyScope(scope string) {
g.hookMu.Lock()
defer g.hookMu.Unlock()
g.mu.Lock()
active := g.scopeCounts[scope] > 0
activeHook, idleHook := g.onScopeActive, g.onScopeIdle
g.mu.Unlock()
if activeHook == nil && idleHook == nil {
return
}
if g.scopeNotified[scope] == active {
return
}
if active {
g.scopeNotified[scope] = true
activeHook(scope)
} else {
delete(g.scopeNotified, scope)
idleHook(scope)
}
}
// SetScopeHooks installs the scope activation callbacks, replaying currently
// active scopes so a resolver layered on after registrations still registers
// its topics. Multi-node deployments use this to answer endpoint resolves only
// while they hold registrations for a scope.
func (g *Registry) SetScopeHooks(active, idle func(scope string)) {
g.hookMu.Lock()
defer g.hookMu.Unlock()
g.mu.Lock()
g.onScopeActive = active
g.onScopeIdle = idle
scopes := make([]string, 0, len(g.scopeCounts))
for scope, n := range g.scopeCounts {
if n > 0 {
scopes = append(scopes, scope)
}
}
g.mu.Unlock()
if active != nil {
for _, s := range scopes {
if !g.scopeNotified[s] {
g.scopeNotified[s] = true
active(s)
}
}
}
}
// 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.
func (g *Registry) Register(r *Registration) error {
key := regKey{r.APIKey, normalizeDeployment(r.Deployment)}
g.mu.Lock()
old := g.regs[r.WorkerID]
if old != nil {
g.removeLocked(old)
}
g.regs[r.WorkerID] = r
g.byKey[key] = append(g.byKey[key], r)
g.scopeCounts[r.APIKey]++
g.mu.Unlock()
g.notifyScope(r.APIKey)
if old != nil {
if old.APIKey != r.APIKey {
g.notifyScope(old.APIKey)
}
old.close()
}
return nil
}
// 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)}
if regs := g.byKey[key]; len(regs) > 0 {
if i := slices.Index(regs, r); i != -1 {
regs = slices.Delete(regs, i, i+1)
}
if len(regs) == 0 {
delete(g.byKey, key)
} else {
g.byKey[key] = regs
}
}
g.scopeCounts[r.APIKey]--
if g.scopeCounts[r.APIKey] == 0 {
delete(g.scopeCounts, r.APIKey)
}
}
// Deregister removes exactly this registration; it is a no-op when a newer
// epoch has already superseded it.
func (g *Registry) Deregister(r *Registration) {
g.mu.Lock()
if g.regs[r.WorkerID] != r {
g.mu.Unlock()
return
}
g.removeLocked(r)
g.mu.Unlock()
g.notifyScope(r.APIKey)
r.close()
}
// ValidateAttach checks an attach without adopting the connection, so the
// attach response can be written before stream frames may flow.
func (g *Registry) ValidateAttach(workerID, instanceID, token string) error {
g.mu.Lock()
r, ok := g.regs[workerID]
g.mu.Unlock()
if !ok {
return ErrUnknownWorker
}
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.
func (g *Registry) BeginAttach(workerID, instanceID, token string) (*AttachTicket, error) {
g.mu.Lock()
r, ok := g.regs[workerID]
g.mu.Unlock()
if !ok {
return nil, ErrUnknownWorker
}
return r.beginAttach(instanceID, token)
}
// Candidates returns the registrations for (api key, deployment segment).
func (g *Registry) Candidates(apiKey, deployment string) []*Registration {
g.mu.Lock()
defer g.mu.Unlock()
return slices.Clone(g.byKey[regKey{apiKey, normalizeDeployment(deployment)}])
}
// SingleAPIKey returns the api key when every registration shares one - the OSS
// resolution for unauthenticated requests to public endpoints. ok is false when
// zero or multiple keys are present.
func (g *Registry) SingleAPIKey() (string, bool) {
g.mu.Lock()
defer g.mu.Unlock()
var key string
for _, r := range g.regs {
if key == "" {
key = r.APIKey
} else if key != r.APIKey {
return "", false
}
}
return key, key != ""
}
+47
View File
@@ -0,0 +1,47 @@
package endpoint
import (
"testing"
"github.com/stretchr/testify/require"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
)
// 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.
func TestRegistrySupersede(t *testing.T) {
g := NewRegistry()
manifest, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{{
Path: "/x", Methods: []string{"GET"}, Public: true,
}})
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(),
}
}
oldReg := mk("i-1")
require.NoError(t, g.Register(oldReg))
newReg := mk("i-2")
require.NoError(t, g.Register(newReg))
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production"))
require.Error(t, g.ValidateAttach("AW_1", "i-1", "ATT_i-1"), "old epoch must not attach")
require.NoError(t, g.ValidateAttach("AW_1", "i-2", "ATT_i-2"))
// the old control connection tears down after the new one registered
g.Deregister(oldReg)
require.Equal(t, []*Registration{newReg}, g.Candidates("key", "production"),
"the retiring epoch must not deregister its successor")
g.Deregister(newReg)
require.Empty(t, g.Candidates("key", "production"))
}
+239
View File
@@ -0,0 +1,239 @@
// 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.
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
}
type streamQueue struct {
chunks [][]byte
eof bool // send an EOF frame after the last chunk drains
}
func newScheduler(wire WireConn) *scheduler {
s := &scheduler{
wire: wire,
queues: make(map[uint32]*streamQueue),
}
s.cond = sync.NewCond(&s.mu)
go s.writeLoop()
return s
}
// 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 []byte) error {
s.mu.Lock()
defer s.mu.Unlock()
for !s.closed && s.queuedBytes+int64(len(chunk)) > connBufferBudget && s.queuedBytes > 0 {
s.cond.Wait()
}
if s.closed {
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))
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))
}
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.
func (s *scheduler) next() (*livekit.AgentHttp_Frame, bool) {
s.mu.Lock()
defer s.mu.Unlock()
for {
if s.closed {
return nil, false
}
if len(s.control) > 0 {
f := s.control[0]
s.control = s.control[1:]
return f, true
}
if len(s.rr) > 0 {
if s.rrIdx >= len(s.rr) {
s.rrIdx = 0
}
id := s.rr[s.rrIdx]
q := s.queues[id]
var f *livekit.AgentHttp_Frame
if len(q.chunks) > 0 {
chunk := q.chunks[0]
q.chunks = q.chunks[1:]
s.queuedBytes -= int64(len(chunk))
f = &livekit.AgentHttp_Frame{
StreamId: id,
Message: &livekit.AgentHttp_Frame_Data{Data: chunk},
}
}
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, 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, true
}
s.cond.Wait()
}
}
func (s *scheduler) writeLoop() {
for {
f, ok := s.next()
if !ok {
return
}
_ = s.wire.SetWriteDeadline(time.Now().Add(writeStallTimeout))
if err := s.wire.WriteFrame(f); err != nil {
s.close(err)
return
}
}
}
+341
View File
@@ -0,0 +1,341 @@
// 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
}
chunk := make([]byte, n)
copy(chunk, p[:n])
if err := s.conn.sched.enqueueData(s.id, chunk); err != nil {
return total, err
}
s.conn.noteActivity(n)
p = p[n:]
total += int(n)
}
return total, nil
}
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()
}
+98
View File
@@ -0,0 +1,98 @@
// 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 (
"fmt"
"regexp"
"strings"
)
// Template is a compiled starlette-style path template. The compilation mirrors
// starlette's compile_path exactly: {name} or {name:convertor} segments become the
// convertor's pattern, everything else is matched literally, anchored on both ends.
// Workers declare templates through FastAPI routes, so any divergence from
// starlette's semantics would make the server route requests the worker-side
// router then refuses.
type Template struct {
raw string
re *regexp.Regexp
}
// convertor patterns copied verbatim from starlette's convertors.py. The uuid
// pattern deliberately makes every hyphen optional, so 32 bare hex characters
// match too.
var convertorPatterns = map[string]string{
"str": `[^/]+`,
"path": `.*`,
"int": `[0-9]+`,
"float": `[0-9]+(?:\.[0-9]+)?`,
"uuid": `[0-9a-fA-F]{8}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{12}`,
}
// starlette's PARAM_REGEX
var paramRegex = regexp.MustCompile(`\{([a-zA-Z_][a-zA-Z0-9_]*)(:[a-zA-Z_][a-zA-Z0-9_]*)?\}`)
// ParseTemplate compiles a starlette path template. Custom convertors are
// rejected: only the five built-ins may travel over the wire.
func ParseTemplate(path string) (*Template, error) {
if !strings.HasPrefix(path, "/") {
return nil, fmt.Errorf("path template must start with '/': %q", path)
}
var pattern strings.Builder
pattern.WriteString("^")
idx := 0
seen := map[string]bool{}
for _, m := range paramRegex.FindAllStringSubmatchIndex(path, -1) {
start, end := m[0], m[1]
name := path[m[2]:m[3]]
convertor := "str"
if m[4] != -1 {
convertor = path[m[4]+1 : m[5]] // skip the ':'
}
convPattern, ok := convertorPatterns[convertor]
if !ok {
return nil, fmt.Errorf("unknown path convertor %q in template %q", convertor, path)
}
if seen[name] {
return nil, fmt.Errorf("duplicated param name %q in template %q", name, path)
}
seen[name] = true
pattern.WriteString(regexp.QuoteMeta(path[idx:start]))
pattern.WriteString("(?:")
pattern.WriteString(convPattern)
pattern.WriteString(")")
idx = end
}
pattern.WriteString(regexp.QuoteMeta(path[idx:]))
pattern.WriteString("$")
re, err := regexp.Compile(pattern.String())
if err != nil {
return nil, fmt.Errorf("invalid path template %q: %w", path, err)
}
return &Template{raw: path, re: re}, nil
}
func (t *Template) Match(path string) bool {
return t.re.MatchString(path)
}
func (t *Template) String() string {
return t.raw
}
+123
View File
@@ -0,0 +1,123 @@
package endpoint
import (
"net/http"
"testing"
"github.com/stretchr/testify/require"
"github.com/livekit/protocol/livekit"
)
func TestTemplateStarletteSemantics(t *testing.T) {
cases := []struct {
template string
path string
match bool
}{
{"/token", "/token", true},
{"/token", "/token/", false},
{"/token", "/Token", false},
{"/users/{id}", "/users/42", true},
{"/users/{id}", "/users/42/posts", false},
{"/users/{id}", "/users/", false},
{"/users/{id:int}", "/users/42", true},
{"/users/{id:int}", "/users/4x2", false},
{"/files/{p:path}", "/files/a/b/c.txt", true},
{"/files/{p:path}", "/files/", true},
{"/price/{v:float}", "/price/1.25", true},
{"/price/{v:float}", "/price/1.", false},
{"/obj/{u:uuid}", "/obj/123e4567-e89b-12d3-a456-426614174000", true},
// starlette's uuid convertor makes every hyphen optional
{"/obj/{u:uuid}", "/obj/123e4567e89b12d3a456426614174000", true},
{"/obj/{u:uuid}", "/obj/123e4567", false},
{"/a/{x}/b/{y}", "/a/1/b/2", true},
{"/a/{x}/b/{y}", "/a/1/c/2", false},
}
for _, c := range cases {
tpl, err := ParseTemplate(c.template)
require.NoError(t, err, c.template)
require.Equal(t, c.match, tpl.Match(c.path), "%s vs %s", c.template, c.path)
}
}
func TestTemplateRejectsCustomConvertors(t *testing.T) {
_, err := ParseTemplate("/x/{id:slug}")
require.Error(t, err)
_, err = ParseTemplate("/x/{a}/{a}")
require.Error(t, err)
_, err = ParseTemplate("no-slash")
require.Error(t, err)
}
func ep(path string, methods []string, public bool) *livekit.AgentHttp_AgentEndpoint {
return &livekit.AgentHttp_AgentEndpoint{Path: path, Methods: methods, Public: public}
}
func TestManifestFullPartialSemantics(t *testing.T) {
// POST /x registered after GET /x must still serve POSTs (starlette scans
// for a FULL match before settling for the PARTIAL 405)
m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
ep("/x", []string{"GET"}, true),
ep("/x", []string{"POST"}, true),
})
require.NoError(t, err)
r, res := m.Match("/x", http.MethodPost, false)
require.Equal(t, MatchFull, res)
require.Contains(t, r.Methods, "POST")
// the manifest carries the app's methods verbatim: FastAPI does not imply
// HEAD from GET, so neither does the matcher
_, res = m.Match("/x", http.MethodHead, false)
require.Equal(t, MatchPartial, res)
// PARTIAL only when no route serves the method
_, res = m.Match("/x", http.MethodDelete, false)
require.Equal(t, MatchPartial, res)
_, res = m.Match("/nope", http.MethodGet, false)
require.Equal(t, MatchNone, res)
}
func TestManifestRedirectSlashes(t *testing.T) {
m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
ep("/hook", []string{"POST"}, true),
})
require.NoError(t, err)
alt, ok := m.RedirectSlashes("/hook/", http.MethodPost, false)
require.True(t, ok)
require.Equal(t, "/hook", alt)
_, ok = m.RedirectSlashes("/other/", http.MethodPost, false)
require.False(t, ok)
}
func TestManifestWebSocketRoutes(t *testing.T) {
m, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
{Path: "/ws", Kind: livekit.AgentHttp_AEK_WEBSOCKET, Public: true},
ep("/http", []string{"GET"}, true),
})
require.NoError(t, err)
_, res := m.Match("/ws", http.MethodGet, true)
require.Equal(t, MatchFull, res)
_, res = m.Match("/ws", http.MethodGet, false)
require.Equal(t, MatchNone, res)
_, res = m.Match("/http", http.MethodGet, true)
require.Equal(t, MatchNone, res)
}
func TestManifestValidation(t *testing.T) {
_, err := ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", nil, false)})
require.Error(t, err, "http endpoint without methods")
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{ep("/x", []string{"get"}, false)})
require.Error(t, err, "lowercase method")
_, err = ParseManifest([]*livekit.AgentHttp_AgentEndpoint{
{Path: "/ws", Kind: livekit.AgentHttp_AEK_WEBSOCKET, Methods: []string{"GET"}},
})
require.Error(t, err, "websocket route with methods")
}
+109
View File
@@ -0,0 +1,109 @@
// 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
}
func (s Settings) Proto() *livekit.AgentHttp_AgentEndpointSettings {
return &livekit.AgentHttp_AgentEndpointSettings{
Protocol: s.Protocol,
AttachToken: s.AttachToken,
DataConnectionCount: s.DataConnCount,
}
}
+59 -7
View File
@@ -155,6 +155,12 @@ type WorkerRegistration struct {
Permissions *livekit.ParticipantPermission
ClientIP string
Deployment string
// agent HTTP endpoints data plane
Endpoints []*livekit.AgentHttp_AgentEndpoint
InstanceID string
EndpointProtocol uint32
EndpointSettings *livekit.AgentHttp_AgentEndpointSettings
}
func MakeWorkerRegistration() WorkerRegistration {
@@ -166,10 +172,17 @@ func MakeWorkerRegistration() WorkerRegistration {
var _ WorkerSignalHandler = (*WorkerRegisterer)(nil)
// EndpointSettingsFunc validates a registration's endpoint manifest and returns the
// negotiated data-plane settings (including the attach token). Returning an error
// fails the registration. It runs only when the registration declares endpoints AND
// a data-plane protocol version; a nil func rejects such registrations.
type EndpointSettingsFunc func(req *livekit.RegisterWorkerRequest) (*livekit.AgentHttp_AgentEndpointSettings, error)
type WorkerRegisterer struct {
WorkerPingHandler
serverInfo *livekit.ServerInfo
deadline time.Time
serverInfo *livekit.ServerInfo
deadline time.Time
endpointSettings EndpointSettingsFunc
registration WorkerRegistration
registered bool
@@ -184,6 +197,13 @@ func NewWorkerRegisterer(conn SignalConn, serverInfo *livekit.ServerInfo, base W
}
}
// WithEndpointSettings enables the HTTP endpoints data plane for registrations that
// declare endpoints.
func (h *WorkerRegisterer) WithEndpointSettings(f EndpointSettingsFunc) *WorkerRegisterer {
h.endpointSettings = f
return h
}
func (h *WorkerRegisterer) Deadline() time.Time {
return h.deadline
}
@@ -221,13 +241,28 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er
h.registration.JobType = req.GetType()
h.registration.Permissions = permissions
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")
}
settings, err := h.endpointSettings(req)
if err != nil {
return err
}
h.registration.EndpointSettings = settings
}
h.registered = true
_, err := h.conn.WriteServerMessage(&livekit.ServerMessage{
Message: &livekit.ServerMessage_Register{
Register: &livekit.RegisterWorkerResponse{
WorkerId: h.registration.ID,
ServerInfo: h.serverInfo,
WorkerId: h.registration.ID,
ServerInfo: h.serverInfo,
EndpointSettings: h.registration.EndpointSettings,
},
},
})
@@ -248,9 +283,11 @@ type Worker struct {
cancel context.CancelFunc
closed chan struct{}
mu sync.Mutex
load float32
status livekit.WorkerStatus
mu sync.Mutex
load float32
status livekit.WorkerStatus
statusSeq uint64
draining bool
runningJobs map[livekit.JobID]*livekit.Job
availability map[livekit.JobID]chan *livekit.AvailabilityResponse
@@ -581,15 +618,30 @@ func (w *Worker) HandleUpdateWorker(update *livekit.UpdateWorkerStatus) error {
w.mu.Lock()
defer w.mu.Unlock()
// status updates may interleave across connections; never regress newer state
if seq := update.GetSeq(); seq != 0 {
if seq <= w.statusSeq {
return nil
}
w.statusSeq = seq
}
if status := update.Status; status != nil && w.status != *status {
w.status = *status
w.Logger().Debugw("worker status changed", "status", w.status)
}
w.load = update.GetLoad()
w.draining = update.GetDraining()
return nil
}
func (w *Worker) Draining() bool {
w.mu.Lock()
defer w.mu.Unlock()
return w.draining
}
func (w *Worker) HandleMigrateJob(req *livekit.MigrateJobRequest) error {
// TODO(theomonnom): On OSS this is not implemented
// We could maybe just move a specific job to another worker
+502
View File
@@ -0,0 +1,502 @@
// 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_test
import (
"bufio"
"bytes"
"context"
"crypto/rand"
"crypto/sha256"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/require"
"github.com/livekit/livekit-server/pkg/agent"
"github.com/livekit/livekit-server/pkg/agent/endpoint"
"github.com/livekit/livekit-server/pkg/agent/endpoint/client"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/routing"
"github.com/livekit/livekit-server/pkg/service"
"github.com/livekit/protocol/auth"
"github.com/livekit/protocol/livekit"
"github.com/livekit/psrpc"
)
const (
testKey = "test"
testSecret = "verysecretsecret"
)
type endpointStack struct {
t *testing.T
ts *httptest.Server
svc *service.AgentService
}
func newEndpointStack(t *testing.T, endpointsCfg agent.EndpointsConfig) *endpointStack {
localNode, err := routing.NewLocalNode(nil)
require.NoError(t, err)
keyProvider := auth.NewSimpleKeyProvider(testKey, testSecret)
svc, err := service.NewAgentService(
&config.Config{
Region: "test",
Keys: map[string]string{testKey: testSecret},
Agents: agent.Config{TargetLoad: agent.DefaultTargetLoad, Endpoints: endpointsCfg},
},
localNode,
psrpc.NewLocalMessageBus(),
keyProvider,
)
require.NoError(t, err)
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)
}))
t.Cleanup(ts.Close)
t.Cleanup(func() { svc.DrainConnections(time.Millisecond, true) })
return &endpointStack{t: t, ts: ts, svc: svc}
}
func (s *endpointStack) wsURL() string {
return "ws" + strings.TrimPrefix(s.ts.URL, "http") + "/agent"
}
func (s *endpointStack) startWorker(target string, deployment string, endpoints []*livekit.AgentHttp_AgentEndpoint) *client.Worker {
w := client.New(client.Config{
ServerURL: s.wsURL(),
APIKey: testKey,
APISecret: testSecret,
AgentName: "test-agent",
Deployment: deployment,
Endpoints: endpoints,
TargetAddr: strings.TrimPrefix(target, "http://"),
})
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
require.NoError(s.t, w.Start(ctx))
s.t.Cleanup(w.Close)
return w
}
func (s *endpointStack) clientToken(t *testing.T) string {
at := auth.NewAccessToken(testKey, testSecret).SetVideoGrant(&auth.VideoGrant{RoomJoin: true, Room: "x"})
tok, err := at.ToJWT()
require.NoError(t, err)
return tok
}
func httpEP(path string, methods []string, public bool) *livekit.AgentHttp_AgentEndpoint {
return &livekit.AgentHttp_AgentEndpoint{Path: path, Methods: methods, Public: public}
}
// newTargetApp is the local app the worker bridges into; it never listens on a
// port reachable through the stack, only via the tunnel.
func newTargetApp(t *testing.T, mux *http.ServeMux) *httptest.Server {
app := httptest.NewServer(mux)
t.Cleanup(app.Close)
return app
}
func TestAgentEndpointsCorrectnessGate(t *testing.T) {
bigDown := make([]byte, 4<<20)
_, _ = rand.Read(bigDown)
mux := http.NewServeMux()
mux.HandleFunc("GET /json", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"ok":true}`))
})
mux.HandleFunc("POST /upload", func(w http.ResponseWriter, r *http.Request) {
sum := sha256.New()
n, err := io.Copy(sum, r.Body)
if err != nil {
http.Error(w, err.Error(), 500)
return
}
fmt.Fprintf(w, "%d:%x", n, sum.Sum(nil))
})
mux.HandleFunc("GET /big", func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write(bigDown)
})
mux.HandleFunc("GET /sse", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
f := w.(http.Flusher)
for i := 0; i < 5; i++ {
fmt.Fprintf(w, "data: event-%d\n\n", i)
f.Flush()
time.Sleep(150 * time.Millisecond)
}
})
app := newTargetApp(t, mux)
stack := newEndpointStack(t, agent.EndpointsConfig{})
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
httpEP("/json", []string{"GET"}, true),
httpEP("/upload", []string{"POST"}, true),
httpEP("/big", []string{"GET"}, true),
httpEP("/sse", []string{"GET"}, true),
})
base := stack.ts.URL + "/agents/production"
t.Run("json round trip", func(t *testing.T) {
resp, err := http.Get(base + "/json")
require.NoError(t, err)
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
require.Equal(t, 200, resp.StatusCode)
require.JSONEq(t, `{"ok":true}`, string(body))
})
t.Run("upload byte exact", func(t *testing.T) {
up := make([]byte, 8<<20)
_, _ = rand.Read(up)
resp, err := http.Post(base+"/upload", "application/octet-stream", bytes.NewReader(up))
require.NoError(t, err)
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
require.Equal(t, 200, resp.StatusCode)
require.Equal(t, fmt.Sprintf("%d:%x", len(up), sha256.Sum256(up)), string(body))
})
t.Run("download byte exact", func(t *testing.T) {
resp, err := http.Get(base + "/big")
require.NoError(t, err)
body, err := io.ReadAll(resp.Body)
resp.Body.Close()
require.NoError(t, err)
require.True(t, bytes.Equal(bigDown, body))
})
t.Run("sse events arrive one at a time", func(t *testing.T) {
start := time.Now()
resp, err := http.Get(base + "/sse")
require.NoError(t, err)
defer resp.Body.Close()
br := bufio.NewReader(resp.Body)
var arrivals []time.Duration
for {
line, err := br.ReadString('\n')
if err != nil {
break
}
if strings.HasPrefix(line, "data:") {
arrivals = append(arrivals, time.Since(start))
}
}
require.Len(t, arrivals, 5)
// incremental delivery: the first event arrives well before the last is
// even written (5 x 150ms); buffering the whole body would collapse gaps
require.Less(t, arrivals[0], 450*time.Millisecond)
require.Greater(t, arrivals[4]-arrivals[0], 300*time.Millisecond)
})
t.Run("32 concurrent requests", func(t *testing.T) {
errCh := make(chan error, 32)
for i := 0; i < 32; i++ {
go func() {
resp, err := http.Get(base + "/json")
if err == nil {
io.Copy(io.Discard, resp.Body)
resp.Body.Close()
if resp.StatusCode != 200 {
err = fmt.Errorf("status %d", resp.StatusCode)
}
}
errCh <- err
}()
}
for i := 0; i < 32; i++ {
require.NoError(t, <-errCh)
}
})
}
func TestAgentEndpointsStatusMapping(t *testing.T) {
mux := http.NewServeMux()
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) })
app := newTargetApp(t, mux)
stack := newEndpointStack(t, agent.EndpointsConfig{})
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
httpEP("/hook", []string{"POST"}, true),
httpEP("/private", []string{"GET"}, false),
})
base := stack.ts.URL + "/agents/production"
t.Run("404 unknown path", func(t *testing.T) {
resp, _ := http.Get(base + "/nope")
resp.Body.Close()
require.Equal(t, 404, resp.StatusCode)
})
t.Run("405 wrong method", func(t *testing.T) {
resp, _ := http.Get(base + "/hook")
resp.Body.Close()
require.Equal(t, 405, resp.StatusCode)
})
t.Run("401 non-public without token", func(t *testing.T) {
resp, _ := http.Get(base + "/private")
resp.Body.Close()
require.Equal(t, 401, resp.StatusCode)
})
t.Run("200 non-public with token", func(t *testing.T) {
req, _ := http.NewRequest("GET", base+"/private", nil)
req.Header.Set("Authorization", "Bearer "+stack.clientToken(t))
resp, err := http.DefaultClient.Do(req)
require.NoError(t, err)
resp.Body.Close()
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.Body.Close()
require.Equal(t, 503, resp.StatusCode)
})
t.Run("307 slash redirect", func(t *testing.T) {
c := &http.Client{CheckRedirect: func(req *http.Request, via []*http.Request) error {
return http.ErrUseLastResponse
}}
resp, err := c.Post(base+"/hook/", "text/plain", nil)
require.NoError(t, err)
resp.Body.Close()
require.Equal(t, 307, resp.StatusCode)
require.Equal(t, "/agents/production/hook", resp.Header.Get("Location"))
})
}
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{})
mux := http.NewServeMux()
mux.HandleFunc("GET /drip", func(w http.ResponseWriter, r *http.Request) {
f := w.(http.Flusher)
buf := make([]byte, 64<<10)
for {
if _, err := w.Write(buf); err != nil {
return
}
f.Flush()
select {
case <-r.Context().Done():
return
default:
}
}
})
mux.HandleFunc("GET /quick", func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte("ok"))
})
app := newTargetApp(t, mux)
stack := newEndpointStack(t, agent.EndpointsConfig{DataConnCount: 1})
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
httpEP("/drip", []string{"GET"}, true),
httpEP("/quick", []string{"GET"}, true),
})
base := stack.ts.URL + "/agents/production"
// a stalled client: open /drip, read a little, then stop reading entirely
resp, err := http.Get(base + "/drip")
require.NoError(t, err)
defer resp.Body.Close()
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)
time.Sleep(500 * time.Millisecond)
// sibling streams on the SAME connection must proceed at full speed
for i := 0; i < 5; i++ {
start := time.Now()
q, err := http.Get(base + "/quick")
require.NoError(t, err)
body, _ := io.ReadAll(q.Body)
q.Body.Close()
require.Equal(t, "ok", string(body))
require.Less(t, time.Since(start), 2*time.Second,
"sibling stream stalled behind a blocked heavy stream")
}
}
func TestAgentEndpointsRetrySafety(t *testing.T) {
var hits atomic.Int32
mux := http.NewServeMux()
mux.HandleFunc("GET /json", func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
_, _ = w.Write([]byte("ok"))
})
app := newTargetApp(t, mux)
// a worker whose local app is unreachable REFUSES streams; the front must
// retry on the healthy worker exactly once
stack := newEndpointStack(t, agent.EndpointsConfig{})
deadTarget := "127.0.0.1:1" // nothing listens
eps := []*livekit.AgentHttp_AgentEndpoint{httpEP("/json", []string{"GET"}, true)}
broken := client.New(client.Config{
ServerURL: stack.wsURL(), APIKey: testKey, APISecret: testSecret,
AgentName: "test-agent", Deployment: "production",
Endpoints: eps, TargetAddr: deadTarget,
})
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
require.NoError(t, broken.Start(ctx))
t.Cleanup(broken.Close)
stack.startWorker(app.URL, "production", eps)
// 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")
require.NoError(t, err)
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
if resp.StatusCode == 200 && string(body) == "ok" {
okCount++
}
}
require.Equal(t, 12, okCount, "REFUSED streams must fall through to the healthy worker")
require.EqualValues(t, 12, hits.Load())
}
func TestAgentEndpointsTruncationAborts(t *testing.T) {
// a worker connection dying mid-response must abort the client connection,
// never expose a clean-looking short body
mux := http.NewServeMux()
release := make(chan struct{})
mux.HandleFunc("GET /partial", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Length", "1000000")
_, _ = w.Write(make([]byte, 1000))
w.(http.Flusher).Flush()
<-release
})
app := newTargetApp(t, mux)
stack := newEndpointStack(t, agent.EndpointsConfig{})
w := stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
httpEP("/partial", []string{"GET"}, true),
})
resp, err := http.Get(stack.ts.URL + "/agents/production/partial")
require.NoError(t, err)
defer resp.Body.Close()
head := make([]byte, 1000)
_, err = io.ReadFull(resp.Body, head)
require.NoError(t, err)
w.Close() // kill the worker mid-stream
close(release)
_, err = io.ReadAll(resp.Body)
require.Error(t, err, "truncated response must not read as clean EOF")
}
func TestAgentEndpointsNoLocalListenerContract(t *testing.T) {
// the front never routes undeclared paths: the worker-local health/info
// routes are unreachable by construction
mux := http.NewServeMux()
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("local")) })
app := newTargetApp(t, mux)
stack := newEndpointStack(t, agent.EndpointsConfig{})
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
httpEP("/declared", []string{"GET"}, true),
})
for _, path := range []string{"/", "/worker", "/undeclared"} {
resp, err := http.Get(stack.ts.URL + "/agents/production" + path)
require.NoError(t, err)
resp.Body.Close()
require.Equal(t, 404, resp.StatusCode, path)
}
}
// dial helper kept for upgrade tests once the SDK lands; avoids unused imports
var _ = net.Dialer{}
// the upgrade exchange and the raw bidirectional session both ride one stream;
// this guards the front's hijack path and the no-half-close rule for upgrades
func TestAgentEndpointsWebSocketUpgrade(t *testing.T) {
upgrader := websocket.Upgrader{}
mux := http.NewServeMux()
mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
c, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
defer c.Close()
for {
mt, msg, err := c.ReadMessage()
if err != nil {
return
}
if err := c.WriteMessage(mt, append([]byte("echo:"), msg...)); err != nil {
return
}
}
})
app := newTargetApp(t, mux)
stack := newEndpointStack(t, agent.EndpointsConfig{})
stack.startWorker(app.URL, "production", []*livekit.AgentHttp_AgentEndpoint{
{Path: "/ws", Kind: livekit.AgentHttp_AEK_WEBSOCKET, Public: true},
})
wsURL := "ws" + strings.TrimPrefix(stack.ts.URL, "http") + "/agents/production/ws"
c, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
require.NoError(t, err)
if resp != nil {
resp.Body.Close()
}
defer c.Close()
for i := 0; i < 5; i++ {
payload := fmt.Sprintf("msg-%d", i)
require.NoError(t, c.WriteMessage(websocket.TextMessage, []byte(payload)))
_, echoed, err := c.ReadMessage()
require.NoError(t, err)
require.Equal(t, "echo:"+payload, string(echoed))
}
// a larger frame exercises credit flow through the raw pump
big := bytes.Repeat([]byte("x"), 256<<10)
require.NoError(t, c.WriteMessage(websocket.BinaryMessage, big))
_, echoed, err := c.ReadMessage()
require.NoError(t, err)
require.Equal(t, len(big)+5, len(echoed))
}
+253 -5
View File
@@ -26,9 +26,11 @@ 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"
"github.com/livekit/livekit-server/pkg/agent/endpoint"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/routing"
"github.com/livekit/livekit-server/pkg/rtc"
@@ -112,8 +114,11 @@ func DispatchAgentWorkerSignal(c agent.SignalConn, h agent.WorkerSignalHandler,
return true
}
func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, registration agent.WorkerRegistration, l logger.Logger) (r agent.WorkerRegistration, ok bool) {
func HandshakeAgentWorker(c agent.SignalConn, serverInfo *livekit.ServerInfo, registration agent.WorkerRegistration, l logger.Logger, opts ...func(*agent.WorkerRegisterer)) (r agent.WorkerRegistration, ok bool) {
wr := agent.NewWorkerRegisterer(c, serverInfo, registration)
for _, opt := range opts {
opt(wr)
}
if err := c.SetReadDeadline(wr.Deadline()); err != nil {
return
}
@@ -138,8 +143,12 @@ type AgentService struct {
type AgentHandler struct {
agentServer rpc.AgentInternalServer
mu sync.Mutex
logger logger.Logger
// the server's only configured api key, when exactly one exists: the
// unauthenticated scope for public endpoints regardless of which node holds
// the workers
singleAPIKey string
mu sync.Mutex
logger logger.Logger
serverInfo *livekit.ServerInfo
workers map[string]*agent.Worker
@@ -147,6 +156,9 @@ type AgentHandler struct {
keyProvider auth.KeyProvider
targetLoad float32
endpointRegistry *endpoint.Registry
endpointsConfig agent.EndpointsConfig
namespaceWorkers map[workerKey][]*agent.Worker
roomKeyCount int
publisherKeyCount int
@@ -199,10 +211,34 @@ func NewAgentService(
agent.PublisherAgentTopic,
agent.ParticipantAgentTopic,
)
s.AgentHandler.endpointsConfig = conf.Agents.Endpoints
if len(conf.Keys) == 1 {
for key := range conf.Keys {
s.AgentHandler.singleAPIKey = key
}
}
return s, nil
}
// EndpointFront is the /agents/{deployment}/{path...} handler backed by this
// node's attached workers. Project scope comes from validated grants when a
// token is present; unauthenticated requests reach public endpoints only.
func (s *AgentService) EndpointFront() http.Handler {
front := endpoint.NewFront(s.endpointRegistry, func(r *http.Request) (string, bool) {
if claims := GetGrants(r.Context()); claims != nil {
return GetAPIKey(r.Context()), true
}
// unauthenticated: with a single configured key the scope is
// unambiguous even when this node holds no registrations (multi-node)
return s.singleAPIKey, false
}, s.logger)
front.WithSingleKeyFallback()
return front
}
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 {
// bound the size of a single signalling frame so an oversized message is
// rejected by the transport before being fully buffered in memory. This
@@ -211,13 +247,157 @@ func (s *AgentService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if s.signalMessageSizeLimit > 0 {
conn.SetReadLimit(s.signalMessageSizeLimit)
}
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
HandleEndpointAttach(s.endpointRegistry, NewEndpointWireConn(conn), s.wireParams())
return
}
sigConn := NewWSSignalConnection(conn, s.signalMessageSizeLimit)
defer sigConn.Close()
s.HandleConnection(r.Context(), sigConn, registration)
}
}
// 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
}
// 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 {
b, err := proto.Marshal(f)
if err != nil {
return err
}
c.writeMu.Lock()
defer c.writeMu.Unlock()
return c.ws.WriteMessage(websocket.BinaryMessage, b)
}
func (c *endpointWireConn) ReadFrame() (*livekit.AgentHttp_Frame, error) {
for {
_ = c.ws.SetReadDeadline(time.Now().Add(endpointWireIdleTimeout))
mt, b, err := c.ws.ReadMessage()
if err != nil {
return nil, err
}
if mt != websocket.BinaryMessage {
continue
}
f := &livekit.AgentHttp_Frame{}
if err := proto.Unmarshal(b, f); err != nil {
return nil, err
}
return f, nil
}
}
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.
func (s *AgentService) wireParams() endpoint.WireParams {
p := endpoint.WireParams{
CreditWindow: s.endpointsConfig.CreditWindow,
ConnectionWindow: s.endpointsConfig.ConnectionWindow,
MaxFrameSize: s.endpointsConfig.MaxFrameSize,
MaxStreamsPerConn: s.endpointsConfig.MaxStreamsPerConn,
}.WithDefaults()
// data frames ride the same websocket read limit as signalling: keep the
// frame size safely under it so a full frame can never kill the wire
if lim := s.signalMessageSizeLimit; lim > 0 && int64(p.MaxFrameSize) > lim/2 {
p.MaxFrameSize = uint32(lim / 2)
}
return p
}
// 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) {
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(), 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,
@@ -240,11 +420,37 @@ func NewAgentHandler(
roomTopic: roomTopic,
publisherTopic: publisherTopic,
participantTopic: participantTopic,
endpointRegistry: endpoint.NewRegistry(),
}
}
// endpointSettings validates a registration's manifest and negotiates the
// data-plane settings, minting the epoch's attach token.
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")
}
if _, err := endpoint.ParseManifest(req.GetEndpoints()); err != nil {
return nil, err
}
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
}
func (h *AgentHandler) HandleConnection(ctx context.Context, conn agent.SignalConn, registration agent.WorkerRegistration) {
registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger)
registration, ok := HandshakeAgentWorker(conn, h.serverInfo, registration, h.logger, func(wr *agent.WorkerRegisterer) {
wr.WithEndpointSettings(h.endpointSettings)
})
if !ok {
return
}
@@ -255,15 +461,57 @@ 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)
handlerWorker := &agentHandlerWorker{h, worker}
for ok := true; ok; {
ok = DispatchAgentWorkerSignal(conn, handlerWorker, worker.Logger())
}
if endpointReg != nil {
h.endpointRegistry.Deregister(endpointReg)
}
h.deregisterWorker(worker)
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 {
return nil
}
manifest, err := endpoint.ParseManifest(w.Endpoints)
if err != nil {
// validated during the handshake; a failure here is a programming error
w.Logger().Errorw("endpoint manifest failed to re-parse", err)
return nil
}
reg := &endpoint.Registration{
WorkerID: w.ID,
InstanceID: w.InstanceID,
APIKey: w.APIKey(),
Deployment: w.Deployment,
Manifest: manifest,
Settings: endpoint.Settings{
Protocol: settings.GetProtocol(),
AttachToken: settings.GetAttachToken(),
DataConnCount: settings.GetDataConnectionCount(),
},
Logger: w.Logger(),
Load: w.Load,
Draining: w.Draining,
}
if err := h.endpointRegistry.Register(reg); err != nil {
w.Logger().Errorw("failed to register endpoints", err)
return nil
}
w.Logger().Infow("endpoints registered",
"deployment", w.Deployment, "routes", len(w.Endpoints))
return reg
}
func (h *AgentHandler) registerWorker(w *agent.Worker) {
h.mu.Lock()
+6 -2
View File
@@ -76,8 +76,12 @@ func (m *APIKeyAuthMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request,
authToken = authHeader[len(bearerPrefix):]
} else {
// attempt to find from request header
authToken = r.FormValue(accessTokenParam)
// attempt to find from the query string. FormValue would also parse
// url-encoded POST bodies, consuming the body of any request that gets
// proxied further (agent HTTP endpoints)
if r.URL != nil {
authToken = r.URL.Query().Get(accessTokenParam)
}
}
if authToken != "" {
+2
View File
@@ -40,6 +40,7 @@ import (
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/utils/xtwirp"
"github.com/livekit/livekit-server/pkg/agent/endpoint"
"github.com/livekit/livekit-server/pkg/config"
"github.com/livekit/livekit-server/pkg/routing"
"github.com/livekit/livekit-server/version"
@@ -150,6 +151,7 @@ func NewLivekitServer(conf *config.Config,
rtcService.SetupRoutes(mux)
whipService.SetupRoutes(mux)
mux.Handle("/agent", agentService)
mux.Handle(endpoint.PathPrefix, agentService.EndpointFront())
mux.HandleFunc("/", s.defaultHandler)
s.httpServer = &http.Server{