Skip to content

MCP Support in Workspace Client #981

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from

Conversation

aravind-segu
Copy link
Contributor

What changes are proposed in this pull request?

  • Anthropic recently released MCP Servers and we are currently adding support for First Party (MCP Servers on Control Plane) and Third Party (MCP Servers on Databricks Apps).
  • In order to authenticate to MCP servers, the MCP Python-sdk provides us with a OAuthTokenProvider class that we need to implement
  • This pr uses the existing WorkspaceClient to populate the token in the Authorization headers of the OAuthTokenProvider so we can use it with the streamable_http client.

How is this tested?

Added Unit tests.
E2E Test

import asyncio
from datetime import datetime, timedelta
from typing import Optional

from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from databricks import sdk

workspace_client = sdk.WorkspaceClient(profile="dogfood")
async def query_mcp():
    async with streamablehttp_client(
        url="https://uc-mcp-server-aravind-6051921418418893.staging.aws.databricksapps.com/mcp/",
        auth=workspace_client.mcp.get_oauth_provider(),
    ) as (read_stream, write_stream, _):
        async with ClientSession(read_stream, write_stream) as session:
            await session.initialize()
            tools = await session.list_tools()
            print(tools)


async def main():
    await query_mcp()


if __name__ == "__main__":
    asyncio.run(main())

Running the above code with a databricks-cli authentication returns successfully
image

Copy link

Please ensure that the NEXT_CHANGELOG.md file is updated with any relevant changes.
If this is not necessary for your PR, please include the following in your PR description:
NO_CHANGELOG=true
and rerun the job.

Copy link

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-py

Inputs:

  • PR number: 981
  • Commit SHA: fd36b9f892a7dee99a1e212014da8149d70c5519

Checks will be approved automatically on success.

@renaudhartert-db
Copy link
Contributor

Closing this PR as per offline discussions.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants