Improve code coverage of mcp-agent

pyproject.toml

@@ -71,6 +71,8 @@ dev = [
     "pytest>=7.4.0",
     "pytest-asyncio>=0.21.1",
     "boto3-stubs[bedrock-runtime]>=1.37.23",
+    "trio>=0.30.0",
+    "pytest-cov>=6.1.1",
     "httpx>=0.28.1",
 ]
 

src/mcp_agent/executor/workflow_signal.py

@@ -225,11 +225,11 @@ async def wait_for_signal(
         self, signal, timeout_seconds: int | None = None
     ) -> SignalValueT:
         event = asyncio.Event()
-        unique_name = str(uuid.uuid4())
+        unique_signal_name = f"{signal.name}_{uuid.uuid4()}"
 
         registration = SignalRegistration(
             signal_name=signal.name,
-            unique_name=unique_name,
+            unique_name=unique_signal_name,
             workflow_id=signal.workflow_id,
             run_id=signal.run_id,
         )
@@ -257,20 +257,24 @@ async def wait_for_signal(
                     self._pending_signals[signal.name] = [
                         ps
                         for ps in self._pending_signals[signal.name]
-                        if ps.registration.unique_name != unique_name
+                        if ps.registration.unique_name != unique_signal_name
                     ]
                     if not self._pending_signals[signal.name]:
                         del self._pending_signals[signal.name]
 
     def on_signal(self, signal_name):
         def decorator(func):
+            unique_signal_name = f"{signal_name}_{uuid.uuid4()}"
+
             async def wrapped(value: SignalValueT):
                 if asyncio.iscoroutinefunction(func):
                     await func(value)
                 else:
                     func(value)
 
-            self._handlers.setdefault(signal_name, []).append(wrapped)
+            self._handlers.setdefault(signal_name, []).append(
+                [unique_signal_name, wrapped]
+            )
             return wrapped
 
         return decorator

src/mcp_agent/logging/transport.py

@@ -273,6 +273,7 @@ def __init__(self, transport: EventTransport | None = None):
         self.listeners: Dict[str, EventListener] = {}
         self._task: asyncio.Task | None = None
         self._running = False
+        self.init_queue()
 
     def init_queue(self):
         if self._running:

src/mcp_agent/mcp/mcp_aggregator.py

@@ -187,6 +187,7 @@ async def close(self):
                 not self.connection_persistence
                 or not self._persistent_connection_manager
             ):
+                self.initialized = False
                 return
 
             try:
@@ -979,8 +980,9 @@ def getter(item: NamespacedPrompt):
         else:
             raise ValueError(f"Unsupported capability: {capability}")
 
-        # Search across all servers
-        for srv_name, items in capability_map.items():
+        # Search servers in the order of self.server_names
+        for srv_name in self.server_names:
+            items = capability_map.get(srv_name, [])
             for item in items:
                 if getter(item) == name:
                     return srv_name, name

src/mcp_agent/utils/pydantic_type_serializer.py

@@ -181,6 +181,14 @@ def serialize_type(typ: Any) -> Dict[str, Any]:
         origin = get_origin(typ)
         if origin is not None:
             args = get_args(typ)
+            # Special handling for Literal: store raw values, not types
+            if origin is Literal:
+                return {
+                    "kind": "generic",
+                    "origin": "Literal",
+                    "literal_values": [make_serializable(a) for a in args],
+                    "repr": str(typ),
+                }
             serialized_args = [
                 PydanticTypeSerializer.serialize_type(arg) for arg in args
             ]
@@ -292,7 +300,7 @@ def _serialize_validators(model_class: Type[BaseModel]) -> List[Dict[str, Any]]:
                 name,
                 validator,
             ) in model_class.__pydantic_decorators__.field_validators.items():
-                field_names = [str(f) for f in validator.fields]
+                field_names = [str(f) for f in validator.info.fields]
                 validators.append(
                     {
                         "type": "field_validator",
@@ -382,9 +390,11 @@ def _serialize_fields(model_class: Type[BaseModel]) -> Dict[str, Dict[str, Any]]
                     "description": make_serializable(
                         getattr(field_info, "description", None)
                     ),
-                    "required": make_serializable(
-                        getattr(field_info, "required", True)
-                    ),
+                    "required": getattr(
+                        field_info,
+                        "is_required",
+                        lambda: getattr(field_info, "required", True),
+                    )(),
                 }
 
                 # Add constraints if defined
@@ -411,10 +421,10 @@ def _serialize_fields(model_class: Type[BaseModel]) -> Dict[str, Dict[str, Any]]
                 else:
                     default = make_serializable(default)
 
+                # Use type_ if available (Pydantic v2), else fallback to Any
+                attr_type = getattr(private_attr, "type_", Any)
                 private_attrs[name] = {
-                    "type": PydanticTypeSerializer.serialize_type(
-                        private_attr.annotation
-                    ),
+                    "type": PydanticTypeSerializer.serialize_type(attr_type),
                     "default": default,
                 }
 
@@ -436,7 +446,18 @@ def _serialize_config(model_class: Type[BaseModel]) -> Dict[str, Any]:
         else:
             return config_dict
 
-        # Extract serializable config values
+        # If config_source is a dict or ConfigDict (Pydantic v2), just copy its items
+        if isinstance(config_source, dict):
+            for key, value in config_source.items():
+                if not str(key).startswith("_"):
+                    try:
+                        json.dumps({key: value})
+                        config_dict[key] = value
+                    except (TypeError, OverflowError):
+                        config_dict[key] = str(value)
+            return config_dict
+
+        # Otherwise, use inspect.getmembers (for class-based config)
         for key, value in inspect.getmembers(config_source):
             if (
                 not key.startswith("_")
@@ -516,6 +537,12 @@ def deserialize_type(serialized: Dict[str, Any]) -> Any:
         elif kind == "generic":
             # Handle generics like List[int], Dict[str, Model], etc.
             origin_name = serialized["origin"]
+
+            # Special handling for Literal: use literal_values if present
+            if origin_name == "Literal" and "literal_values" in serialized:
+                literal_values = serialized["literal_values"]
+                return Literal.__getitem__(tuple(literal_values))
+
             args = [
                 PydanticTypeSerializer.deserialize_type(arg)
                 for arg in serialized["args"]
@@ -621,15 +648,13 @@ def reconstruct_model(serialized: Dict[str, Any]) -> Type[BaseModel]:
             # Get the field type
             field_type = PydanticTypeSerializer.deserialize_type(field_info["type"])
 
-            # Handle default values
+            # Determine if the field is required
+            is_required = field_info.get("required", True)
             default = field_info.get("default", ...)
-            if default is None and not field_info.get("required", True):
-                default = None
-
-            # Handle default factories
             default_factory = field_info.get("default_factory")
+
+            # This logic ensures that fields with a default or default_factory are not required
             if default_factory:
-                # This is a simplification - ideally we'd recreate the actual factory
                 if default_factory == "list":
                     default_factory = list
                 elif default_factory == "dict":
@@ -658,25 +683,61 @@ def reconstruct_model(serialized: Dict[str, Any]) -> Type[BaseModel]:
 
             # Add the field definition
             if constraints or default_factory:
+                # If there is a default_factory, always use default=... and set default_factory
                 field_definitions[field_name] = (
                     field_type,
                     Field(
-                        default=default if default_factory is None else ...,
+                        default=... if default_factory is not None else default,
                         default_factory=default_factory,
                         **constraints,
                     ),
                 )
             else:
-                field_definitions[field_name] = (field_type, default)
+                if is_required:
+                    field_definitions[field_name] = (field_type, Field(default=...))
+                else:
+                    field_definitions[field_name] = (
+                        field_type,
+                        Field(
+                            default=default,
+                        ),
+                    )
 
         # Create model config
         model_config = ConfigDict(**config_dict) if config_dict else None
 
-        # Create the basic model
+        # Collect private attributes to pass to create_model
+        private_attr_kwargs = {}
+        if "__private_attrs__" in fields:
+            for name, attr_info in fields["__private_attrs__"].items():
+                default = attr_info.get("default")
+                if default == "None":
+                    default = None
+                private_attr_kwargs[name] = PrivateAttr(default=default)
+
+        # Create the basic model, including private attributes in the class namespace
         reconstructed_model = create_model(
-            name, __config__=model_config, **field_definitions
+            name, __config__=model_config, **field_definitions, **private_attr_kwargs
         )
 
+        # Patch __init__ to ensure private attributes are initialized on instance
+        private_attrs = getattr(reconstructed_model, "__private_attributes__", {})
+        if private_attrs:
+            orig_init = reconstructed_model.__init__
+
+            def _init_with_private_attrs(self, *args, **kwargs):
+                orig_init(self, *args, **kwargs)
+                for attr_name, private_attr in private_attrs.items():
+                    # Only set if not already set
+                    if not hasattr(self, attr_name):
+                        default = private_attr.default
+                        # If default is ... (Ellipsis), treat as None
+                        if default is ...:
+                            default = None
+                        setattr(self, attr_name, default)
+
+            reconstructed_model.__init__ = _init_with_private_attrs
+
         # Add validators (this gets complex and may require exec/eval)
         if validators:
             for validator in validators:
@@ -719,17 +780,6 @@ def reconstruct_model(serialized: Dict[str, Any]) -> Type[BaseModel]:
                         except Exception as e:
                             logger.error(f"Error recreating validator: {e}")
 
-        # Add private attributes (simplified)
-        if "__private_attrs__" in fields:
-            for name, attr_info in fields["__private_attrs__"].items():
-                attr_type = PydanticTypeSerializer.deserialize_type(attr_info["type"])
-                default = attr_info.get("default")
-                if default == "None":
-                    default = None
-                reconstructed_model.__private_attributes__[name] = PrivateAttr(
-                    default=default, annotation=attr_type
-                )
-
         return reconstructed_model
 
     @classmethod

src/mcp_agent/workflows/llm/augmented_llm.py

@@ -253,7 +253,12 @@ def __init__(
 
             self.agent = Agent(
                 name=self.name,
-                instruction=self.instruction,
+                # Only pass instruction if it's not None
+                **(
+                    {"instruction": self.instruction}
+                    if self.instruction is not None
+                    else {}
+                ),
                 server_names=server_names or [],
                 llm=self,
             )

src/mcp_agent/workflows/llm/augmented_llm_anthropic.py

@@ -953,27 +953,27 @@ def anthropic_content_to_mcp_content(
         mcp_content.append(TextContent(type="text", text=content))
     else:
         for block in content:
-            if block.type == "text":
-                mcp_content.append(TextContent(type="text", text=block.text))
-            elif block.type == "image":
+            # Handle pydantic models (ContentBlock) and dict blocks
+            if isinstance(block, BaseModel):
+                block_type = block.type
+                block_text = block.text
+            else:
+                block_type = block["type"]
+                block_text = block["text"]
+
+            if block_type == "text":
+                mcp_content.append(TextContent(type="text", text=block_text))
+            elif block_type == "image":
                 raise NotImplementedError("Image content conversion not implemented")
-            elif block.type == "tool_use":
-                # Best effort to convert a tool use to text (since there's no ToolUseContent)
-                mcp_content.append(
-                    TextContent(
-                        type="text",
-                        text=to_string(block),
-                    )
-                )
-            elif block.type == "tool_result":
-                # Best effort to convert a tool result to text (since there's no ToolResultContent)
+            elif block_type == "tool_use" or block_type == "tool_result":
+                # Best effort to convert a tool use and tool result to text (since there's no ToolUseContent or ToolResultContent)
                 mcp_content.append(
                     TextContent(
                         type="text",
                         text=to_string(block),
                     )
                 )
-            elif block.type == "document":
+            elif block_type == "document":
                 raise NotImplementedError("Document content conversion not implemented")
             else:
                 # Last effort to convert the content to a string