diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index fa273fd..074fdcb 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -22,6 +22,7 @@ import sys from typing import Any from typing import Optional from typing import TextIO +from typing import Union import anyio from pydantic import BaseModel @@ -65,6 +66,7 @@ class StreamableHTTPServerParams(BaseModel): See MCP SSE Client documentation for more details. https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py """ + url: str headers: dict[str, Any] | None = None timeout: float = 5 @@ -139,7 +141,9 @@ class MCPSessionManager: def __init__( self, - connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams, + connection_params: Union[ + StdioServerParameters, SseServerParams, StreamableHTTPServerParams + ], errlog: TextIO = sys.stderr, ): """Initializes the MCP session manager. @@ -184,7 +188,9 @@ class MCPSessionManager: url=self._connection_params.url, headers=self._connection_params.headers, timeout=timedelta(seconds=self._connection_params.timeout), - sse_read_timeout=timedelta(seconds=self._connection_params.sse_read_timeout), + sse_read_timeout=timedelta( + seconds=self._connection_params.sse_read_timeout + ), terminate_on_close=self._connection_params.terminate_on_close, ) else: diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 0ad6d04..40fadfb 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import logging from typing import Optional diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index 1bc6b3c..f4837ec 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import logging import sys from typing import List @@ -83,7 +85,9 @@ class MCPToolset(BaseToolset): def __init__( self, *, - connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams, + connection_params: ( + StdioServerParameters | SseServerParams | StreamableHTTPServerParams + ), tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, errlog: TextIO = sys.stderr, ):