ziya 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

app/middleware/streaming.py

@@ -80,13 +80,19 @@ class StreamingMiddleware(BaseHTTPMiddleware):
         self._recent_lines = []
         accumulated_content = ""
         accumulated_chunks = []  # Track all chunks for better preservation
+
+        # Code block state tracking for frontend
+        frontend_code_block_state = {
+            'in_block': False,
+            'block_type': None
+        }
         
         # Limits for preserved content to prevent context bloat
         MAX_PRESERVED_TOOLS = 10
         MAX_TOOL_OUTPUT_LENGTH = 5000
         successful_tool_outputs = []  # Track successful tool executions
         tool_sequence_count = 0
-        content_buffer = ""  # Buffer to hold content while checking for tool calls
+        # content_buffer = ""  # Disabled buffering for real-time streaming
         partial_response_preserved = False
         try:
             async for chunk in original_iterator:
@@ -138,7 +144,7 @@ class StreamingMiddleware(BaseHTTPMiddleware):
                                 chunk_content = json.dumps(json_obj)
                                 
                                 # DEBUGGING: Check if JSON serialization changed size
-                                if len(chunk_content) != chunk_size and json_obj.get('type') == 'tool_execution':
+                                if len(chunk_content) != chunk_size and json_obj.get('type') in ['tool_execution', 'tool_display']:
                                     logger.warning(f"🔍 JSON_SIZE_CHANGE: Original {chunk_size} -> Serialized {len(chunk_content)} chars")
                                 
                                 yield f"data: {json.dumps(json_obj)}\n\n"
@@ -146,46 +152,18 @@ class StreamingMiddleware(BaseHTTPMiddleware):
                                 # If it's not valid JSON, just pass it as a string
                                 content = str(chunk)
                                 
-                                # Buffer content to check for tool calls
-                                content_buffer += content
-                                
-                                # Check if we have a complete tool call or if we should flush the buffer
-                                if self._should_flush_buffer(content_buffer):
-                                    # If this contains a complete tool call, execute it and send the result
-                                    if self._contains_complete_tool_call(content_buffer):
-                                        # First, send the complete tool call to the frontend as JSON
-                                        yield f"data: {json.dumps({'tool_call': content_buffer})}\n\n"
-                                        
-                                        try:
-                                            # Execute the tool call
-                                            logger.info(f"Executing tool call in streaming middleware: {content_buffer[:100]}...")
-                                            from app.mcp.consolidated import execute_mcp_tools_with_status
-                                            tool_result = await execute_mcp_tools_with_status(content_buffer)
-                                            logger.info(f"Tool execution result: {tool_result[:100]}...")
-                                            
-                                            # Send the tool result to the frontend
-                                            yield f"data: {json.dumps({'tool_result': tool_result})}\n\n"
-                                        except Exception as tool_error:
-                                            logger.error(f"Error executing tool call: {tool_error}")
-                                            # Send error message
-                                            error_msg = f"\n\n```tool:error\n❌ **Tool Error:** {str(tool_error)}\n```\n\n"
-                                            yield f"data: {json.dumps({'tool_error': error_msg})}\n\n"
-                                        
-                                        # Clear buffer
-                                        content_buffer = ""
-                                    else:
-                                        # Send buffered content as JSON, not raw content
+                                # Immediately pass through tool_start messages without buffering
+                                if '"type": "tool_start"' in content:
+                                    # First flush any buffered content
+                                    if content_buffer:
                                         yield f"data: {json.dumps({'content': content_buffer})}\n\n"
                                         content_buffer = ""
-                                elif self._contains_partial(content_buffer):
-                                    # Still accumulate partial content
-                                    accumulated_content += content
-                                    # Hold the content in buffer, don't send yet
+                                    # Then send the tool_start message
+                                    yield content
                                     continue
-                                else:
-                                    # Safe to send immediately as raw content
-                                    yield f"data: {content}\n\n"
-                                    content_buffer = ""
+                                
+                                # Send content immediately for real-time streaming
+                                yield f"data: {json.dumps({'content': content})}\\n\\n"
                         
                         # Log chunk content preview
                         if len(chunk) > 200:
@@ -247,6 +225,16 @@ class StreamingMiddleware(BaseHTTPMiddleware):
                             # For simple string content
                             content = str(raw_content)
                             
+                            # Immediately pass through tool_start messages without buffering
+                            if '"type": "tool_start"' in content:
+                                # First flush any buffered content
+                                if content_buffer:
+                                    yield f"data: {json.dumps({'content': content_buffer})}\n\n"
+                                    content_buffer = ""
+                                # Then send the tool_start message
+                                yield content
+                                continue
+                            
                             # Buffer content to check for tool calls
                             content_buffer += content
                             
@@ -354,7 +342,7 @@ class StreamingMiddleware(BaseHTTPMiddleware):
                                 logger.warning(f"🔍 LARGE_CHUNK_DETECTED: {chunk_size} chars - monitoring for truncation")
                             
                             # Check if this is a tool result chunk
-                            if "tool_execution" in chunk or "tool_result" in chunk:
+                            if "tool_execution" in chunk or "tool_display" in chunk or "tool_result" in chunk:
                                 logger.info(f"🔍 TOOL_RESULT_CHUNK: size={chunk_size}, content preview: {chunk[:100]}...")
                             
                             # Check if it might be JSON
@@ -368,7 +356,7 @@ class StreamingMiddleware(BaseHTTPMiddleware):
                                 # DEBUGGING: Track large JSON objects
                                 if len(chunk_content) > 5000:
                                     logger.warning(f"🔍 MIDDLEWARE_LARGE_JSON: {len(chunk_content)} chars, type={json_obj.get('type')}")
-                                    if json_obj.get('type') == 'tool_execution':
+                                if json_obj.get('type') in ['tool_execution', 'tool_display']:
                                         result_size = len(json_obj.get('result', ''))
                                         logger.warning(f"🔍 MIDDLEWARE_TOOL_RESULT: tool={json_obj.get('tool_name')}, result_size={result_size}")
                                         if result_size == 0:

app/routes/file_validation.py

@@ -0,0 +1,35 @@
+"""
+File validation API endpoint.
+"""
+
+import os
+from typing import List
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
+
+router = APIRouter()
+
+class FileValidationRequest(BaseModel):
+    files: List[str]
+
+class FileValidationResponse(BaseModel):
+    existingFiles: List[str]
+    missingFiles: List[str]
+
+@router.post("/api/files/validate", response_model=FileValidationResponse)
+async def validate_files(request: FileValidationRequest):
+    """Validate which files exist and return lists of existing/missing files."""
+    
+    existing_files = []
+    missing_files = []
+    
+    for file_path in request.files:
+        if os.path.exists(file_path):
+            existing_files.append(file_path)
+        else:
+            missing_files.append(file_path)
+    
+    return FileValidationResponse(
+        existingFiles=existing_files,
+        missingFiles=missing_files
+    )

app/routes/mcp_routes.py

@@ -283,11 +283,15 @@ async def get_shell_config():
             server_env = server_config.get("env", {})
             
             # Extract allowed commands from environment configuration
+            # Always start with the full default command list
             allowed_commands = DEFAULT_SHELL_CONFIG["allowedCommands"].copy()
-            if "ALLOW_COMMANDS" in server_env:
+            
+            # Only override if environment explicitly provides a DIFFERENT set
+            if "ALLOW_COMMANDS" in server_env and server_env["ALLOW_COMMANDS"].strip():
                 # If environment override exists, use it instead
                 env_commands = [cmd.strip() for cmd in server_env["ALLOW_COMMANDS"].split(",") if cmd.strip()]
-                if env_commands:
+                # Only replace if the environment list is substantially different (not just a subset)
+                if env_commands and len(env_commands) >= len(DEFAULT_SHELL_CONFIG["allowedCommands"]) * 0.8:
                     allowed_commands = env_commands
                     logger.info(f"Using environment override commands: {allowed_commands}")
             
@@ -302,14 +306,20 @@ async def get_shell_config():
             
             return {
                 "enabled": True,
+                "allowedCommands": allowed_commands,
                 "gitOperationsEnabled": git_operations_enabled,
-                "safeGitOperations": git_operations
+                "safeGitOperations": git_operations,
+                "timeout": timeout
             }
         else:
             # Shell server not connected, return default config with enabled=False
