Completely drop RootModel from types module (#1910)
This commit is contained in:
Marcelo Trylesinski
GitHub
parent
f4672c5084
commit
5fdd48a0d3
36 changed files with 640 additions and 800 deletions
|
|
@ -179,6 +179,60 @@ app = Starlette(routes=[Mount("/", app=mcp.streamable_http_app(json_response=Tru
|
|||
|
||||
**Note:** DNS rebinding protection is automatically enabled when `host` is `127.0.0.1`, `localhost`, or `::1`. This now happens in `sse_app()` and `streamable_http_app()` instead of the constructor.
|
||||
|
||||
### Replace `RootModel` by union types with `TypeAdapter` validation
|
||||
|
||||
The following union types are no longer `RootModel` subclasses:
|
||||
|
||||
- `ClientRequest`
|
||||
- `ServerRequest`
|
||||
- `ClientNotification`
|
||||
- `ServerNotification`
|
||||
- `ClientResult`
|
||||
- `ServerResult`
|
||||
- `JSONRPCMessage`
|
||||
|
||||
This means you can no longer access `.root` on these types or use `model_validate()` directly on them. Instead, use the provided `TypeAdapter` instances for validation.
|
||||
|
||||
**Before (v1):**
|
||||
|
||||
```python
|
||||
from mcp.types import ClientRequest, ServerNotification
|
||||
|
||||
# Using RootModel.model_validate()
|
||||
request = ClientRequest.model_validate(data)
|
||||
actual_request = request.root # Accessing the wrapped value
|
||||
|
||||
notification = ServerNotification.model_validate(data)
|
||||
actual_notification = notification.root
|
||||
```
|
||||
|
||||
**After (v2):**
|
||||
|
||||
```python
|
||||
from mcp.types import client_request_adapter, server_notification_adapter
|
||||
|
||||
# Using TypeAdapter.validate_python()
|
||||
request = client_request_adapter.validate_python(data)
|
||||
# No .root access needed - request is the actual type
|
||||
|
||||
notification = server_notification_adapter.validate_python(data)
|
||||
# No .root access needed - notification is the actual type
|
||||
```
|
||||
|
||||
**Available adapters:**
|
||||
|
||||
| Union Type | Adapter |
|
||||
|------------|---------|
|
||||
| `ClientRequest` | `client_request_adapter` |
|
||||
| `ServerRequest` | `server_request_adapter` |
|
||||
| `ClientNotification` | `client_notification_adapter` |
|
||||
| `ServerNotification` | `server_notification_adapter` |
|
||||
| `ClientResult` | `client_result_adapter` |
|
||||
| `ServerResult` | `server_result_adapter` |
|
||||
| `JSONRPCMessage` | `jsonrpc_message_adapter` |
|
||||
|
||||
All adapters are exported from `mcp.types`.
|
||||
|
||||
### Resource URI type changed from `AnyUrl` to `str`
|
||||
|
||||
The `uri` field on resource-related types now uses `str` instead of Pydantic's `AnyUrl`. This aligns with the [MCP specification schema](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/draft/schema.ts) which defines URIs as plain strings (`uri: string`) without strict URL validation. This change allows relative paths like `users/me` that were previously rejected.
|
||||
|
|
|
|||
|
|
@ -125,18 +125,23 @@ extend-exclude = ["README.md"]
|
|||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"C4", # flake8-comprehensions
|
||||
"C90", # mccabe
|
||||
"D212", # pydocstyle: multi-line docstring summary should start at the first line
|
||||
"E", # pycodestyle
|
||||
"F", # pyflakes
|
||||
"I", # isort
|
||||
"PERF", # Perflint
|
||||
"PL", # Pylint
|
||||
"UP", # pyupgrade
|
||||
"C4", # flake8-comprehensions
|
||||
"C90", # mccabe
|
||||
"D212", # pydocstyle: multi-line docstring summary should start at the first line
|
||||
"E", # pycodestyle
|
||||
"F", # pyflakes
|
||||
"I", # isort
|
||||
"PERF", # Perflint
|
||||
"PL", # Pylint
|
||||
"UP", # pyupgrade
|
||||
"TID251", # https://docs.astral.sh/ruff/rules/banned-api/
|
||||
]
|
||||
ignore = ["PERF203", "PLC0415", "PLR0402"]
|
||||
|
||||
[tool.ruff.lint.flake8-tidy-imports.banned-api]
|
||||
"pydantic.RootModel".msg = "Use `pydantic.TypeAdapter` instead."
|
||||
|
||||
|
||||
[tool.ruff.lint.mccabe]
|
||||
max-complexity = 24 # Default is 10
|
||||
|
||||
|
|
|
|||
|
|
@ -242,7 +242,7 @@ class ExperimentalTaskHandlers:
|
|||
def handles_request(request: types.ServerRequest) -> bool:
|
||||
"""Check if this handler handles the given request type."""
|
||||
return isinstance(
|
||||
request.root,
|
||||
request,
|
||||
types.GetTaskRequest | types.GetTaskPayloadRequest | types.ListTasksRequest | types.CancelTaskRequest,
|
||||
)
|
||||
|
||||
|
|
@ -259,7 +259,7 @@ class ExperimentalTaskHandlers:
|
|||
types.ClientResult | types.ErrorData
|
||||
)
|
||||
|
||||
match responder.request.root:
|
||||
match responder.request:
|
||||
case types.GetTaskRequest(params=params):
|
||||
response = await self.get_task(ctx, params)
|
||||
client_response = client_response_type.validate_python(response)
|
||||
|
|
@ -281,7 +281,7 @@ class ExperimentalTaskHandlers:
|
|||
await responder.respond(client_response)
|
||||
|
||||
case _: # pragma: no cover
|
||||
raise ValueError(f"Unhandled request type: {type(responder.request.root)}")
|
||||
raise ValueError(f"Unhandled request type: {type(responder.request)}")
|
||||
|
||||
|
||||
# Backwards compatibility aliases
|
||||
|
|
|
|||
|
|
@ -92,15 +92,13 @@ class ExperimentalClientFeatures:
|
|||
_meta = types.RequestParams.Meta(**meta)
|
||||
|
||||
return await self._session.send_request(
|
||||
types.ClientRequest(
|
||||
types.CallToolRequest(
|
||||
params=types.CallToolRequestParams(
|
||||
name=name,
|
||||
arguments=arguments,
|
||||
task=types.TaskMetadata(ttl=ttl),
|
||||
_meta=_meta,
|
||||
),
|
||||
)
|
||||
types.CallToolRequest(
|
||||
params=types.CallToolRequestParams(
|
||||
name=name,
|
||||
arguments=arguments,
|
||||
task=types.TaskMetadata(ttl=ttl),
|
||||
_meta=_meta,
|
||||
),
|
||||
),
|
||||
types.CreateTaskResult,
|
||||
)
|
||||
|
|
@ -115,10 +113,8 @@ class ExperimentalClientFeatures:
|
|||
GetTaskResult containing the task status and metadata
|
||||
"""
|
||||
return await self._session.send_request(
|
||||
types.ClientRequest(
|
||||
types.GetTaskRequest(
|
||||
params=types.GetTaskRequestParams(task_id=task_id),
|
||||
)
|
||||
types.GetTaskRequest(
|
||||
params=types.GetTaskRequestParams(task_id=task_id),
|
||||
),
|
||||
types.GetTaskResult,
|
||||
)
|
||||
|
|
@ -142,10 +138,8 @@ class ExperimentalClientFeatures:
|
|||
The task result, validated against result_type
|
||||
"""
|
||||
return await self._session.send_request(
|
||||
types.ClientRequest(
|
||||
types.GetTaskPayloadRequest(
|
||||
params=types.GetTaskPayloadRequestParams(task_id=task_id),
|
||||
)
|
||||
types.GetTaskPayloadRequest(
|
||||
params=types.GetTaskPayloadRequestParams(task_id=task_id),
|
||||
),
|
||||
result_type,
|
||||
)
|
||||
|
|
@ -164,9 +158,7 @@ class ExperimentalClientFeatures:
|
|||
"""
|
||||
params = types.PaginatedRequestParams(cursor=cursor) if cursor else None
|
||||
return await self._session.send_request(
|
||||
types.ClientRequest(
|
||||
types.ListTasksRequest(params=params),
|
||||
),
|
||||
types.ListTasksRequest(params=params),
|
||||
types.ListTasksResult,
|
||||
)
|
||||
|
||||
|
|
@ -180,10 +172,8 @@ class ExperimentalClientFeatures:
|
|||
CancelTaskResult with the updated task state
|
||||
"""
|
||||
return await self._session.send_request(
|
||||
types.ClientRequest(
|
||||
types.CancelTaskRequest(
|
||||
params=types.CancelTaskRequestParams(task_id=task_id),
|
||||
)
|
||||
types.CancelTaskRequest(
|
||||
params=types.CancelTaskRequestParams(task_id=task_id),
|
||||
),
|
||||
types.CancelTaskResult,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -122,13 +122,7 @@ class ClientSession(
|
|||
sampling_capabilities: types.SamplingCapability | None = None,
|
||||
experimental_task_handlers: ExperimentalTaskHandlers | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
read_stream,
|
||||
write_stream,
|
||||
types.ServerRequest,
|
||||
types.ServerNotification,
|
||||
read_timeout_seconds=read_timeout_seconds,
|
||||
)
|
||||
super().__init__(read_stream, write_stream, read_timeout_seconds=read_timeout_seconds)
|
||||
self._client_info = client_info or DEFAULT_CLIENT_INFO
|
||||
self._sampling_callback = sampling_callback or _default_sampling_callback
|
||||
self._sampling_capabilities = sampling_capabilities
|
||||
|
|
@ -143,6 +137,14 @@ class ClientSession(
|
|||
# Experimental: Task handlers (use defaults if not provided)
|
||||
self._task_handlers = experimental_task_handlers or ExperimentalTaskHandlers()
|
||||
|
||||
@property
|
||||
def _receive_request_adapter(self) -> TypeAdapter[types.ServerRequest]:
|
||||
return types.server_request_adapter
|
||||
|
||||
@property
|
||||
def _receive_notification_adapter(self) -> TypeAdapter[types.ServerNotification]:
|
||||
return types.server_notification_adapter
|
||||
|
||||
async def initialize(self) -> types.InitializeResult:
|
||||
sampling = (
|
||||
(self._sampling_capabilities or types.SamplingCapability())
|
||||
|
|
@ -167,20 +169,18 @@ class ClientSession(
|
|||
)
|
||||
|
||||
result = await self.send_request(
|
||||
types.ClientRequest(
|
||||
types.InitializeRequest(
|
||||
params=types.InitializeRequestParams(
|
||||
protocol_version=types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=types.ClientCapabilities(
|
||||
sampling=sampling,
|
||||
elicitation=elicitation,
|
||||
experimental=None,
|
||||
roots=roots,
|
||||
tasks=self._task_handlers.build_capability(),
|
||||
),
|
||||
client_info=self._client_info,
|
||||
types.InitializeRequest(
|
||||
params=types.InitializeRequestParams(
|
||||
protocol_version=types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=types.ClientCapabilities(
|
||||
sampling=sampling,
|
||||
elicitation=elicitation,
|
||||
experimental=None,
|
||||
roots=roots,
|
||||
tasks=self._task_handlers.build_capability(),
|
||||
),
|
||||
)
|
||||
client_info=self._client_info,
|
||||
),
|
||||
),
|
||||
types.InitializeResult,
|
||||
)
|
||||
|
|
@ -190,7 +190,7 @@ class ClientSession(
|
|||
|
||||
self._server_capabilities = result.capabilities
|
||||
|
||||
await self.send_notification(types.ClientNotification(types.InitializedNotification()))
|
||||
await self.send_notification(types.InitializedNotification())
|
||||
|
||||
return result
|
||||
|
||||
|
|
@ -218,10 +218,7 @@ class ClientSession(
|
|||
|
||||
async def send_ping(self) -> types.EmptyResult:
|
||||
"""Send a ping request."""
|
||||
return await self.send_request(
|
||||
types.ClientRequest(types.PingRequest()),
|
||||
types.EmptyResult,
|
||||
)
|
||||
return await self.send_request(types.PingRequest(), types.EmptyResult)
|
||||
|
||||
async def send_progress_notification(
|
||||
self,
|
||||
|
|
@ -232,14 +229,12 @@ class ClientSession(
|
|||
) -> None:
|
||||
"""Send a progress notification."""
|
||||
await self.send_notification(
|
||||
types.ClientNotification(
|
||||
types.ProgressNotification(
|
||||
params=types.ProgressNotificationParams(
|
||||
progress_token=progress_token,
|
||||
progress=progress,
|
||||
total=total,
|
||||
message=message,
|
||||
),
|
||||
types.ProgressNotification(
|
||||
params=types.ProgressNotificationParams(
|
||||
progress_token=progress_token,
|
||||
progress=progress,
|
||||
total=total,
|
||||
message=message,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
|
@ -247,11 +242,7 @@ class ClientSession(
|
|||
async def set_logging_level(self, level: types.LoggingLevel) -> types.EmptyResult:
|
||||
"""Send a logging/setLevel request."""
|
||||
return await self.send_request( # pragma: no cover
|
||||
types.ClientRequest(
|
||||
types.SetLevelRequest(
|
||||
params=types.SetLevelRequestParams(level=level),
|
||||
)
|
||||
),
|
||||
types.SetLevelRequest(params=types.SetLevelRequestParams(level=level)),
|
||||
types.EmptyResult,
|
||||
)
|
||||
|
||||
|
|
@ -261,10 +252,7 @@ class ClientSession(
|
|||
Args:
|
||||
params: Full pagination parameters including cursor and any future fields
|
||||
"""
|
||||
return await self.send_request(
|
||||
types.ClientRequest(types.ListResourcesRequest(params=params)),
|
||||
types.ListResourcesResult,
|
||||
)
|
||||
return await self.send_request(types.ListResourcesRequest(params=params), types.ListResourcesResult)
|
||||
|
||||
async def list_resource_templates(
|
||||
self, *, params: types.PaginatedRequestParams | None = None
|
||||
|
|
@ -275,28 +263,28 @@ class ClientSession(
|
|||
params: Full pagination parameters including cursor and any future fields
|
||||
"""
|
||||
return await self.send_request(
|
||||
types.ClientRequest(types.ListResourceTemplatesRequest(params=params)),
|
||||
types.ListResourceTemplatesRequest(params=params),
|
||||
types.ListResourceTemplatesResult,
|
||||
)
|
||||
|
||||
async def read_resource(self, uri: str | AnyUrl) -> types.ReadResourceResult:
|
||||
"""Send a resources/read request."""
|
||||
return await self.send_request(
|
||||
types.ClientRequest(types.ReadResourceRequest(params=types.ReadResourceRequestParams(uri=str(uri)))),
|
||||
types.ReadResourceRequest(params=types.ReadResourceRequestParams(uri=str(uri))),
|
||||
types.ReadResourceResult,
|
||||
)
|
||||
|
||||
async def subscribe_resource(self, uri: str | AnyUrl) -> types.EmptyResult:
|
||||
"""Send a resources/subscribe request."""
|
||||
return await self.send_request( # pragma: no cover
|
||||
types.ClientRequest(types.SubscribeRequest(params=types.SubscribeRequestParams(uri=str(uri)))),
|
||||
types.SubscribeRequest(params=types.SubscribeRequestParams(uri=str(uri))),
|
||||
types.EmptyResult,
|
||||
)
|
||||
|
||||
async def unsubscribe_resource(self, uri: str | AnyUrl) -> types.EmptyResult:
|
||||
"""Send a resources/unsubscribe request."""
|
||||
return await self.send_request( # pragma: no cover
|
||||
types.ClientRequest(types.UnsubscribeRequest(params=types.UnsubscribeRequestParams(uri=str(uri)))),
|
||||
types.UnsubscribeRequest(params=types.UnsubscribeRequestParams(uri=str(uri))),
|
||||
types.EmptyResult,
|
||||
)
|
||||
|
||||
|
|
@ -316,10 +304,8 @@ class ClientSession(
|
|||
_meta = types.RequestParams.Meta(**meta)
|
||||
|
||||
result = await self.send_request(
|
||||
types.ClientRequest(
|
||||
types.CallToolRequest(
|
||||
params=types.CallToolRequestParams(name=name, arguments=arguments, _meta=_meta),
|
||||
)
|
||||
types.CallToolRequest(
|
||||
params=types.CallToolRequestParams(name=name, arguments=arguments, _meta=_meta),
|
||||
),
|
||||
types.CallToolResult,
|
||||
request_read_timeout_seconds=read_timeout_seconds,
|
||||
|
|
@ -364,17 +350,15 @@ class ClientSession(
|
|||
params: Full pagination parameters including cursor and any future fields
|
||||
"""
|
||||
return await self.send_request(
|
||||
types.ClientRequest(types.ListPromptsRequest(params=params)),
|
||||
types.ListPromptsRequest(params=params),
|
||||
types.ListPromptsResult,
|
||||
)
|
||||
|
||||
async def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> types.GetPromptResult:
|
||||
"""Send a prompts/get request."""
|
||||
return await self.send_request(
|
||||
types.ClientRequest(
|
||||
types.GetPromptRequest(
|
||||
params=types.GetPromptRequestParams(name=name, arguments=arguments),
|
||||
)
|
||||
types.GetPromptRequest(
|
||||
params=types.GetPromptRequestParams(name=name, arguments=arguments),
|
||||
),
|
||||
types.GetPromptResult,
|
||||
)
|
||||
|
|
@ -391,14 +375,12 @@ class ClientSession(
|
|||
context = types.CompletionContext(arguments=context_arguments)
|
||||
|
||||
return await self.send_request(
|
||||
types.ClientRequest(
|
||||
types.CompleteRequest(
|
||||
params=types.CompleteRequestParams(
|
||||
ref=ref,
|
||||
argument=types.CompletionArgument(**argument),
|
||||
context=context,
|
||||
),
|
||||
)
|
||||
types.CompleteRequest(
|
||||
params=types.CompleteRequestParams(
|
||||
ref=ref,
|
||||
argument=types.CompletionArgument(**argument),
|
||||
context=context,
|
||||
),
|
||||
),
|
||||
types.CompleteResult,
|
||||
)
|
||||
|
|
@ -410,7 +392,7 @@ class ClientSession(
|
|||
params: Full pagination parameters including cursor and any future fields
|
||||
"""
|
||||
result = await self.send_request(
|
||||
types.ClientRequest(types.ListToolsRequest(params=params)),
|
||||
types.ListToolsRequest(params=params),
|
||||
types.ListToolsResult,
|
||||
)
|
||||
|
||||
|
|
@ -423,7 +405,7 @@ class ClientSession(
|
|||
|
||||
async def send_roots_list_changed(self) -> None: # pragma: no cover
|
||||
"""Send a roots/list_changed notification."""
|
||||
await self.send_notification(types.ClientNotification(types.RootsListChangedNotification()))
|
||||
await self.send_notification(types.RootsListChangedNotification())
|
||||
|
||||
async def _received_request(self, responder: RequestResponder[types.ServerRequest, types.ClientResult]) -> None:
|
||||
ctx = RequestContext[ClientSession, Any](
|
||||
|
|
@ -440,7 +422,7 @@ class ClientSession(
|
|||
return None
|
||||
|
||||
# Core request handling
|
||||
match responder.request.root:
|
||||
match responder.request:
|
||||
case types.CreateMessageRequest(params=params):
|
||||
with responder:
|
||||
# Check if this is a task-augmented request
|
||||
|
|
@ -469,7 +451,7 @@ class ClientSession(
|
|||
|
||||
case types.PingRequest(): # pragma: no cover
|
||||
with responder:
|
||||
return await responder.respond(types.ClientResult(root=types.EmptyResult()))
|
||||
return await responder.respond(types.EmptyResult())
|
||||
|
||||
case _: # pragma: no cover
|
||||
pass # Task requests handled above by _task_handlers
|
||||
|
|
@ -486,7 +468,7 @@ class ClientSession(
|
|||
async def _received_notification(self, notification: types.ServerNotification) -> None:
|
||||
"""Handle notifications from the server."""
|
||||
# Process specific notification types
|
||||
match notification.root:
|
||||
match notification:
|
||||
case types.LoggingMessageNotification(params=params):
|
||||
await self._logging_callback(params)
|
||||
case types.ElicitCompleteNotification(params=params):
|
||||
|
|
|
|||
|
|
@ -2,7 +2,8 @@ import logging
|
|||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import AnyUrl, BaseModel, Field, RootModel, ValidationError
|
||||
# TODO(Marcelo): We should drop the `RootModel`.
|
||||
from pydantic import AnyUrl, BaseModel, Field, RootModel, ValidationError # noqa: TID251
|
||||
from starlette.datastructures import FormData, QueryParams
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import RedirectResponse, Response
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from dataclasses import dataclass
|
|||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, RootModel, ValidationError
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
|
|
@ -14,11 +14,9 @@ from mcp.server.auth.provider import OAuthAuthorizationServerProvider, Registrat
|
|||
from mcp.server.auth.settings import ClientRegistrationOptions
|
||||
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
|
||||
|
||||
|
||||
class RegistrationRequest(RootModel[OAuthClientMetadata]):
|
||||
# this wrapper is a no-op; it's just to separate out the types exposed to the
|
||||
# provider from what we use in the HTTP handler
|
||||
root: OAuthClientMetadata
|
||||
# this alias is a no-op; it's just to separate out the types exposed to the
|
||||
# provider from what we use in the HTTP handler
|
||||
RegistrationRequest = OAuthClientMetadata
|
||||
|
||||
|
||||
class RegistrationErrorResponse(BaseModel):
|
||||
|
|
@ -35,6 +33,7 @@ class RegistrationHandler:
|
|||
# Implements dynamic client registration as defined in https://datatracker.ietf.org/doc/html/rfc7591#section-3.1
|
||||
try:
|
||||
# Parse request body as JSON
|
||||
# TODO(Marcelo): This is unnecessary. We should use `request.body()`.
|
||||
body = await request.json()
|
||||
client_metadata = OAuthClientMetadata.model_validate(body)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import time
|
|||
from dataclasses import dataclass
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, RootModel, ValidationError
|
||||
from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, TypeAdapter, ValidationError
|
||||
from starlette.requests import Request
|
||||
|
||||
from mcp.server.auth.errors import stringify_pydantic_error
|
||||
|
|
@ -40,18 +40,8 @@ class RefreshTokenRequest(BaseModel):
|
|||
resource: str | None = Field(None, description="Resource indicator for the token")
|
||||
|
||||
|
||||
class TokenRequest(
|
||||
RootModel[
|
||||
Annotated[
|
||||
AuthorizationCodeRequest | RefreshTokenRequest,
|
||||
Field(discriminator="grant_type"),
|
||||
]
|
||||
]
|
||||
):
|
||||
root: Annotated[
|
||||
AuthorizationCodeRequest | RefreshTokenRequest,
|
||||
Field(discriminator="grant_type"),
|
||||
]
|
||||
TokenRequest = Annotated[AuthorizationCodeRequest | RefreshTokenRequest, Field(discriminator="grant_type")]
|
||||
token_request_adapter = TypeAdapter[TokenRequest](TokenRequest)
|
||||
|
||||
|
||||
class TokenErrorResponse(BaseModel):
|
||||
|
|
@ -62,11 +52,10 @@ class TokenErrorResponse(BaseModel):
|
|||
error_uri: AnyHttpUrl | None = None
|
||||
|
||||
|
||||
class TokenSuccessResponse(RootModel[OAuthToken]):
|
||||
# this is just a wrapper over OAuthToken; the only reason we do this
|
||||
# is to have some separation between the HTTP response type, and the
|
||||
# type returned by the provider
|
||||
root: OAuthToken
|
||||
# this is just an alias over OAuthToken; the only reason we do this
|
||||
# is to have some separation between the HTTP response type, and the
|
||||
# type returned by the provider
|
||||
TokenSuccessResponse = OAuthToken
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -107,7 +96,8 @@ class TokenHandler:
|
|||
|
||||
try:
|
||||
form_data = await request.form()
|
||||
token_request = TokenRequest.model_validate(dict(form_data)).root
|
||||
# TODO(Marcelo): Can someone check if this `dict()` wrapper is necessary?
|
||||
token_request = token_request_adapter.validate_python(dict(form_data))
|
||||
except ValidationError as validation_error: # pragma: no cover
|
||||
return self.response(
|
||||
TokenErrorResponse(
|
||||
|
|
@ -186,12 +176,7 @@ class TokenHandler:
|
|||
# Exchange authorization code for tokens
|
||||
tokens = await self.provider.exchange_authorization_code(client_info, auth_code)
|
||||
except TokenError as e:
|
||||
return self.response(
|
||||
TokenErrorResponse(
|
||||
error=e.error,
|
||||
error_description=e.error_description,
|
||||
)
|
||||
)
|
||||
return self.response(TokenErrorResponse(error=e.error, error_description=e.error_description))
|
||||
|
||||
case RefreshTokenRequest(): # pragma: no cover
|
||||
refresh_token = await self.provider.load_refresh_token(client_info, token_request.refresh_token)
|
||||
|
|
@ -229,11 +214,6 @@ class TokenHandler:
|
|||
# Exchange refresh token for new tokens
|
||||
tokens = await self.provider.exchange_refresh_token(client_info, refresh_token, scopes)
|
||||
except TokenError as e:
|
||||
return self.response(
|
||||
TokenErrorResponse(
|
||||
error=e.error,
|
||||
error_description=e.error_description,
|
||||
)
|
||||
)
|
||||
return self.response(TokenErrorResponse(error=e.error, error_description=e.error_description))
|
||||
|
||||
return self.response(TokenSuccessResponse(root=tokens))
|
||||
return self.response(tokens)
|
||||
|
|
|
|||
|
|
@ -49,7 +49,7 @@ class ExperimentalServerSessionFeatures:
|
|||
GetTaskResult containing the task status
|
||||
"""
|
||||
return await self._session.send_request(
|
||||
types.ServerRequest(types.GetTaskRequest(params=types.GetTaskRequestParams(task_id=task_id))),
|
||||
types.GetTaskRequest(params=types.GetTaskRequestParams(task_id=task_id)),
|
||||
types.GetTaskResult,
|
||||
)
|
||||
|
||||
|
|
@ -68,7 +68,7 @@ class ExperimentalServerSessionFeatures:
|
|||
The task result, validated against result_type
|
||||
"""
|
||||
return await self._session.send_request(
|
||||
types.ServerRequest(types.GetTaskPayloadRequest(params=types.GetTaskPayloadRequestParams(task_id=task_id))),
|
||||
types.GetTaskPayloadRequest(params=types.GetTaskPayloadRequestParams(task_id=task_id)),
|
||||
result_type,
|
||||
)
|
||||
|
||||
|
|
@ -120,13 +120,11 @@ class ExperimentalServerSessionFeatures:
|
|||
require_task_augmented_elicitation(client_caps)
|
||||
|
||||
create_result = await self._session.send_request(
|
||||
types.ServerRequest(
|
||||
types.ElicitRequest(
|
||||
params=types.ElicitRequestFormParams(
|
||||
message=message,
|
||||
requested_schema=requested_schema,
|
||||
task=types.TaskMetadata(ttl=ttl),
|
||||
)
|
||||
types.ElicitRequest(
|
||||
params=types.ElicitRequestFormParams(
|
||||
message=message,
|
||||
requested_schema=requested_schema,
|
||||
task=types.TaskMetadata(ttl=ttl),
|
||||
)
|
||||
),
|
||||
types.CreateTaskResult,
|
||||
|
|
@ -185,21 +183,19 @@ class ExperimentalServerSessionFeatures:
|
|||
validate_tool_use_result_messages(messages)
|
||||
|
||||
create_result = await self._session.send_request(
|
||||
types.ServerRequest(
|
||||
types.CreateMessageRequest(
|
||||
params=types.CreateMessageRequestParams(
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
system_prompt=system_prompt,
|
||||
include_context=include_context,
|
||||
temperature=temperature,
|
||||
stop_sequences=stop_sequences,
|
||||
metadata=metadata,
|
||||
model_preferences=model_preferences,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
task=types.TaskMetadata(ttl=ttl),
|
||||
)
|
||||
types.CreateMessageRequest(
|
||||
params=types.CreateMessageRequestParams(
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
system_prompt=system_prompt,
|
||||
include_context=include_context,
|
||||
temperature=temperature,
|
||||
stop_sequences=stop_sequences,
|
||||
metadata=metadata,
|
||||
model_preferences=model_preferences,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
task=types.TaskMetadata(ttl=ttl),
|
||||
)
|
||||
),
|
||||
types.CreateTaskResult,
|
||||
|
|
|
|||
|
|
@ -39,7 +39,6 @@ from mcp.types import (
|
|||
Result,
|
||||
SamplingCapability,
|
||||
SamplingMessage,
|
||||
ServerNotification,
|
||||
Task,
|
||||
TaskMetadata,
|
||||
TaskStatusNotification,
|
||||
|
|
@ -156,17 +155,15 @@ class ServerTaskContext:
|
|||
"""Send a task status notification to the client."""
|
||||
task = self._ctx.task
|
||||
await self._session.send_notification(
|
||||
ServerNotification(
|
||||
TaskStatusNotification(
|
||||
params=TaskStatusNotificationParams(
|
||||
task_id=task.task_id,
|
||||
status=task.status,
|
||||
status_message=task.status_message,
|
||||
created_at=task.created_at,
|
||||
last_updated_at=task.last_updated_at,
|
||||
ttl=task.ttl,
|
||||
poll_interval=task.poll_interval,
|
||||
)
|
||||
TaskStatusNotification(
|
||||
params=TaskStatusNotificationParams(
|
||||
task_id=task.task_id,
|
||||
status=task.status,
|
||||
status_message=task.status_message,
|
||||
created_at=task.created_at,
|
||||
last_updated_at=task.last_updated_at,
|
||||
ttl=task.ttl,
|
||||
poll_interval=task.poll_interval,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,14 +6,7 @@ from types import GenericAlias
|
|||
from typing import Annotated, Any, cast, get_args, get_origin, get_type_hints
|
||||
|
||||
import pydantic_core
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
RootModel,
|
||||
WithJsonSchema,
|
||||
create_model,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, Field, WithJsonSchema, create_model
|
||||
from pydantic.fields import FieldInfo
|
||||
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind
|
||||
from typing_extensions import is_typeddict
|
||||
|
|
@ -484,6 +477,8 @@ def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:
|
|||
|
||||
def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]:
|
||||
"""Create a RootModel for dict[str, T] types."""
|
||||
# TODO(Marcelo): We should not rely on RootModel for this.
|
||||
from pydantic import RootModel # noqa: TID251
|
||||
|
||||
class DictModel(RootModel[dict_annotation]):
|
||||
pass
|
||||
|
|
@ -495,9 +490,7 @@ def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]:
|
|||
return DictModel
|
||||
|
||||
|
||||
def _convert_to_content(
|
||||
result: Any,
|
||||
) -> Sequence[ContentBlock]:
|
||||
def _convert_to_content(result: Any) -> Sequence[ContentBlock]:
|
||||
"""Convert a result to a sequence of content objects.
|
||||
|
||||
Note: This conversion logic comes from previous versions of FastMCP and is being
|
||||
|
|
|
|||
|
|
@ -141,16 +141,14 @@ class ExperimentalHandlers:
|
|||
message=f"Task not found: {req.params.task_id}",
|
||||
)
|
||||
)
|
||||
return ServerResult(
|
||||
GetTaskResult(
|
||||
task_id=task.task_id,
|
||||
status=task.status,
|
||||
status_message=task.status_message,
|
||||
created_at=task.created_at,
|
||||
last_updated_at=task.last_updated_at,
|
||||
ttl=task.ttl,
|
||||
poll_interval=task.poll_interval,
|
||||
)
|
||||
return GetTaskResult(
|
||||
task_id=task.task_id,
|
||||
status=task.status,
|
||||
status_message=task.status_message,
|
||||
created_at=task.created_at,
|
||||
last_updated_at=task.last_updated_at,
|
||||
ttl=task.ttl,
|
||||
poll_interval=task.poll_interval,
|
||||
)
|
||||
|
||||
self._request_handlers[GetTaskRequest] = _default_get_task
|
||||
|
|
@ -158,29 +156,29 @@ class ExperimentalHandlers:
|
|||
# Register get_task_result handler if not already registered
|
||||
if GetTaskPayloadRequest not in self._request_handlers:
|
||||
|
||||
async def _default_get_task_result(req: GetTaskPayloadRequest) -> ServerResult:
|
||||
async def _default_get_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
|
||||
ctx = self._server.request_context
|
||||
result = await support.handler.handle(req, ctx.session, ctx.request_id)
|
||||
return ServerResult(result)
|
||||
return result
|
||||
|
||||
self._request_handlers[GetTaskPayloadRequest] = _default_get_task_result
|
||||
|
||||
# Register list_tasks handler if not already registered
|
||||
if ListTasksRequest not in self._request_handlers:
|
||||
|
||||
async def _default_list_tasks(req: ListTasksRequest) -> ServerResult:
|
||||
async def _default_list_tasks(req: ListTasksRequest) -> ListTasksResult:
|
||||
cursor = req.params.cursor if req.params else None
|
||||
tasks, next_cursor = await support.store.list_tasks(cursor)
|
||||
return ServerResult(ListTasksResult(tasks=tasks, next_cursor=next_cursor))
|
||||
return ListTasksResult(tasks=tasks, next_cursor=next_cursor)
|
||||
|
||||
self._request_handlers[ListTasksRequest] = _default_list_tasks
|
||||
|
||||
# Register cancel_task handler if not already registered
|
||||
if CancelTaskRequest not in self._request_handlers:
|
||||
|
||||
async def _default_cancel_task(req: CancelTaskRequest) -> ServerResult:
|
||||
async def _default_cancel_task(req: CancelTaskRequest) -> CancelTaskResult:
|
||||
result = await cancel_task(support.store, req.params.task_id)
|
||||
return ServerResult(result)
|
||||
return result
|
||||
|
||||
self._request_handlers[CancelTaskRequest] = _default_cancel_task
|
||||
|
||||
|
|
@ -201,9 +199,9 @@ class ExperimentalHandlers:
|
|||
logger.debug("Registering handler for ListTasksRequest")
|
||||
wrapper = create_call_wrapper(func, ListTasksRequest)
|
||||
|
||||
async def handler(req: ListTasksRequest) -> ServerResult:
|
||||
async def handler(req: ListTasksRequest) -> ListTasksResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
return result
|
||||
|
||||
self._request_handlers[ListTasksRequest] = handler
|
||||
return func
|
||||
|
|
@ -226,9 +224,9 @@ class ExperimentalHandlers:
|
|||
logger.debug("Registering handler for GetTaskRequest")
|
||||
wrapper = create_call_wrapper(func, GetTaskRequest)
|
||||
|
||||
async def handler(req: GetTaskRequest) -> ServerResult:
|
||||
async def handler(req: GetTaskRequest) -> GetTaskResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
return result
|
||||
|
||||
self._request_handlers[GetTaskRequest] = handler
|
||||
return func
|
||||
|
|
@ -252,9 +250,9 @@ class ExperimentalHandlers:
|
|||
logger.debug("Registering handler for GetTaskPayloadRequest")
|
||||
wrapper = create_call_wrapper(func, GetTaskPayloadRequest)
|
||||
|
||||
async def handler(req: GetTaskPayloadRequest) -> ServerResult:
|
||||
async def handler(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
return result
|
||||
|
||||
self._request_handlers[GetTaskPayloadRequest] = handler
|
||||
return func
|
||||
|
|
@ -278,9 +276,9 @@ class ExperimentalHandlers:
|
|||
logger.debug("Registering handler for CancelTaskRequest")
|
||||
wrapper = create_call_wrapper(func, CancelTaskRequest)
|
||||
|
||||
async def handler(req: CancelTaskRequest) -> ServerResult:
|
||||
async def handler(req: CancelTaskRequest) -> CancelTaskResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
return result
|
||||
|
||||
self._request_handlers[CancelTaskRequest] = handler
|
||||
return func
|
||||
|
|
|
|||
|
|
@ -271,10 +271,10 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
result = await wrapper(req)
|
||||
# Handle both old style (list[Prompt]) and new style (ListPromptsResult)
|
||||
if isinstance(result, types.ListPromptsResult):
|
||||
return types.ServerResult(result)
|
||||
return result
|
||||
else:
|
||||
# Old style returns list[Prompt]
|
||||
return types.ServerResult(types.ListPromptsResult(prompts=result))
|
||||
return types.ListPromptsResult(prompts=result)
|
||||
|
||||
self.request_handlers[types.ListPromptsRequest] = handler
|
||||
return func
|
||||
|
|
@ -289,7 +289,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
async def handler(req: types.GetPromptRequest):
|
||||
prompt_get = await func(req.params.name, req.params.arguments)
|
||||
return types.ServerResult(prompt_get)
|
||||
return prompt_get
|
||||
|
||||
self.request_handlers[types.GetPromptRequest] = handler
|
||||
return func
|
||||
|
|
@ -309,10 +309,10 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
result = await wrapper(req)
|
||||
# Handle both old style (list[Resource]) and new style (ListResourcesResult)
|
||||
if isinstance(result, types.ListResourcesResult):
|
||||
return types.ServerResult(result)
|
||||
return result
|
||||
else:
|
||||
# Old style returns list[Resource]
|
||||
return types.ServerResult(types.ListResourcesResult(resources=result))
|
||||
return types.ListResourcesResult(resources=result)
|
||||
|
||||
self.request_handlers[types.ListResourcesRequest] = handler
|
||||
return func
|
||||
|
|
@ -325,7 +325,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
async def handler(_: Any):
|
||||
templates = await func()
|
||||
return types.ServerResult(types.ListResourceTemplatesResult(resource_templates=templates))
|
||||
return types.ListResourceTemplatesResult(resource_templates=templates)
|
||||
|
||||
self.request_handlers[types.ListResourceTemplatesRequest] = handler
|
||||
return func
|
||||
|
|
@ -376,18 +376,12 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
)
|
||||
for content_item in contents
|
||||
]
|
||||
return types.ServerResult(
|
||||
types.ReadResourceResult(
|
||||
contents=contents_list,
|
||||
)
|
||||
)
|
||||
return types.ReadResourceResult(contents=contents_list)
|
||||
case _: # pragma: no cover
|
||||
raise ValueError(f"Unexpected return type from read_resource: {type(result)}")
|
||||
|
||||
return types.ServerResult( # pragma: no cover
|
||||
types.ReadResourceResult(
|
||||
contents=[content],
|
||||
)
|
||||
return types.ReadResourceResult( # pragma: no cover
|
||||
contents=[content],
|
||||
)
|
||||
|
||||
self.request_handlers[types.ReadResourceRequest] = handler
|
||||
|
|
@ -401,7 +395,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
async def handler(req: types.SetLevelRequest):
|
||||
await func(req.params.level)
|
||||
return types.ServerResult(types.EmptyResult())
|
||||
return types.EmptyResult()
|
||||
|
||||
self.request_handlers[types.SetLevelRequest] = handler
|
||||
return func
|
||||
|
|
@ -414,7 +408,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
async def handler(req: types.SubscribeRequest):
|
||||
await func(req.params.uri)
|
||||
return types.ServerResult(types.EmptyResult())
|
||||
return types.EmptyResult()
|
||||
|
||||
self.request_handlers[types.SubscribeRequest] = handler
|
||||
return func
|
||||
|
|
@ -427,7 +421,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
async def handler(req: types.UnsubscribeRequest):
|
||||
await func(req.params.uri)
|
||||
return types.ServerResult(types.EmptyResult())
|
||||
return types.EmptyResult()
|
||||
|
||||
self.request_handlers[types.UnsubscribeRequest] = handler
|
||||
return func
|
||||
|
|
@ -452,7 +446,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
for tool in result.tools:
|
||||
validate_and_warn_tool_name(tool.name)
|
||||
self._tool_cache[tool.name] = tool
|
||||
return types.ServerResult(result)
|
||||
return result
|
||||
else:
|
||||
# Old style returns list[Tool]
|
||||
# Clear and refresh the entire tool cache
|
||||
|
|
@ -460,20 +454,18 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
for tool in result:
|
||||
validate_and_warn_tool_name(tool.name)
|
||||
self._tool_cache[tool.name] = tool
|
||||
return types.ServerResult(types.ListToolsResult(tools=result))
|
||||
return types.ListToolsResult(tools=result)
|
||||
|
||||
self.request_handlers[types.ListToolsRequest] = handler
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def _make_error_result(self, error_message: str) -> types.ServerResult:
|
||||
"""Create a ServerResult with an error CallToolResult."""
|
||||
return types.ServerResult(
|
||||
types.CallToolResult(
|
||||
content=[types.TextContent(type="text", text=error_message)],
|
||||
is_error=True,
|
||||
)
|
||||
def _make_error_result(self, error_message: str) -> types.CallToolResult:
|
||||
"""Create a CallToolResult with an error."""
|
||||
return types.CallToolResult(
|
||||
content=[types.TextContent(type="text", text=error_message)],
|
||||
is_error=True,
|
||||
)
|
||||
|
||||
async def _get_cached_tool_definition(self, tool_name: str) -> types.Tool | None:
|
||||
|
|
@ -541,10 +533,10 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
unstructured_content: UnstructuredContent
|
||||
maybe_structured_content: StructuredContent | None
|
||||
if isinstance(results, types.CallToolResult):
|
||||
return types.ServerResult(results)
|
||||
return results
|
||||
elif isinstance(results, types.CreateTaskResult):
|
||||
# Task-augmented execution returns task info instead of result
|
||||
return types.ServerResult(results)
|
||||
return results
|
||||
elif isinstance(results, tuple) and len(results) == 2:
|
||||
# tool returned both structured and unstructured content
|
||||
unstructured_content, maybe_structured_content = cast(CombinationContent, results)
|
||||
|
|
@ -572,12 +564,10 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
return self._make_error_result(f"Output validation error: {e.message}")
|
||||
|
||||
# result
|
||||
return types.ServerResult(
|
||||
types.CallToolResult(
|
||||
content=list(unstructured_content),
|
||||
structured_content=maybe_structured_content,
|
||||
is_error=False,
|
||||
)
|
||||
return types.CallToolResult(
|
||||
content=list(unstructured_content),
|
||||
structured_content=maybe_structured_content,
|
||||
is_error=False,
|
||||
)
|
||||
except UrlElicitationRequiredError:
|
||||
# Re-raise UrlElicitationRequiredError so it can be properly handled
|
||||
|
|
@ -627,12 +617,10 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
async def handler(req: types.CompleteRequest):
|
||||
completion = await func(req.params.ref, req.params.argument, req.params.context)
|
||||
return types.ServerResult(
|
||||
types.CompleteResult(
|
||||
completion=completion
|
||||
if completion is not None
|
||||
else types.Completion(values=[], total=None, has_more=None),
|
||||
)
|
||||
return types.CompleteResult(
|
||||
completion=completion
|
||||
if completion is not None
|
||||
else types.Completion(values=[], total=None, has_more=None),
|
||||
)
|
||||
|
||||
self.request_handlers[types.CompleteRequest] = handler
|
||||
|
|
@ -694,11 +682,11 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
):
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
match message:
|
||||
case RequestResponder(request=types.ClientRequest(root=req)) as responder:
|
||||
case RequestResponder() as responder:
|
||||
with responder:
|
||||
await self._handle_request(message, req, session, lifespan_context, raise_exceptions)
|
||||
case types.ClientNotification(root=notify):
|
||||
await self._handle_notification(notify)
|
||||
await self._handle_request(
|
||||
message, responder.request, session, lifespan_context, raise_exceptions
|
||||
)
|
||||
case Exception(): # pragma: no cover
|
||||
logger.error(f"Received exception from stream: {message}")
|
||||
await session.send_log_message(
|
||||
|
|
@ -708,6 +696,8 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
)
|
||||
if raise_exceptions:
|
||||
raise message
|
||||
case _:
|
||||
await self._handle_notification(message)
|
||||
|
||||
for warning in w: # pragma: no cover
|
||||
logger.info("Warning: %s: %s", warning.category.__name__, warning.message)
|
||||
|
|
@ -715,7 +705,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
async def _handle_request(
|
||||
self,
|
||||
message: RequestResponder[types.ClientRequest, types.ServerResult],
|
||||
req: types.ClientRequestType,
|
||||
req: types.ClientRequest,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
raise_exceptions: bool,
|
||||
|
|
@ -803,4 +793,4 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
|
||||
async def _ping_handler(request: types.PingRequest) -> types.ServerResult:
|
||||
return types.ServerResult(types.EmptyResult())
|
||||
return types.EmptyResult()
|
||||
|
|
|
|||
|
|
@ -92,7 +92,7 @@ class ServerSession(
|
|||
init_options: InitializationOptions,
|
||||
stateless: bool = False,
|
||||
) -> None:
|
||||
super().__init__(read_stream, write_stream, types.ClientRequest, types.ClientNotification)
|
||||
super().__init__(read_stream, write_stream)
|
||||
self._stateless = stateless
|
||||
self._initialization_state = (
|
||||
InitializationState.Initialized if stateless else InitializationState.NotInitialized
|
||||
|
|
@ -104,6 +104,14 @@ class ServerSession(
|
|||
](0)
|
||||
self._exit_stack.push_async_callback(lambda: self._incoming_message_stream_reader.aclose())
|
||||
|
||||
@property
|
||||
def _receive_request_adapter(self) -> types.TypeAdapter[types.ClientRequest]:
|
||||
return types.client_request_adapter
|
||||
|
||||
@property
|
||||
def _receive_notification_adapter(self) -> types.TypeAdapter[types.ClientNotification]:
|
||||
return types.client_notification_adapter
|
||||
|
||||
@property
|
||||
def client_params(self) -> types.InitializeRequestParams | None:
|
||||
return self._client_params # pragma: no cover
|
||||
|
|
@ -162,29 +170,27 @@ class ServerSession(
|
|||
await super()._receive_loop()
|
||||
|
||||
async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]):
|
||||
match responder.request.root:
|
||||
match responder.request:
|
||||
case types.InitializeRequest(params=params):
|
||||
requested_version = params.protocol_version
|
||||
self._initialization_state = InitializationState.Initializing
|
||||
self._client_params = params
|
||||
with responder:
|
||||
await responder.respond(
|
||||
types.ServerResult(
|
||||
types.InitializeResult(
|
||||
protocol_version=requested_version
|
||||
if requested_version in SUPPORTED_PROTOCOL_VERSIONS
|
||||
else types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=self._init_options.capabilities,
|
||||
server_info=types.Implementation(
|
||||
name=self._init_options.server_name,
|
||||
title=self._init_options.title,
|
||||
description=self._init_options.description,
|
||||
version=self._init_options.server_version,
|
||||
website_url=self._init_options.website_url,
|
||||
icons=self._init_options.icons,
|
||||
),
|
||||
instructions=self._init_options.instructions,
|
||||
)
|
||||
types.InitializeResult(
|
||||
protocol_version=requested_version
|
||||
if requested_version in SUPPORTED_PROTOCOL_VERSIONS
|
||||
else types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=self._init_options.capabilities,
|
||||
server_info=types.Implementation(
|
||||
name=self._init_options.server_name,
|
||||
title=self._init_options.title,
|
||||
description=self._init_options.description,
|
||||
version=self._init_options.server_version,
|
||||
website_url=self._init_options.website_url,
|
||||
icons=self._init_options.icons,
|
||||
),
|
||||
instructions=self._init_options.instructions,
|
||||
)
|
||||
)
|
||||
self._initialization_state = InitializationState.Initialized
|
||||
|
|
@ -198,7 +204,7 @@ class ServerSession(
|
|||
async def _received_notification(self, notification: types.ClientNotification) -> None:
|
||||
# Need this to avoid ASYNC910
|
||||
await anyio.lowlevel.checkpoint()
|
||||
match notification.root:
|
||||
match notification:
|
||||
case types.InitializedNotification():
|
||||
self._initialization_state = InitializationState.Initialized
|
||||
case _:
|
||||
|
|
@ -214,14 +220,12 @@ class ServerSession(
|
|||
) -> None:
|
||||
"""Send a log message notification."""
|
||||
await self.send_notification(
|
||||
types.ServerNotification(
|
||||
types.LoggingMessageNotification(
|
||||
params=types.LoggingMessageNotificationParams(
|
||||
level=level,
|
||||
data=data,
|
||||
logger=logger,
|
||||
),
|
||||
)
|
||||
types.LoggingMessageNotification(
|
||||
params=types.LoggingMessageNotificationParams(
|
||||
level=level,
|
||||
data=data,
|
||||
logger=logger,
|
||||
),
|
||||
),
|
||||
related_request_id,
|
||||
)
|
||||
|
|
@ -229,10 +233,8 @@ class ServerSession(
|
|||
async def send_resource_updated(self, uri: str | AnyUrl) -> None: # pragma: no cover
|
||||
"""Send a resource updated notification."""
|
||||
await self.send_notification(
|
||||
types.ServerNotification(
|
||||
types.ResourceUpdatedNotification(
|
||||
params=types.ResourceUpdatedNotificationParams(uri=str(uri)),
|
||||
)
|
||||
types.ResourceUpdatedNotification(
|
||||
params=types.ResourceUpdatedNotificationParams(uri=str(uri)),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -322,21 +324,19 @@ class ServerSession(
|
|||
validate_sampling_tools(client_caps, tools, tool_choice)
|
||||
validate_tool_use_result_messages(messages)
|
||||
|
||||
request = types.ServerRequest(
|
||||
types.CreateMessageRequest(
|
||||
params=types.CreateMessageRequestParams(
|
||||
messages=messages,
|
||||
system_prompt=system_prompt,
|
||||
include_context=include_context,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
stop_sequences=stop_sequences,
|
||||
metadata=metadata,
|
||||
model_preferences=model_preferences,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
),
|
||||
)
|
||||
request = types.CreateMessageRequest(
|
||||
params=types.CreateMessageRequestParams(
|
||||
messages=messages,
|
||||
system_prompt=system_prompt,
|
||||
include_context=include_context,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
stop_sequences=stop_sequences,
|
||||
metadata=metadata,
|
||||
model_preferences=model_preferences,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
),
|
||||
)
|
||||
metadata_obj = ServerMessageMetadata(related_request_id=related_request_id)
|
||||
|
||||
|
|
@ -358,7 +358,7 @@ class ServerSession(
|
|||
if self._stateless:
|
||||
raise StatelessModeNotSupported(method="list_roots")
|
||||
return await self.send_request(
|
||||
types.ServerRequest(types.ListRootsRequest()),
|
||||
types.ListRootsRequest(),
|
||||
types.ListRootsResult,
|
||||
)
|
||||
|
||||
|
|
@ -406,13 +406,11 @@ class ServerSession(
|
|||
if self._stateless:
|
||||
raise StatelessModeNotSupported(method="elicitation")
|
||||
return await self.send_request(
|
||||
types.ServerRequest(
|
||||
types.ElicitRequest(
|
||||
params=types.ElicitRequestFormParams(
|
||||
message=message,
|
||||
requested_schema=requested_schema,
|
||||
),
|
||||
)
|
||||
types.ElicitRequest(
|
||||
params=types.ElicitRequestFormParams(
|
||||
message=message,
|
||||
requested_schema=requested_schema,
|
||||
),
|
||||
),
|
||||
types.ElicitResult,
|
||||
metadata=ServerMessageMetadata(related_request_id=related_request_id),
|
||||
|
|
@ -445,14 +443,12 @@ class ServerSession(
|
|||
if self._stateless:
|
||||
raise StatelessModeNotSupported(method="elicitation")
|
||||
return await self.send_request(
|
||||
types.ServerRequest(
|
||||
types.ElicitRequest(
|
||||
params=types.ElicitRequestURLParams(
|
||||
message=message,
|
||||
url=url,
|
||||
elicitation_id=elicitation_id,
|
||||
),
|
||||
)
|
||||
types.ElicitRequest(
|
||||
params=types.ElicitRequestURLParams(
|
||||
message=message,
|
||||
url=url,
|
||||
elicitation_id=elicitation_id,
|
||||
),
|
||||
),
|
||||
types.ElicitResult,
|
||||
metadata=ServerMessageMetadata(related_request_id=related_request_id),
|
||||
|
|
@ -461,7 +457,7 @@ class ServerSession(
|
|||
async def send_ping(self) -> types.EmptyResult: # pragma: no cover
|
||||
"""Send a ping request."""
|
||||
return await self.send_request(
|
||||
types.ServerRequest(types.PingRequest()),
|
||||
types.PingRequest(),
|
||||
types.EmptyResult,
|
||||
)
|
||||
|
||||
|
|
@ -475,30 +471,28 @@ class ServerSession(
|
|||
) -> None:
|
||||
"""Send a progress notification."""
|
||||
await self.send_notification(
|
||||
types.ServerNotification(
|
||||
types.ProgressNotification(
|
||||
params=types.ProgressNotificationParams(
|
||||
progress_token=progress_token,
|
||||
progress=progress,
|
||||
total=total,
|
||||
message=message,
|
||||
),
|
||||
)
|
||||
types.ProgressNotification(
|
||||
params=types.ProgressNotificationParams(
|
||||
progress_token=progress_token,
|
||||
progress=progress,
|
||||
total=total,
|
||||
message=message,
|
||||
),
|
||||
),
|
||||
related_request_id,
|
||||
)
|
||||
|
||||
async def send_resource_list_changed(self) -> None: # pragma: no cover
|
||||
"""Send a resource list changed notification."""
|
||||
await self.send_notification(types.ServerNotification(types.ResourceListChangedNotification()))
|
||||
await self.send_notification(types.ResourceListChangedNotification())
|
||||
|
||||
async def send_tool_list_changed(self) -> None: # pragma: no cover
|
||||
"""Send a tool list changed notification."""
|
||||
await self.send_notification(types.ServerNotification(types.ToolListChangedNotification()))
|
||||
await self.send_notification(types.ToolListChangedNotification())
|
||||
|
||||
async def send_prompt_list_changed(self) -> None: # pragma: no cover
|
||||
"""Send a prompt list changed notification."""
|
||||
await self.send_notification(types.ServerNotification(types.PromptListChangedNotification()))
|
||||
await self.send_notification(types.PromptListChangedNotification())
|
||||
|
||||
async def send_elicit_complete(
|
||||
self,
|
||||
|
|
@ -516,10 +510,8 @@ class ServerSession(
|
|||
related_request_id: Optional ID of the request that triggered this
|
||||
"""
|
||||
await self.send_notification(
|
||||
types.ServerNotification(
|
||||
types.ElicitCompleteNotification(
|
||||
params=types.ElicitCompleteNotificationParams(elicitation_id=elicitation_id)
|
||||
)
|
||||
types.ElicitCompleteNotification(
|
||||
params=types.ElicitCompleteNotificationParams(elicitation_id=elicitation_id)
|
||||
),
|
||||
related_request_id,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any, Generic, Protocol, TypeVar
|
|||
import anyio
|
||||
import httpx
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import Self
|
||||
|
||||
from mcp.shared.exceptions import McpError
|
||||
|
|
@ -179,8 +179,6 @@ class BaseSession(
|
|||
self,
|
||||
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
|
||||
write_stream: MemoryObjectSendStream[SessionMessage],
|
||||
receive_request_type: type[ReceiveRequestT],
|
||||
receive_notification_type: type[ReceiveNotificationT],
|
||||
# If none, reading will never time out
|
||||
read_timeout_seconds: float | None = None,
|
||||
) -> None:
|
||||
|
|
@ -188,8 +186,6 @@ class BaseSession(
|
|||
self._write_stream = write_stream
|
||||
self._response_streams = {}
|
||||
self._request_id = 0
|
||||
self._receive_request_type = receive_request_type
|
||||
self._receive_notification_type = receive_notification_type
|
||||
self._session_read_timeout_seconds = read_timeout_seconds
|
||||
self._in_flight = {}
|
||||
self._progress_callbacks = {}
|
||||
|
|
@ -264,11 +260,7 @@ class BaseSession(
|
|||
self._progress_callbacks[request_id] = progress_callback
|
||||
|
||||
try:
|
||||
jsonrpc_request = JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=request_id,
|
||||
**request_data,
|
||||
)
|
||||
jsonrpc_request = JSONRPCRequest(jsonrpc="2.0", id=request_id, **request_data)
|
||||
|
||||
await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
|
||||
|
||||
|
|
@ -339,26 +331,30 @@ class BaseSession(
|
|||
session_message = SessionMessage(message=jsonrpc_response)
|
||||
await self._write_stream.send(session_message)
|
||||
|
||||
@property
|
||||
def _receive_request_adapter(self) -> TypeAdapter[ReceiveRequestT]:
|
||||
"""Each subclass must provide its own request adapter."""
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def _receive_notification_adapter(self) -> TypeAdapter[ReceiveNotificationT]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def _receive_loop(self) -> None:
|
||||
async with (
|
||||
self._read_stream,
|
||||
self._write_stream,
|
||||
):
|
||||
async with self._read_stream, self._write_stream:
|
||||
try:
|
||||
async for message in self._read_stream:
|
||||
if isinstance(message, Exception): # pragma: no cover
|
||||
await self._handle_incoming(message)
|
||||
elif isinstance(message.message, JSONRPCRequest):
|
||||
try:
|
||||
validated_request = self._receive_request_type.model_validate(
|
||||
validated_request = self._receive_request_adapter.validate_python(
|
||||
message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
by_name=False,
|
||||
)
|
||||
responder = RequestResponder(
|
||||
request_id=message.message.id,
|
||||
request_meta=validated_request.root.params.meta
|
||||
if validated_request.root.params
|
||||
else None,
|
||||
request_meta=validated_request.params.meta if validated_request.params else None,
|
||||
request=validated_request,
|
||||
session=self,
|
||||
on_complete=lambda r: self._in_flight.pop(r.request_id, None),
|
||||
|
|
@ -369,10 +365,10 @@ class BaseSession(
|
|||
|
||||
if not responder._completed: # type: ignore[reportPrivateUsage]
|
||||
await self._handle_incoming(responder)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
# For request validation errors, send a proper JSON-RPC error
|
||||
# response instead of crashing the server
|
||||
logging.warning(f"Failed to validate request: {e}")
|
||||
logging.warning("Failed to validate request", exc_info=True)
|
||||
logging.debug(f"Message that failed validation: {message.message}")
|
||||
error_response = JSONRPCError(
|
||||
jsonrpc="2.0",
|
||||
|
|
@ -388,34 +384,31 @@ class BaseSession(
|
|||
|
||||
elif isinstance(message.message, JSONRPCNotification):
|
||||
try:
|
||||
notification = self._receive_notification_type.model_validate(
|
||||
notification = self._receive_notification_adapter.validate_python(
|
||||
message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
|
||||
by_name=False,
|
||||
)
|
||||
# Handle cancellation notifications
|
||||
if isinstance(notification.root, CancelledNotification):
|
||||
cancelled_id = notification.root.params.request_id
|
||||
if isinstance(notification, CancelledNotification):
|
||||
cancelled_id = notification.params.request_id
|
||||
if cancelled_id in self._in_flight: # pragma: no branch
|
||||
await self._in_flight[cancelled_id].cancel()
|
||||
else:
|
||||
# Handle progress notifications callback
|
||||
if isinstance(notification.root, ProgressNotification): # pragma: no cover
|
||||
progress_token = notification.root.params.progress_token
|
||||
if isinstance(notification, ProgressNotification): # pragma: no cover
|
||||
progress_token = notification.params.progress_token
|
||||
# If there is a progress callback for this token,
|
||||
# call it with the progress information
|
||||
if progress_token in self._progress_callbacks:
|
||||
callback = self._progress_callbacks[progress_token]
|
||||
try:
|
||||
await callback(
|
||||
notification.root.params.progress,
|
||||
notification.root.params.total,
|
||||
notification.root.params.message,
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
"Progress callback raised an exception: %s",
|
||||
e,
|
||||
notification.params.progress,
|
||||
notification.params.total,
|
||||
notification.params.message,
|
||||
)
|
||||
except Exception:
|
||||
logging.exception("Progress callback raised an exception")
|
||||
await self._received_notification(notification)
|
||||
await self._handle_incoming(notification)
|
||||
except Exception: # pragma: no cover
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from collections.abc import Callable
|
|||
from datetime import datetime
|
||||
from typing import Annotated, Any, Final, Generic, Literal, TypeAlias, TypeVar
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel, TypeAdapter
|
||||
from pydantic import BaseModel, ConfigDict, Field, FileUrl, TypeAdapter
|
||||
from pydantic.alias_generators import to_camel
|
||||
|
||||
LATEST_PROTOCOL_VERSION = "2025-11-25"
|
||||
|
|
@ -1631,7 +1631,7 @@ class ElicitCompleteNotification(
|
|||
params: ElicitCompleteNotificationParams
|
||||
|
||||
|
||||
ClientRequestType: TypeAlias = (
|
||||
ClientRequest = (
|
||||
PingRequest
|
||||
| InitializeRequest
|
||||
| CompleteRequest
|
||||
|
|
@ -1650,23 +1650,17 @@ ClientRequestType: TypeAlias = (
|
|||
| ListTasksRequest
|
||||
| CancelTaskRequest
|
||||
)
|
||||
client_request_adapter = TypeAdapter[ClientRequest](ClientRequest)
|
||||
|
||||
|
||||
class ClientRequest(RootModel[ClientRequestType]):
|
||||
pass
|
||||
|
||||
|
||||
ClientNotificationType: TypeAlias = (
|
||||
ClientNotification = (
|
||||
CancelledNotification
|
||||
| ProgressNotification
|
||||
| InitializedNotification
|
||||
| RootsListChangedNotification
|
||||
| TaskStatusNotification
|
||||
)
|
||||
|
||||
|
||||
class ClientNotification(RootModel[ClientNotificationType]):
|
||||
pass
|
||||
client_notification_adapter = TypeAdapter[ClientNotification](ClientNotification)
|
||||
|
||||
|
||||
# Type for elicitation schema - a JSON Schema dict
|
||||
|
|
@ -1760,7 +1754,7 @@ class ElicitationRequiredErrorData(MCPModel):
|
|||
"""List of URL mode elicitations that must be completed."""
|
||||
|
||||
|
||||
ClientResultType: TypeAlias = (
|
||||
ClientResult = (
|
||||
EmptyResult
|
||||
| CreateMessageResult
|
||||
| CreateMessageResultWithTools
|
||||
|
|
@ -1772,13 +1766,10 @@ ClientResultType: TypeAlias = (
|
|||
| CancelTaskResult
|
||||
| CreateTaskResult
|
||||
)
|
||||
client_result_adapter = TypeAdapter[ClientResult](ClientResult)
|
||||
|
||||
|
||||
class ClientResult(RootModel[ClientResultType]):
|
||||
pass
|
||||
|
||||
|
||||
ServerRequestType: TypeAlias = (
|
||||
ServerRequest = (
|
||||
PingRequest
|
||||
| CreateMessageRequest
|
||||
| ListRootsRequest
|
||||
|
|
@ -1788,13 +1779,10 @@ ServerRequestType: TypeAlias = (
|
|||
| ListTasksRequest
|
||||
| CancelTaskRequest
|
||||
)
|
||||
server_request_adapter = TypeAdapter[ServerRequest](ServerRequest)
|
||||
|
||||
|
||||
class ServerRequest(RootModel[ServerRequestType]):
|
||||
pass
|
||||
|
||||
|
||||
ServerNotificationType: TypeAlias = (
|
||||
ServerNotification = (
|
||||
CancelledNotification
|
||||
| ProgressNotification
|
||||
| LoggingMessageNotification
|
||||
|
|
@ -1805,13 +1793,10 @@ ServerNotificationType: TypeAlias = (
|
|||
| ElicitCompleteNotification
|
||||
| TaskStatusNotification
|
||||
)
|
||||
server_notification_adapter = TypeAdapter[ServerNotification](ServerNotification)
|
||||
|
||||
|
||||
class ServerNotification(RootModel[ServerNotificationType]):
|
||||
pass
|
||||
|
||||
|
||||
ServerResultType: TypeAlias = (
|
||||
ServerResult = (
|
||||
EmptyResult
|
||||
| InitializeResult
|
||||
| CompleteResult
|
||||
|
|
@ -1828,7 +1813,4 @@ ServerResultType: TypeAlias = (
|
|||
| CancelTaskResult
|
||||
| CreateTaskResult
|
||||
)
|
||||
|
||||
|
||||
class ServerResult(RootModel[ServerResultType]):
|
||||
pass
|
||||
server_result_adapter = TypeAdapter[ServerResult](ServerResult)
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from starlette.routing import Route
|
|||
from mcp import ClientSession, types
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.types import ClientNotification, RootsListChangedNotification
|
||||
from mcp.types import RootsListChangedNotification
|
||||
from tests.test_helpers import wait_for_server
|
||||
|
||||
|
||||
|
|
@ -135,9 +135,7 @@ async def test_non_compliant_notification_response(non_sdk_server: None, non_sdk
|
|||
await session.initialize()
|
||||
|
||||
# The test server returns a 204 instead of the expected 202
|
||||
await session.send_notification(
|
||||
ClientNotification(RootsListChangedNotification(method="notifications/roots/list_changed"))
|
||||
)
|
||||
await session.send_notification(RootsListChangedNotification(method="notifications/roots/list_changed"))
|
||||
|
||||
if returned_exception: # pragma: no cover
|
||||
pytest.fail(f"Server encountered an exception: {returned_exception}")
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from unittest.mock import patch
|
|||
|
||||
import anyio
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import BaseSession, RequestId, SendResultT
|
||||
|
|
@ -23,20 +24,23 @@ async def test_send_request_stream_cleanup():
|
|||
) -> None: # pragma: no cover
|
||||
pass
|
||||
|
||||
@property
|
||||
def _receive_request_adapter(self) -> TypeAdapter[Any]:
|
||||
return TypeAdapter(object) # pragma: no cover
|
||||
|
||||
@property
|
||||
def _receive_notification_adapter(self) -> TypeAdapter[Any]:
|
||||
return TypeAdapter(object) # pragma: no cover
|
||||
|
||||
# Create streams
|
||||
write_stream_send, write_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
read_stream_send, read_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
|
||||
# Create the session
|
||||
session = TestSession(
|
||||
read_stream_receive,
|
||||
write_stream_send,
|
||||
object, # Request type doesn't matter for this test
|
||||
object, # Notification type doesn't matter for this test
|
||||
)
|
||||
session = TestSession(read_stream_receive, write_stream_send)
|
||||
|
||||
# Create a test request
|
||||
request = ClientRequest(PingRequest())
|
||||
request = PingRequest()
|
||||
|
||||
# Patch the _write_stream.send method to raise an exception
|
||||
async def mock_send(*args: Any, **kwargs: Any):
|
||||
|
|
|
|||
|
|
@ -12,8 +12,6 @@ from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
|||
from mcp.types import (
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
CallToolResult,
|
||||
ClientNotification,
|
||||
ClientRequest,
|
||||
Implementation,
|
||||
InitializedNotification,
|
||||
InitializeRequest,
|
||||
|
|
@ -22,8 +20,9 @@ from mcp.types import (
|
|||
JSONRPCRequest,
|
||||
JSONRPCResponse,
|
||||
ServerCapabilities,
|
||||
ServerResult,
|
||||
TextContent,
|
||||
client_notification_adapter,
|
||||
client_request_adapter,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -41,24 +40,22 @@ async def test_client_session_initialize():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
assert isinstance(request, InitializeRequest)
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(
|
||||
logging=None,
|
||||
resources=None,
|
||||
tools=None,
|
||||
experimental=None,
|
||||
prompts=None,
|
||||
),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
instructions="The server instructions.",
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(
|
||||
logging=None,
|
||||
resources=None,
|
||||
tools=None,
|
||||
experimental=None,
|
||||
prompts=None,
|
||||
),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
instructions="The server instructions.",
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -74,7 +71,7 @@ async def test_client_session_initialize():
|
|||
session_notification = await client_to_server_receive.receive()
|
||||
jsonrpc_notification = session_notification.message
|
||||
assert isinstance(jsonrpc_notification, JSONRPCNotification)
|
||||
initialized_notification = ClientNotification.model_validate(
|
||||
initialized_notification = client_notification_adapter.validate_python(
|
||||
jsonrpc_notification.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
|
||||
|
|
@ -109,7 +106,7 @@ async def test_client_session_initialize():
|
|||
|
||||
# Check that the client sent the initialized notification
|
||||
assert initialized_notification
|
||||
assert isinstance(initialized_notification.root, InitializedNotification)
|
||||
assert isinstance(initialized_notification, InitializedNotification)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
@ -126,18 +123,16 @@ async def test_client_session_custom_client_info():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_client_info = request.root.params.client_info
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_client_info = request.params.client_info
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -185,18 +180,16 @@ async def test_client_session_default_client_info():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_client_info = request.root.params.client_info
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_client_info = request.params.client_info
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -238,21 +231,19 @@ async def test_client_session_version_negotiation_success():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
assert isinstance(request, InitializeRequest)
|
||||
|
||||
# Verify client sent the latest protocol version
|
||||
assert request.root.params.protocol_version == LATEST_PROTOCOL_VERSION
|
||||
assert request.params.protocol_version == LATEST_PROTOCOL_VERSION
|
||||
|
||||
# Server responds with a supported older version
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version="2024-11-05",
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version="2024-11-05",
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -295,18 +286,16 @@ async def test_client_session_version_negotiation_failure():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
assert isinstance(request, InitializeRequest)
|
||||
|
||||
# Server responds with an unsupported version
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version="2020-01-01", # Unsupported old version
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version="2020-01-01", # Unsupported old version
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -349,18 +338,16 @@ async def test_client_capabilities_default():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_capabilities = request.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -422,18 +409,16 @@ async def test_client_capabilities_with_custom_callbacks():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_capabilities = request.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -503,18 +488,16 @@ async def test_client_capabilities_with_sampling_tools():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_capabilities = request.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -572,17 +555,15 @@ async def test_get_server_capabilities():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
assert isinstance(request, InitializeRequest)
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=expected_capabilities,
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=expected_capabilities,
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -639,17 +620,15 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
assert isinstance(request, InitializeRequest)
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
# Answer initialization request
|
||||
|
|
@ -678,9 +657,7 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
|
|||
assert "_meta" in jsonrpc_request.params
|
||||
assert jsonrpc_request.params["_meta"] == meta
|
||||
|
||||
result = ServerResult(
|
||||
CallToolResult(content=[TextContent(type="text", text="Called successfully")], is_error=False)
|
||||
)
|
||||
result = CallToolResult(content=[TextContent(type="text", text="Called successfully")], is_error=False)
|
||||
|
||||
# Send the tools/call result
|
||||
await server_to_client_send.send(
|
||||
|
|
|
|||
|
|
@ -11,14 +11,13 @@ from mcp.shared.context import RequestContext
|
|||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
ClientRequest,
|
||||
Implementation,
|
||||
InitializeRequest,
|
||||
InitializeResult,
|
||||
JSONRPCRequest,
|
||||
JSONRPCResponse,
|
||||
ServerCapabilities,
|
||||
ServerResult,
|
||||
client_request_adapter,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -36,18 +35,16 @@ async def test_client_capabilities_without_tasks():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_capabilities = request.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -108,18 +105,16 @@ async def test_client_capabilities_with_tasks():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_capabilities = request.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -190,18 +185,16 @@ async def test_client_capabilities_auto_built_from_handlers():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_capabilities = request.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
@ -268,18 +261,16 @@ async def test_client_capabilities_with_task_augmented_handlers():
|
|||
session_message = await client_to_server_receive.receive()
|
||||
jsonrpc_request = session_message.message
|
||||
assert isinstance(jsonrpc_request, JSONRPCRequest)
|
||||
request = ClientRequest.model_validate(
|
||||
request = client_request_adapter.validate_python(
|
||||
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
)
|
||||
assert isinstance(request.root, InitializeRequest)
|
||||
received_capabilities = request.root.params.capabilities
|
||||
assert isinstance(request, InitializeRequest)
|
||||
received_capabilities = request.params.capabilities
|
||||
|
||||
result = ServerResult(
|
||||
InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
result = InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
)
|
||||
|
||||
async with server_to_client_send:
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue