tracer-cloud--opensre/surfaces/interactive_shell/command_registry/session_cmds/resume.py

"""Session resume slash command and helpers: /resume."""

from __future__ import annotations

from rich.console import Console
from rich.markup import escape

from core.agent_harness.session import SessionManager
from surfaces.interactive_shell.command_registry.session_cmds.resume_rendering import (
    render_resumed_session_history,
)
from surfaces.interactive_shell.runtime import Session
from surfaces.interactive_shell.ui import DIM, ERROR, HIGHLIGHT, WARNING
from surfaces.interactive_shell.ui.components.choice_menu import (
    repl_choose_one,
    repl_tty_interactive,
)
from surfaces.interactive_shell.ui.components.time_format import format_repl_timestamp


def _record_resume_slash(
    session: Session,
    args: list[str],
    *,
    ok: bool = True,
    picked_id: str | None = None,
) -> None:
    """Record /resume in the active session file after identity is settled."""
    if picked_id:
        text = f"/resume {picked_id[:8]}"
    elif args:
        text = f"/resume {' '.join(args)}"
    else:
        text = "/resume"
    session.record("slash", text, ok=ok)


def _interactive_resume_menu(session: Session, console: Console) -> bool:
    """Show a numbered list of recent sessions and resume the selected one."""
    from core.agent_harness.session import default_session_repo

    entries = [
        e for e in default_session_repo().load_recent(10) if e["session_id"] != session.session_id
    ]
    if not entries:
        console.print(f"[{DIM}]No previous sessions to resume.[/]")
        return True

    choices: list[tuple[str, str]] = []
    for entry in entries:
        sid = entry["session_id"]
        short_id = sid[:8]
        name = entry.get("name") or f"[{short_id}]"
        started_str = format_repl_timestamp(entry.get("started_at"), style="compact")
        label = f"{name[:40]:<40}  {short_id}  {started_str}"
        choices.append((sid, label))
    choices.append(("done", "done"))

    picked = repl_choose_one(title="resume session", breadcrumb="/resume", choices=choices)
    if picked is None or picked == "done":
        return True

    slash_command = f"/resume {picked[:8]}"
    if not _do_resume(picked, session, console, slash_command=slash_command):
        _record_resume_slash(session, [], picked_id=picked, ok=False)
    return True


def _apply_resume_data(
    data: dict,
    session: Session,
    console: Console,
    *,
    slash_command: str | None = None,
) -> bool:
    """Apply loaded session data into the running session and print a summary."""
    messages = data.get("cli_agent_messages") or []
    context = data.get("accumulated_context") or {}
    history = data.get("history") or []
    has_snapshot = data.get("has_snapshot", False)
    sid = data.get("session_id", "")
    short_id = sid[:8] if len(sid) >= 8 else sid
    name = data.get("name") or ""

    if not messages and not context:
        console.print(
            f"[{DIM}]session {short_id} has no conversation to resume "
            "(no chat turns or context found).[/]"
        )
        if not data.get("turn_details") and not has_snapshot:
            console.print(
                f"[{DIM}]tip: turn_detail records are only written when prompt logging is enabled.[/]"
            )
        if slash_command:
            session.record("slash", slash_command, ok=False)
        return True

    existing = session.agent.messages
    if existing:
        console.print(
            f"[{WARNING}]current session has {len(existing)} messages — "
            "they will be replaced by the resumed context.[/]"
        )

    manager = SessionManager.for_session(session)
    manager.rebind_for_resume(
        session,
        session_id=sid,
        started_at=data.get("started_at"),
    )
    manager.restore_context(session, data)

    source = "snapshot" if has_snapshot else "turn records"
    name_str = f" · {escape(name)}" if name else ""
    console.print(
        f"[{HIGHLIGHT}]resumed session {short_id}{name_str}[/] "
        f"[{DIM}]({len(messages)} messages in context from {source})[/]"
    )

    render_resumed_session_history(
        console,
        history=history,
        turn_details=data.get("turn_details") or [],
        messages=list(messages),
    )

    if context:
        console.print(
            f"[{DIM}]accumulated context restored:[/] "
            + ", ".join(f"{escape(k)}={escape(str(v))}" for k, v in sorted(context.items()))
        )

    if slash_command:
        session.record("slash", slash_command)

    return True


def _lookup_resume_session_data(
    prefix: str,
    session: Session,
    console: Console,
) -> dict | None:
    """Resolve a session to resume by ID prefix or name substring."""
    from core.agent_harness.session import default_session_repo

    repo = default_session_repo()
    data = repo.load_session(prefix)
    if data is None and len(prefix) >= 3:
        candidates = [
            e
            for e in repo.load_recent(20)
            if prefix.lower() in (e.get("name") or "").lower()
            and e["session_id"] != session.session_id
        ]
        if len(candidates) == 1:
            data = repo.load_session(candidates[0]["session_id"])
        elif len(candidates) > 1:
            console.print(
                f"[{WARNING}]'{escape(prefix)}' matches {len(candidates)} sessions by name — "
                "use a session ID prefix or be more specific.[/]"
            )
            return None

    if data is not None:
        return data

    n = repo.count_prefix_matches(prefix)
    if n > 1:
        console.print(
            f"[{WARNING}]ambiguous prefix '{escape(prefix)}' matches {n} sessions — "
            "use more characters.[/]"
        )
    else:
        console.print(f"[{ERROR}]session '{escape(prefix)}' not found.[/]")
    return None


def _do_resume(
    prefix: str,
    session: Session,
    console: Console,
    *,
    slash_command: str | None = None,
) -> bool:
    """Load session by ID prefix and restore context into the running session."""
    data = _lookup_resume_session_data(prefix, session, console)
    if data is None:
        return False
    return _apply_resume_data(data, session, console, slash_command=slash_command)


def resume_session_by_prefix(
    prefix: str,
    session: Session,
    console: Console,
    *,
    slash_command: str | None = None,
) -> bool:
    """Load session by ID prefix and restore context into the running session."""
    return _do_resume(prefix, session, console, slash_command=slash_command)


def _cmd_resume(session: Session, console: Console, args: list[str]) -> bool:
    if not args and repl_tty_interactive():
        return _interactive_resume_menu(session, console)

    if not args:
        console.print(f"[{DIM}]usage: /resume <session-id-prefix>[/]")
        console.print(f"[{DIM}]run /sessions to list session IDs.[/]")
        _record_resume_slash(session, args)
        return True

    prefix = args[0].strip()
    session_prefix = prefix.split(":", 1)[0]

    if session.session_id.startswith(session_prefix) and ":" not in prefix:
        console.print(
            f"[{DIM}]session {session_prefix[:8]} is the current session — "
            "run /sessions to pick a previous one.[/]"
        )
        _record_resume_slash(session, args)
        return True

    data = _lookup_resume_session_data(prefix, session, console)
    if data is None:
        _record_resume_slash(session, args, ok=False)
        return True

    slash_command = f"/resume {' '.join(args)}" if args else "/resume"
    _apply_resume_data(data, session, console, slash_command=slash_command)
    return True