codeprobe 0.1.0__py3-none-any.whl

codeprobe/__init__.py

@@ -0,0 +1,3 @@
+"""codeprobe — Benchmark AI coding agents against your own codebase."""
+
+__version__ = "0.1.0a3"

codeprobe/__main__.py

@@ -0,0 +1,5 @@
+"""Allow ``python -m codeprobe`` and ``pipx run codeprobe``."""
+
+from codeprobe.cli import main
+
+main()

codeprobe/adapters/__init__.py

@@ -0,0 +1,19 @@
+"""Agent adapters — Protocol + built-in implementations."""
+
+from codeprobe.adapters.protocol import (
+    AdapterError,
+    AdapterExecutionError,
+    AdapterSetupError,
+    AgentAdapter,
+    AgentConfig,
+    AgentOutput,
+)
+
+__all__ = [
+    "AdapterError",
+    "AdapterExecutionError",
+    "AdapterSetupError",
+    "AgentAdapter",
+    "AgentConfig",
+    "AgentOutput",
+]

codeprobe/adapters/_base.py

@@ -0,0 +1,192 @@
+"""Shared base for agent adapters — eliminates duplicated run/preflight logic."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+import time
+from abc import abstractmethod
+
+from codeprobe.adapters.protocol import (
+    AdapterSetupError,
+    AgentConfig,
+    AgentOutput,
+)
+
+# Only these env vars are forwarded to agent subprocesses.
+# Keeps secrets (OPENAI_API_KEY, AWS_SECRET_*, etc.) out of the child
+# unless explicitly listed here.
+_ADAPTER_ENV_WHITELIST: frozenset[str] = frozenset(
+    {
+        # System essentials
+        "PATH",
+        "HOME",
+        "LANG",
+        "TERM",
+        "TMPDIR",
+        "LC_ALL",
+        # Codeprobe sandbox signal (eval harness sets this)
+        "CODEPROBE_SANDBOX",
+        # Agent-specific API keys (required by the adapters)
+        "ANTHROPIC_API_KEY",
+        "CLAUDE_CONFIG_DIR",
+        "GITHUB_TOKEN",
+        "OPENAI_API_KEY",
+        "COPILOT_API_KEY",
+        # Python toolchain
+        "VIRTUAL_ENV",
+        "PYTHONPATH",
+        # Node/npm (for copilot CLI)
+        "NODE_PATH",
+        "NPM_CONFIG_PREFIX",
+        # Go toolchain
+        "GOPATH",
+        "GOROOT",
+        # Rust toolchain
+        "CARGO_HOME",
+        "RUSTUP_HOME",
+    }
+)
+
+
+def _adapter_safe_env(extra: dict[str, str] | None = None) -> dict[str, str]:
+    """Build a filtered environment for agent subprocesses.
+
+    Only passes whitelisted vars — prevents leaking secrets from parent env.
+    """
+    env = {k: v for k, v in os.environ.items() if k in _ADAPTER_ENV_WHITELIST}
+    if extra:
+        env.update(extra)
+    return env
+
+
+class BaseAdapter:
+    """Base class for CLI-based agent adapters.
+
+    Subclasses set ``_binary_name`` and ``_install_hint``, then implement
+    ``build_command``.  The Protocol requires ``name``, ``preflight``, and
+    ``run``; ``find_binary`` and ``build_command`` are BaseAdapter helpers.
+    """
+
+    _binary_name: str
+    _install_hint: str
+
+    @property
+    def name(self) -> str:
+        return self._binary_name
+
+    def find_binary(self) -> str | None:
+        return shutil.which(self._binary_name)
+
+    def _require_binary(self) -> str:
+        """Return binary path or raise AdapterSetupError."""
+        binary = self.find_binary()
+        if binary is None:
+            raise AdapterSetupError(f"{self._binary_name} CLI not found")
+        return binary
+
+    def preflight(self, config: AgentConfig) -> list[str]:
+        issues: list[str] = []
+        if self.find_binary() is None:
+            issues.append(self._install_hint)
+        return issues
+
+    def isolate_session(self, slot_id: int) -> dict[str, str]:
+        """Default: no session isolation env overrides."""
+        return {}
+
+    @abstractmethod
+    def build_command(self, prompt: str, config: AgentConfig) -> list[str]: ...
+
+    def parse_output(
+        self, result: subprocess.CompletedProcess[str], duration: float
+    ) -> AgentOutput:
+        """Convert subprocess result to AgentOutput.
+
+        Subclasses override to extract tokens, cost, etc. from agent output.
+        """
+        return AgentOutput(
+            stdout=result.stdout,
+            stderr=result.stderr or None,
+            exit_code=result.returncode,
+            duration_seconds=duration,
+        )
+
+    def _write_mcp_config(self, config: AgentConfig) -> str | None:
+        """Write MCP config to a temp file if present. Returns path or None.
+
+        Expands ``${VAR}`` references in string values from the environment
+        so experiment.json can reference secrets without hardcoding them.
+        """
+        if not config.mcp_config:
+            return None
+        expanded = json.loads(os.path.expandvars(json.dumps(config.mcp_config)))
+        tmp = tempfile.NamedTemporaryFile(
+            mode="w", suffix=".json", prefix="codeprobe-mcp-", delete=False
+        )
+        json.dump(expanded, tmp)
+        tmp.close()
+        return tmp.name
+
+    def run(
+        self,
+        prompt: str,
+        config: AgentConfig,
+        session_env: dict[str, str] | None = None,
+    ) -> AgentOutput:
+        cmd = self.build_command(prompt, config)
+        mcp_tmpfile: str | None = None
+
+        # Find and track MCP temp file for cleanup
+        for flag in ("--mcp-config", "--additional-mcp-config"):
+            if flag in cmd:
+                idx = cmd.index(flag)
+                if idx + 1 < len(cmd):
+                    path = cmd[idx + 1]
+                    if path.startswith(tempfile.gettempdir()):
+                        mcp_tmpfile = path
+
+        start = time.monotonic()
+
+        try:
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                timeout=config.timeout_seconds,
+                cwd=config.cwd,
+                env=_adapter_safe_env(session_env),
+            )
+        except subprocess.TimeoutExpired as exc:
+            duration = time.monotonic() - start
+            return AgentOutput(
+                stdout=exc.stdout if isinstance(exc.stdout, str) else "",
+                stderr=exc.stderr if isinstance(exc.stderr, str) else None,
+                exit_code=-1,
+                duration_seconds=duration,
+                error=f"Agent timed out after {config.timeout_seconds}s",
+            )
+        except FileNotFoundError as exc:
+            raise AdapterSetupError(f"Binary not found at runtime: {exc}") from exc
+        finally:
+            if mcp_tmpfile:
+                try:
+                    os.unlink(mcp_tmpfile)
+                except OSError:
+                    pass
+
+        duration = time.monotonic() - start
+
+        try:
+            return self.parse_output(result, duration)
+        except Exception as exc:
+            return AgentOutput(
+                stdout=result.stdout,
+                stderr=result.stderr or None,
+                exit_code=result.returncode,
+                duration_seconds=duration,
+                error=f"Output parse failed: {exc}",
+            )

codeprobe/adapters/aider.py

@@ -0,0 +1,79 @@
+"""Aider CLI agent adapter."""
+
+from __future__ import annotations
+
+import re
+import subprocess
+
+from codeprobe.adapters._base import BaseAdapter
+from codeprobe.adapters.protocol import AgentConfig, AgentOutput
+
+# Matches: "Tokens: 1.2k sent, 856 received. Cost: $0.0034 message, $0.0034 session."
+# Token counts may be plain integers or k-suffixed floats (e.g., 45.3k).
+_TOKEN_RE = re.compile(r"Tokens:\s*([\d.]+k?)\s*sent,\s*([\d.]+k?)\s*received")
+_COST_RE = re.compile(r"Cost:\s*\$([\d.]+)\s*message")
+
+
+def _parse_token_value(raw: str) -> int:
+    """Parse a token count string like '1.2k' or '856' into an integer."""
+    raw = raw.strip().lower()
+    if raw.endswith("k"):
+        return int(float(raw[:-1]) * 1000)
+    return int(float(raw))
+
+
+class AiderAdapter(BaseAdapter):
+    """Adapter for Aider CLI (aider --message)."""
+
+    _binary_name = "aider"
+    _install_hint = "Aider CLI not found. Install with: pip install aider-chat"
+
+    def build_command(self, prompt: str, config: AgentConfig) -> list[str]:
+        binary = self._require_binary()
+        cmd = [binary, "--message", prompt, "--yes-always", "--no-git"]
+
+        if config.model:
+            cmd.extend(["--model", config.model])
+
+        return cmd
+
+    def parse_output(
+        self, result: subprocess.CompletedProcess[str], duration: float
+    ) -> AgentOutput:
+        """Parse Aider CLI output for token counts and cost.
+
+        Aider prints a summary line like:
+            Tokens: 1.2k sent, 856 received. Cost: $0.0034 message, $0.0034 session.
+
+        This may appear in stdout or stderr. We search both.
+        """
+        combined = (result.stdout or "") + "\n" + (result.stderr or "")
+
+        input_tokens: int | None = None
+        output_tokens: int | None = None
+        cost_usd: float | None = None
+        cost_model = "unknown"
+        cost_source = "unavailable"
+
+        token_match = _TOKEN_RE.search(combined)
+        if token_match:
+            input_tokens = _parse_token_value(token_match.group(1))
+            output_tokens = _parse_token_value(token_match.group(2))
+
+        cost_match = _COST_RE.search(combined)
+        if cost_match:
+            cost_usd = float(cost_match.group(1))
+            cost_model = "per_token"
+            cost_source = "log_parsed"
+
+        return AgentOutput(
+            stdout=result.stdout,
+            stderr=result.stderr or None,
+            exit_code=result.returncode,
+            duration_seconds=duration,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cost_usd=cost_usd,
+            cost_model=cost_model,
+            cost_source=cost_source,
+        )

codeprobe/adapters/claude.py

@@ -0,0 +1,109 @@
+"""Claude Code agent adapter."""
+
+from __future__ import annotations
+
+import json
+import re
+import subprocess
+import tempfile
+from pathlib import Path
+
+from codeprobe.adapters._base import BaseAdapter
+from codeprobe.adapters.protocol import (
+    ALLOWED_PERMISSION_MODES,
+    AgentConfig,
+    AgentOutput,
+)
+from codeprobe.adapters.telemetry import JsonStdoutCollector
+from codeprobe.core.sandbox import is_sandboxed
+
+# Claude CLI accepts aliases (sonnet, opus, haiku) or short model IDs
+# (claude-sonnet-4-6) but NOT full API model IDs with date suffixes
+# (claude-sonnet-4-6-20250514). Strip the date suffix when present.
+_API_MODEL_DATE_SUFFIX = re.compile(r"(-\d{8})$")
+
+
+def _normalize_model_for_cli(model: str) -> str:
+    """Normalize a model identifier for the Claude CLI.
+
+    Strips date suffixes from full API model IDs so the CLI can resolve them.
+    Aliases like 'sonnet' or 'haiku' pass through unchanged.
+    """
+    return _API_MODEL_DATE_SUFFIX.sub("", model)
+
+
+class ClaudeAdapter(BaseAdapter):
+    """Adapter for Claude Code CLI (claude -p)."""
+
+    _binary_name = "claude"
+    _install_hint = "Claude CLI not found. Install from https://claude.ai/download"
+
+    def __init__(self) -> None:
+        self._collector = JsonStdoutCollector()
+
+    def preflight(self, config: AgentConfig) -> list[str]:
+        issues = super().preflight(config)
+        if config.permission_mode == "dangerously_skip" and not is_sandboxed():
+            issues.append(
+                "permission_mode='dangerously_skip' requires a sandboxed environment "
+                "(Docker container or CODEPROBE_SANDBOX=1)"
+            )
+        return issues
+
+    def build_command(self, prompt: str, config: AgentConfig) -> list[str]:
+        binary = self._require_binary()
+        cmd = [binary, "-p", prompt, "--output-format", "json"]
+
+        if config.model:
+            cmd.extend(["--model", _normalize_model_for_cli(config.model)])
+
+        if config.permission_mode == "dangerously_skip":
+            cmd.append("--dangerously-skip-permissions")
+        elif config.permission_mode != "default":
+            if config.permission_mode not in ALLOWED_PERMISSION_MODES:
+                raise ValueError(
+                    f"Unsafe permission_mode: {config.permission_mode!r}. "
+                    f"Allowed: {', '.join(sorted(ALLOWED_PERMISSION_MODES))}"
+                )
+            cmd.extend(["--permission-mode", config.permission_mode])
+
+        mcp_path = self._write_mcp_config(config)
+        if mcp_path:
+            cmd.extend(["--mcp-config", mcp_path])
+
+        return cmd
+
+    def isolate_session(self, slot_id: int) -> dict[str, str]:
+        """Return a per-slot CLAUDE_CONFIG_DIR for session isolation."""
+        config_dir = (
+            Path(tempfile.gettempdir()) / "codeprobe-claude" / f"slot-{slot_id}"
+        )
+        config_dir.mkdir(parents=True, exist_ok=True)
+        return {"CLAUDE_CONFIG_DIR": str(config_dir)}
+
+    def parse_output(
+        self, result: subprocess.CompletedProcess[str], duration: float
+    ) -> AgentOutput:
+        """Parse Claude CLI JSON envelope into AgentOutput."""
+        usage = self._collector.collect(result.stdout)
+
+        # Extract content text from the JSON envelope
+        try:
+            envelope = json.loads(result.stdout)
+            stdout_text = envelope.get("result", result.stdout)
+        except (json.JSONDecodeError, ValueError):
+            stdout_text = result.stdout
+
+        return AgentOutput(
+            stdout=stdout_text,
+            stderr=result.stderr or None,
+            exit_code=result.returncode,
+            duration_seconds=duration,
+            input_tokens=usage.input_tokens,
+            output_tokens=usage.output_tokens,
+            cache_read_tokens=usage.cache_read_tokens,
+            cost_usd=usage.cost_usd,
+            cost_model=usage.cost_model,
+            cost_source=usage.cost_source,
+            error=usage.error,
+        )

codeprobe/adapters/codex.py

@@ -0,0 +1,135 @@
+"""OpenAI Codex agent adapter — API-based (no CLI subprocess)."""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+
+from codeprobe.adapters.protocol import (
+    AdapterExecutionError,
+    AdapterSetupError,
+    AgentConfig,
+    AgentOutput,
+)
+from codeprobe.adapters.telemetry import ApiResponseCollector
+
+logger = logging.getLogger(__name__)
+
+
+def _usage_fields(
+    usage: object | None, input_attr: str, output_attr: str
+) -> tuple[int | None, int | None]:
+    """Extract input/output token counts from an API usage object."""
+    if usage is None:
+        return None, None
+    return getattr(usage, input_attr, None), getattr(usage, output_attr, None)
+
+
+class CodexAdapter:
+    """Adapter for OpenAI Codex API.
+
+    Tries the Responses API (responses.create) first. If the model is not
+    available on that endpoint (NotFoundError), falls back to the Chat
+    Completions API (chat.completions.create).
+    """
+
+    def __init__(self) -> None:
+        self._collector = ApiResponseCollector()
+
+    @property
+    def name(self) -> str:
+        return "codex"
+
+    def preflight(self, config: AgentConfig) -> list[str]:
+        issues: list[str] = []
+        try:
+            import openai  # noqa: F401
+        except ImportError:
+            issues.append(
+                "openai SDK not found. Install with: pip install codeprobe[codex]"
+            )
+            return issues
+        if not os.environ.get("OPENAI_API_KEY"):
+            issues.append("OPENAI_API_KEY environment variable not set")
+        return issues
+
+    def isolate_session(self, slot_id: int) -> dict[str, str]:
+        """Codex uses API calls — no session-level isolation needed."""
+        return {}
+
+    def run(
+        self,
+        prompt: str,
+        config: AgentConfig,
+        session_env: dict[str, str] | None = None,
+    ) -> AgentOutput:
+        try:
+            import openai
+        except ImportError:
+            raise AdapterSetupError(
+                "openai SDK not installed. Run: pip install codeprobe[codex]"
+            )
+
+        client = openai.OpenAI()
+        model = config.model or "codex-mini-latest"
+        start = time.monotonic()
+
+        try:
+            try:
+                response = client.responses.create(model=model, input=prompt)
+                stdout = response.output_text or ""
+                input_tokens, output_tokens = _usage_fields(
+                    response.usage, "input_tokens", "output_tokens"
+                )
+            except openai.NotFoundError:
+                logger.info(
+                    "Model %s not found on Responses API, falling back to "
+                    "Chat Completions API",
+                    model,
+                )
+                try:
+                    response = client.chat.completions.create(
+                        model=model,
+                        messages=[{"role": "user", "content": prompt}],
+                    )
+                except openai.NotFoundError as exc:
+                    raise AdapterExecutionError(
+                        f"Model {model!r} not available on Responses or "
+                        f"Chat Completions API: {exc}"
+                    ) from exc
+                content = (
+                    response.choices[0].message.content if response.choices else None
+                )
+                stdout = content or ""
+                input_tokens, output_tokens = _usage_fields(
+                    response.usage, "prompt_tokens", "completion_tokens"
+                )
+        except openai.AuthenticationError as exc:
+            raise AdapterSetupError(f"OPENAI_API_KEY invalid: {exc}") from exc
+        except openai.RateLimitError as exc:
+            raise AdapterExecutionError(f"Rate limited: {exc}") from exc
+        except openai.APIError as exc:
+            raise AdapterExecutionError(f"OpenAI API error: {exc}") from exc
+
+        duration = time.monotonic() - start
+
+        usage = self._collector.collect(
+            stdout,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            model=model,
+        )
+
+        return AgentOutput(
+            stdout=stdout,
+            stderr=None,
+            exit_code=0,
+            duration_seconds=duration,
+            input_tokens=usage.input_tokens,
+            output_tokens=usage.output_tokens,
+            cost_usd=usage.cost_usd,
+            cost_model=usage.cost_model,
+            cost_source=usage.cost_source,
+            error=usage.error,
+        )

codeprobe/adapters/copilot.py

@@ -0,0 +1,104 @@
+"""GitHub Copilot CLI agent adapter."""
+
+from __future__ import annotations
+
+import json
+import logging
+import subprocess
+
+from codeprobe.adapters._base import BaseAdapter
+from codeprobe.adapters.protocol import AgentConfig, AgentOutput
+from codeprobe.adapters.telemetry import NdjsonStreamCollector
+
+logger = logging.getLogger(__name__)
+
+
+class CopilotAdapter(BaseAdapter):
+    """Adapter for GitHub Copilot CLI."""
+
+    _binary_name = "copilot"
+    _install_hint = (
+        "Copilot CLI not found. Install from https://github.com/github/copilot-cli"
+    )
+
+    def __init__(self) -> None:
+        self._collector = NdjsonStreamCollector()
+
+    def preflight(self, config: AgentConfig) -> list[str]:
+        return super().preflight(config)
+
+    def build_command(self, prompt: str, config: AgentConfig) -> list[str]:
+        binary = self._require_binary()
+        cmd = [binary, "--prompt", prompt, "--output-format", "json"]
+
+        # Non-interactive mode requires --allow-all-tools for tool auto-approval
+        cmd.append("--allow-all-tools")
+
+        if config.model:
+            cmd.extend(["--model", config.model])
+
+        mcp_path = self._write_mcp_config(config)
+        if mcp_path:
+            cmd.extend(["--additional-mcp-config", f"@{mcp_path}"])
+
+        return cmd
+
+    def parse_output(
+        self, result: subprocess.CompletedProcess[str], duration: float
+    ) -> AgentOutput:
+        """Parse Copilot CLI NDJSON output for token data.
+
+        Requires Copilot CLI 1.0.4+ with --output-format json which emits
+        NDJSON lines containing "assistant.message" events with outputTokens.
+
+        Input tokens are extracted from NDJSON ``usage`` events when available,
+        falling back to Copilot process log parsing.
+        """
+        raw = result.stdout or ""
+        usage = self._collector.collect(raw)
+
+        # Extract content text from NDJSON events.
+        # On JSON parse failure, the except clause resets to empty,
+        # and the fallback below uses raw output — matching original behavior.
+        result_text_parts: list[str] = []
+        try:
+            for line in raw.strip().splitlines():
+                if not line.strip():
+                    continue
+                obj = json.loads(line)
+                event_type = obj.get("type", "")
+                if event_type == "assistant.message":
+                    content = obj.get("data", {}).get("content", "")
+                    if content:
+                        result_text_parts.append(content)
+                elif event_type == "result":
+                    content = obj.get("data", {}).get("content", "")
+                    if content:
+                        result_text_parts.append(content)
+        except json.JSONDecodeError:
+            logger.warning(
+                "ndjson_parse_fallback: failed to parse Copilot NDJSON output, "
+                "falling back to raw stdout"
+            )
+            result_text_parts = []
+        stdout_text = "\n".join(result_text_parts) if result_text_parts else raw
+
+        # When NDJSON parsing fell back to raw output, surface the fallback
+        # in the error field so callers can detect degraded telemetry.
+        error = usage.error
+        if not result_text_parts and raw:
+            fallback_msg = "ndjson_parse_fallback: raw stdout used as output"
+            error = f"{error}; {fallback_msg}" if error else fallback_msg
+
+        return AgentOutput(
+            stdout=stdout_text,
+            stderr=result.stderr or None,
+            exit_code=result.returncode,
+            duration_seconds=duration,
+            input_tokens=usage.input_tokens,
+            output_tokens=usage.output_tokens,
+            cost_usd=usage.cost_usd,
+            cost_model=usage.cost_model,
+            cost_source=usage.cost_source,
+            error=error,
+        )