Add custom transport support (fixes #9) #82
Conversation
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #82 +/- ##
==========================================
- Coverage 58.32% 57.95% -0.37%
==========================================
Files 53 54 +1
Lines 4393 4626 +233
Branches 747 789 +42
==========================================
+ Hits 2562 2681 +119
- Misses 1608 1707 +99
- Partials 223 238 +15 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
WalkthroughAdds a custom transport registry with YAML-driven registration and factory integration, a WebSocket example and docs, transport batch APIs and response collation, realistic-phase strategy expansion with batch generation and out-of-order helpers, phase-aware and batch fuzzing in ProtocolFuzzer, tests, and multiple formatting tweaks. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant User
participant ConfigLoader
participant Importer as Importlib
participant Registry as CustomTransportRegistry
participant Factory as TransportFactory
Note over ConfigLoader,Registry: Config-driven custom transport loading
User->>ConfigLoader: apply_config_file(path)
ConfigLoader->>ConfigLoader: parse YAML -> custom_transports
loop each transport entry
ConfigLoader->>Importer: import module / resolve factory
Importer-->>ConfigLoader: class / factory
ConfigLoader->>Registry: register(name, class, description, schema, factory)
Registry-->>ConfigLoader: ack
end
Note over Factory: Creation prefers custom registry
User->>Factory: create_transport("scheme", url_or_endpoint, kwargs)
Factory->>Registry: create_transport(scheme, url_or_endpoint, kwargs)
alt registered
Registry-->>User: transport instance
else not registered
Factory-->>User: fallback builtin transport or raise
end
sequenceDiagram
autonumber
participant Fuzzer as ProtocolFuzzer
participant Strategies as StrategyManager
participant Transport
Note over Fuzzer,Transport: Batch fuzzing flow
Fuzzer->>Strategies: generate_batch_request(phase)
Strategies-->>Fuzzer: batch [req / notification ...]
alt generate_only
Fuzzer-->>Fuzzer: return generated batch
else send batch
Fuzzer->>Transport: send_batch_request(batch)
Transport-->>Fuzzer: responses[]
Fuzzer->>Transport: collate_batch_responses(batch, responses)
Transport-->>Fuzzer: id→response map
Fuzzer-->>Fuzzer: create batch FuzzDataResult
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Assessment against linked issues
Assessment against linked issues: Out-of-scope changes
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 14
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mcp_fuzzer/transport/factory.py (1)
21-35: Support two-arg custom transports and unify error messageIn the two-argument branch (mcp_fuzzer/transport/factory.py lines 21–35), try
custom_registry.create_transport(...)before falling back to built-ins, and update theValueErrorto list both built-ins and registered customs:# Back-compat path: two-argument usage if endpoint is not None: - key = url_or_protocol.strip().lower() - mapping = { + key = url_or_protocol.strip().lower() + # Try custom transports first + try: + return custom_registry.create_transport(key, endpoint, **kwargs) + except KeyError: + pass + mapping = { "http": HTTPTransport, "https": HTTPTransport, "streamablehttp": StreamableHTTPTransport, "sse": SSETransport, "stdio": StdioTransport, } try: transport_cls = mapping[key] except KeyError: - raise ValueError(f"Unsupported protocol: {url_or_protocol}") + raise ValueError( + f"Unsupported protocol: {url_or_protocol}. " + f"Supported: http, https, sse, stdio, streamablehttp; " + f"custom: {', '.join(sorted(custom_registry.list_transports().keys()))}" + ) return transport_cls(endpoint, **kwargs)
🧹 Nitpick comments (18)
examples/README-custom-transports.md (3)
69-75: Minor style polish for the bullet list.Add parallel punctuation for consistency.
-- Connection management -- JSON-RPC request/response handling -- Error handling and timeouts -- Streaming support -- Configuration schema definition -- Registration with MCP Fuzzer +- Connection management. +- JSON-RPC request/response handling. +- Error handling and timeouts. +- Streaming support. +- Configuration schema definition. +- Registration with MCP Fuzzer.
131-135: Tiny grammar tweak.“In Python path” → “on the Python path.”
-1. **Import Errors**: Ensure your transport module is in Python path +1. **Import Errors**: Ensure your transport module is on the Python path
140-142: Consider documenting env-based logging and secrets.Point readers to environment variables for secrets/log level (tie-in with the YAML example).
-Enable debug logging to troubleshoot issues: +Enable debug logging to troubleshoot issues (avoid hardcoding secrets; prefer environment variables):examples/config/custom-transport-config.yaml (1)
12-16: Avoid plaintext secrets in examples; prefer env expansion.Recommend illustrating environment variable expansion to discourage committing tokens.
- headers: - Authorization: "Bearer your-token-here" - X-API-Key: "your-api-key" + headers: + # Prefer env vars; implement expansion in config_loader if not present. + Authorization: "${AUTH_BEARER:?}" + X-API-Key: "${API_KEY:?}"I can add opt-in ${VAR} expansion to config_loader and update docs/tests—want a patch?
docs/custom-transports.md (2)
333-336: Tiny grammar nit.“in the Python path” → “on the Python path”.
-2. **Registration Failures**: Check that your transport class inherits from `TransportProtocol` -3. **Connection Issues**: Verify endpoint URLs and network connectivity +2. **Registration Failures**: Check that your transport class inherits from `TransportProtocol`. +3. **Connection Issues**: Verify endpoint URLs and network connectivity.
11-17: Soften “Type Safety” claim or cross-link validation hook.Unless config_loader enforces config_schema, consider softening language or adding a note pointing to where schema validation occurs.
tests/unit/transport/test_custom.py (3)
4-4: Remove unused mocks import.
MockandAsyncMockare not used.Apply:
-from unittest.mock import Mock, AsyncMock
171-176: Avoid mutating private state in tests.Accessing
registry._transports.clear()couples tests to internals. Prefer an explicit API.Proposed addition outside this file (mcp_fuzzer/transport/custom.py):
class CustomTransportRegistry: ... def clear(self) -> None: self._transports.clear()Then update here:
- from mcp_fuzzer.transport.custom import registry - registry._transports.clear() + from mcp_fuzzer.transport.custom import registry + registry.clear()
201-221: Add coverage for 2-arg back-compat with custom protocols.Current tests don't cover
create_transport("custom", "endpoint"). The factory’s 2-arg path currently rejects customs (see separate comment).Suggested test:
@@ def test_custom_transport_via_factory(self): """Test creating custom transport via factory.""" register_custom_transport( name="factory_mock", transport_class=MockTransport, description="Factory mock transport", ) transport = create_transport("factory_mock://test-endpoint") assert isinstance(transport, MockTransport) assert transport.endpoint == "test-endpoint" + def test_custom_transport_via_factory_two_args(self): + """Back-compat: (protocol, endpoint) for custom transports.""" + register_custom_transport( + name="factory_mock", + transport_class=MockTransport, + description="Factory mock transport", + ) + transport = create_transport("factory_mock", "test-endpoint") + assert isinstance(transport, MockTransport) + assert transport.endpoint == "test-endpoint"mcp_fuzzer/transport/factory.py (3)
40-46: Redundant custom-scheme rescue.
urlparsealready captures unknown schemes; this manual extraction is unnecessary.Apply:
- # Handle custom schemes that urlparse doesn't recognize - if not scheme and "://" in url_or_protocol: - # Extract scheme manually for custom transports - scheme_part = url_or_protocol.split("://")[0].lower() - if custom_registry.list_transports().get(scheme_part): - scheme = scheme_part
8-10: Unused logger.
loggeris defined but not used.Apply:
-import logging - -logger = logging.getLogger(__name__)Or add debug logs where helpful.
94-98: Deterministic error message.Sort custom names for stable ordering (helps tests/logs).
- f"custom: {', '.join(custom_registry.list_transports().keys())}" + f"custom: {', '.join(sorted(custom_registry.list_transports().keys()))}"mcp_fuzzer/config_loader.py (2)
275-313: Arbitrary import from YAML is risky; consider a plugin allowlist.Loading any module/class from config can execute untrusted code.
Options:
- Prefer entry points (importlib.metadata.entry_points) with a dedicated group, e.g., "mcp_fuzzer.transports".
- Maintain an allowlist of module prefixes.
- Gate with an explicit opt-in flag or environment variable.
Add clear docs about the trust boundary.
136-138: Escalate severity on config load failures.
logger.warningmay hide real problems. Considerlogger.errorto surface misconfigurations.- except Exception as e: - logger.warning(f"Error loading configuration file: {str(e)}") + except Exception as e: + logger.error(f"Error loading configuration file: {str(e)}") return Falsetests/integration/test_custom_transport.py (2)
188-201: Restore patched method after test to avoid side effects
transport.send_requestis monkey-patched but never restored. Restore it in atry/finally.- transport.send_request = mock_tools_request - tools = await transport.get_tools() - assert tools == [{"name": "integration_tool"}] + transport.send_request = mock_tools_request + try: + tools = await transport.get_tools() + assert tools == [{"name": "integration_tool"}] + finally: + transport.send_request = original_send_request
150-174: Minor: ensure async test runs across environmentsIf pytest-asyncio is not auto-configured, add
@pytest.mark.asyncioto async tests for portability.examples/custom_websocket_transport.py (2)
27-31: Don’t print on import failure; raise a clear ImportError insteadAvoid stdout noise in libraries.
-try: - import websockets -except ImportError: - print("websockets package is required. Install with: pip install websockets") - raise +try: + import websockets +except ImportError as e: + raise ImportError( + "websockets package is required. Install with: pip install websockets" + ) from e
339-342: Clarify usage hintAfter normalization, both of these work; show both.
- print("Usage in MCP Fuzzer:") - print(" transport = create_transport('websocket://localhost:8080/mcp')") + print("Usage in MCP Fuzzer:") + print(" transport = create_transport('websocket://localhost:8080/mcp') # custom scheme") + print(" transport = create_transport('ws://localhost:8080/mcp') # direct ws URL")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (11)
docs/custom-transports.md(1 hunks)examples/README-custom-transports.md(1 hunks)examples/config/custom-transport-config.yaml(1 hunks)examples/custom_websocket_transport.py(1 hunks)mcp_fuzzer/config_loader.py(6 hunks)mcp_fuzzer/transport/__init__.py(1 hunks)mcp_fuzzer/transport/custom.py(1 hunks)mcp_fuzzer/transport/factory.py(3 hunks)pyproject.toml(1 hunks)tests/integration/test_custom_transport.py(1 hunks)tests/unit/transport/test_custom.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
tests/unit/transport/test_custom.py (3)
mcp_fuzzer/transport/base.py (3)
TransportProtocol(13-101)stream_request(38-50)get_tools(59-80)mcp_fuzzer/transport/custom.py (7)
CustomTransportRegistry(15-142)create_transport(113-142)list_transports(105-111)register(21-56)unregister(58-71)get_transport_class(73-87)get_transport_info(89-103)mcp_fuzzer/transport/factory.py (1)
create_transport(13-98)
mcp_fuzzer/transport/__init__.py (3)
mcp_fuzzer/transport/streamable_http.py (1)
StreamableHTTPTransport(41-417)mcp_fuzzer/transport/custom.py (5)
create_transport(113-142)CustomTransportRegistry(15-142)register_custom_transport(149-167)create_custom_transport(170-181)list_custom_transports(184-190)mcp_fuzzer/transport/factory.py (1)
create_transport(13-98)
tests/integration/test_custom_transport.py (5)
mcp_fuzzer/config_loader.py (2)
apply_config_file(104-138)load_custom_transports(275-312)mcp_fuzzer/transport/custom.py (3)
create_transport(113-142)register_custom_transport(149-167)list_custom_transports(184-190)mcp_fuzzer/transport/factory.py (1)
create_transport(13-98)mcp_fuzzer/transport/base.py (3)
TransportProtocol(13-101)stream_request(38-50)get_tools(59-80)mcp_fuzzer/client/tool_client.py (1)
ToolClient(23-487)
mcp_fuzzer/config_loader.py (2)
mcp_fuzzer/transport/custom.py (1)
register_custom_transport(149-167)mcp_fuzzer/exceptions.py (1)
ConfigFileError(105-108)
examples/custom_websocket_transport.py (3)
mcp_fuzzer/transport/base.py (2)
TransportProtocol(13-101)get_tools(59-80)mcp_fuzzer/exceptions.py (1)
ConnectionError(17-20)mcp_fuzzer/transport/custom.py (1)
register_custom_transport(149-167)
mcp_fuzzer/transport/factory.py (1)
mcp_fuzzer/transport/custom.py (2)
list_transports(105-111)create_transport(113-142)
mcp_fuzzer/transport/custom.py (2)
mcp_fuzzer/transport/base.py (1)
TransportProtocol(13-101)mcp_fuzzer/transport/factory.py (1)
create_transport(13-98)
🪛 LanguageTool
examples/README-custom-transports.md
[grammar] ~69-~69: There might be a mistake here.
Context: ...t demonstrates: - Connection management - JSON-RPC request/response handling - Err...
(QB_NEW_EN)
[grammar] ~70-~70: There might be a mistake here.
Context: ...ent - JSON-RPC request/response handling - Error handling and timeouts - Streaming ...
(QB_NEW_EN)
[grammar] ~71-~71: There might be a mistake here.
Context: ...e handling - Error handling and timeouts - Streaming support - Configuration schema...
(QB_NEW_EN)
[grammar] ~72-~72: There might be a mistake here.
Context: ...andling and timeouts - Streaming support - Configuration schema definition - Regist...
(QB_NEW_EN)
[grammar] ~73-~73: There might be a mistake here.
Context: ...upport - Configuration schema definition - Registration with MCP Fuzzer ### Runnin...
(QB_NEW_EN)
[grammar] ~131-~131: There might be a mistake here.
Context: ...rors**: Ensure your transport module is in Python path 2. *Registration Failures...
(QB_NEW_EN)
[grammar] ~132-~132: There might be a mistake here.
Context: ...ures**: Verify your class inherits from TransportProtocol 3. Connection Issues: Check endpoints and...
(QB_NEW_EN)
[grammar] ~158-~158: There might be a mistake here.
Context: ...ts: 1. Check the main documentation at docs/custom-transports.md 2. Review the WebSocket example for impleme...
(QB_NEW_EN)
docs/custom-transports.md
[grammar] ~13-~13: There might be a mistake here.
Context: ...anaging custom transport implementations 2. Factory Integration: Automatic discove...
(QB_NEW_EN)
[grammar] ~14-~14: There might be a mistake here.
Context: ...y and instantiation of custom transports 3. Configuration Support: Declarative con...
(QB_NEW_EN)
[grammar] ~15-~15: There might be a mistake here.
Context: ...figuration of custom transports via YAML 4. Type Safety: Full type checking and va...
(QB_NEW_EN)
[grammar] ~334-~334: There might be a mistake here.
Context: ...that your transport class inherits from TransportProtocol 3. Connection Issues: Verify endpoint URL...
(QB_NEW_EN)
🔇 Additional comments (5)
pyproject.toml (1)
79-79: pytest-asyncio auto-mode: good addition.This aligns with the new async-heavy tests. Extras already require pytest-asyncio>=0.21.0, so compatibility looks fine.
mcp_fuzzer/transport/__init__.py (1)
5-12: Public API surface looks good.Re-exports are coherent; StreamableHTTPTransport and custom-transport helpers are properly exposed.
Also applies to: 19-25
docs/custom-transports.md (1)
115-123: Doc example likely broken until factory two-arg custom fix lands.Both “Via URL Scheme” and “Via Configuration (protocol+endpoint)” rely on create_transport. The two-arg variant for custom protocols will error without the proposed factory patch.
Also applies to: 133-139
tests/unit/transport/test_custom.py (1)
248-287: Protocol compliance test looks solid.tests/integration/test_custom_transport.py (1)
126-139: Assert ConfigFileError in invalid transport test
Replace the genericExceptionwithConfigFileErrorand import it from the loader module.- with pytest.raises(Exception): # Should raise ConfigFileError - load_custom_transports(config_data) + from mcp_fuzzer.config_loader import ConfigFileError + with pytest.raises(ConfigFileError): + load_custom_transports(config_data)
| async def send_request( | ||
| self, method: str, params: Optional[Dict[str, Any]] = None | ||
| ) -> Dict[str, Any]: | ||
| """Send a JSON-RPC request and return the response.""" | ||
| # Implement request sending logic | ||
| # Must return a JSON-RPC response dictionary | ||
| pass | ||
|
|
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Contract mismatch: docs say “return response dict” but example returns only result.
Please standardize the contract for TransportProtocol.send_request across docs and examples. If the contract is “return result”, adjust the text; if it’s “return full JSON-RPC response”, adjust the example.
Option A (return result) — change prose:
- """Send a JSON-RPC request and return the response."""
- # Implement request sending logic
- # Must return a JSON-RPC response dictionary
+ """Send a JSON-RPC request and return the result object."""
+ # Implement request sending logic
+ # Must return the JSON-RPC result (not the envelope)Option B (return full response) — change example:
- if "error" in response:
+ if "error" in response:
raise Exception(f"Server error: {response['error']}")
-
- return response.get("result", {})
+ return responseAlso applies to: 176-205
🏁 Script executed:
#!/usr/bin/env bash
# Print relevant sections to verify contract in docs/custom-transports.md
sed -n '44,51p;176,205p' docs/custom-transports.mdLength of output: 1192
Align send_request return contract in docs and examples: Docstring states “return a JSON-RPC response dictionary” but the WebSocket example returns only response.get("result", {}). Choose one:
- Update the docstring to specify returning the result object (not the envelope), or
- Modify the example to
return responseinstead ofresponse.get("result", {}).
Apply the chosen change consistently in both snippets (lines 44–51 and 176–205 of docs/custom-transports.md).
🤖 Prompt for AI Agents
In docs/custom-transports.md around lines 44–51 (and also update the matching
example at lines 176–205), the docstring currently says "return a JSON-RPC
response dictionary" but the WebSocket example returns only the result object;
change the docstring to state that send_request should return the JSON-RPC
result object (i.e., the "result" field, not the full envelope), and ensure both
code snippets and any explanatory text consistently return and reference the
result object (or, alternatively, if you prefer the full envelope, change both
examples to return `response` instead) — apply this change to both snippets so
the documentation and examples match.
mcp_fuzzer/transport/custom.py
Outdated
| def register( | ||
| self, | ||
| name: str, | ||
| transport_class: Type[TransportProtocol], | ||
| description: str = "", | ||
| config_schema: Optional[Dict[str, Any]] = None, | ||
| factory_function: Optional[Callable] = None, | ||
| ) -> None: | ||
| """Register a custom transport. | ||
| Args: | ||
| name: Unique name for the transport | ||
| transport_class: The transport class to register | ||
| description: Human-readable description | ||
| config_schema: JSON schema for transport configuration | ||
| factory_function: Optional factory function to create transport instances | ||
| Raises: | ||
| ValueError: If transport name is already registered | ||
| """ | ||
| if name in self._transports: | ||
| raise ValueError(f"Transport '{name}' is already registered") | ||
|
|
||
| if not issubclass(transport_class, TransportProtocol): | ||
| raise ValueError( | ||
| f"Transport class {transport_class} must inherit from TransportProtocol" | ||
| ) | ||
|
|
||
| self._transports[name] = { | ||
| "class": transport_class, | ||
| "description": description, | ||
| "config_schema": config_schema, | ||
| "factory": factory_function, | ||
| } | ||
|
|
||
| logger.info(f"Registered custom transport: {name}") |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Normalize transport names to lowercase to avoid duplicate registrations and lookup mismatches
Factory lowercases schemes; the registry should too.
- def register(
+ def register(
self,
name: str,
transport_class: Type[TransportProtocol],
description: str = "",
config_schema: Optional[Dict[str, Any]] = None,
factory_function: Optional[Callable] = None,
) -> None:
@@
- if name in self._transports:
+ key = name.strip().lower()
+ if key in self._transports:
raise ValueError(f"Transport '{name}' is already registered")
if not issubclass(transport_class, TransportProtocol):
raise ValueError(
f"Transport class {transport_class} must inherit from TransportProtocol"
)
- self._transports[name] = {
+ self._transports[key] = {
"class": transport_class,
"description": description,
"config_schema": config_schema,
"factory": factory_function,
}
- logger.info(f"Registered custom transport: {name}")
+ logger.info(f"Registered custom transport: {key}")Apply the same normalization in unregister, get_transport_class, and get_transport_info:
- def unregister(self, name: str) -> None:
+ def unregister(self, name: str) -> None:
@@
- if name not in self._transports:
+ key = name.strip().lower()
+ if key not in self._transports:
raise KeyError(f"Transport '{name}' is not registered")
- del self._transports[name]
- logger.info(f"Unregistered custom transport: {name}")
+ del self._transports[key]
+ logger.info(f"Unregistered custom transport: {key}")- def get_transport_class(self, name: str) -> Type[TransportProtocol]:
+ def get_transport_class(self, name: str) -> Type[TransportProtocol]:
@@
- if name not in self._transports:
+ key = name.strip().lower()
+ if key not in self._transports:
raise KeyError(f"Transport '{name}' is not registered")
- return self._transports[name]["class"]
+ return self._transports[key]["class"]- def get_transport_info(self, name: str) -> Dict[str, Any]:
+ def get_transport_info(self, name: str) -> Dict[str, Any]:
@@
- if name not in self._transports:
+ key = name.strip().lower()
+ if key not in self._transports:
raise KeyError(f"Transport '{name}' is not registered")
- return self._transports[name].copy()
+ return self._transports[key].copy()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def register( | |
| self, | |
| name: str, | |
| transport_class: Type[TransportProtocol], | |
| description: str = "", | |
| config_schema: Optional[Dict[str, Any]] = None, | |
| factory_function: Optional[Callable] = None, | |
| ) -> None: | |
| """Register a custom transport. | |
| Args: | |
| name: Unique name for the transport | |
| transport_class: The transport class to register | |
| description: Human-readable description | |
| config_schema: JSON schema for transport configuration | |
| factory_function: Optional factory function to create transport instances | |
| Raises: | |
| ValueError: If transport name is already registered | |
| """ | |
| if name in self._transports: | |
| raise ValueError(f"Transport '{name}' is already registered") | |
| if not issubclass(transport_class, TransportProtocol): | |
| raise ValueError( | |
| f"Transport class {transport_class} must inherit from TransportProtocol" | |
| ) | |
| self._transports[name] = { | |
| "class": transport_class, | |
| "description": description, | |
| "config_schema": config_schema, | |
| "factory": factory_function, | |
| } | |
| logger.info(f"Registered custom transport: {name}") | |
| def register( | |
| self, | |
| name: str, | |
| transport_class: Type[TransportProtocol], | |
| description: str = "", | |
| config_schema: Optional[Dict[str, Any]] = None, | |
| factory_function: Optional[Callable] = None, | |
| ) -> None: | |
| """Register a custom transport. | |
| Args: | |
| name: Unique name for the transport | |
| transport_class: The transport class to register | |
| description: Human-readable description | |
| config_schema: JSON schema for transport configuration | |
| factory_function: Optional factory function to create transport instances | |
| Raises: | |
| ValueError: If transport name is already registered | |
| """ | |
| key = name.strip().lower() | |
| if key in self._transports: | |
| raise ValueError(f"Transport '{name}' is already registered") | |
| if not issubclass(transport_class, TransportProtocol): | |
| raise ValueError( | |
| f"Transport class {transport_class} must inherit from TransportProtocol" | |
| ) | |
| self._transports[key] = { | |
| "class": transport_class, | |
| "description": description, | |
| "config_schema": config_schema, | |
| "factory": factory_function, | |
| } | |
| logger.info(f"Registered custom transport: {key}") | |
| def unregister(self, name: str) -> None: | |
| key = name.strip().lower() | |
| if key not in self._transports: | |
| raise KeyError(f"Transport '{name}' is not registered") | |
| del self._transports[key] | |
| logger.info(f"Unregistered custom transport: {key}") | |
| def get_transport_class(self, name: str) -> Type[TransportProtocol]: | |
| key = name.strip().lower() | |
| if key not in self._transports: | |
| raise KeyError(f"Transport '{name}' is not registered") | |
| return self._transports[key]["class"] | |
| def get_transport_info(self, name: str) -> Dict[str, Any]: | |
| key = name.strip().lower() | |
| if key not in self._transports: | |
| raise KeyError(f"Transport '{name}' is not registered") | |
| return self._transports[key].copy() |
🤖 Prompt for AI Agents
In mcp_fuzzer/transport/custom.py around lines 21 to 56, transport names are
stored as-is which can conflict with scheme lowercasing; normalize the name to
lowercase on registration (use name = name.lower()) before checking/assigning,
and make the same change in unregister, get_transport_class, and
get_transport_info so every lookup uses the lowercased key; ensure validation
and error messages reference the normalized name and keep logging using the
normalized value.
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (1)
99-146: Always use method='initialize' for InitializeRequest (avoid flaky handshakes)Randomizing the method can make initialization fail nondeterministically. Keep it valid in the “realistic” phase.
- # Use realistic method names - method_options = [ - "initialize", # Correct method - method_names().example(), - method_names().example(), - ] - return { "jsonrpc": "2.0", - "id": random.choice(id_options), - "method": random.choice(method_options), + "id": random.choice(id_options), + "method": "initialize", "params": {Optionally ensure id is not None for requests:
- "id": random.choice(id_options), + "id": random.randint(1, 1_000_000),mcp_fuzzer/transport/base.py (1)
13-28: Normalize non-dict results in stdio and sse transports
In mcp_fuzzer/transport/stdio.py and transport/sse.py, both send_request implementations doreturn response.get("result", response) # and return data.get("result", data)but if the JSON-RPC “result” is a list or scalar, they return a non-dict value, contradicting the Dict[str, Any] return type and HTTP’s fallback. Change to wrap non-dict results, for example:
result = response.get("result", response) return result if isinstance(result, dict) else {"result": result}Apply at stdio.py (line ~196&211) and sse.py (line ~56&72).
mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py (1)
229-247: Batch responses are mis-validated (collated dict goes down single-response path).server_response is a collated dict for batches, but the code only treats lists as batch; collated dicts fall into single-envelope invariants. Result: missed batch validation.
Apply:
- if isinstance(server_response, list): + # Batch: either a raw list of responses or a collated mapping {id: response} + if isinstance(server_response, list) or ( + isinstance(server_response, dict) + and "jsonrpc" not in server_response + ): try: # Handle batch responses with timeout to prevent hanging batch = await asyncio.wait_for( verify_batch_responses(server_response), timeout=self.BATCH_VALIDATION_TIMEOUT, )
♻️ Duplicate comments (4)
docs/custom-transports.md (1)
44-51: Contract alignment for send_request looks goodDoc now clearly states send_request returns the JSON-RPC result object (not the envelope). This resolves the earlier mismatch with the example.
examples/README-custom-transports.md (1)
82-86: Registration-in-process guidance fixedImport-based registration in the same process is now correctly documented.
mcp_fuzzer/config_loader.py (2)
181-184: Resolved: protocol enum includes https.This addresses prior feedback to add "https" to the enum.
229-256: Harden schema for custom_transports: forbid unknown keys (and consider factory support).Add additionalProperties: false at both the custom_transports object and each transport entry. Optional: support a dotted-path "factory" to register a factory_function from config.
Apply:
"custom_transports": { "type": "object", "description": "Configuration for custom transport mechanisms", "patternProperties": { "^[a-zA-Z][a-zA-Z0-9_]*$": { "type": "object", "properties": { "module": { "type": "string", "description": "Python module containing transport" }, "class": { "type": "string", "description": "Transport class name", }, "description": { "type": "string", "description": "Human-readable description", }, + "factory": { + "type": "string", + "description": "Dotted path to a factory function (e.g., package.module.create_transport)" + }, "config_schema": { "type": "object", "description": "JSON schema for transport config", } }, + "additionalProperties": false, "required": ["module", "class"], } }, + "additionalProperties": false },
🧹 Nitpick comments (14)
docs/custom-transports.md (2)
13-16: Minor punctuation for list itemsAdd trailing periods to improve readability.
-1. **Transport Registry**: A centralized registry for managing custom transport implementations -2. **Factory Integration**: Automatic discovery and instantiation of custom transports -3. **Configuration Support**: Declarative configuration of custom transports via YAML -4. **Type Safety**: Full type checking and validation for custom transport implementations +1. **Transport Registry**: A centralized registry for managing custom transport implementations. +2. **Factory Integration**: Automatic discovery and instantiation of custom transports. +3. **Configuration Support**: Declarative configuration of custom transports via YAML. +4. **Type Safety**: Full type checking and validation for custom transport implementations.
257-271: Tighten example schemaConsider prohibiting unknown keys to catch config typos.
config_schema = { "type": "object", "properties": { "url": {"type": "string"}, "timeout": {"type": "number", "default": 30.0}, "retries": {"type": "integer", "default": 3}, "headers": { "type": "object", "additionalProperties": {"type": "string"} } - }, + }, + "additionalProperties": False, "required": ["url"] }examples/README-custom-transports.md (1)
20-43: Clarify Quick Start return contractsAdd one line so new implementers match the documented contracts (result for send_request; envelope for send_raw).
class MyCustomTransport(TransportProtocol): def __init__(self, endpoint: str, **kwargs): # Initialize your transport pass - async def send_request(self, method: str, params=None): - # Implement JSON-RPC request sending + async def send_request(self, method: str, params=None): + # Implement JSON-RPC request sending; return the JSON-RPC result object (dict) pass async def send_raw(self, payload): - # Implement raw payload sending + # Implement raw payload sending; return the full JSON-RPC response envelope (must include 'id') passmcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (2)
149-166: Avoid emitting null cursorsPrefer omitting the cursor key rather than sending null.
- "params": { - "cursor": random.choice(cursor_options) if random.random() < 0.3 else None, - }, + "params": ( + {"cursor": random.choice(cursor_options)} + if random.random() < 0.3 + else {} + ),Apply the same change to prompts/list.
Also applies to: 227-243
12-13: Avoid runtime use of Hypothesis .example()
Hypothesis is declared in the project’s main dependencies, so every runtime invocation ofstrategies.example()will import and execute fuzzing logic—this can slow startup and yield non-deterministic results in production. Refactor to use a dedicated deterministic generator or relocate this fuzz-driven logic into tests.mcp_fuzzer/transport/base.py (1)
133-171: Optional: relax collate to tolerate id-less responses (defensive)With the above fix, this becomes less critical, but you can still fall back to positional mapping when ids are missing.
I can provide a safe fallback mapping (by request order for id-less responses) if desired.
mcp_fuzzer/config_loader.py (2)
275-313: Config-driven custom transport loading: add optional factory support and clearer errors.Support a dotted-path factory in config; pre-resolve it and pass to register_custom_transport. Improves parity with registry API and docs.
Apply:
def load_custom_transports(config_data: Dict[str, Any]) -> None: @@ for transport_name, transport_config in custom_transports.items(): try: # Import the module module_path = transport_config["module"] class_name = transport_config["class"] module = importlib.import_module(module_path) transport_class = getattr(module, class_name) # Register the transport description = transport_config.get("description", "") config_schema = transport_config.get("config_schema") + factory_fn = None + factory_path = transport_config.get("factory") + if factory_path: + try: + mod_path, attr = factory_path.rsplit(".", 1) + except ValueError as ve: + raise ConfigFileError( + f"Invalid factory path '{factory_path}'; expected 'module.attr'" + ) from ve + fmod = importlib.import_module(mod_path) + factory_fn = getattr(fmod, attr) register_custom_transport( name=transport_name, transport_class=transport_class, description=description, config_schema=config_schema, + factory_function=factory_fn, )
88-99: Specify UTF-8 when reading YAML.Prevents locale-dependent decoding issues.
Apply:
- with open(file_path, "r") as f: + with open(file_path, "r", encoding="utf-8") as f:mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py (2)
361-368: Widen type of server_response to include collated mappings.Aligns annotations with actual values from _send_fuzzed_request (dict mapping for batch).
Apply:
- server_response: Optional[Union[Dict[str, Any], List[Any]]], + server_response: Optional[Union[Dict[str, Any], List[Any], Dict[Any, Dict[str, Any]]]],
573-602: Match _create_batch_fuzz_result annotation to possible collated mapping.Batch path may pass a collated dict; reflect in the signature.
Apply:
- server_response: Optional[Union[Dict[str, Any], List[Any]]], + server_response: Optional[Union[Dict[str, Any], List[Any], Dict[Any, Dict[str, Any]]]],Also consider widening FuzzDataResult.server_response and fuzz_data types in mcp_fuzzer/types.py to accept batch variants (Union[List[...], Dict[...]]). I can open a follow-up PR if helpful.
mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (4)
82-98: Strengthen typing and avoid name ambiguityReturn type is a callable or None; also clarify phase literal.
-from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Callable @@ - def get_protocol_fuzzer_method(protocol_type: str, phase: str = "aggressive"): + def get_protocol_fuzzer_method( + protocol_type: str, + phase: str = "aggressive", + ) -> Optional[Callable[[], Dict[str, Any]]]:
103-111: Use Optional[List[str]] for None default and validate boundsMinor typing fix and guard for min/max inversion.
- def generate_batch_request( - protocol_types: List[str] = None, + def generate_batch_request( + protocol_types: Optional[List[str]] = None, phase: str = "aggressive", min_batch_size: int = 2, max_batch_size: int = 5, include_notifications: bool = True, ) -> List[Dict[str, Any]]: @@ - batch_size = random.randint(min_batch_size, max_batch_size) + if min_batch_size > max_batch_size: + min_batch_size, max_batch_size = max_batch_size, min_batch_size + batch_size = random.randint(min_batch_size, max_batch_size)
216-224: Ensure out-of-order batch actually has IDs to reorderDisable notifications for this helper so every entry gets an ID before reassignment.
- batch = ProtocolStrategies.generate_batch_request( - protocol_types, max_batch_size=3 - ) + batch = ProtocolStrategies.generate_batch_request( + protocol_types, max_batch_size=3, include_notifications=False + )
159-171: Handle Infinity and non-scalar IDs in strict JSON mode
Defaultjson.dumps(used throughout the codebase) will emitInfinityforfloat("inf")and serialize objects as nested JSON—both of which strict JSON encoders or downstream consumers may reject. In “fuzz” mode these edge cases are fine, but wrap or stringify them (or passallow_nan=False/ filter out non-scalar IDs) when operating in “strict” mode.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (12)
docs/custom-transports.md(1 hunks)examples/README-custom-transports.md(1 hunks)examples/custom_websocket_transport.py(1 hunks)mcp_fuzzer/config_loader.py(6 hunks)mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py(5 hunks)mcp_fuzzer/fuzz_engine/strategy/realistic/__init__.py(2 hunks)mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py(1 hunks)mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py(2 hunks)mcp_fuzzer/transport/base.py(1 hunks)mcp_fuzzer/transport/custom.py(1 hunks)mcp_fuzzer/transport/factory.py(3 hunks)tests/integration/test_custom_transport.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- mcp_fuzzer/transport/custom.py
- tests/integration/test_custom_transport.py
- mcp_fuzzer/transport/factory.py
🧰 Additional context used
🧬 Code graph analysis (6)
mcp_fuzzer/transport/base.py (5)
examples/custom_websocket_transport.py (1)
send_raw(178-208)mcp_fuzzer/transport/sse.py (1)
send_raw(27-75)mcp_fuzzer/transport/streamable_http.py (1)
send_raw(173-278)mcp_fuzzer/transport/http.py (1)
send_raw(125-159)mcp_fuzzer/transport/stdio.py (1)
send_raw(198-211)
mcp_fuzzer/fuzz_engine/strategy/realistic/__init__.py (1)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (9)
fuzz_list_resources_request_realistic(149-165)fuzz_read_resource_request_realistic(168-186)fuzz_subscribe_request_realistic(189-205)fuzz_unsubscribe_request_realistic(208-224)fuzz_list_prompts_request_realistic(227-242)fuzz_get_prompt_request_realistic(245-270)fuzz_list_roots_request_realistic(273-280)fuzz_set_level_request_realistic(283-294)fuzz_complete_request_realistic(297-320)
mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (3)
mcp_fuzzer/fuzz_engine/strategy/realistic/tool_strategy.py (1)
fuzz_tool_arguments_realistic(140-270)mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (10)
fuzz_initialize_request_realistic(100-146)fuzz_list_resources_request_realistic(149-165)fuzz_read_resource_request_realistic(168-186)fuzz_subscribe_request_realistic(189-205)fuzz_unsubscribe_request_realistic(208-224)fuzz_list_prompts_request_realistic(227-242)fuzz_get_prompt_request_realistic(245-270)fuzz_list_roots_request_realistic(273-280)fuzz_set_level_request_realistic(283-294)fuzz_complete_request_realistic(297-320)mcp_fuzzer/fuzz_engine/strategy/aggressive/protocol_type_strategy.py (2)
fuzz_initialize_request_aggressive(113-237)get_protocol_fuzzer_method(455-476)
mcp_fuzzer/config_loader.py (2)
mcp_fuzzer/transport/custom.py (1)
register_custom_transport(153-171)mcp_fuzzer/exceptions.py (1)
ConfigFileError(105-108)
mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py (3)
mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (2)
get_protocol_fuzzer_method(83-100)generate_batch_request(103-182)mcp_fuzzer/transport/base.py (3)
send_batch_request(103-131)collate_batch_responses(133-170)send_raw(21-22)mcp_fuzzer/types.py (1)
FuzzDataResult(15-27)
examples/custom_websocket_transport.py (3)
mcp_fuzzer/transport/base.py (9)
TransportProtocol(13-170)connect(30-32)disconnect(34-36)send_request(15-18)send_raw(21-22)send_notification(25-28)_stream_request(53-57)get_tools(59-80)call_tool(82-101)mcp_fuzzer/exceptions.py (1)
ConnectionError(17-20)mcp_fuzzer/transport/custom.py (1)
register_custom_transport(153-171)
🪛 LanguageTool
docs/custom-transports.md
[grammar] ~13-~13: There might be a mistake here.
Context: ...anaging custom transport implementations 2. Factory Integration: Automatic discove...
(QB_NEW_EN)
[grammar] ~14-~14: There might be a mistake here.
Context: ...y and instantiation of custom transports 3. Configuration Support: Declarative con...
(QB_NEW_EN)
[grammar] ~15-~15: There might be a mistake here.
Context: ...figuration of custom transports via YAML 4. Type Safety: Full type checking and va...
(QB_NEW_EN)
[grammar] ~334-~334: There might be a mistake here.
Context: ...that your transport class inherits from TransportProtocol 3. Connection Issues: Verify endpoint URL...
(QB_NEW_EN)
examples/README-custom-transports.md
[grammar] ~69-~69: There might be a mistake here.
Context: ...t demonstrates: - Connection management - JSON-RPC request/response handling - Err...
(QB_NEW_EN)
[grammar] ~70-~70: There might be a mistake here.
Context: ...ent - JSON-RPC request/response handling - Error handling and timeouts - Streaming ...
(QB_NEW_EN)
[grammar] ~71-~71: There might be a mistake here.
Context: ...e handling - Error handling and timeouts - Streaming support - Configuration schema...
(QB_NEW_EN)
[grammar] ~72-~72: There might be a mistake here.
Context: ...andling and timeouts - Streaming support - Configuration schema definition - Regist...
(QB_NEW_EN)
[grammar] ~73-~73: There might be a mistake here.
Context: ...upport - Configuration schema definition - Registration with MCP Fuzzer ### Runnin...
(QB_NEW_EN)
[grammar] ~129-~129: There might be a mistake here.
Context: ...rors**: Ensure your transport module is in Python path 2. *Registration Failures...
(QB_NEW_EN)
[grammar] ~130-~130: There might be a mistake here.
Context: ...ures**: Verify your class inherits from TransportProtocol 3. Connection Issues: Check endpoints and...
(QB_NEW_EN)
[grammar] ~156-~156: There might be a mistake here.
Context: ...ts: 1. Check the main documentation at docs/custom-transports.md 2. Review the WebSocket example for impleme...
(QB_NEW_EN)
🔇 Additional comments (8)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (1)
297-320: Parameter naming and log levels conform to MCP spec
No changes needed—completion/completeuses the singularargument, andlogging/setLevelaccepts lowercase syslog levels (e.g. “debug”).mcp_fuzzer/fuzz_engine/strategy/realistic/__init__.py (1)
16-30: Re-exports look consistentPublic API extended cleanly; names match definitions.
Also applies to: 32-53
mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py (1)
312-360: Keep batch as collated dict or list — OK — but make intent explicit.Returning a collated mapping is fine; the above fix ensures invariants run. No code change needed here beyond the validation fix.
examples/custom_websocket_transport.py (3)
72-76: Good normalization for bare endpoints.Prepending ws:// when no scheme is present matches the registry’s behavior.
169-171: Prefer project TimeoutError if available.If mcp_fuzzer.exceptions defines TimeoutError, import and raise it for consistency (parallel to ConnectionError). Otherwise, keep built-in.
Would you like me to adjust imports to use mcp_fuzzer.exceptions.TimeoutError if it exists?
309-316: Factory mapping is solid.Handles websocket:// scheme and bare endpoints; good interoperability with the registry/factory path.
mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (2)
184-201: Initialize strategy dispatch looks goodClear phase-based resolution with an aggressive fallback.
235-244: ToolStrategies phase switch is cleanConcise and readable two-phase dispatch.
| async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]: | ||
| """Send a raw payload and return the response.""" | ||
| # Implement raw payload sending logic | ||
| pass | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Document that send_raw must return an envelope with id (batch compatibility)
Batch collation depends on id being present in responses. Make this explicit so custom transports don’t break batching.
Apply this diff in both the abstract example and the WebSocket example docstrings:
- async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]:
- """Send a raw payload and return the response."""
+ async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+ """Send a raw payload and return the full JSON-RPC response envelope.
+
+ Must include the 'id' field for requests (required for batch collation).
+ If the server returns only a result object, wrap it as:
+ {"id": <request_id>, "result": <result>}.
+ """Also applies to: 206-217
🤖 Prompt for AI Agents
In docs/custom-transports.md around lines 52-56 (and also update the docstring
at lines 206-217 and the WebSocket example docstring), the abstract send_raw
example omits the requirement that implementations return an envelope containing
an "id", which breaks batch collation; update those docstrings to explicitly
state that send_raw must return a Dict[str, Any] envelope that includes an "id"
field (and any other response payload), e.g. "Returns: a response envelope dict
containing at minimum an 'id' key for batch compatibility", so custom transport
examples clearly require returning the id in the response.
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (1)
100-139: Initialize: good to hardcode "initialize"; add missing capabilities and avoid.example()/float/null IDs.
- Capabilities: include resources/prompts/logging/completion so downstream calls aren’t gated off by missing client capabilities.
- Sampling values via
strategy.example()at runtime is slow and flaky; prefer deterministic options.- JSON-RPC ids: avoid floats and nulls for the “realistic” phase; use ints and UUID strings.
Apply:
@@ - protocol_versions = [ - protocol_version_strings().example(), - protocol_version_strings().example(), - DEFAULT_PROTOCOL_VERSION, # Current MCP version - "2024-10-01", # Another valid date format - "1.0.0", # Semantic version - ] + protocol_versions = [ + DEFAULT_PROTOCOL_VERSION, + "2025-03-01", + "2024-10-01", + "1.0.0", + "0.9.0", + ] @@ - id_options = [ - json_rpc_id_values().example(), - json_rpc_id_values().example(), - json_rpc_id_values().example(), - 1, - 2, - 3, # Simple integers - "req-001", - "req-002", # String IDs - ] + id_options = [ + "req-" + uuid4().hex, + "req-" + uuid4().hex[:8], + 1, 2, 3, + "req-001", "req-002", + ] @@ - "capabilities": { - "elicitation": {}, - "experimental": {}, - "roots": {"listChanged": True}, - "sampling": {}, - }, + "capabilities": { + "roots": {"listChanged": True}, + "resources": {}, + "prompts": {}, + "logging": {}, + "completion": {}, + "sampling": {}, + "elicitation": {}, + "experimental": {}, + },Add import:
from uuid import uuid4
♻️ Duplicate comments (2)
docs/custom-transports.md (2)
52-56: Clarify send_raw return contract to match built-ins (return result object).Docs currently say “return the response,” while built-in transports return normalized results. Align to “return result.”
Apply:
- async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]: - """Send a raw payload and return the response.""" + async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]: + """Send a raw payload and return the result object. + + Implementations should unwrap JSON-RPC envelopes: + - If the server returns {"result": ...}, return that result (wrap scalars as {"result": value}). + - If the server returns a bare result, wrap it as {"result": value}. + - Raise on "error". + """
206-217: Update WebSocket send_raw example to return normalized result (not full envelope).This keeps custom transports consistent with SSE/STDIO behavior.
- async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]: + async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]: """Send raw payload over WebSocket.""" if not self.websocket: await self.connect() await self.websocket.send(json.dumps(payload)) response_text = await asyncio.wait_for( self.websocket.recv(), timeout=self.timeout ) - return json.loads(response_text) + response = json.loads(response_text) + if "error" in response: + raise Exception(f"Server error: {response['error']}") + result = response.get("result", response) + return result if isinstance(result, dict) else {"result": result}
🧹 Nitpick comments (8)
mcp_fuzzer/transport/sse.py (1)
56-57: De-duplicate normalization via a small helper.Normalization logic is now duplicated here and in stdio. Consider factoring into a shared helper to prevent drift.
Example within this file:
class SSETransport(TransportProtocol): + @staticmethod + def _normalize_result(payload: Dict[str, Any]) -> Dict[str, Any]: + result = payload.get("result", payload) + return result if isinstance(result, dict) else {"result": result} @@ - result = data.get("result", data) - return result if isinstance(result, dict) else {"result": result} + return SSETransport._normalize_result(data) @@ - result = data.get("result", data) - return result if isinstance(result, dict) else {"result": result} + return SSETransport._normalize_result(data)Also applies to: 73-74
mcp_fuzzer/transport/stdio.py (1)
196-198: Extract a shared result normalizer to avoid duplication.Same suggestion as SSE: centralize in base or a small utility to keep semantics identical across transports.
Minimal local helper:
class StdioTransport(TransportProtocol): + @staticmethod + def _normalize_result(payload: Dict[str, Any]) -> Dict[str, Any]: + result = payload.get("result", payload) + return result if isinstance(result, dict) else {"result": result} @@ - result = response.get("result", response) - return result if isinstance(result, dict) else {"result": result} + return StdioTransport._normalize_result(response) @@ - result = response.get("result", response) - return result if isinstance(result, dict) else {"result": result} + return StdioTransport._normalize_result(response)Also applies to: 212-214
docs/custom-transports.md (1)
141-151: Document batch contract briefly.If batch APIs require envelopes with ids, note that send_batch_request should return envelopes and that send_raw/send_request return results. Prevents confusion for custom transport authors.
Would you like a short paragraph added here describing the batch hooks and expected return shapes?
mcp_fuzzer/config_loader.py (1)
147-279: Enforce schema validation at load time — after loadingconfig_datainmcp_fuzzer/config_loader.py, call the existingjsonschema_validate(config_data, <schema>)(e.g. frommcp_fuzzer.fuzz_engine.invariants) to fail fast on typos such as"config_schemaa".mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (4)
142-159: Don’t send “cursor”: null; omit the key when not present.Explicit nulls can cause avoidable 4xxs; use optional field emission.
- return { - "jsonrpc": "2.0", - "id": random.randint(1, 1000), - "method": "resources/list", - "params": { - "cursor": random.choice(cursor_options) if random.random() < 0.3 else None, - }, - } + params: Dict[str, Any] = {} + if random.random() < 0.3: + params["cursor"] = random.choice([c for c in cursor_options if c]) + return { + "jsonrpc": "2.0", + "id": random.randint(1, 1000), + "method": "resources/list", + "params": params, + }
220-236: Prompts list: omit optional “cursor” instead of sending null.Same pattern as resources/list.
- return { - "jsonrpc": "2.0", - "id": random.randint(1, 1000), - "method": "prompts/list", - "params": { - "cursor": random.choice(cursor_options) if random.random() < 0.3 else None, - }, - } + params: Dict[str, Any] = {} + if random.random() < 0.3: + params["cursor"] = random.choice([c for c in cursor_options if c]) + return { + "jsonrpc": "2.0", + "id": random.randint(1, 1000), + "method": "prompts/list", + "params": params, + }
238-264: Prompts get: omit “arguments” key when None.Cleaner and more compatible across servers.
- return { - "jsonrpc": "2.0", - "id": random.randint(1, 1000), - "method": "prompts/get", - "params": { - "name": random.choice(name_options), - "arguments": random.choice(argument_options), - }, - } + name = random.choice(name_options) + args = random.choice(argument_options) + params: Dict[str, Any] = {"name": name} + if args is not None: + params["arguments"] = args + return { + "jsonrpc": "2.0", + "id": random.randint(1, 1000), + "method": "prompts/get", + "params": params, + }
290-313: completion/complete: ensure cursor is within prefix bounds and avoid mutating shared dicts.Copy the chosen argument and clamp cursor.
- return { - "jsonrpc": "2.0", - "id": random.randint(1, 1000), - "method": "completion/complete", - "params": { - "ref": random.choice(ref_options), - "argument": random.choice(argument_options), - }, - } + ref = random.choice(ref_options) + arg = dict(random.choice(argument_options)) + if "cursor" in arg: + arg["cursor"] = min(arg["cursor"], len(arg.get("prefix", ""))) + return { + "jsonrpc": "2.0", + "id": random.randint(1, 1000), + "method": "completion/complete", + "params": {"ref": ref, "argument": arg}, + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (25)
docs/custom-transports.md(1 hunks)docs/process-management.md(10 hunks)examples/README.md(2 hunks)examples/config/custom-transport-config.yaml(1 hunks)examples/config/mcp-fuzzer.yaml(1 hunks)examples/config/mcp-fuzzer.yml(1 hunks)examples/custom_websocket_transport.py(1 hunks)mcp_fuzzer/cli/runner.py(1 hunks)mcp_fuzzer/config.py(1 hunks)mcp_fuzzer/config_loader.py(7 hunks)mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py(6 hunks)mcp_fuzzer/fuzz_engine/fuzzer/tool_fuzzer.py(9 hunks)mcp_fuzzer/fuzz_engine/runtime/__init__.py(1 hunks)mcp_fuzzer/fuzz_engine/runtime/manager.py(4 hunks)mcp_fuzzer/fuzz_engine/runtime/watchdog.py(6 hunks)mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py(2 hunks)mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py(2 hunks)mcp_fuzzer/reports/safety_reporter.py(1 hunks)mcp_fuzzer/safety_system/policy.py(3 hunks)mcp_fuzzer/safety_system/safety.py(2 hunks)mcp_fuzzer/transport/custom.py(1 hunks)mcp_fuzzer/transport/sse.py(2 hunks)mcp_fuzzer/transport/stdio.py(2 hunks)tests/integration/test_custom_transport.py(1 hunks)tests/unit/transport/test_custom.py(1 hunks)
✅ Files skipped from review due to trivial changes (12)
- docs/process-management.md
- examples/config/mcp-fuzzer.yaml
- mcp_fuzzer/cli/runner.py
- mcp_fuzzer/config.py
- mcp_fuzzer/reports/safety_reporter.py
- examples/config/mcp-fuzzer.yml
- mcp_fuzzer/safety_system/safety.py
- mcp_fuzzer/fuzz_engine/runtime/init.py
- mcp_fuzzer/fuzz_engine/runtime/manager.py
- mcp_fuzzer/fuzz_engine/runtime/watchdog.py
- mcp_fuzzer/safety_system/policy.py
- mcp_fuzzer/fuzz_engine/fuzzer/tool_fuzzer.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tests/integration/test_custom_transport.py
- examples/custom_websocket_transport.py
- mcp_fuzzer/transport/custom.py
- tests/unit/transport/test_custom.py
- mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py
- examples/config/custom-transport-config.yaml
- mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py
🧰 Additional context used
🧬 Code graph analysis (1)
mcp_fuzzer/config_loader.py (3)
mcp_fuzzer/transport/custom.py (1)
register_custom_transport(157-175)mcp_fuzzer/config.py (2)
update(101-103)get(93-95)mcp_fuzzer/exceptions.py (1)
ConfigFileError(105-108)
🪛 LanguageTool
examples/README.md
[grammar] ~251-~251: There might be a mistake here.
Context: ...ts: 1. Check the main documentation at docs/custom-transports.md 2. Review the WebSocket example for impleme...
(QB_NEW_EN)
🔇 Additional comments (8)
mcp_fuzzer/transport/sse.py (2)
56-57: Consistent result shaping — good change.Returning a dict consistently avoids downstream shape checks and aligns with stdio transport.
73-74: Same here — consistent JSON-RPC result normalization.Keeps consumers simple across SSE and STDIO.
examples/README.md (1)
155-158: No action required:create_transportis correctly re-exported inmcp_fuzzer.transport
The import in the example is valid—create_transportis defined inmcp_fuzzer/transport/factory.py(line 13) and re-exported viafrom .factory import create_transport(line 6) and included in__all__(line 20) ofmcp_fuzzer/transport/__init__.py.mcp_fuzzer/transport/stdio.py (2)
196-198: Result normalization — LGTM.Matches SSE behavior; simplifies callers that expect dicts.
212-214: Same normalization for send_raw — LGTM.Keeps return type stable for callers.
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (3)
266-274: LGTM for roots/list.Minimal and correct JSON-RPC shape.
276-288: LGTM for logging/setLevel.Level set is constrained to expected values.
201-218: Ignore subscription-id correlation: MCPresources/unsubscriberequests take only auriparameter inparams(no subscription id/handle required) (modelcontextprotocol.io)Likely an incorrect or invalid review comment.
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (2)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (2)
161-180: Block SSRF and sensitive file reads in resources/readIncludes http/https and sensitive paths (/etc/hosts, /var/log). This risks outbound calls and environment leakage during fuzzing; contradicts “Preserve safety” in #36.
Apply:
- uri_options = [ - "file:///home/user/documents/readme.txt", - "file:///etc/hosts", - "file:///var/log/application.log", - "file:///tmp/session-data.json", - "http://localhost:8080/api/data", - "https://api.example.com/v1/users", - ] + uri_options = [ + "file:///tmp/mcp-fuzzer/readme.txt", + "file:///tmp/mcp-fuzzer/application.log", + "file:///tmp/mcp-fuzzer/session-data.json", + ]Follow-up: consider a module-level allowlist and a global “deny outbound” safety gate.
182-199: Remove http URI from subscribe; keep to local file rootsSame SSRF concern; keep subscriptions local-only.
Apply:
- uri_options = [ - "file:///home/user/documents/", - "file:///var/log/", - "file:///tmp/notifications/", - "http://localhost:8080/api/events", - ] + uri_options = [ + "file:///tmp/mcp-fuzzer/", + "file:///tmp/mcp-fuzzer/notifications/", + ]
🧹 Nitpick comments (8)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (6)
42-57: Prefer integers/strings for JSON‑RPC ids; avoid floats in “realistic” phaseMost servers accept ints/strings; floats often trigger validation quirks. Recommend dropping floats here.
Apply:
- return st.one_of( - st.none(), - st.text(min_size=1, max_size=50), - st.integers(), - st.floats(allow_nan=False, allow_infinity=False), - ) + return st.one_of( + st.none(), + st.text(min_size=1, max_size=50), + st.integers(), + )
99-140: Good: fixed Initialize method; tighten randomness and avoid strategy.example()Calling Hypothesis
example()and mixing withrandomreduces reproducibility/shrinking. Either (a) make this a proper Hypothesis strategy returning Dict, or (b) drop Hypothesis from this function and use a small vetted pool.Apply minimal change (option b):
- protocol_versions = [ - protocol_version_strings().example(), - protocol_version_strings().example(), - DEFAULT_PROTOCOL_VERSION, # Current MCP version - "2024-10-01", # Another valid date format - "1.0.0", # Semantic version - ] + protocol_versions = [ + DEFAULT_PROTOCOL_VERSION, + "2024-11-05", + "2024-10-01", + "1.0.0", + "0.9.0", + ] @@ - id_options = [ - json_rpc_id_values().example(), - json_rpc_id_values().example(), - json_rpc_id_values().example(), - 1, - 2, - 3, # Simple integers - "req-001", - "req-002", # String IDs - ] + id_options = [ + 1, 2, 3, 42, 99, + "req-001", "req-002", "init-123", + ]If you want option (a), I can propose a
st.fixed_dictionaries(...)strategy.
290-314: completion/complete: verify payload shape; also keep URIs localEnsure
ref/argumentfields match the MCP completion schema (names and nesting), and avoid non-local URIs.Apply URI safety tweak:
- {"type": "ref/resource", "uri": "file:///home/user/docs/"}, + {"type": "ref/resource", "uri": "file:///tmp/mcp-fuzzer/docs/"},If the schema differs, I can propose a corrected structure once confirmed.
9-13: Centralize URI safety; add simple helper and allowlistTo prevent regressions, create a tiny helper and reuse across functions.
Add this helper (outside the shown ranges):
# Helper to keep URIs local-only SAFE_FILE_URIS = [ "file:///tmp/mcp-fuzzer/", "file:///tmp/mcp-fuzzer/notifications/", "file:///tmp/mcp-fuzzer/readme.txt", "file:///tmp/mcp-fuzzer/application.log", "file:///tmp/mcp-fuzzer/session-data.json", "file:///tmp/mcp-fuzzer/docs/", ] def pick_safe_uri(prefix_only: bool = False) -> str: candidates = [u for u in SAFE_FILE_URIS if (u.endswith('/') if prefix_only else True)] return random.choice(candidates)Then replace the inline uri_options with
[pick_safe_uri(...), ...].Also applies to: 99-140, 142-314
99-140: Option: make “realistic” generators proper Hypothesis strategiesIf you want them shrinkable and composable with the rest of #9, convert each to
st.fixed_dictionaries(...)strategies and expose both “strategy_” and “build_” APIs.I can send a concrete rewrite for two exemplars (
initialize,prompts/get) if you want to trial it in CI.
59-97: Refine or remove themethod_namesgenerator
This function (defined at protocol_type_strategy.py:59) is only exercised by its unit test (test_realistic_strategies.py) and produces incomplete prefixes (e.g."notifications/"). Either eliminate it or constrain the prefix list to concrete, spec-compliant methods (or wire it into the fuzz engine where it’s meant to be used).mcp_fuzzer/transport/factory.py (1)
35-39: Keep error list deterministic across both error paths.Two-arg error sorts custom names; single-URL error (fixed above) should do the same for consistency and stable snapshots.
The diff above already replaces the unsorted join with a sorted one.
tests/integration/test_custom_transport.py (1)
183-192: Remove redundant reassignment of original_send_request.You assign original_send_request twice; the second assignment at Line 190 is redundant.
Apply this diff:
- original_send_request = transport.send_request - transport.send_request = mock_tools_request + transport.send_request = mock_tools_request
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py(7 hunks)mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py(2 hunks)mcp_fuzzer/transport/factory.py(4 hunks)tests/integration/test_custom_transport.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- mcp_fuzzer/fuzz_engine/fuzzer/protocol_fuzzer.py
🧰 Additional context used
🧬 Code graph analysis (2)
tests/integration/test_custom_transport.py (5)
mcp_fuzzer/config_loader.py (2)
apply_config_file(104-138)load_custom_transports(282-331)mcp_fuzzer/transport/factory.py (1)
create_transport(10-91)mcp_fuzzer/transport/custom.py (5)
create_transport(121-150)register_custom_transport(157-175)list_transports(113-119)unregister(63-77)list_custom_transports(192-198)mcp_fuzzer/transport/base.py (10)
TransportProtocol(13-170)connect(30-32)disconnect(34-36)send_request(15-18)send_raw(21-22)send_notification(25-28)_stream_request(53-57)stream_request(38-50)get_tools(59-80)call_tool(82-101)mcp_fuzzer/client/tool_client.py (1)
ToolClient(23-487)
mcp_fuzzer/transport/factory.py (1)
mcp_fuzzer/transport/custom.py (2)
create_transport(121-150)list_transports(113-119)
🪛 GitHub Actions: Tests
tests/integration/test_custom_transport.py
[error] 108-108: Test failure in tests/integration/test_custom_transport.py::TestCustomTransportConfiguration::test_config_file_custom_transport_loading: ValueError when creating transport for 'integration_test://test-endpoint' - Unsupported URL scheme: none. Supported: http, https, sse, stdio, streamablehttp, custom: integration_test.
mcp_fuzzer/transport/factory.py
[error] 87-87: ValueError: Unsupported URL scheme: none. Supported: http, https, sse, stdio, streamablehttp, custom: integration_test.
🔇 Additional comments (6)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (5)
16-39: Version strategy looks goodCovers both date-based and semver; safe day bounds are a nice touch.
142-159: LGTM: resources/list realistic paramsOptional cursor with low probability matches real usage.
220-236: LGTM: prompts/listLooks fine; optional cursor usage is reasonable.
266-274: LGTM: roots/listSpec-compliant and minimal.
276-288: Confirm allowed logging levels match MCPSet includes "warning" and "critical"; some specs use "warn"/"fatal". Please verify and normalize if needed.
If changes are required:
- level_options = ["debug", "info", "warning", "error", "critical"] + level_options = ["debug", "info", "warn", "error", "fatal"]I can map synonyms to the canonical set once confirmed.
tests/integration/test_custom_transport.py (1)
79-111: Test failure is due to underscore scheme parsing in factory; fix there, not here.The ValueError ("Unsupported URL scheme: none") originates from urlparse not accepting "_" in schemes. After the factory fallback is added, these tests should pass as-is.
Run the test after applying the factory patch; no changes required in the test code.
Also applies to: 235-237
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
mcp_fuzzer/transport/factory.py (1)
95-97: Sort custom transports in the final error for stable output.Keeps parity with the two-arg error path and avoids flaky assertions.
Apply:
- f"custom: {', '.join(custom_registry.list_transports().keys())}" + f"custom: {', '.join(sorted(custom_registry.list_transports().keys()))}"
🧹 Nitpick comments (1)
mcp_fuzzer/transport/factory.py (1)
46-52: Trim and bound manual scheme extraction.Guard against stray whitespace and ensure single split.
Apply:
- scheme_part = url_or_protocol.split("://")[0].lower() + scheme_part = url_or_protocol.split("://", 1)[0].strip().lower()
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
mcp_fuzzer/transport/factory.py(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
mcp_fuzzer/transport/factory.py (1)
mcp_fuzzer/transport/custom.py (2)
create_transport(121-150)list_transports(113-119)
🔇 Additional comments (4)
mcp_fuzzer/transport/factory.py (4)
6-6: Custom registry import looks good.Clear aliasing and no obvious cyclic import risk from the provided context.
20-24: Correct precedence for back-compat path (custom before built-ins).This enables intentional overrides while preserving fallback.
35-39: Deterministic error message (sorted custom list).Good call sorting for stable tests/logs.
53-59: Prefer custom dispatch first for single-URL path.Matches two-arg path behavior; sensible override point.
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (5)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (2)
197-203: Drop HTTP and sensitive paths from unsubscribe URIs; keep to safe tmp roots.Keep fuzzing hermetic and avoid accidental outbound requests or privileged reads.
- uri_options = [ - "file:///home/user/documents/", - "file:///var/log/", - "file:///tmp/notifications/", - "http://localhost:8080/api/events", - ] + uri_options = [ + "file:///tmp/mcp-fuzzer/", + "file:///tmp/mcp-fuzzer/notifications/", + ]
249-258: Don’t send"arguments": nullin prompts/get.Per spec, omit optional fields rather than setting null.
- return { - "jsonrpc": "2.0", - "id": random.randint(1, 1000), - "method": "prompts/get", - "params": { - "name": random.choice(name_options), - "arguments": random.choice(argument_options), - }, - } + req = { + "jsonrpc": "2.0", + "id": random.randint(1, 1000), + "method": "prompts/get", + } + params = {"name": random.choice(name_options)} + args = random.choice(argument_options) + if args is not None: + params["arguments"] = args + req["params"] = params + return reqmcp_fuzzer/transport/factory.py (3)
46-52: LGTM!The fallback logic correctly addresses the RFC 3986 compliance issue where schemes with underscores (like
integration_test://) are not recognized byurlparse. The manual extraction and registry lookup provide proper handling for non-RFC compliant custom schemes.
53-59: LGTM!The custom transport scheme checking logic correctly prioritizes custom transports over built-in schemes, maintaining backward compatibility while enabling extensibility.
94-98: LGTM!The updated error message consistently includes both built-in schemes and sorted custom transport names across both code paths, improving developer experience.
🧹 Nitpick comments (2)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (2)
110-120: Avoid Hypothesis.example()andNoneIDs in “realistic” initialize.Mixing Hypothesis
.example()with runtime random makes runs less reproducible and may yieldid=None, which many servers reject in practice. Prefer simple int/string IDs here and generate versions without Hypothesis.Apply:
- protocol_versions = [ - protocol_version_strings().example(), - protocol_version_strings().example(), - DEFAULT_PROTOCOL_VERSION, # Current MCP version - "2024-10-01", # Another valid date format - "1.0.0", # Semantic version - ] + protocol_versions = [ + DEFAULT_PROTOCOL_VERSION, # Current MCP version + "2024-11-05", # Known spec date + "2024-10-01", + "1.0.0", + f"{random.randint(2020,2030):04d}-{random.randint(1,12):02d}-{random.randint(1,28):02d}", + ] - id_options = [ - json_rpc_id_values().example(), - json_rpc_id_values().example(), - json_rpc_id_values().example(), - 1, - 2, - 3, # Simple integers - "req-001", - "req-002", # String IDs - ] + id_options = [ + 1, 2, 3, # Simple integers + "req-001", "req-002", # String IDs + random.randint(1000, 10_000), + f"req-{random.randint(1, 999):03d}", + ]
270-281: Include all RFC-5424 syslog severities in level_options.
Add the missing"notice","alert", and"emergency"entries alongside the existing values to fully align with the MCP spec.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py(2 hunks)mcp_fuzzer/transport/factory.py(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
mcp_fuzzer/transport/factory.py (1)
mcp_fuzzer/transport/custom.py (2)
create_transport(121-150)list_transports(113-119)
🔇 Additional comments (9)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (6)
125-125: Fix to explicit "initialize" method looks good.Using a fixed "initialize" method improves realism for the InitializeRequest.
141-158: List resources request looks solid.Optional cursor with a low probability keeps requests realistic.
160-176: Read resource URIs are now sandboxed.Only file:// under /tmp/mcp-fuzzer — avoids SSRF and sensitive paths.
178-193: Subscribe request is safe and realistic.Local file roots only; good alignment with prior guidance.
214-229: Prompts list realism LGTM.Optional cursor use mirrors typical pagination flows.
260-267: Roots list request looks correct.mcp_fuzzer/transport/factory.py (3)
6-6: LGTM!The import statement correctly brings in the custom registry for transport management.
20-24: LGTM!The custom transport check in the two-argument path correctly prioritizes custom transports before falling back to built-in mappings, which aligns with the expected behavior for extensibility.
35-39: LGTM!The enhanced error message now includes both built-in protocols and sorted custom transport names, providing better visibility into available options.
| def fuzz_complete_request_realistic() -> Dict[str, Any]: | ||
| """Generate realistic CompleteRequest for testing valid behavior.""" | ||
| ref_options = [ | ||
| {"type": "ref/prompt", "name": "code-review"}, | ||
| {"type": "ref/resource", "uri": "file:///tmp/mcp-fuzzer/docs/"}, | ||
| {"type": "ref/function", "name": "analyze_code"}, | ||
| ] | ||
|
|
||
| argument_options = [ | ||
| {"prefix": "import ", "suffix": ""}, | ||
| {"prefix": "def ", "suffix": "(", "cursor": 4}, | ||
| {"prefix": "class ", "suffix": ":", "cursor": 6}, | ||
| {"prefix": "SELECT ", "suffix": " FROM", "cursor": 7}, | ||
| ] | ||
|
|
||
| return { | ||
| "jsonrpc": "2.0", | ||
| "id": random.randint(1, 1000), | ||
| "method": "completion/complete", | ||
| "params": { | ||
| "ref": random.choice(ref_options), | ||
| "argument": random.choice(argument_options), | ||
| }, | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Double-check completion request schema (field names/types).
Validate that ref and argument match the current completion/complete schema.
🌐 Web query:
In MCP, what is the expected request schema for "completion/complete" (field names and allowed values for ref and argument)?
💡 Result:
Expected JSON-RPC request (params) for completion/complete
-
Envelope: JSON‑RPC 2.0 with method "completion/complete". (spec.modelcontextprotocol.io, modelcontextprotocol.io)
-
params object fields:
- ref (object) — required; reference being completed. Allowed ref.type values:
- "ref/prompt": include ref.name (string) for the prompt name. (modelcontextprotocol.io, spec.modelcontextprotocol.io)
- "ref/resource": include ref.uri (string) for the resource URI. (modelcontextprotocol.io, spec.modelcontextprotocol.io)
- argument (object) — required:
- name (string) — argument name. (modelcontextprotocol.io, spec.modelcontextprotocol.io)
- value (string) — current/partial value to complete. (modelcontextprotocol.io, spec.modelcontextprotocol.io)
- context (object) — optional:
- arguments (object) — optional mapping of previously-resolved argument names → values to provide context. (modelcontextprotocol.io, spec.modelcontextprotocol.io)
- ref (object) — required; reference being completed. Allowed ref.type values:
Example (params portion):
{
"ref": { "type": "ref/prompt", "name": "code_review" },
"argument": { "name": "language", "value": "py" },
"context": { "arguments": { "framework": "flask" } }
} (spec.modelcontextprotocol.io, modelcontextprotocol.io)
(Responses return a completion object with values[], optional total, and hasMore.) (spec.modelcontextprotocol.io)
Citations:
- 1: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
- 2: https://modelcontextprotocol.io/specification/2024-11-05/server/utilities/completion?utm_source=openai
- 3: https://modelcontextprotocol.io/specification/2024-11-05/server/utilities/completion?utm_source=openai
- 4: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
- 5: https://modelcontextprotocol.io/specification/2024-11-05/server/utilities/completion?utm_source=openai
- 6: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
- 7: https://modelcontextprotocol.io/specification/2024-11-05/server/utilities/completion?utm_source=openai
- 8: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
- 9: https://modelcontextprotocol.io/specification/2024-11-05/server/utilities/completion?utm_source=openai
- 10: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
- 11: https://modelcontextprotocol.io/specification/draft/server/utilities/completion?utm_source=openai
- 12: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
- 13: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
- 14: https://modelcontextprotocol.io/specification/draft/server/utilities/completion?utm_source=openai
- 15: https://spec.modelcontextprotocol.io/specification/2025-03-26/server/utilities/completion/?utm_source=openai
Remove unsupported ref type & align argument schema
- The
ref_optionsentry{"type": "ref/function", ...}is invalid: spec only allows"ref/prompt"and"ref/resource". Remove or replace it. argument_optionsmust produce objects with"name"(string) and"value"(string) per the completion/complete schema—drop the"prefix","suffix", and"cursor"fields.
🤖 Prompt for AI Agents
In mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py around
lines 284 to 307, the ref_options list includes an unsupported {"type":
"ref/function", ...} and argument_options uses invalid keys
("prefix","suffix","cursor"); update ref_options to only include "ref/prompt"
and "ref/resource" entries (remove or replace the function entry) and change
argument_options to produce objects with "name" (string) and "value" (string)
keys only so the returned params["argument"] conforms to the completion/complete
schema.
…mplete RFC-5424 syslog levels
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
Documentation PreviewYour documentation changes have been built successfully! Preview URL: https://Agent-Hellboy.github.io/mcp-server-fuzzer/ The documentation will be automatically deployed to the main site when this PR is merged. This comment was automatically generated by the documentation workflow. |
There was a problem hiding this comment.
Actionable comments posted: 6
♻️ Duplicate comments (6)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (5)
16-24: SSRF-safe local URIs and helper — LGTMUsing a strict file:// allowlist under /tmp and a small picker keeps fuzzing local-only and predictable. Nice follow-through on earlier SSRF guidance.
Also applies to: 27-40
198-208: Subscribe restricted to directory URIs — LGTMPrefix-only ensures you subscribe to roots/dirs, aligning with spec intent and prior SSRF guidance.
210-220: Unsubscribe shape and safe URIs — LGTM
240-269: Exclude null arguments in prompts/get — LGTMConditional inclusion of "arguments" matches the spec. Thanks for addressing this.
304-327: completion/complete params don’t match spec (invalid ref type and argument schema)Realistic phase should be spec-conformant. Remove unsupported ref type "ref/function" and emit argument objects with "name" and "value" fields.
Apply:
def fuzz_complete_request_realistic() -> Dict[str, Any]: """Generate realistic CompleteRequest for testing valid behavior.""" ref_options = [ {"type": "ref/prompt", "name": "code-review"}, {"type": "ref/resource", "uri": pick_safe_uri(prefix_only=True)}, - {"type": "ref/function", "name": "analyze_code"}, ] - argument_options = [ - {"prefix": "import ", "suffix": ""}, - {"prefix": "def ", "suffix": "(", "cursor": 4}, - {"prefix": "class ", "suffix": ":", "cursor": 6}, - {"prefix": "SELECT ", "suffix": " FROM", "cursor": 7}, - ] + argument_options = [ + {"name": "language", "value": "py"}, + {"name": "symbol", "value": "def "}, + {"name": "type", "value": "class "}, + {"name": "query", "value": "SELECT "}, + ] return { "jsonrpc": "2.0", "id": random.randint(1, 1000), "method": "completion/complete", "params": { "ref": random.choice(ref_options), "argument": random.choice(argument_options), }, }Optionally add an occasional "context": {"arguments": {...}} once ref/argument are corrected. I can wire that up if you want.
mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (1)
38-80: Direct refs in PROTOCOL_STRATEGIES fix the prior None-call pitfall — LGTMReplacing eager lambda calls with resolved callables prevents accidental invocation of None. Good.
🧹 Nitpick comments (9)
tests/integration/test_custom_transport.py (4)
82-90: Use a real JSON Schema for config_schema (optional).This block isn’t a JSON Schema; consider a minimal valid schema to prevent future confusion.
custom_transports: integration_test: module: "tests.integration.test_custom_transport" class: "IntegrationTestTransport" description: "Integration test transport" - config_schema: - timeout: 30 - retries: 3 + config_schema: + type: object + properties: + timeout: { type: number, default: 30 } + retries: { type: integer, default: 3 } + additionalProperties: false
143-151: Mark async test for pytest-asyncio.Add the marker to avoid accidental sync execution in environments without auto mode.
+ @pytest.mark.asyncio async def test_full_transport_lifecycle(self):
248-256: Mark async test for pytest-asyncio.Same here.
+ @pytest.mark.asyncio async def test_transport_with_mcp_client(self):
30-47: Stable test IDs (optional).Hard-coding "id": 1 can mask bugs in ID handling; incrementing an internal counter would be closer to real behavior. Optional for this stub.
examples/custom_websocket_transport.py (3)
82-82: Log the normalized URL.Log self.url to reflect the actual connection target.
- logger.info(f"Initialized WebSocket transport for {url}") + logger.info(f"Initialized WebSocket transport for {self.url}")
65-71: Type the websocket attribute (optional).Aids readability and tooling.
- self.websocket = None + self.websocket: Optional[websockets.WebSocketClientProtocol] = None
240-287: Consider releasing the lock after send for streams.Holding _lock while consuming the stream serializes all requests. If acceptable, keep; otherwise, send under lock, then release and read without it using an id-scoped filter.
mcp_fuzzer/config_loader.py (2)
138-139: Preserve stack traces for easier diagnosis.Use logger.exception to include traceback.
- logger.error(f"Error loading configuration file: {str(e)}") + logger.exception("Error loading configuration file")
336-340: Also log the traceback when a custom transport fails to load.Improves debuggability.
- logger.error(f"Failed to load custom transport '{transport_name}': {e}") + logger.exception(f"Failed to load custom transport '{transport_name}'")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
examples/custom_websocket_transport.py(1 hunks)mcp_fuzzer/config_loader.py(7 hunks)mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py(3 hunks)mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py(2 hunks)mcp_fuzzer/transport/base.py(1 hunks)tests/integration/test_custom_transport.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- mcp_fuzzer/transport/base.py
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-08T11:50:01.877Z
Learnt from: Agent-Hellboy
PR: Agent-Hellboy/mcp-server-fuzzer#25
File: mcp_fuzzer/strategy/tool_strategies.py:105-139
Timestamp: 2025-08-08T11:50:01.877Z
Learning: Repo Agent-Hellboy/mcp-server-fuzzer: In mcp_fuzzer/strategy/tool_strategies.py, ToolStrategies._generate_aggressive_integer can return None in the "special" branch despite an int annotation. Maintainership decision: ignore for now and handle in issue #18; do not re-flag this as blocking in interim reviews.
Applied to files:
mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py
🧬 Code graph analysis (1)
tests/integration/test_custom_transport.py (4)
mcp_fuzzer/transport/factory.py (1)
create_transport(10-98)mcp_fuzzer/transport/custom.py (5)
create_transport(121-150)register_custom_transport(157-175)list_transports(113-119)unregister(63-77)list_custom_transports(192-198)mcp_fuzzer/transport/base.py (10)
TransportProtocol(13-177)connect(30-32)disconnect(34-36)send_request(15-18)send_raw(21-22)send_notification(25-28)_stream_request(53-57)stream_request(38-50)get_tools(59-80)call_tool(82-101)mcp_fuzzer/client/tool_client.py (1)
ToolClient(23-487)
🔇 Additional comments (10)
mcp_fuzzer/config_loader.py (2)
182-185: LGTM: protocol enum includes https and streamablehttp.Schema now matches the factory’s supported protocols.
88-91: Good: YAML loader uses UTF-8 and safe_load.Solid defaults.
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (6)
167-184: ListResources realistic generator — LGTM
186-196: ReadResource uses safe URI — LGTM
222-238: ListPrompts realistic generator — LGTM
271-279: ListRoots minimal request — LGTM
281-302: logging/setLevel: accepted levels match spec
level_options already includes exactly the valid MCP values (emergency, alert, critical, error, warning, notice, info, debug), so no changes needed.
128-134: InitializeRequest defaults and capability keys are validProtocol versions, IDs, and client capability keys (
elicitation,experimental) align with the MCP schema’s allowed keys.mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (2)
190-207: Initialize request delegation — LGTM
209-231: Out-of-order batch generation — LGTMKeeps at least one duplicate and reorders IDs predictably.
| # Receive until matching id; ignore notifications and unrelated messages | ||
| while True: | ||
| response_text = await asyncio.wait_for( | ||
| self.websocket.recv(), | ||
| timeout=self.timeout | ||
| ) | ||
| response = json.loads(response_text) | ||
| if isinstance(response, dict) and response.get("id") == request_id: | ||
| logger.debug(f"Received WebSocket response for {method} (id={request_id})") | ||
| if "error" in response: | ||
| error_msg = f"Server error: {response['error']}" | ||
| logger.error(error_msg) | ||
| raise Exception(error_msg) | ||
| return response | ||
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) | ||
|
|
||
| except asyncio.TimeoutError: |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add an overall deadline while waiting for the matching response.
Avoids an indefinite loop if unrelated messages keep arriving.
- # Receive until matching id; ignore notifications and unrelated messages
- while True:
- response_text = await asyncio.wait_for(
- self.websocket.recv(),
- timeout=self.timeout
- )
+ # Receive until matching id with an overall deadline
+ loop = asyncio.get_running_loop()
+ deadline = loop.time() + self.timeout
+ while True:
+ remaining = max(0, deadline - loop.time())
+ if remaining == 0:
+ raise asyncio.TimeoutError()
+ response_text = await asyncio.wait_for(
+ self.websocket.recv(),
+ timeout=remaining
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Receive until matching id; ignore notifications and unrelated messages | |
| while True: | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), | |
| timeout=self.timeout | |
| ) | |
| response = json.loads(response_text) | |
| if isinstance(response, dict) and response.get("id") == request_id: | |
| logger.debug(f"Received WebSocket response for {method} (id={request_id})") | |
| if "error" in response: | |
| error_msg = f"Server error: {response['error']}" | |
| logger.error(error_msg) | |
| raise Exception(error_msg) | |
| return response | |
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) | |
| except asyncio.TimeoutError: | |
| # Receive until matching id with an overall deadline | |
| loop = asyncio.get_running_loop() | |
| deadline = loop.time() + self.timeout | |
| while True: | |
| remaining = max(0, deadline - loop.time()) | |
| if remaining == 0: | |
| raise asyncio.TimeoutError() | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), | |
| timeout=remaining | |
| ) | |
| response = json.loads(response_text) | |
| if isinstance(response, dict) and response.get("id") == request_id: | |
| logger.debug(f"Received WebSocket response for {method} (id={request_id})") | |
| if "error" in response: | |
| error_msg = f"Server error: {response['error']}" | |
| logger.error(error_msg) | |
| raise Exception(error_msg) | |
| return response | |
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) |
🤖 Prompt for AI Agents
In examples/custom_websocket_transport.py around lines 153 to 169, the receive
loop waits indefinitely for a matching response and can be starved by unrelated
messages; add an overall deadline by recording a start timestamp before the loop
and computing remaining time for each asyncio.wait_for call (or break/raise when
time elapsed), so each iteration passes max(0, deadline - elapsed) as the
timeout; when the overall deadline expires, log a timeout error and raise
asyncio.TimeoutError (or a descriptive exception) instead of looping forever.
| except asyncio.TimeoutError: | ||
| logger.error(f"WebSocket request timeout for method: {method}") | ||
| raise TimeoutError(f"Request timeout: {method}") |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Use project TimeoutError for consistency (if available).
Prefer mcp_fuzzer.exceptions.TimeoutError to the built-in for uniform handling.
Run to confirm:
🏁 Script executed:
#!/bin/bash
rg -nP 'class\s+TimeoutError\b|from\s+mcp_fuzzer\.exceptions\s+import\s+TimeoutError' -C2Length of output: 91
🏁 Script executed:
#!/bin/bash
# locate any exceptions modules under mcp_fuzzer
find . -type f -path '*/mcp_fuzzer/exceptions*.py'
# search for TimeoutError definitions in these modules
find . -type f -path '*/mcp_fuzzer/exceptions*.py' -exec grep -Hn 'TimeoutError' {} \;Length of output: 380
Use project’s RequestTimeoutError instead of built-in TimeoutError
In examples/custom_websocket_transport.py (lines 169–171), replace
raise TimeoutError(f"Request timeout: {method}")with
from mcp_fuzzer.exceptions import RequestTimeoutError
…
raise RequestTimeoutError(f"Request timeout: {method}")🤖 Prompt for AI Agents
In examples/custom_websocket_transport.py around lines 169 to 171, the code
catches asyncio.TimeoutError and raises the built-in TimeoutError; replace this
with raising the project's RequestTimeoutError and add the required import.
Change the raise to use RequestTimeoutError(f"Request timeout: {method}") and
add from mcp_fuzzer.exceptions import RequestTimeoutError at the module top (or
an existing import block) so the custom exception is available.
| async with self._lock: | ||
| if not self._connected: | ||
| await self.connect() | ||
|
|
||
| try: | ||
| logger.debug("Sending raw WebSocket payload") | ||
| await asyncio.wait_for( | ||
| self.websocket.send(json.dumps(payload)), | ||
| timeout=self.timeout | ||
| ) | ||
|
|
||
| # If this is a notification (no id), do not wait for a response | ||
| req_id = payload.get("id") if isinstance(payload, dict) else None | ||
| if req_id is None: | ||
| return {} | ||
|
|
||
| # Otherwise, wait for the matching response | ||
| while True: | ||
| response_text = await asyncio.wait_for( | ||
| self.websocket.recv(), | ||
| timeout=self.timeout | ||
| ) | ||
| response = json.loads(response_text) | ||
| if isinstance(response, dict) and response.get("id") == req_id: | ||
| return response | ||
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) | ||
|
|
||
| except Exception as e: |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Apply the same overall deadline logic to send_raw.
Mirror the bounded wait to prevent hangs on busy connections.
- while True:
- response_text = await asyncio.wait_for(
- self.websocket.recv(),
- timeout=self.timeout
- )
+ loop = asyncio.get_running_loop()
+ deadline = loop.time() + self.timeout
+ while True:
+ remaining = max(0, deadline - loop.time())
+ if remaining == 0:
+ raise asyncio.TimeoutError()
+ response_text = await asyncio.wait_for(
+ self.websocket.recv(),
+ timeout=remaining
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async with self._lock: | |
| if not self._connected: | |
| await self.connect() | |
| try: | |
| logger.debug("Sending raw WebSocket payload") | |
| await asyncio.wait_for( | |
| self.websocket.send(json.dumps(payload)), | |
| timeout=self.timeout | |
| ) | |
| # If this is a notification (no id), do not wait for a response | |
| req_id = payload.get("id") if isinstance(payload, dict) else None | |
| if req_id is None: | |
| return {} | |
| # Otherwise, wait for the matching response | |
| while True: | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), | |
| timeout=self.timeout | |
| ) | |
| response = json.loads(response_text) | |
| if isinstance(response, dict) and response.get("id") == req_id: | |
| return response | |
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) | |
| except Exception as e: | |
| async def send_raw(self, payload: Dict[str, Any]) -> Dict[str, Any]: | |
| async with self._lock: | |
| if not self._connected: | |
| await self.connect() | |
| try: | |
| logger.debug("Sending raw WebSocket payload") | |
| await asyncio.wait_for( | |
| self.websocket.send(json.dumps(payload)), | |
| timeout=self.timeout | |
| ) | |
| # If this is a notification (no id), do not wait for a response | |
| req_id = payload.get("id") if isinstance(payload, dict) else None | |
| if req_id is None: | |
| return {} | |
| # Otherwise, wait for the matching response with an overall deadline | |
| loop = asyncio.get_running_loop() | |
| deadline = loop.time() + self.timeout | |
| while True: | |
| remaining = max(0, deadline - loop.time()) | |
| if remaining == 0: | |
| raise asyncio.TimeoutError() | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), | |
| timeout=remaining | |
| ) | |
| response = json.loads(response_text) | |
| if isinstance(response, dict) and response.get("id") == req_id: | |
| return response | |
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) | |
| except Exception as e: | |
| # existing error handling… | |
| raise |
| # Register the transport | ||
| description = transport_config.get("description", "") | ||
| config_schema = transport_config.get("config_schema") | ||
| factory_fn = None | ||
| factory_path = transport_config.get("factory") | ||
| if factory_path: | ||
| try: | ||
| mod_path, attr = factory_path.rsplit(".", 1) | ||
| except ValueError as ve: | ||
| raise ConfigFileError( | ||
| f"Invalid factory path '{factory_path}'; expected 'module.attr'" | ||
| ) from ve | ||
| fmod = importlib.import_module(mod_path) | ||
| factory_fn = getattr(fmod, attr) | ||
| if not callable(factory_fn): | ||
| raise ConfigFileError(f"Factory '{factory_path}' is not callable") | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Validate config_schema type before registration.
Prevent bad configs from slipping through.
description = transport_config.get("description", "")
- config_schema = transport_config.get("config_schema")
+ config_schema = transport_config.get("config_schema")
+ if config_schema is not None and not isinstance(config_schema, dict):
+ raise ConfigFileError(
+ f"'config_schema' for {transport_name} must be an object"
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Register the transport | |
| description = transport_config.get("description", "") | |
| config_schema = transport_config.get("config_schema") | |
| factory_fn = None | |
| factory_path = transport_config.get("factory") | |
| if factory_path: | |
| try: | |
| mod_path, attr = factory_path.rsplit(".", 1) | |
| except ValueError as ve: | |
| raise ConfigFileError( | |
| f"Invalid factory path '{factory_path}'; expected 'module.attr'" | |
| ) from ve | |
| fmod = importlib.import_module(mod_path) | |
| factory_fn = getattr(fmod, attr) | |
| if not callable(factory_fn): | |
| raise ConfigFileError(f"Factory '{factory_path}' is not callable") | |
| # Register the transport | |
| description = transport_config.get("description", "") | |
| config_schema = transport_config.get("config_schema") | |
| if config_schema is not None and not isinstance(config_schema, dict): | |
| raise ConfigFileError( | |
| f"'config_schema' for {transport_name} must be an object" | |
| ) | |
| factory_fn = None | |
| factory_path = transport_config.get("factory") | |
| if factory_path: | |
| try: | |
| mod_path, attr = factory_path.rsplit(".", 1) | |
| except ValueError as ve: | |
| raise ConfigFileError( | |
| f"Invalid factory path '{factory_path}'; expected 'module.attr'" | |
| ) from ve | |
| fmod = importlib.import_module(mod_path) | |
| factory_fn = getattr(fmod, attr) | |
| if not callable(factory_fn): | |
| raise ConfigFileError(f"Factory '{factory_path}' is not callable") |
🤖 Prompt for AI Agents
In mcp_fuzzer/config_loader.py around lines 306 to 322, config_schema is read
from transport_config but not validated, allowing invalid types to be
registered; update the code to validate config_schema immediately after
retrieval: if config_schema is not None, ensure it is the expected type (e.g., a
dict/mapping), and if not raise ConfigFileError with a clear message including
the transport name or factory_path; do this before any registration or further
use so only correctly typed schemas are accepted.
| @staticmethod | ||
| def generate_batch_request( | ||
| protocol_types: Optional[List[str]] = None, | ||
| phase: str = "aggressive", | ||
| min_batch_size: int = 2, | ||
| max_batch_size: int = 5, | ||
| include_notifications: bool = True, | ||
| ) -> List[Dict[str, Any]]: | ||
| """ | ||
| Generate a batch of JSON-RPC requests/notifications. | ||
| Args: | ||
| protocol_types: List of protocol types to include (None for random | ||
| selection) | ||
| phase: Fuzzing phase (realistic or aggressive) | ||
| min_batch_size: Minimum number of requests in batch | ||
| max_batch_size: Maximum number of requests in batch | ||
| include_notifications: Whether to include notifications (no ID) | ||
| Returns: | ||
| List of JSON-RPC requests/notifications forming a batch | ||
| """ | ||
| if protocol_types is None: | ||
| # Default protocol types for batching | ||
| protocol_types = [ | ||
| "InitializeRequest", | ||
| "ListResourcesRequest", | ||
| "ReadResourceRequest", | ||
| "ListPromptsRequest", | ||
| "GetPromptRequest", | ||
| "ListRootsRequest", | ||
| "SetLevelRequest", | ||
| "CompleteRequest", | ||
| ] | ||
|
|
||
| if min_batch_size > max_batch_size: | ||
| min_batch_size, max_batch_size = max_batch_size, min_batch_size | ||
| batch_size = random.randint(min_batch_size, max_batch_size) | ||
| batch = [] | ||
|
|
||
| for i in range(batch_size): | ||
| # Randomly select protocol type | ||
| protocol_type = random.choice(protocol_types) | ||
|
|
||
| # Get the fuzzer method for this type and phase | ||
| fuzzer_method = ProtocolStrategies.get_protocol_fuzzer_method( | ||
| protocol_type, phase | ||
| ) | ||
| if not fuzzer_method: | ||
| continue | ||
|
|
||
| # Generate the request | ||
| request = fuzzer_method() | ||
|
|
||
| # Sometimes make it a notification (no ID) if enabled | ||
| if include_notifications and random.random() < 0.3: | ||
| request.pop("id", None) | ||
|
|
||
| # Add some ID edge cases occasionally | ||
| if "id" in request and random.random() < 0.2: | ||
| # Edge case IDs | ||
| edge_cases = [ | ||
| None, # Valid for notifications | ||
| "", # Empty string | ||
| 0, # Zero | ||
| -1, # Negative | ||
| "duplicate_id", # Same ID as another in batch | ||
| float("inf"), # Infinity | ||
| {"nested": "object"}, # Object as ID | ||
| ] | ||
| request["id"] = random.choice(edge_cases) | ||
|
|
||
| batch.append(request) | ||
|
|
||
| # Ensure we have at least one request with ID for responses | ||
| if batch and not any( | ||
| "id" in req and req.get("id") is not None for req in batch | ||
| ): | ||
| # Add an ID to the first request | ||
| if batch: | ||
| batch[0]["id"] = random.randint(1, 1000) | ||
|
|
||
| return batch | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Don’t inject invalid JSON-RPC IDs during the realistic phase
Edge-case IDs like Infinity and objects aren’t valid JSON and will break encoders/transports; null id is not a notification. Gate these to the aggressive phase only and clarify the comment.
Apply:
- # Add some ID edge cases occasionally
- if "id" in request and random.random() < 0.2:
- # Edge case IDs
+ # Add some ID edge cases occasionally (aggressive only)
+ if (
+ phase == ProtocolStrategies.AGGRESSIVE_PHASE
+ and "id" in request
+ and random.random() < 0.2
+ ):
+ # Edge case IDs (some intentionally invalid for stress)
edge_cases = [
- None, # Valid for notifications
+ None, # Null id is allowed by JSON-RPC but is NOT a notification
"", # Empty string
0, # Zero
-1, # Negative
"duplicate_id", # Same ID as another in batch
float("inf"), # Infinity
{"nested": "object"}, # Object as ID
]
request["id"] = random.choice(edge_cases)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @staticmethod | |
| def generate_batch_request( | |
| protocol_types: Optional[List[str]] = None, | |
| phase: str = "aggressive", | |
| min_batch_size: int = 2, | |
| max_batch_size: int = 5, | |
| include_notifications: bool = True, | |
| ) -> List[Dict[str, Any]]: | |
| """ | |
| Generate a batch of JSON-RPC requests/notifications. | |
| Args: | |
| protocol_types: List of protocol types to include (None for random | |
| selection) | |
| phase: Fuzzing phase (realistic or aggressive) | |
| min_batch_size: Minimum number of requests in batch | |
| max_batch_size: Maximum number of requests in batch | |
| include_notifications: Whether to include notifications (no ID) | |
| Returns: | |
| List of JSON-RPC requests/notifications forming a batch | |
| """ | |
| if protocol_types is None: | |
| # Default protocol types for batching | |
| protocol_types = [ | |
| "InitializeRequest", | |
| "ListResourcesRequest", | |
| "ReadResourceRequest", | |
| "ListPromptsRequest", | |
| "GetPromptRequest", | |
| "ListRootsRequest", | |
| "SetLevelRequest", | |
| "CompleteRequest", | |
| ] | |
| if min_batch_size > max_batch_size: | |
| min_batch_size, max_batch_size = max_batch_size, min_batch_size | |
| batch_size = random.randint(min_batch_size, max_batch_size) | |
| batch = [] | |
| for i in range(batch_size): | |
| # Randomly select protocol type | |
| protocol_type = random.choice(protocol_types) | |
| # Get the fuzzer method for this type and phase | |
| fuzzer_method = ProtocolStrategies.get_protocol_fuzzer_method( | |
| protocol_type, phase | |
| ) | |
| if not fuzzer_method: | |
| continue | |
| # Generate the request | |
| request = fuzzer_method() | |
| # Sometimes make it a notification (no ID) if enabled | |
| if include_notifications and random.random() < 0.3: | |
| request.pop("id", None) | |
| # Add some ID edge cases occasionally | |
| if "id" in request and random.random() < 0.2: | |
| # Edge case IDs | |
| edge_cases = [ | |
| None, # Valid for notifications | |
| "", # Empty string | |
| 0, # Zero | |
| -1, # Negative | |
| "duplicate_id", # Same ID as another in batch | |
| float("inf"), # Infinity | |
| {"nested": "object"}, # Object as ID | |
| ] | |
| request["id"] = random.choice(edge_cases) | |
| batch.append(request) | |
| # Ensure we have at least one request with ID for responses | |
| if batch and not any( | |
| "id" in req and req.get("id") is not None for req in batch | |
| ): | |
| # Add an ID to the first request | |
| if batch: | |
| batch[0]["id"] = random.randint(1, 1000) | |
| return batch | |
| # Sometimes make it a notification (no ID) if enabled | |
| if include_notifications and random.random() < 0.3: | |
| request.pop("id", None) | |
| # Add some ID edge cases occasionally (aggressive only) | |
| if ( | |
| phase == ProtocolStrategies.AGGRESSIVE_PHASE | |
| and "id" in request | |
| and random.random() < 0.2 | |
| ): | |
| # Edge case IDs (some intentionally invalid for stress) | |
| edge_cases = [ | |
| None, # Null id is allowed by JSON-RPC but is NOT a notification | |
| "", # Empty string | |
| 0, # Zero | |
| -1, # Negative | |
| "duplicate_id", # Same ID as another in batch | |
| float("inf"), # Infinity | |
| {"nested": "object"}, # Object as ID | |
| ] | |
| request["id"] = random.choice(edge_cases) | |
| batch.append(request) |
| await transport.connect() | ||
|
|
||
| # Test tool calling through transport | ||
| result = await transport.call_tool("test_tool", {"arg": "value"}) | ||
| assert "result" in result |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Ensure cleanup: disconnect the transport.
Prevent lingering connections in case of assertion failures.
- await transport.connect()
-
- # Test tool calling through transport
- result = await transport.call_tool("test_tool", {"arg": "value"})
- assert "result" in result
+ await transport.connect()
+ try:
+ # Test tool calling through transport
+ result = await transport.call_tool("test_tool", {"arg": "value"})
+ assert "result" in result
+ finally:
+ await transport.disconnect()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| await transport.connect() | |
| # Test tool calling through transport | |
| result = await transport.call_tool("test_tool", {"arg": "value"}) | |
| assert "result" in result | |
| await transport.connect() | |
| try: | |
| # Test tool calling through transport | |
| result = await transport.call_tool("test_tool", {"arg": "value"}) | |
| assert "result" in result | |
| finally: | |
| await transport.disconnect() |
🤖 Prompt for AI Agents
In tests/integration/test_custom_transport.py around lines 267 to 271, the test
opens a transport connection but doesn't ensure it is disconnected on failure;
wrap the connect/call/assert sequence in a try/finally (or use a context
manager) and call transport.disconnect() in the finally block so the transport
is always disconnected even if assertions fail.
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (2)
12-14: Fix relative import path (likely off-by-one).
from ....config import DEFAULT_PROTOCOL_VERSIONclimbs above the package root; ifconfig.pylives undermcp_fuzzer/config.py, this should be three dots.-from ....config import DEFAULT_PROTOCOL_VERSION +from ...config import DEFAULT_PROTOCOL_VERSION
124-164: Remove unsupportedelicitationcapability key
As of the MCP spec (2025-03-26), top-level ClientCapabilities only includeroots,sampling, andexperimental; remove theelicitationentry to align with the spec.mcp_fuzzer/transport/base.py (2)
59-77: Fix inconsistent response contract: handle both envelope and unwrapped results in get_toolsToday get_tools assumes send_request returns the unwrapped result; WebSocketTransport returns a full JSON-RPC envelope. Make get_tools robust to both shapes.
try: - response = await self.send_request("tools/list") + response = await self.send_request("tools/list") logging.debug("Raw server response: %s", response) if not isinstance(response, dict): logging.warning( "Server response is not a dictionary. Got type: %s", type(response), ) return [] - if "tools" not in response: - logging.warning( - "Server response missing 'tools' key. Keys present: %s", - list(response.keys()), - ) - return [] - tools = response["tools"] + # Accept either unwrapped {'tools': [...]} or envelope {'result': {'tools': [...]}} + tools_container = None + if "tools" in response: + tools_container = response + elif "result" in response and isinstance(response["result"], dict) and "tools" in response["result"]: + tools_container = response["result"] + if tools_container is None: + logging.warning( + "Server response missing 'tools' key. Keys present: %s", + list(response.keys()), + ) + return [] + tools = tools_container["tools"] logging.info("Found %d tools from server", len(tools)) return tools
82-101: Preserve safety metadata regardless of envelope vs unwrapped resultEnsure _meta is injected into the actual result payload if transports return envelopes.
result = await self.send_request("tools/call", params) - if safety_sanitized and isinstance(result, dict): - if "_meta" not in result: - result["_meta"] = {} - result["_meta"]["safety_sanitized"] = True - result["_meta"]["original_arguments"] = arguments - result["_meta"]["sanitized_arguments"] = sanitized_arguments + if safety_sanitized: + container = result + if isinstance(container, dict) and "result" in container and isinstance(container["result"], dict): + container = container["result"] + if isinstance(container, dict): + container.setdefault("_meta", {}) + container["_meta"]["safety_sanitized"] = True + container["_meta"]["original_arguments"] = arguments + container["_meta"]["sanitized_arguments"] = sanitized_arguments return result
♻️ Duplicate comments (5)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (2)
240-269: prompts/get omits null arguments — thanks for fixing.Only includes “arguments” when present; aligns with spec and prior feedback.
304-327: completion/complete schema mismatch: unsupported ref type and invalid argument shape.Spec allows ref.type “ref/prompt” or “ref/resource”; and argument must be {name, value} (strings). Remove “ref/function” and use name/value.
- ref_options = [ - {"type": "ref/prompt", "name": "code-review"}, - {"type": "ref/resource", "uri": pick_safe_uri(prefix_only=True)}, - {"type": "ref/function", "name": "analyze_code"}, - ] + ref_options = [ + {"type": "ref/prompt", "name": "code-review"}, + {"type": "ref/resource", "uri": pick_safe_uri(prefix_only=True)}, + ] @@ - argument_options = [ - {"prefix": "import ", "suffix": ""}, - {"prefix": "def ", "suffix": "(", "cursor": 4}, - {"prefix": "class ", "suffix": ":", "cursor": 6}, - {"prefix": "SELECT ", "suffix": " FROM", "cursor": 7}, - ] + argument_options = [ + {"name": "language", "value": "py"}, + {"name": "framework", "value": "flask"}, + {"name": "query", "value": "SEL"}, + {"name": "filename", "value": "app"}, + ]mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (1)
38-80: Mapping now avoids lambda-of-None trap — nice.Storing direct callables/None fixes the earlier crash path.
mcp_fuzzer/transport/base.py (1)
126-133: Good fix: preserve request id in batch responses for collationThe normalization and id backfill resolve the prior collation issue. Nice.
tests/integration/test_custom_transport.py (1)
72-78: Nice: using public registry API for cleanupAvoids coupling tests to internals. Good change.
Also applies to: 136-142, 205-211, 241-247
🧹 Nitpick comments (9)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (2)
27-39: Guard against empty candidate set.If all non-dir URIs are later removed,
random.choice(candidates)can raise IndexError. Add a safe fallback.- candidates = [ - u for u in SAFE_FILE_URIS if (u.endswith("/") if prefix_only else True) - ] - return random.choice(candidates) + candidates = [u for u in SAFE_FILE_URIS if (u.endswith("/") if prefix_only else True)] + if not candidates: + return "file:///tmp/mcp-fuzzer/" + return random.choice(candidates)
281-301: logging/setLevel: align levels to common set.Many servers accept a narrower set (e.g., trace/debug/info/warn/error/fatal). Consider restricting to improve “realistic” success rate.
- level_options = [ - "debug", - "info", - "notice", - "warning", - "error", - "critical", - "alert", - "emergency", - ] + level_options = [ + "trace", + "debug", + "info", + "warn", + "error", + "fatal", + ]mcp_fuzzer/transport/base.py (2)
134-137: Avoid emitting response entries for failed notifications (no id)Appending {"id": None, "error": ...} creates unmatched-ID warnings during collation. Skip appending for notifications; just log.
- except Exception as e: - logging.warning(f"Failed to send batch request: {e}") - responses.append({"error": str(e), "id": request.get("id")}) + except Exception as e: + logging.warning(f"Failed to send batch request: {e}") + req_id = request.get("id") + if req_id is not None: + responses.append({"error": str(e), "id": req_id})
159-166: Handle duplicate response ids deterministicallyIf a server sends multiple responses for the same id, keep the first and log duplicates to aid debugging.
for response in responses: response_id = response.get("id") if response_id in expected_responses: - collated[response_id] = response + if response_id in collated: + logging.warning("Duplicate response for ID %s; keeping first and discarding later one.", response_id) + else: + collated[response_id] = response else: # Unmatched response - could be error or notification response logging.warning(f"Received response with unmatched ID: {response_id}")examples/custom_websocket_transport.py (4)
82-83: Log normalized URLReport the effective URL after normalization.
- logger.info(f"Initialized WebSocket transport for {url}") + logger.info(f"Initialized WebSocket transport for {self.url}")
216-219: Map asyncio timeouts to built-in TimeoutError for consistencysend_request already maps timeouts; do the same here.
- except Exception as e: + except asyncio.TimeoutError: + logger.error("Raw WebSocket send timeout") + raise TimeoutError("Raw request timeout") + except Exception as e: logger.error(f"Raw WebSocket send failed: {e}") raise
339-343: Clarify schema: accept bare endpointsMirror constructor/factory behavior in the schema description.
- "url": { - "type": "string", - "description": "WebSocket URL (ws:// or wss://)" - }, + "url": { + "type": "string", + "description": "WebSocket URL (ws:// or wss://) or bare endpoint (host[:port]/path)" + },
146-168: Consider a background reader + dispatcher to avoid dropping notificationsCurrent design consumes and ignores out-of-band messages. A background read loop with per-id waiters (and a notification callback) prevents loss and enables concurrent requests.
Also applies to: 272-287
tests/integration/test_custom_transport.py (1)
80-90: Use “config” for defaults; reserve “config_schema” for schemasThe YAML block appears to pass defaults, not a schema. Rename for clarity and loader compatibility.
- config_schema: - timeout: 30 - retries: 3 + config: + timeout: 30 + retries: 3
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
examples/custom_websocket_transport.py(1 hunks)mcp_fuzzer/config_loader.py(7 hunks)mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py(3 hunks)mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py(2 hunks)mcp_fuzzer/transport/base.py(1 hunks)tests/integration/test_custom_transport.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- mcp_fuzzer/config_loader.py
🔇 Additional comments (11)
mcp_fuzzer/fuzz_engine/strategy/realistic/protocol_type_strategy.py (7)
16-24: Local-only URIs: good hardening.Constrained to file:// under /tmp; aligns with prior SSRF feedback.
167-184: resources/list realism LGTM.Optional cursor with reasonable formats.
186-196: resources/read uses safe URIs.Solid; stays local.
198-207: resources/subscribe directory-only URIs.Good constraint for subscriptions.
210-219: resources/unsubscribe parity with subscribe.Matches subscribe shape and safety.
222-237: prompts/list realism LGTM.Optional cursor handled.
271-278: roots/list minimal params OK.mcp_fuzzer/fuzz_engine/strategy/strategy_manager.py (3)
82-104: Fuzzer resolver fallback is sensible.Graceful fallback to aggressive resolver for unknown types.
190-207: Initialize wrapper: clean two-phase dispatch.
209-232: Out-of-order batch helper LGTM.Small batch, deterministic non-sequential IDs; good for exercising collation.
tests/integration/test_custom_transport.py (1)
143-151: No action needed—asyncio auto mode is enabled
tox.ini already setsasyncio_mode = auto, so async tests run without@pytest.mark.asyncio.
| try: | ||
| logger.debug(f"Sending WebSocket request: {method} (id={request_id})") | ||
| await asyncio.wait_for( | ||
| self.websocket.send(json.dumps(payload)), | ||
| timeout=self.timeout | ||
| ) | ||
|
|
||
| # Receive until matching id; ignore notifications and unrelated messages | ||
| while True: | ||
| response_text = await asyncio.wait_for( | ||
| self.websocket.recv(), | ||
| timeout=self.timeout | ||
| ) | ||
| response = json.loads(response_text) | ||
| if isinstance(response, dict) and response.get("id") == request_id: | ||
| logger.debug(f"Received WebSocket response for {method} (id={request_id})") | ||
| if "error" in response: | ||
| error_msg = f"Server error: {response['error']}" | ||
| logger.error(error_msg) | ||
| raise Exception(error_msg) | ||
| return response | ||
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) |
There was a problem hiding this comment.
Unwrap JSON-RPC envelopes in send_request to align with TransportProtocol usage
Base helpers (e.g., get_tools) expect unwrapped results. Return response["result"] when present.
- if "error" in response:
+ if "error" in response:
error_msg = f"Server error: {response['error']}"
logger.error(error_msg)
raise Exception(error_msg)
- return response
+ return response.get("result", response)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| logger.debug(f"Sending WebSocket request: {method} (id={request_id})") | |
| await asyncio.wait_for( | |
| self.websocket.send(json.dumps(payload)), | |
| timeout=self.timeout | |
| ) | |
| # Receive until matching id; ignore notifications and unrelated messages | |
| while True: | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), | |
| timeout=self.timeout | |
| ) | |
| response = json.loads(response_text) | |
| if isinstance(response, dict) and response.get("id") == request_id: | |
| logger.debug(f"Received WebSocket response for {method} (id={request_id})") | |
| if "error" in response: | |
| error_msg = f"Server error: {response['error']}" | |
| logger.error(error_msg) | |
| raise Exception(error_msg) | |
| return response | |
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) | |
| try: | |
| logger.debug(f"Sending WebSocket request: {method} (id={request_id})") | |
| await asyncio.wait_for( | |
| self.websocket.send(json.dumps(payload)), | |
| timeout=self.timeout | |
| ) | |
| # Receive until matching id; ignore notifications and unrelated messages | |
| while True: | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), | |
| timeout=self.timeout | |
| ) | |
| response = json.loads(response_text) | |
| if isinstance(response, dict) and response.get("id") == request_id: | |
| logger.debug(f"Received WebSocket response for {method} (id={request_id})") | |
| if "error" in response: | |
| error_msg = f"Server error: {response['error']}" | |
| logger.error(error_msg) | |
| raise Exception(error_msg) | |
| return response.get("result", response) | |
| logger.debug("Ignoring out-of-band WebSocket message: %s", response) |
🤖 Prompt for AI Agents
In examples/custom_websocket_transport.py around lines 146 to 167, the
send_request currently returns the full JSON-RPC envelope; change it to unwrap
and return response["result"] when the response is a dict with a "result" key
(while preserving the existing error handling that raises on "error" and
continuing to ignore out-of-band messages). Ensure you still check for matching
id, log receipt, raise with the server error if "error" exists, and return
response["result"] instead of the entire envelope when "result" is present so
downstream helpers receive the unwrapped value.
| @staticmethod | ||
| def generate_batch_request( | ||
| protocol_types: Optional[List[str]] = None, | ||
| phase: str = "aggressive", | ||
| min_batch_size: int = 2, | ||
| max_batch_size: int = 5, | ||
| include_notifications: bool = True, | ||
| ) -> List[Dict[str, Any]]: | ||
| """ | ||
| Generate a batch of JSON-RPC requests/notifications. | ||
| Args: | ||
| protocol_types: List of protocol types to include (None for random | ||
| selection) | ||
| phase: Fuzzing phase (realistic or aggressive) | ||
| min_batch_size: Minimum number of requests in batch | ||
| max_batch_size: Maximum number of requests in batch | ||
| include_notifications: Whether to include notifications (no ID) | ||
| Returns: | ||
| List of JSON-RPC requests/notifications forming a batch | ||
| """ | ||
| if protocol_types is None: | ||
| # Default protocol types for batching | ||
| protocol_types = [ | ||
| "InitializeRequest", | ||
| "ListResourcesRequest", | ||
| "ReadResourceRequest", | ||
| "ListPromptsRequest", | ||
| "GetPromptRequest", | ||
| "ListRootsRequest", | ||
| "SetLevelRequest", | ||
| "CompleteRequest", | ||
| ] | ||
|
|
||
| if min_batch_size > max_batch_size: | ||
| min_batch_size, max_batch_size = max_batch_size, min_batch_size | ||
| batch_size = random.randint(min_batch_size, max_batch_size) | ||
| batch = [] | ||
|
|
||
| for i in range(batch_size): | ||
| # Randomly select protocol type | ||
| protocol_type = random.choice(protocol_types) | ||
|
|
||
| # Get the fuzzer method for this type and phase | ||
| fuzzer_method = ProtocolStrategies.get_protocol_fuzzer_method( | ||
| protocol_type, phase | ||
| ) | ||
| if not fuzzer_method: | ||
| continue | ||
|
|
||
| # Generate the request | ||
| request = fuzzer_method() | ||
|
|
||
| # Sometimes make it a notification (no ID) if enabled | ||
| if include_notifications and random.random() < 0.3: | ||
| request.pop("id", None) | ||
|
|
||
| # Add some ID edge cases occasionally | ||
| if "id" in request and random.random() < 0.2: | ||
| # Edge case IDs | ||
| edge_cases = [ | ||
| None, # Valid for notifications | ||
| "", # Empty string | ||
| 0, # Zero | ||
| -1, # Negative | ||
| "duplicate_id", # Same ID as another in batch | ||
| float("inf"), # Infinity | ||
| {"nested": "object"}, # Object as ID | ||
| ] | ||
| request["id"] = random.choice(edge_cases) | ||
|
|
||
| batch.append(request) | ||
|
|
||
| # Ensure we have at least one request with ID for responses | ||
| if batch and not any( | ||
| "id" in req and req.get("id") is not None for req in batch | ||
| ): | ||
| # Add an ID to the first request | ||
| if batch: | ||
| batch[0]["id"] = random.randint(1, 1000) | ||
|
|
||
| return batch | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Gate ID edge-cases to aggressive phase; include subscribe/unsubscribe; drop Infinity.
Keep “realistic” batches valid JSON-RPC; avoid non‑JSON values like Infinity.
@@
- if protocol_types is None:
+ if protocol_types is None:
# Default protocol types for batching
protocol_types = [
"InitializeRequest",
"ListResourcesRequest",
"ReadResourceRequest",
+ "SubscribeRequest",
+ "UnsubscribeRequest",
"ListPromptsRequest",
"GetPromptRequest",
"ListRootsRequest",
"SetLevelRequest",
"CompleteRequest",
]
@@
- # Add some ID edge cases occasionally
- if "id" in request and random.random() < 0.2:
+ # Add some ID edge cases occasionally (aggressive only)
+ if (
+ phase == ProtocolStrategies.AGGRESSIVE_PHASE
+ and "id" in request
+ and random.random() < 0.2
+ ):
# Edge case IDs
edge_cases = [
- None, # Valid for notifications
+ None, # Null ID (still a request)
"", # Empty string
0, # Zero
-1, # Negative
"duplicate_id", # Same ID as another in batch
- float("inf"), # Infinity
+ (2**63 - 1), # Very large integer, still JSON
{"nested": "object"}, # Object as ID
]
request["id"] = random.choice(edge_cases)Also applies to: 129-139, 163-176
| async def test_transport_with_mcp_client(self): | ||
| """Test using custom transport with MCP client.""" | ||
| # Register custom transport | ||
| register_custom_transport( | ||
| name="client_test", transport_class=IntegrationTestTransport | ||
| ) | ||
|
|
||
| # Create transport | ||
| transport = create_transport("client_test://mcp-server") | ||
|
|
||
| # Import and create MCP client (this would normally be done) | ||
| # This is a simplified test - in real usage, you'd use the full client | ||
| from mcp_fuzzer.client.tool_client import ToolClient | ||
|
|
||
| # Test that the client can be created with custom transport | ||
| client = ToolClient(transport) | ||
| assert client.transport == transport | ||
|
|
||
| # Test basic functionality | ||
| await transport.connect() | ||
|
|
||
| # Test tool calling through transport | ||
| result = await transport.call_tool("test_tool", {"arg": "value"}) |
There was a problem hiding this comment.
Mark async test with pytest.mark.asyncio (or enable auto mode)
Same issue as above.
- async def test_transport_with_mcp_client(self):
+ @pytest.mark.asyncio
+ async def test_transport_with_mcp_client(self):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def test_transport_with_mcp_client(self): | |
| """Test using custom transport with MCP client.""" | |
| # Register custom transport | |
| register_custom_transport( | |
| name="client_test", transport_class=IntegrationTestTransport | |
| ) | |
| # Create transport | |
| transport = create_transport("client_test://mcp-server") | |
| # Import and create MCP client (this would normally be done) | |
| # This is a simplified test - in real usage, you'd use the full client | |
| from mcp_fuzzer.client.tool_client import ToolClient | |
| # Test that the client can be created with custom transport | |
| client = ToolClient(transport) | |
| assert client.transport == transport | |
| # Test basic functionality | |
| await transport.connect() | |
| # Test tool calling through transport | |
| result = await transport.call_tool("test_tool", {"arg": "value"}) | |
| @pytest.mark.asyncio | |
| async def test_transport_with_mcp_client(self): | |
| """Test using custom transport with MCP client.""" | |
| # Register custom transport | |
| register_custom_transport( | |
| name="client_test", transport_class=IntegrationTestTransport | |
| ) | |
| # Create transport | |
| transport = create_transport("client_test://mcp-server") | |
| # Import and create MCP client (this would normally be done) | |
| # This is a simplified test - in real usage, you'd use the full client | |
| from mcp_fuzzer.client.tool_client import ToolClient | |
| # Test that the client can be created with custom transport | |
| client = ToolClient(transport) | |
| assert client.transport == transport | |
| # Test basic functionality | |
| await transport.connect() | |
| # Test tool calling through transport | |
| result = await transport.call_tool("test_tool", {"arg": "value"}) |
🤖 Prompt for AI Agents
In tests/integration/test_custom_transport.py around lines 248 to 270, the async
test function test_transport_with_mcp_client is not marked for asyncio; add the
pytest.mark.asyncio decorator above the function (and import pytest at top of
file if missing) so pytest can run async tests, or alternatively enable
pytest-asyncio’s auto mode in pytest.ini — ensure the decorator is applied
directly to the test function and any necessary pytest-asyncio dependency is
available.
fixes #9
fixes #36
Summary by CodeRabbit
New Features
Improvements
Documentation
Tests
Chores