Update app.py

app.py

@@ -94,13 +94,18 @@ class WebhookHandler(BaseHTTPRequestHandler):
         if event == 'issues' and WebhookHandler.manager_instance:
             action = payload.get('action')
             logger.info(f"Issue action: {action}")
-            if action in ['opened', 'reopened', 'closed', 'assigned']:
+            if action in ['opened', 'reopened', 'closed', 'assigned', 'edited']: # Handle edited issues too
                 # Ensure the event loop is running in the webhook thread if needed
-                loop = asyncio.get_event_loop()
-                asyncio.run_coroutine_threadsafe(
-                    WebhookHandler.manager_instance.handle_webhook_event(event, action, payload),
-                    loop
-                )
+                # Get the loop associated with the asyncio thread
+                loop = asyncio.get_event_loop_policy().get_event_loop()
+                if loop.is_running():
+                    asyncio.run_coroutine_threadsafe(
+                        WebhookHandler.manager_instance.handle_webhook_event(event, action, payload),
+                        loop
+                    )
+                else:
+                     logger.error("Asyncio event loop is not running in the target thread for webhook.")
+
         elif event == 'ping':
              logger.info("Received GitHub webhook ping.")
         else:
@@ -130,8 +135,9 @@ class IssueManager:
         self.issue_clusters: Dict[int, List[int]] = {}  # Store clusters: {cluster_id: [issue_index1, issue_index2]}
         self.issue_list_for_clustering: List[dict] = [] # Store issues in list order for clustering index mapping
         # self._init_local_models() # Consider lazy loading or conditional loading
-        self.ws_clients: List[websockets.WebSocketClientProtocol] = []
+        self.ws_clients: List[websockets.WebSocketServerProtocol] = [] # Use WebSocketServerProtocol
         self.code_editors: Dict[int, OTCodeEditor] = {}  # Store code editors for each issue
+        self.main_loop = asyncio.get_event_loop() # Store ref to main loop if needed elsewhere
 
     # Placeholder for local model initialization - implement actual loading if needed
     def _init_local_models(self):
@@ -146,7 +152,7 @@ class IssueManager:
         return hashlib.md5(content.encode()).hexdigest()
 
     @lru_cache(maxsize=100)
-    async def cached_suggestion(self, issue_hash: str, model: str) -> str:
+    async def cached_suggestion(self, issue_hash: str, model_key: str) -> str:
         # Find the issue corresponding to the hash (inefficient, improve if needed)
         found_issue = None
         for issue in self.issues.values():
@@ -155,9 +161,13 @@ class IssueManager:
                 break
         if not found_issue:
             return "Error: Issue not found for the given hash."
+        if model_key not in HF_MODELS:
+             return f"Error: Invalid model key: {model_key}"
+
+        logger.info(f"Cache miss or first request for issue hash {issue_hash}. Requesting suggestion from {model_key}.")
+        # Pass the actual issue dict and model_key to suggest_resolution
+        return await self.suggest_resolution(found_issue, model_key)
 
-        logger.info(f"Cache miss or first request for issue hash {issue_hash}. Requesting suggestion from {model}.")
-        return await self.suggest_resolution(found_issue, model)
 
     async def handle_webhook_event(self, event: str, action: str, payload: dict):
         logger.info(f"Processing webhook event: {event}, action: {action}")
@@ -171,31 +181,53 @@ class IssueManager:
             logger.warning("Webhook issue data missing 'number'.")
             return
 
+        needs_update = False
         if action == 'closed':
             logger.info(f"Removing closed issue {issue_number} from active list.")
-            self.issues.pop(issue_number, None)
+            if self.issues.pop(issue_number, None):
+                 needs_update = True
             # Optionally remove associated editor, etc.
             self.code_editors.pop(issue_number, None)
         elif action in ['opened', 'reopened', 'edited']: # Handle edited issues too
              logger.info(f"Adding/Updating issue {issue_number} from webhook.")
-             self.issues[issue_number] = issue_data
+             self.issues[issue_number] = self._process_issue_data(issue_data) # Use helper
+             needs_update = True
              # Potentially trigger re-clustering or update specific issue details
         elif action == 'assigned':
              logger.info(f"Issue {issue_number} assigned to {payload.get('assignee', {}).get('login', 'N/A')}")
-             self.issues[issue_number] = issue_data # Update issue data
+             if issue_number in self.issues:
+                 self.issues[issue_number] = self._process_issue_data(issue_data) # Update issue data
+                 needs_update = True
+             else: # Issue might not be in our list if it wasn't open initially
+                 self.issues[issue_number] = self._process_issue_data(issue_data)
+                 needs_update = True
+
         else:
              logger.info(f"Ignoring action '{action}' for issue {issue_number}.")
 
-        # Consider triggering a UI update after handling the webhook
-        # This might involve re-crawling or just updating the specific issue
-        await self.broadcast_issue_update() # Example function to notify clients
+        # Trigger a UI update notification via WebSocket if data changed
+        if needs_update:
+            await self.broadcast_issue_update()
+
+    def _process_issue_data(self, issue_data: dict) -> dict:
+         """Helper to structure issue data consistently."""
+         return {
+             "id": issue_data['number'],
+             "title": issue_data.get('title', 'No Title'),
+             "body": issue_data.get('body', ''),
+             "state": issue_data.get('state', 'unknown'),
+             "labels": [label['name'] for label in issue_data.get('labels', [])],
+             "assignee": issue_data.get('assignee', {}).get('login') if issue_data.get('assignee') else None,
+             "url": issue_data.get('html_url', '#')
+             # Add other relevant fields if needed
+         }
 
     async def crawl_issues(self, repo_url: str, github_token: str, hf_token: str) -> Tuple[List[List], go.Figure, str]:
         """
         Crawls issues, updates internal state, performs clustering, and returns data for UI update.
         """
-        if not repo_url or not github_token or not hf_token:
-             return [], go.Figure(), "Error: Repository URL, GitHub Token, and HF Token are required."
+        if not repo_url or not hf_token: # GitHub token is optional for public repos
+             return [], go.Figure(), "Error: Repository URL and HF Token are required."
 
         logger.info(f"Starting issue crawl for {repo_url}")
         self.repo_url = repo_url
@@ -209,31 +241,39 @@ class IssueManager:
             logger.error(f"Invalid GitHub URL format: {repo_url}")
             return [], go.Figure(), "Error: Invalid GitHub URL format. Use https://github.com/owner/repo"
         owner, repo_name = match.groups()
-        api_url = f"{GITHUB_API}/{owner}/{repo_name}/issues?state=open" # Fetch only open issues
+        api_url = f"{GITHUB_API}/{owner}/{repo_name}/issues?state=open&per_page=100" # Fetch only open issues, max 100
 
         headers = {
-            "Authorization": f"token {github_token}",
             "Accept": "application/vnd.github.v3+json"
         }
+        if github_token:
+             headers["Authorization"] = f"token {github_token}"
+
 
         try:
+            all_issues_data = []
+            page = 1
+            logger.info(f"Fetching issues from {api_url}...")
             async with aiohttp.ClientSession(headers=headers) as session:
-                async with session.get(api_url) as response:
-                    response.raise_for_status() # Raise exception for bad status codes
-                    issues_data = await response.json()
-                    logger.info(f"Fetched {len(issues_data)} open issues.")
-                    for issue in issues_data:
-                        issue_number = issue['number']
-                        self.issues[issue_number] = {
-                            "id": issue_number,
-                            "title": issue.get('title', 'No Title'),
-                            "body": issue.get('body', ''),
-                            "state": issue.get('state', 'unknown'),
-                            "labels": [label['name'] for label in issue.get('labels', [])],
-                            "assignee": issue.get('assignee', {}).get('login') if issue.get('assignee') else None,
-                            "url": issue.get('html_url', '#')
-                            # Add other relevant fields if needed
-                        }
+                 while True:
+                    paginated_url = f"{api_url}&page={page}"
+                    async with session.get(paginated_url) as response:
+                        response.raise_for_status() # Raise exception for bad status codes
+                        issues_page_data = await response.json()
+                        if not issues_page_data: # No more issues on this page
+                            break
+                        logger.info(f"Fetched page {page} with {len(issues_page_data)} issues.")
+                        all_issues_data.extend(issues_page_data)
+                        # Check Link header for next page (more robust pagination)
+                        if 'next' not in response.headers.get('Link', ''):
+                            break
+                        page += 1
+
+
+            logger.info(f"Total issues fetched: {len(all_issues_data)}")
+            for issue_data in all_issues_data:
+                issue_number = issue_data['number']
+                self.issues[issue_number] = self._process_issue_data(issue_data) # Use helper
 
             if not self.issues:
                  logger.info("No open issues found.")
@@ -279,10 +319,12 @@ class IssueManager:
                  error_msg = f"Error: Repository not found at {repo_url}. Check the URL."
              elif e.status == 401:
                  error_msg = "Error: Invalid GitHub token or insufficient permissions."
+             elif e.status == 403:
+                  error_msg = "Error: GitHub API rate limit exceeded or forbidden access. Try adding a GitHub token."
              return [], go.Figure(), error_msg
         except GitCommandError as e:
