Skip to content

Commit edda922

Browse files
mescannecopybara-github
authored andcommitted
feat: add audience and prompt as configurable for OAuth flows
Merge #2738 Some OAuth servers require audience such as [Jira](https://developer.atlassian.com/cloud/confluence/oauth-2-3lo-apps/). This change allows the configuration of audience and prompt (if it needs to be changed) and adds some tests. This is for issue [2755](#2755). Resolves #2755 Happy to provide changes/updates if needed. COPYBARA_INTEGRATE_REVIEW=#2738 from mescanne:oauth-audience-prompt 87ce110 PiperOrigin-RevId: 802850034
1 parent a1679da commit edda922

File tree

3 files changed

+32
-2
lines changed

3 files changed

+32
-2
lines changed

src/google/adk/auth/auth_credential.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ class OAuth2Auth(BaseModelWithConfig):
7979
refresh_token: Optional[str] = None
8080
expires_at: Optional[int] = None
8181
expires_in: Optional[int] = None
82+
audience: Optional[str] = None
8283

8384

8485
class ServiceAccountCredential(BaseModelWithConfig):

src/google/adk/auth/auth_handler.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,9 +188,16 @@ def generate_auth_uri(
188188
scope=" ".join(scopes),
189189
redirect_uri=auth_credential.oauth2.redirect_uri,
190190
)
191+
params = {
192+
"access_type": "offline",
193+
"prompt": "consent",
194+
}
195+
if auth_credential.oauth2.audience:
196+
params["audience"] = auth_credential.oauth2.audience
191197
uri, state = client.create_authorization_url(
192-
url=authorization_endpoint, access_type="offline", prompt="consent"
198+
url=authorization_endpoint, **params
193199
)
200+
194201
exchanged_auth_credential = auth_credential.model_copy(deep=True)
195202
exchanged_auth_credential.oauth2.auth_uri = uri
196203
exchanged_auth_credential.oauth2.state = state

tests/unittests/auth/test_auth_handler.py

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,10 @@ def __init__(
6161
self.state = state
6262

6363
def create_authorization_url(self, url, **kwargs):
64-
return f"{url}?client_id={self.client_id}&scope={self.scope}", "mock_state"
64+
params = f"client_id={self.client_id}&scope={self.scope}"
65+
if kwargs.get("audience"):
66+
params += f"&audience={kwargs.get('audience')}"
67+
return f"{url}?{params}", "mock_state"
6568

6669
def fetch_token(
6770
self,
@@ -225,8 +228,27 @@ def test_generate_auth_uri_oauth2(self, auth_config):
225228
"https://example.com/oauth2/authorize"
226229
)
227230
assert "client_id=mock_client_id" in result.oauth2.auth_uri
231+
assert "audience" not in result.oauth2.auth_uri
228232
assert result.oauth2.state == "mock_state"
229233

234+
@patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
235+
def test_generate_auth_uri_with_audience_and_prompt(
236+
self, openid_auth_scheme, oauth2_credentials
237+
):
238+
"""Test generating an auth URI with audience and prompt."""
239+
oauth2_credentials.oauth2.audience = "test_audience"
240+
exchanged = oauth2_credentials.model_copy(deep=True)
241+
242+
config = AuthConfig(
243+
auth_scheme=openid_auth_scheme,
244+
raw_auth_credential=oauth2_credentials,
245+
exchanged_auth_credential=exchanged,
246+
)
247+
handler = AuthHandler(config)
248+
result = handler.generate_auth_uri()
249+
250+
assert "audience=test_audience" in result.oauth2.auth_uri
251+
230252
@patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
231253
def test_generate_auth_uri_openid(
232254
self, openid_auth_scheme, oauth2_credentials

0 commit comments

Comments
 (0)