feat: embedded MCP server for messaging with agent-zero
Rafael Uzarowski committed
May 19, 2025 at 13:27 UTC
6beb4ec939a28c3bdc1e2dde46210409322d311a
6 files changed
+303
-16
agent.py
+20
-2
@@ -1,8 +1,9 @@
1
import asyncio
2
from collections import OrderedDict
3
from dataclasses import dataclass, field
4
-from datetime import datetime
4
+from datetime import datetime, timezone
5
from typing import Any, Awaitable, Coroutine, Optional, Dict, TypedDict
6
+from enum import Enum
7
import uuid
8
import models
9
@@ -20,6 +21,12 @@ from typing import Callable
21
from python.helpers.localization import Localization
22
23
24
+class AgentContextType(Enum):
25
+ USER = "user"
26
+ TASK = "task"
27
+ MCP = "mcp"
28
+
29
+
30
class AgentContext:
31
32
_contexts: dict[str, "AgentContext"] = {}
@@ -35,6 +42,8 @@ class AgentContext:
42
paused: bool = False,
43
streaming_agent: "Agent|None" = None,
44
created_at: datetime | None = None,
45
+ type: AgentContextType = AgentContextType.USER,
46
+ last_message: datetime | None = None,
47
):
48
# build context
49
self.id = id or str(uuid.uuid4())
@@ -45,9 +54,12 @@ class AgentContext:
54
self.paused = paused
55
self.streaming_agent = streaming_agent
56
self.task: DeferredTask | None = None
48
- self.created_at = created_at or datetime.now()
57
+ self.created_at = created_at or datetime.now(timezone.utc)
58
+ self.type = type
59
AgentContext._counter += 1
60
self.no = AgentContext._counter
61
+ # set to start of unix epoch
62
+ self.last_message = last_message or datetime.now(timezone.utc)
63
64
existing = self._contexts.get(self.id, None)
65
if existing:
@@ -84,6 +96,11 @@ class AgentContext:
96
"log_version": len(self.log.updates),
97
"log_length": len(self.log.logs),
98
"paused": self.paused,
99
+ "last_message": (
100
+ Localization.get().serialize_datetime(self.last_message)
101
+ if self.last_message else Localization.get().serialize_datetime(datetime.fromtimestamp(0))
102
+ ),
103
+ "type": self.type.value,
104
}
105
106
def get_created_at(self):
@@ -468,6 +485,7 @@ class Agent:
485
def hist_add_message(
486
self, ai: bool, content: history.MessageContent, tokens: int = 0
487
):
488
+ self.last_message = datetime.now(timezone.utc)
489
return self.history.add_message(ai=ai, content=content, tokens=tokens)
490
491
def hist_add_user_message(self, message: UserMessage, intervention: bool = False):
python/helpers/job_loop.py
+24
-2
@@ -2,19 +2,41 @@ import asyncio
2
from python.helpers.task_scheduler import TaskScheduler
3
from python.helpers.print_style import PrintStyle
4
from python.helpers import errors
5
+from agent import AgentContext, AgentContextType
6
+from datetime import datetime, timezone, timedelta
7
+from python.helpers.persist_chat import remove_chat
8
9
10
async def run_loop():
11
while True:
12
+ # scheduler tick
13
try:
14
await scheduler_tick()
15
except Exception as e:
16
PrintStyle().error(errors.format_error(e))
13
- await asyncio.sleep(60) # TODO! - if we lower it under 1min, it can run a 5min job multiple times in it's target minute
17
+
18
+ # cleanup tmp chats
19
+ try:
20
+ await cleanup_tmp_chats()
21
+ except Exception as e:
22
+ PrintStyle().error(errors.format_error(e))
23
+
24
+ await asyncio.sleep(60) # TODO! - if we lower it under 1min, it can run a 5min job multiple times in it's target minute
25
26
27
async def scheduler_tick():
28
# Get the task scheduler instance and print detailed debug info
29
scheduler = TaskScheduler.get()
30
# Run the scheduler tick
20
- await scheduler.tick()
\ No newline at end of file
31
+ await scheduler.tick()
32
+
33
+
34
+async def cleanup_tmp_chats():
35
+ contexts = list(AgentContext._contexts.values())
36
+ for context in contexts:
37
+ if context.type == AgentContextType.MCP:
38
+ if context.last_message < datetime.now(timezone.utc) - timedelta(hours=1):
39
+ PrintStyle().debug(f"MCP chat {context.id} - {context.last_message} - cleaning up")
40
+ context.reset()
41
+ AgentContext.remove(context.id)
42
+ remove_chat(context.id)
python/helpers/mcp_server.py
new
+213
@@ -0,0 +1,213 @@
1
+from asyncio import current_task
2
+import os
3
+from typing import Annotated, Literal, Union
4
+from urllib.parse import urlparse
5
+from openai import BaseModel
6
+from pydantic import Field
7
+import uuid
8
+import asyncio
9
+from fastmcp import FastMCP
10
+
11
+from agent import AgentContext, AgentContextType, UserMessage
12
+from python.helpers.persist_chat import save_tmp_chat, remove_chat
13
+from initialize import initialize
14
+from python.helpers.print_style import PrintStyle
15
+from python.helpers.task_scheduler import DeferredTask
16
+
17
+_PRINTER = PrintStyle(italic=True, font_color="green", padding=False)
18
+
19
+
20
+mcp_server: FastMCP = FastMCP(
21
+ name="Agent Zero integrated MCP Server",
22
+ instructions="""
23
+ This server connects you to the Agent Zero instance running on the remote server.
24
+ It exposes tools to interact with the remote Agent Zero instance.
25
+ """,
26
+)
27
+
28
+
29
+class ToolResponse(BaseModel):
30
+ status: Literal["success"] = Field(description="The status of the response", default="success")
31
+ response: str = Field(description="The response from the remote Agent Zero Instance")
32
+ chat_id: str = Field(description="The id of the chat this message belongs to.")
33
+
34
+
35
+class ToolError(BaseModel):
36
+ status: Literal["error"] = Field(description="The status of the response", default="error")
37
+ error: str = Field(description="The error message from the remote Agent Zero Instance")
38
+ chat_id: str = Field(description="The id of the chat this message belongs to.")
39
+
40
+
41
+SEND_MESSAGE_DESCRIPTION = """
42
+Send a message to the remote Agent Zero Instance.
43
+This tool is used to send a message to the remote Agent Zero Instance connected remotely via MCP.
44
+"""
45
+
46
+
47
+@mcp_server.tool(
48
+ name="send_message",
49
+ description=SEND_MESSAGE_DESCRIPTION,
50
+ tags={"agent_zero", "chat", "remote", "communication", "dialogue", "sse", "send", "message", "start", "new", "continue"},
51
+ annotations={
52
+ "remote": True,
53
+ "readOnlyHint": False,
54
+ "destructiveHint": False,
55
+ "idempotentHint": False,
56
+ "openWorldHint": False,
57
+ "title": SEND_MESSAGE_DESCRIPTION,
58
+ },
59
+)
60
+async def send_message(
61
+ message: Annotated[str, Field(description="The message to send to the remote Agent Zero Instance", title="message")],
62
+ attachments: Annotated[list[str], Field(
63
+ description="Optional: A list of attachments (file paths or web urls) to send to the remote Agent Zero Instance with the message. Default: Empty list",
64
+ title="attachments",
65
+ )] | None = None,
66
+ chat_id: Annotated[str, Field(
67
+ description="Optional: ID of the chat. Used to continue a chat. This value is returned in response to sending previous message. Default: Empty string",
68
+ title="chat_id",
69
+ )] | None = None,
70
+ persistent_chat: Annotated[bool, Field(
71
+ description="Optional: Whether to use a persistent chat. If true, the chat will be saved and can be continued later. Default: False.",
72
+ title="persistent_chat",
73
+ )] | None = None,
74
+) -> Annotated[Union[ToolResponse, ToolError], Field(description="The response from the remote Agent Zero Instance", title="response")]:
75
+ context: AgentContext | None = None
76
+ if chat_id:
77
+ context = AgentContext.get(chat_id)
78
+ if not context:
79
+ return ToolError(error="Chat not found", chat_id=chat_id)
80
+ else:
81
+ # If the chat is found, we use the persistent chat flag to determine
82
+ # whether we should save the chat or delete it afterwards
83
+ # If we continue a conversation, it must be persistent
84
+ persistent_chat = True
85
+ else:
86
+ config = initialize()
87
+ context = AgentContext(config=config, type=AgentContextType.MCP)
88
+
89
+ if not message:
90
+ return ToolError(error="Message is required", chat_id=context.id if persistent_chat else "")
91
+
92
+ try:
93
+ response = await _run_chat(context, message, attachments)
94
+ if not persistent_chat:
95
+ context.reset()
96
+ AgentContext.remove(context.id)
97
+ remove_chat(context.id)
98
+ return ToolResponse(response=response, chat_id=context.id if persistent_chat else "")
99
+ except Exception as e:
100
+ return ToolError(error=str(e), chat_id=context.id if persistent_chat else "")
101
+
102
+
103
+FINISH_CHAT_DESCRIPTION = """
104
+Finish a chat with the remote Agent Zero Instance.
105
+This tool is used to finish a persistent chat (send_message with persistent_chat=True) with the remote Agent Zero Instance connected remotely via MCP.
106
+If you want to continue the chat, use the send_message tool instead.
107
+Always use this tool to finish persistent chat conversations with remote Agent Zero.
108
+"""
109
+
110
+
111
+@mcp_server.tool(
112
+ name="finish_chat",
113
+ description=FINISH_CHAT_DESCRIPTION,
114
+ tags={"agent_zero", "chat", "remote", "communication", "dialogue", "sse", "finish", "close", "end", "stop"},
115
+ annotations={
116
+ "remote": True,
117
+ "readOnlyHint": False,
118
+ "destructiveHint": True,
119
+ "idempotentHint": False,
120
+ "openWorldHint": False,
121
+ "title": FINISH_CHAT_DESCRIPTION,
122
+ },
123
+)
124
+async def finish_chat(
125
+ chat_id: Annotated[str, Field(
126
+ description="ID of the chat to be finished. This value is returned in response to sending previous message.",
127
+ title="chat_id",
128
+ )]
129
+) -> Annotated[Union[ToolResponse, ToolError], Field(description="The response from the remote Agent Zero Instance", title="response")]:
130
+ if not chat_id:
131
+ return ToolError(error="Chat ID is required", chat_id="")
132
+
133
+ context = AgentContext.get(chat_id)
134
+ if not context:
135
+ return ToolError(error="Chat not found", chat_id=chat_id)
136
+ else:
137
+ context.reset()
138
+ AgentContext.remove(context.id)
139
+ remove_chat(context.id)
140
+ return ToolResponse(response="Chat finished", chat_id=chat_id)
141
+
142
+
143
+async def _run_chat(context: AgentContext, message: str, attachments: list[str] | None = None):
144
+ async def _run_chat_wrapper(context: AgentContext, message: str, attachments: list[str] | None = None):
145
+ # the agent instance - init in try block
146
+ agent = None
147
+
148
+ try:
149
+ _PRINTER.print("MCP Chat message received")
150
+
151
+ agent = context.streaming_agent or context.agent0
152
+
153
+ # Pcurrent_taskhment filenames for logging
154
+ attachment_filenames = []
155
+ if attachments:
156
+ for attachment in attachments:
157
+ if os.path.exists(attachment):
158
+ attachment_filenames.append(attachment)
159
+ else:
160
+ try:
161
+ url = urlparse(attachment)
162
+ if url.scheme in ["http", "https", "ftp", "ftps", "sftp"]:
163
+ attachment_filenames.append(attachment)
164
+ else:
165
+ _PRINTER.print(f"Skipping attachment: [{attachment}]")
166
+ except Exception:
167
+ _PRINTER.print(f"Skipping attachment: [{attachment}]")
168
+
169
+ _PRINTER.print("User message:")
170
+ _PRINTER.print(f"> {message}")
171
+ if attachment_filenames:
172
+ _PRINTER.print("Attachments:")
173
+ for filename in attachment_filenames:
174
+ _PRINTER.print(f"- {filename}")
175
+
176
+ # Log the message with message_id and attachments
177
+ context.log.log(
178
+ type="user",
179
+ heading="User message",
180
+ content=message,
181
+ kvps={"attachments": attachment_filenames},
182
+ id=str(uuid.uuid4()),
183
+ )
184
+
185
+ agent.hist_add_user_message(
186
+ UserMessage(
187
+ message=message,
188
+ system_message=[],
189
+ attachments=attachment_filenames))
190
+
191
+ # Persist after setting up the context but before running the agent
192
+ save_tmp_chat(context)
193
+
194
+ result = await agent.monologue()
195
+
196
+ # Success
197
+ _PRINTER.print(f"MCP Chat message completed: {result}")
198
+ save_tmp_chat(context)
199
+
200
+ return result
201
+
202
+ except Exception as e:
203
+ # Error
204
+ _PRINTER.print(f"MCP Chat message failed: {e}")
205
+ if agent:
206
+ agent.handle_critical_exception(e)
207
+
208
+ raise RuntimeError(f"MCP Chat message failed: {e}") from e
209
+
210
+ deferred_task = DeferredTask(thread_name="mcp_chat_" + context.id)
211
+ deferred_task.start_task(_run_chat_wrapper, context, message, attachments)
212
+ asyncio.create_task(asyncio.sleep(0.1)) # Ensure background execution doesn't exit immediately on async await
213
+ return await deferred_task.result()
python/helpers/persist_chat.py
+12
-1
@@ -2,7 +2,7 @@ from collections import OrderedDict
2
from datetime import datetime
3
from typing import Any
4
import uuid
5
-from agent import Agent, AgentConfig, AgentContext
5
+from agent import Agent, AgentConfig, AgentContext, AgentContextType
6
from python.helpers import files, history
7
import json
8
from initialize import initialize
@@ -109,6 +109,11 @@ def _serialize_context(context: AgentContext):
109
context.created_at.isoformat() if context.created_at
110
else datetime.fromtimestamp(0).isoformat()
111
),
112
+ "type": context.type.value,
113
+ "last_message": (
114
+ context.last_message.isoformat() if context.last_message
115
+ else datetime.fromtimestamp(0).isoformat()
116
+ ),
117
"agents": agents,
118
"streaming_agent": (
119
context.streaming_agent.number if context.streaming_agent else 0
@@ -154,6 +159,12 @@ def _deserialize_context(data):
159
data.get("created_at", datetime.fromtimestamp(0).isoformat())
160
)
161
),
162
+ type=AgentContextType(data.get("type", AgentContextType.USER.value)),
163
+ last_message=(
164
+ datetime.fromisoformat(
165
+ data.get("last_message", datetime.fromtimestamp(0).isoformat())
166
+ )
167
+ ),
168
log=log,
169
paused=False,
170
# agent0=agent0,
requirements.txt
+4
-1
@@ -24,7 +24,7 @@ newspaper3k==0.2.8
24
paramiko==3.5.0
25
playwright==1.52.0
26
pypdf==4.3.1
27
-python-dotenv==1.0.1
27
+python-dotenv==1.1.0
28
pytz==2024.2
29
sentence-transformers==3.0.1
30
tiktoken==0.8.0
@@ -32,3 +32,6 @@ unstructured==0.15.13
32
unstructured-client==0.25.9
33
webcolors==24.6.0
34
crontab==1.0.1
35
+fastmcp==2.3.4
36
+mcp==1.9.0
37
+a2wsgi==1.10.8
run_ui.py
+30
-10
@@ -20,19 +20,20 @@ from python.helpers.print_style import PrintStyle
20
from python.helpers.task_scheduler import TaskScheduler
21
from python.helpers.defer import DeferredTask
22
23
+
24
# Set the new timezone to 'UTC'
25
os.environ["TZ"] = "UTC"
26
# Apply the timezone change
27
time.tzset()
28
29
# initialize the internal Flask server
29
-app = Flask("app", static_folder=get_abs_path("./webui"), static_url_path="/")
30
-app.config["JSON_SORT_KEYS"] = False # Disable key sorting in jsonify
30
+webapp = Flask("app", static_folder=get_abs_path("./webui"), static_url_path="/")
31
+webapp.config["JSON_SORT_KEYS"] = False # Disable key sorting in jsonify
32
33
lock = threading.Lock()
34
34
-# Set up basic authentication
35
-basic_auth = BasicAuth(app)
35
+# Set up basic authentication for UI and API but not MCP
36
+basic_auth = BasicAuth(webapp)
37
38
39
def is_loopback_address(address):
@@ -123,7 +124,7 @@ def requires_auth(f):
124
125
126
# handle default address, load index
126
-@app.route("/", methods=["GET"])
127
+@webapp.route("/", methods=["GET"])
128
@requires_auth
129
async def serve_index():
130
gitinfo = None
@@ -147,6 +148,10 @@ def run():
148
# Suppress only request logs but keep the startup messages
149
from werkzeug.serving import WSGIRequestHandler
150
from werkzeug.serving import make_server
151
+ from werkzeug.middleware.dispatcher import DispatcherMiddleware
152
+ from a2wsgi import ASGIMiddleware, WSGIMiddleware
153
+ from fastmcp.server.http import create_sse_app
154
+ from python.helpers.mcp_server import mcp_server as mcp_server_instance
155
156
PrintStyle().print("Starting job loop...")
157
job_loop = DeferredTask().start_task(run_loop)
@@ -227,9 +232,28 @@ def run():
232
# initialize and register API handlers
233
handlers = load_classes_from_folder("python/api", "*.py", ApiHandler)
234
for handler in handlers:
230
- register_api_handler(app, handler)
235
+ register_api_handler(webapp, handler)
236
+
237
+ mcp_app = create_sse_app(
238
+ server=mcp_server_instance,
239
+ message_path=mcp_server_instance.settings.message_path,
240
+ sse_path=mcp_server_instance.settings.sse_path,
241
+ auth_server_provider=mcp_server_instance._auth_server_provider,
242
+ auth_settings=mcp_server_instance.settings.auth,
243
+ debug=mcp_server_instance.settings.debug,
244
+ routes=mcp_server_instance._additional_http_routes,
245
+ middleware=None
246
+ )
247
+
248
+ # add the webapp and mcp to the app
249
+ app = DispatcherMiddleware(webapp, {
250
+ "/mcp": ASGIMiddleware(app=mcp_app),
251
+ })
252
+ PrintStyle().debug("Registered middleware for MCP")
253
254
try:
255
+ PrintStyle().debug(f"Starting server at {host}:{port}...")
256
+
257
server = make_server(
258
host=host,
259
port=port,
@@ -259,10 +283,6 @@ def run():
283
process.set_server(server)
284
server.log_startup()
285
server.serve_forever()
262
- # Run Flask app
263
- # app.run(
264
- # request_handler=NoRequestLoggingWSGIRequestHandler, port=port, host=host
265
- # )
286
finally:
287
# Clean up tunnel if it was started
288
if tunnel: