Refactor task architecture: separate pure state from server integration

This refactor eliminates circular imports and simplifies the task API:

Architecture changes:
- Pure task state (TaskContext, TaskStore, helpers) stays in shared/experimental/tasks/
- Server integration (ServerTaskContext, TaskResultHandler, TaskSupport, Experimental)
  moves to server/experimental/
- Empty __init__.py files with absolute imports only

New simplified API:
- server.experimental.enable_tasks() - one-line setup, auto-registers handlers
- ctx.experimental.run_task(work) - spawns work, auto-completes/fails
- ServerTaskContext.elicit()/create_message() - queues requests properly

Key improvements:
- No TYPE_CHECKING hacks or circular import workarounds
- ServerTaskContext reuses session._build_*_request() helpers (no duplication)
- TaskSupport manages task_group lifecycle
- run_task() handles task creation, spawning, and completion automatically

Test changes:
- Removed tests for old internals (test_response_routing, test_elicitation_flow, etc.)
- Added test_run_task_flow.py for new user flow
- Fixed remaining tests to use new API (removed notify= params, updated imports)

Author: maxisbey

examples/servers/simple-task-interactive/mcp_simple_task_interactive/server.py

@@ -1,65 +1,28 @@
-"""Simple interactive task server demonstrating elicitation and sampling."""
+"""Simple interactive task server demonstrating elicitation and sampling.
+
+This example shows the simplified task API where:
+- server.experimental.enable_tasks() sets up all infrastructure
+- ctx.experimental.run_task() handles task lifecycle automatically
+- ServerTaskContext.elicit() and ServerTaskContext.create_message() queue requests properly
+"""
 
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
 from typing import Any
 
-import anyio
 import click
 import mcp.types as types
 import uvicorn
-from anyio.abc import TaskGroup
+from mcp.server.experimental.task_context import ServerTaskContext
 from mcp.server.lowlevel import Server
-from mcp.server.session import ServerSession
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
-from mcp.shared.experimental.tasks import (
-    InMemoryTaskMessageQueue,
-    InMemoryTaskStore,
-    TaskResultHandler,
-    TaskSession,
-    task_execution,
-)
 from starlette.applications import Starlette
 from starlette.routing import Mount
 
+server = Server("simple-task-interactive")
 
-@dataclass
-class AppContext:
-    task_group: TaskGroup
-    store: InMemoryTaskStore
-    queue: InMemoryTaskMessageQueue
-    handler: TaskResultHandler
-    # Track sessions that have been configured (session ID -> bool)
-    configured_sessions: dict[int, bool]
-
-
-@asynccontextmanager
-async def lifespan(server: Server[AppContext, Any]) -> AsyncIterator[AppContext]:
-    store = InMemoryTaskStore()
-    queue = InMemoryTaskMessageQueue()
-    handler = TaskResultHandler(store, queue)
-    async with anyio.create_task_group() as tg:
-        yield AppContext(
-            task_group=tg,
-            store=store,
-            queue=queue,
-            handler=handler,
-            configured_sessions={},
-        )
-    store.cleanup()
-    queue.cleanup()
-
-
-server: Server[AppContext, Any] = Server("simple-task-interactive", lifespan=lifespan)
-
-
-def ensure_handler_configured(session: ServerSession, app: AppContext) -> None:
-    """Ensure the task result handler is configured for this session (once)."""
-    session_id = id(session)
-    if session_id not in app.configured_sessions:
-        session.add_response_router(app.handler)
-        app.configured_sessions[session_id] = True
+# Enable task support - this auto-registers all handlers
+server.experimental.enable_tasks()
 
 
 @server.list_tools()
@@ -84,129 +47,73 @@ async def list_tools() -> list[types.Tool]:
 
 
 @server.call_tool()
-async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[types.TextContent] | types.CreateTaskResult:
+async def handle_call_tool(name: str, arguments: dict[str, Any]) -> types.CallToolResult | types.CreateTaskResult:
     ctx = server.request_context
-    app = ctx.lifespan_context
 
-    # Validate task mode
+    # Validate task mode - this tool requires task augmentation
     ctx.experimental.validate_task_mode(types.TASK_REQUIRED)
 
-    # Ensure handler is configured for response routing
-    ensure_handler_configured(ctx.session, app)
-
-    # Create task
-    metadata = ctx.experimental.task_metadata
-    assert metadata is not None
-    task = await app.store.create_task(metadata)
-
     if name == "confirm_delete":
         filename = arguments.get("filename", "unknown.txt")
         print(f"\n[Server] confirm_delete called for '{filename}'")
