chia-blockchain 2.5.0rc1__py3-none-any.whl → 2.5.1__py3-none-any.whl

chia/daemon/client.py

@@ -3,14 +3,16 @@ from __future__ import annotations
 import asyncio
 import json
 import ssl
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any, AsyncIterator, Dict, List, Optional
+from typing import Any, Optional
 
 import aiohttp
 
 from chia.util.ints import uint32
 from chia.util.json_util import dict_to_json_str
+from chia.util.task_referencer import create_referenced_task
 from chia.util.ws_message import WsRpcMessage, create_payload_dict
 
 
@@ -23,32 +25,42 @@ class DaemonProxy:
         max_message_size: int = 50 * 1000 * 1000,
     ):
         self._uri = uri
-        self._request_dict: Dict[str, asyncio.Event] = {}
-        self.response_dict: Dict[str, WsRpcMessage] = {}
+        self._request_dict: dict[str, asyncio.Event] = {}
+        self.response_dict: dict[str, WsRpcMessage] = {}
         self.ssl_context = ssl_context
         self.heartbeat = heartbeat
         self.client_session: Optional[aiohttp.ClientSession] = None
         self.websocket: Optional[aiohttp.ClientWebSocketResponse] = None
         self.max_message_size = max_message_size
 
-    def format_request(self, command: str, data: Dict[str, Any]) -> WsRpcMessage:
+    def format_request(self, command: str, data: dict[str, Any]) -> WsRpcMessage:
         request = create_payload_dict(command, data, "client", "daemon")
         return request
 
-    async def start(self) -> None:
-        try:
-            self.client_session = aiohttp.ClientSession()
-            self.websocket = await self.client_session.ws_connect(
-                self._uri,
-                autoclose=True,
-                autoping=True,
-                heartbeat=self.heartbeat,
-                ssl=self.ssl_context if self.ssl_context is not None else True,
-                max_msg_size=self.max_message_size,
-            )
-        except Exception:
+    async def start(self, wait_for_start: bool = False) -> None:
+        self.client_session = aiohttp.ClientSession()
+
+        connect_backoff = 2
+        while (self.websocket is None or self.websocket.closed) and connect_backoff <= 60:
+            try:
+                self.websocket = await self.client_session.ws_connect(
+                    self._uri,
+                    autoclose=True,
+                    autoping=True,
+                    heartbeat=self.heartbeat,
+                    ssl=self.ssl_context if self.ssl_context is not None else True,
+                    max_msg_size=self.max_message_size,
+                )
+                break
+            except aiohttp.ClientError:
+                if not wait_for_start:
+                    break
+                await asyncio.sleep(connect_backoff)
+                connect_backoff *= 2
+
+        if self.websocket is None or self.websocket.closed:
             await self.close()
-            raise
+            raise Exception("Failed to connect to daemon")
 
         async def listener_task() -> None:
             try:
@@ -56,7 +68,7 @@ class DaemonProxy:
             finally:
                 await self.close()
 
-        asyncio.create_task(listener_task())
+        create_referenced_task(listener_task(), known_unreferenced=True)
         await asyncio.sleep(1)
 
     async def listener(self) -> None:
@@ -80,7 +92,7 @@ class DaemonProxy:
         string = dict_to_json_str(request)
         if self.websocket is None or self.websocket.closed:
             raise Exception("Websocket is not connected")
-        asyncio.create_task(self.websocket.send_str(string))
+        create_referenced_task(self.websocket.send_str(string), known_unreferenced=True)
         try:
             await asyncio.wait_for(self._request_dict[request_id].wait(), timeout=30)
             self._request_dict.pop(request_id)
@@ -92,13 +104,13 @@ class DaemonProxy:
             raise Exception(f"No response from daemon for request_id: {request_id}")
 
     async def get_version(self) -> WsRpcMessage:
-        data: Dict[str, Any] = {}
+        data: dict[str, Any] = {}
         request = self.format_request("get_version", data)
         response = await self._get(request)
         return response
 
     async def get_network_info(self) -> WsRpcMessage:
