From fc0f4f0c268b746ffbd4df36dc10563cb82f050d Mon Sep 17 00:00:00 2001 From: Tien Le Date: Tue, 6 Oct 2026 14:13:08 +0000 Subject: [PATCH] Forward external Databricks auth to Codex MCP proxies --- src/ucode/agents/codex.py | 6 +++++- tests/test_agent_codex.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 75606811d..f6019eb7f 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -627,7 +627,11 @@ def managed_mcp_uses_managed_file() -> bool: def managed_mcp_entry(argv: list[str]) -> dict: """A ``[mcp_servers.]`` stdio entry from the ``ug mcp-proxy`` argv (same as user scope).""" - return {"command": argv[0], "args": list(argv[1:])} + return { + "command": argv[0], + "args": list(argv[1:]), + "env_vars": ["DATABRICKS_BEARER", "DATABRICKS_BEARER_COMMAND"], + } def managed_mcp_http_entry(url: str, client_id: str) -> dict: diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index c810a1b6f..a68913e2e 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -1929,6 +1929,10 @@ def test_writes_mcp_servers_preserving_model_keys(self, monkeypatch): assert doc["model_provider"] == "Databricks" assert doc["mcp_servers"]["system-ai-github"]["command"] == "/opt/ug" assert doc["mcp_servers"]["system-ai-github"]["args"][:2] == ["mcp-proxy", "--url"] + assert doc["mcp_servers"]["system-ai-github"]["env_vars"] == [ + "DATABRICKS_BEARER", + "DATABRICKS_BEARER_COMMAND", + ] assert captured["owned_paths"] == [["mcp_servers"]] assert captured["tool"] == "codex" @@ -2041,6 +2045,16 @@ def test_token_provider_uses_custom_oauth_when_configured(self, monkeypatch): class TestWriteUserMcpServers: """Batched user-scope `[mcp_servers]` writes for the workspace-managed reconcile path.""" + def test_passes_external_auth_by_name_without_persisting_credentials(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "test-bearer-not-for-persistence") + monkeypatch.setenv("DATABRICKS_BEARER_COMMAND", "test-token-command") + + assert codex.managed_mcp_entry(["ug", "mcp-proxy", "https://ws/svc"]) == { + "command": "ug", + "args": ["mcp-proxy", "https://ws/svc"], + "env_vars": ["DATABRICKS_BEARER", "DATABRICKS_BEARER_COMMAND"], + } + def test_adds_and_preserves_other_tables(self, tmp_path, monkeypatch): path = tmp_path / "config.toml" path.write_text('model = "gpt-5"\n\n[mcp_servers.mine]\ncommand = "x"\nargs = []\n') @@ -2058,6 +2072,7 @@ def test_adds_and_preserves_other_tables(self, tmp_path, monkeypatch): assert dict(doc["mcp_servers"]["system-ai-github"]) == { "command": "ug", "args": ["mcp-proxy", "https://ws/svc"], + "env_vars": ["DATABRICKS_BEARER", "DATABRICKS_BEARER_COMMAND"], } def test_removes_named_entries_only(self, tmp_path, monkeypatch):