Skip to content

websocket: Make unsubscribe helpers accept the request ID returned by subscribe #701

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.

Sign up for GitHub

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

Merged
michaelhly merged 4 commits into from
Aug 27, 2026
Merged

websocket: Make unsubscribe helpers accept the request ID returned by subscribe #701

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9033bde
websocket: Make unsubscribe helpers accept the request ID returned by…
kartsan03 Aug 26, 2026
b8a3e79
websocket: Prefer server subscription IDs when unsubscribing
kartsan03 Aug 26, 2026
a3dabda
Update src/solana/rpc/websocket_api.py
michaelhly Aug 27, 2026
5f8c1fb
Update README.md
michaelhly Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,9 @@ async def main():
asyncio.run(main())
```

`*_unsubscribe()` accepts either the server-assigned ID (`first_resp[0].result`, as above)
or the request ID returned by the matching `*_subscribe()` helper.
Comment thread
michaelhly marked this conversation as resolved.
Outdated

## 🔨 Development

### Setup
Expand Down
133 changes: 79 additions & 54 deletions src/solana/rpc/websocket_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@
from __future__ import annotations

import itertools
import logging
from collections.abc import Callable, Sequence
from contextlib import asynccontextmanager
from collections.abc import Sequence
from typing import Any, AsyncIterator, Dict, List, Optional, Union, cast

from solders.account_decoder import UiDataSliceConfig
Expand Down Expand Up @@ -43,7 +44,7 @@
VoteUnsubscribe,
batch_to_json,
)
from solders.rpc.responses import Notification
from solders.rpc.responses import Notification, SignatureNotification
from solders.rpc.responses import SubscriptionError as SoldersSubscriptionError
from solders.rpc.responses import SubscriptionResult, parse_websocket_message
from solders.signature import Signature
Expand All @@ -62,6 +63,8 @@
_TX_ENCODING_TO_SOLDERS,
)

logger = logging.getLogger(__name__)


class SubscriptionError(Exception):
"""Raise when subscribing to an RPC feed fails."""
Expand Down Expand Up @@ -89,6 +92,7 @@ def __init__(self, *args, **kwargs):
self.subscriptions: Dict[int, Body] = {}
self.sent_subscriptions: Dict[int, Body] = {}
self.failed_subscriptions = {}
self.request_ids_to_subscriptions: Dict[int, int] = {}
self.request_counter = itertools.count()

def increment_counter_and_get_id(self) -> int:
Expand All @@ -103,14 +107,14 @@ async def send_request(self, message: Union[Body, List[Body]]) -> None:
Args:
message: The request(s) to send.
"""
if isinstance(message, list):
to_send = batch_to_json(message)
for req in message:
self.sent_subscriptions[req.id] = req
else:
to_send = message.to_json()
self.sent_subscriptions[message.id] = message
reqs = message if isinstance(message, list) else [message]
to_send = batch_to_json(reqs) if isinstance(message, list) else message.to_json()
Comment thread
michaelhly marked this conversation as resolved.
Outdated
for req in reqs:
self.sent_subscriptions[req.id] = req
await self.send(to_send)
for req in reqs:
if hasattr(req, "subscription_id"):
self._forget_subscription(req.subscription_id)

async def recv( # type: ignore
self,
Expand Down Expand Up @@ -154,12 +158,10 @@ async def account_unsubscribe(
"""Unsubscribe from account notifications.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = AccountUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(AccountUnsubscribe, subscription)

async def logs_subscribe(
self,
Expand All @@ -186,12 +188,10 @@ async def logs_unsubscribe(
"""Unsubscribe from transaction logging.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = LogsUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(LogsUnsubscribe, subscription)

async def block_subscribe(
self,
Expand Down Expand Up @@ -233,12 +233,10 @@ async def block_unsubscribe(
"""Unsubscribe from blocks.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = BlockUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(BlockUnsubscribe, subscription)

async def program_subscribe( # pylint: disable=too-many-arguments
self,
Expand Down Expand Up @@ -290,12 +288,10 @@ async def program_unsubscribe(
"""Unsubscribe from program account notifications.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = ProgramUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(ProgramUnsubscribe, subscription)

async def signature_subscribe(
self,
Expand All @@ -322,12 +318,10 @@ async def signature_unsubscribe(
"""Unsubscribe from signature notifications.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = SignatureUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(SignatureUnsubscribe, subscription)

async def slot_subscribe(self) -> int:
"""Subscribe to receive notification anytime a slot is processed by the validator."""
Expand All @@ -343,12 +337,10 @@ async def slot_unsubscribe(
"""Unsubscribe from slot notifications.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = SlotUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(SlotUnsubscribe, subscription)

async def slots_updates_subscribe(self) -> int:
"""Subscribe to receive a notification from the validator on a variety of updates on every slot."""
Expand All @@ -364,12 +356,10 @@ async def slots_updates_unsubscribe(
"""Unsubscribe from slot update notifications.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = SlotsUpdatesUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(SlotsUpdatesUnsubscribe, subscription)

async def root_subscribe(self) -> int:
"""Subscribe to receive notification anytime a new root is set by the validator."""
Expand All @@ -385,12 +375,10 @@ async def root_unsubscribe(
"""Unsubscribe from root notifications.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = RootUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(RootUnsubscribe, subscription)

async def vote_subscribe(self) -> int:
"""Subscribe to receive notification anytime a new vote is observed in gossip."""
Expand All @@ -406,12 +394,10 @@ async def vote_unsubscribe(
"""Unsubscribe from vote notifications.

Args:
subscription: ID of subscription to cancel.
subscription: Request ID from the matching ``subscribe`` helper, or the
server-assigned ID from the confirmation (``recv()[0].result``).
"""
req_id = self.increment_counter_and_get_id()
req = VoteUnsubscribe(subscription, req_id)
await self.send_request(req)
del self.subscriptions[subscription]
await self._unsubscribe(VoteUnsubscribe, subscription)

def _process_rpc_response(self, raw: str) -> List[Union[Notification, SubscriptionResult]]:
parsed = parse_websocket_message(raw)
Expand All @@ -422,8 +408,47 @@ def _process_rpc_response(self, raw: str) -> List[Union[Notification, Subscripti
raise SubscriptionError(item, subscription)
if isinstance(item, SubscriptionResult):
self.subscriptions[item.result] = self.sent_subscriptions[item.id]
self.request_ids_to_subscriptions[item.id] = item.result
Comment thread
michaelhly marked this conversation as resolved.
elif isinstance(item, SignatureNotification):
# Signature subscriptions expire server-side after the notification.
self._forget_subscription(item.subscription)
return cast(List[Union[Notification, SubscriptionResult]], parsed)

async def _unsubscribe(self, constructor: Callable[[int, int], Body], subscription: int) -> None:
"""Send an unsubscribe request after resolving the server-assigned subscription ID."""
server_subscription = self._resolve_server_subscription(subscription)
req_id = self.increment_counter_and_get_id()
req = constructor(server_subscription, req_id)
await self.send_request(req)
self._forget_subscription(server_subscription)

def _resolve_server_subscription(self, subscription: int) -> int:
"""Translate a subscribe request ID into its server-assigned subscription ID.

Server-assigned IDs already recorded in ``self.subscriptions`` win, so the README
form (``recv()[0].result``) is never rewritten through the request-ID map. An
unconfirmed request ID is sent as-is.
"""
if subscription in self.subscriptions:
return subscription
mapped = self.request_ids_to_subscriptions.get(subscription)
if mapped is not None:
return mapped
logger.warning(
"Unsubscribe target %s is not a known server subscription ID or confirmed request ID",
subscription,
)
return subscription

def _forget_subscription(self, server_subscription: int) -> None:
"""Drop local bookkeeping for a server-assigned subscription ID."""
self.subscriptions.pop(server_subscription, None)

@michaelhly michaelhly Aug 26, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: swapping del for pop(..., None) removes the only local signal that the unsubscribe target isn't a live subscription. Unsubscribing before confirmation (or twice) now silently sends an invalid subscription ID; the server's -32602 comes back as a SoldersSubscriptionError and is raised as SubscriptionError out of an unrelated later recv(), while the real subscription keeps streaming. Dropping the KeyError is the right call, but a returned bool or a logging.warning would let the caller learn the unsubscribe was a no-op at the call site.

self.request_ids_to_subscriptions = {
req_id: sub_id
for req_id, sub_id in self.request_ids_to_subscriptions.items()
if sub_id != server_subscription
}


@asynccontextmanager
async def connect(uri: str = "ws://localhost:8900", **kwargs: Any) -> AsyncIterator[SolanaWsClientProtocol]:
Expand Down
Loading