Upload 2 files

main_app.py

@@ -140,35 +140,37 @@ if 'server_started' not in st.session_state:
     else:
         logger.warning("⚠️ セッション管理サーバー起動失敗 - フォールバックモードで動作")
 
-query_params = st.experimental_get_query_params()
-temp_token = query_params.get("temp_token", [None])[0]
 
+SPACE_URL = "https://huggingface.co/spaces/sirochild/mari-chat-3"
+LOGIN_URL = f"{SPACE_URL}/oauth/huggingface/login"
+ME_URL = f"{SPACE_URL}/me"
 def get_current_user_id():
-    if temp_token:
-        try:
-            user_info = requests.get(f"http://localhost:8000/user?temp_token={temp_token}").json()
-            user_id = user_info.get("id")
-            is_hf_user = True
-        except Exception as e:
-            st.warning(f"HF API エラー: {e}. 匿名モードを使用します。")
-            user_id = str(uuid.uuid4())
-            is_hf_user = False
-    else:
-        user_id = str(uuid.uuid4())
-        is_hf_user = False
-    return user_id, is_hf_user
-
-user_id, is_hf_user = get_current_user_id()
-st.write(f"ユーザー ID: {user_id} (HF ログイン: {is_hf_user})")
+    # --- セッションステートにユーザー情報を保存 ---
+    if "user" not in st.session_state:
+        st.session_state.user = None
 
-if not is_hf_user:
-    if st.button("HuggingFaceでログイン"):
-        # FastAPI /login から auth_url を取得
-        res = requests.get("http://localhost:8000/login").json()
-        auth_url = res.get("auth_url")
-        # ユーザーを OAuth ページにリダイレクト
-        st.markdown(f'<meta http-equiv="refresh" content="0; url={auth_url}">', unsafe_allow_html=True)
+    # --- ログインしてなければボタンを表示 ---
+    if st.session_state.user is None:
+        st.markdown(f"[👉 Hugging Faceでログイン]({LOGIN_URL})")
 
+    else:
+        st.success(f"ログイン中: {st.session_state.user['preferred_username']}")
+        st.json(st.session_state.user)
+
+    # --- クエリパラメータに code があれば認証処理 ---
+    query_params = st.query_params
+    if "code" in query_params:
+        code = query_params["code"]
+        # /me にアクセスしてユーザー情報を取得
+        try:
+            resp = requests.get(ME_URL, cookies={"hf_oauth_code": code})
+            if resp.status_code == 200:
+                st.session_state.user = resp.json()
+                st.experimental_rerun()
+            else:
+                st.error("認証失敗しました")
+        except Exception as e:
+            st.error(f"エラー: {e}")
 
 # --- 必要なモジュールのインポート ---
 

session_api_server.py

@@ -12,6 +12,7 @@ import os
 import time
 import secrets
 import base64
+from huggingface_hub import attach_huggingface_oauth, parse_huggingface_oauth
 from datetime import datetime, timedelta
 from typing import Optional, Dict, Any
 import logging
@@ -36,72 +37,18 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-HF_CLIENT_ID = os.getenv("OAUTH_CLIENT_ID")
-HF_CLIENT_SECRET = os.getenv("OAUTH_CLIENT_SECRET")
-HF_REDIRECT_URI = os.getenv("OAUTH_REDIRECT_URI", "https://sirochild-mari-chat-3.hf.space/callback")
-HF_SCOPES = os.getenv("HF_OAUTH_SCOPES", "openid profile email")
+# OAuthルートを追加
+attach_huggingface_oauth(app)
 
-# 一時トークンとアクセストークンのマッピング
-TEMP_STORE = {}  # temp_token -> hf_access_token
-STATE_STORE = {} # state の管理
-
-@app.get("/login")
-def login():
-    state = secrets.token_urlsafe(16)
-    STATE_STORE[state] = True
-    auth_url = (
-        f"https://huggingface.co/oauth/authorize"
-        f"?client_id={HF_CLIENT_ID}"
-        f"&redirect_uri={HF_REDIRECT_URI}"
-        f"&response_type=code"
-        f"&scope={HF_SCOPES}"
-        f"&state={state}"
-    )
-    return RedirectResponse(auth_url)
-
-@app.get("/callback")
-def callback(code: str, state: str):
-    if state not in STATE_STORE:
-        return JSONResponse({"error": "Invalid state"}, status_code=400)
-    del STATE_STORE[state]
-
-    # Basic 認証ヘッダー作成
-    auth_header = base64.b64encode(f"{HF_CLIENT_ID}:{HF_CLIENT_SECRET}".encode()).decode()
-
-    token_res = requests.post(
-    "https://huggingface.co/oauth/token",
-    data={
-        "grant_type": "authorization_code",
-        "code": code,
-        "redirect_uri": HF_REDIRECT_URI,
-        "client_id": HF_CLIENT_ID,
-        "client_secret": HF_CLIENT_SECRET,
+@app.get("/me")
+def me(request: Request):
+    oauth_info = parse_huggingface_oauth(request)
+    if oauth_info is None:
+        return {"msg": "Not logged in"}
+    return {
+        "username": oauth_info.user_info.preferred_username,
+        "email": oauth_info.user_info.email,
     }
-).json()
-
-    access_token = token_res.get("access_token")
-    if not access_token:
-        return JSONResponse({"error": "Failed to get access token"}, status_code=400)
-
-    # 一時トークン発行
-    temp_token = secrets.token_urlsafe(16)
-    TEMP_STORE[temp_token] = access_token
-
-    # Streamlit に temp_token を渡す
-    return RedirectResponse(f"https://sirochild-mari-chat-3.hf.space/?temp_token={temp_token}")
-
-@app.get("/user")
-def get_user(temp_token: str):
-    access_token = TEMP_STORE.get(temp_token)
-    if not access_token:
-        return JSONResponse({"error": "Invalid or expired token"}, status_code=400)
-
-    user_res = requests.get(
-        "https://huggingface.co/oauth/userinfo",
-        headers={"Authorization": f"Bearer {access_token}"}
-    ).json()
-
-    return JSONResponse(user_res)
 
 
 class SessionManager: