Skip to content

feat: support execution timeout for MCP tools in toolkit#1244

Open
owenliang wants to merge 2 commits intoagentscope-ai:mainfrom
owenliang:main
Open

feat: support execution timeout for MCP tools in toolkit#1244
owenliang wants to merge 2 commits intoagentscope-ai:mainfrom
owenliang:main

Conversation

@owenliang
Copy link
Copy Markdown

PR Title

feat(toolkit): add execution_timeout parameter to register_mcp_tools()

AgentScope Version

0.1.5

Description

Background: MCP tool function execution may block indefinitely due to network or server issues, requiring timeout control.

Purpose: Add execution_timeout parameter to register_mcp_tools() method to allow users to set execution timeout for MCP tool functions.

Changes:

  • _toolkit.py line 841: Add execution_timeout: float | None = None parameter to method signature
  • _toolkit.py lines 875-877: Add parameter documentation
  • _toolkit.py line 940: Pass execution_timeout to get_callable_function()

How to test:

from agentscope.tool import Toolkit
from agentscope.mcp import MCPClient

toolkit = Toolkit()
# Set 30 seconds timeout
toolkit.register_mcp_tools(mcp_client, execution_timeout=30.0)

Checklist

  • Code has been formatted with pre-commit run --all-files command
  • All tests are passing
  • Docstrings are in Google style
  • Related documentation has been updated (e.g. links, examples, etc.)
  • Code is ready for review

Add execution_timeout parameter to register_mcp_tools() to control
MCP tool function execution timeout.
@gemini-code-assist
Copy link
Copy Markdown
Contributor

Summary of Changes

Hello @owenliang, I'm Gemini Code Assist1! 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 robustness of MCP tool function execution by integrating a configurable timeout mechanism. It addresses potential issues where tool functions might block indefinitely due to external factors like network problems, providing users with better control over their application's responsiveness and stability.

Highlights

  • New Parameter Added: Introduced an execution_timeout parameter to the register_mcp_client function, allowing users to specify a timeout for MCP tool function execution.
  • Documentation Update: Updated the docstring for register_mcp_client to include detailed information about the new execution_timeout parameter, its type, default value, and purpose.
  • Timeout Propagation: Modified the internal logic to pass the execution_timeout value to the get_callable_function call, ensuring the timeout is applied during tool function retrieval and execution.
Changelog
  • src/agentscope/tool/_toolkit.py
    • Added execution_timeout: float | None = None parameter to the register_mcp_client function signature.
    • Updated the docstring for register_mcp_client to describe the new execution_timeout parameter.
    • Passed the execution_timeout parameter to the mcp_client.get_callable_function call.
Activity
  • The pull request author, owenliang, created this PR to add execution timeout support for MCP tools.
  • The PR description includes a background, purpose, detailed changes, and testing instructions.
  • A checklist is provided, indicating the author's progress on formatting, testing, documentation, and readiness for review.
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.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a timeout mechanism for MCP tool execution by adding an execution_timeout parameter to register_mcp_client. This is a valuable addition for handling potential blocking issues.

My review has identified a couple of points that need attention:

  • The docstring for the new parameter doesn't fully align with the project's style guide for optional arguments.
  • A more significant issue is the absence of unit tests for this new feature, which is a mandatory requirement according to the repository's contribution guidelines.

Please address these points to ensure the new feature is robust and well-documented.

- Update HttpStatefulClient to use httpx.AsyncClient with streamable_http_client

- Update HttpStatelessClient to use httpx.AsyncClient with streamable_http_client

- Pass client_kwargs directly to httpx.AsyncClient for flexible configuration

- Fix execution_timeout docstring from defaults to to optional

- Add execution_timeout test cases for both SSE and Streamable HTTP clients

- Add slow_tool helper for timeout testing with time-based assertions

- Disable slow_tool in existing tests to avoid schema mismatch
@owenliang
Copy link
Copy Markdown
Author


Summary

This PR introduces two major changes to the MCP client implementation:

1. feat: Add execution_timeout parameter to register_mcp_client

Following the previous PR that added execution_timeout support, this change includes comprehensive test coverage:

  • Added test_execution_timeout_with_register_mcp_client for both SSE and Streamable HTTP transports
  • Tests verify that timeout occurs around 1 second (with 0.5s tolerance) instead of waiting for the full 3-second tool execution
  • Added slow_tool helper function to simulate long-running operations
  • Updated existing tests to disable slow_tool to avoid schema assertion failures

2. refactor: Migrate Streamable HTTP client to MCP v2 API

Breaking Change: Updated both HttpStatefulClient and HttpStatelessClient to use the new MCP v2 API:

  • Replaced deprecated streamablehttp_client with streamable_http_client
  • Now requires httpx.AsyncClient for HTTP configuration
  • client_kwargs are now passed directly to httpx.AsyncClient initialization, enabling flexible configuration (auth, cookies, verify, etc.)
  • Fixed execution_timeout docstring from defaults to to optional

Files Changed

File Changes
src/agentscope/mcp/_http_stateful_client.py Migrated to v2 API
src/agentscope/mcp/_http_stateless_client.py Migrated to v2 API + client_kwargs passthrough
src/agentscope/tool/_toolkit.py Docstring fix
tests/mcp_sse_client_test.py Added timeout tests
tests/mcp_streamable_http_client_test.py Added timeout tests

Testing

  • All 6 MCP-related tests pass
  • 241 total tests pass (80 failures are due to missing optional dependencies, unrelated to this change)
  • No deprecation warnings from MCP SDK

@owenliang
Copy link
Copy Markdown
Author

Additionally, a bug causing the MCP streamable HTTP client to hang was discovered. The issue has already been fixed, but a new version of the MCP Python package has not yet been released. We need to update the dependency promptly once it's available; otherwise, if the MCP server encounters an error during an MCP tool call, the call will hang indefinitely: modelcontextprotocol/python-sdk#1577

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant