From 5c214060fd2f5104509815e8434885ad5784e4cd Mon Sep 17 00:00:00 2001 From: Phil Merrell Date: Wed, 16 Sep 2026 16:41:55 -0600 Subject: [PATCH] feat(attachments): hold a turn's inline attachments to the Memory event quota (PR-6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stored unit is the message, not the file: a turn's inline attachments are persisted as one AgentCore Memory event, base64-inflated 4/3 into a blob payload and bounded by the 10 MB event quota, so ~7.5 MB of raw bytes per turn is the break point. Past it create_message re-raises SessionException — a hole in history. The per-file 4 MB gate and the SPA's 5-file cap do not protect: ~1.3–1.4% of prod attachment turns exceed it, several with only 3–4 files (document-context-offload-validation.md, Claim 7). - `_apply_inline_byte_budget`: first-fit in attachment order against INLINE_ATTACHMENTS_MAX_TOTAL_BYTES (default 7.5 MB, derivation in the constant's comment); images count. Trimmed files join the existing oversized-note path with their own wording ("send in a follow-up"). - `_apply_message_file_cap`: the server side of the SPA's MAX_FILES_PER_MESSAGE, applied to both `files` and `file_upload_ids` before the S3 fetch. The resolver's silent `[:5]` becomes a note to the user; the app_api's dead `max_files_per_message` now reads the same shared constant. - One content-free EMF metric, AttachmentTurnOverQuota (Bytes) in AgentCoreStack/Compaction, plus a count-only log line. - ATTACHMENT_TURN_GUARD_ENABLED=false restores the prior behaviour. Co-Authored-By: Claude Fable 5.1 --- backend/src/apis/app_api/files/service.py | 8 +- backend/src/apis/inference_api/chat/routes.py | 212 +++++++++- backend/src/apis/shared/feature_flags.py | 20 + .../src/apis/shared/files/file_resolver.py | 18 +- backend/src/apis/shared/files/models.py | 25 ++ .../test_attachment_turn_guard.py | 382 ++++++++++++++++++ backend/tests/shared/test_files.py | 15 + .../FILE_UPLOAD_IMPLEMENTATION.md | 5 +- 8 files changed, 674 insertions(+), 11 deletions(-) create mode 100644 backend/tests/apis/inference_api/test_attachment_turn_guard.py diff --git a/backend/src/apis/app_api/files/service.py b/backend/src/apis/app_api/files/service.py index 2c4fc027f..2d697ef10 100644 --- a/backend/src/apis/app_api/files/service.py +++ b/backend/src/apis/app_api/files/service.py @@ -36,6 +36,7 @@ is_allowed_mime_type, is_presentation_file, ALLOWED_MIME_TYPES, + MAX_FILES_PER_MESSAGE, ) from .sheet_preview import ( MAX_WORKBOOK_BYTES, @@ -182,9 +183,10 @@ def __init__( "FILE_UPLOAD_MAX_SIZE_BYTES_PRESENTATION", 25 * 1024 * 1024 # 25MB ) ) - self.max_files_per_message = max_files_per_message or int( - os.environ.get("FILE_UPLOAD_MAX_FILES_PER_MESSAGE", 5) - ) + # Single source of truth is the shared constant (the inference API + # enforces it per message; see ``_apply_message_file_cap``). Kept on + # the service so callers can read the effective limit. + self.max_files_per_message = max_files_per_message or MAX_FILES_PER_MESSAGE self.user_quota_bytes = user_quota_bytes or int( os.environ.get("FILE_UPLOAD_USER_QUOTA_BYTES", 1024 * 1024 * 1024) # 1GB ) diff --git a/backend/src/apis/inference_api/chat/routes.py b/backend/src/apis/inference_api/chat/routes.py index bf83762f8..f32a8e0cf 100644 --- a/backend/src/apis/inference_api/chat/routes.py +++ b/backend/src/apis/inference_api/chat/routes.py @@ -27,10 +27,15 @@ from apis.inference_api.runtime_health import ping_payload from apis.shared.feature_flags import ( agents_enabled, + attachment_turn_guard_enabled, mid_turn_steering_enabled, skills_enabled, ) from apis.shared.files.file_resolver import get_file_resolver +from apis.shared.files.models import ( + INLINE_ATTACHMENTS_MAX_TOTAL_BYTES, + MAX_FILES_PER_MESSAGE, +) from apis.shared.models.managed_models import list_managed_models from apis.shared.quota import ( QuotaExceededEvent, @@ -821,6 +826,110 @@ def _partition_attachments( return inline, tabular, presentations, oversized +def _apply_message_file_cap( + direct_files: list, + upload_ids: list, + max_files: int, +) -> tuple[list, list, list, int]: + """Hold a message to ``max_files`` attachments across both request paths. + + Returns ``(direct_files, upload_ids, dropped_names, dropped_total)``. + Direct ``files`` come first (they are already in the request body), then + ``file_upload_ids`` fill whatever budget remains. Attachment order is + kept, so the first N the user attached are the N that survive. + + The cap is applied to the upload IDs *before* they are resolved: the old + resolver default truncated silently after the fact, and letting every ID + through just to name the losers would fan out one S3 read per ID a client + chose to send. IDs beyond the budget are therefore counted, not named — + ``dropped_names`` holds the direct files (names known) and + ``dropped_total`` counts both. ``max_files <= 0`` disables the cap. + """ + if max_files <= 0: + return direct_files, upload_ids, [], 0 + + kept_direct = direct_files[:max_files] + dropped_names = [f.filename for f in direct_files[max_files:]] + id_budget = max(0, max_files - len(kept_direct)) + kept_ids = upload_ids[:id_budget] + dropped_total = len(dropped_names) + (len(upload_ids) - len(kept_ids)) + return kept_direct, kept_ids, dropped_names, dropped_total + + +def _apply_inline_byte_budget( + inline: list, + max_total_bytes: int, +) -> tuple[list, list, int]: + """Hold the inline set (documents *and* images) to one message's byte + budget. Returns ``(kept, over_budget, requested_bytes)``. + + Why this exists: the turn's inline attachments are persisted as one + AgentCore Memory event, and past ~7.5 MB of raw bytes that write fails + with ``SessionException`` — a hole in history, not a degraded turn. See + ``INLINE_ATTACHMENTS_MAX_TOTAL_BYTES`` for the derivation. + + Policy — first-fit in attachment order: walk the files as the user + attached them, keep each one that still fits, and move any that would + push the running total over the budget to ``over_budget``. Earlier + attachments win, and a later, smaller file that still fits rides along + rather than being punished for a large neighbour. Order within both + lists is the attachment order, so the marker text and the guidance note + are deterministic (they land in the cacheable prefix on later turns). + + Images count toward the budget: they are part of the same message and + the same event, even though the per-file document gate skips them. + ``max_total_bytes <= 0`` disables the budget. + """ + requested = sum(_estimate_decoded_size(f) for f in inline) + if max_total_bytes <= 0: + return list(inline), [], requested + + kept: list = [] + over: list = [] + running = 0 + for file in inline: + size = _estimate_decoded_size(file) + if running + size > max_total_bytes: + over.append(file) + continue + running += size + kept.append(file) + return kept, over, requested + + +def _emit_attachment_over_quota_metric( + requested_bytes: int, + cap_bytes: int, + inline_count: int, + dropped_count: int, +) -> None: + """One content-free EMF record in ``AgentCoreStack/Compaction`` when a + turn's inline attachments had to be trimmed to the byte budget. Never + raises. ``AttachmentTurnOverQuota`` carries the requested bytes so the + rate *and* the size distribution of over-quota turns are measurable + (spec §4E put the rate at ~1.3–1.4% of attachment turns from a proxy; + this is the direct count). + """ + try: + from apis.shared.observability.emf import emit_emf_metrics + from apis.shared.observability.prompt_cache import prompt_cache_observability_enabled + + if not prompt_cache_observability_enabled(): + return + emit_emf_metrics( + "AgentCoreStack/Compaction", + metrics={"AttachmentTurnOverQuota": requested_bytes}, + properties={ + "capBytes": cap_bytes, + "inlineFileCount": inline_count, + "droppedFileCount": dropped_count, + }, + units={"AttachmentTurnOverQuota": "Bytes"}, + ) + except Exception as e: # noqa: BLE001 + logger.debug("AttachmentTurnOverQuota EMF skipped: %s", e) + + def _attachment_marker_names(all_files: list, oversized_inline: list) -> list: """Filenames for the ``[Attached files: …]`` marker on the user message. @@ -852,10 +961,21 @@ def _build_attachment_guidance( diverted_presentations: list, oversized_inline: list, enabled_tools: list | None, + over_budget: list | None = None, + dropped_over_count_names: list[str] | None = None, + dropped_over_count_total: int = 0, + max_files: int = 0, ) -> str: """Return a short markdown addendum describing how attachments will be handled, to append to the user's message so the agent (and the user) both understand why a file isn't inline. + + ``oversized_inline`` is the per-file case (the file itself is too big; + the fix is a smaller file). ``over_budget`` is the aggregate case (each + file is fine, together they exceed one message's budget; the fix is a + follow-up message). They get separate sentences because the remedy + differs. ``dropped_over_count_*`` describe files beyond the per-message + count cap: names where known (direct ``files``), a count otherwise. """ parts: list[str] = [] @@ -908,6 +1028,32 @@ def _build_attachment_guidance( f"and use the Spreadsheet Analysis tool._" ) + if over_budget: + names = ", ".join(f"`{f.filename}`" for f in over_budget) + parts.append( + f"_Attached file(s) {names} were skipped because this message's " + f"attachments together exceed the combined size limit for a " + f"single message. Send them in a follow-up message._" + ) + + if dropped_over_count_total > 0: + limit = f"{max_files} file" + ("s" if max_files != 1 else "") + if dropped_over_count_names: + names = ", ".join(f"`{n}`" for n in dropped_over_count_names) + unnamed = dropped_over_count_total - len(dropped_over_count_names) + tail = f" and {unnamed} more" if unnamed > 0 else "" + parts.append( + f"_Only the first {limit} per message are attached; " + f"{names}{tail} were not. Send them in a follow-up message._" + ) + else: + noun = "file was" if dropped_over_count_total == 1 else "files were" + parts.append( + f"_Only the first {limit} per message are attached; " + f"{dropped_over_count_total} more {noun} not. " + f"Send them in a follow-up message._" + ) + return "\n\n".join(parts) @@ -1532,14 +1678,39 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g # budget; we skip them inline and surface a note instead of # letting Bedrock reject the turn. all_files = list(input_data.files) if input_data.files else [] + upload_ids_to_resolve = list(input_data.file_upload_ids or []) + + # Per-message file count (spec §4E / PR-6). Applied here, before the + # S3 fetch, so a sixth file is reported to the user instead of silently + # truncated by the resolver — and so a client cannot fan out unbounded + # S3 reads. With the guard off, the resolver's own backstop (5) applies + # exactly as it did before. + turn_guard_on = attachment_turn_guard_enabled() + dropped_over_count_names: list[str] = [] + dropped_over_count_total = 0 + if turn_guard_on: + ( + all_files, + upload_ids_to_resolve, + dropped_over_count_names, + dropped_over_count_total, + ) = _apply_message_file_cap(all_files, upload_ids_to_resolve, MAX_FILES_PER_MESSAGE) + if dropped_over_count_total: + logger.warning( + "Dropped %d attachment(s) over the %d-per-message cap", + dropped_over_count_total, + MAX_FILES_PER_MESSAGE, + ) - if input_data.file_upload_ids: + if upload_ids_to_resolve: try: file_resolver = get_file_resolver() resolved_files = await file_resolver.resolve_files( user_id=user_id, - upload_ids=input_data.file_upload_ids, - max_files=5, # Bedrock document limit + upload_ids=upload_ids_to_resolve, + # Already capped above when the guard is on; the resolver's + # own backstop is the pre-guard behaviour. + max_files=None if turn_guard_on else 5, ) for rf in resolved_files: all_files.append( @@ -1602,7 +1773,36 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g f"{[(f.filename, _estimate_decoded_size(f)) for f in oversized_inline]}" ) - attachment_marker_names = _attachment_marker_names(all_files, oversized_inline) + # Aggregate budget for the turn (spec §4E / PR-6): the inline set is one + # persisted message, and a message over ~7.5 MB raw fails the AgentCore + # Memory write with SessionException. Trim first-fit in attachment order; + # the trimmed files join the oversized note path, never the exception. + over_budget_inline: list = [] + if turn_guard_on and files_to_send: + files_to_send, over_budget_inline, requested_inline_bytes = _apply_inline_byte_budget( + files_to_send, INLINE_ATTACHMENTS_MAX_TOTAL_BYTES + ) + if over_budget_inline: + logger.warning( + "Attachment turn over quota: requested_bytes=%d cap_bytes=%d " + "inline_files=%d dropped_files=%d", + requested_inline_bytes, + INLINE_ATTACHMENTS_MAX_TOTAL_BYTES, + len(files_to_send) + len(over_budget_inline), + len(over_budget_inline), + ) + _emit_attachment_over_quota_metric( + requested_bytes=requested_inline_bytes, + cap_bytes=INLINE_ATTACHMENTS_MAX_TOTAL_BYTES, + inline_count=len(files_to_send) + len(over_budget_inline), + dropped_count=len(over_budget_inline), + ) + + # Both classes were dropped from the turn entirely; the marker must not + # promise a card for either. + attachment_marker_names = _attachment_marker_names( + all_files, oversized_inline + over_budget_inline + ) # Pre-create session metadata so OAuth interrupts and other state can # attach to the session row from turn one. Best-effort; on failure the @@ -2697,6 +2897,10 @@ def _session_title_sse() -> Optional[str]: diverted_presentations, oversized_inline, effective_enabled_tools, + over_budget=over_budget_inline, + dropped_over_count_names=dropped_over_count_names, + dropped_over_count_total=dropped_over_count_total, + max_files=MAX_FILES_PER_MESSAGE, ) # When multiple spreadsheets are visible, ship the full inventory # up front so the agent can disambiguate intentionally instead of diff --git a/backend/src/apis/shared/feature_flags.py b/backend/src/apis/shared/feature_flags.py index 31555cb95..35f017136 100644 --- a/backend/src/apis/shared/feature_flags.py +++ b/backend/src/apis/shared/feature_flags.py @@ -352,3 +352,23 @@ def ask_user_question_enabled() -> bool: every time it flipped. """ return os.environ.get("ASK_USER_QUESTION_ENABLED", "").strip().lower() != "false" + + +def attachment_turn_guard_enabled() -> bool: + """Whether a turn's attachments are held to the per-message file count and + the aggregate inline-bytes budget before the message is built. + + Covers ``_apply_message_file_cap`` and ``_apply_inline_byte_budget`` in + the inference API chat route (docs/specs/document-context-offload.md §4E, + PR-6). **Default ON with a kill switch** (house style): unset or empty + resolves to enabled; only the literal ``"false"`` (case-insensitive) + disables. + + While off the route behaves as before this shipped: the ``file_upload_ids`` + resolver silently truncates at five, direct ``files`` are uncounted, and + a turn whose attachments sum past the AgentCore Memory event quota fails + at ``create_message`` with a ``SessionException``. The tuning knobs + (``INLINE_ATTACHMENTS_MAX_TOTAL_BYTES``, ``FILE_UPLOAD_MAX_FILES_PER_MESSAGE``) + live in ``apis.shared.files.models``. + """ + return os.environ.get("ATTACHMENT_TURN_GUARD_ENABLED", "").strip().lower() != "false" diff --git a/backend/src/apis/shared/files/file_resolver.py b/backend/src/apis/shared/files/file_resolver.py index 6ce4a7b5f..7af143e95 100644 --- a/backend/src/apis/shared/files/file_resolver.py +++ b/backend/src/apis/shared/files/file_resolver.py @@ -53,7 +53,7 @@ async def resolve_files( self, user_id: str, upload_ids: List[str], - max_files: int = 5 + max_files: Optional[int] = 5 ) -> List[ResolvedFileContent]: """ Resolve upload IDs to file content objects. @@ -61,7 +61,11 @@ async def resolve_files( Args: user_id: Owner user ID (for authorization) upload_ids: List of upload IDs to resolve - max_files: Maximum files to process (Bedrock limit is 5) + max_files: Backstop on how many IDs are fetched. ``None`` fetches + every ID given. The chat route applies the per-message cap + *before* calling this so it can tell the user which files + were left out; this cap only exists so no caller can fan out + an unbounded number of S3 reads by accident. Returns: List of ResolvedFileContent objects with base64-encoded bytes @@ -71,7 +75,15 @@ async def resolve_files( """ resolved_files = [] - for upload_id in upload_ids[:max_files]: + if max_files is not None and len(upload_ids) > max_files: + logger.warning( + "resolve_files truncating %d upload ID(s) to %d", + len(upload_ids), + max_files, + ) + upload_ids = upload_ids[:max_files] + + for upload_id in upload_ids: try: file_content = await self._resolve_single_file(user_id, upload_id) if file_content: diff --git a/backend/src/apis/shared/files/models.py b/backend/src/apis/shared/files/models.py index 18cf79344..e92d43dd4 100644 --- a/backend/src/apis/shared/files/models.py +++ b/backend/src/apis/shared/files/models.py @@ -165,6 +165,31 @@ def is_presentation_file(filename: str, mime_type: str) -> bool: os.environ.get("INLINE_DOCUMENT_MAX_BYTES", 4 * 1024 * 1024) # 4MB ) +# A turn's inline attachments are persisted as ONE message, and the message — +# not the file — is what AgentCore Memory bounds. Anything over the SDK's +# ~72 KB conversational limit is written as a base64 ``blob`` payload, so raw +# attachment bytes inflate by 4/3 on the way in and are then held to the +# 10 MB event quota. 10 MB × 3/4 = 7.5 MB of raw bytes per turn is the break +# point. Above it ``create_message`` raises ``SessionException`` — a hole in +# history — which is strictly worse than the per-file oversized note, so the +# turn is trimmed to this budget *before* it is built. Prod measurement +# (docs/specs/document-context-offload-validation.md, Claim 7): ~1.3–1.4% of +# 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. +INLINE_ATTACHMENTS_MAX_TOTAL_BYTES = int( + os.environ.get("INLINE_ATTACHMENTS_MAX_TOTAL_BYTES", 7_500_000) # 7.5MB +) + +# Files per message. The SPA enforces the same number client-side +# (``MAX_FILES_PER_MESSAGE`` in file-upload.service.ts); this is the server +# side of it, shared by the ``file_upload_ids`` resolver and the direct +# ``files`` path so a sixth file is reported to the user instead of silently +# truncated. ``0`` (or any non-positive value) disables the count cap. +MAX_FILES_PER_MESSAGE = int( + os.environ.get("FILE_UPLOAD_MAX_FILES_PER_MESSAGE", 5) +) + # ============================================================================= # Database Models (stored in DynamoDB) diff --git a/backend/tests/apis/inference_api/test_attachment_turn_guard.py b/backend/tests/apis/inference_api/test_attachment_turn_guard.py new file mode 100644 index 000000000..76faf6de5 --- /dev/null +++ b/backend/tests/apis/inference_api/test_attachment_turn_guard.py @@ -0,0 +1,382 @@ +"""A turn's attachments are held to one message's budget — never to +``SessionException`` (docs/specs/document-context-offload.md §4E, PR-6). + +The stored unit is the MESSAGE, not the file. A turn's inline attachments are +persisted as one AgentCore Memory event; anything over the SDK's ~72 KB +conversational limit is written as a base64 ``blob`` (4/3 inflation) bounded by +the 10 MB event quota, so ~7.5 MB of raw attachment bytes per turn is the break +point. Past it, ``create_message`` re-raises ``SessionException`` — a hole in +history. The per-file 4 MB gate and the SPA's 5-file cap do not protect: prod +measurement (validation doc, Claim 7) found ~1.3–1.4% of attachment turns over +quota, several with only 3–4 files. + +Two helpers, composed in the chat route: + +- ``_apply_message_file_cap`` — the server side of the SPA's + ``MAX_FILES_PER_MESSAGE``, applied *before* the S3 fetch so a sixth file is + reported to the user instead of silently truncated by the resolver. +- ``_apply_inline_byte_budget`` — first-fit in attachment order against + ``INLINE_ATTACHMENTS_MAX_TOTAL_BYTES``; images count. +""" + +import base64 +import math + +import pytest + +from apis.inference_api.chat.routes import ( + _apply_inline_byte_budget, + _apply_message_file_cap, + _attachment_marker_names, + _build_attachment_guidance, + _emit_attachment_over_quota_metric, + _estimate_decoded_size, + _partition_attachments, +) +from apis.shared.feature_flags import attachment_turn_guard_enabled +from apis.shared.files import models as file_models + +XLSX_MIME = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" +EVENT_QUOTA_BYTES = 10_000_000 # AgentCore Memory event payload quota + + +def _b64_of_size(n: int) -> str: + """Base64 text whose decoded size is exactly ``n`` bytes.""" + return base64.b64encode(b"x" * n).decode() + + +class _Attachment: + """Minimal stand-in for FileContent — the helpers read these three.""" + + def __init__(self, filename: str, content_type: str = "application/pdf", size: int = 0): + self.filename = filename + self.content_type = content_type + self.bytes = _b64_of_size(size) + + +MB = 1024 * 1024 + + +class TestInlineByteBudget: + def test_under_the_budget_keeps_everything_in_order(self): + files = [_Attachment("a.pdf", size=3 * MB), _Attachment("b.pdf", size=3 * MB)] + kept, over, requested = _apply_inline_byte_budget(files, 7_500_000) + + assert kept == files + assert over == [] + assert requested == 6 * MB + + def test_over_the_budget_moves_the_file_that_breaks_it(self): + files = [_Attachment("a.pdf", size=4 * MB), _Attachment("b.pdf", size=4 * MB)] + kept, over, requested = _apply_inline_byte_budget(files, 7_500_000) + + assert [f.filename for f in kept] == ["a.pdf"] + assert [f.filename for f in over] == ["b.pdf"] + assert requested == 8 * MB + + def test_first_fit_earlier_files_win_and_a_later_small_file_still_rides(self): + # A(4) fits, B(4) would break the budget, C(1) still fits after A. + files = [ + _Attachment("a.pdf", size=4 * MB), + _Attachment("b.pdf", size=4 * MB), + _Attachment("c.pdf", size=1 * MB), + ] + kept, over, _ = _apply_inline_byte_budget(files, 7_500_000) + + assert [f.filename for f in kept] == ["a.pdf", "c.pdf"] + assert [f.filename for f in over] == ["b.pdf"] + + def test_order_is_deterministic_attachment_order_in_both_lists(self): + files = [ + _Attachment("z.pdf", size=3 * MB), + _Attachment("y.pdf", size=3 * MB), + _Attachment("x.pdf", size=3 * MB), + _Attachment("w.pdf", size=3 * MB), + ] + kept, over, _ = _apply_inline_byte_budget(files, 7_500_000) + + assert [f.filename for f in kept] == ["z.pdf", "y.pdf"] + assert [f.filename for f in over] == ["x.pdf", "w.pdf"] + + def test_images_count_toward_the_budget(self): + # The per-file document gate skips images; the message-level budget + # must not, because the image bytes are in the same persisted event. + files = [ + _Attachment("photo.png", content_type="image/png", size=4 * MB), + _Attachment("report.pdf", size=4 * MB), + ] + kept, over, _ = _apply_inline_byte_budget(files, 7_500_000) + + assert [f.filename for f in kept] == ["photo.png"] + assert [f.filename for f in over] == ["report.pdf"] + + def test_a_single_file_over_the_budget_is_moved_not_raised(self): + files = [_Attachment("huge.png", content_type="image/png", size=8 * MB)] + kept, over, _ = _apply_inline_byte_budget(files, 7_500_000) + + assert kept == [] + assert over == files + + def test_exactly_at_the_budget_is_allowed(self): + files = [_Attachment("a.pdf", size=7_500_000)] + kept, over, _ = _apply_inline_byte_budget(files, 7_500_000) + assert kept == files and over == [] + + def test_non_positive_budget_disables_the_cap(self): + files = [_Attachment("a.pdf", size=4 * MB), _Attachment("b.pdf", size=4 * MB)] + kept, over, requested = _apply_inline_byte_budget(files, 0) + + assert kept == files + assert over == [] + assert requested == 8 * MB + + def test_empty_input(self): + assert _apply_inline_byte_budget([], 7_500_000) == ([], [], 0) + + +class TestNothingReachesTheSessionExceptionPath: + """The property the whole PR exists for: whatever survives the budget is + writable as one AgentCore Memory event. + + The SDK serialises the message as base64 inside a ``blob`` payload. If the + kept set's encoded size is under the 10 MB event quota, ``create_message`` + cannot fail on size, so the turn degrades to the guidance note rather than + to ``SessionException``. + """ + + @pytest.mark.parametrize( + "sizes", + [ + [4 * MB, 4 * MB], # the classic two-file case + [4 * MB, 4 * MB, 4 * MB, 4 * MB, 4 * MB], # SPA-legal five files + [3 * MB, 3 * MB, 3 * MB], # 3 files, each fine + [2_600_000, 2_600_000, 2_600_000], # the "3–4 files" prod shape + [7_500_000, 1], + [1] * 5 + [4 * MB] * 5, + ], + ) + def test_kept_set_encodes_under_the_event_quota(self, sizes): + files = [_Attachment(f"f{i}.pdf", size=s) for i, s in enumerate(sizes)] + default_cap = file_models.INLINE_ATTACHMENTS_MAX_TOTAL_BYTES + kept, over, _ = _apply_inline_byte_budget(files, default_cap) + + encoded = sum(len(f.bytes) for f in kept) + assert encoded <= math.ceil(default_cap * 4 / 3) + 4 * len(kept) # padding slack + assert encoded <= EVENT_QUOTA_BYTES + # Nothing is lost silently: every input is in exactly one bucket. + assert sorted(f.filename for f in kept + over) == sorted(f.filename for f in files) + + def test_partition_then_budget_chain_on_files_that_each_pass_the_per_file_gate(self): + # Each file is under INLINE_DOCUMENT_MAX_BYTES (4 MiB) so the + # per-file gate keeps all three; only the aggregate budget catches + # the turn. This is exactly the hole the per-file gate left open. + per_file = file_models.INLINE_DOCUMENT_MAX_BYTES - 1024 + files = [_Attachment(f"f{i}.pdf", size=per_file) for i in range(3)] + sheet = _Attachment("data.xlsx", content_type=XLSX_MIME, size=6 * MB) + + inline, tabular, _, oversized = _partition_attachments(files + [sheet]) + assert oversized == [] and inline == files and tabular == [sheet] + + kept, over, requested = _apply_inline_byte_budget( + inline, file_models.INLINE_ATTACHMENTS_MAX_TOTAL_BYTES + ) + assert [f.filename for f in kept] == ["f0.pdf"] + assert [f.filename for f in over] == ["f1.pdf", "f2.pdf"] + assert requested == 3 * per_file + # The diverted spreadsheet never counted: it is not in the message. + 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 + + +class TestMessageFileCap: + def test_under_the_cap_is_untouched(self): + direct = [_Attachment("a.pdf")] + ids = ["u1", "u2"] + assert _apply_message_file_cap(direct, ids, 5) == (direct, ids, [], 0) + + def test_upload_ids_beyond_the_cap_are_counted_not_fetched(self): + # The old resolver truncated at 5 silently; now the route knows. + ids = [f"u{i}" for i in range(7)] + kept_direct, kept_ids, names, total = _apply_message_file_cap([], ids, 5) + + assert kept_direct == [] + assert kept_ids == ids[:5] + assert names == [] + assert total == 2 + + def test_direct_files_beyond_the_cap_are_named(self): + direct = [_Attachment(f"d{i}.pdf") for i in range(6)] + kept_direct, kept_ids, names, total = _apply_message_file_cap(direct, [], 5) + + assert kept_direct == direct[:5] + assert kept_ids == [] + assert names == ["d5.pdf"] + assert total == 1 + + def test_direct_files_come_first_and_ids_fill_the_remainder(self): + direct = [_Attachment("d0.pdf"), _Attachment("d1.pdf")] + ids = ["u0", "u1", "u2", "u3", "u4"] + kept_direct, kept_ids, names, total = _apply_message_file_cap(direct, ids, 5) + + assert kept_direct == direct + assert kept_ids == ["u0", "u1", "u2"] + assert names == [] + assert total == 2 + + def test_direct_overflow_leaves_no_budget_for_ids(self): + direct = [_Attachment(f"d{i}.pdf") for i in range(6)] + kept_direct, kept_ids, names, total = _apply_message_file_cap(direct, ["u0"], 5) + + assert len(kept_direct) == 5 + assert kept_ids == [] + assert names == ["d5.pdf"] + assert total == 2 + + def test_non_positive_cap_disables(self): + direct = [_Attachment(f"d{i}.pdf") for i in range(9)] + ids = [f"u{i}" for i in range(9)] + assert _apply_message_file_cap(direct, ids, 0) == (direct, ids, [], 0) + + def test_default_cap_matches_the_spa(self): + assert file_models.MAX_FILES_PER_MESSAGE == 5 + + +class TestGuidanceText: + def test_aggregate_case_names_the_files_and_says_follow_up(self): + over = [_Attachment("b.pdf"), _Attachment("c.pdf")] + text = _build_attachment_guidance([], [], [], None, over_budget=over) + + assert "`b.pdf`, `c.pdf`" in text + assert "together exceed the combined size limit" in text + assert "follow-up message" in text + + def test_aggregate_and_per_file_notes_are_different_sentences(self): + # The remedy differs: per-file → smaller file; aggregate → next message. + huge = [_Attachment("huge.pdf")] + over = [_Attachment("b.pdf")] + text = _build_attachment_guidance([], [], huge, None, over_budget=over) + + per_file, aggregate = text.split("\n\n") + assert "`huge.pdf`" in per_file and "Try a smaller file" in per_file + assert "`b.pdf`" in aggregate and "follow-up message" in aggregate + assert "`b.pdf`" not in per_file + + def test_count_cap_note_with_names(self): + text = _build_attachment_guidance( + [], [], [], None, + dropped_over_count_names=["f.pdf"], dropped_over_count_total=1, max_files=5, + ) + assert "Only the first 5 files per message are attached" in text + assert "`f.pdf` were not" in text + + def test_count_cap_note_with_names_and_unnamed_ids(self): + text = _build_attachment_guidance( + [], [], [], None, + dropped_over_count_names=["f.pdf"], dropped_over_count_total=3, max_files=5, + ) + assert "`f.pdf` and 2 more were not" in text + + def test_count_cap_note_count_only_when_ids_were_dropped(self): + text = _build_attachment_guidance( + [], [], [], None, dropped_over_count_total=2, max_files=5, + ) + assert "2 more files were not" in text + assert "follow-up message" in text + + def test_count_cap_note_singular(self): + text = _build_attachment_guidance( + [], [], [], None, dropped_over_count_total=1, max_files=5, + ) + assert "1 more file was not" in text + + def test_silent_when_nothing_was_trimmed(self): + assert _build_attachment_guidance([], [], [], None) == "" + assert _build_attachment_guidance( + [], [], [], None, over_budget=[], dropped_over_count_names=[], dropped_over_count_total=0 + ) == "" + + def test_existing_notes_are_unchanged_by_the_new_defaults(self): + huge = [_Attachment("huge.pdf")] + assert _build_attachment_guidance([], [], huge, None) == _build_attachment_guidance( + [], [], huge, None, over_budget=None, dropped_over_count_total=0 + ) + + +class TestMarkerNames: + def test_over_budget_files_are_excluded_like_oversized_ones(self): + # Both were dropped from the turn; the SPA must not rebuild a card. + a, b, c = _Attachment("a.pdf"), _Attachment("b.pdf"), _Attachment("c.pdf") + assert _attachment_marker_names([a, b, c], [c] + [b]) == ["a.pdf"] + + +class TestKillSwitch: + @pytest.mark.parametrize("value", [None, "", "true", "TRUE", "yes", "0"]) + def test_default_on(self, monkeypatch, value): + if value is None: + monkeypatch.delenv("ATTACHMENT_TURN_GUARD_ENABLED", raising=False) + else: + monkeypatch.setenv("ATTACHMENT_TURN_GUARD_ENABLED", value) + assert attachment_turn_guard_enabled() is True + + @pytest.mark.parametrize("value", ["false", "FALSE", " False "]) + def test_only_literal_false_disables(self, monkeypatch, value): + monkeypatch.setenv("ATTACHMENT_TURN_GUARD_ENABLED", value) + assert attachment_turn_guard_enabled() is False + + +class TestOverQuotaMetric: + def test_emits_one_content_free_record_in_bytes(self, monkeypatch): + from apis.shared.observability import emf, prompt_cache + + calls = [] + monkeypatch.setattr(prompt_cache, "prompt_cache_observability_enabled", lambda: True) + monkeypatch.setattr( + emf, "emit_emf_metrics", + lambda namespace, metrics, properties=None, units=None: calls.append( + (namespace, metrics, properties, units) + ), + ) + _emit_attachment_over_quota_metric( + requested_bytes=8 * MB, cap_bytes=7_500_000, inline_count=2, dropped_count=1 + ) + + assert calls == [( + "AgentCoreStack/Compaction", + {"AttachmentTurnOverQuota": 8 * MB}, + {"capBytes": 7_500_000, "inlineFileCount": 2, "droppedFileCount": 1}, + {"AttachmentTurnOverQuota": "Bytes"}, + )] + # Content-free: no filename, no user, no session in the record. + assert not any("name" in k.lower() and "file" not in k.lower() for k in calls[0][2]) + + def test_silenced_with_the_observability_layer(self, monkeypatch): + from apis.shared.observability import emf, prompt_cache + + monkeypatch.setattr(prompt_cache, "prompt_cache_observability_enabled", lambda: False) + monkeypatch.setattr( + emf, "emit_emf_metrics", lambda *a, **k: pytest.fail("must not emit") + ) + _emit_attachment_over_quota_metric(1, 1, 1, 1) + + def test_never_raises(self, monkeypatch): + from apis.shared.observability import emf, prompt_cache + + monkeypatch.setattr(prompt_cache, "prompt_cache_observability_enabled", lambda: True) + + def boom(*a, **k): + raise RuntimeError("emf down") + + monkeypatch.setattr(emf, "emit_emf_metrics", boom) + _emit_attachment_over_quota_metric(1, 1, 1, 1) + + +class TestEstimateDecodedSize: + 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 diff --git a/backend/tests/shared/test_files.py b/backend/tests/shared/test_files.py index 1a5dddc73..01a13cf1e 100644 --- a/backend/tests/shared/test_files.py +++ b/backend/tests/shared/test_files.py @@ -132,3 +132,18 @@ async def test_resolve_max_files(self, file_repository, s3_bucket, aws): resolver._file_repository = file_repository files = await resolver.resolve_files("u1", [f"f{i}" for i in range(10)], max_files=3) assert len(files) == 3 + + @pytest.mark.asyncio + async def test_resolve_files_no_cap_when_max_files_is_none(self, file_repository, s3_bucket, aws): + # The chat route applies the per-message cap itself (so it can tell + # the user which files were left out) and passes None here. + import boto3 + from apis.shared.files.file_resolver import FileResolver + s3 = boto3.client("s3", region_name="us-east-1") + for i in range(7): + s3.put_object(Bucket=s3_bucket, Key=f"uploads/u1/f{i}", Body=b"x") + await file_repository.create_file(_make_file(f"f{i}")) + resolver = FileResolver(s3_client=s3) + resolver._file_repository = file_repository + files = await resolver.resolve_files("u1", [f"f{i}" for i in range(7)], max_files=None) + assert len(files) == 7 diff --git a/docs/feature-summaries/FILE_UPLOAD_IMPLEMENTATION.md b/docs/feature-summaries/FILE_UPLOAD_IMPLEMENTATION.md index 0758fa432..a52ca56fc 100644 --- a/docs/feature-summaries/FILE_UPLOAD_IMPLEMENTATION.md +++ b/docs/feature-summaries/FILE_UPLOAD_IMPLEMENTATION.md @@ -206,8 +206,11 @@ PK: USER#{userId}, SK: QUOTA | `S3_USER_FILES_BUCKET_NAME` | S3 bucket name | | `DYNAMODB_USER_FILES_TABLE_NAME` | DynamoDB table name | | `FILE_UPLOAD_MAX_SIZE_BYTES` | Max file size (default 4MB) | -| `FILE_UPLOAD_MAX_FILES_PER_MESSAGE` | Max files per message (default 5) | +| `FILE_UPLOAD_MAX_FILES_PER_MESSAGE` | Max files per message (default 5). Enforced server-side by the inference API across both `files` and `file_upload_ids`; files beyond the cap are reported to the user in the attachment note, not silently dropped. `0` disables. | | `FILE_UPLOAD_USER_QUOTA_BYTES` | User quota (default 1GB) | +| `INLINE_DOCUMENT_MAX_BYTES` | Per-file inline document cap (default 4MB); larger non-tabular files are skipped with a note | +| `INLINE_ATTACHMENTS_MAX_TOTAL_BYTES` | Per-turn budget for all inline attachments, images included (default 7.5MB = the 10MB AgentCore Memory event quota ÷ base64's 4/3). Files that would push the turn over it are skipped with a note instead of failing the history write. `0` disables. | +| `ATTACHMENT_TURN_GUARD_ENABLED` | Kill switch for the two caps above (default on; only `false` disables) | ## Quota Management