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)