NexQuant/rdagent/app/rl/ui/app.py

"""
RL Post-training Timeline Viewer
Hierarchical view: Session > Loop > Stage > Events

Run:
    streamlit run rdagent/app/rl/ui/app.py
"""

import os
from pathlib import Path

import streamlit as st
from streamlit import session_state as state

from rdagent.app.rl.ui.components import render_session, render_summary
from rdagent.app.rl.ui.config import ALWAYS_VISIBLE_TYPES, OPTIONAL_TYPES
from rdagent.app.rl.ui.data_loader import get_summary, get_valid_sessions, load_session
from rdagent.app.rl.ui.rl_summary import render_job_summary
from rdagent.core.utils import safe_resolve_path

DEFAULT_LOG_BASE = "log/"


def _safe_resolve(user_input: str | None, safe_root: Path) -> Path:
    safe_root = safe_root.expanduser().resolve()
    if not user_input:
        return safe_root
    if "\x00" in user_input:
        raise ValueError("Invalid path: contains null byte")
    try:
        normalized = os.path.normpath(user_input.strip())
        drive, _ = os.path.splitdrive(normalized)
        if drive:
            raise ValueError("Absolute paths with drive letters are not allowed")
        if os.path.isabs(normalized):
            raise ValueError("Absolute paths are not allowed")
        joined = safe_root / normalized
        return safe_resolve_path(joined, safe_root)
    except (OSError, ValueError) as exc:
        raise ValueError(f"Invalid path outside of allowed root: {user_input}") from exc


def get_job_options(base_path: Path, safe_root: Path | None = None) -> list[str]:
    """
    Scan directory and return job options list.

    Security: Validates base_path to prevent path traversal attacks.
    If safe_root is provided, validates against it; otherwise uses CWD.
    """
    options = []
    has_root_tasks = False
    job_dirs = []

    # Security fix: Validate base_path to prevent path traversal
    try:
        base_path_resolved = base_path.expanduser().resolve()  # nosec B614 — validated against safe_root below via relative_to()

        if safe_root is not None:
            safe_root_resolved = safe_root.expanduser().resolve()
            # Reconstruct from trusted root to break taint chain.
            base_path_resolved = safe_root_resolved / base_path_resolved.relative_to(safe_root_resolved)
        else:
            cwd_resolved = Path.cwd().resolve()
            base_path_resolved = cwd_resolved / base_path_resolved.relative_to(cwd_resolved)
    except (OSError, ValueError, RuntimeError):
        # Path is outside allowed root, reject it
        return options

    if not base_path_resolved.exists():
        return options

    for d in base_path_resolved.iterdir():
        if not d.is_dir():
            continue
        if (d / "__session__").exists():
            has_root_tasks = True
        else:
            try:
                if any((sub / "__session__").exists() for sub in d.iterdir() if sub.is_dir()):
                    job_dirs.append(d.name)
            except PermissionError:
                pass

    job_dirs.sort(reverse=True)
    options.extend(job_dirs)
    if has_root_tasks:
        options.append(". (Current)")

    return options


def main():
    st.set_page_config(layout="wide", page_title="RL Timeline", page_icon="🤖")

    with st.sidebar:
        view_mode = st.radio("View Mode", ["Job Summary", "Single Task"], horizontal=True)
        st.divider()

        default_log = os.environ.get("RL_LOG_PATH", DEFAULT_LOG_BASE)
        safe_root = Path(default_log).expanduser().resolve()
        job_folder = str(safe_root)
        selected_types = ALWAYS_VISIBLE_TYPES.copy()
        is_root_job = False

        if view_mode == "Job Summary":
            st.header("Job")
            base_folder = st.text_input("Base Folder", value=default_log, key="base_folder_input")
            try:
                base_path = _safe_resolve(base_folder, safe_root)
            except ValueError as e:
                st.error(str(e))
                return

            job_options = get_job_options(base_path, safe_root)  # nosec B614 – validated by _safe_resolve
            if job_options:
                selected_job = st.selectbox("Select Job", job_options, key="job_select")
                if selected_job.startswith("."):
                    job_folder = str(base_path)
                    is_root_job = True
                else:
                    job_folder = str(base_path / selected_job)
                state.selected_job_folder = job_folder
            else:
                st.warning("No jobs found in this directory")
                job_folder = str(base_path)

            if st.button("Refresh", type="primary", key="refresh_job"):
                st.rerun()
        else:
            st.header("Session")
            default_path = getattr(state, "selected_job_folder", default_log)
            log_folder = st.text_input("Log Folder", value=default_path)
            try:
                log_path = _safe_resolve(log_folder, safe_root)
            except ValueError as e:
                st.error(str(e))
                return

            sessions = get_valid_sessions(log_path)
            if not sessions:
                st.warning("No valid sessions found")
                return

            selected_session = st.selectbox("Session", sessions)

            if st.button("Load", type="primary") or "session" not in state:
                with st.spinner("Loading..."):
                    state.session = load_session(log_path / selected_session)
                    state.session_name = selected_session

            st.divider()

            st.subheader("Show More")
            selected_types = ALWAYS_VISIBLE_TYPES.copy()
            for event_type, (label, default) in OPTIONAL_TYPES.items():
                if st.toggle(label, value=default, key=f"toggle_{event_type}"):
                    selected_types.append(event_type)

            st.divider()

            if "session" in state:
                summary = get_summary(state.session)
                st.subheader("Summary")
                st.metric("Loops", summary.get("loop_count", 0))
                st.metric("LLM Calls", summary.get("llm_call_count", 0))
                success = summary.get("docker_success", 0)
                fail = summary.get("docker_fail", 0)
                st.metric("Docker", f"{success}✓ / {fail}✗")

    if view_mode == "Job Summary":
        st.title("📊 RL Job Summary")
        try:
            job_path = _safe_resolve(job_folder, safe_root)
        except ValueError as e:
            st.warning(str(e))
            return
        if job_path.exists():
            render_job_summary(job_path, safe_root, is_root=is_root_job)
        else:
            st.warning(f"Job folder not found: {job_folder}")
        return

    st.title("🤖 RL Timeline Viewer")

    if "session" not in state:
        st.info("Select a session and click **Load** to view")
        return

    session = state.session
    summary = get_summary(session)
    render_summary(summary)
    st.divider()
    render_session(session, selected_types)


if __name__ == "__main__":
    main()