From 1a24915b65235143263df833a812398d73d2b53d Mon Sep 17 00:00:00 2001 From: Amber Agent Date: Wed, 1 Apr 2026 16:19:49 +0000 Subject: [PATCH] fix(runner): update GITHUB_TOKEN for gh CLI after credential refresh Install a gh CLI wrapper script in /tmp/.ambient-bin and prepend that directory to PATH before the Claude Code subprocess spawns. The wrapper reads /tmp/.ambient_github_token on every invocation so mid-run credential refreshes propagate to gh CLI operations even after the subprocess environment is fixed at spawn time. This mirrors the existing git credential helper approach already used for git operations. Co-Authored-By: Claude Sonnet 4.6 --- .../ambient_runner/platform/auth.py | 70 +++++++- .../tests/test_shared_session_credentials.py | 159 ++++++++++++++++++ 2 files changed, 228 insertions(+), 1 deletion(-) diff --git a/components/runners/ambient-runner/ambient_runner/platform/auth.py b/components/runners/ambient-runner/ambient_runner/platform/auth.py index f43aa0636..fbf3b02e0 100755 --- a/components/runners/ambient-runner/ambient_runner/platform/auth.py +++ b/components/runners/ambient-runner/ambient_runner/platform/auth.py @@ -42,6 +42,13 @@ _GITHUB_TOKEN_FILE = Path("/tmp/.ambient_github_token") _GITLAB_TOKEN_FILE = Path("/tmp/.ambient_gitlab_token") +# Directory containing wrapper scripts that shadow real binaries so that +# mid-run credential refreshes are visible to already-spawned subprocesses. +# Prepended to PATH before the CLI subprocess is spawned so every child +# process inherits it. +_WRAPPER_BIN_DIR = "/tmp/.ambient-bin" +_GH_WRAPPER_PATH = f"{_WRAPPER_BIN_DIR}/gh" + # --------------------------------------------------------------------------- # Vertex AI credential validation (shared across all bridges) @@ -404,9 +411,10 @@ async def populate_runtime_credentials(context: RunnerContext) -> None: if github_creds.get("email"): git_user_email = github_creds["email"] - # Configure git identity and credential helper + # Configure git identity, credential helper, and gh wrapper await configure_git_identity(git_user_name, git_user_email) install_git_credential_helper() + install_gh_wrapper() if auth_failures: raise PermissionError( @@ -583,6 +591,66 @@ async def populate_mcp_server_credentials(context: RunnerContext) -> None: """ _credential_helper_installed = False # reset on every new process / deployment +_gh_wrapper_installed = False # reset on every new process / deployment + + +def install_gh_wrapper() -> None: + """Write a gh CLI wrapper and prepend its directory to PATH (once per process). + + The wrapper reads GITHUB_TOKEN from /tmp/.ambient_github_token so that + mid-run credential refreshes propagate into already-spawned CLI subprocesses. + Git operations are handled by the git credential helper; this wrapper covers + gh CLI calls (gh pr create, gh issue, etc.) which read GITHUB_TOKEN directly + from their environment (fixed at subprocess spawn time). + """ + global _gh_wrapper_installed + if _gh_wrapper_installed: + return + + import shutil + import stat + + # Find the real gh binary, excluding our wrapper dir to avoid a circular lookup. + current_path = os.environ.get("PATH", "") + search_path = ":".join(p for p in current_path.split(":") if p != _WRAPPER_BIN_DIR) + real_gh = shutil.which("gh", path=search_path) + if not real_gh: + logger.debug("gh CLI not found; skipping gh wrapper installation") + return + + wrapper_dir = Path(_WRAPPER_BIN_DIR) + try: + wrapper_dir.mkdir(parents=True, exist_ok=True) + wrapper_path = Path(_GH_WRAPPER_PATH) + wrapper_script = f"""\ +#!/bin/sh +# Ambient gh CLI wrapper. +# Reads GITHUB_TOKEN from the token file so mid-run credential refreshes +# propagate into already-spawned CLI subprocesses (subprocess env is fixed +# at creation time; the file is updated by the runner on every refresh). +if [ -f "/tmp/.ambient_github_token" ]; then + _token=$(cat /tmp/.ambient_github_token 2>/dev/null) + if [ -n "$_token" ]; then + GITHUB_TOKEN="$_token" + export GITHUB_TOKEN + fi +fi +exec {real_gh} "$@" +""" + wrapper_path.write_text(wrapper_script) + wrapper_path.chmod( + stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH + ) # 755 + + # Prepend wrapper dir to PATH so our wrapper takes precedence over the + # real gh binary for any subprocess spawned after this point. + if _WRAPPER_BIN_DIR not in current_path.split(":"): + os.environ["PATH"] = f"{_WRAPPER_BIN_DIR}:{current_path}" + + _gh_wrapper_installed = True + logger.info("Installed gh wrapper at %s", _GH_WRAPPER_PATH) + except Exception as e: + logger.warning(f"Failed to install gh wrapper: {e}") def install_git_credential_helper() -> None: diff --git a/components/runners/ambient-runner/tests/test_shared_session_credentials.py b/components/runners/ambient-runner/tests/test_shared_session_credentials.py index de73d363e..16a948b70 100644 --- a/components/runners/ambient-runner/tests/test_shared_session_credentials.py +++ b/components/runners/ambient-runner/tests/test_shared_session_credentials.py @@ -4,6 +4,7 @@ import os from http.server import BaseHTTPRequestHandler, HTTPServer from io import BytesIO +from pathlib import Path from threading import Thread from unittest.mock import AsyncMock, MagicMock, patch from urllib.error import HTTPError @@ -13,8 +14,11 @@ from ambient_runner.platform.auth import ( _GITHUB_TOKEN_FILE, _GITLAB_TOKEN_FILE, + _GH_WRAPPER_PATH, + _WRAPPER_BIN_DIR, _fetch_credential, clear_runtime_credentials, + install_gh_wrapper, populate_runtime_credentials, sanitize_user_context, ) @@ -710,3 +714,158 @@ async def test_returns_success_on_successful_refresh(self): assert result.get("isError") is None or result.get("isError") is False assert "successfully" in result["content"][0]["text"].lower() + + +# --------------------------------------------------------------------------- +# gh CLI wrapper (mid-run GITHUB_TOKEN refresh support) +# --------------------------------------------------------------------------- + + +class TestGhWrapper: + """Tests for gh CLI wrapper that propagates mid-run credential refreshes. + + The CLI subprocess is spawned once and its environment is fixed at that + point. Updating os.environ["GITHUB_TOKEN"] later does not propagate into + the subprocess. A wrapper script placed in PATH before the subprocess + spawns reads the token file on every gh invocation so the latest token is + always used. + """ + + import ambient_runner.platform.auth as _auth_module + + def _cleanup(self): + import shutil + + shutil.rmtree(_WRAPPER_BIN_DIR, ignore_errors=True) + _GITHUB_TOKEN_FILE.unlink(missing_ok=True) + + def _reset_installed_flag(self): + import ambient_runner.platform.auth as _auth + + _auth._gh_wrapper_installed = False + + def test_install_gh_wrapper_creates_wrapper_script(self): + """install_gh_wrapper writes an executable wrapper at _GH_WRAPPER_PATH.""" + self._cleanup() + self._reset_installed_flag() + try: + with patch("shutil.which", return_value="/usr/bin/gh"): + install_gh_wrapper() + + wrapper = Path(_GH_WRAPPER_PATH) + assert wrapper.exists(), "Wrapper script should be created" + content = wrapper.read_text() + assert "/tmp/.ambient_github_token" in content + assert "/usr/bin/gh" in content + # Must be executable + assert wrapper.stat().st_mode & 0o111, "Wrapper must be executable" + finally: + self._cleanup() + self._reset_installed_flag() + + def test_install_gh_wrapper_prepends_wrapper_dir_to_path(self): + """install_gh_wrapper prepends _WRAPPER_BIN_DIR to PATH.""" + self._cleanup() + self._reset_installed_flag() + original_path = os.environ.get("PATH", "") + try: + with patch("shutil.which", return_value="/usr/bin/gh"): + install_gh_wrapper() + + new_path = os.environ.get("PATH", "") + assert new_path.startswith(f"{_WRAPPER_BIN_DIR}:"), ( + "Wrapper dir should be prepended to PATH" + ) + finally: + os.environ["PATH"] = original_path + self._cleanup() + self._reset_installed_flag() + + def test_install_gh_wrapper_skips_when_gh_not_found(self): + """install_gh_wrapper does nothing when gh binary is not on PATH.""" + self._cleanup() + self._reset_installed_flag() + try: + with patch("shutil.which", return_value=None): + install_gh_wrapper() + + assert not Path(_GH_WRAPPER_PATH).exists(), ( + "Wrapper should not be created when gh is absent" + ) + finally: + self._cleanup() + self._reset_installed_flag() + + def test_install_gh_wrapper_is_idempotent(self): + """Calling install_gh_wrapper twice does not duplicate PATH entries.""" + self._cleanup() + self._reset_installed_flag() + original_path = os.environ.get("PATH", "") + try: + with patch("shutil.which", return_value="/usr/bin/gh"): + install_gh_wrapper() + install_gh_wrapper() # second call should be a no-op + + path_parts = os.environ.get("PATH", "").split(":") + count = path_parts.count(_WRAPPER_BIN_DIR) + assert count == 1, ( + f"Wrapper dir should appear exactly once in PATH, got {count}" + ) + finally: + os.environ["PATH"] = original_path + self._cleanup() + self._reset_installed_flag() + + def test_wrapper_script_reads_token_file_before_env_var(self): + """Wrapper script prioritises the token file over the env var (mid-run refresh).""" + self._cleanup() + self._reset_installed_flag() + try: + with patch("shutil.which", return_value="/usr/bin/gh"): + install_gh_wrapper() + + content = Path(_GH_WRAPPER_PATH).read_text() + # Token file read must appear before the exec line + file_read_pos = content.index("/tmp/.ambient_github_token") + exec_pos = content.index("exec /usr/bin/gh") + assert file_read_pos < exec_pos, ( + "Token file read must precede exec in wrapper script" + ) + finally: + self._cleanup() + self._reset_installed_flag() + + @pytest.mark.asyncio + async def test_populate_installs_gh_wrapper(self): + """populate_runtime_credentials installs the gh wrapper.""" + self._cleanup() + self._reset_installed_flag() + original_path = os.environ.get("PATH", "") + try: + with ( + patch("ambient_runner.platform.auth._fetch_credential") as mock_fetch, + patch("shutil.which", return_value="/usr/bin/gh"), + ): + + async def _creds(ctx, ctype): + if ctype == "github": + return { + "token": "gh-token", + "userName": "u", + "email": "u@e.com", + } + return {} + + mock_fetch.side_effect = _creds + ctx = _make_context() + await populate_runtime_credentials(ctx) + + assert Path(_GH_WRAPPER_PATH).exists(), ( + "populate_runtime_credentials should install the gh wrapper" + ) + finally: + os.environ["PATH"] = original_path + self._cleanup() + self._reset_installed_flag() + for key in ["GITHUB_TOKEN", "GIT_USER_NAME", "GIT_USER_EMAIL"]: + os.environ.pop(key, None)