Merge pull request #2439 from jlowin/oidc-proxy-extra-params

Add extra_authorize_params and extra_token_params to OIDCProxy

src/fastmcp/server/auth/oidc_proxy.py

@@ -222,6 +222,9 @@ class OIDCProxy(OAuthProxy):
         token_endpoint_auth_method: str | None = None,
         # Consent screen configuration
         require_authorization_consent: bool = True,
+        # Extra parameters
+        extra_authorize_params: dict[str, str] | None = None,
+        extra_token_params: dict[str, str] | None = None,
     ) -> None:
         """Initialize the OIDC proxy provider.
 
@@ -259,6 +262,11 @@ class OIDCProxy(OAuthProxy):
                 When True, users see a consent screen before being redirected to the upstream IdP.
                 When False, authorization proceeds directly without user confirmation.
                 SECURITY WARNING: Only disable for local development or testing environments.
+            extra_authorize_params: Additional parameters to forward to the upstream authorization endpoint.
+                Useful for provider-specific parameters like prompt=consent or access_type=offline.
+                Example: {"prompt": "consent", "access_type": "offline"}
+            extra_token_params: Additional parameters to forward to the upstream token endpoint.
+                Useful for provider-specific parameters during token exchange.
         """
         if not config_url:
             raise ValueError("Missing required config URL")
@@ -335,10 +343,24 @@ class OIDCProxy(OAuthProxy):
         if redirect_path:
             init_kwargs["redirect_path"] = redirect_path
 
+        # Build extra params, merging audience with user-provided params
+        # User params override audience if there's a conflict
+        final_authorize_params: dict[str, str] = {}
+        final_token_params: dict[str, str] = {}
+
         if audience:
-            extra_params = {"audience": audience}
-            init_kwargs["extra_authorize_params"] = extra_params
-            init_kwargs["extra_token_params"] = extra_params
+            final_authorize_params["audience"] = audience
+            final_token_params["audience"] = audience
+
+        if extra_authorize_params:
+            final_authorize_params.update(extra_authorize_params)
+        if extra_token_params:
+            final_token_params.update(extra_token_params)
+
+        if final_authorize_params:
+            init_kwargs["extra_authorize_params"] = final_authorize_params
+        if final_token_params:
+            init_kwargs["extra_token_params"] = final_token_params
 
         super().__init__(**init_kwargs)  # ty: ignore[invalid-argument-type]
 

tests/server/auth/test_oidc_proxy.py

@@ -783,3 +783,96 @@ class TestOIDCProxyInitialization:
             validate_proxy(mock_get, proxy, oidc_config)
             assert proxy._extra_authorize_params == {"audience": "test-audience"}
             assert proxy._extra_token_params == {"audience": "test-audience"}
+
+    def test_extra_authorize_params_initialization(self, valid_oidc_configuration_dict):
+        """Test extra authorize params initialization."""
+        with patch(
+            "fastmcp.server.auth.oidc_proxy.OIDCConfiguration.get_oidc_configuration"
+        ) as mock_get:
+            oidc_config = OIDCConfiguration.model_validate(
+                valid_oidc_configuration_dict
+            )
+            mock_get.return_value = oidc_config
+
+            proxy = OIDCProxy(
+                config_url=TEST_CONFIG_URL,
+                client_id=TEST_CLIENT_ID,
+                client_secret=TEST_CLIENT_SECRET,
+                base_url=TEST_BASE_URL,
+                jwt_signing_key="test-secret",
+                extra_authorize_params={
+                    "prompt": "consent",
+                    "access_type": "offline",
+                },
+            )
+
+            validate_proxy(mock_get, proxy, oidc_config)
+
+            assert proxy._extra_authorize_params == {
+                "prompt": "consent",
+                "access_type": "offline",
+            }
+            # Token params should be empty since we didn't set them
+            assert proxy._extra_token_params == {}
+
+    def test_extra_token_params_initialization(self, valid_oidc_configuration_dict):
+        """Test extra token params initialization."""
+        with patch(
+            "fastmcp.server.auth.oidc_proxy.OIDCConfiguration.get_oidc_configuration"
+        ) as mock_get:
+            oidc_config = OIDCConfiguration.model_validate(
+                valid_oidc_configuration_dict
+            )
+            mock_get.return_value = oidc_config
+
+            proxy = OIDCProxy(
+                config_url=TEST_CONFIG_URL,
+                client_id=TEST_CLIENT_ID,
+                client_secret=TEST_CLIENT_SECRET,
+                base_url=TEST_BASE_URL,
+                jwt_signing_key="test-secret",
+                extra_token_params={"custom_param": "custom_value"},
+            )
+
+            validate_proxy(mock_get, proxy, oidc_config)
+
+            # Authorize params should be empty since we didn't set them
+            assert proxy._extra_authorize_params == {}
+            assert proxy._extra_token_params == {"custom_param": "custom_value"}
+
+    def test_extra_params_merge_with_audience(self, valid_oidc_configuration_dict):
+        """Test that extra params merge with audience, with user params taking precedence."""
+        with patch(
+            "fastmcp.server.auth.oidc_proxy.OIDCConfiguration.get_oidc_configuration"
+        ) as mock_get:
+            oidc_config = OIDCConfiguration.model_validate(
+                valid_oidc_configuration_dict
+            )
+            mock_get.return_value = oidc_config
+
+            proxy = OIDCProxy(
+                config_url=TEST_CONFIG_URL,
+                client_id=TEST_CLIENT_ID,
+                client_secret=TEST_CLIENT_SECRET,
+                base_url=TEST_BASE_URL,
+                audience="original-audience",
+                jwt_signing_key="test-secret",
+                extra_authorize_params={
+                    "prompt": "consent",
+                    "audience": "overridden-audience",  # Should override the audience param
+                },
+                extra_token_params={"custom": "value"},
+            )
+
+            validate_proxy(mock_get, proxy, oidc_config)
+
+            # User's extra_authorize_params should override audience
+            assert proxy._extra_authorize_params == {
+                "audience": "overridden-audience",
+                "prompt": "consent",
+            }
+            # Token params should have both audience (from audience param) and custom
+            assert proxy._extra_token_params == {
+                "audience": "original-audience",
+                "custom": "value",
+            }