feat: schedule automatic undelegation on AML check failure

magicblock-api/Cargo.toml

@@ -9,6 +9,7 @@ edition.workspace = true
 
 [dependencies]
 anyhow = { workspace = true }
+async-trait = { workspace = true }
 borsh = "1.5.3"
 fd-lock = { workspace = true }
 tracing = { workspace = true }

magicblock-api/src/magic_validator.rs

@@ -19,8 +19,10 @@ use magicblock_aperture::{
     state::{NodeContext, SharedState},
 };
 use magicblock_chainlink::{
-    config::ChainlinkConfig, remote_account_provider::Endpoints, ProdChainlink,
-    ProdInnerChainlink,
+    config::ChainlinkConfig,
+    fetch_cloner::{UndelegationScheduleRequest, UndelegationScheduler},
+    remote_account_provider::Endpoints,
+    ProdChainlink, ProdInnerChainlink,
 };
 use magicblock_committor_service::{
     config::ChainConfig, BaseIntentCommittor, CommittorService,
@@ -99,6 +101,29 @@ use crate::{
 type InnerChainlinkImpl = ProdInnerChainlink<ChainlinkCloner>;
 
 type ChainlinkImpl = ProdChainlink<ChainlinkCloner>;
+/// Bridges chainlink's [`UndelegationScheduler`] to the committor service so a
+/// delegated clone rejected by AML is undelegated on the base layer.
+struct CommittorUndelegationScheduler(Arc<CommittorService>);
+
+#[async_trait::async_trait]
+impl UndelegationScheduler for CommittorUndelegationScheduler {
+    async fn schedule_undelegation(
+        &self,
+        request: UndelegationScheduleRequest,
+    ) -> magicblock_chainlink::errors::ChainlinkResult<()> {
+        let pubkey = request.pubkey;
+        self.0
+            .schedule_undelegation(pubkey, request.account)
+            .await
+            .map_err(|err| format!("committor response channel closed: {err}"))
+            .and_then(|result| result.map_err(|err| err.to_string()))
+            .map_err(|message| {
+                magicblock_chainlink::errors::ChainlinkError::FailedToScheduleUndelegationAfterAmlRejection(
+                    pubkey, message,
+                )
+            })
+    }
+}
 
 // -----------------
 // MagicValidator
@@ -238,6 +263,7 @@ impl MagicValidator {
                 &ledger.latest_block().clone(),
                 &accountsdb,
                 shared_chain_slot.clone(),
+                committor_service.clone(),
             )
             .await?,
         );
@@ -495,6 +521,7 @@ impl MagicValidator {
         latest_block: &LatestBlock,
         accountsdb: &Arc<AccountsDb>,
         chain_slot: Option<Arc<AtomicU64>>,
+        committor_service: Option<Arc<CommittorService>>,
     ) -> ApiResult<ChainlinkImpl> {
         if Self::replication_mode_uses_disabled_chainlink(
             &config.validator.replication_mode,
@@ -550,6 +577,10 @@ impl MagicValidator {
             &config.chainlink,
             config.storage.as_path(),
             chain_slot.unwrap_or_default(),
+            committor_service.map(|committor_service| {
+                Arc::new(CommittorUndelegationScheduler(committor_service))
+                    as Arc<dyn UndelegationScheduler>
+            }),
         )
         .await?;
 

magicblock-chainlink/Cargo.toml

@@ -58,6 +58,7 @@ url = { workspace = true }
 [dev-dependencies]
 assert_matches = { workspace = true }
 magicblock-chainlink = { path = ".", features = ["dev-context"] }
+tempfile = { workspace = true }
 
 [features]
 default = []

magicblock-chainlink/src/chainlink/errors.rs

@@ -67,6 +67,11 @@ pub enum ChainlinkError {
     #[error("Failed to perform Range risk check: {0}")]
     RangeRisk(#[from] RiskError),
 
+    #[error(
+        "Failed to schedule undelegation for {0} after AML rejection: {1}"
+    )]
+    FailedToScheduleUndelegationAfterAmlRejection(Pubkey, String),
+
     #[error("Chainlink is disabled for non-primary mode")]
     DisabledForNonPrimaryMode,
 }

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

@@ -8,6 +8,7 @@ use std::{
     time::Duration,
 };
 
+use async_trait::async_trait;
 use dlp_api::{
     pda::delegation_record_pda_from_delegated_account, state::DelegationRecord,
 };
@@ -85,6 +86,23 @@ use crate::{
     },
 };
 
+#[derive(Clone)]
+pub struct UndelegationScheduleRequest {
+    pub pubkey: Pubkey,
+    pub account: AccountSharedData,
+}
+
+/// Schedules an undelegation when a delegated clone is rejected by AML before
+/// it enters the local bank. Implemented outside chainlink (e.g. by the API
+/// layer bridging to the committor service) to keep chainlink decoupled.
+#[async_trait]
+pub trait UndelegationScheduler: Send + Sync {
+    async fn schedule_undelegation(
+        &self,
+        request: UndelegationScheduleRequest,
+    ) -> ChainlinkResult<()>;
+}
+
 pub struct FetchCloner<T, U, V, C>
 where
     T: ChainRpcClient,
@@ -138,6 +156,10 @@ where
 
     /// Risk checker for post-delegation action addresses.
     risk_service: Option<Arc<RiskService>>,
+
+    /// Schedules undelegation when post-delegation action AML checks reject
+    /// a delegated clone before it enters the local bank.
+    undelegation_scheduler: Option<Arc<dyn UndelegationScheduler>>,
 }
 
 /// Negative-cache capacity for known-empty eATAs.
@@ -177,6 +199,7 @@ where
                 .pending_operation_timeout_ms
                 .clone(),
             risk_service: self.risk_service.clone(),
+            undelegation_scheduler: self.undelegation_scheduler.clone(),
         }
     }
 }
@@ -198,6 +221,7 @@ where
         subscription_updates_rx: mpsc::Receiver<ForwardedSubscriptionUpdate>,
         allowed_programs: Option<Vec<AllowedProgram>>,
         risk_service: Option<Arc<RiskService>>,
+        undelegation_scheduler: Option<Arc<dyn UndelegationScheduler>>,
     ) -> Arc<Self> {
         let validator_pubkey = validator_keypair.pubkey();
         let blacklisted_accounts = blacklisted_accounts(&validator_pubkey);
@@ -225,6 +249,7 @@ where
                 FETCH_CLONE_OPERATION_TIMEOUT.as_millis() as u64,
             )),
             risk_service,
+            undelegation_scheduler,
         });
 
         let accounts_bank_for_eviction = accounts_bank.clone();
@@ -510,16 +535,49 @@ where
             ));
         }
 
-        self.ensure_delegation_action_dependencies(
-            request.pubkey,
-            request.account.remote_slot(),
-            &request.delegation_actions,
-        )
-        .await?;
+        if let Err(err) = self
+            .ensure_delegation_action_dependencies(
+                request.pubkey,
+                request.account.remote_slot(),
+                &request.delegation_actions,
+            )
+            .await
+        {
+            if matches!(
+                err,
+                ChainlinkError::RangeRisk(
+                    magicblock_aml::RiskError::HighRiskAddresses(_)
+                )
+            ) {
+                self.schedule_undelegation_after_aml_rejection(&request)
+                    .await?;
+            }
+            return Err(err);
+        }
 
         Ok(self.clone_account_with_ownership(request).await?)
     }
 
+    async fn schedule_undelegation_after_aml_rejection(
+        &self,
+        request: &AccountCloneRequest,
+    ) -> ChainlinkResult<()> {
+        let Some(scheduler) = self.undelegation_scheduler.as_ref() else {
+            warn!(
+                pubkey = %request.pubkey,
+                "AML rejected post-delegation actions but undelegation scheduler is unavailable"
+            );
+            return Ok(());
+        };
+
+        scheduler
+            .schedule_undelegation(UndelegationScheduleRequest {
+                pubkey: request.pubkey,
+                account: request.account.clone(),
+            })
+            .await
+    }
+
     fn normalize_unresolved_dlp_clone_request(
         &self,
         request: &mut AccountCloneRequest,