amaye15 commited on
Commit
527a73c
·
1 Parent(s): 57e813e
Files changed (1) hide show
  1. app/main.py +129 -204
app/main.py CHANGED
@@ -9,117 +9,87 @@ import logging
9
  from contextlib import asynccontextmanager
10
 
11
  from fastapi import FastAPI, Depends # Import FastAPI itself
12
- # --- Import necessary items from database.py ---
13
  from .database import connect_db, disconnect_db, database, metadata, users
14
  from .api import router as api_router
15
  from . import schemas, auth, dependencies
16
- from .websocket import manager
17
- # --- Import SQLAlchemy helpers for DDL generation ---
18
  from sqlalchemy.schema import CreateTable
19
- from sqlalchemy.dialects import sqlite # Assuming SQLite, adjust if needed
20
 
21
- # Configure logging
22
  logging.basicConfig(level=logging.INFO)
23
  logger = logging.getLogger(__name__)
24
 
25
- # Base URL for the API
26
- API_BASE_URL = "http://127.0.0.1:7860/api" # Adjust if needed
27
 
28
- # --- FastAPI Lifespan Event ---
29
  @asynccontextmanager
30
  async def lifespan(app: FastAPI):
31
- # Startup
32
  logger.info("Application startup: Connecting DB...")
33
  await connect_db()
34
  logger.info("Application startup: DB Connected. Checking/Creating tables...")
35
- if database.is_connected: # Proceed only if connection succeeded
36
  try:
37
- # 1. Check if table exists using the async connection
38
- # For SQLite, check the sqlite_master table
39
  check_query = "SELECT name FROM sqlite_master WHERE type='table' AND name=:table_name;"
40
  table_exists = await database.fetch_one(query=check_query, values={"table_name": users.name})
41
-
42
  if not table_exists:
43
  logger.info(f"Table '{users.name}' not found, attempting creation using async connection...")
44
-
45
- # 2. Generate the CREATE TABLE statement using SQLAlchemy DDL compiler
46
- # We need a dialect object for the compiler
47
- # Infer dialect from the database URL if possible, default to SQLite for this example
48
- dialect = sqlite.dialect() # Or determine dynamically based on database.url
49
  create_table_stmt = str(CreateTable(users).compile(dialect=dialect))
50
- logger.debug(f"Generated CREATE TABLE statement: {create_table_stmt}")
51
-
52
- # 3. Execute the CREATE TABLE statement via the async database connection
53
  await database.execute(query=create_table_stmt)
54
  logger.info(f"Table '{users.name}' created successfully via async connection.")
55
-
56
- # 4. Optional: Verify again immediately
57
  table_exists_after = await database.fetch_one(query=check_query, values={"table_name": users.name})
58
- if table_exists_after:
59
- logger.info(f"Table '{users.name}' verified after creation.")
60
- else:
61
- logger.error(f"Table '{users.name}' verification FAILED after creation attempt!")
62
-
63
  else:
64
  logger.info(f"Table '{users.name}' already exists (checked via async connection).")
65
-
66
  except Exception as db_setup_err:
67
  logger.exception(f"CRITICAL error during async DB table setup: {db_setup_err}")
68
- # Consider whether to halt startup here depending on severity
69
  else:
70
  logger.error("CRITICAL: Database connection failed, skipping table setup.")
71
-
72
  logger.info("Application startup: DB setup phase complete.")
73
  yield
74
- # Shutdown
75
  logger.info("Application shutdown: Disconnecting DB...")
76
  await disconnect_db()
77
  logger.info("Application shutdown: DB Disconnected.")
78
 
79
- # Create the main FastAPI app instance that Gradio will use
80
- # We attach our API routes to this instance.
81
- app = FastAPI(lifespan=lifespan)
82
- app.include_router(api_router, prefix="/api") # Mount API routes under /api
83
-
84
- # --- Gradio UI Definition ---
85
-
86
- # Store websocket connection globally (or within a class) for the Gradio app instance
87
- # This is tricky because Gradio re-runs functions. State management is key.
88
- # We'll connect the WebSocket *after* login and store the connection task/info in gr.State.
89
 
