From cb4483539fee28b98c3df78dd0e2791f1fd23021 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9o=20Monnom?= Date: Tue, 6 Oct 2026 01:49:17 -0700 Subject: [PATCH] agent endpoint: optional session hook for text-kind routes FrontParams.ResolveSession is a SessionResolver the embedder may install, mirroring Identity: cloud installs one backed by the session registry and OSS leaves it nil. It runs once per request, only when a matched declaration is AEK_TEXT, with the request, the agent and deployment, and a bounded prefix of the body (SessionPeek, 4 KiB by default) read ahead and spliced back so the worker, or a fallback, still gets the whole body. The id it reports goes into the preamble's session_id for the text-kind declarations; with no resolver, or none reported, the route is forwarded unchanged and the worker resolves the session from the body itself. --- pkg/agent/endpoint/attempt.go | 7 +++ pkg/agent/endpoint/front.go | 80 +++++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+) diff --git a/pkg/agent/endpoint/attempt.go b/pkg/agent/endpoint/attempt.go index 9f799ce75..10e1e854b 100644 --- a/pkg/agent/endpoint/attempt.go +++ b/pkg/agent/endpoint/attempt.go @@ -43,6 +43,9 @@ type attempt struct { target string requestID string granted bool + // sessionID is the agent session the front placed the request in; empty + // when it placed none. Only a text-kind declaration is told. + sessionID string body io.Reader pools *bridgePools preamble *livekit.AgentHttp_StreamPreamble @@ -111,6 +114,10 @@ func (a *attempt) newPreamble() *livekit.AgentHttp_StreamPreamble { func (a *attempt) setRoute(w routeWorker) { a.preamble.Route = w.raw a.preamble.Kind = w.kind + a.preamble.SessionId = "" + if w.kind == livekit.AgentHttp_AEK_TEXT { + a.preamble.SessionId = a.sessionID + } } // requestHeader builds the header set the worker sees. diff --git a/pkg/agent/endpoint/front.go b/pkg/agent/endpoint/front.go index 3f8d6c22c..29a4b5482 100644 --- a/pkg/agent/endpoint/front.go +++ b/pkg/agent/endpoint/front.go @@ -16,6 +16,7 @@ package endpoint import ( "bufio" + "bytes" "context" "errors" "fmt" @@ -59,6 +60,9 @@ const ( // worker cannot hold a client open by trickling them forever. maxInformationalHeads = 8 + // DefaultSessionPeek bounds the request body prefix a SessionResolver sees. + DefaultSessionPeek = 4 << 10 + maxResponseHeadSize = 1 << 20 responseBufSize = 8 << 10 ) @@ -132,6 +136,16 @@ type Front struct { // false leaves them to the URL. type Identity func(r *http.Request) (agentName, deployment string, ok bool) +// SessionResolver places a request on an AEK_TEXT route in an agent session. +// It runs once per request, after access is resolved and before a worker is +// picked, and only when a matched declaration is of the text kind: that kind is +// the worker's consent to the front reading its body. body is the first bytes +// of the request body, at most SessionPeek; the worker receives the whole body +// regardless. Reporting false leaves the session to the worker, which resolves +// it from the body itself. Cloud installs one backed by the session registry; +// OSS leaves it nil, and a text route is then forwarded unchanged. +type SessionResolver func(r *http.Request, agentName, deployment string, body []byte) (sessionID string, ok bool) + // FrontParams configures a Front. Fields are read on every request once the // Front is serving, so none may change after construction. type FrontParams struct { @@ -139,6 +153,13 @@ type FrontParams struct { Logger logger.Logger Identity Identity + // ResolveSession places text-kind requests in an agent session. nil + // forwards them unchanged. + ResolveSession SessionResolver + // SessionPeek bounds the body prefix ResolveSession sees. 0 means + // DefaultSessionPeek. + SessionPeek int + // MaxTunnelsPerWorker caps one worker's concurrent WebSocket tunnels. 0 // means half its session's MaxStreams. MaxTunnelsPerWorker int @@ -252,6 +273,21 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } + var sessionID string + if f.params.ResolveSession != nil && declaresText(matched) { + peek, err := peekBody(r, f.sessionPeek()) + if err != nil { + // the client's body failed before any worker saw the request, so + // there is nothing to place and nothing to retry + f.params.Logger.Debugw("agent endpoint request body failed before dispatch", "error", err, "requestID", reqID) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + if id, ok := f.params.ResolveSession(r, agentName, deployment, peek); ok { + sessionID = id + } + } + var bodyConsumed atomic.Int64 a := &attempt{ req: r, @@ -259,6 +295,7 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { target: requestTarget(escPath, r.URL.RawQuery), requestID: reqID, granted: granted, + sessionID: sessionID, webSocket: isWebSocketUpgrade(r), pools: f.pools, } @@ -320,6 +357,49 @@ func (f *Front) ServeHTTP(w http.ResponseWriter, r *http.Request) { f.writeUnavailable(w, "no worker could serve the request") } +func (f *Front) sessionPeek() int { + if f.params.SessionPeek > 0 { + return f.params.SessionPeek + } + return DefaultSessionPeek +} + +// declaresText reports whether any declaration the request matched is of the +// text kind. Declarers of one route may disagree, so the session is resolved +// once and handed only to the workers that asked for it (setRoute). +func declaresText(matched []routeWorker) bool { + for _, w := range matched { + if w.kind == livekit.AgentHttp_AEK_TEXT { + return true + } + } + return false +} + +// peekBody reads the first n bytes of the request body, waiting for them or +// for the body to end, and splices them back in front of the rest, so r.Body +// still reads whole for the worker or a fallback. A body shorter than declared +// is not an error here; the framing reports it to the worker as usual. +func peekBody(r *http.Request, n int) ([]byte, error) { + cl := bodyLength(r) + if cl == 0 { + return nil, nil + } + if cl > 0 && cl < int64(n) { + n = int(cl) + } + buf := make([]byte, n) + read, err := io.ReadFull(r.Body, buf) + if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) { + return nil, err + } + r.Body = struct { + io.Reader + io.Closer + }{io.MultiReader(bytes.NewReader(buf[:read]), r.Body), r.Body} + return buf[:read], nil +} + // matchDeployment resolves a path against the deployment's merged route table. // matched is the dispatch test: a route can be decided while nothing is left to // serve it, and an undecidable table dispatches with no route at all.