Better initialization

agent.py

@@ -395,6 +395,8 @@ class Agent:
                 await self.call_extensions("monologue_end", loop_data=self.loop_data)  # type: ignore
 
     async def prepare_prompt(self, loop_data: LoopData) -> ChatPromptTemplate:
+        self.context.log.set_progress("Building prompt")
+
         # call extensions before setting prompts
         await self.call_extensions("message_loop_prompts_before", loop_data=loop_data)
 
@@ -720,18 +722,10 @@ class Agent:
                 if mcp_tool_candidate:
                     tool = mcp_tool_candidate
             except ImportError:
-                 # Get context safely
-                 current_context = AgentContext.first()
-                 if current_context:
-                    current_context.log.log(type="warning", content="MCP helper module not found. Skipping MCP tool lookup.", temp=True)
-                 PrintStyle(background_color="black", font_color="yellow", padding=True).print(
+                PrintStyle(background_color="black", font_color="yellow", padding=True).print(
                     "MCP helper module not found. Skipping MCP tool lookup."
                  )
             except Exception as e:
-                # Get context safely
-                current_context = AgentContext.first()
-                if current_context:
-                    current_context.log.log(type="warning", content=f"Failed to get MCP tool '{tool_name}': {e}", temp=True)
                 PrintStyle(background_color="black", font_color="red", padding=True).print(
                     f"Failed to get MCP tool '{tool_name}': {e}"
                 )

initialize.py

@@ -1,16 +1,10 @@
-import asyncio
-import json
 import models
 from agent import AgentConfig, ModelConfig
-from python.helpers import dotenv, files, rfc_exchange, runtime, settings, docker, log, defer
-import subprocess
-import shutil
+from python.helpers import runtime, settings, defer
 from python.helpers.print_style import PrintStyle
-from python.helpers.mcp_handler import initialize_mcp
 
 
-
-def initialize():
+def initialize_agent():
     current_settings = settings.get_settings()
 
     # chat model from user settings
@@ -76,16 +70,16 @@ def initialize():
     )
 
     # update SSH and docker settings
-    set_runtime_config(config, current_settings)
+    _set_runtime_config(config, current_settings)
 
     # update config with runtime args
-    args_override(config)
+    _args_override(config)
 
     # initialize MCP in deferred task to prevent blocking the main thread
     # async def initialize_mcp_async(mcp_servers_config: str):
     #     return initialize_mcp(mcp_servers_config)
     # defer.DeferredTask(thread_name="mcp-initializer").start_task(initialize_mcp_async, config.mcp_servers)
-    initialize_mcp(config.mcp_servers)
+    # initialize_mcp(config.mcp_servers)
 
     # import python.helpers.mcp_handler as mcp_helper
     # import agent as agent_helper
@@ -108,8 +102,25 @@ def initialize():
     # return config object
     return config
 
+def initialize_chats():
+    from python.helpers import persist_chat
+    async def initialize_chats_async():
+        persist_chat.load_tmp_chats()
+    return defer.DeferredTask().start_task(initialize_chats_async)
+
+def initialize_mcp():
+    set = settings.get_settings()
+    async def initialize_mcp_async():
+        from python.helpers.mcp_handler import initialize_mcp as _initialize_mcp
+        return _initialize_mcp(set["mcp_servers"])
+    return defer.DeferredTask().start_task(initialize_mcp_async)
+
+def initialize_job_loop():
+    from python.helpers.job_loop import run_loop
+    return defer.DeferredTask("JobLoop").start_task(run_loop)
+
 
-def args_override(config):
+def _args_override(config):
     # update config with runtime args
     for key, value in runtime.args.items():
         if hasattr(config, key):
@@ -130,7 +141,7 @@ def args_override(config):
             setattr(config, key, value)
 
 
-def set_runtime_config(config: AgentConfig, set: settings.Settings):
+def _set_runtime_config(config: AgentConfig, set: settings.Settings):
     ssh_conf = settings.get_runtime_config(set)
     for key, value in ssh_conf.items():
         if hasattr(config, key):

python/extensions/system_prompt/_10_system_prompt.py

@@ -16,7 +16,8 @@ class SystemPrompt(Extension):
 
         system_prompt.append(main)
         system_prompt.append(tools)
-        system_prompt.append(mcp_tools)
+        if mcp_tools:
+            system_prompt.append(mcp_tools)
 
 
 def get_main_prompt(agent: Agent):
@@ -31,4 +32,14 @@ def get_tools_prompt(agent: Agent):
 
 
 def get_mcp_tools_prompt(agent: Agent):
-    return MCPConfig.get_instance().get_tools_prompt()
+    mcp_config = MCPConfig.get_instance()
+    if mcp_config.servers:
+        pre_progress = agent.context.log.progress
+        agent.context.log.set_progress("Collecting MCP tools") # MCP might be initializing, better inform via progress bar
+        # import time
+        # time.sleep(10)
+        tools = MCPConfig.get_instance().get_tools_prompt()
+        agent.context.log.set_progress(pre_progress) # return original progress
+        return tools
+    return ""
+        

python/helpers/api.py

@@ -5,7 +5,7 @@ from typing import Union, TypedDict, Dict, Any
 from attr import dataclass
 from flask import Request, Response, jsonify, Flask
 from agent import AgentContext
-from initialize import initialize
+from initialize import initialize_agent
 from python.helpers.print_style import PrintStyle
 from python.helpers.errors import format_error
 from werkzeug.serving import make_server
@@ -77,8 +77,8 @@ class ApiHandler:
                 first = AgentContext.first()
                 if first:
                     return first
-                return AgentContext(config=initialize())
+                return AgentContext(config=initialize_agent())
             got = AgentContext.get(ctxid)
             if got:
                 return got
-            return AgentContext(config=initialize(), id=ctxid)
+            return AgentContext(config=initialize_agent(), id=ctxid)

python/helpers/defer.py

@@ -10,12 +10,12 @@ class EventLoopThread:
     _instances = {}
     _lock = threading.Lock()
 
-    def __init__(self, thread_name: str = "default") -> None:
+    def __init__(self, thread_name: str = "Background") -> None:
         """Initialize the event loop thread."""
         self.thread_name = thread_name
         self._start()
 
-    def __new__(cls, thread_name: str = "default"):
+    def __new__(cls, thread_name: str = "Background"):
         with cls._lock:
             if thread_name not in cls._instances:
                 instance = super(EventLoopThread, cls).__new__(cls)
@@ -59,7 +59,7 @@ class ChildTask:
 class DeferredTask:
     def __init__(
         self,
-        thread_name: str = "default",
+        thread_name: str = "Background",
     ):
         self.event_loop_thread = EventLoopThread(thread_name)
         self._future: Optional[Future] = None
@@ -72,6 +72,7 @@ class DeferredTask:
         self.args = args
         self.kwargs = kwargs
         self._start_task()
+        return self
 
     def __del__(self):
         self.kill()

python/helpers/mcp_handler.py

@@ -62,20 +62,16 @@ def initialize_mcp(mcp_servers_config: str):
         except Exception as e:
             from agent import AgentContext
 
-            first_context = AgentContext.first()  # TODO replace with better reporting
-            if first_context:
-                (
-                    first_context.log.log(
-                        type="warning",
-                        content=f"Failed to update MCP settings: {e}",
-                        temp=False,
-                    )
-                )
-            (
-                PrintStyle(
-                    background_color="black", font_color="red", padding=True
-                ).print(f"Failed to update MCP settings: {e}")
+            AgentContext.log_to_all(
+                type="warning",
+                content=f"Failed to update MCP settings: {e}",
+                temp=False,
             )
+            
+            PrintStyle(
+                background_color="black", font_color="red", padding=True
+            ).print(f"Failed to update MCP settings: {e}")
+            
 
 
 class MCPTool(Tool):
@@ -373,9 +369,10 @@ class MCPConfig(BaseModel):
 
     @classmethod
     def get_instance(cls) -> "MCPConfig":
-        if cls.__instance is None:
-            cls.__instance = cls(servers_list=[])
-        return cls.__instance
+        # with cls.__lock:
+            if cls.__instance is None:
+                cls.__instance = cls(servers_list=[])
+            return cls.__instance
 
     @classmethod
     def wait_for_lock(cls):
@@ -690,6 +687,11 @@ class MCPConfig(BaseModel):
 
     def get_tools_prompt(self, server_name: str = "") -> str:
         """Get a prompt for all tools"""
+
+        # just to wait for pending initialization
+        with self.__lock:
+            pass
+
         prompt = '## "Remote (MCP Server) Agent Tools" available:\n\n'
         server_names = []
         for server in self.servers:

python/helpers/mcp_server.py

@@ -8,7 +8,7 @@ from fastmcp import FastMCP
 
 from agent import AgentContext, AgentContextType, UserMessage
 from python.helpers.persist_chat import save_tmp_chat, remove_chat
-from initialize import initialize
+from initialize import initialize_agent
 from python.helpers.print_style import PrintStyle
 from python.helpers import settings
 from starlette.middleware import Middleware
@@ -139,7 +139,7 @@ async def send_message(
             # If we continue a conversation, it must be persistent
             persistent_chat = True
     else:
-        config = initialize()
+        config = initialize_agent()
         context = AgentContext(config=config, type=AgentContextType.MCP)
 
     if not message:

python/helpers/persist_chat.py

@@ -5,7 +5,7 @@ import uuid
 from agent import Agent, AgentConfig, AgentContext, AgentContextType
 from python.helpers import files, history
 import json
-from initialize import initialize
+from initialize import initialize_agent
 
 from python.helpers.log import Log, LogItem
 
@@ -146,7 +146,7 @@ def _serialize_log(log: Log):
 
 
 def _deserialize_context(data):
-    config = initialize()
+    config = initialize_agent()
     log = _deserialize_log(data.get("log", None))
 
     context = AgentContext(

python/helpers/settings.py

@@ -965,9 +965,9 @@ def _apply_settings(previous: Settings | None):
     global _settings
     if _settings:
         from agent import AgentContext
-        from initialize import initialize
+        from initialize import initialize_agent
 
-        config = initialize()
+        config = initialize_agent()
         for ctx in AgentContext._contexts.values():
             ctx.config = config  # reinitialize context config with new settings
             # apply config to agents

python/helpers/task_scheduler.py

@@ -16,7 +16,7 @@ from crontab import CronTab
 from pydantic import BaseModel, Field, PrivateAttr
 
 from agent import Agent, AgentContext, UserMessage
-from initialize import initialize
+from initialize import initialize_agent
 from python.helpers.persist_chat import save_tmp_chat
 from python.helpers.print_style import PrintStyle
 from python.helpers.defer import DeferredTask
@@ -713,7 +713,7 @@ class TaskScheduler:
         if not task.context_id:
             raise ValueError(f"Task {task.name} has no context ID")
 
-        config = initialize()
+        config = initialize_agent()
         context: AgentContext = AgentContext(config, id=task.context_id, name=task.name)
         # context.id = task.context_id
         # initial name before renaming is same as task name

run_cli.py

@@ -8,7 +8,7 @@ from python.helpers.print_style import PrintStyle
 from python.helpers.files import read_file
 from python.helpers import files
 import python.helpers.timed_input as timed_input
-from initialize import initialize
+from initialize import initialize_agent
 from python.helpers.dotenv import load_dotenv
 
 
@@ -102,7 +102,7 @@ def run():
     load_dotenv()
 
     # initialize context
-    config = initialize()
+    config = initialize_agent()
     context = AgentContext(config)
 
     # Start the key capture thread for user intervention during agent streaming

run_ui.py

@@ -8,15 +8,13 @@ import threading
 import signal
 from flask import Flask, request, Response
 from flask_basicauth import BasicAuth
+import initialize
 from python.helpers import errors, files, git, mcp_server
 from python.helpers.files import get_abs_path
-from python.helpers import persist_chat, runtime, dotenv, process
-from python.helpers.cloudflare_tunnel import CloudflareTunnel
+from python.helpers import runtime, dotenv, process
 from python.helpers.extract_tools import load_classes_from_folder
 from python.helpers.api import ApiHandler
-from python.helpers.job_loop import run_loop
 from python.helpers.print_style import PrintStyle
-from python.helpers.defer import DeferredTask
 
 
 # Set the new timezone to 'UTC'
@@ -148,11 +146,6 @@ def run():
     from werkzeug.serving import make_server
     from werkzeug.middleware.dispatcher import DispatcherMiddleware
     from a2wsgi import ASGIMiddleware, WSGIMiddleware
-    from fastmcp.server.http import create_sse_app
-    from python.helpers.mcp_server import mcp_server as mcp_server_instance
-
-    PrintStyle().print("Starting job loop...")
-    job_loop = DeferredTask().start_task(run_loop)
 
     PrintStyle().print("Starting server...")
 
@@ -165,33 +158,6 @@ def run():
     host = (
         runtime.get_arg("host") or dotenv.get_dotenv_value("WEB_UI_HOST") or "localhost"
     )
-    use_cloudflare = (
-        runtime.get_arg("cloudflare_tunnel")
-        or dotenv.get_dotenv_value("USE_CLOUDFLARE", "false").lower()
-    ) == "true"
-
-    tunnel = None
-
-    try:
-        # Initialize and start Cloudflare tunnel if enabled
-        if use_cloudflare and port:
-            try:
-                tunnel = CloudflareTunnel(port)
-                tunnel.start()
-            except Exception as e:
-                PrintStyle().error(f"Failed to start Cloudflare tunnel: {e}")
-                PrintStyle().print("Continuing without tunnel...")
-
-        # # initialize contexts from persisted chats - moved to async task
-        # persist_chat.load_tmp_chats()
-
-        # # reload scheduler
-        # scheduler = TaskScheduler.get()
-        # asyncio.run(scheduler.reload())
-
-    except Exception as e:
-        PrintStyle().error(errors.format_error(e))
-
     server = None
 
     def register_api_handler(app, handler: type[ApiHandler]):
@@ -243,53 +209,35 @@ def run():
     )
     PrintStyle().debug("Registered middleware for MCP and MCP token")
 
-    try:
-        PrintStyle().debug(f"Starting server at {host}:{port}...")
-
-        server = make_server(
-            host=host,
-            port=port,
-            app=app,
-            request_handler=NoRequestLoggingWSGIRequestHandler,
-            threaded=True,
-        )
-
-        printer = PrintStyle()
+    PrintStyle().debug(f"Starting server at {host}:{port}...")
 
-        def signal_handler(sig=None, frame=None):
-            nonlocal tunnel, server, printer
-            with lock:
-                printer.print("Caught signal, stopping server...")
-                if server:
-                    server.shutdown()
-                process.stop_server()
-                if tunnel:
-                    tunnel.stop()
-                    tunnel = None
-                printer.print("Server stopped")
-                sys.exit(0)
-
-        signal.signal(signal.SIGINT, signal_handler)
-        signal.signal(signal.SIGTERM, signal_handler)
-
-        process.set_server(server)
-        server.log_startup()
-
-        # Start init_a0 in a background thread when server starts
-        import threading
+    server = make_server(
+        host=host,
+        port=port,
+        app=app,
+        request_handler=NoRequestLoggingWSGIRequestHandler,
+        threaded=True,
+    )
+    process.set_server(server)
+    server.log_startup()
 
-        threading.Thread(target=init_a0, daemon=True).start()
+    # Start init_a0 in a background thread when server starts
+    # threading.Thread(target=init_a0, daemon=True).start()
+    init_a0()
 
-        server.serve_forever()
-    finally:
-        # Clean up tunnel if it was started
-        if tunnel:
-            tunnel.stop()
+    # run the server
+    server.serve_forever()
 
 
 def init_a0():
-    # initialize contexts from persisted chats
-    persist_chat.load_tmp_chats()
+    # initialize contexts and MCP
+    init_chats = initialize.initialize_chats()
+    initialize.initialize_mcp()
+    # start job loop
+    initialize.initialize_job_loop()
+
+    # only wait for init chats, otherwise they would seem to dissapear for a while on restart
+    init_chats.result_sync()
 
 
 # run the internal server