90
- # --- Helper functions for Gradio calling the API ---
 
 
91
 
 
92
  async def make_api_request(method: str, endpoint: str, **kwargs):
93
  async with httpx.AsyncClient() as client:
94
  url = f"{API_BASE_URL}{endpoint}"
95
  try:
96
  response = await client.request(method, url, **kwargs)
97
- response.raise_for_status() # Raise exception for 4xx/5xx errors
98
  return response.json()
99
  except httpx.RequestError as e:
100
  logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}")
101
  return {"error": f"Network error contacting API: {e}"}
102
  except httpx.HTTPStatusError as e:
103
  logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}")
104
- try:
105
- detail = e.response.json().get("detail", e.response.text)
106
- except json.JSONDecodeError:
107
- detail = e.response.text
108
  return {"error": f"API Error: {detail}"}
109
  except Exception as e:
110
  logger.error(f"Unexpected error during API call: {e}")
111
  return {"error": f"An unexpected error occurred: {str(e)}"}
112
 
113
  # --- WebSocket handling within Gradio ---
114
-
115
- async def listen_to_websockets(token: str, notification_state: list):
116
- """Connects to WS and updates state list when a message arrives."""
117
  ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}"
118
  logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.")
119
 
120
  if not token:
121
  logger.warning(f"[{ws_listener_id}] No token provided. Listener task exiting.")
122
- return notification_state
 
123
 
124
  ws_url_base = API_BASE_URL.replace("http", "ws")
125
  ws_url = f"{ws_url_base}/ws/{token}"
@@ -128,76 +98,62 @@ async def listen_to_websockets(token: str, notification_state: list):
128
  try:
129
  async with websockets.connect(ws_url, open_timeout=15.0) as websocket:
130
  logger.info(f"[{ws_listener_id}] WebSocket connected successfully to {ws_url}")
131
-
132
- # --- REMOVE or COMMENT OUT this debug line ---
133
- # await asyncio.sleep(0.5)
134
- # logger.info(f"[{ws_listener_id}] Connections according to manager after connect: {manager.active_connections}")
135
- # --- End removal ---
136
-
137
  while True:
138
  try:
139
- message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0) # e.g., 5 min timeout
140
  logger.info(f"[{ws_listener_id}] Received raw message: {message_str}")
141
  try:
142
  message_data = json.loads(message_str)
143
  logger.info(f"[{ws_listener_id}] Parsed message data: {message_data}")
144
 
145
  if message_data.get("type") == "new_user":
146
- notification = schemas.Notification(**message_data)
147
- logger.info(f"[{ws_listener_id}] Processing 'new_user' notification: {notification.message}")
148
- notification_state.insert(0, notification.message)
149
- logger.info(f"[{ws_listener_id}] State list updated. New length: {len(notification_state)}. Content: {notification_state[:5]}")
150
- if len(notification_state) > 10:
151
- notification_state.pop()
 
 
 
 
 
 
 
 
152
  else:
153
  logger.warning(f"[{ws_listener_id}] Received message of unknown type: {message_data.get('type')}")
154
-
155
- except json.JSONDecodeError:
156
- logger.error(f"[{ws_listener_id}] Failed to decode JSON from WebSocket message: {message_str}")
157
- except Exception as parse_err:
158
- logger.error(f"[{ws_listener_id}] Error processing received message: {parse_err}")
159
-
160
- except asyncio.TimeoutError:
161
- logger.debug(f"[{ws_listener_id}] WebSocket recv timed out, continuing loop.")
162
- continue
163
- except websockets.ConnectionClosedOK:
164
- logger.info(f"[{ws_listener_id}] WebSocket connection closed normally.")
165
- break
166
- except websockets.ConnectionClosedError as e:
167
- logger.error(f"[{ws_listener_id}] WebSocket connection closed with error: {e}")
168
- break
169
- except Exception as e:
170
- logger.error(f"[{ws_listener_id}] Error in WebSocket listener receive loop: {e}")
171
- await asyncio.sleep(1)
172
-
173
- except asyncio.TimeoutError: # Catch timeout during initial connect
174
- logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out: {ws_url}")
175
- except websockets.exceptions.InvalidURI:
176
- logger.error(f"[{ws_listener_id}] Invalid WebSocket URI: {ws_url}")
177
- except websockets.exceptions.WebSocketException as e: # Catch connection errors
178
- logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}")
179
- except Exception as e:
180
- logger.error(f"[{ws_listener_id}] Unexpected error in WebSocket listener task: {e}")
181
 
