Python之Streamlit登录页面实现

extra-streamlit-components==0.1.81
streamlit==1.61.1

import time
from datetime import datetime

import requests
import streamlit as st

from utils.CookieManager import COOKIE_NAME, cookie_manager
from utils.Config import config


def _do_login(username, password):
    """调用后端登录接口换取 utoken 并写入本地 cookie。契约见模块顶部注释。"""
    if not username or not password:
        st.session_state["login_error"] = "用户名和密码不能为空"
        return

    base_url = config.get("sparker").get("base_url")
    url = f"{base_url}/login"
    try:
        resp = requests.post(
            url,
            json={"username": username, "password": password},
            timeout=10,
        )
    except requests.RequestException:
        st.session_state["login_error"] = "服务器错误"
        return

    # 非 200：优先显示响应体中的 detail 字段
    if resp.status_code != 200:
        try:
            detail = resp.json().get("detail")
        except ValueError:
            detail = None
        st.session_state["login_error"] = detail or f"登录失败（HTTP {resp.status_code}）"
        return

    try:
        utoken = resp.json().get("utoken")
    except ValueError:
        st.session_state["login_error"] = "登录失败：响应不是合法 JSON"
        return
    if not utoken:
        st.session_state["login_error"] = "登录失败：响应中未包含 utoken"
        return

    # 置 session_state.utoken 并 rerun 跳转主页。cookie 的 set 作为关浏览器重开后的
    # 恢复来源。set 是异步 component 调用，rerun 可能截断指令下发导致 cookie 写不进；
    # 尝试在 set 与 rerun 之间留一段时间，给前端执行 set 写 cookie 的机会。
    cookie_manager.set(
        COOKIE_NAME,
        utoken,
        expires_at=datetime(2026, 12, 31),
    )
    st.session_state.utoken = utoken
    st.session_state.pop("login_error", None)


# 三列布局：标题、提示、表单统一放进中间列，保证左边缘对齐
left, form_col, right = st.columns([1, 2, 1])
with form_col:
    st.markdown("### 登录")
    st.caption("请输入用户名和密码")

    with st.form("login_form"):
        username = st.text_input("用户名", placeholder="请输入用户名")
        password = st.text_input("密码", type="password", placeholder="请输入密码")
        submitted = st.form_submit_button("登录", use_container_width=True)

    if submitted:
        _do_login(username, password)

    if st.session_state.get("login_error"):
        st.error(st.session_state["login_error"])

import time

import streamlit as st

from utils.CookieManager import COOKIE_NAME, cookie_manager

# 居中窄屏布局（聊天界面）
st.set_page_config(
    layout="centered",
    initial_sidebar_state="expanded"
)

login_page = st.Page("layout/login.py", title="Log in", icon=":material/login:")

dashboard = st.Page("layout/dashboard.py", title="Dashboard", icon=":material/dashboard:", default=True)
bugs = st.Page("layout/bugs.py", title="Bug reports", icon=":material/bug_report:")
alerts = st.Page("layout/alerts.py", title="System alerts", icon=":material/notification_important:")

search = st.Page("layout/search.py", title="Search", icon=":material/search:")
history = st.Page("layout/history.py", title="History", icon=":material/history:")

# 登录态以 session_state.utoken 为准（同会话即时生效、可立即跳转）。
# 关浏览器重开后，从本地 cookie 恢复 utoken 到 session_state（CookieManager 读取
# 为异步 component 回传，首帧可能读到空、第二帧才读到真实值，属该机制固有滞后）。
if "utoken" not in st.session_state:
    cookie_utoken = cookie_manager.get(COOKIE_NAME)
    if cookie_utoken:
        st.session_state.utoken = cookie_utoken

if st.session_state.get("utoken"):
    pg = st.navigation(
        {
            "Reports": [dashboard, bugs, alerts],
            "Tools": [search, history],
        }
    )
    # 侧边栏底部退出按钮：清 session_state + 删 cookie，立即回登录页
    with st.sidebar:
        if st.button("退出登录", use_container_width=True):
            st.session_state.pop("utoken", None)
            try:
                cookie_manager.delete(COOKIE_NAME, key=f"del_utoken_{int(time.time() * 1000)}")
                time.sleep(1)
            except KeyError:
                pass
            st.rerun()
else:
    pg = st.navigation([login_page])

pg.run()