atomsh 0.1.0__py3-none-any.whl

atomsh/__init__.py

@@ -0,0 +1,6 @@
+"""Atomsh — your autonomous coding agent for science.
+
+Installed as the `atomsh` command; entry point is `atomsh.cli:main`.
+"""
+
+__version__ = "0.1.0"

atomsh/agent.py

@@ -0,0 +1,147 @@
+"""The agent loop: model → tool calls → tool results → model, until it stops."""
+
+import json
+from pathlib import Path
+
+from . import tools as toolkit
+from .client import AtomGPT, AtomGPTError
+from .config import MAX_STEPS
+from .interrupt import escape_watch
+from .permissions import ALLOW, Permissions
+from .prompt import system_prompt
+
+DIM = "\033[2m"
+BOLD = "\033[1m"
+RESET = "\033[0m"
+
+
+class Agent:
+    """Drives one conversation against one model with one tool set."""
+
+    def __init__(self, client: AtomGPT, model: str, session,
+                 permissions: Permissions = None, root: Path = None,
+                 color: bool = True, remote=None):
+        self.client = client
+        self.model = model
+        self.session = session
+        self.root = root or Path.cwd()
+        self.permissions = permissions or Permissions("ask", self.root)
+        self.color = color
+        # An MCPClient, when --materials is on. Its tools sit alongside the
+        # local ones; the model does not need to know which is which.
+        self.remote = remote
+        self.remote_tools = set()
+        self.tool_schema = list(toolkit.SCHEMA)
+        if remote is not None:
+            remote_schema = remote.schema()
+            self.remote_tools = {t["function"]["name"] for t in remote_schema}
+            self.tool_schema += remote_schema
+        if not self.session.messages:
+            self.session.messages.append({
+                "role": "system",
+                "content": system_prompt(str(self.root),
+                                         materials=remote is not None),
+            })
+
+    def _dim(self, text: str) -> str:
+        return f"{DIM}{text}{RESET}" if self.color else text
+
+    def run(self, user_text: str) -> str:
+        """Run one user turn to completion. Returns the final assistant text."""
+        self.session.messages.append({"role": "user", "content": user_text})
+        final = ""
+
+        for _ in range(MAX_STEPS):
+            # `last` doubles as "have we printed anything this step": stripped
+            # channel markers leave stray blank lines, so leading whitespace is
+            # swallowed and the closing newline is only added if one is missing.
+            state = {"last": ""}
+
+            def on_text(piece, state=state):
+                if not state["last"] and not piece.strip():
+                    return
+                state["last"] = piece
+                print(piece, end="", flush=True)
+
+            try:
+                with escape_watch() as cancelled:
+                    result = self.client.stream(
+                        self.session.messages, self.tool_schema, self.model,
+                        on_text=on_text, cancelled=cancelled,
+                    )
+            except AtomGPTError as e:
+                print(f"\n{self._dim('error:')} {e}")
+                return ""
+
+            message = result["message"]
+            if state["last"] and not state["last"].endswith("\n"):
+                print()
+            self.session.messages.append(message)
+            self.session.save()
+
+            if result.get("finish_reason") == "cancelled":
+                print(self._dim("  interrupted"))
+                return message.get("content") or ""
+
+            calls = message.get("tool_calls")
+            if not calls:
+                final = message.get("content") or ""
+                if not final and not state["last"]:
+                    print(self._dim("(no response — try rephrasing)"))
+                return final
+
+            for call in calls:
+                self._run_tool(call)
+            self.session.save()
+
+        print(self._dim(f"stopped after {MAX_STEPS} steps"))
+        return final
+
+    def _run_tool(self, call: dict) -> None:
+        """Execute one tool call and append its result to the conversation."""
+        name = (call.get("function") or {}).get("name") or ""
+        raw_args = (call.get("function") or {}).get("arguments") or "{}"
+        try:
+            args = json.loads(raw_args) if raw_args.strip() else {}
+        except ValueError:
+            args = {}
+
+        handler = toolkit.HANDLERS.get(name)
+        if handler is None and name not in self.remote_tools:
+            return self._reply(call, f"Error: no such tool {name!r}.")
+
+        decision, reason = self.permissions.check(name, args)
+        if decision != ALLOW:
+            print(self._dim(f"  ✗ {name} — {reason}"))
+            return self._reply(call, f"Denied: {reason}")
+
+        print(self._dim(f"  · {name}({self._summarize(args)})"))
+        try:
+            if handler is None:
+                output = self.remote.call(name, args)
+            else:
+                output = handler(self.root, **args)
+        except TypeError as e:
+            output = f"Error: bad arguments for {name}: {e}"
+        except Exception as e:  # a tool must never kill the session
+            output = f"Error: {name} raised {type(e).__name__}: {e}"
+        self._reply(call, output)
+
+    def _reply(self, call: dict, content: str) -> None:
+        self.session.messages.append({
+            "role": "tool",
+            "tool_call_id": call.get("id") or "",
+            "content": content,
+        })
+
+    @staticmethod
+    def _summarize(args: dict) -> str:
+        """One-line rendering of tool arguments for the activity log."""
+        parts = []
+        for key, value in args.items():
+            text = value if isinstance(value, str) else json.dumps(value)
+            text = text.replace("\n", " ")
+            if len(text) > 60:
+                text = text[:60] + "…"
+            parts.append(f"{key}={text}")
+        return ", ".join(parts)

