Python: MCP long-running task support in Python (#6319)

* MCP long-running task support in Python * Fix pyupgrade and AGENTS.md reconnect description - pyupgrade: drop forward-reference string annotations in _mcp.py (Python 3.10+ resolves them natively now that MCPTaskOptions is defined before use). - AGENTS.md: align reconnect description with current behavior. Phase 1 (initial tools/call) does NOT retry on connection loss; raises 'connection lost; task state unknown' instead, so a server that accepted the request but lost the response cannot start the operation twice. Phase 2 (tasks/get / tasks/result) still reconnects once against the same task_id. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix bandit nosec marker for CI pipeline * Address PR feedbacks * Clarifiied comments and addressed more PR feedbacks. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-06-04 17:04:55 -07:00
parent 01fc518b29
commit bf4ad48cf2
7 changed files with 1894 additions and 44 deletions
@@ -13,9 +13,12 @@ The Model Context Protocol (MCP) is an open standard for connecting AI agents to
 | **Agent as MCP Server** | [`agent_as_mcp_server.py`](agent_as_mcp_server.py) | Shows how to expose an Agent Framework agent as an MCP server that other AI applications can connect to |
 | **API Key Authentication** | [`mcp_api_key_auth.py`](mcp_api_key_auth.py) | Demonstrates API key authentication with MCP servers using `header_provider`, runtime invocation kwargs, and a command-line API key argument |
 | **GitHub Integration with PAT** | [`mcp_github_pat.py`](mcp_github_pat.py) | Demonstrates connecting to GitHub's MCP server using Personal Access Token (PAT) authentication |
+| **Long-Running Task** | [`mcp_long_running_task.py`](mcp_long_running_task.py) | Demonstrates transparent SEP-2663 long-running task handling for MCP tools that advertise `taskSupport=required`. Self-spawns a stdio MCP child server |

 ## Prerequisites

+Most samples in this folder use OpenAI:
+
 - `OPENAI_API_KEY` environment variable
 - `OPENAI_CHAT_MODEL` environment variable

@@ -23,3 +26,8 @@ Run `mcp_api_key_auth.py` with the MCP API key as the first command-line argumen

 For `mcp_github_pat.py`:
 - `GITHUB_PAT` - Your GitHub Personal Access Token (create at https://github.com/settings/tokens)
+
+For `mcp_long_running_task.py` (uses Azure OpenAI via Entra-ID):
+- Run `az login` once
+- `AZURE_OPENAI_ENDPOINT` - your Azure OpenAI resource endpoint, e.g. `https://<resource>.openai.azure.com/`
+- `AZURE_OPENAI_CHAT_MODEL` (or `AZURE_OPENAI_MODEL`) - the deployment name (e.g. `gpt-4o-mini`)
@@ -0,0 +1,181 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+MCP Long-Running Task (SEP-2663) Example
+
+Demonstrates that ``MCPStdioTool`` transparently drives the MCP long-running
+task lifecycle for tools that advertise ``execution.taskSupport == "required"``.
+The agent observes a single function-call result; the framework handles the
+``tools/call`` → ``tasks/get`` (polled) → ``tasks/result`` sequence in the
+background.
+
+Run it as a single file. The script doubles as both the client and the stdio
+MCP child server (the child branch is selected via ``--server``):
+
+    python mcp_long_running_task.py
+
+Requirements:
+- Azure CLI sign-in (``az login``) — used for Entra-ID auth against Azure OpenAI.
+- ``AZURE_OPENAI_ENDPOINT`` — your Azure OpenAI resource endpoint, e.g.
+  ``https://<resource>.openai.azure.com/``.
+- ``AZURE_OPENAI_CHAT_MODEL`` (or ``AZURE_OPENAI_MODEL``) — the deployment name,
+  e.g. ``gpt-4o-mini``.
+
+This sample uses the lower-level ``mcp.server.lowlevel.Server`` so it can:
+1. Advertise a tool with ``execution=ToolExecution(taskSupport="required")``.
+2. Enable the SDK's experimental task support for the ``tasks/*`` lifecycle.
+"""
+
+import asyncio
+import sys
+from datetime import timedelta
+from typing import Any
+
+from agent_framework import Agent, MCPStdioTool, MCPTaskOptions
+from agent_framework.openai import OpenAIChatClient
+from azure.identity import AzureCliCredential
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+# ---------------------------------------------------------------------------
+# MCP stdio server (child-process branch)
+# ---------------------------------------------------------------------------
+
+
+async def _run_server() -> None:
+    """Run a minimal stdio MCP server exposing one long-running tool."""
+    import mcp.types as types
+    from mcp.server.lowlevel import Server
+    from mcp.server.stdio import stdio_server
+
+    server: Server[Any, Any] = Server("mcp-long-running-task-demo")
+    # Auto-registers handlers for tasks/get, tasks/result, tasks/cancel, tasks/list
+    # backed by an in-memory store.
+    server.experimental.enable_tasks()
+
+    @server.list_tools()
+    async def _list_tools() -> list[types.Tool]:  # pyright: ignore[reportUnusedFunction]
+        return [
+            types.Tool(
+                name="slow_summary",
+                description=(
+                    "Produces a short summary of the supplied text after simulating several seconds of expensive work."
+                ),
+                inputSchema={
+                    "type": "object",
+                    "properties": {
+                        "text": {
+                            "type": "string",
+                            "description": "Text to summarize.",
+                        }
+                    },
+                    "required": ["text"],
+                },
+                # Advertise that this tool MUST be invoked via the task lifecycle.
+                execution=types.ToolExecution(taskSupport="required"),
+            )
+        ]
+
+    @server.call_tool()
+    async def _call_tool(name: str, arguments: dict[str, Any]) -> Any:  # pyright: ignore[reportUnusedFunction]
+        if name != "slow_summary":
+            raise ValueError(f"Unknown tool: {name}")
+
+        ctx = server.request_context
+
+        async def _work(task: Any) -> types.CallToolResult:
+            await task.update_status("Thinking...")
+            await asyncio.sleep(15.0)
+            text: str = (arguments.get("text") or "").strip()
+            words = text.split()
+            preview = " ".join(words[:6]) + ("..." if len(words) > 6 else "")
+            summary = (
+                f"Summarized {len(words)} word(s). First few words: '{preview}'."
+                if words
+                else "No input text was provided."
+            )
+            return types.CallToolResult(
+                content=[types.TextContent(type="text", text=summary)],
+                isError=False,
+            )
+
+        if not ctx.experimental.is_task:
+            # Client invoked the tool without task augmentation. Return a hard
+            # error so a misconfigured client surfaces the problem clearly.
+            return types.CallToolResult(
+                content=[
+                    types.TextContent(
+                        type="text",
+                        text="'slow_summary' must be invoked as a task.",
+                    )
+                ],
+                isError=True,
+            )
+
+        return await ctx.experimental.run_task(_work)
+
+    async with stdio_server() as (read_stream, write_stream):
+        await server.run(read_stream, write_stream, server.create_initialization_options())
+
+
+# ---------------------------------------------------------------------------
+# Agent client (default branch)
+# ---------------------------------------------------------------------------
+
+
+async def _run_client() -> None:
+    mcp_tool = MCPStdioTool(
+        name="LongRunningDemo",
+        description="Demo MCP server exposing a tool that advertises taskSupport=required.",
+        command=sys.executable,
+        args=[__file__, "--server"],
+        # Optional: cap individual tasks at two minutes. The server may apply its
+        # own default if this is omitted.
+        task_options=MCPTaskOptions(default_ttl=timedelta(minutes=2)),
+    )
+
+    async with Agent(
+        client=OpenAIChatClient(credential=AzureCliCredential()),
+        name="LROAgent",
+        instructions=(
+            "You are a helpful assistant. Use the slow_summary tool when the user "
+            "asks for a summary. Wait for the result and present it directly."
+        ),
+        tools=mcp_tool,
+    ) as agent:
+        prompt = (
+            "Please summarize the following text using your slow_summary tool: "
+            "'The Model Context Protocol lets language models talk to external "
+            "tools and resources through a small JSON-RPC surface.'"
+        )
+
+        print("=== run() ===")
+        print(f"User: {prompt}")
+        response = await agent.run(prompt)
+        print(f"Agent: {response.text}\n")
+
+        print("=== run(stream=True) ===")
+        print(f"User: {prompt}")
+        print("Agent: ", end="", flush=True)
+        async for update in agent.run(prompt, stream=True):
+            if update.text:
+                print(update.text, end="", flush=True)
+        print()
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    if len(sys.argv) > 1 and sys.argv[1] == "--server":
+        asyncio.run(_run_server())
+        return
+    asyncio.run(_run_client())
+
+
+if __name__ == "__main__":
+    main()