uc-mcp-proxy 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- uc_mcp_proxy-0.1.0/.github/workflows/ci.yml +29 -0
- uc_mcp_proxy-0.1.0/.github/workflows/publish.yml +23 -0
- uc_mcp_proxy-0.1.0/.gitignore +31 -0
- uc_mcp_proxy-0.1.0/CLAUDE.md +36 -0
- uc_mcp_proxy-0.1.0/LICENSE +21 -0
- uc_mcp_proxy-0.1.0/PKG-INFO +106 -0
- uc_mcp_proxy-0.1.0/README.md +81 -0
- uc_mcp_proxy-0.1.0/pyproject.toml +72 -0
- uc_mcp_proxy-0.1.0/src/uc_mcp_proxy/__init__.py +3 -0
- uc_mcp_proxy-0.1.0/src/uc_mcp_proxy/__main__.py +97 -0
- uc_mcp_proxy-0.1.0/src/uc_mcp_proxy/py.typed +0 -0
- uc_mcp_proxy-0.1.0/tests/__init__.py +0 -0
- uc_mcp_proxy-0.1.0/tests/conftest.py +165 -0
- uc_mcp_proxy-0.1.0/tests/integration/__init__.py +0 -0
- uc_mcp_proxy-0.1.0/tests/integration/test_proxy.py +135 -0
- uc_mcp_proxy-0.1.0/tests/unit/__init__.py +0 -0
- uc_mcp_proxy-0.1.0/tests/unit/test_auth.py +85 -0
- uc_mcp_proxy-0.1.0/tests/unit/test_bridge.py +116 -0
- uc_mcp_proxy-0.1.0/tests/unit/test_cli.py +67 -0
- uc_mcp_proxy-0.1.0/uv.lock +1262 -0
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [master]
|
|
6
|
+
pull_request:
|
|
7
|
+
branches: [master]
|
|
8
|
+
|
|
9
|
+
jobs:
|
|
10
|
+
test:
|
|
11
|
+
runs-on: ubuntu-latest
|
|
12
|
+
strategy:
|
|
13
|
+
matrix:
|
|
14
|
+
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
|
15
|
+
steps:
|
|
16
|
+
- uses: actions/checkout@v4
|
|
17
|
+
|
|
18
|
+
- uses: astral-sh/setup-uv@v4
|
|
19
|
+
with:
|
|
20
|
+
python-version: ${{ matrix.python-version }}
|
|
21
|
+
|
|
22
|
+
- name: Install dependencies
|
|
23
|
+
run: uv sync
|
|
24
|
+
|
|
25
|
+
- name: Run unit tests
|
|
26
|
+
run: uv run pytest -m unit -v --cov
|
|
27
|
+
|
|
28
|
+
- name: Build package
|
|
29
|
+
run: uv build
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
tags:
|
|
6
|
+
- "v*"
|
|
7
|
+
|
|
8
|
+
jobs:
|
|
9
|
+
publish:
|
|
10
|
+
runs-on: ubuntu-latest
|
|
11
|
+
environment: pypi
|
|
12
|
+
permissions:
|
|
13
|
+
id-token: write
|
|
14
|
+
steps:
|
|
15
|
+
- uses: actions/checkout@v4
|
|
16
|
+
|
|
17
|
+
- uses: astral-sh/setup-uv@v4
|
|
18
|
+
|
|
19
|
+
- name: Build package
|
|
20
|
+
run: uv build
|
|
21
|
+
|
|
22
|
+
- name: Publish to PyPI
|
|
23
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
*.egg-info/
|
|
6
|
+
*.egg
|
|
7
|
+
dist/
|
|
8
|
+
build/
|
|
9
|
+
|
|
10
|
+
# Virtual environments
|
|
11
|
+
.venv/
|
|
12
|
+
venv/
|
|
13
|
+
env/
|
|
14
|
+
|
|
15
|
+
# Testing
|
|
16
|
+
.coverage
|
|
17
|
+
.pytest_cache/
|
|
18
|
+
htmlcov/
|
|
19
|
+
|
|
20
|
+
# IDE
|
|
21
|
+
.idea/
|
|
22
|
+
.vscode/
|
|
23
|
+
*.swp
|
|
24
|
+
*.swo
|
|
25
|
+
|
|
26
|
+
# OS
|
|
27
|
+
.DS_Store
|
|
28
|
+
Thumbs.db
|
|
29
|
+
|
|
30
|
+
# uv
|
|
31
|
+
.python-version
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
# uc-mcp-proxy
|
|
2
|
+
|
|
3
|
+
MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth.
|
|
4
|
+
|
|
5
|
+
## Commands
|
|
6
|
+
|
|
7
|
+
- `uv sync` — install all dependencies (including dev)
|
|
8
|
+
- `uv run pytest -m unit -v` — run unit tests only
|
|
9
|
+
- `uv run pytest -m integration -v` — run integration tests only
|
|
10
|
+
- `uv run pytest -v` — run all tests
|
|
11
|
+
- `uv run pytest -m unit --cov -v` — unit tests with coverage
|
|
12
|
+
- `uv build` — build sdist + wheel into `dist/`
|
|
13
|
+
|
|
14
|
+
## Architecture
|
|
15
|
+
|
|
16
|
+
Single-module package in `src/uc_mcp_proxy/`:
|
|
17
|
+
|
|
18
|
+
- `__main__.py` — CLI entry point, `DatabricksAuth` (httpx auth flow), `bridge()` (bidirectional stdio↔HTTP stream copy), `run()` (async main)
|
|
19
|
+
- `__init__.py` — re-exports `DatabricksAuth`
|
|
20
|
+
|
|
21
|
+
The proxy bridges an MCP stdio transport to a remote Streamable HTTP MCP server, injecting Databricks OAuth tokens on every request via `DatabricksAuth`.
|
|
22
|
+
|
|
23
|
+
## Testing
|
|
24
|
+
|
|
25
|
+
Tests live in `tests/` with two marker categories:
|
|
26
|
+
|
|
27
|
+
- `unit` — pure unit tests, no external dependencies, fast
|
|
28
|
+
- `integration` — full proxy flow tests with mocked transports
|
|
29
|
+
|
|
30
|
+
All new code must have unit tests. Maintain ≥75% coverage (`fail_under = 75` in pyproject.toml).
|
|
31
|
+
|
|
32
|
+
## Code Style
|
|
33
|
+
|
|
34
|
+
- Use `from __future__ import annotations` in all modules
|
|
35
|
+
- Type hints on all public functions
|
|
36
|
+
- Keep imports sorted: stdlib → third-party → local
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Tanner Wendland
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: uc-mcp-proxy
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth
|
|
5
|
+
Project-URL: Homepage, https://github.com/IceRhymers/uc-mcp-proxy
|
|
6
|
+
Project-URL: Repository, https://github.com/IceRhymers/uc-mcp-proxy
|
|
7
|
+
Project-URL: Issues, https://github.com/IceRhymers/uc-mcp-proxy/issues
|
|
8
|
+
Author: Tanner Wendland
|
|
9
|
+
License-Expression: MIT
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Keywords: databricks,mcp,oauth,proxy,unity-catalog
|
|
12
|
+
Classifier: Development Status :: 4 - Beta
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
19
|
+
Requires-Python: >=3.10
|
|
20
|
+
Requires-Dist: anyio
|
|
21
|
+
Requires-Dist: databricks-sdk>=0.30.0
|
|
22
|
+
Requires-Dist: httpx
|
|
23
|
+
Requires-Dist: mcp>=1.8
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
|
|
26
|
+
# uc-mcp-proxy
|
|
27
|
+
|
|
28
|
+
MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth.
|
|
29
|
+
|
|
30
|
+
Lets any MCP client that speaks **stdio** (e.g. Claude Desktop, Claude Code) connect to a remote **Streamable HTTP** MCP server hosted on Databricks Apps — handling OAuth authentication automatically.
|
|
31
|
+
|
|
32
|
+
## Installation
|
|
33
|
+
|
|
34
|
+
```bash
|
|
35
|
+
# Run directly (no install needed)
|
|
36
|
+
uvx uc-mcp-proxy --url https://<workspace>.databricks.com/apps/<app>/mcp
|
|
37
|
+
|
|
38
|
+
# Or install globally
|
|
39
|
+
uv tool install uc-mcp-proxy
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
Requires Python 3.10+.
|
|
43
|
+
|
|
44
|
+
## Usage
|
|
45
|
+
|
|
46
|
+
### Claude Desktop / Claude Code (`.mcp.json`)
|
|
47
|
+
|
|
48
|
+
Add to your MCP client configuration:
|
|
49
|
+
|
|
50
|
+
```json
|
|
51
|
+
{
|
|
52
|
+
"mcpServers": {
|
|
53
|
+
"unity-catalog": {
|
|
54
|
+
"type": "stdio",
|
|
55
|
+
"command": "uvx",
|
|
56
|
+
"args": [
|
|
57
|
+
"uc-mcp-proxy",
|
|
58
|
+
"--url", "https://<workspace>.databricks.com/apps/<app>/mcp"
|
|
59
|
+
]
|
|
60
|
+
}
|
|
61
|
+
}
|
|
62
|
+
}
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
### CLI
|
|
66
|
+
|
|
67
|
+
```bash
|
|
68
|
+
uc-mcp-proxy --url <MCP_SERVER_URL> [--profile <DATABRICKS_PROFILE>] [--auth-type <AUTH_TYPE>]
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
| Flag | Description |
|
|
72
|
+
|------|-------------|
|
|
73
|
+
| `--url` | **(required)** Remote MCP server URL |
|
|
74
|
+
| `--profile` | Databricks CLI profile name (uses default if omitted) |
|
|
75
|
+
| `--auth-type` | Databricks auth type, e.g. `databricks-cli` |
|
|
76
|
+
|
|
77
|
+
## How It Works
|
|
78
|
+
|
|
79
|
+
1. Starts an MCP **stdio** server (stdin/stdout)
|
|
80
|
+
2. Connects to the remote MCP server via **Streamable HTTP**
|
|
81
|
+
3. Injects a fresh Databricks OAuth token on every HTTP request
|
|
82
|
+
4. Bridges messages bidirectionally between the two transports
|
|
83
|
+
|
|
84
|
+
## OAuth Authentication
|
|
85
|
+
|
|
86
|
+
Authentication is handled by the [Databricks SDK](https://docs.databricks.com/dev-tools/sdk-python.html), which supports multiple auth methods:
|
|
87
|
+
|
|
88
|
+
- **Databricks CLI** (`databricks-cli`) — uses tokens from `~/.databrickscfg`
|
|
89
|
+
- **OAuth U2M** — browser-based login flow
|
|
90
|
+
- **PAT** — personal access tokens
|
|
91
|
+
- **Azure / GCP / AWS** — cloud-native identity
|
|
92
|
+
|
|
93
|
+
The SDK auto-detects the method, or you can force one with `--auth-type`.
|
|
94
|
+
|
|
95
|
+
## Development
|
|
96
|
+
|
|
97
|
+
```bash
|
|
98
|
+
uv sync # install dependencies
|
|
99
|
+
uv run pytest -m unit -v # run unit tests
|
|
100
|
+
uv run pytest -m integration -v # run integration tests
|
|
101
|
+
uv build # build package
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
## License
|
|
105
|
+
|
|
106
|
+
MIT
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
# uc-mcp-proxy
|
|
2
|
+
|
|
3
|
+
MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth.
|
|
4
|
+
|
|
5
|
+
Lets any MCP client that speaks **stdio** (e.g. Claude Desktop, Claude Code) connect to a remote **Streamable HTTP** MCP server hosted on Databricks Apps — handling OAuth authentication automatically.
|
|
6
|
+
|
|
7
|
+
## Installation
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
# Run directly (no install needed)
|
|
11
|
+
uvx uc-mcp-proxy --url https://<workspace>.databricks.com/apps/<app>/mcp
|
|
12
|
+
|
|
13
|
+
# Or install globally
|
|
14
|
+
uv tool install uc-mcp-proxy
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
Requires Python 3.10+.
|
|
18
|
+
|
|
19
|
+
## Usage
|
|
20
|
+
|
|
21
|
+
### Claude Desktop / Claude Code (`.mcp.json`)
|
|
22
|
+
|
|
23
|
+
Add to your MCP client configuration:
|
|
24
|
+
|
|
25
|
+
```json
|
|
26
|
+
{
|
|
27
|
+
"mcpServers": {
|
|
28
|
+
"unity-catalog": {
|
|
29
|
+
"type": "stdio",
|
|
30
|
+
"command": "uvx",
|
|
31
|
+
"args": [
|
|
32
|
+
"uc-mcp-proxy",
|
|
33
|
+
"--url", "https://<workspace>.databricks.com/apps/<app>/mcp"
|
|
34
|
+
]
|
|
35
|
+
}
|
|
36
|
+
}
|
|
37
|
+
}
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
### CLI
|
|
41
|
+
|
|
42
|
+
```bash
|
|
43
|
+
uc-mcp-proxy --url <MCP_SERVER_URL> [--profile <DATABRICKS_PROFILE>] [--auth-type <AUTH_TYPE>]
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
| Flag | Description |
|
|
47
|
+
|------|-------------|
|
|
48
|
+
| `--url` | **(required)** Remote MCP server URL |
|
|
49
|
+
| `--profile` | Databricks CLI profile name (uses default if omitted) |
|
|
50
|
+
| `--auth-type` | Databricks auth type, e.g. `databricks-cli` |
|
|
51
|
+
|
|
52
|
+
## How It Works
|
|
53
|
+
|
|
54
|
+
1. Starts an MCP **stdio** server (stdin/stdout)
|
|
55
|
+
2. Connects to the remote MCP server via **Streamable HTTP**
|
|
56
|
+
3. Injects a fresh Databricks OAuth token on every HTTP request
|
|
57
|
+
4. Bridges messages bidirectionally between the two transports
|
|
58
|
+
|
|
59
|
+
## OAuth Authentication
|
|
60
|
+
|
|
61
|
+
Authentication is handled by the [Databricks SDK](https://docs.databricks.com/dev-tools/sdk-python.html), which supports multiple auth methods:
|
|
62
|
+
|
|
63
|
+
- **Databricks CLI** (`databricks-cli`) — uses tokens from `~/.databrickscfg`
|
|
64
|
+
- **OAuth U2M** — browser-based login flow
|
|
65
|
+
- **PAT** — personal access tokens
|
|
66
|
+
- **Azure / GCP / AWS** — cloud-native identity
|
|
67
|
+
|
|
68
|
+
The SDK auto-detects the method, or you can force one with `--auth-type`.
|
|
69
|
+
|
|
70
|
+
## Development
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
uv sync # install dependencies
|
|
74
|
+
uv run pytest -m unit -v # run unit tests
|
|
75
|
+
uv run pytest -m integration -v # run integration tests
|
|
76
|
+
uv build # build package
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
## License
|
|
80
|
+
|
|
81
|
+
MIT
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "uc-mcp-proxy"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth"
|
|
5
|
+
requires-python = ">=3.10"
|
|
6
|
+
dependencies = [
|
|
7
|
+
"mcp>=1.8",
|
|
8
|
+
"databricks-sdk>=0.30.0",
|
|
9
|
+
"httpx",
|
|
10
|
+
"anyio",
|
|
11
|
+
]
|
|
12
|
+
license = "MIT"
|
|
13
|
+
readme = "README.md"
|
|
14
|
+
authors = [
|
|
15
|
+
{ name = "Tanner Wendland" },
|
|
16
|
+
]
|
|
17
|
+
keywords = ["mcp", "databricks", "oauth", "proxy", "unity-catalog"]
|
|
18
|
+
classifiers = [
|
|
19
|
+
"Development Status :: 4 - Beta",
|
|
20
|
+
"License :: OSI Approved :: MIT License",
|
|
21
|
+
"Programming Language :: Python :: 3",
|
|
22
|
+
"Programming Language :: Python :: 3.10",
|
|
23
|
+
"Programming Language :: Python :: 3.11",
|
|
24
|
+
"Programming Language :: Python :: 3.12",
|
|
25
|
+
"Programming Language :: Python :: 3.13",
|
|
26
|
+
]
|
|
27
|
+
|
|
28
|
+
[dependency-groups]
|
|
29
|
+
dev = [
|
|
30
|
+
"pytest>=8.0",
|
|
31
|
+
"pytest-cov>=5.0",
|
|
32
|
+
"pytest-randomly>=3.15",
|
|
33
|
+
"anyio[trio]>=4.0",
|
|
34
|
+
]
|
|
35
|
+
|
|
36
|
+
[project.scripts]
|
|
37
|
+
uc-mcp-proxy = "uc_mcp_proxy.__main__:main"
|
|
38
|
+
|
|
39
|
+
[project.urls]
|
|
40
|
+
Homepage = "https://github.com/IceRhymers/uc-mcp-proxy"
|
|
41
|
+
Repository = "https://github.com/IceRhymers/uc-mcp-proxy"
|
|
42
|
+
Issues = "https://github.com/IceRhymers/uc-mcp-proxy/issues"
|
|
43
|
+
|
|
44
|
+
[build-system]
|
|
45
|
+
requires = ["hatchling"]
|
|
46
|
+
build-backend = "hatchling.build"
|
|
47
|
+
|
|
48
|
+
[tool.hatch.build.targets.wheel]
|
|
49
|
+
packages = ["src/uc_mcp_proxy"]
|
|
50
|
+
|
|
51
|
+
[tool.pytest.ini_options]
|
|
52
|
+
testpaths = ["tests"]
|
|
53
|
+
markers = [
|
|
54
|
+
"unit: Pure unit tests with no external dependencies",
|
|
55
|
+
"integration: Tests that exercise the full proxy flow with mocked transports",
|
|
56
|
+
]
|
|
57
|
+
addopts = "-v --tb=short --strict-markers"
|
|
58
|
+
pythonpath = ["src"]
|
|
59
|
+
|
|
60
|
+
[tool.coverage.run]
|
|
61
|
+
source = ["uc_mcp_proxy"]
|
|
62
|
+
branch = true
|
|
63
|
+
omit = ["tests/*"]
|
|
64
|
+
|
|
65
|
+
[tool.coverage.report]
|
|
66
|
+
show_missing = true
|
|
67
|
+
fail_under = 75
|
|
68
|
+
exclude_lines = [
|
|
69
|
+
"pragma: no cover",
|
|
70
|
+
"if __name__ == .__main__.",
|
|
71
|
+
"if TYPE_CHECKING:",
|
|
72
|
+
]
|
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
"""MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import asyncio
|
|
7
|
+
from typing import Generator, AsyncGenerator
|
|
8
|
+
|
|
9
|
+
import anyio
|
|
10
|
+
import httpx
|
|
11
|
+
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
|
12
|
+
from databricks.sdk import WorkspaceClient
|
|
13
|
+
from mcp.client.streamable_http import streamablehttp_client
|
|
14
|
+
from mcp.server.stdio import stdio_server
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class DatabricksAuth(httpx.Auth):
|
|
18
|
+
"""httpx Auth that injects fresh Databricks OAuth tokens per-request.
|
|
19
|
+
|
|
20
|
+
Calls ``WorkspaceClient.config.authenticate()`` on every request to obtain
|
|
21
|
+
a current OAuth bearer token, ensuring tokens are never stale.
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
def __init__(self, client: WorkspaceClient) -> None:
|
|
25
|
+
self._client = client
|
|
26
|
+
|
|
27
|
+
def _apply_headers(self, request: httpx.Request) -> None:
|
|
28
|
+
headers = self._client.config.authenticate()
|
|
29
|
+
request.headers.update(headers)
|
|
30
|
+
# Also forward the token so the Databricks App can use per-user identity
|
|
31
|
+
auth_value = headers.get("Authorization", "")
|
|
32
|
+
if auth_value.startswith("Bearer "):
|
|
33
|
+
request.headers["X-Forwarded-Access-Token"] = auth_value[len("Bearer "):]
|
|
34
|
+
|
|
35
|
+
def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
|
|
36
|
+
self._apply_headers(request)
|
|
37
|
+
yield request
|
|
38
|
+
|
|
39
|
+
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
|
|
40
|
+
self._apply_headers(request)
|
|
41
|
+
yield request
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
async def copy_stream(source: MemoryObjectReceiveStream, dest: MemoryObjectSendStream) -> None:
|
|
45
|
+
"""Copy all messages from source to dest, closing dest when source is exhausted."""
|
|
46
|
+
try:
|
|
47
|
+
async for message in source:
|
|
48
|
+
await dest.send(message)
|
|
49
|
+
finally:
|
|
50
|
+
await dest.aclose()
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
async def bridge(
|
|
54
|
+
stdio_read: MemoryObjectReceiveStream,
|
|
55
|
+
stdio_write: MemoryObjectSendStream,
|
|
56
|
+
http_read: MemoryObjectReceiveStream,
|
|
57
|
+
http_write: MemoryObjectSendStream,
|
|
58
|
+
) -> None:
|
|
59
|
+
"""Bidirectional bridge between stdio and HTTP stream pairs."""
|
|
60
|
+
async with anyio.create_task_group() as tg:
|
|
61
|
+
tg.start_soon(copy_stream, stdio_read, http_write)
|
|
62
|
+
tg.start_soon(copy_stream, http_read, stdio_write)
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
async def run(url: str, profile: str | None = None, auth_type: str | None = None) -> None:
|
|
66
|
+
"""Run the proxy: bridge stdio transport to Streamable HTTP with Databricks OAuth."""
|
|
67
|
+
kwargs: dict = {}
|
|
68
|
+
if profile:
|
|
69
|
+
kwargs["profile"] = profile
|
|
70
|
+
if auth_type:
|
|
71
|
+
kwargs["auth_type"] = auth_type
|
|
72
|
+
client = WorkspaceClient(**kwargs)
|
|
73
|
+
auth = DatabricksAuth(client)
|
|
74
|
+
|
|
75
|
+
async with stdio_server() as (stdio_read, stdio_write):
|
|
76
|
+
async with streamablehttp_client(url, auth=auth) as (
|
|
77
|
+
http_read,
|
|
78
|
+
http_write,
|
|
79
|
+
_get_session_id,
|
|
80
|
+
):
|
|
81
|
+
await bridge(stdio_read, stdio_write, http_read, http_write)
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def main() -> None:
|
|
85
|
+
"""CLI entry point: parse args and run the proxy."""
|
|
86
|
+
parser = argparse.ArgumentParser(
|
|
87
|
+
description="MCP stdio-to-Streamable-HTTP proxy with Databricks OAuth",
|
|
88
|
+
)
|
|
89
|
+
parser.add_argument("--url", required=True, help="Remote MCP server URL")
|
|
90
|
+
parser.add_argument("--profile", default=None, help="Databricks CLI profile")
|
|
91
|
+
parser.add_argument("--auth-type", default=None, help="Databricks auth type (e.g. databricks-cli)")
|
|
92
|
+
args = parser.parse_args()
|
|
93
|
+
asyncio.run(run(args.url, args.profile, args.auth_type))
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
if __name__ == "__main__": # pragma: no cover
|
|
97
|
+
main()
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,165 @@
|
|
|
1
|
+
"""Shared fixtures for uc-mcp-proxy tests."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import pytest
|
|
6
|
+
import httpx
|
|
7
|
+
import anyio
|
|
8
|
+
from contextlib import asynccontextmanager
|
|
9
|
+
from unittest.mock import MagicMock
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
# ---------------------------------------------------------------------------
|
|
13
|
+
# Databricks SDK fixtures
|
|
14
|
+
# ---------------------------------------------------------------------------
|
|
15
|
+
|
|
16
|
+
@pytest.fixture
|
|
17
|
+
def mock_workspace_client():
|
|
18
|
+
"""Mock WorkspaceClient with OAuth token access.
|
|
19
|
+
|
|
20
|
+
Simulates the SDK's authenticate() method which returns
|
|
21
|
+
fresh auth headers on each call.
|
|
22
|
+
"""
|
|
23
|
+
client = MagicMock()
|
|
24
|
+
client.config.host = "https://test-workspace.cloud.databricks.com"
|
|
25
|
+
client.config.authenticate.return_value = {
|
|
26
|
+
"Authorization": "Bearer test-oauth-token"
|
|
27
|
+
}
|
|
28
|
+
return client
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
# ---------------------------------------------------------------------------
|
|
32
|
+
# Stream fixtures (real anyio memory streams, not mocks)
|
|
33
|
+
# ---------------------------------------------------------------------------
|
|
34
|
+
|
|
35
|
+
@pytest.fixture
|
|
36
|
+
def memory_stream_pair():
|
|
37
|
+
"""Factory: create a (send, recv) memory object stream pair.
|
|
38
|
+
|
|
39
|
+
Usage:
|
|
40
|
+
send, recv = memory_stream_pair(buffer=16)
|
|
41
|
+
"""
|
|
42
|
+
def _make(buffer: int = 16):
|
|
43
|
+
return anyio.create_memory_object_stream(buffer)
|
|
44
|
+
return _make
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
@pytest.fixture
|
|
48
|
+
def stdio_streams(memory_stream_pair):
|
|
49
|
+
"""Simulated stdio transport streams.
|
|
50
|
+
|
|
51
|
+
Returns (test_send, proxy_read, proxy_write, test_recv):
|
|
52
|
+
- test_send: test writes here to simulate Claude Code input
|
|
53
|
+
- proxy_read: proxy reads from here (stdio read side)
|
|
54
|
+
- proxy_write: proxy writes here (stdio write side)
|
|
55
|
+
- test_recv: test reads here to verify proxy output to Claude Code
|
|
56
|
+
"""
|
|
57
|
+
test_send, proxy_read = memory_stream_pair()
|
|
58
|
+
proxy_write, test_recv = memory_stream_pair()
|
|
59
|
+
return test_send, proxy_read, proxy_write, test_recv
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
@pytest.fixture
|
|
63
|
+
def http_streams(memory_stream_pair):
|
|
64
|
+
"""Simulated HTTP transport streams.
|
|
65
|
+
|
|
66
|
+
Returns (test_send, proxy_read, proxy_write, test_recv):
|
|
67
|
+
- test_send: test writes here to simulate remote server responses
|
|
68
|
+
- proxy_read: proxy reads from here (HTTP read side)
|
|
69
|
+
- proxy_write: proxy writes here (HTTP write side)
|
|
70
|
+
- test_recv: test reads here to verify proxy sent to remote server
|
|
71
|
+
"""
|
|
72
|
+
test_send, proxy_read = memory_stream_pair()
|
|
73
|
+
proxy_write, test_recv = memory_stream_pair()
|
|
74
|
+
return test_send, proxy_read, proxy_write, test_recv
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
# ---------------------------------------------------------------------------
|
|
78
|
+
# Transport mock fixtures
|
|
79
|
+
# ---------------------------------------------------------------------------
|
|
80
|
+
|
|
81
|
+
@pytest.fixture
|
|
82
|
+
def mock_stdio_server(stdio_streams):
|
|
83
|
+
"""Async context manager replacing mcp.server.stdio.stdio_server.
|
|
84
|
+
|
|
85
|
+
Yields (read_stream, write_stream) from the proxy's perspective.
|
|
86
|
+
"""
|
|
87
|
+
_, proxy_read, proxy_write, _ = stdio_streams
|
|
88
|
+
|
|
89
|
+
@asynccontextmanager
|
|
90
|
+
async def _mock():
|
|
91
|
+
yield (proxy_read, proxy_write)
|
|
92
|
+
|
|
93
|
+
return _mock
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
@pytest.fixture
|
|
97
|
+
def mock_http_client(http_streams):
|
|
98
|
+
"""Async context manager replacing mcp.client.streamable_http.streamablehttp_client.
|
|
99
|
+
|
|
100
|
+
Yields (read_stream, write_stream, get_session_id) from the proxy's perspective.
|
|
101
|
+
Captures the auth kwarg for auth verification.
|
|
102
|
+
"""
|
|
103
|
+
_, proxy_read, proxy_write, _ = http_streams
|
|
104
|
+
captured = {}
|
|
105
|
+
|
|
106
|
+
@asynccontextmanager
|
|
107
|
+
async def _mock(url, *, auth=None, terminate_on_close=True, **kwargs):
|
|
108
|
+
captured["url"] = url
|
|
109
|
+
captured["auth"] = auth
|
|
110
|
+
captured["terminate_on_close"] = terminate_on_close
|
|
111
|
+
yield (proxy_read, proxy_write, lambda: "mock-session-id")
|
|
112
|
+
|
|
113
|
+
_mock.captured = captured
|
|
114
|
+
return _mock
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
# ---------------------------------------------------------------------------
|
|
118
|
+
# Auth fixtures
|
|
119
|
+
# ---------------------------------------------------------------------------
|
|
120
|
+
|
|
121
|
+
@pytest.fixture
|
|
122
|
+
def databricks_auth(mock_workspace_client):
|
|
123
|
+
"""Pre-configured DatabricksAuth instance."""
|
|
124
|
+
from uc_mcp_proxy import DatabricksAuth
|
|
125
|
+
return DatabricksAuth(mock_workspace_client)
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
# ---------------------------------------------------------------------------
|
|
129
|
+
# Sample message fixtures
|
|
130
|
+
# ---------------------------------------------------------------------------
|
|
131
|
+
|
|
132
|
+
@pytest.fixture
|
|
133
|
+
def sample_jsonrpc_request():
|
|
134
|
+
"""A sample JSON-RPC request dict (tools/list)."""
|
|
135
|
+
return {
|
|
136
|
+
"jsonrpc": "2.0",
|
|
137
|
+
"id": 1,
|
|
138
|
+
"method": "tools/list",
|
|
139
|
+
"params": {},
|
|
140
|
+
}
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
@pytest.fixture
|
|
144
|
+
def sample_jsonrpc_response():
|
|
145
|
+
"""A sample JSON-RPC response dict (tools/list result)."""
|
|
146
|
+
return {
|
|
147
|
+
"jsonrpc": "2.0",
|
|
148
|
+
"id": 1,
|
|
149
|
+
"result": {
|
|
150
|
+
"tools": [
|
|
151
|
+
{
|
|
152
|
+
"name": "chat_postmessage",
|
|
153
|
+
"description": "Sends a message to a Slack channel",
|
|
154
|
+
"inputSchema": {
|
|
155
|
+
"type": "object",
|
|
156
|
+
"properties": {
|
|
157
|
+
"channel": {"type": "string"},
|
|
158
|
+
"text": {"type": "string"},
|
|
159
|
+
},
|
|
160
|
+
"required": ["channel"],
|
|
161
|
+
},
|
|
162
|
+
}
|
|
163
|
+
]
|
|
164
|
+
},
|
|
165
|
+
}
|
|
File without changes
|