using callback to make agent stateless

Author: daltonnyx

AgentCrew/__init__.py

@@ -1 +1 @@
-__version__ = "0.6.6-3"
+__version__ = "0.6.6-4"

AgentCrew/main_docker.py

@@ -284,6 +284,7 @@ def setup_agents(services, config_path, remoting_provider=None, model_id=None):
             if remoting_provider:
                 agent.set_custom_system_prompt(agent_manager.get_remote_system_prompt())
                 agent.is_remoting_mode = True
+                agent.activate()
         agent_manager.register_agent(agent)
 
     from AgentCrew.modules.mcpclient.tool import register as mcp_register
@@ -468,13 +469,13 @@ def chat(provider, agent_config, mcp_config, memory_llm):
 
         ui = ConsoleUI(message_handler)
         ui.start()
-    except SystemExit:
-        from AgentCrew.modules.mcpclient import MCPSessionManager
-
-        MCPSessionManager.get_instance().cleanup()
     except Exception as e:
         print(traceback.format_exc())
         click.echo(f"❌ Error: {str(e)}", err=True)
+    finally:
+        from AgentCrew.modules.mcpclient import MCPSessionManager
+
+        MCPSessionManager.get_instance().cleanup()
 
 
 @cli.command()
@@ -562,13 +563,13 @@ def a2a_server(
         click.echo(f"Starting A2A server on {host}:{port}")
         click.echo(f"Available agents: {', '.join(agent_manager.agents.keys())}")
         server.start()
-    except SystemExit:
-        from AgentCrew.modules.mcpclient import MCPSessionManager
-
-        MCPSessionManager.get_instance().cleanup()
     except Exception as e:
         print(traceback.format_exc())
         click.echo(f"❌ Error: {str(e)}", err=True)
+    finally:
+        from AgentCrew.modules.mcpclient import MCPSessionManager
+
+        MCPSessionManager.get_instance().cleanup()
 
 
 @cli.command()

AgentCrew/modules/a2a/task_manager.py

@@ -82,8 +82,6 @@ async def on_send_message(
                 )
             )
 
-        agent.activate()
-
         # Generate task ID from message
         task_id = (
             request.params.message.taskId or f"task_{request.params.message.messageId}"
@@ -195,6 +193,9 @@ async def _process_agent_task(self, agent: LocalAgent, task: Task):
             if task.id not in self.task_history:
                 raise ValueError("Task history is not existed")
 
+            input_tokens = 0
+            output_tokens = 0
+
             async def _process_task():
                 # Process with agent
 
@@ -203,12 +204,21 @@ async def _process_task():
                 response_message = ""
                 thinking_content = ""
                 thinking_signature = ""
+                tool_uses = []
+
+                def process_result(_tool_uses, _input_tokens, _output_tokens):
+                    nonlocal tool_uses, input_tokens, output_tokens
+                    tool_uses = _tool_uses
+                    input_tokens += _input_tokens
+                    output_tokens += _output_tokens
 
                 async for (
                     response_message,
                     chunk_text,
                     thinking_chunk,
-                ) in agent.process_messages(self.task_history[task.id]):
+                ) in agent.process_messages(
+                    self.task_history[task.id], callback=process_result
+                ):
                     # Update current response
                     if response_message:
                         current_response = response_message
@@ -260,8 +270,6 @@ async def _process_task():
                                 )
                             )
 
