Better initialization
frdel committed
Jun 7, 2025 at 22:28 UTC
7528c79fc1ae86ca95114e6c6b27e66d4633a3da
13 files changed
+100
-133
agent.py
+3
-9
@@ -395,6 +395,8 @@ class Agent:
395
await self.call_extensions("monologue_end", loop_data=self.loop_data) # type: ignore
396
397
async def prepare_prompt(self, loop_data: LoopData) -> ChatPromptTemplate:
398
+ self.context.log.set_progress("Building prompt")
399
+
400
# call extensions before setting prompts
401
await self.call_extensions("message_loop_prompts_before", loop_data=loop_data)
402
@@ -720,18 +722,10 @@ class Agent:
722
if mcp_tool_candidate:
723
tool = mcp_tool_candidate
724
except ImportError:
723
- # Get context safely
724
- current_context = AgentContext.first()
725
- if current_context:
726
- current_context.log.log(type="warning", content="MCP helper module not found. Skipping MCP tool lookup.", temp=True)
727
- PrintStyle(background_color="black", font_color="yellow", padding=True).print(
725
+ PrintStyle(background_color="black", font_color="yellow", padding=True).print(
726
"MCP helper module not found. Skipping MCP tool lookup."
727
)
728
except Exception as e:
731
- # Get context safely
732
- current_context = AgentContext.first()
733
- if current_context:
734
- current_context.log.log(type="warning", content=f"Failed to get MCP tool '{tool_name}': {e}", temp=True)
729
PrintStyle(background_color="black", font_color="red", padding=True).print(
730
f"Failed to get MCP tool '{tool_name}': {e}"
731
)
initialize.py
+24
-13
@@ -1,16 +1,10 @@
1
-import asyncio
2
-import json
1
import models
2
from agent import AgentConfig, ModelConfig
5
-from python.helpers import dotenv, files, rfc_exchange, runtime, settings, docker, log, defer
6
-import subprocess
7
-import shutil
3
+from python.helpers import runtime, settings, defer
4
from python.helpers.print_style import PrintStyle
9
-from python.helpers.mcp_handler import initialize_mcp
5
6
12
-
13
-def initialize():
7
+def initialize_agent():
8
current_settings = settings.get_settings()
9
10
# chat model from user settings
@@ -76,16 +70,16 @@ def initialize():
70
)
71
72
# update SSH and docker settings
79
- set_runtime_config(config, current_settings)
73
+ _set_runtime_config(config, current_settings)
74
75
# update config with runtime args
82
- args_override(config)
76
+ _args_override(config)
77
78
# initialize MCP in deferred task to prevent blocking the main thread
79
# async def initialize_mcp_async(mcp_servers_config: str):
80
# return initialize_mcp(mcp_servers_config)
81
# defer.DeferredTask(thread_name="mcp-initializer").start_task(initialize_mcp_async, config.mcp_servers)
88
- initialize_mcp(config.mcp_servers)
82
+ # initialize_mcp(config.mcp_servers)
83
84
# import python.helpers.mcp_handler as mcp_helper
85
# import agent as agent_helper
@@ -108,8 +102,25 @@ def initialize():
102
# return config object
103
return config
104
105
+def initialize_chats():
106
+ from python.helpers import persist_chat
107
+ async def initialize_chats_async():
108
+ persist_chat.load_tmp_chats()
109
+ return defer.DeferredTask().start_task(initialize_chats_async)
110
+
111
+def initialize_mcp():
112
+ set = settings.get_settings()
113
+ async def initialize_mcp_async():
114
+ from python.helpers.mcp_handler import initialize_mcp as _initialize_mcp
115
+ return _initialize_mcp(set["mcp_servers"])
116
+ return defer.DeferredTask().start_task(initialize_mcp_async)
117
+
118
+def initialize_job_loop():
119
+ from python.helpers.job_loop import run_loop
120
+ return defer.DeferredTask("JobLoop").start_task(run_loop)
121
+
122
112
-def args_override(config):
123
+def _args_override(config):
124
# update config with runtime args
125
for key, value in runtime.args.items():
126
if hasattr(config, key):
@@ -130,7 +141,7 @@ def args_override(config):
141
setattr(config, key, value)
142
143
133
-def set_runtime_config(config: AgentConfig, set: settings.Settings):
144
+def _set_runtime_config(config: AgentConfig, set: settings.Settings):
145
ssh_conf = settings.get_runtime_config(set)
146
for key, value in ssh_conf.items():
147
if hasattr(config, key):
python/extensions/system_prompt/_10_system_prompt.py
+13
-2
@@ -16,7 +16,8 @@ class SystemPrompt(Extension):
16
17
system_prompt.append(main)
18
system_prompt.append(tools)
19
- system_prompt.append(mcp_tools)
19
+ if mcp_tools:
20
+ system_prompt.append(mcp_tools)
21
22
23
def get_main_prompt(agent: Agent):
@@ -31,4 +32,14 @@ def get_tools_prompt(agent: Agent):
32
33
34
def get_mcp_tools_prompt(agent: Agent):
34
- return MCPConfig.get_instance().get_tools_prompt()
35
+ mcp_config = MCPConfig.get_instance()
36
+ if mcp_config.servers:
37
+ pre_progress = agent.context.log.progress
38
+ agent.context.log.set_progress("Collecting MCP tools") # MCP might be initializing, better inform via progress bar
39
+ # import time
40
+ # time.sleep(10)
41
+ tools = MCPConfig.get_instance().get_tools_prompt()
42
+ agent.context.log.set_progress(pre_progress) # return original progress
43
+ return tools
44
+ return ""
45
+
python/helpers/api.py
+3
-3
@@ -5,7 +5,7 @@ from typing import Union, TypedDict, Dict, Any
5
from attr import dataclass
6
from flask import Request, Response, jsonify, Flask
7
from agent import AgentContext
8
-from initialize import initialize
8
+from initialize import initialize_agent
9
from python.helpers.print_style import PrintStyle
10
from python.helpers.errors import format_error
11
from werkzeug.serving import make_server
@@ -77,8 +77,8 @@ class ApiHandler:
77
first = AgentContext.first()
78
if first:
79
return first
80
- return AgentContext(config=initialize())
80
+ return AgentContext(config=initialize_agent())
81
got = AgentContext.get(ctxid)
82
if got:
83
return got
84
- return AgentContext(config=initialize(), id=ctxid)
84
+ return AgentContext(config=initialize_agent(), id=ctxid)
python/helpers/cloudflare_tunnel._py
renamed
python/helpers/defer.py
+4
-3
@@ -10,12 +10,12 @@ class EventLoopThread:
10
_instances = {}
11
_lock = threading.Lock()
12
13
- def __init__(self, thread_name: str = "default") -> None:
13
+ def __init__(self, thread_name: str = "Background") -> None:
14
"""Initialize the event loop thread."""
15
self.thread_name = thread_name
16
self._start()
17
18
- def __new__(cls, thread_name: str = "default"):
18
+ def __new__(cls, thread_name: str = "Background"):
19
with cls._lock:
20
if thread_name not in cls._instances:
21
instance = super(EventLoopThread, cls).__new__(cls)
@@ -59,7 +59,7 @@ class ChildTask:
59
class DeferredTask:
60
def __init__(
61
self,
62
- thread_name: str = "default",
62
+ thread_name: str = "Background",
63
):
64
self.event_loop_thread = EventLoopThread(thread_name)
65
self._future: Optional[Future] = None
@@ -72,6 +72,7 @@ class DeferredTask:
72
self.args = args
73
self.kwargs = kwargs
74
self._start_task()
75
+ return self
76
77
def __del__(self):
78
self.kill()
python/helpers/mcp_handler.py
+18
-16
@@ -62,20 +62,16 @@ def initialize_mcp(mcp_servers_config: str):
62
except Exception as e:
63
from agent import AgentContext
64
65
- first_context = AgentContext.first() # TODO replace with better reporting
66
- if first_context:
67
- (
68
- first_context.log.log(
69
- type="warning",
70
- content=f"Failed to update MCP settings: {e}",
71
- temp=False,
72
- )
73
- )
74
- (
75
- PrintStyle(
76
- background_color="black", font_color="red", padding=True
77
- ).print(f"Failed to update MCP settings: {e}")
65
+ AgentContext.log_to_all(
66
+ type="warning",
67
+ content=f"Failed to update MCP settings: {e}",
68
+ temp=False,
69
)
70
+
71
+ PrintStyle(
72
+ background_color="black", font_color="red", padding=True
73
+ ).print(f"Failed to update MCP settings: {e}")
74
+
75
76
77
class MCPTool(Tool):
@@ -373,9 +369,10 @@ class MCPConfig(BaseModel):
369
370
@classmethod
371
def get_instance(cls) -> "MCPConfig":
376
- if cls.__instance is None:
377
- cls.__instance = cls(servers_list=[])
378
- return cls.__instance
372
+ # with cls.__lock:
373
+ if cls.__instance is None:
374
+ cls.__instance = cls(servers_list=[])
375
+ return cls.__instance
376
377
@classmethod
378
def wait_for_lock(cls):
@@ -690,6 +687,11 @@ class MCPConfig(BaseModel):
687
688
def get_tools_prompt(self, server_name: str = "") -> str:
689
"""Get a prompt for all tools"""
690
+
691
+ # just to wait for pending initialization
692
+ with self.__lock:
693
+ pass
694
+
695
prompt = '## "Remote (MCP Server) Agent Tools" available:\n\n'
696
server_names = []
697
for server in self.servers:
python/helpers/mcp_server.py
+2
-2
@@ -8,7 +8,7 @@ from fastmcp import FastMCP
8
9
from agent import AgentContext, AgentContextType, UserMessage
10
from python.helpers.persist_chat import save_tmp_chat, remove_chat
11
-from initialize import initialize
11
+from initialize import initialize_agent
12
from python.helpers.print_style import PrintStyle
13
from python.helpers import settings
14
from starlette.middleware import Middleware
@@ -139,7 +139,7 @@ async def send_message(
139
# If we continue a conversation, it must be persistent
140
persistent_chat = True
141
else:
142
- config = initialize()
142
+ config = initialize_agent()
143
context = AgentContext(config=config, type=AgentContextType.MCP)
144
145
if not message:
python/helpers/persist_chat.py
+2
-2
@@ -5,7 +5,7 @@ import uuid
5
from agent import Agent, AgentConfig, AgentContext, AgentContextType
6
from python.helpers import files, history
7
import json
8
-from initialize import initialize
8
+from initialize import initialize_agent
9
10
from python.helpers.log import Log, LogItem
11
@@ -146,7 +146,7 @@ def _serialize_log(log: Log):
146
147
148
def _deserialize_context(data):
149
- config = initialize()
149
+ config = initialize_agent()
150
log = _deserialize_log(data.get("log", None))
151
152
context = AgentContext(
python/helpers/settings.py
+2
-2
@@ -965,9 +965,9 @@ def _apply_settings(previous: Settings | None):
965
global _settings
966
if _settings:
967
from agent import AgentContext
968
- from initialize import initialize
968
+ from initialize import initialize_agent
969
970
- config = initialize()
970
+ config = initialize_agent()
971
for ctx in AgentContext._contexts.values():
972
ctx.config = config # reinitialize context config with new settings
973
# apply config to agents
python/helpers/task_scheduler.py
+2
-2
@@ -16,7 +16,7 @@ from crontab import CronTab
16
from pydantic import BaseModel, Field, PrivateAttr
17
18
from agent import Agent, AgentContext, UserMessage
19
-from initialize import initialize
19
+from initialize import initialize_agent
20
from python.helpers.persist_chat import save_tmp_chat
21
from python.helpers.print_style import PrintStyle
22
from python.helpers.defer import DeferredTask
@@ -713,7 +713,7 @@ class TaskScheduler:
713
if not task.context_id:
714
raise ValueError(f"Task {task.name} has no context ID")
715
716
- config = initialize()
716
+ config = initialize_agent()
717
context: AgentContext = AgentContext(config, id=task.context_id, name=task.name)
718
# context.id = task.context_id
719
# initial name before renaming is same as task name
run_cli.py
+2
-2
@@ -8,7 +8,7 @@ from python.helpers.print_style import PrintStyle
8
from python.helpers.files import read_file
9
from python.helpers import files
10
import python.helpers.timed_input as timed_input
11
-from initialize import initialize
11
+from initialize import initialize_agent
12
from python.helpers.dotenv import load_dotenv
13
14
@@ -102,7 +102,7 @@ def run():
102
load_dotenv()
103
104
# initialize context
105
- config = initialize()
105
+ config = initialize_agent()
106
context = AgentContext(config)
107
108
# Start the key capture thread for user intervention during agent streaming
run_ui.py
+25
-77
@@ -8,15 +8,13 @@ import threading
8
import signal
9
from flask import Flask, request, Response
10
from flask_basicauth import BasicAuth
11
+import initialize
12
from python.helpers import errors, files, git, mcp_server
13
from python.helpers.files import get_abs_path
13
-from python.helpers import persist_chat, runtime, dotenv, process
14
-from python.helpers.cloudflare_tunnel import CloudflareTunnel
14
+from python.helpers import runtime, dotenv, process
15
from python.helpers.extract_tools import load_classes_from_folder
16
from python.helpers.api import ApiHandler
17
-from python.helpers.job_loop import run_loop
17
from python.helpers.print_style import PrintStyle
19
-from python.helpers.defer import DeferredTask
18
19
20
# Set the new timezone to 'UTC'
@@ -148,11 +146,6 @@ def run():
146
from werkzeug.serving import make_server
147
from werkzeug.middleware.dispatcher import DispatcherMiddleware
148
from a2wsgi import ASGIMiddleware, WSGIMiddleware
151
- from fastmcp.server.http import create_sse_app
152
- from python.helpers.mcp_server import mcp_server as mcp_server_instance
153
-
154
- PrintStyle().print("Starting job loop...")
155
- job_loop = DeferredTask().start_task(run_loop)
149
150
PrintStyle().print("Starting server...")
151
@@ -165,33 +158,6 @@ def run():
158
host = (
159
runtime.get_arg("host") or dotenv.get_dotenv_value("WEB_UI_HOST") or "localhost"
160
)
168
- use_cloudflare = (
169
- runtime.get_arg("cloudflare_tunnel")
170
- or dotenv.get_dotenv_value("USE_CLOUDFLARE", "false").lower()
171
- ) == "true"
172
-
173
- tunnel = None
174
-
175
- try:
176
- # Initialize and start Cloudflare tunnel if enabled
177
- if use_cloudflare and port:
178
- try:
179
- tunnel = CloudflareTunnel(port)
180
- tunnel.start()
181
- except Exception as e:
182
- PrintStyle().error(f"Failed to start Cloudflare tunnel: {e}")
183
- PrintStyle().print("Continuing without tunnel...")
184
-
185
- # # initialize contexts from persisted chats - moved to async task
186
- # persist_chat.load_tmp_chats()
187
-
188
- # # reload scheduler
189
- # scheduler = TaskScheduler.get()
190
- # asyncio.run(scheduler.reload())
191
-
192
- except Exception as e:
193
- PrintStyle().error(errors.format_error(e))
194
-
161
server = None
162
163
def register_api_handler(app, handler: type[ApiHandler]):
@@ -243,53 +209,35 @@ def run():
209
)
210
PrintStyle().debug("Registered middleware for MCP and MCP token")
211
246
- try:
247
- PrintStyle().debug(f"Starting server at {host}:{port}...")
248
-
249
- server = make_server(
250
- host=host,
251
- port=port,
252
- app=app,
253
- request_handler=NoRequestLoggingWSGIRequestHandler,
254
- threaded=True,
255
- )
256
-
257
- printer = PrintStyle()
212
+ PrintStyle().debug(f"Starting server at {host}:{port}...")
213
259
- def signal_handler(sig=None, frame=None):
260
- nonlocal tunnel, server, printer
261
- with lock:
262
- printer.print("Caught signal, stopping server...")
263
- if server:
264
- server.shutdown()
265
- process.stop_server()
266
- if tunnel:
267
- tunnel.stop()
268
- tunnel = None
269
- printer.print("Server stopped")
270
- sys.exit(0)
271
-
272
- signal.signal(signal.SIGINT, signal_handler)
273
- signal.signal(signal.SIGTERM, signal_handler)
274
-
275
- process.set_server(server)
276
- server.log_startup()
277
-
278
- # Start init_a0 in a background thread when server starts
279
- import threading
214
+ server = make_server(
215
+ host=host,
216
+ port=port,
217
+ app=app,
218
+ request_handler=NoRequestLoggingWSGIRequestHandler,
219
+ threaded=True,
220
+ )
221
+ process.set_server(server)
222
+ server.log_startup()
223
281
- threading.Thread(target=init_a0, daemon=True).start()
224
+ # Start init_a0 in a background thread when server starts
225
+ # threading.Thread(target=init_a0, daemon=True).start()
226
+ init_a0()
227
283
- server.serve_forever()
284
- finally:
285
- # Clean up tunnel if it was started
286
- if tunnel:
287
- tunnel.stop()
228
+ # run the server
229
+ server.serve_forever()
230
231
232
def init_a0():
291
- # initialize contexts from persisted chats
292
- persist_chat.load_tmp_chats()
233
+ # initialize contexts and MCP
234
+ init_chats = initialize.initialize_chats()
235
+ initialize.initialize_mcp()
236
+ # start job loop
237
+ initialize.initialize_job_loop()
238
+
239
+ # only wait for init chats, otherwise they would seem to dissapear for a while on restart
240
+ init_chats.result_sync()
241
242
243
# run the internal server