diff --git a/security/authgraph/aidl/vts/functional/sink.rs b/security/authgraph/aidl/vts/functional/sink.rs
index bb357b8..a331eef 100644
--- a/security/authgraph/aidl/vts/functional/sink.rs
+++ b/security/authgraph/aidl/vts/functional/sink.rs
@@ -29,11 +29,11 @@
 }
 
 /// Perform mainline AuthGraph key exchange with the provided sink and local implementation.
-/// Return the agreed AES keys in plaintext.
+/// Return the agreed AES keys in plaintext, together with the session ID.
 pub fn test_mainline(
     local_source: &mut ke::AuthGraphParticipant,
     sink: binder::Strong<dyn IAuthGraphKeyExchange>,
-) -> [key::AesKey; 2] {
+) -> ([key::AesKey; 2], Vec<u8>) {
     // Step 1: create an ephemeral ECDH key at the (local) source.
     let source_init_info = local_source
         .create()
@@ -113,7 +113,7 @@
         Ok(array) => array,
         Err(_) => panic!("wrong number of decrypted shared key arcs"),
     };
-    decrypted_shared_keys_array
+    (decrypted_shared_keys_array, sink_info.sessionId)
 }
 
 /// Perform mainline AuthGraph key exchange with the provided sink, but provide an invalid
diff --git a/security/authgraph/aidl/vts/functional/source.rs b/security/authgraph/aidl/vts/functional/source.rs
index a1e76b3..019e1e8 100644
--- a/security/authgraph/aidl/vts/functional/source.rs
+++ b/security/authgraph/aidl/vts/functional/source.rs
@@ -29,11 +29,11 @@
 }
 
 /// Perform mainline AuthGraph key exchange with the provided source.
-/// Return the agreed AES keys in plaintext.
+/// Return the agreed AES keys in plaintext, together with the session ID.
 pub fn test_mainline(
     local_sink: &mut ke::AuthGraphParticipant,
     source: binder::Strong<dyn IAuthGraphKeyExchange>,
-) -> [key::AesKey; 2] {
+) -> ([key::AesKey; 2], Vec<u8>) {
     // Step 1: create an ephemeral ECDH key at the (remote) source.
     let source_init_info = source
         .create()
@@ -120,7 +120,7 @@
         Ok(array) => array,
         Err(_) => panic!("wrong number of decrypted shared key arcs"),
     };
-    decrypted_shared_keys_array
+    (decrypted_shared_keys_array, source_info.sessionId)
 }
 
 /// Perform mainline AuthGraph key exchange with the provided source, but provide an invalid session