o c۷ik@sdZddlmZddlmZddlmZmZddlZddl m Z ddl m Z ddl mZdd lmZdd lmZmZdd lmZeGd d d ZdS)z TaskSupport - Configuration for experimental task support. This module provides the TaskSupport class which encapsulates all the infrastructure needed for task-augmented requests: store, queue, and handler. ) AsyncIterator)asynccontextmanager) dataclassfieldN) TaskGroup)TaskResultHandler) ServerSession)InMemoryTaskStore)InMemoryTaskMessageQueueTaskMessageQueue) TaskStorec@seZdZUdZeed<eed<eddZe ed<edddZ e dBed <dd d Z e d e fd dZed edfddZded dfddZedddZdS) TaskSupporta Configuration for experimental task support. Encapsulates the task store, message queue, result handler, and task group for spawning background work. When enabled on a server, this automatically: - Configures response routing for each session - Provides default handlers for task operations - Manages a task group for background task execution Example: # Simple in-memory setup server.experimental.enable_tasks() # Custom store/queue for distributed systems server.experimental.enable_tasks( store=RedisTaskStore(redis_url), queue=RedisTaskMessageQueue(redis_url), ) storequeueF)inithandlerN)rdefault _task_groupreturncCst|j|j|_dS)z/Create the result handler from store and queue.N)rrrrselfrZ/home/ubuntu/vllm_env/lib/python3.10/site-packages/mcp/server/experimental/task_support.py __post_init__3szTaskSupport.__post_init__cCs|jdur td|jS)zGet the task group for spawning background work. Raises: RuntimeError: If not within a run() context Nz7TaskSupport not running. Ensure Server.run() is active.)r RuntimeErrorrrrr task_group7s zTaskSupport.task_groupc Csdt4IdH}||_zdVWd|_nd|_wWdIdHdS1IdHs+wYdS)a" Run the task support lifecycle. This creates a task group for spawning background task work. Called automatically by Server.run(). Usage: async with task_support.run(): # Task group is now available ... N)anyiocreate_task_groupr)rtgrrrrunBs .zTaskSupport.runsessioncCs||jdS)ao Configure a session for task support. This registers the result handler as a response router so that responses to queued requests (elicitation, sampling) are routed back to the waiting resolvers. Called automatically by Server.run() for each new session. Args: session: The session to configure N)add_response_routerr)rr rrrconfigure_sessionVs zTaskSupport.configure_sessioncCs|ttdS)a Create in-memory task support. Suitable for development, testing, and single-process servers. For distributed systems, provide custom store and queue implementations. Returns: TaskSupport configured with in-memory store and queue )rr)r r )clsrrr in_memoryes zTaskSupport.in_memory)rN)rr )__name__ __module__ __qualname____doc__r __annotations__r rrrrrrpropertyrrrrrr" classmethodr$rrrrr s   r )r(collections.abcr contextlibr dataclassesrrr anyio.abcr+mcp.server.experimental.task_result_handlerrmcp.server.sessionr2mcp.shared.experimental.tasks.in_memory_task_storer +mcp.shared.experimental.tasks.message_queuer r #mcp.shared.experimental.tasks.storer r rrrrs