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

autobyteus/llm/llm_factory.py

@@ -3,7 +3,7 @@ import logging
 import inspect
 
 from autobyteus.llm.autobyteus_provider import AutobyteusModelProvider
-from autobyteus.llm.models import LLMModel
+from autobyteus.llm.models import LLMModel, ModelInfo, ProviderModelGroup
 from autobyteus.llm.providers import LLMProvider
 from autobyteus.llm.utils.llm_config import LLMConfig, TokenPricingConfig
 from autobyteus.llm.base_llm import BaseLLM
@@ -36,21 +36,6 @@ class LLMFactory(metaclass=SingletonMeta):
             LLMFactory._initialize_registry()
             LLMFactory._initialized = True
 
-    @staticmethod
-    def _clear_model_class_attributes():
-        """
-        Clear all LLMModel instances that were set as class attributes on the LLMModel class.
-        This is necessary for reinitialization to avoid 'model already exists' errors.
-        """
-        # Get all attributes of the LLMModel class
-        for attr_name in list(vars(LLMModel).keys()):
-            attr_value = getattr(LLMModel, attr_name)
-            # Check if the attribute is an instance of LLMModel
-            if isinstance(attr_value, LLMModel):
-                logger.debug(f"Removing LLMModel class attribute: {attr_name}")
-                # Delete the attribute to avoid 'model already exists' errors during reinitialization
-                delattr(LLMModel, attr_name)
-
     @staticmethod
     def reinitialize():
         """
@@ -66,9 +51,6 @@ class LLMFactory(metaclass=SingletonMeta):
         try:
             logger.info("Reinitializing LLM model registry...")
             
-            # Clear all LLMModel instances set as class attributes
-            LLMFactory._clear_model_class_attributes()
-            
             # Reset the initialized flag
             LLMFactory._initialized = False
             
@@ -94,7 +76,7 @@ class LLMFactory(metaclass=SingletonMeta):
         supported_models = [
             # OPENAI Provider Models
             LLMModel(
-                name="GPT_4o_API",
+                name="gpt-4o",
                 value="gpt-4o",
                 provider=LLMProvider.OPENAI,
                 llm_class=OpenAILLM,
@@ -106,7 +88,7 @@ class LLMFactory(metaclass=SingletonMeta):
                 )
             ),
             LLMModel(
-                name="o3_API",
+                name="o3",
                 value="o3",
                 provider=LLMProvider.OPENAI,
                 llm_class=OpenAILLM,
@@ -116,7 +98,7 @@ class LLMFactory(metaclass=SingletonMeta):
                 )
             ),
             LLMModel(
-                name="o4_MINI_API",
+                name="o4-mini",
                 value="o4-mini",
                 provider=LLMProvider.OPENAI,
                 llm_class=OpenAILLM,
@@ -127,7 +109,7 @@ class LLMFactory(metaclass=SingletonMeta):
             ),
             # MISTRAL Provider Models
             LLMModel(
-                name="MISTRAL_LARGE_API",
+                name="mistral-large",
                 value="mistral-large-latest",
                 provider=LLMProvider.MISTRAL,
                 llm_class=MistralLLM,
@@ -138,28 +120,48 @@ class LLMFactory(metaclass=SingletonMeta):
             ),
             # ANTHROPIC Provider Models
             LLMModel(
-                name="CLAUDE_3_7_SONNET_API",
-                value="claude-3-7-sonnet-20250219",
+                name="claude-4-opus",
+                value="claude-opus-4-20250514",
+                provider=LLMProvider.ANTHROPIC,
+                llm_class=ClaudeLLM,
+                canonical_name="claude-4-opus",
+                default_config=LLMConfig(
+                    pricing_config=TokenPricingConfig(15.00, 75.00)
+                )
+            ),
+            LLMModel(
+                name="bedrock-claude-4-opus",
+                value="anthropic.claude-opus-4-20250514-v1:0",
                 provider=LLMProvider.ANTHROPIC,
                 llm_class=ClaudeLLM,
-                canonical_name="claude-3.7",
+                canonical_name="claude-4-opus",
+                default_config=LLMConfig(
+                    pricing_config=TokenPricingConfig(15.00, 75.00)
+                )
+            ),
+            LLMModel(
+                name="claude-4-sonnet",
+                value="claude-sonnet-4-20250514",
+                provider=LLMProvider.ANTHROPIC,
+                llm_class=ClaudeLLM,
+                canonical_name="claude-4-sonnet",
                 default_config=LLMConfig(
                     pricing_config=TokenPricingConfig(3.00, 15.00)
                 )
             ),
             LLMModel(
-                name="BEDROCK_CLAUDE_3_7_SONNET_API",
-                value="anthropic.claude-3-7-sonnet-20250219-v1:0",
+                name="bedrock-claude-4-sonnet",
+                value="anthropic.claude-sonnet-4-20250514-v1:0",
                 provider=LLMProvider.ANTHROPIC,
                 llm_class=ClaudeLLM,
-                canonical_name="claude-3.7",
+                canonical_name="claude-4-sonnet",
                 default_config=LLMConfig(
                     pricing_config=TokenPricingConfig(3.00, 15.00)
                 )
             ),
             # DEEPSEEK Provider Models
             LLMModel(
-                name="DEEPSEEK_CHAT_API",
+                name="deepseek-chat",
                 value="deepseek-chat",
                 provider=LLMProvider.DEEPSEEK,
                 llm_class=DeepSeekLLM,
@@ -172,7 +174,7 @@ class LLMFactory(metaclass=SingletonMeta):
             ),
             # Adding deepseek-reasoner support
             LLMModel(
-                name="DEEPSEEK_REASONER_API",
+                name="deepseek-reasoner",
                 value="deepseek-reasoner",
                 provider=LLMProvider.DEEPSEEK,
                 llm_class=DeepSeekLLM,
@@ -185,28 +187,48 @@ class LLMFactory(metaclass=SingletonMeta):
             ),
             # GEMINI Provider Models
             LLMModel(
-                name="GEMINI_1_5_PRO_API",
-                value="gemini-1-5-pro",
+                name="gemini-2.5-pro",
+                value="gemini-2.5-pro",
+                provider=LLMProvider.GEMINI,
+                llm_class=OpenAILLM,
+                canonical_name="gemini-2.5-pro",
+                default_config=LLMConfig(
+                    pricing_config=TokenPricingConfig(2.50, 10.00)
+                )
+            ),
+            LLMModel(
+                name="gemini-2.5-flash",
+                value="gemini-2.5-flash",
                 provider=LLMProvider.GEMINI,
                 llm_class=OpenAILLM,
-                canonical_name="gemini-1.5-pro",
+                canonical_name="gemini-2.5-flash",
                 default_config=LLMConfig(
-                    pricing_config=TokenPricingConfig(1.25, 5.00)
+                    pricing_config=TokenPricingConfig(0.15, 0.60)
                 )
             ),
             LLMModel(
-                name="GEMINI_1_5_FLASH_API",
-                value="gemini-1-5-flash",
+                name="gemini-2.0-flash",
+                value="gemini-2.0-flash",
                 provider=LLMProvider.GEMINI,
                 llm_class=OpenAILLM,
-                canonical_name="gemini-1.5-flash",
+                canonical_name="gemini-2.0-flash",
+                default_config=LLMConfig(
+                    pricing_config=TokenPricingConfig(0.1, 0.40)
+                )
+            ),
+            LLMModel(
+                name="gemini-2.0-flash-lite",
+                value="gemini-2.0-flash-lite",
+                provider=LLMProvider.GEMINI,
+                llm_class=OpenAILLM,
+                canonical_name="gemini-2.0-flash-lite",
                 default_config=LLMConfig(
                     pricing_config=TokenPricingConfig(0.075, 0.30)
                 )
             ),
             # GROK Provider Models
             LLMModel(
-                name="GROK_2_1212_API",
+                name="grok-2-1212",
                 value="grok-2-1212",
                 provider=LLMProvider.GROK,
                 llm_class=GrokLLM,
@@ -228,7 +250,18 @@ class LLMFactory(metaclass=SingletonMeta):
     def register_model(model: LLMModel):
         """
         Register a new LLM model, storing it under its provider category.
+        If a model with the same name already exists, it will be replaced.
         """
+        # Using a flat list of all models to check for existing model by name
+        all_models = [m for models in LLMFactory._models_by_provider.values() for m in models]
+        
+        for existing_model in all_models:
+            if existing_model.name == model.name:
+                logger.warning(f"Model with name '{model.name}' is being redefined.")
+                # Remove the old model from its provider list
+                LLMFactory._models_by_provider[existing_model.provider].remove(existing_model)
+                break
+
         models = LLMFactory._models_by_provider.setdefault(model.provider, [])
         models.append(model)
 
@@ -238,7 +271,7 @@ class LLMFactory(metaclass=SingletonMeta):
         Create an LLM instance for the specified model identifier.
         
         Args:
-            model_identifier (str): The model name or value to create an instance for.
+            model_identifier (str): The model name to create an instance for.
             llm_config (Optional[LLMConfig]): Configuration for the LLM. If None,
                                              the model's default configuration is used.
         
@@ -251,7 +284,7 @@ class LLMFactory(metaclass=SingletonMeta):
         LLMFactory.ensure_initialized()
         for models in LLMFactory._models_by_provider.values():
             for model_instance in models:
-                if model_instance.value == model_identifier or model_instance.name == model_identifier:
+                if model_instance.name == model_identifier:
                     return model_instance.create_llm(llm_config)
         raise ValueError(f"Unsupported model: {model_identifier}")
 
@@ -296,7 +329,7 @@ class LLMFactory(metaclass=SingletonMeta):
         Get the canonical name for a model by its name.
         
         Args:
-            model_name (str): The model name (e.g., "GPT_4o_API")
+            model_name (str): The model name (e.g., "gpt_4o")
             
         Returns:
             Optional[str]: The canonical name if found, None otherwise
@@ -308,4 +341,35 @@ class LLMFactory(metaclass=SingletonMeta):
                     return model_instance.canonical_name
         return None
 
+    @staticmethod
+    def get_models_grouped_by_provider() -> List[ProviderModelGroup]:
+        """
+        Returns a list of all providers, each with a list of its available models,
+        sorted by provider name and model name. Providers with no models are included
+        with an empty model list.
+        """
+        LLMFactory.ensure_initialized()
+        result: List[ProviderModelGroup] = []
+        # Sort all providers from the enum by name for consistent order
+        all_providers_sorted = sorted(list(LLMProvider), key=lambda p: p.name)
+        
+        for provider in all_providers_sorted:
+            # Get models for the current provider, defaults to [] if none are registered
+            models = LLMFactory._models_by_provider.get(provider, [])
+            
+            # Sort the models for this provider by name
+            sorted_models = sorted(models, key=lambda model: model.name)
+            
+            model_infos = [
+                ModelInfo(name=model.name, canonical_name=model.canonical_name)
+                for model in sorted_models
+            ]
+            
+            result.append(ProviderModelGroup(
+                provider=provider.name,
+                models=model_infos
+            ))
+            
+        return result
+
 default_llm_factory = LLMFactory()

autobyteus/llm/models.py

@@ -1,5 +1,6 @@
 import logging
 from typing import TYPE_CHECKING, Type, Optional, List, Iterator
+from dataclasses import dataclass
 
 from autobyteus.llm.providers import LLMProvider
 from autobyteus.llm.utils.llm_config import LLMConfig
@@ -10,6 +11,18 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+@dataclass
+class ModelInfo:
+    """A simple data structure for essential model information."""
+    name: str
+    canonical_name: str
+
+@dataclass
+class ProviderModelGroup:
+    """A data structure to group models by their provider."""
+    provider: str
+    models: List[ModelInfo]
+
 class LLMModelMeta(type):
     """
     Metaclass for LLMModel to make it iterable and support item access like Enums.
