From c33770f96f9ddb7769e4609d4b69a9e24c7e4947 Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Mon, 29 Jun 2026 11:15:17 +0200 Subject: [PATCH 1/5] feat(sts): send RFC 8707 resource and audience on token exchange The token-propagation plugin hardcoded nil resource/audience on every STS token exchange, so issued tokens could not be scoped to a target backend. Read KAGENT_TOKEN_RESOURCE / KAGENT_TOKEN_AUDIENCE and pass them through to the exchange via a WithExchangeTarget option, omitting unset values. Signed-off-by: QuentinBisson --- go/adk/pkg/runner/adapter.go | 20 ++++++++- go/adk/pkg/sts/plugin.go | 31 +++++++++++--- go/adk/pkg/sts/plugin_test.go | 80 +++++++++++++++++++++++++++++++++++ go/core/pkg/env/kagent.go | 14 ++++++ 4 files changed, 138 insertions(+), 7 deletions(-) diff --git a/go/adk/pkg/runner/adapter.go b/go/adk/pkg/runner/adapter.go index 0441f778c0..93a9d264c2 100644 --- a/go/adk/pkg/runner/adapter.go +++ b/go/adk/pkg/runner/adapter.go @@ -123,6 +123,24 @@ func buildTokenPropagationPlugin(ctx context.Context, log logr.Logger) (*sts.Tok return nil, fmt.Errorf("failed to initialize STS integration: %w", err) } + var opts []sts.Option + if resource, audience := exchangeTarget(); resource != nil || audience != nil { + opts = append(opts, sts.WithExchangeTarget(resource, audience)) + } + log.Info("Enabling STS token propagation plugin", "wellKnownURI", stsWellKnownURI) - return sts.NewTokenPropagationPlugin(integration, log), nil + return sts.NewTokenPropagationPlugin(integration, log, opts...), nil +} + +// exchangeTarget reads the RFC 8707 resource and RFC 8693 audience from the +// environment. An unset value is returned as nil so it is omitted from the +// token-exchange request. +func exchangeTarget() (resource, audience any) { + if r := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_RESOURCE")); r != "" { + resource = r + } + if a := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_AUDIENCE")); a != "" { + audience = a + } + return resource, audience } diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index 39fe26f5fe..eb29a0d874 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -38,17 +38,36 @@ type TokenPropagationPlugin struct { mu sync.RWMutex logger logr.Logger bufferSeconds int64 + resource any // RFC 8707 resource indicator passed to the STS exchange + audience any // RFC 8693 audience passed to the STS exchange +} + +// Option configures a TokenPropagationPlugin. +type Option func(*TokenPropagationPlugin) + +// WithExchangeTarget sets the RFC 8707 resource and RFC 8693 audience sent on +// token-exchange requests. A nil value is omitted from the request. Strings or +// string slices are accepted, matching TokenExchangeRequest.Resource/Audience. +func WithExchangeTarget(resource, audience any) Option { + return func(p *TokenPropagationPlugin) { + p.resource = resource + p.audience = audience + } } // 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 { - return &TokenPropagationPlugin{ +func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, opts ...Option) *TokenPropagationPlugin { + p := &TokenPropagationPlugin{ integration: integration, tokenCache: make(map[string]*TokenCacheEntry), logger: logger.WithName("sts-plugin"), bufferSeconds: 5, } + for _, opt := range opts { + opt(p) + } + return p } // getCachedToken retrieves a valid cached token for the session. @@ -179,10 +198,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..53bda2d99a 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -139,6 +139,86 @@ func TestBeforeRunCallback_ReusesCachedDynamicActorTokenForExchange(t *testing.T } } +func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts []Option + wantResource string + wantAudience string + }{ + { + name: "configured target is sent", + opts: []Option{WithExchangeTarget("https://mcp.example.com", "mcp-backend")}, + wantResource: "https://mcp.example.com", + wantAudience: "mcp-backend", + }, + { + name: "no target leaves resource and audience unset", + opts: nil, + wantResource: "", + wantAudience: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var gotResource, gotAudience string + 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 { + t.Fatalf("ParseForm() error = %v", err) + } + gotResource = r.FormValue("resource") + gotAudience = 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.opts...) + 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) + } + + if gotResource != tt.wantResource { + t.Fatalf("resource = %q, want %q", gotResource, tt.wantResource) + } + if gotAudience != tt.wantAudience { + t.Fatalf("audience = %q, want %q", gotAudience, tt.wantAudience) + } + }) + } +} + 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 5d158b2060..532591a0bd 100644 --- a/go/core/pkg/env/kagent.go +++ b/go/core/pkg/env/kagent.go @@ -76,4 +76,18 @@ var ( "Well-known endpoint for the Security Token Service (STS) used for token exchange.", ComponentAgentRuntime, ) + + KagentTokenResource = RegisterStringVar( + "KAGENT_TOKEN_RESOURCE", + "", + "RFC 8707 resource indicator sent on STS token-exchange requests to scope the issued token to a target backend.", + ComponentAgentRuntime, + ) + + KagentTokenAudience = RegisterStringVar( + "KAGENT_TOKEN_AUDIENCE", + "", + "RFC 8693 audience sent on STS token-exchange requests. Alternate to KAGENT_TOKEN_RESOURCE for servers that key on audience.", + ComponentAgentRuntime, + ) ) From c516180f37a3354265bb38a6814e614a30d5c7cf Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Mon, 29 Jun 2026 13:07:54 +0200 Subject: [PATCH 2/5] refactor(sts): constrain exchange target to string; fix test data race Address review feedback on the resource/audience plumbing: - WithExchangeTarget takes string instead of any, so callers cannot pass a type the STS client silently drops; empty values are omitted. - The resource/audience test captures the token-exchange form over a channel and waits with a timeout, removing the data race and the unreliable t.Fatalf from the server handler goroutine. Signed-off-by: QuentinBisson --- go/adk/pkg/runner/adapter.go | 23 +++++------------------ go/adk/pkg/sts/plugin.go | 15 ++++++++++----- go/adk/pkg/sts/plugin_test.go | 34 +++++++++++++++++++++++++--------- 3 files changed, 40 insertions(+), 32 deletions(-) diff --git a/go/adk/pkg/runner/adapter.go b/go/adk/pkg/runner/adapter.go index 93a9d264c2..213c89b3a9 100644 --- a/go/adk/pkg/runner/adapter.go +++ b/go/adk/pkg/runner/adapter.go @@ -123,24 +123,11 @@ func buildTokenPropagationPlugin(ctx context.Context, log logr.Logger) (*sts.Tok return nil, fmt.Errorf("failed to initialize STS integration: %w", err) } - var opts []sts.Option - if resource, audience := exchangeTarget(); resource != nil || audience != nil { - opts = append(opts, sts.WithExchangeTarget(resource, audience)) - } + // RFC 8707 resource / RFC 8693 audience scope the exchanged token to a + // backend. Empty values are omitted by WithExchangeTarget. + resource := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_RESOURCE")) + audience := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_AUDIENCE")) log.Info("Enabling STS token propagation plugin", "wellKnownURI", stsWellKnownURI) - return sts.NewTokenPropagationPlugin(integration, log, opts...), nil -} - -// exchangeTarget reads the RFC 8707 resource and RFC 8693 audience from the -// environment. An unset value is returned as nil so it is omitted from the -// token-exchange request. -func exchangeTarget() (resource, audience any) { - if r := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_RESOURCE")); r != "" { - resource = r - } - if a := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_AUDIENCE")); a != "" { - audience = a - } - return resource, audience + return sts.NewTokenPropagationPlugin(integration, log, sts.WithExchangeTarget(resource, audience)), nil } diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index eb29a0d874..b1302a774d 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -46,12 +46,17 @@ type TokenPropagationPlugin struct { type Option func(*TokenPropagationPlugin) // WithExchangeTarget sets the RFC 8707 resource and RFC 8693 audience sent on -// token-exchange requests. A nil value is omitted from the request. Strings or -// string slices are accepted, matching TokenExchangeRequest.Resource/Audience. -func WithExchangeTarget(resource, audience any) Option { +// token-exchange requests. Empty values are omitted from the request, so an +// unset target leaves the exchange unscoped. Values are single strings, which +// the STS client always serializes; multi-valued resources are out of scope. +func WithExchangeTarget(resource, audience string) Option { return func(p *TokenPropagationPlugin) { - p.resource = resource - p.audience = audience + if resource != "" { + p.resource = resource + } + if audience != "" { + p.audience = audience + } } } diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 53bda2d99a..9c0fc77602 100644 --- a/go/adk/pkg/sts/plugin_test.go +++ b/go/adk/pkg/sts/plugin_test.go @@ -166,7 +166,15 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - var gotResource, gotAudience string + 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" { @@ -181,10 +189,10 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { return } if err := r.ParseForm(); err != nil { - t.Fatalf("ParseForm() error = %v", err) + gotForm <- exchangeForm{err: err} + } else { + gotForm <- exchangeForm{resource: r.FormValue("resource"), audience: r.FormValue("audience")} } - gotResource = r.FormValue("resource") - gotAudience = r.FormValue("audience") _ = json.NewEncoder(w).Encode(map[string]any{ "access_token": "access-token", "issued_token_type": string(TokenTypeJWT), @@ -209,11 +217,19 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { t.Fatalf("BeforeRunCallback() error = %v", err) } - if gotResource != tt.wantResource { - t.Fatalf("resource = %q, want %q", gotResource, tt.wantResource) - } - if gotAudience != tt.wantAudience { - t.Fatalf("audience = %q, want %q", gotAudience, tt.wantAudience) + 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") } }) } From c3f233b76ca207e624476b8beb171e4f432f160b Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Tue, 7 Jul 2026 19:17:44 +0200 Subject: [PATCH 3/5] refactor(sts): drop options pattern, take resource/audience as params Replace the WithExchangeTarget functional option with plain resource and audience string parameters on NewTokenPropagationPlugin, and type the plugin fields as string instead of any. Empty values are converted to nil at the exchange call site so unset targets are omitted from the request. Signed-off-by: QuentinBisson --- go/adk/pkg/runner/adapter.go | 6 ++--- go/adk/pkg/sts/plugin.go | 41 ++++++++++++++--------------------- go/adk/pkg/sts/plugin_test.go | 13 ++++++----- 3 files changed, 26 insertions(+), 34 deletions(-) diff --git a/go/adk/pkg/runner/adapter.go b/go/adk/pkg/runner/adapter.go index 213c89b3a9..153adf76dc 100644 --- a/go/adk/pkg/runner/adapter.go +++ b/go/adk/pkg/runner/adapter.go @@ -106,7 +106,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 } defaultSTSConfig := sts.DefaultSTSConfig(stsWellKnownURI) @@ -124,10 +124,10 @@ func buildTokenPropagationPlugin(ctx context.Context, log logr.Logger) (*sts.Tok } // RFC 8707 resource / RFC 8693 audience scope the exchanged token to a - // backend. Empty values are omitted by WithExchangeTarget. + // backend. Empty values are omitted from the exchange request. resource := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_RESOURCE")) audience := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_AUDIENCE")) log.Info("Enabling STS token propagation plugin", "wellKnownURI", stsWellKnownURI) - return sts.NewTokenPropagationPlugin(integration, log, sts.WithExchangeTarget(resource, audience)), nil + return sts.NewTokenPropagationPlugin(integration, log, resource, audience), nil } diff --git a/go/adk/pkg/sts/plugin.go b/go/adk/pkg/sts/plugin.go index b1302a774d..67327ad6b6 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -38,41 +38,32 @@ type TokenPropagationPlugin struct { mu sync.RWMutex logger logr.Logger bufferSeconds int64 - resource any // RFC 8707 resource indicator passed to the STS exchange - audience any // RFC 8693 audience passed to the STS exchange + resource string // RFC 8707 resource indicator sent on the STS exchange; empty omits it + audience string // RFC 8693 audience sent on the STS exchange; empty omits it } -// Option configures a TokenPropagationPlugin. -type Option func(*TokenPropagationPlugin) - -// WithExchangeTarget sets the RFC 8707 resource and RFC 8693 audience sent on -// token-exchange requests. Empty values are omitted from the request, so an -// unset target leaves the exchange unscoped. Values are single strings, which -// the STS client always serializes; multi-valued resources are out of scope. -func WithExchangeTarget(resource, audience string) Option { - return func(p *TokenPropagationPlugin) { - if resource != "" { - p.resource = resource - } - if audience != "" { - p.audience = audience - } +// omitEmpty returns nil for an empty string so the STS client omits the +// parameter entirely rather than sending it with an empty value. +func omitEmpty(s string) any { + if s == "" { + return nil } + return s } // 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, opts ...Option) *TokenPropagationPlugin { - p := &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, } - for _, opt := range opts { - opt(p) - } - return p } // getCachedToken retrieves a valid cached token for the session. @@ -203,8 +194,8 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) subjectToken, TokenTypeJWT, actorToken, - p.resource, - p.audience, + omitEmpty(p.resource), + omitEmpty(p.audience), "", // scope "", // requestedTokenType ) diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 9c0fc77602..3742e55809 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(), "", "") 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(), "", "") for _, sessionID := range []string{"sess-one", "sess-two"} { ctx := context.WithValue(context.Background(), kagentmodels.BearerTokenKey, "subject-token") if _, err := plugin.BeforeRunCallback(&fakeInvocationContext{ @@ -144,19 +144,20 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { tests := []struct { name string - opts []Option + resource string + audience string wantResource string wantAudience string }{ { name: "configured target is sent", - opts: []Option{WithExchangeTarget("https://mcp.example.com", "mcp-backend")}, + resource: "https://mcp.example.com", + audience: "mcp-backend", wantResource: "https://mcp.example.com", wantAudience: "mcp-backend", }, { name: "no target leaves resource and audience unset", - opts: nil, wantResource: "", wantAudience: "", }, @@ -208,7 +209,7 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { t.Fatalf("NewSTSIntegration() error = %v", err) } - plugin := NewTokenPropagationPlugin(integration, logr.Discard(), tt.opts...) + 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, From 3d00a01a9a1af2b47531d5cf4d07fda2e00b49b5 Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Tue, 7 Jul 2026 19:21:59 +0200 Subject: [PATCH 4/5] refactor(sts): rename token-exchange env vars to KAGENT_STS_* Align the Go env var names with the Python runtime: KAGENT_TOKEN_RESOURCE/ AUDIENCE become KAGENT_STS_RESOURCE/AUDIENCE so the names make clear they belong to the STS token exchange. Signed-off-by: QuentinBisson --- go/adk/pkg/runner/adapter.go | 4 ++-- go/core/pkg/env/kagent.go | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/go/adk/pkg/runner/adapter.go b/go/adk/pkg/runner/adapter.go index 153adf76dc..25be2c68ba 100644 --- a/go/adk/pkg/runner/adapter.go +++ b/go/adk/pkg/runner/adapter.go @@ -125,8 +125,8 @@ func buildTokenPropagationPlugin(ctx context.Context, log logr.Logger) (*sts.Tok // RFC 8707 resource / RFC 8693 audience scope the exchanged token to a // backend. Empty values are omitted from the exchange request. - resource := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_RESOURCE")) - audience := strings.TrimSpace(os.Getenv("KAGENT_TOKEN_AUDIENCE")) + resource := strings.TrimSpace(os.Getenv("KAGENT_STS_RESOURCE")) + audience := strings.TrimSpace(os.Getenv("KAGENT_STS_AUDIENCE")) log.Info("Enabling STS token propagation plugin", "wellKnownURI", stsWellKnownURI) return sts.NewTokenPropagationPlugin(integration, log, resource, audience), nil diff --git a/go/core/pkg/env/kagent.go b/go/core/pkg/env/kagent.go index 532591a0bd..b22b5e895a 100644 --- a/go/core/pkg/env/kagent.go +++ b/go/core/pkg/env/kagent.go @@ -77,17 +77,17 @@ var ( ComponentAgentRuntime, ) - KagentTokenResource = RegisterStringVar( - "KAGENT_TOKEN_RESOURCE", + 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, ) - KagentTokenAudience = RegisterStringVar( - "KAGENT_TOKEN_AUDIENCE", + KagentSTSAudience = RegisterStringVar( + "KAGENT_STS_AUDIENCE", "", - "RFC 8693 audience sent on STS token-exchange requests. Alternate to KAGENT_TOKEN_RESOURCE for servers that key on audience.", + "RFC 8693 audience sent on STS token-exchange requests. Alternate to KAGENT_STS_RESOURCE for servers that key on audience.", ComponentAgentRuntime, ) ) From 63f6c1ef9e5009cab8d9513b0614f3f3cfb30cbc Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Wed, 8 Jul 2026 10:50:03 +0200 Subject: [PATCH 5/5] refactor(sts): type exchange resource/audience as []string RFC 8707 resource and RFC 8693 audience are both repeatable, so model them as []string on the STS request instead of any. buildFormData adds one form value per entry and an empty slice sends nothing, so the omitEmpty helper and the type switch are gone. The runtime reads comma-separated KAGENT_STS_RESOURCE / KAGENT_STS_AUDIENCE into the slices. Signed-off-by: QuentinBisson --- go/adk/pkg/runner/adapter.go | 20 ++++++++++--- go/adk/pkg/sts/client.go | 37 ++++++++---------------- go/adk/pkg/sts/client_test.go | 53 ++++++++++++++++++++++++++++++++++- go/adk/pkg/sts/integration.go | 8 +++--- go/adk/pkg/sts/models.go | 8 +++--- go/adk/pkg/sts/plugin.go | 19 ++++--------- go/adk/pkg/sts/plugin_test.go | 12 ++++---- 7 files changed, 99 insertions(+), 58 deletions(-) diff --git a/go/adk/pkg/runner/adapter.go b/go/adk/pkg/runner/adapter.go index 67b027a213..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) @@ -145,10 +145,22 @@ func buildTokenPropagationPlugin(ctx context.Context, log logr.Logger) (*sts.Tok } // RFC 8707 resource / RFC 8693 audience scope the exchanged token to a - // backend. Empty values are omitted from the exchange request. - resource := strings.TrimSpace(os.Getenv("KAGENT_STS_RESOURCE")) - audience := strings.TrimSpace(os.Getenv("KAGENT_STS_AUDIENCE")) + // 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, 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 67327ad6b6..a630016a63 100644 --- a/go/adk/pkg/sts/plugin.go +++ b/go/adk/pkg/sts/plugin.go @@ -38,24 +38,15 @@ type TokenPropagationPlugin struct { mu sync.RWMutex logger logr.Logger bufferSeconds int64 - resource string // RFC 8707 resource indicator sent on the STS exchange; empty omits it - audience string // RFC 8693 audience sent on the STS exchange; empty omits it -} - -// omitEmpty returns nil for an empty string so the STS client omits the -// parameter entirely rather than sending it with an empty value. -func omitEmpty(s string) any { - if s == "" { - return nil - } - return s + 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. // 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 { +func NewTokenPropagationPlugin(integration *STSIntegration, logger logr.Logger, resource, audience []string) *TokenPropagationPlugin { return &TokenPropagationPlugin{ integration: integration, tokenCache: make(map[string]*TokenCacheEntry), @@ -194,8 +185,8 @@ func (p *TokenPropagationPlugin) BeforeRunCallback(ctx agent.InvocationContext) subjectToken, TokenTypeJWT, actorToken, - omitEmpty(p.resource), - omitEmpty(p.audience), + p.resource, + p.audience, "", // scope "", // requestedTokenType ) diff --git a/go/adk/pkg/sts/plugin_test.go b/go/adk/pkg/sts/plugin_test.go index 3742e55809..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{ @@ -144,15 +144,15 @@ func TestBeforeRunCallback_SendsResourceAndAudience(t *testing.T) { tests := []struct { name string - resource string - audience string + resource []string + audience []string wantResource string wantAudience string }{ { name: "configured target is sent", - resource: "https://mcp.example.com", - audience: "mcp-backend", + resource: []string{"https://mcp.example.com"}, + audience: []string{"mcp-backend"}, wantResource: "https://mcp.example.com", wantAudience: "mcp-backend", },