refactor: split RequestContext between server and client (#1987)

This commit is contained in:
Marcelo Trylesinski 2026-02-03 14:35:07 +01:00 • committed by GitHub
commit b1f7eec3cd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 210 additions and 212 deletions

View file

@ -11,8 +11,10 @@ Use cases:
- Server polls client's task status via tasks/get, tasks/result, etc.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol
from typing import TYPE_CHECKING, Protocol
from pydantic import TypeAdapter
@ -32,7 +34,7 @@ class GetTaskHandlerFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.GetTaskRequestParams,
) -> types.GetTaskResult | types.ErrorData: ... # pragma: no branch
@ -45,7 +47,7 @@ class GetTaskResultHandlerFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.GetTaskPayloadRequestParams,
) -> types.GetTaskPayloadResult | types.ErrorData: ... # pragma: no branch
@ -58,7 +60,7 @@ class ListTasksHandlerFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData: ... # pragma: no branch
@ -71,7 +73,7 @@ class CancelTaskHandlerFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData: ... # pragma: no branch
@ -88,7 +90,7 @@ class TaskAugmentedSamplingFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData: ... # pragma: no branch
@ -106,14 +108,14 @@ class TaskAugmentedElicitationFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.ElicitRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData: ... # pragma: no branch
async def default_get_task_handler(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.GetTaskRequestParams,
) -> types.GetTaskResult | types.ErrorData:
return types.ErrorData(
@ -123,7 +125,7 @@ async def default_get_task_handler(
async def default_get_task_result_handler(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.GetTaskPayloadRequestParams,
) -> types.GetTaskPayloadResult | types.ErrorData:
return types.ErrorData(
@ -133,7 +135,7 @@ async def default_get_task_result_handler(
async def default_list_tasks_handler(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData:
return types.ErrorData(
@ -143,7 +145,7 @@ async def default_list_tasks_handler(
async def default_cancel_task_handler(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData:
return types.ErrorData(
@ -153,7 +155,7 @@ async def default_cancel_task_handler(
async def default_task_augmented_sampling(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData:
@ -164,7 +166,7 @@ async def default_task_augmented_sampling(
async def default_task_augmented_elicitation(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.ElicitRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData:
@ -248,7 +250,7 @@ class ExperimentalTaskHandlers:
async def handle_request(
self,
ctx: RequestContext["ClientSession", Any],
ctx: RequestContext[ClientSession],
responder: RequestResponder[types.ServerRequest, types.ClientResult],
) -> None:
"""Handle a task-related request from the server.

View file

@ -1,3 +1,5 @@
from __future__ import annotations
import logging
from typing import Any, Protocol
@ -22,7 +24,7 @@ logger = logging.getLogger("client")
class SamplingFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: ... # pragma: no branch
@ -30,22 +32,19 @@ class SamplingFnT(Protocol):
class ElicitationFnT(Protocol):
async def __call__(
self,
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.ElicitRequestParams,
) -> types.ElicitResult | types.ErrorData: ... # pragma: no branch
class ListRootsFnT(Protocol):
async def __call__(
self, context: RequestContext["ClientSession", Any]
self, context: RequestContext[ClientSession]
) -> types.ListRootsResult | types.ErrorData: ... # pragma: no branch
class LoggingFnT(Protocol):
async def __call__(
self,
params: types.LoggingMessageNotificationParams,
) -> None: ... # pragma: no branch
async def __call__(self, params: types.LoggingMessageNotificationParams) -> None: ... # pragma: no branch
class MessageHandlerFnT(Protocol):
@ -62,7 +61,7 @@ async def _default_message_handler(
async def _default_sampling_callback(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData:
return types.ErrorData(
@ -72,7 +71,7 @@ async def _default_sampling_callback(
async def _default_elicitation_callback(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.ElicitRequestParams,
) -> types.ElicitResult | types.ErrorData:
return types.ErrorData( # pragma: no cover
@ -82,7 +81,7 @@ async def _default_elicitation_callback(
async def _default_list_roots_callback(
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
) -> types.ListRootsResult | types.ErrorData:
return types.ErrorData(
code=types.INVALID_REQUEST,
@ -153,10 +152,7 @@ class ClientSession(
else None
)
elicitation = (
types.ElicitationCapability(
form=types.FormElicitationCapability(),
url=types.UrlElicitationCapability(),
)
types.ElicitationCapability(form=types.FormElicitationCapability(), url=types.UrlElicitationCapability())
if self._elicitation_callback is not _default_elicitation_callback
else None
)
@ -414,12 +410,7 @@ class ClientSession(
await self.send_notification(types.RootsListChangedNotification())
async def _received_request(self, responder: RequestResponder[types.ServerRequest, types.ClientResult]) -> None:
ctx = RequestContext[ClientSession, Any](
request_id=responder.request_id,
meta=responder.request_meta,
session=self,
lifespan_context=None,
)
ctx = RequestContext[ClientSession](request_id=responder.request_id, meta=responder.request_meta, session=self)
# Delegate to experimental task handler if applicable
if self._task_handlers.handles_request(responder.request):