zksync-sso 0.1.0 → 0.3.0

package/src/abi/index.ts

@@ -1,4 +1,5 @@
 export { AAFactoryAbi } from "./AAFactory.js";
-export { GuardianRecoveryModuleAbi } from "./GuardianRecoveryModule.js";
+export { GuardianRecoveryValidatorAbi } from "./GuardianRecoveryValidator.js";
 export { SessionKeyValidatorAbi } from "./SessionKeyValidator.js";
+export { SsoAccountAbi } from "./SsoAccount.js";
 export { WebAuthValidatorAbi } from "./WebAuthValidator.js";

package/src/client/ecdsa/decorators/wallet.ts

@@ -1,6 +1,6 @@
 import { type Account, bytesToHex, type Chain, type ExactPartial, formatTransaction, type RpcTransaction, type Transport, type WalletActions } from "viem";
-import { deployContract, getAddresses, getChainId, sendRawTransaction, signMessage, signTypedData, writeContract } from "viem/actions";
-import { signTransaction, type TransactionRequestEIP712, type ZksyncEip712Meta } from "viem/zksync";
+import { deployContract, getAddresses, getCallsStatus, getCapabilities, getChainId, prepareAuthorization, sendCalls, sendRawTransaction, showCallsStatus, signAuthorization, signMessage, signTransaction, signTypedData, waitForCallsStatus, writeContract } from "viem/actions";
+import { type TransactionRequestEIP712, type ZksyncEip712Meta } from "viem/zksync";
 
 import { getTransactionWithPaymasterData } from "../../../paymaster/index.js";
 import { sendEip712Transaction } from "../../session/actions/sendEip712Transaction.js";
@@ -65,9 +65,16 @@ export function zksyncSsoEcdsaWalletActions<
       return signTransaction(client, {
         ...args,
         unformattedTxWithPaymaster,
-      } as any);
+      });
     },
     signTypedData: (args) => signTypedData(client, args),
     writeContract: (args) => writeContract(client, args),
+    signAuthorization: (args) => signAuthorization(client, args),
+    getCallsStatus: (args) => getCallsStatus(client, args),
+    getCapabilities: (args) => getCapabilities(client, args),
+    prepareAuthorization: (args) => prepareAuthorization(client, args),
+    sendCalls: (args) => sendCalls(client, args),
+    showCallsStatus: (args) => showCallsStatus(client, args),
+    waitForCallsStatus: (args) => waitForCallsStatus(client, args),
   };
 }

package/src/client/index.ts

@@ -1,5 +1,122 @@
-export { deployAccount as deployEcdsaAccount, fetchAccount as fetchEcdsaAccount } from "./ecdsa/actions/account.js";
-export { deployAccount, fetchAccount } from "./passkey/actions/account.js";
+export { fetchAccount as fetchEcdsaAccount } from "./ecdsa/actions/account.js";
+export { deployAccount, fetchAccount, fetchAccount as fetchPasskeyAccount } from "./passkey/actions/account.js";
 export * from "./recovery/actions/recovery.js";
 export * from "./recovery/client.js";
 export * from "./session/client.js";
