Support dataset session paths

app.py

@@ -5,6 +5,11 @@ os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
 os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
 os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
 os.environ.setdefault("GRADIO_SSR_MODE", "false")
+if not os.environ.get("HF_TOKEN"):
+    for _token_key in ("HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_HUB_TOKEN"):
+        if os.environ.get(_token_key):
+            os.environ["HF_TOKEN"] = os.environ[_token_key]
+            break
 
 print("codex-traces startup: env configured", flush=True)
 
@@ -27,6 +32,8 @@ APP_TITLE = "Codex Traces"
 DATASET_REPO = "Mike0021/codex-sessions"
 REPO_TYPE = "dataset"
 FILE_SUFFIX = "_rollout.jsonl"
+SESSION_PATH_PREFIX = "sessions/rollout-"
+SESSION_PATH_SUFFIX = ".jsonl"
 PAGE_SIZE = 120
 MAX_OUTPUT_CHARS = 22000
 MAX_MESSAGE_CHARS = 30000
@@ -81,16 +88,30 @@ def _hub_token() -> str | None:
 
 
 def _display_name(file_name: str) -> str:
+    if file_name.startswith(SESSION_PATH_PREFIX) and file_name.endswith(SESSION_PATH_SUFFIX):
+        return file_name[len(SESSION_PATH_PREFIX) : -len(SESSION_PATH_SUFFIX)]
     return file_name[: -len(FILE_SUFFIX)] if file_name.endswith(FILE_SUFFIX) else file_name
 
 
-def _file_name(session_name: str | None) -> str | None:
+def _is_rollout_file(file_name: str) -> bool:
+    return (
+        file_name.startswith(SESSION_PATH_PREFIX)
+        and file_name.endswith(SESSION_PATH_SUFFIX)
+    ) or file_name.endswith(FILE_SUFFIX)
+
+
+def _file_candidates(session_name: str | None) -> list[str]:
     if not session_name:
-        return None
+        return []
     name = str(session_name).strip()
-    if name.endswith(FILE_SUFFIX):
-        return name
-    return f"{name}{FILE_SUFFIX}"
+    if name.endswith(".jsonl"):
+        return [name]
+
+    candidates = [
+        f"{SESSION_PATH_PREFIX}{name}{SESSION_PATH_SUFFIX}",
+        f"{name}{FILE_SUFFIX}",
+    ]
+    return list(dict.fromkeys(candidates))
 
 
 def _fmt_number(value: int | None) -> str:
@@ -167,7 +188,7 @@ def _list_sessions_cached() -> tuple[tuple[str, ...], str]:
         files = list_repo_files(DATASET_REPO, repo_type=REPO_TYPE, token=_hub_token())
     except Exception as exc:
         return (), f"{type(exc).__name__}: {exc}"
-    sessions = tuple(sorted(_display_name(f) for f in files if f.endswith(FILE_SUFFIX)))
+    sessions = tuple(sorted({_display_name(f) for f in files if _is_rollout_file(f)}))
     if not sessions:
         return (), "No *_rollout.jsonl files were found in the dataset."
     return sessions, ""
@@ -613,14 +634,21 @@ def render_session(session_name: str | None, query: str | None = "", page: int |
     session_name = session_name or choices[0]
     if session_name not in choices and _display_name(str(session_name)) in choices:
         session_name = _display_name(str(session_name))
-    file_name = _file_name(session_name)
-    if file_name is None:
+    candidates = _file_candidates(session_name)
+    if not candidates:
         return _render_error("No session selected", "Choose a session from the dropdown.")
 
-    try:
-        data = _load_session(file_name)
-    except Exception as exc:
-        return _render_error("Could not load session", f"{type(exc).__name__}: {exc}")
+    last_error: Exception | None = None
+    data = None
+    for file_name in candidates:
+        try:
+            data = _load_session(file_name)
+            break
+        except Exception as exc:
+            last_error = exc
+    if data is None:
+        details = f"{type(last_error).__name__}: {last_error}" if last_error else "No file candidates were available."
+        return _render_error("Could not load session", details)
 
     query = query or ""
     filtered = _filter_events(data["events"], query)