✨ Implement wake/sake/getAbi LSP method

wake/development/utils.py

@@ -1162,6 +1162,34 @@ def _get_storage_layout(
         return SolcOutputStorageLayout.model_validate(obj._storage_layout)
 
 
+def get_info_from_explorer(addr: str, chain_id: int) -> Dict[str, Any]:
+    u = urlparse(chain_explorer_urls[chain_id].url)
+    config = get_config()
+    api_key = config.api_keys.get(".".join(u.netloc.split(".")[:-1]), None)
+    if api_key is None:
+        raise ValueError(f"Contract not found and API key for {u.netloc} not provided")
+
+    url = (
+        chain_explorer_urls[chain_id].api_url
+        + f"?module=contract&action=getsourcecode&address={addr}&apikey={api_key}"
+    )
+    req = Request(
+        url,
+        headers={
+            "Content-Type": "application/json",
+            "User-Agent": f"wake/{get_package_version('eth-wake')}",
+        },
+    )
+
+    with urlopen(req) as response:
+        parsed = json.loads(response.read().decode("utf-8"))
+
+    if parsed["status"] != "1":
+        raise ValueError(f"Request to {u.netloc} failed: {parsed['result']}")
+
+    return parsed["result"][0]
+
+
 @functools.lru_cache(maxsize=64)
 def _get_storage_layout_from_explorer(
     addr: str, chain_id: int

wake/lsp/methods.py

@@ -139,3 +139,4 @@ class RequestMethodEnum(StrEnum):
     SAKE_SET_LABEL = "wake/sake/setLabel"
     SAKE_GET_BALANCES = "wake/sake/getBalances"
     SAKE_SET_BALANCES = "wake/sake/setBalances"
+    SAKE_GET_ABI = "wake/sake/getAbi"

wake/lsp/sake.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import re
 from collections import ChainMap, defaultdict
 from functools import wraps
@@ -13,10 +14,11 @@
 from wake.config import WakeConfig
 from wake.development.call_trace import CallTrace
 from wake.development.chain_interfaces import AnvilChainInterface
-from wake.development.core import RequestType
+from wake.development.core import RequestType, get_fqn_from_address
 from wake.development.globals import set_config
 from wake.development.json_rpc import JsonRpcError
 from wake.development.transactions import TransactionStatusEnum
+from wake.development.utils import get_logic_contract, chain_explorer_urls, get_info_from_explorer
 from wake.lsp.context import LspContext
 from wake.lsp.exceptions import LspError
 from wake.lsp.lsp_data_model import LspModel
@@ -159,6 +161,18 @@ class SakeLoadStateParams(SakeParams):
     chain_dump: str
 
 
+class SakeGetAbiParams(SakeParams):
+    address: str
+
+
+class SakeGetAbiResult(SakeResult):
+    name: str
+    abi: List
+    proxy_name: Optional[str]
+    proxy_abi: Optional[List]
+    implementation_address: Optional[str]
+
+
 def chain_connected(f):
     @wraps(f)
     async def wrapper(context: SakeContext, params: SakeParams, *args, **kwargs):
@@ -180,6 +194,7 @@ class SakeContext:
     lsp_context: LspContext
     chains: Dict[str, Tuple[Chain, ContextManager]]
     compilation: Dict[str, ContractInfo]
+    name_by_fqn: Dict[str, str]
     abi_by_fqn: Dict[
         str, Dict[Union[bytes, Literal["constructor", "fallback", "receive"]], List]
     ]  # fqn -> ABI
@@ -189,6 +204,7 @@ def __init__(self, lsp_context: LspContext):
         self.lsp_context = lsp_context
         self.chains = {}
         self.compilation = {}
+        self.name_by_fqn = {}
         self.abi_by_fqn = {}
         self.libraries = {}
 
@@ -332,6 +348,7 @@ async def compile(self) -> SakeCompilationResult:
                 skipped_source_units,
             ) = await self.lsp_context.compiler.bytecode_compile()
 
+            self.name_by_fqn.clear()
             self.abi_by_fqn.clear()
             self.libraries.clear()
             fqn_by_metadata: Dict[bytes, str] = {}
@@ -367,6 +384,7 @@ async def compile(self) -> SakeCompilationResult:
                     self.libraries[lib_id] = fqn
 
                 assert info.abi is not None
+                self.name_by_fqn[fqn] = contract_node["name"]
                 self.abi_by_fqn[fqn] = {}
                 for item in info.abi:
                     if item["type"] == "function":
@@ -707,3 +725,54 @@ async def set_label(self, params: SakeSetLabelParams) -> SakeResult:
             return SakeResult(success=True)
         except Exception as e:
             raise LspError(ErrorCodes.InternalError, str(e)) from None
+
+    @chain_connected
+    async def get_abi(self, params: SakeGetAbiParams) -> SakeGetAbiResult:
+        def info_from_address(address: str) -> Tuple[str, List]:
+            fqn = get_fqn_from_address(logic_contract.address, "latest", chain)
+            if fqn is not None:
+                name = self.name_by_fqn[fqn]
+                abi = list(self.abi_by_fqn[fqn].values())
+            elif chain._forked_chain_id in chain_explorer_urls:
+                try:
+                    info = get_info_from_explorer(address, chain._forked_chain_id)
+                    name = info["ContractName"]
+                    abi = json.loads(info["ABI"])
+                except Exception:
+                    name = ""
+                    abi = []
+            else:
+                name = ""
+                abi = []
+
+            return name, abi
+
+        chain = self.chains[params.session_id][0]
+
+        try:
+            contract = Account(params.address, chain=chain)
+            logic_contract = get_logic_contract(contract)
+            name, abi = info_from_address(str(logic_contract.address))
+
+            if contract != logic_contract:
+                return SakeGetAbiResult(
+                    success=True,
+                    name=name,
+                    abi=abi,
+                    proxy_name=None,
+                    proxy_abi=None,
+                    implementation_address=None,
+                )
+            else:
+                proxy_name, proxy_abi = info_from_address(params.address)
+
+                return SakeGetAbiResult(
+                    success=True,
+                    name=name,
+                    abi=abi,
+                    proxy_name=proxy_name,
+                    proxy_abi=proxy_abi,
+                    implementation_address=str(logic_contract.address),
+                )
+        except Exception as e:
+            raise LspError(ErrorCodes.InternalError, str(e)) from None

wake/lsp/server.py

@@ -134,6 +134,7 @@
     SakeContext,
     SakeCreateChainParams,
     SakeDeployParams,
+    SakeGetAbiParams,
     SakeGetBalancesParams,
     SakeLoadStateParams,
     SakeParams,
@@ -403,6 +404,14 @@ def __init__(
                 ),
                 SakeSetBalancesParams,
             ),
+            RequestMethodEnum.SAKE_GET_ABI: (
+                lambda params: (
+                    self.__sake_context.get_abi(  # pyright: ignore reportAttributeAccessIssue
+                        params,
+                    )
+                ),
+                SakeGetAbiParams,
+            ),
         }
 
         self.__notification_mapping = {