-            logger.error(f"Git clone error: {e}")
-            return [], go.Figure(), f"Error cloning repository: {e}"
+            logger.error(f"Git clone error: {e}") # Should not happen if not cloning
+            return [], go.Figure(), f"Error related to Git: {e}"
         except Exception as e:
             logger.exception(f"An unexpected error occurred during issue crawl: {e}") # Log full traceback
             return [], go.Figure(), f"An unexpected error occurred: {e}"
@@ -297,8 +339,29 @@ class IssueManager:
 
     def _generate_stats_plot(self, severity_counts: Dict[str, int]) -> go.Figure:
         """Generates a Plotly bar chart for issue severity distribution."""
-        severities = list(severity_counts.keys())
-        counts = list(severity_counts.values())
+        # Filter out severities with 0 counts
+        filtered_counts = {k: v for k, v in severity_counts.items() if v > 0}
+        if not filtered_counts:
+             # Return an empty figure with a message if no issues found
+             fig = go.Figure()
+             fig.update_layout(
+                 title="Issue Severity Distribution",
+                 xaxis = {"visible": False},
+                 yaxis = {"visible": False},
+                 annotations = [{
+                     "text": "No issues found to display statistics.",
+                     "xref": "paper",
+                     "yref": "paper",
+                     "showarrow": False,
+                     "font": {"size": 16}
+                 }],
+                 plot_bgcolor='rgba(0,0,0,0)',
+                 paper_bgcolor='rgba(0,0,0,0)'
+             )
+             return fig
+
+        severities = list(filtered_counts.keys())
+        counts = list(filtered_counts.values())
 
         fig = px.bar(
             x=severities,
@@ -312,15 +375,18 @@ class IssueManager:
                  'Medium': '#FACC15',   # Yellow
                  'Low': '#84CC16',      # Lime
                  'Unknown': '#6B7280'   # Gray
-            }
+            },
+            text=counts # Display counts on bars
         )
         fig.update_layout(
-            showlegend=False, # Hide legend if coloring by severity directly
+            xaxis_title=None, # Cleaner look
             yaxis_title="Number of Issues",
-            xaxis_title="Severity Level",
             plot_bgcolor='rgba(0,0,0,0)', # Transparent background
-            paper_bgcolor='rgba(0,0,0,0)'
+            paper_bgcolor='rgba(0,0,0,0)',
+            showlegend=False,
+            xaxis={'categoryorder':'array', 'categoryarray':['Critical', 'High', 'Medium', 'Low', 'Unknown']} # Order bars
         )
+        fig.update_traces(textposition='outside')
         return fig
 
     async def _cluster_similar_issues(self):
@@ -332,9 +398,12 @@ class IssueManager:
 
         logger.info("Generating embeddings for clustering...")
         try:
-            embeddings = await self._generate_embeddings([f"{i.get('title','')} {i.get('body','')}" for i in self.issue_list_for_clustering])
-            if not embeddings or len(embeddings) != len(self.issue_list_for_clustering):
-                 logger.error("Failed to generate valid embeddings for all issues.")
+            # Combine title and body for embedding generation
+            texts_to_embed = [f"{i.get('title','')} - {i.get('body','')[:500]}" for i in self.issue_list_for_clustering] # Limit body length
+            embeddings = await self._generate_embeddings(texts_to_embed)
+
+            if embeddings is None or not isinstance(embeddings, list) or len(embeddings) != len(self.issue_list_for_clustering):
+                 logger.error(f"Failed to generate valid embeddings. Expected {len(self.issue_list_for_clustering)}, got {len(embeddings) if embeddings else 'None'}.")
                  self.issue_clusters = {}
                  return
 
@@ -347,14 +416,16 @@ class IssueManager:
             clusters = clusterer.fit_predict(embeddings)
 
             self.issue_clusters = {}
+            noise_count = 0
             for i, cluster_id in enumerate(clusters):
                 if cluster_id == -1: # HDBSCAN uses -1 for noise points
+                    noise_count += 1
                     continue # Skip noise points
                 if cluster_id not in self.issue_clusters:
                     self.issue_clusters[cluster_id] = []
                 self.issue_clusters[cluster_id].append(i) # Store original index
 
-            logger.info(f"Clustering complete. Found {len(self.issue_clusters)} clusters (excluding noise).")
+            logger.info(f"Clustering complete. Found {len(self.issue_clusters)} clusters with {noise_count} noise points.")
 
         except Exception as e:
             logger.exception(f"Error during issue clustering: {e}")
@@ -372,21 +443,35 @@ class IssueManager:
         headers = {"Authorization": f"Bearer {self.hf_token}"}
 
         logger.info(f"Requesting embeddings from {api_url} for {len(texts)} texts.")
-        async with aiohttp.ClientSession(headers=headers) as session:
+        # Add timeout to the request
+        timeout = aiohttp.ClientTimeout(total=60) # 60 seconds timeout
+        async with aiohttp.ClientSession(headers=headers, timeout=timeout) as session:
             try:
-                response = await session.post(api_url, json={"inputs": texts})
+                # Add wait_for_model=True if using serverless inference endpoints
+                payload = {"inputs": texts, "options": {"wait_for_model": True}}
+                response = await session.post(api_url, json=payload)
                 response.raise_for_status()
                 result = await response.json()
                 # Check if the result is a list of embeddings (floats)
                 if isinstance(result, list) and all(isinstance(emb, list) for emb in result):
                      logger.info(f"Successfully received {len(result)} embeddings.")
                      return result
+                elif isinstance(result, dict) and 'error' in result:
+                     logger.error(f"HF Inference API returned an error: {result['error']}")
+                     return None
                 else:
                      logger.error(f"Unexpected embedding format received: {type(result)}. Full response: {result}")
                      return None
             except aiohttp.ClientResponseError as e:
                  logger.error(f"HF Inference API request failed: {e.status} {e.message}")
-                 logger.error(f"Response body: {await e.response.text()}")
+                 try:
+                     error_body = await e.response.text()
+                     logger.error(f"Response body: {error_body[:500]}") # Log first 500 chars
+                 except Exception as read_err:
+                     logger.error(f"Could not read error response body: {read_err}")
+                 return None
+            except asyncio.TimeoutError:
+                 logger.error(f"HF Inference API request timed out after {timeout.total} seconds.")
                  return None
             except Exception as e:
                  logger.exception(f"An unexpected error occurred during embedding generation: {e}")
@@ -405,32 +490,29 @@ class IssueManager:
         model_id = HF_MODELS[model_key]
         logger.info(f"Generating patch for issue {issue_number} using model {model_id}")
 
-        # --- Context Gathering (Simplified) ---
-        # In a real scenario, this needs to be much smarter:
-        # - Identify relevant files based on issue text, stack traces, etc.
-        # - Potentially use git history or blame to find relevant code sections.
-        # For now, we'll use a placeholder or skip context if too complex.
-        context = "Context gathering not implemented. Provide code snippets in the issue description."
+        # --- Context Gathering (Placeholder) ---
+        context = "Context gathering not implemented. Provide relevant code snippets in the issue description if possible."
         # context = await self._get_code_context(issue_number) # Uncomment if implemented
 
         # --- Prompt Engineering ---
-        prompt = f"""You are an expert programmer tasked with fixing a bug described in a GitHub issue.
-Analyze the following issue and provide a code patch in standard `diff` format.
-Focus only on the necessary changes to resolve the problem described.
-Explain your reasoning briefly before the patch.
+        prompt = f"""You are an expert programmer analyzing a GitHub issue. Your task is to generate a code patch in standard `diff` format to fix the described problem.
 
-## Issue Title: {issue.get('title', 'N/A')}
-## Issue Body:
+## Issue Details:
+### Title: {issue.get('title', 'N/A')}
+### Body:
 {issue.get('body', 'N/A')}
+### Labels: {', '.join(issue.get('labels', []))}
 
 ## Relevant Code Context (if available):
 {context}
 
 ## Instructions:
