feat: add LLMWithGateway for enterprise OAuth support

Add LLMWithGateway subclass to support enterprise API gateways with
OAuth 2.0 authentication.

Key features:
- OAuth 2.0 token fetch and automatic refresh
- Thread-safe token caching with TTL
- Custom header injection for gateway-specific requirements
- Template variable replacement for flexible configuration
- Fully generic implementation (no vendor lock-in)

Implementation approach:
- Separate LLMWithGateway class (addresses PR #963 feedback from neubig)
- Focused feature set for OAuth + custom headers (no over-engineering)
- Comprehensive test coverage

This replaces the previous approach of modifying the main LLM class,
keeping the codebase cleaner and more maintainable.

Example usage:
```python
llm = LLMWithGateway(
    model="gpt-4",
    base_url=os.environ["GATEWAY_BASE_URL"],
    gateway_auth_url=os.environ["GATEWAY_AUTH_URL"],
    gateway_auth_headers={
        "X-Client-Id": os.environ["GATEWAY_CLIENT_ID"],
        "X-Client-Secret": os.environ["GATEWAY_CLIENT_SECRET"],
    },
    gateway_auth_body={"grant_type": "client_credentials"},
    custom_headers={"X-Gateway-Key": os.environ["GATEWAY_API_KEY"]},
)
```

Files added:
- openhands-sdk/openhands/sdk/llm/llm_with_gateway.py (new class)
- tests/sdk/llm/test_llm_with_gateway.py (comprehensive tests)

Author: ak684

openhands-sdk/openhands/sdk/llm/__init__.py

@@ -1,6 +1,7 @@
 from openhands.sdk.llm.llm import LLM
 from openhands.sdk.llm.llm_registry import LLMRegistry, RegistryEvent
 from openhands.sdk.llm.llm_response import LLMResponse
+from openhands.sdk.llm.llm_with_gateway import LLMWithGateway
 from openhands.sdk.llm.message import (
     ImageContent,
     Message,
@@ -23,6 +24,7 @@
 __all__ = [
     "LLMResponse",
     "LLM",
+    "LLMWithGateway",
     "LLMRegistry",
     "RouterLLM",
     "RegistryEvent",

openhands-sdk/openhands/sdk/llm/llm_with_gateway.py

@@ -0,0 +1,333 @@
+"""LLM subclass with enterprise gateway support.
+
+This module provides LLMWithGateway, which extends the base LLM class to support
+OAuth 2.0 authentication flows and custom headers for enterprise API gateways.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from typing import Any
+
+import httpx
+from pydantic import Field, PrivateAttr
+
+from openhands.sdk.llm.llm import LLM
+from openhands.sdk.logger import get_logger
+
+logger = get_logger(__name__)
+
+__all__ = ["LLMWithGateway"]
+
+
+class LLMWithGateway(LLM):
+    """LLM subclass with enterprise gateway support.
+
+    Supports OAuth 2.0 token exchange with configurable headers and bodies.
+    Designed for enterprise API gateways that require:
+    1. Initial OAuth call to get a bearer token
+    2. Bearer token included in subsequent LLM API calls
+    3. Custom headers for routing/authentication
+
+    Example usage:
+        llm = LLMWithGateway(
+            model="gpt-4",
+            base_url="https://gateway.company.com/llm/v1",
+            gateway_auth_url="https://gateway.company.com/oauth/token",
+            gateway_auth_headers={
+                "X-Client-Id": os.environ["GATEWAY_CLIENT_ID"],
+                "X-Client-Secret": os.environ["GATEWAY_CLIENT_SECRET"],
+            },
+            gateway_auth_body={"grant_type": "client_credentials"},
+            custom_headers={"X-Gateway-Key": os.environ["GATEWAY_API_KEY"]},
+        )
+    """
+
+    # OAuth configuration
+    gateway_auth_url: str | None = Field(
+        default=None,
+        description="Identity provider URL to fetch gateway tokens (OAuth endpoint).",
+    )
+    gateway_auth_method: str = Field(
+        default="POST",
+        description="HTTP method for identity provider requests.",
+    )
+    gateway_auth_headers: dict[str, str] | None = Field(
+        default=None,
+        description="Headers to include when calling the identity provider.",
+    )
+    gateway_auth_body: dict[str, Any] | None = Field(
+        default=None,
+        description="JSON body to include when calling the identity provider.",
+    )
+    gateway_auth_token_path: str = Field(
+        default="access_token",
+        description="Dot-notation path to the token in the OAuth response (e.g., 'access_token' or 'data.token').",
+    )
+    gateway_auth_token_ttl: int | None = Field(
+        default=None,
+        description="Token TTL in seconds. If not set, defaults to 300s (5 minutes).",
+    )
+
+    # Token header configuration
+    gateway_token_header: str = Field(
+        default="Authorization",
+        description="Header name for the gateway token (defaults to 'Authorization').",
+    )
+    gateway_token_prefix: str = Field(
+        default="Bearer ",
+        description="Prefix prepended to the token (e.g., 'Bearer ').",
+    )
+
+    # Custom headers for all requests
+    custom_headers: dict[str, str] | None = Field(
+        default=None,
+        description="Custom headers to include with every LLM request.",
+    )
+
+    # Private fields for token management
+    _gateway_lock: threading.Lock = PrivateAttr(default_factory=threading.Lock)
+    _gateway_token: str | None = PrivateAttr(default=None)
+    _gateway_token_expiry: float | None = PrivateAttr(default=None)
+
+    def model_post_init(self, __context: Any) -> None:
+        """Initialize private fields after model validation."""
+        super().model_post_init(__context)
+        self._gateway_lock = threading.Lock()
+        self._gateway_token = None
+        self._gateway_token_expiry = None
+
+    def _completion(
+        self, messages: list[dict], **kwargs
+    ) -> Any:  # Returns ModelResponse
+        """Override to inject gateway authentication before calling LiteLLM."""
+        self._prepare_gateway_call(kwargs)
+        return super()._completion(messages, **kwargs)
+
+    def _responses(
+        self, messages: list[dict], **kwargs
+    ) -> Any:  # Returns ResponsesAPIResponse
+        """Override to inject gateway authentication before calling LiteLLM."""
+        self._prepare_gateway_call(kwargs)
+        return super()._responses(messages, **kwargs)
+
+    def _prepare_gateway_call(self, call_kwargs: dict[str, Any]) -> None:
+        """Augment LiteLLM kwargs with gateway headers and token.
+
+        This method:
+        1. Fetches/refreshes OAuth token if needed
+        2. Adds token to headers
+        3. Adds custom headers
+        4. Performs basic template variable replacement
+        """
+        if not self.gateway_auth_url and not self.custom_headers:
+            return
+
+        # Start with existing headers
+        headers: dict[str, str] = {}
+        existing_headers = call_kwargs.get("extra_headers")
+        if isinstance(existing_headers, dict):
+            headers.update(existing_headers)
+
+        # Add custom headers (with template replacement)
+        if self.custom_headers:
+            rendered_headers = self._render_templates(self.custom_headers)
+            if isinstance(rendered_headers, dict):
+                headers.update(rendered_headers)
+
+        # Add gateway token if OAuth is configured
+        if self.gateway_auth_url:
+            token_headers = self._get_gateway_token_headers()
+            if token_headers:
+                headers.update(token_headers)
+
+        # Set headers on the call
+        if headers:
+            call_kwargs["extra_headers"] = headers
+
+    def _get_gateway_token_headers(self) -> dict[str, str]:
+        """Get headers containing the gateway token."""
+        token = self._ensure_gateway_token()
+        if not token:
+            return {}
+
+        header_name = self.gateway_token_header
+        prefix = self.gateway_token_prefix
+        value = f"{prefix}{token}" if prefix else token
+        return {header_name: value}
+
+    def _ensure_gateway_token(self) -> str | None:
+        """Ensure we have a valid gateway token, refreshing if needed.
+
+        Returns:
+            Valid gateway token, or None if gateway auth is not configured.
+        """
+        if not self.gateway_auth_url:
+            return None
+
+        # Fast path: check if current token is still valid (with 5s buffer)
+        now = time.time()
+        if (
+            self._gateway_token
+            and self._gateway_token_expiry
+            and now < self._gateway_token_expiry - 5
+        ):
+            return self._gateway_token
+
+        # Slow path: acquire lock and refresh token
+        with self._gateway_lock:
+            # Double-check after acquiring lock
+            if (
+                self._gateway_token
+                and self._gateway_token_expiry
+                and time.time() < self._gateway_token_expiry - 5
+            ):
+                return self._gateway_token
+
+            # Refresh token
+            return self._refresh_gateway_token()
+
+    def _refresh_gateway_token(self) -> str:
+        """Fetch a new gateway token from the identity provider.
+
+        This method is called while holding _gateway_lock.
+
+        Returns:
+            Fresh gateway token.
+
+        Raises:
+            Exception: If token fetch fails.
+        """
+        method = self.gateway_auth_method.upper()
+        headers = self._render_templates(self.gateway_auth_headers or {})
+        body = self._render_templates(self.gateway_auth_body or {})
+
+        logger.debug(
+            f"Fetching gateway token from {self.gateway_auth_url} "
+            f"(method={method})"
+        )
+
+        try:
+            response = httpx.request(
+                method,
+                self.gateway_auth_url,
+                headers=headers if isinstance(headers, dict) else None,
+                json=body if isinstance(body, dict) else None,
+                timeout=self.timeout or 30,
+            )
+            response.raise_for_status()
+        except Exception as exc:
+            logger.error(f"Gateway auth request failed: {exc}")
+            raise
+
+        try:
+            payload = response.json()
+        except Exception as exc:
+            logger.error(f"Failed to parse gateway auth response JSON: {exc}")
+            raise
+
+        # Extract token from response
+        token_path = self.gateway_auth_token_path
+        token_value = self._extract_from_path(payload, token_path)
+        if not isinstance(token_value, str) or not token_value.strip():
+            raise ValueError(
+                f"Gateway auth response did not contain token at path "
+                f'"{token_path}". Response: {payload}'
+            )
+
+        # Determine TTL
+        ttl_seconds = float(self.gateway_auth_token_ttl or 300)
+
+        # Update cache
+        self._gateway_token = token_value.strip()
+        self._gateway_token_expiry = time.time() + max(ttl_seconds, 1.0)
+
+        logger.info(
+            f"Gateway token refreshed successfully (expires in {ttl_seconds}s)"
+        )
+        return self._gateway_token
+
+    def _render_templates(self, value: Any) -> Any:
+        """Replace template variables in strings with actual values.
+
+        Supports:
+        - {{llm_model}} -> self.model
+        - {{llm_base_url}} -> self.base_url
+        - {{llm_api_key}} -> self.api_key (if set)
+
+        Args:
+            value: String, dict, list, or other value to render.
+
+        Returns:
+            Value with templates replaced.
+        """
+        if isinstance(value, str):
+            replacements: dict[str, str] = {
+                "{{llm_model}}": self.model,
+                "{{llm_base_url}}": self.base_url or "",
+            }
+            if self.api_key:
+                replacements["{{llm_api_key}}"] = self.api_key.get_secret_value()
+
+            result = value
+            for placeholder, actual in replacements.items():
+                result = result.replace(placeholder, actual)
+            return result
+
+        if isinstance(value, dict):
+            return {k: self._render_templates(v) for k, v in value.items()}
+
+        if isinstance(value, list):
+            return [self._render_templates(v) for v in value]
+
+        return value
+
+    @staticmethod
+    def _extract_from_path(payload: Any, path: str) -> Any:
+        """Extract a value from nested dict/list using dot notation.
+
+        Examples:
+            _extract_from_path({"a": {"b": "value"}}, "a.b") -> "value"
+            _extract_from_path({"data": [{"token": "x"}]}, "data.0.token") -> "x"
+
+        Args:
+            payload: Dict or list to traverse.
+            path: Dot-separated path (e.g., "data.token" or "items.0.value").
+
+        Returns:
+            Value at the specified path.
+
+        Raises:
+            ValueError: If path cannot be traversed.
+        """
+        current = payload
+        if not path:
+            return current
+
+        for part in path.split("."):
+            if isinstance(current, dict):
+                current = current.get(part)
+                if current is None:
+                    raise ValueError(
+                        f'Key "{part}" not found in response while traversing path "{path}".'
+                    )
+            elif isinstance(current, list):
+                try:
+                    index = int(part)
+                except (ValueError, TypeError):
+                    raise ValueError(
+                        f'Invalid list index "{part}" while traversing response path "{path}".'
+                    ) from None
+                try:
+                    current = current[index]
+                except (IndexError, TypeError):
+                    raise ValueError(
+                        f'Index {index} out of range while traversing response path "{path}".'
+                    ) from None
+            else:
+                raise ValueError(
+                    f'Cannot traverse path "{path}"; segment "{part}" not found or not accessible.'
+                )
+
+        return current