diff --git a/surfsense_backend/app/routes/google_drive_add_connector_route.py b/surfsense_backend/app/routes/google_drive_add_connector_route.py index 1c72443b4..5cab5b02b 100644 --- a/surfsense_backend/app/routes/google_drive_add_connector_route.py +++ b/surfsense_backend/app/routes/google_drive_add_connector_route.py @@ -155,6 +155,7 @@ async def connect_drive(space_id: int, user: User = Depends(current_active_user) async def reauth_drive( space_id: int, connector_id: int, + thread_id: str | None = None, user: User = Depends(current_active_user), session: AsyncSession = Depends(get_async_session), ): @@ -193,9 +194,10 @@ async def reauth_drive( flow = get_google_flow() state_manager = get_state_manager() - state_encoded = state_manager.generate_secure_state( - space_id, user.id, connector_id=connector_id - ) + extra: dict = {"connector_id": connector_id} + if thread_id: + extra["thread_id"] = thread_id + state_encoded = state_manager.generate_secure_state(space_id, user.id, **extra) auth_url, _ = flow.authorization_url( access_type="offline", @@ -282,6 +284,7 @@ async def drive_callback( user_id = UUID(data["user_id"]) space_id = data["space_id"] reauth_connector_id = data.get("connector_id") + reauth_thread_id = data.get("thread_id") logger.info( f"Processing Google Drive callback for user {user_id}, space {space_id}" @@ -352,6 +355,10 @@ async def drive_callback( logger.info( f"Re-authenticated Google Drive connector {db_connector.id} for user {user_id}" ) + if reauth_thread_id: + return RedirectResponse( + url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/new-chat/{reauth_thread_id}" + ) return RedirectResponse( url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/new-chat?modal=connectors&tab=all&success=true&connector=google-drive-connector&connectorId={db_connector.id}" )