Coverage for mcpgateway/models.py: 99%

206 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2025-07-09 11:03 +0100

1# -*- coding: utf-8 -*- 

2"""MCP Protocol Type Definitions. 

3 

4Copyright 2025 

5SPDX-License-Identifier: Apache-2.0 

6Authors: Mihai Criveti 

7 

8This module defines all core MCP protocol types according to the specification. 

9It includes: 

10 - Message content types (text, image, resource) 

11 - Tool definitions and schemas 

12 - Resource types and templates 

13 - Prompt structures 

14 - Protocol initialization types 

15 - Sampling message types 

16 - Capability definitions 

17""" 

18 

19# Standard 

20from datetime import datetime 

21from enum import Enum 

22from typing import Any, Dict, List, Literal, Optional, Union 

23 

24# Third-Party 

25from pydantic import AnyHttpUrl, AnyUrl, BaseModel, ConfigDict, Field 

26 

27 

28class Role(str, Enum): 

29 """Message role in conversations. 

30 

31 Attributes: 

32 ASSISTANT (str): Indicates the assistant's role. 

33 USER (str): Indicates the user's role. 

34 """ 

35 

36 ASSISTANT = "assistant" 

37 USER = "user" 

38 

39 

40class LogLevel(str, Enum): 

41 """Standard syslog severity levels as defined in RFC 5424. 

42 

43 Attributes: 

44 DEBUG (str): Debug level. 

45 INFO (str): Informational level. 

46 NOTICE (str): Notice level. 

47 WARNING (str): Warning level. 

48 ERROR (str): Error level. 

49 CRITICAL (str): Critical level. 

50 ALERT (str): Alert level. 

51 EMERGENCY (str): Emergency level. 

52 """ 

53 

54 DEBUG = "debug" 

55 INFO = "info" 

56 NOTICE = "notice" 

57 WARNING = "warning" 

58 ERROR = "error" 

59 CRITICAL = "critical" 

60 ALERT = "alert" 

61 EMERGENCY = "emergency" 

62 

63 

64# Base content types 

65class TextContent(BaseModel): 

66 """Text content for messages. 

67 

68 Attributes: 

69 type (Literal["text"]): The fixed content type identifier for text. 

70 text (str): The actual text message. 

71 """ 

72 

73 type: Literal["text"] 

74 text: str 

75 

76 

77class JSONContent(BaseModel): 

78 """JSON content for messages. 

79 Attributes: 

80 type (Literal["text"]): The fixed content type identifier for text. 

81 json (dict): The actual text message. 

82 """ 

83 

84 type: Literal["text"] 

85 text: dict 

86 

87 

88class ImageContent(BaseModel): 

89 """Image content for messages. 

90 

91 Attributes: 

92 type (Literal["image"]): The fixed content type identifier for images. 

93 data (bytes): The binary data of the image. 

94 mime_type (str): The MIME type (e.g. "image/png") of the image. 

95 """ 

96 

97 type: Literal["image"] 

98 data: bytes 

99 mime_type: str 

100 

101 

102class ResourceContent(BaseModel): 

103 """Resource content that can be embedded. 

104 

105 Attributes: 

106 type (Literal["resource"]): The fixed content type identifier for resources. 

107 uri (str): The URI identifying the resource. 

108 mime_type (Optional[str]): The MIME type of the resource, if known. 

109 text (Optional[str]): A textual representation of the resource, if applicable. 

110 blob (Optional[bytes]): Binary data of the resource, if applicable. 

111 """ 

112 

113 type: Literal["resource"] 

114 uri: str 

115 mime_type: Optional[str] = None 

116 text: Optional[str] = None 

117 blob: Optional[bytes] = None 

118 

119 

120ContentType = Union[TextContent, JSONContent, ImageContent, ResourceContent] 

121 

122 

123# Reference types - needed early for completion 

124class PromptReference(BaseModel): 

125 """Reference to a prompt or prompt template. 

126 

127 Attributes: 

128 type (Literal["ref/prompt"]): The fixed reference type identifier for prompts. 

129 name (str): The unique name of the prompt. 

130 """ 

131 

132 type: Literal["ref/prompt"] 

133 name: str 

134 

135 

136class ResourceReference(BaseModel): 

137 """Reference to a resource or resource template. 

138 

139 Attributes: 

140 type (Literal["ref/resource"]): The fixed reference type identifier for resources. 

141 uri (str): The URI of the resource. 

142 """ 

143 

144 type: Literal["ref/resource"] 

145 uri: str 

146 

147 

148# Completion types 

149class CompleteRequest(BaseModel): 

150 """Request for completion suggestions. 

151 

152 Attributes: 

153 ref (Union[PromptReference, ResourceReference]): A reference to a prompt or resource. 

154 argument (Dict[str, str]): A dictionary containing arguments for the completion. 

155 """ 

156 

157 ref: Union[PromptReference, ResourceReference] 

158 argument: Dict[str, str] 

159 

160 

161class CompleteResult(BaseModel): 

162 """Result for a completion request. 

163 

164 Attributes: 

165 completion (Dict[str, Any]): A dictionary containing the completion results. 

166 """ 

167 

168 completion: Dict[str, Any] = Field(..., description="Completion results") 

169 

170 

171# Implementation info 

172class Implementation(BaseModel): 

173 """MCP implementation information. 

174 

175 Attributes: 

176 name (str): The name of the implementation. 

177 version (str): The version of the implementation. 

178 """ 

179 

180 name: str 

181 version: str 

182 

183 

184# Model preferences 

185class ModelHint(BaseModel): 

186 """Hint for model selection. 

187 

188 Attributes: 

189 name (Optional[str]): An optional hint for the model name. 

190 """ 

191 

192 name: Optional[str] = None 

193 

194 

195class ModelPreferences(BaseModel): 

196 """Server preferences for model selection. 

197 

198 Attributes: 

199 cost_priority (float): Priority for cost efficiency (0 to 1). 

200 speed_priority (float): Priority for speed (0 to 1). 

201 intelligence_priority (float): Priority for intelligence (0 to 1). 

202 hints (List[ModelHint]): A list of model hints. 

203 """ 

204 

205 cost_priority: float = Field(ge=0, le=1) 

206 speed_priority: float = Field(ge=0, le=1) 

207 intelligence_priority: float = Field(ge=0, le=1) 

208 hints: List[ModelHint] = [] 

209 

210 

211# Capability types 

212class ClientCapabilities(BaseModel): 

213 """Capabilities that a client may support. 

214 

215 Attributes: 

216 roots (Optional[Dict[str, bool]]): Capabilities related to root management. 

217 sampling (Optional[Dict[str, Any]]): Capabilities related to LLM sampling. 

218 experimental (Optional[Dict[str, Dict[str, Any]]]): Experimental capabilities. 

219 """ 

220 

221 roots: Optional[Dict[str, bool]] = None 

222 sampling: Optional[Dict[str, Any]] = None 

223 experimental: Optional[Dict[str, Dict[str, Any]]] = None 

224 

225 

226class ServerCapabilities(BaseModel): 

227 """Capabilities that a server may support. 

228 

229 Attributes: 

230 prompts (Optional[Dict[str, bool]]): Capability for prompt support. 

231 resources (Optional[Dict[str, bool]]): Capability for resource support. 

232 tools (Optional[Dict[str, bool]]): Capability for tool support. 

233 logging (Optional[Dict[str, Any]]): Capability for logging support. 

234 experimental (Optional[Dict[str, Dict[str, Any]]]): Experimental capabilities. 

235 """ 

236 

237 prompts: Optional[Dict[str, bool]] = None 

238 resources: Optional[Dict[str, bool]] = None 

239 tools: Optional[Dict[str, bool]] = None 

240 logging: Optional[Dict[str, Any]] = None 

241 experimental: Optional[Dict[str, Dict[str, Any]]] = None 

242 

243 

244# Initialization types 

245class InitializeRequest(BaseModel): 

246 """Initial request sent from the client to the server. 

247 

248 Attributes: 

249 protocol_version (str): The protocol version (alias: protocolVersion). 

250 capabilities (ClientCapabilities): The client's capabilities. 

251 client_info (Implementation): The client's implementation information (alias: clientInfo). 

252 

253 Note: 

254 The alias settings allow backward compatibility with older Pydantic versions. 

255 """ 

256 

257 protocol_version: str = Field(..., alias="protocolVersion") 

258 capabilities: ClientCapabilities 

