subprocess uploading

.gitignore

@@ -15,4 +15,5 @@ wheels/
 *.csv
 /logs
 hf-agent-leaderboard/
-.cursor/
+.cursor/
+session_logs/

agent/config.py

@@ -21,6 +21,7 @@ class Config(BaseModel):
     mcpServers: dict[str, MCPServerConfig] = {}
     save_sessions: bool = True
     session_dataset_repo: str = "smolagents/hf-agent-sessions"
+    auto_save_interval: int = 3  # Save every N user turns (0 = disabled)
 
 
 def substitute_env_vars(obj: Any) -> Any:

agent/core/agent_loop.py

@@ -255,6 +255,11 @@ class Handlers:
                 data={"history_size": len(session.context_manager.items)},
             )
         )
+
+        # Increment turn counter and check for auto-save
+        session.increment_turn()
+        await session.auto_save_if_needed()
+
         return final_response
 
     @staticmethod
@@ -414,15 +419,13 @@ class Handlers:
     @staticmethod
     async def shutdown(session: Session) -> bool:
         """Handle shutdown (like shutdown in codex.rs:1329)"""
-        # Save session trajectory if enabled
+        # Save session trajectory if enabled (fire-and-forget, returns immediately)
         if session.config.save_sessions:
-            print("💾 Saving session trajectory...")
+            print("💾 Saving session...")
             repo_id = session.config.session_dataset_repo
-            url = await session.push_to_dataset(repo_id)
-            if url:
-                print(f"✅ Session saved to: {url}")
-            else:
-                print("❌ Failed to save session")
+            local_path = 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"))
@@ -484,26 +487,47 @@ async def submission_loop(
     session = Session(event_queue, config=config, tool_router=tool_router)
     print("Agent loop started")
 
-    # Main processing loop
-    async with tool_router:
-        # Emit ready event after initialization
-        await session.send_event(
-            Event(event_type="ready", data={"message": "Agent initialized"})
+    # Retry any failed uploads from previous sessions (fire-and-forget)
+    if config and config.save_sessions:
+        Session.retry_failed_uploads_detached(
+            directory="session_logs", repo_id=config.session_dataset_repo
         )
 
-        while session.is_running:
-            submission = await submission_queue.get()
+    try:
+        # Main processing loop
+        async with tool_router:
+            # Emit ready event after initialization
+            await session.send_event(
+                Event(event_type="ready", data={"message": "Agent initialized"})
+            )
 
-            try:
-                should_continue = await process_submission(session, submission)
-                if not should_continue:
+            while session.is_running:
+                submission = await submission_queue.get()
+
+                try:
+                    should_continue = await process_submission(session, submission)
+                    if not should_continue:
+                        break
+                except asyncio.CancelledError:
+                    print("\n⚠️  Agent loop cancelled")
                     break
-            except asyncio.CancelledError:
-                break
-            except Exception as e:
-                print(f"❌ Error in agent loop: {e}")
-                await session.send_event(
-                    Event(event_type="error", data={"error": str(e)})
-                )
+                except Exception as e:
+                    print(f"❌ Error in agent loop: {e}")
+                    await session.send_event(
+                        Event(event_type="error", data={"error": str(e)})
+                    )
 
-    print("🛑 Agent loop exited")
+        print("🛑 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...")
+            try:
+                local_path = session.save_and_upload_detached(
+                    session.config.session_dataset_repo
+                )
+                if local_path:
+                    print("✅ Emergency save successful, upload in progress")
+            except Exception as e:
+                print(f"❌ Emergency save failed: {e}")

agent/core/session.py

@@ -1,9 +1,12 @@
 import asyncio
 import json
+import subprocess
+import sys
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
+from pathlib import Path
 from typing import Any, Optional
 
 from litellm import get_max_tokens
@@ -60,6 +63,8 @@ class Session:
         # Session trajectory logging
         self.logged_events: list[dict] = []
         self.session_start_time = datetime.now().isoformat()
+        self.turn_count: int = 0
+        self.last_auto_save_turn: int = 0
 
     async def send_event(self, event: Event) -> None:
         """Send event back to client and log to trajectory"""
@@ -79,6 +84,26 @@ class Session:
         if self.current_task and not self.current_task.done():
             self.current_task.cancel()
 
+    def increment_turn(self) -> None:
+        """Increment turn counter (called after each user interaction)"""
+        self.turn_count += 1
+
+    async def auto_save_if_needed(self) -> None:
+        """Check if auto-save should trigger and save if so (completely non-blocking)"""
+        if not self.config.save_sessions:
+            return
+
+        interval = self.config.auto_save_interval
+        if interval <= 0:
+            return
+
+        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})...")
+            # Fire-and-forget save - returns immediately
+            self.save_and_upload_detached(self.config.session_dataset_repo)
+            self.last_auto_save_turn = self.turn_count
+
     def get_trajectory(self) -> dict:
         """Serialize complete session trajectory for logging"""
         return {
@@ -90,51 +115,121 @@ class Session:
             "events": self.logged_events,
         }
 
