@mcp-ts/sdk 2.3.3 → 2.4.0

package/src/server/mcp/storage-oauth-provider.ts

@@ -5,7 +5,8 @@ import type {
     OAuthClientMetadata,
     OAuthTokens
 } from "@modelcontextprotocol/sdk/shared/auth.js";
-import { sessions, type Session } from "../storage/index.js";
+import { nanoid } from "nanoid";
+import { sessions, type SessionCredentials } from "../storage/index.js";
 import {
     DEFAULT_CLIENT_NAME,
     DEFAULT_CLIENT_URI,
@@ -13,8 +14,9 @@ import {
     DEFAULT_POLICY_URI,
     SOFTWARE_ID,
     SOFTWARE_VERSION,
-    TOKEN_EXPIRY_BUFFER_MS,
+    STATE_EXPIRATION_MS,
 } from '../../shared/constants.js';
+import { formatOAuthState, parseOAuthState } from '../../shared/utils.js';
 
 /**
  * Extension of OAuthClientProvider interface with additional methods
@@ -29,8 +31,6 @@ export interface AgentsOAuthProvider extends OAuthClientProvider {
     ): Promise<{ valid: boolean; serverId?: string; error?: string }>;
     consumeState(state: string): Promise<void>;
     deleteCodeVerifier(): Promise<void>;
-    isTokenExpired(): boolean;
-    setTokenExpiresAt(expiresAt: number): void;
 }
 
 export interface StorageOAuthClientProviderOptions {
@@ -66,7 +66,6 @@ export class StorageOAuthClientProvider implements AgentsOAuthProvider {
     private _authUrl: string | undefined;
     private _clientId: string | undefined;
     private onRedirectCallback?: (url: string) => void;
-    private tokenExpiresAt?: number;
 
     /**
      * Creates a new session-backed OAuth provider
@@ -110,32 +109,32 @@ export class StorageOAuthClientProvider implements AgentsOAuthProvider {
     }
 
     /**
-     * Loads OAuth data from the session store
+     * Loads OAuth credentials from the session store
      * @private
      */
-    private async getSessionData(): Promise<Session> {
-        const data = await sessions.get(this.userId, this.sessionId);
+    private async getCredentials(): Promise<SessionCredentials> {
+        const data = await sessions.getCredentials(this.userId, this.sessionId);
         if (!data) {
-            return {} as Session;
+            return { userId: this.userId, sessionId: this.sessionId };
         }
         return data;
     }
 
     /**
-     * Saves OAuth data to the session store
-     * @param data - Partial OAuth data to save
+     * Saves OAuth credentials to the session store
+     * @param data - Partial OAuth credentials to save
      * @private
      * @throws Error if session doesn't exist (session must be created by controller layer)
      */
-    private async saveSessionData(data: Partial<Session>): Promise<void> {
-        await sessions.update(this.userId, this.sessionId, data);
+    private async patchCredentials(data: Partial<SessionCredentials>): Promise<void> {
+        await sessions.patchCredentials(this.userId, this.sessionId, data);
     }
 
     /**
      * Retrieves stored OAuth client information
      */
     async clientInformation(): Promise<OAuthClientInformationMixed | undefined> {
-        const data = await this.getSessionData();
+        const data = await this.getCredentials();
 
         if (data.clientId && !this._clientId) {
             this._clientId = data.clientId;
@@ -159,7 +158,7 @@ export class StorageOAuthClientProvider implements AgentsOAuthProvider {
      * Stores OAuth client information
      */
     async saveClientInformation(clientInformation: OAuthClientInformationFull): Promise<void> {
-        await this.saveSessionData({
+        await this.patchCredentials({
             clientInformation,
             clientId: clientInformation.client_id
         });
@@ -170,13 +169,7 @@ export class StorageOAuthClientProvider implements AgentsOAuthProvider {
      * Stores OAuth tokens
      */
     async saveTokens(tokens: OAuthTokens): Promise<void> {
-        const data: Partial<Session> = { tokens };
-
-        if (tokens.expires_in) {
-            this.tokenExpiresAt = Date.now() + (tokens.expires_in * 1000) - TOKEN_EXPIRY_BUFFER_MS;
-        }
-
-        await this.saveSessionData(data);
+        await this.patchCredentials({ tokens });
     }
 
     get authUrl() {
@@ -184,25 +177,67 @@ export class StorageOAuthClientProvider implements AgentsOAuthProvider {
     }
 
     async state(): Promise<string> {
-        return this.sessionId;
+        const nonce = nanoid(32);
+        await this.patchCredentials({
+            oauthState: {
+                nonce,
+                sessionId: this.sessionId,
+                serverId: this.serverId,
+                createdAt: Date.now(),
+            },
+            codeVerifier: null,
+        });
+        return formatOAuthState(nonce, this.sessionId);
     }
 
-    async checkState(_state: string): Promise<{ valid: boolean; serverId?: string; error?: string }> {
-        const data = await sessions.get(this.userId, this.sessionId);
+    async checkState(state: string): Promise<{ valid: boolean; serverId?: string; error?: string }> {
+        const parsed = parseOAuthState(state);
+        if (!parsed) {
+            return { valid: false, error: "Invalid OAuth state" };
+        }
+
+        if (parsed.sessionId !== this.sessionId) {
+            return { valid: false, error: "OAuth state mismatch" };
+        }
+
+        const data = await sessions.getCredentials(this.userId, parsed.sessionId);
 
         if (!data) {
             return { valid: false, error: "Session not found" };
         }
 
-        return { valid: true, serverId: this.serverId };
+        const oauthState = data.oauthState;
+        if (!oauthState) {
+            return { valid: false, error: "OAuth state not found" };
+        }
+
+        if (
+            oauthState.nonce !== parsed.nonce ||
+            oauthState.sessionId !== parsed.sessionId ||
+            oauthState.serverId !== this.serverId
+        ) {
+            return { valid: false, error: "OAuth state mismatch" };
+        }
+
+        if (Date.now() - oauthState.createdAt > STATE_EXPIRATION_MS) {
+            return { valid: false, error: "OAuth state expired" };
+        }
+
+        return { valid: true, serverId: oauthState.serverId };
     }
 
-    async consumeState(_state: string): Promise<void> {
-        // No-op
+    async consumeState(state: string): Promise<void> {
+        const result = await this.checkState(state);
+        if (!result.valid) {
+            throw new Error(result.error || "Invalid OAuth state");
+        }
+
+        await this.patchCredentials({ oauthState: null });
     }
 
     async redirectToAuthorization(authUrl: URL): Promise<void> {
         this._authUrl = authUrl.toString();
+        await sessions.update(this.userId, this.sessionId, { authUrl: this._authUrl });
         if (this.onRedirectCallback) {
             this.onRedirectCallback(authUrl.toString());
         }
@@ -214,26 +249,31 @@ export class StorageOAuthClientProvider implements AgentsOAuthProvider {
         if (scope === "all") {
             await sessions.delete(this.userId, this.sessionId);
         } else {
-            const updates: Partial<Session> = {};
+            const updates: Partial<SessionCredentials> = {};
 
             if (scope === "client") {
-                updates.clientInformation = undefined;
-                updates.clientId = undefined;
+                updates.clientInformation = null;
+                updates.clientId = null;
             } else if (scope === "tokens") {
-                updates.tokens = undefined;
+                updates.tokens = null;
             } else if (scope === "verifier") {
-                updates.codeVerifier = undefined;
+                updates.codeVerifier = null;
             }
-            await this.saveSessionData(updates);
+            await this.patchCredentials(updates);
         }
     }
 
     async saveCodeVerifier(verifier: string): Promise<void> {
-        await this.saveSessionData({ codeVerifier: verifier });
+        const data = await this.getCredentials();
+        if (data.codeVerifier) {
+            return;
+        }
+
+        await this.patchCredentials({ codeVerifier: verifier });
     }
 
     async codeVerifier(): Promise<string> {
-        const data = await this.getSessionData();
+        const data = await this.getCredentials();
 
         if (data.clientId && !this._clientId) {
             this._clientId = data.clientId;
@@ -246,27 +286,16 @@ export class StorageOAuthClientProvider implements AgentsOAuthProvider {
     }
 
     async deleteCodeVerifier(): Promise<void> {
-        await this.saveSessionData({ codeVerifier: undefined });
+        await this.patchCredentials({ codeVerifier: null });
     }
 
     async tokens(): Promise<OAuthTokens | undefined> {
-        const data = await this.getSessionData();
+        const data = await this.getCredentials();
 
         if (data.clientId && !this._clientId) {
             this._clientId = data.clientId;
         }
 
-        return data.tokens;
-    }
-
-    isTokenExpired(): boolean {
-        if (!this.tokenExpiresAt) {
-            return false;
-        }
-        return Date.now() >= this.tokenExpiresAt;
-    }
-
-    setTokenExpiresAt(expiresAt: number): void {
-        this.tokenExpiresAt = expiresAt;
+        return data.tokens ?? undefined;
     }
 }

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

@@ -1,7 +1,13 @@
 import { promises as fs } from 'fs';
 import * as path from 'path';
-import type { SessionStore, Session, SetClientOptions } from './types.js';
+import type { SessionStore, Session, SessionCredentials } from './types.js';
 import { generateSessionId } from '../../shared/utils.js';
+import {
+    mergeSessionUpdate,
+    normalizeNewSession,
+    normalizeStoredSession,
+    isSessionExpired,
+} from './session-lifecycle.js';
 
 /**
  * File system implementation of SessionStore
@@ -10,6 +16,7 @@ import { generateSessionId } from '../../shared/utils.js';
 export class FileStorageBackend implements SessionStore {
     private filePath: string;
     private memoryCache: Map<string, Session> | null = null;
+    private credentialsCache: Map<string, SessionCredentials> | null = null;
     private initialized = false;
 
     /**
@@ -35,15 +42,23 @@ export class FileStorageBackend implements SessionStore {
             const json = JSON.parse(data);
 
             this.memoryCache = new Map();
-            if (Array.isArray(json)) {
-                json.forEach((s: Session) => {
-                    this.memoryCache!.set(this.getSessionKey(s.userId || 'unknown', s.sessionId), s);
+            this.credentialsCache = new Map();
+            if (Array.isArray(json.sessions)) {
+                json.sessions.forEach((s: Session) => {
+                    const session = normalizeStoredSession(s);
+                    this.memoryCache!.set(this.getSessionKey(session.userId || 'unknown', session.sessionId), session);
+                });
+            }
+            if (Array.isArray(json.credentials)) {
+                json.credentials.forEach((c: SessionCredentials) => {
+                    this.credentialsCache!.set(this.getSessionKey(c.userId, c.sessionId), c);
                 });
             }
         } catch (error: any) {
             if (error.code === 'ENOENT') {
                 // File does not exist, initialize empty
                 this.memoryCache = new Map();
+                this.credentialsCache = new Map();
                 await this.flush();
             } else {
                 console.error('[FileStorage] Failed to load sessions:', error);
@@ -60,9 +75,11 @@ export class FileStorageBackend implements SessionStore {
     }
 
     private async flush(): Promise<void> {
-        if (!this.memoryCache) return;
-        const sessions = Array.from(this.memoryCache.values());
-        await fs.writeFile(this.filePath, JSON.stringify(sessions, null, 2), 'utf-8');
+        if (!this.memoryCache || !this.credentialsCache) return;
+        await fs.writeFile(this.filePath, JSON.stringify({
+            sessions: Array.from(this.memoryCache.values()),
+            credentials: Array.from(this.credentialsCache.values()),
+        }, null, 2), 'utf-8');
     }
 
     private getSessionKey(userId: string, sessionId: string): string {
@@ -73,7 +90,7 @@ export class FileStorageBackend implements SessionStore {
         return generateSessionId();
     }
 
-    async create(session: Session, ttl?: number): Promise<void> {
+    async create(session: Session): Promise<void> {
         await this.ensureInitialized();
         const { sessionId, userId } = session;
         if (!sessionId || !userId) throw new Error('userId and sessionId required');
@@ -83,12 +100,11 @@ export class FileStorageBackend implements SessionStore {
             throw new Error(`Session ${sessionId} already exists`);
         }
 
-        this.memoryCache!.set(sessionKey, session);
+        this.memoryCache!.set(sessionKey, normalizeNewSession(session));
         await this.flush();
-        // Note: TTL is ignored in file backend - sessions don't auto-expire
     }
 
-    async update(userId: string, sessionId: string, data: Partial<Session>, ttl?: number): Promise<void> {
+    async update(userId: string, sessionId: string, data: Partial<Session>): Promise<void> {
         await this.ensureInitialized();
         if (!userId || !sessionId) throw new Error('userId and sessionId required');
 
@@ -99,14 +115,22 @@ export class FileStorageBackend implements SessionStore {
             throw new Error(`Session ${sessionId} not found`);
         }
 
-        const updated = {
-            ...current,
-            ...data
-        };
+        const updated = mergeSessionUpdate(current, data);
 
         this.memoryCache!.set(sessionKey, updated);
         await this.flush();
-        // Note: TTL is ignored in file backend - sessions don't auto-expire
+    }
+
+    async patchCredentials(userId: string, sessionId: string, data: Partial<SessionCredentials>): Promise<void> {
+        await this.ensureInitialized();
+        const sessionKey = this.getSessionKey(userId, sessionId);
+        if (!this.memoryCache!.has(sessionKey)) {
+            throw new Error(`Session ${sessionId} not found`);
+        }
+
+        const current = this.credentialsCache!.get(sessionKey) ?? { sessionId, userId };
+        this.credentialsCache!.set(sessionKey, { ...current, ...data, sessionId, userId });
+        await this.flush();
     }
 
     async get(userId: string, sessionId: string): Promise<Session | null> {
@@ -115,6 +139,23 @@ export class FileStorageBackend implements SessionStore {
         return this.memoryCache!.get(sessionKey) || null;
     }
 
+    async getCredentials(userId: string, sessionId: string): Promise<SessionCredentials | null> {
+        await this.ensureInitialized();
+        const sessionKey = this.getSessionKey(userId, sessionId);
+        if (!this.memoryCache!.has(sessionKey)) return null;
+        return this.credentialsCache!.get(sessionKey) ?? { 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 list(userId: string): Promise<Session[]> {
         await this.ensureInitialized();
         return Array.from(this.memoryCache!.values()).filter(s => s.userId === userId);
@@ -130,7 +171,9 @@ export class FileStorageBackend implements SessionStore {
     async delete(userId: string, sessionId: string): Promise<void> {
         await this.ensureInitialized();
         const sessionKey = this.getSessionKey(userId, sessionId);
-        if (this.memoryCache!.delete(sessionKey)) {
+        const deleted = this.memoryCache!.delete(sessionKey);
+        this.credentialsCache!.delete(sessionKey);
+        if (deleted) {
             await this.flush();
         }
     }
@@ -143,12 +186,25 @@ export class FileStorageBackend implements SessionStore {
     async clearAll(): Promise<void> {
         await this.ensureInitialized();
         this.memoryCache!.clear();
+        this.credentialsCache!.clear();
         await this.flush();
     }
 
     async cleanupExpired(): Promise<void> {
-        // Could implement TTL check here using createdAt
         await this.ensureInitialized();
+        let changed = false;
+
+        for (const [key, session] of this.memoryCache!.entries()) {
+            if (!isSessionExpired(session)) continue;
+
+            this.memoryCache!.delete(key);
+            this.credentialsCache!.delete(key);
+            changed = true;
+        }
+
+        if (changed) {
+            await this.flush();
+        }
     }
 
     async disconnect(): Promise<void> {

package/src/server/storage/index.ts

@@ -63,27 +63,25 @@ function createSessionMutationEvent(prop: PropertyKey, args: any[]): SessionMuta
     const timestamp = Date.now();
 
     if (prop === 'create') {
-        const [session, ttl] = args as [Session, number | undefined];
+        const [session] = args as [Session];
         if (!session?.userId || !session?.sessionId) return null;
         return {
             type: 'create',
             userId: session.userId,
             sessionId: session.sessionId,
             session,
-            ttl,
             timestamp,
         };
     }
 
     if (prop === 'update') {
-        const [userId, sessionId, patch, ttl] = args as [string, string, Partial<Session>, number | undefined];
+        const [userId, sessionId, patch] = args as [string, string, Partial<Session>];
         if (!userId || !sessionId) return null;
         return {
             type: 'update',
             userId,
             sessionId,
             patch,
-            ttl,
             timestamp,
         };
     }

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

@@ -1,5 +1,6 @@
-import type { SessionStore, Session, SetClientOptions } from './types.js';
+import type { SessionStore, Session, SessionCredentials } from './types.js';
 import { generateSessionId } from '../../shared/utils.js';
+import { isSessionExpired, mergeSessionUpdate, normalizeNewSession } from './session-lifecycle.js';
 
 /**
  * In-memory implementation of SessionStore
@@ -8,6 +9,7 @@ import { generateSessionId } from '../../shared/utils.js';
 export class MemoryStorageBackend implements SessionStore {
     // Map<userId:sessionId, Session>
     private sessions = new Map<string, Session>();
+    private credentials = new Map<string, SessionCredentials>();
 
     // Map<userId, Set<sessionId>>
     private userIdSessions = new Map<string, Set<string>>();
@@ -26,7 +28,7 @@ export class MemoryStorageBackend 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');
 
@@ -35,17 +37,16 @@ export class MemoryStorageBackend implements SessionStore {
             throw new Error(`Session ${sessionId} already exists`);
         }
 
-        this.sessions.set(sessionKey, session);
+        this.sessions.set(sessionKey, normalizeNewSession(session));
 
         // Update index
         if (!this.userIdSessions.has(userId)) {
             this.userIdSessions.set(userId, new Set());
         }
         this.userIdSessions.get(userId)!.add(sessionId);
-        // Note: TTL is ignored in memory backend - sessions don't auto-expire
     }
 
-    async update(userId: string, sessionId: string, data: Partial<Session>, ttl?: number): Promise<void> {
+    async update(userId: string, sessionId: string, data: Partial<Session>): Promise<void> {
         if (!userId || !sessionId) throw new Error('userId and sessionId required');
 
         const sessionKey = this.getSessionKey(userId, sessionId);
@@ -55,13 +56,19 @@ export class MemoryStorageBackend implements SessionStore {
             throw new Error(`Session ${sessionId} not found`);
         }
 
-        const updated = {
-            ...current,
-            ...data
-        };
+        const updated = mergeSessionUpdate(current, data);
 
         this.sessions.set(sessionKey, updated);
-        // Note: TTL is ignored in memory backend - sessions don't auto-expire
+    }
+
+    async patchCredentials(userId: string, sessionId: string, data: Partial<SessionCredentials>): Promise<void> {
+        const sessionKey = this.getSessionKey(userId, sessionId);
+        if (!this.sessions.has(sessionKey)) {
+            throw new Error(`Session ${sessionId} not found`);
+        }
+
+        const current = this.credentials.get(sessionKey) ?? { sessionId, userId };
+        this.credentials.set(sessionKey, { ...current, ...data, sessionId, userId });
     }
 
 
@@ -70,6 +77,22 @@ export class MemoryStorageBackend implements SessionStore {
         return this.sessions.get(sessionKey) || null;
     }
 
+    async getCredentials(userId: string, sessionId: string): Promise<SessionCredentials | null> {
+        const sessionKey = this.getSessionKey(userId, sessionId);
+        if (!this.sessions.has(sessionKey)) return null;
+        return this.credentials.get(sessionKey) ?? { 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 set = this.userIdSessions.get(userId);
         return set ? Array.from(set) : [];
@@ -92,6 +115,7 @@ export class MemoryStorageBackend implements SessionStore {
     async delete(userId: string, sessionId: string): Promise<void> {
         const sessionKey = this.getSessionKey(userId, sessionId);
         this.sessions.delete(sessionKey);
+        this.credentials.delete(sessionKey);
 
         const set = this.userIdSessions.get(userId);
         if (set) {
@@ -108,13 +132,25 @@ export class MemoryStorageBackend implements SessionStore {
 
     async clearAll(): Promise<void> {
         this.sessions.clear();
+        this.credentials.clear();
         this.userIdSessions.clear();
     }
 
     async cleanupExpired(): Promise<void> {
-        // In-memory doesn't implement TTL automatically, 
-        // but we could check createdAt + TTL here if needed.
-        // For now, no-op.
+        for (const [key, session] of this.sessions.entries()) {
+            if (!isSessionExpired(session)) continue;
+
+            this.sessions.delete(key);
+            this.credentials.delete(key);
+
+            const set = this.userIdSessions.get(session.userId);
+            if (set) {
+                set.delete(session.sessionId);
+                if (set.size === 0) {
+                    this.userIdSessions.delete(session.userId);
+                }
+            }
+        }
     }
 
     async disconnect(): Promise<void> {