feat: add audience and prompt as configurable for OAuth flows

Merge https://github.com/google/adk-python/pull/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](https://github.com/google/adk-python/issues/2755).

Resolves #2755

Happy to provide changes/updates if needed.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2738 from mescanne:oauth-audience-prompt 87ce1100792d9156ada2a004bcfaf2fe5fc69602
PiperOrigin-RevId: 802850034
This commit is contained in:
Mark Scannell
2025-09-03 21:53:54 -07:00
committed by Copybara-Service
parent a1679dae3f
commit edda922791
3 changed files with 32 additions and 2 deletions
+1
View File
@@ -79,6 +79,7 @@ class OAuth2Auth(BaseModelWithConfig):
refresh_token: Optional[str] = None
expires_at: Optional[int] = None
expires_in: Optional[int] = None
audience: Optional[str] = None
class ServiceAccountCredential(BaseModelWithConfig):
+8 -1
View File
@@ -188,9 +188,16 @@ class AuthHandler:
scope=" ".join(scopes),
redirect_uri=auth_credential.oauth2.redirect_uri,
)
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, 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
exchanged_auth_credential.oauth2.state = state
+23 -1
View File
@@ -61,7 +61,10 @@ class MockOAuth2Session:
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("audience"):
params += f"&audience={kwargs.get('audience')}"
return f"{url}?{params}", "mock_state"
def fetch_token(
self,
@@ -225,8 +228,27 @@ class TestGenerateAuthUri:
"https://example.com/oauth2/authorize"
)
assert "client_id=mock_client_id" 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"
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
@patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
def test_generate_auth_uri_openid(
self, openid_auth_scheme, oauth2_credentials