-1.  Analyze the issue and the context.
-2.  Determine the code changes needed.
-3.  Provide the changes as a Git diff block (```diff ... ```).
-4.  If you cannot determine a patch, explain why.
+1. Carefully analyze the issue description and context.
+2. Identify the specific code changes required to resolve the issue.
+3. Generate a patch containing *only* the necessary code modifications.
+4. Format the patch strictly according to the standard `diff` format, enclosed in a ```diff ... ``` block.
+5. Provide a brief explanation *before* the diff block explaining the reasoning behind the changes.
+6. If a patch cannot be reasonably determined from the provided information, state that clearly instead of generating an incorrect patch.
 
 ## Patch Suggestion:
 """
@@ -441,39 +523,60 @@ Explain your reasoning briefly before the patch.
         payload = {
             "inputs": prompt,
             "parameters": { # Adjust parameters as needed
-                "max_new_tokens": 1024, # Max length of the generated patch + explanation
-                "temperature": 0.3,   # Lower temperature for more deterministic code
+                "max_new_tokens": 1536, # Increased max tokens for potentially larger patches
+                "temperature": 0.2,   # Low temperature for deterministic code generation
                 "return_full_text": False, # Only get the generated part
-                "do_sample": True,
-                "top_p": 0.9,
-            }
+                "do_sample": False, # Turn off sampling for more deterministic output with low temp
+                # "top_p": 0.9, # Less relevant when do_sample=False
+            },
+             "options": {"wait_for_model": True} # For serverless endpoints
         }
+        timeout = aiohttp.ClientTimeout(total=120) # 2 minutes timeout for generation
 
         try:
-            async with aiohttp.ClientSession(headers=headers) as session:
+            async with aiohttp.ClientSession(headers=headers, timeout=timeout) as session:
                 async with session.post(api_url, json=payload) as response:
                     response.raise_for_status()
                     result = await response.json()
-                    if result and isinstance(result, list):
+                    if result and isinstance(result, list) and 'generated_text' in result[0]:
                         generated_text = result[0].get('generated_text', '')
                         logger.info(f"Received patch suggestion from {model_id}")
-                        # Basic extraction of diff block (improve if needed)
+                        # Improved extraction of diff block
                         diff_match = re.search(r"```diff\n(.*?)```", generated_text, re.DOTALL)
                         explanation = generated_text.split("```diff")[0].strip()
-                        patch = diff_match.group(1).strip() if diff_match else "No diff block found in response."
+                        if not explanation: # Handle cases where explanation might be missing
+                             explanation = "No explanation provided."
+
+                        if diff_match:
+                            patch = diff_match.group(1).strip()
+                            # Basic validation: check if patch contains diff markers like + or -
+                            if not re.search(r'^\s*[+-]', patch, re.MULTILINE):
+                                 patch = f"AI generated response, but no standard diff markers (+/-) found:\n---\n{patch}\n---"
+                                 logger.warning("Generated patch lacks standard diff markers.")
+                        else:
+                            patch = "No diff block found in the AI response."
+                            logger.warning("No diff block found in patch suggestion response.")
+
 
                         return {
                             "explanation": explanation,
                             "patch": patch,
                             "model_used": model_id
                         }
+                    elif isinstance(result, dict) and 'error' in result:
+                        logger.error(f"HF Inference API returned an error for patch generation: {result['error']}")
+                        return {"error": f"AI model returned an error: {result['error']}"}
                     else:
-                        logger.error(f"Unexpected response format from {model_id}: {result}")
+                        logger.error(f"Unexpected response format from {model_id} for patch: {result}")
                         return {"error": "Received unexpected response format from AI model."}
         except aiohttp.ClientResponseError as e:
             logger.error(f"HF Inference API request failed for patch generation: {e.status} {e.message}")
-            logger.error(f"Response body: {await e.response.text()}")
+            error_body = await e.response.text()
+            logger.error(f"Response body: {error_body[:500]}")
             return {"error": f"AI model request failed ({e.status}). Check model availability and HF token."}
+        except asyncio.TimeoutError:
+            logger.error(f"HF Inference API request for patch generation timed out after {timeout.total} seconds.")
+            return {"error": "AI model request timed out."}
         except Exception as e:
             logger.exception(f"Error generating code patch: {e}")
             return {"error": f"An unexpected error occurred: {e}"}
@@ -504,44 +607,61 @@ Explain your reasoning briefly before the patch.
         model_id = HF_MODELS[model_key]
         logger.info(f"Requesting resolution suggestion for issue {issue.get('id','N/A')} using {model_id}")
 
-        prompt = f"""Analyze the following GitHub issue and provide a concise, step-by-step suggestion on how to resolve it. Focus on the technical steps required.
+        prompt = f"""Analyze the following GitHub issue and provide a concise, step-by-step suggestion on how to resolve it. Focus on the technical steps required. Be clear and actionable.
 
-## Issue Title: {issue.get('title', 'N/A')}
-## Issue Body:
+## Issue Details:
+### Title: {issue.get('title', 'N/A')}
+### Body:
 {issue.get('body', 'N/A')}
-## Labels: {', '.join(issue.get('labels', []))}
+### Labels: {', '.join(issue.get('labels', []))}
 
 ## Suggested Resolution Steps:
+1.  **Understand the Root Cause:** [Briefly explain the likely cause based on the description]
+2.  **Identify Files:** [Suggest specific files or modules likely involved]
+3.  **Implement Changes:** [Describe the necessary code modifications or additions]
+4.  **Test Thoroughly:** [Mention specific testing approaches needed]
+5.  **Create Pull Request:** [Standard final step]
+
+Provide details for steps 1-4 based on the issue.
 """
         api_url = f"{HF_INFERENCE_API}/{model_id}"
         headers = {"Authorization": f"Bearer {self.hf_token}"}
         payload = {
             "inputs": prompt,
             "parameters": {
-                "max_new_tokens": 512,
-                "temperature": 0.7, # Higher temp for more creative suggestions
+                "max_new_tokens": 768, # Increased token limit
+                "temperature": 0.6, # Slightly lower temp for more focused suggestions
                 "return_full_text": False,
                 "do_sample": True,
                 "top_p": 0.95,
-            }
+            },
+            "options": {"wait_for_model": True} # For serverless endpoints
         }
+        timeout = aiohttp.ClientTimeout(total=90) # 90 seconds timeout
 
         try:
-            async with aiohttp.ClientSession(headers=headers) as session:
+            async with aiohttp.ClientSession(headers=headers, timeout=timeout) as session:
                 async with session.post(api_url, json=payload) as response:
                     response.raise_for_status()
                     result = await response.json()
-                    if result and isinstance(result, list):
+                    if result and isinstance(result, list) and 'generated_text' in result[0]:
                          suggestion = result[0].get('generated_text', 'No suggestion generated.')
                          logger.info(f"Received suggestion from {model_id}")
                          return suggestion.strip()
+                    elif isinstance(result, dict) and 'error' in result:
+                         logger.error(f"HF Inference API returned an error for suggestion: {result['error']}")
+                         return f"Error: AI model returned an error: {result['error']}"
                     else:
                          logger.error(f"Unexpected response format from {model_id} for suggestion: {result}")
                          return "Error: Received unexpected response format from AI model."
         except aiohttp.ClientResponseError as e:
             logger.error(f"HF Inference API request failed for suggestion: {e.status} {e.message}")
-            logger.error(f"Response body: {await e.response.text()}")
+            error_body = await e.response.text()
+            logger.error(f"Response body: {error_body[:500]}")
             return f"Error: AI model request failed ({e.status}). Check model availability and HF token."
+        except asyncio.TimeoutError:
+            logger.error(f"HF Inference API request for suggestion timed out after {timeout.total} seconds.")
+            return "Error: AI model request timed out."
         except Exception as e:
             logger.exception(f"Error suggesting resolution: {e}")
             return f"An unexpected error occurred: {e}"
@@ -555,37 +675,54 @@ Explain your reasoning briefly before the patch.
             if not self.ws_clients:
                 continue
 
+            # Create payload within the loop to get current status
             status_payload = json.dumps({
                 "type": "collaboration_status",
                 "collaborators": self.collaborators
             })
+            if self.collaborators: # Only broadcast if there's someone to report
+                 logger.debug(f"Broadcasting status: {status_payload}")
+
             # Use asyncio.gather to send concurrently, handling potential errors
             results = await asyncio.gather(
                 *[client.send(status_payload) for client in self.ws_clients],
                 return_exceptions=True # Don't let one failed send stop others
             )
             # Log any errors that occurred during broadcast
+            active_clients = []
             for i, result in enumerate(results):
+                 client = self.ws_clients[i]
                  if isinstance(result, Exception):
-                     logger.warning(f"Failed to send status to client {i}: {result}")
+                     logger.warning(f"Failed to send status to client {getattr(client, 'client_id', client.remote_address)}: {result}. Removing client.")
+                     # Schedule removal in the main loop to avoid modifying list while iterating
+                     # self.main_loop.call_soon_threadsafe(self.remove_ws_client, client) # Requires main_loop ref
+                 else:
+                     active_clients.append(client)
+            # self.ws_clients = active_clients # Update list - Careful with concurrency here
 
 
     async def handle_code_editor_update(self, issue_num: int, delta: str, client_id: str):
         """Applies a delta from one client and broadcasts it to others."""
         if issue_num not in self.code_editors:
-            logger.warning(f"Received code update for non-existent editor for issue {issue_num}")
-            return # Or initialize editor: self.code_editors[issue_num] = OTCodeEditor(...)
+            # Initialize editor if it doesn't exist (e.g., first edit)
+            # This requires knowing the initial content, which might be tricky.
+            # For now, log a warning. A better approach might be needed.
+            logger.warning(f"Received code update for non-existent editor instance for issue {issue_num}. Ignoring.")
+            # Alternative: self.code_editors[issue_num] = OTCodeEditor(initial_content="...")
+            return
 
         try:
             # Apply the delta to the server-side authoritative state
