Optimize MCP tools with integration tests and bug fixes by QuentinAmbard · Pull Request #394 · databricks-solutions/ai-dev-kit

QuentinAmbard · 2026-03-30T08:22:16Z

Summary

This PR optimizes MCP tools for better AI efficiency through docstring compression, tool consolidation, and migrates Knowledge Assistant operations to the Python SDK.

Key Changes

1. Tool Consolidation (77 → 44 tools, 43% reduction)

Module	Before	After	Consolidated Into
pipelines.py	10	2	`manage_pipeline`, `manage_pipeline_run`
volume_files.py	6	1	`manage_volume_files`
vector_search.py	8	4	`manage_vs_endpoint`, `manage_vs_index`, `query_vs_index`, `manage_vs_data`
genie.py	5	2	`manage_genie`, `ask_genie`
aibi_dashboards.py	4	1	`manage_dashboard`
serving.py	3	1	`manage_serving_endpoint`
apps.py	3	1	`manage_app`
file.py	2	1	`manage_workspace_files`
sql.py	6	5	`manage_warehouse` replaces list/get_best
lakebase.py	8	4	`manage_lakebase_database`, `manage_lakebase_branch`, `manage_lakebase_sync`, `generate_lakebase_credential`

2. Docstring Optimization (~89% reduction in documentation size)

Reduced verbosity across all 18 tool files
Maintained all functional information while being concise
Added skill references to complex tools for detailed documentation

3. Knowledge Assistant SDK Migration

Migrated ka_create, ka_get, ka_sync_sources to Python SDK
Keep ka_update using raw API 2.1 due to SDK FieldMask bug
Fixed name sanitization in find_by_name (spaces→underscores)
Fixed ka_create_or_update to lookup by name when no tile_id provided
Added default description for knowledge sources (API requirement)

4. Bug Fixes (discovered via integration tests)

Component	Bug	Fix
Apps	`deploy_app` returned `Wait[AppDeployment]` object	Use `wait_obj.response` to get actual deployment object
Genie	Used manager methods that don't exist in SDK	Use SDK methods directly
Workspace Files	Directory upload failed	Fixed upload for directories
KA find_by_name	Name not sanitized before lookup	Sanitize names before API lookup
KA create_or_update	Didn't lookup by name, caused ALREADY_EXISTS	Added name lookup when no tile_id
KA knowledge source	Empty description rejected by API	Provide default description

5. Integration Test Infrastructure

Added comprehensive integration tests for all MCP tools
Added test runner script with parallel execution and per-suite timeouts
Improved test reliability with quota skip handling

Running Integration Tests

cd databricks-mcp-server

# Run all integration tests
python tests/integration/run_tests.py --all

# Run specific test suites
python tests/integration/run_tests.py --suites genie,pipelines,jobs

# Run agent bricks tests
python -m pytest tests/integration/agent_bricks -m integration -v

Test plan

All integration tests pass
Verify AI can still use tools correctly with compressed docstrings
Test KA create/update/get operations with SDK migration
Test KA name lookup and create_or_update on existing KA

🤖 Generated with Claude Code

- Reduce docstring verbosity across all 18 tool files (~89% reduction) - Keep all functional information while being concise - Add skill references to complex tools (dashboards, vector search, genie, jobs, pipelines, lakebase, unity catalog, serving, apps, agent bricks) - Maintain human readability with bullet points and structure - Preserve critical warnings (ASK USER FIRST, CONFIRM WITH USER) - Keep return format hints for AI parsing Net reduction: 1,843 lines across 18 files 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

- agent_bricks.py: Add context for description, instructions, volume_path, examples - genie.py: Add context for table_identifiers, description, sample_questions, serialized_space - jobs.py: Add context for tasks, job_clusters, environments, schedule, git_source - lakebase.py: Add context for source_branch, ttl_seconds, is_protected, autoscaling params, and sync source/target table names - pipelines.py: Add context for root_path, workspace_file_paths, extra_settings, full_refresh 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Tool consolidations: - pipelines.py: 10→2 (manage_pipeline, manage_pipeline_run) - volume_files.py: 6→1 (manage_volume_files) - aibi_dashboards.py: 4→1 (manage_dashboard) - vector_search.py: 8→4 (manage_vs_endpoint, manage_vs_index, query_vs_index, manage_vs_data) - genie.py: 5→2 (manage_genie, ask_genie) - serving.py: 3→1 (manage_serving_endpoint) - apps.py: 3→1 (manage_app) - file.py: 2→1 (manage_workspace_files) - sql.py: 6→5 (manage_warehouse replaces list/get_best) - lakebase.py: 8→4 (manage_lakebase_database, manage_lakebase_branch, manage_lakebase_sync, generate_lakebase_credential) Key patterns: - All consolidated tools use an action parameter - Each action has required params documented in docstring - Error messages specify which params are required - Hot paths (query_vs_index, ask_genie) kept separate for clarity - All skills updated with action tables and examples 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Test infrastructure: - Add comprehensive integration tests for all MCP tools - Add test runner script with parallel execution support - Add fixtures for workspace, catalog, and resource cleanup - Add test resources (PDFs, SQL files, app configs) Bug fixes in databricks-tools-core: - Fix workspace file upload for directories - Fix job notebook path handling - Fix vector search index operations - Fix apps API responses - Fix dashboard widget handling - Fix agent bricks manager listing Bug fixes in MCP server tools: - Add quota skip handling for apps test - Fix genie space operations - Fix lakebase database operations - Fix compute cluster lifecycle handling - Fix dashboard operations 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

The Databricks SDK's w.apps.deploy() returns a Wait[AppDeployment] object, not an AppDeployment directly. The previous code passed the Wait object to _deployment_to_dict(), which caused getattr() to return None for all attributes since the Wait object doesn't have them. This fix uses wait_obj.response to get the actual AppDeployment object before converting it to a dictionary. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

QuentinAmbard · 2026-04-01T07:18:06Z

Recommended workflow:

Merge PR Optimize MCP tools with integration tests and bug fixes #394 first (MCP tools + apps fix + integration tests)
Rebase PR Enhance AI/BI dashboard skill with comprehensive widget documentation #374 on the updated main - Git will automatically drop the MCP commits since they're already in main
PR Enhance AI/BI dashboard skill with comprehensive widget documentation #374's diff will then only show dashboard skill improvements

- Add tools summary table at top of MCP Tools section - Change code blocks from python syntax to plain text - Add "# MCP Tool: <name>" comments to clarify these are tool calls, not Python code - Move Supporting Tools table to main tools table 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Remove garbage characters from widget documentation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

- Use w.genie.trash_space() in _delete_genie_resource - Add _find_space_by_name() using SDK's list_spaces with pagination - Use w.genie.update_space() and w.genie.create_space() for space management - Use w.genie.get_space() with include_serialized_space in _get_genie_space - Fix validation to allow space_id for updates without display_name 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

- Add per-suite timeout in run_tests.py (10 min default, configurable) - Improve apps test with better cleanup and assertions - Add skip logic for quota-exceeded scenarios 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

…cumentation Add detailed parameter documentation to all 9 Unity Catalog MCP tools: - manage_uc_objects: Document parameters by object_type (catalog/schema/volume/function) - manage_uc_grants: Add privilege lists per securable type - manage_uc_storage: Detail credential and external_location parameters - manage_uc_connections: Document connection_type options and create_foreign_catalog - manage_uc_tags: Detail set_tags/unset_tags/query parameters - manage_uc_security_policies: Document row filter and column mask parameters - manage_uc_monitors: Detail monitor creation and refresh parameters - manage_uc_sharing: Document share/recipient/provider resource types - manage_metric_views: Detail dimension/measure format and query parameters 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Add clear instructions requiring users to: 0. Review the databricks-aibi-dashboards skill for widget JSON structure 1. Call get_table_stats_and_schema() for table schemas 2. Call execute_sql() to test EVERY query before use This prevents widgets from showing errors due to untested queries. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

calreynolds

slack

- Migrate ka_create, ka_get, ka_sync_sources to use Python SDK - Keep ka_update using raw API 2.1 due to SDK FieldMask bug (converts snake_case to camelCase but API expects snake_case) - Fix find_by_name to sanitize names (spaces→underscores) before lookup - Fix ka_create_or_update to lookup by name when no tile_id provided, preventing ALREADY_EXISTS errors on repeated calls - Update MCP tool layer to use new flat response format - Map SDK state values (ACTIVE, CREATING, FAILED) to endpoint_status - Add integration test for updating existing KA via create_or_update 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

- Provide default description for knowledge sources when not specified (API requires non-empty knowledge_source.description) - Move KA update test to after endpoint is ONLINE (update requires ACTIVE state) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Quentin Ambard and others added 5 commits March 30, 2026 10:20

QuentinAmbard mentioned this pull request Apr 1, 2026

Enhance AI/BI dashboard skill with comprehensive widget documentation #374

Open

4 tasks

Quentin Ambard and others added 6 commits April 1, 2026 09:23

Fix typo in aibi_dashboards.py docstring

fe7ea10

Remove garbage characters from widget documentation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

calreynolds self-requested a review April 1, 2026 14:49

calreynolds requested changes Apr 1, 2026

View reviewed changes

QuentinAmbard mentioned this pull request Apr 2, 2026

Improve code execution and synthetic data generation skills #403

Open

4 tasks

Quentin Ambard and others added 2 commits April 2, 2026 12:31

QuentinAmbard changed the title ~~Optimize MCP tool docstrings for token efficiency~~ Optimize MCP tools and migrate KA to Python SDK Apr 2, 2026

QuentinAmbard changed the title ~~Optimize MCP tools and migrate KA to Python SDK~~ Optimize MCP tools with integration tests and bug fixes Apr 2, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Optimize MCP tools with integration tests and bug fixes#394

Optimize MCP tools with integration tests and bug fixes#394
QuentinAmbard wants to merge 13 commits intomainfrom
mcp-tool-token-optimization

QuentinAmbard commented Mar 30, 2026 •

edited

Loading

Uh oh!

QuentinAmbard commented Apr 1, 2026

Uh oh!

calreynolds left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

QuentinAmbard commented Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Key Changes

Running Integration Tests

Test plan

Uh oh!

QuentinAmbard commented Apr 1, 2026

Uh oh!

calreynolds left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

QuentinAmbard commented Mar 30, 2026 •

edited

Loading