Cherry-picked today's changes from HF Space development:

agent/context_manager/manager.py

@@ -243,6 +243,25 @@ class ContextManager:
 
         return False
 
+    def truncate_to_user_message(self, user_message_index: int) -> bool:
+        """Truncate history to just before the Nth user message (0-indexed).
+
+        Removes that user message and everything after it.
+        System message (index 0) is never removed.
+
+        Returns True if the target user message was found and removed.
+        """
+        count = 0
+        for i, msg in enumerate(self.items):
+            if i == 0:
+                continue  # skip system message
+            if getattr(msg, "role", None) == "user":
+                if count == user_message_index:
+                    self.items = self.items[:i]
+                    return True
+                count += 1
+        return False
+
     async def compact(
         self, model_name: str, tool_specs: list[dict] | None = None
     ) -> None:

agent/prompts/system_prompt_v3.yaml

@@ -15,6 +15,12 @@ system_prompt: |
 
   The sub-agent knows how to use github_find_examples, github_read_file, explore_hf_docs, fetch_hf_docs, hf_inspect_dataset, and hf_papers. Be specific in your task description.
 
+  When researching an ML task, include a SOTA check: tell the research sub-agent to search for recent papers on the task or technique to find what approaches, architectures, and hyperparameters are currently achieving the best results. This prevents you from using outdated methods when better ones exist.
+
+  ```
+  research({"task": "Find SOTA approaches for [task]. Search recent papers for best-performing methods, key hyperparameters, and tricks. Also find working code examples using current TRL/Transformers APIs.", "context": "User wants to [goal]."})
+  ```
+
   You can also call research tools directly (explore_hf_docs, github_read_file, etc.) for quick lookups.
 
   Skip research only for trivial non-code operations.

agent/tools/papers_tool.py

@@ -2,11 +2,14 @@
 HF Papers Tool — Discover papers, read their contents, and find linked resources.
 
 Operations: trending, search, paper_details, read_paper,
