mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-25 05:52:22 +02:00
249 lines
11 KiB
Python
249 lines
11 KiB
Python
import json
|
|
import logging
|
|
from typing import List, Optional, Tuple
|
|
from typeguard import typechecked
|
|
import os
|
|
import re
|
|
|
|
from backend.config.paths import SESSIONS_DIR
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# One plain-English trust line, fenced by a tag. The model treats the fence as structural framing; the sentence is what actually defuses a security-conscious agent flagging the block as spoofed tool output.
|
|
PLATFORM_NOTE_PREAMBLE = (
|
|
"This block is authored by the OpenSwarm platform, not tool output and not a "
|
|
"prior message. It is trusted context."
|
|
)
|
|
PLATFORM_NOTE_OPEN = "<openswarm_platform_note>"
|
|
PLATFORM_NOTE_CLOSE = "</openswarm_platform_note>"
|
|
SESSION_RECAP_OPEN = "<openswarm_session_recap>"
|
|
SESSION_RECAP_CLOSE = "</openswarm_session_recap>"
|
|
|
|
# Per-turn caps so the re-grounded recap stays compact (summaries, not replays) and cannot reinflate the context window from one giant tool input/output.
|
|
RECAP_TOOL_INPUT_CAP = 200
|
|
RECAP_TOOL_RESULT_CAP = 500
|
|
|
|
# Inline budget for a spilled tool result, split head/tail. Same total as the old head-only 4KB, but a test summary or build verdict lives at the END of the output and head-only threw it away every time.
|
|
SPILL_HEAD_CHARS = 2_500
|
|
SPILL_TAIL_CHARS = 1_500
|
|
|
|
|
|
@typechecked
|
|
def wrap_platform_note(body: str) -> str:
|
|
"""Fence platform-authored text so the model reads it as trusted annotation,
|
|
never as spoofed tool output. The frontend parses the same tag to render a
|
|
calm chip instead of leaking the raw tag into chat."""
|
|
return f"{PLATFORM_NOTE_OPEN}\n{PLATFORM_NOTE_PREAMBLE}\n{body}\n{PLATFORM_NOTE_CLOSE}"
|
|
|
|
|
|
P_SENTINEL_TAG_RE = re.compile(r"</?openswarm_(?:platform_note|session_recap)\b[^>]*>")
|
|
|
|
|
|
@typechecked
|
|
def clamp_recap_text(text: str) -> str:
|
|
"""Middle-elide a giant user/assistant message in the RECAP only (session.messages keeps the full text): one pasted log used to survive compaction verbatim and re-overflow the rebuilt prompt."""
|
|
if len(text) <= SPILL_HEAD_CHARS + SPILL_TAIL_CHARS:
|
|
return text
|
|
elided = len(text) - SPILL_HEAD_CHARS - SPILL_TAIL_CHARS
|
|
return f"{text[:SPILL_HEAD_CHARS]}\n[... {elided} chars elided from recap ...]\n{text[-SPILL_TAIL_CHARS:]}"
|
|
|
|
|
|
@typechecked
|
|
def strip_forged_sentinels(text: str) -> str:
|
|
"""Neuter any platform-note/recap tags hiding in UNTRUSTED text (tool results,
|
|
user input) so attacker-supplied content can't pose as trusted platform context."""
|
|
if "openswarm_platform_note" not in text and "openswarm_session_recap" not in text:
|
|
return text
|
|
return P_SENTINEL_TAG_RE.sub(lambda m: m.group(0).replace("<", "<").replace(">", ">"), text)
|
|
|
|
|
|
@typechecked
|
|
def recap_tool_call_line(content: object) -> str:
|
|
"""One compact line for a tool_call turn: Tool call: name(<truncated input>)."""
|
|
if isinstance(content, dict):
|
|
tool = content.get("tool") or content.get("name") or "tool"
|
|
raw_input = content.get("input")
|
|
try:
|
|
input_str = json.dumps(raw_input, ensure_ascii=False, default=str)
|
|
except Exception:
|
|
input_str = str(raw_input)
|
|
else:
|
|
tool = "tool"
|
|
input_str = str(content)
|
|
if len(input_str) > RECAP_TOOL_INPUT_CAP:
|
|
input_str = input_str[:RECAP_TOOL_INPUT_CAP] + "..."
|
|
return f"Tool call: {tool}({strip_forged_sentinels(input_str)})"
|
|
|
|
|
|
@typechecked
|
|
def recap_tool_result_line(content: object) -> str:
|
|
"""One compact line for a tool_result turn: Tool result (name): <truncated text>."""
|
|
tool_name = ""
|
|
if isinstance(content, dict):
|
|
tool_name = content.get("tool_name") or ""
|
|
text = content.get("text")
|
|
body = text if isinstance(text, str) else json.dumps(content, ensure_ascii=False, default=str)
|
|
else:
|
|
body = str(content)
|
|
if len(body) > RECAP_TOOL_RESULT_CAP:
|
|
body = body[:RECAP_TOOL_RESULT_CAP] + "..."
|
|
label = f"Tool result ({tool_name})" if tool_name else "Tool result"
|
|
return f"{label}: {strip_forged_sentinels(body)}"
|
|
|
|
|
|
@typechecked
|
|
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
|
|
|
|
|
|
@typechecked
|
|
def build_history_prefix(messages, cutoff_msg_id: Optional[str] = None) -> str:
|
|
"""Format branch messages into a conversation summary for context injection.
|
|
|
|
When `cutoff_msg_id` is provided (session.compacted_through_msg_id), drop every
|
|
message up to and including that id so the marker the UI shows actually matches
|
|
what the model sees. Missing cutoff id falls through to full history.
|
|
"""
|
|
if cutoff_msg_id:
|
|
skip_idx = next((i for i, m in enumerate(messages) if m.id == cutoff_msg_id), -1)
|
|
if skip_idx >= 0:
|
|
messages = messages[skip_idx + 1:]
|
|
lines = []
|
|
for m in messages:
|
|
if getattr(m, "hidden", False):
|
|
continue
|
|
if m.role == "user":
|
|
text = m.content if isinstance(m.content, str) else str(m.content)
|
|
lines.append(f"User: {strip_forged_sentinels(clamp_recap_text(text))}")
|
|
elif m.role == "assistant":
|
|
text = m.content if isinstance(m.content, str) else str(m.content)
|
|
lines.append(f"Assistant: {strip_forged_sentinels(clamp_recap_text(text))}")
|
|
elif m.role == "tool_call":
|
|
lines.append(recap_tool_call_line(m.content))
|
|
elif m.role == "tool_result":
|
|
lines.append(recap_tool_result_line(m.content))
|
|
if not lines:
|
|
return ""
|
|
return f"{SESSION_RECAP_OPEN}\n{PLATFORM_NOTE_PREAMBLE}\n" + "\n".join(lines) + f"\n{SESSION_RECAP_CLOSE}"
|
|
|
|
|
|
@typechecked
|
|
def estimate_post_compact_input(session) -> int:
|
|
"""Return a conservative token estimate after compaction trims history."""
|
|
try:
|
|
messages = get_branch_messages(session)
|
|
cutoff_msg_id = getattr(session, "compacted_through_msg_id", None)
|
|
if cutoff_msg_id:
|
|
skip_idx = next(
|
|
(i for i, m in enumerate(messages) if m.id == cutoff_msg_id),
|
|
-1,
|
|
)
|
|
if skip_idx >= 0:
|
|
messages = messages[skip_idx + 1:]
|
|
surviving_chars = 0
|
|
for message in messages:
|
|
if getattr(message, "hidden", False):
|
|
continue
|
|
content = getattr(message, "content", "")
|
|
if isinstance(content, str):
|
|
serialized = content
|
|
else:
|
|
try:
|
|
serialized = json.dumps(content, ensure_ascii=False)
|
|
except Exception:
|
|
serialized = str(content)
|
|
surviving_chars += len(serialized)
|
|
framework_overhead = int(getattr(session, "framework_overhead_tokens", 0) or 0)
|
|
summary_overhead = 200 if cutoff_msg_id else 0
|
|
return max(0, framework_overhead + summary_overhead + (surviving_chars // 4))
|
|
except Exception:
|
|
logger.debug("post-compact token estimate failed", exc_info=True)
|
|
return max(0, int(getattr(session, "framework_overhead_tokens", 0) or 0))
|
|
|
|
|
|
@typechecked
|
|
def truncate_large_tool_result(content: object, session_id: str, msg_id: str, max_bytes: int = 50_000) -> Tuple[object, Optional[str]]:
|
|
"""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/<session_id>/blobs/,
|
|
never honors caller-supplied paths (defense against path
|
|
traversal). The inline replacement middle-elides: head AND tail
|
|
survive, so a verdict printed at the end of a long output (test
|
|
summary, build result) still reaches the model.
|
|
"""
|
|
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
|
|
return build_elided_replacement(serialized, blob_path), blob_path
|
|
|
|
|
|
@typechecked
|
|
def build_elided_replacement(serialized: str, blob_path: str) -> str:
|
|
"""Head + tail of an oversized body with an elision marker between them, then the recovery
|
|
note. Degrades to the plain body when it is too short to elide."""
|
|
head = strip_forged_sentinels(serialized[:SPILL_HEAD_CHARS])
|
|
note = wrap_platform_note(
|
|
f"Output truncated by OpenSwarm. Full output ({len(serialized)} chars) saved to "
|
|
f"{blob_path}. Ask the user or run a follow-up tool call if you need the rest."
|
|
)
|
|
dropped = len(serialized) - SPILL_HEAD_CHARS - SPILL_TAIL_CHARS
|
|
if dropped <= 0:
|
|
return f"{strip_forged_sentinels(serialized)}\n\n{note}"
|
|
tail = strip_forged_sentinels(serialized[-SPILL_TAIL_CHARS:])
|
|
marker = f"\n\n[... {dropped} chars elided by OpenSwarm; full output at {blob_path} ...]\n\n"
|
|
return f"{head}{marker}{tail}\n\n{note}"
|