Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,23 @@ permissions:
jobs:
test:
runs-on: ubuntu-latest
env:
# uv.lock pins packages to Databricks' internal pypi-proxy, which hosted
# GitHub runners can't reach (downloads time out). Re-resolve against
# public PyPI at CI time: UV_INDEX_URL + a `uv lock` before install.
UV_INDEX_URL: https://pypi.org/simple
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
- run: uv lock
- run: uv run pytest --ignore=tests/test_e2e.py --ignore=tests/test_e2e_tracing.py

e2e:
runs-on: ubuntu-latest
env:
# See the test job: hosted runners can't reach the internal pypi-proxy,
# so resolve against public PyPI (paired with the `uv lock` step below).
UV_INDEX_URL: https://pypi.org/simple
UCODE_TEST_WORKSPACE: ${{ secrets.UCODE_TEST_WORKSPACE }}
DATABRICKS_HOST: ${{ secrets.UCODE_TEST_WORKSPACE }}
# DATABRICKS_BEARER is the CI escape hatch: `databricks auth token`
Expand Down Expand Up @@ -47,6 +56,7 @@ jobs:
opencode-ai
@github/copilot
@earendil-works/pi-coding-agent
- run: uv lock
- run: uv tool install .
# Redirect stdin so any interactive `databricks auth login --no-browser`
# fallback EOFs instead of hanging the runner. With DATABRICKS_BEARER
Expand Down
9 changes: 7 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -89,8 +89,13 @@ Options are shown in this order:
- Managed Databricks MCPs (Vector Search, UC Functions, etc.)
- Custom MCP server URL

Discovered external MCP connections are listed directly. MCP auth uses a Databricks token that
`ucode` sets when launching each tool.
Discovered external MCP connections are listed directly.

Every Databricks MCP server is registered as a local **stdio** server that runs `ucode mcp-proxy`
— a small bridge (shipped with `ucode`) between the coding tool and the Databricks
streamable-HTTP MCP endpoint. The proxy mints a fresh OAuth token from your Databricks CLI profile
on every request, so MCP auth is handled uniformly for every client and never expires mid-session.
The coding tool starts and stops the proxy as a child process; there's nothing extra to run.

To set up an agent and its MCP server(s) in one command, pass `--mcp` with fully-qualified
service name(s) to `ucode configure`:
Expand Down
4 changes: 4 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@ readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"databricks-sql-connector>=3.6.0",
# `ucode mcp-proxy` bridges a client's stdio MCP transport to a Databricks
# streamable-HTTP MCP endpoint, injecting a freshly-minted OAuth bearer per
# request. Uses the official MCP SDK's stdio server + streamable-HTTP client.
"mcp>=1.28.0",
"questionary>=2.0.0",
"tomlkit>=0.13.0",
"typer>=0.12.0",
Expand Down
18 changes: 10 additions & 8 deletions src/ucode/agents/copilot.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,25 +99,27 @@ def build_runtime_env(workspace: str, model: str, token: str) -> dict[str, str]:
return env


def build_mcp_server_entry(url: str) -> dict:
def build_mcp_server_entry(argv: list[str]) -> dict:
# A `local` MCP server runs a stdio command; `command`/`args` split the
# argv. ucode registers the `ucode mcp-proxy ...` bridge here so Copilot
# never speaks HTTP+bearer directly — the proxy handles token refresh. The
# OAUTH_TOKEN env Copilot still injects at launch is for MODEL auth, not MCP.
return {
"type": "http",
"url": url,
"headers": {
"Authorization": "Bearer ${OAUTH_TOKEN}",
},
"type": "local",
"command": argv[0],
"args": list(argv[1:]),
"tools": ["*"],
}


def write_mcp_server_config(name: str, url: str) -> bool:
def write_mcp_server_config(name: str, argv: list[str]) -> bool:
backup_existing_file(COPILOT_MCP_CONFIG_PATH, COPILOT_MCP_BACKUP_PATH)
existing = read_json_safe(COPILOT_MCP_CONFIG_PATH)
mcp_servers = existing.get("mcpServers")
if not isinstance(mcp_servers, dict):
mcp_servers = {}
removed = name in mcp_servers
mcp_servers[name] = build_mcp_server_entry(url)
mcp_servers[name] = build_mcp_server_entry(argv)
existing["mcpServers"] = mcp_servers
write_json_file(COPILOT_MCP_CONFIG_PATH, existing)
return removed
Expand Down
17 changes: 8 additions & 9 deletions src/ucode/agents/opencode.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
OPENCODE_CONFIG_DIR = OPENCODE_XDG_CONFIG_HOME / "opencode"
OPENCODE_CONFIG_PATH = OPENCODE_CONFIG_DIR / "opencode.json"
OPENCODE_BACKUP_PATH = APP_DIR / "opencode-config.backup.json"
OPENCODE_MCP_AUTH_HEADER_VALUE = "Bearer {env:OAUTH_TOKEN}"

