|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +from collections.abc import Iterable |
| 4 | +from typing import TYPE_CHECKING, Any, Generic, Literal |
| 5 | + |
| 6 | +from pydantic import AnyUrl, BaseModel |
| 7 | + |
| 8 | +from mcp.server.context import LifespanContextT, RequestT, ServerRequestContext |
| 9 | +from mcp.server.elicitation import ( |
| 10 | + ElicitationResult, |
| 11 | + ElicitSchemaModelT, |
| 12 | + UrlElicitationResult, |
| 13 | + elicit_with_validation, |
| 14 | +) |
| 15 | +from mcp.server.elicitation import ( |
| 16 | + elicit_url as _elicit_url, |
| 17 | +) |
| 18 | +from mcp.server.lowlevel.helper_types import ReadResourceContents |
| 19 | + |
| 20 | +if TYPE_CHECKING: |
| 21 | + from mcp.server.mcpserver.server import MCPServer |
| 22 | + |
| 23 | + |
| 24 | +class Context(BaseModel, Generic[LifespanContextT, RequestT]): |
| 25 | + """Context object providing access to MCP capabilities. |
| 26 | +
|
| 27 | + This provides a cleaner interface to MCP's RequestContext functionality. |
| 28 | + It gets injected into tool and resource functions that request it via type hints. |
| 29 | +
|
| 30 | + To use context in a tool function, add a parameter with the Context type annotation: |
| 31 | +
|
| 32 | + ```python |
| 33 | + @server.tool() |
| 34 | + async def my_tool(x: int, ctx: Context) -> str: |
| 35 | + # Log messages to the client |
| 36 | + await ctx.info(f"Processing {x}") |
| 37 | + await ctx.debug("Debug info") |
| 38 | + await ctx.warning("Warning message") |
| 39 | + await ctx.error("Error message") |
| 40 | +
|
| 41 | + # Report progress |
| 42 | + await ctx.report_progress(50, 100) |
| 43 | +
|
| 44 | + # Access resources |
| 45 | + data = await ctx.read_resource("resource://data") |
| 46 | +
|
| 47 | + # Get request info |
| 48 | + request_id = ctx.request_id |
| 49 | + client_id = ctx.client_id |
| 50 | +
|
| 51 | + return str(x) |
| 52 | + ``` |
| 53 | +
|
| 54 | + The context parameter name can be anything as long as it's annotated with Context. |
| 55 | + The context is optional - tools that don't need it can omit the parameter. |
| 56 | + """ |
| 57 | + |
| 58 | + _request_context: ServerRequestContext[LifespanContextT, RequestT] | None |
| 59 | + _mcp_server: MCPServer | None |
| 60 | + |
| 61 | + def __init__( |
| 62 | + self, |
| 63 | + *, |
| 64 | + request_context: ServerRequestContext[LifespanContextT, RequestT] | None = None, |
| 65 | + mcp_server: MCPServer | None = None, |
| 66 | + # TODO(Marcelo): We should drop this kwargs parameter. |
| 67 | + **kwargs: Any, |
| 68 | + ): |
| 69 | + super().__init__(**kwargs) |
| 70 | + self._request_context = request_context |
| 71 | + self._mcp_server = mcp_server |
| 72 | + |
| 73 | + @property |
| 74 | + def mcp_server(self) -> MCPServer: |
| 75 | + """Access to the MCPServer instance.""" |
| 76 | + if self._mcp_server is None: # pragma: no cover |
| 77 | + raise ValueError("Context is not available outside of a request") |
| 78 | + return self._mcp_server # pragma: no cover |
| 79 | + |
| 80 | + @property |
| 81 | + def request_context(self) -> ServerRequestContext[LifespanContextT, RequestT]: |
| 82 | + """Access to the underlying request context.""" |
| 83 | + if self._request_context is None: # pragma: no cover |
| 84 | + raise ValueError("Context is not available outside of a request") |
| 85 | + return self._request_context |
| 86 | + |
| 87 | + async def report_progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: |
| 88 | + """Report progress for the current operation. |
| 89 | +
|
| 90 | + Args: |
| 91 | + progress: Current progress value (e.g., 24) |
| 92 | + total: Optional total value (e.g., 100) |
| 93 | + message: Optional message (e.g., "Starting render...") |
| 94 | + """ |
| 95 | + progress_token = self.request_context.meta.get("progress_token") if self.request_context.meta else None |
| 96 | + |
| 97 | + if progress_token is None: # pragma: no cover |
| 98 | + return |
| 99 | + |
| 100 | + await self.request_context.session.send_progress_notification( |
| 101 | + progress_token=progress_token, |
| 102 | + progress=progress, |
| 103 | + total=total, |
| 104 | + message=message, |
| 105 | + related_request_id=self.request_id, |
| 106 | + ) |
| 107 | + |
| 108 | + async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]: |
| 109 | + """Read a resource by URI. |
| 110 | +
|
| 111 | + Args: |
| 112 | + uri: Resource URI to read |
| 113 | +
|
| 114 | + Returns: |
| 115 | + The resource content as either text or bytes |
| 116 | + """ |
| 117 | + assert self._mcp_server is not None, "Context is not available outside of a request" |
| 118 | + return await self._mcp_server.read_resource(uri, self) |
| 119 | + |
| 120 | + async def elicit( |
| 121 | + self, |
| 122 | + message: str, |
| 123 | + schema: type[ElicitSchemaModelT], |
| 124 | + ) -> ElicitationResult[ElicitSchemaModelT]: |
| 125 | + """Elicit information from the client/user. |
| 126 | +
|
| 127 | + This method can be used to interactively ask for additional information from the |
| 128 | + client within a tool's execution. The client might display the message to the |
| 129 | + user and collect a response according to the provided schema. If the client |
| 130 | + is an agent, it might decide how to handle the elicitation -- either by asking |
| 131 | + the user or automatically generating a response. |
| 132 | +
|
| 133 | + Args: |
| 134 | + message: Message to present to the user |
| 135 | + schema: A Pydantic model class defining the expected response structure. |
| 136 | + According to the specification, only primitive types are allowed. |
| 137 | +
|
| 138 | + Returns: |
| 139 | + An ElicitationResult containing the action taken and the data if accepted |
| 140 | +
|
| 141 | + Note: |
| 142 | + Check the result.action to determine if the user accepted, declined, or cancelled. |
| 143 | + The result.data will only be populated if action is "accept" and validation succeeded. |
| 144 | + """ |
| 145 | + |
| 146 | + return await elicit_with_validation( |
| 147 | + session=self.request_context.session, |
| 148 | + message=message, |
| 149 | + schema=schema, |
| 150 | + related_request_id=self.request_id, |
| 151 | + ) |
| 152 | + |
| 153 | + async def elicit_url( |
| 154 | + self, |
| 155 | + message: str, |
| 156 | + url: str, |
| 157 | + elicitation_id: str, |
| 158 | + ) -> UrlElicitationResult: |
| 159 | + """Request URL mode elicitation from the client. |
| 160 | +
|
| 161 | + This directs the user to an external URL for out-of-band interactions |
| 162 | + that must not pass through the MCP client. Use this for: |
| 163 | + - Collecting sensitive credentials (API keys, passwords) |
| 164 | + - OAuth authorization flows with third-party services |
| 165 | + - Payment and subscription flows |
| 166 | + - Any interaction where data should not pass through the LLM context |
| 167 | +
|
| 168 | + The response indicates whether the user consented to navigate to the URL. |
| 169 | + The actual interaction happens out-of-band. When the elicitation completes, |
| 170 | + call `ctx.session.send_elicit_complete(elicitation_id)` to notify the client. |
| 171 | +
|
| 172 | + Args: |
| 173 | + message: Human-readable explanation of why the interaction is needed |
| 174 | + url: The URL the user should navigate to |
| 175 | + elicitation_id: Unique identifier for tracking this elicitation |
| 176 | +
|
| 177 | + Returns: |
| 178 | + UrlElicitationResult indicating accept, decline, or cancel |
| 179 | + """ |
| 180 | + return await _elicit_url( |
| 181 | + session=self.request_context.session, |
| 182 | + message=message, |
| 183 | + url=url, |
| 184 | + elicitation_id=elicitation_id, |
| 185 | + related_request_id=self.request_id, |
| 186 | + ) |
| 187 | + |
| 188 | + async def log( |
| 189 | + self, |
| 190 | + level: Literal["debug", "info", "warning", "error"], |
| 191 | + message: str, |
| 192 | + *, |
| 193 | + logger_name: str | None = None, |
| 194 | + extra: dict[str, Any] | None = None, |
| 195 | + ) -> None: |
| 196 | + """Send a log message to the client. |
| 197 | +
|
| 198 | + Args: |
| 199 | + level: Log level (debug, info, warning, error) |
| 200 | + message: Log message |
| 201 | + logger_name: Optional logger name |
| 202 | + extra: Optional dictionary with additional structured data to include |
| 203 | + """ |
| 204 | + |
| 205 | + if extra: |
| 206 | + log_data = {"message": message, **extra} |
| 207 | + else: |
| 208 | + log_data = message |
| 209 | + |
| 210 | + await self.request_context.session.send_log_message( |
| 211 | + level=level, |
| 212 | + data=log_data, |
| 213 | + logger=logger_name, |
| 214 | + related_request_id=self.request_id, |
| 215 | + ) |
| 216 | + |
| 217 | + @property |
| 218 | + def client_id(self) -> str | None: |
| 219 | + """Get the client ID if available.""" |
| 220 | + return self.request_context.meta.get("client_id") if self.request_context.meta else None # pragma: no cover |
| 221 | + |
| 222 | + @property |
| 223 | + def request_id(self) -> str: |
| 224 | + """Get the unique ID for this request.""" |
| 225 | + return str(self.request_context.request_id) |
| 226 | + |
| 227 | + @property |
| 228 | + def session(self): |
| 229 | + """Access to the underlying session for advanced usage.""" |
| 230 | + return self.request_context.session |
| 231 | + |
| 232 | + async def close_sse_stream(self) -> None: |
| 233 | + """Close the SSE stream to trigger client reconnection. |
| 234 | +
|
| 235 | + This method closes the HTTP connection for the current request, triggering |
| 236 | + client reconnection. Events continue to be stored in the event store and will |
| 237 | + be replayed when the client reconnects with Last-Event-ID. |
| 238 | +
|
| 239 | + Use this to implement polling behavior during long-running operations - |
| 240 | + the client will reconnect after the retry interval specified in the priming event. |
| 241 | +
|
| 242 | + Note: |
| 243 | + This is a no-op if not using StreamableHTTP transport with event_store. |
| 244 | + The callback is only available when event_store is configured. |
| 245 | + """ |
| 246 | + if self._request_context and self._request_context.close_sse_stream: # pragma: no cover |
| 247 | + await self._request_context.close_sse_stream() |
| 248 | + |
| 249 | + async def close_standalone_sse_stream(self) -> None: |
| 250 | + """Close the standalone GET SSE stream to trigger client reconnection. |
| 251 | +
|
| 252 | + This method closes the HTTP connection for the standalone GET stream used |
| 253 | + for unsolicited server-to-client notifications. The client SHOULD reconnect |
| 254 | + with Last-Event-ID to resume receiving notifications. |
| 255 | +
|
| 256 | + Note: |
| 257 | + This is a no-op if not using StreamableHTTP transport with event_store. |
| 258 | + Currently, client reconnection for standalone GET streams is NOT |
| 259 | + implemented - this is a known gap. |
| 260 | + """ |
| 261 | + if self._request_context and self._request_context.close_standalone_sse_stream: # pragma: no cover |
| 262 | + await self._request_context.close_standalone_sse_stream() |
| 263 | + |
| 264 | + # Convenience methods for common log levels |
| 265 | + async def debug(self, message: str, *, logger_name: str | None = None, extra: dict[str, Any] | None = None) -> None: |
| 266 | + """Send a debug log message.""" |
| 267 | + await self.log("debug", message, logger_name=logger_name, extra=extra) |
| 268 | + |
| 269 | + async def info(self, message: str, *, logger_name: str | None = None, extra: dict[str, Any] | None = None) -> None: |
| 270 | + """Send an info log message.""" |
| 271 | + await self.log("info", message, logger_name=logger_name, extra=extra) |
| 272 | + |
| 273 | + async def warning( |
| 274 | + self, message: str, *, logger_name: str | None = None, extra: dict[str, Any] | None = None |
| 275 | + ) -> None: |
| 276 | + """Send a warning log message.""" |
| 277 | + await self.log("warning", message, logger_name=logger_name, extra=extra) |
| 278 | + |
| 279 | + async def error(self, message: str, *, logger_name: str | None = None, extra: dict[str, Any] | None = None) -> None: |
| 280 | + """Send an error log message.""" |
| 281 | + await self.log("error", message, logger_name=logger_name, extra=extra) |
0 commit comments