feat(auth): Add native support for id_token in OAuth2 credentials

[anubhav756]

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

2. Or, if no issue exists, describe the change:

Problem

When performing authentication flows via OAUTH2 or OPEN_ID_CONNECT, the native OAuth2Token response from identity providers (like Google OAuth) often includes an id_token alongside the access_token and refresh_token. However, the ADK's update_credential_with_tokens utility explicitly drops the id_token, preventing agents and tools from verifying user identity or extracting OIDC claims securely. Furthermore, the OAuth2Auth model does not have a designated field for id_token.

Solution

1. Added an id_token: Optional[str] = None field to the OAuth2Auth pydantic model in auth_credential.py.
2. Updated update_credential_with_tokens in oauth2_credential_util.py to correctly extract and map tokens.get("id_token") into the OAuth2Auth credential object.
3. Updated the relevant unit tests to ensure id_token is asserted and preserved during credential updates.

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

Summary of passed pytest results:

$ pytest tests/unittests/auth/test_oauth2_credential_util.py
======================= test session starts =======================
platform darwin -- Python 3.11.9, pytest-9.0.1, pluggy-1.6.0
collected 9 items

tests/unittests/auth/test_oauth2_credential_util.py ......... [100%]
======================== 9 passed in 0.05s ========================

[gemini-code-assist]

Summary of Changes

Hello @anubhav756, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the authentication system by introducing native support for 'id_token' within OAuth2 credentials. Previously, 'id_token' values, crucial for user identity verification in OpenID Connect flows, were dropped. The changes ensure that these tokens are now correctly captured and stored, enabling downstream components to securely verify user identity and extract OIDC claims.

Highlights

• OAuth2Auth Model Enhancement: Added an optional 'id_token' field to the OAuth2Auth Pydantic model, allowing it to store identity tokens received during authentication flows.
• ID Token Preservation: Modified the 'update_credential_with_tokens' utility to correctly extract and store the 'id_token' from OAuth2 token responses into the credential object, preventing its previous omission.
• Unit Test Coverage: Updated existing unit tests to ensure that the 'id_token' is properly handled, preserved, and asserted during credential updates, confirming the new functionality.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/google/adk/auth/auth_credential.py
  • Added 'id_token: Optional[str] = None' to the 'OAuth2Auth' class.
• src/google/adk/auth/oauth2_credential_util.py
  • Assigned 'tokens.get("id_token")' to 'auth_credential.oauth2.id_token' in 'update_credential_with_tokens'.
• tests/unittests/auth/test_oauth2_credential_util.py
  • Added 'id_token': 'new_id_token' to the 'OAuth2Token' dictionary in 'test_update_credential_with_tokens'.
  • Added an assertion 'assert credential.oauth2.id_token == "new_id_token"' in 'test_update_credential_with_tokens'.

Activity

• The author, anubhav756, initiated this pull request to address the omission of 'id_token' in OAuth2 credential handling.
• The pull request includes a clear problem description and a detailed solution plan.
• Unit tests have been updated and confirmed to pass locally, demonstrating the correctness of the changes.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for id_token within the OAuth2 authentication flow, which is a valuable addition for OpenID Connect scenarios. The changes are implemented cleanly and consistently. The OAuth2Auth model is correctly updated to include the optional id_token field, the update_credential_with_tokens utility is properly modified to handle the new token, and the unit tests are expanded to verify this new behavior. The implementation is straightforward and correct.