@@ -24,15 +37,12 @@ class LLMModelMeta(type):
         from autobyteus.llm.llm_factory import LLMFactory
         LLMFactory.ensure_initialized()
 
-        for attr_name in dir(cls):
-            if not attr_name.startswith('_'):  # Skip private/dunder attributes
-                attr_value = getattr(cls, attr_name)
-                if isinstance(attr_value, cls):  # Check if it's an LLMModel instance
-                    yield attr_value
+        for models in LLMFactory._models_by_provider.values():
+            yield from models
 
     def __getitem__(cls, name_or_value: str) -> 'LLMModel':
         """
-        Allows dictionary-like access to LLMModel instances by name (e.g., 'GPT_4o_API')
+        Allows dictionary-like access to LLMModel instances by name (e.g., 'gpt-4o')
         or by value (e.g., 'gpt-4o').
         Search is performed by name first, then by value.
         """
@@ -40,13 +50,12 @@ class LLMModelMeta(type):
         from autobyteus.llm.llm_factory import LLMFactory
         LLMFactory.ensure_initialized()
 
-        # 1. Try to find by name first (e.g., LLMModel['GPT_4o_API'])
-        if hasattr(cls, name_or_value):
-            attribute = getattr(cls, name_or_value)
-            if isinstance(attribute, cls):
-                return attribute
+        # 1. Try to find by name first
+        for model in cls:
+            if model.name == name_or_value:
+                return model
         
