open-swe/tests/test_proxy_auth.py

"""Tests for GitHub proxy auth configuration."""

from __future__ import annotations

import base64
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest

from agent.integrations.langsmith import _configure_github_proxy


class TestSandboxFactoryLoading:
    def test_create_sandbox_loads_only_selected_provider(self) -> None:
        with (
            patch("agent.utils.sandbox.import_module") as mock_import_module,
            patch.dict("os.environ", {"SANDBOX_TYPE": "local"}),
        ):
            module = MagicMock()
            module.create_local_sandbox.return_value = MagicMock(id="local")
            mock_import_module.return_value = module

            from agent.utils.sandbox import create_sandbox

            sandbox = create_sandbox("existing")

        assert sandbox.id == "local"
        mock_import_module.assert_called_once_with("agent.integrations.local")
        module.create_local_sandbox.assert_called_once_with("existing")


class TestConfigureGithubProxy:
    """Tests for _configure_github_proxy payload shape and error handling."""

    def test_sends_correct_payload_shape(self) -> None:
        """Verify the PATCH request uses opaque headers with correct structure."""
        token = "ghs_testtoken123"
        expected_basic = base64.b64encode(f"x-access-token:{token}".encode()).decode()

        with (
            patch("agent.integrations.langsmith.httpx.Client") as mock_client_cls,
            patch.dict("os.environ", {"LANGSMITH_API_KEY": "ls-api-key"}),
        ):
            mock_client = MagicMock()
            mock_response = MagicMock()
            mock_response.raise_for_status = MagicMock()
            mock_client.patch.return_value = mock_response
            mock_client_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
            mock_client_cls.return_value.__exit__ = MagicMock(return_value=False)

            _configure_github_proxy("sandbox-abc123", token)

            mock_client.patch.assert_called_once()
            call_kwargs = mock_client.patch.call_args
            payload = call_kwargs.kwargs["json"]

            assert "proxy_config" in payload
            rules = payload["proxy_config"]["rules"]
            assert len(rules) == 2

            api_rule = rules[0]
            assert api_rule["name"] == "github-api"
            assert api_rule["match_hosts"] == ["api.github.com"]
            api_headers = api_rule["headers"]
            assert len(api_headers) == 1
            assert api_headers[0]["name"] == "Authorization"
            assert api_headers[0]["type"] == "opaque"
            assert api_headers[0]["value"] == f"Bearer {token}"

            web_rule = rules[1]
            assert web_rule["name"] == "github"
            assert web_rule["match_hosts"] == ["github.com", "*.github.com"]

            headers = web_rule["headers"]
            assert len(headers) == 1
            assert headers[0]["name"] == "Authorization"
            assert headers[0]["type"] == "opaque"
            assert headers[0]["value"] == f"Basic {expected_basic}"

    def test_sends_to_correct_url(self) -> None:
        """Verify the PATCH hits the right endpoint."""
        with (
            patch("agent.integrations.langsmith.httpx.Client") as mock_client_cls,
            patch.dict(
                "os.environ",
                {
                    "LANGSMITH_ENDPOINT": "https://test.api.smith.langchain.com",
                    "LANGSMITH_API_KEY": "api-key",
                },
            ),
        ):
            mock_client = MagicMock()
            mock_response = MagicMock()
            mock_response.raise_for_status = MagicMock()
            mock_client.patch.return_value = mock_response
            mock_client_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
            mock_client_cls.return_value.__exit__ = MagicMock(return_value=False)

            _configure_github_proxy("sandbox-xyz", "token")

            url = mock_client.patch.call_args.args[0]
            assert url == "https://test.api.smith.langchain.com/v2/sandboxes/boxes/sandbox-xyz"

    def test_sends_api_key_header(self) -> None:
        """Verify the PATCH includes the LangSmith API key."""
        with (
            patch("agent.integrations.langsmith.httpx.Client") as mock_client_cls,
            patch.dict("os.environ", {"LANGSMITH_API_KEY": "my-api-key"}),
        ):
            mock_client = MagicMock()
            mock_response = MagicMock()
            mock_response.raise_for_status = MagicMock()
            mock_client.patch.return_value = mock_response
            mock_client_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
            mock_client_cls.return_value.__exit__ = MagicMock(return_value=False)

            _configure_github_proxy("sandbox-abc", "token")

            headers = mock_client.patch.call_args.kwargs["headers"]
            assert headers == {"X-API-Key": "my-api-key"}

    def test_retries_transient_http_error(self) -> None:
        """Transient proxy API errors should be retried on the same sandbox."""
        request = httpx.Request(
            "PATCH", "https://api.smith.langchain.com/v2/sandboxes/boxes/sandbox-abc"
        )
        response = httpx.Response(503, request=request)
        transient_error = httpx.HTTPStatusError(
            "Server error",
            request=request,
            response=response,
        )
        with (
            patch("agent.integrations.langsmith.httpx.Client") as mock_client_cls,
            patch("agent.integrations.langsmith.time.sleep") as mock_sleep,
            patch.dict("os.environ", {"LANGSMITH_API_KEY": "api-key"}),
        ):
            mock_client = MagicMock()
            failed_response = MagicMock()
            failed_response.raise_for_status.side_effect = transient_error
            successful_response = MagicMock()
            successful_response.raise_for_status = MagicMock()
            mock_client.patch.side_effect = [failed_response, successful_response]
            mock_client_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
            mock_client_cls.return_value.__exit__ = MagicMock(return_value=False)

            _configure_github_proxy("sandbox-abc", "token")

            assert mock_client.patch.call_count == 2
            mock_sleep.assert_called_once()

    def test_raises_on_non_retryable_http_error(self) -> None:
        """Non-retryable HTTP errors should propagate without retrying."""
        request = httpx.Request(
            "PATCH", "https://api.smith.langchain.com/v2/sandboxes/boxes/sandbox-abc"
        )
        response = httpx.Response(400, request=request)
        error = httpx.HTTPStatusError("Bad request", request=request, response=response)
        with (
            patch("agent.integrations.langsmith.httpx.Client") as mock_client_cls,
            patch("agent.integrations.langsmith.time.sleep") as mock_sleep,
            patch.dict("os.environ", {"LANGSMITH_API_KEY": "api-key"}),
        ):
            mock_client = MagicMock()
            failed_response = MagicMock()
            failed_response.raise_for_status.side_effect = error
            mock_client.patch.return_value = failed_response
            mock_client_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
            mock_client_cls.return_value.__exit__ = MagicMock(return_value=False)

            with pytest.raises(httpx.HTTPStatusError):
                _configure_github_proxy("sandbox-abc", "token")

            mock_client.patch.assert_called_once()
            mock_sleep.assert_not_called()


class TestCreateSandboxWithProxy:
    """Tests for _create_sandbox_with_proxy token source selection."""

    @pytest.mark.asyncio
    async def test_uses_installation_token_for_langsmith(self) -> None:
        """Installation token should be used for proxy auth on langsmith sandboxes."""
        with (
            patch(
                "agent.server.get_github_app_installation_token_with_expiry",
                new_callable=AsyncMock,
                return_value=("ghs_install", None),
            ),
            patch("agent.server.create_sandbox") as mock_create,
            patch("agent.server._configure_github_proxy") as mock_proxy,
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith", "LANGSMITH_API_KEY": "ls-key"}),
        ):
            mock_create.return_value = MagicMock(id="sandbox-123")

            from agent.server import _create_sandbox_with_proxy

            await _create_sandbox_with_proxy()

            mock_create.assert_called_once_with(snapshot_id=None)
            mock_proxy.assert_called_once_with("sandbox-123", "ghs_install")

    @pytest.mark.asyncio
    async def test_skips_proxy_for_non_langsmith(self) -> None:
        """Non-langsmith sandboxes should skip proxy configuration."""
        with (
            patch("agent.server.create_sandbox") as mock_create,
            patch("agent.server._configure_github_proxy") as mock_proxy,
            patch.dict("os.environ", {"SANDBOX_TYPE": "daytona"}),
        ):
            mock_create.return_value = MagicMock(id="sandbox-456")

            from agent.server import _create_sandbox_with_proxy

            await _create_sandbox_with_proxy()

            mock_create.assert_called_once_with(snapshot_id=None)
            mock_proxy.assert_not_called()

    @pytest.mark.asyncio
    async def test_raises_when_no_installation_token_for_langsmith(self) -> None:
        """Should raise ValueError when installation token is unavailable for langsmith."""
        with (
            patch("agent.server.create_sandbox") as mock_create,
            patch(
                "agent.server.get_github_app_installation_token_with_expiry",
                new_callable=AsyncMock,
                return_value=(None, None),
            ),
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
        ):
            mock_create.return_value = MagicMock(id="sandbox-789")

            from agent.server import _create_sandbox_with_proxy

            with pytest.raises(ValueError, match="installation token is unavailable"):
                await _create_sandbox_with_proxy()


