MCP enhancements: improves server config and adds support for all transport types (stdio, streamable-http)

docs/source/workflows/mcp/mcp-client.md

@@ -32,8 +32,12 @@ class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
     Function which connects to a Model Context Protocol (MCP) server and wraps the selected tool as an AIQ function.
     """
     # Add your custom configuration parameters here
-    url: HttpUrl = Field(description="The URL of the MCP server")
+    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')")
+    command: str | None = Field(default=None, description="The command to run for stdio mode (e.g. 'mcp-server')")
+    args: list[str] | None = Field(default=None, description="Additional arguments for the stdio command")
+    env: dict[str, str] | None = Field(default=None, description="Environment variables to set for the stdio process")
     description: str | None = Field(
         default=None,
         description="""
@@ -42,9 +46,14 @@ class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
         """
     )
 ```
-In addition to the URL of the server, the configuration also takes as a parameter the name of the MCP tool you want to use as an AIQ toolkit function. This is required because MCP servers can serve multiple tools, and for this wrapper we want to maintain a one-to-one relationship between AIQ toolkit functions and MCP tools. This means that if you want to include multiple tools from an MCP server you will configure multiple `mcp_tool_wrappers`.
 
-For example:
+The configuration supports two client types:
+
+1. **SSE (Server-Sent Events)**: The default client type that connects to an MCP server over HTTP. This is the primary and recommended mode for most use cases.
+2. **STDIO**: A mode that launches the MCP server as a `subprocess` and communicates with it through standard input/output.
+
+### SSE Mode Configuration
+For SSE mode, you only need to specify the server URL and the tool name:
 
 ```yaml
 functions:
@@ -56,10 +65,28 @@ functions:
     _type: mcp_tool_wrapper
     url: "http://localhost:8080/sse"
     mcp_tool_name: tool_b
-  mcp_tool_c:
+```
+
+### STDIO Mode Configuration
+For STDIO mode, you need to specify the command to run and any additional arguments or environment variables:
+
+```yaml
+functions:
+  github_mcp:
     _type: mcp_tool_wrapper
-    url: "http://localhost:8080/sse"
-    mcp_tool_name: tool_c
+    client_type: stdio
+    command: "docker"
+    args: [
+      "run",
+      "-i",
+      "--rm",
+      "-e",
+      "GITHUB_PERSONAL_ACCESS_TOKEN",
+      "ghcr.io/github/github-mcp-server"
+    ]
+    env:
+      GITHUB_PERSONAL_ACCESS_TOKEN: "${input:github_token}"
+    mcp_tool_name: "github_tool"
 ```
 
 The final configuration parameter (the `description`) is optional, and should only be used if the description provided by the MCP server is not sufficient, or if there is no description provided by the server.
@@ -101,6 +128,29 @@ aiq run --config_file examples/simple_calculator/configs/config-mcp-date.yml --i
 ```
 This will use the `mcp_time_tool` function to get the current hour of the day from the MCP server.
 
+### Using STDIO Mode
+Alternatively, you can run the same example using stdio mode with the `config-mcp-date-stdio.yml` configuration:
+
+```yaml
+functions:
+  mcp_time_tool:
+    _type: mcp_tool_wrapper
+    client_type: stdio
+    command: "python"
+    args: ["-m", "mcp_server_time", "--local-timezone=America/Los_Angeles"]
+    mcp_tool_name: get_current_time
+    description: "Returns the current date and time from the MCP server"
+```
+
+This configuration launches the MCP server directly as a `subprocess` instead of connecting to a running server. Run it with:
+```bash
+aiq run --config_file examples/simple_calculator/configs/config-mcp-date-stdio.yml --input "Is the product of 2 * 4 greater than the current hour of the day?"
+```
+Ensure that MCP server time package is installed in your environment before running the workflow.
+```bash
+uv pip install mcp-server-time
+```
+
 ## Displaying MCP Tools
 The `aiq info mcp` command can be used to list the tools served by an MCP server.
 ```bash

examples/simple_calculator/src/aiq_simple_calculator/configs/config-mcp-date-stdio.yml

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This config file shows how to use the MCP server to get the current date and time.
+# Here the workflow acts as a MCP client and connects to the MCP server running
+# on the specified URL (defaults to `http://localhost:8080/sse`).
+
+general:
+  use_uvloop: true
+
+functions:
+  mcp_time:
+    _type: mcp_client
+    server:
+      transport: stdio
+      command: "python"
+      args: ["-m", "mcp_server_time", "--local-timezone=America/Los_Angeles"]
+    # load 1 of 2 tools and override tool attributes
+    tool_filter:
+      get_current_time:
+        alias: get_current_time_mcp_tool
+        description: "Returns the current date and time from the MCP server"
+
+  mcp_math:
+    _type: mcp_client
+    server:
+      transport: sse
+      url: "http://localhost:9901/sse"
+    # load all tools and use the names and descriptions from the server
+
+llms:
+  nim_llm:
+    _type: nim
+    model_name: meta/llama-3.1-70b-instruct
+    temperature: 0.0
+    max_tokens: 1024
+  openai_llm:
+    _type: openai
+    model_name: gpt-3.5-turbo
+    max_tokens: 2000
+
+workflow:
+  _type: react_agent
+  # ideally we should be able to fetch all dynamic tools without enumerating them here
+  tool_names: ["get_current_time_mcp_tool", "calculator_multiply", "calculator_divide", "calculator_subtract", "calculator_inequality"]
+  llm_name: nim_llm
+  verbose: true
+  retry_parsing_errors: true
+  max_retries: 3

pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   "httpx~=0.27",
   "jinja2~=3.1",
   "jsonpath-ng~=1.7",
-  "mcp>=1.0.0",
+  "mcp~=1.8",
   "networkx~=3.4",
   "numpy~=1.26",
   "openinference-semantic-conventions~=0.1.14",

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

@@ -19,8 +19,6 @@
 import anyio
 import click
 
-from aiq.tool.mcp.mcp_client import MCPBuilder
-
 # Suppress verbose logs from mcp.client.sse and httpx
 logging.getLogger("mcp.client.sse").setLevel(logging.WARNING)
 logging.getLogger("httpx").setLevel(logging.WARNING)
@@ -31,12 +29,14 @@ def format_tool(tool):
     description = getattr(tool, 'description', '')
     input_schema = getattr(tool, 'input_schema', None) or getattr(tool, 'inputSchema', None)
 
-    schema_str = None
-    if input_schema:
-        if hasattr(input_schema, "schema_json"):
-            schema_str = input_schema.schema_json(indent=2)
-        else:
-            schema_str = str(input_schema)
+    # Normalize schema to JSON string
+    if hasattr(input_schema, "schema_json"):
+        schema_str = input_schema.schema_json(indent=2)
+    elif isinstance(input_schema, dict):
+        schema_str = json.dumps(input_schema, indent=2)
+    else:
+        # Final fallback: attempt to dump stringified version wrapped as JSON string
+        schema_str = json.dumps({"raw": str(input_schema)}, indent=2)
 
     return {
         "name": name,
@@ -57,26 +57,63 @@ def print_tool(tool_dict, detail=False):
         click.echo("-" * 60)
 
 
-async def list_tools_and_schemas(url, tool_name=None):
-    builder = MCPBuilder(url=url)
+async def list_tools_and_schemas(command, url, tool_name=None, transport='sse', args=None, env=None):
+    if args is None:
+        args = []
+    from aiq.tool.mcp.mcp_client_base import MCPSSEClient
+    from aiq.tool.mcp.mcp_client_base import MCPStdioClient
+    from aiq.tool.mcp.mcp_client_base import MCPStreamableHTTPClient
+
     try:
-        if tool_name:
-            tool = await builder.get_tool(tool_name)
-            return [format_tool(tool)]
-        else:
-            tools = await builder.get_tools()
-            return [format_tool(tool) for tool in tools.values()]
+        if transport == 'stdio':
+            client = MCPStdioClient(command=command, args=args, env=env)
+        elif transport == 'streamable-http':
+            client = MCPStreamableHTTPClient(url=url)
+        else:  # sse
+            client = MCPSSEClient(url=url)
+
+        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 MCPBuilder: {e}", err=True)
+        click.echo(f"[ERROR] Failed to fetch tools via MCP client: {e}", err=True)
         return []
 
 
-async def list_tools_direct(url, tool_name=None):
+async def list_tools_direct(command, url, tool_name=None, transport='sse', args=None, env=None):
+    if args is None:
+        args = []
     from mcp import ClientSession
     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:
-        async with sse_client(url=url) as (read, write):
+        if transport == 'stdio':
+
+            def get_stdio_client():
+                return stdio_client(server=StdioServerParameters(command=command, args=args, env=env))
+
+            client = get_stdio_client
+        elif transport == 'streamable-http':
+
+            def get_streamable_http_client():
+                return streamablehttp_client(url=url)
+
+            client = get_streamable_http_client
+        else:  # transport
+
+            def get_sse_client():
+                return sse_client(url=url)
+
+            client = get_sse_client
+
+        async with client() as (read, write):
             async with ClientSession(read, write) as session:
                 await session.initialize()
                 response = await session.list_tools()
@@ -98,20 +135,50 @@ async def list_tools_direct(url, tool_name=None):
 
 @click.group(invoke_without_command=True, help="List tool names (default), or show details with --detail or --tool.")
 @click.option('--direct', is_flag=True, help='Bypass MCPBuilder and use direct MCP protocol')
-@click.option('--url', default='http://localhost:9901/sse', show_default=True, help='MCP server URL')
+@click.option('--url',
+              default='http://localhost:9901/sse',
+              show_default=True,
+              help='For SSE/StreamableHTTP: MCP server URL (e.g. http://localhost:8080/sse)')
+@click.option('--transport',
+              type=click.Choice(['sse', 'stdio', 'streamable-http']),
+              default='sse',
+              show_default=True,
+              help='Type of client to use')
+@click.option('--command', help='For stdio: The command to run (e.g. mcp-server)')
+@click.option('--args', help='For stdio: Additional arguments for the command (space-separated)')
+@click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
 @click.option('--tool', default=None, help='Get details for a specific tool by name')
 @click.option('--detail', is_flag=True, help='Show full details for all tools')
 @click.option('--json-output', is_flag=True, help='Output tool metadata in JSON format')
 @click.pass_context
-def list_mcp(ctx, direct, url, tool, detail, json_output):
+def list_mcp(ctx, direct, url, transport, command, args, env, tool, detail, json_output):
     """
     List tool names (default). Use --detail for full output. If --tool is provided,
     always show full output for that tool.
     """
     if ctx.invoked_subcommand is not None:
         return
+
+    if transport == 'stdio':
+        if not command:
+            click.echo("[ERROR] --command is required when using stdio client type", err=True)
+            return
+
+    if transport in ['sse', 'streamable-http']:
+        if not url:
+            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 or streamable-http client type",
+                err=True)
+            return
+
+    stdio_args = args.split() if args else []
+    stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
+
     fetcher = list_tools_direct if direct else list_tools_and_schemas
-    tools = anyio.run(fetcher, url, tool)
+    tools = anyio.run(fetcher, command, url, tool, transport, stdio_args, stdio_env)
 
     if json_output:
         click.echo(json.dumps(tools, indent=2))