Improve auth redirect type safety (#47)

src/app/(centered)/authorized/page.tsx

@@ -12,8 +12,7 @@ export default function Page() {
   const [isLoading, setIsLoading] = useState(true);
   const isLoadingRef = useRef(true);
   const router = useRouter();
-
-  const { state, ide } = useAuthState();
+  const authState = useAuthState();
 
   useEffect(() => {
     if (typeof isSignedIn !== 'undefined' && isLoadingRef.current) {
@@ -27,18 +26,22 @@ export default function Page() {
       let path;
 
       if (!isSignedIn) {
-        path = state ? `/sign-in?state=${state}&ide=${ide}` : '/sign-in';
+        path = authState.params
+          ? `/sign-in?${authState.params.toString()}`
+          : '/sign-in';
       } else if (!orgId) {
-        path = state ? `/select-org/${state}?ide=${ide}` : '/select-org';
+        path = authState.params
+          ? `/select-org?${authState.params.toString()}`
+          : '/select-org';
       } else {
-        path = state
-          ? `/extension/sign-in?state=${state}&ide=${ide}`
+        path = authState.params
+          ? `/extension/sign-in?${authState.params.toString()}`
           : '/dashboard';
       }
 
       setTimeout(() => router.push(path), 1000);
     }
-  }, [router, isLoading, isSignedIn, orgId, state, ide]);
+  }, [router, isLoading, isSignedIn, orgId, authState.params]);
 
   return (
     <div className="flex flex-row items-center gap-2">

src/app/(centered)/extension/sign-in/DeepLink.tsx

@@ -13,17 +13,17 @@ import { TraeLogo } from './TraeLogo';
 
 interface DeepLinkProps {
   editor: string;
-  deepLinkUrl: string;
+  editorRedirect: string;
 }
 
-export const DeepLink = ({ editor, deepLinkUrl }: DeepLinkProps) => {
+export const DeepLink = ({ editor, editorRedirect }: DeepLinkProps) => {
   const [redirectAttempted, setRedirectAttempted] = useState(false);
 
   useEffect(() => {
-    window.location.href = deepLinkUrl;
+    window.location.href = editorRedirect;
     const timer = setTimeout(() => setRedirectAttempted(true), 2000);
     return () => clearTimeout(timer);
-  }, [deepLinkUrl]);
+  }, [editorRedirect]);
 
   const ides = useMemo(
     () => [
@@ -31,34 +31,34 @@ export const DeepLink = ({ editor, deepLinkUrl }: DeepLinkProps) => {
         editor: 'vscode',
         title: 'Visual Studio Code',
         logo: VSCodeLogo,
-        href: deepLinkUrl.replace(`${editor}://`, 'vscode://'),
+        href: editorRedirect.replace(`${editor}://`, 'vscode://'),
       },
       {
         editor: 'vscode-insiders',
         title: 'Visual Studio Code (Insiders)',
         logo: VSCodeInsidersLogo,
-        href: deepLinkUrl.replace(`${editor}://`, 'vscode-insiders://'),
+        href: editorRedirect.replace(`${editor}://`, 'vscode-insiders://'),
       },
       {
         editor: 'cursor',
         title: 'Cursor',
         logo: CursorLogo,
-        href: deepLinkUrl.replace(`${editor}://`, 'cursor://'),
+        href: editorRedirect.replace(`${editor}://`, 'cursor://'),
       },
       {
         editor: 'windsurf',
         title: 'Windsurf',
         logo: WindsurfLogo,
-        href: deepLinkUrl.replace(`${editor}://`, 'windsurf://'),
+        href: editorRedirect.replace(`${editor}://`, 'windsurf://'),
       },
       {
         editor: 'trae',
         title: 'Trae',
         logo: TraeLogo,
-        href: deepLinkUrl.replace(`${editor}://`, 'trae://'),
+        href: editorRedirect.replace(`${editor}://`, 'trae://'),
       },
     ],
-    [editor, deepLinkUrl],
+    [editor, editorRedirect],
   );
 
   const currentIde = ides.find((ide) => ide.editor === editor) ?? ides[0]!;
@@ -70,7 +70,7 @@ export const DeepLink = ({ editor, deepLinkUrl }: DeepLinkProps) => {
           <div className="text-sm text-muted-foreground">
             Redirecting to {currentIde.title}...
           </div>
-          <Link href={deepLinkUrl} title={currentIde.title}>
+          <Link href={editorRedirect} title={currentIde.title}>
             <currentIde.logo />
           </Link>
         </div>

src/app/(centered)/extension/sign-in/page.tsx

@@ -1,22 +1,29 @@
 import { redirect } from 'next/navigation';
 import { auth } from '@clerk/nextjs/server';
 
+import { AuthStateParam } from '@/types';
 import { getSignInToken } from '@/actions/auth';
-import { EXTENSION_EDITOR, EXTENSION_URL } from '@/lib/constants';
+import { EXTENSION_EDITOR, EXTENSION_URI_SCHEME } from '@/lib/constants';
 
 import { DeepLink } from './DeepLink';
 
 type Props = {
-  searchParams: Promise<{ state?: string; ide?: string }>;
+  searchParams: Promise<{ state?: string; auth_redirect?: string }>;
 };
 
 export default async function Page(props: Props) {
-  const { state, ide } = await props.searchParams;
+  const { state, auth_redirect: authRedirect = EXTENSION_URI_SCHEME } =
+    await props.searchParams;
 
   if (!state) {
     redirect(`/sign-in`);
   }
 
+  const authParams = new URLSearchParams({
+    [AuthStateParam.State]: state,
+    [AuthStateParam.AuthRedirect]: authRedirect,
+  });
+
   const { userId, orgId } = await auth();
 
   const code = userId
@@ -24,30 +31,28 @@ export default async function Page(props: Props) {
     : undefined;
 
   if (!code) {
-    redirect(`/sign-in?state=${state}&ide=${ide}`);
+    redirect(`/sign-in?${authParams.toString()}`);
   }
 
   if (!orgId) {
-    redirect(`/select-org/${state}?ide=${ide}`);
+    redirect(`/select-org?${authParams.toString()}`);
   }
 
-  let editor;
-  let deepLinkUrl;
+  let editor = EXTENSION_EDITOR;
+  let editorRedirect = EXTENSION_URI_SCHEME;
 
   try {
     const params = new URLSearchParams({ state, code });
 
-    deepLinkUrl = new URL(
+    editorRedirect = new URL(
       `/auth/clerk/callback?${params.toString()}`,
-      ide ?? EXTENSION_URL,
+      authRedirect,
     ).toString();
 
-    editor = new URL(deepLinkUrl).protocol.slice(0, -1);
+    editor = new URL(editorRedirect).protocol.slice(0, -1);
   } catch (_) {
     // Use the defaults if we can't parse the URL.
-    editor = EXTENSION_EDITOR;
-    deepLinkUrl = EXTENSION_URL;
   }
 
-  return <DeepLink editor={editor} deepLinkUrl={deepLinkUrl} />;
+  return <DeepLink editor={editor} editorRedirect={editorRedirect} />;
 }

src/app/(centered)/select-org/SelectOrg.tsx

@@ -4,9 +4,12 @@ import Link from 'next/link';
 import { OrganizationList, useOrganizationList } from '@clerk/nextjs';
 import { LoaderCircle } from 'lucide-react';
 
+import { useAuthState } from '@/hooks/useAuthState';
 import { Button } from '@/components/ui';
 
-export const SelectOrg = ({ state }: { state?: string }) => {
+export const SelectOrg = () => {
+  const authState = useAuthState();
+
   const { isLoaded, userMemberships, userSuggestions, userInvitations } =
     useOrganizationList({
       userMemberships: true,
@@ -14,8 +17,8 @@ export const SelectOrg = ({ state }: { state?: string }) => {
       userSuggestions: true,
     });
 
-  const redirectUrl = state
-    ? `/extension/sign-in?state=${state}`
+  const redirectUrl = authState.params
+    ? `/extension/sign-in?${authState.params.toString()}`
     : '/dashboard';
 
   const isLoading =

src/app/(centered)/select-org/[[...state]]/page.tsx

@@ -1,10 +0,0 @@
-import { SelectOrg } from './SelectOrg';
-
-type Props = {
-  params: Promise<{ state: string }>;
-};
-
-export default async function Page({ params }: Props) {
-  const { state } = await params;
-  return <SelectOrg state={state} />;
-}

src/app/(centered)/select-org/page.tsx

@@ -0,0 +1,5 @@
+import { SelectOrg } from './SelectOrg';
+
+export default async function Page() {
+  return <SelectOrg />;
+}

src/hooks/useAuthState.ts

@@ -4,40 +4,52 @@ import { useCallback } from 'react';
 import { useSessionStorage, useMount } from 'react-use';
 import { useSearchParams } from 'next/navigation';
 
-import { EXTENSION_URL } from '@/lib/constants';
+import { AuthStateParam, type AuthState } from '@/types';
+import { EXTENSION_URI_SCHEME } from '@/lib/constants';
 
-export type AuthState = {
-  state?: string;
-  ide?: string;
-};
+export const useAuthState = () => {
+  const searchParams = useSearchParams();
 
-export type SetAuthState = (state: AuthState) => void;
+  const [state, setState] = useSessionStorage<string | undefined>(
+    AuthStateParam.State,
+    searchParams.get(AuthStateParam.State) ?? undefined,
+  );
 
-export const useAuthState = (): AuthState & { set: SetAuthState } => {
-  const [state, setState] = useSessionStorage<string | undefined>('state');
-  const [ide, setIde] = useSessionStorage<string | undefined>('ide');
+  const [authRedirect, setAuthRedirect] = useSessionStorage<string | undefined>(
+    AuthStateParam.AuthRedirect,
+    searchParams.get(AuthStateParam.AuthRedirect) ?? undefined,
+  );
 
   const set = useCallback(
     (state: AuthState) => {
       setState(state.state);
-      setIde(state.ide);
+      setAuthRedirect(state.authRedirect);
     },
-    [setState, setIde],
+    [setState, setAuthRedirect],
   );
 
-  return { state, ide, set };
+  const params = state
+    ? new URLSearchParams({
+        [AuthStateParam.State]: state,
+        [AuthStateParam.AuthRedirect]: authRedirect ?? EXTENSION_URI_SCHEME,
+      })
+    : undefined;
+
+  const authState: AuthState = {
+    state,
+    authRedirect,
+    params,
+  };
+
+  return { ...authState, set };
 };
 
 export const useSetAuthState = () => {
-  const { set } = useAuthState();
-  const searchParams = useSearchParams();
+  const { state, authRedirect = EXTENSION_URI_SCHEME, set } = useAuthState();
 
   useMount(() => {
-    const state = searchParams.get('state');
-    const ide = searchParams.get('ide');
-
     if (state) {
-      set({ state, ide: ide ?? EXTENSION_URL });
+      set({ state, authRedirect });
     }
   });
 };

src/lib/constants.ts

@@ -6,4 +6,4 @@ export const EXTENSION_NAME = 'roo-cline';
 
 export const EXTENSION_IDENTIFIER = `${EXTENSION_PUBLISHER}.${EXTENSION_NAME}`;
 
-export const EXTENSION_URL = `vscode://${EXTENSION_IDENTIFIER}`;
+export const EXTENSION_URI_SCHEME = `${EXTENSION_EDITOR}://${EXTENSION_IDENTIFIER}`;

src/types/auth.ts

@@ -0,0 +1,10 @@
+export enum AuthStateParam {
+  State = 'state',
+  AuthRedirect = 'auth_redirect',
+}
+
+export type AuthState = {
+  state?: string;
+  authRedirect?: string;
+  params?: URLSearchParams;
+};

src/types/index.ts

@@ -1,3 +1,4 @@
 export * from './api';
+export * from './auth';
 export * from './org';
 export * from './time-period';