diff --git a/backend/src/agents/builtin_tools/document_read_tool.py b/backend/src/agents/builtin_tools/document_read_tool.py index e99ab8b9..637cdf11 100644 --- a/backend/src/agents/builtin_tools/document_read_tool.py +++ b/backend/src/agents/builtin_tools/document_read_tool.py @@ -23,6 +23,14 @@ * One call can never re-inject a whole document: page ranges are capped by ``max_pages`` (default 8, hard cap 20), pattern results by match count, text by the workspace read bound. +* A pattern scan is bounded in *time* as well as in output. The regex is + model-supplied, and a nested-quantifier pattern makes Python's backtracking + engine run exponentially — measured at 14 s for a 28-character line, against + extracted PDF lines of 60–100 — with no way to cancel a running + ``re.search``. ``catastrophic_pattern`` refuses that family up front (the + pattern is searched literally, and the payload says so) and + ``DOCUMENT_READ_PATTERN_BUDGET_SECONDS`` bounds the walk for anything it + does not catch. * **Content-free record per call** in ``AgentCoreStack/Compaction`` (``DocumentRead``, ``DocumentReadPages``, ``DocumentReadBytes``; properties ``mode`` / ``format``) — the same namespace as the compaction cut and @@ -141,6 +149,10 @@ async def document_read( page_range: 1-indexed inclusive page range for PDFs, "start-end" or a single page "5". pattern: Regular expression (case-insensitive) to search for. + Do not nest one unbounded quantifier inside another — write + "X+" rather than "(X+)+" — because such a pattern can take + exponential time; it is searched as literal text instead and + the result says so. max_pages: Cap on pages returned by page_range (default 8, hard cap 20). offset: Byte offset to continue a truncated text read. diff --git a/backend/src/agents/main_agent/session/document_context.py b/backend/src/agents/main_agent/session/document_context.py index 82e240c3..257148fe 100644 --- a/backend/src/agents/main_agent/session/document_context.py +++ b/backend/src/agents/main_agent/session/document_context.py @@ -14,9 +14,9 @@ keyed by Bedrock's ``document.format`` enum (``pdf`` / ``docx`` / …) plus ``image``. -Token numbers are the compaction estimator's heuristics (bytes/4 for a -document, a flat figure per image), not Bedrock counts: Bedrock reports no -per-block usage. They are comparable across rows and against +Token numbers are heuristics (``document_tokens.estimate_document_tokens``: +pages x the per-page image estimate for PDFs, bytes/4 otherwise; a flat figure +per image), not Bedrock counts: Bedrock reports no per-block usage. They are comparable across rows and against ``contextBreakdown.messages``, which is the measured total they are a share of. """ @@ -24,7 +24,8 @@ from typing import Any, Dict, List, Optional -from agents.main_agent.session.compaction_policy import CHARS_PER_TOKEN, IMAGE_TOKEN_ESTIMATE +from agents.main_agent.session.compaction_policy import IMAGE_TOKEN_ESTIMATE +from apis.shared.files.document_tokens import estimate_document_tokens #: A restore-time stand-in for a document that is no longer inline — today the #: contentless placeholder ``_strip_document_bytes`` writes; from PR-3 the @@ -101,8 +102,10 @@ def summarize_document_context(messages: Optional[List[Dict[str, Any]]]) -> Opti if size is not None: count += 1 attached_here += 1 - tokens += size // CHARS_PER_TOKEN fmt = str((block.get("document") or {}).get("format") or "unknown") + tokens += estimate_document_tokens( + fmt, ((block.get("document") or {}).get("source") or {}).get("bytes") + ) mime[fmt] = mime.get(fmt, 0) + 1 continue size = _inline_bytes(block, "image") @@ -123,7 +126,10 @@ def summarize_document_context(messages: Optional[List[Dict[str, Any]]]) -> Opti inner_size = _inline_bytes(inner, "document") if inner_size is not None: slices += 1 - slice_tokens += inner_size // CHARS_PER_TOKEN + inner_doc = inner.get("document") or {} + slice_tokens += estimate_document_tokens( + inner_doc.get("format"), (inner_doc.get("source") or {}).get("bytes") + ) if prompt: last_prompt_attachments = attached_here diff --git a/backend/src/agents/main_agent/session/document_offload.py b/backend/src/agents/main_agent/session/document_offload.py index f89e9bd9..c4ac0661 100644 --- a/backend/src/agents/main_agent/session/document_offload.py +++ b/backend/src/agents/main_agent/session/document_offload.py @@ -14,9 +14,12 @@ subject: on the attach turn and the next ``pin_turns - 1`` turns; while the incoming or previous prompt names it; while a ``document_read`` result for it sits in the recent turns. -* **Worth it** — the document's estimated tokens (bytes/4) are at least +* **Worth it** — the document's estimated tokens are at least ``DOCUMENT_OFFLOAD_MIN_TOKENS``; below that the re-write costs more than the - eviction saves (the ``clear_at_least`` idea). + eviction saves (the ``clear_at_least`` idea). The estimate is + ``document_tokens.estimate_document_tokens`` — pages x the per-page image + estimate for a PDF, bytes/4 otherwise — *not* bytes/4 for everything, which + under-counted PDFs ~14x and held large scans below the floor forever. * **Free or unavoidable** — decided by the session manager, which owns the cache-gap facts: the prompt cache has expired since the last turn, the model/agent prefix changed, or the last turn's input exceeded the @@ -50,6 +53,7 @@ from agents.main_agent.multimodal.file_sanitizer import FileSanitizer from agents.main_agent.session.compaction_policy import CHARS_PER_TOKEN +from apis.shared.files.document_tokens import estimate_document_tokens logger = logging.getLogger(__name__) @@ -91,10 +95,23 @@ def session_bucket(session_id: str) -> int: def offload_enabled_for(session_id: Optional[str]) -> bool: - """Kill switch and rollout bucket together. Sessions below the percent - are the treated arm; the rest keep today's inline-forever behavior.""" + """Kill switch, the ``document_read`` gate, and the rollout bucket + together. Sessions below the percent are the treated arm; the rest keep + today's inline-forever behavior. + + ``DOCUMENT_READ_ENABLED=false`` disables this path too. The live offload + is the one place bytes leave the prefix *optionally* — restore has to drop + them, this does not — so evicting a document while its only recovery path + is switched off would be strictly worse than the pre-offload world, in + which live bytes never left. Spec §5: "a digest that points at a tool + nobody has is no better than today's placeholder." + """ + from apis.shared.feature_flags import document_read_enabled + if not document_offload_enabled() or not session_id: return False + if not document_read_enabled(): + return False return session_bucket(session_id) < rollout_percent() @@ -153,6 +170,17 @@ def _inline_document(block: Any) -> Optional[Tuple[str, str, int]]: return str(doc.get("name", "")), str(doc.get("format", "")), len(raw) +def _inline_document_bytes(block: Any) -> Optional[bytes]: + """The block's raw bytes, for weighing it (``estimate_document_tokens``).""" + if not isinstance(block, dict): + return None + doc = block.get("document") + if not isinstance(doc, dict): + return None + raw = (doc.get("source") or {}).get("bytes") if isinstance(doc.get("source"), dict) else None + return raw if isinstance(raw, (bytes, bytearray)) else None + + def pinned_document_names( messages: Sequence[Dict[str, Any]], incoming_prompt: Any = None, @@ -208,10 +236,11 @@ class Candidate: name: str format: str size: int - - @property - def tokens(self) -> int: - return self.size // CHARS_PER_TOKEN + #: Estimated prefix weight — pages x the per-page image estimate for PDFs, + #: bytes/4 otherwise. This, not ``size``, is what the ``min_tokens`` floor + #: and the ledger's ``documentTokens`` are measured against, so the row and + #: the eviction decision always read the same quantity. + tokens: int = 0 def candidate_documents( @@ -231,9 +260,12 @@ def candidate_documents( if not inline: continue name, fmt, size = inline - if name in pinned or size // CHARS_PER_TOKEN < min_tokens: + if name in pinned: + continue + tokens = estimate_document_tokens(fmt, _inline_document_bytes(block)) + if tokens < min_tokens: continue - out.append(Candidate(mi, bi, name, fmt, size)) + out.append(Candidate(mi, bi, name, fmt, size, tokens)) return out @@ -264,6 +296,7 @@ def offload_documents( restore path's own matcher and renderer so the bytes equal what a cold restore would produce. Unmatched candidates stay inline. Never raises.""" from agents.main_agent.session.document_rehydration import digest_for, load_session_documents, match_document + from apis.shared.feature_flags import document_read_enabled from apis.shared.files.document_digest import render_digest result = OffloadResult() @@ -293,7 +326,10 @@ def offload_documents( if digest is None: result.skipped_unmatched += 1 continue - text = render_digest(digest, filename=meta.filename, upload_id=meta.upload_id) + text = render_digest( + digest, filename=meta.filename, upload_id=meta.upload_id, + include_handle=document_read_enabled(), + ) content[cand.block_index] = {"text": text} used.add(meta.upload_id) result.offloaded += 1 @@ -340,10 +376,10 @@ def age_document_slices( inline = _inline_document(inner) if not inline: continue - name, _fmt, size = inline + name, fmt, _size = inline + tokens += estimate_document_tokens(fmt, _inline_document_bytes(inner)) inner_content[ii] = {"text": slice_stub(name)} count += 1 - tokens += size // CHARS_PER_TOKEN return count, tokens diff --git a/backend/src/agents/main_agent/session/document_rehydration.py b/backend/src/agents/main_agent/session/document_rehydration.py index 5991fcc3..d6cd5418 100644 --- a/backend/src/agents/main_agent/session/document_rehydration.py +++ b/backend/src/agents/main_agent/session/document_rehydration.py @@ -27,6 +27,12 @@ re-write, accepted and recorded (``document_rehydrated`` carries the digest tokens, so the change is visible). +When ``DOCUMENT_READ_ENABLED=false`` has taken the tool away the digest is +still rendered — it is strictly more than the placeholder — but without the +``upload_id`` handle, so the model is never invited to call a tool it does +not have. The live offload path stops entirely in that case; see +``document_offload.offload_enabled_for``. + Everything is synchronous: it runs inside ``TurnBasedSessionManager.initialize`` under the Strands agent constructor. It never raises — every failure falls back to the placeholder for that block. @@ -43,6 +49,7 @@ from agents.main_agent.multimodal.file_sanitizer import FileSanitizer from agents.main_agent.session.compaction_policy import CHARS_PER_TOKEN +from apis.shared.files.document_tokens import estimate_document_tokens logger = logging.getLogger(__name__) @@ -166,9 +173,16 @@ def rehydrate_documents( an upload row) or, failing that, the placeholder. Never raises.""" from apis.shared.files.document_digest import render_digest + from apis.shared.feature_flags import document_read_enabled + out = copy.deepcopy(messages) result = RehydrationResult(messages=out) enabled = document_rehydration_enabled() and bool(session_id) + # Restore has to drop the bytes either way (Bedrock rejects duplicate + # document names), so a digest is still strictly better than the + # placeholder when the tool is off — but it must not advertise a handle + # the model cannot use. + handle = document_read_enabled() candidates: Optional[List[Any]] = None used: Set[str] = set() @@ -198,7 +212,10 @@ def rehydrate_documents( digest = digest_for(meta, bytes(raw) if isinstance(raw, (bytes, bytearray)) else b"") if digest is not None: was_lazy = not (isinstance(getattr(meta, "digest", None), dict) and meta.digest.get("status") == "ready") - replacement = render_digest(digest, filename=meta.filename, upload_id=meta.upload_id) + replacement = render_digest( + digest, filename=meta.filename, upload_id=meta.upload_id, + include_handle=handle, + ) used.add(meta.upload_id) result.rehydrated += 1 result.digest_tokens += len(replacement) // CHARS_PER_TOKEN @@ -212,7 +229,10 @@ def rehydrate_documents( if replacement is None: replacement = placeholder_text(name, fmt, size) result.stripped += 1 - result.stripped_tokens += size // CHARS_PER_TOKEN + # Same estimator the row and the offloader use, so the + # ``document_stripped`` ledger event is comparable to + # ``document_offload`` / ``document_rehydrated``. + result.stripped_tokens += estimate_document_tokens(fmt, raw) content[idx] = {"text": replacement} return result diff --git a/backend/src/agents/main_agent/session/hooks/context_attribution.py b/backend/src/agents/main_agent/session/hooks/context_attribution.py index c997ec4b..0afc1523 100644 --- a/backend/src/agents/main_agent/session/hooks/context_attribution.py +++ b/backend/src/agents/main_agent/session/hooks/context_attribution.py @@ -25,6 +25,21 @@ every turn afterward is pure arithmetic against the free, authoritative ``projected_input_tokens``. +**Why the split is not computed while an attachment is in context.** +``toolTokens`` is a *residual* between two independently sourced numbers — +``full`` (Strands' projection for the upcoming request) and ``no_tools`` (our own +CountTokens call) — so any disagreement between them about how a content block +is counted lands wholly in it. Bedrock understands a PDF page as an image *and* +a text layer; when the two sources do not agree on that, the document's entire +weight is attributed to tools. Measured on dev 2026-09-16 (session +``61de2256``): a call reported ``toolTokens`` of **106,756** where the session's +real tools prefix was **12,516** — a difference of 94,240 against a document +measured at ~94,485, i.e. the whole document. The split is therefore skipped on +any turn whose context carries inline document or image bytes, and taken on a +later clean turn instead. An absent ``prefixTokens`` reads "not tracked" (the +ledger's convention); a wrong one silently corrupts every share computed from +it. + Best-effort: any failure is swallowed so context attribution can never break a model call. For non-Bedrock models ``count_tokens`` falls back to a heuristic, so the numbers are approximate there. @@ -42,6 +57,31 @@ _BREAKDOWN_ATTR = "_context_attribution_breakdown" # latest per-turn breakdown dict +def _has_inline_attachment(messages: Any) -> bool: + """Whether any message carries inline ``document`` / ``image`` bytes. + + The condition under which ``toolTokens`` cannot be trusted — see the module + docstring. Cheap: a walk over content blocks, no decoding.""" + if not isinstance(messages, list): + return False + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + continue + for block in content: + if not isinstance(block, dict): + continue + for key in ("document", "image"): + payload = block.get(key) + if isinstance(payload, dict): + source = payload.get("source") + if isinstance(source, dict) and isinstance( + source.get("bytes"), (bytes, bytearray) + ): + return True + return False + + def get_context_breakdown(agent: Any) -> Optional[dict]: """Return the latest context breakdown stashed on ``agent``, or ``None``. @@ -90,6 +130,11 @@ async def _compute(self, event: BeforeModelCallEvent) -> None: full = event.projected_input_tokens split = getattr(agent, _SPLIT_ATTR, None) + if split is None and _has_inline_attachment(agent.messages): + # Untrustworthy residual (see module docstring) — leave the split + # uncomputed and try again on a turn without inline bytes. + logger.debug("Context attribution deferred: inline attachment in context") + return if split is None: system_tokens = await model.count_tokens( messages=[], diff --git a/backend/src/apis/shared/files/document_digest.py b/backend/src/apis/shared/files/document_digest.py index ef1aef0e..c7b47467 100644 --- a/backend/src/apis/shared/files/document_digest.py +++ b/backend/src/apis/shared/files/document_digest.py @@ -22,7 +22,9 @@ The rendered form (``render_digest``) is what enters context later. It is hard-capped at ``DOCUMENT_DIGEST_MAX_TOKENS`` (default 1,500, chars/4): the -outline is trimmed from the end first, then the abstract. Every digest +outline is trimmed from the end first, then the abstract — escaped *before* it +is cut, so the cap is measured on what is actually rendered, and dropped +entirely rather than allowed to overshoot. Every digest records its rendered token estimate so the cap is a stored fact per file. Content policy: the digest carries model prose about the user's document @@ -287,6 +289,24 @@ async def generate_abstract(outline: DocumentDigest, sample: str, model_id: str # --------------------------------------------------------------------------- +def _truncate_escaped(text: str, limit: int) -> str: + """Cut an already-escaped string to ``limit`` characters without splitting + an entity — ``&`` must never be left as ``&am``. + + Escaping *after* slicing (what this replaces) silently broke the budget: + each ``&`` becomes five characters, so a slice measured on the raw text + could render up to 5x longer. Measured: an abstract of ``&`` rendered 388 + tokens against a 100-token budget. + """ + if limit <= 0: + return "" + cut = text[:limit] + amp = cut.rfind("&") + if amp != -1 and ";" not in cut[amp:]: + cut = cut[:amp] + return cut.rstrip() + + def _xml_escape(text: str) -> str: return text.replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """) @@ -301,15 +321,23 @@ def render_digest( filename: str, upload_id: str, budget_tokens: int = DOCUMENT_DIGEST_MAX_TOKENS, + include_handle: bool = True, ) -> str: """The ```` block, trimmed to ``budget_tokens``. Sections are dropped from the end first (the model can always ask ``document_read`` for more), then the abstract is truncated. The opening tag always fits: it is the ``document_read`` handle. + + ``include_handle=False`` omits ``upload_id``. Callers pass it when + ``DOCUMENT_READ_ENABLED=false`` has taken the tool away: the outline and + abstract are still strictly more than the pre-PR-3 placeholder, but + advertising a retrieval id for a tool the model does not have would invite + a call that cannot be made. See ``feature_flags.document_read_enabled``. """ + handle = f'upload_id="{upload_id}" ' if include_handle else "" header = ( - f' str: if len(text) > budget_chars and abstract_line: room = budget_chars - len(_build("", kept)) - len(" ") - 1 if room > 20 and digest.abstract: - trimmed = _xml_escape(digest.abstract[: max(0, room - 1)].rstrip()) + "…" + # Escape first, then cut: the budget is measured on what is + # actually rendered, not on the raw text it came from. + trimmed = _truncate_escaped(_xml_escape(digest.abstract), room - 1) + "…" text = _build(f" {trimmed}", kept) else: text = _build("", kept) + if len(text) > budget_chars: + # The cap is hard. Drop the abstract entirely rather than overshoot; + # only the opening tag — the ``document_read`` handle — is allowed to + # survive a budget this small. + text = _build("", kept) return text diff --git a/backend/src/apis/shared/files/document_read.py b/backend/src/apis/shared/files/document_read.py index fc90e6dc..b72a27a4 100644 --- a/backend/src/apis/shared/files/document_read.py +++ b/backend/src/apis/shared/files/document_read.py @@ -40,6 +40,7 @@ import logging import os import re +import time import uuid import zipfile from dataclasses import dataclass, field @@ -76,6 +77,14 @@ DOCUMENT_READ_MAX_MATCHES = int(os.environ.get("DOCUMENT_READ_MAX_MATCHES", 40)) #: Index mode lists a snippet for at most this many pages. DOCUMENT_READ_INDEX_PAGES = int(os.environ.get("DOCUMENT_READ_INDEX_PAGES", 40)) +#: Wall-clock bound on one pattern scan. The backstop behind +#: ``catastrophic_pattern``: a single ``re.search`` cannot be cancelled, but the +#: walk across lines and pages can be stopped, so a merely-slow pattern over a +#: long document degrades to a partial answer instead of hanging the turn to the +#: 600 s SSE timeout. ``0`` disables the bound. +DOCUMENT_READ_PATTERN_BUDGET_SECONDS = float( + os.environ.get("DOCUMENT_READ_PATTERN_BUDGET_SECONDS", 2.0) +) #: Characters kept per matching / snippet line. _LINE_CHARS = 200 _SNIPPET_CHARS = 90 @@ -398,19 +407,22 @@ def _pdf_pages(raw: bytes, pages: Tuple[int, int], limit: int, base: Dict[str, A def _pdf_pattern(raw: bytes, pattern: str, base: Dict[str, Any]) -> DocumentReadResult: - regex = _compile(pattern) + regex, note = _compile(pattern) + clock = _Budget(DOCUMENT_READ_PATTERN_BUDGET_SECONDS) pdf = _open_pdf(raw) + pages_searched = 0 try: count = len(pdf) matches: List[Dict[str, Any]] = [] pages_matched: List[int] = [] for index in range(count): - if len(matches) >= DOCUMENT_READ_MAX_MATCHES: + if len(matches) >= DOCUMENT_READ_MAX_MATCHES or clock.out_of_time(): break + pages_searched = index + 1 lines = _pdf_page_text(pdf, index).splitlines() page_no = index + 1 hit = False - for entry in _grep_lines(lines, regex, DOCUMENT_READ_MAX_MATCHES - len(matches)): + for entry in _grep_lines(lines, regex, DOCUMENT_READ_MAX_MATCHES - len(matches), clock): entry["page"] = page_no matches.append(entry) hit = True @@ -429,6 +441,15 @@ def _pdf_pattern(raw: bytes, pattern: str, base: Dict[str, Any]) -> DocumentRead "matches": matches, "hint": "Call again with page_range to read the matching pages at full fidelity.", } + if note: + payload["pattern_note"] = note + if clock.expired: + payload["timed_out"] = True + payload["pages_searched"] = pages_searched + payload["hint"] = ( + f"The search ran out of time after {pages_searched} of {count} pages; these are " + "the matches found so far. Use a simpler pattern, or read a page_range directly." + ) return DocumentReadResult(mode="pattern", payload=payload, pages_returned=0, format="pdf") @@ -525,18 +546,130 @@ def _docx_read(raw: bytes, pattern: Optional[str], offset: int, base: Dict[str, # --------------------------------------------------------------------------- -def _compile(pattern: str) -> "re.Pattern[str]": +def _scan_flags(text: str): + """Yield ``(index, char)`` for the regex-significant characters of ``text`` + — skipping escaped characters and the contents of ``[...]`` classes, where + ``*``, ``+``, ``|`` and parentheses are literals.""" + i, in_class, n = 0, False, len(text) + while i < n: + ch = text[i] + if ch == "\\": + i += 2 + continue + if in_class: + if ch == "]": + in_class = False + i += 1 + continue + if ch == "[": + in_class = True + i += 1 + continue + yield i, ch + i += 1 + + +def _unbounded_quantifier_at(pattern: str, i: int) -> bool: + """Is there an unbounded quantifier (``*``, ``+`` or ``{n,}``) at ``i``?""" + if i >= len(pattern): + return False + if pattern[i] in "*+": + return True + if pattern[i] == "{": + close = pattern.find("}", i) + return close != -1 and pattern[i + 1: close].endswith(",") + return False + + +def _has_unbounded_quantifier(body: str) -> bool: + return any( + _unbounded_quantifier_at(body, i) for i, ch in _scan_flags(body) if ch in "*+{" + ) + + +def catastrophic_pattern(pattern: str) -> bool: + """Whether ``pattern`` has the nested-quantifier shape that makes Python's + backtracking engine run in exponential time. + + ``(a+)+``, ``(\\w+\\s*)*``, ``(a+|b){2,}`` — an unbounded quantifier applied to + a group that itself contains one. Measured on the real engine: ``(a+)+$`` + against ``"a"*24 + "!"`` takes 0.87 s, 26 takes 3.5 s, 28 takes 14 s. + Extracted PDF lines run 60–100 characters, so such a pattern never returns. + + Deliberately narrow. It flags only the unambiguous family, so ordinary + patterns — ``\\d+``, ``(invoice|receipt)``, ``(foo)+`` — are never refused. + Shapes it does not catch (overlapping literal alternations like ``(a|a)+``, + backreference blowups) are bounded by the wall-clock budget instead, not by + this test. + """ + if not pattern: + return False + starts: List[int] = [] + for i, ch in _scan_flags(pattern): + if ch == "(": + starts.append(i) + elif ch == ")" and starts: + body = pattern[starts.pop() + 1: i] + if _unbounded_quantifier_at(pattern, i + 1) and _has_unbounded_quantifier(body): + return True + return False + + +def _compile(pattern: str) -> Tuple["re.Pattern[str]", Optional[str]]: + """``(regex, note)``. The note is non-``None`` when the pattern was searched + literally instead of as a regex, and is carried to the model on the payload + so a silently different result is never presented as the requested one. + + Two reasons to fall back, both degrading rather than erroring: the pattern + does not compile, or it has the nested-quantifier shape that would hang the + turn (``catastrophic_pattern``).""" + if catastrophic_pattern(pattern): + logger.warning("document_read: refusing catastrophic pattern, searching literally") + return re.compile(re.escape(pattern), re.IGNORECASE), ( + "This pattern nests one unbounded quantifier inside another, which can take " + "exponential time to match, so it was searched as literal text instead. " + "Rewrite it without the nesting (for example '\\w+' rather than '(\\w+)+')." + ) try: - return re.compile(pattern, re.IGNORECASE) + return re.compile(pattern, re.IGNORECASE), None except re.error: - return re.compile(re.escape(pattern), re.IGNORECASE) + return re.compile(re.escape(pattern), re.IGNORECASE), ( + "This pattern is not a valid regular expression, so it was searched as literal text." + ) + + +class _Budget: + """Wall-clock bound on one pattern scan. + The backstop behind ``catastrophic_pattern``: it cannot interrupt a single + ``re.search`` — CPython exposes no way to cancel one, and the scan runs in a + worker thread where signals are unavailable — but it stops the *walk*, so a + merely-slow pattern over a 200-page document cannot run away. A scan that + runs out reports ``timed_out`` with the matches it already has. + """ + + __slots__ = ("_deadline", "expired") + + def __init__(self, seconds: float = 0.0) -> None: + self._deadline = (time.monotonic() + seconds) if seconds > 0 else None + self.expired = False -def _grep_lines(lines: Sequence[str], regex: "re.Pattern[str]", budget: int) -> List[Dict[str, Any]]: + def out_of_time(self) -> bool: + if self._deadline is not None and time.monotonic() > self._deadline: + self.expired = True + return self.expired + + +def _grep_lines( + lines: Sequence[str], + regex: "re.Pattern[str]", + budget: int, + clock: Optional[_Budget] = None, +) -> List[Dict[str, Any]]: """Matching lines with ``_CONTEXT_LINES`` of context, 1-indexed.""" out: List[Dict[str, Any]] = [] for i, line in enumerate(lines): - if len(out) >= budget: + if len(out) >= budget or (clock is not None and clock.out_of_time()): break if not regex.search(line): continue @@ -547,9 +680,10 @@ def _grep_lines(lines: Sequence[str], regex: "re.Pattern[str]", budget: int) -> def _text_pattern(text: str, pattern: str, base: Dict[str, Any], unit: str) -> DocumentReadResult: - regex = _compile(pattern) + regex, note = _compile(pattern) + clock = _Budget(DOCUMENT_READ_PATTERN_BUDGET_SECONDS) lines = text.splitlines() - matches = _grep_lines(lines, regex, DOCUMENT_READ_MAX_MATCHES) + matches = _grep_lines(lines, regex, DOCUMENT_READ_MAX_MATCHES, clock) payload = { **base, "mode": "pattern", @@ -560,6 +694,14 @@ def _text_pattern(text: str, pattern: str, base: Dict[str, Any], unit: str) -> D "truncated": len(matches) >= DOCUMENT_READ_MAX_MATCHES, "matches": matches, } + if note: + payload["pattern_note"] = note + if clock.expired: + payload["timed_out"] = True + payload["hint"] = ( + "The search ran out of time; these are the matches found so far. " + "Use a simpler pattern, or read the document directly." + ) return DocumentReadResult(mode="pattern", payload=payload, format=base.get("format")) diff --git a/backend/src/apis/shared/files/document_tokens.py b/backend/src/apis/shared/files/document_tokens.py new file mode 100644 index 00000000..ea5f12b4 --- /dev/null +++ b/backend/src/apis/shared/files/document_tokens.py @@ -0,0 +1,98 @@ +"""Token weight of an inline attachment — PDFs are dual-encoded, bytes/4 is not. + +``docs/specs/document-context-offload.md`` §6.1 reports a document's weight as +``documentTokens`` and §4C decides whether an eviction is worth its prefix +re-write from the same number. Both used the compaction estimator's ``bytes/4``, +which is **wrong for PDFs by an order of magnitude**. + +Why: Bedrock understands a PDF page as an image *and* an extracted text layer +(the dual encoding §3 of the spec relies on, and the reason the design returns +native blocks rather than flattened text). The image channel dominates, and it +has nothing to do with the file's byte size — a 27 KB 60-page PDF and a 27 MB +60-page scan cost about the same. ``bytes/4`` sees only the bytes. + +Measured on dev 2026-09-16 (session ``61de2256``): a 27,578-byte, 60-page PDF +produced a **109.1K-token** first write against a 14.6K static prefix, so the +document was ≈94K tokens. ``bytes/4`` reported **6,894** — low by ~14×. At that +error the spec's "document share of the prefix" reads ~6% where the truth was +~86%, which inverts the ship/abandon call in the evaluation spec §4.2. + +The model here: **pages x ``PDF_PAGE_TOKEN_ESTIMATE``**, floored by ``bytes/4`` +so a byte-heavy PDF is never scored below the old number. The default (1,500) is +``compaction_policy.IMAGE_TOKEN_ESTIMATE``'s own rationale — Anthropic image +tokens are ~(w*h)/750, and a letter page at ~1000x1100 lands there. Against the +measurement above it is 1,500 vs 1,573 actual per page: ~5% low, deliberately, +because this number gates an eviction and over-stating it would evict documents +whose re-write has not earned it. + +Still a heuristic — Bedrock reports no per-block usage — but one whose dominant +term is now the right term. Non-PDF formats keep ``bytes/4``: Bedrock extracts +them to text, so there is no image channel to miss. + +Never raises: an unparseable PDF falls back to ``bytes/4``, which is exactly +today's behaviour. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Optional + +logger = logging.getLogger(__name__) + +#: Same heuristic as ``compaction_policy.CHARS_PER_TOKEN``; duplicated because +#: ``apis.shared`` must not import the agent layer. +CHARS_PER_TOKEN = 4 + +#: Estimated tokens for one rendered PDF page (the image channel). Env-tunable +#: so the number can be re-fit from measured rows without a deploy. +PDF_PAGE_TOKEN_ESTIMATE = int(os.environ.get("PDF_PAGE_TOKEN_ESTIMATE", 1_500)) + + +def pdf_page_count(raw: bytes) -> Optional[int]: + """Page count of ``raw``, or ``None`` if it will not open. + + Costs ~0.1 ms even for a 200-page file (measured), so callers on the + per-turn path do not need to memoize it. + """ + if not isinstance(raw, (bytes, bytearray)) or not raw: + return None + try: + import io + + import pypdfium2 as pdfium + + pdf = pdfium.PdfDocument(io.BytesIO(bytes(raw))) + try: + return len(pdf) + finally: + pdf.close() + except Exception: # noqa: BLE001 - estimator must never raise + logger.debug("pdf_page_count: unparseable PDF, falling back to bytes/4", exc_info=True) + return None + + +def estimate_document_tokens(fmt: Any, raw: Any) -> int: + """Estimated tokens an inline document block costs in the prefix. + + PDFs: ``max(pages * PDF_PAGE_TOKEN_ESTIMATE, bytes/4)``. Everything else: + ``bytes/4``. + """ + if not isinstance(raw, (bytes, bytearray)): + return 0 + byte_tokens = len(raw) // CHARS_PER_TOKEN + if str(fmt or "").lower() != "pdf": + return byte_tokens + pages = pdf_page_count(raw) + if not pages: + return byte_tokens + return max(pages * PDF_PAGE_TOKEN_ESTIMATE, byte_tokens) + + +__all__ = [ + "CHARS_PER_TOKEN", + "PDF_PAGE_TOKEN_ESTIMATE", + "estimate_document_tokens", + "pdf_page_count", +] diff --git a/backend/src/apis/shared/files/models.py b/backend/src/apis/shared/files/models.py index 3930b1ea..81b8b491 100644 --- a/backend/src/apis/shared/files/models.py +++ b/backend/src/apis/shared/files/models.py @@ -177,8 +177,18 @@ def is_presentation_file(filename: str, mime_type: str) -> bool: # attachment turns exceed it, several with only 3–4 files, so the per-file # cap above and the SPA's 5-file cap do not protect on their own. # ``0`` (or any non-positive value) disables the aggregate budget. +# +# Why 7.0 MB and not the 7.5 MB the arithmetic above suggests: base64 of N raw +# bytes is ``4*ceil(N/3)``, so 7,500,000 encodes to **exactly 10,000,000** — the +# quota itself, with nothing left for the event's JSON envelope (role, content +# keys, the prompt text block, per-file metadata, the wrapper). A guard whose +# default sits precisely on the break point it exists to stay under does not +# prevent the failure it was written for. 7,000,000 encodes to 9,333,336 and +# leaves ~666 KB of headroom, which comfortably covers the envelope while still +# admitting every attachment turn measured in prod (p90 cluster 2.58 MB, largest +# legitimate 29.89 MB — already over either number and correctly trimmed). INLINE_ATTACHMENTS_MAX_TOTAL_BYTES = int( - os.environ.get("INLINE_ATTACHMENTS_MAX_TOTAL_BYTES", 7_500_000) # 7.5MB + os.environ.get("INLINE_ATTACHMENTS_MAX_TOTAL_BYTES", 7_000_000) # 7.0MB ) # Files per message. The SPA enforces the same number client-side diff --git a/backend/tests/agents/main_agent/session/test_context_attribution_hook.py b/backend/tests/agents/main_agent/session/test_context_attribution_hook.py index 6d18321e..03123e24 100644 --- a/backend/tests/agents/main_agent/session/test_context_attribution_hook.py +++ b/backend/tests/agents/main_agent/session/test_context_attribution_hook.py @@ -164,3 +164,85 @@ async def test_count_failure_is_swallowed_and_yields_no_breakdown(self): def test_get_context_breakdown_is_none_when_absent(self): agent = FakeAgent(FakeModel(), messages=[]) assert get_context_breakdown(agent) is None + + +class TestInlineAttachmentGuard: + """``toolTokens`` is a residual between two independently sourced counts + (``full`` from Strands' projection, ``no_tools`` from our CountTokens call), + so any disagreement about how a content block is counted lands wholly in it. + + Measured on dev 2026-09-16 (session ``61de2256``): a call reported + ``toolTokens`` of 106,756 where the session's real tools prefix was 12,516 — + a 94,240 difference against a document measured at ~94,485, i.e. the entire + document attributed to tools. The split is therefore not computed while + inline bytes are in context. + """ + + def _doc_message(self): + return { + "role": "user", + "content": [ + {"text": "what does it say?"}, + {"document": {"format": "pdf", "name": "d_pdf", "source": {"bytes": b"%PDF-1.4"}}}, + ], + } + + @pytest.mark.asyncio + async def test_no_split_is_computed_while_a_document_is_inline(self): + model = FakeModel() + agent = FakeAgent(model, messages=[self._doc_message()]) + await ContextAttributionHook()._on_before_model_call(_event(agent, projected=100_000)) + + assert get_context_breakdown(agent) is None, "a contaminated split must not be published" + assert model.calls == [], "and it must not pay for CountTokens to compute one" + + @pytest.mark.asyncio + async def test_an_image_counts_too(self): + model = FakeModel() + agent = FakeAgent(model, messages=[{ + "role": "user", + "content": [{"image": {"format": "png", "source": {"bytes": b"\x89PNG"}}}], + }]) + await ContextAttributionHook()._on_before_model_call(_event(agent, projected=50_000)) + assert get_context_breakdown(agent) is None + + @pytest.mark.asyncio + async def test_a_digest_is_not_an_attachment_so_the_split_is_taken(self): + """The offload turns the document into text, which counts normally — + so an attachment session still gets ``prefixTokens`` from turn 2.""" + model = FakeModel(system=100, per_msg=10, tool_overhead=500) + agent = FakeAgent(model, messages=[{ + "role": "user", + "content": [{"text": ''}], + }]) + await ContextAttributionHook()._on_before_model_call(_event(agent, projected=650)) + assert _parts(get_context_breakdown(agent)) == {"system": 100, "tools": 540, "messages": 10} + + @pytest.mark.asyncio + async def test_a_later_clean_turn_computes_the_split(self): + """Deferred, not abandoned: the same agent takes the split once the + attachment has left the live context.""" + model = FakeModel(system=100, per_msg=10, tool_overhead=500) + agent = FakeAgent(model, messages=[self._doc_message()]) + hook = ContextAttributionHook() + await hook._on_before_model_call(_event(agent, projected=100_000)) + assert get_context_breakdown(agent) is None + + agent.messages = [{"role": "user", "content": [{"text": "follow-up"}]}] + await hook._on_before_model_call(_event(agent, projected=650)) + assert _parts(get_context_breakdown(agent)) == {"system": 100, "tools": 540, "messages": 10} + + @pytest.mark.asyncio + async def test_a_cached_split_is_still_used_when_a_document_arrives_later(self): + """The guard only defers the *computation*. An agent that already has a + trustworthy split keeps reporting against it.""" + model = FakeModel(system=100, per_msg=10, tool_overhead=500) + agent = FakeAgent(model, messages=[{"role": "user", "content": [{"text": "hi"}]}]) + hook = ContextAttributionHook() + await hook._on_before_model_call(_event(agent, projected=650)) + + agent.messages = [{"role": "user", "content": [{"text": "hi"}]}, self._doc_message()] + await hook._on_before_model_call(_event(agent, projected=95_000)) + parts = _parts(get_context_breakdown(agent)) + assert parts["system"] == 100 and parts["tools"] == 540 + assert parts["messages"] == 95_000 - 640, "the document lands in messages, where it belongs" diff --git a/backend/tests/agents/main_agent/session/test_document_read_flag_coupling.py b/backend/tests/agents/main_agent/session/test_document_read_flag_coupling.py new file mode 100644 index 00000000..dcf347c2 --- /dev/null +++ b/backend/tests/agents/main_agent/session/test_document_read_flag_coupling.py @@ -0,0 +1,126 @@ +"""DOCUMENT_READ_ENABLED must reach every path that promises the tool. + +The gap this guards: the flag was read only at the tool-injection gate +(``_document_tools_gate``). Neither the restore path nor the live offload +consulted it, so pulling the one kill switch an operator would actually reach +for left restore emitting ```` handles for a tool +that would not be injected — and left the live path still evicting bytes, which +is strictly worse than the pre-offload world, where live bytes never left. + +Rules pinned here: + +* restore still renders a digest (it has to drop the bytes either way, and a + digest beats the placeholder) but **without** the ``upload_id`` handle; +* the live offload **stops entirely** — bytes stay inline; +* with the flag on, everything is byte-identical to before. +""" + +from __future__ import annotations + +from typing import Any, Dict + +import pytest + +from agents.main_agent.session import document_offload as do +from agents.main_agent.session import document_rehydration as rh +from apis.shared.files.document_digest import DigestSection, DocumentDigest, render_digest +from apis.shared.files.models import FileMetadata, FileStatus + +RAW = b"P" * 40_000 + + +def _meta() -> FileMetadata: + digest = DocumentDigest( + status="ready", format="pdf", unit="page", count=10, + abstract="An abstract.", sections=[DigestSection(start=1, title="Intro")], + ) + return FileMetadata( + upload_id="u1", user_id="user-1", filename="Policy.pdf", mime_type="application/pdf", + size_bytes=len(RAW), status=FileStatus.READY, s3_key="k/u1", s3_bucket="b", + session_id="s1", digest=digest.to_item(), + ) + + +class _Repo: + def list_session_files_sync(self, session_id, status=None): + return [_meta()] + + def update_file_digest_sync(self, *a, **k): + pass + + +@pytest.fixture(autouse=True) +def repo(monkeypatch): + monkeypatch.setattr("apis.shared.files.repository.get_file_upload_repository", lambda: _Repo()) + + +def _doc_block() -> Dict[str, Any]: + return {"document": {"format": "pdf", "name": "Policy_pdf", "source": {"bytes": RAW}}} + + +def _conversation() -> list: + """Attach on turn 1, three turns elapsed — unpinned and offloadable.""" + return [ + {"role": "user", "content": [{"text": "here"}, _doc_block()]}, + {"role": "assistant", "content": [{"text": "a"}]}, + {"role": "user", "content": [{"text": "q2"}]}, + {"role": "assistant", "content": [{"text": "b"}]}, + {"role": "user", "content": [{"text": "q3"}]}, + {"role": "assistant", "content": [{"text": "c"}]}, + ] + + +class TestRenderHandle: + def test_default_is_unchanged(self): + d = DocumentDigest.from_item(_meta().digest) + assert 'upload_id="u1"' in render_digest(d, filename="Policy.pdf", upload_id="u1") + + def test_handle_can_be_omitted_without_losing_content(self): + d = DocumentDigest.from_item(_meta().digest) + text = render_digest(d, filename="Policy.pdf", upload_id="u1", include_handle=False) + assert "upload_id" not in text + # ...but the content the model actually reasons from survives. + assert "An abstract." in text and "Intro" in text and 'pages="10"' in text + + +class TestRestorePath: + def test_tool_on_renders_the_handle(self, monkeypatch): + monkeypatch.delenv("DOCUMENT_READ_ENABLED", raising=False) + res = rh.rehydrate_documents([{"role": "user", "content": [_doc_block()]}], + session_id="s1", user_id="user-1") + assert 'upload_id="u1"' in res.messages[0]["content"][0]["text"] + + def test_tool_off_keeps_the_digest_but_drops_the_handle(self, monkeypatch): + monkeypatch.setenv("DOCUMENT_READ_ENABLED", "false") + res = rh.rehydrate_documents([{"role": "user", "content": [_doc_block()]}], + session_id="s1", user_id="user-1") + text = res.messages[0]["content"][0]["text"] + assert res.rehydrated == 1, "a digest still beats the placeholder" + assert "upload_id" not in text, "must not advertise a tool that is not injected" + assert "An abstract." in text + + +class TestLiveOffloadPath: + def test_tool_on_offloads(self, monkeypatch): + monkeypatch.delenv("DOCUMENT_READ_ENABLED", raising=False) + assert do.offload_enabled_for("s1") is True + msgs = _conversation() + cands = do.candidate_documents(msgs, do.pinned_document_names(msgs, None)) + assert do.offload_documents(msgs, cands, session_id="s1", user_id="user-1").offloaded == 1 + + def test_tool_off_is_disabled_for_every_session(self, monkeypatch): + monkeypatch.setenv("DOCUMENT_READ_ENABLED", "false") + assert do.offload_enabled_for("s1") is False + # ...even for a session the rollout bucket would otherwise treat. + monkeypatch.setenv("DOCUMENT_OFFLOAD_ROLLOUT_PERCENT", "100") + assert all(do.offload_enabled_for(f"session-{i}") is False for i in range(50)) + + def test_tool_off_leaves_the_bytes_inline(self, monkeypatch): + """The regression in one assertion: the live path must never evict a + document whose only recovery path is switched off.""" + monkeypatch.setenv("DOCUMENT_READ_ENABLED", "false") + msgs = _conversation() + if do.offload_enabled_for("s1"): # pragma: no cover - the gate above + cands = do.candidate_documents(msgs, do.pinned_document_names(msgs, None)) + do.offload_documents(msgs, cands, session_id="s1", user_id="user-1") + assert msgs[0]["content"][1]["document"]["source"]["bytes"] == RAW diff --git a/backend/tests/apis/inference_api/test_attachment_turn_guard.py b/backend/tests/apis/inference_api/test_attachment_turn_guard.py index 76faf6de..08c8bca4 100644 --- a/backend/tests/apis/inference_api/test_attachment_turn_guard.py +++ b/backend/tests/apis/inference_api/test_attachment_turn_guard.py @@ -187,10 +187,14 @@ def test_partition_then_budget_chain_on_files_that_each_pass_the_per_file_gate(s assert sum(len(f.bytes) for f in kept) <= EVENT_QUOTA_BYTES def test_default_cap_is_the_documented_derivation(self): - # 10 MB event quota × 3/4 (base64) = 7.5 MB raw. Encoding exactly the - # cap lands on the quota, not over it. - assert file_models.INLINE_ATTACHMENTS_MAX_TOTAL_BYTES == 7_500_000 - assert len(_b64_of_size(7_500_000)) == EVENT_QUOTA_BYTES + # 10 MB event quota × 3/4 (base64) = 7.5 MB raw is the *break point*, + # not a safe cap: encoding exactly that lands ON the quota, and the + # event's JSON envelope (role, content keys, the prompt text block, + # per-file metadata, the wrapper) is then added on top. The default + # therefore sits below it, with room for the envelope. + assert len(_b64_of_size(7_500_000)) == EVENT_QUOTA_BYTES, "the break point is real" + assert file_models.INLINE_ATTACHMENTS_MAX_TOTAL_BYTES == 7_000_000 + assert len(_b64_of_size(file_models.INLINE_ATTACHMENTS_MAX_TOTAL_BYTES)) < EVENT_QUOTA_BYTES class TestMessageFileCap: @@ -380,3 +384,43 @@ def test_matches_real_decoded_length(self): for n in (0, 1, 2, 3, 4, 100, 7_500_000): f = _Attachment("x", size=n) assert _estimate_decoded_size(f) == n + + +class TestBudgetHeadroom: + """The guard exists so a turn's inline attachments never exceed AgentCore's + 10 MB *event* quota — past which ``create_message`` raises + ``SessionException`` and leaves a hole in history. + + Its first default did not achieve that. base64 of N raw bytes is + ``4*ceil(N/3)``, so 7,500,000 encoded to **exactly 10,000,000** — the quota + itself, with nothing left for the event's JSON envelope. A guard whose + default sits precisely on the break point it exists to stay under does not + prevent the failure it was written for. + """ + + QUOTA_BYTES = 10_000_000 + + def _encoded(self, raw_bytes: int) -> int: + return 4 * math.ceil(raw_bytes / 3) + + def test_the_default_leaves_room_for_the_event_envelope(self): + from apis.shared.files.models import INLINE_ATTACHMENTS_MAX_TOTAL_BYTES as cap + + encoded = self._encoded(cap) + headroom = self.QUOTA_BYTES - encoded + assert headroom > 0, f"budget encodes to {encoded:,} against a {self.QUOTA_BYTES:,} quota" + # Enough for the role, content keys, prompt text and per-file metadata + # many times over, without being so conservative it trims real turns + # (prod p90 attachment cluster is 2.58 MB). + assert headroom >= 250_000, f"only {headroom:,} bytes of headroom" + + def test_the_encoding_math_matches_base64(self): + """Pin the 4/3 inflation the budget is derived from, against the real + encoder rather than the arithmetic in a comment.""" + for raw in (3, 3_000, 7_000_000): + assert len(base64.b64encode(b"\0" * raw)) == self._encoded(raw) + + def test_a_turn_at_the_budget_still_fits_the_quota(self): + from apis.shared.files.models import INLINE_ATTACHMENTS_MAX_TOTAL_BYTES as cap + + assert len(base64.b64encode(b"\0" * cap)) < self.QUOTA_BYTES diff --git a/backend/tests/shared/test_document_digest.py b/backend/tests/shared/test_document_digest.py index 4299e8f3..b0f9a0b8 100644 --- a/backend/tests/shared/test_document_digest.py +++ b/backend/tests/shared/test_document_digest.py @@ -194,3 +194,54 @@ def test_kill_switch(self, monkeypatch): assert dd.document_digest_enabled() is False monkeypatch.setenv("DOCUMENT_DIGEST_ENABLED", "") assert dd.document_digest_enabled() is True + + + +class TestRenderBudgetIsHard: + """``render_digest`` escaped the abstract *after* slicing it to the room + that was left, so ``&`` -> ``&`` could render up to 5x longer than the + slice it was measured on. Measured: an all-``&`` abstract rendered 388 + tokens against a 100-token budget. The cap is the contract (spec PR-2 + decision #15 — ``digest.tokens`` is a stored fact per file), so it is now + measured on the rendered text, and the abstract is dropped rather than + allowed to overshoot. + """ + + def _digest(self, abstract: str) -> dd.DocumentDigest: + return dd.DocumentDigest( + status="ready", format="pdf", unit="page", count=10, + abstract=abstract, sections=[dd.DigestSection(start=1, title="Intro")], + ) + + @pytest.mark.parametrize( + "abstract", + [ + "&" * 4000, # worst case: every char expands 5x + " & " * 400, + "R&D spending rose. " * 300, + "Ordinary prose with no entities at all. " * 200, + "& already-escaped-looking text " * 200, + ], + ids=["all-amps", "angle-and-amp", "r-and-d", "plain-prose", "looks-escaped"], + ) + @pytest.mark.parametrize("budget", [40, 100, 400, 1_500]) + def test_never_exceeds_the_budget(self, abstract, budget): + text = dd.render_digest(self._digest(abstract), filename="a.pdf", upload_id="u", budget_tokens=budget) + assert dd.estimate_tokens(text) <= budget, dd.estimate_tokens(text) + + def test_entities_are_never_split(self): + text = dd.render_digest(self._digest("R&D " * 500), filename="a.pdf", upload_id="u", budget_tokens=60) + # A truncated "&" would leave a bare "&" or a fragment like "&am". + assert text.count("&") == text.count("&") + + def test_the_handle_always_survives(self): + """Only the opening tag may outlive a budget this small — it is what + makes the document retrievable at all.""" + text = dd.render_digest(self._digest("&" * 4000), filename="a.pdf", upload_id="u-keep", budget_tokens=1) + assert 'upload_id="u-keep"' in text + + def test_a_normal_digest_is_unchanged(self): + digest = self._digest("A short, ordinary abstract of the policy.") + text = dd.render_digest(digest, filename="a.pdf", upload_id="u", budget_tokens=1_500) + assert "A short, ordinary abstract of the policy." in text + assert "Intro" in text diff --git a/backend/tests/shared/test_document_read_pattern_safety.py b/backend/tests/shared/test_document_read_pattern_safety.py new file mode 100644 index 00000000..f983c0b1 --- /dev/null +++ b/backend/tests/shared/test_document_read_pattern_safety.py @@ -0,0 +1,153 @@ +"""ReDoS safety for ``document_read`` pattern mode. + +The defect: ``_compile`` passed a model-supplied regex straight to +``re.compile``. A nested-quantifier pattern then runs in exponential time — +measured on the real engine, ``(a+)+$`` against ``"a"*n + "!"`` takes 0.87 s at +n=24, 3.5 s at 26, **14 s at 28**. Extracted PDF lines run 60–100 characters, so +such a pattern never returns; the search happens in ``asyncio.to_thread``, which +CPython gives no way to cancel, so the turn hangs to the 600 s SSE timeout. + +Two defences, tested here: + +* ``catastrophic_pattern`` refuses the nested-quantifier family up front and the + pattern is searched literally instead — the same graceful degradation an + invalid regex already got — with a note carried to the model so a different + result is never presented as the requested one. +* ``_Budget`` bounds the walk across lines and pages, so a *merely slow* pattern + the detector does not flag degrades to a partial answer instead of hanging. +""" + +from __future__ import annotations + +import time + +import pytest + +from apis.shared.files import document_read as dr + +from .test_document_read import build_pdf + +# Shapes that make the backtracking engine go exponential. +CATASTROPHIC = [ + "(a+)+$", + r"(\w+\s*)+", + r"(\s*\w+)*", + "(a*)*", + "(a+|b){2,}", + r"([a-z]+)+@", + "(x+)+y", +] + +# Patterns a model would plausibly write against a real document. None may be +# refused: a false positive silently turns a regex search into a literal one. +LEGITIMATE = [ + r"\d+", + r"\w+\s+total", + "(invoice|receipt)", + "(foo)+", + r"section \d+", + r"(?:\d{4})-(?:\d{2})", + "retention", + r"[A-Z]{2,}\s", + r"\$[\d,]+\.\d{2}", + "(a|b)+", + r"a{2,5}", + r"\(\w+\)+", # escaped parens are literals, not a group + r"[(]a+[)]+", # ...and so are parens inside a character class +] + + +class TestCatastrophicPattern: + @pytest.mark.parametrize("pattern", CATASTROPHIC) + def test_flags_nested_quantifiers(self, pattern): + assert dr.catastrophic_pattern(pattern) is True + + @pytest.mark.parametrize("pattern", LEGITIMATE) + def test_never_flags_a_legitimate_pattern(self, pattern): + assert dr.catastrophic_pattern(pattern) is False + + @pytest.mark.parametrize("pattern", ["", "(", ")", "(((", "a{", "[a+)+", "\\"]) + def test_malformed_input_never_raises(self, pattern): + assert dr.catastrophic_pattern(pattern) in (True, False) + + +class TestCompileFallback: + def test_catastrophic_pattern_is_searched_literally_with_a_note(self): + regex, note = dr._compile("(a+)+$") + assert note is not None and "exponential" in note + # Literal: it matches its own text, and not an exponential input. + assert regex.search("see (a+)+$ here") is not None + assert regex.search("aaaaaaaa!") is None + + def test_the_measured_attack_is_now_instant(self): + """n=28 took 14 s before this change; n=80 must be immediate now.""" + regex, _ = dr._compile("(a+)+$") + t0 = time.monotonic() + regex.search("a" * 80 + "!") + assert time.monotonic() - t0 < 0.05 + + def test_invalid_regex_still_degrades_to_literal(self): + regex, note = dr._compile("([") + assert note is not None and "not a valid regular expression" in note + assert regex.search("a ([ bracket") is not None + + def test_a_good_pattern_compiles_as_a_regex_with_no_note(self): + regex, note = dr._compile(r"section \d+") + assert note is None + assert regex.search("SECTION 44 here") is not None # IGNORECASE preserved + + +class TestBudget: + def test_disabled_budget_never_expires(self): + clock = dr._Budget(0) + assert clock.out_of_time() is False + + def test_expires_and_latches(self): + clock = dr._Budget(0.01) + time.sleep(0.02) + assert clock.out_of_time() is True + assert clock.expired is True + + def test_grep_stops_on_an_exhausted_budget(self): + clock = dr._Budget(0.01) + time.sleep(0.02) + import re + out = dr._grep_lines(["match"] * 100, re.compile("match"), 40, clock) + assert out == [] + + def test_grep_without_a_budget_is_unchanged(self): + import re + out = dr._grep_lines(["alpha", "beta", "alpha"], re.compile("alpha"), 40) + assert [m["line"] for m in out] == [1, 3] + + +class TestEndToEnd: + def test_pdf_pattern_carries_the_note_and_still_answers(self): + raw = build_pdf(["the (a+)+$ literal lives here", "nothing on this page"]) + res = dr._pdf_pattern(raw, "(a+)+$", {"filename": "d.pdf", "format": "pdf"}) + assert res.payload["pattern_note"] + assert res.payload["match_count"] == 1 + assert res.payload["pages_matched"] == [1] + + def test_pdf_pattern_over_a_long_document_finishes_promptly(self): + """The regression in one assertion: a catastrophic pattern against a + 60-page document used to never return.""" + raw = build_pdf([f"SECTION {i} " + "a" * 70 for i in range(1, 61)]) + t0 = time.monotonic() + res = dr._pdf_pattern(raw, r"(\w+\s*)+$", {"filename": "d.pdf", "format": "pdf"}) + assert time.monotonic() - t0 < 5.0 + assert res.payload["pattern_note"] + + def test_text_pattern_reports_a_timeout_with_partial_matches(self, monkeypatch): + monkeypatch.setattr(dr, "DOCUMENT_READ_PATTERN_BUDGET_SECONDS", 1e-9) + text = "\n".join(f"line {i} alpha" for i in range(500)) + res = dr._text_pattern(text, "alpha", {"filename": "d.txt", "format": "txt"}, "line") + assert res.payload["timed_out"] is True + assert "ran out of time" in res.payload["hint"] + + def test_a_normal_search_reports_no_timeout_and_no_note(self): + raw = build_pdf([f"SECTION {i} retention clause" for i in range(1, 21)]) + res = dr._pdf_pattern(raw, "retention", {"filename": "d.pdf", "format": "pdf"}) + assert "timed_out" not in res.payload + assert "pattern_note" not in res.payload + assert res.payload["match_count"] == 20 diff --git a/backend/tests/shared/test_document_tokens.py b/backend/tests/shared/test_document_tokens.py new file mode 100644 index 00000000..46a3563d --- /dev/null +++ b/backend/tests/shared/test_document_tokens.py @@ -0,0 +1,77 @@ +"""Document token estimation (apis/shared/files/document_tokens.py). + +The regression this guards: ``bytes/4`` scored a 27,578-byte, 60-page PDF at +6,894 tokens when the measured first write was ~94K (dev session ``61de2256``, +2026-09-16). Bedrock dual-encodes each PDF page as an image, so the dominant +term is page count, not file size. Fixtures are hand-built PDFs — never user +files — so the page-driven estimate is exercised for real. +""" + +from __future__ import annotations + +import io + +import pytest + +from apis.shared.files import document_tokens as dt + +from .test_document_read import build_pdf + + +class TestPdfPageCount: + def test_counts_pages(self): + assert dt.pdf_page_count(build_pdf([f"p{i}" for i in range(1, 13)])) == 12 + + @pytest.mark.parametrize("raw", [b"", b"not a pdf", b"%PDF-1.4\ngarbage", None, 12, "str"]) + def test_unparseable_is_none_never_raises(self, raw): + assert dt.pdf_page_count(raw) is None + + +class TestEstimateDocumentTokens: + def test_pdf_is_driven_by_pages_not_bytes(self): + """The defect, stated as a test: a byte-small PDF with many pages.""" + raw = build_pdf([f"SECTION {i}" for i in range(1, 61)]) + tokens = dt.estimate_document_tokens("pdf", raw) + assert tokens == 60 * dt.PDF_PAGE_TOKEN_ESTIMATE + # ...and that is far above what bytes/4 would have said. + assert tokens > 10 * (len(raw) // dt.CHARS_PER_TOKEN) + + def test_byte_heavy_pdf_is_never_scored_below_bytes_over_four(self): + """max(), not replace: a scanned PDF whose bytes dominate keeps the + old floor so the estimate can only move up.""" + raw = build_pdf(["one page"]) + b"%" + b"\0" * 400_000 + assert dt.estimate_document_tokens("pdf", raw) == len(raw) // dt.CHARS_PER_TOKEN + + @pytest.mark.parametrize("fmt", ["docx", "txt", "md", "html", "", None]) + def test_non_pdf_keeps_bytes_over_four(self, fmt): + """Bedrock extracts these to text — no image channel to miss.""" + raw = b"x" * 40_000 + assert dt.estimate_document_tokens(fmt, raw) == 10_000 + + def test_pdf_format_match_is_case_insensitive(self): + raw = build_pdf(["a", "b"]) + assert dt.estimate_document_tokens("PDF", raw) == dt.estimate_document_tokens("pdf", raw) + + def test_unparseable_pdf_falls_back_to_bytes(self): + raw = b"not a pdf at all" * 100 + assert dt.estimate_document_tokens("pdf", raw) == len(raw) // dt.CHARS_PER_TOKEN + + @pytest.mark.parametrize("raw", [None, "", 0, [], {}]) + def test_non_bytes_is_zero(self, raw): + assert dt.estimate_document_tokens("pdf", raw) == 0 + + def test_env_tunable(self, monkeypatch): + """The constant is re-fittable from measured rows without a deploy.""" + monkeypatch.setattr(dt, "PDF_PAGE_TOKEN_ESTIMATE", 2_000) + assert dt.estimate_document_tokens("pdf", build_pdf(["a", "b", "c"])) == 6_000 + + +class TestAgainstTheDevMeasurement: + def test_sixty_page_pdf_lands_near_the_measured_write(self): + """Dev session 61de2256: 60 pages measured ~94K tokens. The estimate + must be the same order of magnitude, and deliberately a little under + (it gates an eviction).""" + raw = build_pdf([f"SECTION {i} clause text" for i in range(1, 61)]) + tokens = dt.estimate_document_tokens("pdf", raw) + measured = 94_000 + assert 0.7 * measured <= tokens <= measured, tokens diff --git a/docs/kaizen/review-queue.md b/docs/kaizen/review-queue.md index cab4258f..18e20347 100644 --- a/docs/kaizen/review-queue.md +++ b/docs/kaizen/review-queue.md @@ -5,6 +5,42 @@ Items added by `kaizen-research`, consumed by `kaizen-review-prep`. ## Open +### [2026-09-18] Document offload — rollout percent: what it actually gates (less than it looks), and why 100% is the reasonable call +- **Source**: validation sweep of the merged epic (#1137–#1143, #1145) on develop @ 8911d85f, 2026-09-16/17 — static review, 41 probes, and a live dev clickthrough (session `61de2256`). `docs/specs/document-offload-evaluation.md` is still **`Status: Draft`, no harness**, and its §2 quality gate is written as a **veto** ("no cost result, however good, ships a confirmed quality regression"). Meanwhile `DOCUMENT_OFFLOAD_ROLLOUT_PERCENT` **defaults to 100**, which is that same spec's §4.2 **"Ship"** state. +- **Surface**: backend — `session/document_offload.py` (`rollout_percent`, `offload_enabled_for`, bucket = `crc32(session_id) % 100`); the eval harness per evaluation spec §2/§4.2; no infra wiring exists for any of the epic's flags, so this is a runtime env var on inference-api, not a CDK change. +- **Effort × Impact**: S × H — the decision is one env var; the cost of not making it is that the causal question can never be answered on prod data. +- **Subtracts**: no — it is the condition under which the epic's defaults may stay on in prod. +- **Status**: open, but **smaller than first written — corrected 2026-09-18 after reading the gate.** `offload_enabled_for` (kill switch + bucket) is consulted in exactly two places: the head-of-turn live offload (`turn_based_session_manager.py:1337`) and restore-path slice ageing (`:293`). **PR-3 — the digest-on-restore that is the epic's actual behaviour change — is NOT bucketed**; it rides `DOCUMENT_REHYDRATE_ENABLED`, a plain global flag. So do PR-1 (`DOCUMENT_READ_ENABLED`), PR-2 (`DOCUMENT_DIGEST_ENABLED`) and PR-6 (`ATTACHMENT_TURN_GUARD_ENABLED`). The percent therefore splits only the *least* consequential half of the epic, and it was never a control arm for PRs 1–3 — those can only ever be measured before/after a deploy, whatever the percent says. + - **Recommendation: ship at 100%.** The evaluation spec's B-vs-C arm measures PR-4's incremental effect, and readout item 1 predicts PR-4 fires ≈0 times (restore reaches the document first — `AGENT_CACHE_BYPASS`). Splitting the fleet to A/B a mechanism that may not run buys little; **counting `document_offload` events at 100% answers the same question more cheaply.** If the count turns out to be high rather than ~0, the events carry their own `documentTokens` and `cacheGapSeconds`, so PR-4's share is still attributable from the ledger without arms. + - **What 100% genuinely gives up**: clean causal separation of PR-4's cost effect from PR-3's, *if* PR-4 turns out to fire often. Accept it as a recorded decision rather than a default nobody chose. + - **The real rollback lever is not the percent.** For the change users would actually notice, it is `DOCUMENT_REHYDRATE_ENABLED=false` (back to the pre-epic contentless placeholder); `DOCUMENT_READ_ENABLED=false` now also stops the live offload (#1147), and `DOCUMENT_OFFLOAD_ENABLED=false` stops only PR-4. + - **Still required either way**: the §2 quality veto has not run. Note the epic's §5 rule that PRs 1–3 are a correctness fix (quality should go **up**) while only PR-4 is the cost trade, so A→B and B→C must be scored separately or a strip-fix win masks an offload regression. + +### [2026-09-18] Document offload — the prod readout: five numbers, what each decides, and the one that is probably zero +- **Source**: same sweep. Live dev run confirmed the full lifecycle (turn 1 **full** → turn 2 **digest-only** → turn 3 **retrieved** via `document_read(page_range="44")`, 170 ms → slice aged on schedule at `DOCUMENT_SLICE_MAX_TURNS=2`), and measured the prefix collapsing **109.1K → ~15K** on one session. +- **Surface**: backend — `C#` rows (`hasDocuments` / `documentCount` / `documentTokens` / `documentDigests` / `documentSlices` / `documentReads`), `compactionEvents[].kind` ∈ {`document_stripped`, `document_rehydrated`, `document_offload`}, `S#` rollups (`fullDocumentCalls`, `digestOnlyCalls`, `documentReadCalls`, `documentReadPages`), `F#` feedback rows (#1142), EMF in `AgentCoreStack/Compaction`, `GET /admin/costs/sessions/{id}/calls`. +- **Effort × Impact**: S × H — every field already ships; this is a query session, not a build. +- **Subtracts**: no. +- **Status**: open — run after the epic reaches prod and has a week of attachment volume (~50 attachment sessions/week at the audited rate, so plan 2–4 weeks for anything cost-shaped). + 1. **Is PR-4 doing anything at all?** Count `document_offload` events. **Predicted ≈ 0** — measured in dev: `document_read` is itself an `extra_tool`, so the agent is rebuilt every turn, restore runs every turn, and the *restore* path (PR-3) reaches the document before the live offload ever can. The admin anatomy page diagnoses this itself (`AGENT_CACHE_BYPASS`). If it is ~0, then **all** the benefit is PR-3 and `DOCUMENT_OFFLOAD_PIN_TURNS` is inert code — which changes what to keep, and per evaluation spec §4.2 "a win with zero offload events is a confound, not a result." + 2. **Did the recovery path reach users?** `document_stripped` should fall to the unmatched residue (direct base64 attachments, deleted files) and `document_rehydrated` should carry the rest. Re-upload rate (same filename ≥2× in a session) **14% → target <3%**; this is the one metric PRs 1–3 move on their own. + 3. **Is the digest the right size?** `documentReadCalls` per attachment session against the health band **≈0.3–3**. Near zero means the digest is answering unaided (suspicious — or the model is guessing); >3/turn means it is too thin. + 4. **Was the offload ever wrong?** `document_offload` events with `cacheGapSeconds` under the cache TTL: **target zero**; any non-`over_ceiling` / non-`prefix_changed` one is the paid-when-free guard breaking. + 5. **The outcome signal, for the first time.** Down-thumb rate split by turn class (**full** / **digest-only** / **retrieved**) from the `F#` rows joined on `(sessionId, messageId)`. This is what PR-7 was built for and what every previous cost change on this platform shipped without. Report with `n` per class; per response-feedback spec §9 it is a *comparison between arms*, never a quality score. + +### [2026-09-18] Document offload — recalibrate `PDF_PAGE_TOKEN_ESTIMATE` from prod rows (the six findings are all fixed) +- **Source**: same sweep. #1147 fixed two of six findings — `documentTokens` was counting bytes, not PDF pages (**13.7× low**: 6,894 reported against ~94K measured, because Bedrock dual-encodes each page as an image), and `DOCUMENT_READ_ENABLED` did not reach the restore or offload paths. +- **Surface**: backend — `apis/shared/files/document_tokens.py` (`PDF_PAGE_TOKEN_ESTIMATE`, env-tunable); `apis/shared/files/document_read.py` (`_compile`/`_grep_lines`); `apis/shared/files/models.py` (`INLINE_ATTACHMENTS_MAX_TOTAL_BYTES`); `apis/shared/files/document_digest.py` (`render_digest`); `apis/app_api/admin/costs/` (`prefixTokens`). +- **Effort × Impact**: S × M each. +- **Subtracts**: no. +- **Status**: **all six findings fixed** (`82e0587d`, `10e511e5`, `71b0fdfb`, `e196a00e` — all on PR #1147). What remains is the one thing that needs prod data: + - **✅ FIXED 2026-09-18 (`71b0fdfb`)** — ReDoS in `document_read` pattern mode. `catastrophic_pattern` refuses the nested-quantifier family up front (searched literally instead, with a `pattern_note` on the payload; 7/7 dangerous shapes flagged, 0/13 false positives), and `DOCUMENT_READ_PATTERN_BUDGET_SECONDS` (default 2.0) bounds the walk for shapes it does not catch. **Residual, by construction:** a single `re.search` still cannot be interrupted — CPython exposes no cancel and the scan runs in a worker thread where signals are unavailable — so the budget is a backstop, not a guarantee. The complete fix is the `regex` module's native `timeout=`; it is installed transitively but **not declared**, so adopting it is a dependency decision needing approval. Re-open only if `timed_out` payloads show up in prod. + - **Recalibrate the per-page constant.** The 1,500 default was fitted to a single dev PDF (1,500 vs 1,573 actual/page, deliberately ~5% low because it gates an eviction). Once prod has attach turns, fit it against `cacheWriteInputTokens` minus `prefixTokens.system + prefixTokens.tools` on rows where `documentsAttached = 1`. It is an env var — no deploy. + - **✅ FIXED 2026-09-18 (`e196a00e`)** — `INLINE_ATTACHMENTS_MAX_TOTAL_BYTES` default 7,500,000 → **7,000,000** (encodes to 9,333,336, ~666 KB of headroom under the 10 MB event quota). Still far above prod's p90 attachment cluster (2.58 MB), so no real turn is newly trimmed. Watch `AttachmentTurnOverQuota` in prod anyway — a `SessionException` hole in history would mean the quota is stricter than the decimal reading. + - **✅ FIXED 2026-09-18 (`e196a00e`)** — the digest cap is now hard: escape first, cut on an entity-safe boundary (`&` never left as `&am`), and drop the abstract entirely rather than overshoot. Only the opening tag — the `document_read` handle — may outlive a budget too small for anything else. + - **✅ DIAGNOSED + GUARDED 2026-09-18 (`e196a00e`)** — not an off-by-one. `toolTokens` is a **residual** between two independently sourced counts, `full` (Strands' projection) minus `no_tools` (our own CountTokens call), so any disagreement about how a content block is counted lands wholly in it. The arithmetic is conclusive: 106,756 − 12,516 = **94,240**, against a document measured at ~94,485 — the entire document attributed to tools, because Bedrock counts a PDF page as an image *and* a text layer and the two sources did not agree. The counting disagreement is Bedrock/Strands behaviour we do not control, so the split is simply **not computed while inline document or image bytes are in context** and is taken on a later clean turn; absent reads "not tracked", which a wrong number does not. A digest is text, so attachment sessions still get the field from turn 2. **Residual question for the prod readout:** whether `full` and `no_tools` disagree on *images* too — if `prefixTokens` goes missing on image-only sessions for more than the attach turn, that is the tell. + - **`compaction_policy._block_tokens` has the same PDF blind spot** #1147 fixed elsewhere, and was deliberately left alone: it rescales to measured history and drives cut thresholds, so changing it moves when compaction fires. Its own decision. + ### [2026-09-16] Compaction stack (#1125 → #1128 → #1129 → #1131 → #1132) — run the quality-veto eval BEFORE the new defaults reach prod - **Source**: Phil-initiated — `docs/specs/compaction-model-relative-thresholds.md` §5 (the gate the spec sets for itself) + `docs/specs/compaction-over-threshold-cache-spiral.md` §4.3 (the veto eval) + `docs/specs/agentcore-evaluations-spike-findings.md` (what the managed evaluation service already supplies). Stack built 2026-09-15/16; the 2026-09-15 replay of 20 heavy Sonnet 5 sessions priced 100k/25k at **$73.70 vs $102.64 actual** input-side. - **Surface**: backend — `session/compaction_policy.py` (floor-seeking cut), `session/compaction_summary.py` (8k summary via Nova Micro), `session/turn_based_session_manager.py` (`apply_pending_compaction`); eval harness per the spiral spec §4.3 (constraint retention / revision continuity / reference lookup), fixed-threshold arm as control. @@ -45,7 +81,7 @@ Items added by `kaizen-research`, consumed by `kaizen-review-prep`. - **Surface**: backend — `document_read` tool gated on the session having an attachment (ids kept out of `INJECTED_TOOL_IDS`), `DocumentDigest` at upload on `FileMetadata`, `_strip_document_bytes` → digest + live handle on restore; analytics: `hasDocuments` / `documentTokens` on `MessageMetadata` (the spec's PR-7) plus whatever else makes the cost-effectiveness of the next change *decidable* — per-session attachment tokens vs. answer quality signal, `document_read` call counts and page volumes, digest-vs-full-document turn shares. - **Effort × Impact**: L × H - **Subtracts**: partially — retires the contentless restore placeholder (a correctness bug: a returning user's document is silently lost) and the 14% re-upload rate. -- **Status**: open — task chip spawned 2026-09-16 ("Explore document offloading with actionable cost analytics"). Ask of the exploration: sequence PRs 1–3 together as the correctness fix, pull the analytics (PR-7) forward so the cost work (PR-4) is measured from day one, and align the per-call fields with the compaction ledger (#1130) and `AgentCoreStack/Compaction` so one anatomy page explains a session's cost. +- **Status**: **BUILT 2026-09-16** — shipped as PRs #1137–#1143 (+#1145); analytics were pulled forward into PR-1 as asked, and the per-call fields do align with the compaction ledger and `AgentCoreStack/Compaction`. Validated 2026-09-16/17 (41 probes + a live dev clickthrough); two defects found and fixed in #1147. What remains is **measurement, not build** — see the three [2026-09-18] entries at the top of this queue. Original ask, for the record: task chip spawned 2026-09-16 ("Explore document offloading with actionable cost analytics"). Ask of the exploration: sequence PRs 1–3 together as the correctness fix, pull the analytics (PR-7) forward so the cost work (PR-4) is measured from day one, and align the per-call fields with the compaction ledger (#1130) and `AgentCoreStack/Compaction` so one anatomy page explains a session's cost. ### [2026-09-11] ✅ ANSWERED: `bedrock_cache_points_supported()` STAYS — and it is wrong in the opposite direction from the one we suspected