monkeybot 2.1.1__py3-none-any.whl

monkeybot/core/context/tool_result_ingress.py

@@ -0,0 +1,290 @@
+"""Sanitize, cap, and redact tool/MCP payloads before they enter chat history."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_TOOL_RESULT_MAX_CHARS = 32_768
+_DEFAULT_SUMMARY_TOOL_RESULT_MAX_CHARS = 8_192
+# Only applies to denylisted blob keys (data/base64/…), not arbitrary text fields
+# like browser ``tree``. Whole-result size is handled by spill + TOOL_RESULT_MAX_CHARS.
+_DEFAULT_BLOB_JSON_FIELD_MAX_CHARS = 512
+_MIN_BASE64_RUN = 800
+
+_DATA_URL_RE = re.compile(
+    r"data:(?:image|audio|video|application)/[\w.+-]+;base64,[\sA-Za-z0-9+/=]{200,}",
+    re.IGNORECASE,
+)
+_LONG_B64_RE = re.compile(
+    r"(?<![A-Za-z0-9+/=])(?:[A-Za-z0-9+/]{4}){200,}(?:[A-Za-z0-9+/]{0,3}={0,3})(?![A-Za-z0-9+/=])"
+)
+
+_REDACT_JSON_KEYS = frozenset(
+    {
+        "data",
+        "base64",
+        "image_data",
+        "content_base64",
+        "blob",
+        "binary",
+        "image",
+        "screenshot",
+        "audio",
+        "bytes",
+        "payload",
+    }
+)
+
+# Workspace tools return faithful text/JSON; skip ingress JSON field redaction (see CoreToolExecutor).
+SANITIZE_SKIP_TOOL_NAMES = frozenset(
+    {
+        "read_file",
+        "write_file",
+        "replace_in_file",
+        "glob",
+        "search_memory",
+        "list_skills",
+    }
+)
+
+
+def skip_tool_result_sanitize(tool_name: str) -> bool:
+    """True when tool results must not pass through ``sanitize_tool_result_text``."""
+    return tool_name in SANITIZE_SKIP_TOOL_NAMES
+
+
+def _int_from_env(name: str, default: int) -> int:
+    raw = os.environ.get(name, "").strip()
+    if not raw:
+        return default
+    try:
+        return max(0, int(raw))
+    except ValueError:
+        logger.warning("invalid env var %s=%r; using default %s", name, raw, default)
+        return default
+
+
+def sanitize_enabled_from_env() -> bool:
+    raw = os.environ.get("MONKEYBOT_TOOL_RESULT_SANITIZE", "1").strip().lower()
+    return raw not in ("0", "false", "no", "off")
+
+
+def tool_result_max_chars_from_env() -> int:
+    return _int_from_env("MONKEYBOT_TOOL_RESULT_MAX_CHARS", _DEFAULT_TOOL_RESULT_MAX_CHARS)
+
+
+def summary_tool_result_max_chars_from_env() -> int:
+    return _int_from_env(
+        "MONKEYBOT_SUMMARY_TOOL_RESULT_MAX_CHARS",
+        _DEFAULT_SUMMARY_TOOL_RESULT_MAX_CHARS,
+    )
+
+
+def json_field_max_chars_from_env() -> int:
+    """Max length for *denylisted* JSON blob fields before omission.
+
+    Does **not** apply to ordinary text fields (``tree``, ``result``, etc.).
+    Env: ``MONKEYBOT_TOOL_RESULT_JSON_FIELD_MAX`` (legacy name kept).
+    """
+    return _int_from_env(
+        "MONKEYBOT_TOOL_RESULT_JSON_FIELD_MAX", _DEFAULT_BLOB_JSON_FIELD_MAX_CHARS
+    )
+
+
+def _looks_like_base64(value: str) -> bool:
+    """Return True when *value* is likely base64, including short JSON field payloads."""
+    if not value:
+        return False
+    if _is_plausible_base64_run(value):
+        return True
+    stripped = value.strip()
+    if len(stripped) < 16:
+        return False
+    if any(
+        ch not in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=\n\r"
+        for ch in stripped
+    ):
+        return False
+    sample = stripped[:512]
+    if len(set(sample.strip("="))) <= 2:
+        return False
+    if stripped.rstrip().endswith("=") or any(ch in "+/" for ch in sample):
+        return True
+    return len(set(sample)) >= 6
+
+
+def _is_plausible_base64_run(text: str) -> bool:
+    if len(text) < _MIN_BASE64_RUN:
+        return False
+    sample = text[:4096]
+    if len(set(sample.strip("="))) <= 2:
+        return False
+    if any(ch in "+/" for ch in sample):
+        return True
+    if sample.rstrip().endswith("="):
+        return True
+    allowed = sum(1 for ch in sample if ch.isalnum() or ch in "+/=\n\r")
+    return allowed / max(1, len(sample)) >= 0.98 and len(set(sample)) >= 8
+
+
+def _replace_long_b64(match: re.Match[str]) -> str:
+    run = match.group(0)
+    if not _is_plausible_base64_run(run):
+        return run
+    return f"[omitted ~{len(run)} char base64 run]"
+
+
+def _should_redact_blob_field(key: str, value: str, *, max_field: int) -> bool:
+    """True only for denylisted blob keys that look like binary/base64 or are oversized."""
+    if key.lower() not in _REDACT_JSON_KEYS:
+        return False
+    if _looks_like_base64(value):
+        return True
+    return max_field > 0 and len(value) > max_field
+
+
+def _redact_json_value(key: str, value: Any, *, max_field: int) -> Any:
+    if isinstance(value, str):
+        if _should_redact_blob_field(key, value, max_field=max_field):
+            return f"[omitted {len(value)} chars from {key!r}]"
+        return value
+    if isinstance(value, dict):
+        return {k: _redact_json_value(str(k), v, max_field=max_field) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_redact_json_value(key, item, max_field=max_field) for item in value]
+    return value
+
+
+def _redact_json_text(text: str, *, max_field: int) -> str:
+    stripped = text.lstrip()
+    if not stripped.startswith(("{", "[")):
+        return text
+    try:
+        parsed = json.loads(text)
+    except json.JSONDecodeError:
+        return text
+    redacted = _redact_json_value("", parsed, max_field=max_field)
+    try:
+        return json.dumps(redacted, ensure_ascii=False, default=str)
+    except TypeError:
+        return text
+
+
+def sanitize_tool_result_text(text: str, *, enabled: bool | None = None) -> str:
+    """Strip embedded binary/base64 payloads from tool result strings.
+
+    Ordinary long text fields (e.g. browser ``tree``) are left intact; size is
+    controlled later by spill (``MONKEYBOT_SPILL_MIN_CHARS``) and
+    ``MONKEYBOT_TOOL_RESULT_MAX_CHARS``. Only denylisted blob keys are
+    field-redacted here.
+    """
+    if not text:
+        return text
+    if enabled is None:
+        enabled = sanitize_enabled_from_env()
+    if not enabled:
+        return text
+
+    max_field = json_field_max_chars_from_env()
+    out = _DATA_URL_RE.sub(
+        lambda m: f"data:...;base64,[omitted {len(m.group(0))} chars]",
+        text,
+    )
+    out = _LONG_B64_RE.sub(_replace_long_b64, out)
+    return _redact_json_text(out, max_field=max_field)
+
+
+def truncate_with_note(text: str, max_chars: int, *, label: str = "tool result") -> str:
+    """Hard-cap ``text`` with a trailing omission note."""
+    if max_chars <= 0 or len(text) <= max_chars:
+        return text
+    omitted = len(text) - max_chars
+    return f"{text[:max_chars]}\n[{label} truncated — {omitted} chars omitted]"
+
+
+def cap_tool_result_text(text: str, *, max_chars: int | None = None) -> str:
+    """Apply the inline history hard ceiling after sanitization/spill."""
+    limit = tool_result_max_chars_from_env() if max_chars is None else max_chars
+    return truncate_with_note(text, limit, label="tool result")
+
+
+def summarize_tool_result_text(text: str, *, max_chars: int | None = None) -> str:
+    """Sanitize and cap text fed into context summarization."""
+    limit = summary_tool_result_max_chars_from_env() if max_chars is None else max_chars
+    cleaned = sanitize_tool_result_text(text)
+    return truncate_with_note(cleaned, limit, label="summarized tool result")
+
+
+def format_mcp_binary_block(
+    *,
+    kind: str,
+    data: object,
+    mime: str | None,
+) -> str:
+    """Placeholder for MCP image/audio/resource blocks."""
+    size = len(str(data or ""))
+    mime_label = mime or f"{kind}/*"
+    return (
+        f"[{kind} omitted from tool text: mime={mime_label}, "
+        f"~{size} base64 chars — use a vision-capable provider or a text-returning tool]"
+    )
+
+
+def dump_model_or_str(obj: Any) -> str:
+    """JSON-dump ``obj.model_dump()`` when available and serializable, else ``str(obj)``."""
+    md = getattr(obj, "model_dump", None)
+    if callable(md):
+        try:
+            dumped = md(mode="python", by_alias=True)
+            return json.dumps(dumped, default=str)
+        except TypeError as exc:
+            logger.debug("model_dump not JSON-serializable, falling back to str(): %s", exc)
+    return str(obj)
+
+
+def format_mcp_content_block(block: Any, *, sanitize: bool = True) -> str:
+    """Turn one MCP content block into a safe string."""
+    blk_type = getattr(block, "type", None)
+    if blk_type == "text":
+        txt = getattr(block, "text", None)
+        if txt is None:
+            return ""
+        text = str(txt)
+        return sanitize_tool_result_text(text) if sanitize else text
+    if blk_type == "image":
+        data = getattr(block, "data", None) or ""
+        mime = getattr(block, "mimeType", None) or getattr(block, "mime_type", None) or "image/png"
+        return format_mcp_binary_block(kind="image", data=data, mime=str(mime))
+    if blk_type == "audio":
+        data = getattr(block, "data", None) or ""
+        mime = getattr(block, "mimeType", None) or getattr(block, "mime_type", None) or "audio/*"
+        return format_mcp_binary_block(kind="audio", data=data, mime=str(mime))
+    if blk_type == "resource":
+        uri = getattr(block, "uri", None) or getattr(block, "resource", None) or ""
+        mime = getattr(block, "mimeType", None) or getattr(block, "mime_type", None) or "resource"
+        return f"[resource omitted from tool text: uri={uri!s}, mime={mime}]"
+
+    text = dump_model_or_str(block)
+    return sanitize_tool_result_text(text) if sanitize else text
+
+
+__all__ = [
+    "SANITIZE_SKIP_TOOL_NAMES",
+    "cap_tool_result_text",
+    "dump_model_or_str",
+    "format_mcp_binary_block",
+    "format_mcp_content_block",
+    "sanitize_enabled_from_env",
+    "sanitize_tool_result_text",
+    "skip_tool_result_sanitize",
+    "summarize_tool_result_text",
+    "summary_tool_result_max_chars_from_env",
+    "tool_result_max_chars_from_env",
+    "truncate_with_note",
+]

monkeybot/core/context/tool_shapers.py

@@ -0,0 +1,361 @@
+"""Content-aware shaping for large tool outputs (no ML, pure Python)."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+from collections.abc import Sequence
+from typing import Any, Literal
+
+from monkeybot.core.context.common import ContextPressureTier, text_from_blocks
+from monkeybot.core.context.tool_output_policy import ToolOutputBudget, resolve_tool_budget
+from monkeybot.core.llm.provider import Message
+from monkeybot.core.logging_utils import kv
+from monkeybot.core.types.content_blocks import ContentBlock, Text, ToolResponse
+
+logger = logging.getLogger(__name__)
+
+ContentType = Literal["json", "logs", "code", "prose"]
+
+_DEFAULT_LOG_HEAD_LINES = 40
+_DEFAULT_LOG_TAIL_LINES = 40
+_DEFAULT_MAX_ARRAY_ITEMS = 50
+_DEFAULT_ERROR_PATTERNS = (
+    r"(?i)\berror\b",
+    r"(?i)\bfatal\b",
+    r"(?i)traceback",
+    r"(?i)\bexception\b",
+    r"(?i)\bfailed\b",
+)
+_CODE_MARKERS = re.compile(
+    r"(^\s*(def |class |import |from .+ import |function |const |let |var |#include ))",
+    re.MULTILINE,
+)
+
+
+def _int_from_env(name: str, default: int) -> int:
+    raw = os.environ.get(name, "").strip()
+    if not raw:
+        return default
+    try:
+        return max(1, int(raw))
+    except ValueError:
+        logger.warning(
+            "invalid env var value %s",
+            kv(name=name, value=raw, default=default),
+        )
+        return default
+
+
+def log_head_lines_from_env() -> int:
+    return _int_from_env("MONKEYBOT_SHAPE_LOG_HEAD_LINES", _DEFAULT_LOG_HEAD_LINES)
+
+
+def log_tail_lines_from_env() -> int:
+    return _int_from_env("MONKEYBOT_SHAPE_LOG_TAIL_LINES", _DEFAULT_LOG_TAIL_LINES)
+
+
+def max_array_items_from_env() -> int:
+    return _int_from_env("MONKEYBOT_SHAPE_MAX_ARRAY_ITEMS", _DEFAULT_MAX_ARRAY_ITEMS)
+
+
+def classify_content(text: str, *, tool_name: str, hint: str | None = None) -> ContentType:
+    """Route tool text to a shaper category."""
+    if hint and hint != "auto":
+        return hint  # type: ignore[return-value]
+
+    stripped = text.lstrip()
+    if tool_name in ("run_command",):
+        return "logs"
+    if tool_name in ("web_search", "search_memory", "task") and stripped.startswith(("{", "[")):
+        return "json"
+    if tool_name in ("read_file", "write_file", "replace_in_file", "glob"):
+        return "code"
+
+    if stripped.startswith(("{", "[")):
+        try:
+            json.loads(text)
+            return "json"
+        except json.JSONDecodeError:
+            pass
+
+    lines = text.splitlines()
+    if len(lines) >= 8:
+        duplicate_run = sum(
+            1 for i in range(1, len(lines)) if lines[i] == lines[i - 1] and lines[i].strip()
+        )
+        if duplicate_run >= max(3, len(lines) // 10):
+            return "logs"
+        if any(_CODE_MARKERS.search(line) for line in lines[:20]):
+            return "code"
+
+    if len(lines) >= 30 and not stripped.startswith("{"):
+        return "logs"
+
+    return "prose"
+
+
+def _compiled_keep_patterns(patterns: tuple[str, ...]) -> tuple[re.Pattern[str], ...]:
+    src = patterns or _DEFAULT_ERROR_PATTERNS
+    out: list[re.Pattern[str]] = []
+    for pat in src:
+        try:
+            out.append(re.compile(pat))
+        except re.error:
+            continue
+    return tuple(out)
+
+
+def _line_matches_keep(line: str, patterns: tuple[re.Pattern[str], ...]) -> bool:
+    return any(rx.search(line) for rx in patterns)
+
+
+def shape_logs(
+    text: str,
+    *,
+    max_lines: int | None,
+    keep_patterns: tuple[str, ...],
+    collapse_repeated: bool,
+) -> str:
+    lines = text.splitlines()
+    if len(lines) <= 1:
+        return text
+
+    compiled = _compiled_keep_patterns(keep_patterns)
+    if collapse_repeated:
+        collapsed: list[str] = []
+        prev: str | None = None
+        repeat_count = 0
+        for line in lines:
+            if line == prev and line.strip():
+                repeat_count += 1
+                continue
+            if repeat_count > 0 and prev is not None:
+                collapsed.append(f"{prev}  … (repeated {repeat_count + 1}×)")
+                repeat_count = 0
+            if prev is not None:
+                collapsed.append(prev)
+            prev = line
+        if repeat_count > 0 and prev is not None:
+            collapsed.append(f"{prev}  … (repeated {repeat_count + 1}×)")
+        elif prev is not None:
+            collapsed.append(prev)
+        lines = collapsed
+
+    cap = max_lines
+    if cap is None or len(lines) <= cap:
+        return "\n".join(lines)
+
+    head_n = min(log_head_lines_from_env(), max(1, cap // 2))
+    tail_n = min(log_tail_lines_from_env(), max(1, cap - head_n))
+    keep_indices = set(range(min(head_n, len(lines))))
+    keep_indices.update(range(max(0, len(lines) - tail_n), len(lines)))
+    for i, line in enumerate(lines):
+        if _line_matches_keep(line, compiled):
+            keep_indices.add(i)
+
+    if len(keep_indices) >= cap:
+        ordered = sorted(keep_indices)
+        selected = [lines[i] for i in ordered[:cap]]
+        omitted = len(lines) - len(selected)
+        if omitted > 0:
+            selected.append(f"… (+{omitted} more lines omitted by log shaper)")
+        return "\n".join(selected)
+
+    ordered = sorted(keep_indices)
+    selected = [lines[i] for i in ordered]
+    omitted = len(lines) - len(ordered)
+    if omitted > 0:
+        insert_at = len(selected)
+        for i in range(len(lines)):
+            if i not in keep_indices:
+                insert_at = min(insert_at, len(selected))
+                break
+        marker = f"… (+{omitted} lines omitted by log shaper)"
+        mid = len(selected) // 2
+        selected = [*selected[:mid], marker, *selected[mid:]]
+    return "\n".join(selected)
+
+
+def _json_dedup_key(item: Any) -> str:
+    try:
+        return json.dumps(item, sort_keys=True, default=str)
+    except TypeError:
+        return repr(item)
+
+
+def _json_row_is_anomaly(item: Any) -> bool:
+    if isinstance(item, dict):
+        blob = json.dumps(item, sort_keys=True).lower()
+        return any(k in blob for k in ("error", "exception", "fail", "fatal"))
+    if isinstance(item, str):
+        lower = item.lower()
+        return any(k in lower for k in ("error", "exception", "fail", "fatal"))
+    return False
+
+
+def _shape_json_value(value: Any, *, max_array_items: int, depth: int = 0) -> Any:
+    if depth > 6:
+        return value
+    if isinstance(value, list):
+        seen: set[str] = set()
+        kept: list[Any] = []
+        anomalies: list[Any] = []
+        for item in value:
+            if _json_row_is_anomaly(item):
+                anomalies.append(_shape_json_value(item, max_array_items=max_array_items, depth=depth + 1))
+                continue
+            key = _json_dedup_key(item)
+            if key in seen:
+                continue
+            seen.add(key)
+            kept.append(_shape_json_value(item, max_array_items=max_array_items, depth=depth + 1))
+        merged = anomalies + kept
+        if len(merged) > max_array_items:
+            trimmed = merged[:max_array_items]
+            trimmed.append(f"… (+{len(merged) - max_array_items} more items)")
+            return trimmed
+        return merged
+    if isinstance(value, dict):
+        out: dict[str, Any] = {}
+        for k, v in value.items():
+            if v is None or v == "" or v == [] or v == {}:
+                continue
+            out[str(k)] = _shape_json_value(v, max_array_items=max_array_items, depth=depth + 1)
+        return out
+    return value
+
+
+def shape_json(text: str, *, max_array_items: int | None) -> str:
+    try:
+        parsed = json.loads(text)
+    except json.JSONDecodeError:
+        return text
+    cap = max_array_items if max_array_items is not None else max_array_items_from_env()
+    shaped = _shape_json_value(parsed, max_array_items=cap)
+    return json.dumps(shaped, indent=2, ensure_ascii=False, default=str)
+
+
+def exceeds_tool_output_budget(text: str, *, tool_name: str, budget: ToolOutputBudget) -> bool:
+    """True when source text exceeds configured per-tool caps (safe to shape at append)."""
+    if not text.strip():
+        return False
+    lines = text.splitlines()
+    if budget.max_output_lines is not None and len(lines) > budget.max_output_lines:
+        return True
+    if budget.max_array_items is not None and tool_name in ("web_search", "search_memory", "task"):
+        stripped = text.lstrip()
+        if stripped.startswith(("[", "{")):
+            try:
+                parsed = json.loads(text)
+            except json.JSONDecodeError:
+                return False
+            if isinstance(parsed, list) and len(parsed) > budget.max_array_items:
+                return True
+    if budget.collapse_repeated and len(lines) > 30:
+        repeats = sum(
+            1 for i in range(1, len(lines)) if lines[i] == lines[i - 1] and lines[i].strip()
+        )
+        if repeats >= 5:
+            return True
+    return False
+
+
+def shape_tool_text(
+    text: str,
+    *,
+    tool_name: str,
+    budget: ToolOutputBudget | None,
+    pressure_tier: ContextPressureTier | None = None,
+) -> str:
+    """Apply content-aware shaping when policy or context pressure warrants it."""
+    if not text or not text.strip():
+        return text
+
+    has_policy = budget is not None
+    pressure_shape = pressure_tier in ("moderate", "aggressive")
+    if not has_policy and not pressure_shape:
+        return text
+
+    hint = budget.content_type if budget is not None else None
+    content_type = classify_content(text, tool_name=tool_name, hint=hint)
+
+    if content_type == "code":
+        return text
+    if content_type == "prose" and not pressure_shape:
+        return text
+
+    if content_type == "json":
+        cap = budget.max_array_items if budget is not None else None
+        return shape_json(text, max_array_items=cap)
+
+    if content_type == "logs":
+        return shape_logs(
+            text,
+            max_lines=budget.max_output_lines if budget is not None else None,
+            keep_patterns=budget.keep_patterns if budget is not None else (),
+            collapse_repeated=bool(budget and budget.collapse_repeated),
+        )
+
+    return text
+
+
+def shape_messages_tool_results(
+    messages: Sequence[Message],
+    *,
+    protect_recent: int,
+    pressure_tier: ContextPressureTier | None,
+) -> list[Message]:
+    """Shape tool results in older history rows when context pressure is elevated.
+
+    The most recent ``protect_recent`` messages are left untouched so the model
+    always sees fresh tool output verbatim.
+    """
+    if pressure_tier not in ("moderate", "aggressive"):
+        return list(messages)
+    cutoff = max(0, len(messages) - protect_recent)
+    out: list[Message] = []
+    for i, msg in enumerate(messages):
+        if i >= cutoff or msg.role != "user":
+            out.append(msg)
+            continue
+        new_content: list[ContentBlock] = []
+        changed = False
+        for block in msg.content:
+            if not isinstance(block, ToolResponse) or block.is_error:
+                new_content.append(block)
+                continue
+            text = text_from_blocks(list(block.result))
+            budget = resolve_tool_budget(block.tool_name)
+            shaped = shape_tool_text(
+                text,
+                tool_name=block.tool_name,
+                budget=budget,
+                pressure_tier=pressure_tier,
+            )
+            if shaped == text:
+                new_content.append(block)
+                continue
+            changed = True
+            new_content.append(
+                ToolResponse(
+                    id=block.id,
+                    tool_name=block.tool_name,
+                    result=[Text(text=shaped)],
+                    is_error=block.is_error,
+                )
+            )
+        out.append(Message(role=msg.role, content=new_content) if changed else msg)
+    return out
+
+
+__all__ = [
+    "classify_content",
+    "exceeds_tool_output_budget",
+    "shape_json",
+    "shape_logs",
+    "shape_messages_tool_results",
+    "shape_tool_text",
+]