Update app.py

app.py

@@ -4,23 +4,29 @@ from datetime import datetime
 from collections import deque
 
 # --- GLOBAL SERVER STATE (In-Memory Only) ---
-# We use a simple dictionary to store the last 50 messages for active channels.
-# When the Space restarts or sleeps, all data is wiped.
 class ChatServer:
     def __init__(self):
-        # channels = { "channel_name": deque([messages], maxlen=50) }
         self.channels = {"Main": deque(maxlen=50)}
-        # active_users = { "username": {"last_seen": timestamp, "profile_pic": url} }
         self.active_users = {}
-        # dms = { "user1-user2": deque(maxlen=50) }
         self.dms = {}
+        # Track which users are in which room to delete empty rooms
+        self.user_rooms = {} 
 
-    def prune_users(self):
-        """Remove users who haven't polled in 30 seconds."""
+    def prune_server(self):
+        """Remove inactive users and delete empty custom rooms."""
         now = time.time()
         expired = [u for u, data in self.active_users.items() if now - data['last_seen'] > 30]
         for u in expired:
-            del self.active_users[u]
+            if u in self.active_users:
+                del self.active_users[u]
+            if u in self.user_rooms:
+                del self.user_rooms[u]
+
+        # Delete rooms with no active occupants (except Main)
+        active_rooms = set(self.user_rooms.values())
+        rooms_to_delete = [r for r in self.channels.keys() if r != "Main" and r not in active_rooms]
+        for r in rooms_to_delete:
+            del self.channels[r]
 
     def broadcast(self, channel, user, message, profile_pic):
         if channel not in self.channels:
@@ -70,7 +76,7 @@ def get_messages_html(messages):
     return html
 
 def get_active_users_html(current_user):
-    server.prune_users()
+    server.prune_server()
     html = "<div style='display: flex; flex-direction: column; gap: 5px;'>"
     for user, data in server.active_users.items():
         status = "🟢" if user != current_user else "👤 (You)"
@@ -84,15 +90,16 @@ def get_active_users_html(current_user):
     html += "</div>"
     return html
 
+# --- CSS STYLES ---
+custom_css = """
+#chat-container { height: 500px; overflow-y: auto; border: 1px solid #ddd; padding: 15px; border-radius: 8px; background: var(--body-background-fill); }
+.user-list { border-left: 1px solid #ddd; padding-left: 15px; }
+#msg-input input { border-radius: 20px !important; }
+"""
+
 # --- GRADIO INTERFACE ---
 
-with gr.Blocks(css="""
-    #chat-container { height: 500px; overflow-y: auto; border: 1px solid #ddd; padding: 15px; border-radius: 8px; background: var(--body-background-fill); }
-    .user-list { border-left: 1px solid #ddd; padding-left: 15px; }
-    #msg-input input { border-radius: 20px !important; }
-""") as demo:
-    
-    # HF OAuth session data
+with gr.Blocks() as demo:
     user_session = gr.State(None)
     current_room = gr.State("Main")
     dm_target = gr.State(None)
@@ -104,7 +111,6 @@ with gr.Blocks(css="""
             login_btn = gr.LoginButton()
 
     with gr.Row(visible=False) as chat_ui:
-        # Sidebar
         with gr.Column(scale=2, variant="panel"):
             gr.Markdown("### 🏠 Rooms")
             room_list = gr.Radio(["Main"], value="Main", label="Join a Channel")
@@ -118,7 +124,6 @@ with gr.Blocks(css="""
             gr.Markdown("### 👥 Online Now")
             active_users_box = gr.HTML()
 
-        # Main Chat Area
         with gr.Column(scale=8):
             room_title = gr.Markdown("## # Main")
             chat_display = gr.HTML(elem_id="chat-container")
@@ -132,8 +137,6 @@ with gr.Blocks(css="""
                 )
                 send_btn = gr.Button("Send", scale=1, variant="primary")
 
-    # --- LOGIC ---
-
     def on_load(profile: gr.OAuthProfile | None):
         if profile is None:
             return gr.update(visible=False), None
@@ -141,17 +144,20 @@ with gr.Blocks(css="""
         username = profile.username
         pic = profile.profile_image or "https://huggingface.co/front/assets/huggingface_logo-noborder.svg"
         
-        # Register user
         server.active_users[username] = {"last_seen": time.time(), "pic": pic}
+        server.user_rooms[username] = "Main"
         return gr.update(visible=True), {"name": username, "pic": pic}
 
     demo.load(on_load, None, [chat_ui, user_session])
 
     def refresh_chat(session, room, dm_user):
-        if not session: return "", ""
+        if not session: return "", "", ""
         
-        # Heartbeat
         server.active_users[session['name']]['last_seen'] = time.time()
+        server.user_rooms[session['name']] = room if not dm_user else "DM"
+        
+        # Update room list options dynamically
+        available_rooms = list(server.channels.keys())
         
         if dm_user:
             dm_key = "-".join(sorted([session['name'], dm_user]))
@@ -168,7 +174,6 @@ with gr.Blocks(css="""
 
     def send_msg(text, session, room, dm_user):
         if not text or not session: return ""
-        
         if dm_user:
             server.send_dm(session['name'], dm_user, text, session['pic'])
         else:
@@ -179,9 +184,11 @@ with gr.Blocks(css="""
     msg_input.submit(send_msg, [msg_input, user_session, current_room, dm_target], [msg_input])
 
     def create_and_join(name, current_options):
-        if not name: return gr.update(), "Main"
+        if not name: return gr.update(), "Main", None
         clean_name = name.strip().replace(" ", "-")
-        new_options = list(set(current_options + [clean_name]))
+        if clean_name not in server.channels:
+            server.channels[clean_name] = deque(maxlen=50)
+        new_options = list(server.channels.keys())
         return gr.update(choices=new_options, value=clean_name), clean_name, None
 
     create_room_btn.click(create_and_join, [new_room_name, room_list], [room_list, current_room, dm_target])
@@ -197,4 +204,4 @@ with gr.Blocks(css="""
     
     start_dm_btn.click(start_dm, dm_user_input, [current_room, dm_target])
 
-demo.launch()
+demo.launch(css=custom_css)