259 client_info: Implementation = Field(..., alias="clientInfo") 

260 

261 model_config = ConfigDict( 

262 populate_by_name=True, 

263 ) 

264 

265 

266class InitializeResult(BaseModel): 

267 """Server's response to the initialization request. 

268 

269 Attributes: 

270 protocol_version (str): The protocol version used. 

271 capabilities (ServerCapabilities): The server's capabilities. 

272 server_info (Implementation): The server's implementation information. 

273 instructions (Optional[str]): Optional instructions for the client. 

274 """ 

275 

276 protocol_version: str = Field(..., alias="protocolVersion") 

277 capabilities: ServerCapabilities = Field(..., alias="capabilities") 

278 server_info: Implementation = Field(..., alias="serverInfo") 

279 instructions: Optional[str] = Field(None, alias="instructions") 

280 

281 model_config = ConfigDict( 

282 populate_by_name=True, 

283 ) 

284 

285 

286# Message types 

287class Message(BaseModel): 

288 """A message in a conversation. 

289 

290 Attributes: 

291 role (Role): The role of the message sender. 

292 content (ContentType): The content of the message. 

293 """ 

294 

295 role: Role 

296 content: ContentType 

297 

298 

299class SamplingMessage(BaseModel): 

300 """A message used in LLM sampling requests. 

301 

302 Attributes: 

303 role (Role): The role of the sender. 

304 content (ContentType): The content of the sampling message. 

305 """ 

306 

307 role: Role 

308 content: ContentType 

309 

310 

311# Sampling types for the client features 

312class CreateMessageResult(BaseModel): 

313 """Result from a sampling/createMessage request. 

314 

315 Attributes: 

316 content (Union[TextContent, ImageContent]): The generated content. 

317 model (str): The model used for generating the content. 

318 role (Role): The role associated with the content. 

319 stop_reason (Optional[str]): An optional reason for why sampling stopped. 

320 """ 

321 

322 content: Union[TextContent, ImageContent] 

323 model: str 

324 role: Role 

325 stop_reason: Optional[str] = None 

326 

327 

328# Prompt types 

329class PromptArgument(BaseModel): 

330 """An argument that can be passed to a prompt. 

331 

332 Attributes: 

333 name (str): The name of the argument. 

334 description (Optional[str]): An optional description of the argument. 

335 required (bool): Whether the argument is required. Defaults to False. 

336 """ 

337 

338 name: str 

339 description: Optional[str] = None 

340 required: bool = False 

341 

342 

343class Prompt(BaseModel): 

344 """A prompt template offered by the server. 

345 

346 Attributes: 

347 name (str): The unique name of the prompt. 

348 description (Optional[str]): A description of the prompt. 

349 arguments (List[PromptArgument]): A list of expected prompt arguments. 

350 """ 

351 

352 name: str 

353 description: Optional[str] = None 

354 arguments: List[PromptArgument] = [] 

355 

356 

357class PromptResult(BaseModel): 

358 """Result of rendering a prompt template. 

359 

360 Attributes: 

361 messages (List[Message]): The list of messages produced by rendering the prompt. 

362 description (Optional[str]): An optional description of the rendered result. 

363 """ 

364 

365 messages: List[Message] 

366 description: Optional[str] = None 

367 

368 

369# Tool types 

370class Tool(BaseModel): 

371 """A tool that can be invoked. 

372 

373 Attributes: 

374 name (str): The unique name of the tool. 

375 url (AnyHttpUrl): The URL of the tool. 

376 description (Optional[str]): A description of the tool. 

377 integrationType (str): The integration type of the tool (e.g. MCP or REST). 

378 requestType (str): The HTTP method used to invoke the tool (GET, POST, PUT, DELETE, SSE, STDIO). 

379 headers (Dict[str, Any]): A JSON object representing HTTP headers. 

380 input_schema (Dict[str, Any]): A JSON Schema for validating the tool's input. 

381 annotations (Optional[Dict[str, Any]]): Tool annotations for behavior hints. 

382 auth_type (Optional[str]): The type of authentication used ("basic", "bearer", or None). 

383 auth_username (Optional[str]): The username for basic authentication. 

384 auth_password (Optional[str]): The password for basic authentication. 

385 auth_token (Optional[str]): The token for bearer authentication. 

386 """ 

