Fix: Additional fixes for device/login flow session lifecycle coherency issues.

Updated unit tests accordingly.

Author: mikedarcy

credenza/api/session/storage/session_store.py

@@ -149,8 +149,9 @@ def create_session(self,
             session_json = self.crypto_codec.encrypt(session_json)
 
         session_key = access_token if use_access_token_as_session_key else self.generate_session_key()
-        self.map_session(session_key, session_id)
-        self.backend.setex(self._key(session_id), session_json, int(expires_at - now))
+        ttl = int(expires_at - now)
+        self.map_session(session_key, session_id, ttl)
+        self.backend.setex(self._key(session_id), session_json, ttl)
 
         logger.debug(f"Created session {session_id} (realm={realm})")
         return session_key, session_data
@@ -178,18 +179,21 @@ def get_session_by_session_key(self, session_key) -> (str or None, SessionData o
         return session_id, session
 
     def update_session(self, session_id, session_data: SessionData):
-        session_data.updated_at = time.time()
-        session_data.expires_at = session_data.updated_at + self.ttl
+        now = time.time()
+        session_data.updated_at = now
+        if session_data.expires_at < (now + self.ttl):
+            session_data.expires_at = (now + self.ttl)
+        ttl = int(session_data.expires_at - now)
         session_key = self.get_session_key_for_session_id(session_id)
         if self.crypto_codec:
-            session_data = self.crypto_codec.encrypt(json.dumps(session_data.to_dict(), separators=(",", ":")))
+            session_json = self.crypto_codec.encrypt(json.dumps(session_data.to_dict(), separators=(",", ":")))
         else:
-            session_data = json.dumps(session_data.to_dict(), separators=(",", ":"))
-
-        self.map_session(session_key, session_id)
-        self.backend.setex(self._key(session_id), session_data, self.ttl)
+            session_json = json.dumps(session_data.to_dict(), separators=(",", ":"))
+        self.map_session(session_key, session_id, ttl)
+        self.backend.setex(self._key(session_id), session_json, ttl)
 
         logger.debug(f"Updated session {session_id}")
+        return session_key, session_data
 
     def delete_session(self, session_id):
         session = self.get_session_data(session_id)
@@ -233,9 +237,6 @@ def tag_session_metadata(self, session_id: str, metadata: dict, scope: str = "sy
         target = getattr(session.session_metadata, scope)
         target.update(metadata)
 
-        # Update timestamp
-        session.updated_at = time.time()
-
         # Save updated session back as dict
         self.update_session(session_id, session)
         logger.debug(f"Tagged session {session_id} metadata[{scope}]: {metadata}")

credenza/refresh/refresh_worker.py

@@ -46,7 +46,7 @@ def run_refresh_worker(app):
             sub = session.userinfo.get("sub")
             sys_metadata = session.session_metadata.system or {}
             refresh_expires_at = sys_metadata.get("refresh_expires_at")
-            if refresh_expires_at and now > refresh_expires_at:
+            if refresh_expires_at and now > (refresh_expires_at - session_expiry_threshold):
                 audit_event("refresh_expired", session_id=sid)
                 revoke_tokens(sid, session)
                 store.delete_session(sid)
@@ -68,14 +68,8 @@ def run_refresh_worker(app):
             if allow_auto_refresh:
                 modified = bool(refresh_additional_tokens(sid, session)) or modified
 
-            # Just bump session TTL if allowed and it is about to expire and hasn't otherwise been modified
-            if not modified and allow_auto_refresh:
-                ttl = store.get_ttl(sid)
-                if 0 < ttl < session_expiry_threshold:
-                    modified = True
-
             if modified:
                 store.update_session(sid, session)
-                audit_event("device_session_extended", session_id=sid, user=user, sub=sub, realm=realm)
+                audit_event("device_session_updated", session_id=sid, user=user, sub=sub, realm=realm)
 
         time.sleep(interval)

credenza/rest/session.py

@@ -50,13 +50,11 @@ def get_session():
     realm = session.realm
 
     if request.method == "PUT":
-        # Extend session lifetime
-        session.updated_at = now
-        session.expires_at = now + store.ttl
-
         # Enforce max refreshable lifetime
+        session_expiry_threshold = current_app.config.get("SESSION_EXPIRY_THRESHOLD", 300)
         refresh_expires_at = session.session_metadata.system.get("refresh_expires_at")
-        if refresh_expires_at and now > refresh_expires_at:
+        if refresh_expires_at and now > (refresh_expires_at - session_expiry_threshold):
+            revoke_tokens(sid, session)
             store.delete_session(sid)
             audit_event("refresh_expired", session_id=sid)
             abort(401, "Session has expired and can no longer be refreshed")
@@ -70,8 +68,13 @@ def get_session():
             provider = get_augmentation_provider(realm)
             provider.enrich_userinfo(session.userinfo, session.additional_tokens)
 
-        store.update_session(sid, session)
-        audit_event("session_extended", session_id=sid, user=user, sub=sub, realm=realm)
+        skey, session_data = store.update_session(sid, session)
+        audit_event("session_updated",
+                    session_id=sid,
+                    user=user,
+                    sub=sub,
+                    realm=realm,
+                    expires_at=datetime.fromtimestamp(session_data.expires_at, timezone.utc).isoformat())
 
     response = make_session_response(sid, session)
     return make_json_response(response)

test/refresh/test_refresh_worker.py

@@ -175,9 +175,9 @@ def refresh_access_token(self, refresh_token):
             run_refresh_worker(app)
 
     # Verify audit events
-    #   - one session refresh success
+    #   - one session update success
     assert any(
-        ev == "device_session_extended" and kw.get("session_id") == sid
+        ev == "device_session_updated" and kw.get("session_id") == sid
         for ev, kw in audit_calls
     ), f"Missing success event in {audit_calls}"
 
@@ -209,52 +209,6 @@ def refresh_access_token(self, refresh_token):
     # The "fail" block should have been removed entirely
     assert "fail" not in new_sess.additional_tokens
 
-def test_device_session_refresh(app, store, base_session, factory, profiles, audit_calls, frozen_time, monkeypatch):
-    sid = "S3"
-    now = int(frozen_time)
-    app.config["OIDC_CLIENT_FACTORY"] = factory
-    app.config["OIDC_IDP_PROFILES"] = profiles
-
-    # Prepare a device session that has “expired” to force the extension path
-    sess = copy.deepcopy(base_session)
-    sess.session_metadata.system = {
-        "device_session": True,
-        "allow_automatic_refresh": True,
-    }
-    sess.expires_at = now - 1
-
-    # Drive exactly one session with TTL below token_expiry_threshold
-    monkeypatch.setattr(store, "list_session_ids", lambda: [sid])
-    monkeypatch.setattr(store, "get_session_data", lambda s: sess)
-    monkeypatch.setattr(store, "get_ttl", lambda s: 10)
-
-    # Stub out update_session to mirror real behavior
-    updated = []
-    def fake_update_session(session_id, session_data):
-        # frozen_time == time.time() in this test
-        session_data.updated_at = frozen_time
-        session_data.expires_at = session_data.updated_at + store.ttl
-        updated.append((session_id, session_data))
-
-    monkeypatch.setattr(store, "update_session", fake_update_session)
-
-    # Run one iteration (should raise StopIteration at loop end)
-    with app.app_context():
-        with pytest.raises(StopIteration):
-            run_refresh_worker(app)
-
-    # We expect exactly one session-extension audit event
-    assert any(ev == "device_session_extended" for ev, _ in audit_calls), \
-        f"got audit events {audit_calls}"
-
-    # update_session should have been called once for our SID
-    assert len(updated) == 1
-    called_sid, new_sess = updated[0]
-    assert called_sid == sid
-
-    # expires_at must equal frozen_time + store.ttl
-    assert new_sess.expires_at == now + store.ttl
-
 def test_device_access_token_refresh(app,
                                      store,
                                      base_session,
@@ -321,8 +275,8 @@ def refresh_access_token(self, refresh_token):
     events = [ev for ev, _ in audit_calls]
     assert "access_token_refreshed" in events
 
-    # And the worker should then have extended the session
-    assert "device_session_extended" in events
+    # And the worker should then have emitted the device session updated event
+    assert "device_session_updated" in events
 
     # We must have called update_session exactly once
     assert len(updated) == 1

test/rest/test_session.py

@@ -81,19 +81,29 @@ def test_get_session_invalid_bearer(monkeypatch, client):
     assert resp.status_code == 404
 
 def test_put_session_extend(client, store, frozen_time):
-    sid, session = sm.get_current_session()
+    # Capture original values before the PUT
+    before = store.get_session_data("fake_current_sid")
+    old_expires_at = before.expires_at
+
     resp = client.put("/session")
     assert resp.status_code == 200
-    updated = store.get_session_data(sid)
-    # Expires_at should have been set to frozen_time + ttl
-    assert updated.expires_at == frozen_time + store.ttl
-    # updated_at should match frozen_time
-    assert updated.updated_at == frozen_time
 
-def test_put_session_expired(client):
-    # simulate that refresh_expires_at is in the past
+    after = store.get_session_data("fake_current_sid")
+
+    # update_session sets updated_at to time.time() (frozen_time in tests)
+    assert after.updated_at == pytest.approx(frozen_time, abs=1)
+
+    # update_session sets expires_at = max(current_expires_at, now + ttl)
+    expected_expires = max(old_expires_at, frozen_time + store.ttl)
+    assert after.expires_at == pytest.approx(expected_expires, abs=1)
+
+
+def test_put_session_expired(client, app, monkeypatch):
     sid, sess = sm.get_current_session()
-    sess.session_metadata.system["refresh_expires_at"] = int(time.time()) - 1
+    sess.session_metadata.system["refresh_expires_at"] = int(time.time()) + 60
+    monkeypatch.setattr(sm, "revoke_tokens", lambda sid, session: None)
+    with app.app_context():
+        app.config["SESSION_EXPIRY_THRESHOLD"] = 300
     resp = client.put("/session")
     assert resp.status_code == 401
 
@@ -118,11 +128,14 @@ def test_put_session_with_refresh_access_token(client,
     sess.access_token  = "old_access_token"
     sess.id_token      = "old_id_token"
 
+    # Record current expires_at so we can assert the max() behavior
+    original_expires_at = sess.expires_at
+
     # Stub get_current_session -> (sid, sess)
     monkeypatch.setattr(sm,"get_current_session", lambda: (sid, sess))
 
     # Stub out map_session
-    monkeypatch.setattr(store, "map_session", lambda session_key, session_id: "dummy-map-key")
+    monkeypatch.setattr(store, "map_session", lambda session_key, session_id, ttl: "dummy-map-key")
 
     # Configure our dummy OIDC client factory in Flask config
     class DummyClient:
@@ -154,7 +167,7 @@ def get_client(self, realm, native_client=False):
 
     # Audit events
     assert any(ev == "access_token_refreshed" for ev, _ in audit_calls), audit_calls
-    assert any(ev == "session_extended" for ev, _ in audit_calls), audit_calls
+    assert any(ev == "session_updated" for ev, _ in audit_calls), audit_calls
 
     # Fetch the persisted session from the store
     updated = store.get_session_data(sid)
@@ -169,8 +182,10 @@ def get_client(self, realm, native_client=False):
     assert meta["token_expires_at"]   == now + 3600
     assert meta["refresh_expires_at"] == now + 7200
 
-    # Finally, ensure update_session bumped the session TTL
-    assert updated.expires_at == pytest.approx(now + store.ttl, abs=1)
+    # update_session should have bumped updated_at and applied max() for expires_at
+    assert updated.updated_at == pytest.approx(now, abs=1)
+    expected_expires = max(original_expires_at, now + store.ttl)
+    assert updated.expires_at == pytest.approx(expected_expires, abs=1)
 
 
 def test_put_session_with_refresh_access_token_failure(client,
@@ -193,11 +208,13 @@ def test_put_session_with_refresh_access_token_failure(client,
     sess.access_token  = "old_access_token"
     sess.id_token      = "old_id_token"
 
+    original_expires_at = sess.expires_at
+
     # Patch get_current_session
     monkeypatch.setattr(sm, "get_current_session", lambda: (sid, sess))
 
     # Patch map_session
-    monkeypatch.setattr(store, "map_session", lambda session_key, session_id: "dummy-map-key")
+    monkeypatch.setattr(store, "map_session", lambda session_key, session_id, ttl: "dummy-map-key")
 
     # Configure dummy client that fails
     class DummyFailClient:
@@ -219,7 +236,7 @@ def get_client(self, realm, native_client=False):
 
     # Check audit includes the failure
     assert any(ev == "access_token_refresh_failed" for ev, _ in audit_calls), audit_calls
-    assert any(ev == "session_extended" for ev, _ in audit_calls), audit_calls
+    assert any(ev == "session_updated" for ev, _ in audit_calls), audit_calls
 
     # Fetch stored session
     updated = store.get_session_data(sid)
@@ -234,8 +251,10 @@ def get_client(self, realm, native_client=False):
     assert meta["token_expires_at"]   == now - 1
     assert meta["refresh_expires_at"] == now + 600
 
-    # But TTL was still bumped by update_session
-    assert updated.expires_at == pytest.approx(now + store.ttl, abs=1)
+    # update_session should have bumped updated_at and applied max() for expires_at
+    assert updated.updated_at == pytest.approx(now, abs=1)
+    expected_expires = max(original_expires_at, now + store.ttl)
+    assert updated.expires_at == pytest.approx(expected_expires, abs=1)
 
 
 def test_put_session_additional_tokens_refresh(client,
@@ -250,9 +269,6 @@ def test_put_session_additional_tokens_refresh(client,
     threshold = 500
 
     # Prepare a session with four additional token blocks:
-    #    - "good" and "fail" are expired by > threshold
-    #    - "not_due" is far in the future
-    #    - "no_rt" has no refresh_token
     sess = copy.deepcopy(base_session)
     sess.additional_tokens = {
         "good":    {"refresh_token": "rt_good", "expires_at": now - threshold - 1},
@@ -265,14 +281,16 @@ def test_put_session_additional_tokens_refresh(client,
     sess.expires_at = now + 10000
 
     # Stub get_current_session so GET/PUT /session uses our SID and session
-    monkeypatch.setattr(sm,"get_current_session", lambda: (sid, sess))
+    monkeypatch.setattr(sm, "get_current_session", lambda: (sid, sess))
 
-    # Stub out map_session
-    monkeypatch.setattr(store, "map_session", lambda session_key, session_id: "dummy-map-key")
+    monkeypatch.setattr(store, "map_session", lambda session_key, session_id, ttl: "dummy-map-key")
 
-    # Capture calls to update_session
-    updated = []
-    monkeypatch.setattr(store, "update_session", lambda s, sd: updated.append((s, sd)))
+    # Capture calls to update_session, and make it return (session_key, session_data)
+    updated_calls = []
+    def _update_session(s, sd):
+        updated_calls.append((s, sd))
+        return ("dummy-map-key", sd)
+    monkeypatch.setattr(store, "update_session", _update_session)
 
     # Provide a DummyClient that succeeds for "rt_good" and raises for "rt_fail"
     class DummyClient:
@@ -297,7 +315,6 @@ def get_client(self, realm, native_client=False):
 
     with app.app_context():
         app.config["TOKEN_EXPIRY_THRESHOLD"] = 300
-        # Inject our dummy factory
         app.config["OIDC_CLIENT_FACTORY"] = DummyFactory()
 
     # Perform the PUT /session
@@ -310,8 +327,8 @@ def get_client(self, realm, native_client=False):
     assert "additional_token_refresh_failed"  in events, audit_calls
 
     # Only the "good" path should have resulted in update_session
-    assert len(updated) == 1
-    called_sid, new_sess = updated[0]
+    assert len(updated_calls) == 1
+    called_sid, new_sess = updated_calls[0]
     assert called_sid == sid
 
     # Verify the "good" block was updated correctly
@@ -434,4 +451,4 @@ def test_make_session_response_legacy(app, store, base_session):
     assert abs(since_dt.timestamp() - base_session.created_at) < 1
     assert abs(expires_dt.timestamp() - base_session.expires_at) < 1
 
-    assert resp["seconds_remaining"] == store.get_ttl("sid123")
+    assert resp["seconds_remaining"] == store.get_ttl("sid123")