Update contract clients for authentication

Author: schoibs-virtuals

virtuals_acp/client.py

@@ -5,15 +5,16 @@
 import signal
 import sys
 import threading
-
+import jwt
+import socketio
+import requests
 import time
+
 from datetime import datetime, timezone, timedelta
 from importlib.metadata import version
 from typing import List, Optional, Union, Dict, Any, Callable
-
-import requests
-import socketio
 from web3 import Web3
+from requests.auth import AuthBase
 
 from virtuals_acp.account import ACPAccount
 from virtuals_acp.configs.configs import (
@@ -49,11 +50,6 @@
 
 
 
-import jwt
-import requests
-
-from requests.auth import AuthBase
-
 class BearerAuth(AuthBase):
     def __init__(self, get_access_token: Callable[[], str]):
         self._get_access_token = get_access_token
@@ -70,7 +66,8 @@ def clear_token(self):
 
 
 class ACPApiClient:
-    def __init__(self, acp_url: str, wallet_address: str, require_auth: bool = False):
+    def __init__(self, acp_contract_client: BaseAcpContractClient, acp_url: str, wallet_address: str, require_auth: bool = False):
+        self.acp_contract_client = acp_contract_client
         self.base_url = f"{acp_url}/api"
         self.wallet_address = wallet_address
         self.require_auth = require_auth
@@ -130,13 +127,17 @@ def get_access_token(self) -> str:
                 needs_refresh = True
 
         if not needs_refresh:
-            return self.access_token
+            # Access token is still valid
+            if self.access_token:
+                return self.access_token
+            else:
+                raise Exception("Access token needs refreshing!")
 
         self.access_token = self.refresh_token()
         return self.access_token
 
     def refresh_token(self) -> str:
-        challenge = self.get_auth_challenge(self.wallet_address)
+        challenge = self.get_auth_challenge()
         signature = self.acp_contract_client.sign_typed_data(challenge)
 
         verified = self.verify_auth_challenge(
@@ -204,8 +205,8 @@ def __init__(
                     "All contract clients must have the same agent wallet address"
                 )
 
-        self.acp_client = ACPApiClient(self.acp_url, self.wallet_address)
-        self.no_auth_acp_client = ACPApiClient(self.acp_url, self.wallet_address, require_auth=False)
+        self.acp_client = ACPApiClient(self.acp_contract_client, self.acp_url, self.wallet_address)
+        self.no_auth_acp_client = ACPApiClient(self.acp_contract_client, self.acp_url, self.wallet_address, require_auth=False)
 
         # Socket.IO setup
         self.on_new_task = on_new_task
@@ -235,7 +236,6 @@ def init(self):
         logger.info(f"Initializing socket")
 
         try:
-            # TODO: auth needs to include access token now
             auth_data = {
                 "walletAddress": self.wallet_address,
                 "accessToken": self.acp_client.get_access_token()

virtuals_acp/contract_clients/base_contract_client.py

@@ -121,7 +121,11 @@ def validate_session_key_on_chain(
             )
 
     @abstractmethod
-    def get_acp_version(self) -> str:
+    def get_asset_manager_address(self) -> str:
+        pass
+
+    @abstractmethod
+    def sign_typed_data(self, typed_data: dict[str, Any]) -> str:
         pass
 
     def _build_user_operation(
@@ -151,7 +155,7 @@ def handle_operation(self, trx_data: List[OperationPayload], chain_id: Optional[
 
     @abstractmethod
     def get_job_id(
-        self, receipt: Dict[str, Any], client_address: str, provider_address: str
+        self, response: Dict[str, Any], client_address: str, provider_address: str
     ) -> int:
         """Abstract method to retrieve a job ID from a transaction hash and related addresses."""
         pass

virtuals_acp/contract_clients/contract_client.py

@@ -6,6 +6,7 @@
 from typing import Dict, Any, Optional, List
 
 from eth_account import Account
+from eth_account.messages import encode_typed_data
 from web3 import Web3
 
 from virtuals_acp.alchemy import AlchemyAccountKit
@@ -237,4 +238,22 @@ def perform_x402_request(
             raise ACPError("Failed to perform X402 request", e)
 
     def get_asset_manager_address(self) -> str:
-        raise ACPError("Not Supported")
+        raise ACPError("Not Supported")
+
+    def sign_typed_data(self, typed_data: dict[str, Any]) -> str:
+        domain = typed_data["domain"]
+        types = typed_data["types"]
+        primary_type = typed_data["primaryType"]
+        message = typed_data["message"]
+
+        # encode_typed_data expects (domain_data, types, primary_type, message_data)
+        # It handles EIP-712 hashing internally
+        signable = encode_typed_data(
+            domain,
+            types,
+            primary_type,
+            message,
+        )
+
+        signed = self.account.sign_message(signable)
+        return signed.signature.hex()

virtuals_acp/contract_clients/contract_client_v2.py

@@ -4,6 +4,7 @@
 from typing import Dict, Any, List, Optional
 
 from eth_account import Account
+from eth_account.messages import encode_typed_data
 from web3 import Web3
 
 from virtuals_acp.abis.job_manager import JOB_MANAGER_ABI
@@ -198,4 +199,22 @@ def get_x402_payment_details(self, job_id: int) -> AcpJobX402PaymentDetails:
             raise ACPError("Failed to get X402 payment details", e)
 
     def get_asset_manager_address(self) -> str:
-        return self.memo_manager_contract.functions.assetManager().call()
+        return self.memo_manager_contract.functions.assetManager().call()
+
+    def sign_typed_data(self, typed_data: dict[str, Any]) -> str:
+        domain = typed_data["domain"]
+        types = typed_data["types"]
+        primary_type = typed_data["primaryType"]
+        message = typed_data["message"]
+
+        # encode_typed_data expects (domain_data, types, primary_type, message_data)
+        # It handles EIP-712 hashing internally
+        signable = encode_typed_data(
+            domain,
+            types,
+            primary_type,
+            message,
+        )
+
+        signed = self.account.sign_message(signable)
+        return signed.signature.hex()

virtuals_acp/web3.py

@@ -3,6 +3,7 @@
 from web3 import Web3
 from virtuals_acp.abis.erc20_abi import ERC20_ABI
 
+# TODO: implement wrapper methods in base_contract_client
 
 def getERC20Balance(
     public_client: Web3,