fix(documents): PDF-aware documentTokens, and couple DOCUMENT_READ_ENABLED to the paths that promise the tool

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.

backend/src/agents/main_agent/session/document_context.py

@@ -14,17 +14,18 @@
 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.
 """
 
 from __future__ import annotations
 
 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
 

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
 
 

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

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=[],