class _DummyAgent:
    def with_config(self, config):
        return self


class TestRefreshProxyOnSandboxReuse:
    """Tests for refreshing GitHub proxy auth on sandbox reuse."""

    @staticmethod
    def _execution_config() -> dict:
        return {
            "configurable": {
                "__is_for_execution__": True,
                "thread_id": "thread-123",
                "repo": {"owner": "langchain-ai", "name": "open-swe"},
            },
            "metadata": {},
        }

    @pytest.mark.asyncio
    async def test_refreshes_proxy_for_cached_langsmith_sandbox(self) -> None:
        """Cached sandboxes should get a fresh proxy token before git operations."""
        config = self._execution_config()
        mock_sandbox = MagicMock(id="sandbox-cached")

        with (
            patch(
                "agent.server.resolve_github_token",
                new_callable=AsyncMock,
                return_value=("ghp", None),
            ),
            patch(
                "agent.server.get_sandbox_id_from_metadata",
                new_callable=AsyncMock,
                return_value="sandbox-cached",
            ),
            patch(
                "agent.server.get_github_app_installation_token_with_expiry",
                new_callable=AsyncMock,
                return_value=("ghs_fresh", None),
            ),
            patch("agent.server._configure_github_proxy") as mock_proxy,
            patch(
                "agent.server.aresolve_sandbox_work_dir",
                new_callable=AsyncMock,
                return_value="/workspace",
            ),
            patch(
                "agent.server.check_or_recreate_sandbox",
                new_callable=AsyncMock,
                return_value=mock_sandbox,
            ),
            patch("agent.server.make_model", return_value=MagicMock()),
            patch("agent.server.construct_system_prompt", return_value="prompt"),
            patch("agent.server.create_deep_agent", return_value=_DummyAgent()),
            patch.dict(
                "agent.server.SANDBOX_BACKENDS",
                {"thread-123": mock_sandbox},
                clear=True,
            ),
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
        ):
            from agent.server import get_agent

            await get_agent(config)

            mock_proxy.assert_called_once_with("sandbox-cached", "ghs_fresh")

    @pytest.mark.asyncio
    async def test_refreshes_proxy_when_reconnecting_to_existing_langsmith_sandbox(self) -> None:
        """Reconnected sandboxes should also get a fresh proxy token."""
        config = self._execution_config()
        mock_sandbox = MagicMock(id="sandbox-existing")

        with (
            patch(
                "agent.server.resolve_github_token",
                new_callable=AsyncMock,
                return_value=("ghp", None),
            ),
            patch(
                "agent.server.get_sandbox_id_from_metadata",
                new_callable=AsyncMock,
                return_value="sandbox-existing",
            ),
            patch("agent.server.create_sandbox", return_value=mock_sandbox) as mock_create,
            patch(
                "agent.server.get_github_app_installation_token_with_expiry",
                new_callable=AsyncMock,
                return_value=("ghs_fresh", None),
            ),
            patch("agent.server._configure_github_proxy") as mock_proxy,
            patch(
                "agent.server.aresolve_sandbox_work_dir",
                new_callable=AsyncMock,
                return_value="/workspace",
            ),
            patch("agent.server.make_model", return_value=MagicMock()),
            patch("agent.server.construct_system_prompt", return_value="prompt"),
            patch("agent.server.create_deep_agent", return_value=_DummyAgent()),
            patch.dict("agent.server.SANDBOX_BACKENDS", {}, clear=True),
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
        ):
            from agent.server import get_agent

            await get_agent(config)

            mock_create.assert_called_once_with("sandbox-existing")
            mock_proxy.assert_called_once_with("sandbox-existing", "ghs_fresh")

    @pytest.mark.asyncio
    async def test_proxy_refresh_failure_recreates_sandbox(self) -> None:
        """A stale sandbox whose proxy cannot be patched should be replaced."""
        mock_sandbox = MagicMock(id="sandbox-stale")
        replacement_sandbox = MagicMock(id="sandbox-replacement")
        request = httpx.Request(
            "PATCH", "https://api.smith.langchain.com/v2/sandboxes/boxes/sandbox-stale"
        )
        response = httpx.Response(400, request=request)

        with (
            patch(
                "agent.server.get_github_app_installation_token_with_expiry",
                new_callable=AsyncMock,
                return_value=("ghs_fresh", None),
            ),
            patch(
                "agent.server._configure_github_proxy",
                side_effect=httpx.HTTPStatusError(
                    "Bad request",
                    request=request,
                    response=response,
                ),
            ) as mock_proxy,
            patch(
                "agent.server._recreate_sandbox",
                new_callable=AsyncMock,
                return_value=replacement_sandbox,
            ) as mock_recreate,
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
        ):
            from agent.server import _refresh_github_proxy_or_recreate

            sandbox = await _refresh_github_proxy_or_recreate(mock_sandbox, "thread-123")

            assert sandbox is replacement_sandbox
            mock_proxy.assert_called_once_with("sandbox-stale", "ghs_fresh")
            mock_recreate.assert_awaited_once_with(
                "thread-123",
                github_proxy_token=None,
                github_proxy_repositories=None,
                repo=None,
            )

    @pytest.mark.asyncio
    async def test_starts_stopped_langsmith_sandbox_before_proxy_refresh(self) -> None:
        """Proxy config requires a running LangSmith sandbox."""
        inner_sandbox = MagicMock(name="sandbox-stopped")
        inner_sandbox.name = "sandbox-stopped"
        inner_sandbox._client.get_sandbox_status.return_value = MagicMock(status="stopped")

        with (
            patch(
                "agent.server.get_github_app_installation_token_with_expiry",
                new_callable=AsyncMock,
                return_value=("ghs_fresh", None),
            ),
            patch("agent.server._configure_github_proxy") as mock_proxy,
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
        ):
            from agent.server import LangSmithSandbox, _refresh_github_proxy

            sandbox_backend = object.__new__(LangSmithSandbox)
            sandbox_backend._sandbox = inner_sandbox

            await _refresh_github_proxy(sandbox_backend)

            inner_sandbox._client.get_sandbox_status.assert_called_once_with("sandbox-stopped")
            inner_sandbox.start.assert_called_once_with()
            mock_proxy.assert_called_once_with("sandbox-stopped", "ghs_fresh")

    @pytest.mark.asyncio
    async def test_skips_start_for_ready_langsmith_sandbox_before_proxy_refresh(self) -> None:
        """Ready sandboxes can be patched without starting again."""
        inner_sandbox = MagicMock(name="sandbox-ready")
        inner_sandbox.name = "sandbox-ready"
        inner_sandbox._client.get_sandbox_status.return_value = MagicMock(status="ready")

        with (
            patch(
                "agent.server.get_github_app_installation_token_with_expiry",
                new_callable=AsyncMock,
                return_value=("ghs_fresh", None),
            ),
            patch("agent.server._configure_github_proxy") as mock_proxy,
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
        ):
            from agent.server import LangSmithSandbox, _refresh_github_proxy

            sandbox_backend = object.__new__(LangSmithSandbox)
            sandbox_backend._sandbox = inner_sandbox

            await _refresh_github_proxy(sandbox_backend)

            inner_sandbox._client.get_sandbox_status.assert_called_once_with("sandbox-ready")
            inner_sandbox.start.assert_not_called()
            mock_proxy.assert_called_once_with("sandbox-ready", "ghs_fresh")