atomsh/auth.py

@@ -0,0 +1,206 @@
+"""Authentication against atomgpt.org.
+
+The preferred flow is `atomsh login`: an OAuth 2.1 authorization-code
+exchange with PKCE against atomgpt.org, using a loopback redirect. The
+authorization server hands back the user's existing atomgpt.org API key as the
+access token, so the result is a Bearer credential usable against /api.
+
+`atomsh login --key` is the fallback for headless machines, where opening a
+browser is not possible.
+"""
+
+import base64
+import hashlib
+import json
+import os
+import secrets
+import socket
+import threading
+import urllib.parse
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+import httpx
+
+from .browser import open_url
+from .config import (
+    API_URL,
+    AUTHORIZE_URL,
+    AUTH_FILE,
+    CLIENT_NAME,
+    REGISTER_URL,
+    TOKEN_URL,
+)
+
+
+class AuthError(Exception):
+    """Login failed, or the stored credential is not usable."""
+
+
+# ── stored credential ────────────────────────────────────────────────────────
+
+def load_token() -> str:
+    """Return the stored token, or None. ATOMSH_API_KEY wins if set."""
+    env = os.environ.get("ATOMSH_API_KEY")
+    if env:
+        return env.strip()
+    try:
+        with open(AUTH_FILE, encoding="utf-8") as fh:
+            return json.load(fh).get("access_token")
+    except (OSError, ValueError):
+        return None
+
+
+def save_token(token: str) -> None:
+    """Persist the token 0600, creating the config dir if needed."""
+    AUTH_FILE.parent.mkdir(parents=True, exist_ok=True)
+    # Create with restrictive permissions before writing the secret.
+    fd = os.open(AUTH_FILE, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "w", encoding="utf-8") as fh:
+        json.dump({"access_token": token}, fh)
+
+
+def forget_token() -> bool:
+    """Delete the stored credential. Returns whether there was one."""
+    try:
+        AUTH_FILE.unlink()
+        return True
+    except OSError:
+        return False
+
+
+def whoami(token: str) -> dict:
+    """Validate a token and return the account it belongs to.
+
+    /api/models is the cheapest authenticated endpoint that exists on every
+    deployment, so it doubles as the credential check.
+    """
+    try:
+        r = httpx.get(
+            f"{API_URL}/models",
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=30,
+        )
+    except httpx.HTTPError as e:
+        raise AuthError(f"could not reach {API_URL}: {e}") from e
+    if r.status_code in (401, 403):
+        raise AuthError("token rejected by atomgpt.org (401/403)")
+    r.raise_for_status()
+    models = [m.get("id") for m in (r.json().get("data") or [])]
+    return {"models": models}
+
+
+# ── OAuth login ──────────────────────────────────────────────────────────────
+
+def _pkce_pair() -> tuple:
+    verifier = base64.urlsafe_b64encode(secrets.token_bytes(64)).decode().rstrip("=")
+    digest = hashlib.sha256(verifier.encode()).digest()
+    challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=")
+    return verifier, challenge
+
+
+class _CallbackHandler(BaseHTTPRequestHandler):
+    """Single-shot handler that captures ?code=&state= from the redirect."""
+
+    result = None
+
+    def do_GET(self):  # noqa: N802 — name fixed by BaseHTTPRequestHandler
+        params = urllib.parse.parse_qs(urllib.parse.urlparse(self.path).query)
+        _CallbackHandler.result = {k: v[0] for k, v in params.items()}
+        ok = "code" in _CallbackHandler.result
+        body = (
+            "<h2>atomsh is connected.</h2><p>You can close this tab.</p>"
+            if ok
+            else "<h2>Authorization failed.</h2><p>Return to the terminal.</p>"
+        )
+        payload = f"<!doctype html><meta charset=utf-8><body style='font-family:system-ui;margin:80px auto;max-width:420px'>{body}</body>".encode()
+        self.send_response(200 if ok else 400)
+        self.send_header("Content-Type", "text/html; charset=utf-8")
+        self.send_header("Content-Length", str(len(payload)))
+        self.end_headers()
+        self.wfile.write(payload)
+
+    def log_message(self, *args):
+        """Silence the default stderr access log."""
+
+
+def _free_port() -> int:
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def login_oauth(open_browser: bool = True, timeout: int = 300) -> str:
+    """Run the browser login and return the access token.
+
+    Starts a loopback listener, registers this client, sends the user to the
+    consent screen, then exchanges the returned code for a token.
+    """
+    port = _free_port()
+    redirect_uri = f"http://127.0.0.1:{port}/callback"
+    verifier, challenge = _pkce_pair()
+    state = secrets.token_urlsafe(16)
+
+    try:
+        reg = httpx.post(
+            REGISTER_URL,
+            json={"client_name": CLIENT_NAME, "redirect_uris": [redirect_uri]},
+            timeout=30,
+        )
+        reg.raise_for_status()
+        client_id = reg.json()["client_id"]
+    except (httpx.HTTPError, KeyError, ValueError) as e:
+        raise AuthError(f"client registration failed: {e}") from e
+
+    url = AUTHORIZE_URL + "?" + urllib.parse.urlencode({
+        "response_type": "code",
+        "client_id": client_id,
+        "redirect_uri": redirect_uri,
+        "code_challenge": challenge,
+        "code_challenge_method": "S256",
+        "state": state,
+        "scope": "mcp",
+    })
+
+    server = HTTPServer(("127.0.0.1", port), _CallbackHandler)
+    server.timeout = timeout
+    _CallbackHandler.result = None
+
+    print("Open this URL to authorize atomsh:\n")
+    print(f"  {url}\n")
+    if open_browser:
+        # In a thread: launching a Windows browser from WSL can block for
+        # seconds, and the callback listener should already be waiting.
+        threading.Thread(target=open_url, args=(url,), daemon=True).start()
+        print("Trying to open it in your browser…")
+    print(f"Waiting for authorization (Ctrl-C to cancel, {timeout}s timeout).")
+
+    server.handle_request()  # blocks until the redirect arrives or times out
+    server.server_close()
+
+    result = _CallbackHandler.result
+    if not result:
+        raise AuthError("timed out waiting for the browser redirect")
+    if result.get("error"):
+        raise AuthError(f"authorization denied: {result['error']}")
+    if result.get("state") != state:
+        raise AuthError("state mismatch — aborting")
+
+    try:
+        tok = httpx.post(
+            TOKEN_URL,
+            data={
+                "grant_type": "authorization_code",
+                "code": result["code"],
+                "redirect_uri": redirect_uri,
+                "client_id": client_id,
+                "code_verifier": verifier,
+            },
+            timeout=30,
+        )
+        tok.raise_for_status()
+        token = tok.json()["access_token"]
+    except (httpx.HTTPError, KeyError, ValueError) as e:
+        raise AuthError(f"token exchange failed: {e}") from e
+
+    save_token(token)
+    return token

