autobyteus 1.1.1__py3-none-any.whl → 1.1.3__py3-none-any.whl

autobyteus/agent/bootstrap_steps/__init__.py

@@ -9,6 +9,7 @@ from .agent_runtime_queue_initialization_step import AgentRuntimeQueueInitializa
 from .workspace_context_initialization_step import WorkspaceContextInitializationStep
 # ToolInitializationStep is no longer a bootstrap step.
 from .system_prompt_processing_step import SystemPromptProcessingStep
+from .mcp_server_prewarming_step import McpServerPrewarmingStep
 # LLMConfigFinalizationStep and LLMInstanceCreationStep removed.
 
 __all__ = [
@@ -16,4 +17,5 @@ __all__ = [
     "AgentRuntimeQueueInitializationStep", # UPDATED
     "WorkspaceContextInitializationStep",
     "SystemPromptProcessingStep",
+    "McpServerPrewarmingStep",
 ]

autobyteus/agent/bootstrap_steps/agent_bootstrapper.py

@@ -6,6 +6,7 @@ from .base_bootstrap_step import BaseBootstrapStep
 from .agent_runtime_queue_initialization_step import AgentRuntimeQueueInitializationStep
 from .workspace_context_initialization_step import WorkspaceContextInitializationStep
 from .system_prompt_processing_step import SystemPromptProcessingStep
+from .mcp_server_prewarming_step import McpServerPrewarmingStep
 from autobyteus.agent.events import AgentReadyEvent
 
 if TYPE_CHECKING:
@@ -31,6 +32,7 @@ class AgentBootstrapper:
             self.bootstrap_steps: List[BaseBootstrapStep] = [
                 AgentRuntimeQueueInitializationStep(),
                 WorkspaceContextInitializationStep(),
+                McpServerPrewarmingStep(),
                 SystemPromptProcessingStep(),
             ]
             logger.debug("AgentBootstrapper initialized with default steps.")

autobyteus/agent/bootstrap_steps/mcp_server_prewarming_step.py

@@ -0,0 +1,71 @@
+# file: autobyteus/autobyteus/agent/bootstrap_steps/mcp_server_prewarming_step.py
+import logging
+from typing import TYPE_CHECKING, Set
+
+from .base_bootstrap_step import BaseBootstrapStep
+from autobyteus.tools.mcp.config_service import McpConfigService
+from autobyteus.tools.mcp.server_instance_manager import McpServerInstanceManager
+from autobyteus.tools.tool_category import ToolCategory
+
+if TYPE_CHECKING:
+    from autobyteus.agent.context import AgentContext
+    from autobyteus.agent.phases import AgentPhaseManager
+
+logger = logging.getLogger(__name__)
+
+class McpServerPrewarmingStep(BaseBootstrapStep):
+    """
+    Bootstrap step to eagerly start all MCP servers associated with an agent's tools.
+    This ensures servers are running and ready before the agent becomes idle.
+    """
+
+    def __init__(self):
+        self._config_service = McpConfigService()
+        self._instance_manager = McpServerInstanceManager()
+        logger.debug("McpServerPrewarmingStep initialized.")
+
+    async def execute(self,
+                      context: 'AgentContext',
+                      phase_manager: 'AgentPhaseManager') -> bool:
+        agent_id = context.agent_id
+        logger.info(f"Agent '{agent_id}': Executing McpServerPrewarmingStep.")
+
+        # 1. Find all unique server IDs by inspecting tool definitions.
+        mcp_server_ids: Set[str] = set()
+        for tool in context.config.tools:
+            # This is the new, superior check. It relies on abstract metadata, not concrete types.
+            if tool.definition and tool.definition.category == ToolCategory.MCP:
+                # This is the new, superior way to get the server_id.
+                # It does not rely on private attributes of the tool instance.
+                server_id = tool.definition.metadata.get("mcp_server_id")
+                if server_id:
+                    mcp_server_ids.add(server_id)
+
+        if not mcp_server_ids:
+            logger.debug(f"Agent '{agent_id}': No MCP tools found. Skipping MCP server pre-warming.")
+            return True
+
+        logger.info(f"Agent '{agent_id}': Found {len(mcp_server_ids)} unique MCP server IDs to pre-warm: {mcp_server_ids}")
+
+        # 2. For each server ID, unconditionally start its server instance for this agent.
+        for server_id in mcp_server_ids:
+            try:
+                config = self._config_service.get_config(server_id)
+                if not config:
+                    logger.warning(f"Agent '{agent_id}': Could not find config for server_id '{server_id}' used by a tool. Cannot pre-warm.")
+                    continue
+
+                logger.info(f"Agent '{agent_id}': Pre-warming MCP server '{server_id}'.")
+                # Get the instance for this agent, which creates it if it doesn't exist.
+                server_instance = self._instance_manager.get_server_instance(agent_id, server_id)
+                # Explicitly connect to start the server process.
+                await server_instance.connect()
+                logger.info(f"Agent '{agent_id}': Successfully connected to pre-warmed MCP server '{server_id}'.")
+
+            except Exception as e:
+                error_message = f"Agent '{agent_id}': Failed to pre-warm MCP server '{server_id}': {e}"
+                logger.error(error_message, exc_info=True)
+                # A failure to pre-warm a server is a critical bootstrap failure.
+                return False
+        
+        return True

autobyteus/agent/llm_response_processor/provider_aware_tool_usage_processor.py

@@ -1,11 +1,12 @@
 # file: autobyteus/autobyteus/agent/llm_response_processor/provider_aware_tool_usage_processor.py
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List
 
 from .base_processor import BaseLLMResponseProcessor
+from autobyteus.agent.events import PendingToolInvocationEvent
+from autobyteus.agent.tool_invocation import ToolInvocation
 from autobyteus.tools.usage.parsers import ProviderAwareToolUsageParser
 from autobyteus.tools.usage.parsers.exceptions import ToolUsageParseException
-from autobyteus.agent.events import PendingToolInvocationEvent
 
 if TYPE_CHECKING:
     from autobyteus.agent.context import AgentContext
@@ -17,21 +18,20 @@ logger = logging.getLogger(__name__)
 class ProviderAwareToolUsageProcessor(BaseLLMResponseProcessor):
     """
     A "master" tool usage processor that uses a high-level parser from the
-    `tools` module to extract tool invocations, and then enqueues the
-    necessary agent events based on the parsed results.
+    `tools` module to extract tool invocations. It then ensures each invocation
+    has a session-unique ID before enqueuing the necessary agent events.
     """
+    INVOCATION_COUNTS_KEY = "agent_tool_invocation_counts"
+
     def __init__(self):
         self._parser = ProviderAwareToolUsageParser()
         logger.debug("ProviderAwareToolUsageProcessor initialized.")
 
-    @classmethod
-    def get_name(cls) -> str:
-        return "provider_aware_tool_usage"
-
     async def process_response(self, response: 'CompleteResponse', context: 'AgentContext', triggering_event: 'LLMCompleteResponseReceivedEvent') -> bool:
         """
-        Uses a ProviderAwareToolUsageParser to get a list of tool invocations,
-        and then enqueues a PendingToolInvocationEvent for each one.
+        Uses a ProviderAwareToolUsageParser to get tool invocations, makes their
+        IDs unique within the agent's session, and then enqueues a
+        PendingToolInvocationEvent for each one.
         Propagates ToolUsageParseException if parsing fails.
         """
         try:
@@ -44,9 +44,38 @@ class ProviderAwareToolUsageProcessor(BaseLLMResponseProcessor):
         if not tool_invocations:
             return False
 
-        logger.info(f"Agent '{context.agent_id}': Parsed {len(tool_invocations)} tool invocations. Enqueuing events.")
+        # --- NEW LOGIC FOR UNIQUE ID GENERATION ---
+        
+        # Ensure the counter map exists in the agent's state's custom data
+        if self.INVOCATION_COUNTS_KEY not in context.custom_data:
+            context.custom_data[self.INVOCATION_COUNTS_KEY] = {}
+        
+        invocation_counts = context.custom_data[self.INVOCATION_COUNTS_KEY]
+        
+        processed_invocations: List[ToolInvocation] = []
+
         for invocation in tool_invocations:
-            logger.info(f"Agent '{context.agent_id}' ({self.get_name()}) identified tool invocation: {invocation.name}. Enqueuing event.")
+            base_id = invocation.id
+            
+            # Get the current count for this base ID, default to 0
+            count = invocation_counts.get(base_id, 0)
+            
+            # Create the new session-unique ID
+            unique_id = f"{base_id}_{count}"
+            
+            # Update the invocation's ID in-place
+            invocation.id = unique_id
+            
+            # Increment the counter for the next time this base ID is seen
+            invocation_counts[base_id] = count + 1
+            
+            processed_invocations.append(invocation)
+
+        # --- END NEW LOGIC ---
+
+        logger.info(f"Agent '{context.agent_id}': Parsed {len(processed_invocations)} tool invocations. Enqueuing events with unique IDs.")
+        for invocation in processed_invocations:
+            logger.info(f"Agent '{context.agent_id}' ({self.get_name()}) identified tool invocation: {invocation.name} with unique ID {invocation.id}. Enqueuing event.")
             await context.input_event_queues.enqueue_tool_invocation_request(
                 PendingToolInvocationEvent(tool_invocation=invocation)
             )

autobyteus/agent/runtime/agent_runtime.py

@@ -120,10 +120,7 @@ class AgentRuntime:
         await self.phase_manager.notify_shutdown_initiated() 
         await self._worker.stop(timeout=timeout) 
         
-        if self.context.llm_instance and hasattr(self.context.llm_instance, 'cleanup'):
-            cleanup_func = self.context.llm_instance.cleanup
-            if asyncio.iscoroutinefunction(cleanup_func): await cleanup_func()
-            else: cleanup_func()
+        # LLM instance cleanup is now handled by the AgentWorker before its loop closes.
         
         await self.phase_manager.notify_final_shutdown_complete() 
         logger.info(f"AgentRuntime for '{agent_id}' stop() method completed.")

autobyteus/agent/runtime/agent_worker.py

@@ -15,6 +15,7 @@ from autobyteus.agent.events import (
 from autobyteus.agent.events import WorkerEventDispatcher
 from autobyteus.agent.runtime.agent_thread_pool_manager import AgentThreadPoolManager 
 from autobyteus.agent.bootstrap_steps.agent_bootstrapper import AgentBootstrapper
+from autobyteus.agent.shutdown_steps import AgentShutdownOrchestrator
 
 if TYPE_CHECKING:
     from autobyteus.agent.context import AgentContext
@@ -140,19 +141,19 @@ class AgentWorker:
 
     async def async_run(self) -> None:
         agent_id = self.context.agent_id
-        logger.info(f"AgentWorker '{agent_id}' async_run(): Starting.")
-        
-        # --- Direct Initialization ---
-        initialization_successful = await self._initialize()
-        if not initialization_successful:
-            logger.critical(f"AgentWorker '{agent_id}' failed to initialize. Worker is shutting down.")
-            if self._async_stop_event and not self._async_stop_event.is_set():
-                self._async_stop_event.set()
-            return
-        
-        # --- Main Event Loop ---
-        logger.info(f"AgentWorker '{agent_id}' initialized successfully. Entering main event loop.")
         try:
+            logger.info(f"AgentWorker '{agent_id}' async_run(): Starting.")
+            
+            # --- Direct Initialization ---
+            initialization_successful = await self._initialize()
+            if not initialization_successful:
+                logger.critical(f"AgentWorker '{agent_id}' failed to initialize. Worker is shutting down.")
+                if self._async_stop_event and not self._async_stop_event.is_set():
+                    self._async_stop_event.set()
+                return
+
+            # --- Main Event Loop ---
+            logger.info(f"AgentWorker '{agent_id}' initialized successfully. Entering main event loop.")
             while not self._async_stop_event.is_set(): 
                 try:
                     queue_event_tuple = await asyncio.wait_for(
@@ -176,25 +177,57 @@ class AgentWorker:
             logger.error(f"Fatal error in AgentWorker '{agent_id}' async_run() loop: {e}", exc_info=True)
         finally:
             logger.info(f"AgentWorker '{agent_id}' async_run() loop has finished.")
+            # --- Shutdown sequence moved here, inside the original task's finally block ---
+            logger.info(f"AgentWorker '{agent_id}': Running shutdown sequence on worker loop.")
+            orchestrator = AgentShutdownOrchestrator()
+            cleanup_successful = await orchestrator.run(self.context)
+
+            if not cleanup_successful:
+                logger.critical(f"AgentWorker '{agent_id}': Shutdown resource cleanup failed. The agent may not have shut down cleanly.")
+            else:
+                logger.info(f"AgentWorker '{agent_id}': Shutdown resource cleanup completed successfully.")
+            logger.info(f"AgentWorker '{agent_id}': Shutdown sequence completed.")
 
-    async def _signal_internal_stop(self):
-        if self._async_stop_event and not self._async_stop_event.is_set():
-            self._async_stop_event.set()
-            if self.context.state.input_event_queues:
-                await self.context.state.input_event_queues.enqueue_internal_system_event(AgentStoppedEvent())
 
     async def stop(self, timeout: float = 10.0) -> None:
+        """
+        Gracefully stops the worker by signaling its event loop to terminate,
+        then waiting for the thread to complete its cleanup and exit.
+        """
         if not self._is_active or self._stop_initiated:
             return
+        
+        agent_id = self.context.agent_id
+        logger.info(f"AgentWorker '{agent_id}': Stop requested.")
         self._stop_initiated = True
-        if self.get_worker_loop() and self._async_stop_event:
-            future = asyncio.run_coroutine_threadsafe(self._signal_internal_stop(), self.get_worker_loop())
-            try: future.result(timeout=1.0)
-            except Exception: pass
+
+        # Schedule a coroutine on the worker's loop to set the stop event.
+        if self.get_worker_loop():
+            def _coro_factory():
+                async def _signal_coro():
+                    if self._async_stop_event and not self._async_stop_event.is_set():
+                        self._async_stop_event.set()
+                        if self.context.state.input_event_queues:
+                            await self.context.state.input_event_queues.enqueue_internal_system_event(AgentStoppedEvent())
+                return _signal_coro()
+            
+            future = self.schedule_coroutine_on_worker_loop(_coro_factory)
+            try:
+                # Wait for the signal to be processed.
+                future.result(timeout=max(1.0, timeout-1))
+            except Exception as e:
+                logger.error(f"AgentWorker '{agent_id}': Error signaling stop event: {e}", exc_info=True)
+
+        # Wait for the main thread future to complete.
         if self._thread_future:
-            try: await asyncio.wait_for(asyncio.wrap_future(self._thread_future), timeout=timeout)
-            except asyncio.TimeoutError: logger.warning(f"Timeout waiting for worker thread of '{self.context.agent_id}'.")
+            try:
+                await asyncio.wait_for(asyncio.wrap_future(self._thread_future), timeout=timeout)
+                logger.info(f"AgentWorker '{agent_id}': Worker thread has terminated.")
+            except asyncio.TimeoutError:
+                logger.warning(f"AgentWorker '{agent_id}': Timeout waiting for worker thread to terminate.")
+        
         self._is_active = False
 
+
     def is_alive(self) -> bool:
         return self._thread_future is not None and not self._thread_future.done()

autobyteus/agent/shutdown_steps/__init__.py

@@ -0,0 +1,17 @@
+# file: autobyteus/autobyteus/agent/shutdown_steps/__init__.py
+"""
+Defines individual, self-contained steps for the agent shutdown process.
+These steps are orchestrated by the AgentShutdownOrchestrator.
+"""
+
+from .base_shutdown_step import BaseShutdownStep
+from .llm_instance_cleanup_step import LLMInstanceCleanupStep
+from .mcp_server_cleanup_step import McpServerCleanupStep
+from .agent_shutdown_orchestrator import AgentShutdownOrchestrator
+
+__all__ = [
+    "BaseShutdownStep",
+    "LLMInstanceCleanupStep",
+    "McpServerCleanupStep",
+    "AgentShutdownOrchestrator",
+]

autobyteus/agent/shutdown_steps/agent_shutdown_orchestrator.py

@@ -0,0 +1,63 @@
+# file: autobyteus/autobyteus/agent/shutdown_steps/agent_shutdown_orchestrator.py
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+from .base_shutdown_step import BaseShutdownStep
+from .llm_instance_cleanup_step import LLMInstanceCleanupStep
+from .mcp_server_cleanup_step import McpServerCleanupStep
+
+if TYPE_CHECKING:
+    from autobyteus.agent.context import AgentContext
+
+logger = logging.getLogger(__name__)
+
+class AgentShutdownOrchestrator:
+    """
+    Orchestrates the agent's shutdown process by executing a sequence of
+    self-contained cleanup steps.
+    """
+    def __init__(self, steps: Optional[List[BaseShutdownStep]] = None):
+        """
+        Initializes the AgentShutdownOrchestrator.
+
+        Args:
+            steps: An optional list of shutdown steps to execute. If not provided,
+                   a default sequence will be used.
+        """
+        if steps is None:
+            self.shutdown_steps: List[BaseShutdownStep] = [
+                LLMInstanceCleanupStep(),
+                McpServerCleanupStep(),
+            ]
+            logger.debug("AgentShutdownOrchestrator initialized with default steps.")
+        else:
+            self.shutdown_steps = steps
+            logger.debug(f"AgentShutdownOrchestrator initialized with {len(steps)} custom steps.")
+
+    async def run(self, context: 'AgentContext') -> bool:
+        """
+        Executes the configured sequence of shutdown steps.
+
+        Args:
+            context: The agent's context.
+
+        Returns:
+            True if all steps completed successfully, False otherwise.
+        """
+        agent_id = context.agent_id
+        logger.info(f"Agent '{agent_id}': AgentShutdownOrchestrator starting execution.")
+
+        for step_index, step_instance in enumerate(self.shutdown_steps):
+            step_name = step_instance.__class__.__name__
+            logger.debug(f"Agent '{agent_id}': Executing shutdown step {step_index + 1}/{len(self.shutdown_steps)}: {step_name}")
+            
+            success = await step_instance.execute(context)
+            
+            if not success:
+                error_message = f"Shutdown step {step_name} failed."
+                logger.error(f"Agent '{agent_id}': {error_message} Halting shutdown orchestration.")
+                # The step itself is responsible for detailed error logging.
+                return False
+
+        logger.info(f"Agent '{agent_id}': All shutdown steps completed successfully.")
+        return True

autobyteus/agent/shutdown_steps/base_shutdown_step.py

@@ -0,0 +1,33 @@
+# file: autobyteus/autobyteus/agent/shutdown_steps/base_shutdown_step.py
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from autobyteus.agent.context import AgentContext
+
+logger = logging.getLogger(__name__)
+
+class BaseShutdownStep(ABC):
+    """
+    Abstract base class for individual steps in the agent shutdown process.
+    Each step is responsible for a specific part of the cleanup and
+    for reporting its success or failure.
+    """
+
+    @abstractmethod
+    async def execute(self, context: 'AgentContext') -> bool:
+        """
+        Executes the shutdown step.
+
+        Args:
+            context: The agent's context, providing access to state and resources.
+
+        Returns:
+            True if the step completed successfully, False otherwise.
+            If False, the step is expected to have handled its own detailed logging.
+        """
+        raise NotImplementedError("Subclasses must implement the 'execute' method.")
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}>"

autobyteus/agent/shutdown_steps/llm_instance_cleanup_step.py

@@ -0,0 +1,45 @@
+# file: autobyteus/autobyteus/agent/shutdown_steps/llm_instance_cleanup_step.py
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from .base_shutdown_step import BaseShutdownStep
+
+if TYPE_CHECKING:
+    from autobyteus.agent.context import AgentContext
+
+logger = logging.getLogger(__name__)
+
+class LLMInstanceCleanupStep(BaseShutdownStep):
+    """
+    Shutdown step for cleaning up the agent's LLM instance.
+    """
+    def __init__(self):
+        logger.debug("LLMInstanceCleanupStep initialized.")
+
+    async def execute(self, context: 'AgentContext') -> bool:
+        agent_id = context.agent_id
+        logger.info(f"Agent '{agent_id}': Executing LLMInstanceCleanupStep.")
+
+        llm_instance = context.llm_instance
+        if not llm_instance:
+            logger.debug(f"Agent '{agent_id}': No LLM instance found in context. Skipping cleanup.")
+            return True
+
+        if hasattr(llm_instance, 'cleanup') and callable(getattr(llm_instance, 'cleanup')):
+            try:
+                logger.info(f"Agent '{agent_id}': Running LLM instance cleanup for '{llm_instance.__class__.__name__}'.")
+                cleanup_func = llm_instance.cleanup
+                if asyncio.iscoroutinefunction(cleanup_func):
+                    await cleanup_func()
+                else:
+                    cleanup_func()
+                logger.info(f"Agent '{agent_id}': LLM instance cleanup completed successfully.")
+                return True
+            except Exception as e:
+                error_message = f"Agent '{agent_id}': Error during LLM instance cleanup: {e}"
+                logger.error(error_message, exc_info=True)
+                return False
+        else:
+            logger.debug(f"Agent '{agent_id}': LLM instance of type '{llm_instance.__class__.__name__}' does not have a 'cleanup' method. Skipping.")
+            return True

autobyteus/agent/shutdown_steps/mcp_server_cleanup_step.py

@@ -0,0 +1,32 @@
+# file: autobyteus/autobyteus/agent/shutdown_steps/mcp_server_cleanup_step.py
+import logging
+from typing import TYPE_CHECKING
+
+from .base_shutdown_step import BaseShutdownStep
+from autobyteus.tools.mcp.server_instance_manager import McpServerInstanceManager
+
+if TYPE_CHECKING:
+    from autobyteus.agent.context import AgentContext
+
+logger = logging.getLogger(__name__)
+
+class McpServerCleanupStep(BaseShutdownStep):
+    """
+    Shutdown step for cleaning up all MCP server instances associated with an agent.
+    """
+    def __init__(self):
+        self._instance_manager = McpServerInstanceManager()
+        logger.debug("McpServerCleanupStep initialized.")
+
+    async def execute(self, context: 'AgentContext') -> bool:
+        agent_id = context.agent_id
+        logger.info(f"Agent '{agent_id}': Executing McpServerCleanupStep.")
+
+        try:
+            await self._instance_manager.cleanup_mcp_server_instances_for_agent(agent_id)
+            logger.info(f"Agent '{agent_id}': MCP server instance cleanup completed successfully.")
+            return True
+        except Exception as e:
+            error_message = f"Agent '{agent_id}': Critical failure during McpServerCleanupStep: {e}"
+            logger.error(error_message, exc_info=True)
+            return False