Add oauth support for JSON-string passed configs

Author: timesler-atlassian

mcp_compressor/main.py

@@ -25,11 +25,12 @@
 import keyring.errors
 import psutil
 import typer
+from click.core import ParameterSource
 from cryptography.fernet import Fernet
 from fastmcp import FastMCP
 from fastmcp.client.auth import OAuth
-from fastmcp.client.transports import MCPConfigTransport, SSETransport, StdioTransport, StreamableHttpTransport
-from fastmcp.mcp_config import MCPConfig
+from fastmcp.client.transports import SSETransport, StdioTransport, StreamableHttpTransport
+from fastmcp.mcp_config import MCPConfig, RemoteMCPServer, StdioMCPServer
 from fastmcp.server import create_proxy
 from fastmcp.server.providers.proxy import ProxyClient
 from key_value.aio.protocols import AsyncKeyValue
@@ -70,6 +71,7 @@ def _version_callback(value: bool) -> None:
 
 @app.command()
 def main(
+    ctx: typer.Context,
     command_or_url_list: Annotated[
         list[str],
         typer.Argument(
@@ -223,6 +225,23 @@ def main(
         raise typer.BadParameter(str(exc), param_hint="'COMMAND_OR_URL'") from exc
 
     if parsed_config is not None:
+        conflicting_config_options = [
+            option_name
+            for option_name in ("cwd", "env_list", "header_list", "timeout")
+            if ctx.get_parameter_source(option_name) == ParameterSource.COMMANDLINE
+        ]
+        if conflicting_config_options:
+            joined_options = ", ".join(
+                f"--{name.removesuffix('_list').replace('_', '-')}" for name in conflicting_config_options
+            )
+            raise typer.BadParameter(
+                (
+                    f"JSON MCP config input cannot be combined with {joined_options}; configure those values inside "
+                    "the JSON instead."
+                ),
+                param_hint="'COMMAND_OR_URL'",
+            )
+
         _config, config_server_name = parsed_config
         resolved_server_name = server_name or config_server_name
 
@@ -423,17 +442,13 @@ async def _server(
     parsed_config = _parse_single_server_mcp_config(command_or_url_list)
     if parsed_config is not None:
         config, _config_server_name = parsed_config
-        transport_type: Literal["stdio", "http", "sse"] = _infer_transport_type(
-            str(next(iter(config.mcpServers.values())).model_dump().get("url") or command_or_url)
-        )
-        transport = MCPConfigTransport(config)
+        transport, transport_type = _get_single_server_transport_from_mcp_config(config=config)
         logger.info("Loaded single-server MCP config JSON")
     else:
         transport_type = _infer_transport_type(command_or_url)
         logger.info(f"Inferred transport type: {transport_type}")
 
         # Handle different transport types
-        transport: TransportType
         if transport_type == "stdio":
             transport = _get_stdio_transport(
                 command=command_or_url_list[0], args=command_or_url_list[1:], cwd=cwd, env_list=env_list
@@ -657,6 +672,57 @@ def _parse_single_server_mcp_config(command_or_url_list: list[str]) -> tuple[MCP
     return config, server_name
 
 
+def _get_single_server_transport_from_mcp_config(
+    config: MCPConfig,
+) -> tuple[TransportType, Literal["stdio", "http", "sse"]]:
+    """Create a transport for a validated single-server MCP config.
+
+    Single-server config JSON is self-contained: transport settings come only from the config.
+    Remote configs default to the same OAuth flow used by direct URL inputs unless explicit auth is provided.
+    """
+    server_config = next(iter(config.mcpServers.values()))
+
+    if isinstance(server_config, StdioMCPServer):
+        return (
+            _get_stdio_transport(
+                command=server_config.command,
+                args=server_config.args,
+                cwd=server_config.cwd,
+                env_list=[f"{key}={value}" for key, value in server_config.env.items()] or None,
+            ),
+            "stdio",
+        )
+
+    if isinstance(server_config, RemoteMCPServer):
+        transport_type = "sse" if (server_config.transport == "sse") else _infer_transport_type(server_config.url)
+        auth = server_config.auth
+        if auth in (None, "oauth"):
+            auth = OAuth(mcp_url=server_config.url, token_storage=_build_token_storage())
+
+        headers = {key: _interpolate_string(value) for key, value in server_config.headers.items()}
+        if transport_type == "http":
+            return (
+                StreamableHttpTransport(
+                    url=server_config.url,
+                    headers=headers,
+                    auth=auth,
+                    sse_read_timeout=server_config.sse_read_timeout,
+                ),
+                "http",
+            )
+        return (
+            SSETransport(
+                url=server_config.url,
+                headers=headers,
+                auth=auth,
+                sse_read_timeout=server_config.sse_read_timeout,
+            ),
+            "sse",
+        )
+
+    raise ValueError("Unsupported single-server MCP config type.")
+
+
 def _interpolate_string(value: str) -> str:
     """Interpolate environment variables in a single string.
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mcp-compressor"
-version = "0.2.9"
+version = "0.2.10"
 description = "An MCP server wrapper for reducing tokens consumed by MCP tools."
 authors = [{ name = "Tim Esler", email = "tesler@atlassian.com" }]
 readme = "README.md"
@@ -32,6 +32,7 @@ dependencies = [
     "toons>=0.5.3",
     "typer>=0.16.0",
     "uvicorn>=0.30.0",
+    "click>=8.3.1",
 ]
 
 [project.urls]

tests/test_main.py

@@ -7,8 +7,9 @@
 from typing import Any, cast
 
 import pytest
+from fastmcp.client.auth.bearer import BearerAuth
 from fastmcp.client.auth.oauth import ClientNotFoundError, OAuth
-from fastmcp.client.transports import MCPConfigTransport, SSETransport, StdioTransport, StreamableHttpTransport
+from fastmcp.client.transports import SSETransport, StdioTransport, StreamableHttpTransport
 from fastmcp.exceptions import McpError
 from typer.testing import CliRunner
 
@@ -19,6 +20,7 @@
     suppress_recoverable_oauth_traceback_logging,
 )
 from mcp_compressor.main import (
+    _get_single_server_transport_from_mcp_config,
     _get_sse_transport,
     _get_stdio_transport,
     _get_streamable_http_transport,
@@ -120,6 +122,51 @@ def test_parse_single_server_mcp_config_rejects_multiple_servers() -> None:
 # Tests for transport creation functions
 
 
+def test_get_single_server_transport_from_mcp_config_remote_defaults_to_oauth(monkeypatch: pytest.MonkeyPatch) -> None:
+    config_json = '{"mcpServers": {"weather": {"url": "https://example.com/mcp"}}}'
+    parsed = _parse_single_server_mcp_config([config_json])
+    assert parsed is not None
+    config, _ = parsed
+
+    token_storage = object()
+    monkeypatch.setattr(main_module, "_build_token_storage", lambda: token_storage)
+
+    transport, transport_type = _get_single_server_transport_from_mcp_config(config=config)
+
+    assert transport_type == "http"
+    assert isinstance(transport, StreamableHttpTransport)
+    assert isinstance(transport.auth, OAuth)
+    assert transport.auth.mcp_url == "https://example.com/mcp"
+    assert transport.auth._token_storage is token_storage
+
+
+def test_get_single_server_transport_from_mcp_config_remote_preserves_explicit_auth() -> None:
+    config_json = '{"mcpServers": {"weather": {"url": "https://example.com/mcp", "auth": "abc"}}}'
+    parsed = _parse_single_server_mcp_config([config_json])
+    assert parsed is not None
+    config, _ = parsed
+
+    transport, transport_type = _get_single_server_transport_from_mcp_config(config=config)
+
+    assert transport_type == "http"
+    assert isinstance(transport, StreamableHttpTransport)
+    assert isinstance(transport.auth, BearerAuth)
+    assert transport.auth.token.get_secret_value() == "abc"
+
+
+def test_get_single_server_transport_from_mcp_config_sse_uses_only_config_timeout() -> None:
+    config_json = '{"mcpServers": {"weather": {"url": "https://example.com/sse", "transport": "sse"}}}'
+    parsed = _parse_single_server_mcp_config([config_json])
+    assert parsed is not None
+    config, _ = parsed
+
+    transport, transport_type = _get_single_server_transport_from_mcp_config(config=config)
+
+    assert transport_type == "sse"
+    assert isinstance(transport, SSETransport)
+    assert transport.sse_read_timeout is None
+
+
 def test_get_stdio_transport(tmp_path) -> None:
     """Test that stdio transport is created with correct parameters."""
     transport = _get_stdio_transport(
@@ -296,7 +343,34 @@ async def fake_async_main(**kwargs: Any) -> None:
     assert async_main_called is False
 
 
-async def test_server_uses_mcp_config_transport_for_single_server_json(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.parametrize(
+    ("extra_args", "expected_option"),
+    [
+        (["--cwd", "."], "--cwd"),
+        (["--env", "FOO=bar"], "--env"),
+        (["--header", "Authorization=Bearer abc"], "--header"),
+        (["--timeout", "30"], "--timeout"),
+    ],
+)
+def test_single_server_mcp_config_rejects_conflicting_transport_options(
+    runner: CliRunner, monkeypatch: pytest.MonkeyPatch, extra_args: list[str], expected_option: str
+) -> None:
+    config_json = '{"mcpServers": {"weather": {"url": "https://example.com/mcp"}}}'
+    async_main_called = False
+
+    async def fake_async_main(**kwargs: Any) -> None:
+        nonlocal async_main_called
+        async_main_called = True
+
+    monkeypatch.setattr(main_module, "_async_main", fake_async_main)
+    result = runner.invoke(app, [*extra_args, config_json])
+
+    assert result.exit_code != 0
+    assert expected_option in _strip_ansi(result.output)
+    assert async_main_called is False
+
+
+async def test_server_uses_single_server_config_transport_directly(monkeypatch: pytest.MonkeyPatch) -> None:
     config_json = '{"mcpServers": {"weather": {"command": "uvx", "args": ["mcp-weather"]}}}'
     captured: dict[str, Any] = {}
 
@@ -334,7 +408,9 @@ async def get_compression_stats(self) -> dict[str, int]:
     ) as mcp:
         assert mcp is fake_mcp
 
-    assert isinstance(captured["transport"], MCPConfigTransport)
+    assert isinstance(captured["transport"], StdioTransport)
+    assert captured["transport"].command == "uvx"
+    assert captured["transport"].args == ["mcp-weather"]
     assert captured["compressed_tools_kwargs"]["server_name"] == "weather"
 
 
@@ -576,9 +652,6 @@ def test_version_flag(runner: CliRunner) -> None:
 
 def test_version_short_flag(runner: CliRunner) -> None:
     """-V should be an alias for --version."""
-
-    from mcp_compressor.main import app
-
     result = runner.invoke(app, ["-V"])
     assert result.exit_code == 0
     expected_version = importlib.metadata.version("mcp-compressor")