MCP: properly pass along databricks profile (#4040)

## Changes

Make sure all tools and CLI invocations use the same profile and
workspace ID.

## Why
<!-- Why are these changes needed? Provide the context that the reviewer
might be missing.
For example, were there any decisions behind the change that are not
reflected in the code itself? -->

## Tests
<!-- How have you tested the changes? -->

<!-- If your PR needs to be included in the release notes for next
release,
add a separate entry in NEXT_CHANGELOG.md as part of your PR. -->

Author: fjakobs

experimental/apps-mcp/lib/mcp/middleware_test.go

@@ -224,7 +224,7 @@ func TestServerMiddleware(t *testing.T) {
 		Name:    "test-server",
 		Version: "1.0.0",
 	}
-	server := mcp.NewServer(impl, nil)
+	server := mcp.NewServer(impl, nil, nil)
 
 	var executionOrder []string
 
@@ -264,7 +264,7 @@ func TestServerSessionPersistence(t *testing.T) {
 		Name:    "test-server",
 		Version: "1.0.0",
 	}
-	server := mcp.NewServer(impl, nil)
+	server := mcp.NewServer(impl, nil, nil)
 
 	// Add middleware that increments a counter
 	server.AddMiddlewareFunc(func(ctx *mcp.MiddlewareContext, next mcp.NextFunc) (*mcp.CallToolResult, error) {

experimental/apps-mcp/lib/mcp/server.go

@@ -30,11 +30,15 @@ type serverTool struct {
 }
 
 // NewServer creates a new MCP server.
-func NewServer(impl *Implementation, options any) *Server {
+// If sess is nil, a new session will be created.
+func NewServer(impl *Implementation, options any, sess *session.Session) *Server {
+	if sess == nil {
+		sess = session.NewSession()
+	}
 	return &Server{
 		impl:    impl,
 		tools:   make(map[string]*serverTool),
-		session: session.NewSession(),
+		session: sess,
 	}
 }
 

experimental/apps-mcp/lib/middlewares/databricks_client.go

@@ -9,13 +9,15 @@ import (
 	"github.com/databricks/cli/experimental/apps-mcp/lib/mcp"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/prompts"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/session"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 )
 
 const (
-	DatabricksClientKey = "databricks_client"
+	DatabricksClientKey  = "databricks_client"
+	DatabricksProfileKey = "databricks_profile"
 )
 
 func NewDatabricksClientMiddleware(unauthorizedToolNames []string) mcp.Middleware {
@@ -40,8 +42,41 @@ func NewDatabricksClientMiddleware(unauthorizedToolNames []string) mcp.Middlewar
 	})
 }
 
-func MustGetApiClient(ctx context.Context) (*httpclient.ApiClient, error) {
-	w := MustGetDatabricksClient(ctx)
+func GetDatabricksProfile(ctx context.Context) string {
+	sess, err := session.GetSession(ctx)
+	if err != nil {
+		return ""
+	}
+	profile, ok := sess.Get(DatabricksProfileKey)
+	if !ok {
+		return ""
+	}
+	return profile.(string)
+}
+
+// GetAvailableProfiles returns all available profiles from ~/.databrickscfg.
+func GetAvailableProfiles(ctx context.Context) profile.Profiles {
+	profiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles)
+	if err != nil {
+		// If we can't load profiles, return empty list (config file might not exist)
+		return profile.Profiles{}
+	}
+	return profiles
+}
+
+func MustGetApiClient(ctx context.Context) *httpclient.ApiClient {
+	client, err := GetApiClient(ctx)
+	if err != nil {
+		panic(err)
+	}
+	return client
+}
+
+func GetApiClient(ctx context.Context) (*httpclient.ApiClient, error) {
+	w, err := GetDatabricksClient(ctx)
+	if err != nil {
+		return nil, err
+	}
 	clientCfg, err := config.HTTPClientConfigFromConfig(w.Config)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create HTTP client config: %w", err)
@@ -64,28 +99,36 @@ func GetDatabricksClient(ctx context.Context) (*databricks.WorkspaceClient, erro
 	}
 	w, ok := sess.Get(DatabricksClientKey)
 	if !ok {
-		return nil, errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil))
+		return nil, newAuthError(ctx)
 	}
 	return w.(*databricks.WorkspaceClient), nil
 }
 
 func checkAuth(ctx context.Context) (*databricks.WorkspaceClient, error) {
 	w, err := databricks.NewWorkspaceClient()
 	if err != nil {
-		return nil, wrapAuthError(err)
+		return nil, WrapAuthError(ctx, err)
 	}
 
 	_, err = w.CurrentUser.Me(ctx)
 	if err != nil {
-		return nil, wrapAuthError(err)
+		return nil, WrapAuthError(ctx, err)
 	}
 
 	return w, nil
 }
 