-                # Get final result
-                tool_uses, input_tokens, output_tokens = agent.get_process_result()
                 if tool_uses and len(tool_uses) > 0:
                     # Add thinking content as a separate message if available
                     thinking_data = (

AgentCrew/modules/agents/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import AsyncGenerator, Tuple, Dict, List, Optional, Any
+from typing import AsyncGenerator, Tuple, Dict, List, Optional, Any, Callable
 from enum import Enum
 
 
@@ -77,7 +77,9 @@ def calculate_usage_cost(self, input_tokens, output_tokens) -> float:
 
     @abstractmethod
     async def process_messages(
-        self, messages: Optional[List[Dict[str, Any]]] = None
+        self,
+        messages: Optional[List[Dict[str, Any]]] = None,
+        callback: Optional[Callable] = None,
     ) -> AsyncGenerator:
         """
         Process messages using this agent.

AgentCrew/modules/agents/local_agent.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 import os
-from typing import Dict, Any, List, Optional
+from typing import Dict, Any, List, Optional, Callable
 from AgentCrew.modules.llm.base import BaseLLMService
 from AgentCrew.modules.llm.message import MessageTransformer
 from AgentCrew.modules.agents.base import BaseAgent, MessageType
@@ -385,7 +385,11 @@ def update_llm_service(self, new_llm_service: BaseLLMService) -> bool:
 
         return True
 
-    async def process_messages(self, messages: Optional[List[Dict[str, Any]]] = None):
+    async def process_messages(
+        self,
+        messages: Optional[List[Dict[str, Any]]] = None,
+        callback: Optional[Callable] = None,
+    ):
         """
         Process messages using this agent.
 
@@ -400,9 +404,9 @@ async def process_messages(self, messages: Optional[List[Dict[str, Any]]] = None
         )
 
         assistant_response = ""
-        self.tool_uses = []
-        self.input_tokens_usage = 0
-        self.output_tokens_usage = 0
+        _tool_uses = []
+        _input_tokens_usage = 0
+        _output_tokens_usage = 0
         # Ensure the first message is a system message with the agent's prompt
         if not messages:
             final_messages = list(self.history)
@@ -452,15 +456,23 @@ async def process_messages(self, messages: Optional[List[Dict[str, Any]]] = None
                         chunk_text,
                         thinking_chunk,
                     ) = self.llm.process_stream_chunk(
-                        chunk, assistant_response, self.tool_uses
+                        chunk, assistant_response, _tool_uses
                     )
+                    yield (assistant_response, chunk_text, thinking_chunk)
+
                     if tool_uses:
-                        self.tool_uses = tool_uses
+                        _tool_uses = tool_uses
                     if chunk_input_tokens > 0:
-                        self.input_tokens_usage = chunk_input_tokens
+                        _input_tokens_usage = chunk_input_tokens
                     if chunk_output_tokens > 0:
-                        self.output_tokens_usage = chunk_output_tokens
-                    yield (assistant_response, chunk_text, thinking_chunk)
+                        _output_tokens_usage = chunk_output_tokens
+            if callback:
+                callback(_tool_uses, _input_tokens_usage, _output_tokens_usage)
+            else:
+                self.tool_uses = _tool_uses
+                self.input_tokens_usage = _input_tokens_usage
+                self.output_tokens_usage = _output_tokens_usage
+
         except GeneratorExit as e:
             logger.warning(f"Stream processing interrupted: {e}")
         finally:

AgentCrew/modules/agents/remote_agent.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, List, Optional, Tuple
+from typing import Callable, Dict, Any, List, Optional, Tuple
 from uuid import uuid4
 
 from pydantic import ValidationError
@@ -95,7 +95,11 @@ def configure_think(self, think_setting):
     def calculate_usage_cost(self, input_tokens, output_tokens) -> float:
         return 0.0
 
-    async def process_messages(self, messages: Optional[List[Dict[str, Any]]] = None):
+    async def process_messages(
+        self,
+        messages: Optional[List[Dict[str, Any]]] = None,
+        callback: Optional[Callable] = None,
+    ):
         if not self.client or not self.agent_card:
             raise ValidationError(
                 f"RemoteAgent '{self.name}' not properly initialized."

AgentCrew/modules/chat/message/handler.py

@@ -195,9 +195,15 @@ async def get_assistant_response(
         # Create a reference to the streaming generator
         self.stream_generator = None
 
+        def process_result(_tool_uses, _input_tokens, _output_tokens):
+            nonlocal tool_uses, input_tokens, output_tokens
+            tool_uses = _tool_uses
+            input_tokens += _input_tokens
+            output_tokens += _output_tokens
+
         try:
             # Store the generator in a variable so we can properly close it if needed
-            self.stream_generator = self.agent.process_messages()
+            self.stream_generator = self.agent.process_messages(callback=process_result)
 
             async for (
                 assistant_response,
@@ -239,12 +245,6 @@ async def get_assistant_response(
                         time.sleep(0.5)
                     self._notify("response_chunk", (chunk_text, assistant_response))
 
-            tool_uses, input_tokens_in_turn, output_tokens_in_turn = (
-                self.agent.get_process_result()
-            )
-            input_tokens += input_tokens_in_turn
-            output_tokens += output_tokens_in_turn
-
             # Handle tool use if needed
             if not has_stop_interupted and tool_uses and len(tool_uses) > 0:
                 # Add thinking content as a separate message if available

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "agentcrew-ai"
-version = "0.6.6-3"
+version = "0.6.6-4"
 requires-python = ">=3.12"
 classifiers = [
     "Programming Language :: Python :: 3",