Merge branch 'marvin/public_keys' into marvin/public_account_id

.github/workflows/ci.yml

@@ -83,7 +83,7 @@ jobs:
           RISC0_SKIP_BUILD: "1"
         run: cargo clippy -p "*programs" -- -D warnings
 
-  unit-tests:
+  tests:
     runs-on: ubuntu-latest
     timeout-minutes: 60
     steps:
@@ -99,12 +99,13 @@ jobs:
         run: rustup install
 
       - name: Install nextest
-        run: cargo install cargo-nextest
+        run: cargo install --locked cargo-nextest
 
-      - name: Run unit tests
+      - name: Run tests
         env:
           RISC0_DEV_MODE: "1"
-        run: cargo nextest run --no-fail-fast
+          RUST_LOG: "info"
+        run: cargo nextest run --no-fail-fast -- --skip tps_test
 
   valid-proof-test:
     runs-on: ubuntu-latest
@@ -123,31 +124,8 @@ jobs:
 
       - name: Test valid proof
         env:
-          NSSA_WALLET_HOME_DIR: ./integration_tests/configs/debug/wallet
           RUST_LOG: "info"
-        run: cargo run --bin integration_tests -- ./integration_tests/configs/debug/ test_success_private_transfer_to_another_owned_account
-
-  integration-tests:
-    runs-on: ubuntu-latest
-    timeout-minutes: 120
-    steps:
-      - uses: actions/checkout@v5
-        with:
-          ref: ${{ github.head_ref }}
-
-      - uses: ./.github/actions/install-system-deps
-
-      - uses: ./.github/actions/install-risc0
-
-      - name: Install active toolchain
-        run: rustup install
-
-      - name: Run integration tests
-        env:
-          NSSA_WALLET_HOME_DIR: ./integration_tests/configs/debug/wallet
-          RUST_LOG: "info"
-          RISC0_DEV_MODE: "1"
-        run: cargo run --bin integration_tests -- ./integration_tests/configs/debug/ all
+        run: cargo test -p integration_tests -- --exact private::private_transfer_to_owned_account
 
   artifacts:
     runs-on: ubuntu-latest

Cargo.toml

@@ -16,10 +16,10 @@ members = [
   "program_methods/guest",
   "test_program_methods",
   "test_program_methods/guest",
-  "integration_tests/proc_macro_test_attribute",
   "examples/program_deployment",
   "examples/program_deployment/methods",
   "examples/program_deployment/methods/guest",
+  "bedrock_client",
 ]
 
 [workspace.dependencies]
@@ -34,6 +34,7 @@ sequencer_rpc = { path = "sequencer_rpc" }
 sequencer_runner = { path = "sequencer_runner" }
 wallet = { path = "wallet" }
 test_program_methods = { path = "test_program_methods" }
+bedrock_client = { path = "bedrock_client" }
 
 tokio = { version = "1.28.2", features = [
   "net",
@@ -76,6 +77,11 @@ chrono = "0.4.41"
 borsh = "1.5.7"
 base58 = "0.2.0"
 itertools = "0.14.0"
+url = "2.5.4"
+
+logos-blockchain-common-http-client = { git = "https://github.com/logos-blockchain/logos-blockchain.git" }
+logos-blockchain-key-management-system-service = { git = "https://github.com/logos-blockchain/logos-blockchain.git" }
+logos-blockchain-core = { git = "https://github.com/logos-blockchain/logos-blockchain.git" }
 
 rocksdb = { version = "0.24.0", default-features = false, features = [
   "snappy",
@@ -94,4 +100,4 @@ actix-web = { version = "=4.1.0", default-features = false, features = [
   "macros",
 ] }
 clap = { version = "4.5.42", features = ["derive", "env"] }
-reqwest = { version = "0.11.16", features = ["json"] }
+reqwest = { version = "0.12", features = ["json", "rustls-tls", "stream"] }

README.md

@@ -155,6 +155,8 @@ cargo install --path wallet --force
 
 Run `wallet help` to check everything went well.
 
+Some completion scripts exists, see the [completions](./completions/README.md) folder.
+
 ## Tutorial
 
 This tutorial walks you through creating accounts and executing NSSA programs in both public and private contexts.

bedrock_client/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "bedrock_client"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+reqwest.workspace = true
+anyhow.workspace = true
+logos-blockchain-common-http-client.workspace = true
+logos-blockchain-core.workspace = true

bedrock_client/src/lib.rs

@@ -0,0 +1,32 @@
+use anyhow::Result;
+pub use logos_blockchain_common_http_client::{BasicAuthCredentials, CommonHttpClient, Error};
+use logos_blockchain_core::mantle::SignedMantleTx;
+use reqwest::{Client, Url};
+
+// Simple wrapper
+// maybe extend in the future for our purposes
+pub struct BedrockClient {
+    http_client: CommonHttpClient,
+    node_url: Url,
+}
+
+impl BedrockClient {
+    pub fn new(auth: Option<BasicAuthCredentials>, node_url: Url) -> Result<Self> {
+        let client = Client::builder()
+                //Add more fields if needed
+                .timeout(std::time::Duration::from_secs(60))
+                .build()?;
+
+        let http_client = CommonHttpClient::new_with_client(client, auth);
+        Ok(Self {
+            http_client,
+            node_url,
+        })
+    }
+
+    pub async fn post_transaction(&self, tx: SignedMantleTx) -> Result<(), Error> {
+        self.http_client
+            .post_transaction(self.node_url.clone(), tx)
+            .await
+    }
+}

ci_scripts/deploy.sh

@@ -52,7 +52,7 @@ if [ -d ".git" ]; then
   git reset --hard origin/main
 else
   echo "Cloning repository..."
-  git clone https://github.com/vacp2p/nescience-testnet.git .
+  git clone https://github.com/logos-blockchain/lssa.git .
   git checkout main
 fi
 

common/src/block.rs

@@ -23,7 +23,7 @@ pub type BlockHash = [u8; 32];
 pub type BlockId = u64;
 pub type TimeStamp = u64;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)]
 pub struct BlockHeader {
     pub block_id: BlockId,
     pub prev_block_hash: BlockHash,
@@ -32,18 +32,26 @@ pub struct BlockHeader {
     pub signature: nssa::Signature,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)]
 pub struct BlockBody {
     pub transactions: Vec<EncodedTransaction>,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)]
+pub enum BedrockStatus {
+    Pending,
+    Safe,
+    Finalized,
+}
+
+#[derive(Debug, BorshSerialize, BorshDeserialize)]
 pub struct Block {
     pub header: BlockHeader,
     pub body: BlockBody,
+    pub bedrock_status: BedrockStatus,
 }
 
-#[derive(Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
 pub struct HashableBlockData {
     pub block_id: BlockId,
     pub prev_block_hash: BlockHash,
@@ -52,7 +60,7 @@ pub struct HashableBlockData {
 }
 
 impl HashableBlockData {
-    pub fn into_block(self, signing_key: &nssa::PrivateKey) -> Block {
+    pub fn into_pending_block(self, signing_key: &nssa::PrivateKey) -> Block {
         let data_bytes = borsh::to_vec(&self).unwrap();
         let signature = nssa::Signature::new(signing_key, &data_bytes);
         let hash = OwnHasher::hash(&data_bytes);
@@ -67,6 +75,7 @@ impl HashableBlockData {
             body: BlockBody {
                 transactions: self.transactions,
             },
+            bedrock_status: BedrockStatus::Pending,
         }
     }
 }

common/src/sequencer_client.rs

@@ -44,8 +44,10 @@ impl SequencerClient {
     ) -> Result<Self> {
         Ok(Self {
             client: Client::builder()
-                //Add more fiedls if needed
+                // Add more fields if needed
                 .timeout(std::time::Duration::from_secs(60))
+                // Should be kept in sync with server keep-alive settings
+                .pool_idle_timeout(std::time::Duration::from_secs(5))
                 .build()?,
             sequencer_addr,
             basic_auth,
@@ -60,6 +62,10 @@ impl SequencerClient {
         let request =
             rpc_primitives::message::Request::from_payload_version_2_0(method.to_string(), payload);
 
+        log::debug!(
+            "Calling method {method} with payload {request:?} to sequencer at {}",
+            self.sequencer_addr
+        );
         let mut call_builder = self.client.post(&self.sequencer_addr);
 
         if let Some((username, password)) = &self.basic_auth {

common/src/test_utils.rs

@@ -30,7 +30,7 @@ pub fn produce_dummy_block(
         transactions,
     };
 
-    block_data.into_block(&sequencer_sign_key_for_testing())
+    block_data.into_pending_block(&sequencer_sign_key_for_testing())
 }
 
 pub fn produce_dummy_empty_transaction() -> EncodedTransaction {

completions/README.md

@@ -0,0 +1,135 @@
+# Wallet CLI Completion
+
+Completion scripts for the LSSA `wallet` command. 
+
+## ZSH
+
+Works with both vanilla zsh and oh-my-zsh.
+
+### Features
+
+- Full completion for all wallet subcommands
+- Contextual option completion for each command
+- Dynamic account ID completion via `wallet account list`
+- Descriptions for all commands and options
+
+Note that only accounts created by the user auto-complete.
+Preconfigured accounts and accounts only with `/` (no number) are not completed.
+
+e.g.:
+
+```
+▶ wallet account list 
+Preconfigured Public/Gj1mJy5W7J5pfmLRujmQaLfLMWidNxQ6uwnhb666ZwHw,
+Preconfigured Public/BLgCRDXYdQPMMWVHYRFGQZbgeHx9frkipa8GtpG2Syqy,
+Preconfigured Private/3oCG8gqdKLMegw4rRfyaMQvuPHpcASt7xwttsmnZLSkw,
+Preconfigured Private/AKTcXgJ1xoynta1Ec7y6Jso1z1JQtHqd7aPQ1h9er6xX,
+/ Public/8DstRgMQrB2N9a7ymv98RDDbt8nctrP9ZzaNRSpKDZSu,
+/0 Public/2gJJjtG9UivBGEhA1Jz6waZQx1cwfYupC5yvKEweHaeH,
+/ Private/Bcv15B36bs1VqvQAdY6ZGFM1KioByNQQsB92KTNAx6u2
+```
+
+Only `Public/2gJJjtG9UivBGEhA1Jz6waZQx1cwfYupC5yvKEweHaeH` is used for completion.
+
+### Supported Commands
+
+| Command                | Description                                                 |
+|------------------------|-------------------------------------------------------------|
+| `wallet auth-transfer` | Authenticated transfer (init, send)                         |
+| `wallet chain-info`    | Chain info queries (current-block-id, block, transaction)   |
+| `wallet account`       | Account management (get, list, new, sync-private)           |
+| `wallet pinata`        | Piñata faucet (claim)                                       |
+| `wallet token`         | Token operations (new, send)                                |
+| `wallet amm`           | AMM operations (new, swap, add-liquidity, remove-liquidity) |
+| `wallet check-health`  | Health check                                                |
+
+### Installation
+
+#### Vanilla Zsh
+
+1. Create a completions directory:
+
+   ```sh
+   mkdir -p ~/.zsh/completions
+   ```
+
+2. Copy the completion file:
+
+   ```sh
+   cp ./zsh/_wallet ~/.zsh/completions/
+   ```
+
+3. Add to your `~/.zshrc` (before any `compinit` call, or add these lines if you don't have one):
+
+   ```sh
+   fpath=(~/.zsh/completions $fpath)
+   autoload -Uz compinit && compinit
+   ```
+
+4. Reload your shell:
+
+   ```sh
+   exec zsh
+   ```
+
+#### Oh-My-Zsh
+
+1. Create the plugin directory and copy the file:
+
+   ```sh
+   mkdir -p ~/.oh-my-zsh/custom/plugins/wallet
+   cp _wallet ~/.oh-my-zsh/custom/plugins/wallet/
+   ```
+
+2. Add `wallet` to your plugins array in `~/.zshrc`:
+
+   ```sh
+   plugins=(... wallet)
+   ```
+
+3. Reload your shell:
+
+   ```sh
+   exec zsh
+   ```
+
+### Requirements
+
+The completion script calls `wallet account list` to dynamically fetch account IDs. Ensure the `wallet` command is in your `$PATH`.
+
+### Usage
+
+```sh
+# Main commands
+wallet <TAB>
+
+# Account subcommands
+wallet account <TAB>
+
+# Options for auth-transfer send
+wallet auth-transfer send --<TAB>
+
+# Account types when creating
+wallet account new <TAB>
+# Shows: public  private
+
+# Account IDs (fetched dynamically)
+wallet account get --account-id <TAB>
+# Shows: Public/...  Private/...
+```
+
+## Troubleshooting
+
+### Completions not appearing
+
+1. Check that `compinit` is called in your `.zshrc`
+2. Rebuild the completion cache:
+
+   ```sh
+   rm -f ~/.zcompdump*
+   exec zsh
+   ```
+
+### Account IDs not completing
+
+Ensure `wallet account list` works from your command line.

completions/zsh/_wallet

@@ -0,0 +1,435 @@
+#compdef wallet
+
+# Zsh completion script for the wallet CLI
+# See instructions in ../README.md
+
+_wallet() {
+    local -a commands
+    local -a subcommands
+    local curcontext="$curcontext" state line
+    typeset -A opt_args
+
+    _arguments -C \
+        '(-c --continuous-run)'{-c,--continuous-run}'[Continuous run flag]' \
+        '--auth[Basic authentication in the format user or user\:password]:auth:' \
+        '1: :->command' \
+        '*:: :->args'
+
+    case $state in
+        command)
+            commands=(
+                'auth-transfer:Authenticated transfer subcommand'
+                'chain-info:Generic chain info subcommand'
+                'account:Account view and sync subcommand'
+                'pinata:Pinata program interaction subcommand'
+                'token:Token program interaction subcommand'
+                'amm:AMM program interaction subcommand'
+                'check-health:Check the wallet can connect to the node and builtin local programs match the remote versions'
+                'config:Command to setup config, get and set config fields'
+                'restore-keys:Restoring keys from given password at given depth'
+                'deploy-program:Deploy a program'
+                'help:Print help message or the help of the given subcommand(s)'
+            )
+            _describe -t commands 'wallet commands' commands
+            ;;
+        args)
+            case $line[1] in
+                auth-transfer)
+                    _wallet_auth_transfer
+                    ;;
+                chain-info)
+                    _wallet_chain_info
+                    ;;
+                account)
+                    _wallet_account
+                    ;;
+                pinata)
+                    _wallet_pinata
+                    ;;
+                token)
+                    _wallet_token
+                    ;;
+                amm)
+                    _wallet_amm
+                    ;;
+                config)
+                    _wallet_config
+                    ;;
+                restore-keys)
+                    _wallet_restore_keys
+                    ;;
+                deploy-program)
+                    _wallet_deploy_program
+                    ;;
+                help)
+                    _wallet_help
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# auth-transfer subcommand
+_wallet_auth_transfer() {
+    local -a subcommands
+    
+    _arguments -C \
+        '1: :->subcommand' \
+        '*:: :->args'
+
+    case $state in
+        subcommand)
+            subcommands=(
+                'init:Initialize account under authenticated transfer program'
+                'send:Send native tokens from one account to another with variable privacy'
+                'help:Print this message or the help of the given subcommand(s)'
+            )
+            _describe -t subcommands 'auth-transfer subcommands' subcommands
+            ;;
+        args)
+            case $line[1] in
+                init)
+                    _arguments \
+                        '--account-id[Account ID to initialize]:account_id:_wallet_account_ids'
+                    ;;
+                send)
+                    _arguments \
+                        '--from[Source account ID]:from_account:_wallet_account_ids' \
+                        '--to[Destination account ID (for owned accounts)]:to_account:_wallet_account_ids' \
+                        '--to-npk[Destination nullifier public key (for foreign private accounts)]:npk:' \
+                        '--to-ipk[Destination viewing public key (for foreign private accounts)]:ipk:' \
+                        '--amount[Amount of native tokens to send]:amount:'
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# chain-info subcommand
+_wallet_chain_info() {
+    local -a subcommands
+    
+    _arguments -C \
+        '1: :->subcommand' \
+        '*:: :->args'
+
+    case $state in
+        subcommand)
+            subcommands=(
+                'current-block-id:Get current block id from sequencer'
+                'block:Get block at id from sequencer'
+                'transaction:Get transaction at hash from sequencer'
+                'help:Print this message or the help of the given subcommand(s)'
+            )
+            _describe -t subcommands 'chain-info subcommands' subcommands
+            ;;
+        args)
+            case $line[1] in
+                block)
+                    _arguments \
+                        '--id[Block ID to retrieve]:block_id:'
+                    ;;
+                transaction)
+                    _arguments \
+                        '--hash[Transaction hash to retrieve]:tx_hash:'
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# account subcommand
+_wallet_account() {
+    local -a subcommands
+    
+    _arguments -C \
+        '1: :->subcommand' \
+        '*:: :->args'
+
+    case $state in
+        subcommand)
+            subcommands=(
+                'get:Get account data'
+                'list:List all accounts'
+                'ls:List all accounts (alias for list)'
+                'new:Produce new public or private account'
+                'sync-private:Sync private accounts'
+                'help:Print this message or the help of the given subcommand(s)'
+            )
+            _describe -t subcommands 'account subcommands' subcommands
+            ;;
+        args)
+            case $line[1] in
+                get)
+                    _arguments \
+                        '(-r --raw)'{-r,--raw}'[Get raw account data]' \
+                        '(-k --keys)'{-k,--keys}'[Display keys (pk for public accounts, npk/ipk for private accounts)]' \
+                        '(-a --account-id)'{-a,--account-id}'[Account ID to query]:account_id:_wallet_account_ids'
+                    ;;
+                list|ls)
+                    _arguments \
+                        '(-l --long)'{-l,--long}'[Display detailed account information]'
+                    ;;
+                new)
+                    _arguments -C \
+                        '1: :->account_type' \
+                        '*:: :->new_args'
+                    case $state in
+                        account_type)
+                            compadd public private
+                            ;;
+                        new_args)
+                            _arguments \
+                                '--cci[Chain index of a parent node]:chain_index:'
+                            ;;
+                    esac
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# pinata subcommand
+_wallet_pinata() {
+    local -a subcommands
+    
+    _arguments -C \
+        '1: :->subcommand' \
+        '*:: :->args'
+
+    case $state in
+        subcommand)
+            subcommands=(
+                'claim:Claim tokens from the Piñata faucet'
+                'help:Print this message or the help of the given subcommand(s)'
+            )
+            _describe -t subcommands 'pinata subcommands' subcommands
+            ;;
+        args)
+            case $line[1] in
+                claim)
+                    _arguments \
+                        '--to[Destination account ID to receive claimed tokens]:to_account:_wallet_account_ids'
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# token subcommand
+_wallet_token() {
+    local -a subcommands
+
+    _arguments -C \
+        '1: :->subcommand' \
+        '*:: :->args'
+
+    case $state in
+        subcommand)
+            subcommands=(
+                'new:Produce a new token'
+                'send:Send tokens from one account to another with variable privacy'
+                'burn:Burn tokens on holder, modify definition'
+                'mint:Mint tokens on holder, modify definition'
+                'help:Print this message or the help of the given subcommand(s)'
+            )
+            _describe -t subcommands 'token subcommands' subcommands
+            ;;
+        args)
+            case $line[1] in
+                new)
+                    _arguments \
+                        '--name[Token name]:name:' \
+                        '--total-supply[Total supply of tokens to mint]:total_supply:' \
+                        '--definition-account-id[Account ID for token definition]:definition_account:_wallet_account_ids' \
+                        '--supply-account-id[Account ID to receive initial supply]:supply_account:_wallet_account_ids'
+                    ;;
+                send)
+                    _arguments \
+                        '--from[Source holding account ID]:from_account:_wallet_account_ids' \
+                        '--to[Destination holding account ID (for owned accounts)]:to_account:_wallet_account_ids' \
+                        '--to-npk[Destination nullifier public key (for foreign private accounts)]:npk:' \
+                        '--to-ipk[Destination viewing public key (for foreign private accounts)]:ipk:' \
+                        '--amount[Amount of tokens to send]:amount:'
+                    ;;
+                burn)
+                    _arguments \
+                        '--definition[Definition account ID]:definition_account:_wallet_account_ids' \
+                        '--holder[Holder account ID]:holder_account:_wallet_account_ids' \
+                        '--amount[Amount of tokens to burn]:amount:'
+                    ;;
+                mint)
+                    _arguments \
+                        '--definition[Definition account ID]:definition_account:_wallet_account_ids' \
+                        '--holder[Holder account ID (for owned accounts)]:holder_account:_wallet_account_ids' \
+                        '--holder-npk[Holder nullifier public key (for foreign private accounts)]:npk:' \
+                        '--holder-ipk[Holder viewing public key (for foreign private accounts)]:ipk:' \
+                        '--amount[Amount of tokens to mint]:amount:'
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# amm subcommand
+_wallet_amm() {
+    local -a subcommands
+    
+    _arguments -C \
+        '1: :->subcommand' \
+        '*:: :->args'
+
+    case $state in
+        subcommand)
+            subcommands=(
+                'new:Create a new liquidity pool'
+                'swap:Swap tokens using the AMM'
+                'add-liquidity:Add liquidity to an existing pool'
+                'remove-liquidity:Remove liquidity from a pool'
+                'help:Print this message or the help of the given subcommand(s)'
+            )
+            _describe -t subcommands 'amm subcommands' subcommands
+            ;;
+        args)
+            case $line[1] in
+                new)
+                    _arguments \
+                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
+                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
+                        '--user-holding-lp[User LP token holding account ID]:holding_lp:_wallet_account_ids' \
+                        '--balance-a[Amount of token A to deposit]:balance_a:' \
+                        '--balance-b[Amount of token B to deposit]:balance_b:'
+                    ;;
+                swap)
+                    _arguments \
+                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
+                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
+                        '--amount-in[Amount of tokens to swap]:amount_in:' \
+                        '--min-amount-out[Minimum tokens expected in return]:min_amount_out:' \
+                        '--token-definition[Definition ID of the token being provided]:token_def:'
+                    ;;
+                add-liquidity)
+                    _arguments \
+                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
+                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
+                        '--user-holding-lp[User LP token holding account ID]:holding_lp:_wallet_account_ids' \
+                        '--max-amount-a[Maximum amount of token A to deposit]:max_amount_a:' \
+                        '--max-amount-b[Maximum amount of token B to deposit]:max_amount_b:' \
+                        '--min-amount-lp[Minimum LP tokens to receive]:min_amount_lp:'
+                    ;;
+                remove-liquidity)
+                    _arguments \
+                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
+                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
+                        '--user-holding-lp[User LP token holding account ID]:holding_lp:_wallet_account_ids' \
+                        '--balance-lp[Amount of LP tokens to burn]:balance_lp:' \
+                        '--min-amount-a[Minimum token A to receive]:min_amount_a:' \
+                        '--min-amount-b[Minimum token B to receive]:min_amount_b:'
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# config subcommand
+_wallet_config() {
+    local -a subcommands
+    local -a config_keys
+
+    config_keys=(
+        'all'
+        'override_rust_log'
+        'sequencer_addr'
+        'seq_poll_timeout_millis'
+        'seq_tx_poll_max_blocks'
+        'seq_poll_max_retries'
+        'seq_block_poll_max_amount'
+        'initial_accounts'
+        'basic_auth'
+    )
+
+    _arguments -C \
+        '1: :->subcommand' \
+        '*:: :->args'
+
+    case $state in
+        subcommand)
+            subcommands=(
+                'get:Getter of config fields'
+                'set:Setter of config fields'
+                'description:Prints description of corresponding field'
+                'help:Print this message or the help of the given subcommand(s)'
+            )
+            _describe -t subcommands 'config subcommands' subcommands
+            ;;
+        args)
+            case $line[1] in
+                get|description)
+                    compadd -a config_keys
+                    ;;
+                set)
+                    _arguments \
+                        '1:key:compadd -a config_keys' \
+                        '2:value:'
+                    ;;
+            esac
+            ;;
+    esac
+}
+
+# restore-keys subcommand
+_wallet_restore_keys() {
+    _arguments \
+        '(-d --depth)'{-d,--depth}'[How deep in tree accounts may be]:depth:'
+}
+
+# deploy-program subcommand
+_wallet_deploy_program() {
+    _arguments \
+        '1:binary filepath:_files'
+}
+
+# help subcommand
+_wallet_help() {
+    local -a commands
+    commands=(
+        'auth-transfer:Authenticated transfer subcommand'
+        'chain-info:Generic chain info subcommand'
+        'account:Account view and sync subcommand'
+        'pinata:Pinata program interaction subcommand'
+        'token:Token program interaction subcommand'
+        'amm:AMM program interaction subcommand'
+        'check-health:Check the wallet can connect to the node'
+        'config:Command to setup config, get and set config fields'
+        'restore-keys:Restoring keys from given password at given depth'
+        'deploy-program:Deploy a program'
+    )
+    _describe -t commands 'wallet commands' commands
+}
+
+# Helper function to complete account IDs
+# Uses `wallet account list` to get available accounts
+# Only includes accounts with /N prefix (where N is a number)
+_wallet_account_ids() {
+    local -a accounts
+    local line
+    
+    # Try to get accounts from wallet account list command
+    # Filter to lines starting with /N (numbered accounts) and extract the account ID
+    if command -v wallet &>/dev/null; then
+        while IFS= read -r line; do
+            # Remove trailing comma if present and add to array
+            [[ -n "$line" ]] && accounts+=("${line%,}")
+        done < <(wallet account list 2>/dev/null | grep '^/[0-9]' | awk '{print $2}')
+    fi
+    
+    # Provide type prefixes as fallback if command fails or returns nothing
+    if (( ${#accounts} == 0 )); then
+        compadd -S '' -- 'Public/' 'Private/'
+        return
+    fi
+    
+    _multi_parts / accounts
+}
+
+_wallet "$@"

examples/program_deployment/src/bin/run_hello_world.rs

@@ -3,7 +3,7 @@ use nssa::{
     program::Program,
     public_transaction::{Message, WitnessSet},
 };
-use wallet::{WalletCore, helperfunctions::fetch_config};
+use wallet::WalletCore;
 
 // Before running this example, compile the `hello_world.rs` guest program with:
 //
@@ -24,11 +24,8 @@ use wallet::{WalletCore, helperfunctions::fetch_config};
 
 #[tokio::main]
 async fn main() {
-    // Load wallet config and storage
-    let wallet_config = fetch_config().await.unwrap();
-    let wallet_core = WalletCore::start_from_config_update_chain(wallet_config)
-        .await
-        .unwrap();
+    // Initialize wallet
+    let wallet_core = WalletCore::from_env().unwrap();
 
     // Parse arguments
     // First argument is the path to the program binary

examples/program_deployment/src/bin/run_hello_world_private.rs

@@ -1,5 +1,5 @@
 use nssa::{AccountId, program::Program};
-use wallet::{PrivacyPreservingAccount, WalletCore, helperfunctions::fetch_config};
+use wallet::{PrivacyPreservingAccount, WalletCore};
 
 // Before running this example, compile the `hello_world.rs` guest program with:
 //
@@ -22,11 +22,8 @@ use wallet::{PrivacyPreservingAccount, WalletCore, helperfunctions::fetch_config
 
 #[tokio::main]
 async fn main() {
-    // Load wallet config and storage
-    let wallet_config = fetch_config().await.unwrap();
-    let wallet_core = WalletCore::start_from_config_update_chain(wallet_config)
-        .await
-        .unwrap();
+    // Initialize wallet
+    let wallet_core = WalletCore::from_env().unwrap();
 
     // Parse arguments
     // First argument is the path to the program binary
@@ -53,7 +50,7 @@ async fn main() {
     wallet_core
         .send_privacy_preserving_tx(
             accounts,
-            &Program::serialize_instruction(greeting).unwrap(),
+            Program::serialize_instruction(greeting).unwrap(),
             &program.into(),
         )
         .await

examples/program_deployment/src/bin/run_hello_world_through_tail_call.rs

@@ -3,7 +3,7 @@ use nssa::{
     program::Program,
     public_transaction::{Message, WitnessSet},
 };
-use wallet::{WalletCore, helperfunctions::fetch_config};
+use wallet::WalletCore;
 
 // Before running this example, compile the `simple_tail_call.rs` guest program with:
 //
@@ -24,11 +24,8 @@ use wallet::{WalletCore, helperfunctions::fetch_config};
 
 #[tokio::main]
 async fn main() {
-    // Load wallet config and storage
-    let wallet_config = fetch_config().await.unwrap();
-    let wallet_core = WalletCore::start_from_config_update_chain(wallet_config)
-        .await
-        .unwrap();
+    // Initialize wallet
+    let wallet_core = WalletCore::from_env().unwrap();
 
     // Parse arguments
     // First argument is the path to the program binary

examples/program_deployment/src/bin/run_hello_world_through_tail_call_private.rs

@@ -4,7 +4,7 @@ use nssa::{
     AccountId, ProgramId, privacy_preserving_transaction::circuit::ProgramWithDependencies,
     program::Program,
 };
-use wallet::{PrivacyPreservingAccount, WalletCore, helperfunctions::fetch_config};
+use wallet::{PrivacyPreservingAccount, WalletCore};
 
 // Before running this example, compile the `simple_tail_call.rs` guest program with:
 //
@@ -25,11 +25,8 @@ use wallet::{PrivacyPreservingAccount, WalletCore, helperfunctions::fetch_config
 
 #[tokio::main]
 async fn main() {
-    // Load wallet config and storage
-    let wallet_config = fetch_config().await.unwrap();
-    let wallet_core = WalletCore::start_from_config_update_chain(wallet_config)
-        .await
-        .unwrap();
+    // Initialize wallet
+    let wallet_core = WalletCore::from_env().unwrap();
 
     // Parse arguments
     // First argument is the path to the simple_tail_call program binary
@@ -61,7 +58,7 @@ async fn main() {
     wallet_core
         .send_privacy_preserving_tx(
             accounts,
-            &Program::serialize_instruction(instruction).unwrap(),
+            Program::serialize_instruction(instruction).unwrap(),
             &program_with_dependencies,
         )
         .await

examples/program_deployment/src/bin/run_hello_world_with_authorization.rs

@@ -3,7 +3,7 @@ use nssa::{
     program::Program,
     public_transaction::{Message, WitnessSet},
 };
-use wallet::{WalletCore, helperfunctions::fetch_config};
+use wallet::WalletCore;
 
 // Before running this example, compile the `hello_world_with_authorization.rs` guest program with:
 //
@@ -26,11 +26,8 @@ use wallet::{WalletCore, helperfunctions::fetch_config};
 
 #[tokio::main]
 async fn main() {
-    // Load wallet config and storage
-    let wallet_config = fetch_config().await.unwrap();
-    let wallet_core = WalletCore::start_from_config_update_chain(wallet_config)
-        .await
-        .unwrap();
+    // Initialize wallet
+    let wallet_core = WalletCore::from_env().unwrap();
 
     // Parse arguments
     // First argument is the path to the program binary
@@ -50,7 +47,7 @@ async fn main() {
 
     // Load signing keys to provide authorization
     let signing_key = wallet_core
-        .storage
+        .storage()
         .user_data
         .get_pub_account_signing_key(&account_id)
         .expect("Input account should be a self owned public account");

examples/program_deployment/src/bin/run_hello_world_with_authorization_through_tail_call_with_pda.rs

@@ -4,7 +4,7 @@ use nssa::{
     public_transaction::{Message, WitnessSet},
 };
 use nssa_core::program::PdaSeed;
-use wallet::{WalletCore, helperfunctions::fetch_config};
+use wallet::WalletCore;
 
 // Before running this example, compile the `simple_tail_call.rs` guest program with:
 //
@@ -27,11 +27,8 @@ const PDA_SEED: PdaSeed = PdaSeed::new([37; 32]);
 
 #[tokio::main]
 async fn main() {
-    // Load wallet config and storage
-    let wallet_config = fetch_config().await.unwrap();
-    let wallet_core = WalletCore::start_from_config_update_chain(wallet_config)
-        .await
-        .unwrap();
+    // Initialize wallet
+    let wallet_core = WalletCore::from_env().unwrap();
 
     // Parse arguments
     // First argument is the path to the program binary

examples/program_deployment/src/bin/run_hello_world_with_move_function.rs

@@ -1,6 +1,6 @@
 use clap::{Parser, Subcommand};
 use nssa::{PublicTransaction, program::Program, public_transaction};
-use wallet::{PrivacyPreservingAccount, WalletCore, helperfunctions::fetch_config};
+use wallet::{PrivacyPreservingAccount, WalletCore};
 
 // Before running this example, compile the `hello_world_with_move_function.rs` guest program with:
 //
@@ -62,11 +62,8 @@ async fn main() {
     let bytecode: Vec<u8> = std::fs::read(cli.program_path).unwrap();
     let program = Program::new(bytecode).unwrap();
 
-    // Load wallet config and storage
-    let wallet_config = fetch_config().await.unwrap();
-    let wallet_core = WalletCore::start_from_config_update_chain(wallet_config)
-        .await
-        .unwrap();
+    // Initialize wallet
+    let wallet_core = WalletCore::from_env().unwrap();
 
     match cli.command {
         Command::WritePublic {
@@ -104,7 +101,7 @@ async fn main() {
             wallet_core
                 .send_privacy_preserving_tx(
                     accounts,
-                    &Program::serialize_instruction(instruction).unwrap(),
+                    Program::serialize_instruction(instruction).unwrap(),
                     &program.into(),
                 )
                 .await
@@ -145,7 +142,7 @@ async fn main() {
             wallet_core
                 .send_privacy_preserving_tx(
                     accounts,
-                    &Program::serialize_instruction(instruction).unwrap(),
+                    Program::serialize_instruction(instruction).unwrap(),
                     &program.into(),
                 )
                 .await

integration_tests/Cargo.toml

@@ -11,16 +11,14 @@ sequencer_runner.workspace = true
 wallet.workspace = true
 common.workspace = true
 key_protocol.workspace = true
-proc_macro_test_attribute = { path = "./proc_macro_test_attribute" }
 
-clap = { workspace = true, features = ["derive", "env"] }
 anyhow.workspace = true
 env_logger.workspace = true
 log.workspace = true
-actix.workspace = true
 actix-web.workspace = true
 base64.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
 hex.workspace = true
 tempfile.workspace = true
 borsh.workspace = true
+futures.workspace = true

integration_tests/configs/sequencer/sequencer_config.json

@@ -1,12 +1,12 @@
 {
-    "home": "./sequencer",
+    "home": "",
     "override_rust_log": null,
     "genesis_id": 1,
     "is_genesis_random": true,
     "max_num_tx_in_block": 20,
     "mempool_max_size": 10000,
     "block_create_timeout_millis": 10000,
-    "port": 3040,
+    "port": 0,
     "initial_accounts": [
         {
             "account_id": "6iArKUXxhUJqS7kCaPNhwMWt3ro71PDyBj7jwAyE2VQV",
@@ -155,4 +155,4 @@
         37,
         37
     ]
-}
+}

integration_tests/configs/wallet/wallet_config.json

@@ -1,10 +1,11 @@
 {
   "override_rust_log": null,
-  "sequencer_addr": "http://127.0.0.1:3040",
+  "sequencer_addr": "",
   "seq_poll_timeout_millis": 12000,
   "seq_tx_poll_max_blocks": 5,
   "seq_poll_max_retries": 5,
   "seq_block_poll_max_amount": 100,
+  "basic_auth": null,
   "initial_accounts": [
     {
       "Public": {
@@ -542,6 +543,5 @@
         }
       }
     }
-  ],
-  "basic_auth": null
+  ]
 }

integration_tests/proc_macro_test_attribute/Cargo.toml

@@ -1,9 +0,0 @@
-[package]
-name = "proc_macro_test_attribute"
-version = "0.1.0"
-edition = "2024"
-
-[dependencies]
-
-[lib]
-proc-macro = true

integration_tests/proc_macro_test_attribute/src/lib.rs

@@ -1,49 +0,0 @@
-extern crate proc_macro;
-
-use proc_macro::*;
-
-#[proc_macro_attribute]
-pub fn nssa_integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
-    let input = item.to_string();
-
-    let fn_keyword = "fn ";
-    let fn_keyword_alternative = "fn\n";
-
-    let mut start_opt = None;
-    let mut fn_name = String::new();
-
-    if let Some(start) = input.find(fn_keyword) {
-        start_opt = Some(start);
-    } else if let Some(start) = input.find(fn_keyword_alternative) {
-        start_opt = Some(start);
-    }
-
-    if let Some(start) = start_opt {
-        let rest = &input[start + fn_keyword.len()..];
-        if let Some(end) = rest.find(|c: char| c == '(' || c.is_whitespace()) {
-            let name = &rest[..end];
-            fn_name = name.to_string();
-        }
-    } else {
-        println!("ERROR: keyword fn not found");
-    }
-
-    let extension = format!(
-        r#"
-    {input}
-
-    function_map.insert("{fn_name}".to_string(), |home_dir: PathBuf| Box::pin(async {{
-        let res = pre_test(home_dir).await.unwrap();
-
-        info!("Waiting for first block creation");
-        tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
-        {fn_name}().await;
-
-        post_test(res).await;
-    }}));
-    "#
-    );
-
-    extension.parse().unwrap()
-}

integration_tests/src/lib.rs

@@ -1,38 +1,25 @@
-use std::path::PathBuf;
+//! This library contains common code for integration tests.
+
+use std::{net::SocketAddr, path::PathBuf, sync::LazyLock};
 
 use actix_web::dev::ServerHandle;
-use anyhow::Result;
+use anyhow::{Context as _, Result};
 use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
-use clap::Parser;
 use common::{
     sequencer_client::SequencerClient,
     transaction::{EncodedTransaction, NSSATransaction},
 };
-use log::{info, warn};
+use futures::FutureExt as _;
+use log::debug;
 use nssa::PrivacyPreservingTransaction;
 use nssa_core::Commitment;
 use sequencer_core::config::SequencerConfig;
-use sequencer_runner::startup_sequencer;
 use tempfile::TempDir;
 use tokio::task::JoinHandle;
+use wallet::{WalletCore, config::WalletConfigOverrides};
 
-use crate::test_suite_map::{prepare_function_map, tps_test};
-
-#[macro_use]
-extern crate proc_macro_test_attribute;
-
-pub mod test_suite_map;
-
-mod tps_test_utils;
-
-#[derive(Parser, Debug)]
-#[clap(version)]
-struct Args {
-    /// Path to configs
-    home_dir: PathBuf,
-    /// Test name
-    test_name: String,
-}
+// TODO: Remove this and control time from tests
+pub const TIME_TO_WAIT_FOR_BLOCK_SECONDS: u64 = 12;
 
 pub const ACC_SENDER: &str = "BLgCRDXYdQPMMWVHYRFGQZbgeHx9frkipa8GtpG2Syqy";
 pub const ACC_RECEIVER: &str = "Gj1mJy5W7J5pfmLRujmQaLfLMWidNxQ6uwnhb666ZwHw";
@@ -40,104 +27,181 @@ pub const ACC_RECEIVER: &str = "Gj1mJy5W7J5pfmLRujmQaLfLMWidNxQ6uwnhb666ZwHw";
 pub const ACC_SENDER_PRIVATE: &str = "3oCG8gqdKLMegw4rRfyaMQvuPHpcASt7xwttsmnZLSkw";
 pub const ACC_RECEIVER_PRIVATE: &str = "AKTcXgJ1xoynta1Ec7y6Jso1z1JQtHqd7aPQ1h9er6xX";
 
-pub const TIME_TO_WAIT_FOR_BLOCK_SECONDS: u64 = 12;
-
 pub const NSSA_PROGRAM_FOR_TEST_DATA_CHANGER: &str = "data_changer.bin";
 
-fn make_public_account_input_from_str(account_id: &str) -> String {
+static LOGGER: LazyLock<()> = LazyLock::new(env_logger::init);
+
+/// Test context which sets up a sequencer and a wallet for integration tests.
+///
+/// It's memory and logically safe to create multiple instances of this struct in parallel tests,
+/// as each instance uses its own temporary directories for sequencer and wallet data.
+pub struct TestContext {
+    sequencer_server_handle: ServerHandle,
+    sequencer_loop_handle: JoinHandle<Result<()>>,
+    sequencer_client: SequencerClient,
+    wallet: WalletCore,
+    _temp_sequencer_dir: TempDir,
+    _temp_wallet_dir: TempDir,
+}
+
+impl TestContext {
+    /// Create new test context.
+    pub async fn new() -> Result<Self> {
+        let manifest_dir = env!("CARGO_MANIFEST_DIR");
+
+        let sequencer_config_path =
+            PathBuf::from(manifest_dir).join("configs/sequencer/sequencer_config.json");
+
+        let sequencer_config = SequencerConfig::from_path(&sequencer_config_path)
+            .context("Failed to create sequencer config from file")?;
+
+        Self::new_with_sequencer_config(sequencer_config).await
+    }
+
+    /// Create new test context with custom sequencer config.
+    ///
+    /// `home` and `port` fields of the provided config will be overridden to meet tests parallelism
+    /// requirements.
+    pub async fn new_with_sequencer_config(sequencer_config: SequencerConfig) -> Result<Self> {
+        // Ensure logger is initialized only once
+        *LOGGER;
+
+        debug!("Test context setup");
+
+        let (sequencer_server_handle, sequencer_addr, sequencer_loop_handle, temp_sequencer_dir) =
+            Self::setup_sequencer(sequencer_config)
+                .await
+                .context("Failed to setup sequencer")?;
+
+        // Convert 0.0.0.0 to 127.0.0.1 for client connections
+        // When binding to port 0, the server binds to 0.0.0.0:<random_port>
+        // but clients need to connect to 127.0.0.1:<port> to work reliably
+        let sequencer_addr = if sequencer_addr.ip().is_unspecified() {
+            format!("http://127.0.0.1:{}", sequencer_addr.port())
+        } else {
+            format!("http://{sequencer_addr}")
+        };
+
+        let (wallet, temp_wallet_dir) = Self::setup_wallet(sequencer_addr.clone())
+            .await
+            .context("Failed to setup wallet")?;
+
+        let sequencer_client =
+            SequencerClient::new(sequencer_addr).context("Failed to create sequencer client")?;
+
+        Ok(Self {
+            sequencer_server_handle,
+            sequencer_loop_handle,
+            sequencer_client,
+            wallet,
+            _temp_sequencer_dir: temp_sequencer_dir,
+            _temp_wallet_dir: temp_wallet_dir,
+        })
+    }
+
+    async fn setup_sequencer(
+        mut config: SequencerConfig,
+    ) -> Result<(ServerHandle, SocketAddr, JoinHandle<Result<()>>, TempDir)> {
+        let temp_sequencer_dir =
+            tempfile::tempdir().context("Failed to create temp dir for sequencer home")?;
+
+        debug!(
+            "Using temp sequencer home at {:?}",
+            temp_sequencer_dir.path()
+        );
+        config.home = temp_sequencer_dir.path().to_owned();
+        // Setting port to 0 lets the OS choose a free port for us
+        config.port = 0;
+
+        let (sequencer_server_handle, sequencer_addr, sequencer_loop_handle) =
+            sequencer_runner::startup_sequencer(config).await?;
+
+        Ok((
+            sequencer_server_handle,
+            sequencer_addr,
+            sequencer_loop_handle,
+            temp_sequencer_dir,
+        ))
+    }
+
+    async fn setup_wallet(sequencer_addr: String) -> Result<(WalletCore, TempDir)> {
+        let manifest_dir = env!("CARGO_MANIFEST_DIR");
+        let wallet_config_source_path =
+            PathBuf::from(manifest_dir).join("configs/wallet/wallet_config.json");
+
+        let temp_wallet_dir =
+            tempfile::tempdir().context("Failed to create temp dir for wallet home")?;
+
+        let config_path = temp_wallet_dir.path().join("wallet_config.json");
+        std::fs::copy(&wallet_config_source_path, &config_path)
+            .context("Failed to copy wallet config to temp dir")?;
+
+        let storage_path = temp_wallet_dir.path().join("storage.json");
+        let config_overrides = WalletConfigOverrides {
+            sequencer_addr: Some(sequencer_addr),
+            ..Default::default()
+        };
+
+        let wallet = WalletCore::new_init_storage(
+            config_path,
+            storage_path,
+            Some(config_overrides),
+            "test_pass".to_owned(),
+        )
+        .context("Failed to init wallet")?;
+        wallet
+            .store_persistent_data()
+            .await
+            .context("Failed to store wallet persistent data")?;
+
+        Ok((wallet, temp_wallet_dir))
+    }
+
+    /// Get reference to the wallet.
+    pub fn wallet(&self) -> &WalletCore {
+        &self.wallet
+    }
+
+    /// Get mutable reference to the wallet.
+    pub fn wallet_mut(&mut self) -> &mut WalletCore {
+        &mut self.wallet
+    }
+
+    /// Get reference to the sequencer client.
+    pub fn sequencer_client(&self) -> &SequencerClient {
+        &self.sequencer_client
+    }
+}
+
+impl Drop for TestContext {
+    fn drop(&mut self) {
+        debug!("Test context cleanup");
+
+        let Self {
+            sequencer_server_handle,
+            sequencer_loop_handle,
+            sequencer_client: _,
+            wallet: _,
+            _temp_sequencer_dir,
+            _temp_wallet_dir,
+        } = self;
+
+        sequencer_loop_handle.abort();
+
+        // Can't wait here as Drop can't be async, but anyway stop signal should be sent
+        sequencer_server_handle.stop(true).now_or_never();
+    }
+}
+
+pub fn format_public_account_id(account_id: &str) -> String {
     format!("Public/{account_id}")
 }
 
-fn make_private_account_input_from_str(account_id: &str) -> String {
+pub fn format_private_account_id(account_id: &str) -> String {
     format!("Private/{account_id}")
 }
 
-#[allow(clippy::type_complexity)]
-pub async fn pre_test(
-    home_dir: PathBuf,
-) -> Result<(ServerHandle, JoinHandle<Result<()>>, TempDir)> {
-    wallet::cli::execute_setup("test_pass".to_owned()).await?;
-
-    let home_dir_sequencer = home_dir.join("sequencer");
-
-    let mut sequencer_config =
-        sequencer_runner::config::from_file(home_dir_sequencer.join("sequencer_config.json"))
-            .unwrap();
-
-    let temp_dir_sequencer = replace_home_dir_with_temp_dir_in_configs(&mut sequencer_config);
-
-    let (seq_http_server_handle, sequencer_loop_handle) =
-        startup_sequencer(sequencer_config).await?;
-
-    Ok((
-        seq_http_server_handle,
-        sequencer_loop_handle,
-        temp_dir_sequencer,
-    ))
-}
-
-pub fn replace_home_dir_with_temp_dir_in_configs(
-    sequencer_config: &mut SequencerConfig,
-) -> TempDir {
-    let temp_dir_sequencer = tempfile::tempdir().unwrap();
-
-    sequencer_config.home = temp_dir_sequencer.path().to_path_buf();
-
-    temp_dir_sequencer
-}
-
-#[allow(clippy::type_complexity)]
-pub async fn post_test(residual: (ServerHandle, JoinHandle<Result<()>>, TempDir)) {
-    let (seq_http_server_handle, sequencer_loop_handle, _) = residual;
-
-    info!("Cleanup");
-
-    sequencer_loop_handle.abort();
-    seq_http_server_handle.stop(true).await;
-
-    let wallet_home = wallet::helperfunctions::get_home().unwrap();
-    let persistent_data_home = wallet_home.join("storage.json");
-
-    // Removing persistent accounts after run to not affect other executions
-    // Not necessary an error, if fails as there is tests for failure scenario
-    let _ = std::fs::remove_file(persistent_data_home)
-        .inspect_err(|err| warn!("Failed to remove persistent data with err {err:#?}"));
-
-    // At this point all of the references to sequencer_core must be lost.
-    // So they are dropped and tempdirs will be dropped too,
-}
-
-pub async fn main_tests_runner() -> Result<()> {
-    env_logger::init();
-
-    let args = Args::parse();
-    let Args {
-        home_dir,
-        test_name,
-    } = args;
-
-    let function_map = prepare_function_map();
-
-    match test_name.as_str() {
-        "all" => {
-            // Tests that use default config
-            for (_, fn_pointer) in function_map {
-                fn_pointer(home_dir.clone()).await;
-            }
-            // Run TPS test with its own specific config
-            tps_test().await;
-        }
-        _ => {
-            let fn_pointer = function_map.get(&test_name).expect("Unknown test name");
-
-            fn_pointer(home_dir.clone()).await;
-        }
-    }
-
-    Ok(())
-}
-
-async fn fetch_privacy_preserving_tx(
+pub async fn fetch_privacy_preserving_tx(
     seq_client: &SequencerClient,
     tx_hash: String,
 ) -> PrivacyPreservingTransaction {
@@ -161,7 +225,7 @@ async fn fetch_privacy_preserving_tx(
     }
 }
 
-async fn verify_commitment_is_in_state(
+pub async fn verify_commitment_is_in_state(
     commitment: Commitment,
     seq_client: &SequencerClient,
 ) -> bool {
@@ -173,15 +237,15 @@ async fn verify_commitment_is_in_state(
 
 #[cfg(test)]
 mod tests {
-    use crate::{make_private_account_input_from_str, make_public_account_input_from_str};
+    use super::{format_private_account_id, format_public_account_id};
 
     #[test]
     fn correct_account_id_from_prefix() {
         let account_id1 = "cafecafe";
         let account_id2 = "deadbeaf";
 
-        let account_id1_pub = make_public_account_input_from_str(account_id1);
-        let account_id2_priv = make_private_account_input_from_str(account_id2);
+        let account_id1_pub = format_public_account_id(account_id1);
+        let account_id2_priv = format_private_account_id(account_id2);
 
         assert_eq!(account_id1_pub, "Public/cafecafe".to_string());
         assert_eq!(account_id2_priv, "Private/deadbeaf".to_string());

integration_tests/src/main.rs

@@ -1,15 +0,0 @@
-use anyhow::Result;
-use integration_tests::main_tests_runner;
-
-pub const NUM_THREADS: usize = 8;
-
-fn main() -> Result<()> {
-    actix::System::with_tokio_rt(|| {
-        tokio::runtime::Builder::new_multi_thread()
-            .worker_threads(NUM_THREADS)
-            .enable_all()
-            .build()
-            .unwrap()
-    })
-    .block_on(main_tests_runner())
-}

integration_tests/tests/account.rs

@@ -0,0 +1,28 @@
+use anyhow::Result;
+use integration_tests::{ACC_SENDER, TestContext};
+use log::info;
+use nssa::program::Program;
+use tokio::test;
+
+#[test]
+async fn get_existing_account() -> Result<()> {
+    let ctx = TestContext::new().await?;
+
+    let account = ctx
+        .sequencer_client()
+        .get_account(ACC_SENDER.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        account.program_owner,
+        Program::authenticated_transfer_program().id()
+    );
+    assert_eq!(account.balance, 10000);
+    assert!(account.data.is_empty());
+    assert_eq!(account.nonce, 0);
+
+    info!("Successfully retrieved account with correct details");
+
+    Ok(())
+}

integration_tests/tests/amm.rs

@@ -0,0 +1,405 @@
+use std::time::Duration;
+
+use anyhow::Result;
+use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, format_public_account_id};
+use log::info;
+use tokio::test;
+use wallet::cli::{
+    Command, SubcommandReturnValue,
+    account::{AccountSubcommand, NewSubcommand},
+    programs::{amm::AmmProgramAgnosticSubcommand, token::TokenProgramAgnosticSubcommand},
+};
+
+#[test]
+async fn amm_public() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create new account for the token definition
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id_1,
+    } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for the token supply holder
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id_1,
+    } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for receiving a token transaction
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id_1,
+    } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for the token definition
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id_2,
+    } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for the token supply holder
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id_2,
+    } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for receiving a token transaction
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id_2,
+    } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new token
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_public_account_id(&definition_account_id_1.to_string()),
+        supply_account_id: format_public_account_id(&supply_account_id_1.to_string()),
+        name: "A NAM1".to_string(),
+        total_supply: 37,
+    };
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Transfer 7 tokens from `supply_acc` to the account at account_id `recipient_account_id_1`
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from: format_public_account_id(&supply_account_id_1.to_string()),
+        to: Some(format_public_account_id(
+            &recipient_account_id_1.to_string(),
+        )),
+        to_npk: None,
+        to_ipk: None,
+        amount: 7,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Create new token
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_public_account_id(&definition_account_id_2.to_string()),
+        supply_account_id: format_public_account_id(&supply_account_id_2.to_string()),
+        name: "A NAM2".to_string(),
+        total_supply: 37,
+    };
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Transfer 7 tokens from `supply_acc` to the account at account_id `recipient_account_id_2`
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from: format_public_account_id(&supply_account_id_2.to_string()),
+        to: Some(format_public_account_id(
+            &recipient_account_id_2.to_string(),
+        )),
+        to_npk: None,
+        to_ipk: None,
+        amount: 7,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("=================== SETUP FINISHED ===============");
+
+    // Create new AMM
+
+    // Setup accounts
+    // Create new account for the user holding lp
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: user_holding_lp,
+    } = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Send creation tx
+    let subcommand = AmmProgramAgnosticSubcommand::New {
+        user_holding_a: format_public_account_id(&recipient_account_id_1.to_string()),
+        user_holding_b: format_public_account_id(&recipient_account_id_2.to_string()),
+        user_holding_lp: format_public_account_id(&user_holding_lp.to_string()),
+        balance_a: 3,
+        balance_b: 3,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::AMM(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let user_holding_a_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_1.to_string())
+        .await?
+        .account;
+
+    let user_holding_b_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_2.to_string())
+        .await?
+        .account;
+
+    let user_holding_lp_acc = ctx
+        .sequencer_client()
+        .get_account(user_holding_lp.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
+        4
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_b_acc.data[33..].try_into().unwrap()),
+        4
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_lp_acc.data[33..].try_into().unwrap()),
+        3
+    );
+
+    info!("=================== AMM DEFINITION FINISHED ===============");
+
+    // Make swap
+
+    let subcommand = AmmProgramAgnosticSubcommand::Swap {
+        user_holding_a: format_public_account_id(&recipient_account_id_1.to_string()),
+        user_holding_b: format_public_account_id(&recipient_account_id_2.to_string()),
+        amount_in: 2,
+        min_amount_out: 1,
+        token_definition: definition_account_id_1.to_string(),
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::AMM(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let user_holding_a_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_1.to_string())
+        .await?
+        .account;
+
+    let user_holding_b_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_2.to_string())
+        .await?
+        .account;
+
+    let user_holding_lp_acc = ctx
+        .sequencer_client()
+        .get_account(user_holding_lp.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
+        2
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_b_acc.data[33..].try_into().unwrap()),
+        5
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_lp_acc.data[33..].try_into().unwrap()),
+        3
+    );
+
+    info!("=================== FIRST SWAP FINISHED ===============");
+
+    // Make swap
+
+    let subcommand = AmmProgramAgnosticSubcommand::Swap {
+        user_holding_a: format_public_account_id(&recipient_account_id_1.to_string()),
+        user_holding_b: format_public_account_id(&recipient_account_id_2.to_string()),
+        amount_in: 2,
+        min_amount_out: 1,
+        token_definition: definition_account_id_2.to_string(),
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::AMM(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let user_holding_a_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_1.to_string())
+        .await?
+        .account;
+
+    let user_holding_b_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_2.to_string())
+        .await?
+        .account;
+
+    let user_holding_lp_acc = ctx
+        .sequencer_client()
+        .get_account(user_holding_lp.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
+        4
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_b_acc.data[33..].try_into().unwrap()),
+        3
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_lp_acc.data[33..].try_into().unwrap()),
+        3
+    );
+
+    info!("=================== SECOND SWAP FINISHED ===============");
+
+    // Add liquidity
+
+    let subcommand = AmmProgramAgnosticSubcommand::AddLiquidity {
+        user_holding_a: format_public_account_id(&recipient_account_id_1.to_string()),
+        user_holding_b: format_public_account_id(&recipient_account_id_2.to_string()),
+        user_holding_lp: format_public_account_id(&user_holding_lp.to_string()),
+        min_amount_lp: 1,
+        max_amount_a: 2,
+        max_amount_b: 2,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::AMM(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let user_holding_a_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_1.to_string())
+        .await?
+        .account;
+
+    let user_holding_b_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_2.to_string())
+        .await?
+        .account;
+
+    let user_holding_lp_acc = ctx
+        .sequencer_client()
+        .get_account(user_holding_lp.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
+        3
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_b_acc.data[33..].try_into().unwrap()),
+        1
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_lp_acc.data[33..].try_into().unwrap()),
+        4
+    );
+
+    info!("=================== ADD LIQ FINISHED ===============");
+
+    // Remove liquidity
+
+    let subcommand = AmmProgramAgnosticSubcommand::RemoveLiquidity {
+        user_holding_a: format_public_account_id(&recipient_account_id_1.to_string()),
+        user_holding_b: format_public_account_id(&recipient_account_id_2.to_string()),
+        user_holding_lp: format_public_account_id(&user_holding_lp.to_string()),
+        balance_lp: 2,
+        min_amount_a: 1,
+        min_amount_b: 1,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::AMM(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let user_holding_a_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_1.to_string())
+        .await?
+        .account;
+
+    let user_holding_b_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_2.to_string())
+        .await?
+        .account;
+
+    let user_holding_lp_acc = ctx
+        .sequencer_client()
+        .get_account(user_holding_lp.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
+        5
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_b_acc.data[33..].try_into().unwrap()),
+        4
+    );
+
+    assert_eq!(
+        u128::from_le_bytes(user_holding_lp_acc.data[33..].try_into().unwrap()),
+        2
+    );
+
+    info!("Success!");
+
+    Ok(())
+}

integration_tests/tests/auth_transfer/main.rs

@@ -0,0 +1,2 @@
+mod private;
+mod public;

integration_tests/tests/auth_transfer/private.rs

@@ -0,0 +1,417 @@
+use std::time::Duration;
+
+use anyhow::{Context as _, Result};
+use integration_tests::{
+    ACC_RECEIVER, ACC_RECEIVER_PRIVATE, ACC_SENDER, ACC_SENDER_PRIVATE,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx,
+    format_private_account_id, format_public_account_id, verify_commitment_is_in_state,
+};
+use log::info;
+use nssa::{AccountId, program::Program};
+use nssa_core::{NullifierPublicKey, encryption::shared_key_derivation::Secp256k1Point};
+use tokio::test;
+use wallet::cli::{
+    Command, SubcommandReturnValue,
+    account::{AccountSubcommand, NewSubcommand},
+    programs::native_token_transfer::AuthTransferSubcommand,
+};
+
+#[test]
+async fn private_transfer_to_owned_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let from: AccountId = ACC_SENDER_PRIVATE.parse()?;
+    let to: AccountId = ACC_RECEIVER_PRIVATE.parse()?;
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&from.to_string()),
+        to: Some(format_private_account_id(&to.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let new_commitment1 = ctx
+        .wallet()
+        .get_private_account_commitment(&from)
+        .context("Failed to get private account commitment for sender")?;
+    assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await);
+
+    let new_commitment2 = ctx
+        .wallet()
+        .get_private_account_commitment(&to)
+        .context("Failed to get private account commitment for receiver")?;
+    assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await);
+
+    info!("Successfully transferred privately to owned account");
+
+    Ok(())
+}
+
+#[test]
+async fn private_transfer_to_foreign_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let from: AccountId = ACC_SENDER_PRIVATE.parse()?;
+    let to_npk = NullifierPublicKey([42; 32]);
+    let to_npk_string = hex::encode(to_npk.0);
+    let to_ipk = Secp256k1Point::from_scalar(to_npk.0);
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&from.to_string()),
+        to: None,
+        to_npk: Some(to_npk_string),
+        to_ipk: Some(hex::encode(to_ipk.0)),
+        amount: 100,
+    });
+
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash } = result else {
+        anyhow::bail!("Expected PrivacyPreservingTransfer return value");
+    };
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let new_commitment1 = ctx
+        .wallet()
+        .get_private_account_commitment(&from)
+        .context("Failed to get private account commitment for sender")?;
+
+    let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash.clone()).await;
+    assert_eq!(tx.message.new_commitments[0], new_commitment1);
+
+    assert_eq!(tx.message.new_commitments.len(), 2);
+    for commitment in tx.message.new_commitments.into_iter() {
+        assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await);
+    }
+
+    info!("Successfully transferred privately to foreign account");
+
+    Ok(())
+}
+
+#[test]
+async fn deshielded_transfer_to_public_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let from: AccountId = ACC_SENDER_PRIVATE.parse()?;
+    let to: AccountId = ACC_RECEIVER.parse()?;
+
+    // Check initial balance of the private sender
+    let from_acc = ctx
+        .wallet()
+        .get_account_private(&from)
+        .context("Failed to get sender's private account")?;
+    assert_eq!(from_acc.balance, 10000);
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&from.to_string()),
+        to: Some(format_public_account_id(&to.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let from_acc = ctx
+        .wallet()
+        .get_account_private(&from)
+        .context("Failed to get sender's private account")?;
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&from)
+        .context("Failed to get private account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    let acc_2_balance = ctx
+        .sequencer_client()
+        .get_account_balance(to.to_string())
+        .await?;
+
+    assert_eq!(from_acc.balance, 9900);
+    assert_eq!(acc_2_balance.balance, 20100);
+
+    info!("Successfully deshielded transfer to public account");
+
+    Ok(())
+}
+
+#[test]
+async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let from: AccountId = ACC_SENDER_PRIVATE.parse()?;
+
+    // Create a new private account
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None }));
+
+    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: to_account_id,
+    } = sub_ret
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Get the keys for the newly created account
+    let (to_keys, _) = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .get_private_account(&to_account_id)
+        .cloned()
+        .context("Failed to get private account")?;
+
+    // Send to this account using claiming path (using npk and ipk instead of account ID)
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&from.to_string()),
+        to: None,
+        to_npk: Some(hex::encode(to_keys.nullifer_public_key.0)),
+        to_ipk: Some(hex::encode(to_keys.incoming_viewing_public_key.0)),
+        amount: 100,
+    });
+
+    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash } = sub_ret else {
+        anyhow::bail!("Expected PrivacyPreservingTransfer return value");
+    };
+
+    let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash.clone()).await;
+
+    // Sync the wallet to claim the new account
+    let command = Command::Account(AccountSubcommand::SyncPrivate {});
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let new_commitment1 = ctx
+        .wallet()
+        .get_private_account_commitment(&from)
+        .context("Failed to get private account commitment for sender")?;
+    assert_eq!(tx.message.new_commitments[0], new_commitment1);
+
+    assert_eq!(tx.message.new_commitments.len(), 2);
+    for commitment in tx.message.new_commitments.into_iter() {
+        assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await);
+    }
+
+    let to_res_acc = ctx
+        .wallet()
+        .get_account_private(&to_account_id)
+        .context("Failed to get recipient's private account")?;
+    assert_eq!(to_res_acc.balance, 100);
+
+    info!("Successfully transferred using claiming path");
+
+    Ok(())
+}
+
+#[test]
+async fn shielded_transfer_to_owned_private_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let from: AccountId = ACC_SENDER.parse()?;
+    let to: AccountId = ACC_RECEIVER_PRIVATE.parse()?;
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(&from.to_string()),
+        to: Some(format_private_account_id(&to.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let acc_to = ctx
+        .wallet()
+        .get_account_private(&to)
+        .context("Failed to get receiver's private account")?;
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&to)
+        .context("Failed to get receiver's commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    let acc_from_balance = ctx
+        .sequencer_client()
+        .get_account_balance(from.to_string())
+        .await?;
+
+    assert_eq!(acc_from_balance.balance, 9900);
+    assert_eq!(acc_to.balance, 20100);
+
+    info!("Successfully shielded transfer to owned private account");
+
+    Ok(())
+}
+
+#[test]
+async fn shielded_transfer_to_foreign_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let to_npk = NullifierPublicKey([42; 32]);
+    let to_npk_string = hex::encode(to_npk.0);
+    let to_ipk = Secp256k1Point::from_scalar(to_npk.0);
+    let from: AccountId = ACC_SENDER.parse()?;
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(&from.to_string()),
+        to: None,
+        to_npk: Some(to_npk_string),
+        to_ipk: Some(hex::encode(to_ipk.0)),
+        amount: 100,
+    });
+
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash } = result else {
+        anyhow::bail!("Expected PrivacyPreservingTransfer return value");
+    };
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash).await;
+
+    let acc_1_balance = ctx
+        .sequencer_client()
+        .get_account_balance(from.to_string())
+        .await?;
+
+    assert!(
+        verify_commitment_is_in_state(
+            tx.message.new_commitments[0].clone(),
+            ctx.sequencer_client()
+        )
+        .await
+    );
+
+    assert_eq!(acc_1_balance.balance, 9900);
+
+    info!("Successfully shielded transfer to foreign account");
+
+    Ok(())
+}
+
+#[test]
+#[ignore = "Flaky, TODO: #197"]
+async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // NOTE: This test needs refactoring - continuous run mode doesn't work well with TestContext
+    // The original implementation spawned wallet::cli::execute_continuous_run() in background
+    // but this conflicts with TestContext's wallet management
+
+    let from: AccountId = ACC_SENDER_PRIVATE.parse()?;
+
+    // Create a new private account
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None }));
+    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: to_account_id,
+    } = sub_ret
+    else {
+        anyhow::bail!("Failed to register account");
+    };
+
+    // Get the newly created account's keys
+    let (to_keys, _) = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .get_private_account(&to_account_id)
+        .cloned()
+        .context("Failed to get private account")?;
+
+    // Send transfer using nullifier and incoming viewing public keys
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&from.to_string()),
+        to: None,
+        to_npk: Some(hex::encode(to_keys.nullifer_public_key.0)),
+        to_ipk: Some(hex::encode(to_keys.incoming_viewing_public_key.0)),
+        amount: 100,
+    });
+
+    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash } = sub_ret else {
+        anyhow::bail!("Failed to send transaction");
+    };
+
+    let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash.clone()).await;
+
+    info!("Waiting for next blocks to check if continuous run fetches account");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify commitments are in state
+    assert_eq!(tx.message.new_commitments.len(), 2);
+    for commitment in tx.message.new_commitments.into_iter() {
+        assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await);
+    }
+
+    // Verify receiver account balance
+    let to_res_acc = ctx
+        .wallet()
+        .get_account_private(&to_account_id)
+        .context("Failed to get receiver account")?;
+
+    assert_eq!(to_res_acc.balance, 100);
+
+    Ok(())
+}
+
+#[test]
+async fn initialize_private_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None }));
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Init {
+        account_id: format_private_account_id(&account_id.to_string()),
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Syncing private accounts");
+    let command = Command::Account(AccountSubcommand::SyncPrivate {});
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&account_id)
+        .context("Failed to get private account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    let account = ctx
+        .wallet()
+        .get_account_private(&account_id)
+        .context("Failed to get private account")?;
+
+    assert_eq!(
+        account.program_owner,
+        Program::authenticated_transfer_program().id()
+    );
+    assert_eq!(account.balance, 0);
+    assert!(account.data.is_empty());
+
+    info!("Successfully initialized private account");
+
+    Ok(())
+}

integration_tests/tests/auth_transfer/public.rs

@@ -0,0 +1,248 @@
+use std::time::Duration;
+
+use anyhow::Result;
+use integration_tests::{
+    ACC_RECEIVER, ACC_SENDER, TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, format_public_account_id,
+};
+use log::info;
+use nssa::program::Program;
+use tokio::test;
+use wallet::cli::{
+    Command, SubcommandReturnValue,
+    account::{AccountSubcommand, NewSubcommand},
+    programs::native_token_transfer::AuthTransferSubcommand,
+};
+
+#[test]
+async fn successful_transfer_to_existing_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(ACC_SENDER),
+        to: Some(format_public_account_id(ACC_RECEIVER)),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Checking correct balance move");
+    let acc_1_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_SENDER.to_string())
+        .await?;
+    let acc_2_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_RECEIVER.to_string())
+        .await?;
+
+    info!("Balance of sender: {acc_1_balance:#?}");
+    info!("Balance of receiver: {acc_2_balance:#?}");
+
+    assert_eq!(acc_1_balance.balance, 9900);
+    assert_eq!(acc_2_balance.balance, 20100);
+
+    Ok(())
+}
+
+#[test]
+pub async fn successful_transfer_to_new_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None }));
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command)
+        .await
+        .unwrap();
+
+    let new_persistent_account_id = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .account_ids()
+        .map(ToString::to_string)
+        .find(|acc_id| acc_id != ACC_SENDER && acc_id != ACC_RECEIVER)
+        .expect("Failed to find newly created account in the wallet storage");
+
+    if new_persistent_account_id == String::new() {
+        panic!("Failed to produce new account, not present in persistent accounts");
+    }
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(ACC_SENDER),
+        to: Some(format_public_account_id(&new_persistent_account_id)),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Checking correct balance move");
+    let acc_1_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_SENDER.to_string())
+        .await?;
+    let acc_2_balance = ctx
+        .sequencer_client()
+        .get_account_balance(new_persistent_account_id)
+        .await?;
+
+    info!("Balance of sender: {acc_1_balance:#?}");
+    info!("Balance of receiver: {acc_2_balance:#?}");
+
+    assert_eq!(acc_1_balance.balance, 9900);
+    assert_eq!(acc_2_balance.balance, 100);
+
+    Ok(())
+}
+
+#[test]
+async fn failed_transfer_with_insufficient_balance() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(ACC_SENDER),
+        to: Some(format_public_account_id(ACC_RECEIVER)),
+        to_npk: None,
+        to_ipk: None,
+        amount: 1000000,
+    });
+
+    let failed_send = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await;
+    assert!(failed_send.is_err());
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Checking balances unchanged");
+    let acc_1_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_SENDER.to_string())
+        .await?;
+    let acc_2_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_RECEIVER.to_string())
+        .await?;
+
+    info!("Balance of sender: {acc_1_balance:#?}");
+    info!("Balance of receiver: {acc_2_balance:#?}");
+
+    assert_eq!(acc_1_balance.balance, 10000);
+    assert_eq!(acc_2_balance.balance, 20000);
+
+    Ok(())
+}
+
+#[test]
+async fn two_consecutive_successful_transfers() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // First transfer
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(ACC_SENDER),
+        to: Some(format_public_account_id(ACC_RECEIVER)),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Checking correct balance move after first transfer");
+    let acc_1_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_SENDER.to_string())
+        .await?;
+    let acc_2_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_RECEIVER.to_string())
+        .await?;
+
+    info!("Balance of sender: {acc_1_balance:#?}");
+    info!("Balance of receiver: {acc_2_balance:#?}");
+
+    assert_eq!(acc_1_balance.balance, 9900);
+    assert_eq!(acc_2_balance.balance, 20100);
+
+    info!("First TX Success!");
+
+    // Second transfer
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(ACC_SENDER),
+        to: Some(format_public_account_id(ACC_RECEIVER)),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Checking correct balance move after second transfer");
+    let acc_1_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_SENDER.to_string())
+        .await?;
+    let acc_2_balance = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_RECEIVER.to_string())
+        .await?;
+
+    info!("Balance of sender: {acc_1_balance:#?}");
+    info!("Balance of receiver: {acc_2_balance:#?}");
+
+    assert_eq!(acc_1_balance.balance, 9800);
+    assert_eq!(acc_2_balance.balance, 20200);
+
+    info!("Second TX Success!");
+
+    Ok(())
+}
+
+#[test]
+async fn initialize_public_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None }));
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Init {
+        account_id: format_public_account_id(&account_id.to_string()),
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Checking correct execution");
+    let account = ctx
+        .sequencer_client()
+        .get_account(account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        account.program_owner,
+        Program::authenticated_transfer_program().id()
+    );
+    assert_eq!(account.balance, 0);
+    assert_eq!(account.nonce, 1);
+    assert!(account.data.is_empty());
+
+    info!("Successfully initialized public account");
+
+    Ok(())
+}

integration_tests/tests/config.rs

@@ -0,0 +1,33 @@
+use anyhow::Result;
+use integration_tests::TestContext;
+use log::info;
+use tokio::test;
+use wallet::cli::{Command, config::ConfigSubcommand};
+
+#[test]
+async fn modify_config_field() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let old_seq_poll_timeout_millis = ctx.wallet().config().seq_poll_timeout_millis;
+
+    // Change config field
+    let command = Command::Config(ConfigSubcommand::Set {
+        key: "seq_poll_timeout_millis".to_string(),
+        value: "1000".to_string(),
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let new_seq_poll_timeout_millis = ctx.wallet().config().seq_poll_timeout_millis;
+    assert_eq!(new_seq_poll_timeout_millis, 1000);
+
+    // Return how it was at the beginning
+    let command = Command::Config(ConfigSubcommand::Set {
+        key: "seq_poll_timeout_millis".to_string(),
+        value: old_seq_poll_timeout_millis.to_string(),
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Successfully modified and restored config field");
+
+    Ok(())
+}

integration_tests/tests/keys_restoration.rs

@@ -0,0 +1,217 @@
+use std::{str::FromStr, time::Duration};
+
+use anyhow::Result;
+use integration_tests::{
+    ACC_SENDER, ACC_SENDER_PRIVATE, TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext,
+    format_private_account_id, format_public_account_id, verify_commitment_is_in_state,
+};
+use key_protocol::key_management::key_tree::chain_index::ChainIndex;
+use log::info;
+use nssa::{AccountId, program::Program};
+use tokio::test;
+use wallet::cli::{
+    Command, SubcommandReturnValue,
+    account::{AccountSubcommand, NewSubcommand},
+    programs::native_token_transfer::AuthTransferSubcommand,
+};
+
+#[test]
+async fn restore_keys_from_seed() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let from: AccountId = ACC_SENDER_PRIVATE.parse()?;
+
+    // Create first private account at root
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+        cci: Some(ChainIndex::root()),
+    }));
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: to_account_id1,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create second private account at /0
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+        cci: Some(ChainIndex::from_str("/0")?),
+    }));
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: to_account_id2,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Send to first private account
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&from.to_string()),
+        to: Some(format_private_account_id(&to_account_id1.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 100,
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    // Send to second private account
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&from.to_string()),
+        to: Some(format_private_account_id(&to_account_id2.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 101,
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let from: AccountId = ACC_SENDER.parse()?;
+
+    // Create first public account at root
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
+        cci: Some(ChainIndex::root()),
+    }));
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: to_account_id3,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create second public account at /0
+    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
+        cci: Some(ChainIndex::from_str("/0")?),
+    }));
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: to_account_id4,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Send to first public account
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(&from.to_string()),
+        to: Some(format_public_account_id(&to_account_id3.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 102,
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    // Send to second public account
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(&from.to_string()),
+        to: Some(format_public_account_id(&to_account_id4.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 103,
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Preparation complete, performing keys restoration");
+
+    // Restore keys from seed
+    wallet::cli::execute_keys_restoration(ctx.wallet_mut(), 10).await?;
+
+    // Verify restored private accounts
+    let acc1 = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .private_key_tree
+        .get_node(to_account_id1)
+        .expect("Acc 1 should be restored");
+
+    let acc2 = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .private_key_tree
+        .get_node(to_account_id2)
+        .expect("Acc 2 should be restored");
+
+    // Verify restored public accounts
+    let _acc3 = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .public_key_tree
+        .get_node(to_account_id3)
+        .expect("Acc 3 should be restored");
+
+    let _acc4 = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .public_key_tree
+        .get_node(to_account_id4)
+        .expect("Acc 4 should be restored");
+
+    assert_eq!(
+        acc1.value.1.program_owner,
+        Program::authenticated_transfer_program().id()
+    );
+    assert_eq!(
+        acc2.value.1.program_owner,
+        Program::authenticated_transfer_program().id()
+    );
+
+    assert_eq!(acc1.value.1.balance, 100);
+    assert_eq!(acc2.value.1.balance, 101);
+
+    info!("Tree checks passed, testing restored accounts can transact");
+
+    // Test that restored accounts can send transactions
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_private_account_id(&to_account_id1.to_string()),
+        to: Some(format_private_account_id(&to_account_id2.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 10,
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from: format_public_account_id(&to_account_id3.to_string()),
+        to: Some(format_public_account_id(&to_account_id4.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 11,
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify commitments exist for private accounts
+    let comm1 = ctx
+        .wallet()
+        .get_private_account_commitment(&to_account_id1)
+        .expect("Acc 1 commitment should exist");
+    let comm2 = ctx
+        .wallet()
+        .get_private_account_commitment(&to_account_id2)
+        .expect("Acc 2 commitment should exist");
+
+    assert!(verify_commitment_is_in_state(comm1, ctx.sequencer_client()).await);
+    assert!(verify_commitment_is_in_state(comm2, ctx.sequencer_client()).await);
+
+    // Verify public account balances
+    let acc3 = ctx
+        .sequencer_client()
+        .get_account_balance(to_account_id3.to_string())
+        .await?;
+    let acc4 = ctx
+        .sequencer_client()
+        .get_account_balance(to_account_id4.to_string())
+        .await?;
+
+    assert_eq!(acc3.balance, 91); // 102 - 11
+    assert_eq!(acc4.balance, 114); // 103 + 11
+
+    info!("Successfully restored keys and verified transactions");
+
+    Ok(())
+}

integration_tests/tests/pinata.rs

@@ -0,0 +1,175 @@
+use std::time::Duration;
+
+use anyhow::{Context as _, Result};
+use common::PINATA_BASE58;
+use integration_tests::{
+    ACC_SENDER, ACC_SENDER_PRIVATE, TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext,
+    format_private_account_id, format_public_account_id, verify_commitment_is_in_state,
+};
+use log::info;
+use tokio::test;
+use wallet::cli::{
+    Command, SubcommandReturnValue,
+    account::{AccountSubcommand, NewSubcommand},
+    programs::{
+        native_token_transfer::AuthTransferSubcommand, pinata::PinataProgramAgnosticSubcommand,
+    },
+};
+
+#[test]
+async fn claim_pinata_to_existing_public_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let pinata_prize = 150;
+    let command = Command::Pinata(PinataProgramAgnosticSubcommand::Claim {
+        to: format_public_account_id(ACC_SENDER),
+    });
+
+    let pinata_balance_pre = ctx
+        .sequencer_client()
+        .get_account_balance(PINATA_BASE58.to_string())
+        .await?
+        .balance;
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Checking correct balance move");
+    let pinata_balance_post = ctx
+        .sequencer_client()
+        .get_account_balance(PINATA_BASE58.to_string())
+        .await?
+        .balance;
+
+    let winner_balance_post = ctx
+        .sequencer_client()
+        .get_account_balance(ACC_SENDER.to_string())
+        .await?
+        .balance;
+
+    assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize);
+    assert_eq!(winner_balance_post, 10000 + pinata_prize);
+
+    info!("Successfully claimed pinata to public account");
+
+    Ok(())
+}
+
+#[test]
+async fn claim_pinata_to_existing_private_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let pinata_prize = 150;
+    let command = Command::Pinata(PinataProgramAgnosticSubcommand::Claim {
+        to: format_private_account_id(ACC_SENDER_PRIVATE),
+    });
+
+    let pinata_balance_pre = ctx
+        .sequencer_client()
+        .get_account_balance(PINATA_BASE58.to_string())
+        .await?
+        .balance;
+
+    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash: _ } = result else {
+        anyhow::bail!("Expected PrivacyPreservingTransfer return value");
+    };
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    info!("Syncing private accounts");
+    let command = Command::Account(AccountSubcommand::SyncPrivate {});
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&ACC_SENDER_PRIVATE.parse()?)
+        .context("Failed to get private account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    let pinata_balance_post = ctx
+        .sequencer_client()
+        .get_account_balance(PINATA_BASE58.to_string())
+        .await?
+        .balance;
+
+    assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize);
+
+    info!("Successfully claimed pinata to existing private account");
+
+    Ok(())
+}
+
+#[test]
+async fn claim_pinata_to_new_private_account() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let pinata_prize = 150;
+
+    // Create new private account
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: winner_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    let winner_account_id_formatted = format_private_account_id(&winner_account_id.to_string());
+
+    // Initialize account under auth transfer program
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Init {
+        account_id: winner_account_id_formatted.clone(),
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&winner_account_id)
+        .context("Failed to get private account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    // Claim pinata to the new private account
+    let command = Command::Pinata(PinataProgramAgnosticSubcommand::Claim {
+        to: winner_account_id_formatted,
+    });
+
+    let pinata_balance_pre = ctx
+        .sequencer_client()
+        .get_account_balance(PINATA_BASE58.to_string())
+        .await?
+        .balance;
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&winner_account_id)
+        .context("Failed to get private account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    let pinata_balance_post = ctx
+        .sequencer_client()
+        .get_account_balance(PINATA_BASE58.to_string())
+        .await?
+        .balance;
+
+    assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize);
+
+    info!("Successfully claimed pinata to new private account");
+
+    Ok(())
+}

integration_tests/tests/program_deployment.rs

@@ -0,0 +1,64 @@
+use std::{path::PathBuf, time::Duration};
+
+use anyhow::Result;
+use integration_tests::{
+    NSSA_PROGRAM_FOR_TEST_DATA_CHANGER, TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext,
+};
+use log::info;
+use nssa::{AccountId, program::Program};
+use tokio::test;
+use wallet::cli::Command;
+
+#[test]
+async fn deploy_and_execute_program() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    let manifest_dir = env!("CARGO_MANIFEST_DIR");
+    let binary_filepath: PathBuf = PathBuf::from(manifest_dir)
+        .join("../artifacts/test_program_methods")
+        .join(NSSA_PROGRAM_FOR_TEST_DATA_CHANGER);
+
+    let command = Command::DeployProgram {
+        binary_filepath: binary_filepath.clone(),
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // The program is the data changer and takes one account as input.
+    // We pass an uninitialized account and we expect after execution to be owned by the data
+    // changer program (NSSA account claiming mechanism) with data equal to [0] (due to program
+    // logic)
+    let bytecode = std::fs::read(binary_filepath)?;
+    let data_changer = Program::new(bytecode)?;
+    let account_id: AccountId = "11".repeat(16).parse()?;
+    let message = nssa::public_transaction::Message::try_new(
+        data_changer.id(),
+        vec![account_id],
+        vec![],
+        vec![0],
+    )?;
+    let witness_set = nssa::public_transaction::WitnessSet::for_message(&message, &[]);
+    let transaction = nssa::PublicTransaction::new(message, witness_set);
+    let _response = ctx.sequencer_client().send_tx_public(transaction).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let post_state_account = ctx
+        .sequencer_client()
+        .get_account(account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(post_state_account.program_owner, data_changer.id());
+    assert_eq!(post_state_account.balance, 0);
+    assert_eq!(post_state_account.data.as_ref(), &[0]);
+    assert_eq!(post_state_account.nonce, 0);
+
+    info!("Successfully deployed and executed program");
+
+    Ok(())
+}

integration_tests/tests/token.rs

@@ -0,0 +1,968 @@
+use std::time::Duration;
+
+use anyhow::{Context as _, Result};
+use integration_tests::{
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, format_private_account_id,
+    format_public_account_id, verify_commitment_is_in_state,
+};
+use key_protocol::key_management::key_tree::chain_index::ChainIndex;
+use log::info;
+use nssa::program::Program;
+use tokio::test;
+use wallet::cli::{
+    Command, SubcommandReturnValue,
+    account::{AccountSubcommand, NewSubcommand},
+    programs::token::TokenProgramAgnosticSubcommand,
+};
+
+#[test]
+async fn create_and_transfer_public_token() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create new account for the token definition
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for the token supply holder
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for receiving a token transaction
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new token
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_public_account_id(&definition_account_id.to_string()),
+        supply_account_id: format_public_account_id(&supply_account_id.to_string()),
+        name: "A NAME".to_string(),
+        total_supply: 37,
+    };
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Check the status of the token definition account
+    let definition_acc = ctx
+        .sequencer_client()
+        .get_account(definition_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(definition_acc.program_owner, Program::token().id());
+    // The data of a token definition account has the following layout:
+    // [ 0x00 || name (6 bytes) || total supply (little endian 16 bytes) || metadata id (32 bytes)]
+    assert_eq!(
+        definition_acc.data.as_ref(),
+        &[
+            0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+        ]
+    );
+
+    // Check the status of the token holding account with the total supply
+    let supply_acc = ctx
+        .sequencer_client()
+        .get_account(supply_account_id.to_string())
+        .await?
+        .account;
+
+    // The account must be owned by the token program
+    assert_eq!(supply_acc.program_owner, Program::token().id());
+    // The data of a token holding account has the following layout:
+    // [ 0x01 || corresponding_token_definition_id (32 bytes) || balance (little endian 16 bytes) ]
+    // First byte of the data equal to 1 means it's a token holding account
+    assert_eq!(supply_acc.data.as_ref()[0], 1);
+    // Bytes from 1 to 33 represent the id of the token this account is associated with
+    assert_eq!(
+        &supply_acc.data.as_ref()[1..33],
+        definition_account_id.to_bytes()
+    );
+    assert_eq!(u128::from_le_bytes(supply_acc.data[33..].try_into()?), 37);
+
+    // Transfer 7 tokens from supply_acc to recipient_account_id
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from: format_public_account_id(&supply_account_id.to_string()),
+        to: Some(format_public_account_id(&recipient_account_id.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 7,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Check the status of the supply account after transfer
+    let supply_acc = ctx
+        .sequencer_client()
+        .get_account(supply_account_id.to_string())
+        .await?
+        .account;
+    assert_eq!(supply_acc.program_owner, Program::token().id());
+    assert_eq!(supply_acc.data[0], 1);
+    assert_eq!(&supply_acc.data[1..33], definition_account_id.to_bytes());
+    assert_eq!(u128::from_le_bytes(supply_acc.data[33..].try_into()?), 30);
+
+    // Check the status of the recipient account after transfer
+    let recipient_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id.to_string())
+        .await?
+        .account;
+    assert_eq!(recipient_acc.program_owner, Program::token().id());
+    assert_eq!(recipient_acc.data[0], 1);
+    assert_eq!(&recipient_acc.data[1..33], definition_account_id.to_bytes());
+    assert_eq!(u128::from_le_bytes(recipient_acc.data[33..].try_into()?), 7);
+
+    // Burn 3 tokens from recipient_acc
+    let subcommand = TokenProgramAgnosticSubcommand::Burn {
+        definition: format_public_account_id(&definition_account_id.to_string()),
+        holder: format_public_account_id(&recipient_account_id.to_string()),
+        amount: 3,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Check the status of the token definition account after burn
+    let definition_acc = ctx
+        .sequencer_client()
+        .get_account(definition_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        definition_acc.data.as_ref(),
+        &[
+            0, 65, 32, 78, 65, 77, 69, 34, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+        ]
+    );
+
+    // Check the status of the recipient account after burn
+    let recipient_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(u128::from_le_bytes(recipient_acc.data[33..].try_into()?), 4);
+
+    // Mint 10 tokens at recipient_acc
+    let subcommand = TokenProgramAgnosticSubcommand::Mint {
+        definition: format_public_account_id(&definition_account_id.to_string()),
+        holder: Some(format_public_account_id(&recipient_account_id.to_string())),
+        holder_npk: None,
+        holder_ipk: None,
+        amount: 10,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Check the status of the token definition account after mint
+    let definition_acc = ctx
+        .sequencer_client()
+        .get_account(definition_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        definition_acc.data.as_ref(),
+        &[
+            0, 65, 32, 78, 65, 77, 69, 44, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+        ]
+    );
+
+    // Check the status of the recipient account after mint
+    let recipient_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        u128::from_le_bytes(recipient_acc.data[33..].try_into()?),
+        14
+    );
+
+    info!("Successfully created and transferred public token");
+
+    Ok(())
+}
+
+#[test]
+async fn create_and_transfer_token_with_private_supply() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create new account for the token definition (public)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for the token supply holder (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new account for receiving a token transaction (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create new token
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_public_account_id(&definition_account_id.to_string()),
+        supply_account_id: format_private_account_id(&supply_account_id.to_string()),
+        name: "A NAME".to_string(),
+        total_supply: 37,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Check the status of the token definition account
+    let definition_acc = ctx
+        .sequencer_client()
+        .get_account(definition_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(definition_acc.program_owner, Program::token().id());
+    assert_eq!(
+        definition_acc.data.as_ref(),
+        &[
+            0, 65, 32, 78, 65, 77, 69, 37, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+        ]
+    );
+
+    let new_commitment1 = ctx
+        .wallet()
+        .get_private_account_commitment(&supply_account_id)
+        .context("Failed to get supply account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await);
+
+    // Transfer 7 tokens from supply_acc to recipient_account_id
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from: format_private_account_id(&supply_account_id.to_string()),
+        to: Some(format_private_account_id(&recipient_account_id.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 7,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let new_commitment1 = ctx
+        .wallet()
+        .get_private_account_commitment(&supply_account_id)
+        .context("Failed to get supply account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await);
+
+    let new_commitment2 = ctx
+        .wallet()
+        .get_private_account_commitment(&recipient_account_id)
+        .context("Failed to get recipient account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await);
+
+    // Burn 3 tokens from recipient_acc
+    let subcommand = TokenProgramAgnosticSubcommand::Burn {
+        definition: format_public_account_id(&definition_account_id.to_string()),
+        holder: format_private_account_id(&recipient_account_id.to_string()),
+        amount: 3,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Check the token definition account after burn
+    let definition_acc = ctx
+        .sequencer_client()
+        .get_account(definition_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        definition_acc.data.as_ref(),
+        &[
+            0, 65, 32, 78, 65, 77, 69, 34, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+        ]
+    );
+
+    let new_commitment2 = ctx
+        .wallet()
+        .get_private_account_commitment(&recipient_account_id)
+        .context("Failed to get recipient account commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await);
+
+    // Check the recipient account balance after burn
+    let recipient_acc = ctx
+        .wallet()
+        .get_account_private(&recipient_account_id)
+        .context("Failed to get recipient account")?;
+
+    assert_eq!(
+        u128::from_le_bytes(recipient_acc.data[33..].try_into()?),
+        4 // 7 - 3
+    );
+
+    info!("Successfully created and transferred token with private supply");
+
+    Ok(())
+}
+
+#[test]
+async fn create_token_with_private_definition() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create token definition account (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
+            cci: Some(ChainIndex::root()),
+        })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create supply account (public)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
+            cci: Some(ChainIndex::root()),
+        })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create token with private definition
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_private_account_id(&definition_account_id.to_string()),
+        supply_account_id: format_public_account_id(&supply_account_id.to_string()),
+        name: "A NAME".to_string(),
+        total_supply: 37,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify private definition commitment
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&definition_account_id)
+        .context("Failed to get definition commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    // Verify supply account
+    let supply_acc = ctx
+        .sequencer_client()
+        .get_account(supply_account_id.to_string())
+        .await?
+        .account;
+
+    assert_eq!(supply_acc.program_owner, Program::token().id());
+    assert_eq!(supply_acc.data.as_ref()[0], 1);
+    assert_eq!(u128::from_le_bytes(supply_acc.data[33..].try_into()?), 37);
+
+    // Create private recipient account
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id_private,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create public recipient account
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id_public,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Mint to public account
+    let subcommand = TokenProgramAgnosticSubcommand::Mint {
+        definition: format_private_account_id(&definition_account_id.to_string()),
+        holder: Some(format_public_account_id(
+            &recipient_account_id_public.to_string(),
+        )),
+        holder_npk: None,
+        holder_ipk: None,
+        amount: 10,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify definition account has updated supply
+    let definition_acc = ctx
+        .wallet()
+        .get_account_private(&definition_account_id)
+        .context("Failed to get definition account")?;
+
+    assert_eq!(
+        u128::from_le_bytes(definition_acc.data[7..23].try_into()?),
+        47 // 37 + 10
+    );
+
+    // Verify public recipient received tokens
+    let recipient_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id_public.to_string())
+        .await?
+        .account;
+
+    assert_eq!(
+        u128::from_le_bytes(recipient_acc.data[33..].try_into()?),
+        10
+    );
+
+    // Mint to private account
+    let subcommand = TokenProgramAgnosticSubcommand::Mint {
+        definition: format_private_account_id(&definition_account_id.to_string()),
+        holder: Some(format_private_account_id(
+            &recipient_account_id_private.to_string(),
+        )),
+        holder_npk: None,
+        holder_ipk: None,
+        amount: 5,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify private recipient commitment
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&recipient_account_id_private)
+        .context("Failed to get recipient commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    // Verify private recipient balance
+    let recipient_acc_private = ctx
+        .wallet()
+        .get_account_private(&recipient_account_id_private)
+        .context("Failed to get private recipient account")?;
+
+    assert_eq!(
+        u128::from_le_bytes(recipient_acc_private.data[33..].try_into()?),
+        5
+    );
+
+    info!("Successfully created token with private definition and minted to both account types");
+
+    Ok(())
+}
+
+#[test]
+async fn create_token_with_private_definition_and_supply() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create token definition account (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create supply account (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create token with both private definition and supply
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_private_account_id(&definition_account_id.to_string()),
+        supply_account_id: format_private_account_id(&supply_account_id.to_string()),
+        name: "A NAME".to_string(),
+        total_supply: 37,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify definition commitment
+    let definition_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&definition_account_id)
+        .context("Failed to get definition commitment")?;
+    assert!(verify_commitment_is_in_state(definition_commitment, ctx.sequencer_client()).await);
+
+    // Verify supply commitment
+    let supply_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&supply_account_id)
+        .context("Failed to get supply commitment")?;
+    assert!(verify_commitment_is_in_state(supply_commitment, ctx.sequencer_client()).await);
+
+    // Verify supply balance
+    let supply_acc = ctx
+        .wallet()
+        .get_account_private(&supply_account_id)
+        .context("Failed to get supply account")?;
+
+    assert_eq!(u128::from_le_bytes(supply_acc.data[33..].try_into()?), 37);
+
+    // Create recipient account
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Transfer tokens
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from: format_private_account_id(&supply_account_id.to_string()),
+        to: Some(format_private_account_id(&recipient_account_id.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 7,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify both commitments updated
+    let supply_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&supply_account_id)
+        .context("Failed to get supply commitment")?;
+    assert!(verify_commitment_is_in_state(supply_commitment, ctx.sequencer_client()).await);
+
+    let recipient_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&recipient_account_id)
+        .context("Failed to get recipient commitment")?;
+    assert!(verify_commitment_is_in_state(recipient_commitment, ctx.sequencer_client()).await);
+
+    // Verify balances
+    let supply_acc = ctx
+        .wallet()
+        .get_account_private(&supply_account_id)
+        .context("Failed to get supply account")?;
+    assert_eq!(u128::from_le_bytes(supply_acc.data[33..].try_into()?), 30);
+
+    let recipient_acc = ctx
+        .wallet()
+        .get_account_private(&recipient_account_id)
+        .context("Failed to get recipient account")?;
+    assert_eq!(u128::from_le_bytes(recipient_acc.data[33..].try_into()?), 7);
+
+    info!("Successfully created and transferred token with both private definition and supply");
+
+    Ok(())
+}
+
+#[test]
+async fn shielded_token_transfer() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create token definition account (public)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create supply account (public)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create recipient account (private) for shielded transfer
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create token
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_public_account_id(&definition_account_id.to_string()),
+        supply_account_id: format_public_account_id(&supply_account_id.to_string()),
+        name: "A NAME".to_string(),
+        total_supply: 37,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Perform shielded transfer: public supply -> private recipient
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from: format_public_account_id(&supply_account_id.to_string()),
+        to: Some(format_private_account_id(&recipient_account_id.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 7,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify supply account balance
+    let supply_acc = ctx
+        .sequencer_client()
+        .get_account(supply_account_id.to_string())
+        .await?
+        .account;
+    assert_eq!(u128::from_le_bytes(supply_acc.data[33..].try_into()?), 30);
+
+    // Verify recipient commitment exists
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&recipient_account_id)
+        .context("Failed to get recipient commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    // Verify recipient balance
+    let recipient_acc = ctx
+        .wallet()
+        .get_account_private(&recipient_account_id)
+        .context("Failed to get recipient account")?;
+    assert_eq!(u128::from_le_bytes(recipient_acc.data[33..].try_into()?), 7);
+
+    info!("Successfully performed shielded token transfer");
+
+    Ok(())
+}
+
+#[test]
+async fn deshielded_token_transfer() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create token definition account (public)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create supply account (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create recipient account (public) for deshielded transfer
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Public { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create token with private supply
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_public_account_id(&definition_account_id.to_string()),
+        supply_account_id: format_private_account_id(&supply_account_id.to_string()),
+        name: "A NAME".to_string(),
+        total_supply: 37,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Perform deshielded transfer: private supply -> public recipient
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from: format_private_account_id(&supply_account_id.to_string()),
+        to: Some(format_public_account_id(&recipient_account_id.to_string())),
+        to_npk: None,
+        to_ipk: None,
+        amount: 7,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Verify supply account commitment exists
+    let new_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&supply_account_id)
+        .context("Failed to get supply commitment")?;
+    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+
+    // Verify supply balance
+    let supply_acc = ctx
+        .wallet()
+        .get_account_private(&supply_account_id)
+        .context("Failed to get supply account")?;
+    assert_eq!(u128::from_le_bytes(supply_acc.data[33..].try_into()?), 30);
+
+    // Verify recipient balance
+    let recipient_acc = ctx
+        .sequencer_client()
+        .get_account(recipient_account_id.to_string())
+        .await?
+        .account;
+    assert_eq!(u128::from_le_bytes(recipient_acc.data[33..].try_into()?), 7);
+
+    info!("Successfully performed deshielded token transfer");
+
+    Ok(())
+}
+
+#[test]
+async fn token_claiming_path_with_private_accounts() -> Result<()> {
+    let mut ctx = TestContext::new().await?;
+
+    // Create token definition account (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: definition_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create supply account (private)
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: supply_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Create token
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id: format_private_account_id(&definition_account_id.to_string()),
+        supply_account_id: format_private_account_id(&supply_account_id.to_string()),
+        name: "A NAME".to_string(),
+        total_supply: 37,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Create new private account for claiming path
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(NewSubcommand::Private { cci: None })),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount {
+        account_id: recipient_account_id,
+    } = result
+    else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+
+    // Get keys for foreign mint (claiming path)
+    let (holder_keys, _) = ctx
+        .wallet()
+        .storage()
+        .user_data
+        .get_private_account(&recipient_account_id)
+        .cloned()
+        .context("Failed to get private account keys")?;
+
+    // Mint using claiming path (foreign account)
+    let subcommand = TokenProgramAgnosticSubcommand::Mint {
+        definition: format_private_account_id(&definition_account_id.to_string()),
+        holder: None,
+        holder_npk: Some(hex::encode(holder_keys.nullifer_public_key.0)),
+        holder_ipk: Some(hex::encode(holder_keys.incoming_viewing_public_key.0)),
+        amount: 9,
+    };
+
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    // Sync to claim the account
+    let command = Command::Account(AccountSubcommand::SyncPrivate {});
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+
+    // Verify commitment exists
+    let recipient_commitment = ctx
+        .wallet()
+        .get_private_account_commitment(&recipient_account_id)
+        .context("Failed to get recipient commitment")?;
+    assert!(verify_commitment_is_in_state(recipient_commitment, ctx.sequencer_client()).await);
+
+    // Verify balance
+    let recipient_acc = ctx
+        .wallet()
+        .get_account_private(&recipient_account_id)
+        .context("Failed to get recipient account")?;
+    assert_eq!(u128::from_le_bytes(recipient_acc.data[33..].try_into()?), 9);
+
+    info!("Successfully minted tokens using claiming path");
+
+    Ok(())
+}

integration_tests/tests/tps.rs

@@ -1,6 +1,9 @@
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
+use anyhow::Result;
+use integration_tests::TestContext;
 use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder;
+use log::info;
 use nssa::{
     Account, AccountId, PrivacyPreservingTransaction, PrivateKey, PublicKey, PublicTransaction,
     privacy_preserving_transaction::{self as pptx, circuit},
@@ -13,6 +16,78 @@ use nssa_core::{
     encryption::IncomingViewingPublicKey,
 };
 use sequencer_core::config::{AccountInitialData, CommitmentsInitialData, SequencerConfig};
+use tokio::test;
+
+// TODO: Make a proper benchmark instead of an ad-hoc test
+#[test]
+pub async fn tps_test() -> Result<()> {
+    let num_transactions = 300 * 5;
+    let target_tps = 12;
+
+    let tps_test = TpsTestManager::new(target_tps, num_transactions);
+    let ctx = TestContext::new_with_sequencer_config(tps_test.generate_sequencer_config()).await?;
+
+    let target_time = tps_test.target_time();
+    info!(
+        "TPS test begin. Target time is {target_time:?} for {num_transactions} transactions ({target_tps} TPS)"
+    );
+
+    let txs = tps_test.build_public_txs();
+    let now = Instant::now();
+
+    let mut tx_hashes = vec![];
+    for (i, tx) in txs.into_iter().enumerate() {
+        let tx_hash = ctx
+            .sequencer_client()
+            .send_tx_public(tx)
+            .await
+            .unwrap()
+            .tx_hash;
+        info!("Sent tx {i}");
+        tx_hashes.push(tx_hash);
+    }
+
+    for (i, tx_hash) in tx_hashes.iter().enumerate() {
+        loop {
+            if now.elapsed().as_millis() > target_time.as_millis() {
+                panic!("TPS test failed by timeout");
+            }
+
+            let tx_obj = ctx
+                .sequencer_client()
+                .get_transaction_by_hash(tx_hash.clone())
+                .await
+                .inspect_err(|err| {
+                    log::warn!(
+                        "Failed to get transaction by hash {tx_hash:#?} with error: {err:#?}"
+                    )
+                });
+
+            if let Ok(tx_obj) = tx_obj
+                && tx_obj.transaction.is_some()
+            {
+                info!("Found tx {i} with hash {tx_hash}");
+                break;
+            }
+        }
+    }
+    let time_elapsed = now.elapsed().as_secs();
+
+    let tx_processed = tx_hashes.len();
+    let actual_tps = tx_processed as u64 / time_elapsed;
+    info!("Processed {tx_processed} transactions in {time_elapsed:?} ({actual_tps} TPS)",);
+
+    assert_eq!(tx_processed, num_transactions);
+
+    assert!(
+        time_elapsed <= target_time.as_secs(),
+        "Elapsed time {time_elapsed:?} exceeded target time {target_time:?}"
+    );
+
+    info!("TPS test finished successfully");
+
+    Ok(())
+}
 
 pub(crate) struct TpsTestManager {
     public_keypairs: Vec<(PrivateKey, AccountId)>,
@@ -32,7 +107,7 @@ impl TpsTestManager {
                 let account_id = AccountId::from(&public_key);
                 (private_key, account_id)
             })
-            .collect::<Vec<_>>();
+            .collect();
         Self {
             public_keypairs,
             target_tps,
@@ -72,7 +147,7 @@ impl TpsTestManager {
     /// Generates a sequencer configuration with initial balance in a number of public accounts.
     /// The transactions generated with the function `build_public_txs` will be valid in a node
     /// started with the config from this method.
-    pub(crate) fn generate_tps_test_config(&self) -> SequencerConfig {
+    pub(crate) fn generate_sequencer_config(&self) -> SequencerConfig {
         // Create public public keypairs
         let initial_public_accounts = self
             .public_keypairs
@@ -110,6 +185,7 @@ impl TpsTestManager {
             initial_accounts: initial_public_accounts,
             initial_commitments: vec![initial_commitment],
             signing_key: [37; 32],
+            bedrock_config: None,
         }
     }
 }
@@ -118,7 +194,7 @@ impl TpsTestManager {
 /// it may take a while to run. In normal execution of the node this transaction will be accepted
 /// only once. Disabling the node's nullifier uniqueness check allows to submit this transaction
 /// multiple times with the purpose of testing the node's processing performance.
-#[allow(unused)]
+#[expect(dead_code, reason = "No idea if we need this, should we remove it?")]
 fn build_privacy_transaction() -> PrivacyPreservingTransaction {
     let program = Program::authenticated_transfer_program();
     let sender_nsk = [1; 32];
@@ -159,16 +235,16 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction {
         ]],
     );
     let (output, proof) = circuit::execute_and_prove(
-        &[sender_pre, recipient_pre],
-        &Program::serialize_instruction(balance_to_move).unwrap(),
-        &[1, 2],
-        &[0xdeadbeef1, 0xdeadbeef2],
-        &[
+        vec![sender_pre, recipient_pre],
+        Program::serialize_instruction(balance_to_move).unwrap(),
+        vec![1, 2],
+        vec![0xdeadbeef1, 0xdeadbeef2],
+        vec![
             (sender_npk.clone(), sender_ss),
             (recipient_npk.clone(), recipient_ss),
         ],
-        &[sender_nsk],
-        &[Some(proof)],
+        vec![sender_nsk],
+        vec![Some(proof)],
         &program.into(),
     )
     .unwrap();

key_protocol/src/key_protocol_core/mod.rs

@@ -165,6 +165,14 @@ impl NSSAUserData {
                 .map(Into::into)
         }
     }
+
+    pub fn account_ids(&self) -> impl Iterator<Item = &nssa::AccountId> {
+        self.default_pub_account_signing_keys
+            .keys()
+            .chain(self.public_key_tree.account_id_map.keys())
+            .chain(self.default_user_private_accounts.keys())
+            .chain(self.private_key_tree.account_id_map.keys())
+    }
 }
 
 impl Default for NSSAUserData {

nssa/Cargo.toml

@@ -24,8 +24,9 @@ risc0-binfmt = "3.0.2"
 
 [dev-dependencies]
 test_program_methods.workspace = true
-hex-literal = "1.0.0"
 env_logger.workspace = true
+hex-literal = "1.0.0"
+test-case = "3.3.1"
 
 [features]
 default = []

nssa/core/src/account.rs

@@ -15,9 +15,8 @@ pub type Nonce = u128;
 
 /// Account to be used both in public and private contexts
 #[derive(
-    Clone, Default, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
+    Debug, Default, Clone, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
 )]
-#[cfg_attr(any(feature = "host", test), derive(Debug))]
 pub struct Account {
     pub program_owner: ProgramId,
     pub balance: u128,
@@ -25,8 +24,7 @@ pub struct Account {
     pub nonce: Nonce,
 }
 
-#[derive(Clone, Eq, PartialEq, Serialize, Deserialize)]
-#[cfg_attr(any(feature = "host", test), derive(Debug))]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
 pub struct AccountWithMetadata {
     pub account: Account,
     pub is_authorized: bool,
@@ -45,6 +43,7 @@ impl AccountWithMetadata {
 }
 
 #[derive(
+    Debug,
     Default,
     Copy,
     Clone,
@@ -56,7 +55,7 @@ impl AccountWithMetadata {
     BorshSerialize,
     BorshDeserialize,
 )]
-#[cfg_attr(any(feature = "host", test), derive(Debug, PartialOrd, Ord))]
+#[cfg_attr(any(feature = "host", test), derive(PartialOrd, Ord))]
 pub struct AccountId {
     value: [u8; 32],
 }

nssa/core/src/account/data.rs

@@ -5,8 +5,7 @@ use serde::{Deserialize, Serialize};
 
 pub const DATA_MAX_LENGTH_IN_BYTES: usize = 100 * 1024; // 100 KiB
 
-#[derive(Default, Clone, PartialEq, Eq, Serialize, BorshSerialize)]
-#[cfg_attr(any(feature = "host", test), derive(Debug))]
+#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, BorshSerialize)]
 pub struct Data(Vec<u8>);
 
 impl Data {

nssa/core/src/nullifier.rs

@@ -4,8 +4,8 @@ use serde::{Deserialize, Serialize};
 
 use crate::{Commitment, account::AccountId};
 
-#[derive(Serialize, Deserialize, PartialEq, Eq)]
-#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, Hash))]
+#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
+#[cfg_attr(any(feature = "host", test), derive(Clone, Hash))]
 pub struct NullifierPublicKey(pub [u8; 32]);
 
 impl From<&NullifierPublicKey> for AccountId {

nssa/core/src/program.rs

@@ -30,6 +30,20 @@ impl PdaSeed {
     }
 }
 
+pub fn compute_authorized_pdas(
+    caller_program_id: Option<ProgramId>,
+    pda_seeds: &[PdaSeed],
+) -> HashSet<AccountId> {
+    caller_program_id
+        .map(|caller_program_id| {
+            pda_seeds
+                .iter()
+                .map(|pda_seed| AccountId::from((&caller_program_id, pda_seed)))
+                .collect()
+        })
+        .unwrap_or_default()
+}
+
 impl From<(&ProgramId, &PdaSeed)> for AccountId {
     fn from(value: (&ProgramId, &PdaSeed)) -> Self {
         use risc0_zkvm::sha::{Impl, Sha256};
@@ -93,6 +107,13 @@ impl AccountPostState {
         }
     }
 
+    /// Creates a post state that requests ownership of the account
+    /// if the account's program owner is the default program ID.
+    pub fn new_claimed_if_default(account: Account) -> Self {
+        let claim = account.program_owner == DEFAULT_PROGRAM_ID;
+        Self { account, claim }
+    }
+
     /// Returns `true` if this post state requests that the account
     /// be claimed (owned) by the executing program.
     pub fn requires_claim(&self) -> bool {
@@ -108,6 +129,11 @@ impl AccountPostState {
     pub fn account_mut(&mut self) -> &mut Account {
         &mut self.account
     }
+
+    /// Consumes the post state and returns the underlying account
+    pub fn into_account(self) -> Account {
+        self.account
+    }
 }
 
 #[derive(Serialize, Deserialize, Clone)]

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -1,11 +1,11 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, VecDeque};
 
 use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::{
     MembershipProof, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput,
     PrivacyPreservingCircuitOutput, SharedSecretKey,
     account::AccountWithMetadata,
-    program::{InstructionData, ProgramId, ProgramOutput},
+    program::{ChainedCall, InstructionData, ProgramId, ProgramOutput},
 };
 use risc0_zkvm::{ExecutorEnv, InnerReceipt, Receipt, default_prover};
 
@@ -43,27 +43,44 @@ impl From<Program> for ProgramWithDependencies {
 }
 
 /// Generates a proof of the execution of a NSSA program inside the privacy preserving execution
-/// circuit
+/// circuit.
 #[expect(clippy::too_many_arguments, reason = "TODO: fix later")]
 pub fn execute_and_prove(
-    pre_states: &[AccountWithMetadata],
-    instruction_data: &InstructionData,
-    visibility_mask: &[u8],
-    private_account_nonces: &[u128],
-    private_account_keys: &[(NullifierPublicKey, SharedSecretKey)],
-    private_account_nsks: &[NullifierSecretKey],
-    private_account_membership_proofs: &[Option<MembershipProof>],
+    pre_states: Vec<AccountWithMetadata>,
+    instruction_data: InstructionData,
+    visibility_mask: Vec<u8>,
+    private_account_nonces: Vec<u128>,
+    private_account_keys: Vec<(NullifierPublicKey, SharedSecretKey)>,
+    private_account_nsks: Vec<NullifierSecretKey>,
+    private_account_membership_proofs: Vec<Option<MembershipProof>>,
     program_with_dependencies: &ProgramWithDependencies,
 ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> {
-    let mut program = &program_with_dependencies.program;
-    let dependencies = &program_with_dependencies.dependencies;
-    let mut instruction_data = instruction_data.clone();
-    let mut pre_states = pre_states.to_vec();
+    let ProgramWithDependencies {
+        program,
+        dependencies,
+    } = program_with_dependencies;
     let mut env_builder = ExecutorEnv::builder();
     let mut program_outputs = Vec::new();
 
-    for _i in 0..MAX_NUMBER_CHAINED_CALLS {
-        let inner_receipt = execute_and_prove_program(program, &pre_states, &instruction_data)?;
+    let initial_call = ChainedCall {
+        program_id: program.id(),
+        instruction_data: instruction_data.clone(),
+        pre_states,
+        pda_seeds: vec![],
+    };
+
+    let mut chained_calls = VecDeque::from_iter([(initial_call, program)]);
+    let mut chain_calls_counter = 0;
+    while let Some((chained_call, program)) = chained_calls.pop_front() {
+        if chain_calls_counter >= MAX_NUMBER_CHAINED_CALLS {
+            return Err(NssaError::MaxChainedCallsDepthExceeded);
+        }
+
+        let inner_receipt = execute_and_prove_program(
+            program,
+            &chained_call.pre_states,
+            &chained_call.instruction_data,
+        )?;
 
         let program_output: ProgramOutput = inner_receipt
             .journal
@@ -76,39 +93,23 @@ pub fn execute_and_prove(
         // Prove circuit.
         env_builder.add_assumption(inner_receipt);
 
-        // TODO: Remove when multi-chain calls are supported in the circuit
-        assert!(program_output.chained_calls.len() <= 1);
-        // TODO: Modify when multi-chain calls are supported in the circuit
-        if let Some(next_call) = program_output.chained_calls.first() {
-            program = dependencies
-                .get(&next_call.program_id)
+        for new_call in program_output.chained_calls.into_iter().rev() {
+            let next_program = dependencies
+                .get(&new_call.program_id)
                 .ok_or(NssaError::InvalidProgramBehavior)?;
-            instruction_data = next_call.instruction_data.clone();
-            // Build post states with metadata for next call
-            let mut post_states_with_metadata = Vec::new();
-            for (pre, post) in program_output
-                .pre_states
-                .iter()
-                .zip(program_output.post_states)
-            {
-                let mut post_with_metadata = pre.clone();
-                post_with_metadata.account = post.account().clone();
-                post_states_with_metadata.push(post_with_metadata);
-            }
-
-            pre_states = next_call.pre_states.clone();
-        } else {
-            break;
+            chained_calls.push_front((new_call, next_program));
         }
+
+        chain_calls_counter += 1;
     }
 
     let circuit_input = PrivacyPreservingCircuitInput {
         program_outputs,
-        visibility_mask: visibility_mask.to_vec(),
-        private_account_nonces: private_account_nonces.to_vec(),
-        private_account_keys: private_account_keys.to_vec(),
-        private_account_nsks: private_account_nsks.to_vec(),
-        private_account_membership_proofs: private_account_membership_proofs.to_vec(),
+        visibility_mask,
+        private_account_nonces,
+        private_account_keys,
+        private_account_nsks,
+        private_account_membership_proofs,
         program_id: program_with_dependencies.program.id(),
     };
 
@@ -215,13 +216,13 @@ mod tests {
         let shared_secret = SharedSecretKey::new(&esk, &recipient_keys.ivk());
 
         let (output, proof) = execute_and_prove(
-            &[sender, recipient],
-            &Program::serialize_instruction(balance_to_move).unwrap(),
-            &[0, 2],
-            &[0xdeadbeef],
-            &[(recipient_keys.npk(), shared_secret)],
-            &[],
-            &[None],
+            vec![sender, recipient],
+            Program::serialize_instruction(balance_to_move).unwrap(),
+            vec![0, 2],
+            vec![0xdeadbeef],
+            vec![(recipient_keys.npk(), shared_secret)],
+            vec![],
+            vec![None],
             &Program::authenticated_transfer_program().into(),
         )
         .unwrap();
@@ -311,16 +312,16 @@ mod tests {
         let shared_secret_2 = SharedSecretKey::new(&esk_2, &recipient_keys.ivk());
 
         let (output, proof) = execute_and_prove(
-            &[sender_pre.clone(), recipient],
-            &Program::serialize_instruction(balance_to_move).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![sender_pre.clone(), recipient],
+            Program::serialize_instruction(balance_to_move).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (sender_keys.npk(), shared_secret_1),
                 (recipient_keys.npk(), shared_secret_2),
             ],
-            &[sender_keys.nsk],
-            &[commitment_set.get_proof_for(&commitment_sender), None],
+            vec![sender_keys.nsk],
+            vec![commitment_set.get_proof_for(&commitment_sender), None],
             &program.into(),
         )
         .unwrap();

nssa/src/program.rs

@@ -226,6 +226,15 @@ mod tests {
             }
         }
 
+        pub fn changer_claimer() -> Self {
+            use test_program_methods::{CHANGER_CLAIMER_ELF, CHANGER_CLAIMER_ID};
+
+            Program {
+                id: CHANGER_CLAIMER_ID,
+                elf: CHANGER_CLAIMER_ELF.to_vec(),
+            }
+        }
+
         pub fn noop() -> Self {
             use test_program_methods::{NOOP_ELF, NOOP_ID};
 
@@ -235,6 +244,17 @@ mod tests {
             }
         }
 
+        pub fn malicious_authorization_changer() -> Self {
+            use test_program_methods::{
+                MALICIOUS_AUTHORIZATION_CHANGER_ELF, MALICIOUS_AUTHORIZATION_CHANGER_ID,
+            };
+
+            Program {
+                id: MALICIOUS_AUTHORIZATION_CHANGER_ID,
+                elf: MALICIOUS_AUTHORIZATION_CHANGER_ELF.to_vec(),
+            }
+        }
+
         pub fn modified_transfer_program() -> Self {
             use test_program_methods::MODIFIED_TRANSFER_ELF;
             // This unwrap won't panic since the `MODIFIED_TRANSFER_ELF` comes from risc0 build of

nssa/src/public_transaction/transaction.rs

@@ -4,7 +4,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use log::debug;
 use nssa_core::{
     account::{Account, AccountId, AccountWithMetadata},
-    program::{ChainedCall, DEFAULT_PROGRAM_ID, PdaSeed, ProgramId, validate_execution},
+    program::{ChainedCall, DEFAULT_PROGRAM_ID, validate_execution},
 };
 use sha2::{Digest, digest::FixedOutput};
 
@@ -119,7 +119,7 @@ impl PublicTransaction {
                 return Err(NssaError::MaxChainedCallsDepthExceeded);
             }
 
-            // Check the `program_id` corresponds to a deployed program
+            // Check that the `program_id` corresponds to a deployed program
             let Some(program) = state.programs().get(&chained_call.program_id) else {
                 return Err(NssaError::InvalidInput("Unknown program".into()));
             };
@@ -135,12 +135,14 @@ impl PublicTransaction {
                 chained_call.program_id, program_output
             );
 
-            let authorized_pdas =
-                self.compute_authorized_pdas(&caller_program_id, &chained_call.pda_seeds);
+            let authorized_pdas = nssa_core::program::compute_authorized_pdas(
+                caller_program_id,
+                &chained_call.pda_seeds,
+            );
 
             for pre in &program_output.pre_states {
                 let account_id = pre.account_id;
-                // Check that the program output pre_states coinicide with the values in the public
+                // Check that the program output pre_states coincide with the values in the public
                 // state or with any modifications to those values during the chain of calls.
                 let expected_pre = state_diff
                     .get(&account_id)
@@ -198,22 +200,23 @@ impl PublicTransaction {
             chain_calls_counter += 1;
         }
 
-        Ok(state_diff)
-    }
-
-    fn compute_authorized_pdas(
-        &self,
-        caller_program_id: &Option<ProgramId>,
-        pda_seeds: &[PdaSeed],
-    ) -> HashSet<AccountId> {
-        if let Some(caller_program_id) = caller_program_id {
-            pda_seeds
-                .iter()
-                .map(|pda_seed| AccountId::from((caller_program_id, pda_seed)))
-                .collect()
-        } else {
-            HashSet::new()
+        // Check that all modified uninitialized accounts where claimed
+        for post in state_diff.iter().filter_map(|(account_id, post)| {
+            let pre = state.get_account_by_id(account_id);
+            if pre.program_owner != DEFAULT_PROGRAM_ID {
+                return None;
+            }
+            if pre == *post {
+                return None;
+            }
+            Some(post)
+        }) {
+            if post.program_owner == DEFAULT_PROGRAM_ID {
+                return Err(NssaError::InvalidProgramBehavior);
+            }
         }
+
+        Ok(state_diff)
     }
 }
 

nssa/src/state.rs

@@ -504,6 +504,7 @@ pub mod tests {
             self.insert_program(Program::chain_caller());
             self.insert_program(Program::amm());
             self.insert_program(Program::claimer());
+            self.insert_program(Program::changer_claimer());
             self
         }
 
@@ -865,13 +866,13 @@ pub mod tests {
         let epk = EphemeralPublicKey::from_scalar(esk);
 
         let (output, proof) = circuit::execute_and_prove(
-            &[sender, recipient],
-            &Program::serialize_instruction(balance_to_move).unwrap(),
-            &[0, 2],
-            &[0xdeadbeef],
-            &[(recipient_keys.npk(), shared_secret)],
-            &[],
-            &[None],
+            vec![sender, recipient],
+            Program::serialize_instruction(balance_to_move).unwrap(),
+            vec![0, 2],
+            vec![0xdeadbeef],
+            vec![(recipient_keys.npk(), shared_secret)],
+            vec![],
+            vec![None],
             &Program::authenticated_transfer_program().into(),
         )
         .unwrap();
@@ -912,16 +913,16 @@ pub mod tests {
         let epk_2 = EphemeralPublicKey::from_scalar(esk_2);
 
         let (output, proof) = circuit::execute_and_prove(
-            &[sender_pre, recipient_pre],
-            &Program::serialize_instruction(balance_to_move).unwrap(),
-            &[1, 2],
-            &new_nonces,
-            &[
+            vec![sender_pre, recipient_pre],
+            Program::serialize_instruction(balance_to_move).unwrap(),
+            vec![1, 2],
+            new_nonces.to_vec(),
+            vec![
                 (sender_keys.npk(), shared_secret_1),
                 (recipient_keys.npk(), shared_secret_2),
             ],
-            &[sender_keys.nsk],
-            &[state.get_proof_for_commitment(&sender_commitment), None],
+            vec![sender_keys.nsk],
+            vec![state.get_proof_for_commitment(&sender_commitment), None],
             &program.into(),
         )
         .unwrap();
@@ -965,13 +966,13 @@ pub mod tests {
         let epk = EphemeralPublicKey::from_scalar(esk);
 
         let (output, proof) = circuit::execute_and_prove(
-            &[sender_pre, recipient_pre],
-            &Program::serialize_instruction(balance_to_move).unwrap(),
-            &[1, 0],
-            &[new_nonce],
-            &[(sender_keys.npk(), shared_secret)],
-            &[sender_keys.nsk],
-            &[state.get_proof_for_commitment(&sender_commitment)],
+            vec![sender_pre, recipient_pre],
+            Program::serialize_instruction(balance_to_move).unwrap(),
+            vec![1, 0],
+            vec![new_nonce],
+            vec![(sender_keys.npk(), shared_secret)],
+            vec![sender_keys.nsk],
+            vec![state.get_proof_for_commitment(&sender_commitment)],
             &program.into(),
         )
         .unwrap();
@@ -1179,13 +1180,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1206,13 +1207,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1233,13 +1234,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account],
-            &Program::serialize_instruction(()).unwrap(),
-            &[0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account],
+            Program::serialize_instruction(()).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1260,13 +1261,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account],
-            &Program::serialize_instruction(vec![0]).unwrap(),
-            &[0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account],
+            Program::serialize_instruction(vec![0]).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1289,13 +1290,13 @@ pub mod tests {
         let large_data: Vec<u8> = vec![0; nssa_core::account::data::DATA_MAX_LENGTH_IN_BYTES + 1];
 
         let result = execute_and_prove(
-            &[public_account],
-            &Program::serialize_instruction(large_data).unwrap(),
-            &[0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account],
+            Program::serialize_instruction(large_data).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.to_owned().into(),
         );
 
@@ -1316,13 +1317,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account],
-            &Program::serialize_instruction(()).unwrap(),
-            &[0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account],
+            Program::serialize_instruction(()).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1352,13 +1353,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account_1, public_account_2],
-            &Program::serialize_instruction(()).unwrap(),
-            &[0, 0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account_1, public_account_2],
+            Program::serialize_instruction(()).unwrap(),
+            vec![0, 0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1379,13 +1380,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account],
-            &Program::serialize_instruction(()).unwrap(),
-            &[0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account],
+            Program::serialize_instruction(()).unwrap(),
+            vec![0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1415,13 +1416,13 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[public_account_1, public_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[0, 0],
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account_1, public_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![0, 0],
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1453,13 +1454,13 @@ pub mod tests {
         // Setting only one visibility mask for a circuit execution with two pre_state accounts.
         let visibility_mask = [0];
         let result = execute_and_prove(
-            &[public_account_1, public_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &visibility_mask,
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account_1, public_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            visibility_mask.to_vec(),
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1486,11 +1487,11 @@ pub mod tests {
         // Setting only one nonce for an execution with two private accounts.
         let private_account_nonces = [0xdeadbeef1];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &private_account_nonces,
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            private_account_nonces.to_vec(),
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1500,8 +1501,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -1530,13 +1531,13 @@ pub mod tests {
             SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
         )];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &private_account_keys,
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            private_account_keys.to_vec(),
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -1563,11 +1564,11 @@ pub mod tests {
         // Setting no second commitment proof.
         let private_account_membership_proofs = [Some((0, vec![]))];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1577,8 +1578,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &private_account_membership_proofs,
+            vec![sender_keys.nsk],
+            private_account_membership_proofs.to_vec(),
             &program.into(),
         );
 
@@ -1605,11 +1606,11 @@ pub mod tests {
         // Setting no auth key for an execution with one non default private accounts.
         let private_account_nsks = [];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1619,8 +1620,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &private_account_nsks,
-            &[],
+            private_account_nsks.to_vec(),
+            vec![],
             &program.into(),
         );
 
@@ -1663,13 +1664,13 @@ pub mod tests {
         let private_account_nsks = [recipient_keys.nsk];
         let private_account_membership_proofs = [Some((0, vec![]))];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &private_account_keys,
-            &private_account_nsks,
-            &private_account_membership_proofs,
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            private_account_keys.to_vec(),
+            private_account_nsks.to_vec(),
+            private_account_membership_proofs.to_vec(),
             &program.into(),
         );
 
@@ -1701,11 +1702,11 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1715,8 +1716,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -1749,11 +1750,11 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1763,8 +1764,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -1796,11 +1797,11 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1810,8 +1811,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -1843,11 +1844,11 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1857,8 +1858,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -1888,11 +1889,11 @@ pub mod tests {
         );
 
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1902,8 +1903,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -1927,13 +1928,13 @@ pub mod tests {
 
         let visibility_mask = [0, 3];
         let result = execute_and_prove(
-            &[public_account_1, public_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &visibility_mask,
-            &[],
-            &[],
-            &[],
-            &[],
+            vec![public_account_1, public_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            visibility_mask.to_vec(),
+            vec![],
+            vec![],
+            vec![],
+            vec![],
             &program.into(),
         );
 
@@ -1961,11 +1962,11 @@ pub mod tests {
         // accounts.
         let private_account_nonces = [0xdeadbeef1, 0xdeadbeef2, 0xdeadbeef3];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &private_account_nonces,
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            private_account_nonces.to_vec(),
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -1975,8 +1976,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -2017,13 +2018,13 @@ pub mod tests {
             ),
         ];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &[1, 2],
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &private_account_keys,
-            &[sender_keys.nsk],
-            &[Some((0, vec![]))],
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            vec![1, 2],
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            private_account_keys.to_vec(),
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
             &program.into(),
         );
 
@@ -2053,11 +2054,11 @@ pub mod tests {
         let private_account_nsks = [sender_keys.nsk, recipient_keys.nsk];
         let private_account_membership_proofs = [Some((0, vec![])), Some((1, vec![]))];
         let result = execute_and_prove(
-            &[private_account_1, private_account_2],
-            &Program::serialize_instruction(10u128).unwrap(),
-            &visibility_mask,
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1, private_account_2],
+            Program::serialize_instruction(10u128).unwrap(),
+            visibility_mask.to_vec(),
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (
                     sender_keys.npk(),
                     SharedSecretKey::new(&[55; 32], &sender_keys.ivk()),
@@ -2067,8 +2068,8 @@ pub mod tests {
                     SharedSecretKey::new(&[56; 32], &recipient_keys.ivk()),
                 ),
             ],
-            &private_account_nsks,
-            &private_account_membership_proofs,
+            private_account_nsks.to_vec(),
+            private_account_membership_proofs.to_vec(),
             &program.into(),
         );
 
@@ -2149,16 +2150,16 @@ pub mod tests {
         let private_account_membership_proofs = [Some((1, vec![])), Some((1, vec![]))];
         let shared_secret = SharedSecretKey::new(&[55; 32], &sender_keys.ivk());
         let result = execute_and_prove(
-            &[private_account_1.clone(), private_account_1],
-            &Program::serialize_instruction(100u128).unwrap(),
-            &visibility_mask,
-            &[0xdeadbeef1, 0xdeadbeef2],
-            &[
+            vec![private_account_1.clone(), private_account_1],
+            Program::serialize_instruction(100u128).unwrap(),
+            visibility_mask.to_vec(),
+            vec![0xdeadbeef1, 0xdeadbeef2],
+            vec![
                 (sender_keys.npk(), shared_secret),
                 (sender_keys.npk(), shared_secret),
             ],
-            &private_account_nsks,
-            &private_account_membership_proofs,
+            private_account_nsks.to_vec(),
+            private_account_membership_proofs.to_vec(),
             &program.into(),
         );
 
@@ -3941,8 +3942,9 @@ pub mod tests {
         assert_eq!(to_post, expected_to_post);
     }
 
-    #[test]
-    fn test_private_chained_call() {
+    #[test_case::test_case(1; "single call")]
+    #[test_case::test_case(2; "two calls")]
+    fn test_private_chained_call(number_of_calls: u32) {
         // Arrange
         let chain_caller = Program::chain_caller();
         let auth_transfers = Program::authenticated_transfer_program();
@@ -3978,7 +3980,7 @@ pub mod tests {
         let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
             amount,
             Program::authenticated_transfer_program().id(),
-            1,
+            number_of_calls,
             None,
         );
 
@@ -3999,14 +4001,14 @@ pub mod tests {
         let to_new_nonce = 0xdeadbeef2;
 
         let from_expected_post = Account {
-            balance: initial_balance - amount,
+            balance: initial_balance - number_of_calls as u128 * amount,
             nonce: from_new_nonce,
             ..from_account.account.clone()
         };
         let from_expected_commitment = Commitment::new(&from_keys.npk(), &from_expected_post);
 
         let to_expected_post = Account {
-            balance: amount,
+            balance: number_of_calls as u128 * amount,
             nonce: to_new_nonce,
             ..to_account.account.clone()
         };
@@ -4014,13 +4016,13 @@ pub mod tests {
 
         // Act
         let (output, proof) = execute_and_prove(
-            &[to_account, from_account],
-            &Program::serialize_instruction(instruction).unwrap(),
-            &[1, 1],
-            &[from_new_nonce, to_new_nonce],
-            &[(from_keys.npk(), to_ss), (to_keys.npk(), from_ss)],
-            &[from_keys.nsk, to_keys.nsk],
-            &[
+            vec![to_account, from_account],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![1, 1],
+            vec![from_new_nonce, to_new_nonce],
+            vec![(from_keys.npk(), to_ss), (to_keys.npk(), from_ss)],
+            vec![from_keys.nsk, to_keys.nsk],
+            vec![
                 state.get_proof_for_commitment(&from_commitment),
                 state.get_proof_for_commitment(&to_commitment),
             ],
@@ -4255,13 +4257,13 @@ pub mod tests {
 
         // Execute and prove the circuit with the authorized account but no commitment proof
         let (output, proof) = execute_and_prove(
-            std::slice::from_ref(&authorized_account),
-            &Program::serialize_instruction(balance).unwrap(),
-            &[1],
-            &[nonce],
-            &[(private_keys.npk(), shared_secret)],
-            &[private_keys.nsk],
-            &[None],
+            vec![authorized_account],
+            Program::serialize_instruction(balance).unwrap(),
+            vec![1],
+            vec![nonce],
+            vec![(private_keys.npk(), shared_secret)],
+            vec![private_keys.nsk],
+            vec![None],
             &program.into(),
         )
         .unwrap();
@@ -4308,13 +4310,13 @@ pub mod tests {
 
         // Step 2: Execute claimer program to claim the account with authentication
         let (output, proof) = execute_and_prove(
-            std::slice::from_ref(&authorized_account),
-            &Program::serialize_instruction(balance).unwrap(),
-            &[1],
-            &[nonce],
-            &[(private_keys.npk(), shared_secret)],
-            &[private_keys.nsk],
-            &[None],
+            vec![authorized_account.clone()],
+            Program::serialize_instruction(balance).unwrap(),
+            vec![1],
+            vec![nonce],
+            vec![(private_keys.npk(), shared_secret)],
+            vec![private_keys.nsk],
+            vec![None],
             &claimer_program.into(),
         )
         .unwrap();
@@ -4356,16 +4358,174 @@ pub mod tests {
 
         // Step 3: Try to execute noop program with authentication but without initialization
         let res = execute_and_prove(
-            std::slice::from_ref(&account_metadata),
-            &Program::serialize_instruction(()).unwrap(),
-            &[1],
-            &[nonce2],
-            &[(private_keys.npk(), shared_secret2)],
-            &[private_keys.nsk],
-            &[None],
+            vec![account_metadata],
+            Program::serialize_instruction(()).unwrap(),
+            vec![1],
+            vec![nonce2],
+            vec![(private_keys.npk(), shared_secret2)],
+            vec![private_keys.nsk],
+            vec![None],
             &noop_program.into(),
         );
 
         assert!(matches!(res, Err(NssaError::CircuitProvingError(_))));
     }
+
+    #[test]
+    fn test_public_changer_claimer_no_data_change_no_claim_succeeds() {
+        let initial_data = [];
+        let mut state =
+            V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
+        let account_id = AccountId::new([1; 32]);
+        let program_id = Program::changer_claimer().id();
+        // Don't change data (None) and don't claim (false)
+        let instruction: (Option<Vec<u8>>, bool) = (None, false);
+
+        let message =
+            public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
+                .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        let tx = PublicTransaction::new(message, witness_set);
+
+        let result = state.transition_from_public_transaction(&tx);
+
+        // Should succeed - no changes made, no claim needed
+        assert!(result.is_ok());
+        // Account should remain default/unclaimed
+        assert_eq!(state.get_account_by_id(&account_id), Account::default());
+    }
+
+    #[test]
+    fn test_public_changer_claimer_data_change_no_claim_fails() {
+        let initial_data = [];
+        let mut state =
+            V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
+        let account_id = AccountId::new([1; 32]);
+        let program_id = Program::changer_claimer().id();
+        // Change data but don't claim (false) - should fail
+        let new_data = vec![1, 2, 3, 4, 5];
+        let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);
+
+        let message =
+            public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
+                .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        let tx = PublicTransaction::new(message, witness_set);
+
+        let result = state.transition_from_public_transaction(&tx);
+
+        // Should fail - cannot modify data without claiming the account
+        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+    }
+
+    #[test]
+    fn test_private_changer_claimer_no_data_change_no_claim_succeeds() {
+        let program = Program::changer_claimer();
+        let sender_keys = test_private_account_keys_1();
+        let private_account =
+            AccountWithMetadata::new(Account::default(), true, &sender_keys.npk());
+        // Don't change data (None) and don't claim (false)
+        let instruction: (Option<Vec<u8>>, bool) = (None, false);
+
+        let result = execute_and_prove(
+            vec![private_account],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![1],
+            vec![2],
+            vec![(
+                sender_keys.npk(),
+                SharedSecretKey::new(&[3; 32], &sender_keys.ivk()),
+            )],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
+            &program.into(),
+        );
+
+        // Should succeed - no changes made, no claim needed
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_private_changer_claimer_data_change_no_claim_fails() {
+        let program = Program::changer_claimer();
+        let sender_keys = test_private_account_keys_1();
+        let private_account =
+            AccountWithMetadata::new(Account::default(), true, &sender_keys.npk());
+        // Change data but don't claim (false) - should fail
+        let new_data = vec![1, 2, 3, 4, 5];
+        let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);
+
+        let result = execute_and_prove(
+            vec![private_account],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![1],
+            vec![2],
+            vec![(
+                sender_keys.npk(),
+                SharedSecretKey::new(&[3; 32], &sender_keys.ivk()),
+            )],
+            vec![sender_keys.nsk],
+            vec![Some((0, vec![]))],
+            &program.into(),
+        );
+
+        // Should fail - cannot modify data without claiming the account
+        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
+    }
+
+    #[test]
+    fn test_malicious_authorization_changer_should_fail_in_privacy_preserving_circuit() {
+        // Arrange
+        let malicious_program = Program::malicious_authorization_changer();
+        let auth_transfers = Program::authenticated_transfer_program();
+        let sender_keys = test_public_account_keys_1();
+        let recipient_keys = test_private_account_keys_1();
+
+        let sender_account = AccountWithMetadata::new(
+            Account {
+                program_owner: auth_transfers.id(),
+                balance: 100,
+                ..Default::default()
+            },
+            false,
+            sender_keys.account_id(),
+        );
+        let recipient_account =
+            AccountWithMetadata::new(Account::default(), true, &recipient_keys.npk());
+
+        let recipient_commitment =
+            Commitment::new(&recipient_keys.npk(), &recipient_account.account);
+        let state = V02State::new_with_genesis_accounts(
+            &[(sender_account.account_id, sender_account.account.balance)],
+            std::slice::from_ref(&recipient_commitment),
+        )
+        .with_test_programs();
+
+        let balance_to_transfer = 10u128;
+        let instruction = (balance_to_transfer, auth_transfers.id());
+
+        let recipient_esk = [3; 32];
+        let recipient = SharedSecretKey::new(&recipient_esk, &recipient_keys.ivk());
+
+        let mut dependencies = HashMap::new();
+        dependencies.insert(auth_transfers.id(), auth_transfers);
+        let program_with_deps = ProgramWithDependencies::new(malicious_program, dependencies);
+
+        let recipient_new_nonce = 0xdeadbeef1;
+
+        // Act - execute the malicious program - this should fail during proving
+        let result = execute_and_prove(
+            vec![sender_account, recipient_account],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![0, 1],
+            vec![recipient_new_nonce],
+            vec![(recipient_keys.npk(), recipient)],
+            vec![recipient_keys.nsk],
+            vec![state.get_proof_for_commitment(&recipient_commitment)],
+            &program_with_deps,
+        );
+
+        // Assert - should fail because the malicious program tries to manipulate is_authorized
+        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
+    }
 }

program_methods/guest/src/bin/pinata.rs

@@ -77,7 +77,7 @@ fn main() {
         instruction_words,
         vec![pinata, winner],
         vec![
-            AccountPostState::new(pinata_post),
+            AccountPostState::new_claimed_if_default(pinata_post),
             AccountPostState::new(winner_post),
         ],
     );

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -1,12 +1,18 @@
-use std::collections::HashMap;
+use std::{
+    collections::{HashMap, HashSet, VecDeque, hash_map::Entry},
+    convert::Infallible,
+};
 
 use nssa_core::{
-    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Nullifier,
-    NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
-    account::{Account, AccountId, AccountWithMetadata},
+    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, MembershipProof,
+    Nullifier, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput,
+    PrivacyPreservingCircuitOutput, SharedSecretKey,
+    account::{Account, AccountId, AccountWithMetadata, Nonce},
     compute_digest_for_path,
-    encryption::Ciphertext,
-    program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, validate_execution},
+    program::{
+        AccountPostState, ChainedCall, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, ProgramId,
+        ProgramOutput, validate_execution,
+    },
 };
 use risc0_zkvm::{guest::env, serde::to_vec};
 
@@ -18,118 +24,224 @@ fn main() {
         private_account_keys,
         private_account_nsks,
         private_account_membership_proofs,
-        mut program_id,
+        program_id,
     } = env::read();
 
-    let mut pre_states: Vec<AccountWithMetadata> = Vec::new();
-    let mut state_diff: HashMap<AccountId, Account> = HashMap::new();
+    let execution_state = ExecutionState::derive_from_outputs(program_id, program_outputs);
 
-    let num_calls = program_outputs.len();
-    if num_calls > MAX_NUMBER_CHAINED_CALLS {
-        panic!("Max chained calls depth is exceeded");
-    }
+    let output = compute_circuit_output(
+        execution_state,
+        &visibility_mask,
+        &private_account_nonces,
+        &private_account_keys,
+        &private_account_nsks,
+        &private_account_membership_proofs,
+    );
 
-    let Some(last_program_call) = program_outputs.last() else {
-        panic!("Program outputs is empty")
-    };
+    env::commit(&output);
+}
 
-    if !last_program_call.chained_calls.is_empty() {
-        panic!("Call stack is incomplete");
-    }
+/// State of the involved accounts before and after program execution.
+struct ExecutionState {
+    pre_states: Vec<AccountWithMetadata>,
+    post_states: HashMap<AccountId, Account>,
+}
 
-    for window in program_outputs.windows(2) {
-        let caller = &window[0];
-        let callee = &window[1];
-
-        if caller.chained_calls.len() > 1 {
-            panic!("Privacy Multi-chained calls are not supported yet");
-        }
-
-        // TODO: Modify when multi-chain calls are supported in the circuit
-        let Some(caller_chained_call) = &caller.chained_calls.first() else {
-            panic!("Expected chained call");
+impl ExecutionState {
+    /// Validate program outputs and derive the overall execution state.
+    pub fn derive_from_outputs(program_id: ProgramId, program_outputs: Vec<ProgramOutput>) -> Self {
+        let Some(first_output) = program_outputs.first() else {
+            panic!("No program outputs provided");
         };
 
-        // Check that instruction data in caller is the instruction data in callee
-        if caller_chained_call.instruction_data != callee.instruction_data {
-            panic!("Invalid instruction data");
-        }
-
-        // Check that account pre_states in caller are the ones in calle
-        if caller_chained_call.pre_states != callee.pre_states {
-            panic!("Invalid pre states");
-        }
-    }
-
-    for (i, program_output) in program_outputs.iter().enumerate() {
-        let mut program_output = program_output.clone();
-
-        // Check that `program_output` is consistent with the execution of the corresponding
-        // program.
-        let program_output_words =
-            &to_vec(&program_output).expect("program_output must be serializable");
-        env::verify(program_id, program_output_words)
-            .expect("program output must match the program's execution");
-
-        // Check that the program is well behaved.
-        // See the # Programs section for the definition of the `validate_execution` method.
-        if !validate_execution(
-            &program_output.pre_states,
-            &program_output.post_states,
+        let initial_call = ChainedCall {
             program_id,
-        ) {
-            panic!("Bad behaved program");
-        }
+            instruction_data: first_output.instruction_data.clone(),
+            pre_states: first_output.pre_states.clone(),
+            pda_seeds: Vec::new(),
+        };
+        let mut chained_calls = VecDeque::from_iter([(initial_call, None)]);
 
-        // The invoked program claims the accounts with default program id.
-        for post in program_output
-            .post_states
-            .iter_mut()
-            .filter(|post| post.requires_claim())
-        {
-            // The invoked program can only claim accounts with default program id.
-            if post.account().program_owner == DEFAULT_PROGRAM_ID {
-                post.account_mut().program_owner = program_id;
-            } else {
-                panic!("Cannot claim an initialized account")
+        let mut execution_state = ExecutionState {
+            pre_states: Vec::new(),
+            post_states: HashMap::new(),
+        };
+
+        let mut program_outputs_iter = program_outputs.into_iter();
+        let mut chain_calls_counter = 0;
+
+        while let Some((chained_call, caller_program_id)) = chained_calls.pop_front() {
+            assert!(
+                chain_calls_counter <= MAX_NUMBER_CHAINED_CALLS,
+                "Max chained calls depth is exceeded"
+            );
+
+            let Some(program_output) = program_outputs_iter.next() else {
+                panic!("Insufficient program outputs for chained calls");
+            };
+
+            // Check that instruction data in chained call is the instruction data in program output
+            assert_eq!(
+                chained_call.instruction_data, program_output.instruction_data,
+                "Mismatched instruction data between chained call and program output"
+            );
+
+            // Check that `program_output` is consistent with the execution of the corresponding
+            // program.
+            let program_output_words =
+                &to_vec(&program_output).expect("program_output must be serializable");
+            env::verify(chained_call.program_id, program_output_words).unwrap_or_else(
+                |_: Infallible| unreachable!("Infallible error is never constructed"),
+            );
+
+            // Check that the program is well behaved.
+            // See the # Programs section for the definition of the `validate_execution` method.
+            let execution_valid = validate_execution(
+                &program_output.pre_states,
+                &program_output.post_states,
+                chained_call.program_id,
+            );
+            assert!(execution_valid, "Bad behaved program");
+
+            for next_call in program_output.chained_calls.iter().rev() {
+                chained_calls.push_front((next_call.clone(), Some(chained_call.program_id)));
             }
+
+            let authorized_pdas = nssa_core::program::compute_authorized_pdas(
+                caller_program_id,
+                &chained_call.pda_seeds,
+            );
+            execution_state.validate_and_sync_states(
+                chained_call.program_id,
+                authorized_pdas,
+                program_output.pre_states,
+                program_output.post_states,
+            );
+            chain_calls_counter += 1;
         }
 
-        for (pre, post) in program_output
+        assert!(
+            program_outputs_iter.next().is_none(),
+            "Inner call without a chained call found",
+        );
+
+        // Check that all modified uninitialized accounts were claimed
+        for (account_id, post) in execution_state
             .pre_states
             .iter()
-            .zip(&program_output.post_states)
+            .filter(|a| a.account.program_owner == DEFAULT_PROGRAM_ID)
+            .map(|a| {
+                let post = execution_state
+                    .post_states
+                    .get(&a.account_id)
+                    .expect("Post state must exist for pre state");
+                (a, post)
+            })
+            .filter(|(pre_default, post)| pre_default.account != **post)
+            .map(|(pre, post)| (pre.account_id, post))
         {
-            if let Some(account_pre) = state_diff.get(&pre.account_id) {
-                if account_pre != &pre.account {
-                    panic!("Invalid input");
-                }
-            } else {
-                pre_states.push(pre.clone());
-            }
-            state_diff.insert(pre.account_id, post.account().clone());
+            assert_ne!(
+                post.program_owner, DEFAULT_PROGRAM_ID,
+                "Account {account_id:?} was modified but not claimed"
+            );
         }
 
-        // TODO: Modify when multi-chain calls are supported in the circuit
-        if let Some(next_chained_call) = &program_output.chained_calls.first() {
-            program_id = next_chained_call.program_id;
-        } else if i != program_outputs.len() - 1 {
-            panic!("Inner call without a chained call found")
-        };
+        execution_state
     }
 
-    let n_accounts = pre_states.len();
-    if visibility_mask.len() != n_accounts {
-        panic!("Invalid visibility mask length");
+    /// Validate program pre and post states and populate the execution state.
+    fn validate_and_sync_states(
+        &mut self,
+        program_id: ProgramId,
+        authorized_pdas: HashSet<AccountId>,
+        pre_states: Vec<AccountWithMetadata>,
+        post_states: Vec<AccountPostState>,
+    ) {
+        for (pre, mut post) in pre_states.into_iter().zip(post_states) {
+            let pre_account_id = pre.account_id;
+            let post_states_entry = self.post_states.entry(pre.account_id);
+            match &post_states_entry {
+                Entry::Occupied(occupied) => {
+                    // Ensure that new pre state is the same as known post state
+                    assert_eq!(
+                        occupied.get(),
+                        &pre.account,
+                        "Inconsistent pre state for account {pre_account_id:?}",
+                    );
+
+                    let previous_is_authorized = self
+                        .pre_states
+                        .iter()
+                        .find(|acc| acc.account_id == pre_account_id)
+                        .map(|acc| acc.is_authorized)
+                        .unwrap_or_else(|| {
+                            panic!(
+                                "Pre state must exist in execution state for account {pre_account_id:?}",
+                            )
+                        });
+
+                    let is_authorized =
+                        previous_is_authorized || authorized_pdas.contains(&pre_account_id);
+
+                    assert_eq!(
+                        pre.is_authorized, is_authorized,
+                        "Inconsistent authorization for account {pre_account_id:?}",
+                    );
+                }
+                Entry::Vacant(_) => {
+                    self.pre_states.push(pre);
+                }
+            }
+
+            if post.requires_claim() {
+                // The invoked program can only claim accounts with default program id.
+                if post.account().program_owner == DEFAULT_PROGRAM_ID {
+                    post.account_mut().program_owner = program_id;
+                } else {
+                    panic!("Cannot claim an initialized account {pre_account_id:?}");
+                }
+            }
+
+            post_states_entry.insert_entry(post.into_account());
+        }
     }
 
-    // These lists will be the public outputs of this circuit
-    // and will be populated next.
-    let mut public_pre_states: Vec<AccountWithMetadata> = Vec::new();
-    let mut public_post_states: Vec<Account> = Vec::new();
-    let mut ciphertexts: Vec<Ciphertext> = Vec::new();
-    let mut new_commitments: Vec<Commitment> = Vec::new();
-    let mut new_nullifiers: Vec<(Nullifier, CommitmentSetDigest)> = Vec::new();
+    /// Get an iterator over pre and post states of each account involved in the execution.
+    pub fn into_states_iter(
+        mut self,
+    ) -> impl ExactSizeIterator<Item = (AccountWithMetadata, Account)> {
+        self.pre_states.into_iter().map(move |pre| {
+            let post = self
+                .post_states
+                .remove(&pre.account_id)
+                .expect("Account from pre states should exist in state diff");
+            (pre, post)
+        })
+    }
+}
+
+fn compute_circuit_output(
+    execution_state: ExecutionState,
+    visibility_mask: &[u8],
+    private_account_nonces: &[Nonce],
+    private_account_keys: &[(NullifierPublicKey, SharedSecretKey)],
+    private_account_nsks: &[NullifierSecretKey],
+    private_account_membership_proofs: &[Option<MembershipProof>],
+) -> PrivacyPreservingCircuitOutput {
+    let mut output = PrivacyPreservingCircuitOutput {
+        public_pre_states: Vec::new(),
+        public_post_states: Vec::new(),
+        ciphertexts: Vec::new(),
+        new_commitments: Vec::new(),
+        new_nullifiers: Vec::new(),
+    };
+
+    let states_iter = execution_state.into_states_iter();
+    assert_eq!(
+        visibility_mask.len(),
+        states_iter.len(),
+        "Invalid visibility mask length"
+    );
 
     let mut private_nonces_iter = private_account_nonces.iter();
     let mut private_keys_iter = private_account_keys.iter();
@@ -137,141 +249,156 @@ fn main() {
     let mut private_membership_proofs_iter = private_account_membership_proofs.iter();
 
     let mut output_index = 0;
-    for i in 0..n_accounts {
-        match visibility_mask[i] {
+    for (visibility_mask, (pre_state, post_state)) in
+        visibility_mask.iter().copied().zip(states_iter)
+    {
+        match visibility_mask {
             0 => {
                 // Public account
-                public_pre_states.push(pre_states[i].clone());
-
-                let mut post = state_diff.get(&pre_states[i].account_id).unwrap().clone();
-
-                if post.program_owner == DEFAULT_PROGRAM_ID {
-                    // Claim account
-                    post.program_owner = program_id;
-                }
-                public_post_states.push(post);
+                output.public_pre_states.push(pre_state);
+                output.public_post_states.push(post_state);
             }
             1 | 2 => {
-                let new_nonce = private_nonces_iter.next().expect("Missing private nonce");
-                let (npk, shared_secret) = private_keys_iter.next().expect("Missing keys");
+                let Some((npk, shared_secret)) = private_keys_iter.next() else {
+                    panic!("Missing private account key");
+                };
 
-                if AccountId::from(npk) != pre_states[i].account_id {
-                    panic!("AccountId mismatch");
-                }
+                assert_eq!(
+                    AccountId::from(npk),
+                    pre_state.account_id,
+                    "AccountId mismatch"
+                );
 
-                if visibility_mask[i] == 1 {
+                let new_nullifier = if visibility_mask == 1 {
                     // Private account with authentication
-                    let nsk = private_nsks_iter.next().expect("Missing nsk");
+
+                    let Some(nsk) = private_nsks_iter.next() else {
+                        panic!("Missing private account nullifier secret key");
+                    };
 
                     // Verify the nullifier public key
-                    let expected_npk = NullifierPublicKey::from(nsk);
-                    if &expected_npk != npk {
-                        panic!("Nullifier public key mismatch");
-                    }
+                    assert_eq!(
+                        npk,
+                        &NullifierPublicKey::from(nsk),
+                        "Nullifier public key mismatch"
+                    );
 
                     // Check pre_state authorization
-                    if !pre_states[i].is_authorized {
-                        panic!("Pre-state not authorized");
-                    }
+                    assert!(
+                        pre_state.is_authorized,
+                        "Pre-state not authorized for authenticated private account"
+                    );
 
-                    let membership_proof_opt = private_membership_proofs_iter
-                        .next()
-                        .expect("Missing membership proof");
-                    let (nullifier, set_digest) = membership_proof_opt
-                        .as_ref()
-                        .map(|membership_proof| {
-                            // Compute commitment set digest associated with provided auth path
-                            let commitment_pre = Commitment::new(npk, &pre_states[i].account);
-                            let set_digest =
-                                compute_digest_for_path(&commitment_pre, membership_proof);
+                    let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
+                        panic!("Missing membership proof");
+                    };
 
-                            // Compute update nullifier
-                            let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
-                            (nullifier, set_digest)
-                        })
-                        .unwrap_or_else(|| {
-                            if pre_states[i].account != Account::default() {
-                                panic!("Found new private account with non default values.");
-                            }
-
-                            // Compute initialization nullifier
-                            let nullifier = Nullifier::for_account_initialization(npk);
-                            (nullifier, DUMMY_COMMITMENT_HASH)
-                        });
-                    new_nullifiers.push((nullifier, set_digest));
+                    compute_nullifier_and_set_digest(
+                        membership_proof_opt.as_ref(),
+                        &pre_state.account,
+                        npk,
+                        nsk,
+                    )
                 } else {
                     // Private account without authentication
-                    if pre_states[i].account != Account::default() {
-                        panic!("Found new private account with non default values.");
-                    }
 
-                    if pre_states[i].is_authorized {
-                        panic!("Found new private account marked as authorized.");
-                    }
+                    assert_eq!(
+                        pre_state.account,
+                        Account::default(),
+                        "Found new private account with non default values",
+                    );
+
+                    assert!(
+                        !pre_state.is_authorized,
+                        "Found new private account marked as authorized."
+                    );
+
+                    let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
+                        panic!("Missing membership proof");
+                    };
 
-                    let membership_proof_opt = private_membership_proofs_iter
-                        .next()
-                        .expect("Missing membership proof");
                     assert!(
                         membership_proof_opt.is_none(),
                         "Membership proof must be None for unauthorized accounts"
                     );
+
                     let nullifier = Nullifier::for_account_initialization(npk);
-                    new_nullifiers.push((nullifier, DUMMY_COMMITMENT_HASH));
-                }
+                    (nullifier, DUMMY_COMMITMENT_HASH)
+                };
+                output.new_nullifiers.push(new_nullifier);
 
                 // Update post-state with new nonce
-                let mut post_with_updated_values =
-                    state_diff.get(&pre_states[i].account_id).unwrap().clone();
-                post_with_updated_values.nonce = *new_nonce;
-
-                if post_with_updated_values.program_owner == DEFAULT_PROGRAM_ID {
-                    // Claim account
-                    post_with_updated_values.program_owner = program_id;
-                }
+                let mut post_with_updated_nonce = post_state;
+                let Some(new_nonce) = private_nonces_iter.next() else {
+                    panic!("Missing private account nonce");
+                };
+                post_with_updated_nonce.nonce = *new_nonce;
 
                 // Compute commitment
-                let commitment_post = Commitment::new(npk, &post_with_updated_values);
+                let commitment_post = Commitment::new(npk, &post_with_updated_nonce);
 
                 // Encrypt and push post state
                 let encrypted_account = EncryptionScheme::encrypt(
-                    &post_with_updated_values,
+                    &post_with_updated_nonce,
                     shared_secret,
                     &commitment_post,
                     output_index,
                 );
 
-                new_commitments.push(commitment_post);
-                ciphertexts.push(encrypted_account);
+                output.new_commitments.push(commitment_post);
+                output.ciphertexts.push(encrypted_account);
                 output_index += 1;
             }
             _ => panic!("Invalid visibility mask value"),
         }
     }
 
-    if private_nonces_iter.next().is_some() {
-        panic!("Too many nonces");
-    }
+    assert!(private_nonces_iter.next().is_none(), "Too many nonces");
 
-    if private_keys_iter.next().is_some() {
-        panic!("Too many private account keys");
-    }
+    assert!(
+        private_keys_iter.next().is_none(),
+        "Too many private account keys"
+    );
 
-    if private_nsks_iter.next().is_some() {
-        panic!("Too many private account authentication keys");
-    }
+    assert!(
+        private_nsks_iter.next().is_none(),
+        "Too many private account nullifier secret keys"
+    );
 
-    if private_membership_proofs_iter.next().is_some() {
-        panic!("Too many private account membership proofs");
-    }
+    assert!(
+        private_membership_proofs_iter.next().is_none(),
+        "Too many private account membership proofs"
+    );
 
-    let output = PrivacyPreservingCircuitOutput {
-        public_pre_states,
-        public_post_states,
-        ciphertexts,
-        new_commitments,
-        new_nullifiers,
-    };
-
-    env::commit(&output);
+    output
+}
+
+fn compute_nullifier_and_set_digest(
+    membership_proof_opt: Option<&MembershipProof>,
+    pre_account: &Account,
+    npk: &NullifierPublicKey,
+    nsk: &NullifierSecretKey,
+) -> (Nullifier, CommitmentSetDigest) {
+    membership_proof_opt
+        .as_ref()
+        .map(|membership_proof| {
+            // Compute commitment set digest associated with provided auth path
+            let commitment_pre = Commitment::new(npk, pre_account);
+            let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
+
+            // Compute update nullifier
+            let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
+            (nullifier, set_digest)
+        })
+        .unwrap_or_else(|| {
+            assert_eq!(
+                *pre_account,
+                Account::default(),
+                "Found new private account with non default values"
+            );
+
+            // Compute initialization nullifier
+            let nullifier = Nullifier::for_account_initialization(npk);
+            (nullifier, DUMMY_COMMITMENT_HASH)
+        })
 }

sequencer_core/Cargo.toml

@@ -13,9 +13,16 @@ mempool.workspace = true
 base58.workspace = true
 anyhow.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 tempfile.workspace = true
 chrono.workspace = true
 log.workspace = true
+bedrock_client.workspace = true
+logos-blockchain-key-management-system-service.workspace = true
+logos-blockchain-core.workspace = true
+rand.workspace = true
+reqwest.workspace = true
+borsh.workspace = true
 
 [features]
 default = []

sequencer_core/src/block_settlement_client.rs

@@ -0,0 +1,117 @@
+use std::{fs, path::Path};
+
+use anyhow::{Context, Result, anyhow};
+use bedrock_client::BedrockClient;
+use common::block::HashableBlockData;
+use logos_blockchain_core::mantle::{
+    MantleTx, Op, OpProof, SignedMantleTx, Transaction, TxHash, ledger,
+    ops::channel::{ChannelId, MsgId, inscribe::InscriptionOp},
+};
+use logos_blockchain_key_management_system_service::keys::{
+    ED25519_SECRET_KEY_SIZE, Ed25519Key, Ed25519PublicKey,
+};
+
+use crate::config::BedrockConfig;
+
+/// A component that posts block data to logos blockchain
+pub struct BlockSettlementClient {
+    bedrock_client: BedrockClient,
+    bedrock_signing_key: Ed25519Key,
+    bedrock_channel_id: ChannelId,
+    last_message_id: MsgId,
+}
+
+impl BlockSettlementClient {
+    pub fn try_new(home: &Path, config: &BedrockConfig) -> Result<Self> {
+        let bedrock_signing_key = load_or_create_signing_key(&home.join("bedrock_signing_key"))
+            .context("Failed to load or create signing key")?;
+        let bedrock_channel_id = ChannelId::from(config.channel_id);
+        let bedrock_client = BedrockClient::new(None, config.node_url.clone())
+            .context("Failed to initialize bedrock client")?;
+        let channel_genesis_msg = MsgId::from([0; 32]);
+        Ok(Self {
+            bedrock_client,
+            bedrock_signing_key,
+            bedrock_channel_id,
+            last_message_id: channel_genesis_msg,
+        })
+    }
+
+    /// Create and sign a transaction for inscribing data
+    pub fn create_inscribe_tx(&self, data: Vec<u8>) -> (SignedMantleTx, MsgId) {
+        let verifying_key_bytes = self.bedrock_signing_key.public_key().to_bytes();
+        let verifying_key =
+            Ed25519PublicKey::from_bytes(&verifying_key_bytes).expect("valid ed25519 public key");
+
+        let inscribe_op = InscriptionOp {
+            channel_id: self.bedrock_channel_id,
+            inscription: data,
+            parent: self.last_message_id,
+            signer: verifying_key,
+        };
+        let inscribe_op_id = inscribe_op.id();
+
+        let ledger_tx = ledger::Tx::new(vec![], vec![]);
+
+        let inscribe_tx = MantleTx {
+            ops: vec![Op::ChannelInscribe(inscribe_op)],
+            ledger_tx,
+            // Altruistic test config
+            storage_gas_price: 0,
+            execution_gas_price: 0,
+        };
+
+        let tx_hash = inscribe_tx.hash();
+        let signature_bytes = self
+            .bedrock_signing_key
+            .sign_payload(tx_hash.as_signing_bytes().as_ref())
+            .to_bytes();
+        let signature =
+            logos_blockchain_key_management_system_service::keys::Ed25519Signature::from_bytes(
+                &signature_bytes,
+            );
+
+        let signed_mantle_tx = SignedMantleTx {
+            ops_proofs: vec![OpProof::Ed25519Sig(signature)],
+            ledger_tx_proof: empty_ledger_signature(&tx_hash),
+            mantle_tx: inscribe_tx,
+        };
+        (signed_mantle_tx, inscribe_op_id)
+    }
+
+    /// Post a transaction to the node and wait for inclusion
+    pub async fn post_and_wait(&mut self, block_data: &HashableBlockData) -> Result<u64> {
+        let inscription_data = borsh::to_vec(&block_data)?;
+        let (tx, new_msg_id) = self.create_inscribe_tx(inscription_data);
+
+        // Post the transaction
+        self.bedrock_client.post_transaction(tx).await?;
+
+        self.last_message_id = new_msg_id;
+
+        Ok(block_data.block_id)
+    }
+}
+
+/// Load signing key from file or generate a new one if it doesn't exist
+fn load_or_create_signing_key(path: &Path) -> Result<Ed25519Key> {
+    if path.exists() {
+        let key_bytes = fs::read(path)?;
+        let key_array: [u8; ED25519_SECRET_KEY_SIZE] = key_bytes
+            .try_into()
+            .map_err(|_| anyhow!("Found key with incorrect length"))?;
+        Ok(Ed25519Key::from_bytes(&key_array))
+    } else {
+        let mut key_bytes = [0u8; ED25519_SECRET_KEY_SIZE];
+        rand::RngCore::fill_bytes(&mut rand::thread_rng(), &mut key_bytes);
+        fs::write(path, key_bytes)?;
+        Ok(Ed25519Key::from_bytes(&key_bytes))
+    }
+}
+
+fn empty_ledger_signature(
+    tx_hash: &TxHash,
+) -> logos_blockchain_key_management_system_service::keys::ZkSignature {
+    logos_blockchain_key_management_system_service::keys::ZkKey::multi_sign(&[], tx_hash.as_ref())
+        .expect("multi-sign with empty key set works")
+}

sequencer_core/src/block_store.rs

@@ -46,7 +46,7 @@ impl SequencerBlockStore {
     }
 
     pub fn get_block_at_id(&self, id: u64) -> Result<Block> {
-        Ok(self.dbio.get_block(id)?.into_block(&self.signing_key))
+        Ok(self.dbio.get_block(id)?)
     }
 
     pub fn put_block_at_id(&mut self, block: Block) -> Result<()> {
@@ -113,7 +113,7 @@ mod tests {
             transactions: vec![],
         };
 
-        let genesis_block = genesis_block_hashable_data.into_block(&signing_key);
+        let genesis_block = genesis_block_hashable_data.into_pending_block(&signing_key);
         // Start an empty node store
         let mut node_store =
             SequencerBlockStore::open_db_with_genesis(path, Some(genesis_block), signing_key)

sequencer_core/src/config.rs

@@ -1,5 +1,11 @@
-use std::path::PathBuf;
+use std::{
+    fs::File,
+    io::BufReader,
+    path::{Path, PathBuf},
+};
 
+use anyhow::Result;
+use reqwest::Url;
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -42,4 +48,23 @@ pub struct SequencerConfig {
     pub initial_commitments: Vec<CommitmentsInitialData>,
     /// Sequencer own signing key
     pub signing_key: [u8; 32],
+    /// Bedrock configuration options
+    pub bedrock_config: Option<BedrockConfig>,
+}
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct BedrockConfig {
+    /// Bedrock channel ID
+    pub channel_id: [u8; 32],
+    /// Bedrock Url
+    pub node_url: Url,
+}
+
+impl SequencerConfig {
+    pub fn from_path(config_home: &Path) -> Result<SequencerConfig> {
+        let file = File::open(config_home)?;
+        let reader = BufReader::new(file);
+
+        Ok(serde_json::from_reader(reader)?)
+    }
 }

sequencer_core/src/lib.rs

@@ -13,8 +13,9 @@ use log::warn;
 use mempool::{MemPool, MemPoolHandle};
 use serde::{Deserialize, Serialize};
 
-use crate::block_store::SequencerBlockStore;
+use crate::{block_settlement_client::BlockSettlementClient, block_store::SequencerBlockStore};
 
+mod block_settlement_client;
 pub mod block_store;
 pub mod config;
 
@@ -24,6 +25,7 @@ pub struct SequencerCore {
     mempool: MemPool<EncodedTransaction>,
     sequencer_config: SequencerConfig,
     chain_height: u64,
+    block_settlement_client: Option<BlockSettlementClient>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -51,7 +53,7 @@ impl SequencerCore {
         };
 
         let signing_key = nssa::PrivateKey::try_new(config.signing_key).unwrap();
-        let genesis_block = hashable_data.into_block(&signing_key);
+        let genesis_block = hashable_data.into_pending_block(&signing_key);
 
         // Sequencer should panic if unable to open db,
         // as fixing this issue may require actions non-native to program scope
@@ -87,12 +89,18 @@ impl SequencerCore {
         state.add_pinata_program(PINATA_BASE58.parse().unwrap());
 
         let (mempool, mempool_handle) = MemPool::new(config.mempool_max_size);
+        let block_settlement_client = config.bedrock_config.as_ref().map(|bedrock_config| {
+            BlockSettlementClient::try_new(&config.home, bedrock_config)
+                .expect("Block settlement client should be constructible")
+        });
+
         let mut this = Self {
             state,
             block_store,
             mempool,
             chain_height: config.genesis_id,
             sequencer_config: config,
+            block_settlement_client,
         };
 
         this.sync_state_with_stored_blocks();
@@ -137,9 +145,21 @@ impl SequencerCore {
         Ok(tx)
     }
 
+    pub async fn produce_new_block_and_post_to_settlement_layer(&mut self) -> Result<u64> {
+        let block_data = self.produce_new_block_with_mempool_transactions()?;
+
+        if let Some(block_settlement) = self.block_settlement_client.as_mut() {
+            block_settlement.post_and_wait(&block_data).await?;
+            log::info!("Posted block data to Bedrock");
+        }
+
+        Ok(self.chain_height)
+    }
+
     /// Produces new block from transactions in mempool
-    pub fn produce_new_block_with_mempool_transactions(&mut self) -> Result<u64> {
+    pub fn produce_new_block_with_mempool_transactions(&mut self) -> Result<HashableBlockData> {
         let now = Instant::now();
+
         let new_block_height = self.chain_height + 1;
 
         let mut valid_transactions = vec![];
@@ -167,8 +187,6 @@ impl SequencerCore {
 
         let curr_time = chrono::Utc::now().timestamp_millis() as u64;
 
-        let num_txs_in_block = valid_transactions.len();
-
         let hashable_data = HashableBlockData {
             block_id: new_block_height,
             transactions: valid_transactions,
@@ -176,7 +194,9 @@ impl SequencerCore {
             timestamp: curr_time,
         };
 
-        let block = hashable_data.into_block(self.block_store.signing_key());
+        let block = hashable_data
+            .clone()
+            .into_pending_block(self.block_store.signing_key());
 
         self.block_store.put_block_at_id(block)?;
 
@@ -194,11 +214,10 @@ impl SequencerCore {
         // ```
         log::info!(
             "Created block with {} transactions in {} seconds",
-            num_txs_in_block,
+            hashable_data.transactions.len(),
             now.elapsed().as_secs()
         );
-
-        Ok(self.chain_height)
+        Ok(hashable_data)
     }
 
     pub fn state(&self) -> &nssa::V02State {
@@ -277,6 +296,7 @@ mod tests {
             initial_accounts,
             initial_commitments: vec![],
             signing_key: *sequencer_sign_key_for_testing().value(),
+            bedrock_config: None,
         }
     }
 
@@ -618,9 +638,9 @@ mod tests {
         let tx = common::test_utils::produce_dummy_empty_transaction();
         mempool_handle.push(tx).await.unwrap();
 
-        let block_id = sequencer.produce_new_block_with_mempool_transactions();
-        assert!(block_id.is_ok());
-        assert_eq!(block_id.unwrap(), genesis_height + 1);
+        let block = sequencer.produce_new_block_with_mempool_transactions();
+        assert!(block.is_ok());
+        assert_eq!(block.unwrap().block_id, genesis_height + 1);
     }
 
     #[tokio::test]
@@ -657,7 +677,8 @@ mod tests {
         // Create block
         let current_height = sequencer
             .produce_new_block_with_mempool_transactions()
-            .unwrap();
+            .unwrap()
+            .block_id;
         let block = sequencer
             .block_store
             .get_block_at_id(current_height)
@@ -696,7 +717,8 @@ mod tests {
         mempool_handle.push(tx.clone()).await.unwrap();
         let current_height = sequencer
             .produce_new_block_with_mempool_transactions()
-            .unwrap();
+            .unwrap()
+            .block_id;
         let block = sequencer
             .block_store
             .get_block_at_id(current_height)
@@ -707,7 +729,8 @@ mod tests {
         mempool_handle.push(tx.clone()).await.unwrap();
         let current_height = sequencer
             .produce_new_block_with_mempool_transactions()
-            .unwrap();
+            .unwrap()
+            .block_id;
         let block = sequencer
             .block_store
             .get_block_at_id(current_height)
@@ -742,7 +765,8 @@ mod tests {
             mempool_handle.push(tx.clone()).await.unwrap();
             let current_height = sequencer
                 .produce_new_block_with_mempool_transactions()
-                .unwrap();
+                .unwrap()
+                .block_id;
             let block = sequencer
                 .block_store
                 .get_block_at_id(current_height)

sequencer_rpc/src/net_utils.rs

@@ -1,4 +1,4 @@
-use std::{io, sync::Arc};
+use std::{io, net::SocketAddr, sync::Arc};
 
 use actix_cors::Cors;
 use actix_web::{App, Error as HttpError, HttpResponse, HttpServer, http, middleware, web};
@@ -42,25 +42,24 @@ fn get_cors(cors_allowed_origins: &[String]) -> Cors {
         .max_age(3600)
 }
 
-#[allow(clippy::too_many_arguments)]
 pub fn new_http_server(
     config: RpcConfig,
     seuquencer_core: Arc<Mutex<SequencerCore>>,
     mempool_handle: MemPoolHandle<EncodedTransaction>,
-) -> io::Result<actix_web::dev::Server> {
+) -> io::Result<(actix_web::dev::Server, SocketAddr)> {
     let RpcConfig {
         addr,
         cors_allowed_origins,
         limits_config,
     } = config;
-    info!(target:NETWORK, "Starting http server at {addr}");
+    info!(target:NETWORK, "Starting HTTP server at {addr}");
     let handler = web::Data::new(JsonHandler {
         sequencer_state: seuquencer_core.clone(),
         mempool_handle,
     });
 
     // HTTP server
-    Ok(HttpServer::new(move || {
+    let http_server = HttpServer::new(move || {
         App::new()
             .wrap(get_cors(&cors_allowed_origins))
             .app_data(handler.clone())
@@ -70,6 +69,14 @@ pub fn new_http_server(
     })
     .bind(addr)?
     .shutdown_timeout(SHUTDOWN_TIMEOUT_SECS)
-    .disable_signals()
-    .run())
+    .disable_signals();
+
+    let [addr] = http_server
+        .addrs()
+        .try_into()
+        .expect("Exactly one address bound is expected for sequencer HTTP server");
+
+    info!(target:NETWORK, "HTTP server started at {addr}");
+
+    Ok((http_server.run(), addr))
 }

sequencer_rpc/src/process.rs

@@ -388,6 +388,7 @@ mod tests {
             initial_accounts,
             initial_commitments: vec![],
             signing_key: *sequencer_sign_key_for_testing().value(),
+            bedrock_config: None,
         }
     }
 

sequencer_runner/Cargo.toml

@@ -10,7 +10,6 @@ sequencer_rpc.workspace = true
 
 clap = { workspace = true, features = ["derive", "env"] }
 anyhow.workspace = true
-serde_json.workspace = true
 env_logger.workspace = true
 log.workspace = true
 actix.workspace = true

sequencer_runner/Dockerfile

@@ -7,6 +7,7 @@ RUN apt-get update && apt-get install -y \
     libssl-dev \
     libclang-dev \
     clang \
+    curl \
     && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /sequencer_runner
@@ -31,6 +32,14 @@ RUN cargo build --release --bin sequencer_runner
 # Strip debug symbols to reduce binary size
 RUN strip /sequencer_runner/target/release/sequencer_runner
 
+# Install r0vm
+RUN curl -L https://risczero.com/install | bash
+ENV PATH="/root/.cargo/bin:/root/.risc0/bin:${PATH}"
+RUN rzup install
+RUN cp "$(which r0vm)" /usr/local/bin/r0vm
+RUN test -x /usr/local/bin/r0vm
+RUN r0vm --version
+
 # Runtime stage - minimal image
 FROM debian:trixie-slim
 
@@ -47,6 +56,9 @@ RUN useradd -m -u 1000 -s /bin/bash sequencer_user && \
 # Copy binary from builder
 COPY --from=builder --chown=sequencer_user:sequencer_user /sequencer_runner/target/release/sequencer_runner /usr/local/bin/sequencer_runner
 
+# Copy r0vm binary from builder
+COPY --from=builder --chown=sequencer_user:sequencer_user /usr/local/bin/r0vm /usr/local/bin/r0vm
+
 # Copy entrypoint script
 COPY sequencer_runner/docker-entrypoint.sh /docker-entrypoint.sh
 RUN chmod +x /docker-entrypoint.sh
@@ -71,6 +83,9 @@ HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
 # Run the application
 ENV RUST_LOG=info
 
+# Set explicit location for r0vm binary
+ENV RISC0_SERVER_PATH=/usr/local/bin/r0vm
+
 USER root
 
 ENTRYPOINT ["/docker-entrypoint.sh"]

sequencer_runner/configs/debug/sequencer_config.json

@@ -154,5 +154,9 @@
         37,
         37,
         37
-    ]
-}
+    ],
+    "bedrock_config": {
+        "channel_id": [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1],
+        "node_url": "http://localhost:8080"
+    }
+}

sequencer_runner/src/config.rs

@@ -1,11 +0,0 @@
-use std::{fs::File, io::BufReader, path::PathBuf};
-
-use anyhow::Result;
-use sequencer_core::config::SequencerConfig;
-
-pub fn from_file(config_home: PathBuf) -> Result<SequencerConfig> {
-    let file = File::open(config_home)?;
-    let reader = BufReader::new(file);
-
-    Ok(serde_json::from_reader(reader)?)
-}

sequencer_runner/src/lib.rs

@@ -1,4 +1,4 @@
-use std::{path::PathBuf, sync::Arc};
+use std::{net::SocketAddr, path::PathBuf, sync::Arc};
 
 use actix_web::dev::ServerHandle;
 use anyhow::Result;
@@ -9,8 +9,6 @@ use sequencer_core::{SequencerCore, config::SequencerConfig};
 use sequencer_rpc::new_http_server;
 use tokio::{sync::Mutex, task::JoinHandle};
 
-pub mod config;
-
 pub const RUST_LOG: &str = "RUST_LOG";
 
 #[derive(Parser, Debug)]
@@ -22,7 +20,7 @@ struct Args {
 
 pub async fn startup_sequencer(
     app_config: SequencerConfig,
-) -> Result<(ServerHandle, JoinHandle<Result<()>>)> {
+) -> Result<(ServerHandle, SocketAddr, JoinHandle<Result<()>>)> {
     let block_timeout = app_config.block_create_timeout_millis;
     let port = app_config.port;
 
@@ -32,7 +30,7 @@ pub async fn startup_sequencer(
 
     let seq_core_wrapped = Arc::new(Mutex::new(sequencer_core));
 
-    let http_server = new_http_server(
+    let (http_server, addr) = new_http_server(
         RpcConfig::with_port(port),
         Arc::clone(&seq_core_wrapped),
         mempool_handle,
@@ -52,7 +50,9 @@ pub async fn startup_sequencer(
             let id = {
                 let mut state = seq_core_wrapped.lock().await;
 
-                state.produce_new_block_with_mempool_transactions()?
+                state
+                    .produce_new_block_and_post_to_settlement_layer()
+                    .await?
             };
 
             info!("Block with id {id} created");
@@ -61,7 +61,7 @@ pub async fn startup_sequencer(
         }
     });
 
-    Ok((http_server_handle, main_loop_handle))
+    Ok((http_server_handle, addr, main_loop_handle))
 }
 
 pub async fn main_runner() -> Result<()> {
@@ -70,7 +70,7 @@ pub async fn main_runner() -> Result<()> {
     let args = Args::parse();
     let Args { home_dir } = args;
 
-    let app_config = config::from_file(home_dir.join("sequencer_config.json"))?;
+    let app_config = SequencerConfig::from_path(&home_dir.join("sequencer_config.json"))?;
 
     if let Some(ref rust_log) = app_config.override_rust_log {
         info!("RUST_LOG env var set to {rust_log:?}");
@@ -81,7 +81,7 @@ pub async fn main_runner() -> Result<()> {
     }
 
     // ToDo: Add restart on failures
-    let (_, main_loop_handle) = startup_sequencer(app_config).await?;
+    let (_, _, main_loop_handle) = startup_sequencer(app_config).await?;
 
     main_loop_handle.await??;
 

storage/src/lib.rs

@@ -1,6 +1,6 @@
 use std::{path::Path, sync::Arc};
 
-use common::block::{Block, HashableBlockData};
+use common::block::Block;
 use error::DbError;
 use rocksdb::{
     BoundColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode, MultiThreaded, Options,
@@ -26,6 +26,8 @@ pub const DB_META_FIRST_BLOCK_IN_DB_KEY: &str = "first_block_in_db";
 pub const DB_META_LAST_BLOCK_IN_DB_KEY: &str = "last_block_in_db";
 /// Key base for storing metainformation which describe if first block has been set
 pub const DB_META_FIRST_BLOCK_SET_KEY: &str = "first_block_set";
+/// Key base for storing metainformation about the last finalized block on Bedrock
+pub const DB_META_LAST_FINALIZED_BLOCK_ID: &str = "last_finalized_block_id";
 
 /// Key base for storing snapshot which describe block id
 pub const DB_SNAPSHOT_BLOCK_ID_KEY: &str = "block_id";
@@ -75,6 +77,7 @@ impl RocksDBIO {
             dbio.put_meta_first_block_in_db(block)?;
             dbio.put_meta_is_first_block_set()?;
             dbio.put_meta_last_block_in_db(block_id)?;
+            dbio.put_meta_last_finalized_block_id(None)?;
 
             Ok(dbio)
         } else {
@@ -232,6 +235,28 @@ impl RocksDBIO {
         Ok(())
     }
 
+    pub fn put_meta_last_finalized_block_id(&self, block_id: Option<u64>) -> DbResult<()> {
+        let cf_meta = self.meta_column();
+        self.db
+            .put_cf(
+                &cf_meta,
+                borsh::to_vec(&DB_META_LAST_FINALIZED_BLOCK_ID).map_err(|err| {
+                    DbError::borsh_cast_message(
+                        err,
+                        Some("Failed to serialize DB_META_LAST_FINALIZED_BLOCK_ID".to_string()),
+                    )
+                })?,
+                borsh::to_vec(&block_id).map_err(|err| {
+                    DbError::borsh_cast_message(
+                        err,
+                        Some("Failed to serialize last block id".to_string()),
+                    )
+                })?,
+            )
+            .map_err(|rerr| DbError::rocksdb_cast_message(rerr, None))?;
+        Ok(())
+    }
+
     pub fn put_meta_is_first_block_set(&self) -> DbResult<()> {
         let cf_meta = self.meta_column();
         self.db
@@ -269,7 +294,7 @@ impl RocksDBIO {
                         Some("Failed to serialize block id".to_string()),
                     )
                 })?,
-                borsh::to_vec(&HashableBlockData::from(block)).map_err(|err| {
+                borsh::to_vec(&block).map_err(|err| {
                     DbError::borsh_cast_message(
                         err,
                         Some("Failed to serialize block data".to_string()),
@@ -280,7 +305,7 @@ impl RocksDBIO {
         Ok(())
     }
 
-    pub fn get_block(&self, block_id: u64) -> DbResult<HashableBlockData> {
+    pub fn get_block(&self, block_id: u64) -> DbResult<Block> {
         let cf_block = self.block_column();
         let res = self
             .db
@@ -296,14 +321,12 @@ impl RocksDBIO {
             .map_err(|rerr| DbError::rocksdb_cast_message(rerr, None))?;
 
         if let Some(data) = res {
-            Ok(
-                borsh::from_slice::<HashableBlockData>(&data).map_err(|serr| {
-                    DbError::borsh_cast_message(
-                        serr,
-                        Some("Failed to deserialize block data".to_string()),
-                    )
-                })?,
-            )
+            Ok(borsh::from_slice::<Block>(&data).map_err(|serr| {
+                DbError::borsh_cast_message(
+                    serr,
+                    Some("Failed to deserialize block data".to_string()),
+                )
+            })?)
         } else {
             Err(DbError::db_interaction_error(
                 "Block on this id not found".to_string(),

test_program_methods/guest/src/bin/changer_claimer.rs

@@ -0,0 +1,38 @@
+use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs};
+
+type Instruction = (Option<Vec<u8>>, bool);
+
+/// A program that optionally modifies the account data and optionally claims it.
+fn main() {
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: (data_opt, should_claim),
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let [pre] = match pre_states.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    let account_pre = &pre.account;
+    let mut account_post = account_pre.clone();
+
+    // Update data if provided
+    if let Some(data) = data_opt {
+        account_post.data = data
+            .try_into()
+            .expect("provided data should fit into data limit");
+    }
+
+    // Claim or not based on the boolean flag
+    let post_state = if should_claim {
+        AccountPostState::new_claimed(account_post)
+    } else {
+        AccountPostState::new(account_post)
+    };
+
+    write_nssa_outputs(instruction_words, vec![pre], vec![post_state]);
+}

test_program_methods/guest/src/bin/malicious_authorization_changer.rs

@@ -0,0 +1,53 @@
+use nssa_core::{
+    account::AccountWithMetadata,
+    program::{
+        AccountPostState, ChainedCall, ProgramId, ProgramInput, read_nssa_inputs,
+        write_nssa_outputs_with_chained_call,
+    },
+};
+use risc0_zkvm::serde::to_vec;
+
+type Instruction = (u128, ProgramId);
+
+/// A malicious test program that attempts to change authorization status.
+/// It accepts two accounts and executes a native token transfer program via chain call,
+/// but sets the `is_authorized` field of the first account to true.
+fn main() {
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: (balance, transfer_program_id),
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let [sender, receiver] = match pre_states.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    // Maliciously set is_authorized to true for the first account
+    let authorised_sender = AccountWithMetadata {
+        is_authorized: true,
+        ..sender.clone()
+    };
+
+    let instruction_data = to_vec(&balance).unwrap();
+
+    let chained_call = ChainedCall {
+        program_id: transfer_program_id,
+        instruction_data,
+        pre_states: vec![authorised_sender.clone(), receiver.clone()],
+        pda_seeds: vec![],
+    };
+
+    write_nssa_outputs_with_chained_call(
+        instruction_words,
+        vec![sender.clone(), receiver.clone()],
+        vec![
+            AccountPostState::new(sender.account),
+            AccountPostState::new(receiver.account),
+        ],
+        vec![chained_call],
+    );
+}

wallet/Cargo.toml

@@ -14,7 +14,7 @@ serde_json.workspace = true
 env_logger.workspace = true
 log.workspace = true
 serde.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["macros"] }
 clap.workspace = true
 base64.workspace = true
 bytemuck.workspace = true
@@ -25,6 +25,7 @@ rand.workspace = true
 itertools.workspace = true
 sha2.workspace = true
 futures.workspace = true
+risc0-zkvm.workspace = true
 async-stream = "0.3.6"
 indicatif = { version = "0.18.3", features = ["improved_unicode"] }
-risc0-zkvm.workspace = true
+optfield = "0.4.0"

wallet/src/cli/account.rs

@@ -3,7 +3,7 @@ use base58::ToBase58;
 use clap::Subcommand;
 use itertools::Itertools as _;
 use key_protocol::key_management::key_tree::chain_index::ChainIndex;
-use nssa::{Account, program::Program};
+use nssa::{Account, PublicKey, program::Program};
 use serde::Serialize;
 
 use crate::{
@@ -20,6 +20,9 @@ pub enum AccountSubcommand {
         /// Flag to get raw account data
         #[arg(short, long)]
         raw: bool,
+        /// Display keys (pk for public accounts, npk/ipk for private accounts)
+        #[arg(short, long)]
+        keys: bool,
         /// Valid 32 byte base58 string with privacy prefix
         #[arg(short, long)]
         account_id: String,
@@ -64,13 +67,20 @@ impl WalletSubcommand for NewSubcommand {
             NewSubcommand::Public { cci } => {
                 let (account_id, chain_index) = wallet_core.create_new_account_public(cci);
 
+                let private_key = wallet_core
+                    .storage
+                    .user_data
+                    .get_pub_account_signing_key(&account_id)
+                    .unwrap();
+
+                let public_key = PublicKey::new_from_private_key(private_key);
+
                 println!(
                     "Generated new account with account_id Public/{account_id} at path {chain_index}"
                 );
+                println!("With pk {}", hex::encode(public_key.value()));
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::RegisterAccount { account_id })
             }
@@ -93,9 +103,7 @@ impl WalletSubcommand for NewSubcommand {
                     hex::encode(key.incoming_viewing_public_key.to_bytes())
                 );
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::RegisterAccount { account_id })
             }
@@ -205,7 +213,11 @@ impl WalletSubcommand for AccountSubcommand {
         wallet_core: &mut WalletCore,
     ) -> Result<SubcommandReturnValue> {
         match self {
-            AccountSubcommand::Get { raw, account_id } => {
+            AccountSubcommand::Get {
+                raw,
+                keys,
+                account_id,
+            } => {
                 let (account_id, addr_kind) = parse_addr_with_privacy_prefix(&account_id)?;
 
                 let account_id = account_id.parse()?;
@@ -219,9 +231,43 @@ impl WalletSubcommand for AccountSubcommand {
                         .ok_or(anyhow::anyhow!("Private account not found in storage"))?,
                 };
 
+                // Helper closure to display keys for the account
+                let display_keys = |wallet_core: &WalletCore| -> Result<()> {
+                    match addr_kind {
+                        AccountPrivacyKind::Public => {
+                            let private_key = wallet_core
+                                .storage
+                                .user_data
+                                .get_pub_account_signing_key(&account_id)
+                                .ok_or(anyhow::anyhow!("Public account not found in storage"))?;
+
+                            let public_key = PublicKey::new_from_private_key(private_key);
+                            println!("pk {}", hex::encode(public_key.value()));
+                        }
+                        AccountPrivacyKind::Private => {
+                            let (key, _) = wallet_core
+                                .storage
+                                .user_data
+                                .get_private_account(&account_id)
+                                .ok_or(anyhow::anyhow!("Private account not found in storage"))?;
+
+                            println!("npk {}", hex::encode(key.nullifer_public_key.0));
+                            println!(
+                                "ipk {}",
+                                hex::encode(key.incoming_viewing_public_key.to_bytes())
+                            );
+                        }
+                    }
+                    Ok(())
+                };
+
                 if account == Account::default() {
                     println!("Account is Uninitialized");
 
+                    if keys {
+                        display_keys(wallet_core)?;
+                    }
+
                     return Ok(SubcommandReturnValue::Empty);
                 }
 
@@ -236,6 +282,10 @@ impl WalletSubcommand for AccountSubcommand {
                 println!("{description}");
                 println!("{json_view}");
 
+                if keys {
+                    display_keys(wallet_core)?;
+                }
+
                 Ok(SubcommandReturnValue::Empty)
             }
             AccountSubcommand::New(new_subcommand) => {
@@ -257,9 +307,7 @@ impl WalletSubcommand for AccountSubcommand {
                 {
                     wallet_core.last_synced_block = curr_last_block;
 
-                    let path = wallet_core.store_persistent_data().await?;
-
-                    println!("Stored persistent data at {path:#?}");
+                    wallet_core.store_persistent_data().await?;
                 } else {
                     wallet_core.sync_to_block(curr_last_block).await?;
                 }
@@ -294,7 +342,7 @@ impl WalletSubcommand for AccountSubcommand {
                                 .iter()
                                 .map(|(id, chain_index)| format!("{chain_index} Private/{id}")),
                         )
-                        .format(",\n");
+                        .format("\n");
 
                     println!("{accounts}");
                     return Ok(SubcommandReturnValue::Empty);

wallet/src/cli/chain.rs

@@ -19,7 +19,7 @@ pub enum ChainSubcommand {
     /// Get transaction at hash from sequencer
     Transaction {
         /// hash - valid 32 byte hex string
-        #[arg(short, long)]
+        #[arg(short = 't', long)]
         hash: String,
     },
 }

wallet/src/cli/config.rs

@@ -10,7 +10,13 @@ use crate::{
 #[derive(Subcommand, Debug, Clone)]
 pub enum ConfigSubcommand {
     /// Getter of config fields
-    Get { key: String },
+    Get {
+        /// Print all config fields
+        #[arg(short, long)]
+        all: bool,
+        /// Config field key to get
+        key: Option<String>,
+    },
     /// Setter of config fields
     Set { key: String, value: String },
     /// Prints description of corresponding field
@@ -23,58 +29,66 @@ impl WalletSubcommand for ConfigSubcommand {
         wallet_core: &mut WalletCore,
     ) -> Result<SubcommandReturnValue> {
         match self {
-            ConfigSubcommand::Get { key } => match key.as_str() {
-                "all" => {
+            ConfigSubcommand::Get { all, key } => {
+                if all {
                     let config_str =
                         serde_json::to_string_pretty(&wallet_core.storage.wallet_config)?;
 
                     println!("{config_str}");
-                }
-                "override_rust_log" => {
-                    if let Some(value) = &wallet_core.storage.wallet_config.override_rust_log {
-                        println!("{value}");
-                    } else {
-                        println!("Not set");
+                } else if let Some(key) = key {
+                    match key.as_str() {
+                        "override_rust_log" => {
+                            if let Some(value) =
+                                &wallet_core.storage.wallet_config.override_rust_log
+                            {
+                                println!("{value}");
+                            } else {
+                                println!("Not set");
+                            }
+                        }
+                        "sequencer_addr" => {
+                            println!("{}", wallet_core.storage.wallet_config.sequencer_addr);
+                        }
+                        "seq_poll_timeout_millis" => {
+                            println!(
+                                "{}",
+                                wallet_core.storage.wallet_config.seq_poll_timeout_millis
+                            );
+                        }
+                        "seq_tx_poll_max_blocks" => {
+                            println!(
+                                "{}",
+                                wallet_core.storage.wallet_config.seq_tx_poll_max_blocks
+                            );
+                        }
+                        "seq_poll_max_retries" => {
+                            println!("{}", wallet_core.storage.wallet_config.seq_poll_max_retries);
+                        }
+                        "seq_block_poll_max_amount" => {
+                            println!(
+                                "{}",
+                                wallet_core.storage.wallet_config.seq_block_poll_max_amount
+                            );
+                        }
+                        "initial_accounts" => {
+                            println!("{:#?}", wallet_core.storage.wallet_config.initial_accounts);
+                        }
+                        "basic_auth" => {
+                            if let Some(basic_auth) = &wallet_core.storage.wallet_config.basic_auth
+                            {
+                                println!("{basic_auth}");
+                            } else {
+                                println!("Not set");
+                            }
+                        }
+                        _ => {
+                            println!("Unknown field");
+                        }
                     }
+                } else {
+                    println!("Please provide a key or use --all flag");
                 }
-                "sequencer_addr" => {
-                    println!("{}", wallet_core.storage.wallet_config.sequencer_addr);
-                }
-                "seq_poll_timeout_millis" => {
-                    println!(
-                        "{}",
-                        wallet_core.storage.wallet_config.seq_poll_timeout_millis
-                    );
-                }
-                "seq_tx_poll_max_blocks" => {
-                    println!(
-                        "{}",
-                        wallet_core.storage.wallet_config.seq_tx_poll_max_blocks
-                    );
-                }
-                "seq_poll_max_retries" => {
-                    println!("{}", wallet_core.storage.wallet_config.seq_poll_max_retries);
-                }
-                "seq_block_poll_max_amount" => {
-                    println!(
-                        "{}",
-                        wallet_core.storage.wallet_config.seq_block_poll_max_amount
-                    );
-                }
-                "initial_accounts" => {
-                    println!("{:#?}", wallet_core.storage.wallet_config.initial_accounts);
-                }
-                "basic_auth" => {
-                    if let Some(basic_auth) = &wallet_core.storage.wallet_config.basic_auth {
-                        println!("{basic_auth}");
-                    } else {
-                        println!("Not set");
-                    }
-                }
-                _ => {
-                    println!("Unknown field");
-                }
-            },
+            }
             ConfigSubcommand::Set { key, value } => {
                 match key.as_str() {
                     "override_rust_log" => {
@@ -108,9 +122,7 @@ impl WalletSubcommand for ConfigSubcommand {
                     }
                 }
 
-                let path = wallet_core.store_config_changes().await?;
-
-                println!("Stored changed config at {path:#?}");
+                wallet_core.store_config_changes().await?
             }
             ConfigSubcommand::Description { key } => match key.as_str() {
                 "override_rust_log" => {

wallet/src/cli/mod.rs

@@ -15,7 +15,6 @@ use crate::{
             pinata::PinataProgramAgnosticSubcommand, token::TokenProgramAgnosticSubcommand,
         },
     },
-    helperfunctions::{fetch_config, fetch_persistent_storage, merge_auth_config},
 };
 
 pub mod account;
@@ -97,43 +96,22 @@ pub enum SubcommandReturnValue {
     SyncedToBlock(u64),
 }
 
-pub async fn execute_subcommand(command: Command) -> Result<SubcommandReturnValue> {
-    execute_subcommand_with_auth(command, None).await
-}
-
-pub async fn execute_subcommand_with_auth(
+pub async fn execute_subcommand(
+    wallet_core: &mut WalletCore,
     command: Command,
-    auth: Option<String>,
 ) -> Result<SubcommandReturnValue> {
-    if fetch_persistent_storage().await.is_err() {
-        println!("Persistent storage not found, need to execute setup");
-
-        let password = read_password_from_stdin()?;
-        execute_setup_with_auth(password, auth.clone()).await?;
-    }
-
-    let wallet_config = fetch_config().await?;
-    let wallet_config = merge_auth_config(wallet_config, auth.clone())?;
-    let mut wallet_core = WalletCore::start_from_config_update_chain(wallet_config).await?;
-
     let subcommand_ret = match command {
         Command::AuthTransfer(transfer_subcommand) => {
-            transfer_subcommand
-                .handle_subcommand(&mut wallet_core)
-                .await?
+            transfer_subcommand.handle_subcommand(wallet_core).await?
         }
         Command::ChainInfo(chain_subcommand) => {
-            chain_subcommand.handle_subcommand(&mut wallet_core).await?
+            chain_subcommand.handle_subcommand(wallet_core).await?
         }
         Command::Account(account_subcommand) => {
-            account_subcommand
-                .handle_subcommand(&mut wallet_core)
-                .await?
+            account_subcommand.handle_subcommand(wallet_core).await?
         }
         Command::Pinata(pinata_subcommand) => {
-            pinata_subcommand
-                .handle_subcommand(&mut wallet_core)
-                .await?
+            pinata_subcommand.handle_subcommand(wallet_core).await?
         }
         Command::CheckHealth {} => {
             let remote_program_ids = wallet_core
@@ -165,18 +143,15 @@ pub async fn execute_subcommand_with_auth(
 
             SubcommandReturnValue::Empty
         }
-        Command::Token(token_subcommand) => {
-            token_subcommand.handle_subcommand(&mut wallet_core).await?
-        }
-        Command::AMM(amm_subcommand) => amm_subcommand.handle_subcommand(&mut wallet_core).await?,
+        Command::Token(token_subcommand) => token_subcommand.handle_subcommand(wallet_core).await?,
+        Command::AMM(amm_subcommand) => amm_subcommand.handle_subcommand(wallet_core).await?,
         Command::Config(config_subcommand) => {
-            config_subcommand
-                .handle_subcommand(&mut wallet_core)
-                .await?
+            config_subcommand.handle_subcommand(wallet_core).await?
         }
         Command::RestoreKeys { depth } => {
             let password = read_password_from_stdin()?;
-            execute_keys_restoration_with_auth(password, depth, auth).await?;
+            wallet_core.reset_storage(password)?;
+            execute_keys_restoration(wallet_core, depth).await?;
 
             SubcommandReturnValue::Empty
         }
@@ -200,14 +175,7 @@ pub async fn execute_subcommand_with_auth(
     Ok(subcommand_ret)
 }
 
-pub async fn execute_continuous_run() -> Result<()> {
-    execute_continuous_run_with_auth(None).await
-}
-pub async fn execute_continuous_run_with_auth(auth: Option<String>) -> Result<()> {
-    let config = fetch_config().await?;
-    let config = merge_auth_config(config, auth)?;
-    let mut wallet_core = WalletCore::start_from_config_update_chain(config.clone()).await?;
-
+pub async fn execute_continuous_run(wallet_core: &mut WalletCore) -> Result<()> {
     loop {
         let latest_block_num = wallet_core
             .sequencer_client
@@ -217,7 +185,7 @@ pub async fn execute_continuous_run_with_auth(auth: Option<String>) -> Result<()
         wallet_core.sync_to_block(latest_block_num).await?;
 
         tokio::time::sleep(std::time::Duration::from_millis(
-            config.seq_poll_timeout_millis,
+            wallet_core.config().seq_poll_timeout_millis,
         ))
         .await;
     }
@@ -233,34 +201,7 @@ pub fn read_password_from_stdin() -> Result<String> {
     Ok(password.trim().to_string())
 }
 
-pub async fn execute_setup(password: String) -> Result<()> {
-    execute_setup_with_auth(password, None).await
-}
-
-pub async fn execute_setup_with_auth(password: String, auth: Option<String>) -> Result<()> {
-    let config = fetch_config().await?;
-    let config = merge_auth_config(config, auth)?;
-    let wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password).await?;
-
-    wallet_core.store_persistent_data().await?;
-
-    Ok(())
-}
-
-pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> {
-    execute_keys_restoration_with_auth(password, depth, None).await
-}
-
-pub async fn execute_keys_restoration_with_auth(
-    password: String,
-    depth: u32,
-    auth: Option<String>,
-) -> Result<()> {
-    let config = fetch_config().await?;
-    let config = merge_auth_config(config, auth)?;
-    let mut wallet_core =
-        WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?;
-
+pub async fn execute_keys_restoration(wallet_core: &mut WalletCore, depth: u32) -> Result<()> {
     wallet_core
         .storage
         .user_data

wallet/src/cli/programs/native_token_transfer.rs

@@ -69,9 +69,7 @@ impl WalletSubcommand for AuthTransferSubcommand {
 
                         println!("Transaction data is {transfer_tx:?}");
 
-                        let path = wallet_core.store_persistent_data().await?;
-
-                        println!("Stored persistent accounts at {path:#?}");
+                        wallet_core.store_persistent_data().await?;
                     }
                     AccountPrivacyKind::Private => {
                         let account_id = account_id.parse()?;
@@ -96,9 +94,7 @@ impl WalletSubcommand for AuthTransferSubcommand {
                             )?;
                         }
 
-                        let path = wallet_core.store_persistent_data().await?;
-
-                        println!("Stored persistent accounts at {path:#?}");
+                        wallet_core.store_persistent_data().await?;
                     }
                 }
 
@@ -337,9 +333,7 @@ impl WalletSubcommand for NativeTokenTransferProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -381,9 +375,7 @@ impl WalletSubcommand for NativeTokenTransferProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -421,9 +413,7 @@ impl WalletSubcommand for NativeTokenTransferProgramSubcommandShielded {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -454,9 +444,7 @@ impl WalletSubcommand for NativeTokenTransferProgramSubcommandShielded {
 
                 let tx_hash = res.tx_hash;
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -500,9 +488,7 @@ impl WalletSubcommand for NativeTokenTransferProgramSubcommand {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -520,9 +506,7 @@ impl WalletSubcommand for NativeTokenTransferProgramSubcommand {
 
                 println!("Transaction data is {transfer_tx:?}");
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::Empty)
             }

wallet/src/cli/programs/pinata.rs

@@ -168,9 +168,7 @@ impl WalletSubcommand for PinataProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }

wallet/src/cli/programs/token.rs

@@ -740,9 +740,7 @@ impl WalletSubcommand for TokenProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -790,9 +788,7 @@ impl WalletSubcommand for TokenProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -831,9 +827,7 @@ impl WalletSubcommand for TokenProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -872,9 +866,7 @@ impl WalletSubcommand for TokenProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -923,9 +915,7 @@ impl WalletSubcommand for TokenProgramSubcommandPrivate {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -971,9 +961,7 @@ impl WalletSubcommand for TokenProgramSubcommandDeshielded {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1009,9 +997,7 @@ impl WalletSubcommand for TokenProgramSubcommandDeshielded {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1047,9 +1033,7 @@ impl WalletSubcommand for TokenProgramSubcommandDeshielded {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1102,9 +1086,7 @@ impl WalletSubcommand for TokenProgramSubcommandShielded {
                     println!("Transaction data is {:?}", tx.message);
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1140,9 +1122,7 @@ impl WalletSubcommand for TokenProgramSubcommandShielded {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1178,9 +1158,7 @@ impl WalletSubcommand for TokenProgramSubcommandShielded {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1216,9 +1194,7 @@ impl WalletSubcommand for TokenProgramSubcommandShielded {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1262,9 +1238,7 @@ impl WalletSubcommand for TokenProgramSubcommandShielded {
                     println!("Transaction data is {:?}", tx.message);
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1323,9 +1297,7 @@ impl WalletSubcommand for CreateNewTokenProgramSubcommand {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1371,9 +1343,7 @@ impl WalletSubcommand for CreateNewTokenProgramSubcommand {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }
@@ -1419,9 +1389,7 @@ impl WalletSubcommand for CreateNewTokenProgramSubcommand {
                     )?;
                 }
 
-                let path = wallet_core.store_persistent_data().await?;
-
-                println!("Stored persistent accounts at {path:#?}");
+                wallet_core.store_persistent_data().await?;
 
                 Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
             }

wallet/src/config.rs

@@ -1,11 +1,17 @@
-use std::str::FromStr;
+use std::{
+    io::{BufReader, Write as _},
+    path::Path,
+    str::FromStr,
+};
 
+use anyhow::{Context as _, Result};
 use key_protocol::key_management::{
     KeyChain,
     key_tree::{
         chain_index::ChainIndex, keys_private::ChildKeysPrivate, keys_public::ChildKeysPublic,
     },
 };
+use log::warn;
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -105,6 +111,25 @@ pub struct PersistentStorage {
     pub last_synced_block: u64,
 }
 
+impl PersistentStorage {
+    pub fn from_path(path: &Path) -> Result<Self> {
+        match std::fs::File::open(path) {
+            Ok(file) => {
+                let storage_content = BufReader::new(file);
+                Ok(serde_json::from_reader(storage_content)?)
+            }
+            Err(err) => match err.kind() {
+                std::io::ErrorKind::NotFound => {
+                    anyhow::bail!("Not found, please setup roots from config command beforehand");
+                }
+                _ => {
+                    anyhow::bail!("IO error {err:#?}");
+                }
+            },
+        }
+    }
+}
+
 impl InitialAccountData {
     pub fn account_id(&self) -> nssa::AccountId {
         match &self {
@@ -172,9 +197,11 @@ pub struct GasConfig {
     pub gas_limit_runtime: u64,
 }
 
+#[optfield::optfield(pub WalletConfigOverrides, rewrap, attrs = (derive(Debug, Default)))]
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct WalletConfig {
     /// Override rust log (env var logging level)
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub override_rust_log: Option<String>,
     /// Sequencer URL
     pub sequencer_addr: String,
@@ -189,6 +216,7 @@ pub struct WalletConfig {
     /// Initial accounts for wallet
     pub initial_accounts: Vec<InitialAccountData>,
     /// Basic authentication credentials
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub basic_auth: Option<BasicAuth>,
 }
 
@@ -748,3 +776,98 @@ impl Default for WalletConfig {
         }
     }
 }
+
+impl WalletConfig {
+    pub fn from_path_or_initialize_default(config_path: &Path) -> Result<WalletConfig> {
+        match std::fs::File::open(config_path) {
+            Ok(file) => {
+                let reader = std::io::BufReader::new(file);
+                Ok(serde_json::from_reader(reader)?)
+            }
+            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
+                println!("Config not found, setting up default config");
+
+                let config_home = config_path.parent().ok_or_else(|| {
+                    anyhow::anyhow!(
+                        "Could not get parent directory of config file at {config_path:#?}"
+                    )
+                })?;
+                std::fs::create_dir_all(config_home)?;
+
+                println!("Created configs dir at path {config_home:#?}");
+
+                let mut file = std::fs::OpenOptions::new()
+                    .write(true)
+                    .create(true)
+                    .truncate(true)
+                    .open(config_path)?;
+
+                let config = WalletConfig::default();
+                let default_config_serialized = serde_json::to_vec_pretty(&config).unwrap();
+
+                file.write_all(&default_config_serialized)?;
+
+                println!("Configs set up");
+                Ok(config)
+            }
+            Err(err) => Err(err).context("IO error"),
+        }
+    }
+
+    pub fn apply_overrides(&mut self, overrides: WalletConfigOverrides) {
+        let WalletConfig {
+            override_rust_log,
+            sequencer_addr,
+            seq_poll_timeout_millis,
+            seq_tx_poll_max_blocks,
+            seq_poll_max_retries,
+            seq_block_poll_max_amount,
+            initial_accounts,
+            basic_auth,
+        } = self;
+
+        let WalletConfigOverrides {
+            override_rust_log: o_override_rust_log,
+            sequencer_addr: o_sequencer_addr,
+            seq_poll_timeout_millis: o_seq_poll_timeout_millis,
+            seq_tx_poll_max_blocks: o_seq_tx_poll_max_blocks,
+            seq_poll_max_retries: o_seq_poll_max_retries,
+            seq_block_poll_max_amount: o_seq_block_poll_max_amount,
+            initial_accounts: o_initial_accounts,
+            basic_auth: o_basic_auth,
+        } = overrides;
+
+        if let Some(v) = o_override_rust_log {
+            warn!("Overriding wallet config 'override_rust_log' to {v:#?}");
+            *override_rust_log = v;
+        }
+        if let Some(v) = o_sequencer_addr {
+            warn!("Overriding wallet config 'sequencer_addr' to {v}");
+            *sequencer_addr = v;
+        }
+        if let Some(v) = o_seq_poll_timeout_millis {
+            warn!("Overriding wallet config 'seq_poll_timeout_millis' to {v}");
+            *seq_poll_timeout_millis = v;
+        }
+        if let Some(v) = o_seq_tx_poll_max_blocks {
+            warn!("Overriding wallet config 'seq_tx_poll_max_blocks' to {v}");
+            *seq_tx_poll_max_blocks = v;
+        }
+        if let Some(v) = o_seq_poll_max_retries {
+            warn!("Overriding wallet config 'seq_poll_max_retries' to {v}");
+            *seq_poll_max_retries = v;
+        }
+        if let Some(v) = o_seq_block_poll_max_amount {
+            warn!("Overriding wallet config 'seq_block_poll_max_amount' to {v}");
+            *seq_block_poll_max_amount = v;
+        }
+        if let Some(v) = o_initial_accounts {
+            warn!("Overriding wallet config 'initial_accounts' to {v:#?}");
+            *initial_accounts = v;
+        }
+        if let Some(v) = o_basic_auth {
+            warn!("Overriding wallet config 'basic_auth' to {v:#?}");
+            *basic_auth = v;
+        }
+    }
+}

wallet/src/helperfunctions.rs

@@ -7,23 +7,22 @@ use nssa::Account;
 use nssa_core::account::Nonce;
 use rand::{RngCore, rngs::OsRng};
 use serde::Serialize;
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 use crate::{
     HOME_DIR_ENV_VAR,
     config::{
-        BasicAuth, InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic,
-        PersistentAccountDataPrivate, PersistentAccountDataPublic, PersistentStorage, WalletConfig,
+        InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic,
+        PersistentAccountDataPrivate, PersistentAccountDataPublic, PersistentStorage,
     },
 };
 
 /// Get home dir for wallet. Env var `NSSA_WALLET_HOME_DIR` must be set before execution to succeed.
-pub fn get_home_nssa_var() -> Result<PathBuf> {
+fn get_home_nssa_var() -> Result<PathBuf> {
     Ok(PathBuf::from_str(&std::env::var(HOME_DIR_ENV_VAR)?)?)
 }
 
 /// Get home dir for wallet. Env var `HOME` must be set before execution to succeed.
-pub fn get_home_default_path() -> Result<PathBuf> {
+fn get_home_default_path() -> Result<PathBuf> {
     std::env::home_dir()
         .map(|path| path.join(".nssa").join("wallet"))
         .ok_or(anyhow::anyhow!("Failed to get HOME"))
@@ -38,96 +37,20 @@ pub fn get_home() -> Result<PathBuf> {
     }
 }
 
-/// Fetch config from default home
-pub async fn fetch_config() -> Result<WalletConfig> {
-    let config_home = get_home()?;
-    let mut config_needs_setup = false;
-
-    let config = match tokio::fs::OpenOptions::new()
-        .read(true)
-        .open(config_home.join("wallet_config.json"))
-        .await
-    {
-        Ok(mut file) => {
-            let mut config_contents = vec![];
-            file.read_to_end(&mut config_contents).await?;
-
-            serde_json::from_slice(&config_contents)?
-        }
-        Err(err) => match err.kind() {
-            std::io::ErrorKind::NotFound => {
-                config_needs_setup = true;
-
-                println!("Config not found, setting up default config");
-
-                WalletConfig::default()
-            }
-            _ => anyhow::bail!("IO error {err:#?}"),
-        },
-    };
-
-    if config_needs_setup {
-        tokio::fs::create_dir_all(&config_home).await?;
-
-        println!("Created configs dir at path {config_home:#?}");
-
-        let mut file = tokio::fs::OpenOptions::new()
-            .write(true)
-            .create(true)
-            .truncate(true)
-            .open(config_home.join("wallet_config.json"))
-            .await?;
-
-        let default_config_serialized =
-            serde_json::to_vec_pretty(&WalletConfig::default()).unwrap();
-
-        file.write_all(&default_config_serialized).await?;
-
-        println!("Configs setted up");
-    }
-
-    Ok(config)
+/// Fetch config path from default home
+pub fn fetch_config_path() -> Result<PathBuf> {
+    let home = get_home()?;
+    let config_path = home.join("wallet_config.json");
+    Ok(config_path)
 }
 
-/// Parse CLI auth string and merge with config auth, prioritizing CLI
-pub fn merge_auth_config(
-    mut config: WalletConfig,
-    cli_auth: Option<String>,
-) -> Result<WalletConfig> {
-    if let Some(auth_str) = cli_auth {
-        let cli_auth_config: BasicAuth = auth_str.parse()?;
-
-        if config.basic_auth.is_some() {
-            println!("Warning: CLI auth argument takes precedence over config basic-auth");
-        }
-
-        config.basic_auth = Some(cli_auth_config);
-    }
-    Ok(config)
-}
-
-/// Fetch data stored at home
+/// Fetch path to data storage from default home
 ///
 /// File must be created through setup beforehand.
-pub async fn fetch_persistent_storage() -> Result<PersistentStorage> {
+pub fn fetch_persistent_storage_path() -> Result<PathBuf> {
     let home = get_home()?;
     let accs_path = home.join("storage.json");
-    let mut storage_content = vec![];
-
-    match tokio::fs::File::open(accs_path).await {
-        Ok(mut file) => {
-            file.read_to_end(&mut storage_content).await?;
-            Ok(serde_json::from_slice(&storage_content)?)
-        }
-        Err(err) => match err.kind() {
-            std::io::ErrorKind::NotFound => {
-                anyhow::bail!("Not found, please setup roots from config command beforehand");
-            }
-            _ => {
-                anyhow::bail!("IO error {err:#?}");
-            }
-        },
-    }
+    Ok(accs_path)
 }
 
 /// Produces data for storage

wallet/src/lib.rs

@@ -1,6 +1,6 @@
 use std::{path::PathBuf, sync::Arc};
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
 use chain_storage::WalletChainStore;
 use common::{
@@ -25,10 +25,8 @@ pub use privacy_preserving_tx::PrivacyPreservingAccount;
 use tokio::io::AsyncWriteExt;
 
 use crate::{
-    config::PersistentStorage,
-    helperfunctions::{
-        fetch_persistent_storage, get_home, produce_data_for_storage, produce_random_nonces,
-    },
+    config::{PersistentStorage, WalletConfigOverrides},
+    helperfunctions::{produce_data_for_storage, produce_random_nonces},
     poller::TxPoller,
 };
 
@@ -124,91 +122,133 @@ impl TokenHolding {
 }
 
 pub struct WalletCore {
-    pub storage: WalletChainStore,
-    pub poller: TxPoller,
+    config_path: PathBuf,
+    storage: WalletChainStore,
+    storage_path: PathBuf,
+    poller: TxPoller,
+    // TODO: Make all fields private
     pub sequencer_client: Arc<SequencerClient>,
     pub last_synced_block: u64,
 }
 
 impl WalletCore {
-    pub async fn start_from_config_update_chain(config: WalletConfig) -> Result<Self> {
-        let basic_auth = config
-            .basic_auth
-            .as_ref()
-            .map(|auth| (auth.username.clone(), auth.password.clone()));
-        let client = Arc::new(SequencerClient::new_with_auth(
-            config.sequencer_addr.clone(),
-            basic_auth,
-        )?);
-        let tx_poller = TxPoller::new(config.clone(), client.clone());
+    /// Construct wallet using [`HOME_DIR_ENV_VAR`] env var for paths or user home dir if not set.
+    pub fn from_env() -> Result<Self> {
+        let config_path = helperfunctions::fetch_config_path()?;
+        let storage_path = helperfunctions::fetch_persistent_storage_path()?;
 
+        Self::new_update_chain(config_path, storage_path, None)
+    }
+
+    pub fn new_update_chain(
+        config_path: PathBuf,
+        storage_path: PathBuf,
+        config_overrides: Option<WalletConfigOverrides>,
+    ) -> Result<Self> {
         let PersistentStorage {
             accounts: persistent_accounts,
             last_synced_block,
-        } = fetch_persistent_storage().await?;
+        } = PersistentStorage::from_path(&storage_path)
+            .with_context(|| format!("Failed to read persistent storage at {storage_path:#?}"))?;
 
-        let storage = WalletChainStore::new(config, persistent_accounts)?;
-
-        Ok(Self {
-            storage,
-            poller: tx_poller,
-            sequencer_client: client.clone(),
+        Self::new(
+            config_path,
+            storage_path,
+            config_overrides,
+            |config| WalletChainStore::new(config, persistent_accounts),
             last_synced_block,
-        })
+        )
     }
 
-    pub async fn start_from_config_new_storage(
-        config: WalletConfig,
+    pub fn new_init_storage(
+        config_path: PathBuf,
+        storage_path: PathBuf,
+        config_overrides: Option<WalletConfigOverrides>,
         password: String,
     ) -> Result<Self> {
+        Self::new(
+            config_path,
+            storage_path,
+            config_overrides,
+            |config| WalletChainStore::new_storage(config, password),
+            0,
+        )
+    }
+
+    fn new(
+        config_path: PathBuf,
+        storage_path: PathBuf,
+        config_overrides: Option<WalletConfigOverrides>,
+        storage_ctor: impl FnOnce(WalletConfig) -> Result<WalletChainStore>,
+        last_synced_block: u64,
+    ) -> Result<Self> {
+        let mut config = WalletConfig::from_path_or_initialize_default(&config_path)
+            .with_context(|| format!("Failed to deserialize wallet config at {config_path:#?}"))?;
+        if let Some(config_overrides) = config_overrides {
+            config.apply_overrides(config_overrides);
+        }
+
         let basic_auth = config
             .basic_auth
             .as_ref()
             .map(|auth| (auth.username.clone(), auth.password.clone()));
-        let client = Arc::new(SequencerClient::new_with_auth(
+        let sequencer_client = Arc::new(SequencerClient::new_with_auth(
             config.sequencer_addr.clone(),
             basic_auth,
         )?);
-        let tx_poller = TxPoller::new(config.clone(), client.clone());
+        let tx_poller = TxPoller::new(config.clone(), Arc::clone(&sequencer_client));
 
-        let storage = WalletChainStore::new_storage(config, password)?;
+        let storage = storage_ctor(config)?;
 
         Ok(Self {
+            config_path,
+            storage_path,
             storage,
             poller: tx_poller,
-            sequencer_client: client.clone(),
-            last_synced_block: 0,
+            sequencer_client,
+            last_synced_block,
         })
     }
 
-    /// Store persistent data at home
-    pub async fn store_persistent_data(&self) -> Result<PathBuf> {
-        let home = get_home()?;
-        let storage_path = home.join("storage.json");
+    /// Get configuration with applied overrides
+    pub fn config(&self) -> &WalletConfig {
+        &self.storage.wallet_config
+    }
 
-        let data = produce_data_for_storage(&self.storage.user_data, self.last_synced_block);
-        let storage = serde_json::to_vec_pretty(&data)?;
+    /// Get storage
+    pub fn storage(&self) -> &WalletChainStore {
+        &self.storage
+    }
 
-        let mut storage_file = tokio::fs::File::create(storage_path.as_path()).await?;
-        storage_file.write_all(&storage).await?;
-
-        info!("Stored data at {storage_path:#?}");
-
-        Ok(storage_path)
+    /// Reset storage
+    pub fn reset_storage(&mut self, password: String) -> Result<()> {
+        self.storage = WalletChainStore::new_storage(self.storage.wallet_config.clone(), password)?;
+        Ok(())
     }
 
     /// Store persistent data at home
-    pub async fn store_config_changes(&self) -> Result<PathBuf> {
-        let home = get_home()?;
-        let config_path = home.join("wallet_config.json");
+    pub async fn store_persistent_data(&self) -> Result<()> {
+        let data = produce_data_for_storage(&self.storage.user_data, self.last_synced_block);
+        let storage = serde_json::to_vec_pretty(&data)?;
+
+        let mut storage_file = tokio::fs::File::create(&self.storage_path).await?;
+        storage_file.write_all(&storage).await?;
+
+        println!("Stored persistent accounts at {:#?}", self.storage_path);
+
+        Ok(())
+    }
+
+    /// Store persistent data at home
+    pub async fn store_config_changes(&self) -> Result<()> {
         let config = serde_json::to_vec_pretty(&self.storage.wallet_config)?;
 
-        let mut config_file = tokio::fs::File::create(config_path.as_path()).await?;
+        let mut config_file = tokio::fs::File::create(&self.config_path).await?;
         config_file.write_all(&config).await?;
 
-        info!("Stored data at {config_path:#?}");
+        info!("Stored data at {:#?}", self.config_path);
 
-        Ok(config_path)
+        Ok(())
     }
 
     pub fn create_new_account_public(
@@ -335,7 +375,7 @@ impl WalletCore {
     pub async fn send_privacy_preserving_tx(
         &self,
         accounts: Vec<PrivacyPreservingAccount>,
-        instruction_data: &InstructionData,
+        instruction_data: InstructionData,
         program: &ProgramWithDependencies,
     ) -> Result<(SendTxResponse, Vec<SharedSecretKey>), ExecutionFailureKind> {
         self.send_privacy_preserving_tx_with_pre_check(accounts, instruction_data, program, |_| {
@@ -347,7 +387,7 @@ impl WalletCore {
     pub async fn send_privacy_preserving_tx_with_pre_check(
         &self,
         accounts: Vec<PrivacyPreservingAccount>,
-        instruction_data: &InstructionData,
+        instruction_data: InstructionData,
         program: &ProgramWithDependencies,
         tx_pre_check: impl FnOnce(&[&Account]) -> Result<(), ExecutionFailureKind>,
     ) -> Result<(SendTxResponse, Vec<SharedSecretKey>), ExecutionFailureKind> {
@@ -363,16 +403,16 @@ impl WalletCore {
 
         let private_account_keys = acc_manager.private_account_keys();
         let (output, proof) = nssa::privacy_preserving_transaction::circuit::execute_and_prove(
-            &pre_states,
+            pre_states,
             instruction_data,
-            acc_manager.visibility_mask(),
-            &produce_random_nonces(private_account_keys.len()),
-            &private_account_keys
+            acc_manager.visibility_mask().to_vec(),
+            produce_random_nonces(private_account_keys.len()),
+            private_account_keys
                 .iter()
                 .map(|keys| (keys.npk.clone(), keys.ssk))
                 .collect::<Vec<_>>(),
-            &acc_manager.private_account_auth(),
-            &acc_manager.private_account_membership_proofs(),
+            acc_manager.private_account_auth(),
+            acc_manager.private_account_membership_proofs(),
             &program.to_owned(),
         )
         .unwrap();

wallet/src/main.rs

@@ -1,36 +1,65 @@
-use anyhow::Result;
+use anyhow::{Context as _, Result};
 use clap::{CommandFactory as _, Parser as _};
-use tokio::runtime::Builder;
-use wallet::cli::{Args, execute_continuous_run_with_auth, execute_subcommand_with_auth};
-
-pub const NUM_THREADS: usize = 2;
+use wallet::{
+    WalletCore,
+    cli::{Args, execute_continuous_run, execute_subcommand, read_password_from_stdin},
+    config::WalletConfigOverrides,
+    helperfunctions::{fetch_config_path, fetch_persistent_storage_path},
+};
 
 // TODO #169: We have sample configs for sequencer, but not for wallet
 // TODO #168: Why it requires config as a directory? Maybe better to deduce directory from config
 // file path?
 // TODO #172: Why it requires config as env var while sequencer_runner accepts as
 // argument?
-fn main() -> Result<()> {
-    let runtime = Builder::new_multi_thread()
-        .worker_threads(NUM_THREADS)
-        .enable_all()
-        .build()
-        .unwrap();
-
-    let args = Args::parse();
+#[tokio::main]
+async fn main() -> Result<()> {
+    let Args {
+        continuous_run,
+        auth,
+        command,
+    } = Args::parse();
 
     env_logger::init();
 
-    runtime.block_on(async move {
-        if let Some(command) = args.command {
-            let _output = execute_subcommand_with_auth(command, args.auth).await?;
-            Ok(())
-        } else if args.continuous_run {
-            execute_continuous_run_with_auth(args.auth).await
+    let config_path = fetch_config_path().context("Could not fetch config path")?;
+    let storage_path =
+        fetch_persistent_storage_path().context("Could not fetch persistent storage path")?;
+
+    // Override basic auth if provided via CLI
+    let config_overrides = WalletConfigOverrides {
+        basic_auth: auth.map(|auth| auth.parse()).transpose()?.map(Some),
+        ..Default::default()
+    };
+
+    if let Some(command) = command {
+        let mut wallet = if !storage_path.exists() {
+            // TODO: Maybe move to `WalletCore::from_env()` or similar?
+
+            println!("Persistent storage not found, need to execute setup");
+
+            let password = read_password_from_stdin()?;
+            let wallet = WalletCore::new_init_storage(
+                config_path,
+                storage_path,
+                Some(config_overrides),
+                password,
+            )?;
+
+            wallet.store_persistent_data().await?;
+            wallet
         } else {
-            let help = Args::command().render_long_help();
-            println!("{help}");
-            Ok(())
-        }
-    })
+            WalletCore::new_update_chain(config_path, storage_path, Some(config_overrides))?
+        };
+        let _output = execute_subcommand(&mut wallet, command).await?;
+        Ok(())
+    } else if continuous_run {
+        let mut wallet =
+            WalletCore::new_update_chain(config_path, storage_path, Some(config_overrides))?;
+        execute_continuous_run(&mut wallet).await
+    } else {
+        let help = Args::command().render_long_help();
+        println!("{help}");
+        Ok(())
+    }
 }

wallet/src/program_facades/native_token_transfer/deshielded.rs

@@ -19,7 +19,7 @@ impl NativeTokenTransfer<'_> {
                     PrivacyPreservingAccount::PrivateOwned(from),
                     PrivacyPreservingAccount::Public(to),
                 ],
-                &instruction_data,
+                instruction_data,
                 &program.into(),
                 tx_pre_check,
             )

wallet/src/program_facades/native_token_transfer/private.rs

@@ -15,11 +15,10 @@ impl NativeTokenTransfer<'_> {
         let instruction: u128 = 0;
 
         self.0
-            .send_privacy_preserving_tx_with_pre_check(
+            .send_privacy_preserving_tx(
                 vec![PrivacyPreservingAccount::PrivateOwned(from)],
-                &Program::serialize_instruction(instruction).unwrap(),
+                Program::serialize_instruction(instruction).unwrap(),
                 &Program::authenticated_transfer_program().into(),
-                |_| Ok(()),
             )
             .await
             .map(|(resp, secrets)| {
@@ -47,7 +46,7 @@ impl NativeTokenTransfer<'_> {
                         ipk: to_ipk,
                     },
                 ],
-                &instruction_data,
+                instruction_data,
                 &program.into(),
                 tx_pre_check,
             )
@@ -74,7 +73,7 @@ impl NativeTokenTransfer<'_> {
                     PrivacyPreservingAccount::PrivateOwned(from),
                     PrivacyPreservingAccount::PrivateOwned(to),
                 ],
-                &instruction_data,
+                instruction_data,
                 &program.into(),
                 tx_pre_check,
             )

wallet/src/program_facades/native_token_transfer/shielded.rs

@@ -20,7 +20,7 @@ impl NativeTokenTransfer<'_> {
                     PrivacyPreservingAccount::Public(from),
                     PrivacyPreservingAccount::PrivateOwned(to),
                 ],
-                &instruction_data,
+                instruction_data,
                 &program.into(),
                 tx_pre_check,
             )
@@ -52,7 +52,7 @@ impl NativeTokenTransfer<'_> {
                         ipk: to_ipk,
                     },
                 ],
-                &instruction_data,
+                instruction_data,
                 &program.into(),
                 tx_pre_check,
             )