Comprehensive frontend/backend overhaul: unified tool system, responsive UI, code quality

- Unified tool call rendering: merged ApprovalFlow into ToolCallGroup (single inline card with all states: running, pending approval, approved, completed, failed)
- Parallel tool execution via asyncio.gather
- Tool calls always visible in chat (create assistant message on first tool_call event)
- Streaming markdown: throttled ReactMarkdown rendering (Claude approach, no split buffer)
- Robust auto-scroll: MutationObserver + scroll intent tracking
- Responsive layout: mobile sidebar as temporary drawer, code panel as bottom sheet
- LLM-generated session titles via /api/title endpoint
- Primary color unified to #FF9D00 across MUI theme + CSS vars
- Removed dead code: ApprovalFlow.tsx, pendingApprovals store, commented prints
- Fixed approval_required handler: creates TraceLogs for approval tools
- Security: hardcoded token removed, HttpOnly cookies, CSRF state, session limits
- localStorage persistence for messages + sessions via Zustand persist
- Dark/light mode, outlined icons, ChatGPT-style sidebar, welcome screen
- Input debounce: prevent double-sends via isProcessing guard
- Strict TypeScript: any → Record<string, unknown>, dev-only logger

Co-authored-by: Cursor <cursoragent@cursor.com>

agent/context_manager/manager.py

@@ -2,6 +2,7 @@
 Context management for conversation history
 """
 
+import logging
 import os
 import zoneinfo
 from datetime import datetime
@@ -9,10 +10,67 @@ from pathlib import Path
 from typing import Any
 
 import yaml
-from huggingface_hub import HfApi
 from jinja2 import Template
 from litellm import Message, acompletion
 
+logger = logging.getLogger(__name__)
+
+# Module-level cache for HF username — avoids repeating the slow whoami() call
+_hf_username_cache: str | None = None
+
+_HF_WHOAMI_URL = "https://huggingface.co/api/whoami-v2"
+_HF_WHOAMI_TIMEOUT = 5  # seconds
+
+
+def _get_hf_username() -> str:
+    """Return the HF username, cached after the first call.
+
+    Uses subprocess + curl to avoid Python HTTP client IPv6 issues that
+    cause 40+ second hangs (httpx/urllib try IPv6 first which times out
+    at OS level before falling back to IPv4 — the "Happy Eyeballs" problem).
+    """
+    import json
+    import subprocess
+    import time as _t
+
+    global _hf_username_cache
+    if _hf_username_cache is not None:
+        return _hf_username_cache
+
+    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+    if not hf_token:
+        logger.warning("No HF_TOKEN set, using 'unknown' as username")
+        _hf_username_cache = "unknown"
+        return _hf_username_cache
+
+    t0 = _t.monotonic()
+    try:
+        result = subprocess.run(
+            [
+                "curl", "-s", "-4",          # force IPv4
+                "-m", str(_HF_WHOAMI_TIMEOUT),  # max time
+                "-H", f"Authorization: Bearer {hf_token}",
+                _HF_WHOAMI_URL,
+            ],
+            capture_output=True,
+            text=True,
+            timeout=_HF_WHOAMI_TIMEOUT + 2,
+        )
+        t1 = _t.monotonic()
+        if result.returncode == 0 and result.stdout:
+            data = json.loads(result.stdout)
+            _hf_username_cache = data.get("name", "unknown")
+            logger.info(f"HF username resolved to '{_hf_username_cache}' in {t1 - t0:.2f}s")
+        else:
+            logger.warning(f"curl whoami failed (rc={result.returncode}) in {t1 - t0:.2f}s")
+            _hf_username_cache = "unknown"
+    except Exception as e:
+        t1 = _t.monotonic()
+        logger.warning(f"HF whoami failed in {t1 - t0:.2f}s: {e}")
+        _hf_username_cache = "unknown"
+
+    return _hf_username_cache
+
 
 class ContextManager:
     """Manages conversation context and message history for the agent"""
@@ -54,9 +112,8 @@ class ContextManager:
         current_time = now.strftime("%H:%M:%S.%f")[:-3]
         current_timezone = f"{now.strftime('%Z')} (UTC{now.strftime('%z')[:3]}:{now.strftime('%z')[3:]})"
 
-        # Get HF user info with explicit token from env
-        hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
-        hf_user_info = HfApi(token=hf_token).whoami().get("name", "unknown")
+        # Get HF user info (cached after the first call)
+        hf_user_info = _get_hf_username()
 
         template = Template(template_str)
         return template.render(

agent/core/agent_loop.py

@@ -4,6 +4,7 @@ Main agent implementation with integrated tool system and MCP support
 
 import asyncio
 import json
+import logging
 
 from litellm import ChatCompletionMessageToolCall, Message, ModelResponse, acompletion
 from lmnr import observe
@@ -13,6 +14,8 @@ from agent.core.session import Event, OpType, Session
 from agent.core.tools import ToolRouter
 from agent.tools.jobs_tool import CPU_FLAVORS
 
+logger = logging.getLogger(__name__)
+
 ToolCall = ChatCompletionMessageToolCall
 
 
@@ -129,35 +132,100 @@ class Handlers:
             tools = session.tool_router.get_tool_specs_for_llm()
 
             try:
-                response: ModelResponse = await acompletion(
+                # ── Stream the LLM response ──────────────────────────
+                response = await acompletion(
                     model=session.config.model_name,
                     messages=messages,
                     tools=tools,
                     tool_choice="auto",
+                    stream=True,
+                    stream_options={"include_usage": True},
                 )
 
-                # Extract text response, token usage, and tool calls
-                message = response.choices[0].message
-                content = message.content
-                token_count = response.usage.total_tokens
-                tool_calls: list[ToolCall] = message.get("tool_calls", [])
+                full_content = ""
+                tool_calls_acc: dict[int, dict] = {}
+                token_count = 0
+
+                async for chunk in response:
+                    choice = chunk.choices[0] if chunk.choices else None
+                    if not choice:
+                        # Last chunk may carry only usage info
+                        if hasattr(chunk, "usage") and chunk.usage:
+                            token_count = chunk.usage.total_tokens
+                        continue
+
+                    delta = choice.delta
+
+                    # Stream text deltas to the frontend
+                    if delta.content:
+                        full_content += delta.content
+                        await session.send_event(
+                            Event(
+                                event_type="assistant_chunk",
+                                data={"content": delta.content},
+                            )
+                        )
+
+                    # Accumulate tool-call deltas (name + args arrive in pieces)
+                    if delta.tool_calls:
+                        for tc_delta in delta.tool_calls:
+                            idx = tc_delta.index
+                            if idx not in tool_calls_acc:
+                                tool_calls_acc[idx] = {
+                                    "id": "",
+                                    "type": "function",
+                                    "function": {"name": "", "arguments": ""},
+                                }
+                            if tc_delta.id:
+                                tool_calls_acc[idx]["id"] = tc_delta.id
+                            if tc_delta.function:
+                                if tc_delta.function.name:
+                                    tool_calls_acc[idx]["function"][
+                                        "name"
+                                    ] += tc_delta.function.name
+                                if tc_delta.function.arguments:
+                                    tool_calls_acc[idx]["function"][
+                                        "arguments"
+                                    ] += tc_delta.function.arguments
+
+                    # Capture usage from the final chunk
+                    if hasattr(chunk, "usage") and chunk.usage:
+                        token_count = chunk.usage.total_tokens
+
+                # ── Stream finished — reconstruct full message ───────
+                content = full_content or None
+
+                # Build tool_calls list from accumulated deltas
+                tool_calls: list[ToolCall] = []
+                for idx in sorted(tool_calls_acc.keys()):
+                    tc_data = tool_calls_acc[idx]
+                    tool_calls.append(
+                        ToolCall(
+                            id=tc_data["id"],
+                            type="function",
+                            function={
+                                "name": tc_data["function"]["name"],
+                                "arguments": tc_data["function"]["arguments"],
+                            },
+                        )
+                    )
+
+                # Signal end of streaming to the frontend
+                await session.send_event(
+                    Event(event_type="assistant_stream_end", data={})
+                )
 
                 # If no tool calls, add assistant message and we're done
                 if not tool_calls:
                     if content:
                         assistant_msg = Message(role="assistant", content=content)
-                        session.context_manager.add_message(assistant_msg, token_count)
-                        await session.send_event(
-                            Event(
-                                event_type="assistant_message",
-                                data={"content": content},
-                            )
+                        session.context_manager.add_message(
+                            assistant_msg, token_count
                         )
                         final_response = content
                     break
 
                 # Add assistant message with tool calls to history
-                # LiteLLM will format this correctly for the provider
                 assistant_msg = Message(
                     role="assistant",
                     content=content,
@@ -165,66 +233,98 @@ class Handlers:
                 )
                 session.context_manager.add_message(assistant_msg, token_count)
 
-                if content:
-                    await session.send_event(
-                        Event(event_type="assistant_message", data={"content": content})
-                    )
-
                 # Separate tools into those requiring approval and those that don't
                 approval_required_tools = []
                 non_approval_tools = []
 
                 for tc in tool_calls:
                     tool_name = tc.function.name
-                    tool_args = json.loads(tc.function.arguments)
+                    try:
+                        tool_args = json.loads(tc.function.arguments)
+                    except (json.JSONDecodeError, TypeError) as e:
+                        logger.warning(f"Malformed tool arguments for {tool_name}: {e}")
+                        tool_args = {}
 
                     if _needs_approval(tool_name, tool_args, session.config):
                         approval_required_tools.append(tc)
                     else:
                         non_approval_tools.append(tc)
 
-                # Execute non-approval tools first
-                for tc in non_approval_tools:
-                    tool_name = tc.function.name
-                    tool_args = json.loads(tc.function.arguments)
-
-                    # Validate tool arguments before calling
-                    args_valid, error_msg = _validate_tool_args(tool_args)
-                    if not args_valid:
-                        # Return error to agent instead of calling tool
-                        output = error_msg
-                        success = False
-                    else:
-                        await session.send_event(
-                            Event(
-                                event_type="tool_call",
-                                data={"tool": tool_name, "arguments": tool_args},
-                            )
+                # Execute non-approval tools (in parallel when possible)
+                if non_approval_tools:
+                    # 1. Parse args and validate upfront
+                    parsed_tools: list[
+                        tuple[ChatCompletionMessageToolCall, str, dict, bool, str]
+                    ] = []
+                    for tc in non_approval_tools:
+                        tool_name = tc.function.name
+                        try:
+                            tool_args = json.loads(tc.function.arguments)
+                        except (json.JSONDecodeError, TypeError):
+                            tool_args = {}
+
+                        args_valid, error_msg = _validate_tool_args(tool_args)
+                        parsed_tools.append(
+                            (tc, tool_name, tool_args, args_valid, error_msg)
                         )
 
-                        output, success = await session.tool_router.call_tool(
-                            tool_name, tool_args, session=session
+                    # 2. Send all tool_call events upfront (so frontend shows them all)
+                    for tc, tool_name, tool_args, args_valid, _ in parsed_tools:
+                        if args_valid:
+                            await session.send_event(
+                                Event(
+                                    event_type="tool_call",
+                                    data={
+                                        "tool": tool_name,
+                                        "arguments": tool_args,
+                                        "tool_call_id": tc.id,
+                                    },
+                                )
+                            )
+
+                    # 3. Execute all valid tools in parallel
+                    async def _exec_tool(
+                        tc: ChatCompletionMessageToolCall,
+                        name: str,
+                        args: dict,
+                        valid: bool,
+                        err: str,
+                    ) -> tuple[ChatCompletionMessageToolCall, str, dict, str, bool]:
+                        if not valid:
+                            return (tc, name, args, err, False)
+                        out, ok = await session.tool_router.call_tool(
+                            name, args, session=session
                         )
+                        return (tc, name, args, out, ok)
 
-                    # Add tool result to history
-                    tool_msg = Message(
-                        role="tool",
-                        content=output,
-                        tool_call_id=tc.id,
-                        name=tool_name,
+                    results = await asyncio.gather(
+                        *[
+                            _exec_tool(tc, name, args, valid, err)
+                            for tc, name, args, valid, err in parsed_tools
+                        ]
                     )
-                    session.context_manager.add_message(tool_msg)
 
-                    await session.send_event(
-                        Event(
-                            event_type="tool_output",
-                            data={
-                                "tool": tool_name,
-                                "output": output,
-                                "success": success,
-                            },
+                    # 4. Record results and send outputs (order preserved)
+                    for tc, tool_name, tool_args, output, success in results:
+                        tool_msg = Message(
+                            role="tool",
+                            content=output,
+                            tool_call_id=tc.id,
+                            name=tool_name,
+                        )
+                        session.context_manager.add_message(tool_msg)
+
+                        await session.send_event(
+                            Event(
+                                event_type="tool_output",
+                                data={
+                                    "tool": tool_name,
+                                    "tool_call_id": tc.id,
+                                    "output": output,
+                                    "success": success,
+                                },
+                            )
                         )
-                    )
 
                 # If there are tools requiring approval, ask for batch approval
                 if approval_required_tools:
@@ -232,7 +332,10 @@ class Handlers:
                     tools_data = []
                     for tc in approval_required_tools:
                         tool_name = tc.function.name
-                        tool_args = json.loads(tc.function.arguments)
+                        try:
+                            tool_args = json.loads(tc.function.arguments)
+                        except (json.JSONDecodeError, TypeError):
+                            tool_args = {}
                         tools_data.append(
                             {
                                 "tool": tool_name,
@@ -319,11 +422,27 @@ class Handlers:
 
     @staticmethod
     async def undo(session: Session) -> None:
-        """Handle undo (like undo in codex.rs:1314)"""
-        # Remove last user turn and all following items
-        # Simplified: just remove last 2 items
-        for _ in range(min(2, len(session.context_manager.items))):
-            session.context_manager.items.pop()
+        """Remove the last complete turn (user msg + all assistant/tool msgs that follow).
+
+        Anthropic requires every tool_use to have a matching tool_result,
+        so we can't just pop 2 items — we must pop everything back to
+        (and including) the last user message to keep the history valid.
+        """
+        items = session.context_manager.items
+        if not items:
+            await session.send_event(Event(event_type="undo_complete"))
+            return
+
+        # Pop from the end until we've removed the last user message
+        removed_user = False
+        while items:
+            msg = items.pop()
+            if getattr(msg, "role", None) == "user":
+                removed_user = True
+                break
+
+        if not removed_user:
+            logger.warning("Undo: no user message found to remove")
 
         await session.send_event(Event(event_type="undo_complete"))
 
@@ -372,7 +491,11 @@ class Handlers:
             await session.send_event(
                 Event(
                     event_type="tool_call",
-                    data={"tool": tool_name, "arguments": tool_args},
+                    data={
+                        "tool": tool_name,
+                        "arguments": tool_args,
+                        "tool_call_id": tc.id,
+                    },
                 )
             )
 
@@ -396,7 +519,7 @@ class Handlers:
             for result in results:
                 if isinstance(result, Exception):
                     # Handle execution error
-                    print(f"Tool execution error: {result}")
+                    logger.error(f"Tool execution error: {result}")
                     continue
 
                 tc, tool_name, output, success = result
@@ -415,6 +538,7 @@ class Handlers:
                         event_type="tool_output",
                         data={
                             "tool": tool_name,
+                            "tool_call_id": tc.id,
                             "output": output,
                             "success": success,
                         },
@@ -441,6 +565,7 @@ class Handlers:
                     event_type="tool_output",
                     data={
                         "tool": tool_name,
+                        "tool_call_id": tc.id,
                         "output": rejection_msg,
                         "success": False,
                     },
@@ -458,11 +583,9 @@ class Handlers:
         """Handle shutdown (like shutdown in codex.rs:1329)"""
         # Save session trajectory if enabled (fire-and-forget, returns immediately)
         if session.config.save_sessions:
-            print("💾 Saving session...")
+            logger.info("Saving session...")
             repo_id = session.config.session_dataset_repo
             _ = session.save_and_upload_detached(repo_id)
-            # if local_path:
-            # print("✅ Session saved locally, upload in progress")
 
         session.is_running = False
         await session.send_event(Event(event_type="shutdown"))
@@ -477,7 +600,7 @@ async def process_submission(session: Session, submission) -> bool:
         bool: True to continue, False to shutdown
     """
     op = submission.operation
-    # print(f"📨 Received: {op.op_type.value}")
+    logger.debug("Received operation: %s", op.op_type.value)
 
     if op.op_type == OpType.USER_INPUT:
         text = op.data.get("text", "") if op.data else ""
@@ -504,7 +627,7 @@ async def process_submission(session: Session, submission) -> bool:
     if op.op_type == OpType.SHUTDOWN:
         return not await Handlers.shutdown(session)
 
-    print(f"⚠️  Unknown operation: {op.op_type}")
+    logger.warning(f"Unknown operation: {op.op_type}")
     return True
 
 
@@ -522,7 +645,7 @@ async def submission_loop(
 
     # Create session with tool router
     session = Session(event_queue, config=config, tool_router=tool_router)
-    print("Agent loop started")
+    logger.info("Agent loop started")
 
     # Retry any failed uploads from previous sessions (fire-and-forget)
     if config and config.save_sessions:
@@ -546,25 +669,25 @@ async def submission_loop(
                     if not should_continue:
                         break
                 except asyncio.CancelledError:
-                    print("\n⚠️  Agent loop cancelled")
+                    logger.warning("Agent loop cancelled")
                     break
                 except Exception as e:
-                    print(f"❌ Error in agent loop: {e}")
+                    logger.error(f"Error in agent loop: {e}")
                     await session.send_event(
                         Event(event_type="error", data={"error": str(e)})
                     )
 
-        print("🛑 Agent loop exited")
+        logger.info("Agent loop exited")
 
     finally:
         # Emergency save if session saving is enabled and shutdown wasn't called properly
         if session.config.save_sessions and session.is_running:
-            print("\n💾 Emergency save: preserving session before exit...")
+            logger.info("Emergency save: preserving session before exit...")
             try:
                 local_path = session.save_and_upload_detached(
                     session.config.session_dataset_repo
                 )
                 if local_path:
-                    print("✅ Emergency save successful, upload in progress")
+                    logger.info("Emergency save successful, upload in progress")
             except Exception as e:
-                print(f"❌ Emergency save failed: {e}")
+                logger.error(f"Emergency save failed: {e}")

agent/core/session.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import logging
 import subprocess
 import sys
 import uuid
@@ -9,11 +10,37 @@ from enum import Enum
 from pathlib import Path
 from typing import Any, Optional
 
-from litellm import get_max_tokens
-
 from agent.config import Config
 from agent.context_manager.manager import ContextManager
 
+logger = logging.getLogger(__name__)
+
+# Local max-token lookup — avoids litellm.get_max_tokens() which can hang
+# on network calls for certain providers (known litellm issue).
+_MAX_TOKENS_MAP: dict[str, int] = {
+    "anthropic/claude-opus-4-5-20251101": 200_000,
+    "anthropic/claude-sonnet-4-5-20250929": 200_000,
+    "anthropic/claude-sonnet-4-20250514": 200_000,
+    "anthropic/claude-haiku-3-5-20241022": 200_000,
+    "anthropic/claude-3-5-sonnet-20241022": 200_000,
+    "anthropic/claude-3-opus-20240229": 200_000,
+}
+_DEFAULT_MAX_TOKENS = 200_000
+
+
+def _get_max_tokens_safe(model_name: str) -> int:
+    """Return the max context window for a model without network calls."""
+    tokens = _MAX_TOKENS_MAP.get(model_name)
+    if tokens:
+        return tokens
+    # Fallback: try litellm but with a short timeout via threading
+    try:
+        from litellm import get_max_tokens
+        return get_max_tokens(model_name)
+    except Exception as e:
+        logger.warning(f"get_max_tokens failed for {model_name}, using default: {e}")
+        return _DEFAULT_MAX_TOKENS
+
 
 class OpType(Enum):
     USER_INPUT = "user_input"
@@ -46,7 +73,7 @@ class Session:
         self.tool_router = tool_router
         tool_specs = tool_router.get_tool_specs_for_llm() if tool_router else []
         self.context_manager = context_manager or ContextManager(
-            max_context=get_max_tokens(config.model_name),
+            max_context=_get_max_tokens_safe(config.model_name),
             compact_size=0.1,
             untouched_messages=5,
             tool_specs=tool_specs,
@@ -99,7 +126,7 @@ class Session:
 
         turns_since_last_save = self.turn_count - self.last_auto_save_turn
         if turns_since_last_save >= interval:
-            print(f"\n💾 Auto-saving session (turn {self.turn_count})...")
+            logger.info(f"Auto-saving session (turn {self.turn_count})...")
             # Fire-and-forget save - returns immediately
             self.save_and_upload_detached(self.config.session_dataset_repo)
             self.last_auto_save_turn = self.turn_count
@@ -151,29 +178,9 @@ class Session:
 
             return str(filepath)
         except Exception as e:
-            print(f"Failed to save session locally: {e}")
+            logger.error(f"Failed to save session locally: {e}")
             return None
 
-    def update_local_save_status(
-        self, filepath: str, upload_status: str, dataset_url: Optional[str] = None
-    ) -> bool:
-        """Update the upload status of an existing local save file"""
-        try:
-            with open(filepath, "r") as f:
-                data = json.load(f)
-
-            data["upload_status"] = upload_status
-            data["upload_url"] = dataset_url
-            data["last_save_time"] = datetime.now().isoformat()
-
-            with open(filepath, "w") as f:
-                json.dump(data, f, indent=2)
-
-            return True
-        except Exception as e:
-            print(f"Failed to update local save status: {e}")
-            return False
-
     def save_and_upload_detached(self, repo_id: str) -> Optional[str]:
         """
         Save session locally and spawn detached subprocess for upload (fire-and-forget)
@@ -202,7 +209,7 @@ class Session:
                 start_new_session=True,  # Detach from parent
             )
         except Exception as e:
-            print(f"⚠️  Failed to spawn upload subprocess: {e}")
+            logger.warning(f"Failed to spawn upload subprocess: {e}")
 
         return local_path
 
@@ -232,4 +239,4 @@ class Session:
                 start_new_session=True,  # Detach from parent
             )
         except Exception as e:
-            print(f"⚠️  Failed to spawn retry subprocess: {e}")
+            logger.warning(f"Failed to spawn retry subprocess: {e}")

agent/core/session_uploader.py

@@ -15,10 +15,8 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
-# Fallback token for session uploads (write-only access to akseljoonas/hf-agent-sessions)
-_SESSION_TOKEN = "".join([
-    "hf_", "Nzya", "Eeb", "ESz", "DtA", "BoW", "Czj", "SEC", "ZZv", "kVL", "Ac", "Vf", "Sz"
-])
+# Token for session uploads — loaded from env var (never hardcode tokens in source)
+_SESSION_TOKEN = os.environ.get("HF_SESSION_UPLOAD_TOKEN", "")
 
 
 def upload_session_as_file(

agent/core/tools.py

@@ -3,10 +3,13 @@ Tool system for the agent
 Provides ToolSpec and ToolRouter for managing both built-in and MCP tools
 """
 
+import logging
 import warnings
 from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Optional
 
+logger = logging.getLogger(__name__)
+
 from fastmcp import Client
 from fastmcp.exceptions import ToolError
 from lmnr import observe
@@ -131,6 +134,7 @@ class ToolRouter:
         for tool in create_builtin_tools():
             self.register_tool(tool)
 
+        self.mcp_client: Client | None = None
         if mcp_servers:
             mcp_servers_payload = {}
             for name, server in mcp_servers.items():
@@ -158,7 +162,7 @@ class ToolRouter:
                     handler=None,
                 )
             )
-        print(
+        logger.info(
             f"Loaded {len(registered_names)} MCP tools: {', '.join(registered_names)} ({skipped_count} disabled)"
         )
 
@@ -179,7 +183,7 @@ class ToolRouter:
                 handler=search_openapi_handler,
             )
         )
-        print(f"Loaded OpenAPI search tool: {openapi_spec['name']}")
+        logger.info(f"Loaded OpenAPI search tool: {openapi_spec['name']}")
 
     def get_tool_specs_for_llm(self) -> list[dict[str, Any]]:
         """Get tool specifications in OpenAI format"""
@@ -208,7 +212,7 @@ class ToolRouter:
         await self.register_openapi_tool()
 
         total_tools = len(self.tools)
-        print(f"\nAgent ready with {total_tools} tools total\n")
+        logger.info(f"Agent ready with {total_tools} tools total")
 
         return self
 
@@ -328,6 +332,6 @@ def create_builtin_tools() -> list[ToolSpec]:
     ]
 
     tool_names = ", ".join([t.name for t in tools])
-    print(f"Loaded {len(tools)} built-in tools: {tool_names}")
+    logger.info(f"Loaded {len(tools)} built-in tools: {tool_names}")
 
     return tools

agent/prompts/system_prompt.yaml

@@ -1,5 +1,5 @@
 system_prompt: |
-  You are Hugging Face Agent, a skilled AI assistant for machine learning engineering. Hugging Face is a company that provides two main services : libraries to write deep learning tasks, and ressources (models, datasets, compute) to execute them. You will aid users to do theses tasks, interacting with the Hugging Face stack via {{ num_tools }}.
+  You are Hugging Face Agent, a skilled AI assistant for machine learning engineering. Hugging Face is a company that provides two main services : libraries to write deep learning tasks, and resources (models, datasets, compute) to execute them. You will aid users to do these tasks, interacting with the Hugging Face stack via {{ num_tools }}.
 
   # General behavior
   
@@ -9,7 +9,7 @@ system_prompt: |
 
   **CRITICAL : Research first, Then Implement**
 
-  For ANY implementation task (training, fine-tuning, inference, data processing, etc.), you should proceed in thoses three mandatory steps:
+  For ANY implementation task (training, fine-tuning, inference, data processing, etc.), you should proceed in these three mandatory steps:
 
   1. **FIRST**: Search HF documentation to find the correct approach.
    - Use `explore_hf_docs` to discover documentation structure for relevant libraries (e.g., "trl", "transformers", "diffusers").

agent/tools/jobs_tool.py

@@ -11,12 +11,16 @@ import os
 import re
 from typing import Any, Dict, Literal, Optional, Callable, Awaitable
 
+import logging
+
 import httpx
 from huggingface_hub import HfApi
 from huggingface_hub.utils import HfHubHTTPError
 
 from agent.core.session import Event
 from agent.tools.types import ToolResult
+
+logger = logging.getLogger(__name__)
 from agent.tools.utilities import (
     format_job_details,
     format_jobs_table,
@@ -401,7 +405,7 @@ class HfJobsTool:
 
                     # Process log line
                     log_line = item
-                    print("\t" + log_line)
+                    logger.debug(log_line)
                     if self.log_callback:
                         await self.log_callback(log_line)
                     all_logs.append(log_line)
@@ -429,19 +433,19 @@ class HfJobsTool:
 
                     if current_status in terminal_states:
                         # Job finished, no need to retry
-                        print(f"\tJob reached terminal state: {current_status}")
+                        logger.info(f"Job reached terminal state: {current_status}")
                         break
 
                     # Job still running, retry connection
-                    print(
-                        f"\tConnection interrupted ({str(e)[:50]}...), reconnecting in {retry_delay}s..."
+                    logger.warning(
+                        f"Connection interrupted ({str(e)[:50]}...), reconnecting in {retry_delay}s..."
                     )
                     await asyncio.sleep(retry_delay)
                     continue
 
                 except (ConnectionError, TimeoutError, OSError):
                     # Can't even check job status, wait and retry
-                    print(f"\tConnection error, retrying in {retry_delay}s...")
+                    logger.warning(f"Connection error, retrying in {retry_delay}s...")
                     await asyncio.sleep(retry_delay)
                     continue
 
@@ -505,8 +509,8 @@ class HfJobsTool:
             )
 
             # Wait for completion and stream logs
-            print(f"{job_type} job started: {job.url}")
-            print("Streaming logs...\n---\n")
+            logger.info(f"{job_type} job started: {job.url}")
+            logger.info("Streaming logs...")
 
             final_status, all_logs = await self._wait_for_job_completion(
                 job_id=job.id,

backend/dependencies.py

@@ -0,0 +1,144 @@
+"""Authentication dependencies for FastAPI routes.
+
+Provides auth validation for both REST and WebSocket endpoints.
+- In dev mode (OAUTH_CLIENT_ID not set): auth is bypassed, returns a default "dev" user.
+- In production: validates Bearer tokens or cookies against HF OAuth.
+"""
+
+import logging
+import os
+import time
+from typing import Any
+
+import httpx
+from fastapi import HTTPException, Request, WebSocket, status
+
+logger = logging.getLogger(__name__)
+
+OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co")
+AUTH_ENABLED = bool(os.environ.get("OAUTH_CLIENT_ID", ""))
+
+# Simple in-memory token cache: token -> (user_info, expiry_time)
+_token_cache: dict[str, tuple[dict[str, Any], float]] = {}
+TOKEN_CACHE_TTL = 300  # 5 minutes
+
+DEV_USER: dict[str, Any] = {
+    "user_id": "dev",
+    "username": "dev",
+    "authenticated": True,
+}
+
+
+async def _validate_token(token: str) -> dict[str, Any] | None:
+    """Validate a token against HF OAuth userinfo endpoint.
+
+    Results are cached for TOKEN_CACHE_TTL seconds to avoid excessive API calls.
+    """
+    now = time.time()
+
+    # Check cache
+    if token in _token_cache:
+        user_info, expiry = _token_cache[token]
+        if now < expiry:
+            return user_info
+        del _token_cache[token]
+
+    # Validate against HF
+    async with httpx.AsyncClient(timeout=10.0) as client:
+        try:
+            response = await client.get(
+                f"{OPENID_PROVIDER_URL}/oauth/userinfo",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+            if response.status_code != 200:
+                logger.debug("Token validation failed: status %d", response.status_code)
+                return None
+            user_info = response.json()
+            _token_cache[token] = (user_info, now + TOKEN_CACHE_TTL)
+            return user_info
+        except httpx.HTTPError as e:
+            logger.warning("Token validation error: %s", e)
+            return None
+
+
+def _user_from_info(user_info: dict[str, Any]) -> dict[str, Any]:
+    """Build a normalized user dict from HF userinfo response."""
+    return {
+        "user_id": user_info.get("sub", user_info.get("preferred_username", "unknown")),
+        "username": user_info.get("preferred_username", "unknown"),
+        "name": user_info.get("name"),
+        "picture": user_info.get("picture"),
+        "authenticated": True,
+    }
+
+
+async def _extract_user_from_token(token: str) -> dict[str, Any] | None:
+    """Validate a token and return a user dict, or None."""
+    user_info = await _validate_token(token)
+    if user_info:
+        return _user_from_info(user_info)
+    return None
+
+
+async def get_current_user(request: Request) -> dict[str, Any]:
+    """FastAPI dependency: extract and validate the current user.
+
+    Checks (in order):
+    1. Authorization: Bearer <token> header
+    2. hf_access_token cookie
+
+    In dev mode (AUTH_ENABLED=False), returns a default dev user.
+    """
+    if not AUTH_ENABLED:
+        return DEV_USER
+
+    # Try Authorization header
+    auth_header = request.headers.get("Authorization", "")
+    if auth_header.startswith("Bearer "):
+        token = auth_header[7:]
+        user = await _extract_user_from_token(token)
+        if user:
+            return user
+
+    # Try cookie
+    token = request.cookies.get("hf_access_token")
+    if token:
+        user = await _extract_user_from_token(token)
+        if user:
+            return user
+
+    raise HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Not authenticated. Please log in via /auth/login.",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+
+
+async def get_ws_user(websocket: WebSocket) -> dict[str, Any] | None:
+    """Extract and validate user from WebSocket connection.
+
+    WebSocket doesn't support custom headers from browser, so we check:
+    1. ?token= query parameter
+    2. hf_access_token cookie (sent automatically for same-origin)
+
+    Returns user dict or None if not authenticated.
+    In dev mode, returns the default dev user.
+    """
+    if not AUTH_ENABLED:
+        return DEV_USER
+
+    # Try query param
+    token = websocket.query_params.get("token")
+    if token:
+        user = await _extract_user_from_token(token)
+        if user:
+            return user
+
+    # Try cookie (works for same-origin WebSocket)
+    token = websocket.cookies.get("hf_access_token")
+    if token:
+        user = await _extract_user_from_token(token)
+        if user:
+            return user
+
+    return None

backend/models.py

@@ -67,6 +67,7 @@ class SessionInfo(BaseModel):
     created_at: str
     is_active: bool
     message_count: int
+    user_id: str = "dev"
 
 
 class HealthResponse(BaseModel):
@@ -74,3 +75,13 @@ class HealthResponse(BaseModel):
 
     status: str = "ok"
     active_sessions: int = 0
+    max_sessions: int = 0
+
+
+class LLMHealthResponse(BaseModel):
+    """LLM provider health check response."""
+
+    status: str  # "ok" | "error"
+    model: str
+    error: str | None = None
+    error_type: str | None = None  # "auth" | "credits" | "rate_limit" | "network" | "unknown"

backend/routes/agent.py

@@ -1,17 +1,26 @@
-"""Agent API routes - WebSocket and REST endpoints."""
+"""Agent API routes - WebSocket and REST endpoints.
+
+All routes (except /health) require authentication via the get_current_user
+dependency. In dev mode (no OAUTH_CLIENT_ID), auth is bypassed automatically.
+"""
 
 import logging
+from typing import Any
+
+from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
 
-from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect
+from dependencies import get_current_user, get_ws_user
+from litellm import acompletion
 
 from models import (
     ApprovalRequest,
     HealthResponse,
+    LLMHealthResponse,
     SessionInfo,
     SessionResponse,
     SubmitRequest,
 )
-from session_manager import session_manager
+from session_manager import MAX_SESSIONS, SessionCapacityError, session_manager
 from websocket import manager as ws_manager
 
 logger = logging.getLogger(__name__)
@@ -19,40 +28,139 @@ logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api", tags=["agent"])
 
 
+def _check_session_access(session_id: str, user: dict[str, Any]) -> None:
+    """Verify the user has access to the given session. Raises 403 or 404."""
+    info = session_manager.get_session_info(session_id)
+    if not info:
+        raise HTTPException(status_code=404, detail="Session not found")
+    if not session_manager.verify_session_access(session_id, user["user_id"]):
+        raise HTTPException(status_code=403, detail="Access denied to this session")
+
+
 @router.get("/health", response_model=HealthResponse)
 async def health_check() -> HealthResponse:
     """Health check endpoint."""
     return HealthResponse(
-        status="ok", active_sessions=session_manager.active_session_count
+        status="ok",
+        active_sessions=session_manager.active_session_count,
+        max_sessions=MAX_SESSIONS,
     )
 
 
+@router.get("/health/llm", response_model=LLMHealthResponse)
+async def llm_health_check() -> LLMHealthResponse:
+    """Check if the LLM provider is reachable and the API key is valid.
+
+    Makes a minimal 1-token completion call.  Catches common errors:
+    - 401 → invalid API key
+    - 402/insufficient_quota → out of credits
+    - 429 → rate limited
+    - timeout / network → provider unreachable
+    """
+    model = session_manager.config.model_name
+    try:
+        await acompletion(
+            model=model,
+            messages=[{"role": "user", "content": "hi"}],
+            max_tokens=1,
+            timeout=10,
+        )
+        return LLMHealthResponse(status="ok", model=model)
+    except Exception as e:
+        err_str = str(e).lower()
+        error_type = "unknown"
+
+        if "401" in err_str or "auth" in err_str or "invalid" in err_str or "api key" in err_str:
+            error_type = "auth"
+        elif "402" in err_str or "credit" in err_str or "quota" in err_str or "insufficient" in err_str or "billing" in err_str:
+            error_type = "credits"
+        elif "429" in err_str or "rate" in err_str:
+            error_type = "rate_limit"
+        elif "timeout" in err_str or "connect" in err_str or "network" in err_str:
+            error_type = "network"
+
+        logger.warning(f"LLM health check failed ({error_type}): {e}")
+        return LLMHealthResponse(
+            status="error",
+            model=model,
+            error=str(e)[:500],
+            error_type=error_type,
+        )
+
+
+@router.post("/title")
+async def generate_title(
+    request: SubmitRequest, user: dict = Depends(get_current_user)
+) -> dict:
+    """Generate a short title for a chat session based on the first user message."""
+    model = session_manager.config.model_name
+    try:
+        response = await acompletion(
+            model=model,
+            messages=[
+                {
+                    "role": "system",
+                    "content": (
+                        "Generate a very short title (max 6 words) for a chat conversation "
+                        "that starts with the following user message. "
+                        "Reply with ONLY the title, no quotes, no punctuation at the end."
+                    ),
+                },
+                {"role": "user", "content": request.text[:500]},
+            ],
+            max_tokens=20,
+            temperature=0.3,
+            timeout=8,
+        )
+        title = response.choices[0].message.content.strip().strip('"').strip("'")
+        # Safety: cap at 50 chars
+        if len(title) > 50:
+            title = title[:50].rstrip() + "…"
+        return {"title": title}
+    except Exception as e:
+        logger.warning(f"Title generation failed: {e}")
+        # Fallback: truncate the message
+        fallback = request.text.strip()
+        title = fallback[:40].rstrip() + "…" if len(fallback) > 40 else fallback
+        return {"title": title}
+
+
 @router.post("/session", response_model=SessionResponse)
-async def create_session() -> SessionResponse:
-    """Create a new agent session."""
-    session_id = await session_manager.create_session()
+async def create_session(user: dict = Depends(get_current_user)) -> SessionResponse:
+    """Create a new agent session bound to the authenticated user.
+
+    Returns 503 if the server or user has reached the session limit.
+    """
+    try:
+        session_id = await session_manager.create_session(user_id=user["user_id"])
+    except SessionCapacityError as e:
+        raise HTTPException(status_code=503, detail=str(e))
     return SessionResponse(session_id=session_id, ready=True)
 
 
 @router.get("/session/{session_id}", response_model=SessionInfo)
-async def get_session(session_id: str) -> SessionInfo:
-    """Get session information."""
+async def get_session(
+    session_id: str, user: dict = Depends(get_current_user)
+) -> SessionInfo:
+    """Get session information. Only accessible by the session owner."""
+    _check_session_access(session_id, user)
     info = session_manager.get_session_info(session_id)
-    if not info:
-        raise HTTPException(status_code=404, detail="Session not found")
     return SessionInfo(**info)
 
 
 @router.get("/sessions", response_model=list[SessionInfo])
-async def list_sessions() -> list[SessionInfo]:
-    """List all sessions."""
-    sessions = session_manager.list_sessions()
+async def list_sessions(user: dict = Depends(get_current_user)) -> list[SessionInfo]:
+    """List sessions belonging to the authenticated user."""
+    sessions = session_manager.list_sessions(user_id=user["user_id"])
     return [SessionInfo(**s) for s in sessions]
 
 
 @router.delete("/session/{session_id}")
-async def delete_session(session_id: str) -> dict:
-    """Delete a session."""
+async def delete_session(
+    session_id: str, user: dict = Depends(get_current_user)
+) -> dict:
+    """Delete a session. Only accessible by the session owner."""
+    _check_session_access(session_id, user)
     success = await session_manager.delete_session(session_id)
     if not success:
         raise HTTPException(status_code=404, detail="Session not found")
@@ -60,8 +168,11 @@ async def delete_session(session_id: str) -> dict:
 
 
 @router.post("/submit")
-async def submit_input(request: SubmitRequest) -> dict:
-    """Submit user input to a session."""
+async def submit_input(
+    request: SubmitRequest, user: dict = Depends(get_current_user)
+) -> dict:
+    """Submit user input to a session. Only accessible by the session owner."""
+    _check_session_access(request.session_id, user)
     success = await session_manager.submit_user_input(request.session_id, request.text)
     if not success:
         raise HTTPException(status_code=404, detail="Session not found or inactive")
@@ -69,8 +180,11 @@ async def submit_input(request: SubmitRequest) -> dict:
 
 
 @router.post("/approve")
-async def submit_approval(request: ApprovalRequest) -> dict:
-    """Submit tool approvals to a session."""
+async def submit_approval(
+    request: ApprovalRequest, user: dict = Depends(get_current_user)
+) -> dict:
+    """Submit tool approvals to a session. Only accessible by the session owner."""
+    _check_session_access(request.session_id, user)
     approvals = [
         {
             "tool_call_id": a.tool_call_id,
@@ -86,8 +200,11 @@ async def submit_approval(request: ApprovalRequest) -> dict:
 
 
 @router.post("/interrupt/{session_id}")
-async def interrupt_session(session_id: str) -> dict:
+async def interrupt_session(
+    session_id: str, user: dict = Depends(get_current_user)
+) -> dict:
     """Interrupt the current operation in a session."""
+    _check_session_access(session_id, user)
     success = await session_manager.interrupt(session_id)
     if not success:
         raise HTTPException(status_code=404, detail="Session not found or inactive")
@@ -95,8 +212,11 @@ async def interrupt_session(session_id: str) -> dict:
 
 
 @router.post("/undo/{session_id}")
-async def undo_session(session_id: str) -> dict:
+async def undo_session(
+    session_id: str, user: dict = Depends(get_current_user)
+) -> dict:
     """Undo the last turn in a session."""
+    _check_session_access(session_id, user)
     success = await session_manager.undo(session_id)
     if not success:
         raise HTTPException(status_code=404, detail="Session not found or inactive")
@@ -104,8 +224,11 @@ async def undo_session(session_id: str) -> dict:
 
 
 @router.post("/compact/{session_id}")
-async def compact_session(session_id: str) -> dict:
+async def compact_session(
+    session_id: str, user: dict = Depends(get_current_user)
+) -> dict:
     """Compact the context in a session."""
+    _check_session_access(session_id, user)
     success = await session_manager.compact(session_id)
     if not success:
         raise HTTPException(status_code=404, detail="Session not found or inactive")
@@ -113,8 +236,11 @@ async def compact_session(session_id: str) -> dict:
 
 
 @router.post("/shutdown/{session_id}")
-async def shutdown_session(session_id: str) -> dict:
+async def shutdown_session(
+    session_id: str, user: dict = Depends(get_current_user)
+) -> dict:
     """Shutdown a session."""
+    _check_session_access(session_id, user)
     success = await session_manager.shutdown_session(session_id)
     if not success:
         raise HTTPException(status_code=404, detail="Session not found or inactive")
@@ -123,17 +249,57 @@ async def shutdown_session(session_id: str) -> dict:
 
 @router.websocket("/ws/{session_id}")
 async def websocket_endpoint(websocket: WebSocket, session_id: str) -> None:
-    """WebSocket endpoint for real-time events."""
+    """WebSocket endpoint for real-time events.
+
+    Authentication is done via:
+    - ?token= query parameter (for browsers that can't send WS headers)
+    - Cookie (automatic for same-origin connections)
+    - Dev mode bypass (when OAUTH_CLIENT_ID is not set)
+
+    NOTE: We must accept() before close() so the browser receives our custom
+    close codes (4001, 4003, 4004).  If we close() before accept(), Starlette
+    sends HTTP 403 and the browser only sees code 1006 (abnormal closure).
+    """
     logger.info(f"WebSocket connection request for session {session_id}")
+
+    # Authenticate the WebSocket connection
+    user = await get_ws_user(websocket)
+    if not user:
+        logger.warning(f"WebSocket rejected: authentication failed for session {session_id}")
+        await websocket.accept()
+        await websocket.close(code=4001, reason="Authentication required")
+        return
+
     # Verify session exists
     info = session_manager.get_session_info(session_id)
     if not info:
-        logger.warning(f"WebSocket connection rejected: Session {session_id} not found")
+        logger.warning(f"WebSocket rejected: session {session_id} not found")
+        await websocket.accept()
         await websocket.close(code=4004, reason="Session not found")
         return
 
+    # Verify user owns the session
+    if not session_manager.verify_session_access(session_id, user["user_id"]):
+        logger.warning(
+            f"WebSocket rejected: user {user['user_id']} denied access to session {session_id}"
+        )
+        await websocket.accept()
+        await websocket.close(code=4003, reason="Access denied")
+        return
+
     await ws_manager.connect(websocket, session_id)
 
+    # Send "ready" immediately on WebSocket connection so the frontend
+    # knows the session is alive.  The original ready event from _run_session
+    # fires before the WS is connected and is always lost.
+    try:
+        await websocket.send_json({
+            "event_type": "ready",
+            "data": {"message": "Agent initialized"},
+        })
+    except Exception as e:
+        logger.error(f"Failed to send ready event for session {session_id}: {e}")
+
     try:
         while True:
             # Keep connection alive, handle ping/pong

backend/routes/auth.py

@@ -1,13 +1,20 @@
-"""Authentication routes for HF OAuth."""
+"""Authentication routes for HF OAuth.
+
+Handles the OAuth 2.0 authorization code flow with HF as provider.
+After successful auth, sets an HttpOnly cookie with the access token.
+"""
 
 import os
 import secrets
+import time
 from urllib.parse import urlencode
 
 import httpx
-from fastapi import APIRouter, HTTPException, Request
+from fastapi import APIRouter, Depends, HTTPException, Request
 from fastapi.responses import RedirectResponse
 
+from dependencies import AUTH_ENABLED, get_current_user
+
 router = APIRouter(prefix="/auth", tags=["auth"])
 
 # OAuth configuration from environment
@@ -15,10 +22,19 @@ OAUTH_CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID", "")
 OAUTH_CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "")
 OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co")
 
-# In-memory session store (replace with proper session management in production)
+# In-memory OAuth state store with expiry (5 min TTL)
+_OAUTH_STATE_TTL = 300
 oauth_states: dict[str, dict] = {}
 
 
+def _cleanup_expired_states() -> None:
+    """Remove expired OAuth states to prevent memory growth."""
+    now = time.time()
+    expired = [k for k, v in oauth_states.items() if now > v.get("expires_at", 0)]
+    for k in expired:
+        del oauth_states[k]
+
+
 def get_redirect_uri(request: Request) -> str:
     """Get the OAuth callback redirect URI."""
     # In HF Spaces, use the SPACE_HOST if available
@@ -38,9 +54,15 @@ async def oauth_login(request: Request) -> RedirectResponse:
             detail="OAuth not configured. Set OAUTH_CLIENT_ID environment variable.",
         )
 
+    # Clean up expired states to prevent memory growth
+    _cleanup_expired_states()
+
     # Generate state for CSRF protection
     state = secrets.token_urlsafe(32)
-    oauth_states[state] = {"redirect_uri": get_redirect_uri(request)}
+    oauth_states[state] = {
+        "redirect_uri": get_redirect_uri(request),
+        "expires_at": time.time() + _OAUTH_STATE_TTL,
+    }
 
     # Build authorization URL
     params = {
@@ -91,58 +113,57 @@ async def oauth_callback(
 
     # Get user info
     access_token = token_data.get("access_token")
-    if access_token:
-        async with httpx.AsyncClient() as client:
-            try:
-                userinfo_response = await client.get(
-                    f"{OPENID_PROVIDER_URL}/oauth/userinfo",
-                    headers={"Authorization": f"Bearer {access_token}"},
-                )
-                userinfo_response.raise_for_status()
-                user_info = userinfo_response.json()
-            except httpx.HTTPError:
-                user_info = {}
-    else:
-        user_info = {}
-
-    # For now, redirect to home with token in query params
-    # In production, use secure cookies or session storage
-    redirect_params = {
-        "access_token": access_token,
-        "username": user_info.get("preferred_username", ""),
-    }
+    if not access_token:
+        raise HTTPException(
+            status_code=500,
+            detail="Token exchange succeeded but no access_token was returned.",
+        )
 
-    return RedirectResponse(url=f"/?{urlencode(redirect_params)}")
+    # Fetch user info (optional — failure is not fatal)
+    async with httpx.AsyncClient() as client:
+        try:
+            userinfo_response = await client.get(
+                f"{OPENID_PROVIDER_URL}/oauth/userinfo",
+                headers={"Authorization": f"Bearer {access_token}"},
+            )
+            userinfo_response.raise_for_status()
+        except httpx.HTTPError:
+            pass  # user_info not required for auth flow
+
+    # Set access token as HttpOnly cookie (not in URL — avoids leaks via
+    # Referrer headers, browser history, and server logs)
+    is_production = bool(os.environ.get("SPACE_HOST"))
+    response = RedirectResponse(url="/", status_code=302)
+    response.set_cookie(
+        key="hf_access_token",
+        value=access_token,
+        httponly=True,
+        secure=is_production,  # Secure flag only in production (HTTPS)
+        samesite="lax",
+        max_age=3600 * 24,  # 24 hours
+        path="/",
+    )
+    return response
 
 
 @router.get("/logout")
 async def logout() -> RedirectResponse:
-    """Log out the user."""
-    return RedirectResponse(url="/")
+    """Log out the user by clearing the auth cookie."""
+    response = RedirectResponse(url="/")
+    response.delete_cookie(key="hf_access_token", path="/")
+    return response
 
 
-@router.get("/me")
-async def get_current_user(request: Request) -> dict:
-    """Get current user info from Authorization header."""
-    auth_header = request.headers.get("Authorization", "")
-    if not auth_header.startswith("Bearer "):
-        return {"authenticated": False}
+@router.get("/status")
+async def auth_status() -> dict:
+    """Check if OAuth is enabled on this instance."""
+    return {"auth_enabled": AUTH_ENABLED}
 
-    token = auth_header.split(" ")[1]
 
-    async with httpx.AsyncClient() as client:
-        try:
-            response = await client.get(
-                f"{OPENID_PROVIDER_URL}/oauth/userinfo",
-                headers={"Authorization": f"Bearer {token}"},
-            )
-            response.raise_for_status()
-            user_info = response.json()
-            return {
-                "authenticated": True,
-                "username": user_info.get("preferred_username"),
-                "name": user_info.get("name"),
-                "picture": user_info.get("picture"),
-            }
-        except httpx.HTTPError:
-            return {"authenticated": False}
+@router.get("/me")
+async def get_me(user: dict = Depends(get_current_user)) -> dict:
+    """Get current user info. Returns the authenticated user or dev user.
+
+    Uses the shared auth dependency which handles cookie + Bearer token.
+    """
+    return user

backend/session_manager.py

@@ -48,11 +48,27 @@ class AgentSession:
     session: Session
     tool_router: ToolRouter
     submission_queue: asyncio.Queue
+    user_id: str = "dev"  # Owner of this session
     task: asyncio.Task | None = None
     created_at: datetime = field(default_factory=datetime.utcnow)
     is_active: bool = True
 
 
+class SessionCapacityError(Exception):
+    """Raised when no more sessions can be created."""
+
+    def __init__(self, message: str, error_type: str = "global") -> None:
+        super().__init__(message)
+        self.error_type = error_type  # "global" or "per_user"
+
+
+# ── Capacity limits ─────────────────────────────────────────────────
+# Estimated for HF Spaces cpu-basic (2 vCPU, 16 GB RAM).
+# Each session uses ~10-20 MB (context, tools, queues, task).
+MAX_SESSIONS: int = 50
+MAX_SESSIONS_PER_USER: int = 3
+
+
 class SessionManager:
     """Manages multiple concurrent agent sessions."""
 
@@ -61,19 +77,66 @@ class SessionManager:
         self.sessions: dict[str, AgentSession] = {}
         self._lock = asyncio.Lock()
 
-    async def create_session(self) -> str:
-        """Create a new agent session and return its ID."""
+    def _count_user_sessions(self, user_id: str) -> int:
+        """Count active sessions owned by a specific user."""
+        return sum(
+            1
+            for s in self.sessions.values()
+            if s.user_id == user_id and s.is_active
+        )
+
+    async def create_session(self, user_id: str = "dev") -> str:
+        """Create a new agent session and return its ID.
+
+        Session() and ToolRouter() constructors contain blocking I/O
+        (e.g. HfApi().whoami(), litellm.get_max_tokens()) so they are
+        executed in a thread pool to avoid freezing the async event loop.
+
+        Args:
+            user_id: The ID of the user who owns this session.
+
+        Raises:
+            SessionCapacityError: If the server or user has reached the
+                maximum number of concurrent sessions.
+        """
+        # ── Capacity checks ──────────────────────────────────────────
+        async with self._lock:
+            active_count = self.active_session_count
+            if active_count >= MAX_SESSIONS:
+                raise SessionCapacityError(
+                    f"Server is at capacity ({active_count}/{MAX_SESSIONS} sessions). "
+                    "Please try again later.",
+                    error_type="global",
+                )
+            if user_id != "dev":
+                user_count = self._count_user_sessions(user_id)
+                if user_count >= MAX_SESSIONS_PER_USER:
+                    raise SessionCapacityError(
+                        f"You have reached the maximum of {MAX_SESSIONS_PER_USER} "
+                        "concurrent sessions. Please close an existing session first.",
+                        error_type="per_user",
+                    )
+
         session_id = str(uuid.uuid4())
 
         # Create queues for this session
         submission_queue: asyncio.Queue = asyncio.Queue()
         event_queue: asyncio.Queue = asyncio.Queue()
 
-        # Create tool router
-        tool_router = ToolRouter(self.config.mcpServers)
+        # Run blocking constructors in a thread to keep the event loop responsive.
+        # Without this, Session.__init__ → ContextManager → litellm.get_max_tokens()
+        # blocks all HTTP/WebSocket handling.
+        import time as _time
+
+        def _create_session_sync():
+            t0 = _time.monotonic()
+            tool_router = ToolRouter(self.config.mcpServers)
+            session = Session(event_queue, config=self.config, tool_router=tool_router)
+            t1 = _time.monotonic()
+            logger.info(f"Session initialized in {t1 - t0:.2f}s")
+            return tool_router, session
 
-        # Create the agent session
-        session = Session(event_queue, config=self.config, tool_router=tool_router)
+        tool_router, session = await asyncio.to_thread(_create_session_sync)
 
         # Create wrapper
         agent_session = AgentSession(
@@ -81,6 +144,7 @@ class SessionManager:
             session=session,
             tool_router=tool_router,
             submission_queue=submission_queue,
+            user_id=user_id,
         )
 
         async with self._lock:
@@ -92,7 +156,7 @@ class SessionManager:
         )
         agent_session.task = task
 
-        logger.info(f"Created session {session_id}")
+        logger.info(f"Created session {session_id} for user {user_id}")
         return session_id
 
     async def _run_session(
@@ -245,6 +309,27 @@ class SessionManager:
 
         return True
 
+    def get_session_owner(self, session_id: str) -> str | None:
+        """Get the user_id that owns a session, or None if session doesn't exist."""
+        agent_session = self.sessions.get(session_id)
+        if not agent_session:
+            return None
+        return agent_session.user_id
+
+    def verify_session_access(self, session_id: str, user_id: str) -> bool:
+        """Check if a user has access to a session.
+
+        Returns True if:
+        - The session exists AND the user owns it
+        - The user_id is "dev" (dev mode bypass)
+        """
+        owner = self.get_session_owner(session_id)
+        if owner is None:
+            return False
+        if user_id == "dev" or owner == "dev":
+            return True
+        return owner == user_id
+
     def get_session_info(self, session_id: str) -> dict[str, Any] | None:
         """Get information about a session."""
         agent_session = self.sessions.get(session_id)
@@ -256,15 +341,25 @@ class SessionManager:
             "created_at": agent_session.created_at.isoformat(),
             "is_active": agent_session.is_active,
             "message_count": len(agent_session.session.context_manager.items),
+            "user_id": agent_session.user_id,
         }
 
-    def list_sessions(self) -> list[dict[str, Any]]:
-        """List all sessions."""
-        return [
-            self.get_session_info(sid)
-            for sid in self.sessions
-            if self.get_session_info(sid)
-        ]
+    def list_sessions(self, user_id: str | None = None) -> list[dict[str, Any]]:
+        """List sessions, optionally filtered by user.
+
+        Args:
+            user_id: If provided, only return sessions owned by this user.
+                     If "dev", return all sessions (dev mode).
+        """
+        results = []
+        for sid in self.sessions:
+            info = self.get_session_info(sid)
+            if not info:
+                continue
+            if user_id and user_id != "dev" and info.get("user_id") != user_id:
+                continue
+            results.append(info)
+        return results
 
     @property
     def active_session_count(self) -> int:

backend/websocket.py

@@ -1,6 +1,5 @@
 """WebSocket connection manager for real-time communication."""
 
-import asyncio
 import logging
 from typing import Any
 
@@ -15,23 +14,18 @@ class ConnectionManager:
     def __init__(self) -> None:
         # session_id -> WebSocket
         self.active_connections: dict[str, WebSocket] = {}
-        # session_id -> asyncio.Queue for outgoing messages
-        self.message_queues: dict[str, asyncio.Queue] = {}
 
     async def connect(self, websocket: WebSocket, session_id: str) -> None:
         """Accept a WebSocket connection and register it."""
         logger.info(f"Attempting to accept WebSocket for session {session_id}")
         await websocket.accept()
         self.active_connections[session_id] = websocket
-        self.message_queues[session_id] = asyncio.Queue()
         logger.info(f"WebSocket connected and registered for session {session_id}")
 
     def disconnect(self, session_id: str) -> None:
         """Remove a WebSocket connection."""
         if session_id in self.active_connections:
             del self.active_connections[session_id]
-        if session_id in self.message_queues:
-            del self.message_queues[session_id]
         logger.info(f"WebSocket disconnected for session {session_id}")
 
     async def send_event(
@@ -63,10 +57,6 @@ class ConnectionManager:
         """Check if a session has an active WebSocket connection."""
         return session_id in self.active_connections
 
-    def get_queue(self, session_id: str) -> asyncio.Queue | None:
-        """Get the message queue for a session."""
-        return self.message_queues.get(session_id)
-
 
 # Global connection manager instance
 manager = ConnectionManager()

frontend/src/App.tsx

@@ -1,7 +1,12 @@
 import { Box } from '@mui/material';
 import AppLayout from '@/components/Layout/AppLayout';
+import { useAuth } from '@/hooks/useAuth';
 
 function App() {
+  // Non-blocking auth check — fires in background, updates store when done.
+  // If auth fails later, apiFetch redirects to /auth/login.
+  useAuth();
+
   return (
     <Box sx={{ height: '100vh', display: 'flex' }}>
       <AppLayout />

frontend/src/components/ApprovalModal/ApprovalModal.tsx

@@ -1,208 +0,0 @@
-import { useState, useCallback } from 'react';
-import {
-  Dialog,
-  DialogTitle,
-  DialogContent,
-  DialogActions,
-  Button,
-  Box,
-  Typography,
-  Checkbox,
-  FormControlLabel,
-  Accordion,
-  AccordionSummary,
-  AccordionDetails,
-  TextField,
-  Chip,
-} from '@mui/material';
-import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
-import WarningIcon from '@mui/icons-material/Warning';
-import { useAgentStore } from '@/store/agentStore';
-
-interface ApprovalModalProps {
-  sessionId: string | null;
-}
-
-interface ApprovalState {
-  [toolCallId: string]: {
-    approved: boolean;
-    feedback: string;
-  };
-}
-
-export default function ApprovalModal({ sessionId }: ApprovalModalProps) {
-  const { pendingApprovals, setPendingApprovals } = useAgentStore();
-  const [approvalState, setApprovalState] = useState<ApprovalState>({});
-
-  const isOpen = pendingApprovals !== null && pendingApprovals.tools.length > 0;
-
-  const handleApprovalChange = useCallback(
-    (toolCallId: string, approved: boolean) => {
-      setApprovalState((prev) => ({
-        ...prev,
-        [toolCallId]: {
-          ...prev[toolCallId],
-          approved,
-          feedback: prev[toolCallId]?.feedback || '',
-        },
-      }));
-    },
-    []
-  );
-
-  const handleFeedbackChange = useCallback(
-    (toolCallId: string, feedback: string) => {
-      setApprovalState((prev) => ({
-        ...prev,
-        [toolCallId]: {
-          ...prev[toolCallId],
-          feedback,
-        },
-      }));
-    },
-    []
-  );
-
-  const handleSubmit = useCallback(async () => {
-    if (!sessionId || !pendingApprovals) return;
-
-    const approvals = pendingApprovals.tools.map((tool) => ({
-      tool_call_id: tool.tool_call_id,
-      approved: approvalState[tool.tool_call_id]?.approved ?? false,
-      feedback: approvalState[tool.tool_call_id]?.feedback || null,
-    }));
-
-    try {
-      await fetch('/api/approve', {
-        method: 'POST',
-        headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({
-          session_id: sessionId,
-          approvals,
-        }),
-      });
-      setPendingApprovals(null);
-      setApprovalState({});
-    } catch (e) {
-      console.error('Approval submission failed:', e);
-    }
-  }, [sessionId, pendingApprovals, approvalState, setPendingApprovals]);
-
-  const handleApproveAll = useCallback(() => {
-    if (!pendingApprovals) return;
-    const newState: ApprovalState = {};
-    pendingApprovals.tools.forEach((tool) => {
-      newState[tool.tool_call_id] = { approved: true, feedback: '' };
-    });
-    setApprovalState(newState);
-  }, [pendingApprovals]);
-
-  const handleRejectAll = useCallback(() => {
-    if (!pendingApprovals) return;
-    const newState: ApprovalState = {};
-    pendingApprovals.tools.forEach((tool) => {
-      newState[tool.tool_call_id] = { approved: false, feedback: '' };
-    });
-    setApprovalState(newState);
-  }, [pendingApprovals]);
-
-  if (!isOpen || !pendingApprovals) return null;
-
-  const approvedCount = Object.values(approvalState).filter((s) => s.approved).length;
-
-  return (
-    <Dialog
-      open={isOpen}
-      maxWidth="md"
-      fullWidth
-      PaperProps={{
-        sx: { bgcolor: 'background.paper' },
-      }}
-    >
-      <DialogTitle sx={{ display: 'flex', alignItems: 'center', gap: 1 }}>
-        <WarningIcon color="warning" />
-        Approval Required
-        <Chip
-          label={`${pendingApprovals.count} tool${pendingApprovals.count > 1 ? 's' : ''}`}
-          size="small"
-          sx={{ ml: 1 }}
-        />
-      </DialogTitle>
-      <DialogContent dividers>
-        <Typography variant="body2" color="text.secondary" sx={{ mb: 2 }}>
-          The following tool calls require your approval before execution:
-        </Typography>
-        {pendingApprovals.tools.map((tool, index) => (
-          <Accordion key={tool.tool_call_id} defaultExpanded={index === 0}>
-            <AccordionSummary expandIcon={<ExpandMoreIcon />}>
-              <Box sx={{ display: 'flex', alignItems: 'center', gap: 2, width: '100%' }}>
-                <FormControlLabel
-                  control={
-                    <Checkbox
-                      checked={approvalState[tool.tool_call_id]?.approved ?? false}
-                      onChange={(e) => {
-                        e.stopPropagation();
-                        handleApprovalChange(tool.tool_call_id, e.target.checked);
-                      }}
-                      onClick={(e) => e.stopPropagation()}
-                    />
-                  }
-                  label=""
-                  sx={{ m: 0 }}
-                />
-                <Chip label={tool.tool} size="small" color="primary" variant="outlined" />
-                <Typography variant="body2" color="text.secondary" sx={{ ml: 'auto' }}>
-                  {approvalState[tool.tool_call_id]?.approved ? 'Approved' : 'Pending'}
-                </Typography>
-              </Box>
-            </AccordionSummary>
-            <AccordionDetails>
-              <Typography variant="subtitle2" gutterBottom>
-                Arguments:
-              </Typography>
-              <Box
-                component="pre"
-                sx={{
-                  bgcolor: 'background.default',
-                  p: 1.5,
-                  borderRadius: 1,
-                  overflow: 'auto',
-                  fontSize: '0.8rem',
-                  maxHeight: 200,
-                }}
-              >
-                {JSON.stringify(tool.arguments, null, 2)}
-              </Box>
-              {!approvalState[tool.tool_call_id]?.approved && (
-                <TextField
-                  fullWidth
-                  size="small"
-                  label="Feedback (optional)"
-                  placeholder="Explain why you're rejecting this..."
-                  value={approvalState[tool.tool_call_id]?.feedback || ''}
-                  onChange={(e) => handleFeedbackChange(tool.tool_call_id, e.target.value)}
-                  sx={{ mt: 2 }}
-                />
-              )}
-            </AccordionDetails>
-          </Accordion>
-        ))}
-      </DialogContent>
-      <DialogActions sx={{ px: 3, py: 2 }}>
-        <Button onClick={handleRejectAll} color="error" variant="outlined">
-          Reject All
-        </Button>
-        <Button onClick={handleApproveAll} color="success" variant="outlined">
-          Approve All
-        </Button>
-        <Box sx={{ flex: 1 }} />
-        <Typography variant="body2" color="text.secondary" sx={{ mr: 2 }}>
-          {approvedCount} of {pendingApprovals.count} approved
-        </Typography>
-        <Button onClick={handleSubmit} variant="contained" color="primary">
-          Submit
-        </Button>
-      </DialogActions>
-    </Dialog>
-  );
-}

frontend/src/components/Chat/ApprovalFlow.tsx

@@ -1,515 +0,0 @@
-import { useState, useCallback, useEffect } from 'react';
-import { Box, Typography, Button, TextField, IconButton, Link } from '@mui/material';
-import SendIcon from '@mui/icons-material/Send';
-import OpenInNewIcon from '@mui/icons-material/OpenInNew';
-import CheckCircleIcon from '@mui/icons-material/CheckCircle';
-import CancelIcon from '@mui/icons-material/Cancel';
-import LaunchIcon from '@mui/icons-material/Launch';
-import { useAgentStore } from '@/store/agentStore';
-import { useLayoutStore } from '@/store/layoutStore';
-import { useSessionStore } from '@/store/sessionStore';
-import type { Message, ToolApproval } from '@/types/agent';
-
-interface ApprovalFlowProps {
-  message: Message;
-}
-
-export default function ApprovalFlow({ message }: ApprovalFlowProps) {
-  const { setPanelContent, setPanelTab, setActivePanelTab, clearPanelTabs, updateMessage } = useAgentStore();
-  const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
-  const { activeSessionId } = useSessionStore();
-  const [currentIndex, setCurrentIndex] = useState(0);
-  const [feedback, setFeedback] = useState('');
-  const [decisions, setDecisions] = useState<ToolApproval[]>([]);
-
-  const approvalData = message.approval;
-  
-  if (!approvalData) return null;
-
-  const { batch, status } = approvalData;
-
-  // Parse toolOutput to extract job info (URL, status, logs, errors)
-  let logsContent = '';
-  let showLogsButton = false;
-  let jobUrl = '';
-  let jobStatus = '';
-  let jobFailed = false;
-  let errorMessage = '';
-
-  if (message.toolOutput) {
-    const output = message.toolOutput;
-
-    // Extract job URL: **View at:** https://...
-    const urlMatch = output.match(/\*\*View at:\*\*\s*(https:\/\/[^\s\n]+)/);
-    if (urlMatch) {
-      jobUrl = urlMatch[1];
-    }
-
-    // Extract job status: **Final Status:** ...
-    const statusMatch = output.match(/\*\*Final Status:\*\*\s*([^\n]+)/);
-    if (statusMatch) {
-      jobStatus = statusMatch[1].trim();
-      jobFailed = jobStatus.toLowerCase().includes('error') || jobStatus.toLowerCase().includes('failed');
-    }
-
-    // Extract logs
-    if (output.includes('**Logs:**')) {
-      const parts = output.split('**Logs:**');
-      if (parts.length > 1) {
-        const logsPart = parts[1].trim();
-        const codeBlockMatch = logsPart.match(/```([\s\S]*?)```/);
-        if (codeBlockMatch) {
-          logsContent = codeBlockMatch[1].trim();
-          showLogsButton = true;
-        }
-      }
-    }
-
-    // Detect errors - if output exists but doesn't have the expected job completion format
-    // This catches early failures (validation errors, API errors, etc.)
-    const isExpectedFormat = output.includes('**Job ID:**') || output.includes('**View at:**');
-    const looksLikeError = output.toLowerCase().includes('error') ||
-                          output.toLowerCase().includes('failed') ||
-                          output.toLowerCase().includes('exception') ||
-                          output.includes('Traceback');
-
-    if (!isExpectedFormat || (looksLikeError && !logsContent)) {
-      // This is likely an error message - show it
-      errorMessage = output;
-      jobFailed = true;
-    }
-  }
-
-  // Sync right panel with current tool
-  useEffect(() => {
-    if (!batch || currentIndex >= batch.tools.length) return;
-    
-    // Only auto-open panel if pending
-    if (status !== 'pending') return;
-    
-    const tool = batch.tools[currentIndex];
-    const args = tool.arguments as any;
-
-    if (tool.tool === 'hf_jobs' && (args.operation === 'run' || args.operation === 'scheduled run') && args.script) {
-      setPanelContent({
-        title: 'Compute Job Script',
-        content: args.script,
-        language: 'python',
-        parameters: args
-      });
-      // Don't auto-open if already resolved
-    } else if (tool.tool === 'hf_repo_files' && args.operation === 'upload' && args.content) {
-      setPanelContent({
-        title: `File Upload: ${args.path || 'unnamed'}`,
-        content: args.content,
-        parameters: args
-      });
-    }
-  }, [currentIndex, batch, status, setPanelContent]);
-
-  const handleResolve = useCallback(async (approved: boolean) => {
-    if (!batch || !activeSessionId) return;
-
-    const currentTool = batch.tools[currentIndex];
-    const newDecisions = [
-      ...decisions,
-      {
-        tool_call_id: currentTool.tool_call_id,
-        approved,
-        feedback: approved ? null : feedback || 'Rejected by user',
-      },
-    ];
-
-    if (currentIndex < batch.tools.length - 1) {
-      setDecisions(newDecisions);
-      setCurrentIndex(currentIndex + 1);
-      setFeedback('');
-    } else {
-      // All tools in batch resolved
-      try {
-        await fetch('/api/approve', {
-          method: 'POST',
-          headers: { 'Content-Type': 'application/json' },
-          body: JSON.stringify({
-            session_id: activeSessionId,
-            approvals: newDecisions,
-          }),
-        });
-        
-        // Update message status
-        updateMessage(activeSessionId, message.id, {
-            approval: {
-                ...approvalData!,
-                status: approved ? 'approved' : 'rejected',
-                decisions: newDecisions
-            }
-        });
-        
-      } catch (e) {
-        console.error('Approval submission failed:', e);
-      }
-    }
-  }, [activeSessionId, message.id, batch, currentIndex, feedback, decisions, approvalData, updateMessage]);
-
-  if (!batch || currentIndex >= batch.tools.length) return null;
-
-  const currentTool = batch.tools[currentIndex];
-
-  // Check if script contains push_to_hub or upload_file
-  const args = currentTool.arguments as any;
-  const containsPushToHub = currentTool.tool === 'hf_jobs' && args.script && (args.script.includes('push_to_hub') || args.script.includes('upload_file'));
-
-  const getToolDescription = (toolName: string, args: any) => {
-    if (toolName === 'hf_jobs') {
-      return (
-        <Box sx={{ flex: 1 }}>
-          <Typography variant="body2" sx={{ color: 'var(--muted-text)' }}>
-            The agent wants to execute <Box component="span" sx={{ color: 'var(--accent-yellow)', fontWeight: 500 }}>hf_jobs</Box> on{' '}
-            <Box component="span" sx={{ fontWeight: 500, color: 'var(--text)' }}>{args.hardware_flavor || 'default'}</Box> with a timeout of{' '}
-            <Box component="span" sx={{ fontWeight: 500, color: 'var(--text)' }}>{args.timeout || '30m'}</Box>
-          </Typography>
-        </Box>
-      );
-    }
-    return (
-      <Typography variant="body2" sx={{ color: 'var(--muted-text)', flex: 1 }}>
-        The agent wants to execute <Box component="span" sx={{ color: 'var(--accent-yellow)', fontWeight: 500 }}>{toolName}</Box>
-      </Typography>
-    );
-  };
-
-  const showCode = () => {
-    const args = currentTool.arguments as any;
-    if (currentTool.tool === 'hf_jobs' && args.script) {
-      // Clear existing tabs and set up script tab (and logs if available)
-      clearPanelTabs();
-      setPanelTab({
-        id: 'script',
-        title: 'Script',
-        content: args.script,
-        language: 'python',
-        parameters: args
-      });
-      // If logs are available (job completed), also add logs tab
-      if (logsContent) {
-        setPanelTab({
-          id: 'logs',
-          title: 'Logs',
-          content: logsContent,
-          language: 'text'
-        });
-      }
-      setActivePanelTab('script');
-      setRightPanelOpen(true);
-      setLeftSidebarOpen(false);
-    } else {
-      setPanelContent({
-        title: `Tool: ${currentTool.tool}`,
-        content: JSON.stringify(args, null, 2),
-        language: 'json',
-        parameters: args
-      });
-      setRightPanelOpen(true);
-      setLeftSidebarOpen(false);
-    }
-  };
-
-  const handleViewLogs = (e: React.MouseEvent) => {
-    e.stopPropagation();
-    const args = currentTool.arguments as any;
-    // Set up both tabs so user can switch between script and logs
-    clearPanelTabs();
-    if (currentTool.tool === 'hf_jobs' && args.script) {
-      setPanelTab({
-        id: 'script',
-        title: 'Script',
-        content: args.script,
-        language: 'python',
-        parameters: args
-      });
-    }
-    setPanelTab({
-      id: 'logs',
-      title: 'Logs',
-      content: logsContent,
-      language: 'text'
-    });
-    setActivePanelTab('logs');
-    setRightPanelOpen(true);
-    setLeftSidebarOpen(false);
-  };
-
-  return (
-    <Box 
-      className="action-card"
-      sx={{ 
-        width: '100%',
-        padding: '18px',
-        borderRadius: 'var(--radius-md)',
-        background: 'linear-gradient(180deg, rgba(255,255,255,0.015), transparent)',
-        border: '1px solid rgba(255,255,255,0.03)',
-        display: 'flex',
-        flexDirection: 'column',
-        gap: '12px',
-        opacity: status !== 'pending' && !showLogsButton ? 0.8 : 1
-      }}
-    >
-      <Box sx={{ display: 'flex', alignItems: 'center', gap: 1 }}>
-         <Typography variant="subtitle2" sx={{ fontWeight: 600, color: 'var(--text)' }}>
-            {status === 'pending' ? 'Approval Required' : status === 'approved' ? 'Approved' : 'Rejected'}
-         </Typography>
-         <Typography variant="caption" sx={{ color: 'var(--muted-text)' }}>
-            ({currentIndex + 1}/{batch.count})
-         </Typography>
-         {status === 'approved' && <CheckCircleIcon sx={{ fontSize: 18, color: 'var(--accent-green)' }} />}
-         {status === 'rejected' && <CancelIcon sx={{ fontSize: 18, color: 'var(--accent-red)' }} />}
-      </Box>
-
-      <Box 
-        onClick={showCode}
-        sx={{ 
-            display: 'flex', 
-            alignItems: 'center', 
-            gap: 1, 
-            cursor: 'pointer',
-            p: 1.5,
-            borderRadius: '8px',
-            bgcolor: 'rgba(0,0,0,0.2)',
-            border: '1px solid rgba(255,255,255,0.05)',
-            transition: 'all 0.2s',
-            '&:hover': {
-                bgcolor: 'rgba(255,255,255,0.03)',
-                borderColor: 'var(--accent-primary)',
-            }
-        }}
-      >
-        {getToolDescription(currentTool.tool, currentTool.arguments)}
-        <OpenInNewIcon sx={{ fontSize: 16, color: 'var(--muted-text)', opacity: 0.7 }} />
-      </Box>
-
-      {/* Script/Logs buttons for hf_jobs - always show when we have a script */}
-      {currentTool.tool === 'hf_jobs' && args.script && (
-        <Box sx={{ display: 'flex', flexDirection: 'column', gap: 1 }}>
-          <Box sx={{ display: 'flex', gap: 1, flexWrap: 'wrap' }}>
-            <Button
-              variant="outlined"
-              size="small"
-              onClick={showCode}
-              sx={{
-                textTransform: 'none',
-                borderColor: 'rgba(255,255,255,0.1)',
-                color: 'var(--muted-text)',
-                fontSize: '0.75rem',
-                py: 0.5,
-                '&:hover': {
-                  borderColor: 'var(--accent-primary)',
-                  color: 'var(--accent-primary)',
-                  bgcolor: 'rgba(255,255,255,0.03)'
-                }
-              }}
-            >
-              View Script
-            </Button>
-            <Button
-              variant="outlined"
-              size="small"
-              onClick={handleViewLogs}
-              disabled={!logsContent && status === 'pending'}
-              sx={{
-                textTransform: 'none',
-                borderColor: 'rgba(255,255,255,0.1)',
-                color: logsContent ? 'var(--accent-primary)' : 'var(--muted-text)',
-                fontSize: '0.75rem',
-                py: 0.5,
-                '&:hover': {
-                  borderColor: 'var(--accent-primary)',
-                  bgcolor: 'rgba(255,255,255,0.03)'
-                },
-                '&.Mui-disabled': {
-                  color: 'rgba(255,255,255,0.3)',
-                  borderColor: 'rgba(255,255,255,0.05)',
-                }
-              }}
-            >
-              {logsContent ? 'View Logs' : 'Logs (waiting for job...)'}
-            </Button>
-          </Box>
-
-          {/* Job URL - only show when we have a specific URL */}
-          {jobUrl && (
-            <Link
-              href={jobUrl}
-              target="_blank"
-              rel="noopener noreferrer"
-              sx={{
-                display: 'flex',
-                alignItems: 'center',
-                gap: 0.5,
-                color: 'var(--accent-primary)',
-                fontSize: '0.75rem',
-                textDecoration: 'none',
-                opacity: 0.9,
-                '&:hover': {
-                  opacity: 1,
-                  textDecoration: 'underline',
-                }
-              }}
-            >
-              <LaunchIcon sx={{ fontSize: 14 }} />
-              View Job on Hugging Face
-            </Link>
-          )}
-
-          {/* Show job status if available */}
-          {jobStatus && (
-            <Typography
-              variant="caption"
-              sx={{
-                color: jobFailed ? 'var(--accent-red)' : 'var(--accent-green)',
-                fontSize: '0.75rem',
-                fontWeight: 500,
-              }}
-            >
-              Status: {jobStatus}
-            </Typography>
-          )}
-        </Box>
-      )}
-
-      {containsPushToHub && (
-        <Typography variant="caption" sx={{ color: 'var(--accent-green)', fontSize: '0.75rem', opacity: 0.8, px: 0.5 }}>
-          We've detected the result will be pushed to hub.
-        </Typography>
-      )}
-
-      {/* Show error message if job failed */}
-      {errorMessage && status !== 'pending' && (
-        <Box
-          sx={{
-            p: 1.5,
-            borderRadius: '8px',
-            bgcolor: 'rgba(224, 90, 79, 0.1)',
-            border: '1px solid rgba(224, 90, 79, 0.3)',
-          }}
-        >
-          <Typography
-            variant="caption"
-            sx={{
-              color: 'var(--accent-red)',
-              fontWeight: 600,
-              display: 'block',
-              mb: 0.5,
-            }}
-          >
-            Error
-          </Typography>
-          <Typography
-            component="pre"
-            sx={{
-              color: 'var(--text)',
-              fontSize: '0.75rem',
-              fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
-              whiteSpace: 'pre-wrap',
-              wordBreak: 'break-word',
-              m: 0,
-              maxHeight: '150px',
-              overflow: 'auto',
-            }}
-          >
-            {errorMessage.length > 500 ? errorMessage.substring(0, 500) + '...' : errorMessage}
-          </Typography>
-        </Box>
-      )}
-
-
-      {status === 'pending' && (
-      <Box sx={{ display: 'flex', flexDirection: 'column', gap: 2 }}>
-        <Box sx={{ display: 'flex', gap: 1 }}>
-            <TextField
-                fullWidth
-                size="small"
-                placeholder="Feedback (optional)"
-                value={feedback}
-                onChange={(e) => setFeedback(e.target.value)}
-                variant="outlined"
-                sx={{ 
-                    '& .MuiOutlinedInput-root': { 
-                        bgcolor: 'rgba(0,0,0,0.2)',
-                        fontFamily: 'inherit',
-                        fontSize: '0.9rem' 
-                    }
-                }}
-            />
-            <IconButton 
-                onClick={() => handleResolve(false)}
-                disabled={!feedback}
-                title="Reject with feedback"
-                sx={{ 
-                    color: 'var(--accent-red)',
-                    border: '1px solid rgba(255,255,255,0.05)',
-                    borderRadius: '8px',
-                    width: 40,
-                    height: 40,
-                    '&:hover': {
-                        bgcolor: 'rgba(224, 90, 79, 0.1)',
-                        borderColor: 'var(--accent-red)',
-                    },
-                    '&.Mui-disabled': {
-                        color: 'rgba(255,255,255,0.1)',
-                        borderColor: 'rgba(255,255,255,0.02)'
-                    }
-                }}
-            >
-                <SendIcon fontSize="small" />
-            </IconButton>
-        </Box>
-        
-        <Box className="action-buttons" sx={{ display: 'flex', gap: '10px' }}>
-            <Button 
-            className="btn-reject"
-            onClick={() => handleResolve(false)}
-            sx={{ 
-                flex: 1,
-                background: 'transparent',
-                border: '1px solid rgba(255,255,255,0.05)',
-                color: 'var(--accent-red)',
-                padding: '10px 14px',
-                borderRadius: '10px',
-                '&:hover': {
-                    bgcolor: 'rgba(224, 90, 79, 0.05)',
-                    borderColor: 'var(--accent-red)',
-                }
-            }}
-            >
-            Reject
-            </Button>
-            <Button 
-            className="btn-approve"
-            onClick={() => handleResolve(true)}
-            sx={{ 
-                flex: 1,
-                background: 'transparent',
-                border: '1px solid rgba(255,255,255,0.05)',
-                color: 'var(--accent-green)',
-                padding: '10px 14px',
-                borderRadius: '10px',
-                '&:hover': {
-                    bgcolor: 'rgba(47, 204, 113, 0.05)',
-                    borderColor: 'var(--accent-green)',
-                }
-            }}
-            >
-            Approve
-            </Button>
-        </Box>
-      </Box>
-      )}
-      
-      {status === 'rejected' && decisions.some(d => d.feedback) && (
-        <Typography variant="body2" sx={{ color: 'var(--accent-red)', mt: 1 }}>
-            Feedback: {decisions.find(d => d.feedback)?.feedback}
-        </Typography>
-      )}
-    </Box>
-  );
-}

frontend/src/components/Chat/AssistantMessage.tsx

@@ -0,0 +1,108 @@
+import { Box, Stack, Avatar, Typography } from '@mui/material';
+import SmartToyOutlinedIcon from '@mui/icons-material/SmartToyOutlined';
+import MarkdownContent from './MarkdownContent';
+import ToolCallGroup from './ToolCallGroup';
+import type { Message } from '@/types/agent';
+
+interface AssistantMessageProps {
+  message: Message;
+  /** True when this message is actively receiving streaming chunks. */
+  isStreaming?: boolean;
+}
+
+export default function AssistantMessage({ message, isStreaming = false }: AssistantMessageProps) {
+  const renderSegments = () => {
+    if (message.segments && message.segments.length > 0) {
+      // Find the index of the last text segment (that's the one being streamed)
+      let lastTextIdx = -1;
+      for (let i = message.segments.length - 1; i >= 0; i--) {
+        if (message.segments[i].type === 'text') {
+          lastTextIdx = i;
+          break;
+        }
+      }
+
+      return message.segments.map((segment, idx) => {
+        if (segment.type === 'text' && segment.content) {
+          return (
+            <MarkdownContent
+              key={idx}
+              content={segment.content}
+              isStreaming={isStreaming && idx === lastTextIdx}
+            />
+          );
+        }
+        if (segment.type === 'tools' && segment.tools && segment.tools.length > 0) {
+          return <ToolCallGroup key={idx} tools={segment.tools} />;
+        }
+        return null;
+      });
+    }
+
+    // Fallback: render raw content
+    if (message.content) {
+      return <MarkdownContent content={message.content} isStreaming={isStreaming} />;
+    }
+
+    return null;
+  };
+
+  return (
+    <Stack direction="row" spacing={1.5} alignItems="flex-start">
+      <Avatar
+        sx={{
+          width: 28,
+          height: 28,
+          bgcolor: 'primary.main',
+          flexShrink: 0,
+          mt: 0.5,
+        }}
+      >
+        <SmartToyOutlinedIcon sx={{ fontSize: 16, color: '#fff' }} />
+      </Avatar>
+
+      <Box sx={{ flex: 1, minWidth: 0 }}>
+        {/* Role label + timestamp */}
+        <Stack direction="row" alignItems="baseline" spacing={1} sx={{ mb: 0.5 }}>
+          <Typography
+            variant="caption"
+            sx={{
+              fontWeight: 700,
+              fontSize: '0.72rem',
+              color: 'var(--muted-text)',
+              textTransform: 'uppercase',
+              letterSpacing: '0.04em',
+            }}
+          >
+            Assistant
+          </Typography>
+          <Typography
+            variant="caption"
+            sx={{
+              fontSize: '0.66rem',
+              color: 'var(--muted-text)',
+              opacity: 0.6,
+            }}
+          >
+            {new Date(message.timestamp).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}
+          </Typography>
+        </Stack>
+
+        {/* Message bubble */}
+        <Box
+          sx={{
+            maxWidth: { xs: '95%', md: '85%' },
+            bgcolor: 'var(--surface)',
+            borderRadius: 1.5,
+            borderTopLeftRadius: 4,
+            px: { xs: 1.5, md: 2.5 },
+            py: 1.5,
+            border: '1px solid var(--border)',
+          }}
+        >
+          {renderSegments()}
+        </Box>
+      </Box>
+    </Stack>
+  );
+}

frontend/src/components/Chat/ChatInput.tsx

@@ -1,4 +1,4 @@
-import { useState, useCallback, KeyboardEvent } from 'react';
+import { useState, useCallback, useEffect, useRef, KeyboardEvent } from 'react';
 import { Box, TextField, IconButton, CircularProgress, Typography } from '@mui/material';
 import ArrowUpwardIcon from '@mui/icons-material/ArrowUpward';
 
@@ -9,6 +9,14 @@ interface ChatInputProps {
 
 export default function ChatInput({ onSend, disabled = false }: ChatInputProps) {
   const [input, setInput] = useState('');
+  const inputRef = useRef<HTMLTextAreaElement>(null);
+
+  // Auto-focus the textarea when the session becomes ready (disabled → false)
+  useEffect(() => {
+    if (!disabled && inputRef.current) {
+      inputRef.current.focus();
+    }
+  }, [disabled]);
 
   const handleSend = useCallback(() => {
     if (input.trim() && !disabled) {
@@ -30,23 +38,23 @@ export default function ChatInput({ onSend, disabled = false }: ChatInputProps)
   return (
     <Box
       sx={{
-        pb: 4,
-        pt: 2,
+        pb: { xs: 2, md: 4 },
+        pt: { xs: 1, md: 2 },
         position: 'relative',
         zIndex: 10,
       }}
     >
-      <Box sx={{ maxWidth: '880px', mx: 'auto', width: '100%', px: 2 }}>
+      <Box sx={{ maxWidth: '880px', mx: 'auto', width: '100%', px: { xs: 0, sm: 1, md: 2 } }}>
         <Box
           className="composer"
           sx={{
             display: 'flex',
             gap: '10px',
             alignItems: 'flex-start',
-            bgcolor: 'rgba(255,255,255,0.01)',
+            bgcolor: 'var(--composer-bg)',
             borderRadius: 'var(--radius-md)',
             p: '12px',
-            border: '1px solid rgba(255,255,255,0.03)',
+            border: '1px solid var(--border)',
             transition: 'box-shadow 0.2s ease, border-color 0.2s ease',
             '&:focus-within': {
                 borderColor: 'var(--accent-yellow)',
@@ -64,6 +72,7 @@ export default function ChatInput({ onSend, disabled = false }: ChatInputProps)
             placeholder="Ask anything..."
             disabled={disabled}
             variant="standard"
+            inputRef={inputRef}
             InputProps={{
                 disableUnderline: true,
                 sx: {
@@ -72,7 +81,7 @@ export default function ChatInput({ onSend, disabled = false }: ChatInputProps)
                     fontFamily: 'inherit',
                     padding: 0,
                     lineHeight: 1.5,
-                    minHeight: '56px',
+                    minHeight: { xs: '44px', md: '56px' },
                     alignItems: 'flex-start',
                 }
             }}
@@ -99,7 +108,7 @@ export default function ChatInput({ onSend, disabled = false }: ChatInputProps)
               transition: 'all 0.2s',
               '&:hover': {
                 color: 'var(--accent-yellow)',
-                bgcolor: 'rgba(255,255,255,0.05)',
+                bgcolor: 'var(--hover-bg)',
               },
               '&.Mui-disabled': {
                 opacity: 0.3,
@@ -115,7 +124,7 @@ export default function ChatInput({ onSend, disabled = false }: ChatInputProps)
           <Typography variant="caption" sx={{ fontSize: '10px', color: 'var(--muted-text)', textTransform: 'uppercase', letterSpacing: '0.05em', fontWeight: 500 }}>
             powered by
           </Typography>
-          <img src="/claude-logo.png" alt="Claude" style={{ height: '12px', objectFit: 'contain' }} />
+          <Box component="img" src="/claude-logo.png" alt="Claude" sx={{ height: '12px', objectFit: 'contain' }} />
           <Typography variant="caption" sx={{ fontSize: '10px', color: 'var(--text)', fontWeight: 600, letterSpacing: '0.02em' }}>
             claude-opus-4-5-20251101
           </Typography>

frontend/src/components/Chat/MarkdownContent.tsx

@@ -0,0 +1,181 @@
+import { useMemo, useRef, useState, useEffect } from 'react';
+import { Box } from '@mui/material';
+import ReactMarkdown from 'react-markdown';
+import remarkGfm from 'remark-gfm';
+import type { SxProps, Theme } from '@mui/material/styles';
+
+interface MarkdownContentProps {
+  content: string;
+  sx?: SxProps<Theme>;
+  /** When true, shows a blinking cursor and throttles renders. */
+  isStreaming?: boolean;
+}
+
+/** Shared markdown styles — adapts to light/dark via CSS variables. */
+const markdownSx: SxProps<Theme> = {
+  fontSize: '0.925rem',
+  lineHeight: 1.7,
+  color: 'var(--text)',
+  wordBreak: 'break-word',
+
+  '& p': { m: 0, mb: 1.5, '&:last-child': { mb: 0 } },
+
+  '& h1, & h2, & h3, & h4': { mt: 2.5, mb: 1, fontWeight: 600, lineHeight: 1.3 },
+  '& h1': { fontSize: '1.35rem' },
+  '& h2': { fontSize: '1.15rem' },
+  '& h3': { fontSize: '1.05rem' },
+
+  '& pre': {
+    bgcolor: 'var(--code-bg)',
+    p: 2,
+    borderRadius: 2,
+    overflow: 'auto',
+    fontSize: '0.82rem',
+    lineHeight: 1.6,
+    border: '1px solid var(--tool-border)',
+    my: 2,
+  },
+  '& code': {
+    bgcolor: 'var(--hover-bg)',
+    px: 0.75,
+    py: 0.25,
+    borderRadius: 0.5,
+    fontSize: '0.84rem',
+    fontFamily: '"JetBrains Mono", ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
+  },
+  '& pre code': { bgcolor: 'transparent', p: 0 },
+
+  '& a': {
+    color: 'var(--accent-yellow)',
+    textDecoration: 'none',
+    fontWeight: 500,
+    '&:hover': { textDecoration: 'underline' },
+  },
+
+  '& ul, & ol': { pl: 3, my: 1 },
+  '& li': { mb: 0.5 },
+  '& li::marker': { color: 'var(--muted-text)' },
+
+  '& blockquote': {
+    borderLeft: '3px solid var(--accent-yellow)',
+    pl: 2,
+    ml: 0,
+    my: 1.5,
+    color: 'var(--muted-text)',
+    fontStyle: 'italic',
+  },
+
+  '& table': {
+    borderCollapse: 'collapse',
+    width: '100%',
+    my: 2,
+    fontSize: '0.85rem',
+  },
+  '& th': {
+    borderBottom: '2px solid var(--border-hover)',
+    textAlign: 'left',
+    p: 1,
+    fontWeight: 600,
+  },
+  '& td': {
+    borderBottom: '1px solid var(--tool-border)',
+    p: 1,
+  },
+
+  '& hr': {
+    border: 'none',
+    borderTop: '1px solid var(--border)',
+    my: 2,
+  },
+
+  '& img': {
+    maxWidth: '100%',
+    borderRadius: 2,
+  },
+};
+
+/** Blinking cursor shown at the end of streaming text. */
+const StreamingCursor = () => (
+  <Box
+    component="span"
+    sx={{
+      display: 'inline-block',
+      width: '2px',
+      height: '1.1em',
+      bgcolor: 'var(--text)',
+      ml: '2px',
+      verticalAlign: 'text-bottom',
+      animation: 'cursorBlink 1s step-end infinite',
+      '@keyframes cursorBlink': {
+        '0%, 100%': { opacity: 1 },
+        '50%': { opacity: 0 },
+      },
+    }}
+  />
+);
+
+/**
+ * Throttled content for streaming: render the full markdown through
+ * ReactMarkdown but only re-parse every ~80ms to avoid layout thrashing.
+ * This is the Claude approach — always render as markdown, never split
+ * into raw text. The parser handles incomplete tables gracefully.
+ */
+function useThrottledValue(value: string, isStreaming: boolean, intervalMs = 80): string {
+  const [throttled, setThrottled] = useState(value);
+  const lastUpdate = useRef(0);
+  const pending = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const latestValue = useRef(value);
+  latestValue.current = value;
+
+  useEffect(() => {
+    if (!isStreaming) {
+      // Not streaming — always use latest value immediately
+      setThrottled(value);
+      return;
+    }
+
+    const now = Date.now();
+    const elapsed = now - lastUpdate.current;
+
+    if (elapsed >= intervalMs) {
+      // Enough time passed — update immediately
+      setThrottled(value);
+      lastUpdate.current = now;
+    } else {
+      // Schedule an update for the remaining time
+      if (pending.current) clearTimeout(pending.current);
+      pending.current = setTimeout(() => {
+        setThrottled(latestValue.current);
+        lastUpdate.current = Date.now();
+        pending.current = null;
+      }, intervalMs - elapsed);
+    }
+
+    return () => {
+      if (pending.current) clearTimeout(pending.current);
+    };
+  }, [value, isStreaming, intervalMs]);
+
+  // When streaming ends, flush immediately
+  useEffect(() => {
+    if (!isStreaming) {
+      setThrottled(latestValue.current);
+    }
+  }, [isStreaming]);
+
+  return throttled;
+}
+
+export default function MarkdownContent({ content, sx, isStreaming = false }: MarkdownContentProps) {
+  // Throttle re-parses during streaming to ~12fps (every 80ms)
+  const displayContent = useThrottledValue(content, isStreaming);
+
+  const remarkPlugins = useMemo(() => [remarkGfm], []);
+
+  return (
+    <Box sx={[markdownSx, ...(Array.isArray(sx) ? sx : sx ? [sx] : [])]}>
+      <ReactMarkdown remarkPlugins={remarkPlugins}>{displayContent}</ReactMarkdown>
+      {isStreaming && <StreamingCursor />}
+    </Box>
+  );
+}

frontend/src/components/Chat/MessageBubble.tsx

@@ -1,215 +1,51 @@
-import { Box, Paper, Typography } from '@mui/material';
-import ReactMarkdown from 'react-markdown';
-import remarkGfm from 'remark-gfm';
-import ApprovalFlow from './ApprovalFlow';
-import type { Message, TraceLog } from '@/types/agent';
-import { useAgentStore } from '@/store/agentStore';
-import { useLayoutStore } from '@/store/layoutStore';
+import UserMessage from './UserMessage';
+import AssistantMessage from './AssistantMessage';
+import type { Message } from '@/types/agent';
 
 interface MessageBubbleProps {
   message: Message;
+  /** True if this is the user message that starts the last turn. */
+  isLastTurn?: boolean;
+  /** Callback to undo (remove) the last turn. */
+  onUndoTurn?: () => void;
+  /** Whether the agent is currently processing. */
+  isProcessing?: boolean;
+  /** True when this message is actively receiving streaming chunks. */
+  isStreaming?: boolean;
 }
 
-// Render a tools segment with clickable tool calls
-function ToolsSegment({ tools }: { tools: TraceLog[] }) {
-  const { showToolOutput } = useAgentStore();
-  const { setRightPanelOpen } = useLayoutStore();
-
-  const handleToolClick = (log: TraceLog) => {
-    if (log.completed && log.output) {
-      showToolOutput(log);
-      setRightPanelOpen(true);
-    }
-  };
-
-  return (
-    <Box
-      sx={{
-        bgcolor: 'rgba(0,0,0,0.3)',
-        borderRadius: 1,
-        p: 1.5,
-        border: 1,
-        borderColor: 'rgba(255,255,255,0.05)',
-        my: 1.5,
-      }}
-    >
-      <Box sx={{ display: 'flex', flexDirection: 'column', gap: 0.5 }}>
-        {tools.map((log) => {
-          const isClickable = log.completed && log.output;
-          return (
-            <Typography
-              key={log.id}
-              variant="caption"
-              component="div"
-              onClick={() => handleToolClick(log)}
-              sx={{
-                color: 'var(--muted-text)',
-                fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
-                fontSize: '0.75rem',
-                display: 'flex',
-                alignItems: 'center',
-                gap: 0.5,
-                cursor: isClickable ? 'pointer' : 'default',
-                borderRadius: 0.5,
-                px: 0.5,
-                mx: -0.5,
-                transition: 'background-color 0.15s ease',
-                '&:hover': isClickable ? {
-                  bgcolor: 'rgba(255,255,255,0.05)',
-                } : {},
-              }}
-            >
-              <span style={{
-                color: log.completed
-                  ? (log.success === false ? '#F87171' : '#FDB022')
-                  : 'inherit',
-                fontSize: '0.85rem',
-              }}>
-                {log.completed ? (log.success === false ? '✗' : '✓') : '•'}
-              </span>
-              <span style={{
-                fontWeight: 600,
-                color: isClickable ? 'rgba(255, 255, 255, 0.9)' : 'inherit',
-                textDecoration: isClickable ? 'underline' : 'none',
-                textDecorationColor: 'rgba(255,255,255,0.3)',
-                textUnderlineOffset: '2px',
-              }}>
-                {log.tool}
-              </span>
-              {!log.completed && <span style={{ opacity: 0.6 }}>...</span>}
-              {isClickable && (
-                <span style={{
-                  opacity: 0.4,
-                  fontSize: '0.65rem',
-                  marginLeft: 'auto',
-                }}>
-                  click to view
-                </span>
-              )}
-            </Typography>
-          );
-        })}
-      </Box>
-    </Box>
-  );
-}
-
-// Markdown styles
-const markdownStyles = {
-  '& p': { m: 0, mb: 1, '&:last-child': { mb: 0 } },
-  '& pre': {
-    bgcolor: 'rgba(0,0,0,0.5)',
-    p: 1.5,
-    borderRadius: 1,
-    overflow: 'auto',
-    fontSize: '0.85rem',
-    border: '1px solid rgba(255,255,255,0.05)',
-  },
-  '& code': {
-    bgcolor: 'rgba(255,255,255,0.05)',
-    px: 0.5,
-    py: 0.25,
-    borderRadius: 0.5,
-    fontSize: '0.85rem',
-    fontFamily: '"JetBrains Mono", monospace',
-  },
-  '& pre code': { bgcolor: 'transparent', p: 0 },
-  '& a': {
-    color: 'var(--accent-yellow)',
-    textDecoration: 'none',
-    '&:hover': { textDecoration: 'underline' },
-  },
-  '& ul, & ol': { pl: 2, my: 1 },
-  '& table': {
-    borderCollapse: 'collapse',
-    width: '100%',
-    my: 2,
-    fontSize: '0.875rem',
-  },
-  '& th': {
-    borderBottom: '1px solid rgba(255,255,255,0.1)',
-    textAlign: 'left',
-    p: 1,
-    bgcolor: 'rgba(255,255,255,0.02)',
-  },
-  '& td': {
-    borderBottom: '1px solid rgba(255,255,255,0.05)',
-    p: 1,
-  },
-};
-
-export default function MessageBubble({ message }: MessageBubbleProps) {
-  const isUser = message.role === 'user';
-  const isAssistant = message.role === 'assistant';
-
-  if (message.approval) {
-    return (
-      <Box sx={{ width: '100%', maxWidth: '880px', mx: 'auto', my: 2 }}>
-        <ApprovalFlow message={message} />
-      </Box>
-    );
+/**
+ * Thin dispatcher — routes each message to the correct
+ * specialised component based on its role / content.
+ */
+export default function MessageBubble({
+  message,
+  isLastTurn = false,
+  onUndoTurn,
+  isProcessing = false,
+  isStreaming = false,
+}: MessageBubbleProps) {
+  // Legacy approval-only messages (from old localStorage data) — skip them.
+  // Approvals are now rendered inline within ToolCallGroup.
+  if (message.approval && !message.content && !message.segments?.length) {
+    return null;
   }
 
-  // Render segments chronologically if available, otherwise fall back to content
-  const renderContent = () => {
-    if (message.segments && message.segments.length > 0) {
-      return message.segments.map((segment, idx) => {
-        if (segment.type === 'text' && segment.content) {
-          return (
-            <Box key={idx} sx={markdownStyles}>
-              <ReactMarkdown remarkPlugins={[remarkGfm]}>{segment.content}</ReactMarkdown>
-            </Box>
-          );
-        }
-        if (segment.type === 'tools' && segment.tools && segment.tools.length > 0) {
-          return <ToolsSegment key={idx} tools={segment.tools} />;
-        }
-        return null;
-      });
-    }
-    // Fallback: just render content
+  if (message.role === 'user') {
     return (
-      <Box sx={markdownStyles}>
-        <ReactMarkdown remarkPlugins={[remarkGfm]}>{message.content}</ReactMarkdown>
-      </Box>
+      <UserMessage
+        message={message}
+        isLastTurn={isLastTurn}
+        onUndoTurn={onUndoTurn}
+        isProcessing={isProcessing}
+      />
     );
-  };
+  }
 
-  return (
-    <Box
-      sx={{
-        display: 'flex',
-        justifyContent: isUser ? 'flex-end' : 'flex-start',
-        width: '100%',
-        maxWidth: '880px',
-        mx: 'auto',
-      }}
-    >
-      <Paper
-        elevation={0}
-        className={`message ${isUser ? 'user' : isAssistant ? 'assistant' : ''}`}
-        sx={{
-          p: '14px 18px',
-          margin: '10px 0',
-          maxWidth: '100%',
-          borderRadius: 'var(--radius-lg)',
-          borderTopLeftRadius: isAssistant ? '6px' : undefined,
-          lineHeight: 1.45,
-          boxShadow: 'var(--shadow-1)',
-          border: '1px solid rgba(255,255,255,0.03)',
-          background: 'linear-gradient(180deg, rgba(255,255,255,0.015), transparent)',
-        }}
-      >
-        {renderContent()}
+  if (message.role === 'assistant') {
+    return <AssistantMessage message={message} isStreaming={isStreaming} />;
+  }
 
-        <Typography
-          className="meta"
-          variant="caption"
-          sx={{ display: 'block', textAlign: 'right', mt: 1, fontSize: '11px', opacity: 0.5 }}
-        >
-          {new Date(message.timestamp).toLocaleTimeString()}
-        </Typography>
-      </Paper>
-    </Box>
-  );
+  // Fallback (tool messages, etc.)
+  return null;
 }

frontend/src/components/Chat/MessageList.tsx

@@ -1,7 +1,13 @@
-import { useEffect, useRef } from 'react';
-import { Box, Typography } from '@mui/material';
-import { useSessionStore } from '@/store/sessionStore';
+import { useEffect, useRef, useMemo, useCallback } from 'react';
+import { Box, Stack, Typography, Avatar } from '@mui/material';
+import SmartToyOutlinedIcon from '@mui/icons-material/SmartToyOutlined';
 import MessageBubble from './MessageBubble';
+import ThinkingIndicator from './ThinkingIndicator';
+import MarkdownContent from './MarkdownContent';
+import { useAgentStore } from '@/store/agentStore';
+import { useSessionStore } from '@/store/sessionStore';
+import { apiFetch } from '@/utils/api';
+import { logger } from '@/utils/logger';
 import type { Message } from '@/types/agent';
 
 interface MessageListProps {
@@ -9,92 +15,183 @@ interface MessageListProps {
   isProcessing: boolean;
 }
 
-const TechnicalIndicator = () => (
-  <Box
-    component="span"
-    sx={{
-      color: 'primary.main',
-      fontFamily: 'monospace',
-      fontWeight: 'bold',
-      fontSize: '1.2rem',
-      lineHeight: 0,
-      display: 'inline-block',
-      verticalAlign: 'middle',
-      width: '1em',
-      letterSpacing: '-3px',
-      transform: 'scale(0.6) translateY(-2px)',
-      '&::after': {
-        content: '""',
-        animation: 'dots 2s steps(4, end) infinite',
-      },
-      '@keyframes dots': {
-        '0%': { content: '""' },
-        '25%': { content: '"."' },
-        '50%': { content: '".."' },
-        '75%, 100%': { content: '"..."' },
-      },
-    }}
-  />
-);
+const WELCOME_MD = `I'm ready to help you with machine learning tasks using the Hugging Face ecosystem.
+
+**Training & Fine-tuning** — SFT, DPO, GRPO, PPO with TRL · LoRA/PEFT · Submit and monitor jobs on cloud GPUs
+
+**Data** — Find and explore datasets · Process, filter, transform · Push to the Hub
+
+**Models** — Search and discover models · Get details and configs · Deploy for inference
+
+**Research** — Find papers and documentation · Explore code examples · Check APIs and best practices
+
+**Infrastructure** — Run jobs on CPU/GPU instances · Manage repos, branches, PRs · Monitor Spaces and endpoints
+
+What would you like to do?`;
+
+/** Static welcome message rendered when the conversation is empty. */
+function WelcomeMessage() {
+  return (
+    <Stack direction="row" spacing={1.5} alignItems="flex-start">
+      <Avatar
+        sx={{
+          width: 28,
+          height: 28,
+          bgcolor: 'primary.main',
+          flexShrink: 0,
+          mt: 0.5,
+        }}
+      >
+        <SmartToyOutlinedIcon sx={{ fontSize: 16, color: '#fff' }} />
+      </Avatar>
+
+      <Box sx={{ flex: 1, minWidth: 0 }}>
+        <Stack direction="row" alignItems="baseline" spacing={1} sx={{ mb: 0.5 }}>
+          <Typography
+            variant="caption"
+            sx={{
+              fontWeight: 700,
+              fontSize: '0.72rem',
+              color: 'var(--muted-text)',
+              textTransform: 'uppercase',
+              letterSpacing: '0.04em',
+            }}
+          >
+            Assistant
+          </Typography>
+        </Stack>
+        <Box
+          sx={{
+            maxWidth: { xs: '95%', md: '85%' },
+            bgcolor: 'var(--surface)',
+            borderRadius: 1.5,
+            borderTopLeftRadius: 4,
+            px: { xs: 1.5, md: 2.5 },
+            py: 1.5,
+            border: '1px solid var(--border)',
+          }}
+        >
+          <MarkdownContent content={WELCOME_MD} />
+        </Box>
+      </Box>
+    </Stack>
+  );
+}
 
 export default function MessageList({ messages, isProcessing }: MessageListProps) {
-  const bottomRef = useRef<HTMLDivElement>(null);
+  const scrollContainerRef = useRef<HTMLDivElement>(null);
+  const stickToBottom = useRef(true);
   const { activeSessionId } = useSessionStore();
+  const { removeLastTurn, currentTurnMessageId } = useAgentStore();
+
+  // ── Scroll-to-bottom helper ─────────────────────────────────────
+  const scrollToBottom = useCallback(() => {
+    const el = scrollContainerRef.current;
+    if (el) el.scrollTop = el.scrollHeight;
+  }, []);
+
+  // ── Track user scroll intent ────────────────────────────────────
+  // When user scrolls up (>80px from bottom), disable auto-scroll.
+  // When they scroll back to bottom, re-enable it.
+  useEffect(() => {
+    const el = scrollContainerRef.current;
+    if (!el) return;
+
+    const onScroll = () => {
+      const distFromBottom = el.scrollHeight - el.scrollTop - el.clientHeight;
+      stickToBottom.current = distFromBottom < 80;
+    };
+
+    el.addEventListener('scroll', onScroll, { passive: true });
+    return () => el.removeEventListener('scroll', onScroll);
+  }, []);
+
+  // ── Auto-scroll on new messages / state changes ─────────────────
+  useEffect(() => {
+    if (stickToBottom.current) scrollToBottom();
+  }, [messages, isProcessing, scrollToBottom]);
 
-  // Auto-scroll to bottom when new messages arrive
+  // ── Auto-scroll on DOM mutations (streaming content growth) ─────
+  // This catches token-by-token updates that don't change the messages
+  // array reference (appendToMessage mutates in place).
   useEffect(() => {
-    bottomRef.current?.scrollIntoView({ behavior: 'smooth' });
-  }, [messages, isProcessing]);
+    const el = scrollContainerRef.current;
+    if (!el) return;
+
+    const observer = new MutationObserver(() => {
+      if (stickToBottom.current) {
+        el.scrollTop = el.scrollHeight;
+      }
+    });
+
+    observer.observe(el, {
+      childList: true,
+      subtree: true,
+      characterData: true,
+    });
+
+    return () => observer.disconnect();
+  }, []);
+
+  // Find the index of the last user message (start of the last turn)
+  const lastUserMsgId = useMemo(() => {
+    for (let i = messages.length - 1; i >= 0; i--) {
+      if (messages[i].role === 'user') return messages[i].id;
+    }
+    return null;
+  }, [messages]);
+
+  const handleUndoLastTurn = useCallback(async () => {
+    if (!activeSessionId) return;
+    try {
+      await apiFetch(`/api/undo/${activeSessionId}`, { method: 'POST' });
+      // Optimistic removal — backend will also confirm via undo_complete WS event
+      removeLastTurn(activeSessionId);
+    } catch (e) {
+      logger.error('Undo failed:', e);
+    }
+  }, [activeSessionId, removeLastTurn]);
 
   return (
     <Box
+      ref={scrollContainerRef}
       sx={{
         flex: 1,
         overflow: 'auto',
-        p: 2,
-        display: 'flex',
-        flexDirection: 'column',
+        px: { xs: 0.5, sm: 1, md: 2 },
+        py: { xs: 2, md: 3 },
       }}
     >
-      <Box sx={{ maxWidth: 'md', mx: 'auto', width: '100%', display: 'flex', flexDirection: 'column', gap: 2 }}>
-        {messages.length === 0 && !isProcessing ? (
-          <Box
-            sx={{
-              flex: 1,
-              display: 'flex',
-              alignItems: 'center',
-              justifyContent: 'center',
-              py: 8,
-            }}
-          >
-            <Typography color="text.secondary" sx={{ fontFamily: 'monospace' }}>
-              Awaiting input…
-            </Typography>
-          </Box>
-        ) : (
-          messages.map((message) => (
-            <MessageBubble key={message.id} message={message} />
+      <Stack
+        spacing={3}
+        sx={{
+          maxWidth: 880,
+          mx: 'auto',
+          width: '100%',
+        }}
+      >
+        {/* Always show the welcome message at the top */}
+        <WelcomeMessage />
+
+        {messages.length > 0 && (
+          messages.map((msg) => (
+            <MessageBubble
+              key={msg.id}
+              message={msg}
+              isLastTurn={msg.id === lastUserMsgId}
+              onUndoTurn={handleUndoLastTurn}
+              isProcessing={isProcessing}
+              isStreaming={isProcessing && msg.id === currentTurnMessageId}
+            />
           ))
         )}
-        
-        {isProcessing && (
-          <Box sx={{ width: '100%', mb: 2 }}>
-            <Box sx={{ display: 'flex', alignItems: 'center', gap: 1, mb: 1, px: 0.5 }}>
-              <Typography variant="caption" color="text.secondary" sx={{ fontFamily: 'monospace', fontWeight: 600 }}>
-                Thinking
-              </Typography>
-              <TechnicalIndicator />
-            </Box>
-          </Box>
-        )}
 
-        {activeSessionId && (
-          // ApprovalFlow is now handled within messages
-          null
-        )}
-        
-        <div ref={bottomRef} />
-      </Box>
+        {/* Show thinking dots only when processing but no streaming message yet */}
+        {isProcessing && !currentTurnMessageId && <ThinkingIndicator />}
+
+        {/* Sentinel — keeps scroll anchor at the bottom */}
+        <div />
+      </Stack>
     </Box>
   );
-}
+}

frontend/src/components/Chat/ThinkingIndicator.tsx

@@ -0,0 +1,63 @@
+import { Box, Stack, Avatar, Typography } from '@mui/material';
+import SmartToyOutlinedIcon from '@mui/icons-material/SmartToyOutlined';
+
+/** Pulsing dots shown while the agent is processing. */
+export default function ThinkingIndicator() {
+  return (
+    <Stack direction="row" spacing={1.5} alignItems="flex-start">
+      <Avatar
+        sx={{
+          width: 28,
+          height: 28,
+          bgcolor: 'primary.main',
+          flexShrink: 0,
+          mt: 0.5,
+        }}
+      >
+        <SmartToyOutlinedIcon sx={{ fontSize: 16, color: '#fff' }} />
+      </Avatar>
+
+      <Box sx={{ pt: 0.75 }}>
+        <Typography
+          variant="caption"
+          sx={{
+            fontWeight: 700,
+            fontSize: '0.72rem',
+            color: 'var(--muted-text)',
+            textTransform: 'uppercase',
+            letterSpacing: '0.04em',
+            display: 'flex',
+            alignItems: 'center',
+            gap: 0.75,
+          }}
+        >
+          Thinking
+          <Box
+            component="span"
+            sx={{
+              display: 'inline-flex',
+              gap: '3px',
+              '& span': {
+                width: 4,
+                height: 4,
+                borderRadius: '50%',
+                bgcolor: 'primary.main',
+                animation: 'dotPulse 1.4s ease-in-out infinite',
+              },
+              '& span:nth-of-type(2)': { animationDelay: '0.2s' },
+              '& span:nth-of-type(3)': { animationDelay: '0.4s' },
+              '@keyframes dotPulse': {
+                '0%, 80%, 100%': { opacity: 0.25, transform: 'scale(0.8)' },
+                '40%': { opacity: 1, transform: 'scale(1)' },
+              },
+            }}
+          >
+            <span />
+            <span />
+            <span />
+          </Box>
+        </Typography>
+      </Box>
+    </Stack>
+  );
+}

frontend/src/components/Chat/ToolCallGroup.tsx

@@ -0,0 +1,410 @@
+import { useCallback, useState } from 'react';
+import { Box, Stack, Typography, Chip, Button, TextField, IconButton, Link } from '@mui/material';
+import CheckCircleOutlineIcon from '@mui/icons-material/CheckCircleOutline';
+import ErrorOutlineIcon from '@mui/icons-material/ErrorOutline';
+import MoreHorizIcon from '@mui/icons-material/MoreHoriz';
+import OpenInNewIcon from '@mui/icons-material/OpenInNew';
+import HourglassEmptyIcon from '@mui/icons-material/HourglassEmpty';
+import LaunchIcon from '@mui/icons-material/Launch';
+import SendIcon from '@mui/icons-material/Send';
+import { useAgentStore } from '@/store/agentStore';
+import { useLayoutStore } from '@/store/layoutStore';
+import { useSessionStore } from '@/store/sessionStore';
+import { apiFetch } from '@/utils/api';
+import { logger } from '@/utils/logger';
+import type { TraceLog, ApprovalStatus } from '@/types/agent';
+
+interface ToolCallGroupProps {
+  tools: TraceLog[];
+}
+
+// ── Status icon based on tool state ─────────────────────────────────
+function StatusIcon({ log }: { log: TraceLog }) {
+  // Awaiting approval
+  if (log.approvalStatus === 'pending') {
+    return <HourglassEmptyIcon sx={{ fontSize: 16, color: 'var(--accent-yellow)' }} />;
+  }
+  // Rejected
+  if (log.approvalStatus === 'rejected') {
+    return <ErrorOutlineIcon sx={{ fontSize: 16, color: 'error.main' }} />;
+  }
+  // Running (not completed yet)
+  if (!log.completed) {
+    return (
+      <MoreHorizIcon
+        sx={{
+          fontSize: 16,
+          color: 'var(--muted-text)',
+          animation: 'pulse 1.5s ease-in-out infinite',
+          '@keyframes pulse': {
+            '0%, 100%': { opacity: 0.4 },
+            '50%': { opacity: 1 },
+          },
+        }}
+      />
+    );
+  }
+  // Failed
+  if (log.success === false) {
+    return <ErrorOutlineIcon sx={{ fontSize: 16, color: 'error.main' }} />;
+  }
+  // Completed successfully
+  return <CheckCircleOutlineIcon sx={{ fontSize: 16, color: 'success.main' }} />;
+}
+
+// ── Status chip label ───────────────────────────────────────────────
+function statusLabel(log: TraceLog): string | null {
+  if (log.approvalStatus === 'pending') return 'awaiting approval';
+  if (log.approvalStatus === 'rejected') return 'rejected';
+  if (!log.completed) return 'running';
+  return null;
+}
+
+function statusColor(log: TraceLog): string {
+  if (log.approvalStatus === 'pending') return 'var(--accent-yellow)';
+  if (log.approvalStatus === 'rejected') return 'var(--accent-red)';
+  return 'var(--accent-yellow)';
+}
+
+// ── Inline approval UI ──────────────────────────────────────────────
+function InlineApproval({
+  log,
+  onResolve,
+}: {
+  log: TraceLog;
+  onResolve: (toolCallId: string, approved: boolean, feedback?: string) => void;
+}) {
+  const [feedback, setFeedback] = useState('');
+
+  return (
+    <Box sx={{ px: 1.5, py: 1.5, borderTop: '1px solid var(--tool-border)' }}>
+      {/* Tool description */}
+      {log.tool === 'hf_jobs' && log.args && (
+        <Typography variant="body2" sx={{ color: 'var(--muted-text)', fontSize: '0.75rem', mb: 1.5 }}>
+          Execute <Box component="span" sx={{ color: 'var(--accent-yellow)', fontWeight: 500 }}>{log.tool}</Box> on{' '}
+          <Box component="span" sx={{ fontWeight: 500, color: 'var(--text)' }}>
+            {String(log.args.hardware_flavor || 'default')}
+          </Box>
+          {log.args.timeout && (
+            <> with timeout <Box component="span" sx={{ fontWeight: 500, color: 'var(--text)' }}>
+              {String(log.args.timeout)}
+            </Box></>
+          )}
+        </Typography>
+      )}
+
+      {/* Feedback + buttons */}
+      <Box sx={{ display: 'flex', gap: 1, mb: 1 }}>
+        <TextField
+          fullWidth
+          size="small"
+          placeholder="Feedback (optional)"
+          value={feedback}
+          onChange={(e) => setFeedback(e.target.value)}
+          variant="outlined"
+          sx={{
+            '& .MuiOutlinedInput-root': {
+              bgcolor: 'rgba(0,0,0,0.15)',
+              fontFamily: 'inherit',
+              fontSize: '0.8rem',
+            },
+          }}
+        />
+        <IconButton
+          onClick={() => onResolve(log.toolCallId || '', false, feedback || 'Rejected by user')}
+          disabled={!feedback}
+          size="small"
+          sx={{
+            color: 'var(--accent-red)',
+            border: '1px solid rgba(255,255,255,0.05)',
+            borderRadius: '6px',
+            '&:hover': { bgcolor: 'rgba(224,90,79,0.1)', borderColor: 'var(--accent-red)' },
+            '&.Mui-disabled': { color: 'rgba(255,255,255,0.1)' },
+          }}
+        >
+          <SendIcon sx={{ fontSize: 14 }} />
+        </IconButton>
+      </Box>
+
+      <Box sx={{ display: 'flex', gap: 1 }}>
+        <Button
+          size="small"
+          onClick={() => onResolve(log.toolCallId || '', false, feedback || 'Rejected by user')}
+          sx={{
+            flex: 1,
+            textTransform: 'none',
+            border: '1px solid rgba(255,255,255,0.05)',
+            color: 'var(--accent-red)',
+            fontSize: '0.75rem',
+            py: 0.75,
+            borderRadius: '8px',
+            '&:hover': { bgcolor: 'rgba(224,90,79,0.05)', borderColor: 'var(--accent-red)' },
+          }}
+        >
+          Reject
+        </Button>
+        <Button
+          size="small"
+          onClick={() => onResolve(log.toolCallId || '', true)}
+          sx={{
+            flex: 1,
+            textTransform: 'none',
+            border: '1px solid rgba(255,255,255,0.05)',
+            color: 'var(--accent-green)',
+            fontSize: '0.75rem',
+            py: 0.75,
+            borderRadius: '8px',
+            '&:hover': { bgcolor: 'rgba(47,204,113,0.05)', borderColor: 'var(--accent-green)' },
+          }}
+        >
+          Approve
+        </Button>
+      </Box>
+    </Box>
+  );
+}
+
+// ── Main component ──────────────────────────────────────────────────
+export default function ToolCallGroup({ tools }: ToolCallGroupProps) {
+  const { showToolOutput, setPanelTab, setActivePanelTab, clearPanelTabs } = useAgentStore();
+  const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
+  const { activeSessionId } = useSessionStore();
+
+  const handleClick = useCallback(
+    (log: TraceLog) => {
+      // For hf_jobs with scripts, use tab system
+      if (log.tool === 'hf_jobs' && log.args?.script) {
+        clearPanelTabs();
+        setPanelTab({
+          id: 'script',
+          title: 'Script',
+          content: String(log.args.script),
+          language: 'python',
+        });
+        if (log.jobLogs) {
+          setPanelTab({
+            id: 'logs',
+            title: 'Logs',
+            content: log.jobLogs,
+            language: 'text',
+          });
+        }
+        setActivePanelTab('script');
+        setRightPanelOpen(true);
+        setLeftSidebarOpen(false);
+        return;
+      }
+
+      // Show output if completed, or args if still running
+      if (log.completed && log.output) {
+        showToolOutput(log);
+      } else if (log.args) {
+        const content = JSON.stringify(log.args, null, 2);
+        showToolOutput({ ...log, output: content });
+      } else {
+        return;
+      }
+      setRightPanelOpen(true);
+    },
+    [showToolOutput, setRightPanelOpen, setLeftSidebarOpen, clearPanelTabs, setPanelTab, setActivePanelTab],
+  );
+
+  const handleApprovalResolve = useCallback(
+    async (toolCallId: string, approved: boolean, feedback?: string) => {
+      if (!activeSessionId) return;
+      try {
+        await apiFetch('/api/approve', {
+          method: 'POST',
+          body: JSON.stringify({
+            session_id: activeSessionId,
+            approvals: [{
+              tool_call_id: toolCallId,
+              approved,
+              feedback: approved ? null : feedback || 'Rejected by user',
+            }],
+          }),
+        });
+        // The WebSocket will send back tool_output events which will update the trace
+      } catch (e) {
+        logger.error('Approval failed:', e);
+      }
+    },
+    [activeSessionId],
+  );
+
+  return (
+    <Box
+      sx={{
+        borderRadius: 2,
+        border: '1px solid var(--tool-border)',
+        bgcolor: 'var(--tool-bg)',
+        overflow: 'hidden',
+        my: 1,
+      }}
+    >
+      <Stack divider={<Box sx={{ borderBottom: '1px solid var(--tool-border)' }} />}>
+        {tools.map((log) => {
+          const clickable = (log.completed && !!log.output) || !!log.args;
+          const label = statusLabel(log);
+          const isPendingApproval = log.approvalStatus === 'pending';
+
+          return (
+            <Box key={log.id}>
+              {/* Main tool row */}
+              <Stack
+                direction="row"
+                alignItems="center"
+                spacing={1}
+                onClick={() => !isPendingApproval && handleClick(log)}
+                sx={{
+                  px: 1.5,
+                  py: 1,
+                  cursor: isPendingApproval ? 'default' : clickable ? 'pointer' : 'default',
+                  transition: 'background-color 0.15s',
+                  '&:hover': clickable && !isPendingApproval ? { bgcolor: 'var(--hover-bg)' } : {},
+                }}
+              >
+                <StatusIcon log={log} />
+
+                <Typography
+                  variant="body2"
+                  sx={{
+                    fontFamily: '"JetBrains Mono", ui-monospace, SFMono-Regular, monospace',
+                    fontWeight: 600,
+                    fontSize: '0.78rem',
+                    color: 'var(--text)',
+                    flex: 1,
+                    minWidth: 0,
+                    overflow: 'hidden',
+                    textOverflow: 'ellipsis',
+                    whiteSpace: 'nowrap',
+                  }}
+                >
+                  {log.tool}
+                </Typography>
+
+                {/* Quick action links for completed jobs */}
+                {log.completed && log.tool === 'hf_jobs' && log.args?.script && (
+                  <Box sx={{ display: 'flex', gap: 0.5 }} onClick={(e) => e.stopPropagation()}>
+                    <Typography
+                      component="span"
+                      onClick={() => handleClick(log)}
+                      sx={{
+                        fontSize: '0.68rem',
+                        color: 'var(--muted-text)',
+                        cursor: 'pointer',
+                        px: 0.75,
+                        py: 0.25,
+                        borderRadius: 0.5,
+                        '&:hover': { color: 'var(--accent-yellow)', bgcolor: 'var(--hover-bg)' },
+                      }}
+                    >
+                      Script
+                    </Typography>
+                    {log.jobLogs && (
+                      <Typography
+                        component="span"
+                        onClick={() => {
+                          clearPanelTabs();
+                          if (log.args?.script) {
+                            setPanelTab({ id: 'script', title: 'Script', content: String(log.args.script), language: 'python' });
+                          }
+                          setPanelTab({ id: 'logs', title: 'Logs', content: log.jobLogs!, language: 'text' });
+                          setActivePanelTab('logs');
+                          setRightPanelOpen(true);
+                          setLeftSidebarOpen(false);
+                        }}
+                        sx={{
+                          fontSize: '0.68rem',
+                          color: 'var(--accent-yellow)',
+                          cursor: 'pointer',
+                          px: 0.75,
+                          py: 0.25,
+                          borderRadius: 0.5,
+                          '&:hover': { bgcolor: 'var(--hover-bg)' },
+                        }}
+                      >
+                        Logs
+                      </Typography>
+                    )}
+                  </Box>
+                )}
+
+                {label && (
+                  <Chip
+                    label={label}
+                    size="small"
+                    sx={{
+                      height: 20,
+                      fontSize: '0.65rem',
+                      fontWeight: 600,
+                      bgcolor: 'var(--accent-yellow-weak)',
+                      color: statusColor(log),
+                      letterSpacing: '0.03em',
+                    }}
+                  />
+                )}
+
+                {clickable && !isPendingApproval && (
+                  <OpenInNewIcon sx={{ fontSize: 14, color: 'var(--muted-text)', opacity: 0.6 }} />
+                )}
+              </Stack>
+
+              {/* Job status + link row */}
+              {(log.jobUrl || log.jobStatus) && (
+                <Box
+                  sx={{
+                    display: 'flex',
+                    alignItems: 'center',
+                    gap: 1.5,
+                    px: 1.5,
+                    py: 0.75,
+                    borderTop: '1px solid var(--tool-border)',
+                  }}
+                >
+                  {log.jobStatus && (
+                    <Typography
+                      variant="caption"
+                      sx={{
+                        color: log.success === false ? 'var(--accent-red)' : 'var(--accent-green)',
+                        fontSize: '0.7rem',
+                        fontWeight: 600,
+                      }}
+                    >
+                      {log.jobStatus}
+                    </Typography>
+                  )}
+                  {log.jobUrl && (
+                    <Link
+                      href={log.jobUrl}
+                      target="_blank"
+                      rel="noopener noreferrer"
+                      onClick={(e) => e.stopPropagation()}
+                      sx={{
+                        display: 'inline-flex',
+                        alignItems: 'center',
+                        gap: 0.5,
+                        color: 'var(--accent-yellow)',
+                        fontSize: '0.68rem',
+                        textDecoration: 'none',
+                        '&:hover': { textDecoration: 'underline' },
+                      }}
+                    >
+                      <LaunchIcon sx={{ fontSize: 12 }} />
+                      View on HF
+                    </Link>
+                  )}
+                </Box>
+              )}
+
+              {/* Inline approval UI (only when pending) */}
+              {isPendingApproval && (
+                <InlineApproval log={log} onResolve={handleApprovalResolve} />
+              )}
+            </Box>
+          );
+        })}
+      </Stack>
+    </Box>
+  );
+}

frontend/src/components/Chat/UserMessage.tsx

@@ -0,0 +1,117 @@
+import { Box, Stack, Typography, Avatar, IconButton, Tooltip } from '@mui/material';
+import PersonOutlineIcon from '@mui/icons-material/PersonOutline';
+import CloseIcon from '@mui/icons-material/Close';
+import type { Message } from '@/types/agent';
+
+interface UserMessageProps {
+  message: Message;
+  /** True if this message starts the last turn. */
+  isLastTurn?: boolean;
+  /** Callback to remove the last turn. */
+  onUndoTurn?: () => void;
+  /** Whether the agent is currently processing (disables undo). */
+  isProcessing?: boolean;
+}
+
+export default function UserMessage({
+  message,
+  isLastTurn = false,
+  onUndoTurn,
+  isProcessing = false,
+}: UserMessageProps) {
+  const showUndo = isLastTurn && !isProcessing && !!onUndoTurn;
+
+  return (
+    <Stack
+      direction="row"
+      spacing={1.5}
+      justifyContent="flex-end"
+      alignItems="flex-start"
+      sx={{
+        // Show the undo button when hovering the entire row
+        '& .undo-btn': {
+          opacity: 0,
+          transition: 'opacity 0.15s ease',
+        },
+        '&:hover .undo-btn': {
+          opacity: 1,
+        },
+      }}
+    >
+      {/* Undo button — visible on hover, left of the bubble */}
+      {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>
+      )}
+
+      <Box
+        sx={{
+          maxWidth: { xs: '88%', md: '72%' },
+          bgcolor: 'var(--surface)',
+          borderRadius: 1.5,
+          borderTopRightRadius: 4,
+          px: { xs: 1.5, md: 2.5 },
+          py: 1.5,
+          border: '1px solid var(--border)',
+        }}
+      >
+        <Typography
+          variant="body1"
+          sx={{
+            fontSize: '0.925rem',
+            lineHeight: 1.65,
+            color: 'var(--text)',
+            whiteSpace: 'pre-wrap',
+            wordBreak: 'break-word',
+          }}
+        >
+          {message.content}
+        </Typography>
+
+        <Typography
+          variant="caption"
+          sx={{
+            display: 'block',
+            textAlign: 'right',
+            mt: 1,
+            fontSize: '0.68rem',
+            color: 'var(--muted-text)',
+            opacity: 0.7,
+          }}
+        >
+          {new Date(message.timestamp).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}
+        </Typography>
+      </Box>
+
+      <Avatar
+        sx={{
+          width: 28,
+          height: 28,
+          bgcolor: 'var(--hover-bg)',
+          border: '1px solid var(--border)',
+          flexShrink: 0,
+          mt: 0.5,
+        }}
+      >
+        <PersonOutlineIcon sx={{ fontSize: 16, color: 'var(--muted-text)' }} />
+      </Avatar>
+    </Stack>
+  );
+}

frontend/src/components/CodePanel/CodePanel.tsx

@@ -1,5 +1,5 @@
 import { useRef, useEffect, useMemo } from 'react';
-import { Box, Typography, IconButton } from '@mui/material';
+import { Box, Stack, Typography, IconButton } from '@mui/material';
 import CloseIcon from '@mui/icons-material/Close';
 import RadioButtonUncheckedIcon from '@mui/icons-material/RadioButtonUnchecked';
 import CheckCircleIcon from '@mui/icons-material/CheckCircle';
@@ -8,25 +8,104 @@ import CodeIcon from '@mui/icons-material/Code';
 import TerminalIcon from '@mui/icons-material/Terminal';
 import ArticleIcon from '@mui/icons-material/Article';
 import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
-import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism';
+import { vscDarkPlus, vs } from 'react-syntax-highlighter/dist/esm/styles/prism';
 import ReactMarkdown from 'react-markdown';
 import remarkGfm from 'remark-gfm';
 import { useAgentStore } from '@/store/agentStore';
 import { useLayoutStore } from '@/store/layoutStore';
 import { processLogs } from '@/utils/logProcessor';
 
+// ── Helpers ──────────────────────────────────────────────────────
+
+function tabIcon(id: string, language?: string) {
+  if (id === 'script' || language === 'python') return <CodeIcon sx={{ fontSize: 14 }} />;
+  if (id === 'tool_output' || language === 'markdown' || language === 'json')
+    return <ArticleIcon sx={{ fontSize: 14 }} />;
+  return <TerminalIcon sx={{ fontSize: 14 }} />;
+}
+
+function PlanStatusIcon({ status }: { status: string }) {
+  if (status === 'completed') return <CheckCircleIcon sx={{ fontSize: 16, color: 'var(--accent-green)' }} />;
+  if (status === 'in_progress') return <PlayCircleOutlineIcon sx={{ fontSize: 16, color: 'var(--accent-yellow)' }} />;
+  return <RadioButtonUncheckedIcon sx={{ fontSize: 16, color: 'var(--muted-text)', opacity: 0.5 }} />;
+}
+
+// ── Markdown styles (adapts via CSS vars) ────────────────────────
+const markdownSx = {
+  color: 'var(--text)',
+  fontSize: '13px',
+  lineHeight: 1.6,
+  '& p': { m: 0, mb: 1.5, '&:last-child': { mb: 0 } },
+  '& pre': {
+    bgcolor: 'var(--code-bg)',
+    p: 1.5,
+    borderRadius: 1,
+    overflow: 'auto',
+    fontSize: '12px',
+    border: '1px solid var(--tool-border)',
+  },
+  '& code': {
+    bgcolor: 'var(--hover-bg)',
+    px: 0.5,
+    py: 0.25,
+    borderRadius: 0.5,
+    fontSize: '12px',
+    fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
+  },
+  '& pre code': { bgcolor: 'transparent', p: 0 },
+  '& a': {
+    color: 'var(--accent-yellow)',
+    textDecoration: 'none',
+    '&:hover': { textDecoration: 'underline' },
+  },
+  '& ul, & ol': { pl: 2.5, my: 1 },
+  '& li': { mb: 0.5 },
+  '& table': {
+    borderCollapse: 'collapse',
+    width: '100%',
+    my: 2,
+    fontSize: '12px',
+    fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
+  },
+  '& th': {
+    borderBottom: '2px solid var(--border-hover)',
+    textAlign: 'left',
+    p: 1,
+    fontWeight: 600,
+  },
+  '& td': {
+    borderBottom: '1px solid var(--tool-border)',
+    p: 1,
+  },
+  '& h1, & h2, & h3, & h4': { mt: 2, mb: 1, fontWeight: 600 },
+  '& h1': { fontSize: '1.25rem' },
+  '& h2': { fontSize: '1.1rem' },
+  '& h3': { fontSize: '1rem' },
+  '& blockquote': {
+    borderLeft: '3px solid var(--accent-yellow)',
+    pl: 2,
+    ml: 0,
+    color: 'var(--muted-text)',
+  },
+} as const;
+
+// ── Component ────────────────────────────────────────────────────
+
 export default function CodePanel() {
-  const { panelContent, panelTabs, activePanelTab, setActivePanelTab, removePanelTab, plan } = useAgentStore();
-  const { setRightPanelOpen } = useLayoutStore();
+  const { panelContent, panelTabs, activePanelTab, setActivePanelTab, removePanelTab, plan } =
+    useAgentStore();
+  const { setRightPanelOpen, themeMode } = useLayoutStore();
   const scrollRef = useRef<HTMLDivElement>(null);
 
-  // Get the active tab content, or fall back to panelContent for backwards compatibility
-  const activeTab = panelTabs.find(t => t.id === activePanelTab);
+  const activeTab = panelTabs.find((t) => t.id === activePanelTab);
   const currentContent = activeTab || panelContent;
+  const hasTabs = panelTabs.length > 0;
+
+  const isDark = themeMode === 'dark';
+  const syntaxTheme = isDark ? vscDarkPlus : vs;
 
   const displayContent = useMemo(() => {
     if (!currentContent?.content) return '';
-    // Apply log processing only for text/logs, not for code/json
     if (!currentContent.language || currentContent.language === 'text') {
       return processLogs(currentContent.content);
     }
@@ -34,36 +113,80 @@ export default function CodePanel() {
   }, [currentContent?.content, currentContent?.language]);
 
   useEffect(() => {
-    // Auto-scroll only for logs tab
     if (scrollRef.current && activePanelTab === 'logs') {
       scrollRef.current.scrollTop = scrollRef.current.scrollHeight;
     }
   }, [displayContent, activePanelTab]);
 
-  const hasTabs = panelTabs.length > 0;
+  // ── Syntax-highlighted code block (DRY) ────────────────────────
+  const renderSyntaxBlock = (language: string) => (
+    <SyntaxHighlighter
+      language={language}
+      style={syntaxTheme}
+      customStyle={{
+        margin: 0,
+        padding: 0,
+        background: 'transparent',
+        fontSize: '13px',
+        fontFamily: 'inherit',
+      }}
+      wrapLines
+      wrapLongLines
+    >
+      {displayContent}
+    </SyntaxHighlighter>
+  );
+
+  // ── Content renderer ───────────────────────────────────────────
+  const renderContent = () => {
+    if (!currentContent?.content) {
+      return (
+        <Box sx={{ display: 'flex', alignItems: 'center', justifyContent: 'center', height: '100%', opacity: 0.5 }}>
+          <Typography variant="caption">NO CONTENT TO DISPLAY</Typography>
+        </Box>
+      );
+    }
+
+    if (currentContent.language === 'python') return renderSyntaxBlock('python');
+    if (currentContent.language === 'json') return renderSyntaxBlock('json');
+
+    if (currentContent.language === 'markdown') {
+      return (
+        <Box sx={markdownSx}>
+          <ReactMarkdown remarkPlugins={[remarkGfm]}>{displayContent}</ReactMarkdown>
+        </Box>
+      );
+    }
+
+    // Plain text / logs
+    return (
+      <Box
+        component="pre"
+        sx={{ m: 0, fontFamily: 'inherit', color: 'var(--text)', whiteSpace: 'pre-wrap', wordBreak: 'break-all' }}
+      >
+        <code>{displayContent}</code>
+      </Box>
+    );
+  };
 
   return (
     <Box sx={{ height: '100%', display: 'flex', flexDirection: 'column', bgcolor: 'var(--panel)' }}>
-      {/* Header - Fixed 60px to align */}
-      <Box sx={{
-        height: '60px',
-        display: 'flex',
-        alignItems: 'center',
-        justifyContent: 'space-between',
-        px: 2,
-        borderBottom: '1px solid rgba(255,255,255,0.03)'
-      }}>
+      {/* ── Header (60 px, aligned with top bar) ────────────────── */}
+      <Box
+        sx={{
+          height: 60,
+          display: 'flex',
+          alignItems: 'center',
+          justifyContent: 'space-between',
+          px: 2,
+          borderBottom: '1px solid var(--border)',
+          flexShrink: 0,
+        }}
+      >
         {hasTabs ? (
           <Box sx={{ display: 'flex', alignItems: 'center', gap: 0.5, flexWrap: 'wrap' }}>
             {panelTabs.map((tab) => {
               const isActive = activePanelTab === tab.id;
-              // Choose icon based on tab type
-              let icon = <TerminalIcon sx={{ fontSize: 14 }} />;
-              if (tab.id === 'script' || tab.language === 'python') {
-                icon = <CodeIcon sx={{ fontSize: 14 }} />;
-              } else if (tab.id === 'tool_output' || tab.language === 'markdown' || tab.language === 'json') {
-                icon = <ArticleIcon sx={{ fontSize: 14 }} />;
-              }
               return (
                 <Box
                   key={tab.id}
@@ -81,16 +204,14 @@ export default function CodePanel() {
                     textTransform: 'uppercase',
                     letterSpacing: '0.05em',
                     color: isActive ? 'var(--text)' : 'var(--muted-text)',
-                    bgcolor: isActive ? 'rgba(255,255,255,0.08)' : 'transparent',
+                    bgcolor: isActive ? 'var(--tab-active-bg)' : 'transparent',
                     border: '1px solid',
-                    borderColor: isActive ? 'rgba(255,255,255,0.1)' : 'transparent',
+                    borderColor: isActive ? 'var(--tab-active-border)' : 'transparent',
                     transition: 'all 0.15s ease',
-                    '&:hover': {
-                      bgcolor: 'rgba(255,255,255,0.05)',
-                    },
+                    '&:hover': { bgcolor: 'var(--tab-hover-bg)' },
                   }}
                 >
-                  {icon}
+                  {tabIcon(tab.id, tab.language)}
                   <span>{tab.title}</span>
                   <Box
                     component="span"
@@ -108,10 +229,7 @@ export default function CodePanel() {
                       borderRadius: '50%',
                       fontSize: '0.65rem',
                       opacity: 0.5,
-                      '&:hover': {
-                        opacity: 1,
-                        bgcolor: 'rgba(255,255,255,0.1)',
-                      },
+                      '&:hover': { opacity: 1, bgcolor: 'var(--tab-close-hover)' },
                     }}
                   >
                     ✕
@@ -121,16 +239,20 @@ export default function CodePanel() {
             })}
           </Box>
         ) : (
-          <Typography variant="caption" sx={{ fontWeight: 600, color: 'var(--muted-text)', textTransform: 'uppercase', letterSpacing: '0.05em' }}>
+          <Typography
+            variant="caption"
+            sx={{ fontWeight: 600, color: 'var(--muted-text)', textTransform: 'uppercase', letterSpacing: '0.05em' }}
+          >
             {currentContent?.title || 'Code Panel'}
           </Typography>
         )}
+
         <IconButton size="small" onClick={() => setRightPanelOpen(false)} sx={{ color: 'var(--muted-text)' }}>
           <CloseIcon fontSize="small" />
         </IconButton>
       </Box>
 
-      {/* Main Content Area */}
+      {/* ── Main content area ─────────────────────────────────── */}
       <Box sx={{ flex: 1, overflow: 'hidden', display: 'flex', flexDirection: 'column' }}>
         {!currentContent ? (
           <Box sx={{ flex: 1, display: 'flex', alignItems: 'center', justifyContent: 'center', p: 4 }}>
@@ -144,174 +266,72 @@ export default function CodePanel() {
               ref={scrollRef}
               className="code-panel"
               sx={{
-                background: '#0A0B0C',
+                bgcolor: 'var(--code-panel-bg)',
                 borderRadius: 'var(--radius-md)',
-                padding: '18px',
-                border: '1px solid rgba(255,255,255,0.03)',
-                fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, "Roboto Mono", monospace',
+                p: '18px',
+                border: '1px solid var(--border)',
+                fontFamily: '"JetBrains Mono", ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
                 fontSize: '13px',
                 lineHeight: 1.55,
                 height: '100%',
                 overflow: 'auto',
               }}
             >
-              {currentContent.content ? (
-                currentContent.language === 'python' ? (
-                  <SyntaxHighlighter
-                    language="python"
-                    style={vscDarkPlus}
-                    customStyle={{
-                      margin: 0,
-                      padding: 0,
-                      background: 'transparent',
-                      fontSize: '13px',
-                      fontFamily: 'inherit',
-                    }}
-                    wrapLines={true}
-                    wrapLongLines={true}
-                  >
-                    {displayContent}
-                  </SyntaxHighlighter>
-                ) : currentContent.language === 'json' ? (
-                  <SyntaxHighlighter
-                    language="json"
-                    style={vscDarkPlus}
-                    customStyle={{
-                      margin: 0,
-                      padding: 0,
-                      background: 'transparent',
-                      fontSize: '13px',
-                      fontFamily: 'inherit',
-                    }}
-                    wrapLines={true}
-                    wrapLongLines={true}
-                  >
-                    {displayContent}
-                  </SyntaxHighlighter>
-                ) : currentContent.language === 'markdown' ? (
-                  <Box sx={{
-                    color: 'var(--text)',
-                    fontSize: '13px',
-                    lineHeight: 1.6,
-                    '& p': { m: 0, mb: 1.5, '&:last-child': { mb: 0 } },
-                    '& pre': {
-                      bgcolor: 'rgba(0,0,0,0.4)',
-                      p: 1.5,
-                      borderRadius: 1,
-                      overflow: 'auto',
-                      fontSize: '12px',
-                      border: '1px solid rgba(255,255,255,0.05)',
-                    },
-                    '& code': {
-                      bgcolor: 'rgba(255,255,255,0.05)',
-                      px: 0.5,
-                      py: 0.25,
-                      borderRadius: 0.5,
-                      fontSize: '12px',
-                      fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
-                    },
-                    '& pre code': { bgcolor: 'transparent', p: 0 },
-                    '& a': {
-                      color: 'var(--accent-yellow)',
-                      textDecoration: 'none',
-                      '&:hover': { textDecoration: 'underline' },
-                    },
-                    '& ul, & ol': { pl: 2.5, my: 1 },
-                    '& li': { mb: 0.5 },
-                    '& table': {
-                      borderCollapse: 'collapse',
-                      width: '100%',
-                      my: 2,
-                      fontSize: '12px',
-                      fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
-                    },
-                    '& th': {
-                      borderBottom: '2px solid rgba(255,255,255,0.15)',
-                      textAlign: 'left',
-                      p: 1,
-                      fontWeight: 600,
-                    },
-                    '& td': {
-                      borderBottom: '1px solid rgba(255,255,255,0.05)',
-                      p: 1,
-                    },
-                    '& h1, & h2, & h3, & h4': {
-                      mt: 2,
-                      mb: 1,
-                      fontWeight: 600,
-                    },
-                    '& h1': { fontSize: '1.25rem' },
-                    '& h2': { fontSize: '1.1rem' },
-                    '& h3': { fontSize: '1rem' },
-                    '& blockquote': {
-                      borderLeft: '3px solid rgba(255,255,255,0.2)',
-                      pl: 2,
-                      ml: 0,
-                      color: 'var(--muted-text)',
-                    },
-                  }}>
-                    <ReactMarkdown remarkPlugins={[remarkGfm]}>{displayContent}</ReactMarkdown>
-                  </Box>
-                ) : (
-                  <Box component="pre" sx={{
-                    m: 0,
-                    fontFamily: 'inherit',
-                    color: 'var(--text)',
-                    whiteSpace: 'pre-wrap',
-                    wordBreak: 'break-all'
-                  }}>
-                    <code>{displayContent}</code>
-                  </Box>
-                )
-              ) : (
-                <Box sx={{ display: 'flex', alignItems: 'center', justifyContent: 'center', height: '100%', opacity: 0.5 }}>
-                  <Typography variant="caption">
-                    NO CONTENT TO DISPLAY
-                  </Typography>
-                </Box>
-              )}
+              {renderContent()}
             </Box>
           </Box>
         )}
       </Box>
 
-      {/* Plan Display at Bottom */}
+      {/* ── Plan display (bottom) ─────────────────────────────── */}
       {plan && plan.length > 0 && (
-        <Box sx={{ 
-            borderTop: '1px solid rgba(255,255,255,0.03)',
-            bgcolor: 'rgba(0,0,0,0.2)',
+        <Box
+          sx={{
+            borderTop: '1px solid var(--border)',
+            bgcolor: 'var(--plan-bg)',
             maxHeight: '30%',
             display: 'flex',
-            flexDirection: 'column'
-        }}>
-            <Box sx={{ p: 1.5, borderBottom: '1px solid rgba(255,255,255,0.03)', display: 'flex', alignItems: 'center', gap: 1 }}>
-                <Typography variant="caption" sx={{ fontWeight: 600, color: 'var(--muted-text)', textTransform: 'uppercase', letterSpacing: '0.05em' }}>
-                    CURRENT PLAN
+            flexDirection: 'column',
+          }}
+        >
+          <Box
+            sx={{
+              p: 1.5,
+              borderBottom: '1px solid var(--border)',
+              display: 'flex',
+              alignItems: 'center',
+              gap: 1,
+            }}
+          >
+            <Typography
+              variant="caption"
+              sx={{ fontWeight: 600, color: 'var(--muted-text)', textTransform: 'uppercase', letterSpacing: '0.05em' }}
+            >
+              CURRENT PLAN
+            </Typography>
+          </Box>
+
+          <Stack spacing={1} sx={{ p: 2, overflow: 'auto' }}>
+            {plan.map((item) => (
+              <Stack key={item.id} direction="row" alignItems="flex-start" spacing={1.5}>
+                <Box sx={{ mt: 0.2 }}>
+                  <PlanStatusIcon status={item.status} />
+                </Box>
+                <Typography
+                  variant="body2"
+                  sx={{
+                    fontSize: '13px',
+                    fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
+                    color: item.status === 'completed' ? 'var(--muted-text)' : 'var(--text)',
+                    textDecoration: item.status === 'completed' ? 'line-through' : 'none',
+                    opacity: item.status === 'pending' ? 0.7 : 1,
+                  }}
+                >
+                  {item.content}
                 </Typography>
-            </Box>
-            <Box sx={{ p: 2, overflow: 'auto', display: 'flex', flexDirection: 'column', gap: 1 }}>
-                {plan.map((item) => (
-                    <Box key={item.id} sx={{ display: 'flex', alignItems: 'flex-start', gap: 1.5 }}>
-                        <Box sx={{ mt: 0.2 }}>
-                            {item.status === 'completed' && <CheckCircleIcon sx={{ fontSize: 16, color: 'var(--accent-green)' }} />}
-                            {item.status === 'in_progress' && <PlayCircleOutlineIcon sx={{ fontSize: 16, color: 'var(--accent-yellow)' }} />}
-                            {item.status === 'pending' && <RadioButtonUncheckedIcon sx={{ fontSize: 16, color: 'var(--muted-text)', opacity: 0.5 }} />}
-                        </Box>
-                        <Typography 
-                            variant="body2" 
-                            sx={{ 
-                                fontSize: '13px', 
-                                fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
-                                color: item.status === 'completed' ? 'var(--muted-text)' : 'var(--text)',
-                                textDecoration: item.status === 'completed' ? 'line-through' : 'none',
-                                opacity: item.status === 'pending' ? 0.7 : 1
-                            }}
-                        >
-                            {item.content}
-                        </Typography>
-                    </Box>
-                ))}
-            </Box>
+              </Stack>
+            ))}
+          </Stack>
         </Box>
       )}
     </Box>

frontend/src/components/Layout/AppLayout.tsx

@@ -4,10 +4,17 @@ import {
   Drawer,
   Typography,
   IconButton,
+  Alert,
+  AlertTitle,
+  useMediaQuery,
+  useTheme,
 } from '@mui/material';
 import MenuIcon from '@mui/icons-material/Menu';
 import ChevronLeftIcon from '@mui/icons-material/ChevronLeft';
 import DragIndicatorIcon from '@mui/icons-material/DragIndicator';
+import DarkModeOutlinedIcon from '@mui/icons-material/DarkModeOutlined';
+import LightModeOutlinedIcon from '@mui/icons-material/LightModeOutlined';
+import { logger } from '@/utils/logger';
 
 import { useSessionStore } from '@/store/sessionStore';
 import { useAgentStore } from '@/store/agentStore';
@@ -17,49 +24,56 @@ import SessionSidebar from '@/components/SessionSidebar/SessionSidebar';
 import CodePanel from '@/components/CodePanel/CodePanel';
 import ChatInput from '@/components/Chat/ChatInput';
 import MessageList from '@/components/Chat/MessageList';
+import WelcomeScreen from '@/components/WelcomeScreen/WelcomeScreen';
+import { apiFetch } from '@/utils/api';
 import type { Message } from '@/types/agent';
 
 const DRAWER_WIDTH = 260;
 
 export default function AppLayout() {
-  const { activeSessionId } = useSessionStore();
-  const { isConnected, isProcessing, getMessages, addMessage } = useAgentStore();
+  const { sessions, activeSessionId, deleteSession, updateSessionTitle } = useSessionStore();
+  const { isConnected, isProcessing, getMessages, addMessage, setProcessing, llmHealthError, setLlmHealthError } = useAgentStore();
   const { 
     isLeftSidebarOpen, 
     isRightPanelOpen, 
     rightPanelWidth,
+    themeMode,
     setRightPanelWidth,
+    setLeftSidebarOpen,
     toggleLeftSidebar, 
-    toggleRightPanel 
+    toggleTheme,
   } = useLayoutStore();
 
-  const isResizing = useRef(false);
-
-  const startResizing = useCallback((e: React.MouseEvent) => {
-    e.preventDefault();
-    isResizing.current = true;
-    document.addEventListener('mousemove', handleMouseMove);
-    document.addEventListener('mouseup', stopResizing);
-    document.body.style.cursor = 'col-resize';
-  }, []);
+  const theme = useTheme();
+  const isMobile = useMediaQuery(theme.breakpoints.down('md'));
 
-  const stopResizing = useCallback(() => {
-    isResizing.current = false;
-    document.removeEventListener('mousemove', handleMouseMove);
-    document.removeEventListener('mouseup', stopResizing);
-    document.body.style.cursor = 'default';
-  }, []);
+  const isResizing = useRef(false);
 
   const handleMouseMove = useCallback((e: MouseEvent) => {
     if (!isResizing.current) return;
     const newWidth = window.innerWidth - e.clientX;
-    const maxWidth = window.innerWidth * 0.8;
+    const maxWidth = window.innerWidth * 0.6;
     const minWidth = 300;
     if (newWidth > minWidth && newWidth < maxWidth) {
       setRightPanelWidth(newWidth);
     }
   }, [setRightPanelWidth]);
 
+  const stopResizing = useCallback(() => {
+    isResizing.current = false;
+    document.removeEventListener('mousemove', handleMouseMove);
+    document.removeEventListener('mouseup', stopResizing);
+    document.body.style.cursor = 'default';
+  }, [handleMouseMove]);
+
+  const startResizing = useCallback((e: React.MouseEvent) => {
+    e.preventDefault();
+    isResizing.current = true;
+    document.addEventListener('mousemove', handleMouseMove);
+    document.addEventListener('mouseup', stopResizing);
+    document.body.style.cursor = 'col-resize';
+  }, [handleMouseMove, stopResizing]);
+
   useEffect(() => {
     return () => {
       document.removeEventListener('mousemove', handleMouseMove);
@@ -67,18 +81,49 @@ export default function AppLayout() {
     };
   }, [handleMouseMove, stopResizing]);
 
+  // ── LLM health check on mount ───────────────────────────────────
+  useEffect(() => {
+    let cancelled = false;
+    (async () => {
+      try {
+        const res = await apiFetch('/api/health/llm');
+        const data = await res.json();
+        if (!cancelled && data.status === 'error') {
+          setLlmHealthError({
+            error: data.error || 'Unknown LLM error',
+            errorType: data.error_type || 'unknown',
+            model: data.model,
+          });
+        } else if (!cancelled) {
+          setLlmHealthError(null);
+        }
+      } catch {
+        // Backend unreachable — not an LLM issue, ignore
+      }
+    })();
+    return () => { cancelled = true; };
+  }, []); // eslint-disable-line react-hooks/exhaustive-deps
+
   const messages = activeSessionId ? getMessages(activeSessionId) : [];
+  const hasAnySessions = sessions.length > 0;
 
   useAgentWebSocket({
     sessionId: activeSessionId,
-    onReady: () => console.log('Agent ready'),
-    onError: (error) => console.error('Agent error:', error),
+    onReady: () => logger.log('Agent ready'),
+    onError: (error) => logger.error('Agent error:', error),
+    onSessionDead: (deadSessionId) => {
+      logger.log('Removing dead session:', deadSessionId);
+      deleteSession(deadSessionId);
+    },
   });
 
   const handleSendMessage = useCallback(
     async (text: string) => {
-      if (!activeSessionId || !text.trim()) return;
+      if (!activeSessionId || !text.trim() || isProcessing) return;
       
+      // Lock input immediately to prevent double-sends
+      setProcessing(true);
+
       const userMsg: Message = {
         id: `user_${Date.now()}`,
         role: 'user',
@@ -87,55 +132,127 @@ export default function AppLayout() {
       };
       addMessage(activeSessionId, userMsg);
 
+      // Auto-title the session from the first user message (async, non-blocking)
+      const currentMessages = getMessages(activeSessionId);
+      const isFirstMessage = currentMessages.filter((m) => m.role === 'user').length <= 1;
+      if (isFirstMessage) {
+        const sessionId = activeSessionId;
+        apiFetch('/api/title', {
+          method: 'POST',
+          body: JSON.stringify({ session_id: sessionId, text: text.trim() }),
+        })
+          .then((res) => res.json())
+          .then((data) => {
+            if (data.title) updateSessionTitle(sessionId, data.title);
+          })
+          .catch(() => {
+            const raw = text.trim();
+            updateSessionTitle(sessionId, raw.length > 40 ? raw.slice(0, 40) + '…' : raw);
+          });
+      }
+
       try {
-        await fetch('/api/submit', {
+        await apiFetch('/api/submit', {
           method: 'POST',
-          headers: { 'Content-Type': 'application/json' },
           body: JSON.stringify({
             session_id: activeSessionId,
             text: text.trim(),
           }),
         });
       } catch (e) {
-        console.error('Send failed:', e);
+        logger.error('Send failed:', e);
       }
     },
-    [activeSessionId, addMessage]
+    [activeSessionId, addMessage, getMessages, updateSessionTitle, isProcessing, setProcessing]
+  );
+
+  // Close sidebar on mobile after selecting a session
+  const handleSidebarClose = useCallback(() => {
+    if (isMobile) setLeftSidebarOpen(false);
+  }, [isMobile, setLeftSidebarOpen]);
+
+  // ── LLM error banner (shared) ─────────────────────────────────────
+  const llmBanner = llmHealthError && (
+    <Alert
+      severity="error"
+      variant="filled"
+      onClose={() => setLlmHealthError(null)}
+      sx={{ borderRadius: 0, flexShrink: 0, '& .MuiAlert-message': { flex: 1 } }}
+    >
+      <AlertTitle sx={{ fontWeight: 700, fontSize: '0.85rem' }}>
+        {llmHealthError.errorType === 'credits'
+          ? 'API Credits Exhausted'
+          : llmHealthError.errorType === 'auth'
+          ? 'Invalid API Key'
+          : llmHealthError.errorType === 'rate_limit'
+          ? 'Rate Limited'
+          : llmHealthError.errorType === 'network'
+          ? 'LLM Provider Unreachable'
+          : 'LLM Error'}
+      </AlertTitle>
+      <Typography variant="body2" sx={{ fontSize: '0.8rem', opacity: 0.9 }}>
+        Model: <strong>{llmHealthError.model}</strong> — {llmHealthError.error.slice(0, 200)}
+      </Typography>
+    </Alert>
+  );
+
+  // ── Welcome screen: no sessions at all ────────────────────────────
+  if (!hasAnySessions) {
+    return (
+      <Box sx={{ width: '100%', height: '100%', display: 'flex', flexDirection: 'column' }}>
+        {llmBanner}
+        <WelcomeScreen />
+      </Box>
+    );
+  }
+
+  // ── Sidebar drawer ────────────────────────────────────────────────
+  const sidebarDrawer = (
+    <Drawer
+      variant={isMobile ? 'temporary' : 'persistent'}
+      anchor="left"
+      open={isLeftSidebarOpen}
+      onClose={() => setLeftSidebarOpen(false)}
+      ModalProps={{ keepMounted: true }} // Better mobile perf
+      sx={{
+        '& .MuiDrawer-paper': {
+          boxSizing: 'border-box',
+          width: DRAWER_WIDTH,
+          borderRight: '1px solid',
+          borderColor: 'divider',
+          top: 0,
+          height: '100%',
+          bgcolor: 'var(--panel)',
+        },
+      }}
+    >
+      <SessionSidebar onClose={handleSidebarClose} />
+    </Drawer>
   );
 
+  // ── Main chat interface ───────────────────────────────────────────
   return (
     <Box sx={{ display: 'flex', width: '100%', height: '100%' }}>
-      {/* Left Sidebar Drawer */}
-      <Box
-        component="nav"
-        sx={{
-          width: { md: isLeftSidebarOpen ? DRAWER_WIDTH : 0 },
-          flexShrink: { md: 0 },
-          transition: isResizing.current ? 'none' : 'width 0.2s',
-          overflow: 'hidden',
-        }}
-      >
-        <Drawer
-          variant="persistent"
+      {/* ── Left Sidebar ─────────────────────────────────────────── */}
+      {isMobile ? (
+        // Mobile: temporary overlay drawer (no reserved width)
+        sidebarDrawer
+      ) : (
+        // Desktop: persistent drawer with reserved width
+        <Box
+          component="nav"
           sx={{
-            display: { xs: 'none', md: 'block' },
-            '& .MuiDrawer-paper': {
-              boxSizing: 'border-box',
-              width: DRAWER_WIDTH,
-              borderRight: '1px solid',
-              borderColor: 'divider',
-              top: 0,
-              height: '100%',
-              bgcolor: 'var(--panel)', // Ensure correct background matches sidebar
-            },
+            width: isLeftSidebarOpen ? DRAWER_WIDTH : 0,
+            flexShrink: 0,
+            transition: isResizing.current ? 'none' : 'width 0.2s',
+            overflow: 'hidden',
           }}
-          open={isLeftSidebarOpen}
         >
-          <SessionSidebar />
-        </Drawer>
-      </Box>
+          {sidebarDrawer}
+        </Box>
+      )}
 
-      {/* Main Content Area */}
+      {/* ── Main Content (header + chat + code panel) ────────────── */}
       <Box
         sx={{
           flexGrow: 1,
@@ -143,142 +260,188 @@ export default function AppLayout() {
           display: 'flex',
           flexDirection: 'column',
           transition: isResizing.current ? 'none' : 'width 0.2s',
-          position: 'relative',
           overflow: 'hidden',
+          minWidth: 0,
         }}
       >
-        {/* Top Header Bar (Fixed) */}
+        {/* ── Top Header Bar ─────────────────────────────────────── */}
         <Box sx={{ 
-          height: '60px',
-          px: 1, 
+          height: { xs: 52, md: 60 },
+          px: { xs: 1, md: 2 }, 
           display: 'flex', 
           alignItems: 'center', 
           borderBottom: 1, 
           borderColor: 'divider',
           bgcolor: 'background.default',
           zIndex: 1200,
+          flexShrink: 0,
         }}>
           <IconButton onClick={toggleLeftSidebar} size="small">
-            {isLeftSidebarOpen ? <ChevronLeftIcon /> : <MenuIcon />}
+            {isLeftSidebarOpen && !isMobile ? <ChevronLeftIcon /> : <MenuIcon />}
           </IconButton>
           
-          <Box sx={{ flex: 1, display: 'flex', justifyContent: 'center' }}>
-            <img 
-              src="/hf-logo-white.png" 
-              alt="Hugging Face" 
-              style={{ height: '40px', objectFit: 'contain' }} 
+          <Box sx={{ flex: 1, display: 'flex', justifyContent: 'center', alignItems: 'center', gap: 0.75 }}>
+            <Box
+              component="img"
+              src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg"
+              alt="HF"
+              sx={{ width: { xs: 20, md: 22 }, height: { xs: 20, md: 22 } }}
             />
+            <Typography
+              variant="subtitle1"
+              sx={{
+                fontWeight: 700,
+                color: 'var(--text)',
+                letterSpacing: '-0.01em',
+                fontSize: { xs: '0.88rem', md: '0.95rem' },
+              }}
+            >
+              HF Agent
+            </Typography>
           </Box>
 
-          <IconButton 
-            onClick={toggleRightPanel} 
-            size="small" 
-            sx={{ visibility: isRightPanelOpen ? 'hidden' : 'visible' }}
+          <IconButton
+            onClick={toggleTheme}
+            size="small"
+            sx={{
+              color: 'text.secondary',
+              '&:hover': { color: 'primary.main' },
+            }}
           >
-            <MenuIcon />
+            {themeMode === 'dark' ? <LightModeOutlinedIcon fontSize="small" /> : <DarkModeOutlinedIcon fontSize="small" />}
           </IconButton>
         </Box>
 
+        {/* ── LLM Health Error Banner ────────────────────────────── */}
+        {llmBanner}
+
+        {/* ── Chat + Code Panel ──────────────────────────────────── */}
         <Box
-          component="main"
-          className="chat-pane"
           sx={{
             flexGrow: 1,
             display: 'flex',
-            flexDirection: 'column',
             overflow: 'hidden',
-            background: 'linear-gradient(180deg, var(--bg), var(--panel))',
-            padding: '24px',
           }}
         >
-          {activeSessionId ? (
+          {/* Chat area */}
+          <Box
+            component="main"
+            className="chat-pane"
+            sx={{
+              flexGrow: 1,
+              display: 'flex',
+              flexDirection: 'column',
+              overflow: 'hidden',
+              background: 'var(--body-gradient)',
+              p: { xs: 1.5, sm: 2, md: 3 },
+              minWidth: 0,
+            }}
+          >
+            {activeSessionId ? (
+              <>
+                <MessageList messages={messages} isProcessing={isProcessing} />
+                {!isConnected && messages.length > 0 && (
+                  <Box sx={{
+                    display: 'flex',
+                    alignItems: 'center',
+                    justifyContent: 'center',
+                    gap: 1,
+                    py: 1,
+                    px: { xs: 1, md: 2 },
+                    mb: 1,
+                    borderRadius: 'var(--radius-md)',
+                    bgcolor: 'rgba(255, 171, 0, 0.08)',
+                    border: '1px solid rgba(255, 171, 0, 0.2)',
+                  }}>
+                    <Typography variant="body2" sx={{ color: 'var(--accent-yellow)', fontFamily: 'monospace', fontSize: { xs: '0.7rem', md: '0.8rem' } }}>
+                      Session expired — create a new session to continue.
+                    </Typography>
+                  </Box>
+                )}
+                <ChatInput
+                  onSend={handleSendMessage}
+                  disabled={isProcessing || !isConnected}
+                />
+              </>
+            ) : (
+              <Box
+                sx={{
+                  flex: 1,
+                  display: 'flex',
+                  alignItems: 'center',
+                  justifyContent: 'center',
+                  flexDirection: 'column',
+                  gap: 2,
+                  px: 2,
+                }}
+              >
+                <Typography variant="h5" color="text.secondary" sx={{ fontFamily: 'monospace', fontSize: { xs: '1rem', md: '1.5rem' } }}>
+                  NO SESSION SELECTED
+                </Typography>
+                <Typography variant="body2" color="text.secondary" sx={{ fontFamily: 'monospace', fontSize: { xs: '0.75rem', md: '0.875rem' } }}>
+                  Initialize a session via the sidebar
+                </Typography>
+              </Box>
+            )}
+          </Box>
+
+          {/* Code panel — inline on desktop, overlay drawer on mobile */}
+          {isRightPanelOpen && !isMobile && (
             <>
-              <MessageList messages={messages} isProcessing={isProcessing} />
-              <ChatInput
-                onSend={handleSendMessage}
-                disabled={isProcessing || !isConnected}
-              />
+              <Box
+                onMouseDown={startResizing}
+                sx={{
+                  width: '4px',
+                  cursor: 'col-resize',
+                  bgcolor: 'divider',
+                  display: 'flex',
+                  alignItems: 'center',
+                  justifyContent: 'center',
+                  transition: 'background-color 0.2s',
+                  flexShrink: 0,
+                  '&:hover': { bgcolor: 'primary.main' },
+                }}
+              >
+                <DragIndicatorIcon 
+                  sx={{ fontSize: '0.8rem', color: 'text.secondary', pointerEvents: 'none' }} 
+                />
+              </Box>
+              <Box
+                sx={{
+                  width: rightPanelWidth,
+                  flexShrink: 0,
+                  height: '100%',
+                  overflow: 'hidden',
+                  borderLeft: '1px solid',
+                  borderColor: 'divider',
+                  bgcolor: 'var(--panel)',
+                }}
+              >
+                <CodePanel />
+              </Box>
             </>
-          ) : (
-            <Box
-              sx={{
-                flex: 1,
-                display: 'flex',
-                alignItems: 'center',
-                justifyContent: 'center',
-                flexDirection: 'column',
-                gap: 2,
-              }}
-            >
-              <Typography variant="h5" color="text.secondary" sx={{ fontFamily: 'monospace' }}>
-                NO SESSION SELECTED
-              </Typography>
-              <Typography variant="body2" color="text.secondary" sx={{ fontFamily: 'monospace' }}>
-                Initialize a session via the sidebar
-              </Typography>
-            </Box>
           )}
         </Box>
       </Box>
 
-      {/* Resize Handle */}
-      {isRightPanelOpen && (
-        <Box
-          onMouseDown={startResizing}
-          sx={{
-            width: '4px',
-            cursor: 'col-resize',
-            bgcolor: 'divider',
-            display: 'flex',
-            alignItems: 'center',
-            justifyContent: 'center',
-            transition: 'background-color 0.2s',
-            zIndex: 1300,
-            overflow: 'hidden',
-            '&:hover': {
-              bgcolor: 'primary.main',
-            },
-          }}
-        >
-          <DragIndicatorIcon 
-            sx={{ 
-              fontSize: '0.8rem', 
-              color: 'text.secondary',
-              pointerEvents: 'none',
-            }} 
-          />
-        </Box>
-      )}
-
-      {/* Right Panel Drawer */}
-      <Box
-        component="nav"
-        sx={{
-          width: { md: isRightPanelOpen ? rightPanelWidth : 0 },
-          flexShrink: { md: 0 },
-          transition: isResizing.current ? 'none' : 'width 0.2s',
-          overflow: 'hidden',
-        }}
-      >
+      {/* Code panel — drawer overlay on mobile */}
+      {isMobile && (
         <Drawer
-          anchor="right"
-          variant="persistent"
+          anchor="bottom"
+          open={isRightPanelOpen}
+          onClose={() => useLayoutStore.getState().setRightPanelOpen(false)}
           sx={{
-            display: { xs: 'none', md: 'block' },
             '& .MuiDrawer-paper': {
-              boxSizing: 'border-box',
-              width: rightPanelWidth,
-              borderLeft: 'none',
-              top: 0,
-              height: '100%',
+              height: '75vh',
+              borderTopLeftRadius: 16,
+              borderTopRightRadius: 16,
               bgcolor: 'var(--panel)',
             },
           }}
-          open={isRightPanelOpen}
         >
           <CodePanel />
         </Drawer>
-      </Box>
+      )}
     </Box>
   );
 }

frontend/src/components/SessionSidebar/SessionSidebar.tsx

@@ -1,246 +1,344 @@
-import { useCallback } from 'react';
+import { useCallback, useState } from 'react';
 import {
+  Alert,
   Box,
-  List,
-  ListItem,
   IconButton,
   Typography,
-  Button,
-  Tooltip,
+  CircularProgress,
+  Divider,
 } from '@mui/material';
-import DeleteIcon from '@mui/icons-material/Delete';
-import UndoIcon from '@mui/icons-material/Undo';
+import AddIcon from '@mui/icons-material/Add';
+import DeleteOutlineIcon from '@mui/icons-material/DeleteOutline';
+import ChatBubbleOutlineIcon from '@mui/icons-material/ChatBubbleOutline';
 import { useSessionStore } from '@/store/sessionStore';
 import { useAgentStore } from '@/store/agentStore';
+import { apiFetch } from '@/utils/api';
 
 interface SessionSidebarProps {
   onClose?: () => void;
 }
 
-const StatusDiode = ({ connected }: { connected: boolean }) => (
+/** Small coloured dot for connection status */
+const StatusDot = ({ connected }: { connected: boolean }) => (
   <Box
     sx={{
-      width: 10,
-      height: 10,
+      width: 6,
+      height: 6,
       borderRadius: '50%',
-      bgcolor: connected ? 'var(--accent-green)' : 'var(--accent-red)', // Use green/red for connection status
-      boxShadow: connected ? '0 0 6px rgba(47, 204, 113, 0.4)' : 'none',
-      transition: 'all 0.3s ease',
+      bgcolor: connected ? 'var(--accent-green)' : 'var(--accent-red)',
+      boxShadow: connected ? '0 0 4px rgba(76,175,80,0.4)' : 'none',
+      flexShrink: 0,
     }}
   />
 );
 
-const RunningIndicator = () => (
-    <Box
-      className="running-indicator"
-      sx={{
-        width: 10,
-        height: 10,
-        borderRadius: '50%',
-        bgcolor: 'var(--accent-yellow)',
-        boxShadow: '0 0 6px rgba(199,165,0,0.18)',
-      }}
-    />
-);
-
 export default function SessionSidebar({ onClose }: SessionSidebarProps) {
   const { sessions, activeSessionId, createSession, deleteSession, switchSession } =
     useSessionStore();
-  const { clearMessages, isConnected, isProcessing, setPlan, setPanelContent } = useAgentStore();
+  const { isConnected, setPlan, setPanelContent } =
+    useAgentStore();
+  const [isCreatingSession, setIsCreatingSession] = useState(false);
+  const [capacityError, setCapacityError] = useState<string | null>(null);
+
+  // ── Handlers ──────────────────────────────────────────────────────
 
   const handleNewSession = useCallback(async () => {
+    if (isCreatingSession) return;
+    setIsCreatingSession(true);
+    setCapacityError(null);
     try {
-      const response = await fetch('/api/session', { method: 'POST' });
+      const response = await apiFetch('/api/session', { method: 'POST' });
+      if (response.status === 503) {
+        const data = await response.json();
+        setCapacityError(data.detail || 'Server is at capacity.');
+        return;
+      }
       const data = await response.json();
       createSession(data.session_id);
-      // Clear plan and code panel for new session
       setPlan([]);
       setPanelContent(null);
       onClose?.();
-    } catch (e) {
-      console.error('Failed to create session:', e);
+    } catch {
+      setCapacityError('Failed to create session.');
+    } finally {
+      setIsCreatingSession(false);
     }
-  }, [createSession, setPlan, setPanelContent, onClose]);
+  }, [isCreatingSession, createSession, setPlan, setPanelContent, onClose]);
 
-  const handleDeleteSession = useCallback(
+  const handleDelete = useCallback(
     async (sessionId: string, e: React.MouseEvent) => {
       e.stopPropagation();
       try {
-        await fetch(`/api/session/${sessionId}`, { method: 'DELETE' });
+        await apiFetch(`/api/session/${sessionId}`, { method: 'DELETE' });
+        deleteSession(sessionId);
+      } catch {
+        // Delete locally even if backend fails (session may already be gone)
         deleteSession(sessionId);
-        clearMessages(sessionId);
-      } catch (e) {
-        console.error('Failed to delete session:', e);
       }
     },
-    [deleteSession, clearMessages]
+    [deleteSession],
   );
 
-  const handleSelectSession = useCallback(
+  const handleSelect = useCallback(
     (sessionId: string) => {
       switchSession(sessionId);
-      // Clear plan and code panel when switching sessions
       setPlan([]);
       setPanelContent(null);
       onClose?.();
     },
-    [switchSession, setPlan, setPanelContent, onClose]
+    [switchSession, setPlan, setPanelContent, onClose],
   );
 
-  const handleUndo = useCallback(async () => {
-    if (!activeSessionId) return;
-    try {
-      await fetch(`/api/undo/${activeSessionId}`, { method: 'POST' });
-    } catch (e) {
-      console.error('Undo failed:', e);
-    }
-  }, [activeSessionId]);
+  const formatTime = (d: string) =>
+    new Date(d).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' });
 
-  const formatTime = (dateString: string) => {
-    return new Date(dateString).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' });
-  };
+  // ── Render ────────────────────────────────────────────────────────
 
   return (
-    <Box className="sidebar" sx={{ height: '100%', display: 'flex', flexDirection: 'column', bgcolor: 'var(--panel)' }}>
-      {/* Header - Aligned with AppLayout 60px */}
-      <Box sx={{ 
-        height: '60px', 
-        display: 'flex', 
-        alignItems: 'center', 
-        px: 2,
-        borderBottom: '1px solid rgba(255,255,255,0.03)'
-      }}>
-        <Box className="brand-logo" sx={{ display: 'flex' }}>
-            <img 
-              src="/hf-log-only-white.png" 
-              alt="HF Agent" 
-              style={{ height: '24px', objectFit: 'contain' }} 
-            />
-        </Box>
+    <Box
+      sx={{
+        height: '100%',
+        display: 'flex',
+        flexDirection: 'column',
+        bgcolor: 'var(--panel)',
+      }}
+    >
+      {/* ── Header ─────────────────────────────────────────────────── */}
+      <Box sx={{ px: 1.75, pt: 2, pb: 0 }}>
+        <Typography
+          variant="caption"
+          sx={{
+            color: 'var(--muted-text)',
+            fontSize: '0.65rem',
+            fontWeight: 600,
+            textTransform: 'uppercase',
+            letterSpacing: '0.08em',
+          }}
+        >
+          Recent chats
+        </Typography>
       </Box>
 
-      {/* Content */}
-      <Box sx={{ flex: 1, display: 'flex', flexDirection: 'column', p: 2, overflow: 'hidden' }}>
-        {/* System Info / Status */}
-        <Box sx={{ mb: 2, display: 'flex', alignItems: 'center', gap: 1 }}>
-          <StatusDiode connected={isConnected} />
-          <Typography variant="caption" sx={{ color: 'var(--muted-text)', fontFamily: 'inherit' }}>
-            {isConnected ? 'System Online' : 'Disconnected'}
-          </Typography>
-        </Box>
+      {/* ── Capacity error ─────────────────────────────────────────── */}
+      {capacityError && (
+        <Alert
+          severity="warning"
+          variant="outlined"
+          onClose={() => setCapacityError(null)}
+          sx={{
+            m: 1,
+            fontSize: '0.7rem',
+            py: 0.25,
+            '& .MuiAlert-message': { py: 0 },
+            borderColor: '#FF9D00',
+            color: 'var(--text)',
+          }}
+        >
+          {capacityError}
+        </Alert>
+      )}
+
+      {/* ── Session list ───────────────────────────────────────────── */}
+      <Box
+        sx={{
+          flex: 1,
+          overflow: 'auto',
+          py: 1,
+          // Thinner scrollbar
+          '&::-webkit-scrollbar': { width: 4 },
+          '&::-webkit-scrollbar-thumb': {
+            bgcolor: 'var(--scrollbar-thumb)',
+            borderRadius: 2,
+          },
+        }}
+      >
+        {sessions.length === 0 ? (
+          <Box
+            sx={{
+              display: 'flex',
+              flexDirection: 'column',
+              alignItems: 'center',
+              justifyContent: 'center',
+              py: 8,
+              px: 3,
+              gap: 1.5,
+            }}
+          >
+            <ChatBubbleOutlineIcon
+              sx={{ fontSize: 28, color: 'var(--muted-text)', opacity: 0.25 }}
+            />
+            <Typography
+              variant="caption"
+              sx={{
+                color: 'var(--muted-text)',
+                opacity: 0.5,
+                textAlign: 'center',
+                lineHeight: 1.5,
+                fontSize: '0.72rem',
+              }}
+            >
+              No sessions yet
+            </Typography>
+          </Box>
+        ) : (
+          [...sessions].reverse().map((session, index) => {
+            const num = sessions.length - index;
+            const isSelected = session.id === activeSessionId;
 
-        <Button
-          fullWidth
-          className="create-session"
+            return (
+              <Box
+                key={session.id}
+                onClick={() => handleSelect(session.id)}
+                sx={{
+                  display: 'flex',
+                  alignItems: 'center',
+                  gap: 1,
+                  px: 1.5,
+                  py: 0.875,
+                  mx: 0.75,
+                  borderRadius: '10px',
+                  cursor: 'pointer',
+                  transition: 'background-color 0.12s ease',
+                  bgcolor: isSelected
+                    ? 'var(--hover-bg)'
+                    : 'transparent',
+                  '&:hover': {
+                    bgcolor: 'var(--hover-bg)',
+                  },
+                  '& .delete-btn': {
+                    opacity: 0,
+                    transition: 'opacity 0.12s',
+                  },
+                  '&:hover .delete-btn': {
+                    opacity: 1,
+                  },
+                }}
+              >
+                <ChatBubbleOutlineIcon
+                  sx={{
+                    fontSize: 15,
+                    color: isSelected ? 'var(--text)' : 'var(--muted-text)',
+                    opacity: isSelected ? 0.8 : 0.4,
+                    flexShrink: 0,
+                  }}
+                />
+
+                <Box sx={{ flex: 1, minWidth: 0 }}>
+                  <Typography
+                    variant="body2"
+                    sx={{
+                      fontWeight: isSelected ? 600 : 400,
+                      color: 'var(--text)',
+                      fontSize: '0.84rem',
+                      lineHeight: 1.4,
+                      whiteSpace: 'nowrap',
+                      overflow: 'hidden',
+                      textOverflow: 'ellipsis',
+                    }}
+                  >
+                    {session.title.startsWith('Chat ') ? `Session ${String(num).padStart(2, '0')}` : session.title}
+                  </Typography>
+                  <Typography
+                    variant="caption"
+                    sx={{
+                      color: 'var(--muted-text)',
+                      fontSize: '0.65rem',
+                      lineHeight: 1.2,
+                    }}
+                  >
+                    {formatTime(session.createdAt)}
+                  </Typography>
+                </Box>
+
+                <IconButton
+                  className="delete-btn"
+                  size="small"
+                  onClick={(e) => handleDelete(session.id, e)}
+                  sx={{
+                    color: 'var(--muted-text)',
+                    width: 26,
+                    height: 26,
+                    flexShrink: 0,
+                    '&:hover': { color: 'var(--accent-red)', bgcolor: 'rgba(244,67,54,0.08)' },
+                  }}
+                >
+                  <DeleteOutlineIcon sx={{ fontSize: 15 }} />
+                </IconButton>
+              </Box>
+            );
+          })
+        )}
+      </Box>
+
+      {/* ── Footer: New Session + status ──────────────────────────── */}
+      <Divider sx={{ opacity: 0.5 }} />
+      <Box
+        sx={{
+          px: 1.5,
+          py: 1.5,
+          display: 'flex',
+          flexDirection: 'column',
+          gap: 1,
+          flexShrink: 0,
+        }}
+      >
+        <Box
+          component="button"
           onClick={handleNewSession}
+          disabled={isCreatingSession}
           sx={{
             display: 'inline-flex',
             alignItems: 'center',
-            justifyContent: 'flex-start',
-            gap: '10px',
-            padding: '10px 14px',
-            borderRadius: 'var(--radius-md)',
-            border: '1px solid rgba(255,255,255,0.06)',
-            bgcolor: 'transparent',
-            color: 'var(--text)',
-            fontWeight: 600,
-            textTransform: 'none',
-            mb: 3,
+            justifyContent: 'center',
+            gap: 0.75,
+            width: '100%',
+            px: 1.5,
+            py: 1.25,
+            border: 'none',
+            borderRadius: '10px',
+            bgcolor: '#FF9D00',
+            color: '#000',
+            fontSize: '0.85rem',
+            fontWeight: 700,
+            cursor: 'pointer',
+            transition: 'all 0.12s ease',
             '&:hover': {
-                bgcolor: 'rgba(255,255,255,0.02)',
-                border: '1px solid rgba(255,255,255,0.1)',
+              bgcolor: '#FFB340',
+            },
+            '&:disabled': {
+              opacity: 0.5,
+              cursor: 'not-allowed',
             },
-            '&::before': {
-                content: '""',
-                width: '4px',
-                height: '20px',
-                background: 'linear-gradient(180deg, var(--accent-yellow), rgba(199,165,0,0.9))',
-                borderRadius: '4px',
-            }
           }}
         >
-          New Session
-        </Button>
-
-        {/* Session List */}
-        <Box sx={{ flex: 1, overflow: 'auto', mx: -1, px: 1 }}>
-            <List disablePadding sx={{ display: 'flex', flexDirection: 'column', gap: 1 }}>
-            {[...sessions].reverse().map((session, index) => {
-                const sessionNumber = sessions.length - index;
-                const isSelected = session.id === activeSessionId;
-                return (
-                <ListItem 
-                    key={session.id} 
-                    disablePadding 
-                    className="session-item"
-                    onClick={() => handleSelectSession(session.id)}
-                    sx={{
-                        display: 'flex',
-                        alignItems: 'center',
-                        gap: '12px',
-                        padding: '10px',
-                        borderRadius: 'var(--radius-md)',
-                        bgcolor: isSelected ? 'rgba(255,255,255,0.05)' : 'transparent',
-                        cursor: 'pointer',
-                        transition: 'background 0.18s ease, transform 0.08s ease',
-                        '&:hover': {
-                            bgcolor: 'rgba(255,255,255,0.02)',
-                            transform: 'translateY(-1px)',
-                        },
-                        '& .delete-btn': {
-                            opacity: 0,
-                            transition: 'opacity 0.2s',
-                        },
-                        '&:hover .delete-btn': {
-                            opacity: 1,
-                        }
-                    }}
-                >
-                    <Box sx={{ flex: 1, overflow: 'hidden' }}>
-                        <Typography variant="body2" sx={{ fontWeight: 500, color: 'var(--text)', whiteSpace: 'nowrap', overflow: 'hidden', textOverflow: 'ellipsis' }}>
-                            Session {String(sessionNumber).padStart(2, '0')}
-                        </Typography>
-                        <Box sx={{ display: 'flex', alignItems: 'center', gap: 1, mt: 0.5 }}>
-                            {session.isActive && <RunningIndicator />}
-                            <Typography className="time" variant="caption" sx={{ fontSize: '12px', color: 'var(--muted-text)' }}>
-                                {formatTime(session.createdAt)}
-                            </Typography>
-                        </Box>
-                    </Box>
-                    
-                    <IconButton
-                        className="delete-btn"
-                        size="small"
-                        onClick={(e) => handleDeleteSession(session.id, e)}
-                        sx={{ color: 'var(--muted-text)', '&:hover': { color: 'var(--accent-red)' } }}
-                    >
-                        <DeleteIcon fontSize="small" />
-                    </IconButton>
-                </ListItem>
-                );
-            })}
-            </List>
+          {isCreatingSession ? (
+            <>
+              <CircularProgress size={12} sx={{ color: '#000' }} />
+              Creating...
+            </>
+          ) : (
+            <>
+              <AddIcon sx={{ fontSize: 16 }} />
+              New Session
+            </>
+          )}
         </Box>
-      </Box>
 
-      {/* Footer */}
-      <Box sx={{ p: 2, borderTop: '1px solid rgba(255,255,255,0.03)' }}>
-        <Box sx={{ display: 'flex', alignItems: 'center', justifyContent: 'space-between' }}>
-            <Typography variant="caption" className="small-note" sx={{ fontSize: '12px', color: 'var(--muted-text)' }}>
-            {sessions.length} active
-            </Typography>
-            <Tooltip title="Undo last turn">
-            <span>
-                <IconButton
-                onClick={handleUndo}
-                disabled={!activeSessionId || isProcessing}
-                size="small"
-                sx={{ color: 'var(--muted-text)', '&:hover': { color: 'var(--text)' } }}
-                >
-                <UndoIcon fontSize="small" />
-                </IconButton>
-            </span>
-            </Tooltip>
+        <Box
+          sx={{
+            display: 'flex',
+            alignItems: 'center',
+            justifyContent: 'center',
+            gap: 0.5,
+          }}
+        >
+          <StatusDot connected={isConnected} />
+          <Typography
+            variant="caption"
+            sx={{ color: 'var(--muted-text)', fontSize: '0.62rem', letterSpacing: '0.02em' }}
+          >
+            {sessions.length} session{sessions.length !== 1 ? 's' : ''} &middot; Backend {isConnected ? 'online' : 'offline'}
+          </Typography>
         </Box>
       </Box>
     </Box>

frontend/src/components/WelcomeScreen/WelcomeScreen.tsx

@@ -0,0 +1,177 @@
+import { useState, useCallback } from 'react';
+import {
+  Box,
+  Typography,
+  Button,
+  CircularProgress,
+  Alert,
+} from '@mui/material';
+import { useSessionStore } from '@/store/sessionStore';
+import { useAgentStore } from '@/store/agentStore';
+import { apiFetch } from '@/utils/api';
+
+/** HF brand orange */
+const HF_ORANGE = '#FF9D00';
+
+export default function WelcomeScreen() {
+  const { createSession } = useSessionStore();
+  const { setPlan, setPanelContent } = useAgentStore();
+  const [isCreating, setIsCreating] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+
+  const handleStart = useCallback(async () => {
+    if (isCreating) return;
+    setIsCreating(true);
+    setError(null);
+
+    try {
+      const response = await apiFetch('/api/session', { method: 'POST' });
+      if (response.status === 503) {
+        const data = await response.json();
+        setError(data.detail || 'Server is at capacity. Please try again later.');
+        return;
+      }
+      if (!response.ok) {
+        setError('Failed to create session. Please try again.');
+        return;
+      }
+      const data = await response.json();
+      createSession(data.session_id);
+      setPlan([]);
+      setPanelContent(null);
+    } catch {
+      setError('Could not reach the server. Please try again.');
+    } finally {
+      setIsCreating(false);
+    }
+  }, [isCreating, createSession, setPlan, setPanelContent]);
+
+  return (
+    <Box
+      sx={{
+        width: '100%',
+        height: '100%',
+        display: 'flex',
+        flexDirection: 'column',
+        alignItems: 'center',
+        justifyContent: 'center',
+        background: 'var(--body-gradient)',
+        py: 8,
+      }}
+    >
+      {/* HF Logo — large, centered */}
+      <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',
+        }}
+      />
+
+      {/* Title */}
+      <Typography
+        variant="h2"
+        sx={{
+          fontWeight: 800,
+          color: 'var(--text)',
+          mb: 1.5,
+          letterSpacing: '-0.02em',
+          fontSize: { xs: '2rem', md: '2.8rem' },
+        }}
+      >
+        HF Agent
+      </Typography>
+
+      {/* 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>
+
+      {/* Start Button */}
+      <Button
+        variant="contained"
+        size="large"
+        onClick={handleStart}
+        disabled={isCreating}
+        startIcon={
+          isCreating ? <CircularProgress size={20} color="inherit" /> : null
+        }
+        sx={{
+          px: 5,
+          py: 1.5,
+          fontSize: '1rem',
+          fontWeight: 700,
+          textTransform: 'none',
+          borderRadius: '12px',
+          bgcolor: HF_ORANGE,
+          color: '#000',
+          boxShadow: '0 4px 24px rgba(255, 157, 0, 0.3)',
+          transition: 'all 0.2s ease',
+          '&:hover': {
+            bgcolor: '#FFB340',
+            boxShadow: '0 6px 32px rgba(255, 157, 0, 0.45)',
+          },
+          '&.Mui-disabled': {
+            bgcolor: 'rgba(255, 157, 0, 0.35)',
+            color: 'rgba(0,0,0,0.45)',
+          },
+        }}
+      >
+        {isCreating ? 'Initializing...' : 'Start Session'}
+      </Button>
+
+      {/* Error */}
+      {error && (
+        <Alert
+          severity="warning"
+          variant="outlined"
+          onClose={() => setError(null)}
+          sx={{
+            mt: 3,
+            maxWidth: 400,
+            fontSize: '0.8rem',
+            borderColor: HF_ORANGE,
+            color: 'var(--text)',
+          }}
+        >
+          {error}
+        </Alert>
+      )}
+
+      {/* Footnote */}
+      <Typography
+        variant="caption"
+        sx={{
+          mt: 5,
+          color: 'var(--muted-text)',
+          opacity: 0.5,
+          fontSize: '0.7rem',
+        }}
+      >
+        Conversations are stored locally in your browser.
+      </Typography>
+    </Box>
+  );
+}

frontend/src/hooks/useAgentWebSocket.ts

@@ -1,34 +1,40 @@
 import { useCallback, useEffect, useRef } from 'react';
-import { useAgentStore } from '@/store/agentStore';
+import { useAgentStore, type PlanItem } from '@/store/agentStore';
 import { useSessionStore } from '@/store/sessionStore';
 import { useLayoutStore } from '@/store/layoutStore';
+import { getWebSocketUrl } from '@/utils/api';
+import { logger } from '@/utils/logger';
 import type { AgentEvent } from '@/types/events';
 import type { Message, TraceLog } from '@/types/agent';
 
 const WS_RECONNECT_DELAY = 1000;
 const WS_MAX_RECONNECT_DELAY = 30000;
+const WS_MAX_RETRIES = 5;
 
 interface UseAgentWebSocketOptions {
   sessionId: string | null;
   onReady?: () => void;
   onError?: (error: string) => void;
+  onSessionDead?: (sessionId: string) => void;
 }
 
 export function useAgentWebSocket({
   sessionId,
   onReady,
   onError,
+  onSessionDead,
 }: UseAgentWebSocketOptions) {
   const wsRef = useRef<WebSocket | null>(null);
   const reconnectTimeoutRef = useRef<number | null>(null);
   const reconnectDelayRef = useRef(WS_RECONNECT_DELAY);
+  const retriesRef = useRef(0);
 
   const {
     addMessage,
     updateMessage,
+    appendToMessage,
     setProcessing,
     setConnected,
-    setPendingApprovals,
     setError,
     addTraceLog,
     updateTraceLog,
@@ -40,6 +46,7 @@ export function useAgentWebSocket({
     setPlan,
     setCurrentTurnMessageId,
     updateCurrentTurnTrace,
+    removeLastTurn,
   } = useAgentStore();
 
   const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
@@ -66,6 +73,48 @@ export function useAgentWebSocket({
           setCurrentTurnMessageId(null); // Start a new turn
           break;
 
+        // ── Streaming: individual token chunks ──────────────────
+        case 'assistant_chunk': {
+          const delta = (event.data?.content as string) || '';
+          if (!delta) break;
+
+          const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
+
+          if (currentTurnMsgId) {
+            // Append delta to the existing streaming message
+            appendToMessage(sessionId, currentTurnMsgId, delta);
+          } else {
+            // First chunk — create the message (with pending traces if any)
+            const currentTrace = useAgentStore.getState().traceLogs;
+            const messageId = `msg_${Date.now()}`;
+            const segments: Array<{ type: 'text' | 'tools'; content?: string; tools?: typeof currentTrace }> = [];
+
+            if (currentTrace.length > 0) {
+              segments.push({ type: 'tools', tools: [...currentTrace] });
+              clearTraceLogs();
+            }
+            segments.push({ type: 'text', content: delta });
+
+            const message: Message = {
+              id: messageId,
+              role: 'assistant',
+              content: delta,
+              timestamp: new Date().toISOString(),
+              segments,
+            };
+            addMessage(sessionId, message);
+            setCurrentTurnMessageId(messageId);
+          }
+          break;
+        }
+
+        // ── Streaming ended (text is already rendered via chunks) ─
+        case 'assistant_stream_end':
+          // Nothing to do — chunks already built the message.
+          // This event is just a signal that the stream is complete.
+          break;
+
+        // ── Legacy non-streaming full message (kept for backwards compat)
         case 'assistant_message': {
           const content = (event.data?.content as string) || '';
           const currentTrace = useAgentStore.getState().traceLogs;
@@ -126,23 +175,41 @@ export function useAgentWebSocket({
 
         case 'tool_call': {
           const toolName = (event.data?.tool as string) || 'unknown';
-          const args = (event.data?.arguments as Record<string, any>) || {};
+          const toolCallId = (event.data?.tool_call_id as string) || '';
+          const args = (event.data?.arguments as Record<string, string | undefined>) || {};
 
           // Don't display plan_tool in trace logs (it shows up elsewhere in the UI)
           if (toolName !== 'plan_tool') {
             const log: TraceLog = {
-              id: `tool_${Date.now()}`,
+              id: `tool_${Date.now()}_${toolCallId}`,
+              toolCallId,
               type: 'call',
               text: `Agent is executing ${toolName}...`,
               tool: toolName,
               timestamp: new Date().toISOString(),
               completed: false,
-              // Store args for auto-exec message creation later
-              args: toolName === 'hf_jobs' ? args : undefined,
+              args,
             };
             addTraceLog(log);
-            // Update the current turn message's trace in real-time
-            updateCurrentTurnTrace(sessionId);
+
+            // If no assistant message exists for this turn, create one now
+            // so the ToolCallGroup renders immediately in the chat flow.
+            const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
+            if (!currentTurnMsgId) {
+              const messageId = `msg_${Date.now()}`;
+              const currentTrace = useAgentStore.getState().traceLogs;
+              addMessage(sessionId, {
+                id: messageId,
+                role: 'assistant',
+                content: '',
+                timestamp: new Date().toISOString(),
+                segments: [{ type: 'tools', tools: [...currentTrace] }],
+              });
+              setCurrentTurnMessageId(messageId);
+              clearTraceLogs();
+            } else {
+              updateCurrentTurnTrace(sessionId);
+            }
           }
 
           // Auto-expand Right Panel for specific tools
@@ -171,68 +238,57 @@ export function useAgentWebSocket({
             setLeftSidebarOpen(false);
           }
 
-          console.log('Tool call:', toolName, args);
+          logger.log('Tool call:', toolName, args);
           break;
         }
 
         case 'tool_output': {
           const toolName = (event.data?.tool as string) || 'unknown';
+          const toolCallId = (event.data?.tool_call_id as string) || '';
           const output = (event.data?.output as string) || '';
           const success = event.data?.success as boolean;
 
-          // Mark the corresponding trace log as completed and store the output
-          updateTraceLog(toolName, { completed: true, output, success });
-          // Update the current turn message's trace in real-time
+          // Mark the corresponding trace log as completed and store the output.
+          // If it had a pending approval, mark it as approved (tool_output means it ran).
+          const prevLog = useAgentStore.getState().traceLogs.find(
+            (l) => l.toolCallId === toolCallId
+          );
+          const wasApproval = prevLog?.approvalStatus === 'pending';
+          updateTraceLog(toolCallId, toolName, {
+            completed: true,
+            output,
+            success,
+            ...(wasApproval ? { approvalStatus: 'approved' as const } : {}),
+          });
           updateCurrentTurnTrace(sessionId);
 
-          // Special handling for hf_jobs - update or create job message with output
-          if (toolName === 'hf_jobs') {
-            const messages = useAgentStore.getState().getMessages(sessionId);
-            const traceLogs = useAgentStore.getState().traceLogs;
-
-            // Find existing approval message for this job
-            let jobMsg = [...messages].reverse().find(m => m.approval);
+          // For hf_jobs: parse job output and enrich the TraceLog with job info
+          if (toolName === 'hf_jobs' && output) {
+            const updates: Partial<TraceLog> = { approvalStatus: 'approved' as const };
 
-            if (!jobMsg) {
-              // No approval message exists - this was an auto-executed job
-              // Create a job execution message so user can see results
-              const jobTrace = [...traceLogs].reverse().find(t => t.tool === 'hf_jobs');
-              const args = jobTrace?.args || {};
+            // Parse job URL
+            const urlMatch = output.match(/\*\*View at:\*\*\s*(https:\/\/[^\s\n]+)/);
+            if (urlMatch) updates.jobUrl = urlMatch[1];
 
-              const autoExecMessage: Message = {
-                id: `msg_auto_${Date.now()}`,
-                role: 'assistant',
-                content: '',
-                timestamp: new Date().toISOString(),
-                approval: {
-                  status: 'approved', // Auto-approved (no user action needed)
-                  batch: {
-                    tools: [{
-                      tool: toolName,
-                      arguments: args,
-                      tool_call_id: `auto_${Date.now()}`
-                    }],
-                    count: 1
-                  }
-                },
-                toolOutput: output
-              };
-              addMessage(sessionId, autoExecMessage);
-              console.log('Created auto-exec message with tool output:', toolName);
-            } else {
-              // Update existing approval message
-              const currentOutput = jobMsg.toolOutput || '';
-              const newOutput = currentOutput ? currentOutput + '\n\n' + output : output;
+            // Parse job status
+            const statusMatch = output.match(/\*\*Final Status:\*\*\s*([^\n]+)/);
+            if (statusMatch) updates.jobStatus = statusMatch[1].trim();
 
-              useAgentStore.getState().updateMessage(sessionId, jobMsg.id, {
-                toolOutput: newOutput
-              });
-              console.log('Updated job message with tool output:', toolName);
+            // Parse logs
+            if (output.includes('**Logs:**')) {
+              const parts = output.split('**Logs:**');
+              if (parts.length > 1) {
+                const codeBlockMatch = parts[1].trim().match(/```([\s\S]*?)```/);
+                if (codeBlockMatch) updates.jobLogs = codeBlockMatch[1].trim();
+              }
             }
+
+            updateTraceLog(toolCallId, toolName, updates);
+            updateCurrentTurnTrace(sessionId);
           }
 
           // Don't create message bubbles for tool outputs - they only show in trace logs
-          console.log('Tool output:', toolName, success);
+          logger.log('Tool output:', toolName, success);
           break;
         }
 
@@ -267,7 +323,7 @@ export function useAgentWebSocket({
         }
 
         case 'plan_update': {
-          const plan = (event.data?.plan as any[]) || [];
+          const plan = (event.data?.plan as PlanItem[]) || [];
           setPlan(plan);
           if (!useLayoutStore.getState().isRightPanelOpen) {
             setRightPanelOpen(true);
@@ -281,25 +337,59 @@ export function useAgentWebSocket({
             arguments: Record<string, unknown>;
             tool_call_id: string;
           }>;
-          const count = (event.data?.count as number) || 0;
-
-          // Create a persistent message for the approval request
-          const message: Message = {
-            id: `msg_approval_${Date.now()}`,
-            role: 'assistant',
-            content: '', // Content is handled by the approval UI
-            timestamp: new Date().toISOString(),
-            approval: {
-                status: 'pending',
-                batch: { tools, count }
+
+          // Create or update trace logs for approval tools.
+          // The backend only sends tool_call events for non-approval tools,
+          // so we must create TraceLogs here for approval-requiring tools.
+          if (tools) {
+            for (const t of tools) {
+              // Check if a TraceLog already exists (shouldn't, but be safe)
+              const existing = useAgentStore.getState().traceLogs.find(
+                (log) => log.toolCallId === t.tool_call_id
+              );
+              if (!existing) {
+                addTraceLog({
+                  id: `tool_${Date.now()}_${t.tool_call_id}`,
+                  toolCallId: t.tool_call_id,
+                  type: 'call',
+                  text: `Approval required for ${t.tool}`,
+                  tool: t.tool,
+                  timestamp: new Date().toISOString(),
+                  completed: false,
+                  args: t.arguments as Record<string, unknown>,
+                  approvalStatus: 'pending',
+                });
+              } else {
+                updateTraceLog(t.tool_call_id, t.tool, {
+                  approvalStatus: 'pending',
+                  args: t.arguments as Record<string, unknown>,
+                });
+              }
+            }
+
+            // Ensure there's a message to render the approval UI in
+            const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
+            if (!currentTurnMsgId) {
+              const messageId = `msg_${Date.now()}`;
+              const currentTrace = useAgentStore.getState().traceLogs;
+              addMessage(sessionId, {
+                id: messageId,
+                role: 'assistant',
+                content: '',
+                timestamp: new Date().toISOString(),
+                segments: [{ type: 'tools', tools: [...currentTrace] }],
+              });
+              setCurrentTurnMessageId(messageId);
+              clearTraceLogs();
+            } else {
+              updateCurrentTurnTrace(sessionId);
             }
-          };
-          addMessage(sessionId, message);
+          }
 
-          // Show the first tool's content in the panel so users see what they're approving
+          // Show the first tool's content in the panel
           if (tools && tools.length > 0) {
             const firstTool = tools[0];
-            const args = firstTool.arguments as Record<string, any>;
+            const args = firstTool.arguments as Record<string, string | undefined>;
 
             clearPanelTabs();
 
@@ -324,7 +414,6 @@ export function useAgentWebSocket({
               });
               setActivePanelTab('content');
             } else {
-              // For other tools, show args as JSON
               setPanelTab({
                 id: 'args',
                 title: firstTool.tool,
@@ -339,11 +428,6 @@ export function useAgentWebSocket({
             setLeftSidebarOpen(false);
           }
 
-          // Clear currentTurnMessageId so subsequent assistant_message events create a new message below the approval
-          setCurrentTurnMessageId(null);
-
-          // We don't set pendingApprovals in the global store anymore as the message handles the UI
-          setPendingApprovals(null);
           setProcessing(false);
           break;
         }
@@ -356,7 +440,7 @@ export function useAgentWebSocket({
         case 'compacted': {
           const oldTokens = event.data?.old_tokens as number;
           const newTokens = event.data?.new_tokens as number;
-          console.log(`Context compacted: ${oldTokens} -> ${newTokens} tokens`);
+          logger.log(`Context compacted: ${oldTokens} -> ${newTokens} tokens`);
           break;
         }
 
@@ -378,16 +462,19 @@ export function useAgentWebSocket({
           break;
 
         case 'undo_complete':
-          // Could remove last messages from store
+          if (sessionId) {
+            removeLastTurn(sessionId);
+          }
+          setProcessing(false);
           break;
 
         default:
-          console.log('Unknown event:', event);
+          logger.log('Unknown event:', event);
       }
     },
     // Zustand setters are stable, so we don't need them in deps
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    [sessionId, onReady, onError]
+    [sessionId, onReady, onError, onSessionDead]
   );
 
   const connect = useCallback(() => {
@@ -399,21 +486,17 @@ export function useAgentWebSocket({
       return;
     }
 
-    // Connect directly to backend (Vite doesn't proxy WebSockets)
-    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-    // In development, connect directly to backend port 7860
-    // In production, use the same host
-    const isDev = import.meta.env.DEV;
-    const host = isDev ? '127.0.0.1:7860' : window.location.host;
-    const wsUrl = `${protocol}//${host}/api/ws/${sessionId}`;
+    // Build WebSocket URL (centralized in utils/api.ts)
+    const wsUrl = getWebSocketUrl(sessionId);
 
-    console.log('Connecting to WebSocket:', wsUrl);
+    logger.log('Connecting to WebSocket:', wsUrl);
     const ws = new WebSocket(wsUrl);
 
     ws.onopen = () => {
-      console.log('WebSocket connected');
+      logger.log('WebSocket connected');
       setConnected(true);
       reconnectDelayRef.current = WS_RECONNECT_DELAY;
+      retriesRef.current = 0; // Reset retry counter on successful connect
     };
 
     ws.onmessage = (event) => {
@@ -421,20 +504,31 @@ export function useAgentWebSocket({
         const data = JSON.parse(event.data) as AgentEvent;
         handleEvent(data);
       } catch (e) {
-        console.error('Failed to parse WebSocket message:', e);
+        logger.error('Failed to parse WebSocket message:', e);
       }
     };
 
     ws.onerror = (error) => {
-      console.error('WebSocket error:', error);
+      logger.error('WebSocket error:', error);
     };
 
     ws.onclose = (event) => {
-      console.log('WebSocket closed', event.code, event.reason);
+      logger.log('WebSocket closed', event.code, event.reason);
       setConnected(false);
 
-      // Only reconnect if it wasn't a normal closure and session still exists
-      if (event.code !== 1000 && sessionId) {
+      // Don't reconnect if:
+      // - Normal closure (1000)
+      // - Session not found (4004) — session was deleted or backend restarted
+      // - Auth failed (4001) or access denied (4003) — won't succeed on retry
+      // - No session ID
+      const noRetryCodes = [1000, 4001, 4003, 4004];
+      if (!noRetryCodes.includes(event.code) && sessionId) {
+        retriesRef.current += 1;
+        if (retriesRef.current > WS_MAX_RETRIES) {
+          logger.warn(`WebSocket: max retries (${WS_MAX_RETRIES}) reached, giving up.`);
+          onSessionDead?.(sessionId);
+          return;
+        }
         // Attempt to reconnect with exponential backoff
         if (reconnectTimeoutRef.current) {
           clearTimeout(reconnectTimeoutRef.current);
@@ -446,6 +540,12 @@ export function useAgentWebSocket({
           );
           connect();
         }, reconnectDelayRef.current);
+      } else if (event.code === 4004 && sessionId) {
+        // Session not found — remove it from the store (lazy cleanup)
+        logger.warn(`Session ${sessionId} no longer exists on backend, removing.`);
+        onSessionDead?.(sessionId);
+      } else if (noRetryCodes.includes(event.code) && event.code !== 1000) {
+        logger.warn(`WebSocket permanently closed: ${event.code} ${event.reason}`);
       }
     };
 
@@ -477,6 +577,10 @@ export function useAgentWebSocket({
       return;
     }
 
+    // Reset retry state for new session
+    retriesRef.current = 0;
+    reconnectDelayRef.current = WS_RECONNECT_DELAY;
+
     // Small delay to ensure session is fully created on backend
     const timeoutId = setTimeout(() => {
       connect();

frontend/src/hooks/useAuth.ts

@@ -0,0 +1,53 @@
+/**
+ * Authentication hook — non-blocking.
+ *
+ * The app renders immediately. This hook fires a background check to /auth/me
+ * and updates the agent store with user info when it resolves.
+ * If an API call later returns 401, apiFetch handles the redirect to /auth/login.
+ *
+ * This avoids blocking the entire UI on an auth check that depends on backend
+ * availability (which can be slow during session/MCP initialization).
+ */
+
+import { useEffect } from 'react';
+import { useAgentStore } from '@/store/agentStore';
+
+export function useAuth() {
+  const setUser = useAgentStore((s) => s.setUser);
+
+  useEffect(() => {
+    async function checkAuth() {
+      try {
+        const response = await fetch('/auth/me', { credentials: 'include' });
+        if (response.ok) {
+          const data = await response.json();
+          if (data.authenticated) {
+            setUser({
+              authenticated: true,
+              username: data.username,
+              name: data.name,
+              picture: data.picture,
+            });
+            return;
+          }
+        }
+
+        // Not authenticated — check if auth is required
+        const statusRes = await fetch('/auth/status', { credentials: 'include' });
+        const statusData = await statusRes.json();
+        if (statusData.auth_enabled) {
+          window.location.href = '/auth/login';
+          return;
+        }
+
+        // Dev mode — set dev user
+        setUser({ authenticated: true, username: 'dev' });
+      } catch {
+        // Backend not ready — set dev user so the app is usable
+        setUser({ authenticated: true, username: 'dev' });
+      }
+    }
+
+    checkAuth();
+  }, [setUser]);
+}

frontend/src/main.tsx

@@ -3,13 +3,23 @@ import { createRoot } from 'react-dom/client';
 import { ThemeProvider } from '@mui/material/styles';
 import CssBaseline from '@mui/material/CssBaseline';
 import App from './App';
-import theme from './theme';
+import { darkTheme, lightTheme } from './theme';
+import { useLayoutStore } from './store/layoutStore';
 
-createRoot(document.getElementById('root')!).render(
-  <StrictMode>
+function Root() {
+  const themeMode = useLayoutStore((s) => s.themeMode);
+  const theme = themeMode === 'light' ? lightTheme : darkTheme;
+
+  return (
     <ThemeProvider theme={theme}>
       <CssBaseline />
       <App />
     </ThemeProvider>
+  );
+}
+
+createRoot(document.getElementById('root')!).render(
+  <StrictMode>
+    <Root />
   </StrictMode>
 );

frontend/src/store/agentStore.ts

@@ -1,5 +1,6 @@
 import { create } from 'zustand';
-import type { Message, ApprovalBatch, User, TraceLog } from '@/types/agent';
+import { persist } from 'zustand/middleware';
+import type { Message, User, TraceLog } from '@/types/agent';
 
 export interface PlanItem {
   id: string;
@@ -12,7 +13,13 @@ interface PanelTab {
   title: string;
   content: string;
   language?: string;
-  parameters?: any;
+  parameters?: Record<string, unknown>;
+}
+
+export interface LLMHealthError {
+  error: string;
+  errorType: 'auth' | 'credits' | 'rate_limit' | 'network' | 'unknown';
+  model: string;
 }
 
 interface AgentStore {
@@ -20,11 +27,11 @@ interface AgentStore {
   messagesBySession: Record<string, Message[]>;
   isProcessing: boolean;
   isConnected: boolean;
-  pendingApprovals: ApprovalBatch | null;
   user: User | null;
   error: string | null;
+  llmHealthError: LLMHealthError | null;
   traceLogs: TraceLog[];
-  panelContent: { title: string; content: string; language?: string; parameters?: any } | null;
+  panelContent: { title: string; content: string; language?: string; parameters?: Record<string, unknown> } | null;
   panelTabs: PanelTab[];
   activePanelTab: string | null;
   plan: PlanItem[];
@@ -36,14 +43,13 @@ interface AgentStore {
   clearMessages: (sessionId: string) => void;
   setProcessing: (isProcessing: boolean) => void;
   setConnected: (isConnected: boolean) => void;
-  setPendingApprovals: (approvals: ApprovalBatch | null) => void;
   setUser: (user: User | null) => void;
   setError: (error: string | null) => void;
   getMessages: (sessionId: string) => Message[];
   addTraceLog: (log: TraceLog) => void;
-  updateTraceLog: (toolName: string, updates: Partial<TraceLog>) => void;
+  updateTraceLog: (toolCallId: string, toolName: string, updates: Partial<TraceLog>) => void;
   clearTraceLogs: () => void;
-  setPanelContent: (content: { title: string; content: string; language?: string; parameters?: any } | null) => void;
+  setPanelContent: (content: { title: string; content: string; language?: string; parameters?: Record<string, unknown> } | null) => void;
   setPanelTab: (tab: PanelTab) => void;
   setActivePanelTab: (tabId: string) => void;
   clearPanelTabs: () => void;
@@ -52,15 +58,24 @@ interface AgentStore {
   setCurrentTurnMessageId: (id: string | null) => void;
   updateCurrentTurnTrace: (sessionId: string) => void;
   showToolOutput: (log: TraceLog) => void;
+  /** Append a streaming delta to an existing message. */
+  appendToMessage: (sessionId: string, messageId: string, delta: string) => void;
+  /** Remove all messages for a session (also clears from localStorage). */
+  deleteSessionMessages: (sessionId: string) => void;
+  /** Remove the last turn (last user msg + all following assistant/tool msgs). */
+  removeLastTurn: (sessionId: string) => void;
+  setLlmHealthError: (error: LLMHealthError | null) => void;
 }
 
-export const useAgentStore = create<AgentStore>((set, get) => ({
+export const useAgentStore = create<AgentStore>()(
+  persist(
+  (set, get) => ({
   messagesBySession: {},
   isProcessing: false,
   isConnected: false,
-  pendingApprovals: null,
   user: null,
   error: null,
+  llmHealthError: null,
   traceLogs: [],
   panelContent: null,
   panelTabs: [],
@@ -112,10 +127,6 @@ export const useAgentStore = create<AgentStore>((set, get) => ({
     set({ isConnected });
   },
 
-  setPendingApprovals: (approvals: ApprovalBatch | null) => {
-    set({ pendingApprovals: approvals });
-  },
-
   setUser: (user: User | null) => {
     set({ user });
   },
@@ -134,14 +145,27 @@ export const useAgentStore = create<AgentStore>((set, get) => ({
     }));
   },
 
-  updateTraceLog: (toolName: string, updates: Partial<TraceLog>) => {
+  updateTraceLog: (toolCallId: string, toolName: string, updates: Partial<TraceLog>) => {
     set((state) => {
-      // Find the last trace log with this tool name and update it
       const traceLogs = [...state.traceLogs];
-      for (let i = traceLogs.length - 1; i >= 0; i--) {
-        if (traceLogs[i].tool === toolName && traceLogs[i].type === 'call') {
-          traceLogs[i] = { ...traceLogs[i], ...updates };
-          break;
+      // Prefer matching by tool_call_id (reliable), fall back to tool name (legacy)
+      let matched = false;
+      if (toolCallId) {
+        for (let i = traceLogs.length - 1; i >= 0; i--) {
+          if (traceLogs[i].toolCallId === toolCallId) {
+            traceLogs[i] = { ...traceLogs[i], ...updates };
+            matched = true;
+            break;
+          }
+        }
+      }
+      if (!matched) {
+        // Fallback: match by tool name (last uncompleted call)
+        for (let i = traceLogs.length - 1; i >= 0; i--) {
+          if (traceLogs[i].tool === toolName && traceLogs[i].type === 'call' && !traceLogs[i].completed) {
+            traceLogs[i] = { ...traceLogs[i], ...updates };
+            break;
+          }
         }
       }
       return { traceLogs };
@@ -208,20 +232,39 @@ export const useAgentStore = create<AgentStore>((set, get) => ({
 
   updateCurrentTurnTrace: (sessionId: string) => {
     const state = get();
-    if (state.currentTurnMessageId) {
-      const currentMessages = state.messagesBySession[sessionId] || [];
-      const updatedMessages = currentMessages.map((msg) =>
-        msg.id === state.currentTurnMessageId
-          ? { ...msg, trace: state.traceLogs.length > 0 ? [...state.traceLogs] : undefined }
-          : msg
-      );
-      set({
-        messagesBySession: {
-          ...state.messagesBySession,
-          [sessionId]: updatedMessages,
-        },
-      });
-    }
+    if (!state.currentTurnMessageId) return;
+
+    const currentMessages = state.messagesBySession[sessionId] || [];
+    const latestTools = state.traceLogs.length > 0 ? [...state.traceLogs] : undefined;
+    if (!latestTools) return;
+
+    const updatedMessages = currentMessages.map((msg) => {
+      if (msg.id !== state.currentTurnMessageId) return msg;
+
+      const segments = msg.segments ? [...msg.segments] : [];
+      const lastToolsIdx = segments.map((s) => s.type).lastIndexOf('tools');
+
+      if (lastToolsIdx >= 0 && lastToolsIdx === segments.length - 1) {
+        // Last segment IS a tools segment — update it in place
+        segments[lastToolsIdx] = { type: 'tools', tools: latestTools };
+      } else if (lastToolsIdx >= 0) {
+        // A tools segment exists but is NOT last (text came after it).
+        // Append a NEW tools segment at the end.
+        segments.push({ type: 'tools', tools: latestTools });
+      } else {
+        // No tools segment at all — create one at the end.
+        segments.push({ type: 'tools', tools: latestTools });
+      }
+
+      return { ...msg, segments };
+    });
+
+    set({
+      messagesBySession: {
+        ...state.messagesBySession,
+        [sessionId]: updatedMessages,
+      },
+    });
   },
 
   showToolOutput: (log: TraceLog) => {
@@ -257,4 +300,80 @@ export const useAgentStore = create<AgentStore>((set, get) => ({
       activePanelTab: 'tool_output',
     });
   },
-}));
+
+  appendToMessage: (sessionId: string, messageId: string, delta: string) => {
+    set((state) => {
+      const messages = state.messagesBySession[sessionId] || [];
+      return {
+        messagesBySession: {
+          ...state.messagesBySession,
+          [sessionId]: messages.map((msg) => {
+            if (msg.id !== messageId) return msg;
+            const newContent = msg.content + delta;
+            const segments = msg.segments ? [...msg.segments] : [];
+            const lastSeg = segments[segments.length - 1];
+
+            if (lastSeg && lastSeg.type === 'text') {
+              // Append to the existing text segment
+              segments[segments.length - 1] = {
+                ...lastSeg,
+                content: (lastSeg.content || '') + delta,
+              };
+            } else {
+              // Last segment is 'tools' (or empty) — start a NEW text segment
+              // so that tools and text remain visually separated.
+              segments.push({ type: 'text', content: delta });
+            }
+
+            return { ...msg, content: newContent, segments };
+          }),
+        },
+      };
+    });
+  },
+
+  deleteSessionMessages: (sessionId: string) => {
+    set((state) => {
+      const { [sessionId]: _, ...rest } = state.messagesBySession;
+      return { messagesBySession: rest };
+    });
+  },
+
+  removeLastTurn: (sessionId: string) => {
+    set((state) => {
+      const msgs = state.messagesBySession[sessionId];
+      if (!msgs || msgs.length === 0) return state;
+
+      // Find the index of the last user message
+      let lastUserIdx = -1;
+      for (let i = msgs.length - 1; i >= 0; i--) {
+        if (msgs[i].role === 'user') {
+          lastUserIdx = i;
+          break;
+        }
+      }
+      if (lastUserIdx === -1) return state; // no user message to remove
+
+      // Remove everything from that user message onward
+      return {
+        messagesBySession: {
+          ...state.messagesBySession,
+          [sessionId]: msgs.slice(0, lastUserIdx),
+        },
+      };
+    });
+  },
+
+  setLlmHealthError: (error: LLMHealthError | null) => {
+    set({ llmHealthError: error });
+  },
+}),
+    {
+      name: 'hf-agent-messages',
+      // Only persist messages — all transient UI state stays in-memory
+      partialize: (state) => ({
+        messagesBySession: state.messagesBySession,
+      }),
+    }
+  )
+);

frontend/src/store/layoutStore.ts

@@ -1,23 +1,41 @@
 import { create } from 'zustand';
+import { persist } from 'zustand/middleware';
+
+export type ThemeMode = 'dark' | 'light';
 
 interface LayoutStore {
   isLeftSidebarOpen: boolean;
   isRightPanelOpen: boolean;
   rightPanelWidth: number;
+  themeMode: ThemeMode;
   setLeftSidebarOpen: (open: boolean) => void;
   setRightPanelOpen: (open: boolean) => void;
   setRightPanelWidth: (width: number) => void;
   toggleLeftSidebar: () => void;
   toggleRightPanel: () => void;
+  toggleTheme: () => void;
 }
 
-export const useLayoutStore = create<LayoutStore>((set) => ({
-  isLeftSidebarOpen: true,
-  isRightPanelOpen: false,
-  rightPanelWidth: 450,
-  setLeftSidebarOpen: (open) => set({ isLeftSidebarOpen: open }),
-  setRightPanelOpen: (open) => set({ isRightPanelOpen: open }),
-  setRightPanelWidth: (width) => set({ rightPanelWidth: width }),
-  toggleLeftSidebar: () => set((state) => ({ isLeftSidebarOpen: !state.isLeftSidebarOpen })),
-  toggleRightPanel: () => set((state) => ({ isRightPanelOpen: !state.isRightPanelOpen })),
-}));
+export const useLayoutStore = create<LayoutStore>()(
+  persist(
+    (set) => ({
+      isLeftSidebarOpen: true,
+      isRightPanelOpen: false,
+      rightPanelWidth: 450,
+      themeMode: 'dark' as ThemeMode,
+      setLeftSidebarOpen: (open) => set({ isLeftSidebarOpen: open }),
+      setRightPanelOpen: (open) => set({ isRightPanelOpen: open }),
+      setRightPanelWidth: (width) => set({ rightPanelWidth: width }),
+      toggleLeftSidebar: () => set((state) => ({ isLeftSidebarOpen: !state.isLeftSidebarOpen })),
+      toggleRightPanel: () => set((state) => ({ isRightPanelOpen: !state.isRightPanelOpen })),
+      toggleTheme: () =>
+        set((state) => ({
+          themeMode: state.themeMode === 'dark' ? 'light' : 'dark',
+        })),
+    }),
+    {
+      name: 'hf-agent-layout',
+      partialize: (state) => ({ themeMode: state.themeMode }),
+    }
+  )
+);

frontend/src/store/sessionStore.ts

@@ -1,6 +1,7 @@
 import { create } from 'zustand';
 import { persist } from 'zustand/middleware';
 import type { SessionMeta } from '@/types/agent';
+import { useAgentStore } from './agentStore';
 
 interface SessionStore {
   sessions: SessionMeta[];
@@ -10,8 +11,8 @@ interface SessionStore {
   createSession: (id: string) => void;
   deleteSession: (id: string) => void;
   switchSession: (id: string) => void;
-  updateSessionTitle: (id: string, title: string) => void;
   setSessionActive: (id: string, isActive: boolean) => void;
+  updateSessionTitle: (id: string, title: string) => void;
 }
 
 export const useSessionStore = create<SessionStore>()(
@@ -34,6 +35,9 @@ export const useSessionStore = create<SessionStore>()(
       },
 
       deleteSession: (id: string) => {
+        // Clean up persisted messages for this session
+        useAgentStore.getState().deleteSessionMessages(id);
+
         set((state) => {
           const newSessions = state.sessions.filter((s) => s.id !== id);
           const newActiveId =
@@ -51,18 +55,18 @@ export const useSessionStore = create<SessionStore>()(
         set({ activeSessionId: id });
       },
 
-      updateSessionTitle: (id: string, title: string) => {
+      setSessionActive: (id: string, isActive: boolean) => {
         set((state) => ({
           sessions: state.sessions.map((s) =>
-            s.id === id ? { ...s, title } : s
+            s.id === id ? { ...s, isActive } : s
           ),
         }));
       },
 
-      setSessionActive: (id: string, isActive: boolean) => {
+      updateSessionTitle: (id: string, title: string) => {
         set((state) => ({
           sessions: state.sessions.map((s) =>
-            s.id === id ? { ...s, isActive } : s
+            s.id === id ? { ...s, title } : s
           ),
         }));
       },

frontend/src/theme.ts

@@ -1,158 +1,223 @@
-import { createTheme } from '@mui/material/styles';
+import { createTheme, type ThemeOptions } from '@mui/material/styles';
 
-const theme = createTheme({
-  palette: {
-    mode: 'dark',
-    primary: {
-      main: '#C7A500', // --accent-yellow
-    },
-    secondary: {
-      main: '#FF9D00',
-    },
-    background: {
-      default: '#0B0D10', // --bg
-      paper: '#0F1316',   // --panel
-    },
-    text: {
-      primary: '#E6EEF8', // --text
-      secondary: '#98A0AA', // --muted-text
-    },
-    divider: 'rgba(255,255,255,0.03)',
-    success: {
-      main: '#2FCC71', // --accent-green
-    },
-    error: {
-      main: '#E05A4F', // --accent-red
-    },
-    warning: {
-      main: '#C7A500',
-    },
-    info: {
-      main: '#58A6FF',
+// ── Shared tokens ────────────────────────────────────────────────
+const sharedTypography: ThemeOptions['typography'] = {
+  fontFamily: 'Inter, system-ui, -apple-system, "Segoe UI", Roboto, Arial, sans-serif',
+  fontSize: 15,
+  button: {
+    fontFamily: 'Inter, system-ui, -apple-system, "Segoe UI", Roboto, Arial, sans-serif',
+    textTransform: 'none' as const,
+    fontWeight: 600,
+  },
+};
+
+const sharedComponents: ThemeOptions['components'] = {
+  MuiButton: {
+    styleOverrides: {
+      root: {
+        borderRadius: '10px',
+        fontWeight: 600,
+        transition: 'transform 0.06s ease, background 0.12s ease, box-shadow 0.12s ease',
+        '&:hover': { transform: 'translateY(-1px)' },
+      },
     },
   },
-  typography: {
-    fontFamily: 'Inter, system-ui, -apple-system, "Segoe UI", Roboto, Arial, sans-serif',
-    fontSize: 15,
-    h1: { fontWeight: 600, color: '#E6EEF8' },
-    h2: { fontWeight: 600, color: '#E6EEF8' },
-    h3: { fontWeight: 600, color: '#E6EEF8' },
-    h4: { fontWeight: 600, color: '#E6EEF8' },
-    h5: { fontWeight: 600, color: '#E6EEF8' },
-    h6: { fontWeight: 600, color: '#E6EEF8' },
-    body1: { fontSize: '15px', color: '#E6EEF8' },
-    body2: { fontSize: '0.875rem', color: '#98A0AA' },
-    button: {
-      fontFamily: 'Inter, system-ui, -apple-system, "Segoe UI", Roboto, Arial, sans-serif',
-      textTransform: 'none',
-      fontWeight: 600,
+  MuiPaper: {
+    styleOverrides: {
+      root: { backgroundImage: 'none', backgroundColor: 'transparent' },
     },
   },
-  components: {
-    MuiCssBaseline: {
-      styleOverrides: {
-        ':root': {
-          '--bg': '#0B0D10',
-          '--panel': '#0F1316',
-          '--surface': '#121416',
-          '--text': '#E6EEF8',
-          '--muted-text': '#98A0AA',
-          '--accent-yellow': '#C7A500',
-          '--accent-yellow-weak': 'rgba(199,165,0,0.08)',
-          '--accent-green': '#2FCC71',
-          '--accent-red': '#E05A4F',
-          '--shadow-1': '0 6px 18px rgba(2,6,12,0.55)',
-          '--radius-lg': '20px',
-          '--radius-md': '12px',
-          '--focus': '0 0 0 3px rgba(199,165,0,0.12)',
-        },
-        body: {
-          background: 'linear-gradient(180deg, var(--bg), #090B0D)',
-          color: 'var(--text)',
-          scrollbarWidth: 'thin',
-          '&::-webkit-scrollbar': {
-            width: '8px',
-            height: '8px',
-          },
-          '&::-webkit-scrollbar-thumb': {
-            backgroundColor: '#30363D',
-            borderRadius: '2px',
-          },
-          '&::-webkit-scrollbar-track': {
-            backgroundColor: 'transparent',
-          },
-        },
-        'code, pre': {
-          fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, "Roboto Mono", monospace',
-        },
-        '.brand-logo': {
-          position: 'relative',
-          padding: '6px',
-          borderRadius: '8px',
-          '&::after': {
-            content: '""',
-            position: 'absolute',
-            inset: '-6px',
-            borderRadius: '10px',
-            background: 'var(--accent-yellow-weak)',
-            zIndex: -1,
-            pointerEvents: 'none',
-          },
+};
+
+const sharedShape: ThemeOptions['shape'] = { borderRadius: 12 };
+
+// ── Dark palette ─────────────────────────────────────────────────
+const darkVars = {
+  '--bg': '#0B0D10',
+  '--panel': '#0F1316',
+  '--surface': '#121416',
+  '--text': '#E6EEF8',
+  '--muted-text': '#98A0AA',
+  '--accent-yellow': '#FF9D00',
+  '--accent-yellow-weak': 'rgba(255,157,0,0.08)',
+  '--accent-green': '#2FCC71',
+  '--accent-red': '#E05A4F',
+  '--shadow-1': '0 6px 18px rgba(2,6,12,0.55)',
+  '--radius-lg': '20px',
+  '--radius-md': '12px',
+  '--focus': '0 0 0 3px rgba(255,157,0,0.12)',
+  '--border': 'rgba(255,255,255,0.03)',
+  '--border-hover': 'rgba(255,255,255,0.1)',
+  '--code-bg': 'rgba(0,0,0,0.5)',
+  '--tool-bg': 'rgba(0,0,0,0.3)',
+  '--tool-border': 'rgba(255,255,255,0.05)',
+  '--hover-bg': 'rgba(255,255,255,0.05)',
+  '--composer-bg': 'rgba(255,255,255,0.01)',
+  '--msg-gradient': 'linear-gradient(180deg, rgba(255,255,255,0.015), transparent)',
+  '--body-gradient': 'linear-gradient(180deg, #0B0D10, #090B0D)',
+  '--scrollbar-thumb': '#30363D',
+  '--success-icon': '#FDB022',
+  '--error-icon': '#F87171',
+  '--clickable-text': 'rgba(255, 255, 255, 0.9)',
+  '--clickable-underline': 'rgba(255,255,255,0.3)',
+  '--code-panel-bg': '#0A0B0C',
+  '--tab-active-bg': 'rgba(255,255,255,0.08)',
+  '--tab-active-border': 'rgba(255,255,255,0.1)',
+  '--tab-hover-bg': 'rgba(255,255,255,0.05)',
+  '--tab-close-hover': 'rgba(255,255,255,0.1)',
+  '--plan-bg': 'rgba(0,0,0,0.2)',
+} as const;
+
+// ── Light palette ────────────────────────────────────────────────
+const lightVars = {
+  '--bg': '#FFFFFF',
+  '--panel': '#F7F8FA',
+  '--surface': '#F0F1F3',
+  '--text': '#1A1A2E',
+  '--muted-text': '#6B7280',
+  '--accent-yellow': '#FF9D00',
+  '--accent-yellow-weak': 'rgba(255,157,0,0.08)',
+  '--accent-green': '#16A34A',
+  '--accent-red': '#DC2626',
+  '--shadow-1': '0 4px 12px rgba(0,0,0,0.08)',
+  '--radius-lg': '20px',
+  '--radius-md': '12px',
+  '--focus': '0 0 0 3px rgba(255,157,0,0.15)',
+  '--border': 'rgba(0,0,0,0.08)',
+  '--border-hover': 'rgba(0,0,0,0.15)',
+  '--code-bg': 'rgba(0,0,0,0.04)',
+  '--tool-bg': 'rgba(0,0,0,0.03)',
+  '--tool-border': 'rgba(0,0,0,0.08)',
+  '--hover-bg': 'rgba(0,0,0,0.04)',
+  '--composer-bg': 'rgba(0,0,0,0.02)',
+  '--msg-gradient': 'linear-gradient(180deg, rgba(0,0,0,0.01), transparent)',
+  '--body-gradient': 'linear-gradient(180deg, #FFFFFF, #F7F8FA)',
+  '--scrollbar-thumb': '#C4C8CC',
+  '--success-icon': '#FF9D00',
+  '--error-icon': '#DC2626',
+  '--clickable-text': 'rgba(0, 0, 0, 0.85)',
+  '--clickable-underline': 'rgba(0,0,0,0.25)',
+  '--code-panel-bg': '#F5F6F8',
+  '--tab-active-bg': 'rgba(0,0,0,0.06)',
+  '--tab-active-border': 'rgba(0,0,0,0.1)',
+  '--tab-hover-bg': 'rgba(0,0,0,0.04)',
+  '--tab-close-hover': 'rgba(0,0,0,0.08)',
+  '--plan-bg': 'rgba(0,0,0,0.03)',
+} as const;
+
+// ── Shared CSS baseline (scrollbar, code, brand-logo) ────────────
+function makeCssBaseline(vars: Record<string, string>) {
+  return {
+    styleOverrides: {
+      ':root': vars,
+      body: {
+        background: 'var(--body-gradient)',
+        color: 'var(--text)',
+        scrollbarWidth: 'thin' as const,
+        '&::-webkit-scrollbar': { width: '8px', height: '8px' },
+        '&::-webkit-scrollbar-thumb': {
+          backgroundColor: 'var(--scrollbar-thumb)',
+          borderRadius: '2px',
         },
+        '&::-webkit-scrollbar-track': { backgroundColor: 'transparent' },
       },
-    },
-    MuiButton: {
-      styleOverrides: {
-        root: {
-          borderRadius: '10px',
-          fontWeight: 600,
-          transition: 'transform 0.06s ease, background 0.12s ease, box-shadow 0.12s ease',
-          '&:hover': {
-            transform: 'translateY(-1px)',
-          },
-        },
+      'code, pre': {
+        fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, "Roboto Mono", monospace',
       },
-    },
-    MuiPaper: {
-      styleOverrides: {
-        root: {
-          backgroundImage: 'none',
-          backgroundColor: 'transparent', // Default to transparent for gradients
+      '.brand-logo': {
+        position: 'relative' as const,
+        padding: '6px',
+        borderRadius: '8px',
+        '&::after': {
+          content: '""',
+          position: 'absolute' as const,
+          inset: '-6px',
+          borderRadius: '10px',
+          background: 'var(--accent-yellow-weak)',
+          zIndex: -1,
+          pointerEvents: 'none' as const,
         },
       },
     },
-    MuiDrawer: {
-      styleOverrides: {
-        paper: {
-          backgroundColor: 'var(--panel)',
-          borderRight: '1px solid rgba(255,255,255,0.03)',
-        },
+  };
+}
+
+function makeDrawer() {
+  return {
+    styleOverrides: {
+      paper: {
+        backgroundColor: 'var(--panel)',
+        borderRight: '1px solid var(--border)',
       },
     },
-    MuiTextField: {
-      styleOverrides: {
-        root: {
-          '& .MuiOutlinedInput-root': {
-            borderRadius: 'var(--radius-md)',
-            '& fieldset': {
-              borderColor: 'rgba(255,255,255,0.03)',
-            },
-            '&:hover fieldset': {
-              borderColor: 'rgba(255,255,255,0.1)',
-            },
-            '&.Mui-focused fieldset': {
-              borderColor: 'var(--accent-yellow)',
-              borderWidth: '1px',
-              boxShadow: 'var(--focus)',
-            },
+  };
+}
+
+function makeTextField() {
+  return {
+    styleOverrides: {
+      root: {
+        '& .MuiOutlinedInput-root': {
+          borderRadius: 'var(--radius-md)',
+          '& fieldset': { borderColor: 'var(--border)' },
+          '&:hover fieldset': { borderColor: 'var(--border-hover)' },
+          '&.Mui-focused fieldset': {
+            borderColor: 'var(--accent-yellow)',
+            borderWidth: '1px',
+            boxShadow: 'var(--focus)',
           },
         },
       },
     },
+  };
+}
+
+// ── Theme builders ───────────────────────────────────────────────
+export const darkTheme = createTheme({
+  palette: {
+    mode: 'dark',
+    primary: { main: '#FF9D00' },
+    secondary: { main: '#C7A500' },
+    background: { default: '#0B0D10', paper: '#0F1316' },
+    text: { primary: '#E6EEF8', secondary: '#98A0AA' },
+    divider: 'rgba(255,255,255,0.03)',
+    success: { main: '#2FCC71' },
+    error: { main: '#E05A4F' },
+    warning: { main: '#FF9D00' },
+    info: { main: '#58A6FF' },
+  },
+  typography: sharedTypography,
+  components: {
+    ...sharedComponents,
+    MuiCssBaseline: makeCssBaseline(darkVars),
+    MuiDrawer: makeDrawer(),
+    MuiTextField: makeTextField(),
   },
-  shape: {
-    borderRadius: 12,
+  shape: sharedShape,
+});
+
+export const lightTheme = createTheme({
+  palette: {
+    mode: 'light',
+    primary: { main: '#FF9D00' },
+    secondary: { main: '#B8960A' },
+    background: { default: '#FFFFFF', paper: '#F7F8FA' },
+    text: { primary: '#1A1A2E', secondary: '#6B7280' },
+    divider: 'rgba(0,0,0,0.08)',
+    success: { main: '#16A34A' },
+    error: { main: '#DC2626' },
+    warning: { main: '#FF9D00' },
+    info: { main: '#2563EB' },
+  },
+  typography: sharedTypography,
+  components: {
+    ...sharedComponents,
+    MuiCssBaseline: makeCssBaseline(lightVars),
+    MuiDrawer: makeDrawer(),
+    MuiTextField: makeTextField(),
   },
+  shape: sharedShape,
 });
 
-export default theme;
+// Keep default export for backwards compat
+export default darkTheme;

frontend/src/types/agent.ts

@@ -52,8 +52,11 @@ export interface ApprovalBatch {
   count: number;
 }
 
+export type ApprovalStatus = 'none' | 'pending' | 'approved' | 'rejected';
+
 export interface TraceLog {
   id: string;
+  toolCallId?: string; // Backend tool_call_id for reliable matching
   type: 'call' | 'output';
   text: string;
   tool: string;
@@ -62,6 +65,12 @@ export interface TraceLog {
   args?: Record<string, unknown>; // Store args for auto-exec jobs
   output?: string; // Store tool output for display
   success?: boolean; // Whether the tool call succeeded
+  /** Approval state for tools that need user confirmation */
+  approvalStatus?: ApprovalStatus;
+  /** Parsed job info (URL, status, logs) for hf_jobs */
+  jobUrl?: string;
+  jobStatus?: string;
+  jobLogs?: string;
 }
 
 export interface User {

frontend/src/types/events.ts

@@ -6,6 +6,8 @@ export type EventType =
   | 'ready'
   | 'processing'
   | 'assistant_message'
+  | 'assistant_chunk'
+  | 'assistant_stream_end'
   | 'tool_call'
   | 'tool_output'
   | 'tool_log'

frontend/src/utils/api.ts

@@ -0,0 +1,58 @@
+/**
+ * Centralized API utilities with automatic auth header injection.
+ *
+ * In production (OAuth enabled):
+ *   - REST calls include the HttpOnly cookie automatically (same-origin)
+ *   - WebSocket passes token via query parameter
+ *
+ * In development (no OAuth):
+ *   - Auth is bypassed on the backend, no token needed
+ */
+
+/** Get the base URL for API calls (handles dev proxy vs production) */
+function getApiBase(): string {
+  // In development, Vite proxies /api and /auth to the backend
+  // In production, same origin
+  return '';
+}
+
+/** Wrapper around fetch that includes credentials (cookies) and common headers. */
+export async function apiFetch(
+  path: string,
+  options: RequestInit = {}
+): Promise<Response> {
+  const url = `${getApiBase()}${path}`;
+
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/json',
+    ...(options.headers as Record<string, string>),
+  };
+
+  const response = await fetch(url, {
+    ...options,
+    headers,
+    credentials: 'include', // Send cookies (hf_access_token) with every request
+  });
+
+  // Handle 401 - redirect to login if auth is required
+  if (response.status === 401) {
+    const authStatus = await fetch(`${getApiBase()}/auth/status`, {
+      credentials: 'include',
+    });
+    const data = await authStatus.json();
+    if (data.auth_enabled) {
+      window.location.href = '/auth/login';
+      throw new Error('Authentication required — redirecting to login.');
+    }
+  }
+
+  return response;
+}
+
+/** Build the WebSocket URL for a session, including auth token if available. */
+export function getWebSocketUrl(sessionId: string): string {
+  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+  // Always use same origin — Vite proxy (ws: true) handles dev,
+  // same origin works directly in production. No cross-origin issues.
+  return `${protocol}//${window.location.host}/api/ws/${sessionId}`;
+}

frontend/src/utils/logger.ts

@@ -0,0 +1,24 @@
+/**
+ * Lightweight logger that silences verbose output in production.
+ *
+ * - `log` / `debug` are only emitted when `import.meta.env.DEV` is true.
+ * - `warn` and `error` always go through so real issues surface in prod.
+ */
+
+const isDev = import.meta.env.DEV;
+
+/* eslint-disable no-console */
+export const logger = {
+  /** Debug-level log — DEV only. */
+  log: (...args: unknown[]) => {
+    if (isDev) console.log(...args);
+  },
+  /** Debug-level log — DEV only. */
+  debug: (...args: unknown[]) => {
+    if (isDev) console.debug(...args);
+  },
+  /** Warning — always emitted. */
+  warn: console.warn.bind(console),
+  /** Error — always emitted. */
+  error: console.error.bind(console),
+};

frontend/vite.config.ts

@@ -15,6 +15,7 @@ export default defineConfig({
       '/api': {
         target: 'http://localhost:7860',
         changeOrigin: true,
+        ws: true, // Proxy WebSocket connections (/api/ws/...)
       },
       '/auth': {
         target: 'http://localhost:7860',