diff --git a/experimental/apps-mcp/lib/mcp/middleware_test.go b/experimental/apps-mcp/lib/mcp/middleware_test.go index 739526da9b..1abdb5b79f 100644 --- a/experimental/apps-mcp/lib/mcp/middleware_test.go +++ b/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) { diff --git a/experimental/apps-mcp/lib/mcp/server.go b/experimental/apps-mcp/lib/mcp/server.go index 2c677dfe99..7e9d79f7b1 100644 --- a/experimental/apps-mcp/lib/mcp/server.go +++ b/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, } } diff --git a/experimental/apps-mcp/lib/middlewares/databricks_client.go b/experimental/apps-mcp/lib/middlewares/databricks_client.go index b7a9bf7146..4190b22db9 100644 --- a/experimental/apps-mcp/lib/middlewares/databricks_client.go +++ b/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,7 +99,7 @@ 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 } @@ -72,20 +107,28 @@ func GetDatabricksClient(ctx context.Context) (*databricks.WorkspaceClient, erro 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)) +} diff --git a/experimental/apps-mcp/lib/middlewares/warehouse.go b/experimental/apps-mcp/lib/middlewares/warehouse.go index 680b26d3ef..9bf1b0a071 100644 --- a/experimental/apps-mcp/lib/middlewares/warehouse.go +++ b/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 } diff --git a/experimental/apps-mcp/lib/prompts/auth_error.tmpl b/experimental/apps-mcp/lib/prompts/auth_error.tmpl index 96f2be6162..e2fbda8675 100644 --- a/experimental/apps-mcp/lib/prompts/auth_error.tmpl +++ b/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. diff --git a/experimental/apps-mcp/lib/providers/clitools/configure_auth.go b/experimental/apps-mcp/lib/providers/clitools/configure_auth.go index 00a40713cf..087283e644 100644 --- a/experimental/apps-mcp/lib/providers/clitools/configure_auth.go +++ b/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 -} diff --git a/experimental/apps-mcp/lib/providers/clitools/configure_auth_test.go b/experimental/apps-mcp/lib/providers/clitools/configure_auth_test.go index da2df50c78..7fdd851284 100644 --- a/experimental/apps-mcp/lib/providers/clitools/configure_auth_test.go +++ b/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) }) } diff --git a/experimental/apps-mcp/lib/providers/clitools/explore.go b/experimental/apps-mcp/lib/providers/clitools/explore.go index bc9550e012..7f6bf986ac 100644 --- a/experimental/apps-mcp/lib/providers/clitools/explore.go +++ b/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 diff --git a/experimental/apps-mcp/lib/providers/clitools/invoke_databricks_cli.go b/experimental/apps-mcp/lib/providers/clitools/invoke_databricks_cli.go index 7906c43339..85143438d2 100644 --- a/experimental/apps-mcp/lib/providers/clitools/invoke_databricks_cli.go +++ b/experimental/apps-mcp/lib/providers/clitools/invoke_databricks_cli.go @@ -17,8 +17,12 @@ 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() @@ -26,6 +30,9 @@ func InvokeDatabricksCLI(ctx context.Context, command []string, workingDirectory 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() diff --git a/experimental/apps-mcp/lib/providers/clitools/provider.go b/experimental/apps-mcp/lib/providers/clitools/provider.go index 7330f3ac2a..0da7f25a34 100644 --- a/experimental/apps-mcp/lib/providers/clitools/provider.go +++ b/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 } diff --git a/experimental/apps-mcp/lib/server/server.go b/experimental/apps-mcp/lib/server/server.go index 8bf7bfeb42..3b5bf65cb8 100644 --- a/experimental/apps-mcp/lib/server/server.go +++ b/experimental/apps-mcp/lib/server/server.go @@ -31,11 +31,8 @@ func NewServer(ctx context.Context, cfg *mcp.Config) *Server { Version: build.GetInfo().Version, } - server := mcpsdk.NewServer(impl, nil) sess := session.NewSession() - - // Set enabled capabilities for this MCP server - sess.Set(session.CapabilitiesDataKey, []string{"apps"}) + server := mcpsdk.NewServer(impl, nil, sess) tracker, err := trajectory.NewTracker(ctx, sess, cfg) if err != nil { @@ -50,6 +47,9 @@ func NewServer(ctx context.Context, cfg *mcp.Config) *Server { sess.SetTracker(tracker) + // Set enabled capabilities for this MCP server + sess.Set(session.CapabilitiesDataKey, []string{"apps"}) + return &Server{ server: server, config: cfg, @@ -85,9 +85,6 @@ func (s *Server) RegisterTools(ctx context.Context) error { func (s *Server) registerCLIToolsProvider(ctx context.Context) error { log.Info(ctx, "Registering CLI tools provider") - // Add session to context - ctx = session.WithSession(ctx, s.session) - provider, err := clitools.NewProvider(ctx, s.config, s.session) if err != nil { return err diff --git a/experimental/apps-mcp/templates/appkit/databricks_template_schema.json b/experimental/apps-mcp/templates/appkit/databricks_template_schema.json index dc9f86e7ce..110e5b182f 100644 --- a/experimental/apps-mcp/templates/appkit/databricks_template_schema.json +++ b/experimental/apps-mcp/templates/appkit/databricks_template_schema.json @@ -12,13 +12,18 @@ "description": "SQL Warehouse ID", "order": 2 }, + "profile": { + "type": "string", + "description": "Profile Name", + "default": "", + "order": 3 + }, "app_description": { "type": "string", "description": "App Description (Optional)", "default": "A Databricks App powered by Databricks AppKit", - "order": 3 + "order": 4 } - }, - "success_message": "\nYour new project has been created in the '{{.project_name}}' directory!" + "success_message": "\nYour new project has been created in the '{{.project_name}}' directory!\nYOU MUST read {{.project_name}}/CLAUDE.md immediately. It is STRONGLY RECOMMENDED to immediately run `npm install`, run `npm run dev` in the background, and open http://localhost:8000 in your browser before making changes to the app." } diff --git a/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/.env.tmpl b/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/.env.tmpl index 5d09570971..be54897988 100644 --- a/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/.env.tmpl +++ b/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/.env.tmpl @@ -1,4 +1,4 @@ -DATABRICKS_HOST={{workspace_host}} +{{if ne .profile ""}}DATABRICKS_CONFIG_PROFILE={{.profile}}{{else}}DATABRICKS_HOST={{workspace_host}}{{end}} DATABRICKS_WAREHOUSE_ID={{.sql_warehouse_id}} DATABRICKS_APP_PORT=8000 DATABRICKS_APP_NAME=minimal diff --git a/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/CLAUDE.md b/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/CLAUDE.md index 2984603e28..b239b0a4ba 100644 --- a/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/CLAUDE.md +++ b/experimental/apps-mcp/templates/appkit/template/{{.project_name}}/CLAUDE.md @@ -5,6 +5,38 @@ TypeScript full-stack template powered by **Databricks AppKit** with tRPC for ad - config/queries/: SQL query files for analytics - shared/: Shared TypeScript types +## NPM Scripts + +### Development +- `npm run dev` - Start dev server with hot reload (use during active development) +- `npm start` - Start production server (requires `npm run build` first) + +### Build +- `npm run build` - Full build (server + client) - use before deployment +- `npm run build:server` - Compile server TypeScript only +- `npm run build:client` - Compile and bundle React app only + +### Code Quality +- `npm run typecheck` - Type-check without building (fast validation) +- `npm run lint` - Check for linting issues +- `npm run lint:fix` - Auto-fix linting issues +- `npm run format` - Check code formatting +- `npm run format:fix` - Auto-format code with Prettier + +### Testing +- `npm test` - Run unit tests (vitest) + smoke test +- `npm run test:e2e` - Run all Playwright tests +- `npm run test:e2e:ui` - Run Playwright with interactive UI (for debugging) +- `npm run test:smoke` - Run smoke test only (quick validation) + +### Utility +- `npm run clean` - Remove all build artifacts and node_modules (requires `npm install` after) + +**Common workflows:** +- Development: `npm run dev` → make changes → `npm run typecheck` → `npm run lint:fix` +- Pre-commit: `npm run typecheck && npm run lint:fix && npm run format:fix && npm test` +- Pre-deploy: `npm run build && npm start` (test locally) → `npm test` + ## App Naming Constraints App names must not exceed 30 characters total (including target prefix).