fix: retry transient LangSmith proxy config failures (#1413)

agent/integrations/langsmith.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import base64
 import logging
 import os
+import time
 from abc import ABC, abstractmethod
 from typing import Any
 
@@ -20,6 +21,10 @@ DEFAULT_SANDBOX_VCPUS = 2
 DEFAULT_SANDBOX_MEM_BYTES = 7936 * 1024**2  # 7936 MiB ("large" tier cap)
 DEFAULT_SANDBOX_IDLE_TTL_SECONDS = 10 * 60  # 10 minutes
 DEFAULT_SANDBOX_DELETE_AFTER_STOP_SECONDS = 24 * 60 * 60  # 24 hours
+PROXY_CONFIG_MAX_ATTEMPTS = 3
+PROXY_CONFIG_TIMEOUT_SECONDS = 10.0
+PROXY_CONFIG_RETRY_DELAYS_SECONDS = (0.5, 1.0)
+PROXY_CONFIG_RETRYABLE_STATUS_CODES = frozenset({408, 409, 425, 429, 500, 502, 503, 504, 529})
 
 
 def _get_langsmith_api_key() -> str | None:
@@ -95,6 +100,25 @@ def _github_proxy_rules(github_token: str) -> list[dict[str, Any]]:
     ]
 
 
+def _retry_after_seconds(response: httpx.Response | None) -> float | None:
+    if response is None:
+        return None
+    raw = response.headers.get("Retry-After")
+    if not raw:
+        return None
+    try:
+        delay = float(raw)
+    except ValueError:
+        return None
+    return max(delay, 0.0)
+
+
+def _is_retryable_proxy_config_error(exc: BaseException) -> bool:
+    if isinstance(exc, httpx.HTTPStatusError):
+        return exc.response.status_code in PROXY_CONFIG_RETRYABLE_STATUS_CODES
+    return isinstance(exc, httpx.TransportError)
+
+
 def _configure_github_proxy(sandbox_name: str, github_token: str) -> None:
     """Configure sandbox proxy to inject GitHub auth for GitHub traffic.
 
@@ -113,13 +137,39 @@ def _configure_github_proxy(sandbox_name: str, github_token: str) -> None:
     langsmith_endpoint = os.environ.get("LANGSMITH_ENDPOINT", "https://api.smith.langchain.com")
     url = f"{langsmith_endpoint}/v2/sandboxes/boxes/{sandbox_name}"
     payload = {"proxy_config": {"rules": _github_proxy_rules(github_token)}}
-    with httpx.Client() as client:
-        response = client.patch(
-            url,
-            json=payload,
-            headers={"X-API-Key": api_key},
-        )
-        response.raise_for_status()
+    with httpx.Client(timeout=PROXY_CONFIG_TIMEOUT_SECONDS) as client:
+        for attempt in range(PROXY_CONFIG_MAX_ATTEMPTS):
+            try:
+                response = client.patch(
+                    url,
+                    json=payload,
+                    headers={"X-API-Key": api_key},
+                )
+                response.raise_for_status()
+                break
+            except Exception as exc:
+                if attempt == PROXY_CONFIG_MAX_ATTEMPTS - 1 or not _is_retryable_proxy_config_error(
+                    exc
+                ):
+                    raise
+                retry_after = (
+                    _retry_after_seconds(exc.response)
+                    if isinstance(exc, httpx.HTTPStatusError)
+                    else None
+                )
+                delay = (
+                    retry_after
+                    or PROXY_CONFIG_RETRY_DELAYS_SECONDS[
+                        min(attempt, len(PROXY_CONFIG_RETRY_DELAYS_SECONDS) - 1)
+                    ]
+                )
+                logger.warning(
+                    "Failed to configure GitHub proxy for sandbox %s (%s); retrying in %.1fs",
+                    sandbox_name,
+                    type(exc).__name__,
+                    delay,
+                )
+                time.sleep(delay)
     logger.info("Configured GitHub proxy for sandbox %s", sandbox_name)
 
 

tests/test_proxy_auth.py

@@ -101,22 +101,61 @@ class TestConfigureGithubProxy:
             headers = mock_client.patch.call_args.kwargs["headers"]
             assert headers == {"X-API-Key": "my-api-key"}
 
-    def test_raises_on_http_error(self) -> None:
-        """Verify HTTP errors propagate."""
+    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()
-            mock_client.patch.side_effect = httpx.HTTPStatusError(
-                "Server error", request=MagicMock(), response=MagicMock(status_code=500)
-            )
+            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."""