Port client notification dispatch to SDK v2 unwrapped notifications

This commit is contained in:
Jeremiah Lowin 2026-07-05 21:14:19 -04:00
commit 9151bfefd5
No known key found for this signature in database
2 changed files with 52 additions and 48 deletions

View file

@ -31,47 +31,50 @@ class MessageHandler:
# handle all messages
await self.on_message(message)
match message:
# requests
case RequestResponder():
# handle all requests
# TODO(ty): remove when ty supports match statement narrowing
await self.on_request(message) # type: ignore[arg-type] # ty:ignore[invalid-argument-type]
# SDK v2 delivers server-to-client requests wrapped in a
# RequestResponder (with the request unwrapped on `.request`) and
# notifications unwrapped (the monolith notification model itself, no
# `.root` wrapper). `ServerNotification`/`ServerRequest` are UnionTypes,
# so they can't appear in class match patterns — branch on the concrete
# models directly.
if isinstance(message, RequestResponder):
# handle all requests
# ty doesn't narrow the generic RequestResponder cleanly here.
await self.on_request(message) # type: ignore[arg-type] # ty:ignore[invalid-argument-type]
# handle specific requests
# TODO(ty): remove type ignores when ty supports match statement narrowing
match message.request.root: # type: ignore[union-attr] # ty:ignore[unresolved-attribute]
case mcp_types.PingRequest():
await self.on_ping(message.request.root) # type: ignore[union-attr] # ty:ignore[unresolved-attribute]
case mcp_types.ListRootsRequest():
await self.on_list_roots(message.request.root) # type: ignore[union-attr] # ty:ignore[unresolved-attribute]
case mcp_types.CreateMessageRequest():
await self.on_create_message(message.request.root) # type: ignore[union-attr] # ty:ignore[unresolved-attribute]
# handle specific requests
request = message.request
match request:
case mcp_types.PingRequest():
await self.on_ping(request)
case mcp_types.ListRootsRequest():
await self.on_list_roots(request)
case mcp_types.CreateMessageRequest():
await self.on_create_message(request)
# notifications
case mcp_types.ServerNotification():
# handle all notifications
await self.on_notification(message)
elif isinstance(message, Exception):
await self.on_exception(message)
# handle specific notifications
match message.root:
case mcp_types.CancelledNotification():
await self.on_cancelled(message.root)
case mcp_types.ProgressNotification():
await self.on_progress(message.root)
case mcp_types.LoggingMessageNotification():
await self.on_logging_message(message.root)
case mcp_types.ToolListChangedNotification():
await self.on_tool_list_changed(message.root)
case mcp_types.ResourceListChangedNotification():
await self.on_resource_list_changed(message.root)
case mcp_types.PromptListChangedNotification():
await self.on_prompt_list_changed(message.root)
case mcp_types.ResourceUpdatedNotification():
await self.on_resource_updated(message.root)
else:
# notifications (unwrapped monolith models)
await self.on_notification(message)
case Exception():
await self.on_exception(message)
# handle specific notifications
match message:
case mcp_types.CancelledNotification():
await self.on_cancelled(message)
case mcp_types.ProgressNotification():
await self.on_progress(message)
case mcp_types.LoggingMessageNotification():
await self.on_logging_message(message)
case mcp_types.ToolListChangedNotification():
await self.on_tool_list_changed(message)
case mcp_types.ResourceListChangedNotification():
await self.on_resource_list_changed(message)
case mcp_types.PromptListChangedNotification():
await self.on_prompt_list_changed(message)
case mcp_types.ResourceUpdatedNotification():
await self.on_resource_updated(message)
async def on_message(self, message: Message) -> None:
pass

View file

@ -33,11 +33,11 @@ class TaskNotificationHandler(MessageHandler):
async def dispatch(self, message: Message) -> None:
"""Dispatch messages, including task status notifications."""
if isinstance(message, mcp_types.ServerNotification):
if isinstance(message.root, TaskStatusNotification):
client = self._client_ref()
if client:
client._handle_task_status_notification(message.root)
# SDK v2 delivers notifications unwrapped (no `.root` wrapper).
if isinstance(message, TaskStatusNotification):
client = self._client_ref()
if client:
client._handle_task_status_notification(message)
await super().dispatch(message)
@ -178,8 +178,9 @@ class Task(abc.ABC, Generic[TaskResultT]):
self._check_client_connected()
if self._is_immediate:
# Return synthetic completed status
now = datetime.now(timezone.utc)
# Return synthetic completed status. SDK v2 types the task
# timestamps as ISO 8601 strings.
now = datetime.now(timezone.utc).isoformat()
return GetTaskResult(
task_id=self._task_id,
status="completed",