387 

388 name: str 

389 url: AnyHttpUrl 

390 description: Optional[str] = None 

391 integrationType: str = "MCP" 

392 requestType: str = "SSE" 

393 headers: Dict[str, Any] = Field(default_factory=dict) 

394 input_schema: Dict[str, Any] = Field(default_factory=lambda: {"type": "object", "properties": {}}) 

395 annotations: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Tool annotations for behavior hints") 

396 auth_type: Optional[str] = None 

397 auth_username: Optional[str] = None 

398 auth_password: Optional[str] = None 

399 auth_token: Optional[str] = None 

400 

401 

402class ToolResult(BaseModel): 

403 """Result of a tool invocation. 

404 

405 Attributes: 

406 content (List[ContentType]): A list of content items returned by the tool. 

407 is_error (bool): Flag indicating if the tool call resulted in an error. 

408 """ 

409 

410 content: List[ContentType] 

411 is_error: bool = False 

412 

413 

414# Resource types 

415class Resource(BaseModel): 

416 """A resource available from the server. 

417 

418 Attributes: 

419 uri (str): The unique URI of the resource. 

420 name (str): The human-readable name of the resource. 

421 description (Optional[str]): A description of the resource. 

422 mime_type (Optional[str]): The MIME type of the resource. 

423 size (Optional[int]): The size of the resource. 

424 """ 

425 

426 uri: str 

427 name: str 

428 description: Optional[str] = None 

429 mime_type: Optional[str] = None 

430 size: Optional[int] = None 

431 

432 

433class ResourceTemplate(BaseModel): 

434 """A template for constructing resource URIs. 

435 

436 Attributes: 

437 uri_template (str): The URI template string. 

438 name (str): The unique name of the template. 

439 description (Optional[str]): A description of the template. 

440 mime_type (Optional[str]): The MIME type associated with the template. 

441 """ 

442 

443 uri_template: str 

444 name: str 

445 description: Optional[str] = None 

446 mime_type: Optional[str] = None 

447 

448 

449class ListResourceTemplatesResult(BaseModel): 

450 """The server's response to a resources/templates/list request from the client. 

451 

452 Attributes: 

453 meta (Optional[Dict[str, Any]]): Reserved property for metadata. 

454 next_cursor (Optional[str]): Pagination cursor for the next page of results. 

455 resource_templates (List[ResourceTemplate]): List of resource templates. 

456 """ 

457 

458 meta: Optional[Dict[str, Any]] = Field( 

459 None, alias="_meta", description="This result property is reserved by the protocol to allow clients and servers to attach additional metadata to their responses." 

460 ) 

461 next_cursor: Optional[str] = Field(None, description="An opaque token representing the pagination position after the last returned result.\nIf present, there may be more results available.") 

462 resource_templates: List[ResourceTemplate] = Field(default_factory=list, description="List of resource templates available on the server") 

463 

464 model_config = ConfigDict( 

465 populate_by_name=True, 

466 ) 

467 

468 

469# Root types 

470class FileUrl(AnyUrl): 

471 """A specialized URL type for local file-scheme resources. 

472 

473 Key characteristics 

474 ------------------- 

475 * Scheme restricted - only the "file" scheme is permitted 

476 (e.g. file:///path/to/file.txt). 

477 * No host required - "file" URLs typically omit a network host; 

478 therefore, the host component is not mandatory. 

479 * String-friendly equality - developers naturally expect 

480 FileUrl("file:///data") == "file:///data" to evaluate True. 

481 AnyUrl (Pydantic) does not implement that, so we override 

482 __eq__ to compare against plain strings transparently. 

483 Hash semantics are kept consistent by delegating to the parent class. 

484 

485 Examples 

486 -------- 

487 >>> url = FileUrl("file:///etc/hosts") 

488 >>> url.scheme 

489 'file' 

490 >>> url == "file:///etc/hosts" 

491 True 

492 >>> {"path": url} # hashable 

493 {'path': FileUrl('file:///etc/hosts')} 

494 

495 Notes 

496 ----- 

497 The override does not interfere with comparisons to other 

498 AnyUrl/FileUrl instances; those still use the superclass 

499 implementation. 

500 """ 

501 

502 # Restrict to the "file" scheme and omit host requirement 

503 allowed_schemes = {"file"} 

