From 0f0e67e3a34dd51e4bb4036af611979537184f35 Mon Sep 17 00:00:00 2001 From: Gagan Deep Date: Thu, 30 Jul 2026 14:28:18 +0530 Subject: [PATCH] fix: async wrapping, stdout pollution, URL policy scoping, sdist contents Four defects found by a launch-readiness audit of 0.3.1, each verified against the published wheel. - @shield_tool always built a sync wrapper, so inspect.iscoroutinefunction() returned False on a shielded async def. LangChain and MCP branch on that to decide whether to await, so the tool body never ran and the model received a coroutine repr as its observation. Coroutine functions and async generators are now wrapped in kind. - The decorator printed to stdout on every allowed call. stdout is the transport for MCP stdio servers, so this corrupted the JSON-RPC stream. Blocks now log at WARNING on the "modelfuzz" logger with structured modelfuzz_tool / modelfuzz_rule / modelfuzz_reason fields; allows log at DEBUG. Nothing touches stdout. - URLAllowList was applied to every argument, so http_post(url, body) blocked its own legitimate calls -- body is not a URL. It now returns None for values that are not URLs. While rewriting it: reject non-http(s) schemes (file://api.internal.com/etc/passwd previously passed) and compare hostnames case-insensitively via urlparse().hostname. - No [tool.hatch.build.targets.sdist] meant hatchling swept the working tree, so the 0.3.1 sdist shipped .claude/settings.local.json. Now an allowlist. PolicyResult carries the originating Violation so the log can name the rule. Tests 25 -> 60; decorator.py and engine.py to 100%. --- .gitignore | 5 ++ CHANGELOG.md | 12 ++++ pyproject.toml | 26 ++++++- src/modelfuzz/decorator.py | 67 +++++++++++++++--- src/modelfuzz/engine.py | 3 +- src/modelfuzz/rules.py | 102 ++++++++++++++++++---------- tests/test_decorator.py | 135 +++++++++++++++++++++++++++++++++++++ tests/test_rules.py | 90 +++++++++++++++++++++++++ uv.lock | 2 +- 9 files changed, 396 insertions(+), 46 deletions(-) diff --git a/.gitignore b/.gitignore index bcbb7e1..07400d2 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,8 @@ build/ .pytest_cache/ .ruff_cache/ .DS_Store +.claude/ +.coverage +.coverage.* +htmlcov/ +.env diff --git a/CHANGELOG.md b/CHANGELOG.md index b559d89..a8869b4 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 1888bf5..c32c31f 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/modelfuzz/decorator.py b/src/modelfuzz/decorator.py index 88b2490..69cd6b7 100644 --- a/src/modelfuzz/decorator.py +++ b/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 diff --git a/src/modelfuzz/engine.py b/src/modelfuzz/engine.py index 8c88a66..9d7e41d 100644 --- a/src/modelfuzz/engine.py +++ b/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) diff --git a/src/modelfuzz/rules.py b/src/modelfuzz/rules.py index 47c065a..7667ab8 100644 --- a/src/modelfuzz/rules.py +++ b/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: diff --git a/tests/test_decorator.py b/tests/test_decorator.py index cbdc561..e8445dc 100644 --- a/tests/test_decorator.py +++ b/tests/test_decorator.py @@ -1,5 +1,9 @@ """Tests for the ModelFuzz shield_tool decorator.""" +import asyncio +import inspect +import logging + import pytest from modelfuzz import ModelFuzzBlockError, shield_tool @@ -11,6 +15,26 @@ def send_email(to: str, subject: str, body: str) -> str: return f"Email sent to {to} with subject '{subject}'." +@shield_tool +def send_email_bare(to: str, body: str) -> str: + """Pretend tool decorated with the bare form.""" + return f"Sent to {to}." + + +@shield_tool() +async def fetch_url(url: str, note: str = "clean") -> str: + """Pretend async tool.""" + await asyncio.sleep(0) + return f"fetched {url}" + + +@shield_tool() +async def stream_rows(query: str): + """Pretend async-generator tool.""" + for row in ("a", "b"): + yield f"{query}:{row}" + + class TestShieldToolDecorator: """Tests for the shield_tool decorator.""" @@ -30,3 +54,114 @@ def test_decorator_preserves_metadata(self): """Assert that the decorator preserves the function's metadata.""" assert send_email.__name__ == "send_email" assert send_email.__doc__ == "Pretend tool that sends an email." + + def test_blocks_violation_passed_as_keyword(self): + """Keyword arguments are checked, not just positional ones.""" + with pytest.raises(ModelFuzzBlockError): + send_email(to="alice@example.com", subject="Hello", body="my password is 12345") + + def test_blocks_violation_in_mixed_args(self): + """A violation in a keyword arg is caught alongside clean positional args.""" + with pytest.raises(ModelFuzzBlockError): + send_email("alice@example.com", "Hello", body="the secret is out") + + def test_bare_decorator_form_blocks(self): + """The bare @shield_tool form applies the default engine.""" + with pytest.raises(ModelFuzzBlockError): + send_email_bare("alice@example.com", "my password is 12345") + assert send_email_bare("alice@example.com", "hi") == "Sent to alice@example.com." + + +class TestStdoutIsNotTouched: + """A library must not write to stdout: it is the MCP stdio transport.""" + + def test_allowed_call_writes_nothing_to_stdout(self, capsys): + """An allowed call leaves stdout completely empty.""" + send_email("alice@example.com", "Hello", "Hi Alice, how are you?") + captured = capsys.readouterr() + assert captured.out == "" + assert captured.err == "" + + def test_blocked_call_writes_nothing_to_stdout(self, capsys): + """A blocked call also leaves stdout empty -- it logs instead.""" + with pytest.raises(ModelFuzzBlockError): + send_email("alice@example.com", "Hello", "my password is 12345") + assert capsys.readouterr().out == "" + + +class TestBlockIsLogged: + """Blocks are the audit record, so they must be logged.""" + + def test_block_logs_warning_with_structured_fields(self, caplog): + """A block emits a WARNING carrying the tool, rule and reason.""" + with ( + caplog.at_level(logging.WARNING, logger="modelfuzz"), + pytest.raises(ModelFuzzBlockError), + ): + send_email("alice@example.com", "Hello", "my password is 12345") + + assert len(caplog.records) == 1 + record = caplog.records[0] + assert record.levelno == logging.WARNING + assert record.modelfuzz_tool == "send_email" + assert record.modelfuzz_rule == "SensitiveDataFilter" + assert "password" in record.modelfuzz_reason + + def test_allowed_call_logs_no_warning(self, caplog): + """An allowed call produces no WARNING.""" + with caplog.at_level(logging.WARNING, logger="modelfuzz"): + send_email("alice@example.com", "Hello", "Hi Alice, how are you?") + assert caplog.records == [] + + +class TestAsyncTools: + """Coroutine functions must stay coroutine functions once shielded.""" + + def test_wrapped_coroutine_is_still_a_coroutine_function(self): + """Frameworks branch on this to decide whether to await.""" + assert inspect.iscoroutinefunction(fetch_url) + assert asyncio.iscoroutinefunction(fetch_url) + + def test_async_safe_call_returns_the_real_value(self): + """An awaited safe call returns the tool's value, not a coroutine.""" + result = asyncio.run(fetch_url("http://example.com")) + assert result == "fetched http://example.com" + + def test_async_malicious_call_raises(self): + """A violation in an async tool raises ModelFuzzBlockError.""" + with pytest.raises(ModelFuzzBlockError): + asyncio.run(fetch_url("http://example.com", note="my password is 12345")) + + def test_async_block_happens_before_the_body_runs(self): + """The body must never execute when a policy trips.""" + ran = [] + + @shield_tool() + async def tool(payload: str) -> str: + ran.append(payload) + return "done" + + with pytest.raises(ModelFuzzBlockError): + asyncio.run(tool("the secret is out")) + assert ran == [] + + def test_wrapped_async_generator_is_still_an_async_generator(self): + """Async-generator tools keep their kind too.""" + assert inspect.isasyncgenfunction(stream_rows) + + def test_async_generator_safe_call_yields(self): + """A safe async-generator call yields its rows.""" + + async def collect(): + return [row async for row in stream_rows("q")] + + assert asyncio.run(collect()) == ["q:a", "q:b"] + + def test_async_generator_malicious_call_raises(self): + """A violation in an async-generator tool raises on iteration.""" + + async def collect(): + return [row async for row in stream_rows("my password is 12345")] + + with pytest.raises(ModelFuzzBlockError): + asyncio.run(collect()) diff --git a/tests/test_rules.py b/tests/test_rules.py index 9e57b11..b9f3686 100644 --- a/tests/test_rules.py +++ b/tests/test_rules.py @@ -42,6 +42,96 @@ def test_blocks_userinfo_trick(self, url_allowlist: URLAllowList): assert violation is not None assert "userinfo trick" in violation.reason + def test_blocks_sibling_domain(self, url_allowlist: URLAllowList): + """The subdomain boundary is a dot: evilapi.internal.com is not allowed.""" + violation = url_allowlist("http://evilapi.internal.com") + assert violation is not None + assert "not in allowlist" in violation.reason + + def test_allows_real_subdomain(self, url_allowlist: URLAllowList): + """A genuine subdomain is allowed.""" + assert url_allowlist("https://sub.api.internal.com/v1") is None + + def test_allows_uppercase_host(self, url_allowlist: URLAllowList): + """Hostnames are case-insensitive per DNS.""" + assert url_allowlist("https://API.INTERNAL.COM/v1") is None + + def test_allows_host_with_port(self, url_allowlist: URLAllowList): + """A port does not change the host.""" + assert url_allowlist("https://api.internal.com:8443/v1") is None + + def test_allows_trailing_dot_host(self, url_allowlist: URLAllowList): + """A fully-qualified trailing dot is the same host.""" + assert url_allowlist("https://api.internal.com./v1") is None + + def test_blocks_disallowed_scheme(self, url_allowlist: URLAllowList): + """Only http/https by default, even for an allowlisted host.""" + violation = url_allowlist("file://api.internal.com/etc/passwd") + assert violation is not None + assert "scheme not allowed" in violation.reason + + def test_blocks_malformed_url_that_claims_to_be_one(self, url_allowlist: URLAllowList): + """A string with '://' but no usable host fails closed.""" + violation = url_allowlist("http://") + assert violation is not None + assert "Invalid URL" in violation.reason + + def test_blocks_unparseable_url(self, url_allowlist: URLAllowList): + """A URL-looking string that urlparse rejects fails closed.""" + violation = url_allowlist("http://[oops") + assert violation is not None + assert "Invalid URL" in violation.reason + + def test_blocks_url_with_port_but_no_host(self, url_allowlist: URLAllowList): + """A netloc that parses but yields no hostname fails closed.""" + violation = url_allowlist("https://:8080/path") + assert violation is not None + assert "Invalid URL" in violation.reason + + +class TestURLAllowListIgnoresNonURLs: + """The policy governs URLs only, so it can coexist on multi-arg tools.""" + + @pytest.fixture + def url_allowlist(self) -> URLAllowList: + """Fixture for a URLAllowList with 'api.internal.com' allowed.""" + return URLAllowList(allowed_domains=["api.internal.com"]) + + @pytest.mark.parametrize( + "value", + ["hello world", "", "just some prose about api.internal.com", "a/b/c"], + ) + def test_allows_non_url_strings(self, url_allowlist: URLAllowList, value: str): + """Prose and paths are not URLs and are not this policy's business.""" + assert url_allowlist(value) is None + + @pytest.mark.parametrize("value", [30, None, 3.5, True, {"a": 1}, ["x"], b"bytes"]) + def test_allows_non_string_values(self, url_allowlist: URLAllowList, value: object): + """Non-string arguments are passed through untouched.""" + assert url_allowlist(value) is None + + def test_guards_a_multi_argument_tool(self, url_allowlist: URLAllowList): + """The regression that motivated this: http_post(url, body, timeout).""" + from modelfuzz import ModelFuzzBlockError, PolicyEngine, shield_tool + + engine = PolicyEngine([url_allowlist]) + + @shield_tool(engine=engine) + def http_post(url: str, body: str, timeout: int = 30) -> str: + return f"posted to {url}" + + # A legitimate call is not blocked by its own non-URL arguments. + assert http_post("https://api.internal.com/v1", "hello world") == ( + "posted to https://api.internal.com/v1" + ) + assert http_post("https://api.internal.com/v1", "hi", timeout=5) == ( + "posted to https://api.internal.com/v1" + ) + + # A disallowed host is still blocked. + with pytest.raises(ModelFuzzBlockError): + http_post("http://evil.com/exfil", "hello world") + class TestSensitiveDataFilter: """Tests for the SensitiveDataFilter policy.""" diff --git a/uv.lock b/uv.lock index 9976ffd..01b8939 100644 --- a/uv.lock +++ b/uv.lock @@ -286,7 +286,7 @@ wheels = [ [[package]] name = "modelfuzz" -version = "0.2.1" +version = "0.3.2" source = { editable = "." } dependencies = [ { name = "typer" },