fix: use working_set for queries at the current height not archival

clients/sov-celestia/src/client_state/execution.rs

@@ -224,10 +224,11 @@ where
 
     upgraded_client_state.zero_custom_fields();
 
-    // Creates new Sovereign client parameters. The `genesis_da_height` and
-    // `genesis_state_root` are considered immutable properties of the client.
-    // Changing them implies creating a client that could potentially be compatible
-    // with other rollups.
+    // Creates new Sovereign client parameters.
+    //
+    // The `genesis_state_root` are considered the unique and immutable
+    // properties of the client. Changing them implies creating a client
+    // that could potentially be compatible with other rollups.
     let new_sovereign_params = client_state
         .sovereign_params
         .update_on_upgrade(upgraded_client_state.sovereign_params);

mocks/src/relayer/handle/rollup.rs

@@ -1,8 +1,7 @@
 use async_trait::async_trait;
-use ibc_core::client::context::ClientValidationContext;
 use ibc_core::client::types::Height;
 use ibc_core::handler::types::events::IbcEvent;
-use ibc_core::host::types::path::{ClientConsensusStatePath, Path};
+use ibc_core::host::types::path::Path;
 use ibc_core::host::ValidationContext;
 use ibc_query::core::channel::QueryPacketCommitmentRequest;
 use ibc_query::core::client::{QueryClientStateRequest, QueryConsensusStateRequest};
@@ -42,24 +41,33 @@ where
             QueryReq::HostConsensusState(height) => {
                 QueryResp::HostConsensusState(ibc_ctx.host_consensus_state(&height).unwrap().into())
             }
-            QueryReq::ClientState(client_id) => QueryResp::ClientState(
-                ibc_ctx
-                    .get_client_validation_context()
-                    .client_state(&client_id)
-                    .unwrap()
-                    .into(),
-            ),
-            QueryReq::ConsensusState(client_id, height) => QueryResp::ConsensusState(
-                ibc_ctx
-                    .get_client_validation_context()
-                    .consensus_state(&ClientConsensusStatePath::new(
-                        client_id,
-                        height.revision_number(),
-                        height.revision_height(),
-                    ))
-                    .unwrap()
-                    .into(),
-            ),
+            QueryReq::ClientState(client_id) => {
+                let req = QueryClientStateRequest {
+                    client_id,
+                    query_height: None,
+                };
+
+                let resp = ibc_ctx
+                    .ibc
+                    .client_state(req, &mut ibc_ctx.working_set.borrow_mut())
+                    .unwrap();
+
+                QueryResp::ClientState(resp.client_state)
+            }
+            QueryReq::ConsensusState(client_id, height) => {
+                let req = QueryConsensusStateRequest {
+                    client_id,
+                    consensus_height: Some(height),
+                    query_height: None,
+                };
+
+                let resp = ibc_ctx
+                    .ibc
+                    .consensus_state(req, &mut ibc_ctx.working_set.borrow_mut())
+                    .unwrap();
+
+                QueryResp::ConsensusState(resp.consensus_state)
+            }
             QueryReq::Header(target_height, trusted_height) => {
                 let sov_header = self.obtain_ibc_header(target_height, trusted_height);
 

mocks/src/sovereign/rollup.rs

@@ -1,7 +1,5 @@
 //! Defines a mock rollup structure with default configurations and specs
 
-use std::cell::RefCell;
-use std::rc::Rc;
 use std::sync::{Arc, Mutex};
 
 use ibc_client_tendermint::types::Header;
@@ -108,9 +106,7 @@ where
     }
 
     pub fn ibc_ctx<'a>(&'a self, working_set: &'a mut WorkingSet<S>) -> IbcContext<'a, S> {
-        let shared_working_set = Rc::new(RefCell::new(working_set));
-
-        IbcContext::new(&self.runtime.ibc, shared_working_set.clone())
+        IbcContext::new(&self.runtime.ibc, working_set)
     }
 
     pub fn obtain_ibc_header(&self, target_height: Height, trusted_height: Height) -> SovTmHeader {

mocks/src/tests/client.rs

@@ -45,7 +45,7 @@ async fn test_create_client_on_sov() {
         .src_chain_ctx()
         .query(QueryReq::ValueWithProof(
             Path::ClientState(ClientStatePath(rly.dst_client_id().clone())),
-            Some(client_state.latest_height().increment()), // increment as the proof is available as of the next height
+            None,
         ))
         .await
     {
@@ -64,7 +64,7 @@ async fn test_create_client_on_sov() {
                 revision_number: client_state.latest_height().revision_number(),
                 revision_height: client_state.latest_height().revision_height(),
             }),
-            Some(client_state.latest_height().increment()), // increment as the proof is available as of the next height
+            None,
         ))
         .await
     {

modules/sov-ibc/src/context.rs

@@ -44,11 +44,11 @@ impl<'a, S> IbcContext<'a, S>
 where
     S: Spec,
 {
-    pub fn new(
-        ibc: &'a Ibc<S>,
-        working_set: Rc<RefCell<&'a mut WorkingSet<S>>>,
-    ) -> IbcContext<'a, S> {
-        IbcContext { ibc, working_set }
+    pub fn new(ibc: &'a Ibc<S>, working_set: &'a mut WorkingSet<S>) -> IbcContext<'a, S> {
+        IbcContext {
+            ibc,
+            working_set: Rc::new(RefCell::new(working_set)),
+        }
     }
 
     /// Check that the context slot number matches the host height that IBC modules view.

modules/sov-ibc/src/rpc/helpers.rs

@@ -1,9 +1,5 @@
-use std::cell::RefCell;
-use std::rc::Rc;
-
 use borsh::BorshSerialize;
 use ibc_core::client::types::Height;
-use ibc_core::host::ValidationContext;
 use jsonrpsee::core::RpcResult;
 use sov_ibc_transfer::to_jsonrpsee_error;
 use sov_modules_api::{Spec, StateMap, WorkingSet};
@@ -13,22 +9,27 @@ use crate::context::IbcContext;
 use crate::Ibc;
 
 impl<S: Spec> Ibc<S> {
-    /// Determines the query height to use for the given request. If the query
-    /// height is not provided, it queries the host for the current height.
-    pub(super) fn determine_query_height(
+    pub(super) fn handle_request<F, Response>(
         &self,
-        query_height: Option<Height>,
+        request_height: Option<Height>,
         working_set: &mut WorkingSet<S>,
-    ) -> RpcResult<Height> {
-        match query_height {
-            Some(height) => Ok(height),
+        method: F,
+    ) -> Response
+    where
+        F: FnOnce(&IbcContext<'_, S>) -> Response,
+    {
+        match request_height {
+            Some(h) => {
+                // note: IbcContext doesn't take ownership of the archival_working_set.
+                // So it can't be returned from the this method. Instead, a closure is
+                // passed to the method and the result is returned from the closure.
+                let mut archival_working_set = working_set.get_archival_at(h.revision_height());
+                let ibc_ctx = IbcContext::new(self, &mut archival_working_set);
+                method(&ibc_ctx)
+            }
             None => {
-                let ibc_ctx = IbcContext {
-                    ibc: self,
-                    working_set: Rc::new(RefCell::new(working_set)),
-                };
-
-                ibc_ctx.host_height().map_err(to_jsonrpsee_error)
+                let ibc_ctx = IbcContext::new(self, working_set);
+                method(&ibc_ctx)
             }
         }
     }