llm.py

import json
import logging
import os
import re

# seed_generators was deleted in commit 87d8602 together with src/seed_bank/
# and src/predefined_seeds/, but the eight @tool functions below were left
# calling it. Every LLM seed-generation tool call therefore died with
# "name 'seed_generators' is not defined" -- visible in the campaign logs as
# "Error running tool create_json" -- and seeding fell back to the
# deterministic mapper on every tool. Restored from b1b657e.
from . import seed_generators
try:
    from langchain_openai import ChatOpenAI
    from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
    from langchain_core.tools import tool
except ImportError as _e:
    _LANGCHAIN_IMPORT_ERROR = _e
    ChatOpenAI = None
    HumanMessage = SystemMessage = ToolMessage = None
    def tool(func):
        return func
else:
    _LANGCHAIN_IMPORT_ERROR = None



logger = logging.getLogger(__name__)


class SeedProvenanceError(RuntimeError):
    """TOFU_REQUIRE_LLM_SEEDS is set but no LLM-generated seeds are available."""


def _no_llm() -> bool:
    """TOFU_NO_LLM: skip BOTH LLM calls (selection + seed-gen). Selection returns the
    always-on BASE_MANDATORY_SANITIZERS; seeds come from the deterministic local
    fallback mapping. Zero API cost -- used by the no-LLM executability/first-pass probe
    so LLM spend only happens on tools that actually execute."""
    return os.environ.get("TOFU_NO_LLM", "").strip().lower() not in ("", "0", "false", "no")


# ---- ablation switches (eval/run_ablation.sh). Each removes ONE LLM stage and
# holds the other fixed (cached), so a delta is attributable to that stage.
def _flag(name: str) -> bool:
    return os.environ.get(name, "").strip().lower() not in ("", "0", "false", "no")


def _no_llm_seeds() -> bool:
    """TOFU_NO_LLM_SEEDS: seeds come from the deterministic local fallback mapping
    (schema boundaries + seed-bank payloads); sanitizer selection is unchanged."""
    return _no_llm() or _flag("TOFU_NO_LLM_SEEDS")


def _no_llm_sanitizers():
    """TOFU_NO_LLM_SANITIZERS=mandatory|all: skip LLM selection. 'mandatory' = only
    BASE_MANDATORY_SANITIZERS; 'all' = every AVAILABLE_SANITIZERS entry. Returns the
    mode or None. A value that is neither is an error, never silently ignored --
    a typo would otherwise run the full-ToFu arm under an ablation's name."""
    if _no_llm():
        return "mandatory"
    v = os.environ.get("TOFU_NO_LLM_SANITIZERS", "").strip().lower()
    if v in ("", "0", "false", "no"):
        return None
    if v not in ("mandatory", "all"):
        raise ValueError(f"TOFU_NO_LLM_SANITIZERS={v!r}: expected 'mandatory' or 'all'")
    return v


def _cache_readonly() -> bool:
    """TOFU_LLM_CACHE_READONLY: replay the cache but never write it. Ablation arms
    run after the full arm; a miss in one of them must not change what later arms
    replay."""
    return _flag("TOFU_LLM_CACHE_READONLY")


def _require_llm_seeds() -> bool:
    """Strict seed provenance (per-repo deep campaign): use cached seeds only when
    they are tagged seeds_provenance == "llm"; regenerate anything else live; and
    if the live LLM call fails, abort the tool instead of fuzzing it with the
    local fallback mapping. Off by default -- normal runs keep the fallback."""
    return os.environ.get("TOFU_REQUIRE_LLM_SEEDS", "").strip().lower() not in ("", "0", "false", "no")


_FENCE_RE = re.compile(r"```(?:[A-Za-z0-9_+-]*)\s*\n?(.*?)```", re.DOTALL)


def _coerce_text(content) -> str:
    """Flatten a LangChain message payload to plain text.

    Most backends hand back a str, but some return a list of content
    blocks; treating that list as a string is what produced the
    ``'list' object has no attribute 'strip'`` class of failure.
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for b in content:
            if isinstance(b, dict):
                parts.append(b.get("text") or b.get("content") or "")
            elif isinstance(b, str):
                parts.append(b)
        return "".join(parts)
    return str(content)


def _balanced_span(text: str, open_ch: str, close_ch: str):
    """First balanced ``open_ch``..``close_ch`` span, ignoring delimiters
    that occur inside JSON string literals. Returns None if unbalanced."""
    depth = 0
    start = -1
    in_str = False
    esc = False
    for i, ch in enumerate(text):
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
            continue
        if ch == '"':
            in_str = True
            continue
        if ch == open_ch:
            if depth == 0:
                start = i
            depth += 1
        elif ch == close_ch and depth:
            depth -= 1
            if depth == 0:
                return text[start:i + 1]
    return None



# A backslash that does not start a valid JSON escape (\" \\ \/ \b \f \n \r \t \uXXXX).
# LLM-written adversarial seeds are full of them: Windows traversal paths
# (..\..\windows\win.ini) typed with single backslashes are invalid JSON.
_ESCAPE_RE = re.compile(r'\\(["\\/bfnrt]|u[0-9a-fA-F]{4})|\\')


def _repair_json_text(text: str) -> str:
    """Double every backslash that does not begin a valid JSON escape. Scans left
    to right and consumes valid escapes whole, so an existing \\\\ pair is kept."""
    return _ESCAPE_RE.sub(lambda m: m.group(0) if m.group(1) else "\\\\", text)


def _salvage_objects(text: str) -> list:
    """Parse each top-level {...} object on its own (escape-repaired, control
    characters allowed) and keep the ones that parse. Used only when the reply as
    a whole cannot be parsed, so one broken seed does not discard the rest."""
    out, i = [], 0
    while True:
        j = text.find("{", i)
        if j < 0:
            break
        span = _balanced_span(text[j:], "{", "}")
        if not span:
            break
        for cand in (span, _repair_json_text(span)):
            try:
                v = json.loads(cand, strict=False)
            except Exception:
                continue
            if isinstance(v, dict):
                out.append(v)
            break
        i = j + len(span)
    return out

def _extract_json(content, expect=None, repair=False, info=None):
    """Recover a JSON value from an LLM reply that may wrap it in prose
    and/or a markdown fence.

    Small local models (llama3.1:8b, qwen2.5, mistral, ...) routinely
    ignore a "raw JSON only" instruction and answer with an explanation,
    a fenced block, and a trailing summary. The code this replaces did
    ``s.strip().strip("```json").strip("```")``, which is wrong twice
    over: ``str.strip(chars)`` removes a *character set* rather than a
    substring, and nothing there survives prose outside the fence. Every
    reply from such a model therefore died with "Expecting value: line 1
    column 1 (char 0)" and the caller silently fell back to the
    deterministic default sanitizers -- i.e. the LLM stage of the
    pipeline never ran, while the logs looked healthy.

    ``expect`` may be ``dict``, ``list``, or ``None`` for either.
    Raises ValueError if nothing parsable is present.
    """
    text = _coerce_text(content).strip()
    if not text:
        raise ValueError("empty LLM response")

    candidates = [text]
    candidates.extend(m.strip() for m in _FENCE_RE.findall(text))

    pairs = [("{", "}", dict), ("[", "]", list)]
    for open_ch, close_ch, typ in pairs:
        if expect is not None and typ is not expect:
            continue
        for base in list(candidates):
            span = _balanced_span(base, open_ch, close_ch)
            if span and span not in candidates:
                candidates.append(span)

    last_err = None
    for cand in candidates:
        if not cand:
            continue
        try:
            data = json.loads(cand)
        except Exception as e:
            last_err = e
            continue
        if expect is not None and not isinstance(data, expect):
            last_err = ValueError(
                "parsed a %s, expected %s" % (type(data).__name__, expect.__name__)
            )
            continue
        return data
    if repair:
        # Seed generation only (selection parsing is unchanged). Each step keeps
        # the model's own content; nothing is invented. ``info["repaired"]``
        # records which step was needed so recovered replies stay auditable.
        rcands = []
        for base in candidates:
            if not base:
                continue
            fixed = _repair_json_text(base)
            rcands += [("control_chars", base), ("escapes", fixed)]
            for open_ch, close_ch, typ in pairs:
                if expect is None or typ is expect:
                    sp = _balanced_span(fixed, open_ch, close_ch)
                    if sp:
                        rcands.append(("escapes", sp))
        for how, cand in rcands:
            try:
                data = json.loads(cand, strict=False)
            except Exception:
                continue
            if expect is None or isinstance(data, expect):
                if info is not None:
                    info["repaired"] = how
                return data
        if expect in (None, list):
            objs = _salvage_objects(text)
            if objs:
                if info is not None:
                    info["repaired"] = f"salvaged:{len(objs)}"
                return objs
    raise ValueError("no parsable JSON in LLM response (%s)" % (last_err,))


@tool
def create_pdf(content: str, filename: str) -> str:
    """
    Creates a valid PDF document with the specified content and saves it.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_pdf(content, filename)


@tool
def create_docx(content: str, filename: str) -> str:
    """
    Creates a valid Word (DOCX) document with the specified content and saves it.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_docx(content, filename)


@tool
def create_xlsx(sheet_data: list, filename: str) -> str:
    """
    Creates a valid Excel (XLSX) spreadsheet with the specified table data and saves it.
    sheet_data is a list of rows, where each row is a list of cell values.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_xlsx(sheet_data, filename)


@tool
def create_pptx(content: str, filename: str) -> str:
    """
    Creates a valid PowerPoint (PPTX) presentation with the specified content and saves it.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_pptx(content, filename)


@tool
def create_zip(files: dict, filename: str) -> str:
    """
    Creates a valid ZIP archive containing the specified file names and their text contents and saves it.
    files is a dictionary mapping sub-filename to its text content.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_zip(files, filename)


@tool
def create_xml(xml_content: str, filename: str) -> str:
    """
    Creates a valid XML file with the specified content and saves it.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_xml(xml_content, filename)


@tool
def create_html(html_content: str, filename: str) -> str:
    """
    Creates a valid HTML file with the specified content and saves it.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_html(html_content, filename)


@tool
def create_json(json_content: str, filename: str) -> str:
    """
    Creates a valid JSON file with the specified content and saves it.
    Returns the reference format: 'predefined_seed:<filename>'
    """
    return seed_generators.create_json(json_content, filename)



def _validate_seeds(parsed_seeds, properties, tool_name=""):
    """Drop seed objects that cannot possibly be arguments for this tool.

    BUG 5 (found 2026-09-10, appium generate-double-tap-code). The model
    answered the final turn with an array of its OWN tool-call envelopes
    rather than an array of tool arguments:

        {"name": "create_json",
         "parameters": {"filename": "test_case.json",
                        "json_content": "{\"language\": \"java\", ...}"}}

    Those top-level keys are "name" and "parameters", which match ZERO
    properties of the tool schema. input_parser._try_parse_json requires
    matched > 0, so parse() correctly falls through to FuzzedDataProvider
    and the envelope's own bytes come back through
    ConsumeUnicodeNoSurrogates as UTF-16 mojibake -- the exact artefact the
    parse-order fix was meant to eliminate. libFuzzer then collapsed all 12
    near-identical envelopes to corp: 1.

    The stage still logged LLM_SEEDGEN_OK n_seeds=12, i.e. a healthy marker
    over a dead stage: the same failure family as the four before it.

    A seed is kept iff it is a dict sharing at least one key with the
    schema's declared properties -- deliberately the SAME predicate
    input_parser uses, so "accepted here" means "honoured there". With no
    declared properties there is nothing to check against, so everything is
    kept and the caller is told validation was skipped.

    Returns (kept, rejected, checked).
    """
    if not isinstance(parsed_seeds, list):
        return [], [], False
    if not properties:
        return [s for s in parsed_seeds if isinstance(s, dict)], [], False
    keys = set(properties)
    kept, rejected = [], []
    for s in parsed_seeds:
        if isinstance(s, dict) and keys & set(s):
            kept.append(s)
        else:
            rejected.append(s)
    return kept, rejected, True


