agent.py

"""
agent.py -- Tool-using agent core for the GLM-5.2 personal workspace Space.

This module is deliberately self-contained and network-free (the only
network call in the whole app -- the actual chat completion -- is made by
app.py using the caller's own OpenAI-compatible client). Everything here is
pure local logic:

  - a system prompt that teaches the model a strict, easy-to-parse
    tool-call protocol
  - a small set of real tools: python execution, shell execution, file
    read/write/list/delete, zip packaging, pip install
  - a per-user sandboxed workspace directory with path-traversal protection
  - simple JSON-file backed conversation memory (persists to /data if the
    Space has persistent storage attached, otherwise to a local folder
    that lives only as long as the current container)
  - a parser that pulls a tool call out of the model's raw reply text

app.py drives the actual multi-step loop; this file only exposes the
building blocks so that loop stays short and readable.
"""

import os
import re
import sys
import json
import uuid
import shutil
import zipfile
import subprocess
from pathlib import Path
from typing import Optional, Dict, Any, List, Tuple


# --------------------------------------------------------------------------
# Tunables -- safe to adjust these constants as your needs change
# --------------------------------------------------------------------------

MAX_AGENT_STEPS = 12                       # hard cap on tool round-trips per user message
MAX_TOOL_OUTPUT_CHARS = 8000               # truncate stdout/stderr fed back to the model
MAX_READ_FILE_CHARS = 20000                # truncate read_file output
MAX_FILE_WRITE_BYTES = 8 * 1024 * 1024     # 8MB per write_file call
MAX_UPLOAD_BYTES = 32 * 1024 * 1024        # 32MB per user upload
MAX_HISTORY_ITEMS_KEPT_ON_DISK = 400       # trim the on-disk transcript beyond this
MAX_HISTORY_ITEMS_SENT_TO_MODEL = 40       # trim what's replayed into the next API call
DEFAULT_TOOL_TIMEOUT = 60
MAX_TOOL_TIMEOUT = 240
ULIMIT_V_KB = 3_000_000                    # ~2.9GB address-space cap per tool execution


# --------------------------------------------------------------------------
# Storage roots
# --------------------------------------------------------------------------

def _pick_storage_root() -> Path:
    """Prefer Hugging Face Spaces persistent storage (/data) when it exists
    and is writable, otherwise fall back to a folder next to the app that
    only lives as long as the current container (reset on rebuild/restart).
    """
    data_dir = Path("/data")
    try:
        if data_dir.is_dir() and os.access(str(data_dir), os.W_OK):
            root = data_dir / "glm_workspace"
            root.mkdir(parents=True, exist_ok=True)
            return root
    except OSError:
        pass
    root = Path(__file__).resolve().parent / "workspace_store"
    root.mkdir(parents=True, exist_ok=True)
    return root


STORAGE_ROOT = _pick_storage_root()
IS_PERSISTENT = str(STORAGE_ROOT).startswith("/data")


def _slugify(raw: str) -> str:
    raw = (raw or "").strip().lower()
    slug = re.sub(r"[^a-z0-9_-]+", "-", raw).strip("-")
    return slug or "user"


def user_key_from_userinfo(user_info: Optional[dict]) -> str:
    """Stable, filesystem-safe key identifying this person's private
    workspace, derived from their HF OAuth profile."""
    if not user_info:
        return "user"
    raw = (
        user_info.get("preferred_username")
        or user_info.get("sub")
        or user_info.get("name")
        or "user"
    )
    return _slugify(str(raw))


def get_user_dir(user_key: str) -> Path:
    d = STORAGE_ROOT / _slugify(user_key)
    (d / "files").mkdir(parents=True, exist_ok=True)
    return d


def get_files_dir(user_key: str) -> Path:
    return get_user_dir(user_key) / "files"


def get_history_path(user_key: str) -> Path:
    return get_user_dir(user_key) / "chat_history.json"


