From bc0f9d0192bf143f7331c98b29f8e1f2f86c2e6c Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 26 Aug 2025 13:33:26 +0000 Subject: [PATCH 1/4] feat: Add audience and prompt to OAuth2 flow This change adds the ability to specify an `audience` and `prompt` when generating an OAuth2 authorization URL. This is necessary for some OAuth2 providers that require these parameters. The following changes were made: - Added optional `audience` and `prompt` fields to the `OAuth2Auth` model. - Updated the `AuthHandler` to include the `audience` and `prompt` parameters in the authorization URL. - Added a unit test to verify that the `audience` and `prompt` parameters are correctly included in the authorization URL. --- src/google/adk/auth/auth_credential.py | 2 ++ src/google/adk/auth/auth_handler.py | 9 ++++++- tests/unittests/auth/test_auth_handler.py | 29 ++++++++++++++++++++++- 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/src/google/adk/auth/auth_credential.py b/src/google/adk/auth/auth_credential.py index 34d04dde93..4e4a5fef42 100644 --- a/src/google/adk/auth/auth_credential.py +++ b/src/google/adk/auth/auth_credential.py @@ -79,6 +79,8 @@ class OAuth2Auth(BaseModelWithConfig): refresh_token: Optional[str] = None expires_at: Optional[int] = None expires_in: Optional[int] = None + audience: Optional[str] = None + prompt: Optional[str] = None class ServiceAccountCredential(BaseModelWithConfig): diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index 2e2a9a074f..323052365b 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -188,8 +188,15 @@ def generate_auth_uri( scope=" ".join(scopes), redirect_uri=auth_credential.oauth2.redirect_uri, ) + params = {"access_type": "offline"} + if auth_credential.oauth2.prompt: + params["prompt"] = auth_credential.oauth2.prompt + else: + params["prompt"] = "consent" + if auth_credential.oauth2.audience: + params["audience"] = auth_credential.oauth2.audience uri, state = client.create_authorization_url( - url=authorization_endpoint, access_type="offline", prompt="consent" + url=authorization_endpoint, **params ) exchanged_auth_credential = auth_credential.model_copy(deep=True) exchanged_auth_credential.oauth2.auth_uri = uri diff --git a/tests/unittests/auth/test_auth_handler.py b/tests/unittests/auth/test_auth_handler.py index 2a65f7795f..483765959d 100644 --- a/tests/unittests/auth/test_auth_handler.py +++ b/tests/unittests/auth/test_auth_handler.py @@ -61,7 +61,12 @@ def __init__( self.state = state def create_authorization_url(self, url, **kwargs): - return f"{url}?client_id={self.client_id}&scope={self.scope}", "mock_state" + params = f"client_id={self.client_id}&scope={self.scope}" + if kwargs.get("prompt"): + params += f"&prompt={kwargs.get('prompt')}" + if kwargs.get("audience"): + params += f"&audience={kwargs.get('audience')}" + return f"{url}?{params}", "mock_state" def fetch_token( self, @@ -225,8 +230,30 @@ def test_generate_auth_uri_oauth2(self, auth_config): "https://example.com/oauth2/authorize" ) assert "client_id=mock_client_id" in result.oauth2.auth_uri + assert "prompt=consent" in result.oauth2.auth_uri + assert "audience" not in result.oauth2.auth_uri assert result.oauth2.state == "mock_state" + @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session) + def test_generate_auth_uri_with_audience_and_prompt( + self, openid_auth_scheme, oauth2_credentials + ): + """Test generating an auth URI with audience and prompt.""" + oauth2_credentials.oauth2.audience = "test_audience" + oauth2_credentials.oauth2.prompt = "test_prompt" + exchanged = oauth2_credentials.model_copy(deep=True) + + config = AuthConfig( + auth_scheme=openid_auth_scheme, + raw_auth_credential=oauth2_credentials, + exchanged_auth_credential=exchanged, + ) + handler = AuthHandler(config) + result = handler.generate_auth_uri() + + assert "audience=test_audience" in result.oauth2.auth_uri + assert "prompt=test_prompt" in result.oauth2.auth_uri + @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session) def test_generate_auth_uri_openid( self, openid_auth_scheme, oauth2_credentials From 9a38de49928293d13b148d3de9342bf0caac8f90 Mon Sep 17 00:00:00 2001 From: Mark Scannell Date: Wed, 27 Aug 2025 09:25:53 +0100 Subject: [PATCH 2/4] Resolved a style issue --- src/google/adk/auth/auth_handler.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index 323052365b..dc31d78afa 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -189,10 +189,7 @@ def generate_auth_uri( redirect_uri=auth_credential.oauth2.redirect_uri, ) params = {"access_type": "offline"} - if auth_credential.oauth2.prompt: - params["prompt"] = auth_credential.oauth2.prompt - else: - params["prompt"] = "consent" + params["prompt"] = auth_credential.oauth2.prompt or "consent" if auth_credential.oauth2.audience: params["audience"] = auth_credential.oauth2.audience uri, state = client.create_authorization_url( From d15b764392b510a7d6f6d06fa12dc4ed98bf1569 Mon Sep 17 00:00:00 2001 From: Mark Scannell Date: Fri, 29 Aug 2025 19:42:21 +0100 Subject: [PATCH 3/4] Removed prompt configuration. --- tests/unittests/auth/test_auth_handler.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/unittests/auth/test_auth_handler.py b/tests/unittests/auth/test_auth_handler.py index 483765959d..0a2a2f7802 100644 --- a/tests/unittests/auth/test_auth_handler.py +++ b/tests/unittests/auth/test_auth_handler.py @@ -62,8 +62,6 @@ def __init__( def create_authorization_url(self, url, **kwargs): params = f"client_id={self.client_id}&scope={self.scope}" - if kwargs.get("prompt"): - params += f"&prompt={kwargs.get('prompt')}" if kwargs.get("audience"): params += f"&audience={kwargs.get('audience')}" return f"{url}?{params}", "mock_state" @@ -230,7 +228,6 @@ def test_generate_auth_uri_oauth2(self, auth_config): "https://example.com/oauth2/authorize" ) assert "client_id=mock_client_id" in result.oauth2.auth_uri - assert "prompt=consent" in result.oauth2.auth_uri assert "audience" not in result.oauth2.auth_uri assert result.oauth2.state == "mock_state" @@ -240,7 +237,6 @@ def test_generate_auth_uri_with_audience_and_prompt( ): """Test generating an auth URI with audience and prompt.""" oauth2_credentials.oauth2.audience = "test_audience" - oauth2_credentials.oauth2.prompt = "test_prompt" exchanged = oauth2_credentials.model_copy(deep=True) config = AuthConfig( @@ -252,7 +248,6 @@ def test_generate_auth_uri_with_audience_and_prompt( result = handler.generate_auth_uri() assert "audience=test_audience" in result.oauth2.auth_uri - assert "prompt=test_prompt" in result.oauth2.auth_uri @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session) def test_generate_auth_uri_openid( From af95b2244610711e43c025cd9795f061348759af Mon Sep 17 00:00:00 2001 From: Mark Scannell Date: Mon, 1 Sep 2025 08:28:27 +0100 Subject: [PATCH 4/4] Removed prompt --- src/google/adk/auth/auth_credential.py | 1 - src/google/adk/auth/auth_handler.py | 7 +++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/google/adk/auth/auth_credential.py b/src/google/adk/auth/auth_credential.py index 4e4a5fef42..bc91d48f79 100644 --- a/src/google/adk/auth/auth_credential.py +++ b/src/google/adk/auth/auth_credential.py @@ -80,7 +80,6 @@ class OAuth2Auth(BaseModelWithConfig): expires_at: Optional[int] = None expires_in: Optional[int] = None audience: Optional[str] = None - prompt: Optional[str] = None class ServiceAccountCredential(BaseModelWithConfig): diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index dc31d78afa..7a51a71e29 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -188,13 +188,16 @@ def generate_auth_uri( scope=" ".join(scopes), redirect_uri=auth_credential.oauth2.redirect_uri, ) - params = {"access_type": "offline"} - params["prompt"] = auth_credential.oauth2.prompt or "consent" + params = { + "access_type": "offline", + "prompt": "consent", + } if auth_credential.oauth2.audience: params["audience"] = auth_credential.oauth2.audience uri, state = client.create_authorization_url( url=authorization_endpoint, **params ) + exchanged_auth_credential = auth_credential.model_copy(deep=True) exchanged_auth_credential.oauth2.auth_uri = uri exchanged_auth_credential.oauth2.state = state