From c1140caab88b544bd68d31f71551cd876bb1d09e Mon Sep 17 00:00:00 2001 From: Mohit Paddhariya Date: Sun, 14 Jun 2026 18:14:15 +0530 Subject: [PATCH 1/5] feat(session): add session_title with /rename and auto-title generation Sessions now carry a human-readable title so they can be found again, the way ChatGPT and Claude name conversations. - Add a persisted session_title field to Session (in get_trajectory, reset on /new) plus a _title_user_set flag so an explicit rename is never clobbered by auto-titling. - New agent/core/title.py: generate_conversation_title asks the active model for a 3-6 word title (tiny token budget, no telemetry/billing impact) and falls back to a pure-Python slug of the first user message on any error. Long secret-like tokens are stripped before a title reaches disk. - Auto-title fires once, fire-and-forget, after the first completed turn, so a slow or failing title never blocks or breaks a turn. - Add the /rename command and list it in /help. Part of #325. --- agent/core/agent_loop.py | 50 +++++++++ agent/core/session.py | 10 ++ agent/core/title.py | 170 +++++++++++++++++++++++++++++++ agent/main.py | 22 ++++ agent/utils/terminal_display.py | 1 + tests/unit/test_session_title.py | 147 ++++++++++++++++++++++++++ tests/unit/test_title.py | 100 ++++++++++++++++++ 7 files changed, 500 insertions(+) create mode 100644 agent/core/title.py create mode 100644 tests/unit/test_session_title.py create mode 100644 tests/unit/test_title.py diff --git a/agent/core/agent_loop.py b/agent/core/agent_loop.py index 6ed3a4e33..a5ba789a9 100644 --- a/agent/core/agent_loop.py +++ b/agent/core/agent_loop.py @@ -1171,6 +1171,41 @@ async def _call_llm_non_streaming( ) +async def _generate_and_set_title(session: "Session", final_response: str | None) -> None: + """Generate a conversation title and attach it to the session. + + Runs as a fire-and-forget task after the first turn. Any failure is + swallowed so it can never break the turn that spawned it. + """ + try: + from agent.core.title import ( + extract_first_user_text, + generate_conversation_title, + ) + + first_user_text = extract_first_user_text(session.context_manager.items) + if not first_user_text: + return + title = await generate_conversation_title( + session.config.model_name, + session.hf_token, + first_user_text, + final_response, + ) + # The user may have renamed (or a /new fired) while we were awaiting. + if not title or session._title_user_set or session.session_title: + return + session.session_title = title + await session.send_event( + Event( + event_type="conversation_title", + data={"title": title, "session_id": session.session_id}, + ) + ) + except Exception as e: # noqa: BLE001 + logger.debug("Auto-title task failed: %s", e) + + class Handlers: """Handler functions for each operation type""" @@ -1880,6 +1915,21 @@ async def _exec_tool( ) ) + # Auto-title the conversation once, after the very first completed + # turn, unless the user already named it via /rename. Fire-and-forget + # so a slow or failing title never delays the turn. + if ( + session.turn_count == 0 + and not session.session_title + and not session._title_user_set + ): + asyncio.create_task( + _generate_and_set_title( + session, + final_response if isinstance(final_response, str) else None, + ) + ) + # Increment turn counter and check for auto-save session.increment_turn() await session.auto_save_if_needed() diff --git a/agent/core/session.py b/agent/core/session.py index b578a0b51..d8b0c972a 100644 --- a/agent/core/session.py +++ b/agent/core/session.py @@ -117,6 +117,7 @@ def __init__( hf_username: str | None = None, user_plan: str | None = None, persistence_store: Any | None = None, + session_title: str | None = None, ): self.hf_token: Optional[str] = hf_token self.user_id: Optional[str] = user_id @@ -139,6 +140,12 @@ def __init__( ) self.event_queue = event_queue self.session_id = session_id or str(uuid.uuid4()) + # Human-readable conversation title (auto-generated after the first + # turn, or set explicitly via ``/rename``). Surfaces in ``/resume`` and + # the saved filename. ``_title_user_set`` records an explicit rename so + # auto-titling never clobbers a name the user chose. + self.session_title: str | None = session_title + self._title_user_set: bool = session_title is not None self.inference_billing_session_id: str | None = None self.config = config self.is_running = True @@ -458,6 +465,8 @@ def start_new_conversation(self) -> dict[str, Any]: self.context_manager.running_context_usage = 0 self.session_id = str(uuid.uuid4()) + self.session_title = None + self._title_user_set = False self.inference_billing_session_id = None self.session_start_time = datetime.now().astimezone().isoformat() self.turn_count = 0 @@ -552,6 +561,7 @@ def get_trajectory(self) -> dict: usage_metrics = self.usage_metrics or {} return { "session_id": self.session_id, + "session_title": self.session_title, "user_id": self.user_id, "hf_username": self.hf_username, "session_start_time": self.session_start_time, diff --git a/agent/core/title.py b/agent/core/title.py new file mode 100644 index 000000000..2061a4110 --- /dev/null +++ b/agent/core/title.py @@ -0,0 +1,170 @@ +"""Conversation title generation for CLI sessions. + +A session gets a short human-readable title so it can be found again in the +``/resume`` picker and in its on-disk filename, the way ChatGPT and Claude name +conversations. The title is generated once, right after the first turn: + +* ``generate_conversation_title`` asks the active model for a 3-6 word title. +* On any failure (network, billing, empty reply) it falls back to + ``fallback_title`` — a pure-Python summary of the first user message, so a + session always ends up with *some* readable name and the call never blocks or + breaks a turn. + +The call deliberately does NOT pass a ``session`` to telemetry: it uses a tiny +token budget and no reasoning effort so it can't meaningfully move cost or trip +the YOLO/usage-pause logic. +""" + +from __future__ import annotations + +import logging +import re +from typing import Any + +logger = logging.getLogger(__name__) + +# Keep titles short and skimmable in a one-line picker row. +MAX_TITLE_WORDS = 6 +MAX_TITLE_CHARS = 50 + +# Strip anything that looks like a long opaque secret (tokens, keys, base64 +# blobs) before it can land in a persisted title or a filename. The title is +# derived from the first user message, which may contain a pasted credential. +_SECRET_RUN = re.compile(r"\b[A-Za-z0-9_\-]{30,}\b") + + +def _collapse(text: str) -> str: + """Collapse all whitespace to single spaces and trim.""" + return " ".join(str(text).split()) + + +def _strip_secrets(text: str) -> str: + """Drop long credential-like tokens so they never reach disk via a title.""" + return _collapse(_SECRET_RUN.sub("", text)) + + +def _cap(text: str, max_words: int = MAX_TITLE_WORDS, max_chars: int = MAX_TITLE_CHARS) -> str: + words = text.split() + if len(words) > max_words: + text = " ".join(words[:max_words]) + if len(text) > max_chars: + text = text[: max_chars - 1].rstrip() + "…" + return text + + +def fallback_title(first_user_text: str | None) -> str: + """Readable title derived from the first user message, no LLM required. + + Returns an empty string when there's nothing usable to title from. + """ + cleaned = _strip_secrets(first_user_text or "") + if not cleaned: + return "" + return _cap(cleaned) + + +def extract_first_user_text(items: Any) -> str: + """Pull the first user message's text out of a context-manager item list. + + Handles both dict-shaped messages and litellm ``Message`` objects, and + string or block-list ``content``. + """ + for item in items or []: + role = item.get("role") if isinstance(item, dict) else getattr(item, "role", None) + if role != "user": + continue + content = ( + item.get("content") if isinstance(item, dict) else getattr(item, "content", None) + ) + if isinstance(content, str): + text = content + elif isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict): + value = block.get("text") or block.get("content") + if isinstance(value, str): + parts.append(value) + elif isinstance(block, str): + parts.append(block) + text = " ".join(parts) + else: + text = "" + text = _collapse(text) + if text: + return text + return "" + + +def _clean_model_title(raw: str) -> str: + """Normalize a model-produced title: one line, no quotes/trailing punctuation.""" + title = _collapse(raw) + title = title.splitlines()[0] if title else "" + title = title.strip().strip("\"'“”‘’`") + title = title.rstrip(".!?,;: ") + return _cap(_strip_secrets(title)) + + +_SYSTEM_PROMPT = ( + "You generate a very short title summarizing a conversation. " + "Reply with ONLY the title: 3 to 6 words, Title Case, no quotes, " + "no trailing punctuation." +) + + +async def generate_conversation_title( + model_name: str, + hf_token: str | None, + first_user_text: str, + first_assistant_text: str | None = None, +) -> str: + """Ask the active model for a short title for the conversation. + + Falls back to ``fallback_title(first_user_text)`` on any error or empty + result. Never raises. + """ + fallback = fallback_title(first_user_text) + user_text = _collapse(first_user_text or "") + if not user_text: + return fallback + + try: + from litellm import acompletion + + from agent.core.llm_params import _resolve_llm_params + + params = _resolve_llm_params(model_name, hf_token, reasoning_effort=None) + assistant = _collapse(first_assistant_text or "") or "(no reply yet)" + messages = [ + {"role": "system", "content": _SYSTEM_PROMPT}, + { + "role": "user", + "content": ( + f"First user message:\n{user_text[:1000]}\n\n" + f"Assistant reply:\n{assistant[:1000]}\n\n" + "Return ONLY the title." + ), + }, + ] + response = await acompletion( + messages=messages, + max_completion_tokens=24, + stream=False, + **params, + ) + raw = "" + if getattr(response, "choices", None): + raw = response.choices[0].message.content or "" + title = _clean_model_title(raw) + return title or fallback + except Exception as e: # noqa: BLE001 — a bad title must never break a turn + try: + from agent.core.hf_access import is_inference_billing_error + + if is_inference_billing_error(e): + logger.debug("Auto-title skipped (billing): %s", e) + else: + logger.debug("Auto-title generation failed: %s", e) + except Exception: + logger.debug("Auto-title generation failed: %s", e) + return fallback diff --git a/agent/main.py b/agent/main.py index 86a6de25e..46fa8a381 100644 --- a/agent/main.py +++ b/agent/main.py @@ -455,6 +455,10 @@ def _cancel_event(): else: console.print("[dim]Started new chat.[/dim]") turn_complete_event.set() + elif event.event_type == "conversation_title": + title = (event.data or {}).get("title") + if title: + console.print(f"[dim]Titled this session:[/dim] [cyan]{title}[/cyan]") elif event.event_type == "resume_complete": data = event.data or {} path = data.get("path", "?") @@ -991,6 +995,24 @@ async def _handle_slash_command( ), ) + if command == "/rename": + session = session_holder[0] if session_holder else None + if session is None: + get_console().print("[bold red]No active session to rename.[/bold red]") + return None + new_title = arg.strip() + if not new_title: + get_console().print( + "[dim]Usage: /rename — give the current session a title.[/dim]" + ) + return None + session.session_title = new_title + session._title_user_set = True + get_console().print( + f"[green]Renamed session to[/green] [cyan]{new_title}[/cyan]." + ) + return None + if command == "/model": console = get_console() if not arg: diff --git a/agent/utils/terminal_display.py b/agent/utils/terminal_display.py index 45850e893..ef3b896a0 100644 --- a/agent/utils/terminal_display.py +++ b/agent/utils/terminal_display.py @@ -459,6 +459,7 @@ def print_yolo_approve(count: int) -> None: ("/undo", "", "Undo last turn"), ("/compact", "", "Compact context window"), ("/resume", "[index|id|path]", "Pick up from ./session_logs"), + ("/rename", "", "Rename the current session"), ("/model", "[id]", "Show available models or switch"), ( "/effort", diff --git a/tests/unit/test_session_title.py b/tests/unit/test_session_title.py new file mode 100644 index 000000000..e6c9ee76e --- /dev/null +++ b/tests/unit/test_session_title.py @@ -0,0 +1,147 @@ +import asyncio +from types import SimpleNamespace + +import pytest +from litellm import Message + +import agent.main as main_mod +from agent.core.agent_loop import _generate_and_set_title +from agent.core.session import Event, Session + + +class _FakeConfig: + model_name = "openai/gpt-5.5:fal-ai" + save_sessions = False + session_dataset_repo = "fake/repo" + auto_save_interval = 1 + heartbeat_interval_s = 60 + max_iterations = 10 + yolo_mode = False + confirm_cpu_jobs = False + auto_file_upload = False + reasoning_effort = None + share_traces = False + personal_trace_repo_template = None + mcpServers: dict = {} + + +class _FakeContext: + def __init__(self) -> None: + self.items = [ + Message(role="system", content="system prompt"), + Message(role="user", content="train a model"), + ] + self.model_max_tokens = 200_000 + self.running_context_usage = 0 + self.on_message_added = None + + +def _make_session() -> Session: + return Session( + event_queue=asyncio.Queue(), + config=_FakeConfig(), + tool_router=None, + context_manager=_FakeContext(), + hf_token=None, + user_id="user-a", + local_mode=True, + ) + + +def test_get_trajectory_includes_session_title_default_none(): + session = _make_session() + traj = session.get_trajectory() + assert "session_title" in traj + assert traj["session_title"] is None + + +def test_get_trajectory_reflects_set_title(): + session = _make_session() + session.session_title = "my experiment" + assert session.get_trajectory()["session_title"] == "my experiment" + + +def test_start_new_conversation_resets_title(): + session = _make_session() + session.session_title = "named run" + session._title_user_set = True + session.start_new_conversation() + assert session.session_title is None + assert session._title_user_set is False + + +@pytest.mark.asyncio +async def test_rename_command_sets_title(): + session = _make_session() + result = await main_mod._handle_slash_command( + "/rename my-experiment", + config=session.config, + session_holder=[session], + submission_queue=asyncio.Queue(), + submission_id=[0], + ) + assert result is None + assert session.session_title == "my-experiment" + assert session._title_user_set is True + + +def _fake_response(content): + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=content))] + ) + + +async def _drain_events(queue: asyncio.Queue) -> list[Event]: + events: list[Event] = [] + while not queue.empty(): + events.append(await queue.get()) + return events + + +@pytest.mark.asyncio +async def test_auto_title_sets_title_and_emits_event(monkeypatch): + async def fake_acompletion(**kwargs): + return _fake_response("Train A Model") + + monkeypatch.setattr("litellm.acompletion", fake_acompletion) + session = _make_session() # context has a "train a model" user message + + await _generate_and_set_title(session, "here is how") + + assert session.session_title == "Train A Model" + events = await _drain_events(session.event_queue) + titled = [e for e in events if e.event_type == "conversation_title"] + assert titled, "expected a conversation_title event" + assert titled[0].data["title"] == "Train A Model" + + +@pytest.mark.asyncio +async def test_auto_title_does_not_override_user_rename(monkeypatch): + async def fake_acompletion(**kwargs): + return _fake_response("Auto Generated Name") + + monkeypatch.setattr("litellm.acompletion", fake_acompletion) + session = _make_session() + session.session_title = "my-name" + session._title_user_set = True + + await _generate_and_set_title(session, "here is how") + + assert session.session_title == "my-name" # user's name preserved + events = await _drain_events(session.event_queue) + assert not [e for e in events if e.event_type == "conversation_title"] + + +@pytest.mark.asyncio +async def test_rename_command_empty_arg_is_noop(): + session = _make_session() + result = await main_mod._handle_slash_command( + "/rename", + config=session.config, + session_holder=[session], + submission_queue=asyncio.Queue(), + submission_id=[0], + ) + assert result is None + assert session.session_title is None + assert session._title_user_set is False diff --git a/tests/unit/test_title.py b/tests/unit/test_title.py new file mode 100644 index 000000000..1a2f93582 --- /dev/null +++ b/tests/unit/test_title.py @@ -0,0 +1,100 @@ +from types import SimpleNamespace + +import pytest + +from agent.core.title import ( + extract_first_user_text, + fallback_title, + generate_conversation_title, +) + + +def test_fallback_title_collapses_and_caps_words(): + out = fallback_title(" fine-tune a model on squad data please now ") + assert out == "fine-tune a model on squad data" # capped at 6 words + + +def test_fallback_title_empty_input(): + assert fallback_title("") == "" + assert fallback_title(None) == "" + assert fallback_title(" ") == "" + + +def test_fallback_title_strips_long_secrets(): + secret = "hf_" + "a" * 40 + out = fallback_title(f"use token {secret} to login") + assert secret not in out + assert "use token" in out + + +def test_extract_first_user_text_from_dicts(): + items = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "first task here"}, + {"role": "user", "content": "second"}, + ] + assert extract_first_user_text(items) == "first task here" + + +def test_extract_first_user_text_from_blocks(): + # Block-list content appears in dict-shaped messages (e.g. from saved JSON); + # litellm's Message only accepts string content. + items = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": [{"type": "text", "text": "block one"}]}, + ] + assert extract_first_user_text(items) == "block one" + + +def test_extract_first_user_text_none_when_no_user(): + items = [{"role": "system", "content": "sys"}] + assert extract_first_user_text(items) == "" + + +def _fake_response(content): + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=content))] + ) + + +@pytest.mark.asyncio +async def test_generate_title_cleans_model_output(monkeypatch): + async def fake_acompletion(**kwargs): + return _fake_response('"Fine-Tune Llama On SQuAD."') + + monkeypatch.setattr("litellm.acompletion", fake_acompletion) + out = await generate_conversation_title( + "openai/gpt-5.5", None, "help me fine-tune llama", "sure" + ) + assert out == "Fine-Tune Llama On SQuAD" # quotes + trailing dot stripped + + +@pytest.mark.asyncio +async def test_generate_title_falls_back_on_empty(monkeypatch): + async def fake_acompletion(**kwargs): + return _fake_response("") + + monkeypatch.setattr("litellm.acompletion", fake_acompletion) + out = await generate_conversation_title( + "openai/gpt-5.5", None, "process my dataset", None + ) + assert out == fallback_title("process my dataset") + + +@pytest.mark.asyncio +async def test_generate_title_falls_back_on_error(monkeypatch): + async def boom(**kwargs): + raise RuntimeError("network down") + + monkeypatch.setattr("litellm.acompletion", boom) + out = await generate_conversation_title( + "openai/gpt-5.5", None, "run inference on a model", None + ) + assert out == fallback_title("run inference on a model") + + +@pytest.mark.asyncio +async def test_generate_title_empty_user_text_returns_fallback(): + # No LLM call should be needed when there's nothing to title from. + out = await generate_conversation_title("openai/gpt-5.5", None, "", None) + assert out == "" From 5d0cf8e180d927d2d93941595554b39d384e1995 Mon Sep 17 00:00:00 2001 From: Mohit Paddhariya Date: Sun, 14 Jun 2026 18:23:43 +0530 Subject: [PATCH 2/5] feat(session): store session logs under XDG data dir with override Session logs were written to ./session_logs relative to the launch directory, so they scattered across folders and /resume only saw the current cwd's logs. - Add resolve_session_log_dir(config): config.session_log_dir > the ML_INTERN_SESSION_DIR env var > legacy ./session_logs if it already exists > XDG default ($XDG_DATA_HOME/ml-intern/sessions, defaulting to ~/.local/share/ml-intern/sessions). - Route save_trajectory_local, the /resume picker, and the failed-upload retry scanner through the one resolver so reads and writes never diverge. - Add a session_log_dir config field; keep ./session_logs as a back-compat fallback so existing checkouts are undisturbed. Part of #325. --- agent/config.py | 4 + agent/core/agent_loop.py | 9 ++- agent/core/session.py | 42 +++++++++- agent/main.py | 9 ++- agent/utils/terminal_display.py | 2 +- tests/unit/test_session_dir.py | 139 ++++++++++++++++++++++++++++++++ 6 files changed, 196 insertions(+), 9 deletions(-) create mode 100644 tests/unit/test_session_dir.py diff --git a/agent/config.py b/agent/config.py index c6784db92..12cff4b49 100644 --- a/agent/config.py +++ b/agent/config.py @@ -27,6 +27,10 @@ class Config(BaseModel): mcpServers: dict[str, MCPServerConfig] = {} save_sessions: bool = True session_dataset_repo: str = "smolagents/ml-intern-sessions" + # Where session logs are stored locally. None = resolve at runtime via + # agent.core.session.resolve_session_log_dir (XDG data dir, overridable by + # the ML_INTERN_SESSION_DIR env var, with a legacy ./session_logs fallback). + session_log_dir: str | None = None # Per-user private dataset that mirrors each session in Claude Code JSONL # format so the HF Agent Trace Viewer auto-renders it # (https://huggingface.co/changelog/agent-trace-viewer). Created private diff --git a/agent/core/agent_loop.py b/agent/core/agent_loop.py index a5ba789a9..15766fd2a 100644 --- a/agent/core/agent_loop.py +++ b/agent/core/agent_loop.py @@ -37,7 +37,12 @@ with_prompt_cache_params, with_prompt_caching, ) -from agent.core.session import DEFAULT_SESSION_LOG_DIR, Event, OpType, Session +from agent.core.session import ( + Event, + OpType, + Session, + resolve_session_log_dir, +) from agent.core.tools import ToolRouter from agent.core.usage_thresholds import ( USAGE_THRESHOLD_TOOL_NAME, @@ -2644,7 +2649,7 @@ async def submission_loop( # to publish to the user's HF dataset gets a fresh attempt on next run. if config and config.save_sessions: Session.retry_failed_uploads_detached( - directory=str(DEFAULT_SESSION_LOG_DIR), + directory=str(resolve_session_log_dir(config)), repo_id=config.session_dataset_repo, personal_repo_id=session._personal_trace_repo_id(), ) diff --git a/agent/core/session.py b/agent/core/session.py index d8b0c972a..e3772ca5d 100644 --- a/agent/core/session.py +++ b/agent/core/session.py @@ -29,6 +29,41 @@ DEFAULT_SESSION_LOG_DIR = Path("session_logs") +# Env var to override where session logs are stored (highest precedence after +# an explicit config value). +SESSION_DIR_ENV_VAR = "ML_INTERN_SESSION_DIR" + + +def resolve_session_log_dir(config: Any | None = None) -> Path: + """Resolve where CLI session logs are read from and written to. + + Precedence (first match wins): + 1. ``config.session_log_dir`` if set. + 2. ``$ML_INTERN_SESSION_DIR`` env var. + 3. Legacy ``./session_logs`` if that directory already exists in the + current working directory (back-compat for existing checkouts). + 4. XDG default: ``$XDG_DATA_HOME/ml-intern/sessions`` + (``~/.local/share/ml-intern/sessions`` when XDG_DATA_HOME is unset). + + Every reader and writer routes through this single helper so reads and + writes never diverge. + """ + if config is not None: + configured = getattr(config, "session_log_dir", None) + if configured: + return Path(configured).expanduser() + + env_dir = os.environ.get(SESSION_DIR_ENV_VAR) + if env_dir: + return Path(env_dir).expanduser() + + if DEFAULT_SESSION_LOG_DIR.exists(): + return DEFAULT_SESSION_LOG_DIR + + xdg_data_home = os.environ.get("XDG_DATA_HOME") + base = Path(xdg_data_home).expanduser() if xdg_data_home else Path.home() / ".local" / "share" + return base / "ml-intern" / "sessions" + def _format_usd(value: Any) -> str: if isinstance(value, bool): @@ -576,7 +611,7 @@ def get_trajectory(self) -> dict: def save_trajectory_local( self, - directory: str = str(DEFAULT_SESSION_LOG_DIR), + directory: str | None = None, upload_status: str = "pending", dataset_url: Optional[str] = None, ) -> Optional[str]: @@ -584,7 +619,8 @@ def save_trajectory_local( Save trajectory to local JSON file as backup with upload status Args: - directory: Directory to save logs (default: "session_logs") + directory: Directory to save logs. When None, resolved from config + via ``resolve_session_log_dir`` (XDG path or override). upload_status: Status of upload attempt ("pending", "success", "failed") dataset_url: URL of dataset if upload succeeded @@ -592,6 +628,8 @@ def save_trajectory_local( Path to saved file if successful, None otherwise """ try: + if directory is None: + directory = str(resolve_session_log_dir(self.config)) log_dir = Path(directory) log_dir.mkdir(parents=True, exist_ok=True) diff --git a/agent/main.py b/agent/main.py index 46fa8a381..f63745de3 100644 --- a/agent/main.py +++ b/agent/main.py @@ -873,6 +873,7 @@ async def get_user_input(prompt_session: PromptSession) -> str: async def _resume_picker( arg: str, prompt_session: PromptSession | None, + config=None, ) -> Path | None: """Resolve a session log path via ``arg`` or interactive selection. @@ -884,13 +885,13 @@ async def _resume_picker( list_session_logs, resolve_session_log_arg, ) - from agent.core.session import DEFAULT_SESSION_LOG_DIR + from agent.core.session import resolve_session_log_dir console = get_console() - directory = DEFAULT_SESSION_LOG_DIR + directory = resolve_session_log_dir(config) entries = list_session_logs(directory) if not entries: - console.print(f"[yellow]No session logs found in ./{directory}.[/yellow]") + console.print(f"[yellow]No session logs found in {directory}.[/yellow]") return None if arg: @@ -984,7 +985,7 @@ async def _handle_slash_command( "[bold red]No active session to restore into.[/bold red]" ) return None - selected_path = await _resume_picker(arg, prompt_session) + selected_path = await _resume_picker(arg, prompt_session, config) if selected_path is None: return None submission_id[0] += 1 diff --git a/agent/utils/terminal_display.py b/agent/utils/terminal_display.py index ef3b896a0..b8a5ecf10 100644 --- a/agent/utils/terminal_display.py +++ b/agent/utils/terminal_display.py @@ -458,7 +458,7 @@ def print_yolo_approve(count: int) -> None: ("/clear", "", "Clear terminal and start fresh"), ("/undo", "", "Undo last turn"), ("/compact", "", "Compact context window"), - ("/resume", "[index|id|path]", "Pick up from ./session_logs"), + ("/resume", "[index|id|path]", "Pick up a saved session"), ("/rename", "", "Rename the current session"), ("/model", "[id]", "Show available models or switch"), ( diff --git a/tests/unit/test_session_dir.py b/tests/unit/test_session_dir.py new file mode 100644 index 000000000..745d639be --- /dev/null +++ b/tests/unit/test_session_dir.py @@ -0,0 +1,139 @@ +import asyncio +from pathlib import Path +from types import SimpleNamespace + +from litellm import Message + +from agent.core.session import ( + DEFAULT_SESSION_LOG_DIR, + SESSION_DIR_ENV_VAR, + Session, + resolve_session_log_dir, +) +from agent.core.session_resume import list_session_logs + + +class _FakeConfig: + model_name = "openai/gpt-5.5" + save_sessions = False + session_dataset_repo = "fake/repo" + session_log_dir = None + auto_save_interval = 1 + heartbeat_interval_s = 60 + max_iterations = 10 + yolo_mode = False + confirm_cpu_jobs = False + auto_file_upload = False + reasoning_effort = None + share_traces = False + personal_trace_repo_template = None + mcpServers: dict = {} + + +class _FakeContext: + def __init__(self) -> None: + self.items = [ + Message(role="system", content="system prompt"), + Message(role="user", content="train a model"), + ] + self.model_max_tokens = 200_000 + self.running_context_usage = 0 + self.on_message_added = None + + +def _make_session() -> Session: + return Session( + event_queue=asyncio.Queue(), + config=_FakeConfig(), + tool_router=None, + context_manager=_FakeContext(), + hf_token=None, + user_id="user-a", + local_mode=True, + ) + + +def _clear_env(monkeypatch): + monkeypatch.delenv(SESSION_DIR_ENV_VAR, raising=False) + monkeypatch.delenv("XDG_DATA_HOME", raising=False) + + +def test_resolve_uses_config_value_first(monkeypatch, tmp_path): + monkeypatch.setenv(SESSION_DIR_ENV_VAR, str(tmp_path / "env")) + config = SimpleNamespace(session_log_dir=str(tmp_path / "cfg")) + assert resolve_session_log_dir(config) == tmp_path / "cfg" + + +def test_resolve_uses_env_when_config_unset(monkeypatch, tmp_path): + monkeypatch.setenv(SESSION_DIR_ENV_VAR, str(tmp_path / "env")) + config = SimpleNamespace(session_log_dir=None) + assert resolve_session_log_dir(config) == tmp_path / "env" + + +def test_resolve_uses_legacy_dir_when_present(monkeypatch, tmp_path): + _clear_env(monkeypatch) + monkeypatch.chdir(tmp_path) + (tmp_path / "session_logs").mkdir() + # config unset, no env, legacy ./session_logs exists in cwd + assert resolve_session_log_dir(SimpleNamespace(session_log_dir=None)) == ( + DEFAULT_SESSION_LOG_DIR + ) + + +def test_resolve_falls_back_to_xdg(monkeypatch, tmp_path): + _clear_env(monkeypatch) + monkeypatch.chdir(tmp_path) # no ./session_logs here + monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "xdgdata")) + expected = tmp_path / "xdgdata" / "ml-intern" / "sessions" + assert resolve_session_log_dir(None) == expected + + +def test_resolve_xdg_default_home(monkeypatch, tmp_path): + _clear_env(monkeypatch) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path / "home")) + expected = tmp_path / "home" / ".local" / "share" / "ml-intern" / "sessions" + assert resolve_session_log_dir(None) == expected + + +def test_save_then_list_round_trip_via_resolver(monkeypatch, tmp_path): + # No legacy dir, XDG pointed at tmp -> save writes there, list finds it. + _clear_env(monkeypatch) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "xdgdata")) + + session = _make_session() + saved = session.save_trajectory_local() # directory=None -> resolver + resolved = resolve_session_log_dir(session.config) + + assert saved is not None + assert Path(saved).parent == resolved + assert resolved == tmp_path / "xdgdata" / "ml-intern" / "sessions" + + entries = list_session_logs(resolved) + assert any(Path(e.path) == Path(saved) for e in entries) + + +def test_save_respects_explicit_directory(monkeypatch, tmp_path): + # Back-compat: an explicit directory still works and overrides resolution. + _clear_env(monkeypatch) + session = _make_session() + explicit = tmp_path / "explicit_logs" + saved = session.save_trajectory_local(directory=str(explicit)) + assert saved is not None + assert Path(saved).parent == explicit + + +def test_precedence_config_over_env_over_legacy(monkeypatch, tmp_path): + # All three present -> config wins + monkeypatch.chdir(tmp_path) + (tmp_path / "session_logs").mkdir() + monkeypatch.setenv(SESSION_DIR_ENV_VAR, str(tmp_path / "env")) + config = SimpleNamespace(session_log_dir=str(tmp_path / "cfg")) + assert resolve_session_log_dir(config) == tmp_path / "cfg" + # Remove config -> env wins over legacy + config = SimpleNamespace(session_log_dir=None) + assert resolve_session_log_dir(config) == tmp_path / "env" + # Remove env -> legacy wins + monkeypatch.delenv(SESSION_DIR_ENV_VAR) + assert resolve_session_log_dir(config) == DEFAULT_SESSION_LOG_DIR From ed98e8c8203194dac54f9f9d9c56b3c14bf39542 Mon Sep 17 00:00:00 2001 From: Mohit Paddhariya Date: Sun, 14 Jun 2026 18:35:13 +0530 Subject: [PATCH 3/5] feat(session): use readable titled filenames for session logs Splice a slugified session title into the log filename so saved sessions are human-scannable: session___.json, falling back to the legacy session__.json when there's no usable title. - Add slugify() in title.py (secret-stripped, lowercase, length-capped). - Factor filename construction into Session._session_log_filename. - When a title is set (auto-title or /rename), rename the active log file in place via apply_title_to_local_file so even a single-turn session ends up titled, preserving the original timestamp and leaving no orphan/dup. - Keep the session_ prefix + .json suffix so the upload-retry glob matches. Part of #325. --- agent/core/agent_loop.py | 3 + agent/core/session.py | 60 +++++++++++-- agent/core/title.py | 18 ++++ agent/main.py | 2 + tests/unit/test_session_filename.py | 128 ++++++++++++++++++++++++++++ 5 files changed, 206 insertions(+), 5 deletions(-) create mode 100644 tests/unit/test_session_filename.py diff --git a/agent/core/agent_loop.py b/agent/core/agent_loop.py index 15766fd2a..6ea722e81 100644 --- a/agent/core/agent_loop.py +++ b/agent/core/agent_loop.py @@ -1201,6 +1201,9 @@ async def _generate_and_set_title(session: "Session", final_response: str | None if not title or session._title_user_set or session.session_title: return session.session_title = title + # Rename the active log so even a single-turn session gets a titled + # filename (falls back to next-save naming if no file exists yet). + session.apply_title_to_local_file() await session.send_event( Event( event_type="conversation_title", diff --git a/agent/core/session.py b/agent/core/session.py index e3772ca5d..6abd7d320 100644 --- a/agent/core/session.py +++ b/agent/core/session.py @@ -609,6 +609,59 @@ def get_trajectory(self) -> dict: "tools": tools, } + def _session_log_filename(self, timestamp: str) -> str: + """Build the local log filename for this session. + + When a title is set, splice its slug in for human-scannable names: + ``session___.json``. The ``session_`` prefix and + ``.json`` suffix are kept so the upload-retry glob ('session_*.json') + still matches; the uuid8 + timestamp keep titled names collision-free + even when two sessions share a title. Falls back to the legacy + ``session__.json`` when there's no usable slug. + """ + slug = "" + if self.session_title: + from agent.core.title import slugify + + slug = slugify(self.session_title) + if slug: + return f"session_{slug}_{self.session_id[:8]}_{timestamp}.json" + return f"session_{self.session_id}_{timestamp}.json" + + def apply_title_to_local_file(self) -> None: + """Reflect the current ``session_title`` in the active log filename. + + Called when a title is set (auto-title or ``/rename``). If a log file + already exists on disk it is renamed in place so even a single-turn + session ends up with a titled filename; otherwise the cached path is + cleared so the next save picks up the title. Safe no-op on any error. + """ + old = self._local_save_path + if not old: + return + old_path = Path(old) + if not old_path.exists(): + self._local_save_path = None + return + # Preserve the original timestamp suffix (last two underscore tokens) + # so chronological ordering is stable; regenerate if it doesn't parse. + parts = old_path.stem.split("_") + tail = "_".join(parts[-2:]) + timestamp = ( + tail + if len(parts) >= 2 and parts[-2].isdigit() and parts[-1].isdigit() + else datetime.now().strftime("%Y%m%d_%H%M%S") + ) + new_path = old_path.with_name(self._session_log_filename(timestamp)) + if new_path == old_path: + return + try: + old_path.rename(new_path) + self._local_save_path = str(new_path) + except OSError as e: + logger.debug("Could not rename log to titled name: %s", e) + self._local_save_path = None + def save_trajectory_local( self, directory: str | None = None, @@ -659,11 +712,8 @@ def save_trajectory_local( if self._local_save_path and Path(self._local_save_path).parent == log_dir: filepath = Path(self._local_save_path) else: - filename = ( - f"session_{self.session_id}_" - f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" - ) - filepath = log_dir / filename + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filepath = log_dir / self._session_log_filename(timestamp) self._local_save_path = str(filepath) # Atomic-ish write: stage to .tmp then rename so a crash mid-write diff --git a/agent/core/title.py b/agent/core/title.py index 2061a4110..7a2888e39 100644 --- a/agent/core/title.py +++ b/agent/core/title.py @@ -52,6 +52,24 @@ def _cap(text: str, max_words: int = MAX_TITLE_WORDS, max_chars: int = MAX_TITLE return text +_SLUG_MAX_CHARS = 40 + + +def slugify(text: str | None) -> str: + """Filesystem-safe lowercase slug from a title. + + Strips secret-like tokens, lowercases, collapses any run of non-alphanumeric + characters to a single dash, and caps length. Returns an empty string when + there's nothing usable (e.g. an all-symbol or empty title), so callers can + fall back to the bare filename pattern. + """ + cleaned = _strip_secrets(text or "").lower() + slug = re.sub(r"[^a-z0-9]+", "-", cleaned).strip("-") + if len(slug) > _SLUG_MAX_CHARS: + slug = slug[:_SLUG_MAX_CHARS].rstrip("-") + return slug + + def fallback_title(first_user_text: str | None) -> str: """Readable title derived from the first user message, no LLM required. diff --git a/agent/main.py b/agent/main.py index f63745de3..94057c612 100644 --- a/agent/main.py +++ b/agent/main.py @@ -1009,6 +1009,8 @@ async def _handle_slash_command( return None session.session_title = new_title session._title_user_set = True + # Rename the active log file to reflect the new title immediately. + session.apply_title_to_local_file() get_console().print( f"[green]Renamed session to[/green] [cyan]{new_title}[/cyan]." ) diff --git a/tests/unit/test_session_filename.py b/tests/unit/test_session_filename.py new file mode 100644 index 000000000..56bec425d --- /dev/null +++ b/tests/unit/test_session_filename.py @@ -0,0 +1,128 @@ +import asyncio +import re +from pathlib import Path + +from litellm import Message + +from agent.core.session import Session +from agent.core.session_resume import list_session_logs +from agent.core.title import slugify + + +class _FakeConfig: + model_name = "openai/gpt-5.5" + save_sessions = False + session_dataset_repo = "fake/repo" + session_log_dir = None + auto_save_interval = 1 + heartbeat_interval_s = 60 + max_iterations = 10 + yolo_mode = False + confirm_cpu_jobs = False + auto_file_upload = False + reasoning_effort = None + share_traces = False + personal_trace_repo_template = None + mcpServers: dict = {} + + +class _FakeContext: + def __init__(self) -> None: + self.items = [ + Message(role="system", content="system prompt"), + Message(role="user", content="train a model"), + ] + self.model_max_tokens = 200_000 + self.running_context_usage = 0 + self.on_message_added = None + + +def _make_session() -> Session: + return Session( + event_queue=asyncio.Queue(), + config=_FakeConfig(), + tool_router=None, + context_manager=_FakeContext(), + hf_token=None, + user_id="user-a", + local_mode=True, + ) + + +_TITLED = re.compile(r"^session_[a-z0-9-]+_[0-9a-f]{8}_\d{8}_\d{6}\.json$") +_LEGACY = re.compile(r"^session_[0-9a-f-]{36}_\d{8}_\d{6}\.json$") + + +def test_titled_filename_when_title_set(tmp_path): + session = _make_session() + session.session_title = "Fine-Tune Llama On SQuAD" + saved = session.save_trajectory_local(directory=str(tmp_path)) + name = Path(saved).name + assert _TITLED.match(name), name + assert slugify(session.session_title) in name + assert name.startswith("session_") and name.endswith(".json") + + +def test_legacy_filename_when_no_title(tmp_path): + session = _make_session() + saved = session.save_trajectory_local(directory=str(tmp_path)) + name = Path(saved).name + assert _LEGACY.match(name), name + + +def test_unsafe_or_empty_title_falls_back_to_legacy(tmp_path): + for bad in ("!!!", "✨✨✨", " "): + session = _make_session() + session.session_title = bad + saved = session.save_trajectory_local(directory=str(tmp_path)) + name = Path(saved).name + assert _LEGACY.match(name), (bad, name) + + +def test_title_change_renames_existing_file_in_place(tmp_path): + session = _make_session() + first = session.save_trajectory_local(directory=str(tmp_path)) + assert Path(first).exists() + + # Setting a title renames the existing file in place (no orphan, no dup). + session.session_title = "my run" + session.apply_title_to_local_file() + renamed = session._local_save_path + + assert renamed != first + assert not Path(first).exists() + assert Path(renamed).exists() + assert "my-run" in Path(renamed).name + assert len(list(tmp_path.glob("session_*.json"))) == 1 + + # Further saves overwrite the same titled file. + again = session.save_trajectory_local(directory=str(tmp_path)) + assert again == renamed + assert len(list(tmp_path.glob("session_*.json"))) == 1 + + +def test_rename_preserves_original_timestamp(tmp_path): + session = _make_session() + first = session.save_trajectory_local(directory=str(tmp_path)) + ts = "_".join(Path(first).stem.split("_")[-2:]) + session.session_title = "my run" + session.apply_title_to_local_file() + assert ts in Path(session._local_save_path).name + + +def test_apply_title_no_file_yet_is_noop(tmp_path): + session = _make_session() + session.session_title = "my run" + session.apply_title_to_local_file() # nothing saved yet -> no crash + assert session._local_save_path is None + # The next save still produces a titled filename. + saved = session.save_trajectory_local(directory=str(tmp_path)) + assert "my-run" in Path(saved).name + + +def test_list_session_logs_parses_titled_files(tmp_path): + session = _make_session() + session.session_title = "Fine-Tune Llama On SQuAD" + saved = session.save_trajectory_local(directory=str(tmp_path)) + entries = list_session_logs(tmp_path) + assert any(Path(e.path) == Path(saved) for e in entries) From 1e57363d8d0fe5e97026dd0ec06322b0d506e90e Mon Sep 17 00:00:00 2001 From: Mohit Paddhariya Date: Sun, 14 Jun 2026 19:36:54 +0530 Subject: [PATCH 4/5] feat(cli): arrow-key resume picker with consistent sort and titles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the 'type a session number' /resume prompt with a type-to-filter arrow-key picker (prompt_toolkit, no new dependency) and surface session titles, while fixing the listing's order and duplicate rows. - New agent/utils/session_picker.py: fzf-style picker — filter by title/preview/model, up/down to move, Enter to select, Esc to cancel, with a scrolling viewport for long lists. /resume stays as a non-interactive fast path and headless fallback. - Show the session title in each row (preview fallback for older logs) and in the resume confirmation; skip slash-command first messages in previews. - Sort and display now use one canonical, tz-aware timestamp (session_end_time then start_time, mtime as tiebreaker/fallback), so the list is genuinely newest-first and labels never look scrambled. session_end_time is now tz-aware. - Dedupe the listing by session_id (newest kept) so a resumed-and-continued conversation — which forks the save path into a new-timestamp file — no longer shows up multiple times. Namespace the legacy no-session_id fallback so it can't wrong-merge. - /rename now persists immediately via Session.persist_title: rename the active log in place AND refresh the persisted session_title, or save a fresh titled log when none exists yet (e.g. right after a resume forked the path), so /resume reflects the new name without waiting for the next turn. - Carry session_title across resume. Closes #325. --- agent/core/agent_loop.py | 5 +- agent/core/session.py | 54 +++++++-- agent/core/session_resume.py | 69 ++++++++--- agent/main.py | 33 +++--- agent/utils/session_picker.py | 175 ++++++++++++++++++++++++++++ tests/unit/test_session_filename.py | 61 ++++++++++ tests/unit/test_session_picker.py | 61 ++++++++++ tests/unit/test_session_resume.py | 116 +++++++++++++++++- 8 files changed, 528 insertions(+), 46 deletions(-) create mode 100644 agent/utils/session_picker.py create mode 100644 tests/unit/test_session_picker.py diff --git a/agent/core/agent_loop.py b/agent/core/agent_loop.py index 6ea722e81..b5b3e5ad9 100644 --- a/agent/core/agent_loop.py +++ b/agent/core/agent_loop.py @@ -1201,9 +1201,8 @@ async def _generate_and_set_title(session: "Session", final_response: str | None if not title or session._title_user_set or session.session_title: return session.session_title = title - # Rename the active log so even a single-turn session gets a titled - # filename (falls back to next-save naming if no file exists yet). - session.apply_title_to_local_file() + # Persist the title so even a single-turn session is titled on disk. + session.persist_title() await session.send_event( Event( event_type="conversation_title", diff --git a/agent/core/session.py b/agent/core/session.py index 6abd7d320..aac728c58 100644 --- a/agent/core/session.py +++ b/agent/core/session.py @@ -600,7 +600,7 @@ def get_trajectory(self) -> dict: "user_id": self.user_id, "hf_username": self.hf_username, "session_start_time": self.session_start_time, - "session_end_time": datetime.now().isoformat(), + "session_end_time": datetime.now().astimezone().isoformat(), "model_name": self.config.model_name, "total_cost_usd": total_cost_usd, "usage_metrics": usage_metrics, @@ -629,12 +629,14 @@ def _session_log_filename(self, timestamp: str) -> str: return f"session_{self.session_id}_{timestamp}.json" def apply_title_to_local_file(self) -> None: - """Reflect the current ``session_title`` in the active log filename. + """Reflect the current ``session_title`` in the active log file. Called when a title is set (auto-title or ``/rename``). If a log file - already exists on disk it is renamed in place so even a single-turn - session ends up with a titled filename; otherwise the cached path is - cleared so the next save picks up the title. Safe no-op on any error. + already exists on disk it is renamed to a titled filename AND its + persisted ``session_title`` field is refreshed, so even a single-turn + session ends up titled both on disk and in the JSON that ``/resume`` + reads. If no file exists yet the cached path is cleared so the next + save picks up the title. Safe no-op on any error. """ old = self._local_save_path if not old: @@ -653,15 +655,45 @@ def apply_title_to_local_file(self) -> None: else datetime.now().strftime("%Y%m%d_%H%M%S") ) new_path = old_path.with_name(self._session_log_filename(timestamp)) - if new_path == old_path: - return try: - old_path.rename(new_path) - self._local_save_path = str(new_path) - except OSError as e: - logger.debug("Could not rename log to titled name: %s", e) + if new_path != old_path: + old_path.rename(new_path) + self._local_save_path = str(new_path) + # Refresh the persisted title inside the file so readers (e.g. + # /resume) show the new name, not just the renamed file. Operate on + # the file path directly so an explicitly-located log is updated in + # place regardless of config-resolved directories. + target = Path(self._local_save_path) + with open(target) as f: + data = json.load(f) + data["session_title"] = self.session_title + tmp_path = target.with_suffix(target.suffix + ".tmp") + with open(tmp_path, "w") as f: + json.dump(data, f, indent=2) + tmp_path.replace(target) + except (OSError, ValueError) as e: + logger.debug("Could not retitle log file: %s", e) self._local_save_path = None + def persist_title(self) -> None: + """Persist the current ``session_title`` to disk right away. + + Renames/updates the existing log when there is one; otherwise saves a + fresh titled log so ``/resume`` reflects the new name even when no file + exists yet for this session — e.g. right after a resume forked the save + path, or before the first turn of a session that already has restored + content. Does nothing for an empty session or when saving is disabled. + """ + self.apply_title_to_local_file() + if self._local_save_path is not None or not self.config.save_sessions: + return + has_content = any( + getattr(item, "role", None) != "system" + for item in self.context_manager.items + ) + if has_content: + self.save_trajectory_local() + def save_trajectory_local( self, directory: str | None = None, diff --git a/agent/core/session_resume.py b/agent/core/session_resume.py index ac7d335f6..457dc82f4 100644 --- a/agent/core/session_resume.py +++ b/agent/core/session_resume.py @@ -33,6 +33,7 @@ class SessionLogEntry: message_count: int preview: str mtime: float + session_title: str | None = None def _message_preview(content: Any, max_chars: int = 72) -> str: @@ -61,11 +62,30 @@ def _first_user_preview(messages: list[Any]) -> str: for raw in messages: if isinstance(raw, dict) and raw.get("role") == "user": preview = _message_preview(raw.get("content")) - if preview: + # Skip slash commands (e.g. a leading "/model ...") so the preview + # reflects a real prompt rather than command noise. + if preview and not preview.startswith("/"): return preview return "(no user prompt preview)" +def _sort_timestamp(entry: "SessionLogEntry") -> datetime: + """Single canonical, tz-aware timestamp for both sorting and display. + + Prefers ``session_end_time``, then ``session_start_time``; naive values are + normalized to local tz. Falls back to the file mtime when neither parses, so + unparseable entries keep a sensible (non-collapsed) order. + """ + for ts in (entry.session_end_time, entry.session_start_time): + if isinstance(ts, str) and ts: + try: + dt = datetime.fromisoformat(ts) + except ValueError: + continue + return dt if dt.tzinfo else dt.astimezone() + return datetime.fromtimestamp(entry.mtime).astimezone() + + def list_session_logs( directory: Path = DEFAULT_SESSION_LOG_DIR, ) -> list[SessionLogEntry]: @@ -87,9 +107,13 @@ def list_session_logs( session_id = data.get("session_id") if not isinstance(session_id, str) or not session_id: - session_id = path.stem + # Namespace the fallback so a corrupted/legacy log with no + # session_id can never collide with another file's real id when + # the listing is deduped by session_id below. + session_id = f"legacy:{path.stem}" stat = path.stat() + title = data.get("session_title") entries.append( SessionLogEntry( path=path, @@ -100,27 +124,36 @@ def list_session_logs( message_count=len(messages), preview=_first_user_preview(messages), mtime=stat.st_mtime, + session_title=title if isinstance(title, str) and title else None, ) ) - entries.sort(key=lambda item: item.mtime, reverse=True) - return entries + # Sort and display use the SAME timestamp so the visible order never looks + # scrambled (the old code sorted by mtime but displayed session_end_time). + # mtime is only a tiebreaker for entries sharing a timestamp. + entries.sort(key=lambda e: (_sort_timestamp(e), e.mtime), reverse=True) + + # Collapse multiple on-disk files for the same conversation to one entry. + # A resumed continuation reuses its session_id but forks the save path, so + # continuing writes a new-timestamp file while the original remains — + # without this, /resume shows the same conversation two (or more) times. + # Entries are newest-first, so the first seen per id is the latest state. + deduped: dict[str, SessionLogEntry] = {} + for entry in entries: + deduped.setdefault(entry.session_id, entry) + return list(deduped.values()) def format_session_log_entry(index: int, entry: SessionLogEntry) -> str: - timestamp = entry.session_end_time or entry.session_start_time - label = "unknown time" - if isinstance(timestamp, str) and timestamp: - try: - label = datetime.fromisoformat(timestamp).strftime("%Y-%m-%d %H:%M") - except ValueError: - label = timestamp[:16] + label = _sort_timestamp(entry).astimezone().strftime("%Y-%m-%d %H:%M") short_id = entry.session_id[:8] model = entry.model_name or "unknown model" + # Lead with the human-readable title; fall back to the first-prompt preview + # for older logs that predate session titles. + heading = entry.session_title or entry.preview return ( - f"{index:>2}. {label} {short_id} " - f"{entry.message_count} msgs {model}\n" - f" {entry.preview}" + f"{index:>2}. {heading}\n" + f" {label} {short_id} {entry.message_count} msgs {model}" ) @@ -242,6 +275,13 @@ def restore_session_from_log(session: Any, path: Path) -> dict[str, Any]: saved_user_id = data.get("user_id") is_continuation = saved_user_id == session.user_id + # Carry the saved title across the resume so the conversation keeps its + # name and auto-titling doesn't re-fire on it. + saved_title = data.get("session_title") + if isinstance(saved_title, str) and saved_title: + session.session_title = saved_title + session._title_user_set = True + if is_continuation: if isinstance(saved_session_id, str) and saved_session_id: session.session_id = saved_session_id @@ -283,6 +323,7 @@ def restore_session_from_log(session: Any, path: Path) -> dict[str, Any]: "restored_count": len(restored_messages), "dropped_count": dropped_count, "model_name": session.config.model_name, + "session_title": session.session_title, "invalid_saved_model": invalid_saved_model, "forked": not is_continuation, "had_redacted_content": _has_redacted_content(raw_messages), diff --git a/agent/main.py b/agent/main.py index 94057c612..a053e8b21 100644 --- a/agent/main.py +++ b/agent/main.py @@ -468,9 +468,11 @@ def _cancel_event(): invalid_model = data.get("invalid_saved_model") forked = bool(data.get("forked", False)) redacted = bool(data.get("had_redacted_content", False)) + title = data.get("session_title") verb = "Forked from" if forked else "Resumed" + titled = f" [cyan]{title}[/cyan]" if title else "" console.print( - f"[green]{verb}[/green] {path} " + f"[green]{verb}[/green]{titled} {path} " f"([cyan]{count}[/cyan] messages, " f"model [cyan]{model}[/cyan])." ) @@ -900,30 +902,26 @@ async def _resume_picker( console.print(f"[bold red]No matching session log:[/bold red] {arg}") return selected - console.print() - console.print("[bold]Saved sessions[/bold]") - for index, entry in enumerate(entries, start=1): - console.print(format_session_log_entry(index, entry)) - console.print() - if prompt_session is None: + # Headless/non-interactive: print the list so it's still visible, but + # there's no TTY to drive the picker. + console.print() + console.print("[bold]Saved sessions[/bold]") + for index, entry in enumerate(entries, start=1): + console.print(format_session_log_entry(index, entry)) console.print("[yellow]Cannot prompt for a selection here.[/yellow]") return None + from agent.utils.session_picker import pick_session_interactive + try: - choice = await prompt_session.prompt_async( - "Select session number (blank to cancel): " - ) + selected = await pick_session_interactive(entries) except (EOFError, KeyboardInterrupt): console.print("[dim]Resume cancelled.[/dim]") return None - choice = choice.strip() - if not choice: + if selected is None: console.print("[dim]Resume cancelled.[/dim]") return None - selected = resolve_session_log_arg(choice, entries, directory) - if selected is None: - console.print(f"[bold red]Invalid selection:[/bold red] {choice}") return selected @@ -1009,8 +1007,9 @@ async def _handle_slash_command( return None session.session_title = new_title session._title_user_set = True - # Rename the active log file to reflect the new title immediately. - session.apply_title_to_local_file() + # Persist the new title immediately so /resume reflects it even when no + # file exists yet (e.g. right after a resume forked the save path). + session.persist_title() get_console().print( f"[green]Renamed session to[/green] [cyan]{new_title}[/cyan]." ) diff --git a/agent/utils/session_picker.py b/agent/utils/session_picker.py new file mode 100644 index 000000000..9e17e41e0 --- /dev/null +++ b/agent/utils/session_picker.py @@ -0,0 +1,175 @@ +"""Type-to-filter interactive picker for ``/resume``. + +A small fzf-style picker built on prompt_toolkit (already a dependency): type to +filter the saved sessions by title/preview/model, arrow keys to move, Enter to +select, Esc/Ctrl+C to cancel. Falls back are handled by the caller; this module +only renders the picker and returns the chosen ``Path`` (or ``None``). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from agent.core.session_resume import SessionLogEntry + +_MAX_VISIBLE = 8 + + +def _entry_haystack(entry: "SessionLogEntry") -> str: + """Lowercased text a filter query is matched against.""" + parts = [ + entry.session_title or "", + entry.preview or "", + entry.model_name or "", + entry.session_id or "", + ] + return " ".join(parts).lower() + + +def filter_session_entries( + entries: list["SessionLogEntry"], query: str +) -> list["SessionLogEntry"]: + """Return entries matching ``query`` (order preserved). + + Whitespace-separated terms are ANDed; each must be a substring of the + entry's title/preview/model/id. An empty query matches everything. + """ + terms = query.lower().split() + if not terms: + return list(entries) + out = [] + for entry in entries: + hay = _entry_haystack(entry) + if all(term in hay for term in terms): + out.append(entry) + return out + + +def _row_label(entry: "SessionLogEntry") -> tuple[str, str]: + """Return (heading, meta) display strings for one entry.""" + from agent.core.session_resume import _sort_timestamp + + heading = entry.session_title or entry.preview or "(untitled session)" + label = _sort_timestamp(entry).astimezone().strftime("%Y-%m-%d %H:%M") + model = entry.model_name or "unknown model" + meta = f"{label} · {entry.message_count} msgs · {model}" + return heading, meta + + +async def pick_session_interactive( + entries: list["SessionLogEntry"], +) -> Path | None: + """Run the type-to-filter picker; return the chosen path or None.""" + if not entries: + return None + + from prompt_toolkit.application import Application + from prompt_toolkit.key_binding import KeyBindings + from prompt_toolkit.keys import Keys + from prompt_toolkit.layout import Layout + from prompt_toolkit.layout.containers import HSplit, Window + from prompt_toolkit.layout.controls import FormattedTextControl + from prompt_toolkit.styles import Style + + state = {"query": "", "sel": 0, "offset": 0} + + def matches() -> list["SessionLogEntry"]: + return filter_session_entries(entries, state["query"]) + + def get_fragments(): + rows = matches() + if state["sel"] >= len(rows): + state["sel"] = max(0, len(rows) - 1) + # Keep the selection inside the visible viewport. + if state["sel"] < state["offset"]: + state["offset"] = state["sel"] + elif state["sel"] >= state["offset"] + _MAX_VISIBLE: + state["offset"] = state["sel"] - _MAX_VISIBLE + 1 + + frags: list[tuple[str, str]] = [] + frags.append( + ("class:header", "Resume a session (type to filter · ↑↓ · Enter · Esc)\n") + ) + frags.append(("class:filter", f"\nfilter > {state['query']}\n\n")) + + if not rows: + frags.append(("class:meta", " no matching sessions\n")) + else: + window = rows[state["offset"] : state["offset"] + _MAX_VISIBLE] + if state["offset"] > 0: + frags.append(("class:meta", " ↑ more\n")) + for i, entry in enumerate(window): + row_idx = state["offset"] + i + selected = row_idx == state["sel"] + heading, meta = _row_label(entry) + marker = "❯ " if selected else " " + head_style = "class:selected" if selected else "" + frags.append((head_style, f"{marker}{heading}\n")) + frags.append(("class:meta", f" {meta}\n")) + if state["offset"] + _MAX_VISIBLE < len(rows): + frags.append(("class:meta", " ↓ more\n")) + + frags.append(("class:meta", f"\n{len(rows)} of {len(entries)} match\n")) + return frags + + kb = KeyBindings() + + @kb.add("up") + @kb.add("c-p") + def _(event): + state["sel"] = max(0, state["sel"] - 1) + + @kb.add("down") + @kb.add("c-n") + def _(event): + rows = matches() + state["sel"] = min(len(rows) - 1, state["sel"] + 1) if rows else 0 + + @kb.add("enter") + def _(event): + rows = matches() + result = rows[state["sel"]].path if rows else None + event.app.exit(result=result) + + @kb.add("c-c") + @kb.add("escape") + def _(event): + event.app.exit(result=None) + + @kb.add("backspace") + def _(event): + state["query"] = state["query"][:-1] + state["sel"] = 0 + + @kb.add("c-u") + def _(event): + state["query"] = "" + state["sel"] = 0 + + @kb.add(Keys.Any) + def _(event): + if event.data and event.data.isprintable(): + state["query"] += event.data + state["sel"] = 0 + + style = Style.from_dict( + { + "header": "bold", + "filter": "ansicyan", + "selected": "bold ansicyan", + "meta": "ansibrightblack", + } + ) + + control = FormattedTextControl(get_fragments, focusable=True, show_cursor=False) + layout = Layout(HSplit([Window(content=control, wrap_lines=True)])) + app = Application( + layout=layout, + key_bindings=kb, + style=style, + full_screen=False, + mouse_support=False, + ) + return await app.run_async() diff --git a/tests/unit/test_session_filename.py b/tests/unit/test_session_filename.py index 56bec425d..9933e3cc2 100644 --- a/tests/unit/test_session_filename.py +++ b/tests/unit/test_session_filename.py @@ -101,6 +101,28 @@ def test_title_change_renames_existing_file_in_place(tmp_path): assert len(list(tmp_path.glob("session_*.json"))) == 1 +def test_rename_refreshes_persisted_title(tmp_path): + import json + + from agent.core.session_resume import list_session_logs + + session = _make_session() + session.save_trajectory_local(directory=str(tmp_path)) + + session.session_title = "my run" + session.apply_title_to_local_file() + renamed = session._local_save_path + + # The JSON content's session_title is updated, not just the filename. + data = json.loads(Path(renamed).read_text()) + assert data["session_title"] == "my run" + # And /resume's listing reads that refreshed title. + entry = next( + e for e in list_session_logs(tmp_path) if Path(e.path) == Path(renamed) + ) + assert entry.session_title == "my run" + + def test_rename_preserves_original_timestamp(tmp_path): session = _make_session() first = session.save_trajectory_local(directory=str(tmp_path)) @@ -126,3 +148,42 @@ def test_list_session_logs_parses_titled_files(tmp_path): saved = session.save_trajectory_local(directory=str(tmp_path)) entries = list_session_logs(tmp_path) assert any(Path(e.path) == Path(saved) for e in entries) + + +def test_persist_title_creates_file_when_none_exists(tmp_path): + # Mirrors /rename right after a resume: no file cached, but content exists. + session = _make_session() + session.config.save_sessions = True + session.config.session_log_dir = str(tmp_path) + session.session_title = "demo session" + session._title_user_set = True + session.persist_title() + files = list(tmp_path.glob("session_*.json")) + assert len(files) == 1 + assert "demo-session" in files[0].name + import json + assert json.loads(files[0].read_text())["session_title"] == "demo session" + + +def test_persist_title_renames_existing_file(tmp_path): + session = _make_session() + session.config.save_sessions = True + session.config.session_log_dir = str(tmp_path) + session.save_trajectory_local(directory=str(tmp_path)) + session.session_title = "renamed run" + session.persist_title() + files = list(tmp_path.glob("session_*.json")) + assert len(files) == 1 # renamed in place, no second file + assert "renamed-run" in files[0].name + + +def test_persist_title_noop_for_empty_session(tmp_path): + from litellm import Message + + session = _make_session() + session.config.save_sessions = True + session.config.session_log_dir = str(tmp_path) + session.context_manager.items = [Message(role="system", content="sys")] + session.session_title = "nothing here" + session.persist_title() + assert list(tmp_path.glob("session_*.json")) == [] diff --git a/tests/unit/test_session_picker.py b/tests/unit/test_session_picker.py new file mode 100644 index 000000000..1cbc70395 --- /dev/null +++ b/tests/unit/test_session_picker.py @@ -0,0 +1,61 @@ +from pathlib import Path + +from agent.core.session_resume import SessionLogEntry +from agent.utils.session_picker import _row_label, filter_session_entries + + +def _entry(session_id, title=None, preview="", model=None): + return SessionLogEntry( + path=Path(f"{session_id}.json"), + session_id=session_id, + session_start_time="2026-01-01T00:00:00", + session_end_time="2026-01-01T00:05:00", + model_name=model, + message_count=3, + preview=preview, + mtime=1_700_000_000.0, + session_title=title, + ) + + +def _entries(): + return [ + _entry("a", title="Fine-Tune Llama On SQuAD", model="gpt-4o-mini"), + _entry("b", title="Dataset Audit Helper", model="gpt-5.5"), + _entry("c", preview="APK threat analysis platform", model="gpt-5.5"), + ] + + +def test_empty_query_returns_all(): + entries = _entries() + assert filter_session_entries(entries, "") == entries + assert filter_session_entries(entries, " ") == entries + + +def test_filter_matches_title_case_insensitive(): + out = filter_session_entries(_entries(), "llama") + assert [e.session_id for e in out] == ["a"] + + +def test_filter_matches_preview_and_model(): + assert [e.session_id for e in filter_session_entries(_entries(), "apk")] == ["c"] + # model substring + assert {e.session_id for e in filter_session_entries(_entries(), "gpt-5.5")} == { + "b", + "c", + } + + +def test_filter_terms_are_anded(): + out = filter_session_entries(_entries(), "dataset helper") + assert [e.session_id for e in out] == ["b"] + assert filter_session_entries(_entries(), "dataset llama") == [] + + +def test_row_label_uses_title_then_preview(): + head, meta = _row_label(_entry("a", title="My Title", model="m")) + assert head == "My Title" + head2, _ = _row_label(_entry("b", preview="some prompt", model="m")) + assert head2 == "some prompt" + head3, _ = _row_label(_entry("c", model="m")) + assert head3 == "(untitled session)" diff --git a/tests/unit/test_session_resume.py b/tests/unit/test_session_resume.py index 1a33f5ae2..89ccda9d4 100644 --- a/tests/unit/test_session_resume.py +++ b/tests/unit/test_session_resume.py @@ -23,14 +23,17 @@ def _write_session_log( user_id: str | None = "user-a", extra_messages: list[dict] | None = None, events: list[dict] | None = None, + session_title: str | None = None, + end_time: str = "2026-01-01T00:05:00", ) -> Path: directory.mkdir(exist_ok=True) path = directory / name payload = { "session_id": session_id, + "session_title": session_title, "user_id": user_id, "session_start_time": "2026-01-01T00:00:00", - "session_end_time": "2026-01-01T00:05:00", + "session_end_time": end_time, "model_name": ROUTER_GPT_55, "messages": [ {"role": "system", "content": "old system"}, @@ -64,6 +67,8 @@ def __init__(self, *, user_id: str | None = "user-a") -> None: self.session_id = "current-session" self.session_start_time = "2026-01-02T00:00:00" self.user_id = user_id + self.session_title: str | None = None + self._title_user_set = False self.logged_events: list[dict] = [] self._local_save_path: str | None = None self.turn_count = 0 @@ -382,3 +387,112 @@ def test_resolve_session_log_arg_accepts_index_and_id_prefix(tmp_path): assert session_resume.resolve_session_log_arg("1", entries, log_dir) == newer assert session_resume.resolve_session_log_arg("abc", entries, log_dir) == older assert session_resume.resolve_session_log_arg("nope", entries, log_dir) is None + + +def test_list_populates_session_title_with_preview_fallback(tmp_path): + log_dir = tmp_path / "session_logs" + _write_session_log( + log_dir, "titled.json", session_id="s1", content="train a model", + mtime=time.time(), session_title="Fine-Tune Llama", end_time="2026-02-02T10:00:00", + ) + _write_session_log( + log_dir, "untitled.json", session_id="s2", content="process data", + mtime=time.time() - 5, end_time="2026-02-01T10:00:00", + ) + by_id = {e.session_id: e for e in session_resume.list_session_logs(log_dir)} + assert by_id["s1"].session_title == "Fine-Tune Llama" + assert by_id["s2"].session_title is None + assert by_id["s2"].preview == "process data" # preview still available + + +def test_preview_skips_slash_command_first_message(tmp_path): + log_dir = tmp_path / "session_logs" + _write_session_log( + log_dir, "cmd.json", session_id="s1", content="/model openai/gpt-5.5", + mtime=time.time(), + extra_messages=[{"role": "user", "content": "actually fine-tune llama"}], + ) + entry = session_resume.list_session_logs(log_dir)[0] + assert entry.preview == "actually fine-tune llama" + + +def test_sort_and_display_agree_newest_first(tmp_path): + log_dir = tmp_path / "session_logs" + # Deliberately give the OLDER end_time a NEWER mtime to expose the old + # mtime-sort-vs-end_time-display mismatch; the fix sorts by end_time. + _write_session_log( + log_dir, "a.json", session_id="old", content="x", + mtime=time.time(), end_time="2026-01-01T09:00:00", + ) + _write_session_log( + log_dir, "b.json", session_id="new", content="y", + mtime=time.time() - 100, end_time="2026-03-01T09:00:00", + ) + entries = session_resume.list_session_logs(log_dir) + # Newest end_time first regardless of mtime. + assert [e.session_id for e in entries] == ["new", "old"] + # Displayed labels are in the same (descending) order. + labels = [session_resume._sort_timestamp(e) for e in entries] + assert labels == sorted(labels, reverse=True) + + +def test_sort_timestamp_falls_back_to_mtime(tmp_path): + entry = session_resume.SessionLogEntry( + path=Path("x.json"), session_id="s", session_start_time="not-a-date", + session_end_time="also-bad", model_name=None, message_count=1, + preview="p", mtime=1_700_000_000.0, + ) + ts = session_resume._sort_timestamp(entry) + assert ts.tzinfo is not None # tz-aware, no TypeError + + +def test_format_entry_shows_title(tmp_path): + log_dir = tmp_path / "session_logs" + _write_session_log( + log_dir, "t.json", session_id="s1", content="train a model", + mtime=time.time(), session_title="My Cool Run", + ) + entry = session_resume.list_session_logs(log_dir)[0] + out = session_resume.format_session_log_entry(1, entry) + assert "My Cool Run" in out + + +def test_restore_carries_session_title(tmp_path): + log_dir = tmp_path / "session_logs" + path = _write_session_log( + log_dir, "t.json", session_id="s1", content="hello", + mtime=time.time(), session_title="Resumed Run", user_id="user-a", + ) + session = _FakeSession(user_id="user-a") + result = session_resume.restore_session_from_log(session, path) + assert session.session_title == "Resumed Run" + assert session._title_user_set is True + assert result["session_title"] == "Resumed Run" + + +def test_list_dedupes_same_session_id_keeping_newest(tmp_path): + log_dir = tmp_path / "session_logs" + # Same session_id (a resumed continuation), two files, different end_times. + _write_session_log( + log_dir, "old.json", session_id="dup", content="first", + mtime=time.time() - 100, session_title="Old Title", + end_time="2026-01-01T09:00:00", + ) + _write_session_log( + log_dir, "new.json", session_id="dup", content="continued", + mtime=time.time(), session_title="New Title", + end_time="2026-03-01T09:00:00", + ) + # A genuinely different session that happens to share a title. + _write_session_log( + log_dir, "other.json", session_id="other", content="z", + mtime=time.time() - 50, session_title="New Title", + end_time="2026-02-01T09:00:00", + ) + entries = session_resume.list_session_logs(log_dir) + ids = [e.session_id for e in entries] + assert ids.count("dup") == 1 # collapsed to one + assert "other" in ids # distinct session kept + # The kept "dup" entry is the newest (New Title). + dup = next(e for e in entries if e.session_id == "dup") + assert dup.session_title == "New Title" From fd3ae7c6eaf212c6ceed8716054dd3fb380bf148 Mon Sep 17 00:00:00 2001 From: Mohit Paddhariya Date: Sun, 14 Jun 2026 20:53:27 +0530 Subject: [PATCH 5/5] fix(session): harden auto-title, secret scrubbing, and log storage - Auto-title race: snapshot a per-conversation epoch when spawning the fire-and-forget title task and bail if /new or /resume rotated the conversation during the title LLM call, so a stale title can't be stamped onto a different session. - Fire the auto-title trigger from the usage-threshold / YOLO / abandon completion paths too, so a first turn that paused for an approval still gets titled instead of staying permanently untitled. - Hold a strong reference to the title task so it can't be GC'd mid-await. - Scrub titles through redact.scrub_string on the auto-title path, the /rename path, and the persisted JSON/filename, closing the gap where AWS key ids, Bearer tokens, and NAME=value dumps reached disk via session_title. - Escape titles before rendering them as Rich markup. - Union-read a legacy ./session_logs alongside the resolved dir so pre-XDG-migration sessions stay visible in /resume, and drop the divergent cwd-relative directory defaults that bypassed the resolver. - Guard the title rename and trajectory save with a lock so a heartbeat save on the worker thread can't resurrect a pre-title log file. --- agent/core/agent_loop.py | 88 ++++++++-- agent/core/session.py | 173 +++++++++++++------ agent/core/session_resume.py | 50 +++++- agent/core/title.py | 27 ++- agent/main.py | 32 +++- tests/unit/test_usage_threshold_approvals.py | 6 + 6 files changed, 292 insertions(+), 84 deletions(-) diff --git a/agent/core/agent_loop.py b/agent/core/agent_loop.py index b5b3e5ad9..ddb00179d 100644 --- a/agent/core/agent_loop.py +++ b/agent/core/agent_loop.py @@ -1176,11 +1176,29 @@ async def _call_llm_non_streaming( ) -async def _generate_and_set_title(session: "Session", final_response: str | None) -> None: +# Strong references to in-flight auto-title tasks. asyncio only holds a weak +# reference to a bare create_task result, so without this the task could be +# GC'd mid-await and the title silently dropped (mirrors telemetry.py's +# _heartbeat_tasks pattern). +_title_tasks: set[asyncio.Task] = set() + + +async def _generate_and_set_title( + session: "Session", + final_response: str | None, + origin_session_id: str | None = None, + origin_epoch: int | None = None, +) -> None: """Generate a conversation title and attach it to the session. Runs as a fire-and-forget task after the first turn. Any failure is swallowed so it can never break the turn that spawned it. + + ``origin_session_id`` / ``origin_epoch`` snapshot the conversation identity + at spawn time; if they're omitted they default to the session's current + values. After the (multi-second) title LLM call we bail unless the session + is still the same conversation — otherwise a ``/new`` or ``/resume`` issued + during the await would let us stamp this title onto a different one. """ try: from agent.core.title import ( @@ -1188,6 +1206,11 @@ async def _generate_and_set_title(session: "Session", final_response: str | None generate_conversation_title, ) + if origin_session_id is None: + origin_session_id = session.session_id + if origin_epoch is None: + origin_epoch = session._conversation_epoch + first_user_text = extract_first_user_text(session.context_manager.items) if not first_user_text: return @@ -1197,8 +1220,15 @@ async def _generate_and_set_title(session: "Session", final_response: str | None first_user_text, final_response, ) - # The user may have renamed (or a /new fired) while we were awaiting. - if not title or session._title_user_set or session.session_title: + # Bail if the user renamed, or a /new or /resume rotated the + # conversation, while we were awaiting the title. + if ( + not title + or session._title_user_set + or session.session_title + or session.session_id != origin_session_id + or session._conversation_epoch != origin_epoch + ): return session.session_title = title # Persist the title so even a single-turn session is titled on disk. @@ -1213,6 +1243,35 @@ async def _generate_and_set_title(session: "Session", final_response: str | None logger.debug("Auto-title task failed: %s", e) +def _maybe_spawn_auto_title(session: "Session", final_response: Any) -> None: + """Spawn the one-shot auto-title task if this is the first untitled turn. + + Called from every turn-completion path (normal and the usage-threshold / + YOLO / abandon resume paths) so a first turn that paused for an approval + still gets titled. Snapshots the conversation identity and keeps a strong + reference to the task. Never raises — a title must never break a turn. + """ + try: + if ( + session.turn_count != 0 + or session.session_title + or session._title_user_set + ): + return + task = asyncio.create_task( + _generate_and_set_title( + session, + final_response if isinstance(final_response, str) else None, + session.session_id, + session._conversation_epoch, + ) + ) + _title_tasks.add(task) + task.add_done_callback(_title_tasks.discard) + except Exception as e: # noqa: BLE001 + logger.debug("Auto-title spawn skipped: %s", e) + + class Handlers: """Handler functions for each operation type""" @@ -1258,6 +1317,9 @@ async def _abandon_pending_approval(session: Session) -> None: }, ) ) + # First turn may complete here (paused for an approval, then + # the user continued); title it before turn_count increments. + _maybe_spawn_auto_title(session, final_response) session.increment_turn() await session.auto_save_if_needed() return @@ -1925,17 +1987,7 @@ async def _exec_tool( # Auto-title the conversation once, after the very first completed # turn, unless the user already named it via /rename. Fire-and-forget # so a slow or failing title never delays the turn. - if ( - session.turn_count == 0 - and not session.session_title - and not session._title_user_set - ): - asyncio.create_task( - _generate_and_set_title( - session, - final_response if isinstance(final_response, str) else None, - ) - ) + _maybe_spawn_auto_title(session, final_response) # Increment turn counter and check for auto-save session.increment_turn() @@ -2079,6 +2131,10 @@ async def _exec_usage_threshold_approval( }, ) ) + # First turn may complete here (paused for an approval); title it + # before turn_count increments so it isn't left permanently + # untitled. + _maybe_spawn_auto_title(session, final_response) session.increment_turn() await session.auto_save_if_needed() return @@ -2210,6 +2266,10 @@ async def _exec_yolo_budget_approval( }, ) ) + # First turn may complete here (paused for an approval); title it + # before turn_count increments so it isn't left permanently + # untitled. + _maybe_spawn_auto_title(session, final_response) session.increment_turn() await session.auto_save_if_needed() return diff --git a/agent/core/session.py b/agent/core/session.py index aac728c58..82ecca5f1 100644 --- a/agent/core/session.py +++ b/agent/core/session.py @@ -4,6 +4,7 @@ import os import subprocess import sys +import threading import uuid from dataclasses import dataclass from datetime import datetime @@ -65,6 +66,42 @@ def resolve_session_log_dir(config: Any | None = None) -> Path: return base / "ml-intern" / "sessions" +def legacy_session_log_dirs(resolved: Any | None = None) -> list[Path]: + """Extra legacy dirs to union-read so pre-XDG-migration logs stay visible. + + Returns ``[./session_logs]`` when that cwd-relative legacy dir exists and is + not already the resolved/primary dir; otherwise an empty list. Reads only — + writes always go to the resolved dir, so this never moves or mutates the + legacy logs. + """ + legacy = DEFAULT_SESSION_LOG_DIR + if not legacy.exists(): + return [] + if resolved is not None: + try: + if legacy.resolve() == Path(resolved).resolve(): + return [] + except OSError: + pass + return [legacy] + + +def _scrub_title(title: str | None) -> str | None: + """Best-effort secret scrub for a title before it is written to disk. + + Defence-in-depth: the title sources already scrub at the source, but this + guards the on-disk/uploaded copy against any future un-sanitised source. + """ + if not title: + return title + try: + from agent.core.redact import scrub_string + + return scrub_string(title) + except Exception: + return title + + def _format_usd(value: Any) -> str: if isinstance(value, bool): return "$0.00" @@ -181,6 +218,15 @@ def __init__( # auto-titling never clobbers a name the user chose. self.session_title: str | None = session_title self._title_user_set: bool = session_title is not None + # Bumped on every conversation rotation (``/new`` reset, resume) so a + # fire-and-forget auto-title task spawned for one conversation can + # detect that it finished too late and refuse to stamp a stale title + # onto a different one. + self._conversation_epoch: int = 0 + # Serialises on-disk log mutations between the event-loop thread (title + # rename) and the heartbeat worker thread (trajectory save) so a save + # can't resurrect a file the rename just moved away. + self._save_lock = threading.Lock() self.inference_billing_session_id: str | None = None self.config = config self.is_running = True @@ -502,6 +548,9 @@ def start_new_conversation(self) -> dict[str, Any]: self.session_id = str(uuid.uuid4()) self.session_title = None self._title_user_set = False + # Rotate so any auto-title task still in flight for the previous + # conversation bails instead of titling this fresh one. + self._conversation_epoch += 1 self.inference_billing_session_id = None self.session_start_time = datetime.now().astimezone().isoformat() self.turn_count = 0 @@ -596,7 +645,7 @@ def get_trajectory(self) -> dict: usage_metrics = self.usage_metrics or {} return { "session_id": self.session_id, - "session_title": self.session_title, + "session_title": _scrub_title(self.session_title), "user_id": self.user_id, "hf_username": self.hf_username, "session_start_time": self.session_start_time, @@ -638,42 +687,47 @@ def apply_title_to_local_file(self) -> None: reads. If no file exists yet the cached path is cleared so the next save picks up the title. Safe no-op on any error. """ - old = self._local_save_path - if not old: - return - old_path = Path(old) - if not old_path.exists(): - self._local_save_path = None - return - # Preserve the original timestamp suffix (last two underscore tokens) - # so chronological ordering is stable; regenerate if it doesn't parse. - parts = old_path.stem.split("_") - tail = "_".join(parts[-2:]) - timestamp = ( - tail - if len(parts) >= 2 and parts[-2].isdigit() and parts[-1].isdigit() - else datetime.now().strftime("%Y%m%d_%H%M%S") - ) - new_path = old_path.with_name(self._session_log_filename(timestamp)) - try: - if new_path != old_path: - old_path.rename(new_path) - self._local_save_path = str(new_path) - # Refresh the persisted title inside the file so readers (e.g. - # /resume) show the new name, not just the renamed file. Operate on - # the file path directly so an explicitly-located log is updated in - # place regardless of config-resolved directories. - target = Path(self._local_save_path) - with open(target) as f: - data = json.load(f) - data["session_title"] = self.session_title - tmp_path = target.with_suffix(target.suffix + ".tmp") - with open(tmp_path, "w") as f: - json.dump(data, f, indent=2) - tmp_path.replace(target) - except (OSError, ValueError) as e: - logger.debug("Could not retitle log file: %s", e) - self._local_save_path = None + # Hold the save lock across the read of _local_save_path, the rename, + # and the rewrite so a heartbeat save on the worker thread can't read + # the pre-rename path and recreate it as an orphan after we move it. + with self._save_lock: + old = self._local_save_path + if not old: + return + old_path = Path(old) + if not old_path.exists(): + self._local_save_path = None + return + # Preserve the original timestamp suffix (last two underscore + # tokens) so chronological ordering is stable; regenerate if it + # doesn't parse. + parts = old_path.stem.split("_") + tail = "_".join(parts[-2:]) + timestamp = ( + tail + if len(parts) >= 2 and parts[-2].isdigit() and parts[-1].isdigit() + else datetime.now().strftime("%Y%m%d_%H%M%S") + ) + new_path = old_path.with_name(self._session_log_filename(timestamp)) + try: + if new_path != old_path: + old_path.rename(new_path) + self._local_save_path = str(new_path) + # Refresh the persisted title inside the file so readers (e.g. + # /resume) show the new name, not just the renamed file. Operate + # on the file path directly so an explicitly-located log is + # updated in place regardless of config-resolved directories. + target = Path(self._local_save_path) + with open(target) as f: + data = json.load(f) + data["session_title"] = _scrub_title(self.session_title) + tmp_path = target.with_suffix(target.suffix + ".tmp") + with open(tmp_path, "w") as f: + json.dump(data, f, indent=2) + tmp_path.replace(target) + except (OSError, ValueError) as e: + logger.debug("Could not retitle log file: %s", e) + self._local_save_path = None def persist_title(self) -> None: """Persist the current ``session_title`` to disk right away. @@ -741,19 +795,29 @@ def save_trajectory_local( # the same file instead of creating a new timestamped file every # minute. The timestamp in the filename is kept for first-save # ordering; subsequent saves just rewrite that file. - if self._local_save_path and Path(self._local_save_path).parent == log_dir: - filepath = Path(self._local_save_path) - else: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filepath = log_dir / self._session_log_filename(timestamp) - self._local_save_path = str(filepath) - - # Atomic-ish write: stage to .tmp then rename so a crash mid-write - # doesn't leave a truncated JSON that breaks the retry scanner. - tmp_path = filepath.with_suffix(filepath.suffix + ".tmp") - with open(tmp_path, "w") as f: - json.dump(trajectory, f, indent=2) - tmp_path.replace(filepath) + # + # Hold the save lock across path selection and the write so a + # concurrent title rename (event-loop thread) and this save + # (heartbeat worker thread) can't interleave and leave an orphan + # copy at a path the rename moved away from. + with self._save_lock: + if ( + self._local_save_path + and Path(self._local_save_path).parent == log_dir + ): + filepath = Path(self._local_save_path) + else: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filepath = log_dir / self._session_log_filename(timestamp) + self._local_save_path = str(filepath) + + # Atomic-ish write: stage to .tmp then rename so a crash + # mid-write doesn't leave a truncated JSON that breaks the retry + # scanner. + tmp_path = filepath.with_suffix(filepath.suffix + ".tmp") + with open(tmp_path, "w") as f: + json.dump(trajectory, f, indent=2) + tmp_path.replace(filepath) return str(filepath) except Exception as e: @@ -871,7 +935,7 @@ def save_and_upload_detached(self, repo_id: str) -> Optional[str]: @staticmethod def retry_failed_uploads_detached( - directory: str = str(DEFAULT_SESSION_LOG_DIR), + directory: Optional[str] = None, repo_id: Optional[str] = None, *, personal_repo_id: Optional[str] = None, @@ -881,7 +945,9 @@ def retry_failed_uploads_detached( (fire-and-forget). Args: - directory: Directory containing session logs + directory: Directory containing session logs. When ``None``, + resolved via ``resolve_session_log_dir`` (XDG path or override) + so it never diverges from where saves are written. repo_id: Target dataset repo ID for the shared org/KPI upload. personal_repo_id: Per-user dataset for Claude-Code-format retries. ``None`` skips the personal retry pass. @@ -889,6 +955,9 @@ def retry_failed_uploads_detached( if not repo_id and not personal_repo_id: return + if directory is None: + directory = str(resolve_session_log_dir()) + try: uploader_script = Path(__file__).parent / "session_uploader.py" diff --git a/agent/core/session_resume.py b/agent/core/session_resume.py index 457dc82f4..cbd894598 100644 --- a/agent/core/session_resume.py +++ b/agent/core/session_resume.py @@ -5,6 +5,7 @@ import json import logging import re +from collections.abc import Iterable from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -14,7 +15,6 @@ from agent.core.model_ids import strip_huggingface_model_prefix from agent.core.model_switcher import is_valid_model_id -from agent.core.session import DEFAULT_SESSION_LOG_DIR logger = logging.getLogger(__name__) @@ -86,10 +86,8 @@ def _sort_timestamp(entry: "SessionLogEntry") -> datetime: return datetime.fromtimestamp(entry.mtime).astimezone() -def list_session_logs( - directory: Path = DEFAULT_SESSION_LOG_DIR, -) -> list[SessionLogEntry]: - """Return readable session logs under ``directory``, newest first.""" +def _read_session_log_entries(directory: Path) -> list[SessionLogEntry]: + """Read every readable ``*.json`` log in one directory (no sort/dedupe).""" if not directory.exists(): return [] @@ -127,6 +125,32 @@ def list_session_logs( session_title=title if isinstance(title, str) and title else None, ) ) + return entries + + +def list_session_logs( + directory: Path, + *, + extra_dirs: Iterable[Path] = (), +) -> list[SessionLogEntry]: + """Return readable session logs, newest first, deduped by ``session_id``. + + ``extra_dirs`` are union-read alongside ``directory`` (e.g. a legacy + cwd-relative ``./session_logs`` left behind by the XDG migration) so old + sessions stay visible regardless of launch cwd. A session present in more + than one dir collapses to its newest file via the dedupe below. + """ + entries: list[SessionLogEntry] = [] + seen_dirs: set[Path] = set() + for d in [directory, *extra_dirs]: + try: + key = d.resolve() + except OSError: + key = d + if key in seen_dirs: + continue + seen_dirs.add(key) + entries.extend(_read_session_log_entries(d)) # Sort and display use the SAME timestamp so the visible order never looks # scrambled (the old code sorted by mtime but displayed session_end_time). @@ -145,12 +169,16 @@ def list_session_logs( def format_session_log_entry(index: int, entry: SessionLogEntry) -> str: + from rich.markup import escape + label = _sort_timestamp(entry).astimezone().strftime("%Y-%m-%d %H:%M") short_id = entry.session_id[:8] model = entry.model_name or "unknown model" # Lead with the human-readable title; fall back to the first-prompt preview - # for older logs that predate session titles. - heading = entry.session_title or entry.preview + # for older logs that predate session titles. Escape so a title/preview + # containing Rich markup (e.g. "[red]") can't inject styling or raise when + # this string is printed via the console. + heading = escape(entry.session_title or entry.preview) return ( f"{index:>2}. {heading}\n" f" {label} {short_id} {entry.message_count} msgs {model}" @@ -160,7 +188,7 @@ def format_session_log_entry(index: int, entry: SessionLogEntry) -> str: def resolve_session_log_arg( arg: str, entries: list[SessionLogEntry], - directory: Path = DEFAULT_SESSION_LOG_DIR, + directory: Path, ) -> Path | None: """Resolve ``/resume `` as index, path, filename, or session id prefix.""" value = arg.strip() @@ -275,6 +303,12 @@ def restore_session_from_log(session: Any, path: Path) -> dict[str, Any]: saved_user_id = data.get("user_id") is_continuation = saved_user_id == session.user_id + # Rotate the conversation epoch so an auto-title task still in flight for + # the pre-resume conversation bails instead of stamping its title onto this + # one (a forked resume keeps no field the title-task guard would otherwise + # catch). getattr keeps test doubles without the attr working. + session._conversation_epoch = getattr(session, "_conversation_epoch", 0) + 1 + # Carry the saved title across the resume so the conversation keeps its # name and auto-titling doesn't re-fire on it. saved_title = data.get("session_title") diff --git a/agent/core/title.py b/agent/core/title.py index 7a2888e39..ded95097d 100644 --- a/agent/core/title.py +++ b/agent/core/title.py @@ -21,14 +21,16 @@ import re from typing import Any +from agent.core.redact import scrub_string + logger = logging.getLogger(__name__) # Keep titles short and skimmable in a one-line picker row. MAX_TITLE_WORDS = 6 MAX_TITLE_CHARS = 50 -# Strip anything that looks like a long opaque secret (tokens, keys, base64 -# blobs) before it can land in a persisted title or a filename. The title is +# Generic fallback for long opaque secrets (tokens, keys, base64 blobs) that the +# shared redactor's prefix-anchored patterns don't recognise. The title is # derived from the first user message, which may contain a pasted credential. _SECRET_RUN = re.compile(r"\b[A-Za-z0-9_\-]{30,}\b") @@ -39,8 +41,25 @@ def _collapse(text: str) -> str: def _strip_secrets(text: str) -> str: - """Drop long credential-like tokens so they never reach disk via a title.""" - return _collapse(_SECRET_RUN.sub("", text)) + """Drop credential-like tokens so they never reach disk via a title. + + The title is the one trajectory field that bypasses the save-time + ``redact.scrub`` pass, so it must scrub itself. We run the project's shared + ``scrub_string`` first — it catches structured secrets the generic heuristic + misses (AWS key ids, ``Bearer`` tokens, ``NAME=value`` env dumps, sk-/hf_/ + gh_ keys) — then strip any remaining long opaque run. + """ + return _collapse(_SECRET_RUN.sub("", scrub_string(str(text)))) + + +def strip_title_secrets(text: str | None) -> str: + """Scrub secrets from an explicitly-provided title (e.g. ``/rename``). + + Unlike the auto-title path this does not cap length — the user chose the + name — it only removes credentials so they can't reach a filename or the + persisted JSON title. + """ + return _strip_secrets(text or "") def _cap(text: str, max_words: int = MAX_TITLE_WORDS, max_chars: int = MAX_TITLE_CHARS) -> str: diff --git a/agent/main.py b/agent/main.py index a053e8b21..ae9e64bc4 100644 --- a/agent/main.py +++ b/agent/main.py @@ -458,7 +458,11 @@ def _cancel_event(): elif event.event_type == "conversation_title": title = (event.data or {}).get("title") if title: - console.print(f"[dim]Titled this session:[/dim] [cyan]{title}[/cyan]") + from rich.markup import escape + + console.print( + f"[dim]Titled this session:[/dim] [cyan]{escape(title)}[/cyan]" + ) elif event.event_type == "resume_complete": data = event.data or {} path = data.get("path", "?") @@ -470,7 +474,12 @@ def _cancel_event(): redacted = bool(data.get("had_redacted_content", False)) title = data.get("session_title") verb = "Forked from" if forked else "Resumed" - titled = f" [cyan]{title}[/cyan]" if title else "" + if title: + from rich.markup import escape + + titled = f" [cyan]{escape(title)}[/cyan]" + else: + titled = "" console.print( f"[green]{verb}[/green]{titled} {path} " f"([cyan]{count}[/cyan] messages, " @@ -887,11 +896,15 @@ async def _resume_picker( list_session_logs, resolve_session_log_arg, ) - from agent.core.session import resolve_session_log_dir + from agent.core.session import legacy_session_log_dirs, resolve_session_log_dir console = get_console() directory = resolve_session_log_dir(config) - entries = list_session_logs(directory) + # Union-read any legacy ./session_logs left behind by the XDG migration so + # pre-upgrade sessions stay visible even when launched from a different cwd. + entries = list_session_logs( + directory, extra_dirs=legacy_session_log_dirs(directory) + ) if not entries: console.print(f"[yellow]No session logs found in {directory}.[/yellow]") return None @@ -999,7 +1012,14 @@ async def _handle_slash_command( if session is None: get_console().print("[bold red]No active session to rename.[/bold red]") return None - new_title = arg.strip() + from rich.markup import escape + + from agent.core.title import strip_title_secrets + + # Scrub credentials from the explicit name: like the auto-title path, + # the title reaches a filename and the persisted JSON, which bypass the + # trajectory-wide redactor. + new_title = strip_title_secrets(arg) if not new_title: get_console().print( "[dim]Usage: /rename — give the current session a title.[/dim]" @@ -1011,7 +1031,7 @@ async def _handle_slash_command( # file exists yet (e.g. right after a resume forked the save path). session.persist_title() get_console().print( - f"[green]Renamed session to[/green] [cyan]{new_title}[/cyan]." + f"[green]Renamed session to[/green] [cyan]{escape(new_title)}[/cyan]." ) return None diff --git a/tests/unit/test_usage_threshold_approvals.py b/tests/unit/test_usage_threshold_approvals.py index 521decd4a..ea5222f95 100644 --- a/tests/unit/test_usage_threshold_approvals.py +++ b/tests/unit/test_usage_threshold_approvals.py @@ -28,6 +28,12 @@ def __init__(self, *, continuation="continue_agent"): self.events: list[Event] = [] self.turn_count = 0 self.auto_saved = False + # Title fields a real Session carries; the completion path may spawn the + # one-shot auto-title task, which no-ops here (no user message to title). + self.session_id = "usage-fake" + self.session_title = None + self._title_user_set = False + self._conversation_epoch = 0 async def send_event(self, event: Event): self.events.append(event)