refactor: split RequestContext between server and client (#1987)
This commit is contained in:
parent
050aeb6456
commit
b1f7eec3cd
2 changed files with 210 additions and 212 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in a new issue