Skip to content

fix(models): support impersonated ADC credentials in claude_code and codex_cli #706

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ikenwan-dev wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(models): support impersonated ADC credentials in claude_code and codex_cli #706

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
42 changes: 36 additions & 6 deletions evalbench/generators/models/claude_code.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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"'
)
Expand Down Expand Up @@ -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.
Expand Down
16 changes: 15 additions & 1 deletion evalbench/generators/models/codex_cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
135 changes: 135 additions & 0 deletions evalbench/test/claude_code_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
73 changes: 73 additions & 0 deletions evalbench/test/codex_cli_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading