fix review

web_app/api/user.py

@@ -309,27 +309,20 @@ async def allow_notification(
     response_description="Status of withdrawal operations",
 )
 async def withdraw_all(wallet_id: str) -> dict:
-   """
-   Withdraws all supported tokens from the user's contract.
-
-   :param wallet_id: The wallet ID of the user.
-   :return: detail: "Successfully initiated withdrawals for all tokens"
-   """
-   # Get user's contract address
-   contract_address = await get_user_contract(wallet_id)
-   if not contract_address:
+    """
+    Withdraws all supported tokens from the user's contract.
+
+    :param wallet_id: The wallet ID of the user.
+    :return: detail: "Successfully initiated withdrawals for all tokens"
+    """
+    # Get user's contract address
+    contract_address = user_db.get_contract_address_by_wallet_id(wallet_id)
+    if not contract_address:
        raise HTTPException(status_code=404, detail="Contract not found")
 
-   try:
-       # Perform withdrawals
-       results = await CLIENT.withdraw_all(contract_address)
-       return {
+    # Perform withdrawals
+    results = await CLIENT.withdraw_all(contract_address)
+    return {
            "detail": "Successfully initiated withdrawals for all tokens",
            "results": results
-       }
-   except Exception as e:
-       logger.error(f"Contract withdrawal failed: {str(e)}")
-       raise HTTPException(
-           status_code=500, 
-           detail=f"Failed to withdraw tokens: {str(e)}"
-       )
+    }

web_app/contract_tools/blockchain_call.py

@@ -362,22 +362,25 @@ async def withdraw_all(self, contract_address: str) -> dict[str, str]:
         contract_addr_int = self._convert_address(contract_address)
         results = {}
 
-        for token in TokenParams.tokens():  
+        for token in TokenParams.tokens():
+            token_symbol = token.name
+
             try:
                 token_addr_int = self._convert_address(token.address)  
-                token_symbol = token.name
+
+            except ValueError as e:
+                logger.error(f"Invalid address format for {token_symbol}: {str(e)}")
+                results[token_symbol] = f"Failed: Invalid address format"
+                continue
+
+            try:
                 logger.info(f"Withdrawing {token_symbol} from contract {contract_address}")
                 await self._func_call(
                     addr=contract_addr_int,
                     selector="withdraw",
                     calldata=[token_addr_int, 0] 
                 )
                 results[token_symbol] = "Success"
-
-            except ValueError as e:
-                logger.error(f"Invalid address format for {token_symbol}: {str(e)}")
-                results[token_symbol] = f"Failed: Invalid address format"
-
             except Exception as e:
                 logger.error(f"Error withdrawing {token_symbol}: {e}")
                 results[token_symbol] = f"Failed: {e}"