Update ui_components.py

ui_components.py

@@ -112,7 +112,41 @@ class UIComponents:
             #
             # On Spaces, Gradio auto-configures OAuth to the Space subdomain, so we intentionally do NOT pass redirect_url.
             # If running locally, OAuth via HF won't work; use HF_TOKEN instead.
-            self.login_button = gr.LoginButton(value="Sign in for Inference")
+            # Default: let Gradio handle redirect on Spaces
+            login_kwargs: Dict[str, Any] = {"value": "Sign in for Inference"}
+
+            # Opt-in: allow forcing a normalized redirect if needed
+            if os.getenv("FORCE_HF_REDIRECT", "0") == "1":
+                raw_host = os.getenv("SPACE_HOST", "").strip()
+                if raw_host:
+                    # Normalize host (strip scheme and trailing slashes)
+                    norm_host = raw_host
+                    for prefix in ("https://", "http://"):
+                        if norm_host.startswith(prefix):
+                            norm_host = norm_host[len(prefix):]
+                    norm_host = norm_host.rstrip("/")
+                    if norm_host:
+                        login_kwargs["redirect_url"] = f"https://{norm_host}/login/callback"
+
+            self.login_button = gr.LoginButton(**login_kwargs)
+
+            # Manual token fallback for cases where OAuth redirect is blocked/mismatched
+            with gr.Row():
+                self.token_input = gr.Textbox(
+                    label="Or paste an HF access token",
+                    placeholder="hf_...",
+                    type="password",
+                    scale=3
+                )
+                self.use_token_btn = gr.Button("Use token", variant="secondary")
+
+            # Status/debug info
+            server_seen_host = os.getenv("SPACE_HOST") or "(not on Spaces)"
+            self.login_status = gr.Markdown(
+                f"Server detected host: `{server_seen_host}`\n\nIf OAuth fails, paste a token from https://huggingface.co/settings/tokens.",
+                visible=True,
+                container=True
+            )
     
     def _create_provider_model_selection(self):
         """Create provider and model selection dropdowns with defaults"""
@@ -253,6 +287,24 @@ class UIComponents:
                     logger.info("✅ HuggingFace Inference client updated with OAuth token")
                 except Exception as e:
                     logger.error(f"❌ Failed to update HF client: {e}")
+            
+            # Update UI status optimistically
+            return gr.Markdown("✅ Logged in via Hugging Face OAuth", visible=True)
+
+        def handle_manual_token(token_text: str):
+            token_text = (token_text or "").strip()
+            if not token_text:
+                return gr.Markdown("❌ Please paste a valid HF token", visible=True)
+            os.environ["HF_TOKEN"] = token_text
+            try:
+                self.mcp_client.hf_client = OpenAI(
+                    base_url="https://router.huggingface.co/v1",
+                    api_key=token_text
+                )
+                return gr.Markdown("✅ Token set. Inference provider ready.", visible=True)
+            except Exception as e:
+                logger.error(f"Failed to set HF client with manual token: {e}")
+                return gr.Markdown(f"❌ Failed to initialize client: {e}", visible=True)
                     
         # Provider selection with auto-model loading
         def handle_provider_change(provider_id):
@@ -418,11 +470,21 @@ class UIComponents:
                 label=f"Active Servers ({len(server_choices)} loaded)"
             )
 
-        # Connect OAuth
-        demo.load(
-            fn=handle_oauth_profile,
-            outputs=[]
-        )
+        # Connect OAuth button explicitly and provide manual token fallback
+        if hasattr(self, "login_button"):
+            try:
+                self.login_button.login(
+                    fn=handle_oauth_profile,
+                    outputs=[self.login_status]
+                )
+            except Exception as e:
+                logger.warning(f"LoginButton.login hookup failed: {e}")
+        if hasattr(self, "use_token_btn"):
+            self.use_token_btn.click(
+                fn=handle_manual_token,
+                inputs=[self.token_input],
+                outputs=[self.login_status]
+            )
         
         # Connect provider/model dropdowns with auto-selection on load
         demo.load(