diff --git a/pr_agent/algo/ai_handlers/cloud_auth.py b/pr_agent/algo/ai_handlers/cloud_auth.py index 74f282be5f..fe77f40a4b 100644 --- a/pr_agent/algo/ai_handlers/cloud_auth.py +++ b/pr_agent/algo/ai_handlers/cloud_auth.py @@ -379,6 +379,15 @@ def _require_litellm_interface(interface, name: str, methods=()): ) +def _rebind_to_globals(native, native_globals): + """Rebind a native function onto captured globals, keeping its own defaults.""" + bound = FunctionType( + native.__code__, native_globals, native.__name__, native.__defaults__, native.__closure__, + ) + bound.__kwdefaults__ = copy.copy(native.__kwdefaults__) + return bound + + _azure_oidc_request = ContextVar("pr_agent_azure_oidc_request", default=None) _azure_ad_responses_request = ContextVar("pr_agent_azure_ad_responses_request", default=None) _azure_oidc_bridge = None @@ -433,8 +442,7 @@ def set_cache(key, *args, **kwargs): # and refresh/TTL behavior. Only the outer exchange's inputs are request-owned. native_globals["os"] = SimpleNamespace(getenv=captured_getenv) native_globals["azure_ad_cache"] = SimpleNamespace(get_cache=get_cache, set_cache=set_cache) - exchange = FunctionType(native.__code__, native_globals, native.__name__, native.__defaults__, native.__closure__) - exchange.__kwdefaults__ = copy.copy(native.__kwdefaults__) + exchange = _rebind_to_globals(native, native_globals) return exchange( selector, azure_client_id=environment["AZURE_CLIENT_ID"], @@ -476,11 +484,7 @@ def cached(tenant_id, client_id, client_secret, scope): # Retain native secret-selector resolution, but never use the native # authority-free provider cache. Credentials retain authority on refresh. globals_view = {**original_entra.__globals__, "_cached_entra_id_token_provider": cached} - factory = FunctionType( - original_entra.__code__, globals_view, original_entra.__name__, - original_entra.__defaults__, original_entra.__closure__, - ) - factory.__kwdefaults__ = copy.copy(original_entra.__kwdefaults__) + factory = _rebind_to_globals(original_entra, globals_view) return factory(tenant_id, client_id, client_secret, scope) def password(client_id, azure_username, azure_password, scope="https://cognitiveservices.azure.com/.default"): @@ -571,11 +575,7 @@ def getenv(name, default=None): native_globals["os"] = SimpleNamespace(getenv=getenv) native_globals["get_secret_str"] = get_secret_str native_globals["get_azure_ad_token_from_oidc"] = exchange - bound = FunctionType( - original_resolver.__code__, native_globals, original_resolver.__name__, - original_resolver.__defaults__, original_resolver.__closure__, - ) - bound.__kwdefaults__ = copy.copy(original_resolver.__kwdefaults__) + bound = _rebind_to_globals(original_resolver, native_globals) return bound(litellm_params) @wraps(original_resolver) @@ -619,11 +619,7 @@ def initialize(*args, **kwargs): ): return original_initialize(*args, **kwargs) native_globals = _azure_oidc_companion_globals(original_initialize.__globals__, context) - factory = FunctionType( - original_initialize.__code__, native_globals, original_initialize.__name__, - original_initialize.__defaults__, original_initialize.__closure__, - ) - factory.__kwdefaults__ = copy.copy(original_initialize.__kwdefaults__) + factory = _rebind_to_globals(original_initialize, native_globals) return factory(*args, **kwargs) @wraps(original_cache_key) @@ -1212,6 +1208,14 @@ def _raw_guard_has_header_only_auth(provider, params, headers): return not ((tenant and client and secret) or (username and password and client)) +def _require_http_auth_signature(original): + """Reject a native HTTP validator that no longer exposes the request auth parameters.""" + signature = inspect.signature(original) + if not {"api_key", "headers", "litellm_params"}.issubset(signature.parameters): + raise RuntimeError("LiteLLM's native HTTP authentication interface is incompatible") + return signature + + def _install_raw_api_key_guard_override_bridge(provider): """Undo generated auth headers without reopening a native key resolver.""" if provider not in ("azure_ai", "ragflow", "xai"): @@ -1225,9 +1229,7 @@ def _install_raw_api_key_guard_override_bridge(provider): marker = "_pr_agent_original_raw_override_validate_environment" if getattr(original, marker, None) is not None: return - signature = inspect.signature(original) - if not {"api_key", "headers", "litellm_params"}.issubset(signature.parameters): - raise RuntimeError("LiteLLM's native HTTP authentication interface is incompatible") + signature = _require_http_auth_signature(original) def validate_environment(self, *args, **kwargs): if _raw_api_key_guard_provider.get() != provider or type(self) is not config: @@ -1281,9 +1283,7 @@ def _install_raw_api_key_guard_bridge(): marker = "_pr_agent_original_raw_validate_environment" if getattr(original, marker, None) is not None: return - signature = inspect.signature(original) - if not {"api_key", "headers", "litellm_params"}.issubset(signature.parameters): - raise RuntimeError("LiteLLM's native HTTP authentication interface is incompatible") + signature = _require_http_auth_signature(original) def validate_environment(self, *args, **kwargs): provider = _raw_api_key_guard_provider.get()