Track user session authenticated through upstream auth sessions

This will help us avoid clearing upstream authorization sessions that
might still be useful to keep around for OIDC Backchannel Logouts

crates/storage-pg/.sqlx/query-5a6b91660e4c12b4a1fe2cad08e727a305cbe4029cd4cebd5ecc274e3e32f533.json

@@ -0,0 +1,16 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n                UPDATE upstream_oauth_authorization_sessions\n                SET consumed_at = $1,\n                    user_session_id = $2\n                WHERE upstream_oauth_authorization_session_id = $3\n            ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Timestamptz",
+        "Uuid",
+        "Uuid"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "5a6b91660e4c12b4a1fe2cad08e727a305cbe4029cd4cebd5ecc274e3e32f533"
+}

crates/storage-pg/.sqlx/query-689ffbfc5137ec788e89062ad679bbe6b23a8861c09a7246dc1659c28f12bf8d.json

@@ -1,15 +0,0 @@
-{
-  "db_name": "PostgreSQL",
-  "query": "\n                UPDATE upstream_oauth_authorization_sessions\n                SET consumed_at = $1\n                WHERE upstream_oauth_authorization_session_id = $2\n            ",
-  "describe": {
-    "columns": [],
-    "parameters": {
-      "Left": [
-        "Timestamptz",
-        "Uuid"
-      ]
-    },
-    "nullable": []
-  },
-  "hash": "689ffbfc5137ec788e89062ad679bbe6b23a8861c09a7246dc1659c28f12bf8d"
-}

crates/storage-pg/migrations/20260121103025_upstream_oauth_track_user_session.sql

@@ -0,0 +1,11 @@
+-- Copyright 2026 Element Creations Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
+-- Please see LICENSE files in the repository root for full details.
+
+-- Start tracking the associated `user_session` directly on the authorization session
+-- This will be backfilled in a separate migration rolling in the next version
+ALTER TABLE upstream_oauth_authorization_sessions
+    ADD COLUMN user_session_id UUID
+    REFERENCES user_sessions (user_session_id)
+    ON DELETE SET NULL;

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

@@ -167,11 +167,24 @@ mod tests {
         assert!(!session.is_consumed());
         assert_eq!(session.link_id(), Some(link.id));
 
-        let session = repo
-            .upstream_oauth_session()
-            .consume(&clock, session)
+        // We need to create a user and start a browser session to consume the session
+        let user = repo
+            .user()
+            .add(&mut rng, &clock, "john".to_owned())
             .await
             .unwrap();
+        let browser_session = repo
+            .browser_session()
+            .add(&mut rng, &clock, &user, None)
+            .await
+            .unwrap();
+
+        let session = repo
+            .upstream_oauth_session()
+            .consume(&clock, session, &browser_session)
+            .await
+            .unwrap();
+
         // Reload the session
         let session = repo
             .upstream_oauth_session()
@@ -181,11 +194,6 @@ mod tests {
             .expect("session to be found in the database");
         assert!(session.is_consumed());
 
-        let user = repo
-            .user()
-            .add(&mut rng, &clock, "john".to_owned())
-            .await
-            .unwrap();
         repo.upstream_oauth_link()
             .associate_to_user(&link, &user)
             .await

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

@@ -8,8 +8,8 @@
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use mas_data_model::{
-    Clock, UpstreamOAuthAuthorizationSession, UpstreamOAuthAuthorizationSessionState,
-    UpstreamOAuthLink, UpstreamOAuthProvider,
+    BrowserSession, Clock, UpstreamOAuthAuthorizationSession,
+    UpstreamOAuthAuthorizationSessionState, UpstreamOAuthLink, UpstreamOAuthProvider,
 };
 use mas_storage::{
     Page, Pagination,
@@ -375,15 +375,18 @@ impl UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'_> {
         &mut self,
         clock: &dyn Clock,
         upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession,
+        browser_session: &BrowserSession,
     ) -> Result<UpstreamOAuthAuthorizationSession, Self::Error> {
         let consumed_at = clock.now();
         sqlx::query!(
             r#"
                 UPDATE upstream_oauth_authorization_sessions
-                SET consumed_at = $1
-                WHERE upstream_oauth_authorization_session_id = $2
+                SET consumed_at = $1,
+                    user_session_id = $2
+                WHERE upstream_oauth_authorization_session_id = $3
             "#,
             consumed_at,
+            Uuid::from(browser_session.id),
             Uuid::from(upstream_oauth_authorization_session.id),
         )
         .traced()