-    async def push_to_dataset(self, repo_id: str) -> Optional[str]:
+    def save_trajectory_local(
+        self,
+        directory: str = "session_logs",
+        upload_status: str = "pending",
+        dataset_url: Optional[str] = None,
+    ) -> Optional[str]:
         """
-        Push session trajectory to Hugging Face dataset
+        Save trajectory to local JSON file as backup with upload status
 
         Args:
-            repo_id: HuggingFace dataset repo ID (e.g. 'username/dataset-name')
+            directory: Directory to save logs (default: "session_logs")
+            upload_status: Status of upload attempt ("pending", "success", "failed")
+            dataset_url: URL of dataset if upload succeeded
 
         Returns:
-            URL to the uploaded file if successful, None otherwise
+            Path to saved file if successful, None otherwise
         """
         try:
-            import os
-
-            from datasets import Dataset
+            log_dir = Path(directory)
+            log_dir.mkdir(parents=True, exist_ok=True)
 
-            # Get trajectory data
             trajectory = self.get_trajectory()
 
-            # Convert to dataset row format
-            row = {
-                "session_id": trajectory["session_id"],
-                "session_start_time": trajectory["session_start_time"],
-                "session_end_time": trajectory["session_end_time"],
-                "model_name": trajectory["model_name"],
-                "messages": json.dumps(trajectory["messages"]),
-                "events": json.dumps(trajectory["events"]),
-            }
+            # Add upload metadata
+            trajectory["upload_status"] = upload_status
+            trajectory["upload_url"] = dataset_url
+            trajectory["last_save_time"] = datetime.now().isoformat()
+
+            filename = f"session_{self.session_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
+            filepath = log_dir / filename
 
-            # Try to load existing dataset and append
-            try:
-                from datasets import load_dataset
+            with open(filepath, "w") as f:
+                json.dump(trajectory, f, indent=2)
 
-                existing_dataset = load_dataset(repo_id, split="train")
-                new_dataset = Dataset.from_dict(
-                    {k: list(existing_dataset[k]) + [v] for k, v in row.items()}
-                )
-            except Exception:
-                # Dataset doesn't exist yet, create new one
-                new_dataset = Dataset.from_dict({k: [v] for k, v in row.items()})
+            return str(filepath)
+        except Exception as e:
+            print(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)
 
-            # Push to hub
-            new_dataset.push_to_hub(repo_id, private=True, token=os.getenv("HF_TOKEN"))
+            data["upload_status"] = upload_status
+            data["upload_url"] = dataset_url
+            data["last_save_time"] = datetime.now().isoformat()
 
-            return f"https://huggingface.co/datasets/{repo_id}"
+            with open(filepath, "w") as f:
+                json.dump(data, f, indent=2)
 
+            return True
         except Exception as e:
-            print(f"Failed to push session to dataset: {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)
+
+        Args:
+            repo_id: HuggingFace dataset repo ID
+
+        Returns:
+            Path to local save file
+        """
+        # Save locally first (fast, synchronous)
+        local_path = self.save_trajectory_local(upload_status="pending")
+        if not local_path:
             return None
+
+        # Spawn detached subprocess for upload (fire-and-forget)
+        try:
+            uploader_script = Path(__file__).parent / "session_uploader.py"
+
+            # Use Popen with detached process
+            subprocess.Popen(
+                [sys.executable, str(uploader_script), "upload", local_path, repo_id],
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                start_new_session=True,  # Detach from parent
+            )
+        except Exception as e:
+            print(f"⚠️  Failed to spawn upload subprocess: {e}")
+
+        return local_path
+
+    @staticmethod
+    def retry_failed_uploads_detached(
+        directory: str = "session_logs", repo_id: Optional[str] = None
+    ) -> None:
+        """
+        Spawn detached subprocess to retry failed/pending uploads (fire-and-forget)
+
+        Args:
+            directory: Directory containing session logs
+            repo_id: Target dataset repo ID
+        """
+        if not repo_id:
+            return
+
+        try:
+            uploader_script = Path(__file__).parent / "session_uploader.py"
+
+            # Spawn detached subprocess for retry
+            subprocess.Popen(
+                [sys.executable, str(uploader_script), "retry", directory, repo_id],
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                start_new_session=True,  # Detach from parent
+            )
+        except Exception as e:
+            print(f"⚠️  Failed to spawn retry subprocess: {e}")

agent/core/session_uploader.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+"""
+Standalone script for uploading session trajectories to HuggingFace.
+This runs as a separate process to avoid blocking the main agent.
+"""
+
+import json
+import os
+import sys
+from pathlib import Path
+
+
+def upload_session_to_dataset(session_file: str, repo_id: str, max_retries: int = 3):
+    """Upload a single session file to HuggingFace dataset"""
+    try:
+        from datasets import Dataset, load_dataset
+    except ImportError:
+        print("Error: datasets library not available", file=sys.stderr)
+        return False
+
+    try:
+        # Load session data
+        with open(session_file, "r") as f:
+            data = json.load(f)
+
+        # Check if already uploaded
+        upload_status = data.get("upload_status")
+        if upload_status == "success":
+            return True
+
+        # Prepare row for upload
+        row = {
+            "session_id": data["session_id"],
+            "session_start_time": data["session_start_time"],
+            "session_end_time": data["session_end_time"],
+            "model_name": data["model_name"],
+            "messages": json.dumps(data["messages"]),
+            "events": json.dumps(data["events"]),
+        }
+
+        hf_token = os.getenv("HF_TOKEN")
+        if not hf_token:
+            # Update status to failed
+            data["upload_status"] = "failed"
+            with open(session_file, "w") as f:
+                json.dump(data, f, indent=2)
+            return False
+
+        # Try to load existing dataset and append
+        for attempt in range(max_retries):
+            try:
+                try:
+                    existing_dataset = load_dataset(repo_id, split="train")
+                    new_dataset = Dataset.from_dict(
+                        {k: list(existing_dataset[k]) + [v] for k, v in row.items()}
+                    )
+                except Exception:
+                    # Dataset doesn't exist yet, create new one
+                    new_dataset = Dataset.from_dict({k: [v] for k, v in row.items()})
+
+                # Push to hub
+                new_dataset.push_to_hub(repo_id, private=True, token=hf_token)
+
+                # Update status to success
+                data["upload_status"] = "success"
+                data["upload_url"] = f"https://huggingface.co/datasets/{repo_id}"
+                with open(session_file, "w") as f:
+                    json.dump(data, f, indent=2)
+
+                return True
+
+            except Exception:
+                if attempt < max_retries - 1:
+                    import time
+
+                    wait_time = 2**attempt
+                    time.sleep(wait_time)
+                else:
+                    # Final attempt failed
+                    data["upload_status"] = "failed"
+                    with open(session_file, "w") as f:
+                        json.dump(data, f, indent=2)
+                    return False
+
+    except Exception as e:
+        print(f"Error uploading session: {e}", file=sys.stderr)
+        return False
+
+
+def retry_failed_uploads(directory: str, repo_id: str):
+    """Retry all failed/pending uploads in a directory"""
+    log_dir = Path(directory)
+    if not log_dir.exists():
+        return
+
+    session_files = list(log_dir.glob("session_*.json"))
+
+    for filepath in session_files:
+        try:
+            with open(filepath, "r") as f:
+                data = json.load(f)
+
+            upload_status = data.get("upload_status", "unknown")
+
+            # Only retry pending or failed uploads
+            if upload_status in ["pending", "failed"]:
+                upload_session_to_dataset(str(filepath), repo_id)
+
+        except Exception:
+            pass
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3:
+        print("Usage: session_uploader.py <command> <args...>")
+        sys.exit(1)
+
+    command = sys.argv[1]
+
+    if command == "upload":
+        # python session_uploader.py upload <session_file> <repo_id>
+        if len(sys.argv) < 4:
+            print("Usage: session_uploader.py upload <session_file> <repo_id>")
+            sys.exit(1)
+        session_file = sys.argv[2]
+        repo_id = sys.argv[3]
+        success = upload_session_to_dataset(session_file, repo_id)
+        sys.exit(0 if success else 1)
+
+    elif command == "retry":
+        # python session_uploader.py retry <directory> <repo_id>
+        if len(sys.argv) < 4:
+            print("Usage: session_uploader.py retry <directory> <repo_id>")
+            sys.exit(1)
+        directory = sys.argv[2]
+        repo_id = sys.argv[3]
+        retry_failed_uploads(directory, repo_id)
+        sys.exit(0)
+
+    else:
+        print(f"Unknown command: {command}")
+        sys.exit(1)

agent/main.py

@@ -415,8 +415,7 @@ async def main():
     )
     await submission_queue.put(shutdown_submission)
 
-    # Wait for tasks to complete (longer timeout to allow for session save)
-    await asyncio.wait_for(agent_task, timeout=30.0)
+    await asyncio.wait_for(agent_task, timeout=5.0)
     listener_task.cancel()
 
     print("✨ Goodbye!\n")