+
+import type { Account, Address, Chain, Client, Hash, Hex, Prettify, TransactionReceipt, Transport } from "viem";
+import { concat, getAddress, keccak256, parseEventLogs, toHex } from "viem";
+import { waitForTransactionReceipt, writeContract } from "viem/actions";
+import { getGeneralPaymasterInput } from "viem/zksync";
+
+import { AAFactoryAbi } from "../abi/AAFactory.js";
+import { encodeModuleData, encodeSession } from "../utils/encoding.js";
+import { noThrow } from "../utils/helpers.js";
+import type { SessionConfig } from "../utils/session.js";
+import { type DeployAccountPasskeyArgs, encodePasskeyModuleData } from "./passkey/actions/account.js";
+
+export type DeployAccountPaymasterArgs = {
+  location: Address; // Paymaster address used to pay the fees of creating accounts
+  input?: Hex; // Input for paymaster (if provided)
+};
+export type DeployAccountSessionArgs = {
+  location: Address; // module address
+  initialSession?: SessionConfig;
+};
+
+export type DeployAccountArgs = {
+  accountFactory: Address;
+  installNoDataModules: Address[];
+  owners: Address[]; // ECDSA owners (can be empty)
+  sessionModule?: DeployAccountSessionArgs;
+  paymaster?: DeployAccountPaymasterArgs;
+  passkeyModule?: DeployAccountPasskeyArgs;
+  uniqueAccountId?: string; // Unique account ID, can be omitted if you don't need it
+  onTransactionSent?: (hash: Hash) => void;
+};
+
+export type DeployAccountReturnType = {
+  address: Address;
+  transactionReceipt: TransactionReceipt;
+};
+
+export const deployModularAccount = async <
+  transport extends Transport,
+  chain extends Chain,
+  account extends Account,
+>(
+  client: Client<transport, chain, account>, // Account deployer (any viem client)
+  args: Prettify<DeployAccountArgs>,
+): Promise<DeployAccountReturnType> => {
+  const uniqueIds: Hex[] = [];
+  if (args.uniqueAccountId) {
+    uniqueIds.push(toHex(args.uniqueAccountId));
+  }
+  if (args.owners) {
+    uniqueIds.push(...args.owners);
+  }
+
+  const modules = [];
+  if (args.sessionModule) {
+    modules.push(encodeModuleData({
+      address: args.sessionModule.location,
+      parameters: args.sessionModule.initialSession ? encodeSession(args.sessionModule.initialSession) : "0x",
+    }));
+  }
+  if (args.passkeyModule) {
+    uniqueIds.push(toHex(args.passkeyModule.credentialId));
+    modules.push(await encodePasskeyModuleData(args.passkeyModule));
+  }
+
+  if (args.installNoDataModules) {
+    args.installNoDataModules.forEach((moduleNoDataInstall) => {
+      modules.push(encodeModuleData({
+        address: moduleNoDataInstall,
+        parameters: "0x",
+      }));
+    });
+  }
+
+  let deployProxyArgs = {
+    account: client.account!,
+    chain: client.chain!,
+    address: args.accountFactory,
+    abi: AAFactoryAbi,
+    functionName: "deployProxySsoAccount",
+    args: [
+      keccak256(toHex(concat(uniqueIds))),
+      modules,
+      args.owners,
+    ],
+  } as any;
+
+  if (args.paymaster) {
+    deployProxyArgs = {
+      ...deployProxyArgs,
+      paymaster: args.paymaster.location,
+      paymasterInput: args.paymaster.input ?? getGeneralPaymasterInput({ innerInput: "0x" }),
+    };
+  }
+
+  const transactionHash = await writeContract(client, deployProxyArgs);
+  if (args.onTransactionSent) {
+    noThrow(() => args.onTransactionSent?.(transactionHash));
+  }
+
+  const transactionReceipt = await waitForTransactionReceipt(client, { hash: transactionHash });
+  if (transactionReceipt.status !== "success") throw new Error("Account deployment transaction reverted");
+
+  const accountCreatedEvent = parseEventLogs({ abi: AAFactoryAbi, logs: transactionReceipt.logs })
+    .find((log) => log && log.eventName === "AccountCreated");
+
+  if (!accountCreatedEvent) {
+    throw new Error("No contract address in transaction receipt");
+  }
+
+  const { accountAddress } = accountCreatedEvent.args;
+
+  return {
+    address: getAddress(accountAddress),
+    transactionReceipt: transactionReceipt,
+  };
+};

package/src/client/passkey/account.ts

@@ -1,22 +1,32 @@
 import type { Address } from "abitype";
-import { type CustomSource, type Hash, hashMessage, hashTypedData, type Hex, type LocalAccount } from "viem";
+import { type Chain, type CustomSource, type Hash, hashMessage, hashTypedData, type Hex, type LocalAccount, type Transport } from "viem";
 import { toAccount } from "viem/accounts";
 import { serializeTransaction, type ZksyncTransactionSerializableEIP712 } from "viem/zksync";
 
 import { getEip712Domain } from "../utils/getEip712Domain.js";
+import type { PasskeyRequiredContracts } from "./client.js";
 
-export type ToPasskeyAccountParameters = {
+export type ToPasskeyAccountParameters<
+  transport extends Transport = Transport,
+  chain extends Chain = Chain,
+> = {
   /** Address of the deployed Account's Contract implementation. */
   address: Address;
   sign: (parameters: { hash: Hash }) => Promise<Hex>;
+  chain: NonNullable<chain>;
+  transport: transport;
+  contracts: PasskeyRequiredContracts;
 };
 
 export type PasskeyAccount = LocalAccount<"ssoPasskeyAccount"> & {
   sign: NonNullable<CustomSource["sign"]>;
 };
 
-export function toPasskeyAccount(
-  parameters: ToPasskeyAccountParameters,
+export function toPasskeyAccount<
+  transport extends Transport = Transport,
+  chain extends Chain = Chain,
+>(
+  parameters: ToPasskeyAccountParameters<transport, chain>,
 ): PasskeyAccount {
   const { address, sign } = parameters;
 
@@ -28,6 +38,11 @@ export function toPasskeyAccount(
         hash: hashMessage(message),
       });
     },
+    async signTypedData(typedData) {
+      return sign({
+        hash: hashTypedData(typedData),
+      });
+    },
     async signTransaction(transaction) {
       const signableTransaction = {
         ...transaction,
@@ -45,11 +60,6 @@ export function toPasskeyAccount(
         }),
       });
     },
