Skip to content

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

@github-actions
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.

@github-actions
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