-
Notifications
You must be signed in to change notification settings - Fork 4k
feat: add protocol version override support for client session initialization #2652
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
a6cfd86
9ca38c1
333250b
dce39ec
ce45982
bbef3cf
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -366,6 +366,13 @@ async def main(): | |
| derived).""" | ||
|
|
||
| _entered: bool = field(init=False, default=False) | ||
| protocol_version_override: str | None = None | ||
| """Pin the legacy `initialize` handshake to a specific handshake-era version. | ||
|
|
||
| Only meaningful with `mode='legacy'` or `mode='auto'` (where it skips `server/discover` | ||
| and negotiates directly); raises at construction with any other `mode`, since a version | ||
| pin already fixes the negotiated version. Must be a member of `HANDSHAKE_PROTOCOL_VERSIONS`. | ||
| `None` (the default) negotiates the latest version each `mode` would otherwise pick.""" | ||
| _session: ClientSession | None = field(init=False, default=None) | ||
| _exit_stack: AsyncExitStack | None = field(init=False, default=None) | ||
| _connect: _Connector = field(init=False, repr=False, compare=False) | ||
|
|
@@ -383,6 +390,23 @@ def __post_init__(self) -> None: | |
| f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}" | ||
| ) | ||
|
|
||
| if self.protocol_version_override is not None: | ||
| if self.protocol_version_override not in HANDSHAKE_PROTOCOL_VERSIONS: | ||
| hint = ( | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P2: The Prompt for AI agents |
||
| f" ({self.protocol_version_override!r} is a modern version; mode='auto' already negotiates it)" | ||
| if self.protocol_version_override in MODERN_PROTOCOL_VERSIONS | ||
| else "" | ||
| ) | ||
| raise ValueError( | ||
| "protocol_version_override must be one of " | ||
| f"{list(HANDSHAKE_PROTOCOL_VERSIONS)}; got {self.protocol_version_override!r}{hint}" | ||
| ) | ||
| if self.mode not in ("legacy", "auto"): | ||
| raise ValueError( | ||
| f"protocol_version_override has no effect with mode={self.mode!r} " | ||
| "(a version pin already fixes the negotiated version); use mode='legacy' or mode='auto'" | ||
| ) | ||
|
|
||
| self._folded_extensions = _fold_extensions(self.extensions) | ||
|
|
||
| srv = self.server | ||
|
|
@@ -424,7 +448,12 @@ def __post_init__(self) -> None: | |
|
|
||
| async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession: | ||
| """Enter the resolved connector and return an un-entered ClientSession.""" | ||
| dispatcher = await self._connect(exit_stack, self.mode, self.raise_exceptions) | ||
| # An override on mode='auto' skips discovery and drives `initialize()` directly | ||
| # (see `negotiate_auto`), so the in-proc connector must hand back the legacy, | ||
| # stream-backed dispatcher for this combination too, not the handshake-less | ||
| # DirectDispatcher it otherwise picks for every non-'legacy' mode. | ||
| connect_mode = "legacy" if self.mode == "auto" and self.protocol_version_override is not None else self.mode | ||
| dispatcher = await self._connect(exit_stack, connect_mode, self.raise_exceptions) | ||
| message_handler = self.message_handler | ||
| if self._response_cache is not None: | ||
| message_handler = _evicting_message_handler(self._response_cache, self.message_handler) | ||
|
|
@@ -455,9 +484,12 @@ async def __aenter__(self) -> Client: | |
| session = await exit_stack.enter_async_context(session) | ||
|
|
||
| if self.mode == "legacy": | ||
| await session.initialize() | ||
| if self.protocol_version_override is not None: | ||
| await session.initialize(protocol_version=self.protocol_version_override) | ||
| else: | ||
| await session.initialize() | ||
| elif self.mode == "auto": | ||
| await negotiate_auto(session) | ||
| await negotiate_auto(session, protocol_version=self.protocol_version_override) | ||
|
cubic-dev-ai[bot] marked this conversation as resolved.
|
||
| else: | ||
| session.adopt(self.prior_discover or _synthesize_discover(self.mode)) | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -648,15 +648,14 @@ def _build_capabilities(self, version: str) -> types.ClientCapabilities: | |
| sampling=sampling, elicitation=elicitation, experimental=None, extensions=extensions, roots=roots | ||
| ) | ||
|
|
||
| async def initialize(self) -> types.InitializeResult: | ||
| async def initialize(self, protocol_version: str = LATEST_HANDSHAKE_VERSION) -> types.InitializeResult: | ||
| if self._initialize_result is not None: | ||
| return self._initialize_result | ||
| result = await self.send_request( | ||
| types.InitializeRequest( | ||
| params=types.InitializeRequestParams( | ||
| protocol_version=LATEST_HANDSHAKE_VERSION, | ||
| # The handshake negotiates only legacy versions, where no claim is active. | ||
| capabilities=self._build_capabilities(LATEST_HANDSHAKE_VERSION), | ||
| protocol_version=protocol_version, | ||
| capabilities=self._build_capabilities(protocol_version), | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P2: Older protocol overrides still advertise form and URL elicitation capabilities when the callback is configured. Gate each capability by the requested protocol version; otherwise the server can use an advertisement for features unavailable in the selected protocol. Prompt for AI agents |
||
| client_info=self._client_info, | ||
| ), | ||
| ), | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.