-            self.code_editors[issue_num].apply_delta(json.loads(delta))
+            # Assuming apply_delta modifies the internal state correctly
+            parsed_delta = json.loads(delta) # Parse delta once
+            self.code_editors[issue_num].apply_delta(parsed_delta)
             logger.info(f"Applied delta for issue {issue_num} from client {client_id}")
 
             # Broadcast the delta to all *other* connected clients
             update_payload = json.dumps({
                 "type": "code_update",
                 "issue_num": issue_num,
-                "delta": delta # Send the original delta
+                "delta": delta # Send the original delta JSON string
             })
 
             tasks = []
@@ -600,7 +737,8 @@ Explain your reasoning briefly before the patch.
                  # Log errors during broadcast
                  for i, result in enumerate(results):
                      if isinstance(result, Exception):
-                         logger.warning(f"Failed to broadcast code update to client {i}: {result}")
+                         failed_client = tasks[i].__self__ # Get client from task
+                         logger.warning(f"Failed to broadcast code update to client {getattr(failed_client, 'client_id', 'N/A')}: {result}")
 
         except json.JSONDecodeError:
             logger.error(f"Received invalid JSON delta for issue {issue_num}: {delta}")
@@ -619,8 +757,20 @@ Explain your reasoning briefly before the patch.
              return_exceptions=True
          )
          for i, result in enumerate(results):
+             client = self.ws_clients[i]
              if isinstance(result, Exception):
-                 logger.warning(f"Failed to send issue update notification to client {i}: {result}")
+                 logger.warning(f"Failed to send issue update notification to client {getattr(client, 'client_id', client.remote_address)}: {result}")
+
+    def remove_ws_client(self, client_to_remove: websockets.WebSocketServerProtocol):
+        """Safely removes a client from the list and collaborator dict."""
+        client_id = getattr(client_to_remove, 'client_id', None)
+        if client_to_remove in self.ws_clients:
+            self.ws_clients.remove(client_to_remove)
+            logger.info(f"Removed WebSocket client: {client_id or client_to_remove.remote_address}")
+        if client_id and client_id in self.collaborators:
+            del self.collaborators[client_id]
+            logger.info(f"Removed collaborator {client_id}.")
+        # No need to broadcast here, the periodic task will reflect the change
 
 
 # ========== Gradio UI Definition ==========
@@ -633,47 +783,62 @@ def create_ui(manager: IssueManager):
         if issue_num is None or issue_num not in manager.issues:
             return "<p>Select an issue from the board to see details.</p>"
         issue = manager.issues[issue_num]
-        # Convert markdown body to HTML
-        html_body = markdown2.markdown(issue.get('body', '*No description provided.*'))
+        # Convert markdown body to HTML using markdown2 with extras
+        html_body = markdown2.markdown(
+             issue.get('body', '*No description provided.*'),
+             extras=["fenced-code-blocks", "tables", "strike", "task_list"]
+        )
         # Basic styling
         preview_html = f"""
-        <div style="border: 1px solid #e5e7eb; padding: 15px; border-radius: 8px; background-color: #f9fafb;">
-            <h4><a href='{issue.get('url', '#')}' target='_blank' style='color: #6d28d9; text-decoration: none;'>#{issue['id']} - {issue.get('title', 'N/A')}</a></h4>
+        <div style="border: 1px solid #e5e7eb; padding: 15px; border-radius: 8px; background-color: #f9fafb; font-family: 'Inter', sans-serif;">
+            <h4 style="margin-top: 0; margin-bottom: 10px;">
+                <a href='{issue.get('url', '#')}' target='_blank' style='color: #6d28d9; text-decoration: none; font-weight: 600;'>
+                    #{issue['id']} - {issue.get('title', 'N/A')}
+                </a>
+            </h4>
             <hr style='margin: 10px 0; border-top: 1px solid #e5e7eb;'>
-            <p><strong>State:</strong> {issue.get('state', 'N/A')} | <strong>Assignee:</strong> {issue.get('assignee', 'None')}</p>
-            <p><strong>Labels:</strong> {' | '.join(f'<span style=\'background-color: #eee; padding: 2px 5px; border-radius: 4px; font-size: 0.9em;\'>{l}</span>' for l in issue.get('labels', [])) or 'None'}</p>
-            <div style="margin-top: 10px; max-height: 300px; overflow-y: auto; border-top: 1px dashed #ccc; padding-top: 10px;">
+            <p style="font-size: 0.9em; color: #4b5563; margin-bottom: 5px;">
+                <strong>State:</strong> <span style="background-color: #ddd; padding: 1px 4px; border-radius: 3px;">{issue.get('state', 'N/A')}</span> |
+                <strong>Assignee:</strong> {issue.get('assignee', 'None')}
+            </p>
+            <p style="font-size: 0.9em; color: #4b5563; margin-bottom: 10px;">
+                <strong>Labels:</strong> {' | '.join(f'<span style=\'background-color: #eee; padding: 2px 5px; border-radius: 4px; font-size: 0.9em; display: inline-block; margin-right: 3px;\'>{l}</span>' for l in issue.get('labels', [])) or 'None'}
+            </p>
+            <div style="margin-top: 10px; max-height: 350px; overflow-y: auto; border-top: 1px dashed #ccc; padding-top: 10px; font-size: 0.95em; line-height: 1.5;">
                 {html_body}
             </div>
         </div>
         """
         return preview_html
 
-    async def get_ai_suggestion(issue_num: Optional[int], model_key: str) -> str:
-         """Wrapper to get AI suggestion for the chat."""
+    async def get_ai_suggestion_wrapper(issue_num: Optional[int], model_key: str) -> str:
+         """Wrapper to get AI suggestion for the chat display."""
          if issue_num is None or issue_num not in manager.issues:
              return "Please select a valid issue first."
-         issue = manager.issues[issue_num]
-         issue_hash = manager._get_issue_hash(issue) # Use hash for caching
          # Use cached_suggestion which handles the actual API call via lru_cache
-         suggestion = await manager.cached_suggestion(issue_hash, HF_MODELS[model_key])
-         # Format for chat
-         return f"**Suggestion based on {model_key}:**\n\n{suggestion}"
-
-    async def get_ai_patch(issue_num: Optional[int], model_key: str) -> str:
-         """Wrapper to get AI patch for the chat."""
+         # Note: cached_suggestion needs the issue *hash* and model *ID*
+         issue = manager.issues[issue_num]
+         issue_hash = manager._get_issue_hash(issue)
+         # Pass model_key directly, cached_suggestion will resolve it
+         suggestion = await manager.cached_suggestion(issue_hash, model_key)
+         # Format for display
+         return f"**Suggestion based on {model_key}:**\n\n---\n{suggestion}"
+
+    async def get_ai_patch_wrapper(issue_num: Optional[int], model_key: str) -> str:
+         """Wrapper to get AI patch for the chat display."""
          if issue_num is None or issue_num not in manager.issues:
              return "Please select a valid issue first."
          result = await manager.generate_code_patch(issue_num, model_key)
          if "error" in result:
              return f"**Error generating patch:** {result['error']}"
          else:
-             # Format for chat display
+             # Format for chat display using Markdown code block
              return f"""**Patch Suggestion from {result.get('model_used', model_key)}:**
 
 **Explanation:**
 {result.get('explanation', 'N/A')}
 
+---
 **Patch:**
 ```diff
 {result.get('patch', 'N/A')}
@@ -695,19 +860,18 @@ def create_ui(manager: IssueManager):
                 with gr.Row():
                     github_token = gr.Textbox(label="GitHub Token (Optional)", type="password", info="Required for private repos or higher rate limits.", elem_id="github_token")
                     hf_token = gr.Textbox(label="Hugging Face Token", type="password", info="Required for AI model interactions.", elem_id="hf_token")
-            with gr.Column(scale=1):
-                # Removed language select as code editor handles it
+            with gr.Column(scale=1, min_width=250):
                 model_select = gr.Dropdown(choices=list(HF_MODELS.keys()), value="Mistral-8x7B",
                                            label="🤖 Select AI Model", info="Choose the AI for suggestions and patches.", elem_id="model_select")
                 crawl_btn = gr.Button("🛰️ Scan Repository Issues", variant="primary", icon="🔍", elem_id="crawl_btn")
-                status_output = gr.Textbox(label="Status", interactive=False, lines=1, placeholder="Status updates will appear here...", elem_id="status_output")
+                status_output = gr.Textbox(label="Status", interactive=False, lines=1, max_lines=1, placeholder="Status updates will appear here...", elem_id="status_output")
 
 
         # --- Main Tabs ---
         with gr.Tabs(elem_id="main-tabs"):
             # --- Issue Board Tab ---
             with gr.Tab("📋 Issue Board", id="board", elem_id="tab-board"):
-                with gr.Row():
+                with gr.Row(equal_height=False):
                     with gr.Column(scale=3):
                          gr.Markdown("### Open Issues")
                          issue_list = gr.Dataframe(
@@ -716,15 +880,17 @@ def create_ui(manager: IssueManager):
                              interactive=True,
                              height=500,
                              wrap=True, # Wrap long titles
-                             elem_id="issue_list_df"
+                             elem_id="issue_list_df",
+                             overflow_row_behaviour='paginate', # Paginate if needed
+                             max_rows=20 # Show 20 rows per page
                          )
-                    with gr.Column(scale=2):
+                    with gr.Column(scale=2, min_width=350):
                         gr.Markdown("### Issue Severity")
                         stats_plot = gr.Plot(elem_id="stats_plot")
                         # Placeholder for collaborators - updated via JS
                         collab_status = gr.HTML("""
-                        <div style="margin-top: 20px; border: 1px solid #e5e7eb; padding: 10px; border-radius: 8px;">
-                            <h4 style="margin-bottom: 5px; color: #374151;">👥 Active Collaborators</h4>
+                        <div style="margin-top: 20px; border: 1px solid #e5e7eb; padding: 10px; border-radius: 8px; background-color: #f9fafb;">
+                            <h4 style="margin-bottom: 5px; color: #374151; font-size: 1em;">👥 Active Collaborators</h4>
                             <div id="collab-list" style="font-size: 0.9em; max-height: 100px; overflow-y: auto;">
                                 Connecting...
                             </div>
@@ -735,8 +901,8 @@ def create_ui(manager: IssueManager):
             with gr.Tab("💻 Resolution Studio", id="studio", elem_id="tab-studio"):
                 with gr.Row():
                     # Left Column: Issue Details & AI Tools
-                    with gr.Column(scale=1):
-                         gr.Markdown("### Selected Issue")
+                    with gr.Column(scale=1, min_width=400):
+                         gr.Markdown("### Selected Issue Details")
                          # Hidden number input to store selected issue ID
                          selected_issue_id = gr.Number(label="Selected Issue ID", visible=False, precision=0, elem_id="selected_issue_id")
                          issue_preview_html = gr.HTML(
@@ -750,23 +916,26 @@ def create_ui(manager: IssueManager):
                              # Add placeholders for other buttons if needed
                              # test_btn = gr.Button("🧪 Create Tests (Future)", icon="🔬", interactive=False)
                              # impact_btn = gr.Button("📊 Impact Analysis (Future)", icon="📈", interactive=False)
-                         chat_output_display = gr.Textbox(label="AI Output", lines=10, interactive=False, placeholder="AI suggestions and patches will appear here...", elem_id="ai_output_display")
+                         # Use Markdown for better formatting of AI output
+                         ai_output_display = gr.Markdown(value="*AI suggestions and patches will appear here...*", elem_id="ai_output_display")
 
 
-                    # Right Column: Code Editor & Chat (removed chat interface)
-                    with gr.Column(scale=2):
+                    # Right Column: Code Editor
+                    with gr.Column(scale=2, min_width=500):
                         gr.Markdown("### Collaborative Code Editor")
                         # Use the imported code_editor component
                         # We'll update its value dynamically when an issue is selected
+                        # Ensure the code_editor component itself handles language setting based on input dict keys or a prop
                         code_edit_component = code_editor(
                             label="Code Editor",
-                            # Initial value can be empty or a placeholder message
-                            value={"main.py": "# Select an issue to load relevant code (placeholder)"},
-                            # Language is set dynamically if needed, or defaults
-                            language="python", # Default language
-                            elem_id="code_editor_component"
+                            value={"placeholder.txt": "# Select an issue to load relevant code (placeholder)"},
+                            # language="python", # Let component handle language if possible
+                            elem_id="code_editor_component",
+                            # Ensure the component has a reasonable height
+                            # height=600 # This might need to be set via CSS or component props
                         )
                         # Hidden input to trigger code editor updates from server->client WS messages
+                        # This is a fallback if direct JS manipulation of the editor isn't feasible
                         code_editor_update_trigger = gr.Textbox(visible=False, elem_id="code-editor-update-trigger")
 
 
@@ -789,65 +958,88 @@ def create_ui(manager: IssueManager):
             fn=manager.crawl_issues,
             inputs=[repo_url, github_token, hf_token],
             outputs=[issue_list, stats_plot, status_output],
-            api_name="crawl_issues" # For API access if needed
+            api_name="crawl_issues", # For API access if needed
+            show_progress="full"
         )
 
         # 2. Issue Selection in Dataframe
         async def handle_issue_select(evt: gr.SelectData):
             """Handles issue selection: updates preview, loads code (placeholder)."""
-            if evt.index[0] is None: # No row selected
+            if evt.index[0] is None or evt.value is None: # No row selected or value missing
+                logger.info("Issue deselected or invalid selection event.")
                 return {
-                    selected_issue_id: None,
+                    selected_issue_id: gr.update(value=None),
                     issue_preview_html: "<p style='color: #6b7280;'>Select an issue from the table.</p>",
-                    # Reset code editor or show placeholder
-                    code_edit_component: gr.update(value={"placeholder.txt": "# Select an issue to load code."})
+                    code_edit_component: gr.update(value={"placeholder.txt": "# Select an issue to load code."}),
+                    ai_output_display: "*AI suggestions and patches will appear here...*" # Reset AI output
                 }
 
-            selected_id = int(evt.value[0]) # Get ID from the first column ('ID') of the selected row
-            logger.info(f"Issue selected: ID {selected_id}")
+            try:
+                selected_id = int(evt.value[0]) # Get ID from the first column ('ID') of the selected row value
+                logger.info(f"Issue selected: ID {selected_id}")
+
+                # Update the hidden ID field
+                updates = {selected_issue_id: selected_id}
+
+                # Generate and update the HTML preview
+                preview_html = generate_issue_preview(selected_id)
+                updates[issue_preview_html] = preview_html
+
+                # --- Code Loading Logic (Placeholder) ---
+                # This needs real implementation: Find relevant files for the issue
+                # and load their content into the editor component's value format.
+                # Example: Fetch files related to the issue (needs implementation)
+                # files_content = await fetch_relevant_code_for_issue(selected_id)
+                # For now, create a placeholder editor state and potentially initialize server-side OT editor
+                files_content = {
+                     f"issue_{selected_id}_code.py": f"# Code related to issue {selected_id}\n# (Replace with actual file content)\n\nprint('Hello from issue {selected_id}')",
+                     "README.md": f"# Issue {selected_id}\n\nDetails about the issue..."
+                }
+                # Initialize or update the server-side OT document for this issue
+                # This assumes OTCodeEditor takes initial content dictionary
+                manager.code_editors[selected_id] = OTCodeEditor(value=files_content)
+                logger.info(f"Initialized/Reset OT editor state for issue {selected_id}")
 
-            # Update the hidden ID field
-            updates = {selected_issue_id: selected_id}
+                updates[code_edit_component] = gr.update(value=files_content)
+                # --- End Placeholder ---
 
-            # Generate and update the HTML preview
-            preview_html = generate_issue_preview(selected_id)
-            updates[issue_preview_html] = preview_html
+                # Reset AI output display
+                updates[ai_output_display] = "*AI suggestions and patches will appear here...*"
 
-            # --- Code Loading Logic (Placeholder) ---
-            # This needs real implementation: Find relevant files for the issue
-            # and load their content into the editor component's value format.
-            # Example: Fetch files related to the issue (needs implementation)
-            # files_content = await fetch_relevant_code_for_issue(selected_id)
-            files_content = {
-                 f"issue_{selected_id}_code.py": f"# Code related to issue {selected_id}\n# (Replace with actual file content)\n\nprint('Hello from issue {selected_id}')",
-                 "README.md": f"# Issue {selected_id}\n\nDetails about the issue..."
-            }
-            updates[code_edit_component] = gr.update(value=files_content)
-            # --- End Placeholder ---
+                return updates
+            except (ValueError, TypeError, IndexError) as e:
+                 logger.error(f"Error processing selection event data: {evt.value}. Error: {e}")
+                 return { # Return updates to reset state gracefully
+                     selected_issue_id: gr.update(value=None),
+                     issue_preview_html: "<p style='color: red;'>Error processing selection. Please try again.</p>",
+                     code_edit_component: gr.update(value={"error.txt": "# Error processing selection"}),
+                     ai_output_display: "*Error processing selection.*"
+                 }
 
-            return updates
 
         issue_list.select(
             fn=handle_issue_select,
             inputs=[], # Event data is passed automatically
-            outputs=[selected_issue_id, issue_preview_html, code_edit_component],
+            outputs=[selected_issue_id, issue_preview_html, code_edit_component, ai_output_display],
             show_progress="minimal"
         )
 
         # 3. Suggest Resolution Button Click
         suggest_btn.click(
-            fn=get_ai_suggestion,
+            fn=get_ai_suggestion_wrapper,
             inputs=[selected_issue_id, model_select],
-            outputs=[chat_output_display],
-            api_name="suggest_resolution"
+            outputs=[ai_output_display], # Output to Markdown component
+            api_name="suggest_resolution",
+            show_progress="full"
         )
 
         # 4. Generate Patch Button Click
         patch_btn.click(
-            fn=get_ai_patch,
+            fn=get_ai_patch_wrapper,
             inputs=[selected_issue_id, model_select],
-            outputs=[chat_output_display],
-            api_name="generate_patch"
+            outputs=[ai_output_display], # Output to Markdown component
+            api_name="generate_patch",
+            show_progress="full"
         )
 
         # 5. Code Editor Change (User typing) -> Send update via WebSocket
@@ -865,25 +1057,38 @@ def create_ui(manager: IssueManager):
             client_id = f"client_{hashlib.sha1(os.urandom(16)).hexdigest()[:8]}"
             logger.info(f"Generated Client ID for WebSocket: {client_id}")
 
+            # FIX: Escape all literal curly braces `{` -> `{{` and `}` -> `}}` in the JS code
+            # Variables like {ws_port} and {client_id} remain single-braced.
             return f"""
             <script>
-                // Ensure this runs only once
+                // Ensure this runs only once per page load
                 if (!window.collabWs) {{
                     console.log('Initializing WebSocket connection...');
-                    const wsUrl = `ws://localhost:{ws_port}`; // Use localhost for local Gradio run
-                    // For Hugging Face Spaces, you need to use the public WS endpoint:
-                    // const wsUrl = `wss://YOUR_SPACE_NAME.hf.space/ws`; // Adjust if using custom domain/port mapping
+                    // Determine WebSocket URL based on environment (local vs HF Space)
+                    let wsUrl;
+                    if (window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1') {{
+                        wsUrl = `ws://localhost:{ws_port}`;
+                    }} else {{
+                        // Assume running on HF Space or similar, use wss://
+                        // Construct the WS URL based on the current window location (more robust)
+                        const spaceHost = window.location.host;
+                        wsUrl = `wss://${{spaceHost}}/ws`; // Standard path for Gradio WS proxy
+                        console.log('Detected non-local environment, using secure WebSocket URL:', wsUrl);
+                    }}
+
 
                     window.collabWs = new WebSocket(wsUrl);
                     window.clientId = '{client_id}'; // Store client ID globally for this session
+                    window.aceEditorInstance = null; // Store reference to the Ace editor
+                    window.currentIssueId = null; // Track the currently selected issue ID in the editor
 
                     window.collabWs.onopen = function(event) {{
                         console.log('WebSocket connection established.');
-                        // Optionally send a join message
-                        window.collabWs.send(JSON.stringify({{ type: 'join', clientId: window.clientId }}));
-                        // Initial update for collaborator list (optional)
+                        // Send join message with client ID
+                        window.sendWsMessage({{ type: 'join', clientId: window.clientId }});
+                        // Update collaborator list display
                         const collabListDiv = document.getElementById('collab-list');
-                        if (collabListDiv) collabListDiv.innerHTML = 'Connected.';
+                        if (collabListDiv) collabListDiv.innerHTML = 'Connected. Waiting for status...';
                     }};
 
                     window.collabWs.onmessage = function(event) {{
@@ -891,61 +1096,54 @@ def create_ui(manager: IssueManager):
                         try {{
                             const data = JSON.parse(event.data);
 
+                            // --- Collaboration Status Update ---
                             if (data.type === 'collaboration_status') {{
                                 const collabListDiv = document.getElementById('collab-list');
                                 if (collabListDiv) {{
-                                    if (Object.keys(data.collaborators).length > 0) {{
-                                        collabListDiv.innerHTML = Object.entries(data.collaborators)
-                                            .map(([id, info]) => `<div class="collab-item">${info.name || id}: ${info.status || 'Idle'}</div>`)
+                                    const collaborators = data.collaborators || {{}};
+                                    const activeCollaborators = Object.entries(collaborators)
+                                        .filter(([id, info]) => id !== window.clientId); // Exclude self
+
+                                    if (activeCollaborators.length > 0) {{
+                                        collabListDiv.innerHTML = activeCollaborators
+                                            .map(([id, info]) => `<div class="collab-item" style="margin-bottom: 3px;">${{info.name || id}}: ${{info.status || 'Idle'}}</div>`)
                                             .join('');
                                     }} else {{
-                                        collabListDiv.innerHTML = 'No other collaborators active.';
+                                        collabListDiv.innerHTML = 'You are the only active user.';
                                     }}
                                 }}
+                            // --- Code Update from another client ---
                             }} else if (data.type === 'code_update') {{
-                                console.log('Received code update delta for issue:', data.issue_num);
-                                // Find the Gradio Textbox used as a trigger
-                                const triggerTextbox = document.getElementById('code-editor-update-trigger').querySelector('textarea');
-                                if (triggerTextbox) {{
-                                    // Set its value to the received delta (JSON string)
-                                    // This change event will be picked up by Gradio if a .change() listener is attached
-                                    // However, directly manipulating the Ace editor instance is more reliable if possible.
-                                    // For now, we assume the code_editor component handles incoming deltas internally
-                                    // or provides a JS API. If not, this trigger approach is a fallback.
-                                    triggerTextbox.value = JSON.stringify(data); // Pass full data
-                                    // Manually dispatch an input event to ensure Gradio detects the change
-                                    triggerTextbox.dispatchEvent(new Event('input', {{ bubbles: true }}));
-                                    console.log('Triggered Gradio update for code editor.');
-
-                                    // --- Ideal approach: Directly update Ace Editor ---
-                                    // This requires the code_editor component to expose its Ace instance
-                                    // or provide a JS function like `window.updateCodeEditor(issueNum, delta)`
-                                    /*
-                                    if (window.ace && window.aceEditors && window.aceEditors[data.issue_num]) {{
-                                        const editor = window.aceEditors[data.issue_num];
-                                        editor.getSession().getDocument().applyDeltas([JSON.parse(data.delta)]);
-                                        console.log('Applied delta directly to Ace editor for issue:', data.issue_num);
-                                    }} else {{
-                                        console.warn('Ace editor instance not found for issue:', data.issue_num);
+                                console.log('Received code update delta for issue:', data.issue_num, 'from client:', data.clientId);
+                                // Apply delta only if it's for the currently viewed issue and not from self
+                                if (window.aceEditorInstance && data.issue_num === window.currentIssueId && data.clientId !== window.clientId) {{
+                                    try {{
+                                        const delta = JSON.parse(data.delta); // Parse the delta string
+                                        // Prevent triggering the 'change' listener we set up
+                                        window.aceEditorInstance.getSession().getDocument().applyDeltas([delta]);
+                                        console.log('Applied remote delta to Ace editor for issue:', data.issue_num);
+                                    }} catch (e) {{
+                                        console.error('Failed to parse or apply remote delta:', e, data.delta);
                                     }}
-                                    */
                                 }} else {{
-                                    console.error('Code editor update trigger textbox not found.');
+                                    console.log('Ignoring remote delta (wrong issue, self, or editor not ready). Current:', window.currentIssueId, 'Received:', data.issue_num);
                                 }}
+                            // --- Issue List Update Notification ---
                             }} else if (data.type === 'issues_updated') {{
                                 console.log('Received issues updated notification.');
-                                // Optionally trigger a refresh or show a notification
-                                // Example: Update status bar
-                                const statusBar = document.getElementById('status_output').querySelector('textarea');
+                                // Update status bar to prompt user action
+                                const statusBar = document.getElementById('status_output')?.querySelector('textarea');
                                 if (statusBar) {{
-                                    statusBar.value = 'Issue list updated. Refresh may be needed.';
+                                    const timestamp = new Date().toLocaleTimeString();
+                                    statusBar.value = `[${{timestamp}}] Issue list updated on server. Click "Scan Repository Issues" to refresh.`;
+                                    // Manually dispatch input event for Gradio
                                     statusBar.dispatchEvent(new Event('input', {{ bubbles: true }}));
                                 }}
-                                // More robust: Trigger the crawl button's click event via JS? (Can be complex)
+                                // Optionally add a more visible notification element
                             }}
 
                         }} catch (e) {{
-                            console.error('Failed to parse WebSocket message or update UI:', e);
+                            console.error('Failed to parse WebSocket message or update UI:', e, event.data);
                         }}
                     }};
 
@@ -953,7 +1151,8 @@ def create_ui(manager: IssueManager):
                         console.warn('WebSocket connection closed:', event.code, event.reason);
                         const collabListDiv = document.getElementById('collab-list');
                         if (collabListDiv) collabListDiv.innerHTML = '<span style="color: red;">Disconnected</span>';
-                        // Implement reconnection logic if needed
+                        // Implement basic reconnection logic (optional)
+                        // setTimeout(initWebSocket, 5000); // Attempt to reconnect after 5 seconds
                     }};
 
                     window.collabWs.onerror = function(error) {{
@@ -962,99 +1161,128 @@ def create_ui(manager: IssueManager):
                         if (collabListDiv) collabListDiv.innerHTML = '<span style="color: red;">Connection Error</span>';
                     }};
 
-                    // Function to send messages (e.g., code changes)
+                    // Function to send messages
                     window.sendWsMessage = function(message) {{
                         if (window.collabWs && window.collabWs.readyState === WebSocket.OPEN) {{
                             window.collabWs.send(JSON.stringify(message));
                         }} else {{
-                            console.error('WebSocket not connected. Cannot send message.');
+                            console.error('WebSocket not connected. Cannot send message:', message);
                         }}
                     }};
 
                     // --- JS Integration with Code Editor Component ---
-                    // This part is CRUCIAL and depends heavily on how the `code_editor`
-                    // component is implemented (e.g., using Ace Editor).
-                    // We need to:
-                    // 1. Get the editor instance(s).
-                    // 2. Attach a listener to its 'change' event (which provides deltas).
-                    // 3. When a change occurs, send the delta via `sendWsMessage`.
-
-                    // Example assuming Ace Editor and the component stores instances:
+                    // Tries to find the Ace editor instance and attach listeners.
+                    // Needs to run after the component is rendered.
                     function setupCodeEditorListener() {{
-                        // This needs to run *after* the Gradio component is rendered
-                        // and the editor is initialized. Using setTimeout is a common hack.
-                        setTimeout(() => {{
-                            const editorElement = document.querySelector('#code_editor_component'); // Find editor container
-                            // Find the actual Ace instance (this depends on the component's structure)
-                            // This is a GUESS - inspect the component's HTML/JS to find the correct way
-                            let aceEditor;
-                            if (window.ace && editorElement) {{
-                                // Try common ways Ace is attached
-                                aceEditor = window.ace.edit(editorElement.querySelector('.ace_editor')); // Common pattern
-                                // Or maybe the component stores it globally?
-                                // aceEditor = window.activeAceEditor;
-                            }}
+                        // Check if Ace is loaded
+                        if (typeof ace === 'undefined') {{
+                            console.warn('Ace editor library not found. Retrying...');
+                            setTimeout(setupCodeEditorListener, 1000);
+                            return;
+                        }}
 
-                            if (aceEditor) {{
-                                console.log('Ace Editor instance found. Attaching change listener.');
-                                aceEditor.getSession().on('change', function(delta) {{
-                                    // Only send changes made by the user (ignore programmatic changes)
-                                    if (aceEditor.curOp && aceEditor.curOp.command.name) {{
-                                        console.log('Code changed by user:', delta);
-                                        const issueIdElem = document.getElementById('selected_issue_id').querySelector('input');
-                                        const currentIssueId = issueIdElem ? parseInt(issueIdElem.value, 10) : null;
-
-                                        if (currentIssueId !== null && !isNaN(currentIssueId)) {{
-                                            window.sendWsMessage({{
-                                                type: 'code_update',
-                                                issue_num: currentIssueId,
-                                                delta: JSON.stringify(delta), // Send delta as JSON string
-                                                clientId: window.clientId
-                                            }});
-                                        }} else {{
-                                           console.warn('No valid issue selected, cannot send code update.');
-                                        }}
-                                    }}
-                                }});
-                            }} else {{
-                                console.warn('Could not find Ace Editor instance to attach listener. Collaboration may not work.');
-                                // Retry after a delay?
-                                // setTimeout(setupCodeEditorListener, 2000);
-                            }}
-                        }}, 1500); // Wait 1.5 seconds for Gradio/Ace to initialize
-                    }}
+                        const editorElement = document.querySelector('#code_editor_component .ace_editor');
+                        if (!editorElement) {{
+                             console.warn('Ace editor element not found yet. Retrying...');
+                             setTimeout(setupCodeEditorListener, 1000); // Retry if element not ready
+                             return;
+                        }}
 
-                    // Call setup after initial load and potentially after issue selection changes
-                    // if the editor instance is recreated.
-                    setupCodeEditorListener();
+                        // Get Ace editor instance (may vary based on component implementation)
+                        try {{
+                            window.aceEditorInstance = ace.edit(editorElement);
+                            console.log('Ace Editor instance found:', window.aceEditorInstance);
+                        }} catch (e) {{
+                             console.error('Failed to initialize Ace editor instance:', e);
+                             // Maybe the instance is already created and stored differently?
+                             // Look for common patterns if direct init fails.
+                             if (editorElement.env && editorElement.env.editor) {{
+                                 window.aceEditorInstance = editorElement.env.editor;
+                                 console.log('Found Ace Editor instance via element.env.editor');
+                             }} else {{
+                                 console.error('Could not find Ace editor instance. Collaboration may not work.');
+                                 return; // Stop if editor cannot be found
+                             }}
+                        }}
 
-                    // Re-attach listener if the editor component updates (e.g., on issue select)
-                    // This requires observing changes to the component's container
-                     const observer = new MutationObserver((mutationsList, observer) => {{
+
+                        if (window.aceEditorInstance) {{
+                            console.log('Attaching change listener to Ace editor.');
+                            window.aceEditorInstance.getSession().on('change', function(delta) {{
+                                // Check if the change was initiated by the current user
+                                // `aceEditor.curOp` is one way, but might not always be reliable.
+                                // A simpler check: ignore deltas if they originate from applyDeltas (often used for remote changes)
+                                const isUserChange = !delta.ignore; // Add an 'ignore' flag when applying remote deltas
+
+                                // More robust check: Use internal flags if available
+                                const isUserOriginated = window.aceEditorInstance.curOp && window.aceEditorInstance.curOp.command.name;
+
+                                if (isUserOriginated) {{
+                                    // console.log('Code changed by user:', delta);
+                                    // Get the current issue ID from the hidden Gradio input
+                                    const issueIdInput = document.querySelector('#selected_issue_id input');
+                                    const currentIssueIdStr = issueIdInput ? issueIdInput.value : null;
+                                    window.currentIssueId = currentIssueIdStr ? parseInt(currentIssueIdStr, 10) : null;
+
+                                    if (window.currentIssueId !== null && !isNaN(window.currentIssueId)) {{
+                                        // Send the delta via WebSocket
+                                        window.sendWsMessage({{
+                                            type: 'code_update',
+                                            issue_num: window.currentIssueId,
+                                            delta: JSON.stringify(delta), // Send delta as JSON string
+                                            clientId: window.clientId
+                                        }});
+                                    }} else {{
+                                       // console.warn('No valid issue selected, cannot send code update.');
+                                    }}
+                                }} else {{
+                                    // console.log('Ignoring programmatic change delta:', delta);
+                                }}
+                            }});
+                            console.log('Ace editor listener attached.');
+                        }} else {{
+                            console.error('Ace editor instance is null after setup attempt.');
+                        }}
+                    }} // end setupCodeEditorListener
+
+                    // --- Initialization and Observation ---
+                    // Use MutationObserver to detect when the editor component is added/changed in the DOM.
+                    const observerTarget = document.body; // Observe the whole body or a closer parent container
+                    const observer = new MutationObserver((mutationsList, observer) => {{
                          for(const mutation of mutationsList) {{
-                             if (mutation.type === 'childList' && mutation.addedNodes.length > 0) {{
-                                 // Check if the editor element was re-added/modified significantly
+                             if (mutation.type === 'childList') {{
+                                 // Check if the specific editor element we need is now present
                                  if (document.querySelector('#code_editor_component .ace_editor')) {{
-                                     console.log("Code editor component updated, re-attaching listener...");
-                                     setupCodeEditorListener();
-                                     break; // Assume we only need to re-attach once per mutation batch
+                                     console.log("Ace editor element detected in DOM, attempting setup...");
+                                     // Debounce setup call slightly to ensure editor is fully ready
+                                     clearTimeout(window.setupEditorTimeout);
+                                     window.setupEditorTimeout = setTimeout(setupCodeEditorListener, 200);
+                                     // Optionally disconnect observer if setup only needs to happen once,
+                                     // but re-observing might be needed if component re-renders completely.
+                                     // observer.disconnect();
+                                     // break; // Found it, no need to check other mutations in this batch
                                  }}
                              }}
                          }}
                      }});
-                     const targetNode = document.getElementById('code_editor_component');
-                     if(targetNode) {{
-                        observer.observe(targetNode, {{ childList: true, subtree: true }});
-                     }}
+
+                    // Start observing the target node for configured mutations
+                    if(observerTarget) {{
+                        console.log("Starting MutationObserver to detect code editor element.");
+                        observer.observe(observerTarget, {{ childList: true, subtree: true }});
+                    }}
+
+                    // Initial attempt to set up listener after a short delay, in case element exists on load
+                    setTimeout(setupCodeEditorListener, 500);
 
 
                 }} else {{
-                    console.log('WebSocket connection already initialized.');
+                    console.log('WebSocket connection appears to be already initialized.');
                 }}
             </script>
             """
 
