Allow OAuth U2M logins to send RFC 8707 resource indicators

cmd/auth/login.go

@@ -150,6 +150,7 @@ a new profile is created.
 	var skipWorkspace bool
 	var scopes string
 	var clientID string
+	var resources []string
 	cmd.Flags().DurationVar(&loginTimeout, "timeout", defaultTimeout,
 		"Timeout for completing login challenge in the browser")
 	cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false,
@@ -162,6 +163,8 @@ a new profile is created.
 		"Comma-separated list of OAuth scopes to request (defaults to 'all-apis')")
 	cmd.Flags().StringVar(&clientID, "client-id", "",
 		"OAuth client ID to use for U2M authentication")
+	cmd.Flags().StringArrayVar(&resources, "resource", nil,
+		"RFC 8707 resource indicator to scope the login to (repeatable). Requires --host.")
 
 	cmd.PreRunE = profileHostConflictCheck
 
@@ -279,6 +282,9 @@ a new profile is created.
 		if clientID == "" {
 			clientID = u2mClientIDFromProfile(existingProfile)
 		}
+		if len(resources) == 0 {
+			resources = u2mResourcesFromProfile(existingProfile)
+		}
 
 		// If no host is available from any source, use the discovery flow
 		// via login.databricks.com.
@@ -331,6 +337,9 @@ a new profile is created.
 		if len(scopesList) > 0 {
 			persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList))
 		}
+		if len(resources) > 0 {
+			persistentAuthOpts = append(persistentAuthOpts, u2m.WithResources(resources))
+		}
 		persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...)
 		if err != nil {
 			return err
@@ -422,6 +431,11 @@ a new profile is created.
 			if err != nil {
 				return err
 			}
+			// `resources` is a CLI-only profile key (not an SDK config attribute), so
+			// it is written separately from SaveToProfile.
+			if err := databrickscfg.SaveResourcesToProfile(ctx, profileName, env.Get(ctx, "DATABRICKS_CONFIG_FILE"), resources); err != nil {
+				return err
+			}
 		}
 
 		if err := storeLoginToken(ctx, tokenStore, mode, oauthArgument, token); err != nil {
@@ -643,6 +657,10 @@ var discoveryIncompatibleFlags = []string{
 	"workspace-id",
 	"configure-cluster",
 	"configure-serverless",
+	// A resource indicator scopes the login to a protected resource on a
+	// specific workspace's /oidc, which the login.databricks.com discovery
+	// flow does not target.
+	"resource",
 }
 
 // validateDiscoveryFlagCompatibility returns an error if any flags that require
@@ -825,6 +843,15 @@ func u2mClientIDFromProfile(p *profile.Profile) string {
 	return p.ClientID
 }
 
+// u2mResourcesFromProfile returns the RFC 8707 resource indicators saved on a
+// databricks-cli-auth profile, so `login`/`auth token` re-request the same ones.
+func u2mResourcesFromProfile(p *profile.Profile) []string {
+	if p == nil || p.AuthType != authTypeDatabricksCLI {
+		return nil
+	}
+	return splitScopes(p.Resources)
+}
+
 // promptForWorkspaceSelection lists workspaces for a SPOG account and lets the
 // user pick one. Returns the selected workspace ID or empty string if skipped.
 // This is best-effort: errors are returned to the caller for logging, not shown

cmd/auth/login_test.go

@@ -516,6 +516,40 @@ func TestU2MClientIDFromProfile(t *testing.T) {
 	}
 }
 
