Better support for await

lib/core/modules/coderunner/codeRunner.js

@@ -1,4 +1,5 @@
 const RunCode = require('./runCode.js');
+const Utils = require('../../../utils/utils');
 
 class CodeRunner {
   constructor(options) {
@@ -59,19 +60,23 @@ class CodeRunner {
     this.runCode.registerVar(varName, code);
   }
 
-  evalCode(code, cb, forConsoleOnly = false) {
+  async evalCode(code, cb, forConsoleOnly = false) {
     cb = cb || function() {};
     const awaitIdx = code.indexOf('await');
+    let awaiting = false;
+
     if (awaitIdx > -1) {
-      if (awaitIdx < 2) {
-        let end = code.length;
-        if (code[end - 1] === ';') {
-          end--; // Remove the `;` because we add function calls
-        }
-        code = code.substring(5, end); // remove await keyword
+      awaiting = true;
+      const instructions = Utils.compact(code.split(';'));
+      const last = instructions.pop();
+
+      if (!last.trim().startsWith('return')) {
+        instructions.push(`return ${last}`); 
       } else {
-        code = `(async function() {${code}})();`;
+        instructions.push(last); 
       }
+      
+      code = `(async function() {${instructions.join(';')}})();`;
     }
     let result = this.runCode.doEval(code);
 
@@ -80,11 +85,16 @@ class CodeRunner {
       this.ipc.broadcast("runcode:newCommand", {code});
     }
 
-    if (result instanceof Promise && !result.on) {
-      return result.then((value) => { cb(null, value); }).catch(cb);
+    if (!awaiting) {
+      return cb(null, result);
     }
 
-    cb(null, result);
+    try {
+      const value = await result;
+      cb(null, value);
+    } catch (error) {
+      cb(error);  
+    }
   }
 }