# --------------------------------------------------------------------------
# History persistence
# --------------------------------------------------------------------------

def load_history(user_key: str) -> List[Dict[str, Any]]:
    p = get_history_path(user_key)
    if not p.exists():
        return []
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
        return data if isinstance(data, list) else []
    except (json.JSONDecodeError, OSError):
        return []


def save_history(user_key: str, history: List[Dict[str, Any]]) -> None:
    trimmed = history[-MAX_HISTORY_ITEMS_KEPT_ON_DISK:]
    p = get_history_path(user_key)
    try:
        p.write_text(json.dumps(trimmed, ensure_ascii=False), encoding="utf-8")
    except OSError:
        pass


def clear_history(user_key: str) -> None:
    p = get_history_path(user_key)
    try:
        if p.exists():
            p.unlink()
    except OSError:
        pass


def history_to_api_messages(history: List[Dict[str, Any]]) -> List[Dict[str, str]]:
    """Collapse stored history into the plain role/content pairs the chat
    completion API expects, trimmed to a recent window so context doesn't
    grow without bound."""
    recent = history[-MAX_HISTORY_ITEMS_SENT_TO_MODEL:]
    return [{"role": item["role"], "content": item["content"]} for item in recent]


# --------------------------------------------------------------------------
# Path safety
# --------------------------------------------------------------------------

class UnsafePathError(ValueError):
    pass


def safe_join(base: Path, rel_path: str) -> Path:
    """Resolve rel_path against base, refusing anything that would escape
    the workspace. A leading "/" is treated as "root of the workspace"
    (stripped, not rejected) since every tool path is meant to be
    workspace-relative anyway; ".." is still blocked below."""
    rel_path = (rel_path or "").strip().lstrip("/")
    if not rel_path or rel_path == ".":
        return base.resolve()
    candidate = (base / rel_path).resolve()
    base_resolved = base.resolve()
    if candidate != base_resolved and base_resolved not in candidate.parents:
        raise UnsafePathError(f"Path '{rel_path}' escapes the workspace")
    return candidate


def safe_upload_name(filename: str, target_dir: Path) -> str:
    base = Path(filename or "upload.bin").name
    stem = Path(base).stem
    suffix = re.sub(r"[^A-Za-z0-9.]+", "", Path(base).suffix)[:12]
    stem_slug = re.sub(r"[^A-Za-z0-9_.-]+", "_", stem).strip("_") or "upload"
    candidate = f"{stem_slug}{suffix}"
    if not (target_dir / candidate).exists():
        return candidate
    return f"{stem_slug}_{uuid.uuid4().hex[:6]}{suffix}"


def truncate(text: str, limit: int) -> str:
    if not text:
        return text or ""
    if len(text) <= limit:
        return text
    cut = len(text) - limit
    return text[:limit] + f"\n...[truncated {cut} more characters]..."


# --------------------------------------------------------------------------
# Sandboxed subprocess execution
# --------------------------------------------------------------------------

_SECRET_NAME_RE = re.compile(r"(TOKEN|SECRET|KEY|PASSWORD|CREDENTIAL)", re.IGNORECASE)


def _sanitized_env() -> Dict[str, str]:
    """Strip anything that looks like a credential so a tool call can never
    accidentally echo the Space's own secrets (e.g. HF_TOKEN) back into the
    conversation, where they'd be sent to the inference provider as plain
    text."""
    return {k: v for k, v in os.environ.items() if not _SECRET_NAME_RE.search(k)}


def _ulimit_prefix(cpu_cap: int) -> str:
    # Applied via the shell's own `ulimit` builtin *after* bash has already
    # exec'd (not via subprocess's preexec_fn, which runs arbitrary Python
    # between fork() and exec() and can deadlock in a multi-threaded async
    # server). Limits are inherited across the following exec.
    return f"ulimit -v {ULIMIT_V_KB} 2>/dev/null; ulimit -t {cpu_cap} 2>/dev/null; "


def _run_subprocess(cmd: List[str], cwd: Path, timeout: int) -> Dict[str, Any]:
    try:
        proc = subprocess.run(
            cmd,
            cwd=str(cwd),
            capture_output=True,
            text=True,
            timeout=timeout,
            env=_sanitized_env(),
        )
        return {
            "success": proc.returncode == 0,
            "exit_code": proc.returncode,
            "stdout": truncate(proc.stdout or "", MAX_TOOL_OUTPUT_CHARS),
            "stderr": truncate(proc.stderr or "", MAX_TOOL_OUTPUT_CHARS),
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False, "exit_code": None, "stdout": "",
            "stderr": f"Timed out after {timeout}s.",
        }
    except FileNotFoundError as e:
        return {"success": False, "exit_code": None, "stdout": "", "stderr": str(e)}


def _run_argv(argv: List[str], cwd: Path, timeout: int) -> Dict[str, Any]:
    cpu_cap = max(int(timeout) + 10, 10)
    bash_cmd = ["/bin/bash", "-c", _ulimit_prefix(cpu_cap) + 'exec "$@"', "--", *argv]
    return _run_subprocess(bash_cmd, cwd, timeout)


def _run_shell_str(command: str, cwd: Path, timeout: int) -> Dict[str, Any]:
    cpu_cap = max(int(timeout) + 10, 10)
    bash_cmd = ["/bin/bash", "-lc", _ulimit_prefix(cpu_cap) + command]
    return _run_subprocess(bash_cmd, cwd, timeout)


_SHELL_DENYLIST = [
    r"rm\s+-rf\s+/(?!\S)",                  # rm -rf /
    r"rm\s+-rf\s+/\*",                      # rm -rf /*
    r":\(\)\s*\{\s*:\|\s*:&\s*\}\s*;\s*:",  # classic fork bomb
    r"\bmkfs\b",
    r"\bdd\s+.*of=/dev/",
    r"\bshutdown\b",
    r"\breboot\b",
    r">\s*/dev/sd[a-z]",
]
_SHELL_DENYLIST_RE = re.compile("|".join(_SHELL_DENYLIST), re.IGNORECASE)


def _clean_timeout(args: dict) -> int:
    try:
        t = int(args.get("timeout") or DEFAULT_TOOL_TIMEOUT)
    except (TypeError, ValueError):
        t = DEFAULT_TOOL_TIMEOUT
    return max(1, min(t, MAX_TOOL_TIMEOUT))


# --------------------------------------------------------------------------
# Tool implementations -- each takes (args, files_dir) and returns a dict
# with at least {"success": bool, "stdout": str, "stderr": str}
# --------------------------------------------------------------------------

def tool_run_python(args: dict, files_dir: Path) -> Dict[str, Any]:
    code = args.get("code")
    if not code or not isinstance(code, str):
        return {"success": False, "stdout": "", "stderr": "Missing 'code' string."}
    timeout = _clean_timeout(args)
    script_path = files_dir / f".tmp_{uuid.uuid4().hex}.py"
    try:
        script_path.write_text(code, encoding="utf-8")
        return _run_argv([sys.executable, "-u", script_path.name], files_dir, timeout)
    finally:
        try:
            script_path.unlink(missing_ok=True)
        except OSError:
            pass


def tool_run_shell(args: dict, files_dir: Path) -> Dict[str, Any]:
    command = args.get("command")
    if not command or not isinstance(command, str):
        return {"success": False, "stdout": "", "stderr": "Missing 'command' string."}
    if _SHELL_DENYLIST_RE.search(command):
        return {
            "success": False, "stdout": "",
            "stderr": "Blocked: this command matches a destructive pattern disabled in this workspace.",
        }
    timeout = _clean_timeout(args)
    return _run_shell_str(command, files_dir, timeout)


def tool_write_file(args: dict, files_dir: Path) -> Dict[str, Any]:
    path = args.get("path")
    content = args.get("content", "")
    if not path or not isinstance(path, str):
        return {"success": False, "stdout": "", "stderr": "Missing 'path'."}
    if isinstance(content, (dict, list)):
        content = json.dumps(content, ensure_ascii=False, indent=2)
    content = "" if content is None else str(content)
    if len(content.encode("utf-8", errors="ignore")) > MAX_FILE_WRITE_BYTES:
        return {"success": False, "stdout": "", "stderr": f"File too large (max {MAX_FILE_WRITE_BYTES} bytes)."}
    try:
        target = safe_join(files_dir, path)
    except UnsafePathError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(content, encoding="utf-8")
    return {"success": True, "stdout": f"Wrote {len(content)} chars to {path}", "stderr": "", "file": path}


def tool_read_file(args: dict, files_dir: Path) -> Dict[str, Any]:
    path = args.get("path")
    if not path or not isinstance(path, str):
        return {"success": False, "stdout": "", "stderr": "Missing 'path'."}
    try:
        target = safe_join(files_dir, path)
    except UnsafePathError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}
    if not target.exists() or not target.is_file():
        return {"success": False, "stdout": "", "stderr": f"No such file: {path}"}
    try:
        text = target.read_text(encoding="utf-8", errors="replace")
    except OSError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}
    return {"success": True, "stdout": truncate(text, MAX_READ_FILE_CHARS), "stderr": ""}


def tool_list_files(args: dict, files_dir: Path) -> Dict[str, Any]:
    path = args.get("path", ".") or "."
    try:
        target = safe_join(files_dir, path)
    except UnsafePathError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}
    if not target.exists() or not target.is_dir():
        return {"success": False, "stdout": "", "stderr": f"No such directory: {path}"}
    entries = []
    for p in sorted(target.rglob("*")):
        if p.name.startswith(".tmp_") or p.name == "chat_history.json":
            continue
        rel = p.relative_to(files_dir)
        kind = "dir" if p.is_dir() else "file"
        size = str(p.stat().st_size) if p.is_file() else ""
        entries.append(f"{kind:4}  {size:>10}  {rel}")
        if len(entries) >= 500:
            entries.append("... [truncated: more than 500 entries]")
            break
    return {"success": True, "stdout": "\n".join(entries) or "(empty)", "stderr": ""}


def tool_delete_file(args: dict, files_dir: Path) -> Dict[str, Any]:
    path = args.get("path")
    if not path or not isinstance(path, str):
        return {"success": False, "stdout": "", "stderr": "Missing 'path'."}
    try:
        target = safe_join(files_dir, path)
    except UnsafePathError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}
    if target == files_dir.resolve():
        return {"success": False, "stdout": "", "stderr": "Refusing to delete the whole workspace."}
    if not target.exists():
        return {"success": False, "stdout": "", "stderr": f"No such path: {path}"}
    try:
        if target.is_dir():
            shutil.rmtree(target)
        else:
            target.unlink()
    except OSError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}
    return {"success": True, "stdout": f"Deleted {path}", "stderr": ""}


def tool_make_zip(args: dict, files_dir: Path) -> Dict[str, Any]:
    output = (args.get("output") or f"archive_{uuid.uuid4().hex[:8]}.zip").strip()
    if not output.lower().endswith(".zip"):
        output += ".zip"
    try:
        out_target = safe_join(files_dir, output)
    except UnsafePathError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}

    paths = args.get("paths") or None
    try:
        if paths:
            resolved = [safe_join(files_dir, p) for p in paths]
        else:
            resolved = [
                p for p in files_dir.iterdir()
                if p.name != Path(output).name
                and not p.name.startswith(".tmp_")
                and p.name != "chat_history.json"
            ]
    except UnsafePathError as e:
        return {"success": False, "stdout": "", "stderr": str(e)}

    out_target.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with zipfile.ZipFile(out_target, "w", zipfile.ZIP_DEFLATED) as zf:
        for item in resolved:
            if not item.exists():
                continue
            if item.is_file():
                zf.write(item, item.relative_to(files_dir))
                count += 1
            else:
                for sub in item.rglob("*"):
                    if sub.is_file():
                        zf.write(sub, sub.relative_to(files_dir))
                        count += 1
    if count == 0:
        out_target.unlink(missing_ok=True)
        return {"success": False, "stdout": "", "stderr": "Nothing to zip (no matching files found)."}
    size = out_target.stat().st_size
    return {"success": True, "stdout": f"Created {output} ({count} files, {size} bytes)", "stderr": "", "file": output}


def tool_pip_install(args: dict, files_dir: Path) -> Dict[str, Any]:
    package = args.get("package")
    if not package or not isinstance(package, str):
        return {"success": False, "stdout": "", "stderr": "Missing 'package'."}
    if re.search(r"[;&|`$\n]", package):
        return {"success": False, "stdout": "", "stderr": "Invalid characters in package spec."}
    return _run_argv(
        [sys.executable, "-m", "pip", "install", "--quiet",
         "--disable-pip-version-check", "--break-system-packages", package],
        files_dir, 180,
    )


TOOLS = {
    "run_python": tool_run_python,
    "run_shell": tool_run_shell,
    "write_file": tool_write_file,
    "read_file": tool_read_file,
    "list_files": tool_list_files,
    "delete_file": tool_delete_file,
    "make_zip": tool_make_zip,
    "pip_install": tool_pip_install,
}


def execute_tool(name: str, args: dict, files_dir: Path) -> Dict[str, Any]:
    fn = TOOLS.get(name)
    if not fn:
        return {"success": False, "stdout": "", "stderr": f"Unknown tool '{name}'. Available: {', '.join(TOOLS)}"}
    if not isinstance(args, dict):
        return {"success": False, "stdout": "", "stderr": "Tool 'args' must be a JSON object."}
    try:
        return fn(args, files_dir)
    except Exception as e:  # noqa: BLE001 -- surface tool bugs to the model instead of crashing the turn
        return {"success": False, "stdout": "", "stderr": f"Tool crashed: {e}"}


# --------------------------------------------------------------------------
# Tool-call parsing
# --------------------------------------------------------------------------

_TAGGED_BLOCK_RE = re.compile(r"```tool_call\s*\n(.*?)```", re.DOTALL | re.IGNORECASE)
_ANY_BLOCK_RE = re.compile(r"```[a-zA-Z_-]*\s*\n(.*?)```", re.DOTALL)


def find_tool_call(text: str) -> Tuple[Optional[dict], Optional[str], str]:
    """Look for a tool call in the model's raw turn text.

    Returns (call, parse_error, display_text):
      - call: {"tool": str, "args": dict} if a valid call was found, else None
      - parse_error: a message to feed back to the model if a tool_call-shaped
        block was found but was malformed, else None
      - display_text: the turn text with the raw tool-call block removed, i.e.
        whatever visible reasoning the model wrote for this turn
    """
    m = _TAGGED_BLOCK_RE.search(text)
    if m:
        raw = m.group(1).strip()
        display = (text[: m.start()] + text[m.end():]).strip()
        try:
            data = json.loads(raw)
        except json.JSONDecodeError as e:
            return (
                None,
                f"Your tool_call block wasn't valid JSON ({e}). Send a single "
                f'```tool_call block with valid JSON: {{"tool": "...", "args": {{...}}}}.',
                display,
            )
        if not isinstance(data, dict) or "tool" not in data:
            return None, "Your tool_call JSON must be an object with a 'tool' key and an 'args' object.", display
        return {"tool": data.get("tool"), "args": data.get("args") or {}}, None, display

    # Fallback: tolerate a mis-tagged fenced block (e.g. ```json) as long as
    # it parses into the same shape -- models sometimes drift on the tag.
    for bm in _ANY_BLOCK_RE.finditer(text):
        raw = bm.group(1).strip()
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if isinstance(data, dict) and data.get("tool") in TOOLS:
            display = (text[: bm.start()] + text[bm.end():]).strip()
            return {"tool": data.get("tool"), "args": data.get("args") or {}}, None, display

    return None, None, text.strip()


