From 9be0d61d8bfc6307760310fe8509ada8f27362a0 Mon Sep 17 00:00:00 2001 From: Iyanumajekodunmi756 Date: Wed, 24 Jun 2026 11:52:46 +0100 Subject: [PATCH] feat: Add update_orchestrator function to allow users to rotate their registered orchestrator - Implements update_orchestrator(env, user, new_orchestrator, name) function - Requires user authentication via user.require_auth() - Loads existing UserAccount and validates orchestrator is registered - Removes old OrchestratorOwner reverse-lookup entry from persistent storage - Sets new orchestrator address and updates reverse lookup - Emits UpdateOrchEvent with user, old_orchestrator, and new_orchestrator - Guards: panics if active_tasks_count > 0 (same as withdraw) - Guards: panics if no orchestrator is currently registered - Adds comprehensive tests covering: - Successful orchestrator rotation - Blocked when task is active - Blocked when no orchestrator registered - Uses #[contractevent] macro for UpdateOrchEvent Fixes #8 --- contracts/agent-vault/src/lib.rs | 60 ++++++++++++++++++++ contracts/agent-vault/src/tests.rs | 91 ++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+) diff --git a/contracts/agent-vault/src/lib.rs b/contracts/agent-vault/src/lib.rs index 42f0bb0..6da2456 100644 --- a/contracts/agent-vault/src/lib.rs +++ b/contracts/agent-vault/src/lib.rs @@ -39,6 +39,14 @@ pub struct RegOrchEvent { pub orchestrator: Address, } +#[contractevent] +pub struct UpdateOrchEvent { + #[topic] + pub user: Address, + pub old_orchestrator: Address, + pub new_orchestrator: Address, +} + #[contractevent] pub struct TaskNewEvent { #[topic] @@ -278,6 +286,58 @@ impl AgentVault { ); } + /// Update a user's registered orchestrator to a new address. + /// Requires user authentication and no active tasks. + pub fn update_orchestrator(env: Env, user: Address, new_orchestrator: Address, name: String) { + user.require_auth(); + + let mut account: UserAccount = env + .storage() + .persistent() + .get(&DataKey::User(user.clone())) + .expect("No account"); + + let old_orchestrator = account.orchestrator.expect( + "No orchestrator registered. Use register_orchestrator instead.", + ); + + assert!( + account.active_tasks_count == 0, + "Cannot update orchestrator while tasks are active" + ); + + // Remove old reverse lookup + env.storage() + .persistent() + .remove(&DataKey::OrchestratorOwner(old_orchestrator.clone())); + + // Update orchestrator in account + account.orchestrator = Some(new_orchestrator.clone()); + account.orchestrator_name = name.clone(); + env.storage() + .persistent() + .set(&DataKey::User(user.clone()), &account); + + // Set new reverse lookup + env.storage() + .persistent() + .set(&DataKey::OrchestratorOwner(new_orchestrator.clone()), &user); + + UpdateOrchEvent { + user: user.clone(), + old_orchestrator, + new_orchestrator: new_orchestrator.clone(), + } + .publish(&env); + log!( + &env, + "update_orchestrator user={} old_orch={} new_orch={}", + user, + old_orchestrator, + new_orchestrator + ); + } + // ── Task lifecycle ──────────────────────────────────────────────────────── /// Orchestrator creates a task, locking plan_cost from user's available balance. diff --git a/contracts/agent-vault/src/tests.rs b/contracts/agent-vault/src/tests.rs index c4f92ab..872ab55 100644 --- a/contracts/agent-vault/src/tests.rs +++ b/contracts/agent-vault/src/tests.rs @@ -250,6 +250,97 @@ fn test_register_orchestrator_twice_panics() { .register_orchestrator(&user, &orchestrator2, &name); } +// ── 4.5. Update Orchestrator Tests ─────────────────────────────────────────── + +#[test] +fn test_update_orchestrator_success() { + let test_env = setup_test(); + test_env.client.init(&test_env.admin, &test_env.usdc_sac); + + let user = Address::generate(&test_env.env); + let orchestrator1 = Address::generate(&test_env.env); + let orchestrator2 = Address::generate(&test_env.env); + let name1 = soroban_sdk::String::from_str(&test_env.env, "OldOrchestrator"); + let name2 = soroban_sdk::String::from_str(&test_env.env, "NewOrchestrator"); + + // Register initial orchestrator + test_env + .client + .register_orchestrator(&user, &orchestrator1, &name1); + + // Update to new orchestrator + test_env + .client + .update_orchestrator(&user, &orchestrator2, &name2); + + // Verify new orchestrator is stored + let account = test_env.client.get_account(&user).unwrap(); + assert_eq!(account.orchestrator.unwrap(), orchestrator2); + assert_eq!(account.orchestrator_name, name2); + + // Verify old reverse lookup is removed + assert_eq!( + test_env + .client + .get_orchestrator_owner(&orchestrator1), + None + ); + + // Verify new reverse lookup is set + assert_eq!( + test_env + .client + .get_orchestrator_owner(&orchestrator2) + .unwrap(), + user + ); +} + +#[test] +#[should_panic(expected = "Cannot update orchestrator while tasks are active")] +fn test_update_orchestrator_blocked_active_task() { + let test_env = setup_test(); + test_env.client.init(&test_env.admin, &test_env.usdc_sac); + + let user = Address::generate(&test_env.env); + let orchestrator1 = Address::generate(&test_env.env); + let orchestrator2 = Address::generate(&test_env.env); + let name1 = soroban_sdk::String::from_str(&test_env.env, "OldOrchestrator"); + let name2 = soroban_sdk::String::from_str(&test_env.env, "NewOrchestrator"); + + test_env.token_admin_client.mint(&user, &1000); + test_env.client.deposit(&user, &600); + + // Register initial orchestrator + test_env + .client + .register_orchestrator(&user, &orchestrator1, &name1); + + // Create a task to set active_tasks_count = 1 + test_env.client.create_task(&orchestrator1, &100); + + // Attempt to update orchestrator while task is active + test_env + .client + .update_orchestrator(&user, &orchestrator2, &name2); +} + +#[test] +#[should_panic(expected = "No orchestrator registered")] +fn test_update_orchestrator_no_orchestrator_registered() { + let test_env = setup_test(); + test_env.client.init(&test_env.admin, &test_env.usdc_sac); + + let user = Address::generate(&test_env.env); + let orchestrator = Address::generate(&test_env.env); + let name = soroban_sdk::String::from_str(&test_env.env, "MyOrchestrator"); + + // Try to update without registering first + test_env + .client + .update_orchestrator(&user, &orchestrator, &name); +} + // ── 5. Create Task Tests ───────────────────────────────────────────────────── #[test]