From 6a4aff424f7d16d469c1067e525a1c8f5e0209fc Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Mon, 27 Oct 2025 06:53:31 +0000
Subject: [PATCH] Add count value to the to_device iter item.

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/api/client/sync/v3.rs                      | 1 +
 src/api/client/sync/v5/extensions/to_device.rs | 3 ++-
 src/service/users/device.rs                    | 4 ++--
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/api/client/sync/v3.rs b/src/api/client/sync/v3.rs
index ea37d7d5..f5a54f3b 100644
--- a/src/api/client/sync/v3.rs
+++ b/src/api/client/sync/v3.rs
@@ -375,6 +375,7 @@ async fn build_sync_events(
 	let to_device_events = services
 		.users
 		.get_to_device_events(sender_user, sender_device, Some(since), Some(next_batch))
+		.map(at!(1))
 		.collect::<Vec<_>>();
 
 	let device_one_time_keys_count = services
diff --git a/src/api/client/sync/v5/extensions/to_device.rs b/src/api/client/sync/v5/extensions/to_device.rs
index a87e275a..2717ab9f 100644
--- a/src/api/client/sync/v5/extensions/to_device.rs
+++ b/src/api/client/sync/v5/extensions/to_device.rs
@@ -1,6 +1,6 @@
 use futures::StreamExt;
 use ruma::api::client::sync::sync_events::v5::response;
-use tuwunel_core::{self, Result};
+use tuwunel_core::{self, Result, at};
 
 use super::{Connection, SyncInfo};
 
@@ -17,6 +17,7 @@ pub(super) async fn collect(
 	let events: Vec<_> = services
 		.users
 		.get_to_device_events(sender_user, sender_device, None, Some(conn.next_batch))
+		.map(at!(1))
 		.collect()
 		.await;
 
diff --git a/src/service/users/device.rs b/src/service/users/device.rs
index a640623c..ccfb1399 100644
--- a/src/service/users/device.rs
+++ b/src/service/users/device.rs
@@ -319,7 +319,7 @@ pub fn get_to_device_events<'a>(
 	device_id: &'a DeviceId,
 	since: Option<u64>,
 	to: Option<u64>,
-) -> impl Stream<Item = Raw<AnyToDeviceEvent>> + Send + 'a {
+) -> impl Stream<Item = (u64, Raw<AnyToDeviceEvent>)> + Send + 'a {
 	type Key<'a> = (&'a UserId, &'a DeviceId, u64);
 
 	let from = (user_id, device_id, since.map_or(0, |since| since.saturating_add(1)));
@@ -331,7 +331,7 @@ pub fn get_to_device_events<'a>(
 		.ready_take_while(move |((user_id_, device_id_, count), _): &(Key<'_>, _)| {
 			user_id == *user_id_ && device_id == *device_id_ && to.is_none_or(|to| *count <= to)
 		})
-		.map(at!(1))
+		.map(|((_, _, count), event)| (count, event))
 }
 
 #[implement(super::Service)]