-        print(f"[Server] Task created: {task.taskId}")
-
-        async def do_confirm() -> None:
-            async with task_execution(task.taskId, app.store) as task_ctx:
-                task_session = TaskSession(
-                    session=ctx.session,
-                    task_id=task.taskId,
-                    store=app.store,
-                    queue=app.queue,
-                )
-
-                print("[Server] Sending elicitation request to client...")
-                result = await task_session.elicit(
-                    message=f"Are you sure you want to delete '{filename}'?",
-                    requestedSchema={
-                        "type": "object",
-                        "properties": {"confirm": {"type": "boolean"}},
-                        "required": ["confirm"],
-                    },
-                )
-
-                print(f"[Server] Received elicitation response: action={result.action}, content={result.content}")
-                if result.action == "accept" and result.content:
-                    confirmed = result.content.get("confirm", False)
-                    text = f"Deleted '{filename}'" if confirmed else "Deletion cancelled"
-                else:
-                    text = "Deletion cancelled"
-
-                print(f"[Server] Completing task with result: {text}")
-                await task_ctx.complete(
-                    types.CallToolResult(content=[types.TextContent(type="text", text=text)]),
-                    notify=True,
-                )
-
-        app.task_group.start_soon(do_confirm)
+
+        async def do_confirm(task: ServerTaskContext) -> types.CallToolResult:
+            print(f"[Server] Task {task.task_id} starting elicitation...")
+
+            result = await task.elicit(
+                message=f"Are you sure you want to delete '{filename}'?",
+                requestedSchema={
+                    "type": "object",
+                    "properties": {"confirm": {"type": "boolean"}},
+                    "required": ["confirm"],
+                },
+            )
+
+            print(f"[Server] Received elicitation response: action={result.action}, content={result.content}")
+
+            if result.action == "accept" and result.content:
+                confirmed = result.content.get("confirm", False)
+                text = f"Deleted '{filename}'" if confirmed else "Deletion cancelled"
+            else:
+                text = "Deletion cancelled"
+
+            print(f"[Server] Completing task with result: {text}")
+            return types.CallToolResult(content=[types.TextContent(type="text", text=text)])
+
+        # run_task creates the task, spawns work, returns CreateTaskResult immediately
+        return await ctx.experimental.run_task(do_confirm)
 
     elif name == "write_haiku":
         topic = arguments.get("topic", "nature")
         print(f"\n[Server] write_haiku called for topic '{topic}'")
-        print(f"[Server] Task created: {task.taskId}")
-
-        async def do_haiku() -> None:
-            async with task_execution(task.taskId, app.store) as task_ctx:
-                task_session = TaskSession(
-                    session=ctx.session,
-                    task_id=task.taskId,
-                    store=app.store,
-                    queue=app.queue,
-                )
-
-                print("[Server] Sending sampling request to client...")
-                result = await task_session.create_message(
-                    messages=[
-                        types.SamplingMessage(
-                            role="user",
-                            content=types.TextContent(type="text", text=f"Write a haiku about {topic}"),
-                        )
-                    ],
-                    max_tokens=50,
-                )
-
-                haiku = "No response"
-                if isinstance(result.content, types.TextContent):
-                    haiku = result.content.text
-
-                print(f"[Server] Received sampling response: {haiku[:50]}...")
-                print("[Server] Completing task with haiku")
-                await task_ctx.complete(
-                    types.CallToolResult(content=[types.TextContent(type="text", text=f"Haiku:\n{haiku}")]),
-                    notify=True,
-                )
-
-        app.task_group.start_soon(do_haiku)
-
-    return types.CreateTaskResult(task=task)
-
-
-@server.experimental.get_task()
-async def handle_get_task(request: types.GetTaskRequest) -> types.GetTaskResult:
-    app = server.request_context.lifespan_context
-    task = await app.store.get_task(request.params.taskId)
-    if task is None:
-        raise ValueError(f"Task {request.params.taskId} not found")
-    return types.GetTaskResult(
-        taskId=task.taskId,
-        status=task.status,
-        statusMessage=task.statusMessage,
-        createdAt=task.createdAt,
-        lastUpdatedAt=task.lastUpdatedAt,
-        ttl=task.ttl,
-        pollInterval=task.pollInterval,
-    )
 
+        async def do_haiku(task: ServerTaskContext) -> types.CallToolResult:
+            print(f"[Server] Task {task.task_id} starting sampling...")
 
-@server.experimental.get_task_result()
-async def handle_get_task_result(
-    request: types.GetTaskPayloadRequest,
-) -> types.GetTaskPayloadResult:
-    ctx = server.request_context
-    app = ctx.lifespan_context
+            result = await task.create_message(
+                messages=[
+                    types.SamplingMessage(
+                        role="user",
+                        content=types.TextContent(type="text", text=f"Write a haiku about {topic}"),
+                    )
+                ],
+                max_tokens=50,
+            )
+
+            haiku = "No response"
+            if isinstance(result.content, types.TextContent):
+                haiku = result.content.text
 
-    # Ensure handler is configured for this session
-    ensure_handler_configured(ctx.session, app)
+            print(f"[Server] Received sampling response: {haiku[:50]}...")
+            return types.CallToolResult(content=[types.TextContent(type="text", text=f"Haiku:\n{haiku}")])
 
-    return await app.handler.handle(request, ctx.session, ctx.request_id)
+        return await ctx.experimental.run_task(do_haiku)
+
+    else:
+        return types.CallToolResult(
+            content=[types.TextContent(type="text", text=f"Unknown tool: {name}")],
+            isError=True,
+        )
 
 
 def create_app(session_manager: StreamableHTTPSessionManager) -> Starlette:

examples/servers/simple-task/mcp_simple_task/server.py

@@ -12,7 +12,8 @@
 from anyio.abc import TaskGroup
 from mcp.server.lowlevel import Server
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
-from mcp.shared.experimental.tasks import InMemoryTaskStore, task_execution
+from mcp.shared.experimental.tasks.helpers import task_execution
+from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
 from starlette.applications import Starlette
 from starlette.routing import Mount
 

src/mcp/server/experimental/__init__.py

@@ -0,0 +1,11 @@
+"""
+Server-side experimental features.
+
+WARNING: These APIs are experimental and may change without notice.
+
+Import directly from submodules:
+- mcp.server.experimental.task_context.ServerTaskContext
+- mcp.server.experimental.task_support.TaskSupport
+- mcp.server.experimental.task_result_handler.TaskResultHandler
+- mcp.server.experimental.request_context.Experimental
+"""