feat: blueprint pre-registration (#625)

* feat!: preregistration for Tangle blueprints

* feat!: configs control preregistration and Eigenlayer support

* fix(ci): workflow debugging

* fix: clippy

* fix: better variable names

* fix: tangle incredible squaring tests

---------

Co-authored-by: drewstone <[email protected]>

.github/workflows/ci.yml

@@ -139,12 +139,14 @@ jobs:
           # if test == "gadget_context_derive", add --features="networking", else ""
           if [[ "${{ matrix.package }}" == "gadget-context-derive" ]]; then
             echo "cargo_nextest_args=--features networking" >> $GITHUB_ENV
+          elif [[ "${{ matrix.package }}" == "incredible-squaring-blueprint" || "${{ matrix.package }}" == "incredible-squaring-blueprint-eigenlayer" ]]; then
+            echo "cargo_nextest_args=--profile serial" >> $GITHUB_ENV
           else
-            echo "cargo_nextest_args=" >> $GITHUB_ENV
+            echo "cargo_nextest_args=--profile ci" >> $GITHUB_ENV
           fi
 
       - name: tests
-        run: cargo nextest run --profile ci --package ${{ matrix.package }} ${{ env.cargo_nextest_args }}
+        run: cargo nextest run --package ${{ matrix.package }} ${{ env.cargo_nextest_args }}
 
       # TODO: nextest doesn't support doc tests yet (https://github.com/nextest-rs/nextest/issues/16)
       - name: doc tests

blueprints/examples/src/tests.rs

@@ -140,7 +140,7 @@ async fn test_periodic_web_poller() -> Result<()> {
     let harness = TangleTestHarness::setup(temp_dir).await?;
 
     // Setup service
-    let (mut test_env, service_id) = harness.setup_services().await?;
+    let (mut test_env, service_id, _blueprint_id) = harness.setup_services(false).await?;
 
     // Add the web poller job
     test_env.add_job(crate::periodic_web_poller::constructor("*/5 * * * * *"));
@@ -172,7 +172,7 @@ async fn test_raw_tangle_events() -> Result<()> {
     let env = harness.env().clone();
 
     // Setup service
-    let (mut test_env, service_id) = harness.setup_services().await?;
+    let (mut test_env, service_id, _blueprint_id) = harness.setup_services(false).await?;
 
     // Add the raw tangle events job
     test_env.add_job(crate::raw_tangle_events::constructor(env.clone()).await?);

blueprints/incredible-squaring-eigenlayer/src/tests.rs

@@ -33,8 +33,19 @@ sol!(
 );
 
 #[tokio::test(flavor = "multi_thread")]
-#[allow(clippy::needless_return)]
 async fn test_eigenlayer_incredible_squaring_blueprint() {
+    run_eigenlayer_incredible_squaring_test(true, 1).await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_eigenlayer_pre_register_incredible_squaring_blueprint() {
+    run_eigenlayer_incredible_squaring_test(false, 1).await;
+}
+
+async fn run_eigenlayer_incredible_squaring_test(
+    exit_after_registration: bool,
+    expected_responses: usize,
+) {
     setup_log();
 
     // Initialize test harness
@@ -48,7 +59,6 @@ async fn test_eigenlayer_incredible_squaring_blueprint() {
     let task_manager_address = deploy_task_manager(&harness).await;
 
     // Spawn Task Spawner and Task Response Listener
-    let num_successful_responses_required = 3;
     let successful_responses = Arc::new(Mutex::new(0));
     let successful_responses_clone = successful_responses.clone();
     let response_listener_address =
@@ -91,14 +101,20 @@ async fn test_eigenlayer_incredible_squaring_blueprint() {
     let x_square_eigen = XsquareEigenEventHandler::new(contract.clone(), eigen_client_context);
 
     let mut test_env = EigenlayerBLSTestEnv::new(
-        EigenlayerBLSConfig::new(Default::default(), Default::default()),
+        EigenlayerBLSConfig::new(Default::default(), Default::default())
+            .with_exit_after_register(exit_after_registration),
         env.clone(),
     )
     .unwrap();
     test_env.add_job(initialize_task);
     test_env.add_job(x_square_eigen);
     test_env.add_background_service(aggregator_context);
 
+    if exit_after_registration {
+        // Run the runner once to register, since pre-registration is enabled
+        test_env.run_runner().await.unwrap();
+    }
+
     tokio::spawn(async move {
         test_env.run_runner().await.unwrap();
     });
@@ -110,7 +126,7 @@ async fn test_eigenlayer_incredible_squaring_blueprint() {
     let timeout_duration = Duration::from_secs(300);
     let result = wait_for_responses(
         successful_responses.clone(),
-        num_successful_responses_required,
+        expected_responses,
         timeout_duration,
     )
     .await;
@@ -123,13 +139,13 @@ async fn test_eigenlayer_incredible_squaring_blueprint() {
 
     match result {
         Ok(Ok(())) => {
-            info!("Test completed successfully with {num_successful_responses_required} tasks responded to.");
+            info!("Test completed successfully with {expected_responses} tasks responded to.");
         }
         _ => {
             panic!(
                 "Test failed with {} successful responses out of {} required",
                 successful_responses_clone.lock().unwrap(),
-                num_successful_responses_required
+                expected_responses
             );
         }
     }

blueprints/incredible-squaring/foundry.toml

@@ -2,6 +2,7 @@
 src = "contracts/src"
 out = "contracts/out"
 libs = ["dependencies"]
+remappings = ["incredible-squaring/=contracts/src/"]
 auto_detect_remappings = true
 solc_version = "0.8.20"
 

blueprints/incredible-squaring/remappings.txt

@@ -1,2 +1 @@
-incredible-squaring/=contracts/src/
 tnt-core/=dependencies/tnt-core-0.1.0/src/

blueprints/incredible-squaring/src/tests.rs

@@ -5,6 +5,7 @@ use blueprint_sdk::testing::utils::harness::TestHarness;
 use blueprint_sdk::testing::utils::runner::TestEnv;
 use blueprint_sdk::testing::utils::tangle::{InputValue, OutputValue, TangleTestHarness};
 use color_eyre::Result;
+use std::time::Duration;
 
 #[tokio::test]
 async fn test_incredible_squaring() -> Result<()> {
@@ -27,7 +28,7 @@ async fn test_incredible_squaring() -> Result<()> {
         .unwrap();
 
     // Setup service
-    let (mut test_env, service_id) = harness.setup_services().await?;
+    let (mut test_env, service_id, _blueprint_id) = harness.setup_services(false).await?;
     test_env.add_job(handler);
 
     tokio::spawn(async move {
@@ -47,3 +48,51 @@ async fn test_incredible_squaring() -> Result<()> {
     assert_eq!(results.service_id, service_id);
     Ok(())
 }
+
+#[tokio::test]
+async fn test_pre_registration_incredible_squaring() -> Result<()> {
+    setup_log();
+
+    // Initialize test harness (node, keys, deployment)
+    let temp_dir = tempfile::TempDir::new()?;
+    let harness = TangleTestHarness::setup(temp_dir).await?;
+    let env = harness.env().clone();
+
+    // Create blueprint-specific context
+    let blueprint_ctx = MyContext {
+        env: env.clone(),
+        call_id: None,
+    };
+
+    // Initialize event handler
+    let handler = XsquareEventHandler::new(&env.clone(), blueprint_ctx)
+        .await
+        .unwrap();
+
+    // Setup service, but we don't register yet
+    let (mut test_env, _, blueprint_id) = harness.setup_services(true).await?;
+    test_env.add_job(handler);
+
+    // Run once for pre-registration
+    test_env.run_runner().await.unwrap();
+    let service_id = harness.request_service(blueprint_id).await.unwrap();
+
+    tokio::spawn(async move {
+        // Run again to actually run the service, now that we have registered
+        test_env.run_runner().await.unwrap();
+    });
+
+    tokio::time::sleep(Duration::from_secs(2)).await;
+
+    // Execute job and verify result
+    let _results = harness
+        .execute_job(
+            service_id,
+            0,
+            vec![InputValue::Uint64(5)],
+            vec![OutputValue::Uint64(25)],
+        )
+        .await?;
+
+    Ok(())
+}

crates/runners/core/src/config.rs

@@ -9,6 +9,12 @@ pub trait BlueprintConfig: Send + Sync + 'static {
     async fn requires_registration(&self, _env: &GadgetConfiguration) -> Result<bool, RunnerError> {
         Ok(true)
     }
+    /// Controls whether the runner should exit after registration
+    ///
+    /// Returns true if the runner should exit after registration, false if it should continue
+    fn should_exit_after_registration(&self) -> bool {
+        true // By default, runners exit after registration
+    }
 }
 
 impl BlueprintConfig for () {}

crates/runners/core/src/runner.rs

@@ -48,6 +48,9 @@ impl BlueprintRunner {
     pub async fn run(&mut self) -> Result<(), Error> {
         if self.config.requires_registration(&self.env).await? {
             self.config.register(&self.env).await?;
+            if self.config.should_exit_after_registration() {
+                return Ok(());
+            }
         }
 
         let mut background_receivers = Vec::new();

crates/runners/eigenlayer/src/bls.rs

@@ -20,15 +20,27 @@ use gadget_utils::evm::get_provider_http;
 pub struct EigenlayerBLSConfig {
     earnings_receiver_address: Address,
     delegation_approver_address: Address,
+    exit_after_register: bool,
 }
 
 impl EigenlayerBLSConfig {
+    /// Creates a new [`EigenlayerBLSConfig`] with the given earnings receiver and delegation approver addresses.
+    ///
+    /// By default, a Runner created with this config will exit after registration (Pre-Registration). To change
+    /// this, use [`EigenlayerBLSConfig::with_exit_after_register`] passing false.
     pub fn new(earnings_receiver_address: Address, delegation_approver_address: Address) -> Self {
         Self {
             earnings_receiver_address,
             delegation_approver_address,
+            exit_after_register: true,
         }
     }
+
+    /// Sets whether the Runner should exit after registration or continue with execution.
+    pub fn with_exit_after_register(mut self, should_exit_after_registration: bool) -> Self {
+        self.exit_after_register = should_exit_after_registration;
+        self
+    }
 }
 
 #[async_trait::async_trait]
@@ -45,6 +57,10 @@ impl BlueprintConfig for EigenlayerBLSConfig {
     async fn requires_registration(&self, env: &GadgetConfiguration) -> Result<bool, Error> {
         requires_registration_bls_impl(env).await
     }
+
+    fn should_exit_after_registration(&self) -> bool {
+        self.exit_after_register
+    }
 }
 
 async fn requires_registration_bls_impl(env: &GadgetConfiguration) -> Result<bool, Error> {

crates/runners/tangle/src/tangle.rs

@@ -46,16 +46,35 @@ impl Default for PriceTargets {
 #[derive(Clone, Default)]
 pub struct TangleConfig {
     pub price_targets: PriceTargets,
+    pub exit_after_register: bool,
+}
+
+impl TangleConfig {
+    pub fn new(price_targets: PriceTargets) -> Self {
+        Self {
+            price_targets,
+            exit_after_register: true,
+        }
+    }
+
+    pub fn with_exit_after_register(mut self, should_exit_after_registration: bool) -> Self {
+        self.exit_after_register = should_exit_after_registration;
+        self
+    }
 }
 
 #[async_trait::async_trait]
 impl BlueprintConfig for TangleConfig {
+    async fn register(&self, env: &GadgetConfiguration) -> Result<(), Error> {
+        register_impl(self.price_targets.clone(), vec![], env).await
+    }
+
     async fn requires_registration(&self, env: &GadgetConfiguration) -> Result<bool, Error> {
         requires_registration_impl(env).await
     }
 
-    async fn register(&self, env: &GadgetConfiguration) -> Result<(), Error> {
-        register_impl(self.clone().price_targets, vec![], env).await
+    fn should_exit_after_registration(&self) -> bool {
+        self.exit_after_register
     }
 }
 

crates/testing-utils/tangle/src/harness.rs

@@ -1,3 +1,4 @@
+use crate::node::transactions::join_operators;
 use crate::Error;
 use crate::{
     keys::inject_tangle_key,
@@ -212,25 +213,66 @@ impl TangleTestHarness {
     }
 
     /// Sets up a complete service environment with initialized event handlers
-    pub async fn setup_services(&self) -> Result<(TangleTestEnv, u64), Error> {
+    ///
+    /// # Returns
+    /// A tuple of the test environment, the service ID, and the blueprint ID i.e., (test_env, service_id, blueprint_id)
+    ///
+    /// # Note
+    /// The Service ID will always be 0 if automatic registration is disabled, as there is not yet a service to have an ID
+    pub async fn setup_services(
+        &self,
+        exit_after_registration: bool,
+    ) -> Result<(TangleTestEnv, u64, u64), Error> {
         // Deploy blueprint
         let blueprint_id = self.deploy_blueprint().await?;
 
+        // Join operators
+        join_operators(&self.client, &self.sr25519_signer)
+            .await
+            .map_err(|e| Error::Setup(e.to_string()))?;
+
         // Setup operator and get service
+        let preferences = self.get_default_operator_preferences();
+        let service_id = if !exit_after_registration {
+            setup_operator_and_service(
+                &self.client,
+                &self.sr25519_signer,
+                blueprint_id,
+                preferences,
+                !exit_after_registration,
+            )
+            .await
+            .map_err(|e| Error::Setup(e.to_string()))?
+        } else {
+            0
+        };
+
+        let config = TangleConfig::new(PriceTargets::default())
+            .with_exit_after_register(exit_after_registration);
+
+        // Create and spawn test environment
+        let test_env = TangleTestEnv::new(config, self.env().clone())?;
+
+        Ok((test_env, service_id, blueprint_id))
+    }
+
+    /// Requests a service with the given blueprint and returns the newly created service ID
+    ///
+    /// This function does not register for a service, it only requests service for a blueprint
+    /// that has already been registered to.
+    pub async fn request_service(&self, blueprint_id: u64) -> Result<u64, Error> {
         let preferences = self.get_default_operator_preferences();
         let service_id = setup_operator_and_service(
             &self.client,
             &self.sr25519_signer,
             blueprint_id,
             preferences,
+            false,
         )
         .await
         .map_err(|e| Error::Setup(e.to_string()))?;
 
-        // Create and spawn test environment
-        let test_env = TangleTestEnv::new(TangleConfig::default(), self.env().clone())?;
-
-        Ok((test_env, service_id))
+        Ok(service_id)
     }
 
     /// Executes a job and verifies its output matches the expected result