closecode-ai 0.1.0__py3-none-any.whl

search.py

@@ -0,0 +1,130 @@
+"""Code search tools: glob (find files by pattern) and grep (search contents).
+
+Read-only, so both are available in plan mode too — unlike bash, which plan
+mode blocks entirely. Skips dependency/build directories so results stay
+relevant and fast.
+"""
+
+import fnmatch
+import os
+import re
+from pathlib import Path
+from typing import Optional
+
+from langchain_core.tools import tool
+
+_root: Optional[Path] = None
+
+# Directories never descended into during search.
+SKIP_DIRS = {
+    ".git",
+    "node_modules",
+    "__pycache__",
+    ".venv",
+    "venv",
+    ".cache",
+    "dist",
+    "build",
+    ".idea",
+    ".vscode",
+    "target",
+}
+
+MAX_RESULTS = 50
+MAX_LINE_LEN = 160
+
+
+def bind_search_root(path: str) -> None:
+    """Call once at startup — searches are confined to this root."""
+    global _root
+    _root = Path(path).resolve()
+
+
+def _require_root() -> Path:
+    if _root is None:
+        raise RuntimeError("Search root not bound. Call bind_search_root() before running the agent.")
+    _root.mkdir(parents=True, exist_ok=True)
+    return _root
+
+
+def _walk_files(base: Path):
+    """Yield files under base, pruning skipped directories."""
+    for dirpath, dirnames, filenames in os.walk(base):
+        dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")]
+        for name in filenames:
+            if name.startswith("."):
+                continue
+            yield Path(dirpath) / name
+
+
+@tool
+def glob(pattern: str, path: str = ".") -> str:
+    """Find files matching a glob pattern (e.g. "**/*.py", "src/*.ts").
+    Searches under path (default: working directory root), skipping
+    dependency and build directories. Prefer this over bash find/ls."""
+    root = _require_root()
+    base = (root / path).resolve()
+    if root not in base.parents and base != root:
+        return f"Path '{path}' escapes the working directory."
+    if not base.is_dir():
+        return f"Not a directory: {path}"
+    matches = []
+    for f in _walk_files(base):
+        rel = f.relative_to(base).as_posix()
+        if fnmatch.fnmatchcase(rel, pattern) or fnmatch.fnmatchcase(f.name, pattern):
+            matches.append(rel)
+            if len(matches) >= MAX_RESULTS:
+                break
+    matches.sort()
+    if not matches:
+        return f"No files match '{pattern}' under {path}."
+    suffix = f"\n…truncated at {MAX_RESULTS}, narrow the pattern." if len(matches) >= MAX_RESULTS else ""
+    return "\n".join(matches) + suffix
+
+
+@tool
+def grep(pattern: str, path: str = ".", include: str = "*", max_results: int = 50) -> str:
+    """Search file contents with a regex pattern. path: where to search
+    (default: working directory root). include: file glob to restrict to
+    (e.g. "*.py"). Returns file:line: match, newest relevant first —
+    prefer this over bash grep/rg, and it's available in plan mode."""
+    root = _require_root()
+    base = (root / path).resolve()
+    if root not in base.parents and base != root:
+        return f"Path '{path}' escapes the working directory."
+    if not base.is_dir():
+        return f"Not a directory: {path}"
+    try:
+        rx = re.compile(pattern)
+    except re.error as e:
+        return f"Invalid regex: {e}"
+    safe_max = max(1, min(int(max_results), 200))
+    hits = []
+    for f in _walk_files(base):
+        if not (fnmatch.fnmatchcase(f.relative_to(base).as_posix(), include)
+                or fnmatch.fnmatchcase(f.name, include)):
+            continue
+        try:
+            text = f.read_text(errors="strict")
+        except Exception:
+            continue  # binary or unreadable — skip
+        if "\x00" in text[:8192]:
+            continue
+        for i, line in enumerate(text.splitlines(), 1):
+            if rx.search(line):
+                rel = f.relative_to(root).as_posix()
+                line = line.strip()
+                if len(line) > MAX_LINE_LEN:
+                    line = line[:MAX_LINE_LEN] + "…"
+                hits.append(f"{rel}:{i}: {line}")
+                if len(hits) >= safe_max:
+                    break
+        if len(hits) >= safe_max:
+            break
+    if not hits:
+        return f"No matches for '{pattern}' under {path}."
+    suffix = f"\n…truncated at {safe_max}, narrow the pattern." if len(hits) >= safe_max else ""
+    return "\n".join(hits) + suffix
+
+
+SEARCH_TOOLS = [glob, grep]

session.py

@@ -0,0 +1,197 @@
+
+import json
+import sqlite3
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from langchain_core.messages import messages_from_dict, messages_to_dict
+
+SESSIONS_DIR = Path("./sessions")
+DB_PATH = SESSIONS_DIR / "sessions.db"
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS sessions (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    name TEXT NOT NULL DEFAULT '',
+    model TEXT NOT NULL DEFAULT '',
+    mode TEXT NOT NULL DEFAULT 'build',
+    summary TEXT NOT NULL DEFAULT '',
+    created_at REAL NOT NULL,
+    updated_at REAL NOT NULL,
+    message_count INTEGER NOT NULL DEFAULT 0
+);
+
+CREATE TABLE IF NOT EXISTS messages (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    session_id INTEGER NOT NULL,
+    seq INTEGER NOT NULL,
+    payload TEXT NOT NULL,
+    FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE,
+    UNIQUE (session_id, seq)
+);
+
+CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, seq);
+"""
+
+
+@dataclass
+class SessionInfo:
+    id: int
+    name: str
+    model: str
+    mode: str
+    summary: str
+    created_at: float
+    updated_at: float
+    message_count: int
+
+
+def _connect() -> sqlite3.Connection:
+    conn = sqlite3.connect(DB_PATH)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def _init() -> None:
+    SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
+    with _connect() as conn:
+        conn.executescript(_SCHEMA)
+
+
+def _migrate_from_json() -> None:
+    """Fold any legacy session_*.json files into the DB once, then leave them."""
+    with _connect() as conn:
+        count = conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0]
+        if count:
+            return
+    files = sorted(SESSIONS_DIR.glob("session_*.json"))
+    for path in files:
+        try:
+            msgs = messages_from_dict(json.loads(path.read_text()))
+        except Exception:
+            continue
+        if not msgs:
+            continue
+        mtime = path.stat().st_mtime
+        with _connect() as conn:
+            cur = conn.execute(
+                "INSERT INTO sessions (name, model, mode, summary, created_at, updated_at, message_count) "
+                "VALUES (?, ?, ?, ?, ?, ?, ?)",
+                ("", "", "build", "", mtime, mtime, len(msgs)),
+            )
+            sid = cur.lastrowid
+            for seq, m in enumerate(msgs):
+                conn.execute(
+                    "INSERT INTO messages (session_id, seq, payload) VALUES (?, ?, ?)",
+                    (sid, seq, json.dumps(messages_to_dict([m])[0])),
+                )
+            conn.commit()
+        name = _derive_name(msgs)
+        if name:
+            conn.execute("UPDATE sessions SET name = ? WHERE id = ?", (name, sid))
+            conn.commit()
+
+
+def _derive_name(messages: list) -> str:
+    """Short, human-readable name from the first user message."""
+    for m in messages:
+        if getattr(m, "type", "") == "human":
+            text = str(getattr(m, "content", "") or "").strip()
+            text = " ".join(text.split())
+            return text[:48]
+    return ""
+
+
+def new_session(model: str = "", mode: str = "build") -> int:
+    """Create a new session row and return its id."""
+    now = time.time()
+    with _connect() as conn:
+        cur = conn.execute(
+            "INSERT INTO sessions (model, mode, created_at, updated_at, message_count) VALUES (?, ?, ?, ?, 0)",
+            (model, mode, now, now),
+        )
+        conn.commit()
+        return int(cur.lastrowid)
+
+
+def save(session_id: int, messages: list, model: str = "", mode: str = "build") -> None:
+    """Replace a session's stored messages and refresh its metadata."""
+    payloads = [json.dumps(d) for d in messages_to_dict(messages)]
+    now = time.time()
+    with _connect() as conn:
+        conn.execute(
+            "UPDATE sessions SET model = ?, mode = ?, updated_at = ?, message_count = ? WHERE id = ?",
+            (model, mode, now, len(payloads), session_id),
+        )
+        conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
+        conn.executemany(
+            "INSERT INTO messages (session_id, seq, payload) VALUES (?, ?, ?)",
+            [(session_id, seq, p) for seq, p in enumerate(payloads)],
+        )
+        row = conn.execute("SELECT name FROM sessions WHERE id = ?", (session_id,)).fetchone()
+        if row and not row["name"]:
+            name = _derive_name(messages)
+            if name:
+                conn.execute("UPDATE sessions SET name = ? WHERE id = ?", (name, session_id))
+        conn.commit()
+
+
+def load(session_id: int) -> list:
+    """Rehydrate a session's messages in order. Skips any corrupt rows."""
+    with _connect() as conn:
+        rows = conn.execute(
+            "SELECT payload FROM messages WHERE session_id = ? ORDER BY seq", (session_id,)
+        ).fetchall()
+    payloads = []
+    for row in rows:
+        try:
+            payloads.append(json.loads(row["payload"]))
+        except Exception:
+            continue
+    return messages_from_dict(payloads)
+
+
+def get_session(session_id: int) -> Optional[SessionInfo]:
+    with _connect() as conn:
+        row = conn.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
+    if row is None:
+        return None
+    return SessionInfo(**{k: row[k] for k in row.keys()})
+
+
+def list_sessions() -> list[SessionInfo]:
+    with _connect() as conn:
+        rows = conn.execute(
+            "SELECT * FROM sessions ORDER BY updated_at DESC"
+        ).fetchall()
+    return [SessionInfo(**{k: row[k] for k in row.keys()}) for row in rows]
+
+
+def latest_session_id() -> Optional[int]:
+    with _connect() as conn:
+        row = conn.execute(
+            "SELECT id FROM sessions WHERE message_count > 0 ORDER BY updated_at DESC LIMIT 1"
+        ).fetchone()
+    return int(row["id"]) if row else None
+
+
+def delete_session(session_id: int) -> bool:
+    with _connect() as conn:
+        cur = conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
+        conn.commit()
+        return cur.rowcount > 0
+
+
+def rename_session(session_id: int, name: str) -> bool:
+    with _connect() as conn:
+        cur = conn.execute(
+            "UPDATE sessions SET name = ? WHERE id = ?", (name.strip()[:80], session_id)
+        )
+        conn.commit()
+        return cur.rowcount > 0
+
+
+_init()
+_migrate_from_json()

todos.py

@@ -0,0 +1,99 @@
+"""Todo tracking for the agent (Claude Code-style TodoWrite/TodoRead).
+
+Todos live in memory, scoped to the current session — main.py clears the
+store whenever the session changes (/clear, /resume, new session). The
+pattern mirrors tools.py: a module-global store bound once at startup via
+bind_todo_store().
+"""
+
+from typing import Optional
+
+from langchain_core.tools import tool
+
+_store: Optional["TodoStore"] = None
+
+VALID_STATUSES = ("pending", "in_progress", "completed")
+
+
+class TodoStore:
+    def __init__(self):
+        self.items: list[dict] = []
+
+    def set(self, items: list[dict]) -> list[dict]:
+        """Replace the whole list. Sanitizes entries: drops empties,
+        coerces bad statuses to pending, and enforces exactly one
+        in_progress at a time (extras are demoted to pending)."""
+        cleaned = []
+        for it in items:
+            if not isinstance(it, dict):
+                continue
+            content = str(it.get("content", "")).strip()
+            if not content:
+                continue
+            status = it.get("status", "pending")
+            if status not in VALID_STATUSES:
+                status = "pending"
+            cleaned.append(
+                {
+                    "content": content[:200],
+                    "status": status,
+                    "activeForm": str(it.get("activeForm", ""))[:200],
+                }
+            )
+        seen_active = False
+        for it in cleaned:
+            if it["status"] == "in_progress":
+                if seen_active:
+                    it["status"] = "pending"
+                seen_active = True
+        self.items = cleaned
+        return self.items
+
+    def get(self) -> list[dict]:
+        return list(self.items)
+
+    def clear(self) -> None:
+        self.items = []
+
+
+def bind_todo_store(store: TodoStore) -> None:
+    """Call once at startup before the graph runs any tool calls."""
+    global _store
+    _store = store
+
+
+def _require_store() -> TodoStore:
+    if _store is None:
+        raise RuntimeError("TodoStore not bound. Call bind_todo_store() before running the agent.")
+    return _store
+
+
+def format_todos(items: list[dict]) -> str:
+    """One-line-per-todo text form, also used for tool results."""
+    if not items:
+        return "(no todos)"
+    icons = {"pending": "○", "in_progress": "◐", "completed": "✓"}
+    return "\n".join(
+        f"{icons.get(it['status'], '?')} {it['content']}" for it in items
+    )
+
+
+@tool
+def todo_write(todos: list[dict]) -> str:
+    """Create or update the agent's task list. Call with the FULL list every
+    time (it replaces the previous list). Each item: {"content": "do X",
+    "status": "pending"|"in_progress"|"completed", "activeForm": "Doing X"}.
+    Use for any task with 3+ steps: write the plan first, keep exactly one
+    item in_progress, flip each to completed as you finish it, and re-write
+    the list when the plan changes. This list is how the user sees progress."""
+    items = _require_store().set(todos if isinstance(todos, list) else [])
+    return "Todo list updated:\n" + format_todos(items)
+
+
+@tool
+def todo_read() -> str:
+    """Read the current task list with each item's status."""
+    return format_todos(_require_store().get())
+
+
+TODO_TOOLS = [todo_write, todo_read]

token_tracker.py

@@ -0,0 +1,42 @@
+
+from dataclasses import dataclass
+
+
+@dataclass
+class TokenTracker:
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    calls: int = 0
+    calls_without_usage_data: int = 0
+
+    def add_from_message(self, message) -> None:
+        usage = getattr(message, "usage_metadata", None)
+        if usage:
+            self.prompt_tokens += usage.get("input_tokens", 0) or 0
+            self.completion_tokens += usage.get("output_tokens", 0) or 0
+            self.total_tokens += usage.get("total_tokens", 0) or 0
+            self.calls += 1
+            return
+
+        meta = getattr(message, "response_metadata", None) or {}
+        tu = meta.get("token_usage") or meta.get("usage")
+        if tu:
+            p = tu.get("prompt_tokens", 0) or 0
+            c = tu.get("completion_tokens", 0) or 0
+            self.prompt_tokens += p
+            self.completion_tokens += c
+            self.total_tokens += tu.get("total_tokens", p + c) or (p + c)
+            self.calls += 1
+            return
+
+        self.calls += 1
+        self.calls_without_usage_data += 1
+
+    def summary(self) -> str:
+        if self.calls == 0:
+            return "no model calls yet this session"
+        line = f"{self.calls} model calls \u00b7 {self.prompt_tokens} in / {self.completion_tokens} out / {self.total_tokens} total tokens"
+        if self.calls_without_usage_data:
+            line += f"  ({self.calls_without_usage_data} calls reported no usage data \u2014 provider limitation, not a bug)"
+        return line

tools.py

