Merge branch 'fix/broker-keepalive-eventloop-blocking'

Stop event-loop blocking that trips broker WebSocket keepalive timeouts.

aisbf/analytics.py

@@ -23,6 +23,7 @@ Why did the programmer quit his job? Because he didn't get arrays!
 Token Usage Analytics module for AISBF.
 """
 import json
+import asyncio
 from decimal import Decimal
 import csv
 import io
@@ -357,21 +358,64 @@ class Analytics:
         self._latencies[provider_id]['min_ms'] = min(self._latencies[provider_id]['min_ms'], latency_ms)
         self._latencies[provider_id]['max_ms'] = max(self._latencies[provider_id]['max_ms'], latency_ms)
         
-        # Persist to database
+        # Persist to database OFF the event loop.  record_token_usage() opens a MySQL
+        # connection and runs a blocking INSERT; calling it inline here (this method runs
+        # synchronously inside async request handlers) stalled the event loop long enough
+        # for the CoderAI broker WebSocket keepalive to time out (1011) under load and
+        # forced reconnects.  Nothing awaits the result, so fire-and-forget in a worker
+        # thread is safe.
         if tokens_used > 0:
             logger.info(f"Analytics.record_request: Recording to DB - provider={provider_id}, latency_ms={latency_ms}, success={success}, prompt={prompt_tokens}, completion={completion_tokens}, cost={actual_cost}")
+            self._persist_token_usage(
+                provider_id, model_name, tokens_used, user_id, success, latency_ms,
+                error_type, token_id, prompt_tokens, completion_tokens, actual_cost,
+                rotation_id, autoselect_id, analytics_kind,
+            )
+
+    def _persist_token_usage(self, *args):
+        """Write a token_usage row without blocking the caller's event loop.
+
+        When a running loop is present (the normal request path) the blocking DB write
+        is dispatched to a thread pool and NOT awaited — analytics is pure side-effect
+        logging that nothing downstream waits on.  With no loop (sync scripts/tests) it
+        runs inline so behaviour is unchanged there."""
+        provider_id = args[0]
+
+        # args positional order matches record_token_usage(); pull out what the
+        # per-user write also needs.
+        model_name = args[1]
+        tokens_used = args[2]
+        user_id = args[3]
+        token_id = args[7]
+
+        def _write():
             try:
-                self.db.record_token_usage(provider_id, model_name, tokens_used, user_id, success, latency_ms, error_type, token_id, prompt_tokens, completion_tokens, actual_cost, rotation_id, autoselect_id, analytics_kind)
+                self.db.record_token_usage(*args)
                 logger.info(f"✅ Analytics recording completed successfully for {provider_id}")
             except Exception as e:
                 logger.error(f"❌ Analytics recording FAILED for {provider_id}: {e}")
                 import traceback
                 logger.error(f"Traceback: {traceback.format_exc()}")
-
-            # Also record user token usage if this was an API token request
+            # Also record user token usage if this was an API token request.
             if user_id is not None and token_id is not None:
-                self.db.record_user_token_usage(user_id, token_id, provider_id, model_name, tokens_used)
-    
+                try:
+                    self.db.record_user_token_usage(user_id, token_id, provider_id, model_name, tokens_used)
+                except Exception as e:
+                    logger.error(f"❌ Per-user token usage recording FAILED for {provider_id}: {e}")
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+        if loop is None:
+            _write()
+            return
+        executor = getattr(self.db, 'executor', None)
+        # Fire-and-forget: the returned future is intentionally not awaited.  _write()
+        # swallows all exceptions so the future always resolves cleanly (no "exception
+        # never retrieved" warnings).
+        loop.run_in_executor(executor, _write)
+
     def get_provider_stats(
         self,
         provider_id: str,

aisbf/coderai_broker.py

@@ -255,12 +255,19 @@ class CoderAIBroker:
 
     # -------------------------------------------------- shared cache helpers
 
-    def _store_session_cache(self, session: CoderAISession) -> None:
+    def _build_session_cache_payload(self, session: CoderAISession) -> tuple[str, dict]:
+        """Build the (meta_key, payload) for the shared cache – pure in-memory work."""
         payload = self._serialize_session(session)
         payload["metadata"] = dict(payload.get("metadata") or {})
         payload["metadata"]["broker_node_id"] = self._node_id
         payload["metadata"]["connection_state"] = "connected" if not session.closed else "disconnected"
         meta_key = self._session_meta_key(session.provider_id, session.client_id)
+        return meta_key, payload
+
+    def _flush_session_cache(self, meta_key: str, payload: dict) -> None:
+        """Write the session payload + index to the shared cache.  This performs
+        synchronous Redis round-trips and MUST be run off the event loop when on a hot
+        path (see touch())."""
         self._cache.broker_set(meta_key, payload, ttl=SESSION_TTL_SECONDS)
         # Keep a global index of all live session meta keys so list_sessions
         # can enumerate them without a Redis SCAN.
@@ -271,6 +278,12 @@ class CoderAIBroker:
             index.append(meta_key)
             self._cache.broker_set(self._session_index_key(), index, ttl=SESSION_TTL_SECONDS * 10)
 
+    def _store_session_cache(self, session: CoderAISession) -> None:
+        """Synchronous store – used by low-frequency paths (register / periodic
+        heartbeat_loop) where a durable write before returning is fine."""
+        meta_key, payload = self._build_session_cache_payload(session)
+        self._flush_session_cache(meta_key, payload)
+
     def _mark_session_offline_cache(self, provider_id: str, client_id: str) -> None:
         """Write a closed=True tombstone so remote nodes see the disconnect immediately."""
         meta_key = self._session_meta_key(provider_id, client_id)
@@ -392,8 +405,12 @@ class CoderAIBroker:
                 session.metadata.update(self._normalize_hardware_metadata(metadata))
             if capabilities:
                 session.capabilities = dict(capabilities)
-            self._store_session_cache(session)
-            return session
+            meta_key, payload = self._build_session_cache_payload(session)
+        # touch() runs for EVERY inbound WebSocket frame (heartbeat / response). Flush
+        # the shared-cache TTL refresh in a worker thread so its synchronous Redis I/O
+        # can't stall the event loop / WebSocket keepalive under load.
+        await asyncio.to_thread(self._flush_session_cache, meta_key, payload)
+        return session
 
     # ------------------------------------------------------- session queries
 
@@ -720,7 +737,11 @@ class CoderAIBroker:
             return
         await self._record_request_metric(request_id, message)
         # Write to shared cache so the requesting node (possibly remote) can poll it.
-        self._cache.broker_set(self._reply_key(request_id), message, ttl=REPLY_TTL_SECONDS)
+        # Run the synchronous Redis write in a thread so it doesn't block the event loop
+        # (awaited here to preserve ordering; the loop stays free during the I/O).
+        await asyncio.to_thread(
+            self._cache.broker_set, self._reply_key(request_id), message, REPLY_TTL_SECONDS
+        )
         # Also resolve any local Future for the fast-path case.
         await self.resolve_response(message)
 
@@ -742,7 +763,11 @@ class CoderAIBroker:
             return
 
         # Slow path: push chunk to the shared cache list for the remote requester.
-        self._cache.broker_push(self._reply_key(request_id), message, ttl=REPLY_TTL_SECONDS)
+        # Offload the synchronous Redis push so a stream of chunks can't block the loop;
+        # awaited to keep chunk ordering intact.
+        await asyncio.to_thread(
+            self._cache.broker_push, self._reply_key(request_id), message, REPLY_TTL_SECONDS
+        )
 
     async def finish_stream(self, request_id: str) -> None:
         """Remove a streaming request's pending state once the consumer has drained

