Keep Xaman callback state retry-safe

backend/app/main.py

@@ -45,6 +45,7 @@
     create_pending_login_state,
     fail_origin_login_handoff,
     get_or_create_user_from_external_identity,
+    restore_pending_login_state_after_transient_failure,
     validate_pending_login_state,
     validate_origin_login_handoff,
 )
@@ -759,9 +760,9 @@ def identity_callback(
     """
     Browser callback contract: code + state only.
 
-    Browser callback contract remains code + state only.
-    Backend validates pending state, exchanges code with WordPress bridge,
-    resolves/creates CalorieApp user identity, and issues CalorieApp session cookie.
+    Backend atomically reserves the pending state, exchanges code with the
+    WordPress bridge, restores the state only for transient bridge failures,
+    resolves/creates CalorieApp user identity, and issues the session cookie.
     """
     code = payload.code.strip()
     state = payload.state.strip()
@@ -780,7 +781,19 @@ def identity_callback(
 
     try:
         claims = _exchange_code_for_claims(code=code, state=state)
+    except HTTPException as exc:
+        if exc.status_code in {502, 503, 504}:
+            restored = restore_pending_login_state_after_transient_failure(session, state)
+            if not restored:
+                fail_origin_login_handoff(session, state)
+            raise
+        fail_origin_login_handoff(session, state)
+        raise
+    except Exception:
+        fail_origin_login_handoff(session, state)
+        raise
 
+    try:
         user, created = get_or_create_user_from_external_identity(
             session=session,
             provider=_IDENTITY_PROVIDER,

backend/app/services/identity.py

@@ -303,6 +303,23 @@ def consume_pending_login_state(
     return False, reason
 
 
+def restore_pending_login_state_after_transient_failure(
+    session: Session,
+    state: str,
+) -> bool:
+    """Restore a consumed, unexpired state after a retryable bridge failure."""
+    now = utc_now()
+    updated = session.exec(
+        update(PendingLoginStateDB)
+        .where(PendingLoginStateDB.state_hash == hash_login_state(state))
+        .where(PendingLoginStateDB.status == "consumed")
+        .where(PendingLoginStateDB.expires_at >= now)
+        .values(status="pending", consumed_at=None)
+    )
+    session.commit()
+    return updated.rowcount == 1
+
+
 def cleanup_pending_login_states(session: Session) -> None:
     """Delete expired pending login states opportunistically."""
     now = utc_now()

backend/tests/test_identity_endpoints.py

@@ -1242,7 +1242,7 @@ def test_origin_handoff_is_single_use_across_browsers(
             assert replay_claim.status_code == 409
             assert replay_browser.get("/api/identity/me").status_code == 401
 
-    def test_failed_callback_marks_origin_handoff_failed(
+    def test_transient_bridge_failure_keeps_origin_handoff_pending(
         self,
         client: TestClient,
         monkeypatch: pytest.MonkeyPatch,
@@ -1270,7 +1270,43 @@ def _raise_exchange_failure(code: str, state: str):
 
         assert callback.status_code == 502
         assert status.status_code == 200
+        assert status.json()["status"] == "pending"
+
+    def test_non_retryable_bridge_failure_marks_origin_handoff_failed(
+        self,
+        client: TestClient,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        def reject_exchange(code: str, state: str):
+            raise main_module.HTTPException(
+                status_code=400,
+                detail="Authorization code exchange rejected",
+            )
+
+        monkeypatch.setattr(main_module, "_exchange_code_for_claims", reject_exchange)
+        start = client.post("/api/identity/login/start").json()
+
+        callback = client.post(
+            "/api/identity/callback",
+            json={"code": "rejected-code", "state": start["state"]},
+        )
+        status = client.post(
+            "/api/identity/login/status",
+            json={
+                "state": start["state"],
+                "browser_handoff_token": start["browser_handoff_token"],
+            },
+        )
+        replay = client.post(
+            "/api/identity/callback",
+            json={"code": "rejected-code", "state": start["state"]},
+        )
+
+        assert callback.status_code == 400
+        assert status.status_code == 200
         assert status.json()["status"] == "failed"
+        assert replay.status_code == 400
+        assert "already consumed" in replay.json()["detail"]
 
     def test_concurrent_callback_state_use_allows_only_one_success(self, monkeypatch: pytest.MonkeyPatch):
         from concurrent.futures import ThreadPoolExecutor
@@ -1293,20 +1329,44 @@ def call_callback(code_value: str) -> int:
 
         assert sorted(statuses) == [200, 400]
 
-    def test_bridge_exchange_failure_consumes_state_and_retry_fails(self, client: TestClient, monkeypatch: pytest.MonkeyPatch):
-        def _raise_exchange_failure(code: str, state: str):
-            raise main_module.HTTPException(status_code=502, detail="WordPress bridge exchange failed")
+    def test_bridge_exchange_failure_keeps_state_retryable_then_blocks_replay(
+        self,
+        client: TestClient,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        calls = {"count": 0}
+
+        def flaky_exchange(code: str, state: str):
+            calls["count"] += 1
+            if calls["count"] == 1:
+                raise main_module.HTTPException(
+                    status_code=502,
+                    detail="WordPress bridge exchange failed",
+                )
+            return self._stub_claims()
 
-        monkeypatch.setattr(main_module, "_exchange_code_for_claims", _raise_exchange_failure)
+        monkeypatch.setattr(main_module, "_exchange_code_for_claims", flaky_exchange)
 
         state = client.post("/api/identity/login/start").json()["state"]
 
-        failed = client.post("/api/identity/callback", json={"code": "bridge-code", "state": state})
+        failed = client.post(
+            "/api/identity/callback",
+            json={"code": "bridge-code", "state": state},
+        )
         assert failed.status_code == 502
 
-        retried = client.post("/api/identity/callback", json={"code": "bridge-code", "state": state})
-        assert retried.status_code == 400
-        assert "already consumed" in retried.json()["detail"]
+        retried = client.post(
+            "/api/identity/callback",
+            json={"code": "bridge-code", "state": state},
+        )
+        assert retried.status_code == 200
+
+        replay = client.post(
+            "/api/identity/callback",
+            json={"code": "bridge-code", "state": state},
+        )
+        assert replay.status_code == 400
+        assert "already consumed" in replay.json()["detail"]
 
     def test_state_substitution_fails(self, client: TestClient, monkeypatch: pytest.MonkeyPatch):
         monkeypatch.setattr(main_module, "_exchange_code_for_claims", lambda code, state: self._stub_claims())