182
  logger.info(f"[{ws_listener_id}] Listener task finished.")
183
- return notification_state
184
-
185
 
186
  # --- Gradio Interface ---
187
  with gr.Blocks(theme=gr.themes.Soft()) as demo:
188
  # State variables
189
- # Holds the session token after login
190
  auth_token = gr.State(None)
191
- # Holds user info {id, email} after login
192
  user_info = gr.State(None)
193
- # Holds the list of notification messages
194
  notification_list = gr.State([])
195
- # Holds the asyncio task for the WebSocket listener
196
  websocket_task = gr.State(None)
 
 
197
 
198
  # --- UI Components ---
199
  with gr.Tabs() as tabs:
200
- # --- Registration Tab ---
201
  with gr.TabItem("Register", id="register_tab"):
202
  gr.Markdown("## Create a new account")
203
  reg_email = gr.Textbox(label="Email", type="email")
@@ -205,8 +161,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
205
  reg_confirm_password = gr.Textbox(label="Confirm Password", type="password")
206
  reg_button = gr.Button("Register")
207
  reg_status = gr.Textbox(label="Status", interactive=False)
208
-
209
- # --- Login Tab ---
210
  with gr.TabItem("Login", id="login_tab"):
211
  gr.Markdown("## Login to your account")
212
  login_email = gr.Textbox(label="Email", type="email")
@@ -214,172 +168,143 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
214
  login_button = gr.Button("Login")
215
  login_status = gr.Textbox(label="Status", interactive=False)
216
 
217
- # --- Welcome Tab (shown after login) ---
218
  with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab:
219
  gr.Markdown("## Welcome!", elem_id="welcome_header")
220
  welcome_message = gr.Markdown("", elem_id="welcome_message")
221
  logout_button = gr.Button("Logout")
222
- gr.Markdown("---") # Separator
223
  gr.Markdown("## Real-time Notifications")
224
- # Textbox to display notifications, updated periodically
225
  notification_display = gr.Textbox(
226
  label="New User Alerts",
227
  lines=5,
228
  max_lines=10,
229
  interactive=False,
230
- # The `every=1` makes Gradio call the update function every 1 second
231
- # This function will read the `notification_list` state
232
- every=1
233
  )
 
 
234
 
235
- def update_notification_ui(notif_list_state):
236
- # notif_list_state here *is* the Python list from the gr.State object
237
- log_msg = f"UI Update Triggered. State List Length: {len(notif_list_state)}. Content: {notif_list_state[:5]}"
238
- logger.debug(log_msg) # Keep logging
239
- # --- Ensure we return an update object ---
240
- # Join the list items into a string for display
241
- new_value = "\n".join(notif_list_state)
242
- return gr.update(value=new_value) # Explicitly return gr.update()
243
-
244
- notification_display.change( # Use .change with every= setup on the component
245
- fn=update_notification_ui,
246
- inputs=[notification_list], # Read the state
247
- outputs=[notification_display] # Update the component
248
- )
249
 
250
  # --- Event Handlers ---
251
 
252
- # Registration Logic
253
  async def handle_register(email, password, confirm_password):
254
- if not email or not password or not confirm_password:
255
- return gr.update(value="Please fill in all fields.")
256
- if password != confirm_password:
257
- return gr.update(value="Passwords do not match.")
258
- if len(password) < 8:
259
- return gr.update(value="Password must be at least 8 characters long.")
260
-
261
  payload = {"email": email, "password": password}
262
  result = await make_api_request("post", "/register", json=payload)
 
 
263
 