504 host_required = False 

505 

506 def __eq__(self, other): # type: ignore[override] 

507 """Return True when other is an equivalent URL or string. 

508 

509 If other is a str it is coerced with str(self) for comparison; 

510 otherwise defer to AnyUrl's comparison. 

511 

512 Args: 

513 other (Any): The object to compare against. May be a str, FileUrl, or AnyUrl. 

514 

515 Returns: 

516 bool: True if the other value is equal to this URL, either as a string 

517 or as another URL object. False otherwise. 

518 """ 

519 if isinstance(other, str): 519 ↛ 521line 519 didn't jump to line 521 because the condition on line 519 was always true

520 return str(self) == other 

521 return super().__eq__(other) 

522 

523 # Keep hashing behaviour aligned with equality 

524 __hash__ = AnyUrl.__hash__ 

525 

526 

527class Root(BaseModel): 

528 """A root directory or file. 

529 

530 Attributes: 

531 uri (Union[FileUrl, AnyUrl]): The unique identifier for the root. 

532 name (Optional[str]): An optional human-readable name. 

533 """ 

534 

535 uri: Union[FileUrl, AnyUrl] = Field(..., description="Unique identifier for the root") 

536 name: Optional[str] = Field(None, description="Optional human-readable name") 

537 

538 

539# Progress types 

540class ProgressToken(BaseModel): 

541 """Token for associating progress notifications. 

542 

543 Attributes: 

544 value (Union[str, int]): The token value. 

545 """ 

546 

547 value: Union[str, int] 

548 

549 

550class Progress(BaseModel): 

551 """Progress update for long-running operations. 

552 

553 Attributes: 

554 progress_token (ProgressToken): The token associated with the progress update. 

555 progress (float): The current progress value. 

556 total (Optional[float]): The total progress value, if known. 

557 """ 

558 

559 progress_token: ProgressToken 

560 progress: float 

561 total: Optional[float] = None 

562 

563 

564# JSON-RPC types 

565class JSONRPCRequest(BaseModel): 

566 """JSON-RPC 2.0 request. 

567 

568 Attributes: 

569 jsonrpc (Literal["2.0"]): The JSON-RPC version. 

570 id (Optional[Union[str, int]]): The request identifier. 

571 method (str): The method name. 

572 params (Optional[Dict[str, Any]]): The parameters for the request. 

573 """ 

574 

575 jsonrpc: Literal["2.0"] 

576 id: Optional[Union[str, int]] = None 

577 method: str 

578 params: Optional[Dict[str, Any]] = None 

579 

580 

581class JSONRPCResponse(BaseModel): 

582 """JSON-RPC 2.0 response. 

583 

584 Attributes: 

585 jsonrpc (Literal["2.0"]): The JSON-RPC version. 

586 id (Optional[Union[str, int]]): The request identifier. 

587 result (Optional[Any]): The result of the request. 

588 error (Optional[Dict[str, Any]]): The error object if an error occurred. 

589 """ 

590 

591 jsonrpc: Literal["2.0"] 

592 id: Optional[Union[str, int]] = None 

593 result: Optional[Any] = None 

594 error: Optional[Dict[str, Any]] = None 

595 

596 

597class JSONRPCError(BaseModel): 

598 """JSON-RPC 2.0 error. 

599 

600 Attributes: 

601 code (int): The error code. 

602 message (str): A short description of the error. 

603 data (Optional[Any]): Additional data about the error. 

604 """ 

605 

606 code: int 

607 message: str 

608 data: Optional[Any] = None 

609 

610 

611# Transport message types 

612class SSEEvent(BaseModel): 

613 """Server-Sent Events message. 

614 

615 Attributes: 

616 id (Optional[str]): The event identifier. 

617 event (Optional[str]): The event type. 

618 data (str): The event data. 

619 retry (Optional[int]): The retry timeout in milliseconds. 

620 """ 

621 

622 id: Optional[str] = None 

623 event: Optional[str] = None 

624 data: str 

625 retry: Optional[int] = None 

626 

627 

628class WebSocketMessage(BaseModel): 

629 """WebSocket protocol message. 

630 

631 Attributes: 

632 type (str): The type of the WebSocket message. 

633 data (Any): The message data. 

634 """ 

635 

636 type: str 

