diff --git a/extension/src/wallet/index.ts b/extension/src/wallet/index.ts
index 48c32b5..b07a5fe 100644
--- a/extension/src/wallet/index.ts
+++ b/extension/src/wallet/index.ts
@@ -12,9 +12,12 @@ export class Server {
   async setWallet(wallet: string) {
     this.wallet = wallet;
     await wasm;
-    this.client = await new wasm.Client(await wasm.Wallet.create(wallet));
-    this.client.on_notification((notification: any) => {
-      console.debug('got notification for', this.subscribers.size, 'subscribers:', notification);
+    
+    const clientInstance = new wasm.Client(await wasm.Wallet.create(wallet));
+    this.client = clientInstance;
+    
+    clientInstance.on_notification((notification: Notification) => {
+      console.debug('Received notification for', this.subscribers.size, 'subscribers:', notification);
       for (const subscriber of this.subscribers.values()) {
         subscriber.postMessage(notification);
       }
@@ -27,34 +30,47 @@ export class Server {
       return;
     }
 
-    let target;
-    if (sender.origin === self.location.origin) {
-      console.log('Received call from extension', functionName, args);
-      target = this.client;
-    } else {
-      console.log('Received call from content script', functionName, args);
-      target = this.client.frontend();
-    }
-
-    if (!(functionName in target)) {
-      console.error('Attempted to call undefined function', functionName);
-      return;
-    }
+    try {
+      let target;
+      if (sender.origin === self.location.origin) {
+        console.log('Received call from extension', functionName, args);
+        target = this.client;
+      } else {
+        console.log('Received call from content script', functionName, args);
+        target = this.client.frontend();
+      }
 
-    let func = target[functionName as keyof typeof target] as (...args: any) => Promise<any>;
+      if (typeof target[functionName] !== 'function') {
+        console.error('Attempted to call undefined function', functionName);
+        return;
+      }
 
-    console.debug('Calling function', functionName);
+      const func = target[functionName as keyof typeof target] as (...args: any) => Promise<any>;
 
-    const result = await func.apply(target, args);
-    console.debug('Got result', result);
-    return result;
+      console.debug('Calling function', functionName);
+      const result = await func.apply(target, args);
+      console.debug('Got result', result);
+      return result;
+      
+    } catch (error) {
+      console.error('Error calling client function:', error);
+    }
   }
 
   async init() {
-    await (await wasm).default({
+    await this.loadWasmModule();
+    this.setupNotificationListener();
+    this.setupMessageListener();
+  }
+
+  private async loadWasmModule() {
+    await wasm;
+    await wasm.default({
       module_or_path: (await fetch(wasmModuleUrl)).arrayBuffer(),
     });
+  }
 
+  private setupNotificationListener() {
     chrome.runtime.onConnect.addListener(port => {
       if (port.name !== 'notifications') {
         console.warn('Unknown channel type', port.name);
@@ -62,28 +78,44 @@ export class Server {
       }
 
       this.subscribers.add(port);
-      port.onDisconnect.addListener(port => this.subscribers.delete(port));
+      port.onDisconnect.addListener(() => this.removeSubscriber(port));
     });
+  }
 
+  private setupMessageListener() {
     chrome.runtime.onMessage.addListener((message, sender, respond) => {
-      if (message.target !== 'wallet')
-        return false;
-
-      if (guard.isCallRequest(message)) {
-        this.callClientFunction(sender, message.function, ...message.arguments)
-          .then(respond);
-        return true;
-      } else if (guard.isSetWalletRequest(message))
-        this.setWallet(message.wallet);
-      else if (guard.isGetWalletRequest(message))
-        respond(this.wallet);
-      else
-        console.warn('Unknown message', message);
-
-      return false;
+      if (message.target !== 'wallet') return false;
+
+      switch (true) {
+        case guard.isCallRequest(message):
+          this.handleCallRequest(sender, message, respond);
+          return true;
+
+        case guard.isSetWalletRequest(message):
+          this.setWallet(message.wallet);
+          return true;
+
+        case guard.isGetWalletRequest(message):
+          respond(this.wallet);
+          return true;
+
+        default:
+          console.warn('Unknown message', message);
+          return false;
+      }
     });
   }
 
+  private async handleCallRequest(sender: chrome.runtime.MessageSender, message: any, respond: Function) {
+    const result = await this.callClientFunction(sender, message.function, ...message.arguments);
+    respond(result);
+  }
+
+  private removeSubscriber(port: chrome.runtime.Port) {
+    this.subscribers.delete(port);
+    port.onDisconnect.removeListener(() => this.removeSubscriber(port));
+  }
+
   public static async run() {
     new Server().init();
   }