-func wrapAuthError(err error) error {
+func WrapAuthError(ctx context.Context, err error) error {
 	if errors.Is(err, config.ErrCannotConfigureDefault) {
-		return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil))
+		return newAuthError(ctx)
 	}
 	return err
 }
+
+func newAuthError(ctx context.Context) error {
+	// Prepare template data
+	data := map[string]any{
+		"Profiles": GetAvailableProfiles(ctx),
+	}
+	return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", data))
+}

experimental/apps-mcp/lib/middlewares/warehouse.go

@@ -86,7 +86,10 @@ func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) {
 	// first resolve DATABRICKS_WAREHOUSE_ID env variable
 	warehouseID := env.Get(ctx, "DATABRICKS_WAREHOUSE_ID")
 	if warehouseID != "" {
-		w := MustGetDatabricksClient(ctx)
+		w, err := GetDatabricksClient(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("get databricks client: %w", err)
+		}
 		warehouse, err := w.Warehouses.Get(ctx, sql.GetWarehouseRequest{
 			Id: warehouseID,
 		})
@@ -100,7 +103,7 @@ func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) {
 		}, nil
 	}
 
-	apiClient, err := MustGetApiClient(ctx)
+	apiClient, err := GetApiClient(ctx)
 	if err != nil {
 		return nil, err
 	}

experimental/apps-mcp/lib/prompts/auth_error.tmpl

@@ -6,9 +6,14 @@
 Not authenticated to Databricks
 
 I need to know either the Databricks workspace URL or the Databricks profile name.
-You can list the available profiles by running `databricks auth profiles`.
 
-ASK the user which of the configured profiles or databricks workspace URL they want to use.
+The available profiles are:
+
+{{- range .Profiles }}
+- {{ .Name }} ({{ .Host }})
+{{- end }}
+
+IMPORTANT: YOU MUST ASK the user which of the configured profiles or databricks workspace URL they want to use.
 Only then call the `databricks_configure_auth` tool to configure the authentication.
 
 Do not run anything else before authenticating successfully.

experimental/apps-mcp/lib/providers/clitools/configure_auth.go

@@ -9,7 +9,6 @@ import (
 	"github.com/databricks/cli/experimental/apps-mcp/lib/prompts"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/session"
 	"github.com/databricks/databricks-sdk-go"
-	"github.com/databricks/databricks-sdk-go/config"
 )
 
 // ConfigureAuth creates and validates a Databricks workspace client with optional host and profile.
@@ -20,9 +19,8 @@ func ConfigureAuth(ctx context.Context, sess *session.Session, host, profile *st
 		return nil, nil
 	}
 
-	var cfg *databricks.Config
+	cfg := &databricks.Config{}
 	if host != nil || profile != nil {
-		cfg = &databricks.Config{}
 		if host != nil {
 			cfg.Host = *host
 		}
@@ -32,12 +30,7 @@ func ConfigureAuth(ctx context.Context, sess *session.Session, host, profile *st
 	}
 
 	var client *databricks.WorkspaceClient
-	var err error
-	if cfg != nil {
-		client, err = databricks.NewWorkspaceClient(cfg)
-	} else {
-		client, err = databricks.NewWorkspaceClient()
-	}
+	client, err := databricks.NewWorkspaceClient(cfg)
 	if err != nil {
 		return nil, err
 	}
@@ -49,19 +42,14 @@ func ConfigureAuth(ctx context.Context, sess *session.Session, host, profile *st
 				"WorkspaceURL": *host,
 			}))
 		}
-		return nil, wrapAuthError(err)
+		return nil, middlewares.WrapAuthError(ctx, err)
 	}
 
-	// Store client in session data
+	// Store client and profile in session data
 	sess.Set(middlewares.DatabricksClientKey, client)
+	if profile != nil {
+		sess.Set(middlewares.DatabricksProfileKey, *profile)
+	}
 
 	return client, nil
 }
-
-// wrapAuthError wraps configuration errors with helpful messages
-func wrapAuthError(err error) error {
-	if errors.Is(err, config.ErrCannotConfigureDefault) {
-		return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil))
-	}
-	return err
-}

