uwb(uci): Implement CR-461

Switches most session UCI commands from using app provided session id to
a UWBS generated session handle.

Since UWB framework needs to be both 1.0 & 2.0 compatible, introduced a
new `session_token` term to refer to session fields in UCI params that
could either be session_id (FIRA 1.0) or session_handle (FIRA 2.0). The
conversion from `session_id` to `session_handle` is managed inside
uci_manager. JNI layer will continue to use the existing session_id as
identified. However, since we share the same data structures across UCI
core & JNI, there is some leakage of the terminology to the JNI layer.

Bug: 270089497
Test: atest CtsUwbMultiDeviceTestCase_FiraRangingTests
Test: atest libuwb_uci_packets libuwb_core_tests
(cherry picked from https://android-review.googlesource.com/q/commit:f9d9c4a4457183e1fef999aa352f609e0ca9bf4c)
Merged-In: I93107a1d53b814053dfa1ad5b4425a4aba8509f2
Change-Id: I93107a1d53b814053dfa1ad5b4425a4aba8509f2

12 files changed, 792 insertions, 371 deletions

diff --git a/src/rust/uwb_core/src/params/uci_packets.rs b/src/rust/uwb_core/src/params/uci_packets.rs
index 2e0941e..dfd08d6 100644
--- a/src/rust/uwb_core/src/params/uci_packets.rs
+++ b/src/rust/uwb_core/src/params/uci_packets.rs
@@ -38,6 +38,10 @@ use crate::error::Error;
 pub type SessionId = u32;
 /// The type of the sub-session identifier.
 pub type SubSessionId = u32;
+/// The type of the session handle.
+pub type SessionHandle = u32;
+/// Generic type used to represent either a session id or session handle.
+pub type SessionToken = u32;
 
 /// Wrap the original AppConfigTlv type to redact the PII fields when logging.
 #[derive(Clone, PartialEq)]
diff --git a/src/rust/uwb_core/src/proto/mappings.rs b/src/rust/uwb_core/src/proto/mappings.rs
index 0d21c09..2da89f1 100644
--- a/src/rust/uwb_core/src/proto/mappings.rs
+++ b/src/rust/uwb_core/src/proto/mappings.rs
@@ -876,7 +876,7 @@ impl From<SessionRangeData> for ProtoSessionRangeData {
     fn from(item: SessionRangeData) -> Self {
         let mut result = Self::new();
         result.set_sequence_number(item.sequence_number);
-        result.set_session_id(item.session_id);
+        result.set_session_id(item.session_token);
         result.set_current_ranging_interval_ms(item.current_ranging_interval_ms);
         result.set_ranging_measurement_type(item.ranging_measurement_type.into());
         match to_proto_ranging_measurements(item.ranging_measurements) {
diff --git a/src/rust/uwb_core/src/session/session_manager.rs b/src/rust/uwb_core/src/session/session_manager.rs
index e9432ba..d53ecc7 100644
--- a/src/rust/uwb_core/src/session/session_manager.rs
+++ b/src/rust/uwb_core/src/session/session_manager.rs
@@ -294,7 +294,7 @@ impl<T: UciManager> SessionManagerActor<T> {
 
     fn handle_uci_notification(&mut self, notf: UciSessionNotification) {
         match notf {
-            UciSessionNotification::Status { session_id, session_state, reason_code } => {
+            UciSessionNotification::Status { session_token, session_state, reason_code } => {
                 let reason_code = match ReasonCode::try_from(reason_code) {
                     Ok(r) => r,
                     Err(_) => {
@@ -306,21 +306,21 @@ impl<T: UciManager> SessionManagerActor<T> {
                     }
                 };
                 if session_state == SessionState::SessionStateDeinit {
-                    debug!("Session {} is deinitialized", session_id);
-                    let _ = self.active_sessions.remove(&session_id);
+                    debug!("Session {} is deinitialized", session_token);
+                    let _ = self.active_sessions.remove(&session_token);
                     let _ = self.session_notf_sender.send(SessionNotification::SessionState {
-                        session_id,
+                        session_id: session_token,
                         session_state,
                         reason_code,
                     });
                     return;
                 }
 
-                match self.active_sessions.get_mut(&session_id) {
+                match self.active_sessions.get_mut(&session_token) {
                     Some(session) => {
                         session.on_session_status_changed(session_state);
                         let _ = self.session_notf_sender.send(SessionNotification::SessionState {
-                            session_id,
+                            session_id: session_token,
                             session_state,
                             reason_code,
                         });
@@ -328,36 +328,36 @@ impl<T: UciManager> SessionManagerActor<T> {
                     None => {
                         warn!(
                             "Received notification of the unknown Session {}: {:?}, {:?}",
-                            session_id, session_state, reason_code
+                            session_token, session_state, reason_code
                         );
                     }
                 }
             }
             UciSessionNotification::UpdateControllerMulticastList {
-                session_id,
+                session_token,
                 remaining_multicast_list_size: _,
                 status_list,
-            } => match self.active_sessions.get_mut(&session_id) {
+            } => match self.active_sessions.get_mut(&session_token) {
                 Some(session) => session.on_controller_multicast_list_udpated(status_list),
                 None => {
                     warn!(
                         "Received the notification of the unknown Session {}: {:?}",
-                        session_id, status_list
+                        session_token, status_list
                     );
                 }
             },
             UciSessionNotification::SessionInfo(range_data) => {
-                if self.active_sessions.get(&range_data.session_id).is_some() {
+                if self.active_sessions.get(&range_data.session_token).is_some() {
                     let _ = self.session_notf_sender.send(SessionNotification::RangeData {
-                        session_id: range_data.session_id,
+                        session_id: range_data.session_token,
                         range_data,
                     });
                 } else {
                     warn!("Received range data of the unknown Session: {:?}", range_data);
                 }
             }
-            UciSessionNotification::DataCredit { session_id, credit_availability: _ } => {
-                match self.active_sessions.get(&session_id) {
+            UciSessionNotification::DataCredit { session_token, credit_availability: _ } => {
+                match self.active_sessions.get(&session_token) {
                     Some(_) => {
                         /*
                          * TODO(b/270443790): Handle the DataCredit notification in the new
@@ -367,17 +367,17 @@ impl<T: UciManager> SessionManagerActor<T> {
                     None => {
                         warn!(
                             "Received the Data Credit notification for an unknown Session {}",
-                            session_id
+                            session_token
                         );
                     }
                 }
             }
             UciSessionNotification::DataTransferStatus {
-                session_id,
+                session_token,
                 uci_sequence_number: _,
                 status: _,
             } => {
-                match self.active_sessions.get(&session_id) {
+                match self.active_sessions.get(&session_token) {
                     Some(_) => {
                         /*
                          * TODO(b/270443790): Handle the DataTransferStatus notification in the
@@ -387,7 +387,7 @@ impl<T: UciManager> SessionManagerActor<T> {
                     None => {
                         warn!(
                             "Received a Data Transfer Status notification for unknown Session {}",
-                            session_id
+                            session_token
                         );
                     }
                 }
@@ -481,7 +481,7 @@ pub(crate) mod test_utils {
     pub(crate) fn session_range_data(session_id: SessionId) -> SessionRangeData {
         SessionRangeData {
             sequence_number: 1,
-            session_id,
+            session_token: session_id,
             current_ranging_interval_ms: 3,
             ranging_measurement_type: RangingMeasurementType::TwoWay,
             ranging_measurements: RangingMeasurements::ShortAddressTwoWay(vec![
@@ -512,7 +512,7 @@ pub(crate) mod test_utils {
         session_state: SessionState,
     ) -> UciNotification {
         UciNotification::Session(UciSessionNotification::Status {
-            session_id,
+            session_token: session_id,
             session_state,
             reason_code: ReasonCode::StateChangeWithSessionManagementCommands.into(),
         })
@@ -780,7 +780,7 @@ mod tests {
             setup_session_manager(move |uci_manager| {
                 let multicast_list_notf = vec![UciNotification::Session(
                     UciSessionNotification::UpdateControllerMulticastList {
-                        session_id,
+                        session_token: session_id,
                         remaining_multicast_list_size: 1,
                         status_list: vec![ControleeStatus {
                             mac_address: 0x13,
diff --git a/src/rust/uwb_core/src/session/uwb_session.rs b/src/rust/uwb_core/src/session/uwb_session.rs
index 91bd6ce..37fcba4 100644
--- a/src/rust/uwb_core/src/session/uwb_session.rs
+++ b/src/rust/uwb_core/src/session/uwb_session.rs
@@ -178,6 +178,8 @@ impl<T: UciManager> UwbSessionActor<T> {
     async fn initialize(&mut self, params: AppConfigParams) -> Result<Response> {
         debug_assert!(*self.state_receiver.borrow() == SessionState::SessionStateDeinit);
 
+        // TODO(b/279669973): Support CR-461 fully here. Need to wait for session init rsp.
+        // But, that does not seem to be fully plumbed up in session_manager yet.
         self.uci_manager.session_init(self.session_id, self.session_type).await?;
         self.wait_state(SessionState::SessionStateInit).await?;
 
diff --git a/src/rust/uwb_core/src/uci/command.rs b/src/rust/uwb_core/src/uci/command.rs
index dc06dab..e9e5dc0 100644
--- a/src/rust/uwb_core/src/uci/command.rs
+++ b/src/rust/uwb_core/src/uci/command.rs
@@ -20,7 +20,7 @@ use log::error;
 use crate::error::{Error, Result};
 use crate::params::uci_packets::{
     AppConfigTlv, AppConfigTlvType, Controlees, CountryCode, DeviceConfigId, DeviceConfigTlv,
-    ResetConfig, SessionId, SessionType, UpdateMulticastListAction,
+    ResetConfig, SessionId, SessionToken, SessionType, UpdateMulticastListAction,
 };
 use uwb_uci_packets::{build_session_update_controller_multicast_list_cmd, GroupId, MessageType};
 
@@ -44,40 +44,40 @@ pub enum UciCommand {
         session_type: SessionType,
     },
     SessionDeinit {
-        session_id: SessionId,
+        session_token: SessionToken,
     },
     SessionSetAppConfig {
-        session_id: SessionId,
+        session_token: SessionToken,
         config_tlvs: Vec<AppConfigTlv>,
     },
     SessionGetAppConfig {
-        session_id: SessionId,
+        session_token: SessionToken,
         app_cfg: Vec<AppConfigTlvType>,
     },
     SessionGetCount,
     SessionGetState {
-        session_id: SessionId,
+        session_token: SessionToken,
     },
     SessionUpdateControllerMulticastList {
-        session_id: SessionId,
+        session_token: SessionToken,
         action: UpdateMulticastListAction,
         controlees: Controlees,
     },
     SessionUpdateDtTagRangingRounds {
-        session_id: u32,
+        session_token: u32,
         ranging_round_indexes: Vec<u8>,
     },
     SessionQueryMaxDataSize {
-        session_id: SessionId,
+        session_token: SessionToken,
     },
     SessionStart {
-        session_id: SessionId,
+        session_token: SessionToken,
     },
     SessionStop {
-        session_id: SessionId,
+        session_token: SessionToken,
     },
     SessionGetRangingCount {
-        session_id: SessionId,
+        session_token: SessionToken,
     },
     AndroidSetCountryCode {
         country_code: CountryCode,
@@ -99,21 +99,27 @@ impl TryFrom<UciCommand> for uwb_uci_packets::UciControlPacket {
             UciCommand::SessionInit { session_id, session_type } => {
                 uwb_uci_packets::SessionInitCmdBuilder { session_id, session_type }.build().into()
             }
-            UciCommand::SessionDeinit { session_id } => {
-                uwb_uci_packets::SessionDeinitCmdBuilder { session_id }.build().into()
+            UciCommand::SessionDeinit { session_token } => {
+                uwb_uci_packets::SessionDeinitCmdBuilder { session_token }.build().into()
             }
             UciCommand::CoreGetDeviceInfo => {
                 uwb_uci_packets::GetDeviceInfoCmdBuilder {}.build().into()
             }
             UciCommand::CoreGetCapsInfo => uwb_uci_packets::GetCapsInfoCmdBuilder {}.build().into(),
-            UciCommand::SessionGetState { session_id } => {
-                uwb_uci_packets::SessionGetStateCmdBuilder { session_id }.build().into()
-            }
-            UciCommand::SessionUpdateControllerMulticastList { session_id, action, controlees } => {
-                build_session_update_controller_multicast_list_cmd(session_id, action, controlees)
-                    .map_err(|_| Error::BadParameters)?
-                    .into()
+            UciCommand::SessionGetState { session_token } => {
+                uwb_uci_packets::SessionGetStateCmdBuilder { session_token }.build().into()
             }
+            UciCommand::SessionUpdateControllerMulticastList {
+                session_token,
+                action,
+                controlees,
+            } => build_session_update_controller_multicast_list_cmd(
+                session_token,
+                action,
+                controlees,
+            )
+            .map_err(|_| Error::BadParameters)?
+            .into(),
             UciCommand::CoreSetConfig { config_tlvs } => {
                 uwb_uci_packets::SetConfigCmdBuilder { tlvs: config_tlvs }.build().into()
             }
@@ -122,30 +128,31 @@ impl TryFrom<UciCommand> for uwb_uci_packets::UciControlPacket {
             }
             .build()
             .into(),
-            UciCommand::SessionSetAppConfig { session_id, config_tlvs } => {
+            UciCommand::SessionSetAppConfig { session_token, config_tlvs } => {
                 uwb_uci_packets::SessionSetAppConfigCmdBuilder {
-                    session_id,
+                    session_token,
                     tlvs: config_tlvs.into_iter().map(|tlv| tlv.into_inner()).collect(),
                 }
                 .build()
                 .into()
             }
-            UciCommand::SessionGetAppConfig { session_id, app_cfg } => {
+            UciCommand::SessionGetAppConfig { session_token, app_cfg } => {
                 uwb_uci_packets::SessionGetAppConfigCmdBuilder {
-                    session_id,
+                    session_token,
                     app_cfg: app_cfg.into_iter().map(u8::from).collect(),
                 }
                 .build()
                 .into()
             }
-            UciCommand::SessionUpdateDtTagRangingRounds { session_id, ranging_round_indexes } => {
-                uwb_uci_packets::SessionUpdateDtTagRangingRoundsCmdBuilder {
-                    session_id,
-                    ranging_round_indexes,
-                }
-                .build()
-                .into()
+            UciCommand::SessionUpdateDtTagRangingRounds {
+                session_token,
+                ranging_round_indexes,
+            } => uwb_uci_packets::SessionUpdateDtTagRangingRoundsCmdBuilder {
+                session_token,
+                ranging_round_indexes,
             }
+            .build()
+            .into(),
             UciCommand::AndroidGetPowerStats => {
                 uwb_uci_packets::AndroidGetPowerStatsCmdBuilder {}.build().into()
             }
@@ -166,17 +173,17 @@ impl TryFrom<UciCommand> for uwb_uci_packets::UciControlPacket {
                 uwb_uci_packets::DeviceResetCmdBuilder { reset_config }.build().into()
             }
             // UCI Session Control Commands
-            UciCommand::SessionStart { session_id } => {
-                uwb_uci_packets::SessionStartCmdBuilder { session_id }.build().into()
+            UciCommand::SessionStart { session_token } => {
+                uwb_uci_packets::SessionStartCmdBuilder { session_token }.build().into()
             }
-            UciCommand::SessionStop { session_id } => {
-                uwb_uci_packets::SessionStopCmdBuilder { session_id }.build().into()
+            UciCommand::SessionStop { session_token } => {
+                uwb_uci_packets::SessionStopCmdBuilder { session_token }.build().into()
             }
-            UciCommand::SessionGetRangingCount { session_id } => {
-                uwb_uci_packets::SessionGetRangingCountCmdBuilder { session_id }.build().into()
+            UciCommand::SessionGetRangingCount { session_token } => {
+                uwb_uci_packets::SessionGetRangingCountCmdBuilder { session_token }.build().into()
             }
-            UciCommand::SessionQueryMaxDataSize { session_id } => {
-                uwb_uci_packets::SessionQueryMaxDataSizeCmdBuilder { session_id }.build().into()
+            UciCommand::SessionQueryMaxDataSize { session_token } => {
+                uwb_uci_packets::SessionQueryMaxDataSizeCmdBuilder { session_token }.build().into()
             }
         };
         Ok(packet)
@@ -266,27 +273,27 @@ mod tests {
             .into()
         );
 
-        cmd = UciCommand::SessionDeinit { session_id: 1 };
+        cmd = UciCommand::SessionDeinit { session_token: 1 };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
-            uwb_uci_packets::SessionDeinitCmdBuilder { session_id: 1 }.build().into()
+            uwb_uci_packets::SessionDeinitCmdBuilder { session_token: 1 }.build().into()
         );
 
-        cmd = UciCommand::SessionSetAppConfig { session_id: 1, config_tlvs: vec![] };
+        cmd = UciCommand::SessionSetAppConfig { session_token: 1, config_tlvs: vec![] };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
-            uwb_uci_packets::SessionSetAppConfigCmdBuilder { session_id: 1, tlvs: vec![] }
+            uwb_uci_packets::SessionSetAppConfigCmdBuilder { session_token: 1, tlvs: vec![] }
                 .build()
                 .into()
         );
 
-        cmd = UciCommand::SessionGetAppConfig { session_id: 1, app_cfg: vec![] };
+        cmd = UciCommand::SessionGetAppConfig { session_token: 1, app_cfg: vec![] };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
-            uwb_uci_packets::SessionGetAppConfigCmdBuilder { session_id: 1, app_cfg: vec![] }
+            uwb_uci_packets::SessionGetAppConfigCmdBuilder { session_token: 1, app_cfg: vec![] }
                 .build()
                 .into()
         );
@@ -295,15 +302,15 @@ mod tests {
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(packet, uwb_uci_packets::SessionGetCountCmdBuilder {}.build().into());
 
-        cmd = UciCommand::SessionGetState { session_id: 1 };
+        cmd = UciCommand::SessionGetState { session_token: 1 };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
-            uwb_uci_packets::SessionGetStateCmdBuilder { session_id: 1 }.build().into()
+            uwb_uci_packets::SessionGetStateCmdBuilder { session_token: 1 }.build().into()
         );
 
         cmd = UciCommand::SessionUpdateControllerMulticastList {
-            session_id: 1,
+            session_token: 1,
             action: UpdateMulticastListAction::AddControlee,
             controlees: Controlees::NoSessionKey(vec![]),
         };
@@ -321,43 +328,46 @@ mod tests {
         );
 
         cmd = UciCommand::SessionUpdateDtTagRangingRounds {
-            session_id: 1,
+            session_token: 1,
             ranging_round_indexes: vec![0],
         };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
             uwb_uci_packets::SessionUpdateDtTagRangingRoundsCmdBuilder {
-                session_id: 1,
+                session_token: 1,
                 ranging_round_indexes: vec![0]
             }
             .build()
             .into()
         );
 
-        cmd = UciCommand::SessionQueryMaxDataSize { session_id: 1 };
+        cmd = UciCommand::SessionQueryMaxDataSize { session_token: 1 };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
-            uwb_uci_packets::SessionQueryMaxDataSizeCmdBuilder { session_id: 1 }.build().into()
+            uwb_uci_packets::SessionQueryMaxDataSizeCmdBuilder { session_token: 1 }.build().into()
         );
 
-        cmd = UciCommand::SessionStart { session_id: 1 };
+        cmd = UciCommand::SessionStart { session_token: 1 };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
-            uwb_uci_packets::SessionStartCmdBuilder { session_id: 1 }.build().into()
+            uwb_uci_packets::SessionStartCmdBuilder { session_token: 1 }.build().into()
         );
 
-        cmd = UciCommand::SessionStop { session_id: 1 };
+        cmd = UciCommand::SessionStop { session_token: 1 };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
-        assert_eq!(packet, uwb_uci_packets::SessionStopCmdBuilder { session_id: 1 }.build().into());
+        assert_eq!(
+            packet,
+            uwb_uci_packets::SessionStopCmdBuilder { session_token: 1 }.build().into()
+        );
 
-        cmd = UciCommand::SessionGetRangingCount { session_id: 1 };
+        cmd = UciCommand::SessionGetRangingCount { session_token: 1 };
         packet = uwb_uci_packets::UciControlPacket::try_from(cmd.clone()).unwrap();
         assert_eq!(
             packet,
-            uwb_uci_packets::SessionGetRangingCountCmdBuilder { session_id: 1 }.build().into()
+            uwb_uci_packets::SessionGetRangingCountCmdBuilder { session_token: 1 }.build().into()
         );
 
         let country_code: [u8; 2] = [85, 83];
diff --git a/src/rust/uwb_core/src/uci/notification.rs b/src/rust/uwb_core/src/uci/notification.rs
index 8007158..e06a037 100644
--- a/src/rust/uwb_core/src/uci/notification.rs
+++ b/src/rust/uwb_core/src/uci/notification.rs
@@ -23,7 +23,7 @@ use crate::params::uci_packets::{
     ControleeStatus, CreditAvailability, DataRcvStatusCode, DataTransferNtfStatusCode, DeviceState,
     ExtendedAddressDlTdoaRangingMeasurement, ExtendedAddressOwrAoaRangingMeasurement,
     ExtendedAddressTwoWayRangingMeasurement, FiraComponent, RangingMeasurementType, RawUciMessage,
-    SessionId, SessionState, ShortAddressDlTdoaRangingMeasurement,
+    SessionState, SessionToken, ShortAddressDlTdoaRangingMeasurement,
     ShortAddressOwrAoaRangingMeasurement, ShortAddressTwoWayRangingMeasurement, StatusCode,
 };
 
@@ -52,8 +52,8 @@ pub enum CoreNotification {
 pub enum SessionNotification {
     /// SessionStatusNtf equivalent.
     Status {
-        /// SessionId : u32
-        session_id: SessionId,
+        /// SessionToken : u32
+        session_token: SessionToken,
         /// uwb_uci_packets::SessionState.
         session_state: SessionState,
         /// uwb_uci_packets::Reasoncode.
@@ -61,8 +61,8 @@ pub enum SessionNotification {
     },
     /// SessionUpdateControllerMulticastListNtf equivalent.
     UpdateControllerMulticastList {
-        /// SessionId : u32
-        session_id: SessionId,
+        /// SessionToken : u32
+        session_token: SessionToken,
         /// count of controlees: u8
         remaining_multicast_list_size: usize,
         /// list of controlees.
@@ -72,15 +72,15 @@ pub enum SessionNotification {
     SessionInfo(SessionRangeData),
     /// DataCreditNtf equivalent.
     DataCredit {
-        /// SessionId : u32
-        session_id: SessionId,
+        /// SessionToken : u32
+        session_token: SessionToken,
         /// Credit Availability (for sending Data packets on UWB Session)
         credit_availability: CreditAvailability,
     },
     /// DataTransferStatusNtf equivalent.
     DataTransferStatus {
-        /// SessionId : u32
-        session_id: SessionId,
+        /// SessionToken : u32
+        session_token: SessionToken,
         /// Sequence Number: u8
         uci_sequence_number: u8,
         /// Data Transfer Status Code
@@ -95,7 +95,7 @@ pub struct SessionRangeData {
     pub sequence_number: u32,
 
     /// The identifier of the session.
-    pub session_id: SessionId,
+    pub session_token: SessionToken,
 
     /// The current ranging interval setting in the unit of ms.
     pub current_ranging_interval_ms: u32,
@@ -140,7 +140,7 @@ pub enum RangingMeasurements {
 #[derive(Debug, Clone)]
 pub struct DataRcvNotification {
     /// The identifier of the session on which data transfer is happening.
-    pub session_id: SessionId,
+    pub session_token: SessionToken,
 
     /// The status of the data rx.
     pub status: DataRcvStatusCode,
@@ -166,7 +166,7 @@ impl TryFrom<uwb_uci_packets::UciDataPacket> for DataRcvNotification {
     fn try_from(evt: uwb_uci_packets::UciDataPacket) -> std::result::Result<Self, Self::Error> {
         match evt.specialize() {
             uwb_uci_packets::UciDataPacketChild::UciDataRcv(evt) => Ok(DataRcvNotification {
-                session_id: evt.get_session_id(),
+                session_token: evt.get_session_token(),
                 status: evt.get_status(),
                 uci_sequence_num: evt.get_uci_sequence_number(),
                 source_address: UwbAddress::Extended(evt.get_source_mac_address().to_le_bytes()),
@@ -242,13 +242,13 @@ impl TryFrom<uwb_uci_packets::SessionConfigNotification> for SessionNotification
         use uwb_uci_packets::SessionConfigNotificationChild;
         match evt.specialize() {
             SessionConfigNotificationChild::SessionStatusNtf(evt) => Ok(Self::Status {
-                session_id: evt.get_session_id(),
+                session_token: evt.get_session_token(),
                 session_state: evt.get_session_state(),
                 reason_code: evt.get_reason_code(),
             }),
             SessionConfigNotificationChild::SessionUpdateControllerMulticastListNtf(evt) => {
                 Ok(Self::UpdateControllerMulticastList {
-                    session_id: evt.get_session_id(),
+                    session_token: evt.get_session_token(),
                     remaining_multicast_list_size: evt.get_remaining_multicast_list_size() as usize,
                     status_list: evt.get_controlee_status().clone(),
                 })
@@ -270,12 +270,12 @@ impl TryFrom<uwb_uci_packets::SessionControlNotification> for SessionNotificatio
         match evt.specialize() {
             SessionControlNotificationChild::SessionInfoNtf(evt) => evt.try_into(),
             SessionControlNotificationChild::DataCreditNtf(evt) => Ok(Self::DataCredit {
-                session_id: evt.get_session_id(),
+                session_token: evt.get_session_token(),
                 credit_availability: evt.get_credit_availability(),
             }),
             SessionControlNotificationChild::DataTransferStatusNtf(evt) => {
                 Ok(Self::DataTransferStatus {
-                    session_id: evt.get_session_id(),
+                    session_token: evt.get_session_token(),
                     uci_sequence_number: evt.get_uci_sequence_number(),
                     status: evt.get_status(),
                 })
@@ -363,7 +363,7 @@ impl TryFrom<uwb_uci_packets::SessionInfoNtf> for SessionNotification {
         };
         Ok(Self::SessionInfo(SessionRangeData {
             sequence_number: evt.get_sequence_number(),
-            session_id: evt.get_session_id(),
+            session_token: evt.get_session_token(),
             current_ranging_interval_ms: evt.get_current_ranging_interval(),
             ranging_measurement_type: evt.get_ranging_measurement_type(),
             ranging_measurements,
@@ -538,7 +538,7 @@ mod tests {
         let extended_two_way_session_info_ntf =
             uwb_uci_packets::ExtendedMacTwoWaySessionInfoNtfBuilder {
                 sequence_number: 0x10,
-                session_id: 0x11,
+                session_token: 0x11,
                 rcr_indicator: 0x12,
                 current_ranging_interval: 0x13,
                 two_way_ranging_measurements: vec![extended_measurement.clone()],
@@ -556,7 +556,7 @@ mod tests {
             uci_notification_from_extended_two_way_session_info_ntf,
             UciNotification::Session(SessionNotification::SessionInfo(SessionRangeData {
                 sequence_number: 0x10,
-                session_id: 0x11,
+                session_token: 0x11,
                 ranging_measurement_type: uwb_uci_packets::RangingMeasurementType::TwoWay,
                 current_ranging_interval_ms: 0x13,
                 ranging_measurements: RangingMeasurements::ExtendedAddressTwoWay(vec![
@@ -588,7 +588,7 @@ mod tests {
         };
         let short_two_way_session_info_ntf = uwb_uci_packets::ShortMacTwoWaySessionInfoNtfBuilder {
             sequence_number: 0x10,
-            session_id: 0x11,
+            session_token: 0x11,
             rcr_indicator: 0x12,
             current_ranging_interval: 0x13,
             two_way_ranging_measurements: vec![short_measurement.clone()],
@@ -606,7 +606,7 @@ mod tests {
             uci_notification_from_short_two_way_session_info_ntf,
             UciNotification::Session(SessionNotification::SessionInfo(SessionRangeData {
                 sequence_number: 0x10,
-                session_id: 0x11,
+                session_token: 0x11,
                 ranging_measurement_type: uwb_uci_packets::RangingMeasurementType::TwoWay,
                 current_ranging_interval_ms: 0x13,
                 ranging_measurements: RangingMeasurements::ShortAddressTwoWay(vec![
@@ -634,7 +634,7 @@ mod tests {
         let extended_owr_aoa_session_info_ntf =
             uwb_uci_packets::ExtendedMacOwrAoaSessionInfoNtfBuilder {
                 sequence_number: 0x10,
-                session_id: 0x11,
+                session_token: 0x11,
                 rcr_indicator: 0x12,
                 current_ranging_interval: 0x13,
                 owr_aoa_ranging_measurements: vec![extended_measurement.clone()],
@@ -652,7 +652,7 @@ mod tests {
             uci_notification_from_extended_owr_aoa_session_info_ntf,
             UciNotification::Session(SessionNotification::SessionInfo(SessionRangeData {
                 sequence_number: 0x10,
-                session_id: 0x11,
+                session_token: 0x11,
                 ranging_measurement_type: uwb_uci_packets::RangingMeasurementType::OwrAoa,
                 current_ranging_interval_ms: 0x13,
                 ranging_measurements: RangingMeasurements::ExtendedAddressOwrAoa(
@@ -679,7 +679,7 @@ mod tests {
         };
         let short_owr_aoa_session_info_ntf = uwb_uci_packets::ShortMacOwrAoaSessionInfoNtfBuilder {
             sequence_number: 0x10,
-            session_id: 0x11,
+            session_token: 0x11,
             rcr_indicator: 0x12,
             current_ranging_interval: 0x13,
             owr_aoa_ranging_measurements: vec![short_measurement.clone()],
@@ -697,7 +697,7 @@ mod tests {
             uci_notification_from_short_owr_aoa_session_info_ntf,
             UciNotification::Session(SessionNotification::SessionInfo(SessionRangeData {
                 sequence_number: 0x10,
-                session_id: 0x11,
+                session_token: 0x11,
                 ranging_measurement_type: uwb_uci_packets::RangingMeasurementType::OwrAoa,
                 current_ranging_interval_ms: 0x13,
                 ranging_measurements: RangingMeasurements::ShortAddressOwrAoa(short_measurement),
@@ -710,7 +710,7 @@ mod tests {
     #[test]
     fn test_session_notification_casting_from_session_status_ntf() {
         let session_status_ntf = uwb_uci_packets::SessionStatusNtfBuilder {
-            session_id: 0x20,
+            session_token: 0x20,
             session_state: uwb_uci_packets::SessionState::SessionStateActive,
             reason_code: uwb_uci_packets::ReasonCode::StateChangeWithSessionManagementCommands
                 .into(),
@@ -725,7 +725,7 @@ mod tests {
         assert_eq!(
             uci_notification_from_session_status_ntf,
             UciNotification::Session(SessionNotification::Status {
-                session_id: 0x20,
+                session_token: 0x20,
                 session_state: uwb_uci_packets::SessionState::SessionStateActive,
                 reason_code: uwb_uci_packets::ReasonCode::StateChangeWithSessionManagementCommands
                     .into(),
@@ -748,7 +748,7 @@ mod tests {
         };
         let session_update_controller_multicast_list_ntf =
             uwb_uci_packets::SessionUpdateControllerMulticastListNtfBuilder {
-                session_id: 0x32,
+                session_token: 0x32,
                 remaining_multicast_list_size: 0x2,
                 controlee_status: vec![controlee_status.clone(), another_controlee_status.clone()],
             }
@@ -764,7 +764,7 @@ mod tests {
         assert_eq!(
             uci_notification_from_session_update_controller_multicast_list_ntf,
             UciNotification::Session(SessionNotification::UpdateControllerMulticastList {
-                session_id: 0x32,
+                session_token: 0x32,
                 remaining_multicast_list_size: 0x2,
                 status_list: vec![controlee_status, another_controlee_status],
             })
diff --git a/src/rust/uwb_core/src/uci/response.rs b/src/rust/uwb_core/src/uci/response.rs
index b0d5b2b..3c76f97 100644
--- a/src/rust/uwb_core/src/uci/response.rs
+++ b/src/rust/uwb_core/src/uci/response.rs
@@ -17,12 +17,12 @@ use std::convert::{TryFrom, TryInto};
 use crate::error::{Error, Result};
 use crate::params::uci_packets::{
     AppConfigTlv, CapTlv, CoreSetConfigResponse, DeviceConfigTlv, GetDeviceInfoResponse,
-    PowerStats, RawUciMessage, SessionState, SessionUpdateDtTagRangingRoundsResponse,
-    SetAppConfigResponse, StatusCode, UciControlPacket,
+    PowerStats, RawUciMessage, SessionHandle, SessionState,
+    SessionUpdateDtTagRangingRoundsResponse, SetAppConfigResponse, StatusCode, UciControlPacket,
 };
 use crate::uci::error::status_code_to_result;
 
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq)]
 pub(super) enum UciResponse {
     SetLoggerMode,
     SetNotification,
@@ -33,7 +33,7 @@ pub(super) enum UciResponse {
     CoreGetCapsInfo(Result<Vec<CapTlv>>),
     CoreSetConfig(CoreSetConfigResponse),
     CoreGetConfig(Result<Vec<DeviceConfigTlv>>),
-    SessionInit(Result<()>),
+    SessionInit(Result<Option<SessionHandle>>),
     SessionDeinit(Result<()>),
     SessionSetAppConfig(SetAppConfigResponse),
     SessionGetAppConfig(Result<Vec<AppConfigTlv>>),
@@ -154,8 +154,11 @@ impl TryFrom<uwb_uci_packets::SessionConfigResponse> for UciResponse {
         use uwb_uci_packets::SessionConfigResponseChild;
         match evt.specialize() {
             SessionConfigResponseChild::SessionInitRsp(evt) => {
-                Ok(UciResponse::SessionInit(status_code_to_result(evt.get_status())))
+                Ok(UciResponse::SessionInit(status_code_to_result(evt.get_status()).map(|_| None)))
             }
+            SessionConfigResponseChild::SessionInitRsp_V2(evt) => Ok(UciResponse::SessionInit(
+                status_code_to_result(evt.get_status()).map(|_| Some(evt.get_session_handle())),
+            )),
             SessionConfigResponseChild::SessionDeinitRsp(evt) => {
                 Ok(UciResponse::SessionDeinit(status_code_to_result(evt.get_status())))
             }
diff --git a/src/rust/uwb_core/src/uci/uci_hal.rs b/src/rust/uwb_core/src/uci/uci_hal.rs
index 1153872..151a808 100644
--- a/src/rust/uwb_core/src/uci/uci_hal.rs
+++ b/src/rust/uwb_core/src/uci/uci_hal.rs
@@ -32,7 +32,7 @@ pub type UciHalPacket = Vec<u8>;
 /// this trait and inject into the library.
 /// Note: Each method should be completed in 1000 ms.
 #[async_trait]
-pub trait UciHal: 'static + Send {
+pub trait UciHal: 'static + Send + Sync {
     /// Open the UCI HAL and power on the UWB Subsystem.
     ///
     /// All the other API should be called after the open() completes successfully. Once the method
diff --git a/src/rust/uwb_core/src/uci/uci_logger.rs b/src/rust/uwb_core/src/uci/uci_logger.rs
index 5023d7a..7045b43 100644
--- a/src/rust/uwb_core/src/uci/uci_logger.rs
+++ b/src/rust/uwb_core/src/uci/uci_logger.rs
@@ -74,10 +74,12 @@ fn filter_uci_command(cmd: UciControlPacket) -> UciControlPacket {
         UciControlPacketChild::UciCommand(control_cmd) => match control_cmd.specialize() {
             UciCommandChild::SessionConfigCommand(session_cmd) => match session_cmd.specialize() {
                 SessionConfigCommandChild::SessionSetAppConfigCmd(set_config_cmd) => {
-                    let session_id = set_config_cmd.get_session_id();
+                    let session_token = set_config_cmd.get_session_token();
                     let tlvs = set_config_cmd.get_tlvs().to_owned();
                     let filtered_tlvs = tlvs.into_iter().map(filter_tlv).collect();
-                    SessionSetAppConfigCmdBuilder { session_id, tlvs: filtered_tlvs }.build().into()
+                    SessionSetAppConfigCmdBuilder { session_token, tlvs: filtered_tlvs }
+                        .build()
+                        .into()
                 }
                 _ => session_cmd.into(),
             },
@@ -214,7 +216,7 @@ mod tests {
     #[test]
     fn test_log_command_filter() -> Result<()> {
         let set_config_cmd = UciCommand::SessionSetAppConfig {
-            session_id: 0x1,
+            session_token: 0x1,
             config_tlvs: vec![
                 // Filtered to 0-filled of same length
                 AppConfigTlv { cfg_id: AppConfigTlvType::VendorId, v: vec![0, 1, 2] }.into(),
diff --git a/src/rust/uwb_core/src/uci/uci_manager.rs b/src/rust/uwb_core/src/uci/uci_manager.rs
index 9d05224..9d24c24 100644
--- a/src/rust/uwb_core/src/uci/uci_manager.rs
+++ b/src/rust/uwb_core/src/uci/uci_manager.rs
@@ -13,11 +13,12 @@
 // limitations under the License.
 
 use std::convert::TryInto;
+use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
-use log::{debug, error, warn};
-use tokio::sync::{mpsc, oneshot};
+use log::{debug, error, info, warn};
+use tokio::sync::{mpsc, oneshot, Mutex};
 
 use crate::uci::command::UciCommand;
 //use crate::uci::error::{Error, Result};
@@ -26,13 +27,13 @@ use crate::params::uci_packets::{
     AppConfigTlv, AppConfigTlvType, CapTlv, Controlees, CoreSetConfigResponse, CountryCode,
     CreditAvailability, DeviceConfigId, DeviceConfigTlv, DeviceState, FiraComponent,
     GetDeviceInfoResponse, GroupId, MessageType, PowerStats, RawUciMessage, ResetConfig, SessionId,
-    SessionState, SessionType, SessionUpdateDtTagRangingRoundsResponse, SetAppConfigResponse,
-    UciDataPacket, UciDataPacketHal, UpdateMulticastListAction,
+    SessionState, SessionToken, SessionType, SessionUpdateDtTagRangingRoundsResponse,
+    SetAppConfigResponse, UciDataPacket, UciDataPacketHal, UpdateMulticastListAction,
 };
 use crate::params::utils::bytes_to_u64;
 use crate::uci::message::UciMessage;
 use crate::uci::notification::{
-    CoreNotification, DataRcvNotification, SessionNotification, UciNotification,
+    CoreNotification, DataRcvNotification, SessionNotification, SessionRangeData, UciNotification,
 };
 use crate::uci::response::UciResponse;
 use crate::uci::timeout_uci_hal::TimeoutUciHal;
@@ -150,6 +151,11 @@ pub trait UciManager: 'static + Send + Sync + Clone {
 #[derive(Clone)]
 pub struct UciManagerImpl {
     cmd_sender: mpsc::UnboundedSender<(UciManagerCmd, oneshot::Sender<Result<UciResponse>>)>,
+
+    // FIRA version 2 introduces a UWBS generated session handle to use as identifier for all
+    // session related commands. This map stores the app provided session id to UWBS generated
+    // session handle mapping if provided, else reuses session id.
+    session_id_to_token_map: Arc<Mutex<HashMap<SessionId, SessionToken>>>,
 }
 
 impl UciManagerImpl {
@@ -160,10 +166,18 @@ impl UciManagerImpl {
         logger_mode: UciLoggerMode,
     ) -> Self {
         let (cmd_sender, cmd_receiver) = mpsc::unbounded_channel();
-        let mut actor = UciManagerActor::new(hal, logger, logger_mode, cmd_receiver);
+        let session_id_to_token_map: Arc<Mutex<HashMap<SessionId, SessionToken>>> =
+            Arc::new(Mutex::new(HashMap::new()));
+        let mut actor = UciManagerActor::new(
+            hal,
+            logger,
+            logger_mode,
+            cmd_receiver,
+            session_id_to_token_map.clone(),
+        );
         tokio::spawn(async move { actor.run().await });
 
-        Self { cmd_sender }
+        Self { cmd_sender, session_id_to_token_map }
     }
 
     // Send the |cmd| to the UciManagerActor.
@@ -177,6 +191,15 @@ impl UciManagerImpl {
             }
         }
     }
+
+    async fn get_session_token(&self, session_id: &SessionId) -> Result<SessionToken> {
+        self.session_id_to_token_map
+            .lock()
+            .await
+            .get(session_id)
+            .ok_or(Error::BadParameters)
+            .copied()
+    }
 }
 
 #[async_trait]
@@ -293,14 +316,15 @@ impl UciManager for UciManagerImpl {
     async fn session_init(&self, session_id: SessionId, session_type: SessionType) -> Result<()> {
         let cmd = UciCommand::SessionInit { session_id, session_type };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
-            Ok(UciResponse::SessionInit(resp)) => resp,
+            Ok(UciResponse::SessionInit(resp)) => resp.map(|_| {}),
             Ok(_) => Err(Error::Unknown),
             Err(e) => Err(e),
         }
     }
 
     async fn session_deinit(&self, session_id: SessionId) -> Result<()> {
-        let cmd = UciCommand::SessionDeinit { session_id };
+        let cmd =
+            UciCommand::SessionDeinit { session_token: self.get_session_token(&session_id).await? };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionDeinit(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -313,7 +337,10 @@ impl UciManager for UciManagerImpl {
         session_id: SessionId,
         config_tlvs: Vec<AppConfigTlv>,
     ) -> Result<SetAppConfigResponse> {
-        let cmd = UciCommand::SessionSetAppConfig { session_id, config_tlvs };
+        let cmd = UciCommand::SessionSetAppConfig {
+            session_token: self.get_session_token(&session_id).await?,
+            config_tlvs,
+        };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionSetAppConfig(resp)) => Ok(resp),
             Ok(_) => Err(Error::Unknown),
@@ -326,7 +353,10 @@ impl UciManager for UciManagerImpl {
         session_id: SessionId,
         app_cfg: Vec<AppConfigTlvType>,
     ) -> Result<Vec<AppConfigTlv>> {
-        let cmd = UciCommand::SessionGetAppConfig { session_id, app_cfg };
+        let cmd = UciCommand::SessionGetAppConfig {
+            session_token: self.get_session_token(&session_id).await?,
+            app_cfg,
+        };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionGetAppConfig(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -344,7 +374,9 @@ impl UciManager for UciManagerImpl {
     }
 
     async fn session_get_state(&self, session_id: SessionId) -> Result<SessionState> {
-        let cmd = UciCommand::SessionGetState { session_id };
+        let cmd = UciCommand::SessionGetState {
+            session_token: self.get_session_token(&session_id).await?,
+        };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionGetState(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -367,8 +399,11 @@ impl UciManager for UciManagerImpl {
             warn!("Number of controlees should be between 1 to 8");
             return Err(Error::BadParameters);
         }
-        let cmd =
-            UciCommand::SessionUpdateControllerMulticastList { session_id, action, controlees };
+        let cmd = UciCommand::SessionUpdateControllerMulticastList {
+            session_token: self.get_session_token(&session_id).await?,
+            action,
+            controlees,
+        };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionUpdateControllerMulticastList(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -381,7 +416,10 @@ impl UciManager for UciManagerImpl {
         session_id: u32,
         ranging_round_indexes: Vec<u8>,
     ) -> Result<SessionUpdateDtTagRangingRoundsResponse> {
-        let cmd = UciCommand::SessionUpdateDtTagRangingRounds { session_id, ranging_round_indexes };
+        let cmd = UciCommand::SessionUpdateDtTagRangingRounds {
+            session_token: self.get_session_token(&session_id).await?,
+            ranging_round_indexes,
+        };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionUpdateDtTagRangingRounds(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -390,7 +428,9 @@ impl UciManager for UciManagerImpl {
     }
 
     async fn session_query_max_data_size(&self, session_id: SessionId) -> Result<u16> {
-        let cmd = UciCommand::SessionQueryMaxDataSize { session_id };
+        let cmd = UciCommand::SessionQueryMaxDataSize {
+            session_token: self.get_session_token(&session_id).await?,
+        };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionQueryMaxDataSize(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -399,7 +439,8 @@ impl UciManager for UciManagerImpl {
     }
 
     async fn range_start(&self, session_id: SessionId) -> Result<()> {
-        let cmd = UciCommand::SessionStart { session_id };
+        let cmd =
+            UciCommand::SessionStart { session_token: self.get_session_token(&session_id).await? };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionStart(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -408,7 +449,8 @@ impl UciManager for UciManagerImpl {
     }
 
     async fn range_stop(&self, session_id: SessionId) -> Result<()> {
-        let cmd = UciCommand::SessionStop { session_id };
+        let cmd =
+            UciCommand::SessionStop { session_token: self.get_session_token(&session_id).await? };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionStop(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -417,7 +459,9 @@ impl UciManager for UciManagerImpl {
     }
 
     async fn range_get_ranging_count(&self, session_id: SessionId) -> Result<usize> {
-        let cmd = UciCommand::SessionGetRangingCount { session_id };
+        let cmd = UciCommand::SessionGetRangingCount {
+            session_token: self.get_session_token(&session_id).await?,
+        };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {
             Ok(UciResponse::SessionGetRangingCount(resp)) => resp,
             Ok(_) => Err(Error::Unknown),
@@ -474,7 +518,7 @@ impl UciManager for UciManagerImpl {
         let dest_mac_address =
             bytes_to_u64(dest_mac_address_bytes).ok_or(Error::BadParameters).unwrap();
         let data_snd_packet = uwb_uci_packets::UciDataSndBuilder {
-            session_id,
+            session_token: self.get_session_token(&session_id).await?,
             dest_mac_address,
             dest_fira_component,
             uci_sequence_number,
@@ -500,7 +544,7 @@ struct UciManagerActor<T: UciHal, U: UciLogger> {
 
     // Set to true when |hal| is opened successfully.
     is_hal_opened: bool,
-    // Receive response, notification and data packets from |hal|. Only used when |hal| is opened
+    // Receive response, notification and data packets from |mut hal|. Only used when |hal| is opened
     // successfully.
     packet_receiver: mpsc::UnboundedReceiver<UciHalPacket>,
     // Defrag the UCI packets.
@@ -512,11 +556,11 @@ struct UciManagerActor<T: UciHal, U: UciLogger> {
 
     // Store per-session CreditAvailability. This should be initialized when a UWB session becomes
     // ACTIVE, and updated every time a Data packet fragment is sent or a DataCreditNtf is received.
-    data_credit_map: HashMap<SessionId, CreditAvailability>,
+    data_credit_map: HashMap<SessionToken, CreditAvailability>,
 
     // Store the Uci Data packet fragments to be sent to the UWBS, keyed by the SessionId. This
     // helps to retrieve the next packet fragment to be sent, when the UWBS is ready to accept it.
-    data_packet_fragments_map: HashMap<SessionId, VecDeque<UciDataPacketHal>>,
+    data_packet_fragments_map: HashMap<SessionToken, VecDeque<UciDataPacketHal>>,
 
     // The timeout of waiting for the notification of device ready notification.
     wait_device_status_timeout: PinSleep,
@@ -541,6 +585,14 @@ struct UciManagerActor<T: UciHal, U: UciLogger> {
     session_notf_sender: mpsc::UnboundedSender<SessionNotification>,
     vendor_notf_sender: mpsc::UnboundedSender<RawUciMessage>,
     data_rcv_notf_sender: mpsc::UnboundedSender<DataRcvNotification>,
+
+    // Used to store the last init session id to help map the session handle sent
+    // in session int response can be correctly mapped.
+    last_init_session_id: Option<SessionId>,
+    // FIRA version 2 introduces a UWBS generated session handle to use as identifier for all
+    // session related commands. This map stores the app provided session id to UWBS generated
+    // session handle mapping if provided, else reuses session id.
+    session_id_to_token_map: Arc<Mutex<HashMap<SessionId, SessionToken>>>,
 }
 
 impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
@@ -552,6 +604,7 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
             UciManagerCmd,
             oneshot::Sender<Result<UciResponse>>,
         )>,
+        session_id_to_token_map: Arc<Mutex<HashMap<SessionId, SessionToken>>>,
     ) -> Self {
         Self {
             hal: TimeoutUciHal::new(hal),
@@ -572,6 +625,8 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
             session_notf_sender: mpsc::unbounded_channel().0,
             vendor_notf_sender: mpsc::unbounded_channel().0,
             data_rcv_notf_sender: mpsc::unbounded_channel().0,
+            last_init_session_id: None,
+            session_id_to_token_map,
         }
     }
 
@@ -619,6 +674,60 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
         }
     }
 
+    async fn insert_session_token(&self, session_id: SessionId, session_token: SessionToken) {
+        self.session_id_to_token_map.lock().await.insert(session_id, session_token);
+    }
+
+    async fn remove_session_token(&self, session_token: &SessionToken) {
+        self.session_id_to_token_map.lock().await.retain(|_, val| *val != *session_token);
+    }
+
+    async fn get_session_id(&self, session_token: &SessionToken) -> Result<SessionId> {
+        self.session_id_to_token_map
+            .lock()
+            .await
+            .iter()
+            .find_map(|(key, &val)| if val == *session_token { Some(key) } else { None })
+            .ok_or(Error::BadParameters)
+            .copied()
+    }
+
+    fn save_session_id_if_init_cmd(&mut self, cmd: &UciCommand) {
+        // Store the last init session id to help map the session handle sent
+        // in session init response.
+        if let UciCommand::SessionInit { session_id, .. } = cmd {
+            self.last_init_session_id = Some(*session_id);
+        }
+    }
+
+    async fn store_session_token_if_init_resp(&mut self, resp: &UciResponse) -> Result<()> {
+        // Store the session_id to session_token mapping for this new session.
+        if let UciResponse::SessionInit(session_init_resp) = resp {
+            let session_id = match self.last_init_session_id.take() {
+                Some(session_id) => session_id,
+                None => {
+                    return Err(Error::Unknown);
+                }
+            };
+            if let Ok(opt_session_handle) = session_init_resp {
+                let session_handle = match opt_session_handle {
+                    // Session Handle provided by UWBS, use as token for further commands.
+                    Some(session_handle) => {
+                        info!(
+                            "session handle: {:?} provided for session id: {:?}",
+                            session_handle, session_id
+                        );
+                        *session_handle
+                    }
+                    // Session Handle not provided by UWBS, reuse session id as token for further commands.
+                    None => session_id,
+                };
+                self.insert_session_token(session_id, session_handle).await;
+            }
+        }
+        Ok(())
+    }
+
     async fn handle_cmd(
         &mut self,
         cmd: UciManagerCmd,
@@ -697,6 +806,8 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
             UciManagerCmd::SendUciCommand { cmd } => {
                 debug_assert!(self.uci_cmd_retryer.is_none());
 
+                self.save_session_id_if_init_cmd(&cmd);
+
                 // Remember that this command is a raw UCI command, we'll use this later
                 // to send a raw UCI response.
                 if let UciCommand::RawUciCmd { mt: _, gid, oid, payload: _ } = cmd.clone() {
@@ -756,11 +867,11 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
 
     async fn retry_uci_data_snd(&mut self) {
         if let Some(mut uci_data_snd_retryer) = self.uci_data_snd_retryer.take() {
-            let data_packet_session_id = uci_data_snd_retryer.data_packet_session_id;
+            let data_packet_session_token = uci_data_snd_retryer.data_packet_session_token;
             if !uci_data_snd_retryer.could_retry() {
                 error!(
                     "Out of retries for Uci DataSnd packet, last DataSnd packet session_id:{}",
-                    data_packet_session_id
+                    data_packet_session_token
                 );
                 return;
             }
@@ -772,7 +883,7 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
                 Err(e) => {
                     error!(
                         "DataSnd packet fragment session_id:{} retry failed with error:{}",
-                        data_packet_session_id, e
+                        data_packet_session_token, e
                     );
                 }
             }
@@ -793,14 +904,14 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
 
     async fn handle_data_snd_packet(&mut self, data_snd_packet: UciDataSnd) -> Result<UciResponse> {
         // Verify that there's an entry for the Session in the CreditAvailability map.
-        let data_packet_session_id = data_snd_packet.get_session_id();
+        let data_packet_session_token = data_snd_packet.get_session_token();
         let data_packet_sequence_number = data_snd_packet.get_uci_sequence_number();
 
-        if !self.data_credit_map.contains_key(&data_packet_session_id) {
+        if !self.data_credit_map.contains_key(&data_packet_session_token) {
             error!(
-                "DataSnd packet session_id:{}, sequence_number:{} cannot be sent as unknown \
+                "DataSnd packet session_token:{}, sequence_number:{} cannot be sent as unknown \
                 credit availability for the session",
-                data_packet_session_id, data_packet_sequence_number
+                data_packet_session_token, data_packet_sequence_number
             );
             return Err(Error::PacketTxError);
         }
@@ -809,13 +920,13 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
         let mut packet_fragments: Vec<UciDataPacketHal> = data_snd_packet.into();
         if packet_fragments.is_empty() {
             error!(
-                "DataSnd packet session_id:{}, sequence number:{} could not be split into fragments",
-                data_packet_session_id, data_packet_sequence_number
+                "DataSnd packet session_token:{}, sequence number:{} could not be split into fragments",
+                data_packet_session_token, data_packet_sequence_number
             );
             return Err(Error::PacketTxError);
         }
 
-        match self.data_packet_fragments_map.get_mut(&data_packet_session_id) {
+        match self.data_packet_fragments_map.get_mut(&data_packet_session_token) {
             Some(q) => {
                 for p in packet_fragments.drain(..) {
                     q.push_back(p);
@@ -823,29 +934,29 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
             }
             None => {
                 error!(
-                    "DataSnd packet fragments map not found for session_id:{}",
-                    data_packet_session_id
+                    "DataSnd packet fragments map not found for session_token:{}",
+                    data_packet_session_token
                 );
                 return Err(Error::PacketTxError);
             }
         }
 
-        self.send_data_packet_fragment(data_packet_session_id).await
+        self.send_data_packet_fragment(data_packet_session_token).await
     }
 
     async fn send_data_packet_fragment(
         &mut self,
-        data_packet_session_id: SessionId,
+        data_packet_session_token: SessionToken,
     ) -> Result<UciResponse> {
         // Check if a credit is available before sending this data packet fragment. If not, return
         // for now, and send this packet later when the credit becomes available (indicated by
         // receiving a DataCreditNtf).
-        let credit = self.data_credit_map.get(&data_packet_session_id);
+        let credit = self.data_credit_map.get(&data_packet_session_token);
         if credit.is_none() {
             error!(
-                "DataSnd packet fragment cannot be sent for session_id:{} as unknown \
+                "DataSnd packet fragment cannot be sent for session_token:{} as unknown \
                 credit availability for the session",
-                data_packet_session_id
+                data_packet_session_token
             );
             return Err(Error::PacketTxError);
         }
@@ -855,7 +966,7 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
 
         // We have credit available, let's send the packet to UWBS.
         let hal_data_packet_fragment =
-            match self.data_packet_fragments_map.get_mut(&data_packet_session_id) {
+            match self.data_packet_fragments_map.get_mut(&data_packet_session_token) {
                 Some(q) => {
                     match q.pop_front() {
                         Some(p) => p,
@@ -873,21 +984,22 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
         // Create and save a retryer for sending this data packet fragment.
         self.uci_data_snd_retryer = Some(UciDataSndRetryer {
             data_packet: hal_data_packet_fragment.clone(),
-            data_packet_session_id,
+            data_packet_session_token,
             retry_count: MAX_RETRY_COUNT,
         });
 
         let result = self.hal.send_packet(hal_data_packet_fragment.to_vec()).await;
         if result.is_err() {
             error!(
-                "Result {:?} of sending data packet fragment SessionId: {} to HAL",
-                result, data_packet_session_id
+                "Result {:?} of sending data packet fragment SessionToken: {} to HAL",
+                result, data_packet_session_token
             );
             return Err(Error::PacketTxError);
         }
 
         // Update the map after the successful write.
-        self.data_credit_map.insert(data_packet_session_id, CreditAvailability::CreditNotAvailable);
+        self.data_credit_map
+            .insert(data_packet_session_token, CreditAvailability::CreditNotAvailable);
         Ok(UciResponse::SendUciData(Ok(())))
     }
 
@@ -954,6 +1066,10 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
             self.retry_uci_cmd().await;
             return;
         }
+        if let Err(_e) = self.store_session_token_if_init_resp(&resp).await {
+            error!("Session init response received without a sesson id stored! Something has gone badly wrong: {:?}", resp);
+            return;
+        }
 
         if let Some(uci_cmd_retryer) = self.uci_cmd_retryer.take() {
             uci_cmd_retryer.send_result(Ok(resp));
@@ -986,51 +1102,63 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
                 }
                 let _ = self.core_notf_sender.send(core_notf);
             }
-            UciNotification::Session(session_notf) => {
-                if let SessionNotification::Status { session_id, session_state, reason_code: _ } =
-                    session_notf
-                {
-                    self.handle_session_state_notification(session_id, session_state).await;
-                }
-                if let SessionNotification::DataCredit { session_id, credit_availability } =
-                    session_notf
-                {
-                    if !self.data_credit_map.contains_key(&session_id) {
-                        // Currently just log, as this is unexpected (the entry should exist once
-                        // the ranging session is Active and be removed once it is Idle).
-                        debug!(
-                            "Received a DataCreditNtf for non-existent session_id: {}",
-                            session_id
-                        );
+            UciNotification::Session(orig_session_notf) => {
+                let mod_session_notf = {
+                    match self
+                        .replace_session_token_with_session_id(orig_session_notf.clone())
+                        .await
+                    {
+                        Ok(session_notf) => session_notf,
+                        Err(e) => {
+                            error!("Failed to find corresponding session id, discarding session notification {:?}: {:?}", orig_session_notf, e);
+                            return;
+                        }
                     }
-                    self.data_credit_map.insert(session_id, credit_availability);
-                    if credit_availability == CreditAvailability::CreditAvailable {
-                        if let Err(e) = self.send_data_packet_fragment(session_id).await {
-                            error!(
-                                "Sending data packet fragment failed with Err:{}, after a\
-                                   DataCreditNtf is received, for sessionId:{}",
-                                e, session_id
+                };
+                match orig_session_notf {
+                    SessionNotification::Status {
+                        session_token,
+                        session_state,
+                        reason_code: _,
+                    } => self.handle_session_state_notification(session_token, session_state).await,
+                    SessionNotification::DataCredit { session_token, credit_availability } => {
+                        if !self.data_credit_map.contains_key(&session_token) {
+                            // Currently just log, as this is unexpected (the entry should exist once
+                            // the ranging session is Active and be removed once it is Idle).
+                            debug!(
+                                "Received a DataCreditNtf for non-existent session_token: {}",
+                                session_token
                             );
                         }
-                    } else {
-                        // Log as this should usually not happen (it's not an error).
-                        debug!(
-                            "Received a DataCreditNtf with no credit available for session_id:{}",
-                            session_id
+                        self.data_credit_map.insert(session_token, credit_availability);
+                        if credit_availability == CreditAvailability::CreditAvailable {
+                            if let Err(e) = self.send_data_packet_fragment(session_token).await {
+                                error!(
+                                    "Sending data packet fragment failed with Err:{}, after a\
+                                   DataCreditNtf is received, for session_token:{}",
+                                    e, session_token
+                                );
+                            }
+                        } else {
+                            // Log as this should usually not happen (it's not an error).
+                            debug!(
+                            "Received a DataCreditNtf with no credit available for session_token:{}",
+                            session_token
                         );
+                        }
+                        return; // We consume these here and don't need to send to upper layer.
                     }
-                    return; // We consume these here and don't need to send to upper layer.
-                }
-                if let SessionNotification::DataTransferStatus {
-                    session_id: _,
-                    uci_sequence_number: _,
-                    status: _,
-                } = session_notf
-                {
-                    // Reset the UciDataSnd Retryer since we received a DataTransferStatusNtf.
-                    let _ = self.uci_data_snd_retryer.take();
+                    SessionNotification::DataTransferStatus {
+                        session_token: _,
+                        uci_sequence_number: _,
+                        status: _,
+                    } => {
+                        // Reset the UciDataSnd Retryer since we received a DataTransferStatusNtf.
+                        let _ = self.uci_data_snd_retryer.take();
+                    }
+                    _ => {}
                 }
-                let _ = self.session_notf_sender.send(session_notf);
+                let _ = self.session_notf_sender.send(mod_session_notf);
             }
             UciNotification::Vendor(vendor_notf) => {
                 let _ = self.vendor_notf_sender.send(vendor_notf);
@@ -1038,26 +1166,82 @@ impl<T: UciHal, U: UciLogger> UciManagerActor<T, U> {
         }
     }
 
+    // Modify session_token field in all session related notifications with session id.
+    // TODO: Sharing of structs across UCI (PDL) & JNI layer like this makes this ugly. Ideally
+    // the struct sent to JNI layer should only contain |session_id| and at uci layer
+    // it could be |session_id| or |session_handle|.
+    async fn replace_session_token_with_session_id(
+        &self,
+        session_notification: SessionNotification,
+    ) -> Result<SessionNotification> {
+        match session_notification {
+            SessionNotification::Status { session_token, session_state, reason_code } => {
+                Ok(SessionNotification::Status {
+                    session_token: self.get_session_id(&session_token).await?,
+                    session_state,
+                    reason_code,
+                })
+            }
+            SessionNotification::UpdateControllerMulticastList {
+                session_token,
+                remaining_multicast_list_size,
+                status_list,
+            } => Ok(SessionNotification::UpdateControllerMulticastList {
+                session_token: self.get_session_id(&session_token).await?,
+                remaining_multicast_list_size,
+                status_list,
+            }),
+            SessionNotification::SessionInfo(session_range_data) => {
+                Ok(SessionNotification::SessionInfo(SessionRangeData {
+                    sequence_number: session_range_data.sequence_number,
+                    session_token: self.get_session_id(&session_range_data.session_token).await?,
+                    current_ranging_interval_ms: session_range_data.current_ranging_interval_ms,
+                    ranging_measurement_type: session_range_data.ranging_measurement_type,
+                    ranging_measurements: session_range_data.ranging_measurements,
+                    rcr_indicator: session_range_data.rcr_indicator,
+                    raw_ranging_data: session_range_data.raw_ranging_data,
+                }))
+            }
+            SessionNotification::DataTransferStatus {
+                session_token,
+                uci_sequence_number,
+                status,
+            } => Ok(SessionNotification::DataTransferStatus {
+                session_token: self.get_session_id(&session_token).await?,
+                uci_sequence_number,
+                status,
+            }),
+            SessionNotification::DataCredit { session_token, credit_availability } => {
+                Ok(SessionNotification::DataCredit {
+                    session_token: self.get_session_id(&session_token).await?,
+                    credit_availability,
+                })
+            }
+        }
+    }
+
     async fn handle_session_state_notification(
         &mut self,
-        session_id: SessionId,
+        session_token: SessionToken,
         session_state: SessionState,
     ) {
         match session_state {
             SessionState::SessionStateInit => {
-                if let Err(e) = self.hal.notify_session_initialized(session_id).await {
+                if let Err(e) = self.hal.notify_session_initialized(session_token).await {
                     warn!("notify_session_initialized() failed: {:?}", e);
                 }
             }
             SessionState::SessionStateActive => {
-                self.data_credit_map.insert(session_id, CreditAvailability::CreditAvailable);
-                self.data_packet_fragments_map.insert(session_id, VecDeque::new());
+                self.data_credit_map.insert(session_token, CreditAvailability::CreditAvailable);
+                self.data_packet_fragments_map.insert(session_token, VecDeque::new());
             }
             SessionState::SessionStateIdle => {
-                self.data_credit_map.remove(&session_id);
-                self.data_packet_fragments_map.remove(&session_id);
+                self.data_credit_map.remove(&session_token);
+                self.data_packet_fragments_map.remove(&session_token);
+            }
+            SessionState::SessionStateDeinit => {
+                self.remove_session_token(&session_token).await;
             }
-            _ => {}
         }
     }
 
@@ -1126,7 +1310,7 @@ struct UciDataSndRetryer {
     // sessions, as there can be outstanding DataSnd packet fragments across them. We could do an
     // alternative implementation of sending all of them.
     data_packet: UciDataPacketHal,
-    data_packet_session_id: u32,
+    data_packet_session_token: SessionToken,
     retry_count: usize,
 }
 
@@ -1169,11 +1353,12 @@ enum UciManagerCmd {
     },
 }
 
-#[cfg(test)]
+#[cfg(any(test))]
 mod tests {
     use super::*;
 
     use bytes::Bytes;
+    use tokio::macros::support::Future;
     use uwb_uci_packets::{SessionGetCountCmdBuilder, SessionGetCountRspBuilder};
 
     use crate::params::uci_packets::{
@@ -1210,24 +1395,14 @@ mod tests {
         bytes
     }
 
-    async fn setup_uci_manager_with_open_hal<F>(
-        setup_hal_fn: F,
-        uci_logger_mode: UciLoggerMode,
-        log_sender: mpsc::UnboundedSender<UciLogEvent>,
-    ) -> (UciManagerImpl, MockUciHal)
-    where
-        F: FnOnce(&mut MockUciHal),
-    {
-        init_test_logging();
-
-        // Open the hal.
-        let mut hal = MockUciHal::new();
+    fn setup_hal_for_open(hal: &mut MockUciHal) {
+        // Setup Open the hal.
         let notf = into_uci_hal_packets(uwb_uci_packets::DeviceStatusNtfBuilder {
             device_state: uwb_uci_packets::DeviceState::DeviceStateReady,
         });
         hal.expected_open(Some(notf), Ok(()));
 
-        // Get the device info.
+        // Setup Get the device info.
         let cmd = UciCommand::CoreGetDeviceInfo;
         let resp = into_uci_hal_packets(uwb_uci_packets::GetDeviceInfoRspBuilder {
             status: uwb_uci_packets::StatusCode::UciStatusOk,
@@ -1238,14 +1413,31 @@ mod tests {
             vendor_spec_info: vec![0x1, 0x2],
         });
         hal.expected_send_command(cmd, resp, Ok(()));
+    }
+
+    async fn setup_uci_manager_with_open_hal<F, Fut>(
+        setup_hal_fn: F,
+        uci_logger_mode: UciLoggerMode,
+        log_sender: mpsc::UnboundedSender<UciLogEvent>,
+    ) -> (UciManagerImpl, MockUciHal)
+    where
+        F: FnOnce(MockUciHal) -> Fut,
+        Fut: Future<Output = ()>,
+    {
+        init_test_logging();
 
-        setup_hal_fn(&mut hal);
+        let mut hal = MockUciHal::new();
+        // Open the hal.
+        setup_hal_for_open(&mut hal);
 
         // Verify open_hal() is working.
         let uci_manager =
             UciManagerImpl::new(hal.clone(), MockUciLogger::new(log_sender), uci_logger_mode);
         let result = uci_manager.open_hal().await;
         assert!(result.is_ok());
+        assert!(hal.wait_expected_calls_done().await);
+
+        setup_hal_fn(hal.clone()).await;
 
         (uci_manager, hal)
     }
@@ -1267,7 +1459,7 @@ mod tests {
     #[tokio::test]
     async fn test_close_hal_explicitly() {
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            |hal| {
+            |mut hal| async move {
                 hal.expected_close(Ok(()));
             },
             UciLoggerMode::Disabled,
@@ -1283,7 +1475,7 @@ mod tests {
     #[tokio::test]
     async fn test_close_hal_when_exit() {
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            |hal| {
+            |mut hal| async move {
                 // UciManager should close the hal if the hal is still opened when exit.
                 hal.expected_close(Ok(()));
             },
@@ -1312,12 +1504,11 @@ mod tests {
     #[tokio::test]
     async fn test_device_reset_ok() {
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::DeviceReset { reset_config: ResetConfig::UwbsReset };
                 let resp = into_uci_hal_packets(uwb_uci_packets::DeviceResetRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
                 });
-
                 hal.expected_send_command(cmd, resp, Ok(()));
             },
             UciLoggerMode::Disabled,
@@ -1341,7 +1532,7 @@ mod tests {
         let vendor_spec_info_clone = vendor_spec_info.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::CoreGetDeviceInfo;
                 let resp = into_uci_hal_packets(uwb_uci_packets::GetDeviceInfoRspBuilder {
                     status,
@@ -1377,7 +1568,7 @@ mod tests {
         let tlv_clone = tlv.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::CoreGetCapsInfo;
                 let resp = into_uci_hal_packets(uwb_uci_packets::GetCapsInfoRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
@@ -1408,7 +1599,7 @@ mod tests {
         let config_status_clone = config_status.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::CoreSetConfig { config_tlvs: vec![tlv_clone] };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SetConfigRspBuilder {
                     status,
@@ -1435,7 +1626,7 @@ mod tests {
         let tlv_clone = tlv.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::CoreGetConfig { cfg_id: vec![cfg_id] };
                 let resp = into_uci_hal_packets(uwb_uci_packets::GetConfigRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
@@ -1455,47 +1646,111 @@ mod tests {
         assert!(mock_hal.wait_expected_calls_done().await);
     }
 
-    #[tokio::test]
-    async fn test_session_init_ok() {
-        let session_id = 0x123;
+    fn setup_hal_for_session_initialize(
+        hal: &mut MockUciHal,
+        session_type: SessionType,
+        session_id: u32,
+        session_token: u32,
+    ) {
+        // Setup for hal open.
+        setup_hal_for_open(hal);
+
+        // Setup session init.
+        let cmd = UciCommand::SessionInit { session_id, session_type };
+        let mut resp = if session_id == session_token {
+            into_uci_hal_packets(uwb_uci_packets::SessionInitRspBuilder {
+                status: uwb_uci_packets::StatusCode::UciStatusOk,
+            })
+        } else {
+            // This is testing FIRA v2 flow where a session handle is provided by UWBS.
+            into_uci_hal_packets(uwb_uci_packets::SessionInitRsp_V2Builder {
+                status: uwb_uci_packets::StatusCode::UciStatusOk,
+                session_handle: session_token,
+            })
+        };
+        let mut notf = into_uci_hal_packets(uwb_uci_packets::SessionStatusNtfBuilder {
+            session_token,
+            session_state: uwb_uci_packets::SessionState::SessionStateInit,
+            reason_code: uwb_uci_packets::ReasonCode::StateChangeWithSessionManagementCommands
+                .into(),
+        });
+        resp.append(&mut notf);
+        hal.expected_send_command(cmd, resp, Ok(()));
+        hal.expected_notify_session_initialized(session_token, Ok(()));
+    }
+
+    async fn setup_uci_manager_with_session_initialized<F, Fut>(
+        setup_hal_fn: F,
+        uci_logger_mode: UciLoggerMode,
+        log_sender: mpsc::UnboundedSender<UciLogEvent>,
+        session_id: u32,
+        session_token: u32,
+    ) -> (UciManagerImpl, MockUciHal)
+    where
+        F: FnOnce(MockUciHal) -> Fut,
+        Fut: Future<Output = ()>,
+    {
         let session_type = SessionType::FiraRangingSession;
 
-        let (uci_manager, mut mock_hal) =
-            setup_uci_manager_with_open_hal(
-                move |hal| {
-                    let cmd = UciCommand::SessionInit { session_id, session_type };
-                    let mut resp = into_uci_hal_packets(uwb_uci_packets::SessionInitRspBuilder {
-                        status: uwb_uci_packets::StatusCode::UciStatusOk,
-                    });
-                    let mut notf = into_uci_hal_packets(uwb_uci_packets::SessionStatusNtfBuilder {
-                        session_id,
-                        session_state: uwb_uci_packets::SessionState::SessionStateInit,
-                        reason_code:
-                            uwb_uci_packets::ReasonCode::StateChangeWithSessionManagementCommands
-                                .into(),
-                    });
-                    resp.append(&mut notf);
+        init_test_logging();
 
-                    hal.expected_send_command(cmd, resp, Ok(()));
-                    hal.expected_notify_session_initialized(session_id, Ok(()));
-                },
-                UciLoggerMode::Disabled,
-                mpsc::unbounded_channel::<UciLogEvent>().0,
-            )
-            .await;
+        let mut hal = MockUciHal::new();
+        setup_hal_for_session_initialize(&mut hal, session_type, session_id, session_token);
 
+        // Verify open_hal() is working.
+        let uci_manager =
+            UciManagerImpl::new(hal.clone(), MockUciLogger::new(log_sender), uci_logger_mode);
+        let result = uci_manager.open_hal().await;
+        assert!(result.is_ok());
+
+        // Verify session is initialized.
         let result = uci_manager.session_init(session_id, session_type).await;
         assert!(result.is_ok());
+        assert!(hal.wait_expected_calls_done().await);
+
+        setup_hal_fn(hal.clone()).await;
+
+        (uci_manager, hal)
+    }
+
+    #[tokio::test]
+    async fn test_session_init_ok() {
+        let session_id = 0x123;
+        let session_token = 0x123;
+        let (_, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |_hal| async move {},
+            UciLoggerMode::Disabled,
+            mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
+        )
+        .await;
+        assert!(mock_hal.wait_expected_calls_done().await);
+    }
+
+    #[tokio::test]
+    async fn test_session_init_v2_ok() {
+        let session_id = 0x123;
+        let session_token = 0x321; // different session handle
+        let (_, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |_hal| async move {},
+            UciLoggerMode::Disabled,
+            mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
+        )
+        .await;
         assert!(mock_hal.wait_expected_calls_done().await);
     }
 
     #[tokio::test]
     async fn test_session_deinit_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                let cmd = UciCommand::SessionDeinit { session_id };
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd = UciCommand::SessionDeinit { session_token };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionDeinitRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
                 });
@@ -1504,6 +1759,34 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
+        )
+        .await;
+
+        let result = uci_manager.session_deinit(session_id).await;
+        assert!(result.is_ok());
+        assert!(mock_hal.wait_expected_calls_done().await);
+    }
+
+    #[tokio::test]
+    async fn test_session_deinit_v2_ok() {
+        let session_id = 0x123;
+        let session_token = 0x321; // different session handle
+
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd = UciCommand::SessionDeinit { session_token };
+                let resp = into_uci_hal_packets(uwb_uci_packets::SessionDeinitRspBuilder {
+                    status: uwb_uci_packets::StatusCode::UciStatusOk,
+                });
+
+                hal.expected_send_command(cmd, resp, Ok(()));
+            },
+            UciLoggerMode::Disabled,
+            mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1515,13 +1798,49 @@ mod tests {
     #[tokio::test]
     async fn test_session_set_app_config_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
         let config_tlv = AppConfigTlv::new(AppConfigTlvType::DeviceType, vec![0x12, 0x34, 0x56]);
         let config_tlv_clone = config_tlv.clone();
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            |hal| {
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd = UciCommand::SessionSetAppConfig {
+                    session_token,
+                    config_tlvs: vec![config_tlv_clone],
+                };
+                let resp = into_uci_hal_packets(uwb_uci_packets::SessionSetAppConfigRspBuilder {
+                    status: uwb_uci_packets::StatusCode::UciStatusOk,
+                    cfg_status: vec![],
+                });
+
+                hal.expected_send_command(cmd, resp, Ok(()));
+            },
+            UciLoggerMode::Disabled,
+            mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
+        )
+        .await;
+
+        let expected_result =
+            SetAppConfigResponse { status: StatusCode::UciStatusOk, config_status: vec![] };
+        let result =
+            uci_manager.session_set_app_config(session_id, vec![config_tlv]).await.unwrap();
+        assert_eq!(result, expected_result);
+        assert!(mock_hal.wait_expected_calls_done().await);
+    }
+
+    #[tokio::test]
+    async fn test_session_set_app_config_v2_ok() {
+        let session_id = 0x123;
+        let session_token = 0x321;
+        let config_tlv = AppConfigTlv::new(AppConfigTlvType::DeviceType, vec![0x12, 0x34, 0x56]);
+        let config_tlv_clone = config_tlv.clone();
+
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
                 let cmd = UciCommand::SessionSetAppConfig {
-                    session_id,
+                    session_token,
                     config_tlvs: vec![config_tlv_clone],
                 };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionSetAppConfigRspBuilder {
@@ -1533,6 +1852,8 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1547,13 +1868,15 @@ mod tests {
     #[tokio::test]
     async fn test_session_get_app_config_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
         let config_id = AppConfigTlvType::DeviceType;
         let tlv = AppConfigTlv::new(AppConfigTlvType::DeviceType, vec![0x12, 0x34, 0x56]);
         let tlv_clone = tlv.clone();
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                let cmd = UciCommand::SessionGetAppConfig { session_id, app_cfg: vec![config_id] };
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd =
+                    UciCommand::SessionGetAppConfig { session_token, app_cfg: vec![config_id] };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionGetAppConfigRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
                     tlvs: vec![tlv_clone.into_inner()],
@@ -1563,6 +1886,8 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1577,7 +1902,7 @@ mod tests {
         let session_count = 5;
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::SessionGetCount;
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionGetCountRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
@@ -1599,11 +1924,12 @@ mod tests {
     #[tokio::test]
     async fn test_session_get_state_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
         let session_state = SessionState::SessionStateActive;
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                let cmd = UciCommand::SessionGetState { session_id };
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd = UciCommand::SessionGetState { session_token };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionGetStateRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
                     session_state,
@@ -1613,6 +1939,8 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1624,14 +1952,15 @@ mod tests {
     #[tokio::test]
     async fn test_session_update_controller_multicast_list_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
         let action = UpdateMulticastListAction::AddControlee;
         let controlee = Controlee { short_address: 0x4567, subsession_id: 0x90ab };
         let controlee_clone = controlee.clone();
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
                 let cmd = UciCommand::SessionUpdateControllerMulticastList {
-                    session_id,
+                    session_token,
                     action,
                     controlees: Controlees::NoSessionKey(vec![controlee_clone]),
                 };
@@ -1645,6 +1974,8 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1661,15 +1992,18 @@ mod tests {
 
     #[tokio::test]
     async fn test_set_active_dt_tag_ranging_rounds() {
+        let session_id = 0x123;
+        let session_token = 0x123;
+
         let ranging_rounds = SessionUpdateDtTagRangingRoundsResponse {
             status: StatusCode::UciStatusErrorRoundIndexNotActivated,
             ranging_round_indexes: vec![3],
         };
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
                 let cmd = UciCommand::SessionUpdateDtTagRangingRounds {
-                    session_id: 1,
+                    session_token,
                     ranging_round_indexes: vec![3, 5],
                 };
                 let resp = into_uci_hal_packets(
@@ -1683,10 +2017,13 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
-        let result = uci_manager.session_update_dt_tag_ranging_rounds(1, vec![3, 5]).await.unwrap();
+        let result =
+            uci_manager.session_update_dt_tag_ranging_rounds(session_id, vec![3, 5]).await.unwrap();
 
         assert_eq!(result, ranging_rounds);
         assert!(mock_hal.wait_expected_calls_done().await);
@@ -1695,10 +2032,11 @@ mod tests {
     #[tokio::test]
     async fn test_range_start_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                let cmd = UciCommand::SessionStart { session_id };
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd = UciCommand::SessionStart { session_token };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionStartRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
                 });
@@ -1707,6 +2045,8 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1718,10 +2058,11 @@ mod tests {
     #[tokio::test]
     async fn test_range_stop_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                let cmd = UciCommand::SessionStop { session_id };
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd = UciCommand::SessionStop { session_token };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionStopRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
                 });
@@ -1730,6 +2071,8 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1741,11 +2084,12 @@ mod tests {
     #[tokio::test]
     async fn test_range_get_ranging_count_ok() {
         let session_id = 0x123;
+        let session_token = 0x123;
         let count = 3;
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                let cmd = UciCommand::SessionGetRangingCount { session_id };
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
+            |mut hal| async move {
+                let cmd = UciCommand::SessionGetRangingCount { session_token };
                 let resp =
                     into_uci_hal_packets(uwb_uci_packets::SessionGetRangingCountRspBuilder {
                         status: uwb_uci_packets::StatusCode::UciStatusOk,
@@ -1756,6 +2100,8 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
@@ -1770,7 +2116,7 @@ mod tests {
         let country_code_clone = country_code.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::AndroidSetCountryCode { country_code: country_code_clone };
                 let resp = into_uci_hal_packets(uwb_uci_packets::AndroidSetCountryCodeRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
@@ -1800,7 +2146,7 @@ mod tests {
         let power_stats_clone = power_stats.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::AndroidGetPowerStats;
                 let resp = into_uci_hal_packets(uwb_uci_packets::AndroidGetPowerStatsRspBuilder {
                     stats: power_stats_clone,
@@ -1829,7 +2175,7 @@ mod tests {
         let resp_payload_clone = resp_payload.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd { mt, gid, oid, payload: cmd_payload_clone };
                 let resp = into_uci_hal_packets(uwb_uci_packets::UciVendor_F_ResponseBuilder {
                     opcode: oid as u8,
@@ -1863,7 +2209,7 @@ mod tests {
         let cfg_status = vec![app_config];
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd { mt, gid, oid, payload: cmd_payload_clone };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionSetAppConfigRspBuilder {
                     status,
@@ -1897,7 +2243,7 @@ mod tests {
         let cfg_status = vec![app_config];
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd { mt, gid, oid, payload: cmd_payload_clone };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionSetAppConfigRspBuilder {
                     status,
@@ -1932,7 +2278,7 @@ mod tests {
         let resp_payload_clone = resp_payload.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd {
                     mt: cmd_mt.into(),
                     gid: gid.into(),
@@ -1971,7 +2317,7 @@ mod tests {
         resp_payload_expected.extend(resp_payload_fragment_2.clone().into_iter());
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd {
                     mt: cmd_mt.into(),
                     gid: gid.into(),
@@ -2024,7 +2370,7 @@ mod tests {
         let cfg_status = vec![app_config];
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd { mt, gid, oid, payload: cmd_payload_clone };
                 let resp = into_uci_hal_packets(uwb_uci_packets::SessionSetAppConfigRspBuilder {
                     status,
@@ -2054,7 +2400,7 @@ mod tests {
         let cmd_payload = vec![0x11, 0x22, 0x33, 0x44];
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |_hal| {},
+            move |_hal| async {},
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
         )
@@ -2077,7 +2423,7 @@ mod tests {
         let cmd_payload = vec![0x11, 0x22, 0x33, 0x44];
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |_hal| {},
+            |_hal| async move {},
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
         )
@@ -2098,13 +2444,13 @@ mod tests {
         let oid: u8 = 0x3; // SESSION_SET_APP_CONFIG
         let cmd_payload = vec![0x11, 0x22, 0x33, 0x44];
         let cmd_payload_clone = cmd_payload.clone();
-        let session_id = 0x123;
+        let session_token = 0x123;
         let resp_mt: u8 = 0x2;
         let resp_payload = vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08];
         let resp_payload_clone = resp_payload.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd {
                     mt: cmd_mt.into(),
                     gid: gid.into(),
@@ -2114,7 +2460,7 @@ mod tests {
                 let raw_resp = build_uci_packet(resp_mt, 0, gid, oid, resp_payload_clone);
                 let mut responses =
                     into_uci_hal_packets(uwb_uci_packets::SessionStatusNtfBuilder {
-                        session_id,
+                        session_token,
                         session_state: uwb_uci_packets::SessionState::SessionStateInit,
                         reason_code:
                             uwb_uci_packets::ReasonCode::StateChangeWithSessionManagementCommands
@@ -2149,7 +2495,7 @@ mod tests {
         let resp_payload = vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08];
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::RawUciCmd {
                     mt: cmd_mt.into(),
                     gid: gid.into(),
@@ -2171,27 +2517,73 @@ mod tests {
         assert!(mock_hal.wait_expected_calls_done().await);
     }
 
-    fn setup_active_session(hal: &mut MockUciHal, session_id: u32) {
-        // First setup the Session to be in Active state.
-        let cmd = UciCommand::SessionStart { session_id };
+    // TODO(b/276320369): Listing down the Data Packet Rx scenarios below, will add unit tests
+    // for them in subsequent CLs.
+    #[tokio::test]
+    async fn test_data_packet_recv_ok() {}
+
+    #[tokio::test]
+    async fn test_data_packet_recv_fragmented_packet_ok() {}
+
+    fn setup_hal_for_session_active(
+        hal: &mut MockUciHal,
+        session_type: SessionType,
+        session_id: u32,
+        session_token: u32,
+    ) {
+        // Setup session init.
+        setup_hal_for_session_initialize(hal, session_type, session_id, session_token);
+
+        // Setup session active.
+        let cmd = UciCommand::SessionStart { session_token };
         let mut responses = into_uci_hal_packets(uwb_uci_packets::SessionStartRspBuilder {
             status: uwb_uci_packets::StatusCode::UciStatusOk,
         });
         responses.append(&mut into_uci_hal_packets(uwb_uci_packets::SessionStatusNtfBuilder {
-            session_id,
+            session_token,
             session_state: SessionState::SessionStateActive,
             reason_code: 0, /* ReasonCode::StateChangeWithSessionManagementCommands */
         }));
         hal.expected_send_command(cmd, responses, Ok(()));
     }
 
-    // TODO(b/276320369): Listing down the Data Packet Rx scenarios below, will add unit tests
-    // for them in subsequent CLs.
-    #[tokio::test]
-    async fn test_data_packet_recv_ok() {}
+    async fn setup_uci_manager_with_session_active<F, Fut>(
+        setup_hal_fn: F,
+        uci_logger_mode: UciLoggerMode,
+        log_sender: mpsc::UnboundedSender<UciLogEvent>,
+        session_id: u32,
+        session_token: u32,
+    ) -> (UciManagerImpl, MockUciHal)
+    where
+        F: FnOnce(MockUciHal) -> Fut,
+        Fut: Future<Output = ()>,
+    {
+        let session_type = SessionType::FiraRangingSession;
 
-    #[tokio::test]
-    async fn test_data_packet_recv_fragmented_packet_ok() {}
+        init_test_logging();
+
+        let mut hal = MockUciHal::new();
+        setup_hal_for_session_active(&mut hal, session_type, session_id, session_token);
+
+        // Verify open_hal() is working.
+        let uci_manager =
+            UciManagerImpl::new(hal.clone(), MockUciLogger::new(log_sender), uci_logger_mode);
+        let result = uci_manager.open_hal().await;
+        assert!(result.is_ok());
+
+        // Verify session is initialized.
+        let result = uci_manager.session_init(session_id, session_type).await;
+        assert!(result.is_ok());
+
+        // Verify session is started.
+        let result = uci_manager.range_start(session_id).await;
+        assert!(result.is_ok());
+        assert!(hal.wait_expected_calls_done().await);
+
+        setup_hal_fn(hal.clone()).await;
+
+        (uci_manager, hal)
+    }
 
     #[tokio::test]
     async fn test_data_packet_send_ok() {
@@ -2201,6 +2593,7 @@ mod tests {
         let dpf = 0x1;
         let oid = 0x0;
         let session_id = 0x5;
+        let session_token = 0x5;
         let dest_mac_address = vec![0xa0, 0xb0, 0xc0, 0xd0, 0xa1, 0xb1, 0xc1, 0xd1];
         let dest_fira_component = FiraComponent::Host;
         let uci_sequence_number = 0xa;
@@ -2215,20 +2608,18 @@ mod tests {
         ];
         let status = DataTransferNtfStatusCode::UciDataTransferStatusRepetitionOk;
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                setup_active_session(hal, session_id);
-
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_active(
+            |mut hal| async move {
                 // Now setup the notifications that should be received after a Data packet send.
                 let data_packet_snd =
                     build_uci_packet(mt_data, pbf, dpf, oid, expected_data_snd_payload);
                 let mut ntfs = into_uci_hal_packets(uwb_uci_packets::DataCreditNtfBuilder {
-                    session_id,
+                    session_token,
                     credit_availability: CreditAvailability::CreditAvailable,
                 });
                 ntfs.append(&mut into_uci_hal_packets(
                     uwb_uci_packets::DataTransferStatusNtfBuilder {
-                        session_id,
+                        session_token,
                         uci_sequence_number,
                         status,
                     },
@@ -2237,12 +2628,11 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
-        let result = uci_manager.range_start(session_id).await;
-        assert!(result.is_ok());
-
         let result = uci_manager
             .send_data_packet(
                 session_id,
@@ -2268,6 +2658,7 @@ mod tests {
         let dpf = 0x1;
         let oid = 0x0;
         let session_id = 0x5;
+        let session_token = 0x5;
         let dest_mac_address = vec![0xa0, 0xb0, 0xc0, 0xd0, 0xa1, 0xb1, 0xc1, 0xd1];
         let dest_fira_component = FiraComponent::Host;
         let uci_sequence_number = 0xa;
@@ -2294,10 +2685,8 @@ mod tests {
             }
         }
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                setup_active_session(hal, session_id);
-
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_active(
+            |mut hal| async move {
                 // Expected data packet fragment #1 (UCI Header + Initial App data bytes).
                 let data_packet_snd_fragment_1 = build_uci_packet(
                     mt_data,
@@ -2307,7 +2696,7 @@ mod tests {
                     expected_data_snd_payload_fragment_1,
                 );
                 let ntfs = into_uci_hal_packets(uwb_uci_packets::DataCreditNtfBuilder {
-                    session_id,
+                    session_token,
                     credit_availability: CreditAvailability::CreditAvailable,
                 });
                 hal.expected_send_packet(data_packet_snd_fragment_1, ntfs, Ok(()));
@@ -2321,12 +2710,12 @@ mod tests {
                     expected_data_snd_payload_fragment_2,
                 );
                 let mut ntfs = into_uci_hal_packets(uwb_uci_packets::DataCreditNtfBuilder {
-                    session_id,
+                    session_token,
                     credit_availability: CreditAvailability::CreditAvailable,
                 });
                 ntfs.append(&mut into_uci_hal_packets(
                     uwb_uci_packets::DataTransferStatusNtfBuilder {
-                        session_id,
+                        session_token,
                         uci_sequence_number,
                         status,
                     },
@@ -2335,12 +2724,11 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
-        let result = uci_manager.range_start(session_id).await;
-        assert!(result.is_ok());
-
         let result = uci_manager
             .send_data_packet(
                 session_id,
@@ -2362,6 +2750,7 @@ mod tests {
         let dpf = 0x1;
         let oid = 0x0;
         let session_id = 0x5;
+        let session_token = 0x5;
         let dest_mac_address = vec![0xa0, 0xb0, 0xc0, 0xd0, 0xa1, 0xb1, 0xc1, 0xd1];
         let dest_fira_component = FiraComponent::Host;
         let uci_sequence_number = 0xa;
@@ -2376,10 +2765,8 @@ mod tests {
         ];
         let status = DataTransferNtfStatusCode::UciDataTransferStatusRepetitionOk;
 
-        let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
-                setup_active_session(hal, session_id);
-
+        let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_active(
+            |mut hal| async move {
                 // Setup receiving a CORE_GENERIC_ERROR_NTF with STATUS_COMMAND_RETRY after a
                 // failed Data packet send attempt.
                 let data_packet_snd =
@@ -2392,12 +2779,12 @@ mod tests {
                 // Setup the notifications that should be received after the Data packet send
                 // is successfully retried.
                 let mut ntfs = into_uci_hal_packets(uwb_uci_packets::DataCreditNtfBuilder {
-                    session_id,
+                    session_token,
                     credit_availability: CreditAvailability::CreditAvailable,
                 });
                 ntfs.append(&mut into_uci_hal_packets(
                     uwb_uci_packets::DataTransferStatusNtfBuilder {
-                        session_id,
+                        session_token,
                         uci_sequence_number,
                         status,
                     },
@@ -2406,12 +2793,11 @@ mod tests {
             },
             UciLoggerMode::Disabled,
             mpsc::unbounded_channel::<UciLogEvent>().0,
+            session_id,
+            session_token,
         )
         .await;
 
-        let result = uci_manager.range_start(session_id).await;
-        assert!(result.is_ok());
-
         let result = uci_manager
             .send_data_packet(
                 session_id,
@@ -2469,7 +2855,7 @@ mod tests {
     #[tokio::test]
     async fn test_session_get_count_retry_no_response() {
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::SessionGetCount;
                 hal.expected_send_command(cmd, vec![], Ok(()));
             },
@@ -2486,7 +2872,7 @@ mod tests {
     #[tokio::test]
     async fn test_session_get_count_timeout() {
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::SessionGetCount;
                 hal.expected_send_command(cmd, vec![], Err(Error::Timeout));
             },
@@ -2503,7 +2889,7 @@ mod tests {
     #[tokio::test]
     async fn test_session_get_count_retry_too_many_times() {
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::SessionGetCount;
                 let retry_resp = into_uci_hal_packets(uwb_uci_packets::SessionGetCountRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusCommandRetry,
@@ -2529,7 +2915,7 @@ mod tests {
         let session_count = 5;
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::SessionGetCount;
                 let retry_resp = into_uci_hal_packets(uwb_uci_packets::SessionGetCountRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusCommandRetry,
@@ -2558,7 +2944,7 @@ mod tests {
     async fn test_log_manager_interaction() {
         let (log_sender, mut log_receiver) = mpsc::unbounded_channel::<UciLogEvent>();
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_open_hal(
-            move |hal| {
+            |mut hal| async move {
                 let cmd = UciCommand::SessionGetCount;
                 let resp1 = into_uci_hal_packets(uwb_uci_packets::SessionGetCountRspBuilder {
                     status: uwb_uci_packets::StatusCode::UciStatusOk,
diff --git a/src/rust/uwb_uci_packets/src/lib.rs b/src/rust/uwb_uci_packets/src/lib.rs
index c5bb3f5..ce63bfb 100644
--- a/src/rust/uwb_uci_packets/src/lib.rs
+++ b/src/rust/uwb_uci_packets/src/lib.rs
@@ -716,7 +716,7 @@ impl PacketDefrager {
 #[allow(dead_code)]
 #[derive(Debug, Clone)]
 pub struct ParsedDiagnosticNtfPacket {
-    session_id: u32,
+    session_token: u32,
     sequence_number: u32,
     frame_reports: Vec<ParsedFrameReport>,
 }
@@ -733,7 +733,7 @@ pub struct ParsedFrameReport {
 }
 
 pub fn parse_diagnostics_ntf(evt: AndroidRangeDiagnosticsNtf) -> Result<ParsedDiagnosticNtfPacket> {
-    let session_id = evt.get_session_id();
+    let session_token = evt.get_session_token();
     let sequence_number = evt.get_sequence_number();
     let mut parsed_frame_reports = Vec::new();
     for report in evt.get_frame_reports() {
@@ -773,7 +773,7 @@ pub fn parse_diagnostics_ntf(evt: AndroidRangeDiagnosticsNtf) -> Result<ParsedDi
         });
     }
     Ok(ParsedDiagnosticNtfPacket {
-        session_id,
+        session_token,
         sequence_number,
         frame_reports: parsed_frame_reports,
     })
@@ -821,7 +821,7 @@ pub fn write_controlee_2_0_32byte(controlee: &Controlee_V2_0_32_Byte_Version) ->
 /// This function can build the packet with/without message control, which
 /// is indicated by action parameter.
 pub fn build_session_update_controller_multicast_list_cmd(
-    session_id: u32,
+    session_token: u32,
     action: UpdateMulticastListAction,
     controlees: Controlees,
 ) -> Result<SessionUpdateControllerMulticastListCmd> {
@@ -855,7 +855,7 @@ pub fn build_session_update_controller_multicast_list_cmd(
         _ => return Err(Error::InvalidPacketError),
     }
     Ok(SessionUpdateControllerMulticastListCmdBuilder {
-        session_id,
+        session_token,
         action,
         payload: Some(controlees_buf.freeze()),
     }
@@ -903,9 +903,12 @@ mod tests {
         let frame_report =
             FrameReport { uwb_msg_id: 1, action: 1, antenna_set: 1, frame_report_tlvs: tlvs };
         frame_reports.push(frame_report);
-        let packet =
-            AndroidRangeDiagnosticsNtfBuilder { session_id: 1, sequence_number: 1, frame_reports }
-                .build();
+        let packet = AndroidRangeDiagnosticsNtfBuilder {
+            session_token: 1,
+            sequence_number: 1,
+            frame_reports,
+        }
+        .build();
         let mut parsed_packet = parse_diagnostics_ntf(packet).unwrap();
         let parsed_frame_report = parsed_packet.frame_reports.pop().unwrap();
         assert_eq!(rssi_vec, parsed_frame_report.rssi);
diff --git a/src/rust/uwb_uci_packets/uci_packets.pdl b/src/rust/uwb_uci_packets/uci_packets.pdl
index 848195a..cf065d8 100644
--- a/src/rust/uwb_uci_packets/uci_packets.pdl
+++ b/src/rust/uwb_uci_packets/uci_packets.pdl
@@ -517,7 +517,7 @@ packet UciDataPacket {
 }
 
 packet UciDataSnd : UciDataPacket (data_packet_format = DATA_SND, message_type = DATA) {
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     dest_mac_address: 64,
     dest_fira_component: FiraComponent,
     uci_sequence_number: 8,
@@ -526,7 +526,7 @@ packet UciDataSnd : UciDataPacket (data_packet_format = DATA_SND, message_type =
 }
 
 packet UciDataRcv : UciDataPacket (data_packet_format = DATA_RCV, message_type = DATA) {
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     status: DataRcvStatusCode,
     uci_sequence_number: 32,
     source_mac_address: 64,
@@ -733,6 +733,17 @@ test SessionInitCmd {
     "\x21\x00\x00\x05\x00\x00\x00\x01\x02\x03\x04\x01",
 }
 
+// FIRA version 2 introduces a new version of SESSION_INIT_RSP which
+// includes UWBS generated session handle.
+packet SessionInitRsp_V2 : SessionConfigResponse (opcode = 0x0) { //SESSION_INIT
+    status: StatusCode,
+    session_handle: 32,
+}
+
+test SessionInitRsp_V2 {
+    "\x41\x00\x00\x01\x00\x00\x00\x11\x00\x00\x00\x01",
+}
+
 packet SessionInitRsp : SessionConfigResponse (opcode = 0x0) { //SESSION_INIT
     status: StatusCode,
 }
@@ -742,7 +753,7 @@ test SessionInitRsp {
 }
 
 packet SessionDeinitCmd : SessionConfigCommand (opcode = 0x1) { //SESSION_DEINIT
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
 }
 
 test SessionDeinitCmd {
@@ -758,7 +769,7 @@ test SessionDeinitRsp {
 }
 
 packet SessionStatusNtf : SessionConfigNotification (opcode = 0x2) { //SESSION_STATUS_NTF
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     session_state: SessionState,
     // TODO(b/272775225): Switch back to the enum type ReasonCode, once PDL supports defining a
     // range inside an enum (for the vendor-specific space), in b/267339120.
@@ -777,7 +788,7 @@ struct AppConfigTlv {
 }
 
 packet SessionSetAppConfigCmd : SessionConfigCommand (opcode = 0x3) { //SESSION_SET_APP_CONFIG
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     _count_(tlvs): 8,
     tlvs: AppConfigTlv[]
 }
@@ -802,7 +813,7 @@ test SessionSetAppConfigRsp {
 }
 
 packet SessionGetAppConfigCmd : SessionConfigCommand (opcode = 0x4) { //SESSION_GET_APP_CONFIG
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     _count_(app_cfg): 8,
     app_cfg: 8[], // AppConfigTlvType (Infra does not allow array of enums)
 }
@@ -838,7 +849,7 @@ test SessionGetCountRsp {
 }
 
 packet SessionGetStateCmd : SessionConfigCommand (opcode = 0x6) { //SESSION_GET_STATE
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
 }
 
 test SessionGetStateCmd {
@@ -855,7 +866,7 @@ test SessionGetStateRsp {
 }
 
 packet SessionUpdateDtTagRangingRoundsCmd : SessionConfigCommand (opcode = 0x9) { //SESSION_UPDATE_ACTIVE_ROUNDS_DT_TAG
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     _count_(ranging_round_indexes): 8,
     ranging_round_indexes: 8[],
 }
@@ -899,7 +910,7 @@ enum UpdateMulticastListAction: 8 {
 }
 
 packet SessionUpdateControllerMulticastListCmd : SessionConfigCommand (opcode = 0x7) { //SESSION_UPDATE_CONTROLLER_MULTICAST_LIST
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     action: UpdateMulticastListAction,
     _payload_,
 }
@@ -934,7 +945,7 @@ struct ControleeStatus {
 }
 
 packet SessionUpdateControllerMulticastListNtf : SessionConfigNotification (opcode = 0x7) { //SESSION_UPDATE_CONTROLLER_MULTICAST_LIST
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     remaining_multicast_list_size: 8,
     _count_(controlee_status): 8,
     controlee_status: ControleeStatus[],
@@ -945,7 +956,7 @@ test SessionUpdateControllerMulticastListNtf {
 }
 
 packet DataCreditNtf : SessionControlNotification (opcode = 0x04) { // SESSION_DATA_CREDIT_NTF
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     credit_availability: CreditAvailability,
 }
 
@@ -954,7 +965,7 @@ test DataCreditNtf {
 }
 
 packet DataTransferStatusNtf : SessionControlNotification (opcode = 0x05) { // SESSION_DATA_TRANSFER_STATUS_NTF
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     uci_sequence_number: 8,
     status: DataTransferNtfStatusCode,
     // TODO(b/269779288): Add the tx_count field for implementing the DATA_REPETITION added in CR490.
@@ -965,7 +976,7 @@ test DataTransferStatusNtf {
 }
 
 packet SessionQueryMaxDataSizeCmd : SessionConfigCommand (opcode = 0xB) { //QUERY_MAX_DATA_SIZE
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
 }
 
 test SessionQueryMaxDataSizeCmd {
@@ -973,7 +984,7 @@ test SessionQueryMaxDataSizeCmd {
 }
 
 packet SessionQueryMaxDataSizeRsp : SessionConfigResponse (opcode = 0xB) { //QUER_MAX_DATA_SIZE
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     max_data_size: 16,
 }
 
@@ -982,7 +993,7 @@ test SessionQueryMaxDataSizeRsp {
 }
 
 packet SessionStartCmd : SessionControlCommand (opcode = 0x0) { //RANGE_START
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
 }
 
 test SessionStartCmd {
@@ -1070,7 +1081,7 @@ enum RangingMeasurementType : 8 {
 
 packet SessionInfoNtf : SessionControlNotification (opcode = 0x0) { // SESSION_INFO
     sequence_number: 32,
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     rcr_indicator: 8,
     current_ranging_interval: 32,
     ranging_measurement_type: RangingMeasurementType,
@@ -1141,7 +1152,7 @@ test ExtendedMacOwrAoaSessionInfoNtf {
 }
 
 packet SessionStopCmd : SessionControlCommand (opcode = 0x1) { // SESSION_STOP
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
 }
 
 test SessionStopCmd {
@@ -1157,7 +1168,7 @@ test SessionStopRsp {
 }
 
 packet SessionGetRangingCountCmd : SessionControlCommand (opcode = 0x3) { // SESSION_GET_RANGING_COUNT
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
 }
 
 test SessionGetRangingCountCmd {
@@ -1276,7 +1287,7 @@ struct FrameReport {
 }
 
 packet AndroidRangeDiagnosticsNtf : AndroidNotification (opcode = 0x2) { //FIRA_RANGE_DIAGNOSTICS
-    session_id: 32,
+    session_token: 32, // Session ID or Session Handle (based on UWBS version)
     sequence_number: 32,
     _count_(frame_reports): 8,
     frame_reports: FrameReport[],