From 225d7b7542ea7886a96fbe539383e01eda10e05c Mon Sep 17 00:00:00 2001
From: Leo Vasanko <leo@git>
Date: Mon, 14 Jul 2025 16:10:02 -0600
Subject: [PATCH] Refactoring reset and session tokens, currently broken.

---
 API.md                                        | 370 ++++++++++++++++++
 frontend/src/App.vue                          |  31 +-
 ...edentialView.vue => AddCredentialView.vue} |  36 +-
 frontend/src/stores/auth.js                   |  11 +-
 frontend/src/utils/passkey.js                 |   3 +
 passkey/fastapi/api_handlers.py               | 321 ++++++++-------
 passkey/fastapi/main.py                       |  88 +----
 passkey/fastapi/reset_handlers.py             | 112 +++---
 .../{session_manager.py => session.py}        |  50 +++
 passkey/fastapi/ws_handlers.py                |  52 ++-
 passkey/util/jwt.py                           |  42 +-
 11 files changed, 786 insertions(+), 330 deletions(-)
 create mode 100644 API.md
 rename frontend/src/components/{AddDeviceCredentialView.vue => AddCredentialView.vue} (60%)
 rename passkey/fastapi/{session_manager.py => session.py} (65%)

diff --git a/API.md b/API.md
new file mode 100644
index 0000000..ed049ab
--- /dev/null
+++ b/API.md
@@ -0,0 +1,370 @@
+# PassKey Auth API Documentation
+
+This document describes all API endpoints available in the PassKey Auth FastAPI application.
+
+## Base URL
+- **Development**: `http://localhost:4401`
+- All API endpoints are prefixed with `/auth`
+
+## Authentication
+The API uses JWT tokens stored in HTTP-only cookies for session management. Some endpoints require authentication via session cookies.
+
+---
+
+## HTTP API Endpoints
+
+### User Management
+
+#### `POST /auth/user-info`
+Get detailed information about the current authenticated user and their credentials.
+
+**Authentication**: Required (session cookie)
+
+**Response**:
+```json
+{
+  "status": "success",
+  "user": {
+    "user_id": "string (UUID)",
+    "user_name": "string",
+    "created_at": "string (ISO 8601)",
+    "last_seen": "string (ISO 8601)",
+    "visits": "number"
+  },
+  "credentials": [
+    {
+      "credential_id": "string (hex)",
+      "aaguid": "string (UUID)",
+      "created_at": "string (ISO 8601)",
+      "last_used": "string (ISO 8601) | null",
+      "last_verified": "string (ISO 8601) | null",
+      "sign_count": "number",
+      "is_current_session": "boolean"
+    }
+  ],
+  "aaguid_info": "object (AAGUID information)"
+}
+```
+
+**Error Response**:
+```json
+{
+  "error": "Not authenticated" | "Failed to get user info: <error_message>"
+}
+```
+
+---
+
+### Session Management
+
+#### `POST /auth/logout`
+Log out the current user by clearing the session cookie.
+
+**Authentication**: Not required
+
+**Response**:
+```json
+{
+  "status": "success",
+  "message": "Logged out successfully"
+}
+```
+
+#### `POST /auth/set-session`
+Set session cookie using JWT token from request body or Authorization header.
+
+**Authentication**: Not required
+
+**Request Body** (alternative to Authorization header):
+```json
+{
+  "token": "string (JWT token)"
+}
+```
+
+**Headers** (alternative to request body):
+```
+Authorization: Bearer <JWT_token>
+```
+
+**Response**:
+```json
+{
+  "status": "success",
+  "message": "Session cookie set successfully",
+  "user_id": "string (UUID)"
+}
+```
+
+**Error Response**:
+```json
+{
+  "error": "No session token provided" | "Invalid or expired session token" | "Failed to set session: <error_message>"
+}
+```
+
+#### `GET /auth/forward-auth`
+Verification endpoint for use with Caddy forward_auth or Nginx auth_request.
+
+**Authentication**: Required (session cookie)
+
+**Success Response**:
+- Status: `204 No Content`
+- Headers: `x-auth-user-id: <user_id>`
+
+**Error Response**:
+- Status: `401 Unauthorized`
+- Returns authentication app HTML page
+- Headers: `www-authenticate: PrivateToken`
+
+---
+
+### Credential Management
+
+#### `POST /auth/delete-credential`
+Delete a specific passkey credential for the current user.
+
+**Authentication**: Required (session cookie)
+
+**Request Body**:
+```json
+{
+  "credential_id": "string (hex-encoded credential ID)"
+}
+```
+
+**Response**:
+```json
+{
+  "status": "success",
+  "message": "Credential deleted successfully"
+}
+```
+
+**Error Response**:
+```json
+{
+  "error": "Not authenticated" | "credential_id is required" | "Invalid credential_id format" | "Credential not found or access denied" | "Cannot delete current session credential" | "Cannot delete last remaining credential" | "Failed to delete credential: <error_message>"
+}
+```
+
+---
+
+### Device Addition
+
+#### `POST /auth/create-device-link`
+Generate a device addition link for authenticated users to add new passkeys to their account.
+
+**Authentication**: Required (session cookie)
+
+**Response**:
+```json
+{
+  "status": "success",
+  "message": "Device addition link generated successfully",
+  "addition_link": "string (URL)",
+  "expires_in_hours": 24
+}
+```
+
+**Error Response**:
+```json
+{
+  "error": "Authentication required" | "Failed to create device addition link: <error_message>"
+}
+```
+
+#### `POST /auth/validate-device-token`
+Validate a device addition token and return associated user information.
+
+**Authentication**: Not required
+
+**Request Body**:
+```json
+{
+  "token": "string (device addition token)"
+}
+```
+
+**Response**:
+```json
+{
+  "status": "success",
+  "valid": true,
+  "user_id": "string (UUID)",
+  "user_name": "string",
+  "token": "string (device addition token)"
+}
+```
+
+**Error Response**:
+```json
+{
+  "error": "Device addition token is required" | "Invalid or expired device addition token" | "Device addition token has expired" | "Failed to validate device addition token: <error_message>"
+}
+```
+
+---
+
+### Static File Serving
+
+#### `GET /auth/{passphrase}`
+Handle passphrase-based authentication redirect with cookie setting.
+
+**Parameters**:
+- `passphrase`: String matching pattern `^\w+(\.\w+){2,}$` (e.g., "word1.word2.word3")
+
+**Response**:
+- Status: `303 See Other`
+- Redirect to: `/`
+- Sets temporary cookie: `auth-token` (expires in 2 seconds)
+
+#### `GET /auth`
+Serve the main authentication app.
+
+**Response**: Returns the main `index.html` file for the authentication SPA.
+
+#### `GET /auth/assets/{path}`
+Serve static assets (CSS, JS, images) for the authentication app.
+
+#### `GET /{path:path}`
+Catch-all route for SPA routing. Serves `index.html` for all non-API routes when requesting HTML content.
+
+**Response**:
+- For HTML requests: Returns `index.html`
+- For non-HTML requests: Returns `404 Not Found` JSON response
+
+---
+
+## WebSocket API Endpoints
+
+All WebSocket endpoints are mounted under `/auth/ws/`.
+
+### Registration
+
+#### `WS /auth/ws/register_new`
+Register a new user with a new passkey credential.
+
+**Flow**:
+1. Client connects to WebSocket
+2. Server sends registration options
+3. Client performs WebAuthn ceremony and sends response
+4. Server validates and creates new user + credential
+5. Server sends JWT token for session establishment
+
+**Server Messages**:
+```json
+// Registration options
+{
+  "rp": { "id": "localhost", "name": "Passkey Auth" },
+  "user": { "id": "base64", "name": "string", "displayName": "string" },
+  "challenge": "base64",
+  "pubKeyCredParams": [...],
+  "timeout": 60000,
+  "attestation": "none",
+  "authenticatorSelection": {...}
+}
+
+// Success response
+{
+  "status": "success",
+  "message": "User registered successfully",
+  "token": "string (JWT)"
+}
+
+// Error response
+{
+  "status": "error",
+  "message": "error description"
+}
+```
+
+#### `WS /auth/ws/add_credential`
+Add a new passkey credential to an existing authenticated user.
+
+**Authentication**: Required (session cookie)
+
+**Flow**:
+1. Client connects with valid session
+2. Server sends registration options for existing user
+3. Client performs WebAuthn ceremony and sends response
+4. Server validates and adds new credential
+5. Server sends success confirmation
+
+#### `WS /auth/ws/add_device_credential`
+Add a new passkey credential using a device addition token.
+
+**Flow**:
+1. Client connects and sends device addition token
+2. Server validates token and sends registration options
+3. Client performs WebAuthn ceremony and sends response
+4. Server validates, adds credential, and cleans up token
+5. Server sends JWT token for session establishment
+
+**Initial Client Message**:
+```json
+{
+  "token": "string (device addition token)"
+}
+```
+
+### Authentication
+
+#### `WS /auth/ws/authenticate`
+Authenticate using existing passkey credentials.
+
+**Flow**:
+1. Client connects to WebSocket
+2. Server sends authentication options
+3. Client performs WebAuthn ceremony and sends response
+4. Server validates credential and updates usage stats
+5. Server sends JWT token for session establishment
+
+**Server Messages**:
+```json
+// Authentication options
+{
+  "challenge": "base64",
+  "timeout": 60000,
+  "rpId": "localhost",
+  "allowCredentials": [...] // Optional, for non-discoverable credentials
+}
+
+// Success response
+{
+  "status": "success",
+  "message": "Authentication successful",
+  "token": "string (JWT)"
+}
+
+// Error response
+{
+  "status": "error",
+  "message": "error description"
+}
+```
+
+---
+
+## Error Handling
+
+All endpoints return consistent error responses:
+
+```json
+{
+  "error": "string (error description)"
+}
+```
+
+## Security Features
+
+- **HTTP-only Cookies**: Session tokens are stored in secure, HTTP-only cookies
+- **CSRF Protection**: SameSite cookie attributes prevent CSRF attacks
+- **Token Validation**: All JWT tokens are validated and automatically refreshed
+- **Credential Isolation**: Users can only access and modify their own credentials
+- **Time-based Expiration**: Device addition tokens expire after 24 hours
+- **Rate Limiting**: WebSocket connections are limited and validated
+
+## CORS and Headers
+
+The application includes appropriate CORS headers and security headers for production use with reverse proxies like Caddy or Nginx.
diff --git a/frontend/src/App.vue b/frontend/src/App.vue
index d2e6af6..2f252b4 100644
--- a/frontend/src/App.vue
+++ b/frontend/src/App.vue
@@ -5,39 +5,36 @@
     <RegisterView v-if="store.currentView === 'register'" />
     <ProfileView v-if="store.currentView === 'profile'" />
     <DeviceLinkView v-if="store.currentView === 'device-link'" />
-    <AddDeviceCredentialView v-if="store.currentView === 'add-device-credential'" />
+    <AddCredentialView v-if="store.currentView === 'add-credential'" />
   </div>
 </template>
 
 <script setup>
-import { onMounted, ref } from 'vue'
+import { onMounted } from 'vue'
 import { useAuthStore } from '@/stores/auth'
 import StatusMessage from '@/components/StatusMessage.vue'
 import LoginView from '@/components/LoginView.vue'
 import RegisterView from '@/components/RegisterView.vue'
 import ProfileView from '@/components/ProfileView.vue'
 import DeviceLinkView from '@/components/DeviceLinkView.vue'
-import AddDeviceCredentialView from '@/components/AddDeviceCredentialView.vue'
-import { getCookie } from './utils/helpers'
+import AddCredentialView from '@/components/AddCredentialView.vue'
 
 const store = useAuthStore()
-let isLoggedIn
 
 onMounted(async () => {
-  if (getCookie('auth-token')) {
-    store.currentView = 'add-device-credential'
-    return
+  // Check for device addition session first
+  try {
+    await store.loadUserInfo()
+  } catch (error) {
+    console.log('Failed to load user info:', error)
+    store.currentView = 'login'
   }
-  isLoggedIn = await store.validateStoredToken()
-  if (isLoggedIn) {
-    // User is logged in, load their data and go to profile
-    try {
-      await store.loadUserInfo()
+  if (store.currentCredentials.length) {
+    // User is logged in, go to profile
       store.currentView = 'profile'
-    } catch (error) {
-      console.error('Failed to load user info:', error)
-      store.currentView = 'login'
-    }
+  } else if (store.currentUser) {
+    // User is logged in via reset link, allow adding a credential
+    store.currentView = 'add-credential'
   } else {
     // User is not logged in, show login
     store.currentView = 'login'
diff --git a/frontend/src/components/AddDeviceCredentialView.vue b/frontend/src/components/AddCredentialView.vue
similarity index 60%
rename from frontend/src/components/AddDeviceCredentialView.vue
rename to frontend/src/components/AddCredentialView.vue
index 89e217c..d339514 100644
--- a/frontend/src/components/AddDeviceCredentialView.vue
+++ b/frontend/src/components/AddCredentialView.vue
@@ -15,30 +15,42 @@
 
 <script setup>
 import { useAuthStore } from '@/stores/auth'
-import { registerWithToken } from '@/utils/passkey'
+import { registerWithSession } from '@/utils/passkey'
 import { ref, onMounted } from 'vue'
-import { getCookie } from '@/utils/helpers'
 
 const authStore = useAuthStore()
-const token = ref(null)
+const hasDeviceSession = ref(false)
 
 // Check existing session on app load
-onMounted(() => {
-  // Check for 'auth-token' cookie
-  token.value = getCookie('auth-token')
-  if (!token.value) {
-    authStore.showMessage('No registration token cookie found.', 'error')
+onMounted(async () => {
+  try {
+    // Check if we have a device addition session
+    const response = await fetch('/auth/device-session-check', {
+      credentials: 'include'
+    })
+    const data = await response.json()
+
+    if (data.device_addition_session) {
+      hasDeviceSession.value = true
+    } else {
+      authStore.showMessage('No device addition session found.', 'error')
+      authStore.currentView = 'login'
+    }
+  } catch (error) {
+    authStore.showMessage('Failed to check device addition session.', 'error')
     authStore.currentView = 'login'
-    return
   }
-  // Delete the cookie
-  document.cookie = 'auth-token=; Max-Age=0; path=/'
 })
 
 function register() {
+  if (!hasDeviceSession.value) {
+    authStore.showMessage('No valid device addition session', 'error')
+    return
+  }
+
   authStore.isLoading = true
   authStore.showMessage('Starting registration...', 'info')
-  registerWithToken(token.value).finally(() => {
+  registerWithSession().finally(() => {
     authStore.isLoading = false
   }).then(() => {
     authStore.showMessage('Passkey registered successfully!', 'success', 2000)
diff --git a/frontend/src/stores/auth.js b/frontend/src/stores/auth.js
index 5f39815..869c717 100644
--- a/frontend/src/stores/auth.js
+++ b/frontend/src/stores/auth.js
@@ -30,15 +30,6 @@ export const useAuthStore = defineStore('auth', {
         }, duration)
       }
     },
-    async validateStoredToken() {
-      try {
-        const response = await fetch('/auth/validate-token')
-        const result = await response.json()
-        return result.status === 'success'
-      } catch (error) {
-        return false
-      }
-    },
     async setSessionCookie(sessionToken) {
       const response = await fetch('/auth/set-session', {
         method: 'POST',
@@ -84,7 +75,7 @@ export const useAuthStore = defineStore('auth', {
       }
     },
     async loadUserInfo() {
-      const response = await fetch('/auth/user-info')
+      const response = await fetch('/auth/user-info', {method: 'POST'})
       const result = await response.json()
       if (result.error) throw new Error(`Server: ${result.error}`)
 
diff --git a/frontend/src/utils/passkey.js b/frontend/src/utils/passkey.js
index d4135d0..503e1ed 100644
--- a/frontend/src/utils/passkey.js
+++ b/frontend/src/utils/passkey.js
@@ -24,6 +24,9 @@ export async function registerCredential() {
 export async function registerWithToken(token) {
   return register('/auth/ws/add_device_credential', { token })
 }
+export async function registerWithSession() {
+  return register('/auth/ws/add_device_credential_session')
+}
 
 export async function authenticateUser() {
   const ws = await aWebSocket('/auth/ws/authenticate')
diff --git a/passkey/fastapi/api_handlers.py b/passkey/fastapi/api_handlers.py
index 787598f..ae66cbf 100644
--- a/passkey/fastapi/api_handlers.py
+++ b/passkey/fastapi/api_handlers.py
@@ -8,12 +8,12 @@ This module contains all the HTTP API endpoints for:
 - Login/logout functionality
 """
 
-from fastapi import Request, Response
+from fastapi import FastAPI, Request, Response
 
 from .. import aaguid
 from ..db import sql
 from ..util.jwt import refresh_session_token, validate_session_token
-from .session_manager import (
+from .session import (
     clear_session_cookie,
     get_current_user,
     get_session_token_from_bearer,
@@ -22,98 +22,169 @@ from .session_manager import (
 )
 
 
-async def get_user_info(request: Request) -> dict:
-    """Get user information and credentials from session cookie."""
-    try:
-        user = await get_current_user(request)
-        if not user:
-            return {"error": "Not authenticated"}
+def register_api_routes(app: FastAPI):
+    """Register all API routes on the FastAPI app."""
 
-        # Get current session credential ID
-        current_credential_id = None
-        session_token = get_session_token_from_cookie(request)
-        if session_token:
+    @app.post("/auth/user-info")
+    async def api_user_info(request: Request, response: Response):
+        """Get user information and credentials from session cookie."""
+        try:
+            user = await get_current_user(request)
+            if not user:
+                return {"error": "Not authenticated"}
+
+            # Get current session credential ID
+            current_credential_id = None
+            session_token = get_session_token_from_cookie(request)
+            if session_token:
+                token_data = validate_session_token(session_token)
+                if token_data:
+                    current_credential_id = token_data.get("credential_id")
+
+            # Get all credentials for the user
+            credential_ids = await sql.get_user_credentials(user.user_id)
+
+            credentials = []
+            user_aaguids = set()
+
+            for cred_id in credential_ids:
+                stored_cred = await sql.get_credential_by_id(cred_id)
+
+                # Convert AAGUID to string format
+                aaguid_str = str(stored_cred.aaguid)
+                user_aaguids.add(aaguid_str)
+
+                # Check if this is the current session credential
+                is_current_session = current_credential_id == stored_cred.credential_id
+
+                credentials.append(
+                    {
+                        "credential_id": stored_cred.credential_id.hex(),
+                        "aaguid": aaguid_str,
+                        "created_at": stored_cred.created_at.isoformat(),
+                        "last_used": stored_cred.last_used.isoformat()
+                        if stored_cred.last_used
+                        else None,
+                        "last_verified": stored_cred.last_verified.isoformat()
+                        if stored_cred.last_verified
+                        else None,
+                        "sign_count": stored_cred.sign_count,
+                        "is_current_session": is_current_session,
+                    }
+                )
+
+            # Get AAGUID information for only the AAGUIDs that the user has
+            aaguid_info = aaguid.filter(user_aaguids)
+
+            # Sort credentials by creation date (earliest first, most recently created last)
+            credentials.sort(key=lambda cred: cred["created_at"])
+
+            return {
+                "status": "success",
+                "user": {
+                    "user_id": str(user.user_id),
+                    "user_name": user.user_name,
+                    "created_at": user.created_at.isoformat()
+                    if user.created_at
+                    else None,
+                    "last_seen": user.last_seen.isoformat() if user.last_seen else None,
+                    "visits": user.visits,
+                },
+                "credentials": credentials,
+                "aaguid_info": aaguid_info,
+            }
+        except Exception as e:
+            return {"error": f"Failed to get user info: {str(e)}"}
+
+    @app.post("/auth/logout")
+    async def api_logout(response: Response):
+        """Log out the current user by clearing the session cookie."""
+        clear_session_cookie(response)
+        return {"status": "success", "message": "Logged out successfully"}
+
+    @app.post("/auth/set-session")
+    async def api_set_session(request: Request, response: Response):
+        """Set session cookie using JWT token from request body or Authorization header."""
+        try:
+            session_token = await get_session_token_from_bearer(request)
+
+            if not session_token:
+                return {"error": "No session token provided"}
+
+            # Validate the session token
             token_data = validate_session_token(session_token)
-            if token_data:
-                current_credential_id = token_data.get("credential_id")
+            if not token_data:
+                return {"error": "Invalid or expired session token"}
 
-        # Get all credentials for the user
-        credential_ids = await sql.get_user_credentials(user.user_id)
+            # Set the HTTP-only cookie
+            set_session_cookie(response, session_token)
 
-        credentials = []
-        user_aaguids = set()
+            return {
+                "status": "success",
+                "message": "Session cookie set successfully",
+                "user_id": str(token_data["user_id"]),
+            }
 
-        for cred_id in credential_ids:
-            stored_cred = await sql.get_credential_by_id(cred_id)
+        except Exception as e:
+            return {"error": f"Failed to set session: {str(e)}"}
 
-            # Convert AAGUID to string format
-            aaguid_str = str(stored_cred.aaguid)
-            user_aaguids.add(aaguid_str)
+    @app.post("/auth/delete-credential")
+    async def api_delete_credential(request: Request):
+        """Delete a specific credential for the current user."""
+        try:
+            user = await get_current_user(request)
+            if not user:
+                return {"error": "Not authenticated"}
+
+            # Get the credential ID from the request body
+            try:
+                body = await request.json()
+                credential_id = body.get("credential_id")
+                if not credential_id:
+                    return {"error": "credential_id is required"}
+            except Exception:
+                return {"error": "Invalid request body"}
+
+            # Convert credential_id from hex string to bytes
+            try:
+                credential_id_bytes = bytes.fromhex(credential_id)
+            except ValueError:
+                return {"error": "Invalid credential_id format"}
+
+            # First, verify the credential belongs to the current user
+            try:
+                stored_cred = await sql.get_credential_by_id(credential_id_bytes)
+                if stored_cred.user_id != user.user_id:
+                    return {"error": "Credential not found or access denied"}
+            except ValueError:
+                return {"error": "Credential not found"}
 
             # Check if this is the current session credential
-            is_current_session = current_credential_id == stored_cred.credential_id
+            session_token = get_session_token_from_cookie(request)
+            if session_token:
+                token_data = validate_session_token(session_token)
+                if (
+                    token_data
+                    and token_data.get("credential_id") == credential_id_bytes
+                ):
+                    return {"error": "Cannot delete current session credential"}
 
-            credentials.append(
-                {
-                    "credential_id": stored_cred.credential_id.hex(),
-                    "aaguid": aaguid_str,
-                    "created_at": stored_cred.created_at.isoformat(),
-                    "last_used": stored_cred.last_used.isoformat()
-                    if stored_cred.last_used
-                    else None,
-                    "last_verified": stored_cred.last_verified.isoformat()
-                    if stored_cred.last_verified
-                    else None,
-                    "sign_count": stored_cred.sign_count,
-                    "is_current_session": is_current_session,
-                }
-            )
+            # Get user's remaining credentials count
+            remaining_credentials = await sql.get_user_credentials(user.user_id)
+            if len(remaining_credentials) <= 1:
+                return {"error": "Cannot delete last remaining credential"}
 
-        # Get AAGUID information for only the AAGUIDs that the user has
-        aaguid_info = aaguid.filter(user_aaguids)
+            # Delete the credential
+            await sql.delete_user_credential(credential_id_bytes)
 
-        # Sort credentials by creation date (earliest first, most recently created last)
-        credentials.sort(key=lambda cred: cred["created_at"])
+            return {"status": "success", "message": "Credential deleted successfully"}
 
-        return {
-            "status": "success",
-            "user": {
-                "user_id": str(user.user_id),
-                "user_name": user.user_name,
-                "created_at": user.created_at.isoformat() if user.created_at else None,
-                "last_seen": user.last_seen.isoformat() if user.last_seen else None,
-                "visits": user.visits,
-            },
-            "credentials": credentials,
-            "aaguid_info": aaguid_info,
-        }
-    except Exception as e:
-        return {"error": f"Failed to get user info: {str(e)}"}
+        except Exception as e:
+            return {"error": f"Failed to delete credential: {str(e)}"}
 
 
-async def refresh_token(request: Request, response: Response) -> dict:
-    """Refresh the session token."""
-    try:
-        session_token = get_session_token_from_cookie(request)
-        if not session_token:
-            return {"error": "No session token found"}
-
-        # Validate and refresh the token
-        new_token = refresh_session_token(session_token)
-
-        if new_token:
-            set_session_cookie(response, new_token)
-            return {"status": "success", "refreshed": True}
-        else:
-            clear_session_cookie(response)
-            return {"error": "Invalid or expired session token"}
-
-    except Exception as e:
-        return {"error": f"Failed to refresh token: {str(e)}"}
-
-
-async def validate_token(request: Request) -> dict:
-    """Validate a session token and return user info."""
+async def validate_token(request: Request, response: Response) -> dict:
+    """Validate a session token and return user info. Also refreshes the token if valid."""
     try:
         session_token = get_session_token_from_cookie(request)
         if not session_token:
@@ -122,11 +193,18 @@ async def validate_token(request: Request) -> dict:
         # Validate the session token
         token_data = validate_session_token(session_token)
         if not token_data:
+            clear_session_cookie(response)
             return {"error": "Invalid or expired session token"}
 
+        # Refresh the token if valid
+        new_token = refresh_session_token(session_token)
+        if new_token:
+            set_session_cookie(response, new_token)
+
         return {
             "status": "success",
             "valid": True,
+            "refreshed": bool(new_token),
             "user_id": str(token_data["user_id"]),
             "credential_id": token_data["credential_id"].hex(),
             "issued_at": token_data["issued_at"],
@@ -135,86 +213,3 @@ async def validate_token(request: Request) -> dict:
 
     except Exception as e:
         return {"error": f"Failed to validate token: {str(e)}"}
-
-
-async def logout(response: Response) -> dict:
-    """Log out the current user by clearing the session cookie."""
-    clear_session_cookie(response)
-    return {"status": "success", "message": "Logged out successfully"}
-
-
-async def set_session(request: Request, response: Response) -> dict:
-    """Set session cookie using JWT token from request body or Authorization header."""
-    try:
-        session_token = await get_session_token_from_bearer(request)
-
-        if not session_token:
-            return {"error": "No session token provided"}
-
-        # Validate the session token
-        token_data = validate_session_token(session_token)
-        if not token_data:
-            return {"error": "Invalid or expired session token"}
-
-        # Set the HTTP-only cookie
-        set_session_cookie(response, session_token)
-
-        return {
-            "status": "success",
-            "message": "Session cookie set successfully",
-            "user_id": str(token_data["user_id"]),
-        }
-
-    except Exception as e:
-        return {"error": f"Failed to set session: {str(e)}"}
-
-
-async def delete_credential(request: Request) -> dict:
-    """Delete a specific credential for the current user."""
-    try:
-        user = await get_current_user(request)
-        if not user:
-            return {"error": "Not authenticated"}
-
-        # Get the credential ID from the request body
-        try:
-            body = await request.json()
-            credential_id = body.get("credential_id")
-            if not credential_id:
-                return {"error": "credential_id is required"}
-        except Exception:
-            return {"error": "Invalid request body"}
-
-        # Convert credential_id from hex string to bytes
-        try:
-            credential_id_bytes = bytes.fromhex(credential_id)
-        except ValueError:
-            return {"error": "Invalid credential_id format"}
-
-        # First, verify the credential belongs to the current user
-        try:
-            stored_cred = await sql.get_credential_by_id(credential_id_bytes)
-            if stored_cred.user_id != user.user_id:
-                return {"error": "Credential not found or access denied"}
-        except ValueError:
-            return {"error": "Credential not found"}
-
-        # Check if this is the current session credential
-        session_token = get_session_token_from_cookie(request)
-        if session_token:
-            token_data = validate_session_token(session_token)
-            if token_data and token_data.get("credential_id") == credential_id_bytes:
-                return {"error": "Cannot delete current session credential"}
-
-        # Get user's remaining credentials count
-        remaining_credentials = await sql.get_user_credentials(user.user_id)
-        if len(remaining_credentials) <= 1:
-            return {"error": "Cannot delete last remaining credential"}
-
-        # Delete the credential
-        await sql.delete_user_credential(credential_id_bytes)
-
-        return {"status": "success", "message": "Credential deleted successfully"}
-
-    except Exception as e:
-        return {"error": f"Failed to delete credential: {str(e)}"}
diff --git a/passkey/fastapi/main.py b/passkey/fastapi/main.py
index 7a9548d..fe1c528 100644
--- a/passkey/fastapi/main.py
+++ b/passkey/fastapi/main.py
@@ -18,25 +18,18 @@ from fastapi import (
     Request,
     Response,
 )
-from fastapi import (
-    Path as FastAPIPath,
-)
 from fastapi.responses import (
     FileResponse,
-    RedirectResponse,
+    JSONResponse,
 )
 from fastapi.staticfiles import StaticFiles
 
 from ..db import sql
 from .api_handlers import (
-    delete_credential,
-    get_user_info,
-    logout,
-    refresh_token,
-    set_session,
+    register_api_routes,
     validate_token,
 )
-from .reset_handlers import create_device_addition_link, validate_device_addition_token
+from .reset_handlers import register_reset_routes
 from .ws_handlers import ws_app
 
 STATIC_DIR = Path(__file__).parent.parent / "frontend-build"
@@ -48,34 +41,23 @@ async def lifespan(app: FastAPI):
     yield
 
 
-app = FastAPI(title="Passkey Auth", lifespan=lifespan)
+app = FastAPI(lifespan=lifespan)
 
 # Mount the WebSocket subapp
 app.mount("/auth/ws", ws_app)
 
-
-@app.get("/auth/user-info")
-async def api_get_user_info(request: Request):
-    """Get user information and credentials from session cookie."""
-    return await get_user_info(request)
-
-
-@app.post("/auth/refresh-token")
-async def api_refresh_token(request: Request, response: Response):
-    """Refresh the session token."""
-    return await refresh_token(request, response)
-
-
-@app.get("/auth/validate-token")
-async def api_validate_token(request: Request):
-    """Validate a session token and return user info."""
-    return await validate_token(request)
+# Register API routes
+register_api_routes(app)
+register_reset_routes(app)
 
 
 @app.get("/auth/forward-auth")
 async def forward_authentication(request: Request):
     """A verification endpoint to use with Caddy forward_auth or Nginx auth_request."""
-    result = await validate_token(request)
+    # Create a dummy response object for internal validation (we won't use it for cookies)
+    response = Response()
+
+    result = await validate_token(request, response)
     if result.get("status") != "success":
         # Serve the index.html of the authentication app if not authenticated
         return FileResponse(
@@ -91,52 +73,6 @@ async def forward_authentication(request: Request):
     )
 
 
-@app.post("/auth/logout")
-async def api_logout(response: Response):
-    """Log out the current user by clearing the session cookie."""
-    return await logout(response)
-
-
-@app.post("/auth/set-session")
-async def api_set_session(request: Request, response: Response):
-    """Set session cookie using JWT token from request body or Authorization header."""
-    return await set_session(request, response)
-
-
-@app.post("/auth/delete-credential")
-async def api_delete_credential(request: Request):
-    """Delete a specific credential for the current user."""
-    return await delete_credential(request)
-
-
-@app.post("/auth/create-device-link")
-async def api_create_device_link(request: Request):
-    """Create a device addition link for the authenticated user."""
-    return await create_device_addition_link(request)
-
-
-@app.post("/auth/validate-device-token")
-async def api_validate_device_token(request: Request):
-    """Validate a device addition token."""
-    return await validate_device_addition_token(request)
-
-
-@app.get("/auth/{passphrase}")
-async def reset_authentication(
-    passphrase: str = FastAPIPath(pattern=r"^\w+(\.\w+){2,}$"),
-):
-    response = RedirectResponse(url="/", status_code=303)
-    response.set_cookie(
-        key="auth-token",
-        value=passphrase,
-        httponly=False,
-        secure=True,
-        samesite="strict",
-        max_age=2,
-    )
-    return response
-
-
 # Serve static files
 app.mount(
     "/auth/assets", StaticFiles(directory=STATIC_DIR / "assets"), name="static assets"
@@ -154,7 +90,7 @@ async def redirect_to_index():
 async def spa_handler(request: Request, path: str):
     """Serve the Vue SPA for all routes (except API and static)"""
     if "text/html" not in request.headers.get("accept", ""):
-        return Response(content="Not Found", status_code=404)
+        return JSONResponse({"error": "Not Found"}, status_code=404)
     return FileResponse(STATIC_DIR / "index.html")
 
 
diff --git a/passkey/fastapi/reset_handlers.py b/passkey/fastapi/reset_handlers.py
index 5317ecf..eb5f88e 100644
--- a/passkey/fastapi/reset_handlers.py
+++ b/passkey/fastapi/reset_handlers.py
@@ -9,73 +9,87 @@ This module provides endpoints for authenticated users to:
 
 from datetime import datetime, timedelta
 
-from fastapi import Request
+from fastapi import FastAPI, Path, Request
+from fastapi.responses import RedirectResponse
 
 from ..db import sql
 from ..util.passphrase import generate
-from .session_manager import get_current_user
+from .session import get_current_user
 
 
-async def create_device_addition_link(request: Request) -> dict:
-    """Create a device addition link for the authenticated user."""
-    try:
-        # Require authentication
-        user = await get_current_user(request)
-        if not user:
-            return {"error": "Authentication required"}
+def register_reset_routes(app: FastAPI):
+    """Register all device addition/reset routes on the FastAPI app."""
 
-        # Generate a human-readable token
-        token = generate(n=4, sep=".")  # e.g., "able-ocean-forest-dawn"
+    @app.post("/auth/create-device-link")
+    async def api_create_device_link(request: Request):
+        """Create a device addition link for the authenticated user."""
+        try:
+            # Require authentication
+            user = await get_current_user(request)
+            if not user:
+                return {"error": "Authentication required"}
 
-        # Create reset token in database
-        await sql.create_reset_token(user.user_id, token)
+            # Generate a human-readable token
+            token = generate(n=4, sep=".")  # e.g., "able-ocean-forest-dawn"
 
-        # Generate the device addition link with pretty URL
-        addition_link = f"{request.headers.get('origin', '')}/auth/{token}"
+            # Create reset token in database
+            await sql.create_reset_token(user.user_id, token)
 
-        return {
-            "status": "success",
-            "message": "Device addition link generated successfully",
-            "addition_link": addition_link,
-            "expires_in_hours": 24,
-        }
+            # Generate the device addition link with pretty URL
+            addition_link = f"{request.headers.get('origin', '')}/auth/{token}"
 
-    except Exception as e:
-        return {"error": f"Failed to create device addition link: {str(e)}"}
+            return {
+                "status": "success",
+                "message": "Device addition link generated successfully",
+                "addition_link": addition_link,
+                "expires_in_hours": 24,
+            }
 
+        except Exception as e:
+            return {"error": f"Failed to create device addition link: {str(e)}"}
 
-async def validate_device_addition_token(request: Request) -> dict:
-    """Validate a device addition token and return user info."""
-    try:
-        body = await request.json()
-        token = body.get("token")
+    @app.get("/auth/device-session-check")
+    async def check_device_session(request: Request):
+        """Check if the current session is for device addition."""
+        from .session import is_device_addition_session
 
-        if not token:
-            return {"error": "Device addition token is required"}
+        is_device_session = await is_device_addition_session(request)
+        return {"device_addition_session": is_device_session}
 
-        # Get reset token
-        reset_token = await sql.get_reset_token(token)
-        if not reset_token:
-            return {"error": "Invalid or expired device addition token"}
+    @app.get("/auth/{passphrase}")
+    async def reset_authentication(
+        passphrase: str = Path(pattern=r"^\w+(\.\w+){2,}$"),
+    ):
+        try:
+            # Get reset token to validate it exists and get user_id
+            reset_token = await sql.get_reset_token(passphrase)
+            if not reset_token:
+                # Token doesn't exist, redirect to home
+                return RedirectResponse(url="/", status_code=303)
 
-        # Check if token is expired (24 hours)
-        expiry_time = reset_token.created_at + timedelta(hours=24)
-        if datetime.now() > expiry_time:
-            return {"error": "Device addition token has expired"}
+            # Check if token is expired (24 hours)
+            expiry_time = reset_token.created_at + timedelta(hours=24)
+            if datetime.now() > expiry_time:
+                # Token expired, clean it up and redirect to home
+                await sql.delete_reset_token(passphrase)
+                return RedirectResponse(url="/", status_code=303)
 
-        # Get user info
-        user = await sql.get_user_by_id(reset_token.user_id)
+            # Create a device addition session token for the user
+            from ..util.jwt import create_device_addition_token
 
-        return {
-            "status": "success",
-            "valid": True,
-            "user_id": str(user.user_id),
-            "user_name": user.user_name,
-            "token": token,
-        }
+            session_token = create_device_addition_token(reset_token.user_id)
 
-    except Exception as e:
-        return {"error": f"Failed to validate device addition token: {str(e)}"}
+            # Create response and set session cookie
+            response = RedirectResponse(url="/auth/", status_code=303)
+            from .session import set_session_cookie
+
+            set_session_cookie(response, session_token)
+
+            return response
+
+        except Exception:
+            # On any error, redirect to home
+            return RedirectResponse(url="/", status_code=303)
 
 
 async def use_device_addition_token(token: str) -> dict:
diff --git a/passkey/fastapi/session_manager.py b/passkey/fastapi/session.py
similarity index 65%
rename from passkey/fastapi/session_manager.py
rename to passkey/fastapi/session.py
index be94a3c..570ead9 100644
--- a/passkey/fastapi/session_manager.py
+++ b/passkey/fastapi/session.py
@@ -96,3 +96,53 @@ async def get_user_from_cookie_string(cookie_header: str) -> UUID | None:
         return None
 
     return token_data["user_id"]
+
+
+async def is_device_addition_session(request: Request) -> bool:
+    """Check if the current session is for device addition."""
+    session_token = request.cookies.get(COOKIE_NAME)
+    if not session_token:
+        return False
+
+    token_data = validate_session_token(session_token)
+    if not token_data:
+        return False
+
+    return token_data.get("device_addition", False)
+
+
+async def get_device_addition_user_id(request: Request) -> UUID | None:
+    """Get user ID from device addition session."""
+    session_token = request.cookies.get(COOKIE_NAME)
+    if not session_token:
+        return None
+
+    token_data = validate_session_token(session_token)
+    if not token_data or not token_data.get("device_addition"):
+        return None
+
+    return token_data.get("user_id")
+
+
+async def get_device_addition_user_id_from_cookie(cookie_header: str) -> UUID | None:
+    """Parse cookie header and return user ID if valid device addition session exists."""
+    if not cookie_header:
+        return None
+
+    # Parse cookies from header (simple implementation)
+    cookies = {}
+    for cookie in cookie_header.split(";"):
+        cookie = cookie.strip()
+        if "=" in cookie:
+            name, value = cookie.split("=", 1)
+            cookies[name] = value
+
+    session_token = cookies.get(COOKIE_NAME)
+    if not session_token:
+        return None
+
+    token_data = validate_session_token(session_token)
+    if not token_data or not token_data.get("device_addition"):
+        return None
+
+    return token_data["user_id"]
diff --git a/passkey/fastapi/ws_handlers.py b/passkey/fastapi/ws_handlers.py
index 17eabb7..1422b65 100644
--- a/passkey/fastapi/ws_handlers.py
+++ b/passkey/fastapi/ws_handlers.py
@@ -20,7 +20,7 @@ from ..db import sql
 from ..db.sql import User
 from ..sansio import Passkey
 from ..util.jwt import create_session_token
-from .session_manager import get_user_from_cookie_string
+from .session import get_user_from_cookie_string
 
 # Create a FastAPI subapp for WebSocket endpoints
 ws_app = FastAPI()
@@ -181,6 +181,56 @@ async def websocket_add_device_credential(ws: WebSocket, token: str):
         await ws.send_json({"error": "Internal Server Error"})
 
 
+@ws_app.websocket("/add_device_credential_session")
+async def websocket_add_device_credential_session(ws: WebSocket):
+    """Add a new credential for an existing user via device addition session."""
+    await ws.accept()
+    origin = ws.headers.get("origin")
+    try:
+        # Get device addition user ID from session cookie
+        cookie_header = ws.headers.get("cookie", "")
+        from .session import get_device_addition_user_id_from_cookie
+
+        user_id = await get_device_addition_user_id_from_cookie(cookie_header)
+
+        if not user_id:
+            await ws.send_json({"error": "No valid device addition session found"})
+            return
+
+        # Get user information
+        user = await sql.get_user_by_id(user_id)
+        if not user:
+            await ws.send_json({"error": "User not found"})
+            return
+
+        # WebAuthn registration
+        # Fetch challenge IDs for the user
+        challenge_ids = await sql.get_user_credentials(user_id)
+
+        credential = await register_chat(
+            ws, user_id, user.user_name, challenge_ids, origin=origin
+        )
+
+        # Store the new credential in the database
+        await sql.create_credential_for_user(credential)
+
+        await ws.send_json(
+            {
+                "status": "success",
+                "user_id": str(user_id),
+                "credential_id": credential.credential_id.hex(),
+                "message": "New credential added successfully via device addition session",
+            }
+        )
+    except ValueError as e:
+        await ws.send_json({"error": str(e)})
+    except WebSocketDisconnect:
+        pass
+    except Exception:
+        logging.exception("Internal Server Error")
+        await ws.send_json({"error": "Internal Server Error"})
+
+
 @ws_app.websocket("/authenticate")
 async def websocket_authenticate(ws: WebSocket):
     await ws.accept()
diff --git a/passkey/util/jwt.py b/passkey/util/jwt.py
index 40eb0a8..4bfac25 100644
--- a/passkey/util/jwt.py
+++ b/passkey/util/jwt.py
@@ -58,6 +58,28 @@ class JWTManager:
 
         return jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
 
+    def create_token_without_credential(self, user_id: UUID) -> str:
+        """
+        Create a JWT token for device addition (without credential ID).
+
+        Args:
+            user_id: The user's UUID
+
+        Returns:
+            JWT token string for device addition
+        """
+        now = datetime.now()
+        payload = {
+            "user_id": str(user_id),
+            "credential_id": None,  # No credential for device addition
+            "device_addition": True,  # Flag to indicate this is for device addition
+            "iat": now,
+            "exp": now + timedelta(hours=2),  # Shorter expiry for device addition
+            "iss": "passkeyauth",
+        }
+
+        return jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
+
     def validate_token(self, token: str) -> Optional[dict]:
         """
         Validate a JWT token and return the payload.
@@ -76,12 +98,23 @@ class JWTManager:
                 issuer="passkeyauth",
             )
 
-            return {
+            result = {
                 "user_id": UUID(payload["user_id"]),
-                "credential_id": bytes.fromhex(payload["credential_id"]),
                 "issued_at": payload["iat"],
                 "expires_at": payload["exp"],
             }
+
+            # Handle credential_id for regular tokens vs device addition tokens
+            if payload.get("credential_id") is not None:
+                result["credential_id"] = bytes.fromhex(payload["credential_id"])
+            else:
+                result["credential_id"] = None
+
+            # Add device addition flag if present
+            if payload.get("device_addition"):
+                result["device_addition"] = True
+
+            return result
         except jwt.ExpiredSignatureError:
             return None
         except jwt.InvalidTokenError:
@@ -122,6 +155,11 @@ def create_session_token(user_id: UUID, credential_id: bytes) -> str:
     return get_jwt_manager().create_token(user_id, credential_id)
 
 
+def create_device_addition_token(user_id: UUID) -> str:
+    """Create a token for device addition."""
+    return get_jwt_manager().create_token_without_credential(user_id)
+
+
 def validate_session_token(token: str) -> Optional[dict]:
     """Validate a session token."""
     return get_jwt_manager().validate_token(token)