264
- if "error" in result:
265
- return gr.update(value=f"Registration failed: {result['error']}")
266
- else:
267
- # Optionally switch to login tab after successful registration
268
- return gr.update(value=f"Registration successful for {result.get('email')}! Please log in.")
269
-
270
- reg_button.click(
271
- handle_register,
272
- inputs=[reg_email, reg_password, reg_confirm_password],
273
- outputs=[reg_status]
274
- )
275
 
276
  # Login Logic
277
- async def handle_login(email, password, current_task):
278
- if not email or not password:
279
- return gr.update(value="Please enter email and password."), None, None, None, gr.update(visible=False), current_task
 
 
 
 
 
 
 
 
280
 
281
  payload = {"email": email, "password": password}
282
  result = await make_api_request("post", "/login", json=payload)
283
 
284
  if "error" in result:
285
- return gr.update(value=f"Login failed: {result['error']}"), None, None, None, gr.update(visible=False), current_task
 
286
  else:
287
  token = result.get("access_token")
288
- # Fetch user details using the token
289
- user_data = await dependencies.get_optional_current_user(token) # Use dependency directly
290
-
291
  if not user_data:
292
- # This shouldn't happen if login succeeded, but check anyway
293
- return gr.update(value="Login succeeded but failed to fetch user data."), None, None, None, gr.update(visible=False), current_task
294
 
295
- # Cancel any existing websocket listener task before starting a new one
296
  if current_task and not current_task.done():
297
  current_task.cancel()
298
- try:
299
- await current_task # Wait for cancellation
300
- except asyncio.CancelledError:
301
- logger.info("Previous WebSocket task cancelled.")
302
 
303
- # Start the WebSocket listener task in the background
304
- # We pass the notification_list state *object* itself, which the task will modify
305
- new_task = asyncio.create_task(listen_to_websockets(token, notification_list.value)) # Pass the list
306
 
307
- # Update state and UI
308
  welcome_msg = f"Welcome, {user_data.email}!"
309
- # Switch tabs and show welcome message
310
  return (
311
  gr.update(value="Login successful!"), # login_status
312
  token, # auth_token state
313
- user_data.model_dump(), # user_info state (store as dict)
314
  gr.update(selected="welcome_tab"), # Switch Tabs
315
  gr.update(visible=True), # Make welcome tab visible
316
  gr.update(value=welcome_msg), # Update welcome message markdown
317
- new_task # websocket_task state
 
318
  )
319
 
 
320
  login_button.click(
321
  handle_login,
322
- inputs=[login_email, login_password, websocket_task],
323
- outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task]
324
  )
325
 
326
-
327
- # Function to update the notification display based on the state
328
- # This function is triggered by the `every=1` on the notification_display Textbox
329
- def update_notification_ui(notif_list_state):
330
- # Join the list items into a string for display
331
- return "\n".join(notif_list_state)
332
-
333
- notification_display.change( # Use .change with every= setup on the component
 
 
 
 
 
334
  fn=update_notification_ui,
335
- inputs=[notification_list], # Read the state
336
- outputs=[notification_display] # Update the component
337
  )
338
 
 
 
 
 
339
 
340
  # Logout Logic
 
341
  async def handle_logout(current_task):
342
- # Cancel the websocket listener task if it's running
343
  if current_task and not current_task.done():
344
  current_task.cancel()
345
- try:
346
- await current_task
347
- except asyncio.CancelledError:
348
- logger.info("WebSocket task cancelled on logout.")
349
-
350
- # Clear state and switch back to login tab
351
- return (
352
- None, # Clear auth_token
353
- None, # Clear user_info
354
- [], # Clear notifications
355
- None, # Clear websocket_task
356
- gr.update(selected="login_tab"),# Switch Tabs
357
- gr.update(visible=False), # Hide welcome tab
358
- gr.update(value=""), # Clear welcome message
359
- gr.update(value="") # Clear login status
360
- )
361
-
362
  logout_button.click(
363
  handle_logout,
364
  inputs=[websocket_task],
365
  outputs=[
366
- auth_token,
367
- user_info,
368
- notification_list,
369
- websocket_task,
370
- tabs,
371
- welcome_tab,
372
- welcome_message,
373
- login_status
374
  ]
375
  )
376
 
