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
'''

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:
                        db = self._ensure_session_db()
                        if db is not None:
                            terminal_status = self._run_statuses.get(run_id, {}).get(
                                "status", "failed"
                            )
                            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
'''


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")
    destination.parent.mkdir(parents=True, exist_ok=True)
    content = content.replace(BEFORE, AFTER, 1)
    content = content.replace(RUNS_BEFORE, RUNS_AFTER, 1)
    destination.write_text(
        content.replace(RUN_CLOSE_BEFORE, RUN_CLOSE_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())