diff --git a/proxy_components/engine_tiflash/src/proxy_utils/engine_ext.rs b/proxy_components/engine_tiflash/src/proxy_utils/engine_ext.rs
index d2ebd94aa89..d443d563a5e 100644
--- a/proxy_components/engine_tiflash/src/proxy_utils/engine_ext.rs
+++ b/proxy_components/engine_tiflash/src/proxy_utils/engine_ext.rs
@@ -45,6 +45,10 @@ impl PageStorageExt {
         self.helper().get_wb_size(wb) as usize
     }
 
+    pub fn write_batch_count(&self, wb: RawVoidPtr) -> usize {
+        self.helper().get_wb_count(wb) != 0
+    }
+
     pub fn write_batch_is_empty(&self, wb: RawVoidPtr) -> bool {
         self.helper().is_wb_empty(wb) != 0
     }
diff --git a/proxy_components/engine_tiflash/src/ps_engine/ps_write_batch.rs b/proxy_components/engine_tiflash/src/ps_engine/ps_write_batch.rs
index 1a4df26d655..231fc02282f 100644
--- a/proxy_components/engine_tiflash/src/ps_engine/ps_write_batch.rs
+++ b/proxy_components/engine_tiflash/src/ps_engine/ps_write_batch.rs
@@ -54,9 +54,7 @@ impl ElementaryWriteBatch for PSElementWriteBatch {
     }
 
     fn count(&self) -> usize {
-        // FIXME
-        // TODO
-        0
+        self.ps_ext.write_batch_count(self.ps_wb.ptr)
     }
 
     fn is_empty(&self) -> bool {
diff --git a/proxy_components/mock-engine-store/src/mock_store/mock_engine_store_server.rs b/proxy_components/mock-engine-store/src/mock_store/mock_engine_store_server.rs
index 7e7f68bab8e..7a222edd0da 100644
--- a/proxy_components/mock-engine-store/src/mock_store/mock_engine_store_server.rs
+++ b/proxy_components/mock-engine-store/src/mock_store/mock_engine_store_server.rs
@@ -400,6 +400,7 @@ pub fn gen_engine_store_server_helper(
             fn_wb_put_page: Some(ffi_mockps_wb_put_page),
             fn_wb_del_page: Some(ffi_mockps_wb_del_page),
             fn_get_wb_size: Some(ffi_mockps_get_wb_size),
+            fn_get_wb_count: Some(ffi_mockps_get_wb_count),
             fn_is_wb_empty: Some(ffi_mockps_is_wb_empty),
             fn_handle_merge_wb: Some(ffi_mockps_handle_merge_wb),
             fn_handle_clear_wb: Some(ffi_mockps_handle_clear_wb),
diff --git a/proxy_components/mock-engine-store/src/mock_store/mock_page_storage.rs b/proxy_components/mock-engine-store/src/mock_store/mock_page_storage.rs
index 437425d4b43..eb76ac110f3 100644
--- a/proxy_components/mock-engine-store/src/mock_store/mock_page_storage.rs
+++ b/proxy_components/mock-engine-store/src/mock_store/mock_page_storage.rs
@@ -115,6 +115,11 @@ pub unsafe extern "C" fn ffi_mockps_get_wb_size(wb: RawVoidPtr) -> u64 {
     wb.data.len() as u64
 }
 
+pub unsafe extern "C" fn ffi_mockps_get_wb_count(wb: RawVoidPtr) -> u64 {
+    let wb = <&mut MockPSWriteBatch as From<RawVoidPtr>>::from(wb);
+    wb.data.len() as u64
+}
+
 pub unsafe extern "C" fn ffi_mockps_is_wb_empty(wb: RawVoidPtr) -> u8 {
     let wb = <&mut MockPSWriteBatch as From<RawVoidPtr>>::from(wb);
     u8::from(wb.data.is_empty())
diff --git a/proxy_components/proxy_ffi/src/engine_store_helper_impls.rs b/proxy_components/proxy_ffi/src/engine_store_helper_impls.rs
index cbb357a0e48..e510913096e 100644
--- a/proxy_components/proxy_ffi/src/engine_store_helper_impls.rs
+++ b/proxy_components/proxy_ffi/src/engine_store_helper_impls.rs
@@ -531,6 +531,11 @@ impl EngineStoreServerHelper {
         unsafe { (self.ps.fn_get_wb_size.into_inner())(wb) }
     }
 
+    pub fn get_wb_count(&self, wb: RawVoidPtr) -> usize {
+        debug_assert!(self.ps.fn_get_wb_count.is_some());
+        unsafe { (self.ps.fn_get_wb_count.into_inner())(wb) }
+    }
+
     pub fn is_wb_empty(&self, wb: RawVoidPtr) -> u8 {
         debug_assert!(self.ps.fn_is_wb_empty.is_some());
         unsafe { (self.ps.fn_is_wb_empty.into_inner())(wb) }
diff --git a/proxy_components/proxy_ffi/src/interfaces.rs b/proxy_components/proxy_ffi/src/interfaces.rs
index ca2e9337c6f..560a3ef4be2 100644
--- a/proxy_components/proxy_ffi/src/interfaces.rs
+++ b/proxy_components/proxy_ffi/src/interfaces.rs
@@ -519,6 +519,8 @@ pub mod root {
             >,
             pub fn_get_wb_size:
                 ::std::option::Option<unsafe extern "C" fn(arg1: root::DB::RawVoidPtr) -> u64>,
+            pub fn_get_wb_count:
+                ::std::option::Option<unsafe extern "C" fn(arg1: root::DB::RawVoidPtr) -> u64>,
             pub fn_is_wb_empty:
                 ::std::option::Option<unsafe extern "C" fn(arg1: root::DB::RawVoidPtr) -> u8>,
             pub fn_handle_merge_wb: ::std::option::Option<
@@ -789,7 +791,7 @@ pub mod root {
                 arg3: root::DB::RawVoidPtr,
             ) -> u32;
         }
-        pub const RAFT_STORE_PROXY_VERSION: u64 = 9679186549381427051;
+        pub const RAFT_STORE_PROXY_VERSION: u64 = 4376098841976428098;
         pub const RAFT_STORE_PROXY_MAGIC_NUMBER: u32 = 324508639;
     }
 }
diff --git a/raftstore-proxy/ffi/src/RaftStoreProxyFFI/@version b/raftstore-proxy/ffi/src/RaftStoreProxyFFI/@version
index 59c3d0bca37..7be83642c27 100644
--- a/raftstore-proxy/ffi/src/RaftStoreProxyFFI/@version
+++ b/raftstore-proxy/ffi/src/RaftStoreProxyFFI/@version
@@ -1,3 +1,3 @@
 #pragma once
 #include <cstdint>
-namespace DB { constexpr uint64_t RAFT_STORE_PROXY_VERSION = 9679186549381427051ull; }
\ No newline at end of file
+namespace DB { constexpr uint64_t RAFT_STORE_PROXY_VERSION = 4376098841976428098ull; }
\ No newline at end of file
diff --git a/raftstore-proxy/ffi/src/RaftStoreProxyFFI/ProxyFFI.h b/raftstore-proxy/ffi/src/RaftStoreProxyFFI/ProxyFFI.h
index 3dd1feba12a..0de3827fd3f 100644
--- a/raftstore-proxy/ffi/src/RaftStoreProxyFFI/ProxyFFI.h
+++ b/raftstore-proxy/ffi/src/RaftStoreProxyFFI/ProxyFFI.h
@@ -297,6 +297,7 @@ struct PageStorageInterfaces {
   void (*fn_wb_put_page)(RawVoidPtr, BaseBuffView, BaseBuffView);
   void (*fn_wb_del_page)(RawVoidPtr, BaseBuffView);
   uint64_t (*fn_get_wb_size)(RawVoidPtr);
+  uint64_t (*fn_get_wb_count)(RawVoidPtr);
   uint8_t (*fn_is_wb_empty)(RawVoidPtr);
   void (*fn_handle_merge_wb)(RawVoidPtr, RawVoidPtr);
   void (*fn_handle_clear_wb)(RawVoidPtr);