Streamlit Chat UI: Session-State, JWT Decode, SSE Consumption
What? (Concept Overview)
The Streamlit frontend (frontend/streamlit_app.py) consumes the FastAPI /chat/stream SSE endpoint from inside a browser-rendered chat UI. It demonstrates the single-conversation-per-user pattern: on login, fetch the user’s most recent conversation, restore it (or start a new one), and keep that conversation_id across the whole session. JWT decode is done client-side to read the id/sub claim — without exposing the secret — for fast UI placement.
Project Context
This file wires Streamlit 1.x + requests against the FastAPI backend. The login flow posts to /auth/login (root-level), stores the JWT in st.session_state.token, decodes the JWT payload (middle part) for customer_id, fetches /customers/{id}/conversations for the sidebar, and restores the most recent conversation. The chat input triggers a streaming requests.get(..., stream=True) against /chat/stream parsing each data: {...} line.
How? (Quick Reference Blocks)
3.1 Session-State Setup + JWT Decode
# frontend/streamlit_app.py — Session initialization
import streamlit as st
import requests
import base64
import json
import uuid
# Session-state vars that survive across reruns
if "token" not in st.session_state:
st.session_state.token = None
if "messages" not in st.session_state:
st.session_state.messages = []
if "customer_id" not in st.session_state:
st.session_state.customer_id = None
if "conversation_id" not in st.session_state:
st.session_state.conversation_id = None
def fetch_current_user():
"""Decode JWT payload (middle part) to retrieve customer_id."""
if not st.session_state.token:
return
try:
token_parts = st.session_state.token.split(".")
if len(token_parts) > 1:
payload_part = token_parts[1] + "=" * (-len(token_parts[1]) % 4) # pad
payload = json.loads(base64.b64decode(payload_part).decode("utf-8"))
real_id = payload.get("id") or payload.get("sub")
if real_id and str(real_id).isdigit():
st.session_state.customer_id = int(real_id)
except Exception as e:
print(f"User ID Error: {e}")3.2 Single-Conversation Restore on Login
# frontend/streamlit_app.py — fetch_active_conversation
def fetch_active_conversation():
"""Enforce single active conversation per user."""
headers = {"Authorization": f"Bearer {st.session_state.token}"}
response = requests.get(
f"{API_BASE_URL}/customers/{st.session_state.customer_id}/conversations",
headers=headers, timeout=5,
)
if response.status_code == 200:
conversations = response.json().get("conversations", [])
if conversations:
# pick most-recently-updated
conversations.sort(
key=lambda x: x.get("last_updated", ""), reverse=True)
st.session_state.conversation_id = conversations[0]["conversation_id"]
st.toast(f"📂 Resumed Conversation #{st.session_state.conversation_id}")
else:
new_id = int(str(uuid.uuid4().int)[:6])
st.session_state.conversation_id = new_id
st.toast(f"🆕 Started New Conversation #{new_id}")3.3 SSE Stream Consumer
# frontend/streamlit_app.py — chat input handler
if prompt := st.chat_input("How can I help you today?"):
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
status = st.status("📡 Tracing API Lifecycle...", expanded=True)
try:
headers = {
"Authorization": f"Bearer {st.session_state.token}",
"Accept": "text/event-stream",
}
stream_url = API_BASE_URL.replace("/api/v1", "") + "/chat/stream"
response = requests.get(
stream_url,
headers=headers,
params={
"message": prompt,
"customer_id": st.session_state.customer_id,
"conversation_id": st.session_state.conversation_id,
},
stream=True,
timeout=120,
)
bot_reply = ""
final_metadata = {}
if response.status_code == 200:
for line in response.iter_lines():
if line:
decoded = line.decode("utf-8")
if decoded.startswith("data: "):
data_str = decoded[6:].strip()
if data_str == "[DONE]":
break
event_data = json.loads(data_str)
etype = event_data.get("type")
if etype == "log":
status.write(f"📝 {event_data['content']}")
elif etype == "status":
status.write(f"🔀 {event_data['step']}: {event_data['content']}")
elif etype == "response":
bot_reply = event_data.get("content", "")
final_metadata = event_data.get("metadata", {})
new_id = event_data.get("conversation_id")
if new_id and new_id != st.session_state.conversation_id:
st.session_state.conversation_id = new_id
st.markdown(bot_reply)
st.session_state.messages.append(
{"role": "assistant", "content": bot_reply}
)
elif response.status_code == 401:
st.error("Session Expired. Please logout and login again.")
except Exception as e:
st.error(f"Connection Error: {e}")Why? (Parameter Breakdown
st.session_stateover module-level globals — Streamlit reruns the script top-to-bottom on every widget change. Module globals don’t survive a rerun;st.session_statedoes. Without it, the JWT evaporates after the first click.- Client-side JWT decode (without verification) — Reading
payload.get("id")for UI placement is display only. The server still validates the signature on every request. Verification on the client is unnecessary (and impossible without the secret). This is standard for OAuth2 frontends. base64padding trick (= * (-len % 4)) — JWT payload is base64url with stripped=padding. Python’s base64 decoder requires padding; the negative-mod math (-3 % 4 == 1) gives the exact pad count.- Single-conversation restore — UX call: rather than spawning a new conversation per browser session, restore the user’s last chat. Eliminates the “where did my chat go?” support complaint. New conversation is generated ONLY when the user has none history.
stream=True+iter_lines()for SSE —requests.get(..., stream=True)keeps the connection open;iter_lines()yields decoded lines as they arrive from the FastAPI server. Withoutstream=True, the body would buffer until the entire response completes — no live updates.Accept: text/event-streamon the request — Documents the protocol to anyone inspecting; doesn’t actually affect FastAPI streaming logic but makes curl traces self-explanatory.if new_id and new_id != st.session_state.conversation_id— Server-side, the coordinator can upgrade the conversation_id (e.g., when an archived conversation is recreated). UI must mirror that.
Common Pitfalls
- Calling JWT decode with
verify_signature=Truewithout the secret — would raiseInvalidSignatureError. Only decode for display, never for auth decisions. - Not handling the 401 mid-stream — Server rejects with 401 if the token expired; UI must show re-login prompt. Without catching the status code, the loop silently completes with an empty bot reply.
Real-World Interview Prep
Q1: Why is the JWT decode done client-side here? Is it a security issue?
A: Yes and no. JWTs are signed, not encrypted — anyone can decode the payload. That’s by design. The security model is: (1) the client can READ the payload (its sub, scopes, exp) but (2) cannot modify it without invalidating the signature. So decoding client-side for UI is safe. NEVER trust the client-decoded payload for server-side access control — always verify the signature on the server. The code here uses the decoded customer_id for chat-history requests, which the server re-validates anyway; no security bypass is possible.
Q2: Why use stream=True over WebSocket for the chat?
A: Three reasons. (1) Chat is server-to-client only; SSE fits that pattern. (2) The FastAPI backend already exposes /chat/stream because SSE works for browser clients without needing a WebSocket upgrade. (3) SSE auto-reconnects on socket close — if the browser drops, it re-issues the GET on its own with the last event-id. WebSockets need explicit reconnect logic. For React-based chats SSE is usually the right choice; only switch to WebSockets for bi-directional flows (e.g., live collaborative editing).
Q3: How would you handle a scenario where the user’s token expires mid-conversation?
A: The backend returns 401; the UI shows “Session Expired” + a logout/login form. The chat-state stays in st.session_state so user doesn’t lose messages — but st.session_state.token resets to None, triggering a rerun that renders the login screen. A nicer flow uses a refresh token to silently re-auth without the user noticing; the current architecture uses only an access token (3000-min default = ~50 hours), so an explicit re-login is the user-facing workaround.
Top-to-Bottom Code Walkthrough (frontend/streamlit_app.py)
Streamlit’s mental model: everything is a script that re-runs top to bottom every time the user clicks anything. State must live in st.session_state. Here’s how the frontend is wired in the project.
Imports
import streamlit as st— the only Streamlit-specific import needed for the UI.import httpx— for both the JSONPOSTendpoints and the SSEGET /chat/stream.from urllib.parse import urlparse— to safely extract origin tokens from cookies/tokens.import json— for parsing SSE event payloads.
st.set_page_config(...)
Sets browser-tab title and favicon. Must be the first st.* call.
init_session_state()
A function called explicitly at script start. It uses if "key" not in st.session_state: to idempotently initialise state. Keys initialised:
"messages"— list of{"role": "user"|"assistant", "content": "..."}dicts."conversation_id"— int. 0 means start a new conversation on first message."token"— the JWT from login."auth_headers"— prebuilt{"Authorization": "Bearer ..."}so every request reuses it.
Login form
with st.form("login"):
username = st.text_input("Email")
password = st.text_input("Password", type="password")
submitted = st.form_submit_button("Login")
if submitted:
res = httpx.post(f"{API_BASE_URL}/auth/login", data={"username": username, "password": password})
if res.status_code == 200:
st.session_state.token = res.json()["access_token"]
st.session_state.user = decode_jwt_no_verify(st.session_state.token) # for client-side display
st.rerun()
else:
st.error("Invalid credentials")The httpx.post sends form-encoded because FastAPI’s OAuth2PasswordRequestForm accepts form-encoded. The JWT is stored client-side in st.session_state (memory only, not cookies).
decode_jwt_no_verify(token) helper
Just decodes the payload WITHOUT verifying the signature — used only to extract user.id and email for display. Never trust this decoded payload for any access decision. Real verification always happens on the server.
Chat input handler
if user_input := st.chat_input("Ask a question..."):
st.session_state.messages.append({"role": "user", "content": user_input})
with httpx.stream("GET", f"{API_BASE_URL}/chat/stream",
params={"message": user_input, "customer_id": user_id, "conversation_id": conv_id},
headers=st.session_state.auth_headers) as r:
for line in r.iter_lines():
if line.startswith("data: "):
payload = json.loads(line[6:])
# Render: status events as captions, response events as chat bubbles
if payload.get("type") == "status":
st.caption(f"🔄 {payload['step']} ...")
elif payload.get("type") == "response":
st.session_state.messages.append({"role":"assistant", "content": payload["content"]})
st.chat_message("assistant").write(payload["content"])
elif payload.get("type") == "log":
st.text(payload["content"])The SSE iter_lines blocks on the first line; as the backend posts events, Streamlit re-renders. Each st.chat_message("assistant").write(...) adds a message bubble.
Re-render loop
After each handler step, st.rerun() is not called — Streamlit’s chat_message automatically appends to the visual output on the next re-render. State is preserved in st.session_state.
Common Pitfalls
Forgetting to populate st.session_state keys on first load raises KeyError. Always have an init_session_state() call.
Sending JSON instead of form-encoded to /auth/login — the server expects application/x-www-form-urlencoded. Use httpx.post(..., data=...) not json=.
Holding the SSE connection open across Streamlit re-renders — each rerun opens a new connection. Keep the SSE iterate loop inside a single if user_input block to keep it scoped.
Real-World Interview Prep
Q1: Why is Streamlit’s re-run-on-click model actually a feature for this use case?
A: No client-side state machine to write. The UI is a function of st.session_state — re-renders are pure. You don’t need React/Vue and the front-end dev workflow; one Python file with st.write calls.
Q2: How do you preserve chat scroll position during re-renders?
A: Streamlit’s st.chat_message handles scroll-to-bottom automatically. Manual control: use st.components.v1.html with a <script> that scrolls after each render.
Q3: Why no JWT in cookies — just in st.session_state?
A: st.session_state is client memory, so a hard refresh logs the user out. For a more permanent session, store the JWT in document.cookie via st.components.v1.html. For multi-tab state-sharing across devices — out of scope; switch to React + FastAPI.