feat: schedule automatic undelegation on AML check failure

magicblock-account-cloner/src/lib.rs

@@ -214,6 +214,12 @@ impl ChainlinkCloner {
         InstructionUtils::set_program_authority_instruction(program, authority)
     }
 
+    fn schedule_undelegation_ix(pubkey: Pubkey) -> Instruction {
+        InstructionUtils::schedule_cloned_account_undelegation_instruction(
+            pubkey,
+        )
+    }
+
     // -----------------
     // Clone Fields Helper
     // -----------------
@@ -240,11 +246,16 @@ impl ChainlinkCloner {
         let fields = Self::clone_fields(request);
         let actions: Vec<Instruction> =
             request.delegation_actions.clone().into();
+        let clone_actions = if request.needs_undelegation {
+            Vec::new()
+        } else {
+            actions.clone()
+        };
         let clone_ix = Self::clone_ix(
             request.pubkey,
             request.account.data().to_vec(),
             fields,
-            actions.clone(),
+            clone_actions,
         );
 
         // TODO(#625): Re-enable frequency commits when proper limits are in place:
@@ -255,7 +266,9 @@ impl ChainlinkCloner {
         // To re-enable, uncomment the following and use `ixs` instead of `[clone_ix]`:
         // let ixs = self.maybe_add_crank_commits_ix(request, clone_ix);
         let mut ixs = vec![clone_ix];
-        if !actions.is_empty() {
+        if request.needs_undelegation {
+            ixs.push(Self::schedule_undelegation_ix(request.pubkey));
+        } else if !actions.is_empty() {
             ixs.push(Self::post_delegation_action_ix(request.pubkey, actions));
         }
 
@@ -314,14 +327,20 @@ impl ChainlinkCloner {
 
         let actions: Vec<Instruction> =
             request.delegation_actions.clone().into();
+        let clone_actions = if request.needs_undelegation {
+            Vec::new()
+        } else {
+            actions.clone()
+        };
 
         // Continue txs for remaining chunks
         let mut offset = MAX_INLINE_DATA_SIZE;
         while offset < data.len() {
             let end = (offset + MAX_INLINE_DATA_SIZE).min(data.len());
             let chunk = data[offset..end].to_vec();
             let is_last = end == data.len();
-            let final_without_actions = is_last && actions.is_empty();
+            let final_without_actions =
+                is_last && actions.is_empty() && !request.needs_undelegation;
 
             let continue_ix = Self::clone_continue_ix(
                 request.pubkey,
@@ -334,19 +353,20 @@ impl ChainlinkCloner {
             offset = end;
         }
 
-        if !actions.is_empty() {
+        if request.needs_undelegation || !actions.is_empty() {
             let continue_ix = Self::clone_continue_ix(
                 request.pubkey,
                 data.len() as u32,
                 Vec::new(),
                 true,
-                actions.clone(),
-            );
-            let action_ix =
-                Self::post_delegation_action_ix(request.pubkey, actions);
-            txs.push(
-                self.create_signed_tx(&[continue_ix, action_ix], blockhash),
+                clone_actions,
             );
+            let ix = if request.needs_undelegation {
+                Self::schedule_undelegation_ix(request.pubkey)
+            } else {
+                Self::post_delegation_action_ix(request.pubkey, actions)
+            };
+            txs.push(self.create_signed_tx(&[continue_ix, ix], blockhash));
         }
 
         txs
@@ -740,6 +760,7 @@ mod tests {
             commit_frequency_ms: None,
             delegation_actions: DelegationActions::from(actions),
             delegated_to_other: None,
+            needs_undelegation: false,
         }
     }
 
@@ -849,6 +870,11 @@ mod tests {
                 assert_eq!(executor_pubkey, pubkey);
                 assert_eq!(executor_actions, actions);
             }
+            PostDelegationActionExecutorInstruction::ScheduleUndelegation {
+                ..
+            } => {
+                panic!("expected execute instruction")
+            }
         }
     }
 
@@ -895,6 +921,11 @@ mod tests {
                 assert_eq!(executor_pubkey, pubkey);
                 assert_eq!(executor_actions, actions);
             }
+            PostDelegationActionExecutorInstruction::ScheduleUndelegation {
+                ..
+            } => {
+                panic!("expected execute instruction")
+            }
         }
     }
 }

magicblock-chainlink/src/chainlink/fetch_cloner/ata_projection.rs

@@ -217,6 +217,7 @@ where
         commit_frequency_ms: None,
         delegation_actions: delegation_actions.clone(),
         delegated_to_other: None,
+        needs_undelegation: false,
     })
 }
 
@@ -593,6 +594,7 @@ where
             commit_frequency_ms,
             delegation_actions: actions.unwrap_or_default(),
             delegated_to_other,
+            needs_undelegation: false,
         });
     }
 

magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs

@@ -516,12 +516,24 @@ where
             ));
         }
 
-        self.ensure_delegation_action_dependencies(
-            request.pubkey,
-            request.account.remote_slot(),
-            &request.delegation_actions,
-        )
-        .await?;
+        match self
+            .ensure_delegation_action_dependencies(
+                request.pubkey,
+                request.account.remote_slot(),
+                &request.delegation_actions,
+            )
+            .await
+        {
+            Ok(()) => {}
+            Err(ChainlinkError::RangeRisk(
+                magicblock_aml::RiskError::HighRiskAddresses(_),
+            )) => {
+                request.needs_undelegation = true;
+            }
+            Err(err) => {
+                return Err(err);
+            }
+        }
 
         Ok(self.clone_account_with_ownership(request).await?)
     }
@@ -887,6 +899,7 @@ where
                         commit_frequency_ms,
                         delegation_actions: raw_delegation_actions,
                         delegated_to_other,
+                        needs_undelegation: false,
                     },
                 )
                 .await
@@ -2600,6 +2613,7 @@ where
                 commit_frequency_ms: None,
                 delegation_actions: DelegationActions::default(),
                 delegated_to_other: None,
+                needs_undelegation: false,
             })
             .await?;
         Ok(())

magicblock-chainlink/src/chainlink/fetch_cloner/pipeline.rs

@@ -134,6 +134,7 @@ fn classify_single_account(
                             commit_frequency_ms: None,
                             delegation_actions: DelegationActions::default(),
                             delegated_to_other: None,
+                            needs_undelegation: false,
                         });
                     }
                 }
@@ -378,6 +379,7 @@ where
                 commit_frequency_ms,
                 delegation_actions,
                 delegated_to_other,
+                needs_undelegation: false,
             });
             if cleanup_delegated_subscription {
                 if cleanup_undelegation_tracking {

magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs

@@ -4454,6 +4454,7 @@ async fn test_post_delegation_actions_reject_non_delegated_clone_target() {
                 commit_frequency_ms: None,
                 delegation_actions: actions,
                 delegated_to_other: None,
+                needs_undelegation: false,
             },
         )
         .await
@@ -4505,6 +4506,7 @@ async fn test_dlp_owned_clone_without_actions_clears_stale_delegated_flag() {
                 commit_frequency_ms: None,
                 delegation_actions: DelegationActions::default(),
                 delegated_to_other: None,
+                needs_undelegation: false,
             },
         )
         .await
@@ -4556,6 +4558,7 @@ async fn test_dlp_owned_magic_fee_vault_without_actions_remains_delegated() {
                 commit_frequency_ms: None,
                 delegation_actions: DelegationActions::default(),
                 delegated_to_other: None,
+                needs_undelegation: false,
             },
         )
         .await
@@ -4643,6 +4646,7 @@ async fn test_post_delegation_actions_refresh_writable_dependency_before_target(
                 commit_frequency_ms: None,
                 delegation_actions: actions,
                 delegated_to_other: None,
+                needs_undelegation: false,
             },
         )
         .await
@@ -4718,6 +4722,7 @@ async fn test_post_delegation_actions_execute_once_across_remote_slots() {
                     commit_frequency_ms: None,
                     delegation_actions: actions.clone(),
                     delegated_to_other: None,
+                    needs_undelegation: false,
                 },
             )
             .await
@@ -4788,6 +4793,7 @@ async fn test_delegated_clone_does_not_override_active_local_target() {
                 commit_frequency_ms: None,
                 delegation_actions: DelegationActions::default(),
                 delegated_to_other: None,
+                needs_undelegation: false,
             },
         )
         .await

magicblock-chainlink/src/chainlink/mod.rs

@@ -221,7 +221,7 @@ impl<T: ChainRpcClient, U: ChainPubsubClient, V: AccountsBank, C: Cloner>
         accounts_bank,
         cloner,
         config,
-        chainlink_config
+        chainlink_config,
     ))]
     pub async fn try_new_from_endpoints(
         endpoints: &Endpoints,

magicblock-chainlink/src/cloner/errors.rs

@@ -40,4 +40,7 @@ pub enum ClonerError {
 
     #[error("Failed to evict account {0} : {1:?}")]
     FailedToEvictAccount(Pubkey, Box<ClonerError>),
+
+    #[error("Failed to schedule undelegation {0} : {1:?}")]
+    FailedToScheduleUndelegation(Pubkey, Box<ClonerError>),
 }

magicblock-chainlink/src/cloner/mod.rs

@@ -59,6 +59,9 @@ pub struct AccountCloneRequest {
     /// this contains that validator's pubkey. None if account is not
     /// delegated to another validator.
     pub delegated_to_other: Option<Pubkey>,
+    /// Account that need to be undelegated (e.g. due to AML risk) after cloning.
+    /// Is only true if there are actions with risky signers.
+    pub needs_undelegation: bool,
 }
 
 #[async_trait]

magicblock-chainlink/tests/utils/test_context.rs

@@ -5,6 +5,7 @@ use std::{
     time::{Duration, Instant},
 };
 
+use magicblock_aml::RiskService;
 use magicblock_chainlink::{
     accounts_bank::mock::AccountsBankStub,
     errors::ChainlinkResult,
@@ -55,6 +56,13 @@ pub struct TestContext {
 
 impl TestContext {
     pub async fn init(slot: Slot) -> Self {
+        Self::init_with_services(slot, None).await
+    }
+
+    pub async fn init_with_services(
+        slot: Slot,
+        risk_service: Option<Arc<RiskService>>,
+    ) -> Self {
         let (rpc_client, pubsub_client) = {
             let rpc_client =
                 ChainRpcClientMockBuilder::new().slot(slot).build();
@@ -102,7 +110,7 @@ impl TestContext {
                             validator_keypair.insecure_clone(),
                             rx,
                             None,
-                            None,
+                            risk_service,
                         )),
                         Some(provider),
                     )

magicblock-magic-program-api/src/instruction.rs

@@ -373,4 +373,14 @@ pub enum PostDelegationActionExecutorInstruction {
         cloned_account_pubkey: Pubkey,
         actions: Vec<Instruction>,
     },
+
+    /// Schedules undelegation immediately after a matching delegated clone
+    /// instruction in the same transaction.
+    ///
+    /// # Account references
+    /// - **0.**   `[SIGNER]`  Validator authority
+    /// - **1.**   `[]`        Delegated clone target
+    /// - **2.**   `[]`        Instructions sysvar
+    /// - **3.**   `[WRITE]`   Magic Context account
+    ScheduleUndelegation { cloned_account_pubkey: Pubkey },
 }