+            # But still check for environment timeout override
+            default_config = get_default_shell_config()
+            timeout = int(os.environ.get("COMMAND_TIMEOUT", default_config["timeout"]))
             return {
-                **get_default_shell_config(),
-                "enabled": False
+                **default_config,
+                "enabled": False,
+                "timeout": timeout
             }
         
     except Exception as e:
@@ -341,7 +351,7 @@ async def update_shell_config(config: ShellConfig):
         
         # Create new shell server configuration
         new_shell_config = {
-            "command": ["python", "-u", "mcp_servers/shell_server.py"],
+            "command": ["python", "-u", "app/mcp_servers/shell_server.py"],
             "enabled": config.enabled,
             "env": {
                 "ALLOW_COMMANDS": ",".join(config.allowedCommands),
@@ -352,17 +362,38 @@ async def update_shell_config(config: ShellConfig):
         }
         
         if config.enabled:
+            # Update the server configuration in MCP manager before restarting
+            mcp_manager.server_configs["shell"] = new_shell_config
+            
             # Restart the shell server with new configuration
+            logger.info(f"Attempting to restart shell server with config: {new_shell_config}")
             success = await mcp_manager.restart_server("shell", new_shell_config)
             
+            if not success:
+                # Get detailed error information
+                shell_client = mcp_manager.clients.get("shell")
+                error_details = "Unknown error"
+                if shell_client and shell_client.process:
+                    try:
+                        stderr_output = shell_client.process.stderr.read() if shell_client.process.stderr else ""
+                        if stderr_output:
+                            error_details = f"Shell server stderr: {stderr_output}"
+                    except Exception as e:
+                        error_details = f"Could not read shell server error: {str(e)}"
+                
+                logger.error(f"Shell server restart failed. Details: {error_details}")
+                return {"success": False, "message": f"Failed to restart shell server: {error_details}"}
+            
             if success:
                 logger.info(f"Shell server restarted with new config: {config.allowedCommands}")
+                
+                # Invalidate tools cache to ensure fresh tool list with updated shell tools
+                mcp_manager.invalidate_tools_cache()
+                
                 return {
                     "success": True, 
                     "message": f"Shell server updated instantly. Basic commands: {', '.join(config.allowedCommands)}, Git operations: {'enabled' if config.gitOperationsEnabled else 'disabled'}"
                 }
-            else:
-                return {"success": False, "message": "Failed to restart shell server"}
         else:
             # Disable shell server by disconnecting it
             if "shell" in mcp_manager.clients:
@@ -375,6 +406,9 @@ async def update_shell_config(config: ShellConfig):
                 mcp_manager.server_configs["shell"]["enabled"] = False
                 logger.info("Shell server marked as disabled in server configs")
             
+            # Invalidate tools cache to ensure fresh tool list without shell tools
+            mcp_manager.invalidate_tools_cache()
+            
             return {"success": True, "message": "Shell server disabled instantly"}
         
     except Exception as e:
@@ -421,6 +455,10 @@ async def toggle_server(request: ServerToggleRequest):
             if server_config:
                 success = await mcp_manager.restart_server(request.server_name, server_config)
                 message = f"{request.server_name} server enabled and restarted" if success else f"Failed to restart {request.server_name} server"
+                
+                # Invalidate tools cache to ensure fresh tool list with newly enabled server tools
+                if success:
+                    mcp_manager.invalidate_tools_cache()
             else:
                 message = f"No configuration found for {request.server_name} server"
         else:
@@ -429,6 +467,14 @@ async def toggle_server(request: ServerToggleRequest):
                 await mcp_manager.clients[request.server_name].disconnect()
                 del mcp_manager.clients[request.server_name]
                 logger.info(f"{request.server_name} server disabled")
+            
+            # Update server config to mark as disabled
+            if request.server_name in mcp_manager.server_configs:
+                mcp_manager.server_configs[request.server_name]["enabled"] = False
+            
+            # Invalidate tools cache to ensure fresh tool list without disabled server tools
+            mcp_manager.invalidate_tools_cache()
+            
             message = f"{request.server_name} server disabled"
         
         return {"success": True, "message": message}