Deploy Agent-Zero with HF adaptations (main branch) and security fixes

helpers/api.py

@@ -2,62 +2,226 @@ from abc import abstractmethod
 import json
 import threading
 import os
-from typing import Union, TypedDict, Dict, Any
-from flask import Request, Response, Flask, request
-from python.helpers.print_style import PrintStyle
-from python.helpers.errors import format_error
+from functools import wraps
+from pathlib import Path
+from typing import Union, Dict, Any
+from flask import (
+    Request,
+    Response,
+    jsonify,
+    Flask,
+    session,
+    request,
+    send_file,
+    redirect,
+    url_for,
+)
+from werkzeug.wrappers.response import Response as BaseResponse
+from agent import AgentContext
+from helpers.print_style import PrintStyle
+from helpers.errors import format_error
+from helpers import files, cache
+
+ThreadLockType = Union[threading.Lock, threading.RLock]
+
+CACHE_AREA = "api_handlers(api)"
 
 Input = dict
-Output = Union[Dict[str, Any], Response, TypedDict]
+Output = Union[Dict[str, Any], Response]
+
 
 class ApiHandler:
-    def __init__(self, app: Flask, thread_lock: threading.Lock):
+    def __init__(self, app: Flask, thread_lock: ThreadLockType):
         self.app = app
         self.thread_lock = thread_lock
 
     @classmethod
-    def requires_loopback(cls) -> bool: return False
+    def requires_loopback(cls) -> bool:
+        return False
+
     @classmethod
-    def requires_api_key(cls) -> bool: return False
+    def requires_api_key(cls) -> bool:
+        return False
+
     @classmethod
     def requires_auth(cls) -> bool:
-        return False if os.getenv("HF_SPACE") == "true" else True
+        return True
+
     @classmethod
-    def get_methods(cls) -> list[str]: return ["POST", "GET"]
+    def get_methods(cls) -> list[str]:
+        return ["POST"]
+
     @classmethod
     def requires_csrf(cls) -> bool:
-        return False if os.getenv("HF_SPACE") == "true" else cls.requires_auth()
+        return cls.requires_auth()
 
     @abstractmethod
-    async def process(self, input: Input, request: Request) -> Output: pass
+    async def process(self, input: Input, request: Request) -> Output:
+        pass
 
     async def handle_request(self, request: Request) -> Response:
         try:
             input_data: Input = {}
             if request.is_json:
-                if request.data: input_data = request.get_json()
+                try:
+                    if request.data:
+                        input_data = request.get_json()
+                except Exception as e:
+                    PrintStyle().print(f"Error parsing JSON: {str(e)}")
+                    input_data = {}
             else:
-                input_data = request.args.to_dict()
-                if request.method == "POST":
-                    input_data.update({"data": request.get_data(as_text=True)})
+                input_data = {}
 
             output = await self.process(input_data, request)
 
-            if isinstance(output, Response): return output
-            return Response(response=json.dumps(output), status=200, mimetype="application/json")
+            if isinstance(output, Response):
+                return output
+            else:
+                response_json = json.dumps(output)
+                return Response(
+                    response=response_json, status=200, mimetype="application/json"
+                )
         except Exception as e:
             error = format_error(e)
             PrintStyle.error(f"API error: {error}")
             return Response(response=error, status=500, mimetype="text/plain")
 