-        # 2. If not found by name, iterate and find by value (e.g., LLMModel['gpt-4o'])
+        # 2. If not found by name, iterate and find by value
         for model in cls:
             if model.value == name_or_value:
                 return model
@@ -64,17 +73,14 @@ class LLMModelMeta(type):
         LLMFactory.ensure_initialized()
         
         count = 0
-        for attr_name in dir(cls):
-            if not attr_name.startswith('_'):
-                attr_value = getattr(cls, attr_name)
-                if isinstance(attr_value, cls):
-                    count += 1
+        for models in LLMFactory._models_by_provider.values():
+            count += len(models)
         return count
 
 class LLMModel(metaclass=LLMModelMeta):
     """
     Represents a single model's metadata:
-      - name (str): A human-readable label, e.g. "GPT-4 Official" 
+      - name (str): A human-readable label, e.g. "gpt-4o"
       - value (str): A unique identifier used in code or APIs, e.g. "gpt-4o"
       - canonical_name (str): A shorter, standardized reference name for prompts, e.g. "gpt-4o" or "claude-3.7"
       - provider (LLMProvider): The provider enum 
@@ -94,12 +100,6 @@ class LLMModel(metaclass=LLMModelMeta):
         canonical_name: str,
         default_config: Optional[LLMConfig] = None
     ):
-        # Validate name doesn't already exist as a class attribute
-        if hasattr(LLMModel, name):
-            existing_model = getattr(LLMModel, name)
-            if isinstance(existing_model, LLMModel):
-                logger.warning(f"Model with name '{name}' is being redefined. This is expected during reinitialization.")
-            
         self._name = name
         self._value = value
         self._canonical_name = canonical_name
@@ -107,16 +107,12 @@ class LLMModel(metaclass=LLMModelMeta):
         self.llm_class = llm_class
         self.default_config = default_config if default_config else LLMConfig()
 
-        # Set this instance as a class attribute, making LLMModel.MODEL_NAME available.
-        logger.debug(f"Setting LLMModel class attribute: {name}")
-        setattr(LLMModel, name, self)
-
     @property
     def name(self) -> str:
         """
         A friendly or descriptive name for this model (could appear in UI).
         This is the key used for `LLMModel['MODEL_NAME']` access.
