Update app.py

app.py

@@ -36,27 +36,25 @@ def run_bot():
 threading.Thread(target=run_bot).start()
 
 # Gradio ------------------------------------------------------------------------------------------------------------
-def validate_link(request: gr.Request) -> bool:
+def validate_link(request: gr.Request):
     url_str = str(request.url)
     query_params = parse_qs(urlparse(url_str).query)
     user_id = query_params.get('user_id', [None])[0]
     token = query_params.get('token', [None])[0]
 
     if user_id is None or token is None:
-        return False
+        return False, None, None
 
     if int(user_id) not in user_tokens or user_tokens[int(user_id)] != token:
-        return False
+        return False, None, None
 
-    return True
+    return True, user_id, token
 
 def hello(profile: gr.OAuthProfile | None, request: gr.Request) -> str:
-    if not validate_link(request):
-        return "❌ Invalid link. Please use the link provided by the bot."
+    is_valid, user_id, token = validate_link(request)
 
-    url_str = str(request.url)
-    query_params = parse_qs(urlparse(url_str).query)
-    user_id = query_params.get('user_id', [None])[0]
+    if not is_valid:
+        return "❌ Invalid link or expired token."
 
     if profile is None:
         return f"❌ Not logged in. User ID: {user_id}"
@@ -65,16 +63,17 @@ def hello(profile: gr.OAuthProfile | None, request: gr.Request) -> str:
 
 with gr.Blocks() as demo:
     gr.Markdown("# Gradio OAuth Space")
-    login_button = gr.LoginButton()
     m1 = gr.Markdown()
 
-    def check_link_and_load():
-        if validate_link(gr.Request()):
-            demo.load(hello, inputs=None, outputs=m1)
+    def conditional_login_button(request: gr.Request):
+        is_valid, _, _ = validate_link(request)
+        if is_valid:
+            return gr.LoginButton()
         else:
-            m1.update("❌ Invalid link. Please use the link provided by the bot.")
+            return gr.Markdown("❌ Invalid link or expired token.")
 
-    check_link_and_load()
+    login_button = conditional_login_button(gr.Request())
+    demo.load(hello, inputs=None, outputs=m1)
 
     def check_login_status():
         try:
@@ -89,6 +88,7 @@ with gr.Blocks() as demo:
         else:
             return f"Logged in as {session.get('username', 'Unknown')}"
 
-    login_button.click(check_login_wrapper, inputs=None, outputs=m1)
+    if isinstance(login_button, gr.LoginButton):
+        login_button.click(check_login_wrapper, inputs=None, outputs=m1)
 
-demo.launch()
+demo.launch()