closecode-ai 0.1.0__py3-none-any.whl

main.py

@@ -0,0 +1,520 @@
+
+
+import asyncio
+import contextlib
+import os
+import sys
+import time
+import warnings
+
+warnings.filterwarnings(
+    "ignore",
+    message="Core Pydantic V1 functionality isn't compatible with Python 3.14",
+    category=UserWarning,
+)
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from langchain_core.messages import HumanMessage, SystemMessage  # noqa: E402
+
+from langgraph.errors import GraphRecursionError  # noqa: E402
+
+import session  # noqa: E402
+import ui  # noqa: E402
+from agent import SYSTEM_PROMPT, build_graph  # noqa: E402
+from guardrails import check_user_input, redact_message  # noqa: E402
+from harness import Harness  # noqa: E402
+from llm import DEFAULT_MODEL, KNOWN_MODELS, fetch_openrouter_models, get_llm, resolve_model_arg  # noqa: E402
+from mcp_tools import get_git_repo_path, get_git_tools  # noqa: E402
+from modes import filter_tools_for_mode, mode_system_note  # noqa: E402
+from search import SEARCH_TOOLS, bind_search_root  # noqa: E402
+from todos import TODO_TOOLS, TodoStore, bind_todo_store  # noqa: E402
+from token_tracker import TokenTracker  # noqa: E402
+from tools import LOCAL_TOOLS, bind_harness  # noqa: E402
+
+
+# Module-global todo store, bound to the todo tools once at startup and
+# cleared whenever the session changes (new session, /resume, /clear).
+todo_store = TodoStore()
+
+
+def system_message_for(mode: str) -> SystemMessage:
+    return SystemMessage(content=SYSTEM_PROMPT + mode_system_note(mode))
+
+
+def set_system_message(messages: list, mode: str) -> None:
+    """Ensure a session's first message is the current system prompt, in the
+    right mode note, without clobbering a non-system first message."""
+    if messages and getattr(messages[0], "type", "") == "system":
+        messages[0] = system_message_for(mode)
+    else:
+        messages.insert(0, system_message_for(mode))
+
+
+def _looks_like_tool_json(text: str) -> bool:
+    """True when a streamed buffer is a tool call the model is writing out as
+    JSON (qwen2.5-coder via Ollama does this instead of native tool_calls).
+    Such content is re-parsed into a real tool call in agent.py, so we keep
+    it out of the chat panel entirely."""
+    stripped = text.lstrip()
+    return stripped.startswith("{") and '"name"' in stripped and '"arguments"' in stripped
+
+
+async def run_turn(graph, messages: list, token_tracker: TokenTracker) -> list:
+    """Streams one user turn via astream_events. Prints tool calls/results
+    as they happen, streams the final text response token-by-token into a
+    live-updating panel, and tracks token usage along the way.
+
+    Also races the event stream against an Esc keypress (watched on a
+    background thread by ui.EscListener) so the user can bail out of a turn
+    that's stuck, taking too long, or headed somewhere they don't want it
+    to go, without killing the whole process.
+    """
+    ui.print_thinking()
+
+    live = None
+    text_buffer = ""
+    final_messages = messages
+    turn_start = time.monotonic()
+    interrupted = False
+
+    loop = asyncio.get_running_loop()
+    stop_event = asyncio.Event()
+    esc = ui.EscListener(loop, stop_event)
+    esc.start()
+
+    agen = graph.astream_events({"messages": messages}, version="v2")
+
+    try:
+        while True:
+            next_task = asyncio.ensure_future(agen.__anext__())
+            stop_task = asyncio.ensure_future(stop_event.wait())
+            try:
+                done, _pending = await asyncio.wait(
+                    {next_task, stop_task}, return_when=asyncio.FIRST_COMPLETED
+                )
+            except asyncio.CancelledError:
+                next_task.cancel()
+                stop_task.cancel()
+                raise
+
+            if stop_task in done:
+                interrupted = True
+                next_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError, Exception):
+                    await next_task
+                with contextlib.suppress(Exception):
+                    await agen.aclose()
+                break
+
+            stop_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await stop_task
+
+            try:
+                event = next_task.result()
+            except StopAsyncIteration:
+                break
+
+            kind = event["event"]
+
+            if kind == "on_chat_model_stream":
+                chunk = event["data"]["chunk"]
+                if getattr(chunk, "content", None):
+                    text_buffer += chunk.content
+                    if _looks_like_tool_json(text_buffer):
+                        continue
+                    if live is None:
+                        live = ui.stream_start()
+                    ui.stream_update(live, text_buffer)
+
+            elif kind == "on_chat_model_end":
+                output = event["data"].get("output")
+                if output is not None:
+                    token_tracker.add_from_message(output)
+                    _blocked, reason = redact_message(output)
+                    if reason:
+                        ui.print_notice(
+                            f"Guardrail \u2014 blocked {reason}. The flagged content was "
+                            "removed from conversation history.",
+                            style="bold red",
+                        )
+                if live is not None:
+                    ui.stream_stop(live)
+                    live = None
+                text_buffer = ""
+
+            elif kind == "on_tool_start":
+                ui.print_tool_call(event["name"], event["data"].get("input") or {})
+
+            elif kind == "on_tool_end":
+                output = event["data"].get("output")
+                content = getattr(output, "content", None)
+                if content is None:
+                    content = str(output)
+                ui.print_tool_result(str(content))
+                if event.get("name") == "todo_write":
+                    ui.print_todos(todo_store.get())
+
+            elif kind == "on_chain_end" and event.get("name") == "LangGraph":
+                output = event["data"].get("output")
+                if output and "messages" in output:
+                    final_messages = output["messages"]
+
+    except GraphRecursionError:
+        # This used to be the visible symptom of the sandbox/workdir bug:
+        # write_file silently failing on an absolute-looking path, bash
+        # correctly showing an empty dir, and the model retrying the same
+        # broken step until the recursion cap kicked in. That root cause is
+        # fixed in harness.py now. If this still fires, it means the agent
+        # is genuinely stuck in a loop for some other reason (e.g. a
+        # command that keeps failing for a real, external reason) — surface
+        # that plainly instead of a raw traceback, rather than papering
+        # over it by just raising the limit.
+        ui.print_notice(
+            "Stopped: the agent hit the step limit for this turn without finishing "
+            "(likely repeating a failing action). Check the tool calls/results above "
+            "for what kept failing, then try again or rephrase the task.",
+            style="bold red",
+        )
+        return messages
+    except Exception as e:
+        detail = str(e) or repr(e)
+        cause = getattr(e, "__cause__", None)
+        if cause and str(cause) not in detail:
+            detail = f"{detail} (caused by: {cause})"
+        ui.print_notice(
+            f"Error during this turn [{type(e).__name__}]: {detail}", style="bold red"
+        )
+        return messages
+    finally:
+        esc.stop()
+        if live is not None:
+            ui.stream_stop(live)
+
+    if interrupted:
+        ui.print_notice(
+            "Interrupted (Esc) \u2014 turn stopped early; any tool call already in "
+            "flight may still finish on its own. History kept up to the last "
+            "completed step.",
+            style="yellow",
+        )
+        return final_messages
+
+    ui.print_turn_complete(time.monotonic() - turn_start)
+    return final_messages
+
+
+def save_key_to_dotenv(key: str) -> None:
+    """Persist an API key into .env (creating the file if needed), replacing
+    any existing OPENROUTER_API_KEY line. .env is gitignored, so the key
+    never lands in version control."""
+    path = ".env"
+    lines: list[str] = []
+    if os.path.exists(path):
+        with open(path) as f:
+            lines = f.read().splitlines()
+    updated = False
+    for i, line in enumerate(lines):
+        if line.strip().startswith("OPENROUTER_API_KEY="):
+            lines[i] = f"OPENROUTER_API_KEY={key}"
+            updated = True
+    if not updated:
+        if lines and lines[-1].strip():
+            lines.append("")
+        lines.append(f"OPENROUTER_API_KEY={key}")
+    with open(path, "w") as f:
+        f.write("\n".join(lines) + "\n")
+
+
+def ensure_api_key() -> str:
+    """Make sure an OpenRouter API key is available. If OPENROUTER_API_KEY
+    isn't set (env or .env), prompt the user to paste one — hidden input —
+    and offer to save it to .env for next time. Exits if no key is given,
+    since the agent can't call a model without one."""
+    key = os.environ.get("OPENROUTER_API_KEY", "").strip()
+    if key:
+        return key
+    ui.print_notice("No OPENROUTER_API_KEY found in env or .env.", style="yellow")
+    key = ui.prompt_api_key()
+    if not key:
+        ui.print_notice(
+            "No API key provided — the agent can't run without one. "
+            "Set OPENROUTER_API_KEY and restart.",
+            style="bold red",
+        )
+        raise SystemExit(1)
+    os.environ["OPENROUTER_API_KEY"] = key
+    if ui.confirm_save_key():
+        save_key_to_dotenv(key)
+        ui.print_notice("Saved to .env", style="green")
+    return key
+
+
+COMPACT_PROMPT = """You are summarizing a coding-assistant session so work can continue
+in a fresh context window. Write a dense, structured summary covering:
+
+1. What the user was trying to accomplish (the overall goal)
+2. What was actually done — files created/modified, commands run, key decisions
+3. Current state — what's working, what's unfinished or broken
+4. Anything the user explicitly asked to remember, plus useful context
+   (paths, model choices, config values) needed to continue seamlessly
+
+Be concrete: name files, functions, and decisions. Skip greetings and
+small talk. Write it so another agent could pick up exactly where this
+one left off."""
+
+
+async def compact_messages(messages: list, model_override: str = None) -> str:
+    """Summarize the conversation with the plain LLM (no tools) and return
+    the summary text. Long tool outputs are truncated so the summarizer
+    itself doesn't blow the context window."""
+    llm = get_llm(model_override)
+    lines = []
+    for m in messages:
+        role = getattr(m, "type", "?")
+        content = m.content if isinstance(m.content, str) else str(m.content)
+        if len(content) > 2000:
+            content = content[:2000] + "…[truncated]"
+        lines.append(f"[{role}] {content}")
+    convo = "\n\n".join(lines)
+    summary = await llm.ainvoke(
+        [SystemMessage(content=COMPACT_PROMPT), HumanMessage(content=convo)]
+    )
+    return summary.content if isinstance(summary.content, str) else str(summary.content)
+
+
+async def main():
+    ensure_api_key()
+
+    workdir = os.environ.get("AGENT_WORKDIR", "./sandbox")
+    auto_approve = os.environ.get("AGENT_AUTO_APPROVE", "false").lower() == "true"
+    use_git = os.environ.get("AGENT_ENABLE_GIT", "false").lower() == "true"
+
+    harness = Harness(workdir=workdir, auto_approve=auto_approve, confirm_fn=ui.confirm)
+    bind_harness(harness)
+    bind_todo_store(todo_store)
+    bind_search_root(str(harness.workdir))
+
+    all_tools = list(LOCAL_TOOLS) + SEARCH_TOOLS + TODO_TOOLS
+    if use_git:
+        repo_path = get_git_repo_path()
+        try:
+            git_tools = await get_git_tools(repo_path)
+            all_tools.extend(git_tools)
+        except Exception as e:
+            ui.print_notice(f"Could not load git MCP tools, continuing without them: {e}", style="yellow")
+
+    mode = "build"
+    model_override = None
+    token_tracker = TokenTracker()
+    # Models shown by the most recent /models listing — /model <number>
+    # resolves against this. Starts as the curated shortlist so numbers
+    # work even before /models is ever run.
+    listed_models: list = list(KNOWN_MODELS)
+
+    model_name = model_override or os.environ.get("HF_MODEL_ID") or os.environ.get("OLLAMA_MODEL", DEFAULT_MODEL)
+
+    # Start fresh, or --continue resume the most-recently-used non-empty session
+    # from the SQLite store (metadata restores its mode/model too).
+    resume = "--continue" in sys.argv
+    current_id = session.latest_session_id() if resume else None
+    if current_id is not None:
+        loaded = session.load(current_id)
+        if loaded:
+            info = session.get_session(current_id)
+            if info is not None:
+                mode = info.mode or mode
+                if info.model and not model_override:
+                    model_override = info.model
+                    model_name = info.model
+            ui.print_notice(
+                f"Resumed session #{current_id} ({info.name if info else '(unnamed)'}, {len(loaded)} msgs)",
+                style="cyan",
+            )
+        else:
+            current_id = None
+
+    if current_id is None:
+        current_id = session.new_session(model=model_name, mode=mode)
+        loaded = None
+
+    messages = list(loaded) if loaded is not None else [system_message_for(mode)]
+
+    graph = build_graph(filter_tools_for_mode(all_tools, mode), model_override)
+    ui.set_context(mode, model_name)
+    ui.print_banner(model_name, str(harness.workdir), [t.name for t in all_tools])
+
+    while True:
+        try:
+            raw = ui.user_prompt(mode)
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+
+        if not raw:
+            continue
+        if raw.lower() in {"exit", "quit"}:
+            break
+
+        if raw.startswith("/"):
+            parts = raw[1:].strip().split(maxsplit=1)
+            cmd = parts[0].lower() if parts else ""
+            arg = parts[1] if len(parts) > 1 else ""
+
+            if cmd == "plan":
+                mode = "plan"
+                graph = build_graph(filter_tools_for_mode(all_tools, mode), model_override)
+                set_system_message(messages, mode)
+                ui.set_context(mode, model_name)
+                ui.print_notice("Switched to PLAN mode \u2014 read-only tools only.", style="magenta")
+            elif cmd == "build":
+                mode = "build"
+                graph = build_graph(filter_tools_for_mode(all_tools, mode), model_override)
+                set_system_message(messages, mode)
+                ui.set_context(mode, model_name)
+                ui.print_notice("Switched to BUILD mode \u2014 all tools enabled.", style="blue")
+            elif cmd == "model":
+                if not arg:
+                    ui.print_notice("Usage: /model <number|model-id>  (see /models)", style="yellow")
+                else:
+                    new_model = resolve_model_arg(arg, listed_models)
+                    if arg.strip().isdigit() and new_model == arg.strip():
+                        ui.print_notice(
+                            f"No model #{arg.strip()} in the current list — run /models "
+                            "first (or pass a full OpenRouter model id).",
+                            style="yellow",
+                        )
+                    else:
+                        model_override = new_model
+                        model_name = new_model
+                        graph = build_graph(filter_tools_for_mode(all_tools, mode), model_override)
+                        ui.set_context(mode, model_name)
+                        ui.print_notice(f"Switched model to {new_model}", style="cyan")
+            elif cmd == "models":
+                parts = arg.split()
+                force = "--refresh" in parts
+                query = " ".join(p for p in parts if p != "--refresh").strip().lower()
+                ui.print_notice("Fetching model list from OpenRouter…", style="dim")
+                models, source = fetch_openrouter_models(force_refresh=force)
+                if query:
+                    models = [m for m in models
+                              if query in m[0].lower() or query in m[1].lower()]
+                    if not models:
+                        ui.print_notice(f"No models match '{query}'.", style="yellow")
+                        continue
+                listed_models = models
+                ui.print_models(models, model_name, source=source, query=query or None)
+            elif cmd == "key":
+                key = ui.prompt_api_key()
+                if not key:
+                    ui.print_notice("No key entered — keeping the current one.", style="yellow")
+                else:
+                    os.environ["OPENROUTER_API_KEY"] = key
+                    if ui.confirm_save_key():
+                        save_key_to_dotenv(key)
+                        ui.print_notice("Saved to .env", style="green")
+                    graph = build_graph(filter_tools_for_mode(all_tools, mode), model_override)
+                    ui.print_notice("API key updated.", style="green")
+            elif cmd == "sessions":
+                ui.print_sessions(session.list_sessions())
+            elif cmd == "resume":
+                if not arg:
+                    ui.print_notice("Usage: /resume <session-id>  (see /sessions)", style="yellow")
+                    continue
+                try:
+                    sid = int(arg)
+                except ValueError:
+                    ui.print_notice("Usage: /resume <session-id>  (see /sessions)", style="yellow")
+                    continue
+                info = session.get_session(sid)
+                loaded = session.load(sid) if info else None
+                if info is None or not loaded:
+                    ui.print_notice(f"No resumable session #{sid}. See /sessions.", style="yellow")
+                    continue
+                if info.mode:
+                    mode = info.mode
+                if info.model and not model_override:
+                    model_override = info.model
+                    model_name = info.model
+                graph = build_graph(filter_tools_for_mode(all_tools, mode), model_override)
+                current_id = sid
+                messages = list(loaded)
+                todo_store.clear()
+                set_system_message(messages, mode)
+                ui.set_context(mode, model_name)
+                ui.print_notice(
+                    f"Resumed session #{sid} ({info.name or '(unnamed)'}, {len(messages)} msgs)",
+                    style="cyan",
+                )
+            elif cmd == "delete":
+                if not arg:
+                    ui.print_notice("Usage: /delete <session-id>  (see /sessions)", style="yellow")
+                    continue
+                try:
+                    sid = int(arg)
+                except ValueError:
+                    ui.print_notice("Usage: /delete <session-id>  (see /sessions)", style="yellow")
+                    continue
+                if sid == current_id:
+                    ui.print_notice("Can't delete the active session. Resume a different one first.", style="yellow")
+                    continue
+                if session.delete_session(sid):
+                    ui.print_notice(f"Deleted session #{sid}.", style="yellow")
+                else:
+                    ui.print_notice(f"No session #{sid}. See /sessions.", style="yellow")
+            elif cmd == "usage":
+                ui.print_token_usage(token_tracker.summary())
+            elif cmd == "clear":
+                messages = [system_message_for(mode)]
+                todo_store.clear()
+                ui.print_notice("History cleared.")
+            elif cmd == "compact":
+                if len(messages) <= 2:
+                    ui.print_notice("Nothing to compact yet.", style="yellow")
+                else:
+                    ui.print_notice("Compacting conversation…", style="dim")
+                    old_count = len(messages)
+                    try:
+                        summary = await compact_messages(messages, model_override)
+                    except Exception as e:
+                        ui.print_notice(f"Compaction failed: {e}", style="bold red")
+                        continue
+                    messages = [
+                        system_message_for(mode),
+                        HumanMessage(
+                            content="[Summary of earlier conversation]\n" + summary
+                        ),
+                    ]
+                    session.save(current_id, messages, model=model_name, mode=mode)
+                    ui.print_notice(
+                        f"Compacted {old_count} messages into a summary.",
+                        style="green",
+                    )
+            elif cmd == "help":
+                ui.print_help()
+            else:
+                ui.print_notice(f"Unknown command: /{cmd}  (try /help)", style="yellow")
+            continue
+
+        guard = check_user_input(raw)
+        if guard is not None:
+            ui.print_notice(f"Guardrail \u2014 {guard}", style="bold red")
+            continue
+
+        messages.append(HumanMessage(content=raw))
+        ui.print_user_message(raw)
+        messages = await run_turn(graph, messages, token_tracker)
+        session.save(current_id, messages, model=model_name, mode=mode)
+
+
+def cli_entry():
+    """Entry point for the `closecode` console script (see pyproject.toml
+    [project.scripts]). Installed via pip/pipx, runs the async main()."""
+    asyncio.run(main())
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

mcp_tools.py

@@ -0,0 +1,27 @@
+import os
+from langchain_mcp_adapters.client import MultiServerMCPClient
+
+
+def build_mcp_client(repo_path: str) -> MultiServerMCPClient:
+    return MultiServerMCPClient(
+        {
+            "git": {
+                "command": "python",
+                "args": ["-m", "mcp_server_git", "--repository", repo_path],
+                "transport": "stdio",
+            }
+        }
+    )
+
+
+async def get_git_tools(repo_path: str):
+    """Returns a list of LangChain-compatible tools backed by mcp-server-git.
+    Call this once at startup (it's async — see main.py)."""
+    client = build_mcp_client(repo_path)
+    return await client.get_tools()
+
+
+def get_git_repo_path() -> str:
+    """Where the git MCP server should operate. Defaults to the same
+    directory as the agent's sandbox, but can be pointed elsewhere."""
+    return os.environ.get("GIT_REPO_PATH", os.environ.get("AGENT_WORKDIR", "./sandbox"))

modes.py

@@ -0,0 +1,36 @@
+
+PLAN_MODE_BLOCKED_EXACT = {"bash", "write_file", "edit_file"}
+PLAN_MODE_BLOCKED_KEYWORDS = [
+    "commit", "push", "checkout", "reset", "merge", "rebase",
+    "add", "rm", "delete", "remove", "stash", "write", "create_branch",
+]
+
+
+def filter_tools_for_mode(tools: list, mode: str) -> list:
+    if mode == "build":
+        return list(tools)
+
+    kept = []
+    for t in tools:
+        name = t.name.lower()
+        if name in PLAN_MODE_BLOCKED_EXACT:
+            continue
+        if any(keyword in name for keyword in PLAN_MODE_BLOCKED_KEYWORDS):
+            continue
+        kept.append(t)
+    return kept
+
+
+def mode_system_note(mode: str) -> str:
+    if mode == "plan":
+        return (
+            "\n\nYou are currently in PLAN MODE. Tools that write files, edit files, "
+            "run shell commands, or change git state are not available to you right now. "
+            "Explore and read what you need, then respond with a clear, concrete, "
+            "step-by-step plan of the changes you would make. Do not claim to have made "
+            "changes you did not actually make — you can't, in this mode."
+        )
+    return (
+        "\n\nYou are currently in BUILD MODE. You have full access to read, write, edit, "
+        "run commands, and make git changes as needed to complete the task."
+    )