aisbf/database.py

@@ -26,6 +26,7 @@ import json
 import hashlib
 import time
 import copy
+import contextlib
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Any
 from datetime import datetime, timedelta
@@ -55,10 +56,12 @@ def _verify_password(password: str, stored_hash: str) -> bool:
 
 try:
     import mysql.connector as _mysql_connector
+    from mysql.connector import pooling as _mysql_pooling
     MYSQL_AVAILABLE = True
 except ImportError:
     MYSQL_AVAILABLE = False
     _mysql_connector = None
+    _mysql_pooling = None
 
 logger = logging.getLogger(__name__)
 
@@ -90,16 +93,24 @@ class _MySQLConnectionWrapper:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type is None:
-            self._conn.commit()
+            try:
+                self._conn.commit()
+            except Exception:
+                pass
         else:
             try:
                 self._conn.rollback()
             except Exception:
                 pass
-        # Connection is intentionally left open: cursor and conn variables in the
-        # calling function remain valid after the with-block exits (matching SQLite's
-        # context-manager behaviour).  The connection is closed by GC when the caller
-        # function returns and conn goes out of scope.
+        # Release the connection deterministically instead of waiting for GC.  For a
+        # pooled connection, close() returns it to the pool for reuse (avoiding a fresh
+        # TCP+auth handshake per query on the event loop); for a direct fallback
+        # connection it closes the socket.  No caller uses conn/cursor after the
+        # with-block, so this is safe.
+        try:
+            self._conn.close()
+        except Exception:
+            pass
         return False
 
     # Forward attribute access so the wrapper can be used directly as well