@@ -0,0 +1,137 @@
+
+import os
+from typing import Optional
+
+import requests
+from langchain_core.tools import tool
+
+from harness import Harness
+
+_harness: Optional[Harness] = None
+
+
+def bind_harness(harness: Harness) -> None:
+    """Call this once at startup before the graph runs any tool calls."""
+    global _harness
+    _harness = harness
+
+
+def _require_harness() -> Harness:
+    if _harness is None:
+        raise RuntimeError("Harness not bound. Call bind_harness() before running the agent.")
+    return _harness
+
+
+@tool
+def bash(command: str, timeout: int = 30) -> str:
+    """Run a shell command in the sandboxed working directory and return its combined stdout/stderr.
+    timeout: max seconds to wait (default 30). Increase this for slow operations like
+    package installs or builds (e.g. timeout=120 for a two-minute cap) — the units are
+    SECONDS, not milliseconds. Capped at 300s (5 minutes) regardless of what's requested."""
+    safe_timeout = max(1, min(int(timeout), 300))
+    return _require_harness().run_bash(command, timeout=safe_timeout)
+
+
+@tool
+def read_file(path: str) -> str:
+    """Read a text file's contents. Path is relative to the agent's working directory."""
+    return _require_harness().read_file(path)
+
+
+@tool
+def write_file(path: str, content: str) -> str:
+    """Write text content to a file, creating parent directories if needed. Path is relative to the working directory."""
+    return _require_harness().write_file(path, content)
+
+
+@tool
+def list_dir(path: str = ".") -> str:
+    """List files and subdirectories at a given path (default: the working directory root). Directories are marked with a trailing '/'."""
+    return _require_harness().list_dir(path)
+
+
+@tool
+def edit_file(path: str, old_text: str, new_text: str) -> str:
+    """Replace one exact occurrence of old_text with new_text in a file. Use this instead of write_file for small changes — it's cheaper and fails safely if old_text isn't found or isn't unique, rather than risking an overwrite of the wrong content."""
+    return _require_harness().edit_file(path, old_text, new_text)
+
+
+@tool
+def run_tests(command: str = "pytest") -> str:
+    """Run the project's test suite (default command: 'pytest') and report PASSED/FAILED with output. Use this to verify a change actually works, not just that it was written."""
+    return _require_harness().run_tests(command)
+
+
+_TAVILY_API_URL = "https://api.tavily.com/search"
+
+
+@tool
+def tavily_search(query: str, max_results: int = 5) -> str:
+    """Search the web using Tavily and return the top results (title, url, content snippet, relevance score). Handy for checking current API docs, finding library versions, resolving error messages, or verifying facts that changed after your training cutoff. Requires a TAVILY_API_KEY in .env."""
+    api_key = os.environ.get("TAVILY_API_KEY")
+    if not api_key:
+        return "Error: TAVILY_API_KEY is not set. Add it to the .env file and restart."
+    safe_max = max(1, min(int(max_results), 10))
+    try:
+        resp = requests.post(
+            _TAVILY_API_URL,
+            json={"api_key": api_key, "query": query, "max_results": safe_max, "search_depth": "basic"},
+            timeout=30,
+        )
+        resp.raise_for_status()
+        data = resp.json()
+    except Exception as e:
+        return f"Tavily search failed: {e}"
+
+    results = data.get("results") or []
+    if not results:
+        return f"No results found for: {query}"
+    lines = [f"Query: {query}\n"]
+    for i, r in enumerate(results, 1):
+        title = r.get("title", "")
+        url = r.get("url", "")
+        content = (r.get("content") or "").strip()
+        lines.append(f"{i}. {title}\n   URL: {url}\n   {content}")
+    return "\n".join(lines)
+
+
+# Local (non-MCP) tools. main.py combines this with any MCP-provided tools
+# (e.g. git) before binding to the model.
+LOCAL_TOOLS = [bash, read_file, write_file, list_dir, edit_file, run_tests, tavily_search]
+
+# Tools considered safe in "plan" mode: read-only, no filesystem/shell
+# mutation. bash is excluded entirely even though some commands are
+# harmless (e.g. `ls`) — there's no reliable way to tell a read-only shell
+# command from a destructive one without actually parsing it, so plan mode
+# blocks bash outright rather than trying to guess.
+_PLAN_SAFE_LOCAL_NAMES = {"read_file", "list_dir", "tavily_search"}
+
+# Heuristic for filtering MCP tools (e.g. git) in plan mode: block anything
+# whose name suggests it mutates state. This is a name-based guess, not a
+# guarantee — if you add other MCP servers, sanity-check their tool names
+# fall into one of these buckets as expected.
+_MUTATING_KEYWORDS = ("commit", "push", "reset", "checkout", "branch", "merge", "add", "rm", "stash", "revert", "rebase", "delete", "write", "create")
+
+
+def filter_tools_for_mode(tools: list, mode: str) -> list:
+    """mode == 'build': every tool is available.
+    mode == 'plan': only read-only tools are available, so the agent can
+    explore and reason about a task but literally cannot call anything
+    that writes to disk, runs arbitrary shell, or mutates git state —
+    enforced by never binding those tools to the model at all, which is a
+    stronger guarantee than trusting the model to just not call them."""
+    if mode == "build":
+        return tools
+
+    filtered = []
+    for t in tools:
+        name = getattr(t, "name", "")
+        if name in _PLAN_SAFE_LOCAL_NAMES:
+            filtered.append(t)
+        elif name in {"bash", "write_file", "edit_file", "run_tests"}:
+            continue
+        elif not any(keyword in name.lower() for keyword in _MUTATING_KEYWORDS):
+            # Likely an MCP tool (e.g. git status/diff/log) that doesn't
+            # match a known-mutating keyword — allow it through.
+            filtered.append(t)
+    return filtered