-        # Inject the JavaScript into the Gradio app
+        # Inject the JavaScript into the Gradio app when it loads
         app.load(_js=web_socket_js(WS_PORT), fn=None, inputs=None, outputs=None)
 
     return app
@@ -1065,20 +1293,22 @@ async def handle_ws_connection(websocket: websockets.WebSocketServerProtocol, pa
     """Handles incoming WebSocket connections and messages."""
     client_id = None # Initialize client_id for this connection
     manager.ws_clients.append(websocket)
-    logger.info(f"WebSocket client connected: {websocket.remote_address}")
+    logger.info(f"WebSocket client connected: {websocket.remote_address} (Total: {len(manager.ws_clients)})")
     try:
         async for message in websocket:
             try:
                 data = json.loads(message)
                 msg_type = data.get("type")
-                logger.debug(f"Received WS message: {data}") # Log received message content
+                # logger.debug(f"Received WS message: {data}") # Log received message content
 
                 if msg_type == "join":
                      client_id = data.get("clientId", f"anon_{websocket.id}")
                      setattr(websocket, 'client_id', client_id) # Associate ID with socket object
                      manager.collaborators[client_id] = {"name": client_id, "status": "Connected"} # Add to collaborators
                      logger.info(f"Client {client_id} joined.")
-                     # Don't await broadcast here, let the periodic task handle it
+                     # Trigger immediate broadcast in case it's needed, but don't await
+                     asyncio.create_task(manager.broadcast_collaboration_status())
+
 
                 elif msg_type == "code_update":
                     issue_num = data.get("issue_num")
@@ -1096,77 +1326,124 @@ async def handle_ws_connection(websocket: websockets.WebSocketServerProtocol, pa
                      if sender_id and sender_id in manager.collaborators:
                          manager.collaborators[sender_id]["status"] = status
                          logger.info(f"Client {sender_id} status updated: {status}")
-                         # Don't await broadcast here
+                         # Trigger broadcast, don't await
+                         asyncio.create_task(manager.broadcast_collaboration_status())
+
 
                 else:
-                    logger.warning(f"Unknown WebSocket message type received: {msg_type}")
+                    logger.warning(f"Unknown WebSocket message type received: {msg_type} from {client_id or websocket.remote_address}")
 
             except json.JSONDecodeError:
-                logger.error(f"Received invalid JSON over WebSocket: {message}")
+                logger.error(f"Received invalid JSON over WebSocket from {client_id or websocket.remote_address}: {message}")
             except Exception as e:
-                logger.exception(f"Error processing WebSocket message: {e}")
+                logger.exception(f"Error processing WebSocket message from {client_id or websocket.remote_address}: {e}")
 
     except ConnectionClosed as e:
-        logger.info(f"WebSocket client disconnected: {websocket.remote_address} (Code: {e.code}, Reason: {e.reason})")
+        logger.info(f"WebSocket client {client_id or websocket.remote_address} disconnected: (Code: {e.code}, Reason: {e.reason})")
     except Exception as e:
-        logger.exception(f"Unexpected error in WebSocket handler: {e}")
+        logger.exception(f"Unexpected error in WebSocket handler for {client_id or websocket.remote_address}: {e}")
     finally:
-        logger.info(f"Cleaning up connection for client {client_id if client_id else websocket.remote_address}")
-        manager.ws_clients.remove(websocket)
-        if client_id and client_id in manager.collaborators:
-             del manager.collaborators[client_id] # Remove collaborator on disconnect
-             logger.info(f"Removed collaborator {client_id}.")
-        # Don't await broadcast here
+        logger.info(f"Cleaning up connection for client {client_id or websocket.remote_address}")
+        # Use the safe removal method, ensuring it runs in the correct context if needed
+        manager.remove_ws_client(websocket)
+        # Trigger a final status broadcast
+        asyncio.create_task(manager.broadcast_collaboration_status())
+
 
 async def start_websocket_server(manager: IssueManager, port: int):
     """Starts the WebSocket server."""
     # Pass manager instance to the connection handler factory
-    handler = lambda ws, path: handle_ws_connection(ws, path, manager)
-    async with websockets.serve(handler, "localhost", port):
-        logger.info(f"WebSocket server started on ws://localhost:{port}")
-        await asyncio.Future()  # Run forever
+    handler_with_manager = lambda ws, path: handle_ws_connection(ws, path, manager)
+    try:
+        # Set ping interval and timeout to keep connections alive and detect broken ones
+        server = await websockets.serve(
+            handler_with_manager,
+            "0.0.0.0", # Bind to all interfaces
+            port,
+            ping_interval=20, # Send pings every 20 seconds
+            ping_timeout=20   # Wait 20 seconds for pong response
+        )
+        logger.info(f"WebSocket server started on ws://0.0.0.0:{port}")
+        await asyncio.Future()  # Run forever until cancelled
+    except OSError as e:
+         logger.error(f"Failed to start WebSocket server on port {port}: {e}. Port might be in use.")
+         # Exit or handle error appropriately
+         raise
+    except Exception as e:
+         logger.exception(f"Unexpected error starting WebSocket server: {e}")
+         raise
+
 
 def run_webhook_server(manager: IssueManager, port: int):
     """Starts the HTTP webhook server in a separate thread."""
     WebhookHandler.manager_instance = manager # Pass manager instance to the class
-    server_address = ("", port)
-    httpd = HTTPServer(server_address, WebhookHandler)
-    logger.info(f"Webhook HTTP server started on port {port}")
-    httpd.serve_forever()
+    server_address = ("0.0.0.0", port) # Bind to all interfaces
+    try:
+        httpd = HTTPServer(server_address, WebhookHandler)
+        logger.info(f"Webhook HTTP server started on [http://0.0.0.0](http://0.0.0.0):{port}")
+        httpd.serve_forever()
+    except OSError as e:
+        logger.error(f"Failed to start Webhook server on port {port}: {e}. Port might be in use.")
+        # Exit or handle error appropriately
+        # Consider signaling the main thread to stop
+    except Exception as e:
+        logger.exception(f"Unexpected error in Webhook server: {e}")
 
 
 # ========== Main Execution ==========
 if __name__ == "__main__":
     # --- Setup ---
     manager = IssueManager()
+    main_event_loop = asyncio.get_event_loop()
+    manager.main_loop = main_event_loop # Ensure manager has loop reference if needed
 
     # --- Start Background Servers ---
-    # 1. Webhook Server (HTTP)
+    # 1. Webhook Server (HTTP) - Runs in its own thread
     webhook_thread = threading.Thread(target=run_webhook_server, args=(manager, WEBHOOK_PORT), daemon=True)
     webhook_thread.start()
 
-    # 2. WebSocket Server (Runs in main asyncio loop)
-    # We need to run the WebSocket server and the collaborator status broadcast
-    # within an asyncio event loop.
+    # 2. WebSocket Server & Broadcast Task (Asyncio) - Run in main thread's event loop
     async def main_async_tasks():
         # Start the periodic broadcast task
         broadcast_task = asyncio.create_task(manager.broadcast_collaboration_status())
         # Start the WebSocket server
         websocket_server_task = asyncio.create_task(start_websocket_server(manager, WS_PORT))
-        await asyncio.gather(broadcast_task, websocket_server_task)
+        # Keep tasks running
+        await asyncio.gather(broadcast_task, websocket_server_task, return_exceptions=True)
 
-    # Run the asyncio tasks in a separate thread
-    asyncio_thread = threading.Thread(target=lambda: asyncio.run(main_async_tasks()), daemon=True)
-    asyncio_thread.start()
-
-    # --- Create and Launch Gradio App ---
+    # --- Create Gradio App ---
+    # Must be created before starting the loop if it interacts with async functions directly
     app = create_ui(manager)
+
+    # --- Launch Gradio App & Async Tasks ---
+    # Gradio's launch method handles the event loop integration when run directly.
+    # It will run the asyncio tasks alongside the FastAPI/Uvicorn server.
+    app.queue() # Enable queue for handling multiple requests/long-running tasks
     app.launch(
         # share=True, # Enable for public access (use with caution)
-        server_name="0.0.0.0", # Bind to all interfaces for accessibility in containers/networks
+        server_name="0.0.0.0", # Bind to all interfaces for accessibility
         server_port=7860,      # Default Gradio port
-        favicon_path="[https://huggingface.co/front/assets/huggingface_logo-noborder.svg](https://huggingface.co/front/assets/huggingface_logo-noborder.svg)"
+        favicon_path="[https://huggingface.co/front/assets/huggingface_logo-noborder.svg](https://huggingface.co/front/assets/huggingface_logo-noborder.svg)",
+        # Let Gradio manage the asyncio loop
+        # asyncio_task=main_async_tasks() # This might conflict with launch's loop management
     )
 
-    logger.info("Gradio app launched. Webhook and WebSocket servers running in background.")
+    # If launch() blocks, the asyncio tasks need to be run differently,
+    # potentially starting the loop manually before launch or using threading
+    # for the asyncio part if launch() isn't compatible.
+    # However, modern Gradio often handles this integration better.
+
+    # Start the asyncio tasks *after* defining the app but *before* launch blocks,
+    # or ensure launch itself runs them. Gradio's `launch` usually handles the loop.
+    # Let's rely on Gradio's launch to manage the loop and potentially run our tasks.
+    # If WS/broadcast doesn't work, we may need to manually manage the loop/threading.
+
+    # A common pattern if launch() blocks and doesn't run background async tasks:
+    # asyncio_thread = threading.Thread(target=lambda: asyncio.run(main_async_tasks()), daemon=True)
+    # asyncio_thread.start()
+    # app.launch(...)
+
+    logger.info("Gradio app launched. Webhook server running in background thread.")
+    # The asyncio tasks (WebSocket, broadcast) should be running via Gradio's event loop.
+