feat(proof): enhance aggregation request handling with prover isolation (#682)

- Updated `AggregationRequestKey` to include an optional prover address for isolating aggregate proofs per prover.
- Modified `batch_handler` and `shasta_batch_handler` to utilize the new `new_with_image_id_and_prover` method for creating aggregation keys, allowing for more granular proof management based on the prover.

Author: smtmfft

host/src/server/api/v3/proof/batch_handler.rs

@@ -120,10 +120,11 @@ async fn batch_handler(
     let result = if batch_request.aggregate {
         prove_aggregation(
             &actor,
-            AggregationRequestKey::new_with_image_id(
+            AggregationRequestKey::new_with_image_id_and_prover(
                 batch_request.proof_type,
                 sub_batch_ids.clone(),
                 image_id.clone(),
+                batch_request.prover.to_string(),
             )
             .into(),
             AggregationRequestEntity::new(

host/src/server/api/v3/proof/shasta_handler.rs

@@ -112,68 +112,67 @@ async fn shasta_batch_handler(
         sub_batch_ids,
     ) = process_shasta_batch(&shasta_request, &image_id);
 
-    let result = if shasta_request.aggregate {
-        prove_aggregation(
-            &actor,
-            RequestKey::ShastaAggregation(AggregationRequestKey::new_with_image_id(
-                shasta_request.proof_type,
-                sub_batch_ids.clone(),
-                image_id.clone(),
-            )),
-            RequestEntity::ShastaAggregation(AggregationRequestEntity::new(
-                sub_batch_ids,
-                vec![],
-                shasta_request.proof_type,
-                shasta_request.prover_args.clone(),
-            )),
-            sub_request_keys,
-            sub_request_entities,
-        )
-        .await
+    // Run input step first to reuse cached guest input (both aggregate and non-aggregate)
+    let statuses =
+        prove_many(&actor, sub_input_request_keys, sub_input_request_entities).await?;
+    let is_all_sub_success = statuses
+        .iter()
+        .all(|status| matches!(status, raiko_reqpool::Status::Success { .. }));
+    let result = if !is_all_sub_success {
+        Ok(raiko_reqpool::Status::Registered)
     } else {
-        let statuses =
-            prove_many(&actor, sub_input_request_keys, sub_input_request_entities).await?;
-        let is_all_sub_success = statuses
+        let guest_inputs: Vec<_> = statuses
             .iter()
-            .all(|status| matches!(status, raiko_reqpool::Status::Success { .. }));
-        if !is_all_sub_success {
-            Ok(raiko_reqpool::Status::Registered)
+            .map(|s| match s {
+                raiko_reqpool::Status::Success { proof, .. } => proof.proof.clone().unwrap(),
+                _ => unreachable!(),
+            })
+            .collect();
+        let sub_request_entities_with_input: Vec<_> = sub_request_entities
+            .iter()
+            .zip(guest_inputs)
+            .map(|(entity, guest_input)| match entity {
+                raiko_reqpool::RequestEntity::ShastaProof(e) => {
+                    let mut prover_args = e.prover_args().clone();
+                    prover_args.insert(
+                        PROVER_ARG_SHASTA_GUEST_INPUT.to_string(),
+                        encode_guest_input_str_to_prover_arg_value(&guest_input).expect("wrap"),
+                    );
+                    ShastaProofRequestEntity::new_with_guest_input_entity(
+                        e.guest_input_entity().clone(),
+                        *e.proof_type(),
+                        prover_args,
+                    )
+                    .into()
+                }
+                _ => unreachable!(),
+            })
+            .collect();
+
+        if shasta_request.aggregate {
+            prove_aggregation(
+                &actor,
+                RequestKey::ShastaAggregation(AggregationRequestKey::new_with_image_id_and_prover(
+                    shasta_request.proof_type,
+                    sub_batch_ids.clone(),
+                    image_id.clone(),
+                    shasta_request.prover.to_string(),
+                )),
+                RequestEntity::ShastaAggregation(AggregationRequestEntity::new(
+                    sub_batch_ids,
+                    vec![],
+                    shasta_request.proof_type,
+                    shasta_request.prover_args.clone(),
+                )),
+                sub_request_keys,
+                sub_request_entities_with_input,
+            )
+            .await
         } else {
-            let guest_inputs_of_entities = statuses
-                .iter()
-                .map(|status| match status {
-                    // get saved guest input and pass down to real prover
-                    raiko_reqpool::Status::Success { proof, .. } => proof.proof.clone().unwrap(),
-                    _ => unreachable!("is_all_sub_success checked"),
-                })
-                .collect::<Vec<_>>();
-            let sub_request_entities = sub_request_entities
-                .iter()
-                .zip(guest_inputs_of_entities)
-                .to_owned()
-                .map(|(entity, guest_input)| match entity {
-                    raiko_reqpool::RequestEntity::ShastaProof(request_entity) => {
-                        let mut prover_args = request_entity.prover_args().clone();
-                        prover_args.insert(
-                            PROVER_ARG_SHASTA_GUEST_INPUT.to_string(),
-                            encode_guest_input_str_to_prover_arg_value(&guest_input)
-                                .expect("failed to wrap shasta_guest_input string"),
-                        );
-                        ShastaProofRequestEntity::new_with_guest_input_entity(
-                            request_entity.guest_input_entity().clone(),
-                            *request_entity.proof_type(),
-                            prover_args,
-                        )
-                        .into()
-                    }
-                    _ => unreachable!("Invalid request entity"),
-                })
-                .collect::<Vec<_>>();
-            prove_many(&actor, sub_request_keys, sub_request_entities)
+            prove_many(&actor, sub_request_keys, sub_request_entities_with_input)
                 .await
-                .map(|statuses| {
-                    statuses
-                        .into_iter()
+                .map(|s| {
+                    s.into_iter()
                         .next()
                         .unwrap_or_else(|| raiko_reqpool::Status::Failed {
                             error: "No status returned".to_string(),

reqpool/src/request.rs

@@ -209,6 +209,8 @@ pub struct AggregationRequestKey {
     block_numbers: Vec<u64>,
     /// The image ID for zk provers (optional)
     image_id: Option<ImageId>,
+    /// The prover address (optional); when set, isolates aggregate proofs per prover
+    prover_address: Option<String>,
 }
 
 impl AggregationRequestKey {
@@ -217,6 +219,7 @@ impl AggregationRequestKey {
             proof_type,
             block_numbers,
             image_id: None,
+            prover_address: None,
         }
     }
 
@@ -229,6 +232,22 @@ impl AggregationRequestKey {
             proof_type,
             block_numbers,
             image_id: Some(image_id.clone()),
+            prover_address: None,
+        }
+    }
+
+    /// Create an aggregation key with image ID and prover address for prover isolation.
+    pub fn new_with_image_id_and_prover(
+        proof_type: ProofType,
+        block_numbers: Vec<u64>,
+        image_id: ImageId,
+        prover_address: String,
+    ) -> Self {
+        Self {
+            proof_type,
+            block_numbers,
+            image_id: Some(image_id.clone()),
+            prover_address: Some(prover_address),
         }
     }
 }