Now all providers type works correctly

aisbf/database.py

@@ -141,6 +141,26 @@ class DatabaseManager:
                 )
             ''')
 
+            # Migration: Add user_id column to token_usage if it doesn't exist
+            # This handles databases created before the user_id column was added
+            try:
+                if self.db_type == 'sqlite':
+                    cursor.execute("PRAGMA table_info(token_usage)")
+                    columns = [row[1] for row in cursor.fetchall()]
+                    if 'user_id' not in columns:
+                        cursor.execute('ALTER TABLE token_usage ADD COLUMN user_id INTEGER')
+                        logger.info("Migration: Added user_id column to token_usage table")
+                else:  # mysql
+                    cursor.execute("""
+                        SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS
+                        WHERE TABLE_NAME = 'token_usage' AND COLUMN_NAME = 'user_id'
+                    """)
+                    if not cursor.fetchone():
+                        cursor.execute('ALTER TABLE token_usage ADD COLUMN user_id INTEGER')
+                        logger.info("Migration: Added user_id column to token_usage table")
+            except Exception as e:
+                logger.warning(f"Migration check for token_usage.user_id: {e}")
+
             # Create indexes for token_usage
             try:
                 cursor.execute('''

aisbf/handlers.py

@@ -589,10 +589,84 @@ class RequestHandler:
                 # This is more reliable than checking response iterability which can cause false positives
                 is_google_stream = provider_config.type == 'google'
                 is_kiro_stream = provider_config.type == 'kiro'
+                is_kilo_stream = provider_config.type in ('kilo', 'kilocode')
                 logger.info(f"Is Google streaming response: {is_google_stream} (provider type: {provider_config.type})")
                 logger.info(f"Is Kiro streaming response: {is_kiro_stream} (provider type: {provider_config.type})")
+                logger.info(f"Is Kilo streaming response: {is_kilo_stream} (provider type: {provider_config.type})")
 
-                if is_kiro_stream:
+                if is_kilo_stream:
+                    # Handle Kilo/KiloCode streaming response
+                    # Kilo returns an async generator that yields OpenAI-compatible SSE bytes directly
+                    # We parse these and pass through with minimal processing
+                    
+                    accumulated_response_text = ""  # Track full response for token counting
+                    chunk_count = 0
+                    tool_calls_from_stream = []  # Track tool calls from stream
+
+                    async for chunk in response:
+                        chunk_count += 1
+                        try:
+                            logger.debug(f"Kilo chunk type: {type(chunk)}")
+
+                            # Parse SSE chunk to extract JSON data
+                            chunk_data = None
+                            
+                            if isinstance(chunk, bytes):
+                                try:
+                                    chunk_str = chunk.decode('utf-8')
+                                    # May contain multiple SSE lines
+                                    for sse_line in chunk_str.split('\n'):
+                                        sse_line = sse_line.strip()
+                                        if sse_line.startswith('data: '):
+                                            data_str = sse_line[6:].strip()
+                                            if data_str and data_str != '[DONE]':
+                                                try:
+                                                    chunk_data = json.loads(data_str)
+                                                except json.JSONDecodeError:
+                                                    pass
+                                except (UnicodeDecodeError, Exception) as e:
+                                    logger.warning(f"Failed to parse Kilo bytes chunk: {e}")
+                            elif isinstance(chunk, str):
+                                if chunk.startswith('data: '):
+                                    data_str = chunk[6:].strip()
+                                    if data_str and data_str != '[DONE]':
+                                        try:
+                                            chunk_data = json.loads(data_str)
+                                        except json.JSONDecodeError:
+                                            pass
+                            
+                            if chunk_data:
+                                # Extract content and tool calls from chunk
+                                choices = chunk_data.get('choices', [])
+                                if choices:
+                                    delta = choices[0].get('delta', {})
+                                    
+                                    # Track content
+                                    delta_content = delta.get('content', '')
+                                    if delta_content:
+                                        accumulated_response_text += delta_content
+                                    
+                                    # Track tool calls
+                                    delta_tool_calls = delta.get('tool_calls', [])
+                                    if delta_tool_calls:
+                                        for tc in delta_tool_calls:
+                                            tool_calls_from_stream.append(tc)
+                            
+                            # Pass through the chunk as-is
+                            if isinstance(chunk, bytes):
+                                yield chunk
+                            elif isinstance(chunk, str):
+                                yield chunk.encode('utf-8')
+                            else:
+                                yield f"data: {json.dumps(chunk)}\n\n".encode('utf-8')
+
+                        except Exception as chunk_error:
+                            logger.warning(f"Error processing Kilo chunk: {chunk_error}")
+                            continue
+
+                    logger.info(f"Kilo streaming processed {chunk_count} chunks total")
+
+                elif is_kiro_stream:
                     # Handle Kiro streaming response
                     # Kiro returns an async generator that yields OpenAI-compatible SSE strings directly
                     # We need to parse these and handle tool calls properly
@@ -2753,9 +2827,10 @@ class RotationHandler:
         import json
         logger = logging.getLogger(__name__)
         
-        # Check if this is a Google provider based on configuration
+        # Check if this is a Google or Kilo provider based on configuration
         is_google_provider = provider_type == 'google'
-        logger.info(f"Creating streaming response for provider type: {provider_type}, is_google: {is_google_provider}")
+        is_kilo_provider = provider_type in ('kilo', 'kilocode')
+        logger.info(f"Creating streaming response for provider type: {provider_type}, is_google: {is_google_provider}, is_kilo: {is_kilo_provider}")
         
         # Generate system_fingerprint for this request
         # If seed is present in request, generate unique fingerprint per request
@@ -3031,6 +3106,47 @@ class RotationHandler:
                         }]
                     }
                     yield f"data: {json.dumps(final_chunk)}\n\n".encode('utf-8')
+                elif is_kilo_provider:
+                    # Handle Kilo/KiloCode streaming response
+                    # Kilo returns an async generator that yields OpenAI-compatible SSE bytes
+                    accumulated_response_text = ""
+                    chunk_count = 0
+
+                    async for chunk in response:
+                        chunk_count += 1
+                        try:
+                            # Pass through the chunk as-is (already in SSE format)
+                            if isinstance(chunk, bytes):
+                                # Parse to track content for token counting
+                                try:
+                                    chunk_str = chunk.decode('utf-8')
+                                    for sse_line in chunk_str.split('\n'):
+                                        sse_line = sse_line.strip()
+                                        if sse_line.startswith('data: '):
+                                            data_str = sse_line[6:].strip()
+                                            if data_str and data_str != '[DONE]':
+                                                try:
+                                                    chunk_data = json.loads(data_str)
+                                                    choices = chunk_data.get('choices', [])
+                                                    if choices:
+                                                        delta = choices[0].get('delta', {})
+                                                        delta_content = delta.get('content', '')
+                                                        if delta_content:
+                                                            accumulated_response_text += delta_content
+                                                except json.JSONDecodeError:
+                                                    pass
+                                except (UnicodeDecodeError, Exception):
+                                    pass
+                                yield chunk
+                            elif isinstance(chunk, str):
+                                yield chunk.encode('utf-8')
+                            else:
+                                yield f"data: {json.dumps(chunk)}\n\n".encode('utf-8')
+                        except Exception as chunk_error:
+                            logger.warning(f"Error processing Kilo chunk: {chunk_error}")
+                            continue
+
+                    logger.info(f"Kilo streaming processed {chunk_count} chunks total")
                 else:
                     # Handle OpenAI/Anthropic/Kiro streaming responses
                     # Some providers return async generators, others return sync iterables

config/providers.json

