Implement RPC calls in Jupyter/Colab

boa/__init__.py

@@ -44,6 +44,16 @@ def set_env(new_env):
     Env._singleton = new_env
 
 
+def set_browser_env(address=None):
+    """Set the environment to use the browser's network in Jupyter/Colab"""
+    # import locally because jupyter is generally not installed
+    from boa.integrations.jupyter import BrowserRPC, BrowserSigner
+
+    env = NetworkEnv(rpc=BrowserRPC())
+    env.set_eoa(BrowserSigner(address))
+    set_env(env)
+
+
 def set_network_env(url):
     """Set the environment to use a custom network URL"""
     set_env(NetworkEnv.from_url(url))

boa/integrations/jupyter/__init__.py

@@ -1,6 +1,6 @@
+from boa.integrations.jupyter.browser import BrowserRPC, BrowserSigner
 from boa.integrations.jupyter.constants import PLUGIN_NAME
 from boa.integrations.jupyter.handlers import setup_handlers
-from boa.integrations.jupyter.signer import BrowserSigner
 
 
 def load_jupyter_server_extension(server_app):
@@ -16,4 +16,9 @@ def load_jupyter_server_extension(server_app):
 _load_jupyter_server_extension = load_jupyter_server_extension
 
 
-__all__ = [BrowserSigner, load_jupyter_server_extension, _load_jupyter_server_extension]
+__all__ = [
+    BrowserSigner,
+    BrowserRPC,
+    load_jupyter_server_extension,
+    _load_jupyter_server_extension,
+]

boa/integrations/jupyter/browser.py

@@ -0,0 +1,176 @@
+"""
+This module implements the BrowserSigner class, which is used to sign transactions
+in IPython/JupyterLab/Google Colab.
+"""
+import json
+import logging
+from asyncio import get_running_loop, sleep
+from itertools import chain
+from multiprocessing.shared_memory import SharedMemory
+from os import urandom
+from typing import Any, Awaitable
+
+import nest_asyncio
+from IPython.display import Javascript, display
+
+from boa.rpc import RPC, RPCError
+
+from .constants import (
+    ADDRESS_TIMEOUT_MESSAGE,
+    CALLBACK_TOKEN_BYTES,
+    CALLBACK_TOKEN_TIMEOUT,
+    NUL,
+    PLUGIN_NAME,
+    RPC_TIMEOUT_MESSAGE,
+    SHARED_MEMORY_LENGTH,
+    TRANSACTION_TIMEOUT_MESSAGE,
+)
+from .utils import convert_frontend_dict, install_jupyter_javascript_triggers
+
+try:
+    from google.colab.output import eval_js as colab_eval_js
+except ImportError:
+    colab_eval_js = None  # not in Google Colab, use SharedMemory instead
+
+
+nest_asyncio.apply()
+
+
+class BrowserSigner:
+    """
+    A BrowserSigner is a class that can be used to sign transactions in IPython/JupyterLab.
+    """
+
+    def __init__(self, address=None):
+        """
+        Create a BrowserSigner instance.
+        :param address: The account address. If not provided, it will be requested from the browser.
+        """
+        if address:
+            self.address = address
+        else:
+            self.address = _javascript_call(
+                "loadSigner", timeout_message=ADDRESS_TIMEOUT_MESSAGE
+            )
+
+    def send_transaction(self, tx_data: dict) -> dict:
+        """
+        Implements the Account class' send_transaction method.
+        It executes a Javascript snippet that requests the user's signature for the transaction.
+        Then, it waits for the signature to be received via the API.
+        :param tx_data: The transaction data to sign.
+        :return: The signed transaction data.
+        """
+        sign_data = _javascript_call(
+            "signTransaction", tx_data, timeout_message=TRANSACTION_TIMEOUT_MESSAGE
+        )
+        return convert_frontend_dict(sign_data)
+
+
+class BrowserRPC(RPC):
+    """
+    An RPC object that sends requests to the browser via Javascript.
+    """
+
+    @property
+    def identifier(self) -> str:
+        return type(self).__name__  # every instance does the same
+
+    @property
+    def name(self):
+        return self.identifier
+
+    def fetch(self, method: str, params: Any) -> Any:
+        if method == "eth_getTransactionReceipt":
+            # we do the polling in the browser to avoid too many callbacks
+            # each callback generates currently 10px empty space in the frontend
+            timeout_ms = CALLBACK_TOKEN_TIMEOUT.total_seconds() * 1000
+            return _javascript_call(
+                "waitForTransactionReceipt",
+                params,
+                timeout_ms,
+                timeout_message=RPC_TIMEOUT_MESSAGE,
+            )
+
+        return _javascript_call(
+            "rpc", method, params, timeout_message=RPC_TIMEOUT_MESSAGE
+        )
+
+    def fetch_multi(self, payloads: list[tuple[str, Any]]) -> list[Any]:
+        return _javascript_call(
+            "multiRpc", payloads, timeout_message=RPC_TIMEOUT_MESSAGE
+        )
+
+
+def _javascript_call(js_func: str, *args, timeout_message: str) -> Any:
+    """
+    This function attempts to call a Javascript function in the browser and then
+    wait for the result to be sent back to the API.
+    - Inside Google Colab, it uses the eval_js function to call the Javascript function.
+    - Outside, it uses a SharedMemory object and polls until the frontend called our API.
+    A custom timeout message is useful for user feedback.
+    :param snippet: A function that given a token and some kwargs, returns a Javascript snippet.
+    :param kwargs: The arguments to pass to the Javascript snippet.
+    :return: The result of the Javascript snippet sent to the API.
+    """
+    install_jupyter_javascript_triggers()
+
+    token = _generate_token()
+    args_str = ", ".join(json.dumps(p) for p in chain([token], args))
+    js_code = f"window._titanoboa.{js_func}({args_str})"
+    # logging.warning(f"Calling {js_func} with {args_str}")
+
+    if colab_eval_js:
+        result = colab_eval_js(js_code)
+        return _parse_js_result(json.loads(result))
+
+    memory = SharedMemory(name=token, create=True, size=SHARED_MEMORY_LENGTH)
+    logging.info(f"Waiting for {token}")
+    try:
+        memory.buf[:1] = NUL
+        display(Javascript(js_code))
+        return _wait_buffer_set(memory.buf, timeout_message)
+    finally:
+        memory.unlink()  # get rid of the SharedMemory object after it's been used
+
+
+def _generate_token():
+    """Generate a secure unique token to identify the SharedMemory object."""
+    return f"{PLUGIN_NAME}_{urandom(CALLBACK_TOKEN_BYTES).hex()}"
+
+
+def _wait_buffer_set(buffer: memoryview, timeout_message: str) -> Any:
+    """
+    Wait for the SharedMemory object to be filled with data.
+    :param buffer: The buffer to wait for.
+    :param timeout_message: The message to show if the timeout is reached.
+    :return: The contents of the buffer.
+    """
+
+    async def _async_wait(deadline: float) -> Awaitable[dict[str, Any]]:
+        inner_loop = get_running_loop()
+        while buffer.tobytes().startswith(NUL):
+            if inner_loop.time() > deadline:
+                raise TimeoutError(timeout_message)
+            await sleep(0.01)
+
+        message_bytes = buffer.tobytes().split(NUL)[0]
+        return json.loads(message_bytes.decode())
+
+    loop = get_running_loop()
+    future = _async_wait(deadline=loop.time() + CALLBACK_TOKEN_TIMEOUT.total_seconds())
+    task = loop.create_task(future)
+    loop.run_until_complete(task)
+    return _parse_js_result(task.result())
+
+
+def _parse_js_result(result: dict) -> Any:
+    if "data" in result:
+        return result["data"]
+
+    # raise the error in the Jupyter cell so that the user can see it
+    error = result["error"]
+    error = error.get("info", error).get("error", error)
+    raise RPCError(
+        message=error.get("message", error), code=error.get("code", "CALLBACK_ERROR")
+    )

boa/integrations/jupyter/constants.py

@@ -1,13 +1,14 @@
 from datetime import timedelta
 
-CALLBACK_TOKEN_TIMEOUT = timedelta(minutes=3)
-MEMORY_LENGTH = 50 * 1024  # Size of the shared memory object
-CALLBACK_TOKEN_BYTES = 32
 NUL = b"\0"
+CALLBACK_TOKEN_TIMEOUT = timedelta(minutes=4)  # matches default TransactionSettings
+SHARED_MEMORY_LENGTH = 50 * 1024 + len(NUL)  # Size of the shared memory object
+CALLBACK_TOKEN_BYTES = 32
 ETHERS_JS_URL = "https://cdnjs.cloudflare.com/ajax/libs/ethers/6.9.0/ethers.umd.min.js"
 PLUGIN_NAME = "titanoboa_jupyterlab"
 TOKEN_REGEX = rf"{PLUGIN_NAME}_[0-9a-fA-F]{{{CALLBACK_TOKEN_BYTES * 2}}}"
 TRANSACTION_TIMEOUT_MESSAGE = (
     "Timeout waiting for user to confirm transaction in the browser wallet plug-in."
 )
 ADDRESS_TIMEOUT_MESSAGE = "Timeout loading browser browser wallet plug-in."
+RPC_TIMEOUT_MESSAGE = "Timeout waiting for response from RPC."

boa/integrations/jupyter/handlers.py

@@ -10,7 +10,7 @@
 from jupyter_server.utils import url_path_join
 from tornado.web import authenticated
 
-from boa.integrations.jupyter.constants import PLUGIN_NAME, TOKEN_REGEX
+from boa.integrations.jupyter.constants import NUL, PLUGIN_NAME, TOKEN_REGEX
 
 
 class CallbackHandler(APIHandler):
@@ -37,11 +37,13 @@ def post(self, token: str):
             return self.finish({"error": error})
 
         try:
-            body += b"\0"  # mark the end of the buffer
-            memory.buf[: len(body)] = body
+            memory.buf[: len(body) + 1] = body + NUL
         except ValueError:
             self.set_status(HTTPStatus.REQUEST_ENTITY_TOO_LARGE)
-            error = f"Request body has {len(body)} bytes, but only {memory.size} are allowed"
+            max_len = memory.size - len(NUL)
+            error = (
+                f"Request body has {len(body)} bytes, but only {max_len} are allowed"
+            )
             return self.finish({"error": error})
         finally:
             memory.close()

boa/integrations/jupyter/jupyter.js

@@ -3,62 +3,96 @@
  * BrowserSigner to the frontend.
  */
 (() => {
+    let provider; // cache the provider to avoid re-creating it every time
     const getEthersProvider = () => {
+        if (provider) return provider;
         const {ethereum} = window;
         if (!ethereum) {
-            throw new Error('No Ethereum browser plugin found');
+            throw new Error('No Ethereum plugin found. Please authorize the site on your browser wallet.');
         }
-        return new ethers.BrowserProvider(ethereum);
+        return provider = new ethers.BrowserProvider(ethereum);
     };
+
+    /** Stringify data, converting big ints to strings */
     const stringify = (data) => JSON.stringify(data, (_, v) => (typeof v === 'bigint' ? v.toString() : v));
+
+    /** Get the value of a cookie with the given name */
     const getCookie = (name) => (document.cookie.match(`\\b${name}=([^;]*)\\b`))?.[1];
-    const parsePromise = promise =>
-        promise.then(data => ({data})).catch(e => {
-            console.error(e.stack || e.message);
-            return {error: e.message};
-        });
+
+    /** Converts a success/failed promise into an object with either a data or error field */
+    const parsePromise = promise => promise.then(data => ({data})).catch(error => ({
+        error: Object.keys(error).length ? error : {
+            message: error.message, // the default error object doesn't have enumerable properties
+            stack: error.stack
+        }
+    }));
+
+    /** Async sleep for the given time */
+    const sleep = time => new Promise(resolve => setTimeout(resolve, time));
 
     const colab = window.colab ?? window.google?.colab; // in the parent window or in an iframe
     /** Calls the callback endpoint with the given token and body */
     async function callbackAPI(token, body) {
         const headers = {['X-XSRFToken']: getCookie('_xsrf')};
-        const init = {method: 'POST', body: stringify(body), headers};
+        const init = {method: 'POST', body, headers};
         const url = `../titanoboa_jupyterlab/callback/${token}`;
         const response = await fetch(url, init);
         return response.text();
     }
 
     /** Load the signer via ethers user */
     const loadSigner = async () => {
-        const provider = getEthersProvider();
-        console.log(`Loading the user's signer`);
-        const signer = await provider.getSigner();
+        const signer = await getEthersProvider().getSigner();
         return signer.getAddress();
     };
 
     /** Sign a transaction via ethers */
     async function signTransaction(transaction) {
-        const provider = getEthersProvider();
-        console.log('Starting to sign transaction');
-        const signer = await provider.getSigner();
+        const signer = await getEthersProvider().getSigner();
         return signer.sendTransaction(transaction);
     }
 
+    /** Call an RPC method via ethers */
+    const rpc = (method, params) => getEthersProvider().send(method, params);
+
+    /** Wait until the transaction is mined */
+    const waitForTransactionReceipt = async (params, timeout, wait = 1000) => {
+        try {
+            const result = await rpc('eth_getTransactionReceipt', params);
+            if (result) {
+                return result;
+            }
+        } catch (err) { // ignore "server error" (happens while transaction is mined)
+            if (err?.info?.error?.code !== -32603) {
+                throw err;
+            }
+        }
+        if (timeout < wait) {
+            throw new Error('Timeout waiting for transaction receipt');
+        }
+        await sleep(wait);
+        return waitForTransactionReceipt(params, timeout - wait, wait);
+    };
+
+    /** Call multiple RPCs in sequence */
+    const multiRpc = (payloads) => payloads.reduce(
+        async (previousPromise, [method, params]) => [...await previousPromise, await rpc(method, params)],
+        [],
+    );
+
     /** Call the backend when the given function is called, handling errors */
     const handleCallback = func => async (token, ...args) => {
-        const body = await parsePromise(func(...args));
-        if (colab) {
-            // Colab expects the response to be JSON
-            return JSON.stringify(body);
-        }
-        const responseText = await callbackAPI(token, body);
-        console.log(`Callback ${token} => ${responseText}`);
+        const body = stringify(await parsePromise(func(...args)));
+        // console.log(`Boa: ${func.name}(${args.map(a => JSON.stringify(a)).join(',')}) = ${body};`);
+        return colab ? body : callbackAPI(token, body);
     };
 
-    console.log(`Registering Boa callbacks`);
     // expose functions to window, so they can be called from the BrowserSigner
     window._titanoboa = {
         loadSigner: handleCallback(loadSigner),
-        signTransaction: handleCallback(signTransaction)
+        signTransaction: handleCallback(signTransaction),
+        waitForTransactionReceipt: handleCallback(waitForTransactionReceipt),
+        rpc: handleCallback(rpc),
+        multiRpc: handleCallback(multiRpc),
     };
 })();