atomsh/browser.py

@@ -0,0 +1,71 @@
+"""Opening a URL in the user's browser, including from WSL.
+
+Python's webbrowser module shells out to xdg-open, which on a headless Linux
+box (WSL included) prints a dozen "not found" lines to stderr before giving
+up. That noise is worse than useless during login, so this module picks an
+opener deliberately and keeps every child process quiet.
+"""
+
+import os
+import shutil
+import subprocess
+import webbrowser
+
+
+def is_wsl() -> bool:
+    """Whether we are running under WSL, where the browser lives on Windows."""
+    try:
+        with open("/proc/version", encoding="utf-8", errors="ignore") as fh:
+            return "microsoft" in fh.read().lower()
+    except OSError:
+        return False
+
+
+def _run(cmd: list) -> bool:
+    """Launch a command with its output discarded. True if it started."""
+    try:
+        subprocess.run(cmd, stdout=subprocess.DEVNULL,
+                       stderr=subprocess.DEVNULL, timeout=20, check=False)
+        return True
+    except (OSError, subprocess.SubprocessError):
+        return False
+
+
+def _candidates(url: str) -> list:
+    cmds = []
+    if is_wsl():
+        # Hand the URL to Windows. explorer.exe exits non-zero even when it
+        # works, so "did not raise" is the success test, not the exit code.
+        cmds += [
+            ["wslview", url],
+            ["explorer.exe", url],
+            ["powershell.exe", "-NoProfile", "-Command", "Start-Process", url],
+        ]
+    cmds += [["xdg-open", url], ["open", url]]
+    return cmds
+
+
+def _webbrowser_quiet(url: str) -> bool:
+    """Last resort: webbrowser.open with the child's stderr muted."""
+    try:
+        devnull = os.open(os.devnull, os.O_WRONLY)
+    except OSError:
+        return webbrowser.open(url)
+    saved = os.dup(2)
+    try:
+        os.dup2(devnull, 2)
+        return webbrowser.open(url)
+    except Exception:
+        return False
+    finally:
+        os.dup2(saved, 2)
+        os.close(saved)
+        os.close(devnull)
+
+
+def open_url(url: str) -> bool:
+    """Open `url` in a browser. False means the user has to do it themselves."""
+    for cmd in _candidates(url):
+        if shutil.which(cmd[0]) and _run(cmd):
+            return True
+    return _webbrowser_quiet(url)