-        data: Dict[str, Any] = {}
+        data: dict[str, Any] = {}
         request = self.format_request("get_network_info", data)
         response = await self._get(request)
         return response
@@ -124,7 +136,7 @@ class DaemonProxy:
         return False
 
     async def is_keyring_locked(self) -> bool:
-        data: Dict[str, Any] = {}
+        data: dict[str, Any] = {}
         request = self.format_request("is_keyring_locked", data)
         response = await self._get(request)
         if "is_keyring_locked" in response["data"]:
@@ -152,7 +164,7 @@ class DaemonProxy:
         request = self.format_request("exit", {})
         return await self._get(request)
 
-    async def get_keys_for_plotting(self, fingerprints: Optional[List[uint32]] = None) -> WsRpcMessage:
+    async def get_keys_for_plotting(self, fingerprints: Optional[list[uint32]] = None) -> WsRpcMessage:
         data = {"fingerprints": fingerprints} if fingerprints else {}
         request = self.format_request("get_keys_for_plotting", data)
         response = await self._get(request)
@@ -160,7 +172,12 @@ class DaemonProxy:
 
 
 async def connect_to_daemon(
-    self_hostname: str, daemon_port: int, max_message_size: int, ssl_context: ssl.SSLContext, heartbeat: int
+    self_hostname: str,
+    daemon_port: int,
+    max_message_size: int,
+    ssl_context: ssl.SSLContext,
+    heartbeat: int,
+    wait_for_start: bool = False,
 ) -> DaemonProxy:
     """
     Connect to the local daemon.
@@ -172,12 +189,13 @@ async def connect_to_daemon(
         max_message_size=max_message_size,
         heartbeat=heartbeat,
     )
-    await client.start()
+
+    await client.start(wait_for_start=wait_for_start)
     return client
 
 
 async def connect_to_daemon_and_validate(
-    root_path: Path, config: Dict[str, Any], quiet: bool = False
+    root_path: Path, config: dict[str, Any], quiet: bool = False, wait_for_start: bool = False
 ) -> Optional[DaemonProxy]:
     """
     Connect to the local daemon and do a ping to ensure that something is really
@@ -199,6 +217,7 @@ async def connect_to_daemon_and_validate(
             max_message_size=daemon_max_message_size,
             ssl_context=ssl_context,
             heartbeat=daemon_heartbeat,
+            wait_for_start=wait_for_start,
         )
         r = await connection.ping()
 