SPEC: ToolSpec = {
"binary": "opencode",
Expand Down Expand Up @@ -193,25 +192,25 @@ def write_tool_config(
return state, token


def build_mcp_server_entry(url: str) -> dict:
def build_mcp_server_entry(argv: list[str]) -> dict:
# A `local` MCP server runs a command over stdio; `command` is the full
# argv. ucode registers the `ucode mcp-proxy ...` bridge here so OpenCode
# never speaks HTTP+bearer directly — the proxy mints fresh tokens itself.
return {
"type": "remote",
"url": url,
"type": "local",
"command": list(argv),
"enabled": True,
"headers": {
"Authorization": OPENCODE_MCP_AUTH_HEADER_VALUE,
},
}


def write_mcp_server_config(name: str, url: str) -> bool:
def write_mcp_server_config(name: str, argv: list[str]) -> bool:
backup_existing_file(OPENCODE_CONFIG_PATH, OPENCODE_BACKUP_PATH)
existing = read_json_safe(OPENCODE_CONFIG_PATH)
mcp_servers = existing.get("mcp")
if not isinstance(mcp_servers, dict):
mcp_servers = {}
removed = name in mcp_servers
mcp_servers[name] = build_mcp_server_entry(url)
mcp_servers[name] = build_mcp_server_entry(argv)
existing["mcp"] = mcp_servers
write_json_file(OPENCODE_CONFIG_PATH, existing)
return removed
Expand Down
37 changes: 37 additions & 0 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -868,6 +868,43 @@ def mcp_web_search_cmd() -> None:
serve()


@app.command("mcp-proxy", hidden=True)
def mcp_proxy_cmd(
url: Annotated[
str,
typer.Option("--url", help="Databricks streamable-HTTP MCP endpoint to forward to."),
],
host: Annotated[
str | None,
typer.Option(
"--host", help="Workspace URL for token minting. Defaults to the saved workspace."
),
] = None,
profile: Annotated[
str | None, typer.Option("--profile", help="Databricks CLI profile.")
] = None,
use_pat: Annotated[
bool, typer.Option("--use-pat", help="Use the profile's static PAT instead of OAuth.")
] = False,
) -> None:
"""Bridge a coding agent's stdio MCP transport to a Databricks MCP endpoint.

Each configured client spawns this as a local stdio MCP server (see
`ucode configure mcp`); it forwards messages to ``--url`` and injects a
freshly-minted OAuth bearer on every upstream request, so the token never
expires mid-session. Not meant for interactive use — the agent manages this
process's lifecycle."""
from ucode.mcp_proxy import serve

state = load_state()
workspace = host or state.get("workspace")
if not workspace:
print_err("No workspace configured. Run `ucode configure` first.")
raise typer.Exit(1)
profile = profile or state.get("profile")
serve(url, workspace, profile, use_pat=use_pat or bool(state.get("use_pat")))


@app.command("auth-token", hidden=True)
def auth_token_cmd(
host: Annotated[
Expand Down
21 changes: 21 additions & 0 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1138,6 +1138,27 @@ def build_auth_token_argv(
return argv


def build_mcp_proxy_argv(
url: str, workspace: str, profile: str | None = None, *, use_pat: bool = False
) -> list[str]:
"""Argv for the stdio MCP bridge: `ucode mcp-proxy --url ... --host ...`.

Every coding agent registers this single command as a local stdio MCP
server instead of a per-client HTTP endpoint with a bearer header. The proxy
forwards to ``url`` and mints a fresh OAuth token on each upstream request,
so tokens never expire mid-session — the client only ever spawns a process,
which keeps registration uniform across CLIs that disagree on HTTP-auth
syntax. Like `build_auth_token_argv`, this resolves the absolute `ucode`
path and passes plain arguments (no shell), so it runs identically on every
platform."""
argv = [_ucode_binary(), "mcp-proxy", "--url", url, "--host", workspace.rstrip("/")]
if profile:
argv += ["--profile", profile]
if use_pat:
argv.append("--use-pat")
return argv


def build_auth_shell_command(
workspace: str, profile: str | None = None, *, use_pat: bool = False
) -> str:
Expand Down
Loading
Loading