@@ -174,13 +174,18 @@
     "kilo": {
       "id": "kilo",
       "name": "KiloCode",
-      "endpoint": "https://kilocode.ai/api/openrouter",
-      "type": "openai",
-      "api_key_required": true,
-      "api_key": "YOUR_KILO_API_KEY",
+      "endpoint": "https://kilo.ai/api/openrouter/v1",
+      "type": "kilo",
+      "api_key_required": false,
+      "api_key": "",
       "nsfw": false,
       "privacy": false,
-      "rate_limit": 0
+      "rate_limit": 0,
+      "kilo_config": {
+        "_comment": "Uses Kilo OAuth2 Device Authorization Grant flow",
+        "credentials_file": "~/.kilo_credentials.json",
+        "api_base": "https://api.kilo.ai"
+      }
     },
     "perplexity": {
       "id": "perplexity",

main.py

@@ -42,11 +42,14 @@ import time
 import logging
 import sys
 import os
+import signal
+import atexit
 import argparse
 import secrets
 import hashlib
 import asyncio
 import httpx
+import multiprocessing
 from logging.handlers import RotatingFileHandler
 from datetime import datetime, timedelta
 from collections import defaultdict
@@ -1040,6 +1043,44 @@ async def startup_event():
     logger.info(f"Available rotations: {list(config.rotations.keys()) if config else []}")
     logger.info(f"Available autoselect: {list(config.autoselect.keys()) if config else []}")
 
+def _cleanup_multiprocessing_children():
+    """Terminate any lingering multiprocessing child processes."""
+    try:
+        active_children = multiprocessing.active_children()
+        if active_children:
+            logger.info(f"Terminating {len(active_children)} multiprocessing child process(es)...")
+            for child in active_children:
+                logger.debug(f"  Terminating child process: {child.name} (PID {child.pid})")
+                child.terminate()
+            # Give them a moment to terminate gracefully
+            for child in active_children:
+                child.join(timeout=2)
+            # Force kill any still alive
+            for child in multiprocessing.active_children():
+                logger.warning(f"  Force killing child process: {child.name} (PID {child.pid})")
+                child.kill()
+    except Exception as e:
+        logger.warning(f"Error cleaning up multiprocessing children: {e}")
+
+
+def _signal_handler(signum, frame):
+    """Handle SIGINT/SIGTERM for clean shutdown including multiprocessing children."""
+    sig_name = signal.Signals(signum).name
+    logger.info(f"Received {sig_name}, shutting down...")
+    _cleanup_multiprocessing_children()
+    # Re-raise the signal so uvicorn can handle its own shutdown
+    signal.signal(signum, signal.SIG_DFL)
+    os.kill(os.getpid(), signum)
+
+
+# Register signal handlers for clean shutdown
+signal.signal(signal.SIGINT, _signal_handler)
+signal.signal(signal.SIGTERM, _signal_handler)
+
+# Also register atexit handler as a safety net
+atexit.register(_cleanup_multiprocessing_children)
+
+
 @app.on_event("shutdown")
 async def shutdown_event():
     """Cleanup on shutdown"""
@@ -1050,6 +1091,9 @@ async def shutdown_event():
         logger.info("Shutting down TOR hidden service...")
         tor_service.disconnect()
         logger.info("TOR hidden service shutdown complete")
+    
+    # Cleanup multiprocessing children (sentence-transformers, torch, etc.)
+    _cleanup_multiprocessing_children()
 
 # Authentication middleware
 @app.middleware("http")

templates/dashboard/providers.html

@@ -152,7 +152,7 @@ function renderProviderDetails(key) {
     if (isKiloProvider && !provider.kilo_config) {
         provider.kilo_config = {
             credentials_file: '~/.kilo_credentials.json',
-            api_base: 'https://kilocode.ai/api/openrouter'
+            api_base: 'https://api.kilo.ai/api/gateway'
         };
     }
     
@@ -253,7 +253,7 @@ function renderProviderDetails(key) {
                 
                 <div class="form-group">
                     <label>API Base URL</label>
-                    <input type="text" value="${kiloConfig.api_base || 'https://kilocode.ai/api/openrouter'}" readonly style="background: #0f2840; cursor: not-allowed;" placeholder="https://kilocode.ai/api/openrouter">
+                    <input type="text" value="${kiloConfig.api_base || 'https://api.kilo.ai/api/gateway'}" readonly style="background: #0f2840; cursor: not-allowed;" placeholder="https://api.kilo.ai/api/gateway">
                     <small style="color: #a0a0a0; display: block; margin-top: 5px;">Kilocode API base URL (fixed)</small>
                 </div>
                 
@@ -591,7 +591,7 @@ function updateNewProviderDefaults() {
         'ollama': 'Ollama local provider. No API key required by default. Endpoint: http://localhost:11434/api',
         'kiro': 'Kiro (Amazon Q Developer) provider. Uses Kiro credentials (IDE, CLI, or direct tokens). Endpoint: https://q.us-east-1.amazonaws.com',
         'claude': 'Claude Code provider. Uses OAuth2 authentication (browser-based login). Endpoint: https://api.anthropic.com/v1',
-        'kilocode': 'Kilocode provider. Uses OAuth2 Device Authorization Grant. Endpoint: https://kilocode.ai/api/openrouter'
+        'kilocode': 'Kilocode provider. Uses OAuth2 Device Authorization Grant. Endpoint: https://api.kilo.ai/api/gateway'
     };
     
     descriptionEl.textContent = descriptions[providerType] || 'Standard provider configuration.';
@@ -653,11 +653,11 @@ function confirmAddProvider() {
             credentials_file: '~/.claude_credentials.json'
         };
     } else if (providerType === 'kilocode') {
-        newProvider.endpoint = 'https://kilocode.ai/api/openrouter';
+        newProvider.endpoint = 'https://api.kilo.ai/api/gateway';
         newProvider.name = key + ' (Kilocode OAuth2)';
         newProvider.kilo_config = {
             credentials_file: '~/.kilo_credentials.json',
-            api_base: 'https://kilocode.ai/api/openrouter'
+            api_base: 'https://api.kilo.ai/api/gateway'
         };
     } else if (providerType === 'openai') {
         newProvider.endpoint = 'https://api.openai.com/v1';
@@ -772,12 +772,12 @@ function updateProviderType(key, newType) {
         providersData[key].api_key_required = false;
         providersData[key].kilo_config = {
             credentials_file: '~/.kilo_credentials.json',
-            api_base: 'https://kilocode.ai/api/openrouter'
+            api_base: 'https://api.kilo.ai/api/gateway'
         };
         delete providersData[key].kiro_config;
         delete providersData[key].claude_config;
         // Set endpoint for kilocode (fixed, not modifiable)
-        providersData[key].endpoint = 'https://kilocode.ai/api/openrouter';
+        providersData[key].endpoint = 'https://api.kilo.ai/api/gateway';
     } else if (newType !== 'kiro' && newType !== 'claude' && newType !== 'kilocode' && (oldType === 'kiro' || oldType === 'claude' || oldType === 'kilocode')) {
         // Transitioning FROM kiro/claude/kilocode: remove special configs, set api_key_required to true
         providersData[key].api_key_required = true;