fix: key STS token cache by acting subject, not session alone

go/adk/pkg/a2a/executor.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"iter"
 	"os"
-	"strings"
 
 	a2atype "github.com/a2aproject/a2a-go/v2/a2a"
 	"github.com/a2aproject/a2a-go/v2/a2asrv"
@@ -273,23 +272,11 @@ func extractSessionName(message *a2atype.Message) string {
 // withBearerToken extracts the Bearer token from the incoming A2A request's
 // Authorization header and stores it in ctx for API key passthrough.
 func withBearerToken(ctx context.Context) context.Context {
-	callCtx, ok := a2asrv.CallContextFrom(ctx)
-	if !ok {
+	token := models.BearerFromCallContext(ctx)
+	if token == "" {
 		return ctx
 	}
-	meta := callCtx.ServiceParams()
-	if meta == nil {
-		return ctx
-	}
-	vals, ok := meta.Get("authorization")
-	if !ok || len(vals) == 0 || vals[0] == "" {
-		return ctx
-	}
-	parts := strings.Fields(strings.TrimSpace(vals[0]))
-	if len(parts) >= 2 && strings.EqualFold(parts[0], "Bearer") {
-		return context.WithValue(ctx, models.BearerTokenKey, parts[1])
-	}
-	return ctx
+	return context.WithValue(ctx, models.BearerTokenKey, token)
 }
 
 // dropPreAppendedDecisionFromHistory removes a pre-appended HITL decision

go/adk/pkg/models/base.go

@@ -1,13 +1,16 @@
 package models
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net"
 	"net/http"
 	"strings"
 	"time"
 
+	"github.com/a2aproject/a2a-go/v2/a2asrv"
+	"github.com/kagent-dev/kagent/go/adk/pkg/constants"
 	"google.golang.org/genai"
 )
 
@@ -80,6 +83,39 @@ var BearerTokenKey = &contextKey{}
 
 type contextKey struct{}
 
+// BearerFromCallContext returns the bearer token carried by the A2A call
+// context's Authorization header, or "" when there is none.
+func BearerFromCallContext(ctx context.Context) string {
+	callCtx, ok := a2asrv.CallContextFrom(ctx)
+	if !ok {
+		return ""
+	}
+	meta := callCtx.ServiceParams()
+	if meta == nil {
+		return ""
+	}
+	vals, ok := meta.Get(constants.AuthorizationHeader)
+	if !ok || len(vals) == 0 {
+		return ""
+	}
+	parts := strings.Fields(strings.TrimSpace(vals[0]))
+	if len(parts) >= 2 && strings.EqualFold(parts[0], "Bearer") {
+		return parts[1]
+	}
+	return ""
+}
+
+// BearerTokenFromContext returns the credential the request authenticates with.
+// It prefers the value stored under BearerTokenKey and falls back to the A2A
+// call context, which reaches callers whose context was not threaded through
+// the executor.
+func BearerTokenFromContext(ctx context.Context) string {
+	if token, ok := ctx.Value(BearerTokenKey).(string); ok && token != "" {
+		return token
+	}
+	return BearerFromCallContext(ctx)
+}
+
 // headerTransport wraps an http.RoundTripper and adds custom headers to all requests
 type headerTransport struct {
 	base    http.RoundTripper

go/adk/pkg/sts/integration.go

@@ -9,6 +9,11 @@ import (
 // GetSubjectTokenFunc is a function type for extracting subject tokens.
 // It receives the bearer token (from Authorization header) and should return
 // the subject token for STS exchange, or empty string if not available.
+//
+// It must be a pure function of bearerToken. TokenPropagationPlugin caches the
+// exchange under a hash of the bearer, so an implementation that mints or
+// fetches a token returns one value per call while the cache keeps serving the
+// first for the entry's lifetime.
 type GetSubjectTokenFunc func(bearerToken string) string
 
 // DefaultGetSubjectToken extracts the JWT token from the Authorization header.