@oglofus/auth 1.0.1 → 1.1.1

package/dist/plugins/email-otp.js

@@ -1,17 +1,31 @@
+import { addSeconds, cloneWithout, createId, createNumericCode, secretHash, secretVerify } from "../core/utils.js";
+import { ensureFields } from "../core/validators.js";
 import { AuthError } from "../errors/index.js";
 import { createIssue } from "../issues/index.js";
 import { errorOperation, successOperation } from "../types/results.js";
-import { addSeconds, createId, createNumericCode, secretHash, secretVerify, } from "../core/utils.js";
-import { cloneWithout } from "../core/utils.js";
-import { ensureFields } from "../core/validators.js";
 const PENDING_PREFIX = "pending:";
+const EMAIL_OTP_REQUEST_POLICY = { limit: 3, windowSeconds: 300 };
+const OTP_VERIFY_POLICY = { limit: 10, windowSeconds: 300 };
 const isPendingUserId = (value) => value.startsWith(PENDING_PREFIX);
+const createRateLimitedError = (retryAfterSeconds) => new AuthError("RATE_LIMITED", "Too many requests.", 429, [], retryAfterSeconds === undefined ? undefined : { retryAfterSeconds });
 export const emailOtpPlugin = (config) => {
     const ttl = config.challengeTtlSeconds ?? 10 * 60;
     const maxAttempts = config.maxAttempts ?? 5;
     const codeLength = config.codeLength ?? 6;
-    const verifyChallenge = async (challengeId, code, now) => {
+    const verifyChallenge = async (challengeId, code, now, request, security, rateLimiter) => {
         const challenge = await config.otp.findChallengeById(challengeId);
+        if (rateLimiter) {
+            const policy = security?.rateLimits?.otpVerify ?? OTP_VERIFY_POLICY;
+            const result = await rateLimiter.consume(request?.ip
+                ? `otpVerify:ip:${request.ip}:identity:${challenge?.email ?? `challenge:${challengeId}`}`
+                : `otpVerify:identity:${challenge?.email ?? `challenge:${challengeId}`}`, policy.limit, policy.windowSeconds);
+            if (!result.allowed) {
+                return {
+                    ok: false,
+                    error: createRateLimitedError(result.retryAfterSeconds),
+                };
+            }
+        }
         if (!challenge) {
             return {
                 ok: false,
@@ -23,9 +37,7 @@ export const emailOtpPlugin = (config) => {
         if (challenge.expiresAt.getTime() <= now.getTime()) {
             return {
                 ok: false,
-                error: new AuthError("OTP_EXPIRED", "OTP has expired.", 400, [
-                    createIssue("OTP expired", ["code"]),
-                ]),
+                error: new AuthError("OTP_EXPIRED", "OTP has expired.", 400, [createIssue("OTP expired", ["code"])]),
             };
         }
         if (challenge.attempts >= maxAttempts) {
@@ -45,9 +57,7 @@ export const emailOtpPlugin = (config) => {
             }
             return {
                 ok: false,
-                error: new AuthError("OTP_INVALID", "Invalid OTP code.", 400, [
-                    createIssue("Invalid code", ["code"]),
-                ]),
+                error: new AuthError("OTP_INVALID", "Invalid OTP code.", 400, [createIssue("Invalid code", ["code"])]),
             };
         }
         const consumed = await config.otp.consumeChallenge(challengeId);
@@ -68,7 +78,7 @@ export const emailOtpPlugin = (config) => {
     return {
         kind: "auth_method",
         method: "email_otp",
-        version: "1.0.0",
+        version: "2.0.0",
         supports: {
             register: true,
         },
@@ -79,6 +89,13 @@ export const emailOtpPlugin = (config) => {
         createApi: (ctx) => ({
             request: async (input, request) => {
                 const email = input.email.trim().toLowerCase();
+                if (ctx.adapters.rateLimiter) {
+                    const policy = ctx.security?.rateLimits?.emailOtpRequest ?? EMAIL_OTP_REQUEST_POLICY;
+                    const limited = await ctx.adapters.rateLimiter.consume(request?.ip ? `emailOtpRequest:ip:${request.ip}:identity:${email}` : `emailOtpRequest:identity:${email}`, policy.limit, policy.windowSeconds);
+                    if (!limited.allowed) {
+                        return errorOperation(createRateLimitedError(limited.retryAfterSeconds));
+                    }
+                }
                 const user = await ctx.adapters.users.findByEmail(email);
                 const code = createNumericCode(codeLength);
                 const challenge = await config.otp.createChallenge({
@@ -116,7 +133,7 @@ export const emailOtpPlugin = (config) => {
             },
         }),
         register: async (ctx, input) => {
-            const verified = await verifyChallenge(input.challengeId, input.code, ctx.now());
+            const verified = await verifyChallenge(input.challengeId, input.code, ctx.now(), ctx.request, ctx.security, ctx.adapters.rateLimiter);
             if (!verified.ok) {
                 return errorOperation(verified.error);
             }
@@ -130,18 +147,14 @@ export const emailOtpPlugin = (config) => {
             if (requiredError) {
                 return errorOperation(requiredError);
             }
-            const payload = cloneWithout(input, [
-                "method",
-                "challengeId",
-                "code",
-            ]);
+            const payload = cloneWithout(input, ["method", "challengeId", "code"]);
             payload.email = verified.email;
             payload.emailVerified = true;
             const user = await ctx.adapters.users.create(payload);
             return successOperation({ user });
         },
         authenticate: async (ctx, input) => {
-            const verified = await verifyChallenge(input.challengeId, input.code, ctx.now());
+            const verified = await verifyChallenge(input.challengeId, input.code, ctx.now(), ctx.request, ctx.security, ctx.adapters.rateLimiter);
             if (!verified.ok) {
                 return errorOperation(verified.error);
             }

package/dist/plugins/index.d.ts

@@ -1,7 +1,8 @@
-export * from "./password.js";
 export * from "./email-otp.js";
 export * from "./magic-link.js";
 export * from "./oauth2.js";
+export * from "./organizations.js";
 export * from "./passkey.js";
+export * from "./password.js";
+export * from "./stripe.js";
 export * from "./two-factor.js";
-export * from "./organizations.js";

package/dist/plugins/index.js

@@ -1,7 +1,8 @@
-export * from "./password.js";
 export * from "./email-otp.js";
 export * from "./magic-link.js";
 export * from "./oauth2.js";
+export * from "./organizations.js";
 export * from "./passkey.js";
+export * from "./password.js";
+export * from "./stripe.js";
 export * from "./two-factor.js";
-export * from "./organizations.js";

package/dist/plugins/magic-link.js

@@ -1,18 +1,14 @@
+import { addSeconds, cloneWithout, createId, createToken, deterministicTokenHash } from "../core/utils.js";
+import { ensureFields } from "../core/validators.js";
 import { AuthError } from "../errors/index.js";
 import { createIssue } from "../issues/index.js";
 import { errorOperation, successOperation } from "../types/results.js";
-import { addSeconds, createId, createToken, secretHash, secretVerify } from "../core/utils.js";
-import { cloneWithout } from "../core/utils.js";
-import { ensureFields } from "../core/validators.js";
+const MAGIC_LINK_REQUEST_POLICY = { limit: 3, windowSeconds: 300 };
+const createRateLimitedError = (retryAfterSeconds) => new AuthError("RATE_LIMITED", "Too many requests.", 429, [], retryAfterSeconds === undefined ? undefined : { retryAfterSeconds });
 export const magicLinkPlugin = (config) => {
     const ttl = config.tokenTtlSeconds ?? 15 * 60;
     const verifyToken = async (token, now) => {
-        const hashed = secretHash(token, "magic-link");
-        // `secretHash` is salted; this branch should use deterministic hash.
-        // We keep compatibility by also trying deterministic fallback below.
-        const deterministicHash = secretHash(token, "deterministic_magic_link");
-        const candidate = (await config.links.findActiveTokenByHash(deterministicHash)) ??
-            (await config.links.findActiveTokenByHash(hashed));
+        const candidate = await config.links.findActiveTokenByHash(deterministicTokenHash(token, "magic_link"));
         if (!candidate) {
             return {
                 ok: false,
@@ -21,15 +17,6 @@ export const magicLinkPlugin = (config) => {
                 ]),
             };
         }
-        // Defensive verify for deterministic hash implementations.
-        if (!secretVerify(token, candidate.tokenHash) && candidate.tokenHash !== deterministicHash) {
-            return {
-                ok: false,
-                error: new AuthError("MAGIC_LINK_INVALID", "Invalid magic link.", 400, [
-                    createIssue("Invalid token", ["token"]),
-                ]),
-            };
-        }
         if (candidate.expiresAt.getTime() <= now.getTime()) {
             return {
                 ok: false,
@@ -55,7 +42,7 @@ export const magicLinkPlugin = (config) => {
     return {
         kind: "auth_method",
         method: "magic_link",
-        version: "1.0.0",
+        version: "2.0.0",
         supports: {
             register: true,
         },
@@ -66,9 +53,16 @@ export const magicLinkPlugin = (config) => {
         createApi: (ctx) => ({
             request: async (input, request) => {
                 const email = input.email.trim().toLowerCase();
+                if (ctx.adapters.rateLimiter) {
+                    const policy = ctx.security?.rateLimits?.magicLinkRequest ?? MAGIC_LINK_REQUEST_POLICY;
+                    const limited = await ctx.adapters.rateLimiter.consume(request?.ip ? `magicLinkRequest:ip:${request.ip}:identity:${email}` : `magicLinkRequest:identity:${email}`, policy.limit, policy.windowSeconds);
+                    if (!limited.allowed) {
+                        return errorOperation(createRateLimitedError(limited.retryAfterSeconds));
+                    }
+                }
                 const user = await ctx.adapters.users.findByEmail(email);
                 const rawToken = createToken();
-                const tokenHash = secretHash(rawToken, "deterministic_magic_link");
+                const tokenHash = deterministicTokenHash(rawToken, "magic_link");
                 const expiresAt = addSeconds(ctx.now(), ttl);
                 const token = await config.links.createToken({
                     userId: user?.id,
@@ -118,10 +112,7 @@ export const magicLinkPlugin = (config) => {
             if (requiredError) {
                 return errorOperation(requiredError);
             }
-            const payload = cloneWithout(input, [
-                "method",
-                "token",
-            ]);
+            const payload = cloneWithout(input, ["method", "token"]);
             payload.email = verified.email;
             payload.emailVerified = true;
             const user = await ctx.adapters.users.create(payload);

package/dist/plugins/oauth2.d.ts

@@ -14,6 +14,12 @@ export type ArcticAuthorizationCodeClient = {
 } | {
     validateAuthorizationCode(code: string, codeVerifier: string | null): Promise<OAuth2Tokens>;
 };
+export type OAuth2AuthorizationCodeExchangeInput = {
+    authorizationCode: string;
+    codeVerifier?: string;
+    redirectUri: string;
+};
+export type OAuth2AuthorizationCodeExchange = (input: OAuth2AuthorizationCodeExchangeInput) => Promise<OAuth2Tokens>;
 export type OAuth2ResolvedProfile<U extends UserBase, P extends string> = {
     provider?: P;
     providerUserId: string;
@@ -22,7 +28,7 @@ export type OAuth2ResolvedProfile<U extends UserBase, P extends string> = {
     profile?: Partial<U>;
 };
 export type OAuth2ProviderConfig<U extends UserBase, P extends string> = {
-    client: ArcticAuthorizationCodeClient;
+    exchangeAuthorizationCode: OAuth2AuthorizationCodeExchange;
     resolveProfile: (input: {
         provider: P;
         tokens: OAuth2Tokens;
@@ -38,4 +44,5 @@ export type OAuth2PluginConfig<U extends UserBase, P extends string, K extends k
     requiredProfileFields?: readonly K[];
     pendingProfileTtlSeconds?: number;
 };
+export declare const arcticAuthorizationCodeExchange: (client: ArcticAuthorizationCodeClient) => OAuth2AuthorizationCodeExchange;
 export declare const oauth2Plugin: <U extends UserBase, P extends string, K extends keyof U = never>(config: OAuth2PluginConfig<U, P, K>) => AuthMethodPlugin<"oauth2", OAuth2AuthenticateInput<P>, OAuth2AuthenticateInput<P>, U>;

package/dist/plugins/oauth2.js

@@ -1,12 +1,14 @@
 import { ArcticFetchError, OAuth2RequestError } from "arctic";
+import { addSeconds, createId } from "../core/utils.js";
+import { ensureFields } from "../core/validators.js";
 import { AuthError } from "../errors/index.js";
 import { createIssue } from "../issues/index.js";
 import { errorOperation, successOperation } from "../types/results.js";
-import { addSeconds, createId } from "../core/utils.js";
-import { ensureFields } from "../core/validators.js";
-const exchangeAuthorizationCode = async (client, authorizationCode, codeVerifier) => {
-    const validateAuthorizationCode = client.validateAuthorizationCode;
-    return validateAuthorizationCode(authorizationCode, codeVerifier);
+export const arcticAuthorizationCodeExchange = (client) => {
+    return async ({ authorizationCode, codeVerifier }) => {
+        const validateAuthorizationCode = client.validateAuthorizationCode;
+        return validateAuthorizationCode(authorizationCode, codeVerifier ?? null);
+    };
 };
 const getAccessToken = (tokens) => {
     try {
@@ -30,12 +32,28 @@ export const oauth2Plugin = (config) => {
     return {
         kind: "auth_method",
         method: "oauth2",
-        version: "1.0.0",
+        version: "2.0.0",
         supports: {
             register: false,
         },
         issueFields: {
-            authenticate: ["provider", "authorizationCode", "redirectUri", "codeVerifier"],
+            authenticate: ["provider", "authorizationCode", "redirectUri", "codeVerifier", "idempotencyKey"],
+        },
+        completePendingProfile: async (_ctx, { record, user }) => {
+            const continuation = record.continuation;
+            if (!continuation ||
+                typeof continuation.provider !== "string" ||
+                typeof continuation.providerUserId !== "string") {
+                return errorOperation(new AuthError("PLUGIN_MISCONFIGURED", "Pending OAuth2 profile is missing provider continuation metadata.", 500));
+            }
+            await config.accounts.linkAccount({
+                userId: user.id,
+                provider: continuation.provider,
+                providerUserId: continuation.providerUserId,
+                accessToken: continuation.accessToken,
+                refreshToken: continuation.refreshToken,
+            });
+            return successOperation(undefined);
         },
         authenticate: async (ctx, input) => {
             const providerConfig = config.providers[input.provider];
@@ -44,6 +62,19 @@ export const oauth2Plugin = (config) => {
                     createIssue("Provider disabled", ["provider"]),
                 ]));
             }
+            if (ctx.adapters.idempotency) {
+                if (!input.idempotencyKey || input.idempotencyKey.trim() === "") {
+                    return errorOperation(new AuthError("INVALID_INPUT", "idempotencyKey is required for OAuth2 callbacks.", 400, [
+                        createIssue("idempotencyKey is required when idempotency is enabled", ["idempotencyKey"]),
+                    ]));
+                }
+                const first = await ctx.adapters.idempotency.checkAndSet(`oauth2:${input.provider}:${input.idempotencyKey}`, ctx.security?.oauth2IdempotencyTtlSeconds ?? 300);
+                if (!first) {
+                    return errorOperation(new AuthError("CONFLICT", "Duplicate OAuth2 callback.", 409, [
+                        createIssue("Duplicate OAuth2 callback", ["idempotencyKey"]),
+                    ]));
+                }
+            }
             if ((providerConfig.pkceRequired ?? true) && !input.codeVerifier) {
                 return errorOperation(new AuthError("INVALID_INPUT", "Missing OAuth2 PKCE code verifier.", 400, [
                     createIssue("codeVerifier is required for this provider", ["codeVerifier"]),
@@ -51,7 +82,11 @@ export const oauth2Plugin = (config) => {
             }
             let tokens;
             try {
-                tokens = await exchangeAuthorizationCode(providerConfig.client, input.authorizationCode, input.codeVerifier ?? null);
+                tokens = await providerConfig.exchangeAuthorizationCode({
+                    authorizationCode: input.authorizationCode,
+                    codeVerifier: input.codeVerifier,
+                    redirectUri: input.redirectUri,
+                });
             }
             catch (error) {
                 if (error instanceof OAuth2RequestError) {
@@ -60,7 +95,9 @@ export const oauth2Plugin = (config) => {
                     ]));
                 }
                 if (error instanceof ArcticFetchError) {
-                    return errorOperation(new AuthError("OAUTH2_EXCHANGE_FAILED", "OAuth2 provider is unreachable right now.", 502, [createIssue("Provider request failed", ["provider"])]));
+                    return errorOperation(new AuthError("OAUTH2_EXCHANGE_FAILED", "OAuth2 provider is unreachable right now.", 502, [
+                        createIssue("Provider request failed", ["provider"]),
+                    ]));
                 }
                 return errorOperation(new AuthError("OAUTH2_EXCHANGE_FAILED", "OAuth2 code exchange failed.", 502));
             }
@@ -89,30 +126,35 @@ export const oauth2Plugin = (config) => {
                 return successOperation({ user: linkedUser });
             }
             const normalizedEmail = exchanged.email?.trim().toLowerCase();
+            if (!normalizedEmail) {
+                return errorOperation(new AuthError("INVALID_INPUT", "OAuth2 provider did not return a usable email.", 400, [
+                    createIssue("Email is required", ["email"]),
+                ]));
+            }
             const accessToken = getAccessToken(tokens);
             const refreshToken = getRefreshToken(tokens);
-            if (normalizedEmail) {
-                const existing = await ctx.adapters.users.findByEmail(normalizedEmail);
-                if (existing) {
-                    await config.accounts.linkAccount({
-                        userId: existing.id,
-                        provider: input.provider,
-                        providerUserId: exchanged.providerUserId,
-                        accessToken,
-                        refreshToken,
-                    });
-                    return successOperation({ user: existing });
-                }
+            const continuation = {
+                provider: input.provider,
+                providerUserId: exchanged.providerUserId,
+                accessToken,
+                refreshToken,
+            };
+            const existing = await ctx.adapters.users.findByEmail(normalizedEmail);
+            if (existing) {
+                await config.accounts.linkAccount({
+                    userId: existing.id,
+                    provider: input.provider,
+                    providerUserId: exchanged.providerUserId,
+                    accessToken,
+                    refreshToken,
+                });
+                return successOperation({ user: existing });
             }
             const profileRecord = {
                 ...exchanged.profile,
+                email: normalizedEmail,
+                emailVerified: exchanged.emailVerified ?? false,
             };
-            if (normalizedEmail) {
-                profileRecord.email = normalizedEmail;
-            }
-            if (exchanged.emailVerified !== undefined) {
-                profileRecord.emailVerified = exchanged.emailVerified;
-            }
             const requiredError = ensureFields(profileRecord, requiredProfileFields.map(String));
             if (requiredError) {
                 const pending = ctx.adapters.pendingProfiles;
@@ -128,7 +170,8 @@ export const oauth2Plugin = (config) => {
                     sourceMethod: "oauth2",
                     email: normalizedEmail,
                     missingFields,
-                    prefill: exchanged.profile ?? {},
+                    prefill: profileRecord,
+                    continuation: continuation,
                 };
                 await pending.create({
                     ...meta,
@@ -137,17 +180,7 @@ export const oauth2Plugin = (config) => {
                 });
                 return errorOperation(new AuthError("PROFILE_COMPLETION_REQUIRED", "Additional profile fields are required.", 400, requiredError.issues, meta), requiredError.issues);
             }
-            if (!normalizedEmail) {
-                return errorOperation(new AuthError("INVALID_INPUT", "OAuth2 provider did not return a usable email.", 400, [
-                    createIssue("Email is required", ["email"]),
-                ]));
-            }
-            const newUserPayload = {
-                ...profileRecord,
-                email: normalizedEmail,
-                emailVerified: exchanged.emailVerified ?? false,
-            };
-            const user = await ctx.adapters.users.create(newUserPayload);
+            const user = await ctx.adapters.users.create(profileRecord);
             await config.accounts.linkAccount({
                 userId: user.id,
                 provider: input.provider,

package/dist/plugins/organizations.d.ts

@@ -1,5 +1,5 @@
 import type { MembershipBase, OrganizationBase, UserBase } from "../types/model.js";
-import type { OrganizationsPluginConfig, OrganizationsPluginApi, DomainPlugin } from "../types/plugins.js";
+import type { DomainPlugin, OrganizationsPluginApi, OrganizationsPluginConfig } from "../types/plugins.js";
 export type OrganizationsPluginOptions<O extends OrganizationBase, Role extends string, M extends MembershipBase<Role>, Permission extends string, Feature extends string, LimitKey extends string, RequiredOrgFields extends keyof O = never> = OrganizationsPluginConfig<O, Role, M, Permission, Feature, LimitKey, RequiredOrgFields> & {
     inviteBaseUrl: string;
     inviteTtlSeconds?: number;

package/dist/plugins/organizations.js

@@ -1,7 +1,7 @@
+import { addSeconds, createId, createToken, deterministicTokenHash } from "../core/utils.js";
 import { AuthError } from "../errors/index.js";
 import { createIssue } from "../issues/index.js";
 import { errorOperation, successOperation } from "../types/results.js";
-import { addSeconds, createId, createToken, secretHash } from "../core/utils.js";
 const normalizeEmail = (value) => value.trim().toLowerCase();
 const findOwnerRoles = (roles) => Object.entries(roles)
     .filter(([, definition]) => definition.system?.owner)
@@ -49,7 +49,7 @@ export const organizationsPlugin = (config) => {
     const plugin = {
         kind: "domain",
         method: "organizations",
-        version: "1.0.0",
+        version: "2.0.0",
         __organizationConfig: config,
         createApi: (ctx) => {
             const ensureActorMembership = async (userId, organizationId) => {
@@ -75,15 +75,34 @@ export const organizationsPlugin = (config) => {
                     return membershipRes;
                 }
                 const resolved = resolveRole(membershipRes.data.role, config.handlers.roles);
+                const stripeApi = ctx.getPluginApi?.("stripe") ?? null;
+                let billingEntitlements = {
+                    features: {},
+                    limits: {},
+                };
+                if (stripeApi) {
+                    const stripeEntitlements = await stripeApi.getEntitlements({
+                        subject: {
+                            kind: "organization",
+                            organizationId,
+                        },
+                    });
+                    if (!stripeEntitlements.ok) {
+                        return stripeEntitlements;
+                    }
+                    billingEntitlements = stripeEntitlements.data;
+                }
                 const featureOverrides = await config.handlers.entitlements.getFeatureOverrides(organizationId);
                 const limitOverrides = await config.handlers.entitlements.getLimitOverrides(organizationId);
                 return successOperation({
                     features: {
                         ...resolved.features,
+                        ...billingEntitlements.features,
                         ...featureOverrides,
                     },
                     limits: {
                         ...resolved.limits,
+                        ...billingEntitlements.limits,
                         ...limitOverrides,
                     },
                 });
@@ -118,9 +137,7 @@ export const organizationsPlugin = (config) => {
                         });
                         return { organization, membership };
                     };
-                    const data = ctx.adapters.withTransaction
-                        ? await ctx.adapters.withTransaction(run)
-                        : await run();
+                    const data = ctx.adapters.withTransaction ? await ctx.adapters.withTransaction(run) : await run();
                     return successOperation(data);
                 },
                 inviteMember: async (input, request) => {
@@ -141,7 +158,7 @@ export const organizationsPlugin = (config) => {
                         return errorOperation(new AuthError("ROLE_INVALID", "Unknown role.", 400));
                     }
                     const rawToken = createToken(32);
-                    const tokenHash = secretHash(rawToken, "org-invite");
+                    const tokenHash = deterministicTokenHash(rawToken, "org_invite");
                     const invite = {
                         id: createId(),
                         organizationId: input.organizationId,
@@ -173,7 +190,7 @@ export const organizationsPlugin = (config) => {
                     });
                 },
                 acceptInvite: async (input) => {
-                    const invite = await config.handlers.invites.findActiveByTokenHash(secretHash(input.token, "org-invite"));
+                    const invite = await config.handlers.invites.findActiveByTokenHash(deterministicTokenHash(input.token, "org_invite"));
                     if (!invite) {
                         return errorOperation(new AuthError("ORGANIZATION_INVITE_INVALID", "Invalid invite token.", 400));
                     }
@@ -194,9 +211,17 @@ export const organizationsPlugin = (config) => {
                     const existing = await config.handlers.memberships.findByUserAndOrganization(input.userId, invite.organizationId);
                     let membership;
                     if (existing) {
-                        membership = await config.handlers.memberships.setRole(existing.id, invite.role);
+                        const updatedRole = await config.handlers.memberships.setRole(existing.id, invite.role);
+                        if (!updatedRole) {
+                            return errorOperation(new AuthError("MEMBERSHIP_NOT_FOUND", "Membership not found.", 404));
+                        }
+                        membership = updatedRole;
                         if (membership.status !== "active") {
-                            membership = await config.handlers.memberships.setStatus(membership.id, "active");
+                            const updatedStatus = await config.handlers.memberships.setStatus(membership.id, "active");
+                            if (!updatedStatus) {
+                                return errorOperation(new AuthError("MEMBERSHIP_NOT_FOUND", "Membership not found.", 404));
+                            }
+                            membership = updatedStatus;
                         }
                     }
                     else {
@@ -212,6 +237,27 @@ export const organizationsPlugin = (config) => {
                         membership,
                     });
                 },
+                setActiveOrganization: async (input, request) => {
+                    const session = await ctx.adapters.sessions.findById(input.sessionId);
+                    if (!session) {
+                        return errorOperation(new AuthError("SESSION_NOT_FOUND", "Session not found.", 404));
+                    }
+                    if (input.organizationId) {
+                        const memberships = await config.handlers.memberships.listByUser(session.userId);
+                        const active = memberships.find((membership) => membership.organizationId === input.organizationId && membership.status === "active");
+                        if (!active) {
+                            return errorOperation(new AuthError("MEMBERSHIP_FORBIDDEN", "No active membership for organization.", 403));
+                        }
+                    }
+                    const updatedSession = await config.handlers.organizationSessions.setActiveOrganization(input.sessionId, input.organizationId);
+                    if (!updatedSession) {
+                        return errorOperation(new AuthError("SESSION_NOT_FOUND", "Session not found.", 404));
+                    }
+                    return successOperation({
+                        sessionId: input.sessionId,
+                        activeOrganizationId: updatedSession.activeOrganizationId ?? null,
+                    });
+                },
                 setMemberRole: async (input, request) => {
                     const actorUserId = request?.userId;
                     if (!actorUserId) {
@@ -248,6 +294,9 @@ export const organizationsPlugin = (config) => {
                         }
                     }
                     const membership = await config.handlers.memberships.setRole(input.membershipId, input.role);
+                    if (!membership) {
+                        return errorOperation(new AuthError("MEMBERSHIP_NOT_FOUND", "Membership not found.", 404));
+                    }
                     return successOperation({ membership });
                 },
                 listMemberships: async (input) => {