|
| 1 | +"""Smart truncation of oversized record bodies (`_truncate_attributes`). |
| 2 | +
|
| 3 | +The legacy path replaces an over-cap body with `{"_truncated": True}`, losing the event's |
| 4 | +type/id and all content. Smart truncation preserves the event shape + partial content so the |
| 5 | +record log stays reconstructable server-side. These pin that contract. |
| 6 | +""" |
| 7 | + |
| 8 | +from orjson import dumps |
| 9 | + |
| 10 | +from oss.src.core.sessions.records.streaming import ( |
| 11 | + MAX_ATTRIBUTES_BYTES, |
| 12 | + _TRUNCATION_MARKER, |
| 13 | + _truncate_attributes, |
| 14 | +) |
| 15 | + |
| 16 | + |
| 17 | +def _size(obj) -> int: |
| 18 | + return len(dumps(obj)) |
| 19 | + |
| 20 | + |
| 21 | +def test_under_budget_returns_unchanged(): |
| 22 | + attrs = {"type": "message", "text": "hi"} |
| 23 | + out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) |
| 24 | + assert out is attrs # untouched, no _truncated marker |
| 25 | + |
| 26 | + |
| 27 | +def test_large_string_field_is_trimmed_but_structure_preserved(): |
| 28 | + big = "x" * (MAX_ATTRIBUTES_BYTES * 2) |
| 29 | + attrs = {"type": "tool_result", "id": "call-1", "output": big} |
| 30 | + out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) |
| 31 | + |
| 32 | + # Discriminator fields survive (unlike the legacy whole-body drop). |
| 33 | + assert out["type"] == "tool_result" |
| 34 | + assert out["id"] == "call-1" |
| 35 | + # The big field is trimmed + marked, and the whole body now fits the cap. |
| 36 | + assert out["output"].endswith(_TRUNCATION_MARKER) |
| 37 | + assert len(out["output"]) < len(big) |
| 38 | + assert _size(out) <= MAX_ATTRIBUTES_BYTES |
| 39 | + # Metadata records what was trimmed. |
| 40 | + assert out["_truncated"]["fields"] == ["output"] |
| 41 | + assert out["_truncated"]["original_bytes"] == _size(attrs) |
| 42 | + |
| 43 | + |
| 44 | +def test_trims_the_largest_of_several_string_fields(): |
| 45 | + attrs = { |
| 46 | + "type": "message", |
| 47 | + "small": "ok", |
| 48 | + "text": "y" * (MAX_ATTRIBUTES_BYTES * 2), |
| 49 | + } |
| 50 | + out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) |
| 51 | + assert out["small"] == "ok" # small field untouched |
| 52 | + assert out["text"].endswith(_TRUNCATION_MARKER) |
| 53 | + assert _size(out) <= MAX_ATTRIBUTES_BYTES |
| 54 | + |
| 55 | + |
| 56 | +def test_non_string_bloat_falls_back_to_discriminator_only(): |
| 57 | + # A huge nested structure with no single big string leaf can't be string-trimmed. |
| 58 | + attrs = { |
| 59 | + "type": "tool_call", |
| 60 | + "id": "call-9", |
| 61 | + "input": {str(i): i for i in range(MAX_ATTRIBUTES_BYTES)}, |
| 62 | + } |
| 63 | + out = _truncate_attributes(attrs, MAX_ATTRIBUTES_BYTES, _size(attrs)) |
| 64 | + assert out["type"] == "tool_call" |
| 65 | + assert out["id"] == "call-9" |
| 66 | + assert out["_truncated"] is True |
| 67 | + assert _size(out) <= MAX_ATTRIBUTES_BYTES |
| 68 | + |
| 69 | + |
| 70 | +def test_non_dict_attributes_fall_back(): |
| 71 | + huge = "z" * (MAX_ATTRIBUTES_BYTES * 2) |
| 72 | + out = _truncate_attributes(huge, MAX_ATTRIBUTES_BYTES, _size(huge)) |
| 73 | + assert out == {"_truncated": True, "_original_bytes": _size(huge)} |
| 74 | + |
| 75 | + |
| 76 | +def test_smart_truncation_flag_is_reachable_from_the_env_object(): |
| 77 | + """The publish path reads `env.agenta.sessions.records.smart_truncation` inside a |
| 78 | + try/except that swallows anything and drops the record. When `SessionsConfig` was not |
| 79 | + attached to `AgentaConfig` this raised AttributeError, so every over-cap record was |
| 80 | + discarded instead of truncated, and the tests above still passed because they call |
| 81 | + `_truncate_attributes` directly and never touch the flag.""" |
| 82 | + from oss.src.utils.env import env |
| 83 | + |
| 84 | + assert isinstance(env.agenta.sessions.records.smart_truncation, bool) |
0 commit comments