diff --git a/cheetahclaws/agent.py b/cheetahclaws/agent.py index e12b648..d13a2fd 100644 --- a/cheetahclaws/agent.py +++ b/cheetahclaws/agent.py @@ -7,7 +7,11 @@ from dataclasses import dataclass, field from typing import Generator -from cheetahclaws.tool_registry import get_tool_schemas +from cheetahclaws.tool_registry import ( + get_active_tool_names, + get_tool_schemas, + normalize_tool_profile, +) from cheetahclaws.tools import execute_tool from cheetahclaws import tools as _tools_init # ensure built-in tools are registered on import from cheetahclaws.providers import stream, AssistantTurn, TextChunk, ThinkingChunk, detect_provider, nim_next_model @@ -161,6 +165,30 @@ def run( session_id=session_id, removed=_before_len - len(state.messages)) + # Derive the model-visible and executable surface from the same source + # for this turn. This prevents a provider from seeing a schema that + # dispatch would reject (or vice versa), and avoids sending optional + # integration schemas on every coding request. + try: + active_profile = normalize_tool_profile(config.get("tool_profile")) + except ValueError as profile_error: + active_profile = "standard" + _log.warn("invalid_tool_profile", + session_id=session_id, + requested=config.get("tool_profile"), + fallback=active_profile, + error=str(profile_error)) + disabled_tools = config.get("disabled_tools") or () + if not isinstance(disabled_tools, (list, tuple, set, frozenset)): + disabled_tools = () + active_tool_schemas = get_tool_schemas(active_profile, disabled_tools) + active_tool_names = get_active_tool_names(active_profile, disabled_tools) + config = { + **config, + "tool_profile": active_profile, + "_active_tool_names": active_tool_names, + } + # ── Quota check — before spending tokens ────────────────────────── # Project this request's INPUT so a single large (tool-heavy) call can't # blow past the cap, then clamp the OUTPUT cap to the remaining headroom @@ -224,7 +252,7 @@ def run( model=config["model"], system=system_prompt, messages=state.messages, - tool_schemas=get_tool_schemas(), + tool_schemas=active_tool_schemas, config=_call_config, ): if isinstance(event, (TextChunk, ThinkingChunk)): @@ -353,7 +381,7 @@ def run( # Auto-nudge: text-only reply when the user clearly wanted # investigation (their message contained an absolute path). # One shot only — see `_nudges_remaining` init above. - if _nudges_remaining > 0 and get_tool_schemas(): + if _nudges_remaining > 0 and active_tool_schemas: _nudges_remaining -= 1 _nudge_msg = ( "[system reminder] You replied with text and no tool " @@ -448,6 +476,12 @@ def run( # Check permissions first (must be sequential — may prompt user) permissions: dict[str, bool] = {} for tc in tool_calls: + if tc["name"] not in active_tool_names: + # Treat a stale/malicious call as an execution error, not a + # permission question. The model never received this schema + # on this turn, so prompting a user for it would be misleading. + permissions[tc["id"]] = True + continue permitted = _check_permission(tc, config) if not permitted: if config.get("permission_mode") == "plan": @@ -477,6 +511,13 @@ def run( def _exec_one(tc): """Execute a single tool call, return (tc, result, permitted).""" tid = tc["id"] + if tc["name"] not in active_tool_names: + return ( + tc, + f"Error: tool '{tc['name']}' is not enabled by the " + f"{active_profile!r} tool profile for this turn.", + True, + ) # Read-only dedup short-circuit: skip the actual execute_tool # call, return the synthetic reminder as the tool result. Marked # `permitted=True` so downstream loop-error counters don't treat diff --git a/cheetahclaws/commands/core.py b/cheetahclaws/commands/core.py index fd8caf6..23e6991 100644 --- a/cheetahclaws/commands/core.py +++ b/cheetahclaws/commands/core.py @@ -144,7 +144,10 @@ def _est(text: str) -> int: tool_tokens = 0 try: from cheetahclaws.tool_registry import get_tool_schemas - tool_tokens = _est(json.dumps(get_tool_schemas())) + tool_tokens = _est(json.dumps(get_tool_schemas( + config.get("tool_profile", "standard"), + config.get("disabled_tools") or (), + ))) except Exception: tool_tokens = 0 diff --git a/cheetahclaws/config.py b/cheetahclaws/config.py index 6b19a3c..93a6cab 100644 --- a/cheetahclaws/config.py +++ b/cheetahclaws/config.py @@ -36,6 +36,20 @@ "thinking_budget": 10000, "custom_base_url": "", # for "custom" provider "max_tool_output": 32000, + # Tool schemas are part of every provider request. Keep the normal coding + # loop small; opt into research/orchestration, or use ``full`` for every + # legacy/plugin/MCP tool. + "tool_profile": "standard", # standard | research | orchestration | full + # Bound input work before a tool result reaches the generic output cap. + "tool_read_max_bytes": 256 * 1024, + "tool_read_scan_max_bytes": 2 * 1024 * 1024, + "tool_read_max_output_chars": 50_000, + "web_fetch_max_bytes": 512 * 1024, + "pdf_extract_max_chars": 50_000, + "pdf_extract_max_pages": 50, + # Read-only cache values are post-truncation and capped independently so a + # single large fetch cannot consume unbounded resident memory. + "max_tool_cache_output": 12_000, "max_agent_depth": 3, "max_concurrent_agents": 3, "session_daily_limit": 10000, # max sessions kept per day in daily/ @@ -155,11 +169,22 @@ def load_config() -> dict: CONFIG_DIR.mkdir(exist_ok=True) SESSIONS_DIR.mkdir(exist_ok=True) cfg = dict(DEFAULTS) + saved_config: dict = {} + had_saved_config = CONFIG_FILE.exists() if CONFIG_FILE.exists(): try: - cfg.update(json.loads(CONFIG_FILE.read_text())) + saved_config = json.loads(CONFIG_FILE.read_text()) + if isinstance(saved_config, dict): + cfg.update(saved_config) + else: + saved_config = {} except Exception: pass + # Existing installations predate profiles and historically exposed every + # registered tool. Preserve that behavior until the user explicitly picks + # a compact profile; only a fresh config starts at ``standard``. + if had_saved_config and "tool_profile" not in saved_config: + cfg["tool_profile"] = "full" # Backward-compat: legacy single api_key → anthropic_api_key if cfg.get("api_key") and not cfg.get("anthropic_api_key"): cfg["anthropic_api_key"] = cfg.pop("api_key") diff --git a/cheetahclaws/tool_registry.py b/cheetahclaws/tool_registry.py index a0d3c19..33bec22 100644 --- a/cheetahclaws/tool_registry.py +++ b/cheetahclaws/tool_registry.py @@ -7,8 +7,9 @@ import hashlib import json +import threading from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, FrozenSet, Iterable, List, Optional @dataclass @@ -21,33 +22,116 @@ class ToolDef: func: callable(params: dict, config: dict) -> str read_only: True if the tool never mutates state concurrent_safe: True if safe to run in parallel with other tools + profiles: optional tool-surface profiles this tool belongs to. ``None`` + applies the built-in classification for known tools; third-party + tools default to ``full`` so they are never exposed accidentally. """ name: str schema: Dict[str, Any] func: Callable[[Dict[str, Any], Dict[str, Any]], str] read_only: bool = False concurrent_safe: bool = False + profiles: Optional[FrozenSet[str]] = None # --------------- internal state --------------- _registry: Dict[str, ToolDef] = {} +# ``standard`` deliberately contains the small, high-frequency coding surface. +# Other profiles extend it rather than making the model choose among every +# optional integration on every turn. Unknown/plugin tools remain opt-in via +# ``full`` unless their author explicitly sets ``ToolDef.profiles``. +_PROFILE_NAMES = frozenset({"standard", "research", "orchestration", "full"}) +_STANDARD_TOOLS = frozenset({ + "Read", "Write", "Edit", "Bash", "Glob", "Grep", "GetDiagnostics", + "AskUserQuestion", "NotebookEdit", + "MemorySave", "MemoryDelete", "MemorySearch", "MemoryList", "MemoryVerify", +}) +_RESEARCH_TOOLS = frozenset({ + "WebFetch", "WebSearch", "Research", "ReadPDF", "ReadImage", + "ReadSpreadsheet", "SummarizeLargeFile", +}) +_ORCHESTRATION_TOOLS = frozenset({ + "Agent", "SendMessage", "CheckAgentResult", "ListAgentTasks", + "ListAgentTypes", "Skill", "SkillList", "TaskCreate", "TaskUpdate", + "TaskGet", "TaskList", "EnterPlanMode", "ExitPlanMode", "SleepTimer", +}) + + +def _default_profiles(name: str) -> FrozenSet[str]: + """Classify first-party tools without forcing every registration to change.""" + if name in _STANDARD_TOOLS: + return frozenset({"standard"}) + if name in _RESEARCH_TOOLS: + return frozenset({"research"}) + if name in _ORCHESTRATION_TOOLS: + return frozenset({"orchestration"}) + return frozenset({"full"}) + + +def normalize_tool_profile(profile: str | None) -> str: + """Return a validated profile name. + + Missing values intentionally select ``standard``: this is the safe and + token-efficient default. A caller needing every legacy integration can + set ``tool_profile=full`` explicitly. + """ + if profile is None: + return "standard" + if not isinstance(profile, str): + raise ValueError("Tool profile must be a string.") + normalized = (profile or "standard").strip().lower() + if normalized not in _PROFILE_NAMES: + choices = ", ".join(sorted(_PROFILE_NAMES)) + raise ValueError(f"Unknown tool profile '{profile}'. Choose one of: {choices}.") + return normalized + + +def _profile_allows(tool: ToolDef, profile: str) -> bool: + if profile == "full": + return True + labels = tool.profiles or _default_profiles(tool.name) + if "standard" in labels: + return True + return profile in labels + # --------------- result cache (read-only tools only) --------------- _CACHE_MAX = 64 # max cached entries _cache: Dict[str, str] = {} # hash → result _cache_order: list[str] = [] # LRU eviction order - - -def _cache_key(name: str, params: Dict[str, Any], session_id: str = "") -> str: - """Create a stable hash from tool name + params + session. +_cache_lock = threading.RLock() +_DEFAULT_CACHE_VALUE_MAX = 12_000 +_CACHE_CONFIG_KEYS = ( + # Path authorization is checked inside the tool function. Include it in + # the cache key so a result authorized under one root cannot bypass a + # stricter root later in the same session. + "allowed_root", "_worktree_cwd", + # These settings change source work or visible content for read-only tools. + "tool_read_max_bytes", "tool_read_scan_max_bytes", "tool_read_max_output_chars", + "web_fetch_max_bytes", "pdf_extract_max_chars", "pdf_extract_max_pages", + "max_tool_cache_output", +) + + +def _cache_key( + name: str, + params: Dict[str, Any], + session_id: str = "", + config: Dict[str, Any] | None = None, +) -> str: + """Create a stable hash from tool name + params + session + output policy. Including the session_id keeps cached results scoped to the originator — in a shared daemon, A's Read of ~/.env never gets handed to B's session. """ + cache_config = { + key: (config or {}).get(key) for key in _CACHE_CONFIG_KEYS + if key in (config or {}) + } raw = json.dumps( - {"n": name, "p": params, "s": session_id}, + {"n": name, "p": params, "s": session_id, "c": cache_config}, sort_keys=True, default=str, ) return hashlib.sha256(raw.encode()).hexdigest()[:16] @@ -55,14 +139,17 @@ def _cache_key(name: str, params: Dict[str, Any], session_id: str = "") -> str: def clear_tool_cache() -> None: """Clear the tool result cache. Called on file writes to invalidate.""" - _cache.clear() - _cache_order.clear() + with _cache_lock: + _cache.clear() + _cache_order.clear() # --------------- public API --------------- def register_tool(tool_def: ToolDef) -> None: """Register a tool, overwriting any existing tool with the same name.""" + if tool_def.profiles is None: + tool_def.profiles = _default_profiles(tool_def.name) _registry[tool_def.name] = tool_def @@ -76,9 +163,81 @@ def get_all_tools() -> List[ToolDef]: return list(_registry.values()) -def get_tool_schemas() -> List[Dict[str, Any]]: - """Return the schemas of all registered tools (for API tool parameter).""" - return [t.schema for t in _registry.values()] +def get_tool_schemas( + profile: str | None = "full", + disabled_tools: Iterable[str] | None = None, +) -> List[Dict[str, Any]]: + """Return schemas visible to the model for one tool-surface profile.""" + active_profile = normalize_tool_profile(profile) + disabled = set(disabled_tools or ()) + return [ + tool.schema for tool in _registry.values() + if tool.name not in disabled and _profile_allows(tool, active_profile) + ] + + +def get_active_tool_names( + profile: str | None = "full", + disabled_tools: Iterable[str] | None = None, +) -> FrozenSet[str]: + """Return the executable counterpart to :func:`get_tool_schemas`.""" + return frozenset( + schema["name"] for schema in get_tool_schemas(profile, disabled_tools) + ) + + +def _effective_output_cap(config: Dict[str, Any], max_output: int) -> int: + """Keep an individual tool result below model-context safety limits.""" + try: + from cheetahclaws.compaction import get_context_limit + model = config.get("model", "") if config else "" + declared_ctx = get_context_limit(model) or 32768 + # Reserve 16K for system prompt + tool schemas + framing + headroom. + # 0.5× for CJK-safety (1 char ≈ 1 token worst case). + safe_ctx = min(declared_ctx, 30000) + effective_max = max(2000, int((safe_ctx - 16000) * 0.5)) + return min(max_output, effective_max) + except Exception: + # Compaction module unavailable in some test contexts — retain the + # static cap rather than failing dispatch. + return max_output + + +def _truncate_result(result: str, params: Dict[str, Any], max_output: int) -> str: + """Trim a result while retaining a useful beginning and ending.""" + if len(result) <= max_output: + return result + first_half = max_output // 2 + last_quarter = max_output // 4 + truncated = len(result) - first_half - last_quarter + file_hint = "" + fpath = (params or {}).get("file_path") if isinstance(params, dict) else None + if fpath and isinstance(fpath, str): + file_hint = ( + f" Tip: this came from `{fpath}` — call " + f"`SummarizeLargeFile(file_path='{fpath}')` to get a " + f"complete chunked + map-reduce summary that fits." + ) + return ( + result[:first_half] + + f"\n[... {truncated} chars truncated to keep total tool " + f"output ≤ {max_output:,} chars (model context safety).\n" + f"{file_hint}]\n" + + result[-last_quarter:] + ) + + +def _cache_put(key: str, value: str) -> None: + """Insert a bounded value and keep the LRU index free of duplicates.""" + with _cache_lock: + if key in _cache: + if key in _cache_order: + _cache_order.remove(key) + _cache[key] = value + _cache_order.append(key) + while len(_cache_order) > _CACHE_MAX: + old = _cache_order.pop(0) + _cache.pop(old, None) def execute_tool( @@ -102,13 +261,31 @@ def execute_tool( if tool is None: return f"Error: tool '{name}' not found." + active_names = (config or {}).get("_active_tool_names") + if active_names is not None and name not in active_names: + profile = (config or {}).get("tool_profile", "standard") + return ( + f"Error: tool '{name}' is not enabled by the {profile!r} tool " + "profile for this turn. Select a profile that includes it and retry." + ) + + output_cap = _effective_output_cap(config or {}, max_output) + # Cache hit for read-only tools (same name + same params + same session). use_cache = tool.read_only if use_cache: sid = (config or {}).get("_session_id", "") or "" - key = _cache_key(name, params, sid) - if key in _cache: - return _cache[key] + key = _cache_key(name, params, sid, config) + with _cache_lock: + cached = _cache.get(key) + if cached is not None: + if key in _cache_order: + _cache_order.remove(key) + _cache_order.append(key) + if cached is not None: + # Cache values are already bounded, but cap again because a later + # call can have a smaller context window than the original one. + return _truncate_result(cached, params, output_cap) else: # Write tools invalidate cache (file content may have changed) if name in ("Write", "Edit", "Bash", "NotebookEdit"): @@ -119,58 +296,19 @@ def execute_tool( except Exception as e: return f"Error executing {name}: {e}" - # Store in cache for read-only tools - if use_cache: - _cache[key] = result - _cache_order.append(key) - # Evict oldest if over limit - while len(_cache_order) > _CACHE_MAX: - old = _cache_order.pop(0) - _cache.pop(old, None) + result = _truncate_result(result, params, output_cap) - # Model-aware truncation: the static 32K-char cap is fine for English - # but blows up CJK content (1 token per char). Cap effective max by the - # model's actual context window so a Bash / Read / WebFetch result - # can never single-handedly overflow the next API call. ~30K-token - # conservative ceiling (handles 32K-context models like qwen2.5-72b - # behind a `custom/` provider that lies about context_limit). - try: - from cheetahclaws.compaction import get_context_limit - model = config.get("model", "") if config else "" - declared_ctx = get_context_limit(model) or 32768 - # Reserve 16K for system prompt + tool schemas + framing + headroom. - # 0.5× for CJK-safety (1 char ≈ 1 token worst case). - safe_ctx = min(declared_ctx, 30000) - effective_max = max(2000, int((safe_ctx - 16000) * 0.5)) - if effective_max < max_output: - max_output = effective_max - except Exception: - # Compaction module unavailable in some test contexts — fall back - # to the static 32K cap rather than crashing. - pass - - if len(result) > max_output: - first_half = max_output // 2 - last_quarter = max_output // 4 - truncated = len(result) - first_half - last_quarter - # Surface a SummarizeLargeFile pointer when the truncated tool - # call had a `file_path` arg — gives the model a path forward - # instead of just losing 50%+ of the content. - file_hint = "" - fpath = (params or {}).get("file_path") if isinstance(params, dict) else None - if fpath and isinstance(fpath, str): - file_hint = ( - f" Tip: this came from `{fpath}` — call " - f"`SummarizeLargeFile(file_path='{fpath}')` to get a " - f"complete chunked + map-reduce summary that fits." - ) - result = ( - result[:first_half] - + f"\n[... {truncated} chars truncated to keep total tool " - f"output ≤ {max_output:,} chars (model context safety).\n" - f"{file_hint}]\n" - + result[-last_quarter:] - ) + # Cache only a bounded post-truncation result. This prevents a single + # pathological read-only response from occupying unbounded process RAM. + if use_cache: + try: + cache_cap = int((config or {}).get( + "max_tool_cache_output", _DEFAULT_CACHE_VALUE_MAX + )) + except (TypeError, ValueError): + cache_cap = _DEFAULT_CACHE_VALUE_MAX + cache_cap = max(1_000, min(output_cap, cache_cap)) + _cache_put(key, _truncate_result(result, params, cache_cap)) return result diff --git a/cheetahclaws/tools/__init__.py b/cheetahclaws/tools/__init__.py index 39e461d..4e85215 100644 --- a/cheetahclaws/tools/__init__.py +++ b/cheetahclaws/tools/__init__.py @@ -448,6 +448,16 @@ def execute_tool( """Dispatch tool execution; ask permission for write/destructive ops.""" cfg = config or {} + # This check must run *before* the registry's read-only cache lookup. A + # cache key can cover normal config values, but it cannot safely encode + # ambient authorization such as the filesystem sandbox environment or the + # credential-path denylist. Keep the wrapper check too as defense in depth + # for callers that invoke the registry directly. + if name == "Read" and inputs.get("file_path"): + denied = _check_path_allowed(inputs["file_path"], cfg) + if denied: + return denied + def _check(desc: str) -> bool: if permission_mode == "accept-all": return True @@ -496,7 +506,12 @@ def _read_with_overflow_check(p: dict, c: dict) -> str: denied = _check_path_allowed(p["file_path"], c) if denied: return denied - result = _read(**p) + result = _read( + **p, + max_bytes=c.get("tool_read_max_bytes", 256 * 1024), + scan_max_bytes=c.get("tool_read_scan_max_bytes", 2 * 1024 * 1024), + max_output_chars=c.get("tool_read_max_output_chars", 50_000), + ) # Skip redirect for already-small results (errors, empty, etc.) if not result or len(result) < 8000: return result @@ -578,7 +593,10 @@ def _read_with_overflow_check(p: dict, c: dict) -> str: name="WebFetch", schema=_schemas["WebFetch"], func=lambda p, c: ( - _webfetch(p["url"], p.get("prompt")) + _webfetch( + p["url"], p.get("prompt"), + max_bytes=c.get("web_fetch_max_bytes", 512 * 1024), + ) if isinstance(p.get("url"), str) and p["url"].strip() else "Error: WebFetch requires a non-empty 'url' " "argument (the URL to fetch)." diff --git a/cheetahclaws/tools/files.py b/cheetahclaws/tools/files.py index a77bb57..17c20e7 100644 --- a/cheetahclaws/tools/files.py +++ b/cheetahclaws/tools/files.py @@ -14,6 +14,39 @@ from cheetahclaws.tool_registry import ToolDef, register_tool +def _extract_page_prefix(page, fitz_module, char_cap: int) -> tuple[str, bool]: + """Extract a page in clipped bands instead of materializing all page text. + + PyMuPDF's plain ``page.get_text()`` builds the complete page string first. + Reading shallow horizontal bands keeps peak extraction work bounded even for + a pathological one-page PDF. A compatibility fallback supports lightweight + fake page objects used by downstream plugins/tests. + """ + rect = getattr(page, "rect", None) + if rect is None or not getattr(rect, "height", 0): + text = page.get_text() + return text[:char_cap], len(text) > char_cap + + chunks: list[str] = [] + captured = 0 + band_height = min(144.0, max(36.0, float(rect.height) / 16.0)) + max_bands = 256 + y = float(rect.y0) + bands_read = 0 + while y < float(rect.y1) and captured < char_cap and bands_read < max_bands: + clip = fitz_module.Rect(rect.x0, y, rect.x1, min(y + band_height, rect.y1)) + text = page.get_text("text", clip=clip) + remaining = char_cap - captured + if len(text) > remaining: + chunks.append(text[:remaining]) + return "".join(chunks), True + chunks.append(text) + captured += len(text) + y += band_height + bands_read += 1 + return "".join(chunks), y < float(rect.y1) + + def _read_pdf(params: dict, config: dict) -> str: """Read text content from a PDF file.""" try: @@ -35,27 +68,69 @@ def _read_pdf(params: dict, config: dict) -> str: return f"Error: not a PDF file: {file_path}" try: - doc = fitz.open(str(p)) - total = len(doc) - - # Parse page range - if pages: - page_list = _parse_page_range(pages, total) - else: - page_list = list(range(min(total, 50))) # default: first 50 pages + try: + char_cap = int(config.get("pdf_extract_max_chars", 50_000)) + except (TypeError, ValueError): + char_cap = 50_000 + try: + page_cap = int(config.get("pdf_extract_max_pages", 50)) + except (TypeError, ValueError): + page_cap = 50 + char_cap = max(1_000, char_cap) + page_cap = max(1, page_cap) - text_parts = [] - for i in page_list: - if 0 <= i < total: + doc = fitz.open(str(p)) + try: + total = len(doc) + + # Parse page range. Explicit page lists are capped too: otherwise + # a request such as ``1-999999`` allocates and extracts far more + # than one agent turn can safely use. + if pages: + page_list, page_range_truncated = _parse_page_range_capped( + pages, total, max_pages=page_cap, + ) + else: + page_list = list(range(min(total, page_cap))) + page_range_truncated = total > page_cap + + text_parts = [] + extracted_chars = 0 + source_truncated = page_range_truncated + for i in page_list: + if extracted_chars >= char_cap: + source_truncated = True + break + if not (0 <= i < total): + continue page = doc[i] - text = page.get_text() - if text.strip(): - text_parts.append(f"--- Page {i+1} ---\n{text.strip()}") - - doc.close() + text, page_truncated = _extract_page_prefix( + page, fitz, char_cap - extracted_chars, + ) + source_truncated = source_truncated or page_truncated + clean_text = text.strip() + if not clean_text: + continue + remaining = char_cap - extracted_chars + if len(clean_text) > remaining: + clean_text = clean_text[:remaining] + source_truncated = True + text_parts.append(f"--- Page {i+1} ---\n{clean_text}") + extracted_chars += len(clean_text) + if extracted_chars >= char_cap: + source_truncated = True + break + finally: + doc.close() if not text_parts: - return f"PDF has {total} pages but no extractable text (may be scanned/image-only)." + message = f"PDF has {total} pages but no extractable text (may be scanned/image-only)." + if source_truncated: + message += ( + f" Extraction also stopped at {char_cap:,} characters or " + f"{page_cap} pages; use a narrower `pages` range." + ) + return message header = f"PDF: {p.name} ({total} pages, showing {len(text_parts)})\n\n" content = "\n\n".join(text_parts) @@ -71,8 +146,12 @@ def _read_pdf(params: dict, config: dict) -> str: if redirect: return redirect - if len(content) > 50000: - content = content[:50000] + f"\n\n[... truncated, {len(content)-50000} chars remaining ...]" + if source_truncated: + content += ( + f"\n\n[... ReadPDF stopped at {char_cap:,} extracted characters " + f"or {page_cap} pages; use a narrower `pages` range or " + "SummarizeLargeFile for complete coverage ...]" + ) return header + content @@ -250,19 +329,46 @@ def _format_table(rows: list[list], title: str, total_hint: str = "") -> str: return "\n".join(lines) -def _parse_page_range(spec: str, total: int) -> list[int]: - """Parse page range like '1-5', '3', '1,3,5-8'.""" +def _parse_page_range_capped( + spec: str, + total: int, + max_pages: int | None = None, +) -> tuple[list[int], bool]: + """Parse a page range and state whether a requested page was omitted.""" pages = [] + seen = set() + + def _add(page: int) -> bool: + # Ignore out-of-range singleton values just as ranges are clamped. + # They must not consume the page budget ahead of valid requests. + if not 0 <= page < total: + return False + if page in seen: + return False + if max_pages is not None and len(pages) >= max_pages: + return True + seen.add(page) + pages.append(page) + return False + for part in spec.split(","): part = part.strip() if "-" in part: a, b = part.split("-", 1) start = max(int(a) - 1, 0) end = min(int(b), total) - pages.extend(range(start, end)) + for page in range(start, end): + if _add(page): + return sorted(pages), True elif part.isdigit(): - pages.append(int(part) - 1) - return sorted(set(pages)) + if _add(int(part) - 1): + return sorted(pages), True + return sorted(pages), False + + +def _parse_page_range(spec: str, total: int, max_pages: int | None = None) -> list[int]: + """Parse page range like '1-5', '3', '1,3,5-8'.""" + return _parse_page_range_capped(spec, total, max_pages)[0] # ── Register ───────────────────────────────────────────────────────────── diff --git a/cheetahclaws/tools/fs.py b/cheetahclaws/tools/fs.py index fa553e0..7f3a6b5 100644 --- a/cheetahclaws/tools/fs.py +++ b/cheetahclaws/tools/fs.py @@ -5,6 +5,14 @@ from pathlib import Path +# A Read call should never materialize an arbitrarily large file (or even an +# arbitrarily long single line) before the agent can apply its output cap. +_LINE_SCAN_CHARS = 64 * 1024 +_DEFAULT_READ_MAX_BYTES = 256 * 1024 +_DEFAULT_READ_SCAN_MAX_BYTES = 2 * 1024 * 1024 +_DEFAULT_READ_MAX_OUTPUT_CHARS = 50_000 + + def _read_preserving_newlines(p: Path) -> str: """Read a text file without newline translation. @@ -41,19 +49,185 @@ def maybe_truncate_diff(diff_text: str, max_lines: int = 80) -> str: # ── Read ───────────────────────────────────────────────────────────────── -def _read(file_path: str, limit: int = None, offset: int = None) -> str: +def _encoded_size(text: str) -> int: + """Approximate original UTF-8 bytes while preserving malformed input.""" + return len(text.encode("utf-8", errors="replace")) + + +def _read_logical_line(handle, capture_bytes: int | None, scan_remaining: int): + """Read one universal-newline line with bounded retained text and I/O. + + ``TextIOWrapper(newline="")`` recognizes LF, CRLF, and legacy CR line + endings while preserving them in the returned string. It is read in small + chunks so a minified file cannot force a whole-line allocation. + """ + chunks: list[str] = [] + captured = 0 + scanned = 0 + saw_data = False + + while True: + remaining = scan_remaining - scanned + if remaining <= 0: + return "".join(chunks), False, False, scanned, False, True + + # UTF-8 needs at most four bytes per codepoint. The small allowance + # prevents one read from materially exceeding the scan budget. + request_chars = min(_LINE_SCAN_CHARS, max(1, remaining // 4)) + if capture_bytes is not None: + # Do not create a 64 KiB decoded fragment merely to keep a tiny + # output prefix. A multi-byte codepoint can over-read by at most a + # few bytes; the captured/returned content below is byte-exact. + request_chars = min(request_chars, max(1, capture_bytes - captured)) + piece = handle.readline(request_chars) + if not piece: + return "".join(chunks), saw_data, True, scanned, False, False + + saw_data = True + scanned += _encoded_size(piece) + if capture_bytes is not None: + room = capture_bytes - captured + if room <= 0: + return "".join(chunks), False, False, scanned, True, False + encoded_piece = piece.encode("utf-8", errors="replace") + if len(encoded_piece) > room: + # Decode only whole codepoints so the rendered result remains + # valid text while never exceeding the source-byte ceiling. + chunks.append(encoded_piece[:room].decode("utf-8", errors="ignore")) + return "".join(chunks), False, False, scanned, True, False + chunks.append(piece) + captured += len(encoded_piece) + + if piece.endswith(("\n", "\r")): + return "".join(chunks), True, False, scanned, False, False + if scanned >= scan_remaining: + return "".join(chunks), False, False, scanned, False, True + + +def _read( + file_path: str, + limit: int = None, + offset: int = None, + max_bytes: int = _DEFAULT_READ_MAX_BYTES, + scan_max_bytes: int = _DEFAULT_READ_SCAN_MAX_BYTES, + max_output_chars: int = _DEFAULT_READ_MAX_OUTPUT_CHARS, +) -> str: + """Stream a numbered text slice with bounded I/O and rendered output.""" p = Path(file_path) if not p.exists(): return f"Error: file not found: {file_path}" if p.is_dir(): return f"Error: {file_path} is a directory" try: - lines = _read_preserving_newlines(p).splitlines(keepends=True) - start = offset or 0 - chunk = lines[start:start + limit] if limit else lines[start:] - if not chunk: + start = max(0, int(offset or 0)) + line_limit = int(limit) if limit else None + byte_limit = max(1, int(max_bytes or _DEFAULT_READ_MAX_BYTES)) + scan_limit = max(1, int(scan_max_bytes or _DEFAULT_READ_SCAN_MAX_BYTES)) + output_limit = max(1, int(max_output_chars or _DEFAULT_READ_MAX_OUTPUT_CHARS)) + rendered: list[str] = [] + rendered_lines = 0 + source_bytes = 0 + scanned_bytes = 0 + rendered_chars = 0 + line_no = 0 + source_budget_hit = False + scan_budget_hit = False + output_budget_hit = False + + with p.open(encoding="utf-8", errors="replace", newline="") as handle: + while True: + if line_limit is not None and rendered_lines >= line_limit: + break + if source_bytes >= byte_limit: + # A file ending exactly at the source ceiling is complete, + # not truncated. Probe one character only to distinguish + # it from a longer file without materializing more input. + source_budget_hit = bool(handle.read(1)) + break + if scanned_bytes >= scan_limit: + scan_budget_hit = True + break + + if line_no < start: + _, ended, eof, consumed, _, scan_hit = _read_logical_line( + handle, None, scan_limit - scanned_bytes, + ) + scanned_bytes += consumed + if not ended: + if scan_hit: + scan_budget_hit = True + if eof or scan_budget_hit: + break + else: + line_no += 1 + if eof: + break + continue + + prefix = f"{line_no + 1:6}\t" + output_room = output_limit - rendered_chars + if output_room <= len(prefix): + output_budget_hit = True + break + source_remaining = byte_limit - source_bytes + output_content_room = output_room - len(prefix) + capture_bytes = min(source_remaining, output_content_room) + source_constrained = source_remaining <= output_content_room + text, ended, eof, consumed, capture_hit, scan_hit = _read_logical_line( + handle, max(1, capture_bytes), scan_limit - scanned_bytes, + ) + scanned_bytes += consumed + if not text and eof: + break + line_no += 1 + source_bytes += _encoded_size(text) + formatted = prefix + text + rendered.append(formatted) + rendered_chars += len(formatted) + rendered_lines += 1 + if capture_hit: + if source_constrained: + source_budget_hit = True + else: + output_budget_hit = True + break + if scan_hit: + scan_budget_hit = True + break + if not ended: + # EOF after a final line with no terminator is still a + # valid logical line; any other incomplete line hit a cap. + if not eof: + scan_budget_hit = True + break + + if scan_budget_hit: + marker = ( + f"[... Read stopped after scanning {scan_limit:,} bytes; " + "use a smaller offset or a narrower file ...]\n" + ) + if not rendered: + return marker + rendered.append("\n" + marker) + elif source_budget_hit: + marker = ( + f"[... Read stopped after {byte_limit:,} source bytes; use " + "offset and limit to request another line range ...]\n" + ) + if not rendered: + return marker + rendered.append("\n" + marker) + elif output_budget_hit: + marker = ( + f"[... Read output capped at {output_limit:,} characters to " + "keep memory and model context bounded ...]\n" + ) + if not rendered: + return marker + rendered.append("\n" + marker) + if not rendered: return "(empty file)" - return "".join(f"{start + i + 1:6}\t{l}" for i, l in enumerate(chunk)) + return "".join(rendered) except Exception as e: return f"Error: {e}" diff --git a/cheetahclaws/tools/web.py b/cheetahclaws/tools/web.py index d6fa873..56fcc5d 100644 --- a/cheetahclaws/tools/web.py +++ b/cheetahclaws/tools/web.py @@ -4,23 +4,69 @@ import re -def _webfetch(url: str, prompt: str = None) -> str: +_DEFAULT_WEB_FETCH_MAX_BYTES = 512 * 1024 + + +def _read_response_bytes(response, max_bytes: int) -> tuple[bytes, bool]: + """Consume at most ``max_bytes`` from a streamed HTTP response.""" + data = bytearray() + truncated = False + for chunk in response.iter_bytes(chunk_size=64 * 1024): + remaining = max_bytes - len(data) + if remaining <= 0: + truncated = True + break + if len(chunk) > remaining: + data.extend(chunk[:remaining]) + truncated = True + break + data.extend(chunk) + + # Content-Length lets us report truncation even when the response happens + # to end exactly at the cap without reading an extra network chunk. + try: + content_length = int(response.headers.get("content-length", "0")) + truncated = truncated or content_length > len(data) + except (TypeError, ValueError): + pass + return bytes(data), truncated + + +def _webfetch( + url: str, + prompt: str = None, + max_bytes: int = _DEFAULT_WEB_FETCH_MAX_BYTES, +) -> str: try: import httpx - r = httpx.get(url, headers={"User-Agent": "NanoClaude/1.0"}, - timeout=30, follow_redirects=True) - r.raise_for_status() - ct = r.headers.get("content-type", "") - if "html" in ct: - text = re.sub(r"]*>.*?", "", r.text, + byte_limit = max(1, int(max_bytes or _DEFAULT_WEB_FETCH_MAX_BYTES)) + with httpx.stream( + "GET", url, + headers={"User-Agent": "NanoClaude/1.0"}, + timeout=30, + follow_redirects=True, + ) as response: + response.raise_for_status() + raw, source_truncated = _read_response_bytes(response, byte_limit) + content_type = response.headers.get("content-type", "") + encoding = response.encoding or "utf-8" + + text = raw.decode(encoding, errors="replace") + if "html" in content_type.lower(): + text = re.sub(r"]*>.*?", "", text, flags=re.DOTALL | re.IGNORECASE) text = re.sub(r"]*>.*?", "", text, flags=re.DOTALL | re.IGNORECASE) text = re.sub(r"<[^>]+>", " ", text) text = re.sub(r"\s+", " ", text).strip() - else: - text = r.text - return text[:25000] + + output = text[:25000] + if source_truncated: + output += ( + f"\n\n[... WebFetch stopped after {byte_limit:,} response bytes " + "to keep memory and latency bounded ...]" + ) + return output except ImportError: return "Error: httpx not installed — run: pip install httpx" except Exception as e: diff --git a/tests/test_agent_tool_profiles.py b/tests/test_agent_tool_profiles.py new file mode 100644 index 0000000..92c8898 --- /dev/null +++ b/tests/test_agent_tool_profiles.py @@ -0,0 +1,83 @@ +"""Agent-level contract: the advertised and executable tool sets match.""" +from __future__ import annotations + +from cheetahclaws import agent +from cheetahclaws.agent import AgentState, run +from cheetahclaws.providers import AssistantTurn + + +def _turn(text="", tool_calls=None): + value = AssistantTurn.__new__(AssistantTurn) + value.text = text + value.tool_calls = tool_calls or [] + value.in_tokens = 1 + value.out_tokens = 1 + value.cache_read_tokens = 0 + value.cache_write_tokens = 0 + return value + + +def _config(**extra): + return { + "model": "test", + "permission_mode": "accept-all", + "_session_id": "tool-profile-test", + **extra, + } + + +def test_standard_profile_sends_only_the_compact_surface(monkeypatch): + seen_schemas = [] + + def fake_stream(**kwargs): + seen_schemas.append(kwargs["tool_schemas"]) + yield _turn("done") + + monkeypatch.setattr(agent, "stream", fake_stream) + + list(run("hello", AgentState(), _config(), "system")) + + names = {schema["name"] for schema in seen_schemas[0]} + assert "Read" in names + assert "MemorySearch" in names + assert "WebFetch" not in names + assert "Agent" not in names + assert "ReadPDF" not in names + + +def test_research_profile_exposes_web_and_document_tools(monkeypatch): + seen_schemas = [] + + def fake_stream(**kwargs): + seen_schemas.append(kwargs["tool_schemas"]) + yield _turn("done") + + monkeypatch.setattr(agent, "stream", fake_stream) + + list(run("hello", AgentState(), _config(tool_profile="research"), "system")) + + names = {schema["name"] for schema in seen_schemas[0]} + assert {"Read", "WebFetch", "WebSearch", "ReadPDF"} <= names + assert "Agent" not in names + + +def test_tool_outside_profile_is_rejected_without_permission_prompt(monkeypatch): + replies = iter([ + _turn(tool_calls=[{ + "id": "stale-web", "name": "WebFetch", + "input": {"url": "https://example.test"}, + }]), + _turn("done"), + ]) + + def fake_stream(**_kwargs): + yield next(replies) + + monkeypatch.setattr(agent, "stream", fake_stream) + state = AgentState() + list(run("hello", state, _config(), "system")) + + tool_results = [m for m in state.messages if m.get("role") == "tool"] + assert len(tool_results) == 1 + assert "not enabled" in tool_results[0]["content"] + assert "standard" in tool_results[0]["content"] diff --git a/tests/test_bounded_tool_io.py b/tests/test_bounded_tool_io.py new file mode 100644 index 0000000..31117ed --- /dev/null +++ b/tests/test_bounded_tool_io.py @@ -0,0 +1,214 @@ +"""Regression tests for bounded input work in high-volume tools.""" +from __future__ import annotations + +import sys +import types + +from cheetahclaws.tools.files import ( + _extract_page_prefix, + _parse_page_range, + _parse_page_range_capped, + _read_pdf, +) +from cheetahclaws.tools.fs import _read +from cheetahclaws.tools.web import _webfetch + + +def test_read_streams_requested_line_range(tmp_path): + path = tmp_path / "lines.txt" + path.write_text("zero\none\ntwo\nthree\n", encoding="utf-8") + + result = _read(str(path), offset=1, limit=2, max_bytes=1_000) + + assert " 2\tone" in result + assert " 3\ttwo" in result + assert "zero" not in result + assert "three" not in result + + +def test_read_caps_a_single_huge_line_without_loading_it(tmp_path): + path = tmp_path / "minified.js" + path.write_bytes(b"x" * 200_000 + b"\nnext\n") + + result = _read(str(path), max_bytes=128) + + assert "Read stopped after 128 source bytes" in result + assert len(result) < 1_000 + + +def test_read_never_returns_more_than_the_source_byte_cap_for_utf8(tmp_path): + path = tmp_path / "emoji.txt" + path.write_text("😀" * 10 + "\n", encoding="utf-8") + + result = _read(str(path), max_bytes=5, scan_max_bytes=100) + visible = result.split("\n[...", 1)[0].split("\t", 1)[1] + + assert len(visible.encode("utf-8")) <= 5 + assert "stopped after 5 source bytes" in result + + +def test_read_does_not_claim_truncation_at_exact_source_byte_eof(tmp_path): + path = tmp_path / "exact.txt" + path.write_bytes(b"exact") + + result = _read(str(path), max_bytes=5, scan_max_bytes=100) + + assert "exact" in result + assert "stopped after" not in result + + +def test_webfetch_streams_and_caps_response(monkeypatch): + class FakeResponse: + headers = {"content-type": "text/html", "content-length": "1000"} + encoding = "utf-8" + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def raise_for_status(self): + return None + + def iter_bytes(self, **_kwargs): + yield b"

" + b"a" * 100 + yield b"b" * 900 + b"

" + + fake_httpx = types.SimpleNamespace(stream=lambda *_args, **_kwargs: FakeResponse()) + monkeypatch.setitem(sys.modules, "httpx", fake_httpx) + + result = _webfetch("https://example.test", max_bytes=128) + + assert "WebFetch stopped after 128 response bytes" in result + assert len(result) < 400 + + +def test_pdf_extract_stops_after_character_cap(monkeypatch, tmp_path): + class FakePage: + def __init__(self, text): + self._text = text + + def get_text(self): + return self._text + + class FakeDoc: + def __init__(self): + self.pages = [FakePage("a" * 900), FakePage("b" * 900)] + self.closed = False + + def __len__(self): + return len(self.pages) + + def __getitem__(self, index): + return self.pages[index] + + def close(self): + self.closed = True + + doc = FakeDoc() + monkeypatch.setitem(sys.modules, "fitz", types.SimpleNamespace(open=lambda _p: doc)) + path = tmp_path / "sample.pdf" + path.write_bytes(b"%PDF-fake") + + result = _read_pdf( + {"file_path": str(path)}, + {"pdf_extract_max_chars": 1_000, "pdf_extract_max_pages": 50}, + ) + + assert doc.closed is True + assert "ReadPDF stopped at 1,000 extracted characters" in result + assert len(result) < 1_500 + + +def test_pdf_page_range_is_bounded_before_building_a_large_list(): + assert _parse_page_range("1-1000000", 1_000_000, max_pages=3) == [0, 1, 2] + + +def test_pdf_page_range_reports_when_explicit_request_is_capped(): + pages, truncated = _parse_page_range_capped("1-3", 3, max_pages=2) + assert pages == [0, 1] + assert truncated is True + + +def test_pdf_page_range_ignores_out_of_range_singletons_before_capping(): + pages, truncated = _parse_page_range_capped("999,1,2", 3, max_pages=2) + assert pages == [0, 1] + assert truncated is False + + +def test_pdf_extract_uses_bounded_clipped_bands(): + class Rect: + def __init__(self, x0, y0, x1, y1): + self.x0, self.y0, self.x1, self.y1 = x0, y0, x1, y1 + self.height = y1 - y0 + + class Page: + rect = Rect(0, 0, 100, 1_000_000_000) + + def __init__(self): + self.calls = [] + + def get_text(self, mode, *, clip): + self.calls.append((mode, clip)) + return "" + + page = Page() + fake_fitz = types.SimpleNamespace(Rect=Rect) + text, truncated = _extract_page_prefix(page, fake_fitz, 1_000) + + assert text == "" + assert truncated is True + assert len(page.calls) == 256 + + +def test_read_stops_seeking_when_offset_exceeds_scan_budget(tmp_path): + path = tmp_path / "many-lines.txt" + path.write_text("row\n" * 100, encoding="utf-8") + + result = _read( + str(path), offset=50, limit=1, + max_bytes=1_000, scan_max_bytes=16, + ) + + assert "stopped after scanning 16 bytes" in result + + +def test_read_keeps_legacy_cr_newline_offsets(tmp_path): + path = tmp_path / "classic-mac.txt" + path.write_bytes(b"first\rsecond\rthird\r") + + result = _read(str(path), offset=1, limit=1, max_bytes=1_000) + + assert " 2\tsecond\r" in result + + +def test_read_caps_rendered_output_for_many_short_lines(tmp_path): + path = tmp_path / "empty-lines.txt" + path.write_bytes(b"\n" * 10_000) + + result = _read(str(path), max_bytes=10_000, max_output_chars=200) + + assert "Read output capped at 200 characters" in result + assert len(result) < 400 + + +def test_read_cache_cannot_bypass_allowed_root(tmp_path): + from cheetahclaws.tool_registry import clear_tool_cache + from cheetahclaws.tools import execute_tool + + path = tmp_path / "visible.txt" + path.write_text("private content\n", encoding="utf-8") + clear_tool_cache() + first = execute_tool( + "Read", {"file_path": str(path)}, permission_mode="accept-all", + config={"allowed_root": str(tmp_path), "_session_id": "root-change"}, + ) + second = execute_tool( + "Read", {"file_path": str(path)}, permission_mode="accept-all", + config={"allowed_root": str(tmp_path / "other"), "_session_id": "root-change"}, + ) + + assert "private content" in first + assert "Error:" in second + assert "private content" not in second diff --git a/tests/test_tool_profile_config.py b/tests/test_tool_profile_config.py new file mode 100644 index 0000000..da5d0e9 --- /dev/null +++ b/tests/test_tool_profile_config.py @@ -0,0 +1,33 @@ +"""Compatibility and validation tests for tool-profile configuration.""" +from __future__ import annotations + +import json + +import pytest + +from cheetahclaws import config as config_module +from cheetahclaws.tool_registry import normalize_tool_profile + + +def test_legacy_saved_config_keeps_full_tool_surface(monkeypatch, tmp_path): + config_file = tmp_path / "config.json" + config_file.write_text(json.dumps({"model": "test"}), encoding="utf-8") + monkeypatch.setattr(config_module, "CONFIG_DIR", tmp_path) + monkeypatch.setattr(config_module, "CONFIG_FILE", config_file) + monkeypatch.setattr(config_module, "SESSIONS_DIR", tmp_path / "sessions") + + assert config_module.load_config()["tool_profile"] == "full" + + +def test_fresh_config_uses_compact_standard_profile(monkeypatch, tmp_path): + monkeypatch.setattr(config_module, "CONFIG_DIR", tmp_path) + monkeypatch.setattr(config_module, "CONFIG_FILE", tmp_path / "missing.json") + monkeypatch.setattr(config_module, "SESSIONS_DIR", tmp_path / "sessions") + + assert config_module.load_config()["tool_profile"] == "standard" + + +@pytest.mark.parametrize("value", [1, ["standard"], {"profile": "full"}]) +def test_invalid_tool_profile_value_is_a_clean_validation_error(value): + with pytest.raises(ValueError): + normalize_tool_profile(value) diff --git a/tests/test_tool_registry.py b/tests/test_tool_registry.py index abd6104..e75f551 100644 --- a/tests/test_tool_registry.py +++ b/tests/test_tool_registry.py @@ -4,8 +4,10 @@ from cheetahclaws.tool_registry import ( ToolDef, + clear_tool_cache, clear_registry, execute_tool, + get_active_tool_names, get_all_tools, get_tool, get_tool_schemas, @@ -15,10 +17,15 @@ @pytest.fixture(autouse=True) def _clean_registry(): - """Reset registry before each test.""" + """Isolate registry tests without leaving later integration tests empty.""" + original_tools = get_all_tools() clear_registry() + clear_tool_cache() yield clear_registry() + clear_tool_cache() + for tool in original_tools: + register_tool(tool) def _make_echo_tool(name: str = "echo", read_only: bool = False) -> ToolDef: @@ -89,6 +96,46 @@ def test_get_tool_schemas(): assert schemas[0]["name"] == "echo" +def test_tool_profiles_filter_schemas_and_names(): + register_tool(ToolDef( + name="core", + schema={"name": "core", "input_schema": {}}, + func=lambda _p, _c: "core", + profiles=frozenset({"standard"}), + )) + register_tool(ToolDef( + name="research_only", + schema={"name": "research_only", "input_schema": {}}, + func=lambda _p, _c: "research", + profiles=frozenset({"research"}), + )) + register_tool(ToolDef( + name="full_only", + schema={"name": "full_only", "input_schema": {}}, + func=lambda _p, _c: "full", + profiles=frozenset({"full"}), + )) + + assert [s["name"] for s in get_tool_schemas("standard")] == ["core"] + assert [s["name"] for s in get_tool_schemas("research")] == [ + "core", "research_only", + ] + assert get_active_tool_names("orchestration") == frozenset({"core"}) + assert {s["name"] for s in get_tool_schemas("full")} == { + "core", "research_only", "full_only", + } + + +def test_tool_profiles_honor_disabled_tools(): + register_tool(ToolDef( + name="core", + schema={"name": "core", "input_schema": {}}, + func=lambda _p, _c: "core", + profiles=frozenset({"standard"}), + )) + assert get_tool_schemas("standard", disabled_tools=["core"]) == [] + + # ------------------------------------------------------------------ # execute_tool # ------------------------------------------------------------------ @@ -138,6 +185,54 @@ def test_no_truncation_when_within_limit(): assert result == "short" +def test_cache_stores_bounded_result_and_reapplies_smaller_cap(): + calls = 0 + + def big_func(params: dict, config: dict) -> str: + nonlocal calls + calls += 1 + return "x" * 20_000 + + register_tool(ToolDef( + name="cached_big", + schema={"name": "cached_big", "input_schema": {}}, + func=big_func, + read_only=True, + )) + + first = execute_tool( + "cached_big", {}, {"max_tool_cache_output": 2_000}, max_output=10_000, + ) + second = execute_tool( + "cached_big", {}, {"max_tool_cache_output": 2_000}, max_output=1_500, + ) + + assert calls == 1 + assert "truncated" in first + assert "truncated" in second + assert len(second) < len(first) + + +def test_cache_key_includes_input_bound_settings(): + calls = 0 + + def config_echo(_params: dict, config: dict) -> str: + nonlocal calls + calls += 1 + return str(config["tool_read_max_bytes"]) + + register_tool(ToolDef( + name="config_sensitive", + schema={"name": "config_sensitive", "input_schema": {}}, + func=config_echo, + read_only=True, + )) + + assert execute_tool("config_sensitive", {}, {"tool_read_max_bytes": 10}) == "10" + assert execute_tool("config_sensitive", {}, {"tool_read_max_bytes": 20}) == "20" + assert calls == 2 + + # ------------------------------------------------------------------ # duplicate register overwrites # ------------------------------------------------------------------