mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-15 17:25:26 -04:00
fix(search): add download budgets to web_fetch with truncation notice and hard ceiling (#3955)
* fix(search): add download budgets to web_fetch with truncation notice and hard ceiling MAX_OUTPUT_CHARS only trims what the agent sees; fetch_webpage_content buffered and cached the entire response body first, so a large or hostile URL could pull arbitrarily many bytes into memory and the content cache. The fetch is now a capped streaming GET (SSRF redirect guard unchanged): a soft default budget (WEB_FETCH_SOFT_MAX_BYTES, 2 MB), a per-call override via full/max_bytes on the web_fetch tool, and a hard ceiling (WEB_FETCH_HARD_MAX_BYTES, 20 MB) that the override can never exceed. When Content-Length already declares a body over the ceiling the fetch is refused before any body bytes are buffered. Truncated results carry truncated/fetched_bytes/total_bytes, the tool output leads with a partial-content notice telling the model how to re-fetch with full=true, and the tool schema documents the flag. A truncated PDF is reported as a budget error since a cut PDF is unparseable. The effective cap is part of the content-cache key so a truncated fetch is never served to a full-budget request. Existing tests that faked httpx.get or the old _get_public_url signature are adapted to the streaming interface; behavior pins are unchanged. Fixes #3812 * fix(search): close compressed-body cap bypass and protect the partial notice Addresses RaresKeY's review on #3955: - Force Accept-Encoding: identity for the capped fetch. With gzip/deflate the wire bytes (and Content-Length) can be a fraction of the decoded body, so a tiny compressed response could pass the hard-cap preflight and then expand past the ceiling in a single decoded chunk before the streamed cap could slice it. Identity makes Content-Length the true body size and keeps each streamed chunk bounded by the network read, so the hard ceiling actually bounds memory. - Lead web_fetch output with the partial-content notice and cap the page title. The notice is the user-facing contract for partial fetches, but the title is untrusted, uncapped page content; placed ahead of the notice a giant title could push it past MAX_OUTPUT_CHARS and drop it. The notice now leads and the title is capped as a second guard. Adds regressions: the fetch advertises identity encoding, and a truncated result with an oversized title still surfaces the partial notice. * fix(search): reject compressed responses that ignore the identity request Requesting Accept-Encoding: identity is not enough on its own: a server can ignore it and still return Content-Encoding: gzip, and httpx.iter_bytes would decode that, so a tiny compressed body could balloon into one decoded chunk far past the hard cap before the streamed loop slices it (and Content-Length, the compressed wire length, makes the preflight and size metadata unreliable). Refuse a non-identity Content-Encoding before reading the body. Adds a regression where the server ignores the identity request and returns gzip; the fetch is refused before any body is decoded.
This commit is contained in:
committed by
GitHub
parent
2fab378c6a
commit
074a1e6eff
+162
-13
@@ -15,6 +15,8 @@ from urllib.parse import urljoin, urlparse
|
||||
import httpx
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
from src.constants import WEB_FETCH_SOFT_MAX_BYTES, WEB_FETCH_HARD_MAX_BYTES
|
||||
|
||||
from .analytics import RateLimitError, error_logger
|
||||
from .cache import (
|
||||
CONTENT_CACHE_DIR,
|
||||
@@ -89,18 +91,128 @@ def _public_http_url(url: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _get_public_url(url: str, headers: dict, timeout: int, max_redirects: int = 5) -> httpx.Response:
|
||||
class BodyTooLargeError(Exception):
|
||||
"""The server declared a body larger than the hard fetch ceiling."""
|
||||
|
||||
def __init__(self, url: str, declared_bytes: int):
|
||||
self.url = url
|
||||
self.declared_bytes = declared_bytes
|
||||
super().__init__(
|
||||
f"response body is {declared_bytes:,} bytes, over the "
|
||||
f"{WEB_FETCH_HARD_MAX_BYTES:,}-byte hard cap"
|
||||
)
|
||||
|
||||
|
||||
class _CappedFetch:
|
||||
"""Result of a size-capped streaming GET.
|
||||
|
||||
Carries just what fetch_webpage_content needs from an httpx.Response,
|
||||
plus the cap bookkeeping: the (possibly truncated) body, whether the
|
||||
cap cut it short, and the size the server declared via Content-Length
|
||||
(wire bytes; None when absent).
|
||||
"""
|
||||
|
||||
__slots__ = ("status_code", "headers", "content", "truncated",
|
||||
"declared_bytes", "encoding", "url")
|
||||
|
||||
def __init__(self, status_code, headers, content, truncated,
|
||||
declared_bytes, encoding, url):
|
||||
self.status_code = status_code
|
||||
self.headers = headers
|
||||
self.content = content
|
||||
self.truncated = truncated
|
||||
self.declared_bytes = declared_bytes
|
||||
self.encoding = encoding
|
||||
self.url = url
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
return self.content.decode(self.encoding or "utf-8", errors="replace")
|
||||
|
||||
def raise_for_status(self):
|
||||
if self.status_code >= 400:
|
||||
request = httpx.Request("GET", self.url)
|
||||
raise httpx.HTTPStatusError(
|
||||
f"HTTP {self.status_code} for {self.url}",
|
||||
request=request,
|
||||
response=httpx.Response(self.status_code, request=request),
|
||||
)
|
||||
|
||||
|
||||
def _get_public_url(url: str, headers: dict, timeout: int, max_redirects: int = 5,
|
||||
max_bytes: int = None) -> "_CappedFetch":
|
||||
"""Capped streaming GET with SSRF-guarded manual redirects.
|
||||
|
||||
The body is streamed and buffering stops at ``max_bytes`` (default: the
|
||||
soft cap), so an oversized resource cannot be pulled into memory or the
|
||||
content cache in full. When Content-Length already declares a body over
|
||||
the hard ceiling, the fetch is refused before any body bytes are read.
|
||||
"""
|
||||
cap = min(max_bytes or WEB_FETCH_SOFT_MAX_BYTES, WEB_FETCH_HARD_MAX_BYTES)
|
||||
current = url
|
||||
for _ in range(max_redirects + 1):
|
||||
if not _public_http_url(current):
|
||||
raise httpx.RequestError("Blocked private/internal URL", request=httpx.Request("GET", current))
|
||||
response = httpx.get(current, headers=headers, timeout=timeout, follow_redirects=False)
|
||||
if response.status_code not in (301, 302, 303, 307, 308):
|
||||
return response
|
||||
location = response.headers.get("location")
|
||||
if not location:
|
||||
return response
|
||||
current = urljoin(str(response.url), location)
|
||||
# Force identity transfer-encoding. With gzip/deflate the wire bytes
|
||||
# (and Content-Length) can be a small fraction of the decoded body, so
|
||||
# a tiny compressed response could pass the hard-cap preflight and then
|
||||
# expand past the ceiling in a single decoded chunk before the streamed
|
||||
# cap below can slice it. Identity makes Content-Length the true body
|
||||
# size and keeps each streamed chunk bounded by the network read.
|
||||
req_headers = dict(headers or {})
|
||||
req_headers["Accept-Encoding"] = "identity"
|
||||
with httpx.stream("GET", current, headers=req_headers, timeout=timeout,
|
||||
follow_redirects=False) as response:
|
||||
if response.status_code in (301, 302, 303, 307, 308):
|
||||
location = response.headers.get("location")
|
||||
if not location:
|
||||
return _CappedFetch(response.status_code, response.headers, b"",
|
||||
False, None, response.encoding, str(response.url))
|
||||
current = urljoin(str(response.url), location)
|
||||
continue
|
||||
|
||||
# A server can ignore the identity request and still return a
|
||||
# compressed body; httpx.iter_bytes would then decode it, and a tiny
|
||||
# gzip can balloon into one decoded chunk far past the cap before we
|
||||
# slice. Refuse a compressed Content-Encoding so the streamed cap
|
||||
# stays a real memory bound (Content-Length is the compressed wire
|
||||
# length here, so the preflight and size metadata are unreliable too).
|
||||
enc = (response.headers.get("content-encoding") or "").strip().lower()
|
||||
if enc and enc != "identity":
|
||||
raise httpx.RequestError(
|
||||
f"Refusing compressed response (Content-Encoding: {enc}) after "
|
||||
"requesting identity: cannot bound decoded body size",
|
||||
request=httpx.Request("GET", current),
|
||||
)
|
||||
|
||||
declared = None
|
||||
raw_len = response.headers.get("content-length")
|
||||
if raw_len and raw_len.isdigit():
|
||||
declared = int(raw_len)
|
||||
# Refuse before buffering anything when the server already tells
|
||||
# us the body exceeds the absolute ceiling (Content-Length is wire
|
||||
# bytes; the decompressed body can only be larger).
|
||||
if declared is not None and declared > WEB_FETCH_HARD_MAX_BYTES:
|
||||
raise BodyTooLargeError(current, declared)
|
||||
|
||||
chunks = []
|
||||
read = 0
|
||||
truncated = False
|
||||
# We requested identity above, so iter_bytes yields the raw body in
|
||||
# network-read-sized chunks (no decompression expansion); the cap
|
||||
# therefore bounds what we actually buffer.
|
||||
for chunk in response.iter_bytes():
|
||||
read += len(chunk)
|
||||
if read > cap:
|
||||
keep = cap - (read - len(chunk))
|
||||
if keep > 0:
|
||||
chunks.append(chunk[:keep])
|
||||
truncated = True
|
||||
break
|
||||
chunks.append(chunk)
|
||||
return _CappedFetch(response.status_code, response.headers,
|
||||
b"".join(chunks), truncated, declared,
|
||||
response.encoding, str(response.url))
|
||||
raise httpx.RequestError("Too many redirects", request=httpx.Request("GET", current))
|
||||
|
||||
# PDF extraction (optional dependency)
|
||||
@@ -222,9 +334,19 @@ def _empty_result(url: str, error: str = "") -> dict:
|
||||
# ----------------------------------------------------------------------
|
||||
# Main content fetcher
|
||||
# ----------------------------------------------------------------------
|
||||
def fetch_webpage_content(url: str, timeout: int = 5, retry_attempt: int = 0) -> dict:
|
||||
"""Fetch and extract meaningful content from a webpage with caching."""
|
||||
cache_key = generate_cache_key(url)
|
||||
def fetch_webpage_content(url: str, timeout: int = 5, retry_attempt: int = 0,
|
||||
max_bytes: int = None) -> dict:
|
||||
"""Fetch and extract meaningful content from a webpage with caching.
|
||||
|
||||
``max_bytes`` raises the download budget per call (clamped to the hard
|
||||
cap); the default is the soft cap. When the body is cut short the result
|
||||
carries ``truncated``/``fetched_bytes``/``total_bytes`` so callers can
|
||||
tell the model the content is partial (#3812).
|
||||
"""
|
||||
effective_cap = min(max_bytes or WEB_FETCH_SOFT_MAX_BYTES, WEB_FETCH_HARD_MAX_BYTES)
|
||||
# The cap is part of the cache identity: a truncated soft-cap fetch must
|
||||
# not be served to a later full-budget request for the same URL.
|
||||
cache_key = generate_cache_key(f"{url}#cap={effective_cap}")
|
||||
cache_file = CONTENT_CACHE_DIR / f"{cache_key}.cache"
|
||||
|
||||
# Check cache
|
||||
@@ -250,15 +372,21 @@ def fetch_webpage_content(url: str, timeout: int = 5, retry_attempt: int = 0) ->
|
||||
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
|
||||
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
|
||||
"Accept-Language": "en-US,en;q=0.5",
|
||||
"Accept-Encoding": "gzip, deflate",
|
||||
# identity so the streamed size cap in _get_public_url stays honest
|
||||
# (a compressed body can decode to far more than Content-Length).
|
||||
"Accept-Encoding": "identity",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
response = _get_public_url(url, headers=headers, timeout=timeout)
|
||||
response = _get_public_url(url, headers=headers, timeout=timeout,
|
||||
max_bytes=effective_cap)
|
||||
|
||||
if response.status_code == 429:
|
||||
raise RateLimitError(f"Rate limit hit for {url} (attempt {retry_attempt})")
|
||||
|
||||
response.raise_for_status()
|
||||
except BodyTooLargeError as e:
|
||||
error_logger.warning(f"Refused oversized body for {url}: {e}")
|
||||
return _empty_result(url, f"TooLarge: {e}")
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_logger.warning(f"HTTP {e.response.status_code} fetching {url}: {e}")
|
||||
return _empty_result(url, f"HTTP {e.response.status_code}: {e}")
|
||||
@@ -269,9 +397,27 @@ def fetch_webpage_content(url: str, timeout: int = 5, retry_attempt: int = 0) ->
|
||||
error_logger.error(str(e))
|
||||
return _empty_result(url, str(e))
|
||||
|
||||
# Size bookkeeping shared by every content branch below. getattr keeps
|
||||
# plain httpx.Response stand-ins (tests) working without the cap fields.
|
||||
_size_fields = {
|
||||
"truncated": getattr(response, "truncated", False),
|
||||
"fetched_bytes": len(response.content),
|
||||
"total_bytes": getattr(response, "declared_bytes", None),
|
||||
}
|
||||
|
||||
# PDF handling
|
||||
content_type = response.headers.get("Content-Type", "").lower()
|
||||
if "application/pdf" in content_type or url.lower().endswith(".pdf"):
|
||||
if _size_fields["truncated"]:
|
||||
# A PDF cut mid-stream is not parseable; unlike text there is no
|
||||
# useful partial result, so report the budget problem instead.
|
||||
_declared = _size_fields["total_bytes"]
|
||||
return _empty_result(
|
||||
url,
|
||||
f"TooLarge: PDF exceeds the {effective_cap:,}-byte fetch budget"
|
||||
+ (f" (size {_declared:,} bytes)" if _declared else "")
|
||||
+ "; retry with a larger budget if it fits under the hard cap",
|
||||
)
|
||||
if pdf_extract_text is None:
|
||||
logger.error("pdfminer.six is not installed; cannot extract PDF text.")
|
||||
pdf_text = ""
|
||||
@@ -295,6 +441,7 @@ def fetch_webpage_content(url: str, timeout: int = 5, retry_attempt: int = 0) ->
|
||||
"js_message": "",
|
||||
"success": bool(pdf_text),
|
||||
"error": "" if pdf_text else "Failed to extract PDF text",
|
||||
**_size_fields,
|
||||
}
|
||||
_cache_result(cache_file, cache_key, result, url)
|
||||
return result
|
||||
@@ -329,6 +476,7 @@ def fetch_webpage_content(url: str, timeout: int = 5, retry_attempt: int = 0) ->
|
||||
"js_message": "",
|
||||
"success": bool(text_body),
|
||||
"error": "" if text_body else "Empty response body",
|
||||
**_size_fields,
|
||||
}
|
||||
_cache_result(cache_file, cache_key, result, url)
|
||||
return result
|
||||
@@ -391,6 +539,7 @@ def fetch_webpage_content(url: str, timeout: int = 5, retry_attempt: int = 0) ->
|
||||
"js_message": js_message,
|
||||
"success": True,
|
||||
"error": "",
|
||||
**_size_fields,
|
||||
}
|
||||
_cache_result(cache_file, cache_key, result, url)
|
||||
return result
|
||||
|
||||
@@ -57,13 +57,23 @@ class WebSearchTool:
|
||||
class WebFetchTool:
|
||||
async def execute(self, content: str, ctx: dict) -> dict:
|
||||
from src.search.content import fetch_webpage_content
|
||||
from src.constants import WEB_FETCH_HARD_MAX_BYTES
|
||||
raw = content.strip()
|
||||
url = ""
|
||||
max_bytes = None
|
||||
if raw.startswith("{"):
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
if isinstance(parsed, dict):
|
||||
url = str(parsed.get("url") or "").strip()
|
||||
# Download-budget override (#3812): "full": true raises the
|
||||
# budget to the hard cap; an explicit max_bytes is clamped
|
||||
# to the hard cap downstream. Default stays the soft cap.
|
||||
if parsed.get("full") is True:
|
||||
max_bytes = WEB_FETCH_HARD_MAX_BYTES
|
||||
mb = parsed.get("max_bytes")
|
||||
if isinstance(mb, int) and mb > 0:
|
||||
max_bytes = mb
|
||||
except json.JSONDecodeError:
|
||||
url = ""
|
||||
if not url:
|
||||
@@ -78,7 +88,7 @@ class WebFetchTool:
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
loop.run_in_executor(None, lambda: fetch_webpage_content(url, timeout=10)),
|
||||
loop.run_in_executor(None, lambda: fetch_webpage_content(url, timeout=10, max_bytes=max_bytes)),
|
||||
timeout=30,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
@@ -94,8 +104,28 @@ class WebFetchTool:
|
||||
return {"error": f"web_fetch: {url}: {err}", "exit_code": 1}
|
||||
return {"error": f"web_fetch: {url}: no readable text content (not HTML, or the page needs JS/login)", "exit_code": 1}
|
||||
|
||||
# Tell the model when the download budget cut the body short and how
|
||||
# to get the rest, instead of silently presenting a partial page as
|
||||
# the whole thing.
|
||||
size_note = ""
|
||||
if result.get("truncated"):
|
||||
fetched = result.get("fetched_bytes") or 0
|
||||
total = result.get("total_bytes")
|
||||
total_txt = f" of {total:,} bytes" if total else ""
|
||||
size_note = (
|
||||
f"[partial content: download stopped at {fetched:,} bytes{total_txt}. "
|
||||
f'Re-call with {{"url": "{url}", "full": true}} to fetch up to '
|
||||
f"{WEB_FETCH_HARD_MAX_BYTES:,} bytes.]\n\n"
|
||||
)
|
||||
|
||||
# The notice must lead the output so the MAX_OUTPUT_CHARS trim below can
|
||||
# never drop it. The title is untrusted, uncapped page content, so a
|
||||
# giant title ahead of the notice could push it out of range; keep the
|
||||
# notice first and cap the title as a second guard.
|
||||
if len(title) > 300:
|
||||
title = title[:300] + "..."
|
||||
header = (f"# {title}\n" if title else "") + f"Source: {url}\n\n"
|
||||
output = header + text
|
||||
output = size_note + header + text
|
||||
if len(output) > MAX_OUTPUT_CHARS:
|
||||
output = output[:MAX_OUTPUT_CHARS] + "\n\n[...truncated]"
|
||||
return {"output": output, "exit_code": 0}
|
||||
|
||||
@@ -65,6 +65,14 @@ MAX_OUTPUT_CHARS = 10_000 # cap for bash/python/web_search/web_fetch outpu
|
||||
MAX_READ_CHARS = 20_000 # cap for read_file / document preview
|
||||
MAX_DIFF_LINES = 400 # cap for edit_file unified-diff display
|
||||
|
||||
# web_fetch response-size policy (#3812). MAX_OUTPUT_CHARS above only trims
|
||||
# what the agent SEES; these caps bound what the server downloads, parses,
|
||||
# and writes to the content cache. The soft cap is the default download
|
||||
# budget; the agent can raise it per call (full/max_bytes) but never past
|
||||
# the hard cap, so a model can't decide to pull a multi-GB file.
|
||||
WEB_FETCH_SOFT_MAX_BYTES = 2_000_000 # default download budget (2 MB)
|
||||
WEB_FETCH_HARD_MAX_BYTES = 20_000_000 # absolute ceiling, even with override (20 MB)
|
||||
|
||||
# API Configuration
|
||||
MAX_CONTEXT_MESSAGES = 90
|
||||
REQUEST_TIMEOUT = 20
|
||||
|
||||
+3
-2
@@ -68,11 +68,12 @@ FUNCTION_TOOL_SCHEMAS = [
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "web_fetch",
|
||||
"description": "Fetch and read the text content of a specific URL the user names (e.g. 'check example.com', 'what's on this page <url>'). Use when you already have a concrete URL/domain. NOT for open-ended searches (use web_search) or 'research X' jobs (use trigger_research).",
|
||||
"description": "Fetch and read the text content of a specific URL the user names (e.g. 'check example.com', 'what's on this page <url>'). Use when you already have a concrete URL/domain. NOT for open-ended searches (use web_search) or 'research X' jobs (use trigger_research). Downloads are size-budgeted; a '[partial content: ...]' notice in the result means the body was cut short and you can re-call with full=true for the rest.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"url": {"type": "string", "description": "The URL or domain to fetch (http/https; a bare domain like example.com is fine)"}
|
||||
"url": {"type": "string", "description": "The URL or domain to fetch (http/https; a bare domain like example.com is fine)"},
|
||||
"full": {"type": "boolean", "description": "Raise the download budget to the hard cap for large pages/files. Use only after a result reported partial content."}
|
||||
},
|
||||
"required": ["url"]
|
||||
}
|
||||
|
||||
@@ -58,7 +58,7 @@ def test_content_fetcher_extracts_og_image_and_body_fallback(module, tmp_path, m
|
||||
|
||||
monkeypatch.setattr(module, "CONTENT_CACHE_DIR", tmp_path)
|
||||
module.content_cache_index.clear()
|
||||
monkeypatch.setattr(module, "_get_public_url", lambda url, headers, timeout: _FakeResponse(html))
|
||||
monkeypatch.setattr(module, "_get_public_url", lambda url, headers, timeout, **kwargs: _FakeResponse(html))
|
||||
|
||||
result = module.fetch_webpage_content("https://example.com/parity-test")
|
||||
|
||||
@@ -82,7 +82,7 @@ def test_fetch_webpage_content_returns_empty_result_on_http_status_error(status_
|
||||
monkeypatch.setattr(
|
||||
service_content,
|
||||
"_get_public_url",
|
||||
lambda url, headers, timeout: _FakeErrorResponse(status_code),
|
||||
lambda url, headers, timeout, **kwargs: _FakeErrorResponse(status_code),
|
||||
)
|
||||
|
||||
result = service_content.fetch_webpage_content(f"https://example.com/status-{status_code}")
|
||||
@@ -119,7 +119,7 @@ def test_fetch_webpage_content_429_takes_distinct_rate_limit_path(tmp_path, monk
|
||||
monkeypatch.setattr(
|
||||
service_content,
|
||||
"_get_public_url",
|
||||
lambda url, headers, timeout: _FakeRateLimitResponse(),
|
||||
lambda url, headers, timeout, **kwargs: _FakeRateLimitResponse(),
|
||||
)
|
||||
|
||||
result = service_content.fetch_webpage_content("https://example.com/rate-limited")
|
||||
|
||||
@@ -904,7 +904,13 @@ def test_web_fetch_guard_blocks_redirect_into_private(monkeypatch):
|
||||
url = "http://public.example/start"
|
||||
headers = {"location": "http://169.254.169.254/latest/meta-data/"}
|
||||
|
||||
monkeypatch.setattr(httpx, "get", lambda url, **kwargs: _Resp())
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def _fake_stream(method, url, **kwargs):
|
||||
yield _Resp()
|
||||
|
||||
monkeypatch.setattr(httpx, "stream", _fake_stream)
|
||||
|
||||
with _pytest.raises(httpx.RequestError) as exc:
|
||||
content._get_public_url("http://public.example/start", headers={}, timeout=5)
|
||||
|
||||
@@ -35,7 +35,7 @@ def _patch_fetch(monkeypatch, text, content_type):
|
||||
monkeypatch.setattr(
|
||||
content_mod,
|
||||
"_get_public_url",
|
||||
lambda url, headers=None, timeout=5: _FakeResponse(text, content_type),
|
||||
lambda url, headers=None, timeout=5, **kwargs: _FakeResponse(text, content_type),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
"""web_fetch download budgets (#3812).
|
||||
|
||||
MAX_OUTPUT_CHARS only trims what the agent sees; these caps bound what the
|
||||
server downloads, parses, and caches. Soft cap by default with a truncation
|
||||
notice, per-call override clamped to the hard cap, and a pre-buffer refusal
|
||||
when Content-Length already exceeds the hard ceiling.
|
||||
"""
|
||||
import json
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from src.constants import WEB_FETCH_SOFT_MAX_BYTES, WEB_FETCH_HARD_MAX_BYTES
|
||||
from services.search import content as content_mod
|
||||
|
||||
|
||||
class _FakeStream:
|
||||
"""Stands in for the httpx.stream(...) context manager."""
|
||||
|
||||
def __init__(self, body: bytes, content_type="text/plain", content_length=None,
|
||||
status_code=200, chunk=8192):
|
||||
self._body = body
|
||||
self._chunk = chunk
|
||||
self.status_code = status_code
|
||||
self.encoding = "utf-8"
|
||||
self.url = "https://example.com/x"
|
||||
self.headers = {"Content-Type": content_type}
|
||||
if content_length is not None:
|
||||
self.headers["content-length"] = str(content_length)
|
||||
self.body_reads = 0
|
||||
|
||||
def iter_bytes(self):
|
||||
for i in range(0, len(self._body), self._chunk):
|
||||
self.body_reads += 1
|
||||
yield self._body[i:i + self._chunk]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_cache(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(content_mod, "CONTENT_CACHE_DIR", tmp_path)
|
||||
monkeypatch.setattr(content_mod, "_cache_result", lambda *a, **k: None)
|
||||
monkeypatch.setattr(content_mod, "_public_http_url", lambda u: True)
|
||||
|
||||
|
||||
def _patch_stream(monkeypatch, fake):
|
||||
@contextmanager
|
||||
def fake_stream(method, url, **kwargs):
|
||||
yield fake
|
||||
monkeypatch.setattr(content_mod.httpx, "stream", fake_stream)
|
||||
return fake
|
||||
|
||||
|
||||
def test_body_under_cap_is_untouched(monkeypatch, no_cache):
|
||||
_patch_stream(monkeypatch, _FakeStream(b"hello world"))
|
||||
r = content_mod.fetch_webpage_content("https://example.com/a.txt")
|
||||
assert r["success"] is True
|
||||
assert r["content"] == "hello world"
|
||||
assert r["truncated"] is False
|
||||
assert r["fetched_bytes"] == len(b"hello world")
|
||||
|
||||
|
||||
def test_body_over_soft_cap_truncates_with_flags(monkeypatch, no_cache):
|
||||
body = b"x" * (WEB_FETCH_SOFT_MAX_BYTES + 50_000)
|
||||
_patch_stream(monkeypatch, _FakeStream(body, content_length=len(body)))
|
||||
r = content_mod.fetch_webpage_content("https://example.com/big.txt")
|
||||
assert r["truncated"] is True
|
||||
assert r["fetched_bytes"] == WEB_FETCH_SOFT_MAX_BYTES
|
||||
assert r["total_bytes"] == len(body)
|
||||
assert len(r["content"]) == WEB_FETCH_SOFT_MAX_BYTES
|
||||
|
||||
|
||||
def test_max_bytes_override_raises_budget(monkeypatch, no_cache):
|
||||
body = b"y" * (WEB_FETCH_SOFT_MAX_BYTES + 50_000)
|
||||
_patch_stream(monkeypatch, _FakeStream(body))
|
||||
r = content_mod.fetch_webpage_content(
|
||||
"https://example.com/big.txt", max_bytes=len(body) + 1
|
||||
)
|
||||
assert r["truncated"] is False
|
||||
assert r["fetched_bytes"] == len(body)
|
||||
|
||||
|
||||
def test_override_is_clamped_to_hard_cap(monkeypatch, no_cache):
|
||||
# Ask for more than the ceiling; the effective budget must be the ceiling.
|
||||
fake = _patch_stream(monkeypatch, _FakeStream(b"z" * 10, chunk=4))
|
||||
r = content_mod.fetch_webpage_content(
|
||||
"https://example.com/a.txt", max_bytes=WEB_FETCH_HARD_MAX_BYTES * 10
|
||||
)
|
||||
assert r["success"] is True
|
||||
# The clamp itself: effective cap recorded in the cache key path is the
|
||||
# hard cap, and a declared body over the ceiling is refused regardless.
|
||||
big = _FakeStream(b"", content_length=WEB_FETCH_HARD_MAX_BYTES + 1)
|
||||
_patch_stream(monkeypatch, big)
|
||||
r = content_mod.fetch_webpage_content(
|
||||
"https://example.com/huge.bin", max_bytes=WEB_FETCH_HARD_MAX_BYTES * 10
|
||||
)
|
||||
assert r["success"] is False
|
||||
assert "TooLarge" in r["error"]
|
||||
assert big.body_reads == 0 # refused before buffering
|
||||
|
||||
|
||||
def test_declared_over_hard_cap_refused_before_buffering(monkeypatch, no_cache):
|
||||
fake = _FakeStream(b"irrelevant", content_length=WEB_FETCH_HARD_MAX_BYTES + 1)
|
||||
_patch_stream(monkeypatch, fake)
|
||||
r = content_mod.fetch_webpage_content("https://example.com/huge.iso")
|
||||
assert r["success"] is False
|
||||
assert "TooLarge" in r["error"]
|
||||
assert fake.body_reads == 0
|
||||
|
||||
|
||||
def test_truncated_pdf_is_an_error_not_garbage(monkeypatch, no_cache):
|
||||
body = b"%PDF-1.4 " + b"p" * (WEB_FETCH_SOFT_MAX_BYTES + 10)
|
||||
_patch_stream(monkeypatch, _FakeStream(body, content_type="application/pdf"))
|
||||
r = content_mod.fetch_webpage_content("https://example.com/big.pdf")
|
||||
assert r["success"] is False
|
||||
assert "TooLarge" in r["error"]
|
||||
|
||||
|
||||
def test_fetch_requests_identity_encoding(monkeypatch, no_cache):
|
||||
# Compressed responses can decode to far more than Content-Length, so the
|
||||
# streamed cap and the hard-cap preflight are only honest when we refuse
|
||||
# transfer compression. Pin that the fetch advertises identity, not gzip.
|
||||
seen = {}
|
||||
|
||||
@contextmanager
|
||||
def fake_stream(method, url, **kwargs):
|
||||
seen["headers"] = kwargs.get("headers") or {}
|
||||
yield _FakeStream(b"hello")
|
||||
monkeypatch.setattr(content_mod.httpx, "stream", fake_stream)
|
||||
|
||||
content_mod.fetch_webpage_content("https://example.com/a.txt")
|
||||
assert seen["headers"].get("Accept-Encoding") == "identity"
|
||||
|
||||
|
||||
def test_rejects_compressed_response_that_ignored_identity(monkeypatch, no_cache):
|
||||
# We request Accept-Encoding: identity, but a server can ignore it and send
|
||||
# gzip anyway. httpx would decode it, so a tiny compressed body could balloon
|
||||
# past the cap in one decoded chunk. Refuse before reading the body.
|
||||
fake = _FakeStream(b"x" * 5000, content_length=40)
|
||||
fake.headers["content-encoding"] = "gzip"
|
||||
_patch_stream(monkeypatch, fake)
|
||||
r = content_mod.fetch_webpage_content("https://example.com/a.txt")
|
||||
assert r["success"] is False
|
||||
assert "Content-Encoding" in r["error"] or "compressed" in r["error"]
|
||||
assert fake.body_reads == 0 # refused before decoding any body
|
||||
|
||||
|
||||
def test_oversized_title_does_not_hide_partial_notice(monkeypatch):
|
||||
# The partial-content notice is the PR's core contract; an untrusted,
|
||||
# oversized page title must not push it past MAX_OUTPUT_CHARS.
|
||||
import asyncio
|
||||
from src.agent_tools.web_tools import WebFetchTool
|
||||
from src.constants import MAX_OUTPUT_CHARS
|
||||
|
||||
def fake_fetch(url, timeout=10, max_bytes=None):
|
||||
return {
|
||||
"content": "partial body",
|
||||
"title": "T" * (MAX_OUTPUT_CHARS + 5_000),
|
||||
"error": "",
|
||||
"truncated": True,
|
||||
"fetched_bytes": WEB_FETCH_SOFT_MAX_BYTES,
|
||||
"total_bytes": 9_000_000,
|
||||
}
|
||||
|
||||
import src.search.content as alias_mod
|
||||
monkeypatch.setattr(alias_mod, "fetch_webpage_content", fake_fetch)
|
||||
|
||||
out = asyncio.run(WebFetchTool().execute(
|
||||
json.dumps({"url": "https://example.com/big.txt"}), ctx={}
|
||||
))
|
||||
assert out["exit_code"] == 0
|
||||
assert out["output"].startswith("[partial content:")
|
||||
assert '"full": true' in out["output"]
|
||||
|
||||
|
||||
def test_tool_layer_emits_partial_notice_and_parses_full(monkeypatch):
|
||||
import asyncio
|
||||
from src.agent_tools.web_tools import WebFetchTool
|
||||
|
||||
calls = {}
|
||||
|
||||
def fake_fetch(url, timeout=10, max_bytes=None):
|
||||
calls["max_bytes"] = max_bytes
|
||||
return {
|
||||
"content": "partial body",
|
||||
"title": "Big File",
|
||||
"error": "",
|
||||
"truncated": True,
|
||||
"fetched_bytes": WEB_FETCH_SOFT_MAX_BYTES,
|
||||
"total_bytes": 5_000_000,
|
||||
}
|
||||
|
||||
import src.search.content as alias_mod
|
||||
monkeypatch.setattr(alias_mod, "fetch_webpage_content", fake_fetch)
|
||||
|
||||
out = asyncio.run(WebFetchTool().execute(
|
||||
json.dumps({"url": "https://example.com/big.txt"}), ctx={}
|
||||
))
|
||||
assert out["exit_code"] == 0
|
||||
assert "[partial content:" in out["output"]
|
||||
assert '"full": true' in out["output"]
|
||||
assert calls["max_bytes"] is None
|
||||
|
||||
asyncio.run(WebFetchTool().execute(
|
||||
json.dumps({"url": "https://example.com/big.txt", "full": True}), ctx={}
|
||||
))
|
||||
assert calls["max_bytes"] == WEB_FETCH_HARD_MAX_BYTES
|
||||
Reference in New Issue
Block a user