experimental/apps-mcp/lib/providers/clitools/configure_auth_test.go

@@ -75,6 +75,8 @@ func TestConfigureAuthWithCustomHost(t *testing.T) {
 }
 
 func TestWrapAuthError(t *testing.T) {
+	ctx := context.Background()
+
 	tests := []struct {
 		name     string
 		err      error
@@ -89,7 +91,7 @@ func TestWrapAuthError(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			wrapped := wrapAuthError(tt.err)
+			wrapped := middlewares.WrapAuthError(ctx, tt.err)
 			assert.Contains(t, wrapped.Error(), tt.expected)
 		})
 	}

experimental/apps-mcp/lib/providers/clitools/explore.go

@@ -8,7 +8,6 @@ import (
 	"github.com/databricks/cli/experimental/apps-mcp/lib/prompts"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/session"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
-	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 )
@@ -21,32 +20,12 @@ func Explore(ctx context.Context) (string, error) {
 		warehouse = nil
 	}
 
-	currentProfile := getCurrentProfile(ctx)
-	profiles := getAvailableProfiles(ctx)
+	currentProfile := middlewares.GetDatabricksProfile(ctx)
+	profiles := middlewares.GetAvailableProfiles(ctx)
 
 	return generateExploreGuidance(ctx, warehouse, currentProfile, profiles), nil
 }
 
-// getCurrentProfile returns the currently active profile name.
-func getCurrentProfile(ctx context.Context) string {
-	// Check DATABRICKS_CONFIG_PROFILE env var
-	profileName := env.Get(ctx, "DATABRICKS_CONFIG_PROFILE")
-	if profileName == "" {
-		return "DEFAULT"
-	}
-	return profileName
-}
-
-// getAvailableProfiles returns all available profiles from ~/.databrickscfg.
-func getAvailableProfiles(ctx context.Context) profile.Profiles {
-	profiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles)
-	if err != nil {
-		// If we can't load profiles, return empty list (config file might not exist)
-		return profile.Profiles{}
-	}
-	return profiles
-}
-
 // generateExploreGuidance creates comprehensive guidance for data exploration.
 func generateExploreGuidance(ctx context.Context, warehouse *sql.EndpointInfo, currentProfile string, profiles profile.Profiles) string {
 	// Build workspace/profile information
@@ -102,6 +81,7 @@ func generateExploreGuidance(ctx context.Context, warehouse *sql.EndpointInfo, c
 		"WarehouseName": warehouseName,
 		"WarehouseID":   warehouseID,
 		"ProfilesInfo":  profilesInfo,
+		"Profile":       currentProfile,
 	}
 
 	// Render base explore template

experimental/apps-mcp/lib/providers/clitools/invoke_databricks_cli.go

@@ -17,15 +17,22 @@ func InvokeDatabricksCLI(ctx context.Context, command []string, workingDirectory
 		return "", errors.New("command is required")
 	}
 
-	workspaceClient := middlewares.MustGetDatabricksClient(ctx)
+	workspaceClient, err := middlewares.GetDatabricksClient(ctx)
+	if err != nil {
+		return "", fmt.Errorf("get databricks client: %w", err)
+	}
 	host := workspaceClient.Config.Host
+	profile := middlewares.GetDatabricksProfile(ctx)
 
 	// GetCLIPath returns the path to the current CLI executable
 	cliPath := common.GetCLIPath()
 	cmd := exec.CommandContext(ctx, cliPath, command...)
 	cmd.Dir = workingDirectory
 	env := os.Environ()
 	env = append(env, "DATABRICKS_HOST="+host)
+	if profile != "" {
+		env = append(env, "DATABRICKS_CONFIG_PROFILE="+profile)
+	}
 	cmd.Env = env
 
 	output, err := cmd.CombinedOutput()

experimental/apps-mcp/lib/providers/clitools/provider.go

@@ -93,7 +93,7 @@ func (p *Provider) RegisterTools(server *mcpsdk.Server) error {
 		},
 		func(ctx context.Context, req *mcpsdk.CallToolRequest, args struct{}) (*mcpsdk.CallToolResult, any, error) {
 			log.Debug(ctx, "explore called")
-			result, err := Explore(session.WithSession(ctx, p.session))
+			result, err := Explore(ctx)
 			if err != nil {
 				return nil, nil, err
 			}