"""Session-memory-backed compaction. This module returns the same ``CompactionResult`` shape as the LLM compact path, but uses the current session-memory notes as the summary and preserves a recent verbatim suffix. If any gate fails it returns ``None`` so the standard LLM compact path remains the fallback. """ from __future__ import annotations import logging import os from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Mapping, Sequence from openspace.services.conversation.messages import ( build_compact_boundary_message, build_compact_summary_message, extract_discovered_tool_names, get_message_uuid, is_compact_boundary_message, ) from openspace.services.memory.session_memory import ( find_message_cursor_index, get_last_summarized_message_id, get_session_memory_path, is_session_memory_empty, read_session_memory, set_last_summarized_message_id, truncate_session_memory_for_compact, wait_for_session_memory_extraction, ) if TYPE_CHECKING: from openspace.services.conversation.compact import CompactionResult from openspace.services.tooling.hooks import HookRegistry from openspace.services.tooling.context import ToolUseContext logger = logging.getLogger(__name__) OPENSPACE_SESSION_MEMORY_COMPACT_ENABLED_ENV = ( "OPENSPACE_SESSION_MEMORY_COMPACT_ENABLED" ) OPENSPACE_DISABLE_SESSION_MEMORY_COMPACT_ENV = ( "OPENSPACE_DISABLE_SESSION_MEMORY_COMPACT" ) OPENSPACE_SM_COMPACT_MIN_TOKENS_ENV = "OPENSPACE_SM_COMPACT_MIN_TOKENS" OPENSPACE_SM_COMPACT_MIN_TEXT_MESSAGES_ENV = ( "OPENSPACE_SM_COMPACT_MIN_TEXT_MESSAGES" ) OPENSPACE_SM_COMPACT_MAX_TOKENS_ENV = "OPENSPACE_SM_COMPACT_MAX_TOKENS" @dataclass(frozen=True, slots=True) class SessionMemoryCompactConfig: """OpenSpace session memory compaction thresholds.""" min_tokens: int = 10_000 min_text_block_messages: int = 5 max_tokens: int = 40_000 DEFAULT_SM_COMPACT_CONFIG = SessionMemoryCompactConfig() _sm_compact_config = DEFAULT_SM_COMPACT_CONFIG def _env_truthy(value: str | None) -> bool: return value is not None and value.lower() in {"1", "true", "yes", "on"} def _env_defined_falsy(value: str | None) -> bool: return value is not None and value.lower() in {"0", "false", "no", "off", ""} def set_session_memory_compact_config( config: SessionMemoryCompactConfig | Mapping[str, int], ) -> None: global _sm_compact_config if isinstance(config, SessionMemoryCompactConfig): _sm_compact_config = config return current = _sm_compact_config _sm_compact_config = SessionMemoryCompactConfig( min_tokens=int(config.get("min_tokens", config.get("minTokens", current.min_tokens))), min_text_block_messages=int( config.get( "min_text_block_messages", config.get("minTextBlockMessages", current.min_text_block_messages), ) ), max_tokens=int(config.get("max_tokens", config.get("maxTokens", current.max_tokens))), ) def get_session_memory_compact_config() -> SessionMemoryCompactConfig: return _sm_compact_config def reset_session_memory_compact_config() -> None: global _sm_compact_config _sm_compact_config = DEFAULT_SM_COMPACT_CONFIG def _config_from_env() -> SessionMemoryCompactConfig: return SessionMemoryCompactConfig( min_tokens=_parse_positive_int( os.environ.get(OPENSPACE_SM_COMPACT_MIN_TOKENS_ENV), _sm_compact_config.min_tokens, ), min_text_block_messages=_parse_positive_int( os.environ.get(OPENSPACE_SM_COMPACT_MIN_TEXT_MESSAGES_ENV), _sm_compact_config.min_text_block_messages, ), max_tokens=_parse_positive_int( os.environ.get(OPENSPACE_SM_COMPACT_MAX_TOKENS_ENV), _sm_compact_config.max_tokens, ), ) def _parse_positive_int(raw: str | None, default: int) -> int: if not raw: return default try: parsed = int(raw) except ValueError: return default return parsed if parsed > 0 else default def should_use_session_memory_compaction() -> bool: """Return whether session-memory compaction is enabled.""" if _env_truthy(os.environ.get(OPENSPACE_DISABLE_SESSION_MEMORY_COMPACT_ENV)): return False enabled = os.environ.get(OPENSPACE_SESSION_MEMORY_COMPACT_ENABLED_ENV) if _env_defined_falsy(enabled): return False if _env_truthy(enabled): return True return True def has_text_blocks(message: Mapping[str, Any]) -> bool: role = message.get("role") content = message.get("content") if role == "assistant": if isinstance(content, str): return bool(content) if isinstance(content, Sequence) and not isinstance( content, (str, bytes, bytearray), ): return any( isinstance(block, Mapping) and block.get("type") == "text" and bool(block.get("text")) for block in content ) if role == "user": if isinstance(content, str): return bool(content) if isinstance(content, Sequence) and not isinstance( content, (str, bytes, bytearray), ): return any( isinstance(block, str) or ( isinstance(block, Mapping) and block.get("type") == "text" and bool(block.get("text")) ) for block in content ) return False def _get_tool_result_ids(message: Mapping[str, Any]) -> list[str]: if message.get("role") == "tool": tool_call_id = message.get("tool_call_id") return [str(tool_call_id)] if tool_call_id else [] if message.get("role") != "user": return [] content = message.get("content") if not isinstance(content, Sequence) or isinstance( content, (str, bytes, bytearray), ): return [] ids: list[str] = [] for block in content: if isinstance(block, Mapping) and block.get("type") == "tool_result": tool_use_id = block.get("tool_use_id") if tool_use_id: ids.append(str(tool_use_id)) return ids def _get_tool_use_ids(message: Mapping[str, Any]) -> list[str]: ids: list[str] = [] tool_calls = message.get("tool_calls") if isinstance(tool_calls, Sequence) and not isinstance( tool_calls, (str, bytes, bytearray), ): for call in tool_calls: if isinstance(call, Mapping): tool_id = call.get("id") if tool_id: ids.append(str(tool_id)) content = message.get("content") if isinstance(content, Sequence) and not isinstance( content, (str, bytes, bytearray), ): for block in content: if isinstance(block, Mapping) and block.get("type") == "tool_use": block_id = block.get("id") if block_id: ids.append(str(block_id)) return ids def _has_tool_use_with_ids( message: Mapping[str, Any], tool_use_ids: set[str], ) -> bool: return any(tool_id in tool_use_ids for tool_id in _get_tool_use_ids(message)) def _assistant_message_identity(message: Mapping[str, Any]) -> str | None: if message.get("role") != "assistant": return None value = message.get("id") if isinstance(value, str) and value: return value meta = message.get("_meta") if isinstance(meta, Mapping): for key in ("response_id", "id", "message_id"): value = meta.get(key) if isinstance(value, str) and value: return value return None def adjust_index_to_preserve_api_invariants( messages: Sequence[Mapping[str, Any]], start_index: int, ) -> int: """Avoid slicing through tool_use/tool_result pairs or split thinking.""" if start_index <= 0 or start_index >= len(messages): return start_index adjusted_index = start_index all_tool_result_ids: list[str] = [] for message in messages[start_index:]: all_tool_result_ids.extend(_get_tool_result_ids(message)) if all_tool_result_ids: tool_use_ids_in_kept = set[str]() for message in messages[adjusted_index:]: tool_use_ids_in_kept.update(_get_tool_use_ids(message)) needed = { tool_id for tool_id in all_tool_result_ids if tool_id not in tool_use_ids_in_kept } for i in range(adjusted_index - 1, -1, -1): if not needed: break message = messages[i] if _has_tool_use_with_ids(message, needed): adjusted_index = i needed.difference_update(_get_tool_use_ids(message)) identities = { identity for identity in ( _assistant_message_identity(message) for message in messages[adjusted_index:] ) if identity } for i in range(adjusted_index - 1, -1, -1): identity = _assistant_message_identity(messages[i]) if identity and identity in identities: adjusted_index = i return adjusted_index def calculate_messages_to_keep_index( messages: Sequence[Mapping[str, Any]], last_summarized_index: int, ) -> int: if not messages: return 0 from openspace.services.conversation.compact import estimate_message_tokens config = _config_from_env() start_index = ( last_summarized_index + 1 if last_summarized_index >= 0 else len(messages) ) total_tokens = 0 text_block_count = 0 for message in messages[start_index:]: total_tokens += estimate_message_tokens([message]) if has_text_blocks(message): text_block_count += 1 if total_tokens >= config.max_tokens: return adjust_index_to_preserve_api_invariants(messages, start_index) if ( total_tokens >= config.min_tokens and text_block_count >= config.min_text_block_messages ): return adjust_index_to_preserve_api_invariants(messages, start_index) boundary_index = -1 for index in range(len(messages) - 1, -1, -1): if is_compact_boundary_message(messages[index]): boundary_index = index break floor = boundary_index + 1 if boundary_index >= 0 else 0 for i in range(start_index - 1, floor - 1, -1): message = messages[i] total_tokens += estimate_message_tokens([message]) if has_text_blocks(message): text_block_count += 1 start_index = i if total_tokens >= config.max_tokens: break if ( total_tokens >= config.min_tokens and text_block_count >= config.min_text_block_messages ): break return adjust_index_to_preserve_api_invariants(messages, start_index) def reset_last_summarized_after_compaction(context: "ToolUseContext | None") -> None: if context is not None: set_last_summarized_message_id(context, None) async def try_session_memory_compaction( messages: list[dict[str, Any]], context: "ToolUseContext | None" = None, *, auto_compact_threshold: int | None = None, hook_registry: "HookRegistry | None" = None, model: str | None = None, ) -> "CompactionResult | None": if not should_use_session_memory_compaction(): return None await wait_for_session_memory_extraction(context) session_id = getattr(context, "session_id", None) if context is not None else None cwd = getattr(context, "cwd", None) if context is not None else None session_dir = getattr(context, "session_dir", None) if context is not None else None if session_dir: memory_path = Path(session_dir).expanduser().resolve() / "session_memory.md" try: session_memory = memory_path.read_text(encoding="utf-8") except FileNotFoundError: session_memory = None except OSError: logger.debug("Failed to read session memory for compact", exc_info=True) return None else: session_memory = read_session_memory(session_id, cwd=cwd) memory_path = get_session_memory_path(session_id, cwd=cwd) if not session_memory: await _emit_sm_compact_event( context, "session_memory_compact_skipped", {"reason": "no_session_memory"}, ) return None if is_session_memory_empty(session_memory): await _emit_sm_compact_event( context, "session_memory_compact_skipped", {"reason": "empty_template"}, ) return None last_summarized = ( get_last_summarized_message_id(context) if context is not None else None ) if last_summarized is not None: last_summarized_index = find_message_cursor_index(messages, last_summarized) if last_summarized_index == -1: await _emit_sm_compact_event( context, "session_memory_compact_skipped", {"reason": "summarized_id_not_found"}, ) return None else: last_summarized_index = len(messages) - 1 await _emit_sm_compact_event( context, "session_memory_compact_resumed_session", {}, ) try: start_index = calculate_messages_to_keep_index( messages, last_summarized_index, ) messages_to_keep = [ dict(message) for message in messages[start_index:] if not is_compact_boundary_message(message) ] result = await _create_compaction_result_from_session_memory( messages=messages, session_memory=session_memory, messages_to_keep=messages_to_keep, memory_path=memory_path, context=context, hook_registry=hook_registry, model=model, ) from openspace.services.conversation.compact import ( build_post_compact_messages, estimate_message_tokens, ) post_messages = build_post_compact_messages(result) post_token_count = estimate_message_tokens(post_messages, model=model) if ( auto_compact_threshold is not None and post_token_count >= auto_compact_threshold ): await _emit_sm_compact_event( context, "session_memory_compact_skipped", { "reason": "threshold_exceeded", "post_compact_token_count": post_token_count, "auto_compact_threshold": auto_compact_threshold, }, ) return None segment_data = await _write_session_memory_compact_segment( context, messages[:start_index], ) compact_was_truncated = bool(result.compact_was_truncated) await _record_session_memory_compact_summary( context, result, segment_data=segment_data, memory_path=memory_path, was_truncated=compact_was_truncated, ) result.post_compact_token_count = post_token_count result.true_post_compact_token_count = post_token_count result.compact_source = "session_memory_compact" result.compact_memory_path = str(memory_path) result.compact_was_truncated = compact_was_truncated await _emit_sm_compact_event( context, "session_memory_compact", { "memory_path": str(memory_path), "messages_kept": len(messages_to_keep), "start_index": start_index, "post_compact_token_count": post_token_count, }, ) reset_last_summarized_after_compaction(context) return result except Exception as exc: logger.debug("Session memory compaction failed: %s", exc, exc_info=True) await _emit_sm_compact_event( context, "session_memory_compact_skipped", {"reason": "error", "error": str(exc)}, ) return None async def _create_compaction_result_from_session_memory( *, messages: list[dict[str, Any]], session_memory: str, messages_to_keep: list[dict[str, Any]], memory_path: Path, context: "ToolUseContext | None", hook_registry: "HookRegistry | None", model: str | None, ) -> "CompactionResult": from openspace.services.conversation.compact import ( CompactionResult, estimate_message_tokens, get_compact_user_summary_message, token_count_from_last_api_response, ) from openspace.services.conversation.attachments import create_post_compact_attachments from openspace.services.tooling.hooks import run_post_compact_hooks pre_compact_token_count = token_count_from_last_api_response(messages) pre_compact_discovered = extract_discovered_tool_names(messages) if context is not None: pre_compact_discovered.update( getattr(context, "discovered_tool_names", set()) or set() ) boundary = build_compact_boundary_message( "auto", pre_compact_token_count, pre_compact_discovered_tools=sorted(pre_compact_discovered), ) truncated, was_truncated = truncate_session_memory_for_compact(session_memory) transcript_path = _transcript_path_for_context(context) summary_content = get_compact_user_summary_message( truncated, True, str(transcript_path) if transcript_path is not None else None, True, ) if was_truncated: summary_content += ( "\n\nSome session memory sections were truncated for length. " f"The full session memory can be viewed at: {memory_path}" ) summary_msg = build_compact_summary_message( summary_content, visible_in_transcript_only=True, ) attachments: list[dict[str, Any]] = [] if context is not None: attachments = await create_post_compact_attachments( context, effective_model=model, messages_to_keep=messages_to_keep, full_compact=False, ) post_hook = await run_post_compact_hooks( hook_registry, {"trigger": "auto", "compact_summary": truncated, "source": "session_memory"}, context, ) post_tokens = estimate_message_tokens([summary_msg, *messages_to_keep, *attachments], model=model) return CompactionResult( boundary_marker=boundary, summary_messages=[summary_msg], attachments=attachments, hook_results=[], messages_to_keep=messages_to_keep if messages_to_keep else None, user_display_message=post_hook.user_display_message, pre_compact_token_count=pre_compact_token_count, post_compact_token_count=post_tokens, true_post_compact_token_count=post_tokens, compact_source="session_memory_compact", compact_memory_path=str(memory_path), compact_was_truncated=bool(was_truncated), ) def _transcript_path_for_context(context: "ToolUseContext | None") -> Path | None: if context is None: return None session_dir = getattr(context, "session_dir", None) if session_dir: candidate = Path(session_dir).expanduser().resolve() return candidate.with_suffix(".messages") if candidate.suffix else candidate.parent / f"{candidate.name}.messages" session_id = getattr(context, "session_id", None) if not session_id: return None from openspace.services.memory.session_memory import get_session_transcript_path return get_session_transcript_path(session_id, cwd=getattr(context, "cwd", None)) async def _write_session_memory_compact_segment( context: "ToolUseContext | None", messages: Sequence[Mapping[str, Any]], ) -> dict[str, Any] | None: if context is None or not messages: return None storage = getattr(context, "session_storage", None) write_segment = getattr(storage, "write_session_transcript_segment", None) if write_segment is None: return None try: result = write_segment( messages, reason="session_memory_compact", task_id=getattr(context, "task_id", None), parent_task_id=getattr(context, "parent_task_id", None), agent_id=getattr(context, "agent_id", None), ) if hasattr(result, "__await__"): result = await result return dict(result) if isinstance(result, Mapping) else None except Exception: logger.debug("Failed to write session-memory compact segment", exc_info=True) return None async def _record_session_memory_compact_summary( context: "ToolUseContext | None", result: "CompactionResult", *, segment_data: Mapping[str, Any] | None, memory_path: Path, was_truncated: bool, ) -> None: if context is None: return storage = getattr(context, "session_storage", None) record = getattr(storage, "record_compact_summary", None) if record is None: return summary_message = result.summary_messages[0] if result.summary_messages else {} segment_ref_id = None if isinstance(segment_data, Mapping): segment_id = segment_data.get("segment_id") if segment_id: segment_ref_id = f"transcript_segment:{storage.session_id}:{segment_id}" try: record_result = record( summary_message_uuid=get_message_uuid(summary_message), compact_source="session_memory_compact", segment_ref_id=segment_ref_id, memory_path=str(memory_path), was_truncated=was_truncated, task_id=getattr(context, "task_id", None), parent_task_id=getattr(context, "parent_task_id", None), agent_id=getattr(context, "agent_id", None), ) if hasattr(record_result, "__await__"): await record_result except Exception: logger.debug("Failed to write session-memory compact summary ref", exc_info=True) async def _emit_sm_compact_event( context: "ToolUseContext | None", event_type: str, data: dict[str, Any], ) -> None: if context is None: return try: await context.emit_event(event_type, data) except Exception: pass