Save which user session created a compat session

This also exposes the user session in the GraphQL API, and allow
filtering on browser session ID on the app session list.

crates/cli/src/commands/manage.rs

@@ -184,7 +184,7 @@ impl Options {
 
                 let compat_session = repo
                     .compat_session()
-                    .add(&mut rng, &clock, &user, device, admin)
+                    .add(&mut rng, &clock, &user, device, None, admin)
                     .await?;
 
                 let token = TokenType::CompatAccessToken.generate(&mut rng);

crates/data-model/src/compat/session.rs

@@ -80,6 +80,7 @@ pub struct CompatSession {
     pub state: CompatSessionState,
     pub user_id: Ulid,
     pub device: Device,
+    pub user_session_id: Option<Ulid>,
     pub created_at: DateTime<Utc>,
     pub is_synapse_admin: bool,
     pub last_active_at: Option<DateTime<Utc>>,

crates/graphql/src/model/browser_sessions.rs

@@ -12,11 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use async_graphql::{Context, Description, Object, ID};
+use async_graphql::{
+    connection::{query, Connection, Edge, OpaqueCursor},
+    Context, Description, Object, ID,
+};
 use chrono::{DateTime, Utc};
-use mas_storage::{user::BrowserSessionRepository, RepositoryAccess};
+use mas_data_model::Device;
+use mas_storage::{
+    app_session::AppSessionFilter, user::BrowserSessionRepository, Pagination, RepositoryAccess,
+};
 
-use super::{NodeType, SessionState, User};
+use super::{
+    AppSession, CompatSession, Cursor, NodeCursor, NodeType, OAuth2Session, PreloadedTotalCount,
+    SessionState, User,
+};
 use crate::state::ContextExt;
 
 /// A browser session represents a logged in user in a browser.
@@ -92,6 +101,97 @@ impl BrowserSession {
     pub async fn last_active_at(&self) -> Option<DateTime<Utc>> {
         self.0.last_active_at
     }
+
+    /// Get the list of both compat and OAuth 2.0 sessions started by this
+    /// browser session, chronologically sorted
+    #[allow(clippy::too_many_arguments)]
+    async fn app_sessions(
+        &self,
+        ctx: &Context<'_>,
+
+        #[graphql(name = "state", desc = "List only sessions in the given state.")]
+        state_param: Option<SessionState>,
+
+        #[graphql(name = "device", desc = "List only sessions for the given device.")]
+        device_param: Option<String>,
+
+        #[graphql(desc = "Returns the elements in the list that come after the cursor.")]
+        after: Option<String>,
+        #[graphql(desc = "Returns the elements in the list that come before the cursor.")]
+        before: Option<String>,
+        #[graphql(desc = "Returns the first *n* elements from the list.")] first: Option<i32>,
+        #[graphql(desc = "Returns the last *n* elements from the list.")] last: Option<i32>,
+    ) -> Result<Connection<Cursor, AppSession, PreloadedTotalCount>, async_graphql::Error> {
+        let state = ctx.state();
+        let mut repo = state.repository().await?;
+
+        query(
+            after,
+            before,
+            first,
+            last,
+            |after, before, first, last| async move {
+                let after_id = after
+                    .map(|x: OpaqueCursor<NodeCursor>| {
+                        x.extract_for_types(&[NodeType::OAuth2Session, NodeType::CompatSession])
+                    })
+                    .transpose()?;
+                let before_id = before
+                    .map(|x: OpaqueCursor<NodeCursor>| {
+                        x.extract_for_types(&[NodeType::OAuth2Session, NodeType::CompatSession])
+                    })
+                    .transpose()?;
+                let pagination = Pagination::try_new(before_id, after_id, first, last)?;
+
+                let device_param = device_param.map(Device::try_from).transpose()?;
+
+                let filter = AppSessionFilter::new().for_browser_session(&self.0);
+
+                let filter = match state_param {
+                    Some(SessionState::Active) => filter.active_only(),
+                    Some(SessionState::Finished) => filter.finished_only(),
+                    None => filter,
+                };
+
+                let filter = match device_param.as_ref() {
+                    Some(device) => filter.for_device(device),
+                    None => filter,
+                };
+
+                let page = repo.app_session().list(filter, pagination).await?;
+
+                let count = if ctx.look_ahead().field("totalCount").exists() {
+                    Some(repo.app_session().count(filter).await?)
+                } else {
+                    None
+                };
+
+                repo.cancel().await?;
+
+                let mut connection = Connection::with_additional_fields(
+                    page.has_previous_page,
+                    page.has_next_page,
+                    PreloadedTotalCount(count),
+                );
+
+                connection
+                    .edges
+                    .extend(page.edges.into_iter().map(|s| match s {
+                        mas_storage::app_session::AppSession::Compat(session) => Edge::new(
+                            OpaqueCursor(NodeCursor(NodeType::CompatSession, session.id)),
+                            AppSession::CompatSession(Box::new(CompatSession::new(*session))),
+                        ),
+                        mas_storage::app_session::AppSession::OAuth2(session) => Edge::new(
+                            OpaqueCursor(NodeCursor(NodeType::OAuth2Session, session.id)),
+                            AppSession::OAuth2Session(Box::new(OAuth2Session(*session))),
+                        ),
+                    }));
+
+                Ok::<_, async_graphql::Error>(connection)
+            },
+        )
+        .await
+    }
 }
 
 /// An authentication records when a user enter their credential in a browser

crates/graphql/src/model/compat_sessions.rs

@@ -18,7 +18,7 @@ use chrono::{DateTime, Utc};
 use mas_storage::{compat::CompatSessionRepository, user::UserRepository};
 use url::Url;
 
-use super::{NodeType, SessionState, User};
+use super::{BrowserSession, NodeType, SessionState, User};
 use crate::state::ContextExt;
 
 /// Lazy-loaded reverse reference.
@@ -125,6 +125,27 @@ impl CompatSession {
         Ok(sso_login.map(CompatSsoLogin))
     }
 
+    /// The browser session which started this session, if any.
+    pub async fn browser_session(
+        &self,
+        ctx: &Context<'_>,
+    ) -> Result<Option<BrowserSession>, async_graphql::Error> {
+        let Some(user_session_id) = self.session.user_session_id else {
+            return Ok(None);
+        };
+
+        let state = ctx.state();
+        let mut repo = state.repository().await?;
+        let browser_session = repo
+            .browser_session()
+            .lookup(user_session_id)
+            .await?
+            .context("Could not load browser session")?;
+        repo.cancel().await?;
+
+        Ok(Some(BrowserSession(browser_session)))
+    }
+
     /// The state of the session.
     pub async fn state(&self) -> SessionState {
         match &self.session.state {

crates/graphql/src/model/mod.rs

@@ -32,7 +32,7 @@ pub use self::{
     node::{Node, NodeType},
     oauth::{OAuth2Client, OAuth2Session},
     upstream_oauth::{UpstreamOAuth2Link, UpstreamOAuth2Provider},
-    users::{User, UserEmail},
+    users::{AppSession, User, UserEmail},
     viewer::{Anonymous, Viewer, ViewerSession},
 };
 

crates/graphql/src/model/users.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use anyhow::Context as _;
 use async_graphql::{
     connection::{query, Connection, Edge, OpaqueCursor},
     Context, Description, Enum, Object, Union, ID,
@@ -544,6 +545,12 @@ impl User {
         #[graphql(name = "device", desc = "List only sessions for the given device.")]
         device_param: Option<String>,
 
+        #[graphql(
+            name = "browserSession",
+            desc = "List only sessions for the given session."
+        )]
+        browser_session_param: Option<ID>,
+
         #[graphql(desc = "Returns the elements in the list that come after the cursor.")]
         after: Option<String>,
         #[graphql(desc = "Returns the elements in the list that come before the cursor.")]
@@ -552,6 +559,7 @@ impl User {
         #[graphql(desc = "Returns the last *n* elements from the list.")] last: Option<i32>,
     ) -> Result<Connection<Cursor, AppSession, PreloadedTotalCount>, async_graphql::Error> {
         let state = ctx.state();
+        let requester = ctx.requester();
         let mut repo = state.repository().await?;
 
         query(
@@ -587,6 +595,38 @@ impl User {
                     None => filter,
                 };
 
+                let maybe_session = match browser_session_param {
+                    Some(id) => {
+                        // This might fail, but we're probably alright with it
+                        let id = NodeType::BrowserSession
+                            .extract_ulid(&id)
+                            .context("Invalid browser_session parameter")?;
+
+                        let Some(session) = repo
+                            .browser_session()
+                            .lookup(id)
+                            .await?
+                            .filter(|u| requester.is_owner_or_admin(u))
+                        else {
+                            // If we couldn't find the session or if the requester can't access it,
+                            // return an empty list
+                            return Ok(Connection::with_additional_fields(
+                                false,
+                                false,
+                                PreloadedTotalCount(Some(0)),
+                            ));
+                        };
+
+                        Some(session)
+                    }
+                    None => None,
+                };
+
+                let filter = match maybe_session {
+                    Some(ref session) => filter.for_browser_session(session),
+                    None => filter,
+                };
+
                 let page = repo.app_session().list(filter, pagination).await?;
 
                 let count = if ctx.look_ahead().field("totalCount").exists() {
@@ -625,7 +665,7 @@ impl User {
 
 /// A session in an application, either a compatibility or an OAuth 2.0 one
 #[derive(Union)]
-enum AppSession {
+pub enum AppSession {
     CompatSession(Box<CompatSession>),
     OAuth2Session(Box<OAuth2Session>),
 }

crates/handlers/src/compat/login.rs

@@ -411,7 +411,7 @@ async fn user_password_login(
 
     let session = repo
         .compat_session()
-        .add(&mut rng, clock, &user, device, false)
+        .add(&mut rng, clock, &user, device, None, false)
         .await?;
 
     Ok((session, user))
@@ -738,7 +738,14 @@ mod tests {
         // Complete the flow by fulfilling it with a session
         let compat_session = repo
             .compat_session()
-            .add(&mut state.rng(), &state.clock, user, device.clone(), false)
+            .add(
+                &mut state.rng(),
+                &state.clock,
+                user,
+                device.clone(),
+                None,
+                false,
+            )
             .await
             .unwrap();
 

crates/handlers/src/compat/login_sso_complete.rs

@@ -208,7 +208,14 @@ pub async fn post(
 
     let compat_session = repo
         .compat_session()
-        .add(&mut rng, &clock, &session.user, device, false)
+        .add(
+            &mut rng,
+            &clock,
+            &session.user,
+            device,
+            Some(&session),
+            false,
+        )
         .await?;
 
     repo.compat_sso_login()

crates/storage-pg/.sqlx/query-04e25c9267bf2eb143a6445345229081e7b386743a93b3833ef8ad9d09972f3b.json

@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                SELECT compat_session_id\n                     , device_id\n                     , user_id\n                     , created_at\n                     , finished_at\n                     , is_synapse_admin\n                     , last_active_at\n                     , last_active_ip as \"last_active_ip: IpAddr\"\n                FROM compat_sessions\n                WHERE compat_session_id = $1\n            ",
+  "query": "\n                SELECT compat_session_id\n                     , device_id\n                     , user_id\n                     , user_session_id\n                     , created_at\n                     , finished_at\n                     , is_synapse_admin\n                     , last_active_at\n                     , last_active_ip as \"last_active_ip: IpAddr\"\n                FROM compat_sessions\n                WHERE compat_session_id = $1\n            ",
   "describe": {
     "columns": [
       {
@@ -20,26 +20,31 @@
       },
       {
         "ordinal": 3,
+        "name": "user_session_id",
+        "type_info": "Uuid"
+      },
+      {
+        "ordinal": 4,
         "name": "created_at",
         "type_info": "Timestamptz"
       },
       {
-        "ordinal": 4,
+        "ordinal": 5,
         "name": "finished_at",
         "type_info": "Timestamptz"
       },
       {
-        "ordinal": 5,
+        "ordinal": 6,
         "name": "is_synapse_admin",
         "type_info": "Bool"
       },
       {
-        "ordinal": 6,
+        "ordinal": 7,
         "name": "last_active_at",
         "type_info": "Timestamptz"
       },
       {
-        "ordinal": 7,
+        "ordinal": 8,
         "name": "last_active_ip: IpAddr",
         "type_info": "Inet"
       }
@@ -53,6 +58,7 @@
       false,
       false,
       false,
+      true,
       false,
       true,
       false,
@@ -60,5 +66,5 @@
       true
     ]
   },
-  "hash": "23c03635d6433099a4353ba0c80b88737724edb16315832891550e29088d02bf"
+  "hash": "04e25c9267bf2eb143a6445345229081e7b386743a93b3833ef8ad9d09972f3b"
 }

crates/storage-pg/.sqlx/query-4070549b235e059eaeccc4751b480ccb30ad5b62d933b4efb03491124a9361ad.json

@@ -0,0 +1,19 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n                INSERT INTO compat_sessions \n                    (compat_session_id, user_id, device_id,\n                     user_session_id, created_at, is_synapse_admin)\n                VALUES ($1, $2, $3, $4, $5, $6)\n            ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Uuid",
+        "Uuid",
+        "Text",
+        "Uuid",
+        "Timestamptz",
+        "Bool"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "4070549b235e059eaeccc4751b480ccb30ad5b62d933b4efb03491124a9361ad"
+}

crates/storage-pg/.sqlx/query-cff3ac0fff62ffdc5640fce08c2ffabc1d89202561b736c5d03b501dfcd8d886.json

@@ -1,18 +0,0 @@
-{
-  "db_name": "PostgreSQL",
-  "query": "\n                INSERT INTO compat_sessions (compat_session_id, user_id, device_id, created_at, is_synapse_admin)\n                VALUES ($1, $2, $3, $4, $5)\n            ",
-  "describe": {
-    "columns": [],
-    "parameters": {
-      "Left": [
-        "Uuid",
-        "Uuid",
-        "Text",
-        "Timestamptz",
-        "Bool"
-      ]
-    },
-    "nullable": []
-  },
-  "hash": "cff3ac0fff62ffdc5640fce08c2ffabc1d89202561b736c5d03b501dfcd8d886"
-}

crates/storage-pg/migrations/20240220150201_compat_sessions_user_sessions_link.sql

@@ -0,0 +1,19 @@
+-- Copyright 2024 The Matrix.org Foundation C.I.C.
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+-- Adds an optional link between the compatibility sessions and the user sessions
+ALTER TABLE compat_sessions
+    ADD COLUMN user_session_id UUID
+        REFERENCES user_sessions (user_session_id)
+        ON DELETE SET NULL;

crates/storage-pg/src/app_session.rs

@@ -102,11 +102,12 @@ impl TryFrom<AppSessionLookup> for AppSession {
             last_active_ip,
         } = value;
 
+        let user_session_id = user_session_id.map(Ulid::from);
+
         match (
             compat_session_id,
             oauth2_session_id,
             oauth2_client_id,
-            user_session_id,
             user_id,
             scope_list,
             device_id,
@@ -116,7 +117,6 @@ impl TryFrom<AppSessionLookup> for AppSession {
                 Some(compat_session_id),
                 None,
                 None,
-                None,
                 Some(user_id),
                 None,
                 Some(device_id),
@@ -140,6 +140,7 @@ impl TryFrom<AppSessionLookup> for AppSession {
                     state,
                     user_id: user_id.into(),
                     device,
+                    user_session_id,
                     created_at,
                     is_synapse_admin,
                     last_active_at,
@@ -153,7 +154,6 @@ impl TryFrom<AppSessionLookup> for AppSession {
                 None,
                 Some(oauth2_session_id),
                 Some(oauth2_client_id),
-                user_session_id,
                 user_id,
                 Some(scope_list),
                 None,
@@ -180,7 +180,7 @@ impl TryFrom<AppSessionLookup> for AppSession {
                     created_at,
                     client_id: oauth2_client_id.into(),
                     user_id: user_id.map(Ulid::from),
-                    user_session_id: user_session_id.map(Ulid::from),
+                    user_session_id,
                     scope,
                     last_active_at,
                     last_active_ip,
@@ -269,6 +269,10 @@ impl<'c> AppSessionRepository for PgAppSessionRepository<'c> {
                     Expr::col((OAuth2Sessions::Table, OAuth2Sessions::FinishedAt)).is_not_null()
                 }
             }))
+            .and_where_option(filter.browser_session().map(|browser_session| {
+                Expr::col((OAuth2Sessions::Table, OAuth2Sessions::UserSessionId))
+                    .eq(Uuid::from(browser_session.id))
+            }))
             .and_where_option(filter.device().map(|device| {
                 Expr::val(device.to_scope_token().to_string()).eq(PgFunc::any(Expr::col((
                     OAuth2Sessions::Table,
@@ -288,7 +292,10 @@ impl<'c> AppSessionRepository for PgAppSessionRepository<'c> {
             )
             .expr_as(Expr::cust("NULL"), AppSessionLookupIden::Oauth2SessionId)
             .expr_as(Expr::cust("NULL"), AppSessionLookupIden::Oauth2ClientId)
-            .expr_as(Expr::cust("NULL"), AppSessionLookupIden::UserSessionId)
+            .expr_as(
+                Expr::col((CompatSessions::Table, CompatSessions::UserSessionId)),
+                AppSessionLookupIden::UserSessionId,
+            )
             .expr_as(
                 Expr::col((CompatSessions::Table, CompatSessions::UserId)),
                 AppSessionLookupIden::UserId,
@@ -329,6 +336,10 @@ impl<'c> AppSessionRepository for PgAppSessionRepository<'c> {
                     Expr::col((CompatSessions::Table, CompatSessions::FinishedAt)).is_not_null()
                 }
             }))
+            .and_where_option(filter.browser_session().map(|browser_session| {
+                Expr::col((CompatSessions::Table, CompatSessions::UserSessionId))
+                    .eq(Uuid::from(browser_session.id))
+            }))
             .and_where_option(filter.device().map(|device| {
                 Expr::col((CompatSessions::Table, CompatSessions::DeviceId)).eq(device.to_string())
             }))
@@ -385,6 +396,10 @@ impl<'c> AppSessionRepository for PgAppSessionRepository<'c> {
                     Expr::col((OAuth2Sessions::Table, OAuth2Sessions::FinishedAt)).is_not_null()
                 }
             }))
+            .and_where_option(filter.browser_session().map(|browser_session| {
+                Expr::col((OAuth2Sessions::Table, OAuth2Sessions::UserSessionId))
+                    .eq(Uuid::from(browser_session.id))
+            }))
             .and_where_option(filter.device().map(|device| {
                 Expr::val(device.to_scope_token().to_string()).eq(PgFunc::any(Expr::col((
                     OAuth2Sessions::Table,
@@ -406,6 +421,10 @@ impl<'c> AppSessionRepository for PgAppSessionRepository<'c> {
                     Expr::col((CompatSessions::Table, CompatSessions::FinishedAt)).is_not_null()
                 }
             }))
+            .and_where_option(filter.browser_session().map(|browser_session| {
+                Expr::col((CompatSessions::Table, CompatSessions::UserSessionId))
+                    .eq(Uuid::from(browser_session.id))
+            }))
             .and_where_option(filter.device().map(|device| {
                 Expr::col((CompatSessions::Table, CompatSessions::DeviceId)).eq(device.to_string())
             }))
@@ -493,7 +512,7 @@ mod tests {
         let device = Device::generate(&mut rng);
         let compat_session = repo
             .compat_session()
-            .add(&mut rng, &clock, &user, device.clone(), false)
+            .add(&mut rng, &clock, &user, device.clone(), None, false)
             .await
             .unwrap();
 

crates/storage-pg/src/compat/mod.rs

@@ -87,7 +87,7 @@ mod tests {
         let device_str = device.as_str().to_owned();
         let session = repo
             .compat_session()
-            .add(&mut rng, &clock, &user, device.clone(), false)
+            .add(&mut rng, &clock, &user, device.clone(), None, false)
             .await
             .unwrap();
         assert_eq!(session.user_id, user.id);
@@ -203,7 +203,7 @@ mod tests {
         let device = Device::generate(&mut rng);
         let sso_login_session = repo
             .compat_session()
-            .add(&mut rng, &clock, &user, device, false)
+            .add(&mut rng, &clock, &user, device, None, false)
             .await
             .unwrap();
 
@@ -291,7 +291,7 @@ mod tests {
         let device = Device::generate(&mut rng);
         let session = repo
             .compat_session()
-            .add(&mut rng, &clock, &user, device, false)
+            .add(&mut rng, &clock, &user, device, None, false)
             .await
             .unwrap();
 
@@ -411,7 +411,7 @@ mod tests {
         let device = Device::generate(&mut rng);
         let session = repo
             .compat_session()
-            .add(&mut rng, &clock, &user, device, false)
+            .add(&mut rng, &clock, &user, device, None, false)
             .await
             .unwrap();
 
@@ -584,7 +584,7 @@ mod tests {
         let device = Device::generate(&mut rng);
         let session = repo
             .compat_session()
-            .add(&mut rng, &clock, &user, device, false)
+            .add(&mut rng, &clock, &user, device, None, false)
             .await
             .unwrap();
 

crates/storage-pg/src/compat/session.rs

@@ -17,7 +17,8 @@ use std::net::IpAddr;
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use mas_data_model::{
-    CompatSession, CompatSessionState, CompatSsoLogin, CompatSsoLoginState, Device, User,
+    BrowserSession, CompatSession, CompatSessionState, CompatSsoLogin, CompatSsoLoginState, Device,
+    User,
 };
 use mas_storage::{
     compat::{CompatSessionFilter, CompatSessionRepository},
@@ -55,6 +56,7 @@ struct CompatSessionLookup {
     compat_session_id: Uuid,
     device_id: String,
     user_id: Uuid,
+    user_session_id: Option<Uuid>,
     created_at: DateTime<Utc>,
     finished_at: Option<DateTime<Utc>>,
     is_synapse_admin: bool,
@@ -83,6 +85,7 @@ impl TryFrom<CompatSessionLookup> for CompatSession {
             id,
             state,
             user_id: value.user_id.into(),
+            user_session_id: value.user_session_id.map(Ulid::from),
             device,
             created_at: value.created_at,
             is_synapse_admin: value.is_synapse_admin,
@@ -100,6 +103,7 @@ struct CompatSessionAndSsoLoginLookup {
     compat_session_id: Uuid,
     device_id: String,
     user_id: Uuid,
+    user_session_id: Option<Uuid>,
     created_at: DateTime<Utc>,
     finished_at: Option<DateTime<Utc>>,
     is_synapse_admin: bool,
@@ -135,6 +139,7 @@ impl TryFrom<CompatSessionAndSsoLoginLookup> for (CompatSession, Option<CompatSs
             state,
             user_id: value.user_id.into(),
             device,
+            user_session_id: value.user_session_id.map(Ulid::from),
             created_at: value.created_at,
             is_synapse_admin: value.is_synapse_admin,
             last_active_at: value.last_active_at,
@@ -214,6 +219,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
                 SELECT compat_session_id
                      , device_id
                      , user_id
+                     , user_session_id
                      , created_at
                      , finished_at
                      , is_synapse_admin
@@ -251,6 +257,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
         clock: &dyn Clock,
         user: &User,
         device: Device,
+        browser_session: Option<&BrowserSession>,
         is_synapse_admin: bool,
     ) -> Result<CompatSession, Self::Error> {
         let created_at = clock.now();
@@ -259,12 +266,15 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
 
         sqlx::query!(
             r#"
-                INSERT INTO compat_sessions (compat_session_id, user_id, device_id, created_at, is_synapse_admin)
-                VALUES ($1, $2, $3, $4, $5)
+                INSERT INTO compat_sessions 
+                    (compat_session_id, user_id, device_id,
+                     user_session_id, created_at, is_synapse_admin)
+                VALUES ($1, $2, $3, $4, $5, $6)
             "#,
             Uuid::from(id),
             Uuid::from(user.id),
             device.as_str(),
+            browser_session.map(|s| Uuid::from(s.id)),
             created_at,
             is_synapse_admin,
         )
@@ -277,6 +287,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
             state: CompatSessionState::default(),
             user_id: user.id,
             device,
+            user_session_id: browser_session.map(|s| s.id),
             created_at,
             is_synapse_admin,
             last_active_at: None,
@@ -350,6 +361,10 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> {
                 Expr::col((CompatSessions::Table, CompatSessions::UserId)),
                 CompatSessionAndSsoLoginLookupIden::UserId,
             )
+            .expr_as(
+                Expr::col((CompatSessions::Table, CompatSessions::UserSessionId)),
+                CompatSessionAndSsoLoginLookupIden::UserSessionId,
+            )
             .expr_as(
                 Expr::col((CompatSessions::Table, CompatSessions::CreatedAt)),
                 CompatSessionAndSsoLoginLookupIden::CreatedAt,

crates/storage-pg/src/iden.rs

@@ -53,6 +53,7 @@ pub enum CompatSessions {
     CompatSessionId,
     UserId,
     DeviceId,
+    UserSessionId,
     CreatedAt,
     FinishedAt,
     IsSynapseAdmin,

crates/storage/src/app_session.rs

@@ -15,7 +15,7 @@
 //! Repositories to interact with all kinds of sessions
 
 use async_trait::async_trait;
-use mas_data_model::{CompatSession, Device, Session, User};
+use mas_data_model::{BrowserSession, CompatSession, Device, Session, User};
 
 use crate::{repository_impl, Page, Pagination};
 
@@ -56,6 +56,7 @@ pub enum AppSession {
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
 pub struct AppSessionFilter<'a> {
     user: Option<&'a User>,
+    browser_session: Option<&'a BrowserSession>,
     state: Option<AppSessionState>,
     device_id: Option<&'a Device>,
 }
@@ -67,7 +68,7 @@ impl<'a> AppSessionFilter<'a> {
         Self::default()
     }
 
-    /// Set the user who owns the compatibility sessions
+    /// Set the user who owns the sessions
     #[must_use]
     pub fn for_user(mut self, user: &'a User) -> Self {
         self.user = Some(user);
@@ -80,6 +81,19 @@ impl<'a> AppSessionFilter<'a> {
         self.user
     }
 
+    /// Set the browser session filter
+    #[must_use]
+    pub fn for_browser_session(mut self, browser_session: &'a BrowserSession) -> Self {
+        self.browser_session = Some(browser_session);
+        self
+    }
+
+    /// Get the browser session filter
+    #[must_use]
+    pub fn browser_session(&self) -> Option<&BrowserSession> {
+        self.browser_session
+    }
+
     /// Set the device ID filter
     #[must_use]
     pub fn for_device(mut self, device_id: &'a Device) -> Self {

crates/storage/src/compat/session.rs

@@ -16,7 +16,7 @@ use std::net::IpAddr;
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
-use mas_data_model::{CompatSession, CompatSsoLogin, Device, User};
+use mas_data_model::{BrowserSession, CompatSession, CompatSsoLogin, Device, User};
 use rand_core::RngCore;
 use ulid::Ulid;
 
@@ -175,6 +175,7 @@ pub trait CompatSessionRepository: Send + Sync {
     /// * `clock`: The clock used to generate timestamps
     /// * `user`: The user to create the compat session for
     /// * `device`: The device ID of this session
+    /// * `browser_session`: The browser session which created this session
     /// * `is_synapse_admin`: Whether the session is a synapse admin session
     ///
     /// # Errors
@@ -186,6 +187,7 @@ pub trait CompatSessionRepository: Send + Sync {
         clock: &dyn Clock,
         user: &User,
         device: Device,
+        browser_session: Option<&BrowserSession>,
         is_synapse_admin: bool,
     ) -> Result<CompatSession, Self::Error>;
 
@@ -261,6 +263,7 @@ repository_impl!(CompatSessionRepository:
         clock: &dyn Clock,
         user: &User,
         device: Device,
+        browser_session: Option<&BrowserSession>,
         is_synapse_admin: bool,
     ) -> Result<CompatSession, Self::Error>;