-        Example: "GPT_4o_API"
+        Example: "gpt-4o"
         """
         return self._name
 

autobyteus/llm/ollama_provider.py

@@ -2,6 +2,7 @@ from autobyteus.llm.models import LLMModel
 from autobyteus.llm.api.ollama_llm import OllamaLLM
 from autobyteus.llm.providers import LLMProvider
 from autobyteus.llm.utils.llm_config import LLMConfig, TokenPricingConfig
+from autobyteus.llm.ollama_provider_resolver import OllamaProviderResolver
 from typing import TYPE_CHECKING
 import os
 import logging
@@ -39,7 +40,7 @@ class OllamaModelProvider:
         try:
             from autobyteus.llm.llm_factory import LLMFactory  # Local import to avoid circular dependency
             
-            ollama_host = os.getenv('OLLAMA_HOST', OllamaLLM.DEFAULT_OLLAMA_HOST)
+            ollama_host = os.getenv('DEFAULT_OLLAMA_HOST', OllamaLLM.DEFAULT_OLLAMA_HOST)
             
             if not OllamaModelProvider.is_valid_url(ollama_host):
                 logger.error(f"Invalid Ollama host URL: {ollama_host}")
@@ -73,11 +74,14 @@ class OllamaModelProvider:
                     model_name = model_info.get('model')
                     if not model_name:
                         continue
+
+                    # Determine the provider based on the model name
+                    provider = OllamaProviderResolver.resolve(model_name)
                         
                     llm_model = LLMModel(
                         name=model_name,
                         value=model_name,
-                        provider=LLMProvider.OLLAMA,
+                        provider=provider,
                         llm_class=OllamaLLM,
                         canonical_name=model_name,  # Use model_name as the canonical_name
                         default_config=LLMConfig(

autobyteus/llm/ollama_provider_resolver.py

@@ -0,0 +1,44 @@
+from autobyteus.llm.providers import LLMProvider
+import logging
+
+logger = logging.getLogger(__name__)
+
+class OllamaProviderResolver:
+    """
+    A utility class to resolve the correct LLMProvider for Ollama models
+    based on keywords in their names. This helps attribute models to their
+    original creators (e.g., Google for 'gemma').
+    """
+    
+    # A mapping from keywords to providers. The list is ordered to handle
+    # potential overlaps, though current keywords are distinct.
+    KEYWORD_PROVIDER_MAP = [
+        (['gemma', 'gemini'], LLMProvider.GEMINI),
+        (['llama'], LLMProvider.GROQ),
+        (['mistral'], LLMProvider.MISTRAL),
+        (['deepseek'], LLMProvider.DEEPSEEK),
+    ]
+
+    @staticmethod
+    def resolve(model_name: str) -> LLMProvider:
+        """
+        Resolves the LLMProvider for a given model name from Ollama.
+        It checks for keywords in the model name and returns the corresponding
+        provider. If no specific provider is found, it defaults to OLLAMA.
+
+        Args:
+            model_name (str): The name of the model discovered from Ollama (e.g., 'gemma:7b').
+
+        Returns:
+            LLMProvider: The resolved provider for the model.
+        """
+        lower_model_name = model_name.lower()
+        
+        for keywords, provider in OllamaProviderResolver.KEYWORD_PROVIDER_MAP:
+            for keyword in keywords:
+                if keyword in lower_model_name:
+                    logger.debug(f"Resolved provider for model '{model_name}' to '{provider.name}' based on keyword '{keyword}'.")
+                    return provider
+                    
+        logger.debug(f"Model '{model_name}' did not match any specific provider keywords. Defaulting to OLLAMA provider.")
+        return LLMProvider.OLLAMA

autobyteus/tools/__init__.py

@@ -10,6 +10,7 @@ from .base_tool import BaseTool
 from .functional_tool import tool # The @tool decorator
 from .parameter_schema import ParameterSchema, ParameterDefinition, ParameterType
 from .tool_config import ToolConfig # Configuration data object, primarily for class-based tools
+from .tool_category import ToolCategory
 
 # --- Re-export specific tools for easier access ---
 
@@ -47,6 +48,7 @@ __all__ = [
     "ParameterDefinition",
     "ParameterType",
     "ToolConfig",
+    "ToolCategory",
 
     # Re-exported functional tool instances
     "ask_user_input",

autobyteus/tools/base_tool.py

@@ -9,10 +9,13 @@ from autobyteus.events.event_emitter import EventEmitter
 from autobyteus.events.event_types import EventType
 
 from .tool_meta import ToolMeta
+from .tool_state import ToolState
+
 if TYPE_CHECKING:
     from autobyteus.agent.context import AgentContext
     from autobyteus.tools.parameter_schema import ParameterSchema
     from autobyteus.tools.tool_config import ToolConfig
+    from .tool_state import ToolState
 
 logger = logging.getLogger('autobyteus')
 
@@ -25,7 +28,10 @@ class BaseTool(ABC, EventEmitter, metaclass=ToolMeta):
         self.agent_id: Optional[str] = None
         # The config is stored primarily for potential use by subclasses or future base features.
         self._config = config
-        logger.debug(f"BaseTool instance initializing for potential class {self.__class__.__name__}")
+        # Add a dedicated state dictionary for the tool instance
+        # CHANGED: Use ToolState class for explicit state management.
+        self.tool_state: 'ToolState' = ToolState()
+        logger.debug(f"BaseTool instance initializing for potential class {self.__class__.__name__}. tool_state initialized.")
 
     @classmethod
     def get_name(cls) -> str:

autobyteus/tools/functional_tool.py

@@ -8,6 +8,7 @@ from autobyteus.tools.base_tool import BaseTool
 from autobyteus.tools.parameter_schema import ParameterSchema, ParameterDefinition, ParameterType
 from autobyteus.tools.tool_config import ToolConfig
 from autobyteus.tools.registry import default_tool_registry, ToolDefinition
+from autobyteus.tools.tool_category import ToolCategory
 
 if TYPE_CHECKING:
     from autobyteus.agent.context import AgentContext 
@@ -28,15 +29,20 @@ class FunctionalTool(BaseTool):
                  config_schema: Optional[ParameterSchema],
                  is_async: bool,
                  expects_context: bool,
+                 expects_tool_state: bool,
                  func_param_names: TypingList[str],
                  instantiation_config: Optional[Dict[str, Any]] = None):
         super().__init__(config=ToolConfig(params=instantiation_config) if instantiation_config else None)
         self._original_func = original_func
         self._is_async = is_async
         self._expects_context = expects_context
+        self._expects_tool_state = expects_tool_state
         self._func_param_names = func_param_names
         self._instantiation_config = instantiation_config or {}
         
+        # This instance has its own state dictionary, inherited from BaseTool's __init__
+        # self.tool_state: Dict[str, Any] = {} # This is now handled by super().__init__()
+        
         # Override instance methods to provide specific schema info
         self.get_name = lambda: name
         self.get_description = lambda: description
@@ -65,6 +71,9 @@ class FunctionalTool(BaseTool):
         
         if self._expects_context:
             call_args['context'] = context
+        
+        if self._expects_tool_state:
+            call_args['tool_state'] = self.tool_state
             
         if self._is_async:
             return await self._original_func(**call_args)
@@ -143,15 +152,19 @@ def _get_parameter_type_from_hint(py_type: Any, param_name: str) -> Tuple[Parame
     logger.warning(f"Unmapped type hint {py_type} (actual_type: {actual_type}) for param '{param_name}'. Defaulting to ParameterType.STRING.")
     return ParameterType.STRING, None
 
-def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList[str], bool, ParameterSchema]:
+def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList[str], bool, bool, ParameterSchema]:
     func_param_names = []
     expects_context = False
+    expects_tool_state = False
     generated_arg_schema = ParameterSchema()
 
     for param_name, param_obj in sig.parameters.items():
         if param_name == "context":
             expects_context = True
-            continue 
+            continue
+        if param_name == "tool_state":
+            expects_tool_state = True
+            continue
         
         func_param_names.append(param_name)
         
@@ -177,7 +190,7 @@ def _parse_signature(sig: inspect.Signature, tool_name: str) -> Tuple[TypingList
         )
         generated_arg_schema.add_parameter(schema_param)
         
-    return func_param_names, expects_context, generated_arg_schema
+    return func_param_names, expects_context, expects_tool_state, generated_arg_schema
 
 # --- The refactored @tool decorator ---
 
@@ -196,7 +209,7 @@ def tool(
         
         sig = inspect.signature(func)
         is_async = inspect.iscoroutinefunction(func)
-        func_param_names, expects_context, gen_arg_schema = _parse_signature(sig, tool_name)
+        func_param_names, expects_context, expects_tool_state, gen_arg_schema = _parse_signature(sig, tool_name)
         
         final_arg_schema = argument_schema if argument_schema is not None else gen_arg_schema
 
@@ -209,6 +222,7 @@ def tool(
                 config_schema=config_schema,
                 is_async=is_async,
                 expects_context=expects_context,
+                expects_tool_state=expects_tool_state,
                 func_param_names=func_param_names,
                 instantiation_config=inst_config.params if inst_config else None
             )
@@ -221,7 +235,8 @@ def tool(
             argument_schema=final_arg_schema,
             config_schema=config_schema,
             custom_factory=factory,
-            tool_class=None
+            tool_class=None,
+            category=ToolCategory.LOCAL
         )
         default_tool_registry.register_tool(tool_def)
         

autobyteus/tools/mcp/call_handlers/stdio_handler.py

@@ -1,5 +1,6 @@
 # file: autobyteus/autobyteus/tools/mcp/call_handlers/stdio_handler.py
 import logging
+import asyncio
 from typing import Dict, Any, cast, TYPE_CHECKING
 
 from .base_handler import McpCallHandler
@@ -11,6 +12,9 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+# A default timeout for STDIO subprocesses to prevent indefinite hangs.
+DEFAULT_STDIO_TIMEOUT = 30  # seconds
+
 class StdioMcpCallHandler(McpCallHandler):
     """Handles MCP tool calls over a stateless STDIO transport."""
 
@@ -23,6 +27,7 @@ class StdioMcpCallHandler(McpCallHandler):
         """
         Creates a new subprocess, establishes a session, and executes the
         requested tool call. It handles 'list_tools' as a special case.
+        Includes a timeout to prevent hanging on unresponsive subprocesses.
         """
         logger.debug(f"Handling STDIO call to tool '{remote_tool_name}' on server '{config.server_id}'.")
         
@@ -39,7 +44,8 @@ class StdioMcpCallHandler(McpCallHandler):
             cwd=stdio_config.cwd
         )
 
-        try:
+        async def _perform_call():
+            """Inner function to be wrapped by the timeout."""
             # The stdio_client context manager provides the read/write streams.
             async with stdio_client(mcp_lib_stdio_params) as (read_stream, write_stream):
                 # The ClientSession is its own context manager that handles initialization.
@@ -54,6 +60,14 @@ class StdioMcpCallHandler(McpCallHandler):
                     
                     logger.debug(f"STDIO call to tool '{remote_tool_name}' on server '{config.server_id}' completed.")
                     return result
+
+        try:
+            return await asyncio.wait_for(_perform_call(), timeout=DEFAULT_STDIO_TIMEOUT)
+        except asyncio.TimeoutError:
+            error_message = (f"MCP call to '{remote_tool_name}' on server '{config.server_id}' timed out "
+                             f"after {DEFAULT_STDIO_TIMEOUT} seconds. The subprocess may have hung.")
+            logger.error(error_message)
+            raise RuntimeError(error_message) from None
         except Exception as e:
             logger.error(
                 f"An error occurred during STDIO tool call to '{remote_tool_name}' on server '{config.server_id}': {e}",