diff --git a/api/streaming.py b/api/streaming.py index 6eeb38926..9f14f1dd0 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -4229,6 +4229,113 @@ def _assign_stable_message_ids(result_messages, *existing_arrays): return stamped +_POST_COMPRESSION_TOOL_RESULT_TOTAL_TOKENS = 4096 +_POST_COMPRESSION_TOOL_RESULT_MIN_SNIPPET_TOKENS = 256 +_POST_COMPRESSION_TOOL_RESULT_MARKER = "[WebUI compressed-context budget:" +_POST_COMPRESSION_TOOL_RESULT_NOTE_RE = re.compile( + r"^\[WebUI compressed-context budget: omitted \d+(?: of \d+)? chars " + r"\(~\d+ rough tokens\) from this tool result; " + r"the full output remains in the visible transcript/tool log\.\]$" +) +_ROUGH_TOKEN_CHARS = 4 + + +def _positive_int_value(value, default: int = 0) -> int: + try: + parsed = int(value or 0) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +def _rough_text_token_count(text: str) -> int: + text = str(text or "") + if not text: + return 0 + return max(1, (len(text) + (_ROUGH_TOKEN_CHARS - 1)) // _ROUGH_TOKEN_CHARS) + + +def _post_compression_tool_result_budget(compressor) -> int: + budget = _positive_int_value( + getattr(compressor, 'tail_token_budget', None), + _POST_COMPRESSION_TOOL_RESULT_TOTAL_TOKENS, + ) + threshold = _positive_int_value(getattr(compressor, 'threshold_tokens', None), 0) + if threshold: + budget = min(budget, max(512, threshold // 2)) + return max(512, budget) + + +def _compressed_context_tool_result_summary(text: str, *, original_tokens: int, keep_tokens: int) -> str: + text = str(text or "") + keep_chars = max(0, int(keep_tokens or 0) * _ROUGH_TOKEN_CHARS) + note_prefix = f"{_POST_COMPRESSION_TOOL_RESULT_MARKER} omitted " + note_suffix = ( + f" (~{int(original_tokens or 0)} rough tokens) from this tool result; " + "the full output remains in the visible transcript/tool log.]" + ) + if keep_chars < (_POST_COMPRESSION_TOOL_RESULT_MIN_SNIPPET_TOKENS * _ROUGH_TOKEN_CHARS): + return f"{note_prefix}{len(text)} chars{note_suffix}" + note_len_for_budget = len(f"{note_prefix}{len(text)} of {len(text)} chars{note_suffix}") + snippet_limit = max(0, keep_chars - note_len_for_budget - 2) + snippet = text[:snippet_limit].rstrip() + if not snippet: + return f"{note_prefix}{len(text)} chars{note_suffix}" + omitted_chars = max(0, len(text) - len(snippet)) + note = f"{note_prefix}{omitted_chars} of {len(text)} chars{note_suffix}" + return f"{snippet}\n\n{note}" + + +def _is_compressed_context_tool_result_summary(text: str) -> bool: + text = str(text or "").strip() + if not text: + return False + if _POST_COMPRESSION_TOOL_RESULT_NOTE_RE.fullmatch(text): + return True + _snippet, separator, note = text.rpartition("\n\n") + return bool(separator and _POST_COMPRESSION_TOOL_RESULT_NOTE_RE.fullmatch(note.strip())) + + +def _hard_prune_post_compression_tool_results(messages, *, compressor=None): + if not messages: + return messages, 0 + budget = _post_compression_tool_result_budget(compressor) + raw_tool_tokens = 0 + replacements = [] + for idx in range(len(messages) - 1, -1, -1): + msg = messages[idx] + if not isinstance(msg, dict) or msg.get('role') != 'tool': + continue + text = _message_text(msg.get('content', '')) + if not text: + continue + token_count = _rough_text_token_count(text) + if _is_compressed_context_tool_result_summary(text): + raw_tool_tokens += token_count + continue + remaining = max(0, budget - raw_tool_tokens) + if token_count <= remaining: + raw_tool_tokens += token_count + continue + replacements.append((idx, text, token_count, remaining)) + # Once a tool row exceeds the residual budget, older tool rows should + # not be allowed to spend that same residual budget again. + raw_tool_tokens = budget + + if not replacements: + return messages, 0 + + pruned = copy.deepcopy(list(messages)) + for idx, text, token_count, keep_tokens in replacements: + msg = pruned[idx] + msg['content'] = _compressed_context_tool_result_summary( + text, + original_tokens=token_count, + keep_tokens=keep_tokens, + ) + return pruned, len(replacements) + + def _prune_context_tool_results_after_compression(agent, context_messages): """Run the active compressor's cheap tool-result pruning on model context. @@ -4236,28 +4343,35 @@ def _prune_context_tool_results_after_compression(agent, context_messages): before producing the final answer. Those completed tail tool results are model-facing context, but they were produced after the compression pass and therefore did not go through the compressor's tool-output pruning. Apply the - same cheap pruning once more after a confirmed compression event. This keeps - the visible transcript untouched while preventing the next turn from seeing - raw post-compression tool dumps. + same cheap pruning once more after a confirmed compression event, then apply + a WebUI hard cap to retained tool-result payloads. This keeps the visible + transcript untouched while preventing the next turn from seeing raw + post-compression tool dumps that the compressor protected as recent tail. """ if not context_messages: return context_messages compressor = getattr(agent, 'context_compressor', None) prune = getattr(compressor, '_prune_old_tool_results', None) - if not callable(prune): - return context_messages - try: - pruned_messages, pruned_count = prune( - copy.deepcopy(context_messages), - protect_tail_count=getattr(compressor, 'protect_last_n', 20), - protect_tail_tokens=getattr(compressor, 'tail_token_budget', None), - ) - except Exception: - logger.debug("post-compression context tool-result pruning failed", exc_info=True) - return context_messages - if not pruned_count: - return context_messages - return _deduplicate_context_messages(pruned_messages) + pruned_messages = context_messages + if callable(prune): + try: + candidate_messages, pruned_count = prune( + copy.deepcopy(context_messages), + protect_tail_count=getattr(compressor, 'protect_last_n', 20), + protect_tail_tokens=getattr(compressor, 'tail_token_budget', None), + ) + if pruned_count: + pruned_messages = _deduplicate_context_messages(candidate_messages) + except Exception: + logger.debug("post-compression context tool-result pruning failed", exc_info=True) + + hard_pruned_messages, hard_pruned_count = _hard_prune_post_compression_tool_results( + pruned_messages, + compressor=compressor, + ) + if hard_pruned_count: + return _deduplicate_context_messages(hard_pruned_messages) + return pruned_messages def _restore_reasoning_metadata(previous_messages, updated_messages): diff --git a/tests/test_auto_compression_card.py b/tests/test_auto_compression_card.py index 4d3641ddb..00860594d 100644 --- a/tests/test_auto_compression_card.py +++ b/tests/test_auto_compression_card.py @@ -4,6 +4,7 @@ from api.compression_anchor import visible_messages_for_anchor from api.models import Session from api.streaming import ( _is_fallback_lifecycle_message, + _message_text, _prune_context_tool_results_after_compression, ) @@ -73,6 +74,136 @@ def test_post_compression_context_prunes_tail_tool_results_with_active_compresso assert context_messages[1]["content"] == "x" * 5000 +def test_post_compression_context_hard_prunes_protected_tail_tool_payload(): + class NoopCompressor: + protect_last_n = 20 + tail_token_budget = 1024 + threshold_tokens = 4096 + + def _prune_old_tool_results(self, messages, protect_tail_count, protect_tail_tokens=None): + return messages, 0 + + compressor = NoopCompressor() + agent = type("Agent", (), {"context_compressor": compressor})() + recent_payload = ("recent protected payload\n" * 1800) + "RECENT_TOOL_PAYLOAD_END" + session = Session( + session_id="budgetproof1", + title="Budget proof", + workspace="/tmp", + model="gpt-4o", + messages=[ + {"role": "assistant", "content": "", "tool_calls": [{"id": "call_recent"}]}, + {"role": "tool", "tool_call_id": "call_recent", "content": recent_payload}, + {"role": "assistant", "content": "Final answer"}, + ], + context_messages=[ + {"role": "assistant", "content": "", "tool_calls": [{"id": "call_recent"}]}, + {"role": "tool", "tool_call_id": "call_recent", "content": recent_payload}, + {"role": "assistant", "content": "Final answer"}, + ], + ) + + session.context_messages = _prune_context_tool_results_after_compression( + agent, + session.context_messages, + ) + + visible_tool = session.messages[1]["content"] + context_tool = session.context_messages[1]["content"] + context_tool_tokens = (len(_message_text(context_tool)) + 3) // 4 + + assert visible_tool == recent_payload + assert "RECENT_TOOL_PAYLOAD_END" in visible_tool + assert context_tool != recent_payload + assert "RECENT_TOOL_PAYLOAD_END" not in context_tool + assert "WebUI compressed-context budget" in context_tool + assert context_tool_tokens <= compressor.tail_token_budget + + +def test_post_compression_context_note_only_replacement_consumes_residual_budget(): + class NoopCompressor: + protect_last_n = 20 + tail_token_budget = 1024 + threshold_tokens = 4096 + + def _prune_old_tool_results(self, messages, protect_tail_count, protect_tail_tokens=None): + return messages, 0 + + compressor = NoopCompressor() + agent = type("Agent", (), {"context_compressor": compressor})() + old_small_payload = "s" * 80 + old_oversize_payload = "h" * 2000 + recent_payload = "r" * 4000 + context_messages = [ + {"role": "tool", "tool_call_id": "call_old_small", "content": old_small_payload}, + {"role": "tool", "tool_call_id": "call_old_oversize", "content": old_oversize_payload}, + {"role": "tool", "tool_call_id": "call_recent", "content": recent_payload}, + ] + + pruned = _prune_context_tool_results_after_compression(agent, context_messages) + + assert pruned[2]["content"] == recent_payload + assert pruned[1]["content"] != old_oversize_payload + assert "WebUI compressed-context budget" in pruned[1]["content"] + assert pruned[0]["content"] != old_small_payload + assert "WebUI compressed-context budget" in pruned[0]["content"] + assert context_messages[0]["content"] == old_small_payload + + +def test_post_compression_context_hard_prune_is_idempotent(): + class NoopCompressor: + protect_last_n = 20 + tail_token_budget = 1024 + threshold_tokens = 4096 + + def _prune_old_tool_results(self, messages, protect_tail_count, protect_tail_tokens=None): + return messages, 0 + + compressor = NoopCompressor() + agent = type("Agent", (), {"context_compressor": compressor})() + context_messages = [ + {"role": "tool", "tool_call_id": f"call_{idx}", "content": (f"tool {idx} payload\n" * 350)} + for idx in range(4) + ] + + once = _prune_context_tool_results_after_compression(agent, context_messages) + twice = _prune_context_tool_results_after_compression(agent, once) + + assert [msg["content"] for msg in twice] == [msg["content"] for msg in once] + assert any("WebUI compressed-context budget" in str(msg.get("content") or "") for msg in once) + assert context_messages[0]["content"] != once[0]["content"] + + +def test_post_compression_context_raw_marker_collision_still_prunes(): + class NoopCompressor: + protect_last_n = 20 + tail_token_budget = 1024 + threshold_tokens = 4096 + + def _prune_old_tool_results(self, messages, protect_tail_count, protect_tail_tokens=None): + return messages, 0 + + compressor = NoopCompressor() + agent = type("Agent", (), {"context_compressor": compressor})() + raw_payload = ( + "tool echoed marker text: [WebUI compressed-context budget: not a generated summary]\n" + + ("raw collision payload\n" * 600) + + "RAW_COLLISION_END" + ) + context_messages = [ + {"role": "tool", "tool_call_id": "call_collision", "content": raw_payload}, + ] + + once = _prune_context_tool_results_after_compression(agent, context_messages) + twice = _prune_context_tool_results_after_compression(agent, once) + + assert once[0]["content"] != raw_payload + assert "RAW_COLLISION_END" not in once[0]["content"] + assert "WebUI compressed-context budget" in once[0]["content"] + assert twice[0]["content"] == once[0]["content"] + assert context_messages[0]["content"] == raw_payload + + def test_auto_compression_running_sse_uses_active_session_running_card(): block = _compressing_listener_block()