main
py 604 lines 22.9 KB
Raw
1 # noqa: D401 (docstrings) – internal helper
2 import asyncio
3 import uuid
4 import atexit
5 import json
6 from typing import Any, List
7 import contextlib
8 import threading
9
10 from helpers import settings, projects
11 from starlette.responses import Response as StarletteResponse
12 from starlette.requests import Request
13
14 # Local imports
15 from helpers.print_style import PrintStyle
16 from agent import AgentContext, UserMessage, AgentContextType
17 from initialize import initialize_agent
18 from helpers.persist_chat import remove_chat
19
20 # Import FastA2A
21 try:
22 from fasta2a import Worker, FastA2A # type: ignore
23 from fasta2a.broker import InMemoryBroker # type: ignore
24 from fasta2a.storage import InMemoryStorage # type: ignore
25 from fasta2a.schema import Message, Artifact, AgentProvider, Skill # type: ignore
26 FASTA2A_AVAILABLE = True
27 except ImportError: # pragma: no cover – library not installed
28 FASTA2A_AVAILABLE = False
29 # Minimal stubs for type checkers when FastA2A is not available
30
31 class Worker: # type: ignore
32 def __init__(self, **kwargs):
33 pass
34
35 async def run_task(self, params):
36 pass
37
38 async def cancel_task(self, params):
39 pass
40
41 def build_message_history(self, history):
42 return []
43
44 def build_artifacts(self, result):
45 return []
46
47 class FastA2A: # type: ignore
48 def __init__(self, **kwargs):
49 pass
50
51 async def __call__(self, scope, receive, send):
52 pass
53
54 class InMemoryBroker: # type: ignore
55 pass
56
57 class InMemoryStorage: # type: ignore
58 async def update_task(self, **kwargs):
59 pass
60
61 Message = Artifact = AgentProvider = Skill = Any # type: ignore
62
63 _PRINTER = PrintStyle(italic=True, font_color="purple", padding=False)
64
65
66 def _enable_streaming_capability(agent_card_body: bytes) -> bytes:
67 """Return an agent-card JSON body with A2A streaming enabled."""
68 agent_card = json.loads(agent_card_body)
69 capabilities = agent_card.get("capabilities")
70 if not isinstance(capabilities, dict):
71 capabilities = {}
72 agent_card["capabilities"] = capabilities
73 capabilities["streaming"] = True
74 return json.dumps(agent_card, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
75
76
77 class AgentZeroFastA2A(FastA2A): # type: ignore[misc]
78 """FastA2A app with Agent Zero defaults layered over library defaults."""
79
80 async def _agent_card_endpoint(self, request: Request) -> StarletteResponse:
81 response = await super()._agent_card_endpoint(request)
82 body = _enable_streaming_capability(response.body)
83 self._agent_card_json_schema = body
84 return StarletteResponse(content=body, media_type="application/json")
85
86
87 class AgentZeroWorker(Worker): # type: ignore[misc]
88 """Agent Zero implementation of FastA2A Worker."""
89
90 def __init__(self, broker, storage):
91 super().__init__(broker=broker, storage=storage)
92 self.storage = storage
93
94 async def run_task(self, params: Any) -> None: # params: TaskSendParams
95 """Execute a task by processing the message through Agent Zero."""
96 context = None
97 try:
98 task_id = params['id']
99 message = params['message']
100
101 _PRINTER.print(f"[A2A] Processing task {task_id} with new temporary context")
102
103 # Convert A2A message to Agent Zero format
104 agent_message = self._convert_message(message)
105
106 # Always create new temporary context for this A2A conversation
107 cfg = initialize_agent()
108 context = AgentContext(cfg, type=AgentContextType.BACKGROUND)
109
110 # Retrieve project from message.metadata (standard A2A pattern)
111 metadata = message.get('metadata', {}) or {}
112 project_name = metadata.get('project')
113
114 # Activate project if specified
115 if project_name:
116 projects.activate_project(context.id, project_name)
117
118 # Log user message so it appears instantly in UI chat window
119 context.log.log(
120 type="user", # type: ignore[arg-type]
121 heading="Remote user message",
122 content=agent_message.message,
123 kvps={"from": "A2A"},
124 )
125
126 # Process message through Agent Zero (includes response)
127 task = context.communicate(agent_message)
128 result_text = await task.result()
129
130 # Build A2A message from result
131 response_message: Message = { # type: ignore
132 'role': 'agent',
133 'parts': [{'kind': 'text', 'text': str(result_text)}],
134 'kind': 'message',
135 'message_id': str(uuid.uuid4())
136 }
137
138 await self.storage.update_task( # type: ignore[attr-defined]
139 task_id=task_id,
140 state='completed',
141 new_messages=[response_message]
142 )
143
144 # Clean up context like non-persistent MCP chats
145 context.reset()
146 AgentContext.remove(context.id)
147 remove_chat(context.id)
148
149 _PRINTER.print(f"[A2A] Completed task {task_id} and cleaned up context")
150
151 except Exception as e:
152 _PRINTER.print(f"[A2A] Error processing task {params.get('id', 'unknown')}: {e}")
153 await self.storage.update_task(
154 task_id=params.get('id', 'unknown'),
155 state='failed'
156 )
157
158 # Clean up context even on failure to prevent resource leaks
159 if context:
160 context.reset()
161 AgentContext.remove(context.id)
162 remove_chat(context.id)
163 _PRINTER.print(f"[A2A] Cleaned up failed context {context.id}")
164
165 async def cancel_task(self, params: Any) -> None: # params: TaskIdParams
166 """Cancel a running task."""
167 task_id = params['id']
168 _PRINTER.print(f"[A2A] Cancelling task {task_id}")
169 await self.storage.update_task(task_id=task_id, state='canceled') # type: ignore[attr-defined]
170
171 # Note: No context cleanup needed since contexts are always temporary and cleaned up in run_task
172
173 def build_message_history(self, history: List[Any]) -> List[Message]: # type: ignore
174 # Not used in this simplified implementation
175 return []
176
177 def build_artifacts(self, result: Any) -> List[Artifact]: # type: ignore
178 # No artifacts for now
179 return []
180
181 def _convert_message(self, a2a_message: Message) -> UserMessage: # type: ignore
182 """Convert A2A message to Agent Zero UserMessage."""
183 # Extract text from message parts
184 text_parts = [part.get('text', '') for part in a2a_message.get('parts', []) if part.get('kind') == 'text']
185 message_text = '\n'.join(text_parts)
186
187 # Extract file attachments
188 attachments = []
189 for part in a2a_message.get('parts', []):
190 if part.get('kind') == 'file':
191 file_info = part.get('file', {})
192 if 'uri' in file_info:
193 attachments.append(file_info['uri'])
194
195 return UserMessage(
196 message=message_text,
197 attachments=attachments
198 )
199
200
201 class DynamicA2AProxy:
202 """Dynamic proxy for FastA2A server that allows reconfiguration."""
203
204 _instance = None
205
206 def __init__(self):
207 self.app = None
208 self.token = ""
209 self._lock = threading.Lock() # Use threading.Lock instead of asyncio.Lock
210 self._startup_done: bool = False
211 self._worker_bg_task: asyncio.Task | None = None
212 self._reconfigure_needed: bool = False # Flag for deferred reconfiguration
213
214 if FASTA2A_AVAILABLE:
215 # Initialize with default token
216 cfg = settings.get_settings()
217 self.token = cfg.get("mcp_server_token", "")
218 self._configure()
219 self._register_shutdown()
220 else:
221 _PRINTER.print("[A2A] FastA2A not available, server will return 503")
222
223 @staticmethod
224 def get_instance():
225 if DynamicA2AProxy._instance is None:
226 DynamicA2AProxy._instance = DynamicA2AProxy()
227 return DynamicA2AProxy._instance
228
229 def reconfigure(self, token: str):
230 """Reconfigure the FastA2A server with new token."""
231 self.token = token
232 if FASTA2A_AVAILABLE:
233 with self._lock:
234 # Mark that reconfiguration is needed - will be done on next request
235 self._reconfigure_needed = True
236 self._startup_done = False # Force restart on next request
237 _PRINTER.print("[A2A] Reconfiguration scheduled for next request")
238
239 def _configure(self):
240 """Configure the FastA2A application with Agent Zero integration."""
241 try:
242 storage = InMemoryStorage() # type: ignore[arg-type]
243 broker = InMemoryBroker() # type: ignore[arg-type]
244
245 # Define Agent Zero's skills
246 skills: List[Skill] = [{ # type: ignore
247 "id": "general_assistance",
248 "name": "General AI Assistant",
249 "description": "Provides general AI assistance including code execution, file management, web browsing, and problem solving",
250 "tags": ["ai", "assistant", "code", "files", "web", "automation"],
251 "examples": [
252 "Write and execute Python code",
253 "Manage files and directories",
254 "Browse the web and extract information",
255 "Solve complex problems step by step",
256 "Install software and manage systems"
257 ],
258 "input_modes": ["text/plain", "application/octet-stream"],
259 "output_modes": ["text/plain", "application/json"]
260 }]
261
262 provider: AgentProvider = { # type: ignore
263 "organization": "Agent Zero",
264 "url": "https://github.com/frdel/agent-zero"
265 }
266
267 # Create new FastA2A app with proper thread safety
268 new_app = AgentZeroFastA2A( # type: ignore
269 storage=storage,
270 broker=broker,
271 name="Agent Zero",
272 description=(
273 "A general AI assistant that can execute code, manage files, browse the web, and "
274 "solve complex problems in an isolated Linux environment."
275 ),
276 version="1.0.0",
277 provider=provider,
278 skills=skills,
279 lifespan=None, # We manage lifespan manually
280 middleware=[], # No middleware - we handle auth in wrapper
281 )
282
283 # Store for later lazy startup (needs active event-loop)
284 self._storage = storage # type: ignore[attr-defined]
285 self._broker = broker # type: ignore[attr-defined]
286 self._worker = AgentZeroWorker(broker=broker, storage=storage) # type: ignore[attr-defined]
287
288 # Atomic update of the app
289 self.app = new_app
290
291 # _PRINTER.print("[A2A] FastA2A server configured successfully")
292
293 except Exception as e:
294 _PRINTER.print(f"[A2A] Failed to configure FastA2A server: {e}")
295 self.app = None
296 raise
297
298 # ---------------------------------------------------------------------
299 # Shutdown handling
300 # ---------------------------------------------------------------------
301
302 def _register_shutdown(self):
303 """Register an atexit hook to gracefully stop worker & task manager."""
304
305 def _sync_shutdown():
306 try:
307 if not self._startup_done or not FASTA2A_AVAILABLE:
308 return
309 loop = asyncio.new_event_loop()
310 loop.run_until_complete(self._async_shutdown())
311 loop.close()
312 except Exception:
313 pass # ignore errors during interpreter shutdown
314
315 atexit.register(_sync_shutdown)
316
317 async def _async_shutdown(self):
318 """Async shutdown: cancel worker task & close task manager."""
319 if self._worker_bg_task and not self._worker_bg_task.done():
320 self._worker_bg_task.cancel()
321 with contextlib.suppress(asyncio.CancelledError):
322 await self._worker_bg_task
323 try:
324 if hasattr(self, 'app') and self.app:
325 await self.app.task_manager.__aexit__(None, None, None) # type: ignore[attr-defined]
326 except Exception:
327 pass
328
329 async def _async_reconfigure(self):
330 """Perform async reconfiguration with proper lifecycle management."""
331 _PRINTER.print("[A2A] Starting async reconfiguration")
332
333 # Shutdown existing components
334 await self._async_shutdown()
335
336 # Reset startup state
337 self._startup_done = False
338 self._worker_bg_task = None
339
340 # Reconfigure with new token
341 self._configure()
342
343 # Restart components
344 await self._startup()
345
346 # Clear reconfiguration flag
347 self._reconfigure_needed = False
348
349 _PRINTER.print("[A2A] Async reconfiguration completed")
350
351 async def _startup(self):
352 """Ensure TaskManager and Worker are running inside current event-loop."""
353 if self._startup_done or not FASTA2A_AVAILABLE:
354 return
355 self._startup_done = True
356
357 # Start task manager
358 await self.app.task_manager.__aenter__() # type: ignore[attr-defined]
359
360 async def _worker_loop():
361 async with self._worker.run(): # type: ignore[attr-defined]
362 await asyncio.Event().wait()
363
364 # fire-and-forget background task – keep reference
365 self._worker_bg_task = asyncio.create_task(_worker_loop())
366 _PRINTER.print("[A2A] Worker & TaskManager started")
367
368 async def __call__(self, scope, receive, send):
369 """ASGI application interface with token-based routing."""
370 if not FASTA2A_AVAILABLE:
371 # FastA2A not available, return 503
372 response = b'HTTP/1.1 503 Service Unavailable\r\n\r\nFastA2A not available'
373 await send({
374 'type': 'http.response.start',
375 'status': 503,
376 'headers': [[b'content-type', b'text/plain']],
377 })
378 await send({
379 'type': 'http.response.body',
380 'body': response,
381 })
382 return
383
384 from helpers import settings
385 cfg = settings.get_settings()
386 if not cfg["a2a_server_enabled"]:
387 response = b'HTTP/1.1 403 Forbidden\r\n\r\nA2A server is disabled'
388 await send({
389 'type': 'http.response.start',
390 'status': 403,
391 'headers': [[b'content-type', b'text/plain']],
392 })
393 await send({
394 'type': 'http.response.body',
395 'body': response,
396 })
397 return
398
399 # Check if reconfiguration is needed
400 if self._reconfigure_needed:
401 try:
402 await self._async_reconfigure()
403 except Exception as e:
404 _PRINTER.print(f"[A2A] Error during reconfiguration: {e}")
405 # Return 503 if reconfiguration failed
406 await send({
407 'type': 'http.response.start',
408 'status': 503,
409 'headers': [[b'content-type', b'text/plain']],
410 })
411 await send({
412 'type': 'http.response.body',
413 'body': b'FastA2A reconfiguration failed',
414 })
415 return
416
417 if self.app is None:
418 # FastA2A not configured, return 503
419 response = b'HTTP/1.1 503 Service Unavailable\r\n\r\nFastA2A not configured'
420 await send({
421 'type': 'http.response.start',
422 'status': 503,
423 'headers': [[b'content-type', b'text/plain']],
424 })
425 await send({
426 'type': 'http.response.body',
427 'body': response,
428 })
429 return
430
431 # Lazy-start background components the first time we get a request
432 if not self._startup_done:
433 try:
434 _PRINTER.print("[A2A] Starting up FastA2A components")
435 await self._startup()
436 except Exception as e:
437 _PRINTER.print(f"[A2A] Error during startup: {e}")
438 # Return 503 if startup failed
439 await send({
440 'type': 'http.response.start',
441 'status': 503,
442 'headers': [[b'content-type', b'text/plain']],
443 })
444 await send({
445 'type': 'http.response.body',
446 'body': b'FastA2A startup failed',
447 })
448 return
449
450 # Handle token-based routing: /a2a/t-{token}/... or /t-{token}/...
451 path = scope.get('path', '')
452
453 # Strip /a2a prefix if present (DispatcherMiddleware doesn't always strip it)
454 if path.startswith('/a2a'):
455 path = path[4:] # Remove '/a2a' prefix
456
457 # Initialize project name
458 project_name = None
459
460 # Check if path matches token pattern /t-{token}/
461 if path.startswith('/t-'):
462 # Extract token from path
463 if '/' in path[3:]:
464 path_parts = path[3:].split('/', 1) # Remove '/t-' prefix
465 request_token = path_parts[0]
466 remaining_path = '/' + path_parts[1] if len(path_parts) > 1 else '/'
467
468 # Check for project pattern /p-{project}/
469 if remaining_path.startswith('/p-'):
470 project_parts = remaining_path[3:].split('/', 1)
471 if project_parts[0]:
472 project_name = project_parts[0]
473 remaining_path = '/' + project_parts[1] if len(project_parts) > 1 else '/'
474 _PRINTER.print(f"[A2A] Extracted project from URL: {project_name}")
475 else:
476 request_token = path[3:]
477 remaining_path = '/'
478
479 # Validate token
480 cfg = settings.get_settings()
481 expected_token = cfg.get("mcp_server_token")
482
483 if expected_token and request_token != expected_token:
484 # Invalid token, return 401
485 await send({
486 'type': 'http.response.start',
487 'status': 401,
488 'headers': [[b'content-type', b'text/plain']],
489 })
490 await send({
491 'type': 'http.response.body',
492 'body': b'Unauthorized',
493 })
494 return
495
496 # If project specified, inject it into the request payload
497 if project_name:
498 # Buffer messages and modify before returning the complete body
499 received_messages = []
500 body_modified = False
501 original_receive = receive
502
503 async def receive_wrapper():
504 nonlocal body_modified
505
506 # Receive and buffer the next message
507 message = await original_receive()
508 received_messages.append(message)
509
510 # When we get the complete body, inject project into JSON
511 if message['type'] == 'http.request' and not message.get('more_body', False) and not body_modified:
512 body_modified = True
513 try:
514 import json
515 # Reconstruct full body from all buffered messages
516 body_parts = [msg.get('body', b'') for msg in received_messages if msg['type'] == 'http.request']
517 full_body = b''.join(body_parts)
518 data = json.loads(full_body)
519
520 # INJECT project into message.metadata (standard A2A pattern)
521 if 'params' in data and 'message' in data['params']:
522 msg_data = data['params']['message']
523 # Initialize metadata if it doesn't exist
524 if 'metadata' not in msg_data or msg_data['metadata'] is None:
525 msg_data['metadata'] = {}
526 msg_data['metadata']['project'] = project_name
527
528 # Serialize back to JSON
529 modified_body = json.dumps(data).encode('utf-8')
530
531 # Return modified message IMMEDIATELY (before FastA2A processes it)
532 return {
533 'type': 'http.request',
534 'body': modified_body,
535 'more_body': False
536 }
537 except Exception as e:
538 _PRINTER.print(f"[A2A] Failed to inject project into payload: {e}")
539
540 return message
541
542 receive = receive_wrapper
543
544 # Update scope with cleaned path
545 scope = dict(scope)
546 scope['path'] = remaining_path
547 else:
548 # No token in path, check other auth methods
549 request = Request(scope, receive=receive)
550
551 cfg = settings.get_settings()
552 expected = cfg.get("mcp_server_token")
553
554 if expected:
555 auth_header = request.headers.get("Authorization", "")
556 api_key = request.headers.get("X-API-KEY") or request.query_params.get("api_key")
557
558 is_authorized = (
559 (auth_header.startswith("Bearer ") and auth_header.split(" ", 1)[1] == expected) or
560 (api_key == expected)
561 )
562
563 if not is_authorized:
564 # No valid auth, return 401
565 await send({
566 'type': 'http.response.start',
567 'status': 401,
568 'headers': [[b'content-type', b'text/plain']],
569 })
570 await send({
571 'type': 'http.response.body',
572 'body': b'Unauthorized',
573 })
574 return
575 else:
576 _PRINTER.print("[A2A] No expected token found in settings")
577
578 # Delegate to FastA2A app with cleaned scope
579 with self._lock:
580 app = self.app
581 if app:
582 await app(scope, receive, send)
583 else:
584 # App not configured, return 503
585 await send({
586 'type': 'http.response.start',
587 'status': 503,
588 'headers': [[b'content-type', b'text/plain']],
589 })
590 await send({
591 'type': 'http.response.body',
592 'body': b'FastA2A app not configured',
593 })
594 return
595
596
597 def is_available():
598 """Check if FastA2A is available and properly configured."""
599 return FASTA2A_AVAILABLE and DynamicA2AProxy.get_instance().app is not None
600
601
602 def get_proxy():
603 """Get the FastA2A proxy instance."""
604 return DynamicA2AProxy.get_instance()