feat: update notifications to get milestone by height (#482)

Co-authored-by: David Salami <31099392+Wizdave97@users.noreply.github.com>
Co-authored-by: David Salami <wizdave97@gmail.com>

Author: royvardhan

modules/consensus/tendermint/prover/src/tests/integration_tests.rs

@@ -128,6 +128,7 @@ mod tests {
 		)?;
 
 		let (milestone_number, milestone) = client.get_latest_milestone().await?;
+		trace!("Latest milestone: number {}", milestone_number);
 
 		let latest_height = client.latest_height().await?;
 		let abci_query: AbciQuery =
@@ -143,6 +144,16 @@ mod tests {
 			return Err("Proof should be present".into());
 		}
 
+		let latest_milestone_at_height =
+			client.get_latest_milestone_at_height(latest_height).await?;
+
+		match latest_milestone_at_height {
+			Some((number, _)) => {
+				trace!("Latest milestone at latest height {}: number {}", latest_height, number)
+			},
+			None => trace!("No milestone found at height {}", latest_height),
+		}
+
 		Ok(())
 	}
 

tesseract/consensus/polygon/src/client.rs

@@ -242,6 +242,144 @@ impl HeimdallClient {
 		Ok(abci_query)
 	}
 
+	/// Retrieves the milestone count at a specific height using ABCI query.
+	///
+	/// This method queries the Heimdall node's ABCI store to get the milestone count
+	/// at the specified height, allowing milestone updates even when syncing.
+	///
+	/// # Arguments
+	///
+	/// * `height` - The height at which to query the milestone count
+	///
+	/// # Returns
+	///
+	/// - `Ok(u64)`: The milestone count at the specified height
+	/// - `Err(ProverError)`: If the request fails or the response cannot be parsed
+	///
+	/// # Errors
+	///
+	/// Returns `ProverError` if:
+	/// - The ABCI query fails
+	/// - The height conversion fails
+	/// - The response cannot be deserialized
+	pub async fn get_milestone_count_at_height(
+		&self,
+		height: u64,
+	) -> Result<Option<u64>, ProverError> {
+		let key = vec![0x83];
+
+		let abci_query: AbciQuery = self
+			.http_client
+			.abci_query(
+				Some("/store/milestone/key".to_string()),
+				key,
+				Some(Height::try_from(height).unwrap()),
+				true,
+			)
+			.await
+			.map_err(|e| ProverError::NetworkError(e.to_string()))?;
+
+		let count_bytes = abci_query.value;
+		if count_bytes.is_empty() {
+			return Ok(None); // No milestones yet
+		}
+
+		// The count is stored as a u64 in big-endian format
+		if count_bytes.len() != 8 {
+			return Err(ProverError::ConversionError(
+				"Invalid milestone count bytes length".to_string(),
+			));
+		}
+
+		let mut bytes = [0u8; 8];
+		bytes.copy_from_slice(&count_bytes[..=7]);
+		let count = u64::from_be_bytes(bytes);
+
+		Ok(Some(count))
+	}
+
+	/// Retrieves the latest milestone at a specific height using ABCI query.
+	///
+	/// This method queries the Heimdall node's ABCI store to get the latest milestone
+	/// at the specified height, allowing milestone updates even when syncing.
+	///
+	/// # Arguments
+	///
+	/// * `height` - The height at which to query the latest milestone
+	///
+	/// # Returns
+	///
+	/// - `Ok(Option<(u64, Milestone)>)`: The milestone number and data if available
+	/// - `Err(ProverError)`: If the request fails or the response cannot be parsed
+	///
+	/// # Errors
+	///
+	/// Returns `ProverError` if:
+	/// - The ABCI query fails
+	/// - The height conversion fails
+	/// - The response cannot be deserialized
+	pub async fn get_latest_milestone_at_height(
+		&self,
+		height: u64,
+	) -> Result<Option<(u64, Milestone)>, ProverError> {
+		let count = self.get_milestone_count_at_height(height).await?;
+
+		match count {
+			Some(count) => {
+				let milestone = self.get_milestone_at_height(count, height).await?;
+				Ok(Some((count, milestone)))
+			},
+			None => Ok(None),
+		}
+	}
+
+	/// Retrieves a specific milestone at a specific height using ABCI query.
+	///
+	/// This method queries the Heimdall node's ABCI store to get a specific milestone
+	/// at the specified height.
+	///
+	/// # Arguments
+	///
+	/// * `milestone_number` - The milestone number to retrieve
+	/// * `height` - The height at which to query the milestone
+	///
+	/// # Returns
+	///
+	/// - `Ok(Milestone)`: The milestone data
+	/// - `Err(ProverError)`: If the request fails or the response cannot be parsed
+	///
+	/// # Errors
+	///
+	/// Returns `ProverError` if:
+	/// - The ABCI query fails
+	/// - The height conversion fails
+	/// - The response cannot be deserialized
+	pub async fn get_milestone_at_height(
+		&self,
+		milestone_number: u64,
+		height: u64,
+	) -> Result<Milestone, ProverError> {
+		let mut key = vec![0x81];
+		key.extend_from_slice(&milestone_number.to_be_bytes());
+
+		let abci_query: AbciQuery = self
+			.http_client
+			.abci_query(
+				Some("/store/milestone/key".to_string()),
+				key,
+				Some(Height::try_from(height).unwrap()),
+				true,
+			)
+			.await
+			.map_err(|e| ProverError::NetworkError(e.to_string()))?;
+
+		let milestone = ismp_polygon::Milestone::proto_decode(&abci_query.value).map_err(|e| {
+			ProverError::ConversionError(format!("Failed to decode milestone: {}", e))
+		})?;
+
+		Ok(milestone)
+	}
+
 	/// Fetches an Ethereum block header from the execution RPC client and converts it to a
 	/// CodecHeader.
 	///

