tracer-cloud--opensre/core/agent_harness/session/lifecycle.py

"""Centralized session lifecycle owner for every surface.

``SessionManager`` is the single component that creates, resolves, rotates,
restores, and flushes :class:`SessionCore` objects (surfaces subclass it).
Surfaces (interactive
shell, gateway, headless) delegate session lifecycle to it instead of each
re-implementing bootstrap + persistence wiring:

- **create** — a fresh session: construct, run the core bootstrap (persistent
  tasks + integration hydration/warm), and open its storage stream.
- **resolve** — load a persisted session by id: construct with that id, run the
  core bootstrap, restore its saved conversation context, and reopen storage.
- **rotate** — close the outgoing session and create a fresh replacement (new
  handle; used by the gateway).
- **open_storage** — open the JSONL stream for an already-bootstrapped handle
  (interactive REPL entry after ``SessionBootstrapSpec``).
- **rotate_in_place** / **rebind_for_resume** — flush + reset the *live* handle
  the REPL already holds (``/new`` / ``/resume``), preserving loop-owned UI
  state instead of releasing it.
- **restore_context** — rehydrate messages / accumulated context / history from
  a persisted session dict.
- **close** — terminal teardown of a discarded handle: flush + release
  resources (cancel warm task, drop background references).

Surface-specific concerns stay with the surface: the shell layers terminal UI
state (theme, grounding providers, prompt history) on top of a manager-created
session; the gateway injects per-chat metadata. Neither re-implements the core
bootstrap, and neither reaches across surfaces to do it.
"""

from __future__ import annotations

import contextlib
import logging
from datetime import datetime
from typing import Any, TypeVar

from core.agent_harness.session.persistence.ports import SessionRepo, SessionStorage

# Import from submodules (not the package __init__) so the session package can
# re-export SessionManager without a circular import.
from core.agent_harness.session.session_core import SessionCore
from platform.common.task_registry import TaskRegistry
from platform.observability.trace.spans import component_span

logger = logging.getLogger(__name__)

# In-place lifecycle methods return the caller's own session type (a surface's
# ``Session`` subclass or a plain ``SessionCore``), so they preserve it.
_S = TypeVar("_S", bound=SessionCore)