377
- # Mount the Gradio app onto the FastAPI app at the root
378
  app = gr.mount_gradio_app(app, demo, path="/")
379
 
380
- # If running this file directly (for local testing)
381
- # Use uvicorn to run the FastAPI app (which now includes Gradio)
382
  if __name__ == "__main__":
383
  import uvicorn
384
- # Use port 7860 as Gradio prefers, host 0.0.0.0 for Docker
385
- uvicorn.run(app, host="0.0.0.0", port=7860)
 
9
  from contextlib import asynccontextmanager
10
 
11
  from fastapi import FastAPI, Depends # Import FastAPI itself
 
12
  from .database import connect_db, disconnect_db, database, metadata, users
13
  from .api import router as api_router
14
  from . import schemas, auth, dependencies
15
+ from .websocket import manager # Import the connection manager instance
16
+
17
  from sqlalchemy.schema import CreateTable
18
+ from sqlalchemy.dialects import sqlite
19
 
 
20
  logging.basicConfig(level=logging.INFO)
21
  logger = logging.getLogger(__name__)
22
 
23
+ API_BASE_URL = "http://127.0.0.1:7860/api"
 
24
 
25
+ # --- Lifespan (remains the same) ---
26
  @asynccontextmanager
27
  async def lifespan(app: FastAPI):
28
+ # ... (same DB setup code) ...
29
  logger.info("Application startup: Connecting DB...")
30
  await connect_db()
31
  logger.info("Application startup: DB Connected. Checking/Creating tables...")
32
+ if database.is_connected:
33
  try:
 
 
34
  check_query = "SELECT name FROM sqlite_master WHERE type='table' AND name=:table_name;"
35
  table_exists = await database.fetch_one(query=check_query, values={"table_name": users.name})
 
36
  if not table_exists:
37
  logger.info(f"Table '{users.name}' not found, attempting creation using async connection...")
38
+ dialect = sqlite.dialect()
 
 
 
 
39
  create_table_stmt = str(CreateTable(users).compile(dialect=dialect))
 
 
 
40
  await database.execute(query=create_table_stmt)
41
  logger.info(f"Table '{users.name}' created successfully via async connection.")
 
 
42
  table_exists_after = await database.fetch_one(query=check_query, values={"table_name": users.name})
43
+ if table_exists_after: logger.info(f"Table '{users.name}' verified after creation.")
44
+ else: logger.error(f"Table '{users.name}' verification FAILED after creation attempt!")
 
 
 
45
  else:
46
  logger.info(f"Table '{users.name}' already exists (checked via async connection).")
 
47
  except Exception as db_setup_err:
48
  logger.exception(f"CRITICAL error during async DB table setup: {db_setup_err}")
 
49
  else:
50
  logger.error("CRITICAL: Database connection failed, skipping table setup.")
 
51
  logger.info("Application startup: DB setup phase complete.")
52
  yield
 
53
  logger.info("Application shutdown: Disconnecting DB...")
54
  await disconnect_db()
55
  logger.info("Application shutdown: DB Disconnected.")
56
 
 
 
 
 
 
 
 
 
 
 
57
 
58
+ # --- FastAPI App Setup (remains the same) ---
59
+ app = FastAPI(lifespan=lifespan)
60
+ app.include_router(api_router, prefix="/api")
61
 
62
+ # --- Helper functions (make_api_request remains the same) ---
63
  async def make_api_request(method: str, endpoint: str, **kwargs):
64
  async with httpx.AsyncClient() as client:
65
  url = f"{API_BASE_URL}{endpoint}"
66
  try:
67
  response = await client.request(method, url, **kwargs)
68
+ response.raise_for_status()
69
  return response.json()
70
  except httpx.RequestError as e:
71
  logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}")
72
  return {"error": f"Network error contacting API: {e}"}
73
  except httpx.HTTPStatusError as e:
74
  logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}")
75
+ try: detail = e.response.json().get("detail", e.response.text)
76
+ except json.JSONDecodeError: detail = e.response.text
 
 
77
  return {"error": f"API Error: {detail}"}
78
  except Exception as e:
