Improve maintainability (#109)

* Improve modularity

* update4 docs

* Add unittest

* resolve comments

* more improvements

* resolve comments

* resolve comments

* Fix review comments

* Fix: Improve transport handling in CLI

Author: Agent-Hellboy

.gitignore

@@ -61,7 +61,6 @@ testing_servers/
 nosetests.xml
 coverage.xml
 *.cover
-*.py
 .hypothesis/
 .pytest_cache/
 cover/

README.md

@@ -39,6 +39,15 @@ If your server conforms to the [MCP schema](https://github.com/modelcontextproto
 - Production Ready: Environment detection and production-safe defaults
 - Intelligent Testing: Hypothesis-based data generation with custom strategies
 
+### Extensibility for Contributors
+MCP Server Fuzzer is designed for easy extension while keeping CLI usage simple:
+
+- **Custom Transports**: Add support for new protocols via config or self-registration (see [docs/transport/custom-transports.md](docs/transport/custom-transports.md)).
+- **Pluggable Safety**: Swap safety providers for custom filtering rules.
+- **Injectable Components**: Advanced users can inject custom clients/reporters for testing or plugins.
+
+The modularity improvements (dependency injection, registries) make it maintainer-friendly without complicating the core CLI experience.
+
 ## Quick Start
 
 ### Installation

docs/architecture/client-architecture.md

@@ -81,8 +81,6 @@ import asyncio
 asyncio.run(fuzz_server())
 ```
 
-For backward compatibility, the old `UnifiedMCPFuzzerClient` name is still available as an alias for `MCPFuzzerClient`.
-
 ## Future Improvements
 
 1. **Plugin Architecture**: Add support for custom fuzzing strategies via plugins

docs/components/safety.md

@@ -311,10 +311,6 @@ mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --no
 ### Advanced Safety Configuration
 
 ```bash
-# Custom safety plugin
-mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" \
-  --safety-plugin my_safety_module.SafetyProvider
-
 # Retry with safety on interrupt
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" \
   --retry-with-safety-on-interrupt
@@ -388,12 +384,10 @@ class CustomSafetySystem(SafetySystem):
 ```bash
 # Use custom safety provider
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" \
-  --safety-plugin my_safety_module.CustomSafetyProvider
 
 # With configuration
 export MCP_FUZZER_SAFETY_CONFIG=/path/to/safety_config.json
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" \
-  --safety-plugin my_safety_module.CustomSafetyProvider
 ```
 
 ## Safety Alerts and Logging

docs/development/contributing.md

@@ -485,9 +485,9 @@ The project uses semantic versioning:
 
 - **MAJOR**: Incompatible API changes
 
-- **MINOR**: New functionality (backward compatible)
+- **MINOR**: New functionality that preserves the public API
 
-- **PATCH**: Bug fixes (backward compatible)
+- **PATCH**: Bug fixes that do not change the API
 
 ### Release Steps
 

docs/development/reference.md

@@ -57,9 +57,8 @@ Notes:
 | Option | Type | Default | Description |
 |--------|------|---------|-------------|
 | `--enable-safety-system` | Flag | False | Enable system-level safety features |
+| `--no-safety` | Flag | False | Disable argument-level safety filtering |
 | `--fs-root` | Path | ~/.mcp_fuzzer | Restrict filesystem operations to specified directory |
-| `--safety-plugin` | String | - | Dotted path to custom safety provider |
-| `--no-safety` | Flag | False | Disable argument-level safety filtering (not recommended) |
 | `--retry-with-safety-on-interrupt` | Flag | False | Retry once with safety system enabled on Ctrl-C |
 
 ### Reporting Options
@@ -1103,7 +1102,7 @@ The safety system focuses on containment and preventing external references duri
 
 - Argument-level filtering (`mcp_fuzzer.safety_system.safety.SafetyFilter`):
   - Blocks URLs and risky commands in tool arguments; recursively sanitizes dicts/lists.
-  - Pluggable provider via `--safety-plugin`; `--no-safety` disables filtering.
+  - CLI provides `--fs-root` to redirect sandbox.
   - `set_fs_root(path)` records a sandbox root (for future path checks).
 
 - System-level blocking (`mcp_fuzzer.safety_system.system_blocker.SystemCommandBlocker`):

docs/getting-started/examples.md

@@ -151,14 +151,12 @@ mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --fs
 
 # Disable argument-level safety (not recommended)
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --no-safety
+
 ```
 
 ### Advanced Safety Configuration
 
 ```bash
-# Custom safety plugin
-mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --safety-plugin my_safety_module.SafetyProvider
-
 # Retry with safety on interrupt
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --retry-with-safety-on-interrupt
 
@@ -925,4 +923,3 @@ mcp-fuzzer --mode tools --phase aggressive --protocol http --endpoint http://loc
 # Monitor performance
 mcp-fuzzer --mode tools --protocol http --endpoint http://localhost:8000 --runs 100 --log-level INFO
 These examples cover the most common use cases and should help you get started with MCP Server Fuzzer. For more advanced configurations and customizations, refer to the [Reference](reference.md), [Architecture](architecture.md), and [Runtime Management](runtime-management.md) documentation.
-

docs/getting-started/getting-started.md

@@ -225,14 +225,15 @@ mcp-fuzzer --mode tools --auth-env --endpoint http://localhost:8000
 ### Basic Safety Features
 
 ```bash
-# Enable safety system
+# Enable safety system (default)
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --enable-safety-system
 
 # Set filesystem root
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --fs-root /tmp/safe_dir
 
 # Disable argument-level safety (not recommended)
 mcp-fuzzer --mode tools --protocol stdio --endpoint "python test_server.py" --no-safety
+
 ```
 
 ### Safety System Features

docs/transport/custom-transports.md

@@ -61,6 +61,11 @@ class MyCustomTransport(TransportProtocol):
         # Implement notification sending logic
         pass
 
+    async def close(self) -> None:
+        """Close the transport connection."""
+        # Cleanup logic
+        pass
+
     async def _stream_request(
         self, payload: Dict[str, Any]
     ) -> AsyncIterator[Dict[str, Any]]:
@@ -71,72 +76,46 @@ class MyCustomTransport(TransportProtocol):
             yield {}
 ```
 
-### 2. Register the Transport
-
-#### Option A: Programmatic Registration
-
-```python
-from mcp_fuzzer.transport import register_custom_transport
-
-# Register your custom transport
-register_custom_transport(
-    name="mytransport",
-    transport_class=MyCustomTransport,
-    description="My custom MCP transport implementation",
-    config_schema={
-        "type": "object",
-        "properties": {
-            "timeout": {"type": "number"},
-            "retries": {"type": "integer"}
-        }
-    }
-)
-```
-
-#### Option B: Configuration File Registration
+### 2. Register Your Transport
 
-Add to your `mcp-fuzzer.yaml` configuration file:
+Add your custom transport to the configuration file:
 
 ```yaml
+# config.yaml
 custom_transports:
-  mytransport:
+  my-custom:
     module: "my_package.transports"
     class: "MyCustomTransport"
-    description: "My custom MCP transport implementation"
-    config:
-      timeout: 30
-      retries: 3
+    # Optional: additional init kwargs
+    extra_arg: "value"
 ```
 
-### 3. Use the Custom Transport
+### 3. Use in CLI
 
-Once registered, your custom transport can be used in several ways:
+Now you can use your custom transport via the CLI:
 
-#### Via URL Scheme
-
-```python
-from mcp_fuzzer.transport import create_transport
-
-# Use custom transport with URL scheme
-transport = create_transport("mytransport://my-endpoint")
+```bash
+mcp-fuzzer --protocol my-custom --config config.yaml --endpoint my-endpoint
 ```
 
-#### Via Factory Function
+## Advanced: Self-Registration with Registry
+
+The transport factory now uses a `TransportRegistry` for built-in transports. For custom transports, you can optionally self-register in your module for even easier extension:
 
 ```python
-from mcp_fuzzer.transport import create_custom_transport
+# In my_package/transports.py
+from mcp_fuzzer.transport.factory import registry
+
+class MyCustomTransport(TransportProtocol):
+    # ... implementation ...
 
-# Create custom transport instance directly
-transport = create_custom_transport("mytransport", "my-endpoint")
+# Self-register (runs when module is imported)
+registry.register("my-custom", MyCustomTransport)
 ```
 
-#### Via Configuration
+This makes extension simpler – no factory changes needed, and CLI usage remains unchanged.
 
-```yaml
-# In mcp-fuzzer.yaml
-protocol: mytransport
-endpoint: "my-endpoint"
-```
+Note: Self-registration is optional; config-based registration (step 2) still works and is recommended for most cases.
 
 ## Example: WebSocket Transport
 

docs/transport/transport-improvements.md

@@ -231,7 +231,7 @@ async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]:
 
 ### For Existing Code
 
-The improvements are backward compatible. Existing code will continue to work without changes:
+Existing code continues to work without changes:
 
 ```python
 # This still works exactly the same