-    def get_context(self, ctxid: str):
-        from agent import AgentContext
-        from initialize import initialize_agent
-        with self.thread_lock:
-            if not ctxid:
-                first = AgentContext.first()
-                if first: return first
-                return AgentContext(config=initialize_agent())
-            got = AgentContext.get(ctxid)
-            if got: return got
-            return AgentContext(config=initialize_agent(), id=ctxid)
+    def use_context(self, ctxid: str, create_if_not_exists: bool = True):
+        from helpers.context_utils import use_context as _use_context
+        return _use_context(self.thread_lock, ctxid, create_if_not_exists)
+
+
+from helpers.network import is_loopback_address
+
+
+def requires_api_key(f):
+    @wraps(f)
+    async def decorated(*args, **kwargs):
+        if os.getenv("HF_SPACE") == "true":
+             return await f(*args, **kwargs)
+             
+        from helpers.settings import get_settings
+        valid_api_key = get_settings()["mcp_server_token"]
+
+        if api_key := request.headers.get("X-API-KEY"):
+            if api_key != valid_api_key:
+                return Response("Invalid API key", 401)
+        elif request.json and request.json.get("api_key"):
+            api_key = request.json.get("api_key")
+            if api_key != valid_api_key:
+                return Response("Invalid API key", 401)
+        else:
+            return Response("API key required", 401)
+        return await f(*args, **kwargs)
+
+    return decorated
+
+
+def requires_loopback(f):
+    @wraps(f)
+    async def decorated(*args, **kwargs):
+        if os.getenv("HF_SPACE") == "true":
+             return await f(*args, **kwargs)
+             
+        if not is_loopback_address(str(request.remote_addr)):
+            return Response("Access denied.", 403, {})
+        return await f(*args, **kwargs)
+
+    return decorated
+
+
+def requires_auth(f):
+    @wraps(f)
+    async def decorated(*args, **kwargs):
+        if os.getenv("HF_SPACE") == "true":
+             return await f(*args, **kwargs)
+             
+        from helpers import login
+        user_pass_hash = login.get_credentials_hash()
+        if not user_pass_hash:
+            return await f(*args, **kwargs)
+        if session.get("authentication") != user_pass_hash:
+            return redirect(url_for("login_handler"))
+        return await f(*args, **kwargs)
+
+    return decorated
+
+
+def csrf_protect(f):
+    @wraps(f)
+    async def decorated(*args, **kwargs):
+        if os.getenv("HF_SPACE") == "true":
+             return await f(*args, **kwargs)
+             
+        from helpers import runtime
+        token = session.get("csrf_token")
+        header = request.headers.get("X-CSRF-Token")
+        cookie = request.cookies.get("csrf_token_" + runtime.get_runtime_id())
+        sent = header or cookie
+        if not token or not sent or token != sent:
+            return Response("CSRF token missing or invalid", 403)
+        return await f(*args, **kwargs)
+
+    return decorated
+
+
+def register_api_route(app: Flask, lock: ThreadLockType) -> None:
+    from helpers.modules import load_classes_from_file
+    from helpers import plugins
+
+    async def _dispatch(path: str) -> BaseResponse:
+        cached = cache.get(CACHE_AREA, path)
+        if cached is not None:
+            return await cached()
+
+        handler_cls: type[ApiHandler] | None = None
+
+        builtin_file = files.get_abs_path(f"api/{path}.py")
+        if files.is_in_dir(builtin_file, files.get_abs_path("api")) and files.exists(
+            builtin_file
+        ):
+            classes = load_classes_from_file(builtin_file, ApiHandler)
+            if classes:
+                handler_cls = classes[0]
+
+        if handler_cls is None and path.startswith("plugins/"):
+            parts = path.split("/", 2)
+            if len(parts) == 3:
+                _, plugin_name, handler_name = parts
+                plugin_dir = plugins.find_plugin_dir(plugin_name)
+                if plugin_dir:
+                    plugin_file = Path(plugin_dir) / "api" / f"{handler_name}.py"
+                    if plugin_file.is_file():
+                        classes = load_classes_from_file(str(plugin_file), ApiHandler)
+                        if classes:
+                            handler_cls = classes[0]
+
+        if handler_cls is None:
+            return Response(f"API endpoint not found: {path}", 404)
+
+        if request.method not in handler_cls.get_methods():
+            return Response(f"Method {request.method} not allowed for: {path}", 405)
+
+        async def call_handler() -> BaseResponse:
+            instance = handler_cls(app, lock)
+            return await instance.handle_request(request=request)
+
+        handler_fn = call_handler
+        if handler_cls.requires_csrf():
+            handler_fn = csrf_protect(handler_fn)
+        if handler_cls.requires_api_key():
+            handler_fn = requires_api_key(handler_fn)
+        if handler_cls.requires_auth():
+            handler_fn = requires_auth(handler_fn)
+        if handler_cls.requires_loopback():
+            handler_fn = requires_loopback(handler_fn)
+
+        cache.add(CACHE_AREA, path, handler_fn)
+        return await handler_fn()
+
+    app.add_url_rule(
+        "/api/<path:path>",
+        "api_dispatch",
+        _dispatch,
+        methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
+    )

helpers/ui_server.py