637 data: Any 

638 

639 

640# Notification types 

641class ResourceUpdateNotification(BaseModel): 

642 """Notification of resource changes. 

643 

644 Attributes: 

645 method (Literal["notifications/resources/updated"]): The notification method. 

646 uri (str): The URI of the updated resource. 

647 """ 

648 

649 method: Literal["notifications/resources/updated"] 

650 uri: str 

651 

652 

653class ResourceListChangedNotification(BaseModel): 

654 """Notification of resource list changes. 

655 

656 Attributes: 

657 method (Literal["notifications/resources/list_changed"]): The notification method. 

658 """ 

659 

660 method: Literal["notifications/resources/list_changed"] 

661 

662 

663class PromptListChangedNotification(BaseModel): 

664 """Notification of prompt list changes. 

665 

666 Attributes: 

667 method (Literal["notifications/prompts/list_changed"]): The notification method. 

668 """ 

669 

670 method: Literal["notifications/prompts/list_changed"] 

671 

672 

673class ToolListChangedNotification(BaseModel): 

674 """Notification of tool list changes. 

675 

676 Attributes: 

677 method (Literal["notifications/tools/list_changed"]): The notification method. 

678 """ 

679 

680 method: Literal["notifications/tools/list_changed"] 

681 

682 

683class CancelledNotification(BaseModel): 

684 """Notification of request cancellation. 

685 

686 Attributes: 

687 method (Literal["notifications/cancelled"]): The notification method. 

688 request_id (Union[str, int]): The ID of the cancelled request. 

689 reason (Optional[str]): An optional reason for cancellation. 

690 """ 

691 

692 method: Literal["notifications/cancelled"] 

693 request_id: Union[str, int] 

694 reason: Optional[str] = None 

695 

696 

697class ProgressNotification(BaseModel): 

698 """Notification of operation progress. 

699 

700 Attributes: 

701 method (Literal["notifications/progress"]): The notification method. 

702 progress_token (ProgressToken): The token associated with the progress. 

703 progress (float): The current progress value. 

704 total (Optional[float]): The total progress value, if known. 

705 """ 

706 

707 method: Literal["notifications/progress"] 

708 progress_token: ProgressToken 

709 progress: float 

710 total: Optional[float] = None 

711 

712 

713class LoggingNotification(BaseModel): 

714 """Notification of log messages. 

715 

716 Attributes: 

717 method (Literal["notifications/message"]): The notification method. 

718 level (LogLevel): The log level of the message. 

719 logger (Optional[str]): The logger name. 

720 data (Any): The log message data. 

721 """ 

722 

723 method: Literal["notifications/message"] 

724 level: LogLevel 

725 logger: Optional[str] = None 

726 data: Any 

727 

728 

729# Federation types 

730class FederatedTool(Tool): 

731 """A tool from a federated gateway. 

732 

733 Attributes: 

734 gateway_id (str): The identifier of the gateway. 

735 gateway_name (str): The name of the gateway. 

736 """ 

737 

738 gateway_id: str 

739 gateway_name: str 

740 

741 

742class FederatedResource(Resource): 

743 """A resource from a federated gateway. 

744 

745 Attributes: 

746 gateway_id (str): The identifier of the gateway. 

747 gateway_name (str): The name of the gateway. 

748 """ 

749 

750 gateway_id: str 

751 gateway_name: str 

752 

753 

754class FederatedPrompt(Prompt): 

755 """A prompt from a federated gateway. 

756 

757 Attributes: 

758 gateway_id (str): The identifier of the gateway. 

759 gateway_name (str): The name of the gateway. 

760 """ 

761 

762 gateway_id: str 

763 gateway_name: str 

764 

765 

766class Gateway(BaseModel): 

767 """A federated gateway peer. 

768 

769 Attributes: 

770 id (str): The unique identifier for the gateway. 

771 name (str): The name of the gateway. 

772 url (AnyHttpUrl): The URL of the gateway. 

773 capabilities (ServerCapabilities): The capabilities of the gateway. 

774 last_seen (Optional[datetime]): Timestamp when the gateway was last seen. 

775 """ 

776 

777 id: str 

778 name: str 

779 url: AnyHttpUrl 

780 capabilities: ServerCapabilities 

781 last_seen: Optional[datetime] = None