fix: Correct the value returned for replicated vs non replicated execution in system api

rs/system_api/src/lib.rs

@@ -3396,15 +3396,18 @@ impl SystemApi for SystemApiImpl {
         let result = match &self.api_type {
             ApiType::Start { .. }
             | ApiType::Init { .. }
-            | ApiType::ReplyCallback { .. }
-            | ApiType::RejectCallback { .. }
             | ApiType::Cleanup { .. }
             | ApiType::PreUpgrade { .. }
-            | ApiType::InspectMessage { .. }
             | ApiType::Update { .. }
             | ApiType::SystemTask { .. }
             | ApiType::ReplicatedQuery { .. } => Ok(1),
-            ApiType::NonReplicatedQuery { .. } => Ok(0),
+            ApiType::ReplyCallback { .. } | ApiType::RejectCallback { .. } => {
+                match self.execution_parameters.execution_mode {
+                    ExecutionMode::NonReplicated => Ok(0),
+                    ExecutionMode::Replicated => Ok(1),
+                }
+            }
+            ApiType::InspectMessage { .. } | ApiType::NonReplicatedQuery { .. } => Ok(0),
         };
         trace_syscall!(self, ic0_in_replicated_execution, result);
         result

rs/system_api/tests/common/mod.rs

@@ -99,6 +99,10 @@ impl ApiTypeBuilder {
         )
     }
 
+    pub fn build_replicated_query_api() -> ApiType {
+        ApiType::replicated_query(UNIX_EPOCH, vec![], user_test_id(1).get())
+    }
+
     pub fn build_non_replicated_query_api() -> ApiType {
         ApiType::non_replicated_query(
             UNIX_EPOCH,
@@ -175,6 +179,27 @@ impl ApiTypeBuilder {
             0.into(),
         )
     }
+
+    pub fn build_inspect_message_api() -> ApiType {
+        ApiType::inspect_message(
+            PrincipalId::new_anonymous(),
+            "test".to_string(),
+            vec![],
+            UNIX_EPOCH,
+        )
+    }
+
+    pub fn build_start_api() -> ApiType {
+        ApiType::start(UNIX_EPOCH)
+    }
+
+    pub fn build_init_api() -> ApiType {
+        ApiType::init(UNIX_EPOCH, vec![], user_test_id(1).get())
+    }
+
+    pub fn build_pre_upgrade_api() -> ApiType {
+        ApiType::pre_upgrade(UNIX_EPOCH, user_test_id(1).get())
+    }
 }
 
 pub fn get_system_api(

rs/system_api/tests/system_api.rs

@@ -2069,3 +2069,40 @@ fn test_save_log_message_keeps_total_log_size_limited() {
     assert_eq!(log.records().len(), initial_records_number + 1);
     assert_le!(log.used_space(), MAX_ALLOWED_CANISTER_LOG_BUFFER_SIZE);
 }
+
+#[test]
+fn in_replicated_execution_works_correctly() {
+    // The following should execute in replicated mode.
+    for api_type in &[
+        ApiTypeBuilder::build_update_api(),
+        ApiTypeBuilder::build_system_task_api(),
+        ApiTypeBuilder::build_start_api(),
+        ApiTypeBuilder::build_init_api(),
+        ApiTypeBuilder::build_pre_upgrade_api(),
+        ApiTypeBuilder::build_replicated_query_api(),
+        ApiTypeBuilder::build_reply_api(Cycles::new(0)),
+        ApiTypeBuilder::build_reject_api(RejectContext::new(RejectCode::CanisterReject, "error")),
+    ] {
+        let cycles_account_manager = CyclesAccountManagerBuilder::new().build();
+        let system_state = SystemStateBuilder::default().build();
+        let api = get_system_api(api_type.clone(), &system_state, cycles_account_manager);
+        assert_eq!(api.ic0_in_replicated_execution(), Ok(1));
+    }
+
+    // The following should execute in non-replicated mode.
+    for api_type in &[
+        ApiTypeBuilder::build_non_replicated_query_api(),
+        ApiTypeBuilder::build_composite_query_api(),
+        ApiTypeBuilder::build_composite_reply_api(Cycles::new(0)),
+        ApiTypeBuilder::build_composite_reject_api(RejectContext::new(
+            RejectCode::CanisterReject,
+            "error",
+        )),
+        ApiTypeBuilder::build_inspect_message_api(),
+    ] {
+        let cycles_account_manager = CyclesAccountManagerBuilder::new().build();
+        let system_state = SystemStateBuilder::default().build();
+        let api = get_system_api(api_type.clone(), &system_state, cycles_account_manager);
+        assert_eq!(api.ic0_in_replicated_execution(), Ok(0));
+    }
+}