# app/main.py import gradio as gr import httpx import websockets import asyncio import json import os import logging from contextlib import asynccontextmanager from fastapi import FastAPI, Depends # Import FastAPI itself from .database import connect_db, disconnect_db, database, metadata, users from .api import router as api_router from . import schemas, auth, dependencies from .websocket import manager # Import the connection manager instance from sqlalchemy.schema import CreateTable from sqlalchemy.dialects import sqlite logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) API_BASE_URL = "http://127.0.0.1:7860/api" # --- Lifespan (remains the same) --- @asynccontextmanager async def lifespan(app: FastAPI): # ... (same DB setup code) ... logger.info("Application startup: Connecting DB...") await connect_db() logger.info("Application startup: DB Connected. Checking/Creating tables...") if database.is_connected: try: check_query = "SELECT name FROM sqlite_master WHERE type='table' AND name=:table_name;" table_exists = await database.fetch_one(query=check_query, values={"table_name": users.name}) if not table_exists: logger.info(f"Table '{users.name}' not found, attempting creation using async connection...") dialect = sqlite.dialect() create_table_stmt = str(CreateTable(users).compile(dialect=dialect)) await database.execute(query=create_table_stmt) logger.info(f"Table '{users.name}' created successfully via async connection.") table_exists_after = await database.fetch_one(query=check_query, values={"table_name": users.name}) if table_exists_after: logger.info(f"Table '{users.name}' verified after creation.") else: logger.error(f"Table '{users.name}' verification FAILED after creation attempt!") else: logger.info(f"Table '{users.name}' already exists (checked via async connection).") except Exception as db_setup_err: logger.exception(f"CRITICAL error during async DB table setup: {db_setup_err}") else: logger.error("CRITICAL: Database connection failed, skipping table setup.") logger.info("Application startup: DB setup phase complete.") yield logger.info("Application shutdown: Disconnecting DB...") await disconnect_db() logger.info("Application shutdown: DB Disconnected.") # --- FastAPI App Setup (remains the same) --- app = FastAPI(lifespan=lifespan) app.include_router(api_router, prefix="/api") # --- Helper functions (make_api_request remains the same) --- async def make_api_request(method: str, endpoint: str, **kwargs): async with httpx.AsyncClient() as client: url = f"{API_BASE_URL}{endpoint}" try: response = await client.request(method, url, **kwargs); response.raise_for_status(); return response.json() except httpx.RequestError as e: logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}"); return {"error": f"Network error: {e}"} except httpx.HTTPStatusError as e: logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}") try: detail = e.response.json().get("detail", e.response.text) except json.JSONDecodeError: detail = e.response.text return {"error": f"API Error: {detail}"} except Exception as e: logger.error(f"Unexpected error during API call: {e}"); return {"error": f"Unexpected error: {str(e)}"} # --- WebSocket handling --- # <<< Pass state objects by reference >>> async def listen_to_websockets(token: str, notification_list_state: gr.State, notification_trigger_state: gr.State): """Connects to WS and updates state list and trigger when a message arrives.""" ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}" logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.") if not token: logger.warning(f"[{ws_listener_id}] No token provided. Task exiting.") return # Just exit, don't need to return state values ws_url_base = API_BASE_URL.replace("http", "ws") ws_url = f"{ws_url_base}/ws/{token}" logger.info(f"[{ws_listener_id}] Attempting to connect: {ws_url}") try: async with websockets.connect(ws_url, open_timeout=15.0) as websocket: logger.info(f"[{ws_listener_id}] WebSocket connected successfully.") while True: try: message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0) logger.info(f"[{ws_listener_id}] Received: {message_str[:100]}...") try: message_data = json.loads(message_str) if message_data.get("type") == "new_user": notification = schemas.Notification(**message_data) logger.info(f"[{ws_listener_id}] Processing 'new_user': {notification.message}") # --- Modify state values directly --- current_list = notification_list_state.value.copy() # Operate on copy current_list.insert(0, notification.message) if len(current_list) > 10: current_list.pop() notification_list_state.value = current_list # Assign back modified copy notification_trigger_state.value += 1 # Increment trigger # --- Log the update --- logger.info(f"[{ws_listener_id}] States updated: list len={len(notification_list_state.value)}, trigger={notification_trigger_state.value}") else: logger.warning(f"[{ws_listener_id}] Unknown message type: {message_data.get('type')}") # ... (error handling for parsing) ... except json.JSONDecodeError: logger.error(f"[{ws_listener_id}] JSON Decode Error: {message_str}") except Exception as parse_err: logger.error(f"[{ws_listener_id}] Message Processing Error: {parse_err}") # ... (error handling for websocket recv/connection) ... except asyncio.TimeoutError: logger.debug(f"[{ws_listener_id}] WebSocket recv timed out."); continue except websockets.ConnectionClosedOK: logger.info(f"[{ws_listener_id}] WebSocket closed normally."); break except websockets.ConnectionClosedError as e: logger.error(f"[{ws_listener_id}] WebSocket closed with error: {e}"); break except Exception as e: logger.error(f"[{ws_listener_id}] Listener loop error: {e}"); await asyncio.sleep(1); break # ... (error handling for websocket connect) ... except asyncio.TimeoutError: logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out.") except websockets.exceptions.InvalidURI: logger.error(f"[{ws_listener_id}] Invalid WebSocket URI.") except websockets.exceptions.WebSocketException as e: logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}") except Exception as e: logger.error(f"[{ws_listener_id}] Unexpected error in listener task: {e}") logger.info(f"[{ws_listener_id}] Listener task finished.") # No need to return state values when task ends # --- Gradio Interface --- with gr.Blocks(theme=gr.themes.Soft()) as demo: # State variables auth_token = gr.State(None) user_info = gr.State(None) notification_list = gr.State([]) websocket_task = gr.State(None) # Add trigger states notification_trigger = gr.State(0) last_polled_trigger = gr.State(0) # State to track last seen trigger # --- UI Components --- with gr.Tabs() as tabs: # --- Registration/Login Tabs (remain the same) --- with gr.TabItem("Register", id="register_tab"): gr.Markdown("## Create a new account") reg_email = gr.Textbox(label="Email", type="email") reg_password = gr.Textbox(label="Password (min 8 chars)", type="password") reg_confirm_password = gr.Textbox(label="Confirm Password", type="password") reg_button = gr.Button("Register") reg_status = gr.Textbox(label="Status", interactive=False) with gr.TabItem("Login", id="login_tab"): gr.Markdown("## Login to your account") login_email = gr.Textbox(label="Email", type="email") login_password = gr.Textbox(label="Password", type="password") login_button = gr.Button("Login") login_status = gr.Textbox(label="Status", interactive=False) # --- Welcome Tab --- with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab: gr.Markdown("## Welcome!", elem_id="welcome_header") welcome_message = gr.Markdown("", elem_id="welcome_message") logout_button = gr.Button("Logout") gr.Markdown("---") gr.Markdown("## Real-time Notifications") notification_display = gr.Textbox( # Visible display label="New User Alerts", lines=5, max_lines=10, interactive=False ) # <<< Hidden component for polling >>> dummy_poller = gr.Number(label="poller", value=0, visible=False, every=1) # --- Event Handlers --- # Registration Logic (remains the same) async def handle_register(email, password, confirm_password): if not email or not password or not confirm_password: return gr.update(value="Please fill fields.") if password != confirm_password: return gr.update(value="Passwords mismatch.") if len(password) < 8: return gr.update(value="Password >= 8 chars.") payload = {"email": email, "password": password} result = await make_api_request("post", "/register", json=payload) if "error" in result: return gr.update(value=f"Register failed: {result['error']}") else: return gr.update(value=f"Register success: {result.get('email')}! Log in.") reg_button.click(handle_register, inputs=[reg_email, reg_password, reg_confirm_password], outputs=[reg_status]) # Login Logic # <<< MODIFIED: Pass/Reset both trigger states >>> async def handle_login(email, password, current_task, current_trigger_val, current_last_poll_val): # Define failure outputs matching the outputs list fail_outputs = (gr.update(value="..."), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val, current_last_poll_val) if not email or not password: return fail_outputs[:1] + (gr.update(value="Enter email/password"),) + fail_outputs[2:] payload = {"email": email, "password": password} result = await make_api_request("post", "/login", json=payload) if "error" in result: return (gr.update(value=f"Login failed: {result['error']}"),) + fail_outputs[1:] else: token = result.get("access_token") user_data = await dependencies.get_optional_current_user(token) if not user_data: return (gr.update(value="Login ok but user fetch failed."),) + fail_outputs[1:] if current_task and not current_task.done(): current_task.cancel() try: await current_task except asyncio.CancelledError: logger.info("Previous WebSocket task cancelled.") # <<< Pass state objects to listener >>> new_task = asyncio.create_task(listen_to_websockets(token, notification_list, notification_trigger)) welcome_msg = f"Welcome, {user_data.email}!" # Reset triggers on successful login return ( gr.update(value="Login successful!"), token, user_data.model_dump(), # status, token, user_info gr.update(selected="welcome_tab"), gr.update(visible=True), gr.update(value=welcome_msg), # UI changes new_task, 0, 0 # websocket_task, notification_trigger, last_polled_trigger ) # <<< MODIFIED: Add last_polled_trigger to inputs/outputs >>> login_button.click( handle_login, inputs=[login_email, login_password, websocket_task, notification_trigger, last_polled_trigger], outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger, last_polled_trigger] ) # <<< Polling function >>> def poll_and_update(current_trigger_value, last_known_trigger, current_notif_list): """ Checks trigger state change and updates UI if needed. """ if current_trigger_value != last_known_trigger: logger.info(f"Polling detected trigger change ({last_known_trigger} -> {current_trigger_value}). Updating UI.") new_value = "\n".join(current_notif_list) # Return new display value AND update last_known_trigger state value return gr.update(value=new_value), current_trigger_value else: # No change, return NoUpdate for display AND existing last_known_trigger value return gr.NoUpdate(), last_known_trigger # <<< Attach polling function to the dummy component's change event >>> dummy_poller.change( fn=poll_and_update, inputs=[notification_trigger, last_polled_trigger, notification_list], outputs=[notification_display, last_polled_trigger], # Update display & last polled state queue=False # Try immediate update ) # Logout Logic # <<< MODIFIED: Reset both trigger states >>> async def handle_logout(current_task): if current_task and not current_task.done(): current_task.cancel() try: await current_task except asyncio.CancelledError: logger.info("WebSocket task cancelled on logout.") # Reset all relevant states return ( None, None, [], None, # auth_token, user_info, notification_list, websocket_task gr.update(selected="login_tab"), gr.update(visible=False), # tabs, welcome_tab gr.update(value=""), gr.update(value=""), # welcome_message, login_status 0, 0 ) # notification_trigger, last_polled_trigger (reset) # <<< MODIFIED: Add last_polled_trigger to outputs >>> logout_button.click( handle_logout, inputs=[websocket_task], outputs=[ auth_token, user_info, notification_list, websocket_task, tabs, welcome_tab, welcome_message, login_status, notification_trigger, last_polled_trigger # Add trigger states here ] ) # Mount Gradio App (remains the same) app = gr.mount_gradio_app(app, demo, path="/") # Run Uvicorn (remains the same) if __name__ == "__main__": import uvicorn uvicorn.run("app.main:app", host="0.0.0.0", port=7860, reload=True)