fix: make a2a chats temporary
Rafael Uzarowski committed
Aug 1, 2025 at 09:18 UTC
40daadc54fc0230df2fd75927bd4ae180b1bd1ff
1 file changed
+21
-9
python/helpers/fasta2a_server.py
+21
-9
@@ -13,6 +13,7 @@ from starlette.requests import Request
13
from python.helpers.print_style import PrintStyle
14
from agent import AgentContext, UserMessage, AgentContextType
15
from initialize import initialize_agent
16
+from python.helpers.persist_chat import remove_chat
17
18
# Import FastA2A
19
try:
@@ -69,22 +70,19 @@ class AgentZeroWorker(Worker): # type: ignore[misc]
70
71
async def run_task(self, params: Any) -> None: # params: TaskSendParams
72
"""Execute a task by processing the message through Agent Zero."""
73
+ context = None
74
try:
75
task_id = params['id']
74
- context_id = params['context_id']
76
message = params['message']
77
77
- _PRINTER.print(f"[A2A] Processing task {task_id} in context {context_id}")
78
+ _PRINTER.print(f"[A2A] Processing task {task_id} with new temporary context")
79
80
# Convert A2A message to Agent Zero format
81
agent_message = self._convert_message(message)
82
82
- # Get or create Agent Zero context
83
- context = AgentContext.get(context_id)
84
- if not context:
85
- # Create new context for this A2A conversation
86
- cfg = initialize_agent()
87
- context = AgentContext(cfg, id=context_id, type=AgentContextType.BACKGROUND)
83
+ # Always create new temporary context for this A2A conversation
84
+ cfg = initialize_agent()
85
+ context = AgentContext(cfg, type=AgentContextType.BACKGROUND)
86
87
# Log user message so it appears instantly in UI chat window
88
context.log.log(
@@ -113,7 +111,12 @@ class AgentZeroWorker(Worker): # type: ignore[misc]
111
new_messages=[response_message]
112
)
113
116
- _PRINTER.print(f"[A2A] Completed task {task_id}")
114
+ # Clean up context like non-persistent MCP chats
115
+ context.reset()
116
+ AgentContext.remove(context.id)
117
+ remove_chat(context.id)
118
+
119
+ _PRINTER.print(f"[A2A] Completed task {task_id} and cleaned up context")
120
121
except Exception as e:
122
_PRINTER.print(f"[A2A] Error processing task {params.get('id', 'unknown')}: {e}")
@@ -122,12 +125,21 @@ class AgentZeroWorker(Worker): # type: ignore[misc]
125
state='failed'
126
)
127
128
+ # Clean up context even on failure to prevent resource leaks
129
+ if context:
130
+ context.reset()
131
+ AgentContext.remove(context.id)
132
+ remove_chat(context.id)
133
+ _PRINTER.print(f"[A2A] Cleaned up failed context {context.id}")
134
+
135
async def cancel_task(self, params: Any) -> None: # params: TaskIdParams
136
"""Cancel a running task."""
137
task_id = params['id']
138
_PRINTER.print(f"[A2A] Cancelling task {task_id}")
139
await self.storage.update_task(task_id=task_id, state='canceled') # type: ignore[attr-defined]
140
141
+ # Note: No context cleanup needed since contexts are always temporary and cleaned up in run_task
142
+
143
def build_message_history(self, history: List[Any]) -> List[Message]: # type: ignore
144
# Not used in this simplified implementation
145
return []