This commit is contained in:
Eugene Yurtsev
2025-03-24 10:16:44 -04:00
parent 71ddda1d09
commit e8b5aadb92
2 changed files with 29 additions and 4 deletions

View File

@@ -42,6 +42,12 @@ Examples:
# Using SSE transport with additional HTTP options # Using SSE transport with additional HTTP options
mcpdoc --yaml sample_config.yaml --follow-redirects --timeout 15 --transport sse --host localhost --port 8080 mcpdoc --yaml sample_config.yaml --follow-redirects --timeout 15 --transport sse --host localhost --port 8080
# Allow fetching from additional domains. The domains hosting the llms.txt files are always allowed.
mcpdoc --yaml sample_config.yaml --allowed-domains https://example.com/ https://another-example.com/
# Allow fetching from any domain
mcpdoc --yaml sample_config.yaml --allowed-domains '*'
""" """
@@ -74,6 +80,12 @@ def parse_args() -> argparse.Namespace:
action="store_true", action="store_true",
help="Whether to follow HTTP redirects", help="Whether to follow HTTP redirects",
) )
parser.add_argument(
"--allowed-domains",
type=str,
nargs="*",
help="Additional allowed domains to fetch documentation from. Use '*' to allow all domains",
)
parser.add_argument( parser.add_argument(
"--timeout", type=float, default=10.0, help="HTTP request timeout in seconds" "--timeout", type=float, default=10.0, help="HTTP request timeout in seconds"
) )
@@ -229,6 +241,7 @@ def main() -> None:
follow_redirects=args.follow_redirects, follow_redirects=args.follow_redirects,
timeout=args.timeout, timeout=args.timeout,
settings=settings, settings=settings,
allowed_domains=args.allowed_domains,
) )
if args.transport == "sse": if args.transport == "sse":

View File

@@ -40,6 +40,7 @@ def create_server(
follow_redirects: bool = False, follow_redirects: bool = False,
timeout: float = 10, timeout: float = 10,
settings: dict | None = None, settings: dict | None = None,
allowed_domains: list[str] | None = None,
) -> FastMCP: ) -> FastMCP:
"""Create the server and generate documentation retrieval tools. """Create the server and generate documentation retrieval tools.
@@ -48,6 +49,10 @@ def create_server(
follow_redirects: Whether to follow HTTP redirects when fetching docs follow_redirects: Whether to follow HTTP redirects when fetching docs
timeout: HTTP request timeout in seconds timeout: HTTP request timeout in seconds
settings: Additional settings to pass to FastMCP settings: Additional settings to pass to FastMCP
allowed_domains: Additional domains to allow fetching from.
Use ['*'] to allow all domains
The domain hosting the llms.txt file is always appended to the list
of allowed domains.
Returns: Returns:
A FastMCP server instance configured with documentation tools A FastMCP server instance configured with documentation tools
@@ -81,7 +86,14 @@ def create_server(
return content return content
# Parse the domain names in the llms.txt URLs # Parse the domain names in the llms.txt URLs
allowed_domains = set(extract_domain(entry["llms_txt"]) for entry in doc_source) domains = set(extract_domain(entry["llms_txt"]) for entry in doc_source)
# Add additional allowed domains if specified
if allowed_domains:
if "*" in allowed_domains:
domains = {"*"} # Special marker for allowing all domains
else:
domains.update(allowed_domains)
@server.tool() @server.tool()
async def fetch_docs(url: str) -> str: async def fetch_docs(url: str) -> str:
@@ -99,11 +111,11 @@ def create_server(
The fetched documentation content converted to markdown, or an error message The fetched documentation content converted to markdown, or an error message
if the request fails or the URL is not from an allowed domain. if the request fails or the URL is not from an allowed domain.
""" """
nonlocal allowed_domains nonlocal domains
if not any(url.startswith(domain) for domain in allowed_domains): if "*" not in domains and not any(url.startswith(domain) for domain in domains):
return ( return (
"Error: URL not allowed. Must start with one of the following domains: " "Error: URL not allowed. Must start with one of the following domains: "
+ ", ".join(allowed_domains) + ", ".join(domains)
) )
try: try: