From af7f210f61485cb28dde51372b85d5ee50507f16 Mon Sep 17 00:00:00 2001 From: Ike Nwankwo Date: Fri, 6 Nov 2026 19:58:01 +0000 Subject: [PATCH] feat(models): forward MCP OAuth scopes in agy_cli generator --- evalbench/generators/models/agy_cli.py | 148 +++++++++++++++++- evalbench/test/agy_cli_test.py | 205 ++++++++++++++++++++++++- 2 files changed, 341 insertions(+), 12 deletions(-) diff --git a/evalbench/generators/models/agy_cli.py b/evalbench/generators/models/agy_cli.py index 38eb5900..036478c9 100644 --- a/evalbench/generators/models/agy_cli.py +++ b/evalbench/generators/models/agy_cli.py @@ -1,4 +1,5 @@ from .agent_cli import AgentCliGenerator +from . import mcp_client from .tool_naming import canonicalize_agy_tool_name, parse_agy_mcp_tool_call import subprocess import os @@ -117,6 +118,7 @@ def __init__(self, querygenerator_config): self.agy_cli_version = self._resolve_agy_version() self._setup_auth() + self._scoped_mcp_servers: dict[str, list[str]] = {} self.setup_config = querygenerator_config.get("setup", {}) if self.setup_config: self._setup_tools() @@ -804,26 +806,159 @@ def _setup_mcp_servers(self, mcp_servers_config: dict): logging.info("Removing stale MCP server configuration: %s", stale) del existing[stale] for server_name, config in mcp_servers_config.items(): - existing[server_name] = self._translate_mcp_config(dict(config)) + existing[server_name] = self._translate_mcp_config( + server_name, dict(config) + ) logging.info("Configured MCP server: %s", server_name) with open(self.mcp_config_path, "w") as f: json.dump(current_config, f, indent=2) - @staticmethod - def _translate_mcp_config(config: dict) -> dict: + def _translate_mcp_config(self, server_name: str, config: dict) -> dict: """Normalizes a cross-harness MCP server config into agy's schema. Maps the common gemini-style ``httpUrl`` alias to ``serverUrl``, agy's native field. ``serverUrl`` and ``url`` are accepted by agy - directly and need no translation. Other fields like - ``authProviderType``, ``oauth.scopes``, and stdio fields pass through - natively. + directly and need no translation. + + When ``authProviderType: google_credentials`` is paired with + ``oauth.scopes``, agy's built-in ADC provider ignores ``oauth.scopes`` + and hardcodes the ``cloud-platform`` scope. In that case, we mint a + scoped Bearer token via ``gcloud`` and inject it into + ``headers.Authorization`` (which overrides agy's built-in token header + while preserving ``x-goog-user-project`` and skipping OAuth discovery). """ if "httpUrl" in config and "serverUrl" not in config: config["serverUrl"] = config.pop("httpUrl") + + scopes = mcp_client.extract_mcp_oauth_scopes(config) + if scopes: + oauth_block = config.get("oauth") + if isinstance(oauth_block, dict): + oauth_block = dict(oauth_block) + oauth_block.pop("scopes", None) + if oauth_block: + config["oauth"] = oauth_block + else: + config.pop("oauth", None) + + if config.get("authProviderType") == "google_credentials": + headers = dict(config.get("headers") or {}) + if "Authorization" not in headers: + if not hasattr(self, "_scoped_mcp_servers"): + self._scoped_mcp_servers = {} + self._scoped_mcp_servers[server_name] = scopes + token = self._fetch_gcloud_access_token(scopes=scopes) + if token: + headers["Authorization"] = f"Bearer {token}" + config["headers"] = headers + else: + logging.warning( + "MCP server '%s' requires google_credentials with " + "scopes %s, but failed to fetch an access token " + "via gcloud.", + server_name, + scopes, + ) + return config + def _fetch_gcloud_access_token( + self, scopes: Optional[list[str]] = None + ) -> str: + """Fetches a Google Cloud access token for scoped MCP `google_credentials` auth.""" + token_env = os.environ.copy() + impersonate_sa = self.env.get( + "CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT" + ) or token_env.get("CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT") + if impersonate_sa: + token_env["CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT"] = ( + impersonate_sa + ) + adc = self.env.get("GOOGLE_APPLICATION_CREDENTIALS") + if impersonate_sa: + if os.path.exists("/etc/evalbench-sa-key/key.json"): + token_env["GOOGLE_APPLICATION_CREDENTIALS"] = ( + "/etc/evalbench-sa-key/key.json" + ) + else: + token_env.pop("GOOGLE_APPLICATION_CREDENTIALS", None) + elif adc and os.path.exists(adc): + token_env["GOOGLE_APPLICATION_CREDENTIALS"] = adc + + commands = [] + if scopes: + mcp_client.validate_oauth_scopes(scopes) + scopes_str = ",".join(scopes) + commands.append( + [ + "gcloud", + "auth", + "application-default", + "print-access-token", + f"--scopes={scopes_str}", + ] + ) + commands.extend( + [ + ["gcloud", "auth", "application-default", "print-access-token"], + ["gcloud", "auth", "print-access-token"], + ] + ) + for cmd in commands: + try: + result = subprocess.run( + cmd, + check=True, + capture_output=True, + text=True, + env=token_env, + ) + token = result.stdout.strip() + if token: + return token + except (subprocess.CalledProcessError, FileNotFoundError): + continue + return "" + + def _refresh_scoped_mcp_tokens(self) -> None: + """Refreshes Bearer tokens in mcp_config.json for scoped MCP servers.""" + scoped_servers = getattr(self, "_scoped_mcp_servers", None) + if not scoped_servers or not os.path.exists(self.mcp_config_path): + return + try: + with open(self.mcp_config_path, "r") as f: + cfg = json.load(f) + except (OSError, json.JSONDecodeError): + return + + mcp_servers = cfg.get("mcpServers") + if not isinstance(mcp_servers, dict): + return + + updated = False + for server_name, scopes in scoped_servers.items(): + server_cfg = mcp_servers.get(server_name) + if not isinstance(server_cfg, dict): + continue + token = self._fetch_gcloud_access_token(scopes=scopes) + if token: + headers = dict(server_cfg.get("headers") or {}) + headers["Authorization"] = f"Bearer {token}" + server_cfg["headers"] = headers + updated = True + + if updated: + try: + with open(self.mcp_config_path, "w") as f: + json.dump(cfg, f, indent=2) + except OSError as e: + logging.warning( + "Failed to refresh scoped MCP token in %s: %s", + self.mcp_config_path, + e, + ) + def _merged_env(self, extra: dict | None = None) -> dict: """Returns the process environment overlaid with the generator's configured env (and an optional per-call ``extra``).""" @@ -915,6 +1050,7 @@ def _execute_cli_command( ) def _run_agy_cli(self, cli_cmd: CLICommand, timeout_seconds=None): + self._refresh_scoped_mcp_tokens() env = self._merged_env(cli_cmd.env) # Registered via --add-dir below, fake_home fallback included: without # that, a scenario with no work_dir writes to /scratch diff --git a/evalbench/test/agy_cli_test.py b/evalbench/test/agy_cli_test.py index 7e9e427a..70c8decb 100644 --- a/evalbench/test/agy_cli_test.py +++ b/evalbench/test/agy_cli_test.py @@ -1287,7 +1287,9 @@ def test_translate_mcp_config_maps_httpurl_to_serverurl(): """The gemini-style ``httpUrl`` alias is rewritten to agy's ``serverUrl`` (left untranslated, agy ignores it and attaches a transportless server with zero tools).""" - out = AgyCliGenerator._translate_mcp_config({"httpUrl": "https://x/mcp"}) + generator = object.__new__(AgyCliGenerator) + generator.env = {} + out = generator._translate_mcp_config("test", {"httpUrl": "https://x/mcp"}) assert out == {"serverUrl": "https://x/mcp"} @@ -1295,8 +1297,11 @@ def test_translate_mcp_config_maps_httpurl_to_serverurl(): def test_translate_mcp_config_does_not_clobber_existing_serverurl(): """When ``serverUrl`` is already present, the canonical value wins and the ``httpUrl`` alias is not mapped over it.""" - out = AgyCliGenerator._translate_mcp_config( - {"httpUrl": "https://alias/mcp", "serverUrl": "https://canonical/mcp"} + generator = object.__new__(AgyCliGenerator) + generator.env = {} + out = generator._translate_mcp_config( + "test", + {"httpUrl": "https://alias/mcp", "serverUrl": "https://canonical/mcp"}, ) assert out["serverUrl"] == "https://canonical/mcp" @@ -1304,17 +1309,205 @@ def test_translate_mcp_config_does_not_clobber_existing_serverurl(): def test_translate_mcp_config_passes_native_fields_through_unchanged(): """stdio fields and other native agy schema keys pass through verbatim -- - no Bearer-token injection or field rewriting (unlike claude_code).""" + no Bearer-token injection or field rewriting when no custom oauth.scopes + are configured.""" + generator = object.__new__(AgyCliGenerator) + generator.env = {} cfg = { "command": "npx", "args": ["-y", "server"], "env": {"K": "V"}, "authProviderType": "google_credentials", - "oauth": {"scopes": ["a", "b"]}, + "oauth": {"clientId": "my-client"}, "headers": {"X": "Y"}, } - assert AgyCliGenerator._translate_mcp_config(dict(cfg)) == cfg + assert generator._translate_mcp_config("test", dict(cfg)) == cfg + + +def test_translate_mcp_config_forwards_oauth_scopes(): + """Because agy ignores oauth.scopes and hardcodes the cloud-platform scope + for google_credentials, _translate_mcp_config mints a scoped Bearer token + via gcloud and injects it into headers.Authorization.""" + generator = object.__new__(AgyCliGenerator) + generator.env = {} + generator._scoped_mcp_servers = {} + + mcp_config = { + "httpUrl": "https://colaboratory.googleapis.com/mcp", + "authProviderType": "google_credentials", + "oauth": { + "scopes": [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/colaboratory", + ] + }, + "headers": { + "X-Goog-User-Project": "my-project" + }, + } + + with patch.object( + generator, "_fetch_gcloud_access_token", return_value="fake_scoped_token" + ) as mock_fetch: + translated = generator._translate_mcp_config("colab-mcp", mcp_config) + + assert "oauth" not in translated + assert translated["authProviderType"] == "google_credentials" + assert translated["serverUrl"] == "https://colaboratory.googleapis.com/mcp" + assert translated["headers"]["Authorization"] == "Bearer fake_scoped_token" + assert translated["headers"]["X-Goog-User-Project"] == "my-project" + assert generator._scoped_mcp_servers["colab-mcp"] == [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/colaboratory", + ] + + mock_fetch.assert_called_once_with( + scopes=[ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/colaboratory", + ] + ) + + +def test_fetch_gcloud_access_token_passes_scopes(): + generator = object.__new__(AgyCliGenerator) + generator.env = {} + + scopes = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/colaboratory", + ] + + with patch("generators.models.agy_cli.subprocess.run") as mock_run: + mock_proc = MagicMock() + mock_proc.stdout = "scoped_token_xyz\n" + mock_run.return_value = mock_proc + + token = generator._fetch_gcloud_access_token(scopes=scopes) + + assert token == "scoped_token_xyz" + first_call_cmd = mock_run.call_args_list[0][0][0] + assert first_call_cmd == [ + "gcloud", + "auth", + "application-default", + "print-access-token", + "--scopes=https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/colaboratory", + ] + + +def test_fetch_gcloud_access_token_uses_gke_key_when_impersonating(monkeypatch): + generator = object.__new__(AgyCliGenerator) + generator.env = { + "GOOGLE_APPLICATION_CREDENTIALS": "/tmp_sessions/impersonated_adc.json", + "CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT": ( + "evalbench-runner@onemcp-evalviewer-dev.iam.gserviceaccount.com" + ), + } + + monkeypatch.setattr( + "generators.models.agy_cli.os.path.exists", + lambda p: p in ( + "/tmp_sessions/impersonated_adc.json", + "/etc/evalbench-sa-key/key.json", + ), + ) + + with patch("generators.models.agy_cli.subprocess.run") as mock_run: + mock_proc = MagicMock() + mock_proc.stdout = "impersonated_scoped_token\n" + mock_run.return_value = mock_proc + + token = generator._fetch_gcloud_access_token( + scopes=["https://www.googleapis.com/auth/colaboratory"] + ) + + assert token == "impersonated_scoped_token" + called_env = mock_run.call_args_list[0].kwargs["env"] + assert ( + called_env["CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT"] + == "evalbench-runner@onemcp-evalviewer-dev.iam.gserviceaccount.com" + ) + assert ( + called_env["GOOGLE_APPLICATION_CREDENTIALS"] + == "/etc/evalbench-sa-key/key.json" + ) + + +def test_fetch_gcloud_access_token_unsets_adc_when_impersonating_on_workload_identity( + monkeypatch, +): + generator = object.__new__(AgyCliGenerator) + generator.env = { + "GOOGLE_APPLICATION_CREDENTIALS": "/tmp_sessions/impersonated_adc.json", + "CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT": ( + "evalbench-runner@onemcp-evalviewer-dev.iam.gserviceaccount.com" + ), + } + + monkeypatch.setattr( + "generators.models.agy_cli.os.path.exists", + lambda p: p == "/tmp_sessions/impersonated_adc.json", + ) + + with patch("generators.models.agy_cli.subprocess.run") as mock_run: + mock_proc = MagicMock() + mock_proc.stdout = "wi_impersonated_scoped_token\n" + mock_run.return_value = mock_proc + + token = generator._fetch_gcloud_access_token( + scopes=["https://www.googleapis.com/auth/colaboratory"] + ) + + assert token == "wi_impersonated_scoped_token" + called_env = mock_run.call_args_list[0].kwargs["env"] + assert ( + called_env["CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT"] + == "evalbench-runner@onemcp-evalviewer-dev.iam.gserviceaccount.com" + ) + assert "GOOGLE_APPLICATION_CREDENTIALS" not in called_env + + +def test_run_agy_cli_refreshes_scoped_mcp_token(mock_run, sandbox): + generator = AgyCliGenerator({}) + generator._scoped_mcp_servers = { + "colab-mcp": ["https://www.googleapis.com/auth/colaboratory"] + } + initial_cfg = { + "mcpServers": { + "colab-mcp": { + "serverUrl": "https://colaboratory.googleapis.com/mcp", + "authProviderType": "google_credentials", + "headers": { + "Authorization": "Bearer old_token", + "X-Goog-User-Project": "my-project", + }, + } + } + } + with open(generator.mcp_config_path, "w") as f: + json.dump(initial_cfg, f) + + with patch.object( + generator, "_fetch_gcloud_access_token", return_value="refreshed_token" + ) as mock_fetch: + generator._run_agy_cli(CLICommand(generator.agy_bin, "hi")) + + mock_fetch.assert_called_once_with( + scopes=["https://www.googleapis.com/auth/colaboratory"] + ) + + with open(generator.mcp_config_path, "r") as f: + updated_cfg = json.load(f) + assert ( + updated_cfg["mcpServers"]["colab-mcp"]["headers"]["Authorization"] + == "Bearer refreshed_token" + ) + assert ( + updated_cfg["mcpServers"]["colab-mcp"]["headers"]["X-Goog-User-Project"] + == "my-project" + ) def test_run_passes_configured_model_flag(mock_run, sandbox):