@@ -213,7 +232,7 @@ async def connect_to_daemon_and_validate(
 
 @asynccontextmanager
 async def acquire_connection_to_daemon(
-    root_path: Path, config: Dict[str, Any], quiet: bool = False
+    root_path: Path, config: dict[str, Any], quiet: bool = False
 ) -> AsyncIterator[Optional[DaemonProxy]]:
     """
     Asynchronous context manager which attempts to create a connection to the daemon.

chia/daemon/keychain_proxy.py

@@ -5,7 +5,7 @@ import logging
 import ssl
 import traceback
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload
+from typing import Any, Literal, Optional, Union, overload
 
 from aiohttp import ClientConnectorError, ClientSession
 from chia_rs import AugSchemeMPL, G1Element, PrivateKey
@@ -31,6 +31,7 @@ from chia.util.errors import (
     KeychainProxyConnectionTimeout,
 )
 from chia.util.keychain import Keychain, KeyData, bytes_to_mnemonic, mnemonic_to_seed
+from chia.util.task_referencer import create_referenced_task
 from chia.util.ws_message import WsRpcMessage
 
 
@@ -71,7 +72,7 @@ class KeychainProxy(DaemonProxy):
         """
         return self.keychain is not None
 
-    def format_request(self, command: str, data: Dict[str, Any]) -> WsRpcMessage:
+    def format_request(self, command: str, data: dict[str, Any]) -> WsRpcMessage:
         """
         Overrides DaemonProxy.format_request() to add keychain-specific RPC params
         """
@@ -98,8 +99,8 @@ class KeychainProxy(DaemonProxy):
         except asyncio.TimeoutError:
             raise KeychainProxyConnectionTimeout()
 
-    async def start(self) -> None:
-        self.keychain_connection_task = asyncio.create_task(self.connect_to_keychain())
+    async def start(self, wait_for_start: bool = False) -> None:
+        self.keychain_connection_task = create_referenced_task(self.connect_to_keychain())
         await self.connection_established.wait()  # wait until connection is established.
 
     async def connect_to_keychain(self) -> None:
@@ -141,7 +142,7 @@ class KeychainProxy(DaemonProxy):
         if self.keychain_connection_task is not None:
             await self.keychain_connection_task
 
-    async def get_response_for_request(self, request_name: str, data: Dict[str, Any]) -> Tuple[WsRpcMessage, bool]:
+    async def get_response_for_request(self, request_name: str, data: dict[str, Any]) -> tuple[WsRpcMessage, bool]:
         request = self.format_request(request_name, data)
         response = await self._get(request)
         success = response["data"].get("success", False)
@@ -252,11 +253,11 @@ class KeychainProxy(DaemonProxy):
             if not success:
                 self.handle_error(response)
 
-    async def get_all_private_keys(self) -> List[Tuple[PrivateKey, bytes]]:
+    async def get_all_private_keys(self) -> list[tuple[PrivateKey, bytes]]:
         """
         Forwards to Keychain.get_all_private_keys()
         """
-        keys: List[Tuple[PrivateKey, bytes]] = []
+        keys: list[tuple[PrivateKey, bytes]] = []
         if self.use_local_keychain():
             keys = self.keychain.get_all_private_keys()
         else:
@@ -415,11 +416,11 @@ class KeychainProxy(DaemonProxy):
                 self.handle_error(response)
         return key_data
 
-    async def get_keys(self, include_secrets: bool = False) -> List[KeyData]:
+    async def get_keys(self, include_secrets: bool = False) -> list[KeyData]:
         """
         Returns all KeyData
         """
-        keys: List[KeyData] = []
+        keys: list[KeyData] = []
         if self.use_local_keychain():
             keys = self.keychain.get_keys(include_secrets)
         else:

chia/daemon/keychain_server.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import logging
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, List, Type
+from typing import Any
 
 from chia_rs import PrivateKey
 
@@ -65,7 +65,7 @@ class GetKeyRequest(Streamable):
 @streamable
 @dataclass(frozen=True)
 class GetKeysResponse(Streamable):
-    keys: List[KeyData]
+    keys: list[KeyData]
 
 
 @streamable
@@ -91,7 +91,7 @@ class GetPublicKeyRequest(Streamable):
 class GetPublicKeyResponse(Streamable):
     key: KeyData
 
-    def to_json_dict(self) -> Dict[str, Any]:
+    def to_json_dict(self) -> dict[str, Any]:
         # Ensure that only approved keys are returned
         approved_keys = ["fingerprint", "public_key", "label"]
         key_dict = self.key.to_json_dict()
@@ -108,9 +108,9 @@ class GetPublicKeysRequest(Streamable):
 @streamable
 @dataclass(frozen=True)
 class GetPublicKeysResponse(Streamable):
-    keys: List[KeyData]
+    keys: list[KeyData]
 
-    def to_json_dict(self) -> Dict[str, Any]:
+    def to_json_dict(self) -> dict[str, Any]:
         # Ensure that only approved keys are returned
         approved_keys = ["fingerprint", "public_key", "label"]
         return {
@@ -149,9 +149,9 @@ class KeychainServer:
     """
 
     _default_keychain: Keychain = field(default_factory=Keychain)
-    _alt_keychains: Dict[str, Keychain] = field(default_factory=dict)
+    _alt_keychains: dict[str, Keychain] = field(default_factory=dict)
 
-    def get_keychain_for_request(self, request: Dict[str, Any]) -> Keychain:
+    def get_keychain_for_request(self, request: dict[str, Any]) -> Keychain:
         """
         Keychain instances can have user and service strings associated with them.
         The keychain backends ultimately point to the same data stores, but the user
@@ -171,7 +171,7 @@ class KeychainServer:
                 self._alt_keychains[key] = keychain
         return keychain
 
-    async def handle_command(self, command: str, data: Dict[str, Any]) -> Dict[str, Any]:
+    async def handle_command(self, command: str, data: dict[str, Any]) -> dict[str, Any]:
         try:
             if command == "add_private_key":
                 data["private"] = True
@@ -208,7 +208,7 @@ class KeychainServer:
             log.exception(e)
             return {"success": False, "error": str(e), "command": command}
 
-    async def add_key(self, request: Dict[str, Any]) -> Dict[str, Any]:
+    async def add_key(self, request: dict[str, Any]) -> dict[str, Any]:
         if self.get_keychain_for_request(request).is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}
 
@@ -245,7 +245,7 @@ class KeychainServer:
 
         return {"success": True, "fingerprint": fingerprint}
 
-    async def check_keys(self, request: Dict[str, Any]) -> Dict[str, Any]:
+    async def check_keys(self, request: dict[str, Any]) -> dict[str, Any]:
         if self.get_keychain_for_request(request).is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}
 
