Pass all the parameters allowed by streamablehttp_client via DatabricksMCPClient for better control over headers, timeout, etc

Open Beginner friendly
#253 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
72/100
Issue type
Feature
Clarity
Clearly specified
Activity status
Quiet
Tech stack
python

Research direction

Start in databricks_mcp/src/databricks_mcp/mcp.py around line 135 and inspect the DatabricksMCPClient methods that use streamablehttp_client. Check the streamablehttp_client parameters, then verify that options such as headers, timeout, and terminate_on_close can be passed through and that the shown list_tools call works without the zstd decoding error.

Written by the indexing model from the issue text.

Description

We need to allow more parameters (e.g. headers, timeout, terminate_on_close) to be passed into streamablehttp_client used by DatabricksMCPClient. Patching the DatabricksMCPClient methods using streamablehttp_client with **kwargs solved the following zstd decoding error.

Running:

from databricks_mcp import DatabricksMCPClient
import nest_asyncio
import asyncio

nest_asyncio.apply()

server_url = "https://e2-demo-field-eng.cloud.databricks.com/api/2.0/mcp/external/conn_aichemy_pubchem"
#server_url = "https://mcp-nitin-1444828305810485.aws.databricksapps.com/mcp"
mcp_client = DatabricksMCPClient(server_url=server_url, workspace_client=ws_client)
mcp_client.list_tools()

Error returned:

Error parsing JSON response
Traceback (most recent call last):
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-588c2647-b519-45f9-bf41-fa08778d1b83/lib/python3.12/site-packages/httpx/_decoders.py", line 185, in decode
    output.write(self.decompressor.decompress(data))
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
zstd.ZstdError: zstd decompressor error: Unknown frame descriptor

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-588c2647-b519-45f9-bf41-fa08778d1b83/lib/python3.12/site-packages/mcp/client/streamable_http.py", line 384, in _handle_json_response
    content = await response.aread()
              ^^^^^^^^^^^^^^^^^^^^^^
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-588c2647-b519-45f9-bf41-fa08778d1b83/lib/python3.12/site-packages/httpx/_models.py", line 979, in aread
    self._content = b"".join([part async for part in self.aiter_bytes()])
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-588c2647-b519-45f9-bf41-fa08778d1b83/lib/python3.12/site-packages/httpx/_models.py", line 998, in aiter_bytes
    decoded = decoder.decode(raw_bytes)
              ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-588c2647-b519-45f9-bf41-fa08778d1b83/lib/python3.12/site-packages/httpx/_decoders.py", line 191, in decode
    raise DecodingError(str(exc)) from exc
httpx.DecodingError: zstd decompressor error: Unknown frame descriptor

Solution:
Adding a specific header solved the above error. Please also pass in kwargs like timeout for greater control.
https://github.com/databricks/databricks-ai-bridge/blob/main/databricks_mcp/src/databricks_mcp/mcp.py#L135

    async def _get_tools_async(self, **kwargs) -> List[Tool]:
        """Fetch tools from the MCP endpoint asynchronously."""
        async with streamablehttp_client(
            url=self.server_url,
            auth=DatabricksOAuthClientProvider(self.client),
            headers={"Accept-Encoding": "identity"},
            **kwargs,
        ) as (read_stream, write_stream, _):
            async with ClientSession(read_stream, write_stream) as session:
                await session.initialize()
                return (await session.list_tools()).tools

mcp_client.list_tools(timeout=60, terminate_on_close=False) now successfully returns:

[Tool(name='search_compounds', title=None, description='Search PubChem database for compounds by name, CAS number, formula, or identifier', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'description': 'Search query (compound name, CAS, formula, or identifier)'}, 'search_type': {'type': 'string', 'enum': ['name', 'smiles', 'inchi', 'sdf', 'cid', 'formula'], 'description': 'Type of search to perform (default: name)'}, 'max_records': {'type': 'number', 'description': 'Maximum number of results (1-10000, default: 100)', 'minimum': 1, 'maximum': 10000}}, 'required': ['query']}, outputSchema=None, icons=None, annotations=None, meta=None, execution=None),
...]
Dominant language
Python
Stars
103
Forks
83
Avg merge
23h 6m
Merged PRs (30d)
81

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from databricks/databricks-ai-bridge

All issues in databricks/databricks-ai-bridge

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.