atlas-iac/services/hermes/scripts/patch_api_server_sessions.py

#!/usr/bin/env python3
"""Add explicit parent lineage to Hermes API-created sessions."""

from __future__ import annotations

import argparse
from pathlib import Path

BEFORE = """        model = body.get("model") or self._model_name
        system_prompt = body.get("system_prompt")
        if system_prompt is not None and not isinstance(system_prompt, str):
            return web.json_response(_openai_error("system_prompt must be a string", code="invalid_system_prompt"), status=400)
        db.create_session(session_id, "api_server", model=str(model) if model else None, system_prompt=system_prompt)
"""
AFTER = """        model = body.get("model") or self._model_name
        system_prompt = body.get("system_prompt")
        if system_prompt is not None and not isinstance(system_prompt, str):
            return web.json_response(_openai_error("system_prompt must be a string", code="invalid_system_prompt"), status=400)

        # API workers are first-class children of the objective that launched
        # them.  Accept a JSON field for normal clients and a header for thin
        # relays that cannot extend their request schema.
        metadata = body.get("metadata")
        metadata_parent = metadata.get("parent_session_id") if isinstance(metadata, dict) else None
        raw_parent = body.get("parent_session_id") or metadata_parent or request.headers.get(
            "X-Hermes-Parent-Session-Id"
        )
        parent_session_id = str(raw_parent).strip() if raw_parent else None
        if parent_session_id:
            if (
                len(parent_session_id) > self._MAX_SESSION_HEADER_LEN
                or re.search(r'[\\r\\n\\x00]', parent_session_id)
                or _is_path_unsafe(parent_session_id)
                or parent_session_id == session_id
            ):
                return web.json_response(_openai_error("Invalid parent session ID", code="invalid_parent_session_id"), status=400)
            if not db.get_session(parent_session_id):
                return web.json_response(_openai_error(f"Parent session not found: {parent_session_id}", code="parent_session_not_found"), status=404)

        db.create_session(
            session_id,
            "api_server",
            model=str(model) if model else None,
            system_prompt=system_prompt,
            parent_session_id=parent_session_id,
        )
"""
RUNS_BEFORE = """        run_id = f"run_{uuid.uuid4().hex}"
        session_id = body.get("session_id") or stored_session_id or run_id
        # Approval queues gate host-side tool execution and must be isolated
"""
RESPONSES_SESSION_BEFORE = """        # Reuse session from previous_response_id chain so the dashboard
        # groups the entire conversation under one session entry.
        session_id = stored_session_id or str(uuid.uuid4())

        # Per-client model routing for /v1/responses (see model_routes).
"""
RESPONSES_SESSION_AFTER = """        # Reuse session from previous_response_id chain so the dashboard
        # groups the entire conversation under one session entry.
        session_id = stored_session_id or str(uuid.uuid4())

        # A trusted relay may identify a conversational surface so state.db and
        # WebUI do not collapse it into an anonymous Api_Server session. Keep
        # the accepted vocabulary narrow: these headers are presentation and
        # routing metadata, never an authorization boundary.
        # X-Hermes-Conversation-Platform was parsed before history compaction.
        conversation_title = request.headers.get(
            "X-Hermes-Conversation-Title", ""
        ).strip()
        if conversation_platform:
            if conversation_platform != "telegram":
                return web.json_response(
                    _openai_error("Unsupported conversation platform"), status=400
                )
            if (
                not gateway_session_key
                or not gateway_session_key.startswith("telegram")
                or len(conversation_title) > 128
                or re.search(r'[\\r\\n\\x00]', conversation_title)
            ):
                return web.json_response(
                    _openai_error("Invalid conversation metadata"), status=400
                )
            db = self._ensure_session_db()
            if db is None:
                return web.json_response(
                    _openai_error("Session database unavailable"), status=503
                )
            origin = json.dumps({
                "platform": "telegram",
                "session_key": gateway_session_key,
            }, separators=(",", ":"))
            db.create_session(
                session_id,
                "api_server",
                model=str(body.get("model") or self._model_name or ""),
                system_prompt=instructions if isinstance(instructions, str) else None,
                session_key=gateway_session_key,
                chat_type="private",
            )
            db.record_gateway_session_peer(
                session_id,
                source="api_server",
                session_key=gateway_session_key,
                chat_type="private",
                display_name="Telegram",
                origin_json=origin,
            )
            if conversation_title:
                db.set_session_title(session_id, conversation_title)

        # Per-client model routing for /v1/responses (see model_routes).
"""
TRUNCATION_BEFORE = """        # Truncation support
        if body.get("truncation") == "auto" and len(conversation_history) > 100:
            conversation_history = conversation_history[-100:]
"""
TRUNCATION_AFTER = """        # Telegram keeps a durable compact summary plus bounded recent turns.
        conversation_platform = request.headers.get(
            "X-Hermes-Conversation-Platform", ""
        ).strip().lower()
        if conversation_platform == "telegram":
            from gateway.platforms.telegram_continuity import compact_telegram_history
            conversation_history = compact_telegram_history(conversation_history)
        elif body.get("truncation") == "auto" and len(conversation_history) > 100:
            conversation_history = conversation_history[-100:]
"""
HISTORY_STORE_BEFORE = """        # Build output items from the current turn only.  AIAgent returns a
"""
HISTORY_STORE_AFTER = """        if conversation_platform == "telegram":
            full_history = compact_telegram_history(full_history)

        # Build output items from the current turn only.  AIAgent returns a
"""

SSE_STORE_BEFORE = """            self._response_store.put(response_id, {
"""

SSE_STORE_AFTER = """            if gateway_session_key and gateway_session_key.startswith("telegram"):
                from gateway.platforms.telegram_continuity import compact_telegram_history
                conversation_history_snapshot = compact_telegram_history(
                    conversation_history_snapshot
                )
            self._response_store.put(response_id, {
"""

RUNS_AFTER = """        run_id = f"run_{uuid.uuid4().hex}"
        session_id = body.get("session_id") or stored_session_id or run_id

        # Persist API-run lineage before the agent starts. Automated callers
        # may omit a parent, so a deployment can provide a narrowly matched
        # default without grouping ordinary interactive conversations.
        metadata = body.get("metadata")
        metadata_parent = metadata.get("parent_session_id") if isinstance(metadata, dict) else None
        raw_parent = body.get("parent_session_id") or metadata_parent or request.headers.get(
            "X-Hermes-Parent-Session-Id"
        )
        default_parent = os.environ.get("HERMES_API_DEFAULT_PARENT_SESSION_ID", "").strip()
        default_prefixes = tuple(
            item.strip()
            for item in os.environ.get("HERMES_API_DEFAULT_PARENT_MATCH_PREFIXES", "").split("||")
            if item.strip()
        )
        if not raw_parent and default_parent and default_prefixes and user_message.startswith(default_prefixes):
            raw_parent = default_parent
        parent_session_id = str(raw_parent).strip() if raw_parent else None

        from gateway.session import _is_path_unsafe
        if parent_session_id:
            if (
                len(parent_session_id) > self._MAX_SESSION_HEADER_LEN
                or re.search(r'[\\r\\n\\x00]', parent_session_id)
                or _is_path_unsafe(parent_session_id)
                or parent_session_id == session_id
            ):
                return web.json_response(_openai_error("Invalid parent session ID", code="invalid_parent_session_id"), status=400)
            db = self._ensure_session_db()
            if db is None:
                return web.json_response(_openai_error("Session database unavailable", code="session_db_unavailable"), status=503)
            if not db.get_session(parent_session_id):
                return web.json_response(_openai_error(f"Parent session not found: {parent_session_id}", code="parent_session_not_found"), status=404)
            existing_session = db.get_session(session_id)
            if existing_session and existing_session.get("parent_session_id") != parent_session_id:
                return web.json_response(
                    _openai_error(
                        "Session belongs to a different parent objective",
                        code="session_parent_conflict",
                    ),
                    status=409,
                )
            if not existing_session:
                db.create_session(
                    session_id,
                    "api_server",
                    model=str(body.get("model") or self._model_name or ""),
                    system_prompt=instructions if isinstance(instructions, str) else None,
                    parent_session_id=parent_session_id,
                )
                incident = re.search(r"(?:for|Analyze) incident ([^\\s.]+)(?:\\.|\\s)", user_message)
                if incident and parent_session_id == default_parent:
                    parts = incident.group(1).split("/")
                    label = " · ".join(parts[1:3]) if len(parts) >= 3 else incident.group(1)
                    db.set_session_title(session_id, f"Sonar · {label} · {session_id[-8:]}")
            else:
                # A deliberate continuation of the same visible worker should
                # return to the active state until this run reaches a terminal
                # status. Cross-parent reuse was rejected above.
                db.reopen_session(session_id)

        # Approval queues gate host-side tool execution and must be isolated
"""

RUN_CLOSE_BEFORE = """            finally:
                # If the asyncio wrapper is cancelled (for example via
"""

RUN_CLOSE_AFTER = """            finally:
                # Parent-linked API workers are durable dashboard sessions.
                # Close them on every terminal run path so the dashboard can
                # distinguish a long model wait from completed or failed work.
                if parent_session_id:
                    try:
                        terminal_status = self._run_statuses.get(run_id, {}).get(
                            "status", "failed"
                        )
                        activity_status = terminal_status if terminal_status in {
                            "completed",
                            "failed",
                            "cancelled",
                        } else "failed"
                        self._record_run_activity(
                            session_id,
                            f"run.{activity_status}",
                        )
                        db = self._ensure_session_db()
                        if db is not None:
                            db.end_session(session_id, f"api_run_{terminal_status}")
                    except Exception:
                        logger.exception(
                            "[api_server] failed to close session %s for run %s",
                            session_id,
                            run_id,
                        )

                # If the asyncio wrapper is cancelled (for example via
"""

EVENT_CALLBACK_SIGNATURE_BEFORE = """    def _make_run_event_callback(self, run_id: str, loop: "asyncio.AbstractEventLoop"):
"""

EVENT_CALLBACK_SIGNATURE_AFTER = '''    _RUN_ACTIVITY_BYTES = 512_000
    _RUN_ACTIVITY_FILES = 256
    _RUN_ACTIVITY_HEARTBEAT_SECONDS = 15.0

    def _record_run_activity(
        self,
        session_id: str,
        event_type: str,
        *,
        tool_name: str = None,
        preview: str = None,
        is_error: bool = False,
    ) -> None:
        """Append a bounded, display-only event without changing agent history."""
        labels = {
            "run.started": "Worker started",
            "run.completed": "Worker completed",
            "run.failed": "Worker failed",
            "run.cancelled": "Worker cancelled",
            "_thinking": "Hermes is reasoning",
            "reasoning.available": "Hermes finished a reasoning step",
            "tool.started": "Tool started",
            "tool.completed": "Tool failed" if is_error else "Tool finished",
            "subagent.tool": "Nested worker tool",
            "subagent.progress": "Nested worker progress",
            "subagent_progress": "Nested worker progress",
            "subagent.complete": "Nested worker finished",
            "subagent.text": "Nested worker produced an update",
            "subagent.thinking": "Nested worker is reasoning",
        }
        label = labels.get(event_type)
        if not session_id or not label:
            return
        now = time.time()
        if event_type in {
            "_thinking",
            "reasoning.available",
            "subagent.progress",
            "subagent_progress",
            "subagent.text",
            "subagent.thinking",
        }:
            # Streaming providers may emit hundreds of token-level progress
            # callbacks. They carry no displayable detail here; keep one live
            # heartbeat per session interval while retaining every persisted
            # assistant message, tool call, and tool result in the session DB.
            heartbeats = getattr(self, "_run_activity_heartbeats", None)
            if heartbeats is None:
                heartbeats = {}
                self._run_activity_heartbeats = heartbeats
            last_heartbeat = float(heartbeats.get(session_id, 0.0))
            if now - last_heartbeat < self._RUN_ACTIVITY_HEARTBEAT_SECONDS:
                return
            heartbeats[session_id] = now
            if len(heartbeats) > self._RUN_ACTIVITY_FILES * 2:
                cutoff = now - 3_600.0
                self._run_activity_heartbeats = {
                    key: value
                    for key, value in heartbeats.items()
                    if value >= cutoff
                }
        # Tool names are bounded identifiers. All previews are arbitrary model
        # or tool text and can contain prompts, arguments, paths, or credentials,
        # so the activity projection must never persist them.
        detail = tool_name if event_type in {
            "tool.started",
            "tool.completed",
            "subagent.tool",
        } else ""
        try:
            detail = redact_sensitive_text(str(detail)).strip()[:480]
            content = f"{label}: {detail}" if detail else label
            entry = {
                "role": "assistant",
                "content": f"Activity · {content}",
                "timestamp": now,
                "activity_event": event_type,
            }
            root = Path(
                os.environ.get("HERMES_HOME", "~/.hermes")
            ).expanduser() / "run-activity"
            if root.is_symlink():
                return
            root.mkdir(mode=0o700, parents=True, exist_ok=True)
            root.chmod(0o700)
            digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest()
            path = root / f"{digest}.jsonl"
            previous = root / f"{digest}.previous.jsonl"
            if path.is_symlink() or previous.is_symlink():
                return
            new_file = not path.exists()
            if path.exists() and path.stat().st_size >= self._RUN_ACTIVITY_BYTES:
                previous.unlink(missing_ok=True)
                path.replace(previous)
            line = (json.dumps(entry, separators=(",", ":")) + "\\n").encode()
            flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND
            flags |= getattr(os, "O_NOFOLLOW", 0)
            fd = os.open(path, flags, 0o600)
            try:
                os.fchmod(fd, 0o600)
                os.write(fd, line)
            finally:
                os.close(fd)
            if new_file:
                files = sorted(
                    root.glob("*.jsonl"),
                    key=lambda item: item.stat().st_mtime,
                )
                for stale in files[:-self._RUN_ACTIVITY_FILES]:
                    if not stale.is_symlink():
                        stale.unlink(missing_ok=True)
        except Exception:
            logger.debug("Could not record run activity", exc_info=True)

    def _make_run_event_callback(
        self,
        run_id: str,
        loop: "asyncio.AbstractEventLoop",
        *,
        session_id: str,
    ):
'''

