From ef7c523f1a8d23c1ca4f06f1c8a86bb5be58eb0a Mon Sep 17 00:00:00 2001 From: ciregenz Date: Sat, 23 May 2026 03:11:58 -0700 Subject: [PATCH] [eric] split: extract history compaction helpers --- backend/apps/agents/agent_manager.py | 171 ++-------------------- backend/apps/agents/history_compaction.py | 166 +++++++++++++++++++++ 2 files changed, 176 insertions(+), 161 deletions(-) create mode 100644 backend/apps/agents/history_compaction.py diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 4fdc6f71..2a453560 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -50,6 +50,13 @@ from backend.apps.agents.tool_catalog import ( _is_fully_denied, ) from backend.apps.agents.aux_llm import _safe_resp_text +from backend.apps.agents.history_compaction import ( + _approx_tokens, + _build_history_prefix, + _get_branch_messages, + _summarize_message_block, + _truncate_large_tool_result, +) logger = logging.getLogger(__name__) @@ -611,60 +618,6 @@ class AgentManager: sections.append(f"[Using skill: {name}]\n\n{content}") return "\n\n".join(sections) - @staticmethod - def _get_branch_messages(session) -> list: - """Return the linear message list for the active branch, walking the branch tree.""" - branch_id = session.active_branch_id or "main" - branch = session.branches.get(branch_id) - - if not branch or not branch.fork_point_message_id: - return [m for m in session.messages if m.branch_id == "main" or m.branch_id == branch_id] - - segments = [] - cur = branch - cur_id = branch_id - visited = set() - while cur and cur.fork_point_message_id: - if cur_id in visited: - break - visited.add(cur_id) - segments.insert(0, {"branch_id": cur_id, "up_to": cur.fork_point_message_id}) - cur_id = cur.parent_branch_id or "main" - cur = session.branches.get(cur_id) - segments.insert(0, {"branch_id": cur_id, "up_to": None}) - - result = [] - for i, seg in enumerate(segments): - fork_msg_id = seg["up_to"] - if fork_msg_id: - fork_idx = next((j for j, m in enumerate(session.messages) if m.id == fork_msg_id), len(session.messages)) - result.extend(m for m in session.messages[:fork_idx] if m.branch_id == seg["branch_id"]) - else: - next_fork = segments[i + 1]["up_to"] if i + 1 < len(segments) else None - if next_fork: - fork_idx = next((j for j, m in enumerate(session.messages) if m.id == next_fork), len(session.messages)) - result.extend(m for m in session.messages[:fork_idx] if m.branch_id == seg["branch_id"]) - else: - result.extend(m for m in session.messages if m.branch_id == seg["branch_id"]) - - if not any(m.branch_id == branch_id for m in result): - result.extend(m for m in session.messages if m.branch_id == branch_id) - return result - - @staticmethod - def _build_history_prefix(messages) -> str: - """Format branch messages into a conversation summary for context injection.""" - lines = [] - for m in messages: - if m.role not in ("user", "assistant") or getattr(m, "hidden", False): - continue - text = m.content if isinstance(m.content, str) else str(m.content) - label = "User" if m.role == "user" else "Assistant" - lines.append(f"{label}: {text}") - if not lines: - return "" - return "\n" + "\n".join(lines) + "\n" - # ------------------------------------------------------------------ # Compaction & token guard (Phase 2) # @@ -679,72 +632,6 @@ class AgentManager: # surfaces from the catch-all # ------------------------------------------------------------------ - @staticmethod - def _approx_tokens(text: str) -> int: - """Conservative chars/4 estimate. Used for the pre-send guard - and the compaction trigger when a precise count_tokens isn't - cheap (or the route isn't Anthropic). Errs slightly high so we - compact a touch earlier than strictly necessary.""" - return max(1, len(text or "") // 4) - - @staticmethod - def _summarize_message_block(messages: list) -> str: - """Programmatic, no-LLM summary of a message slice. Mirrors the - shape of browser_agent._summarize_messages: extracts the original - user task, counts tool calls, captures the last assistant text. - Cheap, deterministic, and never makes a network call, so - compaction itself adds zero latency to the user's turn. - """ - if not messages: - return "" - - initial_task = "" - for m in messages: - if getattr(m, "role", "") == "user": - content = getattr(m, "content", "") - txt = content if isinstance(content, str) else str(content) - if txt.strip(): - initial_task = txt.strip()[:400] - break - - tool_calls_by_name: dict[str, int] = {} - last_tool_results = 0 - last_assistant_text = "" - for m in messages: - role = getattr(m, "role", "") - if role == "tool_call": - content = getattr(m, "content", {}) or {} - name = (content.get("tool") if isinstance(content, dict) else None) or "unknown" - tool_calls_by_name[name] = tool_calls_by_name.get(name, 0) + 1 - elif role == "tool_result": - last_tool_results += 1 - elif role == "assistant": - content = getattr(m, "content", "") - if isinstance(content, str) and content.strip(): - last_assistant_text = content.strip() - elif isinstance(content, list): - for block in content: - if isinstance(block, dict) and block.get("type") == "text": - txt = (block.get("text") or "").strip() - if txt: - last_assistant_text = txt - - parts = [""] - parts.append("[The following is a programmatic summary of earlier turns in this session. Originals are preserved on disk and viewable via the chat UI's compaction drawer.]") - if initial_task: - parts.append(f'Initial user request: "{initial_task}"') - if tool_calls_by_name: - total = sum(tool_calls_by_name.values()) - top = sorted(tool_calls_by_name.items(), key=lambda kv: -kv[1])[:8] - parts.append(f"Tool calls so far ({total} total): " + ", ".join(f"{n}×{c}" for n, c in top)) - if last_tool_results: - parts.append(f"Tool results received: {last_tool_results}") - if last_assistant_text: - parts.append("Last assistant message:") - parts.append(last_assistant_text[:1200]) - parts.append("") - return "\n".join(parts) - def _maybe_compact(self, session: AgentSession, force: bool = False) -> bool: """Run summarizer when ctx_used_pct >= compact_threshold_pct (or force). @@ -757,7 +644,7 @@ class AgentManager: ctx_used = session.tokens.get("input", 0) / max(1, session.context_window) if not force and ctx_used < session.compact_threshold_pct: return False - msgs = self._get_branch_messages(session) + msgs = _get_branch_messages(session) if len(msgs) < 4: return False # Summarize everything up to (but not including) the last 6 @@ -773,44 +660,6 @@ class AgentManager: session.compacted_through_msg_id = last_id return True - @staticmethod - def _truncate_large_tool_result(content: object, session_id: str, msg_id: str, max_bytes: int = 50_000) -> tuple[object, str | None]: - """Spill a large tool_result body to disk, return a truncated - inline replacement plus the on-disk path (or None if untouched). - - Storage is session-scoped under data/sessions//blobs/, - never honors caller-supplied paths (defense against path - traversal). The inline replacement keeps the first 4KB so the - model retains some signal about what was returned. - """ - if not isinstance(content, str): - try: - serialized = json.dumps(content) if not isinstance(content, str) else content - except Exception: - serialized = str(content) - else: - serialized = content - if len(serialized.encode("utf-8")) <= max_bytes: - return content, None - blobs_dir = os.path.join(SESSIONS_DIR, session_id, "blobs") - os.makedirs(blobs_dir, exist_ok=True) - # Sanitize msg_id (it's UUID hex, but be defensive). - safe_msg_id = re.sub(r"[^a-zA-Z0-9_-]", "", str(msg_id))[:64] or "blob" - blob_path = os.path.join(blobs_dir, f"{safe_msg_id}.txt") - try: - with open(blob_path, "w", encoding="utf-8") as f: - f.write(serialized) - except Exception as e: - logger.warning(f"Failed to spill tool result to {blob_path}: {e}") - return content, None - head = serialized[:4_000] - replacement = ( - f"{head}\n\n" - f"[truncated, full output ({len(serialized)} chars) saved to {blob_path}. " - f"Ask the user or run a follow-up tool call if you need the rest.]" - ) - return replacement, blob_path - def _build_prompt_content(self, prompt: str, images: list | None = None, context_paths: list | None = None, forced_tools: list[str] | None = None, attached_skills: list | None = None, api_type: str = "anthropic", model: str = ""): """Build message content for the Anthropic SDK's prompt stream. @@ -1689,7 +1538,7 @@ class AgentManager: # at *write* time (before the next turn ships history to the # SDK) so the bloat never re-enters context. try: - truncated_content, blob_path = self._truncate_large_tool_result( + truncated_content, blob_path = _truncate_large_tool_result( result_msg.content, session.id, result_msg.id ) if blob_path: @@ -2470,7 +2319,7 @@ class AgentManager: if session.needs_fork: session.needs_fork = False elif len(session.messages) > 1: - history = self._build_history_prefix(self._get_branch_messages(session)) + history = _build_history_prefix(_get_branch_messages(session)) if history: if isinstance(prompt_content, str): prompt_content = history + "\n\n" + prompt_content diff --git a/backend/apps/agents/history_compaction.py b/backend/apps/agents/history_compaction.py new file mode 100644 index 00000000..44841895 --- /dev/null +++ b/backend/apps/agents/history_compaction.py @@ -0,0 +1,166 @@ +import json +import logging +import os +import re + +from backend.config.paths import SESSIONS_DIR + +logger = logging.getLogger(__name__) + + +def _get_branch_messages(session) -> list: + """Return the linear message list for the active branch, walking the branch tree.""" + branch_id = session.active_branch_id or "main" + branch = session.branches.get(branch_id) + + if not branch or not branch.fork_point_message_id: + return [m for m in session.messages if m.branch_id == "main" or m.branch_id == branch_id] + + segments = [] + cur = branch + cur_id = branch_id + visited = set() + while cur and cur.fork_point_message_id: + if cur_id in visited: + break + visited.add(cur_id) + segments.insert(0, {"branch_id": cur_id, "up_to": cur.fork_point_message_id}) + cur_id = cur.parent_branch_id or "main" + cur = session.branches.get(cur_id) + segments.insert(0, {"branch_id": cur_id, "up_to": None}) + + result = [] + for i, seg in enumerate(segments): + fork_msg_id = seg["up_to"] + if fork_msg_id: + fork_idx = next((j for j, m in enumerate(session.messages) if m.id == fork_msg_id), len(session.messages)) + result.extend(m for m in session.messages[:fork_idx] if m.branch_id == seg["branch_id"]) + else: + next_fork = segments[i + 1]["up_to"] if i + 1 < len(segments) else None + if next_fork: + fork_idx = next((j for j, m in enumerate(session.messages) if m.id == next_fork), len(session.messages)) + result.extend(m for m in session.messages[:fork_idx] if m.branch_id == seg["branch_id"]) + else: + result.extend(m for m in session.messages if m.branch_id == seg["branch_id"]) + + if not any(m.branch_id == branch_id for m in result): + result.extend(m for m in session.messages if m.branch_id == branch_id) + return result + + +def _build_history_prefix(messages) -> str: + """Format branch messages into a conversation summary for context injection.""" + lines = [] + for m in messages: + if m.role not in ("user", "assistant") or getattr(m, "hidden", False): + continue + text = m.content if isinstance(m.content, str) else str(m.content) + label = "User" if m.role == "user" else "Assistant" + lines.append(f"{label}: {text}") + if not lines: + return "" + return "\n" + "\n".join(lines) + "\n" + + +def _approx_tokens(text: str) -> int: + """Conservative chars/4 estimate. Used for the pre-send guard + and the compaction trigger when a precise count_tokens isn't + cheap (or the route isn't Anthropic). Errs slightly high so we + compact a touch earlier than strictly necessary.""" + return max(1, len(text or "") // 4) + + +def _summarize_message_block(messages: list) -> str: + """Programmatic, no-LLM summary of a message slice. Mirrors the + shape of browser_agent._summarize_messages: extracts the original + user task, counts tool calls, captures the last assistant text. + Cheap, deterministic, and never makes a network call, so + compaction itself adds zero latency to the user's turn. + """ + if not messages: + return "" + + initial_task = "" + for m in messages: + if getattr(m, "role", "") == "user": + content = getattr(m, "content", "") + txt = content if isinstance(content, str) else str(content) + if txt.strip(): + initial_task = txt.strip()[:400] + break + + tool_calls_by_name: dict[str, int] = {} + last_tool_results = 0 + last_assistant_text = "" + for m in messages: + role = getattr(m, "role", "") + if role == "tool_call": + content = getattr(m, "content", {}) or {} + name = (content.get("tool") if isinstance(content, dict) else None) or "unknown" + tool_calls_by_name[name] = tool_calls_by_name.get(name, 0) + 1 + elif role == "tool_result": + last_tool_results += 1 + elif role == "assistant": + content = getattr(m, "content", "") + if isinstance(content, str) and content.strip(): + last_assistant_text = content.strip() + elif isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + txt = (block.get("text") or "").strip() + if txt: + last_assistant_text = txt + + parts = [""] + parts.append("[The following is a programmatic summary of earlier turns in this session. Originals are preserved on disk and viewable via the chat UI's compaction drawer.]") + if initial_task: + parts.append(f'Initial user request: "{initial_task}"') + if tool_calls_by_name: + total = sum(tool_calls_by_name.values()) + top = sorted(tool_calls_by_name.items(), key=lambda kv: -kv[1])[:8] + parts.append(f"Tool calls so far ({total} total): " + ", ".join(f"{n}×{c}" for n, c in top)) + if last_tool_results: + parts.append(f"Tool results received: {last_tool_results}") + if last_assistant_text: + parts.append("Last assistant message:") + parts.append(last_assistant_text[:1200]) + parts.append("") + return "\n".join(parts) + + +def _truncate_large_tool_result(content: object, session_id: str, msg_id: str, max_bytes: int = 50_000) -> tuple[object, str | None]: + """Spill a large tool_result body to disk, return a truncated + inline replacement plus the on-disk path (or None if untouched). + + Storage is session-scoped under data/sessions//blobs/, + never honors caller-supplied paths (defense against path + traversal). The inline replacement keeps the first 4KB so the + model retains some signal about what was returned. + """ + if not isinstance(content, str): + try: + serialized = json.dumps(content) if not isinstance(content, str) else content + except Exception: + serialized = str(content) + else: + serialized = content + if len(serialized.encode("utf-8")) <= max_bytes: + return content, None + blobs_dir = os.path.join(SESSIONS_DIR, session_id, "blobs") + os.makedirs(blobs_dir, exist_ok=True) + # Sanitize msg_id (it's UUID hex, but be defensive). + safe_msg_id = re.sub(r"[^a-zA-Z0-9_-]", "", str(msg_id))[:64] or "blob" + blob_path = os.path.join(blobs_dir, f"{safe_msg_id}.txt") + try: + with open(blob_path, "w", encoding="utf-8") as f: + f.write(serialized) + except Exception as e: + logger.warning(f"Failed to spill tool result to {blob_path}: {e}") + return content, None + head = serialized[:4_000] + replacement = ( + f"{head}\n\n" + f"[truncated, full output ({len(serialized)} chars) saved to {blob_path}. " + f"Ask the user or run a follow-up tool call if you need the rest.]" + ) + return replacement, blob_path