-            find_datasets, find_models, find_collections, find_all_resources
+            find_datasets, find_models, find_collections, find_all_resources,
+            citation_graph, snippet_search, recommend
 """
 
 import asyncio
+import os
 import re
+import time
 from typing import Any
 
 import httpx
@@ -30,6 +33,101 @@ SORT_MAP = {
     "trending": "trendingScore",
 }
 
+# ---------------------------------------------------------------------------
+# Semantic Scholar API
+# ---------------------------------------------------------------------------
+
+S2_API = "https://api.semanticscholar.org"
+S2_API_KEY = os.environ.get("S2_API_KEY")
+S2_HEADERS: dict[str, str] = {"x-api-key": S2_API_KEY} if S2_API_KEY else {}
+S2_TIMEOUT = 12
+_s2_last_request: float = 0.0
+
+# Shared response cache (survives across sessions, keyed by (path, params_tuple))
+_s2_cache: dict[str, Any] = {}
+_S2_CACHE_MAX = 500
+
+
+def _s2_paper_id(arxiv_id: str) -> str:
+    """Convert bare arxiv ID to S2 format."""
+    return f"ARXIV:{arxiv_id}"
+
+
+def _s2_cache_key(path: str, params: dict | None) -> str:
+    """Build a hashable cache key from path + sorted params."""
+    p = tuple(sorted((params or {}).items()))
+    return f"{path}:{p}"
+
+
+async def _s2_request(
+    client: httpx.AsyncClient,
+    method: str,
+    path: str,
+    **kwargs: Any,
+) -> httpx.Response | None:
+    """S2 request with 2 retries on 429/5xx. Rate-limited only when using API key."""
+    global _s2_last_request
+    url = f"{S2_API}{path}"
+    kwargs.setdefault("headers", {}).update(S2_HEADERS)
+    kwargs.setdefault("timeout", S2_TIMEOUT)
+
+    for attempt in range(3):
+        # Rate limit only when authenticated (1 req/s for search, 10 req/s for others)
+        if S2_API_KEY:
+            min_interval = 1.0 if "search" in path else 0.1
+            elapsed = time.monotonic() - _s2_last_request
+            if elapsed < min_interval:
+                await asyncio.sleep(min_interval - elapsed)
+        _s2_last_request = time.monotonic()
+
+        try:
+            resp = await client.request(method, url, **kwargs)
+            if resp.status_code == 429:
+                if attempt < 2:
+                    await asyncio.sleep(60)
+                    continue
+                return None
+            if resp.status_code >= 500:
+                if attempt < 2:
+                    await asyncio.sleep(3)
+                    continue
+                return None
+            return resp
+        except (httpx.RequestError, httpx.HTTPStatusError):
+            if attempt < 2:
+                await asyncio.sleep(3)
+                continue
+            return None
+    return None
+
+
+async def _s2_get_json(
+    client: httpx.AsyncClient, path: str, params: dict | None = None,
+) -> dict | None:
+    """Cached S2 GET returning parsed JSON or None."""
+    key = _s2_cache_key(path, params)
+    if key in _s2_cache:
+        return _s2_cache[key]
+
+    resp = await _s2_request(client, "GET", path, params=params or {})
+    if resp and resp.status_code == 200:
+        data = resp.json()
+        if len(_s2_cache) < _S2_CACHE_MAX:
+            _s2_cache[key] = data
+        return data
+    return None
+
+
+async def _s2_get_paper(
+    client: httpx.AsyncClient, arxiv_id: str, fields: str,
+) -> dict | None:
+    """Fetch a single paper from S2 by arxiv ID. Returns None on failure."""
+    return await _s2_get_json(
+        client,
+        f"/graph/v1/paper/{_s2_paper_id(arxiv_id)}",
+        {"fields": fields},
+    )
+
 
 # ---------------------------------------------------------------------------
 # HTML paper parsing
@@ -193,7 +291,7 @@ def _format_paper_list(
     return "\n".join(lines)
 
 
-def _format_paper_detail(paper: dict) -> str:
+def _format_paper_detail(paper: dict, s2_data: dict | None = None) -> str:
     arxiv_id = paper.get("id", "")
     title = paper.get("title", "Unknown")
     upvotes = paper.get("upvotes", 0)
@@ -205,7 +303,12 @@ def _format_paper_detail(paper: dict) -> str:
     authors = paper.get("authors") or []
 
     lines = [f"# {title}"]
-    lines.append(f"**arxiv_id:** {arxiv_id} | **upvotes:** {upvotes}")
+    meta_parts = [f"**arxiv_id:** {arxiv_id}", f"**upvotes:** {upvotes}"]
+    if s2_data:
+        cites = s2_data.get("citationCount", 0)
+        influential = s2_data.get("influentialCitationCount", 0)
+        meta_parts.append(f"**citations:** {cites} ({influential} influential)")
+    lines.append(" | ".join(meta_parts))
     lines.append(f"https://huggingface.co/papers/{arxiv_id}")
     lines.append(f"https://arxiv.org/abs/{arxiv_id}")
 
@@ -218,16 +321,27 @@ def _format_paper_detail(paper: dict) -> str:
 
     if keywords:
         lines.append(f"**Keywords:** {', '.join(keywords)}")
+    if s2_data and s2_data.get("s2FieldsOfStudy"):
+        fields = [f["category"] for f in s2_data["s2FieldsOfStudy"] if f.get("category")]
+        if fields:
+            lines.append(f"**Fields:** {', '.join(fields)}")
+    if s2_data and s2_data.get("venue"):
+        lines.append(f"**Venue:** {s2_data['venue']}")
     if github:
         lines.append(f"**GitHub:** {github} ({stars} stars)")
 
+    if s2_data and s2_data.get("tldr"):
+        tldr_text = s2_data["tldr"].get("text", "")
+        if tldr_text:
+            lines.append(f"\n## TL;DR\n{tldr_text}")
     if ai_summary:
         lines.append(f"\n## AI Summary\n{ai_summary}")
     if summary:
         lines.append(f"\n## Abstract\n{_truncate(summary, 500)}")
 
     lines.append(
-        "\n**Next:** Use read_paper to read specific sections, or find_all_resources to discover linked datasets/models."
+        "\n**Next:** Use read_paper to read specific sections, find_all_resources for linked datasets/models, "
+        "or citation_graph to trace references and citations."
     )
     return "\n".join(lines)
 
@@ -441,11 +555,101 @@ async def _op_trending(args: dict[str, Any], limit: int) -> ToolResult:
     }
 
 
+def _format_s2_paper_list(papers: list[dict], title: str) -> str:
+    """Format a list of S2 paper results."""
+    lines = [f"# {title}"]
+    lines.append(f"Showing {len(papers)} result(s)\n")
+
+    for i, paper in enumerate(papers, 1):
+        ptitle = paper.get("title") or "(untitled)"
+        year = paper.get("year") or "?"
+        cites = paper.get("citationCount", 0)
+        venue = paper.get("venue") or ""
+        ext_ids = paper.get("externalIds") or {}
+        aid = ext_ids.get("ArXiv", "")
+        tldr = (paper.get("tldr") or {}).get("text", "")
+
+        lines.append(f"### {i}. {ptitle}")
+        meta = [f"Year: {year}", f"Citations: {cites}"]
+        if venue:
+            meta.append(f"Venue: {venue}")
+        if aid:
+            meta.append(f"arxiv_id: {aid}")
+        lines.append(" | ".join(meta))
+        if aid:
+            lines.append(f"https://arxiv.org/abs/{aid}")
+        if tldr:
+            lines.append(f"**TL;DR:** {tldr}")
+        lines.append("")
+
+    lines.append("Use paper_details with arxiv_id for full info, or read_paper to read sections.")
+    return "\n".join(lines)
+
+
+async def _s2_bulk_search(query: str, args: dict[str, Any], limit: int) -> ToolResult | None:
+    """Search via S2 bulk endpoint with filters. Returns None on failure."""
+    params: dict[str, Any] = {
+        "query": query,
+        "limit": limit,
+        "fields": "title,externalIds,year,citationCount,tldr,venue,publicationDate",
+    }
+
+    # Date filter
+    date_from = args.get("date_from", "")
+    date_to = args.get("date_to", "")
+    if date_from or date_to:
+        params["publicationDateOrYear"] = f"{date_from}:{date_to}"
+
+    # Fields of study
+    categories = args.get("categories")
+    if categories:
+        params["fieldsOfStudy"] = categories
+
+    # Min citations
+    min_cites = args.get("min_citations")
+    if min_cites:
+        params["minCitationCount"] = str(min_cites)
+
+    # Sort
+    sort_by = args.get("sort_by")
+    if sort_by and sort_by != "relevance":
+        params["sort"] = f"{sort_by}:desc"
+
+    async with httpx.AsyncClient(timeout=15) as client:
+        resp = await _s2_request(client, "GET", "/graph/v1/paper/search/bulk", params=params)
+        if not resp or resp.status_code != 200:
+            return None
+        data = resp.json()
+
+    papers = data.get("data") or []
+    if not papers:
+        return {
+            "formatted": f"No papers found for '{query}' with the given filters.",
+            "totalResults": 0,
+            "resultsShared": 0,
+        }
+
+    formatted = _format_s2_paper_list(papers[:limit], f"Papers matching '{query}' (Semantic Scholar)")
+    return {
+        "formatted": formatted,
+        "totalResults": data.get("total", len(papers)),
+        "resultsShared": min(limit, len(papers)),
+    }
+
+
 async def _op_search(args: dict[str, Any], limit: int) -> ToolResult:
     query = args.get("query")
     if not query:
         return _error("'query' is required for search operation.")
 
+    # Route to S2 when filters are present
+    use_s2 = any(args.get(k) for k in ("date_from", "date_to", "categories", "min_citations", "sort_by"))
+    if use_s2:
+        result = await _s2_bulk_search(query, args, limit)
+        if result is not None:
+            return result
+        # Fall back to HF search (without filters) if S2 fails
+
     async with httpx.AsyncClient(timeout=15) as client:
         resp = await client.get(
             f"{HF_API}/papers/search", params={"q": query, "limit": limit}
@@ -545,6 +749,108 @@ async def _op_read_paper(args: dict[str, Any], limit: int) -> ToolResult:
     return {"formatted": formatted, "totalResults": 1, "resultsShared": 1}
 
 
+# ---------------------------------------------------------------------------
+# Citation graph (Semantic Scholar)
+# ---------------------------------------------------------------------------
+
+
+def _format_citation_entry(entry: dict, show_context: bool = False) -> str:
+    """Format a single citation/reference entry."""
+    paper = entry.get("citingPaper") or entry.get("citedPaper") or {}
+    title = paper.get("title") or "(untitled)"
+    year = paper.get("year") or "?"
+    cites = paper.get("citationCount", 0)
+    ext_ids = paper.get("externalIds") or {}
+    aid = ext_ids.get("ArXiv", "")
+    influential = " **[influential]**" if entry.get("isInfluential") else ""
+
+    parts = [f"- **{title}** ({year}, {cites} cites){influential}"]
+    if aid:
+        parts[0] += f"  arxiv:{aid}"
+
+    if show_context:
+        intents = entry.get("intents") or []
+        if intents:
+            parts.append(f"  Intent: {', '.join(intents)}")
+        contexts = entry.get("contexts") or []
+        for ctx in contexts[:2]:
+            if ctx:
+                parts.append(f"  > {_truncate(ctx, 200)}")
+
+    return "\n".join(parts)
+
+
+def _format_citation_graph(
+    arxiv_id: str,
+    references: list[dict] | None,
+    citations: list[dict] | None,
+) -> str:
+    lines = [f"# Citation Graph for {arxiv_id}"]
+    lines.append(f"https://arxiv.org/abs/{arxiv_id}\n")
+
+    if references is not None:
+        lines.append(f"## References ({len(references)})")
+        if references:
+            for entry in references:
+                lines.append(_format_citation_entry(entry))
+        else:
+            lines.append("No references found.")
+        lines.append("")
+
+    if citations is not None:
+        lines.append(f"## Citations ({len(citations)})")
+        if citations:
+            for entry in citations:
+                lines.append(_format_citation_entry(entry, show_context=True))
+        else:
+            lines.append("No citations found.")
+        lines.append("")
+
+    lines.append("**Tip:** Use paper_details with an arxiv_id from above to explore further.")
+    return "\n".join(lines)
+
+
+async def _op_citation_graph(args: dict[str, Any], limit: int) -> ToolResult:
+    arxiv_id = _validate_arxiv_id(args)
+    if not arxiv_id:
+        return _error("'arxiv_id' is required for citation_graph.")
+
+    direction = args.get("direction", "both")
+    s2_id = _s2_paper_id(arxiv_id)
+    fields = "title,externalIds,year,citationCount,influentialCitationCount,contexts,intents,isInfluential"
+    params = {"fields": fields, "limit": limit}
+
+    async with httpx.AsyncClient(timeout=15) as client:
+        refs, cites = None, None
+        coros = []
+        if direction in ("references", "both"):
+            coros.append(_s2_get_json(client, f"/graph/v1/paper/{s2_id}/references", params))
+        if direction in ("citations", "both"):
+            coros.append(_s2_get_json(client, f"/graph/v1/paper/{s2_id}/citations", params))
+
+        results = await asyncio.gather(*coros, return_exceptions=True)
+        idx = 0
+        if direction in ("references", "both"):
+            r = results[idx]
+            if isinstance(r, dict):
+                refs = r.get("data", [])
+            idx += 1
+        if direction in ("citations", "both"):
+            r = results[idx]
+            if isinstance(r, dict):
+                cites = r.get("data", [])
+
+    if refs is None and cites is None:
+        return _error(f"Could not fetch citation data for {arxiv_id}. Paper may not be indexed by Semantic Scholar.")
+
+    total = (len(refs) if refs else 0) + (len(cites) if cites else 0)
+    return {
+        "formatted": _format_citation_graph(arxiv_id, refs, cites),
+        "totalResults": total,
+        "resultsShared": total,
+    }
+
+
 async def _op_find_datasets(args: dict[str, Any], limit: int) -> ToolResult:
     arxiv_id = _validate_arxiv_id(args)
     if not arxiv_id:
@@ -703,6 +1009,136 @@ async def _op_find_all_resources(args: dict[str, Any], limit: int) -> ToolResult
     return {"formatted": formatted, "totalResults": total, "resultsShared": total}
 
 
+# ---------------------------------------------------------------------------
+# Snippet search (Semantic Scholar)
+# ---------------------------------------------------------------------------
+
+
+def _format_snippets(snippets: list[dict], query: str) -> str:
+    lines = [f"# Snippet Search: '{query}'"]
+    lines.append(f"Found {len(snippets)} matching passage(s)\n")
+
+    for i, item in enumerate(snippets, 1):
+        paper = item.get("paper") or {}
+        ptitle = paper.get("title") or "(untitled)"
+        year = paper.get("year") or "?"
+        cites = paper.get("citationCount", 0)
+        ext_ids = paper.get("externalIds") or {}
+        aid = ext_ids.get("ArXiv", "")
+
+        snippet = item.get("snippet") or {}
+        text = snippet.get("text", "")
+        section = snippet.get("section") or ""
+
+        lines.append(f"### {i}. {ptitle} ({year}, {cites} cites)")
+        if aid:
+            lines.append(f"arxiv:{aid}")
+        if section:
+            lines.append(f"Section: {section}")
+        if text:
+            lines.append(f"> {_truncate(text, 400)}")
+        lines.append("")
+
+    lines.append("Use paper_details or read_paper with arxiv_id to explore a paper further.")
+    return "\n".join(lines)
+
+
+async def _op_snippet_search(args: dict[str, Any], limit: int) -> ToolResult:
+    query = args.get("query")
+    if not query:
+        return _error("'query' is required for snippet_search.")
+
+    params: dict[str, Any] = {
+        "query": query,
+        "limit": limit,
+        "fields": "title,externalIds,year,citationCount",
+    }
+
+    # Optional filters (same as search)
+    date_from = args.get("date_from", "")
+    date_to = args.get("date_to", "")
+    if date_from or date_to:
+        params["publicationDateOrYear"] = f"{date_from}:{date_to}"
+    if args.get("categories"):
+        params["fieldsOfStudy"] = args["categories"]
+    if args.get("min_citations"):
+        params["minCitationCount"] = str(args["min_citations"])
+
+    async with httpx.AsyncClient(timeout=15) as client:
+        resp = await _s2_request(client, "GET", "/graph/v1/snippet/search", params=params)
+        if not resp or resp.status_code != 200:
+            return _error("Snippet search failed. Semantic Scholar may be unavailable.")
+        data = resp.json()
+
+    snippets = data.get("data") or []
+    if not snippets:
+        return {
+            "formatted": f"No snippets found for '{query}'.",
+            "totalResults": 0,
+            "resultsShared": 0,
+        }
+
+    return {
+        "formatted": _format_snippets(snippets, query),
+        "totalResults": len(snippets),
+        "resultsShared": len(snippets),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Recommendations (Semantic Scholar)
+# ---------------------------------------------------------------------------
+
+
+async def _op_recommend(args: dict[str, Any], limit: int) -> ToolResult:
+    positive_ids = args.get("positive_ids")
+    arxiv_id = _validate_arxiv_id(args)
+
+    if not arxiv_id and not positive_ids:
+        return _error("'arxiv_id' or 'positive_ids' is required for recommend.")
+
+    fields = "title,externalIds,year,citationCount,tldr,venue"
+
+    async with httpx.AsyncClient(timeout=15) as client:
+        if positive_ids and not arxiv_id:
+            # Multi-paper recommendations (POST, not cached)
+            pos = [_s2_paper_id(pid.strip()) for pid in positive_ids.split(",") if pid.strip()]
+            neg_raw = args.get("negative_ids", "")
+            neg = [_s2_paper_id(pid.strip()) for pid in neg_raw.split(",") if pid.strip()] if neg_raw else []
+            resp = await _s2_request(
+                client, "POST", "/recommendations/v1/papers/",
+                json={"positivePaperIds": pos, "negativePaperIds": neg},
+                params={"fields": fields, "limit": limit},
+            )
+            if not resp or resp.status_code != 200:
+                return _error("Recommendation request failed. Semantic Scholar may be unavailable.")
+            data = resp.json()
+        else:
+            # Single-paper recommendations (cached)
+            data = await _s2_get_json(
+                client,
+                f"/recommendations/v1/papers/forpaper/{_s2_paper_id(arxiv_id)}",
+                {"fields": fields, "limit": limit, "from": "recent"},
+            )
+            if not data:
+                return _error("Recommendation request failed. Semantic Scholar may be unavailable.")
+
+    papers = data.get("recommendedPapers") or []
+    if not papers:
+        return {
+            "formatted": "No recommendations found.",
+            "totalResults": 0,
+            "resultsShared": 0,
+        }
+
+    title = f"Recommended papers based on {arxiv_id or positive_ids}"
+    return {
+        "formatted": _format_s2_paper_list(papers[:limit], title),
+        "totalResults": len(papers),
+        "resultsShared": min(limit, len(papers)),
+    }
+
+
 # ---------------------------------------------------------------------------
 # Operation dispatch
 # ---------------------------------------------------------------------------
@@ -712,6 +1148,9 @@ _OPERATIONS = {
     "search": _op_search,
     "paper_details": _op_paper_details,
     "read_paper": _op_read_paper,
+    "citation_graph": _op_citation_graph,
+    "snippet_search": _op_snippet_search,
+    "recommend": _op_recommend,
     "find_datasets": _op_find_datasets,
     "find_models": _op_find_models,
     "find_collections": _op_find_collections,
@@ -726,22 +1165,25 @@ _OPERATIONS = {
 HF_PAPERS_TOOL_SPEC = {
     "name": "hf_papers",
     "description": (
-        "Discover ML research papers, find their linked resources (datasets, models, collections), "
-        "and read paper contents on HuggingFace Hub and arXiv.\n\n"
-        "Use this when exploring a research area, looking for datasets for a task, "
-        "implementing a paper's approach, or trying to improve performance on something. "
-        "Typical flow:\n"
-        "  hf_papers(search/trending) → hf_papers(read_paper) → hf_papers(find_all_resources) → hf_inspect_dataset\n\n"
+        "Discover ML research papers, analyze citations, search paper contents, and find linked resources.\n\n"
+        "Combines HuggingFace Hub, arXiv, and Semantic Scholar. Use for exploring research areas, "
+        "finding datasets for a task, tracing citation chains, or implementing a paper's approach.\n\n"
+        "Typical flows:\n"
+        "  search → read_paper → find_all_resources → hf_inspect_dataset\n"
+        "  search → paper_details → citation_graph → read_paper (trace influence)\n"
+        "  snippet_search → paper_details → read_paper (find specific claims)\n\n"
         "Operations:\n"
         "- trending: Get trending daily papers, optionally filter by topic keyword\n"
-        "- search: Full-text search for papers by query\n"
-        "- paper_details: Get metadata, abstract, AI summary, and github link for a paper\n"
-        "- read_paper: Read paper contents — without section: returns abstract + table of contents; "
-        "with section: returns full section text\n"
+        "- search: Search papers. Uses HF by default (ML-tuned). Add date_from/min_citations/categories to use Semantic Scholar with filters\n"
+        "- paper_details: Metadata, abstract, AI summary, github link\n"
+        "- read_paper: Read paper contents — without section: abstract + TOC; with section: full text\n"
+        "- citation_graph: Get references and citations for a paper with influence flags and citation intents\n"
+        "- snippet_search: Semantic search over full-text passages from 12M+ papers\n"
+        "- recommend: Find similar papers (single paper or positive/negative examples)\n"
         "- find_datasets: Find datasets linked to a paper\n"
         "- find_models: Find models linked to a paper\n"
         "- find_collections: Find collections that include a paper\n"
-        "- find_all_resources: Parallel fetch of datasets + models + collections for a paper (unified view)"
+        "- find_all_resources: Parallel fetch of datasets + models + collections for a paper"
     ),
     "parameters": {
         "type": "object",
@@ -754,36 +1196,69 @@ HF_PAPERS_TOOL_SPEC = {
             "query": {
                 "type": "string",
                 "description": (
-                    "Search query. Required for: search. "
-                    "Optional for: trending (filters results by keyword match on title, summary, and AI-generated keywords)."
+                    "Search query. Required for: search, snippet_search. "
+                    "Optional for: trending (filters by keyword). "
+                    "Supports boolean syntax for Semantic Scholar: '\"exact phrase\" term1 | term2'."
                 ),
             },
             "arxiv_id": {
                 "type": "string",
                 "description": (
                     "ArXiv paper ID (e.g. '2305.18290'). "
-                    "Required for: paper_details, read_paper, find_datasets, find_models, find_collections, find_all_resources. "
-                    "Get IDs from trending or search results first."
+                    "Required for: paper_details, read_paper, citation_graph, find_datasets, find_models, find_collections, find_all_resources. "
+                    "Optional for: recommend (single-paper recs). Get IDs from search results first."
                 ),
             },
             "section": {
                 "type": "string",
                 "description": (
                     "Section name or number to read (e.g. '3', 'Experiments', '4.2'). "
-                    "Optional for: read_paper. Without this, read_paper returns the abstract + table of contents "
-                    "so you can choose which section to read."
+                    "Optional for: read_paper. Without this, returns abstract + TOC."
                 ),
             },
+            "direction": {
+                "type": "string",
+                "enum": ["citations", "references", "both"],
+                "description": "Direction for citation_graph. Default: both.",
+            },
             "date": {
                 "type": "string",
                 "description": "Date in YYYY-MM-DD format. Optional for: trending (defaults to recent papers).",
             },
+            "date_from": {
+                "type": "string",
+                "description": "Start date (YYYY-MM-DD). Triggers Semantic Scholar search. For: search, snippet_search.",
+            },
+            "date_to": {
+                "type": "string",
+                "description": "End date (YYYY-MM-DD). Triggers Semantic Scholar search. For: search, snippet_search.",
+            },
+            "categories": {
+                "type": "string",
+                "description": "Field of study filter (e.g. 'Computer Science'). Triggers Semantic Scholar search.",
+            },
+            "min_citations": {
+                "type": "integer",
+                "description": "Minimum citation count filter. Triggers Semantic Scholar search.",
+            },
+            "sort_by": {
+                "type": "string",
+                "enum": ["relevance", "citationCount", "publicationDate"],
+                "description": "Sort order for Semantic Scholar search. Default: relevance.",
+            },
+            "positive_ids": {
+                "type": "string",
+                "description": "Comma-separated arxiv IDs for multi-paper recommendations. For: recommend.",
+            },
+            "negative_ids": {
+                "type": "string",
+                "description": "Comma-separated arxiv IDs as negative examples. For: recommend.",
+            },
             "sort": {
                 "type": "string",
                 "enum": ["downloads", "likes", "trending"],
                 "description": (
-                    "Sort order for find_datasets and find_models. Default: downloads. "
-                    "Use 'downloads' for most-used, 'likes' for community favorites, 'trending' for recently popular."
+                    "Sort order for find_datasets and find_models. Default: downloads."
                 ),
             },
             "limit": {

agent/tools/research_tool.py

@@ -46,12 +46,22 @@ Your job: explore documentation, code examples, APIs, and repos,
 then return a concise, actionable summary. The main agent will use
 your findings to implement the actual solution.
 
+# Being up to date is critical
+
+Always prioritize finding the most current, state-of-the-art approaches.
+ML moves fast — a method from 6 months ago may already be obsolete.
+
+- Search for **recent papers** (use `hf_papers`) to find SOTA methods, models, and datasets for the task
+- Compare what you find in docs/examples against what recent papers recommend — prefer the newer approach
+- When multiple approaches exist, identify which is SOTA and why (benchmark results, adoption, recency)
+- Include in your findings: what is the current best model, dataset, and method for the task
+
 # Research methodology
 
-1. **Discovery**: Find relevant entry points — example scripts, doc pages, API endpoints
+1. **Discovery**: Find relevant entry points — example scripts, doc pages, API endpoints, **and recent papers for SOTA approaches**
 2. **Tracing**: Follow the chain from entry point to implementation detail
-3. **Analysis**: Identify patterns, current API usage, key dependencies
-4. **Synthesis**: Summarize findings in a structured format
+3. **Analysis**: Identify patterns, current API usage, key dependencies. **Compare against SOTA from recent papers**
+4. **Synthesis**: Summarize findings in a structured format, highlighting what is current best practice vs. outdated
 
 # How to use your tools
 
@@ -75,12 +85,30 @@ your findings to implement the actual solution.
   - DPO: needs "prompt", "chosen", "rejected"
   - GRPO: needs "prompt" only
 
-## Papers
-- `hf_papers`: Search papers, get details, find linked datasets/models
+## Papers & citations
+- `hf_papers(operation="search", query=...)`: Search papers (HF-tuned for ML)
+- `hf_papers(operation="search", query=..., min_citations=50, sort_by="citationCount")`: Find highly-cited papers via Semantic Scholar
+- `hf_papers(operation="search", query=..., date_from="2024-01-01")`: Search with date filter
+- `hf_papers(operation="paper_details", arxiv_id=...)`: Metadata, citations, TL;DR
+- `hf_papers(operation="citation_graph", arxiv_id=...)`: References + citations with influence flags and intents
+- `hf_papers(operation="snippet_search", query=...)`: Semantic search across 12M+ full-text paper passages
+- `hf_papers(operation="recommend", arxiv_id=...)`: Find related papers
 
 ## Hub repo inspection
 - `hf_repo_files`: List/read files in any HF repo (model, dataset, space)
 
+# Paper analysis checklist
+
+When reading a paper, always extract:
+- **Key claims**: What does the paper propose or demonstrate?
+- **Methodology**: Architecture, training setup, key techniques
+- **Results**: Benchmark numbers, comparisons to baselines
+- **Limitations**: What the authors acknowledge or what seems missing
+
+Use `citation_graph` to trace influence: check what a breakthrough paper cites (foundations)
+and who cites it (impact and extensions). Use `snippet_search` to verify claims across
+papers (e.g., "does method X consistently outperform Y?").
+
 # Correct research pattern for ML tasks
 
 ```
@@ -101,11 +129,12 @@ hf_inspect_dataset({"dataset": "org/name", "split": "train", "sample_rows": 3})
 # Output format
 
 Your output MUST include:
+- **SOTA landscape**: Current best models, datasets, and methods for the task (from recent papers). Flag anything outdated.
 - **Key findings**: The most important things you discovered (current API usage, working patterns)
 - **Essential references**: Specific file paths, URLs, function names, doc sections, code snippets
   that the main agent should use directly
 - **Code patterns**: Key imports, configurations, and usage patterns from working examples
-- **Recommendations**: What to do next based on your findings
+- **Recommendations**: What to do next based on your findings, preferring SOTA approaches
 
 Be concise. Your output goes into another agent's context — every token counts.
 Aim for 500-1500 words max. Include actual code snippets from examples you read,

agent/tools/sandbox_tool.py

@@ -12,7 +12,6 @@ a cpu-basic sandbox is auto-created (no approval needed).
 from __future__ import annotations
 
 import asyncio
-import shlex
 import threading
 from typing import Any
 
@@ -49,9 +48,15 @@ async def resolve_sandbox_script(
     if not sandbox or not _looks_like_path(script):
         return None, None
     try:
-        result = await asyncio.to_thread(sandbox.bash, f"cat {shlex.quote(script)}")
+        # Use the read endpoint instead of bash("cat ...") which truncates at 25KB.
+        result = await asyncio.to_thread(sandbox.read, script, limit=100_000)
         if result.success and result.output:
-            return result.output, None
+            # Strip line number prefixes (read returns "N\tcontent" format)
+            lines = []
+            for line in result.output.split("\n"):
+                parts = line.split("\t", 1)
+                lines.append(parts[1] if len(parts) == 2 else line)
+            return "\n".join(lines), None
         return None, f"Failed to read {script} from sandbox: {result.error}"
     except Exception as e:
         return None, f"Failed to read {script} from sandbox: {e}"

backend/main.py

@@ -12,7 +12,8 @@ from fastapi.staticfiles import StaticFiles
 from routes.agent import router as agent_router
 from routes.auth import router as auth_router
 
-load_dotenv()
+# Load .env from project root (parent directory)
+load_dotenv(Path(__file__).parent.parent / ".env")
 
 # Configure logging
 logging.basicConfig(

backend/models.py

@@ -54,6 +54,12 @@ class SubmitRequest(BaseModel):
     text: str
 
 
+class TruncateRequest(BaseModel):
+    """Request to truncate conversation history to before a specific user message."""
+
+    user_message_index: int
+
+
 class SessionResponse(BaseModel):
     """Response when creating a new session."""
 

backend/routes/agent.py

@@ -7,6 +7,7 @@ dependency. In dev mode (no OAUTH_CLIENT_ID), auth is bypassed automatically.
 import asyncio
 import json
 import logging
+import os
 from typing import Any
 
 from dependencies import get_current_user
@@ -25,6 +26,7 @@ from models import (
     SessionInfo,
     SessionResponse,
     SubmitRequest,
+    TruncateRequest,
 )
 from session_manager import MAX_SESSIONS, SessionCapacityError, session_manager
 
@@ -205,13 +207,15 @@ async def create_session(
 
     Returns 503 if the server or user has reached the session limit.
     """
-    # Extract the user's HF token (Bearer header or HttpOnly cookie)
+    # Extract the user's HF token (Bearer header, HttpOnly cookie, or env var)
     hf_token = None
     auth_header = request.headers.get("Authorization", "")
     if auth_header.startswith("Bearer "):
         hf_token = auth_header[7:]
     if not hf_token:
         hf_token = request.cookies.get("hf_access_token")