-    async signTypedData(typedData) {
-      return sign({
-        hash: hashTypedData(typedData),
-      });
-    },
   });
 
   return {

package/src/client/passkey/actions/account.ts

@@ -9,6 +9,35 @@ import { noThrow } from "../../../utils/helpers.js";
 import { base64UrlToUint8Array, getPasskeySignatureFromPublicKeyBytes, getPublicKeyBytesFromPasskeySignature } from "../../../utils/passkey.js";
 import type { SessionConfig } from "../../../utils/session.js";
 
+export type DeployAccountPasskeyArgs = {
+  location: Address; // module address
+  credentialId: string; // Unique id of the passkey public key (base64)
+  credentialPublicKey: Uint8Array; // Public key of the previously registered
+  expectedOrigin?: string; // Expected origin of the passkey
+};
+export const encodePasskeyModuleData = async (
+  args: DeployAccountPasskeyArgs,
+): Promise<Hash> => {
+  let origin: string | undefined = args.expectedOrigin;
+  if (!origin) {
+    try {
+      origin = window.location.origin;
+    } catch {
+      throw new Error("Can't identify expectedOrigin, please provide it manually");
+    }
+  }
+  const passkeyPublicKey = getPublicKeyBytesFromPasskeySignature(args.credentialPublicKey);
+  const encodedPasskeyParameters = encodePasskeyModuleParameters({
+    credentialId: args.credentialId,
+    passkeyPublicKey,
+    expectedOrigin: origin,
+  });
+  return encodeModuleData({
+    address: args.location,
+    parameters: encodedPasskeyParameters,
+  });
+};
+
 /* TODO: try to get rid of most of the contract params like passkey, session */
 /* it should come from factory, not passed manually each time */
 export type DeployAccountArgs = {

package/src/client/passkey/actions/passkey.ts

@@ -1,7 +1,8 @@
 import { startAuthentication, startRegistration } from "@simplewebauthn/browser";
-import { generateAuthenticationOptions, type GenerateAuthenticationOptionsOpts, generateRegistrationOptions, type GenerateRegistrationOptionsOpts, type VerifiedRegistrationResponse, verifyAuthenticationResponse, verifyRegistrationResponse } from "@simplewebauthn/server";
-import { type AuthenticationResponseJSON, type PublicKeyCredentialCreationOptionsJSON, type PublicKeyCredentialRequestOptionsJSON, type RegistrationResponseJSON } from "@simplewebauthn/types";
-import { type Account, type Address, type Chain, type Client, encodeFunctionData, type Hash, type Hex, toBytes, toHex, type TransactionReceipt, type Transport } from "viem";
+import type { AuthenticationResponseJSON, GenerateAuthenticationOptionsOpts, GenerateRegistrationOptionsOpts, PublicKeyCredentialCreationOptionsJSON, PublicKeyCredentialRequestOptionsJSON, RegistrationResponseJSON, VerifiedRegistrationResponse } from "@simplewebauthn/server";
+import { generateAuthenticationOptions, generateRegistrationOptions, verifyAuthenticationResponse, verifyRegistrationResponse } from "@simplewebauthn/server";
+import type { Account, Address, Chain, Client, Hash, Hex, TransactionReceipt, Transport } from "viem";
+import { encodeFunctionData, toBytes, toHex } from "viem";
 import { waitForTransactionReceipt } from "viem/actions";
 import { getGeneralPaymasterInput, sendTransaction } from "viem/zksync";
 
@@ -92,7 +93,9 @@ export const registerNewPasskey = async (args: RegisterNewPasskeyArgs): Promise<
   if (!origin) throw new Error("Can't set origin automatically, please provide it manually in the arguments");
 
   const passkeyRegistrationOptions = "passkeyRegistrationOptions" in args ? args.passkeyRegistrationOptions : await generatePasskeyRegistrationOptions(args);
-  const registrationResponse: RegistrationResponseJSON = await startRegistration(passkeyRegistrationOptions);
+  const registrationResponse: RegistrationResponseJSON = await startRegistration({
+    optionsJSON: passkeyRegistrationOptions,
+  });
   const verification = await verifyRegistrationResponse({
     response: registrationResponse,
     expectedChallenge: passkeyRegistrationOptions.challenge,
@@ -104,8 +107,8 @@ export const registerNewPasskey = async (args: RegisterNewPasskeyArgs): Promise<
     passkeyRegistrationOptions,
     passkeyRegistrationResponse: registrationResponse,
     verificationResponse: verification,
-    credentialPublicKey: verification.registrationInfo.credentialPublicKey,
-    credentialId: verification.registrationInfo.credentialID,
+    credentialPublicKey: verification.registrationInfo.credential.publicKey,
+    credentialId: verification.registrationInfo.credential.id,
   };
 };
 
@@ -123,7 +126,8 @@ export const requestPasskeyAuthentication = async (args: RequestPasskeyAuthentic
   const passkeyAuthenticationOptions = await generatePasskeyAuthenticationOptions({
     challenge: toBytes(args.challenge),
   });
-  const authenticationResponse: AuthenticationResponseJSON = await startAuthentication(passkeyAuthenticationOptions);
+  const optionsJSON: PublicKeyCredentialRequestOptionsJSON = { ...passkeyAuthenticationOptions };
+  const authenticationResponse: AuthenticationResponseJSON = await startAuthentication({ optionsJSON: optionsJSON });
 
   let { rpID, origin } = identifyPasskeyParams();
   rpID = args.rpID || passkeyAuthenticationOptions.rpId || rpID;
@@ -135,10 +139,11 @@ export const requestPasskeyAuthentication = async (args: RequestPasskeyAuthentic
     expectedChallenge: passkeyAuthenticationOptions.challenge,
     expectedOrigin: origin,
     expectedRPID: rpID,
-    authenticator: {
-      credentialPublicKey: args.credentialPublicKey,
-      credentialID: authenticationResponse.id,
+    credential: {
+      id: authenticationResponse.id,
+      publicKey: args.credentialPublicKey,
       counter: 0, // TODO: figure out if this has to be dynamic
+
     },
   });
   if (!verification.verified || !verification.authenticationInfo) throw new Error("Passkey validation failed");

package/src/client/passkey/client.ts

@@ -26,6 +26,9 @@ export function createZksyncPasskeyClient<
 
   const account = toPasskeyAccount({
     address: parameters.address,
+    chain: parameters.chain,
+    contracts: parameters.contracts,
+    transport: parameters.transport,
     sign: async ({ hash }) => {
       const passkeySignature = await requestPasskeyAuthentication({
         challenge: hash,
@@ -35,7 +38,8 @@ export function createZksyncPasskeyClient<
       return passkeyHashSignatureResponseFormat(
         passkeySignature.passkeyAuthenticationResponse.id,
         passkeySignature.passkeyAuthenticationResponse.response,
-        parameters.contracts);
+        parameters.contracts,
+      );
     },
   });
   const client = createClient<transport, chain, Account, rpcSchema>({

package/src/client/passkey/decorators/wallet.ts

@@ -1,5 +1,5 @@
 import { type Account, bytesToHex, type Chain, type ExactPartial, formatTransaction, type RpcTransaction, type Transport, type WalletActions } from "viem";
-import { deployContract, getAddresses, getChainId, sendRawTransaction, signMessage, signTypedData, writeContract } from "viem/actions";
+import { deployContract, getAddresses, getCallsStatus, getCapabilities, getChainId, prepareAuthorization, sendCalls, sendRawTransaction, showCallsStatus, signAuthorization, signMessage, signTypedData, waitForCallsStatus, writeContract } from "viem/actions";
 import { signTransaction, type TransactionRequestEIP712, type ZksyncEip712Meta } from "viem/zksync";
 
 import { getTransactionWithPaymasterData } from "../../../paymaster/index.js";
@@ -65,9 +65,16 @@ export function zksyncSsoPasskeyWalletActions<
       return signTransaction(client, {
         ...args,
         unformattedTxWithPaymaster,
-      } as any);
+      } as any) as any;
     },
     signTypedData: (args) => signTypedData(client, args),
     writeContract: (args) => writeContract(client, args),
+    signAuthorization: (args) => signAuthorization(client, args),
+    getCallsStatus: (args) => getCallsStatus(client, args),
+    getCapabilities: (args) => getCapabilities(client, args),
+    prepareAuthorization: (args) => prepareAuthorization(client, args),
+    sendCalls: (args) => sendCalls(client, args),
+    showCallsStatus: (args) => showCallsStatus(client, args),
+    waitForCallsStatus: (args) => waitForCallsStatus(client, args),
   };
 }

package/src/client/recovery/actions/recovery.ts

@@ -2,7 +2,7 @@ import { type Account, type Address, type Chain, type Client, encodeFunctionData
 import { waitForTransactionReceipt } from "viem/actions";
 import { getGeneralPaymasterInput, sendTransaction } from "viem/zksync";
 
-import { GuardianRecoveryModuleAbi } from "../../../abi/GuardianRecoveryModule.js";
+import { GuardianRecoveryValidatorAbi } from "../../../abi/GuardianRecoveryValidator.js";
 import { noThrow } from "../../../utils/helpers.js";
 
 export type ProposeGuardianArgs = {
@@ -35,7 +35,7 @@ export const proposeGuardian = async <
   }
 
   const callData = encodeFunctionData({
-    abi: GuardianRecoveryModuleAbi,
+    abi: GuardianRecoveryValidatorAbi,
     functionName: "proposeGuardian",
     args: [keccak256(toHex(origin)), args.newGuardian],
   });
@@ -91,7 +91,7 @@ export const confirmGuardian = async <
     }
   }
   const callData = encodeFunctionData({
-    abi: GuardianRecoveryModuleAbi,
+    abi: GuardianRecoveryValidatorAbi,
     functionName: "addGuardian",
     args: [keccak256(toHex(origin)), args.accountToGuard],
   });
