feat: Optional limit on note-transport fetch

Cargo.toml

@@ -41,7 +41,7 @@ miden-node-proto-build           = { default-features = false, version = "0.12"
 miden-node-rpc                   = { version = "0.12" }
 miden-node-store                 = { version = "0.12" }
 miden-node-utils                 = { version = "0.12" }
-miden-note-transport-proto-build = { default-features = false, version = "0.1" }
+miden-note-transport-proto-build = { default-features = false, git = "https://github.com/0xMiden/miden-note-transport", branch = "ev/fetch-limit" }
 miden-remote-prover              = { features = ["concurrent"], version = "0.12" }
 miden-remote-prover-client       = { default-features = false, features = ["tx-prover"], version = "0.12" }
 

crates/rust-client/src/note_transport/generated/nostd/miden_note_transport.rs

@@ -26,6 +26,8 @@ pub struct FetchNotesRequest {
     pub tags: ::prost::alloc::vec::Vec<u32>,
     #[prost(fixed64, tag = "2")]
     pub cursor: u64,
+    #[prost(fixed32, optional, tag = "3")]
+    pub limit: ::core::option::Option<u32>,
 }
 /// API response for fetching notes
 #[derive(Clone, PartialEq, ::prost::Message)]

crates/rust-client/src/note_transport/generated/std/miden_note_transport.rs

@@ -26,6 +26,8 @@ pub struct FetchNotesRequest {
     pub tags: ::prost::alloc::vec::Vec<u32>,
     #[prost(fixed64, tag = "2")]
     pub cursor: u64,
+    #[prost(fixed32, optional, tag = "3")]
+    pub limit: ::core::option::Option<u32>,
 }
 /// API response for fetching notes
 #[derive(Clone, PartialEq, ::prost::Message)]

crates/rust-client/src/note_transport/grpc.rs

@@ -111,9 +111,14 @@ impl GrpcNoteTransportClient {
         &self,
         tags: &[NoteTag],
         cursor: NoteTransportCursor,
+        limit: Option<u32>,
     ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
         let tags_int = tags.iter().map(NoteTag::as_u32).collect();
-        let request = FetchNotesRequest { tags: tags_int, cursor: cursor.value() };
+        let request = FetchNotesRequest {
+            tags: tags_int,
+            cursor: cursor.value(),
+            limit,
+        };
 
         let response = self
             .api()
@@ -199,8 +204,9 @@ impl super::NoteTransportClient for GrpcNoteTransportClient {
         &self,
         tags: &[NoteTag],
         cursor: NoteTransportCursor,
+        limit: Option<u32>,
     ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
-        self.fetch_notes(tags, cursor).await
+        self.fetch_notes(tags, cursor, limit).await
     }
 
     async fn stream_notes(

crates/rust-client/src/note_transport/mod.rs

@@ -69,7 +69,7 @@ impl<AUTH> Client<AUTH> {
         // Get global cursor
         let cursor = self.store.get_note_transport_cursor().await?;
 
-        let update = NoteTransport::new(api).fetch_notes(cursor, note_tags).await?;
+        let update = NoteTransport::new(api).fetch_notes(note_tags, cursor, None).await?;
 
         self.store.apply_note_transport_update(update).await?;
 
@@ -87,7 +87,7 @@ impl<AUTH> Client<AUTH> {
         let note_tags = self.store.get_unique_note_tags().await?;
 
         let update = NoteTransport::new(api)
-            .fetch_notes(NoteTransportCursor::init(), note_tags)
+            .fetch_notes(note_tags, NoteTransportCursor::init(), None)
             .await?;
 
         self.store.apply_note_transport_update(update).await?;
@@ -141,18 +141,22 @@ impl NoteTransport {
     /// Fetch notes for provided note tags with pagination
     ///
     /// Only notes after the provided cursor are requested.
+    /// If a limit is provided, then only a number of notes up to that limit will be downloaded.
     pub(crate) async fn fetch_notes<I>(
         &mut self,
-        cursor: NoteTransportCursor,
         tags: I,
+        cursor: NoteTransportCursor,
+        limit: Option<u32>,
     ) -> Result<NoteTransportUpdate, ClientError>
     where
         I: IntoIterator<Item = NoteTag>,
     {
         let mut note_updates = vec![];
         // Fetch notes
-        let (note_infos, rcursor) =
-            self.api.fetch_notes(&tags.into_iter().collect::<Vec<_>>(), cursor).await?;
+        let (note_infos, rcursor) = self
+            .api
+            .fetch_notes(&tags.into_iter().collect::<Vec<_>>(), cursor, limit)
+            .await?;
         for note_info in &note_infos {
             // e2ee impl hint:
             // for key in self.store.decryption_keys() try
@@ -207,6 +211,7 @@ pub trait NoteTransportClient: Send + Sync {
         &self,
         tag: &[NoteTag],
         cursor: NoteTransportCursor,
+        limit: Option<u32>,
     ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError>;
 
     /// Stream notes for a given tag

crates/rust-client/src/sync/mod.rs

@@ -166,7 +166,7 @@ where
         // TODO We can run both sync_state, fetch_notes futures in parallel
         let note_transport_update = if let Some(mut note_transport) = note_transport {
             let cursor = self.store.get_note_transport_cursor().await?;
-            Some(note_transport.fetch_notes(cursor, note_tags).await?)
+            Some(note_transport.fetch_notes(note_tags, cursor, None).await?)
         } else {
             None
         };

crates/rust-client/src/test_utils/note_transport.rs

@@ -45,33 +45,35 @@ impl MockNoteTransportNode {
         &self,
         tags: &[NoteTag],
         cursor: NoteTransportCursor,
+        limit: Option<u32>,
     ) -> (Vec<NoteInfo>, NoteTransportCursor) {
-        let mut notes = vec![];
-        let mut rcursor = NoteTransportCursor::init();
+        let mut notesc_unlimited = Vec::new();
+
         for tag in tags {
             // Assumes stored notes are ordered by cursor
-            let tnotes = self
-                .notes
-                .get(tag)
-                .map(|pg_notes| {
-                    // Find first element after cursor
-                    if let Some(pos) = pg_notes.iter().position(|(_, tcursor)| *tcursor > cursor) {
-                        &pg_notes[pos..]
-                    } else {
-                        &[]
+            if let Some(tag_notes) = self.notes.get(tag) {
+                // Find first element after cursor
+                if let Some(pos) = tag_notes.iter().position(|(_, tcursor)| *tcursor > cursor) {
+                    for (note, note_cursor) in &tag_notes[pos..] {
+                        notesc_unlimited.push((note.clone(), *note_cursor));
                     }
-                })
-                .map(Vec::from)
-                .unwrap_or_default();
-            rcursor = rcursor.max(
-                tnotes
-                    .iter()
-                    .map(|(_, cursor)| *cursor)
-                    .max()
-                    .unwrap_or(NoteTransportCursor::init()),
-            );
-            notes.extend(tnotes.into_iter().map(|(note, _)| note).collect::<Vec<_>>());
+                }
+            }
         }
+
+        // Sort mixed-tagged notes by cursor
+        notesc_unlimited.sort_by_key(|(_, cursor)| *cursor);
+
+        // Apply limit if specified
+        let limit_usize = limit.map(|l| l as usize);
+        let notesc_limited: Vec<_> =
+            notesc_unlimited.iter().take(limit_usize.unwrap_or(usize::MAX)).collect();
+
+        let rcursor = notesc_limited.last().map(|(_, cursor)| *cursor).unwrap_or(cursor);
+
+        // Extract notes
+        let notes: Vec<NoteInfo> = notesc_limited.iter().map(|(note, _)| note.clone()).collect();
+
         (notes, rcursor)
     }
 }
@@ -105,8 +107,9 @@ impl MockNoteTransportApi {
         &self,
         tags: &[NoteTag],
         cursor: NoteTransportCursor,
+        limit: Option<u32>,
     ) -> (Vec<NoteInfo>, NoteTransportCursor) {
-        self.mock_node.read().get_notes(tags, cursor)
+        self.mock_node.read().get_notes(tags, cursor, limit)
     }
 }
 
@@ -136,8 +139,9 @@ impl NoteTransportClient for MockNoteTransportApi {
         &self,
         tags: &[NoteTag],
         cursor: NoteTransportCursor,
+        limit: Option<u32>,
     ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
-        Ok(self.fetch_notes(tags, cursor))
+        Ok(self.fetch_notes(tags, cursor, limit))
     }
 
     async fn stream_notes(