class SessionManager:
    """Owns the create / resolve / rotate / restore / flush session lifecycle.

    Storage and repo backends are injectable so tests can run against in-memory
    persistence; production surfaces use the shared JSONL singletons (resolved
    lazily to avoid importing the package ``__init__`` from within it).
    """

    def __init__(
        self,
        *,
        storage: SessionStorage | None = None,
        repo: SessionRepo | None = None,
    ) -> None:
        if storage is None or repo is None:
            from core.agent_harness.session import (
                DEFAULT_SESSION_REPO,
                DEFAULT_SESSION_STORAGE,
            )

            storage = storage or DEFAULT_SESSION_STORAGE
            repo = repo or DEFAULT_SESSION_REPO
        self._storage = storage
        self._repo = repo

    @classmethod
    def for_session(cls, session: SessionCore) -> SessionManager:
        """Build a manager bound to a live session's own storage backend.

        The single named construction point for the in-place lifecycle calls
        (``/new`` / ``/resume``) so they all bind to ``session.storage``
        consistently instead of re-passing it at each call site.
        """
        return cls(storage=session.storage)

    # ─── Core bootstrap ──────────────────────────────────────────────────

    def bootstrap(
        self,
        session: _S,
        *,
        hydrate_integrations: bool = True,
        warm_integrations: bool = False,
        persistent_tasks: bool = True,
    ) -> _S:
        """Apply the surface-agnostic startup mutations to ``session``.

        This is the single definition of "a booted session": a persistent task
        registry and hydrated (optionally warmed) integration state. Surface UI
        wiring is layered by the surface after this returns.
        """
        if persistent_tasks:
            session.task_registry = TaskRegistry.persistent()
        # Safe read-only facts (version/env) so agents never need subprocess introspection.
        with component_span("runtime_metadata:bootstrap", session_id=session.session_id):
            session.refresh_runtime_metadata()
        if hydrate_integrations:
            session.hydrate_configured_integrations()
        if warm_integrations:
            session.warm_resolved_integrations()
        return session

    # ─── Lifecycle ───────────────────────────────────────────────────────

    def create(
        self,
        *,
        session_id: str | None = None,
        hydrate_integrations: bool = True,
        warm_integrations: bool = False,
        persistent_tasks: bool = True,
        open_storage: bool = True,
    ) -> SessionCore:
        """Build a fresh session, bootstrap it, and open its storage stream."""
        session = SessionCore(session_id=session_id) if session_id else SessionCore()
        # Align the session's own persistence backend with the manager's, so
        # session.record()/append go through the same storage the manager opens
        # and flushes. Otherwise an injected backend is bypassed by the default
        # JSONL field on SessionCore.
        session.storage = self._storage
        self.bootstrap(
            session,
            hydrate_integrations=hydrate_integrations,
            warm_integrations=warm_integrations,
            persistent_tasks=persistent_tasks,
        )
        if open_storage:
            self.open_storage(session)
        return session

    def open_storage(self, session: _S) -> _S:
        """Open the JSONL stream for an already-bootstrapped session handle.

        The interactive shell bootstraps via ``SessionBootstrapSpec`` first, then
        calls this once it knows the run is an interactive REPL (not a one-shot
        ``initial_input`` path).
        """
        session.storage = self._storage
        self._storage.open_session(session)
        return session

    def resolve(
        self,
        session_id: str,
        *,
        hydrate_integrations: bool = True,
        warm_integrations: bool = True,
        persistent_tasks: bool = True,
    ) -> SessionCore:
        """Load a persisted session by id: bootstrap, restore context, reopen storage."""
        session = self.create(
            session_id=session_id,
            hydrate_integrations=hydrate_integrations,
            warm_integrations=warm_integrations,
            persistent_tasks=persistent_tasks,
            open_storage=False,
        )
        data = self._repo.load_session(session_id)
        self.restore_context(session, data)
        self._storage.reopen_session(session.session_id)
        return session

    def rotate(
        self,
        *,
        old_session_id: str | None = None,
        new_session_id: str | None = None,
        warm_integrations: bool = True,
    ) -> SessionCore:
        """Close the outgoing session (if any) and create its replacement."""
        if old_session_id:
            outgoing = SessionCore(session_id=old_session_id)
            # Reconstructed handle: align its backend with the manager's so the
            # close flush lands on the same storage the manager owns.
            outgoing.storage = self._storage
            self.close(outgoing)
        return self.create(session_id=new_session_id, warm_integrations=warm_integrations)

    def rotate_in_place(self, session: _S) -> _S:
        """Flush the outgoing session file, reset state, and open a new session id.

        Mutates the live ``session`` handle the REPL already holds (``/new``).
        The caller restores any conversation context to carry forward
        (``agent.messages``, ``accumulated_context``, ``resumed_from_name``)
        after this returns.

        Flushes but does not release resources: ``clear()`` resets in-memory
        state (and cancels the warm task) while the loop-owned
        ``prompt_refresh_fn`` is preserved for the continuing REPL.
        """
        session.storage = self._storage
        self._flush(session)
        session.clear()
        self._storage.open_session(session)
        return session

    def rebind_for_resume(
        self,
        session: _S,
        *,
        session_id: str,
        started_at: Any | None = None,
    ) -> _S:
        """Point the live session handle at a persisted id before :meth:`restore_context`.

        Used by the interactive shell ``/resume`` command on the in-process
        session object. When ``session_id`` differs from the current id the
        outgoing file is flushed; otherwise only in-memory state is cleared
        without rotating identity. Either way the live handle is reused, so
        loop-owned ``prompt_refresh_fn`` is preserved (flush, not close).
        """
        session.storage = self._storage
        if session.session_id != session_id:
            self._flush(session)
            session.clear(rotate_identity=False)
            session.session_id = session_id
            if isinstance(started_at, str) and started_at:
                with contextlib.suppress(Exception):
                    session.started_at = datetime.fromisoformat(started_at).timestamp()
            self._storage.reopen_session(session_id)
        else:
            session.clear(rotate_identity=False)
            session.session_id = session_id
        return session

    def restore_context(self, session: _S, data: dict[str, Any] | None) -> _S:
        """Rehydrate conversation messages, accumulated context, and history.

        ``data`` is the persisted session dict from ``SessionRepo.load_session``;
        a ``None`` or empty dict leaves the session untouched.
        """
        if not data:
            return session
        messages = data.get("cli_agent_messages")
        if isinstance(messages, list):
            restored: list[tuple[str, str]] = []
            for item in messages:
                try:
                    role, content = item
                except (TypeError, ValueError):
                    continue
                if role in {"user", "assistant"} and isinstance(content, str) and content:
                    restored.append((role, content))
            session.cli_agent_messages = restored
        context = data.get("accumulated_context")
        if isinstance(context, dict):
            session.accumulated_context = dict(context)
        history = data.get("history")
        if isinstance(history, list):
            session.history = [dict(item) for item in history if isinstance(item, dict)]
        return session

    def close(self, session: SessionCore) -> None:
        """Finalize a session for good: persist buffered state and release resources.

        This is the terminal teardown hook — the session handle is being
        discarded (end of a REPL run, or ``rotate``'s outgoing session). It is
        NOT for the in-place swaps (``/new`` / ``/resume``) which reuse the live
        handle; those call :meth:`rotate_in_place` / :meth:`rebind_for_resume`,
        which flush without releasing loop-owned UI state.

        Persisting is best-effort (a failed flush must not crash teardown);
        the session releases its own resources (:meth:`SessionCore.release_resources`)
        to prevent per-session leaks.
        """
        self._flush(session)
        from platform.observability.trace.spans import emit_thread_boundary

        emit_thread_boundary(session.session_id, name="session_end", phase="session_end")
        session.release_resources()

    @staticmethod
    def _flush(session: SessionCore) -> None:
        """Best-effort persist through the session's own backend.

        Flushes through ``session.storage`` — the backend it recorded turns
        through — so the end-of-session marker lands with the data. A failed
        flush is logged, never raised, so teardown/rotation cannot crash.
        """
        try:
            session.storage.flush(session)
        except OSError:
            logger.debug("[session] flush failed", exc_info=True)


__all__ = ["SessionManager"]