Removed Zustand storage, cleanup up auth storage IO interfaces, implementations, and usage.

This commit is contained in:
Bob Dickinson 2026-07-04 21:56:45 -07:00
commit 3c65bdb4c6

View file

@ -8,11 +8,8 @@ import {
OAuthTokensSchema,
} from "@modelcontextprotocol/sdk/shared/auth.js";
import type { OAuthStorage } from "./storage.js";
import {
type OAuthStore,
type ServerOAuthState,
waitForOAuthStorePersistLoad,
} from "./store.js";
import { type OAuthMemoryStore, type ServerOAuthState } from "./store.js";
import type { OAuthPersistBackend } from "./oauth-persist.js";
import type {
IdpSessionState,
OAuthClientRegistrationKind,
@ -20,47 +17,53 @@ import type {
SaveTokensOptions,
} from "./storage.js";
type OAuthStoreWithPersist = OAuthStore & {
persist: {
hasHydrated: () => boolean;
};
};
/**
* Concrete OAuthStorage implementation parameterized on a Zustand store.
* The store carries the storage adapter (sessionStorage, file, remote HTTP, …),
* so the same body works for browser, Node, and remote environments.
* Concrete OAuthStorage implementation backed by in-memory state and an explicit
* persist backend (file, remote HTTP, sessionStorage, …).
*/
export class OAuthStorageBase implements OAuthStorage {
private readonly store: OAuthStoreWithPersist;
private loadedPromise: Promise<void> | undefined;
private loaded = false;
private loadPromise: Promise<void> | undefined;
private readonly memory: OAuthMemoryStore;
private readonly backend: OAuthPersistBackend;
constructor(store: OAuthStore) {
this.store = store as OAuthStoreWithPersist;
constructor(memory: OAuthMemoryStore, backend: OAuthPersistBackend) {
this.memory = memory;
this.backend = backend;
}
load(): Promise<void> {
if (!this.loadedPromise) {
this.loadedPromise = (async () => {
if (this.store.persist.hasHydrated()) {
return;
}
await waitForOAuthStorePersistLoad(this.store);
})();
if (!this.loadPromise) {
this.loadPromise = this.doLoad();
}
return this.loadedPromise;
return this.loadPromise;
}
private async doLoad(): Promise<void> {
if (this.loaded) {
return;
}
const snapshot = await this.backend.read();
if (snapshot) {
this.memory.replace(snapshot);
}
this.loaded = true;
}
private async ensureLoaded(): Promise<void> {
await this.load();
}
private async persist(): Promise<void> {
await this.backend.write(this.memory.snapshot());
}
async getClientInformation(
serverUrl: string,
isPreregistered?: boolean,
): Promise<OAuthClientInformation | undefined> {
await this.ensureLoaded();
const state = this.store.getState().getServerState(serverUrl);
const state = this.memory.getState().getServerState(serverUrl);
const clientInfo = isPreregistered
? state.preregisteredClientInformation
: state.clientInformation;
@ -72,10 +75,11 @@ export class OAuthStorageBase implements OAuthStorage {
return await OAuthClientInformationSchema.parseAsync(clientInfo);
}
getClientRegistrationKind(
async getClientRegistrationKind(
serverUrl: string,
): OAuthClientRegistrationKind | undefined {
return this.store.getState().getServerState(serverUrl)
): Promise<OAuthClientRegistrationKind | undefined> {
await this.ensureLoaded();
return this.memory.getState().getServerState(serverUrl)
.clientRegistrationKind;
}
@ -85,10 +89,11 @@ export class OAuthStorageBase implements OAuthStorage {
options: SaveClientInformationOptions,
): Promise<void> {
await this.ensureLoaded();
this.store.getState().setServerState(serverUrl, {
this.memory.getState().setServerState(serverUrl, {
clientInformation,
clientRegistrationKind: options.registrationKind,
});
await this.persist();
}
async savePreregisteredClientInformation(
@ -96,10 +101,11 @@ export class OAuthStorageBase implements OAuthStorage {
clientInformation: OAuthClientInformation,
): Promise<void> {
await this.ensureLoaded();
this.store.getState().setServerState(serverUrl, {
this.memory.getState().setServerState(serverUrl, {
preregisteredClientInformation: clientInformation,
clientRegistrationKind: "static",
});
await this.persist();
}
async clearClientInformation(
@ -116,12 +122,13 @@ export class OAuthStorageBase implements OAuthStorage {
updates.clientRegistrationKind = undefined;
}
this.store.getState().setServerState(serverUrl, updates);
this.memory.getState().setServerState(serverUrl, updates);
await this.persist();
}
async getTokens(serverUrl: string): Promise<OAuthTokens | undefined> {
await this.ensureLoaded();
const state = this.store.getState().getServerState(serverUrl);
const state = this.memory.getState().getServerState(serverUrl);
if (!state.tokens) {
return undefined;
}
@ -135,19 +142,22 @@ export class OAuthStorageBase implements OAuthStorage {
options?: SaveTokensOptions,
): Promise<void> {
await this.ensureLoaded();
this.store.getState().setServerState(serverUrl, {
this.memory.getState().setServerState(serverUrl, {
tokens,
...(options?.enterpriseManaged === true && { enterpriseManaged: true }),
});
await this.persist();
}
async clearTokens(serverUrl: string): Promise<void> {
await this.ensureLoaded();
this.store.getState().setServerState(serverUrl, { tokens: undefined });
this.memory.getState().setServerState(serverUrl, { tokens: undefined });
await this.persist();
}
getCodeVerifier(serverUrl: string): string | undefined {
const state = this.store.getState().getServerState(serverUrl);
async getCodeVerifier(serverUrl: string): Promise<string | undefined> {
await this.ensureLoaded();
const state = this.memory.getState().getServerState(serverUrl);
return state.codeVerifier;
}
@ -156,33 +166,37 @@ export class OAuthStorageBase implements OAuthStorage {
codeVerifier: string,
): Promise<void> {
await this.ensureLoaded();
this.store.getState().setServerState(serverUrl, { codeVerifier });
this.memory.getState().setServerState(serverUrl, { codeVerifier });
await this.persist();
}
async clearCodeVerifier(serverUrl: string): Promise<void> {
await this.ensureLoaded();
this.store
.getState()
.setServerState(serverUrl, { codeVerifier: undefined });
this.memory.getState().setServerState(serverUrl, { codeVerifier: undefined });
await this.persist();
}
getScope(serverUrl: string): string | undefined {
const state = this.store.getState().getServerState(serverUrl);
async getScope(serverUrl: string): Promise<string | undefined> {
await this.ensureLoaded();
const state = this.memory.getState().getServerState(serverUrl);
return state.scope;
}
async saveScope(serverUrl: string, scope: string | undefined): Promise<void> {
await this.ensureLoaded();
this.store.getState().setServerState(serverUrl, { scope });
this.memory.getState().setServerState(serverUrl, { scope });
await this.persist();
}
async clearScope(serverUrl: string): Promise<void> {
await this.ensureLoaded();
this.store.getState().setServerState(serverUrl, { scope: undefined });
this.memory.getState().setServerState(serverUrl, { scope: undefined });
await this.persist();
}
getServerMetadata(serverUrl: string): OAuthMetadata | null {
const state = this.store.getState().getServerState(serverUrl);
async getServerMetadata(serverUrl: string): Promise<OAuthMetadata | null> {
await this.ensureLoaded();
const state = this.memory.getState().getServerState(serverUrl);
return state.serverMetadata || null;
}
@ -191,26 +205,27 @@ export class OAuthStorageBase implements OAuthStorage {
metadata: OAuthMetadata,
): Promise<void> {
await this.ensureLoaded();
this.store
.getState()
.setServerState(serverUrl, { serverMetadata: metadata });
this.memory.getState().setServerState(serverUrl, { serverMetadata: metadata });
await this.persist();
}
async clearServerMetadata(serverUrl: string): Promise<void> {
await this.ensureLoaded();
this.store
this.memory
.getState()
.setServerState(serverUrl, { serverMetadata: undefined });
await this.persist();
}
async clear(serverUrl: string): Promise<void> {
await this.ensureLoaded();
this.store.getState().clearServerState(serverUrl);
this.memory.getState().clearServerState(serverUrl);
await this.persist();
}
async getIdpSession(issuer: string): Promise<IdpSessionState | undefined> {
await this.ensureLoaded();
const session = this.store.getState().getIdpSession(issuer);
const session = this.memory.getState().getIdpSession(issuer);
if (
!session.idToken &&
!session.refreshToken &&
@ -226,16 +241,19 @@ export class OAuthStorageBase implements OAuthStorage {
session: Partial<IdpSessionState>,
): Promise<void> {
await this.ensureLoaded();
this.store.getState().setIdpSession(issuer, session);
this.memory.getState().setIdpSession(issuer, session);
await this.persist();
}
async clearIdpSession(issuer: string): Promise<void> {
await this.ensureLoaded();
this.store.getState().clearIdpSession(issuer);
this.memory.getState().clearIdpSession(issuer);
await this.persist();
}
async clearEnterpriseManagedResourceServers(): Promise<void> {
await this.ensureLoaded();
this.store.getState().clearEnterpriseManagedResourceServers();
this.memory.getState().clearEnterpriseManagedResourceServers();
await this.persist();
}
}