@@ -261,7 +261,7 @@ class KeychainServer:
 
         return {"success": True}
 
-    async def delete_all_keys(self, request: Dict[str, Any]) -> Dict[str, Any]:
+    async def delete_all_keys(self, request: dict[str, Any]) -> dict[str, Any]:
         if self.get_keychain_for_request(request).is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}
 
@@ -269,7 +269,7 @@ class KeychainServer:
 
         return {"success": True}
 
-    async def delete_key_by_fingerprint(self, request: Dict[str, Any]) -> Dict[str, Any]:
+    async def delete_key_by_fingerprint(self, request: dict[str, Any]) -> dict[str, Any]:
         if self.get_keychain_for_request(request).is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}
 
@@ -285,7 +285,7 @@ class KeychainServer:
 
         return {"success": True}
 
-    async def run_request(self, request_dict: Dict[str, Any], request_type: Type[Any]) -> Dict[str, Any]:
+    async def run_request(self, request_dict: dict[str, Any], request_type: type[Any]) -> dict[str, Any]:
         keychain = self.get_keychain_for_request(request_dict)
         if keychain.is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}
@@ -314,8 +314,8 @@ class KeychainServer:
                 "error_details": {"message": str(e)},
             }
 
-    async def get_all_private_keys(self, request: Dict[str, Any]) -> Dict[str, Any]:
-        all_keys: List[Dict[str, Any]] = []
+    async def get_all_private_keys(self, request: dict[str, Any]) -> dict[str, Any]:
+        all_keys: list[dict[str, Any]] = []
         if self.get_keychain_for_request(request).is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}
 
@@ -325,8 +325,8 @@ class KeychainServer:
 
         return {"success": True, "private_keys": all_keys}
 
-    async def get_first_private_key(self, request: Dict[str, Any]) -> Dict[str, Any]:
-        key: Dict[str, Any] = {}
+    async def get_first_private_key(self, request: dict[str, Any]) -> dict[str, Any]:
+        key: dict[str, Any] = {}
         if self.get_keychain_for_request(request).is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}
 
@@ -340,7 +340,7 @@ class KeychainServer:
 
         return {"success": True, "private_key": key}
 
-    async def get_key_for_fingerprint(self, request: Dict[str, Any]) -> Dict[str, Any]:
+    async def get_key_for_fingerprint(self, request: dict[str, Any]) -> dict[str, Any]:
         keychain = self.get_keychain_for_request(request)
         if keychain.is_keyring_locked():
             return {"success": False, "error": KEYCHAIN_ERR_LOCKED}