EVENT_CALLBACK_BODY_BEFORE = """        def _callback(event_type: str, tool_name: str = None, preview: str = None, args=None, **kwargs):
            ts = time.time()
            if event_type == "tool.started":
"""

EVENT_CALLBACK_BODY_AFTER = """        def _callback(event_type: str, tool_name: str = None, preview: str = None, args=None, **kwargs):
            ts = time.time()
            self._record_run_activity(
                session_id,
                event_type,
                tool_name=tool_name,
                preview=preview,
                is_error=bool(kwargs.get("is_error", False)),
            )
            if event_type == "tool.started":
"""

EVENT_CALLBACK_END_BEFORE = """            # _thinking and subagent_progress are intentionally not forwarded
"""

EVENT_CALLBACK_END_AFTER = """            elif event_type in {
                "subagent.tool",
                "subagent.progress",
                "subagent_progress",
                "subagent.complete",
                "subagent.text",
                "subagent.thinking",
            }:
                safe_preview = (
                    redact_sensitive_text(str(tool_name or ""))[:480]
                    if event_type == "subagent.tool"
                    else ""
                )
                _push({
                    "event": event_type,
                    "run_id": run_id,
                    "timestamp": ts,
                    "tool": tool_name if event_type == "subagent.tool" else None,
                    "preview": safe_preview,
                })
"""

EVENT_CALLBACK_CALL_BEFORE = """        event_cb = self._make_run_event_callback(run_id, loop)
"""

EVENT_CALLBACK_CALL_AFTER = """        event_cb = self._make_run_event_callback(
            run_id,
            loop,
            session_id=session_id,
        )
        self._record_run_activity(session_id, "run.started")
"""

RUN_SWEEP_BEFORE = """                self._run_streams.pop(run_id, None)
                self._run_streams_created.pop(run_id, None)
                self._active_run_agents.pop(run_id, None)
                self._active_run_tasks.pop(run_id, None)
                self._run_approval_sessions.pop(run_id, None)
"""

RUN_SWEEP_AFTER = """                self._run_streams.pop(run_id, None)
                self._run_streams_created.pop(run_id, None)
                # Stream retention and run lifetime are separate. A long run
                # can legitimately outlive its unconsumed SSE queue; keep the
                # control handles so /stop and approval resolution still work.
                terminal_status = self._run_statuses.get(run_id, {}).get("status")
                if terminal_status in {"completed", "failed", "cancelled"}:
                    self._active_run_agents.pop(run_id, None)
                    self._active_run_tasks.pop(run_id, None)
                    self._run_approval_sessions.pop(run_id, None)
"""


def patch(source: Path, destination: Path) -> None:
    """Apply the narrow session-lineage extension and fail on upstream drift."""
    content = source.read_text(encoding="utf-8")
    if BEFORE not in content:
        raise RuntimeError("Hermes API session patch context changed")
    if RUNS_BEFORE not in content:
        raise RuntimeError("Hermes API runs patch context changed")
    if RUN_CLOSE_BEFORE not in content:
        raise RuntimeError("Hermes API run-close patch context changed")
    if RESPONSES_SESSION_BEFORE not in content:
        raise RuntimeError("Hermes Responses session metadata patch context changed")
    for marker, message in (
        (EVENT_CALLBACK_SIGNATURE_BEFORE, "event callback signature"),
        (EVENT_CALLBACK_BODY_BEFORE, "event callback body"),
        (EVENT_CALLBACK_END_BEFORE, "event callback end"),
        (EVENT_CALLBACK_CALL_BEFORE, "event callback call"),
        (RUN_SWEEP_BEFORE, "run stream sweep"),
    ):
        if marker not in content:
            raise RuntimeError(f"Hermes API {message} patch context changed")
    destination.parent.mkdir(parents=True, exist_ok=True)
    content = content.replace(BEFORE, AFTER, 1)
    content = content.replace(RUNS_BEFORE, RUNS_AFTER, 1)
    content = content.replace(RUN_CLOSE_BEFORE, RUN_CLOSE_AFTER, 1)
    content = content.replace(RESPONSES_SESSION_BEFORE, RESPONSES_SESSION_AFTER, 1)
    content = content.replace(TRUNCATION_BEFORE, TRUNCATION_AFTER, 1)
    content = content.replace(HISTORY_STORE_BEFORE, HISTORY_STORE_AFTER, 1)
    content = content.replace(SSE_STORE_BEFORE, SSE_STORE_AFTER, 1)
    content = content.replace(
        EVENT_CALLBACK_SIGNATURE_BEFORE,
        EVENT_CALLBACK_SIGNATURE_AFTER,
        1,
    )
    content = content.replace(EVENT_CALLBACK_BODY_BEFORE, EVENT_CALLBACK_BODY_AFTER, 1)
    content = content.replace(EVENT_CALLBACK_END_BEFORE, EVENT_CALLBACK_END_AFTER, 1)
    content = content.replace(RUN_SWEEP_BEFORE, RUN_SWEEP_AFTER, 1)
    destination.write_text(
        content.replace(EVENT_CALLBACK_CALL_BEFORE, EVENT_CALLBACK_CALL_AFTER, 1),
        encoding="utf-8",
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", type=Path)
    parser.add_argument("destination", type=Path)
    args = parser.parse_args()
    patch(args.source, args.destination)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())