Fix high-priority chat, persistence, AI worker and chart reliability issues

SECURITY.md

@@ -17,7 +17,7 @@ We encourage users to update to the latest version to receive security updates.
 
 The security of Graphite is a top priority. We appreciate the community's efforts in responsible disclosure and will make every effort to address vulnerabilities in a timely manner.
 
-If you discover a security vulnerability, please report it by opening an issue on the [GitHub repository](https://github.com/dovvnloading/Graphite/issues).
+If you discover a potential security vulnerability, **do not open a public GitHub issue with exploit details**. Instead, use the repository's **Security → Report a vulnerability** option if private vulnerability reporting is enabled. If that option is unavailable, contact the maintainers privately to arrange a secure reporting channel before sharing sensitive technical details.
 
 When reporting a vulnerability, please include the following details:
 

graphite_app/graphite_app/api_provider.py

@@ -3,6 +3,9 @@
 import graphite_config as config
 
 USE_API_MODE = False
+# Bound network waits so user-initiated jobs cannot hang indefinitely.
+OLLAMA_REQUEST_TIMEOUT_SECONDS = 180.0
+API_REQUEST_TIMEOUT_SECONDS = 60.0
 API_PROVIDER_TYPE = None
 API_CLIENT = None
 API_MODELS = {
@@ -76,7 +79,9 @@ def chat(task: str, messages: list, **kwargs) -> dict:
         model = config.OLLAMA_MODELS.get(task)
         if not model:
             raise ValueError(f"No Ollama model configured for task: {task}")
-        return ollama.chat(model=model, messages=messages, **kwargs)
+        return ollama.Client(timeout=OLLAMA_REQUEST_TIMEOUT_SECONDS).chat(
+            model=model, messages=messages, **kwargs
+        )
     else:
         if not API_CLIENT:
             raise RuntimeError("API client not initialized. Configure API settings first.")
@@ -116,7 +121,8 @@ def chat(task: str, messages: list, **kwargs) -> dict:
 
             response = gemini_model.generate_content(
                 contents=gemini_history,
-                generation_config=kwargs
+                generation_config=kwargs,
+                request_options={'timeout': API_REQUEST_TIMEOUT_SECONDS}
             )
 
             return {
@@ -142,7 +148,9 @@ def initialize_api(provider: str, api_key: str, base_url: str = None):
         if not base_url:
             base_url = 'https://api.openai.com/v1'
 
-        API_CLIENT = OpenAI(api_key=api_key, base_url=base_url)
+        API_CLIENT = OpenAI(
+            api_key=api_key, base_url=base_url, timeout=API_REQUEST_TIMEOUT_SECONDS
+        )
 
     elif provider == config.API_PROVIDER_GEMINI:
         try:

graphite_app/graphite_app/graphite_agents.py

@@ -1,9 +1,17 @@
 import ollama
 import json
+import math
 from PySide6.QtCore import QThread, Signal, QPointF
 import graphite_config as config
 import api_provider
 
+MAX_CHART_VALUES = 2000
+MAX_CHART_BINS = 100
+MAX_SANKEY_NODES = 75
+MAX_SANKEY_LINKS = 300
+MAX_CHART_RESPONSE_CHARS = 256_000
+
+
 class ChatWorkerThread(QThread):
     """Run chat generation off the UI thread and emit success/error signals."""
     finished = Signal(str)
@@ -347,12 +355,17 @@ def _normalize_numeric_list(self, values):
         """Convert model-produced values into a clean float list."""
         if not isinstance(values, list):
             raise ValueError("Values must be a list")
+        if len(values) > MAX_CHART_VALUES:
+            raise ValueError(f"Chart exceeds the limit of {MAX_CHART_VALUES} values")
         normalized = []
         for index, value in enumerate(values):
             try:
-                normalized.append(float(value))
-            except (TypeError, ValueError):
-                raise ValueError(f"Value at index {index} is not numeric")
+                number = float(value)
+            except (TypeError, ValueError, OverflowError) as exc:
+                raise ValueError(f"Value at index {index} is not numeric") from exc
+            if not math.isfinite(number):
+                raise ValueError(f"Value at index {index} must be finite")
+            normalized.append(number)
         return normalized
 
     def normalize_chart_payload(self, data, chart_type):
@@ -385,7 +398,9 @@ def normalize_chart_payload(self, data, chart_type):
                 bins = int(float(bins))
             except (TypeError, ValueError):
                 bins = 10
-            normalized['bins'] = max(1, bins)
+            if not 1 <= bins <= MAX_CHART_BINS:
+                raise ValueError(f"Histogram bins must be between 1 and {MAX_CHART_BINS}")
+            normalized['bins'] = bins
 
         if chart_type == 'sankey':
             sankey_data = normalized.get('data', {})
@@ -398,6 +413,11 @@ def normalize_chart_payload(self, data, chart_type):
                 nodes = []
             if not isinstance(links, list):
                 links = []
+            if len(nodes) > MAX_SANKEY_NODES or len(links) > MAX_SANKEY_LINKS:
+                raise ValueError(
+                    f"Sankey charts support at most {MAX_SANKEY_NODES} nodes "
+                    f"and {MAX_SANKEY_LINKS} links"
+                )
 
             clean_nodes = []
             for i, node in enumerate(nodes):
@@ -412,7 +432,9 @@ def normalize_chart_payload(self, data, chart_type):
                     source = int(link.get('source'))
                     target = int(link.get('target'))
                     value = float(link.get('value'))
-                except (TypeError, ValueError):
+                except (TypeError, ValueError, OverflowError):
+                    continue
+                if not math.isfinite(value):
                     continue
                 clean_links.append({'source': source, 'target': target, 'value': value})
 
@@ -494,7 +516,9 @@ def validate_chart_data(self, data, chart_type):
                     return False, "Labels must be a list"
                 if len(data['labels']) != len(data['values']):
                     return False, "Labels and values must have the same length"
-
+                if any(v < 0 for v in data['values']) or sum(data['values']) <= 0:
+                    return False, "Pie charts need non-negative values with a positive total"
+
             # Validate numeric values for non-Sankey charts
             if chart_type != 'sankey':
                 try:
@@ -548,7 +572,10 @@ def get_response(self, text, chart_type):
 
             # Using a more specialized model for code/JSON generation
             response = api_provider.chat(task=config.TASK_CHART, messages=messages)
-            cleaned_response = self.clean_response(response['message']['content'])
+            raw_text = response['message']['content']
+            if len(raw_text) > MAX_CHART_RESPONSE_CHARS:
+                raise ValueError("Chart response is too large to render safely")
+            cleaned_response = self.clean_response(raw_text)
 
             # Parse JSON
             try:

graphite_app/graphite_app/graphite_app.py

@@ -226,16 +226,21 @@ def save_chat(self):
         else:
             QMessageBox.information(self, "Nothing to Save", "The canvas is empty.")
 
+    def is_generation_running(self):
+        """True if any active AI generation task is still running."""
+        return any(
+            (worker := getattr(self, name, None)) is not None and worker.isRunning()
+            for name in ('chat_thread', 'takeaway_thread', 'explainer_thread', 'chart_thread')
+        )
+
     def confirm_session_transition(self):
         """Protect unsaved work before loading, clearing, or closing a session."""
-        for worker_name in ('chat_thread', 'takeaway_thread', 'explainer_thread', 'chart_thread'):
-            worker = getattr(self, worker_name, None)
-            if worker is not None and worker.isRunning():
-                QMessageBox.warning(
-                    self, "Generation in Progress",
-                    "Finish the current AI task before switching or closing chats."
-                )
-                return False
+        if self.is_generation_running():
+            QMessageBox.warning(
+                self, "Generation in Progress",
+                "Finish the current AI task before switching or closing chats."
+            )
+            return False
         if not self.session_manager.has_unsaved_changes():
             return True
         choice = QMessageBox.question(
@@ -446,7 +451,10 @@ def send_message(self):
         message = self.message_input.text().strip()
         if not message:
             return
-
+        if self.is_generation_running():
+            QMessageBox.information(self, "Generation in Progress", "An AI task is already running.")
+            return
+
         # Disable input during processing
         self.message_input.setEnabled(False)
         self.send_button.setEnabled(False)
@@ -459,12 +467,18 @@ def send_message(self):
 
         # Add user message node
         user_node = self.chat_view.scene().add_chat_node(
-            message, 
-            is_user=True, 
+            message,
+            is_user=True,
             parent_node=self.current_node,
             conversation_history=history
         )
-
+        if user_node is None:
+            self.message_input.setEnabled(True)
+            self.send_button.setEnabled(True)
+            self.loading_overlay.hide()
+            QMessageBox.critical(self, "Message Failed", "Could not create the message node.")
+            return
+
         # Update conversation history
         user_node.conversation_history = history + [
             {'role': 'user', 'content': message}
@@ -491,7 +505,23 @@ def handle_response(self, response, user_node):
                 {'role': 'assistant', 'content': response}
             ]
         )
-
+        if ai_node is None:
+            # The user's question is still on the canvas. Keep it recoverable
+            # instead of leaving the input disabled after a rendering failure.
+            self.message_input.setEnabled(True)
+            self.send_button.setEnabled(True)
+            self.loading_overlay.hide()
+            try:
+                self.session_manager.save_current_chat()
+            except Exception as exc:
+                QMessageBox.warning(self, "Auto-save Failed", f"Your canvas is unsaved: {exc}")
+            QMessageBox.critical(
+                self, "Response Failed",
+                "The generated response could not be displayed. "
+                "Your question remains on the canvas."
+            )
+            return
+
         # Update current node and view
         self.current_node = ai_node
         self.message_input.clear()
@@ -502,8 +532,11 @@ def handle_response(self, response, user_node):
         self.send_button.setEnabled(True)
         self.loading_overlay.hide()
 
-        # Auto-save after response
-        self.session_manager.save_current_chat()
+        # A failed save should not leave the UI broken; the canvas remains available.
+        try:
+            self.session_manager.save_current_chat()
+        except Exception as exc:
+            QMessageBox.warning(self, "Auto-save Failed", f"Your canvas is unsaved: {exc}")
 
     def handle_error(self, error_message):
         QMessageBox.critical(self, "Error", f"An error occurred: {error_message}")
@@ -524,6 +557,9 @@ def cleanup_takeaway_thread(self):
 
     def generate_takeaway(self, node):
         """Generate takeaway for the given node"""
+        if self.is_generation_running():
+            QMessageBox.information(self, "Generation in Progress", "Finish the current AI task first.")
+            return
         try:
             # Cleanup any existing thread
             self.cleanup_takeaway_thread()
@@ -585,6 +621,9 @@ def cleanup_explainer_thread(self):
 
     def generate_explainer(self, node):
         """Generate simple explanation for the given node"""
+        if self.is_generation_running():
+            QMessageBox.information(self, "Generation in Progress", "Finish the current AI task first.")
+            return
         try:
             # Cleanup any existing thread
             self.cleanup_explainer_thread()
@@ -638,6 +677,9 @@ def handle_explainer_error(self, error_message):
 
     def generate_chart(self, node, chart_type):
         """Generate chart for the given node"""
+        if self.is_generation_running():
+            QMessageBox.information(self, "Generation in Progress", "Finish the current AI task first.")
+            return
         try:
             # Show loading overlay
             self.loading_overlay.show()

graphite_app/graphite_app/graphite_core.py

@@ -337,7 +337,7 @@ def serialize_frame(self, frame):
         return {
             'nodes': [self.window.chat_view.scene().nodes.index(node) for node in frame.nodes],
             'position': {'x': frame.pos().x(), 'y': frame.pos().y()},
-            'note': frame.note,
+            'note': frame.edit_text if frame.editing else frame.note,
             'size': {
                 'width': frame.rect.width(),
                 'height': frame.rect.height()
@@ -349,7 +349,7 @@ def serialize_frame(self, frame):
     def serialize_note(self, note):
         """Convert a Note to a serializable dictionary"""
         return {
-            'content': note.content,
+            'content': note.edit_text if note.editing else note.content,
             'position': {'x': note.pos().x(), 'y': note.pos().y()},
             'size': {'width': note.width, 'height': note.height},
             'color': note.color,
@@ -402,6 +402,13 @@ def has_unsaved_changes(self):
             return False
         return current != self._saved_snapshot
 
+    def detach_deleted_chat(self, chat_id):
+        """Keep an active deleted chat on the canvas as an unsaved new session."""
+        if self.current_chat_id == chat_id:
+            self.current_chat_id = None
+            # None cannot equal a snapshot, so closing or switching prompts to save.
+            self._saved_snapshot = None
+
     def mark_current_state_saved(self):
         """Record a clean baseline after a new session or successful load."""
         self._saved_snapshot = self.capture_current_snapshot()