@@ -0,0 +1,302 @@
+from dataclasses import dataclass, field
+from datetime import timedelta
+import asyncio
+import logging
+import os
+import secrets
+import threading
+import time
+from typing import Any
+
+from flask import (
+    Flask,
+    Response,
+    redirect,
+    render_template_string,
+    request,
+    send_file,
+    session,
+    url_for,
+)
+from socketio import ASGIApp
+from starlette.applications import Starlette
+from starlette.routing import Mount, Route
+from starlette.responses import JSONResponse
+from uvicorn.middleware.wsgi import WSGIMiddleware
+from werkzeug.wrappers.request import Request as WerkzeugRequest
+import socketio  # type: ignore[import-untyped]
+
+from helpers import dotenv, fasta2a_server, files, git, login, mcp_server, runtime
+from helpers.api import register_api_route, requires_auth
+from helpers.extension import extensible
+from helpers.files import get_abs_path
+from helpers.print_style import PrintStyle
+from helpers.server_startup import StartupMonitor
+from helpers import settings as settings_helper
+from helpers.ws import register_ws_namespace, validate_ws_origin
+from helpers.ws_manager import WsManager, set_shared_ws_manager
+
+
+UPLOAD_LIMIT_BYTES = 5 * 1024 * 1024 * 1024
+
+
+def configure_process_environment() -> None:
+    logging.getLogger().setLevel(logging.WARNING)
+    os.environ["TZ"] = "UTC"
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    if hasattr(time, "tzset"):
+        time.tzset()
+
+
+@dataclass
+class UiServerRuntime:
+    webapp: Flask
+    socketio_server: socketio.AsyncServer
+    ws_manager: WsManager
+    lock: threading.RLock
+    settings_snapshot: dict[str, Any]
+    _routes_registered: bool = False
+    _transport_registered: bool = False
+    _route_handlers: "UiRouteHandlers | None" = field(default=None, init=False)
+
+    @classmethod
+    def create(cls) -> "UiServerRuntime":
+        webapp = Flask("app", static_folder=get_abs_path("./webui"), static_url_path="/")
+        webapp.secret_key = os.getenv("FLASK_SECRET_KEY") or secrets.token_hex(32)
+
+        if os.getenv("HF_SPACE") == "true":
+            from werkzeug.middleware.proxy_fix import ProxyFix
+            webapp.wsgi_app = ProxyFix(webapp.wsgi_app, x_for=1, x_proto=1, x_host=1)
+
+        WerkzeugRequest.max_form_memory_size = UPLOAD_LIMIT_BYTES
+        webapp.config.update(
+            JSON_SORT_KEYS=False,
+            SESSION_COOKIE_NAME="session_" + runtime.get_runtime_id(),
+            SESSION_COOKIE_SAMESITE="Lax",
+            SESSION_PERMANENT=True,
+            PERMANENT_SESSION_LIFETIME=timedelta(days=1),
+            MAX_CONTENT_LENGTH=int(
+                os.getenv("FLASK_MAX_CONTENT_LENGTH", str(UPLOAD_LIMIT_BYTES))
+            ),
+            MAX_FORM_MEMORY_SIZE=int(
+                os.getenv("FLASK_MAX_FORM_MEMORY_SIZE", str(UPLOAD_LIMIT_BYTES))
+            ),
+        )
+
+        lock = threading.RLock()
+        
+        cors_allowed = "*" if os.getenv("HF_SPACE") == "true" else (lambda _origin, environ: validate_ws_origin(environ)[0])
+        
+        socketio_server = socketio.AsyncServer(
+            async_mode="asgi",
+            namespaces="*",
+            cors_allowed_origins=cors_allowed,
+            logger=False,
+            engineio_logger=False,
+            ping_interval=25,
+            ping_timeout=30,
+            max_http_buffer_size=50 * 1024 * 1024,
+        )
+
+        ws_manager = WsManager(socketio_server, lock)
+        set_shared_ws_manager(ws_manager)
+
+        server_runtime = cls(
+            webapp=webapp,
+            socketio_server=socketio_server,
+            ws_manager=ws_manager,
+            lock=lock,
+            settings_snapshot={},
+        )
+        server_runtime.refresh_runtime_settings()
+        return server_runtime
+
+    def refresh_runtime_settings(self) -> None:
+        self.settings_snapshot = settings_helper.get_settings()
+        settings_helper.set_runtime_settings_snapshot(self.settings_snapshot)
+        self.ws_manager.set_server_restart_broadcast(
+            self.settings_snapshot.get("websocket_server_restart_enabled", True)
+        )
+
+    def register_http_routes(self) -> None:
+        if self._routes_registered:
+            return
+
+        handlers = UiRouteHandlers(self)
+        self._route_handlers = handlers
+        self.webapp.add_url_rule(
+            "/login",
+            "login_handler",
+            handlers.login_handler,
+            methods=["GET", "POST"],
+        )
+        self.webapp.add_url_rule(
+            "/logout",
+            "logout_handler",
+            handlers.logout_handler,
+            methods=["GET"],
+        )
+        self.webapp.add_url_rule(
+            "/",
+            "serve_index",
+            handlers.serve_index,
+            methods=["GET"],
+        )
+        self.webapp.add_url_rule(
+            "/plugins/<plugin_name>/<path:asset_path>",
+            "serve_builtin_plugin_asset",
+            handlers.serve_builtin_plugin_asset,
+            methods=["GET"],
+        )
+        self.webapp.add_url_rule(
+            "/usr/plugins/<plugin_name>/<path:asset_path>",
+            "serve_plugin_asset",
+            handlers.serve_plugin_asset,
+            methods=["GET"],
+        )
+        self.webapp.add_url_rule(
+            "/extensions/webui/<path:asset_path>",
+            "serve_extension_asset",
+            handlers.serve_extension_asset,
+            methods=["GET"],
+        )
+        self._routes_registered = True
+
+    def register_transport_handlers(self) -> None:
+        if self._transport_registered:
+            return
+        register_api_route(self.webapp, self.lock)
+        register_ws_namespace(
+            self.socketio_server,
+            self.webapp,
+            self.lock,
+            manager=self.ws_manager,
+        )
+        self._transport_registered = True
+
+    def build_asgi_app(self, startup_monitor: StartupMonitor):
+        with startup_monitor.stage("wsgi.middleware.create"):
+            wsgi_app = WSGIMiddleware(self.webapp)
+
+        with startup_monitor.stage("mcp.proxy.init"):
+            mcp_app = mcp_server.DynamicMcpProxy.get_instance()
+
+        with startup_monitor.stage("a2a.proxy.init"):
+            a2a_app = fasta2a_server.DynamicA2AProxy.get_instance()
+
+        async def health_endpoint(request):
+            return JSONResponse({"status": "ok", "service": "agent-zero"})
+
+        async def api_docs_endpoint(request):
+            from api.api_docs import ApiDocs
+            handler = ApiDocs(self.webapp, self.lock)
+            result = await handler.process({}, None)
+            return JSONResponse(result)
+
+        with startup_monitor.stage("starlette.app.create"):
+            starlette_app = Starlette(
+                routes=[
+                    Route("/health", endpoint=health_endpoint, methods=["GET"]),
+                    Route("/api-docs", endpoint=api_docs_endpoint, methods=["GET"]),
+                    Mount("/mcp", app=mcp_app),
+                    Mount("/a2a", app=a2a_app),
+                    Mount("/", app=wsgi_app),
+                ],
+                lifespan=startup_monitor.lifespan(),
+            )
+
+        with startup_monitor.stage("socketio.asgi.create"):
+            return ASGIApp(self.socketio_server, other_asgi_app=starlette_app)
+
+    def access_log_enabled(self) -> bool:
+        return self.settings_snapshot.get("uvicorn_access_logs_enabled", False)
+
+
+class UiRouteHandlers:
+    def __init__(self, runtime_state: UiServerRuntime) -> None:
+        self.runtime = runtime_state
+
+    @extensible
+    async def login_handler(self):
+        error = None
+        if request.method == "POST":
+            user = dotenv.get_dotenv_value("AUTH_LOGIN")
+            password = dotenv.get_dotenv_value("AUTH_PASSWORD")
+
+            if request.form["username"] == user and request.form["password"] == password:
+                session["authentication"] = login.get_credentials_hash()
+                return redirect(url_for("serve_index"))
+            else:
+                await asyncio.sleep(1)
+                error = "Invalid Credentials. Please try again."
+
+        login_page_content = files.read_file("webui/login.html")
+        return render_template_string(login_page_content, error=error)
+
+    @extensible
+    async def logout_handler(self):
+        session.pop("authentication", None)
+        return redirect(url_for("login_handler"))
+
+    @requires_auth
+    @extensible
+    async def serve_index(self):
+        try:
+            gitinfo = git.get_git_info()
+        except Exception:
+            gitinfo = {
+                "version": "unknown",
+                "commit_time": "unknown",
+            }
+
+        index = files.read_file("webui/index.html")
+        return files.replace_placeholders_text(
+            _content=index,
+            version_no=gitinfo["version"],
+            version_time=gitinfo["commit_time"],
+            runtime_id=runtime.get_runtime_id(),
+            runtime_is_development=("true" if runtime.is_development() else "false"),
+            logged_in=("true" if login.get_credentials_hash() else "false"),
+        )
+
+    @requires_auth
+    async def serve_builtin_plugin_asset(self, plugin_name, asset_path):
+        return await self._serve_plugin_asset(plugin_name, asset_path)
+
+    @requires_auth
+    async def serve_plugin_asset(self, plugin_name, asset_path):
+        return await self._serve_plugin_asset(plugin_name, asset_path)
+
+    @requires_auth
+    async def serve_extension_asset(self, asset_path):
+        exts = files.get_abs_path("extensions/webui")
+        path = files.get_abs_path(exts, asset_path)
+        if not files.is_in_dir(path, exts):
+            return Response("Access denied", 403)
+        return send_file(path)
+
+    @extensible
+    async def _serve_plugin_asset(self, plugin_name, asset_path):
+        from helpers import plugins
+
+        plugin_dir = plugins.find_plugin_dir(plugin_name)
+        if not plugin_dir:
+            return Response("Plugin not found", 404)
+
+        try:
+            asset_file = files.get_abs_path(plugin_dir, asset_path)
+            webui_dir = files.get_abs_path(plugin_dir, "webui")
+            webui_extensions_dir = files.get_abs_path(plugin_dir, "extensions/webui")
+
+            if not files.is_in_dir(str(asset_file), str(webui_dir)) and not files.is_in_dir(
+                str(asset_file), str(webui_extensions_dir)
+            ):
+                return Response("Access denied", 403)
+
+            if not files.is_file(asset_file):
+                return Response("Asset not found", 404)
+
+            return send_file(str(asset_file))
+        except Exception as e:
+            PrintStyle.error(f"Error serving plugin asset: {e}")
+            return Response("Error serving asset", 500)

start_hf.sh

@@ -5,33 +5,27 @@ REPO_URL="https://github.com/JsonLord/agent-zero.git"
 CLONE_DIR="/home/user/app/agent-zero-clone"
 APP_DIR="/home/user/app"
 
-echo "Cloning Agent-Zero (fix-a2a-auth-15797037665282663574 branch)..."
+echo "Cloning Agent-Zero (main branch)..."
 if [ -d "$CLONE_DIR" ]; then
     rm -rf "$CLONE_DIR"
 fi
 
-git clone --branch fix-a2a-auth-15797037665282663574 "$REPO_URL" "$CLONE_DIR"
+git clone "$REPO_URL" "$CLONE_DIR"
 
 echo "Applying adaptations..."
 cp -rn "$CLONE_DIR"/* "$APP_DIR/" 2>/dev/null || true
 
-cp -f /home/user/app/helpers/api.py "$APP_DIR/python/helpers/api.py"
-cp -f /home/user/app/helpers/runtime.py "$APP_DIR/python/helpers/runtime.py"
-cp -f /home/user/app/helpers/settings.py "$APP_DIR/python/helpers/settings.py"
-cp -f /home/user/app/run_ui.py "$APP_DIR/run_ui.py"
-cp -f /home/user/app/api/api_docs.py "$APP_DIR/python/api/api_docs.py"
-
-if [ -f "/home/user/app/helpers/ws.py" ]; then
-    cp -f /home/user/app/helpers/ws.py "$APP_DIR/python/helpers/ws.py"
-fi
+cp -f /home/user/app/helpers/api.py "$APP_DIR/helpers/api.py"
+cp -f /home/user/app/helpers/runtime.py "$APP_DIR/helpers/runtime.py"
+cp -f /home/user/app/helpers/settings.py "$APP_DIR/helpers/settings.py"
+cp -f /home/user/app/helpers/ui_server.py "$APP_DIR/helpers/ui_server.py"
+cp -f /home/user/app/api/api_docs.py "$APP_DIR/api/api_docs.py"
 
 echo "Installing dependencies..."
 if command -v uv > /dev/null; then
     uv pip install --system --no-cache -r "$APP_DIR/requirements.txt"
-    uv pip install --system flask-socketio eventlet
 else
     pip install --no-cache-dir -r "$APP_DIR/requirements.txt"
-    pip install flask-socketio eventlet
 fi
 
 export HF_SPACE=true
@@ -39,4 +33,4 @@ export PORT=7860
 export HOST=0.0.0.0
 
 echo "Starting Agent-Zero..."
-python run_ui.py --host 0.0.0.0 --port 7860 --dockerized=true
+python run_ui.py