@mcp-ts/sdk 2.3.3 → 2.4.0

package/src/server/storage/neon-backend.ts

@@ -1,11 +1,14 @@
-import type { SessionStore, Session } from './types.js';
-import { SESSION_TTL_SECONDS } from '../../shared/constants.js';
+import type { SessionStore, Session, SessionCredentials } from './types.js';
+import type { SessionStatus } from './types.js';
+import { DORMANT_SESSION_EXPIRATION_MS } from '../../shared/constants.js';
 import { generateSessionId } from '../../shared/utils.js';
 import { encryptObject, decryptObject } from './crypto.js';
+import { resolveSessionExpiresAt } from './session-lifecycle.js';
 
 export interface NeonStorageOptions {
     schema?: string;
     table?: string;
+    credentialsTable?: string;
 }
 
 type NeonSql = {
@@ -13,6 +16,7 @@ type NeonSql = {
 };
 
 type NeonSessionRow = {
+    id?: string;
     session_id: string;
     server_id?: string | null;
     server_name?: string | null;
@@ -20,18 +24,27 @@ type NeonSessionRow = {
     transport_type: 'sse' | 'streamable-http';
     callback_url: string;
     created_at: string | Date;
+    updated_at?: string | Date | null;
+    expires_at?: string | Date | null;
     user_id: string;
     headers?: unknown;
-    active?: boolean | null;
+    auth_url?: string | null;
+    status?: SessionStatus | null;
+};
+
+type NeonCredentialsRow = {
+    session_id: string;
+    user_id: string;
     client_information?: unknown;
     tokens?: unknown;
-    code_verifier?: string | null;
+    code_verifier?: unknown;
     client_id?: string | null;
+    oauth_state?: unknown;
 };
 
 export class NeonStorageBackend implements SessionStore {
-    private readonly DEFAULT_TTL = SESSION_TTL_SECONDS;
     private readonly tableName: string;
+    private readonly credentialsTableName: string;
 
     constructor(
         private readonly sql: NeonSql,
@@ -39,23 +52,29 @@ export class NeonStorageBackend implements SessionStore {
     ) {
         const schema = options.schema || 'public';
         const table = options.table || 'mcp_sessions';
+        const credentialsTable = options.credentialsTable || 'mcp_credentials';
         this.tableName = `${this.quoteIdentifier(schema)}.${this.quoteIdentifier(table)}`;
+        this.credentialsTableName = `${this.quoteIdentifier(schema)}.${this.quoteIdentifier(credentialsTable)}`;
     }
 
     async init(): Promise<void> {
+        await this.assertTable(this.tableName, 'mcp_sessions');
+        await this.assertTable(this.credentialsTableName, 'mcp_credentials');
+        console.log('[mcp-ts][Storage] Neon: storage tables verified.');
+    }
+
+    private async assertTable(qualifiedName: string, displayName: string): Promise<void> {
         const [{ exists } = { exists: null }] = await this.sql.query(
             'SELECT to_regclass($1) AS exists',
-            [this.tableName.replace(/"/g, '')]
+            [qualifiedName.replace(/"/g, '')]
         ) as Array<{ exists: string | null }>;
 
         if (!exists) {
             throw new Error(
-                '[NeonStorage] Table "mcp_sessions" not found in your database. ' +
+                `[NeonStorage] Table "${displayName}" not found in your database. ` +
                 'Please create it using the Neon storage guide in docs/storage-backends/neon.md.'
             );
         }
-
-        console.log('[mcp-ts][Storage] Neon: "mcp_sessions" table verified.');
     }
 
     generateSessionId(): string {
@@ -78,22 +97,45 @@ export class NeonStorageBackend implements SessionStore {
             transportType: row.transport_type,
             callbackUrl: row.callback_url,
             createdAt: new Date(row.created_at).getTime(),
+            updatedAt: new Date(row.updated_at ?? row.created_at).getTime(),
+            expiresAt: row.expires_at ? new Date(row.expires_at).getTime() : null,
             userId: row.user_id,
             headers: decryptObject(row.headers),
-            active: row.active ?? false,
-            clientInformation: row.client_information as Session['clientInformation'],
+            authUrl: row.auth_url ?? undefined,
+            status: row.status ?? 'pending',
+        };
+    }
+
+    private mapRowToCredentials(row: NeonCredentialsRow, userId: string, sessionId: string): SessionCredentials {
+        return {
+            sessionId,
+            userId,
+            clientInformation: decryptObject(row.client_information),
             tokens: decryptObject(row.tokens),
-            codeVerifier: row.code_verifier ?? undefined,
+            codeVerifier: decryptObject(row.code_verifier),
             clientId: row.client_id ?? undefined,
+            oauthState: row.oauth_state as SessionCredentials['oauthState'],
         };
     }
 
-    async create(session: Session, ttl?: number): Promise<void> {
+    private hasCredentialData(data: Partial<SessionCredentials>): boolean {
+        return (
+            'clientInformation' in data ||
+            'tokens' in data ||
+            'codeVerifier' in data ||
+            'clientId' in data ||
+            'oauthState' in data
+        );
+    }
+
+    async create(session: Session): Promise<void> {
         const { sessionId, userId } = session;
         if (!sessionId || !userId) throw new Error('userId and sessionId required');
 
-        const effectiveTtl = ttl ?? this.DEFAULT_TTL;
-        const expiresAt = new Date(Date.now() + effectiveTtl * 1000).toISOString();
+        const status = session.status ?? 'pending';
+        const createdAt = new Date(session.createdAt || Date.now()).toISOString();
+        const updatedAt = new Date(session.updatedAt ?? session.createdAt ?? Date.now()).toISOString();
+        const expiresAt = resolveSessionExpiresAt(status, new Date(createdAt).getTime());
 
         try {
             await this.sql.query(
@@ -106,16 +148,14 @@ export class NeonStorageBackend implements SessionStore {
                     transport_type,
                     callback_url,
                     created_at,
+                    updated_at,
                     headers,
-                    active,
-                    client_information,
-                    tokens,
-                    code_verifier,
-                    client_id,
+                    auth_url,
+                    status,
                     expires_at
                 ) VALUES (
                     $1, $2, $3, $4, $5, $6, $7, $8,
-                    $9, $10, $11, $12, $13, $14, $15
+                    $9, $10, $11, $12, $13
                 )`,
                 [
                     sessionId,
@@ -125,14 +165,12 @@ export class NeonStorageBackend implements SessionStore {
                     session.serverUrl,
                     session.transportType,
                     session.callbackUrl,
-                    new Date(session.createdAt || Date.now()).toISOString(),
+                    createdAt,
+                    updatedAt,
                     encryptObject(session.headers),
-                    session.active ?? false,
-                    session.clientInformation,
-                    encryptObject(session.tokens),
-                    session.codeVerifier,
-                    session.clientId,
-                    expiresAt,
+                    session.authUrl ?? null,
+                    status,
+                    expiresAt === null ? null : new Date(expiresAt).toISOString(),
                 ]
             );
         } catch (error: any) {
@@ -141,57 +179,106 @@ export class NeonStorageBackend implements SessionStore {
             }
             throw new Error(`Failed to create session in Neon: ${error.message}`);
         }
+
     }
 
-    async update(userId: string, sessionId: string, data: Partial<Session>, ttl?: number): Promise<void> {
+    async update(userId: string, sessionId: string, data: Partial<Session>): Promise<void> {
         const currentSession = await this.get(userId, sessionId);
         if (!currentSession) {
             throw new Error(`Session ${sessionId} not found for userId ${userId}`);
         }
 
         const updatedSession = { ...currentSession, ...data };
-        const effectiveTtl = ttl ?? this.DEFAULT_TTL;
-        const expiresAt = new Date(Date.now() + effectiveTtl * 1000).toISOString();
-
-        const updatedRows = await this.sql.query(
-            `UPDATE ${this.tableName}
-             SET
-                server_id = $1,
-                server_name = $2,
-                server_url = $3,
-                transport_type = $4,
-                callback_url = $5,
-                active = $6,
-                headers = $7,
-                client_information = $8,
-                tokens = $9,
-                code_verifier = $10,
-                client_id = $11,
-                expires_at = $12,
-                updated_at = now()
-             WHERE user_id = $13 AND session_id = $14
-             RETURNING id`,
+        const status = updatedSession.status ?? 'pending';
+        const expiresAt = resolveSessionExpiresAt(status);
+
+        const shouldUpdateSession = (
+            'serverId' in data ||
+            'serverName' in data ||
+            'serverUrl' in data ||
+            'transportType' in data ||
+            'callbackUrl' in data ||
+            'status' in data ||
+            'headers' in data ||
+            'authUrl' in data
+        );
+
+        if (shouldUpdateSession) {
+            const updatedRows = await this.sql.query(
+                `UPDATE ${this.tableName}
+                 SET
+                    server_id = $1,
+                    server_name = $2,
+                    server_url = $3,
+                    transport_type = $4,
+                    callback_url = $5,
+                    status = $6,
+                    headers = $7,
+                    auth_url = $8,
+                    expires_at = $9,
+                    updated_at = now()
+                 WHERE user_id = $10 AND session_id = $11
+                 RETURNING id`,
+                [
+                    updatedSession.serverId,
+                    updatedSession.serverName,
+                    updatedSession.serverUrl,
+                    updatedSession.transportType,
+                    updatedSession.callbackUrl,
+                    status,
+                    encryptObject(updatedSession.headers),
+                    updatedSession.authUrl ?? null,
+                    expiresAt === null ? null : new Date(expiresAt).toISOString(),
+                    userId,
+                    sessionId,
+                ]
+            ) as Array<{ id: string }>;
+
+            if (updatedRows.length === 0) {
+                throw new Error(`Session ${sessionId} not found for userId ${userId}`);
+            }
+        }
+
+    }
+
+    async patchCredentials(userId: string, sessionId: string, data: Partial<SessionCredentials>): Promise<void> {
+        if (!this.hasCredentialData(data)) return;
+
+        await this.sql.query(
+            `INSERT INTO ${this.credentialsTableName} (
+                user_id,
+                session_id,
+                client_information,
+                tokens,
+                code_verifier,
+                client_id,
+                oauth_state,
+                updated_at
+            ) VALUES ($1, $2, $3, $4, $5, $6, $7, now())
+            ON CONFLICT (user_id, session_id)
+            DO UPDATE SET
+                client_information = CASE WHEN $8 THEN EXCLUDED.client_information ELSE ${this.credentialsTableName}.client_information END,
+                tokens = CASE WHEN $9 THEN EXCLUDED.tokens ELSE ${this.credentialsTableName}.tokens END,
+                code_verifier = CASE WHEN $10 THEN EXCLUDED.code_verifier ELSE ${this.credentialsTableName}.code_verifier END,
+                client_id = CASE WHEN $11 THEN EXCLUDED.client_id ELSE ${this.credentialsTableName}.client_id END,
+                oauth_state = CASE WHEN $12 THEN EXCLUDED.oauth_state ELSE ${this.credentialsTableName}.oauth_state END,
+                updated_at = now()`,
             [
-                updatedSession.serverId,
-                updatedSession.serverName,
-                updatedSession.serverUrl,
-                updatedSession.transportType,
-                updatedSession.callbackUrl,
-                updatedSession.active ?? false,
-                encryptObject(updatedSession.headers),
-                updatedSession.clientInformation,
-                encryptObject(updatedSession.tokens),
-                updatedSession.codeVerifier,
-                updatedSession.clientId,
-                expiresAt,
                 userId,
                 sessionId,
+                'clientInformation' in data ? (data.clientInformation == null ? null : encryptObject(data.clientInformation)) : null,
+                'tokens' in data ? (data.tokens == null ? null : encryptObject(data.tokens)) : null,
+                'codeVerifier' in data ? (data.codeVerifier == null ? null : encryptObject(data.codeVerifier)) : null,
+                'clientId' in data ? (data.clientId ?? null) : null,
+                'oauthState' in data ? (data.oauthState ?? null) : null,
+                'clientInformation' in data,
+                'tokens' in data,
+                'codeVerifier' in data,
+                'clientId' in data,
+                'oauthState' in data,
             ]
-        ) as Array<{ id: string }>;
+        );
 
-        if (updatedRows.length === 0) {
-            throw new Error(`Session ${sessionId} not found for userId ${userId}`);
-        }
     }
 
     async get(userId: string, sessionId: string): Promise<Session | null> {
@@ -200,13 +287,39 @@ export class NeonStorageBackend implements SessionStore {
                 `SELECT * FROM ${this.tableName} WHERE user_id = $1 AND session_id = $2`,
                 [userId, sessionId]
             ) as NeonSessionRow[];
-            return rows[0] ? this.mapRowToSessionData(rows[0]) : null;
+
+            if (!rows[0]) return null;
+
+            return this.mapRowToSessionData(rows[0]);
         } catch (error) {
             console.error('[NeonStorage] Failed to get session:', error);
             return null;
         }
     }
 
+    async getCredentials(userId: string, sessionId: string): Promise<SessionCredentials | null> {
+        try {
+            const credentialRows = await this.sql.query(
+                `SELECT * FROM ${this.credentialsTableName} WHERE user_id = $1 AND session_id = $2`,
+                [userId, sessionId]
+            ) as NeonCredentialsRow[];
+
+            if (credentialRows[0]) {
+                return this.mapRowToCredentials(credentialRows[0], userId, sessionId);
+            }
+
+            const sessionRows = await this.sql.query(
+                `SELECT id FROM ${this.tableName} WHERE user_id = $1 AND session_id = $2`,
+                [userId, sessionId]
+            ) as Array<{ id: string }>;
+
+            return sessionRows[0] ? { sessionId, userId } : null;
+        } catch (error) {
+            console.error('[NeonStorage] Failed to get credentials:', error);
+            return null;
+        }
+    }
+
     async list(userId: string): Promise<Session[]> {
         try {
             const rows = await this.sql.query(
@@ -220,6 +333,17 @@ export class NeonStorageBackend implements SessionStore {
         }
     }
 
+    async clearCredentials(userId: string, sessionId: string): Promise<void> {
+        try {
+            await this.sql.query(
+                `DELETE FROM ${this.credentialsTableName} WHERE user_id = $1 AND session_id = $2`,
+                [userId, sessionId]
+            );
+        } catch (error) {
+            console.error('[NeonStorage] Failed to clear credentials:', error);
+        }
+    }
+
     async delete(userId: string, sessionId: string): Promise<void> {
         try {
             await this.sql.query(
@@ -258,6 +382,7 @@ export class NeonStorageBackend implements SessionStore {
 
     async clearAll(): Promise<void> {
         try {
+            await this.sql.query(`DELETE FROM ${this.credentialsTableName}`);
             await this.sql.query(`DELETE FROM ${this.tableName}`);
         } catch (error) {
             console.error('[NeonStorage] Failed to clear sessions:', error);
@@ -267,9 +392,17 @@ export class NeonStorageBackend implements SessionStore {
     async cleanupExpired(): Promise<void> {
         try {
             await this.sql.query(
-                `DELETE FROM ${this.tableName} WHERE expires_at < $1`,
+                `DELETE FROM ${this.tableName}
+                 WHERE expires_at IS NOT NULL
+                   AND expires_at < $1
+                   AND status <> 'active'`,
                 [new Date().toISOString()]
             );
+            await this.sql.query(
+                `DELETE FROM ${this.tableName}
+                 WHERE status = 'active' AND updated_at < $1`,
+                [new Date(Date.now() - DORMANT_SESSION_EXPIRATION_MS).toISOString()]
+            );
         } catch (error) {
             console.error('[NeonStorage] Failed to cleanup expired sessions:', error);
         }

package/src/server/storage/redis-backend.ts

@@ -1,14 +1,19 @@
 import type { Redis } from 'ioredis';
-import type { SessionStore, Session } from './types.js';
-import { SESSION_TTL_SECONDS } from '../../shared/constants.js';
+import type { SessionStore, Session, SessionCredentials } from './types.js';
 import { generateSessionId } from '../../shared/utils.js';
+import {
+    mergeSessionUpdate,
+    normalizeNewSession,
+    normalizeStoredSession,
+    resolveSessionRedisTtlSeconds,
+} from './session-lifecycle.js';
 
 /**
  * Redis implementation of SessionStore
  */
 export class RedisStorageBackend implements SessionStore {
-    private readonly DEFAULT_TTL = SESSION_TTL_SECONDS;
     private readonly KEY_PREFIX = 'mcp:session:';
+    private readonly CREDENTIALS_KEY_PREFIX = 'mcp:credentials:';
     private readonly USER_ID_KEY_PREFIX = 'mcp:userId:';
     private readonly USER_ID_KEY_SUFFIX = ':sessions';
 
@@ -31,6 +36,10 @@ export class RedisStorageBackend implements SessionStore {
         return `${this.KEY_PREFIX}${userId}:${sessionId}`;
     }
 
+    private getCredentialsKey(userId: string, sessionId: string): string {
+        return `${this.CREDENTIALS_KEY_PREFIX}${userId}:${sessionId}`;
+    }
+
     /**
      * Generates Redis key for tracking all sessions for a user
      * @private
@@ -78,18 +87,18 @@ export class RedisStorageBackend implements SessionStore {
         return generateSessionId();
     }
 
-    async create(session: Session, ttl?: number): Promise<void> {
+    async create(session: Session): Promise<void> {
         const { sessionId, userId } = session;
         if (!sessionId || !userId) throw new Error('userId and sessionId required');
 
         const sessionKey = this.getSessionKey(userId, sessionId);
         const userIdKey = this.getUserIdKey(userId);
-        const effectiveTtl = ttl ?? this.DEFAULT_TTL;
+        const sessionWithLifecycle = normalizeNewSession(session);
+        const effectiveTtl = resolveSessionRedisTtlSeconds(sessionWithLifecycle);
 
-        /** ioredis syntax: set(key, val, 'EX', ttl, 'NX') */
         const result = await this.redis.set(
             sessionKey,
-            JSON.stringify(session),
+            JSON.stringify(sessionWithLifecycle),
             'EX',
             effectiveTtl,
             'NX'
@@ -101,11 +110,10 @@ export class RedisStorageBackend implements SessionStore {
 
         await this.redis.sadd(userIdKey, sessionId);
     }
-    async update(userId: string, sessionId: string, data: Partial<Session>, ttl?: number): Promise<void> {
+    async update(userId: string, sessionId: string, data: Partial<Session>): Promise<void> {
         const sessionKey = this.getSessionKey(userId, sessionId);
-        const effectiveTtl = ttl ?? this.DEFAULT_TTL;
 
-        /** Lua script for atomic parsing, merging, and saving */
+        /** Lua script for atomic get-and-set with expiration refresh. */
         const script = `
             local currentStr = redis.call("GET", KEYS[1])
             if not currentStr then
@@ -113,21 +121,25 @@ export class RedisStorageBackend implements SessionStore {
             end
 
             local current = cjson.decode(currentStr)
-            local updates = cjson.decode(ARGV[1])
+            local updated = cjson.decode(ARGV[1])
 
-            for k,v in pairs(updates) do
-                current[k] = v
-            end
-
-            redis.call("SET", KEYS[1], cjson.encode(current), "EX", ARGV[2])
+            redis.call("SET", KEYS[1], cjson.encode(updated), "EX", ARGV[2])
             return 1
         `;
 
+        const current = await this.get(userId, sessionId);
+        if (!current) {
+            throw new Error(`Session ${sessionId} not found for userId ${userId}`);
+        }
+
+        const updated = mergeSessionUpdate(current, data);
+        const effectiveTtl = resolveSessionRedisTtlSeconds(updated);
+
         const result = await this.redis.eval(
             script,
             1,
             sessionKey,
-            JSON.stringify(data),
+            JSON.stringify(updated),
             effectiveTtl
         );
 
@@ -136,6 +148,29 @@ export class RedisStorageBackend implements SessionStore {
         }
     }
 
+    async patchCredentials(userId: string, sessionId: string, data: Partial<SessionCredentials>): Promise<void> {
+        const sessionKey = this.getSessionKey(userId, sessionId);
+        const credentialsKey = this.getCredentialsKey(userId, sessionId);
+        const sessionExists = await this.redis.exists(sessionKey);
+        if (!sessionExists) {
+            throw new Error(`Session ${sessionId} not found for userId ${userId}`);
+        }
+
+        const session = await this.get(userId, sessionId);
+        const currentTtl = await this.redis.ttl(sessionKey);
+        const effectiveTtl = currentTtl > 0 && session
+            ? currentTtl
+            : resolveSessionRedisTtlSeconds(session ?? { status: 'pending' });
+        const currentStr = await this.redis.get(credentialsKey);
+        const current = currentStr ? JSON.parse(currentStr) as SessionCredentials : { sessionId, userId };
+        const credentials = { ...current, ...data, sessionId, userId };
+        if (effectiveTtl > 0) {
+            await this.redis.set(credentialsKey, JSON.stringify(credentials), 'EX', effectiveTtl);
+        } else {
+            await this.redis.set(credentialsKey, JSON.stringify(credentials));
+        }
+    }
+
     async get(userId: string, sessionId: string): Promise<Session | null> {
         try {
             const sessionKey = this.getSessionKey(userId, sessionId);
@@ -145,14 +180,32 @@ export class RedisStorageBackend implements SessionStore {
                 return null;
             }
 
-            const Session: Session = JSON.parse(sessionDataStr);
-            return Session;
+            const session: Session = JSON.parse(sessionDataStr);
+            return normalizeStoredSession(session);
         } catch (error) {
             console.error('[RedisStorageBackend] Failed to get session:', error);
             return null;
         }
     }
 
+    async getCredentials(userId: string, sessionId: string): Promise<SessionCredentials | null> {
+        const session = await this.get(userId, sessionId);
+        if (!session) return null;
+
+        const credentialsStr = await this.redis.get(this.getCredentialsKey(userId, sessionId));
+        return credentialsStr ? JSON.parse(credentialsStr) as SessionCredentials : { sessionId, userId };
+    }
+
+    async clearCredentials(userId: string, sessionId: string): Promise<void> {
+        await this.patchCredentials(userId, sessionId, {
+            clientInformation: null,
+            tokens: null,
+            codeVerifier: null,
+            clientId: null,
+            oauthState: null,
+        });
+    }
+
     async listIds(userId: string): Promise<string[]> {
         const sessions = await this.list(userId);
         return sessions.map((session) => session.sessionId);
@@ -176,7 +229,9 @@ export class RedisStorageBackend implements SessionStore {
                 await this.redis.srem(userIdKey, ...staleSessionIds);
             }
 
-            return results.filter((session): session is Session => session !== null);
+            return results
+                .filter((session): session is Session => session !== null)
+                .map((session) => normalizeStoredSession(session));
         } catch (error) {
             console.error(`[RedisStorageBackend] Failed to get session data for ${userId}:`, error);
             return [];
@@ -186,10 +241,11 @@ export class RedisStorageBackend implements SessionStore {
     async delete(userId: string, sessionId: string): Promise<void> {
         try {
             const sessionKey = this.getSessionKey(userId, sessionId);
+            const credentialsKey = this.getCredentialsKey(userId, sessionId);
             const userIdKey = this.getUserIdKey(userId);
 
             await this.redis.srem(userIdKey, sessionId);
-            await this.redis.del(sessionKey);
+            await this.redis.del(sessionKey, credentialsKey);
         } catch (error) {
             console.error('[RedisStorageBackend] Failed to remove session:', error);
         }
@@ -206,7 +262,7 @@ export class RedisStorageBackend implements SessionStore {
                     }
 
                     try {
-                        return (JSON.parse(data) as Session).sessionId;
+                        return normalizeStoredSession(JSON.parse(data) as Session).sessionId;
                     } catch (error) {
                         console.error('[RedisStorageBackend] Failed to parse session while listing all session IDs:', error);
                         return null;
@@ -224,8 +280,9 @@ export class RedisStorageBackend implements SessionStore {
     async clearAll(): Promise<void> {
         try {
             const keys = await this.scanKeys(`${this.KEY_PREFIX}*`);
+            const credentialKeys = await this.scanKeys(`${this.CREDENTIALS_KEY_PREFIX}*`);
             const userIdKeys = await this.scanKeys(`${this.USER_ID_KEY_PREFIX}*${this.USER_ID_KEY_SUFFIX}`);
-            const allKeys = [...keys, ...userIdKeys];
+            const allKeys = [...keys, ...credentialKeys, ...userIdKeys];
             if (allKeys.length > 0) {
                 await this.redis.del(...allKeys);
             }
@@ -254,6 +311,7 @@ export class RedisStorageBackend implements SessionStore {
                 const staleSessionIds = sessionIds.filter((_, index) => existenceChecks[index] === 0);
                 if (staleSessionIds.length > 0) {
                     await this.redis.srem(userIdKey, ...staleSessionIds);
+                    await this.redis.del(...staleSessionIds.map((sessionId) => this.getCredentialsKey(userId, sessionId)));
                 }
 
                 const remainingCount = await this.redis.scard(userIdKey);