Fix "aiq info" problems after adding AsyncExitStack

Signed-off-by: Anuradha Karuppiah <[email protected]>

Author: AnuradhaKaruppiah

src/aiq/cli/commands/info/list_mcp.py

@@ -60,19 +60,23 @@ async def list_tools_and_schemas(command, url, tool_name=None, client_type='sse'
         args = []
     from aiq.tool.mcp.mcp_client import MCPSSEClient
     from aiq.tool.mcp.mcp_client import MCPStdioClient
+    from aiq.tool.mcp.mcp_client import MCPStreamableHTTPClient
 
     try:
         if client_type == 'stdio':
             client = MCPStdioClient(command=command, args=args, env=env)
-        else:
+        elif client_type == 'streamable-http':
+            client = MCPStreamableHTTPClient(url=url)
+        else:  # sse
             client = MCPSSEClient(url=url)
 
-        if tool_name:
-            tool = await client.get_tool(tool_name)
-            return [format_tool(tool)]
-        else:
-            tools = await client.get_tools()
-            return [format_tool(tool) for tool in tools.values()]
+        async with client:
+            if tool_name:
+                tool = await client.get_tool(tool_name)
+                return [format_tool(tool)]
+            else:
+                tools = await client.get_tools()
+                return [format_tool(tool) for tool in tools.values()]
     except Exception as e:
         click.echo(f"[ERROR] Failed to fetch tools via MCP client: {e}", err=True)
         return []
@@ -85,6 +89,7 @@ async def list_tools_direct(command, url, tool_name=None, client_type='sse', arg
     from mcp.client.sse import sse_client
     from mcp.client.stdio import StdioServerParameters
     from mcp.client.stdio import stdio_client
+    from mcp.client.streamable_http import streamablehttp_client
 
     try:
         if client_type == 'stdio':
@@ -93,7 +98,13 @@ def get_stdio_client():
                 return stdio_client(server=StdioServerParameters(command=command, args=args, env=env))
 
             client = get_stdio_client
-        else:
+        elif client_type == 'streamable-http':
+
+            def get_streamable_http_client():
+                return streamablehttp_client(url=url)
+
+            client = get_streamable_http_client
+        else:  # sse
 
             def get_sse_client():
                 return sse_client(url=url)
@@ -125,9 +136,9 @@ def get_sse_client():
 @click.option('--url',
               default='http://localhost:9901/sse',
               show_default=True,
-              help='For SSE: MCP server URL (e.g. http://localhost:8080/sse)')
+              help='For SSE/StreamableHTTP: MCP server URL (e.g. http://localhost:8080/sse)')
 @click.option('--client-type',
-              type=click.Choice(['sse', 'stdio']),
+              type=click.Choice(['sse', 'stdio', 'streamable-http']),
               default='sse',
               show_default=True,
               help='Type of client to use')
@@ -151,12 +162,14 @@ def list_mcp(ctx, direct, url, client_type, command, args, env, tool, detail, js
             click.echo("[ERROR] --command is required when using stdio client type", err=True)
             return
 
-    if client_type == 'sse':
+    if client_type in ['sse', 'streamable-http']:
         if not url:
-            click.echo("[ERROR] --url is required when using sse client type", err=True)
+            click.echo("[ERROR] --url is required when using sse or streamable-http client type", err=True)
             return
         if command or args or env:
-            click.echo("[ERROR] --command, --args, and --env are not allowed when using sse client type", err=True)
+            click.echo(
+                "[ERROR] --command, --args, and --env are not allowed when using sse or streamable-http client type",
+                err=True)
             return
 
     stdio_args = args.split() if args else []

src/aiq/tool/mcp/mcp_tool.py

@@ -14,8 +14,8 @@
 # limitations under the License.
 
 import logging
+from typing import Literal
 
-from pydantic import BaseModel
 from pydantic import Field
 from pydantic import HttpUrl
 
@@ -35,7 +35,8 @@ class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
     # Add your custom configuration parameters here
     url: HttpUrl | None = Field(default=None, description="The URL of the MCP server (for SSE mode)")
     mcp_tool_name: str = Field(description="The name of the tool served by the MCP Server that you want to use")
-    client_type: str = Field(default="sse", description="The type of client to use ('sse' or 'stdio')")
+    client_type: Literal["sse", "stdio", "streamable-http"] = Field(default="sse",
+                                                                    description="The type of transport to use")
     command: str | None = Field(default=None,
                                 description="The command to run for stdio mode (e.g. 'docker' or 'python')")
     args: list[str] | None = Field(default=None, description="Additional arguments for the stdio command")
@@ -52,15 +53,15 @@ class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
         """)
 
     def model_post_init(self, __context):
-        """Validate that stdio and SSE properties are mutually exclusive."""
+        """Validate that stdio and SSE/Streamable HTTP properties are mutually exclusive."""
         super().model_post_init(__context)
 
         if self.client_type == 'stdio':
             if self.url is not None:
                 raise ValueError("url should not be set when using stdio client type")
             if not self.command:
                 raise ValueError("command is required when using stdio client type")
-        else:
+        elif self.client_type == 'streamable-http' or self.client_type == 'sse':
             if self.command is not None or self.args is not None or self.env is not None:
                 raise ValueError("command, args, and env should not be set when using sse client type")
             if not self.url:
@@ -83,17 +84,20 @@ async def mcp_tool(config: MCPToolConfig, builder: Builder):  # pylint: disable=
         if not config.command:
             raise ValueError("command is required when using stdio client type")
 
+        source = f"{config.command} {' '.join(config.args) if config.args else ''}"
         client = MCPStdioClient(command=config.command, args=config.args, env=config.env)
     elif config.client_type == 'streamable-http':
         if not config.url:
             raise ValueError("url is required when using streamable-http client type")
 
-        client = MCPStreamableHTTPClient(url=str(config.url))
+        source = str(config.url)
+        client = MCPStreamableHTTPClient(url=source)
     elif config.client_type == 'sse':
         if not config.url:
             raise ValueError("url is required when using sse client type")
 
-        client = MCPSSEClient(url=str(config.url))
+        source = str(config.url)
+        client = MCPSSEClient(url=source)
     else:
         raise ValueError(f"Invalid client type: {config.client_type}")
 
@@ -103,19 +107,12 @@ async def mcp_tool(config: MCPToolConfig, builder: Builder):  # pylint: disable=
         if config.description:
             tool.set_description(description=config.description)
 
-        if config.client_type == "sse" or config.client_type == "streamable-http":
-            source = config.url
-        elif config.client_type == "stdio":
-            source = f"{config.command} {' '.join(config.args) if config.args else ''}"
-        else:
-            raise ValueError(f"Invalid client type: {config.client_type}")
-
         logger.info("Configured to use tool: %s from MCP server at %s", tool.name, source)
 
         def _convert_from_str(input_str: str) -> tool.input_schema:
             return tool.input_schema.model_validate_json(input_str)
 
-        async def _response_fn(tool_input: tool.input_schema) -> str:
+        async def _response_fn(tool_input: tool.input_schema | None = None, **kwargs) -> str:
             # Run the tool, catching any errors and sending to agent for correction
             try:
                 if tool_input: