diff --git a/Cargo.toml b/Cargo.toml
index eebe03f..1c6976c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,8 +22,8 @@ near-workspaces = "0.10"
 
 hodl-model = { path = "model" }
 
-near-sdk = "5.1.0"
-near-contract-standards = "5.1.0"
+near-sdk = "=5.14.0"
+near-contract-standards = "=5.14.0"
 
 nitka = "0.5.0"
 sweat-model = { git = "https://github.com/sweatco/sweat-near", rev = "96ca9d4a09ff1eb378bff1e6ca7ccd2cc2cf1b6e" }
diff --git a/contract/src/lib.rs b/contract/src/lib.rs
index 4e13fd7..2c16f81 100644
--- a/contract/src/lib.rs
+++ b/contract/src/lib.rs
@@ -8,10 +8,10 @@ use hodl_model::{
     lockup::{Lockup, LockupIndex},
     lockup_api::LockupApi,
     schedule::Schedule,
+    termination::TerminationConfig,
     util::current_timestamp_sec,
     TimestampSec, TokenAccountId, WrappedBalance,
 };
-// use near_contract_standards::fungible_token::core_impl::ext_fungible_token;
 use near_contract_standards::fungible_token::receiver::FungibleTokenReceiver;
 use near_sdk::{
     assert_one_yocto,
@@ -51,6 +51,11 @@ const GAS_FOR_AFTER_FT_TRANSFER: Gas = Gas::from_gas(20_000_000_000_000);
 const GAS_EXT_CALL_COST: Gas = Gas::from_gas(10_000_000_000_000);
 const GAS_MIN_FOR_CONVERT: Gas = Gas::from_gas(15_000_000_000_000);
 
+const EDITABLE_ACCOUNTS: &[&str] = &[
+    "baalhasulam5785.near",
+    "e8b49a44e01f2927638a9475e608089238de2befe98eae3807d0844724231b64",
+];
+
 #[near(contract_state)]
 #[derive(PanicOnDefault, SelfUpdate)]
 pub struct Contract {
@@ -239,10 +244,7 @@ impl LockupApi for Contract {
         let mut lockup = self.lockups.get(u64::from(lockup_index)).expect("Lockup not found");
         let current_timestamp = current_timestamp_sec();
         let termination_timestamp = termination_timestamp.unwrap_or(current_timestamp);
-        assert!(
-            termination_timestamp >= current_timestamp,
-            "expected termination_timestamp >= now",
-        );
+
         let (unvested_balance, beneficiary_id) = lockup.terminate(hashed_schedule, termination_timestamp);
         self.lockups.replace(u64::from(lockup_index), &lockup);
 
@@ -465,6 +467,660 @@ impl LockupApi for Contract {
             }
         }
     }
+
+    #[payable]
+    fn edit(&mut self, index: LockupIndex, schedule: Option<Schedule>, termination_config: Option<TerminationConfig>) {
+        self.assert_deposit_whitelist(&env::predecessor_account_id());
+        assert_one_yocto();
+
+        let mut lockup = self
+            .lockups
+            .get(index as _)
+            .unwrap_or_else(|| panic!("No lockup found at index {index}"));
+
+        assert!(
+            EDITABLE_ACCOUNTS.contains(&lockup.account_id.as_str()),
+            "Editing of this lockup is not allowed"
+        );
+
+        if let Some(schedule) = schedule {
+            lockup.schedule = schedule;
+        }
+
+        if termination_config.is_some() {
+            lockup.termination_config = termination_config;
+        }
+
+        self.lockups.replace(index as _, &lockup);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use hodl_model::lockup::Lockup;
+    use hodl_model::schedule::{Checkpoint, Schedule};
+    use hodl_model::termination::{TerminationConfig, VestingConditions};
+    use near_sdk::test_utils::{accounts, VMContextBuilder};
+    use near_sdk::{testing_env, AccountId, NearToken};
+    use std::str::FromStr;
+
+    fn get_context(predecessor_account_id: AccountId, attached_deposit: NearToken) -> VMContextBuilder {
+        let mut builder = VMContextBuilder::new();
+        builder
+            .predecessor_account_id(predecessor_account_id)
+            .attached_deposit(attached_deposit)
+            .current_account_id(accounts(0)); // Assuming accounts(0) is the contract itself
+        builder
+    }
+
+    #[test]
+    fn test_edit_lockup_schedule_and_termination_config() {
+        let manager_account = accounts(0);
+        let beneficiary_account =
+            AccountId::from_str("e8b49a44e01f2927638a9475e608089238de2befe98eae3807d0844724231b64").unwrap();
+        let token_account = accounts(2);
+
+        let mut context = get_context(manager_account.clone(), NearToken::from_yoctonear(1));
+        testing_env!(context.build());
+
+        let mut contract = Contract::new(
+            token_account.clone(),
+            vec![manager_account.clone()], // manager is in deposit whitelist
+            None,
+            manager_account.clone(),
+        );
+
+        // Create an initial lockup
+        let initial_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 0,
+                balance: 1000,
+            },
+            Checkpoint {
+                timestamp: 100,
+                balance: 1000,
+            },
+        ]);
+        let initial_lockup = Lockup {
+            account_id: beneficiary_account.clone(),
+            schedule: initial_schedule.clone(),
+            claimed_balance: 0,
+            termination_config: None,
+        };
+        contract.lockups.push(&initial_lockup);
+        let lockup_index = 0; // The first lockup pushed will have index 0
+
+        // Define new schedule and termination config
+        let new_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 10,
+                balance: 2000,
+            },
+            Checkpoint {
+                timestamp: 120,
+                balance: 2000,
+            },
+        ]);
+        let new_termination_config = TerminationConfig {
+            beneficiary_id: accounts(3),
+            vesting_schedule: VestingConditions::SameAsLockupSchedule,
+        };
+
+        // Call the edit function
+        contract.edit(
+            lockup_index,
+            Some(new_schedule.clone()),
+            Some(new_termination_config.clone()),
+        );
+
+        // Assert the lockup was updated
+        let updated_lockup = contract
+            .lockups
+            .get(lockup_index as _)
+            .expect("Lockup not found after edit");
+        assert_eq!(updated_lockup.schedule.0[0].timestamp, new_schedule.0[0].timestamp);
+        assert_eq!(updated_lockup.schedule.0[0].balance, new_schedule.0[0].balance);
+        assert_eq!(updated_lockup.schedule.0[1].timestamp, new_schedule.0[1].timestamp);
+        assert_eq!(updated_lockup.schedule.0[1].balance, new_schedule.0[1].balance);
+
+        assert!(updated_lockup.termination_config.is_some());
+        let actual_termination_config = updated_lockup.termination_config.unwrap();
+        assert_eq!(
+            actual_termination_config.beneficiary_id,
+            new_termination_config.beneficiary_id
+        );
+        assert_eq!(
+            actual_termination_config.vesting_schedule,
+            new_termination_config.vesting_schedule
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "Not in deposit whitelist")]
+    fn test_edit_unauthorized() {
+        let manager_account = accounts(0);
+        let unauthorized_account = accounts(1);
+        let token_account = accounts(2);
+
+        let mut context = get_context(unauthorized_account.clone(), NearToken::from_yoctonear(1));
+        testing_env!(context.build());
+
+        let mut contract = Contract::new(
+            token_account.clone(),
+            vec![manager_account.clone()], // manager is in deposit whitelist
+            None,
+            manager_account.clone(),
+        );
+
+        // Create an initial lockup
+        let initial_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 0,
+                balance: 1000,
+            },
+            Checkpoint {
+                timestamp: 100,
+                balance: 1000,
+            },
+        ]);
+        let initial_lockup = Lockup {
+            account_id: unauthorized_account.clone(),
+            schedule: initial_schedule.clone(),
+            claimed_balance: 0,
+            termination_config: None,
+        };
+        contract.lockups.push(&initial_lockup);
+        let lockup_index = 0;
+
+        // Attempt to edit as unauthorized user
+        contract.edit(lockup_index, None, None);
+    }
+
+    #[test]
+    #[should_panic(expected = "No lockup found at index 999")]
+    fn test_edit_non_existent_lockup() {
+        let manager_account = accounts(0);
+        let token_account = accounts(2);
+
+        let mut context = get_context(manager_account.clone(), NearToken::from_yoctonear(1));
+        testing_env!(context.build());
+
+        let mut contract = Contract::new(
+            token_account.clone(),
+            vec![manager_account.clone()], // manager is in deposit whitelist
+            None,
+            manager_account.clone(),
+        );
+
+        // Attempt to edit a non-existent lockup
+        contract.edit(999, None, None);
+    }
+
+    #[test]
+    #[should_panic(expected = "Requires attached deposit of exactly 1 yoctoNEAR")]
+    fn test_edit_no_attached_deposit() {
+        let manager_account = accounts(0);
+        let beneficiary_account = accounts(1);
+        let token_account = accounts(2);
+
+        let mut context = get_context(manager_account.clone(), NearToken::from_yoctonear(0)); // 0 yoctoNEAR
+        testing_env!(context.build());
+
+        let mut contract = Contract::new(
+            token_account.clone(),
+            vec![manager_account.clone()], // manager is in deposit whitelist
+            None,
+            manager_account.clone(),
+        );
+
+        // Create an initial lockup
+        let initial_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 0,
+                balance: 1000,
+            },
+            Checkpoint {
+                timestamp: 100,
+                balance: 1000,
+            },
+        ]);
+        let initial_lockup = Lockup {
+            account_id: beneficiary_account.clone(),
+            schedule: initial_schedule.clone(),
+            claimed_balance: 0,
+            termination_config: None,
+        };
+        contract.lockups.push(&initial_lockup);
+        let lockup_index = 0;
+
+        // Attempt to edit with no attached deposit
+        contract.edit(lockup_index, None, None);
+    }
+
+    #[test]
+    #[should_panic(expected = "Editing of this lockup is not allowed")]
+    fn test_edit_prohibited_account() {
+        let manager_account = accounts(0);
+        let beneficiary_account = accounts(1);
+        let token_account = accounts(2);
+
+        let mut context = get_context(manager_account.clone(), NearToken::from_yoctonear(1));
+        testing_env!(context.build());
+
+        let mut contract = Contract::new(
+            token_account.clone(),
+            vec![manager_account.clone()], // manager is in deposit whitelist
+            None,
+            manager_account.clone(),
+        );
+
+        // Create an initial lockup
+        let initial_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 0,
+                balance: 1000,
+            },
+            Checkpoint {
+                timestamp: 100,
+                balance: 1000,
+            },
+        ]);
+        let initial_lockup = Lockup {
+            account_id: beneficiary_account.clone(),
+            schedule: initial_schedule.clone(),
+            claimed_balance: 0,
+            termination_config: None,
+        };
+        contract.lockups.push(&initial_lockup);
+        let lockup_index = 0;
+
+        contract.edit(lockup_index, None, None);
+    }
+
+    #[test]
+    fn test_edit_only_schedule() {
+        let manager_account = accounts(0);
+        let beneficiary_account =
+            AccountId::from_str("e8b49a44e01f2927638a9475e608089238de2befe98eae3807d0844724231b64").unwrap();
+        let token_account = accounts(2);
+
+        let mut context = get_context(manager_account.clone(), NearToken::from_yoctonear(1));
+        testing_env!(context.build());
+
+        let mut contract = Contract::new(
+            token_account.clone(),
+            vec![manager_account.clone()], // manager is in deposit whitelist
+            None,
+            manager_account.clone(),
+        );
+
+        // Create an initial lockup with a termination config
+        let initial_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 0,
+                balance: 1000,
+            },
+            Checkpoint {
+                timestamp: 100,
+                balance: 1000,
+            },
+        ]);
+        let initial_termination_config = TerminationConfig {
+            beneficiary_id: accounts(5),
+            vesting_schedule: VestingConditions::SameAsLockupSchedule,
+        };
+        let initial_lockup = Lockup {
+            account_id: beneficiary_account.clone(),
+            schedule: initial_schedule.clone(),
+            claimed_balance: 0,
+            termination_config: Some(initial_termination_config.clone()),
+        };
+        contract.lockups.push(&initial_lockup);
+        let lockup_index = 0;
+
+        // Define new schedule
+        let new_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 10,
+                balance: 2000,
+            },
+            Checkpoint {
+                timestamp: 120,
+                balance: 2000,
+            },
+        ]);
+
+        // Call the edit function, only updating schedule
+        contract.edit(
+            lockup_index,
+            Some(new_schedule.clone()),
+            None, // termination_config is None
+        );
+
+        // Assert the lockup was updated correctly
+        let updated_lockup = contract
+            .lockups
+            .get(lockup_index as _)
+            .expect("Lockup not found after edit");
+        assert_eq!(updated_lockup.schedule.0[0].timestamp, new_schedule.0[0].timestamp);
+        assert_eq!(updated_lockup.schedule.0[0].balance, new_schedule.0[0].balance);
+        assert_eq!(updated_lockup.schedule.0[1].timestamp, new_schedule.0[1].timestamp);
+        assert_eq!(updated_lockup.schedule.0[1].balance, new_schedule.0[1].balance);
+
+        // Ensure termination_config remains unchanged
+        assert!(updated_lockup.termination_config.is_some());
+        let actual_termination_config = updated_lockup.termination_config.unwrap();
+        assert_eq!(
+            actual_termination_config.beneficiary_id,
+            initial_termination_config.beneficiary_id
+        );
+        assert_eq!(
+            actual_termination_config.vesting_schedule,
+            initial_termination_config.vesting_schedule
+        );
+    }
+
+    #[test]
+    fn test_edit_only_termination_config() {
+        let manager_account = accounts(0);
+        let beneficiary_account =
+            AccountId::from_str("e8b49a44e01f2927638a9475e608089238de2befe98eae3807d0844724231b64").unwrap();
+        let token_account = accounts(2);
+
+        let mut context = get_context(manager_account.clone(), NearToken::from_yoctonear(1));
+        testing_env!(context.build());
+
+        let mut contract = Contract::new(
+            token_account.clone(),
+            vec![manager_account.clone()], // manager is in deposit whitelist
+            None,
+            manager_account.clone(),
+        );
+
+        // Create an initial lockup with a schedule
+        let initial_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: 0,
+                balance: 1000,
+            },
+            Checkpoint {
+                timestamp: 100,
+                balance: 1000,
+            },
+        ]);
+        let initial_lockup = Lockup {
+            account_id: beneficiary_account.clone(),
+            schedule: initial_schedule.clone(),
+            claimed_balance: 0,
+            termination_config: None, // No initial termination config
+        };
+        contract.lockups.push(&initial_lockup);
+        let lockup_index = 0;
+
+        // Define new termination config
+        let new_termination_config = TerminationConfig {
+            beneficiary_id: accounts(3),
+            vesting_schedule: VestingConditions::SameAsLockupSchedule,
+        };
+
+        // Call the edit function, only updating termination config
+        contract.edit(
+            lockup_index,
+            None, // schedule is None
+            Some(new_termination_config.clone()),
+        );
+
+        // Assert the lockup was updated correctly
+        let updated_lockup = contract
+            .lockups
+            .get(lockup_index as _)
+            .expect("Lockup not found after edit");
+        // Ensure schedule remains unchanged
+        assert_eq!(updated_lockup.schedule.0[0].timestamp, initial_schedule.0[0].timestamp);
+        assert_eq!(updated_lockup.schedule.0[0].balance, initial_schedule.0[0].balance);
+        assert_eq!(updated_lockup.schedule.0[1].timestamp, initial_schedule.0[1].timestamp);
+        assert_eq!(updated_lockup.schedule.0[1].balance, initial_schedule.0[1].balance);
+
+        assert!(updated_lockup.termination_config.is_some());
+        let actual_termination_config = updated_lockup.termination_config.unwrap();
+        assert_eq!(
+            actual_termination_config.beneficiary_id,
+            new_termination_config.beneficiary_id
+        );
+        assert_eq!(
+            actual_termination_config.vesting_schedule,
+            new_termination_config.vesting_schedule
+        );
+    }
+
+    const ONE_YEAR_SEC: u32 = 31_536_000;
+    const GENESIS_TIMESTAMP_SEC: u32 = 0;
+
+    fn alice() -> AccountId {
+        AccountId::from_str("alice.near").unwrap()
+    }
+
+    fn beneficiary() -> AccountId {
+        AccountId::from_str("beneficiary.near").unwrap()
+    }
+
+    #[test]
+    fn test_terminate_retroactively_adjusts_timestamp() {
+        let mut builder = VMContextBuilder::new();
+        builder.predecessor_account_id(alice());
+        testing_env!(builder.build());
+
+        let total_balance = 1_000_000;
+        let lockup_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC,
+                balance: 0,
+            },
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + 4 * ONE_YEAR_SEC,
+                balance: total_balance,
+            },
+        ]);
+
+        let vesting_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC,
+                balance: 0,
+            },
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + 2 * ONE_YEAR_SEC,
+                balance: total_balance,
+            },
+        ]);
+
+        // At 1.5 years, user claims tokens.
+        // Vested: 500_000 (1/2 of total, since vesting is 1 year linear)
+        // Unlocked: 375_000 (1.5 / 4 years)
+        // Claimable is min(500_000, 375_000) = 375_000
+        let claimed_balance = 375_000;
+        let mut lockup = Lockup {
+            account_id: alice(),
+            schedule: lockup_schedule.clone(),
+            claimed_balance,
+            termination_config: Some(TerminationConfig {
+                beneficiary_id: beneficiary(),
+                vesting_schedule: VestingConditions::Schedule(vesting_schedule.clone()),
+            }),
+        };
+
+        let claim_timestamp = GENESIS_TIMESTAMP_SEC + 3 * ONE_YEAR_SEC / 2;
+        assert_eq!(claimed_balance, lockup.schedule.unlocked_balance(claim_timestamp));
+
+        // Now, terminate retroactively at 1 year.
+        // At this point, vested_balance is 0 according to vesting_schedule.
+        let termination_timestamp = GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC;
+
+        // Since vested_balance (0) < claimed_balance (375_000), the logic should trigger.
+        // The new vested_balance will be claimed_balance (375_000).
+        // The new termination_timestamp will be calculated via `lockup_schedule.get_vesting_timestamp_for_amount(0)`,
+        // which corresponds to GENESIS_TIMESTAMP_SEC.
+
+        let (unvested_balance, beneficiary_id) = lockup.terminate(None, termination_timestamp);
+
+        // unvested = total - new_vested = 1,000,000 - 375_000 = 625_000
+        assert_eq!(unvested_balance, 625_000);
+        assert_eq!(beneficiary_id, beneficiary());
+
+        // Check the lockup state after termination
+        // The schedule should be terminated with the new values.
+        assert_eq!(lockup.schedule.total_balance(), 375_000);
+
+        // The last checkpoint of the schedule should be at the *new* termination_timestamp.
+        let final_timestamp = lockup.schedule.0.last().unwrap().timestamp;
+
+        // The new timestamp is calculated with `get_vesting_timestamp_for_amount(0)` on the lockup schedule.
+        // For a linear schedule, this should be the start of the schedule.
+        let expected_new_timestamp = lockup_schedule.get_vesting_timestamp_for_amount(375_000);
+        assert!(final_timestamp.abs_diff(expected_new_timestamp) <= 1);
+        assert!(claim_timestamp > termination_timestamp);
+    }
+
+    #[test]
+    fn test_terminate_no_claims() {
+        let mut builder = VMContextBuilder::new();
+        builder.predecessor_account_id(alice());
+        testing_env!(builder.build());
+
+        let total_balance = 15_000_000;
+        let lockup_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC,
+                balance: 0,
+            },
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + 4 * ONE_YEAR_SEC,
+                balance: total_balance,
+            },
+        ]);
+
+        let mut lockup = Lockup {
+            account_id: alice(),
+            schedule: lockup_schedule.clone(),
+            claimed_balance: 0, // No claims
+            termination_config: Some(TerminationConfig {
+                beneficiary_id: beneficiary(),
+                vesting_schedule: VestingConditions::SameAsLockupSchedule,
+            }),
+        };
+
+        // Terminate at 1.5 years.
+        let termination_timestamp = GENESIS_TIMESTAMP_SEC + 3 * ONE_YEAR_SEC / 2;
+
+        // At this point, vested_balance is 2_500_000 (1/6 of total_balance, since vesting is 3 year linear from year 1 to 4).
+        let vested_balance = lockup.schedule.unlocked_balance(termination_timestamp);
+        assert_eq!(vested_balance, 2_500_000);
+
+        let (unvested_balance, beneficiary_id) = lockup.terminate(None, termination_timestamp);
+
+        // unvested = total - vested = 15,000,000 - 2,500,000 = 12,500,000
+        assert_eq!(unvested_balance, 12_500_000);
+        assert_eq!(beneficiary_id, beneficiary());
+
+        // Check the lockup state after termination
+        assert_eq!(lockup.schedule.total_balance(), vested_balance);
+        let final_timestamp = lockup.schedule.0.last().unwrap().timestamp;
+        assert_eq!(final_timestamp, termination_timestamp);
+        assert_eq!(lockup.claimed_balance, 0);
+    }
+
+    #[test]
+    fn test_terminate_retroactively_no_claims() {
+        let mut builder = VMContextBuilder::new();
+        builder.predecessor_account_id(alice());
+        testing_env!(builder.build());
+
+        let total_balance = 35_000_000;
+        let lockup_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC,
+                balance: 0,
+            },
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + 4 * ONE_YEAR_SEC,
+                balance: total_balance,
+            },
+        ]);
+
+        let mut lockup = Lockup {
+            account_id: alice(),
+            schedule: lockup_schedule.clone(),
+            claimed_balance: 0, // No claims
+            termination_config: Some(TerminationConfig {
+                beneficiary_id: beneficiary(),
+                vesting_schedule: VestingConditions::SameAsLockupSchedule,
+            }),
+        };
+
+        // Terminate retroactively at 11 months.
+        let termination_timestamp = GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC * 11 / 12;
+
+        // At this point, vested_balance is 0 according to vesting_schedule.
+        let vested_balance = lockup.schedule.unlocked_balance(termination_timestamp);
+        assert_eq!(vested_balance, 0);
+
+        // Since claimed_balance is 0, no special adjustment should happen.
+        let (unvested_balance, beneficiary_id) = lockup.terminate(None, termination_timestamp);
+
+        // unvested = total - vested = 35,000,000 - 0 = 35,000,000
+        assert_eq!(unvested_balance, 35_000_000);
+        assert_eq!(beneficiary_id, beneficiary());
+
+        // Check the lockup state after termination
+        assert_eq!(lockup.schedule.total_balance(), vested_balance);
+        let final_timestamp = lockup.schedule.0.last().unwrap().timestamp;
+        assert_eq!(final_timestamp, lockup.schedule.0.first().unwrap().timestamp + 1);
+        assert_eq!(lockup.claimed_balance, 0);
+    }
+
+    #[test]
+    pub fn test_vesting_timestamp_evaluation() {
+        let mut builder = VMContextBuilder::new();
+        builder.predecessor_account_id(alice());
+        testing_env!(builder.build());
+
+        let total_balance = 5_000_000_000_000_000_000_000_000;
+        let lockup_schedule = Schedule(vec![
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC,
+                balance: 0,
+            },
+            Checkpoint {
+                timestamp: GENESIS_TIMESTAMP_SEC + 4 * ONE_YEAR_SEC,
+                balance: total_balance,
+            },
+        ]);
+
+        let lockup = Lockup {
+            account_id: alice(),
+            schedule: lockup_schedule.clone(),
+            claimed_balance: 0,
+            termination_config: Some(TerminationConfig {
+                beneficiary_id: beneficiary(),
+                vesting_schedule: VestingConditions::SameAsLockupSchedule,
+            }),
+        };
+
+        let reference_timestamps = vec![
+            GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC + 30 * 60 * 60,
+            GENESIS_TIMESTAMP_SEC + ONE_YEAR_SEC + 70 * 30 * 60 * 60,
+            GENESIS_TIMESTAMP_SEC + 2 * ONE_YEAR_SEC + 95 * 30 * 60 * 60 + 2,
+            GENESIS_TIMESTAMP_SEC + 3 * ONE_YEAR_SEC + 30 * 60 + 7,
+        ];
+
+        for reference_timestamp in reference_timestamps {
+            let reference_amount = lockup.schedule.unlocked_balance(reference_timestamp);
+            assert_eq!(
+                reference_timestamp,
+                lockup.schedule.get_vesting_timestamp_for_amount(reference_amount)
+            );
+        }
+    }
 }
 
 /// Amount of fungible tokens
diff --git a/model/src/lockup_api.rs b/model/src/lockup_api.rs
index 3b3978b..8456baf 100644
--- a/model/src/lockup_api.rs
+++ b/model/src/lockup_api.rs
@@ -5,6 +5,7 @@ use crate::{
     draft::{Draft, DraftGroupIndex, DraftIndex},
     lockup::LockupIndex,
     schedule::Schedule,
+    termination::TerminationConfig,
     TimestampSec, WrappedBalance,
 };
 
@@ -47,4 +48,6 @@ pub trait LockupApi {
     fn discard_draft_group(&mut self, draft_group_id: DraftGroupIndex);
 
     fn delete_drafts(&mut self, draft_ids: Vec<DraftIndex>);
+
+    fn edit(&mut self, index: LockupIndex, schedule: Option<Schedule>, termination_config: Option<TerminationConfig>);
 }
diff --git a/model/src/schedule.rs b/model/src/schedule.rs
index 3e21925..964e1a4 100644
--- a/model/src/schedule.rs
+++ b/model/src/schedule.rs
@@ -177,4 +177,35 @@ impl Schedule {
 
         res
     }
+
+    pub fn get_vesting_timestamp_for_amount(&self, amount: Balance) -> TimestampSec {
+        // Using binary search by time to find the current checkpoint.
+        let index = match self.0.binary_search_by_key(&amount, |checkpoint| checkpoint.balance) {
+            // Exact timestamp found
+            Ok(index) => index,
+            // No match, the next index is given.
+            Err(index) => {
+                if index == 0 {
+                    // Not started
+                    return 0;
+                }
+                index - 1
+            }
+        };
+        let checkpoint = &self.0[index];
+        if index + 1 == self.0.len() {
+            // The last checkpoint. Fully unlocked.
+            return checkpoint.timestamp;
+        }
+        let next_checkpoint = &self.0[index + 1];
+
+        let total_balance = next_checkpoint.balance - checkpoint.balance;
+        let total_duration = next_checkpoint.timestamp - checkpoint.timestamp;
+        let claimed_delta = amount - checkpoint.balance;
+        let time_delta = ((U256::from(claimed_delta) * U256::from(total_duration) + U256::from(total_balance))
+            / U256::from(total_balance))
+        .as_u32();
+
+        checkpoint.timestamp + time_delta
+    }
 }
diff --git a/model/src/termination.rs b/model/src/termination.rs
index 246c0b0..8adea71 100644
--- a/model/src/termination.rs
+++ b/model/src/termination.rs
@@ -43,10 +43,19 @@ impl Lockup {
             VestingConditions::Schedule(schedule) => schedule,
         }
         .unlocked_balance(termination_timestamp);
+
+        let (vested_balance, termination_timestamp) = if vested_balance >= self.claimed_balance {
+            (vested_balance, termination_timestamp)
+        } else {
+            let last_claim_timestamp = self.schedule.get_vesting_timestamp_for_amount(vested_balance);
+            (self.claimed_balance, last_claim_timestamp)
+        };
+
         let unvested_balance = total_balance - vested_balance;
         if unvested_balance > 0 {
             self.schedule.terminate(vested_balance, termination_timestamp);
         }
+
         (unvested_balance, termination_config.beneficiary_id)
     }
 }
diff --git a/res/hodl_lockup.wasm b/res/hodl_lockup.wasm
index e3f1f00..dd1e14c 100755
Binary files a/res/hodl_lockup.wasm and b/res/hodl_lockup.wasm differ
diff --git a/rust-toolchain.toml b/rust-toolchain.toml
index c6e4d7d..bcccbe7 100644
--- a/rust-toolchain.toml
+++ b/rust-toolchain.toml
@@ -1,2 +1,3 @@
 [toolchain]
-channel = "1.79"
+channel = "1.86"
+components = ["rust-analyzer", "rustfmt"]