fix: async wrapping, stdout pollution, URL policy scoping, sdist contents

.gitignore

@@ -7,3 +7,8 @@ build/
 .pytest_cache/
 .ruff_cache/
 .DS_Store
+.claude/
+.coverage
+.coverage.*
+htmlcov/
+.env

CHANGELOG.md

@@ -4,6 +4,18 @@ All notable changes to this project are documented here.
 
 ## [Unreleased]
 
+## [0.3.2] - 2026-07-30
+
+- fix: `@shield_tool` now wraps coroutine functions and async generators in kind. Previously it always produced a sync wrapper, so `inspect.iscoroutinefunction()` returned `False` on a shielded `async def` and frameworks (LangChain, MCP) never awaited it — the model received a coroutine `repr` instead of the tool's result
+- fix: stop writing to stdout on every allowed call. stdout is the transport for MCP stdio servers, and the `print()` corrupted the JSON-RPC stream
+- feat: blocked calls are now logged at `WARNING` on the `modelfuzz` logger with structured `modelfuzz_tool` / `modelfuzz_rule` / `modelfuzz_reason` fields, so a denial leaves an audit record
+- fix: `URLAllowList` returns `None` for values that are not URLs, so it can guard a tool like `http_post(url, body, timeout)` without flagging `body` or `timeout`. Previously any multi-argument tool blocked 100% of its legitimate calls
+- fix: `URLAllowList` rejects non-http(s) schemes, which previously passed through for an allowlisted host (`file://api.internal.com/etc/passwd`)
+- fix: `URLAllowList` compares hostnames case-insensitively and tolerates a trailing dot, via `urlparse().hostname` instead of hand-rolled `netloc` splitting
+- feat: `PolicyResult` carries the originating `Violation`, so callers can see which rule fired
+- fix: declare `[tool.hatch.build.targets.sdist]`. The 0.3.1 sdist shipped `.claude/settings.local.json` because hatchling swept the working tree
+- test: cover the async, stdout, structured-logging, keyword-argument and bare-decorator paths that had no coverage (25 → 60 tests)
+
 ## [0.3.1] - 2026-07-30
 
 - docs: replace the Red-Team Scanner example with real captured scan output, contrasting a weak model breached on the first probe against a resistant model whose refusals are mutated into new variants

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "modelfuzz"
-version = "0.3.1"
+version = "0.3.2"
 description = "Runtime guardrails for AI agents."
 readme = "README.md"
 license = { text = "MIT" }
@@ -52,6 +52,30 @@ build-backend = "hatchling.build"
 [tool.hatch.build.targets.wheel]
 packages = ["src/modelfuzz"]
 
+# Without an explicit sdist target, hatchling sweeps the working tree and ships
+# whatever is lying around (0.3.1 shipped .claude/settings.local.json this way).
+# This is an allowlist: anything not named here stays out of the tarball.
+[tool.hatch.build.targets.sdist]
+include = [
+    "src/modelfuzz",
+    "tests",
+    "demo.py",
+    "pyproject.toml",
+    "README.md",
+    "LICENSE",
+    "CHANGELOG.md",
+    "SECURITY.md",
+    "CONTRIBUTING.md",
+]
+exclude = [
+    ".claude",
+    ".coverage*",
+    "htmlcov",
+    ".env",
+    "uv.lock",
+    "**/__pycache__",
+]
+
 [tool.ruff]
 line-length = 100
 src = ["src", "tests"]

src/modelfuzz/decorator.py

@@ -1,6 +1,8 @@
 """Decorators for shielding agent tool functions."""
 
 import functools
+import inspect
+import logging
 from collections.abc import Callable
 from typing import ParamSpec, TypeVar, overload
 
@@ -11,6 +13,8 @@
 P = ParamSpec("P")
 R = TypeVar("R")
 
+logger = logging.getLogger("modelfuzz")
+
 # Default policy engine for the decorator
 default_policies = [SensitiveDataFilter()]
 _default_engine = PolicyEngine(default_policies)
