refactor: align controller package with StarkNet wallet API spec

packages/controller/src/account.ts

@@ -12,11 +12,13 @@ import {
   KeychainOptions,
   Modal,
   ResponseCodes,
+  createUserRejectedError,
 } from "./types";
 import { AsyncMethodReturns } from "@cartridge/penpal";
 import BaseProvider from "./provider";
 import { toArray } from "./utils";
-import { SIGNATURE } from "@starknet-io/types-js";
+import { Signature } from "@starknet-io/types-js";
+import { KeychainIFrame } from "./iframe";
 
 class ControllerAccount extends WalletAccount {
   private keychain: AsyncMethodReturns<Keychain>;
@@ -84,22 +86,50 @@ class ControllerAccount extends WalletAccount {
       // Session call or Paymaster flow failed.
       // Session not avaialble, manual flow fallback
       this.modal.open();
-      const manualExecute = await this.keychain.execute(
-        calls,
-        undefined,
-        undefined,
-        true,
-        (sessionExecute as ConnectError).error,
-      );
 
-      // Manual call succeeded
-      if (manualExecute.code === ResponseCodes.SUCCESS) {
-        resolve(manualExecute as InvokeFunctionResponse);
-        this.modal.close();
-        return;
+      // Check if modal supports cancellation (is a KeychainIFrame)
+      if (this.modal instanceof KeychainIFrame) {
+        const executePromise = this.keychain.execute(
+          calls,
+          undefined,
+          undefined,
+          true,
+          (sessionExecute as ConnectError).error,
+        );
+
+        const resultPromise = this.modal.withCancellation(executePromise, () =>
+          reject(createUserRejectedError()),
+        );
+
+        const manualExecute = await resultPromise;
+
+        // Manual call succeeded
+        if (manualExecute.code === ResponseCodes.SUCCESS) {
+          resolve(manualExecute as InvokeFunctionResponse);
+          this.modal.close();
+          return;
+        }
+
+        reject((manualExecute as ConnectError).error);
+      } else {
+        // Fallback for non-cancellable modals
+        const manualExecute = await this.keychain.execute(
+          calls,
+          undefined,
+          undefined,
+          true,
+          (sessionExecute as ConnectError).error,
+        );
+
+        // Manual call succeeded
+        if (manualExecute.code === ResponseCodes.SUCCESS) {
+          resolve(manualExecute as InvokeFunctionResponse);
+          this.modal.close();
+          return;
+        }
+
+        reject((manualExecute as ConnectError).error);
       }
-
-      reject((manualExecute as ConnectError).error);
       return;
     });
   }
@@ -112,26 +142,50 @@ class ControllerAccount extends WalletAccount {
    * @returns the signature of the JSON object
    * @throws {Error} if the JSON object is not a valid JSON
    */
-  async signMessage(typedData: TypedData): Promise<SIGNATURE> {
+  async signMessage(typedData: TypedData): Promise<Signature> {
     return new Promise(async (resolve, reject) => {
       const sessionSign = await this.keychain.signMessage(typedData, "", true);
 
       // Session sign succeeded
       if (!("code" in sessionSign)) {
-        resolve(sessionSign as SIGNATURE);
+        resolve(sessionSign as Signature);
         return;
       }
 
       // Session not avaialble, manual flow fallback
       this.modal.open();
-      const manualSign = await this.keychain.signMessage(typedData, "", false);
 
-      if (!("code" in manualSign)) {
-        resolve(manualSign as SIGNATURE);
+      // Check if modal supports cancellation (is a KeychainIFrame)
+      if (this.modal instanceof KeychainIFrame) {
+        const signPromise = this.keychain.signMessage(typedData, "", false);
+
+        const resultPromise = this.modal.withCancellation(signPromise, () =>
+          reject(createUserRejectedError()),
+        );
+
+        const manualSign = await resultPromise;
+
+        if (!("code" in manualSign)) {
+          resolve(manualSign as Signature);
+        } else {
+          reject((manualSign as ConnectError).error);
+        }
+        this.modal.close();
       } else {
-        reject((manualSign as ConnectError).error);
+        // Fallback for non-cancellable modals
+        const manualSign = await this.keychain.signMessage(
+          typedData,
+          "",
+          false,
+        );
+
+        if (!("code" in manualSign)) {
+          resolve(manualSign as Signature);
+        } else {
+          reject((manualSign as ConnectError).error);
+        }
+        this.modal.close();
       }
-      this.modal.close();
     });
   }
 }

packages/controller/src/controller.ts

@@ -23,6 +23,7 @@ import {
   ProfileContextTypeVariant,
   ResponseCodes,
   StarterPack,
+  createUserRejectedError,
 } from "./types";
 import { parseChainId } from "./utils";
 
@@ -32,6 +33,7 @@ export default class ControllerProvider extends BaseProvider {
   private iframes: IFrames;
   private selectedChain: ChainId;
   private chains: Map<ChainId, Chain>;
+  private pendingConnectRejection?: (error: any) => void;
 
   isReady(): boolean {
     return !!this.keychain;
@@ -164,8 +166,11 @@ export default class ControllerProvider extends BaseProvider {
 
     this.iframes.keychain.open();
 
-    try {
-      let response = await this.keychain.connect(
+    return new Promise<WalletAccount | undefined>((resolve, reject) => {
+      // Track this pending connection for cancellation
+      this.pendingConnectRejection = reject;
+
+      this.keychain!.connect(
         // Policy precedence logic:
         // 1. If shouldOverridePresetPolicies is true and policies are provided, use policies
         // 2. Otherwise, if preset is defined, use empty object (let preset take precedence)
@@ -177,27 +182,38 @@ export default class ControllerProvider extends BaseProvider {
             : this.options.policies || {},
         this.rpcUrl(),
         this.options.signupOptions,
-      );
-      if (response.code !== ResponseCodes.SUCCESS) {
-        throw new Error(response.message);
-      }
-
-      response = response as ConnectReply;
-      this.account = new ControllerAccount(
-        this,
-        this.rpcUrl(),
-        response.address,
-        this.keychain,
-        this.options,
-        this.iframes.keychain,
-      );
+      )
+        .then((response) => {
+          // Clear the pending rejection handler
+          this.pendingConnectRejection = undefined;
+
+          if (response.code !== ResponseCodes.SUCCESS) {
+            this.iframes.keychain!.close();
+            reject(new Error(response.message));
+            return;
+          }
+
+          const connectReply = response as ConnectReply;
+          this.account = new ControllerAccount(
+            this,
+            this.rpcUrl(),
+            connectReply.address,
+            this.keychain!,
+            this.options,
+            this.iframes.keychain!,
+          );
 
-      return this.account;
-    } catch (e) {
-      console.log(e);
-    } finally {
-      this.iframes.keychain.close();
-    }
+          this.iframes.keychain!.close();
+          resolve(this.account);
+        })
+        .catch((error) => {
+          // Clear the pending rejection handler
+          this.pendingConnectRejection = undefined;
+          this.iframes.keychain!.close();
+          console.log(error);
+          reject(error);
+        });
+    });
   }
 
   async switchStarknetChain(chainId: string): Promise<boolean> {
@@ -443,6 +459,13 @@ export default class ControllerProvider extends BaseProvider {
     return new KeychainIFrame({
       ...this.options,
       onClose: this.keychain?.reset,
+      onCancel: () => {
+        // User cancelled by clicking outside - reject pending operations
+        if (this.pendingConnectRejection) {
+          this.pendingConnectRejection(createUserRejectedError());
+          this.pendingConnectRejection = undefined;
+        }
+      },
       onConnect: (keychain) => {
         this.keychain = keychain;
       },

packages/controller/src/iframe/base.ts

@@ -14,6 +14,7 @@ export class IFrame<CallSender extends {}> implements Modal {
   private iframe?: HTMLIFrameElement;
   private container?: HTMLDivElement;
   private onClose?: () => void;
+  private onCancel?: () => void;
   private child?: AsyncMethodReturns<CallSender>;
   private closeTimeout?: NodeJS.Timeout;
 
@@ -22,12 +23,14 @@ export class IFrame<CallSender extends {}> implements Modal {
     url,
     preset,
     onClose,
+    onCancel,
     onConnect,
     methods = {},
   }: Pick<ControllerOptions, "preset"> & {
     id: string;
     url: URL;
     onClose?: () => void;
+    onCancel?: () => void;
     onConnect: (child: AsyncMethodReturns<CallSender>) => void;
     methods?: { [key: string]: (...args: any[]) => void };
   }) {
@@ -76,13 +79,16 @@ export class IFrame<CallSender extends {}> implements Modal {
     // Add click event listener to close iframe when clicking outside
     container.addEventListener("click", (e) => {
       if (e.target === container) {
+        // User clicked outside - this is a cancellation
         // Attempting to reset(clear context) for keychain iframe (identified by ID)
         if (id === "controller-keychain" && this.child) {
           // Type assertion for keychain child only
           (this.child as any)
             .reset?.()
             .catch((e: any) => console.error("Error resetting context:", e));
         }
+        // Call onCancel to notify about user cancellation
+        this.onCancel?.();
         this.close();
       }
     });
@@ -129,6 +135,7 @@ export class IFrame<CallSender extends {}> implements Modal {
     }
 
     this.onClose = onClose;
+    this.onCancel = onCancel;
   }
 
   open() {
@@ -201,4 +208,24 @@ export class IFrame<CallSender extends {}> implements Modal {
   isOpen() {
     return this.container?.style.display !== "none";
   }
+
+  // Register a one-time cancellation handler for the current operation
+  withCancellation<T>(operation: Promise<T>, onCancel: () => void): Promise<T> {
+    // Store the original onCancel
+    const originalOnCancel = this.onCancel;
+
+    // Set a one-time handler that includes the provided callback
+    this.onCancel = () => {
+      onCancel();
+      // Restore the original handler
+      this.onCancel = originalOnCancel;
+      // Call the original if it exists
+      originalOnCancel?.();
+    };
+
+    // When the operation completes (success or failure), restore the original handler
+    return operation.finally(() => {
+      this.onCancel = originalOnCancel;
+    });
+  }
 }

packages/controller/src/iframe/keychain.ts

@@ -6,6 +6,7 @@ import { IFrame, IFrameOptions } from "./base";
 type KeychainIframeOptions = IFrameOptions<Keychain> &
   KeychainOptions & {
     version?: string;
+    onCancel?: () => void;
   };
 
 export class KeychainIFrame extends IFrame<Keychain> {