@@ -147,7 +147,7 @@ export const removeGuardian = async <
     }
   }
   const callData = encodeFunctionData({
-    abi: GuardianRecoveryModuleAbi,
+    abi: GuardianRecoveryValidatorAbi,
     functionName: "removeGuardian",
     args: [keccak256(toHex(origin)), args.guardian],
   });

package/src/client/recovery/decorators/wallet.ts

@@ -1,5 +1,5 @@
 import { type Account, bytesToHex, type Chain, formatTransaction, type Transport, type WalletActions } from "viem";
-import { deployContract, getAddresses, getChainId, sendRawTransaction, signMessage, signTypedData, writeContract } from "viem/actions";
+import { deployContract, getAddresses, getCallsStatus, getCapabilities, getChainId, prepareAuthorization, sendCalls, sendRawTransaction, showCallsStatus, signAuthorization, signMessage, signTypedData, waitForCallsStatus, writeContract } from "viem/actions";
 import { signTransaction, type ZksyncEip712Meta } from "viem/zksync";
 
 import { sendEip712Transaction } from "../actions/sendEip712Transaction.js";
@@ -45,8 +45,15 @@ export function zksyncSsoWalletActions<
     },
     signMessage: (args) => signMessage(client, args),
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    signTransaction: (args) => signTransaction(client, args as any),
+    signTransaction: (args) => signTransaction(client, args as any) as any,
     signTypedData: (args) => signTypedData(client, args),
     writeContract: (args) => writeContract(client, args),
+    signAuthorization: (args) => signAuthorization(client, args),
+    getCallsStatus: (args) => getCallsStatus(client, args),
+    getCapabilities: (args) => getCapabilities(client, args),
+    prepareAuthorization: (args) => prepareAuthorization(client, args),
+    sendCalls: (args) => sendCalls(client, args),
+    showCallsStatus: (args) => showCallsStatus(client, args),
+    waitForCallsStatus: (args) => waitForCallsStatus(client, args),
   };
 }

package/src/client/session/actions/session.ts

@@ -1,11 +1,12 @@
 import { type Account, type Address, type Chain, type Client, encodeFunctionData, type Hash, type Hex, type Prettify, type TransactionReceipt, type Transport } from "viem";
-import { waitForTransactionReceipt } from "viem/actions";
+import { readContract, waitForTransactionReceipt } from "viem/actions";
 import { getGeneralPaymasterInput, sendTransaction } from "viem/zksync";
 
 import { SessionKeyValidatorAbi } from "../../../abi/SessionKeyValidator.js";
 import { type CustomPaymasterHandler, getTransactionWithPaymasterData } from "../../../paymaster/index.js";
 import { noThrow } from "../../../utils/helpers.js";
-import type { SessionConfig } from "../../../utils/session.js";
+import type { SessionConfig, SessionState, SessionStateEventCallback } from "../../../utils/session.js";
+import { SessionEventType, SessionStatus } from "../../../utils/session.js";
 
 export type CreateSessionArgs = {
   sessionConfig: SessionConfig;
@@ -119,3 +120,85 @@ export const revokeSession = async <
     transactionReceipt,
   };
 };
+
+export type GetSessionStateArgs = {
+  account: Address;
+  sessionConfig: SessionConfig;
+  contracts: {
+    session: Address; // session module
+  };
+};
+export type GetSessionStateReturnType = {
+  sessionState: SessionState;
+};
+export const getSessionState = async <
+  transport extends Transport,
+  chain extends Chain,
+>(client: Client<transport, chain>, args: Prettify<GetSessionStateArgs>): Promise<Prettify<GetSessionStateReturnType>> => {
+  const sessionState = await readContract(client, {
+    address: args.contracts.session,
+    abi: SessionKeyValidatorAbi,
+    functionName: "sessionState",
+    args: [args.account, args.sessionConfig],
+  });
+
+  return {
+    sessionState: sessionState as SessionState,
+  };
+};
+
+export type CheckSessionStateArgs = {
+  sessionConfig: SessionConfig;
+  sessionState: SessionState;
+  onSessionStateChange: SessionStateEventCallback;
+  sessionNotifyTimeout?: NodeJS.Timeout;
+};
+export type CheckSessionStateReturnType = {
+  newTimeout?: NodeJS.Timeout;
+};
+
+/**
+ * Checks the current session state and sets up expiry notification.
+ * This function will trigger the callback with the session state.
+ */
+export const sessionStateNotify = (args: Prettify<CheckSessionStateArgs>): CheckSessionStateReturnType => {
+  // Generate a session ID for tracking timeouts
+  const { sessionState } = args;
+  const now = BigInt(Math.floor(Date.now() / 1000));
+
+  // Check session status
+  if (sessionState.status === SessionStatus.NotInitialized) { // Not initialized
+    args.onSessionStateChange({
+      type: SessionEventType.Inactive,
+      message: "Session is not initialized",
+    });
+  } else if (sessionState.status === SessionStatus.Closed) { // Closed/Revoked
+    args.onSessionStateChange({
+      type: SessionEventType.Revoked,
+      message: "Session has been revoked",
+    });
+  } else if (args.sessionConfig.expiresAt <= now) {
+    // Session has expired
+    args.onSessionStateChange({
+      type: SessionEventType.Expired,
+      message: "Session has expired",
+    });
+  } else {
+    // Session is active, set up expiry notification
+    const timeToExpiry = Number(args.sessionConfig.expiresAt - now) * 1000;
+    if (args.sessionNotifyTimeout) {
+      clearTimeout(args.sessionNotifyTimeout);
+    }
+    const newTimeout = setTimeout(() => {
+      args.onSessionStateChange({
+        type: SessionEventType.Expired,
+        message: "Session has expired",
+      });
+    }, timeToExpiry);
+    return {
+      newTimeout,
+    };
+  }
+
+  return {};
+};

package/src/client/session/client.ts

@@ -4,8 +4,9 @@ import { zksyncInMemoryNode } from "viem/chains";
 
 import type { CustomPaymasterHandler } from "../../paymaster/index.js";
 import { encodeSessionTx } from "../../utils/encoding.js";
-import type { SessionConfig } from "../../utils/session.js";
+import type { SessionConfig, SessionStateEventCallback } from "../../utils/session.js";
 import { toSessionAccount } from "./account.js";
+import { getSessionState, sessionStateNotify } from "./actions/session.js";
 import { publicActionsRewrite } from "./decorators/publicActionsRewrite.js";
 import { type ZksyncSsoWalletActions, zksyncSsoWalletActions } from "./decorators/wallet.js";
 
@@ -90,10 +91,31 @@ export function createZksyncSessionClient<
       sessionConfig: parameters.sessionConfig,
       contracts: parameters.contracts,
       paymasterHandler: parameters.paymasterHandler,
+      onSessionStateChange: parameters.onSessionStateChange,
+      _sessionNotifyTimeout: undefined as NodeJS.Timeout | undefined,
     }))
     .extend(publicActions)
     .extend(publicActionsRewrite)
     .extend(zksyncSsoWalletActions);
+
+  // Check session state on initialization if callback is provided
+  if (client.onSessionStateChange) {
+    getSessionState(client, {
+      account: client.account.address,
+      sessionConfig: client.sessionConfig,
+      contracts: client.contracts,
+    }).then(({ sessionState }) => {
+      sessionStateNotify({
+        sessionConfig: client.sessionConfig,
+        sessionState,
+        onSessionStateChange: client.onSessionStateChange!,
+        sessionNotifyTimeout: client._sessionNotifyTimeout,
+      });
+    }).catch((error) => {
+      console.error("Failed to get session state on initialization:", error);
+    });
+  }
+
   return client;
 }
 
@@ -105,6 +127,9 @@ type ZksyncSsoSessionData = {
   sessionConfig: SessionConfig;
   contracts: SessionRequiredContracts;
   paymasterHandler?: CustomPaymasterHandler;
+  onSessionStateChange?: SessionStateEventCallback;
+  skipPreTransactionStateValidation?: boolean;
+  _sessionNotifyTimeout?: NodeJS.Timeout;
 };
 
 export type ClientWithZksyncSsoSessionData<
@@ -143,4 +168,6 @@ export interface ZksyncSsoSessionClientConfig<
   key?: string;
   name?: string;
   paymasterHandler?: CustomPaymasterHandler;
+  onSessionStateChange?: SessionStateEventCallback;
+  skipPreTransactionStateValidation?: boolean; // Useful if you want to send session transactions really fast
 }