# SANITIZER -> ADVERSARIAL INPUT CLASS (the paper's G, sec4_design.tex:158-163)
#
# The design says: "Values for the selected arguments are generated using the
# corresponding sanitizer's adversarial input classes", and the thesis sentence
# at sec4_design.tex:10 says the property EXERCISED must be the property JUDGED.
#
# That link did not exist. generate_seeds() never received the selected
# sanitizers, and _local_fallback_mapping picked a seed-bank category by keyword
# match on the ARGUMENT NAME alone -- so the property exercised was chosen by one
# heuristic and the property judged by a separate LLM call, and they could
# disagree silently. Measured 2026-09-10 on excel-reader::get_cell: the LLM
# selected check_path_traversal, the seeds were generic, and check_path_traversal
# never fired. RQ4 asks exactly whether mutation reaches the code paths the
# SELECTED sanitizers are meant to exercise; nothing connected the two.
#
# Substring-matched, longest first, so check_introspection_sql_injection maps to
# sql_injection rather than falling through. Sanitizers with no category are
# COUNTED AND REPORTED rather than silently ignored: how many selected sanitizers
# actually get adversarial inputs is itself an RQ1/RQ4 number, and with 6 seed
# categories against 71 registered sanitizers it is not going to be all of them.
SANITIZER_SEED_CATEGORIES = [
    ("sql",                 "sql_injection"),
    ("xxe",                 "xml_bomb"),
    ("xml",                 "xml_bomb"),
    ("amplification",       "xml_bomb"),
    ("ssrf",                "ssrf_url"),
    ("url",                 "ssrf_url"),
    ("uri",                 "ssrf_url"),
    ("network",             "ssrf_url"),
    ("path_traversal",      "path_traversal"),
    ("jailbreak",           "path_traversal"),
    ("file_access",         "path_traversal"),
    ("sensitive_file",      "path_traversal"),
    ("directory",           "path_traversal"),
    ("symlink",             "path_traversal"),
    ("xss",                 "xss_markdown"),
    ("markdown",            "xss_markdown"),
    ("html",                "xss_markdown"),
    ("redos",               "wildcards"),
    ("regex",               "wildcards"),
    ("wildcard",            "wildcards"),
    ("glob",                "wildcards"),
]


def categories_for_sanitizers(selected, seed_bank):
    """Adversarial input classes for the sanitizers the LLM actually selected.

    Returns (categories, mapped, unmapped) so the caller can report how much of
    the selection got adversarial inputs -- never silently drop the remainder.
    """
    cats, mapped, unmapped = [], [], []
    for name in (selected or []):
        low = name.lower()
        hit = next((c for frag, c in SANITIZER_SEED_CATEGORIES if frag in low), None)
        if hit and hit in (seed_bank or {}):
            mapped.append(name)
            if hit not in cats:
                cats.append(hit)
        else:
            unmapped.append(name)
    return cats, mapped, unmapped


class LLMHandler:
    """LLM backend for sanitizer selection and seed generation.

    Backend is configurable via environment so the same code runs against
    Ollama Cloud or a LOCAL Ollama daemon. Local Ollama serves an
    OpenAI-compatible API at http://localhost:11434/v1 and ignores the key.

        export TOFU_LLM_BASE_URL=http://localhost:11434/v1
        export TOFU_LLM_MODEL=llama3.1:8b
        export OLLAMA_CLOUD_API_KEY=ollama     # any non-empty string locally

    Record whichever model you used in the paper's setup subsection: sanitizer
    selection quality is model-dependent, and VIPER-MCP's own RQ4 showed
    triggered-server counts swinging from 0 to 88 purely by changing backend.
    """
    def __init__(self, api_key: str = None, model_name: str = None, base_url: str = None):
        base_url = base_url or os.environ.get("TOFU_LLM_BASE_URL", "https://ollama.com/v1")
        model_name = model_name or os.environ.get("TOFU_LLM_MODEL", "minimax-m3:cloud")
        # OPENAI_API_KEY first when the endpoint is OpenAI's, so the hosted arm
        # needs no other change. Local Ollama ignores the value entirely.
        api_key = (api_key
                   or (os.environ.get("OPENAI_API_KEY")
                       if "openai.com" in (base_url or "") else None)
                   or os.environ.get("OLLAMA_CLOUD_API_KEY")
                   or os.environ.get("OPENAI_API_KEY")
                   or "ollama")
        timeout = int(os.environ.get("TOFU_LLM_TIMEOUT", "120"))
        self.model_name = model_name
        self.base_url = base_url
        # How many times to re-ask when a reply will not parse as JSON, and
        # whether to request the backend's JSON mode. Both are worth
        # recording in the paper's setup: they change how often the LLM
        # stage runs at all on a small local model.
        self.max_json_attempts = int(os.environ.get("TOFU_LLM_JSON_ATTEMPTS", "3"))
        self.json_mode = os.environ.get("TOFU_LLM_JSON_MODE", "1").lower() not in ("0", "false", "no")
        # Optional replay cache. When TOFU_LLM_CACHE_DIR is set, select_sanitizers
        # and generate_seeds return the recorded response for TOFU_LLM_CACHE_KEY
        # (one file per config) instead of calling the model, and write-through on
        # a miss. A re-run of an already-fuzzed tool then costs ZERO LLM calls; a
        # tool with no cache entry still calls the model and records it. The cache
        # is populated live during a run, or retroactively from run logs + corpus
        # seeds by eval/gen_llm_cache.py. Determinism is unaffected: the cached
        # response is exactly what temperature=0/seed produced.
        self.cache_dir = os.environ.get("TOFU_LLM_CACHE_DIR") or None
        self.cache_key = os.environ.get("TOFU_LLM_CACHE_KEY") or None
        if ChatOpenAI is None:
            raise RuntimeError(
                "langchain_openai is not importable, so no LLM backend can be "
                "constructed (%s). Install the pinned requirements into the "
                "active interpreter before running the campaign." % (_LANGCHAIN_IMPORT_ERROR,)
            )
        # DETERMINISM. This used to be ChatOpenAI(...) with no temperature, which
        # means LangChain's default of 0.7 -- so every sanitizer selection was a
        # SAMPLE. Two identical runs of the same tool could select different
        # sanitizers, no run was reproducible, and any comparison between LLM
        # backends was confounded by sampling noise rather than measuring the
        # backend. Model choice is an experimental variable in this paper, so the
        # sampling has to be pinned before any of those numbers mean anything.
        self.temperature = float(os.environ.get("TOFU_LLM_TEMPERATURE", "0"))
        self.seed = int(os.environ.get("TOFU_LLM_SEED", "20260910"))
        logger.info(f"LLM backend: model={model_name} base_url={base_url} "
                    f"json_mode={self.json_mode} json_attempts={self.max_json_attempts} "
                    f"temperature={self.temperature} seed={self.seed}")

        # REASONING. gemma4:12b is a thinking model (`ollama show` lists the
        # "thinking" capability) and with no cap it will reason until it decides
        # to stop. Measured on the M2 Pro at 17.5 generated tok/s, so ~12k
        # reasoning tokens is ~11 minutes for ONE call -- the 2026-09-10 smoke
        # run hung past the harness's 720s ceiling without completing a single
        # HTTP request.
        #
        # This was MASKED before: Ollama defaulted gemma4 to a 4096-token
        # context on 16 GB, which truncated the thinking at ~2045 tokens and
        # produced "length limit was reached" instead of a hang. Raising the
        # context to 32768 to fix the truncation is what let the thinking run
        # free. The two settings have to move together.
        #
        # Measured against the live daemon, same prompt:
        #   default (no cap)        hangs > 720s
        #   max_tokens=800           4s, 78 completion tokens, correct JSON
        #   reasoning_effort="none"  1s, 19 completion tokens, correct JSON
        #
        # TOFU_LLM_REASONING is an EXPERIMENTAL VARIABLE, not a tuning knob:
        # it changes how the model selects sanitizers, which is exactly what
        # RQ2/RQ3 measure. Record it with every run and state it in the setup
        # subsection alongside model, temperature and seed. "none" is the
        # campaign default; re-run a subset with it unset for the thinking arm.
        self.reasoning = os.environ.get("TOFU_LLM_REASONING", "none")
        # Belt and braces: even with reasoning off, never let one call run away
        # again. 4096 is ~4 minutes of generation at the measured rate.
        self.max_tokens = int(os.environ.get("TOFU_LLM_MAX_TOKENS", "4096"))
        logger.info(f"LLM backend: reasoning_effort={self.reasoning} "
                    f"max_tokens={self.max_tokens}")

        kwargs = dict(api_key=api_key, base_url=base_url, model=model_name,
                      timeout=timeout, max_retries=2, temperature=self.temperature,
                      max_tokens=self.max_tokens)
        extra = {"seed": self.seed}
        if self.reasoning and self.reasoning.lower() not in ("", "default", "auto"):
            extra["reasoning_effort"] = self.reasoning
        # Some backends reject reasoning_effort outright -- OpenAI's non-reasoning
        # models answer 400 "Unrecognized request argument supplied:
        # reasoning_effort" -- and that 400 arrives at INVOKE time, not at
        # construction, so it would kill every call of a campaign rather than
        # failing fast. Probe once here and drop the argument if unsupported.
        self._reasoning_sent = "reasoning_effort" in extra
        self._ctor = dict(api_key=api_key, base_url=base_url, model=model_name,
                          timeout=timeout)
        try:
            # langchain-openai exposes these directly on recent versions.
            self.model = ChatOpenAI(**extra, **kwargs)
        except TypeError:
            # Older versions: pass them through to the backend request body.
            # Ollama's OpenAI-compatible endpoint honours both.
            self.model = ChatOpenAI(model_kwargs=extra, **kwargs)
            logger.info("LLM backend: seed/reasoning passed via model_kwargs "
                        "(older langchain-openai)")

        # temperature=0 is necessary but NOT sufficient for bit-identical output:
        # llama.cpp/Ollama batching and GPU non-determinism can still perturb
        # logits. Record the seed with every run and re-run a tool twice to
        # confirm stability before claiming reproducibility in the paper.

    def _load_seed_bank(self) -> dict:
        """Loads all JSON files from src/seed_bank/ into a dictionary."""
        seed_bank = {}
        bank_dir = os.path.join(os.path.dirname(__file__), "seed_bank")
        if not os.path.exists(bank_dir):
            return seed_bank
            
        for file in os.listdir(bank_dir):
            if file.endswith(".json"):
                try:
                    with open(os.path.join(bank_dir, file), "r") as f:
                        data = json.load(f)
                    category = data.get("category")
                    if category:
                        seed_bank[category] = {
                            "description": data.get("description", ""),
                            "payloads": data.get("payloads", [])
                        }
                except Exception as e:
                    logging.warning(f"Failed to load seed bank file {file}: {e}")
        return seed_bank

    def _local_fallback_mapping(self, input_schema: dict, seed_bank: dict,
                                selected_sanitizers: list = None) -> list:
        """Generates fuzzing seeds locally without LLM API supporting type fuzzing and type confusion."""
        logging.warning("LLM_SEEDGEN_FALLBACK deterministic local mapping in use (no LLM seeds)")
        # Adversarial classes come from the SELECTED sanitizers (the paper's G),
        # not from a keyword match on the argument name. The name still decides
        # WHICH argument carries the payload; the sanitizer decides WHAT payload.
        wanted_cats, _m, _u = categories_for_sanitizers(selected_sanitizers, seed_bank)
        logging.info("[Seed Gen] Using local fallback mapping mechanism with multi-type fuzzing...")
        properties = input_schema.get("properties", {})
        
        # 1. Map properties to default values based on type.
        #
        # ENUMS. This mapper ignored "enum" entirely and filled every string
        # with "test". Measured 2026-09-10 on
        # AppiumTestDistribution/mcp-appium-gestures::generate-double-tap-code,
        # whose `language` is enum ['javascript','java']: 34 of the 37 seeds
        # carried language="test" and the server rejected EVERY ONE with
        #   MCP error -32602: Invalid ... parameters [{"received":"test", ...}]
        # So the tool never executed, no sanitizer could fire, and the run
        # reported a healthy 0 findings. Any tool with an enum-constrained
        # required field was 100% unfuzzable through the deterministic path.
        # (input_parser._generate_value_for_schema already honoured enums, so
        # the FDP path and the seed path disagreed with each other.)
        #
        # The default must be a VALID enum member, otherwise the payloads in
        # every OTHER field never reach the tool body. One deliberate
        # enum-violation seed is still emitted below -- rejecting an invalid
        # enum is legitimate behaviour worth exercising, just not 34 times.
        defaults = {}
        for prop, attrs in properties.items():
            p_type = attrs.get("type", "string")
            enum_vals = attrs.get("enum") or []
            if enum_vals:
                defaults[prop] = enum_vals[0]
            elif p_type in ("integer", "number"):
                defaults[prop] = 1
            elif p_type == "boolean":
                defaults[prop] = True
            elif p_type == "array":
                defaults[prop] = []
            elif p_type == "object":
                defaults[prop] = {}
            else:
                defaults[prop] = "test"
                
        seeds = []
        
        # 2. Generate type-specific boundary values
        for prop, attrs in properties.items():
            p_type = attrs.get("type", "string")

            enum_vals = attrs.get("enum") or []
            if enum_vals:
                # Every declared member, so the tool body is reached through
                # each enum branch, plus exactly ONE out-of-enum probe.
                for val in list(enum_vals) + ["__not_in_enum__"]:
                    seed = defaults.copy()
                    seed[prop] = val
                    seeds.append(seed)
                continue

            if p_type in ("integer", "number"):
                # JSON-compliant boundary numbers
                boundaries = [0, -1, 2147483647, -2147483648, 1e308, -1e308, None]
                for val in boundaries:
                    seed = defaults.copy()
                    seed[prop] = val
                    seeds.append(seed)
            elif p_type == "boolean":
                # Alternate boolean values
                for val in [True, False]:
                    seed = defaults.copy()
                    seed[prop] = val
                    seeds.append(seed)
            elif p_type == "array":
                # Array boundaries and malformations
                variations = [
                    [],
                    [None],
                    ["test"],
                    [1, 2, 3],
                    ["text", 123, None, {}]
                ]
                for val in variations:
                    seed = defaults.copy()
                    seed[prop] = val
                    seeds.append(seed)
            elif p_type == "object":
                # Object boundaries
                variations = [
                    {},
                    {"nested": "val"},
                    {"depth1": {"depth2": {"depth3": "deep"}}}
                ]
                for val in variations:
                    seed = defaults.copy()
                    seed[prop] = val
                    seeds.append(seed)
            elif p_type == "string":
                # Check keywords and inject from seed bank
                prop_lower = prop.lower()
                desc_lower = attrs.get("description", "").lower()
                
                category_to_use = None
                if "sql" in prop_lower or "query" in prop_lower or "sql" in desc_lower or "query" in desc_lower:
                    category_to_use = "sql_injection"
                elif "url" in prop_lower or "uri" in prop_lower or "link" in prop_lower or "url" in desc_lower or "uri" in desc_lower or "link" in desc_lower:
                    category_to_use = "ssrf_url"
                elif "path" in prop_lower or "file" in prop_lower or "dir" in prop_lower or "path" in desc_lower or "file" in desc_lower or "dir" in desc_lower:
                    category_to_use = "path_traversal"
                elif "pattern" in prop_lower or "search" in prop_lower or "like" in prop_lower or "pattern" in desc_lower or "search" in desc_lower or "like" in desc_lower:
                    category_to_use = "wildcards"
                elif "html" in prop_lower or "xml" in prop_lower or "xml" in desc_lower or "html" in desc_lower or "markdown" in desc_lower:
                    category_to_use = "xss_markdown"
                    
                # Sanitizer-derived classes win. If the selection asked for a
                # class the name-heuristic did not pick, still exercise it --
                # otherwise the selected sanitizer has no inputs to fire on.
                if wanted_cats:
                    category_to_use = (category_to_use if category_to_use in wanted_cats
                                       else wanted_cats[0])

                if category_to_use and category_to_use in seed_bank:
                    category_data = seed_bank[category_to_use]
                    for payload in category_data["payloads"][:5]:
                        seed = defaults.copy()
                        seed[prop] = payload
                        seeds.append(seed)
                else:
                    # Fallback string injections
                    for cat in ["path_traversal", "sql_injection"]:
                        if cat in seed_bank:
                            for payload in seed_bank[cat]["payloads"][:2]:
                                seed = defaults.copy()
                                seed[prop] = payload
                                seeds.append(seed)
                                
        # 3. Generate Type Confusion seeds (deliberate type mismatches)
        for prop, attrs in properties.items():
            p_type = attrs.get("type", "string")
            
            confusion_vals = []
            if p_type == "string":
                confusion_vals = [123, True, ["not_a_string"]]
            elif p_type in ("integer", "number"):
                confusion_vals = ["not_a_number", [1]]
            elif p_type == "boolean":
                confusion_vals = ["true", 1]
            elif p_type in ("array", "object"):
                confusion_vals = ["not_a_structure", 99]
                
            for val in confusion_vals:
                seed = defaults.copy()
                seed[prop] = val
                seeds.append(seed)

        # Add basic clean defaults
        seeds.append(defaults)
        
        # Deduplicate seeds
        unique_seeds = []
        seen = []
        for s in seeds:
            s_str = json.dumps(s, sort_keys=True)
            if s_str not in seen:
                seen.append(s_str)
                unique_seeds.append(s)
                
        return unique_seeds

    # ------------------------------------------------------------------ cache
    def _cache_path(self):
        if not (self.cache_dir and self.cache_key):
            return None
        safe = re.sub(r"[^A-Za-z0-9_.-]", "_", self.cache_key)
        return os.path.join(self.cache_dir, safe + ".json")

    def _cache_load(self):
        p = self._cache_path()
        if p and os.path.exists(p):
            try:
                with open(p) as f:
                    return json.load(f)
            except Exception:
                return None
        return None

    def _cache_update(self, **kv):
        """Merge keys into this config's cache file, atomically. Best-effort."""
        p = self._cache_path()
        if not p or _cache_readonly():
            return
        try:
            os.makedirs(os.path.dirname(p), exist_ok=True)
            cur = self._cache_load() or {}
            cur.update(kv)
            tmp = p + ".tmp"
            with open(tmp, "w") as f:
                json.dump(cur, f)
            os.replace(tmp, p)
        except Exception as e:
            logger.debug(f"llm cache write failed for {p}: {e}")

    def generate_seeds(self, tool_name: str, tool_description: str, input_schema: dict,
                       selected_sanitizers: list = None) -> list:
        """Generate fuzzing seeds by letting the LLM select and adapt payloads from a seed bank."""
        if _no_llm_seeds():
            logging.info(f"NO_LLM seed-gen (fallback mapping) tool={tool_name}")
            return self._local_fallback_mapping(input_schema, self._load_seed_bank(), selected_sanitizers)
        cached = self._cache_load()
        if (cached is not None and cached.get("seeds") is not None and _require_llm_seeds()
                and cached.get("seeds_provenance") != "llm"):
            # strict mode: a fallback or unproven entry is not LLM evidence -- regenerate
            logging.info(f"LLM_SEEDGEN_CACHE_SKIP tool={tool_name} key={self.cache_key} "
                         f"provenance={cached.get('seeds_provenance') or 'unknown'} "
                         f"reason=TOFU_REQUIRE_LLM_SEEDS")
            cached = None
        if cached is not None and cached.get("seeds") is not None:
            # provenance: WHO produced these seeds. eval/gen_llm_cache.py
            # retro-builds entries from corpus seed_*.json, which the fuzzer
            # writes for fallback seeds too -- so a hit alone is not evidence
            # that the LLM generated them (2026-09-23: >=12 of 5022 VM entries
            # replay fallback seeds). run_eval.parse_log reads this field.
            logging.info(f"LLM_SEEDGEN_CACHE_HIT tool={tool_name} "
                         f"n_seeds={len(cached['seeds'])} key={self.cache_key} "
                         f"provenance={cached.get('seeds_provenance') or 'unknown'}")
            return cached["seeds"]
        logging.info(f"[LLM] Generation of smart seed corpus for '{tool_name}' using Seed Bank...")
        
        seed_bank = self._load_seed_bank()
        if not seed_bank:
            if _require_llm_seeds():
                raise SeedProvenanceError(f"seed bank empty for '{tool_name}'; refusing the local "
                                          f"fallback because TOFU_REQUIRE_LLM_SEEDS is set")
            logging.warning("Seed bank is empty or not found. Falling back to local mapping.")
            return self._local_fallback_mapping(input_schema, {}, selected_sanitizers)
            
        example_structure = {}
        properties = input_schema.get("properties", {})
        for prop, attrs in properties.items():
            p_type = attrs.get("type", "string")
            example_structure[prop] = f"<{p_type}>"
            
        # Format the seed bank categories for the prompt. Restrict them to the
        # classes the SELECTED sanitizers actually need (the paper's G), so the
        # property exercised is the property judged. Offering the whole bank
        # regardless of selection is what decoupled the two.
        wanted_cats, mapped, unmapped = categories_for_sanitizers(selected_sanitizers, seed_bank)
        if wanted_cats:
            logging.info(
                f"SEEDGEN_CLASSES tool={tool_name} categories={wanted_cats} "
                f"n_sanitizers_mapped={len(mapped)} n_unmapped={len(unmapped)}")
            if unmapped:
                # Honest accounting: with 6 seed-bank categories against 71
                # registered sanitizers, some selections get no adversarial
                # inputs at all. That number belongs in RQ1/RQ4, not in a
                # silent drop.
                logging.warning(
                    f"SEEDGEN_UNMAPPED tool={tool_name} n={len(unmapped)} "
                    f"sanitizers={unmapped} -- no adversarial input class exists "
                    f"for these; their oracle can only fire incidentally")
        bank_for_prompt = {c: seed_bank[c] for c in wanted_cats} if wanted_cats else seed_bank

        seed_bank_desc = ""
        for cat, data in bank_for_prompt.items():
            seed_bank_desc += f"- Category: '{cat}'\n  Description: {data['description']}\n  Sample Payloads: {json.dumps(data['payloads'][:3])}\n\n"

        selected_desc = ""
        if selected_sanitizers:
            selected_desc = (
                "\nThe following security sanitizers have been selected for this tool. "
                "Your test cases MUST exercise the properties they check:\n"
                + "\n".join(f"- {n}" for n in selected_sanitizers) + "\n")
            
        prompt = f"""
You are a security and QA engineering expert. You will generate a JSON array of test cases (inputs) to fuzz the following tool.

Tool Name: {tool_name}
Description: {tool_description}
Schema Contract: {input_schema}
{selected_desc}
We have a structured Seed Bank containing pre-curated payloads categorized by attack vectors:
{seed_bank_desc}

We also have predefined file generator tools that you can call to create valid binary/text files:
- create_pdf
- create_docx
- create_xlsx
- create_pptx
- create_zip
- create_xml
- create_html
- create_json

You should call these file generator tools if you want to test the tool with high-quality custom documents. Each tool will save the file and return a string like "predefined_seed:<filename>".
Use the exact returned string as the value for the file/URI parameter in the test case.

Strict Instructions:
1. Select the relevant Seed Bank categories for fuzzing this tool.
2. Adapt the payloads from those selected categories into the parameters of the tool schema.
3. For non-string parameters:
   - For integers/numbers, include test cases with boundary values (e.g. 0, -1, large numbers).
   - For arrays, include empty arrays [] and heterogeneous/malformed arrays.
   - For booleans, test both true and false.
4. Intentionally include a few type confusion/type mismatch test cases (e.g., passing a string to an integer, or an array to a string parameter) to check type-handling robustness.
5. Every object in the array MUST strictly follow this argument structure: {json.dumps(example_structure)}
6. Return ONLY a valid raw JSON array containing 6 to 12 objects. No explanations, no markdown blocks. Just plain raw JSON.
"""
        messages = [HumanMessage(content=prompt)]
        
        tools = [
            create_pdf,
            create_docx,
            create_xlsx,
            create_pptx,
            create_zip,
            create_xml,
            create_html,
            create_json
        ]
        tools_map = {t.name: t for t in tools}
        model_with_tools = self.model.bind_tools(tools)
        
        max_attempts = 4
        attempt = 1
        
        while attempt <= max_attempts:
            try:
                response = model_with_tools.invoke(messages)
                messages.append(response)
                
                if response.tool_calls:
                    for tool_call in response.tool_calls:
                        t_name = tool_call["name"]
                        t_args = tool_call["args"]
                        t_id = tool_call["id"]
                        
                        logging.info(f"[LLM Seed Gen Tool Call] Invoking {t_name} with {t_args}")
                        if t_name in tools_map:
                            try:
                                result = tools_map[t_name].invoke(t_args)
                            except Exception as e:
                                result = f"Error running tool {t_name}: {e}"
                        else:
                            result = f"Error: Tool {t_name} not found."
                        
                        logging.info(f"[LLM Seed Gen Tool Call] Result: {result}")
                        messages.append(ToolMessage(content=result, tool_call_id=t_id))
                    continue
                
                content = _coerce_text(response.content).strip()
                if not content:
                    messages.append(HumanMessage(content="Please provide the final raw JSON array of test cases now."))
                    attempt += 1
                    continue
                
                pinfo = {}
                parsed_seeds = _extract_json(content, expect=list, repair=True, info=pinfo)

                # Validate BEFORE declaring success. See _validate_seeds.
                kept, rejected, checked = _validate_seeds(
                    parsed_seeds, properties, tool_name)
                if checked and rejected:
                    logging.warning(
                        f"LLM_SEEDGEN_REJECTED tool={tool_name} attempt={attempt} "
                        f"n_rejected={len(rejected)} n_kept={len(kept)} "
                        f"schema_keys={sorted(properties)} "
                        f"first_rejected_keys={sorted(rejected[0]) if isinstance(rejected[0], dict) else type(rejected[0]).__name__}")
                if not kept:
                    raise ValueError(
                        "none of the %d returned objects shares a key with the "
                        "tool schema %s -- you returned tool-call envelopes or "
                        "unrelated objects, not tool arguments"
                        % (len(parsed_seeds), sorted(properties)))

                logging.info(
                    f"LLM_SEEDGEN_OK tool={tool_name} attempt={attempt} "
                    f"n_seeds={len(kept)} n_rejected={len(rejected)} "
                    f"validated={int(checked)} repaired={pinfo.get('repaired', 'no')}")
                self._cache_update(seeds=kept, seeds_provenance="llm")
                return kept
                
            except Exception as e:
                logging.warning(f"Seed generation attempt {attempt} failed ({e}).")
                if "no parsable JSON" in str(e) and locals().get("content"):
                    # keep the evidence: what the model actually wrote
                    logging.warning(f"SEEDGEN_RAW_REPLY tool={tool_name} attempt={attempt} "
                                    f"len={len(content)} reply={content[:2000]!r}")
                if attempt == max_attempts:
                    break
                messages.append(HumanMessage(content=f"Your previous output was invalid: {str(e)}. Please correct it and output ONLY the valid raw JSON array."))
                attempt += 1
                
        # Trigger local fallback mapping if LLM failed
        if _require_llm_seeds():
            raise SeedProvenanceError(f"LLM seed generation failed for '{tool_name}' after "
                                      f"{max_attempts} attempts; refusing the local fallback "
                                      f"because TOFU_REQUIRE_LLM_SEEDS is set")
        return self._local_fallback_mapping(input_schema, seed_bank, selected_sanitizers)



    def select_sanitizers(self, tool_name: str, tool_description: str, input_schema: dict, available_sanitizers: dict) -> tuple:
        """Selects sanitizing functions and their behavioral rules based on docstrings/schemas."""
        logging.info(f"Selecting relevant sanitizers for '{tool_name}'...")
        mode = _no_llm_sanitizers()
        if mode == "mandatory":
            from .sanitizers import BASE_MANDATORY_SANITIZERS
            sel = [x for x in BASE_MANDATORY_SANITIZERS if x in available_sanitizers]
            logging.info(f"NO_LLM selection (mandatory only) tool={tool_name} n_selected={len(sel)}")
            return sel, {}
        if mode == "all":
            sel = list(available_sanitizers.keys())
            logging.info(f"NO_LLM selection (all sanitizers) tool={tool_name} n_selected={len(sel)}")
            return sel, {}
        cached = self._cache_load()
        if cached is not None and cached.get("selected_sanitizers") is not None:
            sel = [x for x in cached["selected_sanitizers"] if x in available_sanitizers]
            rules = {k: v for k, v in (cached.get("rules") or {}).items() if k in sel}
            logger.info(f"LLM_SELECTION_CACHE_HIT tool={tool_name} "
                        f"n_selected={len(sel)} key={self.cache_key}")
            return sel, rules

        sanitizers_format = "\n".join([f"- {name}: {desc}" for name, desc in available_sanitizers.items()])

        prompt = f"""
You are an expert security engineer. Analyze the following tool metadata and select which sanitizers from the provided list are relevant to monitor its inputs. Also, analyze the tool description and docstring to determine what the tool's expected behavior is if a sanitizer triggers.

Tool Name: {tool_name}
Description: {tool_description}
Schema: {input_schema}

Available Sanitizers:
{sanitizers_format}

Strict Instructions:
1. Return ONLY a valid JSON object with two keys:
   - "selected_sanitizers": A JSON array of strings containing the exact names of the selected sanitizers.
   - "rules": A JSON object mapping a selected sanitizer name to its expected behavior rule:
     {{
       "sanitizer_name": {{
         "on_trigger": "allow_success_always" or "allow_success_with_fallback" or "raise_exception",
         "expected_fallback_values": [<expected return values when sanitizer is triggered, e.g. [], "", null, etc.>]
       }}
     }}
     Use "allow_success_always" if the tool is designed to process arbitrary text that might contain invalid elements (like a JSON extractor or lenient parser) and should always succeed normally with any output.
     Use "allow_success_with_fallback" if the tool expects a specific valid structure but gracefully returns a fallback value when invalid.
     Use "raise_exception" if the tool should strictly raise an exception/error on invalid input.
2. If no sanitizers are relevant, return: {{"selected_sanitizers": [], "rules": {{}}}}
3. No code blocks, no markdown formatting. Just raw JSON.
"""
        messages = [HumanMessage(content=prompt)]
        
        # A small local model will happily wrap the object in prose and a
        # markdown fence no matter how the instruction is phrased, so ask
        # for the backend's JSON mode where available, parse tolerantly,
        # and re-ask with a sharper instruction before giving up. The
        # outcome is logged on one greppable line per tool: the fraction
        # of tools where the LLM stage actually ran is a number the
        # evaluation has to report, not assume.
        model = self.model
        json_mode = self.json_mode
        if json_mode:
            try:
                model = self.model.bind(response_format={"type": "json_object"})
            except Exception as e:
                json_mode = False
                logger.debug(f"backend rejected json_object mode ({e}); continuing without it")

        last_err = None
        for attempt in range(1, self.max_json_attempts + 1):
            p = prompt
            if attempt > 1:
                p = prompt + (
                    f"\n\nYour previous reply could not be parsed ({last_err}). "
                    "Reply with the JSON object and nothing else: no explanation, "
                    "no markdown fence, no text before or after it. The first "
                    "character of your reply must be '{' and the last must be '}'."
                )
            try:
                response = model.invoke([HumanMessage(content=p)])
                logging.info(_coerce_text(response.content))
                data = _extract_json(response.content, expect=dict)
                selected = data.get("selected_sanitizers", [])
                rules = data.get("rules", {})
                if not isinstance(selected, list) or not isinstance(rules, dict):
                    raise ValueError(
                        "expected 'selected_sanitizers' to be a list and 'rules' a "
                        f"object, got {type(selected).__name__}/{type(rules).__name__}"
                    )
                valid_selection = [x for x in selected if x in available_sanitizers]
                dropped = [x for x in selected if x not in available_sanitizers]
                valid_rules = {x: r for x, r in rules.items() if x in valid_selection}
                if dropped:
                    logger.warning(
                        f"[{tool_name}] LLM named {len(dropped)} sanitizer(s) that are "
                        f"not in the registry; dropped: {dropped}"
                    )
                logger.info(
                    f"LLM_SELECTION_OK tool={tool_name} attempt={attempt} "
                    f"n_selected={len(valid_selection)} n_dropped={len(dropped)}"
                )
                logging.info(f"Sanitizers selected: {valid_selection}")
                logging.info(f"Sanitizer rules: {valid_rules}")
                self._cache_update(selected_sanitizers=valid_selection, rules=valid_rules)
                return valid_selection, valid_rules
            except Exception as e:
                last_err = e
                logger.warning(
                    f"[{tool_name}] sanitizer-selection attempt "
                    f"{attempt}/{self.max_json_attempts} failed: {e}"
                )
                if self._reasoning_sent and "reasoning_effort" in str(e):
                    # Backend has no reasoning knob. Rebuild without it rather
                    # than burning every remaining attempt on the same 400.
                    self._reasoning_sent = False
                    kw = dict(**self._ctor, max_retries=2,
                              temperature=self.temperature, max_tokens=self.max_tokens)
                    try:
                        self.model = ChatOpenAI(seed=self.seed, **kw)
                    except TypeError:
                        self.model = ChatOpenAI(model_kwargs={"seed": self.seed}, **kw)
                    model = self.model
                    logger.warning(f"[{tool_name}] backend rejected reasoning_effort; "
                                   f"rebuilt the client without it")
                if json_mode:
                    # A backend that rejects response_format only says so at
                    # invoke time, which would otherwise burn every attempt
                    # on the same error. Drop the constraint and retry plain.
                    json_mode = False
                    model = self.model
                    logger.warning(
                        f"[{tool_name}] retrying without json_object mode"
                    )

        logger.error(f"LLM_SELECTION_FAILED tool={tool_name} after "
                     f"{self.max_json_attempts} attempts: {last_err}")
        raise RuntimeError(
            f"LLM API Call Failed during sanitizer selection for '{tool_name}': {last_err}"
        ) from last_err