tesseract/consensus/polygon/src/notification.rs

@@ -26,7 +26,7 @@ pub async fn consensus_notification(
 	let consensus_state: ConsensusState =
 		ConsensusState::decode(&mut &consensus_state_serialized[..])?;
 
-	let trusted_state: TrustedState = consensus_state.tendermint_state.into();
+	let trusted_state: TrustedState = consensus_state.clone().tendermint_state.into();
 
 	let untrusted_header = client.prover.signed_header(latest_height).await?;
 
@@ -42,46 +42,9 @@ pub async fn consensus_notification(
 		true,
 	);
 
-	let (mut milestone_number, mut milestone) = client.prover.get_latest_milestone().await?;
-	let query_height = untrusted_header.header.height.value() - 1;
-	let mut milestone_proof =
-		client.prover.get_milestone_proof(milestone_number, query_height).await?;
-
-	if milestone_proof.value.is_empty() {
-		milestone_number -= 1;
-		milestone_proof = client.prover.get_milestone_proof(milestone_number, query_height).await?;
-
-		milestone = ismp_polygon::Milestone::proto_decode(&milestone_proof.value)
-			.map_err(|e| anyhow::anyhow!("failed to decode milestone: {}", e))?;
-	}
-
-	let maybe_milestone_update = if milestone.end_block > consensus_state.last_finalized_block {
-		let evm_header = client
-			.prover
-			.fetch_header(milestone.end_block)
-			.await?
-			.ok_or_else(|| anyhow::anyhow!("EVM header not found"))?;
-
-		let merkle_proof = milestone_proof
-			.clone()
-			.proof
-			.map(|p| convert_tm_to_ics_merkle_proof::<ICS23HostFunctions>(&p))
-			.transpose()
-			.map_err(|_| anyhow::anyhow!("bad client state proof"))?
-			.ok_or_else(|| anyhow::anyhow!("proof not found"))?;
-
-		let proof = CommitmentProofBytes::try_from(merkle_proof)
-			.map_err(|e| anyhow::anyhow!("bad client state proof: {}", e))?;
-
-		Some(ismp_polygon::MilestoneUpdate {
-			evm_header,
-			milestone_number,
-			ics23_state_proof: proof.into(),
-			milestone,
-		})
-	} else {
-		None
-	};
+	let maybe_milestone_update =
+		build_milestone_update(client, untrusted_header.header.height.value(), &consensus_state)
+			.await?;
 
 	match validator_set_hash_match.is_ok() && next_validator_set_hash_match.is_ok() {
 		true => {
@@ -138,6 +101,16 @@ pub async fn consensus_notification(
 				let matched_height = height;
 				let matched_header = matched_header.expect("Header must be present if found");
 				let next_validators = client.prover.next_validators(matched_height).await?;
+
+				// Also attempt to construct a milestone update corresponding to the matched header
+				// height
+				let maybe_milestone_update = build_milestone_update(
+					client,
+					matched_header.header.height.value(),
+					&consensus_state,
+				)
+				.await?;
+
 				return Ok(Some(PolygonConsensusUpdate {
 					tendermint_proof: CodecConsensusProof::from(&ConsensusProof::new(
 						matched_header.clone(),
@@ -147,7 +120,7 @@ pub async fn consensus_notification(
 							Some(next_validators)
 						},
 					)),
-					milestone_update: None,
+					milestone_update: maybe_milestone_update,
 				}));
 			} else {
 				log::error!(target: "tesseract", "Fatal error, failed to find any header that matches onchain validator set");
@@ -158,6 +131,62 @@ pub async fn consensus_notification(
 	Ok(None)
 }
 
+async fn build_milestone_update(
+	client: &PolygonPosHost,
+	reference_height: u64,
+	consensus_state: &ConsensusState,
+) -> anyhow::Result<Option<ismp_polygon::MilestoneUpdate>> {
+	let query_height = reference_height.saturating_sub(1);
+	let latest_milestone_at_height =
+		client.prover.get_latest_milestone_at_height(query_height).await?;
+
+	let (milestone_number, milestone) = match latest_milestone_at_height {
+		Some((number, milestone)) => (number, milestone),
+		None => {
+			log::warn!(
+				target: "tesseract",
+				"No milestone found at height {}, falling back to current latest",
+				reference_height
+			);
+			return Ok(None);
+		},
+	};
+
+	let milestone_proof = client.prover.get_milestone_proof(milestone_number, query_height).await?;
+
+	if milestone_proof.value.is_empty() {
+		return Ok(None);
+	}
+
+	if milestone.end_block > consensus_state.last_finalized_block {
+		let evm_header = client
+			.prover
+			.fetch_header(milestone.end_block)
+			.await?
+			.ok_or_else(|| anyhow::anyhow!("EVM header not found"))?;
+
+		let merkle_proof = milestone_proof
+			.clone()
+			.proof
+			.map(|p| convert_tm_to_ics_merkle_proof::<ICS23HostFunctions>(&p))
+			.transpose()
+			.map_err(|_| anyhow::anyhow!("bad client state proof"))?
+			.ok_or_else(|| anyhow::anyhow!("proof not found"))?;
+
+		let proof = CommitmentProofBytes::try_from(merkle_proof)
+			.map_err(|e| anyhow::anyhow!("bad client state proof: {}", e))?;
+
+		Ok(Some(ismp_polygon::MilestoneUpdate {
+			evm_header,
+			milestone_number,
+			ics23_state_proof: proof.into(),
+			milestone,
+		}))
+	} else {
+		Ok(None)
+	}
+}
+
 pub fn convert_tm_to_ics_merkle_proof<H>(
 	tm_proof: &ProofOps,
 ) -> Result<MerkleProof, anyhow::Error> {