+func TestU2MResourcesFromProfile(t *testing.T) {
+	tests := []struct {
+		name    string
+		profile *profile.Profile
+		want    []string
+	}{
+		{name: "no profile"},
+		{
+			name:    "implicit auth type",
+			profile: &profile.Profile{Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github"},
+		},
+		{
+			name:    "M2M auth type",
+			profile: &profile.Profile{AuthType: "oauth-m2m", Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github"},
+		},
+		{
+			name:    "U2M auth type single resource",
+			profile: &profile.Profile{AuthType: authTypeDatabricksCLI, Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github"},
+			want:    []string{"https://workspace.test/ai-gateway/mcp/system.ai.github"},
+		},
+		{
+			name:    "U2M auth type multiple resources",
+			profile: &profile.Profile{AuthType: authTypeDatabricksCLI, Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github, https://workspace.test/ai-gateway/mcp/system.ai.slack"},
+			want:    []string{"https://workspace.test/ai-gateway/mcp/system.ai.github", "https://workspace.test/ai-gateway/mcp/system.ai.slack"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, u2mResourcesFromProfile(tt.profile))
+		})
+	}
+}
+
 func TestRunHostDiscovery_NoHost(t *testing.T) {
 	ctx := t.Context()
 	args := &auth.AuthArguments{}
@@ -796,6 +830,12 @@ func TestValidateDiscoveryFlagCompatibility(t *testing.T) {
 			flagVal: "true",
 			wantErr: "--configure-serverless requires --host to be specified",
 		},
+		{
+			name:    "resource is incompatible",
+			setFlag: "resource",
+			flagVal: "https://workspace.test/ai-gateway/mcp-services/system.ai.github",
+			wantErr: "--resource requires --host to be specified",
+		},
 		{
 			name: "no flags set is ok",
 		},
@@ -807,6 +847,7 @@ func TestValidateDiscoveryFlagCompatibility(t *testing.T) {
 			cmd.Flags().String("workspace-id", "", "")
 			cmd.Flags().Bool("configure-cluster", false, "")
 			cmd.Flags().Bool("configure-serverless", false, "")
+			cmd.Flags().StringArray("resource", nil, "")
 
 			if tt.setFlag != "" {
 				require.NoError(t, cmd.Flags().Set(tt.setFlag, tt.flagVal))

cmd/auth/token.go

@@ -267,6 +267,9 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 	if clientID := u2mClientIDFromProfile(existingProfile); clientID != "" {
 		allArgs = append(allArgs, u2m.WithClientID(clientID))
 	}
+	if resources := u2mResourcesFromProfile(existingProfile); len(resources) > 0 {
+		allArgs = append(allArgs, u2m.WithResources(resources))
+	}
 	allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument))
 	persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...)
 	if err != nil {
@@ -423,6 +426,10 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s
 		scopesList = splitScopes(existingProfile.Scopes)
 	}
 
+	// Preserve RFC 8707 resource indicators from the existing profile so the
+	// inline login requests the same resources the user previously configured.
+	resourcesList := u2mResourcesFromProfile(existingProfile)
+
 	oauthArgument, err := loginArgs.ToOAuthArgument()
 	if err != nil {
 		return "", nil, err
@@ -437,6 +444,9 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s
 	if len(scopesList) > 0 {
 		persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList))
 	}
+	if len(resourcesList) > 0 {
+		persistentAuthOpts = append(persistentAuthOpts, u2m.WithResources(resourcesList))
+	}
 	persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...)
 	if err != nil {
 		return "", nil, err
@@ -467,6 +477,11 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s
 	if err != nil {
 		return "", nil, err
 	}
+	// `resources` is a CLI-only profile key (not an SDK config attribute), so
+	// it is written separately from SaveToProfile.
+	if err := databrickscfg.SaveResourcesToProfile(ctx, profileName, env.Get(ctx, "DATABRICKS_CONFIG_FILE"), resourcesList); err != nil {
+		return "", nil, err
+	}
 	if err := storeLoginToken(ctx, tokenStore, mode, oauthArgument, token); err != nil {
 		return "", nil, err
 	}

libs/auth/u2m/persistent_auth.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
@@ -112,6 +113,13 @@ type PersistentAuth struct {
 	// scopes is the list of OAuth scopes to request.
 	scopes []string
 
+	// resources is the list of RFC 8707 resource indicators to send on the
+	// authorization request. The Databricks /oidc authorize endpoint reads
+	// these to scope the login to a specific protected resource (e.g. an AI
+	// Gateway MCP connection), so it can drive that resource's own login
+	// before issuing the authorization code. Empty means an unrestricted login.
+	resources []string
+
 	// disableOfflineAccess controls whether offline_access scope is requested.
 	// When true, offline_access will NOT be automatically added to scopes,
 	// meaning the token will not include a refresh token.
@@ -194,6 +202,15 @@ func WithScopes(scopes []string) PersistentAuthOption {
 	}
 }
 
+// WithResources sets the RFC 8707 resource indicators for the PersistentAuth.
+// Each value is added as a `resource` query parameter on the authorization
+// request.
+func WithResources(resources []string) PersistentAuthOption {
+	return func(a *PersistentAuth) {
+		a.resources = resources
+	}
+}
+
 // WithDisableOfflineAccess controls whether offline_access scope is requested.
 func WithDisableOfflineAccess(disable bool) PersistentAuthOption {
 	return func(a *PersistentAuth) {
@@ -623,10 +640,14 @@ func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) {
 	if err != nil {
 		return nil, fmt.Errorf("fetching OAuth endpoints: %w", err)
 	}
+	authURL, err := appendResources(endpoints.AuthorizationEndpoint, a.resources)
+	if err != nil {
+		return nil, err
+	}
 	return &oauth2.Config{
 		ClientID: a.clientID,
 		Endpoint: oauth2.Endpoint{
-			AuthURL:   endpoints.AuthorizationEndpoint,
+			AuthURL:   authURL,
 			TokenURL:  endpoints.TokenEndpoint,
 			AuthStyle: oauth2.AuthStyleInParams,
 		},
@@ -635,6 +656,26 @@ func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) {
 	}, nil
 }
 
+// appendResources adds RFC 8707 `resource` indicators to an authorization
+// endpoint URL, preserving any query parameters the endpoint already carries.
+// The oauth2 library appends its own parameters (client_id, PKCE, etc.) after
+// these when it builds the final authorization URL.
+func appendResources(authURL string, resources []string) (string, error) {
+	if len(resources) == 0 {
+		return authURL, nil
+	}
+	u, err := url.Parse(authURL)
+	if err != nil {
+		return "", fmt.Errorf("parsing authorization endpoint: %w", err)
+	}
+	q := u.Query()
+	for _, r := range resources {
+		q.Add("resource", r)
+	}
+	u.RawQuery = q.Encode()
+	return u.String(), nil
+}
+
 func (a *PersistentAuth) stateAndPKCE() (string, *authhandler.PKCEParams, error) {
 	verifier, err := a.randomString(64)
 	if err != nil {

libs/auth/u2m/persistent_auth_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -168,6 +169,60 @@ func TestPersistentAuthClientID(t *testing.T) {
 	}
 }
 
+func TestPersistentAuthResources(t *testing.T) {
+	tests := []struct {
+		name string
+		opts []PersistentAuthOption
+		want []string
+	}{
+		{
+			name: "none",
+			want: nil,
+		},
+		{
+			name: "single",
+			opts: []PersistentAuthOption{WithResources([]string{"https://workspace.test/ai-gateway/mcp-services/system.ai.github"})},
+			want: []string{"https://workspace.test/ai-gateway/mcp-services/system.ai.github"},
+		},
+		{
+			name: "multiple",
+			opts: []PersistentAuthOption{WithResources([]string{"https://a.test/r1", "https://b.test/r2"})},
+			want: []string{"https://a.test/r1", "https://b.test/r2"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			arg, err := NewBasicWorkspaceOAuthArgument("https://workspace.test")
+			if err != nil {
+				t.Fatalf("NewBasicWorkspaceOAuthArgument(): %v", err)
+			}
+			opts := append([]PersistentAuthOption{
+				WithOAuthArgument(arg),
+				WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}),
+			}, tt.opts...)
+			p, err := NewPersistentAuth(t.Context(), opts...)
+			if err != nil {
+				t.Fatalf("NewPersistentAuth(): %v", err)
+			}
+			cfg, err := p.oauth2Config()
+			if err != nil {
+				t.Fatalf("oauth2Config(): %v", err)
+			}
+			parsed, err := url.Parse(cfg.Endpoint.AuthURL)
+			if err != nil {
+				t.Fatalf("parsing AuthURL %q: %v", cfg.Endpoint.AuthURL, err)
+			}
+
+			got := parsed.Query()["resource"]
+
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("resource params = %v, want %v (AuthURL=%q)", got, tt.want, cfg.Endpoint.AuthURL)
+			}
+		})
+	}
+}
+
 func TestToken_RefreshesExpiredAccessToken(t *testing.T) {
 	ctx := t.Context()
 	expectedKey := "https://accounts.cloud.databricks.test/oidc/accounts/xyz"

libs/databrickscfg/ops.go

@@ -258,6 +258,30 @@ func SetConfiguredAuthStorage(ctx context.Context, value, configFilePath string)
 	return writeConfigFile(ctx, configFile)
 }
 
+// SaveResourcesToProfile writes (or clears) the `resources` key on a profile
+// section. `resources` is a comma-separated list of RFC 8707 resource indicators
+// requested during U2M login. It is a CLI-only profile key (not a
+// databricks-sdk-go config attribute), so it is written here rather than through
+// SaveToProfile: `databricks auth token` re-reads it to request the same
+// resources, and a change in the requested resources is reflected in the profile.
+// An empty list removes the key. The profile section must already exist.
+func SaveResourcesToProfile(ctx context.Context, profileName, configFilePath string, resources []string) error {
+	configFile, err := loadOrCreateConfigFile(ctx, configFilePath)
+	if err != nil {
+		return err
+	}
+	section, err := configFile.GetSection(profileName)
+	if err != nil {
+		return fmt.Errorf("profile %q not found: %w", profileName, err)
+	}
+	if len(resources) == 0 {
+		section.DeleteKey("resources")
+	} else {
+		section.Key("resources").SetValue(strings.Join(resources, ","))
+	}
+	return writeConfigFile(ctx, configFile)
+}
+
 // ClearDefaultProfile removes the default_profile key from the [__settings__]
 // section if the current default matches the given profile name.
 func ClearDefaultProfile(ctx context.Context, profileName, configFilePath string) error {