diff --git a/cmd/agent-endpoint-client/main.go b/cmd/agent-endpoint-client/main.go new file mode 100644 index 000000000..8755a212e --- /dev/null +++ b/cmd/agent-endpoint-client/main.go @@ -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() +} diff --git a/go.mod b/go.mod index d6a026553..bf3bd71d5 100644 --- a/go.mod +++ b/go.mod @@ -21,7 +21,7 @@ require ( github.com/jxskiss/base62 v1.1.0 github.com/livekit/mageutil v0.0.0-20250511045019-0f1ff63f7731 github.com/livekit/mediatransportutil v0.0.0-20260608063931-a3417d38cda0 - github.com/livekit/protocol v1.50.5-0.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 diff --git a/go.sum b/go.sum index 321500c41..daf981d12 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/pkg/agent/client.go b/pkg/agent/client.go index 788da9846..383729f29 100644 --- a/pkg/agent/client.go +++ b/pkg/agent/client.go @@ -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 { diff --git a/pkg/agent/config.go b/pkg/agent/config.go index ee529884e..8fd32b41e 100644 --- a/pkg/agent/config.go +++ b/pkg/agent/config.go @@ -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"` } diff --git a/pkg/agent/endpoint/client/client.go b/pkg/agent/endpoint/client/client.go new file mode 100644 index 000000000..9d2dee2a5 --- /dev/null +++ b/pkg/agent/endpoint/client/client.go @@ -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() +} diff --git a/pkg/agent/endpoint/conn.go b/pkg/agent/endpoint/conn.go new file mode 100644 index 000000000..9473b6ab3 --- /dev/null +++ b/pkg/agent/endpoint/conn.go @@ -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 + } + } +} diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go new file mode 100644 index 000000000..4754b4b2f --- /dev/null +++ b/pkg/agent/endpoint/front.go @@ -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 +} diff --git a/pkg/agent/endpoint/manifest.go b/pkg/agent/endpoint/manifest.go new file mode 100644 index 000000000..6d827d448 --- /dev/null +++ b/pkg/agent/endpoint/manifest.go @@ -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 +} diff --git a/pkg/agent/endpoint/registry.go b/pkg/agent/endpoint/registry.go new file mode 100644 index 000000000..33dc1c2e2 --- /dev/null +++ b/pkg/agent/endpoint/registry.go @@ -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 != "" +} diff --git a/pkg/agent/endpoint/registry_test.go b/pkg/agent/endpoint/registry_test.go new file mode 100644 index 000000000..fb3e3a67a --- /dev/null +++ b/pkg/agent/endpoint/registry_test.go @@ -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")) +} diff --git a/pkg/agent/endpoint/scheduler.go b/pkg/agent/endpoint/scheduler.go new file mode 100644 index 000000000..682767674 --- /dev/null +++ b/pkg/agent/endpoint/scheduler.go @@ -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 + } + } +} diff --git a/pkg/agent/endpoint/stream.go b/pkg/agent/endpoint/stream.go new file mode 100644 index 000000000..3ec2dc189 --- /dev/null +++ b/pkg/agent/endpoint/stream.go @@ -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() +} diff --git a/pkg/agent/endpoint/template.go b/pkg/agent/endpoint/template.go new file mode 100644 index 000000000..62008de02 --- /dev/null +++ b/pkg/agent/endpoint/template.go @@ -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 +} diff --git a/pkg/agent/endpoint/template_test.go b/pkg/agent/endpoint/template_test.go new file mode 100644 index 000000000..07eb2db75 --- /dev/null +++ b/pkg/agent/endpoint/template_test.go @@ -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") +} diff --git a/pkg/agent/endpoint/wire.go b/pkg/agent/endpoint/wire.go new file mode 100644 index 000000000..5abfaff6e --- /dev/null +++ b/pkg/agent/endpoint/wire.go @@ -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, + } +} diff --git a/pkg/agent/worker.go b/pkg/agent/worker.go index b93c33437..ab8f50d66 100644 --- a/pkg/agent/worker.go +++ b/pkg/agent/worker.go @@ -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 diff --git a/pkg/service/agentendpoint_test.go b/pkg/service/agentendpoint_test.go new file mode 100644 index 000000000..c6d11f2c5 --- /dev/null +++ b/pkg/service/agentendpoint_test.go @@ -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)) +} diff --git a/pkg/service/agentservice.go b/pkg/service/agentservice.go index b5128cb9e..9d338c947 100644 --- a/pkg/service/agentservice.go +++ b/pkg/service/agentservice.go @@ -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() diff --git a/pkg/service/auth.go b/pkg/service/auth.go index 9d8c1e921..53acb11d8 100644 --- a/pkg/service/auth.go +++ b/pkg/service/auth.go @@ -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 != "" { diff --git a/pkg/service/server.go b/pkg/service/server.go index e26569804..38aaf52e2 100644 --- a/pkg/service/server.go +++ b/pkg/service/server.go @@ -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{