@@ -28,7 +32,9 @@ def shield_tool(engine=None):
     """Wrap a tool function so every call is intercepted before execution.
 
     Usable bare (``@shield_tool``) or called (``@shield_tool()`` /
-    ``@shield_tool(engine)``).
+    ``@shield_tool(engine)``). Sync functions, coroutine functions and async
+    generators are each wrapped in kind, so framework introspection
+    (``inspect.iscoroutinefunction`` and friends) keeps working.
 
     Args:
         engine: The policy engine to use for validation. If None, a default
@@ -49,16 +55,61 @@ def decorator(func: Callable[P, R]) -> Callable[P, R]:
     return decorator
 
 
+def _enforce(func: Callable[..., object], actual_engine: PolicyEngine, args, kwargs) -> None:
+    """Check every argument against the engine, raising on the first violation.
+
+    Blocks are logged at WARNING with structured fields so they reach the host
+    application's audit trail. Nothing is written to stdout -- stdout is the
+    transport for MCP stdio servers, and a stray write there corrupts the
+    JSON-RPC stream.
+    """
+    for arg in list(args) + list(kwargs.values()):
+        result = actual_engine.run(arg)
+        if result.allowed:
+            continue
+
+        reason = result.reason or "Call blocked by policy"
+        rule_name = result.violation.rule_name if result.violation else None
+        logger.warning(
+            "ModelFuzz blocked tool call: tool=%s rule=%s reason=%s",
+            func.__name__,
+            rule_name,
+            reason,
+            extra={
+                "modelfuzz_tool": func.__name__,
+                "modelfuzz_rule": rule_name,
+                "modelfuzz_reason": reason,
+            },
+        )
+        raise ModelFuzzBlockError(reason)
+
+
 def _wrap(func: Callable[P, R], actual_engine: PolicyEngine) -> Callable[P, R]:
+    if inspect.iscoroutinefunction(func):
+
+        @functools.wraps(func)
+        async def async_wrapper(*args: P.args, **kwargs: P.kwargs):
+            _enforce(func, actual_engine, args, kwargs)
+            logger.debug("ModelFuzz intercepted: %s", func.__name__)
+            return await func(*args, **kwargs)
+
+        return async_wrapper
+
+    if inspect.isasyncgenfunction(func):
+
+        @functools.wraps(func)
+        async def asyncgen_wrapper(*args: P.args, **kwargs: P.kwargs):
+            _enforce(func, actual_engine, args, kwargs)
+            logger.debug("ModelFuzz intercepted: %s", func.__name__)
+            async for item in func(*args, **kwargs):
+                yield item
+
+        return asyncgen_wrapper
+
     @functools.wraps(func)
     def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-        # Check all arguments against the policy engine
-        for arg in list(args) + list(kwargs.values()):
-            result = actual_engine.run(arg)
-            if not result.allowed:
-                raise ModelFuzzBlockError(result.reason or "Call blocked by policy")
-
-        print(f"ModelFuzz Intercepted: {func.__name__}")
+        _enforce(func, actual_engine, args, kwargs)
+        logger.debug("ModelFuzz intercepted: %s", func.__name__)
         return func(*args, **kwargs)
 
     return wrapper

src/modelfuzz/engine.py

@@ -12,6 +12,7 @@ class PolicyResult:
 
     allowed: bool
     reason: str | None = None
+    violation: Violation | None = None
 
 
 class PolicyEngine:
@@ -32,5 +33,5 @@ def run(self, call: object) -> PolicyResult:
         for policy in self.policies:
             violation = policy(call)
             if violation:
-                return PolicyResult(allowed=False, reason=violation.reason)
+                return PolicyResult(allowed=False, reason=violation.reason, violation=violation)
         return PolicyResult(allowed=True)

src/modelfuzz/rules.py

@@ -12,54 +12,86 @@ class Violation:
     reason: str
 
 
-class URLAllowList:
-    """A policy that ensures URLs are on an allowlist and blocks parsing tricks."""
+DEFAULT_URL_SCHEMES = frozenset({"http", "https"})
 
-    def __init__(self, allowed_domains: list[str]) -> None:
-        self.allowed_domains = allowed_domains
 
-    def __call__(self, url: str) -> Violation | None:
+class URLAllowList:
+    """A policy that ensures URLs are on an allowlist and blocks parsing tricks.
+
+    The policy governs *URLs only*. Arguments that are not URLs -- an email
+    body, a timeout int, None -- are passed through untouched, so a single
+    engine can guard a tool like ``http_post(url, body)`` without flagging
+    ``body``. Note the tradeoff: a bare host with no scheme (``"evil.com"``) is
+    not identifiable as a URL and is therefore allowed through. Pair this with
+    a rule that governs the arguments it does not.
+    """
+
+    def __init__(
+        self,
+        allowed_domains: list[str],
+        allowed_schemes: set[str] | frozenset[str] | None = None,
+    ) -> None:
+        self.allowed_domains = [d.lower().rstrip(".") for d in allowed_domains]
+        self.allowed_schemes = (
+            frozenset(s.lower() for s in allowed_schemes)
+            if allowed_schemes is not None
+            else DEFAULT_URL_SCHEMES
+        )
+
+    def __call__(self, url: object) -> Violation | None:
         """Check if a URL is allowed.
 
         Args:
-            url: The URL to check.
+            url: The value to check. Non-URL values are not governed by this
+                policy and return None.
 
         Returns:
             A Violation object if the URL is blocked, otherwise None.
         """
+        if not isinstance(url, str):
+            return None
+
+        # A string carrying a scheme separator is claiming to be a URL, so a
+        # parse failure from here on must fail closed rather than sail through.
+        looks_like_url = "://" in url
+
         try:
             parsed = urlparse(url)
-            domain = parsed.netloc
-
-            # Block userinfo tricks (e.g., http://api.internal.com@evil.com)
-            if "@" in domain:
-                return Violation(
-                    rule_name="URLAllowList",
-                    reason=f"URL contains userinfo trick: {url}",
-                )
-
-            # Extract the hostname without port
-            hostname = domain.split(":")[0]
-
-            # Check for exact match or valid subdomain
-            is_allowed = any(
-                hostname == allowed or hostname.endswith(f".{allowed}")
-                for allowed in self.allowed_domains
-            )
-
-            if not is_allowed:
-                return Violation(
-                    rule_name="URLAllowList",
-                    reason=f"URL domain not in allowlist: {hostname}",
-                )
+        except Exception:
+            return self._block(f"Invalid URL: {url}") if looks_like_url else None
 
-            return None
+        if not parsed.scheme or not parsed.netloc:
+            return self._block(f"Invalid URL: {url}") if looks_like_url else None
 
-        except Exception:
-            return Violation(
-                rule_name="URLAllowList",
-                reason=f"Invalid URL: {url}",
-            )
+        if parsed.scheme.lower() not in self.allowed_schemes:
+            return self._block(f"URL scheme not allowed: {parsed.scheme}")
+
+        # Block userinfo tricks (e.g., http://api.internal.com@evil.com)
+        if "@" in parsed.netloc:
+            return self._block(f"URL contains userinfo trick: {url}")
+
+        try:
+            hostname = (parsed.hostname or "").rstrip(".")
+        except ValueError:
+            return self._block(f"Invalid URL: {url}")
+
+        if not hostname:
+            return self._block(f"Invalid URL: {url}")
+
+        # Check for exact match or valid subdomain
+        is_allowed = any(
+            hostname == allowed or hostname.endswith(f".{allowed}")
+            for allowed in self.allowed_domains
+        )
+
+        if not is_allowed:
+            return self._block(f"URL domain not in allowlist: {hostname}")
+
+        return None
+
+    @staticmethod
+    def _block(reason: str) -> Violation:
+        return Violation(rule_name="URLAllowList", reason=reason)
 
 
 class SensitiveDataFilter: