lol yolo merge

Copying mcp tools

.gitignore

@@ -13,4 +13,6 @@ wheels/
 .claude/
 *.jsonl
 *.csv
-/logs
+/logs
+hf-agent-leaderboard/
+.cursor/

agent/core/tools.py

@@ -3,7 +3,6 @@ Tool system for the agent
 Provides ToolSpec and ToolRouter for managing both built-in and MCP tools
 """
 
-import subprocess
 import warnings
 from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Optional
@@ -13,12 +12,15 @@ from lmnr import observe
 from mcp.types import EmbeddedResource, ImageContent, TextContent
 
 from agent.config import MCPServerConfig
+from agent.tools.jobs_tool import HF_JOBS_TOOL_SPEC, hf_jobs_handler
 
 # Suppress aiohttp deprecation warning
 warnings.filterwarnings(
     "ignore", category=DeprecationWarning, module="aiohttp.connector"
 )
 
+NOT_ALLOWED_TOOL_NAMES = ["hf_jobs"]
+
 
 def convert_mcp_content_to_string(content: list) -> str:
     """
@@ -104,6 +106,10 @@ class ToolRouter:
     async def register_mcp_tools(self) -> None:
         tools = await self.mcp_client.list_tools()
         for tool in tools:
+            if tool.name in NOT_ALLOWED_TOOL_NAMES:
+                print(f"Skipping not MCP allowed tool: {tool.name}")
+                continue
+            print(f"MCP Tool: {tool.name}")
             self.register_tool(
                 ToolSpec(
                     name=tool.name,
@@ -173,93 +179,14 @@ class ToolRouter:
 # ============================================================================
 
 
-async def bash_handler(arguments: dict[str, Any]) -> tuple[str, bool]:
-    """Execute bash command"""
-    try:
-        command = arguments.get("command", "")
-        result = subprocess.run(
-            command, shell=True, capture_output=True, text=True, timeout=30
-        )
-        output = result.stdout if result.returncode == 0 else result.stderr
-        success = result.returncode == 0
-        return output, success
-    except Exception as e:
-        return f"Error: {str(e)}", False
-
-
-async def read_file_handler(arguments: dict[str, Any]) -> tuple[str, bool]:
-    """Read file contents"""
-    try:
-        path = arguments.get("path", "")
-        with open(path, "r") as f:
-            content = f.read()
-        return content, True
-    except Exception as e:
-        return f"Error reading file: {str(e)}", False
-
-
-async def write_file_handler(arguments: dict[str, Any]) -> tuple[str, bool]:
-    """Write to file"""
-    try:
-        path = arguments.get("path", "")
-        content = arguments.get("content", "")
-        with open(path, "w") as f:
-            f.write(content)
-        return f"Successfully wrote to {path}", True
-    except Exception as e:
-        return f"Error writing file: {str(e)}", False
-
-
 def create_builtin_tools() -> list[ToolSpec]:
     """Create built-in tool specifications"""
+    print(f"Creating built-in tools: {HF_JOBS_TOOL_SPEC['name']}")
     return [
         ToolSpec(
-            name="bash",
-            description="Execute a bash command and return its output",
-            parameters={
-                "type": "object",
-                "properties": {
-                    "command": {
-                        "type": "string",
-                        "description": "The bash command to execute",
-                    }
-                },
-                "required": ["command"],
-            },
-            handler=bash_handler,
-        ),
-        ToolSpec(
-            name="read_file",
-            description="Read the contents of a file",
-            parameters={
-                "type": "object",
-                "properties": {
-                    "path": {
-                        "type": "string",
-                        "description": "Path to the file to read",
-                    }
-                },
-                "required": ["path"],
-            },
-            handler=read_file_handler,
-        ),
-        ToolSpec(
-            name="write_file",
-            description="Write content to a file",
-            parameters={
-                "type": "object",
-                "properties": {
-                    "path": {
-                        "type": "string",
-                        "description": "Path to the file to write",
-                    },
-                    "content": {
-                        "type": "string",
-                        "description": "Content to write to the file",
-                    },
-                },
-                "required": ["path", "content"],
-            },
-            handler=write_file_handler,
+            name=HF_JOBS_TOOL_SPEC["name"],
+            description=HF_JOBS_TOOL_SPEC["description"],
+            parameters=HF_JOBS_TOOL_SPEC["parameters"],
+            handler=hf_jobs_handler,
         ),
     ]

agent/main.py

@@ -3,6 +3,7 @@ Interactive CLI chat with the agent
 """
 
 import asyncio
+import json
 import os
 from dataclasses import dataclass
 from pathlib import Path
@@ -64,8 +65,11 @@ async def event_listener(
                     print(f"\n🤖 Assistant: {content}")
             elif event.event_type == "tool_call":
                 tool_name = event.data.get("tool", "") if event.data else ""
+                arguments = event.data.get("arguments", {}) if event.data else {}
                 if tool_name:
-                    print(f"🔧 Calling tool: {tool_name}")
+                    print(
+                        f"🔧 Calling tool: {tool_name} with arguments: {json.dumps(arguments)[:100]}..."
+                    )
             elif event.event_type == "tool_output":
                 output = event.data.get("output", "") if event.data else ""
                 success = event.data.get("success", False) if event.data else False

agent/tests/__init__.py

@@ -1,3 +0,0 @@
-"""
-Test suite for HF Agent
-"""

agent/tests/unit/test_base.py

@@ -1,10 +0,0 @@
-"""
-Unit tests for base agent components
-"""
-
-
-def test_base_agent_initialization():
-    """Test that BaseAgent can be initialized with config"""
-    # This will fail because BaseAgent is abstract
-    # Subclasses should implement this
-    pass

agent/tools/__init__.py

@@ -0,0 +1,8 @@
+"""
+Hugging Face tools for the agent
+"""
+
+from agent.tools.jobs_tool import HF_JOBS_TOOL_SPEC, HfJobsTool, hf_jobs_handler
+from agent.tools.types import ToolResult
+
+__all__ = ["ToolResult", "HF_JOBS_TOOL_SPEC", "hf_jobs_handler", HfJobsTool]

agent/tools/jobs_tool.py

@@ -0,0 +1,836 @@
+"""
+Hugging Face Jobs Tool - Using huggingface-hub library
+
+Refactored to use official huggingface-hub library instead of custom HTTP client
+"""
+
+import asyncio
+import base64
+from typing import Any, Dict, Literal, Optional
+
+from huggingface_hub import HfApi
+from huggingface_hub.utils import HfHubHTTPError
+
+from agent.tools.types import ToolResult
+from agent.tools.utilities import (
+    format_job_details,
+    format_jobs_table,
+    format_scheduled_job_details,
+    format_scheduled_jobs_table,
+)
+
+# Hardware flavors
+CPU_FLAVORS = ["cpu-basic", "cpu-upgrade", "cpu-performance", "cpu-xl"]
+GPU_FLAVORS = [
+    "sprx8",
+    "zero-a10g",
+    "t4-small",
+    "t4-medium",
+    "l4x1",
+    "l4x4",
+    "l40sx1",
+    "l40sx4",
+    "l40sx8",
+    "a10g-small",
+    "a10g-large",
+    "a10g-largex2",
+    "a10g-largex4",
+    "a100-large",
+    "h100",
+    "h100x8",
+]
+SPECIALIZED_FLAVORS = ["inf2x6"]
+ALL_FLAVORS = CPU_FLAVORS + GPU_FLAVORS + SPECIALIZED_FLAVORS
+
+# Operation names
+OperationType = Literal[
+    "run",
+    "uv",
+    "ps",
+    "logs",
+    "inspect",
+    "cancel",
+    "scheduled run",
+    "scheduled uv",
+    "scheduled ps",
+    "scheduled inspect",
+    "scheduled delete",
+    "scheduled suspend",
+    "scheduled resume",
+]
+
+# Constants
+DEFAULT_LOG_WAIT_SECONDS = 10
+UV_DEFAULT_IMAGE = "ghcr.io/astral-sh/uv:python3.12-bookworm"
+
+
+def _build_uv_command(
+    script: str,
+    with_deps: list[str] | None = None,
+    python: str | None = None,
+    script_args: list[str] | None = None,
+) -> list[str]:
+    """Build UV run command"""
+    parts = ["uv", "run"]
+
+    if with_deps:
+        for dep in with_deps:
+            parts.extend(["--with", dep])
+
+    if python:
+        parts.extend(["-p", python])
+
+    parts.append(script)
+
+    if script_args:
+        parts.extend(script_args)
+
+    return parts
+
+
+def _wrap_inline_script(
+    script: str,
+    with_deps: list[str] | None = None,
+    python: str | None = None,
+    script_args: list[str] | None = None,
+) -> str:
+    """Wrap inline script with base64 encoding to avoid file creation"""
+    encoded = base64.b64encode(script.encode("utf-8")).decode("utf-8")
+    # Build the uv command with stdin (-)
+    uv_command = _build_uv_command("-", with_deps, python, script_args)
+    # Join command parts with proper spacing
+    uv_command_str = " ".join(uv_command)
+    return f'echo "{encoded}" | base64 -d | {uv_command_str}'
+
+
+def _resolve_uv_command(
+    script: str,
+    with_deps: list[str] | None = None,
+    python: str | None = None,
+    script_args: list[str] | None = None,
+) -> list[str]:
+    """Resolve UV command based on script source (URL, inline, or file path)"""
+    # If URL, use directly
+    if script.startswith("http://") or script.startswith("https://"):
+        return _build_uv_command(script, with_deps, python, script_args)
+
+    # If contains newline, treat as inline script
+    if "\n" in script:
+        wrapped = _wrap_inline_script(script, with_deps, python, script_args)
+        return ["/bin/sh", "-lc", wrapped]
+
+    # Otherwise, treat as file path
+    return _build_uv_command(script, with_deps, python, script_args)
+
+
+async def _async_call(func, *args, **kwargs):
+    """Wrap synchronous HfApi calls for async context"""
+    return await asyncio.to_thread(func, *args, **kwargs)
+
+
+def _job_info_to_dict(job_info) -> Dict[str, Any]:
+    """Convert JobInfo object to dictionary for formatting functions"""
+    return {
+        "id": job_info.id,
+        "status": {"stage": job_info.status.stage, "message": job_info.status.message},
+        "command": job_info.command,
+        "createdAt": job_info.created_at.isoformat(),
+        "dockerImage": job_info.docker_image,
+        "spaceId": job_info.space_id,
+        "flavor": job_info.flavor,
+        "owner": {"name": job_info.owner.name},
+    }
+
+
+def _scheduled_job_info_to_dict(scheduled_job_info) -> Dict[str, Any]:
+    """Convert ScheduledJobInfo object to dictionary for formatting functions"""
+    job_spec = scheduled_job_info.job_spec
+
+    # Extract last run and next run from status
+    last_run = None
+    next_run = None
+    if scheduled_job_info.status:
+        if scheduled_job_info.status.last_job:
+            last_run = scheduled_job_info.status.last_job.created_at
+            if last_run:
+                last_run = (
+                    last_run.isoformat()
+                    if hasattr(last_run, "isoformat")
+                    else str(last_run)
+                )
+        if scheduled_job_info.status.next_job_run_at:
+            next_run = scheduled_job_info.status.next_job_run_at
+            next_run = (
+                next_run.isoformat()
+                if hasattr(next_run, "isoformat")
+                else str(next_run)
+            )
+
+    return {
+        "id": scheduled_job_info.id,
+        "schedule": scheduled_job_info.schedule,
+        "suspend": scheduled_job_info.suspend,
+        "lastRun": last_run,
+        "nextRun": next_run,
+        "jobSpec": {
+            "dockerImage": job_spec.docker_image,
+            "spaceId": job_spec.space_id,
+            "command": job_spec.command or [],
+            "flavor": job_spec.flavor or "cpu-basic",
+        },
+    }
+
+
+class HfJobsTool:
+    """Tool for managing Hugging Face compute jobs using huggingface-hub library"""
+
+    def __init__(self, hf_token: Optional[str] = None, namespace: Optional[str] = None):
+        self.api = HfApi(token=hf_token)
+        self.namespace = namespace
+
+    async def execute(self, params: Dict[str, Any]) -> ToolResult:
+        """Execute the specified operation"""
+        operation = params.get("operation")
+        args = params.get("args", {})
+
+        # If no operation provided, return usage instructions
+        if not operation:
+            return self._show_help()
+
+        # Normalize operation name
+        operation = operation.lower()
+
+        # Check if help is requested
+        if args.get("help"):
+            return self._show_operation_help(operation)
+
+        try:
+            # Route to appropriate handler
+            if operation == "run":
+                return await self._run_job(args)
+            elif operation == "uv":
+                return await self._run_uv_job(args)
+            elif operation == "ps":
+                return await self._list_jobs(args)
+            elif operation == "logs":
+                return await self._get_logs(args)
+            elif operation == "inspect":
+                return await self._inspect_job(args)
+            elif operation == "cancel":
+                return await self._cancel_job(args)
+            elif operation == "scheduled run":
+                return await self._scheduled_run(args)
+            elif operation == "scheduled uv":
+                return await self._scheduled_uv(args)
+            elif operation == "scheduled ps":
+                return await self._list_scheduled_jobs(args)
+            elif operation == "scheduled inspect":
+                return await self._inspect_scheduled_job(args)
+            elif operation == "scheduled delete":
+                return await self._delete_scheduled_job(args)
+            elif operation == "scheduled suspend":
+                return await self._suspend_scheduled_job(args)
+            elif operation == "scheduled resume":
+                return await self._resume_scheduled_job(args)
+            else:
+                return {
+                    "formatted": f'Unknown operation: "{operation}"\n\n'
+                    "Available operations:\n"
+                    "- run, uv, ps, logs, inspect, cancel\n"
+                    "- scheduled run, scheduled uv, scheduled ps, scheduled inspect, "
+                    "scheduled delete, scheduled suspend, scheduled resume\n\n"
+                    "Call this tool with no operation for full usage instructions.",
+                    "totalResults": 0,
+                    "resultsShared": 0,
+                    "isError": True,
+                }
+
+        except HfHubHTTPError as e:
+            return {
+                "formatted": f"API Error: {str(e)}",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+        except Exception as e:
+            return {
+                "formatted": f"Error executing {operation}: {str(e)}",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+
+    def _show_help(self) -> ToolResult:
+        """Show usage instructions when tool is called with no arguments"""
+        cpu_flavors_list = ", ".join(CPU_FLAVORS)
+        gpu_flavors_list = ", ".join(GPU_FLAVORS)
+        specialized_flavors_list = ", ".join(SPECIALIZED_FLAVORS)
+
+        hardware_section = f"**CPU:** {cpu_flavors_list}\n"
+        if GPU_FLAVORS:
+            hardware_section += f"**GPU:** {gpu_flavors_list}\n"
+        if SPECIALIZED_FLAVORS:
+            hardware_section += f"**Specialized:** {specialized_flavors_list}"
+
+        usage_text = f"""# HuggingFace Jobs API
+
+Manage compute jobs on Hugging Face infrastructure.
+
+## Available Commands
+
+### Job Management
+- **run** - Run a job with a Docker image
+- **uv** - Run a Python script with UV (inline dependencies)
+- **ps** - List jobs
+- **logs** - Fetch job logs
+- **inspect** - Get detailed job information
+- **cancel** - Cancel a running job
+
+### Scheduled Jobs
+- **scheduled run** - Create a scheduled job
+- **scheduled uv** - Create a scheduled UV job
+- **scheduled ps** - List scheduled jobs
+- **scheduled inspect** - Get scheduled job details
+- **scheduled delete** - Delete a scheduled job
+- **scheduled suspend** - Pause a scheduled job
+- **scheduled resume** - Resume a suspended job
+
+## Examples
+
+### Run a simple job
+Call this tool with:
+```json
+{{
+  "operation": "run",
+  "args": {{
+    "image": "python:3.12",
+    "command": ["python", "-c", "print('Hello from HF Jobs!')"],
+    "flavor": "cpu-basic"
+  }}
+}}
+```
+
+### Run a Python script with UV
+Call this tool with:
+```json
+{{
+  "operation": "uv",
+  "args": {{
+    "script": "import random\\nprint(42 + random.randint(1, 5))"
+  }}
+}}
+```
+
+## Hardware Flavors
+
+{hardware_section}
+
+## Command Format Guidelines
+
+**Array format (default):**
+- Recommended for every command—JSON keeps arguments intact (URLs with `&`, spaces, etc.)
+- Use `["/bin/sh", "-lc", "..."]` when you need shell operators like `&&`, `|`, or redirections
+- Works with any language: Python, bash, node, npm, uv, etc.
+
+**String format (simple cases only):**
+- Still accepted for backwards compatibility, parsed with POSIX shell semantics
+- Rejects shell operators and can mis-handle characters such as `&`; switch to arrays when things turn complex
+- `$HF_TOKEN` stays literal—forward it via `secrets: {{ "HF_TOKEN": "$HF_TOKEN" }}`
+
+### Show command-specific help
+Call this tool with:
+```json
+{{"operation": "<operation>", "args": {{"help": true}}}}
+```
+
+## Tips
+
+- Jobs default to non-detached mode (tail logs for up to {DEFAULT_LOG_WAIT_SECONDS}s or until completion). Set `detach: true` to return immediately.
+- Prefer array commands to avoid shell parsing surprises
+- To access private Hub assets, include `secrets: {{ "HF_TOKEN": "$HF_TOKEN" }}` to inject your auth token.
+"""
+        return {"formatted": usage_text, "totalResults": 1, "resultsShared": 1}
+
+    def _show_operation_help(self, operation: str) -> ToolResult:
+        """Show help for a specific operation"""
+        help_text = f"Help for operation: {operation}\n\nCall with appropriate arguments. Use the main help for examples."
+        return {"formatted": help_text, "totalResults": 1, "resultsShared": 1}
+
+    async def _run_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Run a job using HfApi.run_job()"""
+        try:
+            job = await _async_call(
+                self.api.run_job,
+                image=args.get("image", "python:3.12"),
+                command=args.get("command"),
+                env=args.get("env"),
+                secrets=args.get("secrets"),
+                flavor=args.get("flavor", "cpu-basic"),
+                timeout=args.get("timeout", "30m"),
+                namespace=args.get("namespace") or self.namespace,
+            )
+
+            # If detached, return immediately
+            if args.get("detach", False):
+                response = f"""Job started successfully!
+
+**Job ID:** {job.id}
+**Status:** {job.status.stage}
+**View at:** {job.url}
+
+To check logs, call this tool with `{{"operation": "logs", "args": {{"job_id": "{job.id}"}}}}`
+To inspect, call this tool with `{{"operation": "inspect", "args": {{"job_id": "{job.id}"}}}}`"""
+                return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+            # Not detached - return job info
+            response = f"""Job started: {job.id}
+
+**Status:** {job.status.stage}
+**View logs at:** {job.url}
+
+Note: Logs are being collected. Check the job page for real-time logs.
+"""
+            return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+        except Exception as e:
+            raise Exception(f"Failed to run job: {str(e)}")
+
+    async def _run_uv_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Run UV job with inline script support (no local files needed)"""
+        try:
+            script = args.get("script")
+            if not script:
+                raise ValueError("script is required")
+
+            # Resolve the command based on script type (URL, inline, or file)
+            command = _resolve_uv_command(
+                script=script,
+                with_deps=args.get("with_deps") or args.get("dependencies"),
+                python=args.get("python"),
+                script_args=args.get("script_args"),
+            )
+
+            # Use run_job with UV image instead of run_uv_job
+            job = await _async_call(
+                self.api.run_job,
+                image=UV_DEFAULT_IMAGE,
+                command=command,
+                env=args.get("env"),
+                secrets=args.get("secrets"),
+                flavor=args.get("flavor", "cpu-basic"),
+                timeout=args.get("timeout", "30m"),
+                namespace=args.get("namespace") or self.namespace,
+            )
+
+            response = f"""UV Job started: {job.id}
+
+**Status:** {job.status.stage}
+**View at:** {job.url}
+
+To check logs, call this tool with `{{"operation": "logs", "args": {{"job_id": "{job.id}"}}}}`
+"""
+            return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+        except Exception as e:
+            raise Exception(f"Failed to run UV job: {str(e)}")
+
+    async def _list_jobs(self, args: Dict[str, Any]) -> ToolResult:
+        """List jobs using HfApi.list_jobs()"""
+        jobs_list = await _async_call(
+            self.api.list_jobs, namespace=args.get("namespace") or self.namespace
+        )
+
+        # Filter jobs
+        if not args.get("all", False):
+            jobs_list = [j for j in jobs_list if j.status.stage == "RUNNING"]
+
+        if args.get("status"):
+            status_filter = args["status"].upper()
+            jobs_list = [j for j in jobs_list if status_filter in j.status.stage]
+
+        # Convert JobInfo objects to dicts for formatting
+        jobs_dicts = [_job_info_to_dict(j) for j in jobs_list]
+
+        table = format_jobs_table(jobs_dicts)
+
+        if len(jobs_list) == 0:
+            if args.get("all", False):
+                return {
+                    "formatted": "No jobs found.",
+                    "totalResults": 0,
+                    "resultsShared": 0,
+                }
+            return {
+                "formatted": 'No running jobs found. Use `{"args": {"all": true}}` to show all jobs.',
+                "totalResults": 0,
+                "resultsShared": 0,
+            }
+
+        response = f"**Jobs ({len(jobs_list)} total):**\n\n{table}"
+        return {
+            "formatted": response,
+            "totalResults": len(jobs_list),
+            "resultsShared": len(jobs_list),
+        }
+
+    async def _get_logs(self, args: Dict[str, Any]) -> ToolResult:
+        """Fetch logs using HfApi.fetch_job_logs()"""
+        job_id = args.get("job_id")
+        if not job_id:
+            return {
+                "formatted": "job_id is required",
+                "isError": True,
+                "totalResults": 0,
+                "resultsShared": 0,
+            }
+
+        try:
+            # Fetch logs (returns generator, convert to list)
+            logs_gen = self.api.fetch_job_logs(
+                job_id=job_id, namespace=args.get("namespace") or self.namespace
+            )
+            logs = await _async_call(list, logs_gen)
+
+            if not logs:
+                return {
+                    "formatted": f"No logs available for job {job_id}",
+                    "totalResults": 0,
+                    "resultsShared": 0,
+                }
+
+            log_text = "\n".join(logs)
+            return {
+                "formatted": f"**Logs for {job_id}:**\n\n```\n{log_text}\n```",
+                "totalResults": 1,
+                "resultsShared": 1,
+            }
+
+        except Exception as e:
+            return {
+                "formatted": f"Failed to fetch logs: {str(e)}",
+                "isError": True,
+                "totalResults": 0,
+                "resultsShared": 0,
+            }
+
+    async def _inspect_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Inspect job using HfApi.inspect_job()"""
+        job_id = args.get("job_id")
+        if not job_id:
+            return {
+                "formatted": "job_id is required",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+
+        job_ids = job_id if isinstance(job_id, list) else [job_id]
+
+        jobs = []
+        for jid in job_ids:
+            try:
+                job = await _async_call(
+                    self.api.inspect_job,
+                    job_id=jid,
+                    namespace=args.get("namespace") or self.namespace,
+                )
+                jobs.append(_job_info_to_dict(job))
+            except Exception as e:
+                raise Exception(f"Failed to inspect job {jid}: {str(e)}")
+
+        formatted_details = format_job_details(jobs)
+        response = f"**Job Details** ({len(jobs)} job{'s' if len(jobs) > 1 else ''}):\n\n{formatted_details}"
+
+        return {
+            "formatted": response,
+            "totalResults": len(jobs),
+            "resultsShared": len(jobs),
+        }
+
+    async def _cancel_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Cancel job using HfApi.cancel_job()"""
+        job_id = args.get("job_id")
+        if not job_id:
+            return {
+                "formatted": "job_id is required",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+
+        await _async_call(
+            self.api.cancel_job,
+            job_id=job_id,
+            namespace=args.get("namespace") or self.namespace,
+        )
+
+        response = f"""✓ Job {job_id} has been cancelled.
+
+To verify, call this tool with `{{"operation": "inspect", "args": {{"job_id": "{job_id}"}}}}`"""
+
+        return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+    async def _scheduled_run(self, args: Dict[str, Any]) -> ToolResult:
+        """Create scheduled job using HfApi.create_scheduled_job()"""
+        try:
+            scheduled_job = await _async_call(
+                self.api.create_scheduled_job,
+                image=args.get("image", "python:3.12"),
+                command=args.get("command"),
+                schedule=args.get("schedule"),
+                env=args.get("env"),
+                secrets=args.get("secrets"),
+                flavor=args.get("flavor", "cpu-basic"),
+                timeout=args.get("timeout", "30m"),
+                namespace=args.get("namespace") or self.namespace,
+            )
+
+            scheduled_dict = _scheduled_job_info_to_dict(scheduled_job)
+
+            response = f"""✓ Scheduled job created successfully!
+
+**Scheduled Job ID:** {scheduled_dict["id"]}
+**Schedule:** {scheduled_dict["schedule"]}
+**Suspended:** {"Yes" if scheduled_dict.get("suspend") else "No"}
+**Next Run:** {scheduled_dict.get("nextRun", "N/A")}
+
+To inspect, call this tool with `{{"operation": "scheduled inspect", "args": {{"scheduled_job_id": "{scheduled_dict["id"]}"}}}}`
+To list all, call this tool with `{{"operation": "scheduled ps"}}`"""
+
+            return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+        except Exception as e:
+            raise Exception(f"Failed to create scheduled job: {str(e)}")
+
+    async def _scheduled_uv(self, args: Dict[str, Any]) -> ToolResult:
+        """Create scheduled UV job with inline script support"""
+        try:
+            script = args.get("script")
+            if not script:
+                raise ValueError("script is required")
+
+            schedule = args.get("schedule")
+            if not schedule:
+                raise ValueError("schedule is required")
+
+            # Resolve the command based on script type
+            command = _resolve_uv_command(
+                script=script,
+                with_deps=args.get("with_deps") or args.get("dependencies"),
+                python=args.get("python"),
+                script_args=args.get("script_args"),
+            )
+
+            # Use create_scheduled_job with UV image
+            scheduled_job = await _async_call(
+                self.api.create_scheduled_job,
+                image=UV_DEFAULT_IMAGE,
+                command=command,
+                schedule=schedule,
+                env=args.get("env"),
+                secrets=args.get("secrets"),
+                flavor=args.get("flavor", "cpu-basic"),
+                timeout=args.get("timeout", "30m"),
+                namespace=args.get("namespace") or self.namespace,
+            )
+
+            scheduled_dict = _scheduled_job_info_to_dict(scheduled_job)
+
+            response = f"""✓ Scheduled UV job created successfully!
+
+**Scheduled Job ID:** {scheduled_dict["id"]}
+**Schedule:** {scheduled_dict["schedule"]}
+**Suspended:** {"Yes" if scheduled_dict.get("suspend") else "No"}
+**Next Run:** {scheduled_dict.get("nextRun", "N/A")}
+
+To inspect, call this tool with `{{"operation": "scheduled inspect", "args": {{"scheduled_job_id": "{scheduled_dict["id"]}"}}}}`"""
+
+            return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+        except Exception as e:
+            raise Exception(f"Failed to create scheduled UV job: {str(e)}")
+
+    async def _list_scheduled_jobs(self, args: Dict[str, Any]) -> ToolResult:
+        """List scheduled jobs using HfApi.list_scheduled_jobs()"""
+        scheduled_jobs_list = await _async_call(
+            self.api.list_scheduled_jobs,
+            namespace=args.get("namespace") or self.namespace,
+        )
+
+        # Filter jobs - default: hide suspended jobs unless --all is specified
+        if not args.get("all", False):
+            scheduled_jobs_list = [j for j in scheduled_jobs_list if not j.suspend]
+
+        # Convert to dicts for formatting
+        scheduled_dicts = [_scheduled_job_info_to_dict(j) for j in scheduled_jobs_list]
+
+        table = format_scheduled_jobs_table(scheduled_dicts)
+
+        if len(scheduled_jobs_list) == 0:
+            if args.get("all", False):
+                return {
+                    "formatted": "No scheduled jobs found.",
+                    "totalResults": 0,
+                    "resultsShared": 0,
+                }
+            return {
+                "formatted": 'No active scheduled jobs found. Use `{"args": {"all": true}}` to show suspended jobs.',
+                "totalResults": 0,
+                "resultsShared": 0,
+            }
+
+        response = f"**Scheduled Jobs ({len(scheduled_jobs_list)} total):**\n\n{table}"
+        return {
+            "formatted": response,
+            "totalResults": len(scheduled_jobs_list),
+            "resultsShared": len(scheduled_jobs_list),
+        }
+
+    async def _inspect_scheduled_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Inspect scheduled job using HfApi.inspect_scheduled_job()"""
+        scheduled_job_id = args.get("scheduled_job_id")
+        if not scheduled_job_id:
+            return {
+                "formatted": "scheduled_job_id is required",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+
+        scheduled_job = await _async_call(
+            self.api.inspect_scheduled_job,
+            scheduled_job_id=scheduled_job_id,
+            namespace=args.get("namespace") or self.namespace,
+        )
+
+        scheduled_dict = _scheduled_job_info_to_dict(scheduled_job)
+        formatted_details = format_scheduled_job_details(scheduled_dict)
+
+        return {
+            "formatted": f"**Scheduled Job Details:**\n\n{formatted_details}",
+            "totalResults": 1,
+            "resultsShared": 1,
+        }
+
+    async def _delete_scheduled_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Delete scheduled job using HfApi.delete_scheduled_job()"""
+        scheduled_job_id = args.get("scheduled_job_id")
+        if not scheduled_job_id:
+            return {
+                "formatted": "scheduled_job_id is required",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+
+        await _async_call(
+            self.api.delete_scheduled_job,
+            scheduled_job_id=scheduled_job_id,
+            namespace=args.get("namespace") or self.namespace,
+        )
+
+        return {
+            "formatted": f"✓ Scheduled job {scheduled_job_id} has been deleted.",
+            "totalResults": 1,
+            "resultsShared": 1,
+        }
+
+    async def _suspend_scheduled_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Suspend scheduled job using HfApi.suspend_scheduled_job()"""
+        scheduled_job_id = args.get("scheduled_job_id")
+        if not scheduled_job_id:
+            return {
+                "formatted": "scheduled_job_id is required",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+
+        await _async_call(
+            self.api.suspend_scheduled_job,
+            scheduled_job_id=scheduled_job_id,
+            namespace=args.get("namespace") or self.namespace,
+        )
+
+        response = f"""✓ Scheduled job {scheduled_job_id} has been suspended.
+
+To resume, call this tool with `{{"operation": "scheduled resume", "args": {{"scheduled_job_id": "{scheduled_job_id}"}}}}`"""
+
+        return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+    async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult:
+        """Resume scheduled job using HfApi.resume_scheduled_job()"""
+        scheduled_job_id = args.get("scheduled_job_id")
+        if not scheduled_job_id:
+            return {
+                "formatted": "scheduled_job_id is required",
+                "totalResults": 0,
+                "resultsShared": 0,
+                "isError": True,
+            }
+
+        await _async_call(
+            self.api.resume_scheduled_job,
+            scheduled_job_id=scheduled_job_id,
+            namespace=args.get("namespace") or self.namespace,
+        )
+
+        response = f"""✓ Scheduled job {scheduled_job_id} has been resumed.
+
+To inspect, call this tool with `{{"operation": "scheduled inspect", "args": {{"scheduled_job_id": "{scheduled_job_id}"}}}}`"""
+
+        return {"formatted": response, "totalResults": 1, "resultsShared": 1}
+
+
+# Tool specification for agent registration
+HF_JOBS_TOOL_SPEC = {
+    "name": "hf_jobs",
+    "description": (
+        "Manage Hugging Face CPU/GPU compute jobs. Run commands in Docker containers, "
+        "execute Python scripts with UV. List, schedule and monitor jobs/logs. "
+        "Call this tool with no operation for full usage instructions and examples."
+    ),
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "operation": {
+                "type": "string",
+                "enum": [
+                    "run",
+                    "uv",
+                    "ps",
+                    "logs",
+                    "inspect",
+                    "cancel",
+                    "scheduled run",
+                    "scheduled uv",
+                    "scheduled ps",
+                    "scheduled inspect",
+                    "scheduled delete",
+                    "scheduled suspend",
+                    "scheduled resume",
+                ],
+                "description": (
+                    "Operation to execute. Valid values: [run, uv, ps, logs, inspect, cancel, "
+                    "scheduled run, scheduled uv, scheduled ps, scheduled inspect, scheduled delete, "
+                    "scheduled suspend, scheduled resume]"
+                ),
+            },
+            "args": {
+                "type": "object",
+                "description": "Operation-specific arguments as a JSON object",
+                "additionalProperties": True,
+            },
+        },
+    },
+}
+
+
+async def hf_jobs_handler(arguments: Dict[str, Any]) -> tuple[str, bool]:
+    """Handler for agent tool router"""
+    try:
+        tool = HfJobsTool()
+        result = await tool.execute(arguments)
+        return result["formatted"], not result.get("isError", False)
+    except Exception as e:
+        return f"Error executing HF Jobs tool: {str(e)}", False

agent/tools/types.py

@@ -0,0 +1,16 @@
+"""
+Types for Hugging Face tools
+
+Ported from: hf-mcp-server/packages/mcp/src/types/
+"""
+
+from typing import TypedDict
+
+
+class ToolResult(TypedDict, total=False):
+    """Result returned by HF tool operations"""
+
+    formatted: str
+    totalResults: int
+    resultsShared: int
+    isError: bool

agent/tools/utilities.py

@@ -0,0 +1,142 @@
+"""
+Utility functions for Hugging Face tools
+
+Ported from: hf-mcp-server/packages/mcp/src/jobs/formatters.ts
+"""
+
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+
+def truncate(text: str, max_length: int) -> str:
+    """Truncate a string to a maximum length with ellipsis"""
+    if len(text) <= max_length:
+        return text
+    return text[: max_length - 3] + "..."
+
+
+def format_date(date_str: Optional[str]) -> str:
+    """Format a date string to a readable format"""
+    if not date_str:
+        return "N/A"
+    try:
+        date = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
+        return date.strftime("%Y-%m-%d %H:%M:%S")
+    except Exception:
+        return date_str
+
+
+def format_command(command: Optional[List[str]]) -> str:
+    """Format command array as a single string"""
+    if not command or len(command) == 0:
+        return "N/A"
+    return " ".join(command)
+
+
+def get_image_or_space(job: Dict[str, Any]) -> str:
+    """Get image/space identifier from job"""
+    if job.get("spaceId"):
+        return job["spaceId"]
+    if job.get("dockerImage"):
+        return job["dockerImage"]
+    return "N/A"
+
+
+def format_jobs_table(jobs: List[Dict[str, Any]]) -> str:
+    """Format jobs as a markdown table"""
+    if len(jobs) == 0:
+        return "No jobs found."
+
+    # Calculate dynamic ID column width
+    longest_id_length = max(len(job["id"]) for job in jobs)
+    id_column_width = max(longest_id_length, len("JOB ID"))
+
+    # Define column widths
+    col_widths = {
+        "id": id_column_width,
+        "image": 20,
+        "command": 30,
+        "created": 19,
+        "status": 12,
+    }
+
+    # Build header
+    header = f"| {'JOB ID'.ljust(col_widths['id'])} | {'IMAGE/SPACE'.ljust(col_widths['image'])} | {'COMMAND'.ljust(col_widths['command'])} | {'CREATED'.ljust(col_widths['created'])} | {'STATUS'.ljust(col_widths['status'])} |"
+    separator = f"|{'-' * (col_widths['id'] + 2)}|{'-' * (col_widths['image'] + 2)}|{'-' * (col_widths['command'] + 2)}|{'-' * (col_widths['created'] + 2)}|{'-' * (col_widths['status'] + 2)}|"
+
+    # Build rows
+    rows = []
+    for job in jobs:
+        job_id = job["id"]
+        image = truncate(get_image_or_space(job), col_widths["image"])
+        command = truncate(format_command(job.get("command")), col_widths["command"])
+        created = truncate(format_date(job.get("createdAt")), col_widths["created"])
+        status = truncate(job["status"]["stage"], col_widths["status"])
+
+        rows.append(
+            f"| {job_id.ljust(col_widths['id'])} | {image.ljust(col_widths['image'])} | {command.ljust(col_widths['command'])} | {created.ljust(col_widths['created'])} | {status.ljust(col_widths['status'])} |"
+        )
+
+    return "\n".join([header, separator] + rows)
+
+
+def format_scheduled_jobs_table(jobs: List[Dict[str, Any]]) -> str:
+    """Format scheduled jobs as a markdown table"""
+    if len(jobs) == 0:
+        return "No scheduled jobs found."
+
+    # Calculate dynamic ID column width
+    longest_id_length = max(len(job["id"]) for job in jobs)
+    id_column_width = max(longest_id_length, len("ID"))
+
+    # Define column widths
+    col_widths = {
+        "id": id_column_width,
+        "schedule": 12,
+        "image": 18,
+        "command": 25,
+        "lastRun": 19,
+        "nextRun": 19,
+        "suspend": 9,
+    }
+
+    # Build header
+    header = f"| {'ID'.ljust(col_widths['id'])} | {'SCHEDULE'.ljust(col_widths['schedule'])} | {'IMAGE/SPACE'.ljust(col_widths['image'])} | {'COMMAND'.ljust(col_widths['command'])} | {'LAST RUN'.ljust(col_widths['lastRun'])} | {'NEXT RUN'.ljust(col_widths['nextRun'])} | {'SUSPENDED'.ljust(col_widths['suspend'])} |"
+    separator = f"|{'-' * (col_widths['id'] + 2)}|{'-' * (col_widths['schedule'] + 2)}|{'-' * (col_widths['image'] + 2)}|{'-' * (col_widths['command'] + 2)}|{'-' * (col_widths['lastRun'] + 2)}|{'-' * (col_widths['nextRun'] + 2)}|{'-' * (col_widths['suspend'] + 2)}|"
+
+    # Build rows
+    rows = []
+    for job in jobs:
+        job_id = job["id"]
+        schedule = truncate(job["schedule"], col_widths["schedule"])
+        image = truncate(get_image_or_space(job["jobSpec"]), col_widths["image"])
+        command = truncate(
+            format_command(job["jobSpec"].get("command")), col_widths["command"]
+        )
+        last_run = truncate(format_date(job.get("lastRun")), col_widths["lastRun"])
+        next_run = truncate(format_date(job.get("nextRun")), col_widths["nextRun"])
+        suspend = "Yes" if job.get("suspend") else "No"
+
+        rows.append(
+            f"| {job_id.ljust(col_widths['id'])} | {schedule.ljust(col_widths['schedule'])} | {image.ljust(col_widths['image'])} | {command.ljust(col_widths['command'])} | {last_run.ljust(col_widths['lastRun'])} | {next_run.ljust(col_widths['nextRun'])} | {suspend.ljust(col_widths['suspend'])} |"
+        )
+
+    return "\n".join([header, separator] + rows)
+
+
+def format_job_details(jobs: Any) -> str:
+    """Format job details as JSON in a markdown code block"""
+    import json
+
+    job_array = jobs if isinstance(jobs, list) else [jobs]
+    json_str = json.dumps(job_array, indent=2)
+    return f"```json\n{json_str}\n```"
+
+
+def format_scheduled_job_details(jobs: Any) -> str:
+    """Format scheduled job details as JSON in a markdown code block"""
+    import json
+
+    job_array = jobs if isinstance(jobs, list) else [jobs]
+    json_str = json.dumps(job_array, indent=2)
+    return f"```json\n{json_str}\n```"

pyproject.toml

@@ -19,4 +19,5 @@ dependencies = [
     "lmnr[all]>=0.7.23",
     "transformers>=2.3.0",
     "torch>=2.9.1",
+    "pytest>=9.0.2",
 ]

tests/integration/tools/test_jobs_integration.py

@@ -0,0 +1,452 @@
+#!/usr/bin/env python3
+"""
+Integration tests for refactored HF Jobs Tool
+Tests with real HF API using HF_TOKEN from environment
+"""
+import os
+import sys
+import asyncio
+import time
+
+# Add parent directory to path
+sys.path.insert(0, '.')
+
+from agent.tools.jobs_tool import HfJobsTool
+
+# ANSI color codes for better output
+GREEN = '\033[92m'
+YELLOW = '\033[93m'
+RED = '\033[91m'
+BLUE = '\033[94m'
+RESET = '\033[0m'
+
+
+def print_test(msg):
+    """Print test message in blue"""
+    print(f"{BLUE}[TEST]{RESET} {msg}")
+
+
+def print_success(msg):
+    """Print success message in green"""
+    print(f"{GREEN}✓{RESET} {msg}")
+
+
+def print_warning(msg):
+    """Print warning message in yellow"""
+    print(f"{YELLOW}⚠{RESET} {msg}")
+
+
+def print_error(msg):
+    """Print error message in red"""
+    print(f"{RED}✗{RESET} {msg}")
+
+
+async def test_basic_job_run(tool):
+    """Test running a basic job"""
+    print_test("Running a simple Python job...")
+
+    result = await tool.execute({
+        "operation": "run",
+        "args": {
+            "image": "python:3.12",
+            "command": ["python", "-c", "print('Hello from HF Jobs!')"],
+            "flavor": "cpu-basic",
+            "timeout": "5m",
+            "detach": True  # Don't wait for completion
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to run job: {result['formatted']}")
+        return None
+
+    # Extract job ID from response
+    import re
+    job_id_match = re.search(r'\*\*Job ID:\*\* (\S+)', result['formatted'])
+    if job_id_match:
+        job_id = job_id_match.group(1)
+        print_success(f"Job started with ID: {job_id}")
+        return job_id
+
+    print_error("Could not extract job ID from response")
+    return None
+
+
+async def test_list_jobs(tool):
+    """Test listing jobs"""
+    print_test("Listing running jobs...")
+
+    result = await tool.execute({
+        "operation": "ps",
+        "args": {}
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to list jobs: {result['formatted']}")
+        return False
+
+    print_success(f"Listed jobs: {result['totalResults']} running")
+    if result['totalResults'] > 0:
+        print(f"   {result['formatted'][:200]}...")
+    return True
+
+
+async def test_inspect_job(tool, job_id):
+    """Test inspecting a specific job"""
+    print_test(f"Inspecting job {job_id}...")
+
+    result = await tool.execute({
+        "operation": "inspect",
+        "args": {
+            "job_id": job_id
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to inspect job: {result['formatted']}")
+        return False
+
+    print_success(f"Inspected job successfully")
+    return True
+
+
+async def test_get_logs(tool, job_id):
+    """Test fetching job logs"""
+    print_test(f"Fetching logs for job {job_id}...")
+
+    # Wait a bit for logs to be available
+    await asyncio.sleep(2)
+
+    result = await tool.execute({
+        "operation": "logs",
+        "args": {
+            "job_id": job_id
+        }
+    })
+
+    if result.get("isError"):
+        print_warning(f"Could not fetch logs (might be too early): {result['formatted'][:100]}")
+        return False
+
+    print_success(f"Fetched logs successfully")
+    if "Hello from HF Jobs!" in result['formatted']:
+        print_success("  Found expected output in logs!")
+    return True
+
+
+async def test_cancel_job(tool, job_id):
+    """Test cancelling a job"""
+    print_test(f"Cancelling job {job_id}...")
+
+    result = await tool.execute({
+        "operation": "cancel",
+        "args": {
+            "job_id": job_id
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to cancel job: {result['formatted']}")
+        return False
+
+    print_success(f"Cancelled job successfully")
+    return True
+
+
+async def test_uv_job(tool):
+    """Test running a UV job"""
+    print_test("Running a UV Python script job...")
+
+    result = await tool.execute({
+        "operation": "uv",
+        "args": {
+            "script": "print('Hello from UV!')\nimport sys\nprint(f'Python version: {sys.version}')",
+            "flavor": "cpu-basic",
+            "timeout": "5m",
+            "detach": True
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to run UV job: {result['formatted']}")
+        return None
+
+    # Extract job ID
+    import re
+    job_id_match = re.search(r'UV Job started: (\S+)', result['formatted'])
+    if job_id_match:
+        job_id = job_id_match.group(1)
+        print_success(f"UV job started with ID: {job_id}")
+        return job_id
+
+    print_error("Could not extract job ID from response")
+    return None
+
+
+async def test_list_all_jobs(tool):
+    """Test listing all jobs (including completed)"""
+    print_test("Listing all jobs (including completed)...")
+
+    result = await tool.execute({
+        "operation": "ps",
+        "args": {
+            "all": True
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to list all jobs: {result['formatted']}")
+        return False
+
+    print_success(f"Listed all jobs: {result['totalResults']} total")
+    return True
+
+
+async def test_scheduled_job(tool):
+    """Test creating and managing a scheduled job"""
+    print_test("Creating a scheduled job (daily at midnight)...")
+
+    result = await tool.execute({
+        "operation": "scheduled run",
+        "args": {
+            "image": "python:3.12",
+            "command": ["python", "-c", "print('Scheduled job running!')"],
+            "schedule": "@daily",
+            "flavor": "cpu-basic",
+            "timeout": "5m"
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to create scheduled job: {result['formatted']}")
+        return None
+
+    # Extract scheduled job ID
+    import re
+    job_id_match = re.search(r'\*\*Scheduled Job ID:\*\* (\S+)', result['formatted'])
+    if not job_id_match:
+        print_error("Could not extract scheduled job ID")
+        return None
+
+    scheduled_job_id = job_id_match.group(1)
+    print_success(f"Scheduled job created with ID: {scheduled_job_id}")
+    return scheduled_job_id
+
+
+async def test_list_scheduled_jobs(tool):
+    """Test listing scheduled jobs"""
+    print_test("Listing scheduled jobs...")
+
+    result = await tool.execute({
+        "operation": "scheduled ps",
+        "args": {}
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to list scheduled jobs: {result['formatted']}")
+        return False
+
+    print_success(f"Listed scheduled jobs: {result['totalResults']} active")
+    return True
+
+
+async def test_inspect_scheduled_job(tool, scheduled_job_id):
+    """Test inspecting a scheduled job"""
+    print_test(f"Inspecting scheduled job {scheduled_job_id}...")
+
+    result = await tool.execute({
+        "operation": "scheduled inspect",
+        "args": {
+            "scheduled_job_id": scheduled_job_id
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to inspect scheduled job: {result['formatted']}")
+        return False
+
+    print_success(f"Inspected scheduled job successfully")
+    return True
+
+
+async def test_suspend_scheduled_job(tool, scheduled_job_id):
+    """Test suspending a scheduled job"""
+    print_test(f"Suspending scheduled job {scheduled_job_id}...")
+
+    result = await tool.execute({
+        "operation": "scheduled suspend",
+        "args": {
+            "scheduled_job_id": scheduled_job_id
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to suspend scheduled job: {result['formatted']}")
+        return False
+
+    print_success(f"Suspended scheduled job successfully")
+    return True
+
+
+async def test_resume_scheduled_job(tool, scheduled_job_id):
+    """Test resuming a scheduled job"""
+    print_test(f"Resuming scheduled job {scheduled_job_id}...")
+
+    result = await tool.execute({
+        "operation": "scheduled resume",
+        "args": {
+            "scheduled_job_id": scheduled_job_id
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to resume scheduled job: {result['formatted']}")
+        return False
+
+    print_success(f"Resumed scheduled job successfully")
+    return True
+
+
+async def test_delete_scheduled_job(tool, scheduled_job_id):
+    """Test deleting a scheduled job"""
+    print_test(f"Deleting scheduled job {scheduled_job_id}...")
+
+    result = await tool.execute({
+        "operation": "scheduled delete",
+        "args": {
+            "scheduled_job_id": scheduled_job_id
+        }
+    })
+
+    if result.get("isError"):
+        print_error(f"Failed to delete scheduled job: {result['formatted']}")
+        return False
+
+    print_success(f"Deleted scheduled job successfully")
+    return True
+
+
+async def main():
+    """Run all integration tests"""
+    print("=" * 70)
+    print(f"{BLUE}HF Jobs Tool - Integration Tests{RESET}")
+    print("=" * 70)
+    print()
+
+    # Check for HF_TOKEN
+    hf_token = os.environ.get('HF_TOKEN')
+    if not hf_token:
+        print_error("HF_TOKEN not found in environment variables!")
+        print_warning("Set it with: export HF_TOKEN='your_token_here'")
+        sys.exit(1)
+
+    print_success(f"Found HF_TOKEN (length: {len(hf_token)})")
+    print()
+
+    # Initialize tool with token
+    tool = HfJobsTool(hf_token=hf_token)
+
+    # Track job IDs for cleanup
+    job_ids = []
+    scheduled_job_ids = []
+
+    try:
+        # Test 1: Run basic job
+        print(f"\n{YELLOW}{'=' * 70}{RESET}")
+        print(f"{YELLOW}Test Suite 1: Regular Jobs{RESET}")
+        print(f"{YELLOW}{'=' * 70}{RESET}\n")
+
+        job_id = await test_basic_job_run(tool)
+        if job_id:
+            job_ids.append(job_id)
+
+            # Wait a moment for job to register
+            await asyncio.sleep(1)
+
+            # Test 2: List jobs
+            await test_list_jobs(tool)
+
+            # Test 3: Inspect job
+            await test_inspect_job(tool, job_id)
+
+            # Test 4: Get logs
+            await test_get_logs(tool, job_id)
+
+            # Test 5: Cancel job (cleanup)
+            await test_cancel_job(tool, job_id)
+
+        # Test 6: UV job
+        print()
+        uv_job_id = await test_uv_job(tool)
+        if uv_job_id:
+            job_ids.append(uv_job_id)
+            await asyncio.sleep(1)
+            await test_cancel_job(tool, uv_job_id)
+
+        # Test 7: List all jobs
+        print()
+        await test_list_all_jobs(tool)
+
+        # Test Suite 2: Scheduled Jobs
+        print(f"\n{YELLOW}{'=' * 70}{RESET}")
+        print(f"{YELLOW}Test Suite 2: Scheduled Jobs{RESET}")
+        print(f"{YELLOW}{'=' * 70}{RESET}\n")
+
+        scheduled_job_id = await test_scheduled_job(tool)
+        if scheduled_job_id:
+            scheduled_job_ids.append(scheduled_job_id)
+
+            # Wait a moment for job to register
+            await asyncio.sleep(1)
+
+            # Test scheduled job operations
+            await test_list_scheduled_jobs(tool)
+            print()
+            await test_inspect_scheduled_job(tool, scheduled_job_id)
+            print()
+            await test_suspend_scheduled_job(tool, scheduled_job_id)
+            print()
+            await test_resume_scheduled_job(tool, scheduled_job_id)
+            print()
+
+            # Cleanup: Delete scheduled job
+            await test_delete_scheduled_job(tool, scheduled_job_id)
+
+        # Final summary
+        print(f"\n{YELLOW}{'=' * 70}{RESET}")
+        print(f"{GREEN}✓ All integration tests completed!{RESET}")
+        print(f"{YELLOW}{'=' * 70}{RESET}\n")
+
+        print_success("Refactored implementation works correctly with real HF API")
+        print_success("All 13 operations tested and verified")
+        print()
+        print(f"{BLUE}Summary:{RESET}")
+        print(f"  • Regular jobs: ✓ run, list, inspect, logs, cancel")
+        print(f"  • UV jobs: ✓ run")
+        print(f"  • Scheduled jobs: ✓ create, list, inspect, suspend, resume, delete")
+        print()
+
+    except Exception as e:
+        print_error(f"Test failed with exception: {str(e)}")
+        import traceback
+        traceback.print_exc()
+
+        # Attempt cleanup
+        print(f"\n{YELLOW}Attempting cleanup...{RESET}")
+        for job_id in job_ids:
+            try:
+                await test_cancel_job(tool, job_id)
+            except:
+                pass
+
+        for scheduled_job_id in scheduled_job_ids:
+            try:
+                await test_delete_scheduled_job(tool, scheduled_job_id)
+            except:
+                pass
+
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/unit/tools/test_jobs_tool.py

@@ -0,0 +1,454 @@
+"""
+Tests for HF Jobs Tool
+
+Tests the refactored jobs tool implementation using huggingface-hub library
+"""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from agent.tools.jobs_tool import HfJobsTool, hf_jobs_handler
+
+
+def create_mock_job_info(
+    job_id="test-job-1",
+    stage="RUNNING",
+    command=None,
+    docker_image="python:3.12",
+):
+    """Create a mock JobInfo object"""
+    from huggingface_hub._jobs_api import JobInfo
+
+    if command is None:
+        command = ["echo", "test"]
+
+    return JobInfo(
+        id=job_id,
+        created_at="2024-01-01T00:00:00.000000Z",
+        docker_image=docker_image,
+        space_id=None,
+        command=command,
+        arguments=[],
+        environment={},
+        secrets={},
+        flavor="cpu-basic",
+        status={"stage": stage, "message": None},
+        owner={"id": "123", "name": "test-user", "type": "user"},
+        endpoint="https://huggingface.co",
+        url=f"https://huggingface.co/jobs/test-user/{job_id}",
+    )
+
+
+def create_mock_scheduled_job_info(
+    job_id="sched-job-1",
+    schedule="@daily",
+    suspend=False,
+):
+    """Create a mock ScheduledJobInfo object"""
+    from huggingface_hub._jobs_api import ScheduledJobInfo
+
+    return ScheduledJobInfo(
+        id=job_id,
+        created_at="2024-01-01T00:00:00.000000Z",
+        job_spec={
+            "docker_image": "python:3.12",
+            "space_id": None,
+            "command": ["python", "backup.py"],
+            "arguments": [],
+            "environment": {},
+            "secrets": {},
+            "flavor": "cpu-basic",
+            "timeout": 1800,
+            "tags": None,
+            "arch": None,
+        },
+        schedule=schedule,
+        suspend=suspend,
+        concurrency=False,
+        status={
+            "last_job": None,
+            "next_job_run_at": "2024-01-02T00:00:00.000000Z",
+        },
+        owner={"id": "123", "name": "test-user", "type": "user"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_show_help():
+    """Test that help message is shown when no operation specified"""
+    tool = HfJobsTool()
+    result = await tool.execute({})
+
+    assert "HuggingFace Jobs API" in result["formatted"]
+    assert "Available Commands" in result["formatted"]
+    assert result["totalResults"] == 1
+    assert not result.get("isError", False)
+
+
+@pytest.mark.asyncio
+async def test_show_operation_help():
+    """Test operation-specific help"""
+    tool = HfJobsTool()
+    result = await tool.execute({"operation": "run", "args": {"help": True}})
+
+    assert "Help for operation" in result["formatted"]
+    assert result["totalResults"] == 1
+
+
+@pytest.mark.asyncio
+async def test_invalid_operation():
+    """Test invalid operation handling"""
+    tool = HfJobsTool()
+    result = await tool.execute({"operation": "invalid_op"})
+
+    assert result.get("isError") == True
+    assert "Unknown operation" in result["formatted"]
+
+
+@pytest.mark.asyncio
+async def test_run_job_missing_command():
+    """Test run job with missing required parameter"""
+    tool = HfJobsTool()
+
+    # Mock the HfApi.run_job to raise an error
+    with patch.object(tool.api, "run_job") as mock_run:
+        mock_run.side_effect = Exception("command parameter is required")
+
+        result = await tool.execute(
+            {"operation": "run", "args": {"image": "python:3.12"}}
+        )
+
+        assert result.get("isError") == True
+
+
+@pytest.mark.asyncio
+async def test_list_jobs_mock():
+    """Test list jobs with mock API"""
+    tool = HfJobsTool()
+
+    # Create mock job objects
+    running_job = create_mock_job_info("test-job-1", "RUNNING")
+    completed_job = create_mock_job_info(
+        "test-job-2", "COMPLETED", ["python", "script.py"]
+    )
+
+    # Mock the HfApi.list_jobs method
+    with patch.object(tool.api, "list_jobs") as mock_list:
+        mock_list.return_value = [running_job, completed_job]
+
+        # Test listing only running jobs (default)
+        result = await tool.execute({"operation": "ps"})
+
+        assert not result.get("isError", False)
+        assert "test-job-1" in result["formatted"]
+        assert "test-job-2" not in result["formatted"]  # COMPLETED jobs filtered out
+        assert result["totalResults"] == 1
+        assert result["resultsShared"] == 1
+
+        # Test listing all jobs
+        result = await tool.execute({"operation": "ps", "args": {"all": True}})
+
+        assert not result.get("isError", False)
+        assert "test-job-1" in result["formatted"]
+        assert "test-job-2" in result["formatted"]
+        assert result["totalResults"] == 2
+        assert result["resultsShared"] == 2
+
+
+@pytest.mark.asyncio
+async def test_inspect_job_mock():
+    """Test inspect job with mock API"""
+    tool = HfJobsTool()
+
+    mock_job = create_mock_job_info("test-job-1", "RUNNING")
+
+    with patch.object(tool.api, "inspect_job") as mock_inspect:
+        mock_inspect.return_value = mock_job
+
+        result = await tool.execute(
+            {"operation": "inspect", "args": {"job_id": "test-job-1"}}
+        )
+
+        assert not result.get("isError", False)
+        assert "test-job-1" in result["formatted"]
+        assert "Job Details" in result["formatted"]
+        mock_inspect.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_cancel_job_mock():
+    """Test cancel job with mock API"""
+    tool = HfJobsTool()
+
+    with patch.object(tool.api, "cancel_job") as mock_cancel:
+        mock_cancel.return_value = None
+
+        result = await tool.execute(
+            {"operation": "cancel", "args": {"job_id": "test-job-1"}}
+        )
+
+        assert not result.get("isError", False)
+        assert "cancelled" in result["formatted"]
+        assert "test-job-1" in result["formatted"]
+        mock_cancel.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_run_job_mock():
+    """Test run job with mock API"""
+    tool = HfJobsTool()
+
+    mock_job = create_mock_job_info("new-job-123", "RUNNING")
+
+    with patch.object(tool.api, "run_job") as mock_run:
+        mock_run.return_value = mock_job
+
+        result = await tool.execute(
+            {
+                "operation": "run",
+                "args": {
+                    "image": "python:3.12",
+                    "command": ["python", "-c", "print('test')"],
+                    "flavor": "cpu-basic",
+                    "detach": True,
+                },
+            }
+        )
+
+        assert not result.get("isError", False)
+        assert "new-job-123" in result["formatted"]
+        assert "Job started" in result["formatted"]
+        mock_run.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_run_uv_job_mock():
+    """Test run UV job with mock API"""
+    tool = HfJobsTool()
+
+    mock_job = create_mock_job_info("uv-job-456", "RUNNING")
+
+    with patch.object(tool.api, "run_uv_job") as mock_run:
+        mock_run.return_value = mock_job
+
+        result = await tool.execute(
+            {
+                "operation": "uv",
+                "args": {
+                    "script": "print('Hello UV')",
+                    "flavor": "cpu-basic",
+                },
+            }
+        )
+
+        assert not result.get("isError", False)
+        assert "uv-job-456" in result["formatted"]
+        assert "UV Job started" in result["formatted"]
+        mock_run.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_get_logs_mock():
+    """Test get logs with mock API"""
+    tool = HfJobsTool()
+
+    # Mock fetch_job_logs to return a generator
+    def log_generator():
+        yield "Log line 1"
+        yield "Log line 2"
+        yield "Hello from HF Jobs!"
+
+    with patch.object(tool.api, "fetch_job_logs") as mock_logs:
+        mock_logs.return_value = log_generator()
+
+        result = await tool.execute(
+            {"operation": "logs", "args": {"job_id": "test-job-1"}}
+        )
+
+        assert not result.get("isError", False)
+        assert "Log line 1" in result["formatted"]
+        assert "Hello from HF Jobs!" in result["formatted"]
+
+
+@pytest.mark.asyncio
+async def test_handler():
+    """Test the handler function"""
+    with patch("agent.tools.jobs_tool.HfJobsTool") as MockTool:
+        mock_tool_instance = MockTool.return_value
+        mock_tool_instance.execute = AsyncMock(
+            return_value={
+                "formatted": "Test output",
+                "totalResults": 1,
+                "resultsShared": 1,
+                "isError": False,
+            }
+        )
+
+        output, success = await hf_jobs_handler({"operation": "ps"})
+
+        assert success == True
+        assert "Test output" in output
+
+
+@pytest.mark.asyncio
+async def test_handler_error():
+    """Test handler with error"""
+    with patch("agent.tools.jobs_tool.HfJobsTool") as MockTool:
+        MockTool.side_effect = Exception("Test error")
+
+        output, success = await hf_jobs_handler({})
+
+        assert success == False
+        assert "Error" in output
+
+
+@pytest.mark.asyncio
+async def test_scheduled_jobs_mock():
+    """Test scheduled jobs operations with mock API"""
+    tool = HfJobsTool()
+
+    mock_scheduled_job = create_mock_scheduled_job_info()
+
+    # Test list scheduled jobs
+    with patch.object(tool.api, "list_scheduled_jobs") as mock_list:
+        mock_list.return_value = [mock_scheduled_job]
+
+        result = await tool.execute({"operation": "scheduled ps"})
+
+        assert not result.get("isError", False)
+        assert "sched-job-1" in result["formatted"]
+        assert "Scheduled Jobs" in result["formatted"]
+
+
+@pytest.mark.asyncio
+async def test_create_scheduled_job_mock():
+    """Test create scheduled job with mock API"""
+    tool = HfJobsTool()
+
+    mock_scheduled_job = create_mock_scheduled_job_info()
+
+    with patch.object(tool.api, "create_scheduled_job") as mock_create:
+        mock_create.return_value = mock_scheduled_job
+
+        result = await tool.execute(
+            {
+                "operation": "scheduled run",
+                "args": {
+                    "image": "python:3.12",
+                    "command": ["python", "backup.py"],
+                    "schedule": "@daily",
+                    "flavor": "cpu-basic",
+                },
+            }
+        )
+
+        assert not result.get("isError", False)
+        assert "sched-job-1" in result["formatted"]
+        assert "Scheduled job created" in result["formatted"]
+        mock_create.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_inspect_scheduled_job_mock():
+    """Test inspect scheduled job with mock API"""
+    tool = HfJobsTool()
+
+    mock_scheduled_job = create_mock_scheduled_job_info()
+
+    with patch.object(tool.api, "inspect_scheduled_job") as mock_inspect:
+        mock_inspect.return_value = mock_scheduled_job
+
+        result = await tool.execute(
+            {
+                "operation": "scheduled inspect",
+                "args": {"scheduled_job_id": "sched-job-1"},
+            }
+        )
+
+        assert not result.get("isError", False)
+        assert "sched-job-1" in result["formatted"]
+        assert "Scheduled Job Details" in result["formatted"]
+
+
+@pytest.mark.asyncio
+async def test_suspend_scheduled_job_mock():
+    """Test suspend scheduled job with mock API"""
+    tool = HfJobsTool()
+
+    with patch.object(tool.api, "suspend_scheduled_job") as mock_suspend:
+        mock_suspend.return_value = None
+
+        result = await tool.execute(
+            {
+                "operation": "scheduled suspend",
+                "args": {"scheduled_job_id": "sched-job-1"},
+            }
+        )
+
+        assert not result.get("isError", False)
+        assert "suspended" in result["formatted"]
+        assert "sched-job-1" in result["formatted"]
+
+
+@pytest.mark.asyncio
+async def test_resume_scheduled_job_mock():
+    """Test resume scheduled job with mock API"""
+    tool = HfJobsTool()
+
+    with patch.object(tool.api, "resume_scheduled_job") as mock_resume:
+        mock_resume.return_value = None
+
+        result = await tool.execute(
+            {
+                "operation": "scheduled resume",
+                "args": {"scheduled_job_id": "sched-job-1"},
+            }
+        )
+
+        assert not result.get("isError", False)
+        assert "resumed" in result["formatted"]
+        assert "sched-job-1" in result["formatted"]
+
+
+@pytest.mark.asyncio
+async def test_delete_scheduled_job_mock():
+    """Test delete scheduled job with mock API"""
+    tool = HfJobsTool()
+
+    with patch.object(tool.api, "delete_scheduled_job") as mock_delete:
+        mock_delete.return_value = None
+
+        result = await tool.execute(
+            {
+                "operation": "scheduled delete",
+                "args": {"scheduled_job_id": "sched-job-1"},
+            }
+        )
+
+        assert not result.get("isError", False)
+        assert "deleted" in result["formatted"]
+        assert "sched-job-1" in result["formatted"]
+
+
+@pytest.mark.asyncio
+async def test_list_jobs_with_status_filter():
+    """Test list jobs with status filter"""
+    tool = HfJobsTool()
+
+    running_job = create_mock_job_info("job-1", "RUNNING")
+    completed_job = create_mock_job_info("job-2", "COMPLETED")
+    error_job = create_mock_job_info("job-3", "ERROR")
+
+    with patch.object(tool.api, "list_jobs") as mock_list:
+        mock_list.return_value = [running_job, completed_job, error_job]
+
+        # Filter by status
+        result = await tool.execute(
+            {"operation": "ps", "args": {"all": True, "status": "ERROR"}}
+        )
+
+        assert not result.get("isError", False)
+        assert "job-3" in result["formatted"]
+        assert "job-1" not in result["formatted"]
+        assert result["resultsShared"] == 1

uv.lock

@@ -893,6 +893,7 @@ dependencies = [
     { name = "numpy" },
     { name = "pandas" },
     { name = "pydantic" },
+    { name = "pytest" },
     { name = "python-dotenv" },
     { name = "requests" },
     { name = "tenacity" },
@@ -911,6 +912,7 @@ requires-dist = [
     { name = "numpy", specifier = ">=1.24.0" },
     { name = "pandas", specifier = ">=2.3.3" },
     { name = "pydantic", specifier = ">=2.12.3" },
+    { name = "pytest", specifier = ">=9.0.2" },
     { name = "python-dotenv", specifier = ">=1.2.1" },
     { name = "requests", specifier = ">=2.32.5" },
     { name = "tenacity", specifier = ">=8.0.0" },
@@ -1098,6 +1100,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/59/91/aa6bde563e0085a02a435aa99b49ef75b0a4b062635e606dab23ce18d720/inflection-0.5.1-py2.py3-none-any.whl", hash = "sha256:f38b2b640938a4f35ade69ac3d053042959b62a0f1076a5bbaa1b9526605a8a2", size = 9454, upload-time = "2020-08-22T08:16:27.816Z" },
 ]
 
+[[package]]
+name = "iniconfig"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
+]
+
 [[package]]
 name = "inspect-ai"
 version = "0.3.149"
@@ -2703,6 +2714,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/73/cb/ac7874b3e5d58441674fb70742e6c374b28b0c7cb988d37d991cde47166c/platformdirs-4.5.0-py3-none-any.whl", hash = "sha256:e578a81bb873cbb89a41fcc904c7ef523cc18284b7e3b3ccf06aca1403b7ebd3", size = 18651, upload-time = "2025-10-08T17:44:47.223Z" },
 ]
 
+[[package]]
+name = "pluggy"
+version = "1.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
+]
+
 [[package]]
 name = "ply"
 version = "3.11"
@@ -3064,6 +3084,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/df/80/fc9d01d5ed37ba4c42ca2b55b4339ae6e200b456be3a1aaddf4a9fa99b8c/pyperclip-1.11.0-py3-none-any.whl", hash = "sha256:299403e9ff44581cb9ba2ffeed69c7aa96a008622ad0c46cb575ca75b5b84273", size = 11063, upload-time = "2025-09-26T14:40:36.069Z" },
 ]
 
+[[package]]
+name = "pytest"
+version = "9.0.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"