fix(org-switch): Use Redis to pass org context during JWT generation

Problem: getAdditionalUserInfoClaim couldn't access Redis sessions
(Better Auth stores sessions in Redis, not PostgreSQL), causing
org switch to always fall back to the first organization.

Solution: Store org context in Redis before OAuth re-auth flow:
1. OrgSwitcher calls /api/auth/set-org-context with selected orgId
2. SSO stores {userId → orgId} in Redis with 5 min TTL
3. getAdditionalUserInfoClaim reads from Redis and uses it for tenant_id
4. Redis key is deleted after use (one-time use)

No database schema changes required.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: mjunaidca

apps/sso/src/app/api/auth/set-org-context/route.ts

@@ -0,0 +1,94 @@
+import { NextRequest, NextResponse } from "next/server";
+import { auth } from "@/lib/auth";
+import { headers } from "next/headers";
+import { redis } from "@/lib/redis";
+import { db } from "@/lib/db";
+import { member } from "@/lib/db/schema-export";
+import { eq } from "drizzle-orm";
+
+// Redis key prefix for org context
+const ORG_CONTEXT_PREFIX = "org_context:";
+// TTL: 5 minutes (enough time for OAuth flow to complete)
+const ORG_CONTEXT_TTL = 300;
+
+/**
+ * POST /api/auth/set-org-context
+ *
+ * Stores the user's selected organization in Redis for use during JWT generation.
+ * Called before initiating OAuth re-authentication for org switching.
+ *
+ * This is needed because:
+ * - getAdditionalUserInfoClaim can't access Redis sessions
+ * - We don't want to modify the DB schema
+ * - The OAuth token endpoint is a back-channel request without browser context
+ *
+ * Request body: { organizationId: string }
+ * Security: Validates user is a member of the requested organization.
+ */
+export async function POST(request: NextRequest) {
+  try {
+    // Check if Redis is available
+    if (!redis) {
+      console.error("[set-org-context] Redis not configured");
+      return NextResponse.json(
+        { error: "Service unavailable" },
+        { status: 503 }
+      );
+    }
+
+    // Get current session
+    const session = await auth.api.getSession({
+      headers: await headers(),
+    });
+
+    if (!session?.user) {
+      return NextResponse.json(
+        { error: "Unauthorized" },
+        { status: 401 }
+      );
+    }
+
+    const body = await request.json();
+    const { organizationId } = body;
+
+    if (!organizationId || typeof organizationId !== "string") {
+      return NextResponse.json(
+        { error: "organizationId is required" },
+        { status: 400 }
+      );
+    }
+
+    // Verify user is a member of this organization
+    const memberships = await db
+      .select()
+      .from(member)
+      .where(eq(member.userId, session.user.id));
+
+    const isMember = memberships.some((m: { organizationId: string }) => m.organizationId === organizationId);
+
+    if (!isMember) {
+      return NextResponse.json(
+        { error: "User is not a member of this organization" },
+        { status: 403 }
+      );
+    }
+
+    // Store in Redis with TTL
+    const key = `${ORG_CONTEXT_PREFIX}${session.user.id}`;
+    await redis.set(key, organizationId, { ex: ORG_CONTEXT_TTL });
+
+    console.log("[set-org-context] Stored org context:", session.user.id, "->", organizationId);
+
+    return NextResponse.json({
+      success: true,
+      organizationId,
+      expiresIn: ORG_CONTEXT_TTL,
+    });
+  } catch (error) {
+    console.error("[set-org-context] Error:", error);
+    return NextResponse.json(
+      { error: "Internal server error" },
+      { status: 500 }
+    );
+  }
+}

apps/sso/src/lib/auth.ts

@@ -13,7 +13,7 @@ import { deviceAuthorization } from "better-auth/plugins"; // 014-mcp-oauth-stan
 import { db } from "./db";
 import * as schema from "../../auth-schema"; // Use Better Auth generated schema
 import { member } from "../../auth-schema";
-import { eq, and, inArray, desc } from "drizzle-orm";
+import { eq, and, inArray } from "drizzle-orm";
 import { Resend } from "resend";
 import * as nodemailer from "nodemailer";
 import { TRUSTED_CLIENTS, DEFAULT_ORG_ID } from "./trusted-clients";
@@ -630,40 +630,29 @@ export const auth = betterAuth({
           console.log("[JWT] Organization Names:", organizationNames);
         }
 
-        // Get active organization from user's most recent session
-        // This allows org switcher to update tenant_id in JWT
+        // Get active organization from Redis (set by /api/auth/set-org-context during org switch)
+        // This is needed because getAdditionalUserInfoClaim can't access Redis sessions directly
         let activeOrgId: string | null = null;
 
-        // Query ALL sessions for this user to debug
-        const allUserSessions = await db
-          .select({
-            id: schema.session.id,
-            activeOrganizationId: schema.session.activeOrganizationId,
-            updatedAt: schema.session.updatedAt,
-          })
-          .from(schema.session)
-          .where(eq(schema.session.userId, user.id))
-          .orderBy(desc(schema.session.updatedAt));
-
-        console.log("[JWT] All sessions for user:", user.id);
-        console.log("[JWT] Session count:", allUserSessions.length);
-        allUserSessions.forEach((s: { id: string; activeOrganizationId: string | null; updatedAt: Date | null }, i: number) => {
-          console.log(`[JWT] Session ${i}: id=${s.id?.slice(0, 8)}..., activeOrgId=${s.activeOrganizationId}, updated=${s.updatedAt}`);
-        });
-
-        const userSessions = allUserSessions.slice(0, 1);
-
-        if (userSessions.length > 0 && userSessions[0].activeOrganizationId) {
-          // Verify the active org is one the user actually belongs to
-          if (organizationIds.includes(userSessions[0].activeOrganizationId)) {
-            activeOrgId = userSessions[0].activeOrganizationId;
-            console.log("[JWT] Using activeOrganizationId from session:", activeOrgId);
-          } else {
-            console.log("[JWT] Session activeOrganizationId not in user's orgs, falling back");
-            console.log("[JWT] organizationIds:", organizationIds);
+        // Check Redis for org context (set before OAuth re-auth for org switching)
+        if (redis) {
+          try {
+            const redisOrgId = await redis.get<string>(`org_context:${user.id}`);
+            console.log("[JWT] Redis org_context:", redisOrgId);
+
+            if (redisOrgId && organizationIds.includes(redisOrgId)) {
+              activeOrgId = redisOrgId;
+              console.log("[JWT] Using activeOrganizationId from Redis:", activeOrgId);
+              // Clean up the Redis key after use (one-time use)
+              await redis.del(`org_context:${user.id}`);
+            } else if (redisOrgId) {
+              console.log("[JWT] Redis org_context not in user's orgs, falling back");
+            }
+          } catch (error) {
+            console.error("[JWT] Error reading org context from Redis:", error);
           }
         } else {
-          console.log("[JWT] No session with activeOrganizationId found");
+          console.log("[JWT] Redis not available, skipping org context lookup");
         }
 
         // Use active org if set, otherwise fall back to first org

apps/web/src/components/OrgSwitcher.tsx

@@ -15,24 +15,27 @@ import {
 import { Button } from "@/components/ui/button"
 import { Building2, Check, Loader2 } from "lucide-react"
 
+// SSO URL for API calls
+const SSO_URL = process.env.NEXT_PUBLIC_SSO_URL || "http://localhost:3001"
+
 /**
  * Organization Switcher Component
  *
  * Allows users to switch between organizations they belong to.
  * Implements the enterprise pattern: Identity Session + Tenant-Scoped Tokens
  *
  * How it works:
- * 1. User clicks organization → calls organization.setActive()
- * 2. SSO updates session.activeOrganizationId in database
- * 3. Redirect to SSO login with prompt=none for silent re-auth
- * 4. SSO issues NEW JWT with updated tenant_id from session
- * 5. All subsequent requests use new JWT with new tenant_id
+ * 1. User clicks organization → calls organization.setActive() (updates SSO session)
+ * 2. Calls SSO /api/auth/set-org-context to store orgId in Redis
+ * 3. Redirect to SSO login for re-auth
+ * 4. SSO reads orgId from Redis in getAdditionalUserInfoClaim
+ * 5. SSO issues NEW JWT with updated tenant_id
+ * 6. All subsequent requests use new JWT with new tenant_id
  *
- * Why redirect instead of refresh?
- * - The JWT (taskflow_id_token cookie) is issued at login time
- * - router.refresh() only re-renders React components, doesn't replace JWT
- * - Need full OAuth flow to get new JWT with updated tenant_id claim
- * - prompt=none enables silent re-auth (no login screen shown)
+ * Why Redis for org context?
+ * - getAdditionalUserInfoClaim can't access Redis sessions (Better Auth limitation)
+ * - No DB schema changes needed
+ * - Short TTL (5 min) ensures cleanup
  */
 export function OrgSwitcher() {
   const { user } = useAuth()
@@ -65,20 +68,29 @@ export function OrgSwitcher() {
     try {
       console.log("[OrgSwitcher] Step 1: Calling setActive for org:", orgId)
 
-      // Step 1: Update SSO session's active organization
+      // Step 1: Update SSO session's active organization (for SSO UI consistency)
       const result = await organization.setActive({ organizationId: orgId })
       console.log("[OrgSwitcher] setActive result:", result)
 
-      // Delay to ensure session.activeOrganizationId is committed to database
-      // before starting OAuth flow. The setActive() call updates the session,
-      // and the OAuth flow needs to read the updated session.
-      // 500ms should be sufficient for DB write to complete.
-      await new Promise(resolve => setTimeout(resolve, 500))
+      // Step 2: Store org context in Redis for JWT generation
+      console.log("[OrgSwitcher] Step 2: Storing org context in Redis")
+      const contextResult = await fetch(`${SSO_URL}/api/auth/set-org-context`, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        credentials: "include", // Send SSO session cookies
+        body: JSON.stringify({ organizationId: orgId }),
+      })
+
+      if (!contextResult.ok) {
+        const errorData = await contextResult.json().catch(() => ({}))
+        console.error("[OrgSwitcher] Failed to store org context:", errorData)
+        throw new Error(errorData.error || "Failed to store organization context")
+      }
+
+      console.log("[OrgSwitcher] Org context stored in Redis")
 
-      // Step 2: Re-authenticate to get new JWT with updated tenant_id
-      // Pass orgId explicitly to ensure SSO uses the correct org even if session race condition
-      // The SSO will use this org_id parameter to set tenant_id in the JWT
-      console.log("[OrgSwitcher] Step 2: Initiating OAuth flow for new tokens with orgId:", orgId)
+      // Step 3: Re-authenticate to get new JWT with updated tenant_id
+      console.log("[OrgSwitcher] Step 3: Initiating OAuth flow for new tokens")
       await initiateLogin(orgId)
 
       // Note: initiateLogin sets window.location.href which triggers navigation