From f84b9090add36ff89f9c5bab3780b89668022112 Mon Sep 17 00:00:00 2001 From: Andre Popovitch Date: Mon, 30 Sep 2024 15:20:44 -0500 Subject: [PATCH] Implement additional helpers in ic-nervous-system-agent --- Cargo.lock | 3 + rs/nervous_system/agent/BUILD.bazel | 3 + rs/nervous_system/agent/Cargo.toml | 3 + rs/nervous_system/agent/src/agent_impl.rs | 2 + rs/nervous_system/agent/src/lib.rs | 11 ++- rs/nervous_system/agent/src/nns/governance.rs | 16 ++++ rs/nervous_system/agent/src/nns/mod.rs | 1 + rs/nervous_system/agent/src/pocketic_impl.rs | 2 + rs/nervous_system/agent/src/sns/governance.rs | 7 ++ rs/nervous_system/agent/src/sns/index.rs | 7 ++ rs/nervous_system/agent/src/sns/ledger.rs | 7 ++ rs/nervous_system/agent/src/sns/root.rs | 5 ++ rs/nervous_system/agent/src/sns/swap.rs | 75 +++++++++++++++++++ rs/nns/governance/api/src/lib.rs | 1 + rs/nns/governance/api/src/request_impls.rs | 7 ++ rs/sns/swap/src/lib.rs | 1 + rs/sns/swap/src/request_impls.rs | 19 +++++ 17 files changed, 168 insertions(+), 2 deletions(-) create mode 100644 rs/nervous_system/agent/src/nns/governance.rs create mode 100644 rs/nns/governance/api/src/request_impls.rs create mode 100644 rs/sns/swap/src/request_impls.rs diff --git a/Cargo.lock b/Cargo.lock index 067e8b126a1..c3e3aa5b0bb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9526,9 +9526,12 @@ dependencies = [ "ic-agent", "ic-base-types", "ic-nervous-system-clients", + "ic-nns-common", "ic-nns-constants", + "ic-nns-governance-api", "ic-sns-governance", "ic-sns-root", + "ic-sns-swap", "ic-sns-wasm", "pocket-ic", "serde", diff --git a/rs/nervous_system/agent/BUILD.bazel b/rs/nervous_system/agent/BUILD.bazel index 88fc914ac53..29fcc996251 100644 --- a/rs/nervous_system/agent/BUILD.bazel +++ b/rs/nervous_system/agent/BUILD.bazel @@ -6,10 +6,13 @@ DEPENDENCIES = [ # Keep sorted. "//packages/pocket-ic", "//rs/nervous_system/clients", + "//rs/nns/common", "//rs/nns/constants", + "//rs/nns/governance/api", "//rs/nns/sns-wasm", "//rs/sns/governance", "//rs/sns/root", + "//rs/sns/swap", "//rs/types/base_types", "@crate_index//:anyhow", "@crate_index//:candid", diff --git a/rs/nervous_system/agent/Cargo.toml b/rs/nervous_system/agent/Cargo.toml index 6fd6bf16418..b8e2557c7de 100644 --- a/rs/nervous_system/agent/Cargo.toml +++ b/rs/nervous_system/agent/Cargo.toml @@ -13,11 +13,14 @@ candid = { workspace = true } ic-agent = { workspace = true } ic-base-types = { path = "../../types/base_types" } ic-nervous-system-clients = { path = "../clients" } +ic-nns-governance-api = { path = "../../nns/governance/api" } +ic-nns-common = { path = "../../nns/common" } ic-nns-constants = { path = "../../nns/constants" } ic-sns-wasm = { path = "../../nns/sns-wasm" } ic-sns-governance = { path = "../../sns/governance" } pocket-ic = { path = "../../../packages/pocket-ic" } ic-sns-root = { path = "../../sns/root" } +ic-sns-swap = { path = "../../sns/swap" } serde = { workspace = true } tempfile = { workspace = true } thiserror = { workspace = true } diff --git a/rs/nervous_system/agent/src/agent_impl.rs b/rs/nervous_system/agent/src/agent_impl.rs index 2f050e5444b..6c47afc2319 100644 --- a/rs/nervous_system/agent/src/agent_impl.rs +++ b/rs/nervous_system/agent/src/agent_impl.rs @@ -15,6 +15,8 @@ pub enum AgentCallError { CandidDecode(candid::Error), } +impl crate::sealed::Sealed for Agent {} + impl CallCanisters for Agent { type Error = AgentCallError; async fn call( diff --git a/rs/nervous_system/agent/src/lib.rs b/rs/nervous_system/agent/src/lib.rs index 29bb33a4082..d90b0b4d69f 100644 --- a/rs/nervous_system/agent/src/lib.rs +++ b/rs/nervous_system/agent/src/lib.rs @@ -7,8 +7,15 @@ use candid::Principal; use ic_nervous_system_clients::Request; use std::fmt::Display; -pub trait CallCanisters { - type Error: Display + Send; +// This is used to "seal" the CallCanisters trait so that it cannot be implemented outside of this crate. +// This is useful because it means we can modify the trait in the future without worrying about +// breaking backwards compatibility with implementations outside of this crate. +mod sealed { + pub trait Sealed {} +} + +pub trait CallCanisters: sealed::Sealed { + type Error: Display + Send + std::error::Error + 'static; fn call( &self, canister_id: impl Into + Send, diff --git a/rs/nervous_system/agent/src/nns/governance.rs b/rs/nervous_system/agent/src/nns/governance.rs new file mode 100644 index 00000000000..7f9786153a0 --- /dev/null +++ b/rs/nervous_system/agent/src/nns/governance.rs @@ -0,0 +1,16 @@ +use crate::CallCanisters; +use ic_nns_common::pb::v1::ProposalId; +use ic_nns_constants::GOVERNANCE_CANISTER_ID; +use ic_nns_governance_api::pb::v1::{ + GetNeuronsFundAuditInfoRequest, GetNeuronsFundAuditInfoResponse, +}; + +pub async fn get_neurons_fund_audit_info( + agent: &C, + nns_proposal_id: ProposalId, +) -> Result { + let request = GetNeuronsFundAuditInfoRequest { + nns_proposal_id: Some(nns_proposal_id), + }; + agent.call(GOVERNANCE_CANISTER_ID, request).await +} diff --git a/rs/nervous_system/agent/src/nns/mod.rs b/rs/nervous_system/agent/src/nns/mod.rs index 341405efcb0..280c372d50c 100644 --- a/rs/nervous_system/agent/src/nns/mod.rs +++ b/rs/nervous_system/agent/src/nns/mod.rs @@ -1 +1,2 @@ +pub mod governance; pub mod sns_wasm; diff --git a/rs/nervous_system/agent/src/pocketic_impl.rs b/rs/nervous_system/agent/src/pocketic_impl.rs index cfb9907571e..bbdcb32bb2e 100644 --- a/rs/nervous_system/agent/src/pocketic_impl.rs +++ b/rs/nervous_system/agent/src/pocketic_impl.rs @@ -17,6 +17,8 @@ pub enum PocketIcCallError { CandidDecode(candid::Error), } +impl crate::sealed::Sealed for PocketIc {} + impl CallCanisters for PocketIc { type Error = PocketIcCallError; async fn call( diff --git a/rs/nervous_system/agent/src/sns/governance.rs b/rs/nervous_system/agent/src/sns/governance.rs index 9d5e6e6ced5..8783ac0eed5 100644 --- a/rs/nervous_system/agent/src/sns/governance.rs +++ b/rs/nervous_system/agent/src/sns/governance.rs @@ -28,3 +28,10 @@ impl GovernanceCanister { .await } } + +impl GovernanceCanister { + pub fn new(canister_id: impl Into) -> Self { + let canister_id = canister_id.into(); + Self { canister_id } + } +} diff --git a/rs/nervous_system/agent/src/sns/index.rs b/rs/nervous_system/agent/src/sns/index.rs index 18324670829..ec231e7ef53 100644 --- a/rs/nervous_system/agent/src/sns/index.rs +++ b/rs/nervous_system/agent/src/sns/index.rs @@ -5,3 +5,10 @@ use serde::{Deserialize, Serialize}; pub struct IndexCanister { pub canister_id: PrincipalId, } + +impl IndexCanister { + pub fn new(canister_id: impl Into) -> Self { + let canister_id = canister_id.into(); + Self { canister_id } + } +} diff --git a/rs/nervous_system/agent/src/sns/ledger.rs b/rs/nervous_system/agent/src/sns/ledger.rs index 41a5b0a6f31..d5ff1ceb2b6 100644 --- a/rs/nervous_system/agent/src/sns/ledger.rs +++ b/rs/nervous_system/agent/src/sns/ledger.rs @@ -5,3 +5,10 @@ use serde::{Deserialize, Serialize}; pub struct LedgerCanister { pub canister_id: PrincipalId, } + +impl LedgerCanister { + pub fn new(canister_id: impl Into) -> Self { + let canister_id = canister_id.into(); + Self { canister_id } + } +} diff --git a/rs/nervous_system/agent/src/sns/root.rs b/rs/nervous_system/agent/src/sns/root.rs index a1aa270b9bf..52596deaac0 100644 --- a/rs/nervous_system/agent/src/sns/root.rs +++ b/rs/nervous_system/agent/src/sns/root.rs @@ -10,6 +10,11 @@ pub struct RootCanister { } impl RootCanister { + pub fn new(canister_id: impl Into) -> Self { + let canister_id = canister_id.into(); + Self { canister_id } + } + pub async fn sns_canisters_summary( &self, agent: &C, diff --git a/rs/nervous_system/agent/src/sns/swap.rs b/rs/nervous_system/agent/src/sns/swap.rs index d9080e28e19..48bcf820bc9 100644 --- a/rs/nervous_system/agent/src/sns/swap.rs +++ b/rs/nervous_system/agent/src/sns/swap.rs @@ -1,7 +1,82 @@ use ic_base_types::PrincipalId; +use ic_sns_swap::pb::v1::{ + GetDerivedStateRequest, GetDerivedStateResponse, GetInitRequest, GetInitResponse, + ListSnsNeuronRecipesRequest, ListSnsNeuronRecipesResponse, SnsNeuronRecipe, +}; use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::CallCanisters; #[derive(Clone, Debug, Deserialize, Serialize)] pub struct SwapCanister { pub canister_id: PrincipalId, } + +#[derive(Debug, Error)] +pub enum ListAllSnsNeuronRecipesError { + #[error(transparent)] + CanisterCallError(#[from] CallCanistersError), + #[error("There seem to be too many neuron recipes ({0}).")] + TooManyRecipes(u64), +} + +impl SwapCanister { + pub fn new(canister_id: impl Into) -> Self { + let canister_id = canister_id.into(); + Self { canister_id } + } + + pub async fn get_derived_state( + &self, + agent: &C, + ) -> Result { + agent + .call(self.canister_id, GetDerivedStateRequest {}) + .await + } + + pub async fn get_init(&self, agent: &C) -> Result { + agent.call(self.canister_id, GetInitRequest {}).await + } + + pub async fn list_sns_neuron_recipes( + &self, + agent: &C, + limit: u32, + offset: u64, + ) -> Result { + agent + .call( + self.canister_id, + ListSnsNeuronRecipesRequest { + limit: Some(limit), + offset: Some(offset), + }, + ) + .await + } + + pub async fn list_all_sns_neuron_recipes( + &self, + agent: &C, + ) -> Result, ListAllSnsNeuronRecipesError> { + let mut sns_neuron_recipes: Vec = vec![]; + let batch_size = 10_000_u64; + let num_calls = 100_u64; + for i in 0..num_calls { + let new_sns_neuron_recipes = self + .list_sns_neuron_recipes(agent, batch_size as u32, batch_size * i) + .await + .map_err(ListAllSnsNeuronRecipesError::CanisterCallError)?; + if new_sns_neuron_recipes.sns_neuron_recipes.is_empty() { + return Ok(sns_neuron_recipes); + } else { + sns_neuron_recipes.extend(new_sns_neuron_recipes.sns_neuron_recipes.into_iter()) + } + } + Err(ListAllSnsNeuronRecipesError::TooManyRecipes( + batch_size * num_calls, + )) + } +} diff --git a/rs/nns/governance/api/src/lib.rs b/rs/nns/governance/api/src/lib.rs index 1d8adeafcc0..3369d4770bd 100644 --- a/rs/nns/governance/api/src/lib.rs +++ b/rs/nns/governance/api/src/lib.rs @@ -2,5 +2,6 @@ pub mod bitcoin; pub mod pb; pub mod proposal_submission_helpers; pub mod proposal_validation; +mod request_impls; pub mod subnet_rental; pub mod test_api; diff --git a/rs/nns/governance/api/src/request_impls.rs b/rs/nns/governance/api/src/request_impls.rs new file mode 100644 index 00000000000..d36bf71724e --- /dev/null +++ b/rs/nns/governance/api/src/request_impls.rs @@ -0,0 +1,7 @@ +use ic_nervous_system_clients::Request; + +impl Request for crate::pb::v1::GetNeuronsFundAuditInfoRequest { + type Response = crate::pb::v1::GetNeuronsFundAuditInfoResponse; + const METHOD: &'static str = "get_neurons_fund_audit_info"; + const UPDATE: bool = false; +} diff --git a/rs/sns/swap/src/lib.rs b/rs/sns/swap/src/lib.rs index a5531dec6e5..aa342e0bcb9 100644 --- a/rs/sns/swap/src/lib.rs +++ b/rs/sns/swap/src/lib.rs @@ -4,6 +4,7 @@ pub mod logs; pub mod memory; pub mod neurons_fund; pub mod pb; +mod request_impls; pub mod swap; pub mod swap_builder; pub mod types; diff --git a/rs/sns/swap/src/request_impls.rs b/rs/sns/swap/src/request_impls.rs new file mode 100644 index 00000000000..be943eeeb08 --- /dev/null +++ b/rs/sns/swap/src/request_impls.rs @@ -0,0 +1,19 @@ +use ic_nervous_system_clients::Request; + +impl Request for crate::pb::v1::GetDerivedStateRequest { + type Response = crate::pb::v1::GetDerivedStateResponse; + const METHOD: &'static str = "get_derived_state"; + const UPDATE: bool = false; +} + +impl Request for crate::pb::v1::GetInitRequest { + type Response = crate::pb::v1::GetInitResponse; + const METHOD: &'static str = "get_init"; + const UPDATE: bool = false; +} + +impl Request for crate::pb::v1::ListSnsNeuronRecipesRequest { + type Response = crate::pb::v1::ListSnsNeuronRecipesResponse; + const METHOD: &'static str = "list_sns_neuron_recipes"; + const UPDATE: bool = false; +}