adk-python icon indicating copy to clipboard operation
adk-python copied to clipboard

feat(toolset): add generate_preprocessing_events method to BaseToolset

Open loveyana opened this issue 2 months ago • 3 comments

Add new generate_preprocessing_events method to BaseToolset that allows toolsets to generate events (such as authentication requests) during the preprocessing phase before tool discovery occurs.

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

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

  • Closes: #3331

2. Problem Description:

The current ADK Python implementation has critical defects in the MCP (Model Context Protocol) Toolset OAuth2 authentication flow. MCP Toolset tool discovery occurs during the preprocessing phase using ReadonlyContext, which lacks the request_credential() method necessary for OAuth2 user interaction. This creates a fundamental architectural problem where authentication is required before tools can be discovered, but the authentication mechanism is not available at that stage.

Core Issues:

  • MCP session creation requires authentication before tool discovery
  • ReadonlyContext lacks authentication capability during preprocessing phase
  • Session initialization timing mismatch between authentication and tool discovery
  • MCP Toolsets cannot trigger OAuth2 flows when needed

3. Solution:

Add a new generate_preprocessing_events() method to BaseToolset that allows toolsets to generate authentication events during the preprocessing phase, before tool discovery occurs. This method:

  • Has access to full ToolContext with authentication capabilities
  • Executes at the perfect timing (after request processors, before tool discovery)
  • Maintains full backward compatibility with existing toolsets
  • Provides foundation for MCP Toolset OAuth2 authentication flows

Testing Plan

Please describe the tests that you ran to verify your changes. This is required for all PRs that are not small documentation or typo fixes.

Unit Tests:

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

Added Tests:

  • test_preprocess_calls_toolset_generate_preprocessing_events: Verifies that _preprocess_async calls the new method on toolsets and processes generated events
  • test_preprocess_calls_both_generate_events_and_process_request: Verifies correct call order (generate_preprocessing_events → process_llm_request)

Test Coverage:

  • Method invocation verification
  • Event generation and propagation
  • Integration with existing preprocessing flow
  • Backward compatibility with existing toolsets

Manual End-to-End (E2E) Tests:

Testing the Infrastructure:

  1. Verify that existing toolsets continue to work unchanged (backward compatibility)
  2. Verify that custom toolsets can override generate_preprocessing_events to generate authentication events
  3. Verify that generated events are properly yielded through the preprocessing flow

Future Testing (after MCP Toolset implementation):

  1. Create MCP Toolset with OAuth2 authentication
  2. Verify authentication flow triggers before tool discovery
  3. Verify successful MCP session creation after authentication

Checklist

  • [x] I have read the CONTRIBUTING.md document.
  • [x] I have performed a self-review of my own code.
  • [x] I have commented my code, particularly in hard-to-understand areas.
  • [x] I have added tests that prove my fix is effective or that my feature works.
  • [x] New and existing unit tests pass locally with my changes.
  • [x] I have manually tested my changes end-to-end.
  • [x] Any dependent changes have been merged and published in downstream modules.

Additional context

Architecture Impact: This change provides the foundational infrastructure needed to solve MCP Toolset OAuth2 authentication issues. The implementation follows ADK's existing patterns:

  • Uses the same ToolContext and LlmRequest parameters as process_llm_request
  • Follows the AsyncGenerator[Event, None] pattern used throughout ADK
  • Maintains backward compatibility with default implementation
  • Integrates seamlessly with existing preprocessing flow

Implementation Details:

  • BaseToolset: Added generate_preprocessing_events method with default no-op implementation
  • base_llm_flow.py: Added call to new method before process_llm_request in toolset processing
  • Tests: Added comprehensive unit tests covering method invocation and event generation

Future Work: This PR provides the infrastructure. The next step will be implementing the actual OAuth2 authentication logic in MCP Toolset using this new capability.

Breaking Changes: None - this is a purely additive change with full backward compatibility.

loveyana avatar Oct 30 '25 11:10 loveyana

Summary of Changes

Hello @loveyana, I'm Gemini Code Assist[^1]! 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 introduces a foundational change to the ADK Python framework by adding a generate_preprocessing_events method to the BaseToolset. This new capability allows toolsets to proactively generate and yield events, such as authentication requests, during the initial preprocessing stage of an LLM flow. This directly addresses a long-standing limitation where complex authentication flows, particularly for MCP Toolsets using OAuth2, could not be initiated at the necessary point in the lifecycle due to context restrictions. By providing access to the full ToolContext at the right time, this PR unblocks advanced authentication scenarios and maintains compatibility with existing implementations.

Highlights

  • New Preprocessing Event Generation: Introduced a new generate_preprocessing_events method to BaseToolset which allows toolsets to emit events (like authentication requests) during the preprocessing phase, before tool discovery.
  • Addressing Authentication Gaps: This change resolves a critical architectural issue where ReadonlyContext in the preprocessing phase lacked the necessary request_credential() method for OAuth2 authentication, preventing proper MCP Toolset session creation.
  • Full Context Access: The new method provides access to the full ToolContext with authentication capabilities, executing at the ideal timing (after request processors, before tool discovery).
  • Backward Compatibility: The implementation is purely additive, ensuring full backward compatibility with existing toolsets through a default no-op implementation.
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 :thumbsup: and :thumbsdown: 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.

[^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[bot] avatar Oct 30 '25 11:10 gemini-code-assist[bot]

Hi @loveyana, Thank you for your contribution! We appreciate you taking the time to submit this pull request. Your PR has been received by the team and is currently under review. We will provide feedback as soon as we have an update to share.

ryanaiagent avatar Oct 31 '25 08:10 ryanaiagent

Hi @loveyana, Thank you for your contribution! We appreciate you taking the time to submit this pull request. Your PR has been received by the team and is currently under review. We will provide feedback as soon as we have an update to share.

@ryanaiagent Friendly ping when you have a moment! Happy to help with anything needed to move this forward. Thanks! 😊

loveyana avatar Nov 21 '25 03:11 loveyana

@ryanaiagent Rohit,would you pls help review this PR, and it really blocking the OAuth2 flows of MCP OAuth in our agent use case. I guess it is also helpful others who is using OAuth2 flows of MCP.

cuericlee avatar Dec 03 '25 09:12 cuericlee