disagreement 0.2.0rc1__py3-none-any.whl → 0.4.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. disagreement/__init__.py +2 -4
  2. disagreement/audio.py +42 -5
  3. disagreement/cache.py +43 -4
  4. disagreement/caching.py +121 -0
  5. disagreement/client.py +1682 -1535
  6. disagreement/enums.py +10 -3
  7. disagreement/error_handler.py +5 -1
  8. disagreement/errors.py +1341 -3
  9. disagreement/event_dispatcher.py +3 -5
  10. disagreement/ext/__init__.py +1 -0
  11. disagreement/ext/app_commands/__init__.py +0 -2
  12. disagreement/ext/app_commands/commands.py +0 -2
  13. disagreement/ext/app_commands/context.py +0 -2
  14. disagreement/ext/app_commands/converters.py +2 -4
  15. disagreement/ext/app_commands/decorators.py +5 -7
  16. disagreement/ext/app_commands/handler.py +1 -3
  17. disagreement/ext/app_commands/hybrid.py +0 -2
  18. disagreement/ext/commands/__init__.py +63 -61
  19. disagreement/ext/commands/cog.py +0 -2
  20. disagreement/ext/commands/converters.py +16 -5
  21. disagreement/ext/commands/core.py +728 -563
  22. disagreement/ext/commands/decorators.py +294 -219
  23. disagreement/ext/commands/errors.py +0 -2
  24. disagreement/ext/commands/help.py +0 -2
  25. disagreement/ext/commands/view.py +1 -3
  26. disagreement/gateway.py +632 -586
  27. disagreement/http.py +1362 -1041
  28. disagreement/interactions.py +0 -2
  29. disagreement/models.py +2682 -2263
  30. disagreement/shard_manager.py +0 -2
  31. disagreement/ui/view.py +167 -165
  32. disagreement/voice_client.py +263 -162
  33. {disagreement-0.2.0rc1.dist-info → disagreement-0.4.0.dist-info}/METADATA +33 -6
  34. disagreement-0.4.0.dist-info/RECORD +55 -0
  35. disagreement-0.2.0rc1.dist-info/RECORD +0 -54
  36. {disagreement-0.2.0rc1.dist-info → disagreement-0.4.0.dist-info}/WHEEL +0 -0
  37. {disagreement-0.2.0rc1.dist-info → disagreement-0.4.0.dist-info}/licenses/LICENSE +0 -0
  38. {disagreement-0.2.0rc1.dist-info → disagreement-0.4.0.dist-info}/top_level.txt +0 -0