79
  logger.error(f"Unexpected error during API call: {e}")
80
  return {"error": f"An unexpected error occurred: {str(e)}"}
81
 
82
  # --- WebSocket handling within Gradio ---
83
+ # <<< MODIFIED: Accept trigger state object >>>
84
+ async def listen_to_websockets(token: str, notification_list_state: gr.State, notification_trigger_state: gr.State):
85
+ """Connects to WS and updates state list and trigger when a message arrives."""
86
  ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}"
87
  logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.")
88
 
89
  if not token:
90
  logger.warning(f"[{ws_listener_id}] No token provided. Listener task exiting.")
91
+ # <<< Return original state values >>>
92
+ return notification_list_state.value, notification_trigger_state.value
93
 
94
  ws_url_base = API_BASE_URL.replace("http", "ws")
95
  ws_url = f"{ws_url_base}/ws/{token}"
 
98
  try:
99
  async with websockets.connect(ws_url, open_timeout=15.0) as websocket:
100
  logger.info(f"[{ws_listener_id}] WebSocket connected successfully to {ws_url}")
 
 
 
 
 
 
101
  while True:
102
  try:
103
+ message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0)
104
  logger.info(f"[{ws_listener_id}] Received raw message: {message_str}")
105
  try:
106
  message_data = json.loads(message_str)
107
  logger.info(f"[{ws_listener_id}] Parsed message data: {message_data}")
108
 
109
  if message_data.get("type") == "new_user":
110
+ notification = schemas.Notification(**message_data)
111
+ logger.info(f"[{ws_listener_id}] Processing 'new_user' notification: {notification.message}")
112
+
113
+ # <<< Modify state objects' values >>>
114
+ current_list = notification_list_state.value.copy()
115
+ current_list.insert(0, notification.message)
116
+ if len(current_list) > 10: current_list.pop()
117
+ notification_list_state.value = current_list
118
+ logger.info(f"[{ws_listener_id}] State list updated via state object. New length: {len(notification_list_state.value)}. Content: {notification_list_state.value[:5]}")
119
+
120
+ # <<< Update trigger state object's value >>>
121
+ notification_trigger_state.value += 1
122
+ logger.info(f"[{ws_listener_id}] Incremented notification trigger to {notification_trigger_state.value}")
123
+
124
  else:
125
  logger.warning(f"[{ws_listener_id}] Received message of unknown type: {message_data.get('type')}")
126
+ # ... (error handling for parsing) ...
127
+ except json.JSONDecodeError: logger.error(f"[{ws_listener_id}] Failed to decode JSON: {message_str}")
128
+ except Exception as parse_err: logger.error(f"[{ws_listener_id}] Error processing message: {parse_err}")
129
+ # ... (error handling for websocket recv/connection) ...
130
+ except asyncio.TimeoutError: logger.debug(f"[{ws_listener_id}] WebSocket recv timed out."); continue
131
+ except websockets.ConnectionClosedOK: logger.info(f"[{ws_listener_id}] WebSocket connection closed normally."); break
132
+ except websockets.ConnectionClosedError as e: logger.error(f"[{ws_listener_id}] WebSocket connection closed with error: {e}"); break
133
+ except Exception as e: logger.error(f"[{ws_listener_id}] Error in listener receive loop: {e}"); await asyncio.sleep(1); break # Break on unknown errors too
134
+ # ... (error handling for websocket connect) ...
135
+ except asyncio.TimeoutError: logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out: {ws_url}")
136
+ except websockets.exceptions.InvalidURI: logger.error(f"[{ws_listener_id}] Invalid WebSocket URI: {ws_url}")
137
+ except websockets.exceptions.WebSocketException as e: logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}")
138
+ except Exception as e: logger.error(f"[{ws_listener_id}] Unexpected error in WebSocket listener task: {e}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
139
 
140
  logger.info(f"[{ws_listener_id}] Listener task finished.")
141
+ # <<< Return original state values >>>
142
+ return notification_list_state.value, notification_trigger_state.value
143
 
144
  # --- Gradio Interface ---
145
  with gr.Blocks(theme=gr.themes.Soft()) as demo:
146
  # State variables
 
147
  auth_token = gr.State(None)
 
148
  user_info = gr.State(None)
 
149
  notification_list = gr.State([])
 
150
  websocket_task = gr.State(None)
151
+ # <<< Add Dummy State >>>
152
+ notification_trigger = gr.State(0) # Simple counter
153
 
154
  # --- UI Components ---
155
  with gr.Tabs() as tabs:
156
+ # --- Registration/Login Tabs (remain the same) ---
157
  with gr.TabItem("Register", id="register_tab"):
158
  gr.Markdown("## Create a new account")
159
  reg_email = gr.Textbox(label="Email", type="email")
 
161
  reg_confirm_password = gr.Textbox(label="Confirm Password", type="password")
162
  reg_button = gr.Button("Register")
163
  reg_status = gr.Textbox(label="Status", interactive=False)
 
 
164
  with gr.TabItem("Login", id="login_tab"):
165
  gr.Markdown("## Login to your account")
166
  login_email = gr.Textbox(label="Email", type="email")
 
168
  login_button = gr.Button("Login")
169
  login_status = gr.Textbox(label="Status", interactive=False)
170
 
171
+ # --- Welcome Tab ---
172
  with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab:
173
  gr.Markdown("## Welcome!", elem_id="welcome_header")
174
  welcome_message = gr.Markdown("", elem_id="welcome_message")
175
  logout_button = gr.Button("Logout")
176
+ gr.Markdown("---")
177
  gr.Markdown("## Real-time Notifications")
 
178
  notification_display = gr.Textbox(
179
  label="New User Alerts",
180
  lines=5,
181
  max_lines=10,
182
  interactive=False,
183
+ # <<< REMOVE every=1 >>>
 
 
184
  )
185
+ # <<< Add Dummy Component (Hidden) >>>
186
+ dummy_trigger_output = gr.Textbox(label="trigger", visible=False)
187
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
188
 
189
  # --- Event Handlers ---
190
 
191
+ # Registration Logic (remains the same)
192
  async def handle_register(email, password, confirm_password):
193
+ if not email or not password or not confirm_password: return gr.update(value="Please fill in all fields.")
194
+ if password != confirm_password: return gr.update(value="Passwords do not match.")
195
+ if len(password) < 8: return gr.update(value="Password must be at least 8 characters long.")
 
 
 
 
196
  payload = {"email": email, "password": password}
197
  result = await make_api_request("post", "/register", json=payload)
198
+ if "error" in result: return gr.update(value=f"Registration failed: {result['error']}")
199
+ else: return gr.update(value=f"Registration successful for {result.get('email')}! Please log in.")
200
 
201
+ reg_button.click(handle_register, inputs=[reg_email, reg_password, reg_confirm_password], outputs=[reg_status])
 
 
 
 
 
 
 
 
 
 
202
 
203
  # Login Logic
204
+ # <<< MODIFIED: Pass trigger state, update outputs >>>
205
+ async def handle_login(email, password, current_task, current_trigger_val):
206
+ # --- Output state needs to align with return values ---
207
+ # Returns: login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger
208
+ outputs_tuple = (
209
+ gr.update(value="Please enter email and password."), # login_status
210
+ None, None, None, gr.update(visible=False), None, # auth_token, user_info, tabs, welcome_tab, welcome_message
211
+ current_task, # websocket_task (no change)
212
+ current_trigger_val # notification_trigger (no change)
213
+ )
214
+ if not email or not password: return outputs_tuple
215
 
216
  payload = {"email": email, "password": password}
217
  result = await make_api_request("post", "/login", json=payload)
218
 
219
  if "error" in result:
220
+ outputs_tuple = (gr.update(value=f"Login failed: {result['error']}"), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val)
221
+ return outputs_tuple
222
  else:
223
  token = result.get("access_token")
224
+ user_data = await dependencies.get_optional_current_user(token)
 
 
225
  if not user_data:
226
+ outputs_tuple = (gr.update(value="Login succeeded but failed to fetch user data."), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val)
227
+ return outputs_tuple
228
 
 
229
  if current_task and not current_task.done():
230
  current_task.cancel()
231
+ try: await current_task
232
+ except asyncio.CancelledError: logger.info("Previous WebSocket task cancelled.")
 
 
233
 
234
+ # <<< Pass both state objects to listener >>>
235
+ new_task = asyncio.create_task(listen_to_websockets(token, notification_list, notification_trigger))
 
236
 
 
237
  welcome_msg = f"Welcome, {user_data.email}!"
238
+ # --- Ensure number of return values matches outputs ---
239
  return (
240
  gr.update(value="Login successful!"), # login_status
241
  token, # auth_token state
242
+ user_data.model_dump(), # user_info state
243
  gr.update(selected="welcome_tab"), # Switch Tabs
244
  gr.update(visible=True), # Make welcome tab visible
245
  gr.update(value=welcome_msg), # Update welcome message markdown
246
+ new_task, # websocket_task state
247
+ 0 # Reset notification_trigger state to 0 on login
248
  )
249
 
250
+ # <<< MODIFIED: Add notification_trigger to inputs/outputs >>>
251
  login_button.click(
252
  handle_login,
253
+ inputs=[login_email, login_password, websocket_task, notification_trigger],
254
+ outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger]
255
  )
256
 
257
+ # Function to update the notification display based on the list state
258
+ # <<< Triggered by dummy component now >>>
259
+ def update_notification_ui(notif_list): # Takes the list value directly now
260
+ log_msg = f"UI Update Triggered via Dummy. State List Length: {len(notif_list)}. Content: {notif_list[:5]}"
261
+ logger.info(log_msg) # Use info level to ensure visibility
262
+ new_value = "\n".join(notif_list)
263
+ return gr.update(value=new_value)
264
+
265
+ # <<< Add Event handler for the dummy trigger >>>
266
+ # When the dummy_trigger_output *would* change (because notification_trigger state changed)...
267
+ # ...call the update_notification_ui function.
268
+ # Pass the *current* value of notification_list state as input to the function.
269
+ dummy_trigger_output.change(
270
  fn=update_notification_ui,
271
+ inputs=[notification_list], # Input is the list state we want to display
272
+ outputs=[notification_display] # Output updates the real display textbox
273
  )
274
 
275
+ # <<< Link the dummy trigger state to the dummy output component >>>
276
+ # This makes the dummy_trigger_output.change event fire when notification_trigger changes
277
+ notification_trigger.change(lambda x: x, inputs=notification_trigger, outputs=dummy_trigger_output, queue=False) # Use queue=False for immediate trigger if possible
278
+
279
 
280
  # Logout Logic
281
+ # <<< MODIFIED: Add notification_trigger to outputs >>>
282
  async def handle_logout(current_task):
 
283
  if current_task and not current_task.done():
284
  current_task.cancel()
285
+ try: await current_task
286
+ except asyncio.CancelledError: logger.info("WebSocket task cancelled on logout.")
287
+ # --- Ensure number of return values matches outputs ---
288
+ return ( None, None, [], None, # auth_token, user_info, notification_list, websocket_task
289
+ gr.update(selected="login_tab"), gr.update(visible=False), # tabs, welcome_tab
290
+ gr.update(value=""), gr.update(value=""), # welcome_message, login_status
291
+ 0 ) # notification_trigger (reset to 0)
292
+
293
+ # <<< MODIFIED: Add notification_trigger to outputs >>>
 
 
 
 
 
 
 
 
294
  logout_button.click(
295
  handle_logout,
296
  inputs=[websocket_task],
297
  outputs=[
298
+ auth_token, user_info, notification_list, websocket_task,
299
+ tabs, welcome_tab, welcome_message, login_status,
300
+ notification_trigger # Add trigger to outputs
 
 
 
 
 
301
  ]
302
  )
303
 
304
+ # Mount Gradio App (remains the same)
305
  app = gr.mount_gradio_app(app, demo, path="/")
306
 
307
+ # Run Uvicorn (remains the same)
 
308
  if __name__ == "__main__":
309
  import uvicorn
310
+ uvicorn.run("app.main:app", host="0.0.0.0", port=7860, reload=True)