+    if not hf_token:
+        hf_token = os.environ.get("HF_TOKEN")
 
     try:
         session_id = await session_manager.create_session(
@@ -435,6 +439,18 @@ async def undo_session(session_id: str, user: dict = Depends(get_current_user))
     return {"status": "undo_requested", "session_id": session_id}
 
 
+@router.post("/truncate/{session_id}")
+async def truncate_session(
+    session_id: str, body: TruncateRequest, user: dict = Depends(get_current_user)
+) -> dict:
+    """Truncate conversation to before a specific user message."""
+    _check_session_access(session_id, user)
+    success = await session_manager.truncate(session_id, body.user_message_index)
+    if not success:
+        raise HTTPException(status_code=404, detail="Session not found, inactive, or message index out of range")
+    return {"status": "truncated", "session_id": session_id}
+
+
 @router.post("/compact/{session_id}")
 async def compact_session(
     session_id: str, user: dict = Depends(get_current_user)

backend/routes/auth.py

@@ -10,7 +10,7 @@ import time
 from urllib.parse import urlencode
 
 import httpx
-from dependencies import AUTH_ENABLED, get_current_user
+from dependencies import AUTH_ENABLED, check_org_membership, get_current_user
 from fastapi import APIRouter, Depends, HTTPException, Request
 from fastapi.responses import RedirectResponse
 
@@ -169,3 +169,20 @@ async def get_me(user: dict = Depends(get_current_user)) -> dict:
     Uses the shared auth dependency which handles cookie + Bearer token.
     """
     return user
+
+
+ORG_NAME = "ml-agent-explorers"
+
+
+@router.get("/org-membership")
+async def org_membership(
+    request: Request, user: dict = Depends(get_current_user)
+) -> dict:
+    """Check if the authenticated user belongs to the ml-agent-explorers org."""
+    if not AUTH_ENABLED:
+        return {"is_member": True}
+    token = request.cookies.get("hf_access_token") or ""
+    if not token:
+        return {"is_member": False}
+    is_member = await check_org_membership(token, ORG_NAME)
+    return {"is_member": is_member}

backend/session_manager.py

@@ -319,6 +319,14 @@ class SessionManager:
         operation = Operation(op_type=OpType.UNDO)
         return await self.submit(session_id, operation)
 
+    async def truncate(self, session_id: str, user_message_index: int) -> bool:
+        """Truncate conversation to before a specific user message (direct, no queue)."""
+        async with self._lock:
+            agent_session = self.sessions.get(session_id)
+        if not agent_session or not agent_session.is_active:
+            return False
+        return agent_session.session.context_manager.truncate_to_user_message(user_message_index)
+
     async def compact(self, session_id: str) -> bool:
         """Compact context in a session."""
         operation = Operation(op_type=OpType.COMPACT)

frontend/package-lock.json

@@ -130,7 +130,6 @@
       "integrity": "sha512-H3mcG6ZDLTlYfaSNi0iOKkigqMFvkTKlGUYlD8GW7nNOYRrevuA46iTypPyv+06V3fEmvvazfntkBU34L0azAw==",
       "dev": true,
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "@babel/code-frame": "^7.28.6",
         "@babel/generator": "^7.28.6",
@@ -447,7 +446,6 @@
       "resolved": "https://registry.npmjs.org/@emotion/react/-/react-11.14.0.tgz",
       "integrity": "sha512-O000MLDBDdk/EohJPFUqvnp4qnHeYkVP5B0xEG0D/L7cOKP9kefu2DXn8dj74cQfsEzUqh+sr1RzFqiL1o+PpA==",
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "@babel/runtime": "^7.18.3",
         "@emotion/babel-plugin": "^11.13.5",
@@ -491,7 +489,6 @@
       "resolved": "https://registry.npmjs.org/@emotion/styled/-/styled-11.14.1.tgz",
       "integrity": "sha512-qEEJt42DuToa3gurlH4Qqc1kVpNq8wO8cJtDzU46TjlzWjDlsVyevtYCRijVq3SrHsROS+gVQ8Fnea108GnKzw==",
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "@babel/runtime": "^7.18.3",
         "@emotion/babel-plugin": "^11.13.5",
@@ -1224,7 +1221,6 @@
       "resolved": "https://registry.npmjs.org/@mui/material/-/material-6.5.0.tgz",
       "integrity": "sha512-yjvtXoFcrPLGtgKRxFaH6OQPtcLPhkloC0BML6rBG5UeldR0nPULR/2E2BfXdo5JNV7j7lOzrrLX2Qf/iSidow==",
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "@babel/runtime": "^7.26.0",
         "@mui/core-downloads-tracker": "^6.5.0",
@@ -1919,7 +1915,6 @@
       "resolved": "https://registry.npmjs.org/@types/react/-/react-18.3.27.tgz",
       "integrity": "sha512-cisd7gxkzjBKU2GgdYrTdtQx1SORymWyaAFhaxQPK9bYO9ot3Y5OikQRvY0VYQtvwjeQnizCINJAenh/V7MK2w==",
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "@types/prop-types": "*",
         "csstype": "^3.2.2"
@@ -2005,7 +2000,6 @@
       "integrity": "sha512-npiaib8XzbjtzS2N4HlqPvlpxpmZ14FjSJrteZpPxGUaYPlvhzlzUZ4mZyABo0EFrOWnvyd0Xxroq//hKhtAWg==",
       "dev": true,
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "@typescript-eslint/scope-manager": "8.53.0",
         "@typescript-eslint/types": "8.53.0",
@@ -2272,7 +2266,6 @@
       "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
       "dev": true,
       "license": "MIT",
-      "peer": true,
       "bin": {
         "acorn": "bin/acorn"
       },
@@ -2421,7 +2414,6 @@
         }
       ],
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "baseline-browser-mapping": "^2.9.0",
         "caniuse-lite": "^1.0.30001759",
@@ -2774,7 +2766,6 @@
       "integrity": "sha512-LEyamqS7W5HB3ujJyvi0HQK/dtVINZvd5mAAp9eT5S/ujByGjiZLCzPcHVzuXbpJDJF/cxwHlfceVUDZ2lnSTw==",
       "dev": true,
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "@eslint-community/eslint-utils": "^4.8.0",
         "@eslint-community/regexpp": "^4.12.1",
@@ -4673,7 +4664,6 @@
       "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
       "dev": true,
       "license": "MIT",
-      "peer": true,
       "engines": {
         "node": ">=12"
       },
@@ -4771,7 +4761,6 @@
       "resolved": "https://registry.npmjs.org/react/-/react-18.3.1.tgz",
       "integrity": "sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==",
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "loose-envify": "^1.1.0"
       },
@@ -4784,7 +4773,6 @@
       "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-18.3.1.tgz",
       "integrity": "sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==",
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "loose-envify": "^1.1.0",
         "scheduler": "^0.23.2"
@@ -5269,7 +5257,6 @@
       "integrity": "sha512-hjcS1mhfuyi4WW8IWtjP7brDrG2cuDZukyrYrSauoXGNgx0S7zceP07adYkJycEr56BOUTNPzbInooiN3fn1qw==",
       "dev": true,
       "license": "Apache-2.0",
-      "peer": true,
       "bin": {
         "tsc": "bin/tsc",
         "tsserver": "bin/tsserver"
@@ -5435,7 +5422,6 @@
       "resolved": "https://registry.npmjs.org/use-sync-external-store/-/use-sync-external-store-1.6.0.tgz",
       "integrity": "sha512-Pp6GSwGP/NrPIrxVFAIkOQeyw8lFenOHijQWkUTrDvrF4ALqylP2C/KCkeS9dpUM3KvYRQhna5vt7IL95+ZQ9w==",
       "license": "MIT",
-      "peer": true,
       "peerDependencies": {
         "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0"
       }
@@ -5474,7 +5460,6 @@
       "integrity": "sha512-o5a9xKjbtuhY6Bi5S3+HvbRERmouabWbyUcpXXUA1u+GNUKoROi9byOJ8M0nHbHYHkYICiMlqxkg1KkYmm25Sw==",
       "dev": true,
       "license": "MIT",
-      "peer": true,
       "dependencies": {
         "esbuild": "^0.21.3",
         "postcss": "^8.4.43",

frontend/src/components/Chat/ActivityStatusBar.tsx

@@ -20,11 +20,90 @@ const TOOL_LABELS: Record<string, string> = {
   research: 'Researching',
 };
 
+/** Format raw research log into a clean status label. */
+function formatResearchStatus(raw: string): string {
+  const s = raw.replace(/^▸\s*/, '');
+  const jsonStart = s.indexOf('{');
+  const toolName = jsonStart > 0 ? s.slice(0, jsonStart).trim() : s.trim();
+  let args: Record<string, string> = {};
+  if (jsonStart > 0) {
+    const jsonStr = s.slice(jsonStart);
+    try {
+      const parsed = JSON.parse(jsonStr);
+      for (const [k, v] of Object.entries(parsed)) {
+        if (typeof v === 'string') args[k] = v;
+      }
+    } catch {
+      // JSON is likely truncated — extract complete "key": "value" pairs
+      for (const m of jsonStr.matchAll(/"(\w+)":\s*"([^"]*)"/g)) {
+        args[m[1]] = m[2];
+      }
+      // Also try to extract a truncated value for known keys if not found yet
+      if (!args.query && !args.arxiv_id) {
+        const partial = jsonStr.match(/"(query|arxiv_id)":\s*"([^"]*)/);
+        if (partial) args[partial[1]] = partial[2];
+      }
+    }
+  }
+
+  if (toolName === 'github_find_examples') {
+    const d = (args.keyword) || (args.repo);
+    return d ? `Finding examples: ${d}` : 'Finding examples';
+  }
+  if (toolName === 'github_read_file') {
+    const f = ((args.path) || '').split('/').pop();
+    return f ? `Reading ${f}` : 'Reading file';
+  }
+  if (toolName === 'explore_hf_docs') {
+    const d = (args.endpoint) || (args.query);
+    return d ? `Exploring docs: ${d}` : 'Exploring docs';
+  }
+  if (toolName === 'fetch_hf_docs') {
+    const p = ((args.url) || '').split('/').pop()?.replace(/\.md$/, '');
+    return p ? `Reading docs: ${p}` : 'Fetching docs';
+  }
+  if (toolName === 'hf_inspect_dataset') {
+    const d = args.dataset as string;
+    return d ? `Inspecting dataset: ${d}` : 'Inspecting dataset';
+  }
+  if (toolName === 'hf_papers') {
+    const op = args.operation as string;
+    const detail = (args.query) || (args.arxiv_id) || (args.positive_ids);
+    const opLabels: Record<string, string> = {
+      trending: 'Browsing trending papers',
+      search: 'Searching papers',
+      paper_details: 'Reading paper details',
+      read_paper: 'Reading paper',
+      citation_graph: 'Tracing citations',
+      snippet_search: 'Searching paper passages',
+      recommend: 'Finding similar papers',
+      find_datasets: 'Finding paper datasets',
+      find_models: 'Finding paper models',
+      find_collections: 'Finding paper collections',
+      find_all_resources: 'Finding paper resources',
+    };
+    const base = (op && opLabels[op]) || 'Searching papers';
+    return detail ? `${base}: ${detail}` : base;
+  }
+  if (toolName === 'find_hf_api') {
+    const d = (args.query) || (args.tag);
+    return d ? `Finding API: ${d}` : 'Finding API endpoints';
+  }
+  if (toolName === 'hf_repo_files') {
+    const d = (args.repo_id) || (args.repo);
+    return d ? `Reading ${d} files` : 'Reading repo files';
+  }
+  return 'Researching';
+}
+
 function statusLabel(status: ActivityStatus): string {
   switch (status.type) {
     case 'thinking': return 'Thinking';
     case 'streaming': return 'Writing';
     case 'tool': {
+      if (status.toolName === 'research' && status.description) {
+        return formatResearchStatus(status.description);
+      }
       const base = status.description || TOOL_LABELS[status.toolName] || `Running ${status.toolName}`;
       if (status.toolName === 'bash' && status.description && /install/i.test(status.description)) {
         return `${base} — this can take a few minutes, sit tight`;
@@ -32,6 +111,7 @@ function statusLabel(status: ActivityStatus): string {
       return base;
     }
     case 'waiting-approval': return 'Waiting for approval';
+    case 'cancelled': return 'What should the agent do instead?';
     default: return '';
   }
 }
@@ -59,7 +139,7 @@ export default function ActivityStatusBar() {
           animation: `${shimmer} 4s ease-in-out infinite`,
         }}
       >
-        {label}…
+        {label}{activityStatus.type !== 'cancelled' && '…'}
       </Typography>
     </Box>
   );

frontend/src/components/Chat/ChatInput.tsx

@@ -2,7 +2,7 @@ import { useState, useCallback, useEffect, useRef, KeyboardEvent } from 'react';
 import { Box, TextField, IconButton, CircularProgress, Typography, Menu, MenuItem, ListItemIcon, ListItemText, Chip } from '@mui/material';
 import ArrowUpwardIcon from '@mui/icons-material/ArrowUpward';
 import ArrowDropDownIcon from '@mui/icons-material/ArrowDropDown';
-import CloseIcon from '@mui/icons-material/Close';
+import StopIcon from '@mui/icons-material/Stop';
 import { apiFetch } from '@/utils/api';
 
 // Model configuration
@@ -67,7 +67,6 @@ interface ChatInputProps {
 
 export default function ChatInput({ onSend, onStop, isProcessing = false, disabled = false, placeholder = 'Ask anything...' }: ChatInputProps) {
   const [input, setInput] = useState('');
-  const [stopHovered, setStopHovered] = useState(false);
   const inputRef = useRef<HTMLTextAreaElement>(null);
   const [selectedModelId, setSelectedModelId] = useState<string>(() => {
     try {
@@ -207,20 +206,23 @@ export default function ChatInput({ onSend, onStop, isProcessing = false, disabl
           {isProcessing ? (
             <IconButton
               onClick={onStop}
-              onMouseEnter={() => setStopHovered(true)}
-              onMouseLeave={() => setStopHovered(false)}
               sx={{
                 mt: 1,
-                p: 1,
+                p: 1.5,
                 borderRadius: '10px',
-                color: stopHovered ? 'var(--accent-yellow)' : 'var(--muted-text)',
+                color: 'var(--muted-text)',
                 transition: 'all 0.2s',
+                position: 'relative',
                 '&:hover': {
                   bgcolor: 'var(--hover-bg)',
+                  color: 'var(--accent-red)',
                 },
               }}
             >
-              {stopHovered ? <CloseIcon fontSize="small" /> : <CircularProgress size={20} color="inherit" />}
+              <Box sx={{ position: 'relative', display: 'flex', alignItems: 'center', justifyContent: 'center' }}>
+                <CircularProgress size={28} thickness={3} sx={{ color: 'inherit', position: 'absolute' }} />
+                <StopIcon sx={{ fontSize: 16 }} />
+              </Box>
             </IconButton>
           ) : (
             <IconButton

frontend/src/components/Chat/MessageBubble.tsx

@@ -6,6 +6,7 @@ interface MessageBubbleProps {
   message: UIMessage;
   isLastTurn?: boolean;
   onUndoTurn?: () => void;
+  onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise<void>;
   isProcessing?: boolean;
   isStreaming?: boolean;
   approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
@@ -15,6 +16,7 @@ export default function MessageBubble({
   message,
   isLastTurn = false,
   onUndoTurn,
+  onEditAndRegenerate,
   isProcessing = false,
   isStreaming = false,
   approveTools,
@@ -25,6 +27,7 @@ export default function MessageBubble({
         message={message}
         isLastTurn={isLastTurn}
         onUndoTurn={onUndoTurn}
+        onEditAndRegenerate={onEditAndRegenerate}
         isProcessing={isProcessing}
       />
     );

frontend/src/components/Chat/MessageList.tsx

@@ -10,6 +10,7 @@ interface MessageListProps {
   isProcessing: boolean;
   approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
   onUndoLastTurn: () => void | Promise<void>;
+  onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise<void>;
 }
 
 function getGreeting(): string {
@@ -56,7 +57,7 @@ function WelcomeGreeting() {
   );
 }
 
-export default function MessageList({ messages, isProcessing, approveTools, onUndoLastTurn }: MessageListProps) {
+export default function MessageList({ messages, isProcessing, approveTools, onUndoLastTurn, onEditAndRegenerate }: MessageListProps) {
   const scrollContainerRef = useRef<HTMLDivElement>(null);
   const stickToBottom = useRef(true);
 
@@ -135,6 +136,7 @@ export default function MessageList({ messages, isProcessing, approveTools, onUn
               message={msg}
               isLastTurn={msg.id === lastUserMsgId}
               onUndoTurn={onUndoLastTurn}
+              onEditAndRegenerate={onEditAndRegenerate}
               isProcessing={isProcessing}
               isStreaming={isProcessing && msg.id === lastAssistantId}
               approveTools={approveTools}

frontend/src/components/Chat/ToolCallGroup.tsx

@@ -10,6 +10,7 @@ import BlockIcon from '@mui/icons-material/Block';
 import { useAgentStore } from '@/store/agentStore';
 import { useLayoutStore } from '@/store/layoutStore';
 import { logger } from '@/utils/logger';
+import { RESEARCH_MAX_STEPS } from '@/lib/research-store';
 import type { UIMessage } from 'ai';
 
 // ---------------------------------------------------------------------------
@@ -35,37 +36,153 @@ interface ToolCallGroupProps {
 // Research sub-steps (inline under the research tool row)
 // ---------------------------------------------------------------------------
 
+/** Hook that ticks every second while startedAt is set, returning elapsed seconds. */
+function useElapsed(startedAt: number | null): number | null {
+  const [elapsed, setElapsed] = useState<number | null>(null);
+  useEffect(() => {
+    if (startedAt === null) { setElapsed(null); return; }
+    setElapsed(Math.round((Date.now() - startedAt) / 1000));
+    const id = setInterval(() => setElapsed(Math.round((Date.now() - startedAt) / 1000)), 1000);
+    return () => clearInterval(id);
+  }, [startedAt]);
+  return elapsed;
+}
+
+/** Format token count like the CLI: "12.4k" or "800". */
+function formatTokens(tokens: number): string {
+  return tokens >= 1000 ? `${(tokens / 1000).toFixed(1)}k` : String(tokens);
+}
+
+/** Format elapsed seconds like the CLI: "18s" or "2m 5s". */
+function formatElapsed(seconds: number): string {
+  if (seconds < 60) return `${seconds}s`;
+  return `${Math.floor(seconds / 60)}m ${seconds % 60}s`;
+}
+
+/** Build the research stats chip label. */
+function researchChipLabel(
+  stats: { toolCount: number; tokenCount: number; startedAt: number | null; finalElapsed: number | null },
+  liveElapsed: number | null,
+): string | null {
+  const elapsed = stats.finalElapsed ?? liveElapsed;
+  if (elapsed === null && stats.toolCount === 0) return null;
+  const parts: string[] = [];
+  if (stats.startedAt !== null) parts.push('running');
+  if (stats.toolCount > 0) parts.push(`${stats.toolCount} tools`);
+  if (stats.tokenCount > 0) parts.push(`${formatTokens(stats.tokenCount)} tokens`);
+  if (elapsed !== null) parts.push(formatElapsed(elapsed));
+  return parts.join(' \u00B7 ');
+}
+
+/** Parse JSON args from a step string like "tool_name  {json}" (may be truncated at 80 chars). */
+function parseStepArgs(step: string): Record<string, string> {
+  const jsonStart = step.indexOf('{');
+  if (jsonStart < 0) return {};
+  const jsonStr = step.slice(jsonStart);
+  try {
+    const parsed = JSON.parse(jsonStr);
+    const result: Record<string, string> = {};
+    for (const [k, v] of Object.entries(parsed)) {
+      if (typeof v === 'string') result[k] = v;
+    }
+    return result;
+  } catch {
+    // JSON likely truncated — extract key-value pairs via regex
+    const result: Record<string, string> = {};
+    // Match complete "key": "value" pairs
+    for (const m of jsonStr.matchAll(/"(\w+)":\s*"([^"]*)"/g)) {
+      result[m[1]] = m[2];
+    }
+    // Match truncated trailing value: "key": "value... (no closing quote)
+    if (Object.keys(result).length === 0 || !result.query) {
+      const trunc = jsonStr.match(/"(\w+)":\s*"([^"]+)$/);
+      if (trunc && !result[trunc[1]]) {
+        result[trunc[1]] = trunc[2];
+      }
+    }
+    return result;
+  }
+}
+
 /** Pretty labels for research sub-agent tool calls */
-function formatResearchStep(step: string): { icon: string; label: string } {
-  if (step === 'Starting research sub-agent...') return { icon: '🔍', label: 'Starting research' };
-  if (step === 'Research complete.') return { icon: '✓', label: 'Research complete' };
-  if (step.startsWith('github_find_examples')) return { icon: '📂', label: step.replace('github_find_examples', 'Finding examples') };
+function formatResearchStep(raw: string): { label: string } {
+  // Backend sends logs like "▸ tool_name  {args}" — strip the prefix
+  const step = raw.replace(/^▸\s*/, '');
+  const args = parseStepArgs(step);
+
+  if (step.startsWith('github_find_examples')) {
+    const detail = (args.keyword) || (args.repo);
+    return { label: detail ? `Finding examples: ${detail}` : 'Finding examples' };
+  }
   if (step.startsWith('github_read_file')) {
-    const path = step.match(/\(([^)]+)\)/)?.[1] || '';
+    const path = (args.path) || '';
     const filename = path.split('/').pop() || path;
-    return { icon: '📄', label: `Reading ${filename}` };
+    return { label: filename ? `Reading ${filename}` : 'Reading file' };
+  }
+  if (step.startsWith('explore_hf_docs')) {
+    const endpoint = (args.endpoint) || (args.query);
+    return { label: endpoint ? `Exploring docs: ${endpoint}` : 'Exploring docs' };
+  }
+  if (step.startsWith('fetch_hf_docs')) {
+    const url = (args.url) || '';
+    const page = url.split('/').pop()?.replace(/\.md$/, '');
+    return { label: page ? `Reading docs: ${page}` : 'Fetching docs' };
+  }
+  if (step.startsWith('hf_inspect_dataset')) {
+    const dataset = (args.dataset);
+    return { label: dataset ? `Inspecting dataset: ${dataset}` : 'Inspecting dataset' };
   }
-  if (step.startsWith('explore_hf_docs')) return { icon: '📚', label: step.replace('explore_hf_docs', 'Exploring docs') };
-  if (step.startsWith('fetch_hf_docs')) return { icon: '📖', label: step.replace('fetch_hf_docs', 'Fetching docs') };
-  if (step.startsWith('hf_inspect_dataset')) return { icon: '🗃️', label: step.replace('hf_inspect_dataset', 'Inspecting dataset') };
-  if (step.startsWith('hf_papers')) return { icon: '📑', label: 'Searching papers' };
-  if (step.startsWith('find_hf_api')) return { icon: '🔌', label: 'Finding API endpoints' };
-  if (step.startsWith('hf_repo_files')) return { icon: '📁', label: 'Reading repo files' };
-  return { icon: '→', label: step };
+  if (step.startsWith('hf_papers')) {
+    const op = args.operation as string;
+    const detail = (args.query) || (args.arxiv_id);
+    const opLabels: Record<string, string> = {
+      trending: 'Browsing trending papers',
+      search: 'Searching papers',
+      paper_details: 'Reading paper details',
+      read_paper: 'Reading paper',
+      citation_graph: 'Tracing citations',
+      snippet_search: 'Searching paper snippets',
+      recommend: 'Finding related papers',
+      find_datasets: 'Finding paper datasets',
+      find_models: 'Finding paper models',
+      find_collections: 'Finding paper collections',
+      find_all_resources: 'Finding paper resources',
+    };
+    const base = (op && opLabels[op]) || 'Searching papers';
+    return { label: detail ? `${base}: ${detail}` : base };
+  }
+  if (step.startsWith('find_hf_api')) {
+    const detail = (args.query) || (args.tag);
+    return { label: detail ? `Finding API: ${detail}` : 'Finding API endpoints' };
+  }
+  if (step.startsWith('hf_repo_files')) {
+    const repo = (args.repo_id) || (args.repo);
+    return { label: repo ? `Reading ${repo} files` : 'Reading repo files' };
+  }
+  if (step.startsWith('read')) {
+    const path = (args.path) || '';
+    const filename = path.split('/').pop();
+    return { label: filename ? `Reading ${filename}` : 'Reading file' };
+  }
+  if (step.startsWith('bash')) {
+    const cmd = args.command as string;
+    const short = cmd && cmd.length > 40 ? cmd.slice(0, 40) + '...' : cmd;
+    return { label: short ? `Running: ${short}` : 'Running command' };
+  }
+  return { label: step.replace(/^▸\s*/, '') };
 }
 
+/** Rolling 2-line display of research sub-tool calls — hidden when complete. */
 function ResearchSteps({ steps, isRunning }: { steps: string[]; isRunning: boolean }) {
-  // Filter out the "Starting..." and "complete" meta-steps for the list
-  const toolSteps = steps.filter(
-    s => s !== 'Starting research sub-agent...' && s !== 'Research complete.',
-  );
-  if (toolSteps.length === 0) return null;
+  if (!isRunning) return null;
+  const visible = steps.slice(-RESEARCH_MAX_STEPS);
+  if (visible.length === 0) return null;
 
   return (
     <Box sx={{ pl: 4.5, pr: 1.5, pb: 1, pt: 0.25 }}>
-      {toolSteps.map((step, i) => {
-        const { icon, label } = formatResearchStep(step);
-        const isLast = i === toolSteps.length - 1;
+      {visible.map((step, i) => {
+        const { label } = formatResearchStep(step);
+        const isLast = i === visible.length - 1;
         return (
           <Stack
             key={i}
@@ -74,17 +191,16 @@ function ResearchSteps({ steps, isRunning }: { steps: string[]; isRunning: boole
             spacing={0.75}
             sx={{ py: 0.2 }}
           >
-            <Typography sx={{ fontSize: '0.65rem', lineHeight: 1, width: 14, textAlign: 'center', flexShrink: 0 }}>
-              {isLast && isRunning ? '' : icon}
-            </Typography>
-            {isLast && isRunning && (
+            {isLast ? (
               <CircularProgress size={10} thickness={5} sx={{ color: 'var(--accent-yellow)', flexShrink: 0 }} />
+            ) : (
+              <CheckCircleOutlineIcon sx={{ fontSize: 12, color: 'var(--muted-text)', flexShrink: 0 }} />
             )}
             <Typography
               sx={{
                 fontFamily: '"JetBrains Mono", ui-monospace, SFMono-Regular, monospace',
                 fontSize: '0.68rem',
-                color: isLast && isRunning ? 'var(--text)' : 'var(--muted-text)',
+                color: isLast ? 'var(--text)' : 'var(--muted-text)',
                 overflow: 'hidden',
                 textOverflow: 'ellipsis',
                 whiteSpace: 'nowrap',
@@ -132,8 +248,8 @@ function costLabel(hardware: string): string | null {
 // Visual helpers
 // ---------------------------------------------------------------------------
 
-function StatusIcon({ state, cancelled }: { state: ToolPartState; cancelled?: boolean }) {
-  if (cancelled) {
+function StatusIcon({ state, cancelled, isRejected }: { state: ToolPartState; cancelled?: boolean; isRejected?: boolean }) {
+  if (cancelled || isRejected) {
     return <BlockIcon sx={{ fontSize: 16, color: 'var(--muted-text)' }} />;
   }
   switch (state) {
@@ -397,11 +513,17 @@ function InlineApproval({
 // ---------------------------------------------------------------------------
 
 export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProps) {
-  const { setPanel, lockPanel, getJobUrl, getEditedScript } = useAgentStore();
+  const { setPanel, lockPanel, getJobUrl, getEditedScript, setJobStatus, getJobStatus, setToolError, getToolError, setToolRejected, getToolRejected } = useAgentStore();
   const researchSteps = useAgentStore(s => {
     const activeId = s.activeSessionId;
     return activeId ? (s.sessionStates[activeId]?.researchSteps) : undefined;
   }) ?? EMPTY_STEPS;
+  const researchStats = useAgentStore(s => {
+    const activeId = s.activeSessionId;
+    return activeId ? s.sessionStates[activeId]?.researchStats : undefined;
+  }) ?? { toolCount: 0, tokenCount: 0, startedAt: null, finalElapsed: null };
+  const liveElapsed = useElapsed(researchStats.startedAt);
+  const isProcessing = useAgentStore(s => s.isProcessing);
   const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
 
   // ── Batch approval state ──────────────────────────────────────────
@@ -417,6 +539,9 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
   // Track which toolCallIds we've already submitted so we can detect new approval rounds
   const submittedIdsRef = useRef<Set<string>>(new Set());
 
+  // ── Panel lock state (for auto-follow vs user-selected) ───────────
+  const [lockedToolId, setLockedToolId] = useState<string | null>(null);
+
   // Reset submission state when new (unseen) pending tools arrive — e.g. second approval round
   useEffect(() => {
     if (!isSubmitting || pendingTools.length === 0) return;
@@ -428,6 +553,35 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
     }
   }, [pendingTools, isSubmitting]);
 
+  // Clean up stale decisions for tools that are no longer pending
+  useEffect(() => {
+    const pendingIds = new Set(pendingTools.map(t => t.toolCallId));
+    const decisionIds = Object.keys(decisions);
+    const hasStale = decisionIds.some(id => !pendingIds.has(id));
+    if (hasStale) {
+      setDecisions(prev => {
+        const cleaned = { ...prev };
+        for (const id of decisionIds) {
+          if (!pendingIds.has(id)) delete cleaned[id];
+        }
+        return cleaned;
+      });
+    }
+  }, [pendingTools, decisions]);
+
+  // Persist error states when tools error
+  useEffect(() => {
+    for (const tool of tools) {
+      const currentlyHasError = tool.state === 'output-error';
+      const persistedError = getToolError(tool.toolCallId);
+
+      // Persist error state if we detect it and haven't already
+      if (currentlyHasError && !persistedError) {
+        setToolError(tool.toolCallId, true);
+      }
+    }
+  }, [tools, setToolError, getToolError]);
+
   const { scriptLabelMap, toolDisplayMap } = useMemo(() => {
     const hfJobs = tools.filter(t => t.toolName === 'hf_jobs' && (t.input as Record<string, unknown>)?.script);
     const scriptMap: Record<string, string> = {};
@@ -463,6 +617,10 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
         if (editedScript) {
           logger.log(`Sending edited script for ${toolCallId} (${editedScript.length} chars)`);
         }
+        // Mark tool as rejected if not approved
+        if (!d.approved) {
+          setToolRejected(toolCallId, true);
+        }
         return {
           tool_call_id: toolCallId,
           approved: d.approved,
@@ -482,7 +640,7 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
         setIsSubmitting(false);
       }
     },
-    [approveTools, lockPanel, getEditedScript],
+    [approveTools, lockPanel, getEditedScript, setToolRejected],
   );
 
   const handleApproveAll = useCallback(() => {
@@ -518,8 +676,8 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
     });
   }, []);
 
-  // ── Panel click handler ───────────────────────────────────────────
-  const handleClick = useCallback(
+  // ── Show tool panel (shared logic) ────────────────────────────────
+  const showToolPanel = useCallback(
     (tool: DynamicToolPart) => {
       const args = tool.input as Record<string, unknown> | undefined;
       const displayName = toolDisplayMap[tool.toolCallId] || tool.toolName;
@@ -545,7 +703,13 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
       const inputSection = args ? { content: JSON.stringify(args, null, 2), language: 'json' } : undefined;
 
       const outputText = tool.output ?? (tool.state === 'output-error' ? (tool as Record<string, unknown>).errorText : undefined);
-      if ((tool.state === 'output-available' || tool.state === 'output-error') && outputText) {
+
+      // Determine if tool is still running or has completed
+      const isRunning = tool.state === 'input-available' || tool.state === 'input-streaming' || tool.state === 'approval-responded';
+      const hasCompleted = tool.state === 'output-available' || tool.state === 'output-error' || tool.state === 'output-denied';
+
+      if (outputText) {
+        // Tool has output - show it (regardless of state)
         let language = 'text';
         const content = String(outputText);
         if (content.trim().startsWith('{') || content.trim().startsWith('[')) language = 'json';
@@ -557,14 +721,79 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
         const content = `Tool \`${tool.toolName}\` returned an error with no output message.`;
         setPanel({ title: displayName, output: { content, language: 'markdown' }, input: inputSection }, 'output');
         setRightPanelOpen(true);
-      } else if (args) {
+      } else if (hasCompleted && args) {
+        // Tool completed but has no output - show input as fallback
         setPanel({ title: displayName, output: { content: JSON.stringify(args, null, 2), language: 'json' }, input: inputSection }, 'output');
         setRightPanelOpen(true);
+      } else if (isRunning && args) {
+        // Tool is still running - show running message
+        const content = `Tool \`${tool.toolName}\` is still running...\n\nClick the input tab to view the tool arguments.`;
+        setPanel({ title: displayName, output: { content, language: 'markdown' }, input: inputSection }, 'output');
+        setRightPanelOpen(true);
+      } else if (args) {
+        const runningMessages = [
+          'Crunching numbers and herding tensors...',
+          'Teaching the model some new tricks...',
+          'Consulting the GPU oracle...',
+          'Wrangling data into submission...',
+          'Brewing a fresh batch of predictions...',
+          'Negotiating with the transformer heads...',
+          'Polishing the attention weights...',
+          'Aligning the embedding stars...',
+        ];
+        const funMsg = runningMessages[Math.floor(Math.random() * runningMessages.length)];
+        setPanel({ title: displayName, output: { content: funMsg, language: 'text' }, input: inputSection }, 'output');
+        setRightPanelOpen(true);
       }
     },
     [toolDisplayMap, setPanel, getEditedScript, setRightPanelOpen, setLeftSidebarOpen],
   );
 
+  // ── Panel click handler ───────────────────────────────────────────
+  const handleClick = useCallback(
+    (tool: DynamicToolPart) => {
+      // Toggle lock: if clicking the same tool that's already locked, unlock it
+      if (lockedToolId === tool.toolCallId) {
+        setLockedToolId(null);
+        return;
+      }
+
+      // Lock this tool
+      setLockedToolId(tool.toolCallId);
+
+      // Show the panel
+      showToolPanel(tool);
+    },
+    [lockedToolId, showToolPanel],
+  );
+
+  // ── Auto-follow currently active tool when not locked ─────────────
+  const activeToolIdRef = useRef<string | null>(null);
+
+  useEffect(() => {
+    if (lockedToolId !== null) return; // User has locked a tool, don't auto-follow
+
+    // Find the currently running tool (latest tool that's in progress)
+    const runningTool = tools.slice().reverse().find(t =>
+      t.state === 'input-available' ||
+      t.state === 'input-streaming' ||
+      t.state === 'approval-responded'
+    );
+
+    if (runningTool) {
+      // Track this as the active tool and show its panel
+      activeToolIdRef.current = runningTool.toolCallId;
+      showToolPanel(runningTool);
+    } else if (activeToolIdRef.current) {
+      // No running tool, but we were following one - check if it completed
+      const completedTool = tools.find(t => t.toolCallId === activeToolIdRef.current);
+      if (completedTool && (completedTool.state === 'output-available' || completedTool.state === 'output-error')) {
+        // The tool we were following has completed - update its panel
+        showToolPanel(completedTool);
+      }
+    }
+  }, [tools, lockedToolId, showToolPanel]);
+
   // ── Parse hf_jobs metadata from output ────────────────────────────
   function parseJobMeta(output: unknown): { jobUrl?: string; jobStatus?: string } {
     if (typeof output !== 'string') return {};
@@ -651,25 +880,46 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
           const clickable =
             state === 'output-available' ||
             state === 'output-error' ||
-            !!tool.input;
+            !!tool.input ||
+            (!isProcessing && (state === 'input-available' || state === 'input-streaming'));
           const localDecision = decisions[tool.toolCallId];
 
           const cancelled = isCancelledTool(tool);
-          const displayState = isPending && localDecision
-            ? (localDecision.approved ? 'input-available' : 'output-denied')
-            : state;
-          const label = cancelled ? 'cancelled' : statusLabel(displayState as ToolPartState);
+          const currentlyHasError = state === 'output-error';
+          const persistedError = getToolError(tool.toolCallId);
+          const persistedRejection = getToolRejected(tool.toolCallId);
+
+          // Stale in-progress tools after page reload: treat as completed
+          const stale = !isProcessing && (state === 'input-available' || state === 'input-streaming');
+          const displayState = stale ? 'output-available'
+            : isPending && localDecision
+              ? (localDecision.approved ? 'input-available' : 'output-denied')
+              : state;
+          const isRejected = displayState === 'output-denied' || persistedRejection;
+          const hasError = (persistedError || currentlyHasError) && !isRejected;
+          const label = cancelled ? 'cancelled'
+            : isRejected ? 'rejected'
+            : hasError ? 'error'
+            : statusLabel(displayState as ToolPartState);
 
           // Parse job metadata from hf_jobs output and store
           const jobUrlFromStore = tool.toolName === 'hf_jobs' ? getJobUrl(tool.toolCallId) : undefined;
-          const jobMetaFromOutput = tool.toolName === 'hf_jobs' && tool.state === 'output-available'
-            ? parseJobMeta(tool.output)
+          const jobStatusFromStore = tool.toolName === 'hf_jobs' ? getJobStatus(tool.toolCallId) : undefined;
+
+          const jobMetaFromOutput = tool.toolName === 'hf_jobs' && (tool.output || (tool as Record<string, unknown>).errorText)
+            ? parseJobMeta(tool.output ?? (tool as Record<string, unknown>).errorText)
             : {};
-          
-          // Combine job URL from store (available immediately) with output metadata (available at completion)
+
+          // Store job status if we just parsed it and don't have it stored yet
+          if (tool.toolName === 'hf_jobs' && jobMetaFromOutput.jobStatus && !jobStatusFromStore) {
+            setJobStatus(tool.toolCallId, jobMetaFromOutput.jobStatus);
+          }
+
+          // Combine job URL and status from store (persisted) with output metadata (freshly parsed)
+          // Prefer stored values to ensure they persist across renders
           const jobMeta = {
             jobUrl: jobUrlFromStore || jobMetaFromOutput.jobUrl,
-            jobStatus: jobMetaFromOutput.jobStatus,
+            jobStatus: jobStatusFromStore || jobMetaFromOutput.jobStatus,
           };
 
           return (
@@ -685,15 +935,20 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
                   py: 1,
                   cursor: isPending ? 'default' : clickable ? 'pointer' : 'default',
                   transition: 'background-color 0.15s',
+                  bgcolor: lockedToolId === tool.toolCallId ? 'var(--hover-bg)' : 'transparent',
+                  borderLeft: lockedToolId === tool.toolCallId ? '3px solid var(--accent-yellow)' : '3px solid transparent',
                   '&:hover': clickable && !isPending ? { bgcolor: 'var(--hover-bg)' } : {},
                 }}
               >
                 <StatusIcon
                   cancelled={cancelled}
+                  isRejected={isRejected}
                   state={
-                    (tool.toolName === 'hf_jobs' && jobMeta.jobStatus && ['ERROR', 'FAILED', 'CANCELLED'].includes(jobMeta.jobStatus) && displayState === 'output-available')
+                    hasError
                       ? 'output-error'
-                      : displayState as ToolPartState
+                      : ((tool.toolName === 'hf_jobs' && jobMeta.jobStatus && ['ERROR', 'FAILED', 'CANCELLED'].includes(jobMeta.jobStatus))
+                        ? 'output-error'
+                        : displayState as ToolPartState)
                   }
                 />
 
@@ -715,23 +970,37 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
                 </Typography>
 
                 {/* Status chip (non hf_jobs, or hf_jobs without final status) */}
-                {label && !(tool.toolName === 'hf_jobs' && jobMeta.jobStatus) && (
-                  <Chip
-                    label={label}
-                    size="small"
-                    sx={{
-                      height: 20,
-                      fontSize: '0.65rem',
-                      fontWeight: 600,
-                      bgcolor: cancelled ? 'rgba(255,255,255,0.05)'
-                        : displayState === 'output-error' ? 'rgba(224,90,79,0.12)'
-                        : displayState === 'output-denied' ? 'rgba(255,255,255,0.05)'
-                        : 'var(--accent-yellow-weak)',
-                      color: cancelled ? 'var(--muted-text)' : statusColor(displayState as ToolPartState),
-                      letterSpacing: '0.03em',
-                    }}
-                  />
-                )}
+                {(() => {
+                  // Research tool: override chip label with live stats (but not if cancelled/done)
+                  const researchDone = cancelled || state === 'output-available' || state === 'output-error' || state === 'output-denied';
+                  const researchLabel = tool.toolName === 'research' && !researchDone
+                    ? researchChipLabel(researchStats, liveElapsed)
+                    : (tool.toolName === 'research' && researchDone && researchStats.finalElapsed !== null)
+                      ? researchChipLabel({ ...researchStats, startedAt: null }, null)
+                      : null;
+                  const chipLabel = researchLabel || label;
+                  if (!chipLabel || (tool.toolName === 'hf_jobs' && jobMeta.jobStatus)) return null;
+
+                  return (
+                    <Chip
+                      label={chipLabel}
+                      size="small"
+                      sx={{
+                        height: 20,
+                        fontSize: '0.65rem',
+                        fontWeight: 600,
+                        bgcolor: (cancelled || isRejected) ? 'rgba(255,255,255,0.05)'
+                          : hasError ? 'rgba(224,90,79,0.12)'
+                          : (researchLabel && displayState === 'output-available') ? 'rgba(47,204,113,0.12)'
+                          : 'var(--accent-yellow-weak)',
+                        color: (cancelled || isRejected) ? 'var(--muted-text)'
+                          : hasError ? 'var(--accent-red)'
+                          : statusColor(displayState as ToolPartState),
+                        letterSpacing: '0.03em',
+                      }}
+                    />
+                  );
+                })()}
 
                 {/* HF Jobs: final status chip from job metadata */}
                 {tool.toolName === 'hf_jobs' && jobMeta.jobStatus && (
@@ -785,11 +1054,11 @@ export default function ToolCallGroup({ tools, approveTools }: ToolCallGroupProp
                 )}
               </Stack>
 
-              {/* Research sub-agent steps */}
-              {tool.toolName === 'research' && researchSteps.length > 0 && (
+              {/* Research sub-agent rolling steps (visible only while running) */}
+              {tool.toolName === 'research' && !cancelled && state !== 'output-available' && state !== 'output-error' && state !== 'output-denied' && (
                 <ResearchSteps
                   steps={researchSteps}
-                  isRunning={state === 'input-streaming' || state === 'input-available'}
+                  isRunning={researchStats.startedAt !== null}
                 />
               )}
 

frontend/src/components/Chat/UserMessage.tsx

@@ -1,5 +1,8 @@
-import { Box, Stack, Typography, IconButton, Tooltip } from '@mui/material';
+import { useState, useRef, useEffect } from 'react';
+import { Box, Stack, Typography, IconButton, Tooltip, TextField } from '@mui/material';
 import CloseIcon from '@mui/icons-material/Close';
+import EditIcon from '@mui/icons-material/Edit';
+import CheckIcon from '@mui/icons-material/Check';
 import type { UIMessage } from 'ai';
 import type { MessageMeta } from '@/types/agent';
 
@@ -7,6 +10,7 @@ interface UserMessageProps {
   message: UIMessage;
   isLastTurn?: boolean;
   onUndoTurn?: () => void;
+  onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise<void>;
   isProcessing?: boolean;
 }
 
@@ -21,14 +25,57 @@ export default function UserMessage({
   message,
   isLastTurn = false,
   onUndoTurn,
+  onEditAndRegenerate,
   isProcessing = false,
 }: UserMessageProps) {
   const showUndo = isLastTurn && !isProcessing && !!onUndoTurn;
+  const showEdit = !isProcessing && !!onEditAndRegenerate;
   const text = extractText(message);
   const meta = message.metadata as MessageMeta | undefined;
   const timeStr = meta?.createdAt
     ? new Date(meta.createdAt).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })
     : null;
+
+  const [isEditing, setIsEditing] = useState(false);
+  const [editText, setEditText] = useState(text);
+  const inputRef = useRef<HTMLTextAreaElement>(null);
+
+  useEffect(() => {
+    if (isEditing && inputRef.current) {
+      inputRef.current.focus();
+      inputRef.current.selectionStart = inputRef.current.value.length;
+    }
+  }, [isEditing]);
+
+  const handleStartEdit = () => {
+    setEditText(text);
+    setIsEditing(true);
+  };
+
+  const handleConfirmEdit = () => {
+    const trimmed = editText.trim();
+    if (!trimmed || trimmed === text) {
+      setIsEditing(false);
+      return;
+    }
+    setIsEditing(false);
+    onEditAndRegenerate?.(message.id, trimmed);
+  };
+
+  const handleCancelEdit = () => {
+    setIsEditing(false);
+    setEditText(text);
+  };
+
+  const handleKeyDown = (e: React.KeyboardEvent) => {
+    if (e.key === 'Enter' && !e.shiftKey) {
+      e.preventDefault();
+      handleConfirmEdit();
+    } else if (e.key === 'Escape') {
+      handleCancelEdit();
+    }
+  };
+
   return (
     <Stack
       direction="row"
@@ -36,35 +83,56 @@ export default function UserMessage({
       justifyContent="flex-end"
       alignItems="flex-start"
       sx={{
-        '& .undo-btn': {
+        '& .action-btn': {
           opacity: 0,
           transition: 'opacity 0.15s ease',
         },
-        '&:hover .undo-btn': {
+        '&:hover .action-btn': {
           opacity: 1,
         },
       }}
     >
-      {showUndo && (
-        <Box className="undo-btn" sx={{ display: 'flex', alignItems: 'center', mt: 0.75 }}>
-          <Tooltip title="Remove this turn" placement="left">
-            <IconButton
-              onClick={onUndoTurn}
-              size="small"
-              sx={{
-                width: 24,
-                height: 24,
-                color: 'var(--muted-text)',
-                '&:hover': {
-                  color: 'var(--accent-red)',
-                  bgcolor: 'rgba(244,67,54,0.08)',
-                },
-              }}
-            >
-              <CloseIcon sx={{ fontSize: 14 }} />
-            </IconButton>
-          </Tooltip>
-        </Box>
+      {!isEditing && (showUndo || showEdit) && (
+        <Stack className="action-btn" direction="row" spacing={0.25} sx={{ mt: 0.75 }}>
+          {showEdit && (
+            <Tooltip title="Edit & regenerate" placement="left">
+              <IconButton
+                onClick={handleStartEdit}
+                size="small"
+                sx={{
+                  width: 24,
+                  height: 24,
+                  color: 'var(--muted-text)',
+                  '&:hover': {
+                    color: 'var(--accent-yellow)',
+                    bgcolor: 'rgba(255,157,0,0.08)',
+                  },
+                }}
+              >
+                <EditIcon sx={{ fontSize: 14 }} />
+              </IconButton>
+            </Tooltip>
+          )}
+          {showUndo && (
+            <Tooltip title="Remove this turn" placement="left">
+              <IconButton
+                onClick={onUndoTurn}
+                size="small"
+                sx={{
+                  width: 24,
+                  height: 24,
+                  color: 'var(--muted-text)',
+                  '&:hover': {
+                    color: 'var(--accent-red)',
+                    bgcolor: 'rgba(244,67,54,0.08)',
+                  },
+                }}
+              >
+                <CloseIcon sx={{ fontSize: 14 }} />
+              </IconButton>
+            </Tooltip>
+          )}
+        </Stack>
       )}
 
       <Box
@@ -78,20 +146,66 @@ export default function UserMessage({
           border: '1px solid var(--border)',
         }}
       >
-        <Typography
-          variant="body1"
-          sx={{
-            fontSize: '0.925rem',
-            lineHeight: 1.65,
-            color: 'var(--text)',
-            whiteSpace: 'pre-wrap',
-            wordBreak: 'break-word',
-          }}
-        >
-          {text}
-        </Typography>
+        {isEditing ? (
+          <Stack spacing={1}>
+            <TextField
+              inputRef={inputRef}
+              multiline
+              fullWidth
+              value={editText}
+              onChange={(e) => setEditText(e.target.value)}
+              onKeyDown={handleKeyDown}
+              variant="outlined"
+              size="small"
+              sx={{
+                '& .MuiOutlinedInput-root': {
+                  fontFamily: 'inherit',
+                  fontSize: '0.925rem',
+                  lineHeight: 1.65,
+                  color: 'var(--text)',
+                  '& fieldset': { borderColor: 'var(--accent-yellow)', borderWidth: 1.5 },
+                  '&:hover fieldset': { borderColor: 'var(--accent-yellow)' },
+                  '&.Mui-focused fieldset': { borderColor: 'var(--accent-yellow)' },
+                },
+              }}
+            />
+            <Stack direction="row" spacing={0.5} justifyContent="flex-end">
+              <Tooltip title="Cancel (Esc)">
+                <IconButton
+                  onClick={handleCancelEdit}
+                  size="small"
+                  sx={{ color: 'var(--muted-text)', '&:hover': { color: 'var(--accent-red)' } }}
+                >
+                  <CloseIcon sx={{ fontSize: 16 }} />
+                </IconButton>
+              </Tooltip>
+              <Tooltip title="Confirm (Enter)">
+                <IconButton
+                  onClick={handleConfirmEdit}
+                  size="small"
+                  sx={{ color: 'var(--accent-green)', '&:hover': { bgcolor: 'rgba(47,204,113,0.1)' } }}
+                >
+                  <CheckIcon sx={{ fontSize: 16 }} />
+                </IconButton>
+              </Tooltip>
+            </Stack>
+          </Stack>
+        ) : (
+          <Typography
+            variant="body1"
+            sx={{
+              fontSize: '0.925rem',
+              lineHeight: 1.65,
+              color: 'var(--text)',
+              whiteSpace: 'pre-wrap',
+              wordBreak: 'break-word',
+            }}
+          >
+            {text}
+          </Typography>
+        )}
 
-        {timeStr && (
+        {timeStr && !isEditing && (
           <Typography
             variant="caption"
             sx={{ color: 'var(--muted-text)', mt: 0.5, display: 'block', textAlign: 'right', fontSize: '0.7rem' }}

frontend/src/components/SessionChat.tsx

@@ -5,7 +5,7 @@
  * runs — processing events — but only the active session renders visible
  * UI (MessageList + ChatInput).
  */
-import { useCallback, useEffect, useState } from 'react';
+import { useCallback, useEffect } from 'react';
 import { useAgentChat } from '@/hooks/useAgentChat';
 import { useAgentStore } from '@/store/agentStore';
 import { useSessionStore } from '@/store/sessionStore';
@@ -24,9 +24,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
   const { isConnected, isProcessing, activityStatus, updateSession } = useAgentStore();
   const { updateSessionTitle } = useSessionStore();
 
-  const [wasCancelled, setWasCancelled] = useState(false);
-
-  const { messages, sendMessage, stop, status, undoLastTurn, approveTools } = useAgentChat({
+  const { messages, sendMessage, stop, status, undoLastTurn, editAndRegenerate, approveTools } = useAgentChat({
     sessionId,
     isActive,
     onReady: () => logger.log(`Session ${sessionId} ready`),
@@ -57,11 +55,11 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
     return () => document.removeEventListener('visibilitychange', onVisible);
   }, [isActive, sessionId]);
 
-  // Wrap stop to track cancellation
+  // Wrap stop to show cancelled shimmer
   const handleStop = useCallback(() => {
     stop();
-    setWasCancelled(true);
-  }, [stop]);
+    updateSession(sessionId, { activityStatus: { type: 'cancelled' } });
+  }, [stop, updateSession, sessionId]);
 
   // SDK status is the ground truth — if it's streaming/submitted, agent is busy
   const sdkBusy = status === 'streaming' || status === 'submitted';
@@ -71,12 +69,11 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
     async (text: string) => {
       if (!text.trim() || busy) return;
 
-      setWasCancelled(false);
-      updateSession(sessionId, { isProcessing: true });
+      updateSession(sessionId, { isProcessing: true, activityStatus: { type: 'thinking' } });
       sendMessage({ text: text.trim(), metadata: { createdAt: new Date().toISOString() } });
 
       // Auto-title the session from the first user message
-      const isFirstMessage = messages.filter((m) => m.role === 'user').length <= 1;
+      const isFirstMessage = messages.filter((m) => m.role === 'user').length === 0;
       if (isFirstMessage) {
         apiFetch('/api/title', {
           method: 'POST',
@@ -105,6 +102,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
         isProcessing={busy}
         approveTools={approveTools}
         onUndoLastTurn={undoLastTurn}
+        onEditAndRegenerate={editAndRegenerate}
       />
       <ChatInput
         onSend={handleSendMessage}
@@ -114,9 +112,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
         placeholder={
           activityStatus.type === 'waiting-approval'
             ? 'Approve or reject pending tools first...'
-            : wasCancelled
-              ? 'What should the agent do instead?'
-              : undefined
+            : undefined
         }
       />
     </>

frontend/src/components/SessionSidebar/SessionSidebar.tsx

@@ -2,6 +2,12 @@ import { useCallback, useState } from 'react';
 import {
   Alert,
   Box,
+  Button,
+  Dialog,
+  DialogActions,
+  DialogContent,
+  DialogContentText,
+  DialogTitle,
   IconButton,
   Typography,
   CircularProgress,
@@ -51,20 +57,52 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
     }
   }, [isCreatingSession, createSession, setPlan, clearPanel, onClose]);
 
-  const handleDelete = useCallback(
-    async (sessionId: string, e: React.MouseEvent) => {
+  // -- Delete with dialog confirmation ------------------------------------
+  const [confirmDeleteId, setConfirmDeleteId] = useState<string | null>(null);
+  const [isDeleting, setIsDeleting] = useState(false);
+
+  const handleDeleteClick = useCallback(
+    (sessionId: string, e: React.MouseEvent) => {
       e.stopPropagation();
-      useAgentStore.getState().clearSessionState(sessionId);
-      try {
-        await apiFetch(`/api/session/${sessionId}`, { method: 'DELETE' });
-        deleteSession(sessionId);
-      } catch {
-        deleteSession(sessionId);
-      }
+      setConfirmDeleteId(sessionId);
     },
-    [deleteSession],
+    [],
   );
 
+  const handleDeleteConfirm = useCallback(async () => {
+    if (!confirmDeleteId || isDeleting) return;
+    const sessionId = confirmDeleteId;
+    setIsDeleting(true);
+
+    const isLastSession = sessions.length === 1;
+
+    useAgentStore.getState().clearSessionState(sessionId);
+    try {
+      await apiFetch(`/api/session/${sessionId}`, { method: 'DELETE' });
+      deleteSession(sessionId);
+    } catch {
+      deleteSession(sessionId);
+    }
+
+    // If this was the last session, create a new one
+    if (isLastSession) {
+      try {
+        const response = await apiFetch('/api/session', { method: 'POST' });
+        if (response.ok) {
+          const data = await response.json();
+          createSession(data.session_id);
+          setPlan([]);
+          clearPanel();
+        }
+      } catch (error) {
+        console.error('Failed to create new session after deleting last one:', error);
+      }
+    }
+
+    setIsDeleting(false);
+    setConfirmDeleteId(null);
+  }, [deleteSession, confirmDeleteId, isDeleting, sessions, createSession, setPlan, clearPanel]);
+
   const handleSelect = useCallback(
     (sessionId: string) => {
       switchSession(sessionId);
@@ -181,6 +219,7 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
                   px: 1.5,
                   py: 0.875,
                   mx: 0.75,
+                  mb: 0.2,
                   borderRadius: '10px',
                   cursor: 'pointer',
                   transition: 'background-color 0.12s ease',
@@ -256,7 +295,7 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
                 <IconButton
                   className="delete-btn"
                   size="small"
-                  onClick={(e) => handleDelete(session.id, e)}
+                  onClick={(e) => handleDeleteClick(session.id, e)}
                   sx={{
                     color: 'var(--muted-text)',
                     width: 26,
@@ -328,6 +367,89 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
         </Box>
 
       </Box>
+      {/* Delete confirmation dialog */}
+      <Dialog
+        open={!!confirmDeleteId}
+        onClose={() => !isDeleting && setConfirmDeleteId(null)}
+        slotProps={{
+          backdrop: { sx: { backgroundColor: 'rgba(0,0,0,0.5)', backdropFilter: 'blur(4px)' } },
+        }}
+        PaperProps={{
+          sx: {
+            bgcolor: 'var(--panel)',
+            border: '1px solid var(--border)',
+            borderRadius: 'var(--radius-md)',
+            boxShadow: 'var(--shadow-1)',
+            maxWidth: 340,
+            mx: 2,
+          },
+        }}
+      >
+        <DialogTitle
+          sx={{
+            color: 'var(--text)',
+            fontWeight: 700,
+            fontSize: '0.95rem',
+            pb: 0,
+            pt: 2.5,
+            px: 3,
+          }}
+        >
+          Delete conversation?
+        </DialogTitle>
+        <DialogContent sx={{ px: 3, pt: 1 }}>
+          <DialogContentText
+            sx={{
+              color: 'var(--muted-text)',
+              fontSize: '0.82rem',
+              lineHeight: 1.6,
+            }}
+          >
+            This will permanently remove this conversation and its history.
+          </DialogContentText>
+        </DialogContent>
+        <DialogActions sx={{ px: 3, pb: 2.5, gap: 1 }}>
+          <Button
+            onClick={() => setConfirmDeleteId(null)}
+            size="small"
+            disabled={isDeleting}
+            sx={{
+              color: 'var(--muted-text)',
+              fontSize: '0.82rem',
+              px: 2,
+              '&:hover': { bgcolor: 'var(--hover-bg)' },
+            }}
+          >
+            Cancel
+          </Button>
+          <Button
+            onClick={handleDeleteConfirm}
+            variant="contained"
+            size="small"
+            disabled={isDeleting}
+            startIcon={isDeleting ? <CircularProgress size={16} sx={{ color: '#fff' }} /> : undefined}
+            sx={{
+              fontSize: '0.82rem',
+              px: 2.5,
+              bgcolor: 'var(--accent-red)',
+              color: '#fff',
+              boxShadow: 'none',
+              '&:hover': {
+                bgcolor: 'var(--accent-red)',
+                filter: 'brightness(1.15)',
+                boxShadow: 'none',
+              },
+              '&.Mui-disabled': {
+                bgcolor: 'var(--accent-red)',
+                color: '#fff',
+                opacity: 0.7,
+              },
+            }}
+          >
+            {isDeleting ? 'Deleting...' : 'Delete'}
+          </Button>
+        </DialogActions>
+      </Dialog>
     </Box>
   );
 }

frontend/src/components/WelcomeScreen/WelcomeScreen.tsx

@@ -1,4 +1,4 @@
-import { useState, useCallback, useEffect, useRef } from 'react';
+import { useState, useCallback, useEffect, useRef, type ReactNode } from 'react';
 import {
   Box,
   Typography,
@@ -6,53 +6,236 @@ import {
   CircularProgress,
   Alert,
 } from '@mui/material';
+import CheckCircleIcon from '@mui/icons-material/CheckCircle';
 import OpenInNewIcon from '@mui/icons-material/OpenInNew';
 import GroupAddIcon from '@mui/icons-material/GroupAdd';
+import LoginIcon from '@mui/icons-material/Login';
+import RocketLaunchIcon from '@mui/icons-material/RocketLaunch';
 import { useSessionStore } from '@/store/sessionStore';
 import { useAgentStore } from '@/store/agentStore';
 import { apiFetch } from '@/utils/api';
 import { isInIframe, triggerLogin } from '@/hooks/useAuth';
+import { useOrgMembership } from '@/hooks/useOrgMembership';
 
-/** HF brand orange */
 const HF_ORANGE = '#FF9D00';
+const ORG_JOIN_URL =
+  'https://huggingface.co/organizations/ml-agent-explorers/share/GzPMJUivoFPlfkvFtIqEouZKSytatKQSZT';
 
-const ORG_JOIN_URL = 'https://huggingface.co/organizations/ml-agent-explorers/share/GzPMJUivoFPlfkvFtIqEouZKSytatKQSZT';
-const ORG_JOINED_KEY = 'hf-agent-org-joined';
+// ---------------------------------------------------------------------------
+// ChecklistStep sub-component
+// ---------------------------------------------------------------------------
 
-function hasJoinedOrg(): boolean {
-  try { return localStorage.getItem(ORG_JOINED_KEY) === '1'; } catch { return false; }
+type StepStatus = 'completed' | 'active' | 'locked';
+
+interface ChecklistStepProps {
+  stepNumber: number;
+  title: string;
+  description: string;
+  status: StepStatus;
+  lockedReason?: string;
+  actionLabel?: string;
+  onAction?: () => void;
+  actionIcon?: ReactNode;
+  actionHref?: string;
+  loading?: boolean;
+  isLast?: boolean;
+}
+
+function StepIndicator({ status, stepNumber }: { status: StepStatus; stepNumber: number }) {
+  if (status === 'completed') {
+    return <CheckCircleIcon sx={{ fontSize: 28, color: 'var(--accent-green)' }} />;
+  }
+  return (
+    <Box
+      sx={{
+        width: 28,
+        height: 28,
+        borderRadius: '50%',
+        display: 'flex',
+        alignItems: 'center',
+        justifyContent: 'center',
+        fontSize: '0.8rem',
+        fontWeight: 700,
+        ...(status === 'active'
+          ? { bgcolor: HF_ORANGE, color: '#000' }
+          : { bgcolor: 'transparent', border: '2px solid var(--border)', color: 'var(--muted-text)' }),
+      }}
+    >
+      {stepNumber}
+    </Box>
+  );
 }
 
-function markOrgJoined(): void {
-  try { localStorage.setItem(ORG_JOINED_KEY, '1'); } catch { /* ignore */ }
+function ChecklistStep({
+  stepNumber,
+  title,
+  description,
+  status,
+  lockedReason,
+  actionLabel,
+  onAction,
+  actionIcon,
+  actionHref,
+  loading = false,
+  isLast = false,
+}: ChecklistStepProps) {
+  const btnSx = {
+    px: 3,
+    py: 0.75,
+    fontSize: '0.85rem',
+    fontWeight: 700,
+    textTransform: 'none' as const,
+    borderRadius: '10px',
+    whiteSpace: 'nowrap' as const,
+    textDecoration: 'none',
+    ...(status === 'active'
+      ? {
+          bgcolor: HF_ORANGE,
+          color: '#000',
+          boxShadow: '0 2px 12px rgba(255, 157, 0, 0.25)',
+          '&:hover': { bgcolor: '#FFB340', boxShadow: '0 4px 20px rgba(255, 157, 0, 0.4)' },
+        }
+      : {
+          bgcolor: 'rgba(255,255,255,0.04)',
+          color: 'var(--muted-text)',
+          '&.Mui-disabled': { bgcolor: 'rgba(255,255,255,0.04)', color: 'var(--muted-text)' },
+        }),
+  };
+
+  return (
+    <Box
+      sx={{
+        display: 'flex',
+        alignItems: 'center',
+        gap: 2,
+        px: 3,
+        py: 2.5,
+        borderLeft: '3px solid',
+        borderLeftColor:
+          status === 'completed'
+            ? 'var(--accent-green)'
+            : status === 'active'
+              ? HF_ORANGE
+              : 'transparent',
+        ...(!isLast && { borderBottom: '1px solid var(--border)' }),
+        opacity: status === 'locked' ? 0.55 : 1,
+        transition: 'opacity 0.2s, border-color 0.2s',
+      }}
+    >
+      <StepIndicator status={status} stepNumber={stepNumber} />
+
+      <Box sx={{ flex: 1, minWidth: 0 }}>
+        <Typography
+          variant="subtitle2"
+          sx={{
+            fontWeight: 600,
+            fontSize: '0.92rem',
+            color: status === 'completed' ? 'var(--muted-text)' : 'var(--text)',
+            ...(status === 'completed' && { textDecoration: 'line-through', textDecorationColor: 'var(--muted-text)' }),
+          }}
+        >
+          {title}
+        </Typography>
+        <Typography variant="body2" sx={{ color: 'var(--muted-text)', fontSize: '0.8rem', mt: 0.25, lineHeight: 1.5 }}>
+          {status === 'locked' && lockedReason ? lockedReason : description}
+        </Typography>
+      </Box>
+
+      {status === 'completed' ? (
+        <Typography variant="caption" sx={{ color: 'var(--accent-green)', fontWeight: 600, fontSize: '0.78rem', whiteSpace: 'nowrap' }}>
+          Done
+        </Typography>
+      ) : actionLabel ? (
+        actionHref ? (
+          <Button
+            variant="contained"
+            size="small"
+            component="a"
+            href={actionHref}
+            target="_blank"
+            rel="noopener noreferrer"
+            disabled={status === 'locked'}
+            startIcon={actionIcon}
+            sx={btnSx}
+            onClick={onAction}
+          >
+            {actionLabel}
+          </Button>
+        ) : (
+          <Button
+            variant="contained"
+            size="small"
+            disabled={status === 'locked' || loading}
+            startIcon={loading ? <CircularProgress size={16} color="inherit" /> : actionIcon}
+            onClick={onAction}
+            sx={btnSx}
+          >
+            {loading ? 'Loading...' : actionLabel}
+          </Button>
+        )
+      ) : null}
+    </Box>
+  );
 }
 
+// ---------------------------------------------------------------------------
+// WelcomeScreen
+// ---------------------------------------------------------------------------
+
 export default function WelcomeScreen() {
   const { createSession } = useSessionStore();
   const { setPlan, clearPanel, user } = useAgentStore();
   const [isCreating, setIsCreating] = useState(false);
   const [error, setError] = useState<string | null>(null);
-  const [orgJoined, setOrgJoined] = useState(hasJoinedOrg);
-  const joinLinkOpened = useRef(false);
 
   const inIframe = isInIframe();
-  const isAuthenticated = user?.authenticated;
+  const isAuthenticated = !!user?.authenticated;
   const isDevUser = user?.username === 'dev';
 
-  // Auto-advance when user returns from the join link
+  // Iframe: localStorage-based org tracking (no auth token available)
+  const [iframeOrgJoined, setIframeOrgJoined] = useState(() => {
+    try { return localStorage.getItem('hf-agent-org-joined') === '1'; } catch { return false; }
+  });
+  const joinLinkOpened = useRef(false);
+
+  // Auto-advance when user returns from org join link (iframe only)
   useEffect(() => {
+    if (!inIframe) return;
     const handleVisibility = () => {
       if (document.visibilityState !== 'visible' || !joinLinkOpened.current) return;
       joinLinkOpened.current = false;
-      markOrgJoined();
-      setOrgJoined(true);
+      try { localStorage.setItem('hf-agent-org-joined', '1'); } catch { /* ignore */ }
+      setIframeOrgJoined(true);
     };
-
     document.addEventListener('visibilitychange', handleVisibility);
     return () => document.removeEventListener('visibilitychange', handleVisibility);
-  }, []);
+  }, [inIframe]);
+
+  const isOrgMember = inIframe ? iframeOrgJoined : !!user?.orgMember;
 
-  const tryCreateSession = useCallback(async () => {
+  // Poll for org membership once authenticated (skipped in dev mode and iframe)
+  const popupRef = useOrgMembership(isAuthenticated && !isDevUser && !inIframe && !isOrgMember);
+
+  // ---- Actions ----
+
+  const handleJoinOrg = useCallback(() => {
+    if (inIframe) {
+      // Iframe: open link, track via visibilitychange + localStorage
+      joinLinkOpened.current = true;
+      window.open(ORG_JOIN_URL, '_blank', 'noopener,noreferrer');
+      return;
+    }
+    // Direct: open as popup, auto-close via polling
+    const popup = window.open(ORG_JOIN_URL, 'hf-org-join', 'noopener');
+    if (popup) {
+      popupRef.current = popup;
+    } else {
+      window.open(ORG_JOIN_URL, '_blank', 'noopener,noreferrer');
+    }
+  }, [popupRef, inIframe]);
+
+  const handleStartSession = useCallback(async () => {
+    if (isCreating) return;
     setIsCreating(true);
     setError(null);
 
@@ -80,72 +263,21 @@ export default function WelcomeScreen() {
     } finally {
       setIsCreating(false);
     }
-  }, [createSession, setPlan, clearPanel]);
+  }, [isCreating, createSession, setPlan, clearPanel]);
 
-  const handleStart = useCallback(async () => {
-    if (isCreating) return;
-
-    if (!isAuthenticated && !isDevUser) {
-      if (inIframe) return;
-      triggerLogin();
-      return;
-    }
-
-    await tryCreateSession();
-  }, [isCreating, isAuthenticated, isDevUser, inIframe, tryCreateSession]);
-
-  // Build the direct Space URL for the "open in new tab" link
-  const spaceHost = typeof window !== 'undefined'
-    ? window.location.hostname.includes('.hf.space')
-      ? window.location.origin
-      : `https://smolagents-ml-agent.hf.space`
-    : '';
-
-  // Shared button style
-  const primaryBtnSx = {
-    px: 5,
-    py: 1.5,
-    fontSize: '1rem',
-    fontWeight: 700,
-    textTransform: 'none' as const,
-    borderRadius: '12px',
-    bgcolor: HF_ORANGE,
-    color: '#000',
-    boxShadow: '0 4px 24px rgba(255, 157, 0, 0.3)',
-    textDecoration: 'none',
-    '&:hover': {
-      bgcolor: '#FFB340',
-      boxShadow: '0 6px 32px rgba(255, 157, 0, 0.45)',
-    },
-  };
+  // ---- Step status helpers ----
 
-  // Description block (reused across screens)
-  const description = (
-    <Typography
-      variant="body1"
-      sx={{
-        color: 'var(--muted-text)',
-        maxWidth: 520,
-        mb: 5,
-        lineHeight: 1.8,
-        fontSize: '0.95rem',
-        textAlign: 'center',
-        px: 2,
-        '& strong': { color: 'var(--text)', fontWeight: 600 },
-      }}
-    >
-      A general-purpose AI agent for <strong>machine learning engineering</strong>.
-      It browses <strong>Hugging Face documentation</strong>, manages{' '}
-      <strong>repositories</strong>, launches <strong>training jobs</strong>,
-      and explores <strong>datasets</strong> — all through natural conversation.
-    </Typography>
-  );
+  const signInStatus: StepStatus = isAuthenticated ? 'completed' : 'active';
+  const joinOrgStatus: StepStatus = isOrgMember ? 'completed' : isAuthenticated ? 'active' : 'locked';
+  const startStatus: StepStatus = isAuthenticated && isOrgMember ? 'active' : 'locked';
 
-  // Which screen to show
-  const needsJoin = inIframe && !orgJoined;
-  const showOpenAgent = inIframe && orgJoined;
-  const showSignin = !inIframe && !isAuthenticated && !isDevUser;
-  const showReady = !inIframe && (isAuthenticated || isDevUser);
+  // Space URL for iframe "Open HF Agent" step
+  const spaceHost =
+    typeof window !== 'undefined'
+      ? window.location.hostname.includes('.hf.space')
+        ? window.location.origin
+        : 'https://smolagents-ml-agent.hf.space'
+      : '';
 
   return (
     <Box
@@ -160,12 +292,12 @@ export default function WelcomeScreen() {
         py: 8,
       }}
     >
-      {/* HF Logo */}
+      {/* Logo */}
       <Box
         component="img"
         src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg"
         alt="Hugging Face"
-        sx={{ width: 96, height: 96, mb: 3, display: 'block' }}
+        sx={{ width: 80, height: 80, mb: 2.5, display: 'block' }}
       />
 
       {/* Title */}
@@ -174,120 +306,128 @@ export default function WelcomeScreen() {
         sx={{
           fontWeight: 800,
           color: 'var(--text)',
-          mb: 1.5,
+          mb: 1,
           letterSpacing: '-0.02em',
-          fontSize: { xs: '2rem', md: '2.8rem' },
+          fontSize: { xs: '1.8rem', md: '2.4rem' },
         }}
       >
         HF Agent
       </Typography>
 
-      {/* ── Iframe: join org (first visit only) ──────────────────── */}
-      {needsJoin && (
-        <>
-          <Typography
-            variant="body1"
-            sx={{
-              color: 'var(--muted-text)',
-              maxWidth: 480,
-              mb: 4,
-              lineHeight: 1.8,
-              fontSize: '0.95rem',
-              textAlign: 'center',
-              px: 2,
-              '& strong': { color: 'var(--text)', fontWeight: 600 },
-            }}
-          >
-            Under the hood, this agent uses GPUs, inference APIs, and other paid Hub goodies — but we made them all free for you. Just join <strong>ML Agent Explorers</strong> to get started!
-          </Typography>
-
-          <Button
-            variant="contained"
-            size="large"
-            component="a"
-            href={ORG_JOIN_URL}
-            target="_blank"
-            rel="noopener noreferrer"
-            onClick={() => { joinLinkOpened.current = true; }}
-            startIcon={<GroupAddIcon />}
-            sx={primaryBtnSx}
-          >
-            Join ML Agent Explorers
-          </Button>
-        </>
-      )}
-
-      {/* ── Iframe: already joined → open Space ──────────────────── */}
-      {showOpenAgent && (
-        <>
-          {description}
-          <Button
-            variant="contained"
-            size="large"
-            component="a"
-            href={spaceHost}
-            target="_blank"
-            rel="noopener noreferrer"
-            endIcon={<OpenInNewIcon />}
-            sx={primaryBtnSx}
-          >
-            Open HF Agent
-          </Button>
-        </>
-      )}
-
-      {/* ── Direct: not logged in → sign in ──────────────────────── */}
-      {showSignin && (
-        <>
-          {description}
-          <Button
-            variant="contained"
-            size="large"
-            onClick={() => triggerLogin()}
-            sx={primaryBtnSx}
-          >
-            Sign in with Hugging Face
-          </Button>
+      {/* Description */}
+      <Typography
+        variant="body1"
+        sx={{
+          color: 'var(--muted-text)',
+          maxWidth: 480,
+          mb: 4,
+          lineHeight: 1.7,
+          fontSize: '0.9rem',
+          textAlign: 'center',
+          px: 2,
+          '& strong': { color: 'var(--text)', fontWeight: 600 },
+        }}
+      >
+        A general-purpose AI agent for <strong>machine learning engineering</strong>.
+        It browses <strong>Hugging Face docs</strong>, manages <strong>repos</strong>,
+        launches <strong>training jobs</strong>, and explores <strong>datasets</strong>.
+      </Typography>
 
-          <Typography
-            variant="caption"
-            sx={{
-              mt: 2.5,
-              color: 'var(--muted-text)',
-              fontSize: '0.78rem',
-              textAlign: 'center',
-              maxWidth: 360,
-              lineHeight: 1.6,
-            }}
-          >
-            Make sure to enable access to the <strong>ml-agent-explorers</strong> org when prompted.
-          </Typography>
-        </>
-      )}
+      {/* ── Checklist ──────────────────────────────────────────── */}
+      <Box
+        sx={{
+          width: '100%',
+          maxWidth: 520,
+          bgcolor: 'var(--surface)',
+          border: '1px solid var(--border)',
+          borderRadius: '12px',
+          overflow: 'hidden',
+          mx: 2,
+        }}
+      >
+        {isDevUser ? (
+          /* Dev mode: single step */
+          <ChecklistStep
+            stepNumber={1}
+            title="Start Session"
+            description="Launch an AI agent session for ML engineering."
+            status="active"
+            actionLabel="Start Session"
+            actionIcon={<RocketLaunchIcon sx={{ fontSize: 16 }} />}
+            onAction={handleStartSession}
+            loading={isCreating}
+            isLast
+          />
+        ) : inIframe ? (
+          /* Iframe: 2 steps */
+          <>
+            <ChecklistStep
+              stepNumber={1}
+              title="Join ML Agent Explorers"
+              description="Get free access to GPUs, inference APIs, and Hub resources."
+              status={isOrgMember ? 'completed' : 'active'}
+              actionLabel="Join Organization"
+              actionIcon={<GroupAddIcon sx={{ fontSize: 16 }} />}
+              onAction={handleJoinOrg}
+            />
+            <ChecklistStep
+              stepNumber={2}
+              title="Open HF Agent"
+              description="Open the agent in a full browser tab to get started."
+              status={isOrgMember ? 'active' : 'locked'}
+              lockedReason="Join the organization first."
+              actionLabel="Open HF Agent"
+              actionIcon={<OpenInNewIcon sx={{ fontSize: 16 }} />}
+              actionHref={spaceHost}
+              isLast
+            />
+          </>
+        ) : (
+          /* Direct access: 3 steps */
+          <>
+            <ChecklistStep
+              stepNumber={1}
+              title="Sign in with Hugging Face"
+              description="Authenticate to access GPU resources and model APIs."
+              status={signInStatus}
+              actionLabel="Sign in"
+              actionIcon={<LoginIcon sx={{ fontSize: 16 }} />}
+              onAction={() => triggerLogin()}
+            />
+            <ChecklistStep
+              stepNumber={2}
+              title="Join ML Agent Explorers"
+              description="Get free access to GPUs, inference APIs, and Hub resources."
+              status={joinOrgStatus}
+              lockedReason="Sign in first to continue."
+              actionLabel="Join Organization"
+              actionIcon={<GroupAddIcon sx={{ fontSize: 16 }} />}
+              onAction={handleJoinOrg}
+            />
+            <ChecklistStep
+              stepNumber={3}
+              title="Start Session"
+              description="Launch an AI agent session for ML engineering."
+              status={startStatus}
+              lockedReason="Complete the steps above to continue."
+              actionLabel="Start Session"
+              actionIcon={<RocketLaunchIcon sx={{ fontSize: 16 }} />}
+              onAction={handleStartSession}
+              loading={isCreating}
+              isLast
+            />
+          </>
+        )}
+      </Box>
 
-      {/* ── Direct: authenticated → start session ────────────────── */}
-      {showReady && (
-        <>
-          {description}
-          <Button
-            variant="contained"
-            size="large"
-            onClick={handleStart}
-            disabled={isCreating}
-            startIcon={
-              isCreating ? <CircularProgress size={20} color="inherit" /> : null
-            }
-            sx={{
-              ...primaryBtnSx,
-              '&.Mui-disabled': {
-                bgcolor: 'rgba(255, 157, 0, 0.35)',
-                color: 'rgba(0,0,0,0.45)',
-              },
-            }}
-          >
-            {isCreating ? 'Initializing...' : 'Start Session'}
-          </Button>
-        </>
+      {/* Polling hint when waiting for org join */}
+      {isAuthenticated && !isOrgMember && !isDevUser && !inIframe && (
+        <Typography
+          variant="caption"
+          sx={{ mt: 2, color: 'var(--muted-text)', fontSize: '0.75rem', textAlign: 'center' }}
+        >
+          This page updates automatically when you join the organization.
+        </Typography>
       )}
 
       {/* Error */}
@@ -311,7 +451,7 @@ export default function WelcomeScreen() {
       {/* Footnote */}
       <Typography
         variant="caption"
-        sx={{ mt: 5, color: 'var(--muted-text)', opacity: 0.5, fontSize: '0.7rem' }}
+        sx={{ mt: 4, color: 'var(--muted-text)', opacity: 0.5, fontSize: '0.7rem' }}
       >
         Conversations are stored locally in your browser.
       </Typography>

frontend/src/hooks/useAgentChat.ts

@@ -12,6 +12,7 @@ import { useChat } from '@ai-sdk/react';
 import { type UIMessage, lastAssistantMessageIsCompleteWithApprovalResponses } from 'ai';
 import { SSEChatTransport, type SideChannelCallbacks } from '@/lib/sse-chat-transport';
 import { loadMessages, saveMessages } from '@/lib/chat-message-store';
+import { saveResearch, loadResearch, clearResearch, RESEARCH_MAX_STEPS } from '@/lib/research-store';
 import { llmMessagesToUIMessages } from '@/lib/convert-llm-messages';
 import { apiFetch } from '@/utils/api';
 import { useAgentStore } from '@/store/agentStore';
@@ -86,14 +87,46 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         }
       },
       onToolLog: (tool: string, log: string) => {
-        // Research sub-agent: accumulate steps + update activity status
+        // Research sub-agent: parse stats vs step logs
         if (tool === 'research') {
           const sessState = useAgentStore.getState().getSessionState(sessionId);
-          const steps = [...sessState.researchSteps, log];
-          updateSession(sessionId, {
-            researchSteps: steps,
-            activityStatus: { type: 'tool', toolName: 'research', description: log },
-          });
+          const stats = { ...sessState.researchStats };
+
+          if (log === 'Starting research sub-agent...') {
+            const newStats = { toolCount: 0, tokenCount: 0, startedAt: Date.now(), finalElapsed: null };
+            updateSession(sessionId, {
+              researchSteps: [],
+              researchStats: newStats,
+              activityStatus: { type: 'tool', toolName: 'research', description: log },
+            });
+            saveResearch(sessionId, [], newStats);
+          } else if (log.startsWith('tokens:')) {
+            stats.tokenCount = parseInt(log.slice(7), 10);
+            updateSession(sessionId, { researchStats: stats });
+            saveResearch(sessionId, sessState.researchSteps, stats);
+          } else if (log.startsWith('tools:')) {
+            stats.toolCount = parseInt(log.slice(6), 10);
+            updateSession(sessionId, { researchStats: stats });
+            saveResearch(sessionId, sessState.researchSteps, stats);
+          } else if (log === 'Research complete.') {
+            const elapsed = stats.startedAt
+              ? Math.round((Date.now() - stats.startedAt) / 1000)
+              : null;
+            const doneStats = { ...stats, startedAt: null, finalElapsed: elapsed };
+            updateSession(sessionId, {
+              researchStats: doneStats,
+              activityStatus: { type: 'tool', toolName: 'research', description: log },
+            });
+            clearResearch(sessionId);
+          } else {
+            // Regular tool call step — append (trim to max)
+            const steps = [...sessState.researchSteps, log].slice(-RESEARCH_MAX_STEPS);
+            updateSession(sessionId, {
+              researchSteps: steps,
+              activityStatus: { type: 'tool', toolName: 'research', description: log },
+            });
+            saveResearch(sessionId, steps, stats);
+          }
           return;
         }
 
@@ -235,8 +268,11 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         const updates: Partial<import('@/store/agentStore').PerSessionState> = {
           activityStatus: { type: 'tool', toolName, description },
         };
-        // Clear research steps when a new research call starts
-        if (toolName === 'research') updates.researchSteps = [];
+        // Clear research steps + stats when a new research call starts
+        if (toolName === 'research') {
+          updates.researchSteps = [];
+          updates.researchStats = { toolCount: 0, tokenCount: 0, startedAt: null, finalElapsed: null };
+        }
         updateSession(sessionId, updates);
       },
       onInterrupted: () => { /* no-op — handled by stop() caller */ },
@@ -332,7 +368,7 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         if (msgsRes.ok) {
           const data = await msgsRes.json();
           if (cancelled || !Array.isArray(data) || data.length === 0) return;
-          const uiMsgs = llmMessagesToUIMessages(data, pendingIds);
+          const uiMsgs = llmMessagesToUIMessages(data, pendingIds, chatActionsRef.current.messages);
           if (uiMsgs.length > 0) {
             chat.setMessages(uiMsgs);
             saveMessages(sessionId, uiMsgs);
@@ -344,9 +380,25 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         // results make tools look "done" even when the agent is still
         // mid-turn and about to call more tools.
         if (backendIsProcessing) {
-          updateSession(sessionId, { isProcessing: true, activityStatus: { type: 'thinking' } });
+          // Restore research sub-agent state alongside isProcessing in one
+          // atomic update so the UI never sees isProcessing=false with stale
+          // tool states (which would coerce them to 'output-available').
+          const savedResearch = loadResearch(sessionId);
+          updateSession(sessionId, {
+            isProcessing: true,
+            activityStatus: savedResearch?.stats.startedAt
+              ? { type: 'tool', toolName: 'research', description: 'Resuming research...' }
+              : { type: 'thinking' },
+            ...(savedResearch && {
+              researchSteps: savedResearch.steps,
+              researchStats: savedResearch.stats,
+            }),
+          });
         } else if (pendingIds && pendingIds.size > 0) {
           updateSession(sessionId, { activityStatus: { type: 'waiting-approval' } });
+          clearResearch(sessionId);
+        } else {
+          clearResearch(sessionId);
         }
       } catch {
         /* backend unreachable -- localStorage fallback is fine */
@@ -453,7 +505,7 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
                 // Final hydration to get the complete message state
                 const result = await hydrateMessages();
                 if (result) {
-                  const uiMsgs = llmMessagesToUIMessages(result.data, result.pendingIds);
+                  const uiMsgs = llmMessagesToUIMessages(result.data, result.pendingIds, chatActionsRef.current.messages);
                   if (uiMsgs.length > 0) {
                     chat.setMessages(uiMsgs);
                     saveMessages(sessionId, uiMsgs);
@@ -467,7 +519,7 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
                 stopReconnect();
                 const result = await hydrateMessages();
                 if (result) {
-                  const uiMsgs = llmMessagesToUIMessages(result.data, result.pendingIds);
+                  const uiMsgs = llmMessagesToUIMessages(result.data, result.pendingIds, chatActionsRef.current.messages);
                   if (uiMsgs.length > 0) {
                     chat.setMessages(uiMsgs);
                     saveMessages(sessionId, uiMsgs);
@@ -491,7 +543,7 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
       if (!result) return;
 
       const { data, pendingIds, info } = result;
-      const uiMsgs = llmMessagesToUIMessages(data, pendingIds);
+      const uiMsgs = llmMessagesToUIMessages(data, pendingIds, chatActionsRef.current.messages);
       if (uiMsgs.length > 0) {
         chat.setMessages(uiMsgs);
         saveMessages(sessionId, uiMsgs);
@@ -514,11 +566,14 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         pollTimerRef.current = setInterval(async () => {
           const fresh = await hydrateMessages();
           if (!fresh) return;
-          const msgs = llmMessagesToUIMessages(fresh.data, fresh.pendingIds);
-          if (msgs.length > 0) {
+          const msgs = llmMessagesToUIMessages(fresh.data, fresh.pendingIds, chatActionsRef.current.messages);
+
+          const currentCount = chatActionsRef.current.messages.length;
+          if (msgs.length > currentCount || currentCount === 0) {
             chat.setMessages(msgs);
             saveMessages(sessionId, msgs);
-          }
+          } 
+
           // If backend stopped processing, clean up
           if (fresh.info && !fresh.info.is_processing) {
             updateSession(sessionId, { isProcessing: false });
@@ -542,7 +597,7 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
     if (chat.messages.length !== prevLenRef.current) {
       prevLenRef.current = chat.messages.length;
       saveMessages(sessionId, chat.messages);
-    }
+    } 
   }, [sessionId, chat.messages]);
 
   // -- Undo last turn (REST call + client-side message removal) -----------
@@ -598,6 +653,11 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
       if (hasApproved) {
         updateSession(sessionId, { isProcessing: true });
       }
+
+      // Persist updated tool states so a page refresh during execution
+      // won't restore stale approval-requested state from localStorage.
+      saveMessages(sessionId, chatActionsRef.current.messages);
+
       return true;
     },
     [sessionId, chat, updateSession, setNeedsAttention],
@@ -612,12 +672,52 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
     apiFetch(`/api/interrupt/${sessionId}`, { method: 'POST' }).catch(() => {});
   }, [sessionId, updateSession]);
 
+  // -- Edit message + regenerate from that point ----------------------------
+  const editAndRegenerate = useCallback(async (messageId: string, newText: string) => {
+    try {
+      const msgs = chatActionsRef.current.messages;
+      const setMsgs = chatActionsRef.current.setMessages;
+      if (!setMsgs) return;
+
+      // Find the target message and compute user message index (0-indexed, skipping system)
+      const msgIndex = msgs.findIndex(m => m.id === messageId);
+      if (msgIndex < 0) return;
+
+      let userMsgIndex = 0;
+      for (let i = 0; i < msgIndex; i++) {
+        if (msgs[i].role === 'user') userMsgIndex++;
+      }
+
+      // 1. Truncate backend history
+      const res = await apiFetch(`/api/truncate/${sessionId}`, {
+        method: 'POST',
+        body: JSON.stringify({ user_message_index: userMsgIndex }),
+        headers: { 'Content-Type': 'application/json' },
+      });
+      if (!res.ok) {
+        logger.error('Truncate API returned', res.status);
+        return;
+      }
+
+      // 2. Truncate frontend messages
+      const truncated = msgs.slice(0, msgIndex);
+      setMsgs(truncated);
+      saveMessages(sessionId, truncated);
+
+      // 3. Send the edited message (reuses existing transport + /api/chat)
+      chat.sendMessage({ text: newText, metadata: { createdAt: new Date().toISOString() } });
+    } catch (e) {
+      logger.error('Edit and regenerate failed:', e);
+    }
+  }, [sessionId, chat]);
+
   return {
     messages: chat.messages,
     sendMessage: chat.sendMessage,
     stop,
     status: chat.status,
     undoLastTurn,
+    editAndRegenerate,
     approveTools,
   };
 }

frontend/src/hooks/useOrgMembership.ts

@@ -0,0 +1,45 @@
+/**
+ * Polls backend for org membership status.
+ * When membership is detected, updates the user in the agent store
+ * and closes any org-join popup that was opened.
+ */
+import { useEffect, useRef } from 'react';
+import { useAgentStore } from '@/store/agentStore';
+
+const POLL_INTERVAL_MS = 3000;
+
+/**
+ * @param enabled  Only poll when true (user is authenticated but not yet confirmed as org member)
+ * @returns popupRef — assign `window.open()` result to `.current` so the hook can auto-close it
+ */
+export function useOrgMembership(enabled: boolean) {
+  const user = useAgentStore((s) => s.user);
+  const setUser = useAgentStore((s) => s.setUser);
+  const popupRef = useRef<Window | null>(null);
+
+  useEffect(() => {
+    if (!enabled || user?.orgMember) return;
+
+    let cancelled = false;
+
+    const check = async () => {
+      try {
+        const res = await fetch('/auth/org-membership', { credentials: 'include' });
+        if (!res.ok || cancelled) return;
+        const data = await res.json();
+        if (cancelled) return;
+        if (data.is_member && user) {
+          setUser({ ...user, orgMember: true });
+          try { popupRef.current?.close(); } catch { /* cross-origin or already closed */ }
+          popupRef.current = null;
+        }
+      } catch { /* backend unreachable — skip */ }
+    };
+
+    check();
+    const id = setInterval(check, POLL_INTERVAL_MS);
+    return () => { cancelled = true; clearInterval(id); };
+  }, [enabled, user?.orgMember, user, setUser]);
+
+  return popupRef;
+}

frontend/src/lib/chat-message-store.ts

@@ -38,7 +38,8 @@ function writeAll(map: MessagesMap): void {
 
 export function loadMessages(sessionId: string): UIMessage[] {
   const map = readAll();
-  return map[sessionId] ?? [];
+  const messages = map[sessionId] ?? [];
+  return messages;
 }
 
 export function saveMessages(sessionId: string, messages: UIMessage[]): void {

frontend/src/lib/convert-llm-messages.ts

@@ -16,19 +16,24 @@ interface LLMMessage {
   name?: string | null;
 }
 
-let idCounter = 0;
+// Generate stable IDs based on message position to prevent duplicate renders
+// when the same message is re-converted multiple times (e.g., during polling)
+let uiMessageCounter = 0;
 function nextId(): string {
-  return `msg-${Date.now()}-${++idCounter}`;
+  return `msg-${++uiMessageCounter}`;
 }
 
 /**
  * @param pendingApprovalIds - Set of tool_call_ids that are waiting for approval.
  *   When provided, matching tool calls without results will get state
  *   'approval-requested' instead of 'input-available'.
+ * @param existingUIMessages - Current UI messages to preserve IDs when content matches.
+ *   This prevents React from re-rendering messages with new IDs during polling.
  */
 export function llmMessagesToUIMessages(
   messages: LLMMessage[],
   pendingApprovalIds?: Set<string>,
+  existingUIMessages?: UIMessage[],
 ): UIMessage[] {
   // Build a map of tool_call_id -> tool result for pairing
   const toolResults = new Map<string, { output: string; isError: boolean }>();
@@ -43,13 +48,22 @@ export function llmMessagesToUIMessages(
 
   const uiMessages: UIMessage[] = [];
 
+  // Helper to get existing message ID at a given position if roles match
+  const getExistingId = (index: number, role: 'user' | 'assistant'): string | null => {
+    if (!existingUIMessages || index >= existingUIMessages.length) return null;
+    const existing = existingUIMessages[index];
+    return existing.role === role ? existing.id : null;
+  };
+
   for (const msg of messages) {
     if (msg.role === 'system') continue;
     if (msg.role === 'tool') continue; // handled via tool_calls pairing
 
     if (msg.role === 'user') {
+      // Try to reuse existing ID if the message at this position matches
+      const existingId = getExistingId(uiMessages.length, 'user');
       uiMessages.push({
-        id: nextId(),
+        id: existingId || nextId(),
         role: 'user',
         parts: [{ type: 'text', text: msg.content || '' }],
       });
@@ -109,8 +123,11 @@ export function llmMessagesToUIMessages(
       if (prev && prev.role === 'assistant') {
         prev.parts.push(...parts);
       } else {
+        // Try to reuse existing ID if the message at this position matches
+        const existingId = getExistingId(uiMessages.length, 'assistant');
+        const newId = existingId || nextId();
         uiMessages.push({
-          id: nextId(),
+          id: newId,
           role: 'assistant',
           parts,
         });

frontend/src/lib/research-store.ts

@@ -0,0 +1,56 @@
+/**
+ * Persist research sub-agent state (steps + stats) per session.
+ * Survives page refresh so the rolling display isn't lost mid-research.
+ */
+import type { PerSessionState } from '@/store/agentStore';
+
+/** Max steps to keep in storage and display. Single source of truth. */
+export const RESEARCH_MAX_STEPS = 40;
+
+const STORAGE_KEY = 'hf-agent-research';
+
+type ResearchState = {
+  steps: string[];
+  stats: PerSessionState['researchStats'];
+};
+
+type ResearchMap = Record<string, ResearchState>;
+
+function readAll(): ResearchMap {
+  try {
+    const raw = localStorage.getItem(STORAGE_KEY);
+    return raw ? JSON.parse(raw) : {};
+  } catch {
+    return {};
+  }
+}
+
+function writeAll(map: ResearchMap): void {
+  try {
+    localStorage.setItem(STORAGE_KEY, JSON.stringify(map));
+  } catch { /* quota exceeded — ignore */ }
+}
+
+export function saveResearch(
+  sessionId: string,
+  steps: string[],
+  stats: PerSessionState['researchStats'],
+): void {
+  const map = readAll();
+  map[sessionId] = {
+    steps: steps.slice(-RESEARCH_MAX_STEPS),
+    stats,
+  };
+  writeAll(map);
+}
+
+export function loadResearch(sessionId: string): ResearchState | null {
+  const map = readAll();
+  return map[sessionId] ?? null;
+}
+
+export function clearResearch(sessionId: string): void {
+  const map = readAll();
+  delete map[sessionId];
+  writeAll(map);
+}

frontend/src/lib/sse-chat-transport.ts

@@ -277,7 +277,8 @@ export class SSEChatTransport implements ChatTransport<UIMessage> {
     this.sessionId = sessionId;
     this.sideChannel = sideChannel;
     // Mark as connected immediately — no persistent connection to establish
-    sideChannel.onConnectionChange(true);
+    // Defer to avoid setState during render
+    queueMicrotask(() => sideChannel.onConnectionChange(true));
   }
 
   updateSideChannel(sideChannel: SideChannelCallbacks): void {

frontend/src/store/agentStore.ts

@@ -50,7 +50,8 @@ export type ActivityStatus =
   | { type: 'thinking' }
   | { type: 'tool'; toolName: string; description?: string }
   | { type: 'waiting-approval' }
-  | { type: 'streaming' };
+  | { type: 'streaming' }
+  | { type: 'cancelled' };
 
 /** State that is tracked per-session (each session has its own copy). */
 export interface PerSessionState {
@@ -62,6 +63,8 @@ export interface PerSessionState {
   plan: PlanItem[];
   /** Steps completed by the research sub-agent (tool_log events). */
   researchSteps: string[];
+  /** Live stats from the research sub-agent. */
+  researchStats: { toolCount: number; tokenCount: number; startedAt: number | null; finalElapsed: number | null };
 }
 
 const defaultSessionState: PerSessionState = {
@@ -72,6 +75,7 @@ const defaultSessionState: PerSessionState = {
   panelEditable: false,
   plan: [],
   researchSteps: [],
+  researchStats: { toolCount: 0, tokenCount: 0, startedAt: null, finalElapsed: null },
 };
 
 interface AgentStore {
@@ -101,6 +105,15 @@ interface AgentStore {
   // Job URLs (tool_call_id -> job URL) for HF jobs
   jobUrls: Record<string, string>;
 
+  // Job statuses (tool_call_id -> job status) for HF jobs
+  jobStatuses: Record<string, string>;
+
+  // Tool error states (tool_call_id -> true if errored) - persisted across renders
+  toolErrors: Record<string, boolean>;
+
+  // Tool rejected states (tool_call_id -> true if rejected by user) - persisted across renders
+  rejectedTools: Record<string, boolean>;
+
   // ── Per-session actions ─────────────────────────────────────────────
 
   /** Update a session's state. If it's the active session, also update flat state. */
@@ -138,6 +151,15 @@ interface AgentStore {
 
   setJobUrl: (toolCallId: string, jobUrl: string) => void;
   getJobUrl: (toolCallId: string) => string | undefined;
+
+  setJobStatus: (toolCallId: string, status: string) => void;
+  getJobStatus: (toolCallId: string) => string | undefined;
+
+  setToolError: (toolCallId: string, hasError: boolean) => void;
+  getToolError: (toolCallId: string) => boolean | undefined;
+
+  setToolRejected: (toolCallId: string, isRejected: boolean) => void;
+  getToolRejected: (toolCallId: string) => boolean | undefined;
 }
 
 /**
@@ -159,6 +181,44 @@ function syncSnapshot(
   };
 }
 
+// Load persisted tool errors from localStorage
+function loadToolErrors(): Record<string, boolean> {
+  try {
+    const stored = localStorage.getItem('hf-agent-tool-errors');
+    return stored ? JSON.parse(stored) : {};
+  } catch {
+    return {};
+  }
+}
+
+// Save tool errors to localStorage
+function saveToolErrors(errors: Record<string, boolean>): void {
+  try {
+    localStorage.setItem('hf-agent-tool-errors', JSON.stringify(errors));
+  } catch (e) {
+    console.warn('Failed to persist tool errors:', e);
+  }
+}
+
+// Load persisted rejected tools from localStorage
+function loadRejectedTools(): Record<string, boolean> {
+  try {
+    const stored = localStorage.getItem('hf-agent-rejected-tools');
+    return stored ? JSON.parse(stored) : {};
+  } catch {
+    return {};
+  }
+}
+
+// Save rejected tools to localStorage
+function saveRejectedTools(rejected: Record<string, boolean>): void {
+  try {
+    localStorage.setItem('hf-agent-rejected-tools', JSON.stringify(rejected));
+  } catch (e) {
+    console.warn('Failed to persist rejected tools:', e);
+  }
+}
+
 export const useAgentStore = create<AgentStore>()((set, get) => ({
   sessionStates: {},
   activeSessionId: null,
@@ -178,6 +238,9 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
 
   editedScripts: {},
   jobUrls: {},
+  jobStatuses: {},
+  toolErrors: loadToolErrors(),
+  rejectedTools: loadRejectedTools(),
 
   // ── Per-session state management ──────────────────────────────────
 
@@ -189,7 +252,7 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
     // Apply the processing→idle side effect
     const processingCleared = 'isProcessing' in updates && !updates.isProcessing;
     if (processingCleared) {
-      if (updated.activityStatus.type !== 'waiting-approval') {
+      if (updated.activityStatus.type !== 'waiting-approval' && updated.activityStatus.type !== 'cancelled') {
         updated.activityStatus = { type: 'idle' };
       }
     }
@@ -237,6 +300,7 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
         panelEditable: state.panelEditable,
         plan: state.plan,
         researchSteps: state.sessionStates[state.activeSessionId]?.researchSteps ?? [],
+        researchStats: state.sessionStates[state.activeSessionId]?.researchStats ?? defaultSessionState.researchStats,
       };
     }
 
@@ -267,7 +331,7 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
 
   setProcessing: (isProcessing) => {
     const current = get().activityStatus;
-    const preserveStatus = current.type === 'waiting-approval';
+    const preserveStatus = current.type === 'waiting-approval' || current.type === 'cancelled';
     set({ isProcessing, ...(!isProcessing && !preserveStatus ? { activityStatus: { type: 'idle' } } : {}) });
   },
   setConnected: (isConnected) => set({ isConnected }),
@@ -349,4 +413,38 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
   },
 
   getJobUrl: (toolCallId) => get().jobUrls[toolCallId],
+
+  // ── Job Statuses ────────────────────────────────────────────────────
+
+  setJobStatus: (toolCallId, status) => {
+    set((state) => ({
+      jobStatuses: { ...state.jobStatuses, [toolCallId]: status },
+    }));
+  },
+
+  getJobStatus: (toolCallId) => get().jobStatuses[toolCallId],
+
+  // ── Tool Errors ─────────────────────────────────────────────────────
+
+  setToolError: (toolCallId, hasError) => {
+    set((state) => {
+      const updated = { ...state.toolErrors, [toolCallId]: hasError };
+      saveToolErrors(updated);
+      return { toolErrors: updated };
+    });
+  },
+
+  getToolError: (toolCallId) => get().toolErrors[toolCallId],
+
+  // ── Tool Rejections ──────────────────────────────────────────────────
+
+  setToolRejected: (toolCallId, isRejected) => {
+    set((state) => {
+      const updated = { ...state.rejectedTools, [toolCallId]: isRejected };
+      saveRejectedTools(updated);
+      return { rejectedTools: updated };
+    });
+  },
+
+  getToolRejected: (toolCallId) => get().rejectedTools[toolCallId],
 }));

frontend/src/types/agent.ts

@@ -29,4 +29,5 @@ export interface User {
   username?: string;
   name?: string;
   picture?: string;
+  orgMember?: boolean;
 }