@@ -141,20 +152,67 @@ class DatabaseManager:
         self.database_type = database_type
         self.db_type = self.db_config.get('type', 'sqlite').lower()
         self.executor = get_db_executor()
-        
+        self._mysql_pool = None
+
         if self.db_type == 'mysql':
             if not MYSQL_AVAILABLE:
                 raise ImportError("MySQL connector not available. Install mysql-connector-python.")
+            self._init_mysql_pool()
 
         self._initialize_database()
         logger.info(f"Database initialized: {self.db_type}")
 
+    def _init_mysql_pool(self) -> None:
+        """Create a MySQL connection pool so queries reuse warm connections instead
+        of paying a fresh TCP+auth handshake per call on the event loop.  Pool
+        creation failures are non-fatal: _get_connection() falls back to direct
+        connections so the DB stays usable even if MySQL is briefly unreachable at
+        startup."""
+        if _mysql_pooling is None:
+            return
+        try:
+            pool_size = int(self.db_config.get('mysql_pool_size', 16) or 16)
+        except (TypeError, ValueError):
+            pool_size = 16
+        pool_size = max(1, min(pool_size, 32))  # mysql.connector hard-caps pool_size at 32
+        try:
+            self._mysql_pool = _mysql_pooling.MySQLConnectionPool(
+                pool_name=f"aisbf_{self.database_type}",
+                pool_size=pool_size,
+                pool_reset_session=True,
+                host=self.db_config['mysql_host'],
+                port=self.db_config['mysql_port'],
+                user=self.db_config['mysql_user'],
+                password=self.db_config['mysql_password'],
+                database=self.db_config['mysql_database'],
+            )
+            logger.info(f"MySQL connection pool initialized (size={pool_size})")
+        except Exception as e:
+            self._mysql_pool = None
+            logger.warning(f"MySQL pool init failed ({e}); falling back to per-call connections")
+
     def _get_connection(self):
         """Get a database connection based on the configured type."""
         if self.db_type == 'sqlite':
             db_path = Path(self.db_config['sqlite_path']).expanduser()
             return sqlite3.connect(str(db_path))
         elif self.db_type == 'mysql':
+            # Preferred path: borrow a warm connection from the pool.  ping(reconnect=True)
+            # revives a connection the server dropped (MySQL wait_timeout) so callers never
+            # see a stale socket.  Pool exhaustion or any error falls through to a direct
+            # connection so a request is never rejected just because the pool is busy.
+            if self._mysql_pool is not None:
+                try:
+                    conn = self._mysql_pool.get_connection()
+                    try:
+                        conn.ping(reconnect=True, attempts=2, delay=0)
+                    except Exception:
+                        with contextlib.suppress(Exception):
+                            conn.close()
+                        raise
+                    return _MySQLConnectionWrapper(conn)
+                except Exception as pool_err:
+                    logger.debug(f"MySQL pool unavailable ({pool_err}); using direct connection")
             try:
                 conn = _mysql_connector.connect(
                     host=self.db_config['mysql_host'],

aisbf/routes/coderai_broker.py

@@ -144,7 +144,25 @@ async def _coderai_broker_websocket_impl(websocket: WebSocket, scope_name: str):
             if queued is not None:
                 await websocket.send_text(json.dumps(queued))
 
+    # Process response frames on a dedicated worker so a slow publish_response (cache
+    # writes, metric recording) can never delay the receive loop from reading and
+    # answering the next heartbeat. A single worker draining a FIFO queue preserves
+    # streaming chunk order.
+    publish_queue: asyncio.Queue = asyncio.Queue()
+
+    async def _process_publish_queue() -> None:
+        while True:
+            msg = await publish_queue.get()
+            try:
+                await broker.touch(session.session_id)
+                await broker.publish_response(msg)
+            except Exception:
+                logger.error("CoderAI broker publish handler failed", exc_info=True)
+            finally:
+                publish_queue.task_done()
+
     queue_task = asyncio.create_task(_drain_broker_queue())
+    publish_task = asyncio.create_task(_process_publish_queue())
     try:
         while True:
             raw = await websocket.receive_text()
@@ -171,7 +189,8 @@ async def _coderai_broker_websocket_impl(websocket: WebSocket, scope_name: str):
                 asyncio.create_task(_broker_refresh_models(session.provider_id, owner_user_id))
                 continue
             if op == "heartbeat":
-                await broker.touch(session.session_id, metadata=message.get("payload") or {})
+                # Answer the keepalive immediately; the TTL touch (which does off-loop
+                # cache I/O) happens after so a slow refresh can't delay the reply.
                 await websocket.send_text(json.dumps({
                     "v": 1,
                     "request_id": message.get("request_id"),
@@ -179,17 +198,22 @@ async def _coderai_broker_websocket_impl(websocket: WebSocket, scope_name: str):
                     "event": "heartbeat",
                     "payload": {"ts": int(time.time())},
                 }))
+                await broker.touch(session.session_id, metadata=message.get("payload") or {})
                 continue
-            await broker.touch(session.session_id)
-            await broker.publish_response(message)
+            # Hand response frames to the worker and immediately loop back to reading,
+            # so heartbeats/register are never queued behind response processing.
+            publish_queue.put_nowait(message)
     except WebSocketDisconnect:
         logger.info(f"CoderAI broker disconnected provider={provider_id} client={client_id}")
     except Exception as e:
         logger.error(f"CoderAI broker websocket error: {e}", exc_info=True)
     finally:
         queue_task.cancel()
+        publish_task.cancel()
         with contextlib.suppress(asyncio.CancelledError):
             await queue_task
+        with contextlib.suppress(asyncio.CancelledError):
+            await publish_task
         await broker.unregister(session.session_id)
         await broker.fail_session_requests(session.session_id, f"CoderAI session '{session.session_id}' disconnected")
 

aisbf/utils.py

@@ -22,40 +22,107 @@ Why did the programmer quit his job? Because he didn't get arrays!
 
 Utility functions for AISBF.
 """
+import hashlib
+import logging
+import threading
+from collections import OrderedDict
 from typing import Dict, List, Optional
 from langchain_text_splitters import TokenTextSplitter
 from .config import config
 
+logger = logging.getLogger(__name__)
+
+# tiktoken's BPE tokenizer is CPU-bound and was being run several times per request
+# on the *same* large message list (each pass re-encoding tens of KB), blocking the
+# asyncio event loop.  Two caches remove that cost:
+#   * _ENCODING_BY_MODEL avoids re-resolving/loading the encoding object per call.
+#   * _TOKEN_COUNT_CACHE memoizes the total for identical (model, messages) inputs so
+#     the repeated per-request calls collapse to a single BPE pass.
+# count_messages_tokens is a pure function of its inputs, so memoization is safe.
+_ENCODING_BY_MODEL: Dict[str, object] = {}
+_TOKEN_COUNT_CACHE: "OrderedDict[str, int]" = OrderedDict()
+_TOKEN_COUNT_CACHE_MAX = 1024
+_TOKEN_COUNT_LOCK = threading.Lock()
+
+
+def _get_encoding(model: str):
+    """Return a cached tiktoken encoding for the model (loaded at most once per name)."""
+    enc = _ENCODING_BY_MODEL.get(model)
+    if enc is not None:
+        return enc
+    import tiktoken
+    # OpenAI models use cl100k_base encoding
+    if model.startswith(('gpt-', 'text-', 'davinci-', 'ada-', 'babbage-', 'curie-')):
+        try:
+            enc = tiktoken.encoding_for_model(model)
+        except KeyError:
+            # Fallback to cl100k_base if model not found
+            enc = tiktoken.get_encoding("cl100k_base")
+    else:
+        # Default encoding for other models
+        enc = tiktoken.get_encoding("cl100k_base")
+    _ENCODING_BY_MODEL[model] = enc
+    return enc
+
+
+def _messages_cache_key(messages: List[Dict], model: str) -> str:
+    """Cheap content hash of (model, messages).  Hashing is O(n) but far cheaper than a
+    BPE encode pass, so a cache hit that skips re-encoding is a clear win."""
+    h = hashlib.blake2b(digest_size=16)
+    h.update(model.encode('utf-8', 'ignore'))
+    for msg in messages:
+        content = msg.get('content', '')
+        if not content:
+            h.update(b'\x00')
+            continue
+        if isinstance(content, str):
+            h.update(b'\x01')
+            h.update(content.encode('utf-8', 'ignore'))
+        elif isinstance(content, list):
+            h.update(b'\x02')
+            for item in content:
+                if isinstance(item, dict):
+                    text = item.get('text', '')
+                    if text:
+                        h.update(text.encode('utf-8', 'ignore'))
+                elif isinstance(item, str):
+                    h.update(item.encode('utf-8', 'ignore'))
+                h.update(b'\x1f')
+        h.update(b'\x1e')
+    return h.hexdigest()
+
 
 def count_messages_tokens(messages: List[Dict], model: str) -> int:
     """
     Count the total number of tokens in a list of messages.
-    
-    This function uses tiktoken for accurate token counting.
-    
+
+    This function uses tiktoken for accurate token counting.  Results are memoized on
+    (model, messages) content so the several calls made per request collapse to one
+    BPE pass and stop starving the event loop.
+
     Args:
         messages: List of message dictionaries with 'role' and 'content' keys
         model: Model name to determine encoding
-    
+
     Returns:
         Total token count across all messages
     """
-    import tiktoken
-    import logging
-    logger = logging.getLogger(__name__)
-    
-    # Select encoding based on model
-    # OpenAI models use cl100k_base encoding
-    if model.startswith(('gpt-', 'text-', 'davinci-', 'ada-', 'babbage-', 'curie-')):
-        try:
-            encoding = tiktoken.encoding_for_model(model)
-        except KeyError:
-            # Fallback to cl100k_base if model not found
-            encoding = tiktoken.get_encoding("cl100k_base")
-    else:
-        # Default encoding for other models
-        encoding = tiktoken.get_encoding("cl100k_base")
-    
+    cache_key = None
+    try:
+        cache_key = _messages_cache_key(messages, model)
+    except Exception:
+        # Never let cache-key construction break token counting.
+        cache_key = None
+
+    if cache_key is not None:
+        with _TOKEN_COUNT_LOCK:
+            cached = _TOKEN_COUNT_CACHE.get(cache_key)
+            if cached is not None:
+                _TOKEN_COUNT_CACHE.move_to_end(cache_key)
+                return cached
+
+    encoding = _get_encoding(model)
+
     total_tokens = 0
     for msg in messages:
         content = msg.get('content', '')
@@ -71,7 +138,14 @@ def count_messages_tokens(messages: List[Dict], model: str) -> int:
                             total_tokens += len(encoding.encode(text, disallowed_special=()))
                     elif isinstance(item, str):
                         total_tokens += len(encoding.encode(item, disallowed_special=()))
-    
+
+    if cache_key is not None:
+        with _TOKEN_COUNT_LOCK:
+            _TOKEN_COUNT_CACHE[cache_key] = total_tokens
+            _TOKEN_COUNT_CACHE.move_to_end(cache_key)
+            while len(_TOKEN_COUNT_CACHE) > _TOKEN_COUNT_CACHE_MAX:
+                _TOKEN_COUNT_CACHE.popitem(last=False)
+
     logger.debug(f"Token count for model {model}: {total_tokens}")
     return total_tokens