# --------------------------------------------------------------------------
# System prompt
# --------------------------------------------------------------------------

BASE_SYSTEM_PROMPT = """You are a personal autonomous coding assistant running inside the user's own private, isolated Linux workspace (a Hugging Face Space container only this user can reach). You reason, write code, run it for real, read the results, fix bugs, and hand back finished files -- the same kind of workflow Claude Code or ChatGPT's code interpreter use.

WORKSPACE
All tool paths are relative to your private working directory. Anything you write there can be downloaded by the user directly from this chat. Nothing you do here is visible to anyone else.

TOOL PROTOCOL
To take an action, reply with ONLY one fenced block, exactly like this, and nothing else in that message:

```tool_call
{"tool": "<name>", "args": {...}}
```

Wait for the result before deciding what to do next. Only ONE tool call per message -- never bundle several. If you ever want to show what a tool call looks like as an example without actually running it, do not use a real ```tool_call fence for that -- describe it in words instead, since a real fenced block will actually execute.

AVAILABLE TOOLS
- run_python -- {"code": "...", "timeout": optional seconds (default 60)}. Runs Python 3 in your workspace.
- run_shell -- {"command": "...", "timeout": optional seconds (default 60)}. Runs a bash command in your workspace.
- write_file -- {"path": "relative/path.ext", "content": "..."}. Creates or overwrites a file.
- read_file -- {"path": "relative/path.ext"}.
- list_files -- {"path": optional relative dir, default "."}.
- delete_file -- {"path": "relative/path.ext"}.
- make_zip -- {"output": "name.zip", "paths": optional list of relative paths -- omit to zip everything}.
- pip_install -- {"package": "name"}. Installs a Python package for use in run_python.

WORKING STYLE
1. Before telling the user a piece of code or a file is finished, actually run it with run_python or run_shell and check the output. If it errors, read the error, fix it, and run it again. Keep going until it genuinely works, or you've made a real effort and can explain what's blocking it.
2. When the deliverable is more than a one-liner, or is a set of files, package it with make_zip and name the file in your final answer so the user can download it.
3. When there is nothing left to run, just answer normally in plain text/Markdown with NO tool_call block. That ends your turn and is what the user sees as your reply.
4. You have a limited number of tool calls for this message, so work efficiently rather than exploring aimlessly.
5. Never try to damage or escape the container; stay inside your own workspace, and treat destructive system-wide commands as off-limits.
6. For ordinary conversation that needs no code or files, just answer directly -- you don't have to use a tool just because you can.
"""


def build_system_prompt(user_extra: Optional[str]) -> str:
    if user_extra and user_extra.strip():
        return BASE_SYSTEM_PROMPT + "\n\nADDITIONAL INSTRUCTIONS FROM THE USER:\n" + user_extra.strip()
    return BASE_SYSTEM_PROMPT


def format_tool_result_message(tool: str, result: Dict[str, Any]) -> str:
    lines = [f"[TOOL RESULT: {tool}]", f"success: {result.get('success')}"]
    if result.get("exit_code") is not None:
        lines.append(f"exit_code: {result['exit_code']}")
    stdout = result.get("stdout") or ""
    stderr = result.get("stderr") or ""
    if stdout:
        lines.append("stdout:\n" + stdout)
    if stderr:
        lines.append("stderr:\n" + stderr)
    if not stdout and not stderr:
        lines.append("(no output)")
    return "\n".join(lines)