closecode-ai 0.1.0__py3-none-any.whl

harness.py

@@ -0,0 +1,200 @@
+
+import os
+import signal
+import subprocess
+from pathlib import Path
+from typing import Callable, Optional
+
+from guardrails import check_bash_command, check_write_content
+
+
+class PermissionDenied(Exception):
+    pass
+
+
+class Harness:
+    def __init__(
+        self,
+        workdir: str,
+        auto_approve: bool = False,
+        confirm_fn: Optional[Callable[[str], str]] = None,
+    ):
+        self.workdir = Path(workdir).resolve()
+        self.workdir.mkdir(parents=True, exist_ok=True)
+        self.auto_approve = auto_approve
+        # Defaults to plain input() if no styled confirm function is given —
+        # keeps Harness usable standalone without depending on ui.py.
+        self._confirm_fn = confirm_fn or (
+            lambda action: "allow" if input(f"\n[permission] Allow agent to {action}? [y/N] ").strip().lower() in ("y", "yes") else "deny"
+        )
+
+    def confirm(self, action_description: str) -> bool:
+        if self.auto_approve:
+            return True
+        result = self._confirm_fn(action_description)
+        if result == "always":
+            self.auto_approve = True
+            return True
+        return result == "allow"
+
+    def resolve_path(self, relative_path: str) -> Path:
+        # Normalize before joining. Models very commonly pass absolute-looking
+        # paths (e.g. "/my_project/main.py" — they're thinking "project root",
+        # not "host filesystem root"). That's a real problem with pathlib:
+        # Path("/sandbox") / "/my_project/main.py" DISCARDS the left side
+        # entirely (joining an absolute path onto another replaces it), so
+        # the join silently points outside the sandbox instead of into it.
+        # This is exactly what caused write_file to appear to create files
+        # while bash (which correctly uses the same self.workdir) saw an
+        # empty directory — write_file's target had quietly become
+        # `/my_project/main.py` on the host, not `<workdir>/my_project/main.py`.
+        #
+        # Fix: strip any leading slash/backslash and any Windows drive
+        # prefix before joining, so every path is treated as relative to
+        # the sandbox no matter how the model phrased it. bash and every
+        # file tool then agree on the exact same root in every case.
+        cleaned = relative_path.replace("\\", "/").lstrip("/")
+        if len(cleaned) > 1 and cleaned[1] == ":":  # e.g. "C:/foo" or "C:foo"
+            cleaned = cleaned[2:].lstrip("/")
+
+        target = (self.workdir / cleaned).resolve()
+        if self.workdir != target and self.workdir not in target.parents:
+            raise PermissionDenied(f"Path '{relative_path}' escapes the sandboxed working directory.")
+        return target
+
+    def run_bash(self, command: str, timeout: int = 30) -> str:
+        reason = check_bash_command(command)
+        if reason:
+            return (
+                f"Guardrail blocked this command: {reason}. It looks destructive or "
+                "malicious, so it cannot run \u2014 even if approval were granted. "
+                "Rephrase the command to do the same thing safely, or use a "
+                "different approach."
+            )
+        if not self.confirm(f"run: `{command}`"):
+            return "Permission denied by user."
+        proc = subprocess.Popen(
+            command,
+            shell=True,
+            cwd=self.workdir,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            start_new_session=True,
+        )
+        try:
+            output, _ = proc.communicate(timeout=timeout)
+            output = (output or "").strip() or "(no output)"
+            return output[-4000:]
+        except subprocess.TimeoutExpired:
+            try:
+                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+            proc.communicate()
+            return (
+                f"Command timed out after {timeout}s and was terminated (including any "
+                f"child processes it started). If this was meant to be a long-running "
+                f"process (a server, a watcher), run it in the background instead — see "
+                f"the system prompt for how."
+            )
+
+    def read_file(self, path: str) -> str:
+        try:
+            target = self.resolve_path(path)
+        except PermissionDenied as e:
+            return str(e)
+        if not target.exists():
+            return f"File not found: {path}"
+        if not target.is_file():
+            return f"Not a file: {path}"
+        return target.read_text(errors="replace")[:8000]
+
+    def write_file(self, path: str, content: str) -> str:
+        reason = check_write_content(content)
+        if reason:
+            return (
+                f"Guardrail refused to write this content: {reason} (malicious-code "
+                "indicator). Policy forbids creating malware or exploit code "
+                "\u2014 even inside the sandbox."
+            )
+        try:
+            target = self.resolve_path(path)
+        except PermissionDenied as e:
+            return str(e)
+        if not self.confirm(f"write {len(content)} chars to `{path}`"):
+            return "Permission denied by user."
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(content)
+        return f"Wrote {len(content)} chars to {path}"
+
+    def list_dir(self, path: str = ".") -> str:
+        try:
+            target = self.resolve_path(path)
+        except PermissionDenied as e:
+            return str(e)
+        if not target.exists():
+            return f"Path not found: {path}"
+        if not target.is_dir():
+            return f"Not a directory: {path}"
+        entries = sorted(target.iterdir(), key=lambda p: (p.is_file(), p.name))
+        lines = []
+        for entry in entries:
+            marker = "/" if entry.is_dir() else ""
+            lines.append(f"{entry.name}{marker}")
+        return "\n".join(lines) if lines else "(empty directory)"
+
+    def edit_file(self, path: str, old_text: str, new_text: str) -> str:
+        """Targeted find-and-replace — far cheaper in tokens than rewriting
+        a whole file via write_file, and safer since it fails loudly if the
+        anchor text isn't found or isn't unique, instead of silently
+        clobbering unrelated content."""
+        reason = check_write_content(new_text)
+        if reason:
+            return (
+                f"Guardrail refused this edit: {reason} (malicious-code "
+                "indicator). Policy forbids creating malware or exploit code."
+            )
+        try:
+            target = self.resolve_path(path)
+        except PermissionDenied as e:
+            return str(e)
+        if not target.exists():
+            return f"File not found: {path}"
+        content = target.read_text(errors="replace")
+        count = content.count(old_text)
+        if count == 0:
+            return f"'old_text' not found in {path}. No changes made — check for exact whitespace/formatting differences."
+        if count > 1:
+            return f"'old_text' appears {count} times in {path}. Make it more specific so the edit is unambiguous. No changes made."
+        if not self.confirm(f"replace one occurrence of text in `{path}`"):
+            return "Permission denied by user."
+        updated = content.replace(old_text, new_text, 1)
+        target.write_text(updated)
+        return f"Replaced 1 occurrence in {path}."
+
+    def run_tests(self, command: str = "pytest", timeout: int = 60) -> str:
+        """Runs a test command (default: pytest) and returns its output.
+        Separate from run_bash mainly so the model has a clearly-named,
+        single-purpose action for 'verify my work' rather than free-form
+        shell access every time."""
+        reason = check_bash_command(command)
+        if reason:
+            return f"Guardrail blocked this test command: {reason}. It looks destructive or malicious and cannot run."
+        if not self.confirm(f"run tests with: `{command}`"):
+            return "Permission denied by user."
+        try:
+            result = subprocess.run(
+                command,
+                shell=True,
+                cwd=self.workdir,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+            output = (result.stdout or "") + (result.stderr or "")
+            output = output.strip() or "(no output)"
+            status = "PASSED" if result.returncode == 0 else f"FAILED (exit code {result.returncode})"
+            return f"{status}\n\n{output[-4000:]}"
+        except subprocess.TimeoutExpired:
+            return f"Test command timed out after {timeout}s."

llm.py

@@ -0,0 +1,145 @@
+import os
+from langchain_openrouter import ChatOpenRouter
+
+import json
+import time
+
+import requests
+
+# Exposed so main.py's banner can show the real model name instead of
+# "unknown" — previously main.py only checked HF_MODEL_ID/OPENROUTER_MODEL
+# env vars, but this file hardcodes the model directly, so neither existed.
+DEFAULT_MODEL = "nvidia/nemotron-3.5-lightning:free"
+
+# Curated starter list shown by the /models command. Each entry is
+# (openrouter model id, short note). This is just shortcuts — the user can
+# always switch to any other OpenRouter model id with `/model <id>`.
+# Free-tier availability and pricing change over time; when in doubt check
+# https://openrouter.ai/models before picking a paid one.
+KNOWN_MODELS = [
+    ("nvidia/nemotron-3.5-lightning:free", "default · free · fast, decent tool calling"),
+    ("qwen/qwen-2.5-coder-32b-instruct:free", "free · code-specialized, solid tool use"),
+    ("meta-llama/llama-3.1-8b-instruct:free", "free · fastest, weakest tool calling"),
+    ("qwen/qwen-2.5-72b-instruct", "paid · strong, reliable tool calling"),
+    ("meta-llama/llama-3.1-70b-instruct", "paid · strong, reliable tool calling"),
+]
+
+
+def resolve_model_arg(arg: str, choices: list) -> str:
+    """Turn a /model argument into a model id. A number picks from `choices`
+    (1-based, as shown by the last /models listing); anything else is
+    treated as a raw OpenRouter model id and passed through unchanged."""
+    arg = arg.strip()
+    if arg.isdigit():
+        idx = int(arg) - 1
+        if 0 <= idx < len(choices):
+            return choices[idx][0]
+    return arg
+
+
+# ---------------------------------------------------------------------------
+# Live model list from OpenRouter
+# ---------------------------------------------------------------------------
+
+_MODELS_ENDPOINT = "https://openrouter.ai/api/v1/models"
+_MODELS_CACHE_PATH = os.path.join(
+    os.path.expanduser("~"), ".cache", "closecode", "openrouter_models.json"
+)
+_MODELS_CACHE_TTL = 24 * 3600  # seconds
+
+
+def _load_cached_models():
+    """Return the cached [(id, note)] list if it's still fresh, else None."""
+    try:
+        with open(_MODELS_CACHE_PATH) as f:
+            data = json.load(f)
+        if time.time() - data.get("fetched_at", 0) < _MODELS_CACHE_TTL:
+            models = data.get("models")
+            if models:
+                return [tuple(m) for m in models]
+    except Exception:
+        pass
+    return None
+
+
+def _save_cached_models(models: list) -> None:
+    try:
+        os.makedirs(os.path.dirname(_MODELS_CACHE_PATH), exist_ok=True)
+        with open(_MODELS_CACHE_PATH, "w") as f:
+            json.dump({"fetched_at": time.time(), "models": models}, f)
+    except Exception:
+        pass
+
+
+def _note_for(entry: dict) -> str:
+    """Short human note: display name + free or $/M input pricing."""
+    name = entry.get("name") or entry.get("id", "")
+    pricing = entry.get("pricing") or {}
+    try:
+        prompt_per_token = float(pricing.get("prompt") or 0)
+    except (TypeError, ValueError):
+        prompt_per_token = 0
+    if prompt_per_token == 0:
+        return f"{name} · free"
+    return f"{name} · ${prompt_per_token * 1e6:.2f}/M"
+
+
+def fetch_openrouter_models(force_refresh: bool = False):
+    """Return (models, source) where models is [(id, note)] and source is
+    one of "live", "cache", "fallback".
+
+    Tries the 24h disk cache first, then the OpenRouter API (no auth
+    needed for the public models endpoint). Anything failing — no network,
+    bad response — falls back to the hardcoded KNOWN_MODELS shortlist so
+    /models always shows something useful.
+    """
+    if not force_refresh:
+        cached = _load_cached_models()
+        if cached:
+            return cached, "cache"
+    try:
+        resp = requests.get(_MODELS_ENDPOINT, timeout=15)
+        resp.raise_for_status()
+        items = resp.json().get("data") or []
+        models = []
+        for entry in items:
+            mid = entry.get("id") or ""
+            if not mid:
+                continue
+            try:
+                free = float((entry.get("pricing") or {}).get("prompt") or 0) == 0
+            except (TypeError, ValueError):
+                free = False
+            models.append((mid, _note_for(entry), free))
+        # Free models first (cheapest to try), then alphabetical — stable,
+        # predictable ordering so list numbers don't shuffle randomly.
+        models.sort(key=lambda t: (not t[2], t[0].lower()))
+        models = [(mid, note) for mid, note, _ in models]
+        if models:
+            _save_cached_models(models)
+            return models, "live"
+    except Exception:
+        pass
+    return list(KNOWN_MODELS), "fallback"
+
+
+def get_llm(model_override: str = None):
+
+    token = os.environ.get("OPENROUTER_API_KEY")
+
+    if not token:
+        raise EnvironmentError(
+            "OPENROUTER_API_KEY is not set. "
+            "Get a key at https://openrouter.ai/settings/keys"
+        )
+
+    model = model_override or os.environ.get("OPENROUTER_MODEL", DEFAULT_MODEL)
+
+    llm = ChatOpenRouter(
+        model=model,
+        temperature=0.2,
+        max_tokens=4096,
+        api_key=token,
+    )
+
+    return llm