diff --git a/pkg/agent/config.go b/pkg/agent/config.go index 1bbcfbc25..30be1ceb1 100644 --- a/pkg/agent/config.go +++ b/pkg/agent/config.go @@ -30,4 +30,12 @@ type EndpointsConfig struct { // TunnelDrainTimeout is how long a WebSocket tunnel outlives its worker or // this node starting to drain; 0 takes the endpoint package default. TunnelDrainTimeout time.Duration `yaml:"tunnel_drain_timeout,omitempty"` + + // ThreadsAddress is the AgentThreads gRPC service (host:port) the front + // resolves a text route's thread with, so the request is placed before a + // worker sees it; empty leaves every thread to the worker. + ThreadsAddress string `yaml:"threads_address,omitempty"` + // ThreadLookupTimeout bounds one lookup; 0 takes the endpoint package + // default. + ThreadLookupTimeout time.Duration `yaml:"thread_lookup_timeout,omitempty"` } diff --git a/pkg/agent/endpoint/threads.go b/pkg/agent/endpoint/threads.go new file mode 100644 index 000000000..4642d3a1e --- /dev/null +++ b/pkg/agent/endpoint/threads.go @@ -0,0 +1,121 @@ +// 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 ( + "bytes" + "context" + "encoding/json" + "net/http" + "strings" + "time" + + "github.com/livekit/protocol/livekit/agent" + "github.com/livekit/protocol/logger" + "google.golang.org/grpc" + "google.golang.org/grpc/metadata" +) + +// getThreadMethod is AgentThreads.GetThread of cloud-protocol, called by path +// with livekit/protocol's messages so this package depends on nothing private. +const getThreadMethod = "/cloud_protocol.AgentThreads/GetThread" + +// DefaultThreadLookupTimeout bounds one GetThread on the request path. +const DefaultThreadLookupTimeout = 2 * time.Second + +// NewThreadResolver resolves the thread of an A2A request through the +// AgentThreads gRPC service at conn. The request names its thread as the A2A +// contextId of the message in its body; one that names none is left to the +// worker, which creates it. The lookup runs under the caller's own token, so +// the service applies the caller's scope, and follows a redirect once, so a +// merged thread places the request where the conversation continues. +func NewThreadResolver(conn grpc.ClientConnInterface, timeout time.Duration, log logger.Logger) ThreadResolver { + if timeout <= 0 { + timeout = DefaultThreadLookupTimeout + } + return func(r *http.Request, agentName, deployment string, body []byte) (string, bool) { + contextID := contextIDOf(body) + if contextID == "" { + return "", false + } + token := requestToken(r) + if token == "" { + return "", false + } + ctx, cancel := context.WithTimeout(r.Context(), timeout) + defer cancel() + ctx = metadata.AppendToOutgoingContext(ctx, "authorization", "Bearer "+token) + + thread, err := getThread(ctx, conn, contextID) + if err == nil && thread.GetRedirectTo() != "" { + thread, err = getThread(ctx, conn, thread.GetRedirectTo()) + } + if err != nil { + // the worker resolves it from the body itself, and answers the caller + // with the service's verdict + log.Debugw("agent thread lookup failed; left to the worker", "error", err, "agentName", agentName, "deployment", deployment) + return "", false + } + if thread.GetAgentName() != agentName { + return "", false + } + return thread.GetThreadId(), true + } +} + +func getThread(ctx context.Context, conn grpc.ClientConnInterface, threadID string) (*agent.AgentThread_Thread, error) { + req := &agent.AgentThread_GetRequest{ + Lookup: &agent.AgentThread_GetRequest_ThreadId{ThreadId: threadID}, + } + thread := &agent.AgentThread_Thread{} + if err := conn.Invoke(ctx, getThreadMethod, req, thread); err != nil { + return nil, err + } + return thread, nil +} + +// contextIDOf is the message's contextId in a JSON-RPC A2A body prefix: the +// first string under that key, read token by token so a prefix cut mid-document +// still yields it when it came before the cut. +func contextIDOf(body []byte) string { + dec := json.NewDecoder(bytes.NewReader(body)) + var key bool + for { + tok, err := dec.Token() + if err != nil { + // io.EOF, or the cut: whatever came before is all there is + return "" + } + s, ok := tok.(string) + if !ok { + key = false + continue + } + if key { + return s + } + key = s == "contextId" + } +} + +// requestToken is the caller's LiveKit token as the front reads it: the +// bearer, else access_token on the query for a browser that cannot set +// headers. Empty when there is none. +func requestToken(r *http.Request) string { + if tok, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer "); ok && tok != "" { + return tok + } + return r.URL.Query().Get("access_token") +} diff --git a/pkg/agent/endpoint/threads_test.go b/pkg/agent/endpoint/threads_test.go new file mode 100644 index 000000000..2d808f9ef --- /dev/null +++ b/pkg/agent/endpoint/threads_test.go @@ -0,0 +1,118 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package endpoint + +import ( + "context" + "net" + "net/http" + "strings" + "testing" + + "github.com/livekit/protocol/livekit/agent" + "github.com/livekit/protocol/logger" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +// fakeThreads answers GetThread for the threads it holds, by id, and records +// the authorization it was called with. It serves the method by path, the way +// the resolver calls it, so no generated service stub is needed. +type fakeThreads struct { + threads map[string]*agent.AgentThread_Thread + auth []string +} + +func (f *fakeThreads) serve(t *testing.T) *grpc.ClientConn { + t.Helper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + srv := grpc.NewServer(grpc.UnknownServiceHandler(func(_ any, stream grpc.ServerStream) error { + method, _ := grpc.MethodFromServerStream(stream) + if method != getThreadMethod { + return status.Error(codes.Unimplemented, method) + } + md, _ := metadata.FromIncomingContext(stream.Context()) + f.auth = append(f.auth, strings.Join(md.Get("authorization"), ",")) + req := &agent.AgentThread_GetRequest{} + if err := stream.RecvMsg(req); err != nil { + return err + } + thread, ok := f.threads[req.GetThreadId()] + if !ok { + return status.Error(codes.NotFound, req.GetThreadId()) + } + return stream.SendMsg(thread) + })) + go srv.Serve(lis) + t.Cleanup(srv.Stop) + conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + return conn +} + +// a body naming a thread places the request on it, a merged thread on the one +// it continues in, under the caller's own token; a body naming none, a thread +// of another agent, or a lookup that fails is left to the worker. +func TestThreadResolverPlacesByContextIdUnderTheCallersToken(t *testing.T) { + fake := &fakeThreads{threads: map[string]*agent.AgentThread_Thread{ + "AT_1": {ThreadId: "AT_1", AgentName: "a"}, + "AT_2": {ThreadId: "AT_2", AgentName: "a", RedirectTo: "AT_3"}, + "AT_3": {ThreadId: "AT_3", AgentName: "a"}, + "AT_9": {ThreadId: "AT_9", AgentName: "other"}, + }} + resolve := NewThreadResolver(fake.serve(t), 0, logger.GetLogger()) + body := func(id string) []byte { + return []byte(`{"jsonrpc":"2.0","method":"message/send","params":{"message":{"messageId":"m","contextId":"` + id + `","parts":[{"kind":"text","text":"hi"}]}}}`) + } + req := func(tok string) *http.Request { + r, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "http://node/agents/a/d/a2a", nil) + require.NoError(t, err) + if tok != "" { + r.Header.Set("Authorization", "Bearer "+tok) + } + return r + } + + id, ok := resolve(req("tok"), "a", "d", body("AT_1")) + require.True(t, ok) + require.Equal(t, "AT_1", id) + require.Equal(t, []string{"Bearer tok"}, fake.auth) + + id, ok = resolve(req("tok"), "a", "d", body("AT_2")) + require.True(t, ok, "a merged thread is followed once") + require.Equal(t, "AT_3", id) + + _, ok = resolve(req("tok"), "a", "d", body("AT_missing")) + require.False(t, ok, "an unknown thread is the worker's to refuse") + _, ok = resolve(req("tok"), "a", "d", body("AT_9")) + require.False(t, ok, "another agent's thread is not placed") + _, ok = resolve(req("tok"), "a", "d", []byte(`{"jsonrpc":"2.0","params":{"message":{"messageId":"m"}}}`)) + require.False(t, ok, "no contextId: the worker creates the thread") + _, ok = resolve(req(""), "a", "d", body("AT_1")) + require.False(t, ok, "no token to look it up under") + + // the peek may cut the body: a contextId before the cut still places it + cut := body("AT_1")[:100] + require.Contains(t, string(cut), `"contextId":"AT_1"`) + id, ok = resolve(req("tok"), "a", "d", cut) + require.True(t, ok) + require.Equal(t, "AT_1", id) +}