disagreement/gateway.py CHANGED
@@ -1,586 +1,632 @@
1
- # disagreement/gateway.py
2
-
3
- """
4
- Manages the WebSocket connection to the Discord Gateway.
5
- """
6
-
7
- import asyncio
8
- import logging
9
- import traceback
10
- import aiohttp
11
- import json
12
- import zlib
13
- import time
14
- import random
15
- from typing import Optional, TYPE_CHECKING, Any, Dict
16
-
17
- from .enums import GatewayOpcode, GatewayIntent
18
- from .errors import GatewayException, DisagreementException, AuthenticationError
19
- from .interactions import Interaction
20
-
21
- if TYPE_CHECKING:
22
- from .client import Client # For type hinting
23
- from .event_dispatcher import EventDispatcher
24
- from .http import HTTPClient
25
- from .interactions import Interaction # Added for INTERACTION_CREATE
26
-
27
- # ZLIB Decompression constants
28
- ZLIB_SUFFIX = b"\x00\x00\xff\xff"
29
- MAX_DECOMPRESSION_SIZE = 10 * 1024 * 1024 # 10 MiB, adjust as needed
30
-
31
-
32
- logger = logging.getLogger(__name__)
33
-
34
-
35
- class GatewayClient:
36
- """
37
- Handles the Discord Gateway WebSocket connection, heartbeating, and event dispatching.
38
- """
39
-
40
- def __init__(
41
- self,
42
- http_client: "HTTPClient",
43
- event_dispatcher: "EventDispatcher",
44
- token: str,
45
- intents: int,
46
- client_instance: "Client", # Pass the main client instance
47
- verbose: bool = False,
48
- *,
49
- shard_id: Optional[int] = None,
50
- shard_count: Optional[int] = None,
51
- max_retries: int = 5,
52
- max_backoff: float = 60.0,
53
- ):
54
- self._http: "HTTPClient" = http_client
55
- self._dispatcher: "EventDispatcher" = event_dispatcher
56
- self._token: str = token
57
- self._intents: int = intents
58
- self._client_instance: "Client" = client_instance # Store client instance
59
- self.verbose: bool = verbose
60
- self._shard_id: Optional[int] = shard_id
61
- self._shard_count: Optional[int] = shard_count
62
- self._max_retries: int = max_retries
63
- self._max_backoff: float = max_backoff
64
-
65
- self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
66
- self._loop: asyncio.AbstractEventLoop = asyncio.get_event_loop()
67
- self._heartbeat_interval: Optional[float] = None
68
- self._last_sequence: Optional[int] = None
69
- self._session_id: Optional[str] = None
70
- self._resume_gateway_url: Optional[str] = None
71
-
72
- self._keep_alive_task: Optional[asyncio.Task] = None
73
- self._receive_task: Optional[asyncio.Task] = None
74
-
75
- self._last_heartbeat_sent: Optional[float] = None
76
- self._last_heartbeat_ack: Optional[float] = None
77
-
78
- # For zlib decompression
79
- self._buffer = bytearray()
80
- self._inflator = zlib.decompressobj()
81
-
82
- async def _reconnect(self) -> None:
83
- """Attempts to reconnect using exponential backoff with jitter."""
84
- delay = 1.0
85
- for attempt in range(self._max_retries):
86
- try:
87
- await self.connect()
88
- return
89
- except Exception as e: # noqa: BLE001
90
- if attempt >= self._max_retries - 1:
91
- logger.error(
92
- "Reconnect failed after %s attempts: %s", attempt + 1, e
93
- )
94
- raise
95
- jitter = random.uniform(0, delay)
96
- wait_time = min(delay + jitter, self._max_backoff)
97
- logger.warning(
98
- "Reconnect attempt %s failed: %s. Retrying in %.2f seconds...",
99
- attempt + 1,
100
- e,
101
- wait_time,
102
- )
103
- await asyncio.sleep(wait_time)
104
- delay = min(delay * 2, self._max_backoff)
105
-
106
- async def _decompress_message(
107
- self, message_bytes: bytes
108
- ) -> Optional[Dict[str, Any]]:
109
- """Decompresses a zlib-compressed message from the Gateway."""
110
- self._buffer.extend(message_bytes)
111
-
112
- if len(message_bytes) < 4 or message_bytes[-4:] != ZLIB_SUFFIX:
113
- # Message is not complete or not zlib compressed in the expected way
114
- return None
115
- # Or handle partial messages if Discord ever sends them fragmented like this,
116
- # but typically each binary message is a complete zlib stream.
117
-
118
- try:
119
- decompressed = self._inflator.decompress(self._buffer)
120
- self._buffer.clear() # Reset buffer after successful decompression
121
- return json.loads(decompressed.decode("utf-8"))
122
- except zlib.error as e:
123
- logger.error("Zlib decompression error: %s", e)
124
- self._buffer.clear() # Clear buffer on error
125
- self._inflator = zlib.decompressobj() # Reset inflator
126
- return None
127
- except json.JSONDecodeError as e:
128
- logger.error("JSON decode error after decompression: %s", e)
129
- return None
130
-
131
- async def _send_json(self, payload: Dict[str, Any]):
132
- if self._ws and not self._ws.closed:
133
- if self.verbose:
134
- logger.debug("GATEWAY SEND: %s", payload)
135
- await self._ws.send_json(payload)
136
- else:
137
- logger.warning(
138
- "Gateway send attempted but WebSocket is closed or not available."
139
- )
140
- # raise GatewayException("WebSocket is not connected.")
141
-
142
- async def _heartbeat(self):
143
- """Sends a heartbeat to the Gateway."""
144
- self._last_heartbeat_sent = time.monotonic()
145
- payload = {"op": GatewayOpcode.HEARTBEAT, "d": self._last_sequence}
146
- await self._send_json(payload)
147
- # print("Sent heartbeat.")
148
-
149
- async def _keep_alive(self):
150
- """Manages the heartbeating loop."""
151
- if self._heartbeat_interval is None:
152
- # This should not happen if HELLO was processed correctly
153
- logger.error("Heartbeat interval not set. Cannot start keep_alive.")
154
- return
155
-
156
- try:
157
- while True:
158
- await self._heartbeat()
159
- await asyncio.sleep(
160
- self._heartbeat_interval / 1000
161
- ) # Interval is in ms
162
- except asyncio.CancelledError:
163
- logger.debug("Keep_alive task cancelled.")
164
- except Exception as e:
165
- logger.error("Error in keep_alive loop: %s", e)
166
- # Potentially trigger a reconnect here or notify client
167
- await self._client_instance.close_gateway(code=1000) # Generic close
168
-
169
- async def _identify(self):
170
- """Sends the IDENTIFY payload to the Gateway."""
171
- payload = {
172
- "op": GatewayOpcode.IDENTIFY,
173
- "d": {
174
- "token": self._token,
175
- "intents": self._intents,
176
- "properties": {
177
- "$os": "python", # Or platform.system()
178
- "$browser": "disagreement", # Library name
179
- "$device": "disagreement", # Library name
180
- },
181
- "compress": True, # Request zlib compression
182
- },
183
- }
184
- if self._shard_id is not None and self._shard_count is not None:
185
- payload["d"]["shard"] = [self._shard_id, self._shard_count]
186
- await self._send_json(payload)
187
- logger.info("Sent IDENTIFY.")
188
-
189
- async def _resume(self):
190
- """Sends the RESUME payload to the Gateway."""
191
- if not self._session_id or self._last_sequence is None:
192
- logger.warning("Cannot RESUME: session_id or last_sequence is missing.")
193
- await self._identify() # Fallback to identify
194
- return
195
-
196
- payload = {
197
- "op": GatewayOpcode.RESUME,
198
- "d": {
199
- "token": self._token,
200
- "session_id": self._session_id,
201
- "seq": self._last_sequence,
202
- },
203
- }
204
- await self._send_json(payload)
205
- logger.info(
206
- "Sent RESUME for session %s at sequence %s.",
207
- self._session_id,
208
- self._last_sequence,
209
- )
210
-
211
- async def update_presence(
212
- self,
213
- status: str,
214
- activity_name: Optional[str] = None,
215
- activity_type: int = 0,
216
- since: int = 0,
217
- afk: bool = False,
218
- ):
219
- """Sends the presence update payload to the Gateway."""
220
- payload = {
221
- "op": GatewayOpcode.PRESENCE_UPDATE,
222
- "d": {
223
- "since": since,
224
- "activities": (
225
- [
226
- {
227
- "name": activity_name,
228
- "type": activity_type,
229
- }
230
- ]
231
- if activity_name
232
- else []
233
- ),
234
- "status": status,
235
- "afk": afk,
236
- },
237
- }
238
- await self._send_json(payload)
239
-
240
- async def _handle_dispatch(self, data: Dict[str, Any]):
241
- """Handles DISPATCH events (actual Discord events)."""
242
- event_name = data.get("t")
243
- sequence_num = data.get("s")
244
- raw_event_d_payload = data.get(
245
- "d"
246
- ) # This is the 'd' field from the gateway event
247
-
248
- if sequence_num is not None:
249
- self._last_sequence = sequence_num
250
-
251
- if event_name == "READY": # Special handling for READY
252
- if not isinstance(raw_event_d_payload, dict):
253
- logger.error(
254
- "READY event 'd' payload is not a dict or is missing: %s",
255
- raw_event_d_payload,
256
- )
257
- # Consider raising an error or attempting a reconnect
258
- return
259
- self._session_id = raw_event_d_payload.get("session_id")
260
- self._resume_gateway_url = raw_event_d_payload.get("resume_gateway_url")
261
-
262
- app_id_str = "N/A"
263
- # Store application_id on the client instance
264
- if (
265
- "application" in raw_event_d_payload
266
- and isinstance(raw_event_d_payload["application"], dict)
267
- and "id" in raw_event_d_payload["application"]
268
- ):
269
- app_id_value = raw_event_d_payload["application"]["id"]
270
- self._client_instance.application_id = (
271
- app_id_value # Snowflake can be str or int
272
- )
273
- app_id_str = str(app_id_value)
274
- else:
275
- logger.warning(
276
- "Could not find application ID in READY payload. App commands may not work."
277
- )
278
-
279
- # Parse and store the bot's own user object
280
- if "user" in raw_event_d_payload and isinstance(
281
- raw_event_d_payload["user"], dict
282
- ):
283
- try:
284
- # Assuming Client has a parse_user method that takes user data dict
285
- # and returns a User object, also caching it.
286
- bot_user_obj = self._client_instance.parse_user(
287
- raw_event_d_payload["user"]
288
- )
289
- self._client_instance.user = bot_user_obj
290
- logger.info(
291
- "Gateway READY. Bot User: %s#%s. Session ID: %s. App ID: %s. Resume URL: %s",
292
- bot_user_obj.username,
293
- bot_user_obj.discriminator,
294
- self._session_id,
295
- app_id_str,
296
- self._resume_gateway_url,
297
- )
298
- except Exception as e:
299
- logger.error("Error parsing bot user from READY payload: %s", e)
300
- logger.info(
301
- "Gateway READY (user parse failed). Session ID: %s. App ID: %s. Resume URL: %s",
302
- self._session_id,
303
- app_id_str,
304
- self._resume_gateway_url,
305
- )
306
- else:
307
- logger.warning("Bot user object not found or invalid in READY payload.")
308
- logger.info(
309
- "Gateway READY (no user). Session ID: %s. App ID: %s. Resume URL: %s",
310
- self._session_id,
311
- app_id_str,
312
- self._resume_gateway_url,
313
- )
314
-
315
- await self._dispatcher.dispatch(event_name, raw_event_d_payload)
316
- elif event_name == "INTERACTION_CREATE":
317
- # print(f"GATEWAY RECV INTERACTION_CREATE: {raw_event_d_payload}")
318
- if isinstance(raw_event_d_payload, dict):
319
- interaction = Interaction(
320
- data=raw_event_d_payload, client_instance=self._client_instance
321
- )
322
- await self._dispatcher.dispatch(
323
- "INTERACTION_CREATE", raw_event_d_payload
324
- )
325
- # Dispatch to a new client method that will then call AppCommandHandler
326
- if hasattr(self._client_instance, "process_interaction"):
327
- asyncio.create_task(
328
- self._client_instance.process_interaction(interaction)
329
- ) # type: ignore
330
- else:
331
- logger.warning(
332
- "Client instance does not have process_interaction method for INTERACTION_CREATE."
333
- )
334
- else:
335
- logger.error(
336
- "INTERACTION_CREATE event 'd' payload is not a dict: %s",
337
- raw_event_d_payload,
338
- )
339
- elif event_name == "RESUMED":
340
- logger.info("Gateway RESUMED successfully.")
341
- # RESUMED 'd' payload is often an empty object or debug info.
342
- # Ensure it's a dict for the dispatcher.
343
- event_data_to_dispatch = (
344
- raw_event_d_payload if isinstance(raw_event_d_payload, dict) else {}
345
- )
346
- await self._dispatcher.dispatch(event_name, event_data_to_dispatch)
347
- await self._dispatcher.dispatch(
348
- "SHARD_RESUME", {"shard_id": self._shard_id}
349
- )
350
- elif event_name:
351
- # For other events, ensure 'd' is a dict, or pass {} if 'd' is null/missing.
352
- # Models/parsers in EventDispatcher will need to handle potentially empty dicts.
353
- event_data_to_dispatch = (
354
- raw_event_d_payload if isinstance(raw_event_d_payload, dict) else {}
355
- )
356
- # print(f"GATEWAY RECV EVENT: {event_name} | DATA: {event_data_to_dispatch}")
357
- await self._dispatcher.dispatch(event_name, event_data_to_dispatch)
358
- else:
359
- logger.warning("Received dispatch with no event name: %s", data)
360
-
361
- async def _process_message(self, msg: aiohttp.WSMessage):
362
- """Processes a single message from the WebSocket."""
363
- if msg.type == aiohttp.WSMsgType.TEXT:
364
- try:
365
- data = json.loads(msg.data)
366
- except json.JSONDecodeError:
367
- logger.error("Failed to decode JSON from Gateway: %s", msg.data[:200])
368
- return
369
- elif msg.type == aiohttp.WSMsgType.BINARY:
370
- decompressed_data = await self._decompress_message(msg.data)
371
- if decompressed_data is None:
372
- logger.error(
373
- "Failed to decompress or decode binary message from Gateway."
374
- )
375
- return
376
- data = decompressed_data
377
- elif msg.type == aiohttp.WSMsgType.ERROR:
378
- logger.error(
379
- "WebSocket error: %s",
380
- self._ws.exception() if self._ws else "Unknown WSError",
381
- )
382
- raise GatewayException(
383
- f"WebSocket error: {self._ws.exception() if self._ws else 'Unknown WSError'}"
384
- )
385
- elif msg.type == aiohttp.WSMsgType.CLOSED:
386
- close_code = (
387
- self._ws.close_code
388
- if self._ws and hasattr(self._ws, "close_code")
389
- else "N/A"
390
- )
391
- logger.warning(
392
- "WebSocket connection closed by server. Code: %s", close_code
393
- )
394
- # Raise an exception to signal the closure to the client's main run loop
395
- raise GatewayException(f"WebSocket closed by server. Code: {close_code}")
396
- else:
397
- logger.warning("Received unhandled WebSocket message type: %s", msg.type)
398
- return
399
-
400
- if self.verbose:
401
- logger.debug("GATEWAY RECV: %s", data)
402
- op = data.get("op")
403
- # 'd' payload (event_data) is handled specifically by each opcode handler below
404
-
405
- if op == GatewayOpcode.DISPATCH:
406
- await self._handle_dispatch(data) # _handle_dispatch will extract 'd'
407
- elif op == GatewayOpcode.HEARTBEAT: # Server requests a heartbeat
408
- await self._heartbeat()
409
- elif op == GatewayOpcode.RECONNECT: # Server requests a reconnect
410
- logger.info(
411
- "Gateway requested RECONNECT. Closing and will attempt to reconnect."
412
- )
413
- await self.close(code=4000, reconnect=True)
414
- elif op == GatewayOpcode.INVALID_SESSION:
415
- # The 'd' payload for INVALID_SESSION is a boolean indicating resumability
416
- can_resume = data.get("d") is True
417
- logger.warning(
418
- "Gateway indicated INVALID_SESSION. Resumable: %s", can_resume
419
- )
420
- if not can_resume:
421
- self._session_id = None # Clear session_id to force re-identify
422
- self._last_sequence = None
423
- # Close and reconnect. The connect logic will decide to resume or identify.
424
- await self.close(code=4000 if can_resume else 4009, reconnect=True)
425
- elif op == GatewayOpcode.HELLO:
426
- hello_d_payload = data.get("d")
427
- if (
428
- not isinstance(hello_d_payload, dict)
429
- or "heartbeat_interval" not in hello_d_payload
430
- ):
431
- logger.error(
432
- "HELLO event 'd' payload is invalid or missing heartbeat_interval: %s",
433
- hello_d_payload,
434
- )
435
- await self.close(code=1011) # Internal error, malformed HELLO
436
- return
437
- self._heartbeat_interval = hello_d_payload["heartbeat_interval"]
438
- logger.info(
439
- "Gateway HELLO. Heartbeat interval: %sms.", self._heartbeat_interval
440
- )
441
- # Start heartbeating
442
- if self._keep_alive_task:
443
- self._keep_alive_task.cancel()
444
- self._keep_alive_task = self._loop.create_task(self._keep_alive())
445
-
446
- # Identify or Resume
447
- if self._session_id and self._resume_gateway_url: # Check if we can resume
448
- logger.info("Attempting to RESUME session.")
449
- await self._resume()
450
- else:
451
- logger.info("Performing initial IDENTIFY.")
452
- await self._identify()
453
- elif op == GatewayOpcode.HEARTBEAT_ACK:
454
- self._last_heartbeat_ack = time.monotonic()
455
- # print("Received heartbeat ACK.")
456
- pass # Good, connection is alive
457
- else:
458
- logger.warning(
459
- "Received unhandled Gateway Opcode: %s with data: %s", op, data
460
- )
461
-
462
- async def _receive_loop(self):
463
- """Continuously receives and processes messages from the WebSocket."""
464
- if not self._ws or self._ws.closed:
465
- logger.warning(
466
- "Receive loop cannot start: WebSocket is not connected or closed."
467
- )
468
- return
469
-
470
- try:
471
- async for msg in self._ws:
472
- await self._process_message(msg)
473
- except asyncio.CancelledError:
474
- logger.debug("Receive_loop task cancelled.")
475
- except aiohttp.ClientConnectionError as e:
476
- logger.warning(
477
- "ClientConnectionError in receive_loop: %s. Attempting reconnect.", e
478
- )
479
- await self.close(code=1006, reconnect=True) # Abnormal closure
480
- except Exception as e:
481
- logger.error("Unexpected error in receive_loop: %s", e)
482
- traceback.print_exc()
483
- await self.close(code=1011, reconnect=True)
484
- finally:
485
- logger.info("Receive_loop ended.")
486
- # If the loop ends unexpectedly (not due to explicit close),
487
- # the main client might want to try reconnecting.
488
-
489
- async def connect(self):
490
- """Connects to the Discord Gateway."""
491
- if self._ws and not self._ws.closed:
492
- logger.warning("Gateway already connected or connecting.")
493
- return
494
-
495
- gateway_url = (
496
- self._resume_gateway_url or (await self._http.get_gateway_bot())["url"]
497
- )
498
- if not gateway_url.endswith("?v=10&encoding=json&compress=zlib-stream"):
499
- gateway_url += "?v=10&encoding=json&compress=zlib-stream"
500
-
501
- logger.info("Connecting to Gateway: %s", gateway_url)
502
- try:
503
- await self._http._ensure_session() # Ensure the HTTP client's session is active
504
- assert (
505
- self._http._session is not None
506
- ), "HTTPClient session not initialized after ensure_session"
507
- self._ws = await self._http._session.ws_connect(gateway_url, max_msg_size=0)
508
- logger.info("Gateway WebSocket connection established.")
509
-
510
- if self._receive_task:
511
- self._receive_task.cancel()
512
- self._receive_task = self._loop.create_task(self._receive_loop())
513
-
514
- await self._dispatcher.dispatch(
515
- "SHARD_CONNECT", {"shard_id": self._shard_id}
516
- )
517
-
518
- except aiohttp.ClientConnectorError as e:
519
- raise GatewayException(
520
- f"Failed to connect to Gateway (Connector Error): {e}"
521
- ) from e
522
- except aiohttp.WSServerHandshakeError as e:
523
- if e.status == 401: # Unauthorized during handshake
524
- raise AuthenticationError(
525
- f"Gateway handshake failed (401 Unauthorized): {e.message}. Check your bot token."
526
- ) from e
527
- raise GatewayException(
528
- f"Gateway handshake failed (Status: {e.status}): {e.message}"
529
- ) from e
530
- except Exception as e: # Catch other potential errors during connection
531
- raise GatewayException(
532
- f"An unexpected error occurred during Gateway connection: {e}"
533
- ) from e
534
-
535
- async def close(self, code: int = 1000, *, reconnect: bool = False):
536
- """Closes the Gateway connection."""
537
- logger.info("Closing Gateway connection with code %s...", code)
538
- if self._keep_alive_task and not self._keep_alive_task.done():
539
- self._keep_alive_task.cancel()
540
- try:
541
- await self._keep_alive_task
542
- except asyncio.CancelledError:
543
- pass # Expected
544
-
545
- if self._receive_task and not self._receive_task.done():
546
- current = asyncio.current_task(loop=self._loop)
547
- self._receive_task.cancel()
548
- if self._receive_task is not current:
549
- try:
550
- await self._receive_task
551
- except asyncio.CancelledError:
552
- pass # Expected
553
-
554
- if self._ws and not self._ws.closed:
555
- await self._ws.close(code=code)
556
- logger.info("Gateway WebSocket closed.")
557
-
558
- self._ws = None
559
- # Do not reset session_id, last_sequence, or resume_gateway_url here
560
- # if the close code indicates a resumable disconnect (e.g. 4000-4009, or server-initiated RECONNECT)
561
- # The connect logic will decide whether to resume or re-identify.
562
- # However, if it's a non-resumable close (e.g. Invalid Session non-resumable), clear them.
563
- if code == 4009: # Invalid session, not resumable
564
- logger.info("Clearing session state due to non-resumable invalid session.")
565
- self._session_id = None
566
- self._last_sequence = None
567
- self._resume_gateway_url = None # This might be re-fetched anyway
568
-
569
- await self._dispatcher.dispatch(
570
- "SHARD_DISCONNECT", {"shard_id": self._shard_id}
571
- )
572
-
573
- @property
574
- def latency(self) -> Optional[float]:
575
- """Returns the latency between heartbeat and ACK in seconds."""
576
- if self._last_heartbeat_sent is None or self._last_heartbeat_ack is None:
577
- return None
578
- return self._last_heartbeat_ack - self._last_heartbeat_sent
579
-
580
- @property
581
- def last_heartbeat_sent(self) -> Optional[float]:
582
- return self._last_heartbeat_sent
583
-
584
- @property
585
- def last_heartbeat_ack(self) -> Optional[float]:
586
- return self._last_heartbeat_ack
1
+ """
2
+ Manages the WebSocket connection to the Discord Gateway.
3
+ """
4
+
5
+ import asyncio
6
+ import logging
7
+ import traceback
8
+ import aiohttp
9
+ import json
10
+ import zlib
11
+ import time
12
+ import random
13
+ from typing import Optional, TYPE_CHECKING, Any, Dict
14
+
15
+ from .models import Activity
16
+
17
+ from .enums import GatewayOpcode, GatewayIntent
18
+ from .errors import GatewayException, DisagreementException, AuthenticationError
19
+ from .interactions import Interaction
20
+
21
+ if TYPE_CHECKING:
22
+ from .client import Client # For type hinting
23
+ from .event_dispatcher import EventDispatcher
24
+ from .http import HTTPClient
25
+ from .interactions import Interaction # Added for INTERACTION_CREATE
26
+
27
+ # ZLIB Decompression constants
28
+ ZLIB_SUFFIX = b"\x00\x00\xff\xff"
29
+ MAX_DECOMPRESSION_SIZE = 10 * 1024 * 1024 # 10 MiB, adjust as needed
30
+
31
+
32
+ logger = logging.getLogger(__name__)
33
+
34
+
35
+ class GatewayClient:
36
+ """
37
+ Handles the Discord Gateway WebSocket connection, heartbeating, and event dispatching.
38
+ """
39
+
40
+ def __init__(
41
+ self,
42
+ http_client: "HTTPClient",
43
+ event_dispatcher: "EventDispatcher",
44
+ token: str,
45
+ intents: int,
46
+ client_instance: "Client", # Pass the main client instance
47
+ verbose: bool = False,
48
+ *,
49
+ shard_id: Optional[int] = None,
50
+ shard_count: Optional[int] = None,
51
+ max_retries: int = 5,
52
+ max_backoff: float = 60.0,
53
+ ):
54
+ self._http: "HTTPClient" = http_client
55
+ self._dispatcher: "EventDispatcher" = event_dispatcher
56
+ self._token: str = token
57
+ self._intents: int = intents
58
+ self._client_instance: "Client" = client_instance # Store client instance
59
+ self.verbose: bool = verbose
60
+ self._shard_id: Optional[int] = shard_id
61
+ self._shard_count: Optional[int] = shard_count
62
+ self._max_retries: int = max_retries
63
+ self._max_backoff: float = max_backoff
64
+
65
+ self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
66
+ try:
67
+ self._loop: asyncio.AbstractEventLoop = asyncio.get_running_loop()
68
+ except RuntimeError:
69
+ self._loop = asyncio.new_event_loop()
70
+ asyncio.set_event_loop(self._loop)
71
+ self._heartbeat_interval: Optional[float] = None
72
+ self._last_sequence: Optional[int] = None
73
+ self._session_id: Optional[str] = None
74
+ self._resume_gateway_url: Optional[str] = None
75
+
76
+ self._keep_alive_task: Optional[asyncio.Task] = None
77
+ self._receive_task: Optional[asyncio.Task] = None
78
+
79
+ self._last_heartbeat_sent: Optional[float] = None
80
+ self._last_heartbeat_ack: Optional[float] = None
81
+
82
+ # For zlib decompression
83
+ self._buffer = bytearray()
84
+ self._inflator = zlib.decompressobj()
85
+
86
+ self._member_chunk_requests: Dict[str, asyncio.Future] = {}
87
+
88
+ async def _reconnect(self) -> None:
89
+ """Attempts to reconnect using exponential backoff with jitter."""
90
+ delay = 1.0
91
+ for attempt in range(self._max_retries):
92
+ try:
93
+ await self.connect()
94
+ return
95
+ except Exception as e: # noqa: BLE001
96
+ if attempt >= self._max_retries - 1:
97
+ logger.error(
98
+ "Reconnect failed after %s attempts: %s", attempt + 1, e
99
+ )
100
+ raise
101
+ jitter = random.uniform(0, delay)
102
+ wait_time = min(delay + jitter, self._max_backoff)
103
+ logger.warning(
104
+ "Reconnect attempt %s failed: %s. Retrying in %.2f seconds...",
105
+ attempt + 1,
106
+ e,
107
+ wait_time,
108
+ )
109
+ await asyncio.sleep(wait_time)
110
+ delay = min(delay * 2, self._max_backoff)
111
+
112
+ async def _decompress_message(
113
+ self, message_bytes: bytes
114
+ ) -> Optional[Dict[str, Any]]:
115
+ """Decompresses a zlib-compressed message from the Gateway."""
116
+ self._buffer.extend(message_bytes)
117
+
118
+ if len(message_bytes) < 4 or message_bytes[-4:] != ZLIB_SUFFIX:
119
+ # Message is not complete or not zlib compressed in the expected way
120
+ return None
121
+ # Or handle partial messages if Discord ever sends them fragmented like this,
122
+ # but typically each binary message is a complete zlib stream.
123
+
124
+ try:
125
+ decompressed = self._inflator.decompress(self._buffer)
126
+ self._buffer.clear() # Reset buffer after successful decompression
127
+ return json.loads(decompressed.decode("utf-8"))
128
+ except zlib.error as e:
129
+ logger.error("Zlib decompression error: %s", e)
130
+ self._buffer.clear() # Clear buffer on error
131
+ self._inflator = zlib.decompressobj() # Reset inflator
132
+ return None
133
+ except json.JSONDecodeError as e:
134
+ logger.error("JSON decode error after decompression: %s", e)
135
+ return None
136
+
137
+ async def _send_json(self, payload: Dict[str, Any]):
138
+ if self._ws and not self._ws.closed:
139
+ if self.verbose:
140
+ logger.debug("GATEWAY SEND: %s", payload)
141
+ await self._ws.send_json(payload)
142
+ else:
143
+ logger.warning(
144
+ "Gateway send attempted but WebSocket is closed or not available."
145
+ )
146
+ # raise GatewayException("WebSocket is not connected.")
147
+
148
+ async def _heartbeat(self):
149
+ """Sends a heartbeat to the Gateway."""
150
+ self._last_heartbeat_sent = time.monotonic()
151
+ payload = {"op": GatewayOpcode.HEARTBEAT, "d": self._last_sequence}
152
+ await self._send_json(payload)
153
+
154
+ async def _keep_alive(self):
155
+ """Manages the heartbeating loop."""
156
+ if self._heartbeat_interval is None:
157
+
158
+ logger.error("Heartbeat interval not set. Cannot start keep_alive.")
159
+ return
160
+
161
+ try:
162
+ while True:
163
+ await self._heartbeat()
164
+ await asyncio.sleep(
165
+ self._heartbeat_interval / 1000
166
+ ) # Interval is in ms
167
+ except asyncio.CancelledError:
168
+ logger.debug("Keep_alive task cancelled.")
169
+ except Exception as e:
170
+ logger.error("Error in keep_alive loop: %s", e)
171
+ # Potentially trigger a reconnect here or notify client
172
+ await self._client_instance.close_gateway(code=1000) # Generic close
173
+
174
+ async def _identify(self):
175
+ """Sends the IDENTIFY payload to the Gateway."""
176
+ payload = {
177
+ "op": GatewayOpcode.IDENTIFY,
178
+ "d": {
179
+ "token": self._token,
180
+ "intents": self._intents,
181
+ "properties": {
182
+ "$os": "python", # Or platform.system()
183
+ "$browser": "disagreement", # Library name
184
+ "$device": "disagreement", # Library name
185
+ },
186
+ "compress": True, # Request zlib compression
187
+ },
188
+ }
189
+ if self._shard_id is not None and self._shard_count is not None:
190
+ payload["d"]["shard"] = [self._shard_id, self._shard_count]
191
+ await self._send_json(payload)
192
+ logger.info("Sent IDENTIFY.")
193
+
194
+ async def _resume(self):
195
+ """Sends the RESUME payload to the Gateway."""
196
+ if not self._session_id or self._last_sequence is None:
197
+ logger.warning("Cannot RESUME: session_id or last_sequence is missing.")
198
+ await self._identify() # Fallback to identify
199
+ return
200
+
201
+ payload = {
202
+ "op": GatewayOpcode.RESUME,
203
+ "d": {
204
+ "token": self._token,
205
+ "session_id": self._session_id,
206
+ "seq": self._last_sequence,
207
+ },
208
+ }
209
+ await self._send_json(payload)
210
+ logger.info(
211
+ "Sent RESUME for session %s at sequence %s.",
212
+ self._session_id,
213
+ self._last_sequence,
214
+ )
215
+
216
+ async def update_presence(
217
+ self,
218
+ status: str,
219
+ activity: Optional[Activity] = None,
220
+ *,
221
+ since: int = 0,
222
+ afk: bool = False,
223
+ ) -> None:
224
+ """Sends the presence update payload to the Gateway."""
225
+ payload = {
226
+ "op": GatewayOpcode.PRESENCE_UPDATE,
227
+ "d": {
228
+ "since": since,
229
+ "activities": [activity.to_dict()] if activity else [],
230
+ "status": status,
231
+ "afk": afk,
232
+ },
233
+ }
234
+ await self._send_json(payload)
235
+
236
+ async def request_guild_members(
237
+ self,
238
+ guild_id: str,
239
+ query: str = "",
240
+ limit: int = 0,
241
+ presences: bool = False,
242
+ user_ids: Optional[list[str]] = None,
243
+ nonce: Optional[str] = None,
244
+ ):
245
+ """Sends the request guild members payload to the Gateway."""
246
+ payload = {
247
+ "op": GatewayOpcode.REQUEST_GUILD_MEMBERS,
248
+ "d": {
249
+ "guild_id": guild_id,
250
+ "query": query,
251
+ "limit": limit,
252
+ "presences": presences,
253
+ },
254
+ }
255
+ if user_ids:
256
+ payload["d"]["user_ids"] = user_ids
257
+ if nonce:
258
+ payload["d"]["nonce"] = nonce
259
+
260
+ await self._send_json(payload)
261
+
262
+ async def _handle_dispatch(self, data: Dict[str, Any]):
263
+ """Handles DISPATCH events (actual Discord events)."""
264
+ event_name = data.get("t")
265
+ sequence_num = data.get("s")
266
+ raw_event_d_payload = data.get(
267
+ "d"
268
+ ) # This is the 'd' field from the gateway event
269
+
270
+ if sequence_num is not None:
271
+ self._last_sequence = sequence_num
272
+
273
+ if event_name == "READY": # Special handling for READY
274
+ if not isinstance(raw_event_d_payload, dict):
275
+ logger.error(
276
+ "READY event 'd' payload is not a dict or is missing: %s",
277
+ raw_event_d_payload,
278
+ )
279
+ # Consider raising an error or attempting a reconnect
280
+ return
281
+ self._session_id = raw_event_d_payload.get("session_id")
282
+ self._resume_gateway_url = raw_event_d_payload.get("resume_gateway_url")
283
+
284
+ app_id_str = "N/A"
285
+ # Store application_id on the client instance
286
+ if (
287
+ "application" in raw_event_d_payload
288
+ and isinstance(raw_event_d_payload["application"], dict)
289
+ and "id" in raw_event_d_payload["application"]
290
+ ):
291
+ app_id_value = raw_event_d_payload["application"]["id"]
292
+ self._client_instance.application_id = (
293
+ app_id_value # Snowflake can be str or int
294
+ )
295
+ app_id_str = str(app_id_value)
296
+ else:
297
+ logger.warning(
298
+ "Could not find application ID in READY payload. App commands may not work."
299
+ )
300
+
301
+ # Parse and store the bot's own user object
302
+ if "user" in raw_event_d_payload and isinstance(
303
+ raw_event_d_payload["user"], dict
304
+ ):
305
+ try:
306
+ # Assuming Client has a parse_user method that takes user data dict
307
+ # and returns a User object, also caching it.
308
+ bot_user_obj = self._client_instance.parse_user(
309
+ raw_event_d_payload["user"]
310
+ )
311
+ self._client_instance.user = bot_user_obj
312
+ logger.info(
313
+ "Gateway READY. Bot User: %s#%s. Session ID: %s. App ID: %s. Resume URL: %s",
314
+ bot_user_obj.username,
315
+ bot_user_obj.discriminator,
316
+ self._session_id,
317
+ app_id_str,
318
+ self._resume_gateway_url,
319
+ )
320
+ except Exception as e:
321
+ logger.error("Error parsing bot user from READY payload: %s", e)
322
+ logger.info(
323
+ "Gateway READY (user parse failed). Session ID: %s. App ID: %s. Resume URL: %s",
324
+ self._session_id,
325
+ app_id_str,
326
+ self._resume_gateway_url,
327
+ )
328
+ else:
329
+ logger.warning("Bot user object not found or invalid in READY payload.")
330
+ logger.info(
331
+ "Gateway READY (no user). Session ID: %s. App ID: %s. Resume URL: %s",
332
+ self._session_id,
333
+ app_id_str,
334
+ self._resume_gateway_url,
335
+ )
336
+
337
+ await self._dispatcher.dispatch(event_name, raw_event_d_payload)
338
+ elif event_name == "GUILD_MEMBERS_CHUNK":
339
+ if isinstance(raw_event_d_payload, dict):
340
+ nonce = raw_event_d_payload.get("nonce")
341
+ if nonce and nonce in self._member_chunk_requests:
342
+ future = self._member_chunk_requests[nonce]
343
+ if not future.done():
344
+ # Append members to a temporary list stored on the future object
345
+ if not hasattr(future, "_members"):
346
+ future._members = [] # type: ignore
347
+ future._members.extend(raw_event_d_payload.get("members", [])) # type: ignore
348
+
349
+ # If this is the last chunk, resolve the future
350
+ if (
351
+ raw_event_d_payload.get("chunk_index")
352
+ == raw_event_d_payload.get("chunk_count", 1) - 1
353
+ ):
354
+ future.set_result(future._members) # type: ignore
355
+ del self._member_chunk_requests[nonce]
356
+
357
+ elif event_name == "INTERACTION_CREATE":
358
+
359
+ if isinstance(raw_event_d_payload, dict):
360
+ interaction = Interaction(
361
+ data=raw_event_d_payload, client_instance=self._client_instance
362
+ )
363
+ await self._dispatcher.dispatch(
364
+ "INTERACTION_CREATE", raw_event_d_payload
365
+ )
366
+ # Dispatch to a new client method that will then call AppCommandHandler
367
+ if hasattr(self._client_instance, "process_interaction"):
368
+ asyncio.create_task(
369
+ self._client_instance.process_interaction(interaction)
370
+ ) # type: ignore
371
+ else:
372
+ logger.warning(
373
+ "Client instance does not have process_interaction method for INTERACTION_CREATE."
374
+ )
375
+ else:
376
+ logger.error(
377
+ "INTERACTION_CREATE event 'd' payload is not a dict: %s",
378
+ raw_event_d_payload,
379
+ )
380
+ elif event_name == "RESUMED":
381
+ logger.info("Gateway RESUMED successfully.")
382
+ # RESUMED 'd' payload is often an empty object or debug info.
383
+ # Ensure it's a dict for the dispatcher.
384
+ event_data_to_dispatch = (
385
+ raw_event_d_payload if isinstance(raw_event_d_payload, dict) else {}
386
+ )
387
+ await self._dispatcher.dispatch(event_name, event_data_to_dispatch)
388
+ await self._dispatcher.dispatch(
389
+ "SHARD_RESUME", {"shard_id": self._shard_id}
390
+ )
391
+ elif event_name:
392
+ # For other events, ensure 'd' is a dict, or pass {} if 'd' is null/missing.
393
+ # Models/parsers in EventDispatcher will need to handle potentially empty dicts.
394
+ event_data_to_dispatch = (
395
+ raw_event_d_payload if isinstance(raw_event_d_payload, dict) else {}
396
+ )
397
+
398
+ await self._dispatcher.dispatch(event_name, event_data_to_dispatch)
399
+ else:
400
+ logger.warning("Received dispatch with no event name: %s", data)
401
+
402
+ async def _process_message(self, msg: aiohttp.WSMessage):
403
+ """Processes a single message from the WebSocket."""
404
+ if msg.type == aiohttp.WSMsgType.TEXT:
405
+ try:
406
+ data = json.loads(msg.data)
407
+ except json.JSONDecodeError:
408
+ logger.error("Failed to decode JSON from Gateway: %s", msg.data[:200])
409
+ return
410
+ elif msg.type == aiohttp.WSMsgType.BINARY:
411
+ decompressed_data = await self._decompress_message(msg.data)
412
+ if decompressed_data is None:
413
+ logger.error(
414
+ "Failed to decompress or decode binary message from Gateway."
415
+ )
416
+ return
417
+ data = decompressed_data
418
+ elif msg.type == aiohttp.WSMsgType.ERROR:
419
+ logger.error(
420
+ "WebSocket error: %s",
421
+ self._ws.exception() if self._ws else "Unknown WSError",
422
+ )
423
+ raise GatewayException(
424
+ f"WebSocket error: {self._ws.exception() if self._ws else 'Unknown WSError'}"
425
+ )
426
+ elif msg.type == aiohttp.WSMsgType.CLOSED:
427
+ close_code = (
428
+ self._ws.close_code
429
+ if self._ws and hasattr(self._ws, "close_code")
430
+ else "N/A"
431
+ )
432
+ logger.warning(
433
+ "WebSocket connection closed by server. Code: %s", close_code
434
+ )
435
+ # Raise an exception to signal the closure to the client's main run loop
436
+ raise GatewayException(f"WebSocket closed by server. Code: {close_code}")
437
+ else:
438
+ logger.warning("Received unhandled WebSocket message type: %s", msg.type)
439
+ return
440
+
441
+ if self.verbose:
442
+ logger.debug("GATEWAY RECV: %s", data)
443
+ op = data.get("op")
444
+ # 'd' payload (event_data) is handled specifically by each opcode handler below
445
+
446
+ if op == GatewayOpcode.DISPATCH:
447
+ await self._handle_dispatch(data) # _handle_dispatch will extract 'd'
448
+ elif op == GatewayOpcode.HEARTBEAT: # Server requests a heartbeat
449
+ await self._heartbeat()
450
+ elif op == GatewayOpcode.RECONNECT: # Server requests a reconnect
451
+ logger.info(
452
+ "Gateway requested RECONNECT. Closing and will attempt to reconnect."
453
+ )
454
+ await self.close(code=4000, reconnect=True)
455
+ elif op == GatewayOpcode.INVALID_SESSION:
456
+ # The 'd' payload for INVALID_SESSION is a boolean indicating resumability
457
+ can_resume = data.get("d") is True
458
+ logger.warning(
459
+ "Gateway indicated INVALID_SESSION. Resumable: %s", can_resume
460
+ )
461
+ if not can_resume:
462
+ self._session_id = None # Clear session_id to force re-identify
463
+ self._last_sequence = None
464
+ # Close and reconnect. The connect logic will decide to resume or identify.
465
+ await self.close(code=4000 if can_resume else 4009, reconnect=True)
466
+ elif op == GatewayOpcode.HELLO:
467
+ hello_d_payload = data.get("d")
468
+ if (
469
+ not isinstance(hello_d_payload, dict)
470
+ or "heartbeat_interval" not in hello_d_payload
471
+ ):
472
+ logger.error(
473
+ "HELLO event 'd' payload is invalid or missing heartbeat_interval: %s",
474
+ hello_d_payload,
475
+ )
476
+ await self.close(code=1011) # Internal error, malformed HELLO
477
+ return
478
+ self._heartbeat_interval = hello_d_payload["heartbeat_interval"]
479
+ logger.info(
480
+ "Gateway HELLO. Heartbeat interval: %sms.", self._heartbeat_interval
481
+ )
482
+ # Start heartbeating
483
+ if self._keep_alive_task:
484
+ self._keep_alive_task.cancel()
485
+ self._keep_alive_task = self._loop.create_task(self._keep_alive())
486
+
487
+ # Identify or Resume
488
+ if self._session_id and self._resume_gateway_url: # Check if we can resume
489
+ logger.info("Attempting to RESUME session.")
490
+ await self._resume()
491
+ else:
492
+ logger.info("Performing initial IDENTIFY.")
493
+ await self._identify()
494
+ elif op == GatewayOpcode.HEARTBEAT_ACK:
495
+ self._last_heartbeat_ack = time.monotonic()
496
+ else:
497
+ logger.warning(
498
+ "Received unhandled Gateway Opcode: %s with data: %s", op, data
499
+ )
500
+
501
+ async def _receive_loop(self):
502
+ """Continuously receives and processes messages from the WebSocket."""
503
+ if not self._ws or self._ws.closed:
504
+ logger.warning(
505
+ "Receive loop cannot start: WebSocket is not connected or closed."
506
+ )
507
+ return
508
+
509
+ try:
510
+ async for msg in self._ws:
511
+ await self._process_message(msg)
512
+ except asyncio.CancelledError:
513
+ logger.debug("Receive_loop task cancelled.")
514
+ except aiohttp.ClientConnectionError as e:
515
+ logger.warning(
516
+ "ClientConnectionError in receive_loop: %s. Attempting reconnect.", e
517
+ )
518
+ await self.close(code=1006, reconnect=True) # Abnormal closure
519
+ except Exception as e:
520
+ logger.error("Unexpected error in receive_loop: %s", e)
521
+ traceback.print_exc()
522
+ await self.close(code=1011, reconnect=True)
523
+ finally:
524
+ logger.info("Receive_loop ended.")
525
+ # If the loop ends unexpectedly (not due to explicit close),
526
+ # the main client might want to try reconnecting.
527
+
528
+ async def connect(self):
529
+ """Connects to the Discord Gateway."""
530
+ if self._ws and not self._ws.closed:
531
+ logger.warning("Gateway already connected or connecting.")
532
+ return
533
+
534
+ gateway_url = (
535
+ self._resume_gateway_url or (await self._http.get_gateway_bot())["url"]
536
+ )
537
+ if not gateway_url.endswith("?v=10&encoding=json&compress=zlib-stream"):
538
+ gateway_url += "?v=10&encoding=json&compress=zlib-stream"
539
+
540
+ logger.info("Connecting to Gateway: %s", gateway_url)
541
+ try:
542
+ await self._http._ensure_session() # Ensure the HTTP client's session is active
543
+ assert (
544
+ self._http._session is not None
545
+ ), "HTTPClient session not initialized after ensure_session"
546
+ self._ws = await self._http._session.ws_connect(gateway_url, max_msg_size=0)
547
+ logger.info("Gateway WebSocket connection established.")
548
+
549
+ if self._receive_task:
550
+ self._receive_task.cancel()
551
+ self._receive_task = self._loop.create_task(self._receive_loop())
552
+
553
+ await self._dispatcher.dispatch(
554
+ "SHARD_CONNECT", {"shard_id": self._shard_id}
555
+ )
556
+
557
+ except aiohttp.ClientConnectorError as e:
558
+ raise GatewayException(
559
+ f"Failed to connect to Gateway (Connector Error): {e}"
560
+ ) from e
561
+ except aiohttp.WSServerHandshakeError as e:
562
+ if e.status == 401: # Unauthorized during handshake
563
+ raise AuthenticationError(
564
+ f"Gateway handshake failed (401 Unauthorized): {e.message}. Check your bot token."
565
+ ) from e
566
+ raise GatewayException(
567
+ f"Gateway handshake failed (Status: {e.status}): {e.message}"
568
+ ) from e
569
+ except Exception as e: # Catch other potential errors during connection
570
+ raise GatewayException(
571
+ f"An unexpected error occurred during Gateway connection: {e}"
572
+ ) from e
573
+
574
+ async def close(self, code: int = 1000, *, reconnect: bool = False):
575
+ """Closes the Gateway connection."""
576
+ logger.info("Closing Gateway connection with code %s...", code)
577
+ if self._keep_alive_task and not self._keep_alive_task.done():
578
+ self._keep_alive_task.cancel()
579
+ try:
580
+ await self._keep_alive_task
581
+ except asyncio.CancelledError:
582
+ pass
583
+
584
+ if self._receive_task and not self._receive_task.done():
585
+ current = asyncio.current_task(loop=self._loop)
586
+ self._receive_task.cancel()
587
+ if self._receive_task is not current:
588
+ try:
589
+ await self._receive_task
590
+ except asyncio.CancelledError:
591
+ pass
592
+
593
+ if self._ws and not self._ws.closed:
594
+ await self._ws.close(code=code)
595
+ logger.info("Gateway WebSocket closed.")
596
+
597
+ self._ws = None
598
+ # Do not reset session_id, last_sequence, or resume_gateway_url here
599
+ # if the close code indicates a resumable disconnect (e.g. 4000-4009, or server-initiated RECONNECT)
600
+ # The connect logic will decide whether to resume or re-identify.
601
+ # However, if it's a non-resumable close (e.g. Invalid Session non-resumable), clear them.
602
+ if code == 4009: # Invalid session, not resumable
603
+ logger.info("Clearing session state due to non-resumable invalid session.")
604
+ self._session_id = None
605
+ self._last_sequence = None
606
+ self._resume_gateway_url = None # This might be re-fetched anyway
607
+
608
+ await self._dispatcher.dispatch(
609
+ "SHARD_DISCONNECT", {"shard_id": self._shard_id}
610
+ )
611
+
612
+ @property
613
+ def latency(self) -> Optional[float]:
614
+ """Returns the latency between heartbeat and ACK in seconds."""
615
+ if self._last_heartbeat_sent is None or self._last_heartbeat_ack is None:
616
+ return None
617
+ return self._last_heartbeat_ack - self._last_heartbeat_sent
618
+
619
+ @property
620
+ def latency_ms(self) -> Optional[float]:
621
+ """Returns the latency between heartbeat and ACK in milliseconds."""
622
+ if self._last_heartbeat_sent is None or self._last_heartbeat_ack is None:
623
+ return None
624
+ return (self._last_heartbeat_ack - self._last_heartbeat_sent) * 1000
625
+
626
+ @property
627
+ def last_heartbeat_sent(self) -> Optional[float]:
628
+ return self._last_heartbeat_sent
629
+
630
+ @property
631
+ def last_heartbeat_ack(self) -> Optional[float]:
632
+ return self._last_heartbeat_ack