From b455cddbd990f60d37ce33bea6fe282b9f4f2184 Mon Sep 17 00:00:00 2001 From: Ike Nwankwo Date: Fri, 6 Nov 2026 19:53:56 +0000 Subject: [PATCH] fix(auth): support impersonated ADC credentials in ClaudeCodeGenerator and CodexCliGenerator --- evalbench/generators/models/claude_code.py | 42 ++++++- evalbench/generators/models/codex_cli.py | 16 ++- evalbench/test/claude_code_test.py | 135 +++++++++++++++++++++ evalbench/test/codex_cli_test.py | 73 +++++++++++ 4 files changed, 259 insertions(+), 7 deletions(-) diff --git a/evalbench/generators/models/claude_code.py b/evalbench/generators/models/claude_code.py index 3456ad13..f59242d1 100644 --- a/evalbench/generators/models/claude_code.py +++ b/evalbench/generators/models/claude_code.py @@ -143,8 +143,8 @@ def __init__(self, querygenerator_config): if os.path.exists(fake_adc_path) and "GOOGLE_APPLICATION_CREDENTIALS" not in self.env: self.env["GOOGLE_APPLICATION_CREDENTIALS"] = fake_adc_path - if "CLOUDSDK_CONFIG" not in self.env: - self.env["CLOUDSDK_CONFIG"] = fake_gcloud_dir + if "CLOUDSDK_CONFIG" not in self.env: + self.env["CLOUDSDK_CONFIG"] = fake_gcloud_dir # Copy Claude Code auth credentials from real home to fake home so the # CLI can authenticate in the sandboxed environment. Restricted to a @@ -312,8 +312,23 @@ def _translate_mcp_config(self, server_name: str, config: dict) -> dict: # helper can't produce a token (gcloud/ADC unavailable in Claude # Code's env), it prints nothing and Claude Code falls back to the # static header below. + impersonate_sa = self.env.get( + "CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT" + ) + adc_env_prefix = "" + if impersonate_sa: + if os.path.exists("/etc/evalbench-sa-key/key.json"): + adc_env_prefix = ( + 'GOOGLE_APPLICATION_CREDENTIALS="/etc/evalbench-sa-key/key.json" ' + ) + else: + adc_env_prefix = "env -u GOOGLE_APPLICATION_CREDENTIALS " config.setdefault( - "headersHelper", self._google_credentials_headers_helper(scopes=scopes)) + "headersHelper", + self._google_credentials_headers_helper( + scopes=scopes, adc_env_prefix=adc_env_prefix + ), + ) # Baked static token: the initial value and the fallback for when # headersHelper can't run. headersHelper output takes precedence. @@ -333,6 +348,7 @@ def _translate_mcp_config(self, server_name: str, config: dict) -> dict: @staticmethod def _google_credentials_headers_helper( scopes: Optional[list[str]] = None, + adc_env_prefix: str = "", ) -> str: """Shell command Claude Code executes on each MCP connection to mint a fresh Google bearer token. @@ -349,10 +365,10 @@ def _google_credentials_headers_helper( mcp_client.validate_oauth_scopes(scopes) scopes_str = ",".join(scopes) scoped_cmd = ( - f'gcloud auth application-default print-access-token --scopes="{scopes_str}" 2>/dev/null || ' + f'{adc_env_prefix}gcloud auth application-default print-access-token --scopes="{scopes_str}" 2>/dev/null || ' ) return ( - f'tok="$({scoped_cmd}gcloud auth application-default print-access-token ' + f'tok="$({scoped_cmd}{adc_env_prefix}gcloud auth application-default print-access-token ' '2>/dev/null || gcloud auth print-access-token 2>/dev/null)"; ' '[ -n "$tok" ] && printf \'{"Authorization":"Bearer %s"}\' "$tok"' ) @@ -387,8 +403,22 @@ def _fetch_gcloud_access_token( fake-home config and both commands fail). """ token_env = os.environ.copy() + impersonate_sa = self.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 adc and os.path.exists(adc): + 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 # ADC first (works for these endpoints); user token as a fallback for # any MCP server that happens to accept it. diff --git a/evalbench/generators/models/codex_cli.py b/evalbench/generators/models/codex_cli.py index 5cb0af76..3d88660a 100644 --- a/evalbench/generators/models/codex_cli.py +++ b/evalbench/generators/models/codex_cli.py @@ -663,8 +663,22 @@ def _fetch_gcloud_access_token( if one was configured. """ token_env = os.environ.copy() + impersonate_sa = self.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 adc and os.path.exists(adc): + 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 if scopes is None: diff --git a/evalbench/test/claude_code_test.py b/evalbench/test/claude_code_test.py index a959f081..d84a27a0 100644 --- a/evalbench/test/claude_code_test.py +++ b/evalbench/test/claude_code_test.py @@ -363,3 +363,138 @@ def test_local_run_copies_gcloud_credentials_into_the_sandbox( assert os.path.exists(os.path.join(fake_gcloud, "access_tokens.db")) assert generator.env["GOOGLE_APPLICATION_CREDENTIALS"] == os.path.join( fake_gcloud, "application_default_credentials.json") + + +def test_explicit_adc_does_not_set_cloudsdk_config(tmp_path, monkeypatch): + """When an explicit GOOGLE_APPLICATION_CREDENTIALS file (such as + /tmp_sessions/impersonated_adc.json) is copied into fake_home, CLOUDSDK_CONFIG + must not point at fake_gcloud_dir so gcloud does not load impersonated_adc.json + as its default ADC.""" + impersonated = tmp_path / "impersonated_adc.json" + impersonated.write_text('{"type": "impersonated_service_account"}') + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr( + "generators.models.claude_code.SESSIONS_ROOT", str(tmp_path / "sessions") + ) + monkeypatch.setattr(sys, "argv", ["/evalbench/evalbench/eval_server.py"]) + + generator = ClaudeCodeGenerator( + { + "model": "claude-opus-4-6", + "use_vertex": True, + "session_id": "s1", + "env": { + "GOOGLE_APPLICATION_CREDENTIALS": str(impersonated), + "CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT": ( + "evalbench-runner@onemcp-evalviewer-dev.iam.gserviceaccount.com" + ), + }, + } + ) + + fake_adc = os.path.join( + generator.fake_home, + ".config", + "gcloud", + "application_default_credentials.json", + ) + assert generator.env["GOOGLE_APPLICATION_CREDENTIALS"] == fake_adc + assert os.path.exists(fake_adc) + assert "CLOUDSDK_CONFIG" not in generator.env + + +def test_impersonation_with_gke_key_uses_sa_key_for_gcloud(monkeypatch): + generator = object.__new__(ClaudeCodeGenerator) + 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.claude_code.os.path.exists", + lambda p: p in ( + "/tmp_sessions/impersonated_adc.json", + "/etc/evalbench-sa-key/key.json", + ), + ) + + with patch("generators.models.claude_code.subprocess.run") as mock_run: + mock_proc = MagicMock() + mock_proc.stdout = "impersonated_token_abc\n" + mock_run.return_value = mock_proc + + mcp_config = { + "httpUrl": "https://colaboratory.googleapis.com/mcp", + "authProviderType": "google_credentials", + "oauth": { + "scopes": ["https://www.googleapis.com/auth/colaboratory"] + }, + } + translated = generator._translate_mcp_config("colab-mcp", mcp_config) + + assert ( + 'GOOGLE_APPLICATION_CREDENTIALS="/etc/evalbench-sa-key/key.json" ' + 'gcloud auth application-default print-access-token ' + '--scopes="https://www.googleapis.com/auth/colaboratory"' + in translated["headersHelper"] + ) + 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_impersonation_without_gke_key_unsets_adc_for_gcloud(monkeypatch): + """Under GKE Workload Identity (no /etc/evalbench-sa-key/key.json), gcloud + must not receive GOOGLE_APPLICATION_CREDENTIALS=/tmp_sessions/impersonated_adc.json + when impersonating via CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT.""" + generator = object.__new__(ClaudeCodeGenerator) + generator.env = { + "GOOGLE_APPLICATION_CREDENTIALS": "/tmp_sessions/impersonated_adc.json", + "CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT": ( + "evalbench-runner@onemcp-evalviewer-dev.iam.gserviceaccount.com" + ), + } + + monkeypatch.setenv( + "GOOGLE_APPLICATION_CREDENTIALS", "/tmp_sessions/impersonated_adc.json" + ) + monkeypatch.setattr( + "generators.models.claude_code.os.path.exists", + lambda p: p == "/tmp_sessions/impersonated_adc.json", + ) + + with patch("generators.models.claude_code.subprocess.run") as mock_run: + mock_proc = MagicMock() + mock_proc.stdout = "wi_impersonated_token\n" + mock_run.return_value = mock_proc + + mcp_config = { + "httpUrl": "https://colaboratory.googleapis.com/mcp", + "authProviderType": "google_credentials", + "oauth": { + "scopes": ["https://www.googleapis.com/auth/colaboratory"] + }, + } + translated = generator._translate_mcp_config("colab-mcp", mcp_config) + + assert ( + "env -u GOOGLE_APPLICATION_CREDENTIALS " + "gcloud auth application-default print-access-token " + '--scopes="https://www.googleapis.com/auth/colaboratory"' + in translated["headersHelper"] + ) + 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 diff --git a/evalbench/test/codex_cli_test.py b/evalbench/test/codex_cli_test.py index ba287a60..8179b59d 100644 --- a/evalbench/test/codex_cli_test.py +++ b/evalbench/test/codex_cli_test.py @@ -234,3 +234,76 @@ def test_fetch_gcloud_access_token_rejects_unsafe_scopes(): generator._fetch_gcloud_access_token( scopes=['$(whoami)'] ) + + +def test_fetch_gcloud_access_token_impersonation_with_gke_key(monkeypatch): + generator = object.__new__(CodexCliGenerator) + 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.codex_cli.os.path.exists", + lambda p: p in ( + "/tmp_sessions/impersonated_adc.json", + "/etc/evalbench-sa-key/key.json", + ), + ) + + with patch("generators.models.codex_cli.subprocess.run") as mock_run: + mock_proc = MagicMock() + mock_proc.stdout = "impersonated_token_xyz\n" + mock_run.return_value = mock_proc + + token = generator._fetch_gcloud_access_token( + scopes=["https://www.googleapis.com/auth/colaboratory"] + ) + + assert token == "impersonated_token_xyz" + 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_impersonation_without_gke_key(monkeypatch): + generator = object.__new__(CodexCliGenerator) + generator.env = { + "GOOGLE_APPLICATION_CREDENTIALS": "/tmp_sessions/impersonated_adc.json", + "CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT": ( + "evalbench-runner@onemcp-evalviewer-dev.iam.gserviceaccount.com" + ), + } + + monkeypatch.setenv( + "GOOGLE_APPLICATION_CREDENTIALS", "/tmp_sessions/impersonated_adc.json" + ) + monkeypatch.setattr( + "generators.models.codex_cli.os.path.exists", + lambda p: p == "/tmp_sessions/impersonated_adc.json", + ) + + with patch("generators.models.codex_cli.subprocess.run") as mock_run: + mock_proc = MagicMock() + mock_proc.stdout = "wi_impersonated_token_xyz\n" + mock_run.return_value = mock_proc + + token = generator._fetch_gcloud_access_token( + scopes=["https://www.googleapis.com/auth/colaboratory"] + ) + + assert token == "wi_impersonated_token_xyz" + 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