fix utc datetime import, circular dependency, and add type hints

Author: maxisbey

src/mcp/server/lowlevel/experimental.py

@@ -70,7 +70,7 @@ def decorator(
             logger.debug("Registering handler for ListTasksRequest")
             wrapper = create_call_wrapper(func, ListTasksRequest)
 
-            async def handler(req: ListTasksRequest):
+            async def handler(req: ListTasksRequest) -> ServerResult:
                 result = await wrapper(req)
                 return ServerResult(result)
 
@@ -79,17 +79,23 @@ async def handler(req: ListTasksRequest):
 
         return decorator
 
-    def get_task(self):
+    def get_task(
+        self,
+    ) -> Callable[
+        [Callable[[GetTaskRequest], Awaitable[GetTaskResult]]], Callable[[GetTaskRequest], Awaitable[GetTaskResult]]
+    ]:
         """Register a handler for getting task status.
 
         WARNING: This API is experimental and may change without notice.
         """
 
-        def decorator(func: Callable[[GetTaskRequest], Awaitable[GetTaskResult]]):
+        def decorator(
+            func: Callable[[GetTaskRequest], Awaitable[GetTaskResult]],
+        ) -> Callable[[GetTaskRequest], Awaitable[GetTaskResult]]:
             logger.debug("Registering handler for GetTaskRequest")
             wrapper = create_call_wrapper(func, GetTaskRequest)
 
-            async def handler(req: GetTaskRequest):
+            async def handler(req: GetTaskRequest) -> ServerResult:
                 result = await wrapper(req)
                 return ServerResult(result)
 
@@ -98,17 +104,24 @@ async def handler(req: GetTaskRequest):
 
         return decorator
 
-    def get_task_result(self):
+    def get_task_result(
+        self,
+    ) -> Callable[
+        [Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]],
+        Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
+    ]:
         """Register a handler for getting task results/payload.
 
         WARNING: This API is experimental and may change without notice.
         """
 
-        def decorator(func: Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]):
+        def decorator(
+            func: Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
+        ) -> Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]:
             logger.debug("Registering handler for GetTaskPayloadRequest")
             wrapper = create_call_wrapper(func, GetTaskPayloadRequest)
 
-            async def handler(req: GetTaskPayloadRequest):
+            async def handler(req: GetTaskPayloadRequest) -> ServerResult:
                 result = await wrapper(req)
                 return ServerResult(result)
 
@@ -117,17 +130,24 @@ async def handler(req: GetTaskPayloadRequest):
 
         return decorator
 
-    def cancel_task(self):
+    def cancel_task(
+        self,
+    ) -> Callable[
+        [Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]],
+        Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
+    ]:
         """Register a handler for cancelling tasks.
 
         WARNING: This API is experimental and may change without notice.
         """
 
-        def decorator(func: Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]):
+        def decorator(
+            func: Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
+        ) -> Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]:
             logger.debug("Registering handler for CancelTaskRequest")
             wrapper = create_call_wrapper(func, CancelTaskRequest)
 
-            async def handler(req: CancelTaskRequest):
+            async def handler(req: CancelTaskRequest) -> ServerResult:
                 result = await wrapper(req)
                 return ServerResult(result)
 

src/mcp/shared/context.py

@@ -3,7 +3,7 @@
 
 from typing_extensions import TypeVar
 
-from mcp import McpError
+from mcp.shared.exceptions import McpError
 from mcp.shared.session import BaseSession
 from mcp.types import (
     METHOD_NOT_FOUND,

src/mcp/shared/experimental/tasks/helpers.py

@@ -4,7 +4,7 @@
 
 from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
-from datetime import UTC, datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
@@ -57,7 +57,7 @@ def create_task_state(
     return Task(
         taskId=task_id or generate_task_id(),
         status="working",
-        createdAt=datetime.now(UTC),
+        createdAt=datetime.now(timezone.utc),
         ttl=metadata.ttl,
         pollInterval=500,  # Default 500ms poll interval
     )

src/mcp/shared/experimental/tasks/in_memory_task_store.py

@@ -9,7 +9,7 @@
 """
 
 from dataclasses import dataclass, field
-from datetime import UTC, datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 from mcp.shared.experimental.tasks.helpers import create_task_state, is_terminal
 from mcp.shared.experimental.tasks.store import TaskStore
@@ -51,13 +51,13 @@ def _calculate_expiry(self, ttl_ms: int | None) -> datetime | None:
         """Calculate expiry time from TTL in milliseconds."""
         if ttl_ms is None:
             return None
-        return datetime.now(UTC) + timedelta(milliseconds=ttl_ms)
+        return datetime.now(timezone.utc) + timedelta(milliseconds=ttl_ms)
 
     def _is_expired(self, stored: StoredTask) -> bool:
         """Check if a task has expired."""
         if stored.expires_at is None:
             return False
-        return datetime.now(UTC) >= stored.expires_at
+        return datetime.now(timezone.utc) >= stored.expires_at
 
     def _cleanup_expired(self) -> None:
         """Remove all expired tasks. Called lazily during access operations."""

tests/experimental/tasks/server/test_server.py

@@ -1,6 +1,6 @@
 """Tests for server-side task support (handlers, capabilities, integration)."""
 
-from datetime import UTC, datetime
+from datetime import datetime, timezone
 from typing import Any
 
 import anyio
@@ -54,14 +54,14 @@ async def test_list_tasks_handler() -> None:
         Task(
             taskId="task-1",
             status="working",
-            createdAt=datetime.now(UTC),
+            createdAt=datetime.now(timezone.utc),
             ttl=60000,
             pollInterval=1000,
         ),
         Task(
             taskId="task-2",
             status="completed",
-            createdAt=datetime.now(UTC),
+            createdAt=datetime.now(timezone.utc),
             ttl=60000,
             pollInterval=1000,
         ),
@@ -92,7 +92,7 @@ async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
         return GetTaskResult(
             taskId=request.params.taskId,
             status="working",
-            createdAt=datetime.now(UTC),
+            createdAt=datetime.now(timezone.utc),
             ttl=60000,
             pollInterval=1000,
         )
@@ -140,7 +140,7 @@ async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
         return CancelTaskResult(
             taskId=request.params.taskId,
             status="cancelled",
-            createdAt=datetime.now(UTC),
+            createdAt=datetime.now(timezone.utc),
             ttl=60000,
         )
 
@@ -174,7 +174,7 @@ async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
         return CancelTaskResult(
             taskId=request.params.taskId,
             status="cancelled",
-            createdAt=datetime.now(UTC),
+            createdAt=datetime.now(timezone.utc),
             ttl=None,
         )