diff --git a/go/adk/pkg/runner/adapter.go b/go/adk/pkg/runner/adapter.go index b769d3fe04..40a7bf8087 100644 --- a/go/adk/pkg/runner/adapter.go +++ b/go/adk/pkg/runner/adapter.go @@ -127,7 +127,7 @@ func buildTokenPropagationPlugin(ctx context.Context, log logr.Logger) (*sts.Tok // Propagate-only mode: keep parity with Python by enabling plugin without STS exchange. if stsWellKnownURI == "" { log.Info("Enabling token propagation plugin without STS exchange") - return sts.NewTokenPropagationPlugin(nil, log), nil + return sts.NewTokenPropagationPlugin(nil, log, nil, nil), nil } defaultSTSConfig := sts.DefaultSTSConfig(stsWellKnownURI) @@ -144,6 +144,23 @@ func buildTokenPropagationPlugin(ctx context.Context, log logr.Logger) (*sts.Tok return nil, fmt.Errorf("failed to initialize STS integration: %w", err) } + // RFC 8707 resource / RFC 8693 audience scope the exchanged token to a + // backend. Both are repeatable; empty values are omitted from the request. + resource := splitCSV(os.Getenv("KAGENT_STS_RESOURCE")) + audience := splitCSV(os.Getenv("KAGENT_STS_AUDIENCE")) + log.Info("Enabling STS token propagation plugin", "wellKnownURI", stsWellKnownURI) - return sts.NewTokenPropagationPlugin(integration, log), nil + return sts.NewTokenPropagationPlugin(integration, log, resource, audience), nil +} + +// splitCSV parses a comma-separated value into trimmed, non-empty entries, +// returning nil when none are present so the exchange target stays unset. +func splitCSV(v string) []string { + var out []string + for p := range strings.SplitSeq(v, ",") { + if p = strings.TrimSpace(p); p != "" { + out = append(out, p) + } + } + return out } diff --git a/go/adk/pkg/sts/client.go b/go/adk/pkg/sts/client.go index ccc78e496f..8f17ce216a 100644 --- a/go/adk/pkg/sts/client.go +++ b/go/adk/pkg/sts/client.go @@ -82,27 +82,14 @@ func (c *STSClient) buildFormData(req *TokenExchangeRequest) url.Values { } } - // Add optional parameters - if req.Resource != nil { - switch v := req.Resource.(type) { - case string: - data.Set("resource", v) - case []string: - for _, r := range v { - data.Add("resource", r) - } - } + // Add optional parameters. resource/audience are RFC 8707/8693 repeatable; + // an empty slice sends nothing. + for _, r := range req.Resource { + data.Add("resource", r) } - if req.Audience != nil { - switch v := req.Audience.(type) { - case string: - data.Set("audience", v) - case []string: - for _, a := range v { - data.Add("audience", a) - } - } + for _, a := range req.Audience { + data.Add("audience", a) } if req.Scope != "" { @@ -135,8 +122,8 @@ func (c *STSClient) ExchangeToken( subjectTokenType TokenType, actorToken string, actorTokenType TokenType, - resource any, - audience any, + resource []string, + audience []string, scope string, requestedTokenType TokenType, additionalParameters map[string]any, @@ -217,8 +204,8 @@ func (c *STSClient) Impersonate( ctx context.Context, subjectToken string, subjectTokenType TokenType, - resource any, - audience any, + resource []string, + audience []string, scope string, requestedTokenType TokenType, additionalParameters map[string]any, @@ -244,8 +231,8 @@ func (c *STSClient) Delegate( subjectTokenType TokenType, actorToken string, actorTokenType TokenType, - resource any, - audience any, + resource []string, + audience []string, scope string, requestedTokenType TokenType, additionalParameters map[string]any, diff --git a/go/adk/pkg/sts/client_test.go b/go/adk/pkg/sts/client_test.go index 0ff8cd9764..35fa19e3d3 100644 --- a/go/adk/pkg/sts/client_test.go +++ b/go/adk/pkg/sts/client_test.go @@ -93,7 +93,7 @@ func TestSTSClientDelegateBuildsRequestData(t *testing.T) { "actor-token", TokenTypeJWT, nil, - "https://api.example.com", + []string{"https://api.example.com"}, "read write", "", nil, @@ -103,6 +103,57 @@ func TestSTSClientDelegateBuildsRequestData(t *testing.T) { } } +func TestBuildFormDataResourceAudience(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + resource []string + audience []string + }{ + {name: "single", resource: []string{"https://mcp.example.com"}, audience: []string{"mcp-backend"}}, + {name: "multiple resources", resource: []string{"https://a.example.com", "https://b.example.com"}}, + {name: "audience only", audience: []string{"mcp-backend"}}, + {name: "none omits both keys"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + data := (&STSClient{}).buildFormData(&TokenExchangeRequest{ + GrantType: GrantTypeTokenExchange, + SubjectToken: "subject", + SubjectTokenType: TokenTypeJWT, + Resource: tt.resource, + Audience: tt.audience, + }) + assertFormMulti(t, data, "resource", tt.resource) + assertFormMulti(t, data, "audience", tt.audience) + }) + } +} + +// assertFormMulti checks key carries exactly want in order; an empty want +// means the key must be absent, never present-but-empty. +func assertFormMulti(t *testing.T, form url.Values, key string, want []string) { + t.Helper() + got := form[key] + if len(want) == 0 { + if form.Has(key) { + t.Fatalf("%s = %v, want key absent", key, got) + } + return + } + if len(got) != len(want) { + t.Fatalf("%s = %v, want %v", key, got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("%s[%d] = %q, want %q", key, i, got[i], want[i]) + } + } +} + func TestSTSClientExchangeTokenErrorResponse(t *testing.T) { t.Parallel() srv := newMockSTSClientServer(t, func(w http.ResponseWriter, r *http.Request) { diff --git a/go/adk/pkg/sts/integration.go b/go/adk/pkg/sts/integration.go index 2de4f83713..1f22bb2d97 100644 --- a/go/adk/pkg/sts/integration.go +++ b/go/adk/pkg/sts/integration.go @@ -132,8 +132,8 @@ func (i *STSIntegration) ExchangeToken( ctx context.Context, subjectToken string, subjectTokenType TokenType, - resource any, - audience any, + resource []string, + audience []string, scope string, requestedTokenType TokenType, ) (*TokenExchangeResponse, error) { @@ -153,8 +153,8 @@ func (i *STSIntegration) ExchangeTokenWithActorToken( subjectToken string, subjectTokenType TokenType, actorToken string, - resource any, - audience any, + resource []string, + audience []string, scope string, requestedTokenType TokenType, ) (*TokenExchangeResponse, error) { diff --git a/go/adk/pkg/sts/models.go b/go/adk/pkg/sts/models.go index cfedac328f..acfb26f883 100644 --- a/go/adk/pkg/sts/models.go +++ b/go/adk/pkg/sts/models.go @@ -40,10 +40,10 @@ type TokenExchangeRequest struct { ActorToken string `json:"actor_token,omitempty"` // ActorTokenType is the type of the actor_token (required if ActorToken is set) ActorTokenType TokenType `json:"actor_token_type,omitempty"` - // Resource is the logical name of the target service or resource (optional) - Resource any `json:"resource,omitempty"` // Can be string or []string - // Audience is the logical name of the target service or resource (optional) - Audience any `json:"audience,omitempty"` // Can be string or []string + // Resource holds RFC 8707 resource indicators; repeatable (optional) + Resource []string `json:"resource,omitempty"` + // Audience holds RFC 8693 audiences; repeatable (optional) + Audience []string `json:"audience,omitempty"` // Scope is the scope of the requested token (optional) Scope string `json:"scope,omitempty"` // RequestedTokenType is the type of the requested token (optional) diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index 39fe26f5fe..a630016a63 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -38,16 +38,22 @@ type TokenPropagationPlugin struct { mu sync.RWMutex logger logr.Logger bufferSeconds int64 + resource []string // RFC 8707 resource indicators sent on the STS exchange; empty omits them + audience []string // RFC 8693 audiences sent on the STS exchange; empty omits them } // NewTokenPropagationPlugin creates a new token propagation plugin. // If integration is nil, the plugin will pass through tokens without exchange. -func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger) *TokenPropagationPlugin { +// resource and audience scope the exchanged token to a backend; empty values +// are omitted from the request, leaving the exchange unscoped. +func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, resource, audience []string) *TokenPropagationPlugin { return &TokenPropagationPlugin{ integration: integration, tokenCache: make(map[string]*TokenCacheEntry), logger: logger.WithName("sts-plugin"), bufferSeconds: 5, + resource: resource, + audience: audience, } } @@ -179,10 +185,10 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) subjectToken, TokenTypeJWT, actorToken, - nil, // resource - nil, // audience - "", // scope - "", // requestedTokenType + p.resource, + p.audience, + "", // scope + "", // requestedTokenType ) if err != nil { p.logger.Error(err, "STS token exchange failed, tools may not authenticate", "sessionID", sessionID) diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 8bf15cb0bf..f82e210dab 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -59,7 +59,7 @@ func (f fakeSession) LastUpdateTime() time.Time { return time.Time{} } func TestHeaderProvider_UsesSessionIDMethod(t *testing.T) { t.Parallel() - plugin := NewTokenPropagationPlugin(nil, logr.Discard()) + plugin := NewTokenPropagationPlugin(nil, logr.Discard(), nil, nil) plugin.setCachedToken("sess-123", "token-abc", 0) headers := plugin.HeaderProvider(fakeSessionContext{ @@ -120,7 +120,7 @@ func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T t.Fatalf("NewSTSIntegration() error = %v", err) } - plugin := NewTokenPropagationPlugin(integration, logr.Discard()) + plugin := NewTokenPropagationPlugin(integration, logr.Discard(), nil, nil) for _, sessionID := range []string{"sess-one", "sess-two"} { ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, "subject-token") if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{ @@ -139,6 +139,103 @@ func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T } } +func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + resource []string + audience []string + wantResource string + wantAudience string + }{ + { + name: "configured target is sent", + resource: []string{"https://mcp.example.com"}, + audience: []string{"mcp-backend"}, + wantResource: "https://mcp.example.com", + wantAudience: "mcp-backend", + }, + { + name: "no target leaves resource and audience unset", + wantResource: "", + wantAudience: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + type exchangeForm struct { + resource string + audience string + err error + } + // Buffered so the handler never blocks on send; the value is read + // back on the test goroutine to avoid a data race on the captured form. + gotForm := make(chan exchangeForm, 1) + + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/oauth-authorization-server" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "issuer": srv.URL, + "token_endpoint": srv.URL + "/token", + }) + return + } + if r.URL.Path != "/token" { + http.NotFound(w, r) + return + } + if err := r.ParseForm(); err != nil { + gotForm <- exchangeForm{err: err} + } else { + gotForm <- exchangeForm{resource: r.FormValue("resource"), audience: r.FormValue("audience")} + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access-token", + "issued_token_type": string(TokenTypeJWT), + }) + })) + defer srv.Close() + + integration, err := NewSTSIntegration( + srv.URL+"/.well-known/oauth-authorization-server", + "", nil, nil, 5, true, false, + ) + if err != nil { + t.Fatalf("NewSTSIntegration() error = %v", err) + } + + plugin := NewTokenPropagationPlugin(integration, logr.Discard(), tt.resource, tt.audience) + ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, "subject-token") + if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{ + Context: ctx, + sessionID: "sess-resource", + }); err != nil { + t.Fatalf("BeforeRunCallback() error = %v", err) + } + + select { + case got := <-gotForm: + if got.err != nil { + t.Fatalf("ParseForm() error = %v", got.err) + } + if got.resource != tt.wantResource { + t.Fatalf("resource = %q, want %q", got.resource, tt.wantResource) + } + if got.audience != tt.wantAudience { + t.Fatalf("audience = %q, want %q", got.audience, tt.wantAudience) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for token exchange request") + } + }) + } +} + func TestExtractJWTExpiryUsesUnverifiedClaims(t *testing.T) { t.Parallel() want := time.Now().Add(time.Hour).Unix() diff --git a/go/core/pkg/env/kagent.go b/go/core/pkg/env/kagent.go index a72b6f0028..3de9d00004 100644 --- a/go/core/pkg/env/kagent.go +++ b/go/core/pkg/env/kagent.go @@ -94,4 +94,18 @@ var ( "Well-known endpoint for the Security Token Service (STS) used for token exchange.", ComponentAgentRuntime, ) + + KagentSTSResource = RegisterStringVar( + "KAGENT_STS_RESOURCE", + "", + "RFC 8707 resource indicator sent on STS token-exchange requests to scope the issued token to a target backend.", + ComponentAgentRuntime, + ) + + KagentSTSAudience = RegisterStringVar( + "KAGENT_STS_AUDIENCE", + "", + "RFC 8693 audience sent on STS token-exchange requests. Alternate to KAGENT_STS_RESOURCE for servers that key on audience.", + ComponentAgentRuntime, + ) )